feat: support function name retrieve of a lambda function

Author: chaoming0625

extensions/brainpylib/operator/_cuda.py

@@ -1,88 +1,82 @@
-import sys
 import ctypes
 import ctypes.util
-from numba import cuda as ncuda
-from cffi import FFI
+import sys
 
+from cffi import FFI
+from numba import cuda
 from numba import types
 
+
 class Dl_info(ctypes.Structure):
-    """
-    Structure of the Dl_info returned by the CFFI of dl.dladdr
-    """
+  """
+  Structure of the Dl_info returned by the CFFI of dl.dladdr
+  """
 
-    _fields_ = (
-        ("dli_fname", ctypes.c_char_p),
-        ("dli_fbase", ctypes.c_void_p),
-        ("dli_sname", ctypes.c_char_p),
-        ("dli_saddr", ctypes.c_void_p),
-    )
+  _fields_ = (
+    ("dli_fname", ctypes.c_char_p),
+    ("dli_fbase", ctypes.c_void_p),
+    ("dli_sname", ctypes.c_char_p),
+    ("dli_saddr", ctypes.c_void_p),
+  )
 
 
 # Find the dynamic linker library path. Only works on unix-like os
 libdl_path = ctypes.util.find_library("dl")
 if libdl_path:
-    # Load the dynamic linker dynamically
-    libdl = ctypes.CDLL(libdl_path)
-
-    # Define dladdr to get the pointer to a symbol in a shared
-    # library already loaded.
-    # https://man7.org/linux/man-pages/man3/dladdr.3.html
-    libdl.dladdr.argtypes = (ctypes.c_void_p, ctypes.POINTER(Dl_info))
-    # restype is None as it returns by reference
+  # Load the dynamic linker dynamically
+  libdl = ctypes.CDLL(libdl_path)
+
+  # Define dladdr to get the pointer to a symbol in a shared
+  # library already loaded.
+  # https://man7.org/linux/man-pages/man3/dladdr.3.html
+  libdl.dladdr.argtypes = (ctypes.c_void_p, ctypes.POINTER(Dl_info))
+  # restype is None as it returns by reference
 else:
-    # On Windows it is nontrivial to have libdl, so we disable everything about
-    # it and use other ways to find paths of libraries
-    libdl = None
+  # On Windows it is nontrivial to have libdl, so we disable everything about
+  # it and use other ways to find paths of libraries
+  libdl = None
 
 
 def find_path_of_symbol_in_library(symbol):
-    if libdl is None:
-        raise ValueError("libdl not found.")
-
-    info = Dl_info()
+  if libdl is None:
+    raise ValueError("libdl not found.")
 
-    result = libdl.dladdr(symbol, ctypes.byref(info))
-
-    if result and info.dli_fname:
-        return info.dli_fname.decode(sys.getfilesystemencoding())
-    else:
-        raise ValueError("Cannot determine path of Library.")
+  info = Dl_info()
+  result = libdl.dladdr(symbol, ctypes.byref(info))
+  if result and info.dli_fname:
+    return info.dli_fname.decode(sys.getfilesystemencoding())
+  else:
+    raise ValueError("Cannot determine path of Library.")
 
 
 try:
-    _libcuda = ncuda.driver.find_driver()
-
-    if sys.platform == "win32":
-        libcuda_path = ctypes.util.find_library(_libcuda._name)
-    else:
-        libcuda_path = find_path_of_symbol_in_library(_libcuda.cuMemcpy)
-
-    numba_cffi_loaded = True
+  _libcuda = cuda.driver.find_driver()
+  if sys.platform == "win32":
+    libcuda_path = ctypes.util.find_library(_libcuda._name)
+  else:
+    libcuda_path = find_path_of_symbol_in_library(_libcuda.cuMemcpy)
+  numba_cffi_loaded = True
 except Exception:
-    numba_cffi_loaded = False
+  numba_cffi_loaded = False
 
-if numba_cffi_loaded:
 
