EnzymeAD / EnzymeAD/Enzyme-JAX

MultiRotate Op and MultiSliceOp

Open
#1,949 4 comments 1 reaction 2 assignees Claimed by @wsmoses View on GitHub
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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.