warp-lang 1.10.0__py3-none-macosx_11_0_arm64.whl

warp/_src/jax_experimental/xla_ffi.py

@@ -0,0 +1,658 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import ctypes
+import enum
+
+import jax.numpy as jnp
+import numpy as np
+
+import warp as wp
+
+_wp_module_name_ = "warp.jax_experimental.xla_ffi"
+
+#######################################################################
+# ctypes structures and enums for XLA's FFI API:
+# https://github.com/openxla/xla/blob/a1a5e62fbffa3a3b6c409d72607456cf5b353a22/xla/ffi/api/c_api.h
+#######################################################################
+
+
+# typedef enum {
+#   XLA_FFI_Extension_Metadata = 1,
+# } XLA_FFI_Extension_Type;
+class XLA_FFI_Extension_Type(enum.IntEnum):
+    Metadata = 1
+
+
+# typedef struct XLA_FFI_Extension_Base {
+#   size_t struct_size;
+#   XLA_FFI_Extension_Type type;
+#   struct XLA_FFI_Extension_Base* next;
+# } XLA_FFI_Extension_Base;
+class XLA_FFI_Extension_Base(ctypes.Structure):
+    pass
+
+
+XLA_FFI_Extension_Base._fields_ = [
+    ("struct_size", ctypes.c_size_t),
+    ("type", ctypes.c_int),  # XLA_FFI_Extension_Type
+    ("next", ctypes.POINTER(XLA_FFI_Extension_Base)),
+]
+
+
+# typedef enum {
+#   XLA_FFI_ExecutionStage_INSTANTIATE = 0,
+#   XLA_FFI_ExecutionStage_PREPARE = 1,
+#   XLA_FFI_ExecutionStage_INITIALIZE = 2,
+#   XLA_FFI_ExecutionStage_EXECUTE = 3,
+# } XLA_FFI_ExecutionStage;
+class XLA_FFI_ExecutionStage(enum.IntEnum):
+    INSTANTIATE = 0
+    PREPARE = 1
+    INITIALIZE = 2
+    EXECUTE = 3
+
+
+# typedef enum {
+#   XLA_FFI_DataType_INVALID = 0,
+#   XLA_FFI_DataType_PRED = 1,
+#   XLA_FFI_DataType_S8 = 2,
+#   XLA_FFI_DataType_S16 = 3,
+#   XLA_FFI_DataType_S32 = 4,
+#   XLA_FFI_DataType_S64 = 5,
+#   XLA_FFI_DataType_U8 = 6,
+#   XLA_FFI_DataType_U16 = 7,
+#   XLA_FFI_DataType_U32 = 8,
+#   XLA_FFI_DataType_U64 = 9,
+#   XLA_FFI_DataType_F16 = 10,
+#   XLA_FFI_DataType_F32 = 11,
+#   XLA_FFI_DataType_F64 = 12,
+#   XLA_FFI_DataType_BF16 = 16,
+#   XLA_FFI_DataType_C64 = 15,
+#   XLA_FFI_DataType_C128 = 18,
+#   XLA_FFI_DataType_TOKEN = 17,
+#   XLA_FFI_DataType_F8E5M2 = 19,
+#   XLA_FFI_DataType_F8E3M4 = 29,
+#   XLA_FFI_DataType_F8E4M3 = 28,
+#   XLA_FFI_DataType_F8E4M3FN = 20,
+#   XLA_FFI_DataType_F8E4M3B11FNUZ = 23,
+#   XLA_FFI_DataType_F8E5M2FNUZ = 24,
+#   XLA_FFI_DataType_F8E4M3FNUZ = 25,
+#   XLA_FFI_DataType_F4E2M1FN = 32,
+#   XLA_FFI_DataType_F8E8M0FNU = 33,
+# } XLA_FFI_DataType;
+class XLA_FFI_DataType(enum.IntEnum):
+    INVALID = 0
+    PRED = 1
+    S8 = 2
+    S16 = 3
+    S32 = 4
+    S64 = 5
+    U8 = 6
+    U16 = 7
+    U32 = 8
+    U64 = 9
+    F16 = 10
+    F32 = 11
+    F64 = 12
+    BF16 = 16
+    C64 = 15
+    C128 = 18
+    TOKEN = 17
+    F8E5M2 = 19
+    F8E3M4 = 29
+    F8E4M3 = 28
+    F8E4M3FN = 20
+    F8E4M3B11FNUZ = 23
+    F8E5M2FNUZ = 24
+    F8E4M3FNUZ = 25
+    F4E2M1FN = 32
+    F8E8M0FNU = 33
+
+
+# struct XLA_FFI_Buffer {
+#   size_t struct_size;
+#   XLA_FFI_Extension_Base* extension_start;
+#
+#   XLA_FFI_DataType dtype;
+#   void* data;
+#   int64_t rank;
+#   int64_t* dims;  // length == rank
+# };
+class XLA_FFI_Buffer(ctypes.Structure):
+    _fields_ = (
+        ("struct_size", ctypes.c_size_t),
+        ("extension_start", ctypes.POINTER(XLA_FFI_Extension_Base)),
+        ("dtype", ctypes.c_int),  # XLA_FFI_DataType
+        ("data", ctypes.c_void_p),
+        ("rank", ctypes.c_int64),
+        ("dims", ctypes.POINTER(ctypes.c_int64)),
+    )
+
+
+# typedef enum {
+#   XLA_FFI_ArgType_BUFFER = 1,
+# } XLA_FFI_ArgType;
+class XLA_FFI_ArgType(enum.IntEnum):
+    BUFFER = 1
+
+
+# typedef enum {
+#   XLA_FFI_RetType_BUFFER = 1,
+# } XLA_FFI_RetType;
+class XLA_FFI_RetType(enum.IntEnum):
+    BUFFER = 1
+
+
+# struct XLA_FFI_Args {
+#   size_t struct_size;
+#   XLA_FFI_Extension_Base* extension_start;
+#   int64_t size;
+#   XLA_FFI_ArgType* types;  // length == size
+#   void** args;             // length == size
+# };
+class XLA_FFI_Args(ctypes.Structure):
+    _fields_ = (
+        ("struct_size", ctypes.c_size_t),
+        ("extension_start", ctypes.POINTER(XLA_FFI_Extension_Base)),
+        ("size", ctypes.c_int64),
+        ("types", ctypes.POINTER(ctypes.c_int)),  # XLA_FFI_ArgType*
+        ("args", ctypes.POINTER(ctypes.c_void_p)),
+    )
+
+
+# struct XLA_FFI_Rets {
+#   size_t struct_size;
+#   XLA_FFI_Extension_Base* extension_start;
+#   int64_t size;
+#   XLA_FFI_RetType* types;  // length == size
+#   void** rets;             // length == size
+# };
+class XLA_FFI_Rets(ctypes.Structure):
+    _fields_ = (
+        ("struct_size", ctypes.c_size_t),
+        ("extension_start", ctypes.POINTER(XLA_FFI_Extension_Base)),
+        ("size", ctypes.c_int64),
+        ("types", ctypes.POINTER(ctypes.c_int)),  # XLA_FFI_RetType*
+        ("rets", ctypes.POINTER(ctypes.c_void_p)),
+    )
+
+
+# typedef struct XLA_FFI_ByteSpan {
+#   const char* ptr;
+#   size_t len;
+# } XLA_FFI_ByteSpan;
+class XLA_FFI_ByteSpan(ctypes.Structure):
+    _fields_ = (
+        ("ptr", ctypes.POINTER(ctypes.c_char)),
+        ("len", ctypes.c_size_t),
+    )
+
+
+# typedef struct XLA_FFI_Scalar {
+#   XLA_FFI_DataType dtype;
+#   void* value;
+# } XLA_FFI_Scalar;
+class XLA_FFI_Scalar(ctypes.Structure):
+    _fields_ = (
+        ("dtype", ctypes.c_int),
+        ("value", ctypes.c_void_p),
+    )
+
+
+# typedef struct XLA_FFI_Array {
+#   XLA_FFI_DataType dtype;
+#   size_t size;
+#   void* data;
+# } XLA_FFI_Array;
+class XLA_FFI_Array(ctypes.Structure):
+    _fields_ = (
+        ("dtype", ctypes.c_int),
+        ("size", ctypes.c_size_t),
+        ("data", ctypes.c_void_p),
+    )
+
+
+# typedef enum {
+#   XLA_FFI_AttrType_ARRAY = 1,
+#   XLA_FFI_AttrType_DICTIONARY = 2,
+#   XLA_FFI_AttrType_SCALAR = 3,
+#   XLA_FFI_AttrType_STRING = 4,
+# } XLA_FFI_AttrType;
+class XLA_FFI_AttrType(enum.IntEnum):
+    ARRAY = 1
+    DICTIONARY = 2
+    SCALAR = 3
+    STRING = 4
+
+
+# struct XLA_FFI_Attrs {
+#   size_t struct_size;
+#   XLA_FFI_Extension_Base* extension_start;
+#   int64_t size;
+#   XLA_FFI_AttrType* types;   // length == size
+#   XLA_FFI_ByteSpan** names;  // length == size
+#   void** attrs;              // length == size
+# };
+class XLA_FFI_Attrs(ctypes.Structure):
+    _fields_ = (
+        ("struct_size", ctypes.c_size_t),
+        ("extension_start", ctypes.POINTER(XLA_FFI_Extension_Base)),
+        ("size", ctypes.c_int64),
+        ("types", ctypes.POINTER(ctypes.c_int)),  # XLA_FFI_AttrType*
+        ("names", ctypes.POINTER(ctypes.POINTER(XLA_FFI_ByteSpan))),
+        ("attrs", ctypes.POINTER(ctypes.c_void_p)),
+    )
+
+
+# struct XLA_FFI_Api_Version {
+#   size_t struct_size;
+#   XLA_FFI_Extension_Base* extension_start;
+#   int major_version;  // out
+#   int minor_version;  // out
+# };
+class XLA_FFI_Api_Version(ctypes.Structure):
+    _fields_ = (
+        ("struct_size", ctypes.c_size_t),
+        ("extension_start", ctypes.POINTER(XLA_FFI_Extension_Base)),
+        ("major_version", ctypes.c_int),
+        ("minor_version", ctypes.c_int),
+    )
+
+
+# enum XLA_FFI_Handler_TraitsBits {
+#   // Calls to FFI handler are safe to trace into the command buffer. It means
+#   // that calls to FFI handler always launch exactly the same device operations
+#   // (can depend on attribute values) that can be captured and then replayed.
+#   XLA_FFI_HANDLER_TRAITS_COMMAND_BUFFER_COMPATIBLE = 1u << 0,
+# };
+class XLA_FFI_Handler_TraitsBits(enum.IntEnum):
+    COMMAND_BUFFER_COMPATIBLE = 1 << 0
+
+
+# struct XLA_FFI_Metadata {
+#   size_t struct_size;
+#   XLA_FFI_Api_Version api_version;
+#   XLA_FFI_Handler_Traits traits;
+# };
+class XLA_FFI_Metadata(ctypes.Structure):
+    _fields_ = (
+        ("struct_size", ctypes.c_size_t),
+        ("api_version", XLA_FFI_Api_Version),  # XLA_FFI_Extension_Type
+        ("traits", ctypes.c_uint32),  # XLA_FFI_Handler_Traits
+    )
+
+
+# struct XLA_FFI_Metadata_Extension {
+#   XLA_FFI_Extension_Base extension_base;
+#   XLA_FFI_Metadata* metadata;
+# };
+class XLA_FFI_Metadata_Extension(ctypes.Structure):
+    _fields_ = (
+        ("extension_base", XLA_FFI_Extension_Base),
+        ("metadata", ctypes.POINTER(XLA_FFI_Metadata)),
+    )
+
+
+# typedef enum {
+#   XLA_FFI_Error_Code_OK = 0,
+#   XLA_FFI_Error_Code_CANCELLED = 1,
+#   XLA_FFI_Error_Code_UNKNOWN = 2,
+#   XLA_FFI_Error_Code_INVALID_ARGUMENT = 3,
+#   XLA_FFI_Error_Code_DEADLINE_EXCEEDED = 4,
+#   XLA_FFI_Error_Code_NOT_FOUND = 5,
+#   XLA_FFI_Error_Code_ALREADY_EXISTS = 6,
+#   XLA_FFI_Error_Code_PERMISSION_DENIED = 7,
+#   XLA_FFI_Error_Code_RESOURCE_EXHAUSTED = 8,
+#   XLA_FFI_Error_Code_FAILED_PRECONDITION = 9,
+#   XLA_FFI_Error_Code_ABORTED = 10,
+#   XLA_FFI_Error_Code_OUT_OF_RANGE = 11,
+#   XLA_FFI_Error_Code_UNIMPLEMENTED = 12,
+#   XLA_FFI_Error_Code_INTERNAL = 13,
+#   XLA_FFI_Error_Code_UNAVAILABLE = 14,
+#   XLA_FFI_Error_Code_DATA_LOSS = 15,
+#   XLA_FFI_Error_Code_UNAUTHENTICATED = 16
+# } XLA_FFI_Error_Code;
+class XLA_FFI_Error_Code(enum.IntEnum):
+    OK = 0
+    CANCELLED = 1
+    UNKNOWN = 2
+    INVALID_ARGUMENT = 3
+    DEADLINE_EXCEEDED = 4
+    NOT_FOUND = 5
+    ALREADY_EXISTS = 6
+    PERMISSION_DENIED = 7
+    RESOURCE_EXHAUSTED = 8
+    FAILED_PRECONDITION = 9
+    ABORTED = 10
+    OUT_OF_RANGE = 11
+    UNIMPLEMENTED = 12
+    INTERNAL = 13
+    UNAVAILABLE = 14
+    DATA_LOSS = 15
+    UNAUTHENTICATED = 16
+
+
+# struct XLA_FFI_Error_Create_Args {
+#   size_t struct_size;
+#   XLA_FFI_Extension_Base* extension_start;
+#   const char* message;
+#   XLA_FFI_Error_Code errc;
+# };
+class XLA_FFI_Error_Create_Args(ctypes.Structure):
+    _fields_ = (
+        ("struct_size", ctypes.c_size_t),
+        ("extension_start", ctypes.POINTER(XLA_FFI_Extension_Base)),
+        ("message", ctypes.c_char_p),
+        ("errc", ctypes.c_int),
+    )  # XLA_FFI_Error_Code
+
+
+XLA_FFI_Error_Create = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_Error_Create_Args))
+
+
+# struct XLA_FFI_Stream_Get_Args {
+#   size_t struct_size;
+#   XLA_FFI_Extension_Base* extension_start;
+#   XLA_FFI_ExecutionContext* ctx;
+#   void* stream;  // out
+# };
+class XLA_FFI_Stream_Get_Args(ctypes.Structure):
+    _fields_ = (
+        ("struct_size", ctypes.c_size_t),
+        ("extension_start", ctypes.POINTER(XLA_FFI_Extension_Base)),
+        ("ctx", ctypes.c_void_p),  # XLA_FFI_ExecutionContext*
+        ("stream", ctypes.c_void_p),
+    )  # // out
+
+
+XLA_FFI_Stream_Get = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_Stream_Get_Args))
+
+
+# struct XLA_FFI_DeviceOrdinal_Get {
+#   size_t struct_size;
+#   XLA_FFI_Extension_Base* extension_start;
+#   XLA_FFI_ExecutionContext* ctx;
+#   int32_t device_ordinal;  // out
+# };
+class XLA_FFI_DeviceOrdinal_Get_Args(ctypes.Structure):
+    _fields_ = (
+        ("struct_size", ctypes.c_size_t),
+        ("extension_start", ctypes.POINTER(XLA_FFI_Extension_Base)),
+        ("ctx", ctypes.c_void_p),  # XLA_FFI_ExecutionContext*
+        ("device_ordinal", ctypes.c_int32),
+    )  # // out
+
+
+XLA_FFI_DeviceOrdinal_Get = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_DeviceOrdinal_Get_Args))
+
+
+# struct XLA_FFI_Api {
+#   size_t struct_size;
+#   XLA_FFI_Extension_Base* extension_start;
+#
+#   XLA_FFI_Api_Version api_version;
+#   XLA_FFI_InternalApi* internal_api;
+#
+#   _XLA_FFI_API_STRUCT_FIELD(XLA_FFI_Error_Create);
+#   _XLA_FFI_API_STRUCT_FIELD(XLA_FFI_Error_GetMessage);
+#   _XLA_FFI_API_STRUCT_FIELD(XLA_FFI_Error_Destroy);
+#   _XLA_FFI_API_STRUCT_FIELD(XLA_FFI_Handler_Register);
+#   _XLA_FFI_API_STRUCT_FIELD(XLA_FFI_Stream_Get);
+#   _XLA_FFI_API_STRUCT_FIELD(XLA_FFI_TypeId_Register);
+#   _XLA_FFI_API_STRUCT_FIELD(XLA_FFI_ExecutionContext_Get);
+#   _XLA_FFI_API_STRUCT_FIELD(XLA_FFI_State_Set);
+#   _XLA_FFI_API_STRUCT_FIELD(XLA_FFI_State_Get);
+#   _XLA_FFI_API_STRUCT_FIELD(XLA_FFI_DeviceMemory_Allocate);
+#   _XLA_FFI_API_STRUCT_FIELD(XLA_FFI_DeviceMemory_Free);
+#   _XLA_FFI_API_STRUCT_FIELD(XLA_FFI_ThreadPool_Schedule);
+#   _XLA_FFI_API_STRUCT_FIELD(XLA_FFI_ThreadPool_NumThreads);
+#   _XLA_FFI_API_STRUCT_FIELD(XLA_FFI_Future_Create);
+#   _XLA_FFI_API_STRUCT_FIELD(XLA_FFI_Future_SetAvailable);
+#   _XLA_FFI_API_STRUCT_FIELD(XLA_FFI_Future_SetError);
+#   _XLA_FFI_API_STRUCT_FIELD(XLA_FFI_RunId_Get);
+#   _XLA_FFI_API_STRUCT_FIELD(XLA_FFI_DeviceOrdinal_Get);
+# };
+class XLA_FFI_Api(ctypes.Structure):
+    _fields_ = (
+        ("struct_size", ctypes.c_size_t),
+        ("extension_start", ctypes.POINTER(XLA_FFI_Extension_Base)),
+        ("api_version", XLA_FFI_Api_Version),
+        ("internal_api", ctypes.c_void_p),  # XLA_FFI_InternalApi*
+        ("XLA_FFI_Error_Create", XLA_FFI_Error_Create),  # XLA_FFI_Error_Create
+        ("XLA_FFI_Error_GetMessage", ctypes.c_void_p),  # XLA_FFI_Error_GetMessage
+        ("XLA_FFI_Error_Destroy", ctypes.c_void_p),  # XLA_FFI_Error_Destroy
+        ("XLA_FFI_Handler_Register", ctypes.c_void_p),  # XLA_FFI_Handler_Register
+        ("XLA_FFI_Stream_Get", XLA_FFI_Stream_Get),  # XLA_FFI_Stream_Get
+        ("XLA_FFI_TypeId_Register", ctypes.c_void_p),  # XLA_FFI_TypeId_Register
+        ("XLA_FFI_ExecutionContext_Get", ctypes.c_void_p),  # XLA_FFI_ExecutionContext_Get
+        ("XLA_FFI_State_Set", ctypes.c_void_p),  # XLA_FFI_State_Set
+        ("XLA_FFI_State_Get", ctypes.c_void_p),  # XLA_FFI_State_Get
+        ("XLA_FFI_DeviceMemory_Allocate", ctypes.c_void_p),  # XLA_FFI_DeviceMemory_Allocate
+        ("XLA_FFI_DeviceMemory_Free", ctypes.c_void_p),  # XLA_FFI_DeviceMemory_Free
+        ("XLA_FFI_ThreadPool_Schedule", ctypes.c_void_p),  # XLA_FFI_ThreadPool_Schedule
+        ("XLA_FFI_ThreadPool_NumThreads", ctypes.c_void_p),  # XLA_FFI_ThreadPool_NumThreads
+        ("XLA_FFI_Future_Create", ctypes.c_void_p),  # XLA_FFI_Future_Create
+        ("XLA_FFI_Future_SetAvailable", ctypes.c_void_p),  # XLA_FFI_Future_SetAvailable
+        ("XLA_FFI_Future_SetError", ctypes.c_void_p),  # XLA_FFI_Future_SetError
+        # TODO(chaserileyroberts): Make this return the correct value and not a c_void_p.
+        ("XLA_FFI_RunId_Get", ctypes.c_void_p),  # XLA_FFI_RunId_Get
+        ("XLA_FFI_DeviceOrdinal_Get", XLA_FFI_DeviceOrdinal_Get),  # XLA_FFI_DeviceOrdinal_Get
+    )
+
+
+# struct XLA_FFI_CallFrame {
+#   size_t struct_size;
+#   XLA_FFI_Extension_Base* extension_start;
+#   const XLA_FFI_Api* api;
+#   XLA_FFI_ExecutionContext* ctx;
+#   XLA_FFI_ExecutionStage stage;
+#   XLA_FFI_Args args;
+#   XLA_FFI_Rets rets;
+#   XLA_FFI_Attrs attrs;
+#
+#   // XLA FFI handler implementation can use `future` to signal a result of
+#   // asynchronous computation to the XLA runtime. XLA runtime will keep all
+#   // arguments, results and attributes alive until `future` is completed.
+#   XLA_FFI_Future* future;  // out
+# };
+class XLA_FFI_CallFrame(ctypes.Structure):
+    _fields_ = (
+        ("struct_size", ctypes.c_size_t),
+        ("extension_start", ctypes.POINTER(XLA_FFI_Extension_Base)),
+        ("api", ctypes.POINTER(XLA_FFI_Api)),
+        ("ctx", ctypes.c_void_p),  # XLA_FFI_ExecutionContext*
+        ("stage", ctypes.c_int),  # XLA_FFI_ExecutionStage
+        ("args", XLA_FFI_Args),
+        ("rets", XLA_FFI_Rets),
+        ("attrs", XLA_FFI_Attrs),
+        ("future", ctypes.c_void_p),  # XLA_FFI_Future* // out
+    )
+
+
+_xla_data_type_to_constructor = {
+    # XLA_FFI_DataType.INVALID
+    XLA_FFI_DataType.PRED: jnp.bool,
+    XLA_FFI_DataType.S8: jnp.int8,
+    XLA_FFI_DataType.S16: jnp.int16,
+    XLA_FFI_DataType.S32: jnp.int32,
+    XLA_FFI_DataType.S64: jnp.int64,
+    XLA_FFI_DataType.U8: jnp.uint8,
+    XLA_FFI_DataType.U16: jnp.uint16,
+    XLA_FFI_DataType.U32: jnp.uint32,
+    XLA_FFI_DataType.U64: jnp.uint64,
+    XLA_FFI_DataType.F16: jnp.float16,
+    XLA_FFI_DataType.F32: jnp.float32,
+    XLA_FFI_DataType.F64: jnp.float64,
+    XLA_FFI_DataType.BF16: jnp.bfloat16,
+    XLA_FFI_DataType.C64: jnp.complex64,
+    XLA_FFI_DataType.C128: jnp.complex128,
+    # XLA_FFI_DataType.TOKEN
+    # XLA_FFI_DataType.F4E2M1FN: jnp.float4_e2m1fn.dtype,
+    # XLA_FFI_DataType.F8E8M0FNU: jnp.float8_e8m0fnu.dtype,
+}
+
+# newer types not supported by older versions
+if hasattr(jnp, "float8_e5m2"):
+    _xla_data_type_to_constructor[XLA_FFI_DataType.F8E5M2] = jnp.float8_e5m2
+if hasattr(jnp, "float8_e3m4"):
+    _xla_data_type_to_constructor[XLA_FFI_DataType.F8E3M4] = jnp.float8_e3m4
+if hasattr(jnp, "float8_e4m3"):
+    _xla_data_type_to_constructor[XLA_FFI_DataType.F8E4M3] = jnp.float8_e4m3
+if hasattr(jnp, "float8_e4m3fn"):
+    _xla_data_type_to_constructor[XLA_FFI_DataType.F8E4M3FN] = jnp.float8_e4m3fn
+if hasattr(jnp, "float8_e4m3b11fnuz"):
+    _xla_data_type_to_constructor[XLA_FFI_DataType.F8E4M3B11FNUZ] = jnp.float8_e4m3b11fnuz
+if hasattr(jnp, "float8_e5m2fnuz"):
+    _xla_data_type_to_constructor[XLA_FFI_DataType.F8E5M2FNUZ] = jnp.float8_e5m2fnuz
+if hasattr(jnp, "float8_e4m3fnuz"):
+    _xla_data_type_to_constructor[XLA_FFI_DataType.F8E4M3FNUZ] = jnp.float8_e4m3fnuz
+
+
+########################################################################
+# Helpers for translating between ctypes and python types
+#######################################################################
+
+
+def decode_bytespan(span: XLA_FFI_ByteSpan):
+    len = span.len
+    chars = ctypes.cast(span.ptr, ctypes.POINTER(ctypes.c_char * len))
+    return chars.contents.value.decode("utf-8")
+
+
+def decode_scalar(scalar: XLA_FFI_Scalar):
+    # TODO validate if dtype supported
+    dtype = jnp.dtype(_xla_data_type_to_constructor[scalar.dtype])
+    bytes = ctypes.string_at(scalar.value, dtype.itemsize)
+    return np.frombuffer(bytes, dtype=dtype).reshape(())
+
+
+def decode_array(array: XLA_FFI_Array):
+    # TODO validate if dtype supported
+    dtype = jnp.dtype(_xla_data_type_to_constructor[array.dtype])
+    bytes = ctypes.string_at(array.data, dtype.itemsize * array.size)
+    return np.frombuffer(bytes, dtype=dtype)
+
+
+def decode_attrs(attrs: XLA_FFI_Attrs):
+    result = {}
+    for i in range(attrs.size):
+        attr_name = decode_bytespan(attrs.names[i].contents)
+        attr_type = attrs.types[i]
+        if attr_type == XLA_FFI_AttrType.STRING:
+            bytespan = ctypes.cast(attrs.attrs[i], ctypes.POINTER(XLA_FFI_ByteSpan))
+            attr_value = decode_bytespan(bytespan.contents)
+        elif attr_type == XLA_FFI_AttrType.SCALAR:
+            attr_value = ctypes.cast(attrs.attrs[i], ctypes.POINTER(XLA_FFI_Scalar))
+            attr_value = decode_scalar(attr_value.contents)
+        elif attr_type == XLA_FFI_AttrType.ARRAY:
+            attr_value = ctypes.cast(attrs.attrs[i], ctypes.POINTER(XLA_FFI_Array))
+            attr_value = decode_array(attr_value.contents)
+        elif attr_type == XLA_FFI_AttrType.DICTIONARY:
+            attr_value = ctypes.cast(attrs.attrs[i], ctypes.POINTER(XLA_FFI_Attrs))
+            attr_value = decode_attrs(attr_value.contents)
+        else:
+            raise Exception("Unexpected attr type")
+        result[attr_name] = attr_value
+    return result
+
+
+# error-string to XLA_FFI_Error
+def create_ffi_error(api, errc, message):
+    create_args = XLA_FFI_Error_Create_Args(
+        ctypes.sizeof(XLA_FFI_Error_Create_Args),
+        ctypes.POINTER(XLA_FFI_Extension_Base)(),
+        ctypes.c_char_p(message.encode("utf-8")),
+        errc,
+    )
+    return api.contents.XLA_FFI_Error_Create(create_args)
+
+
+def create_invalid_argument_ffi_error(api, message):
+    return create_ffi_error(api, XLA_FFI_Error_Code.INVALID_ARGUMENT, message)
+
+
+# Extract CUDA stream from XLA_FFI_CallFrame.
+def get_stream_from_callframe(call_frame):
+    api = call_frame.api
+    get_stream_args = XLA_FFI_Stream_Get_Args(
+        ctypes.sizeof(XLA_FFI_Stream_Get_Args), ctypes.POINTER(XLA_FFI_Extension_Base)(), call_frame.ctx, None
+    )
+    api.contents.XLA_FFI_Stream_Get(get_stream_args)
+    # TODO check result
+    return get_stream_args.stream
+
+
+def get_device_ordinal_from_callframe(call_frame):
+    api = call_frame.api
+    get_device_args = XLA_FFI_DeviceOrdinal_Get_Args(
+        ctypes.sizeof(XLA_FFI_DeviceOrdinal_Get_Args), ctypes.POINTER(XLA_FFI_Extension_Base)(), call_frame.ctx, 0
+    )
+    api.contents.XLA_FFI_DeviceOrdinal_Get(get_device_args)
+    return get_device_args.device_ordinal
+
+
+_dtype_from_ffi = {
+    XLA_FFI_DataType.S8: wp.int8,
+    XLA_FFI_DataType.S16: wp.int16,
+    XLA_FFI_DataType.S32: wp.int32,
+    XLA_FFI_DataType.S64: wp.int64,
+    XLA_FFI_DataType.U8: wp.uint8,
+    XLA_FFI_DataType.U16: wp.uint16,
+    XLA_FFI_DataType.U32: wp.uint32,
+    XLA_FFI_DataType.U64: wp.uint64,
+    XLA_FFI_DataType.F16: wp.float16,
+    XLA_FFI_DataType.F32: wp.float32,
+    XLA_FFI_DataType.F64: wp.float64,
+}
+
+
+def dtype_from_ffi(ffi_dtype):
+    return _dtype_from_ffi.get(ffi_dtype)
+
+
+def jax_dtype_from_ffi(ffi_dtype):
+    return _xla_data_type_to_constructor.get(ffi_dtype)
+
+
+# Execution context (stream, stage)
+class ExecutionContext:
+    stage: XLA_FFI_ExecutionStage
+    stream: int
+
+    def __init__(self, callframe: XLA_FFI_CallFrame):
+        self.stage = XLA_FFI_ExecutionStage(callframe.stage)
+        self.stream = get_stream_from_callframe(callframe)
+
+
+class FfiBuffer:
+    dtype: str
+    data: int
+    shape: tuple[int]
+
+    def __init__(self, xla_buffer):
+        # TODO check if valid
+        self.dtype = jnp.dtype(_xla_data_type_to_constructor[xla_buffer.dtype])
+        self.shape = tuple(xla_buffer.dims[i] for i in range(xla_buffer.rank))
+        self.data = xla_buffer.data
+
+    @property
+    def __cuda_array_interface__(self):
+        return {
+            "shape": self.shape,
+            "typestr": self.dtype.char,
+            "data": (self.data, False),
+            "version": 2,
+        }