[GlobalOptimization] Do not hoist fill-like operations

compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.cpp

@@ -14,6 +14,7 @@
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/Builders.h"
+#include "mlir/Transforms/RegionUtils.h"
 
 namespace mlir::iree_compiler::IREE::LinalgExt {
 
@@ -363,6 +364,25 @@ bool isGatherlikeOp(Operation *op) {
   return true;
 }
 
+bool isFillLikeOp(linalg::GenericOp linalgOp) {
+  // Check if there are any non-scalar inputs or non-scalar captures in the
+  // region.
+  for (Value input : linalgOp.getDpsInputs()) {
+    if (isa<ShapedType>(input.getType())) {
+      return false;
+    }
+  }
+
+  bool foundNonScalar = false;
+  visitUsedValuesDefinedAbove(linalgOp.getRegion(), [&](OpOperand *operand) {
+    if (isa<ShapedType>(operand->get().getType())) {
+      foundNonScalar = true;
+    }
+  });
+
+  return !foundNonScalar;
+}
+
 FailureOr<SmallVector<AffineMap>>
 getIGEMMContractionIndexingMaps(linalg::LinalgOp linalgOp) {
   MLIRContext *ctx = linalgOp->getContext();

compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.h

@@ -195,5 +195,9 @@ bool isBroadcastingOp(linalg::LinalgOp op);
 ///     2. `linalg.yield` consumes the result of a `tensor.extract_slice`
 bool isGatherlikeOp(Operation *op);
 
+/// Returns true if the operation is a GenericOp that has no tensor inputs,
+/// either as inputs or as implicit captures.
+bool isFillLikeOp(linalg::GenericOp op);
+
 } // namespace mlir::iree_compiler::IREE::LinalgExt
 #endif // IREE_COMPILER_DIALECT_LINALGEXT_UTILS_UTILS_H_

compiler/src/iree/compiler/ExternalInterfaces/UtilExternalModels.cpp

@@ -12,6 +12,7 @@
 #include "iree/compiler/Dialect/Encoding/IR/EncodingOps.h"
 #include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtDialect.h"
 #include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"
+#include "iree/compiler/Dialect/LinalgExt/Utils/Utils.h"
 #include "iree/compiler/Dialect/Util/IR/UtilTypes.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
@@ -325,7 +326,7 @@ struct HoistableLinalgOpInterface
         }
       }
     }
-    return !linalg::isaFillOpInterface(genericOp).has_value();
+    return !IREE::LinalgExt::isFillLikeOp(genericOp);
   }
   bool isAtomicallyHoistableOp(Operation *) const { return true; }
   bool isOperandHoistable(Operation *, OpOperand *) const { return true; }