compiler: Add PyTorch tensor storage and code generation

.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 \

.github/workflows/pytest-gpu.yaml

@@ -42,8 +42,9 @@ 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"
+          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

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.

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."

devito/arch/compiler.py

@@ -1,4 +1,7 @@
+import ctypes
+import os
 import platform
+import sys
 import time
 import warnings
 from contextlib import suppress
@@ -183,6 +186,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
@@ -192,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'
 
@@ -202,10 +210,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 +259,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 +283,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"))
@@ -300,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
@@ -469,6 +487,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,15 +542,14 @@ 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,
                 # since GCC4.9 and later versions implement OpenMP 4.0, hence
                 # 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):
@@ -572,24 +592,30 @@ 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 llvmm1 and 'openmp' in language:
-                mx = platform.march
-                self.ldflags += [f'-mcpu=apple-{mx}',
-                                 '-fopenmp', f'-L{llvmm1["libs"]}']
+            if not isinstance(platform, AppleArm):
+                self.cflags += ['-march=native']
+            if llvmm1 and self.openmp:
+                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
                 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 +649,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):
@@ -665,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):
 
@@ -682,10 +710,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 +894,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 +906,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 +958,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 +975,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 +1086,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
@@ -1088,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):
     """

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')

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++20'}
+
+    @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