Merge pull request #2730 from devitocodes/locarray

misc: Miscellaneous corner cases fixes

Author: mloubout

.github/workflows/docker-bases.yml

@@ -18,16 +18,16 @@ on:
     inputs:
       cpu:
         type: boolean
-        default: false
+        default: true
       nvidia:
         type: boolean
-        default: false
+        default: true
       amd:
         type: boolean
-        default: false
+        default: true
       intel:
         type: boolean
-        default: false
+        default: true
 
       tags:
         description: "Build compiler bases"
@@ -227,7 +227,9 @@ jobs:
           push: true
           target: "nvc"
           platforms: ${{ matrix.platform }}
-          build-args: "arch=nvc"
+          build-args: |
+            arch=nvc
+            ver=nvhpc-25-7
           # Label (not tag) with runner name for traceability without changing image tags
           labels: builder-runner=${{ runner.name }}
           tags: "devitocodes/bases:nvidia-nvc-${{ matrix.arch }}"
@@ -240,7 +242,9 @@ jobs:
           push: true
           target: "nvcc"
           platforms: ${{ matrix.platform }}
-          build-args: "arch=nvcc"
+          build-args: |
+            arch=nvcc
+            ver=nvhpc-25-7
           labels: builder-runner=${{ runner.name }}
           tags: "devitocodes/bases:nvidia-nvcc-${{ matrix.arch }}"
 
@@ -252,7 +256,9 @@ jobs:
           push: true
           target: "nvc-host"
           platforms: ${{ matrix.platform }}
-          build-args: "arch=nvc-host"
+          build-args: |
+            arch=nvc-host
+            ver=nvhpc-25-7
           labels: builder-runner=${{ runner.name }}
           tags: "devitocodes/bases:cpu-nvc-${{ matrix.arch }}"
 

.github/workflows/docker-devito.yml

@@ -215,7 +215,7 @@ jobs:
             echo "No tags generated, skipping"
             exit 0
           fi
-          for tag in $TAGS; do
+          echo "$TAGS" | while read -r tag; do
             refs=""
             for arch in $ARCHES; do
               refs="$refs devitocodes/devito:${tag}-${arch}"

conftest.py

@@ -35,7 +35,7 @@ def skipif(items, whole_module=False):
     accepted.update({'device', 'device-C', 'device-openmp', 'device-openacc',
                      'device-aomp', 'cpu64-icc', 'cpu64-icx', 'cpu64-nvc',
                      'noadvisor', 'cpu64-arm', 'cpu64-icpx', 'chkpnt'})
-    accepted.update({'nodevice'})
+    accepted.update({'nodevice', 'noomp'})
     unknown = sorted(set(items) - accepted)
     if unknown:
         raise ValueError("Illegal skipif argument(s) `%s`" % unknown)
@@ -86,6 +86,10 @@ def skipif(items, whole_module=False):
              not get_advisor_path()):
             skipit = "Only `icx+advisor` should be tested here"
             break
+        # Slip if not using openmp
+        if i == 'noomp' and 'openmp' not in configuration['language']:
+            skipit = "Must use openmp"
+            break
         # Skip if it won't run on Arm
         if i == 'cpu64-arm' and isinstance(configuration['platform'], Arm):
             skipit = "Arm doesn't support x86-specific instructions"

devito/arch/archinfo.py

@@ -577,7 +577,7 @@ def get_advisor_path():
 @memoized_func
 def get_hip_path():
     # *** First try: via commonly used environment variables
-    for i in ['HIP_HOME']:
+    for i in ['HIP_HOME', 'ROCM_HOME']:
         hip_home = os.environ.get(i)
         if hip_home:
             return hip_home

devito/ir/cgen/printer.py

@@ -349,7 +349,8 @@ def _print_ListInitializer(self, expr):
         return f"{{{', '.join(self._print(i) for i in expr.params)}}}"
 
     def _print_IndexedPointer(self, expr):
-        return f"{expr.base}{''.join(f'[{self._print(i)}]' for i in expr.index)}"
+        base = self._print(expr.base)
+        return f"{base}{''.join(f'[{self._print(i)}]' for i in expr.index)}"
 
     def _print_IntDiv(self, expr):
         lhs = self._print(expr.lhs)

devito/ir/iet/visitors.py

@@ -353,7 +353,10 @@ def _gen_rettype(self, obj):
         elif isinstance(obj, (FieldFromComposite, FieldFromPointer)):
             return self._gen_value(obj.function.base, 0).typename
         else:
-            return None
+            try:
+                return obj._type_.__name__
+            except AttributeError:
+                return None
 
     def _args_decl(self, args):
         """Generate cgen declarations from an iterable of symbols and expressions."""

devito/mpi/distributed.py

@@ -40,7 +40,7 @@
     def cleanup():
         devito_mpi_finalize()
     atexit.register(cleanup)
-except ImportError as e:
+except (RuntimeError, ImportError) as e:
     # Dummy fallback in case mpi4py/MPI aren't available
     class NoneMetaclass(type):
         def __getattr__(self, name):

devito/mpi/routines.py

@@ -1395,7 +1395,7 @@ def _arg_values(self, args=None, **kwargs):
 class AllreduceCall(Call):
 
     def __init__(self, arguments, **kwargs):
-        super().__init__('MPI_Allreduce', arguments)
+        super().__init__('MPI_Allreduce', arguments, **kwargs)
 
 
 class ReductionBuilder(object):
