623 Commits

Author SHA1 Message Date
Peiming Liu
0e1429add5 [mlir][sparse] fix build error.
Reviewed By: aartbik, vzakhari

Differential Revision: https://reviews.llvm.org/D149874
2023-05-04 17:30:13 +00:00
Peiming Liu
36c95ee739 [mlir][sparse] group tensor id and levels into pairs in loop emitter
This addressed some unresolved comments in https://reviews.llvm.org/D142930

Reviewed By: aartbik, wrengr

Differential Revision: https://reviews.llvm.org/D148565
2023-05-04 16:15:42 +00:00
wren romano
55a74af71b [mlir][sparse] Adding a new overload of getOpaquePointerType
(This will be used in future patches, but is split off for easier reviewing)

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D149805
2023-05-03 18:39:20 -07:00
Aart Bik
9a018a7b48 [mlir][sparse] relax constraints on tensor.cast with pre-rewriting
Reviewed By: wrengr

Differential Revision: https://reviews.llvm.org/D149489
2023-05-01 16:03:44 -07:00
Peiming Liu
4ca1932776 [mlir][sparse] fix compiler warning in release build.
Differential Revision: https://reviews.llvm.org/D149607
2023-05-01 19:02:43 +00:00
Peiming Liu
d4db528938 [mlir][sparse] extend unpack operation to support unpacking a batched COO type
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D149103
2023-05-01 18:17:29 +00:00
Aart Bik
dbdb4affa0 [mlir][sparse] avoid slice rewriting when conditions are not met
Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D148964
2023-04-21 16:09:08 -07:00
Aart Bik
86888e420c [mlir][sparse][gpu] generate proper memcpy in/out host and device
The host registration is a convenient way to get CUDA kernels
running, but it may be slow and does not work for all buffer
(like global constants). This revision uses the proper alloc
copy dealloc chains for buffers, using asynchronous chains
to increase overlap. The host registration mechanism is
kept under a flag for the output, just for experimentation
purposes while this project ramps up.

Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D148682
2023-04-21 09:30:42 -07:00
Peiming Liu
a7cfcc686b [mlir][sparse] fix crash when generating coiteration loop with compressed-hi DLT.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D148842
2023-04-20 21:15:49 +00:00
Rahul Kayaith
6089d612a5 [mlir] Prevent implicit downcasting to interfaces
Currently conversions to interfaces may happen implicitly (e.g.
`Attribute -> TypedAttr`), failing a runtime assert if the interface
isn't actually implemented. This change marks the `Interface(ValueT)`
constructor as explicit so that a cast is required.

Where it was straightforward to I adjusted code to not require casts,
otherwise I just made them explicit.

Depends on D148491, D148492

Reviewed By: rriddle

Differential Revision: https://reviews.llvm.org/D148493
2023-04-20 16:31:54 -04:00
Peiming Liu
fd2211d84a use heap memory for position buffer allocated for PackOp.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D148818
2023-04-20 20:26:01 +00:00
Matthias Springer
4c48f016ef [mlir][Affine][NFC] Wrap dialect in "affine" namespace
This cleanup aligns the affine dialect with all the other dialects.

Differential Revision: https://reviews.llvm.org/D148687
2023-04-20 11:19:21 +09:00
Peiming Liu
7864d736cf [mlir][sparse] extend pack operation to support packing a batched COO type
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D148670
2023-04-20 01:35:30 +00:00
Peiming Liu
abd66d918a [mlir][sparse] support iteration over compressed-hi dimension level in loop emitter
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D148668
2023-04-20 00:57:08 +00:00
Peiming Liu
b9589545c4 [mlir][sparse] introduce a new compressed(hi) dimension level type
`compressed(hi)` is similar to `compressed`, but instead of reusing the previous position high as the current position low, it uses a pair of positions for each sparse index.

The patch only introduces the definition (syntax) but does not provide codegen implementation.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D148664
2023-04-18 23:26:11 +00:00
Peiming Liu
6a148c5aa7 [mlir][sparse] enable more sparse convolution kernels.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D147670
2023-04-17 17:43:52 +00:00
Peiming Liu
65aaabc8d1 [mlir][sparse] use macro to build common operations in LoopEmitter
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D147553
2023-04-17 16:57:50 +00:00
Peiming Liu
2cd15925f4 [mlir][sparse] implement index redution on dense level (for CSR)
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D147550
2023-04-17 16:36:31 +00:00
Aart Bik
4889214a48 [mlir][sparse][gpu] generate single module, unique kernel names
This fixes a TODO in the first version.

Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D148406
2023-04-15 17:25:36 -07:00
Peiming Liu
5fd9d80135 [mlir][sparse] extend loop emitter to emit slice driven loops
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D142930
2023-04-13 03:29:40 +00:00
Oleg Shyshkov
75ef84bf52 [mlir][bufferization] Make function boundary type convertion logic dynamic.
Having to choose from only static or dynamic layout for all function is limiting.

