[all-commits] [llvm/llvm-project] b508c5: [MLIR] Add a utility to sort the operands of commu...

[Srishti Srivastava via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: b508c5649f5e21e17e9f5633236ec61c551803af
      https://github.com/llvm/llvm-project/commit/b508c5649f5e21e17e9f5633236ec61c551803af
  Author: srishti-cb <srishti.srivastava at polymagelabs.com>
  Date:   2022-07-30 (Sat, 30 Jul 2022)

  Changed paths:
    A mlir/include/mlir/Transforms/CommutativityUtils.h
    M mlir/lib/Transforms/Utils/CMakeLists.txt
    A mlir/lib/Transforms/Utils/CommutativityUtils.cpp
    A mlir/test/Transforms/test-commutativity-utils.mlir
    M mlir/test/lib/Dialect/Test/TestOps.td
    M mlir/test/lib/Transforms/CMakeLists.txt
    A mlir/test/lib/Transforms/TestCommutativityUtils.cpp
    M mlir/tools/mlir-opt/mlir-opt.cpp

  Log Message:
  -----------
  [MLIR] Add a utility to sort the operands of commutative ops

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 <srishti.srivastava at polymagelabs.com>

Reviewed By: Mogball

Differential Revision: https://reviews.llvm.org/D124750