Commit Graph

378 Commits

Author SHA1 Message Date
Aart Bik
d34661309b [mlir][sparse] fix merge conflict in sddmm test (#66670)
This fixes a merge conflict in
https://github.com/llvm/llvm-project/pull/66563
2023-09-18 09:51:21 -07:00
Aart Bik
6a45339bac [mlir][sparse] refine sparse fusion with empty tensors materialization (#66563)
This is a minor step towards deprecating bufferization.alloc_tensor().
It replaces the examples with tensor.empty() and adjusts the underlying
rewriting logic to prepare for this upcoming change.
2023-09-18 09:01:11 -07:00
Martin Erhart
6bf043e743 [mlir][bufferization] Remove allow-return-allocs and create-deallocs pass options, remove bufferization.escape attribute (#66619)
This commit removes the deallocation capabilities of
one-shot-bufferization. One-shot-bufferization should never deallocate
any memrefs as this should be entirely handled by the
ownership-based-buffer-deallocation pass going forward. This means the
`allow-return-allocs` pass option will default to true now,
`create-deallocs` defaults to false and they, as well as the escape
attribute indicating whether a memref escapes the current region, will
be removed. A new `allow-return-allocs-from-loops` option is added as a
temporary workaround for some bufferization limitations.
2023-09-18 16:44:48 +02: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
Martin Erhart
c199f7dc62 Revert "[mlir][bufferization] Remove allow-return-allocs and create-deallocs pass options, remove bufferization.escape attribute"
This reverts commit 6a91dfedeb.

This caused problems in downstream projects. We are reverting to give
them more time for integration.
2023-09-13 13:53:48 +00:00
Martin Erhart
6a91dfedeb [mlir][bufferization] Remove allow-return-allocs and create-deallocs pass options, remove bufferization.escape attribute
This is the first commit in a series with the goal to rework the
BufferDeallocation pass. Currently, this pass heavily relies on copies
to perform correct deallocations, which leads to very slow code and
potentially high memory usage. Additionally, there are unsupported cases
such as returning memrefs which this series of commits aims to add
support for as well.

This first commit removes the deallocation capabilities of
one-shot-bufferization.One-shot-bufferization should never deallocate any
memrefs as this should be entirely handled by the buffer-deallocation pass
going forward. This means the allow-return-allocs pass option will
default to true now, create-deallocs defaults to false and they, as well
as the escape attribute indicating whether a memref escapes the current region,
will be removed.

The documentation should w.r.t. these pass option changes should also be
updated in this commit.

Reviewed By: springerm

Differential Revision: https://reviews.llvm.org/D156662
2023-09-13 09:30:22 +00:00
Daniil Dudkin
8a6e54c9b3 [mlir][arith] Rename operations: maxfmaximumf, minfminimumf (#65800)
This patch is part of a larger initiative aimed at fixing floating-point `max` and `min` operations in MLIR: https://discourse.llvm.org/t/rfc-fix-floating-point-max-and-min-operations-in-mlir/72671.

This commit addresses Task 1.2 of the mentioned RFC. By renaming these operations, we align their names with LLVM intrinsics that have corresponding semantics.
2023-09-11 22:02:19 -07:00
Yinying Li
898bf539a7 [mlir][sparse] Surface syntax change in parsing
Example: compressed(nonunique, nonordered) or compressed(nonordered, nonunique) instead of compressed_nu_no.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D159366
2023-09-01 19:25:00 +00:00
Aart Bik
b86d3cbc12 [mlir][sparse] complete various FIXMEs in sparse support lib
Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D159245
2023-08-30 21:30:25 -07:00
Matthias Springer
79ff70fda2 [mlir][sparse] Better error handling when bufferizing sparse_tensor ops
sparse_tensor ops cannot be bufferized with One-Shot Bufferize. (They can only be analyzed.) The sparse compiler does the actual lowering to memref. Produce a proper error message instead of crashing.

This fixes #61311.

Differential Revision: https://reviews.llvm.org/D158728
2023-08-25 08:34:05 +02:00
Christian Sigg
00d27fb0f8 [mlir][sparse] Disable sparse_reshape_dot.mlir test with XFAIL
... instead of UNSUPPORTED.
2023-08-24 08:28:42 +02:00
Peiming Liu
fa9c93e30e [mlir][sparse] Remove view-based sparse tensor collapse_shape implementation.
We will migrate to a cleaner and more complete implementation.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D158658
2023-08-23 20:29:28 +00: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
Aart Bik
6eb0332e72 [mlir][sparse] test pack copy behavior
Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D158229
2023-08-17 17:00:30 -07:00
Aart Bik
289f7231f9 [mlir][sparse][gpu] minor code cleanup for sparse gpu ops
Consistent order of ops and related methods.
Also, renamed SpGEMMGetSizeOp to SpMatGetSizeOp
since this is a general utility for sparse matrices,
not specific to GEMM ops only.

Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D157922
2023-08-14 15:08:57 -07:00
Aart Bik
7a8dcf2ec7 [mlir][sparse][gpu] add CHECK test to spGEMM libgen
Reviewed By: K-Wu

Differential Revision: https://reviews.llvm.org/D157675
2023-08-11 09:56:13 -07:00
Tom Eccles
dea33c80d3 [mlir][Transforms] teach CSE about recursive memory effects
Add support for reasoning about operations with recursive memory effects
to CSE. The recursive effects are gathered by a helper function. I
decided to allow returning duplicates from the helper function because
there's no benefit to spending the computation time to remove them in
the existing use case.

Differential Revision: https://reviews.llvm.org/D156805
2023-08-10 09:40:01 +00: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
Peiming Liu
849187c95d [mlir][transform][sparse] introduce sparse tensor transform extension.
Reviewed By: aartbik, ftynse

Differential Revision: https://reviews.llvm.org/D157025
2023-08-04 18:40:07 +00:00
Kun Wu
1e491c425b [mlir][sparse][gpu] add 2:4 spmm prune_and_check flag
Differential Revision: https://reviews.llvm.org/D155909
2023-08-01 18:24:18 +00:00
wren romano
889f4bf264 [mlir][sparse] Improve DimLvlMapParser's handling of variable bindings
This commit comprises a number of related changes:

(1) Reintroduces the semantic distinction between `parseVarUsage` vs `parseVarBinding`, adds documentation explaining the distinction, and adds commentary to the one place that violates the desired/intended semantics.

(2) Improves documentation/commentary about the forward-declaration of level-vars, and about the meaning of the `bool` parameter to `parseLvlSpec`.

(2) Removes the `VarEnv::addVars` method, and instead has `DimLvlMapParser` handle the conversion issues directly.  In particular, the parser now stores and maintains the `{dims,lvls}AndSymbols` arrays, thereby avoiding the O(n^2) behavior of scanning through the entire `VarEnv` for each `parse{Dim,Lvl}Spec` call.  Unfortunately there still remains another source of O(n^2) behavior, namely: the `AsmParser::parseAffineExpr` method will copy the `DimLvlMapParser::{dims,lvls}AndSymbols` arrays into `AffineParser::dimsAndSymbols` on each `parse{Dim,Lvl}Spec` call; but fixing that would require extensive changes to `AffineParser` itself.

Depends On D155532

Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D155533
2023-07-20 15:56:03 -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
K-Wu
e37fc3cc39 [mlir][sparse][gpu] Impl 2:4 SpMM rewrite for linalg op w/ DENSE24 attr
Differential Revision: https://reviews.llvm.org/D154772
2023-07-10 22:36:57 +00:00
Peiming Liu
fc5d8fce7d [mlir][sparse] support dual sparse convolution.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D152601
2023-07-10 16:49:32 +00:00
Aart Bik
03125e6894 [mlir][sparse][gpu] fix missing dealloc
This dealloc was incorrectly removed in
https://reviews.llvm.org/D153173

Reviewed By: K-Wu

Differential Revision: https://reviews.llvm.org/D154564
2023-07-06 09:48:19 -07:00
Kun Wu
be2dd22b8f [mlir][sparse][gpu] reuse CUDA environment handle throughout instance lifetime
Differential Revision: https://reviews.llvm.org/D153173
2023-06-30 21:52:34 +00:00
Aart Bik
b939c015a4 [mlir][sparse] add affine parsing to new surface syntax for STEA
(1) uses the previously introduce API to reuse AffineExpr parser without codedup
(2) solves the look-ahead problem when parsing level spec

Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D154254
2023-06-30 14:48:23 -07: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
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
Peiming Liu
df11a2b41a [mlir][sparse] admit un-sparsifiable operations if all its operands are loaded from dense input
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D153998
2023-06-28 21:27:50 +00:00
Aart Bik
f14c8eb595 [mlir][sparse][gpu] refine SDDMM pattern for cuSPARSE
Old pattern was missing some cases (e.g. swapping the arguments)
but it also allowed too many cases (e.g. non-empty "absent" or
different arguments for add/mul). This fixes the issues.

Reviewed By: K-Wu

Differential Revision: https://reviews.llvm.org/D153487
2023-06-21 18:31:55 -07:00
Peiming Liu
e7df82816b [mlir][sparse] rewrite arith::SelectOp to semiring operations to sparsify it.
Reviewed By: aartbik, K-Wu

Differential Revision: https://reviews.llvm.org/D153397
2023-06-21 21:22:18 +00:00
Aart Bik
5c03c056e0 [mlir][sparse] enhance element-wise fusion heuristics
We prevent merging a sparse-in/dense-out with dense-in
kernels because the result is usuall not sparsifiable.
Dense kernels and sparse kernels are still fused, obviously.

Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D153077
2023-06-15 16:48:40 -07:00
Kun Wu
9167dd46ba [mlir][sparse][gpu] recognizing sddmm pattern in GPU libgen path
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D151582
2023-06-15 23:48:11 +00:00
Aart Bik
65bfd5cb25 [mlir][sparse] proper in-place SDDMM with spy function
This specific operation matches the cuSPARSE SDDMM semantics exactly.

Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D152969
2023-06-15 13:59:38 -07:00
Peiming Liu
faf7cd97d0 [mlir][sparse] merger extension to support sparsifying arith::CmpI/CmpF operation
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D152761
2023-06-15 17:26:50 +00:00
Peiming Liu
83b7f018fd [mlir][sparse] fix crashes when the tensor that defines the loop bound can not be found
Reviewed By: aartbik, K-Wu

Differential Revision: https://reviews.llvm.org/D152877
2023-06-14 20:27:50 +00:00
Peiming Liu
fd68d36109 [mlir][sparse] unifying enterLoopOverTensorAtLvl and enterCoIterationOverTensorsAtLvls
The tensor levels are now explicitly categorized into different `LoopCondKind` to instruct LoopEmitter generate different code for different kinds of condition (e.g., `SparseCond`, `SparseSliceCond`, `SparseAffineIdxCond`, etc)

The process of generating a while loop is now dissembled into three steps and they are dispatched to different LoopCondKind handler.
1. Generate LoopCondition (e.g., `pos <= posHi` for `SparseCond`, `slice.isNonEmpty` for `SparseAffineIdxCond`)
2. Generate LoopBody (e.g., compute the coordinates)
3. Generate ExtraChecks (e.g., `if (onSlice(crd))` for `SparseSliceCond`)

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D152464
2023-06-14 20:03:10 +00:00
Aart Bik
debdf7e0ff [mlir][sparse] refine single condition set up for semi-ring ops
Reviewed By: Peiming, K-Wu

Differential Revision: https://reviews.llvm.org/D152874
2023-06-14 09:23:09 -07:00
Peiming Liu
0258a53521 Brings back "[mlir][sparse] moving inbound check for slice driven loop into before block of the WhileOp"
This reverts commit 07b927902d.

Reviewed By: K-Wu

Differential Revision: https://reviews.llvm.org/D152566
2023-06-09 17:45:46 +00:00
Peiming Liu
07b927902d Revert "[mlir][sparse] moving inbound check for slice driven loop into before block of the WhileOp"
This reverts commit 853d704fd0.

Differential Revision: https://reviews.llvm.org/D152562
2023-06-09 17:21:40 +00:00
Kun Wu
97f4c22b3a [mlir][sparse][gpu] unify dnmat and dnvec handle and ops
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D152465
2023-06-09 17:16:48 +00:00
Peiming Liu
853d704fd0 [mlir][sparse] moving inbound check for slice driven loop into before block of the WhileOp
This patch changes the while loop generated for iterating over a fully reduced sparse level with affine index expression.
Before:
```
cont = true
while (cont) {
  if (inBound()) {
    ....
    cont = true;
  } else {
    cont = false;
  }
}
```
After:
```
while(inBound()) {
  ....
}
```

Reviewed By: K-Wu

Differential Revision: https://reviews.llvm.org/D152463
2023-06-09 17:03:15 +00:00
Kun Wu
8ed59c53de [mlir][sparse][gpu] add sm8.0+ tensor core 2:4 sparsity support
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D151775
2023-06-06 23:13:21 +00:00
Aart Bik
378f1885e3 [mlir][sparse] enhance sparse reduction support
Formerly, we accepted and/prod reductions as a standard
reduction but these change the semantics after sparsification
by not looking at implicit zeros. Therefore, we only accept
standard reductions that are insensitive to implicit vs.
explicit zeros, and leave the more complex reductions to
the sparse_tensor.reduce custom reduction implementation.

Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D151929
2023-06-01 16:30:21 -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