[MLIR] Add a utility to sort the operands of commutative ops

[srishti-pm]

Summary

Added a commutativity utility pattern and a function to populate it. The pattern sorts the operands of an op in ascending order of the "key" associated with each operand iff the op is commutative. This sorting is stable.

The function is intended to be used inside passes to simplify the matching of commutative operations. After the application of the above-mentioned pattern, since the commutative operands now have a deterministic order in which they occur in an op, the matching of large DAGs becomes much simpler, i.e., requires much less number of checks to be written by a user in her/his pattern matching function.

The "key" associated with an operand is the list of the "AncestorKeys" associated with the ancestors of this operand, in a breadth-first order.

The operand of any op is produced by a set of ops and block arguments. Each of these ops and block arguments is called an "ancestor" of this operand.

Now, the "AncestorKey" associated with:

1. A block argument is {type: BLOCK_ARGUMENT, opName: ""}.
2. A non-constant-like op, for example, arith.addi, is {type: NON_CONSTANT_OP, opName: "arith.addi"}.
3. A constant-like op, for example, arith.constant, is {type: CONSTANT_OP, opName: "arith.constant"}.

So, if an operand, say A, was produced as follows:

`<block argument>`  `<block argument>`
             \          /
              \        /
              `arith.subi`           `arith.constant`
                         \            /
                         `arith.addi`
                                |
                           returns `A`

Then, the block arguments and operations present in the backward slice of A, in the breadth-first order are:
arith.addi, arith.subi, arith.constant, <block argument>, and <block argument>.

Thus, the "key" associated with operand A is:

{
 {type: NON_CONSTANT_OP, opName: "arith.addi"},
 {type: NON_CONSTANT_OP, opName: "arith.subi"},
 {type: CONSTANT_OP, opName: "arith.constant"},
 {type: BLOCK_ARGUMENT, opName: ""},
 {type: BLOCK_ARGUMENT, opName: ""}
}

Now, if "keyA" is the key associated with operand A and "keyB" is the key associated with operand B, then:
"keyA" < "keyB" iff:

1. In the first unequal pair of corresponding AncestorKeys, the AncestorKey in operand A is smaller, or,
2. Both the AncestorKeys in every pair are the same and the size of operand A's "key" is smaller.

AncestorKeys of type BLOCK_ARGUMENT are considered the smallest, those of type CONSTANT_OP, the largest, and NON_CONSTANT_OP types come in between. Within the types NON_CONSTANT_OP and CONSTANT_OP, the smaller ones are the ones with smaller op names (lexicographically).

---

Some examples of such a sorting:

Assume that the sorting is being applied to foo.commutative, which is a commutative op.

Example 1:

%1 = foo.const 0
%2 = foo.mul <block argument>, <block argument>
%3 = foo.commutative %1, %2

Here,

1. The key associated with %1 is:

{
 {CONSTANT_OP, "foo.const"}
}

2. The key associated with %2 is:

{
 {NON_CONSTANT_OP, "foo.mul"},
 {BLOCK_ARGUMENT, ""},
 {BLOCK_ARGUMENT, ""}
}

The key of %2 < the key of %1
Thus, the sorted foo.commutative is:

%3 = foo.commutative %2, %1

Example 2:

%1 = foo.const 0
%2 = foo.mul <block argument>, <block argument>
%3 = foo.mul %2, %1
%4 = foo.add %2, %1
%5 = foo.commutative %1, %2, %3, %4

Here,

1. The key associated with %1 is:

{
 {CONSTANT_OP, "foo.const"}
}

2. The key associated with %2 is:

{
 {NON_CONSTANT_OP, "foo.mul"},
 {BLOCK_ARGUMENT, ""}
}

3. The key associated with %3 is:

{
 {NON_CONSTANT_OP, "foo.mul"},
 {NON_CONSTANT_OP, "foo.mul"},
 {CONSTANT_OP, "foo.const"},
 {BLOCK_ARGUMENT, ""},
 {BLOCK_ARGUMENT, ""}
}

4. The key associated with %4 is:

{
 {NON_CONSTANT_OP, "foo.add"},
 {NON_CONSTANT_OP, "foo.mul"},
 {CONSTANT_OP, "foo.const"},
 {BLOCK_ARGUMENT, ""},
 {BLOCK_ARGUMENT, ""}
}

Thus, the sorted foo.commutative is:

%5 = foo.commutative %4, %3, %2, %1

Signed-off-by: Srishti Srivastava <[e-mail]>

[srishti-pm]

Fixing a comment typo and enhancing the commit summary even further.

[mehdi_amini]

Seems like this should be added to canonicalization? The "push constants to the right hand side" is there already.

I also don't understand the complexity of the implementation, I may need an example to understand why you're recursively operating on the producer ops for the operands.
From the revision description: , (1) the operands defined by non-constant-like ops come first, followed by (2) block arguments, and these are followed by (3) the operands defined by constant-like ops which all seems like a fairly local check: a stable-sort on the operands would be deterministic and local to a single operation.

mlir/lib/Transforms/Utils/CommutativityUtils.cpp

236

The string conversion seems unnecessary to me?

272

Is this a leak?

[srishti-pm]

In D124750#3500748, @mehdi_amini wrote:

Seems like this should be added to canonicalization? The "push constants to the right hand side" is there already.

I think this was not added to canonicalization because we wanted it to be an independent utility that can be used if needed and not be used if not needed. Yes, the "push constants to the right hand side" is there already and that's actually the reason why this utility also pushes the constants to the right. I didn't want this utility to do something else and clash with the "push constants to the right hand side" canonicalization if and when both were being used together. So, when we decided what our sorting order will be, we made sure that this order kept constants to the right. For more context, you could refer to this RFC's discussion, starting from this comment: https://discourse.llvm.org/t/mlir-pdl-extending-pdl-pdlinterp-bytecode-to-enable-commutative-matching/60798/12?u=srishtisrivastava.

But, note this:
Right now, the code shows that I am sorting the constants alphabetically (as in, theoretically, we can say that arith.constant comes before tf.Const). I will remove this "alphabetical sorting" among the constants. It isn't consistent with the existing "pushing of constants" and moreover adds unnecessary computation too. I'll ensure that the sorting is stable and pushes all the constants to the right.

