41 Commits

Author SHA1 Message Date
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
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
71cc0f1c04 [mlir][sparse] introduce sparse_tensor::StorageSpecifierType and related operations on it
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D139961
2022-12-14 22:27:22 +00:00
bixia1
974b4bf9fd [mlir][sparse] Add expand_symmetry attribute to the new operator.
The attribute tells the operator to handle symmetric structures for 2D tensors.
By default, the operator assumes the input tensor is not symmetric.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D138230
2022-11-23 16:32:15 -08:00
Peiming Liu
4fa00ce15c [mlir][sparse] extend foreach operation to accept reduction arguments; fix sparse tensor rewriting patterns that do not propagate sparse tensor SSA properly.
This patch re-commit D137468 and D137463, which were reverted by mistakes.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D137579
2022-11-07 21:40:30 +00:00
Stella Stamenova
99171078bb Revert "[mlir][sparse] extend foreach operation to accept reduction arguments."
This reverts commit 53d5d3401120f2aa741a73a5a9ba0ce012ca532c.

This is causing a build failure on the windows mlir bot that was previously hidden by another sparse tensor change that caused failures:

https://lab.llvm.org/buildbot/#/builders/13/builds/28006
2022-11-07 09:34:10 -08:00
bixia1
cf24d49dc8 [mlir][sparse] Add sparse_tensor.sort_coo operator.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D137442
2022-11-07 08:23:51 -08:00
Peiming Liu
53d5d34011 [mlir][sparse] extend foreach operation to accept reduction arguments.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D137463
2022-11-04 23:34:16 +00:00
Aart Bik
80b08b68f2 [mlir][sparse] add a cursor to sparse storage scheme
This prepare a subsequent revision that will generalize
the insertion code generation. Similar to the support lib,
insertions become much easier to perform with some "cursor"
bookkeeping. Note that we, in the long run, could perhaps
avoid storing the "cursor" permanently and use some
retricted-scope solution (alloca?) instead. However,
that puts harder restrictions on insertion-chain operations,
so for now we follow the more straightforward approach.

Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D136800
2022-10-27 11:18:50 -07:00
bixia1
14504cae9a [mlir][sparse] Extend sparse_tensor.push_back to allow push_back a value n times.
Reviewed By: aartbik, Peiming

Differential Revision: https://reviews.llvm.org/D136653
2022-10-26 09:21:03 -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
Aart Bik
a3610359b5 [mlir][sparse] change memref argument to proper SSA components
The indices for insert/compress were previously provided as
a memref<?xindex> with proper rank, since that matched the
argument for the runtime support libary better. However, with
proper codegen coming, providing the indices as SSA values
is much cleaner. This also brings the sparse_tensor.insert
closer to unification with tensor.insert, planned in the
longer run.

Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D134404
2022-09-27 16:37:37 -07:00
bixia1
4329ca61e8 [mlir][sparse] Add sparse_tensor.sort operator.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D134482
2022-09-27 12:00:08 -07:00
Peiming Liu
e08865a12c [mlir][sparse] Introducing a new sparse_tensor.foreach operator.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D134484
2022-09-22 23:49:22 +00:00
Jim Kitchen
07150fece5 [mlir][sparse] Add sparse_tensor.select operation
The new select operation allows filtering of sparse tensors
by conditionally keeping or removing each element. This
can be used to remove negative values or select the upper
triangle of a matrix.

The select op has a single region which operates on a single
value and must return a boolean True to keep or False to drop.

Reviewed by: aartbik

Differential Revision: https://reviews.llvm.org/D133569
2022-09-13 15:22:53 -05:00
Aart Bik
f76dcede3f [mlir][sparse] rename lex_insert into insert
This change goes not impact any semantics yet, but it
is in preparation for implementing the unordered and not-unique
properties. Changing lex_insert to insert is a first step.

Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D133531
2022-09-08 17:26:35 -07:00
Aart Bik
610b09074a [mlir][sparse] change variable dimension to fixed attribute pointers/indices
The "sparsification" pass does not need the ability to use runtime values for
the dimension, so the only source for variability would have been user code.
Restricting the dimension to constants simplifies code generation.

Reviewed By: Peiming, wrengr

Differential Revision: https://reviews.llvm.org/D133458
2022-09-07 16:27:24 -07:00
Peiming Liu
edca72f5bc [mlir][sparse] Refactoring: remove dependence on tuple type when lowering sparse tensors.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D133390
2022-09-07 17:53:48 +00:00
Peiming Liu
c3aeb3e644 [mlir][sparse] Introduce sparse_tensor.storage operator to create a sparse tensor storage tuple
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D133231
2022-09-03 00:08:29 +00:00
Peiming Liu
7ea643c06d [mlir][sparse] Introduce new sparse_tensor.storage_get/set to access memory that stores the handle of a sparse tensor
Introduce new sparse_tensor.storage_get/set to access memory that stores the handle of a sparse tensor. The sparse tensor storage are represented as a tuple, these operation will later be eliminated and the tuple will be flattened after sparse tensor codegen

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D133049
2022-08-31 22:15:15 +00:00
Peiming Liu
de907138ec [mlir][sparse] Add new concatente operator to sparse tensor
See https://www.tensorflow.org/xla/operation_semantics#concatenate for the operator semantics

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D131111
2022-08-08 17:23:43 +00:00
Matthias Springer
27a431f5e9 [mlir][bufferization][NFC] Move sparse_tensor.release to bufferization dialect
This op used to belong to the sparse dialect, but there are use cases for dense bufferization as well. (E.g., when a tensor alloc is returned from a function and should be deallocated at the call site.) This change moves the op to the bufferization dialect, which now has an `alloc_tensor` and a `dealloc_tensor` op.

