[GreedPatternRewriter] Preprocess constants while building worklist when not processing top down

[rriddle]

Summary

This avoids accidentally reversing the order of constants during successive
application, e.g. when running the canonicalizer. This helps reduce the number
of iterations, and also avoids unnecessary changes to input IR.

Fixes #51892

[jpienaar]

Nice, this was a pain point for folks diff'ing across passes, thanks!

mlir/include/mlir/Transforms/FoldUtils.h

51

opportunity

56

Success and failure is a bit dubious for me here. Returning a bool indicating folded & optionally hoisted feels more descriptive even.

mlir/lib/Transforms/Utils/FoldUtils.cpp

82

Could we make referencedDialect.count into a lamba helper such as alreadyFolded to be more self-documenting?

82

So here we could bubble up above not-known folded ops?

mlir/test/Transforms/test-operation-folder.mlir

27

Could we run -test-patterns 2x above too? (this is probably sufficient, but running 2x is probably more explicit here)

[Mogball]

This change is causing some behaviour which I am struggling to describe. I'm trying to make a reproducer. No need to revert (I don't think).

[Mogball]

So there are at least 2 significant changes to behaviour

1. The folder actually sorts commutative ops, but this doesn't trigger the greedy rewriter to run another iteration because tryToFold might still return failure. If -my-pass only has one pattern commutative(a, b) -> foo but the input IR is commutative(b, a), the pattern doesn't match because the folder will re-order the operands but the rewriter will exit.
2. Constants that are already in the IR aren't getting "folded". I.e. foo.my_const could be "folded" into arith.constant.

[rriddle]

In D122692#3421001, @Mogball wrote:

So there are at least 2 significant changes to behaviour

1. The folder actually sorts commutative ops, but this doesn't trigger the greedy rewriter to run another iteration because tryToFold might still return failure. If -my-pass only has one pattern commutative(a, b) -> foo but the input IR is commutative(b, a), the pattern doesn't match because the folder will re-order the operands but the rewriter will exit.

Was this kind of awkwardly relying on the fact that we always did 2 iterations (given that we refolded existing constants)?

2. Constants that are already in the IR aren't getting "folded". I.e. foo.my_const could be "folded" into arith.constant.

The situations where that happens feels kind of wrong. Folding a constant into a different constant feels weird.

mlir/lib/Transforms/Utils/FoldUtils.cpp

82

Much cleaner, thanks.

82

Yep, just ensures that unique'd constants are pooled at the top. Should never cause issues given that constants can't have operands.

[rriddle]

In D122692#3421016, @rriddle wrote:

In D122692#3421001, @Mogball wrote:

So there are at least 2 significant changes to behaviour

1. The folder actually sorts commutative ops, but this doesn't trigger the greedy rewriter to run another iteration because tryToFold might still return failure. If -my-pass only has one pattern commutative(a, b) -> foo but the input IR is commutative(b, a), the pattern doesn't match because the folder will re-order the operands but the rewriter will exit.

Was this kind of awkwardly relying on the fact that we always did 2 iterations (given that we refolded existing constants)?

2. Constants that are already in the IR aren't getting "folded". I.e. foo.my_const could be "folded" into arith.constant.

The situations where that happens feels kind of wrong. Folding a constant into a different constant feels weird.

To expound upon this further, it feels like an implicit canonicalization on the foo.my_const.

[jpienaar]

In D122692#3421019, @rriddle wrote:

In D122692#3421016, @rriddle wrote:

In D122692#3421001, @Mogball wrote:

So there are at least 2 significant changes to behaviour

1. The folder actually sorts commutative ops, but this doesn't trigger the greedy rewriter to run another iteration because tryToFold might still return failure. If -my-pass only has one pattern commutative(a, b) -> foo but the input IR is commutative(b, a), the pattern doesn't match because the folder will re-order the operands but the rewriter will exit.

Was this kind of awkwardly relying on the fact that we always did 2 iterations (given that we refolded existing constants)?

And one of a or b were constant? Then isn't this expected to work? E.g., the constant operand is flipped before pattern matching is attempted?

2. Constants that are already in the IR aren't getting "folded". I.e. foo.my_const could be "folded" into arith.constant.

The situations where that happens feels kind of wrong. Folding a constant into a different constant feels weird.

To expound upon this further, it feels like an implicit canonicalization on the foo.my_const.

Yeah this is unfortunate, but it meant that previously this was seen as canonical behavior. Not ideal I agree. @Mogball does this just affect unit test rather than semantics/end-to-end behavior? (The new behavior probably results in less changes compared to old - in particular there was the case of discardable attributes getting dropped from constants due to this).

[Chia-hungDuan]

In D122692#3421020, @jpienaar wrote:

In D122692#3421019, @rriddle wrote:

In D122692#3421016, @rriddle wrote:

In D122692#3421001, @Mogball wrote:

So there are at least 2 significant changes to behaviour

1. The folder actually sorts commutative ops, but this doesn't trigger the greedy rewriter to run another iteration because tryToFold might still return failure. If -my-pass only has one pattern commutative(a, b) -> foo but the input IR is commutative(b, a), the pattern doesn't match because the folder will re-order the operands but the rewriter will exit.

Was this kind of awkwardly relying on the fact that we always did 2 iterations (given that we refolded existing constants)?

And one of a or b were constant? Then isn't this expected to work? E.g., the constant operand is flipped before pattern matching is attempted?

Add few more contexts here. My understanding is, we changed the timing of adding a constant to the folder,
In the past, we add a constant when it calls tryToFold but now we add a constant when it's pushed to the worklist(before calling tryToFold)

The case we are seeing is a pattern like
def Pat<(FooOp BarOp, ConstantOp), ...>
The ir is (notice the operand order)

