From c0a858996e002526b8ba1e4274b2ec4a1da2d26b Mon Sep 17 00:00:00 2001
From: Andreas Kloeckner <inform@tiker.net>
Date: Tue, 30 Apr 2013 22:10:00 -0400
Subject: [PATCH] Add automatic argument shapes, default_order, and move more
 front-end fucntionality to make_kernel.

---
 examples/hello-loopy.py     |  17 +-
 loopy/__init__.py           |   5 +-
 loopy/codegen/expression.py |  16 +-
 loopy/kernel/__init__.py    |  64 ++-----
 loopy/kernel/creation.py    | 325 ++++++++++++++++++++++++++----------
 loopy/kernel/data.py        |  49 ++++--
 loopy/symbolic.py           |   6 +-
 test/test_loopy.py          |  58 ++++---
 8 files changed, 344 insertions(+), 196 deletions(-)

diff --git a/examples/hello-loopy.py b/examples/hello-loopy.py
index a6b8e6d8b..a35d7272d 100644
--- a/examples/hello-loopy.py
+++ b/examples/hello-loopy.py
@@ -15,14 +15,15 @@ a = cl.array.arange(queue, n, dtype=np.float32)
 # -----------------------------------------------------------------------------
 # generation (loopy bits start here)
 # -----------------------------------------------------------------------------
-knl = lp.make_kernel(ctx.devices[0],
-    "{[i]: 0<=i<n}",   # "loop domain"-- what values does i take?
-    "out[i] = 2*a[i]", # "instructions" to be executed across the domain
-    [ # argument declarations
-        lp.GlobalArg("out", shape="n"),
-        lp.GlobalArg("a", shape="n"),
-        lp.ValueArg("n"),
-        ])
+knl = lp.make_kernel(
+        ctx.devices[0],
+        "{ [i]: 0<=i<n }",
+        "out[i] = 2*a[i]",
+        [ # argument declarations
+            lp.GlobalArg("out"),
+            lp.GlobalArg("a"),
+            lp.ValueArg("n"),
+            ])
 
 # -----------------------------------------------------------------------------
 # transformation
diff --git a/loopy/__init__.py b/loopy/__init__.py
index 41ce634c4..7a1ac8d26 100644
--- a/loopy/__init__.py
+++ b/loopy/__init__.py
@@ -50,6 +50,7 @@ class LoopyAdvisory(UserWarning):
 # {{{ imported user interface
 
 from loopy.kernel.data import (
+        auto_shape, auto_strides,
         ValueArg, ScalarArg, GlobalArg, ArrayArg, ConstantArg, ImageArg,
 
         default_function_mangler, single_arg_function_mangler, opencl_function_mangler,
@@ -72,7 +73,9 @@ from loopy.codegen import generate_code
 from loopy.compiled import CompiledKernel, auto_test_vs_ref
 from loopy.check import check_kernels
 
-__all__ = ["ValueArg", "ScalarArg", "GlobalArg", "ArrayArg", "ConstantArg", "ImageArg",
+__all__ = [
+        "auto_shape", "auto_strides",
+        "ValueArg", "ScalarArg", "GlobalArg", "ArrayArg", "ConstantArg", "ImageArg",
         "LoopKernel",
         "Instruction",
         "default_function_mangler", "single_arg_function_mangler",
diff --git a/loopy/codegen/expression.py b/loopy/codegen/expression.py
index 902ed7ef2..f8b3d6775 100644
--- a/loopy/codegen/expression.py
+++ b/loopy/codegen/expression.py
@@ -379,25 +379,21 @@ class LoopyCCodeMapper(RecursiveMapper):
                 if not isinstance(expr.index, tuple):
                     index_expr = (index_expr,)
 
-                if arg.strides is not None:
-                    ary_strides = arg.strides
-                else:
-                    ary_strides = (1,)
+                if arg.strides is None:
+                    raise RuntimeError("index access to '%s' requires known "
+                            "strides" % arg.name)
 
-                if len(ary_strides) != len(index_expr):
+                if len(arg.strides) != len(index_expr):
                     raise RuntimeError("subscript to '%s' in '%s' has the wrong "
                             "number of indices (got: %d, expected: %d)" % (
                                 expr.aggregate.name, expr,
-                                len(index_expr), len(ary_strides)))
-
-                if len(index_expr) == 0:
-                    return "*" + expr.aggregate.name
+                                len(index_expr), len(arg.strides)))
 
                 from pymbolic.primitives import Subscript
                 return base_impl(
                         Subscript(expr.aggregate, arg.offset+sum(
                             stride*expr_i for stride, expr_i in zip(
-                                ary_strides, index_expr))),
+                                arg.strides, index_expr))),
                         enclosing_prec, type_context)
 
 
