From ecacb97f26aa6517c89e48f41307609f35e12309 Mon Sep 17 00:00:00 2001 From: Mathias Louboutin Date: Fri, 6 Nov 2026 22:11:19 +0000 Subject: [PATCH 01/12] compiler: Add PyTorch tensor storage and code generation --- .github/workflows/pytest-gpu.yaml | 2 +- README.md | 3 + devito/arch/compiler.py | 38 +- devito/core/__init__.py | 9 + devito/core/torch.py | 46 +++ devito/data/allocators.py | 71 +++- devito/data/data.py | 450 ++++++++++++++++++---- devito/data/distributed/exchange.py | 8 +- devito/data/distributed/redistribution.py | 6 +- devito/ir/iet/torch.py | 25 ++ devito/operator/operator.py | 25 +- devito/operator/registry.py | 2 +- devito/passes/iet/definitions.py | 6 + devito/passes/iet/engine.py | 1 + devito/passes/iet/langbase.py | 2 +- devito/passes/iet/languages/targets.py | 21 + devito/passes/iet/languages/torch.py | 33 ++ devito/tools/dtypes_lowering.py | 33 +- devito/types/args.py | 3 + devito/types/dense.py | 69 ++-- devito/types/grid.py | 20 +- devito/types/sparse.py | 6 +- docs/torch.md | 213 ++++++++++ examples/pytorch/__init__.py | 1 + examples/pytorch/devito_layer.py | 154 ++++++++ pyproject.toml | 1 + requirements-testing.txt | 1 + requirements-torch.txt | 1 + tests/test_data_backends.py | 198 ++++++++++ tests/test_torch_allocators.py | 96 +++++ tests/test_torch_codegen.py | 248 ++++++++++++ tests/test_torch_layer.py | 80 ++++ tests/test_torch_openacc.py | 247 ++++++++++++ 33 files changed, 1976 insertions(+), 143 deletions(-) create mode 100644 devito/core/torch.py create mode 100644 devito/ir/iet/torch.py create mode 100644 devito/passes/iet/languages/torch.py create mode 100644 docs/torch.md create mode 100644 examples/pytorch/__init__.py create mode 100644 examples/pytorch/devito_layer.py create mode 100644 requirements-torch.txt create mode 100644 tests/test_data_backends.py create mode 100644 tests/test_torch_allocators.py create mode 100644 tests/test_torch_codegen.py create mode 100644 tests/test_torch_layer.py create mode 100644 tests/test_torch_openacc.py diff --git a/.github/workflows/pytest-gpu.yaml b/.github/workflows/pytest-gpu.yaml index 9bf17dab280..09b7b042c51 100644 --- a/.github/workflows/pytest-gpu.yaml +++ b/.github/workflows/pytest-gpu.yaml @@ -42,7 +42,7 @@ jobs: include: # -------------------- NVIDIA job -------------------- - name: pytest-gpu-acc-nvidia - test_files: "tests/test_adjoint.py tests/test_gpu_common.py tests/test_gpu_openacc.py tests/test_operator.py::TestEstimateMemory" + test_files: "tests/test_adjoint.py tests/test_gpu_common.py tests/test_gpu_openacc.py tests/test_torch_openacc.py tests/test_operator.py::TestEstimateMemory" base: "devitocodes/bases:nvidia-nvc12" runner_label: NVIDIAGPU test_drive_cmd: "nvidia-smi" diff --git a/README.md b/README.md index f9edc094bc2..b1c67a8a344 100644 --- a/README.md +++ b/README.md @@ -99,6 +99,9 @@ have seen in the past. ## Resources +For PyTorch-owned tensors, generated C++ CPU loops, and CUDA tensor pointers +with OpenACC, see the optional [PyTorch backend](docs/torch.md). + To learn how to use Devito, [here](https://github.com/devitocodes/devito/blob/main/examples) is a good place to start, with lots of examples and tutorials. diff --git a/devito/arch/compiler.py b/devito/arch/compiler.py index d3567656367..19d9a3a20d9 100644 --- a/devito/arch/compiler.py +++ b/devito/arch/compiler.py @@ -183,6 +183,10 @@ def __init__(self): ``suffix='4.9'``, the ``gcc-4.9`` will be used as JIT compiler. cpp : bool, optional, default=False If True, JIT compile using a C++ compiler. + std : str, optional, default=None + Override the compiler class default language standard. + openmp : bool, optional, default=False + Enable OpenMP compilation independently of the language. mpi : bool, optional, default=False If True, JIT compile using an MPI compiler. platform : Platform, optional @@ -202,10 +206,14 @@ def __init__(self, **kwargs): else: self._name = maybe_name + std = kwargs.pop('std', None) super().__init__(**kwargs) self.__lookup_cmds__() self._cpp = kwargs.get('cpp', self._default_cpp) + self._std = std + self.openmp = kwargs.get('openmp', False) or 'openmp' in kwargs.get( + 'language', configuration['language']) self.suffix = kwargs.get('suffix') if not kwargs.get('mpi'): @@ -247,6 +255,8 @@ def __new_with__(self, **kwargs): """ return self.__class__(name=self.name, suffix=kwargs.pop('suffix', self.suffix), mpi=kwargs.pop('mpi', configuration['mpi']), + cpp=kwargs.pop('cpp', self._cpp), + std=kwargs.pop('std', self._std), **kwargs) @property @@ -269,7 +279,7 @@ def version(self): @property def std(self): - return self._cxxstd if self._cpp else self._cstd + return self._std or (self._cxxstd if self._cpp else self._cstd) def get_version(self): result, stdout, stderr = call_capture_output((self.cc, "--version")) @@ -469,6 +479,9 @@ def __getstate__(self): def _signature_items(self): return (str(self._cmdline([])), self.src_ext, self.so_ext) + def add_defines(self, defines): + self.defines = filter_ordered(self.defines + as_list(defines)) + def add_include_dirs(self, dirs): self.include_dirs = filter_ordered(self.include_dirs + as_list(dirs)) @@ -521,7 +534,6 @@ def __init_finalize__(self, **kwargs): else: self.cflags = ['-march=native'] + self.cflags - language = kwargs.pop('language', configuration['language']) try: if self.version >= Version("4.9.0"): # Append the openmp flag regardless of the `language` value, @@ -529,7 +541,7 @@ def __init_finalize__(self, **kwargs): # they support `#pragma omp simd` self.ldflags += ['-fopenmp'] except (TypeError, ValueError): - if 'openmp' in language: + if self.openmp: self.ldflags += ['-fopenmp'] def __lookup_cmds__(self): @@ -589,7 +601,7 @@ def __init_finalize__(self, **kwargs): self.cflags += ['-mcpu=native'] else: self.cflags += ['-march=native'] - if 'openmp' in language: + if self.openmp: self.ldflags += ['-fopenmp'] def __lookup_cmds__(self): @@ -623,7 +635,7 @@ def __init_finalize__(self, **kwargs): if not configuration['safe-math']: self.cflags.append('-ffast-math') - if 'openmp' in language: + if self.openmp: self.ldflags += ['-fopenmp'] if isinstance(platform, NvidiaDevice): @@ -682,10 +694,10 @@ def __init_finalize__(self, **kwargs): if language == 'openacc': self.cflags.extend(['-mp', '-acc:gpu']) - elif 'openmp' in language: + elif self.openmp: self.cflags.extend(['-mp=gpu']) elif isinstance(platform, Cpu64): - if 'openmp' in language: + if self.openmp: self.cflags.append('-mp') if isinstance(platform, NvidiaArm): self.cflags.append(f'-mcpu={platform.march}') @@ -866,7 +878,6 @@ class IntelCompiler(Compiler): def __init_finalize__(self, **kwargs): platform = kwargs.pop('platform', configuration['platform']) - language = kwargs.pop('language', configuration['language']) self.cflags.append("-xHost") @@ -879,7 +890,7 @@ def __init_finalize__(self, **kwargs): # Systematically use 512-bit vectors if avx512 is available. self.cflags.append("-qopt-zmm-usage=high") - if 'openmp' in language: + if self.openmp: self.ldflags.append('-qopenmp') if kwargs.get('mpi'): @@ -931,11 +942,9 @@ class IntelKNLCompiler(IntelCompiler): def __init_finalize__(self, **kwargs): IntelCompiler.__init_finalize__(self, **kwargs) - language = kwargs.pop('language', configuration['language']) - self.cflags.append('-xMIC-AVX512') - if language != 'openmp': + if not self.openmp: warning("Running on Intel KNL without OpenMP is highly discouraged") @@ -950,7 +959,7 @@ def __init_finalize__(self, **kwargs): if language == 'sycl': warning(f"Use SyclCompiler (`sycl`) to jit-compile sycl, not {self.name}") - elif 'openmp' in language: + elif self.openmp: # Earlier versions to OneAPI 2023.2.0 (clang17 underneath), have an # OpenMP bug concerning reductions, hence with them we're forced to # use the obsolete -fopenmp @@ -1061,7 +1070,8 @@ def __init_finalize__(self, **kwargs): self._base.__init_finalize__(self, **kwargs) # Update cflags try: - extra_c = environ.get('CFLAGS').split(' ') + flags = 'CXXFLAGS' if self._cpp else 'CFLAGS' + extra_c = environ.get(flags, environ.get('CFLAGS')).split(' ') self.cflags = self.cflags + extra_c except AttributeError: pass diff --git a/devito/core/__init__.py b/devito/core/__init__.py index f226a9ff0d1..958edbdbb1a 100644 --- a/devito/core/__init__.py +++ b/devito/core/__init__.py @@ -31,6 +31,10 @@ DeviceCustomCXXOmpOperator, DeviceNoopCXXOmpOperator, DeviceAdvCXXOmpOperator, DeviceFsgCXXOmpOperator ) +from devito.core.torch import ( + Cpu64AdvTorchOperator, Cpu64CustomTorchOperator, Cpu64FsgTorchOperator, + Cpu64NoopTorchOperator +) from devito.operator.registry import operator_registry # Register CPU Operators @@ -82,6 +86,11 @@ operator_registry.add(PowerCXXAdvCOperator, Power, 'advanced', 'CXX') operator_registry.add(PowerAdvCXXOmpOperator, Power, 'advanced', 'CXXopenmp') +operator_registry.add(Cpu64NoopTorchOperator, Cpu64, 'noop', 'torch') +operator_registry.add(Cpu64AdvTorchOperator, Cpu64, 'advanced', 'torch') +operator_registry.add(Cpu64FsgTorchOperator, Cpu64, 'advanced-fsg', 'torch') +operator_registry.add(Cpu64CustomTorchOperator, Cpu64, 'custom', 'torch') + # Register Device Operators operator_registry.add(DeviceCustomOmpOperator, Device, 'custom', 'C') operator_registry.add(DeviceCustomOmpOperator, Device, 'custom', 'openmp') diff --git a/devito/core/torch.py b/devito/core/torch.py new file mode 100644 index 00000000000..2484c4c5cdf --- /dev/null +++ b/devito/core/torch.py @@ -0,0 +1,46 @@ +"""Torch language specializations of the standard CPU operator pipeline.""" + +from devito.core.cpu import ( + Cpu64AdvOperator, Cpu64CustomOperator, Cpu64FsgOperator, Cpu64NoopOperator +) +from devito.exceptions import InvalidOperator +from devito.ir.iet.torch import TorchCGen +from devito.passes.iet import TorchTarget, linearize +from devito.tools import timed_pass + + +class TorchOperatorMixin: + """Reuse CPU optimization passes with the Torch C++ printer and target.""" + + _ccode_handler = TorchCGen + _Target = TorchTarget + _compiler_options = {'cpp': True, 'std': 'c++17'} + + @classmethod + def _normalize_kwargs(cls, **kwargs): + if kwargs['options'].get('linearize') is False: + raise InvalidOperator("Torch C++ code generation requires linearization") + kwargs['options']['linearize'] = True + return super()._normalize_kwargs(**kwargs) + + +class Cpu64NoopTorchOperator(TorchOperatorMixin, Cpu64NoopOperator): + + @classmethod + @timed_pass(name='specializing.IET') + def _specialize_iet(cls, graph, **kwargs): + graph = super()._specialize_iet(graph, **kwargs) + linearize(graph, **kwargs) + return graph + + +class Cpu64AdvTorchOperator(TorchOperatorMixin, Cpu64AdvOperator): + pass + + +class Cpu64FsgTorchOperator(TorchOperatorMixin, Cpu64FsgOperator): + pass + + +class Cpu64CustomTorchOperator(TorchOperatorMixin, Cpu64CustomOperator): + pass diff --git a/devito/data/allocators.py b/devito/data/allocators.py index d883118f461..38d986d1223 100644 --- a/devito/data/allocators.py +++ b/devito/data/allocators.py @@ -9,7 +9,12 @@ from devito.logger import logger from devito.parameters import configuration -from devito.tools import infer_datasize, is_integer +from devito.tools import dtype_to_torch, infer_datasize, is_integer + +try: + import torch +except ImportError: + torch = None __all__ = [ 'ALLOC_ALIGNED', @@ -18,6 +23,7 @@ 'ALLOC_KNL_MCDRAM', 'ALLOC_NUMA_ANY', 'ALLOC_NUMA_LOCAL', + 'TorchAllocator', 'default_allocator', ] @@ -30,6 +36,9 @@ class AbstractMemoryAllocator: __metaclass__ = abc.ABCMeta + _array_type = np.ndarray + is_Device = False + guaranteed_alignment = 64 """Guaranteed data alignment.""" @@ -342,6 +351,66 @@ def put_local(self): return self._node == 'local' +class TorchAllocator(AbstractMemoryAllocator): + + """PyTorch-owned storage for ordinary Functions and TimeFunctions. + + Parameters + ---------- + device : str, optional + PyTorch device. Defaults to the supplied tensor's device, or CPU. + tensor : torch.Tensor, optional + Existing contiguous allocation, including halo and padding. As with + DataReference, use a no-op initializer to retain its values. + """ + + _array_type = torch.Tensor if torch is not None else None + + def __init__(self, device=None, tensor=None): + if torch is None: + raise ImportError("TorchAllocator requires PyTorch") + self.device = torch.device(device or getattr(tensor, 'device', 'cpu')) + if self.device.type not in ('cpu', 'cuda'): + raise ValueError("TorchAllocator supports CPU and CUDA storage") + if self.device.type == 'cpu': + # PyTorch reports CPU storage without a device index, even for cpu:0. + self.device = torch.device('cpu') + elif self.device.index is None: + self.device = torch.device('cuda', torch.cuda.current_device()) + self.tensor = tensor + + @classmethod + def available(cls): + return torch is not None + + @property + def is_Device(self): + return self.device.type == 'cuda' + + def alloc(self, shape, dtype, padding=0): + if self.tensor is None: + try: + left, right = padding + except TypeError: + left = right = padding + size = int(np.prod(shape)) + storage = torch.empty(size + left + right, device=self.device, + dtype=dtype_to_torch(dtype)) + tensor = storage[left:left + size].reshape(tuple(shape)) + else: + tensor = self.tensor + if tuple(tensor.shape) != tuple(shape): + raise ValueError( + f"External tensor must have allocated shape {tuple(shape)}") + if tensor.dtype != dtype_to_torch(dtype): + raise ValueError("External tensor dtype does not match Function dtype") + return tensor, None + + def free(self, *args): + # PyTorch owns the storage; dropping its last reference releases it. + return + + class DataReference(MemoryAllocator): """ diff --git a/devito/data/data.py b/devito/data/data.py index 38bbc80b040..51892f10d35 100644 --- a/devito/data/data.py +++ b/devito/data/data.py @@ -1,21 +1,33 @@ +import ctypes import weakref from collections.abc import Iterable import numpy as np -from devito.data.allocators import ALLOC_ALIGNED +from devito.data.allocators import ALLOC_ALIGNED, TorchAllocator from devito.data.utils import * +from devito.exceptions import InvalidArgument from devito.logger import warning from devito.parameters import configuration -from devito.tools import as_list, as_tuple, is_integer +from devito.tools import ( + as_list, as_tuple, c_restrict_void_p, dtype_from_torch, is_integer +) -__all__ = ['Data'] +try: + import torch + from torch.utils._pytree import tree_leaves, tree_map +except ImportError: + torch = None +__all__ = ['Data', 'NumpyData'] -class Data(np.ndarray): + +class Data: """ - A numpy.ndarray supporting distributed Dimensions. + Backend-independent indexing, decomposition and MPI behavior. + + Constructing Data directly selects the allocator backend (NumPy by default). Parameters ---------- @@ -26,69 +38,104 @@ class Data(np.ndarray): decomposition : tuple of Decomposition, optional The data decomposition, for each dimension. modulo : tuple of bool, optional - If the i-th entry is True, then the i-th array dimension uses modulo - indexing. + If the i-th entry is True, then the i-th array dimension uses modulo indexing. allocator : MemoryAllocator, optional - Used to allocate memory. Defaults to `ALLOC_ALIGNED`. + Used to allocate memory. Defaults to ``ALLOC_ALIGNED``. distributor : Distributor, optional - The distributor from which the original decomposition was produced. - Note that `decomposition` may differ from `distributor.decomposition`. - padding : int or 2-tuple of ints, optional - The number of points that are allocated before and after the data, - that is in addition to the requested shape. Defaults to 0. + The distributor from which the original decomposition was produced. Note that + the decomposition Parameter above may be different to distributor.decomposition. Notes ----- - NumPy array subclassing is described at: :: - - https://numpy.org/doc/stable/user/basics.subclassing.html - - Any view or copy created from ``self``, for instance via a slice operation - or a universal function ("ufunc" in NumPy jargon), will still be of type - `Data`. + Construction, index translation, view metadata, distributed assignment + and gather are shared. Concrete views implement native indexing and + storage access. NumPy is used for decomposition and MPI index metadata. """ def __new__(cls, shape, dtype, decomposition=None, modulo=None, allocator=ALLOC_ALIGNED, distributor=None, padding=0): - assert len(shape) == len(modulo) - ndarray, memfree_args = allocator.alloc(shape, dtype, padding=padding) - obj = ndarray.view(cls) - obj._allocator = allocator + assert modulo is None or len(shape) == len(modulo) + storage, memfree_args = allocator.alloc(shape, dtype, padding=padding) + obj = cls._from_array(storage, allocator=allocator) obj._memfree_args = memfree_args - obj._decomposition = decomposition or (None,)*len(shape) - obj._modulo = modulo or (False,)*len(shape) - obj._distributor = distributor + obj._initialize_metadata(shape, decomposition, modulo, distributor) + if memfree_args is not None: + weakref.finalize(storage, allocator.free, *memfree_args) + return obj + + @classmethod + def _from_array(cls, storage, allocator=ALLOC_ALIGNED): + """Create a non-owning Data view for allocated storage or an argument.""" + if not isinstance(storage, allocator._array_type): + raise InvalidArgument(f"{type(allocator).__name__} requires " + f"{allocator._array_type.__name__} storage") + if isinstance(storage, Data): + type(storage)._check_layout(storage, allocator) + return storage + if cls is Data: + if isinstance(storage, np.ndarray): + backend = NumpyData + elif torch is not None and isinstance(storage, torch.Tensor): + backend = TorchData + else: + raise InvalidArgument( + f"Unsupported native storage type {type(storage).__name__}") + else: + backend = cls + backend._check_layout(storage, allocator) + obj = backend._from_storage(storage) + obj._allocator = allocator + obj._memfree_args = None + obj._initialize_metadata(storage.shape) + return obj + + @staticmethod + def _is_array(value): + return isinstance(value, np.ndarray) or ( + torch is not None and isinstance(value, torch.Tensor)) + + @staticmethod + def _check_layout(data, allocator): + """Check the native storage constraints imposed by the allocator.""" + return + + def _arg_check(self, dtype): + """Check backend-specific requirements for an Operator argument.""" + return + + def _arg_apply(self, access_mode): + """Complete native storage writes after generated code returns.""" + return + + def __init__(self, *args, **kwargs): + pass + + @classmethod + def _from_storage(cls, storage): + """Wrap native storage without allocating or initializing metadata.""" + raise NotImplementedError("Data backends must provide native storage views") + + def _initialize_metadata(self, shape, decomposition=None, modulo=None, + distributor=None): + self._decomposition = decomposition or (None,)*len(shape) + self._modulo = modulo or (False,)*len(shape) + self._distributor = distributor # This cannot be a property, as Data objects constructed from this # object might not have any `decomposition`, but they would still be # distributed. Hence, in `__array_finalize__` we must copy this value - obj._is_distributed = any(i is not None for i in obj._decomposition) + self._is_distributed = any(i is not None for i in self._decomposition) # Saves the last index used in `__getitem__`. This allows `__array_finalize__` # to reconstruct information about the computed view (e.g., `decomposition`) - obj._index_stash = None + self._index_stash = None # Sanity check -- A Dimension can't be at the same time modulo-iterated # and MPI-distributed - assert all(i is None for i, j in zip(obj._decomposition, obj._modulo, strict=True) - if j is True) - - if memfree_args is not None: - # Release the memory once `ndarray` is gone, and with it every view - # of it -- this Data, but also whatever was handed out by e.g. - # `Function.data` or `Function._data_allocated`, since NumPy anchors - # the `base` chain of all of them on `ndarray`. Releasing it when - # this Data dies instead would leave those views pointing at freed - # memory as soon as the owning Function went out of scope - weakref.finalize(ndarray, allocator.free, *memfree_args) + assert all(i is None for i, j in + zip(self._decomposition, self._modulo, strict=True) if j is True) - return obj - - def __reduce__(self): - warning("Pickling of `Data` objects is not supported. Casting to `numpy.ndarray`") - return self.view(np.ndarray).__reduce__() - - def __array_finalize__(self, obj): + def _finalize_metadata(self, obj): # `self` is the newly created object # `obj` is the object from which `self` was created if obj is None: @@ -151,7 +198,7 @@ def __array_finalize__(self, obj): @property def _local(self): """A view of ``self`` with global indexing disabled.""" - ret = self.view() + ret = self._view() ret._is_distributed = False return ret @@ -201,10 +248,10 @@ def _scattered_exchange(self, glb_idx): return None def __repr__(self): - return super(Data, self._local).__repr__() + return repr(self._raw) def __str__(self): - return super(Data, self._local).__str__() + return str(self._raw) def transpose(self, *axes): """ @@ -227,7 +274,7 @@ def transpose(self, *axes): else tuple(ax % self.ndim for ax in axes) ) - ret = super().transpose(*axes) + ret = self._transpose_local(new_order) ret._decomposition = tuple(self._decomposition[i] for i in new_order) ret._modulo = tuple(self._modulo[i] for i in new_order) ret._is_distributed = any(d is not None for d in ret._decomposition) @@ -241,7 +288,7 @@ def swapaxes(self, axis1, axis2): """ axis1 = axis1 % self.ndim axis2 = axis2 % self.ndim - ret = super().swapaxes(axis1, axis2) + ret = self._swapaxes_local(axis1, axis2) order = list(range(self.ndim)) order[axis1], order[axis2] = order[axis2], order[axis1] ret._decomposition = tuple(self._decomposition[i] for i in order) @@ -268,7 +315,7 @@ def __getitem__(self, glb_idx): return None else: self._index_stash = glb_idx - retval = super().__getitem__(loc_idx) + retval = self._getitem_local(loc_idx) self._index_stash = None return retval @@ -283,13 +330,8 @@ def __setitem__(self, glb_idx, val): if loc_idx is NONLOCAL: # no-op return - elif np.isscalar(val): - if index_is_basic(loc_idx): - # Won't go through `__getitem__` as it's basic indexing mode, - # so we should just propagate `loc_idx` - super().__setitem__(loc_idx, val) - else: - super().__setitem__(glb_idx, val) + elif np.isscalar(val) or (self._is_array(val) and val.ndim == 0): + self._setitem_local(loc_idx, val) elif isinstance(val, Data) and val._is_decomposed: # Lazy import to avoid a circular dependency (see # `_scattered_exchange`). @@ -312,7 +354,7 @@ def __setitem__(self, glb_idx, val): data_global = [] idx_global = [] for j in range(nprocs): - data_global.append(comm.bcast(np.array(val), root=j)) + data_global.append(comm.bcast(val._host_array().copy(), root=j)) idx_global.append(comm.bcast(idx, root=j)) # Set the data: for j in range(nprocs): @@ -320,11 +362,12 @@ def __setitem__(self, glb_idx, val): or data_global[j].size == 0 if not skip: self.__setitem__(idx_global[j], data_global[j]) - elif isinstance(val, np.ndarray) or is_windowed(val): + elif self._is_array(val) or is_windowed(val): if self._is_decomposed: # `val` is replicated, `self` is decomposed -> `val` gets decomposed glb_idx = self._normalize_index(glb_idx) glb_idx, val = self._process_args(glb_idx, val) + loc_idx = self._index_glb_to_loc(glb_idx) val_idx = [index_dist_to_repl(i, dec) for i, dec in zip(glb_idx, self._decomposition, strict=True)] if NONLOCAL in val_idx: @@ -358,22 +401,23 @@ def __setitem__(self, glb_idx, val): else: val_idx = None - super().__setitem__(glb_idx, val if val_idx is None else val[val_idx]) + self._setitem_local(loc_idx, val if val_idx is None else val[val_idx]) elif isinstance(val, Iterable): if self._is_decomposed: raise NotImplementedError("With MPI, data can only be set " "via scalars, numpy arrays or " "other data ") - super().__setitem__(glb_idx, val) + self._setitem_local(loc_idx, val) else: raise ValueError(f"Cannot insert obj of type `{type(val)}` into a Data") def _normalize_index(self, idx): + idx = self._index_array(idx) if isinstance(idx, np.ndarray): # Advanced indexing mode return (idx,) else: - idx = as_tuple(idx) + idx = tuple(self._index_array(i) for i in as_tuple(idx)) if any(i is Ellipsis for i in idx): # Explicitly replace the Ellipsis items = (slice(None),)*(self.ndim - len(idx) + 1) @@ -418,9 +462,8 @@ def _process_args(self, idx, val): val._distributor._rebuild(val.shape, self._distributor.dimensions, self._distributor.comm) - new_val = Data(val.shape, val.dtype.type, - decomposition=val._decomposition, modulo=val._modulo, - distributor=distributor) + new_val = val._new_data(val.shape, decomposition=val._decomposition, + modulo=val._modulo, distributor=distributor) slc = as_tuple([slice(None, None, 1) for j in transform]) new_val[slc] = val[slc] return as_tuple(processed), new_val[as_tuple(transform)] @@ -545,10 +588,10 @@ def as_axes(v): in zip(start, stop, step, strict=True)) if not (self._is_decomposed and self._distributor.nprocs > 1): - return np.array(self[idx]) + return self._copy_native(self[idx]) local = self[idx] - block = np.ascontiguousarray(np.asarray(local)) + block = np.ascontiguousarray(local._host_array()) # Result indices this rank owns, per axis (replicated axis -> full extent) owned = [ np.arange(block.shape[ax], dtype=np.int64) if dec is None @@ -565,12 +608,275 @@ def as_axes(v): len(range(*sl.indices(d.size if d is not None else self.shape[ax]))) for ax, (sl, d) in enumerate(zip(idx, self._decomposition, strict=True)) ) - out = np.empty(shape, dtype=self.dtype) + out = np.empty(shape, dtype=self._numpy_dtype) for blk, idxs in gathered: if blk.size: out[np.ix_(*idxs)] = blk - return out + return self._from_mpi(out) + + def _new_data(self, shape, **kwargs): + kwargs.setdefault('modulo', (False,)*len(shape)) + return type(self)(shape, self._numpy_dtype.type, **kwargs) + + @staticmethod + def _index_array(value): + return value def reset(self): """Set all Data entries to 0.""" self[:] = 0.0 + + +class NumpyData(Data, np.ndarray): + + """NumPy storage implementation of Data. + + NumPy array operations supply local views and the native ctypes pointer. + """ + + @classmethod + def _from_storage(cls, storage): + return storage.view(cls) + + def __reduce__(self): + warning("Pickling of `Data` objects is not supported. Casting to `numpy.ndarray`") + return np.array(self).__reduce__() + + def __array_finalize__(self, obj): + self._finalize_metadata(obj) + + @property + def _raw(self): + return np.ndarray.view(self, np.ndarray) + + @property + def _numpy_dtype(self): + return self.dtype + + @property + def _C_data(self): + return self.ctypes.data_as(c_restrict_void_p) + + def _readonly(self): + view = self.view() + view.setflags(write=False) + return view + + def _host_array(self): + return np.asarray(self) + + def _view(self): + return np.ndarray.view(self) + + def _transpose_local(self, axes): + return np.ndarray.transpose(self, axes) + + def _swapaxes_local(self, axis1, axis2): + return np.ndarray.swapaxes(self, axis1, axis2) + + def _getitem_local(self, index): + return np.ndarray.__getitem__(self, index) + + def _setitem_local(self, index, value): + self._raw[index] = value._raw if isinstance(value, Data) else value + + @staticmethod + def _from_mpi(value): + return value + + @staticmethod + def _copy_native(value): + return np.array(value) + + +if torch is not None: + + class TorchData(Data, torch.Tensor): + """A tensor inheriting Data's modulo, decomposition and MPI semantics. + + Parameters + ---------- + shape : tuple of int + Local allocation shape, including halo and padding. + dtype : numpy.dtype + Compiler dtype used to select the tensor element type. + decomposition : tuple of Decomposition, optional + Global-to-local decomposition for each axis. + modulo : tuple of bool, optional + Axes using circular indexing, such as buffered time levels. + allocator : TorchAllocator, optional + Storage allocator. Defaults to a CPU TorchAllocator. + distributor : Distributor, optional + MPI topology and communicator associated with this allocation. + + Notes + ----- + Views and tensor arithmetic preserve Data metadata. Negative-stride reads + materialize tensors because PyTorch cannot represent negative strides. + """ + + def __new__(cls, shape, dtype, decomposition=None, modulo=None, allocator=None, + distributor=None): + if allocator is None: + allocator = TorchAllocator() + return super().__new__(cls, shape, dtype, + decomposition=decomposition, + modulo=modulo, allocator=allocator, + distributor=distributor) + + @classmethod + def _from_storage(cls, tensor): + return tensor.as_subclass(cls) + + @classmethod + def __torch_function__(cls, func, types, args=(), kwargs=None): + parent = next((i for i in tree_leaves((args, kwargs or {})) + if isinstance(i, cls)), None) + result = super().__torch_function__(func, types, args, kwargs) + + def finalize(value): + if isinstance(value, cls) and value is not parent: + value._finalize_metadata(parent) + return value + + return tree_map(finalize, result) + + @property + def _raw(self): + return self.as_subclass(torch.Tensor) + + @property + def _numpy_dtype(self): + return dtype_from_torch(self.dtype) + + def _wrap(self, tensor): + if tensor.ndim == 0: + return tensor + view = tensor.as_subclass(type(self)) + view._finalize_metadata(self) + return view + + def _view(self): + return self._wrap(self._raw.view(self.shape)) + + def _transpose_local(self, axes): + return self._wrap(self._raw.permute(axes)) + + def _swapaxes_local(self, axis1, axis2): + return self._wrap(self._raw.swapaxes(axis1, axis2)) + + @staticmethod + def _index_array(value): + # Index arrays are control metadata, not numerical field buffers. + if isinstance(value, torch.Tensor): + return value.detach().cpu().numpy() + return value + + def _negative_indices(self, index): + index = index if isinstance(index, tuple) else (index,) + index = index + (slice(None),)*(self.ndim-len(index)) + if not any(isinstance(i, slice) and i.step is not None and i.step < 0 + for i in index): + return None + if not all(is_integer(i) or isinstance(i, slice) for i in index): + raise NotImplementedError("Negative slices combined with " + "advanced indexing") + axes, shape = [], [] + for i, size in zip(index, self.shape, strict=False): + if is_integer(i): + axes.append(torch.tensor([int(i)], device=self.device)) + else: + positions = range(*i.indices(size)) + # Normalize empty endpoints for either step sign + stop = positions.start + len(positions)*positions.step + axis = torch.arange(positions.start, stop, positions.step, + device=self.device) + axes.append(axis) + shape.append(len(positions)) + return axes, tuple(shape) + + def _getitem_local(self, index): + negative = self._negative_indices(index) + if negative is None: + result = self._raw[index] + else: + axes, shape = negative + result = self._raw[torch.meshgrid(*axes, indexing='ij')].reshape(shape) + return self._wrap(result) + + def _setitem_local(self, index, value): + if isinstance(value, Data): + value = value._raw + if isinstance(value, np.ndarray) and any(i < 0 for i in value.strides): + value = value.copy() + value = torch.as_tensor(value, dtype=self.dtype, device=self.device) + negative = self._negative_indices(index) + if negative is None: + self._raw[index] = value + else: + axes, shape = negative + expanded = torch.broadcast_to(value, shape) + expanded = expanded.reshape(tuple(len(a) for a in axes)) + self._raw[torch.meshgrid(*axes, indexing='ij')] = expanded + + def _new_data(self, shape, **kwargs): + kwargs.setdefault('modulo', (False,)*len(shape)) + return type(self)(shape, self._numpy_dtype, + allocator=TorchAllocator(device=self.device), **kwargs) + + def _from_mpi(self, value): + return torch.as_tensor(value, dtype=self.dtype, device=self.device) + + @staticmethod + def _copy_native(value): + if isinstance(value, Data): + value = value._raw + return None if value is None else value.clone() + + def fill(self, value): + self.fill_(value) + + def copy(self): + """Copy storage and view metadata for the standard Data interface.""" + return self.clone() + + @staticmethod + def _check_layout(data, allocator): + if data.layout != torch.strided or not data.is_contiguous(): + raise InvalidArgument("TorchAllocator requires contiguous " + "strided storage") + if data.device != allocator.device: + raise InvalidArgument("Tensor device must match allocator device") + if data.requires_grad or data.is_conj() or data.is_neg(): + raise InvalidArgument("Autograd and lazy tensor views are " + "unsupported") + if data.data_ptr() % allocator.guaranteed_alignment: + raise InvalidArgument("Tensor storage does not satisfy " + "Function alignment") + + def _arg_check(self, dtype): + if self._numpy_dtype != np.dtype(dtype): + raise InvalidArgument("Tensor dtype must match Function dtype") + + @property + def _C_data(self): + if self.is_cuda: + torch.cuda.synchronize(self.device) + return ctypes.cast(self.data_ptr(), c_restrict_void_p) + + def _arg_apply(self, access_mode): + if self.is_cuda: + torch.cuda.synchronize(self.device) + if access_mode is not None and access_mode.is_write: + torch.autograd.graph.increment_version(self) + + def _readonly(self): + # PyTorch has no read-only views; snapshots protect the source storage. + return self.clone() + + def _host_array(self): + if self.is_cuda: + torch.cuda.synchronize(self.device) + return self._raw.detach().cpu().numpy() + + __all__.append('TorchData') diff --git a/devito/data/distributed/exchange.py b/devito/data/distributed/exchange.py index 2e82f3cd74e..78facf3b11b 100644 --- a/devito/data/distributed/exchange.py +++ b/devito/data/distributed/exchange.py @@ -50,11 +50,15 @@ def __init__(self, data, idx): def get(self): """Return `data[idx]` as a NumPy array.""" - return self._plan.get(np.asarray(self._data)) + return self._data._from_mpi(self._plan.get(self._data._host_array())) def put(self, value): """Assign `data[idx] = value`.""" - self._plan.put(np.asarray(self._data), value) + if hasattr(value, '_host_array'): + value = value._host_array() + local = self._data._host_array() + self._plan.put(local, value) + self._data._setitem_local(Ellipsis, local) class _ExchangeKey: diff --git a/devito/data/distributed/redistribution.py b/devito/data/distributed/redistribution.py index 51a77420ba2..f85a266fa5b 100644 --- a/devito/data/distributed/redistribution.py +++ b/devito/data/distributed/redistribution.py @@ -61,7 +61,9 @@ def redistribute_set(data, glb_idx, other): return False layout, gcoords, values = spec - _push(layout, gcoords, values, np.asarray(data)) + local = data._host_array() + _push(layout, gcoords, values, local) + data._setitem_local(Ellipsis, local) return True @@ -110,7 +112,7 @@ def _structured_spec(data, glb_idx, other): mesh = np.meshgrid(*coords_per_axis, indexing='ij') gcoords = {axis: m.reshape(-1) for axis, m in enumerate(mesh)} - values = np.ascontiguousarray(np.asarray(other)).reshape(-1) + values = np.ascontiguousarray(other._host_array()).reshape(-1) layout = Layout(data._distributor, decomposition, global_shape) return layout, gcoords, values diff --git a/devito/ir/iet/torch.py b/devito/ir/iet/torch.py new file mode 100644 index 00000000000..cef9304edf3 --- /dev/null +++ b/devito/ir/iet/torch.py @@ -0,0 +1,25 @@ +"""C++ printer for optimized loops over PyTorch tensor storage.""" + +import cgen as c + +from devito.data import FULL +from devito.ir.iet.visitors import CGen + + +class TorchCGen(CGen): + """Print typed ATen storage views followed by Devito's flattened loops.""" + + def visit_PointerCast(self, o): + cast = super().visit_PointerCast(o) + f = o.function + if not f.is_DiscreteFunction: + return cast + shape = ', '.join( + f'static_cast({self.ccode(f._C_get_field(FULL, d).size)})' + for d in f.dimensions + ) + dtype = self.ccode(f.indexed._C_typedata) + tensor = (f'at::from_blob({cast.data}, {{{shape}}}, ' + f'at::TensorOptions().dtype(c10::CppTypeToScalarType<{dtype}>::value)' + '.device(at::kCPU))') + return c.Initializer(cast.vdecl, f'{tensor}.data_ptr<{dtype}>()') diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 38b2d4dc28f..d07deb3beb0 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -166,6 +166,8 @@ class Operator(Callable): refer to the relevant documentation. """ + _compiler_options = {} + def __new__(cls, expressions, **kwargs): if expressions is None: # Return a dummy Callable. This is exploited by unpickling. Users @@ -520,6 +522,7 @@ def _lower_iet(cls, uiet, **kwargs): # Lower IET to a target-specific IET graph = Graph(iet, **kwargs) + cls._Target.initialize(graph, **kwargs) graph = cls._specialize_iet(graph, **kwargs) # Instrument the IET for C-level profiling @@ -792,7 +795,8 @@ def _postprocess_errors(self, retval, comm=None): def _postprocess_arguments(self, args, **kwargs): """Process runtime arguments upon returning from ``.apply()``.""" for p in self.parameters: - p._arg_apply(args[p.name], alias=kwargs.get(p.name)) + p._arg_apply(args[p.name], alias=kwargs.get(p.name), + access_mode=self._access_modes.get(p)) @cached_property def _known_arguments(self): @@ -852,10 +856,14 @@ def includes(self): def namespaces(self): return OrderedSet(*self._printer._namespaces).union(self._namespaces) + @property + def _ccode_handler(self): + from devito.ir.iet.visitors import CGen + return CGen + @cached_property def ccode(self): - from devito.ir.iet.visitors import CGen - return CGen(printer=self._printer).visit(self) + return self._ccode_handler(printer=self._printer).visit(self) def _jit_compile(self): """ @@ -1743,6 +1751,9 @@ def parse_kwargs(**kwargs): kwargs['language'] = configuration['language'] # `compiler` + backend = operator_selector(**kwargs) + compiler_options = dict(backend._compiler_options) + compiler_options['openmp'] = options['openmp'] compiler = kwargs.get('compiler') if compiler is not None: if not isinstance(compiler, str): @@ -1752,14 +1763,16 @@ def parse_kwargs(**kwargs): kwargs['compiler'] = compiler_registry[compiler](platform=kwargs['platform'], language=kwargs['language'], mpi=configuration['mpi'], - name=compiler) + name=compiler, + **compiler_options) elif any([platform, language]): kwargs['compiler'] =\ configuration['compiler'].__new_with__(platform=kwargs['platform'], language=kwargs['language'], - mpi=configuration['mpi']) + mpi=configuration['mpi'], + **compiler_options) else: - kwargs['compiler'] = configuration['compiler'].__new_with__() + kwargs['compiler'] = configuration['compiler'].__new_with__(**compiler_options) # Make sure compiler and language are compatible if compiler is not None and kwargs['compiler']._cpp and \ diff --git a/devito/operator/registry.py b/devito/operator/registry.py index 35e5646204f..51dc6d34f30 100644 --- a/devito/operator/registry.py +++ b/devito/operator/registry.py @@ -27,7 +27,7 @@ class OperatorRegistry(OrderedDict, metaclass=Singleton): _modes = ('noop', 'advanced', 'advanced-fsg') _languages = ('C', 'CXX', 'openmp', 'Copenmp', 'CXXopenmp', - 'openacc', 'cuda', 'hip', 'sycl') + 'openacc', 'cuda', 'hip', 'sycl', 'torch') _accepted = _modes + tuple(product(_modes, _languages)) def add(self, operator, platform, mode, language='C'): diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index 02e7f4b599d..a45793b65ce 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -9,6 +9,7 @@ import numpy as np +from devito.exceptions import InvalidOperator from devito.ir import ( Block, Call, Definition, DummyExpr, EntryFunction, FindNodes, FindSymbols, Iteration, List, MapExprStmts, Return, Transformer, make_callable @@ -481,6 +482,11 @@ def place_definitions(self, iet, globs=None, **kwargs): iet : Callable The input Iteration/Expression tree. """ + for expr, parents in MapExprStmts().visit(iet).items(): + if any(f.is_DiscreteFunction and f._mem_local for f in expr.functions) and \ + not any(isinstance(i, self.langbb.DeviceIteration) for i in parents): + raise InvalidOperator("Device-resident Functions cannot be " + "accessed by host code") storage = Storage() # Process inline definitions diff --git a/devito/passes/iet/engine.py b/devito/passes/iet/engine.py index 677935462d9..77d372c019a 100644 --- a/devito/passes/iet/engine.py +++ b/devito/passes/iet/engine.py @@ -157,6 +157,7 @@ def apply(self, func, **kwargs): compiler.add_library_dirs(as_tuple(metadata.get('lib_dirs')), rpath=metadata.get('rpath', False)) compiler.add_libraries(as_tuple(metadata.get('libs'))) + compiler.add_defines(as_tuple(metadata.get('defines'))) except KeyError: pass diff --git a/devito/passes/iet/langbase.py b/devito/passes/iet/langbase.py index 7fbbfa84ba7..7b8d7e692d9 100644 --- a/devito/passes/iet/langbase.py +++ b/devito/passes/iet/langbase.py @@ -511,7 +511,7 @@ def _(iet): def _device_pointers(self, iet): functions = FindSymbols().visit(iet) - devfuncs = [f for f in functions if f.is_Array and f._mem_local] + devfuncs = [f for f in functions if f.is_AbstractFunction and f._mem_local] return set(devfuncs) def _is_offloadable(self, iet): diff --git a/devito/passes/iet/languages/targets.py b/devito/passes/iet/languages/targets.py index c998b9315a7..204900920f5 100644 --- a/devito/passes/iet/languages/targets.py +++ b/devito/passes/iet/languages/targets.py @@ -9,6 +9,8 @@ DeviceOmpDataManager, DeviceOmpizer, DeviceOmpOrchestrator, OmpDataManager, Ompizer, OmpOrchestrator, SimdOmpizer ) +from devito.passes.iet.languages.torch import TorchPrinter +from devito.passes.iet.languages.torch import initialize as initialize_torch __all__ = [ 'COmpTarget', @@ -19,6 +21,7 @@ 'DeviceCXXOmpTarget', 'DeviceOmpTarget', 'OmpTarget', + 'TorchTarget', ] @@ -28,6 +31,10 @@ class Target: Orchestrator = None Printer = None + @staticmethod + def initialize(graph, **kwargs): + return + @classmethod def langbb(cls): return cls.Parizer.langbb @@ -90,3 +97,17 @@ class DeviceAccTarget(Target): DataManager = DeviceAccDataManager Orchestrator = AccOrchestrator Printer = AccPrinter + + +class TorchTarget(CXXOmpTarget): + initialize = staticmethod(initialize_torch) + Printer = TorchPrinter + + @classmethod + def langbb(cls): + return CXXOmpTarget.langbb() + + @staticmethod + def Parizer(**kwargs): + parizer = CXXOmpizer if kwargs['options']['openmp'] else CXXSimdOmpizer + return parizer(**kwargs) diff --git a/devito/passes/iet/languages/torch.py b/devito/passes/iet/languages/torch.py new file mode 100644 index 00000000000..cdf84d981b7 --- /dev/null +++ b/devito/passes/iet/languages/torch.py @@ -0,0 +1,33 @@ +"""JIT requirements for code using the PyTorch C++ runtime.""" + +try: + import torch + from torch.utils.cpp_extension import include_paths, library_paths +except ImportError: + torch = None + +from devito.passes.iet.engine import iet_pass +from devito.passes.iet.languages.CXX import CXXPrinter + +__all__ = ['TorchPrinter', 'initialize'] + + +@iet_pass +def initialize(iet, **kwargs): + """Declare the headers and toolchain requirements of the Torch language.""" + if torch is None: + raise ImportError("PyTorch C++ support is required for language='torch'") + + paths = library_paths() + return iet, { + 'includes': ['ATen/ops/from_blob.h'], + 'include_dirs': include_paths(), + 'lib_dirs': paths, + 'libs': ['torch_cpu', 'torch', 'c10'], + 'defines': [f'_GLIBCXX_USE_CXX11_ABI={int(torch.compiled_with_cxx11_abi())}'], + 'rpath': True + } + + +class TorchPrinter(CXXPrinter): + pass diff --git a/devito/tools/dtypes_lowering.py b/devito/tools/dtypes_lowering.py index 02f232e2b8e..f14b8099517 100644 --- a/devito/tools/dtypes_lowering.py +++ b/devito/tools/dtypes_lowering.py @@ -9,6 +9,11 @@ import numpy as np from cgen import dtype_to_ctype as cgen_dtype_to_ctype +try: + import torch +except ImportError: + torch = None + from .utils import as_tuple __all__ = ['int2', 'int3', 'int4', 'float2', 'float3', 'float4', 'double2', # noqa @@ -16,7 +21,8 @@ 'dtype_to_cstr', 'dtype_to_ctype', 'infer_datasize', 'dtype_to_mpitype', 'dtype_len', 'ctypes_to_cstr', 'c_restrict_void_p', 'ctypes_vector_mapper', 'is_external_ctype', 'infer_dtype', 'infer_dtype_from_args', 'extract_dtype', - 'CustomDtype', 'mpi4py_mapper', 'as_fp64_decimal'] + 'CustomDtype', 'mpi4py_mapper', 'as_fp64_decimal', + 'dtype_to_torch', 'dtype_from_torch'] def as_fp64_decimal(v): @@ -412,3 +418,28 @@ def infer_dtype_from_args(args): dtypes = {extract_dtype(i) for i in args} - {None} return infer_dtype(dtypes, vector=True) + + +# PyTorch and NumPy scalar types used by Function storage +_torch_dtypes = {} if torch is None else { + np.dtype('float32'): torch.float32, np.dtype('float64'): torch.float64, + np.dtype('int32'): torch.int32, np.dtype('int64'): torch.int64 +} + + +def dtype_to_torch(dtype): + """Lower a Function dtype to the corresponding PyTorch scalar type.""" + if torch is None: + raise ImportError("PyTorch is required for tensor allocation") + try: + return _torch_dtypes[np.dtype(dtype)] + except KeyError: + raise TypeError(f"Unsupported tensor dtype {dtype}") from None + + +def dtype_from_torch(dtype): + """Return the compiler dtype corresponding to a PyTorch scalar type.""" + try: + return next(n for n, t in _torch_dtypes.items() if t == dtype) + except StopIteration: + raise TypeError(f"Unsupported tensor dtype {dtype}") from None diff --git a/devito/types/args.py b/devito/types/args.py index 9e57339b632..634f50be67f 100644 --- a/devito/types/args.py +++ b/devito/types/args.py @@ -54,6 +54,9 @@ def _arg_apply(self, *args, **kwargs): """ Postprocess arguments upon returning from dynamically executed code. May be called if self's state needs to be updated. + + Array arguments receive the compiler's access mode through the optional + ``access_mode`` keyword so native storage can record writes. """ # By default, this is a no-op return diff --git a/devito/types/dense.py b/devito/types/dense.py index 30f543a47c2..f103eca8077 100644 --- a/devito/types/dense.py +++ b/devito/types/dense.py @@ -1,5 +1,5 @@ from collections import namedtuple -from ctypes import POINTER, Structure, byref, c_int, c_ulong, c_void_p, cast +from ctypes import POINTER, Structure, byref, c_int, c_ulong, c_void_p from functools import cached_property, reduce, wraps from operator import mul from textwrap import dedent, wrap @@ -7,6 +7,7 @@ import numpy as np import sympy +from devito.arch import Device from devito.builtins import assign from devito.data import ( CENTER, DOMAIN, FULL, HALO, LEFT, NOPAD, OWNED, RIGHT, Data, default_allocator @@ -22,8 +23,8 @@ from devito.parameters import configuration from devito.symbolics import FieldFromPointer, IndexedPointer, normalize_args from devito.tools import ( - ReducerMap, as_tuple, c_restrict_void_p, dtype_to_ctype, flatten, humanbytes, - is_integer, memoized_meth, mpi4py_mapper + ReducerMap, as_tuple, c_restrict_void_p, flatten, humanbytes, is_integer, + memoized_meth, mpi4py_mapper ) from devito.types.args import ArgProvider from devito.types.basic import AbstractFunction @@ -83,6 +84,8 @@ def __init_finalize__(self, *args, function=None, **kwargs): self._data = None self._first_touch = kwargs.get('first_touch', configuration['first-touch']) self._allocator = kwargs.get('allocator') or default_allocator() + if self._allocator.is_Device: + self._space = 'local' # Data initialization initializer = kwargs.get('initializer') @@ -101,7 +104,7 @@ def __init_finalize__(self, *args, function=None, **kwargs): elif initializer is None or callable(initializer) or self.alias: # Initialization postponed until the first access to .data self._initializer = initializer - elif isinstance(initializer, (np.ndarray, list, tuple)): + elif isinstance(initializer, (list, tuple)) or Data._is_array(initializer): # Allocate memory and initialize it. Note that we do *not* hold # a reference to the user-provided buffer self._initializer = None @@ -585,7 +588,7 @@ def _data_with_inhalo(self): """ self._is_halo_dirty = True self._halo_exchange() - return np.asarray(self._data[self._mask_inhalo]) + return self._data[self._mask_inhalo]._raw @property @_allocate_memory @@ -608,7 +611,7 @@ def _data_allocated(self): """ self._is_halo_dirty = True self._halo_exchange() - return np.asarray(self._data) + return self._data._raw def _data_buffer(self, **kwargs): """ @@ -650,23 +653,21 @@ def _data_in_region(self, region, dim, side): for d, s, (pl, pr) in zip(self.dimensions, self.shape_allocated, self._padding, strict=True) ] - return np.asarray(self._data[index_array]) + return self._data[index_array] @property @_allocate_memory def data_ro_domain(self): """Read-only view of the domain data values.""" view = self._data._global(self._mask_domain, self._decomposition) - view.setflags(write=False) - return view + return view._readonly() @property @_allocate_memory def data_ro_with_halo(self): """Read-only view of the domain+outhalo data values.""" view = self._data._global(self._mask_outhalo, self._decomposition_outhalo) - view.setflags(write=False) - return view + return view._readonly() _data_ro_with_outhalo = data_ro_with_halo @@ -681,8 +682,7 @@ def _data_ro_with_inhalo(self): This accessor does *not* support global indexing. """ view = self._data[self._mask_inhalo] - view.setflags(write=False) - return np.asarray(view) + return view._readonly()._raw @property @_allocate_memory @@ -695,8 +695,7 @@ def _data_ro_allocated(self): This accessor does *not* support global indexing. """ view = self._data - view.setflags(write=False) - return np.asarray(view) + return view._readonly()._raw @cached_property def local_indices(self): @@ -718,8 +717,8 @@ def local_indices(self): @property def initializer(self): - if isinstance(self._data, np.ndarray): - return self.data_with_halo.view(np.ndarray) + if self._data is not None: + return self.data_with_halo._raw else: return self._initializer @@ -751,12 +750,12 @@ def _C_make_dataobj(self, alias=None, **args): an Operator. """ key = alias or self - data = args[key.name] + data = Data._from_array(args[key.name], allocator=self._allocator) dataobj = byref(self._C_ctype._type_()) - dataobj._obj.data = data.ctypes.data_as(c_restrict_void_p) + dataobj._obj.data = data._C_data dataobj._obj.size = (c_int*self.ndim)(*data.shape) - dataobj._obj.nbytes = data.nbytes + dataobj._obj.nbytes = int(np.prod(data.shape))*data._numpy_dtype.itemsize # MPI-related fields dataobj._obj.npsize = (c_ulong*self.ndim)(*[ @@ -778,11 +777,8 @@ def _C_make_dataobj(self, alias=None, **args): return dataobj def _C_as_ndarray(self, dataobj): - """Cast the data carried by a DiscreteFunction dataobj to an ndarray.""" - shape = tuple(dataobj._obj.size[i] for i in range(self.ndim)) - ctype_1d = dtype_to_ctype(self.dtype) * int(reduce(mul, shape)) - buf = cast(dataobj._obj.data, POINTER(ctype_1d)).contents - return np.frombuffer(buf, dtype=self.dtype).reshape(shape) + """Return the array retained by a DiscreteFunction dataobj.""" + return dataobj._obj.underlying_array @memoized_meth def _C_make_index(self, dim, side=None): @@ -856,7 +852,7 @@ def _halo_exchange(self): # Gather send data data = self._data_in_region(OWNED, d, i) - sendbuf = np.ascontiguousarray(data.view(comm_dtype)) + sendbuf = np.ascontiguousarray(data._host_array().view(comm_dtype)) # Setup recv buffer shape = self._data_in_region(HALO, d, i.flip()).shape @@ -943,12 +939,21 @@ def _arg_check(self, args, intervals, **kwargs): if self.name not in args: raise InvalidArgument(f"No runtime value for `{self.name}`") - data = args[self.name] + data = Data._from_array(args[self.name], allocator=self._allocator) + args[self.name] = data + data._arg_check(self.dtype) + if self._allocator.is_Device: + if not isinstance(args.platform, Device) or args.language != 'openacc': + raise InvalidArgument("Device-resident storage requires OpenACC") + deviceid = self._allocator.device.index + if args.get('deviceid', -1) not in (-1, deviceid): + raise InvalidArgument("OpenACC deviceid must match tensor storage") + args['deviceid'] = deviceid if len(data.shape) != self.ndim: raise InvalidArgument(f"Shape {data.shape} of runtime value `{self.name}` " f"does not match dimensions {self.dimensions}") - if data.dtype != self.dtype: + if data._numpy_dtype != self.dtype: warning(f"Data type {data.dtype} of runtime value `{self.name}` " f"does not match the Function data type {self.dtype}") @@ -960,6 +965,9 @@ def _arg_finalize(self, args, alias=None): key = alias or self return {key.name: self._C_make_dataobj(alias=key, **args)} + def _arg_apply(self, dataobj, alias=None, access_mode=None): + dataobj._obj.underlying_array._arg_apply(access_mode) + # Pickling support @property @@ -1350,8 +1358,9 @@ def _arg_check(self, args, intervals, **kwargs): if args.options['index-mode'] == 'int32' and \ args.options['linearize'] and \ - data.size - 1 >= np.iinfo(np.int32).max: - raise InvalidArgument(f"`{self.name}`, with its {data.size} elements, is too " + int(np.prod(data.shape)) - 1 >= np.iinfo(np.int32).max: + raise InvalidArgument(f"`{self.name}`, with its " + f"{int(np.prod(data.shape))} elements, is too " "big for int32 pointer arithmetic. Consider using the " "'index-mode=int64' option, the save=Buffer(..) " "API (TimeFunction only), or domain " diff --git a/devito/types/grid.py b/devito/types/grid.py index 4568b852bf1..f4ed50faded 100644 --- a/devito/types/grid.py +++ b/devito/types/grid.py @@ -1131,16 +1131,16 @@ class SubDomainSet(MultiSubDomain): >>> op = Operator(eq) >>> summary = op.apply() >>> f.data - Data([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - [0, 1, 1, 1, 1, 1, 1, 1, 1, 0], - [0, 1, 1, 1, 1, 1, 1, 1, 1, 0], - [0, 1, 1, 1, 1, 1, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - [0, 1, 1, 1, 1, 1, 1, 1, 1, 0], - [0, 1, 1, 1, 1, 1, 1, 1, 1, 0], - [0, 1, 1, 1, 1, 1, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], dtype=int32) + array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [0, 1, 1, 1, 1, 1, 1, 1, 1, 0], + [0, 1, 1, 1, 1, 1, 1, 1, 1, 0], + [0, 1, 1, 1, 1, 1, 1, 1, 1, 0], + [0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [0, 1, 1, 1, 1, 1, 1, 1, 1, 0], + [0, 1, 1, 1, 1, 1, 1, 1, 1, 0], + [0, 1, 1, 1, 1, 1, 1, 1, 1, 0], + [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], dtype=int32) """ def __init__(self, **kwargs): diff --git a/devito/types/sparse.py b/devito/types/sparse.py index 6e913646a36..1839f64b881 100644 --- a/devito/types/sparse.py +++ b/devito/types/sparse.py @@ -783,7 +783,8 @@ def _arg_values(self, estimate_memory=False, **kwargs): return values - def _arg_apply(self, dataobj, alias=None): + def _arg_apply(self, dataobj, alias=None, access_mode=None): + super()._arg_apply(dataobj, alias=alias, access_mode=access_mode) key = alias if alias is not None else self if isinstance(key, AbstractSparseFunction): # Gather into `self.data` @@ -2315,7 +2316,8 @@ def _dist_scatter(self, alias=None, data=None): # The implementation in AbstractSparseFunction now relies on us # having a .coordinates property, which we don't have. - def _arg_apply(self, dataobj, *subfuncs, alias=None): + def _arg_apply(self, dataobj, *subfuncs, alias=None, access_mode=None): + DiscreteFunction._arg_apply(self, dataobj, alias=alias, access_mode=access_mode) key = alias if alias is not None else self if isinstance(key, AbstractSparseFunction): # Gather into `self.data` diff --git a/docs/torch.md b/docs/torch.md new file mode 100644 index 00000000000..486cbd3e1c1 --- /dev/null +++ b/docs/torch.md @@ -0,0 +1,213 @@ +# PyTorch allocation and code generation + +Use the standard Devito API with `TorchAllocator` and `language='torch'`: + +```python +from devito import Eq, Grid, Operator, TimeFunction, TorchAllocator + +grid = Grid(shape=(64, 64)) +u = TimeFunction(name='u', grid=grid, space_order=2, allocator=TorchAllocator()) +u.data.fill_(1) +op = Operator(Eq(u.forward, u + 1), language='torch', compiler='gcc') +op.apply(time_M=9) +``` + +`u` is an ordinary Devito `TimeFunction`, and `op` is selected through the usual +operator registry. The allocator chooses PyTorch-owned storage; `.data` and +`.data_with_halo` expose tensor views inheriting Devito's `Data` behavior. +NumPy remains a dependency for +compiler metadata, not numerical tensor storage. + +The torch language uses the existing CPU optimization pipeline and a `CGen` +subclass. It emits non-owning `at::Tensor` views over the argument storage and +obtains typed pointers with `data_ptr()`. Derivative expansion, blocking, +SIMD, and optimized stencil loops remain Devito's. Set +`opt=('advanced', {'openmp': True})` for OpenMP parallel loops. + +Compilation and execution use Devito's standard C++ compiler configuration, +ctypes ABI, JIT cache, argument checks, and profiler. There is no Python operator +wrapper, PyBind module, or separate PyTorch extension build. The printer uses +flat indexing because C variable-length arrays are not part of C++. + +## Installation + +Install PyTorch from the appropriate official CPU or CUDA wheel index, then: + +```sh +python -m pip install -e '.[torch,tests]' +``` + +The optional extra uses the current main branch's dependency ranges. CPU code +generation requires a C++17 compiler. The compiler standard is selected at construction through its +`std` property. The Torch target declares headers, libraries and ABI settings +as IET pass metadata. The existing graph machinery applies those +requirements to the operator-local compiler. Keyword normalization does not +configure the compiler, and the configured template is unchanged. Ninja is not +needed. + +## Existing tensors + +Use the allocator in the same way as `ExternalAllocator`. A supplied tensor must +include the Function's halo and padding. A no-op initializer preserves its values: + +```python +import torch +from devito import Function + +tensor = torch.ones((64, 64)) +f = Function(name='f', grid=grid, space_order=0, + allocator=TorchAllocator(tensor=tensor), initializer=lambda _: None) +Operator(Eq(f, f + 1), language='torch', compiler='gcc').apply() +# tensor now contains 2, with the same allocation. +``` + +A contiguous tensor can also be supplied through the normal runtime override +`op.apply(f=tensor)`. The usual Devito rules for overrides, dimension bounds, +halo layout, and grid spacing apply. Storage must match the Function's dtype, +device, and alignment. No implicit contiguous conversion or device copy occurs. +Supported dtypes are float32, float64, int32, and int64, specified through the +normal Devito `dtype` argument. Tensor initializers are also accepted. + +Storage and code generation remain independent: CPU TorchAllocator allocations +can be used with the existing C language, and the torch printer can operate on +ordinary host allocations. + +## OpenACC and CUDA tensors + +Select the existing OpenACC language and a CUDA allocator: + +```python +u = TimeFunction(name='u', grid=grid, space_order=2, + allocator=TorchAllocator(device='cuda:0')) +op = Operator(Eq(u.forward, u + 1), platform='nvidiaX', + language='openacc', compiler='nvc') +op.apply(time_M=9) +``` + +The allocator marks storage as device-resident. The existing OpenACC passes +emit `deviceptr(...)` clauses and omit host/device mapping and transfer +operations for those Functions. Tensor pointers pass through the normal Devito +data-object ABI. PyTorch owns the allocations throughout their lifetime. + +The allocator checks the target and device, supplies a matching default +`deviceid`, and synchronizes CUDA before and after execution so PyTorch and +OpenACC streams can safely consume each other's results. Written tensors have +their PyTorch version counters updated. CUDA Functions cannot be accessed by +host fallback code. + +Runtime validation requires CUDA-enabled PyTorch, an NVIDIA GPU, and the NVIDIA +HPC SDK OpenACC compiler. GPU code-generation tests can run without these; +the GPU execution test is skipped when they are unavailable. + +## Current limits and tests + +`Data` is the backend-independent base for `NumpyData(Data, np.ndarray)` and +`TorchData(Data, torch.Tensor)`. It owns modulo indexing, decomposition metadata, +global/local index translation, distributed assignment, and gather. Constructing +`Data` directly still selects NumPy by default, preserving the existing API. +Torch views and arithmetic results retain this metadata, so buffered time indices +wrap and distributed views retain Devito's indexing semantics. Private allocation +and inhalo accessors return native arrays, as in the NumPy API, without Devito's +modulo or global indexing. + +Both backends use the same MPI indexing and gather implementation. Data owns +native storage access and communication views; the allocator owns allocations. +CPU tensors expose zero-copy NumPy views to MPI through PyTorch’s standard +`.numpy()` interoperability. CUDA communication uses explicit host staging and +synchronization; +CUDA-aware MPI is not assumed. OpenACC kernels continue to consume GPU tensor +pointers directly, without staging. Install `.[mpi]` and an MPI runtime to use +distributed execution. + +Autograd tensors are not supported. PyTorch has no read-only views, so +`data_ro_*` returns independent snapshots. Negative-stride tensor reads produce +copies because PyTorch cannot represent negative-stride views; direct assignment +through negative slices writes back to the original allocation. Combining +negative strides with advanced tensor indexing is not currently supported. + +For an explicit differentiable integration, see +[the diffusion layer example](../examples/pytorch/devito_layer.py). +That CPU-only `nn.Module` +supplies a verified first-order input adjoint for a fixed linear diffusion +operator. Its native kernels receive detached workspaces; the autograd bridge +does not make arbitrary Operators or physical coefficients differentiable. + +```sh +OMP_NUM_THREADS=2 python -m pytest -q \ + tests/test_torch_allocators.py tests/test_torch_codegen.py +OMP_NUM_THREADS=2 python -m pytest -q tests/test_data_backends.py +``` + +Tests cover the standard allocator and operator API, native JIT execution, +numerical parity for blocked 3D stencils, OpenMP, external tensors, rejected +runtime layouts, shared indexing/decomposition semantics, MPI halo exchange and +gather, and OpenACC device-pointer generation. The shared-backend MPI tests use +the repository's existing multi-process pytest runner. In managed environments +without the CMA syscall, MPICH can use `MPIR_CVAR_CH4_CMA_ENABLE=0 UCX_TLS=self,tcp`. +A skipped CUDA execution test does not establish GPU runtime compatibility. + +## Internal organization + +The backend follows the existing module and argument-provider structure: + +| Module | Torch support | +| --- | --- | +| `devito/data/allocators.py` | `TorchAllocator`, alongside the other allocators; allocation and external ownership | +| `devito/data/data.py` | `TorchData`, alongside `Data` and `NumpyData`; native views and the shared indexing/MPI implementation | +| `devito/tools/dtypes_lowering.py` | Scalar dtype conversion | +| `devito/types/dense.py`, `devito/types/sparse.py` | The existing `_arg_check`, `_arg_finalize` and `_arg_apply` lifecycle | +| `devito/core/torch.py`, `devito/ir/iet/torch.py` | Language selection and C++ printing through the existing optimization pipeline | +| `devito/passes/iet/languages/torch.py` | Runtime headers and JIT requirements declared as IET metadata | + +Allocators return their native array type. `Data` creates the corresponding view +through its shared construction path. Function argument validation wraps native +overrides in non-owning Data views, and the usual data-object ABI retains those +views for the duration of the call. Tensor write synchronization and version +updates run through `_arg_apply`. + +There is no separate storage runtime or backend package. The public API remains +`from devito import TorchAllocator`, tensor `.data`, and +`Operator(..., language='torch')`. + +Tests cover allocation in `test_torch_allocators.py`, C++ generation in +`test_torch_codegen.py`, shared indexing/MPI in `test_data_backends.py`, OpenACC +in `test_torch_openacc.py`, and the example layer in `test_torch_layer.py`. + +## Contribution validation + +Torch tests run in the existing CPU and MPI workflows, which collect `tests/`, +and in the NVIDIA OpenACC job of the existing GPU workflow. PyTorch is installed +through the standard test dependencies, including in the shared Docker image; +there is no separate Torch workflow or workflow-specific dependency installation. +The tests use the same dependency requirements as the main branch. + +```sh +# All CPU tests, including the new torch and OpenACC-host cases +OMP_NUM_THREADS=2 pytest -q -m 'not parallel' --maxfail=0 tests/ + +# All MPI tests; transport flags are needed only where CMA is unavailable +MPIR_CVAR_CH4_CMA_ENABLE=0 UCX_TLS=self,tcp MPIEXEC_TIMEOUT=900 \ + OMP_NUM_THREADS=1 pytest -q -m parallel --maxfail=0 tests/ + +# Actual GCC OpenACC compilation and host execution, without a GPU +OMP_NUM_THREADS=2 pytest -q tests/test_torch_openacc.py -k host + +# NVIDIA OpenACC interoperability; requires CUDA PyTorch, GPU and nvc++ +python -c 'import torch; assert torch.cuda.is_available()' +nvc++ --version +pytest -q tests/test_torch_openacc.py -k cuda +``` + +The OpenACC host cases execute generated code on GCC's `acc_device_host` target +and compare stencil results with PyTorch. They launch a fresh interpreter with +the compiler's system `libgomp` loaded before PyTorch's bundled runtime, which +may lack newer OpenACC symbols. This is a test configuration, not a CUDA-device +emulation or a substitute for the NVIDIA cases. + +CUDA coverage includes float32/float64 stencils, allocator-owned and external +storage, runtime overrides, nondefault producer streams, mixed NumPy/tensor +inputs, and device selection. A second-device case requires two GPUs. The +workflow's CUDA job is opt-in through `workflow_dispatch` and requires a +configured self-hosted runner labelled `nvidiagpu`; it fails its preflight if +CUDA or the NVIDIA compiler is unavailable, rather than reporting an all-skipped +CUDA run as successful. diff --git a/examples/pytorch/__init__.py b/examples/pytorch/__init__.py new file mode 100644 index 00000000000..d32a758f6d8 --- /dev/null +++ b/examples/pytorch/__init__.py @@ -0,0 +1 @@ +"""PyTorch integration examples for the optional torch storage backend.""" diff --git a/examples/pytorch/devito_layer.py b/examples/pytorch/devito_layer.py new file mode 100644 index 00000000000..950819a7ae2 --- /dev/null +++ b/examples/pytorch/devito_layer.py @@ -0,0 +1,154 @@ +"""A differentiable fixed diffusion layer using an explicit discrete adjoint. + +Run from the repository root with: + OMP_NUM_THREADS=2 python -m examples.pytorch.devito_layer + +This CPU, eager-mode example implements first-order input gradients for a +constant-coefficient, uniform-grid, zero-ghost diffusion stencil. It does not +make arbitrary Devito Operators differentiable. +""" + +from functools import lru_cache +from itertools import count +from math import isfinite +from threading import Lock + +import torch +from torch import nn +from torch.autograd.function import once_differentiable + +from devito import ( + Dimension, Eq, Grid, Operator, TimeFunction, TorchAllocator, configuration +) +from devito.tools import dtype_from_torch + +_names = count() + + +class _DiffusionPlan: + + def __init__(self, planes, height, width, dtype, steps, coefficient): + self.grid = Grid(shape=(height, width), extent=(height-1., width-1.)) + self.plane = Dimension(name='plane') + self.shape = (2, planes, height, width) + self.dtype = dtype_from_torch(dtype).type + self.steps = steps + self.lock = Lock() + self.symbol = self._state() + self.operator = Operator( + Eq(self.symbol.forward, self.symbol + coefficient*self.symbol.laplace), + language='torch', compiler='gcc', platform='cpu64', + opt=('advanced', {'openmp': True}) + ) + + def _state(self): + return TimeFunction( + name=f'diffusion_{next(_names)}', grid=self.grid, + dimensions=(self.grid.stepping_dim, self.plane, *self.grid.dimensions), + shape=self.shape, time_order=1, space_order=2, dtype=self.dtype, + allocator=TorchAllocator() + ) + + def run(self, value): + # Cache the compiled operator, never mutable activations. A fresh + # workspace makes two forwards before either backward safe. + with self.lock: + state = self._state() + state.data[0].copy_(value.detach().reshape(self.shape[1:])) + self.operator.apply(**{self.symbol.name: state}, time_M=self.steps-1, + autotune='off') + output = torch.empty(value.shape, dtype=value.dtype, device=value.device) + output.copy_(state.data[self.steps].reshape(value.shape)) + return output + + +@lru_cache(maxsize=16) +def _plan(planes, height, width, dtype, steps, coefficient): + return _DiffusionPlan(planes, height, width, dtype, steps, coefficient) + + +class _DiffusionFunction(torch.autograd.Function): + + @staticmethod + def forward(ctx, value, plan): + ctx.plan = plan + return plan.run(value) + + @staticmethod + @once_differentiable + def backward(ctx, grad_output): + # A = I + coefficient*L is symmetric for this exact five-point + # discretization and zero-ghost boundary treatment. Thus (A^steps)^T + # is A^steps too. This is an explicit adjoint, not autodiff of C++. + return ctx.plan.run(grad_output), None + + +class DevitoDiffusion(nn.Module): + """Fixed diffusion on CPU float32/float64 tensors shaped [N, C, H, W]. + + `coefficient` is alpha*dt on a unit-spaced 2D grid and lies in [0, 1/4]. + Batch and channel planes evolve independently. First-order gradients pass + through to the input; steps/coefficient are fixed Python configuration. + CUDA, distributed execution and higher-order gradients are not supported + by this example layer. + """ + + def __init__(self, steps=3, coefficient=0.1): + super().__init__() + if isinstance(steps, bool) or not isinstance(steps, int) or steps < 1: + raise ValueError('steps must be a positive integer') + if isinstance(coefficient, torch.Tensor): + raise TypeError('coefficient is fixed configuration, not a tensor parameter') + if not isfinite(coefficient) or not 0 <= coefficient <= 0.25: + raise ValueError('coefficient must lie in [0, 0.25]') + self.steps = steps + self.coefficient = float(coefficient) + + def forward(self, value): + if configuration['mpi']: + raise NotImplementedError('This example layer does not support MPI') + if value.device.type != 'cpu': + raise ValueError('This example layer supports CPU tensors only') + if value.dtype not in (torch.float32, torch.float64): + raise TypeError('Expected float32 or float64 tensor storage') + if value.ndim != 4 or min(value.shape) < 1 or min(value.shape[-2:]) < 2: + raise ValueError('Expected [N, C, H, W] with H and W at least 2') + batch, channels, height, width = value.shape + plan = _plan(batch*channels, height, width, value.dtype, + self.steps, self.coefficient) + return _DiffusionFunction.apply(value, plan) + + def extra_repr(self): + return f'steps={self.steps}, coefficient={self.coefficient:g}' + + +def train_demo(iterations=40): + """Train learned convolutions on both sides of a Devito physics layer.""" + torch.manual_seed(7) + inputs = torch.randn(8, 1, 16, 16) + target = nn.functional.avg_pool2d(inputs, kernel_size=3, stride=1, padding=1) + network = nn.Sequential( + nn.Conv2d(1, 4, kernel_size=3, padding=1), + nn.Tanh(), + DevitoDiffusion(steps=2, coefficient=0.1), + nn.Conv2d(4, 1, kernel_size=1), + ) + optimizer = torch.optim.Adam(network.parameters(), lr=0.03) + losses = [] + for _ in range(iterations): + optimizer.zero_grad(set_to_none=True) + loss = nn.functional.mse_loss(network(inputs), target) + loss.backward() + if not torch.isfinite(network[0].weight.grad).all(): + raise RuntimeError('Nonfinite upstream gradient') + losses.append(loss.item()) + optimizer.step() + with torch.no_grad(): + losses.append(nn.functional.mse_loss(network(inputs), target).item()) + return network, losses + + +if __name__ == '__main__': + model, history = train_demo() + print(model) + print(f'Training MSE: {history[0]:.6f} -> {history[-1]:.6f}') diff --git a/pyproject.toml b/pyproject.toml index 7fa3f0c16cd..c9549c81b44 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,6 +59,7 @@ dependencies = { file = ["requirements.txt"] } [tool.setuptools.dynamic.optional-dependencies] # Map each extra to its requirements-*.txt +torch = { file = ["requirements-torch.txt"] } mpi = { file = ["requirements-mpi.txt"] } nvidia = { file = ["requirements-nvidia.txt"] } tests = { file = ["requirements-testing.txt"] } diff --git a/requirements-testing.txt b/requirements-testing.txt index 58fc63050a8..a20b7bfc0ed 100644 --- a/requirements-testing.txt +++ b/requirements-testing.txt @@ -10,3 +10,4 @@ cloudpickle<3.1.3 ipympl<0.10.1 ipykernel<7.0.0 pytest-xdist<3.8.1 +torch>=2.2 diff --git a/requirements-torch.txt b/requirements-torch.txt new file mode 100644 index 00000000000..657e282eebe --- /dev/null +++ b/requirements-torch.txt @@ -0,0 +1 @@ +torch>=2.2 diff --git a/tests/test_data_backends.py b/tests/test_data_backends.py new file mode 100644 index 00000000000..f276fcfb551 --- /dev/null +++ b/tests/test_data_backends.py @@ -0,0 +1,198 @@ +"""Shared Data behavior, tested against both numerical storage backends.""" + +import numpy as np +import pytest + +torch = pytest.importorskip('torch') + +from devito import ( # noqa: E402 + Eq, Function, Grid, Operator, TimeFunction, TorchAllocator +) +from devito.data import Data, Decomposition, NumpyData, TorchData # noqa: E402 + + +@pytest.fixture(params=['numpy', 'torch']) +def allocator(request): + return TorchAllocator() if request.param == 'torch' else None + + +def native(data): + if isinstance(data, Data): + data = data._raw + return data.numpy() if isinstance(data, torch.Tensor) else np.asarray(data) + + +def test_common_base(): + numpy = Data((4,), np.float32, modulo=(False,)) + tensor = Data((4,), np.float32, allocator=TorchAllocator()) + assert isinstance(numpy, (Data, NumpyData, np.ndarray)) + assert isinstance(tensor, Data) + assert isinstance(tensor, TorchData) + assert isinstance(tensor, torch.Tensor) + assert not issubclass(Data, np.ndarray) + + +@pytest.mark.parametrize('allocator', [None, TorchAllocator()]) +def test_host_view_shares_storage(allocator): + f = Function(name='host_view', grid=Grid(shape=(8, 8)), space_order=0, + allocator=allocator) + host = f.data._host_array() + assert isinstance(host, np.ndarray) + host[2, 3] = 7 + assert f.data[2, 3] == 7 + f.data[4, 5] = 9 + assert host[4, 5] == 9 + + +def test_modulo_and_views(allocator): + grid = Grid(shape=(8, 8)) + u = TimeFunction(name='u', grid=grid, allocator=allocator) + u.data[0] = 2 + u.data[1] = 3 + np.testing.assert_equal(native(u.data[4]), np.full((8, 8), 2)) + np.testing.assert_equal(native(u.data[-3]), np.full((8, 8), 3)) + u.data[5, 2, 3] = 7 + assert u.data[1, 2, 3] == 7 + view = u.data[:, 1:5, 2:6] + assert isinstance(view, Data) + assert view._modulo == (True, False, False) + assert view._distributor is u.grid.distributor + view[4, 0, 0] = 9 + assert u.data[0, 1, 2] == 9 + projected = view[1] + assert projected._modulo == (False, False) + np.testing.assert_equal(native((view+1)[2]), native(view[0])+1) + if isinstance(view, torch.Tensor): + result = torch.add(input=view, other=1) + np.testing.assert_equal(native(result[2]), native(view[0])+1) + + +def test_indexing_assignment_and_broadcasting(allocator): + f = Function(name='f', grid=Grid(shape=(8, 8)), space_order=0, + allocator=allocator) + expected = np.arange(64, dtype=np.float32).reshape(8, 8) + f.data[:] = expected + for index in [(slice(None, None, -1), slice(None, None, -2)), + (Ellipsis, 3), ([1, 4, 6], slice(2, 5)), + (np.array([True, False]*4), slice(None))]: + np.testing.assert_equal(native(f.data[index]), expected[index]) + f.data[6:1:-2, 7:2:-2] = np.arange(9, dtype=np.float32).reshape(3, 3) + expected[6:1:-2, 7:2:-2] = np.arange(9).reshape(3, 3) + np.testing.assert_equal(native(f.data), expected) + f.data[:] = np.arange(8, dtype=np.float32) + np.testing.assert_equal(native(f.data), np.broadcast_to(np.arange(8), (8, 8))) + + +def test_decomposition_and_nonlocal_access(allocator): + decomposition = Decomposition([range(4), range(4, 8)], 1) + data = Data((4,), np.float32, decomposition=(decomposition,), + modulo=(False,), **({'allocator': allocator} if allocator else {})) + data._local[:] = np.arange(4, 8, dtype=np.float32) + assert data[2] is None + assert data[5] == 5 + np.testing.assert_equal(native(data._local), np.arange(4, 8)) + view = data[5:7] + assert isinstance(view, Data) + assert view._decomposition[0].loc_abs_min == 0 + assert view[0] == 5 + view[1] = 12 + assert data[6] == 12 + assert native(data[:2]).size == 0 + + +def test_buffered_reverse_assignment_and_gather(allocator): + u = TimeFunction(name='u', grid=Grid(shape=(8, 8)), allocator=allocator) + u.data[0, ::-1, 0] = np.arange(8, dtype=np.float32) + np.testing.assert_equal(native(u.data[0, :, 0]), np.arange(7, -1, -1)) + gathered = u.data_gather() + np.testing.assert_equal(native(gathered), native(u.data)) + + +@pytest.mark.parallel(mode=4) +def test_mpi_indexing_halo_and_gather(allocator, mode): + grid = Grid(shape=(8, 8), extent=(7., 7.)) + f = Function(name='f', grid=grid, space_order=1, allocator=allocator) + values = np.arange(64, dtype=np.float32).reshape(8, 8) + f.data[:] = values + owned = grid.distributor.glb_slices + selection = tuple(owned[d] for d in grid.dimensions) + np.testing.assert_equal(native(f.data._local), values[selection]) + for x, y in [(0, 0), (4, 4), (7, 7)]: + actual = f.data[x, y] + if selection[0].start <= x < selection[0].stop and \ + selection[1].start <= y < selection[1].stop: + assert actual == values[x, y] + else: + assert actual is None + # Existing Function halo orchestration, using backend-specific MPI buffers. + halo = native(f._data_with_inhalo) + padded = np.pad(values, 1) + expected_halo = padded[selection[0].start:selection[0].stop+2, + selection[1].start:selection[1].stop+2] + np.testing.assert_equal(halo, expected_halo) + gathered = f.data_gather(rank=2) + if grid.distributor.myrank == 2: + np.testing.assert_equal(native(gathered), values) + else: + assert gathered is None + reversed_data = f.data[::-1, ::-1] + np.testing.assert_equal(native(reversed_data._local), + values[selection][::-1, ::-1]) + f.data[:] = reversed_data + gathered = f.data_gather() + if grid.distributor.myrank == 0: + np.testing.assert_equal(native(gathered), values[::-1, ::-1]) + # Replicated values follow the same global negative-slice assignment semantics. + f.data[::-1, ::-1] = values + gathered = f.data_gather() + if grid.distributor.myrank == 0: + np.testing.assert_equal(native(gathered), values[::-1, ::-1]) + + +@pytest.mark.parallel(mode=2) +def test_mpi_tensor_kernel(allocator, mode): + grid = Grid(shape=(16, 8), extent=(15., 7.)) + u = TimeFunction(name='u', grid=grid, space_order=2, allocator=allocator) + u.data[0] = 1 + language = 'torch' if allocator else 'C' + op = Operator(Eq(u.forward, u + 0.1*u.laplace), language=language, compiler='gcc') + op.apply(time_M=0) + # MPI-generated stencil execution must agree with the same global reference. + gathered = u.data_gather() + if grid.distributor.myrank == 0: + expected = np.ones((16, 8), dtype=np.float32) + expected[0, :] -= 0.1 + expected[-1, :] -= 0.1 + expected[:, 0] -= 0.1 + expected[:, -1] -= 0.1 + np.testing.assert_allclose(native(gathered)[1], expected, atol=1e-6) + + +def test_private_data_accessors_use_native_indexing(allocator): + u = TimeFunction(name='native_access', grid=Grid(shape=(8, 8)), + space_order=0, padding=0, allocator=allocator) + # Public time buffers wrap, while raw storage follows native array bounds. + u.data[2] = 3 + for data in (u._data_allocated, u._data_with_inhalo, + u._data_ro_allocated, u._data_ro_with_inhalo): + assert not isinstance(data, Data) + with pytest.raises(IndexError): + data[2] + np.testing.assert_equal(native(data[0]), + np.full((8, 8), 3)) + + +@pytest.mark.parametrize('index', [ + (slice(0, 1, -1), slice(None)), + (slice(3, 3, -1), slice(None)), + (slice(2, 1), slice(None, None, -1)), + (slice(None, None, -1), slice(0, 1, -1)) +]) +def test_empty_reverse_slices(allocator, index): + f = Function(name='empty_reverse', grid=Grid(shape=(8, 8)), space_order=0, + allocator=allocator) + expected = np.arange(64, dtype=np.float32).reshape(8, 8) + f.data[:] = expected + np.testing.assert_equal(native(f.data[index]), expected[index]) + f.data[index] = 9 + np.testing.assert_equal(native(f.data), expected) diff --git a/tests/test_torch_allocators.py b/tests/test_torch_allocators.py new file mode 100644 index 00000000000..ed7ada53db5 --- /dev/null +++ b/tests/test_torch_allocators.py @@ -0,0 +1,96 @@ +"""Tensor allocation, external ownership and native argument validation.""" + +import numpy as np +import pytest + +torch = pytest.importorskip('torch') + +from devito import Eq, Function, Grid, Operator, TorchAllocator # noqa: E402 +from devito.exceptions import InvalidArgument # noqa: E402 + + +@pytest.mark.parametrize('dtype', [np.float32, np.float64, np.int32, np.int64]) +def test_allocator(dtype): + f = Function(name='f', grid=Grid(shape=(8, 8)), dtype=dtype, + allocator=TorchAllocator()) + assert type(f) is Function + assert isinstance(f.data, torch.Tensor) + f.data.fill_(3) + assert torch.all(f.data == 3) + data = f._data_allocated + obj = f._C_make_dataobj(**{f.name: data}) + assert obj._obj.data.value == data.data_ptr() + assert obj._obj.underlying_array.data_ptr() == data.data_ptr() + assert f._C_as_ndarray(obj).data_ptr() == data.data_ptr() + snapshot = f.data_ro_domain + snapshot.zero_() + assert torch.all(f.data == 3) + + +def test_external_tensor(monkeypatch): + tensor = torch.arange(64, dtype=torch.float32).reshape(8, 8) + f = Function(name='f', grid=Grid(shape=(8, 8)), space_order=0, + allocator=TorchAllocator(tensor=tensor), initializer=lambda _: None) + assert f.data.data_ptr() == tensor.data_ptr() + op = Operator(Eq(f, 2*f+1), language='torch', compiler='gcc') + expected = 2*tensor + 1 + + def no_numpy(*args, **kwargs): + raise AssertionError('Tensor data must not pass through NumPy') + + with monkeypatch.context() as patch: + patch.setattr(torch.Tensor, 'numpy', no_numpy) + op.apply() + torch.testing.assert_close(tensor, expected) + replacement = torch.zeros_like(tensor) + op.apply(f=replacement) + torch.testing.assert_close(replacement, torch.ones_like(tensor)) + torch.testing.assert_close(tensor, expected) + + +@pytest.mark.parametrize('initializer', [ + torch.ones((8, 8)), np.ones((8, 8)), [[1]*8]*8 +]) +def test_initializers(initializer): + f = Function(name='initialized', grid=Grid(shape=(8, 8)), space_order=0, + allocator=TorchAllocator(), initializer=initializer) + assert torch.all(f.data == 1) + + +@pytest.mark.parametrize('kind', [ + 'dtype', 'strides', 'grad', 'numpy', 'alignment', 'sparse' +]) +def test_invalid_override(kind): + f = Function(name='f', grid=Grid(shape=(8, 8)), space_order=0, + allocator=TorchAllocator()) + op = Operator(Eq(f, f+1), language='torch', compiler='gcc') + if kind == 'dtype': + value = torch.zeros((8, 8), dtype=torch.float64) + elif kind == 'strides': + value = torch.zeros((8, 8)).T + elif kind == 'sparse': + value = torch.zeros((8, 8)).to_sparse() + elif kind == 'grad': + value = torch.zeros((8, 8), requires_grad=True) + elif kind == 'alignment': + value = torch.zeros(65)[1:].reshape(8, 8) + else: + value = np.zeros((8, 8), dtype=np.float32) + with pytest.raises(InvalidArgument): + op.arguments(f=value) + + +def test_default_cuda_device(monkeypatch): + monkeypatch.setattr(torch.cuda, 'current_device', lambda: 2) + allocator = TorchAllocator(device='cuda') + assert allocator.device == torch.device('cuda:2') + # A supplied tensor determines the device unless explicitly overridden. + tensor = torch.empty((8, 8)) + assert TorchAllocator(tensor=tensor).device == tensor.device + + +def test_indexed_cpu_device(): + f = Function(name='cpu_index', grid=Grid(shape=(4, 4)), + allocator=TorchAllocator(device='cpu:0')) + f.data[:] = 2 + torch.testing.assert_close(f.data, torch.full((4, 4), 2.)) diff --git a/tests/test_torch_codegen.py b/tests/test_torch_codegen.py new file mode 100644 index 00000000000..abb2f397cc5 --- /dev/null +++ b/tests/test_torch_codegen.py @@ -0,0 +1,248 @@ +"""Torch language selection, optimized C++ printing and native JIT execution.""" + +import numpy as np +import pytest +from packaging.version import Version + +torch = pytest.importorskip('torch') + +from conftest import _R # noqa: E402 +from devito import ( # noqa: E402 + Constant, Dimension, Eq, Function, Grid, Operator, TimeFunction, TorchAllocator, + configuration, switchconfig +) +from devito.arch import Cpu64 # noqa: E402 +from devito.arch.compiler import ( # noqa: E402 + ClangCompiler, CustomCompiler, GNUCompiler, IntelCompiler, NvidiaCompiler, + OneapiCompiler +) +from devito.core.torch import Cpu64AdvTorchOperator # noqa: E402 +from devito.ir import Dereference, FindNodes # noqa: E402 +from devito.ir.iet.torch import TorchCGen # noqa: E402 +from devito.operator.operator import parse_kwargs # noqa: E402 + + +@pytest.mark.parametrize('opt', ['noop', 'advanced', 'advanced-fsg']) +def test_standard_operator(opt): + u = TimeFunction(name='u', grid=Grid(shape=(8, 8)), allocator=TorchAllocator()) + op = Operator(Eq(u.forward, u+1), language='torch', opt=opt, compiler='gcc') + assert isinstance(op, Operator) + assert op._ccode_handler is TorchCGen + if opt == 'advanced': + assert type(op) is Cpu64AdvTorchOperator + assert op._compiler.cc == 'g++' + assert op._compiler.src_ext == 'cpp' + assert 'at::Tensor' in str(op.ccode) + assert '.data_ptr()' in str(op.ccode) + assert 'for (int time' in str(op.ccode) + assert 'pybind11' not in str(op.ccode) + ptr, version = u._data_allocated.data_ptr(), u._data_allocated._version + op.apply(time_M=3) + torch.testing.assert_close(u.data[0], torch.full((8, 8), 4.)) + assert u._data_allocated.data_ptr() == ptr + assert u._data_allocated._version > version + # Standard ctypes loader and cache, not torch's extension mechanism. + assert op.cfunction is op._cfunction + assert op._lib._name.endswith('.so') + assert str(op._compiler.get_jit_dir()) in op._lib._name + + +@pytest.mark.parametrize('parallel,dtype', [(False, np.float32), (True, np.float64)]) +def test_cpp_stencil(parallel, dtype): + grid = Grid(shape=(12, 12, 12), extent=(11., 11., 11.)) + u = TimeFunction(name='u', grid=grid, dtype=dtype, space_order=2, + allocator=TorchAllocator()) + ref = TimeFunction(name='ref', grid=grid, dtype=dtype, space_order=2) + u.data[0, 6, 6, 6] = 1 + ref.data[0, 6, 6, 6] = 1 + scale = Constant(name='scale', value=0.01, dtype=dtype) + opt = ('advanced', {'blockinner': True, 'openmp': parallel}) + op = Operator(Eq(u.forward, u + scale*u.laplace), language='torch', + opt=opt, compiler='gcc') + reference = Operator(Eq(ref.forward, ref + scale*ref.laplace), + opt=opt, compiler='gcc') + op.apply(time_M=2) + reference.apply(time_M=2) + np.testing.assert_allclose(u.data.numpy(), ref.data, rtol=1e-5, atol=1e-6) + assert any(d.is_Block for d in op.dimensions) + assert '#pragma omp simd' in str(op.ccode) + assert ('#pragma omp parallel' in str(op.ccode)) == parallel + + +def test_allocator_with_c_language(): + # Storage choice is independent of language, as with other allocators. + f = Function(name='f', grid=Grid(shape=(8, 8)), space_order=0, + allocator=TorchAllocator()) + Operator(Eq(f, f+1), language='C', compiler='gcc').apply() + assert torch.all(f.data == 1) + + +@pytest.mark.parametrize('dtype', [np.int32, np.int64]) +def test_integer_codegen(dtype): + f = Function(name='integer', grid=Grid(shape=(8, 8)), space_order=0, + dtype=dtype, allocator=TorchAllocator()) + Operator(Eq(f, f+3), language='torch', compiler='gcc').apply() + assert torch.all(f.data == 3) + + +def test_custom_compiler(monkeypatch): + monkeypatch.setenv('CXX', 'g++') + monkeypatch.setenv('CXXFLAGS', '-O2 -fPIC') + f = Function(name='custom', grid=Grid(shape=(8, 8)), space_order=0, + allocator=TorchAllocator()) + op = Operator(Eq(f, f+1), language='torch', compiler='custom') + assert '-std=c++17' in op._compiler.cflags + assert '-std=c99' not in op._compiler.cflags + op.apply() + assert torch.all(f.data == 1) + + +def test_torch_language_with_host_allocator(): + f = Function(name='host', grid=Grid(shape=(8, 8)), space_order=0) + previous_ext = configuration['compiler'].src_ext + with switchconfig(language='torch'): + assert configuration['compiler'].src_ext == previous_ext + op = Operator(Eq(f, f+1)) + op.apply() + assert configuration['compiler'].src_ext == previous_ext + assert np.all(f.data == 1) + + +def test_preemptive_autotuning_preserves_tensor_state(monkeypatch): + u = TimeFunction(name='tuned', grid=Grid(shape=(12, 12, 12)), + allocator=TorchAllocator()) + op = Operator(Eq(u.forward, u.dx+1), language='torch', compiler='gcc', + opt=('blocking', {'openmp': False, 'blockinner': True})) + + def no_numpy(*args, **kwargs): + raise AssertionError('Autotuning must preserve tensor-owned storage') + + with monkeypatch.context() as patch: + patch.setattr(torch.Tensor, 'numpy', no_numpy) + op.apply(time_M=0, autotune=('basic', 'preemptive')) + + assert op._state['autotuning'][-1]['runs'] > 0 + torch.testing.assert_close(u.data[0], torch.zeros((12, 12, 12))) + torch.testing.assert_close(u.data[1], torch.ones((12, 12, 12))) + + +def test_function_names_do_not_collide_with_tensor_views(): + grid = Grid(shape=(8, 8)) + f = Function(name='field', grid=grid, space_order=0, allocator=TorchAllocator()) + g = Function(name='field_tensor', grid=grid, space_order=0, + allocator=TorchAllocator()) + f.data[:] = 1 + g.data[:] = 2 + op = Operator(Eq(f, f+g), language='torch', compiler='gcc') + op.apply() + torch.testing.assert_close(f.data, torch.full((8, 8), 3.)) + + +@pytest.mark.parametrize('compiler_type,flag', [ + (ClangCompiler, '-fopenmp'), (IntelCompiler, '-qopenmp'), + (OneapiCompiler, '-fopenmp') +]) +def test_cpu_compiler_enables_openmp(monkeypatch, compiler_type, flag): + # Inspect toolchain flags without requiring every vendor compiler locally. + monkeypatch.setattr('devito.arch.compiler.sniff_compiler_version', + lambda cc, **kwargs: Version('16.0')) + compiler = compiler_type(cpp=True, openmp=True, + platform=Cpu64('test', isa='cpp')) + assert flag in compiler.ldflags + assert compiler.src_ext == 'cpp' + + +def test_custom_compiler_honors_cxxflags(monkeypatch): + for name in ('CC', 'CXX', 'CFLAGS', 'LDFLAGS'): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv('CXXFLAGS', '-O2 -fPIC') + compiler = CustomCompiler(cpp=True) + assert '-O2' in compiler.cflags + assert compiler.cflags.index('-O2') > compiler.cflags.index('-O3') + + +@pytest.mark.parametrize('rotate', [False, True]) +def test_thread_private_temporaries(rotate): + grid = Grid(shape=(8, 8, 8)) + x, y, z = grid.dimensions + t = grid.stepping_dim + d = Dimension(name='d') + allocator = TorchAllocator() + c = Function(name='coefficient', grid=grid, shape=(2, 8), dimensions=(d, z), + allocator=allocator) + u = TimeFunction(name='temporary_state', grid=grid, space_order=3, + allocator=allocator) + ref = TimeFunction(name='temporary_reference', grid=grid, space_order=3, + allocator=allocator) + c.data_with_halo[:] = 1 + u.data_with_halo[:] = 1.5 + ref.data_with_halo[:] = 1.5 + expression = Eq(u.forward, + _R(_R(c[0, z]*u[t, x+1, y+1, z] + + c[1, z+1]*u[t, x+1, y+1, z+1]) + + _R(c[0, z]*u[t, x+2, y+2, z] + + c[1, z+1]*u[t, x+2, y+2, z+1]) + + _R(u[t, x, y+1, z+1] + 3*u[t, x+1, y+1, z+1]) + + _R(u[t, x, y+3, z+1] + 3*u[t, x+1, y+3, z+1]))) + reference = Operator(expression, language='C', opt='noop', compiler='gcc') + op = Operator(expression, language='torch', compiler='gcc', + opt=('advanced', {'openmp': True, 'min-storage': True, + 'cire-mingain': 0, 'cire-rotate': rotate})) + assert FindNodes(Dereference).visit(op) + reference.apply(temporary_state=ref, time_M=1) + op.apply(time_M=1) + torch.testing.assert_close(u.data, ref.data) + + +def test_torch_compiler_requirements_are_operator_local(): + template = configuration['compiler'] + previous = (template.cc, template.src_ext, list(template.cflags), + list(template.include_dirs), list(template.libraries)) + f = Function(name='local_toolchain', grid=Grid(shape=(8, 8)), space_order=0, + allocator=TorchAllocator()) + op = Operator(Eq(f, f+1), language='torch', compiler='gcc') + assert op._compiler.src_ext == 'cpp' + assert '-std=c++17' in op._compiler.cflags + assert 'torch_cpu' in op._compiler.libraries + assert previous == (template.cc, template.src_ext, template.cflags, + template.include_dirs, template.libraries) + op.apply() + torch.testing.assert_close(f.data, torch.ones((8, 8))) + + +def test_cpp_only_compiler_accepts_backend_options(monkeypatch): + monkeypatch.setattr('devito.arch.compiler.sniff_compiler_version', + lambda cc, **kwargs: Version('16.0')) + compiler = NvidiaCompiler(cpp=True, openmp=True, + platform=Cpu64('test', isa='cpp')) + assert compiler.src_ext == 'cpp' + assert '-mp' in compiler.cflags + + +def test_normalize_kwargs_leaves_compiler_unchanged(): + kwargs = parse_kwargs(language='torch', compiler='gcc') + compiler = kwargs['compiler'] + attrs = ('cflags', 'defines', 'include_dirs', 'libraries', 'library_dirs', 'ldflags') + previous = {name: list(getattr(compiler, name)) for name in attrs} + normalized = Cpu64AdvTorchOperator._normalize_kwargs(**kwargs) + assert normalized['compiler'] is compiler + assert previous == {name: getattr(compiler, name) for name in attrs} + + +@pytest.mark.parametrize('cpp,expected', [(False, 'c99'), (True, 'c++14')]) +def test_compiler_standard_property(cpp, expected): + compiler = GNUCompiler(cpp=cpp) + assert compiler.std == expected + assert f'-std={compiler.std}' in compiler.cflags + + +def test_compiler_standard_class_override_and_clone(): + class Cxx20Compiler(GNUCompiler): + _cxxstd = 'c++20' + + compiler = Cxx20Compiler(cpp=True) + clone = compiler.__new_with__() + assert clone is not compiler + assert compiler.std == clone.std == 'c++20' + assert '-std=c++20' in clone.cflags + assert clone.cc == compiler.cc diff --git a/tests/test_torch_layer.py b/tests/test_torch_layer.py new file mode 100644 index 00000000000..fa489a8961e --- /dev/null +++ b/tests/test_torch_layer.py @@ -0,0 +1,80 @@ +"""Validate the example layer's explicit discrete adjoint and network use.""" + +import pytest + +torch = pytest.importorskip('torch') + +from examples.pytorch.devito_layer import DevitoDiffusion, train_demo # noqa: E402 + + +def reference(value, steps, coefficient): + channels = value.shape[1] + kernel = value.new_tensor([[0, coefficient, 0], + [coefficient, 1-4*coefficient, coefficient], + [0, coefficient, 0]]) + weights = kernel.expand(channels, 1, 3, 3) + for _ in range(steps): + value = torch.nn.functional.conv2d(value, weights, padding=1, groups=channels) + return value + + +@pytest.mark.parametrize('dtype', [torch.float32, torch.float64]) +def test_forward_and_backward_match_pytorch(dtype): + torch.manual_seed(3) + value = torch.randn(2, 2, 5, 6, dtype=dtype, requires_grad=True) + expected_input = value.detach().clone().requires_grad_() + layer = DevitoDiffusion(steps=3, coefficient=0.1) + output = layer(value) + expected = reference(expected_input, 3, 0.1) + cotangent = torch.randn_like(output) + actual_gradient, = torch.autograd.grad(output, value, cotangent) + expected_gradient, = torch.autograd.grad(expected, expected_input, cotangent) + torch.testing.assert_close(output, expected) + torch.testing.assert_close(actual_gradient, expected_gradient) + + +def test_finite_difference_gradcheck(): + torch.manual_seed(4) + layer = DevitoDiffusion(steps=2, coefficient=0.12) + value = torch.randn(1, 1, 3, 4, dtype=torch.float64, requires_grad=True) + assert torch.autograd.gradcheck(layer, (value,), eps=1e-6, atol=1e-5, rtol=1e-4) + + +def test_multiple_forwards_and_noncontiguous_input(): + layer = DevitoDiffusion(steps=3, coefficient=0.1) + left = torch.randn(2, 2, 6, 5, dtype=torch.float64).transpose(-1, -2) + left.requires_grad_() + right = torch.randn_like(left, requires_grad=True) + left_output, right_output = layer(left), layer(right) + saved = left_output.detach().clone() + layer(torch.zeros_like(left)) + torch.testing.assert_close(left_output, saved) + left_output.sum().backward() + right_output.square().sum().backward() + expected_left = left.detach().clone().requires_grad_() + expected_right = right.detach().clone().requires_grad_() + reference(expected_left, 3, 0.1).sum().backward() + reference(expected_right, 3, 0.1).square().sum().backward() + torch.testing.assert_close(left.grad, expected_left.grad) + torch.testing.assert_close(right.grad, expected_right.grad) + + +def test_network_trains_through_physics(): + network, losses = train_demo() + assert losses[-1] < 0.5*losses[0] + assert network[0].weight.grad.abs().sum() > 0 + assert network[3].weight.grad.abs().sum() > 0 + + +@pytest.mark.parametrize('steps,coefficient', [(0, .1), (2, -.1), (2, .3)]) +def test_invalid_configuration(steps, coefficient): + with pytest.raises(ValueError): + DevitoDiffusion(steps, coefficient) + + +def test_invalid_input(): + layer = DevitoDiffusion() + with pytest.raises(ValueError): + layer(torch.zeros(8, 8)) + with pytest.raises(TypeError): + layer(torch.zeros((1, 1, 8, 8), dtype=torch.int64)) diff --git a/tests/test_torch_openacc.py b/tests/test_torch_openacc.py new file mode 100644 index 00000000000..bf9c61486ef --- /dev/null +++ b/tests/test_torch_openacc.py @@ -0,0 +1,247 @@ +"""Compiled OpenACC host tests and CUDA tensor interoperability tests. + +GCC host execution checks generated OpenACC kernels without requiring a GPU. +CUDA cases require CUDA-enabled PyTorch, an NVIDIA GPU, and the HPC SDK. +""" + +import os +import shutil +import subprocess +import sys +from pathlib import Path + +import numpy as np +import pytest + +torch = pytest.importorskip('torch') + +from devito import ( # noqa: E402 + Eq, Function, Grid, Operator, TimeFunction, TorchAllocator +) +from devito.arch import NVIDIAX # noqa: E402 +from devito.arch.compiler import GNUCompiler, compiler_registry # noqa: E402 +from devito.exceptions import InvalidArgument, InvalidOperator # noqa: E402 +from devito.passes.iet.languages.openacc import AccBB # noqa: E402 +from devito.symbolics import Macro # noqa: E402 + + +class HostOpenACCCompiler(GNUCompiler): + """Use GCC's C++ OpenACC frontend for host-only execution checks.""" + + _default_cpp = True + + def __init_finalize__(self, **kwargs): + super().__init_finalize__(**kwargs) + self.cflags.append('-fopenacc') + self.ldflags.append('-fopenacc') + + +def _reference(value, steps): + weight = value.new_tensor([[0, .1, 0], [.1, .6, .1], [0, .1, 0]])[None, None] + # Compare the discretization at float32 precision, not cuDNN's optional TF32 + with torch.backends.cudnn.flags(allow_tf32=False): + for _ in range(steps): + value = torch.nn.functional.conv2d(value[None, None], weight, padding=1)[0, 0] + return value + + +def _host_case(case, precision): + """Execute an actual GCC OpenACC kernel against tensor-owned host storage.""" + AccBB.mapper[NVIDIAX] = Macro('acc_device_host') + compiler_registry['openacc-host'] = HostOpenACCCompiler + dtype = getattr(np, precision) + grid = Grid(shape=(12, 12), extent=(11., 11.)) + allocator = TorchAllocator() + u = TimeFunction(name='acc_u', grid=grid, dtype=dtype, space_order=2, + allocator=allocator) + if case == 'mixed': + coefficient = Function(name='coefficient', grid=grid, dtype=dtype) + coefficient.data[:] = 2 + expression = Eq(u.forward, u+coefficient) + else: + expression = Eq(u.forward, u+.1*u.laplace) if case == 'stencil' else Eq( + u.forward, u+1) + op = Operator(expression, platform='nvidiaX', language='openacc', + compiler='openacc-host') + if case == 'override': + data = torch.zeros(tuple(u.shape_allocated), dtype=getattr(torch, precision)) + ptr = data.data_ptr() + op.apply(acc_u=data, time_m=0, time_M=4) + assert data.data_ptr() == ptr + assert u._data is None + torch.testing.assert_close(data[u._mask_domain][1], torch.full( + grid.shape, 5., dtype=data.dtype)) + else: + torch.manual_seed(3) + initial = torch.randn(grid.shape, dtype=getattr(torch, precision)) + u.data[0].copy_(initial) + ptr = u._data_allocated.data_ptr() + op.apply(time_M=4) + expected = _reference(initial, 5) if case == 'stencil' else initial + ( + 10 if case == 'mixed' else 5) + torch.testing.assert_close(u.data[5], expected) + assert u._data_allocated.data_ptr() == ptr + assert 'acc parallel loop' in str(op.ccode) + assert op.cfunction is not None + + +@pytest.mark.skipif(shutil.which('gcc') is None or sys.platform != 'linux', + reason='Requires GCC OpenACC on Linux') +@pytest.mark.parametrize('precision', ['float32', 'float64']) +@pytest.mark.parametrize('case', ['update', 'stencil', 'override', 'mixed']) +def test_openacc_host_execution(case, precision): + # PyTorch 2.2 bundles an older libgomp without current GOACC symbols. Load + # the system compiler's runtime first in a fresh interpreter, preserving + # inherited preloads. This is host OpenACC execution, not CUDA validation. + runtime = subprocess.check_output(['gcc', '-print-file-name=libgomp.so.1'], + text=True).strip() + environment = os.environ.copy() + environment['LD_PRELOAD'] = ':'.join(filter(None, [ + runtime, environment.get('LD_PRELOAD') + ])) + environment['DEVITO_LOGGING'] = 'ERROR' + result = subprocess.run([sys.executable, str(Path(__file__).resolve()), + '--host', case, precision], env=environment, + capture_output=True, text=True, timeout=90) + assert result.returncode == 0, result.stdout + result.stderr + + +cuda_available = torch.cuda.is_available() and shutil.which('nvc++') is not None +cuda_required = pytest.mark.skipif(not cuda_available, + reason='Requires CUDA PyTorch, NVIDIA GPU and nvc++') + + +def test_openacc_deviceptr_codegen(): + u = TimeFunction(name='u', grid=Grid(shape=(8, 8)), + allocator=TorchAllocator(device='cuda:0')) + op = Operator(Eq(u.forward, u+1), platform='nvidiaX', + language='openacc', compiler='gcc') + code = str(op.ccode) + assert '#pragma acc parallel loop' in code + assert 'deviceptr(u)' in code + assert 'copyin(' not in code + assert 'copyout(' not in code + assert 'acc enter data' not in code + assert 'acc exit data' not in code + assert u._data is None + + +def test_openacc_host_access_rejected(): + u = TimeFunction(name='u', grid=Grid(shape=(8, 8)), + allocator=TorchAllocator(device='cuda:0')) + with pytest.raises(InvalidOperator, match='host code'): + Operator(Eq(u[0, 0, 0], 1), platform='nvidiaX', + language='openacc', compiler='gcc') + + +def test_openacc_mixed_storage_codegen(): + grid = Grid(shape=(8, 8)) + u = TimeFunction(name='u', grid=grid, allocator=TorchAllocator(device='cuda:0')) + a = Function(name='a', grid=grid) + op = Operator(Eq(u.forward, u+a), platform='nvidiaX', + language='openacc', compiler='gcc') + code = str(op.ccode) + assert 'deviceptr(u)' in code + assert 'present(a)' in code + assert 'copyin(a[' in code + assert 'copyin(u[' not in code + + +@cuda_required +@pytest.mark.parametrize('precision', ['float32', 'float64']) +@pytest.mark.parametrize('case', ['owned', 'external', 'override', 'stencil', 'mixed']) +def test_openacc_cuda_storage(case, precision): + grid = Grid(shape=(12, 12), extent=(11., 11.)) + dtype = getattr(np, precision) + tensor_dtype = getattr(torch, precision) + if case == 'external': + # space_order=0 makes the full external allocation equal to the domain. + data = torch.ones(grid.shape, dtype=tensor_dtype, device='cuda:0') + f = Function(name='f', grid=grid, space_order=0, dtype=dtype, + allocator=TorchAllocator(tensor=data), initializer=lambda _: None) + op = Operator(Eq(f, 2*f+1), platform='nvidiaX', language='openacc', + compiler='nvc') + ptr = data.data_ptr() + op.apply() + torch.testing.assert_close(data, torch.full_like(data, 3)) + assert data.data_ptr() == ptr + return + u = TimeFunction(name='cuda_u', grid=grid, dtype=dtype, space_order=2, + allocator=TorchAllocator(device='cuda:0')) + expression = Eq(u.forward, u+.1*u.laplace) if case == 'stencil' else Eq( + u.forward, u+1) + if case == 'mixed': + coefficient = Function(name='coefficient', grid=grid, dtype=dtype) + coefficient.data[:] = 2 + expression = Eq(u.forward, u+coefficient) + op = Operator(expression, platform='nvidiaX', language='openacc', compiler='nvc') + if case == 'override': + data = torch.zeros(tuple(u.shape_allocated), dtype=tensor_dtype, device='cuda:0') + # Producer work on a nondefault stream must finish before OpenACC runs. + producer = torch.cuda.Stream() + producer.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(producer): + data.fill_(2) + ptr = data.data_ptr() + op.apply(cuda_u=data, time_m=0, time_M=4) + torch.testing.assert_close(data[u._mask_domain][1], torch.full( + grid.shape, 7., dtype=tensor_dtype, device='cuda:0')) + assert data.data_ptr() == ptr + else: + initial = torch.randn(grid.shape, dtype=tensor_dtype, device='cuda:0') + u.data[0].copy_(initial) + ptr = u._data_allocated.data_ptr() + op.apply(time_M=4) + expected = _reference(initial, 5) if case == 'stencil' else initial + ( + 10 if case == 'mixed' else 5) + torch.testing.assert_close(u.data[5], expected) + assert u._data_allocated.data_ptr() == ptr + code = str(op.ccode) + assert 'deviceptr(cuda_u)' in code + assert 'copyin(cuda_u[' not in code + assert 'copyout(cuda_u[' not in code + + +@cuda_required +def test_openacc_cuda_rejects_incorrect_device(): + f = Function(name='f', grid=Grid(shape=(8, 8)), space_order=0, + allocator=TorchAllocator(device='cuda:0')) + op = Operator(Eq(f, f+1), platform='nvidiaX', language='openacc', compiler='nvc') + with pytest.raises(InvalidArgument, match='deviceid'): + op.apply(deviceid=1) + + +@cuda_required +@pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.device_count() < 2, + reason='Requires two CUDA devices') +def test_openacc_cuda_second_device(): + data = torch.ones((8, 8), device='cuda:1') + f = Function(name='f', grid=Grid(shape=(8, 8)), space_order=0, + allocator=TorchAllocator(tensor=data), initializer=lambda _: None) + op = Operator(Eq(f, f+1), platform='nvidiaX', language='openacc', compiler='nvc') + op.apply() + torch.testing.assert_close(data, torch.full_like(data, 2)) + + +@pytest.mark.skipif(not torch.cuda.is_available() or shutil.which('nvc++') is None, + reason='Requires CUDA PyTorch, an NVIDIA GPU and nvc++') +def test_openacc_cuda_execution(): + grid = Grid(shape=(16, 16)) + u = TimeFunction(name='u', grid=grid, allocator=TorchAllocator(device='cuda:0')) + op = Operator(Eq(u.forward, u+1), platform='nvidiaX', + language='openacc', compiler='nvc') + tensor = torch.zeros(tuple(u.shape_allocated), device='cuda:0') + producer = torch.cuda.Stream() + producer.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(producer): + tensor.fill_(2) + ptr = tensor.data_ptr() + op.apply(u=tensor, time_m=0, time_M=3) + assert tensor.data_ptr() == ptr + torch.testing.assert_close(tensor[u._mask_domain][0], + torch.full((16, 16), 6., device='cuda:0')) + + +if __name__ == '__main__': + assert sys.argv[1] == '--host' + _host_case(sys.argv[2], sys.argv[3]) From 2162b9bf2dbf1b8ece89c18f048afe11aa74c840 Mon Sep 17 00:00:00 2001 From: Mathias Louboutin Date: Sat, 7 Nov 2026 18:30:16 +0000 Subject: [PATCH 02/12] arch: Compile Torch kernels with C++20 --- devito/core/torch.py | 2 +- docs/torch.md | 2 +- tests/test_torch_codegen.py | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/devito/core/torch.py b/devito/core/torch.py index 2484c4c5cdf..aaa089621dd 100644 --- a/devito/core/torch.py +++ b/devito/core/torch.py @@ -14,7 +14,7 @@ class TorchOperatorMixin: _ccode_handler = TorchCGen _Target = TorchTarget - _compiler_options = {'cpp': True, 'std': 'c++17'} + _compiler_options = {'cpp': True, 'std': 'c++20'} @classmethod def _normalize_kwargs(cls, **kwargs): diff --git a/docs/torch.md b/docs/torch.md index 486cbd3e1c1..2c73351a3af 100644 --- a/docs/torch.md +++ b/docs/torch.md @@ -38,7 +38,7 @@ python -m pip install -e '.[torch,tests]' ``` The optional extra uses the current main branch's dependency ranges. CPU code -generation requires a C++17 compiler. The compiler standard is selected at construction through its +generation requires a C++20 compiler. The compiler standard is selected at construction through its `std` property. The Torch target declares headers, libraries and ABI settings as IET pass metadata. The existing graph machinery applies those requirements to the operator-local compiler. Keyword normalization does not diff --git a/tests/test_torch_codegen.py b/tests/test_torch_codegen.py index abb2f397cc5..e8cde5ef09c 100644 --- a/tests/test_torch_codegen.py +++ b/tests/test_torch_codegen.py @@ -91,7 +91,7 @@ def test_custom_compiler(monkeypatch): f = Function(name='custom', grid=Grid(shape=(8, 8)), space_order=0, allocator=TorchAllocator()) op = Operator(Eq(f, f+1), language='torch', compiler='custom') - assert '-std=c++17' in op._compiler.cflags + assert '-std=c++20' in op._compiler.cflags assert '-std=c99' not in op._compiler.cflags op.apply() assert torch.all(f.data == 1) @@ -202,7 +202,7 @@ def test_torch_compiler_requirements_are_operator_local(): allocator=TorchAllocator()) op = Operator(Eq(f, f+1), language='torch', compiler='gcc') assert op._compiler.src_ext == 'cpp' - assert '-std=c++17' in op._compiler.cflags + assert '-std=c++20' in op._compiler.cflags assert 'torch_cpu' in op._compiler.libraries assert previous == (template.cc, template.src_ext, template.cflags, template.include_dirs, template.libraries) From 3111356685105eaed2abd995a78ac1246c64a746 Mon Sep 17 00:00:00 2001 From: Mathias Louboutin Date: Sat, 7 Nov 2026 18:30:16 +0000 Subject: [PATCH 03/12] data: Preserve the established NumPy Data representation --- devito/data/data.py | 26 ++++++++++++++++++++++++++ devito/types/grid.py | 20 ++++++++++---------- tests/test_data.py | 17 ++++++++++++++++- 3 files changed, 52 insertions(+), 11 deletions(-) diff --git a/devito/data/data.py b/devito/data/data.py index 51892f10d35..8fa5c750f45 100644 --- a/devito/data/data.py +++ b/devito/data/data.py @@ -645,6 +645,32 @@ def __reduce__(self): def __array_finalize__(self, obj): self._finalize_metadata(obj) + def __repr__(self): + """Preserve the public Data representation using NumPy's array formatter.""" + prefix = f'{Data.__name__}(' + implied = self.dtype in (np.dtype(int), np.dtype(float), np.dtype(complex), + np.dtype(bool)) and self.size > 0 + suffix = ')' if implied else ',' + if self.size or self.shape == (0,): + values = np.array2string(self._raw, separator=', ', prefix=prefix, + suffix=suffix) + else: + values = f'[], shape={self.shape}' + result = f'{prefix}{values}{suffix}' + if implied: + return result + if self.dtype.names or not self.dtype.isnative or self.dtype.kind in 'SUV': + dtype = str(self.dtype) + if not self.dtype.names: + dtype = repr(dtype) + else: + dtype = self.dtype.name + dtype = f'dtype={dtype})' + width = np.get_printoptions()['linewidth'] + last_line = len(result.rsplit('\n', 1)[-1]) + separator = '\n' + ' '*len(prefix) if last_line + len(dtype) + 1 > width else ' ' + return f'{result}{separator}{dtype}' + @property def _raw(self): return np.ndarray.view(self, np.ndarray) diff --git a/devito/types/grid.py b/devito/types/grid.py index f4ed50faded..4568b852bf1 100644 --- a/devito/types/grid.py +++ b/devito/types/grid.py @@ -1131,16 +1131,16 @@ class SubDomainSet(MultiSubDomain): >>> op = Operator(eq) >>> summary = op.apply() >>> f.data - array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - [0, 1, 1, 1, 1, 1, 1, 1, 1, 0], - [0, 1, 1, 1, 1, 1, 1, 1, 1, 0], - [0, 1, 1, 1, 1, 1, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - [0, 1, 1, 1, 1, 1, 1, 1, 1, 0], - [0, 1, 1, 1, 1, 1, 1, 1, 1, 0], - [0, 1, 1, 1, 1, 1, 1, 1, 1, 0], - [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], dtype=int32) + Data([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [0, 1, 1, 1, 1, 1, 1, 1, 1, 0], + [0, 1, 1, 1, 1, 1, 1, 1, 1, 0], + [0, 1, 1, 1, 1, 1, 1, 1, 1, 0], + [0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [0, 1, 1, 1, 1, 1, 1, 1, 1, 0], + [0, 1, 1, 1, 1, 1, 1, 1, 1, 0], + [0, 1, 1, 1, 1, 1, 1, 1, 1, 0], + [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], dtype=int32) """ def __init__(self, **kwargs): diff --git a/tests/test_data.py b/tests/test_data.py index 28698b93442..3b073c0468f 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -8,7 +8,7 @@ PrecomputedSparseFunction, PrecomputedSparseTimeFunction, SparseFunction, SparseTimeFunction, TimeFunction, clear_cache, configuration, switchconfig ) -from devito.data import LEFT, RIGHT, Decomposition, convert_index, loc_data_idx +from devito.data import LEFT, RIGHT, Decomposition, NumpyData, convert_index, loc_data_idx from devito.data.allocators import DataReference, PosixAllocator from devito.data.distributed.layout import Layout from devito.data.distributed.selection import ( @@ -322,6 +322,21 @@ def test_indexing_into_sparse_subfunc_singlempi(self, mode): assert np.allclose(s.coordinates.data[-1, :], coords[-1, :] / 2) + @pytest.mark.parametrize('value', [ + np.arange(12, dtype=np.float32).reshape(3, 4), + np.arange(6, dtype=np.int32).reshape(2, 3), + np.ones((2, 3), dtype=np.float64), + np.empty((0, 2), dtype=np.float32), + np.arange(100, dtype=np.int64), + ]) + def test_representation(self, value): + class Data(np.ndarray): + pass + + data = NumpyData(value.shape, value.dtype.type, modulo=(False,)*value.ndim) + data[:] = value + assert repr(data) == repr(value.view(Data)) + class TestLocDataIDX: """ From f4e2ea0e2b411f0bfb84f02d8b650fa8f3a94edb Mon Sep 17 00:00:00 2001 From: Mathias Louboutin Date: Sat, 7 Nov 2026 18:30:16 +0000 Subject: [PATCH 04/12] tests: Give allocator cases stable parallel test IDs --- tests/test_data_backends.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_data_backends.py b/tests/test_data_backends.py index f276fcfb551..9423b3cf99f 100644 --- a/tests/test_data_backends.py +++ b/tests/test_data_backends.py @@ -32,7 +32,7 @@ def test_common_base(): assert not issubclass(Data, np.ndarray) -@pytest.mark.parametrize('allocator', [None, TorchAllocator()]) +@pytest.mark.parametrize('allocator', [None, TorchAllocator()], ids=['numpy', 'torch']) def test_host_view_shares_storage(allocator): f = Function(name='host_view', grid=Grid(shape=(8, 8)), space_order=0, allocator=allocator) From 79c5b1c8765b8da5d50a2c0bd0ec6dfbe5a203b0 Mon Sep 17 00:00:00 2001 From: Mathias Louboutin Date: Sat, 7 Nov 2026 18:46:33 +0000 Subject: [PATCH 05/12] arch: Honor OpenMP options for Apple Clang --- devito/arch/compiler.py | 2 +- tests/test_torch_codegen.py | 13 ++++++++++++- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/devito/arch/compiler.py b/devito/arch/compiler.py index 19d9a3a20d9..b718e5be77d 100644 --- a/devito/arch/compiler.py +++ b/devito/arch/compiler.py @@ -590,7 +590,7 @@ def __init_finalize__(self, **kwargs): # The compiler can be installed with Homebrew or can be built from scratch. # Check if installed and set compiler flags accordingly llvmm1 = get_m1_llvm_path(language) - if llvmm1 and 'openmp' in language: + if llvmm1 and self.openmp: mx = platform.march self.ldflags += [f'-mcpu=apple-{mx}', '-fopenmp', f'-L{llvmm1["libs"]}'] diff --git a/tests/test_torch_codegen.py b/tests/test_torch_codegen.py index e8cde5ef09c..ac3f26b16a1 100644 --- a/tests/test_torch_codegen.py +++ b/tests/test_torch_codegen.py @@ -11,7 +11,7 @@ Constant, Dimension, Eq, Function, Grid, Operator, TimeFunction, TorchAllocator, configuration, switchconfig ) -from devito.arch import Cpu64 # noqa: E402 +from devito.arch import M1, Cpu64 # noqa: E402 from devito.arch.compiler import ( # noqa: E402 ClangCompiler, CustomCompiler, GNUCompiler, IntelCompiler, NvidiaCompiler, OneapiCompiler @@ -152,6 +152,17 @@ def test_cpu_compiler_enables_openmp(monkeypatch, compiler_type, flag): assert compiler.src_ext == 'cpp' +@pytest.mark.parametrize('openmp', [False, True]) +def test_apple_compiler_honors_openmp_option(monkeypatch, openmp): + monkeypatch.setattr('devito.arch.compiler.sniff_compiler_version', + lambda cc, **kwargs: Version('16.0')) + monkeypatch.setattr(type(M1), 'march', 'm1') + paths = {'libs': '/llvm/lib', 'include': '/llvm/include'} + monkeypatch.setattr('devito.arch.compiler.get_m1_llvm_path', lambda language: paths) + compiler = ClangCompiler(cpp=True, openmp=openmp, language='torch', platform=M1) + assert ('-fopenmp' in compiler.ldflags) == openmp + + def test_custom_compiler_honors_cxxflags(monkeypatch): for name in ('CC', 'CXX', 'CFLAGS', 'LDFLAGS'): monkeypatch.delenv(name, raising=False) From 40c82f6afa26be7dc938038ef8e1350e19c77fb6 Mon Sep 17 00:00:00 2001 From: Mathias Louboutin Date: Sat, 7 Nov 2026 20:47:34 +0000 Subject: [PATCH 06/12] data: Normalize boolean masks by their consumed axes --- devito/data/data.py | 27 +++++++++++++++++------ devito/data/distributed/selection.py | 15 ++----------- devito/data/utils.py | 13 +++++++++++ tests/test_data_backends.py | 32 ++++++++++++++++++++++++++++ 4 files changed, 67 insertions(+), 20 deletions(-) diff --git a/devito/data/data.py b/devito/data/data.py index 8fa5c750f45..2f95ef7a9f4 100644 --- a/devito/data/data.py +++ b/devito/data/data.py @@ -418,12 +418,14 @@ def _normalize_index(self, idx): return (idx,) else: idx = tuple(self._index_array(i) for i in as_tuple(idx)) + axes = sum(index_axis_span(i) for i in idx if i is not Ellipsis) + missing = (slice(None),)*max(0, self.ndim - axes) if any(i is Ellipsis for i in idx): # Explicitly replace the Ellipsis - items = (slice(None),)*(self.ndim - len(idx) + 1) - items = idx[:idx.index(Ellipsis)] + items + idx[idx.index(Ellipsis)+1:] + position = next(n for n, i in enumerate(idx) if i is Ellipsis) + items = idx[:position] + missing + idx[position+1:] else: - items = idx + (slice(None),)*(self.ndim - len(idx)) + items = idx + missing # Normalize slice steps: processed = [slice(i.start, i.stop, 1) if (isinstance(i, slice) and i.step is None) @@ -484,10 +486,21 @@ def _index_glb_to_loc(self, glb_idx): return glb_idx loc_idx = [] - for i, s, mod, dec in zip( - glb_idx, self.shape, self._modulo, self._decomposition, strict=False - ): - if mod is True: + axes = tuple(zip(self.shape, self._modulo, self._decomposition, strict=False)) + axis = 0 + for i in glb_idx: + if axis >= len(axes): + break + span = index_axis_span(i) + if span != 1: + # Boolean masks may cover several axes; new axes cover none + loc_idx.append(i) + axis += span + continue + s, mod, dec = axes[axis] + axis += 1 + is_boolean = isinstance(i, np.ndarray) and i.dtype == np.bool_ + if mod is True and not is_boolean: # Need to wrap index based on modulo v = index_apply_modulo(i, s) elif self._is_distributed is True and dec is not None: diff --git a/devito/data/distributed/selection.py b/devito/data/distributed/selection.py index 34cb5b530c0..ef29443b5bb 100644 --- a/devito/data/distributed/selection.py +++ b/devito/data/distributed/selection.py @@ -13,6 +13,7 @@ import numpy as np +from devito.data.utils import index_axis_span from devito.tools import is_integer __all__ = ['Affine', 'Explicit', 'IndexScalar', 'Selection', 'index_has_array', @@ -93,18 +94,6 @@ def _is_legacy_multidim_basic(idx, ndim): all(is_integer(i) for i in idx)) -def _axis_span(item): - """Number of array axes consumed by one raw index component.""" - if item is np.newaxis: - return 0 - elif isinstance(item, np.ndarray) and item.dtype == np.bool_: - return item.ndim - elif isinstance(item, (list, tuple)): - arr = np.asarray(item) - return arr.ndim if arr.dtype == np.bool_ else 1 - return 1 - - def _expand(idx, ndim): """ Expand a raw index into exactly `ndim` per-axis components, replacing @@ -123,7 +112,7 @@ def _expand(idx, ndim): # Replace Ellipsis with the right number of full slices if any(i is Ellipsis for i in items): pos = next(k for k, i in enumerate(items) if i is Ellipsis) - consumed = sum(_axis_span(i) for i in items if i is not Ellipsis) + consumed = sum(index_axis_span(i) for i in items if i is not Ellipsis) fill = [slice(None)] * (ndim - consumed) items = items[:pos] + fill + items[pos + 1:] diff --git a/devito/data/utils.py b/devito/data/utils.py index 45f6c376bc9..c8ed8d61d04 100644 --- a/devito/data/utils.py +++ b/devito/data/utils.py @@ -9,6 +9,7 @@ 'convert_index', 'flip_idx', 'index_apply_modulo', + 'index_axis_span', 'index_dist_to_repl', 'index_handle_oob', 'index_is_basic', @@ -17,6 +18,18 @@ ] +def index_axis_span(item): + """Number of array axes consumed by one raw index component.""" + if item is np.newaxis: + return 0 + elif isinstance(item, np.ndarray) and item.dtype == np.bool_: + return item.ndim + elif isinstance(item, (list, tuple)): + arr = np.asarray(item) + return arr.ndim if arr.dtype == np.bool_ else 1 + return 1 + + def is_windowed(val): """ Whether `val` provides its values by window rather than from memory, such diff --git a/tests/test_data_backends.py b/tests/test_data_backends.py index 9423b3cf99f..7bb1408a4ed 100644 --- a/tests/test_data_backends.py +++ b/tests/test_data_backends.py @@ -83,6 +83,38 @@ def test_indexing_assignment_and_broadcasting(allocator): np.testing.assert_equal(native(f.data), np.broadcast_to(np.arange(8), (8, 8))) +@pytest.mark.parametrize('case', [ + 'spatial', 'all_time', 'leading', 'ellipsis', 'trailing_ellipsis', 'modulo' +]) +def test_multidimensional_boolean_masks(allocator, case): + u = TimeFunction(name='masked', grid=Grid(shape=(4, 5)), space_order=0, + padding=0, allocator=allocator) + expected = np.arange(40, dtype=np.float32).reshape(2, 4, 5) + u.data[:] = expected + spatial = np.arange(20).reshape(4, 5) % 3 == 0 + leading = np.arange(8).reshape(2, 4) % 2 == 0 + if case == 'spatial': + index = (0, spatial) + elif case == 'all_time': + index = (slice(None), spatial) + elif case == 'leading': + index = (leading, slice(None)) + elif case == 'ellipsis': + index = (Ellipsis, spatial) + elif case == 'trailing_ellipsis': + index = (leading, Ellipsis) + else: + index = (4, spatial) + reference = (0, spatial) if case == 'modulo' else index + if allocator is not None: + index = tuple(torch.from_numpy(i) if isinstance(i, np.ndarray) else i + for i in index) + np.testing.assert_equal(native(u.data[index]), expected[reference]) + u.data[index] = 9 + expected[reference] = 9 + np.testing.assert_equal(native(u.data), expected) + + def test_decomposition_and_nonlocal_access(allocator): decomposition = Decomposition([range(4), range(4, 8)], 1) data = Data((4,), np.float32, decomposition=(decomposition,), From 6ff2665bd83007d7830898fdaeb3aba4832e0816 Mon Sep 17 00:00:00 2001 From: Mathias Louboutin Date: Sun, 8 Nov 2026 02:29:21 +0000 Subject: [PATCH 07/12] arch: Reuse the loaded OpenMP runtime on Apple Silicon --- devito/arch/archinfo.py | 11 ++++++ devito/arch/compiler.py | 3 +- examples/pytorch/devito_layer.py | 2 +- tests/test_arch.py | 22 +++++++++++- tests/test_data_backends.py | 2 +- tests/test_torch_allocators.py | 4 +-- tests/test_torch_codegen.py | 57 ++++++++++++++++++++++++-------- 7 files changed, 81 insertions(+), 20 deletions(-) diff --git a/devito/arch/archinfo.py b/devito/arch/archinfo.py index ccaed53cfe7..7b3f75ae574 100644 --- a/devito/arch/archinfo.py +++ b/devito/arch/archinfo.py @@ -638,6 +638,17 @@ def get_m1_llvm_path(language): if os.path.exists(os.path.join(prefix, 'lib', 'libomp.dylib')): libs = os.path.join(prefix, "lib") include = os.path.join(prefix, "include") + # Reuse an existing LLVM OpenMP runtime. Loading another copy, e.g. + # after importing a numerical package, aborts on first parallel use + dyld = ctypes.CDLL(None) + dyld._dyld_image_count.restype = ctypes.c_uint32 + dyld._dyld_get_image_name.argtypes = [ctypes.c_uint32] + dyld._dyld_get_image_name.restype = ctypes.c_char_p + for i in range(dyld._dyld_image_count()): + name = os.fsdecode(dyld._dyld_get_image_name(i)) + if os.path.basename(name) == 'libomp.dylib': + libs = os.path.dirname(name) + break return {'libs': libs, 'include': include} elif language == "openmp": warning("Apple's arm64 clang found but openmp libraries not found." diff --git a/devito/arch/compiler.py b/devito/arch/compiler.py index b718e5be77d..d5cf541b22f 100644 --- a/devito/arch/compiler.py +++ b/devito/arch/compiler.py @@ -593,7 +593,8 @@ def __init_finalize__(self, **kwargs): if llvmm1 and self.openmp: mx = platform.march self.ldflags += [f'-mcpu=apple-{mx}', - '-fopenmp', f'-L{llvmm1["libs"]}'] + '-fopenmp', f'-L{llvmm1["libs"]}', + f'-Wl,-rpath,{llvmm1["libs"]}'] self.cflags += ['-Xclang', f'-I{llvmm1["include"]}'] else: if platform in [POWER8, POWER9]: diff --git a/examples/pytorch/devito_layer.py b/examples/pytorch/devito_layer.py index 950819a7ae2..8fd7bf7ba57 100644 --- a/examples/pytorch/devito_layer.py +++ b/examples/pytorch/devito_layer.py @@ -37,7 +37,7 @@ def __init__(self, planes, height, width, dtype, steps, coefficient): self.symbol = self._state() self.operator = Operator( Eq(self.symbol.forward, self.symbol + coefficient*self.symbol.laplace), - language='torch', compiler='gcc', platform='cpu64', + language='torch', opt=('advanced', {'openmp': True}) ) diff --git a/tests/test_arch.py b/tests/test_arch.py index 305887d2e2e..64eda1857d8 100644 --- a/tests/test_arch.py +++ b/tests/test_arch.py @@ -1,12 +1,13 @@ from concurrent.futures import ThreadPoolExecutor from contextlib import suppress +from unittest.mock import Mock import pytest from packaging.version import Version from devito import configuration, switchconfig from devito.arch import BLACKWELL, NVIDIAX, platform_registry -from devito.arch.archinfo import AppleArm, get_apple_platform +from devito.arch.archinfo import AppleArm, get_apple_platform, get_m1_llvm_path from devito.arch.compiler import ( CudaCompiler, GNUCompiler, compiler_registry, sniff_compiler_version ) @@ -40,6 +41,25 @@ def test_apple_platform(brand, expected): assert platform.name == expected +@pytest.mark.parametrize('loaded', [ + [], ['/numerical/lib/libomp.dylib'], ['/usr/lib/libSystem.B.dylib'] +]) +def test_apple_openmp_runtime(monkeypatch, loaded): + monkeypatch.setattr('devito.arch.archinfo.sys.platform', 'darwin') + monkeypatch.setattr('devito.arch.archinfo.run', lambda *a, **kw: Mock( + stdout=b'Target: arm64-apple-darwin\nInstalledDir: /llvm/bin\n')) + monkeypatch.setattr('devito.arch.archinfo.os.path.exists', lambda p: True) + dyld = Mock() + dyld._dyld_image_count.return_value = len(loaded) + dyld._dyld_get_image_name.side_effect = [p.encode() for p in loaded] + monkeypatch.setattr('devito.arch.archinfo.ctypes.CDLL', lambda name: dyld) + + paths = get_m1_llvm_path.func('openmp') + assert paths['include'] == '/llvm/include' + assert paths['libs'] == ('/numerical/lib' if '/numerical/lib/libomp.dylib' + in loaded else '/llvm/lib') + + def test_switcharch(): old_compiler = configuration['compiler'] with switchconfig(compiler='gcc-4.9'): diff --git a/tests/test_data_backends.py b/tests/test_data_backends.py index 7bb1408a4ed..61d26f4e675 100644 --- a/tests/test_data_backends.py +++ b/tests/test_data_backends.py @@ -187,7 +187,7 @@ def test_mpi_tensor_kernel(allocator, mode): u = TimeFunction(name='u', grid=grid, space_order=2, allocator=allocator) u.data[0] = 1 language = 'torch' if allocator else 'C' - op = Operator(Eq(u.forward, u + 0.1*u.laplace), language=language, compiler='gcc') + op = Operator(Eq(u.forward, u + 0.1*u.laplace), language=language) op.apply(time_M=0) # MPI-generated stencil execution must agree with the same global reference. gathered = u.data_gather() diff --git a/tests/test_torch_allocators.py b/tests/test_torch_allocators.py index ed7ada53db5..ac6630e7452 100644 --- a/tests/test_torch_allocators.py +++ b/tests/test_torch_allocators.py @@ -32,7 +32,7 @@ def test_external_tensor(monkeypatch): f = Function(name='f', grid=Grid(shape=(8, 8)), space_order=0, allocator=TorchAllocator(tensor=tensor), initializer=lambda _: None) assert f.data.data_ptr() == tensor.data_ptr() - op = Operator(Eq(f, 2*f+1), language='torch', compiler='gcc') + op = Operator(Eq(f, 2*f+1), language='torch') expected = 2*tensor + 1 def no_numpy(*args, **kwargs): @@ -63,7 +63,7 @@ def test_initializers(initializer): def test_invalid_override(kind): f = Function(name='f', grid=Grid(shape=(8, 8)), space_order=0, allocator=TorchAllocator()) - op = Operator(Eq(f, f+1), language='torch', compiler='gcc') + op = Operator(Eq(f, f+1), language='torch') if kind == 'dtype': value = torch.zeros((8, 8), dtype=torch.float64) elif kind == 'strides': diff --git a/tests/test_torch_codegen.py b/tests/test_torch_codegen.py index ac3f26b16a1..a9524775f1a 100644 --- a/tests/test_torch_codegen.py +++ b/tests/test_torch_codegen.py @@ -1,5 +1,9 @@ """Torch language selection, optimized C++ printing and native JIT execution.""" +import subprocess +import sys +from textwrap import dedent + import numpy as np import pytest from packaging.version import Version @@ -22,15 +26,39 @@ from devito.operator.operator import parse_kwargs # noqa: E402 +@pytest.mark.parametrize('first', ['devito', 'torch']) +def test_openmp_runtime_interoperability(first): + # Keep a runtime abort in the child so its diagnostic is reported by pytest + code = f'import {first}\n' + dedent(''' + import torch + from devito import Eq, Function, Grid, Operator, TorchAllocator + + for allocator in (None, TorchAllocator()): + f = Function(name='f', grid=Grid(shape=(128, 128)), space_order=0, + allocator=allocator) + language = 'torch' if allocator else 'openmp' + op = Operator(Eq(f, f+1), language=language, + opt=('advanced', {'openmp': True})) + assert '#pragma omp parallel' in str(op.ccode) + op.apply() + assert (f.data == 1).all() + assert torch.ones((512, 512)).sum().item() == 512**2 + ''') + result = subprocess.run([sys.executable, '-c', code], capture_output=True, + text=True, timeout=120) + assert result.returncode == 0, result.stdout + result.stderr + + @pytest.mark.parametrize('opt', ['noop', 'advanced', 'advanced-fsg']) def test_standard_operator(opt): u = TimeFunction(name='u', grid=Grid(shape=(8, 8)), allocator=TorchAllocator()) - op = Operator(Eq(u.forward, u+1), language='torch', opt=opt, compiler='gcc') + op = Operator(Eq(u.forward, u+1), language='torch', opt=opt) assert isinstance(op, Operator) assert op._ccode_handler is TorchCGen if opt == 'advanced': assert type(op) is Cpu64AdvTorchOperator - assert op._compiler.cc == 'g++' + assert op._compiler.cc == op._compiler.CXX + ( + f'-{op._compiler.suffix}' if op._compiler.suffix else '') assert op._compiler.src_ext == 'cpp' assert 'at::Tensor' in str(op.ccode) assert '.data_ptr()' in str(op.ccode) @@ -43,7 +71,7 @@ def test_standard_operator(opt): assert u._data_allocated._version > version # Standard ctypes loader and cache, not torch's extension mechanism. assert op.cfunction is op._cfunction - assert op._lib._name.endswith('.so') + assert op._lib._name.endswith(op._compiler.so_ext) assert str(op._compiler.get_jit_dir()) in op._lib._name @@ -58,9 +86,9 @@ def test_cpp_stencil(parallel, dtype): scale = Constant(name='scale', value=0.01, dtype=dtype) opt = ('advanced', {'blockinner': True, 'openmp': parallel}) op = Operator(Eq(u.forward, u + scale*u.laplace), language='torch', - opt=opt, compiler='gcc') + opt=opt) reference = Operator(Eq(ref.forward, ref + scale*ref.laplace), - opt=opt, compiler='gcc') + opt=opt) op.apply(time_M=2) reference.apply(time_M=2) np.testing.assert_allclose(u.data.numpy(), ref.data, rtol=1e-5, atol=1e-6) @@ -73,7 +101,7 @@ def test_allocator_with_c_language(): # Storage choice is independent of language, as with other allocators. f = Function(name='f', grid=Grid(shape=(8, 8)), space_order=0, allocator=TorchAllocator()) - Operator(Eq(f, f+1), language='C', compiler='gcc').apply() + Operator(Eq(f, f+1), language='C').apply() assert torch.all(f.data == 1) @@ -81,12 +109,12 @@ def test_allocator_with_c_language(): def test_integer_codegen(dtype): f = Function(name='integer', grid=Grid(shape=(8, 8)), space_order=0, dtype=dtype, allocator=TorchAllocator()) - Operator(Eq(f, f+3), language='torch', compiler='gcc').apply() + Operator(Eq(f, f+3), language='torch').apply() assert torch.all(f.data == 3) def test_custom_compiler(monkeypatch): - monkeypatch.setenv('CXX', 'g++') + monkeypatch.setenv('CXX', configuration['compiler'].CXX) monkeypatch.setenv('CXXFLAGS', '-O2 -fPIC') f = Function(name='custom', grid=Grid(shape=(8, 8)), space_order=0, allocator=TorchAllocator()) @@ -111,7 +139,7 @@ def test_torch_language_with_host_allocator(): def test_preemptive_autotuning_preserves_tensor_state(monkeypatch): u = TimeFunction(name='tuned', grid=Grid(shape=(12, 12, 12)), allocator=TorchAllocator()) - op = Operator(Eq(u.forward, u.dx+1), language='torch', compiler='gcc', + op = Operator(Eq(u.forward, u.dx+1), language='torch', opt=('blocking', {'openmp': False, 'blockinner': True})) def no_numpy(*args, **kwargs): @@ -133,7 +161,7 @@ def test_function_names_do_not_collide_with_tensor_views(): allocator=TorchAllocator()) f.data[:] = 1 g.data[:] = 2 - op = Operator(Eq(f, f+g), language='torch', compiler='gcc') + op = Operator(Eq(f, f+g), language='torch') op.apply() torch.testing.assert_close(f.data, torch.full((8, 8), 3.)) @@ -161,6 +189,7 @@ def test_apple_compiler_honors_openmp_option(monkeypatch, openmp): monkeypatch.setattr('devito.arch.compiler.get_m1_llvm_path', lambda language: paths) compiler = ClangCompiler(cpp=True, openmp=openmp, language='torch', platform=M1) assert ('-fopenmp' in compiler.ldflags) == openmp + assert ('-Wl,-rpath,/llvm/lib' in compiler.ldflags) == openmp def test_custom_compiler_honors_cxxflags(monkeypatch): @@ -195,8 +224,8 @@ def test_thread_private_temporaries(rotate): c[1, z+1]*u[t, x+2, y+2, z+1]) + _R(u[t, x, y+1, z+1] + 3*u[t, x+1, y+1, z+1]) + _R(u[t, x, y+3, z+1] + 3*u[t, x+1, y+3, z+1]))) - reference = Operator(expression, language='C', opt='noop', compiler='gcc') - op = Operator(expression, language='torch', compiler='gcc', + reference = Operator(expression, language='C', opt='noop') + op = Operator(expression, language='torch', opt=('advanced', {'openmp': True, 'min-storage': True, 'cire-mingain': 0, 'cire-rotate': rotate})) assert FindNodes(Dereference).visit(op) @@ -211,7 +240,7 @@ def test_torch_compiler_requirements_are_operator_local(): list(template.include_dirs), list(template.libraries)) f = Function(name='local_toolchain', grid=Grid(shape=(8, 8)), space_order=0, allocator=TorchAllocator()) - op = Operator(Eq(f, f+1), language='torch', compiler='gcc') + op = Operator(Eq(f, f+1), language='torch') assert op._compiler.src_ext == 'cpp' assert '-std=c++20' in op._compiler.cflags assert 'torch_cpu' in op._compiler.libraries @@ -231,7 +260,7 @@ def test_cpp_only_compiler_accepts_backend_options(monkeypatch): def test_normalize_kwargs_leaves_compiler_unchanged(): - kwargs = parse_kwargs(language='torch', compiler='gcc') + kwargs = parse_kwargs(language='torch') compiler = kwargs['compiler'] attrs = ('cflags', 'defines', 'include_dirs', 'libraries', 'library_dirs', 'ldflags') previous = {name: list(getattr(compiler, name)) for name in attrs} From 87a337c7d2dff43ce120798407a0ab4774b16c52 Mon Sep 17 00:00:00 2001 From: Mathias Louboutin Date: Sun, 8 Nov 2026 02:45:27 +0000 Subject: [PATCH 08/12] ci: Preserve failure diagnostics before worker shutdown --- .github/workflows/pytest-core-nompi.yaml | 2 ++ pyproject.toml | 2 +- requirements-testing.txt | 1 + 3 files changed, 4 insertions(+), 1 deletion(-) diff --git a/.github/workflows/pytest-core-nompi.yaml b/.github/workflows/pytest-core-nompi.yaml index 6e4e7232a27..3b3e336328d 100644 --- a/.github/workflows/pytest-core-nompi.yaml +++ b/.github/workflows/pytest-core-nompi.yaml @@ -163,6 +163,7 @@ jobs: run: | pytest \ -v \ + --capture=no \ -k "${{ matrix.test-set }}" \ -m "not parallel" \ --cov \ @@ -248,6 +249,7 @@ jobs: command: | pytest \ -v \ + --capture=no \ -k "${{ matrix.test-set }}" \ -m "not parallel" \ --cov \ diff --git a/pyproject.toml b/pyproject.toml index c9549c81b44..482349bfdd0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -74,7 +74,7 @@ fallback_version = "0+untagged" [tool.pytest.ini_options] python_files = ["test_*.py", "*_example.py", "example_*.py"] -addopts = "--durations=20 --maxfail=5" +addopts = "--durations=20 --maxfail=5 --exclude-warning-annotations" filterwarnings = ["ignore::pytest.PytestCacheWarning"] [tool.ruff] diff --git a/requirements-testing.txt b/requirements-testing.txt index a20b7bfc0ed..5185972ffe4 100644 --- a/requirements-testing.txt +++ b/requirements-testing.txt @@ -10,4 +10,5 @@ cloudpickle<3.1.3 ipympl<0.10.1 ipykernel<7.0.0 pytest-xdist<3.8.1 +pytest-github-actions-annotate-failures>=0.4.2,<0.5 torch>=2.2 From 5d1a810fda74ea397fc928334ad50a35ae9b2d92 Mon Sep 17 00:00:00 2001 From: Mathias Louboutin Date: Sun, 8 Nov 2026 02:55:23 +0000 Subject: [PATCH 09/12] arch: Keep Apple OpenMP runtime for synthetic CPU targets --- devito/arch/compiler.py | 14 ++++++++++---- tests/test_torch_codegen.py | 27 +++++++++++++++------------ 2 files changed, 25 insertions(+), 16 deletions(-) diff --git a/devito/arch/compiler.py b/devito/arch/compiler.py index d5cf541b22f..0deaefdf0ef 100644 --- a/devito/arch/compiler.py +++ b/devito/arch/compiler.py @@ -1,4 +1,5 @@ import platform +import sys import time import warnings from contextlib import suppress @@ -584,18 +585,23 @@ def __init_finalize__(self, **kwargs): '-fopenmp-targets=amdgcn-amd-amdhsa', '-Xopenmp-target=amdgcn-amd-amdhsa'] self.ldflags += [f'-march={platform.march}'] - elif isinstance(platform, AppleArm): + elif isinstance(platform, AppleArm) or sys.platform == 'darwin': + # Synthetic CPU targets must use the host's OpenMP runtime too # NOTE: # Apple Mx supports OpenMP through Apple's LLVM compiler. # The compiler can be installed with Homebrew or can be built from scratch. # Check if installed and set compiler flags accordingly llvmm1 = get_m1_llvm_path(language) + if not isinstance(platform, AppleArm): + self.cflags += ['-march=native'] if llvmm1 and self.openmp: - mx = platform.march - self.ldflags += [f'-mcpu=apple-{mx}', - '-fopenmp', f'-L{llvmm1["libs"]}', + if isinstance(platform, AppleArm): + self.ldflags += [f'-mcpu=apple-{platform.march}'] + self.ldflags += ['-fopenmp', f'-L{llvmm1["libs"]}', f'-Wl,-rpath,{llvmm1["libs"]}'] self.cflags += ['-Xclang', f'-I{llvmm1["include"]}'] + elif self.openmp and not isinstance(platform, AppleArm): + self.ldflags += ['-fopenmp'] else: if platform in [POWER8, POWER9]: # -march isn't supported on power architectures diff --git a/tests/test_torch_codegen.py b/tests/test_torch_codegen.py index a9524775f1a..062d7c1575d 100644 --- a/tests/test_torch_codegen.py +++ b/tests/test_torch_codegen.py @@ -33,16 +33,17 @@ def test_openmp_runtime_interoperability(first): import torch from devito import Eq, Function, Grid, Operator, TorchAllocator - for allocator in (None, TorchAllocator()): - f = Function(name='f', grid=Grid(shape=(128, 128)), space_order=0, - allocator=allocator) - language = 'torch' if allocator else 'openmp' - op = Operator(Eq(f, f+1), language=language, - opt=('advanced', {'openmp': True})) - assert '#pragma omp parallel' in str(op.ccode) - op.apply() - assert (f.data == 1).all() - assert torch.ones((512, 512)).sum().item() == 512**2 + for platform in (None, 'cpu64-dummy'): + for allocator in (None, TorchAllocator()): + f = Function(name='f', grid=Grid(shape=(128, 128)), space_order=0, + allocator=allocator) + language = 'torch' if allocator else 'openmp' + op = Operator(Eq(f, f+1), language=language, platform=platform, + opt=('advanced', {'openmp': True})) + assert '#pragma omp parallel' in str(op.ccode) + op.apply() + assert (f.data == 1).all() + assert torch.ones((512, 512)).sum().item() == 512**2 ''') result = subprocess.run([sys.executable, '-c', code], capture_output=True, text=True, timeout=120) @@ -181,13 +182,15 @@ def test_cpu_compiler_enables_openmp(monkeypatch, compiler_type, flag): @pytest.mark.parametrize('openmp', [False, True]) -def test_apple_compiler_honors_openmp_option(monkeypatch, openmp): +@pytest.mark.parametrize('platform', [M1, Cpu64('test', isa='cpp')]) +def test_apple_compiler_honors_openmp_option(monkeypatch, openmp, platform): + monkeypatch.setattr('devito.arch.compiler.sys.platform', 'darwin') monkeypatch.setattr('devito.arch.compiler.sniff_compiler_version', lambda cc, **kwargs: Version('16.0')) monkeypatch.setattr(type(M1), 'march', 'm1') paths = {'libs': '/llvm/lib', 'include': '/llvm/include'} monkeypatch.setattr('devito.arch.compiler.get_m1_llvm_path', lambda language: paths) - compiler = ClangCompiler(cpp=True, openmp=openmp, language='torch', platform=M1) + compiler = ClangCompiler(cpp=True, openmp=openmp, language='torch', platform=platform) assert ('-fopenmp' in compiler.ldflags) == openmp assert ('-Wl,-rpath,/llvm/lib' in compiler.ldflags) == openmp From 0dee59cb43a6063e9ef64304181ccbfc70cc41af Mon Sep 17 00:00:00 2001 From: Mathias Louboutin Date: Sun, 8 Nov 2026 04:01:59 +0000 Subject: [PATCH 10/12] compiler: Use portable Torch integer types and CPU test targets --- devito/ir/iet/torch.py | 13 +++++++++++-- tests/test_torch_codegen.py | 6 ++++-- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/devito/ir/iet/torch.py b/devito/ir/iet/torch.py index cef9304edf3..d8f0521594c 100644 --- a/devito/ir/iet/torch.py +++ b/devito/ir/iet/torch.py @@ -1,6 +1,7 @@ """C++ printer for optimized loops over PyTorch tensor storage.""" import cgen as c +import numpy as np from devito.data import FULL from devito.ir.iet.visitors import CGen @@ -19,7 +20,15 @@ def visit_PointerCast(self, o): for d in f.dimensions ) dtype = self.ccode(f.indexed._C_typedata) + # ATen specializes fixed-width integer types. NumPy's C spelling (e.g. + # long) need not name the same C++ type as int64_t on every platform + scalar_dtype = (f'{np.dtype(f.dtype).name}_t' + if np.issubdtype(f.dtype, np.integer) else dtype) tensor = (f'at::from_blob({cast.data}, {{{shape}}}, ' - f'at::TensorOptions().dtype(c10::CppTypeToScalarType<{dtype}>::value)' + 'at::TensorOptions().dtype(' + f'c10::CppTypeToScalarType<{scalar_dtype}>::value)' '.device(at::kCPU))') - return c.Initializer(cast.vdecl, f'{tensor}.data_ptr<{dtype}>()') + pointer = f'{tensor}.data_ptr<{scalar_dtype}>()' + if scalar_dtype != dtype: + pointer = f'reinterpret_cast<{dtype} *>({pointer})' + return c.Initializer(cast.vdecl, pointer) diff --git a/tests/test_torch_codegen.py b/tests/test_torch_codegen.py index 062d7c1575d..7a29d5a9ab8 100644 --- a/tests/test_torch_codegen.py +++ b/tests/test_torch_codegen.py @@ -39,7 +39,7 @@ def test_openmp_runtime_interoperability(first): allocator=allocator) language = 'torch' if allocator else 'openmp' op = Operator(Eq(f, f+1), language=language, platform=platform, - opt=('advanced', {'openmp': True})) + opt=('blocking', {'openmp': True})) assert '#pragma omp parallel' in str(op.ccode) op.apply() assert (f.data == 1).all() @@ -110,7 +110,9 @@ def test_allocator_with_c_language(): def test_integer_codegen(dtype): f = Function(name='integer', grid=Grid(shape=(8, 8)), space_order=0, dtype=dtype, allocator=TorchAllocator()) - Operator(Eq(f, f+3), language='torch').apply() + op = Operator(Eq(f, f+3), language='torch') + assert f'data_ptr<{np.dtype(dtype).name}_t>()' in str(op.ccode) + op.apply() assert torch.all(f.data == 3) From 17c69a908d648baee4c793c036212a4e33a9ce75 Mon Sep 17 00:00:00 2001 From: Mathias Louboutin Date: Sun, 8 Nov 2026 04:01:59 +0000 Subject: [PATCH 11/12] ci: Match CUDA wheels and host OpenACC compilation --- .github/workflows/pytest-gpu.yaml | 2 ++ docker/Dockerfile.devito | 5 +++++ tests/test_torch_openacc.py | 14 +++++++++++--- 3 files changed, 18 insertions(+), 3 deletions(-) diff --git a/.github/workflows/pytest-gpu.yaml b/.github/workflows/pytest-gpu.yaml index 09b7b042c51..2426de7782d 100644 --- a/.github/workflows/pytest-gpu.yaml +++ b/.github/workflows/pytest-gpu.yaml @@ -44,6 +44,7 @@ jobs: - name: pytest-gpu-acc-nvidia test_files: "tests/test_adjoint.py tests/test_gpu_common.py tests/test_gpu_openacc.py tests/test_torch_openacc.py tests/test_operator.py::TestEstimateMemory" base: "devitocodes/bases:nvidia-nvc12" + build_args: "--build-arg TORCH_INDEX_URL=https://download.pytorch.org/whl/cu124" runner_label: NVIDIAGPU test_drive_cmd: "nvidia-smi" # Respect CUDA_VISIBLE_DEVICES and also hard-limit Docker to that device. @@ -80,6 +81,7 @@ jobs: file: docker/Dockerfile.devito tag: ${{ matrix.name }} base: ${{ matrix.base }} + args: ${{ matrix.build_args }} - id: probe name: Probe GPU diff --git a/docker/Dockerfile.devito b/docker/Dockerfile.devito index 672daf0a9b4..e39ac936420 100644 --- a/docker/Dockerfile.devito +++ b/docker/Dockerfile.devito @@ -30,6 +30,11 @@ RUN rm -rf /app/devito/.git RUN eval "$MPI4PY_FLAGS /venv/bin/pip install --no-cache-dir --verbose -r /app/devito/requirements-mpi.txt" # Devito +ARG TORCH_INDEX_URL +RUN if [ -n "$TORCH_INDEX_URL" ]; then \ + /venv/bin/pip install --no-cache-dir --index-url "$TORCH_INDEX_URL" \ + -r /app/devito/requirements-torch.txt; \ + fi RUN /venv/bin/pip install --no-cache-dir -e "/app/devito[extras,tests]" && rm -rf ~/.cache/pip FROM $base AS utilities diff --git a/tests/test_torch_openacc.py b/tests/test_torch_openacc.py index bf9c61486ef..8f777bf12e4 100644 --- a/tests/test_torch_openacc.py +++ b/tests/test_torch_openacc.py @@ -16,7 +16,7 @@ torch = pytest.importorskip('torch') from devito import ( # noqa: E402 - Eq, Function, Grid, Operator, TimeFunction, TorchAllocator + Eq, Function, Grid, Operator, TimeFunction, TorchAllocator, configuration ) from devito.arch import NVIDIAX # noqa: E402 from devito.arch.compiler import GNUCompiler, compiler_registry # noqa: E402 @@ -32,8 +32,8 @@ class HostOpenACCCompiler(GNUCompiler): def __init_finalize__(self, **kwargs): super().__init_finalize__(**kwargs) - self.cflags.append('-fopenacc') - self.ldflags.append('-fopenacc') + self.cflags.extend(['-fopenacc', '-foffload=disable']) + self.ldflags.extend(['-fopenacc', '-foffload=disable']) def _reference(value, steps): @@ -111,6 +111,14 @@ def test_openacc_host_execution(case, precision): reason='Requires CUDA PyTorch, NVIDIA GPU and nvc++') +@pytest.mark.skipif(os.environ.get('CI') != 'true' or + configuration['platform'] is not NVIDIAX, + reason='Requires the NVIDIA CI runtime') +def test_openacc_cuda_runtime(): + assert torch.cuda.is_available(), 'NVIDIA CI requires driver-compatible CUDA PyTorch' + assert shutil.which('nvc++') is not None + + def test_openacc_deviceptr_codegen(): u = TimeFunction(name='u', grid=Grid(shape=(8, 8)), allocator=TorchAllocator(device='cuda:0')) From 4b6445606532734e85781e6eb50914faff5efa62 Mon Sep 17 00:00:00 2001 From: Mathias Louboutin Date: Sun, 8 Nov 2026 04:39:42 +0000 Subject: [PATCH 12/12] arch: Bind NVIDIA kernels to their linked runtime --- devito/arch/compiler.py | 15 ++++++++++++++- tests/test_arch.py | 39 ++++++++++++++++++++++++++++++++++++++- 2 files changed, 52 insertions(+), 2 deletions(-) diff --git a/devito/arch/compiler.py b/devito/arch/compiler.py index 0deaefdf0ef..05a43cc5e28 100644 --- a/devito/arch/compiler.py +++ b/devito/arch/compiler.py @@ -1,3 +1,5 @@ +import ctypes +import os import platform import sys import time @@ -197,6 +199,7 @@ def __init__(self): fields = {'cc', 'ld'} linker_opt = '-Wl,' _default_cpp = False + _load_mode = ctypes.DEFAULT_MODE _cxxstd = 'c++14' _cstd = 'c99' @@ -311,7 +314,11 @@ def load(self, soname): The loaded shared object. """ try: - return npct.load_library(str(self.get_jit_dir().joinpath(soname)), '.') + library = self.get_jit_dir().joinpath(soname) + if self._load_mode != ctypes.DEFAULT_MODE: + return ctypes.CDLL(str(library.with_suffix(self.so_ext)), + mode=self._load_mode) + return npct.load_library(str(library), '.') except AttributeError as e: # Some ctypes paths re-wrap dlopen failures, such as a partially # written .so, as AttributeError. Normalize so jit_compile treats @@ -684,6 +691,8 @@ def __lookup_cmds__(self): class NvidiaCompiler(Compiler): _default_cpp = True + # Prefer the linked NVIDIA runtime over globally loaded OpenACC symbols + _load_mode = getattr(os, 'RTLD_DEEPBIND', ctypes.DEFAULT_MODE) def __init_finalize__(self, **kwargs): @@ -1105,6 +1114,10 @@ def __new_with__(self, **kwargs): def _default_cpp(self): return self._base._default_cpp + @property + def _load_mode(self): + return self._base._load_mode + class CompilerRegistry(dict): """ diff --git a/tests/test_arch.py b/tests/test_arch.py index 64eda1857d8..ac05d97243c 100644 --- a/tests/test_arch.py +++ b/tests/test_arch.py @@ -1,5 +1,9 @@ +import ctypes +import os +import shutil from concurrent.futures import ThreadPoolExecutor from contextlib import suppress +from subprocess import check_call from unittest.mock import Mock import pytest @@ -9,7 +13,8 @@ from devito.arch import BLACKWELL, NVIDIAX, platform_registry from devito.arch.archinfo import AppleArm, get_apple_platform, get_m1_llvm_path from devito.arch.compiler import ( - CudaCompiler, GNUCompiler, compiler_registry, sniff_compiler_version + CudaCompiler, CustomCompiler, GNUCompiler, NvidiaCompiler, compiler_registry, + sniff_compiler_version ) @@ -22,6 +27,38 @@ def test_sniff_compiler_version(cc): sniff_compiler_version(cc) +@pytest.mark.skipif(not hasattr(os, 'RTLD_DEEPBIND') or shutil.which('gcc') is None, + reason='Requires ELF deep binding and GCC') +@pytest.mark.parametrize('custom', [False, True]) +def test_compiler_runtime_symbols(tmp_path, monkeypatch, custom): + """A globally loaded runtime must not interpose on NVIDIA kernel dependencies.""" + monkeypatch.setattr('devito.arch.compiler.sniff_compiler_version', + lambda cc, **kwargs: Version('16.0')) + symbol = f'devito_test_runtime_{str(custom).lower()}' + vendor = f'vendor_{custom}' + sources = { + 'global': f'int {symbol}(void) {{ return 17; }}', + vendor: f'int {symbol}(void) {{ return 42; }}', + 'kernel': f'extern int {symbol}(void); int check(void) {{ return {symbol}(); }}' + } + for name, source in sources.items(): + src = tmp_path / f'{name}.c' + src.write_text(source) + args = ['gcc', '-shared', '-fPIC', str(src), '-o', + str(tmp_path / f'lib{name}.so')] + if name == 'kernel': + args += [f'-L{tmp_path}', f'-l{vendor}', f'-Wl,-rpath,{tmp_path}'] + check_call(args) + + global_runtime = ctypes.CDLL(str(tmp_path / 'libglobal.so'), mode=ctypes.RTLD_GLOBAL) + assert getattr(global_runtime, symbol)() == 17 + platform = platform_registry['cpu64-dummy']() + compiler = (CustomCompiler(base=NvidiaCompiler, platform=platform) if custom else + NvidiaCompiler(platform=platform)) + monkeypatch.setattr(compiler, 'get_jit_dir', lambda: tmp_path) + assert compiler.load('libkernel').check() == 42 + + @pytest.mark.parametrize("cc", ['gcc-4.9', 'gcc-11', 'gcc', 'gcc-14', 'gcc-123']) def test_gcc(cc): assert cc in compiler_registry