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

Matthias Springer llvmlistbot at llvm.org
Thu Feb 15 03:21:34 PST 2024


https://github.com/matthias-springer created https://github.com/llvm/llvm-project/pull/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).

>From 93c9fb0d4e0a4590d7e5737483aaaca2488d0053 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm at google.com>
Date: Thu, 15 Feb 2024 11:17:46 +0000
Subject: [PATCH] [mlir][sparse] Fix memory leaks (part 1)

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).
---
 .../Dialect/SparseTensor/CPU/concatenate_dim_0.mlir  |  1 +
 .../Dialect/SparseTensor/CPU/sparse_cast.mlir        | 10 ++++++++++
 .../Dialect/SparseTensor/CPU/sparse_cmp.mlir         |  4 +++-
 .../SparseTensor/CPU/sparse_conv_1d_nwc_wcf.mlir     |  3 +++
 .../Dialect/SparseTensor/CPU/sparse_conv_2d_55.mlir  |  8 +++++++-
 .../Dialect/SparseTensor/CPU/sparse_coo_test.mlir    |  1 +
 .../Dialect/SparseTensor/CPU/sparse_expand.mlir      |  1 +
 .../Dialect/SparseTensor/CPU/sparse_index_dense.mlir |  9 +++++++++
 .../Dialect/SparseTensor/CPU/sparse_matrix_ops.mlir  |  1 +
 .../Dialect/SparseTensor/CPU/sparse_pack_d.mlir      |  3 +++
 .../Dialect/SparseTensor/CPU/sparse_permute.mlir     | 12 ++++++++++++
 .../SparseTensor/CPU/sparse_pooling_nhwc.mlir        |  1 +
 .../SparseTensor/CPU/sparse_reduce_custom.mlir       |  2 ++
 .../SparseTensor/CPU/sparse_reduce_custom_prod.mlir  |  2 ++
 .../SparseTensor/CPU/sparse_reductions_prod.mlir     | 12 ++++++++++++
 .../Dialect/SparseTensor/CPU/sparse_scf_nested.mlir  |  1 +
 .../SparseTensor/CPU/sparse_semiring_select.mlir     |  1 +
 .../Dialect/SparseTensor/CPU/sparse_unary.mlir       |  3 +++
 .../Dialect/SparseTensor/CPU/sparse_vector_ops.mlir  |  3 +++
 19 files changed, 76 insertions(+), 2 deletions(-)

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_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_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_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_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_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_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_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_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_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_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