From ad02966a95686bd2c291cf92ce72a0a01e31c9b3 Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Mon, 20 May 2019 20:01:09 -0500 Subject: [PATCH 1/8] Begin refactoring ArgDescrInferenceMapper --- loopy/kernel/function_interface.py | 72 ++++++++++++++++++++++++++++++ loopy/preprocess.py | 41 ++++++----------- loopy/symbolic.py | 49 -------------------- loopy/transform/callable.py | 6 +-- 4 files changed, 88 insertions(+), 80 deletions(-) diff --git a/loopy/kernel/function_interface.py b/loopy/kernel/function_interface.py index 536fc9735..3bd544917 100644 --- a/loopy/kernel/function_interface.py +++ b/loopy/kernel/function_interface.py @@ -103,6 +103,78 @@ class ArrayArgDescriptor(ImmutableRecord): update_persistent_hash = update_persistent_hash + +def get_arg_descriptor_for_expression(kernel, expr): + """ + :returns: a :class:`ArrayArgDescriptor` or a :class:`ValueArgDescriptor` + describing the argument expression *expr* in *kernel*. + """ + from loopy.symbolic import (SubArrayRef, pw_aff_to_expr, Variable, + SweptInameStrideCollector) + from loopy.kernel.data import TemporaryVariable, ArrayArg + + if isinstance(expr, SubArrayRef): + name = expr.subscript.aggregate.name + arg = kernel.get_arg_descriptor(name) + + if not isinstance(arg, (TemporaryVariable, ArrayArg)): + raise LoopyError("unsupported argument type " + "'%s' of '%s' in call statement" + % (type(arg).__name__, expr.name)) + + aspace = arg.address_space + + from loopy.kernel.array import FixedStrideArrayDimTag as DimTag + from loopy.isl_helpers import simplify_via_aff + sub_dim_tags = [] + sub_shape = [] + + # FIXME This blindly assumes that dim_tag has a stride and + # will not work for non-stride dim tags (e.g. vec or sep). + + # FIXME: This will almost always be nonlinear--when does this + # actually help? Maybe the + linearized_index = simplify_via_aff( + sum( + dim_tag.stride*iname for dim_tag, iname in + zip(arg.dim_tags, expr.subscript.index_tuple))) + + strides_as_dict = SweptInameStrideCollector( + tuple(iname.name for iname in expr.swept_inames) + )(linearized_index) + sub_dim_tags = tuple( + DimTag(strides_as_dict[iname]) for iname in expr.swept_inames) + sub_shape = tuple( + pw_aff_to_expr( + kernel.get_iname_bounds(iname.name).upper_bound_pw_aff)+1 + for iname in expr.swept_inames) + if expr.swept_inames == (): + sub_shape = (1, ) + sub_dim_tags = (DimTag(1),) + + return ArrayArgDescriptor( + address_space=aspace, + dim_tags=sub_dim_tags, + shape=sub_shape) + + elif isinstance(expr, Variable): + arg = kernel.get_arg_descriptor(expr.name) + + if isinstance(arg, (TemporaryVariable, ArrayArg)): + return ArrayArgDescriptor( + address_space=arg.aspace, + dim_tags=arg.dim_tags, + shape=arg.shape) + elif isinstance(arg, ValueArg): + return ValueArgDescriptor() + else: + raise LoopyError("unsupported argument type " + "'%s' of '%s' in call statement" + % (type(arg).__name__, expr.name)) + + else: + return ValueArgDescriptor() + # }}} diff --git a/loopy/preprocess.py b/loopy/preprocess.py index a8dde5792..d03296435 100644 --- a/loopy/preprocess.py +++ b/loopy/preprocess.py @@ -2169,47 +2169,32 @@ class ArgDescrInferenceMapper(RuleAwareIdentityMapper): def map_call(self, expr, expn_state, **kwargs): from pymbolic.primitives import Call, CallWithKwargs - from loopy.kernel.function_interface import ValueArgDescriptor - from loopy.symbolic import ResolvedFunction, SubArrayRef + from loopy.symbolic import ResolvedFunction if not isinstance(expr.function, ResolvedFunction): # ignore if the call is not to a ResolvedFunction return super(ArgDescrInferenceMapper, self).map_call(expr, expn_state) - if isinstance(expr, Call): - kw_parameters = {} - else: - assert isinstance(expr, CallWithKwargs) - kw_parameters = expr.kw_parameters - - # descriptors for the args and kwargs of the Call - arg_id_to_descr = dict((i, par.get_array_arg_descriptor(self.caller_kernel)) - if isinstance(par, SubArrayRef) else (i, ValueArgDescriptor()) - for i, par in tuple(enumerate(expr.parameters)) + - tuple(kw_parameters.items())) - - assignee_id_to_descr = {} + arg_id_to_val = dict(enumerate(expr.parameters)) + if isinstance(expr, CallWithKwargs): + arg_id_to_val.update(expr.kw_parameters) if 'assignees' in kwargs: # If supplied with assignees then this is a CallInstruction assignees = kwargs['assignees'] - assert isinstance(assignees, tuple) - for i, par in enumerate(assignees): - if isinstance(par, SubArrayRef): - assignee_id_to_descr[-i-1] = ( - par.get_array_arg_descriptor(self.caller_kernel)) - else: - assignee_id_to_descr[-i-1] = ValueArgDescriptor() - - # gathering all the descriptors - combined_arg_id_to_descr = arg_id_to_descr.copy() - combined_arg_id_to_descr.update(assignee_id_to_descr) + for i, arg in enumerate(assignees): + arg_id_to_val[-i-1] = arg + + from loopy.kernel.function_interface import get_arg_descriptor_for_expression + arg_id_to_descr = dict( + (arg_id, get_arg_descriptor_for_expression(arg)) + for arg_id, arg in six.iteritems(arg_id_to_val)) # specializing the function according to the parameter description in_knl_callable = self.callables_table[expr.function.name] new_in_knl_callable, self.callables_table = ( in_knl_callable.with_descrs( - combined_arg_id_to_descr, self.caller_kernel, + arg_id_to_descr, self.caller_kernel, self.callables_table)) self.callables_table, new_func_id = ( self.callables_table.with_callable( @@ -2229,7 +2214,7 @@ class ArgDescrInferenceMapper(RuleAwareIdentityMapper): for child in expr.parameters), dict( (key, self.rec(val, expn_state)) - for key, val in six.iteritems(kw_parameters)) + for key, val in six.iteritems(expr.kw_parameters)) ) map_call_with_kwargs = map_call diff --git a/loopy/symbolic.py b/loopy/symbolic.py index 9a64fe4ac..a76f37654 100644 --- a/loopy/symbolic.py +++ b/loopy/symbolic.py @@ -826,55 +826,6 @@ class SubArrayRef(p.Expression): return EvaluatorWithDeficientContext(swept_inames_to_zeros)( self.subscript) - def get_array_arg_descriptor(self, kernel): - """ - Returns the dim_tags, memory scope, shape informations of a - :class:`SubArrayRef` argument in the caller kernel packed into - :class:`ArrayArgDescriptor` for the instance of :class:`SubArrayRef` in - the given *kernel*. - """ - from loopy.kernel.function_interface import ArrayArgDescriptor - - name = self.subscript.aggregate.name - - if name in kernel.temporary_variables: - assert name not in kernel.arg_dict - arg = kernel.temporary_variables[name] - else: - assert name in kernel.arg_dict - arg = kernel.arg_dict[name] - - aspace = arg.address_space - - from loopy.kernel.array import FixedStrideArrayDimTag as DimTag - from loopy.isl_helpers import simplify_via_aff - sub_dim_tags = [] - sub_shape = [] - try: - linearized_index = simplify_via_aff( - sum(dim_tag.stride*iname for dim_tag, iname in - zip(arg.dim_tags, self.subscript.index_tuple))) - except isl.Error: - linearized_index = sum(dim_tag.stride*iname for dim_tag, iname in - zip(arg.dim_tags, self.subscript.index_tuple)) - - strides_as_dict = SweptInameStrideCollector(tuple(iname.name for iname in - self.swept_inames))(linearized_index) - sub_dim_tags = tuple( - DimTag(strides_as_dict[iname]) for iname in self.swept_inames) - sub_shape = tuple( - pw_aff_to_expr( - kernel.get_iname_bounds(iname.name).upper_bound_pw_aff)+1 - for iname in self.swept_inames) - if self.swept_inames == (): - sub_shape = (1, ) - sub_dim_tags = (DimTag(1),) - - return ArrayArgDescriptor( - address_space=aspace, - dim_tags=sub_dim_tags, - shape=sub_shape) - def __getinitargs__(self): return (self.swept_inames, self.subscript) diff --git a/loopy/transform/callable.py b/loopy/transform/callable.py index 953ad5613..135987e06 100644 --- a/loopy/transform/callable.py +++ b/loopy/transform/callable.py @@ -628,8 +628,8 @@ def _match_caller_callee_argument_dimension_for_single_kernel( # Call to a callable kernel can only occur through a # CallInstruction. continue - # getting the caller->callee arg association + # get the caller->callee arg association parameters = insn.expression.parameters[:] kw_parameters = {} if isinstance(insn.expression, CallWithKwargs): @@ -658,7 +658,7 @@ def _match_caller_callee_argument_dimension_for_single_kernel( parameter_shapes.append(_shape_1_if_empty(kw_parameters[pos_to_kw[i]]) .get_array_arg_descriptor(caller_knl).shape) - # inserting the assignees at the required positions. + # insert the assignees at the required positions assignee_write_count = -1 for i, arg in enumerate(callee_knl.args): if arg.is_output_only: @@ -686,7 +686,7 @@ def _match_caller_callee_argument_dimension_for_single_kernel( raise NotImplementedError("Unknown instruction %s." % type(insn)) - # subkernel with instructions adjusted according to the new dimensions. + # subkernel with instructions adjusted according to the new dimensions new_callee_knl = callee_knl.copy(instructions=new_callee_insns) return new_callee_knl -- GitLab From e51a8af5d91609c7355327ff8c67aa665dd8458e Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Tue, 21 May 2019 01:24:22 -0500 Subject: [PATCH 2/8] Fixes for get_arg_descriptor_for_expression --- loopy/kernel/function_interface.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/loopy/kernel/function_interface.py b/loopy/kernel/function_interface.py index 3bd544917..26f90cd46 100644 --- a/loopy/kernel/function_interface.py +++ b/loopy/kernel/function_interface.py @@ -109,13 +109,14 @@ def get_arg_descriptor_for_expression(kernel, expr): :returns: a :class:`ArrayArgDescriptor` or a :class:`ValueArgDescriptor` describing the argument expression *expr* in *kernel*. """ - from loopy.symbolic import (SubArrayRef, pw_aff_to_expr, Variable, + from pymbolic.primitives import Variable + from loopy.symbolic import (SubArrayRef, pw_aff_to_expr, SweptInameStrideCollector) from loopy.kernel.data import TemporaryVariable, ArrayArg if isinstance(expr, SubArrayRef): name = expr.subscript.aggregate.name - arg = kernel.get_arg_descriptor(name) + arg = kernel.get_var_descriptor(name) if not isinstance(arg, (TemporaryVariable, ArrayArg)): raise LoopyError("unsupported argument type " @@ -125,7 +126,7 @@ def get_arg_descriptor_for_expression(kernel, expr): aspace = arg.address_space from loopy.kernel.array import FixedStrideArrayDimTag as DimTag - from loopy.isl_helpers import simplify_via_aff + from loopy.symbolic import simplify_using_aff sub_dim_tags = [] sub_shape = [] @@ -134,7 +135,8 @@ def get_arg_descriptor_for_expression(kernel, expr): # FIXME: This will almost always be nonlinear--when does this # actually help? Maybe the - linearized_index = simplify_via_aff( + linearized_index = simplify_using_aff( + kernel, sum( dim_tag.stride*iname for dim_tag, iname in zip(arg.dim_tags, expr.subscript.index_tuple))) @@ -158,7 +160,7 @@ def get_arg_descriptor_for_expression(kernel, expr): shape=sub_shape) elif isinstance(expr, Variable): - arg = kernel.get_arg_descriptor(expr.name) + arg = kernel.get_var_descriptor(expr.name) if isinstance(arg, (TemporaryVariable, ArrayArg)): return ArrayArgDescriptor( -- GitLab From fe208a40aef35e797d77d98497a355c045f53872 Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Tue, 21 May 2019 01:27:11 -0500 Subject: [PATCH 3/8] Add CallInstruction.arg_id_to_val --- loopy/kernel/instruction.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/loopy/kernel/instruction.py b/loopy/kernel/instruction.py index 0a2079ba5..1a56e8582 100644 --- a/loopy/kernel/instruction.py +++ b/loopy/kernel/instruction.py @@ -22,6 +22,7 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ +import six from six.moves import intern from pytools import ImmutableRecord, memoize_method from loopy.diagnostic import LoopyError @@ -1137,6 +1138,22 @@ class CallInstruction(MultiAssignmentBase): result += "\n" + 10*" " + "if (%s)" % " && ".join(self.predicates) return result + def arg_id_to_val(self): + """:returns: a :class:`dict` mapping argument identifiers (non-negative numbers + for positional arguments, strings for keyword args, and negative numbers + for assignees) to their respective values + """ + + from pymbolic.primitives import CallWithKwargs + arg_id_to_val = dict(enumerate(self.expression.parameters)) + if isinstance(self.expression, CallWithKwargs): + for kw, val in six.iteritems(self.expression.kw_parameters): + arg_id_to_val[kw] = val + for i, arg in enumerate(self.assignees): + arg_id_to_val[-i-1] = arg + + return arg_id_to_val + @property def atomicity(self): # Function calls can impossibly be atomic, and even the result assignment -- GitLab From 1795061095519ab225385152bf241c3b37a1741d Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Tue, 21 May 2019 01:27:40 -0500 Subject: [PATCH 4/8] Fix call site of get_arg_descriptor_for_expression --- loopy/preprocess.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/loopy/preprocess.py b/loopy/preprocess.py index d03296435..54a9204dc 100644 --- a/loopy/preprocess.py +++ b/loopy/preprocess.py @@ -2187,7 +2187,8 @@ class ArgDescrInferenceMapper(RuleAwareIdentityMapper): from loopy.kernel.function_interface import get_arg_descriptor_for_expression arg_id_to_descr = dict( - (arg_id, get_arg_descriptor_for_expression(arg)) + (arg_id, get_arg_descriptor_for_expression( + self.caller_kernel, arg)) for arg_id, arg in six.iteritems(arg_id_to_val)) # specializing the function according to the parameter description -- GitLab From a5b691ff1e107a04fd7271fad47cc1ec0f2d2da8 Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Tue, 21 May 2019 01:28:24 -0500 Subject: [PATCH 5/8] Add FIXME regarding simplify_{via,using}_aff --- loopy/symbolic.py | 1 + 1 file changed, 1 insertion(+) diff --git a/loopy/symbolic.py b/loopy/symbolic.py index a76f37654..d214b5e4f 100644 --- a/loopy/symbolic.py +++ b/loopy/symbolic.py @@ -1635,6 +1635,7 @@ def guarded_pwaff_from_expr(space, expr, vars_to_zero=None): # {{{ simplify using aff +# FIXME: redundant with simplify_via_aff def simplify_using_aff(kernel, expr): inames = get_dependencies(expr) & kernel.all_inames() -- GitLab From 6560e593523eb6b18a835c6f7839ccc820b0ca7b Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Tue, 21 May 2019 01:29:11 -0500 Subject: [PATCH 6/8] Refactor/simplify _match_caller_callee_argument_dimension_for_single_kernel --- loopy/transform/callable.py | 54 +++++++++++-------------------------- 1 file changed, 15 insertions(+), 39 deletions(-) diff --git a/loopy/transform/callable.py b/loopy/transform/callable.py index 135987e06..042990c77 100644 --- a/loopy/transform/callable.py +++ b/loopy/transform/callable.py @@ -34,7 +34,7 @@ from loopy.kernel.instruction import (CallInstruction, MultiAssignmentBase, Assignment, CInstruction, _DataObliviousInstruction) from loopy.symbolic import IdentityMapper, SubstitutionMapper, CombineMapper from loopy.isl_helpers import simplify_via_aff -from loopy.kernel.function_interface import (get_kw_pos_association, +from loopy.kernel.function_interface import ( CallableKernel, ScalarCallable) from loopy.program import Program, ResolvedFunctionMarker from loopy.symbolic import SubArrayRef @@ -616,10 +616,10 @@ class DimChanger(IdentityMapper): def _match_caller_callee_argument_dimension_for_single_kernel( caller_knl, callee_knl): """ - Returns a copy of *caller_knl* with the instance of - :class:`loopy.kernel.function_interface.CallableKernel` addressed by - *callee_function_name* in the *caller_knl* aligned with the argument - dimesnsions required by *caller_knl*. + :returns: a copy of *caller_knl* with the instance of + :class:`loopy.kernel.function_interface.CallableKernel` addressed by + *callee_function_name* in the *caller_knl* aligned with the argument + dimensions required by *caller_knl*. """ for insn in caller_knl.instructions: if not isinstance(insn, CallInstruction) or ( @@ -629,14 +629,6 @@ def _match_caller_callee_argument_dimension_for_single_kernel( # CallInstruction. continue - # get the caller->callee arg association - parameters = insn.expression.parameters[:] - kw_parameters = {} - if isinstance(insn.expression, CallWithKwargs): - kw_parameters = insn.expression.kw_parameters - - assignees = insn.assignees - def _shape_1_if_empty(shape): assert isinstance(shape, tuple) if shape == (): @@ -644,34 +636,18 @@ def _match_caller_callee_argument_dimension_for_single_kernel( else: return shape - parameter_shapes = [] - for par in parameters: - if isinstance(par, SubArrayRef): - parameter_shapes.append( - _shape_1_if_empty( - par.get_array_arg_descriptor(caller_knl).shape)) - else: - parameter_shapes.append((1, )) - - kw_to_pos, pos_to_kw = get_kw_pos_association(callee_knl) - for i in range(len(parameters), len(parameters)+len(kw_parameters)): - parameter_shapes.append(_shape_1_if_empty(kw_parameters[pos_to_kw[i]]) - .get_array_arg_descriptor(caller_knl).shape) - - # insert the assignees at the required positions - assignee_write_count = -1 - for i, arg in enumerate(callee_knl.args): - if arg.is_output_only: - assignee = assignees[-assignee_write_count-1] - parameter_shapes.insert(i, _shape_1_if_empty(assignee - .get_array_arg_descriptor(caller_knl).shape)) - assignee_write_count -= 1 - - callee_arg_to_desired_dim_tag = dict(zip([arg.name for arg in - callee_knl.args], parameter_shapes)) + from loopy.kernel.function_interface import ( + ArrayArgDescriptor, get_arg_descriptor_for_expression) + arg_id_to_shape = {} + for arg_id, arg in six.iteritems(insn.arg_id_to_val()): + arg_descr = get_arg_descriptor_for_expression(caller_knl, arg) + if isinstance(arg_descr, ArrayArgDescriptor): + arg_id_to_shape[arg_id] = _shape_1_if_empty(arg_descr) + dim_changer = DimChanger( callee_knl.arg_dict, - callee_arg_to_desired_dim_tag) + arg_id_to_shape) + new_callee_insns = [] for callee_insn in callee_knl.instructions: if isinstance(callee_insn, MultiAssignmentBase): -- GitLab From 3cf7abe0019d70995185e93daf5081a7c900bf35 Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Thu, 23 May 2019 14:23:00 -0500 Subject: [PATCH 7/8] Add parameter matching FIXME --- loopy/kernel/function_interface.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/loopy/kernel/function_interface.py b/loopy/kernel/function_interface.py index 34d360512..ba01c9011 100644 --- a/loopy/kernel/function_interface.py +++ b/loopy/kernel/function_interface.py @@ -675,6 +675,11 @@ class CallableKernel(InKernelCallable): # {{{ map the arg_descrs so that all the variables are from the callees # perspective + # FIXME: This is ill-formed, because par can be an expression, e.g. + # 2*i+2 or 2*(i+1). A key feature of expression is that structural + # equality and semantic equality are not the same, so even if the + # SubstitutionMapper allowed non-variables, it would have to solve the + # (considerable) problem of expression equivalence. substs = {} for arg, par in zip(self.subkernel.args, expr.parameters): if isinstance(arg, ValueArg): -- GitLab From 7361be5ab66ed86bb859e3c5ae5484e41031354a Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Thu, 23 May 2019 14:33:40 -0500 Subject: [PATCH 8/8] Do not allow passing entire array by name without using SubArrayRef --- loopy/kernel/function_interface.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/loopy/kernel/function_interface.py b/loopy/kernel/function_interface.py index ba01c9011..cf6e92771 100644 --- a/loopy/kernel/function_interface.py +++ b/loopy/kernel/function_interface.py @@ -127,7 +127,8 @@ class ArrayArgDescriptor(ImmutableRecord): def get_arg_descriptor_for_expression(kernel, expr): """ :returns: a :class:`ArrayArgDescriptor` or a :class:`ValueArgDescriptor` - describing the argument expression *expr* in *kernel*. + describing the argument expression *expr* which occurs + in a call in the code of *kernel*. """ from pymbolic.primitives import Variable from loopy.symbolic import (SubArrayRef, pw_aff_to_expr, @@ -183,10 +184,10 @@ def get_arg_descriptor_for_expression(kernel, expr): arg = kernel.get_var_descriptor(expr.name) if isinstance(arg, (TemporaryVariable, ArrayArg)): - return ArrayArgDescriptor( - address_space=arg.aspace, - dim_tags=arg.dim_tags, - shape=arg.shape) + raise LoopyError("may not pass entire array " + "'%s' in call statement in kernel '%s'" + % (expr.name, kernel.name)) + elif isinstance(arg, ValueArg): return ValueArgDescriptor() else: -- GitLab