[mlir][sparse] Fix memory leaks (part 2) (llvm#81979)

matthias-springer · web-flow · commit ccc20b4e5633 · 2024-02-17T11:04:17.000+01:00
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).
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
 }

Original file line number	Diff line number	Diff line change
`@@ -380,6 +380,7 @@ module {`
`380`	`380`	`%dv = sparse_tensor.convert %arg0 : tensor<?xf64, #SparseVector> to tensor<?xf64>`
`381`	`381`	`%3 = vector.transfer_read %dv[%c0], %d0: tensor<?xf64>, vector<32xf64>`
`382`	`382`	`vector.print %3 : vector<32xf64>`
	`383`	`+ bufferization.dealloc_tensor %dv : tensor<?xf64>`
`383`	`384`	`return`
`384`	`385`	`}`
`385`	`386`
`@@ -394,6 +395,7 @@ module {`
`394`	`395`	`%dv = sparse_tensor.convert %arg0 : tensor<?xi32, #SparseVector> to tensor<?xi32>`
`395`	`396`	`%3 = vector.transfer_read %dv[%c0], %d0: tensor<?xi32>, vector<32xi32>`
`396`	`397`	`vector.print %3 : vector<32xi32>`
	`398`	`+ bufferization.dealloc_tensor %dv : tensor<?xi32>`
`397`	`399`	`return`
`398`	`400`	`}`
`399`	`401`
`@@ -403,6 +405,7 @@ module {`
`403`	`405`	`%dm = sparse_tensor.convert %arg0 : tensor<?x?xf64, #DCSR> to tensor<?x?xf64>`
`404`	`406`	`%1 = vector.transfer_read %dm[%c0, %c0], %d0: tensor<?x?xf64>, vector<4x8xf64>`
`405`	`407`	`vector.print %1 : vector<4x8xf64>`
	`408`	`+ bufferization.dealloc_tensor %dm : tensor<?x?xf64>`
`406`	`409`	`return`
`407`	`410`	`}`
`408`	`411`
`@@ -418,6 +421,7 @@ module {`
`418`	`421`	`%2 = vector.transfer_read %1[%c0], %du: memref<?xf64>, vector<16xf64>`
`419`	`422`	`vector.print %2 : vector<16xf64>`
`420`	`423`
	`424`	`+ bufferization.dealloc_tensor %c : tensor<4x4xf64>`
`421`	`425`	`return`
`422`	`426`	`}`
`423`	`427`
`@@ -433,6 +437,7 @@ module {`
`433`	`437`	`%2 = vector.transfer_read %1[%c0], %du: memref<?xi8>, vector<16xi8>`
`434`	`438`	`vector.print %2 : vector<16xi8>`
`435`	`439`
	`440`	`+ bufferization.dealloc_tensor %c : tensor<4x4xi8>`
`436`	`441`	`return`
`437`	`442`	`}`
`438`	`443`
Original file line number	Diff line number	Diff line change
`@@ -227,5 +227,9 @@ func.func @entry() {`
`227`	`227`	`bufferization.dealloc_tensor %CDCD_ret : tensor<?x?x?x?xf32, #CDCD>`
`228`	`228`	`bufferization.dealloc_tensor %DCCD_ret : tensor<?x?x?x?xf32, #DCCD>`
`229`	`229`
	`230`	`+ bufferization.dealloc_tensor %1 : tensor<?x?x?x?xf32>`
	`231`	`+ bufferization.dealloc_tensor %2 : tensor<?x?x?x?xf32>`
	`232`	`+ bufferization.dealloc_tensor %3 : tensor<?x?x?x?xf32>`
	`233`	`+`
`230`	`234`	`return`
`231`	`235`	`}`
Original file line number	Diff line number	Diff line change
`@@ -249,5 +249,8 @@ func.func @entry() {`
`249`	`249`	`bufferization.dealloc_tensor %CCCCC_ret : tensor<?x?x?x?x?xf32, #CCCCC>`
`250`	`250`	`bufferization.dealloc_tensor %CDCDC_ret : tensor<?x?x?x?x?xf32, #CDCDC>`
`251`	`251`
	`252`	`+ bufferization.dealloc_tensor %1 : tensor<?x?x?x?x?xf32>`
	`253`	`+ bufferization.dealloc_tensor %2 : tensor<?x?x?x?x?xf32>`
	`254`	`+`
`252`	`255`	`return`
`253`	`256`	`}`
Original file line number	Diff line number	Diff line change
`@@ -106,6 +106,9 @@ module {`
`106`	`106`	`bufferization.dealloc_tensor %s1 : tensor<2x3x4xf64, #Tensor1>`
`107`	`107`	`bufferization.dealloc_tensor %s2 : tensor<2x3x4xf64, #Tensor2>`
`108`	`108`	`bufferization.dealloc_tensor %s3 : tensor<2x3x4xf64, #Tensor3>`
	`109`	`+ bufferization.dealloc_tensor %d1 : tensor<2x3x4xf32>`
	`110`	`+ bufferization.dealloc_tensor %d2 : tensor<2x3x4xf32>`
	`111`	`+ bufferization.dealloc_tensor %d3 : tensor<2x3x4xf32>`
`109`	`112`
`110`	`113`	`return`
`111`	`114`	`}`