Help with upgrading legacy custom calls to FFI

[dionhaefner]

I maintain a few projects that extend JAX via custom Cython code, including Veros and mpi4jax. These projects have seen test failures when upgrading to JAX 0.8.0 since jax.interpreters.mlir.custom_call is now removed.

I've tried upgrading to jax.ffi but am getting stuck on a few things. Simply replacing custom_call with build_ffi_lowering_function doesn't seem to do the trick (see traceback below).

Full traceback

Traceback (most recent call last):
  File "/Users/dion/.virtualenvs/veros/bin/veros", line 8, in <module>
    sys.exit(cli())
  File "/Users/dion/.virtualenvs/veros/lib/python3.13/site-packages/click/core.py", line 1462, in __call__
    return self.main(*args, **kwargs)
  File "/Users/dion/.virtualenvs/veros/lib/python3.13/site-packages/click/core.py", line 1383, in main
    rv = self.invoke(ctx)
  File "/Users/dion/.virtualenvs/veros/lib/python3.13/site-packages/click/core.py", line 1850, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/Users/dion/.virtualenvs/veros/lib/python3.13/site-packages/click/core.py", line 1246, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/Users/dion/.virtualenvs/veros/lib/python3.13/site-packages/click/core.py", line 814, in invoke
    return callback(*args, **kwargs)
  File "/Users/dion/codes/veros/veros/cli/veros_run.py", line 152, in cli
    return run(setup_file, *args, **kwargs)
  File "/Users/dion/codes/veros/veros/cli/veros_run.py", line 87, in run
    sim.run()
  File "/Users/dion/codes/veros/veros/veros.py", line 355, in run
    self.step(self.state)
  File "/Users/dion/codes/veros/veros/routines.py", line 214, in __call__
    out = self.function(*args, **kwargs)
  File "/Users/dion/codes/veros/veros/veros.py", line 266, in step
    momentum.momentum(state)
  File "/Users/dion/codes/veros/veros/routines.py", line 214, in __call__
    out = self.function(*args, **kwargs)
  File "/Users/dion/codes/veros/veros/core/momentum.py", line 312, in momentum
    friction.friction(state)
  File "/Users/dion/codes/veros/veros/routines.py", line 214, in __call__
    out = self.function(*args, **kwargs)
  File "/Users/dion/codes/veros/veros/core/friction.py", line 987, in friction
    vs.update(implicit_vert_friction(state))
  File "/Users/dion/codes/veros/veros/routines.py", line 348, in __call__
    out = self.function(*args)
  File "/Users/dion/codes/veros/veros/core/friction.py", line 125, in implicit_vert_friction
    res = utilities.solve_implicit(a_tri, b_tri, c_tri, d_tri, water_mask, b_edge=b_tri_edge, edge_mask=edge_mask)
  File "/Users/dion/codes/veros/veros/routines.py", line 348, in __call__
    out = self.function(*args)
  File "/Users/dion/codes/veros/veros/core/utilities.py", line 59, in solve_implicit
    return solve_tridiagonal(a, b, c, d, water_mask, edge_mask)
  File "/Users/dion/codes/veros/veros/routines.py", line 348, in __call__
    out = self.function(*args)
  File "/Users/dion/codes/veros/veros/core/operators.py", line 131, in solve_tridiagonal_jax
    return tdma(a, b, c, d, water_mask, edge_mask)
  File "/Users/dion/codes/veros/veros/core/special/tdma_.py", line 61, in tdma
    return tdma_p.bind(a, b, c, d, system_depths)
jax._src.source_info_util.JaxStackTraceBeforeTransformation: AttributeError: 'jaxlib.mlir._mlir_libs._mlir.ir.BlockArgument' object has no attribute 'avals_out'

