[llvm-branch-commits] [mlir] [mlir][sparse] Fix memory leaks (part 3) (PR #85184)

[Matthias Springer via llvm-branch-commits]

https://github.com/matthias-springer created https://github.com/llvm/llvm-project/pull/85184

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 889cee2838737faf6e037b2d4ca3b785bb1be755 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm at google.com>
Date: Thu, 14 Mar 2024 06:05:58 +0000
Subject: [PATCH] [mlir][sparse] Fix memory leaks (part 3)

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).
---
 .../SparseTensor/CPU/dual_sparse_conv_2d.mlir       |  1 +
 .../Dialect/SparseTensor/CPU/sparse_conv_2d.mlir    |  3 +++
 .../CPU/sparse_conversion_sparse2dense.mlir         | 13 +++++++++++++
 .../CPU/sparse_conversion_sparse2sparse.mlir        |  8 ++++++--
 4 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/dual_sparse_conv_2d.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/dual_sparse_conv_2d.mlir
index 350b5b41dafc00..c645ca6567209c 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/dual_sparse_conv_2d.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/dual_sparse_conv_2d.mlir
@@ -222,6 +222,7 @@ module {
     bufferization.dealloc_tensor %sparse_filter_CD : tensor<3x3xi32, #CDR>
     bufferization.dealloc_tensor %sparse_filter_CSC : tensor<3x3xi32, #CSC>
 
+    bufferization.dealloc_tensor %0 : tensor<6x6xi32>
     bufferization.dealloc_tensor %2 : tensor<6x6xi32, #DCSR>
     bufferization.dealloc_tensor %3 : tensor<6x6xi32, #CSR>
     bufferization.dealloc_tensor %4 : tensor<6x6xi32, #CDR>
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_2d.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_2d.mlir
index f8fb8fdf53e356..55d4caeb7eb30e 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_2d.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_2d.mlir
@@ -273,11 +273,14 @@ module {
     bufferization.dealloc_tensor %sparse_input_CSC : tensor<8x8xi32, #CSC>
     bufferization.dealloc_tensor %sparse_input_CD : tensor<8x8xi32, #CDR>
 
+    bufferization.dealloc_tensor %0 : tensor<6x6xi32>
     bufferization.dealloc_tensor %1 : tensor<6x6xi32, #DCSR>
     bufferization.dealloc_tensor %2 : tensor<6x6xi32, #DCSR>
     bufferization.dealloc_tensor %3 : tensor<6x6xi32, #CSR>
     bufferization.dealloc_tensor %4 : tensor<6x6xi32, #CDR>
     bufferization.dealloc_tensor %5 : tensor<6x6xi32, #CSC>
+    bufferization.dealloc_tensor %6 : tensor<6x6xi32>
+
     return
   }
 }
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion_sparse2dense.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion_sparse2dense.mlir
index 9b05f9bf3a29c2..e145c4542a7bfc 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion_sparse2dense.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion_sparse2dense.mlir
@@ -233,6 +233,19 @@ module {
     // bufferization.dealloc_tensor %s2pp4 : tensor<2x?x?xf64, #Tensor4>
     // bufferization.dealloc_tensor %s2pp5 : tensor<2x?x?xf64, #Tensor5>
     // bufferization.dealloc_tensor %s2pp6 : tensor<2x?x?xf64, #Tensor6>
+
+    bufferization.dealloc_tensor %d2341 : tensor<2x3x4xf64>
+    bufferization.dealloc_tensor %d2342 : tensor<2x3x4xf64>
+    bufferization.dealloc_tensor %d2343 : tensor<2x3x4xf64>
+    bufferization.dealloc_tensor %d2344 : tensor<2x3x4xf64>
+    bufferization.dealloc_tensor %d2345 : tensor<2x3x4xf64>
+    bufferization.dealloc_tensor %d2346 : tensor<2x3x4xf64>
+    bufferization.dealloc_tensor %dp344 : tensor<?x3x4xf64>
+    bufferization.dealloc_tensor %d2p45 : tensor<2x?x4xf64>
+    bufferization.dealloc_tensor %d23p6 : tensor<2x3x?xf64>
+    bufferization.dealloc_tensor %dp3p4 : tensor<?x3x?xf64>
+    bufferization.dealloc_tensor %dpp45 : tensor<?x?x4xf64>
+
     return
   }
 }
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 0f9dfb9da7204f..541f618a0ac3b1 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion_sparse2sparse.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion_sparse2sparse.mlir
@@ -114,12 +114,14 @@ module {
     call @dump(%d31) : (tensor<2x3x4xf64>) -> ()
 
     //
-    // Release sparse tensors.
+    // Release tensors.
     //
     bufferization.dealloc_tensor %t13 : tensor<2x3x4xf64, #Tensor3>
     bufferization.dealloc_tensor %t31 : tensor<2x3x4xf64, #Tensor1>
     bufferization.dealloc_tensor %s1 : tensor<2x3x4xf64, #Tensor1>
     bufferization.dealloc_tensor %s3 : tensor<2x3x4xf64, #Tensor3>
+    bufferization.dealloc_tensor %d13 : tensor<2x3x4xf64>
+    bufferization.dealloc_tensor %d31 : tensor<2x3x4xf64>
 
     return
   }
@@ -167,12 +169,14 @@ module {
     call @dump(%d31) : (tensor<2x3x4xf64>) -> ()
 
     //
-    // Release sparse tensors.
+    // Release tensors.
     //
     bufferization.dealloc_tensor %t13 : tensor<2x3x4xf64, #SingletonTensor3>
     bufferization.dealloc_tensor %t31 : tensor<2x3x4xf64, #SingletonTensor1>
     bufferization.dealloc_tensor %s1 : tensor<2x3x4xf64, #SingletonTensor1>
     bufferization.dealloc_tensor %s3 : tensor<2x3x4xf64, #SingletonTensor3>
+    bufferization.dealloc_tensor %d13 : tensor<2x3x4xf64>
+    bufferization.dealloc_tensor %d31 : tensor<2x3x4xf64>
 
     return
   }