Generalize JIT call fusion using SSA dependencies

[yanzin00]

Summary

This PR adds a generic fuse-jit-calls pass that combines directly SSA-connected enzymexla.jit_call operations into a single LLVM wrapper and replaces the original calls with one fused JIT call.

• For each JIT call, the pass examines its operands to find producer JIT calls and its results to find consumer JIT calls. This discovers the connected SSA dependency graph, including chains and DAGs containing more than two calls.

• Operands defined outside the discovered group become arguments of the fused call. Results used only by another JIT call in the group become internal dependencies and are mapped directly between the cloned wrappers.Results that still have users outside the group remain outputs of the fused call.

• Before modifying the IR, the pass validates that fusion is safe and supported. Currently, all calls must be in the same block, operands must dominate the fused call, and fusion must not move calls across unrelated side effects. The referenced LLVM wrappers must use the supported single-block, pointer-only, non-vararg, void-returning ABI. Calls with unsupported backend, layout, argument, or result metadata are not fused. Each JIT result must currently have a supported flat output_operand_alias.

Test

The pass is tested in: test/lit_tests/mpi/fuse_jit_calls.mlir

The test first runs lower-enzymexla-mpi to lower the MPI operations into enzymexla.jit_call operations and LLVM wrapper functions. It then runs fuse-jit-calls and verifies that the dependent MPI_Irecv and MPI_Wait JIT calls are replaced by one fused JIT call.

IR before fuse-jit-calls

