[Mlir-commits] [mlir] [mlir][scf] Fix ForOpTensorCastFolder on loops that change shape (PR #234422)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Oct 21 16:00:31 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-scf
Author: Jared Hoberock (jaredhoberock)
<details>
<summary>Changes</summary>
`ForOpTensorCastFolder` folds a `tensor.cast` into a `scf.for` without checking what the loop yields, so a non-shape-preserving loop gets a cast to a static type it may not actually have, which is UB. The `matmul_on_tensors` test was one example of such a loop.
This change makes the fold check two things to avoid UB:
1. The loop yields a `tensor.cast` from the static shape, or
2. The loop is shape-preserving
If either is true, the fold is legal.
A loop which happened to fold to the same shape it has at runtime and execute as expected may no longer receive this fold. Such a loop will continue to execute correctly, but it no longer carries a static shape and therefore may be less optimized by downstream passes.
`isShapePreserving` already existed in `LoopCanonicalization.cpp`, so it moves into `SCF.cpp` as a public function so both clients can use it.
Also, `isShapePreserving` now accepts any `DestinationStyleOpInterface` instead of just `tensor.insert_slice` since these ops return their destination tensor with the same type. This also allows the uses in `LoopCanonicalization` to support such ops.
Now, `matmul_on_tensors` yields a cast from the static type. Also adds tests for a nested loop that still folds and two loops that must not fold.
Assisted-by: Claude
---
Full diff: https://github.com/llvm/llvm-project/pull/234422.diff
4 Files Affected:
- (modified) mlir/include/mlir/Dialect/SCF/IR/SCF.h (+6)
- (modified) mlir/lib/Dialect/SCF/IR/SCF.cpp (+41-5)
- (modified) mlir/lib/Dialect/SCF/Transforms/LoopCanonicalization.cpp (-29)
- (modified) mlir/test/Dialect/SCF/canonicalize.mlir (+72-6)
``````````diff
diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCF.h b/mlir/include/mlir/Dialect/SCF/IR/SCF.h
index 455c9c34cb81a3..2440584554959d 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCF.h
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCF.h
@@ -114,6 +114,12 @@ SmallVector<Value> replaceAndCastForOpIterArg(RewriterBase &rewriter,
Value replacement,
const ValueTypeCastFnTy &castFn);
+/// A simple, conservative analysis to determine if the loop is shape
+/// preserving. I.e., the arg-th yielded value has the same shape as the
+/// corresponding region iter_arg: it is the iter_arg, or derives from it
+/// through tied results of destination-style ops and of shape-preserving loops.
+bool isShapePreserving(ForOp forOp, int64_t arg);
+
/// Helper function to compute the difference between two values. This is used
/// by the loop implementations to compute the trip count.
std::optional<llvm::APSInt> computeUbMinusLb(Value lb, Value ub, bool isSigned);
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index ae548cf1d2e10b..4d59f390627f79 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -31,6 +31,7 @@
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/DebugLog.h"
#include <optional>
@@ -974,6 +975,31 @@ mlir::scf::replaceAndCastForOpIterArg(RewriterBase &rewriter, scf::ForOp forOp,
return newResults;
}
+bool mlir::scf::isShapePreserving(ForOp forOp, int64_t arg) {
+ assert(arg < static_cast<int64_t>(forOp.getNumResults()) &&
+ "arg is out of bounds");
+ Value value = forOp.getYieldedValues()[arg];
+ while (value) {
+ if (value == forOp.getRegionIterArgs()[arg])
+ return true;
+ OpResult opResult = dyn_cast<OpResult>(value);
+ if (!opResult)
+ return false;
+
+ value = llvm::TypeSwitch<Operation *, Value>(opResult.getOwner())
+ .Case([&](DestinationStyleOpInterface op) {
+ return op.getTiedOpOperand(opResult)->get();
+ })
+ .Case([&](ForOp forOp) {
+ return isShapePreserving(forOp, opResult.getResultNumber())
+ ? forOp.getInitArgs()[opResult.getResultNumber()]
+ : Value();
+ })
+ .Default(nullptr);
+ }
+ return false;
+}
+
namespace {
/// Fold scf.for iter_arg/result pairs that go through incoming/ougoing
/// a tensor.cast op pair so as to pull the tensor.cast inside the scf.for:
@@ -982,8 +1008,9 @@ namespace {
/// %0 = tensor.cast %t0 : tensor<32x1024xf32> to tensor<?x?xf32>
/// %1 = scf.for %i = %c0 to %c1024 step %c32 iter_args(%iter_t0 = %0)
/// -> (tensor<?x?xf32>) {
-/// %2 = call @do(%iter_t0) : (tensor<?x?xf32>) -> tensor<?x?xf32>
-/// scf.yield %2 : tensor<?x?xf32>
+/// %2 = call @do(%iter_t0) : (tensor<?x?xf32>) -> tensor<32x1024xf32>
+/// %3 = tensor.cast %2 : tensor<32x1024xf32> to tensor<?x?xf32>
+/// scf.yield %3 : tensor<?x?xf32>
/// }
/// use_of(%1)
/// ```
@@ -994,9 +1021,8 @@ namespace {
/// %0 = scf.for %arg2 = %c0 to %c1024 step %c32 iter_args(%arg3 = %arg0)
/// -> (tensor<32x1024xf32>) {
/// %2 = tensor.cast %arg3 : tensor<32x1024xf32> to tensor<?x?xf32>
-/// %3 = call @do(%2) : (tensor<?x?xf32>) -> tensor<?x?xf32>
-/// %4 = tensor.cast %3 : tensor<?x?xf32> to tensor<32x1024xf32>
-/// scf.yield %4 : tensor<32x1024xf32>
+/// %3 = call @do(%2) : (tensor<?x?xf32>) -> tensor<32x1024xf32>
+/// scf.yield %3 : tensor<32x1024xf32>
/// }
/// %1 = tensor.cast %0 : tensor<32x1024xf32> to tensor<?x?xf32>
/// use_of(%1)
@@ -1020,6 +1046,16 @@ struct ForOpTensorCastFolder : public OpRewritePattern<ForOp> {
continue;
if (!std::get<1>(it).hasOneUse())
continue;
+ // The fold casts the yielded value to the static type, which is only
+ // valid if every iteration yields a value of that type.
+ int64_t resultNumber = std::get<1>(it).getResultNumber();
+ auto yieldedCast =
+ op.getYieldedValues()[resultNumber].getDefiningOp<tensor::CastOp>();
+ bool yieldsStaticType =
+ yieldedCast && yieldedCast.getSource().getType() ==
+ incomingCast.getSource().getType();
+ if (!yieldsStaticType && !isShapePreserving(op, resultNumber))
+ continue;
// Create a new ForOp with that iter operand replaced.
rewriter.replaceOp(
diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopCanonicalization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopCanonicalization.cpp
index 6eabf6becf5db7..bb14086086ecca 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopCanonicalization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopCanonicalization.cpp
@@ -21,7 +21,6 @@
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-#include "llvm/ADT/TypeSwitch.h"
namespace mlir {
#define GEN_PASS_DEF_SCFFORLOOPCANONICALIZATION
@@ -31,34 +30,6 @@ namespace mlir {
using namespace mlir;
using namespace mlir::scf;
-/// A simple, conservative analysis to determine if the loop is shape
-/// conserving. I.e., the type of the arg-th yielded value is the same as the
-/// type of the corresponding basic block argument of the loop.
-/// Note: This function handles only simple cases. Expand as needed.
-static bool isShapePreserving(ForOp forOp, int64_t arg) {
- assert(arg < static_cast<int64_t>(forOp.getNumResults()) &&
- "arg is out of bounds");
- Value value = forOp.getYieldedValues()[arg];
- while (value) {
- if (value == forOp.getRegionIterArgs()[arg])
- return true;
- OpResult opResult = dyn_cast<OpResult>(value);
- if (!opResult)
- return false;
-
- using tensor::InsertSliceOp;
- value = llvm::TypeSwitch<Operation *, Value>(opResult.getOwner())
- .Case([&](InsertSliceOp op) { return op.getDest(); })
- .Case([&](ForOp forOp) {
- return isShapePreserving(forOp, opResult.getResultNumber())
- ? forOp.getInitArgs()[opResult.getResultNumber()]
- : Value();
- })
- .Default(nullptr);
- }
- return false;
-}
-
namespace {
/// Fold dim ops of iter_args to dim ops of their respective init args. E.g.:
///
diff --git a/mlir/test/Dialect/SCF/canonicalize.mlir b/mlir/test/Dialect/SCF/canonicalize.mlir
index bd8c35ad882501..6257c33868b491 100644
--- a/mlir/test/Dialect/SCF/canonicalize.mlir
+++ b/mlir/test/Dialect/SCF/canonicalize.mlir
@@ -887,7 +887,7 @@ func.func @replace_duplicate_iter_args(%lb: index, %ub: index, %step: index, %a:
// -----
-func.func private @do(%arg0: tensor<?x?xf32>) -> tensor<?x?xf32>
+func.func private @do(%arg0: tensor<?x?xf32>) -> tensor<32x1024xf32>
func.func @matmul_on_tensors(%t0: tensor<32x1024xf32>) -> tensor<?x?xf32> {
%c0 = arith.constant 0 : index
@@ -895,8 +895,9 @@ func.func @matmul_on_tensors(%t0: tensor<32x1024xf32>) -> tensor<?x?xf32> {
%c1024 = arith.constant 1024 : index
%0 = tensor.cast %t0 : tensor<32x1024xf32> to tensor<?x?xf32>
%1 = scf.for %i = %c0 to %c1024 step %c32 iter_args(%iter_t0 = %0) -> (tensor<?x?xf32>) {
- %2 = func.call @do(%iter_t0) : (tensor<?x?xf32>) -> tensor<?x?xf32>
- scf.yield %2 : tensor<?x?xf32>
+ %2 = func.call @do(%iter_t0) : (tensor<?x?xf32>) -> tensor<32x1024xf32>
+ %3 = tensor.cast %2 : tensor<32x1024xf32> to tensor<?x?xf32>
+ scf.yield %3 : tensor<?x?xf32>
} {some_attr}
return %1 : tensor<?x?xf32>
}
@@ -906,15 +907,80 @@ func.func @matmul_on_tensors(%t0: tensor<32x1024xf32>) -> tensor<?x?xf32> {
// CHECK-NOT: tensor.cast
// CHECK: %[[FOR_RES:.*]] = scf.for {{.*}} iter_args(%[[ITER_T0:.*]] = %[[T0]]) -> (tensor<32x1024xf32>) {
// CHECK: %[[CAST:.*]] = tensor.cast %[[ITER_T0]] : tensor<32x1024xf32> to tensor<?x?xf32>
-// CHECK: %[[DONE:.*]] = func.call @do(%[[CAST]]) : (tensor<?x?xf32>) -> tensor<?x?xf32>
-// CHECK: %[[UNCAST:.*]] = tensor.cast %[[DONE]] : tensor<?x?xf32> to tensor<32x1024xf32>
-// CHECK: scf.yield %[[UNCAST]] : tensor<32x1024xf32>
+// CHECK: %[[DONE:.*]] = func.call @do(%[[CAST]]) : (tensor<?x?xf32>) -> tensor<32x1024xf32>
+// CHECK: scf.yield %[[DONE]] : tensor<32x1024xf32>
// CHECK: } {some_attr}
// CHECK: %[[RES:.*]] = tensor.cast
// CHECK: return %[[RES]] : tensor<?x?xf32>
// -----
+// The inner loop inserts into the tensor it carries, so both loops keep the
+// static shape and the cast moves into the outer loop.
+
+func.func @tensor_cast_into_for_nested_tied_yield(%t0: tensor<32x1024xf32>, %f: f32) -> tensor<?x?xf32> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c32 = arith.constant 32 : index
+ %0 = tensor.cast %t0 : tensor<32x1024xf32> to tensor<?x?xf32>
+ %1 = scf.for %i = %c0 to %c32 step %c1 iter_args(%iter_t0 = %0) -> (tensor<?x?xf32>) {
+ %2 = scf.for %j = %c0 to %c32 step %c1 iter_args(%iter_t1 = %iter_t0) -> (tensor<?x?xf32>) {
+ %3 = tensor.insert %f into %iter_t1[%i, %j] : tensor<?x?xf32>
+ scf.yield %3 : tensor<?x?xf32>
+ }
+ scf.yield %2 : tensor<?x?xf32>
+ }
+ return %1 : tensor<?x?xf32>
+}
+// CHECK-LABEL: func @tensor_cast_into_for_nested_tied_yield
+// CHECK: scf.for {{.*}} -> (tensor<32x1024xf32>) {
+// CHECK: scf.for {{.*}} -> (tensor<32x1024xf32>) {
+
+// -----
+
+// The loop may change the shape it carries, so the cast stays outside it.
+
+func.func private @resize(%arg0: tensor<?x?xf32>) -> tensor<?x?xf32>
+
+func.func @no_tensor_cast_into_for_unknown_yield(%t0: tensor<32x1024xf32>) -> tensor<?x?xf32> {
+ %c0 = arith.constant 0 : index
+ %c32 = arith.constant 32 : index
+ %c1024 = arith.constant 1024 : index
+ %0 = tensor.cast %t0 : tensor<32x1024xf32> to tensor<?x?xf32>
+ %1 = scf.for %i = %c0 to %c1024 step %c32 iter_args(%iter_t0 = %0) -> (tensor<?x?xf32>) {
+ %2 = func.call @resize(%iter_t0) : (tensor<?x?xf32>) -> tensor<?x?xf32>
+ scf.yield %2 : tensor<?x?xf32>
+ }
+ return %1 : tensor<?x?xf32>
+}
+// CHECK-LABEL: func @no_tensor_cast_into_for_unknown_yield
+// CHECK: %[[CAST:.*]] = tensor.cast
+// CHECK: scf.for {{.*}} iter_args(%{{.*}} = %[[CAST]]) -> (tensor<?x?xf32>) {
+
+// -----
+
+// The loop yields a cast from another static type, so the cast stays outside it.
+
+func.func private @reshape(%arg0: tensor<?x?xf32>) -> tensor<16x2048xf32>
+
+func.func @no_tensor_cast_into_for_other_static_yield(%t0: tensor<32x1024xf32>) -> tensor<?x?xf32> {
+ %c0 = arith.constant 0 : index
+ %c32 = arith.constant 32 : index
+ %c1024 = arith.constant 1024 : index
+ %0 = tensor.cast %t0 : tensor<32x1024xf32> to tensor<?x?xf32>
+ %1 = scf.for %i = %c0 to %c1024 step %c32 iter_args(%iter_t0 = %0) -> (tensor<?x?xf32>) {
+ %2 = func.call @reshape(%iter_t0) : (tensor<?x?xf32>) -> tensor<16x2048xf32>
+ %3 = tensor.cast %2 : tensor<16x2048xf32> to tensor<?x?xf32>
+ scf.yield %3 : tensor<?x?xf32>
+ }
+ return %1 : tensor<?x?xf32>
+}
+// CHECK-LABEL: func @no_tensor_cast_into_for_other_static_yield
+// CHECK: %[[CAST:.*]] = tensor.cast
+// CHECK: scf.for {{.*}} iter_args(%{{.*}} = %[[CAST]]) -> (tensor<?x?xf32>) {
+
+// -----
+
// CHECK-LABEL: @cond_prop
func.func @cond_prop(%arg0 : i1) -> index {
%res = scf.if %arg0 -> index {
``````````
</details>
https://github.com/llvm/llvm-project/pull/234422
More information about the Mlir-commits
mailing list