Merge pull request #58 from xsuite/revert-57-revert-56-refactor

Revert "Revert "use get_a_buffer""

examples/ex_array_init_perf.py

@@ -0,0 +1,42 @@
+import xobjects as xo
+import numpy as np
+import time
+
+
+def timeit(code, nn=200):
+    import math
+
+    cc = compile(code, "", "exec")
+    t1 = time.time()
+    for i in range(nn):
+        exec(cc)
+    t2 = time.time()
+    dt = (t2 - t1) / nn
+    oo = int(math.log10(dt) / 3)-1
+    unit = "num kM"[oo + 3]
+    scale = 1000**oo
+    print(f"{nn:4} times {dt/scale:8.3f} {unit}sec  {code}")
+
+
+Arr = xo.Float64[:]
+
+
+class StA(xo.Struct):
+    a = Arr
+
+
+class StS(xo.Struct):
+    s = xo.Float64
+
+
+buf = xo.context_default.new_buffer(2**20)
+al = [1, 2, 3]
+ar = np.array(al)
+
+
+timeit("Arr(al,_buffer=buf)", 1000)
+timeit("Arr(ar,_buffer=buf)", 1000)
+
+timeit("StS(s=1,_buffer=buf)", 1000)
+timeit("StA(a=al,_buffer=buf)", 1000)
+timeit("StA(a=ar,_buffer=buf)", 1000)

setup.py

