Skip to content

[mlir][sparse] added some clarification on sparse tensors ABI #96331

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jun 21, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 13 additions & 1 deletion mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,17 @@ include "mlir/Pass/PassBase.td"
def SparseAssembler : Pass<"sparse-assembler", "ModuleOp"> {
let summary = "Add [dis]assemble operations on external sparse tensors";
let description = [{
A pass that converts public entry methods that use sparse tensors as
Unlike dense tensors, MLIR does **not** provide a direct `_mlir_ciface_`
ABI for passing sparse tensors as arguments from and to external methods
(within MLIR-generated methods, sparse tensors can be freely passed
around, but this eventually uses a bespoke parameter passing format
that is subject to change; like opaque pointers when the sparse runtime
support library is used or the constituent arrays and structs for
direct IR codegen). The sparse assembler pass, however, can be used
to obtain a stable `_mlir_ciface_` API for passing sparse tensors
from and to an external environment, such as Python, PyTorch, or JAX.

The pass converts public entry methods that use sparse tensors as
input parameters and/or output return values into wrapper methods
that [dis]assemble the individual tensors that constitute the actual
storage used externally into MLIR sparse tensors. This pass can be used
Expand All @@ -27,6 +37,8 @@ def SparseAssembler : Pass<"sparse-assembler", "ModuleOp"> {
By default, the pass uses the [dis]assemble operations to input and output
sparse tensors. When the direct-out option is set, however, the output
directly returns the MLIR allocated buffers to the external runtime.

The pass should always run before the actual sparsification passes.
}];
let constructor = "mlir::createSparseAssembler()";
let dependentDialects = [
Expand Down
Loading