%cst = ...
%bar = ...
%foo = Foo(%cst, %bar)

In the past, it will have two runs on the ir because it will add %cst to the folder and veiw it as a success folding. besides, In the first run, it also reorders the operands of FooOp. As a result, in the second round, the pattern can be matched.

And now, the %cst has been added to the folder. When it calls tryToFold, it just returns failure(). So there's no second round (even the operands of FooOp has been reordered). And when we check the IR that failed to match, it shows

%cst = ...
%bar = ...
%foo = Foo(%bar, %cst)

Which make us think why it wasn't matched even the form looks correct.

2. Constants that are already in the IR aren't getting "folded". I.e. foo.my_const could be "folded" into arith.constant.

The situations where that happens feels kind of wrong. Folding a constant into a different constant feels weird.

To expound upon this further, it feels like an implicit canonicalization on the foo.my_const.

Yeah this is unfortunate, but it meant that previously this was seen as canonical behavior. Not ideal I agree. @Mogball does this just affect unit test rather than semantics/end-to-end behavior? (The new behavior probably results in less changes compared to old - in particular there was the case of discardable attributes getting dropped from constants due to this).

[mehdi_amini]

In D122692#3421156, @Chia-hungDuan wrote:

Add few more contexts here. My understanding is, we changed the timing of adding a constant to the folder,
In the past, we add a constant when it calls tryToFold but now we add a constant when it's pushed to the worklist(before calling tryToFold)

The case we are seeing is a pattern like
def Pat<(FooOp BarOp, ConstantOp), ...>
The ir is (notice the operand order)

%cst = ...
%bar = ...
%foo = Foo(%cst, %bar)

In the past, it will have two runs on the ir because it will add %cst to the folder and veiw it as a success folding. besides, In the first run, it also reorders the operands of FooOp. As a result, in the second round, the pattern can be matched.

And now, the %cst has been added to the folder. When it calls tryToFold, it just returns failure(). So there's no second round (even the operands of FooOp has been reordered). And when we check the IR that failed to match, it shows

%cst = ...
%bar = ...
%foo = Foo(%bar, %cst)

Which make us think why it wasn't matched even the form looks correct.

To me this is a breakage of the API contract of applyPatternsAndFoldGreedily: it does not seem right that applyPatternsAndFoldGreedily "converges" and returns while a pattern would still apply.

[rriddle]

In D122692#3421179, @mehdi_amini wrote:

In D122692#3421156, @Chia-hungDuan wrote:

Add few more contexts here. My understanding is, we changed the timing of adding a constant to the folder,
In the past, we add a constant when it calls tryToFold but now we add a constant when it's pushed to the worklist(before calling tryToFold)

The case we are seeing is a pattern like
def Pat<(FooOp BarOp, ConstantOp), ...>
The ir is (notice the operand order)

%cst = ...
%bar = ...
%foo = Foo(%cst, %bar)

In the past, it will have two runs on the ir because it will add %cst to the folder and veiw it as a success folding. besides, In the first run, it also reorders the operands of FooOp. As a result, in the second round, the pattern can be matched.

And now, the %cst has been added to the folder. When it calls tryToFold, it just returns failure(). So there's no second round (even the operands of FooOp has been reordered). And when we check the IR that failed to match, it shows

%cst = ...
%bar = ...
%foo = Foo(%bar, %cst)

Which make us think why it wasn't matched even the form looks correct.

To me this is a breakage of the API contract of applyPatternsAndFoldGreedily: it does not seem right that applyPatternsAndFoldGreedily "converges" and returns while a pattern would still apply.

Agreed here. It seems to me that this has been an underlying issue masked by the way we were processing constants (i.e. unnecessarily refolding).

[mehdi_amini]

Sent out https://reviews.llvm.org/D122870 for a repro.

[mehdi_amini]

In D122692#3421322, @mehdi_amini wrote:

Sent out https://reviews.llvm.org/D122870 for a repro.

Talked to River offline: seems like a good reason to temporarily revert, he'll re-land with a fix.

