[Mlir-commits] [mlir] [mlir][scf] Fix ForOpTensorCastFolder on loops that change shape (PR #234422)
Jared Hoberock
llvmlistbot at llvm.org
Wed Oct 21 15:59:53 PDT 2026
https://github.com/jaredhoberock created https://github.com/llvm/llvm-project/pull/234422
`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
>From 537852d88359060d735c13abcbaf113a5a015fb8 Mon Sep 17 00:00:00 2001
From: Jared Hoberock <jaredhoberock at gmail.com>
Date: Tue, 20 Oct 2026 12:42:44 -0500
Subject: [PATCH] [mlir][scf] Fix ForOpTensorCastFolder on loops that change
shape
`ForOpTensorCastFolder` gives a loop-carried tensor the more static type
its initial value is cast from, and casts the value the loop yields to
that type. It does so without looking at what the loop yields. A loop
whose body yields a tensor of another shape (a call that returns a tensor
of unknown shape, or a cast from another static type) is rewritten to
cast that tensor to a static type it may not have, which is undefined
behavior at run time.
The fold now applies only when every iteration yields the static type:
the yielded value is a `tensor.cast` from the static type, or the loop is
shape preserving at that position. The shape check is the one
`-scf-for-loop-canonicalization` already uses (`isShapePreserving`),
moved into the SCF dialect library and extended from
`tensor.insert_slice` to all destination-style ops, whose tied results
have the shape of their inits. Loops that insert into or update their
iteration argument, such as tiling loops, still fold. The dim folds of
`-scf-for-loop-canonicalization` now also look through destination-style
ops.
A loop whose yielded value has the static shape at run time without
the check proving it (for example a `tensor.empty` sized by `tensor.dim`
of the iteration argument) no longer folds and keeps the dynamic type. A
pipeline that relied on the fold for such a loop sees the dynamic type
downstream, which can affect later steps such as vectorization or static
allocation. `matmul_on_tensors` relied on the fold for a call of unknown
result shape; its callee now returns the static type.
Assisted-by: Claude
---
mlir/include/mlir/Dialect/SCF/IR/SCF.h | 6 ++
mlir/lib/Dialect/SCF/IR/SCF.cpp | 46 +++++++++--
.../SCF/Transforms/LoopCanonicalization.cpp | 29 -------
mlir/test/Dialect/SCF/canonicalize.mlir | 78 +++++++++++++++++--
4 files changed, 119 insertions(+), 40 deletions(-)
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 {
More information about the Mlir-commits
mailing list