[mlir][sparse] fixed memory leak on sparse tensors

[aartbik]

Summary

This was introduced by https://reviews.llvm.org/D134933
This change also cleans up a dump&release method which
had become a misnomer after prior bufferization improvement.

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

 
 // This also checks the singleton->compressed conversion.
 #SingletonTensor3 = #sparse_tensor.encoding<{
   dimLevelType = [ "dense", "dense", "compressed" ]
 }>
 
 module {
   //
-  // Utilities for output and releasing memory.
+  // Utility for output.
   //
   func.func @dump(%arg0: tensor<2x3x4xf64>) {
     %c0 = arith.constant 0 : index
     %d0 = arith.constant -1.0 : f64
     %0 = vector.transfer_read %arg0[%c0, %c0, %c0], %d0: tensor<2x3x4xf64>, vector<2x3x4xf64>
     vector.print %0 : vector<2x3x4xf64>
     return
   }
-  func.func @dumpAndRelease_234(%arg0: tensor<2x3x4xf64>) {
-    call @dump(%arg0) : (tensor<2x3x4xf64>) -> ()
-    return
-  }
 
   //
   // The first test suite (for non-singleton DimLevelTypes).
   //
   func.func @testNonSingleton() {
     //
     // Initialize a 3-dim dense tensor.
     //
func.func @testNonSingleton() {
     %d23 = sparse_tensor.convert %t23 : tensor<2x3x4xf64, #Tensor3> to tensor<2x3x4xf64>
     %d31 = sparse_tensor.convert %t31 : tensor<2x3x4xf64, #Tensor1> to tensor<2x3x4xf64>
 
     //
     // Check round-trip equality.  And release dense tensors.
     //
     // CHECK-COUNT-5: ( ( ( 1, 2, 3, 4 ), ( 5, 6, 7, 8 ), ( 9, 10, 11, 12 ) ), ( ( 13, 14, 15, 16 ), ( 17, 18, 19, 20 ), ( 21, 22, 23, 24 ) ) )
     call @dump(%src) : (tensor<2x3x4xf64>) -> ()
-    call @dumpAndRelease_234(%d13) : (tensor<2x3x4xf64>) -> ()
-    call @dumpAndRelease_234(%d21) : (tensor<2x3x4xf64>) -> ()
-    call @dumpAndRelease_234(%d23) : (tensor<2x3x4xf64>) -> ()
-    call @dumpAndRelease_234(%d31) : (tensor<2x3x4xf64>) -> ()
+    call @dump(%d13) : (tensor<2x3x4xf64>) -> ()
+    call @dump(%d21) : (tensor<2x3x4xf64>) -> ()
+    call @dump(%d23) : (tensor<2x3x4xf64>) -> ()
+    call @dump(%d31) : (tensor<2x3x4xf64>) -> ()
 
     //
     // Release sparse tensors.
     //
     bufferization.dealloc_tensor %t13 : tensor<2x3x4xf64, #Tensor3>
     bufferization.dealloc_tensor %t21 : tensor<2x3x4xf64, #Tensor1>
     bufferization.dealloc_tensor %t23 : tensor<2x3x4xf64, #Tensor3>
     bufferization.dealloc_tensor %t31 : tensor<2x3x4xf64, #Tensor1>
func.func @testSingleton() {
     %d31 = sparse_tensor.convert %t31 : tensor<2x3x4xf64, #SingletonTensor1> to tensor<2x3x4xf64>
     %d32 = sparse_tensor.convert %t32 : tensor<2x3x4xf64, #SingletonTensor2> to tensor<2x3x4xf64>
 
     //
     // Check round-trip equality.  And release dense tensors.
     //
     // CHECK-COUNT-7: ( ( ( 1, 0, 0, 0 ), ( 0, 6, 0, 0 ), ( 0, 0, 11, 0 ) ), ( ( 0, 14, 0, 0 ), ( 0, 0, 0, 20 ), ( 21, 0, 0, 0 ) ) )
     call @dump(%src) : (tensor<2x3x4xf64>) -> ()
-    call @dumpAndRelease_234(%d12) : (tensor<2x3x4xf64>) -> ()
-    call @dumpAndRelease_234(%d13) : (tensor<2x3x4xf64>) -> ()
-    call @dumpAndRelease_234(%d21) : (tensor<2x3x4xf64>) -> ()
-    call @dumpAndRelease_234(%d23) : (tensor<2x3x4xf64>) -> ()
-    call @dumpAndRelease_234(%d31) : (tensor<2x3x4xf64>) -> ()
-    call @dumpAndRelease_234(%d32) : (tensor<2x3x4xf64>) -> ()
+    call @dump(%d12) : (tensor<2x3x4xf64>) -> ()
+    call @dump(%d13) : (tensor<2x3x4xf64>) -> ()
+    call @dump(%d21) : (tensor<2x3x4xf64>) -> ()
+    call @dump(%d23) : (tensor<2x3x4xf64>) -> ()
+    call @dump(%d31) : (tensor<2x3x4xf64>) -> ()
+    call @dump(%d32) : (tensor<2x3x4xf64>) -> ()
 
     //
     // Release sparse tensors.
     //
     bufferization.dealloc_tensor %t12 : tensor<2x3x4xf64, #SingletonTensor2>
     bufferization.dealloc_tensor %t13 : tensor<2x3x4xf64, #SingletonTensor3>
     bufferization.dealloc_tensor %t21 : tensor<2x3x4xf64, #SingletonTensor1>
     bufferization.dealloc_tensor %t23 : tensor<2x3x4xf64, #SingletonTensor3>
     bufferization.dealloc_tensor %t31 : tensor<2x3x4xf64, #SingletonTensor1>
     bufferization.dealloc_tensor %t32 : tensor<2x3x4xf64, #SingletonTensor2>
+    bufferization.dealloc_tensor %s1 : tensor<2x3x4xf64, #SingletonTensor1>
+    bufferization.dealloc_tensor %s2 : tensor<2x3x4xf64, #SingletonTensor2>
+    bufferization.dealloc_tensor %s3 : tensor<2x3x4xf64, #SingletonTensor3>
 
     return
   }
 
   //
   // Main driver.
   //
   func.func @entry() {
     call @testNonSingleton() : () -> ()
     call @testSingleton() : () -> ()
     return
   }
 }