flang/test/Lower/host-associated.f90

 ! CHECK:         %[[VAL_28:.*]] = fir.emboxchar %[[VAL_12]], %[[VAL_4]] : (!fir.ref<!fir.char<1,?>>, index) -> !fir.boxchar<1>
 ! CHECK:         return %[[VAL_28]] : !fir.boxchar<1>
 ! CHECK:       }
 
 ! CHECK-LABEL: func @_QPget_message(
 ! CHECK-SAME:                       %[[VAL_0:.*]]: !fir.ref<!fir.char<1,40>>,
 ! CHECK-SAME:                       %[[VAL_1:.*]]: index,
 ! CHECK-SAME:                       %[[VAL_2:.*]]: tuple<!fir.boxproc<() -> ()>, i64> {fir.char_proc}) -> !fir.boxchar<1> {
-! CHECK:         %[[VAL_3:.*]] = arith.constant 40 : index
-! CHECK:         %[[VAL_4:.*]] = arith.constant 12 : index
-! CHECK:         %[[VAL_5:.*]] = arith.constant false
-! CHECK:         %[[VAL_6:.*]] = arith.constant 1 : index
-! CHECK:         %[[VAL_7:.*]] = arith.constant 32 : i8
-! CHECK:         %[[VAL_8:.*]] = arith.constant 0 : index
+! CHECK-DAG:     %[[VAL_3:.*]] = arith.constant 40 : index
+! CHECK-DAG:     %[[VAL_4:.*]] = arith.constant 12 : index
+! CHECK-DAG:     %[[VAL_5:.*]] = arith.constant false
+! CHECK-DAG:     %[[VAL_6:.*]] = arith.constant 1 : index
+! CHECK-DAG:     %[[VAL_7:.*]] = arith.constant 32 : i8
+! CHECK-DAG:     %[[VAL_8:.*]] = arith.constant 0 : index
 ! CHECK:         %[[VAL_9:.*]] = fir.convert %[[VAL_0]] : (!fir.ref<!fir.char<1,40>>) -> !fir.ref<!fir.char<1,?>>
 ! CHECK:         %[[VAL_10:.*]] = fir.address_of(@_QQcl.{{.*}}) : !fir.ref<!fir.char<1,12>>
 ! CHECK:         %[[VAL_11:.*]] = fir.extract_value %[[VAL_2]], [0 : index] : (tuple<!fir.boxproc<() -> ()>, i64>) -> !fir.boxproc<() -> ()>
 ! CHECK:         %[[VAL_12:.*]] = fir.box_addr %[[VAL_11]] : (!fir.boxproc<() -> ()>) -> (() -> ())
 ! CHECK:         %[[VAL_13:.*]] = fir.extract_value %[[VAL_2]], [1 : index] : (tuple<!fir.boxproc<() -> ()>, i64>) -> i64
 ! CHECK:         %[[VAL_14:.*]] = fir.call @llvm.stacksave() : () -> !fir.ref<i8>
 ! CHECK:         %[[VAL_15:.*]] = fir.alloca !fir.char<1,?>(%[[VAL_13]] : i64) {bindc_name = ".result"}
 ! CHECK:         %[[VAL_16:.*]] = fir.convert %[[VAL_12]] : (() -> ()) -> ((!fir.ref<!fir.char<1,?>>, index) -> !fir.boxchar<1>)

mlir/include/mlir/Transforms/FoldUtils.h

public:
   /// operations generated when folding. If the op was completely folded it is
   /// erased. If it is just updated in place, `inPlaceUpdate` is set to true.
   LogicalResult
   tryToFold(Operation *op,
             function_ref<void(Operation *)> processGeneratedConstants = nullptr,
             function_ref<void(Operation *)> preReplaceAction = nullptr,
             bool *inPlaceUpdate = nullptr);
 
+  /// Tries to fold a pre-existing constant operation. `constValue` represents
+  /// the value of the constant, and can be optionally passed if the value is
+  /// already known (e.g. if the constant was discovered by m_Constant). This is
+  /// purely an optimization oppurtunity for callers that already know the value
+  /// of the constant. Returns success if an existing constant for `op` already
+  /// exists in the folder, in which case `op` is replaced and erased.
+  /// Otherwise, returns failure and `op` is inserted into the folder (and
+  /// hoisted if necessary).
+  LogicalResult foldKnownConstant(Operation *op, Attribute constValue = {});
+
   /// Notifies that the given constant `op` should be remove from this
   /// OperationFolder's internal bookkeeping.
   ///
   /// Note: this method must be called if a constant op is to be deleted
   /// externally to this OperationFolder. `op` must be a constant op.
   void notifyRemoval(Operation *op);
 
   /// Create an operation of specific op type with the given builder,

mlir/lib/Transforms/Utils/FoldUtils.cpp

 LogicalResult OperationFolder::tryToFold(
     Operation *op, function_ref<void(Operation *)> processGeneratedConstants,
     function_ref<void(Operation *)> preReplaceAction, bool *inPlaceUpdate) {
   if (inPlaceUpdate)
     *inPlaceUpdate = false;
 
   // If this is a unique'd constant, return failure as we know that it has
   // already been folded.
-  if (referencedDialects.count(op))
+  if (referencedDialects.count(op)) {
+    // Check to see if we should rehoist, i.e. if a non-constant operation was
+    // inserted before this one.
+    Block *opBlock = op->getBlock();
+    if (&opBlock->front() != op && !referencedDialects.count(op->getPrevNode()))
+      op->moveBefore(&opBlock->front());
     return failure();
+  }
 
   // Try to fold the operation.
   SmallVector<Value, 8> results;
   OpBuilder builder(op);
   if (failed(tryToFold(builder, op, results, processGeneratedConstants)))
     return failure();
 
   // Check to see if the operation was just updated in place.
LogicalResult OperationFolder::tryToFold(
 
   // Replace all of the result values and erase the operation.
   for (unsigned i = 0, e = results.size(); i != e; ++i)
     op->getResult(i).replaceAllUsesWith(results[i]);
   op->erase();
   return success();
 }
 
+LogicalResult OperationFolder::foldKnownConstant(Operation *op,
+                                                 Attribute constValue) {
+  Block *opBlock = op->getBlock();
+
+  // If this is a constant we unique'd, we don't need to fold, but we can check
+  // to see if we should rehoist it.
+  if (referencedDialects.count(op)) {
+    if (&opBlock->front() != op && !referencedDialects.count(op->getPrevNode()))
+      op->moveBefore(&opBlock->front());
+    return failure();
+  }
+
+  // Get the constant value of the op if necessary.
+  if (!constValue) {
+    matchPattern(op, m_Constant(&constValue));
+    assert(constValue && "expected `op` to be a constant");
+  } else {
+    // Ensure that the provided constant was actually correct.
+#ifndef NDEBUG
+    Attribute expectedValue;
+    matchPattern(op, m_Constant(&expectedValue));
+    assert(
+        expectedValue == constValue &&
+        "provided constant value was not the expected value of the constant");
+#endif
+  }
+
+  // Check for an existing constant operation for the attribute value.
+  Region *insertRegion = getInsertionRegion(interfaces, opBlock);
+  auto &uniquedConstants = foldScopes[insertRegion];
+  Operation *&folderConstOp = uniquedConstants[std::make_tuple(
+      op->getDialect(), constValue, *op->result_type_begin())];
+
+  // If there is an existing constant, replace `op`.
+  if (folderConstOp) {
+    op->replaceAllUsesWith(folderConstOp);
+    op->erase();
+    return success();
+  }
+
+  // Otherwise, we insert `op`. If `op` is in the insertion block and is either
+  // already at the front of the block, or the previous operation is already a
+  // constant we unique'd (i.e. one we inserted), then we don't need to do
+  // anything. Otherwise, we move the constant to the insertion block.
+  Block *insertBlock = &insertRegion->front();
+  if (opBlock != insertBlock || (&insertBlock->front() != op &&
+                                 !referencedDialects.count(op->getPrevNode())))
+    op->moveBefore(&insertBlock->front());
+
+  folderConstOp = op;
+  referencedDialects[op].push_back(op->getDialect());
+  return failure();
+}
+
 /// Notifies that the given constant `op` should be remove from this
 /// OperationFolder's internal bookkeeping.
 void OperationFolder::notifyRemoval(Operation *op) {
   // Check to see if this operation is uniqued within the folder.
   auto it = referencedDialects.find(op);
   if (it == referencedDialects.end())
     return;
 

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

 //===- GreedyPatternRewriteDriver.cpp - A greedy rewriter -----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements mlir::applyPatternsAndFoldGreedily.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "mlir/IR/Matchers.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Rewrite/PatternApplicator.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/RegionUtils.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ScopedPrinter.h"
#endif
   bool changed = false;
   unsigned iteration = 0;
   do {
     worklist.clear();
     worklistMap.clear();
 
     if (!config.useTopDownTraversal) {
       // Add operations to the worklist in postorder.
-      for (auto &region : regions)
-        region.walk([this](Operation *op) { addToWorklist(op); });
+      for (auto &region : regions) {
+        region.walk([this](Operation *op) {
+          // If we aren't processing top-down, check for existing constants when
+          // populating the worklist. This avoids accidentally reversing the
+          // constant order during processing.
+          Attribute constValue;
+          if (matchPattern(op, m_Constant(&constValue)))
+            if (succeeded(folder.foldKnownConstant(op, constValue)))
+              return;
+          addToWorklist(op);
+        });
+      }
     } else {
       // Add all nested operations to the worklist in preorder.
       for (auto &region : regions)
         region.walk<WalkOrder::PreOrder>(
             [this](Operation *op) { worklist.push_back(op); });
 
       // Reverse the list so our pop-back loop processes them in-order.
       std::reverse(worklist.begin(), worklist.end());

mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir

func @transfer_read_progressive(%A : memref<?x?xf32>, %base: index) -> vector<3x15xf32> {
   // CHECK:         %[[vec_1d:.*]] = vector.transfer_read %[[A]][%{{.*}}, %[[base]]], %[[C7]] : memref<?x?xf32>, vector<15xf32>
   // CHECK:         memref.store %[[vec_1d]], %[[alloc_casted]][%[[I]]] : memref<3xvector<15xf32>>
   // CHECK:       } else {
   // CHECK:         store %[[splat]], %[[alloc_casted]][%[[I]]] : memref<3xvector<15xf32>>
   // CHECK:       }
   // CHECK:     }
   // CHECK:     %[[cst:.*]] = memref.load %[[alloc]][] : memref<vector<3x15xf32>>
 
-  // FULL-UNROLL: %[[C7:.*]] = arith.constant 7.000000e+00 : f32
-  // FULL-UNROLL: %[[VEC0:.*]] = arith.constant dense<7.000000e+00> : vector<3x15xf32>
-  // FULL-UNROLL: %[[C0:.*]] = arith.constant 0 : index
+  // FULL-UNROLL-DAG: %[[C7:.*]] = arith.constant 7.000000e+00 : f32
+  // FULL-UNROLL-DAG: %[[VEC0:.*]] = arith.constant dense<7.000000e+00> : vector<3x15xf32>
+  // FULL-UNROLL-DAG: %[[C0:.*]] = arith.constant 0 : index
   // FULL-UNROLL: %[[DIM:.*]] = memref.dim %[[A]], %[[C0]] : memref<?x?xf32>
   // FULL-UNROLL: cmpi sgt, %[[DIM]], %[[base]] : index
   // FULL-UNROLL: %[[VEC1:.*]] = scf.if %{{.*}} -> (vector<3x15xf32>) {
   // FULL-UNROLL:   vector.transfer_read %[[A]][%[[base]], %[[base]]], %[[C7]] : memref<?x?xf32>, vector<15xf32>
   // FULL-UNROLL:   vector.insert %{{.*}}, %[[VEC0]] [0] : vector<15xf32> into vector<3x15xf32>
   // FULL-UNROLL:   scf.yield %{{.*}} : vector<3x15xf32>
   // FULL-UNROLL: } else {
   // FULL-UNROLL:   scf.yield %{{.*}} : vector<3x15xf32>

mlir/test/Dialect/Async/async-parallel-for-num-worker-threads.mlir

 // RUN: mlir-opt %s -split-input-file -async-parallel-for=num-workers=-1  \
 // RUN: | FileCheck %s --dump-input=always
 
 // CHECK-LABEL: @num_worker_threads(
 // CHECK:       %[[MEMREF:.*]]: memref<?xf32>
 func @num_worker_threads(%arg0: memref<?xf32>) {
 
-  // CHECK: %[[scalingCstInit:.*]] = arith.constant 8.000000e+00 : f32
-  // CHECK: %[[bracketLowerBound4:.*]] = arith.constant 4 : index
-  // CHECK: %[[scalingCst4:.*]] = arith.constant 4.000000e+00 : f32
-  // CHECK: %[[bracketLowerBound8:.*]] = arith.constant 8 : index
-  // CHECK: %[[scalingCst8:.*]] = arith.constant 2.000000e+00 : f32
-  // CHECK: %[[bracketLowerBound16:.*]] = arith.constant 16 : index
-  // CHECK: %[[scalingCst16:.*]] = arith.constant 1.000000e+00 : f32
-  // CHECK: %[[bracketLowerBound32:.*]] = arith.constant 32 : index
-  // CHECK: %[[scalingCst32:.*]] = arith.constant 8.000000e-01 : f32
-  // CHECK: %[[bracketLowerBound64:.*]] = arith.constant 64 : index
-  // CHECK: %[[scalingCst64:.*]] = arith.constant 6.000000e-01 : f32
-  // CHECK:   %[[workersIndex:.*]] = async.runtime.num_worker_threads : index
-  // CHECK:   %[[inBracket4:.*]] = arith.cmpi sgt, %[[workersIndex]], %[[bracketLowerBound4]] : index
-  // CHECK:   %[[scalingFactor4:.*]] = arith.select %[[inBracket4]], %[[scalingCst4]], %[[scalingCstInit]] : f32
-  // CHECK:   %[[inBracket8:.*]] = arith.cmpi sgt, %[[workersIndex]], %[[bracketLowerBound8]] : index
-  // CHECK:   %[[scalingFactor8:.*]] = arith.select %[[inBracket8]], %[[scalingCst8]], %[[scalingFactor4]] : f32
-  // CHECK:   %[[inBracket16:.*]] = arith.cmpi sgt, %[[workersIndex]], %[[bracketLowerBound16]] : index
-  // CHECK:   %[[scalingFactor16:.*]] = arith.select %[[inBracket16]], %[[scalingCst16]], %[[scalingFactor8]] : f32
-  // CHECK:   %[[inBracket32:.*]] = arith.cmpi sgt, %[[workersIndex]], %[[bracketLowerBound32]] : index
-  // CHECK:   %[[scalingFactor32:.*]] = arith.select %[[inBracket32]], %[[scalingCst32]], %[[scalingFactor16]] : f32
-  // CHECK:   %[[inBracket64:.*]] = arith.cmpi sgt, %[[workersIndex]], %[[bracketLowerBound64]] : index
-  // CHECK:   %[[scalingFactor64:.*]] = arith.select %[[inBracket64]], %[[scalingCst64]], %[[scalingFactor32]] : f32
-  // CHECK:   %[[workersInt:.*]] = arith.index_cast %[[workersIndex]] : index to i32
-  // CHECK:   %[[workersFloat:.*]] = arith.sitofp %[[workersInt]] : i32 to f32
-  // CHECK:   %[[scaledFloat:.*]] = arith.mulf %[[scalingFactor64]], %[[workersFloat]] : f32
-  // CHECK:   %[[scaledInt:.*]] = arith.fptosi %[[scaledFloat]] : f32 to i32
-  // CHECK:   %[[scaledIndex:.*]] = arith.index_cast %[[scaledInt]] : i32 to index
+  // CHECK-DAG: %[[scalingCstInit:.*]] = arith.constant 8.000000e+00 : f32
+  // CHECK-DAG: %[[bracketLowerBound4:.*]] = arith.constant 4 : index
+  // CHECK-DAG: %[[scalingCst4:.*]] = arith.constant 4.000000e+00 : f32
+  // CHECK-DAG: %[[bracketLowerBound8:.*]] = arith.constant 8 : index
+  // CHECK-DAG: %[[scalingCst8:.*]] = arith.constant 2.000000e+00 : f32
+  // CHECK-DAG: %[[bracketLowerBound16:.*]] = arith.constant 16 : index
+  // CHECK-DAG: %[[scalingCst16:.*]] = arith.constant 1.000000e+00 : f32
+  // CHECK-DAG: %[[bracketLowerBound32:.*]] = arith.constant 32 : index
+  // CHECK-DAG: %[[scalingCst32:.*]] = arith.constant 8.000000e-01 : f32
+  // CHECK-DAG: %[[bracketLowerBound64:.*]] = arith.constant 64 : index
+  // CHECK-DAG: %[[scalingCst64:.*]] = arith.constant 6.000000e-01 : f32
+  // CHECK:     %[[workersIndex:.*]] = async.runtime.num_worker_threads : index
+  // CHECK:     %[[inBracket4:.*]] = arith.cmpi sgt, %[[workersIndex]], %[[bracketLowerBound4]] : index
+  // CHECK:     %[[scalingFactor4:.*]] = arith.select %[[inBracket4]], %[[scalingCst4]], %[[scalingCstInit]] : f32
+  // CHECK:     %[[inBracket8:.*]] = arith.cmpi sgt, %[[workersIndex]], %[[bracketLowerBound8]] : index
+  // CHECK:     %[[scalingFactor8:.*]] = arith.select %[[inBracket8]], %[[scalingCst8]], %[[scalingFactor4]] : f32
+  // CHECK:     %[[inBracket16:.*]] = arith.cmpi sgt, %[[workersIndex]], %[[bracketLowerBound16]] : index
+  // CHECK:     %[[scalingFactor16:.*]] = arith.select %[[inBracket16]], %[[scalingCst16]], %[[scalingFactor8]] : f32
+  // CHECK:     %[[inBracket32:.*]] = arith.cmpi sgt, %[[workersIndex]], %[[bracketLowerBound32]] : index
+  // CHECK:     %[[scalingFactor32:.*]] = arith.select %[[inBracket32]], %[[scalingCst32]], %[[scalingFactor16]] : f32
+  // CHECK:     %[[inBracket64:.*]] = arith.cmpi sgt, %[[workersIndex]], %[[bracketLowerBound64]] : index
+  // CHECK:     %[[scalingFactor64:.*]] = arith.select %[[inBracket64]], %[[scalingCst64]], %[[scalingFactor32]] : f32
+  // CHECK:     %[[workersInt:.*]] = arith.index_cast %[[workersIndex]] : index to i32
+  // CHECK:     %[[workersFloat:.*]] = arith.sitofp %[[workersInt]] : i32 to f32
+  // CHECK:     %[[scaledFloat:.*]] = arith.mulf %[[scalingFactor64]], %[[workersFloat]] : f32
+  // CHECK:     %[[scaledInt:.*]] = arith.fptosi %[[scaledFloat]] : f32 to i32
+  // CHECK:     %[[scaledIndex:.*]] = arith.index_cast %[[scaledInt]] : i32 to index
 
   %lb = arith.constant 0 : index
   %ub = arith.constant 100 : index
   %st = arith.constant 1 : index
   scf.parallel (%i) = (%lb) to (%ub) step (%st) {
     %one = arith.constant 1.0 : f32
     memref.store %one, %arg0[%i] : memref<?xf32>
   }
 
   return
 }

mlir/test/Dialect/Linalg/detensorize_if.mlir

^bb2(%6: tensor<i32>):  // pred: ^bb1
   } -> tensor<i32>
   cf.br ^bb3(%8 : tensor<i32>)
 
 ^bb3(%10: tensor<i32>):  // pred: ^bb1
   return %10 : tensor<i32>
 }
 
 // CHECK-LABEL:  func @main()
-// CHECK-NEXT:     arith.constant 0
-// CHECK-NEXT:     arith.constant 10
-// CHECK-NEXT:     cf.br ^[[bb1:.*]](%{{.*}}: i32)
+// CHECK-DAG:     arith.constant 0
+// CHECK-DAG:     arith.constant 10
+// CHECK:         cf.br ^[[bb1:.*]](%{{.*}}: i32)
 // CHECK-NEXT:   ^[[bb1]](%{{.*}}: i32):
 // CHECK-NEXT:     arith.cmpi slt, %{{.*}}, %{{.*}}
 // CHECK-NEXT:     cf.cond_br %{{.*}}, ^[[bb2:.*]](%{{.*}} : i32), ^bb3(%{{.*}} : i32)
 // CHECK-NEXT:   ^[[bb2]](%{{.*}}: i32)
 // CHECK-NEXT:     arith.addi %{{.*}}, %{{.*}}
 // CHECK-NEXT:     cf.br ^[[bb3:.*]](%{{.*}} : i32)
 // CHECK-NEXT:   ^[[bb3]](%{{.*}}: i32)
 // CHECK-NEXT:     tensor.from_elements %{{.*}} : tensor<i32>

 ^bb3(%10: tensor<i32>):  // pred: ^bb1
   cf.br ^bb4(%10 : tensor<i32>)
 
 ^bb4(%11: tensor<i32>):  // pred: ^bb1
   return %11 : tensor<i32>
 }
 
 // CHECK-LABEL:  func @main()
-// CHECK-NEXT:     arith.constant 0
-// CHECK-NEXT:     arith.constant 10
-// CHECK-NEXT:     cf.br ^[[bb1:.*]](%{{.*}}: i32)
+// CHECK-DAG:     arith.constant 0
+// CHECK-DAG:     arith.constant 10
+// CHECK:         cf.br ^[[bb1:.*]](%{{.*}}: i32)
 // CHECK-NEXT:   ^[[bb1]](%{{.*}}: i32):
 // CHECK-NEXT:     arith.cmpi slt, %{{.*}}, %{{.*}}
 // CHECK-NEXT:     cf.cond_br %{{.*}}, ^[[bb2:.*]](%{{.*}} : i32), ^bb3(%{{.*}} : i32)
 // CHECK-NEXT:   ^[[bb2]](%{{.*}}: i32)
 // CHECK-NEXT:     arith.addi %{{.*}}, %{{.*}}
 // CHECK-NEXT:     cf.br ^[[bb3:.*]](%{{.*}} : i32)
 // CHECK-NEXT:   ^[[bb3]](%{{.*}}: i32)
 // CHECK-NEXT:     cf.br ^[[bb4:.*]](%{{.*}} : i32)
^bb2(%6: tensor<i32>):  // pred: ^bb1
   } -> tensor<i32>
   cf.br ^bb3(%8 : tensor<i32>)
 
 ^bb3(%10: tensor<i32>):  // pred: ^bb1
   return %10 : tensor<i32>
 }
 
 // CHECK-LABEL:  func @main()
-// CHECK-NEXT:     arith.constant 0
-// CHECK-NEXT:     arith.constant 10
-// CHECK-NEXT:     cf.br ^[[bb1:.*]](%{{.*}}: i32)
+// CHECK-DAG:     arith.constant 0
+// CHECK-DAG:     arith.constant 10
+// CHECK:         cf.br ^[[bb1:.*]](%{{.*}}: i32)
 // CHECK-NEXT:   ^[[bb1]](%{{.*}}: i32):
 // CHECK-NEXT:     arith.cmpi slt, %{{.*}}, %{{.*}}
 // CHECK-NEXT:     cf.cond_br %{{.*}}, ^[[bb2:.*]](%{{.*}} : i32), ^bb2(%{{.*}} : i32)
 // CHECK-NEXT:   ^[[bb2]](%{{.*}}: i32)
 // CHECK-NEXT:     arith.addi %{{.*}}, %{{.*}}
 // CHECK-NEXT:     cf.br ^[[bb3:.*]](%{{.*}} : i32)
 // CHECK-NEXT:   ^[[bb3]](%{{.*}}: i32)
 // CHECK-NEXT:     tensor.from_elements %{{.*}} : tensor<i32>
 // CHECK-NEXT:     return %{{.*}}
 // CHECK-NEXT:   }

mlir/test/Dialect/Linalg/transform-patterns.mlir

func @aligned_promote_fill(%arg0: memref<?x?xf32, offset: ?, strides: [?, 1]>) {
   %cf = arith.constant 1.0 : f32
   %3 = memref.subview %arg0[%c0, %c0][%c2000, %c4000][%c1, %c1] :
    memref<?x?xf32, offset: ?, strides: [?, 1]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
   linalg.fill { __internal_linalg_transform__ = "_promote_views_aligned_"}
    ins(%cf : f32) outs(%3 : memref<?x?xf32, offset: ?, strides: [?, ?]>)
   return
 }
 // CHECK-LABEL: func @aligned_promote_fill
-// CHECK:   %[[cf:.*]] = arith.constant {{.*}} : f32
+// CHECK:   %[[cf:.*]] = arith.constant 1.{{.*}} : f32
 // CHECK:         %[[s0:.*]] = memref.subview {{.*}}: memref<?x?xf32, #map{{.*}}> to memref<?x?xf32, #map{{.*}}>
 // CHECK:         %[[a0:.*]] = memref.alloc() {alignment = 32 : i64} : memref<32000000xi8>
 // CHECK:         %[[v0:.*]] = memref.view %[[a0]]{{.*}} : memref<32000000xi8> to memref<?x?xf32>
 // CHECK:         %[[l0:.*]] = memref.subview %[[v0]][0, 0] [%{{.*}}, %{{.*}}] [1, 1] : memref<?x?xf32> to memref<?x?xf32, #[[$STRIDED_2D_u_1]]>
 // CHECK:         linalg.fill ins({{.*}} : f32) outs(%[[v0]] : memref<?x?xf32>)
 // CHECK:         memref.copy %[[s0]], %[[l0]] : memref<?x?xf32, #map{{.*}}> to memref<?x?xf32, #map{{.*}}>
 // CHECK:         linalg.fill ins(%[[cf]] : f32) outs(%[[v0]] : memref<?x?xf32>)
 

mlir/test/Dialect/SparseTensor/dense.mlir

 //
 // Test with an all-dense-annotated "sparse" matrix as input and
 // a non-annotated dense matrix as output that is inplacable.
 // This allows updating the dense output in place.
 //
 // CHECK-LABEL:   func @dense2(
 // CHECK-SAME:                 %[[VAL_0:.*]]: tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>>,
 // CHECK-SAME:                 %[[VAL_1:.*]]: tensor<32x16xf32> {linalg.inplaceable = true}) -> tensor<32x16xf32> {
-// CHECK:           %[[VAL_2:.*]] = arith.constant 1.000000e+00 : f32
-// CHECK:           %[[VAL_3:.*]] = arith.constant 32 : index
-// CHECK:           %[[VAL_4:.*]] = arith.constant 16 : index
-// CHECK:           %[[VAL_5:.*]] = arith.constant 0 : index
-// CHECK:           %[[VAL_6:.*]] = arith.constant 1 : index
+// CHECK-DAG:       %[[VAL_2:.*]] = arith.constant 1.000000e+00 : f32
+// CHECK-DAG:       %[[VAL_3:.*]] = arith.constant 32 : index
+// CHECK-DAG:       %[[VAL_4:.*]] = arith.constant 16 : index
+// CHECK-DAG:       %[[VAL_5:.*]] = arith.constant 0 : index
+// CHECK-DAG:       %[[VAL_6:.*]] = arith.constant 1 : index
 // CHECK:           %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>> to memref<?xf32>
 // CHECK:           %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16xf32>
 // CHECK:           scf.for %[[VAL_9:.*]] = %[[VAL_5]] to %[[VAL_3]] step %[[VAL_6]] {
 // CHECK:             scf.for %[[VAL_10:.*]] = %[[VAL_5]] to %[[VAL_4]] step %[[VAL_6]] {
 // CHECK:               %[[VAL_11:.*]] = arith.muli %[[VAL_9]], %[[VAL_4]] : index
 // CHECK:               %[[VAL_12:.*]] = arith.addi %[[VAL_11]], %[[VAL_10]] : index
 // CHECK:               %[[VAL_13:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_12]]] : memref<?xf32>
 // CHECK:               %[[VAL_14:.*]] = arith.addf %[[VAL_13]], %[[VAL_2]] : f32

mlir/test/Dialect/SparseTensor/sparse_scalars.mlir

 }
 
 // CHECK-LABEL:   func @mul(
 // CHECK-SAME:              %[[VAL_0:.*0]]: tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>>,
 // CHECK-SAME:              %[[VAL_1:.*1]]: tensor<f32>,
 // CHECK-SAME:              %[[VAL_2:.*2]]: f32,
 // CHECK-SAME:              %[[VAL_3:.*3]]: f32,
 // CHECK-SAME:              %[[VAL_4:.*4]]: tensor<32x16xf32> {linalg.inplaceable = true}) -> tensor<32x16xf32> {
