[MLIR][SCF] Remove loop invariant arguments of scf.while

[avarmapml]

Summary

• This commit adds a canonicalization pattern on scf.while to remove the loop invariant arguments.
• An argument is considered loop invariant if the iteration argument value is the same as the corresponding one being yielded (at the same position) in both the before/after block of scf.while.
• For the arguments removed, their use within scf.while and their corresponding scf.while's result are replaced with their corresponding initial value.

Signed-off-by: Abhishek Varma <[e-mail]>

[ftynse]

This seems to be based on a wrong assumption about the iteration arguments being always passed from the before to the after block, and in the same order. This is not guaranteed by the op. I'm okay with a pattern that only works in that specific case, but the pattern must check for it explicitly.

mlir/lib/Dialect/SCF/SCF.cpp

2372–2374

There is no guarantee that beforeBlockArgs and afterBlockArgs have the same length. IIRC, llvm::zip is actually zip-shortest, but make sure this assumption is used consciously.

2385

Couldn't this rather directly take the index + 1-th operand of the condition op and check it instead of iterating over the potentially large number of value users. Same below.

2387

There's no relation between the position of the value in the condition op and iter args prescribed by the op semantics.

For example, the following is perfectly valid and is likely the case we want to simplify here. I just swapped the arguments when passing them from the "before" to the "after" block.

scf.while <...> iter_args(%arg0 = ..., %arg1 = ...) {
  scf.condition %cond, %arg1, %arg0
} do {
^bb0(%arg2, %arg3):
  scf.yield %arg3, %arg2
}

And here we may want to remove %arg1.

scf.while <...> iter_args(%arg0 = ..., %arg1 = ...) {
  ...
  scf.condition %cond, %arg1
} do {
^bb0(%arg2):
  %non_invariant = ...
  scf.yield %non_invariant %arg2
}

2406–2408

RAUW in patterns is usually a bad idea. It will work in this case because it's a simple rewrite pattern, but it would be better to find a workaround.

2423

Nit: I got confused by "creating a new block" that refers to a compound statement in the C++ code rather than to the IR block. Using compound statements with RAII is a pretty common idiom in MLIR that doesn't need a special comment.

2441–2443

Again, there is no guarantee that the argument types for before and after blocks match. The pattern didn't check it either, so it feels like this may end up creating invalid IR.

2450–2451

We should rather do rewriter.replaceOp(op, ...).

2452–2453

This will leave some of the vector entries default-initialized (to nullptr, I suppose). I am surprised that mergeBlocks below is fine with that.

2460–2461

The code should rather populate newBefore/AfterBlockArgs properly and avoid doing RAUW above.

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

2387

It's interesting: I implemented this canonicalization on MHLO WhileOp a couple of weeks ago, and it was very trivial (50 lines: https://github.com/tensorflow/tensorflow/commit/710788431c91b23b5aef8fbfcac1f964c724c9b0#diff-1b007bb7b9fd6d952764a854d28c0e7386595284233dfcd35dc897e0217fffaeR5547-R5594 ) because the condition only takes a single i1.

2406–2408

I also didn't directly see how to avoid this in MHLO, but is it a problem for canonicalization?

[rriddle]

mlir/lib/Dialect/SCF/SCF.cpp

2406–2408

but is it a problem for canonicalization

Yes. Any mutation in a pattern outside of the rewriter is essentially undefined behavior.

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

2406–2408

Can you remind me why this is a problem outside of dialect conversion?

2406–2408

(also it'd be great to be able to assert on this)

[rriddle]

mlir/lib/Dialect/SCF/SCF.cpp

2406–2408

a) We often use canonicalization patterns in pattern drivers other than the greedy rewriter.
b) We've had quite a few crashes in the past with mutations not being updated in the internal maps of the greedy rewriter (often nowadays it doesn't create a crash/just a bit of a different work order).

Those are the two major ones, but the mentality of "it works" is really dangerous for this because it isn't guaranteed to work. We could update the driver in some seemingly legal way (which we have before), and end up crashing existing code. It also breaks composability, because now patterns can only be used in very specific use cases/drivers. None of those are desirable to me, and I'd rather expand on the PatternRewriter API. All or nothing for these types of API contracts is the only way to prevent confusion/headaches/crashes down the line.

[avarmapml]

• Added two patterns to take care of various cases that may arise.
• Not using RAUW rather resorting to the rewrite driver.

[avarmapml]

