Targeting JAX with a domain-specific compiler

[connorjward]

Hello,

I would like to use JAX as the target for the compiler of a domain-specific language and I would like to know what the best way is to do that.

To give an example, an expression in the DSL of interest (called UFL) could be

$$ \int f \ v \ \textrm{d}x $$

And in code this would look something like:

f = Coefficient(...)
v = TestFunction(...)
expr = f * v * dx

We would like to be able to generate JAX code from this which would look something like:

def compiled_expr(indata, outdata, gather_map, scatter_map):
    indata_packed = indata[gather_map]
    # 'f' represents a collection of JAX operations of things like einsum, pointwise additions
    #  and transcendental functions like sin and cos
    outdata_packed = f(indata_packed)
    return outdata.at[scatter_map].add(outdata_packed)

Currently the only approach I can think of for doing this is to construct the above function as a string, at which point one can do:

func = eval("""def compiled_expr(...):
...""")
callable = jax.jit(func)

but it would be much nicer if we could construct the expression directly in the IR itself. For example instead of producing "array_1 + array_2" it would be nice to be able to write jax.internals.ir.Add(array_1, array_2) and build up an expression tree. Is something like this possible?

[jakevdp]

This looks really interesting! We don't have any way to programmatically build an IR in a step-by-step fashion like this, but I wonder if it would be possible to do so via normal JAX operations in your context? The supported way to build an expression tree (i.e. a jaxpr) is by wrapping a function with an abstract evaluator like jax.make_jaxpr. Could you perhaps use this approach with the routine that is meant to generate the jax expression?

[connorjward]

Yeah that seems to work! Here's some code I've written that seems to do the right thing:

import jax
import jax.numpy as jnp


A = jnp.full((10, 3), 1)
B = jnp.full((10, 3), 2)


def generate(a, b):
    return add_one(a) + inv(b)


def add_one(a):
    return a + 1


def inv(b):
    return -b


jitted = jax.jit(generate)
traced = jitted.trace(A, B)
lowered = traced.lower()
compiled = lowered.compile()

And traced.jaxpr is

{ lambda ; a:i32[10,3] b:i32[10,3]. let
    c:i32[10,3] = add a 1:i32[]
    d:i32[10,3] = neg b
    e:i32[10,3] = add c d
  in (e,) }

So it seems like we can write our usual traversals calling normal JAX operations. I guess I didn't appreciate that when you JIT things you are really only looking for JAX operations and so the Python callstack isn't relevant. Thanks!