diff --git a/xobjects/context.py b/xobjects/context.py index 506fbce..a2c71e0 100644 --- a/xobjects/context.py +++ b/xobjects/context.py @@ -254,6 +254,7 @@ def add_kernels( extra_classes: Sequence[Type] = (), extra_headers: Sequence[SourceType] = (), compile: bool = True, # noqa + extra_compile_args: Sequence[str] = (), ): """ Adds user-defined kernels to the context. The kernel source @@ -333,6 +334,7 @@ def add_kernels( extra_classes=extra_classes, extra_headers=extra_headers, compile=compile, + extra_compile_args=extra_compile_args, ) self.kernels.update(generated_kernels) @@ -348,6 +350,7 @@ def build_kernels( extra_classes: Sequence[Type], extra_headers: Sequence[SourceType], compile: bool, + extra_compile_args: Sequence[str], ) -> Dict[Tuple[str, tuple], KernelType]: pass diff --git a/xobjects/context_cpu.py b/xobjects/context_cpu.py index e57ee41..7e6f60f 100644 --- a/xobjects/context_cpu.py +++ b/xobjects/context_cpu.py @@ -169,8 +169,8 @@ def add_kernels( specialize=True, apply_to_source=(), save_source_as=None, - extra_compile_args: Sequence[str] = ("-O3", "-Wno-unused-function"), - extra_link_args: Sequence[str] = ("-O3",), + extra_compile_args: Sequence[str] = (), + extra_link_args: Sequence[str] = (), extra_cdef="", extra_classes=(), extra_headers=(), @@ -271,13 +271,16 @@ def build_kernels( specialize=True, apply_to_source=(), save_source_as=None, - extra_compile_args=("-O3", "-Wno-unused-function"), - extra_link_args=("-O3",), + extra_compile_args=(), + extra_link_args=(), extra_cdef="", extra_classes=(), extra_headers=(), compile=True, # noqa ) -> Dict[Tuple[str, tuple], "KernelCpu"]: + extra_compile_args += ("-O3", "-Wno-unused-function") + extra_link_args += ("-O3",) + # Determine names and paths clean_up_so = not module_name module_name = module_name or str(uuid.uuid4().hex) @@ -409,20 +412,25 @@ def compile_kernel( ffi_interface.cdef("int omp_get_max_threads();") # Compile - xtr_compile_args = ["-std=c99"] - xtr_link_args = ["-std=c99"] + xtr_compile_args = ["-std=c99", "-DXO_CONTEXT_CPU"] + xtr_link_args = ["-std=c99", "-DXO_CONTEXT_CPU"] xtr_compile_args += extra_compile_args xtr_link_args += extra_link_args if self.openmp_enabled: xtr_compile_args.append("-fopenmp") xtr_link_args.append("-fopenmp") + xtr_compile_args.append("-DXO_CONTEXT_CPU_OPENMP") + xtr_link_args.append("-DXO_CONTEXT_CPU_OPENMP") # https://mac.r-project.org/openmp/ # on macos comment the above and uncomment the below flags to compile OpenMP with Xcode clang: # xtr_compile_args.append("-Xclang") # xtr_compile_args.append("-fopenmp") # xtr_link_args.append("-lomp") + else: + xtr_compile_args.append("-DXO_CONTEXT_CPU_SERIAL") + xtr_link_args.append("-DXO_CONTEXT_CPU_SERIAL") if os.name == "nt": # windows # TODO: to be handled properly diff --git a/xobjects/context_cupy.py b/xobjects/context_cupy.py index 556162c..fc16e2b 100644 --- a/xobjects/context_cupy.py +++ b/xobjects/context_cupy.py @@ -416,6 +416,7 @@ def build_kernels( specialize=True, apply_to_source=(), save_source_as=None, + extra_compile_args=(), extra_cdef=None, extra_classes=(), extra_headers=(), @@ -454,7 +455,10 @@ def build_kernels( with open(save_source_as, "w") as fid: fid.write(specialized_source) - module = cupy.RawModule(code=specialized_source) + extra_compile_args = (*extra_compile_args, "-DXO_CONTEXT_CUDA") + module = cupy.RawModule( + code=specialized_source, options=extra_compile_args + ) out_kernels = {} for pyname, kernel in kernel_descriptions.items(): diff --git a/xobjects/context_pyopencl.py b/xobjects/context_pyopencl.py index c5277e3..4424891 100644 --- a/xobjects/context_pyopencl.py +++ b/xobjects/context_pyopencl.py @@ -180,6 +180,7 @@ def build_kernels( specialize=True, apply_to_source=(), save_source_as=None, + extra_compile_args=(), extra_cdef=None, extra_classes=(), extra_headers=(), @@ -218,8 +219,13 @@ def build_kernels( with open(save_source_as, "w") as fid: fid.write(specialized_source) + extra_compile_args = ( + *extra_compile_args, + "-cl-std=CL2.0", + "-DXO_CONTEXT_CL", + ) prg = cl.Program(self.context, specialized_source).build( - options="-cl-std=CL2.0", + options=extra_compile_args, ) out_kernels = {} diff --git a/xobjects/hybrid_class.py b/xobjects/hybrid_class.py index 597c223..50b4d68 100644 --- a/xobjects/hybrid_class.py +++ b/xobjects/hybrid_class.py @@ -96,7 +96,7 @@ def default(self, obj): def _build_xofields_dict(bases, data): if "_xofields" in data.keys(): xofields = data["_xofields"].copy() - elif any(map(lambda b: hasattr(b, "_xofields"), bases)): + elif any(hasattr(b, "_xofields") for b in bases): n_filled = 0 for bb in bases: if hasattr(bb, "_xofields") and len(bb._xofields.keys()) > 0: diff --git a/xobjects/struct.py b/xobjects/struct.py index 9d4c48e..54e299f 100644 --- a/xobjects/struct.py +++ b/xobjects/struct.py @@ -473,6 +473,7 @@ def compile_class_kernels( apply_to_source=(), save_source_as=None, extra_classes=(), + extra_compile_args=(), ): if only_if_needed: all_found = True @@ -489,6 +490,7 @@ def compile_class_kernels( extra_classes=[cls] + list(extra_classes), apply_to_source=apply_to_source, save_source_as=save_source_as, + extra_compile_args=extra_compile_args, ) def compile_kernels( @@ -497,6 +499,7 @@ def compile_kernels( apply_to_source=(), save_source_as=None, extra_classes=(), + extra_compile_args=(), ): self.compile_class_kernels( context=self._context, @@ -504,6 +507,7 @@ def compile_kernels( apply_to_source=apply_to_source, save_source_as=save_source_as, extra_classes=extra_classes, + extra_compile_args=extra_compile_args, )