298 Commits

Author SHA1 Message Date
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
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
Peiming Liu
ab99b5d1f6 [mlir][sparse] deduplicate non-unique coordinates unconditionally
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D145621
2023-03-09 21:59:57 +00:00
Peiming Liu
6df483c9a0 [mlir][sparse] add a check test for foreach operation on constant sparse tensor
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D145728
2023-03-09 21:25:37 +00:00
Peiming Liu
41089f86e3 [mlir][sparse] fix bugs when convert coo to coo but with different dim ordering
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D145723
2023-03-09 20:55:03 +00:00
Peiming Liu
55270f56d2 [mlir][sparse] fix a bug in unpack op that used wrong compare predicate.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D145603
2023-03-08 19:52:09 +00:00
Peiming Liu
cc009334eb [mlir][sparse] deduplicate non-unique coordinates when coiterating COO tensors
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D145518
2023-03-07 21:52:38 +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
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
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
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
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
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
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
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
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
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
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
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
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
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
Quentin Colombet
add43ff180 Re-apply "[mlir][SparseTensor] Add a few more tests for sparse vectorization"
This reverts commit 93f40c983e0adbb63cbb7c59814090134d691dd1.

Update the tests to also work on window.
The order in which the `arith.constant`s appear in the output IR is
slightly different between window and linux.
Use CHECK.*-DAG for the constants.

Original message:
These tests cover muli, xor, and, addf, subf, and addi.

The tests themselves are not that interesting, their goal is to provide
code coverage for all the types of reductions currently supported.

NFC

Differential Revision: https://reviews.llvm.org/D141369
2023-01-13 13:02:19 +00:00
Quentin Colombet
93f40c983e Revert "[mlir][SparseTensor] Add a few more tests for sparse vectorization"
This reverts commit 904f2ccc3ba1d3aaf94140aa4595fd41af67d897.

This breaks a window bot. Reverting while I investigate.
https://lab.llvm.org/buildbot/#/builders/13/builds/30748
2023-01-13 10:30:09 +00:00
Quentin Colombet
904f2ccc3b [mlir][SparseTensor] Add a few more tests for sparse vectorization
These tests cover muli, xor, and, addf, subf, and addi.

The tests themselves are not that interesting, their goal is to provide
code coverage for all the types of reductions currently supported.

NFC

Differential Revision: https://reviews.llvm.org/D141369
2023-01-13 09:42:07 +00:00
Peiming Liu
885a1f4316 [mlir][sparse] support parsing slices in sparse tensor encoding attribute
Reviewed By: bixia

Differential Revision: https://reviews.llvm.org/D140712
2023-01-12 22:35:24 +00:00
bixia1
b78b547371 [mlir][sparse] Avoid values buffer reallocation for annotated all dense tensors.
Previously, we rely on the InsertOp to gradually increase the size of the
storage for all sparse tensors. We now allocate the full size values buffer
for annotated all dense tensors when we first allocate the tensor. This avoids
the cost of gradually increasing the buffer and allows accessing the values
buffer as if it were a dense tensor.

Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D141516
2023-01-11 16:31:07 -08:00
bixia1
14aba2084d [mlir][sparse] Improve the rewriting for dense-to-sparse conversion.
Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D141335
2023-01-11 08:27:04 -08:00
bixia1
f3fd739d39 [mlir][sparse] Improve the rewriting for NewOp with dimension ordering.
Previously, we use a temporary tensor with identity ordering. We now use a
temporary tensor with the destination dimension ordering, to enable the use of
sort_coo for sorting the tensor.

Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D141295
2023-01-09 14:40:33 -08:00
bixia1
81e3079d0f [mlir][sparse] Replace sparse_tensor.sort with sparse_tensor.sort_coo for sorting COO tensors.
Add codegen pattern for sparse_tensor.indices_buffer.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D140871
2023-01-05 15:42:57 -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
bixia1
3fdd85da06 [mlir][sparse] Add AOS optimization.
Use an array of structures to represent the indices for the tailing COO region
of a sparse tensor.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D140870
2023-01-04 18:16:04 -08:00
Aart Bik
1c7ffe0c38 [mlir][sparse] add test that combines sparse codegen and lowering to llvm struct
Reviewed By: bixia

Differential Revision: https://reviews.llvm.org/D141006
2023-01-04 12:12:31 -08:00
bixia1
90aa436291 [mlir][sparse] Add layout to the memref for the indices buffers to prepare for the AOS storage optimization for COO regions.
Fix relevant FileCheck tests.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D140742
2023-01-04 07:36:11 -08:00
bixia1
840e2ba336 [mlir][sparse] Use DLT in the mangled function names for insertion.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D140484
2022-12-28 08:21:22 -08:00
Aart Bik
431f6a543e [sparse][mlir][vectorization] add support for shift-by-invariant
Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D140596
2022-12-27 11:07:13 -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
083ddffe47 [mlir][sparse] introduce sparse_tensor::StorageSpecifierToLLVM pass
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D140122
2022-12-22 22:45:15 +00:00
Peiming Liu
a3672add76 [mlir][sparse] avoid unnecessary tmp COO buffer and convert when lowering ConcatentateOp.
When concat along dim 0, and all inputs/outputs are ordered with identity dimension ordering,
the concatenated coordinates will be yield in lexOrder, thus no need to sort.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D140228
2022-12-16 18:26:39 +00:00