[mlir][LICM] Restrict LICM to pure tensor semantics

[CoTinker]

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

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: Longsheng Mou (CoTinker)

Changes

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

---

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

2 Files Affected:

• (modified) mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp (+4-4)
• (modified) mlir/test/Transforms/loop-invariant-subset-hoisting.mlir (+16)

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..edc99d7448bcb 100644
--- a/mlir/test/Transforms/loop-invariant-subset-hoisting.mlir
+++ b/mlir/test/Transforms/loop-invariant-subset-hoisting.mlir
@@ -595,3 +595,19 @@ 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>
+  %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
+}

[CoTinker]

Ping~

[cxy-1993]

Others,LGTM