Commit Graph

758 Commits

Author SHA1 Message Date
Matthias Springer
ae9e1d1df4 [mlir][SparseTensor] Fix incorrect API usage in RewritePatterns
Incorrect API usage was detected by D144552.

Differential Revision: https://reviews.llvm.org/D145166
2023-03-02 17:59:57 +01:00
Peiming Liu
b60cf8c972 [mlir][sparse] support coiteration with fused reshape tensor
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D145091
2023-03-01 20:55:46 +00:00
Peiming Liu
fc126022e8 [mlir][sparse] fuse collapse_shape on sparse tensor with GenericOp.
Instead of always materializing a new sparse tensor after reshape, this patch tries to fuses the reshape (currently only on COO) with GenericOp and coiterates with the reshaped tensors without allocating a new sparse tensor.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D145016
2023-03-01 19:05:48 +00:00
bixia1
2c81d43241 [mlir][sparse] Improve the implementation of sparse_tensor.new for the codegen path.
Rewrite a NewOp into a NewOp of a sorted COO tensor and a ConvertOp for
converting the sorted COO tensor to the destination tensor type.

Codegen a NewOp of a sorted COO tensor to use the new bulk reader API and sort
the elements only when the input is not sorted.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D144504
2023-03-01 07:29:49 -08:00
Peiming Liu
849529ba8a [mlir][sparse] fix performance bug in matmul with a sparse rhs due to suboptimal iteration graphs.
While dense tensors support random accesses, it is critical to visit them in a row-major order for better cache locality. However, we previously consider dense inputs and outputs together when computing constraints for building iteration graph, it could lead us to less efficient iteration graphs.

This patch adds a new `SortMask::kIncludeDenseInput` to treat dense inputs/outputs separately when building iteration graph, thus increasing the chance for use to construct a better iteration graph.

A more fine-grained approach is to treat each input separately.

Note, related to:
 https://github.com/llvm/llvm-project/issues/51651

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D144932
2023-02-28 21:02:17 +00:00
Kohei Yamaguchi
9a29d87538 [mlir][sparse] Add checking parent op of SortOp
Fix crash with segmentation fault caused by setting a parent operator
that is not func::FuncOp with sparse_tensor SortOp.

fixes https://github.com/llvm/llvm-project/issues/59988

Reviewed By: aartbik, wrengr

Differential Revision: https://reviews.llvm.org/D143874
2023-02-27 16:37:02 +01:00
Peiming Liu
85dbb3fc4b [mlir][sparse] support sparse tensor element type conversion in codegen path
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D144578
2023-02-23 17:49:50 +00: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
Peiming Liu
b7d86f3f1c [mlir][sparse] revert optimization for dense->csc conversion.
Eliminates the sort seems make the whole conversion slower (probably because loop rotation leads to bad locality).

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D144517
2023-02-21 21:34:01 +00:00
Peiming Liu
9e8d9316ce [mlir][sparse] allow foreach operation to generate out-of-order loop on non-annotated tensor.
No need for a temp COO and sort even when converting dense -> CSC, we can instead rotate the loop to yield a ordered coordinates at beginning.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D144213
2023-02-16 23:23:20 +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
e2e83f4c8f [mlir][sparse] support coiteration over sparse tensor slices
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D140736
2023-02-15 23:52:22 +00:00
wren romano
ae7942e296 [mlir][sparse] adding SparseTensorType::get{Pointer,Index}Type methods
Depends On D143800

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D143946
2023-02-15 14:37:55 -08:00
wren romano
d950bdc73e [mlir][sparse] misc code cleanup
* Flattening/simplifying some nested conditionals
* const-ifying some local variables

Depends On D143800

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D143949
2023-02-15 13:29:00 -08:00
wren romano
bb4fc6b6d6 [mlir][sparse] Adding SparseTensorType::{operator==, hasSameDimToLvlMap}
Depends On D143800

Reviewed By: aartbik, Peiming

Differential Revision: https://reviews.llvm.org/D144052
2023-02-15 12:05:29 -08:00
Jie Fu
71251e8d4f [mlir] Fix -Wsign-compare in SparseTensorRewriting.cpp and Sparsification.cpp (NFC)
/home/jiefu/llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp:279:33: error: comparison of integers of different signs: 'int64_t' (aka 'long') and 'const mlir::sparse_tensor::Level' (aka 'const unsigned long') [-Werror,-Wsign-compare]
    assert(env.op().getRank(&t) == lvlRank);
           ~~~~~~~~~~~~~~~~~~~~ ^  ~~~~~~~
