9 Commits

Author SHA1 Message Date
Aart Bik
0e34dbb4f4
[mlir][sparse] fix bug with all-dense assembler (#108615)
When only all-dense "sparse" tensors occur in a function prototype, the
assembler would skip the method conversion purely based on input/output
counts. It should rewrite based on the presence of any annotation,
however.
2024-09-13 17:24:48 -07:00
Christian Sigg
a5757c5b65
Switch member calls to isa/dyn_cast/cast/... to free function calls. (#89356)
This change cleans up call sites. Next step is to mark the member
functions deprecated.

See https://mlir.llvm.org/deprecation and
https://discourse.llvm.org/t/preferred-casting-style-going-forward.
2024-04-19 15:58:27 +02:00
Aart Bik
5122a2c232
[mlir][sparse] allow for direct-out passing of sparse tensor buffers (#88327)
In order to support various external frameworks (JAX vs PyTorch) we need
a bit more flexibility in [dis]assembling external buffers to and from
sparse tensors in MLIR land. This PR adds a direct-out option that
avoids the rigid pre-allocated for copy-out semantics.

Note that over time, we expect the [dis]assemble operations to converge
into something that supports all sorts of external frameworks. Until
then, this option helps in experimenting with different options.
2024-04-11 10:07:24 -07:00
Peiming Liu
fc9f1d49aa
[mlir][sparse] use a consistent order between [dis]assembleOp and sto… (#84079)
…rage layout.
2024-03-06 09:57:41 -08:00
Peiming Liu
52b69aa32f
[mlir][sparse] support sparsifying batch levels (#83898) 2024-03-04 14:39:06 -08:00
Peiming Liu
f40ee6e83f
[mlir][sparse] assemble SoA COO correctly. (#82449) 2024-02-20 18:46:34 -08:00
Aart Bik
4d273b948e
[mlir][sparse] ensure [dis]assembler wrapper methods properly inline (#81907) 2024-02-15 11:39:32 -08:00
Aart Bik
d00e6d07b1
[mlir][sparse] refine sparse assembler strategy (#80521)
Rewrite *all* public methods, making original internal, private methods,
and exposing wrappers under the original name. This works a bit better
in practice (when combined with c-interface mechanism of torch-mlir for
example).
2024-02-05 10:48:18 -08:00
Aart Bik
33b463ad99
[mlir][sparse] external entry method wrapper for sparse tensors (#80326)
Similar to the emit_c_interface, this pull request adds a pass that
converts public entry methods that use sparse tensors as input
parameters and/or output return values into wrapper functions that
[dis]assemble the individual tensors that constitute the actual storage
used externally into MLIR sparse tensors. This pass can be used to
prepare the public entry methods of a program that is compiled by the
MLIR sparsifier to interface with an external runtime, e.g., when
passing sparse tensors as numpy arrays from and to Python. Note that
eventual bufferization decisions (e.g. who [de]allocates the underlying
memory) should be resolved in agreement with the external runtime
(Python, PyTorch, JAX, etc.)
2024-02-01 13:32:52 -08:00