@@ -7,7 +7,7 @@
 
 setup(
     name="xobjects",
-    version="0.1.23",
+    version="0.1.24",
     description="In-memory serialization and code generator for CPU and GPU",
     long_description="In-memory serialization and code generator for CPU and GPU",
     author="Riccardo De Maria",

xobjects/__init__.py

@@ -28,7 +28,7 @@
 
 from .specialize_source import specialize_source
 
-from .typeutils import context_default, allocate_on_buffer
+from .typeutils import context_default, get_a_buffer
 
 from .hybrid_class import JEncoder, HybridClass, MetaHybridClass, ThisClass
 

xobjects/array.py

@@ -292,7 +292,7 @@ def _inspect_args(cls, *args):
         - offsets
         - value: None if args contains dimensions else args[0]
         """
-        log.debug(f"get size for {cls} from {args}")
+        # log.debug(f"get size for {cls} from {args}")
         info = Info()
         extra = {}
         if cls._size is not None:

xobjects/capi.py

@@ -420,13 +420,13 @@ def gen_method_member(cls, path, conf):
     lst.append(gen_method_offset(path, conf))
 
     # pointed = gen_c_pointed(Arg(Int8, pointer=True), conf)
-    lst.extend(Ref_get_c_offset("offset",conf))
+    lst.extend(Ref_get_c_offset("offset", conf))
     pointed = gen_c_pointed(retarg, conf)
     lst.append(f" return {pointed};")
     # GPU not handled
-    #lst.append("  char *reloff_p = (char *) obj + offset;")
-    #lst.append("  int64_t  reloff= *(int64_t *) reloff_p;")
-    #lst.append("  return (void*) (reloff_p + reloff);")
+    # lst.append("  char *reloff_p = (char *) obj + offset;")
+    # lst.append("  int64_t  reloff= *(int64_t *) reloff_p;")
+    # lst.append("  return (void*) (reloff_p + reloff);")
     lst.append("}")
     return "\n".join(lst), kernel
 
@@ -443,28 +443,29 @@ def gen_method_switch(cls, path, conf, method):
         extra=method.args,
         ret=method.ret,
     )
-    refname=lasttype.__name__
+    refname = lasttype.__name__
 
     decl = gen_c_decl_from_kernel(kernel, conf)
     lst = [decl + "{"]
-    voidp=gen_pointer("void*",conf)
+    voidp = gen_pointer("void*", conf)
     lst.append(f"  {voidp} member = {refname}_member(obj);")
     lst.append(f"  switch ({refname}_typeid(obj)){{")
     for atype in lasttype._reftypes:
-        atname=atype.__name__
-        targs=[f"({atname}) member"]
+        atname = atype.__name__
+        targs = [f"({atname}) member"]
         for arg in kernel.args[1:]:
             targs.append(f"{arg.name}")
-        targs=','.join(targs)
-        lst.append(f"""\
+        targs = ",".join(targs)
+        lst.append(
+            f"""\
         #ifndef {refname.upper()}_SKIP_{atname.upper()}
         case {refname}_{atname}_t:
             return {atname}_{method.c_name}({targs});
             break;
         #endif"""
         )
-    lst.append ("  }")
-    lst.append ("  return 0;")
+    lst.append("  }")
+    lst.append("  return 0;")
     lst.append("}")
     return "\n".join(lst), kernel
 

xobjects/context.py

@@ -60,7 +60,7 @@ def sort_classes(classes):
     lst = list(classes)
     for cls in lst:
         cldeps = []
-        cllist=[]
+        cllist = []
         if hasattr(cls, "_get_inner_types"):
             cllist.extend(cls._get_inner_types())
         if hasattr(cls, "_depends_on"):
@@ -418,11 +418,13 @@ def get_classes(self):
             classes.append(self.ret.atype)
         return classes
 
+
 class Source:
     def __init__(self, source, name=None):
         self.source = source
         self.name = name
 
+
 class Method:
     def __init__(self, args, c_name, ret):
         self.args = args
@@ -478,7 +480,7 @@ def get_test_contexts():
     ctxstr = os.environ.get("XOBJECTS_TEST_CONTEXTS")
     if ctxstr is None:
         yield xo.ContextCpu()
-        #yield xo.ContextCpu(omp_num_threads=2)
+        # yield xo.ContextCpu(omp_num_threads=2)
         if xo.ContextCupy in xo.context.available:
             yield xo.ContextCupy()
         if xo.ContextPyopencl in xo.context.available:

xobjects/context_cpu.py

@@ -78,14 +78,21 @@ def cdef_from_kernel(kernel, pyname=None):
     signature += ");"
     return signature
 
+
 # order of base classes matters as it defines which __setitem__ is used
 class LinkedArrayCpu(BaseLinkedArray, np.ndarray):
-
     @classmethod
     def _build_view(cls, a):
         assert len(a.shape) == 1
-        return cls(shape=a.shape, dtype=a.dtype, buffer=a.data, offset=0,
-                strides=a.strides, order='C')
+        return cls(
+            shape=a.shape,
+            dtype=a.dtype,
+            buffer=a.data,
+            offset=0,
+            strides=a.strides,
+            order="C",
+        )
+
 
 class ContextCpu(XContext):
     """
@@ -111,7 +118,9 @@ def __init__(self, omp_num_threads=0):
         self.omp_num_threads = omp_num_threads
 
         if self.omp_num_threads > 1:
-            raise NotImplementedError("OpenMP parallelization not yet supported!")
+            raise NotImplementedError(
+                "OpenMP parallelization not yet supported!"
+            )
 
     def _make_buffer(self, capacity):
         return BufferNumpy(capacity=capacity, context=self)
@@ -128,7 +137,7 @@ def add_kernels(
         extra_cdef=None,
         extra_classes=[],
         extra_headers=[],
-        compile=True
+        compile=True,
     ):
 
         """
@@ -255,7 +264,7 @@ def add_kernels(
                 xtr_compile_args.append("-fopenmp")
                 xtr_link_args.append("-fopenmp")
 
-            if os.name == 'nt': #windows
+            if os.name == "nt":  # windows
                 # TODO: to be handled properly
                 xtr_compile_args = []
                 xtr_link_args = []
@@ -297,11 +306,15 @@ def add_kernels(
 
             finally:
                 # Clean temp files
-                files_to_remove = [so_fname, tempfname + ".c", tempfname + ".o"]
+                files_to_remove = [
+                    so_fname,
+                    tempfname + ".c",
+                    tempfname + ".o",
+                ]
 
                 for ff in files_to_remove:
                     if os.path.exists(ff):
-                        if os.name == 'nt' and ff.endswith('.pyd'):
+                        if os.name == "nt" and ff.endswith(".pyd"):
                             # pyd files are protected on windows
                             continue
                         os.remove(ff)
@@ -318,7 +331,9 @@ def add_kernels(
         for pyname, kernel in kernels.items():
             self.kernels[pyname].source = source
             self.kernels[pyname].specialized_source = specialized_source
-            self.kernels[pyname].description.pyname = pyname # TODO: find better implementation?
+            self.kernels[
+                pyname
+            ].description.pyname = pyname  # TODO: find better implementation?
 
     def nparray_to_context_array(self, arr):
         """
@@ -504,7 +519,9 @@ def update_from_nplike(self, offset, dest_dtype, value):
         if dest_dtype != value.dtype:
             value = value.astype(dtype=dest_dtype)  # make a copy
         src = value.view("int8")
-        self.buffer[offset : offset + src.nbytes] = value.flatten().view("int8")
+        self.buffer[offset : offset + src.nbytes] = value.flatten().view(
+            "int8"
+        )
 
     def to_bytearray(self, offset, nbytes):
         """copy in byte array: used in update_from_xbuffer"""
@@ -539,9 +556,9 @@ def to_function_arg(self, arg, value):
                         self.ffi_interface.from_buffer(slice_first_elem.data),
                     )
                 elif hasattr(value, "_shape"):  # xobject array
-                    assert isinstance(value._buffer.context, ContextCpu), (
-                        f"Incompatible context for argument `{arg.name}`."
-                    )
+                    assert isinstance(
+                        value._buffer.context, ContextCpu
+                    ), f"Incompatible context for argument `{arg.name}`."
                     return self.ffi_interface.cast(
                         value._c_type + "*",
                         self.ffi_interface.from_buffer(
@@ -558,9 +575,9 @@ def to_function_arg(self, arg, value):
             if hasattr(arg.atype, "_dtype"):  # it is numerical scalar
                 return arg.atype(value)  # try to return a numpy scalar
             elif hasattr(arg.atype, "_size"):  # it is a compound xobject
-                assert isinstance(value._buffer.context, ContextCpu), (
-                    f"Incompatible context for argument `{arg.name}`."
-                )
+                assert isinstance(
+                    value._buffer.context, ContextCpu
+                ), f"Incompatible context for argument `{arg.name}`."
                 buf = np.frombuffer(value._buffer.buffer, dtype="int8")
                 ptr = buf.ctypes.data + value._offset
                 return self.ffi_interface.cast(arg.atype._c_type, ptr)