[Mlir-commits] [mlir] [mlir][sparse] added some clarification on sparse tensors ABI (PR #96331)

[Aart Bik]

https://github.com/aartbik updated https://github.com/llvm/llvm-project/pull/96331

>From 11076bdbe13af4517b3b608e7f851536e5e218f9 Mon Sep 17 00:00:00 2001
From: Aart Bik <ajcbik at google.com>
Date: Fri, 21 Jun 2024 09:50:32 -0700
Subject: [PATCH 1/2] [mlir][sparse] added some clarification on sparse tensors
 ABI

---
 .../mlir/Dialect/SparseTensor/Transforms/Passes.td | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

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 = [

>From 0a99dbc8e4f0ffe6b359f9c31b4dcbbb027bd8a7 Mon Sep 17 00:00:00 2001
From: Aart Bik <ajcbik at google.com>
Date: Fri, 21 Jun 2024 09:54:01 -0700
Subject: [PATCH 2/2] edit

---
 mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
index 58c9b58352651..8ec18a1e18648 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
@@ -22,9 +22,9 @@ def SparseAssembler : Pass<"sparse-assembler", "ModuleOp"> {
     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.
+    from and to an external environment, such as Python, PyTorch, or JAX.
 
-    The pass  converts public entry methods that use sparse tensors as
+    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