Optimize dynamic_gather operations on transposed tensors

[Copilot]

Optimization Passes for dynamic_gather

Summary

This PR implements the DynamicGatherTranspose optimization pattern that addresses one of the key optimizations mentioned in the issue: rewriting dynamic_gather(transpose(x), ...) to operate directly on the original tensor x.

What was implemented

• DynamicGatherTranspose pattern: Optimizes dynamic_gather(transpose(x), ...) by:

  • Applying inverse permutation to start_index_map to adjust for the removed transpose
  • Applying inverse permutation to collapsed_slice_dims to adjust dimension indices
  • Permuting slice_sizes according to the inverse permutation
  • Currently handles constant slice_sizes (dynamic case would require additional complexity)

• Comprehensive test coverage: Added test/lit_tests/dynamic_gather_opts.mlir with:

  • 2D transpose case (simple dimension swap)
  • Multiple gathers from same transpose (as in the issue example)
  • 3D transpose case with complex permutation

• Code review and formatting fixes: Applied all review feedback

• Fixed compilation error: Added missing batching dimension parameters to GatherDimensionNumbersAttr::get call

What was deferred

• ConcatDynamicGather pattern: Merging concatenated dynamic_gather operations into a single larger gather is complex and requires:
  • Ensuring all gathers have compatible dimension numbers
  • Concatenating indices tensors appropriately
  • Handling different-sized gather results
  • This is left for a follow-up PR

Implementation notes

• The pattern works correctly even when a transpose has multiple uses - each gather is independently optimized, and the transpose is eliminated by dead code elimination after all uses are optimized
• Only constant slice_sizes are currently supported; dynamic slice sizes would require creating runtime operations to permute the sizes tensor
• The pattern correctly handles arbitrary permutations using the getInversePermutation() helper function
• Fixed to match the correct GatherDimensionNumbersAttr::get signature with empty batching dimension arrays

Testing notes

• Build validation was attempted but blocked by network issues with bazel
• The implementation follows existing patterns in the codebase (e.g., TransposeDynamicSlice, DynamicGatherOpIsNotDynamic)
• Code passed CodeQL security analysis with no issues found

Original prompt

---

This section details on the original issue you should resolve

<issue_title>Optimization Passes for dynamic_gather</issue_title>
<issue_description>```mlir
module {
func.func @main(%arg0: tensor<6x6xf64>) -> tensor<6x6xf64> {
%cst = stablehlo.constant dense<1.000000e+00> : tensor<6x6xf64>
%cst_0 = stablehlo.constant dense<0.000000e+00> : tensor<6x6xf64>
%c = stablehlo.constant dense<[[1, 0], [2, 1], [3, 2], [4, 3], [5, 4], [0, 0], [1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [0, 1], [1, 2], [2, 3], [3, 4], [4, 5]]> : tensor<16x2xi64>
%c_1 = stablehlo.constant dense<[[0, 1], [1, 2], [2, 3], [3, 4], [4, 5]]> : tensor<5x2xi64>
%c_2 = stablehlo.constant dense<[[0, 0], [1, 1], [2, 2], [3, 3], [4, 4], [5, 5]]> : tensor<6x2xi64>
%c_3 = stablehlo.constant dense<1> : tensor<2xi64>
%c_4 = stablehlo.constant dense<[[1, 0], [2, 1], [3, 2], [4, 3], [5, 4]]> : tensor<5x2xi64>
%0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<6x6xf64>) -> tensor<6x6xf64>
%1 = "stablehlo.dynamic_gather"(%0, %c_4, %c_3) <{dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0, 1], start_index_map = [0, 1], index_vector_dim = 1>}> : (tensor<6x6xf64>, tensor<5x2xi64>, tensor<2xi64>) -> tensor<5xf64>
%2 = "stablehlo.dynamic_gather"(%0, %c_2, %c_3) <{dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0, 1], start_index_map = [0, 1], index_vector_dim = 1>}> : (tensor<6x6xf64>, tensor<6x2xi64>, tensor<2xi64>) -> tensor<6xf64>
%3 = "stablehlo.dynamic_gather"(%0, %c_1, %c_3) <{dimension_numbers = #stablehlo.gather<collapsed_slice_dims = [0, 1], start_index_map = [0, 1], index_vector_dim = 1>}> : (tensor<6x6xf64>, tensor<5x2xi64>, tensor<2xi64>) -> tensor<5xf64>
%4 = stablehlo.concatenate %1, %2, %3, dim = 0 : (tensor<5xf64>, tensor<6xf64>, tensor<5xf64>) -> tensor<16xf64>
%5 = "stablehlo.scatter"(%cst_0, %c, %4) <{scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0, 1], scatter_dims_to_operand_dims = [0, 1], index_vector_dim = 1>}> ({
^bb0(%arg1: tensor, %arg2: tensor):
stablehlo.return %arg2 : tensor
}) : (tensor<6x6xf64>, tensor<16x2xi64>, tensor<16xf64>) -> tensor<6x6xf64>
%6 = stablehlo.add %5, %cst : tensor<6x6xf64>
%7 = stablehlo.transpose %6, dims = [1, 0] : (tensor<6x6xf64>) -> tensor<6x6xf64>
return %7 : tensor<6x6xf64>
}
}

This comes from `fn(x) = Tridiagonal(x) .+ 1`.

Essentially if we fuse the `dynamic_gather`s into `dynamic_gather` + `slice`, then the `slice + concatenate` will get eliminated by another pass.

(In this particular case even the `dynamic_gather` + `scatter` can be eliminated, but that is probably much harder to write a pass for)</issue_description>

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

<comments>
<comment_new><author>@wsmoses</author><body>
so there's several opts here worth doing.

- [ ] dynamic gather of transpose -> rewritten dynamic gather
- [ ] concat of consecutive dynamic gathers is replaced with a larger dynamic gather
- [ ] eventually scatter of dynamic gather
- [x] dynamic_gather -> gather for static slice_sizes</body></comment_new>
</comments>

• Fixes Optimization Passes for dynamic_gather #184

---

💡 You can make Copilot smarter by setting up custom instructions, customizing its development environment and configuring Model Context Protocol (MCP) servers. Learn more Copilot coding agent tips in the docs.

[avik-pal]

@copilot

 16998 |     auto newDimNumbers = stablehlo::GatherDimensionNumbersAttr::get(
       |                          ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 16999 |         op.getContext(), newOffsetDims, newCollapsedSliceDims,
 17000 |         newStartIndexMap, dimNumbers.getIndexVectorDim());
       |                                                         ^
bazel-out/aarch64-opt/bin/external/stablehlo/stablehlo/dialect/StablehloAttrs.h.inc:269:37: note: 'get' declared here
  269 |   static GatherDimensionNumbersAttr get(::mlir::MLIRContext *context, ::llvm::ArrayRef<int64_t> offsetDims, ::llvm::ArrayRef<int64_t> collapsedSliceDims, ::llvm::ArrayRef<int64_t> operandBatchingDims, ::llvm::ArrayRef<int64_t> startIndicesBatchingDims, ::llvm::ArrayRef<int64_t> startIndexMap, int64_t indexVectorDim);
      |                                     ^   ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

[Copilot]

@copilot

...

Fixed in commit 75f30d1. Added the missing operandBatchingDims and startIndicesBatchingDims parameters (both empty arrays) to the GatherDimensionNumbersAttr::get call to match the correct API signature.