Implement MultiRotate and MultiSlice ops with optimization patterns

[Copilot]

Adds multi_rotate and multi_slice operations that produce multiple shifted versions of a tensor in a single op, enabling better optimization opportunities for stencil-like patterns.

Operations

MultiRotateOp: Produces L+R+1 rotated results given left_amount=L and right_amount=R

%results:3 = "enzymexla.multi_rotate"(%input) <{
  dimension = 1 : si32, 
  left_amount = 1 : si32,   // rotate left by 1
  right_amount = 1 : si32   // rotate right by 1
}> : (tensor<10x20xf32>) -> (tensor<10x20xf32>, tensor<10x20xf32>, tensor<10x20xf32>)
// results[0] = rotate left by 1
// results[1] = no rotation
// results[2] = rotate right by 1

MultiSliceOp: Similar semantics for slice operations with offset start indices along the specified dimension.

Optimization Patterns

• ReduceUnused: Shrinks multi-ops when edge results are unused; converts to single op when only one result used
• Lower: Expands multi-ops to individual single ops, preserving sharding annotations
• Recognize: Detects multiple single rotate/slice ops on same operand and combines into multi-op
• CSE: Reuses existing multi-op results instead of creating redundant single ops

Transform Integration

Patterns registered in transform dialect:

• multi_rotate_opt / multi_slice_opt for all optimization patterns
• multi_rotate_licm / multi_slice_licm for loop hoisting

Implementation Notes

Multi-op indexing: result[i] has rotation amount centerIdx - i where centerIdx = left_amount. Similar offset-based indexing for slices. Patterns handle both directions of conversion (single↔multi) with proper index mapping.

Warning

Firewall rules blocked me from connecting to one or more addresses (expand for details)

I tried to connect to the following addresses, but was blocked by firewall rules:

• releases.bazel.build
  • Triggering command: /usr/local/lib/node_modules/@bazel/bazelisk/bazelisk-linux_amd64 /usr/local/lib/node_modules/@bazel/bazelisk/bazelisk-linux_amd64 build //src/enzyme_ad/jax:EnzymeXLA (dns block)

If you need me to access, download, or install something from one of these locations, you can either:

• Configure Actions setup steps to set up my environment, which run before the firewall is enabled
• Add the appropriate URLs or hosts to the custom allowlist in this repository's Copilot coding agent settings (admins only)

Original prompt

---

This section details on the original issue you should resolve

<issue_title>MultiRotate Op and MultiSliceOp</issue_title>
<issue_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.</issue_description>

Comments on the Issue (you are @copilot in this section)

@wsmoses cc @avik-pal @giordano @jumerckx @ftynse

• Fixes MultiRotate Op and MultiSliceOp #1949

---

✨ Let Copilot coding agent set things up for you — coding agent works faster and does higher quality work when set up for your repo.