/usr/include/assert.h:93:27: note: expanded from macro 'assert'
     (static_cast <bool> (expr)                                         \
                          ^~~~
1 error generated.

/home/jiefu/llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp:788:29: error: comparison of integers of different signs: 'int64_t' (aka 'long') and 'const mlir::sparse_tensor::Dimension' (aka 'const unsigned long') [-Werror,-Wsign-compare]
    assert(srcRTT.getRank() == dimRank);
           ~~~~~~~~~~~~~~~~ ^  ~~~~~~~
/usr/include/assert.h:93:27: note: expanded from macro 'assert'
     (static_cast <bool> (expr)                                         \
                          ^~~~
/home/jiefu/llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp:810:31: error: comparison of integers of different signs: 'int64_t' (aka 'long') and 'const mlir::sparse_tensor::Dimension' (aka 'const unsigned long') [-Werror,-Wsign-compare]
      assert(srcRTT.getRank() == dimRank);
             ~~~~~~~~~~~~~~~~ ^  ~~~~~~~
/usr/include/assert.h:93:27: note: expanded from macro 'assert'
     (static_cast <bool> (expr)                                         \
                          ^~~~
2 errors generated.
2023-02-15 11:58:08 +08: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
81cb70e46e [mlir][sparse] fix a bug in UnpackOp converter.
UnpackOp Converter used to create reallocOp unconditionally, but it might cause issue when the requested memory size is smaller than the actually storage.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D144065
2023-02-15 02:36:00 +00:00
Peiming Liu
ce9ce66b8d [mlir][sparse] fix a memory leakage when converting from a tensor slice
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D143929
2023-02-13 22:44:12 +00:00
Peiming Liu
dc6427d687 [mlir][sparse] implement lowering rules for sparse_tensor::unpack
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D143672
2023-02-11 01:05:46 +00:00
Jim Kitchen
81d0d2b2a0 [mlir][sparse] Sparse reduction in lex order no longer produces dense output
Previously, when performing a reduction on a sparse tensor, the result
would be different depending on iteration order. For expanded access pattern,
an empty row would contribute no entry in the output. For lex ordering, the
identity would end up in the output.

This code changes that behavior and keeps track of whether any entries were
actually reduced in lex ordering, making the output consistent between the
two iteration styles.

Differential Revision: https://reviews.llvm.org/D142050
2023-02-10 13:09:28 -06:00
Matthias Springer
9fa6b3504b [mlir][bufferization] Improve aliasing OpOperand/OpResult property
`getAliasingOpOperands`/`getAliasingOpResults` now encodes OpOperand/OpResult, buffer relation and a degree of certainty. E.g.:
```
// aliasingOpOperands(%r) = {(%t, EQUIV, DEFINITE)}
// aliasingOpResults(%t) = {(%r, EQUIV, DEFINITE)}
%r = tensor.insert %f into %t[%idx] : tensor<?xf32>

// aliasingOpOperands(%r) = {(%t0, EQUIV, MAYBE), (%t1, EQUIV, MAYBE)}
// aliasingOpResults(%t0) = {(%r, EQUIV, MAYBE)}
// aliasingOpResults(%t1) = {(%r, EQUIV, MAYBE)}
%r = arith.select %c, %t0, %t1 : tensor<?xf32>
```

`BufferizableOpInterface::bufferRelation` is removed, as it is now part of `getAliasingOpOperands`/`getAliasingOpResults`.

This change allows for better analysis, in particular wrt. equivalence. This allows additional optimizations and better error checking (which is sometimes overly conservative). Examples:

* EmptyTensorElimination can eliminate `tensor.empty` inside `scf.if` blocks. This requires a modeling of equivalence: It is not a per-OpResult property anymore. Instead, it can be specified for each OpOperand and OpResult. This is important because `tensor.empty` may be eliminated only if all values on the SSA use-def chain to the final consumer (`tensor.insert_slice`) are equivalent.
* The detection of "returning allocs from a block" can be improved. (Addresses a TODO in `assertNoAllocsReturned`.) This allows us to bufferize IR such as "yielding a `tensor.extract_slice` result from an `scf.if` branch", which currently fails to bufferize because the alloc detection is too conservative.
* Better bufferization of loops. Aliases of the iter_arg can be yielded (even if they are not equivalent) without having to realloc and copy the entire buffer on each iteration.

The above-mentioned examples are not yet implemented with this change. This change just improves the BufferizableOpInterface, its implementations and related helper functions, so that better aliasing information is available for each op.

Differential Revision: https://reviews.llvm.org/D142129
2023-02-09 11:35:03 +01:00
bixia1
a150766880 [mlir][sparse] Implement hybrid quick sort for sparse_tensor.sort.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D143227
2023-02-08 14:06:31 -08:00
Peiming Liu
0352690421 [mlir][sparse] make foreach operation support sparse tensor slices.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D140713
2023-02-08 18:58:35 +00:00
Peiming Liu
7a8edea69d [mlir][sparse] fix bug when packing tensor with 32 bit pointer width.
Reviewed By: wrengr

Differential Revision: https://reviews.llvm.org/D143450
2023-02-07 02:00:40 +00:00
Aart Bik
3bd82f30dc [mlir][sparse] compute allocation size_hint
This adds the hint to a number of tensor allocations in codegens,
shaving off quite some time from e.g. reading in sparse matrices
due to zero-reallocation scheme. Note that we can probably provide
hints on all allocations, and refine the heuristics that use them
for general tensors.

Reviewed By: bixia

Differential Revision: https://reviews.llvm.org/D143309
2023-02-06 14:08:53 -08:00
Peiming Liu
7fef8d69cc [mlir][sparse] implement bufferizableOpInterface for sparse_tensor.pack operation
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D143281
2023-02-03 23:55:59 +00:00
Peiming Liu
a41672e16a [mlir][sparse] implement lowering rules for sparse_tensor.pack operation
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D143230
2023-02-03 23:51:36 +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
Aart Bik
e2e6e7a6a3 [mlir][sparse] start using size_hint provided in allocation op
Even though we introduced the size_hint, we never used it.
This is a very first step, using the hint during the codegen path.
Note that we can refine the heuristics. Also, we need to start
adding the hint on all allocation generated for reading tensors,
converting tensors, etc.

Reviewed By: Peiming, bixia

Differential Revision: https://reviews.llvm.org/D143292
2023-02-03 14:02:41 -08:00
Aart Bik
28b92777c7 [mlir][sparse] add some documentation to storage layout (NFC)
in particular, the trailing COO optimization was not
desribed in the general layout description

Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D143284
2023-02-03 11:09:38 -08:00
bixia1
3b1c86cd0f [mlir][sparse] Implement heap sort for sparse_tensor.sort.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D142913
2023-02-02 15:36:38 -08:00
Rahul Kayaith
0ce25b1235 [mlir] Require explicit casts when using TypedValue
Currently `TypedValue` can be constructed directly from `Value`, hiding
errors that could be caught at compile time. For example the following
will compile, but crash/assert at runtime:
```
void foo(TypedValue<IntegerType>);
void bar(TypedValue<FloatType> v) {
  foo(v);
}
```

This change removes the constructors and replaces them with explicit
llvm casts.

Depends on D142852

Reviewed By: rriddle

Differential Revision: https://reviews.llvm.org/D142855
2023-02-01 21:54:53 -05:00
Matthias Springer
1ac248e485 [mlir][bufferization][NFC] Rename getAliasingOpOperand/getAliasingOpResult
* `getAliasingOpOperand` => `getAliasingOpOperands`
* `getAliasingOpResult` => `getAliasingOpResults`

Also a few minor code cleanups and better documentation.

Differential Revision: https://reviews.llvm.org/D142979
2023-02-01 10:07:41 +01:00
bixia1
8550aebd57 [mlir][sparse] Extend sorting function generator to support operand beyond (lo, hi, xs, ys).
This is to prepare for implementing a hybrid quick sort, which switches to heap
sort when the recursive depth exceeds certain limits.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D142731
2023-01-31 10:22:28 -08:00
Matthias Springer
34d65e81e8 [mlir][bufferization] Generalize and rename isMemoryWrite
The name of the method was confusing. It is bufferizesToMemoryWrite, but from the perspective of OpResults.

`bufferizesToMemoryWrite(OpResult)` now supports ops with regions that do not have aliasing OpOperands (such as `scf.if`). These ops no longer need to implement `isMemoryWrite`.

Differential Revision: https://reviews.llvm.org/D141684
2023-01-30 09:34:04 +01: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
7cec4d169d [mlir][sparse] Change the quick sort pivot selection.
Previously, we choose the value at (lo + hi)/2 as a pivot for partitioning the
data in [lo, hi). We now choose the median for the three values at lo, (lo +
hi)/2, and (hi-1) as a pivot to match the std::qsort implementation.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D142679
2023-01-29 17:56:57 -08:00
Aart Bik
9bcef91694 [mlir][sparse] fixed typo
Reviewed By: bixia

Differential Revision: https://reviews.llvm.org/D142784
2023-01-27 15:09:17 -08:00
Peiming Liu
a7bf2e558f [mlir][sparse] refactoring isAdmissibleTensorExp into codegen
This patch moves some utils into CodegenEnv class, it should make the code easier to follow and it eliminates several indirect value assignment that use `ptr**`.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D142040
2023-01-27 19:37:23 +00:00
Matthias Springer
fa639d3b43 [mlir][sparse] Implement BufferizableOpInterface for additional ops
The handling of unknown ops will be tightened in a subsequent change. All sparse_tensor ops should implement BufferizableOpInterface, otherwise, they are treated as "unknown" and additional buffer allocs/copies may be inserted around them.

Differential Revision: https://reviews.llvm.org/D142005
2023-01-27 17:56:55 +01:00
wren romano
9916ab03f1 [mlir][sparse] (re)introducing getRankedTensorType/getMemrefType
The bulk of D142074 seems to have gotten overwritten due to some sort of merge conflict (afaict there's no record of it having been reverted intentionally).  So this commit redoes those changes.  In addition to the original changes, this commit also:
* moves the definition of `getRankedTensorType` (from `Transforms/CodegenUtils.h` to `IR/SparseTensor.h`), so that it can be used by `IR/SparseTensorDialect.cpp`.
* adds `getMemRefType` as another abbreviation.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D142503
2023-01-25 11:29:54 -08:00
Aart Bik
14c0317fef [mlir][sparse] clean vectorization bail-out for VL=0
Fixes https://github.com/llvm/llvm-project/issues/59970

Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D142290
2023-01-23 12:12:05 -08:00
wren romano
255c3f1159 [mlir][sparse] factoring out getRankedTensorType helper function
Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D142074
2023-01-20 19:36:01 -08:00
Peiming Liu
83a50839b7 [mlir][sparse] avoid using mutable descriptor when unnecessary (NFC)
Use SparseTensorDescriptor whenever not calling setters, to avoid needing to create a temporal buffer for simple query purposes.

Reviewed By: bixia, wrengr

Differential Revision: https://reviews.llvm.org/D141953
2023-01-17 20:54:27 +00:00
Kazu Hirata
91682b2631 Remove redundant initialization of std::optional (NFC) 2023-01-14 14:06:18 -08:00
Kazu Hirata
0a81ace004 [mlir] Use std::optional instead of llvm::Optional (NFC)
This patch replaces (llvm::|)Optional< with std::optional<.  I'll post
a separate patch to remove #include "llvm/ADT/Optional.h".

This is part of an effort to migrate from llvm::Optional to
std::optional:

https://discourse.llvm.org/t/deprecating-llvm-optional-x-hasvalue-getvalue-getvalueor/63716
2023-01-14 01:25:58 -08:00
Kazu Hirata
a1fe1f5f77 [mlir] Add #include <optional> (NFC)
This patch adds #include <optional> to those files containing
llvm::Optional<...> or Optional<...>.

I'll post a separate patch to actually replace llvm::Optional with
std::optional.

This is part of an effort to migrate from llvm::Optional to
std::optional:

https://discourse.llvm.org/t/deprecating-llvm-optional-x-hasvalue-getvalue-getvalueor/63716
2023-01-13 21:05:06 -08:00
bixia1
6646664154 [mlir][sparse] Improve ConcatenateOp rewriting for annotated all dense result.
Previously, we rely on InsertOp to add values to the result, in the same way we
add values to a sparse tensor with compressed dimensions. We now direct store
values to the values buffer.

Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D141517
2023-01-13 08:36:54 -08:00
bixia1
ae8b1a1e30 [mlir][sparse] Minor code clean up (NFC).
Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D141554
2023-01-12 15:25:33 -08:00