diff --git a/devito/dimension.py b/devito/dimension.py index ffa380c60d..f8499e395b 100644 --- a/devito/dimension.py +++ b/devito/dimension.py @@ -1,3 +1,4 @@ +import cgen from sympy import Symbol __all__ = ['Dimension', 'x', 'y', 'z', 't', 'p'] @@ -28,6 +29,16 @@ def get_varname(self): self._count += 1 return name + @property + def ccode(self): + """C-level variable name of this dimension""" + return "%s_size" % self.name if self.size is None else "%d" % self.size + + @property + def decl(self): + """Variable declaration for C-level kernel headers""" + return cgen.Value("const int", self.ccode) + # Set of default dimensions for space and time x = Dimension('x') diff --git a/devito/iteration.py b/devito/iteration.py index a87fa9ce97..d51413bf9d 100644 --- a/devito/iteration.py +++ b/devito/iteration.py @@ -12,6 +12,26 @@ __all__ = ['Iteration'] +class IterationBound(object): + """Utility class to encapsulate variable loop bounds and link them + back to the respective Dimension object. + + :param name: Variable name for the open loop bound variable + """ + + def __init__(self, name, dim): + self.name = name + self.dim = dim + + def __repr__(self): + return self.name + + @property + def ccode(self): + """C code for the variable declaration within a kernel signature""" + return cgen.Value('const int', self.name) + + class Iteration(Expression): """Iteration object that encapsualtes a single loop over sympy expressions. @@ -57,13 +77,17 @@ def __init__(self, expressions, dimension, limits, offsets=None): else: self.limits = list((0, limits, 1)) - # Adjust loop limits according to provided offsets - o_min, o_max = 0, 0 + # Replace open limits with variables names + if self.limits[1] is None: + # FIXME: Add dimension size as variable bound. + # Needs further generalisation to support loop blocking. + self.limits[1] = IterationBound("%s_size" % self.dim.name, self.dim) + + # Record offsets to later adjust loop limits accordingly + self.offsets = [0, 0] for off in (offsets or {}): - o_min = min(o_min, int(off)) - o_max = max(o_max, int(off)) - self.limits[0] += -o_min - self.limits[1] -= o_max + self.offsets[0] = min(self.offsets[0], int(off)) + self.offsets[1] = max(self.offsets[1], int(off)) def __repr__(self): str_expr = "\n\t".join([str(s) for s in self.expressions]) @@ -85,7 +109,6 @@ def ccode(self): :returns: :class:`cgen.For` object representing the loop """ - forward = self.limits[1] >= self.limits[0] loop_body = [s.ccode for s in self.expressions] if self.dim.buffered is not None: modulo = self.dim.buffered @@ -94,13 +117,10 @@ def ccode(self): for s, v in self.subs.items()] loop_body = v_subs + loop_body loop_init = cgen.InlineInitializer( - cgen.Value("int", self.index), self.limits[0]) - loop_cond = '%s %s %s' % (self.index, '<' if forward else '>', self.limits[1]) - if self.limits[2] == 1: - loop_inc = '%s%s' % (self.index, '++' if forward else '--') - else: - loop_inc = '%s %s %s' % (self.index, '+=' if forward else '-=', - self.limits[2]) + cgen.Value("int", self.index), "%d + %d" % (self.limits[0], -self.offsets[0])) + loop_cond = '%s %s %s' % (self.index, '<' if self.limits[2] >= 0 else '>', + "%s - %d" % (self.limits[1], self.offsets[1])) + loop_inc = '%s += %s' % (self.index, self.limits[2]) return cgen.For(loop_init, loop_cond, loop_inc, cgen.Block(loop_body)) @property @@ -109,8 +129,11 @@ def signature(self): :returns: List of unique data objects required by the loop """ - signatures = [e.signature for e in self.expressions] - return filter_ordered(chain(*signatures)) + signature = [e.signature for e in self.expressions] + signature = filter_ordered(chain(*signature)) + if isinstance(self.limits[1], IterationBound): + signature += [self.dim] + return signature def indexify(self): """Convert all enclosed stencil expressions to "indexed" format""" diff --git a/devito/stencilkernel.py b/devito/stencilkernel.py index dbc7cf7b2d..7782dc126c 100644 --- a/devito/stencilkernel.py +++ b/devito/stencilkernel.py @@ -1,3 +1,5 @@ +from collections import OrderedDict +from ctypes import c_int from hashlib import sha1 from itertools import chain from os import path @@ -7,6 +9,7 @@ from devito.compiler import (get_compiler_from_env, get_tmp_dir, jit_compile_and_load) +from devito.dimension import Dimension from devito.expression import Expression from devito.interfaces import SymbolicData from devito.iteration import Iteration @@ -66,17 +69,47 @@ def apply(self, *args, **kwargs): """Apply defined stencil kernel to a set of data objects""" if len(args) <= 0: args = self.signature - args = [a.data if isinstance(a, SymbolicData) else a for a in args] - # Check shape of argument data - for arg, v in zip(args, self.signature): - if not isinstance(arg, np.ndarray): - raise TypeError('No array data found for argument %s' % v.name) - if arg.shape != v.shape: - error('Expected argument %s with shape %s, but got shape %s' - % (v.name, v.shape, arg)) - raise ValueError('Argument with wrong shape') + + # Map of required arguments and actual dimension sizes + arguments = OrderedDict([(arg.name, arg) for arg in self.signature]) + dim_sizes = {} + + # Traverse positional args and infer loop sizes for open dimensions + f_args = [f for f in arguments.values() if isinstance(f, SymbolicData)] + for f, arg in zip(f_args, args): + # Ensure we're dealing or deriving numpy arrays + data = f.data if isinstance(f, SymbolicData) else arg + if not isinstance(data, np.ndarray): + error('No array data found for argument %s' % f.name) + arguments[f.name] = data + + # Ensure data dimensions match symbol dimensions + for i, dim in enumerate(f.indices): + # Infer open loop limits + if dim.size is None: + if dim.buffered: + # Check if provided as a keyword arg + size = kwargs.get(dim.name, None) + if size is None: + error("Unknown dimension size, please provide " + "size via Kernel.apply(%s=)" % dim.name) + raise RuntimeError('Dimension of unspecified size') + dim_sizes[dim] = size + elif dim in dim_sizes: + # Ensure size matches previously defined size + assert dim_sizes[dim] == data.shape[i] + else: + # Derive size from grid data shape and store + dim_sizes[dim] = data.shape[i] + else: + assert dim.size == data.shape[i] + # Insert loop size arguments from dimension values + d_args = [d for d in arguments.values() if isinstance(d, Dimension)] + for d in d_args: + arguments[d.name] = dim_sizes[d] + # Invoke kernel function with args - self.cfunction(*args) + self.cfunction(*list(arguments.values())) @property def signature(self): @@ -84,8 +117,8 @@ def signature(self): :returns: List of unique data objects required by the kernel """ - signatures = [e.signature for e in self.expressions] - return filter_ordered(chain(*signatures)) + signature = [e.signature for e in self.expressions] + return filter_ordered(chain(*signature)) @property def ccode(self): @@ -95,11 +128,13 @@ def ccode(self): and Expression objects, and adds the necessary template code around it. """ - header_vars = [c.Pointer(c.POD(v.dtype, '%s_vec' % v.name)) + header_vars = [v.decl if isinstance(v, Dimension) else + c.Pointer(c.POD(v.dtype, '%s_vec' % v.name)) for v in self.signature] header = c.FunctionDeclaration(c.Value('int', self.name), header_vars) - cast_shapes = [(v, ''.join(['[%d]' % d for d in v.shape[1:]])) - for v in self.signature] + functions = [f for f in self.signature if isinstance(f, SymbolicData)] + cast_shapes = [(f, ''.join(["[%s]" % i.ccode for i in f.indices[1:]])) + for f in functions] casts = [c.Initializer(c.POD(v.dtype, '(*%s)%s' % (v.name, shape)), '(%s (*)%s) %s' % (c.dtype_to_ctype(v.dtype), shape, '%s_vec' % v.name)) @@ -146,5 +181,6 @@ def argtypes(self): :returns: A list of ctypes of the matrix parameters and scalar parameters """ - return [np.ctypeslib.ndpointer(dtype=v.dtype, flags='C') + return [c_int if isinstance(v, Dimension) else + np.ctypeslib.ndpointer(dtype=v.dtype, flags='C') for v in self.signature] diff --git a/examples/diffusion/example_diffusion.py b/examples/diffusion/example_diffusion.py index 5ec35aa2ae..2957d3ade8 100644 --- a/examples/diffusion/example_diffusion.py +++ b/examples/diffusion/example_diffusion.py @@ -120,9 +120,6 @@ def execute_devito(ui, spacing=0.01, a=0.5, timesteps=500): # Note: This should be made simpler through the use of defaults u = TimeData(name='u', shape=(nx, ny), time_order=1, space_order=2) u.data[0, :] = ui[:] - u.indices[0].size = timesteps - u.indices[1].size = nx - u.indices[2].size = ny # Derive the stencil according to devito conventions eqn = Eq(u.dt, a * (u.dx2 + u.dy2)) @@ -131,7 +128,7 @@ def execute_devito(ui, spacing=0.01, a=0.5, timesteps=500): # Execute the generated Devito stencil operator tstart = time.time() - op.apply(u) + op.apply(u, t=timesteps) runtime = time.time() - tstart log("Devito: Diffusion with dx=%0.4f, dy=%0.4f, executed %d timesteps in %f seconds" % (spacing, spacing, timesteps, runtime)) diff --git a/tests/test_stencil_arithmetic.py b/tests/test_stencil_arithmetic.py index 614be1efdd..a70570249d 100644 --- a/tests/test_stencil_arithmetic.py +++ b/tests/test_stencil_arithmetic.py @@ -123,3 +123,20 @@ def test_arithmetic_indexed_buffered(i, j, k, l, expr, result): eqn = eval(expr) StencilKernel(eqn)(fa) assert np.allclose(fa.data[1, 1:-1, 1:-1], result[1:-1, 1:-1], rtol=1e-12) + + +@pytest.mark.parametrize('expr, result', [ + ('Eq(a[1, k, l], a[0, k - 1 , l] + 1.)', np.zeros((5, 6)) + 3.), +]) +def test_arithmetic_indexed_open_loops(i, j, k, l, expr, result): + """Test point-wise arithmetic with stencil offsets and open loop + boundaries in indexed expression format""" + k.size = None + l.size = None + a = DenseData(name='a', dimensions=(i, k, l), shape=(3, 5, 6)).indexed + fa = a.function + fa.data[0, :, :] = 2. + + eqn = eval(expr) + StencilKernel(eqn)(fa) + assert np.allclose(fa.data[1, 1:-1, 1:-1], result[1:-1, 1:-1], rtol=1e-12)