diff --git a/devito/passes/clusters/buffering.py b/devito/passes/clusters/buffering.py index 5f72b8d746..ba6f30a1ce 100644 --- a/devito/passes/clusters/buffering.py +++ b/devito/passes/clusters/buffering.py @@ -176,7 +176,12 @@ def callback(self, clusters, prefix): accesses = chain(*[c.scope[v.f] for c in v.clusters]) index_mapper = {i: mds[(v.xd, i)] for i in v.indices} for a in accesses: - subs[a.access] = b.indexed[[index_mapper.get(i, i) for i in a]] + indices = [index_mapper.get(i, i) for i in a] + try: + indices = b._buffer_indices(indices) + except AttributeError: + pass + subs[a.access] = b.indexed[indices] processed = [] for c in clusters: @@ -431,7 +436,7 @@ def generate_buffers(clusters, key, sregistry, options, **kwargs): buffer, = buffers xd = buffer.indices[dim] # The new buffer is derived from `buffer`, so it inherits its padding policy - extra_kwargs = {'is_autopaddable': buffer.is_autopaddable} + extra_kwargs = {'is_autopaddable': buffer.is_autopaddable, 'buffer': buffer} else: size = infer_buffer_size(f, dim, clusters) @@ -459,6 +464,7 @@ def generate_buffers(clusters, key, sregistry, options, **kwargs): dimensions[dimensions.index(dim)] = xd # Finally create the actual buffer + ## Better refactor: stash callback on original `f` cls = callback or Array name = sregistry.make_name(prefix=f'{f.name}b') mapper[f] = cls(name=name, dimensions=dimensions, dtype=f.dtype,