Hi @ftynse
Thank you for your input. I clearly missed out on the entire scf.while construct earlier.
I've tried addressing the same in this new update.
I've added 2 patterns (also demonstrated with necessary test cases) :-

1. Pattern 1: It aims to figure out which of the before block's arguments are loop invariant.
2. Pattern 2: It aims to remove those values from scf.condition which are defined outside the block.

The above two patterns aims to achieve the non-trivial case you enlisted.

It brings up another opportunity/case : A variant of Pattern 2 which would aim to do the analysis for scf.yield.

I believe this would then cover many possible scenarios.

Added two patterns because it made sense to logically separate the two out which in turn makes them extensible.

You may re-review now.

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

2361

Isn't arg1 intended to be in first position here?
Or maybe the example is missing an %arg2 = %a in the iter_args?

I think it'd help if you were naming the variable differently and matching the names of the "after" block arguments with the condition operands name.

2422

I'm not sure why the condOpIndices is needed here? Can't you just fuse this loop with the previous one?

(also in general I prefer to stick with signed int when possible)

2518

2533

afterBlockArgs is unused in this loop.

2545–2546

We should be able to rewrite the while in-place.

I chatted with River offline yesterday, if we're missing APIs on the rewriter, then let's take revisions such as this one as opportunities to add the missing APIs, even if their implementation is "unsafe" (don't allow roll-back). At least the patterns will be written in an optimal way.

[ftynse]

Thanks! There seems to be a simpler way to express the same transformation.

mlir/lib/Dialect/SCF/SCF.cpp

2361

Isn't arg1 intended to be in first position here?

%arg3_after is always %arg0_before. This way, %arg0_before is always %a. So this looks okay to me.

%arg1_after would also work though, which deserves a comment.

Renaming would definitely help.

2403

Nit: drop the number of stack elements from SmallVector unless you have a strong reason why a specific number is used.

2406–2408

Re how to avoid RAUW: if the value being replaced in a block argument, create a new block without this argument then merge in the old block into the new block and provide the replacement value to the merge call. This is messy and should be better supported by the rewriter.

2422

+1.

I would just walk the use-def backwards, this is the natural way to express such things. We know that i-th before argument comes from either i-th iter_arg or from i-th yield operand. So take the i-th yield operand and see if it either equal to the i-th iter_arg directly or if it some k-th argument of the "after" block. In the former case, we have already proven that the i-th iter_arg is loop invariant. In the latter case, since we know that k-th argument of the "after" block comes from the (k+1)-th operand of the condition, we can further check if the (k+1)-th operand of the condition is equal to the i-th argument of the "before" block. If so, the i-th iter_args is also invariant.

The only iteration we need is over the arguments of yield. I would be also tempted to express this as a pattern that removes one argument at a time since the greedy rewriter will keep applying it as long as it matches, but it is more expensive and there is a limit on the number of repeated applications.

2545–2546

Unless I am missing something, we can't rewrite in-place if the number of returned values changes (it does here).

[bondhugula]

I think it's a big missing piece that there isn't a way to replace uses of a block argument in a rewrite pattern other than by RAUW. There are many rewrite patterns in the tree that use RAUW because there isn't another way -- the approach is completely safe now in the greedy rewrite driver as it would only potentially cause the rewrite driver to pick up the updated operations in its next outer loop iteration instead of the same one -- so there is less progress and slower convergence FWIW. (One would be in trouble however if a populate... method is exposed and the pattern gets added in something that's not the pattern rewrite driver.)

