xobjects 0.6.3__tar.gz → 0.6.5__tar.gz

{xobjects-0.6.3/xobjects.egg-info → xobjects-0.6.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xobjects
-Version: 0.6.3
+Version: 0.6.5
 Summary: In-memory serialization and code generator for CPU and GPU
 Home-page: https://xsuite.readthedocs.io/
 Download-URL: https://pypi.python.org/pypi/xobjects

{xobjects-0.6.3 → xobjects-0.6.5}/tests/test_kernel.py

@@ -4,7 +4,9 @@
 # ########################################### #
 
 import cffi
+import pytest
 import sysconfig
+from pathlib import Path
 
 import numpy as np
 

xobjects-0.6.5/xobjects/_version.py

@@ -0,0 +1 @@
+__version__ = "0.6.5"

{xobjects-0.6.3 → xobjects-0.6.5}/xobjects/context_cpu.py

@@ -19,6 +19,60 @@ import scipy as sp
 
 _forbid_compile = False
 _suppress_warnings = False
+allow_no_prebuilt_kernel = False
+
+
+def _class_allows_no_prebuilt_kernel(cls):
+    return (
+        getattr(cls, 'allow_no_prebuilt_kernel', False)
+        or getattr(getattr(cls, '_DressingClass', None),
+                   'allow_no_prebuilt_kernel', False)
+        or getattr(getattr(cls, '_XoStruct', None),
+                   'allow_no_prebuilt_kernel', False)
+    )
+
+
+def allow_no_prebuilt_kernel_enabled(context=None, classes=()):
+    if classes is None:
+        classes = ()
+    elif isinstance(classes, type):
+        classes = (classes,)
+
+    if os.environ.get('XSUITE_ALLOW_NO_PREBUILT_KERNELS') is not None:
+        return True
+    if allow_no_prebuilt_kernel:
+        return True
+    if any(_class_allows_no_prebuilt_kernel(cls) for cls in classes):
+        return True
+    return getattr(context, 'allow_no_prebuilt_kernel', False)
+
+
+def _is_serial_cpu_context(context):
+    if context is None or not hasattr(context, 'openmp_enabled'):
+        return False
+    return context.openmp_enabled is False
+
+
+def require_prebuilt_kernel(context=None, classes=()):
+    return (
+        not allow_no_prebuilt_kernel_enabled(context, classes=classes)
+        and _is_serial_cpu_context(context)
+    )
+
+
+def no_prebuilt_kernel_jit_message():
+    return (
+        'To allow just-in-time compilation instead, as in older Xsuite '
+        'versions, set the environment variable '
+        '`XSUITE_ALLOW_NO_PREBUILT_KERNELS`, set '
+        '`xobjects.context_cpu.allow_no_prebuilt_kernel = True`, or set '
+        '`context.allow_no_prebuilt_kernel = True`. Classes that require '
+        'just-in-time compilation can also define '
+        '`allow_no_prebuilt_kernel = True` as a class attribute. Using '
+        'just-in-time compilation instead of prebuilt kernels may require '
+        'lengthy compilation whenever a different kernel is needed.'
+    )
+
 
 from .context import (
     Kernel,
@@ -150,8 +204,7 @@ class ContextCpu(XContext):
         """
         super().__init__()
         self.omp_num_threads = omp_num_threads
-        if omp_num_threads == 0:
-            self.allow_prebuilt_kernels = True
+        self.allow_prebuilt_kernels = True
 
     def __str__(self):
         if not self.openmp_enabled:
@@ -415,8 +468,18 @@ class ContextCpu(XContext):
                 log.debug(f"cffi def {pyname} {signature}")
 
         if self.openmp_enabled:
-            ffi_interface.cdef("void omp_set_num_threads(int);")
-            ffi_interface.cdef("int omp_get_max_threads();")
+            # The wrapper is needed to ensure that the omp functions are linked
+            ffi_interface.cdef("void xo_omp_set_num_threads(int);")
+            ffi_interface.cdef("int xo_omp_get_max_threads();")
+            specialized_source += """
+                void xo_omp_set_num_threads(int num_threads) {
+                    omp_set_num_threads(num_threads);
+                }
+
+                int xo_omp_get_max_threads(void) {
+                    return omp_get_max_threads();
+                }
+                """
 
         # Compile
         xtr_compile_args = ["-std=c99", "-DXO_CONTEXT_CPU"]
@@ -528,8 +591,8 @@ class ContextCpu(XContext):
         spec.loader.exec_module(module)
 
         if self.openmp_enabled:
-            self.omp_set_num_threads = module.lib.omp_set_num_threads
-            self.omp_get_max_threads = module.lib.omp_get_max_threads
+            self.omp_set_num_threads = module.lib.xo_omp_set_num_threads
+            self.omp_get_max_threads = module.lib.xo_omp_get_max_threads
 
         return module
 
@@ -768,10 +831,8 @@ class BufferNumpy(XBuffer):
         value = nplike_to_numpy(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"
-        )
+        src = value.flatten().view("int8")
+        self.buffer[offset : offset + src.nbytes] = src
 
     def to_bytearray(self, offset, nbytes):
         """copy in byte array: used in update_from_xbuffer"""

{xobjects-0.6.3 → xobjects-0.6.5}/xobjects/struct.py

@@ -518,25 +518,37 @@ class Struct(metaclass=MetaStruct):
         extra_classes=(),
         extra_compile_args=(),
     ):
-
         if context.allow_prebuilt_kernels:
             _print_state = Print.suppress
             Print.suppress = True
             try:
-                from xsuite import (
-                    get_suitable_kernel,
-                    XSK_PREBUILT_KERNELS_LOCATION,
-                )
-
-                kernel_info = get_suitable_kernel(
-                    config={},
-                    tracker_element_classes=[],
-                    classes=list(extra_classes) + [cls],
-                )
-            except ImportError:
-                kernel_info = None
-
-            Print.suppress = _print_state
+                try:
+                    from xsuite import (
+                        get_suitable_kernel,
+                        XSK_PREBUILT_KERNELS_LOCATION,
+                    )
+
+                    kernel_info = get_suitable_kernel(
+                        config={},
+                        tracker_element_classes=[],
+                        classes=list(extra_classes) + [cls],
+                        context=context,
+                    )
+                except ImportError as err:
+                    from . import context_cpu
+
+                    if not context_cpu.require_prebuilt_kernel(
+                            context, classes=list(extra_classes) + [cls]):
+                        kernel_info = None
+                    else:
+                        raise ImportError(
+                            'Xsuite is required to load prebuilt kernels but could '
+                            'not be imported. Please install it with '
+                            f'`pip install xsuite`. '
+                            f'{context_cpu.no_prebuilt_kernel_jit_message()}'
+                        ) from err
+            finally:
+                Print.suppress = _print_state
             if kernel_info:
                 kernels = context.kernels_from_file(
                     module_name=kernel_info["module_name"],

{xobjects-0.6.3 → xobjects-0.6.5}/xobjects/test_helpers.py

@@ -108,6 +108,37 @@ def fix_random_seed(seed: int):
     return decorator
 
 
+def allow_no_prebuilt_kernels(test_function=None, *, skip_when_forbid_compile=True):
+    """Allow JIT compilation for tests that intentionally compile kernels.
+
+    By default, the wrapped test is skipped when compilation is forbidden by
+    ``XOBJECTS_FORBID_COMPILE``. Use ``skip_when_forbid_compile=False`` when
+    the test has more specific ``skip_if_forbid_compile()`` guards inside the
+    test.
+    """
+
+    def decorator(test_function):
+        @wraps(test_function)
+        def wrapper(*args, **kwargs):
+            if skip_when_forbid_compile:
+                skip_if_forbid_compile()
+            old_value = os.environ.get("XSUITE_ALLOW_NO_PREBUILT_KERNELS")
+            os.environ["XSUITE_ALLOW_NO_PREBUILT_KERNELS"] = "1"
+            try:
+                return test_function(*args, **kwargs)
+            finally:
+                if old_value is None:
+                    del os.environ["XSUITE_ALLOW_NO_PREBUILT_KERNELS"]
+                else:
+                    os.environ["XSUITE_ALLOW_NO_PREBUILT_KERNELS"] = old_value
+
+        return wrapper
+
+    if test_function is None:
+        return decorator
+    return decorator(test_function)
+
+
 def skip_if_forbid_compile():
     if os.environ.get("XOBJECTS_FORBID_COMPILE"):
         pytest.skip(

{xobjects-0.6.3 → xobjects-0.6.5/xobjects.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xobjects
-Version: 0.6.3
+Version: 0.6.5
 Summary: In-memory serialization and code generator for CPU and GPU
 Home-page: https://xsuite.readthedocs.io/
 Download-URL: https://pypi.python.org/pypi/xobjects

xobjects-0.6.3/xobjects/_version.py

@@ -1 +0,0 @@
-__version__ = "0.6.3"