I also don't understand the complexity of the implementation, I may need an example to understand why you're recursively operating on the producer ops for the operands.
From the revision description: (1) the operands defined by non-constant-like ops come first, followed by (2) block arguments, and these are followed by (3) the operands defined by constant-like ops` which all seems like a fairly local check: a stable-sort on the operands would be deterministic and local to a single operation.

I do this because, firstly, in the description, if you look below this paragraph, you will see the following:
"And, if two operands come from the same op, the function backtracks and
looks even further to sort them. This backtracking is done over the
backward slice of the operand, in a breadth-first traversal."

So, in essence, we are looking at the entire origin of each of the operands, in a breadth-first fashion, to decide the ordering of the operands. Secondly, we need to sort the producers of the operands so that we have a deterministic sorting (because if the producers are not sorted, we don't know if we are doing the right sorting or not, even if we are looking at the entire backward slice). This is because since the producers are not sorted, the breadth-first traversal becomes non-deterministic to a user who is writing a pattern matching function (say, matchAndRewrite()).

[srishti-pm]

Within constant-like ops, removed the requirement for them being sorted alphabetically. Basically, all constants will be treated as equals by the sorting algorithm and it will not distinguish between, say, arith.constant and tf.Const. This is because multiple canonicalizations exist in various dialects that push the constants to the right but do not make any distinction among constants. So, since we want this utility to not clash with those canonicalizations, this is being done.

[mehdi_amini]

In D124750#3502228, @srishti-pm wrote:

In D124750#3500748, @mehdi_amini wrote:

Seems like this should be added to canonicalization? The "push constants to the right hand side" is there already.

I think this was not added to canonicalization because we wanted it to be an independent utility that can be used if needed and not be used if not needed.

You're telling me "what" while I'm actually more interested in the "why" here?

I also don't understand the complexity of the implementation, I may need an example to understand why you're recursively operating on the producer ops for the operands.
From the revision description: (1) the operands defined by non-constant-like ops come first, followed by (2) block arguments, and these are followed by (3) the operands defined by constant-like ops` which all seems like a fairly local check: a stable-sort on the operands would be deterministic and local to a single operation.

I do this because, firstly, in the description, if you look below this paragraph, you will see the following:
"And, if two operands come from the same op, the function backtracks and
looks even further to sort them. This backtracking is done over the
backward slice of the operand, in a breadth-first traversal."

Same as before: this does not tell me why, can you provide an example where this matters?

[srishti-pm]

In D124750#3502295, @mehdi_amini wrote:

You're telling me "what" while I'm actually more interested in the "why" here?

I'm not sure what your question is, with a "why". Let me think about this a bit. I'll get back to you.

Same as before: this does not tell me why, can you provide an example where this matters?

Sure. This is a bit lengthy. I'm really sorry for that !

So, lets start with some basic understanding here. Let's say I am writing a matchAndRewrite() function, where I take the following INPUT and convert it to the following OUTPUT:

INPUT:
a = div b, c
d = sub e, f
g = add d, a
h = const 0
i = mul h, g

OUTPUT:
i = some_op b, c, e, f

Now, when I'm writing a C++ code to match and rewrite:

If I only sort the i = mul h, g op, I get my canonicalized input as follows:

CANONICALIZED INPUT #1:

a = div b, c
d = sub e, f
g = add d, a
h = const 0
i = mul g, h

So, I'm basically sorting i = mul h, g to i = mul g, h using the utility and then writing the if-else statements to match the CANONICALIZED INPUT #1.
So, a psuedo code will be:

if mul.operand[0].defOp != add OR mul.operand[1].defOp != const 0
  return failure

if add.operand[0].defOp != sub
  if add.operand[0].defOp != div OR add.operand[1].defOp != sub
    return failure
  else
    get the values of b, c, e, and f
else if add.operand[1].defOp == div
  get the values of b, c, e, and f
else
  return failure

rewrite<some_op>(mul, b, c, e, f)

But, if I had sorted the producers as well, my canonicalized input would be:
CANONICALIZED INPUT #2:

a = div b, c
d = sub e, f
g = add a, d
h = const 0
i = mul g, h

and thus my code will reduce to:

if mul.operand[0].defOp !=  add OR mul.operand[1].defOp != const 0
  return failure

if add.operand[0].defOp != div OR add.operand[1].defOp != sub
  return failure

get the values of b, c, e, and f

rewrite<some_op>(mul, b, c, e, f)

So, in essence, we can see that the effort of an end user writing a C++ pattern has reduced if I sort the producers as well. But, one may argue that I could have sorted the add op after seeing it and then my if-else statements would reduce. So, the above illustration doesn't explain why we sort the producers.

The real reason for sorting the producers is that, if such a thing is not done, the sorting and this entire utility will be virtually useless. A deterministic sorting of an op requires its producers to be sorted. Our sorting algorithm is based on the breadth-first traversal of backard slices. On the same level of DAG, the traversal looks at operands from the first to the last. That is how the breadth-first traversal is defined here. Now, if this traversal is non-deterministic, then, the whole use of sorting something will collapse. Maybe, this can be BEST explained with the below example.

If I have this IR:
d = div b, c
s = sub e, f
x = xor k, l
g = add s, d
h = add d, x
i = add g, h

Then, i = add g, h will be sorted to i = add g, h (no change).

But, when I have the below IR (which is functionally the same as the above IR):
d = div b, c
s = sub e, f
x = xor k, l
g = add d, s
h = add d, x
i = add g, h

Then, i = add g, h will be sorted to i = add h, g.

So, we have two functionally same IRs being sorted differently. This is clearly not useful. The sorting depends on what the input IR is. So, even after sorting, we still have functionally same IRs that look different. So, the pattern matcher (a human) still has to write an excessive number of if-else statements to match the input. This is exactly what this sorting was supposed to avoid. This is as good as not having done any sorting at all!

Is the motivation clear now?

[srishti-pm]

Made the sorting strictly stable.

[mehdi_amini]

In D124750#3502401, @srishti-pm wrote:

In D124750#3502295, @mehdi_amini wrote:

You're telling me "what" while I'm actually more interested in the "why" here?

I'm not sure what your question is, with a "why". Let me think about this a bit. I'll get back to you.

My "why?" question is about canonicalization: could this be a canonicalization and if so why / why not? This is an important thing to answer before looking into the discussion below actually:

Same as before: this does not tell me why, can you provide an example where this matters?

Sure. This is a bit lengthy. I'm really sorry for that !

No worry, thanks for being thorough.

...
So, in essence, we can see that the effort of an end user writing a C++ pattern has reduced if I sort the producers as well.

So far you explained "what is canonicalization and why are we canonicalizing", the same rationale applies to "push constant to the right" that we do already in canonicalization, and this is exactly why I asked before whether we could do what you're doing as a canonicalization.

The real reason for sorting the producers is that, if such a thing is not done, the sorting and this entire utility will be virtually useless.

So: I understand that the producers should be sorted for a pattern to apply, but our disconnect earlier is that my usual approach to see canonicalization is to process an entire block/region, and as such we don't work on slices but on operation in order until fix-point. I'm a bit concerned about efficiency of your approach, because when integrated in a framework that actually visit the entire block you would recompute subset of the same slice over and over and re-attempt to sort everything multiple times:

d = add b, c
s = add d, f
g = add s, d
h = add d, x
i = add g, h

Every time the algorithm would consider a commutative op, that is all the op, it would recurse and try to re-sort the current slice, processing the same ops over and over.

A deterministic sorting of an op requires its producers to be sorted.

This isn't clear to me why based on the sorting criteria you provided, but in any case a local sort with fixed-point iteration on an isolated region should converge (hopefully).

Our sorting algorithm is based on the breadth-first traversal of backard slices. On the same level of DAG, the traversal looks at operands from the first to the last. That is how the breadth-first traversal is defined here. Now, if this traversal is non-deterministic, then, the whole use of sorting something will collapse. Maybe, this can be BEST explained with the below example.

If I have this IR:
d = div b, c
s = sub e, f
x = xor k, l
g = add s, d
h = add d, x
i = add g, h

