[Mlir-commits] [mlir] [mlir][LICM] Restrict LICM to pure tensor semantics (PR #129673)

[Longsheng Mou]

https://github.com/CoTinker updated https://github.com/llvm/llvm-project/pull/129673

>From 4c45c54c1ae9cb37ba8cf8ba901f5fb2e32d143f Mon Sep 17 00:00:00 2001
From: Longsheng Mou <longshengmou at gmail.com>
Date: Tue, 4 Mar 2025 16:29:40 +0800
Subject: [PATCH] [mlir][LICM] Restrict LICM to pure tensor semantics

This PR fixes a bug where LICM incorrectly allowed buffer semantics, which could lead to a crash.
---
 .../Utils/LoopInvariantCodeMotionUtils.cpp     |  8 ++++----
 .../loop-invariant-subset-hoisting.mlir        | 18 ++++++++++++++++++
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
index 7460746934a78..cb3f2c52e2116 100644
--- a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
@@ -280,11 +280,11 @@ MatchingSubsets::populateSubsetOpsAtIterArg(LoopLikeOpInterface loopLike,
       if (auto insertionOp =
               dyn_cast<SubsetInsertionOpInterface>(use.getOwner())) {
         // Current implementation expects that the insertionOp implement
-        // the destinationStyleOpInterface as well. Abort if that tha is not
-        // the case
-        if (!isa<DestinationStyleOpInterface>(use.getOwner())) {
+        // the DestinationStyleOpInterface and with pure tensor semantics
+        // as well. Abort if that is not the case.
+        auto dstOp = dyn_cast<DestinationStyleOpInterface>(use.getOwner());
+        if (!dstOp || !dstOp.hasPureTensorSemantics())
           return failure();
-        }
 
         // The value must be used as a destination. (In case of a source, the
         // entire tensor would be read, which would prevent any hoisting.)
diff --git a/mlir/test/Transforms/loop-invariant-subset-hoisting.mlir b/mlir/test/Transforms/loop-invariant-subset-hoisting.mlir
index 3a78287a0dcad..c720b3638fad5 100644
--- a/mlir/test/Transforms/loop-invariant-subset-hoisting.mlir
+++ b/mlir/test/Transforms/loop-invariant-subset-hoisting.mlir
@@ -595,3 +595,21 @@ func.func @hoist_vector_transfer_write_pairs_disjoint_tensor(
   }
   return %1 : tensor<?x?xf32>
 }
+
+// -----
+
+// Ensure that cases with buffer semantics exit gracefully.
+
+// CHECK-LABEL: @hoist_buffer
+func.func @hoist_buffer(%arg0: memref<7x7xf16>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %alloc = memref.alloc() : memref<7x7xf16>
+  // CHECK: scf.for
+  // CHECK:   linalg.copy
+  %0 = scf.for %arg1 = %c0 to %c1 step %c1 iter_args(%arg2 = %alloc) -> (memref<7x7xf16>) {
+    linalg.copy ins(%arg0 : memref<7x7xf16>) outs(%arg2 : memref<7x7xf16>)
+    scf.yield %alloc : memref<7x7xf16>
+  }
+  return
+}