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

Author: matthias-springer

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());

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();