Differential Revision: https://reviews.llvm.org/D148074
2023-04-12 11:02:43 +02:00
wren romano
9d4df97ff0 [mlir][sparse] Canonicalizing arguments to genReshapeDstShape and foreachInSparseConstant
These functions don't need a`PatternRewriter`, they only need an `OpBuilder`.  And, the builder should be the first argument, before the `Location`, to match the style used everywhere else in MLIR.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D148059
2023-04-11 19:11:59 -07:00
Frederik Gossen
4a33c67995 Fix invalid op result access in SparseTensorCodegen
This trigges an assertion introduced in https://reviews.llvm.org/D147883

Differential Revision: https://reviews.llvm.org/D147950
2023-04-10 12:36:26 -04:00
wren romano
686ef4b4a8 [mlir][sparse] Simplifying Merger::expContainsTensor
`expContainsTensor` used to call `expIsTensor` to short-circuit the recursive calls; however, the very first thing `expContainsTensor` does is to check `expIsTensor`, so the short-circuiting code just causes the function to check that condition redundantly.

Depends On D146684

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D146688
2023-04-07 12:32:47 -07:00
Aart Bik
61c1ed86a5 [mlir][sparse][gpu] sparse GPU code generator pipeline setup
Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D147571
2023-04-05 15:12:54 -07:00
Aart Bik
19466ebc7f [mlir][sparse][gpu] a first prototype sparse GPU code generator
This implements a proof-of-concept GPU code generator
to the sparse compiler pipeline, currently only capable
of generating CUDA threads for outermost parallel loops.

The objective, obviously, is to grow this concept
to a full blown GPU code generator, capable of the
right combinaton of code generation as well as exploiting
idiomatic kernels or vector specific libraries (think cuSparse).

Reviewed By: ThomasRaoux

