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

[Peiming]

Summary

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"

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

void mlir::sparse_tensor::buildSparseCompiler(
   pm.addPass(createSparseTensorRewritePass(options.enableRuntimeLibrary));
   pm.addPass(createSparsificationPass(options.sparsificationOptions()));
   if (options.enableRuntimeLibrary)
     pm.addPass(createSparseTensorConversionPass(
         options.sparseTensorConversionOptions()));
   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
   // it to this pipeline.
   pm.addNestedPass<func::FuncOp>(createConvertLinalgToLoopsPass());
   pm.addNestedPass<func::FuncOp>(createConvertVectorToSCFPass());
   pm.addNestedPass<func::FuncOp>(createConvertSCFToCFPass());
   pm.addPass(createLowerAffinePass());

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

func.func @entry() {
     // CHECK: ( 889, 1514, -21, -3431 )
     //
     %v = vector.transfer_read %0[%c0], %i0: tensor<?xi32>, vector<4xi32>
     vector.print %v : vector<4xi32>
 
     // 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

func.func @entry() {
     //
     %v = vector.transfer_read %0[%cst0, %cst0], %f0
           : tensor<?x?xf64>, vector<2x5xf64>
     vector.print %v : vector<2x5xf64>
 
     // 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

func.func @entry() {
     scf.for %i = %c0 to %c5 step %c1 {
       %v = vector.transfer_read %0[%i, %c0], %d0: tensor<?x?xf32>, vector<5xf32>
       vector.print %v : vector<5xf32>
     }
 
     // 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

func.func @entry() {
     // CHECK: ( ( 3548, 3550, 3552, 3554 ), ( 6052, 6053, 6054, 6055 ), ( -56, -63, -70, -77 ), ( -13704, -13709, -13714, -13719 ) )
     //
     %v = vector.transfer_read %0[%c0, %c0], %i0: tensor<?x?xf64>, vector<4x4xf64>
     vector.print %v : vector<4x4xf64>
 
     // 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
   }
 }