[Mlir-commits] [mlir] b6c453c - [mlir][sparse] Fix memory leaks (part 1) (#81843)

llvmlistbot at llvm.org llvmlistbot at llvm.org
Fri Feb 16 00:57:08 PST 2024


Author: Matthias Springer
Date: 2024-02-16T09:57:04+01:00
New Revision: b6c453c13f9dca064ac06b8b2cff80946225aa47

URL: https://github.com/llvm/llvm-project/commit/b6c453c13f9dca064ac06b8b2cff80946225aa47
DIFF: https://github.com/llvm/llvm-project/commit/b6c453c13f9dca064ac06b8b2cff80946225aa47.diff

LOG: [mlir][sparse] Fix memory leaks (part 1) (#81843)

This commit fixes memory leaks in sparse tensor integration tests by
adding `bufferization.dealloc_tensor` ops.

Note: Buffer deallocation will be automated in the future with the
ownership-based buffer deallocation pass, making `dealloc_tensor`
obsolete (only codegen path, not when using the runtime library).

Added: 
    

Modified: 
    mlir/test/Integration/Dialect/SparseTensor/CPU/concatenate_dim_0.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_cast.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_cmp.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_1d_nwc_wcf.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_2d_55.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion_block.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_coo_test.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_dot.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_expand.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_filter_conv2d.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_flatten.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_generate.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_index_dense.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matrix_ops.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matvec.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_mttkrp.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_reduction.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pack_d.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_permute.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pooling_nhwc.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_quantized_matmul.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reduce_custom.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reduce_custom_prod.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reduce_custom_sum.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions_min.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions_prod.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_matmul.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_mm_fusion.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_scf_nested.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_select.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_semiring_select.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_spmm.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_bf16.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_c32.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_f16.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_tanh.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_tensor_mul.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_unary.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_vector_ops.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/concatenate_dim_0.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/concatenate_dim_0.mlir
index 8894b385d3cd93..515f5b88b480d2 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/concatenate_dim_0.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/concatenate_dim_0.mlir
@@ -96,6 +96,7 @@ module {
     %1 = sparse_tensor.values %A : tensor<9x4xf64, #MAT_C_C> to memref<?xf64>
     call @printMemref1dF64(%1) : (memref<?xf64>) -> ()
 
+    bufferization.dealloc_tensor %c : tensor<9x4xf64>
     return
   }
 

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_cast.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_cast.mlir
index 1d546f90ca4fa1..6efe7b334b9845 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_cast.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_cast.mlir
@@ -279,6 +279,16 @@ module {
     bufferization.dealloc_tensor %3 : tensor<10xf32, #SV>
     bufferization.dealloc_tensor %5 : tensor<10xf64, #SV>
     bufferization.dealloc_tensor %7 : tensor<10xf64, #SV>
+    bufferization.dealloc_tensor %c0 : tensor<10xf32>
+    bufferization.dealloc_tensor %c1 : tensor<10xf32>
+    bufferization.dealloc_tensor %c2 : tensor<10xi32>
+    bufferization.dealloc_tensor %c3 : tensor<10xi32>
+    bufferization.dealloc_tensor %c4 : tensor<10xf64>
+    bufferization.dealloc_tensor %c5 : tensor<10xf32>
+    bufferization.dealloc_tensor %c6 : tensor<10xi64>
+    bufferization.dealloc_tensor %c7 : tensor<10xi64>
+    bufferization.dealloc_tensor %c8 : tensor<10xi8>
+    bufferization.dealloc_tensor %c9 : tensor<10xi32>
 
     return
   }

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_cmp.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_cmp.mlir
index 772daf646c772e..035db33fb4b318 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_cmp.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_cmp.mlir
@@ -146,9 +146,11 @@ module {
 
     bufferization.dealloc_tensor %lhs_sp : tensor<4x4xf64, #DCSR>
     bufferization.dealloc_tensor %rhs_sp : tensor<4x4xf64, #DCSR>
+    bufferization.dealloc_tensor %all_dn_out : tensor<4x4xi8>
     bufferization.dealloc_tensor %lhs_sp_out : tensor<4x4xi8, #DCSR>
     bufferization.dealloc_tensor %all_sp_out : tensor<4x4xi8, #DCSR>
-
+    bufferization.dealloc_tensor %lhs_sp_ret : tensor<4x4xi8>
+    bufferization.dealloc_tensor %rhs_sp_ret : tensor<4x4xi8>
     return
   }
 }

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_1d_nwc_wcf.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_1d_nwc_wcf.mlir
index 9295ee8b469c2f..16a67a1458369c 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_1d_nwc_wcf.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_1d_nwc_wcf.mlir
@@ -140,5 +140,8 @@ func.func @entry() {
   bufferization.dealloc_tensor %CCC_ret : tensor<?x?x?xf32, #CCC>
   bufferization.dealloc_tensor %CDC_ret : tensor<?x?x?xf32, #CDC>
 
+  bufferization.dealloc_tensor %1 : tensor<?x?x?xf32>
+  bufferization.dealloc_tensor %2 : tensor<?x?x?xf32>
+
   return
 }

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_2d_55.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_2d_55.mlir
index 9e7bbe71038b11..a7d7d1c5ed3c31 100755
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_2d_55.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_2d_55.mlir
@@ -207,13 +207,19 @@ module {
     %v5 = vector.transfer_read %5[%c0, %c0], %i0 : tensor<6x6xi32>, vector<6x6xi32>
     vector.print %v5 : vector<6x6xi32>
 
-    // Release sparse resources.
+    // Release resources.
     bufferization.dealloc_tensor %input_CSR : tensor<10x10xi32, #CSR>
     bufferization.dealloc_tensor %input_DCSR : tensor<10x10xi32, #DCSR>
     bufferization.dealloc_tensor %filter_CSR : tensor<5x5xi32, #CSR>
     bufferization.dealloc_tensor %sinput_CSR : tensor<10x10xi32, #CSR>
     bufferization.dealloc_tensor %sinput_DCSR : tensor<10x10xi32, #DCSR>
     bufferization.dealloc_tensor %sfilter_CSR : tensor<5x5xi32, #CSR>
+    bufferization.dealloc_tensor %0 : tensor<6x6xi32>
+    bufferization.dealloc_tensor %1 : tensor<6x6xi32>
+    bufferization.dealloc_tensor %2 : tensor<6x6xi32>
+    bufferization.dealloc_tensor %3 : tensor<6x6xi32>
+    bufferization.dealloc_tensor %4 : tensor<6x6xi32>
+    bufferization.dealloc_tensor %5 : tensor<6x6xi32>
 
     return
   }

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion_block.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion_block.mlir
index 6e396d9a84fade..809414ba977d2b 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion_block.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion_block.mlir
@@ -103,6 +103,12 @@ module {
     call @dumpf64(%v3) : (memref<?xf64>) -> ()
     call @dumpf64(%v4) : (memref<?xf64>) -> ()
 
+    // TODO: Fix memory leaks.
+    bufferization.dealloc_tensor %1 : tensor<2x4xf64, #BSR>
+    bufferization.dealloc_tensor %2 : tensor<2x4xf64, #BSR>
+    bufferization.dealloc_tensor %3 : tensor<2x4xf64, #CSR>
+    bufferization.dealloc_tensor %4 : tensor<2x4xf64, #CSC>
+
     return
   }
 }

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_coo_test.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_coo_test.mlir
index 3091bd03a009b0..aaf15ecc681fc2 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_coo_test.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_coo_test.mlir
@@ -201,6 +201,7 @@ module {
     bufferization.dealloc_tensor %C1 : tensor<8x8xf32>
     bufferization.dealloc_tensor %C2 : tensor<8x8xf32>
     bufferization.dealloc_tensor %C3 : tensor<8x8xf32>
+    bufferization.dealloc_tensor %C4 : tensor<8x8xf32>
     bufferization.dealloc_tensor %CSR_A : tensor<8x8xf32, #CSR>
     bufferization.dealloc_tensor %COO_A : tensor<8x8xf32, #SortedCOO>
     bufferization.dealloc_tensor %COO_B : tensor<8x8xf32, #SortedCOO>

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_dot.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_dot.mlir
index 101cb87655ec2b..f7ba4daa245892 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_dot.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_dot.mlir
@@ -84,6 +84,7 @@ module {
     vector.print %noe2 : index
 
     // Release the resources.
+    bufferization.dealloc_tensor %0 : tensor<f32>
     bufferization.dealloc_tensor %s1 : tensor<1024xf32, #SparseVector>
     bufferization.dealloc_tensor %s2 : tensor<1024xf32, #SparseVector>
 

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_expand.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_expand.mlir
index 21b2c040bf6800..93a7ea51ec9c4d 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_expand.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_expand.mlir
@@ -100,6 +100,7 @@ module {
     bufferization.dealloc_tensor %x1 : tensor<8x2xf64, #CSC>
     bufferization.dealloc_tensor %x2 : tensor<2x4xf64, #CSC>
     bufferization.dealloc_tensor %x3 : tensor<8x4xf64, #CSC>
+    bufferization.dealloc_tensor %xc : tensor<8x4xf64>
 
     return
   }

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_filter_conv2d.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_filter_conv2d.mlir
index 5eadc9008c7d71..7478b604ff6715 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_filter_conv2d.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_filter_conv2d.mlir
@@ -116,7 +116,10 @@ module {
 
     // Release the resources.
     bufferization.dealloc_tensor %sparse_filter : tensor<3x3xi32, #DCSR>
+    bufferization.dealloc_tensor %0 : tensor<6x6xi32>
     bufferization.dealloc_tensor %1 : tensor<6x6xi32, #DCSR>
+    bufferization.dealloc_tensor %sparse_ret : tensor<6x6xi32>
+
     return
   }
 }

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_flatten.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_flatten.mlir
index bc677d9a434a2e..55ecac8dbc1874 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_flatten.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_flatten.mlir
@@ -115,6 +115,7 @@ module {
 
     // Release the resources.
     bufferization.dealloc_tensor %a : tensor<7x3x3x3x3x3x5x3xf64, #SparseTensor>
+    bufferization.dealloc_tensor %0 : tensor<7x3xf64>
 
     return
   }

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_generate.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_generate.mlir
index fcc16f5e9cb4f6..c1547033062d37 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_generate.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_generate.mlir
@@ -87,6 +87,7 @@ module {
 
     // Release the resources.
     bufferization.dealloc_tensor %sv : tensor<?xf64, #SparseVector>
+    bufferization.dealloc_tensor %empty : tensor<?xf64>
     memref.dealloc %array : memref<?xi64>
     func.call @rtdrand(%g) : (!Generator) -> ()
 

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_index_dense.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_index_dense.mlir
index 730836c356f3cc..4bb1f3d12871f3 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_index_dense.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_index_dense.mlir
@@ -161,6 +161,7 @@ module {
     %dm = sparse_tensor.convert %m2 : tensor<3x4xi64> to tensor<3x4xi64, #SparseMatrix>
 
     // Setup out tensors.
+    // Note: Constants bufferize to read-only buffers.
     %init_8 = tensor.empty() : tensor<8xi64>
     %init_3_4 = tensor.empty() : tensor<3x4xi64>
 
@@ -208,6 +209,14 @@ module {
     bufferization.dealloc_tensor %dv : tensor<8xi64, #SparseVector>
     bufferization.dealloc_tensor %sm : tensor<3x4xi64, #SparseMatrix>
     bufferization.dealloc_tensor %dm : tensor<3x4xi64, #SparseMatrix>
+    bufferization.dealloc_tensor %0 : tensor<8xi64>
+    bufferization.dealloc_tensor %1 : tensor<8xi64>
+    bufferization.dealloc_tensor %2 : tensor<8xi64>
+    bufferization.dealloc_tensor %3 : tensor<8xi64>
+    bufferization.dealloc_tensor %4 : tensor<3x4xi64>
+    bufferization.dealloc_tensor %5 : tensor<3x4xi64>
+    bufferization.dealloc_tensor %6 : tensor<3x4xi64>
+    bufferization.dealloc_tensor %7 : tensor<3x4xi64>
 
     return
   }

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul.mlir
index 5328fe8f42c914..fa0dbac269b926 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul.mlir
@@ -316,10 +316,19 @@ module {
     bufferization.dealloc_tensor %b2 : tensor<8x4xf64, #DCSR>
     bufferization.dealloc_tensor %b3 : tensor<8x4xf64, #CSR>
     bufferization.dealloc_tensor %b4 : tensor<8x4xf64, #DCSR>
+    bufferization.dealloc_tensor %c1 : tensor<4x4xf64>
+    bufferization.dealloc_tensor %c2 : tensor<4x4xf64>
+    bufferization.dealloc_tensor %c4 : tensor<4x4xf64>
+    bufferization.dealloc_tensor %c5 : tensor<4x4xf64>
+    bufferization.dealloc_tensor %c7 : tensor<4x4xf64>
+    bufferization.dealloc_tensor %c8 : tensor<4x4xf64>
+    bufferization.dealloc_tensor %0 : tensor<4x4xf64>
     bufferization.dealloc_tensor %1 : tensor<4x4xf64, #CSR>
     bufferization.dealloc_tensor %2 : tensor<4x4xf64, #DCSR>
+    bufferization.dealloc_tensor %3 : tensor<4x4xf64>
     bufferization.dealloc_tensor %4 : tensor<4x4xf64, #CSR>
     bufferization.dealloc_tensor %5 : tensor<4x4xf64, #DCSR>
+    bufferization.dealloc_tensor %6 : tensor<4x4xf64>
     bufferization.dealloc_tensor %7 : tensor<4x4xf64, #CSR>
     bufferization.dealloc_tensor %8 : tensor<4x4xf64, #DCSR>
 

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matrix_ops.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matrix_ops.mlir
index 2ce1f2ee1b8280..2cecc242034389 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matrix_ops.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matrix_ops.mlir
@@ -134,6 +134,7 @@ module {
     %dm = sparse_tensor.convert %arg0 : tensor<?x?xf64, #DCSR> to tensor<?x?xf64>
     %u = tensor.cast %dm : tensor<?x?xf64> to tensor<*xf64>
     call @printMemrefF64(%u) : (tensor<*xf64>) -> ()
+    bufferization.dealloc_tensor %dm : tensor<?x?xf64>
     return
   }
 

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matvec.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matvec.mlir
index 940c98cdda34e4..bdad713709afa2 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matvec.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matvec.mlir
@@ -125,6 +125,8 @@ module {
 
     // Release the resources.
     bufferization.dealloc_tensor %a : tensor<?x?xi32, #SparseMatrix>
+    bufferization.dealloc_tensor %b : tensor<?xi32>
+    bufferization.dealloc_tensor %0 : tensor<?xi32>
 
     return
   }

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_mttkrp.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_mttkrp.mlir
index 3b6cda4d16316e..30e620b9d610ba 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_mttkrp.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_mttkrp.mlir
@@ -142,6 +142,9 @@ module {
 
     // Release the resources.
     bufferization.dealloc_tensor %b : tensor<?x?x?xf64, #SparseTensor>
+    bufferization.dealloc_tensor %c : tensor<?x?xf64>
+    bufferization.dealloc_tensor %d : tensor<?x?xf64>
+    bufferization.dealloc_tensor %a : tensor<?x?xf64>
 
     return
   }

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_reduction.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_reduction.mlir
index 5d88bc87d1ce65..17def3f52c003e 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_reduction.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_reduction.mlir
@@ -110,6 +110,8 @@ module {
     bufferization.dealloc_tensor %st1 : tensor<?x?x?xi32, #SparseTensor>
     bufferization.dealloc_tensor %st2 : tensor<?x?x?xi32, #SparseTensor>
     bufferization.dealloc_tensor %0 : tensor<?x?xi32, #SparseMatrix>
+    bufferization.dealloc_tensor %dm : tensor<?x?xi32>
+
     return
   }
 }

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pack_d.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pack_d.mlir
index c818c23bfa0f46..da816c7fbb1172 100755
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pack_d.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pack_d.mlir
@@ -125,6 +125,9 @@ module {
     %v2 = vector.transfer_read %d1[%c0, %c0, %c0], %f0 : tensor<4x3x2xf32>, vector<4x3x2xf32>
     vector.print %v2 : vector<4x3x2xf32>
 
+    bufferization.dealloc_tensor %d0 : tensor<4x3x2xf32>
+    bufferization.dealloc_tensor %d1 : tensor<4x3x2xf32>
+    bufferization.dealloc_tensor %d2 : tensor<4x3x2xf32>
     // FIXME: doing this explicitly crashes runtime
     // bufferization.dealloc_tensor %s0 : tensor<4x3x2xf32, #CCC>
     // bufferization.dealloc_tensor %s1 : tensor<4x3x2xf32, #BatchedCSR>

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_permute.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_permute.mlir
index 43d49b79a5146a..664a86c7ad58fd 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_permute.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_permute.mlir
@@ -99,7 +99,9 @@ module {
     call @dump(%0) : (tensor<2x3x4xf64>) -> ()
     %1 = call @convert1(%b) : (tensor<2x3x4xf64, #S1>) -> tensor<2x3x4xf64>
     call @dump(%1) : (tensor<2x3x4xf64>) -> ()
+    bufferization.dealloc_tensor %0 : tensor<2x3x4xf64>
     bufferization.dealloc_tensor %b : tensor<2x3x4xf64, #S1>
+    bufferization.dealloc_tensor %1 : tensor<2x3x4xf64>
     return
   }
 
@@ -127,7 +129,9 @@ module {
     call @dump(%0) : (tensor<2x3x4xf64>) -> ()
     %2 = call @convert2(%b) : (tensor<2x3x4xf64, #S2>) -> tensor<2x3x4xf64>
     call @dump(%2) : (tensor<2x3x4xf64>) -> ()
+    bufferization.dealloc_tensor %0 : tensor<2x3x4xf64>
     bufferization.dealloc_tensor %b : tensor<2x3x4xf64, #S2>
+    bufferization.dealloc_tensor %2 : tensor<2x3x4xf64>
     return
   }
 
@@ -155,7 +159,9 @@ module {
     call @dump(%0) : (tensor<2x3x4xf64>) -> ()
     %3 = call @convert3(%b) : (tensor<2x3x4xf64, #S3>) -> tensor<2x3x4xf64>
     call @dump(%3) : (tensor<2x3x4xf64>) -> ()
+    bufferization.dealloc_tensor %0 : tensor<2x3x4xf64>
     bufferization.dealloc_tensor %b : tensor<2x3x4xf64, #S3>
+    bufferization.dealloc_tensor %3 : tensor<2x3x4xf64>
     return
   }
 
@@ -183,7 +189,9 @@ module {
     call @dump(%0) : (tensor<2x3x4xf64>) -> ()
     %4 = call @convert4(%b) : (tensor<2x3x4xf64, #S4>) -> tensor<2x3x4xf64>
     call @dump(%4) : (tensor<2x3x4xf64>) -> ()
+    bufferization.dealloc_tensor %0 : tensor<2x3x4xf64>
     bufferization.dealloc_tensor %b : tensor<2x3x4xf64, #S4>
+    bufferization.dealloc_tensor %4 : tensor<2x3x4xf64>
     return
   }
 
@@ -211,7 +219,9 @@ module {
     call @dump(%0) : (tensor<2x3x4xf64>) -> ()
     %5 = call @convert5(%b) : (tensor<2x3x4xf64, #S5>) -> tensor<2x3x4xf64>
     call @dump(%5) : (tensor<2x3x4xf64>) -> ()
+    bufferization.dealloc_tensor %0 : tensor<2x3x4xf64>
     bufferization.dealloc_tensor %b : tensor<2x3x4xf64, #S5>
+    bufferization.dealloc_tensor %5 : tensor<2x3x4xf64>
     return
   }
 
@@ -239,7 +249,9 @@ module {
     call @dump(%0) : (tensor<2x3x4xf64>) -> ()
     %6 = call @convert6(%b) : (tensor<2x3x4xf64, #S6>) -> tensor<2x3x4xf64>
     call @dump(%6) : (tensor<2x3x4xf64>) -> ()
+    bufferization.dealloc_tensor %0 : tensor<2x3x4xf64>
     bufferization.dealloc_tensor %b : tensor<2x3x4xf64, #S6>
+    bufferization.dealloc_tensor %6 : tensor<2x3x4xf64>
     return
   }
 

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pooling_nhwc.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pooling_nhwc.mlir
index aaa8fa1b0207c6..3ce089d7a7cf6b 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pooling_nhwc.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pooling_nhwc.mlir
@@ -87,5 +87,6 @@ func.func @entry() {
   bufferization.dealloc_tensor %in_CCCC : tensor<1x4x4x1xf32, #CCCC>
   bufferization.dealloc_tensor %CCCC_ret : tensor<1x3x3x1xf32, #CCCC>
   bufferization.dealloc_tensor %dense_ret : tensor<1x3x3x1xf32>
+  bufferization.dealloc_tensor %s1 : tensor<1x3x3x1xf32>
   return
 }

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_quantized_matmul.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_quantized_matmul.mlir
index 900262eca14bff..b322d965f2dc7e 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_quantized_matmul.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_quantized_matmul.mlir
@@ -91,6 +91,7 @@ module {
 
     // Release the resources.
     bufferization.dealloc_tensor %sparse_input2 : tensor<3x6xi8, #DCSR>
+    bufferization.dealloc_tensor %0 : tensor<5x6xi32>
 
     return
   }

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reduce_custom.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reduce_custom.mlir
index 3d0677c2a9cb80..17219cde035d34 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reduce_custom.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reduce_custom.mlir
@@ -126,6 +126,7 @@ module {
     %dv = sparse_tensor.convert %arg0 : tensor<?xf64, #SparseVector> to tensor<?xf64>
     %2 = vector.transfer_read %dv[%c0], %d0: tensor<?xf64>, vector<16xf64>
     vector.print %2 : vector<16xf64>
+    bufferization.dealloc_tensor %dv : tensor<?xf64>
     return
   }
 
@@ -140,6 +141,7 @@ module {
     %dm = sparse_tensor.convert %arg0 : tensor<?x?xf64, #CSR> to tensor<?x?xf64>
     %2 = vector.transfer_read %dm[%c0, %c0], %d0: tensor<?x?xf64>, vector<5x5xf64>
     vector.print %2 : vector<5x5xf64>
+    bufferization.dealloc_tensor %dm : tensor<?x?xf64>
     return
   }
 

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reduce_custom_prod.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reduce_custom_prod.mlir
index a90ea807082f1e..6e2c572cf21ba0 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reduce_custom_prod.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reduce_custom_prod.mlir
@@ -101,6 +101,7 @@ module {
     %dv = sparse_tensor.convert %arg0 : tensor<?xf64, #SparseVector> to tensor<?xf64>
     %2 = vector.transfer_read %dv[%c0], %d0: tensor<?xf64>, vector<16xf64>
     vector.print %2 : vector<16xf64>
+    bufferization.dealloc_tensor %dv : tensor<?xf64>
     return
   }
 
@@ -115,6 +116,7 @@ module {
     %dm = sparse_tensor.convert %arg0 : tensor<?x?xf64, #CSR> to tensor<?x?xf64>
     %2 = vector.transfer_read %dm[%c0, %c0], %d0: tensor<?x?xf64>, vector<5x5xf64>
     vector.print %2 : vector<5x5xf64>
+    bufferization.dealloc_tensor %dm : tensor<?x?xf64>
     return
   }
 

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reduce_custom_sum.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reduce_custom_sum.mlir
index b8ce269183b109..80c35676e804b1 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reduce_custom_sum.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reduce_custom_sum.mlir
@@ -176,6 +176,10 @@ module {
 
     // Release the resources.
     bufferization.dealloc_tensor %s0 : tensor<8xi32, #SV>
+    bufferization.dealloc_tensor %0 : tensor<i32>
+    bufferization.dealloc_tensor %1 : tensor<i32>
+    bufferization.dealloc_tensor %2 : tensor<i32>
+    bufferization.dealloc_tensor %3 : tensor<i32>
 
     return
   }

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions.mlir
index 86c9f09602462a..4ad23d1c031238 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions.mlir
@@ -153,6 +153,10 @@ module {
     // Release the resources.
     bufferization.dealloc_tensor %sparse_input_i32 : tensor<32xi32, #SV>
     bufferization.dealloc_tensor %sparse_input_f32 : tensor<32xf32, #SV>
+    bufferization.dealloc_tensor %0 : tensor<i32>
+    bufferization.dealloc_tensor %1 : tensor<f32>
+    bufferization.dealloc_tensor %2 : tensor<i32>
+    bufferization.dealloc_tensor %3 : tensor<i32>
 
     return
   }

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions_min.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions_min.mlir
index 4bde8a46e26c4c..198920024221b0 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions_min.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions_min.mlir
@@ -137,6 +137,12 @@ module {
     bufferization.dealloc_tensor %sv0 : tensor<32xi32, #SV>
     bufferization.dealloc_tensor %sv1 : tensor<32xi32, #SV>
     bufferization.dealloc_tensor %sv2 : tensor<32xi32, #SV>
+    bufferization.dealloc_tensor %0 : tensor<i32>
+    bufferization.dealloc_tensor %1 : tensor<i32>
+    bufferization.dealloc_tensor %2 : tensor<i32>
+    bufferization.dealloc_tensor %3 : tensor<i32>
+    bufferization.dealloc_tensor %4 : tensor<i32>
+    bufferization.dealloc_tensor %5 : tensor<i32>
 
     return
   }

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions_prod.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions_prod.mlir
index 858cdec69d8d9e..7ec30787fea464 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions_prod.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions_prod.mlir
@@ -163,6 +163,7 @@ module {
   }
 
   func.func @entry() {
+    // Note: Constants bufferize to read-only buffers.
     %ri = arith.constant dense< 7   > : tensor<i32>
     %rf = arith.constant dense< 2.0 > : tensor<f32>
 
@@ -269,6 +270,17 @@ module {
     bufferization.dealloc_tensor %s1_i32 : tensor<32xi32, #SV>
     bufferization.dealloc_tensor %s1_f32 : tensor<32xf32, #SV>
     bufferization.dealloc_tensor %s0     : tensor<32xi32, #SV>
+    bufferization.dealloc_tensor %0 : tensor<i32>
+    bufferization.dealloc_tensor %1 : tensor<f32>
+    bufferization.dealloc_tensor %2 : tensor<i32>
+    bufferization.dealloc_tensor %3 : tensor<f32>
+    bufferization.dealloc_tensor %4 : tensor<i32>
+    bufferization.dealloc_tensor %5 : tensor<f32>
+    bufferization.dealloc_tensor %6 : tensor<i32>
+    bufferization.dealloc_tensor %7 : tensor<f32>
+    bufferization.dealloc_tensor %8 : tensor<i32>
+    bufferization.dealloc_tensor %9 : tensor<i32>
+    bufferization.dealloc_tensor %10 : tensor<i32>
 
     return
   }

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_matmul.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_matmul.mlir
index 09f1f11567aeae..3330d2249707fc 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_matmul.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_matmul.mlir
@@ -134,6 +134,9 @@ module {
 
     // Release the resources.
     bufferization.dealloc_tensor %s : tensor<?x?xf32, #SparseMatrix>
+    bufferization.dealloc_tensor %0 : tensor<?x?xf32>
+    bufferization.dealloc_tensor %a : tensor<?x?xf32>
+    bufferization.dealloc_tensor %b : tensor<?x?xf32>
 
     return
   }

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_mm_fusion.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_mm_fusion.mlir
index 1204489ab0b2a8..afaf36c4072c09 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_mm_fusion.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_mm_fusion.mlir
@@ -111,6 +111,7 @@ module {
           %r = arith.mulf %t, %s : f64
           linalg.yield %r : f64
     } -> tensor<8x8xf64>
+    bufferization.dealloc_tensor %2 : tensor<8x8xf64>
     return %3 : tensor<8x8xf64>
   }
 

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_scf_nested.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_scf_nested.mlir
index f945d3bd514d96..439144fedeeb89 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_scf_nested.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_scf_nested.mlir
@@ -74,6 +74,7 @@ module @func_sparse.2 {
     %dm = sparse_tensor.convert %arg0 : tensor<2x3x4xf64, #SparseMatrix> to tensor<2x3x4xf64>
     %0 = vector.transfer_read %dm[%c0, %c0, %c0], %d0: tensor<2x3x4xf64>, vector<2x3x4xf64>
     vector.print %0 : vector<2x3x4xf64>
+    bufferization.dealloc_tensor %dm : tensor<2x3x4xf64>
     return
   }
 

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_select.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_select.mlir
index f7f420c1e4c940..533afb6644aeda 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_select.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_select.mlir
@@ -109,6 +109,7 @@ module {
     %dv = sparse_tensor.convert %arg0 : tensor<?xf64, #SparseVector> to tensor<?xf64>
     %2 = vector.transfer_read %dv[%c0], %d0: tensor<?xf64>, vector<16xf64>
     vector.print %2 : vector<16xf64>
+    bufferization.dealloc_tensor %dv : tensor<?xf64>
     return
   }
 
@@ -123,6 +124,7 @@ module {
     %dm = sparse_tensor.convert %arg0 : tensor<?x?xf64, #CSR> to tensor<?x?xf64>
     %2 = vector.transfer_read %dm[%c0, %c0], %d0: tensor<?x?xf64>, vector<5x5xf64>
     vector.print %2 : vector<5x5xf64>
+    bufferization.dealloc_tensor %dm : tensor<?x?xf64>
     return
   }
 

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_semiring_select.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_semiring_select.mlir
index 8805cabe393e30..6244be0ba7ab64 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_semiring_select.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_semiring_select.mlir
@@ -99,6 +99,7 @@ module {
     bufferization.dealloc_tensor %sl: tensor<5x5xf64, #DCSR>
     bufferization.dealloc_tensor %sr: tensor<5x5xf64, #DCSR>
     bufferization.dealloc_tensor %1:  tensor<5x5xf64, #DCSR>
+    bufferization.dealloc_tensor %r : tensor<5x5xf64>
 
     return
   }

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_spmm.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_spmm.mlir
index 4eb250fbc7d382..573b1a2aac2598 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_spmm.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_spmm.mlir
@@ -115,7 +115,8 @@ module {
 
     // Release the resources.
     bufferization.dealloc_tensor %a : tensor<?x?xf64, #SparseMatrix>
-
+    bufferization.dealloc_tensor %b : tensor<?x?xf64>
+    bufferization.dealloc_tensor %0 : tensor<?x?xf64>
 
     return
   }

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum.mlir
index bc5b6e8899e7e5..e6cbff231024ed 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum.mlir
@@ -100,6 +100,7 @@ module {
 
     // Release the resources.
     bufferization.dealloc_tensor %a : tensor<?x?xf64, #SparseMatrix>
+    bufferization.dealloc_tensor %0 : tensor<f64>
 
     return
   }

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_bf16.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_bf16.mlir
index cc1027ac4ffe7e..ee00a19a412306 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_bf16.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_bf16.mlir
@@ -95,6 +95,7 @@ module {
 
     // Release the resources.
     bufferization.dealloc_tensor %a : tensor<?x?xbf16, #SparseMatrix>
+    bufferization.dealloc_tensor %0 : tensor<bf16>
 
     return
   }

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_c32.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_c32.mlir
index ed4fa6840f4788..5fdf636ef1230a 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_c32.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_c32.mlir
@@ -106,6 +106,7 @@ module {
     vector.print %imag : f64
 
     // Release the resources.
+    bufferization.dealloc_tensor %0 : tensor<complex<f64>>
     bufferization.dealloc_tensor %a : tensor<?x?xcomplex<f64>, #SparseMatrix>
 
     return

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_f16.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_f16.mlir
index 68cfe372946144..6a34695229495d 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_f16.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_f16.mlir
@@ -93,6 +93,7 @@ module {
     vector.print %vf : f32
 
     // Release the resources.
+    bufferization.dealloc_tensor %0 : tensor<f16>
     bufferization.dealloc_tensor %a : tensor<?x?xf16, #SparseMatrix>
 
     return

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_tanh.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_tanh.mlir
index fddfce78c5e6b6..336044d5660057 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_tanh.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_tanh.mlir
@@ -72,6 +72,7 @@ module {
         : tensor<?xf64, #SparseVector> to tensor<?xf64>
     %3 = vector.transfer_read %dv[%c0], %d0: tensor<?xf64>, vector<32xf64>
     vector.print %3 : vector<32xf64>
+    bufferization.dealloc_tensor %dv : tensor<?xf64>
     return
   }
 

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_tensor_mul.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_tensor_mul.mlir
index e5e1566abc59d2..d53b03025f5588 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_tensor_mul.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_tensor_mul.mlir
@@ -126,6 +126,8 @@ module {
     bufferization.dealloc_tensor %sta : tensor<?x?x?xf64, #ST>
     bufferization.dealloc_tensor %stb : tensor<?x?x?xf64, #ST>
     bufferization.dealloc_tensor %0  : tensor<?x?x?xf64, #ST>
+    bufferization.dealloc_tensor %dt : tensor<?x?x?xf64>
+
     return
   }
 }

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_unary.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_unary.mlir
index a811071b90d98b..e03f99253b7845 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_unary.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_unary.mlir
@@ -216,6 +216,7 @@ module {
     %dv = sparse_tensor.convert %arg0 : tensor<?xf64, #SparseVector> to tensor<?xf64>
     %3 = vector.transfer_read %dv[%c0], %d0: tensor<?xf64>, vector<32xf64>
     vector.print %3 : vector<32xf64>
+    bufferization.dealloc_tensor %dv : tensor<?xf64>
     return
   }
 
@@ -231,6 +232,7 @@ module {
     %dv = sparse_tensor.convert %arg0 : tensor<?xi32, #SparseVector> to tensor<?xi32>
     %3 = vector.transfer_read %dv[%c0], %d0: tensor<?xi32>, vector<32xi32>
     vector.print %3 : vector<32xi32>
+    bufferization.dealloc_tensor %dv : tensor<?xi32>
     return
   }
 
@@ -244,6 +246,7 @@ module {
     %dm = sparse_tensor.convert %arg0 : tensor<?x?xf64, #DCSR> to tensor<?x?xf64>
     %3 = vector.transfer_read %dm[%c0, %c0], %d0: tensor<?x?xf64>, vector<4x8xf64>
     vector.print %3 : vector<4x8xf64>
+    bufferization.dealloc_tensor %dm : tensor<?x?xf64>
     return
   }
 

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_vector_ops.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_vector_ops.mlir
index 1dfca563e89b94..d9ca2dca85342a 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_vector_ops.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_vector_ops.mlir
@@ -174,6 +174,7 @@ module {
     %dv = sparse_tensor.convert %arg0 : tensor<?xf64, #SparseVector> to tensor<?xf64>
     %2 = vector.transfer_read %dv[%c0], %d0: tensor<?xf64>, vector<32xf64>
     vector.print %2 : vector<32xf64>
+    bufferization.dealloc_tensor %dv : tensor<?xf64>
     return
   }
 
@@ -253,9 +254,11 @@ module {
     bufferization.dealloc_tensor %sv1_dup : tensor<?xf64, #SparseVector>
     bufferization.dealloc_tensor %sv2 : tensor<?xf64, #SparseVector>
     bufferization.dealloc_tensor %0 : tensor<?xf64, #SparseVector>
+    // Note: No dealloc for %1 because it was inplace!
     bufferization.dealloc_tensor %2 : tensor<?xf64, #SparseVector>
     bufferization.dealloc_tensor %3 : tensor<?xf64, #SparseVector>
     bufferization.dealloc_tensor %4 : tensor<?xf64, #DenseVector>
+    bufferization.dealloc_tensor %5 : tensor<f64>
     return
   }
 }


        


More information about the Mlir-commits mailing list