Differential Revision: https://reviews.llvm.org/D147483
2023-04-05 11:32:06 -07:00
wren romano
34c9c59ce4 [mlir][sparse] Using SparseTensorType in SparsePackOpConverter
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D147465
2023-04-03 16:36:56 -07:00
Peiming Liu
7b86f7c5d4 [mlir][sparse] support sparse bufferization.alloc_tensor with copy argument.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D147358
2023-03-31 22:27:23 +00:00
wren romano
962484aeca [mlir][sparse] Correcting some terminology / naming-scheme errors.
The name "coords" should be used for the complete tuple of Dimension-/Level-many "crd" values associated with a single element.  Whereas the name "coordinates" should only be used for collections of "crd" values which span several elements (e.g., the tensor's coordinates buffer for a single level).

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D147291
2023-03-30 18:26:34 -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
f2696e469a [mlir][sparse] Cleaning up some usage of SparseTensorType
This is a followup to D147192.

Reviewed By: aartbik, Peiming

Differential Revision: https://reviews.llvm.org/D147196
2023-03-30 12:00:00 -07:00
Peiming Liu
c24547e969 [mlir][sparse] avoid creating temporary unordered COO buffer when reshape sparse tensor.
Reviewed By: aartbik, wrengr

Differential Revision: https://reviews.llvm.org/D147192
2023-03-30 01:29:55 +00:00
wren romano
46a384dfbe [mlir][sparse] Preliminary code changes for ExprId, LatPointId, LatSetId newtypes
This commit contains several code changes which are ultimately required for converting the varions `Merger` identifiers from typedefs to newtypes.  The actual implementation of the newtypes themselves has been split off into separate commits, in hopes of simplifying the review process.

Depends On D146561

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D146684
2023-03-29 18:01:56 -07:00
Peiming Liu
33267f4007 [mlir][sparse] convert a sparse tensor slice to sparse tensor correctly.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D147074
2023-03-28 21:39:31 +00:00
Peiming Liu
c44d307c55 [mlir][sparse] add create-sparse-deallocs options to match the create-deallocs in BufferizationOption.
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D147010
2023-03-27 23:18:32 +00:00
wren romano
755285f1e9 [mlir][sparse] Factoring out LoopEmitter::isValidLevel
Depends On D146674

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D146676
2023-03-24 15:51:59 -07:00
wren romano
74c54206d7 [mlir][sparse] moving kInvalidId into "detail" namespace
In the next few commits I will be converting the various Merger identifier typedefs into newtypes; and once that's done, the `kInvalidId` constant will only be used internally and therefore does not need to be part of the public `mlir::sparse_tensor` namespace.

Depends On D146673

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D146674
2023-03-24 15:15:34 -07:00
wren romano
0280681b6d [mlir][sparse] Updating TensorExp ctor to catch unknown TensorExp::Kind
Depends On D146562

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D146673
2023-03-24 15:08:13 -07:00
wren romano
7c7c10a023 [mlir][sparse] Updating the Merger::{exp,lat,set} methods to return const
This helps the `Merger` maintain invariants, as well as clarifying the immutability of the underlying objects (with the one exception of `TensorExp::val`).

Depends On: D146559

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D146083
2023-03-24 14:48:33 -07:00
Peiming Liu
2b21327fee [mlir][sparse] fix crash when using pure constant index in indexing mapping (fixes #61530)
To address https://github.com/llvm/llvm-project/issues/61530

Reviewed By: aartbik, wrengr

Differential Revision: https://reviews.llvm.org/D146563
2023-03-21 23:45:20 +00:00
wren romano
13e9afd16d [mlir][sparse] Adding new Merger::addLat overload
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D146559
2023-03-21 16:22:04 -07:00
wren romano
72455b314f [mlir][sparse] Fixing -Wunused-variable in Sparsification.cpp
Reviewed By: aartbik, Peiming

Differential Revision: https://reviews.llvm.org/D146474
2023-03-20 16:53:19 -07:00
wren romano
1f58ae8066 [mlir][sparse] Making TensorExp::Kind a nested enum-class
This improves namespacing, and follows the pattern used for "Kind" enums elsewhere in MLIR.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D146086
2023-03-20 16:12:31 -07:00
Peiming Liu
1328bb6ef1 [mlir][sparse] extend loop emitter and optimize lattices with the awareness of slice based iteration
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D142929
2023-03-20 22:19:57 +00:00
Peiming Liu
d03805f2ee [mlir][sparse] add merger/topo sort support for slice-based affine sparse index codegen
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D142928
2023-03-20 21:24:10 +00:00
Peiming Liu
ee928fcde2 [mlir][sparse] add new sparisification option for dependent index reduction-based codegen
Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D142927
2023-03-16 20:10:58 +00:00
Jakub Kuderski
a0a76804c4 [ADT] Allow llvm::enumerate to enumerate over multiple ranges
This does not work by a mere composition of `enumerate` and `zip_equal`,
because C++17 does not allow for recursive expansion of structured
bindings.

This implementation uses `zippy` to manage the iteratees and adds the
stream of indices as the first zipped range. Because we have an upfront
assertion that all input ranges are of the same length, we only need to
check if the second range has ended during iteration.

As a consequence of using `zippy`, `enumerate` will now follow the
reference and lifetime semantics of the `zip*` family of functions. The
main difference is that `enumerate` exposes each tuple of references
through a new tuple-like type `enumerate_result`, with the familiar
`.index()` and `.value()` member functions.

Because the `enumerate_result` returned on dereference is a
temporary, enumeration result can no longer be used through an
lvalue ref.

Reviewed By: dblaikie, zero9178

Differential Revision: https://reviews.llvm.org/D144503
2023-03-15 19:34:22 -04: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
wren romano
b60de1dfcc [mlir][sparse] Updating Merger::foreachTensorLoopId to take LatPointId
Since all callsites of `foreachTensorLoopId` would simply look up the `LatPointId` to extract its `BitVector`, it's cleaner to let the `Merger` handle that instead.  This seems to better capture the intent of the `foreachTensorLoopId` method, and improves decoupling (since it removes a place that leaks the implementation detail that we use `BitVector`).

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D146082
2023-03-15 12:27:47 -07:00
Jakub Kuderski
8c258fda1f [ADT][mlir][NFCI] Do not use non-const lvalue-refs with enumerate
Replace references to enumerate results with either result_pairs
(reference wrapper type) or structured bindings. I did not use
structured bindings everywhere as it wasn't clear to me it would
improve readability.

This is in preparation to the switch to zip semantics which won't
support non-const lvalue reference to elements:
https://reviews.llvm.org/D144503.

I chose to use values instead of const lvalue-refs because MLIR is
biased towards avoiding `const` local variables. This won't degrade
performance because currently `result_pair` is cheap to copy (size_t
+ iterator), and in the future, the enumerator iterator dereference
will return temporaries anyway.

Reviewed By: dblaikie

Differential Revision: https://reviews.llvm.org/D146006
2023-03-15 10:43:56 -04:00