[Mlir-commits] [mlir] [mlir][sparse] Expose runtime library option in sparsification-and-bufferization (PR #233217)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Oct 18 02:52:22 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Vito Secona (secona)
<details>
<summary>Changes</summary>
This change adds the enable-runtime-library option to the -sparsification-and-bufferization pass. With this, the -sparsification-and-bufferization pass has all options from SparsificationOptions.
Assisted by Codex for the test checks.
---
Full diff: https://github.com/llvm/llvm-project/pull/233217.diff
3 Files Affected:
- (modified) mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td (+5)
- (modified) mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp (+8-5)
- (added) mlir/test/Dialect/SparseTensor/sparsification_and_bufferization_runtime_lib.mlir (+91)
``````````diff
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
index 0960254a314692..bea8fb4515dd30 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
@@ -498,6 +498,9 @@ def SparsificationAndBufferization : Pass<"sparsification-and-bufferization", "M
"Enable i32 indexing into vectors (for efficient gather/scatter)">,
Option<"enableGPULibgen", "enable-gpu-libgen", "bool", "false",
"Enable GPU acceleration by means of direct library calls">,
+ //
+ // The options below are SparsificationPass' options copied over.
+ //
Option<"sparseEmitStrategy", "sparse-emit-strategy", "mlir::SparseEmitStrategy",
"mlir::SparseEmitStrategy::kFunctional",
"Emit functional code or interfaces (to debug) for sparse loops", [{llvm::cl::values(
@@ -524,6 +527,8 @@ def SparsificationAndBufferization : Pass<"sparsification-and-bufferization", "M
clEnumValN(mlir::SparseParallelizationStrategy::kAnyStorageAnyLoop,
"any-storage-any-loop",
"Enable sparse parallelization for any storage and loop."))}]>,
+ Option<"enableRuntimeLibrary", "enable-runtime-library", "bool",
+ "false", "Enable runtime library for manipulating sparse tensors">,
];
}
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
index b38d43476efb7f..2605ca15188219 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
@@ -59,7 +59,7 @@ class SparsificationAndBufferizationPass
: public impl::SparsificationAndBufferizationBase<
SparsificationAndBufferizationPass> {
public:
- // Private pass options only.
+ // Bufferization, sparsification, and sparse storage options.
SparsificationAndBufferizationPass(
const bufferization::OneShotBufferizationOptions &bufferizationOptions,
const SparsificationOptions &sparsificationOptions,
@@ -68,8 +68,9 @@ class SparsificationAndBufferizationPass
: bufferizationOptions(bufferizationOptions),
sparsificationOptions(sparsificationOptions),
createSparseDeallocs(createSparseDeallocs),
- enableRuntimeLibrary(enableRuntimeLibrary),
- enableBufferInitialization(enableBufferInitialization) {}
+ enableBufferInitialization(enableBufferInitialization) {
+ this->enableRuntimeLibrary = enableRuntimeLibrary;
+ }
// Private pass options and visible pass options.
SparsificationAndBufferizationPass(
const bufferization::OneShotBufferizationOptions &bufferizationOptions,
@@ -81,9 +82,9 @@ class SparsificationAndBufferizationPass
: bufferizationOptions(bufferizationOptions),
sparsificationOptions(sparsificationOptions),
createSparseDeallocs(createSparseDeallocs),
- enableRuntimeLibrary(enableRuntimeLibrary),
enableBufferInitialization(enableBufferInitialization) {
// Set the visible pass options explicitly.
+ this->enableRuntimeLibrary = enableRuntimeLibrary;
vectorLength = vl;
enableVLAVectorization = vla;
enableSIMDIndex32 = index32;
@@ -129,6 +130,9 @@ class SparsificationAndBufferizationPass
// Overrides the default parallelization strategy using user-provided value.
this->sparsificationOptions.parallelizationStrategy = parallelization;
+ // Keep sparsification and sparse storage lowering on the same runtime path.
+ this->sparsificationOptions.enableRuntimeLibrary = enableRuntimeLibrary;
+
// Run enabling transformations.
{
OpPassManager pm("builtin.module");
@@ -207,7 +211,6 @@ class SparsificationAndBufferizationPass
bufferization::OneShotBufferizationOptions bufferizationOptions;
SparsificationOptions sparsificationOptions;
bool createSparseDeallocs;
- bool enableRuntimeLibrary;
bool enableBufferInitialization;
};
diff --git a/mlir/test/Dialect/SparseTensor/sparsification_and_bufferization_runtime_lib.mlir b/mlir/test/Dialect/SparseTensor/sparsification_and_bufferization_runtime_lib.mlir
new file mode 100644
index 00000000000000..7fb4575b1d9347
--- /dev/null
+++ b/mlir/test/Dialect/SparseTensor/sparsification_and_bufferization_runtime_lib.mlir
@@ -0,0 +1,91 @@
+// RUN: mlir-opt %s --sparsification-and-bufferization --canonicalize | FileCheck %s --check-prefix=CODEGEN
+// RUN: mlir-opt %s --sparsification-and-bufferization="enable-runtime-library=false" --canonicalize | FileCheck %s --check-prefix=CODEGEN
+// RUN: mlir-opt %s --sparsification-and-bufferization="enable-runtime-library=true" --canonicalize | FileCheck %s --check-prefix=RUNTIME
+
+// CODEGEN: [[$SV:#.*]] = #sparse_tensor.encoding
+#SV = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }>
+
+// CODEGEN-LABEL: func.func @values(
+// CODEGEN-SAME: %[[POS:[^:]+]]: memref<?xindex>,
+// CODEGEN-SAME: %[[CRD:[^:]+]]: memref<?xindex>,
+// CODEGEN-SAME: %[[VAL:[^:]+]]: memref<?xf64>,
+// CODEGEN-SAME: %[[SPEC:[^:]+]]: !sparse_tensor.storage_specifier<[[$SV]]>) -> memref<?xf64> {
+// CODEGEN-NEXT: %[[SIZE:.*]] = sparse_tensor.storage_specifier.get %[[SPEC]] val_mem_sz : !sparse_tensor.storage_specifier<[[$SV]]>
+// CODEGEN-NEXT: %[[VALUES:.*]] = memref.subview %[[VAL]][0] [%[[SIZE]]] [1] : memref<?xf64> to memref<?xf64>
+// CODEGEN-NEXT: return %[[VALUES]] : memref<?xf64>
+// CODEGEN-NEXT: }
+
+// RUNTIME-LABEL: func.func @values(
+// RUNTIME-SAME: %[[TENSOR:.*]]: !llvm.ptr) -> memref<?xf64> {
+// RUNTIME-NEXT: %[[VALUES:.*]] = call @sparseValuesF64(%[[TENSOR]]) : (!llvm.ptr) -> memref<?xf64>
+// RUNTIME-NEXT: return %[[VALUES]] : memref<?xf64>
+// RUNTIME-NEXT: }
+
+func.func @values(%A: tensor<?xf64, #SV>) -> memref<?xf64> {
+ %values = sparse_tensor.values %A : tensor<?xf64, #SV> to memref<?xf64>
+ return %values : memref<?xf64>
+}
+
+// CODEGEN-LABEL: func.func @empty(
+// CODEGEN-SAME: %[[DIM:[^:]+]]: index) -> (memref<?xindex>, memref<?xindex>, memref<?xf64>, !sparse_tensor.storage_specifier<[[$SV]]>) {
+// CODEGEN: %[[C0:.*]] = arith.constant 0 : index
+// CODEGEN: %[[POS_ALLOC:.*]] = memref.alloc() : memref<16xindex>
+// CODEGEN-NEXT: %[[POS:.*]] = memref.cast %[[POS_ALLOC]] : memref<16xindex> to memref<?xindex>
+// CODEGEN-NEXT: %[[CRD_ALLOC:.*]] = memref.alloc() : memref<16xindex>
+// CODEGEN-NEXT: %[[CRD:.*]] = memref.cast %[[CRD_ALLOC]] : memref<16xindex> to memref<?xindex>
+// CODEGEN-NEXT: %[[VAL_ALLOC:.*]] = memref.alloc() : memref<16xf64>
+// CODEGEN-NEXT: %[[VAL:.*]] = memref.cast %[[VAL_ALLOC]] : memref<16xf64> to memref<?xf64>
+// CODEGEN-NEXT: %[[INIT:.*]] = sparse_tensor.storage_specifier.init : !sparse_tensor.storage_specifier<[[$SV]]>
+// CODEGEN-NEXT: %[[SPEC:.*]] = sparse_tensor.storage_specifier.set %[[INIT]] lvl_sz at 0 with %[[DIM]] : !sparse_tensor.storage_specifier<[[$SV]]>
+// CODEGEN-NEXT: %[[POS_SIZE:.*]] = sparse_tensor.storage_specifier.get %[[SPEC]] pos_mem_sz at 0 : !sparse_tensor.storage_specifier<[[$SV]]>
+// CODEGEN: memref.store %[[C0]], %[[POS1:.*]][%[[POS_SIZE]]] : memref<?xindex>
+// CODEGEN-NEXT: %[[SPEC1:.*]] = sparse_tensor.storage_specifier.set %[[SPEC]] pos_mem_sz at 0 with %[[SIZE1:.*]] : !sparse_tensor.storage_specifier<[[$SV]]>
+// CODEGEN: memref.store %[[C0]], %[[POS2:.*]][%[[SIZE1]]] : memref<?xindex>
+// CODEGEN-NEXT: %[[SPEC2:.*]] = sparse_tensor.storage_specifier.set %[[SPEC1]] pos_mem_sz at 0 with %[[SIZE2:.*]] : !sparse_tensor.storage_specifier<[[$SV]]>
+// CODEGEN-NEXT: return %[[POS2]], %[[CRD]], %[[VAL]], %[[SPEC2]] : memref<?xindex>, memref<?xindex>, memref<?xf64>, !sparse_tensor.storage_specifier<[[$SV]]>
+// CODEGEN-NEXT: }
+
+// RUNTIME-LABEL: func.func @empty(
+// RUNTIME-SAME: %[[DIM:[^:]+]]: index) -> !llvm.ptr {
+// RUNTIME-NEXT: %[[NULL:.*]] = llvm.mlir.zero : !llvm.ptr
+// RUNTIME-DAG: %[[C1_I32:.*]] = arith.constant 1 : i32
+// RUNTIME-DAG: %[[C0_I32:.*]] = arith.constant 0 : i32
+// RUNTIME-DAG: %[[C0:.*]] = arith.constant 0 : index
+// RUNTIME-DAG: %[[COMPRESSED:.*]] = arith.constant 262144 : i64
+// RUNTIME: %[[LVL_TYPES_ALLOC:.*]] = memref.alloca() : memref<1xi64>
+// RUNTIME-NEXT: %[[LVL_TYPES:.*]] = memref.cast %[[LVL_TYPES_ALLOC]] : memref<1xi64> to memref<?xi64>
+// RUNTIME-NEXT: memref.store %[[COMPRESSED]], %[[LVL_TYPES_ALLOC]][%[[C0]]] : memref<1xi64>
+// RUNTIME-NEXT: %[[SIZES_ALLOC:.*]] = memref.alloca() : memref<1xindex>
+// RUNTIME-NEXT: %[[SIZES:.*]] = memref.cast %[[SIZES_ALLOC]] : memref<1xindex> to memref<?xindex>
+// RUNTIME-NEXT: memref.store %[[DIM]], %[[SIZES_ALLOC]][%[[C0]]] : memref<1xindex>
+// RUNTIME-NEXT: %[[MAP_ALLOC:.*]] = memref.alloca() : memref<1xindex>
+// RUNTIME-NEXT: %[[MAP:.*]] = memref.cast %[[MAP_ALLOC]] : memref<1xindex> to memref<?xindex>
+// RUNTIME-NEXT: memref.store %[[C0]], %[[MAP_ALLOC]][%[[C0]]] : memref<1xindex>
+// RUNTIME-NEXT: %[[TENSOR:.*]] = call @newSparseTensor(%[[SIZES]], %[[SIZES]], %[[LVL_TYPES]], %[[MAP]], %[[MAP]], %[[C0_I32]], %[[C0_I32]], %[[C1_I32]], %[[C0_I32]], %[[NULL]]) : (memref<?xindex>, memref<?xindex>, memref<?xi64>, memref<?xindex>, memref<?xindex>, i32, i32, i32, i32, !llvm.ptr) -> !llvm.ptr
+// RUNTIME-NEXT: call @endLexInsert(%[[TENSOR]]) : (!llvm.ptr) -> ()
+// RUNTIME-NEXT: return %[[TENSOR]] : !llvm.ptr
+// RUNTIME-NEXT: }
+
+func.func @empty(%size: index) -> tensor<?xf64, #SV> {
+ %A = tensor.empty(%size) : tensor<?xf64, #SV>
+ return %A : tensor<?xf64, #SV>
+}
+
+// CODEGEN-LABEL: func.func @dealloc(
+// CODEGEN-SAME: %{{[^:]+}}: memref<?xindex>,
+// CODEGEN-SAME: %{{[^:]+}}: memref<?xindex>,
+// CODEGEN-SAME: %{{[^:]+}}: memref<?xf64>,
+// CODEGEN-SAME: %{{[^:]+}}: !sparse_tensor.storage_specifier<[[$SV]]>) {
+// CODEGEN-NEXT: return
+// CODEGEN-NEXT: }
+
+// RUNTIME-LABEL: func.func @dealloc(
+// RUNTIME-SAME: %[[TENSOR:.*]]: !llvm.ptr) {
+// RUNTIME-NEXT: call @delSparseTensor(%[[TENSOR]]) : (!llvm.ptr) -> ()
+// RUNTIME-NEXT: return
+// RUNTIME-NEXT: }
+
+func.func @dealloc(%A: tensor<?xf64, #SV>) {
+ bufferization.dealloc_tensor %A : tensor<?xf64, #SV>
+ return
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/233217
More information about the Mlir-commits
mailing list