Possible Memory Leaks Problem In HCL buffer_add.mlir Unit Test
wyanzhao opened this issue · 3 comments
We identified a possible memory leak in the buffer_add.mlir unit test. The issue seems to be introduced by the hcl.buffer_at pragma. When we run hcl-opt -opt on the following input code:
func @add_buffer_at_axis_0(%A: memref<1024x1024xf32>, %B: memref<1024x1024xf32>, %C: memref<1024x1024xf32>)
{
%s = hcl.create_op_handle "s"
%l1 = hcl.create_loop_handle %s, "i"
%l2 = hcl.create_loop_handle %s, "j"
affine.for %i = 0 to 1024 {
// ...
}
%buf = hcl.buffer_at(%B: memref<1024x1024xf32>, %l1) -> memref<1024xf32>
return
}
We get the following output:
module {
func @add_buffer_at_axis_0(%arg0: memref<1024x1024xf32>, %arg1: memref<1024x1024xf32>, %arg2: memref<1024x1024xf32>) {
affine.for %arg3 = 0 to 1024 {
%0 = memref.alloc() : memref<1024xf32>
%cst = arith.constant 0.000000e+00 : f32
affine.for %arg4 = 0 to 1024 {
affine.store %cst, %0[%arg4] : memref<1024xf32>
} {buffer, loop_name = "j_init", pipeline_ii = 1 : i32}
affine.for %arg4 = 0 to 1024 {
%1 = affine.load %arg0[%arg3, %arg4] : memref<1024x1024xf32>
%cst_0 = arith.constant 1.000000e+00 : f32
%2 = arith.addf %1, %cst_0 : f32
affine.store %2, %0[%arg4] : memref<1024xf32>
} {loop_name = "j"}
affine.for %arg4 = 0 to 1024 {
%1 = affine.load %0[%arg4] : memref<1024xf32>
affine.store %1, %arg1[%arg3, %arg4] : memref<1024x1024xf32>
} {buffer, loop_name = "j_back", pipeline_ii = 1 : i32}
} {loop_name = "i", op_name = "s"}
return
}
}
As you can see, the hcl.buffer_at
transformation generates %0 = memref.alloc()
without generating memref::deallocs
. According to the documentation from MLIR (Buffer Deallocation Internals), deallocs
have to be placed to free allocated memory and avoid memory leaks, this code could cause memory leaks in some platforms.
Thanks @wyanzhao for finding the memory leak issue. This is indeed something we have not considered before. We will look into the code and fix it soon.
After further investigation, I found that the hcl.reuse_at
pragma could also cause a memory leak. These include the blur
, blur5
, blur_x
, and conv2d
functions under the reuse.mlir
unit test, as well as the hcl.buffer_at
related unit tests buffer_conv.mlir
and buffer_gemm.mlir
.
Thank you for your attention to this matter.
Yes, we did not include dealloc
operations for all the alloc
operations for now, and expect the memory to be freed automatically after execution.