1# distutils: language = c++
2
3# Test case for defining a XLA custom call target in Cython, and registering
4# it via the xla_client SWIG API.
5
6from cpython.pycapsule cimport PyCapsule_New
7
8cdef void test_subtract_f32(void* out_ptr, void** data_ptr) nogil:
9  cdef float a = (<float*>(data_ptr[0]))[0]
10  cdef float b = (<float*>(data_ptr[1]))[0]
11  cdef float* out = <float*>(out_ptr)
12  out[0] = a - b
13
14
15cpu_custom_call_targets = {}
16
17cdef register_custom_call_target(fn_name, void* fn):
18  cdef const char* name = "xla._CPU_CUSTOM_CALL_TARGET"
19  cpu_custom_call_targets[fn_name] = PyCapsule_New(fn, name, NULL)
20
21register_custom_call_target(b"test_subtract_f32", <void*>(test_subtract_f32))
22