Then, i = add g, h will be sorted to i = add g, h (no change).

But, when I have the below IR (which is functionally the same as the above IR):
d = div b, c
s = sub e, f
x = xor k, l
g = add d, s
h = add d, x
i = add g, h

Then, i = add g, h will be sorted to i = add h, g.

Why? Again your description of the sort is as follow:

(1) the operands defined by non-constant-like ops come first, followed by
(2) block arguments, and these are followed by
(3) the operands defined by constant-like ops.
In addition to this, within the category (1), the order of operands is alphabetical w.r.t. the dialect name and op name.

In your example:

g = add d, s
h = add d, x
i = add g, h

The operands fits category 1) I believe, and so they should be sorted "alphabetical w.r.t. the dialect name and op name", so a stable sort would never re-order them in any way.

[srishti-pm]

My "why?" question is about canonicalization: could this be a canonicalization and if so why / why not? This is an important thing to answer before looking into the discussion below actually:

I think, yes, it can be a canonicalization. I don't see why not.

[srishti-pm]

So far you explained "what is canonicalization and why are we canonicalizing", the same rationale applies to "push constant to the right" that we do already in canonicalization, and this is exactly why I asked before whether we could do what you're doing as a canonicalization.

You are right. And the answer is yes, we could do this as a canonicalization. In fact, it is a canonicalization because we are converting various forms of an IR to a specific canonicalized form. It just hasn't been added to any op's canonicalizer here.

[srishti-pm]

So: I understand that the producers should be sorted for a pattern to apply, but our disconnect earlier is that my usual approach to see canonicalization is to process an entire block/region, and as such we don't work on slices but on operation in order until fix-point. I'm a bit concerned about efficiency of your approach, because when integrated in a framework that actually visit the entire block you would recompute subset of the same slice over and over and re-attempt to sort everything multiple times:

I completely agree with your concern. I had the same concern myself (refer to discussion starting from here: https://discourse.llvm.org/t/mlir-pdl-extending-pdl-pdlinterp-bytecode-to-enable-commutative-matching/60798/12?u=srishtisrivastava). I think we can find a way to call this utility in the canonicalizers of all the commutative ops and remove the recursion (which won't be needed anymore, given the canonicalization happens from top to bottom). What are your views on this?

[mehdi_amini]

That was the sense of my question about canonicalization indeed :)

[srishti-pm]

Every time the algorithm would consider a commutative op, that is all the op, it would recurse and try to re-sort the current slice, processing the same ops over and over.

Yes, exactly.

[srishti-pm]

Why? Again your description of the sort is as follow:

I think I need to update my commit summary and revision summary. I had thought that this was an intuitive way of explaining the utility. But, I understand that it is quite misleading.

[Mogball]

I need to look at the algorithm in more detail, but I'm not a fan of using a string key. Concatenating strings to make compound keys is not very efficient and potentially brittle. Can you assign unique IDs and use an array of IDs instead?

mlir/include/mlir/Transforms/CommutativityUtils.h

25

Add documentation? Similar to what you have in the patch description.

[Mogball]

On the matter of whether this should be a canonicalization, my concern with this is that if an operation has its own preferred ordering of operands that conflicts with the sort, then this will cause canonicalization to loop infinitely.

It's not actually the canonicalizer pass that moves constants to the right hand size. It's the folder. And it probably shouldn't be the folder that does this. So I'm open to making this part of canonicalization IF the sorted operand order produced by this utility is the canonical order of operands for commutative operations, so that conflicts are not possible.

[jpienaar]

(1) the operands defined by non-constant-like ops come first, followed by (2) block arguments, and these are followed by (3) the operands defined by constant-like ops.

I would have thought block-arguments would come first as we don't know their values, while non-constant-like ops could be folded at some point and then become constant-like. Meaning, they seem closer to constant than block arguments.

+1 to Mehdi's question about just stable sorting based on based on 4 criteria (3 buckets + ordering within (1)) and then we should be able to avoid all the string mangling too as Jeff asked about.

[rriddle]

+1 on all of the other comments, especially related to the use of strings.

In D124750#3503607, @Mogball wrote:

On the matter of whether this should be a canonicalization, my concern with this is that if an operation has its own preferred ordering of operands that conflicts with the sort, then this will cause canonicalization to loop infinitely.

It's not actually the canonicalizer pass that moves constants to the right hand size. It's the folder. And it probably shouldn't be the folder that does this. So I'm open to making this part of canonicalization IF the sorted operand order produced by this utility is the canonical order of operands for commutative operations, so that conflicts are not possible.

We can decide whatever we want the canonical ordering of operands to be for the Commutative trait. We don't have to leave things up to operations if it doesn't make sense.

[srishti-pm]

Improving the documentation of the functionality of this utility to make it more clear.

[srishti-pm]

Fixing a minor typo.

[mehdi_amini]

Thanks for improving the doc! Are you moving this to be used in canonicalization next?
I think a first good step would be to make it a pattern and test it with a pass that applies it in the greedy rewriter. I would land this first and then try to enable this in the canonicalized.

Also, have you thought already about how to get rid of string manipulation?

[srishti-pm]

Thanks for improving the doc! Are you moving this to be used in canonicalization next?
I think a first good step would be to make it a pattern and test it with a pass that applies it in the greedy rewriter. I would land this first and then try to enable this in the canonicalized.

Thanks! I hope it is unambiguous and clear now! Yes, I will do the "moving this to be used in canonicalization" now. Thanks for the suggestion. The plan sounds good.

Also, have you thought already about how to get rid of string manipulation?

No, I haven't. I'm still thinking about this. Meanwhile, I'll also address the other comments given here.

[srishti-pm]

Replying to the various comments given in this revision:

1. Regarding string manipulation:

I need to look at the algorithm in more detail, but I'm not a fan of using a string key. Concatenating strings to make compound keys is not very efficient and potentially brittle. Can you assign unique IDs and use an array of IDs instead?

Sure, I'm currently brainstorming to come up with a way to do this. But, @Mogball, do you remember our discussion of including the attibutes, etc. also in the unique ID/key (https://discourse.llvm.org/t/mlir-pdl-extending-pdl-pdlinterp-bytecode-to-enable-commutative-matching/60798/13)? We had decided that this was something that can be added if and when required but probably won't be included in the first revision of this utility. As in, the first revision will contain a TODO comment describing this enhancement. Will we still want to do this when we are using an array of unique IDs? I wish to understand this so that I know if I have to think of an ID which can be easily extended to include the attributes, etc. of an op.

+1 to Mehdi's question about just stable sorting based on based on 4 criteria (3 buckets + ordering within (1)) and then we should be able to avoid all the string mangling too as Jeff asked about.

Actually, that question of Mehdi is now obsolete (I think). It was based on my incorrect documentation of the "sorting rule". I have now corrected the documentation of the rule (in the revision summary, the commit summary, and the code comments). The rule was finalized in a recent RFC (refer to comments starting from here: https://discourse.llvm.org/t/mlir-pdl-extending-pdl-pdlinterp-bytecode-to-enable-commutative-matching/60798/19). Does this sound okay?

2. Regarding reversing non-constant-like op and block argument order:

I would have thought block-arguments would come first as we don't know their values, while non-constant-like ops could be folded at some point and then become constant-like. Meaning, they seem closer to constant than block arguments.

Sure. This sounds much better actually. I think at some point during the implementation, I had also thought of the same :)
I'll do this. Thanks for the suggestion.