diff --git a/loopy/kernel/__init__.py b/loopy/kernel/__init__.py
index b07a487fc..2cd1bce37 100644
--- a/loopy/kernel/__init__.py
+++ b/loopy/kernel/__init__.py
@@ -54,46 +54,28 @@ class CannotBranchDomainTree(RuntimeError):
 # {{{ loop kernel object
 
 class LoopKernel(Record):
-    """
+    """These correspond more or less directly to arguments of
+    :func:`loopy.make_kernel`.
+
     :ivar device: :class:`pyopencl.Device`
-    :ivar domains: :class:`islpy.BasicSet`
+    :ivar domains: a list of :class:`islpy.BasicSet` instances
     :ivar instructions:
     :ivar args:
     :ivar schedule:
     :ivar name:
-    :ivar preambles: a list of (tag, code) tuples that identify preamble snippets.
-        Each tag's snippet is only included once, at its first occurrence.
-        The preambles will be inserted in order of their tags.
-    :ivar preamble_generators: a list of functions of signature
-        (seen_dtypes, seen_functions) where seen_functions is a set of
-        (name, c_name, arg_dtypes), generating extra entries for `preambles`.
-    :ivar assumptions: the initial implemented_domain, captures assumptions
-        on the parameters. (an isl.Set)
-    :ivar local_sizes: A dictionary from integers to integers, mapping
-        workgroup axes to their sizes, e.g. *{0: 16}* forces axis 0 to be
-        length 16.
+    :ivar preambles:
+    :ivar preamble_generators:
+    :ivar assumptions:
+    :ivar local_sizes:
     :ivar temporary_variables:
     :ivar iname_to_tag:
-    :ivar substitutions: a mapping from substitution names to :class:`SubstitutionRule`
-        objects
-    :ivar function_manglers: list of functions of signature (name, arg_dtypes)
-        returning a tuple (result_dtype, c_name)
-        or a tuple (result_dtype, c_name, arg_dtypes),
-        where c_name is the C-level function to be called.
-    :ivar symbol_manglers: list of functions of signature (name) returning
-        a tuple (result_dtype, c_name), where c_name is the C-level symbol to be
-        evaluated.
-    :ivar defines: a dictionary of replacements to be made in instructions given
-        as strings before parsing. A macro instance intended to be replaced should
-        look like "MACRO" in the instruction code. The expansion given in this
-        parameter is allowed to be a list. In this case, instructions are generated
-        for *each* combination of macro values.
-
-        These defines may also be used in the domain and in argument shapes and
-        strides. They are expanded only upon kernel creation.
+    :ivar function_manglers:
+    :ivar symbol_manglers:
 
     The following arguments are not user-facing:
 
+    :ivar substitutions: a mapping from substitution names to :class:`SubstitutionRule`
+        objects
     :ivar iname_slab_increments: a dictionary mapping inames to (lower_incr,
         upper_incr) tuples that will be separated out in the execution to generate
         'bulk' slabs with fewer conditionals.