Differential Revision: https://reviews.llvm.org/D129985
2022-07-19 09:18:19 +02:00
Jim Kitchen
2b8a4d9ce1 [mlir][sparse] Introduce new reduce op
A new sparse_tensor operation allows for
custom reduction code to be injected during
linalg.generic lowering for sparse tensors.
An identity value is provided to indicate
the starting value of the reduction. A single
block region is required to contain the
custom reduce computation.

Reviewed by: aartbik

Differential Revision: https://reviews.llvm.org/D128004
2022-07-15 15:30:41 -05:00
wren romano
3cf03f1c56 [mlir][sparse] Adding IsSparseTensorPred and updating ops to use it
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D126994
2022-06-03 17:15:31 -07:00
Matthias Springer
6232a8f3d6 [mlir][sparse][NFC] Switch InitOp to bufferization::AllocTensorOp
Now that we have an AllocTensorOp (previously InitTensorOp) in the bufferization dialect, the InitOp in the sparse dialect is no longer needed.

Differential Revision: https://reviews.llvm.org/D126180
2022-06-02 00:03:52 +02:00
River Riddle
fb35cd3baf [mlir][NFC] Update textual references of func to func.func in SparseTensor tests
The special case parsing of `func` operations is being removed.
2022-04-20 22:17:29 -07:00
Jim Kitchen
414ed019ac [mlir][sparse] Introduce new binary and unary op
When the sparse_tensor dialect lowers linalg.generic,
it makes inferences about how the operations should
affect the looping logic. For example, multiplication
is an intersection while addition is a union of two
sparse tensors.

The new binary and unary op separate the looping logic
from the computation by nesting the computation code
inside a block which is merged at the appropriate level
in the lowered looping code.

The binary op can have custom computation code for the
overlap, left, and right sparse overlap regions. The
unary op can have custom computation code for the
present and absent values.

Reviewed by: aartbik

Differential Revision: https://reviews.llvm.org/D121018
2022-03-17 12:31:09 -05:00
Aart Bik
efa15f4178 [mlir][sparse] add ability for sparse tensor output
Rationale:
Although file I/O is a bit alien to MLIR itself, we provide two convenient ways
for sparse tensor I/O. The input part was already there (behind the swiss army
knife sparse_tensor.new). Now we have a sparse_tensor.out to write out data. As
before, the ops are kept vague and may change in the future. For now this
allows us to compare TACO vs MLIR very easily.

Reviewed By: bixia

Differential Revision: https://reviews.llvm.org/D117850
2022-01-21 15:43:29 -08:00
Aart Bik
4f2ec7f983 [mlir][sparse] finalize sparse output in the presence of reductions
This revision implements sparse outputs (from scratch) in all cases where
the loops can be reordered with all but one parallel loops outer. If the
inner parallel loop appears inside one or more reductions loops, then an
access pattern expansion is required (aka. workspaces in TACO speak).

Reviewed By: bixia

Differential Revision: https://reviews.llvm.org/D115091
2021-12-07 10:54:29 -08:00
Aart Bik
f66e5769d4 [mlir][sparse] first version of "truly" dynamic sparse tensors as outputs of kernels
This revision contains all "sparsification" ops and rewriting necessary to support sparse output tensors when the kernel has no reduction (viz. insertions occur in lexicographic order and are "injective"). This will be later generalized to allow reductions too. Also, this first revision only supports sparse 1-d tensors (viz. vectors) as output in the runtime support library. This will be generalized to n-d tensors shortly. But this way, the revision is kept to a manageable size.

Reviewed By: bixia

Differential Revision: https://reviews.llvm.org/D113705
2021-11-15 15:33:32 -08:00
Aart Bik
1e6ef0cfb0 [mlir][sparse] refine trait of sparse_tensor.convert
Rationale:
The currently used trait was demanding that all types are the same
which is not true (since the sparse part may change and the dim sizes
may be relaxed). This revision uses the correct trait and makes the
rank match test explicit in the verify method.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D112576
2021-10-26 14:36:49 -07:00
Aart Bik
9d1db3d4a1 [mlir][sparse] generalize sparse_tensor.convert on static/dynamic dimension sizes
This revison lifts the artificial restriction on having exact matches between
source and destination type shapes. A static size may become dynamic. We still
reject changing a dynamic size into a static size to avoid the need for a
runtime "assert" on the conversion. This revision also refactors some of the
conversion code to share same-content buffers.

Reviewed By: bixia