[Mogball]

I'm open to iterating in tree. Landing this utility first and then try adding it as a canonicalization SGTM.

The string could be replaced with an array of tuples (yuck) for now. An enum (constant, block arg, everything else) plus the OperationName.

Attributes need to be handled somehow possibly by adding the op's dictionary attribute and using a deep compare.

[srishti-pm]

@Mogball @mehdi_amini @jpienaar, sorry there haven't been any updates from my side here for the past 10 or so days. I had been busy in some other tasks. I have started working on this again.

[srishti-pm]

Addressed all the comments.

[Mogball]

I haven't thought too hard about the algorithm itself yet. I'm in the camp of "let's move forward if it works". I have mostly trivial comments.

clang/docs/tools/clang-formatted-files.txt

8451 ↗

(On Diff #436274)

I don't think this file needs to be modified.

mlir/include/mlir/Transforms/CommutativityUtils.h

25

Why do all of these need to be exposed publicly? I think this file should only contain SortCommutativeOperands.

35

using BlockArgumentOrOpKey = std::pair<BlockArgumentOrOpType, StringRef>

The default operator< for std::pair should work for you.

141

ArrayRef<OperandBFS *>

162

Pass these both by ArrayRef

175

ArrayRef<OperandBFS *>

187

ArrayRef<OperandBFS *>

195

SmallVectorImpl

202

ArrayRef<OperandBFS *>

212

Please move the body matchAndRewrite into a source file. It only needs Operation *. And then all the structs/enum/utility functions in the header can be moved there as well.

227

memory leak?

243

sortedOperands(numOperands, nullptr);

276

297

mlir/lib/Transforms/Utils/CommutativityUtils.cpp

18

unused?

52

The doc of a public function shouldn't be repeated above the implementation.

[Mogball]

mlir/include/mlir/Transforms/CommutativityUtils.h

212

This could stand to be a static_assert

258

Could you not change getIndicesOfUnassigned... to populate two lists of operands and pass these to assignSortedPositionTo instead of using a set to track the indices. You could put the operand index inside OperandBFS to keep track.

mlir/lib/Transforms/Utils/CommutativityUtils.cpp

25

This function doesn't seem like it pays for itself -- llvm::any_of?

93

These could all be ArrayRefs since you aren't modifying them

168

drop mlir::

202

This function doesn't seem like it pays for itself.

[srishti-pm]

Minor changes.

[mehdi_amini]

mlir/include/mlir/Transforms/CommutativityUtils.h

212

Can we make this not a template? This will be a code bloat otherwise.

333

There seems to me to be far too much code in the public head: I can't isolate the public API here, if this is just about a pattern then the usual way is to have a "populate" method.

[srishti-pm]

Addressing comments.

[srishti-pm]

Addressed most of the comments. A few remaining.

[srishti-pm]

Increasing pattern benefit + minor typo correction.

[mehdi_amini]

Right now I'm not yet understanding all of the algorithm (haven't spent enough time on it), but I'm mostly concerned by the runtime cost of this normalization.

mlir/lib/Transforms/Utils/CommutativityUtils.cpp

372

378–380

[srishti-pm]

In D124750#3587161, @mehdi_amini wrote:

Right now I'm not yet understanding all of the algorithm (haven't spent enough time on it), but I'm mostly concerned by the runtime cost of this normalization.

I understand your concern. I'll go through my implementation once and optimize things that need it.

[Mogball]

In principle, I think the algorithm is fine. I'm pretty sure you can rewrite bits of it to get rid of the map/sets. I'm only concerned about handling attributes. (e.g. cmpi slt vs cmpi sgt)

[srishti-pm]

Addressed the final comments.

[srishti-pm]

@mehdi_amini, I have made sure that the algorithm is good in terms of both time and space complexity.

@Mogball, "handling attributes (e.g. cmpi slt vs cmpi sgt)" doesn't seem hard to me. I think this algorithm can be extended with ease to take attributes into account. But, I don't think that it should be a part of this revision (I believe you agree) because it seems like an incremental change which should be added only when the need arises. A new user should first get accustomed to this utility (and how its sorts operands with different backward slices) and then the utility can be extended to differentiate between backward slices containing ops with the same name but different attribute values.

mlir/include/mlir/Transforms/CommutativityUtils.h

35

I have added a constructor to BlockArgumentOrOpKey (now renamed to AncestorKey) and thus I think this comment is obsolete now. Hope that this fine. Adding a constructor made the code look cleaner.

227

Fixed this by adding a struct called "Ancestor" which refers to either an op or a block argument.

258

I think that doing this might not be a good idea. It will increase the space complexity unnecessarily (OperandBFS is a big structure) and not help much with the time complexity because the sets of the indices are expected to be small. At least the number of indices will be <= the total number of operands and each element in these sets will occupy very less space (size of unsigned).

mlir/lib/Transforms/Utils/CommutativityUtils.cpp

272

Fixed this by adding a struct called "Ancestor" which refers to either an op or a block argument.

[srishti-pm]

Fixed memory leak.

[Mogball]

mlir/lib/Transforms/Utils/CommutativityUtils.cpp

179

Please remove these. They don't improve readability.

[srishti-pm]

Addressed of all Jeff's comments.

[Mogball]

I'm glad the DenseSets are gone, but my three-ish biggest gripes are:

• The algorithm is conceptually simple, but there is way more code than is necessary to achieve it.
• More comments (excluding "doc" comments) than code is generally not a good sign
• The implementation is still inefficient in a lot of obvious ways.

mlir/lib/Transforms/Utils/CommutativityUtils.cpp

11–12

22–39

This class isn't necessary. Ancestor can just be Operation *. If it's null, then we know it's a block argument (the bool flag is redundant).

41

66–67

My biggest complaint with respect to readability is that there are more comments than code. This is fine if the comment has a big explanation about the algorithm and how keys are represented, especially with a nice ASCII diagram as you have below. But if this constructor had 0 comments except maybe "Only non-constant ops are sorted by name", it would be perfectly fine.

77

Constant ops could be sorted by name as well.

89–97

This should behave the same as the manually written comparator. operator< of std::tuple compares the first element and then the next if the first is equal.

111

Since this is a set of pointers expected to be small, you can use SmallPtrSet for a big efficiency boost (linear scan instead of hashing when small).

149–150

Since you are moving sorted operands into their sorted position and tracking the unsorted range, you shouldn't even need this flag because you will always know the sorted and unsorted subranges.

There are multiple loops in which you iterate over the entire operand list but skip those where this flag is set/unset. In those cases, you can always just iterate the subrange of interest.

156–160

163–178

I would drop these helpers. std::queue will already assert if pop or front are called on an empty queue.

201–218

268

270

This flag is not necessary because you can just check bfsOfOperandsWithKey.size() == 1

278

279–280

