enzyme/test/MLIR/ForwardMode/while.mlir

// RUN: %eopt --enzyme %s | FileCheck %s

module {
  func.func @while(%x : f64) -> f64 {
    %cst = arith.constant 10.000000e+00 : f64
    %c0 = arith.constant 0 : index
    %c1 = arith.constant 1 : index
    %c10 = arith.constant 10 : index

    %r:2 = scf.while (%arg1 = %c0, %arg2 = %cst) : (index, f64) -> (index, f64) {
      %1 = arith.cmpi slt, %arg1, %c10 : index
      scf.condition(%1) %arg1, %arg2 : index, f64
    } do {
    ^bb0(%arg1: index, %arg2: f64):
      %1 = arith.addi %arg1, %c1 : index
      %2 = arith.addf %arg2, %x : f64
      scf.yield %1, %2 : index, f64
    }
    return %r#1 : f64
  }
  func.func @dwhile(%x : f64, %dx : f64) -> f64 {
    %r = enzyme.fwddiff @while(%x, %dx) <{ activity=[#enzyme.activity<enzyme_dup>], ret_activity=[#enzyme.activity<enzyme_dupnoneed>] }> : (f64, f64) -> (f64)
    return %r : f64
  }
  // CHECK: @fwddiffewhile
  // CHECK: (%[[arg0:.+]]: f64, %[[arg1:.+]]: f64) -> f64 {
  // CHECK:   %[[cst:.+]] = arith.constant 0.000000e+00 : f64
  // CHECK:   %[[cst_0:.+]] = arith.constant 1.000000e+01 : f64
  // CHECK:   %[[c0:.+]] = arith.constant 0 : index
  // CHECK:   %[[c1:.+]] = arith.constant 1 : index
  // CHECK:   %[[c10:.+]] = arith.constant 10 : index
  // CHECK:   %[[r0:.+]]:3 = scf.while (%[[arg2:.+]] = %[[c0]], %[[arg3:.+]] = %[[cst_0]], %[[arg4:.+]] = %[[cst]]) : (index, f64, f64) -> (index, f64, f64) {
  // CHECK:     %[[v1:.+]] = arith.cmpi slt, %[[arg2]], %[[c10]] : index
  // CHECK:     scf.condition(%[[v1]]) %[[arg2]], %[[arg3]], %[[arg4]] : index, f64, f64
  // CHECK:   } do {
  // CHECK:   ^bb0(%[[arg2:.+]]: index, %[[arg3:.+]]: f64, %[[arg4:.+]]: f64):
  // CHECK:     %[[v1:.+]] = arith.addi %[[arg2]], %[[c1]] : index
  // CHECK:     %[[v2:.+]] = arith.addf %[[arg4]], %[[arg1]] fastmath<fast> : f64
  // CHECK:     %[[v3:.+]] = arith.addf %[[arg3]], %[[arg0]] : f64
  // CHECK:     scf.yield %[[v1]], %[[v3]], %[[v2]] : index, f64, f64
  // CHECK:   }
  // CHECK:   return %[[r0]]#2 : f64
  // CHECK: }
}