[Mlir-commits] [mlir] 48a73bc - [mlir][sparse] Extract `StorageSpecifierToLLVMPass` from bufferization pipeline (#68635)

[llvmlistbot at llvm.org]

Author: Matthias Springer
Date: 2023-10-09T18:01:50-07:00
New Revision: 48a73bc455c919acc57b28f00dbeae5548877a72

URL: https://github.com/llvm/llvm-project/commit/48a73bc455c919acc57b28f00dbeae5548877a72
DIFF: https://github.com/llvm/llvm-project/commit/48a73bc455c919acc57b28f00dbeae5548877a72.diff

LOG: [mlir][sparse] Extract `StorageSpecifierToLLVMPass` from bufferization pipeline (#68635)

`StorageSpecifierToLLVMPass` does not have to be part of the
bufferization mini pipeline. It can run after the bufferization
pipeline. This is desirable because it keeps the bufferization pipeline
smaller.

Also fix incorrect bufferization API usage: `bufferizeOp` instead of
`bufferizeModuleOp` was used, even though function boundaries were
bufferized.

Added: 
    

Modified: 
    mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
    mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp b/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
index 54069064839eac3..7569413546c0a6e 100644
--- a/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
@@ -42,6 +42,8 @@ void mlir::sparse_tensor::buildSparseCompiler(
       /*enableSIMDIndex32=*/options.force32BitVectorIndices));
   if (options.testBufferizationAnalysisOnly)
     return;
+
+  pm.addPass(createStorageSpecifierToLLVMPass());
   pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
   pm.addNestedPass<func::FuncOp>(
       mlir::bufferization::createFinalizingBufferizePass());

diff  --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
index 6fca8f82e356626..480e18e257277de 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
@@ -39,7 +39,7 @@ namespace sparse_tensor {
 /// Return `true` if one of the given types is a sparse tensor type.
 static bool containsSparseTensor(TypeRange types) {
   for (Type t : types)
-    if (getSparseTensorEncoding(t))
+    if (isa<TensorType>(t) && getSparseTensorEncoding(t))
       return true;
   return false;
 }
@@ -97,7 +97,8 @@ class SparsificationAndBufferizationPass
       return false;
     });
 
-    if (failed(bufferization::bufferizeOp(getOperation(), updatedOptions)))
+    if (failed(bufferization::bufferizeModuleOp(cast<ModuleOp>(getOperation()),
+                                                updatedOptions)))
       return failure();
 
     bufferization::removeBufferizationAttributesInModule(getOperation());
@@ -154,7 +155,6 @@ class SparsificationAndBufferizationPass
         pm.addPass(createSparseTensorCodegenPass(createSparseDeallocs,
                                                  enableBufferInitialization));
         pm.addPass(createSparseBufferRewritePass(enableBufferInitialization));
-        pm.addPass(createStorageSpecifierToLLVMPass());
       }
       if (failed(runPipeline(pm, getOperation())))
         return signalPassFailure();