if we're missing APIs on the rewriter, then let's take revisions such as this one as opportunities to add the missing APIs, even if their implementation is "unsafe" (don't allow roll-back).
At least the patterns will be written in an optimal way.

I couldn't parse your last clause. Is there a typo?

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

2545–2546

Right: I think I was thinking about the first pattern above

[mehdi_amini]

In D116923#3237474, @bondhugula wrote:

if we're missing APIs on the rewriter, then let's take revisions such as this one as opportunities to add the missing APIs, even if their implementation is "unsafe" (don't allow roll-back).
At least the patterns will be written in an optimal way.

I couldn't parse your last clause. Is there a typo?

I meant that we should write patterns using only rewrite APIs, but without sacrificing the performance of the rewrite: so as much "in-place" as possible. When we're limited by the rewriter API surface we likely should add more instead of writing patterns in a convoluted way to fit the rewriter limitations.

[bondhugula]

In D116923#3238276, @mehdi_amini wrote:

In D116923#3237474, @bondhugula wrote:

if we're missing APIs on the rewriter, then let's take revisions such as this one as opportunities to add the missing APIs, even if their implementation is "unsafe" (don't allow roll-back).
At least the patterns will be written in an optimal way.

I couldn't parse your last clause. Is there a typo?

I meant that we should write patterns using only rewrite APIs, but without sacrificing the performance of the rewrite: so as much "in-place" as possible. When we're limited by the rewriter API surface we likely should add more instead of writing patterns in a convoluted way to fit the rewriter limitations.

Right - I agree.

[avarmapml]

Addressed review comments.
Major change is w.r.t the way loop invariants were decided.
This update traverses each scf.yield operands and back traces to query whether the
corresponding before block argument should be removed.

[ftynse]

LGTM with comments addressed. Thanks for iterating!

mlir/lib/Dialect/SCF/SCF.cpp

2431

You do like these loops... There is rarely a need to iterate over a list of block arguments or op results to find the right value. Here, this does the same without loops.

auto yieldOpBlockArg = yieldOpArg.dyn_cast<BlockArgument>();
if (yieldOpBlockArg && yieldOpBlockArg.getOwner() == &afterBlock) {
  Value condOpArg = condOpArgs[yieldOpBlockArg.getArgNumber()];
  // ...
}

2464

Nit: I suppose you could do something like ValueRange(newYieldOpArgs).getTypes() and avoid having an extra vector for newBeforeBlockType, but it's not a big deal.

2468–2475

Nit: you can write directly to newBeforeBlockArgs[i] instead of creating a variable

mlir/test/Dialect/SCF/canonicalize.mlir

869–871

Why isn't the order deterministic here?

872–882

I'd just CHECK, we still see the transformation is correct even if the operations are not on immediately next lines (or if they change the syntax to have multiple lines). This makes tests less brittle.

877

Drop the comment here (there is actually a patch that removes it, which would break this test). We shouldn't test for irrelevant things.
Similarly, don't match for a specific block name, just ^{{.*}}( should be fine.

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

2375

When does it make sense to keep implicitly captured values (%a and %b here) as operands of scf.condition?

2467

This loop isn't trivial and like deserve some doc

mlir/test/Dialect/SCF/canonicalize.mlir

913

You don't need to check for things like that (types), it is covered by the verifier already

In general try to check for the strict minimal amount of things.

[avarmapml]

Addressed review comments.
Two major changes to notice are :-

1. Updates were made to use yield op's block argument number instead of resorting to loops.
2. Test cases CHECK lines were made more compact/strict by not comparing types.

mlir/lib/Dialect/SCF/SCF.cpp

2375

You're correct with this observation but this pattern will remove the invariant arguments from before block and replace their uses within.

And that is where the second pattern in this PR (RemoveLoopInvariantValueYielded) comes into play which eliminates such scf.condition operands and the corresponding after block arguments.

2431

Haha.
This was really useful. It never crossed my mind to check for such available APIs.
Thank you!

2464

This was helpful too - thanks.

mlir/test/Dialect/SCF/canonicalize.mlir

869–871

They're indeed deterministic. I'm using CHECK now for all the lines as you mentioned in another comment that it'd make the test cases less brittle.

913

I understand. I've made the changes for other test lines too. You may check.

[avarmapml]

Hi @mehdi_amini @ftynse
The review comments were useful and I've made the said updates in the revision.
You may re-review/accept/land this revision.

[avarmapml]

Hi @ftynse @mehdi_amini @bondhugula ,

I've addressed all the useful review comments made so far.

If there are no other review comments, can this revision be landed?

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

2375

It seemed to me that we can get "for free" the final form in the first pattern instead of doing "half the work" and having to re-do more expensive work afterward?

2449

I'd like that we first check if there is anything to do before starting to build vectors/maps: the common path (do nothing) should be as fast as possible.

[avarmapml]

Making early exit as early as possible before starting to build vectors/maps.

[avarmapml]

mlir/lib/Dialect/SCF/SCF.cpp

2375

The idea was to logically separate the task into two kinds of patterns that we'd like to canonicalize :-

1. Identify and remove loop invariant before block arguments.
2. Identify and remove loop invariant return operands which in turn affects after block arguments.

Early exit from both the patterns would be beneficial in this case as you rightly pointed out.

[mehdi_amini]

Good enough for me right now! I'll let Alex have another look.

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

2400–2406

Nit: move all the declarations closer to their uses, that is after the early exit line 2422 (for those that aren't necessary before)

2418

Nit: no else after break.

[avarmapml]

Handled review comments.

[avarmapml]

Hi @mehdi_amini @ftynse @bondhugula

If there are no more comments, can this revision be landed soon?

[ftynse]

Please rebase on HEAD and reupload. It looks like there is a non-trivial assertion failure.

[avarmapml]

Rebased to latest HEAD.
createBlock API requires argument location too and hence lead to assertion failure.
Fixed the same and brought a couple more declarations closer to their uses.

[avarmapml]

Hi @ftynse

I've rebased the current revision on HEAD and resolved the assert failures.

mlir/lib/Dialect/SCF/SCF.cpp

for (auto yieldedAndBlockArgs :
           replaced = true;
         }
       }
     }
     return success(replaced);
   }
 };
 