You don't need to make a copy. In fact, I think you should just track the indices.

299–305

This shouldn't be necessary if you're tracking the unsorted subrange. Just quit when it gets to 1 element.

308–309

313

Argument names in comments are not necessary when the passed variable has a descriptive name.

330–354

There is no way you need this much code. A std::swap between the current operand and the first unsorted position should be enough.

361

This is possibly the longest function name I've ever seen. Please make it more concise.

401

This check could be moved into pushAncestor

536–539

You could just be returning a flag to indicate whether any swapping occurred so that you don't have to track before and after.

[Mogball]

mlir/lib/Transforms/Utils/CommutativityUtils.cpp

524

And then you'll never need to check isSorted again.

[srishti-pm]

mlir/lib/Transforms/Utils/CommutativityUtils.cpp

77

The only reason we separated constant ops from the non-constant ops was because the former are canonicalized to the right (as a stable sort) by existing canonicalizations. And, we didn't want our algorithm to conflict with these existing canonicalizations. That is the reason I am not sorting them by name and just keeping them to the right (as a stable sort).

270

.size() is an O(N) operation and that is why I usually try to avoid it. Do you still agree we should use it here? I understand that N is an expectedly small value.

279–280

I agree. But, we had discussed to store operands instead of indices and that's why I did this. I will change this to use indices again (keeping other things unchanged).

330–354

If I do a swap, the sorting will no longer be stable and I believe that there was a discussion that concluded with the fact that "we want stable sorting".

361

Could you please give a suggestion for the name? After a long thought, I came up with this name. It was better than all my other ideas.

[Mogball]

mlir/lib/Transforms/Utils/CommutativityUtils.cpp

77

I know. You can sort them separately from regular ops and also by name and the overall behaviour would be the same.

270

Size is constant time

279–280

I mean you can have list (not a set) of indices to shift

330–354

That's true, but shifting like this is very slow as well. At this point, you might want to give std::stable_sort with a custom comparator that does extra BFS iterations on demand a try.

[srishti-pm]

mlir/lib/Transforms/Utils/CommutativityUtils.cpp

330–354

So, this is what I think:

The number of commutative operands is not expected to be huge. So, we can afford to do shifting. In most cases, we wouldn't have to shift more than 1 or 2 positions. But, the custom comparator might cost us a lot, considering that each BFS could potentially be very large, especially for deep learning models. So, doing the BFS traversals again and again for the same operand, even though caching will be involved, doesn't sound like a good idea to me.

What are your views?

[Mogball]

mlir/lib/Transforms/Utils/CommutativityUtils.cpp

330–354

Very rough estimate: on its own, this function is N. Finding the smallest key is N, and then finding all matching elements is N. This function is called for each operand that needs to be moved, but the number of such operands decreases. So the sort itself averages out to be 3N^2 iterations over the operand list.

Now for traversals, doing BFS on demand inside the comparator doesn't mean it has to restart every time. It would do extra iterations on top of existing iteration results only when needed to break ties. In your case, you do an extra iteration of BFS for all operands if the current smallest key is identical, not just for the ones needed. It's hard to estimate the number of iterations of BFS, but certainly it's more in your case. Using std::stable_sort would also bring the complexity down to N logN

[mehdi_amini]

For this to be a usable canonicalization, it is really the case where the operands are already sorted (no-op) that needs to be heavily optimized (that is no complex data structure to populate, etc.)

[Mogball]

I'm not sure how to check the no-op case without constructing at least a queue. Even assuming only 2 commutative operands, if they look the same depth=1, then the comparison has to iterate.

[srishti-pm]

Used the stable_sort function.

[Mogball]

mlir/lib/Transforms/Utils/CommutativityUtils.cpp

238

Why can't you sort the OperandBFS directly to avoid the hash map?

286

Can you use unique_ptr so that the memory doesn't leak?

[srishti-pm]

mlir/lib/Transforms/Utils/CommutativityUtils.cpp

238

Because comparators are not allowed to modify their input arguments.

[Mogball]

mlir/lib/Transforms/Utils/CommutativityUtils.cpp

238

The arguments have to be const references? If so, just const_cast

[srishti-pm]

mlir/lib/Transforms/Utils/CommutativityUtils.cpp

111

This set isn't expected to be small, right? There can be many ancestors. The only thing that is small in this context is the number of operands, right?

[srishti-pm]

1. Addressed all the comments.
2. Refactored the code.

[srishti-pm]

Updated an outdated comment.

[Mogball]

Got a few small nits but otherwise LGTM. Thanks for all the hard work! This looks really solid now. I haven't thought too hard about the performance of that while loop but it seems good enough to land for now.

mlir/lib/Transforms/Utils/CommutativityUtils.cpp

247

253

256

[mehdi_amini]

. I haven't thought too hard about the performance of that while loop but it seems good enough to land for now.

What's the finality of it? That is: outside of canonicalization what is its purpose?

[Mogball]

I'm referring to the nitty gritty details of the while loop inside the comparator. It looks pretty tight to me right now. If the operands are already sorted, worst-case each operand is compared against only its neighbours. Unfortunately, without extra bookkeeping, BFS will still need to iterate when necessary

[srishti-pm]

Addressed the final NITs.

mlir/include/mlir/Transforms/CommutativityUtils.h

+//===- CommutativityUtils.h - Commutativity utilities -----------*- C++ -*-===//
+//
+// 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 header file declares a function to populate the commutativity utility
+// pattern. This function is intended to be used inside passes to simplify the
+// matching of commutative operations by fixing the order of their operands.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TRANSFORMS_COMMUTATIVITYUTILS_H
+#define MLIR_TRANSFORMS_COMMUTATIVITYUTILS_H
+
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+
+/// Populates the commutativity utility patterns.
+void populateCommutativityUtilsPatterns(RewritePatternSet &patterns);
+
+} // namespace mlir
+
+#endif // MLIR_TRANSFORMS_COMMUTATIVITYUTILS_H

