EnzymeAD / EnzymeAD/Enzyme-JAX
MultiRotate Op and MultiSliceOp
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
A MultiRotate op has the following semantics
rotleft2, rotleft1, rotate0, rotright1, rotright2 = multi_rotate %x, dimension=, left_amount=2, right_amount=2
This has the semantics that rotleft2 = rotate x left by 2, rotleft2 = rotate x left by 1, etc
This should have an EnzymeHLOOpt for the case where the leftmost or rightmost operands are not used, in which case the operand is replace it with a smaller multi_rotate. If only one result is used, it can be replaced with a single rotate op.
A MultiSlice op has the following semantics
sliceleft2, sliceleft1, slice, sliceright2, sliceright2 = multi_slice %x[start0:end0:step0, start1:end1:step1, ....], dimension=, left_amount=2, right_amount=2
Suppose dimension = 0, this has the semantics that sliceleft2 = x[start0:end0-4:step0, start1:end1:step1, ....]; sliceleft1 = x[start0+1:end0-3:step0, start1:end1:step1, ....];
This similarly should have an EnzymeHLOOpt for the case where the leftmost or rightmost operands are not used, in which case the operand is replace it with a smaller multi_slice. If only one result is used, it can be replaced with a single slice op.
For both new operands, also write patterns to lower the multi versions of op into the corresponding single versions (as above), preserving sharding. Additionally write versions to recognize as large a width multi-op as there are available users [aka all results must be used]. Furthermore, implement constant propagation, folders, and CSE for the operations. Finally add an additional CSE for slice and rotate that will re-use an existing multi-op result, if it already computes the result.
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.