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
  • [x] ConstantOp
  • [ ] ~ConvertOp~
  • [ ] ConvolutionOp #151
  • [ ] ~CosineOp~
  • [ ] ~ClzOp~
  • [ ] CustomCallOp
  • [ ] ~DivOp~
  • [ ] DotGeneralOp
  • [ ] DynamicBroadcastInDimOp
  • [ ] DynamicConvOp
  • [ ] DynamicGatherOp
  • [ ] DynamicIotaOp
  • [ ] DynamicPadOp
  • [ ] DynamicReshapeOp
  • [ ] DynamicSliceOp
  • [ ] DynamicUpdateSliceOp
  • [ ] ~ExpOp~
  • [ ] Expm1Op
  • [ ] FftOp
  • [ ] ~FloorOp~
  • [ ] GatherOp
  • [ ] GetDimensionSizeOp
  • [ ] GetTupleElementOp
  • [x] IfOp #151
  • [ ] ~ImagOp~
  • [ ] InfeedOp
  • [ ] IotaOp
  • [ ] IsFiniteOp
  • [ ] ~LogOp~
  • [ ] ~Logp1Op~
  • [ ] ~LogisticOp~
  • [ ] MapOp
  • [ ] ~MaxOp~
  • [ ] ~MinOp~
  • [ ] ~MulOp~
  • [ ] ~NegateOp~
  • [ ] ~NotOp~
  • [ ] OptimizationBarrierOp
  • [ ] ~OrOp~
  • [ ] OutfeedOp
  • [ ] PadOp
  • [ ] PartitionIdOp
  • [ ] ~PopcntOp~
  • [ ] ~PowOp~
  • [ ] ~RealOp~
  • [ ] RecvOp
  • [ ] ReduceOp
  • [ ] ReducePrecisionOp
  • [ ] ReduceScatterOp
  • [ ] ReduceWindowOp
  • [ ] ~RemainderOp~
  • [ ] ReplicaIdOp
  • [ ] ReshapeOp
  • [x] ReverseOp #151
  • [ ] RngOp
  • [ ] RngBitGeneratorOp
  • [ ] ~RoundOp~
  • [ ] ~RoundNearestEvenOp~
  • [ ] ~RsqrtOp~
  • [ ] ScatterOp
  • [ ] SelectOp
  • [ ] SelectAndScatterOp
  • [ ] SendOp
  • [ ] ~ShiftLeftOp~
  • [ ] ~ShiftRightArithmeticOp~
  • [ ] ~ShiftRightLogicalOp~
  • [ ] SignO~p~
  • [ ] ~SineOp~
  • [ ] SliceOp
  • [ ] SortOp
  • [ ] ~SqrtOp~
  • [ ] ~SubtractOp~
  • [ ] ~TanhOp~
  • [x] TransposeOp
  • [ ] TriangularSolveOp
  • [ ] TupleOp
  • [ ] UniformDequantizeOp
  • [ ] UniformQuantizeOp
  • [x] WhileOp #151
  • [ ] ~XorOp~
  • [ ] Deprecated operations in StableHLO
  • [ ] BroadcastOp
  • [ ] CreateTokenOp
  • [ ] CrossReplicaSumOp
  • [ ] DotOp
  • [ ] EinsumOp
  • ~TorchIndexSelectOp~
  • [ ] UnaryEinsumOp
• [ ] 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

[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]