mlir/lib/Transforms/Utils/CMakeLists.txt

 add_mlir_library(MLIRTransformUtils
+  CommutativityUtils.cpp
   ControlFlowSinkUtils.cpp
   DialectConversion.cpp
   FoldUtils.cpp
   GreedyPatternRewriteDriver.cpp
   InliningUtils.cpp
   LoopInvariantCodeMotionUtils.cpp
   RegionUtils.cpp
   SideEffectUtils.cpp

mlir/lib/Transforms/Utils/CommutativityUtils.cpp

+//===- CommutativityUtils.cpp - Commutativity utilities ---------*- C++ -*-===//
+//
+// 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 a commutativity utility pattern and a function to
+// populate this pattern. The function is intended to be used inside passes to
+// simplify the matching of commutative operations by fixing the order of their
+// operands.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Transforms/CommutativityUtils.h"
+
+#include <queue>
+
+using namespace mlir;
+
+/// The possible "types" of ancestors. Here, an ancestor is an op or a block
+/// argument present in the backward slice of a value.
+enum AncestorType {
+  /// Pertains to a block argument.
+  BLOCK_ARGUMENT,
+
+  /// Pertains to a non-constant-like op.
+  NON_CONSTANT_OP,
+
+  /// Pertains to a constant-like op.
+  CONSTANT_OP
+};
+
+/// Stores the "key" associated with an ancestor.
+struct AncestorKey {
+  /// Holds `BLOCK_ARGUMENT`, `NON_CONSTANT_OP`, or `CONSTANT_OP`, depending on
+  /// the ancestor.
+  AncestorType type;
+
+  /// Holds the op name of the ancestor if its `type` is `NON_CONSTANT_OP`.
+  /// Else, holds "".
+  StringRef opName;
+
+  /// Constructor for `AncestorKey`.
+  AncestorKey(Operation *op) {
+    if (!op) {
+      type = BLOCK_ARGUMENT;
+    } else if (!op->hasTrait<OpTrait::ConstantLike>()) {
+      type = NON_CONSTANT_OP;
+      opName = op->getName().getStringRef();
+    } else {
+      type = CONSTANT_OP;
+    }
+  }
+
+  /// Overloaded operator `<` for `AncestorKey`.
+  ///
+  /// AncestorKeys of type `BLOCK_ARGUMENT` are considered the smallest and
+  /// those of type `CONSTANT_OP`, the largest. `NON_CONSTANT_OP` types come in
+  /// between, with the smaller ones being the ones with smaller op names
+  /// (lexicographically).
+  bool operator<(const AncestorKey &key) const {
+    if ((type == BLOCK_ARGUMENT && key.type != BLOCK_ARGUMENT) ||
+        (type == NON_CONSTANT_OP && key.type == CONSTANT_OP))
+      return true;
+    if ((key.type == BLOCK_ARGUMENT && type != BLOCK_ARGUMENT) ||
+        (key.type == NON_CONSTANT_OP && type == CONSTANT_OP))
+      return false;
+    return opName < key.opName;
+  }
+};
+
+/// Stores the BFS traversal information of an operand.
+struct OperandBFS {
+  /// Stores the queue of ancestors of the operand's BFS traversal at a
+  /// particular point in time.
+  std::queue<Operation *> ancestorQueue;
+
+  /// Stores the list of ancestors that have been visited by the BFS traversal
+  /// at a particular point in time.
+  DenseSet<Operation *> visitedAncestors;
+
+  /// Stores the operand's "key". This "key" is defined as a list of the
+  /// "AncestorKeys" associated with the ancestors of this operand, in a
+  /// breadth-first order.
+  ///
+  /// So, if an operand, say `A`, was produced as follows:
+  ///
+  /// `<block argument>`  `<block argument>`
+  ///             \          /
+  ///              \        /
+  ///             `arith.subi`           `arith.constant`
+  ///                       \            /
+  ///                        `arith.addi`
+  ///                              |
+  ///                         returns `A`
+  ///
+  /// Then, the ancestors of `A`, in the breadth-first order are:
+  /// `arith.addi`, `arith.subi`, `arith.constant`, `<block argument>`, and
+  /// `<block argument>`.
+  ///
+  /// Thus, the "key" associated with operand `A` is:
+  /// {
+  ///  {type: `NON_CONSTANT_OP`, opName: "arith.addi"},
+  ///  {type: `NON_CONSTANT_OP`, opName: "arith.subi"},
+  ///  {type: `CONSTANT_OP`, opName: ""},
+  ///  {type: `BLOCK_ARGUMENT`, opName: ""},
+  ///  {type: `BLOCK_ARGUMENT`, opName: ""}
+  /// }
+  SmallVector<AncestorKey, 4> key;
+
+  /// Push an ancestor into the operand's BFS information structure. This
+  /// entails it being pushed into the queue (always) and inserted into the
+  /// "visited ancestors" list (iff it is an op rather than a block argument).
+  void pushAncestor(Operation *op) {
+    ancestorQueue.push(op);
+    if (op)
+      visitedAncestors.insert(op);
+    return;
+  }
+
+  /// Refresh the key.
+  ///
+  /// Refreshing a key entails making it up-to-date with the operand's BFS
+  /// traversal that has happened till that point in time, i.e, appending the
+  /// existing key with the front ancestor's "AncestorKey". Note that a key
+  /// directly reflects the BFS and thus needs to be refreshed during the
+  /// progression of the traversal.
+  void refreshKey() {
+    if (ancestorQueue.empty())
+      return;
+
+    Operation *frontAncestor = ancestorQueue.front();
+    AncestorKey frontAncestorKey(frontAncestor);
+    key.push_back(frontAncestorKey);
+    return;
+  }
+
+  /// Pop the front ancestor, if any, from the queue and then push its adjacent
+  /// unvisited ancestors, if any, to the queue (this is the main body of the
+  /// BFS algorithm).
+  void popFrontAndPushAdjacentUnvisitedAncestors() {
+    if (ancestorQueue.empty())
+      return;
+    Operation *frontAncestor = ancestorQueue.front();
+    ancestorQueue.pop();
+    if (!frontAncestor)
+      return;
+    for (Value operand : frontAncestor->getOperands()) {
+      Operation *operandDefOp = operand.getDefiningOp();
+      if (!operandDefOp || !visitedAncestors.contains(operandDefOp))
+        pushAncestor(operandDefOp);
+    }
+    return;
+  }
+};
+
+/// Sorts the operands of `op` in ascending order of the "key" associated with
+/// each operand iff `op` is commutative.
+///
+/// After the application of this pattern, since the commutative operands now
+/// have a deterministic order in which they occur in an op, the matching of
+/// large DAGs becomes much simpler, i.e., requires much less number of checks
+/// to be written by a user in her/his pattern matching function.
+///
+/// Some examples of such a sorting:
+///
+/// Assume that the sorting is being applied to `foo.commutative`, which is a
+/// commutative op.
+///
+/// Example 1:
+///
+/// %1 = foo.const 0
+/// %2 = foo.mul <block argument>, <block argument>
+/// %3 = foo.commutative %1, %2
+///
+/// Here,
+/// 1. The key associated with %1 is:
+///     `{
+///       {CONSTANT_OP, ""}
+///      }`
+/// 2. The key associated with %2 is:
+///     `{
+///       {NON_CONSTANT_OP, "foo.mul"},
+///       {BLOCK_ARGUMENT, ""},
+///       {BLOCK_ARGUMENT, ""}
+///      }`
+///
+/// The key of %2 < the key of %1
+/// Thus, the sorted `foo.commutative` is:
+/// %3 = foo.commutative %2, %1
+///
+/// Example 2:
+///
+/// %1 = foo.const 0
+/// %2 = foo.mul <block argument>, <block argument>
+/// %3 = foo.mul %2, %1
+/// %4 = foo.add %2, %1
+/// %5 = foo.commutative %1, %2, %3, %4
+///
+/// Here,
+/// 1. The key associated with %1 is:
+///     `{
+///       {CONSTANT_OP, ""}
+///      }`
+/// 2. The key associated with %2 is:
+///     `{
+///       {NON_CONSTANT_OP, "foo.mul"},
+///       {BLOCK_ARGUMENT, ""}
+///      }`
+/// 3. The key associated with %3 is:
+///     `{
+///       {NON_CONSTANT_OP, "foo.mul"},
+///       {NON_CONSTANT_OP, "foo.mul"},
+///       {CONSTANT_OP, ""},
+///       {BLOCK_ARGUMENT, ""},
+///       {BLOCK_ARGUMENT, ""}
+///      }`
+/// 4. The key associated with %4 is:
+///     `{
+///       {NON_CONSTANT_OP, "foo.add"},
+///       {NON_CONSTANT_OP, "foo.mul"},
+///       {CONSTANT_OP, ""},
+///       {BLOCK_ARGUMENT, ""},
+///       {BLOCK_ARGUMENT, ""}
+///      }`
+///
+/// Thus, the sorted `foo.commutative` is:
+/// %5 = foo.commutative %4, %3, %2, %1
+class SortCommutativeOperands : public RewritePattern {
+public:
+  SortCommutativeOperands(MLIRContext *context)
+      : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/5, context) {}
+  LogicalResult matchAndRewrite(Operation *op,
+                                PatternRewriter &rewriter) const override {
+    // Stores the mapping between an operand and its BFS traversal information.
+    DenseMap<Value, OperandBFS *> operandToItsBFSMap;
+
+    // Custom comparator for two commutative operands, which returns true iff
+    // the "key" of `operandA` < the "key" of `operandB`, i.e.,
+    // 1. In the first unequal pair of corresponding AncestorKeys, the
+    // AncestorKey in `operandA` is smaller, or,
+    // 2. Both the AncestorKeys in every pair are the same and the size of
+    // `operandA`'s "key" is smaller.
+    auto CommutativeOperandComparator = [&](Value operandA, Value operandB) {
+      if (operandA == operandB)
+        return false;
+
+      // Iteratively perform the BFS's of both operands until an order among
+      // them can be determined.
+      OperandBFS *bfsA = operandToItsBFSMap.lookup(operandA);
+      OperandBFS *bfsB = operandToItsBFSMap.lookup(operandB);
+      unsigned keyIndex = 0;
+      while (true) {
+        if (bfsA->key.size() <= keyIndex) {
+          if (bfsA->ancestorQueue.empty())
+            return true;
+          bfsA->popFrontAndPushAdjacentUnvisitedAncestors();
+          bfsA->refreshKey();
+        }
+        if (bfsB->key.size() <= keyIndex) {
+          if (bfsB->ancestorQueue.empty())
+            return false;
+          bfsB->popFrontAndPushAdjacentUnvisitedAncestors();
+          bfsB->refreshKey();
+        }
+        if (bfsA->ancestorQueue.empty() || bfsB->ancestorQueue.empty())
+          return bfsA->key.size() < bfsB->key.size();
+        if (bfsA->key[keyIndex] < bfsB->key[keyIndex])
+          return true;
+        if (bfsB->key[keyIndex] < bfsA->key[keyIndex])
+          return false;
+        keyIndex++;
+      }
+    };
+
+    // If `op` is not commutative, do nothing.
+    if (!op->hasTrait<OpTrait::IsCommutative>())
+      return failure();
+
+    // Populate the mapping between the operands and their BFS traversal
+    // information.
+    SmallVector<Value, 2> operands = op->getOperands();
+    for (unsigned i = 0, e = op->getNumOperands(); i < e; i++) {
+      OperandBFS *bfs = new OperandBFS();
+      bfs->pushAncestor(operands[i].getDefiningOp());
+      bfs->refreshKey();
+      operandToItsBFSMap.insert({operands[i], bfs});
+    }
+
+    // Sort the operands.
+    SmallVector<Value, 2> sortedOperands = operands;
+    std::stable_sort(sortedOperands.begin(), sortedOperands.end(),
+                     CommutativeOperandComparator);
+    if (sortedOperands == operands)
+      return failure();
+    rewriter.updateRootInPlace(op, [&] { op->setOperands(sortedOperands); });
+    return success();
+  }
+};
+
+void mlir::populateCommutativityUtilsPatterns(RewritePatternSet &patterns) {
+  patterns.add<SortCommutativeOperands>(patterns.getContext());
+}