Differential Revision: https://reviews.llvm.org/D111915
2021-10-18 13:54:03 -07:00
Aart Bik
a652e5b53a [mlir][sparse] emergency fix after constant -> arith.constant change
Reviewed By: Mogball

Differential Revision: https://reviews.llvm.org/D111743
2021-10-13 10:26:17 -07:00
Aart Bik
35517a251d [mlir][sparse] add init sparse tensor operation
This is the first step towards supporting general sparse tensors as output
of operations. The init sparse tensor is used to materialize an empty sparse
tensor of given shape and sparsity into a subsequent computation (similar to
the dense tensor init operation counterpart).

Example:
  %c = sparse_tensor.init %d1, %d2 : tensor<?x?xf32, #SparseMatrix>
  %0 = linalg.matmul
    ins(%a, %b: tensor<?x?xf32>, tensor<?x?xf32>)
    outs(%c: tensor<?x?xf32, #SparseMatrix>) -> tensor<?x?xf32, #SparseMatrix>

Reviewed By: bixia

Differential Revision: https://reviews.llvm.org/D111684
2021-10-13 09:47:56 -07:00
Mogball
a54f4eae0e [MLIR] Replace std ops with arith dialect ops
Precursor: https://reviews.llvm.org/D110200

Removed redundant ops from the standard dialect that were moved to the
`arith` or `math` dialects.

Renamed all instances of operations in the codebase and in tests.

Reviewed By: rriddle, jpienaar

Differential Revision: https://reviews.llvm.org/D110797
2021-10-13 03:07:03 +00:00
Aart Bik
16b8f4ddae [mlir][sparse] add a "release" operation to sparse tensor dialect
We have several ways to materialize sparse tensors (new and convert) but no explicit operation to release the underlying sparse storage scheme at runtime (other than making an explicit delSparseTensor() library call). To simplify memory management, a sparse_tensor.release operation has been introduced that lowers to the runtime library call while keeping tensors, opague pointers, and memrefs transparent in the initial IR.

*Note* There is obviously some tension between the concept of immutable tensors and memory management methods. This tension is addressed by simply stating that after the "release" call, no further memref related operations are allowed on the tensor value. We expect the design to evolve over time, however, and arrive at a more satisfactory view of tensors and buffers eventually.

Bug:
http://llvm.org/pr52046

Reviewed By: bixia

Differential Revision: https://reviews.llvm.org/D111099
2021-10-05 09:35:59 -07:00
Aart Bik
697ea09d47 [mlir][sparse] add sparse tensor type conversion operation
Introduces a conversion from one (sparse) tensor type to another
(sparse) tensor type. See the operation doc for details. Actual
codegen for all cases is still TBD.

Reviewed By: ThomasRaoux

Differential Revision: https://reviews.llvm.org/D107205
2021-07-31 12:53:31 -07:00
Aart Bik
c2415d67a5 [mlir][sparse] fixed bug in verification
The order of testing in two sparse tensor ops was incorrect,
which could cause an invalid cast (crashing the compiler instead
of reporting the error). This revision fixes that bug.

Reviewed By: gussmith23

Differential Revision: https://reviews.llvm.org/D106841
2021-07-27 08:49:21 -07:00
Aart Bik
36b66ab9ed [mlir][sparse] add support for "simply dynamic" sparse tensor expressions
Slowly we are moving toward full support of sparse tensor *outputs*. First
step was support for all-dense annotated "sparse" tensors. This step adds
support for truly sparse tensors, but only for operations in which the values
of a tensor change, but not the nonzero structure (this was refered to as
"simply dynamic" in the [Bik96] thesis).

Some background text was posted on discourse:
https://llvm.discourse.group/t/sparse-tensors-in-mlir/3389/25

Reviewed By: gussmith23

Differential Revision: https://reviews.llvm.org/D104577
2021-06-22 13:37:32 -07:00
Aart Bik
727a63e0d9 [mlir][sparse] allow all-dense annotated "sparse" tensor output
This is a very careful start with alllowing sparse tensors at the
left-hand-side of tensor index expressions (viz. sparse output).
Note that there is a subtle difference between non-annotated tensors
(dense, remain n-dim, handled by classic bufferization) and all-dense
annotated "sparse" tensors (linearized to 1-dim without overhead
storage, bufferized by sparse compiler, backed by runtime support library).
This revision gently introduces some new IR to facilitate annotated outputs,
to be generalized to truly sparse tensors in the future.

Reviewed By: gussmith23, bixia

Differential Revision: https://reviews.llvm.org/D104074
2021-06-15 14:55:07 -07:00
Aart Bik
96a23911f6 [mlir][sparse] complete migration to sparse tensor type
A very elaborate, but also very fun revision because all
puzzle pieces are finally "falling in place".

1. replaces lingalg annotations + flags with proper sparse tensor types
2. add rigorous verification on sparse tensor type and sparse primitives
3. removes glue and clutter on opaque pointers in favor of sparse tensor types
4. migrates all tests to use sparse tensor types

NOTE: next CL will remove *all* obsoleted sparse code in Linalg

Reviewed By: bixia

Differential Revision: https://reviews.llvm.org/D102095
2021-05-10 12:55:22 -07:00