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

[llvmlistbot at llvm.org]

Author: Longsheng Mou
Date: 2025-03-19T14:17:49+08:00
New Revision: efc31ecd2773ba378598c86b3a90d038075b9d4e

URL: https://github.com/llvm/llvm-project/commit/efc31ecd2773ba378598c86b3a90d038075b9d4e
DIFF: https://github.com/llvm/llvm-project/commit/efc31ecd2773ba378598c86b3a90d038075b9d4e.diff

LOG: [mlir][LICM] Restrict LICM to pure tensor semantics (#129673)

This PR fixes a bug where LICM incorrectly allowed buffer semantics,
which could lead to a crash. Fixes #129416.

Added: 
    

Modified: 
    mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
    mlir/test/Transforms/loop-invariant-subset-hoisting.mlir

Removed: 
    


################################################################################
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
+}