Skip to content

Commit

Permalink
[mlir][sparse] added some clarification on sparse tensors ABI (llvm#9…
Browse files Browse the repository at this point in the history
  • Loading branch information
aartbik authored Jun 21, 2024
1 parent b284ced commit ad39315
Showing 1 changed file with 13 additions and 1 deletion.
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

0 comments on commit ad39315

Please sign in to comment.