-    # functions needed
-    ffi = FFI()
-    ffi.cdef("int cuMemcpy(void* dst, void* src, unsigned int len, int type);")
-    ffi.cdef(
-        "int cuMemcpyAsync(void* dst, void* src, unsigned int len, int type, void* stream);"
-    )
-    ffi.cdef("int cuStreamSynchronize(void* stream);")
-
-    ffi.cdef("int cudaMallocHost(void** ptr, size_t size);")
-    ffi.cdef("int cudaFreeHost(void* ptr);")
-
-    # load libraray
-    # could  ncuda.driver.find_library()
-    libcuda = ffi.dlopen(libcuda_path)
-    cuMemcpy = libcuda.cuMemcpy
-    cuMemcpyAsync = libcuda.cuMemcpyAsync
-    cuStreamSynchronize = libcuda.cuStreamSynchronize
-
-    memcpyHostToHost = types.int32(0)
-    memcpyHostToDevice = types.int32(1)
-    memcpyDeviceToHost = types.int32(2)
-    memcpyDeviceToDevice = types.int32(3)
+if numba_cffi_loaded:
+  # functions needed
+  ffi = FFI()
+  ffi.cdef("int cuMemcpy(void* dst, void* src, unsigned int len, int type);")
+  ffi.cdef("int cuMemcpyAsync(void* dst, void* src, unsigned int len, int type, void* stream);")
+  ffi.cdef("int cuStreamSynchronize(void* stream);")
+  ffi.cdef("int cudaMallocHost(void** ptr, size_t size);")
+  ffi.cdef("int cudaFreeHost(void* ptr);")
+
+  # load libraray
+  # could  ncuda.driver.find_library()
+  libcuda = ffi.dlopen(libcuda_path)
+  cuMemcpy = libcuda.cuMemcpy
+  cuMemcpyAsync = libcuda.cuMemcpyAsync
+  cuStreamSynchronize = libcuda.cuStreamSynchronize
+
+  memcpyHostToHost = types.int32(0)
+  memcpyHostToDevice = types.int32(1)
+  memcpyDeviceToHost = types.int32(2)
+  memcpyDeviceToDevice = types.int32(3)

extensions/brainpylib/operator/cpu_op.py

@@ -3,6 +3,7 @@
 import collections.abc
 import ctypes
 from functools import partial
+from types import LambdaType
 from typing import Callable, Union, Sequence
 
 import jax.numpy as jnp
@@ -15,6 +16,7 @@
 from numba import types
 from numba.core.dispatcher import Dispatcher
 
+_lambda_no = 0
 ctypes.pythonapi.PyCapsule_New.argtypes = [
   ctypes.c_void_p,  # void* pointer
   ctypes.c_char_p,  # const char *name
@@ -52,9 +54,9 @@ def xla_cpu_custom_call_target(output_ptrs, input_ptrs):
     {args_in}
   )
   func_to_call(args_out, args_in)
-    '''.format(args_in=",\n\t".join(args_in),
-               args_out=",\n\t".join(args_out))
-  print(code_string)
+    '''.format(args_in=",\n    ".join(args_in),
+               args_out=",\n    ".join(args_out))
+  # print(code_string)
   exec(compile(code_string.strip(), '', 'exec'), code_scope)
 
   new_f = code_scope['xla_cpu_custom_call_target']
@@ -101,7 +103,9 @@ def register_cpu_op(
     out_shapes: Union[Callable, ShapedArray, Sequence[ShapedArray]]
 ):
   # primitive
-  prim = core.Primitive(func.__name__)
+  prim = core.Primitive(f'_lambda_func{_lambda_no}'
+                        if (isinstance(func, LambdaType) and func.__name__ == "<lambda>")
+                        else func.__name__)
   prim.multiple_results = True
 
   # user defined function
@@ -161,6 +165,8 @@ def bind_primitive(*inputs):
   def abs_eval(*ins):
     return ins
 
+  import brainpy as bp
+  bp.math.set_platform('cpu')
 
   def custom_op(outs, ins):
     y, y1 = outs
@@ -170,6 +176,9 @@ def custom_op(outs, ins):
 
 
   z = jnp.ones((1, 2), dtype=jnp.float32)
-  jit_op = register_cpu_op(custom_op, abs_eval)
+  op = register_cpu_op(custom_op, abs_eval)
+
+  from jax import jit
+  jit_op = jit(op)
 
   print(jit_op(z, z))

extensions/brainpylib/operator/gpu_op.py

@@ -141,7 +141,9 @@ def register_gpu_op(
   if not _cuda.numba_cffi_loaded:
     raise RuntimeError("Numba cffi could not be loaded.")
   # primitive
-  prim = core.Primitive(f'_lambda_func{_lambda_no}' if isinstance(func, LambdaType) else func.__name__)
+  prim = core.Primitive(f'_lambda_func{_lambda_no}'
+                        if (isinstance(func, LambdaType) and func.__name__ == "<lambda>")
+                        else func.__name__)
   prim.multiple_results = True
 
   # user defined function