@@ -121,7 +103,6 @@ class LoopKernel(Record):
                 single_arg_function_mangler,
                 ],
             symbol_manglers=[opencl_symbol_mangler],
-            defines={},
 
             # non-user-facing
             iname_slab_increments={},
@@ -203,25 +184,6 @@ class LoopKernel(Record):
 
         # }}}
 
-        # {{{ expand macros in arg shapes
-
-        from loopy.kernel.data import ShapedArg
-        from loopy.kernel.creation import expand_defines_in_expr
-
-        processed_args = []
-        for arg in args:
-            for arg_name in arg.name.split(","):
-                new_arg = arg.copy(name=arg_name)
-                if isinstance(arg, ShapedArg):
-                    if arg.shape is not None:
-                        new_arg = new_arg.copy(shape=expand_defines_in_expr(arg.shape, defines))
-                    if arg.strides is not None:
-                        new_arg = new_arg.copy(strides=expand_defines_in_expr(arg.strides, defines))
-
-                processed_args.append(new_arg)
-
-        # }}}
-
         index_dtype = np.dtype(index_dtype)
         if index_dtype.kind != 'i':
             raise TypeError("index_dtype must be an integer")
@@ -235,7 +197,7 @@ class LoopKernel(Record):
         Record.__init__(self,
                 device=device, domains=domains,
                 instructions=instructions,
-                args=processed_args,
+                args=args,
                 schedule=schedule,
                 name=name,
                 preambles=preambles,
diff --git a/loopy/kernel/creation.py b/loopy/kernel/creation.py
index 3b706c80f..94bc195a4 100644
--- a/loopy/kernel/creation.py
+++ b/loopy/kernel/creation.py
@@ -28,7 +28,7 @@ THE SOFTWARE.
 
 
 import numpy as np
-from loopy.symbolic import IdentityMapper
+from loopy.symbolic import IdentityMapper, WalkMapper
 from loopy.kernel.data import Instruction, SubstitutionRule
 import islpy as isl
 from islpy import dim_type
@@ -85,76 +85,6 @@ class MakeUnique:
 
 # }}}
 
-# {{{ domain parsing
-
-EMPTY_SET_DIMS_RE = re.compile(r"^\s*\{\s*\:")
-SET_DIMS_RE = re.compile(r"^\s*\{\s*\[([a-zA-Z0-9_, ]+)\]\s*\:")
-
-def _find_inames_in_set(dom_str):
-    empty_match = EMPTY_SET_DIMS_RE.match(dom_str)
-    if empty_match is not None:
-        return set()
-
-    match = SET_DIMS_RE.match(dom_str)
-    if match is None:
-        raise RuntimeError("invalid syntax for domain '%s'" % dom_str)
-
-    result = set(iname.strip() for iname in match.group(1).split(",")
-            if iname.strip())
-
-    return result
-
-EX_QUANT_RE = re.compile(r"\bexists\s+([a-zA-Z0-9])\s*\:")
-
-def _find_existentially_quantified_inames(dom_str):
-    return set(ex_quant.group(1) for ex_quant in EX_QUANT_RE.finditer(dom_str))
-
-def parse_domains(ctx, domains, defines):
-    if isinstance(domains, str):
-        domains = [domains]
-
-    result = []
-    used_inames = set()
-
-    for dom in domains:
-        if isinstance(dom, str):
-            dom, = expand_defines(dom, defines)
-
-            if not dom.lstrip().startswith("["):
-                # i.e. if no parameters are already given
-                parameters = (_gather_isl_identifiers(dom)
-                        - _find_inames_in_set(dom)
-                        - _find_existentially_quantified_inames(dom))
-                dom = "[%s] -> %s" % (",".join(parameters), dom)
-
-            try:
-                dom = isl.BasicSet.read_from_str(ctx, dom)
-            except:
-                print "failed to parse domain '%s'" % dom
-                raise
-        else:
-            assert isinstance(dom, (isl.Set, isl.BasicSet))
-            # assert dom.get_ctx() == ctx
-
-        for i_iname in xrange(dom.dim(dim_type.set)):
-            iname = dom.get_dim_name(dim_type.set, i_iname)
-
-            if iname is None:
-                raise RuntimeError("domain '%s' provided no iname at index "
-                        "%d (redefined iname?)" % (dom, i_iname))
-
-            if iname in used_inames:
-                raise RuntimeError("domain '%s' redefines iname '%s' "
-                        "that is part of a previous domain" % (dom, iname))
-
-            used_inames.add(iname)
-
-        result.append(dom)
-
-    return result
-
-# }}}
-
 # {{{ expand defines
 
 WORD_RE = re.compile(r"\b([a-zA-Z0-9_]+)\b")
@@ -346,6 +276,79 @@ def parse_if_necessary(insn, defines):
 
 # }}}
 
+# {{{ domain parsing
+
+EMPTY_SET_DIMS_RE = re.compile(r"^\s*\{\s*\:")
+SET_DIMS_RE = re.compile(r"^\s*\{\s*\[([a-zA-Z0-9_, ]+)\]\s*\:")
+
+def _find_inames_in_set(dom_str):
+    empty_match = EMPTY_SET_DIMS_RE.match(dom_str)
+    if empty_match is not None:
+        return set()
+
+    match = SET_DIMS_RE.match(dom_str)
+    if match is None:
+        raise RuntimeError("invalid syntax for domain '%s'" % dom_str)
+
+    result = set(iname.strip() for iname in match.group(1).split(",")
+            if iname.strip())
+
+    return result
+
+EX_QUANT_RE = re.compile(r"\bexists\s+([a-zA-Z0-9])\s*\:")
+
+def _find_existentially_quantified_inames(dom_str):
+    return set(ex_quant.group(1) for ex_quant in EX_QUANT_RE.finditer(dom_str))
+
+def parse_domains(ctx, domains, defines):
+    if isinstance(domains, str):
+        domains = [domains]
+
+    result = []
+    used_inames = set()
+
+    for dom in domains:
+        if isinstance(dom, str):
+            dom, = expand_defines(dom, defines)
+
+            if not dom.lstrip().startswith("["):
+                # i.e. if no parameters are already given
+                parameters = (_gather_isl_identifiers(dom)
+                        - _find_inames_in_set(dom)
+                        - _find_existentially_quantified_inames(dom))
+                dom = "[%s] -> %s" % (",".join(parameters), dom)
+
+            try:
+                dom = isl.BasicSet.read_from_str(ctx, dom)
+            except:
+                print "failed to parse domain '%s'" % dom
+                raise
+        else:
+            assert isinstance(dom, (isl.Set, isl.BasicSet))
+            # assert dom.get_ctx() == ctx
+
+        for i_iname in xrange(dom.dim(dim_type.set)):
+            iname = dom.get_dim_name(dim_type.set, i_iname)
+
+            if iname is None:
+                raise RuntimeError("domain '%s' provided no iname at index "
+                        "%d (redefined iname?)" % (dom, i_iname))
+
+            if iname in used_inames:
+                raise RuntimeError("domain '%s' redefines iname '%s' "
+                        "that is part of a previous domain" % (dom, iname))
+
+            used_inames.add(iname)
+
+        result.append(dom)
+
+    return result
+
+# }}}
+
+def guess_kernel_args_if_requested(domains, instructions, kernel_args):
+    return kernel_args
+
 # {{{ tag reduction inames as sequential
 
 def tag_reduction_inames_as_sequential(knl):
@@ -592,23 +595,168 @@ def check_for_reduction_inames_duplication_requests(kernel):
 
 # }}}
 
-# {{{ kernel creation top-level
+# {{{
 
-def make_kernel(device, domains, instructions, kernel_args=[], *args, **kwargs):
-    """User-facing kernel creation entrypoint."""
+def apply_default_order_to_args(kernel, default_order):
+    from loopy.kernel.data import ShapedArg
 
-    for forbidden_kwarg in [
-            "substitutions",
-            "iname_slab_increments",
-            "applied_iname_rewrites",
-            "cache_manager",
-            "isl_context",
-            ]:
-        if forbidden_kwarg in kwargs:
-            raise RuntimeError("'%s' is not part of user-facing interface"
-                    % forbidden_kwarg)
+    processed_args = []
+    for arg in kernel.args:
+        if isinstance(arg, ShapedArg):
+            arg = arg.copy(order=default_order)
+        processed_args.append(arg)
+
+    return kernel.copy(args=processed_args)
+
+# }}}
+
+# {{{ duplicate arguments and expand defines in shapes
+
+def dup_args_and_expand_defines_in_shapes(kernel, defines):
+    from loopy.kernel.data import ShapedArg, auto_shape, auto_strides
+    from loopy.kernel.creation import expand_defines_in_expr
+
+    processed_args = []
+    for arg in kernel.args:
+        for arg_name in arg.name.split(","):
+            new_arg = arg.copy(name=arg_name)
+            if isinstance(arg, ShapedArg):
+                if arg.shape is not None and arg.shape is not auto_shape:
+                    new_arg = new_arg.copy(shape=expand_defines_in_expr(arg.shape, defines))
+                if arg.strides is not None and arg.strides is not auto_strides:
+                    new_arg = new_arg.copy(strides=expand_defines_in_expr(arg.strides, defines))
+
+            processed_args.append(new_arg)
+
+    return kernel.copy(args=processed_args)
+
+# }}}
+
+# {{{ guess argument shapes
+
+class _AccessRangeMapper(WalkMapper):
+    def __init__(self, arg_name):
+        self.arg_name = arg_name
+        self.access_range = None
+
+    def map_subscript(self, expr, domain):
+        WalkMapper.map_subscript(self, expr, domain)
+
+        from pymbolic.primitives import Variable
+        assert isinstance(expr.aggregate, Variable)
+
+        if expr.aggregate.name != self.arg_name:
+            return
+
+        subscript = expr.index
+        if not isinstance(subscript, tuple):
+            subscript = (subscript,)
+
+        from loopy.symbolic import get_dependencies, get_access_range
 
-    defines = kwargs.get("defines", {})
+        if not get_dependencies(subscript) <= set(domain.get_var_dict()):
+            raise RuntimeError("cannot determine access range for '%s': "
+                    "undetermined index in '%s'"
+                    % (self.arg_name, ", ".join(str(i) for i in subscript)))
+
+        access_range = get_access_range(domain, subscript)
+
+        if self.access_range is None:
+            self.access_range = access_range
+        else:
+            if (self.access_range.dim(dim_type.set)
+                    != access_range.dim(dim_type.set)):
+                raise RuntimeError(
+                        "error while determining shape of argument '%s': "
+                        "varying number of indices encountered"
+                        % self.arg_name)
+
+            self.access_range = self.access_range | access_range
+
+def guess_arg_shape_if_requested(kernel, default_order):
+    new_args = []
+
+    from loopy.kernel.data import ShapedArg, auto_shape, auto_strides
+
+    for arg in kernel.args:
+        if isinstance(arg, ShapedArg) and (
+                arg.shape is auto_shape or arg.strides is auto_strides):
+            armap = _AccessRangeMapper(arg.name)
+
+            for insn in kernel.instructions:
+                domain = kernel.get_inames_domain(kernel.insn_inames(insn))
+                armap(insn.assignee, domain)
+                armap(insn.expression, domain)
+
+            if armap.access_range is None:
+                # no subscripts found, let's call it a scalar
+                shape = ()
+            else:
+                from loopy.isl_helpers import static_max_of_pw_aff
+                from loopy.symbolic import pw_aff_to_expr
+
+                shape = tuple(
+                        pw_aff_to_expr(static_max_of_pw_aff(
+                            kernel.cache_manager.dim_max(armap.access_range, i) + 1,
+                            constants_only=False))
+                        for i in xrange(armap.access_range.dim(dim_type.set)))
+
+            if arg.shape is auto_shape:
+                arg = arg.copy(shape=shape)
+            if arg.strides is auto_strides:
+                from loopy.kernel.data import make_strides
+                arg = arg.copy(strides=make_strides(shape, default_order))
+
+        new_args.append(arg)
+
+    return kernel.copy(args=new_args)
+
+# }}}
+
+# {{{ kernel creation top-level
+
+def make_kernel(device, domains, instructions, kernel_args=[], **kwargs):
+    """User-facing kernel creation entrypoint.
+
+    :arg device: :class:`pyopencl.Device`
+    :arg domains: :class:`islpy.BasicSet`
+    :arg instructions:
+    :arg kernel_args:
+
+    The following keyword arguments are recognized:
+
+    :arg preambles: a list of (tag, code) tuples that identify preamble snippets.
+        Each tag's snippet is only included once, at its first occurrence.
+        The preambles will be inserted in order of their tags.
+    :arg preamble_generators: a list of functions of signature
+        (seen_dtypes, seen_functions) where seen_functions is a set of
+        (name, c_name, arg_dtypes), generating extra entries for *preambles*.
+    :arg defines: a dictionary of replacements to be made in instructions given
+        as strings before parsing. A macro instance intended to be replaced should
+        look like "MACRO" in the instruction code. The expansion given in this
+        parameter is allowed to be a list. In this case, instructions are generated
+        for *each* combination of macro values.
+
+        These defines may also be used in the domain and in argument shapes and
+        strides. They are expanded only upon kernel creation.
+    :arg default_order: "C" (default) or "F"
+    :arg function_manglers: list of functions of signature (name, arg_dtypes)
+        returning a tuple (result_dtype, c_name)
+        or a tuple (result_dtype, c_name, arg_dtypes),
+        where c_name is the C-level function to be called.
+    :arg symbol_manglers: list of functions of signature (name) returning
+        a tuple (result_dtype, c_name), where c_name is the C-level symbol to be
+        evaluated.
+    :arg assumptions: the initial implemented_domain, captures assumptions
+        on the parameters. (an isl.Set)
+    :arg local_sizes: A dictionary from integers to integers, mapping
+        workgroup axes to their sizes, e.g. *{0: 16}* forces axis 0 to be
+        length 16.
+    :arg temporary_variables:
+    """
+
+    defines = kwargs.pop("defines", {})
+    default_order = kwargs.pop("default_order", "C")
 
     # {{{ instruction/subst parsing
 
@@ -645,8 +793,10 @@ def make_kernel(device, domains, instructions, kernel_args=[], *args, **kwargs):
 
     domains = parse_domains(isl_context, domains, defines)
 
+    kernel_args = guess_kernel_args_if_requested(domains, instructions, kernel_args)
+
     from loopy.kernel import LoopKernel
-    knl = LoopKernel(device, domains, instructions, kernel_args, *args, **kwargs)
+    knl = LoopKernel(device, domains, instructions, kernel_args, **kwargs)
 
     check_for_nonexistent_iname_deps(knl)
     check_for_reduction_inames_duplication_requests(knl)
@@ -654,6 +804,9 @@ def make_kernel(device, domains, instructions, kernel_args=[], *args, **kwargs):
     knl = tag_reduction_inames_as_sequential(knl)
     knl = create_temporaries(knl)
     knl = expand_cses(knl)
+    knl = dup_args_and_expand_defines_in_shapes(knl, defines)
+    knl = guess_arg_shape_if_requested(knl, default_order)
+    knl = apply_default_order_to_args(knl, default_order)
 
     # -------------------------------------------------------------------------
     # Ordering dependency:
diff --git a/loopy/kernel/data.py b/loopy/kernel/data.py
index 80e53b17d..8248b4ca9 100644
--- a/loopy/kernel/data.py
+++ b/loopy/kernel/data.py
@@ -136,8 +136,26 @@ def parse_tag(tag):
 
 # {{{ arguments
 
+class auto_shape:
+    pass
+
+class auto_strides:
+    pass
+
+def make_strides(shape, order):
+    from pyopencl.compyte.array import (
+            f_contiguous_strides,
+            c_contiguous_strides)
+
+    if order == "F":
+        return f_contiguous_strides(1, shape)
+    elif order == "C":
+        return c_contiguous_strides(1, shape)
+    else:
+        raise ValueError("invalid order: %s" % order)
+
 class ShapedArg(Record):
-    def __init__(self, name, dtype=None, shape=None, strides=None, order="C",
+    def __init__(self, name, dtype=None, shape=None, strides=None, order=None,
             offset=0):
         """
         All of the following are optional. Specify either strides or shape.
@@ -166,23 +184,26 @@ class ShapedArg(Record):
 
             return tuple(parse_if_necessary(xi) for xi in x)
 
-        if strides is not None:
+        if strides == "auto":
+            strides = auto_strides
+        if shape == "auto":
+            shape = auto_shape
+
+        strides_known = strides is not None and strides is not auto_strides
+        shape_known = shape is not None and shape is not auto_shape
+
+        if strides_known:
             strides = process_tuple(strides)
 
-        if shape is not None:
+        if shape_known:
             shape = process_tuple(shape)
 
-        if strides is None and shape is not None:
-            from pyopencl.compyte.array import (
-                    f_contiguous_strides,
-                    c_contiguous_strides)
-
-            if order == "F":
-                strides = f_contiguous_strides(1, shape)
-            elif order == "C":
-                strides = c_contiguous_strides(1, shape)
-            else:
-                raise ValueError("invalid order: %s" % order)
+        if not strides_known and shape_known:
+            if len(shape) == 1:
+                # don't need order to know that
+                strides = (1,)
+            elif order is not None:
+                strides = make_strides(shape, order)
 
         Record.__init__(self,
                 name=name,
diff --git a/loopy/symbolic.py b/loopy/symbolic.py
index b2a06cda6..91f7e4907 100644
--- a/loopy/symbolic.py
+++ b/loopy/symbolic.py
@@ -168,15 +168,15 @@ class IdentityMapper(IdentityMapperBase, IdentityMapperMixin):
     pass
 
 class WalkMapper(WalkMapperBase):
-    def map_reduction(self, expr):
+    def map_reduction(self, expr, *args):
         if not self.visit(expr):
             return
 
-        self.rec(expr.expr)
+        self.rec(expr.expr, *args)
 
     map_tagged_variable = WalkMapperBase.map_variable
 
-    def map_loopy_function_identifier(self, expr):
+    def map_loopy_function_identifier(self, expr, *args):
         self.visit(expr)
 
     map_linear_subscript = WalkMapperBase.map_subscript
diff --git a/test/test_loopy.py b/test/test_loopy.py
index 3694b0390..8c034d6f5 100644
--- a/test/test_loopy.py
+++ b/test/test_loopy.py
@@ -202,14 +202,10 @@ def test_wg_too_small(ctx_factory):
     kernel_gen = lp.generate_loop_schedules(knl)
     kernel_gen = lp.check_kernels(kernel_gen)
 
+    import pytest
     for gen_knl in kernel_gen:
-        try:
+        with pytest.raises(RuntimeError):
             lp.CompiledKernel(ctx, gen_knl).get_code()
-        except RuntimeError, e:
-            assert "implemented and desired" in str(e)
-            pass # expected!
-        else:
-            assert False # expecting an error
 
 
 
@@ -644,14 +640,13 @@ def test_dependent_loop_bounds(ctx_factory):
                 ],
             [
                 "<> row_len = a_rowstarts[i+1] - a_rowstarts[i]",
-                "ax[i] = sum(jj, a_values[a_rowstarts[i]+jj])",
+                "a_sum[i] = sum(jj, a_values[[a_rowstarts[i]+jj]])",
                 ],
             [
-                lp.GlobalArg("a_rowstarts", np.int32),
-                lp.GlobalArg("a_indices", np.int32),
+                lp.GlobalArg("a_rowstarts", np.int32, shape="auto"),
+                lp.GlobalArg("a_indices", np.int32, shape="auto"),
                 lp.GlobalArg("a_values", dtype),
-                lp.GlobalArg("x", dtype),
-                lp.GlobalArg("ax", dtype),
+                lp.GlobalArg("a_sum", dtype, shape="auto"),
                 lp.ValueArg("n", np.int32),
                 ],
             assumptions="n>=1 and row_len>=1")
@@ -676,14 +671,13 @@ def test_dependent_loop_bounds_2(ctx_factory):
             [
                 "<> row_start = a_rowstarts[i]",
                 "<> row_len = a_rowstarts[i+1] - row_start",
-                "ax[i] = sum(jj, a_values[row_start+jj])",
+                "ax[i] = sum(jj, a_values[[row_start+jj]])",
                 ],
             [
-                lp.GlobalArg("a_rowstarts", np.int32),
-                lp.GlobalArg("a_indices", np.int32),
+                lp.GlobalArg("a_rowstarts", np.int32, shape="auto"),
+                lp.GlobalArg("a_indices", np.int32, shape="auto"),
                 lp.GlobalArg("a_values", dtype),
-                lp.GlobalArg("x", dtype),
-                lp.GlobalArg("ax", dtype),
+                lp.GlobalArg("ax", dtype, shape="auto"),
                 lp.ValueArg("n", np.int32),
                 ],
             assumptions="n>=1 and row_len>=1")
@@ -718,7 +712,7 @@ def test_dependent_loop_bounds_3(ctx_factory):
                 "a[i,jj] = 1",
                 ],
             [
-                lp.GlobalArg("a_row_lengths", np.int32),
+                lp.GlobalArg("a_row_lengths", np.int32, shape="auto"),
                 lp.GlobalArg("a", dtype, shape=("n,n"), order="C"),
                 lp.ValueArg("n", np.int32),
                 ])
@@ -1029,17 +1023,35 @@ def test_write_parameter(ctx_factory):
                 ],
             assumptions="n>=1")
 
-    try:
+    import pytest
+    with pytest.raises(RuntimeError):
         lp.CompiledKernel(ctx, knl).get_code()
-    except RuntimeError, e:
-        assert "may not be written" in str(e)
-        pass # expected!
-    else:
-        assert False # expecting an error
 
 
 
 
+def test_arg_shape_guessing(ctx_factory):
+    ctx = ctx_factory()
+
+    knl = lp.make_kernel(ctx.devices[0], [
+            "{[i,j]: 0<=i,j<n }",
+            ],
+            """
+                a = 1.5 + sum((i,j), i*j)
+                b[i, j] = i*j
+                c[i+j, j] = b[j,i]
+                """,
+            [
+                lp.GlobalArg("a", shape=lp.auto_shape),
+                lp.GlobalArg("b", shape=lp.auto_shape),
+                lp.GlobalArg("c", shape=lp.auto_shape),
+                lp.ValueArg("n"),
+                ],
+            assumptions="n>=1")
+
+    print knl
+    print lp.CompiledKernel(ctx, knl).get_highlighted_code()
+
 if __name__ == "__main__":
     import sys
     if len(sys.argv) > 1:
-- 
GitLab