Possible Memory Leaks Problem In HCL buffer_add.mlir Unit Test

[wyanzhao]

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.

[chhzh123]

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.

[wyanzhao]

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.

[chhzh123]

Yes, we did not include dealloc operations for all the alloc operations for now, and expect the memory to be freed automatically after execution.