[mlir][sparse] run canonicalization pass after DenseOpBufferize.

As pointed out by Matthias: "DenseBufferizationPass should be run right after TensorCopyInsertionPass. (Running it after bufferizing the sparse IR is also OK.) The reason for this is that whether copies are needed for not depends on the structure of the program (SSA use-def chains). In particular, running the canonicalizer in-between is problematic because it could introduce new RaW conflicts"

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D136980

mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp

@@ -66,9 +66,9 @@ void mlir::sparse_tensor::buildSparseCompiler(
   else
     pm.addPass(createSparseTensorCodegenPass());
   pm.addPass(createSparseBufferRewritePass());
-  pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
   pm.addPass(createDenseBufferizationPass(
       getBufferizationOptions(/*analysisOnly=*/false)));
+  pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
   pm.addNestedPass<func::FuncOp>(
       mlir::bufferization::createFinalizingBufferizePass());
   // TODO(springerm): Add sparse support to the BufferDeallocation pass and add

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matvec.mlir

@@ -91,10 +91,6 @@ module {
     // Release the resources.
     bufferization.dealloc_tensor %a : tensor<?x?xi32, #SparseMatrix>
 
-    // TODO(springerm): auto release!
-    bufferization.dealloc_tensor %b : tensor<?xi32>
-    bufferization.dealloc_tensor %x : tensor<?xi32>
-
     return
   }
 }

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_mttkrp.mlir

@@ -115,11 +115,6 @@ module {
     // Release the resources.
     bufferization.dealloc_tensor %b : tensor<?x?x?xf64, #SparseTensor>
 
-    // TODO(springerm): auto release!
-    bufferization.dealloc_tensor %a : tensor<?x?xf64>
-    bufferization.dealloc_tensor %c : tensor<?x?xf64>
-    bufferization.dealloc_tensor %d : tensor<?x?xf64>
-
     return
   }
 }

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_matmul.mlir

@@ -108,11 +108,6 @@ module {
     // Release the resources.
     bufferization.dealloc_tensor %s : tensor<?x?xf32, #SparseMatrix>
 
-    // TODO(springerm): auto release!
-    bufferization.dealloc_tensor %x : tensor<?x?xf32>
-    bufferization.dealloc_tensor %a : tensor<?x?xf32>
-    bufferization.dealloc_tensor %b : tensor<?x?xf32>
-
     return
   }
 }

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_spmm.mlir

@@ -90,9 +90,6 @@ module {
     // Release the resources.
     bufferization.dealloc_tensor %a : tensor<?x?xf64, #SparseMatrix>
 
-    // TODO(springerm): auto release!
-    bufferization.dealloc_tensor %b : tensor<?x?xf64>
-    bufferization.dealloc_tensor %x : tensor<?x?xf64>
 
     return
   }