[Mlir-commits] [mlir] [mlir][sparse] added some clarification on sparse tensors ABI (PR #96331)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Jun 21 09:52:04 PDT 2024
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Aart Bik (aartbik)
<details>
<summary>Changes</summary>
---
Full diff: https://github.com/llvm/llvm-project/pull/96331.diff
1 Files Affected:
- (modified) mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td (+13-1)
``````````diff
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
index 75a479b898040..58c9b58352651 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
@@ -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
+ in and from 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
@@ -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 = [
``````````
</details>
https://github.com/llvm/llvm-project/pull/96331
More information about the Mlir-commits
mailing list