168 Commits

Author SHA1 Message Date
Aart Bik
22212ca745
[mlir][sparse] simplify some header code (#70989)
This is a first revision in a small series of changes that removes
duplications between direct encoding methods and sparse tensor type
wrapper methods (in favor of the latter abstraction, since it provides
more safety). The goal is to simply end up with "just" SparseTensorType
2023-11-02 09:31:11 -07:00
Peiming Liu
ef222988b4
[mlir][sparse] implements sparse_tensor.reinterpret_map (#70388) 2023-10-26 16:00:32 -07:00
Yinying Li
b165650aee
[mlir][sparse] Return actual identity map instead of null map (#70365)
Changes:

1. For both dimToLvl and lvlToDim, always returns the actual map instead
of AffineMap() for identity map.
2. Updated custom builder for encoding to have default values.
3. Non-inferable lvlToDim will still return AffineMap() during
inference, so it will be caught by verifier.
2023-10-26 18:30:34 -04:00
Peiming Liu
d808d922b4
[mlir][sparse] introduce sparse_tensor.reinterpret_map operation. (#70378) 2023-10-26 15:04:09 -07:00
Aart Bik
2e2011da38
[mlir][sparse] avoid excessive macro magic (#70276)
The shorthands are not even always shorter and the code is less clear
than when simply written out.
2023-10-25 20:48:37 -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
c780352de9
[mlir][sparse] implement sparse_tensor.lvl operation. (#69993) 2023-10-24 13:23:28 -07:00
Peiming Liu
f0f5fdf73d
[mlir][sparse] introduce sparse_tensor.lvl operation. (#69978) 2023-10-23 15:49:39 -07:00
Peiming Liu
e9fa1fdec9
[mlir][sparse] support CSR/BSR conversion (#69800) 2023-10-20 17:24:34 -07:00
Peiming Liu
6456e0bbbb
[mlir][sparse] implement sparse_tensor.crd_translate operation (#69653) 2023-10-20 16:18:02 -07:00
Peiming Liu
ff21a90e51
[mlir][sparse] introduce sparse_tensor.crd_translate operation (#69630) 2023-10-19 15:42:09 -07:00
Yinying Li
fb5047f524
[mlir][sparse] Remove old syntax (#69624) 2023-10-19 17:33:28 -04:00
Yinying Li
7b9fb1c228
[mlir][sparse] Update verifier for block sparsity and singleton (#69389)
Updates:
1. Verification of block sparsity.
2. Verification of singleton level type can only follow compressed or
loose_compressed levels. And all level types after singleton should be
singleton.
3. Added getBlockSize function.
4. Added an invalid encoding test for an incorrect lvlToDim map that
user provides.
2023-10-19 12:34:18 -04:00
Yinying Li
d4088e7d5f
[mlir][sparse] Populate lvlToDim (#68937)
Updates:
1. Infer lvlToDim from dimToLvl
2. Add more tests for block sparsity
3. Finish TODOs related to lvlToDim, including adding lvlToDim to python
binding

Verification of lvlToDim that user provides will be implemented in the
next PR.
2023-10-17 16:09:39 -04:00
Peiming Liu
761c9dd927
[mlir][sparse] implementating stageSparseOpPass as an interface (#69022) 2023-10-17 10:54:44 -07:00
Peiming Liu
f248d0b28d
[mlir][sparse] implement sparse_tensor.reorder_coo (#68916)
As a side effect of the change, it also unifies the convertOp
implementation between lib/codegen path.
2023-10-12 13:22:45 -07:00
Peiming Liu
0aacc2137a
[mlir][sparse] introduce sparse_tensor.reorder_coo operation (#68827) 2023-10-12 09:42:12 -07:00
Peiming Liu
dda3dc5e38
[mlir][sparse] simplify ConvertOp rewriting rules (#68350)
Canonicalize complex convertOp into multiple stages, such that it can
either be done by a direct conversion or by sorting.
2023-10-11 09:34:11 -07:00
Yinying Li
6280e23124
[mlir][sparse] Print new syntax (#68130)
Printing changes from `#sparse_tensor.encoding<{ lvlTypes = [
"compressed" ] }>` to `map = (d0) -> (d0 : compressed)`. Level
properties, ELL and slice are also supported.
2023-10-04 16:36:05 -04: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
Yinying Li
d2e8517912
[mlir][sparse] Update Enum name for CompressedWithHigh (#67845)
Change CompressedWithHigh to LooseCompressed.
2023-10-02 11:06:40 -04: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
Aart Bik
836411b99f
[mlir][sparse] add lvlToDim field to sparse tensor encoding (#67194)
Note the new surface syntax allows for defining a dimToLvl and lvlToDim
map at once (where usually the latter can be inferred from the former,
but not always). This revision adds storage for the latter, together
with some intial boilerplate. The actual support (inference, validation,
printing, etc.) is still TBD of course.
2023-09-22 15:51:25 -07: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
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
Aart Bik
bb44a6b7bb [mlir][sparse] migrate more to new surface syntax
Replaced the "NEW_SYNTAX" with the more readable "map"
(which we may, or may not keep). Minor improvement in
keyword parsing, migrated a few more examples over.

Reviewed By: Peiming, yinying-lisa-li

Differential Revision: https://reviews.llvm.org/D158325
2023-08-21 12:49:21 -07:00
wren romano
cad4646733 [mlir][sparse] Improve handling of NEW_SYNTAX
Improves the conversion from `DimLvlMap` to STEA, in order to correct rank-mismatch issues in the roundtrip tests.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D157162
2023-08-04 17:53:34 -07:00
wren romano
fdbe9312b1 [mlir][sparse] Adding getters/setters to DimLvlMap
Depends On D156768

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D156770
2023-08-01 12:55:45 -07:00
Peiming Liu
269c82d389 [mlir][sparse] introduce new 2:4 block sparsity level type.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D155128
2023-07-12 23:33:53 +00:00
Matthias Springer
cb7bda2ace [mlir][NFC] Use getConstantIntValue instead of casting to ConstantIndexOp
`getConstantIntValue` extracts constant values from all constant-like ops, not just `arith::ConstantIndexOp`.

Differential Revision: https://reviews.llvm.org/D154356
2023-07-04 14:08:37 +02:00
Aart Bik
6b88c852b6 [mlir][sparse] Start migration to new surface syntax for STEA
We are in the progress of migrating to a much improved surface syntax for the Sparse Tensor Encoding Attribute (STEA).

You can see a preview of this in the StableHLO RFC at

 https://github.com/openxla/stablehlo/blob/main/rfcs/20230210-sparsity.md

//**This design is courtesy Wren Romano.**//

This initial revision
(1) Introduces the first version of a new parser written by Wren Romano
(2) Introduces a simple "migration plan" using NEW_SYNTAX on the STEA, which will allow us to test the new parser with new examples, as well as migrate existing examples over without the need to rewrite them all

This first "drop" merely provides the entry points to parse the new syntax. The parser is still under active development. For example, we need to address the "lookahead" issue when parsing the lvl spec (viz. do we see l0 = d0 or a direct d0). Another larger task is to actually implement "affine" parsing (since the MLIR affine parser is not accessible in other parts of the tree).

EXAMPLE:

Currently, CSR looks like

  #CSR = #sparse_tensor.encoding<{
    lvlTypes = ["dense","compressed"],
    dimToLvl = affine_map<(i,j) -> (i,j)>
  }>

but you can "force" the new parser with

  #CSR = #sparse_tensor.encoding<{
    NEW_SYNTAX =
    (d0, d1) -> (l0 = d0 : dense, l1 = d1 : compressed)
  }>

Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D153997
2023-06-29 11:32:07 -07:00
Aart Bik
11a4f5bdfb [mlir][sparse] minor code changes
Submitting for Wren

Reviewed By: K-Wu

Differential Revision: https://reviews.llvm.org/D153804
2023-06-26 12:57:52 -07:00
wren romano
7a1077baa0 [mlir][sparse] Improving SparseTensorDimSliceAttr methods
This patch makes the following changes to `SparseTensorDimSliceAttr` methods:
* Mark `isDynamic` constexpr.
* Add new helpers `getStatic` and `getStaticString` to avoid repetition.
* Moved the definitions for `getStatic{Offset,Stride,Size}` and `isCompletelyDynamic` out of the class declaration; because there's no benefit to inlining them.
* Changed `parse` to use `kDynamic` rather than literals.
* Changed `verify` to use the `isDynamic` helper.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D150919
2023-05-30 17:30:55 -07:00
wren romano
f58e67dee9 [mlir][sparse] Removing unused helper function
Depends On D151505

Reviewed By: aartbik, Peiming

Differential Revision: https://reviews.llvm.org/D151522
2023-05-30 15:59:26 -07:00
wren romano
af2bec7c4a [mlir][sparse] Adding new STEA::{with,without}DimSlices factories
(These factories are used in downstream code, despite not being used within the MLIR codebase.)

Depends On D151513

Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D151518
2023-05-30 15:53:30 -07: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
wren romano
76647fce13 [mlir][sparse] Combining dimOrdering+higherOrdering fields into dimToLvl
This is a major step along the way towards the new STEA design.  While a great deal of this patch is simple renaming, there are several significant changes as well.  I've done my best to ensure that this patch retains the previous behavior and error-conditions, even though those are at odds with the eventual intended semantics of the `dimToLvl` mapping.  Since the majority of the compiler does not yet support non-permutations, I've also added explicit assertions in places that previously had implicitly assumed it was dealing with permutations.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D151505
2023-05-30 15:19:50 -07:00
Tres Popp
68f58812e3 [mlir] Move casting calls from methods to function calls
The MLIR classes Type/Attribute/Operation/Op/Value support
cast/dyn_cast/isa/dyn_cast_or_null functionality through llvm's doCast
functionality in addition to defining methods with the same name.
This change begins the migration of uses of the method to the
corresponding function call as has been decided as more consistent.

Note that there still exist classes that only define methods directly,
such as AffineExpr, and this does not include work currently to support
a functional cast/isa call.

Context:
- https://mlir.llvm.org/deprecation/ at "Use the free function variants
  for dyn_cast/cast/isa/…"
- Original discussion at https://discourse.llvm.org/t/preferred-casting-style-going-forward/68443

Implementation:
This patch updates all remaining uses of the deprecated functionality in
mlir/. This was done with clang-tidy as described below and further
modifications to GPUBase.td and OpenMPOpsInterfaces.td.

Steps are described per line, as comments are removed by git:
0. Retrieve the change from the following to build clang-tidy with an
   additional check:
   main...tpopp:llvm-project:tidy-cast-check
1. Build clang-tidy
2. Run clang-tidy over your entire codebase while disabling all checks
   and enabling the one relevant one. Run on all header files also.
3. Delete .inc files that were also modified, so the next build rebuilds
   them to a pure state.

```
ninja -C $BUILD_DIR clang-tidy

run-clang-tidy -clang-tidy-binary=$BUILD_DIR/bin/clang-tidy -checks='-*,misc-cast-functions'\
               -header-filter=mlir/ mlir/* -fix

rm -rf $BUILD_DIR/tools/mlir/**/*.inc
```

Differential Revision: https://reviews.llvm.org/D151542
2023-05-26 10:29:55 +02: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
wren romano
d824af535c [mlir][sparse] Improving error-detection for STEA::get{Pos,Crd}Type
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D150995
2023-05-19 16:24:38 -07: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
867e19648d [mlir][sparse] Adding SparseTensorEncodingAttr::getDimSlice
This helps catch segfaults and OOB.

Reviewed By: aartbik, Peiming

Differential Revision: https://reviews.llvm.org/D150917
2023-05-18 17:06:01 -07:00
wren romano
d9a2f89bee [mlir][sparse] Adjusting error message wording, to better match new field names
This is a followup to D150330, split out because it's not purely mechanical.

Depends On D150330

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D150409
2023-05-18 15:02:08 -07:00
Peiming Liu
afe78db770 [mlir][sparse] Make sparse_tensor::StorageLayout publicly available.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D150739
2023-05-18 20:29:46 +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
Tres Popp
c1fa60b4cd [mlir] Update method cast calls to function calls
The MLIR classes Type/Attribute/Operation/Op/Value support
cast/dyn_cast/isa/dyn_cast_or_null functionality through llvm's doCast
functionality in addition to defining methods with the same name.
This change begins the migration of uses of the method to the
corresponding function call as has been decided as more consistent.

Note that there still exist classes that only define methods directly,
such as AffineExpr, and this does not include work currently to support
a functional cast/isa call.

Context:

* https://mlir.llvm.org/deprecation/ at "Use the free function variants for dyn_cast/cast/isa/…"
* Original discussion at https://discourse.llvm.org/t/preferred-casting-style-going-forward/68443

Implementation:
This follows a previous patch that updated calls
`op.cast<T>()-> cast<T>(op)`. However some cases could not handle an
unprefixed `cast` call due to occurrences of variables named cast, or
occurring inside of class definitions which would resolve to the method.
All C++ files that did not work automatically with `cast<T>()` are
updated here to `llvm::cast` and similar with the intention that they
can be easily updated after the methods are removed through a
find-replace.

See https://github.com/llvm/llvm-project/compare/main...tpopp:llvm-project:tidy-cast-check
for the clang-tidy check that is used and then update printed
occurrences of the function to include `llvm::` before.

One can then run the following:
```
ninja -C $BUILD_DIR clang-tidy

run-clang-tidy -clang-tidy-binary=$BUILD_DIR/bin/clang-tidy -checks='-*,misc-cast-functions'\
                 -export-fixes /tmp/cast/casts.yaml mlir/*\
                 -header-filter=mlir/ -fix

rm -rf $BUILD_DIR/tools/mlir/**/*.inc
```

Differential Revision: https://reviews.llvm.org/D150348
2023-05-12 11:21:30 +02:00
Aart Bik
9a018a7b48 [mlir][sparse] relax constraints on tensor.cast with pre-rewriting
Reviewed By: wrengr

Differential Revision: https://reviews.llvm.org/D149489
2023-05-01 16:03:44 -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
b9589545c4 [mlir][sparse] introduce a new compressed(hi) dimension level type
`compressed(hi)` is similar to `compressed`, but instead of reusing the previous position high as the current position low, it uses a pair of positions for each sparse index.

The patch only introduces the definition (syntax) but does not provide codegen implementation.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D148664
2023-04-18 23:26:11 +00:00