Calling taichi kernels via primitives over custom_call

[femtomc]

Has anyone investigated a route for interop between taichi and JAX by utilizing custom_call infrastructure in XLA?

Are there immediate blockers to attempting this?

[sharadmv]

I actually did investigate adding Taichi custom calls but didn't made much progress and have not worked on it for a month or so. I am still struggling to understand how/where in their codebase Taichi launches the kernels. If you happen to know this, please let me know!

If you wanted to investigate this yourself, I think looking at jax-triton is a good starting point since it is a similar type of integration.

[femtomc]

Thank you @sharadmv -- good to know this is somewhat novel ground :)

[chaoming0625]

Dear @sharadmv , are you still working on Taichi? I think Taichi is much more flexible than triton when working on sparse data.

[femtomc]

@chaoming0625 I think the question is whether Taichi exposes a kernel_ptr that you could wrap in a custom_call invocation (if I'm understanding above).

In general, I think we want to wrap the kernel launch invocation in custom call -- we also need to provide shape invariants to that custom call, etc.

[chaoming0625]

Yes, I agree.

[femtomc]

I don't know enough about Taichi to know where to look for such a kernel_ptr, however :)

[chaoming0625]

I opened an issue in Taichi repo. Hoping Taichi developers can give us some insights.

[sharadmv]

Yes my initial struggle with the Taichi was in obtaining the kernel pointer, which seems to be hidden away behind some abstractions. Thank you for following up!