Open penpornk opened 6 months ago
Hello @penpornk , Intel GPU is also following this new FFI design to extend the custom call, but we found some pain points in the current design after finishing a simple PoC:
window
https://github.com/openxla/xla/blob/7491cd61955271984d422b1de4338e8a1b70c09a/xla/client/xla_builder.cc#L2667-L2669 that are widely used in real custom callees. Could Google add the parsing of these members to the thunk process?backend_config_str
will be parsed to attribution map https://github.com/openxla/xla/blob/91a45ef9faf67d440d4d4c8e9dec5fc8c0f5929a/xla/service/gpu/ir_emitter_unnested.cc#L1407-L1418 in thunk process, but in our PoC this is failed because some escape characters in backend_config_str
can't be recognized. We will check whether the newest main branch can work.A follow-up of #1
in previous comments, seems we need to wrap stream
and other needed handlers as int64_t vars if the internal class ServiceExecutableRunOptions
are not preferred anymore.
About "[...] how to reimplement custom calls as FFI calls in jaxlib?", there is ongoing work currently on the CPU backend. Maybe you will find it useful: google/jax#21574
About "[...] how to reimplement custom calls as FFI calls in jaxlib?", there is ongoing work currently on the CPU backend. Maybe you will find it useful: google/jax#21574
Thanks, we are looking into it.
Hi @penpornk , we prepared a commit to reproduce the issue #3 mentioned in https://github.com/openxla/xla/issues/10059#issuecomment-2147712924, could you help to check it? We can submit a public issue if needed, thanks.
Previous XLA custom call API versions pass parameters as
void** buffers
. The new version, typed FFI, allows passing metadata such as data type and shape along with each parameter, which greatly improves the overall programmability and flexibility.Custom call API comparison between versions: https://github.com/openxla/xla/blob/76e1d730d4546b9a15c187102e3ee37734ff8ffa/xla/service/hlo.proto#L51-L111
XLA:GPU already has typed FFI support (part of the Thunks runtime work). Per our goal to share runtime components between XLA:GPU and XLA:CPU, we will be migrating XLA:CPU to use typed FFI for custom calls as well.
See more custom call examples (v1-v4) in issue https://github.com/openxla/xla/issues/8319.