more vibe

Author: wsmoses

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -14607,6 +14607,7 @@ struct GatherOpCanon final
     int64_t affineDim;
     int64_t scale;
     int64_t start;
+    int64_t componentIndex;
   };
 
   LogicalResult
@@ -14650,7 +14651,8 @@ struct GatherOpCanon final
           dnums.getStartIndexMap()[0],
           affineDim,
           static_cast<int64_t>(*scaleVal),
-          static_cast<int64_t>(*startVal)
+          static_cast<int64_t>(*startVal),
+          static_cast<int64_t>(i)
       });
     }
 
@@ -14676,7 +14678,12 @@ struct GatherOpCanon final
 
     int64_t currentScale = 1;
     for (const auto& ds : dimStrides) {
-      int64_t implied_dim_size = baseIndicesTy.getDimSize(ds.affineDim);
+      int64_t implied_dim_size;
+      if (affineIota.implied_dim_sizes.size() > static_cast<size_t>(ds.componentIndex)) {
+        implied_dim_size = affineIota.implied_dim_sizes[ds.componentIndex];
+      } else {
+        implied_dim_size = baseIndicesTy.getDimSize(ds.affineDim);
+      }
       if (ds.scale != currentScale) {
         llvm::errs() << "BAIL affine iota matching: strides not forming a contiguous block (" << ds.scale << " vs " << currentScale << ")\n";
         return failure();

src/enzyme_ad/jax/Utils.cpp

@@ -1866,6 +1866,11 @@ std::optional<AffineIotaLikeTensor> detectAffineIotaLikeTensor(mlir::Value tenso
     AffineIotaLikeTensor result;
     result.starts.push_back(iotaLike->start);
     result.dimensions.push_back(iotaLike->dimension);
+    
+    // Explicitly record the implied dimension size from the original iota shape
+    auto indicesTy = cast<RankedTensorType>(iotaLike->tensorType);
+    result.implied_dim_sizes.push_back(indicesTy.getDimSize(iotaLike->dimension));
+    
     result.scales.push_back(iotaLike->scale);
     result.tensorType = iotaLike->tensorType;
     return result;
@@ -1916,6 +1921,7 @@ std::optional<AffineIotaLikeTensor> detectAffineIotaLikeTensor(mlir::Value tenso
         AffineIotaLikeTensor combined = *lhsAffine;
         combined.starts.insert(combined.starts.end(), rhsAffine->starts.begin(), rhsAffine->starts.end());
         combined.dimensions.insert(combined.dimensions.end(), rhsAffine->dimensions.begin(), rhsAffine->dimensions.end());
+        combined.implied_dim_sizes.insert(combined.implied_dim_sizes.end(), rhsAffine->implied_dim_sizes.begin(), rhsAffine->implied_dim_sizes.end());
         combined.scales.insert(combined.scales.end(), rhsAffine->scales.begin(), rhsAffine->scales.end());
         combined.tensorType = cast<RankedTensorType>(addOp.getType());
         return combined;

src/enzyme_ad/jax/Utils.h

@@ -1114,6 +1114,7 @@ std::optional<IotaLikeTensor> detectIotaLikeTensor(mlir::Value tensor);
 struct AffineIotaLikeTensor {
   llvm::SmallVector<mlir::TypedAttr> starts;
   llvm::SmallVector<int64_t> dimensions;
+  llvm::SmallVector<int64_t> implied_dim_sizes;
   llvm::SmallVector<mlir::TypedAttr> scales;
   mlir::RankedTensorType tensorType;
 };

test/lit_tests/gather_iota_to_slice.mlir

@@ -231,3 +231,35 @@ func.func @gather_iota_negative_stride_offset_to_slice_reverse(%arg0: tensor<10x
 // CHECK-NEXT: %[[SLICE:.+]] = stablehlo.slice %arg0 [1:8:2]
 // CHECK-NEXT: %[[REVERSE:.+]] = stablehlo.reverse %[[SLICE]], dims = [0]
 // CHECK-NEXT: return %[[REVERSE]]
+
+// ============================================================================
+// Tests for gather with multi-dimensional iota reshaped into a flat index
+// ============================================================================
+
+// Reshaped indices: gather with a 2D grid of indices flattened
+func.func @gather_reshaped_iota_to_slice(%arg0: tensor<10xi64>) -> tensor<4xi64> {
+    %iota1 = stablehlo.iota dim = 0 : tensor<2x2xi64>
+    %iota2 = stablehlo.iota dim = 1 : tensor<2x2xi64>
+    %c2 = stablehlo.constant dense<2> : tensor<2x2xi64>
+    %scaled = stablehlo.multiply %iota1, %c2 : tensor<2x2xi64>
+    %added = stablehlo.add %scaled, %iota2 : tensor<2x2xi64>
+    %c_offset = stablehlo.constant dense<1> : tensor<2x2xi64>
+    %indices_2d = stablehlo.add %added, %c_offset : tensor<2x2xi64>
+    // indices_2d is [[1, 2],
+    //                [3, 4]]
+    %indices = stablehlo.reshape %indices_2d : (tensor<2x2xi64>) -> tensor<4x1xi64>
+    
+    %0 = "stablehlo.gather"(%arg0, %indices) {
+      dimension_numbers = #stablehlo.gather<
+        collapsed_slice_dims = [0],
+        start_index_map = [0],
+        index_vector_dim = 1
+      >,
+      slice_sizes = array<i64: 1>
+    } : (tensor<10xi64>, tensor<4x1xi64>) -> tensor<4xi64>
+    return %0 : tensor<4xi64>
+}
+// CHECK-LABEL: func.func @gather_reshaped_iota_to_slice
+// CHECK-NEXT: %[[SLICE:.+]] = stablehlo.slice %arg0 [1:5]
+// CHECK-NEXT: return %[[SLICE]]
+