The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/dion/.virtualenvs/veros/bin/veros", line 8, in <module>
    sys.exit(cli())
             ~~~^^
  File "/Users/dion/.virtualenvs/veros/lib/python3.13/site-packages/click/core.py", line 1462, in __call__
    return self.main(*args, **kwargs)
           ~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "/Users/dion/.virtualenvs/veros/lib/python3.13/site-packages/click/core.py", line 1383, in main
    rv = self.invoke(ctx)
  File "/Users/dion/.virtualenvs/veros/lib/python3.13/site-packages/click/core.py", line 1850, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
                           ~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^
  File "/Users/dion/.virtualenvs/veros/lib/python3.13/site-packages/click/core.py", line 1246, in invoke
    return ctx.invoke(self.callback, **ctx.params)
           ~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/dion/.virtualenvs/veros/lib/python3.13/site-packages/click/core.py", line 814, in invoke
    return callback(*args, **kwargs)
  File "/Users/dion/codes/veros/veros/cli/veros_run.py", line 152, in cli
    return run(setup_file, *args, **kwargs)
  File "/Users/dion/codes/veros/veros/cli/veros_run.py", line 87, in run
    sim.run()
    ~~~~~~~^^
  File "/Users/dion/codes/veros/veros/veros.py", line 355, in run
    self.step(self.state)
    ~~~~~~~~~^^^^^^^^^^^^
  File "/Users/dion/codes/veros/veros/routines.py", line 214, in __call__
    out = self.function(*args, **kwargs)
  File "/Users/dion/codes/veros/veros/veros.py", line 266, in step
    momentum.momentum(state)
    ~~~~~~~~~~~~~~~~~^^^^^^^
  File "/Users/dion/codes/veros/veros/routines.py", line 214, in __call__
    out = self.function(*args, **kwargs)
  File "/Users/dion/codes/veros/veros/core/momentum.py", line 312, in momentum
    friction.friction(state)
    ~~~~~~~~~~~~~~~~~^^^^^^^
  File "/Users/dion/codes/veros/veros/routines.py", line 214, in __call__
    out = self.function(*args, **kwargs)
  File "/Users/dion/codes/veros/veros/core/friction.py", line 987, in friction
    vs.update(implicit_vert_friction(state))
              ~~~~~~~~~~~~~~~~~~~~~~^^^^^^^
  File "/Users/dion/codes/veros/veros/routines.py", line 348, in __call__
    out = self.function(*args)
  File "/Users/dion/codes/veros/veros/core/special/tdma_.py", line 114, in tdma_xla_encode_cpu
    out = custom_call(
        a,
    ...<5 lines>...
        np.int64(stride),
    )
  File "/Users/dion/.virtualenvs/veros/lib/python3.13/site-packages/jax/_src/ffi.py", line 210, in _lowering_op
    kwargs["result_types"] = [mlir.aval_to_ir_type(aval) for aval in ctx.avals_out]
                                                                     ^^^^^^^^^^^^^
AttributeError: 'jaxlib.mlir._mlir_libs._mlir.ir.BlockArgument' object has no attribute 'avals_out'
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

For the sake of this discussion, you can assume that I want to upgrade this file. Here's what I tried so far.

Happy for any pointers!

[danielsuo]

Hello! It seems like you're very close. You might try something like

operands = (a, b, c, d, system_depths,
            as_mhlo_constant(num_systems, np.int64),
            as_mhlo_constant(stride, np.int64),)
out = build_ffi_lowering_function(kernel, result_types=out_types)(ctx, *operands)

I'm not positive that'll work out of the box since build_ffi_lowering_function has some additional logic/checks that the bare mlir.custom_call doesn't have, but let me know.

For reference, you can take a look at build_ffi_lowering_function and jax._src.interpreters.mlir.custom_call to compare.

[dionhaefner]