mlir/test/Transforms/test-commutativity-utils.mlir

+// RUN: mlir-opt %s -test-commutativity-utils | FileCheck %s
+
+// CHECK-LABEL: @test_small_pattern_1
+func.func @test_small_pattern_1(%arg0 : i32) -> i32 {
+  // CHECK-NEXT: %[[ARITH_CONST:.*]] = arith.constant
+  %0 = arith.constant 45 : i32
+
+  // CHECK-NEXT: %[[TEST_ADD:.*]] = "test.addi"
+  %1 = "test.addi"(%arg0, %arg0): (i32, i32) -> i32
+
+  // CHECK-NEXT: %[[ARITH_ADD:.*]] = arith.addi
+  %2 = arith.addi %arg0, %arg0 : i32
+
+  // CHECK-NEXT: %[[ARITH_MUL:.*]] = arith.muli
+  %3 = arith.muli %arg0, %arg0 : i32
+
+  // CHECK-NEXT: %[[RESULT:.*]] = "test.op_commutative"(%[[ARITH_ADD]], %[[ARITH_MUL]], %[[TEST_ADD]], %[[ARITH_CONST]])
+  %result = "test.op_commutative"(%0, %1, %2, %3): (i32, i32, i32, i32) -> i32
+
+  // CHECK-NEXT: return %[[RESULT]]
+  return %result : i32
+}
+
+// CHECK-LABEL: @test_small_pattern_2
+// CHECK-SAME: (%[[ARG0:.*]]: i32
+func.func @test_small_pattern_2(%arg0 : i32) -> i32 {
+  // CHECK-NEXT: %[[TEST_CONST:.*]] = "test.constant"
+  %0 = "test.constant"() {value = 0 : i32} : () -> i32
+
+  // CHECK-NEXT: %[[ARITH_CONST:.*]] = arith.constant
+  %1 = arith.constant 0 : i32
+
+  // CHECK-NEXT: %[[ARITH_ADD:.*]] = arith.addi
+  %2 = arith.addi %arg0, %arg0 : i32
+
+  // CHECK-NEXT: %[[RESULT:.*]] = "test.op_commutative"(%[[ARG0]], %[[ARITH_ADD]], %[[TEST_CONST]], %[[ARITH_CONST]])
+  %result = "test.op_commutative"(%0, %1, %2, %arg0): (i32, i32, i32, i32) -> i32
+
+  // CHECK-NEXT: return %[[RESULT]]
+  return %result : i32
+}
+
+// CHECK-LABEL: @test_large_pattern
+func.func @test_large_pattern(%arg0 : i32, %arg1 : i32) -> i32 {
+  // CHECK-NEXT: arith.divsi
+  %0 = arith.divsi %arg0, %arg1 : i32
+
+  // CHECK-NEXT: arith.divsi
+  %1 = arith.divsi %0, %arg0 : i32
+
+  // CHECK-NEXT: arith.divsi
+  %2 = arith.divsi %1, %arg1 : i32
+
+  // CHECK-NEXT: arith.addi
+  %3 = arith.addi %1, %arg1 : i32
+
+  // CHECK-NEXT: arith.subi
+  %4 = arith.subi %2, %3 : i32
+
+  // CHECK-NEXT: "test.addi"
+  %5 = "test.addi"(%arg0, %arg0): (i32, i32) -> i32
+
+  // CHECK-NEXT: %[[VAL6:.*]] = arith.divsi
+  %6 = arith.divsi %4, %5 : i32
+
+  // CHECK-NEXT: arith.divsi
+  %7 = arith.divsi %1, %arg1 : i32
+
+  // CHECK-NEXT: %[[VAL8:.*]] = arith.muli
+  %8 = arith.muli %1, %arg1 : i32
+
+  // CHECK-NEXT: %[[VAL9:.*]] = arith.subi
+  %9 = arith.subi %7, %8 : i32
+
+  // CHECK-NEXT: "test.addi"
+  %10 = "test.addi"(%arg0, %arg0): (i32, i32) -> i32
+
+  // CHECK-NEXT: %[[VAL11:.*]] = arith.divsi
+  %11 = arith.divsi %9, %10 : i32
+
+  // CHECK-NEXT: %[[VAL12:.*]] = arith.divsi
+  %12 = arith.divsi %6, %arg1 : i32
+
+  // CHECK-NEXT: arith.subi
+  %13 = arith.subi %arg1, %arg0 : i32
+
+  // CHECK-NEXT: "test.op_commutative"(%[[VAL12]], %[[VAL12]], %[[VAL8]], %[[VAL9]])
+  %14 = "test.op_commutative"(%12, %9, %12, %8): (i32, i32, i32, i32) -> i32
+
+  // CHECK-NEXT: %[[VAL15:.*]] = arith.divsi
+  %15 = arith.divsi %13, %14 : i32
+
+  // CHECK-NEXT: %[[VAL16:.*]] = arith.addi
+  %16 = arith.addi %2, %15 : i32
+
+  // CHECK-NEXT: arith.subi
+  %17 = arith.subi %16, %arg1 : i32
+
+  // CHECK-NEXT: "test.addi"
+  %18 = "test.addi"(%arg0, %arg0): (i32, i32) -> i32
+
+  // CHECK-NEXT: %[[VAL19:.*]] = arith.divsi
+  %19 = arith.divsi %17, %18 : i32
+
+  // CHECK-NEXT: "test.addi"
+  %20 = "test.addi"(%arg0, %16): (i32, i32) -> i32
+
+  // CHECK-NEXT: %[[VAL21:.*]] = arith.divsi
+  %21 = arith.divsi %17, %20 : i32
+
+  // CHECK-NEXT: %[[RESULT:.*]] = "test.op_large_commutative"(%[[VAL16]], %[[VAL19]], %[[VAL19]], %[[VAL21]], %[[VAL6]], %[[VAL11]], %[[VAL15]])
+  %result = "test.op_large_commutative"(%16, %6, %11, %15, %19, %21, %19): (i32, i32, i32, i32, i32, i32, i32) -> i32
+
+  // CHECK-NEXT: return %[[RESULT]]
+  return %result : i32
+}