@@ -1422,6 +1422,6 @@ def make(self, dr):
         op = self.mapper[dr.op]
 
         arguments = [inplace, Byref(f), Integer(1), mpitype, op, comm]
-        allreduce = AllreduceCall(arguments)
+        allreduce = AllreduceCall(arguments, writes=f)
 
         return allreduce

devito/passes/clusters/buffering.py

@@ -403,8 +403,13 @@ def generate_buffers(clusters, key, sregistry, options, **kwargs):
         # Finally create the actual buffer
         cls = callback or Array
         name = sregistry.make_name(prefix='%sb' % f.name)
+        # We specify the padding to match the input Function's one, so that
+        # the array can be used in place of the Function with valid strides
+        # Plain Array do not track mapped so we default to no padding
+        padding = 0 if cls is Array else f.padding
         mapper[f] = cls(name=name, dimensions=dimensions, dtype=f.dtype,
-                        grid=f.grid, halo=f.halo, space='mapped', mapped=f, f=f)
+                        padding=padding, grid=f.grid, halo=f.halo,
+                        space='mapped', mapped=f, f=f)
 
     return mapper
 
@@ -436,25 +441,6 @@ def __init__(self, f, b, clusters):
 
         self.indices = extract_indices(f, self.dim, clusters)
 
-        # The IterationSpace within which the buffer will be accessed
-        # NOTE: The `key` is to avoid Clusters including `f` but not directly
-        # using it in an expression, such as HaloTouch Clusters
-        key = lambda c: any(i in c.ispace.dimensions for i in self.bdims)
-        ispaces = {c.ispace for c in clusters if key(c)}
-
-        if len(ispaces) > 1:
-            # Best effort to make buffering work in the presence of multiple
-            # IterationSpaces
-            stamp = Stamp()
-            ispaces = {i.lift(self.bdims, v=stamp) for i in ispaces}
-
-            if len(ispaces) > 1:
-                raise CompilationError("Unsupported `buffering` over different "
-                                       "IterationSpaces")
-
-        assert len(ispaces) == 1, "Unexpected form of `buffering`"
-        self.ispace = ispaces.pop()
-
     def __repr__(self):
         return "Descriptor[%s -> %s]" % (self.f, self.b)
 
@@ -474,6 +460,47 @@ def itdims(self):
         """
         return (self.xd, self.dim.root)
 
+    @cached_property
+    def ispace(self):
+        # The IterationSpace within which the buffer will be accessed
+
+        # NOTE: The `key` is to avoid Clusters including `f` but not directly
+        # using it in an expression, such as HaloTouch Clusters
+        def key(c):
+            bufferdim = any(i in c.ispace.dimensions for i in self.bdims)
+            xd_only = all(d._defines & self.xd._defines for d in c.ispace.dimensions)
+            return bufferdim or xd_only
+
+        ispaces = set()
+        for c in self.clusters:
+            if not key(c):
+                continue
+
+            # Skip wild clusters (e.g. HaloTouch Clusters)
+            if c.is_wild:
+                continue
+
+            # Iterations space and buffering dims
+            edims = [d for d in self.bdims if d not in c.ispace.dimensions]
+            if not edims:
+                ispaces.add(c.ispace)
+            else:
+                # Add all missing buffering dimensions and reorder to
+                # avoid duplicates with different ordering
+                ispaces.add(c.ispace.insert(self.dim, edims).reorder())
+
+        if len(ispaces) > 1:
+            # Best effort to make buffering work in the presence of multiple
+            # IterationSpaces
+            stamp = Stamp()
+            ispaces = {i.lift(self.bdims, v=stamp) for i in ispaces}
+
+            if len(ispaces) > 1:
+                raise CompilationError("Unsupported `buffering` over different "
+                                       "IterationSpaces")
+
+        return ispaces.pop()
+
     @cached_property
     def subdims_mapper(self):
         return {d.root: d for d in self.ispace.itdims if d.is_AbstractSub}

devito/passes/iet/languages/C.py

@@ -5,12 +5,40 @@
 from devito.passes.iet.definitions import DataManager
 from devito.passes.iet.orchestration import Orchestrator
 from devito.passes.iet.langbase import LangBB
+from devito.passes.iet.languages.utils import _atomic_add_split
 from devito.symbolics import c_complex, c_double_complex
+from devito.symbolics.extended_sympy import UnaryOp
 from devito.tools import dtype_to_cstr
 
 __all__ = ['CBB', 'CDataManager', 'COrchestrator']
 
 
+class RealExt(UnaryOp):
+
+    _op = '__real__ '
+
+
+class ImagExt(UnaryOp):
+
+    _op = '__imag__ '
+
+
+def atomic_add(i, pragmas, split=False):
+    # Base case, real reduction
+    if not split:
+        return i._rebuild(pragmas=pragmas)
+
+    # Complex reduction, split using a temp pointer
+    # Transforms lhs += rhs into
+    # {
+    #   pragmas
+    #   __real__ lhs += __real__ rhs;
+    #   pragmas
+    #   __imag__ lhs += __imag__ rhs;
+    # }
+    return _atomic_add_split(i, pragmas, RealExt, ImagExt)
+
+
 class CBB(LangBB):
 
     mapper = {
@@ -29,7 +57,7 @@ class CBB(LangBB):
         'host-free-pin': lambda i:
             Call('free', (i,)),
         'alloc-global-symbol': lambda i, j, k:
-            Call('memcpy', (i, j, k)),
+            Call('memcpy', (i, j, k))
     }