Tracking issue for missing Batch Op Interface

[Pangoraw]

NOTE: Strikethrough ops are deliberately not implemented because the default broadcasting behavior of enzyme batch is enough.

• StableHLO
  • AbsOp
  • AddOp
  • AfterAllOp
  • AllGatherOp
  • AllReduceOp
  • AllToAllOp
  • AndOp
  • Atan2Op
  • BatchNormGradOp
  • BatchNormInferenceOp
  • BatchNormTrainingOp
  • BitcastConvertOp
  • BroadcastInDimOp
  • CaseOp
  • CbrtOp
  • CeilOp
  • CholeskyOp
  • ClampOp
  • CollectiveBroadcastOp
  • CollectivePermuteOp
  • CompareOp
  • ComplexOp
  • CompositeOp
  • ConcatenateOp
  • ConstantOp
  • ConvertOp
  • ConvolutionOp Generic batch op interface #151
  • CosineOp
  • ClzOp
  • CustomCallOp
  • DivOp
  • DotGeneralOp
  • DynamicBroadcastInDimOp
  • DynamicConvOp
  • DynamicGatherOp
  • DynamicIotaOp
  • DynamicPadOp
  • DynamicReshapeOp
  • DynamicSliceOp
  • DynamicUpdateSliceOp
  • ExpOp
  • Expm1Op
  • FftOp
  • FloorOp
  • GatherOp
  • GetDimensionSizeOp
  • IfOp Generic batch op interface #151
  • ImagOp
  • InfeedOp
  • IotaOp
  • IsFiniteOp
  • LogOp
  • Logp1Op
  • LogisticOp
  • MaxOp
  • MinOp
  • MulOp
  • NegateOp
  • NotOp
  • OptimizationBarrierOp
  • OrOp
  • OutfeedOp
  • PadOp
  • PartitionIdOp
  • PopcntOp
  • PowOp
  • RealOp
  • RecvOp
  • ReduceOp
  • ReducePrecisionOp
  • ReduceScatterOp
  • ReduceWindowOp
  • RemainderOp
  • ReplicaIdOp
  • ReshapeOp
  • ReverseOp Generic batch op interface #151
  • RngBitGeneratorOp
  • RoundOp
  • RoundNearestEvenOp
  • RsqrtOp
  • ScatterOp
  • SelectOp
  • SelectAndScatterOp
  • SendOp
  • ShiftLeftOp
  • ShiftRightArithmeticOp
  • ShiftRightLogicalOp
  • SignOp
  • SineOp
  • SliceOp
  • SortOp
  • SqrtOp
  • SubtractOp
  • TanhOp
  • TransposeOp
  • TriangularSolveOp
  • UniformDequantizeOp
  • UniformQuantizeOp
  • WhileOp Generic batch op interface #151
  • XorOp
  • Deprecated operations in StableHLO
    • RngOp
    • GetTupleElementOp
    • BroadcastOp
    • CreateTokenOp
    • CrossReplicaSumOp
    • DotOp
    • EinsumOp
    • TorchIndexSelectOp
    • UnaryEinsumOp
    • TupleOp
    • MapOp
• CHLO
  • Binary Element-wise Operations
    • BroadcastAddOp
    • BroadcastAtan2Op
    • BroadcastDivOp
    • BroadcastMaxOp
    • BroadcastMinOp
    • BroadcastMulOp
    • BroadcastNextAfterOp
    • BroadcastPolygammaOp
    • BroadcastPowOp
    • BroadcastRemOp
    • BroadcastShiftLeftOp
    • BroadcastShiftRightArithmeticOp
    • BroadcastShiftRightLogicalOp
    • BroadcastSubOp
    • BroadcastZetaOp
  • Binary Logical Element-wise Operations
    • BroadcastAndOp
    • BroadcastOrOp
    • BroadcastXorOp
  • Non-broadcasting Binary Operations
    • NextAfterOp
    • PolygammaOp
    • ZetaOp
  • ComplexOp
  • Unary Element-wise Operations
    • AcosOp
    • AcoshOp
    • AsinOp
    • AsinhOp
    • AtanOp
    • AtanhOp
    • BesselI1eOp
    • ConjOp
    • CoshOp
    • SinhOp
    • TanOp
    • ConstantOp (shared with StableHLO_ConstantOp)
    • ConstantLikeOp
    • DigammaOp
    • ErfOp
    • ErfInvOp
    • ErfcOp
    • IsInfOp
    • IsNegInfOp
    • IsPosInfOp
    • LgammaOp
  • BroadcastCompareOp
  • BroadcastSelectOp
  • TopKOp
• EnzymeXLA
  • KernelCallOp
  • JitCallOp
  • GetStreamOp
  • Memref2PointerOp
  • Pointer2MemrefOp
  • AffineScopeOp
  • RotateOp
  • WrapOp
  • ExtendOp
  • CommRegionOp
  • LUFactorizationOp

[mofeing]

... because the default broadcasting behavior of enzyme batch is enough.

Do you mean unrolling or leave the op unchanged?

[Pangoraw]

The batch pass will take the original op:

%0 = stablehlo.add %arg0, %arg1 : tensor<10xf32>

and just prepend the broadcasted dimensions (e.g. 20x4):

%0 = stablehlo.add %arg0, %arg1 : tensor<20x4x10xf32>

This is the default behavior for all ops unless they implement BatchOpInterface.

[wsmoses]

whileop definitely shouldn't be unrolled here in most cases [since it almost always has a number of iterations fixed by a constant aka non data-dependent value]

[avik-pal]

Should CustomCallOp always fallback to the looped version?

[wsmoses]

tenatively I think so, it at minimum feels like a better default [more likely to not error]