[mlir][linalg, tosa] Fix memory leaks in integration tests

[matthias-springer]

Buffers are no longer deallocation by One-Shot Bufferize. This is now done by a separate buffer deallocation pass.

Also fix a bug in the vector.mask folding, which was triggered by -buffer-deallocation-pipeline, which runs the canonicalizer.

[llvmbot]

@llvm/pr-subscribers-mlir-linalg
@llvm/pr-subscribers-mlir-vector
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-tosa

Author: Matthias Springer (matthias-springer)

Changes

Buffers are no longer deallocation by One-Shot Bufferize. This is now done by a separate buffer deallocation pass.

Also fix a bug in the vector.mask folding, which was triggered by -buffer-deallocation-pipeline, which runs the canonicalizer.

---

Full diff: https://github.com/llvm/llvm-project/pull/85366.diff

5 Files Affected:

• (modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+1-1)
• (modified) mlir/test/Dialect/Vector/canonicalize.mlir (+12)
• (modified) mlir/test/Integration/Dialect/Linalg/CPU/mmt4d.mlir (+1-1)
• (modified) mlir/test/Integration/Dialect/Linalg/CPU/test-matmul-masked-vec.mlir (+1-1)
• (modified) mlir/test/Integration/Dialect/Tosa/CPU/test-fully-connected.mlir (+1-1)

diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 75f6220ad8f3fa..696d0dac258467 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -6054,7 +6054,7 @@ LogicalResult MaskOp::fold(FoldAdaptor adaptor,
   maskableOp->dropAllUses();
   maskableOp->moveBefore(getOperation());
 
-  results.push_back(maskableOp->getResult(0));
+  llvm::append_range(results, maskableOp->getResults());
   return success();
 }
 
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 4c73a6271786e6..627ac54cf145bf 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -2483,6 +2483,18 @@ func.func @all_true_vector_mask(%a : vector<3x4xf32>) -> vector<3x4xf32> {
 
 // -----
 
+// CHECK-LABEL: func @all_true_vector_mask_no_result(
+func.func @all_true_vector_mask_no_result(%a : vector<3x4xf32>, %m : memref<3x4xf32>) {
+//   CHECK-NOT:   vector.mask
+//       CHECK:   vector.transfer_write
+  %c0 = arith.constant 0 : index
+  %all_true = vector.constant_mask [3, 4] : vector<3x4xi1>
+  vector.mask %all_true { vector.transfer_write %a, %m[%c0, %c0] : vector<3x4xf32>, memref<3x4xf32> } : vector<3x4xi1>
+  return
+}
+
+// -----
+
 // CHECK-LABEL:   func.func @fold_shape_cast_with_mask(
 // CHECK-SAME:     %[[VAL_0:.*]]: tensor<1x?xf32>) -> vector<1x4xi1> {
 func.func @fold_shape_cast_with_mask(%arg0: tensor<1x?xf32>) -> vector<1x4xi1> {
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/mmt4d.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/mmt4d.mlir
index 8ee4e1fb48fef1..92c7039c849601 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/mmt4d.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/mmt4d.mlir
@@ -1,6 +1,6 @@
 // DEFINE: %{compile} =  mlir-opt %s \
 // DEFINE:    -transform-interpreter -test-transform-dialect-erase-schedule \
-// DEFINE:    -one-shot-bufferize -func-bufferize -cse -canonicalize -convert-vector-to-scf -test-lower-to-llvm -o %t
+// DEFINE:    -one-shot-bufferize="bufferize-function-boundaries" -buffer-deallocation-pipeline -cse -canonicalize -convert-vector-to-scf -test-lower-to-llvm -o %t
 // DEFINE: %{entry_point} = mmt4d
 // DEFINE: %{run} = mlir-cpu-runner %t -e %{entry_point} -entry-point-result=void \
 // DEFINE:    -shared-libs=%mlir_runner_utils,%mlir_c_runner_utils
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/test-matmul-masked-vec.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/test-matmul-masked-vec.mlir
index 0378d638df61ab..a70d794506c483 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/test-matmul-masked-vec.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/test-matmul-masked-vec.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s -transform-interpreter -test-transform-dialect-erase-schedule -one-shot-bufferize -func-bufferize -lower-vector-mask --test-lower-to-llvm | \
+// RUN: mlir-opt %s -transform-interpreter -test-transform-dialect-erase-schedule -one-shot-bufferize="bufferize-function-boundaries" -buffer-deallocation-pipeline -lower-vector-mask --test-lower-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void --shared-libs=%mlir_c_runner_utils,%mlir_runner_utils | \
 // RUN: FileCheck %s
 
diff --git a/mlir/test/Integration/Dialect/Tosa/CPU/test-fully-connected.mlir b/mlir/test/Integration/Dialect/Tosa/CPU/test-fully-connected.mlir
index bf178c826574e4..258b1b4f2fab45 100644
--- a/mlir/test/Integration/Dialect/Tosa/CPU/test-fully-connected.mlir
+++ b/mlir/test/Integration/Dialect/Tosa/CPU/test-fully-connected.mlir
@@ -1,5 +1,5 @@
 // RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(tosa-to-linalg-named,tosa-to-linalg,tosa-to-arith))" | \
-// RUN: mlir-opt -one-shot-bufferize -func-bufferize -test-lower-to-llvm | \
+// RUN: mlir-opt -one-shot-bufferize="bufferize-function-boundaries" -buffer-deallocation-pipeline -test-lower-to-llvm | \
 // RUN: mlir-cpu-runner -O3 -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_runner_utils \
 // RUN: | FileCheck %s

[c-rhodes]

Buffers are no longer deallocation by One-Shot Bufferize. This is now done by a separate buffer deallocation pass.

When did this change? We have ArmSME integration tests in Linalg that use one-shot but don't currently dealloc, I guess they're now leaking memory and should be fixed as well?

[matthias-springer]

This changed around September last year. We posted a PSA and migration guide but you may have missed it.

Btw, I ran check-mlir with the address sanitizer today and I did not see any failing/leaking ArmSME tests.

[c-rhodes]

This changed around September last year. We posted a PSA and migration guide but you may have missed it.

Thanks for link.

Btw, I ran check-mlir with the address sanitizer today and I did not see any failing/leaking ArmSME tests.

Great thanks. Grepping on the existing tests I see there's only a single one using -buffer-deallocation-pipeline. With so few tests using it I can see how it could be missed this is required, my approach at least is usually to look at what other tests are doing. Is it worth updating existing tests to always use this to reduce the likelihood of leaks happening in the future?

[c-rhodes]

left a question in my previous comment, but this LGTM as is cheers.

[eric-k256]

The TOSA change looks fine to me.

[matthias-springer]

There are a few tests in test/Integration/Dialect/Linalg that have -buffer-deallocation-pipeline.

Is it worth updating existing tests to always use this to reduce the likelihood of leaks happening in the future?

Yes, I'm in the process of fixing all memory leaks at the moment and this PR will bring down the number of leaking tests to 5 or so. The remaining ones are sparse compiler tests and they need a different fix (but they will eventually also use the buffer deallocation pipeline). To check for leaks, I built MLIR with -DLLVM_USE_SANITIZER="Address" and ran check-mlir. Ideally, we'd have a build bot that runs with ASAN.

[banach-space]

Btw, I ran check-mlir with the address sanitizer today and I did not see any failing/leaking ArmSME tests.

Just to double check, you will need to use these flags (and have recent qemu-aarch64 available):

-DMLIR_INCLUDE_INTEGRATION_TESTS=On -DMLIR_RUN_ARM_SME_TESTS=On -DARM_EMULATOR_EXECUTABLE=<path-to>/qemu-aarch64

In particular, MLIR_INCLUDE_INTEGRATION_TESTS alone is not sufficient :)

To check for leaks, I built MLIR with -DLLVM_USE_SANITIZER="Address" and ran check-mlir. Ideally, we'd have a build bot that runs with ASAN.

@DavidSpickett Something to consider for clang-aarch64-full-2stage? Or would that be too much potential noise?

[banach-space]

OK, I've just checked and see that both SVE and SME integration tests also require updating.

SVE
Use these flags:

  -DMLIR_INCLUDE_INTEGRATION_TESTS=On
  -DMLIR_RUN_ARM_SVE_TESTS=On

Updated in #86488

SME
Use these flags:

  -DMLIR_INCLUDE_INTEGRATION_TESTS=On
  -DMLIR_RUN_ARM_SME_TESTS=On
  -DARM_EMULATOR_EXECUTABLE=<path-to-emulator> 

These tests need to run under an emulator (there's no hardware) and that seems to prevent the sanitizer from working:

LeakSanitizer does not work under ptrace (strace, gdb, etc)

Updated in