// -----// IR Dump Before FuseJITCallsPass (fuse-jit-calls) //----- //
module {
  llvm.mlir.global external constant @MPI_INT() {addr_space = 0 : i32} : !llvm.ptr
  llvm.mlir.global external constant @MPI_COMM_WORLD() {addr_space = 0 : i32} : !llvm.ptr
  llvm.func @MPI_Irecv(!llvm.ptr, i32, !llvm.ptr, i32, i32, !llvm.ptr, !llvm.ptr) -> i32
  llvm.func @enzymexla_wrapper_MPI_Irecv_MPI_INT(%arg0: !llvm.ptr {enzymexla.memory_effects = ["read", "write", "allocate", "free"]}, %arg1: !llvm.ptr {enzymexla.memory_effects = ["read", "write", "allocate", "free"]}, %arg2: !llvm.ptr {enzymexla.memory_effects = ["read", "write", "allocate", "free"]}, %arg3: !llvm.ptr {enzymexla.memory_effects = ["read", "write", "allocate", "free"]}, %arg4: !llvm.ptr {enzymexla.memory_effects = ["read", "write", "allocate", "free"]}) attributes {enzymexla.memory_effects = ["read", "write", "allocate", "free"]} {
    %0 = llvm.mlir.addressof @MPI_COMM_WORLD : !llvm.ptr
    %1 = llvm.mlir.addressof @MPI_INT : !llvm.ptr
    %2 = llvm.load %arg1 : !llvm.ptr -> i32
    %3 = llvm.load %arg2 : !llvm.ptr -> i32
    %4 = llvm.load %arg3 : !llvm.ptr -> i32
    %5 = llvm.call @MPI_Irecv(%arg0, %2, %1, %3, %4, %0, %arg4) : (!llvm.ptr, i32, !llvm.ptr, i32, i32, !llvm.ptr, !llvm.ptr) -> i32
    llvm.return
  }
  llvm.func @MPI_Wait(!llvm.ptr, !llvm.ptr) -> i32
  llvm.func @enzymexla_wrapper_MPI_Wait(%arg0: !llvm.ptr {enzymexla.memory_effects = ["read", "write", "allocate", "free"]}) attributes {enzymexla.memory_effects = ["read", "write", "allocate", "free"]} {
    %c1_i32 = arith.constant 1 : i32
    %0 = llvm.alloca %c1_i32 x !llvm.array<6 x i32> : (i32) -> !llvm.ptr
    %1 = llvm.call @MPI_Wait(%arg0, %0) : (!llvm.ptr, !llvm.ptr) -> i32
    llvm.return
  }
  func.func @main(%arg0: tensor<5xf64> {enzymexla.memory_effects = ["read", "write", "allocate", "free"], tf.aliasing_output = 0 : i32}) -> tensor<5xf64> attributes {enzymexla.memory_effects = ["read", "write", "allocate", "free"]} {
    %c = stablehlo.constant dense<-1> : tensor<i32>
    %c_0 = stablehlo.constant dense<5> : tensor<i32>
    %c_1 = stablehlo.constant dense<42> : tensor<i32>
    %c_2 = stablehlo.constant dense<1> : tensor<i32>
    %0 = stablehlo.transpose %arg0, dims = [0] : (tensor<5xf64>) -> tensor<5xf64>
    %1:2 = enzymexla.jit_call @enzymexla_wrapper_MPI_Irecv_MPI_INT (%0, %c_0, %c_2, %c_1, %c) {output_operand_aliases = [#stablehlo.output_operand_alias<output_tuple_indices = [0], operand_index = 0, operand_tuple_indices = []>, #stablehlo.output_operand_alias<output_tuple_indices = [1], operand_index = 4, operand_tuple_indices = []>]} : (tensor<5xf64>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>) -> (tensor<5xf64>, tensor<i32>)
    enzymexla.jit_call @enzymexla_wrapper_MPI_Wait (%1#1) : (tensor<i32>) -> ()
    %2 = stablehlo.transpose %1#0, dims = [0] : (tensor<5xf64>) -> tensor<5xf64>
    return %2 : tensor<5xf64>
  }
}

Before fusion, MPI_Irecv and MPI_Wait are represented by separate JIT calls. The request result %1#1 from MPI_Irecv is passed directly to MPI_Wait, while the receive buffer %1#0 remains externally used by stablehlo.transpose.

IR after fuse-jit-calls

// -----// IR Dump After FuseJITCallsPass (fuse-jit-calls) //----- //
module {
  llvm.mlir.global external constant @MPI_INT() {addr_space = 0 : i32} : !llvm.ptr
  llvm.mlir.global external constant @MPI_COMM_WORLD() {addr_space = 0 : i32} : !llvm.ptr
  llvm.func @MPI_Irecv(!llvm.ptr, i32, !llvm.ptr, i32, i32, !llvm.ptr, !llvm.ptr) -> i32
  llvm.func @__enzyme_fused_enzymexla_wrapper_MPI_Irecv_MPI_INT_enzymexla_wrapper_MPI_Wait(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !llvm.ptr, %arg3: !llvm.ptr, %arg4: !llvm.ptr) {
    %c1_i32 = arith.constant 1 : i32
    %0 = llvm.mlir.addressof @MPI_COMM_WORLD : !llvm.ptr
    %1 = llvm.mlir.addressof @MPI_INT : !llvm.ptr
    %2 = llvm.load %arg1 : !llvm.ptr -> i32
    %3 = llvm.load %arg2 : !llvm.ptr -> i32
    %4 = llvm.load %arg3 : !llvm.ptr -> i32
    %5 = llvm.call @MPI_Irecv(%arg0, %2, %1, %3, %4, %0, %arg4) : (!llvm.ptr, i32, !llvm.ptr, i32, i32, !llvm.ptr, !llvm.ptr) -> i32
    %6 = llvm.alloca %c1_i32 x !llvm.array<6 x i32> : (i32) -> !llvm.ptr
    %7 = llvm.call @MPI_Wait(%arg4, %6) : (!llvm.ptr, !llvm.ptr) -> i32
    llvm.return
  }
  llvm.func @enzymexla_wrapper_MPI_Irecv_MPI_INT(%arg0: !llvm.ptr {enzymexla.memory_effects = ["read", "write", "allocate", "free"]}, %arg1: !llvm.ptr {enzymexla.memory_effects = ["read", "write", "allocate", "free"]}, %arg2: !llvm.ptr {enzymexla.memory_effects = ["read", "write", "allocate", "free"]}, %arg3: !llvm.ptr {enzymexla.memory_effects = ["read", "write", "allocate", "free"]}, %arg4: !llvm.ptr {enzymexla.memory_effects = ["read", "write", "allocate", "free"]}) attributes {enzymexla.memory_effects = ["read", "write", "allocate", "free"]} {
    %0 = llvm.mlir.addressof @MPI_COMM_WORLD : !llvm.ptr
    %1 = llvm.mlir.addressof @MPI_INT : !llvm.ptr
    %2 = llvm.load %arg1 : !llvm.ptr -> i32
    %3 = llvm.load %arg2 : !llvm.ptr -> i32
    %4 = llvm.load %arg3 : !llvm.ptr -> i32
    %5 = llvm.call @MPI_Irecv(%arg0, %2, %1, %3, %4, %0, %arg4) : (!llvm.ptr, i32, !llvm.ptr, i32, i32, !llvm.ptr, !llvm.ptr) -> i32
    llvm.return
  }
  llvm.func @MPI_Wait(!llvm.ptr, !llvm.ptr) -> i32
  llvm.func @enzymexla_wrapper_MPI_Wait(%arg0: !llvm.ptr {enzymexla.memory_effects = ["read", "write", "allocate", "free"]}) attributes {enzymexla.memory_effects = ["read", "write", "allocate", "free"]} {
    %c1_i32 = arith.constant 1 : i32
    %0 = llvm.alloca %c1_i32 x !llvm.array<6 x i32> : (i32) -> !llvm.ptr
    %1 = llvm.call @MPI_Wait(%arg0, %0) : (!llvm.ptr, !llvm.ptr) -> i32
    llvm.return
  }
  func.func @main(%arg0: tensor<5xf64> {enzymexla.memory_effects = ["read", "write", "allocate", "free"], tf.aliasing_output = 0 : i32}) -> tensor<5xf64> attributes {enzymexla.memory_effects = ["read", "write", "allocate", "free"]} {
    %c = stablehlo.constant dense<-1> : tensor<i32>
    %c_0 = stablehlo.constant dense<5> : tensor<i32>
    %c_1 = stablehlo.constant dense<42> : tensor<i32>
    %c_2 = stablehlo.constant dense<1> : tensor<i32>
    %0 = stablehlo.transpose %arg0, dims = [0] : (tensor<5xf64>) -> tensor<5xf64>
    %1 = enzymexla.jit_call @__enzyme_fused_enzymexla_wrapper_MPI_Irecv_MPI_INT_enzymexla_wrapper_MPI_Wait (%0, %c_0, %c_2, %c_1, %c) {output_operand_aliases = [#stablehlo.output_operand_alias<output_tuple_indices = [], operand_index = 0, operand_tuple_indices = []>]} : (tensor<5xf64>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<5xf64>
    %2 = stablehlo.transpose %1, dims = [0] : (tensor<5xf64>) -> tensor<5xf64>
    return %2 : tensor<5xf64>
  }
}

After fusion, both wrapper bodies are cloned into one fused LLVM wrapper, and the two original JIT calls are replaced by a single fused JIT call. The MPI request becomes an internal pointer shared between MPI_Irecv and MPI_Wait, so it is removed from the JIT results. The receive buffer remains externally visible, and its alias to operand 0 is preserved.

• The original LLVM wrappers remain temporarily and may be removed by a later symbol DCE pass if they are no longer referenced.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 31.73%. Comparing base (8d04eb0) to head (7522725).

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2625   +/-   ##
=======================================
  Coverage   31.73%   31.73%           
=======================================
  Files         227      227           
  Lines       44908    44908           
=======================================
  Hits        14253    14253           
  Misses      30655    30655           

☔ 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.