Commit Graph

86 Commits

Author SHA1 Message Date
Matthias Springer
77f8297c6f [mlir][sparse] Improve sparse tensor type constraints (#112133)
Sparse tensors are always ranked tensors. Encodings cannot be attached
to unranked tensors. Change the type constraint to `RankedTensorOf`, so
that we generate `TypedValue<RankedTensorType>` instead of
`TypedValue<TensorType>`. This removes the need for type casting in some
cases.

Also improve the verifiers (missing `return` statements) and switch a
few other `AnyTensor` to `AnyRankedTensor`.

This commit is in preparation of a dialect conversion commit that
required fixes in the sparse dialect.
2024-10-13 21:12:38 +02:00
Peiming Liu
785a24f156 [mlir][sparse] introduce sparse_tensor.coiterate operation. (#101100)
This PR introduces `sparse_tensor.coiterate` operation, which represents
a loop that traverses multiple sparse iteration space.
2024-07-31 15:14:27 -07:00
Peiming Liu
12189f8005 [mlir][sparse] introduce sparse_tensor.extract_value operation. (#101219) 2024-07-30 16:26:18 -07:00
Peiming Liu
e276cf0831 [mlir][sparse] introduce sparse_tensor.iterate operation (#88955)
A `sparse_tensor.iterate` iterates over a sparse iteration space
extracted from `sparse_tensor.extract_iteration_space` operation
introduced in https://github.com/llvm/llvm-project/pull/88554.
2024-06-10 10:20:24 -07:00
Peiming Liu
481bd5d416 [mlir][sparse] introduce sparse_tensor.extract_iteration_space operation. (#88554)
A `sparse_tensor.extract_space %tensor at %iterator` extracts a *sparse*
iteration space defined `%tensor`, the operation to traverse the
iteration space will be introduced in following PRs.
2024-04-16 11:32:30 -07:00
Peiming Liu
b9556532c7 Revert "[mlir][sparse] introduce sparse_tensor.iterate operation" (#88953)
Reverts llvm/llvm-project#88807 (merged by mistake)
2024-04-16 11:31:33 -07:00
Peiming Liu
8debcf03c5 [mlir][sparse] introduce sparse_tensor.iterate operation (#88807)
A `sparse_tensor.iterate` iterates over a sparse iteration space
extracted from `sparse_tensor.extract_iteration_space` operation
introduced in https://github.com/llvm/llvm-project/pull/88554.

*DO NOT MERGE* before https://github.com/llvm/llvm-project/pull/88554
2024-04-16 11:31:09 -07:00
Aart Bik
f388a3a446 [mlir][sparse] update doc and examples of the [dis]assemble operations (#88213)
The doc and examples of the [dis]assemble operations did not reflect all
the recent changes on order of the operands. Also clarified some of the
text.
2024-04-10 09:42:12 -07:00
Aart Bik
f3a8af07fa [mlir][sparse] best effort finalization of escaping empty sparse tensors (#85482)
This change lifts the restriction that purely allocated empty sparse
tensors cannot escape the method. Instead it makes a best effort to add
a finalizing operation before the escape.

This assumes that
(1) we never build sparse tensors across method boundaries
    (e.g. allocate in one, insert in other method)
(2) if we have other uses of the empty allocation in the
    same method, we assume that either that op will fail
    or will do the finalization for us.

This is best-effort, but fixes some very obvious missing cases.
2024-03-15 16:43:09 -07:00
Peiming Liu
94e27c265a [mlir][sparse] reuse tensor.insert operation to insert elements into … (#84987)
…a sparse tensor.
2024-03-12 16:59:17 -07:00
Peiming Liu
fc9f1d49aa [mlir][sparse] use a consistent order between [dis]assembleOp and sto… (#84079)
…rage layout.
2024-03-06 09:57:41 -08:00
Aart Bik
8dfa1d878d [mlir][sparse] add roundtrip and invalid tests for sparse_tensor.print (#83349) 2024-02-28 14:52:43 -08:00
Christian Ulmann
7ed96b1c0d [MLIR][LLVM] Remove last typed pointer remnants from tests (#71232)
This commit removes all LLVM dialect typed pointers from the lit tests.
Typed pointers have been deprecated for a while now and it's planned to
soon remove them from the LLVM dialect.

Related PSA:
https://discourse.llvm.org/t/psa-removal-of-typed-pointers-from-the-llvm-dialect/74502
2023-11-04 14:13:31 +01:00
Peiming Liu
d808d922b4 [mlir][sparse] introduce sparse_tensor.reinterpret_map operation. (#70378) 2023-10-26 15:04:09 -07:00
Aart Bik
7e83a1af5d [mlir][sparse] add verification of absent value in sparse_tensor.unary (#70248)
This value should always be a plain contant or something invariant
computed outside the surrounding linalg operation, since there is no
co-iteration defined on anything done in this branch.

Fixes:
https://github.com/llvm/llvm-project/issues/69395
2023-10-25 13:56:43 -07:00
Peiming Liu
f0f5fdf73d [mlir][sparse] introduce sparse_tensor.lvl operation. (#69978) 2023-10-23 15:49:39 -07:00
Peiming Liu
ff21a90e51 [mlir][sparse] introduce sparse_tensor.crd_translate operation (#69630) 2023-10-19 15:42:09 -07:00
Peiming Liu
0aacc2137a [mlir][sparse] introduce sparse_tensor.reorder_coo operation (#68827) 2023-10-12 09:42:12 -07:00
Peiming Liu
0083f8338c [mlir][sparse] renaming sparse_tensor.sort_coo to sparse_tensor.sort (#68161)
Rationale: the operation does not always sort COO tensors (also used for
sparse_tensor.compress for example).
2023-10-03 16:28:25 -07:00
Peiming Liu
6ca47eb49d [mlir][sparse] rename sparse_tensor.(un)pack to sparse_tensor.(dis)as… (#67717)
…semble

Pack/Unpack are overridden in many other places, rename the operations
to avoid confusion.
2023-09-28 11:01:10 -07:00
Yinying Li
256ac4619b [mlir][sparse] Change tests to use new syntax for ELL and slice (#67569)
Examples:

1. `#ELL = #sparse_tensor.encoding<{ lvlTypes = [ "dense", "dense",
"compressed" ], dimToLvl = affine_map<(i,j)[c] -> (c*4*i, i, j)>
}>`
to
`#ELL = #sparse_tensor.encoding<{ map = [s0](d0, d1) -> (d0 * (s0 * 4) :
dense, d0 : dense, d1 : compressed)
}>`

2. `#CSR_SLICE = #sparse_tensor.encoding<{ lvlTypes = [ "dense",
"compressed" ], dimSlices = [ (1, 4, 1), (1, 4, 2) ]
}>`
to
`#CSR_SLICE = #sparse_tensor.encoding<{ map = (d0 :
#sparse_tensor<slice(1, 4, 1)>, d1 : #sparse_tensor<slice(1, 4, 2)>) ->
(d0 : dense, d1 : compressed)
}>`
2023-09-27 19:40:52 -04:00
Peiming Liu
bfa3bc4378 [mlir][sparse] unifies sparse_tensor.sort_coo/sort into one operation. (#66722)
The use cases of the two operations are largely overlapped, let's
simplify it and only use one of them.
2023-09-19 17:02:32 -07:00
Yinying Li
3dc621124f [mlir][sparse] Migrate tests to use new syntax (#66543)
**COO**
`lvlTypes = [ "compressed_nu", "singleton" ]` to `map = (d0, d1) -> (d0
: compressed(nonunique), d1 : singleton)`
`lvlTypes = [ "compressed_nu_no", "singleton_no" ]` to `map = (d0, d1)
-> (d0 : compressed(nonunique, nonordered), d1 : singleton(nonordered))`

**SortedCOO**
`lvlTypes = [ "compressed_nu", "singleton" ]` to `map = (d0, d1) -> (d0
: compressed(nonunique), d1 : singleton)`

**BCOO**
`lvlTypes = [ "dense", "compressed_hi_nu", "singleton" ]` to `map = (d0,
d1, d2) -> (d0 : dense, d1 : compressed(nonunique, high), d2 :
singleton)`

**BCSR**
`lvlTypes = [ "compressed", "compressed", "dense", "dense" ], dimToLvl =
affine_map<(d0, d1) -> (d0 floordiv 2, d1 floordiv 3, d0 mod 2, d1 mod
3)>` to
`map = ( i, j ) ->
      ( i floordiv 2 : compressed,
        j floordiv 3 : compressed,
        i mod 2 : dense,
        j mod 3 : dense
      )`

**Tensor and other supported formats(e.g. CCC, CDC, CCCC)**

Currently, ELL and slice are not supported yet in the new syntax and the
CHECK tests will be updated once printing is set to output the new
syntax.

Previous PRs: #66146, #66309, #66443
2023-09-15 16:12:20 -04:00
Yinying Li
2a07f0fd40 [mlir][sparse] Migrate more tests to use new syntax (#66443)
**Dense**
`lvlTypes = [ "dense", "dense" ]` to `map = (d0, d1) -> (d0 : dense, d1
: dense)`
`lvlTypes = [ "dense", "dense" ], dimToLvl = affine_map<(i,j) -> (j,i)>`
to `map = (d0, d1) -> (d1 : dense, d0 : dense)`

**DCSR**
`lvlTypes = [ "compressed", "compressed" ]` to `map = (d0, d1) -> (d0 :
compressed, d1 : compressed)`

**DCSC**
`lvlTypes = [ "compressed", "compressed" ], dimToLvl = affine_map<(i,j)
-> (j,i)>` to `map = (d0, d1) -> (d1 : compressed, d0 : compressed)`

**Block Row**
`lvlTypes = [ "compressed", "dense" ]` to `map = (d0, d1) -> (d0 :
compressed, d1 : dense)`

**Block Column**
`lvlTypes = [ "compressed", "dense" ], dimToLvl = affine_map<(i,j) ->
(j,i)>` to `map = (d0, d1) -> (d1 : compressed, d0 : dense)`

This is an ongoing effort: #66146, #66309
2023-09-14 23:19:57 +00:00
Yinying Li
e2e429d994 [mlir][sparse] Migrate more tests to new syntax (#66309)
CSR:
`lvlTypes = [ "dense", "compressed" ]` to `map = (d0, d1) -> (d0 :
dense, d1 : compressed)`

CSC:
`lvlTypes = [ "dense", "compressed" ], dimToLvl = affine_map<(d0, d1) ->
(d1, d0)>` to `map = (d0, d1) -> (d1 : dense, d0 : compressed)`

This is an ongoing effort: #66146
2023-09-14 12:21:13 -04:00
Yinying Li
dbe1be9aa4 [mlir][sparse] Migrate tests to use new syntax (#66146)
lvlTypes = [ "compressed" ] to map = (d0) -> (d0 : compressed)
lvlTypes = [ "dense" ] to map = (d0) -> (d0 : dense)
2023-09-13 11:41:25 -04:00
Yinying Li
51ebecf309 [mlir][sparse] Changed sparsity properties to use _ instead of -
Example: compressed-no -> compressed_no

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D158567
2023-08-23 17:00:27 +00:00
Peiming Liu
a63d6a0014 [mlir][sparse] make UnpackOp return the actual filled length of unpacked memory
This might simplify frontend implementation by avoiding recomputation for the same value.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D154244
2023-06-30 21:35:15 +00:00
wren romano
540d5e0ce6 [mlir][sparse] Updating STEA parser/printer to use the name "dimSlices"
Depends On D151505

Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D151513
2023-05-30 15:50:07 -07:00
Peiming Liu
b2e6b73544 [mlir][sparse] extend unpack operation to unpack arbitrary encodings.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D151174
2023-05-23 22:34:01 +00:00
Peiming Liu
de56088866 [mlir][sparse] Support packing external data into arbitrary sparse tensor encoding.
We previously only support packing two array (values and coordinates) into COO tensors.
This patch allows packing inputs into arbitrary sparse tensor format.

It also deletes the "implicit" data canonicalization performed inside sparse compiler,
but instead requires users to canonicalize the data before passing it to the sparse compiler.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D150916
2023-05-19 17:41:49 +00:00
wren romano
a0615d020a [mlir][sparse] Renaming the STEA field dimLevelType to lvlTypes
This commit is part of the migration of towards the new STEA syntax/design.  In particular, this commit includes the following changes:
* Renaming compiler-internal functions/methods:
  * `SparseTensorEncodingAttr::{getDimLevelType => getLvlTypes}`
  * `Merger::{getDimLevelType => getLvlType}` (for consistency)
  * `sparse_tensor::{getDimLevelType => buildLevelType}` (to help reduce confusion vs actual getter methods)
* Renaming external facets to match:
  * the STEA parser and printer
  * the C and Python bindings
  * PyTACO

However, the actual renaming of the `DimLevelType` itself (along with all the "dlt" names) will be handled in a separate commit.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D150330
2023-05-17 14:24:09 -07:00
Peiming Liu
d4db528938 [mlir][sparse] extend unpack operation to support unpacking a batched COO type
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D149103
2023-05-01 18:17:29 +00:00
Peiming Liu
7864d736cf [mlir][sparse] extend pack operation to support packing a batched COO type
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D148670
2023-04-20 01:35:30 +00:00
Peiming Liu
33267f4007 [mlir][sparse] convert a sparse tensor slice to sparse tensor correctly.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D147074
2023-03-28 21:39:31 +00:00
Peiming Liu
6db397a8d4 [mlir][sparse] support dynamic sparse tensor slices.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D141532
2023-03-10 23:12:41 +00:00
Peiming Liu
8237cac612 [mlir][sparse] extend storage specifier operations for slices.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D141641
2023-03-10 18:58:47 +00:00
wren romano
84cd51bb97 [mlir][sparse] Renaming "pointer/index" to "position/coordinate"
The old "pointer/index" names often cause confusion since these names clash with names of unrelated things in MLIR; so this change rectifies this by changing everything to use "position/coordinate" terminology instead.

In addition to the basic terminology, there have also been various conventions for making certain distinctions like: (1) the overall storage for coordinates in the sparse-tensor, vs the particular collection of coordinates of a given element; and (2) particular coordinates given as a `Value` or `TypedValue<MemRefType>`, vs particular coordinates given as `ValueRange` or similar.  I have striven to maintain these distinctions
as follows:

  * "p/c" are used for individual position/coordinate values, when there is no risk of confusion.  (Just like we use "d/l" to abbreviate "dim/lvl".)

  * "pos/crd" are used for individual position/coordinate values, when a longer name is helpful to avoid ambiguity or to form compound names (e.g., "parentPos").  (Just like we use "dim/lvl" when we need a longer form of "d/l".)

    I have also used these forms for a handful of compound names where the old name had been using a three-letter form previously, even though a longer form would be more appropriate.  I've avoided renaming these to use a longer form purely for expediency sake, since changing them would require a cascade of other renamings.  They should be updated to follow the new naming scheme, but that can be done in future patches.

  * "coords" is used for the complete collection of crd values associated with a single element.  In the runtime library this includes both `std::vector` and raw pointer representations.  In the compiler, this is used specifically for buffer variables with C++ type `Value`, `TypedValue<MemRefType>`, etc.

    The bare form "coords" is discouraged, since it fails to make the dim/lvl distinction; so the compound names "dimCoords/lvlCoords" should be used instead.  (Though there may exist a rare few cases where is is appropriate to be intentionally ambiguous about what coordinate-space the coords live in; in which case the bare "coords" is appropriate.)

    There is seldom the need for the pos variant of this notion.  In most circumstances we use the term "cursor", since the same buffer is reused for a 'moving' pos-collection.

  * "dcvs/lcvs" is used in the compiler as the `ValueRange` analogue of "dimCoords/lvlCoords".  (The "vs" stands for "`Value`s".)  I haven't found the need for it, but "pvs" would be the obvious name for a pos-`ValueRange`.

    The old "ind"-vs-"ivs" naming scheme does not seem to have been sustained in more recent code, which instead prefers other mnemonics (e.g., adding "Buf" to the end of the names for `TypeValue<MemRefType>`).  I have cleaned up a lot of these to follow the "coords"-vs-"cvs" naming scheme, though haven't done an exhaustive cleanup.

  * "positions/coordinates" are used for larger collections of pos/crd values; in particular, these are used when referring to the complete sparse-tensor storage components.

    I also prefer to use these unabbreviated names in the documentation, unless there is some specific reason why using the abbreviated forms helps resolve ambiguity.

In addition to making this terminology change, this change also does some cleanup along the way:
  * correcting the dim/lvl terminology in certain places.
  * adding `const` when it requires no other code changes.
  * miscellaneous cleanup that was entailed in order to make the proper distinctions.  Most of these are in CodegenUtils.{h,cpp}

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D144773
2023-03-06 12:23:33 -08:00
Peiming Liu
44ff23d5e4 [mlir][sparse] unconditionally use IndexType for sparse_tensor.specifier
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D144574
2023-02-22 20:21:34 +00:00
bixia1
c2e248c6ae [mlir][sparse] Remove the expansion of symmetric MTX in the sparse tensor storage.
We will support symmetric MTX without expanding the data in the sparse tensor
storage.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D144059
2023-02-16 13:02:17 -08:00
Peiming Liu
c738b430c4 [mlir][sparse] introduce operations to query sparse tensor slice offset/strides at the given dimenion
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D141442
2023-02-16 00:31:11 +00:00
wren romano
f708a549b8 [mlir][sparse] Factoring out SparseTensorType class
This change adds a new `SparseTensorType` class for making the "dim" vs "lvl" distinction more overt, and for abstracting over the differences between sparse-tensors and dense-tensors.  In addition, this change also adds new type aliases `Dimension`, `Level`, and `FieldIndex` to make code more self-documenting.

Although the diff is very large, the majority of the changes are mechanical in nature (e.g., changing types to use the new aliases, updating variable names to match, etc).  Along the way I also made many variables `const` when they could be; the majority of which required only adding the keyword.  A few places had conditional definitions of these variables, requiring actual code changes; however, that was only done when the overall change was extremely local and easy to extract.  All these changes are included in the current patch only because it would be too onerous to split them off into a separate patch.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D143800
2023-02-14 19:17:19 -08:00
Peiming Liu
6dbca86d83 [mlir][sparse] introduce sparse_tensor::unpack operation
An inverse operation of sparse_tenosr::pack

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D143669
2023-02-11 01:01:35 +00:00
Peiming Liu
1f07853f2b [mlir][sparse] introduce sparse_tensor.pack operation
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D143224
2023-02-03 22:30:52 +00:00
bixia1
0c7f1c1520 [mlir][sparse] Extend sparse_tensor.sort with a enum attribute to specify a sorting implementation.
Currently, all the non-stable sorting algorithms are implemented via the
straightforward quick sort. This will be fixed in the following PR.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D142678
2023-01-29 18:34:08 -08:00
bixia1
9bde3d0cc5 [mlir][sparse] Add operator sparse_tensor.indices_buffer.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D140762
2023-01-05 09:35:55 -08:00
Peiming Liu
988733c600 [mlir][sparse] use sparse_tensor::StorageSpecifier to store dim/memSizes
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D140130
2022-12-23 00:47:36 +00:00
Peiming Liu
71cc0f1c04 [mlir][sparse] introduce sparse_tensor::StorageSpecifierType and related operations on it
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D139961
2022-12-14 22:27:22 +00:00
bixia1
974b4bf9fd [mlir][sparse] Add expand_symmetry attribute to the new operator.
The attribute tells the operator to handle symmetric structures for 2D tensors.
By default, the operator assumes the input tensor is not symmetric.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D138230
2022-11-23 16:32:15 -08:00
Peiming Liu
4fa00ce15c [mlir][sparse] extend foreach operation to accept reduction arguments; fix sparse tensor rewriting patterns that do not propagate sparse tensor SSA properly.
This patch re-commit D137468 and D137463, which were reverted by mistakes.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D137579
2022-11-07 21:40:30 +00:00