[Mlir-commits] [mlir] [mlir][sparse] Fix memory leaks (part 2) (PR #81979)
Matthias Springer
llvmlistbot at llvm.org
Fri Feb 16 01:50:52 PST 2024
https://github.com/matthias-springer created https://github.com/llvm/llvm-project/pull/81979
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 19025169b4d0ea761fdf98f7d6e6ed8bf2f50835 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm at google.com>
Date: Fri, 16 Feb 2024 09:49:51 +0000
Subject: [PATCH] [mlir][sparse] Fix memory leaks (part 2)
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/sparse_binary.mlir | 5 +++++
.../SparseTensor/CPU/sparse_conv_2d_nchw_fchw.mlir | 1 +
.../SparseTensor/CPU/sparse_conv_2d_nhwc_hwcf.mlir | 4 ++++
.../Dialect/SparseTensor/CPU/sparse_conv_3d.mlir | 6 ++++++
.../SparseTensor/CPU/sparse_conv_3d_ndhwc_dhwcf.mlir | 3 +++
.../SparseTensor/CPU/sparse_conversion_element.mlir | 3 +++
.../CPU/sparse_dilated_conv_2d_nhwc_hwcf.mlir | 5 +++++
.../Dialect/SparseTensor/CPU/sparse_matmul_slice.mlir | 8 +++++++-
.../CPU/sparse_strided_conv_2d_nhwc_hwcf.mlir | 5 +++++
9 files changed, 39 insertions(+), 1 deletion(-)
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_binary.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_binary.mlir
index da944363cf75c4..59ecbfdef85043 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_binary.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_binary.mlir
@@ -380,6 +380,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
}
@@ -394,6 +395,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
}
@@ -403,6 +405,7 @@ module {
%dm = sparse_tensor.convert %arg0 : tensor<?x?xf64, #DCSR> to tensor<?x?xf64>
%1 = vector.transfer_read %dm[%c0, %c0], %d0: tensor<?x?xf64>, vector<4x8xf64>
vector.print %1 : vector<4x8xf64>
+ bufferization.dealloc_tensor %dm : tensor<?x?xf64>
return
}
@@ -418,6 +421,7 @@ module {
%2 = vector.transfer_read %1[%c0], %du: memref<?xf64>, vector<16xf64>
vector.print %2 : vector<16xf64>
+ bufferization.dealloc_tensor %c : tensor<4x4xf64>
return
}
@@ -433,6 +437,7 @@ module {
%2 = vector.transfer_read %1[%c0], %du: memref<?xi8>, vector<16xi8>
vector.print %2 : vector<16xi8>
+ bufferization.dealloc_tensor %c : tensor<4x4xi8>
return
}
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_2d_nchw_fchw.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_2d_nchw_fchw.mlir
index c7e4ffab7ab3b1..95ce4f1bf48d5d 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_2d_nchw_fchw.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_2d_nchw_fchw.mlir
@@ -207,6 +207,7 @@ func.func @entry() {
bufferization.dealloc_tensor %out2D_nhwc : tensor<?x?x?x?xf32>
bufferization.dealloc_tensor %out2D_nhwc_CCCD : tensor<?x?x?x?xf32>
bufferization.dealloc_tensor %out2D_nhwc_CCCC : tensor<?x?x?x?xf32>
+ bufferization.dealloc_tensor %dense_ret :tensor<?x?x?x?xf32>
bufferization.dealloc_tensor %in2D_nhwc_CCCC : tensor<?x?x?x?xf32, #CCCC>
bufferization.dealloc_tensor %in2D_nhwc_CCCD : tensor<?x?x?x?xf32, #CDCD>
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_2d_nhwc_hwcf.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_2d_nhwc_hwcf.mlir
index afa6a0c0ebbc7e..d0fbce7146fe57 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_2d_nhwc_hwcf.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_2d_nhwc_hwcf.mlir
@@ -227,5 +227,9 @@ func.func @entry() {
bufferization.dealloc_tensor %CDCD_ret : tensor<?x?x?x?xf32, #CDCD>
bufferization.dealloc_tensor %DCCD_ret : tensor<?x?x?x?xf32, #DCCD>
+ bufferization.dealloc_tensor %1 : tensor<?x?x?x?xf32>
+ bufferization.dealloc_tensor %2 : tensor<?x?x?x?xf32>
+ bufferization.dealloc_tensor %3 : tensor<?x?x?x?xf32>
+
return
}
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_3d.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_3d.mlir
index 451d2b87694614..f0a26dc46b056e 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_3d.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_3d.mlir
@@ -348,5 +348,11 @@ func.func @entry() {
bufferization.dealloc_tensor %CDC_ret : tensor<?x?x?xf32, #CDC>
bufferization.dealloc_tensor %DDC_ret : tensor<?x?x?xf32, #DDC>
bufferization.dealloc_tensor %DCC_ret : tensor<?x?x?xf32, #DCC>
+
+ bufferization.dealloc_tensor %1 : tensor<?x?x?xf32>
+ bufferization.dealloc_tensor %2 : tensor<?x?x?xf32>
+ bufferization.dealloc_tensor %3 : tensor<?x?x?xf32>
+ bufferization.dealloc_tensor %4 : tensor<?x?x?xf32>
+
return
}
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_3d_ndhwc_dhwcf.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_3d_ndhwc_dhwcf.mlir
index 2c4b96804ccc70..346a1436928971 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_3d_ndhwc_dhwcf.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conv_3d_ndhwc_dhwcf.mlir
@@ -249,5 +249,8 @@ func.func @entry() {
bufferization.dealloc_tensor %CCCCC_ret : tensor<?x?x?x?x?xf32, #CCCCC>
bufferization.dealloc_tensor %CDCDC_ret : tensor<?x?x?x?x?xf32, #CDCDC>
+ bufferization.dealloc_tensor %1 : tensor<?x?x?x?x?xf32>
+ bufferization.dealloc_tensor %2 : tensor<?x?x?x?x?xf32>
+
return
}
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion_element.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion_element.mlir
index 65a37bc8e731e5..81f68366be28c6 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion_element.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion_element.mlir
@@ -106,6 +106,9 @@ module {
bufferization.dealloc_tensor %s1 : tensor<2x3x4xf64, #Tensor1>
bufferization.dealloc_tensor %s2 : tensor<2x3x4xf64, #Tensor2>
bufferization.dealloc_tensor %s3 : tensor<2x3x4xf64, #Tensor3>
+ bufferization.dealloc_tensor %d1 : tensor<2x3x4xf32>
+ bufferization.dealloc_tensor %d2 : tensor<2x3x4xf32>
+ bufferization.dealloc_tensor %d3 : tensor<2x3x4xf32>
return
}
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_dilated_conv_2d_nhwc_hwcf.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_dilated_conv_2d_nhwc_hwcf.mlir
index af7ee43647581e..b4d40ae084015a 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_dilated_conv_2d_nhwc_hwcf.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_dilated_conv_2d_nhwc_hwcf.mlir
@@ -144,5 +144,10 @@ func.func @entry() {
bufferization.dealloc_tensor %filter2D_nhwc_CDCC : tensor<?x?x?x?xf32, #CDCC>
bufferization.dealloc_tensor %in2D_nhwc_CCCC : tensor<?x?x?x?xf32, #CCCC>
bufferization.dealloc_tensor %in2D_nhwc_CDCC : tensor<?x?x?x?xf32, #CDCC>
+
+ bufferization.dealloc_tensor %dense_ret : tensor<?x?x?x?xf32>
+ bufferization.dealloc_tensor %CCCC_ret : tensor<?x?x?x?xf32>
+ bufferization.dealloc_tensor %CDCC_ret : tensor<?x?x?x?xf32>
+
return
}
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul_slice.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul_slice.mlir
index 85d51931db6c21..96c8a30ade8e42 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul_slice.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul_slice.mlir
@@ -264,7 +264,13 @@ module {
%du = tensor.cast %r : tensor<4x4xf64> to tensor<*xf64>
call @printMemrefF64(%du) : (tensor<*xf64>) -> ()
- // Releases resources (we do not need to deallocate slices).
+ // Releases resources.
+ bufferization.dealloc_tensor %c2 : tensor<4x4xf64>
+ bufferization.dealloc_tensor %c3 : tensor<4x4xf64>
+ bufferization.dealloc_tensor %c4 : tensor<4x4xf64>
+ bufferization.dealloc_tensor %c4_coo : tensor<4x4xf64>
+ bufferization.dealloc_tensor %c4_dyn : tensor<4x4xf64>
+ bufferization.dealloc_tensor %d : tensor<4x4xf64>
bufferization.dealloc_tensor %b1 : tensor<8x4xf64, #CSR>
bufferization.dealloc_tensor %t1 : tensor<8x8xf64, #CSR>
bufferization.dealloc_tensor %b1_coo : tensor<8x4xf64, #COO>
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_strided_conv_2d_nhwc_hwcf.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_strided_conv_2d_nhwc_hwcf.mlir
index 8ee80045afc760..5184083f665d56 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_strided_conv_2d_nhwc_hwcf.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_strided_conv_2d_nhwc_hwcf.mlir
@@ -143,5 +143,10 @@ func.func @entry() {
bufferization.dealloc_tensor %filter2D_nhwc_CDCC : tensor<?x?x?x?xf32, #CDCC>
bufferization.dealloc_tensor %in2D_nhwc_CCCC : tensor<?x?x?x?xf32, #CCCC>
bufferization.dealloc_tensor %in2D_nhwc_CDCC : tensor<?x?x?x?xf32, #CDCC>
+
+ bufferization.dealloc_tensor %dense_ret : tensor<?x?x?x?xf32>
+ bufferization.dealloc_tensor %CCCC_ret : tensor<?x?x?x?xf32>
+ bufferization.dealloc_tensor %CDCC_ret : tensor<?x?x?x?xf32>
+
return
}
More information about the Mlir-commits
mailing list