llvm / llvm/llvm-project

Investigate the ability of `linalg.elementwise` to hold more than one internal operation

Open
#221,197 5 comments 0 reactions 0 assignees View on GitHub
mlir:linalg
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.