+/// Remove loop invariant arguments from `before` block of scf.while.
+/// A before block argument is considered loop invariant if :-
+///   1. i-th yield operand is equal to the i-th while operand.
+///   2. i-th yield operand is k-th after block argument which is (k+1)-th
+///      condition operand AND this (k+1)-th condition operand is equal to i-th
+///      iter argument/while operand.
+/// For the arguments which are removed, their uses inside scf.while
+/// are replaced with their corresponding initial value.
+///
+/// Eg:
+///    INPUT :-
+///    %res = scf.while <...> iter_args(%arg0_before = %a, %arg1_before = %b,
+///                                     ..., %argN_before = %N)
+///           {
+///                ...
+///                scf.condition(%cond) %arg1_before, %arg0_before,
+///                                     %arg2_before, %arg0_before, ...
+///           } do {
+///             ^bb0(%arg1_after, %arg0_after_1, %arg2_after, %arg0_after_2,
+///                  ..., %argK_after):
+///                ...
+///                scf.yield %arg0_after_2, %b, %arg1_after, ..., %argN
+///           }
+///
+///    OUTPUT :-
+///    %res = scf.while <...> iter_args(%arg2_before = %c, ..., %argN_before =
+///                                     %N)
+///           {
+///                ...
+///                scf.condition(%cond) %b, %a, %arg2_before, %a, ...
+///           } do {
+///             ^bb0(%arg1_after, %arg0_after_1, %arg2_after, %arg0_after_2,
+///                  ..., %argK_after):
+///                ...
+///                scf.yield %arg1_after, ..., %argN
+///           }
+///
+///    EXPLANATION:
+///      We iterate over each yield operand.
+///        1. 0-th yield operand %arg0_after_2 is 4-th condition operand
+///           %arg0_before, which in turn is the 0-th iter argument. So we
+///           remove 0-th before block argument and yield operand, and replace
+///           all uses of the 0-th before block argument with its initial value
+///           %a.
+///        2. 1-th yield operand %b is equal to the 1-th iter arg's initial
+///           value. So we remove this operand and the corresponding before
+///           block argument and replace all uses of 1-th before block argument
+///           with %b.
+struct RemoveLoopInvariantArgsFromBeforeBlock
+    : public OpRewritePattern<WhileOp> {
+  using OpRewritePattern<WhileOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(WhileOp op,
+                                PatternRewriter &rewriter) const override {
+    Block &afterBlock = op.getAfter().front();
+    Block::BlockArgListType beforeBlockArgs = op.getBeforeArguments();
+    ConditionOp condOp = op.getConditionOp();
+    OperandRange condOpArgs = condOp.getArgs();
+    Operation *yieldOp = afterBlock.getTerminator();
+    ValueRange yieldOpArgs = yieldOp->getOperands();
+
+    bool canSimplify = false;
+    for (auto it : llvm::enumerate(llvm::zip(op.getOperands(), yieldOpArgs))) {
+      auto index = static_cast<unsigned>(it.index());
+      Value initVal, yieldOpArg;
+      std::tie(initVal, yieldOpArg) = it.value();
+      // If i-th yield operand is equal to the i-th operand of the scf.while,
+      // the i-th before block argument is a loop invariant.
+      if (yieldOpArg == initVal) {
+        canSimplify = true;
+        break;
+      }
+      // If the i-th yield operand is k-th after block argument, then we check
+      // if the (k+1)-th condition op operand is equal to either the i-th before
+      // block argument or the initial value of i-th before block argument. If
+      // the comparison results `true`, i-th before block argument is a loop
+      // invariant.
+      auto yieldOpBlockArg = yieldOpArg.dyn_cast<BlockArgument>();
+      if (yieldOpBlockArg && yieldOpBlockArg.getOwner() == &afterBlock) {
+        Value condOpArg = condOpArgs[yieldOpBlockArg.getArgNumber()];
+        if (condOpArg == beforeBlockArgs[index] || condOpArg == initVal) {
+          canSimplify = true;
+          break;
+        }
+      }
+    }
+
+    if (!canSimplify)
+      return failure();
+
+    SmallVector<Value> newInitArgs, newYieldOpArgs;
+    DenseMap<unsigned, Value> beforeBlockInitValMap;
+    SmallVector<Location> newBeforeBlockArgLocs;
+    for (auto it : llvm::enumerate(llvm::zip(op.getOperands(), yieldOpArgs))) {
+      auto index = static_cast<unsigned>(it.index());
+      Value initVal, yieldOpArg;
+      std::tie(initVal, yieldOpArg) = it.value();
+
+      // If i-th yield operand is equal to the i-th operand of the scf.while,
+      // the i-th before block argument is a loop invariant.
+      if (yieldOpArg == initVal) {
+        beforeBlockInitValMap.insert({index, initVal});
+        continue;
+      } else {
+        // If the i-th yield operand is k-th after block argument, then we check
+        // if the (k+1)-th condition op operand is equal to either the i-th
+        // before block argument or the initial value of i-th before block
+        // argument. If the comparison results `true`, i-th before block
+        // argument is a loop invariant.
+        auto yieldOpBlockArg = yieldOpArg.dyn_cast<BlockArgument>();
+        if (yieldOpBlockArg && yieldOpBlockArg.getOwner() == &afterBlock) {
+          Value condOpArg = condOpArgs[yieldOpBlockArg.getArgNumber()];
+          if (condOpArg == beforeBlockArgs[index] || condOpArg == initVal) {
+            beforeBlockInitValMap.insert({index, initVal});
+            continue;
+          }
+        }
+      }
+      newInitArgs.emplace_back(initVal);
+      newYieldOpArgs.emplace_back(yieldOpArg);
+      newBeforeBlockArgLocs.emplace_back(beforeBlockArgs[index].getLoc());
+    }
+
+    {
+      OpBuilder::InsertionGuard g(rewriter);
+      rewriter.setInsertionPoint(yieldOp);
+      rewriter.replaceOpWithNewOp<YieldOp>(yieldOp, newYieldOpArgs);
+    }
+
+    auto newWhile =
+        rewriter.create<WhileOp>(op.getLoc(), op.getResultTypes(), newInitArgs);
+
+    Block &newBeforeBlock = *rewriter.createBlock(
+        &newWhile.getBefore(), /*insertPt*/ {},
+        ValueRange(newYieldOpArgs).getTypes(), newBeforeBlockArgLocs);
+
+    Block &beforeBlock = op.getBefore().front();
+    SmallVector<Value> newBeforeBlockArgs(beforeBlock.getNumArguments());
+    // For each i-th before block argument we find it's replacement value as :-
+    //   1. If i-th before block argument is a loop invariant, we fetch it's
+    //      initial value from `beforeBlockInitValMap` by querying for key `i`.
+    //   2. Else we fetch j-th new before block argument as the replacement
+    //      value of i-th before block argument.
+    for (unsigned i = 0, j = 0, n = beforeBlock.getNumArguments(); i < n; i++) {
+      // If the index 'i' argument was a loop invariant we fetch it's initial
+      // value from `beforeBlockInitValMap`.
+      if (beforeBlockInitValMap.count(i) != 0)
+        newBeforeBlockArgs[i] = beforeBlockInitValMap[i];
+      else
+        newBeforeBlockArgs[i] = newBeforeBlock.getArgument(j++);
+    }
+
+    rewriter.mergeBlocks(&beforeBlock, &newBeforeBlock, newBeforeBlockArgs);
+    rewriter.inlineRegionBefore(op.getAfter(), newWhile.getAfter(),
+                                newWhile.getAfter().begin());
+
+    rewriter.replaceOp(op, newWhile.getResults());
+    return success();
+  }
+};
+
+/// Remove loop invariant value from result (condition op) of scf.while.
+/// A value is considered loop invariant if the final value yielded by
+/// scf.condition is defined outside of the `before` block. We remove the
+/// corresponding argument in `after` block and replace the use with the value.
+/// We also replace the use of the corresponding result of scf.while with the
+/// value.
+///
+/// Eg:
+///    INPUT :-
+///    %res_input:K = scf.while <...> iter_args(%arg0_before = , ...,
+///                                             %argN_before = %N) {
+///                ...
+///                scf.condition(%cond) %arg0_before, %a, %b, %arg1_before, ...
+///           } do {
+///             ^bb0(%arg0_after, %arg1_after, %arg2_after, ..., %argK_after):
+///                ...
+///                some_func(%arg1_after)
+///                ...
+///                scf.yield %arg0_after, %arg2_after, ..., %argN_after
+///           }
+///
+///    OUTPUT :-
+///    %res_output:M = scf.while <...> iter_args(%arg0 = , ..., %argN = %N) {
+///                ...
+///                scf.condition(%cond) %arg0, %arg1, ..., %argM
+///           } do {
+///             ^bb0(%arg0, %arg3, ..., %argM):
+///                ...
+///                some_func(%a)
+///                ...
+///                scf.yield %arg0, %b, ..., %argN
+///           }
+///
+///     EXPLANATION:
+///       1. The 1-th and 2-th operand of scf.condition are defined outside the
+///          before block of scf.while, so they get removed.
+///       2. %res_input#1's uses are replaced by %a and %res_input#2's uses are
+///          replaced by %b.
+///       3. The corresponding after block argument %arg1_after's uses are
+///          replaced by %a and %arg2_after's uses are replaced by %b.
+struct RemoveLoopInvariantValueYielded : public OpRewritePattern<WhileOp> {
+  using OpRewritePattern<WhileOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(WhileOp op,
+                                PatternRewriter &rewriter) const override {
+    Block &beforeBlock = op.getBefore().front();
+    ConditionOp condOp = op.getConditionOp();
+    OperandRange condOpArgs = condOp.getArgs();
+
+    bool canSimplify = false;
+    for (Value condOpArg : condOpArgs) {
+      // Those values not defined within `before` block will be considered as
+      // loop invariant values. We map the corresponding `index` with their
+      // value.
+      if (condOpArg.getParentBlock() != &beforeBlock) {
+        canSimplify = true;
+        break;
+      }
+    }
+
+    if (!canSimplify)
+      return failure();
+
+    Block::BlockArgListType afterBlockArgs = op.getAfterArguments();
+
+    SmallVector<Value> newCondOpArgs;
+    SmallVector<Type> newAfterBlockType;
+    DenseMap<unsigned, Value> condOpInitValMap;
+    SmallVector<Location> newAfterBlockArgLocs;
+    for (auto it : llvm::enumerate(condOpArgs)) {
+      auto index = static_cast<unsigned>(it.index());
+      Value condOpArg = it.value();
+      // Those values not defined within `before` block will be considered as
+      // loop invariant values. We map the corresponding `index` with their
+      // value.
+      if (condOpArg.getParentBlock() != &beforeBlock) {
+        condOpInitValMap.insert({index, condOpArg});
+      } else {
+        newCondOpArgs.emplace_back(condOpArg);
+        newAfterBlockType.emplace_back(condOpArg.getType());
+        newAfterBlockArgLocs.emplace_back(afterBlockArgs[index].getLoc());
+      }
+    }
+
+    {
+      OpBuilder::InsertionGuard g(rewriter);
+      rewriter.setInsertionPoint(condOp);
+      rewriter.replaceOpWithNewOp<ConditionOp>(condOp, condOp.getCondition(),
+                                               newCondOpArgs);
+    }
+
+    auto newWhile = rewriter.create<WhileOp>(op.getLoc(), newAfterBlockType,
+                                             op.getOperands());
+
+    Block &newAfterBlock =
+        *rewriter.createBlock(&newWhile.getAfter(), /*insertPt*/ {},
+                              newAfterBlockType, newAfterBlockArgLocs);
+
+    Block &afterBlock = op.getAfter().front();
+    // Since a new scf.condition op was created, we need to fetch the new
+    // `after` block arguments which will be used while replacing operations of
+    // previous scf.while's `after` blocks. We'd also be fetching new result
+    // values too.
+    SmallVector<Value> newAfterBlockArgs(afterBlock.getNumArguments());
+    SmallVector<Value> newWhileResults(afterBlock.getNumArguments());
+    for (unsigned i = 0, j = 0, n = afterBlock.getNumArguments(); i < n; i++) {
+      Value afterBlockArg, result;
+      // If index 'i' argument was loop invariant we fetch it's value from the
+      // `condOpInitMap` map.
+      if (condOpInitValMap.count(i) != 0) {
+        afterBlockArg = condOpInitValMap[i];
+        result = afterBlockArg;
+      } else {
+        afterBlockArg = newAfterBlock.getArgument(j);
+        result = newWhile.getResult(j);
+        j++;
+      }
+      newAfterBlockArgs[i] = afterBlockArg;
+      newWhileResults[i] = result;
+    }
+
+    rewriter.mergeBlocks(&afterBlock, &newAfterBlock, newAfterBlockArgs);
+    rewriter.inlineRegionBefore(op.getBefore(), newWhile.getBefore(),
+                                newWhile.getBefore().begin());
+
+    rewriter.replaceOp(op, newWhileResults);
+    return success();
+  }
+};
+
 /// Remove WhileOp results that are also unused in 'after' block.
 ///
 ///  %0:2 = scf.while () : () -> (i32, i64) {
 ///    %condition = "test.condition"() : () -> i1
 ///    %v1 = "test.get_some_value"() : () -> i32
 ///    %v2 = "test.get_some_value"() : () -> i64
 ///    scf.condition(%condition) %v1, %v2 : i32, i64
 ///  } do {
