rename loop -> for_loop

Author: sbrantq

src/enzyme_ad/jax/Passes/LowerEnzymeProbProg.cpp

@@ -2209,17 +2209,17 @@ struct GetFlattenedSamplesFromTraceOpConversion
   }
 };
 
-struct LoopOpConversion : public OpConversionPattern<enzyme::LoopOp> {
+struct ForLoopOpConversion : public OpConversionPattern<enzyme::ForLoopOp> {
   using OpConversionPattern::OpConversionPattern;
 
   std::string backend;
-  LoopOpConversion(std::string backend, TypeConverter &typeConverter,
-                   MLIRContext *context, PatternBenefit benefit = 1)
+  ForLoopOpConversion(std::string backend, TypeConverter &typeConverter,
+                      MLIRContext *context, PatternBenefit benefit = 1)
       : OpConversionPattern(typeConverter, context, benefit), backend(backend) {
   }
 
   LogicalResult
-  matchAndRewrite(enzyme::LoopOp op, OpAdaptor adaptor,
+  matchAndRewrite(enzyme::ForLoopOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     SmallVector<Value> initVals = {adaptor.getLowerBound()};
     initVals.append(adaptor.getInitArgs().begin(), adaptor.getInitArgs().end());
@@ -2365,14 +2365,14 @@ struct LowerProbProgToStableHLOPass
     target.addIllegalOp<enzyme::CholeskySolveOp>();
     target.addIllegalOp<enzyme::DotOp>();
     target.addIllegalOp<enzyme::UnflattenSliceOp>();
-    target.addIllegalOp<enzyme::LoopOp>();
+    target.addIllegalOp<enzyme::ForLoopOp>();
 
     target.addLegalOp<UnrealizedConversionCastOp>();
 
     RewritePatternSet patterns(context);
 
     patterns.add<RandomOpConversion, CholeskySolveOpConversion, DotOpConversion,
-                 UnflattenSliceOpConversion, LoopOpConversion>(
+                 UnflattenSliceOpConversion, ForLoopOpConversion>(
         backend, typeConverter, context);
 
     if (failed(applyPartialConversion(getOperation(), target,

test/lit_tests/probprog/hmc.mlir

@@ -42,7 +42,7 @@ module {
     } attributes {activity = [#enzyme<activity enzyme_active>], ret_activity = [#enzyme<activity enzyme_activenoneed>, #enzyme<activity enzyme_const>]} : (tensor<2xf64>, tensor<f64>) -> (tensor<2xui64>, tensor<2xf64>)
     %9 = "enzyme.broadcast"(%cst_6) <{shape = array<i64: 2>}> : (tensor<f64>) -> tensor<2xf64>
     %10 = "enzyme.broadcast"(%cst) <{shape = array<i64: 2>}> : (tensor<f64>) -> tensor<2xf64>
-    %11:4 = enzyme.loop(%cst_1 : tensor<i64>) to(%cst_5 : tensor<i64>) step(%cst_0 : tensor<i64>) iter_args(%1, %result, %8#1, %8#0 : tensor<2xf64>, tensor<2xf64>, tensor<2xf64>, tensor<2xui64>) -> tensor<2xf64>, tensor<2xf64>, tensor<2xf64>, tensor<2xui64> {
+    %11:4 = enzyme.for_loop(%cst_1 : tensor<i64>) to(%cst_5 : tensor<i64>) step(%cst_0 : tensor<i64>) iter_args(%1, %result, %8#1, %8#0 : tensor<2xf64>, tensor<2xf64>, tensor<2xf64>, tensor<2xui64>) -> tensor<2xf64>, tensor<2xf64>, tensor<2xf64>, tensor<2xui64> {
     ^bb0(%arg3: tensor<i64>, %arg4: tensor<2xf64>, %arg5: tensor<2xf64>, %arg6: tensor<2xf64>, %arg7: tensor<2xui64>):
       %23 = arith.mulf %10, %arg6 : tensor<2xf64>
       %24 = arith.subf %arg5, %23 : tensor<2xf64>

test/lit_tests/probprog/loop.mlir

@@ -23,7 +23,7 @@ module {
     %c1 = stablehlo.constant dense<1> : tensor<i64>
     %init = stablehlo.constant dense<0.0> : tensor<f64>
 
-    %result = enzyme.loop (%c0 : tensor<i64>) to (%n : tensor<i64>) step (%c1 : tensor<i64>)
+    %result = enzyme.for_loop (%c0 : tensor<i64>) to (%n : tensor<i64>) step (%c1 : tensor<i64>)
       iter_args(%init : tensor<f64>)
       -> tensor<f64> {
     ^bb0(%iv: tensor<i64>, %sum_iter: tensor<f64>):
@@ -64,11 +64,11 @@ module {
     %c1 = stablehlo.constant dense<1> : tensor<i64>
     %init = stablehlo.constant dense<0.0> : tensor<f64>
 
-    %result = enzyme.loop (%c0 : tensor<i64>) to (%m : tensor<i64>) step (%c1 : tensor<i64>)
+    %result = enzyme.for_loop (%c0 : tensor<i64>) to (%m : tensor<i64>) step (%c1 : tensor<i64>)
       iter_args(%init : tensor<f64>)
       -> tensor<f64> {
     ^bb0(%i: tensor<i64>, %outer_sum_iter: tensor<f64>):
-      %inner_result = enzyme.loop (%c0 : tensor<i64>) to (%n : tensor<i64>) step (%c1 : tensor<i64>)
+      %inner_result = enzyme.for_loop (%c0 : tensor<i64>) to (%n : tensor<i64>) step (%c1 : tensor<i64>)
         iter_args(%outer_sum_iter : tensor<f64>)
         -> tensor<f64> {
       ^bb1(%j: tensor<i64>, %inner_sum_iter: tensor<f64>):
@@ -108,7 +108,7 @@ module {
     %init_sum = stablehlo.constant dense<0.0> : tensor<f64>
     %init_prod = stablehlo.constant dense<1.0> : tensor<f64>
 
-    %sum, %prod = enzyme.loop (%c0 : tensor<i64>) to (%n : tensor<i64>) step (%c1 : tensor<i64>)
+    %sum, %prod = enzyme.for_loop (%c0 : tensor<i64>) to (%n : tensor<i64>) step (%c1 : tensor<i64>)
       iter_args(%init_sum, %init_prod : tensor<f64>, tensor<f64>)
       -> tensor<f64>, tensor<f64> {
     ^bb0(%iv: tensor<i64>, %s_iter: tensor<f64>, %p_iter: tensor<f64>):
@@ -173,7 +173,7 @@ module {
     %init_sum = stablehlo.constant dense<0.0> : tensor<f64>
     %init_trace = enzyme.initTrace : !enzyme.Trace
 
-    %sum, %trace = enzyme.loop (%c0 : tensor<i64>) to (%n : tensor<i64>) step (%c1 : tensor<i64>)
+    %sum, %trace = enzyme.for_loop (%c0 : tensor<i64>) to (%n : tensor<i64>) step (%c1 : tensor<i64>)
       iter_args(%init_sum, %init_trace : tensor<f64>, !enzyme.Trace)
       -> tensor<f64>, !enzyme.Trace {
     ^bb0(%iv: tensor<i64>, %s_iter: tensor<f64>, %t_iter: !enzyme.Trace):