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

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

Author: CoTinker

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.)

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