LogicalResult matchAndRewrite(WhileOp op,
     rewriter.replaceOpWithNewOp<YieldOp>(yield, newYields);
     return success();
   }
 };
 } // namespace
 
 void WhileOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                           MLIRContext *context) {
-  results.insert<WhileConditionTruth, WhileUnusedResult, WhileCmpCond,
-                 WhileUnusedArg>(context);
+  results.insert<RemoveLoopInvariantArgsFromBeforeBlock,
+                 RemoveLoopInvariantValueYielded, WhileConditionTruth,
+                 WhileCmpCond, WhileUnusedResult>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/SCF/SCFOps.cpp.inc"

mlir/test/Dialect/SCF/canonicalize.mlir

 // CHECK-NEXT:         %[[res:.*]] = scf.while (%[[arg2:.+]] = %{{.*}}) : (i32) -> i32 {
 // CHECK-NEXT:           %[[cmp:.*]] = "test.condition"(%[[arg2]]) : (i32) -> i1
 // CHECK-NEXT:           scf.condition(%[[cmp]]) %[[arg2]] : i32
 // CHECK-NEXT:         } do {
 // CHECK-NEXT:         ^bb0(%[[post:.+]]: i32):  
 // CHECK-NEXT:           %[[next:.+]] = "test.use"(%[[post]]) : (i32) -> i32
 // CHECK-NEXT:           scf.yield %[[next]] : i32
 // CHECK-NEXT:         }
 // CHECK-NEXT:         return %[[res]] : i32
 
 // -----
 
+// CHECK-LABEL: @invariant_loop_args_in_same_order
+// CHECK-SAME: (%[[FUNC_ARG0:.*]]: tensor<i32>)
+func @invariant_loop_args_in_same_order(%f_arg0: tensor<i32>) -> (tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>) {
+  %cst_0 = arith.constant dense<0> : tensor<i32>
+  %cst_1 = arith.constant dense<1> : tensor<i32>
+  %cst_42 = arith.constant dense<42> : tensor<i32>
+
+  %0:5 = scf.while (%arg0 = %cst_0, %arg1 = %f_arg0, %arg2 = %cst_1, %arg3 = %cst_1, %arg4 = %cst_0) : (tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>) -> (tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>) {
+    %1 = arith.cmpi slt, %arg0, %cst_42 : tensor<i32>
+    %2 = tensor.extract %1[] : tensor<i1>
+    scf.condition(%2) %arg0, %arg1, %arg2, %arg3, %arg4 : tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>
+  } do {
+  ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<i32>, %arg4: tensor<i32>): // no predecessors
+    // %arg1 here will get replaced by %cst_1
+    %1 = arith.addi %arg0, %arg1 : tensor<i32>
+    %2 = arith.addi %arg2, %arg3 : tensor<i32>
+    scf.yield %1, %arg1, %2, %2, %arg4 : tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>
+  }
+  return %0#0, %0#1, %0#2, %0#3, %0#4 : tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>
+}
+// CHECK:    %[[CST42:.*]] = arith.constant dense<42>
+// CHECK:    %[[ONE:.*]] = arith.constant dense<1>
+// CHECK:    %[[ZERO:.*]] = arith.constant dense<0>
+// CHECK:    %[[WHILE:.*]]:3 = scf.while (%[[ARG0:.*]] = %[[ZERO]], %[[ARG2:.*]] = %[[ONE]], %[[ARG3:.*]] = %[[ONE]])
+// CHECK:       arith.cmpi slt, %[[ARG0]], %{{.*}}
+// CHECK:       tensor.extract %{{.*}}[]
+// CHECK:       scf.condition(%{{.*}}) %[[ARG0]], %[[ARG2]], %[[ARG3]]
+// CHECK:    } do {
+// CHECK:     ^{{.*}}(%[[ARG0:.*]]: tensor<i32>, %[[ARG2:.*]]: tensor<i32>, %[[ARG3:.*]]: tensor<i32>):
+// CHECK:       %[[VAL0:.*]] = arith.addi %[[ARG0]], %[[FUNC_ARG0]]
+// CHECK:       %[[VAL1:.*]] = arith.addi %[[ARG2]], %[[ARG3]]
+// CHECK:       scf.yield %[[VAL0]], %[[VAL1]], %[[VAL1]]
+// CHECK:    }
+// CHECK:    return %[[WHILE]]#0, %[[FUNC_ARG0]], %[[WHILE]]#1, %[[WHILE]]#2, %[[ZERO]]
+
+// CHECK-LABEL: @while_loop_invariant_argument_different_order
+func @while_loop_invariant_argument_different_order() -> (tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>) {
+  %cst_0 = arith.constant dense<0> : tensor<i32>
+  %cst_1 = arith.constant dense<1> : tensor<i32>
+  %cst_42 = arith.constant dense<42> : tensor<i32>
+
+  %0:6 = scf.while (%arg0 = %cst_0, %arg1 = %cst_1, %arg2 = %cst_1, %arg3 = %cst_1, %arg4 = %cst_0) : (tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>) -> (tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>) {
+    %1 = arith.cmpi slt, %arg0, %cst_42 : tensor<i32>
+    %2 = tensor.extract %1[] : tensor<i1>
+    scf.condition(%2) %arg1, %arg0, %arg2, %arg0, %arg3, %arg4 : tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>
+  } do {
+  ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<i32>, %arg4: tensor<i32>, %arg5: tensor<i32>): // no predecessors
+    %1 = arith.addi %arg0, %cst_1 : tensor<i32>
+    %2 = arith.addi %arg2, %arg3 : tensor<i32>
+    scf.yield %arg3, %arg1, %2, %2, %arg4 : tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>
+  }
+  return %0#0, %0#1, %0#2, %0#3, %0#4, %0#5 : tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>
+}
+// CHECK:    %[[CST42:.*]] = arith.constant dense<42>
+// CHECK:    %[[ONE:.*]] = arith.constant dense<1>
+// CHECK:    %[[ZERO:.*]] = arith.constant dense<0>
+// CHECK:    %[[WHILE:.*]]:2 = scf.while (%[[ARG1:.*]] = %[[ONE]], %[[ARG4:.*]] = %[[ZERO]])
+// CHECK:       arith.cmpi slt, %[[ZERO]], %[[CST42]]
+// CHECK:       tensor.extract %{{.*}}[]
+// CHECK:       scf.condition(%{{.*}}) %[[ARG1]], %[[ARG4]]
+// CHECK:    } do {
+// CHECK:     ^{{.*}}(%{{.*}}: tensor<i32>, %{{.*}}: tensor<i32>):
+// CHECK:       scf.yield %[[ZERO]], %[[ONE]]
+// CHECK:    }
+// CHECK:    return %[[WHILE]]#0, %[[ZERO]], %[[ONE]], %[[ZERO]], %[[ONE]], %[[WHILE]]#1
+
+// -----
+
 // CHECK-LABEL: @while_unused_result
 func @while_unused_result() -> i32 {
   %0:2 = scf.while () : () -> (i32, i64) {
     %condition = "test.condition"() : () -> i1
     %v1 = "test.get_some_value"() : () -> i32
     %v2 = "test.get_some_value"() : () -> i64
     scf.condition(%condition) %v1, %v2 : i32, i64
   } do {