clemisch/jaxtomo

Tomographic projector in JAX

jaxtomo's Issues

FP slow, dominated by MemcpyD2H

When profiling FP, runtime seems to be dominated by MemcpyD2H as seen in the Perfetto UI. The actual computations take only very little time in between.

I think this could be related to JAX calling each step in jax.lax.scan from host, requiring some synchronization at each iteration. This is discussed here and here.

I don't know how to resolve this in pure JAX. Alternatives to current implementation is vmap'ing instead of scan, or using some unroll in scan. IIRC both led to longer runtime.

Ultimately we should avoid the Memcpy at each projection -- which is what I think is happening. We could switch the levels of scan and vmap, i.e. vmap over angles, but scan over detector rows. I would expect that to be slower, but we could try it.

jit around pmap'ed FP & BP is a problem when pmap'ing SIR cost functions

The jit around pmap might lead to performance issues when using pmap on SIR cost functions.

When a user only uses FP or BP on its own, the jit seems to be faster. But it leads to synchronizing all data on every device, which is not desirable if any substantial computation comes afterwards. This could be filtering projections for FBP or computing cost function on projections for SIR.

So it might be better to disable automatic jit of pmap. We could instead provide a function for FBP (combining filtering and BP) and jit around that?