compiler: fix buffering ispace

Author: mloubout

devito/passes/clusters/buffering.py

@@ -403,8 +403,11 @@ 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
         mapper[f] = cls(name=name, dimensions=dimensions, dtype=f.dtype,
-                        grid=f.grid, halo=f.halo, space='mapped', mapped=f, f=f)
+                        padding=f.padding, grid=f.grid, halo=f.halo,
+                        space='mapped', mapped=f, f=f)
 
     return mapper
 
@@ -439,8 +442,27 @@ def __init__(self, f, b, 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)}
+        def key(c):
+            bufferdim = any(i in c.ispace.dimensions for i in self.bdims)
+            timeonly = all(d.is_Time for d in c.ispace.dimensions)
+            return bufferdim or timeonly
+
+        ispaces = set()
+        for c in clusters:
+            if not key(c):
+                continue
+
+            # Skip wild clusters (e.g. HaloTouch Clusters)
+            if c.is_wild:
+                continue
+            # Iterations space and buffering dims
+            ispace = c.ispace
+            edims = [d for d in self.bdims if d not in ispace.dimensions]
+            if not edims:
+                ispaces.add(ispace)
+            else:
+                # Add all missing buffering dimensions
+                ispaces.add(ispace.insert(self.dim, edims).reorder())
 
         if len(ispaces) > 1:
             # Best effort to make buffering work in the presence of multiple
@@ -449,6 +471,7 @@ def __init__(self, f, b, clusters):
             ispaces = {i.lift(self.bdims, v=stamp) for i in ispaces}
 
             if len(ispaces) > 1:
+                print(ispaces)
                 raise CompilationError("Unsupported `buffering` over different "
                                        "IterationSpaces")
 

devito/passes/iet/languages/C.py

@@ -27,6 +27,7 @@ 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
     # {

devito/passes/iet/languages/CXX.py

@@ -80,8 +80,9 @@ 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
-    # Transforns lhs += rhs into
+    # Transforms lhs += rhs into
     # {
     #   pragmas
     #   reinterpret_cast<float*>(&lhs)[0] += std::real(rhs);

devito/passes/iet/linearization.py

@@ -72,7 +72,21 @@ def key1(f, d):
     if f.is_regular:
         # For paddable objects the following holds:
         # `same dim + same halo + same padding_dtype => same (auto-)padding`
-        return (d, f._size_halo[d], f._size_padding[d])
+        if d is f.dimensions[-1]:
+            # Only the last dimension is padded
+            try:
+                if f.padding == f.mapped.padding:
+                    # Padding set from the mapped Function
+                    # e.g. from buffering or fft temp array
+                    pad_key = f.mapped.__padding_dtype__
+                else:
+                    pad_key = f.__padding_dtype__
+            except AttributeError:
+                pad_key = f.__padding_dtype__
+        else:
+            pad_key = None
+
+        return (d, f._size_halo[d], pad_key)
     else:
         return False
 

devito/passes/iet/parpragma.py

@@ -242,9 +242,9 @@ def _make_reductions(self, partree):
             # Implement reduction
             mapper = {partree.root: partree.root._rebuild(reduction=reductions)}
         elif all(i is OpInc for _, _, i in reductions):
-            test2 = not self._support_complex_reduction(self.compiler) and \
-                any(np.iscomplexobj(i.dtype(0)) for i, _, _ in reductions)
-            mapper = {i: self.langbb['atomic'](i, test2) for i in exprs}
+            flag = (not self._support_complex_reduction(self.compiler) and
+                    any(np.iscomplexobj(i.dtype(0)) for i, _, _ in reductions))
+            mapper = {i: self.langbb['atomic'](i, flag) for i in exprs}
         else:
             raise NotImplementedError