-// CHECK:           %[[VAL_5:.*]] = arith.constant 2.200000e+00 : f32
-// CHECK:           %[[VAL_6:.*]] = arith.constant 0 : index
-// CHECK:           %[[VAL_7:.*]] = arith.constant 1 : index
+// CHECK-DAG:       %[[VAL_5:.*]] = arith.constant 2.200000e+00 : f32
+// CHECK-DAG:       %[[VAL_6:.*]] = arith.constant 0 : index
+// CHECK-DAG:       %[[VAL_7:.*]] = arith.constant 1 : index
 // CHECK:           %[[VAL_8:.*]] = arith.addf %[[VAL_2]], %[[VAL_3]] : f32
 // CHECK:           %[[VAL_9:.*]] = sparse_tensor.pointers %[[VAL_0]], %[[VAL_6]] : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
 // CHECK:           %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_0]], %[[VAL_6]] : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
 // CHECK:           %[[VAL_11:.*]] = sparse_tensor.pointers %[[VAL_0]], %[[VAL_7]] : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
 // CHECK:           %[[VAL_12:.*]] = sparse_tensor.indices %[[VAL_0]], %[[VAL_7]] : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
 // CHECK:           %[[VAL_13:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>> to memref<?xf32>
 // CHECK:           %[[VAL_14:.*]] = bufferization.to_memref %[[VAL_1]] : memref<f32>
 // CHECK:           %[[VAL_15:.*]] = bufferization.to_memref %[[VAL_4]] : memref<32x16xf32>

