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

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.

Reviewed By: Peiming

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

diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion_sparse2sparse.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion_sparse2sparse.mlir
index 11bd380..94993f2 100644 (file)
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion_sparse2sparse.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion_sparse2sparse.mlir
@@ -35,7 +35,7 @@
 
 module {
   //
-  // Utilities for output and releasing memory.
+  // Utility for output.
   //
   func.func @dump(%arg0: tensor<2x3x4xf64>) {
     %c0 = arith.constant 0 : index
@@ -44,10 +44,6 @@ module {
     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).
@@ -93,10 +89,10 @@ module {
     //
     // 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.
@@ -160,12 +156,12 @@ module {
     //
     // 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.
@@ -176,6 +172,9 @@ module {
     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
   }