Hoist loop-invariant `if` out of `while` by jumerckx · Pull Request #2586 · EnzymeAD/Enzyme-JAX

jumerckx · 2026-06-25T12:53:23Z

This pass transforms a while loop containing an if-else operation, into an if-else operation where each branch contains the while loop. This fires when the condition of the if-statement is loop-invariant.

On the code in EnzymeAD/Reactant.jl#2907 this manages to move the if out by one level, but that doesn't seem to be enough for further raising.

enzymexlamlir-opt  --enzyme-hlo-opt='licm_passes=true' issue_2907.mlir
module @reactant_stencil attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  func.func @main(%arg0: tensor<6x6xf64> {enzymexla.memory_effects = []}, %arg1: tensor<f64> {enzymexla.memory_effects = []}) -> tensor<f64> attributes {enzymexla.memory_effects = []} {
    %c = stablehlo.constant dense<2> : tensor<i64>
    %c_0 = stablehlo.constant dense<[true, true, true, true, true, false]> : tensor<6xi1>
    %c_1 = stablehlo.constant dense<[false, true, true, true, true, true]> : tensor<6xi1>
    %c_2 = stablehlo.constant dense<1> : tensor<i32>
    %cst = stablehlo.constant dense<4.000000e+00> : tensor<f64>
    %c_3 = stablehlo.constant dense<0> : tensor<i64>
    %c_4 = stablehlo.constant dense<6> : tensor<i64>
    %c_5 = stablehlo.constant dense<1> : tensor<i64>
    %cst_6 = stablehlo.constant dense<0.000000e+00> : tensor<f64>
    %0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<6x6xf64>) -> tensor<6x6xf64>
    %1 = stablehlo.add %cst, %arg1 : tensor<f64>
    %2 = stablehlo.broadcast_in_dim %1, dims = [] : (tensor<f64>) -> tensor<6x6xf64>
    %3 = stablehlo.multiply %2, %arg0 : tensor<6x6xf64>
    %4:2 = stablehlo.while(%iterArg = %c_3, %iterArg_7 = %cst_6) : tensor<i64>, tensor<f64> attributes {enzyme.disable_mincut, enzymexla.symmetric_matrix = [#enzymexla<guaranteed UNKNOWN>, #enzymexla<guaranteed UNKNOWN>, #enzymexla<guaranteed UNKNOWN>, #enzymexla<guaranteed UNKNOWN>, #enzymexla<guaranteed NOTGUARANTEED>, #enzymexla<guaranteed UNKNOWN>, #enzymexla<guaranteed UNKNOWN>]}
    cond {
      %5 = stablehlo.compare LT, %iterArg, %c_4 : (tensor<i64>, tensor<i64>) -> tensor<i1>
      stablehlo.return %5 : tensor<i1>
    } do {
      %5 = stablehlo.add %c_5, %iterArg {enzymexla.bounds = [[1, 6]], enzymexla.non_negative = [#enzymexla<guaranteed NOTGUARANTEED>]} : tensor<i64>
      %6 = stablehlo.convert %5 {enzymexla.bounds = [[1, 6]]} : (tensor<i64>) -> tensor<i32>
      %7 = stablehlo.dynamic_slice %c_1, %iterArg, sizes = [1] : (tensor<6xi1>, tensor<i64>) -> tensor<1xi1>
      %8 = stablehlo.reshape %7 : (tensor<1xi1>) -> tensor<i1>
      %9 = stablehlo.dynamic_slice %c_0, %iterArg, sizes = [1] : (tensor<6xi1>, tensor<i64>) -> tensor<1xi1>
      %10 = stablehlo.reshape %9 : (tensor<1xi1>) -> tensor<i1>
      %11 = stablehlo.subtract %6, %c_2 {enzymexla.bounds = [[0, 5]]} : tensor<i32>
      %12 = "stablehlo.if"(%10) ({
        %13 = stablehlo.add %iterArg, %c {enzymexla.bounds = [[2, 7]]} : tensor<i64>
        %14 = stablehlo.convert %13 {enzymexla.bounds = [[2, 7]]} : (tensor<i64>) -> tensor<i32>
        %15 = stablehlo.subtract %14, %c_2 {enzymexla.bounds = [[1, 6]]} : tensor<i32>
        %16 = "stablehlo.if"(%8) ({
          %17 = stablehlo.convert %iterArg {enzymexla.bounds = [[0, 5]]} : (tensor<i64>) -> tensor<i32>
          %18 = stablehlo.subtract %17, %c_2 {enzymexla.bounds = [[-1, 4]]} : tensor<i32>
          %19:2 = stablehlo.while(%iterArg_8 = %c_3, %iterArg_9 = %iterArg_7) : tensor<i64>, tensor<f64> attributes {enzyme.disable_mincut}
          cond {
            %20 = stablehlo.compare LT, %iterArg_8, %c_4 : (tensor<i64>, tensor<i64>) -> tensor<i1>
            stablehlo.return %20 : tensor<i1>
          } do {
            %20 = stablehlo.add %c_5, %iterArg_8 {enzymexla.bounds = [[1, 6]], enzymexla.non_negative = [#enzymexla<guaranteed NOTGUARANTEED>]} : tensor<i64>
            %21 = stablehlo.convert %20 {enzymexla.bounds = [[1, 6]]} : (tensor<i64>) -> tensor<i32>
            %22 = stablehlo.subtract %21, %c_2 {enzymexla.bounds = [[0, 5]]} : tensor<i32>
            %23 = stablehlo.dynamic_slice %3, %iterArg, %iterArg_8, sizes = [1, 1] : (tensor<6x6xf64>, tensor<i64>, tensor<i64>) -> tensor<1x1xf64>
            %24 = stablehlo.reshape %23 : (tensor<1x1xf64>) -> tensor<f64>
            %25 = stablehlo.dynamic_slice %c_0, %iterArg_8, sizes = [1] : (tensor<6xi1>, tensor<i64>) -> tensor<1xi1>
            %26 = stablehlo.reshape %25 : (tensor<1xi1>) -> tensor<i1>
            %27 = "stablehlo.if"(%26) ({
              %41 = stablehlo.add %iterArg_8, %c {enzymexla.bounds = [[2, 7]]} : tensor<i64>
              %42 = stablehlo.convert %41 {enzymexla.bounds = [[2, 7]]} : (tensor<i64>) -> tensor<i32>
              %43 = stablehlo.subtract %42, %c_2 {enzymexla.bounds = [[1, 6]]} : tensor<i32>
              %44 = stablehlo.dynamic_slice %0, %43, %11, sizes = [1, 1] : (tensor<6x6xf64>, tensor<i32>, tensor<i32>) -> tensor<1x1xf64>
              %45 = stablehlo.reshape %44 : (tensor<1x1xf64>) -> tensor<f64>
              %46 = stablehlo.subtract %24, %45 : tensor<f64>
              stablehlo.return %46 : tensor<f64>
            }, {
              stablehlo.return %24 : tensor<f64>
            }) : (tensor<i1>) -> tensor<f64>
            %28 = stablehlo.dynamic_slice %0, %22, %15, sizes = [1, 1] : (tensor<6x6xf64>, tensor<i32>, tensor<i32>) -> tensor<1x1xf64>
            %29 = stablehlo.reshape %28 : (tensor<1x1xf64>) -> tensor<f64>
            %30 = stablehlo.subtract %27, %29 : tensor<f64>
            %31 = stablehlo.dynamic_slice %c_1, %iterArg_8, sizes = [1] : (tensor<6xi1>, tensor<i64>) -> tensor<1xi1>
            %32 = stablehlo.reshape %31 : (tensor<1xi1>) -> tensor<i1>
            %33 = "stablehlo.if"(%32) ({
              %41 = stablehlo.convert %iterArg_8 {enzymexla.bounds = [[0, 5]]} : (tensor<i64>) -> tensor<i32>
              %42 = stablehlo.subtract %41, %c_2 {enzymexla.bounds = [[-1, 4]]} : tensor<i32>
              %43 = stablehlo.dynamic_slice %0, %42, %11, sizes = [1, 1] : (tensor<6x6xf64>, tensor<i32>, tensor<i32>) -> tensor<1x1xf64>
              %44 = stablehlo.reshape %43 : (tensor<1x1xf64>) -> tensor<f64>
              %45 = stablehlo.subtract %30, %44 : tensor<f64>
              stablehlo.return %45 : tensor<f64>
            }, {
              stablehlo.return %30 : tensor<f64>
            }) : (tensor<i1>) -> tensor<f64>
            %34 = stablehlo.dynamic_slice %0, %22, %18, sizes = [1, 1] : (tensor<6x6xf64>, tensor<i32>, tensor<i32>) -> tensor<1x1xf64>
            %35 = stablehlo.reshape %34 : (tensor<1x1xf64>) -> tensor<f64>
            %36 = stablehlo.subtract %33, %35 : tensor<f64>
            %37 = stablehlo.dynamic_slice %0, %22, %11, sizes = [1, 1] : (tensor<6x6xf64>, tensor<i32>, tensor<i32>) -> tensor<1x1xf64>
            %38 = stablehlo.reshape %37 : (tensor<1x1xf64>) -> tensor<f64>
            %39 = stablehlo.multiply %36, %38 : tensor<f64>
            %40 = stablehlo.add %iterArg_9, %39 : tensor<f64>
            stablehlo.return %20, %40 : tensor<i64>, tensor<f64>
          }
          stablehlo.return %19#1 : tensor<f64>
        }, {
          %17:2 = stablehlo.while(%iterArg_8 = %c_3, %iterArg_9 = %iterArg_7) : tensor<i64>, tensor<f64> attributes {enzyme.disable_mincut}
          cond {
            %18 = stablehlo.compare LT, %iterArg_8, %c_4 : (tensor<i64>, tensor<i64>) -> tensor<i1>
            stablehlo.return %18 : tensor<i1>
          } do {
            %18 = stablehlo.add %c_5, %iterArg_8 {enzymexla.bounds = [[1, 6]], enzymexla.non_negative = [#enzymexla<guaranteed NOTGUARANTEED>]} : tensor<i64>
            %19 = stablehlo.convert %18 {enzymexla.bounds = [[1, 6]]} : (tensor<i64>) -> tensor<i32>
            %20 = stablehlo.subtract %19, %c_2 {enzymexla.bounds = [[0, 5]]} : tensor<i32>
            %21 = stablehlo.dynamic_slice %3, %iterArg, %iterArg_8, sizes = [1, 1] : (tensor<6x6xf64>, tensor<i64>, tensor<i64>) -> tensor<1x1xf64>
            %22 = stablehlo.reshape %21 : (tensor<1x1xf64>) -> tensor<f64>
            %23 = stablehlo.dynamic_slice %c_0, %iterArg_8, sizes = [1] : (tensor<6xi1>, tensor<i64>) -> tensor<1xi1>
            %24 = stablehlo.reshape %23 : (tensor<1xi1>) -> tensor<i1>
            %25 = "stablehlo.if"(%24) ({
              %36 = stablehlo.add %iterArg_8, %c {enzymexla.bounds = [[2, 7]]} : tensor<i64>
              %37 = stablehlo.convert %36 {enzymexla.bounds = [[2, 7]]} : (tensor<i64>) -> tensor<i32>
              %38 = stablehlo.subtract %37, %c_2 {enzymexla.bounds = [[1, 6]]} : tensor<i32>
              %39 = stablehlo.dynamic_slice %0, %38, %11, sizes = [1, 1] : (tensor<6x6xf64>, tensor<i32>, tensor<i32>) -> tensor<1x1xf64>
              %40 = stablehlo.reshape %39 : (tensor<1x1xf64>) -> tensor<f64>
              %41 = stablehlo.subtract %22, %40 : tensor<f64>
              stablehlo.return %41 : tensor<f64>
            }, {
              stablehlo.return %22 : tensor<f64>
            }) : (tensor<i1>) -> tensor<f64>
            %26 = stablehlo.dynamic_slice %0, %20, %15, sizes = [1, 1] : (tensor<6x6xf64>, tensor<i32>, tensor<i32>) -> tensor<1x1xf64>
            %27 = stablehlo.reshape %26 : (tensor<1x1xf64>) -> tensor<f64>
            %28 = stablehlo.subtract %25, %27 : tensor<f64>
            %29 = stablehlo.dynamic_slice %c_1, %iterArg_8, sizes = [1] : (tensor<6xi1>, tensor<i64>) -> tensor<1xi1>
            %30 = stablehlo.reshape %29 : (tensor<1xi1>) -> tensor<i1>
            %31 = "stablehlo.if"(%30) ({
              %36 = stablehlo.convert %iterArg_8 {enzymexla.bounds = [[0, 5]]} : (tensor<i64>) -> tensor<i32>
              %37 = stablehlo.subtract %36, %c_2 {enzymexla.bounds = [[-1, 4]]} : tensor<i32>
              %38 = stablehlo.dynamic_slice %0, %37, %11, sizes = [1, 1] : (tensor<6x6xf64>, tensor<i32>, tensor<i32>) -> tensor<1x1xf64>
              %39 = stablehlo.reshape %38 : (tensor<1x1xf64>) -> tensor<f64>
              %40 = stablehlo.subtract %28, %39 : tensor<f64>
              stablehlo.return %40 : tensor<f64>
            }, {
              stablehlo.return %28 : tensor<f64>
            }) : (tensor<i1>) -> tensor<f64>
            %32 = stablehlo.dynamic_slice %0, %20, %11, sizes = [1, 1] : (tensor<6x6xf64>, tensor<i32>, tensor<i32>) -> tensor<1x1xf64>
            %33 = stablehlo.reshape %32 : (tensor<1x1xf64>) -> tensor<f64>
            %34 = stablehlo.multiply %31, %33 : tensor<f64>
            %35 = stablehlo.add %iterArg_9, %34 : tensor<f64>
            stablehlo.return %18, %35 : tensor<i64>, tensor<f64>
          }
          stablehlo.return %17#1 : tensor<f64>
        }) : (tensor<i1>) -> tensor<f64>
        stablehlo.return %16 : tensor<f64>
      }, {
        %13 = "stablehlo.if"(%8) ({
          %14 = stablehlo.convert %iterArg {enzymexla.bounds = [[0, 5]]} : (tensor<i64>) -> tensor<i32>
          %15 = stablehlo.subtract %14, %c_2 {enzymexla.bounds = [[-1, 4]]} : tensor<i32>
          %16:2 = stablehlo.while(%iterArg_8 = %c_3, %iterArg_9 = %iterArg_7) : tensor<i64>, tensor<f64> attributes {enzyme.disable_mincut}
          cond {
            %17 = stablehlo.compare LT, %iterArg_8, %c_4 : (tensor<i64>, tensor<i64>) -> tensor<i1>
            stablehlo.return %17 : tensor<i1>
          } do {
            %17 = stablehlo.add %c_5, %iterArg_8 {enzymexla.bounds = [[1, 6]], enzymexla.non_negative = [#enzymexla<guaranteed NOTGUARANTEED>]} : tensor<i64>
            %18 = stablehlo.convert %17 {enzymexla.bounds = [[1, 6]]} : (tensor<i64>) -> tensor<i32>
            %19 = stablehlo.subtract %18, %c_2 {enzymexla.bounds = [[0, 5]]} : tensor<i32>
            %20 = stablehlo.dynamic_slice %3, %iterArg, %iterArg_8, sizes = [1, 1] : (tensor<6x6xf64>, tensor<i64>, tensor<i64>) -> tensor<1x1xf64>
            %21 = stablehlo.reshape %20 : (tensor<1x1xf64>) -> tensor<f64>
            %22 = stablehlo.dynamic_slice %c_0, %iterArg_8, sizes = [1] : (tensor<6xi1>, tensor<i64>) -> tensor<1xi1>
            %23 = stablehlo.reshape %22 : (tensor<1xi1>) -> tensor<i1>
            %24 = "stablehlo.if"(%23) ({
              %35 = stablehlo.add %iterArg_8, %c {enzymexla.bounds = [[2, 7]]} : tensor<i64>
              %36 = stablehlo.convert %35 {enzymexla.bounds = [[2, 7]]} : (tensor<i64>) -> tensor<i32>
              %37 = stablehlo.subtract %36, %c_2 {enzymexla.bounds = [[1, 6]]} : tensor<i32>
              %38 = stablehlo.dynamic_slice %0, %37, %11, sizes = [1, 1] : (tensor<6x6xf64>, tensor<i32>, tensor<i32>) -> tensor<1x1xf64>
              %39 = stablehlo.reshape %38 : (tensor<1x1xf64>) -> tensor<f64>
              %40 = stablehlo.subtract %21, %39 : tensor<f64>
              stablehlo.return %40 : tensor<f64>
            }, {
              stablehlo.return %21 : tensor<f64>
            }) : (tensor<i1>) -> tensor<f64>
            %25 = stablehlo.dynamic_slice %c_1, %iterArg_8, sizes = [1] : (tensor<6xi1>, tensor<i64>) -> tensor<1xi1>
            %26 = stablehlo.reshape %25 : (tensor<1xi1>) -> tensor<i1>
            %27 = "stablehlo.if"(%26) ({
              %35 = stablehlo.convert %iterArg_8 {enzymexla.bounds = [[0, 5]]} : (tensor<i64>) -> tensor<i32>
              %36 = stablehlo.subtract %35, %c_2 {enzymexla.bounds = [[-1, 4]]} : tensor<i32>
              %37 = stablehlo.dynamic_slice %0, %36, %11, sizes = [1, 1] : (tensor<6x6xf64>, tensor<i32>, tensor<i32>) -> tensor<1x1xf64>
              %38 = stablehlo.reshape %37 : (tensor<1x1xf64>) -> tensor<f64>
              %39 = stablehlo.subtract %24, %38 : tensor<f64>
              stablehlo.return %39 : tensor<f64>
            }, {
              stablehlo.return %24 : tensor<f64>
            }) : (tensor<i1>) -> tensor<f64>
            %28 = stablehlo.dynamic_slice %0, %19, %15, sizes = [1, 1] : (tensor<6x6xf64>, tensor<i32>, tensor<i32>) -> tensor<1x1xf64>
            %29 = stablehlo.reshape %28 : (tensor<1x1xf64>) -> tensor<f64>
            %30 = stablehlo.subtract %27, %29 : tensor<f64>
            %31 = stablehlo.dynamic_slice %0, %19, %11, sizes = [1, 1] : (tensor<6x6xf64>, tensor<i32>, tensor<i32>) -> tensor<1x1xf64>
            %32 = stablehlo.reshape %31 : (tensor<1x1xf64>) -> tensor<f64>
            %33 = stablehlo.multiply %30, %32 : tensor<f64>
            %34 = stablehlo.add %iterArg_9, %33 : tensor<f64>
            stablehlo.return %17, %34 : tensor<i64>, tensor<f64>
          }
          stablehlo.return %16#1 : tensor<f64>
        }, {
          %14:2 = stablehlo.while(%iterArg_8 = %c_3, %iterArg_9 = %iterArg_7) : tensor<i64>, tensor<f64> attributes {enzyme.disable_mincut}
          cond {
            %15 = stablehlo.compare LT, %iterArg_8, %c_4 : (tensor<i64>, tensor<i64>) -> tensor<i1>
            stablehlo.return %15 : tensor<i1>
          } do {
            %15 = stablehlo.add %c_5, %iterArg_8 {enzymexla.bounds = [[1, 6]], enzymexla.non_negative = [#enzymexla<guaranteed NOTGUARANTEED>]} : tensor<i64>
            %16 = stablehlo.convert %15 {enzymexla.bounds = [[1, 6]]} : (tensor<i64>) -> tensor<i32>
            %17 = stablehlo.subtract %16, %c_2 {enzymexla.bounds = [[0, 5]]} : tensor<i32>
            %18 = stablehlo.dynamic_slice %3, %iterArg, %iterArg_8, sizes = [1, 1] : (tensor<6x6xf64>, tensor<i64>, tensor<i64>) -> tensor<1x1xf64>
            %19 = stablehlo.reshape %18 : (tensor<1x1xf64>) -> tensor<f64>
            %20 = stablehlo.dynamic_slice %c_0, %iterArg_8, sizes = [1] : (tensor<6xi1>, tensor<i64>) -> tensor<1xi1>
            %21 = stablehlo.reshape %20 : (tensor<1xi1>) -> tensor<i1>
            %22 = "stablehlo.if"(%21) ({
              %30 = stablehlo.add %iterArg_8, %c {enzymexla.bounds = [[2, 7]]} : tensor<i64>
              %31 = stablehlo.convert %30 {enzymexla.bounds = [[2, 7]]} : (tensor<i64>) -> tensor<i32>
              %32 = stablehlo.subtract %31, %c_2 {enzymexla.bounds = [[1, 6]]} : tensor<i32>
              %33 = stablehlo.dynamic_slice %0, %32, %11, sizes = [1, 1] : (tensor<6x6xf64>, tensor<i32>, tensor<i32>) -> tensor<1x1xf64>
              %34 = stablehlo.reshape %33 : (tensor<1x1xf64>) -> tensor<f64>
              %35 = stablehlo.subtract %19, %34 : tensor<f64>
              stablehlo.return %35 : tensor<f64>
            }, {
              stablehlo.return %19 : tensor<f64>
            }) : (tensor<i1>) -> tensor<f64>
            %23 = stablehlo.dynamic_slice %c_1, %iterArg_8, sizes = [1] : (tensor<6xi1>, tensor<i64>) -> tensor<1xi1>
            %24 = stablehlo.reshape %23 : (tensor<1xi1>) -> tensor<i1>
            %25 = "stablehlo.if"(%24) ({
              %30 = stablehlo.convert %iterArg_8 {enzymexla.bounds = [[0, 5]]} : (tensor<i64>) -> tensor<i32>
              %31 = stablehlo.subtract %30, %c_2 {enzymexla.bounds = [[-1, 4]]} : tensor<i32>
              %32 = stablehlo.dynamic_slice %0, %31, %11, sizes = [1, 1] : (tensor<6x6xf64>, tensor<i32>, tensor<i32>) -> tensor<1x1xf64>
              %33 = stablehlo.reshape %32 : (tensor<1x1xf64>) -> tensor<f64>
              %34 = stablehlo.subtract %22, %33 : tensor<f64>
              stablehlo.return %34 : tensor<f64>
            }, {
              stablehlo.return %22 : tensor<f64>
            }) : (tensor<i1>) -> tensor<f64>
            %26 = stablehlo.dynamic_slice %0, %17, %11, sizes = [1, 1] : (tensor<6x6xf64>, tensor<i32>, tensor<i32>) -> tensor<1x1xf64>
            %27 = stablehlo.reshape %26 : (tensor<1x1xf64>) -> tensor<f64>
            %28 = stablehlo.multiply %25, %27 : tensor<f64>
            %29 = stablehlo.add %iterArg_9, %28 : tensor<f64>
            stablehlo.return %15, %29 : tensor<i64>, tensor<f64>
          }
          stablehlo.return %14#1 : tensor<f64>
        }) : (tensor<i1>) -> tensor<f64>
        stablehlo.return %13 : tensor<f64>
      }) : (tensor<i1>) -> tensor<f64>
      stablehlo.return %5, %12 : tensor<i64>, tensor<f64>
    }
    return %4#1 : tensor<f64>
  }
}

codecov · 2026-06-25T13:55:34Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 31.08%. Comparing base (3e7a52c) to head (64cd746).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2586      +/-   ##
==========================================
+ Coverage   31.07%   31.08%   +0.01%     
==========================================
  Files         225      225              
  Lines       44504    44509       +5     
==========================================
+ Hits        13830    13837       +7     
+ Misses      30674    30672       -2

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Pangoraw

lgtm, maybe add it here

Enzyme-JAX/src/enzyme_ad/jax/Integrations/c/EnzymeXLA.cpp

Line 649 in 8a258d0

static void addLICMPasses(std::vector<std::string> &list) {

format

wsmoses

lgtm, lets try to get this and the other ones merged together quickly

jumerckx · 2026-07-03T18:43:10Z

I'm afk now but can try merge, and revise the other one later today or early tomorrow

jumerckx · 2026-07-04T21:46:58Z

unclear which comment is preventing merging here, I clicked resolve on all of them but still seeing "Merging is blocked due to pending merge requirements"

giordano · 2026-07-04T21:49:26Z

Required CI jobs are...well...required to merge the PR

jumerckx · 2026-07-04T21:51:45Z

fair 😅
I think before I resolved the comments, those were actually blocking, hadn't noticed the notification had changed, thanks

jumerckx · 2026-07-08T11:58:13Z

Sorry I lost track of this (again).
CI failures are unrelated, could this get merged?

I tried a bit reproducing the TPU issue on Hydra, without success. The macos failure seems to be because of this patch in XLA_PATCHES in workspace.bzl

echo '#include <cstdio>' >> xla/tsl/util/filewrapper.cc
echo 'namespace std { __attribute__((weak)) void __throw_bad_array_new_length() { fprintf(stderr, "erring in throw_bad_array_new_length\\n"); __builtin_trap(); } }' >> xla/tsl/util/filewrapper.cc

Claude proposes replacing this with:

echo '#include <cstdio>' >> xla/tsl/util/filewrapper.cc
echo '#include <version>' >> xla/tsl/util/filewrapper.cc
echo '#if !defined(_LIBCPP_VERSION)' >> xla/tsl/util/filewrapper.cc
echo 'namespace std { __attribute__((weak)) void __throw_bad_array_new_length() { fprintf(stderr, "erring in throw_bad_array_new_length\\n"); __builtin_trap(); } }' >> xla/tsl/util/filewrapper.cc
echo '#endif' >> xla/tsl/util/filewrapper.cc

I'm not sure if this is sensical, and these patches seem to be introduced by enzymead-bot. I'm not really familiar with how all this works.

Pangoraw · 2026-07-08T12:28:29Z

Merging this one, let's track the CI failure separately.

wsmoses reviewed Jun 25, 2026

View reviewed changes

Comment thread src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp Outdated

wsmoses reviewed Jun 25, 2026

View reviewed changes

Comment thread src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

wsmoses approved these changes Jun 25, 2026

View reviewed changes

Comment thread src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

jumerckx force-pushed the jm/while-if branch from ff3a871 to 7e4d3ef Compare June 25, 2026 21:59

jumerckx mentioned this pull request Jun 25, 2026

Extremely slow @trace for a stencil operation EnzymeAD/Reactant.jl#2907

Open

Pangoraw approved these changes Jul 1, 2026

View reviewed changes

jumerckx changed the title ~~Host loop-invariant if out of while~~ Hoist loop-invariant if out of while Jul 1, 2026

jumerckx added 6 commits July 1, 2026 06:13

WhileIfVersioning

6e3fe86

format

add to enzyme-hlo-opt pass

bc961be

test

12da472

rename WhileIfVersioning --> LoopUnswitch

a2eb24b

add threshold for applying

ea90ae3

add to EnzymeXLA addLICMPasses

3958c38

jumerckx force-pushed the jm/while-if branch from 7071f23 to 3958c38 Compare July 1, 2026 11:13

wsmoses reviewed Jul 1, 2026

View reviewed changes

Comment thread src/enzyme_ad/jax/Integrations/c/EnzymeXLA.cpp Outdated

wsmoses reviewed Jul 1, 2026

View reviewed changes

Comment thread src/enzyme_ad/jax/Integrations/c/EnzymeXLA.cpp

loop_unswitch_threshold

3072740

jumerckx mentioned this pull request Jul 3, 2026

Add loop_unswitch_threshold and excluded_passes options EnzymeAD/Reactant.jl#3025

Closed

2 tasks

wsmoses approved these changes Jul 3, 2026

View reviewed changes

Merge branch 'main' into jm/while-if

64cd746

Pangoraw merged commit e850e6b into main Jul 8, 2026
33 of 41 checks passed

Pangoraw deleted the jm/while-if branch July 8, 2026 12:28

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Hoist loop-invariant `if` out of `while`#2586

Hoist loop-invariant `if` out of `while`#2586
Pangoraw merged 8 commits into
mainfrom
jm/while-if

jumerckx commented Jun 25, 2026

Uh oh!

codecov Bot commented Jun 25, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Pangoraw left a comment

Uh oh!

Uh oh!

Uh oh!

wsmoses left a comment

Uh oh!

jumerckx commented Jul 3, 2026

Uh oh!

jumerckx commented Jul 4, 2026

Uh oh!

giordano commented Jul 4, 2026

Uh oh!

jumerckx commented Jul 4, 2026 •

edited

Loading

Uh oh!

jumerckx commented Jul 8, 2026

Uh oh!

Uh oh!

Pangoraw commented Jul 8, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

jumerckx commented Jun 25, 2026

Uh oh!

codecov Bot commented Jun 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Pangoraw left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

wsmoses left a comment

Choose a reason for hiding this comment

Uh oh!

jumerckx commented Jul 3, 2026

Uh oh!

jumerckx commented Jul 4, 2026

Uh oh!

giordano commented Jul 4, 2026

Uh oh!

jumerckx commented Jul 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jumerckx commented Jul 8, 2026

Uh oh!

Uh oh!

Pangoraw commented Jul 8, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

codecov Bot commented Jun 25, 2026 •

edited

Loading

jumerckx commented Jul 4, 2026 •

edited

Loading