Investigate the ability of `linalg.elementwise` to hold more than one internal operation
- Dominant language
- LLVM
- Stars
- 40.5k
- Forks
- 18.7k
- PR merge metrics
- PR metrics pending
Description
`linalg.elementwise` could allow multiple operations for simpler fusion:
```
// Before fusion
%exp = linalg.elementwise ins(%arg0)
%add = linalg.elementwise ins(%exp, %arg1) // %exp is _chain_
%abs = linalg.elementwise ins(%add) // %add is _chain_
%select = linalg.elementwise ins(%arg2, %abs, %arg3) // %abs is _chain_
// After fusion
%select = linalg.elementwise ins(%arg0, %arg1, %arg2, %arg3)
```
_(see #221194 for the syntax above)_
The argument tree is guaranteed to be unique. The first operation takes in _only_ from arguments (as many as required). Subsequent operations take the _chain_ from the previous operation as the first operand and the following arguments when necessary. `select` is different because the first operand is the mask, so the _chain_ goes into the second operand.
Fusion may not work IF **not every operation** in the elementwise sequence:
* is in _canonical form_: chain as the first argument.
* takes extra operands from arguments and **not** the _chain_. [1]
* has compatible projected permutation affine maps.
_[1] `%select = linalg.elementwise ins(%arg2, %abs, %add)` does not form a tree._
Partial fusion should be encouraged. Elementwise fusion algorithms should scan the chain until it can no longer fuse. Greedy algorithms should find as many _"fusion bubbles"_ as possible in a given IR region. Further transforms and canonicalization of remaining operations may encourage nearby _"bubbles"_ to fuse together.
Partially splitting a fused elementwise can also be done trivially by separating the arguments and creating two smaller fused operations (or a single operation in case of peeling).
Contributor guide
Research direction
Start with the linalg.elementwise operation and the syntax described in issue #221194. Investigate how chained operations, operand ordering, projected permutation affine maps, partial fusion, and splitting would be represented and transformed. Done means documenting or implementing a viable design that supports multiple internal operations while preserving valid fusion bubbles and non-fusible cases.
Written by the indexing model from the issue text.
Assessment
- Domain
- compilers
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100