diff --git a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp index c6e5e63dc..430ff922b 100644 --- a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp +++ b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp @@ -6935,6 +6935,246 @@ struct NoopReverse final : OpRewritePattern { } }; +bool check_periodicity(std::vector &sliceStart, + std::vector &sliceStride, std::vector &prev, + int dim, int k) { + if (sliceStart[k] == -1) { + sliceStart[k] = dim; + prev[k] = dim; + return true; + } + if (dim - prev[k] == sliceStride[k]) { + prev[k] = dim; + return true; + } else + return false; +} + +int get_index(const std::vector &strides, const std::vector &dims, + int rank_lower, int rank_upper, int iter) { + int index = 0; + for (int i = rank_lower; i >= rank_upper; i--) { + int mod = iter % dims[i]; + iter = iter / dims[i]; + index += mod * strides[i]; + } + return index; +} + +// Rank 0 to n-1, where 0 is the slowest moving dimension (outermost) +bool isGatherPeriodic(const std::vector &data, + const std::vector &dims, int rank, int x_dim, + std::vector &sliceStart, std::vector &sliceEnd, + std::vector &sliceStride) { + // Calculate total elements and strides + int total_size = 1; + std::vector strides(rank); + for (int i = 0; i < rank; i++) + total_size *= dims[i]; + + strides[rank - 1] = 1; + for (int i = rank - 2; i >= 0; i--) + strides[i] = strides[i + 1] * dims[i + 1]; + + int outer_batch_size = 1; + for (int i = 0; i <= x_dim; i++) { + outer_batch_size *= dims[i]; + } + int inner_batch_size = 1; + for (int i = x_dim + 1; i < rank; i++) { + inner_batch_size *= dims[i]; + } + + int vec_length = dims[x_dim]; + + // Run 2 iterations of each rank to get the stride on each rank + bool strideFound = false; + for (int i = 0; i < outer_batch_size / vec_length; i++) { + int index_outer = get_index(strides, dims, x_dim - 1, 0, i); + for (int j = 0; j < inner_batch_size; j++) { + int index_inner = get_index(strides, dims, rank - 1, x_dim + 1, j); + for (int k = 0; k < vec_length; k++) { + int index = index_outer + k * inner_batch_size + index_inner; + int value = data[index]; + if (sliceStride[k] == -1) { + sliceStride[k] = value; + } else { + sliceStride[k] = value - sliceStride[k]; + strideFound = true; + } + } + if (strideFound) + break; + } + if (strideFound) + break; + } + + // Run all the iterations to check if the strides match + std::vector prev(vec_length, -1); + for (int i = 0; i < outer_batch_size / vec_length; i++) { + int index_outer = get_index(strides, dims, x_dim - 1, 0, i); + for (int j = 0; j < inner_batch_size; j++) { + int index_inner = get_index(strides, dims, rank - 1, x_dim + 1, j); + for (int k = 0; k < vec_length; k++) { + int index = index_outer + k * inner_batch_size + index_inner; + int value = data[index]; + auto res = check_periodicity(sliceStart, sliceStride, prev, value, k); + if (!res) + return false; + } + } + } + for (int k = 0; k < vec_length; k++) { + if(sliceStride[k] >= 0) { + sliceEnd[k] = prev[k] + 1; + } + else if (sliceStride[k] < 0) { + sliceEnd[k] = prev[k]; + sliceStart[k] = sliceStart[k] + 1; + } + } + return true; +} + +/// Converts gather ops to slice ops in case we have a single set of constant +/// indices. +struct GatherToSliceOp final : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(mlir::stablehlo::GatherOp gather, + PatternRewriter &rewriter) const override { + DenseIntElementsAttr index; + // 1. Check following preconditions for converting gather to slice op + // Preconditions: + // i. Check if all dims are collapsed? + // ii. Check if sliceSizes are 1 in each dim (else more complicated and + // there could be a subset that + // could still be transformed to slice op if there is no overlap) + // iii. Based on indexVector dim, check if correspoding values in each dim + // are either constant or strided + // 2. If so, convert the gather to a slice op + // 3. If not, return failure + // 4. To convert to slice operation, + + // Get all the other properties of the gather operation + auto startIndices = gather.getStartIndices(); + auto gatherOperands = gather.getOperands(); + auto gatherDimensionNumbers = gather.getDimensionNumbers(); + auto gatherSliceSizes = gather.getSliceSizes(); + auto gatherIndexVectorDim = gatherDimensionNumbers.getIndexVectorDim(); + auto gatherStartIndexMap = gatherDimensionNumbers.getStartIndexMap(); + auto collapsedSliceDims = gatherDimensionNumbers.getCollapsedSliceDims(); + auto offsetDims = + gatherDimensionNumbers.getOffsetDims(); // Non collapsed dimensions + + ////Check rank of gather operands + // auto gatherOperandsRank = gatherOperands.getType().getRank(); + + //// TODO: Currently only handling simplified case + //// Check collapseSlicedims size equals gatherOperandsRank + // if(collapsedSliceDims.size() != gatherOperandsRank) + // return failure(); + + // Check if sliceSizes are 1 in each dim + if (llvm::any_of(gatherSliceSizes, [](int64_t size) { return size != 1; })) + return failure(); + + // From start indices check if it's constant, else return failure + if (!matchPattern(startIndices, m_Constant(&index))) + return failure(); + + // Currently only handling constant index case with dense elements + std::vector indices; + if (auto denseAttr = llvm::dyn_cast(index)) { + for (auto value : denseAttr.getValues()) { + int64_t intValue = value.getSExtValue(); + indices.push_back(intValue); + } + } else { + return failure(); + } + + // Process indices in row-major order + auto tensorType = index.getType().cast(); + auto rank = tensorType.getRank(); + auto shape = tensorType.getShape(); + + // dims : gatherIndexVectorDim, innermostDim .... outermostDim + int indexVectorSize = shape[gatherIndexVectorDim]; + + std::vector dims(shape.size()); + for (int i = 0; i < shape.size(); i++) { + dims[i] = shape[i]; + } + + std::vector sliceStart(indexVectorSize, -1); + std::vector sliceEnd(indexVectorSize, -1); + std::vector sliceStride(indexVectorSize, -1); + + auto isPeriodic = + isGatherPeriodic(indices, dims, rank, gatherIndexVectorDim, sliceStart, + sliceEnd, sliceStride); + if (!isPeriodic) + return failure(); + + SmallVector sliceStartI64(sliceStart.begin(), sliceStart.end()); + SmallVector sliceEndI64(sliceEnd.begin(), sliceEnd.end()); + SmallVector sliceStrideI64(sliceStride.begin(), + sliceStride.end()); + + // Create the slice type + SmallVector sliceShape(sliceStrideI64.size()); + SmallVector collapsedShape; + SmallVector reverseDims; + bool reverse = false; + for (int i = 0; i < sliceStrideI64.size(); i++) { + sliceShape[i] = sliceEndI64[i] - sliceStartI64[i]; + //Reverse the slice if it's negative + if(sliceShape[i] < 0) { + sliceShape[i] = -sliceShape[i]; + sliceStrideI64[i] = -sliceStrideI64[i]; + auto temp = sliceStartI64[i]; + sliceStartI64[i] = sliceEndI64[i]; + sliceEndI64[i] = temp; + reverseDims.push_back(i); + reverse = true; + } + + if (sliceShape[i] != 1) + collapsedShape.push_back(sliceShape[i]); + } + Type elementType = gather.getType().getElementType(); + auto sliceType = RankedTensorType::get(sliceShape, elementType); + + // Fix the constant dims + for (int i = 0; i < sliceStrideI64.size(); i++) { + if (sliceStrideI64[i] == 0) + sliceStrideI64[i] = 1; + } + + // Creating the slice op + Value sliceOp = rewriter.create( + gather.getLoc(), sliceType, gather.getOperand(), + rewriter.getDenseI64ArrayAttr(sliceStartI64), + rewriter.getDenseI64ArrayAttr(sliceEndI64), + rewriter.getDenseI64ArrayAttr(sliceStrideI64)); + + if(reverse) { + sliceOp = rewriter.create(gather.getLoc(), sliceOp, reverseDims); + } + + // Create result type and reshape operation + auto collapsedType = RankedTensorType::get(collapsedShape, elementType); + Value sliceReshaped = rewriter.create( + gather.getLoc(), collapsedType, sliceOp); + + rewriter.replaceOp(gather, sliceReshaped); + + return success(); + } +}; + /// Converts gather ops to slice ops in case we have a single set of constant /// indices. struct GatherOpCanon final : OpRewritePattern { @@ -8215,6 +8455,7 @@ struct EnzymeHLOOptPass // clang-format off patterns.add< + GatherToSliceOp, BroadcastInDimOpCanon, ChainedDynamicBroadcastInDimCanonicalization, CompareOpCanon, diff --git a/test/lit_tests/gatherToSlice.mlir b/test/lit_tests/gatherToSlice.mlir new file mode 100644 index 000000000..03a6c2ef8 --- /dev/null +++ b/test/lit_tests/gatherToSlice.mlir @@ -0,0 +1,87 @@ +// RUN: mlir-opt -split-input-file -convert-stablehlo-gather-to-slice %s | FileCheck %s + +// Original example: +// %c_1179 = stablehlo.constant dense<"0x00000000000000006B00000000000000070000000000000000..."> : tensor<180x3xi64> +// %2803 = stablehlo.dynamic_update_slice %2715, %2802, %c_1336, %c_1308, %c_1329 : (tensor<1x128x194xf64>, tensor<1x1x180xf64>, tensor, tensor, tensor) -> tensor<1x128x194xf64> +// %2804 = "stablehlo.gather"(%2803, %c_1179) <{dimension_numbers = #stablehlo.gather, indices_are_sorted = false, +// slice_sizes = array}> : (tensor<1x128x194xf64>, tensor<180x3xi64>) -> tensor<180xf64> + +// CHECK-LABEL: func @gather_to_slice +//func.func @gather_to_slice_wrapped_around(%arg0: tensor<1x128x194xf64>) -> tensor<180xf64> { +// %indices = stablehlo.constant dense<"0x00000000000000006B00000000000000070000000000000000000000000000006B00000000000000BA0000000000000000000000000000006B00000000000000B90000000000000000000000000000006B00000000000000B80000000000000000000000000000006B00000000000000B70000000000000000000000000000006B00000000000000B60000000000000000000000000000006B00000000000000B50000000000000000000000000000006B00000000000000B40000000000000000000000000000006B00000000000000B30000000000000000000000000000006B00000000000000B20000000000000000000000000000006B00000000000000B10000000000000000000000000000006B00000000000000B00000000000000000000000000000006B00000000000000AF0000000000000000000000000000006B00000000000000AE0000000000000000000000000000006B00000000000000AD0000000000000000000000000000006B00000000000000AC0000000000000000000000000000006B00000000000000AB0000000000000000000000000000006B00000000000000AA0000000000000000000000000000006B00000000000000A90000000000000000000000000000006B00000000000000A80000000000000000000000000000006B00000000000000A70000000000000000000000000000006B00000000000000A60000000000000000000000000000006B00000000000000A50000000000000000000000000000006B00000000000000A40000000000000000000000000000006B00000000000000A30000000000000000000000000000006B00000000000000A20000000000000000000000000000006B00000000000000A10000000000000000000000000000006B00000000000000A00000000000000000000000000000006B000000000000009F0000000000000000000000000000006B000000000000009E0000000000000000000000000000006B000000000000009D0000000000000000000000000000006B000000000000009C0000000000000000000000000000006B000000000000009B0000000000000000000000000000006B000000000000009A0000000000000000000000000000006B00000000000000990000000000000000000000000000006B00000000000000980000000000000000000000000000006B00000000000000970000000000000000000000000000006B00000000000000960000000000000000000000000000006B00000000000000950000000000000000000000000000006B00000000000000940000000000000000000000000000006B00000000000000930000000000000000000000000000006B00000000000000920000000000000000000000000000006B00000000000000910000000000000000000000000000006B00000000000000900000000000000000000000000000006B000000000000008F0000000000000000000000000000006B000000000000008E0000000000000000000000000000006B000000000000008D0000000000000000000000000000006B000000000000008C0000000000000000000000000000006B000000000000008B0000000000000000000000000000006B000000000000008A0000000000000000000000000000006B00000000000000890000000000000000000000000000006B00000000000000880000000000000000000000000000006B00000000000000870000000000000000000000000000006B00000000000000860000000000000000000000000000006B00000000000000850000000000000000000000000000006B00000000000000840000000000000000000000000000006B00000000000000830000000000000000000000000000006B00000000000000820000000000000000000000000000006B00000000000000810000000000000000000000000000006B00000000000000800000000000000000000000000000006B000000000000007F0000000000000000000000000000006B000000000000007E0000000000000000000000000000006B000000000000007D0000000000000000000000000000006B000000000000007C0000000000000000000000000000006B000000000000007B0000000000000000000000000000006B000000000000007A0000000000000000000000000000006B00000000000000790000000000000000000000000000006B00000000000000780000000000000000000000000000006B00000000000000770000000000000000000000000000006B00000000000000760000000000000000000000000000006B00000000000000750000000000000000000000000000006B00000000000000740000000000000000000000000000006B00000000000000730000000000000000000000000000006B00000000000000720000000000000000000000000000006B00000000000000710000000000000000000000000000006B00000000000000700000000000000000000000000000006B000000000000006F0000000000000000000000000000006B000000000000006E0000000000000000000000000000006B000000000000006D0000000000000000000000000000006B000000000000006C0000000000000000000000000000006B000000000000006B0000000000000000000000000000006B000000000000006A0000000000000000000000000000006B00000000000000690000000000000000000000000000006B00000000000000680000000000000000000000000000006B00000000000000670000000000000000000000000000006B00000000000000660000000000000000000000000000006B00000000000000650000000000000000000000000000006B00000000000000640000000000000000000000000000006B00000000000000630000000000000000000000000000006B00000000000000620000000000000000000000000000006B00000000000000610000000000000000000000000000006B00000000000000600000000000000000000000000000006B000000000000005F0000000000000000000000000000006B000000000000005E0000000000000000000000000000006B000000000000005D0000000000000000000000000000006B000000000000005C0000000000000000000000000000006B000000000000005B0000000000000000000000000000006B000000000000005A0000000000000000000000000000006B00000000000000590000000000000000000000000000006B00000000000000580000000000000000000000000000006B00000000000000570000000000000000000000000000006B00000000000000560000000000000000000000000000006B00000000000000550000000000000000000000000000006B00000000000000540000000000000000000000000000006B00000000000000530000000000000000000000000000006B00000000000000520000000000000000000000000000006B00000000000000510000000000000000000000000000006B00000000000000500000000000000000000000000000006B000000000000004F0000000000000000000000000000006B000000000000004E0000000000000000000000000000006B000000000000004D0000000000000000000000000000006B000000000000004C0000000000000000000000000000006B000000000000004B0000000000000000000000000000006B000000000000004A0000000000000000000000000000006B00000000000000490000000000000000000000000000006B00000000000000480000000000000000000000000000006B00000000000000470000000000000000000000000000006B00000000000000460000000000000000000000000000006B00000000000000450000000000000000000000000000006B00000000000000440000000000000000000000000000006B00000000000000430000000000000000000000000000006B00000000000000420000000000000000000000000000006B00000000000000410000000000000000000000000000006B00000000000000400000000000000000000000000000006B000000000000003F0000000000000000000000000000006B000000000000003E0000000000000000000000000000006B000000000000003D0000000000000000000000000000006B000000000000003C0000000000000000000000000000006B000000000000003B0000000000000000000000000000006B000000000000003A0000000000000000000000000000006B00000000000000390000000000000000000000000000006B00000000000000380000000000000000000000000000006B00000000000000370000000000000000000000000000006B00000000000000360000000000000000000000000000006B00000000000000350000000000000000000000000000006B00000000000000340000000000000000000000000000006B00000000000000330000000000000000000000000000006B00000000000000320000000000000000000000000000006B00000000000000310000000000000000000000000000006B00000000000000300000000000000000000000000000006B000000000000002F0000000000000000000000000000006B000000000000002E0000000000000000000000000000006B000000000000002D0000000000000000000000000000006B000000000000002C0000000000000000000000000000006B000000000000002B0000000000000000000000000000006B000000000000002A0000000000000000000000000000006B00000000000000290000000000000000000000000000006B00000000000000280000000000000000000000000000006B00000000000000270000000000000000000000000000006B00000000000000260000000000000000000000000000006B00000000000000250000000000000000000000000000006B00000000000000240000000000000000000000000000006B00000000000000230000000000000000000000000000006B00000000000000220000000000000000000000000000006B00000000000000210000000000000000000000000000006B00000000000000200000000000000000000000000000006B000000000000001F0000000000000000000000000000006B000000000000001E0000000000000000000000000000006B000000000000001D0000000000000000000000000000006B000000000000001C0000000000000000000000000000006B000000000000001B0000000000000000000000000000006B000000000000001A0000000000000000000000000000006B00000000000000190000000000000000000000000000006B00000000000000180000000000000000000000000000006B00000000000000170000000000000000000000000000006B00000000000000160000000000000000000000000000006B00000000000000150000000000000000000000000000006B00000000000000140000000000000000000000000000006B00000000000000130000000000000000000000000000006B00000000000000120000000000000000000000000000006B00000000000000110000000000000000000000000000006B00000000000000100000000000000000000000000000006B000000000000000F0000000000000000000000000000006B000000000000000E0000000000000000000000000000006B000000000000000D0000000000000000000000000000006B000000000000000C0000000000000000000000000000006B000000000000000B0000000000000000000000000000006B000000000000000A0000000000000000000000000000006B00000000000000090000000000000000000000000000006B000000000000000800000000000000"> : tensor<180x3xi64> +// +// // CHECK: %[[SLICE:.*]] = "stablehlo.slice"(%arg0) +// // CHECK-SAME: start_indices = array +// // CHECK-SAME: limit_indices = array +// // CHECK-SAME: strides = array +// // CHECK: return %[[SLICE]] : tensor<180xf64> +// %result = "stablehlo.gather"(%arg0, %indices) { +// dimension_numbers = #stablehlo.gather< +// collapsed_slice_dims = [0, 1, 2], +// start_index_map = [0, 1, 2], +// index_vector_dim = 1 +// >, +// indices_are_sorted = false, +// slice_sizes = array +// } : (tensor<1x128x194xf64>, tensor<180x3xi64>) -> tensor<180xf64> +// +// return %result : tensor<180xf64> +//} + +// ----- + +// Example of multi dim strided slice op (for reference): +// %1 = "stablehlo.slice"(%arg0) +// start_indices = array +// limit_indices = array +// strides = array +// } : tensor<1x128x194xf64> -> tensor<180xf64> +// + +func.func @gather_to_slice_reverse(%arg0: tensor<1x128x194xf64>) -> tensor<179xf64> { + %indices = stablehlo.constant dense<"0x00000000000000006B00000000000000BA0000000000000000000000000000006B00000000000000B90000000000000000000000000000006B00000000000000B80000000000000000000000000000006B00000000000000B70000000000000000000000000000006B00000000000000B60000000000000000000000000000006B00000000000000B50000000000000000000000000000006B00000000000000B40000000000000000000000000000006B00000000000000B30000000000000000000000000000006B00000000000000B20000000000000000000000000000006B00000000000000B10000000000000000000000000000006B00000000000000B00000000000000000000000000000006B00000000000000AF0000000000000000000000000000006B00000000000000AE0000000000000000000000000000006B00000000000000AD0000000000000000000000000000006B00000000000000AC0000000000000000000000000000006B00000000000000AB0000000000000000000000000000006B00000000000000AA0000000000000000000000000000006B00000000000000A90000000000000000000000000000006B00000000000000A80000000000000000000000000000006B00000000000000A70000000000000000000000000000006B00000000000000A60000000000000000000000000000006B00000000000000A50000000000000000000000000000006B00000000000000A40000000000000000000000000000006B00000000000000A30000000000000000000000000000006B00000000000000A20000000000000000000000000000006B00000000000000A10000000000000000000000000000006B00000000000000A00000000000000000000000000000006B000000000000009F0000000000000000000000000000006B000000000000009E0000000000000000000000000000006B000000000000009D0000000000000000000000000000006B000000000000009C0000000000000000000000000000006B000000000000009B0000000000000000000000000000006B000000000000009A0000000000000000000000000000006B00000000000000990000000000000000000000000000006B00000000000000980000000000000000000000000000006B00000000000000970000000000000000000000000000006B00000000000000960000000000000000000000000000006B00000000000000950000000000000000000000000000006B00000000000000940000000000000000000000000000006B00000000000000930000000000000000000000000000006B00000000000000920000000000000000000000000000006B00000000000000910000000000000000000000000000006B00000000000000900000000000000000000000000000006B000000000000008F0000000000000000000000000000006B000000000000008E0000000000000000000000000000006B000000000000008D0000000000000000000000000000006B000000000000008C0000000000000000000000000000006B000000000000008B0000000000000000000000000000006B000000000000008A0000000000000000000000000000006B00000000000000890000000000000000000000000000006B00000000000000880000000000000000000000000000006B00000000000000870000000000000000000000000000006B00000000000000860000000000000000000000000000006B00000000000000850000000000000000000000000000006B00000000000000840000000000000000000000000000006B00000000000000830000000000000000000000000000006B00000000000000820000000000000000000000000000006B00000000000000810000000000000000000000000000006B00000000000000800000000000000000000000000000006B000000000000007F0000000000000000000000000000006B000000000000007E0000000000000000000000000000006B000000000000007D0000000000000000000000000000006B000000000000007C0000000000000000000000000000006B000000000000007B0000000000000000000000000000006B000000000000007A0000000000000000000000000000006B00000000000000790000000000000000000000000000006B00000000000000780000000000000000000000000000006B00000000000000770000000000000000000000000000006B00000000000000760000000000000000000000000000006B00000000000000750000000000000000000000000000006B00000000000000740000000000000000000000000000006B00000000000000730000000000000000000000000000006B00000000000000720000000000000000000000000000006B00000000000000710000000000000000000000000000006B00000000000000700000000000000000000000000000006B000000000000006F0000000000000000000000000000006B000000000000006E0000000000000000000000000000006B000000000000006D0000000000000000000000000000006B000000000000006C0000000000000000000000000000006B000000000000006B0000000000000000000000000000006B000000000000006A0000000000000000000000000000006B00000000000000690000000000000000000000000000006B00000000000000680000000000000000000000000000006B00000000000000670000000000000000000000000000006B00000000000000660000000000000000000000000000006B00000000000000650000000000000000000000000000006B00000000000000640000000000000000000000000000006B00000000000000630000000000000000000000000000006B00000000000000620000000000000000000000000000006B00000000000000610000000000000000000000000000006B00000000000000600000000000000000000000000000006B000000000000005F0000000000000000000000000000006B000000000000005E0000000000000000000000000000006B000000000000005D0000000000000000000000000000006B000000000000005C0000000000000000000000000000006B000000000000005B0000000000000000000000000000006B000000000000005A0000000000000000000000000000006B00000000000000590000000000000000000000000000006B00000000000000580000000000000000000000000000006B00000000000000570000000000000000000000000000006B00000000000000560000000000000000000000000000006B00000000000000550000000000000000000000000000006B00000000000000540000000000000000000000000000006B00000000000000530000000000000000000000000000006B00000000000000520000000000000000000000000000006B00000000000000510000000000000000000000000000006B00000000000000500000000000000000000000000000006B000000000000004F0000000000000000000000000000006B000000000000004E0000000000000000000000000000006B000000000000004D0000000000000000000000000000006B000000000000004C0000000000000000000000000000006B000000000000004B0000000000000000000000000000006B000000000000004A0000000000000000000000000000006B00000000000000490000000000000000000000000000006B00000000000000480000000000000000000000000000006B00000000000000470000000000000000000000000000006B00000000000000460000000000000000000000000000006B00000000000000450000000000000000000000000000006B00000000000000440000000000000000000000000000006B00000000000000430000000000000000000000000000006B00000000000000420000000000000000000000000000006B00000000000000410000000000000000000000000000006B00000000000000400000000000000000000000000000006B000000000000003F0000000000000000000000000000006B000000000000003E0000000000000000000000000000006B000000000000003D0000000000000000000000000000006B000000000000003C0000000000000000000000000000006B000000000000003B0000000000000000000000000000006B000000000000003A0000000000000000000000000000006B00000000000000390000000000000000000000000000006B00000000000000380000000000000000000000000000006B00000000000000370000000000000000000000000000006B00000000000000360000000000000000000000000000006B00000000000000350000000000000000000000000000006B00000000000000340000000000000000000000000000006B00000000000000330000000000000000000000000000006B00000000000000320000000000000000000000000000006B00000000000000310000000000000000000000000000006B00000000000000300000000000000000000000000000006B000000000000002F0000000000000000000000000000006B000000000000002E0000000000000000000000000000006B000000000000002D0000000000000000000000000000006B000000000000002C0000000000000000000000000000006B000000000000002B0000000000000000000000000000006B000000000000002A0000000000000000000000000000006B00000000000000290000000000000000000000000000006B00000000000000280000000000000000000000000000006B00000000000000270000000000000000000000000000006B00000000000000260000000000000000000000000000006B00000000000000250000000000000000000000000000006B00000000000000240000000000000000000000000000006B00000000000000230000000000000000000000000000006B00000000000000220000000000000000000000000000006B00000000000000210000000000000000000000000000006B00000000000000200000000000000000000000000000006B000000000000001F0000000000000000000000000000006B000000000000001E0000000000000000000000000000006B000000000000001D0000000000000000000000000000006B000000000000001C0000000000000000000000000000006B000000000000001B0000000000000000000000000000006B000000000000001A0000000000000000000000000000006B00000000000000190000000000000000000000000000006B00000000000000180000000000000000000000000000006B00000000000000170000000000000000000000000000006B00000000000000160000000000000000000000000000006B00000000000000150000000000000000000000000000006B00000000000000140000000000000000000000000000006B00000000000000130000000000000000000000000000006B00000000000000120000000000000000000000000000006B00000000000000110000000000000000000000000000006B00000000000000100000000000000000000000000000006B000000000000000F0000000000000000000000000000006B000000000000000E0000000000000000000000000000006B000000000000000D0000000000000000000000000000006B000000000000000C0000000000000000000000000000006B000000000000000B0000000000000000000000000000006B000000000000000A0000000000000000000000000000006B00000000000000090000000000000000000000000000006B000000000000000800000000000000"> : tensor<179x3xi64> + // CHECK-NOT: stablehlo.gather + // CHECK: %[[SLICE:.*]] = stablehlo.slice %[[ARG0]] [0:1, 107:108, 8:187] + // CHECK-SAME: : (tensor<1x128x194xf64>) -> tensor<1x1x179xf64> + // CHECK: %[[REVERSED:.*]] = stablehlo.reverse %[[SLICE]], dims = [2] + // CHECK-SAME: : tensor<1x1x179xf64> + // CHECK: %[[RESHAPED:.*]] = stablehlo.reshape %[[REVERSED]] + // CHECK-SAME: : (tensor<1x1x179xf64>) -> tensor<179xf64> + %result = "stablehlo.gather"(%arg0, %indices) { + dimension_numbers = #stablehlo.gather< + collapsed_slice_dims = [0, 1, 2], + start_index_map = [0, 1, 2], + index_vector_dim = 1 + >, + indices_are_sorted = false, + slice_sizes = array + } : (tensor<1x128x194xf64>, tensor<179x3xi64>) -> tensor<179xf64> + + return %result : tensor<179xf64> +} + +func.func @gather_to_slice_collapse_dims(%arg0: tensor<1x128x194xf64>) -> tensor<3xf64> { + %indices = stablehlo.constant dense<[ + [0, 10, 4], + [0, 10, 5], + [0, 10, 6] +]> : tensor<3x3xi64> + // CHECK-NOT: stablehlo.gather + // CHECK: %[[SLICE:.*]] = stablehlo.slice %[[ARG0]] [0:1, 10:11, 4:7] + // CHECK-SAME: : (tensor<1x128x194xf64>) -> tensor<1x1x3xf64> + // CHECK: %[[RESHAPE:.*]] = stablehlo.reshape %[[SLICE]] + // CHECK-SAME: : (tensor<1x1x3xf64>) -> tensor<3xf64> + %result = "stablehlo.gather"(%arg0, %indices) { + dimension_numbers = #stablehlo.gather< + collapsed_slice_dims = [0, 1, 2], + start_index_map = [0, 1, 2], + index_vector_dim = 1 + >, + indices_are_sorted = false, + slice_sizes = array + } : (tensor<1x128x194xf64>, tensor<3x3xi64>) -> tensor<3xf64> + + return %result : tensor<3xf64> +} +