mlir/test/Dialect/Tensor/bufferize.mlir

func @tensor.from_elements_3d(%f0 : f32) -> tensor<3x2x2xf32> {
   %0 = tensor.from_elements %f0,%f1,%f2,%f3,%f4,%f5,%f6,%f7,%f8,%f9,%f10,%f11
          : tensor<3x2x2xf32>
   return %0 : tensor<3x2x2xf32>
 }
 
 // CHECK-LABEL:   func @tensor.generate(
 // CHECK-SAME:                                       %[[ARG:.*]]: tensor<*xf32>,
 // CHECK-SAME:                                       %[[DYNAMIC_EXTENT:.*]]: index) -> tensor<?xindex> {
-// CHECK:           %[[CASTED:.*]] = bufferization.to_memref %[[ARG]] : memref<*xf32>
 // CHECK-DAG:       %[[C0:.*]] = arith.constant 0 : index
 // CHECK-DAG:       %[[C1:.*]] = arith.constant 1 : index
+// CHECK:           %[[CASTED:.*]] = bufferization.to_memref %[[ARG]] : memref<*xf32>
 // CHECK:           %[[MEMREF:.*]] = memref.alloc(%[[DYNAMIC_EXTENT]]) {{.*}} : memref<?xindex>
 // CHECK:           scf.parallel (%[[I:.*]]) = (%[[C0]]) to (%[[DYNAMIC_EXTENT]]) step (%[[C1]]) {
 // CHECK:             %[[ELEM:.*]] = memref.dim %[[CASTED]], %[[I]] : memref<*xf32>
 // CHECK:             store %[[ELEM]], %[[MEMREF]][%[[I]]] : memref<?xindex>
 // CHECK:             scf.yield
 // CHECK:           }
 // CHECK:           %[[RET:.*]] = bufferization.to_tensor %[[MEMREF]] : memref<?xindex>
 // CHECK:           return %[[RET]] : tensor<?xindex>

