Profiling reactix codebase

[aysunrhn]

The actual computation happens within JAX. The instructions for running the code are handed to XLA, which compiles and runs the code. For this reason, we need to use multiple tools for profiling.

• Wall clock timing
• cProfile (python level profiling)
• JAX XLA trace

[aysunrhn]

Prepared a profiling script based on the transport-model.ipynb notebook.

1. Wall clock timing:

• Compiling takes ~1.25s, each run is ~0.23s. If we're running thousands of solves (like in MCMC) the compile cost becomes irrelevant. If we run it once (like in transport-model.ipynb) it's 5 times the runtime

2. cProfile

• Doesn't tell that much, the entire 0.23s is spent inside the compiled XLA kernel

3. JAX XLA trace

with jax.profiler.trace(str(trace_dir), create_perfetto_link=False):
    solve()

• Built in jax tracer, it record every low level operation (matrix multiplication, loop iterations, memory writes)
• Perfetto renders the trace in an interactive chart so you can zoom in individual second and see exactly which operation was running at each moment
• From what I can understand, reactix uses JAX cleanly, there are no additional costs that come from our side
  • While loop overhead is around ~429ms, this is unavoidable without changing the solver architecture
  • The tolerances are the only lever we can pull, but this would sacrifice accuracy I suppose

[aysunrhn]

One issue we can address in performance is the spike during the while loop:

• There are 3 phases from what I can tell:
  • The initial transient phase
  • Porosity discontinuity (the spike)
  • Post discontinuity
• This is because the reaction model uses a step change porosity:

porosity = jnp.ones(n_cells) * 0.3
porosity = porosity.at[100:].set(0.1)

• We could look into using a sigmoid instead of a step function unless the sharp interface is needed for physical reasons