[mlir][affine] Add fold logic when the affine.yield has IV as operand in the AffineForEmptyLoopFolder #165639

linuxlonelyeagle · 2025-10-30T01:50:40Z

Reland #164064. Fix original PR by reconstructing the return value(replacements) using llvm::to_vector.

Co-authored-by: Jakub Kuderski <[email protected]>

…oldResult> type.

llvmbot · 2025-10-30T01:51:15Z

@llvm/pr-subscribers-mlir-affine

@llvm/pr-subscribers-mlir

Author: lonely eagle (linuxlonelyeagle)

Changes

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

2 Files Affected:

(modified) mlir/lib/Dialect/Affine/IR/AffineOps.cpp (+24-5)
(modified) mlir/test/Dialect/Affine/canonicalize.mlir (+13)

diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index e0a53cd52f143..feb95d91a4b80 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2610,6 +2610,21 @@ static std::optional<uint64_t> getTrivialConstantTripCount(AffineForOp forOp) {
   return ub - lb <= 0 ? 0 : (ub - lb + step - 1) / step;
 }
 
+/// Calculate the constant value of the loop's induction variable for its last
+/// trip.
+static std::optional<int64_t>
+getConstantInductionVarForLastTrip(AffineForOp forOp) {
+  std::optional<uint64_t> tripCount = getTrivialConstantTripCount(forOp);
+  if (!tripCount.has_value())
+    return std::nullopt;
+  if (tripCount.value() == 0)
+    return std::nullopt;
+  int64_t lb = forOp.getConstantLowerBound();
+  int64_t step = forOp.getStepAsInt();
+  int64_t lastTripIv = lb + (tripCount.value() - 1) * step;
+  return lastTripIv;
+}
+
 /// Fold the empty loop.
 static SmallVector<OpFoldResult> AffineForEmptyLoopFolder(AffineForOp forOp) {
   if (!llvm::hasSingleElement(*forOp.getBody()))
@@ -2622,7 +2637,7 @@ static SmallVector<OpFoldResult> AffineForEmptyLoopFolder(AffineForOp forOp) {
     // results.
     return forOp.getInits();
   }
-  SmallVector<Value, 4> replacements;
+  SmallVector<OpFoldResult> replacements;
   auto yieldOp = cast<AffineYieldOp>(forOp.getBody()->getTerminator());
   auto iterArgs = forOp.getRegionIterArgs();
   bool hasValDefinedOutsideLoop = false;
@@ -2630,10 +2645,14 @@ static SmallVector<OpFoldResult> AffineForEmptyLoopFolder(AffineForOp forOp) {
   for (unsigned i = 0, e = yieldOp->getNumOperands(); i < e; ++i) {
     Value val = yieldOp.getOperand(i);
     BlockArgument *iterArgIt = llvm::find(iterArgs, val);
-    // TODO: It should be possible to perform a replacement by computing the
-    // last value of the IV based on the bounds and the step.
-    if (val == forOp.getInductionVar())
+    if (val == forOp.getInductionVar()) {
+      if (auto lastTripIv = getConstantInductionVarForLastTrip(forOp)) {
+        replacements.push_back(IntegerAttr::get(
+            IndexType::get(forOp.getContext()), lastTripIv.value()));
+        continue;
+      }
       return {};
+    }
     if (iterArgIt == iterArgs.end()) {
       // `val` is defined outside of the loop.
       assert(forOp.isDefinedOutsideOfLoop(val) &&
@@ -2656,7 +2675,7 @@ static SmallVector<OpFoldResult> AffineForEmptyLoopFolder(AffineForOp forOp) {
   // out of order.
   if (tripCount.has_value() && tripCount.value() >= 2 && iterArgsNotInOrder)
     return {};
-  return llvm::to_vector_of<OpFoldResult>(replacements);
+  return replacements;
 }
 
 /// Canonicalize the bounds of the given loop.
diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index 1169cd1c29d74..997f23b4bd669 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -609,6 +609,19 @@ func.func @fold_zero_iter_loops(%in : index) -> index {
 
 // -----
 
+// CHECK-LABEL: func @fold_empty_loop_iv
+//  CHECK-SAME:   %[[INIT:.*]]: index
+func.func @fold_empty_loop_iv(%init: index) -> (index, index) {
+  %res:2 = affine.for %i = 0 to 10 step 1 iter_args(%arg0 = %init, %arg1 = %init) -> (index, index) {
+    affine.yield %i, %arg1 : index, index
+  }
+  // CHECK: %[[C9:.*]] = arith.constant 9 : index
+  // CHECK: return %[[C9]], %[[INIT]] : index, index
+  return %res#0, %res#1 : index, index
+}
+
+// -----
+
 // CHECK-DAG: #[[$SET:.*]] = affine_set<(d0, d1)[s0] : (d0 >= 0, -d0 + 1022 >= 0, d1 >= 0, -d1 + s0 - 2 >= 0)>
 
 // CHECK-LABEL: func @canonicalize_affine_if

linuxlonelyeagle · 2025-10-30T01:54:59Z

The type of replacements is now SmallVector. To be honest, I'm not sure if this fixes the issue.

/home/uweigand/sandbox/buildbot/mlir-s390x-linux/llvm-project/mlir/lib/Dialect/Affine/IR/AffineOps.cpp:2678:10: error: could not convert ‘replacements’ from ‘SmallVector<[...],4>’ to ‘SmallVector<[...],6>’
 2678 |   return replacements;
      |          ^~~~~~~~~~~~
      |          |
      |          SmallVector<[...],4>

kuhar

Can you add some context to the PR description? What PR this relands, why it was reverted, what was fixed on top of the original PR...

linuxlonelyeagle · 2025-10-30T17:22:38Z

Can you add some context to the PR description? What PR this relands, why it was reverted, what was fixed on top of the original PR...

Thank you that tell me how to do it.

kuhar · 2025-10-30T17:38:19Z

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

@@ -2656,7 +2675,7 @@ static SmallVector<OpFoldResult> AffineForEmptyLoopFolder(AffineForOp forOp) {
  // out of order.
  if (tripCount.has_value() && tripCount.value() >= 2 && iterArgsNotInOrder)
    return {};
-  return llvm::to_vector_of<OpFoldResult>(replacements);
+  return replacements;


I don't think this is any different than the previous version -- to_vector_of<T> should pick the same small size as SmallVector<T>

llvm-project/llvm/include/llvm/ADT/SmallVector.h

Lines 1318 to 1319 in e63f0f5

template <typename Out, typename R> SmallVector<Out> to_vector_of(R &&Range) {

return {adl_begin(Range), adl_end(Range)};

kuhar · 2025-10-30T18:13:29Z

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

  if (tripCount.has_value() && tripCount.value() >= 2 && iterArgsNotInOrder)
    return {};
-  return replacements;
+  return llvm::to_vector(replacements);


What was the previous error? Can you share the link to the buildbot that complained? I don't understand why this would fix anything.

https://lab.llvm.org/buildbot/#/builders/117/builds/14631 you can see it in the build-unified-tree.

71.321 [1324/4/3756] Building CXX object tools/mlir/lib/Dialect/Ptr/IR/CMakeFiles/obj.MLIRPtrDialect.dir/PtrDialect.cpp.o 71.352 [1323/4/3757] Building CXX object tools/mlir/lib/Dialect/PDL/IR/CMakeFiles/obj.MLIRPDLDialect.dir/PDLTypes.cpp.o 71.394 [1322/4/3758] Building CXX object tools/mlir/lib/Dialect/PDLInterp/IR/CMakeFiles/obj.MLIRPDLInterpDialect.dir/PDLInterp.cpp.o 71.435 [1321/4/3759] Building CXX object tools/mlir/lib/Dialect/Ptr/IR/CMakeFiles/obj.MLIRPtrMemorySpaceInterfaces.dir/MemorySpaceInterfaces.cpp.o 71.448 [1320/4/3760] Linking CXX static library lib/libMLIRPDLDialect.a 71.518 [1319/4/3761] Linking CXX static library lib/libMLIRPtrMemorySpaceInterfaces.a 71.551 [1318/4/3762] Linking CXX static library lib/libMLIRPDLInterpDialect.a 71.597 [1317/4/3763] Building CXX object tools/mlir/lib/Dialect/Affine/IR/CMakeFiles/obj.MLIRAffineDialect.dir/AffineOps.cpp.o FAILED: tools/mlir/lib/Dialect/Affine/IR/CMakeFiles/obj.MLIRAffineDialect.dir/AffineOps.cpp.o CCACHE_CPP2=yes CCACHE_HASHDIR=yes CCACHE_SLOPPINESS=pch_defines,time_macros /usr/bin/ccache /usr/bin/c++ -DGTEST_HAS_RTTI=0 -D_DEBUG -D_GLIBCXX_ASSERTIONS -D_GLIBCXX_USE_CXX11_ABI=1 -D_GNU_SOURCE -D__STDC_CONSTANT_MACROS -D__STDC_FORMAT_MACROS -D__STDC_LIMIT_MACROS -I/home/uweigand/sandbox/buildbot/mlir-s390x-linux/build/tools/mlir/lib/Dialect/Affine/IR -I/home/uweigand/sandbox/buildbot/mlir-s390x-linux/llvm-project/mlir/lib/Dialect/Affine/IR -I/home/uweigand/sandbox/buildbot/mlir-s390x-linux/build/tools/mlir/include -I/home/uweigand/sandbox/buildbot/mlir-s390x-linux/llvm-project/mlir/include -I/home/uweigand/sandbox/buildbot/mlir-s390x-linux/build/include -I/home/uweigand/sandbox/buildbot/mlir-s390x-linux/llvm-project/llvm/include -fPIC -fno-semantic-interposition -fvisibility-inlines-hidden -Werror=date-time -Wall -Wextra -Wno-unused-parameter -Wwrite-strings -Wcast-qual -Wno-missing-field-initializers -pedantic -Wno-long-long -Wimplicit-fallthrough -Wno-uninitialized -Wno-nonnull -Wno-class-memaccess -Wno-redundant-move -Wno-pessimizing-move -Wno-array-bounds -Wno-stringop-overread -Wno-noexcept-type -Wdelete-non-virtual-dtor -Wsuggest-override -Wno-comment -Wno-misleading-indentation -Wctad-maybe-unsupported -fdiagnostics-color -ffunction-sections -fdata-sections -Wundef -Wno-unused-but-set-parameter -Wno-deprecated-copy -O3 -DNDEBUG -fno-exceptions -funwind-tables -fno-rtti -UNDEBUG -std=c++17 -MD -MT tools/mlir/lib/Dialect/Affine/IR/CMakeFiles/obj.MLIRAffineDialect.dir/AffineOps.cpp.o -MF tools/mlir/lib/Dialect/Affine/IR/CMakeFiles/obj.MLIRAffineDialect.dir/AffineOps.cpp.o.d -o tools/mlir/lib/Dialect/Affine/IR/CMakeFiles/obj.MLIRAffineDialect.dir/AffineOps.cpp.o -c /home/uweigand/sandbox/buildbot/mlir-s390x-linux/llvm-project/mlir/lib/Dialect/Affine/IR/AffineOps.cpp /home/uweigand/sandbox/buildbot/mlir-s390x-linux/llvm-project/mlir/lib/Dialect/Affine/IR/AffineOps.cpp: In function ‘llvm::SmallVector<mlir::OpFoldResult> AffineForEmptyLoopFolder(mlir::affine::AffineForOp)’: /home/uweigand/sandbox/buildbot/mlir-s390x-linux/llvm-project/mlir/lib/Dialect/Affine/IR/AffineOps.cpp:2678:10: error: could not convert ‘replacements’ from ‘SmallVector<[...],4>’ to ‘SmallVector<[...],6>’ 2678 | return replacements; | ^~~~~~~~~~~~ | |

My apologies, I misunderstood your earlier point. I thought you wanted me to use to_vector_of<T>.

linuxlonelyeagle and others added 6 commits October 28, 2025 16:51

add loop yield iv logic into AffineForEmptyLoopFolder.

a4e8857

Update mlir/lib/Dialect/Affine/IR/AffineOps.cpp

0fa939c

Co-authored-by: Jakub Kuderski <[email protected]>

update.

e4ebc2b

Process trip count equals zero.

b0ce9ff

rebase main and update code.

c816d09

make replacements from SmallVector<OpFoldResult,4> to SmallVector<OpF…

86c0c0b

…oldResult> type.

linuxlonelyeagle requested a review from matthias-springer October 30, 2025 01:50

llvmbot added mlir:affine mlir labels Oct 30, 2025

linuxlonelyeagle requested a review from kuhar October 30, 2025 01:50

kuhar reviewed Oct 30, 2025

View reviewed changes

use llvm::to_vector return replacements.

22b96ba

kuhar requested changes Oct 30, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[mlir][affine] Add fold logic when the affine.yield has IV as operand in the AffineForEmptyLoopFolder #165639

[mlir][affine] Add fold logic when the affine.yield has IV as operand in the AffineForEmptyLoopFolder #165639

linuxlonelyeagle commented Oct 30, 2025 •

edited

Loading

Uh oh!

llvmbot commented Oct 30, 2025 •

edited

Loading

Uh oh!

linuxlonelyeagle commented Oct 30, 2025

Uh oh!

kuhar left a comment

Uh oh!

linuxlonelyeagle commented Oct 30, 2025

Uh oh!

kuhar Oct 30, 2025

Uh oh!

kuhar Oct 30, 2025

Uh oh!

kuhar Oct 30, 2025

Uh oh!

linuxlonelyeagle Oct 30, 2025

Uh oh!

linuxlonelyeagle Oct 30, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

	template <typename Out, typename R> SmallVector<Out> to_vector_of(R &&Range) {
	return {adl_begin(Range), adl_end(Range)};

[mlir][affine] Add fold logic when the affine.yield has IV as operand in the AffineForEmptyLoopFolder #165639

Are you sure you want to change the base?

[mlir][affine] Add fold logic when the affine.yield has IV as operand in the AffineForEmptyLoopFolder #165639

Conversation

linuxlonelyeagle commented Oct 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Oct 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

linuxlonelyeagle commented Oct 30, 2025

Uh oh!

kuhar left a comment

Choose a reason for hiding this comment

Uh oh!

linuxlonelyeagle commented Oct 30, 2025

Uh oh!

kuhar Oct 30, 2025

Choose a reason for hiding this comment

Uh oh!

kuhar Oct 30, 2025

Choose a reason for hiding this comment

Uh oh!

kuhar Oct 30, 2025

Choose a reason for hiding this comment

Uh oh!

linuxlonelyeagle Oct 30, 2025

Choose a reason for hiding this comment

Uh oh!

linuxlonelyeagle Oct 30, 2025

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

linuxlonelyeagle commented Oct 30, 2025 •

edited

Loading

llvmbot commented Oct 30, 2025 •

edited

Loading