mlir/test/Dialect/Tensor/split-padding.mlir

func @pad_non_zero_sizes(%input: tensor<?x?x8xf32>, %low0: index, %high1: index) -> tensor<?x?x8xf32> {
   %f0 = arith.constant 0.0 : f32
   %0 = tensor.pad %input low[%low0, 0, 0] high[0, %high1, 0] {
   ^bb0(%dim0: index, %dim1: index, %dim2: index):
     tensor.yield %f0 : f32
   } : tensor<?x?x8xf32> to tensor<?x?x8xf32>
   return %0 : tensor<?x?x8xf32>
 }
 
-// CHECK: %[[F0:.+]] = arith.constant 0.000000e+00 : f32
-// CHECK: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[F0:.+]] = arith.constant 0.000000e+00 : f32
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
 // CHECK: %[[EQ0:.+]] = arith.cmpi eq, %[[LOW0]], %[[C0]] : index
 // CHECK: %[[EQ1:.+]] = arith.cmpi eq, %[[HIGH1]], %[[C0]] : index
 // CHECK: %[[AND:.+]] = arith.andi %[[EQ0]], %[[EQ1]] : i1
 // CHECK: %[[IF:.+]] = scf.if %[[AND]] -> (tensor<?x?x8xf32>) {
 // CHECK:   scf.yield %[[INPUT]] : tensor<?x?x8xf32>
 // CHECK: } else {
 // CHECK:   %[[PAD:.+]] = tensor.pad %[[INPUT]] low[%[[LOW0]], 0, 0] high[0, %[[HIGH1]], 0]  {
 // CHECK:   ^bb0(%{{.+}}: index, %{{.+}}: index, %{{.+}}: index):
 // CHECK:     tensor.yield %[[F0]] : f32
 // CHECK:   } : tensor<?x?x8xf32> to tensor<?x?x8xf32>
 // CHECK:   scf.yield %[[PAD]] : tensor<?x?x8xf32>
 // CHECK: }
 // CHECK: return %[[IF]] : tensor<?x?x8xf32>

mlir/test/Transforms/test-operation-folder.mlir

func @test_fold_before_previously_folded_op() -> (i32, i32) {
   // CHECK: %[[CST:.+]] = arith.constant true
   // CHECK-NEXT: "test.cast"(%[[CST]]) : (i1) -> i32
   // CHECK-NEXT: "test.cast"(%[[CST]]) : (i1) -> i32
 
   %0 = "test.cast"() {test_fold_before_previously_folded_op} : () -> (i32)
   %1 = "test.cast"() {test_fold_before_previously_folded_op} : () -> (i32)
   return %0, %1 : i32, i32
 }
+
+func @test_dont_reorder_constants() -> (i32, i32, i32) {
+  // Test that we don't reorder existing constants during folding if it isn't necessary.
+  // CHECK: %[[CST:.+]] = arith.constant 1
+  // CHECK-NEXT: %[[CST:.+]] = arith.constant 2
+  // CHECK-NEXT: %[[CST:.+]] = arith.constant 3
+  %0 = arith.constant 1 : i32
+  %1 = arith.constant 2 : i32
+  %2 = arith.constant 3 : i32
+  return %0, %1, %2 : i32, i32, i32
+}