diff --git a/tests/test_capi.py b/tests/test_capi.py index 72f16c1..1cc59d6 100644 --- a/tests/test_capi.py +++ b/tests/test_capi.py @@ -21,6 +21,15 @@ class Struct1(xo.Struct): field2 = xo.Float64 +class RawValue(xo.RawUnion): + scalar = xo.Float64 + bits = xo.UInt64 + + +class StructRawUnion(xo.Struct): + value = RawValue + + class Struct2(xo.Struct): field1 = xo.Int32 field2 = xo.Float64[:] @@ -243,6 +252,36 @@ def test_struct1(): assert p2[0] == s1.field2 +def test_raw_union_member_getters(): + kernels = StructRawUnion._gen_kernels() + ctx = xo.ContextCpu() + ctx.add_kernels(kernels=kernels) + + source = StructRawUnion._gen_c_api().source + assert "StructRawUnion_get_value_scalar" in source + assert "StructRawUnion_get_value_bits" in source + assert "StructRawUnion_get_value(" not in source + + value = 3.5 + value_bits = np.array([value], dtype=np.float64).view(np.uint64)[0] + s1 = StructRawUnion(value=value) + s2 = StructRawUnion(value=("bits", value_bits)) + raw_value = RawValue(value) + + assert isinstance(s1.value, RawValue) + assert s1.value.scalar == value + assert s1.value.bits == value_bits + assert s2.value.scalar == value + assert s2.value.bits == value_bits + assert raw_value.scalar == value + assert raw_value.bits == value_bits + + assert ctx.kernels.StructRawUnion_get_value_scalar(obj=s1) == value + assert ctx.kernels.StructRawUnion_get_value_bits(obj=s1) == value_bits + assert ctx.kernels.StructRawUnion_get_value_scalar(obj=s2) == value + assert ctx.kernels.StructRawUnion_get_value_bits(obj=s2) == value_bits + + def test_struct2(): s2 = Struct2(field1=2, field2=5) diff --git a/tests/test_general.py b/tests/test_general.py index 07ea1c4..67810c5 100644 --- a/tests/test_general.py +++ b/tests/test_general.py @@ -12,44 +12,44 @@ def test_print_mode_default(capsys): printer = Print() - printer('visible') + printer("visible") - assert capsys.readouterr().out == 'visible\n' + assert capsys.readouterr().out == "visible\n" def test_print_mode_suppressed(capsys): printer = Print() - with xo.settings.override(print_mode='suppress'): - printer('hidden') + with xo.settings.override(print_mode="suppress"): + printer("hidden") - assert capsys.readouterr().out == '' + assert capsys.readouterr().out == "" def test_settings_control_xsuite_printer(capsys): - with xo.settings.override(print_mode='suppress'): - xo._print('hidden') + with xo.settings.override(print_mode="suppress"): + xo._print("hidden") - assert capsys.readouterr().out == '' + assert capsys.readouterr().out == "" def test_python_setting_overrides_environment_default(capsys): - with xo.settings.override(print_mode='print'): - xo._print('visible') + with xo.settings.override(print_mode="print"): + xo._print("visible") - assert capsys.readouterr().out == 'visible\n' + assert capsys.readouterr().out == "visible\n" def test_invalid_print_mode_setting(): - with pytest.raises(ValueError, match='XSUITE_PRINT_MODE.*print.*suppress'): - xo.settings.print_mode = 'invalid' + with pytest.raises(ValueError, match="XSUITE_PRINT_MODE.*print.*suppress"): + xo.settings.print_mode = "invalid" def test_settings_override_restores_value_after_error(): original = xo.settings.print_mode with pytest.raises(RuntimeError): - with xo.settings.override(print_mode='suppress'): - assert xo.settings.print_mode == 'suppress' + with xo.settings.override(print_mode="suppress"): + assert xo.settings.print_mode == "suppress" raise RuntimeError assert xo.settings.print_mode == original @@ -57,84 +57,88 @@ def test_settings_override_restores_value_after_error(): def test_settings_are_discoverable(): expected_settings = { - 'print_mode', - 'progress_indicator', - 'allow_kernel_compilation', - 'force_kernel_compilation', - 'show_kernel_diagnostics', - 'cffi_forbid_compile', - 'cffi_keep_build_files', - 'cuda_backend', - 'cuda_fast_compile', - 'cuda_compiler', + "print_mode", + "progress_indicator", + "allow_kernel_compilation", + "force_kernel_compilation", + "show_kernel_diagnostics", + "cffi_forbid_compile", + "cffi_keep_build_files", + "cuda_backend", + "cuda_fast_compile", + "cuda_compiler", } assert expected_settings <= set(dir(xo.settings)) for name in expected_settings: - assert f'{name}=' in repr(xo.settings) + assert f"{name}=" in repr(xo.settings) assert name in type(xo.settings).__doc__ def test_print_mode_environment_variable_is_startup_default(): environment = os.environ.copy() - environment['XSUITE_PRINT_MODE'] = 'suppress' + environment["XSUITE_PRINT_MODE"] = "suppress" code = ( - 'import xobjects as xo; ' + "import xobjects as xo; " 'assert xo.settings.print_mode == "suppress"; ' - 'xo._print("hidden")') + 'xo._print("hidden")' + ) completed = subprocess.run( - [sys.executable, '-c', code], + [sys.executable, "-c", code], env=environment, capture_output=True, text=True, check=True, ) - assert completed.stdout == '' + assert completed.stdout == "" def test_print_mode_environment_variable_can_be_overridden_in_python(): environment = os.environ.copy() - environment['XSUITE_PRINT_MODE'] = 'suppress' + environment["XSUITE_PRINT_MODE"] = "suppress" code = ( - 'import xobjects as xo; ' + "import xobjects as xo; " 'xo.settings.print_mode = "print"; ' 'assert xo.settings.print_mode == "print"; ' - 'xo._print("visible")') + 'xo._print("visible")' + ) completed = subprocess.run( - [sys.executable, '-c', code], + [sys.executable, "-c", code], env=environment, capture_output=True, text=True, check=True, ) - assert completed.stdout == 'visible\n' + assert completed.stdout == "visible\n" def test_environment_change_after_import_does_not_override_python_setting( - capsys, monkeypatch): - with xo.settings.override(print_mode='print'): - monkeypatch.setenv('XSUITE_PRINT_MODE', 'suppress') - xo._print('visible') + capsys, monkeypatch +): + with xo.settings.override(print_mode="print"): + monkeypatch.setenv("XSUITE_PRINT_MODE", "suppress") + xo._print("visible") - assert capsys.readouterr().out == 'visible\n' + assert capsys.readouterr().out == "visible\n" -@pytest.mark.parametrize('value', ['1', 'true', 'YES', 'on']) +@pytest.mark.parametrize("value", ["1", "true", "YES", "on"]) def test_boolean_environment_true_values(value): environment = os.environ.copy() - environment['XSUITE_ALLOW_KERNEL_COMPILATION'] = value + environment["XSUITE_ALLOW_KERNEL_COMPILATION"] = value code = ( - 'import xobjects as xo; ' - 'assert xo.settings.allow_kernel_compilation is True; ' - 'xo.settings.allow_kernel_compilation = False; ' - 'assert xo.settings.allow_kernel_compilation is False') + "import xobjects as xo; " + "assert xo.settings.allow_kernel_compilation is True; " + "xo.settings.allow_kernel_compilation = False; " + "assert xo.settings.allow_kernel_compilation is False" + ) subprocess.run( - [sys.executable, '-c', code], + [sys.executable, "-c", code], env=environment, capture_output=True, text=True, @@ -142,16 +146,17 @@ def test_boolean_environment_true_values(value): ) -@pytest.mark.parametrize('value', ['0', 'false', 'NO', 'off']) +@pytest.mark.parametrize("value", ["0", "false", "NO", "off"]) def test_boolean_environment_false_values(value): environment = os.environ.copy() - environment['XSUITE_FORCE_KERNEL_COMPILATION'] = value + environment["XSUITE_FORCE_KERNEL_COMPILATION"] = value code = ( - 'import xobjects as xo; ' - 'assert xo.settings.force_kernel_compilation is False') + "import xobjects as xo; " + "assert xo.settings.force_kernel_compilation is False" + ) subprocess.run( - [sys.executable, '-c', code], + [sys.executable, "-c", code], env=environment, capture_output=True, text=True, @@ -161,46 +166,49 @@ def test_boolean_environment_false_values(value): def test_invalid_boolean_environment_value(): environment = os.environ.copy() - environment['XSUITE_CFFI_FORBID_COMPILE'] = 'sometimes' + environment["XSUITE_CFFI_FORBID_COMPILE"] = "sometimes" completed = subprocess.run( - [sys.executable, '-c', 'import xobjects'], + [sys.executable, "-c", "import xobjects"], env=environment, capture_output=True, text=True, ) assert completed.returncode != 0 - assert 'Invalid boolean value' in completed.stderr - assert 'xobjects.settings.cffi_forbid_compile' in completed.stderr - assert 'XSUITE_CFFI_FORBID_COMPILE' in completed.stderr + assert "Invalid boolean value" in completed.stderr + assert "xobjects.settings.cffi_forbid_compile" in completed.stderr + assert "XSUITE_CFFI_FORBID_COMPILE" in completed.stderr def test_runtime_settings_environment_defaults(): environment = os.environ.copy() - environment.update({ - 'XSUITE_PROGRESS_INDICATOR': 'text', - 'XSUITE_FORCE_KERNEL_COMPILATION': 'yes', - 'XSUITE_SHOW_KERNEL_DIAGNOSTICS': 'on', - 'XSUITE_CFFI_FORBID_COMPILE': 'true', - 'XSUITE_CFFI_KEEP_BUILD_FILES': '1', - 'XSUITE_CUDA_BACKEND': 'clang', - 'XSUITE_CUDA_FAST_COMPILE': 'false', - 'XSUITE_CUDA_COMPILER': '/path/to/clang++', - }) + environment.update( + { + "XSUITE_PROGRESS_INDICATOR": "text", + "XSUITE_FORCE_KERNEL_COMPILATION": "yes", + "XSUITE_SHOW_KERNEL_DIAGNOSTICS": "on", + "XSUITE_CFFI_FORBID_COMPILE": "true", + "XSUITE_CFFI_KEEP_BUILD_FILES": "1", + "XSUITE_CUDA_BACKEND": "clang", + "XSUITE_CUDA_FAST_COMPILE": "false", + "XSUITE_CUDA_COMPILER": "/path/to/clang++", + } + ) code = ( - 'import xobjects as xo; ' + "import xobjects as xo; " 'assert xo.settings.progress_indicator == "text"; ' - 'assert xo.settings.force_kernel_compilation is True; ' - 'assert xo.settings.show_kernel_diagnostics is True; ' - 'assert xo.settings.cffi_forbid_compile is True; ' - 'assert xo.settings.cffi_keep_build_files is True; ' + "assert xo.settings.force_kernel_compilation is True; " + "assert xo.settings.show_kernel_diagnostics is True; " + "assert xo.settings.cffi_forbid_compile is True; " + "assert xo.settings.cffi_keep_build_files is True; " 'assert xo.settings.cuda_backend == "clang"; ' - 'assert xo.settings.cuda_fast_compile is False; ' - 'assert xo.settings.cuda_compiler == "/path/to/clang++"') + "assert xo.settings.cuda_fast_compile is False; " + 'assert xo.settings.cuda_compiler == "/path/to/clang++"' + ) subprocess.run( - [sys.executable, '-c', code], + [sys.executable, "-c", code], env=environment, capture_output=True, text=True, @@ -209,7 +217,7 @@ def test_runtime_settings_environment_defaults(): @pytest.mark.parametrize( - 'allow, force, compilation_allowed', + "allow, force, compilation_allowed", [ (False, False, False), (True, False, True), @@ -222,12 +230,14 @@ def test_kernel_compilation_settings(allow, force, compilation_allowed): allow_kernel_compilation=allow, force_kernel_compilation=force, ): - assert xo.context_cpu.kernel_compilation_allowed( - xo.ContextCpu()) is compilation_allowed + assert ( + xo.context_cpu.kernel_compilation_allowed(xo.ContextCpu()) + is compilation_allowed + ) def test_user_context_environment_variable(monkeypatch): - monkeypatch.setenv('XOBJECTS_USER_CONTEXT', 'ContextCpu:auto') + monkeypatch.setenv("XOBJECTS_USER_CONTEXT", "ContextCpu:auto") context = xo.get_user_context() assert context.openmp_enabled @@ -235,8 +245,8 @@ def test_user_context_environment_variable(monkeypatch): def test_test_contexts_environment_variable(monkeypatch): monkeypatch.setenv( - 'XOBJECTS_TEST_CONTEXTS', - 'ContextCpu;ContextCpu:auto', + "XOBJECTS_TEST_CONTEXTS", + "ContextCpu;ContextCpu:auto", ) contexts = list(xo.context.get_test_contexts()) @@ -252,8 +262,8 @@ def test_cffi_forbid_compile_setting(): xo.ContextCpu().build_kernels({}) message = str(err.value) - assert 'xobjects.settings.cffi_forbid_compile' in message - assert 'XSUITE_CFFI_FORBID_COMPILE' in message + assert "xobjects.settings.cffi_forbid_compile" in message + assert "XSUITE_CFFI_FORBID_COMPILE" in message def test_allow_kernel_compilation_decorator_restores_state(): diff --git a/tests/test_raw_union.py b/tests/test_raw_union.py new file mode 100644 index 0000000..93a52ed --- /dev/null +++ b/tests/test_raw_union.py @@ -0,0 +1,78 @@ +# copyright ################################# # +# This file is part of the Xobjects Package. # +# Copyright (c) CERN, 2026. # +# ########################################### # + +import numpy as np +import pytest + +import xobjects as xo + + +class RawValue(xo.RawUnion): + scalar = xo.Float64 + bits = xo.UInt64 + + +class Holder(xo.Struct): + value = RawValue + + +def _float_bits(value): + return np.array([value], dtype=np.float64).view(np.uint64)[0] + + +def test_raw_union_declaration_consumes_member_attributes(): + assert RawValue._members == {"scalar": xo.Float64, "bits": xo.UInt64} + assert RawValue._size == 8 + + assert not hasattr(RawValue, "scalar") + assert not hasattr(RawValue, "bits") + + +def test_raw_union_standalone_instance_reads_explicit_members(): + value = 3.5 + raw_value = RawValue(value) + + assert isinstance(raw_value, RawValue) + assert raw_value.scalar == value + assert raw_value.bits == _float_bits(value) + assert "scalar" in dir(raw_value) + assert "bits" in dir(raw_value) + + with pytest.raises(AttributeError, match="missing"): + raw_value.missing + + +def test_raw_union_struct_field_reads_explicit_members(): + value = 3.5 + holder = Holder(value=value) + + assert isinstance(holder.value, RawValue) + assert holder.value.scalar == value + assert holder.value.bits == _float_bits(value) + + with pytest.raises(TypeError): + float(holder.value) + + +def test_raw_union_constructor_accepts_explicit_member(): + value = 3.5 + value_bits = _float_bits(value) + holder = Holder(value=("bits", value_bits)) + + assert holder.value.scalar == value + assert holder.value.bits == value_bits + + +def test_raw_union_copy_from_same_type(): + source = RawValue(3.5) + holder = Holder(value=0.0) + + holder.value = source + assert holder.value.scalar == source.scalar + assert holder.value.bits == source.bits + + copied = RawValue(source) + assert copied.scalar == source.scalar + assert copied.bits == source.bits diff --git a/xobjects/__init__.py b/xobjects/__init__.py index 6b50a6d..9124c77 100644 --- a/xobjects/__init__.py +++ b/xobjects/__init__.py @@ -16,6 +16,7 @@ UInt8, ) from .array import Array +from .raw_union import RawUnion from .string import String from .struct import Struct, Field, ThisClass from .ref import Ref, UnionRef diff --git a/xobjects/array.py b/xobjects/array.py index 8c723d3..5245a4e 100644 --- a/xobjects/array.py +++ b/xobjects/array.py @@ -414,7 +414,7 @@ def _inspect_args(cls, *args): return info @classmethod - def _from_buffer(cls, buffer, offset=0): + def _from_buffer(cls, buffer, offset=0, container=None): self = object.__new__(cls) self._buffer = buffer self._offset = offset @@ -451,7 +451,7 @@ def _from_buffer(cls, buffer, offset=0): return self @classmethod - def _to_buffer(cls, buffer, offset, value, info=None): + def _to_buffer(cls, buffer, offset, value, info=None, container=None): if info is None: info = cls._inspect_args(value) value = info.value # can be None if value contained shape info diff --git a/xobjects/capi.py b/xobjects/capi.py index af07f9b..1b4ad1e 100644 --- a/xobjects/capi.py +++ b/xobjects/capi.py @@ -9,23 +9,46 @@ from .struct import is_field, is_struct from .array import is_index, is_array from .ref import is_unionref, is_ref +from .raw_union import is_raw_union from .string import is_string -def is_compound(atype): - """Types that are referenced in C with an opaque pointer""" - return is_struct(atype) or is_array(atype) or is_unionref(atype) +def is_compound(xotype): + """Return True for types passed in C as opaque object pointers. + Parameters: + xotype: xobjects type to check. -def is_type(atype): - return is_compound(atype) or is_scalar(atype) or is_string(atype) + Returns: + True if the type is a struct, array, or unionref. + """ + return is_struct(xotype) or is_array(xotype) or is_unionref(xotype) + + +def is_type(xotype): + """Return True for any xobjects type with a C API representation. + + Parameters: + xotype: xobjects type to check. + + Returns: + True if the type can appear in generated C API methods. + """ + return is_compound(xotype) or is_scalar(xotype) or is_string(xotype) def get_inner_type(part): - """type contained in a field, array or ref, else None""" - if is_field(part): # is a field + """Return the contained type for a field, array, or ref. + + Parameters: + part: Path item, such as a field, array, ref, or unionref. + + Returns: + Inner xobjects type, or None for a unionref. + """ + if is_field(part): return part.ftype - elif is_array(part): # is an array + elif is_array(part): return part._itemtype elif is_ref(part): return part._reftype @@ -36,42 +59,78 @@ def get_inner_type(part): def gen_pointer(ctype, conf): - gpumem = conf.get("gpumem", "") - return f"{gpumem}{ctype}" + """Add the configured GPU memory qualifier to a pointer type. + + Parameters: + ctype: C pointer type text. + conf: C API generation options. + + Returns: + Pointer type text with the configured memory qualifier. + """ + gpu_memory_qualifier = conf.get("gpumem", "") + return f"{gpu_memory_qualifier}{ctype}" def gen_c_type_from_arg(arg: Arg, conf): - gpumem = conf.get("gpumem", "") + """Return the C type used for an argument or return value. + + Parameters: + arg: Kernel argument, return argument, or None for void. + conf: C API generation options. + + Returns: + C type text without an argument name. + """ + gpu_memory_qualifier = conf.get("gpumem", "") if arg is None: - cdec = "void" + c_type = "void" else: - cdec = arg.atype._c_type - if cdec.endswith("*"): # handle literal char* - cdec = f"{gpumem}{cdec}" + c_type = arg.atype._c_type + if c_type.endswith("*"): # handle literal char* + c_type = f"{gpu_memory_qualifier}{c_type}" if arg.pointer: - cdec = f"{gpumem}{cdec}*" + c_type = f"{gpu_memory_qualifier}{c_type}*" if arg.const: - cdec = "const " + cdec - return cdec + c_type = "const " + c_type + return c_type def gen_c_arg_from_arg(arg: Arg, conf): - gpumem = conf.get("gpumem", "") - cpurestrict = conf.get("cpurestrict", "") + """Return one C function-argument declaration. + + Parameters: + arg: Kernel argument to declare. + conf: C API generation options. + + Returns: + C argument declaration including the argument name. + """ + gpu_memory_qualifier = conf.get("gpumem", "") + cpu_restrict_qualifier = conf.get("cpurestrict", "") if arg is None: - cdec = "void" + c_type = "void" else: - cdec = arg.atype._c_type + c_type = arg.atype._c_type if arg.pointer: - cdec = f"{gpumem}{cdec}*{cpurestrict}" + c_type = f"{gpu_memory_qualifier}{c_type}*{cpu_restrict_qualifier}" elif is_compound(arg.atype): - cdec = f"{cdec}{cpurestrict}" + c_type = f"{c_type}{cpu_restrict_qualifier}" if arg.const: - cdec = "const " + cdec - return f"{cdec} {arg.name}" + c_type = "const " + c_type + return f"{c_type} {arg.name}" def gen_c_size_from_arg(arg: Arg, conf): + """Return the C-side storage size, in bytes, for an argument. + + Parameters: + arg: Kernel argument, return argument, or None for void. + conf: C API generation options. + + Returns: + Size in bytes, or None for void. + """ if arg is None: return None else: @@ -82,19 +141,38 @@ def gen_c_size_from_arg(arg: Arg, conf): def gen_c_decl_from_kernel(kernel: Kernel, conf): - args = ", ".join([gen_c_arg_from_arg(arg, conf) for arg in kernel.args]) + """Return a C function declaration for a kernel. + + Parameters: + kernel: Kernel description to declare. + conf: C API generation options. + + Returns: + C function declaration without a trailing semicolon. + """ + c_args = ", ".join([gen_c_arg_from_arg(arg, conf) for arg in kernel.args]) if kernel.ret is None: - ret = "void" + return_type = "void" else: - ret = gen_c_type_from_arg(kernel.ret, conf) - gpufun = conf.get("gpufun") - if gpufun is None: - return f"{ret} {kernel.c_name}({args})" + return_type = gen_c_type_from_arg(kernel.ret, conf) + gpu_function_qualifier = conf.get("gpufun") + if gpu_function_qualifier is None: + return f"{return_type} {kernel.c_name}({c_args})" else: - return f"{gpufun} {ret} {kernel.c_name}({args})" + return ( + f"{gpu_function_qualifier} {return_type} {kernel.c_name}({c_args})" + ) def get_layers(parts): + """Count array-index layers in a path. + + Parameters: + parts: Path items to inspect. + + Returns: + Number of path items that have a shape. + """ layers = 0 for part in parts: if hasattr(part, "_shape"): @@ -103,174 +181,282 @@ def get_layers(parts): def int_from_obj(offset, conf): - """Generate code to read the integer at location obj + offset (bytes).""" + """Generate code to read the integer at obj + offset. + + Parameters: + offset: C expression for the byte offset from obj. + conf: C API generation options. + + Returns: + C expression that loads the configured integer type. + """ int_pointer_type = gen_pointer(conf.get("inttype", "int64_t") + "*", conf) char_pointer_type = gen_pointer(conf.get("chartype", "char") + "*", conf) return f"*({int_pointer_type})(({char_pointer_type}) obj+{offset})" def Field_get_c_offset(self, conf): + """Return C code or an integer offset for a field access. + + Parameters: + self: Field object from an access path. + conf: C API generation options. + + Returns: + Static byte offset, or C lines that follow a dynamic reference. + """ if self.is_reference: - doffset = f"offset+{self.offset}" # starts of data - refoffset = int_from_obj(doffset, conf) - return [f" offset+={refoffset};"] + data_offset_expr = f"offset+{self.offset}" + reference_offset_expr = int_from_obj(data_offset_expr, conf) + return [f" offset+={reference_offset_expr};"] else: return self.offset -def Ref_get_c_offset(self, conf): - refoffset = int_from_obj("offset", conf) - return [f" offset+={refoffset};"] +def Ref_get_c_offset(_ref, conf): + """Return C code that follows a relative reference offset. + + Parameters: + _ref: Ref object from an access path. Not used. + conf: C API generation options. + + Returns: + C lines that add the referenced object's relative offset. + """ + reference_offset_expr = int_from_obj("offset", conf) + return [f" offset+={reference_offset_expr};"] + +def Index_get_c_offset(index_part, conf, index_start): + """Return C code that applies an array index to the byte offset. -def Index_get_c_offset(part, conf, icount): - cls = part.cls - inttype = conf.get("inttype", "int64_t") + Parameters: + index_part: Index object from an access path. + conf: C API generation options. + index_start: First C index argument used by this path item. - out = [] - if hasattr(cls, "_strides"): # static shape or 1d dynamic shape - strides = cls._strides + Returns: + C lines that update offset for this index. + """ + array_cls = index_part.cls + int_type = conf.get("inttype", "int64_t") + + code_lines = [] + if hasattr(array_cls, "_strides"): # static shape or 1d dynamic shape + strides = array_cls._strides else: - nd = len(cls._shape) + ndim = len(array_cls._shape) strides = [] - for ii in range(nd): - stride_offset = 8 + (len(cls._dshape_idx) * 8) + (ii * 8) - sname = f"{cls.__name__}_s{ii}" - svalue = int_from_obj(f"offset+{stride_offset}", conf) - out.append(f" {inttype} {sname}={svalue};") - strides.append(sname) - - soffset = "+".join([f"i{ii+icount}*{ss}" for ii, ss in enumerate(strides)]) - if cls._data_offset > 0: - soffset = f"{cls._data_offset}+{soffset}" - if cls._is_static_type: - out.append(f" offset+={soffset};") + for axis in range(ndim): + stride_offset = 8 + (len(array_cls._dshape_idx) * 8) + (axis * 8) + stride_name = f"{array_cls.__name__}_s{axis}" + stride_value = int_from_obj(f"offset+{stride_offset}", conf) + code_lines.append(f" {int_type} {stride_name}={stride_value};") + strides.append(stride_name) + + slot_offset = "+".join( + f"i{axis + index_start}*{stride}" + for axis, stride in enumerate(strides) + ) + if array_cls._data_offset > 0: + slot_offset = f"{array_cls._data_offset}+{slot_offset}" + if array_cls._is_static_type: + code_lines.append(f" offset+={slot_offset};") else: - lookup_field_offset = f"offset+{soffset}" - out.append(f" offset+={int_from_obj(lookup_field_offset, conf)};") - return out + lookup_field_offset = f"offset+{slot_offset}" + code_lines.append( + f" offset+={int_from_obj(lookup_field_offset, conf)};" + ) + return code_lines def gen_method_offset(path, conf): - """return code to obtain offset of the target in bytes""" - inttype = conf.get("inttype", "int64_t") - lst = [f" {inttype} offset=0;"] - offset = 0 - icount = 0 + """Return C code that computes the target byte offset. + + Parameters: + path: Access path from root object to target. + conf: C API generation options. + + Returns: + C statements that initialize and update the offset variable. + """ + int_type = conf.get("inttype", "int64_t") + code_lines = [f" {int_type} offset=0;"] + pending_static_offset = 0 + next_index_arg = 0 for part in path: if is_index(part): - soffset = Index_get_c_offset(part, conf, icount) - icount += len(part.cls._shape) + offset_step = Index_get_c_offset(part, conf, next_index_arg) + next_index_arg += len(part.cls._shape) elif is_field(part): - soffset = Field_get_c_offset(part, conf) + offset_step = Field_get_c_offset(part, conf) elif is_ref(part): - soffset = Ref_get_c_offset(part, conf) + offset_step = Ref_get_c_offset(part, conf) else: - soffset = None + offset_step = None - if type(soffset) is int: - offset += soffset - elif type(soffset) is list: - if offset > 0: - lst.append(f" offset+={offset};") # dump current offset - lst.extend(soffset) # update reference offset - offset = 0 - if offset > 0: - lst.append(f" offset+={offset};") - return "\n".join(lst) + if type(offset_step) is int: + pending_static_offset += offset_step + elif type(offset_step) is list: + if pending_static_offset > 0: + code_lines.append(f" offset+={pending_static_offset};") + code_lines.extend(offset_step) + pending_static_offset = 0 + if pending_static_offset > 0: + code_lines.append(f" offset+={pending_static_offset};") + return "\n".join(code_lines) def gen_fun_kernel(cls, path, action, const, extra, ret, add_nindex=False): - typename = cls._c_type - fields = [] - indices = 0 + """Build the Kernel object for one generated C API method. + + Parameters: + cls: Root xobjects class for the generated method. + path: Access path from root object to target. + action: Method action name, such as get or shape. + const: Whether the root object argument is const. + extra: Extra Kernel arguments after obj and indices. + ret: Kernel return Arg, or None for void. + add_nindex: Append the number of index arguments to action. + + Returns: + Kernel description for the generated C method. + """ + type_name = cls._c_type + field_names = [] + index_count = 0 + for part in path: - if is_field(part): # is field - fields.append(part.name) - elif is_index(part): # is array - indices += len(part.cls._shape) - if add_nindex and indices > 0: - action += str(indices) - fun_name = [typename, action] - if len(fields) > 0: - fun_name.append("_".join(fields)) - fun_name = "_".join(fun_name) + if is_field(part): + field_names.append(part.name) + elif is_index(part): + index_count += len(part.cls._shape) + + if add_nindex and index_count > 0: + action += str(index_count) + + function_name_parts = [type_name, action] + if len(field_names) > 0: + function_name_parts.append("_".join(field_names)) + + function_name = "_".join(function_name_parts) args = [Arg(cls, pointer=False, const=const, name="obj")] - for ii in range(indices): - args.append(Arg(Int64, name=f"i{ii}")) + for index_arg in range(index_count): + args.append(Arg(Int64, name=f"i{index_arg}")) args.extend(extra) - return Kernel(args, c_name=fun_name, ret=ret) + return Kernel(args, c_name=function_name, ret=ret) def gen_c_pointed(target: Arg, conf): - size = gen_c_size_from_arg(target, conf) - ret = gen_c_type_from_arg(target, conf) + """Return C code that reads or points to the target at obj + offset. + + Parameters: + target: Argument describing the target type. + conf: C API generation options. + + Returns: + C expression usable as an lvalue or return value. + """ + target_size = gen_c_size_from_arg(target, conf) + target_c_type = gen_c_type_from_arg(target, conf) if target.pointer or is_compound(target.atype) or is_string(target.atype): - chartype = gen_pointer(conf.get("chartype", "char") + "*", conf) - return f"({ret})(({chartype}) obj+offset)" + char_pointer_type = gen_pointer( + conf.get("chartype", "char") + "*", conf + ) + return f"({target_c_type})(({char_pointer_type}) obj+offset)" - rettype = gen_pointer(ret + "*", conf) - if size == 1: - return f"*(({rettype}) obj+offset)" + target_pointer_type = gen_pointer(target_c_type + "*", conf) + if target_size == 1: + return f"*(({target_pointer_type}) obj+offset)" - chartype = gen_pointer(conf.get("chartype", "char") + "*", conf) - return f"*({rettype})(({chartype}) obj+offset)" + char_pointer_type = gen_pointer(conf.get("chartype", "char") + "*", conf) + return f"*({target_pointer_type})(({char_pointer_type}) obj+offset)" def gen_method_get(cls, path, conf): - lasttype = path[-1] - retarg = Arg(lasttype) + """Generate a C getter for a scalar path target. + + Parameters: + cls: Root xobjects class for the generated method. + path: Access path from root object to scalar target. + conf: C API generation options. + + Returns: + Tuple of C source text and its Kernel description. + """ + target_type = path[-1] + return_arg = Arg(target_type) kernel = gen_fun_kernel( cls, path, const=True, action="get", extra=[], - ret=retarg, + ret=return_arg, ) - decl = gen_c_decl_from_kernel(kernel, conf) + declaration = gen_c_decl_from_kernel(kernel, conf) - lst = [decl + "{"] - lst.append(gen_method_offset(path, conf)) - pointed = gen_c_pointed(retarg, conf) - lst.append(f" return {pointed};") - lst.append("}") - return "\n".join(lst), kernel + code_lines = [declaration + "{"] + code_lines.append(gen_method_offset(path, conf)) + target_expr = gen_c_pointed(return_arg, conf) + code_lines.append(f" return {target_expr};") + code_lines.append("}") + return "\n".join(code_lines), kernel def gen_method_set(cls, path, conf): - lasttype = path[-1] - valarg = Arg(lasttype, name="value") + """Generate a C setter for a scalar path target. + + Parameters: + cls: Root xobjects class for the generated method. + path: Access path from root object to scalar target. + conf: C API generation options. + + Returns: + Tuple of C source text and its Kernel description. + """ + target_type = path[-1] + value_arg = Arg(target_type, name="value") kernel = gen_fun_kernel( cls, path, const=False, action="set", - extra=[valarg], + extra=[value_arg], ret=None, ) - decl = gen_c_decl_from_kernel(kernel, conf) + declaration = gen_c_decl_from_kernel(kernel, conf) - lst = [decl + "{"] - lst.append(gen_method_offset(path, conf)) - pointed = gen_c_pointed(valarg, conf) - lst.append(f" {pointed}=value;") - lst.append("}") - return "\n".join(lst), kernel + code_lines = [declaration + "{"] + code_lines.append(gen_method_offset(path, conf)) + target_expr = gen_c_pointed(value_arg, conf) + code_lines.append(f" {target_expr}=value;") + code_lines.append("}") + return "\n".join(code_lines), kernel def gen_method_getp(cls, path, conf): - lasttype = path[-1] - if lasttype is None: - retarg = Arg(Void, pointer="True") + """Generate a C getter that returns a pointer to the path target. + + Parameters: + cls: Root xobjects class for the generated method. + path: Access path from root object to target. + conf: C API generation options. + + Returns: + Tuple of C source text and its Kernel description. + """ + target_type = path[-1] + if target_type is None: + return_arg = Arg(Void, pointer="True") else: - retarg = Arg(lasttype) - # if is_scalar(lasttype) or is_string(lasttype): - if is_scalar(lasttype): - retarg.pointer = True + return_arg = Arg(target_type) + # if is_scalar(target_type) or is_string(target_type): + if is_scalar(target_type): + return_arg.pointer = True kernel = gen_fun_kernel( cls, @@ -278,23 +464,63 @@ def gen_method_getp(cls, path, conf): const=False, action="getp", extra=[], - ret=retarg, + ret=return_arg, add_nindex=True, ) - decl = gen_c_decl_from_kernel(kernel, conf) + declaration = gen_c_decl_from_kernel(kernel, conf) + + code_lines = [declaration + "{"] + code_lines.append(gen_method_offset(path, conf)) + target_expr = gen_c_pointed(return_arg, conf) + code_lines.append(f" return {target_expr};") + code_lines.append("}") + return "\n".join(code_lines), kernel - lst = [decl + "{"] - lst.append(gen_method_offset(path, conf)) - pointed = gen_c_pointed(retarg, conf) - lst.append(f" return {pointed};") - lst.append("}") - return "\n".join(lst), kernel + +def gen_method_get_raw_union_member(cls, path, member_name, member_type, conf): + """Generate a C getter for one named RawUnion member view.""" + type_name = cls._c_type + field_names = [] + index_count = 0 + for part in path: + if is_field(part): + field_names.append(part.name) + elif is_index(part): + index_count += len(part.cls._shape) + + return_arg = Arg(member_type) + function_name_parts = [type_name, "get"] + if field_names: + function_name_parts.append("_".join(field_names)) + function_name_parts.append(member_name) + args = [Arg(cls, pointer=False, const=True, name="obj")] + for index_arg in range(index_count): + args.append(Arg(Int64, name=f"i{index_arg}")) + kernel = Kernel(args, c_name="_".join(function_name_parts), ret=return_arg) + declaration = gen_c_decl_from_kernel(kernel, conf) + + code_lines = [declaration + "{"] + code_lines.append(gen_method_offset(path, conf)) + target_expr = gen_c_pointed(return_arg, conf) + code_lines.append(f" return {target_expr};") + code_lines.append("}") + return "\n".join(code_lines), kernel def gen_method_len(cls, path, conf): - lasttype = path[-1] + """Generate a C method returning the total array length. - retarg = Arg(Int64) + Parameters: + cls: Root xobjects class for the generated method. + path: Access path from root object to array target. + conf: C API generation options. + + Returns: + Tuple of C source text and its Kernel description. + """ + array_type = path[-1] + + return_arg = Arg(Int64) kernel = gen_fun_kernel( cls, @@ -302,45 +528,55 @@ def gen_method_len(cls, path, conf): const=False, action="len", extra=[], - ret=retarg, + ret=return_arg, add_nindex=True, ) - decl = gen_c_decl_from_kernel(kernel, conf) + declaration = gen_c_decl_from_kernel(kernel, conf) - lst = [decl + "{"] + code_lines = [declaration + "{"] - if lasttype._is_static_shape: - ll = lasttype._get_n_items() - lst.append(f" return {ll};"), kernel + if array_type._is_static_shape: + item_count = array_type._get_n_items() + code_lines.append(f" return {item_count};") else: - lst.append(gen_method_offset(path, conf)) - arrarg = Arg(Int64, pointer=True) - pointed = gen_c_pointed(arrarg, conf) - typearr = gen_pointer("int64_t*", conf) - lst.append(f" {typearr} arr = {pointed};") + code_lines.append(gen_method_offset(path, conf)) + header_arg = Arg(Int64, pointer=True) + header_expr = gen_c_pointed(header_arg, conf) + header_pointer_type = gen_pointer("int64_t*", conf) + code_lines.append(f" {header_pointer_type} arr = {header_expr};") dim_len_idx = 1 terms = [] - for dim_len in lasttype._shape: + for dim_len in array_type._shape: if dim_len: terms.append(str(dim_len)) else: terms.append(f"arr[{dim_len_idx}]") dim_len_idx += 1 length_expr = "*".join(terms) - lst.append(f" return {length_expr};") - lst.append("}") - return "\n".join(lst), kernel + code_lines.append(f" return {length_expr};") + code_lines.append("}") + return "\n".join(code_lines), kernel def gen_method_size(cls, path, conf): - innertype = get_inner_type(path[-1]) - if innertype is None: # cannot determine size + """Generate a C method returning item size in bytes. + + Parameters: + cls: Root xobjects class for the generated method. + path: Access path from root object to container target. + conf: C API generation options. + + Returns: + Tuple of C source text and Kernel, or (None, None) if size is unknown. + """ + inner_type = get_inner_type(path[-1]) + if inner_type is None: # cannot determine size return (None, None) - retarg = Arg(Int64) + return_arg = Arg(Int64) action = "size" layers = get_layers(path) - if layers > 0 and not is_scalar(innertype): + if layers > 0 and not is_scalar(inner_type): action += str(layers) kernel = gen_fun_kernel( @@ -349,23 +585,33 @@ def gen_method_size(cls, path, conf): const=False, action=action, extra=[], - ret=retarg, + ret=return_arg, ) - decl = gen_c_decl_from_kernel(kernel, conf) + declaration = gen_c_decl_from_kernel(kernel, conf) - lst = [decl + "{"] + code_lines = [declaration + "{"] - if innertype._size is None: - lst.append(gen_method_offset(path, conf)) - pointed = gen_c_pointed(retarg, conf) - lst.append(f" return {pointed};") + if inner_type._size is None: + code_lines.append(gen_method_offset(path, conf)) + target_expr = gen_c_pointed(return_arg, conf) + code_lines.append(f" return {target_expr};") else: - lst.append(f" return {innertype._size};") - lst.append("}") - return "\n".join(lst), kernel + code_lines.append(f" return {inner_type._size};") + code_lines.append("}") + return "\n".join(code_lines), kernel def gen_method_shape(cls, path, conf): + """Generate a C method that writes array shape into out_shape. + + Parameters: + cls: Root xobjects class for the generated method. + path: Access path from root object to array target. + conf: C API generation options. + + Returns: + Tuple of C source text and its Kernel description. + """ array_type = path[-1] out_shape_arg = Arg(Int64, name="out_shape", pointer=True) @@ -379,20 +625,20 @@ def gen_method_shape(cls, path, conf): ret=None, add_nindex=True, ) - decl = gen_c_decl_from_kernel(kernel, conf) + declaration = gen_c_decl_from_kernel(kernel, conf) - lst = [decl + "{"] + code_lines = [declaration + "{"] - nd = len(array_type._shape) + ndim = len(array_type._shape) if array_type._is_static_shape: terms = [str(dim) for dim in array_type._shape] else: - lst.append(gen_method_offset(path, conf)) - arrarg = Arg(Int64, pointer=True) - pointed = gen_c_pointed(arrarg, conf) - typearr = gen_pointer("int64_t*", conf) - lst.append(f" {typearr} arr = {pointed};") + code_lines.append(gen_method_offset(path, conf)) + header_arg = Arg(Int64, pointer=True) + header_expr = gen_c_pointed(header_arg, conf) + header_pointer_type = gen_pointer("int64_t*", conf) + code_lines.append(f" {header_pointer_type} arr = {header_expr};") dim_len_idx = 1 terms = [] for dim_len in array_type._shape: @@ -402,19 +648,28 @@ def gen_method_shape(cls, path, conf): terms.append(f"arr[{dim_len_idx}]") dim_len_idx += 1 - for dim_idx in range(nd): - lst.append(f" out_shape[{dim_idx}] = {terms[dim_idx]};") + for dim_idx in range(ndim): + code_lines.append(f" out_shape[{dim_idx}] = {terms[dim_idx]};") - lst.append("}") + code_lines.append("}") - return "\n".join(lst), kernel + return "\n".join(code_lines), kernel def gen_method_nd(cls, path, conf): - "return length of shape" + """Generate a C method returning the number of array dimensions. + + Parameters: + cls: Root xobjects class for the generated method. + path: Access path from root object to array target. + conf: C API generation options. + + Returns: + Tuple of C source text and its Kernel description. + """ array_type = path[-1] - retarg = Arg(UInt32) + return_arg = Arg(UInt32) kernel = gen_fun_kernel( cls, @@ -422,33 +677,60 @@ def gen_method_nd(cls, path, conf): const=False, action="nd", extra=[], - ret=retarg, + ret=return_arg, add_nindex=True, ) - decl = gen_c_decl_from_kernel(kernel, conf) + declaration = gen_c_decl_from_kernel(kernel, conf) - lst = [decl + "{"] + code_lines = [declaration + "{"] - nd = len(array_type._shape) - lst.append(f" return {nd};") - lst.append("}") + ndim = len(array_type._shape) + code_lines.append(f" return {ndim};") + code_lines.append("}") - return "\n".join(lst), kernel + return "\n".join(code_lines), kernel def gen_method_strides(cls, path, conf): - "return strides" + """Placeholder for a future C method returning array strides. + + Parameters: + cls: Root xobjects class for the generated method. + path: Access path from root object to array target. + conf: C API generation options. + + Returns: + Always (None, None); this method is not implemented. + """ return None, None def gen_method_getpos(cls, path, conf): - "return slot position from index and strides" + """Placeholder for a future C method returning a slot position. + + Parameters: + cls: Root xobjects class for the generated method. + path: Access path from root object to array target. + conf: C API generation options. + + Returns: + Always (None, None); this method is not implemented. + """ return None, None def gen_method_typeid(cls, path, conf): - "return typeid of a unionref" - retarg = Arg(Int64) + """Generate a C method returning a unionref type id. + + Parameters: + cls: Root xobjects class for the generated method. + path: Access path from root object to unionref target. + conf: C API generation options. + + Returns: + Tuple of C source text and its Kernel description. + """ + return_arg = Arg(Int64) kernel = gen_fun_kernel( cls, @@ -456,23 +738,32 @@ def gen_method_typeid(cls, path, conf): const=True, action="typeid", extra=[], - ret=retarg, + ret=return_arg, ) - decl = gen_c_decl_from_kernel(kernel, conf) + declaration = gen_c_decl_from_kernel(kernel, conf) - lst = [decl + "{"] + code_lines = [declaration + "{"] - lst.append(gen_method_offset(path, conf)) - pointed = gen_c_pointed(retarg, conf) - lst.append(" offset+=8;") - lst.append(f" return {pointed};") - lst.append("}") - return "\n".join(lst), kernel + code_lines.append(gen_method_offset(path, conf)) + target_expr = gen_c_pointed(return_arg, conf) + code_lines.append(" offset+=8;") + code_lines.append(f" return {target_expr};") + code_lines.append("}") + return "\n".join(code_lines), kernel def gen_method_member(cls, path, conf): - "return typeid of a unionref" - retarg = Arg(Void, pointer=True) + """Generate a C method returning the active unionref member. + + Parameters: + cls: Root xobjects class for the generated method. + path: Access path from root object to unionref target. + conf: C API generation options. + + Returns: + Tuple of C source text and its Kernel description. + """ + return_arg = Arg(Void, pointer=True) kernel = gen_fun_kernel( cls, @@ -480,29 +771,34 @@ def gen_method_member(cls, path, conf): const=True, action="member", extra=[], - ret=retarg, + ret=return_arg, ) - decl = gen_c_decl_from_kernel(kernel, conf) + declaration = gen_c_decl_from_kernel(kernel, conf) - lst = [decl + "{"] + code_lines = [declaration + "{"] - lst.append(gen_method_offset(path, conf)) + code_lines.append(gen_method_offset(path, conf)) - # pointed = gen_c_pointed(Arg(Int8, pointer=True), 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("}") - return "\n".join(lst), kernel + code_lines.extend(Ref_get_c_offset("offset", conf)) + target_expr = gen_c_pointed(return_arg, conf) + code_lines.append(f" return {target_expr};") + code_lines.append("}") + return "\n".join(code_lines), kernel def gen_method_switch(cls, path, conf, method): - "generate switch methods declared in _methods" - lasttype = path[-1] + """Generate a C dispatcher for a unionref method. + + Parameters: + cls: Root unionref class for the generated method. + path: Access path ending at that unionref type. + conf: C API generation options. + method: Method descriptor declared on the unionref. + + Returns: + Tuple of C source text and its Kernel description. + """ + unionref_type = path[-1] kernel = gen_fun_kernel( cls, @@ -512,104 +808,154 @@ def gen_method_switch(cls, path, conf, method): extra=method.args, ret=method.ret, ) - refname = lasttype.__name__ - - decl = gen_c_decl_from_kernel(kernel, conf) - lst = [decl + "{"] - 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"] + unionref_name = unionref_type.__name__ + + declaration = gen_c_decl_from_kernel(kernel, conf) + code_lines = [declaration + "{"] + void_pointer_type = gen_pointer("void*", conf) + code_lines.append( + f" {void_pointer_type} member = {unionref_name}_member(obj);" + ) + code_lines.append(f" switch ({unionref_name}_typeid(obj)){{") + for member_type in unionref_type._reftypes: + member_type_name = member_type.__name__ + call_args = [f"({member_type_name}) member"] for arg in kernel.args[1:]: - targs.append(f"{arg.name}") - targs = ",".join(targs) - lst.append(f"""\ - #ifndef {refname.upper()}_SKIP_{atname.upper()} - case {refname}_{atname}_t: - return {atname}_{method.c_name}({targs}); + call_args.append(f"{arg.name}") + call_args = ",".join(call_args) + code_lines.append(f"""\ + #ifndef {unionref_name.upper()}_SKIP_{member_type_name.upper()} + case {unionref_name}_{member_type_name}_t: + return {member_type_name}_{method.c_name}({call_args}); break; #endif""") - lst.append(" }") - lst.append(f" return{'' if method.ret is None else ' 0'};") - lst.append("}") - return "\n".join(lst), kernel + code_lines.append(" }") + code_lines.append(f" return{'' if method.ret is None else ' 0'};") + code_lines.append("}") + return "\n".join(code_lines), kernel def gen_typedef(cls, conf): - gpumem = conf.get("gpumem", "") - typename = cls._c_type - return f"typedef {gpumem} struct {typename}_s * {typename};" + """Generate the opaque C typedef for a compound type. + Parameters: + cls: xobjects class to typedef. + conf: C API generation options. -def gen_enum(cls, conf): - typename = cls.__name__ - lst = ",".join(f"{typename}_{tt._c_type}_t" for tt in cls._reftypes) - return f"enum {typename}_e{{{lst}}};" + Returns: + C typedef statement. + """ + gpu_memory_qualifier = conf.get("gpumem", "") + type_name = cls._c_type + return ( + f"typedef {gpu_memory_qualifier} struct {type_name}_s * {type_name};" + ) -def methods_from_path(cls, path, conf): - """ - size: all - get, set: lasttype is scalar - getp: all real types +def gen_enum(cls, conf): + """Generate the C enum listing unionref member types. + + Parameters: + cls: Unionref class to describe. + conf: C API generation options. Present for API symmetry. + Returns: + C enum declaration statement. """ - out = [] - lasttype = path[-1] + type_name = cls.__name__ + enum_items = ",".join( + f"{type_name}_{member._c_type}_t" for member in cls._reftypes + ) + return f"enum {type_name}_e{{{enum_items}}};" - if is_scalar(lasttype): - out.append(gen_method_get(cls, path, conf)) - out.append(gen_method_set(cls, path, conf)) - if is_type(lasttype): - out.append(gen_method_getp(cls, path, conf)) +def methods_from_path(cls, path, conf): + """Return generated C methods and kernels needed for one access path. - if is_array(lasttype): - out.append(gen_method_len(cls, path, conf)) - out.append(gen_method_shape(cls, path, conf)) - out.append(gen_method_nd(cls, path, conf)) - # out.append(gen_method_strides(cls, path, conf)) - # out.append(gen_method_getpos(cls, path, conf)) + Parameters: + cls: Root xobjects class for generated methods. + path: Access path from root object to target. + conf: C API generation options. - if is_unionref(lasttype): - out.append(gen_method_typeid(cls, path, conf)) - out.append(gen_method_member(cls, path, conf)) - if cls == lasttype: - for method in lasttype._methods: - out.append(gen_method_switch(cls, path, conf, method)) - return out + Returns: + List of (C source, Kernel) tuples. Stubs may return (None, None). + """ + generated_methods = [] + target_type = path[-1] + + if is_scalar(target_type): + generated_methods.append(gen_method_get(cls, path, conf)) + generated_methods.append(gen_method_set(cls, path, conf)) + + if is_raw_union(target_type): + for member_name, member_type in target_type._members.items(): + generated_methods.append( + gen_method_get_raw_union_member( + cls, path, member_name, member_type, conf + ) + ) + + if is_type(target_type): + generated_methods.append(gen_method_getp(cls, path, conf)) + + if is_array(target_type): + generated_methods.append(gen_method_len(cls, path, conf)) + generated_methods.append(gen_method_shape(cls, path, conf)) + generated_methods.append(gen_method_nd(cls, path, conf)) + # Not yet implemented, only method stubs: + # generated_methods.append(gen_method_strides(cls, path, conf)) + # generated_methods.append(gen_method_getpos(cls, path, conf)) + + if is_unionref(target_type): + generated_methods.append(gen_method_typeid(cls, path, conf)) + generated_methods.append(gen_method_member(cls, path, conf)) + if cls == target_type: + for method in target_type._methods: + generated_methods.append( + gen_method_switch(cls, path, conf, method) + ) + return generated_methods def gen_cdef(cls, conf): - out = [] - out.append(gen_typedef(cls, conf)) + """Generate C type declarations needed before method declarations. + + Parameters: + cls: xobjects class to declare. + conf: C API generation options. + + Returns: + C type declarations as source text. + """ + declarations = [] + declarations.append(gen_typedef(cls, conf)) if is_unionref(cls): - out.append(gen_enum(cls, conf)) - return "\n".join(out) + declarations.append(gen_enum(cls, conf)) + return "\n".join(declarations) def gen_code(cls, paths, conf): - """ - Generate source code, used by add_kernels + """Generate C source for the class C API. - Provide all symbols starting from cls.__name__ + Parameters: + cls: xobjects class to generate for. + paths: Access paths that need generated methods. + conf: C API generation options. - The definition must be completed by inner classes if present + Returns: + Full generated C source text for the class API. """ - - typename = cls.__name__ + type_name = cls.__name__ sources = [] - sources.append(f"#ifndef XOBJ_TYPEDEF_{typename}") - sources.append(f"#define XOBJ_TYPEDEF_{typename}") + sources.append(f"#ifndef XOBJ_TYPEDEF_{type_name}") + sources.append(f"#define XOBJ_TYPEDEF_{type_name}") sources.append(gen_cdef(cls, conf)) - out = [] + generated_methods = [] for path in paths: - out.extend(methods_from_path(cls, path, conf)) + generated_methods.extend(methods_from_path(cls, path, conf)) - for source, _ in out: + for source, _ in generated_methods: if source is not None: sources.append(source) @@ -621,16 +967,22 @@ def gen_code(cls, paths, conf): def gen_kernels(cls, paths, conf): - """ - Generate kernel defintions of the C Api + """Generate Kernel definitions for the class C API. + Parameters: + cls: xobjects class to generate for. + paths: Access paths that need generated kernels. + conf: C API generation options. + + Returns: + Dictionary mapping C function names to Kernel descriptions. """ - out = [] + generated_methods = [] for path in paths: - out.extend(methods_from_path(cls, path, conf)) + generated_methods.extend(methods_from_path(cls, path, conf)) kernels = {} - for _, kernel in out: + for _, kernel in generated_methods: if kernel is not None: kernels[kernel.c_name] = kernel @@ -638,17 +990,21 @@ def gen_kernels(cls, paths, conf): def gen_cdefs(cls, paths, conf): - """ + """Generate C declarations for the class C API. - Generate kernel defintions of the C Api + Parameters: + cls: xobjects class to generate for. + paths: Access paths that need generated declarations. + conf: C API generation options. + Returns: + C declarations as source text. """ - kernels = gen_kernels(cls, paths, conf) - out = [gen_cdef(cls, conf)] + declarations = [gen_cdef(cls, conf)] for _, kernel in kernels.items(): - out.append(gen_c_decl_from_kernel(kernel, conf) + ";") + declarations.append(gen_c_decl_from_kernel(kernel, conf) + ";") - return "\n".join(out) + return "\n".join(declarations) diff --git a/xobjects/context.py b/xobjects/context.py index bd3d65f..478f1f3 100644 --- a/xobjects/context.py +++ b/xobjects/context.py @@ -113,7 +113,7 @@ def classes_from_kernels(kernels): classes = set() for _, kernel in kernels.items(): classes.update(kernel.get_classes()) - return classes + return sorted(classes, key=lambda cls: cls.__name__) def _concatenate_sources(sources, apply_to_source=()): @@ -357,7 +357,7 @@ def build_kernels( def get_installed_c_source_and_library_paths( self, - ) -> tuple[set[Path], set[str], set[Path]]: + ) -> tuple[list[Path], list[str], list[Path]]: """Returns a list of C paths registered in dependent packages. In a package that depends on xobjects, you can register C source and @@ -423,7 +423,7 @@ def get_build_info(): libraries = [libraries] libs.update(libraries) - return sources, libs, lib_paths + return sorted(sources), sorted(libs), sorted(lib_paths) @abstractmethod def nparray_to_context_array(self, arr, copy=False): diff --git a/xobjects/context_cpu.py b/xobjects/context_cpu.py index 53bf29f..fd67bc4 100644 --- a/xobjects/context_cpu.py +++ b/xobjects/context_cpu.py @@ -3,6 +3,7 @@ # Copyright (c) CERN, 2021. # # ########################################### # +import ctypes import importlib.util import logging import os @@ -20,6 +21,23 @@ _suppress_warnings = False +_preloaded_shared_libs_cache = {} + + +def _deduplicate_paths(seq): + return list(dict.fromkeys(Path(item).as_posix() for item in seq)) + + +def _preload_shared_libraries(paths): + # Make symbols visible to later kernel imports, and load now, + # not lazily: potential failure will be immediate and clearer. + mode = os.RTLD_GLOBAL | os.RTLD_NOW + + for path in paths: + key = str(Path(path).expanduser().resolve()) + if key not in _preloaded_shared_libs_cache: + _preloaded_shared_libs_cache[key] = ctypes.CDLL(key, mode=mode) + def _class_allows_kernel_compilation(cls): for candidate in ( @@ -43,8 +61,7 @@ def kernel_compilation_allowed(context=None, classes=()): elif isinstance(classes, type): classes = (classes,) - if (settings.allow_kernel_compilation - or settings.force_kernel_compilation): + if settings.allow_kernel_compilation or settings.force_kernel_compilation: return True if any(_class_allows_kernel_compilation(cls) for cls in classes): return True @@ -183,7 +200,6 @@ class ContextCpu(XContext): Returns: ContextCpu: platform object. - """ _cffi_verbose = False @@ -227,6 +243,10 @@ def add_kernels( save_source_as=None, extra_compile_args: Sequence[str] = (), extra_link_args: Sequence[str] = (), + extra_include_dirs: Sequence[str] = (), + extra_libraries: Sequence[str] = (), + extra_library_dirs: Sequence[str] = (), + preload_libraries: Sequence[str] = (), extra_cdef="", extra_classes=(), extra_headers=(), @@ -252,6 +272,11 @@ def add_kernels( code. Default is ```None```. extra_compile_args: Extra arguments to be passed to the compiler. extra_link_args: Extra arguments to be passed to the linker. + extra_include_dirs: Extra include directories passed to cffi. + extra_libraries: Extra libraries passed to cffi. + extra_library_dirs: Extra library directories passed to cffi. + preload_libraries: Shared libraries to load before importing the + compiled kernel module. extra_cdef: Extra C definitions to be passed to cffi. extra_classes: Extra xobjects classes whose API is needed. extra_headers: Extra headers to be added to the source code. @@ -311,6 +336,10 @@ def add_kernels( save_source_as=save_source_as, extra_compile_args=extra_compile_args, extra_link_args=extra_link_args, + extra_include_dirs=extra_include_dirs, + extra_libraries=extra_libraries, + extra_library_dirs=extra_library_dirs, + preload_libraries=preload_libraries, extra_cdef=extra_cdef, extra_classes=extra_classes, extra_headers=extra_headers, @@ -329,11 +358,22 @@ def build_kernels( save_source_as=None, extra_compile_args=(), extra_link_args=(), + extra_include_dirs=(), + extra_libraries=(), + extra_library_dirs=(), + preload_libraries=(), extra_cdef="", extra_classes=(), extra_headers=(), compile=True, # noqa + compiler_language="c", + extra_source_files=None, ) -> Dict[Tuple[str, tuple], "KernelCpu"]: + # extra_source_files: additional source files compiled as separate translation + # units and linked into the module (cffi's set_source(sources=...)), as opposed + # to `sources`, which are concatenated into the single generated translation unit. + # Use this when a source must not share a TU with the rest (e.g. to avoid symbol/ADL clashes). + extra_source_files = extra_source_files or () extra_compile_args += ("-O3", "-Wno-unused-function") extra_link_args += ("-O3",) @@ -384,6 +424,11 @@ def build_kernels( extra_compile_args, extra_link_args, containing_dir=containing_dir, + compiler_language=compiler_language, + extra_source_files=extra_source_files, + extra_include_dirs=extra_include_dirs, + extra_libraries=extra_libraries, + extra_library_dirs=extra_library_dirs, ) try: @@ -392,6 +437,7 @@ def build_kernels( module_name, kernel_descriptions, containing_dir=containing_dir, + preload_libraries=preload_libraries, ) finally: # Whether loaded successfully or not, delete the so @@ -423,16 +469,26 @@ def kernels_from_file( module_name: str, kernel_descriptions: Dict[str, Kernel], containing_dir=".", + preload_libraries=(), ) -> Dict[Tuple[str, tuple], "KernelCpu"]: """ Import a compiled module `module_name` located in `containing_dir` (by default it is the current working directory), and add the kernels from the module, as defined in `kernel_descriptions`, to the context. - Returns the path to the loaded so file. + + Args: + module_name: Compiled extension module name. + kernel_descriptions: Mapping from Python names to Kernel objects. + containing_dir: Directory containing the compiled module. + preload_libraries: Shared libraries to load before module import. + + Returns: + Mapping from Python kernel names to KernelCpu objects. """ module = self._load_kernel_module( name=module_name, containing_dir=containing_dir, + preload_libraries=preload_libraries, ) out_kernels = {} for pyname, kernel_desc in kernel_descriptions.items(): @@ -454,6 +510,11 @@ def compile_kernel( extra_compile_args, extra_link_args, containing_dir=".", + compiler_language="c", + extra_source_files=(), + extra_include_dirs=(), + extra_libraries=(), + extra_library_dirs=(), ) -> Path: ffi_interface = cffi.FFI() ffi_interface.cdef(cdefs) @@ -482,9 +543,11 @@ def compile_kernel( } """ - # Compile - xtr_compile_args = ["-std=c99", "-DXO_CONTEXT_CPU"] - xtr_link_args = ["-std=c99", "-DXO_CONTEXT_CPU"] + # Compile (Default is C99). compiler_language="c++" builds the TU as C++17 + # (need to link header-only C++ dependencies, i.e. the TPSA library). + std_flag = "-std=c++17" if compiler_language == "c++" else "-std=c99" + xtr_compile_args = [std_flag, "-DXO_CONTEXT_CPU"] + xtr_link_args = [std_flag, "-DXO_CONTEXT_CPU"] xtr_compile_args += extra_compile_args xtr_link_args += extra_link_args @@ -498,11 +561,21 @@ def compile_kernel( xtr_link_args.append("-DXO_CONTEXT_CPU_SERIAL") ( - extra_include_paths, - extra_libraries, - extra_library_paths, + installed_include_paths, + installed_libraries, + installed_library_paths, ) = self.get_installed_c_source_and_library_paths() + include_paths = _deduplicate_paths( + [*installed_include_paths, *extra_include_dirs] + ) + libraries = _deduplicate_paths( + [*installed_libraries, *extra_libraries] + ) + library_paths = _deduplicate_paths( + [*installed_library_paths, *extra_library_dirs] + ) + if os.name == "nt": # windows # TODO: to be handled properly xtr_compile_args = [] @@ -515,9 +588,11 @@ def compile_kernel( ffi_interface.set_source( module_name, specialized_source, - include_dirs=[path.as_posix() for path in extra_include_paths], - libraries=list(extra_libraries), - library_dirs=[path.as_posix() for path in extra_library_paths], + include_dirs=[Path(path).as_posix() for path in include_paths], + libraries=list(libraries), + library_dirs=[Path(path).as_posix() for path in library_paths], + source_extension=".cpp" if compiler_language == "c++" else ".c", + sources=list(extra_source_files), extra_compile_args=xtr_compile_args, extra_link_args=xtr_link_args, ) @@ -533,10 +608,11 @@ def compile_kernel( _print("Done compiling ContextCpu kernels.") return Path(output_file) finally: - # Clean temp files + # Clean temp files (the generated source has the language's extension). if not settings.cffi_keep_build_files: + src_ext = ".cpp" if compiler_language == "c++" else ".c" files_to_remove = [ - module_name + ".c", + module_name + src_ext, module_name + ".o", ] @@ -585,11 +661,18 @@ def _load_kernel_module( self, name: str, containing_dir=".", + preload_libraries=(), ): """ Load a kernel from a stored shared object file. + + Args: + name: Compiled extension module name. + containing_dir: Directory containing the shared object. + preload_libraries: Shared libraries to load before module import. """ so_path = _so_for_module_name(name, containing_dir) + _preload_shared_libraries(preload_libraries) # Import the compiled module spec = importlib.util.spec_from_file_location(name, so_path) module = importlib.util.module_from_spec(spec) diff --git a/xobjects/context_cupy.py b/xobjects/context_cupy.py index 4111b6c..b687b6b 100644 --- a/xobjects/context_cupy.py +++ b/xobjects/context_cupy.py @@ -34,11 +34,13 @@ import cupy import cupyx.scipy import cupyx.scipy.interpolate + with warnings.catch_warnings(): warnings.filterwarnings( - 'ignore', - message=r'cupyx\.jit\.rawkernel is experimental\.', - category=FutureWarning) + "ignore", + message=r"cupyx\.jit\.rawkernel is experimental\.", + category=FutureWarning, + ) import cupyx.scipy.signal import cupyx.scipy.special import cupyx.scipy.stats diff --git a/xobjects/general.py b/xobjects/general.py index a8d57cc..6eaf83f 100644 --- a/xobjects/general.py +++ b/xobjects/general.py @@ -16,7 +16,7 @@ class Print: """ def __call__(self, *args, **kwargs): - if settings.print_mode == 'suppress': + if settings.print_mode == "suppress": return print(*args, **kwargs) diff --git a/xobjects/headers/atomicadd.h b/xobjects/headers/atomicadd.h index 6cf3fa9..4d7b9d1 100644 --- a/xobjects/headers/atomicadd.h +++ b/xobjects/headers/atomicadd.h @@ -56,6 +56,19 @@ DEF_ATOMIC_ADD(uint64_t, u64) DEF_ATOMIC_ADD(float , f32) DEF_ATOMIC_ADD(double , f64) +#ifdef __cplusplus +// _Generic is C11 only, but not supported in C++ (with g++). +GPUFUN int8_t atomicAdd(GPUGLMEM int8_t *addr, int8_t val) { return atomicAdd_i8 (addr, val); } +GPUFUN int16_t atomicAdd(GPUGLMEM int16_t *addr, int16_t val) { return atomicAdd_i16(addr, val); } +GPUFUN int32_t atomicAdd(GPUGLMEM int32_t *addr, int32_t val) { return atomicAdd_i32(addr, val); } +GPUFUN int64_t atomicAdd(GPUGLMEM int64_t *addr, int64_t val) { return atomicAdd_i64(addr, val); } +GPUFUN uint8_t atomicAdd(GPUGLMEM uint8_t *addr, uint8_t val) { return atomicAdd_u8 (addr, val); } +GPUFUN uint16_t atomicAdd(GPUGLMEM uint16_t *addr, uint16_t val) { return atomicAdd_u16(addr, val); } +GPUFUN uint32_t atomicAdd(GPUGLMEM uint32_t *addr, uint32_t val) { return atomicAdd_u32(addr, val); } +GPUFUN uint64_t atomicAdd(GPUGLMEM uint64_t *addr, uint64_t val) { return atomicAdd_u64(addr, val); } +GPUFUN float atomicAdd(GPUGLMEM float *addr, float val) { return atomicAdd_f32(addr, val); } +GPUFUN double atomicAdd(GPUGLMEM double *addr, double val) { return atomicAdd_f64(addr, val); } +#else // Using _Generic to select the right function based on type (since C11). // See https://en.cppreference.com/w/c/language/generic.html #define atomicAdd(addr, val) _Generic((addr), \ @@ -70,6 +83,7 @@ DEF_ATOMIC_ADD(double , f64) float*: atomicAdd_f32, \ double*: atomicAdd_f64 \ )(addr, (val)) +#endif // __cplusplus #endif // XO_CONTEXT_CPU diff --git a/xobjects/raw_union.py b/xobjects/raw_union.py new file mode 100644 index 0000000..6739988 --- /dev/null +++ b/xobjects/raw_union.py @@ -0,0 +1,111 @@ +# copyright ################################# # +# This file is part of the Xobjects Package. # +# Copyright (c) CERN, 2026. # +# ########################################### # + +import numpy as np + +from .array import Array +from .scalar import is_scalar +from .typeutils import Info, allocate_on_buffer + + +class MetaRawUnion(type): + def __new__(cls, name, bases, data): + members = {} + for base in bases: + members.update(getattr(base, "_members", {})) + + for member_name, member_type in list(data.items()): + if member_name.startswith("_"): + continue + if is_scalar(member_type): + members[member_name] = member_type + del data[member_name] + + if members: + size = max(member._size for member in members.values()) + dtype = np.dtype(f"V{size}") + data["_members"] = members + data["_dtype"] = dtype + data["_size"] = size + if "_c_type" not in data: + data["_c_type"] = f"{name}_raw_t" + + return type.__new__(cls, name, bases, data) + + def __getitem__(cls, shape): + return Array.mk_arrayclass(cls, shape) + + def __repr__(cls): + return cls.__name__ + + +class RawUnion(metaclass=MetaRawUnion): + _members = {} + + def __init__(self, value=0, _context=None, _buffer=None, _offset=None): + self._buffer, self._offset = allocate_on_buffer( + self.__class__._size, _context, _buffer, _offset + ) + self.__class__._to_buffer(self._buffer, self._offset, value) + + def __getattr__(self, name): + member_type = self.__class__._members.get(name) + if member_type is None: + raise AttributeError( + f"{self.__class__.__name__} has no RawUnion member {name!r}" + ) + return member_type._from_buffer( + self._buffer, self._offset, container=self + ) + + def __dir__(self): + return sorted(set(super().__dir__()) | set(self.__class__._members)) + + def __repr__(self): + members = ", ".join(self.__class__._members) + return f"<{self.__class__.__name__}; members: {members}>" + + @classmethod + def _member_from_value(cls, value): + if isinstance(value, tuple): + member_name, member_value = value + return cls._members[member_name], member_value + return next(iter(cls._members.values())), value + + @classmethod + def _inspect_args(cls, value=0): + return Info(size=cls._size, value=value) + + @classmethod + def _from_buffer(cls, buffer, offset=0, container=None): + self = object.__new__(cls) + self._buffer = buffer + self._offset = offset + return self + + @classmethod + def _to_buffer(cls, buffer, offset, value, info=None, container=None): + if isinstance(value, cls): + buffer.update_from_xbuffer( + offset, value._buffer, value._offset, cls._size + ) + return + member_type, member_value = cls._member_from_value(value) + data = member_type._dtype.type(member_value).tobytes() + if len(data) < cls._size: + data = data + bytes(cls._size - len(data)) + buffer.update_from_buffer(offset, data) + + @classmethod + def _gen_data_paths(cls, base=None): + paths = [] + if base is None: + base = [] + paths.append(base + [cls]) + return paths + + +def is_raw_union(cls): + return isinstance(cls, MetaRawUnion) and cls is not RawUnion diff --git a/xobjects/ref.py b/xobjects/ref.py index ac10a25..6e68315 100644 --- a/xobjects/ref.py +++ b/xobjects/ref.py @@ -37,7 +37,7 @@ def __init__(self, reftype): self._size = 8 - def _from_buffer(self, buffer, offset=0): + def _from_buffer(self, buffer, offset=0, container=None): refoffset = Int64._from_buffer(buffer, offset) if refoffset == NULLVALUE: return None @@ -45,7 +45,7 @@ def _from_buffer(self, buffer, offset=0): refoffset += offset # from relative to absolute offset return self._reftype._from_buffer(buffer, refoffset) - def _to_buffer(self, buffer, offset, value, info=None): + def _to_buffer(self, buffer, offset, value, info=None, container=None): # Get/set content if value is None: refoffset = NULLVALUE # NULL value @@ -153,7 +153,7 @@ def _type_from_typeid(cls, typeid): # If no match found: raise TypeError(f"Invalid id: {typeid}!") - def _from_buffer(cls, buffer, offset=0): + def _from_buffer(cls, buffer, offset=0, container=None): refoffset, typeid = Int64._array_from_buffer(buffer, offset, 2) if refoffset == NULLVALUE: return None @@ -175,7 +175,7 @@ def _inspect_args(cls, *args): info = Info(size=cls._size) return info - def _to_buffer(cls, buffer, offset, value, info=None): + def _to_buffer(cls, buffer, offset, value, info=None, container=None): if isinstance(value, cls): # binary copy buffer.update_from_xbuffer( offset, value._buffer, value._offset, value._size diff --git a/xobjects/scalar.py b/xobjects/scalar.py index ee6f4e0..41492ec 100644 --- a/xobjects/scalar.py +++ b/xobjects/scalar.py @@ -24,11 +24,11 @@ def __init__(self, dtype, cname): self._size = self._dtype.itemsize self._c_type = cname - def _from_buffer(self, buffer, offset=0): + def _from_buffer(self, buffer, offset=0, container=None): data = buffer.to_bytearray(offset, self._size) return np.frombuffer(data, dtype=self._dtype)[0] - def _to_buffer(self, buffer, offset, value, info=None): + def _to_buffer(self, buffer, offset, value, info=None, container=None): data = self._dtype.type(value).tobytes() buffer.update_from_buffer(offset, data) diff --git a/xobjects/settings.py b/xobjects/settings.py index fd72a35..21dac33 100644 --- a/xobjects/settings.py +++ b/xobjects/settings.py @@ -91,8 +91,8 @@ class Settings: """ def __init__(self): - object.__setattr__(self, '_definitions', {}) - object.__setattr__(self, '_values', {}) + object.__setattr__(self, "_definitions", {}) + object.__setattr__(self, "_values", {}) def _register( self, @@ -105,47 +105,55 @@ def _register( value_type=None, ): if name in self._definitions: - raise ValueError(f'Setting {name!r} is already registered.') + raise ValueError(f"Setting {name!r} is already registered.") definition = { - 'environment_variable': environment_variable, - 'choices': choices, - 'value_type': value_type, + "environment_variable": environment_variable, + "choices": choices, + "value_type": value_type, } self._definitions[name] = definition value = default - if (environment_variable is not None - and environment_variable in os.environ): + if ( + environment_variable is not None + and environment_variable in os.environ + ): environment_value = os.environ[environment_variable] try: - value = (environment_parser(environment_value) - if environment_parser else environment_value) + value = ( + environment_parser(environment_value) + if environment_parser + else environment_value + ) except (TypeError, ValueError) as err: raise ValueError( - f'Invalid value for ' - f'{self._setting_description(name, definition)}: ' - f'{err}') from err + f"Invalid value for " + f"{self._setting_description(name, definition)}: " + f"{err}" + ) from err self._set(name, value) def _set(self, name, value): try: definition = self._definitions[name] except KeyError as err: - raise AttributeError(f'Unknown Xsuite setting {name!r}.') from err + raise AttributeError(f"Unknown Xsuite setting {name!r}.") from err - choices = definition['choices'] - value_type = definition['value_type'] + choices = definition["choices"] + value_type = definition["value_type"] setting_description = self._setting_description(name, definition) if value_type is not None and not isinstance(value, value_type): raise TypeError( - f'Invalid value {value!r} for {setting_description}; ' - f'expected {self._type_name(value_type)}.') + f"Invalid value {value!r} for {setting_description}; " + f"expected {self._type_name(value_type)}." + ) if choices is not None and value not in choices: - expected = ', '.join(repr(choice) for choice in choices) + expected = ", ".join(repr(choice) for choice in choices) raise ValueError( - f'Invalid value {value!r} for {setting_description}; ' - f'expected one of {expected}.') + f"Invalid value {value!r} for {setting_description}; " + f"expected one of {expected}." + ) self._values[name] = value @@ -153,7 +161,7 @@ def __getattr__(self, name): try: return self._values[name] except KeyError as err: - raise AttributeError(f'Unknown Xsuite setting {name!r}.') from err + raise AttributeError(f"Unknown Xsuite setting {name!r}.") from err def __setattr__(self, name, value): self._set(name, value) @@ -164,24 +172,26 @@ def override(self, **kwargs): previous = {} for name, value in kwargs.items(): if name not in self._definitions: - raise AttributeError(f'Unknown Xsuite setting {name!r}.') + raise AttributeError(f"Unknown Xsuite setting {name!r}.") previous[name] = getattr(self, name) # Validate every value before changing any setting. for name, value in kwargs.items(): definition = self._definitions[name] setting_description = self._setting_description(name, definition) - value_type = definition['value_type'] + value_type = definition["value_type"] if value_type is not None and not isinstance(value, value_type): raise TypeError( - f'Invalid value {value!r} for {setting_description}; ' - f'expected {self._type_name(value_type)}.') - choices = definition['choices'] + f"Invalid value {value!r} for {setting_description}; " + f"expected {self._type_name(value_type)}." + ) + choices = definition["choices"] if choices is not None and value not in choices: - expected = ', '.join(repr(choice) for choice in choices) + expected = ", ".join(repr(choice) for choice in choices) raise ValueError( - f'Invalid value {value!r} for {setting_description}; ' - f'expected one of {expected}.') + f"Invalid value {value!r} for {setting_description}; " + f"expected one of {expected}." + ) try: for name, value in kwargs.items(): @@ -192,9 +202,10 @@ def override(self, **kwargs): self._set(name, value) def __repr__(self): - values = ', '.join( - f'{name}={getattr(self, name)!r}' for name in self._values) - return f'Settings({values})' + values = ", ".join( + f"{name}={getattr(self, name)!r}" for name in self._values + ) + return f"Settings({values})" def __dir__(self): return sorted(set(super().__dir__()) | set(self._definitions)) @@ -202,30 +213,31 @@ def __dir__(self): @staticmethod def _type_name(value_type): if isinstance(value_type, tuple): - return ' or '.join(tt.__name__ for tt in value_type) + return " or ".join(tt.__name__ for tt in value_type) return value_type.__name__ @staticmethod def _setting_description(name, definition): - description = f'Python setting xobjects.settings.{name}' - environment_variable = definition['environment_variable'] + description = f"Python setting xobjects.settings.{name}" + environment_variable = definition["environment_variable"] if environment_variable is not None: description += ( - f' or equivalently the environment variable ' - f'{environment_variable}' + f" or equivalently the environment variable " + f"{environment_variable}" ) return description def _parse_boolean(value): normalized = value.strip().lower() - if normalized in ('1', 'true', 'yes', 'on'): + if normalized in ("1", "true", "yes", "on"): return True - if normalized in ('0', 'false', 'no', 'off'): + if normalized in ("0", "false", "no", "off"): return False raise ValueError( - f'Invalid boolean value {value!r}; expected one of 1, 0, true, ' - 'false, yes, no, on, or off.') + f"Invalid boolean value {value!r}; expected one of 1, 0, true, " + "false, yes, no, on, or off." + ) def _parse_choice(value): @@ -239,78 +251,78 @@ def _parse_optional_string(value): settings = Settings() settings._register( - 'print_mode', - default='print', - environment_variable='XSUITE_PRINT_MODE', - choices=('print', 'suppress'), + "print_mode", + default="print", + environment_variable="XSUITE_PRINT_MODE", + choices=("print", "suppress"), environment_parser=_parse_choice, ) settings._register( - 'progress_indicator', - default='tqdm', - environment_variable='XSUITE_PROGRESS_INDICATOR', - choices=('tqdm', 'text', 'suppress'), + "progress_indicator", + default="tqdm", + environment_variable="XSUITE_PROGRESS_INDICATOR", + choices=("tqdm", "text", "suppress"), environment_parser=_parse_choice, ) settings._register( - 'allow_kernel_compilation', + "allow_kernel_compilation", default=False, - environment_variable='XSUITE_ALLOW_KERNEL_COMPILATION', + environment_variable="XSUITE_ALLOW_KERNEL_COMPILATION", choices=(False, True), environment_parser=_parse_boolean, value_type=bool, ) settings._register( - 'force_kernel_compilation', + "force_kernel_compilation", default=False, - environment_variable='XSUITE_FORCE_KERNEL_COMPILATION', + environment_variable="XSUITE_FORCE_KERNEL_COMPILATION", choices=(False, True), environment_parser=_parse_boolean, value_type=bool, ) settings._register( - 'show_kernel_diagnostics', + "show_kernel_diagnostics", default=False, - environment_variable='XSUITE_SHOW_KERNEL_DIAGNOSTICS', + environment_variable="XSUITE_SHOW_KERNEL_DIAGNOSTICS", choices=(False, True), environment_parser=_parse_boolean, value_type=bool, ) settings._register( - 'cffi_forbid_compile', + "cffi_forbid_compile", default=False, - environment_variable='XSUITE_CFFI_FORBID_COMPILE', + environment_variable="XSUITE_CFFI_FORBID_COMPILE", choices=(False, True), environment_parser=_parse_boolean, value_type=bool, ) settings._register( - 'cffi_keep_build_files', + "cffi_keep_build_files", default=False, - environment_variable='XSUITE_CFFI_KEEP_BUILD_FILES', + environment_variable="XSUITE_CFFI_KEEP_BUILD_FILES", choices=(False, True), environment_parser=_parse_boolean, value_type=bool, ) settings._register( - 'cuda_backend', - default='nvrtc', - environment_variable='XSUITE_CUDA_BACKEND', - choices=('nvrtc', 'clang'), + "cuda_backend", + default="nvrtc", + environment_variable="XSUITE_CUDA_BACKEND", + choices=("nvrtc", "clang"), environment_parser=_parse_choice, ) settings._register( - 'cuda_fast_compile', + "cuda_fast_compile", default=True, - environment_variable='XSUITE_CUDA_FAST_COMPILE', + environment_variable="XSUITE_CUDA_FAST_COMPILE", choices=(False, True), environment_parser=_parse_boolean, value_type=bool, ) settings._register( - 'cuda_compiler', + "cuda_compiler", default=None, - environment_variable='XSUITE_CUDA_COMPILER', + environment_variable="XSUITE_CUDA_COMPILER", environment_parser=_parse_optional_string, value_type=(str, type(None)), ) diff --git a/xobjects/string.py b/xobjects/string.py index 9bf4e72..60dc74f 100644 --- a/xobjects/string.py +++ b/xobjects/string.py @@ -61,7 +61,7 @@ def _inspect_args(cls, string_or_int): else: return Info(size=cls._size) - def _to_buffer(cls, buffer, offset, value, info=None): + def _to_buffer(cls, buffer, offset, value, info=None, container=None): # log.debug(f"{cls} to buffer {offset} `{value}`") if info is None: info = cls._inspect_args(value) @@ -87,7 +87,7 @@ def _get_data(cls, buffer, offset): ll = Int64._from_buffer(buffer, offset) return buffer.to_bytearray(offset + 8, ll - 8) - def _from_buffer(cls, buffer, offset=0, encoding="utf8"): + def _from_buffer(cls, buffer, offset=0, encoding="utf8", container=None): # TODO keep in mind that in windows many funcitons returns wchar encoded in utf16 return cls._get_data(buffer, offset).decode(encoding).rstrip("\x00") diff --git a/xobjects/struct.py b/xobjects/struct.py index be36f69..d0b9103 100644 --- a/xobjects/struct.py +++ b/xobjects/struct.py @@ -104,7 +104,9 @@ def __get__(self, instance, cls=None): return self else: ftype, offset = self.get_offset(instance) - return ftype._from_buffer(instance._buffer, offset) + return ftype._from_buffer( + instance._buffer, offset, container=instance + ) def __set__(self, instance, value): """ @@ -120,7 +122,9 @@ def __set__(self, instance, value): self.__get__(instance)._update(value) else: # TODO check if below is really needed ftype, offset = self.get_offset(instance) - ftype._to_buffer(instance._buffer, offset, value) + ftype._to_buffer( + instance._buffer, offset, value, container=instance + ) def get_offset(self, instance): # compatible with info if self.is_reference: @@ -321,7 +325,7 @@ def _post_init(self): pass @classmethod - def _from_buffer(cls, buffer, offset=0): + def _from_buffer(cls, buffer, offset=0, container=None): self = object.__new__(cls) self._buffer = buffer self._offset = offset @@ -336,7 +340,7 @@ def _from_buffer(cls, buffer, offset=0): return self @classmethod - def _to_buffer(cls, buffer, offset, value, info=None): + def _to_buffer(cls, buffer, offset, value, info=None, container=None): if isinstance(value, cls) and not cls._has_refs: # binary copy buffer.update_from_xbuffer( offset, value._buffer, value._offset, value._size @@ -518,8 +522,10 @@ def compile_class_kernels( extra_classes=(), extra_compile_args=(), ): - if (context.allow_prebuilt_kernels - and not settings.force_kernel_compilation): + if ( + context.allow_prebuilt_kernels + and not settings.force_kernel_compilation + ): with settings.override(print_mode="suppress"): try: from xsuite import ( diff --git a/xobjects/union.py b/xobjects/union.py index 6851fbe..b890719 100644 --- a/xobjects/union.py +++ b/xobjects/union.py @@ -106,14 +106,14 @@ def _inspect_args(cls, *args): raise ValueError(f"{value} has wrong number of arguments") @classmethod - def _from_buffer(cls, buffer, offset=0): + def _from_buffer(cls, buffer, offset=0, container=None): self = object.__new__(cls) self._buffer = buffer self._offset = offset return self @classmethod - def _to_buffer(cls, buffer, offset, value, info=None): + def _to_buffer(cls, buffer, offset, value, info=None, container=None): if info is None: info = cls._inspect_args(value) Int64._to_buffer(buffer, offset, info.typeid)