mlir/test/lib/Dialect/Test/TestOps.td

 }
 
 def TestOpWithVariadicResultsAndFolder: TEST_Op<"op_with_variadic_results_and_folder"> {
   let arguments = (ins Variadic<I32>);
   let results = (outs Variadic<I32>);
   let hasFolder = 1;
 }
 
+def TestAddIOp : TEST_Op<"addi"> {
+  let arguments = (ins I32:$op1, I32:$op2);
+  let results = (outs I32);
+}
+
 def TestCommutativeOp : TEST_Op<"op_commutative", [Commutative]> {
   let arguments = (ins I32:$op1, I32:$op2, I32:$op3, I32:$op4);
   let results = (outs I32);
 }
 
+def TestLargeCommutativeOp : TEST_Op<"op_large_commutative", [Commutative]> {
+  let arguments = (ins I32:$op1, I32:$op2, I32:$op3, I32:$op4, I32:$op5, I32:$op6, I32:$op7);
+  let results = (outs I32);
+}
+
 def TestCommutative2Op : TEST_Op<"op_commutative2", [Commutative]> {
   let arguments = (ins I32:$op1, I32:$op2);
   let results = (outs I32);
 }
 
 def TestIdempotentTraitOp
  : TEST_Op<"op_idempotent_trait",
            [SameOperandsAndResultType, NoSideEffect, Idempotent]> {

mlir/test/lib/Transforms/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestTransforms
+  TestCommutativityUtils.cpp
   TestConstantFold.cpp
   TestControlFlowSink.cpp
   TestInlining.cpp
   TestIntRangeInference.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   ADDITIONAL_HEADER_DIRS

mlir/test/lib/Transforms/TestCommutativityUtils.cpp

+//===- TestCommutativityUtils.cpp - Pass to test the commutativity utility-===//
+//
+// 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 pass tests the functionality of the commutativity utility pattern.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Transforms/CommutativityUtils.h"
+
+#include "TestDialect.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+
+namespace {
+
+struct CommutativityUtils
+    : public PassWrapper<CommutativityUtils, OperationPass<func::FuncOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CommutativityUtils)
+
+  StringRef getArgument() const final { return "test-commutativity-utils"; }
+  StringRef getDescription() const final {
+    return "Test the functionality of the commutativity utility";
+  }
+
+  void runOnOperation() override {
+    auto func = getOperation();
+    auto *context = &getContext();
+
+    RewritePatternSet patterns(context);
+    populateCommutativityUtilsPatterns(patterns);
+
+    (void)applyPatternsAndFoldGreedily(func, std::move(patterns));
+  }
+};
+} // namespace
+
+namespace mlir {
+namespace test {
+void registerCommutativityUtils() { PassRegistration<CommutativityUtils>(); }
+} // namespace test
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestReducer();
 void registerTestSpirvEntryPointABIPass();
 void registerTestSpirvModuleCombinerPass();
 void registerTestTraitsPass();
 void registerTosaTestQuantUtilAPIPass();
 void registerVectorizerTestPass();
 
 namespace test {
+void registerCommutativityUtils();
 void registerConvertCallOpPass();
 void registerInliner();
 void registerMemRefBoundCheck();
 void registerPatternsTestPass();
 void registerSimpleParametricTilingPass();
 void registerTestAffineLoopParametricTilingPass();
 void registerTestAliasAnalysisPass();
 void registerTestBuiltinAttributeInterfaces();
void registerTestPasses() {
   registerTestPrintNestingPass();
   registerTestReducer();
   registerTestSpirvEntryPointABIPass();
   registerTestSpirvModuleCombinerPass();
   registerTestTraitsPass();
   registerVectorizerTestPass();
   registerTosaTestQuantUtilAPIPass();
 
+  mlir::test::registerCommutativityUtils();
   mlir::test::registerConvertCallOpPass();
   mlir::test::registerInliner();
   mlir::test::registerMemRefBoundCheck();
   mlir::test::registerPatternsTestPass();
   mlir::test::registerSimpleParametricTilingPass();
   mlir::test::registerTestAffineLoopParametricTilingPass();
   mlir::test::registerTestAliasAnalysisPass();
   mlir::test::registerTestBuiltinAttributeInterfaces();