TMP SPLIT REVIEW ME

Author: inducer

loopy/check.py

@@ -862,10 +862,23 @@ def map_call(self, expr: p.Call, domain: nisl.Set, insn_id: str):
 
             kw_to_pos, _ = get_kw_pos_association(subkernel)
 
-            arg_id_to_arg = self.kernel.id_to_insn[insn_id].arg_id_to_arg()
+            insn = self.kernel.id_to_insn[insn_id]
+            if isinstance(insn, CallInstruction):
+                arg_id_to_arg = insn.arg_id_to_arg()
 
-            kwargs = {k: _mark_variables_from_caller(arg_id_to_arg[kw_to_pos[k]])
-                      for k in subkernel.get_unwritten_value_args()}
+                kwargs = {k: _mark_variables_from_caller(
+                                cast("ArithmeticExpression",
+                                     arg_id_to_arg[kw_to_pos[k]]))
+                          for k in subkernel.get_unwritten_value_args()}
+            else:
+                # The call appears as a plain expression (e.g. on the right
+                # hand side of an assignment); its arguments are found in
+                # *expr* itself.
+                kwargs = {k: _mark_variables_from_caller(
+                                cast("ArithmeticExpression",
+                                     expr.parameters[kw_to_pos[k]]))
+                          for k in subkernel.get_unwritten_value_args()
+                          if kw_to_pos[k] < len(expr.parameters)}
 
             kw_space = nisl.Space.from_names(
                 out=[], param=[*get_dependencies(tuple(kwargs.values())),

loopy/kernel/creation.py

@@ -2500,8 +2500,26 @@ def make_function(
         assumptions_space = nisl.Space.from_names(param=dom0.space.param_names)
         assumptions = nisl.Set.universe(assumptions_space)
     elif isinstance(assumptions, str):
+        # Value args are also legitimate assumption parameters, even when
+        # they do not appear in any domain (e.g. when they only occur in
+        # an array shape). For string declarations the eventual argument
+        # type is not known at this point, so declare all of their names.
+        import re
+
+        param_names = set(get_outer_params(parsed_domains))
+        for arg in kernel_args:
+            if isinstance(arg, str):
+                # kernel_data strings are split on commas up front, so
+                # fragments (e.g. " 2]" of "x: float64[n, 2]") may end up
+                # here; only valid names are useful.
+                arg = arg.strip()
+                if re.fullmatch(r"[a-zA-Z_][a-zA-Z0-9_]*", arg):
+                    param_names.add(arg)
+            elif isinstance(arg, ValueArg):
+                param_names.add(arg.name)
+
         assumptions_set_str = "[%s] -> { : %s}" \
-                % (",".join(s for s in get_outer_params(parsed_domains)),
+                % (",".join(sorted(param_names)),
                     assumptions)
         assumptions = nisl.make_set(assumptions_set_str)
     else:

loopy/transform/callable.py

@@ -204,9 +204,10 @@ def substitute_into_domain(
             param_name: str,
             expr: ArithmeticExpression,
             allowed_param_dims: Collection[str],
-        ):
+        ) -> nisl.Set:
     """
-    :arg allowed_deps: A :class:`list` of :class:`str` that are
+    :arg allowed_param_dims: Names that may be introduced as new parameters
+        of *domain* (i.e. parameters of the caller).
     """
     import pymbolic.primitives as prim
 
@@ -217,24 +218,29 @@ def substitute_into_domain(
 
     # {{{ rename 'param_name' to avoid namespace pollution with allowed_param_dims
 
-    domain = domain.rename_dims([(
-            param_name, UniqueNameGenerator(set(allowed_param_dims))(param_name))])
+    renamed_param_name = UniqueNameGenerator(set(allowed_param_dims))(param_name)
+    domain = domain.rename_dims([(param_name, renamed_param_name)])
 
     # }}}
 
-    domain.add_dims(DimType.param, [
-        dep for dep in get_dependencies(expr)
-        if dep in allowed_param_dims
-    ])
+    new_param_names: list[str] = []
+    for dep in get_dependencies(expr):
+        if dep in allowed_param_dims:
+            new_param_names.append(dep)
+        else:
+            raise ValueError("Augmenting caller's domain "
+                    f"with '{dep}' is not allowed.")
+
+    domain = domain.add_dims(DimType.param, new_param_names)
 
     set_ = isl_set_from_expr(domain.var_pw_affs,
-            prim.Comparison(prim.Variable(param_name),
+            prim.Comparison(prim.Variable(renamed_param_name),
                             "==",
                             expr))
 
     domain = domain & set_
 
-    return domain.project_out([param_name])
+    return domain.project_out([renamed_param_name])
 
 
 def get_valid_domain_param_names(knl):

test/test_reduction.py

@@ -147,7 +147,8 @@ def test_multi_nested_dependent_reduction():
                 lp.ValueArg("ntgts", np.int32),
                 lp.ValueArg("nboxes", np.int32),
                 ],
-            assumptions="ntgts>=1",
+            # 'a' is indexed by 'itgt', so the bounds check needs 'ntgts <= n'
+            assumptions="ntgts>=1 and ntgts <= n",
             target=lp.PyOpenCLTarget())
 
     print(lp.generate_code_v2(knl).device_code())
@@ -179,7 +180,8 @@ def test_recursive_nested_dependent_reduction():
                 lp.ValueArg("ntgts", np.int32),
                 lp.ValueArg("nboxes", np.int32),
                 ],
-            assumptions="ntgts>=1",
+            # 'a' is indexed by 'itgt', so the bounds check needs 'ntgts <= n'
+            assumptions="ntgts>=1 and ntgts <= n",
             target=lp.PyOpenCLTarget())
 
     print(lp.generate_code_v2(knl).device_code())