From ccc20b4e5633de2ba6c0aca68f31598c8875e105 Mon Sep 17 00:00:00 2001 From: Matthias Springer Date: Sat, 17 Feb 2024 11:04:17 +0100 Subject: [PATCH] [mlir][sparse] Fix memory leaks (part 2) (#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). --- .../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 da944363cf75c..59ecbfdef8504 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 to tensor %3 = vector.transfer_read %dv[%c0], %d0: tensor, vector<32xf64> vector.print %3 : vector<32xf64> + bufferization.dealloc_tensor %dv : tensor return } @@ -394,6 +395,7 @@ module { %dv = sparse_tensor.convert %arg0 : tensor to tensor %3 = vector.transfer_read %dv[%c0], %d0: tensor, vector<32xi32> vector.print %3 : vector<32xi32> + bufferization.dealloc_tensor %dv : tensor return } @@ -403,6 +405,7 @@ module { %dm = sparse_tensor.convert %arg0 : tensor to tensor %1 = vector.transfer_read %dm[%c0, %c0], %d0: tensor, vector<4x8xf64> vector.print %1 : vector<4x8xf64> + bufferization.dealloc_tensor %dm : tensor return } @@ -418,6 +421,7 @@ module { %2 = vector.transfer_read %1[%c0], %du: memref, 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, 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 c7e4ffab7ab3b..95ce4f1bf48d5 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 bufferization.dealloc_tensor %out2D_nhwc_CCCD : tensor bufferization.dealloc_tensor %out2D_nhwc_CCCC : tensor + bufferization.dealloc_tensor %dense_ret :tensor bufferization.dealloc_tensor %in2D_nhwc_CCCC : tensor bufferization.dealloc_tensor %in2D_nhwc_CCCD : tensor 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 afa6a0c0ebbc7..d0fbce7146fe5 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 bufferization.dealloc_tensor %DCCD_ret : tensor + bufferization.dealloc_tensor %1 : tensor + bufferization.dealloc_tensor %2 : tensor + bufferization.dealloc_tensor %3 : tensor + 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 451d2b8769461..f0a26dc46b056 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 bufferization.dealloc_tensor %DDC_ret : tensor bufferization.dealloc_tensor %DCC_ret : tensor + + bufferization.dealloc_tensor %1 : tensor + bufferization.dealloc_tensor %2 : tensor + bufferization.dealloc_tensor %3 : tensor + bufferization.dealloc_tensor %4 : tensor + 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 2c4b96804ccc7..346a143692897 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 bufferization.dealloc_tensor %CDCDC_ret : tensor + bufferization.dealloc_tensor %1 : tensor + bufferization.dealloc_tensor %2 : tensor + 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 65a37bc8e731e..81f68366be28c 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 af7ee43647581..b4d40ae084015 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 bufferization.dealloc_tensor %in2D_nhwc_CCCC : tensor bufferization.dealloc_tensor %in2D_nhwc_CDCC : tensor + + bufferization.dealloc_tensor %dense_ret : tensor + bufferization.dealloc_tensor %CCCC_ret : tensor + bufferization.dealloc_tensor %CDCC_ret : tensor + 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 85d51931db6c2..96c8a30ade8e4 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 8ee80045afc76..5184083f665d5 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 bufferization.dealloc_tensor %in2D_nhwc_CCCC : tensor bufferization.dealloc_tensor %in2D_nhwc_CDCC : tensor + + bufferization.dealloc_tensor %dense_ret : tensor + bufferization.dealloc_tensor %CCCC_ret : tensor + bufferization.dealloc_tensor %CDCC_ret : tensor + return }