Commit Graph

59 Commits

Author SHA1 Message Date
Peiming Liu
2cc4b3d07c [mlir][sparse] code cleanup using the assumption that dim2lvl maps ar… (#72894)
…e simplified.
2023-11-20 10:25:42 -08:00
Aart Bik
83cf0dc982 [mlir][sparse] implement direct IR alloc/empty/new for non-permutations (#72585)
This change implements the correct *level* sizes set up for the direct
IR codegen fields in the sparse storage scheme. This brings libgen and
codegen together again.

This is step 3 out of 3 to make sparse_tensor.new work for BSR
2023-11-16 17:17:41 -08:00
Yinying Li
c5a67e16b6 [mlir][sparse] Use variable instead of inlining sparse encoding (#72561)
Example:

#CSR = #sparse_tensor.encoding<{
  map = (d0, d1) -> (d0 : dense, d1 : compressed),
}>

// CHECK: #[[$CSR.*]] = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0
: dense, d1 : compressed) }>
// CHECK-LABEL: func private @sparse_csr(
// CHECK-SAME: tensor<?x?xf32, **#[[$CSR]]**>)
func.func private @sparse_csr(tensor<?x?xf32, #CSR>)
2023-11-16 19:30:21 -05:00
Christian Ulmann
dcae289d3a [MLIR][SparseTensor] Introduce opaque pointers in LLVM dialect lowering (#70570)
This commit changes the SparseTensor LLVM dialect lowering from using
`llvm.ptr<i8>` to `llvm.ptr`. This change ensures that the lowering now
properly relies on opaque pointers, instead of working with already type
erased i8 pointers.
2023-10-31 07:34:49 +01:00
Peiming Liu
f82bee1367 [mlir][sparse] split post-sparsification-rewriting into two passes. (#70727) 2023-10-30 15:22:21 -07:00
Peiming Liu
c780352de9 [mlir][sparse] implement sparse_tensor.lvl operation. (#69993) 2023-10-24 13:23:28 -07:00
Aart Bik
d3af65358d [mlir][sparse] introduce MapRef, unify conversion/codegen for reader (#68360)
This revision introduces a MapRef, which will support a future
generalization beyond permutations (e.g. block sparsity). This revision
also unifies the conversion/codegen paths for the sparse_tensor.new
operation from file (eg. the readers). Note that more unification is
planned as well as general affine dim2lvl and lvl2dim (all marked with
TODOs).
2023-10-06 13:42:01 -07:00
Yinying Li
6280e23124 [mlir][sparse] Print new syntax (#68130)
Printing changes from `#sparse_tensor.encoding<{ lvlTypes = [
"compressed" ] }>` to `map = (d0) -> (d0 : compressed)`. Level
properties, ELL and slice are also supported.
2023-10-04 16:36:05 -04:00
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
Yinying Li
79b9d41bd7 [mlir][sparse] Generalize sparse encoding in check tests (#67476)
For all the mlir tests (except for roundtrip_coding.mlir), change the
check test to use general form of encoding
`#sparse_tensor.encoding<{{{.*}}}>` instead of actual encoding such as
`#sparse_tensor.encoding<{ lvlTypes = [ "compressed", "singleton" ] }>`.
2023-09-26 16:56:06 -04:00
Aart Bik
3e4a8c2c7d [mlir][sparse] remove most bufferization.alloc_tensor ops from sparse (#66847)
The only ones left need actual deprecation in bufferization module.
2023-09-20 09:51:08 -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
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
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
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
wren romano
76647fce13 [mlir][sparse] Combining dimOrdering+higherOrdering fields into dimToLvl
This is a major step along the way towards the new STEA design.  While a great deal of this patch is simple renaming, there are several significant changes as well.  I've done my best to ensure that this patch retains the previous behavior and error-conditions, even though those are at odds with the eventual intended semantics of the `dimToLvl` mapping.  Since the majority of the compiler does not yet support non-permutations, I've also added explicit assertions in places that previously had implicitly assumed it was dealing with permutations.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D151505
2023-05-30 15:19:50 -07:00
wren romano
a0615d020a [mlir][sparse] Renaming the STEA field dimLevelType to lvlTypes
This commit is part of the migration of towards the new STEA syntax/design.  In particular, this commit includes the following changes:
* Renaming compiler-internal functions/methods:
  * `SparseTensorEncodingAttr::{getDimLevelType => getLvlTypes}`
  * `Merger::{getDimLevelType => getLvlType}` (for consistency)
  * `sparse_tensor::{getDimLevelType => buildLevelType}` (to help reduce confusion vs actual getter methods)
* Renaming external facets to match:
  * the STEA parser and printer
  * the C and Python bindings
  * PyTACO

However, the actual renaming of the `DimLevelType` itself (along with all the "dlt" names) will be handled in a separate commit.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D150330
2023-05-17 14:24:09 -07:00
bixia1
6071f6fd67 [mlir][sparse] Fix a problem in handling data type conversion.
Previously, the genCast function generates arith.trunci for converting f32 to
i32. Fix the function to use mlir::convertScalarToDtype to correctly handle
conversion cases beyond index casting.

Add a test case for codegen the sparse_tensor.convert op.

Reviewed By: aartbik, Peiming, wrengr

Differential Revision: https://reviews.llvm.org/D147272
2023-03-30 14:54:53 -07: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
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
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
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
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
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
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
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
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
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
bixia1
2aceadda78 [mlir][sparse] Put the implementation for the insertion operation to subroutines.
Previously, we generated inlined implementation for insert operation and
observed MLIR compile time increase due to the size of the main routine. We now
put the insert operation implementation in subroutines and leave the inlining
decision to the MLIR compiler.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D138957
2022-12-01 13:23:59 -08:00
Aart Bik
70633a8d55 [mlir][sparse] first general insertion implementation with pure codegen
This revision generalizes lowering the sparse_tensor.insert op into actual code that directly operates on the memrefs of a sparse storage scheme. The current insertion strategy does *not* rely on a cursor anymore, with introduces some testing overhead for each insertion (but still proportional to the rank, as before). Over time, we can optimize the code generation, but this version enables us to finish the effort to migrate from library to actual codegen.

Things to do:
(1) carefully deal with (un)ordered and (not)unique
(2) omit overhead when not needed
(3) optimize and specialize
(4) try to avoid the pointer "cleanup" (at HasInserts), and make sure the storage scheme is consistent at every insertion point (so that it can "escape" without concerns).

Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D137457
2022-11-08 13:10:05 -08: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
Aart Bik
0f3e4d1afa [mlir][sparse] lower number of entries op to actual code
works both along runtime path and pure codegen path

Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D136389
2022-10-21 10:48:37 -07:00
bixia1
0128f80167 [mlir][sparse] Fix the codegen for the convert operator to handle hidden nop convert.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D136291
2022-10-20 12:30:21 -07:00
Aart Bik
96cab659a1 [mlir][sparse] end-to-end sparse vector insertion codegen
Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D136275
2022-10-19 12:32:20 -07:00
bixia1
58b449c3bb [mlir][sparse] Replace the folding of nop convert with a codegen rule.
This is to allow the use of a nop convert to express that the sparse tensor
allocated through bufferization::AllocTensorOp will be expanded to sparse
tensor storage by sparse tensor codegen.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D136214
2022-10-19 10:20:47 -07:00
Aart Bik
d22df0ebba [mlir][sparse] refine insertion code
builds SSA cycle for compress insertion loop
adds casting on index mismatch during push_back

Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D136186
2022-10-18 14:16:38 -07:00
Aart Bik
9f596a7c67 [mlir][sparse] implement simple codegen for insertion (and related ops)
This is a proof of concept insertion implementation that sets up
the basic framework and implements it with push backs for just
sparse vectors. It adds insertion/compression through SSA values,
so that we properly update the memref after after pushback operation.

Note that properly using SSA values in sparsification is still TBD
but I will wait until Peiming's loop emitter is in to avoid conflicts.

Reviewed By: wrengr

Differential Revision: https://reviews.llvm.org/D136008
2022-10-17 18:02:08 -07:00
Aart Bik
94e4928bd3 [mlir][sparse] minor edit in doc, removed TAB from test
Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D134924
2022-09-30 09:42:36 -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
Aart Bik
4d06861950 [mlir][sparse] add "sort" to the compress op codegen
This revision also adds convenience methods to test the
dim level type/property (with the codegen being first client)

Reviewed By: bixia

Differential Revision: https://reviews.llvm.org/D134776
2022-09-28 10:41:40 -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
4132bce9e5 [mlir][sparse] Add codegen rule for the push_back operator.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D134372
2022-09-22 09:09:49 -07:00
Aart Bik
3986c86986 [mlir][sparse] partially implement codegen for sparse_tensor.compress
Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D133912
2022-09-15 10:32:33 -07:00
Aart Bik
6607fdf749 [mlir][sparse] add memSizes array to sparse storage format
Rationale:
For every dynamic memref (memref<?xtype>), the stored size really
indicates the capacity and the entry in the memSizes indicates
the actual size. This allows us to use memref's as "vectors".

Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D133724
2022-09-12 14:04:01 -07:00
bixia1
8a583bd53d [mlir][sparse] Add codegen for expand op.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D133454
2022-09-08 14:06:01 -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