Right, I was missing the ctx argument. Passing that with result_types gives this:

  File "/Users/dion/.virtualenvs/veros/lib/python3.13/site-packages/jax/_src/interpreters/mlir.py", line 3228, in custom_call
    assert len(result_layouts) == len(result_types), (
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError: ([(2, 1, 0)], [RankedTensorType(tensor<31x43x15xf64>), RankedTensorType(tensor<15xf64>)])

It seems that the crux of the matter is that build_ffi_lowering_function operates on ctx.avals_in and ctx.avals_out, which in my case is not identical to inputs and outputs of the custom call. For example, the two constants as_mhlo_constant(num_systems, np.int64), as_mhlo_constant(stride, np.int64) are not inputs to the primitive but do get passed to the custom call. I can fix this quite easily for outputs but not passing those constants would be a bigger refactor.

This is the traceback I get now that complains about missing layouts for the two constants:

  File "/Users/dion/.virtualenvs/veros/lib/python3.13/site-packages/jax/_src/interpreters/mlir.py", line 1354, in lower_jaxpr_to_module
    raise ValueError("\n".join(msg_lines) + "\n" +
                     dump_module_message(ctx.module, "verification")) from e
ValueError: Cannot lower jaxpr with verifier errors:
        'stablehlo.custom_call' op Number of operands must match the number of operand layouts, 7 != 5
                at loc("tdma"(callsite("tdma"("/Users/dion/codes/veros/veros/core/special/tdma_.py":61:15 to :53) at callsite("solve_tridiagonal_jax"("/Users/dion/codes/veros/veros/core/operators.py":131:15 to :54) at callsite("VerosKernel.__call__"("/Users/dion/codes/veros/veros/routines.py":348:22 to :42) at callsite("solve_implicit"("/Users/dion/codes/veros/veros/core/utilities.py":59:11 to :63) at callsite("VerosKernel.__call__"("/Users/dion/codes/veros/veros/routines.py":348:22 to :42) at callsite("implicit_vert_friction"("/Users/dion/codes/veros/veros/core/friction.py":125:10 to :114) at callsite("VerosKernel.__call__"("/Users/dion/codes/veros/veros/routines.py":348:22 to :42) at callsite("friction"("/Users/dion/codes/veros/veros/core/friction.py":987:18 to :47) at callsite("VerosRoutine.__call__"("/Users/dion/codes/veros/veros/routines.py":214:30 to :60) at "momentum"("/Users/dion/codes/veros/veros/core/momentum.py":312:8 to :32))))))))))))
        see current operation: %5:2 = "stablehlo.custom_call"(%arg0, %arg1, %arg2, %arg3, %2, %3, %4) <{api_version = 1 : i32, backend_config = "", call_target_name = "tdma_cython_double", called_computations = [], has_side_effect = false, operand_layouts = [dense<[2, 1, 0]> : tensor<3xindex>, dense<[2, 1, 0]> : tensor<3xindex>, dense<[2, 1, 0]> : tensor<3xindex>, dense<[2, 1, 0]> : tensor<3xindex>, dense<[1, 0]> : tensor<2xindex>], result_layouts = [dense<[2, 1, 0]> : tensor<3xindex>, dense<0> : tensor<1xindex>]}> {mhlo.backend_config = {}} : (tensor<31x43x15xf64>, tensor<31x43x15xf64>, tensor<31x43x15xf64>, tensor<31x43x15xf64>, tensor<31x43xi32>, tensor<1xi64>, tensor<1xi64>) -> (tensor<31x43x15xf64>, tensor<15xf64>)
                at loc("tdma"(callsite("tdma"("/Users/dion/codes/veros/veros/core/special/tdma_.py":61:15 to :53) at callsite("solve_tridiagonal_jax"("/Users/dion/codes/veros/veros/core/operators.py":131:15 to :54) at callsite("VerosKernel.__call__"("/Users/dion/codes/veros/veros/routines.py":348:22 to :42) at callsite("solve_implicit"("/Users/dion/codes/veros/veros/core/utilities.py":59:11 to :63) at callsite("VerosKernel.__call__"("/Users/dion/codes/veros/veros/routines.py":348:22 to :42) at callsite("implicit_vert_friction"("/Users/dion/codes/veros/veros/core/friction.py":125:10 to :114) at callsite("VerosKernel.__call__"("/Users/dion/codes/veros/veros/routines.py":348:22 to :42) at callsite("friction"("/Users/dion/codes/veros/veros/core/friction.py":987:18 to :47) at callsite("VerosRoutine.__call__"("/Users/dion/codes/veros/veros/routines.py":214:30 to :60) at "momentum"("/Users/dion/codes/veros/veros/core/momentum.py":312:8 to :32))))))))))))
        'stablehlo.custom_call' op Number of operands must match the number of operand layouts, 7 != 5
                at loc("tdma")
        see current operation: %2:2 = "stablehlo.custom_call"(%arg0, %arg1, %arg2, %arg3, %arg4, %0, %1) <{api_version = 1 : i32, backend_config = "", call_target_name = "tdma_cython_double", called_computations = [], has_side_effect = false, operand_layouts = [dense<[2, 1, 0]> : tensor<3xindex>, dense<[2, 1, 0]> : tensor<3xindex>, dense<[2, 1, 0]> : tensor<3xindex>, dense<[2, 1, 0]> : tensor<3xindex>, dense<[1, 0]> : tensor<2xindex>], result_layouts = [dense<[2, 1, 0]> : tensor<3xindex>, dense<0> : tensor<1xindex>]}> {mhlo.backend_config = {}} : (tensor<31x43x15xf64>, tensor<31x43x15xf64>, tensor<31x43x15xf64>, tensor<31x43x15xf64>, tensor<31x43xi32>, tensor<1xi64>, tensor<1xi64>) -> (tensor<31x43x15xf64>, tensor<15xf64>)
                at loc("tdma")
Define JAX_DUMP_IR_TO to dump the module.

Passing operand_layouts by hand doesn't seem to work either, same error.

[dionhaefner]

I also tried skip_ffi_layout_processing=True which gives:

jax.errors.JaxRuntimeError: NOT_FOUND: No FFI handler registered for tdma_cython_double on a platform Host (canonical host)

Although we do register that target in the same module.

[danielsuo]

Good call re skip_ffi_layout_processing=True. That gets you closer to a bare mlir.custom_call.

Let's see; one thing that jumps out to me is when you call jax.ffi.register_ffi_target, you have api_version=0, which I believe is the old version of the API and might not get registered appropriately with FFI. Perhaps api_version=1? See here.

Another possibility is jax.ffi.regiister_ffi_target doesn't get called for some reason under HAS_CPU_EXT, but presumably this isn't the case since it was working before.

[danielsuo]

Could you share your extension and any logs? Hoping it's just updating signatures, but I'm sure you've already looked.

[dionhaefner]

Here's one, but there are dozens more in mpi4jax: cpu, cuda, xpu.

There are no logs, just a violent crash. I haven't looked too closely yet on how to upgrade since using API version 0 worked fine.

[ezhulenev]

api_version=0 is deprecated and will be gone eventually, the C/C++ part of custom call also should be migrated, but it should not be too hard, you can look for examples in Jax itself

[dionhaefner]

Looks like we'll have to do some upgrading then. Thanks @danielsuo @ezhulenev for the support!

[dionhaefner]

On a closer look, the work to upgrade our routines is more than we can lift at the moment. So this will likely be the end for mpi4jax unless there are outside contributions. Which is generally OK since jax.distributed became a lot more powerful over the years, although there's no 1-1 feature parity.