20 Commits

Author SHA1 Message Date
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
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
Peiming Liu
4176ce61f1
[mlir][sparse] fix logical error when generating sort_coo. (#66690)
To fix issue: https://github.com/llvm/llvm-project/issues/66664
2023-09-18 15:26:01 -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
bixia1
abb05014f9 [mlir][sparse] Modify the pivot selection method for quick sort.
Previously, we choose the median of three values. We now choose the median of
five values when the number of values being sorted exceed a threshold
(currently 100). This is similar to std::sort.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D145534
2023-03-15 13:53:00 -07:00
bixia1
2ef416273f [mlir][sparse] Improve sort operation by generating inlined code to compare values.
Previously, we generate function calls to compare values for sorting. It turns
out that the compiler doesn't inline those function calls. We now directly
generate inlined code. Also, modify the code for comparing values to use less
number of branches.

This improves all sort implementation in general. For arabic-2005.mtx CSR, the
improvement is around 25%.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D145442
2023-03-14 15:14:49 -07:00
bixia1
f6424d11cb [mlir][sparse] Improve quick sort by using a loop to sort the bigger partition.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D145440
2023-03-10 20:43:08 -08: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
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
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
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
Hanhan Wang
0a1569a400 [mlir][NFC] Remove trailing whitespaces from *.td and *.mlir files.
This is generated by running

```
sed --in-place 's/[[:space:]]\+$//' mlir/**/*.td
sed --in-place 's/[[:space:]]\+$//' mlir/**/*.mlir
```

Reviewed By: rriddle, dcaballe

Differential Revision: https://reviews.llvm.org/D138866
2022-11-28 15:26:30 -08:00
bixia1
4f729d5a70 [mlir][sparse] Add rewriting rules for sparse_tensor.sort_coo.
Refactor the rewriting of sparse_tensor.sort to support the implementation of
sparse_tensor.sort_coo.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D137522
2022-11-14 08:48:53 -08:00
bixia1
9b800bf79d [mlir][sparse] Improve the non-stable sort implementation.
Replace the quick sort partition method with one that is more similar to the
method used by C++ std quick sort. This improves the runtime for sorting
sk_2005.mtx by more than 10x.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D137290
2022-11-07 07:38:42 -08:00
bixia1
d45be88736 [mlir][sparse] Implement the rewrite for sparse_tensor.push_back a value n times.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D136654
2022-10-31 08:19:12 -07:00
Aart Bik
86bf62fa64 [mlir][sparse] improve push_back type checking, printing, parsing
Rationale:
Enforces type consistency on parsed and generated IR.

Reviewed By: bixia

Differential Revision: https://reviews.llvm.org/D136132
2022-10-18 09:55:25 -07:00
bixia1
9409bbb2e0 [mlir][sparse] Implement insertion sort for the stable sort operator.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D135182
2022-10-06 09:48:39 -07:00
bixia1
1c835b5a8e [mlir][sparse] Allow the push_back operator to skip capacity check and reallocation.
Add UnitAttr `inbounds` for this purpose.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D134913
2022-09-29 16:38:06 -07:00
bixia1
654bbbde55 [mlir][sparse] Move the implementation of sparse_tensor.push_back to the buffer rewriter.
Reviewed By: aartbik, Peiming

Differential Revision: https://reviews.llvm.org/D134777
2022-09-29 15:06:00 -07:00
bixia1
062e515b70 [mlir][sparse] Add rewrite rule for the sort operator.
Add sparse-buffer-rewrite pass to rewrite sparse primitives on buffers to MLIR
implementation.

Add sparse rewrite rule for the sort operator.

Add FileCheck test and integration test.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D134627
2022-09-29 11:38:19 -07:00