From f719cfb541e3f0938be21e3da1d892e0b1a968ef Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Thu, 5 Aug 2021 17:03:36 -0500 Subject: [PATCH 01/41] add map_domain to loopy init --- loopy/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/loopy/__init__.py b/loopy/__init__.py index a73f83bb9..177fae61c 100644 --- a/loopy/__init__.py +++ b/loopy/__init__.py @@ -76,7 +76,7 @@ affine_map_inames, find_unused_axis_tag, make_reduction_inames_unique, has_schedulable_iname_nesting, get_iname_duplication_options, - add_inames_to_insn, add_inames_for_unused_hw_axes) + add_inames_to_insn, add_inames_for_unused_hw_axes, map_domain) from loopy.transform.instruction import ( find_instructions, map_instructions, @@ -202,7 +202,7 @@ "affine_map_inames", "find_unused_axis_tag", "make_reduction_inames_unique", "has_schedulable_iname_nesting", "get_iname_duplication_options", - "add_inames_to_insn", "add_inames_for_unused_hw_axes", + "add_inames_to_insn", "add_inames_for_unused_hw_axes", "map_domain", "add_prefetch", "change_arg_to_image", "tag_array_axes", "tag_data_axes", From 2780237507258383eb99442e039c9df99fa363d3 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Thu, 5 Aug 2021 17:06:09 -0500 Subject: [PATCH 02/41] copy in map_domain funcs from old branch (add-map-domain-transformation) --- loopy/transform/iname.py | 323 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 323 insertions(+) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index c3b4a42ee..33b52b913 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -72,6 +72,8 @@ .. autofunction:: add_inames_to_insn +.. autofunction:: map_domain + .. autofunction:: add_inames_for_unused_hw_axes """ @@ -1832,6 +1834,327 @@ def add_inames_to_insn(kernel, inames, insn_match): # }}} +# {{{ map_domain + +class _MapDomainMapper(RuleAwareIdentityMapper): + def __init__(self, rule_mapping_context, within, new_inames, substitutions): + super(_MapDomainMapper, self).__init__(rule_mapping_context) + + self.within = within + + self.old_inames = frozenset(substitutions) + self.new_inames = new_inames + + self.substitutions = substitutions + + def map_reduction(self, expr, expn_state): + red_overlap = frozenset(expr.inames) & self.old_inames + arg_ctx_overlap = frozenset(expn_state.arg_context) & self.old_inames + if (red_overlap + and self.within( + expn_state.kernel, + expn_state.instruction)): + if len(red_overlap) != len(self.old_inames): + raise LoopyError("reduction '%s' involves a part " + "of the map domain inames. Reductions must " + "either involve all or none of the map domain " + "inames." % str(expr)) + + if arg_ctx_overlap: + if arg_ctx_overlap == red_overlap: + # All variables are shadowed by context, that's OK. + return super(_MapDomainMapper, self).map_reduction( + expr, expn_state) + else: + raise LoopyError("reduction '%s' has" + "some of the reduction variables affected " + "by the map_domain shadowed by context. " + "Either all or none must be shadowed." + % str(expr)) + + new_inames = list(expr.inames) + for old_iname in self.old_inames: + new_inames.remove(old_iname) + new_inames.extend(self.new_inames) + + from loopy.symbolic import Reduction + return Reduction(expr.operation, tuple(new_inames), + self.rec(expr.expr, expn_state), + expr.allow_simultaneous) + else: + return super(_MapDomainMapper, self).map_reduction(expr, expn_state) + + def map_variable(self, expr, expn_state): + if (expr.name in self.old_inames + and expr.name not in expn_state.arg_context + and self.within( + expn_state.kernel, + expn_state.instruction)): + return self.substitutions[expr.name] + else: + return super(_MapDomainMapper, self).map_variable(expr, expn_state) + + +def _find_aff_subst_from_map(iname, isl_map): + if not isinstance(isl_map, isl.BasicMap): + raise RuntimeError("isl_map must be a BasicMap") + + dt, dim_idx = isl_map.get_var_dict()[iname] + + assert dt == dim_type.in_ + + # Force isl to solve for only this iname on its side of the map, by + # projecting out all other "in" variables. + isl_map = isl_map.project_out(dt, dim_idx+1, isl_map.dim(dt)-(dim_idx+1)) + isl_map = isl_map.project_out(dt, 0, dim_idx) + dim_idx = 0 + + # Convert map to set to avoid "domain of affine expression should be a set". + # The old "in" variable will be the last of the out_dims. + new_dim_idx = isl_map.dim(dim_type.out) + isl_map = isl_map.move_dims( + dim_type.out, isl_map.dim(dim_type.out), + dt, dim_idx, 1) + isl_map = isl_map.range() # now a set + dt = dim_type.set + dim_idx = new_dim_idx + del new_dim_idx + + for cns in isl_map.get_constraints(): + if cns.is_equality() and cns.involves_dims(dt, dim_idx, 1): + coeff = cns.get_coefficient_val(dt, dim_idx) + cns_zeroed = cns.set_coefficient_val(dt, dim_idx, 0) + if cns_zeroed.involves_dims(dt, dim_idx, 1): + # not suitable, constraint still involves dim, perhaps in a div + continue + + if coeff.is_one(): + return -cns_zeroed.get_aff() + elif coeff.is_negone(): + return cns_zeroed.get_aff() + else: + # not suitable, coefficient does not have unit coefficient + continue + + raise LoopyError("no suitable equation for '%s' found" % iname) + + +# TODO to match convention elsewhere, swap 'dt' and 'dim_type' identifiers +# (use dt to abbreviate islpy.dim_type, and use dim_type for variables +# containing a specific dim_type) + +def _find_and_rename_dim(old_map, dim_types, old_name, new_name): + # (This function is only used once here, but do not inline it; it is used many + # times in child branch update-dependencies-during-transformations.) + new_map = old_map.copy() + for dt in dim_types: + new_map = new_map.set_dim_name( + dt, new_map.find_dim_by_name(dt, old_name), new_name) + return new_map + + +@for_each_kernel +def map_domain(kernel, isl_map, within=None): + # FIXME: Express _split_iname_backend in terms of this + # Missing/deleted for now: + # - slab processing + # - priorities processing + # FIXME: Process priorities + # FIXME: Express affine_map_inames in terms of this, deprecate + # FIXME: Document + + # FIXME: Support within + # FIXME: Right now, this requires all inames in a domain (or none) to + # be mapped. That makes this awkward to use. + + # {{{ within processing (disabled for now) + if within is not None: + raise NotImplementedError("within") + + from loopy.match import parse_match + within = parse_match(within) + + # {{{ return the same kernel if no kernel matches + + if not any(within(kernel, insn) for insn in kernel.instructions): + return kernel + + # }}} + + # }}} + + if not isl_map.is_bijective(): + raise LoopyError("isl_map must be bijective") + + new_inames = frozenset(isl_map.get_var_dict(dim_type.out)) + old_inames = frozenset(isl_map.get_var_dict(dim_type.in_)) + + # {{{ solve for representation of old inames in terms of new + + substitutions = {} + var_substitutions = {} + applied_iname_rewrites = kernel.applied_iname_rewrites[:] + + from loopy.symbolic import aff_to_expr + from pymbolic import var + for iname in old_inames: + substitutions[iname] = aff_to_expr( + _find_aff_subst_from_map(iname, isl_map)) + var_substitutions[var(iname)] = aff_to_expr( + _find_aff_subst_from_map(iname, isl_map)) + + applied_iname_rewrites.append(var_substitutions) + del var_substitutions + + # }}} + + from loopy.schedule.checker.utils import ( + add_and_name_isl_dims, + ) + + def process_set(s): + var_dict = s.get_var_dict() + + overlap = old_inames & frozenset(var_dict) + + if not overlap: + # inames in s are not present in transform map, don't change s + return s + + if len(overlap) != len(old_inames): + raise LoopyError("loop domain '%s' involves a part " + "of the map domain inames. Domains must " + "either involve all or none of the map domain " + "inames." % s) + + from loopy.schedule.checker.utils import ( + add_eq_isl_constraint_from_names, + ) + + # {{{ align dims of isl_map and s + + # FIXME: Make this less gross + # FIXME: Make an exported/documented interface of this in islpy + from islpy import _align_dim_type + + map_with_s_domain = isl.Map.from_domain(s) + + # {{{ deal with dims missing from transform map (isl_map) + + # If dims in s are missing from transform map, they need to be added + # so that intersect_domain doesn't remove them. + # Order doesn't matter here because dims will be aligned in the next step. + dims_missing_from_transform_map = list( + set(s.get_var_names(dim_type.set)) - + set(isl_map.get_var_names(dim_type.in_))) + augmented_isl_map = add_and_name_isl_dims( + isl_map, dim_type.in_, dims_missing_from_transform_map) + + # We want these missing inames to map to themselves so that the transform + # has no effect on them. Unfortunatley isl will break if the + # names of the out dims aren't unique, so we will temporariliy rename them + # and then change the names back afterward. + + # FIXME: need better way to make sure proxy dim names are unique + dims_missing_from_transform_map_proxies = [ + d+"__prox" for d in dims_missing_from_transform_map] + assert not set(dims_missing_from_transform_map_proxies) & set( + augmented_isl_map.get_var_dict().keys()) + + augmented_isl_map = add_and_name_isl_dims( + augmented_isl_map, dim_type.out, dims_missing_from_transform_map_proxies) + + # Set proxy iname equal to real iname + for proxy_iname, real_iname in zip( + dims_missing_from_transform_map_proxies, + dims_missing_from_transform_map): + augmented_isl_map = add_eq_isl_constraint_from_names( + augmented_isl_map, proxy_iname, real_iname) + + # }}} + + dim_types = [dim_type.param, dim_type.in_, dim_type.out] + s_names = [ + map_with_s_domain.get_dim_name(dt, i) + for dt in dim_types + for i in range(map_with_s_domain.dim(dt)) + ] + map_names = [ + augmented_isl_map.get_dim_name(dt, i) + for dt in dim_types + for i in range(augmented_isl_map.dim(dt)) + ] + + # (order doesn't matter in s_names/map_names, + # _align_dim_type just converts these to sets + # to determine which names are in both the obj and template, + # not sure why this isn't just handled inside _align_dim_type) + aligned_map = _align_dim_type( + dim_type.param, + augmented_isl_map, map_with_s_domain, False, + map_names, s_names) + aligned_map = _align_dim_type( + dim_type.in_, + aligned_map, map_with_s_domain, False, + map_names, s_names) + + # }}} + + new_s = aligned_map.intersect_domain(s).range() + + # Now rename the proxy dims back to their original names + for proxy_iname, real_iname in zip( + dims_missing_from_transform_map_proxies, + dims_missing_from_transform_map): + new_s = _find_and_rename_dim( + new_s, [dim_type.set], proxy_iname, real_iname) + + return new_s + + # FIXME: Revive _project_out_only_if_all_instructions_in_within + + new_domains = [process_set(dom) for dom in kernel.domains] + + # {{{ update within_inames + + new_insns = [] + for insn in kernel.instructions: + overlap = old_inames & insn.within_inames + if overlap and within(kernel, insn): + if len(overlap) != len(old_inames): + raise LoopyError("instruction '%s' is within only a part " + "of the map domain inames. Instructions must " + "either be within all or none of the map domain " + "inames." % insn.id) + + insn = insn.copy( + within_inames=(insn.within_inames - old_inames) | new_inames) + else: + # leave insn unmodified + pass + + new_insns.append(insn) + + # }}} + + kernel = kernel.copy( + domains=new_domains, + instructions=new_insns, + applied_iname_rewrites=applied_iname_rewrites) + + rule_mapping_context = SubstitutionRuleMappingContext( + kernel.substitutions, kernel.get_var_name_generator()) + ins = _MapDomainMapper(rule_mapping_context, within, + new_inames, substitutions) + + kernel = ins.map_kernel(kernel) + kernel = rule_mapping_context.finish_kernel(kernel) + + return kernel + +# }}} + + @for_each_kernel def add_inames_for_unused_hw_axes(kernel, within=None): """ From e5edf400ca070ca98fd9aed30673383eeac35996 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Thu, 5 Aug 2021 17:17:23 -0500 Subject: [PATCH 03/41] define add_and_name_isl_dims and add_eq_isl_constraint_from_names locally for now (previously were defined in schedule checking utils before this code was moved to independent branch) --- loopy/transform/iname.py | 44 +++++++++++++++++++++++++++++----------- 1 file changed, 32 insertions(+), 12 deletions(-) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index 33b52b913..25cd6e4db 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -1939,10 +1939,38 @@ def _find_aff_subst_from_map(iname, isl_map): raise LoopyError("no suitable equation for '%s' found" % iname) -# TODO to match convention elsewhere, swap 'dt' and 'dim_type' identifiers +# FIXME to match convention elsewhere, swap 'dt' and 'dim_type' identifiers # (use dt to abbreviate islpy.dim_type, and use dim_type for variables # containing a specific dim_type) +def _add_and_name_isl_dims(isl_map, dt, names): + # (This function is also defined in independent, unmerged branch + # statement-instance-order-and-lex-order-map, and used in child branches + # thereof. Once these branches are all merged, it may make sense to move + # this function to a location for more general-purpose machinery. In the + # other branches, this function's name excludes the leading underscore.) + new_idx_start = isl_map.dim(dt) + new_map = isl_map.add_dims(dt, len(names)) + for i, name in enumerate(names): + new_map = new_map.set_dim_name(dt, new_idx_start+i, name) + return new_map + + +def _add_eq_isl_constraint_from_names(isl_map, var1, var2): + # (This function is also defined in independent, unmerged branch + # statement-instance-order-and-lex-order-map, and used in child branches + # thereof. Once these branches are all merged, it may make sense to move + # this function to a location for more general-purpose machinery. In the + # other branches, this function's name excludes the leading underscore.) + + # add constraint var1 = var2 + + return isl_map.add_constraint( + isl.Constraint.eq_from_names( + isl_map.space, + {1: 0, var1: 1, var2: -1})) + + def _find_and_rename_dim(old_map, dim_types, old_name, new_name): # (This function is only used once here, but do not inline it; it is used many # times in child branch update-dependencies-during-transformations.) @@ -2008,10 +2036,6 @@ def map_domain(kernel, isl_map, within=None): # }}} - from loopy.schedule.checker.utils import ( - add_and_name_isl_dims, - ) - def process_set(s): var_dict = s.get_var_dict() @@ -2027,10 +2051,6 @@ def process_set(s): "either involve all or none of the map domain " "inames." % s) - from loopy.schedule.checker.utils import ( - add_eq_isl_constraint_from_names, - ) - # {{{ align dims of isl_map and s # FIXME: Make this less gross @@ -2047,7 +2067,7 @@ def process_set(s): dims_missing_from_transform_map = list( set(s.get_var_names(dim_type.set)) - set(isl_map.get_var_names(dim_type.in_))) - augmented_isl_map = add_and_name_isl_dims( + augmented_isl_map = _add_and_name_isl_dims( isl_map, dim_type.in_, dims_missing_from_transform_map) # We want these missing inames to map to themselves so that the transform @@ -2061,14 +2081,14 @@ def process_set(s): assert not set(dims_missing_from_transform_map_proxies) & set( augmented_isl_map.get_var_dict().keys()) - augmented_isl_map = add_and_name_isl_dims( + augmented_isl_map = _add_and_name_isl_dims( augmented_isl_map, dim_type.out, dims_missing_from_transform_map_proxies) # Set proxy iname equal to real iname for proxy_iname, real_iname in zip( dims_missing_from_transform_map_proxies, dims_missing_from_transform_map): - augmented_isl_map = add_eq_isl_constraint_from_names( + augmented_isl_map = _add_eq_isl_constraint_from_names( augmented_isl_map, proxy_iname, real_iname) # }}} From c6320d849eb981ddc6511794c7d2093c6b283fe1 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Thu, 5 Aug 2021 17:29:44 -0500 Subject: [PATCH 04/41] define ensure_dim_names_match_and_align locally for now (previously was defined in schedule checking utils before this code was moved to independent branch) --- test/test_transform.py | 242 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 242 insertions(+) diff --git a/test/test_transform.py b/test/test_transform.py index 51e7c2636..91db5ceeb 100644 --- a/test/test_transform.py +++ b/test/test_transform.py @@ -594,6 +594,248 @@ def test_nested_substs_in_insns(ctx_factory): lp.auto_test_vs_ref(ref_prg, ctx, t_unit) +# {{{ test_map_domain_vs_split_iname + +def _ensure_dim_names_match_and_align(obj_map, tgt_map): + # (This function is also defined in independent, unmerged branch + # new-dependency-and-nest-constraint-semantics-development, and used in + # child branches thereof. Once these branches are all merged, it may make + # sense to move this function to a location for more general-purpose + # machinery. In the other branches, this function's name excludes the + # leading underscore.) + from islpy import align_spaces + from islpy import dim_type as dt + + # first make sure names match + if not all( + set(obj_map.get_var_names(dt)) == set(tgt_map.get_var_names(dt)) + for dt in + [dt.in_, dt.out, dt.param]): + raise ValueError( + "Cannot align spaces; names don't match:\n%s\n%s" + % (obj_map, tgt_map)) + + return align_spaces(obj_map, tgt_map) + + +def test_map_domain_vs_split_iname(): + + # {{{ Make kernel + + knl = lp.make_kernel( + [ + "[nx,nt] -> {[x, t]: 0 <= x < nx and 0 <= t < nt}", + "[ni] -> {[i]: 0 <= i < ni}", + ], + """ + a[x,t] = b[x,t] {id=stmta} + c[x,t] = d[x,t] {id=stmtc} + e[i] = f[i] + """, + lang_version=(2018, 2), + ) + knl = lp.add_and_infer_dtypes(knl, {"b,d,f": np.float32}) + ref_knl = knl + + # }}} + + # {{{ Apply domain change mapping + + knl_map_dom = ref_knl # loop priority goes away, deps stay + + # Create map_domain mapping: + import islpy as isl + transform_map = isl.BasicMap( + "[nt] -> {[t] -> [t_outer, t_inner]: " + "0 <= t_inner < 32 and " + "32*t_outer + t_inner = t and " + "0 <= 32*t_outer + t_inner < nt}") + + # Call map_domain to transform kernel + knl_map_dom = lp.map_domain(knl_map_dom, transform_map) + + # Prioritize loops (prio should eventually be updated in map_domain?) + knl_map_dom = lp.prioritize_loops(knl_map_dom, "x, t_outer, t_inner") + + # Get a linearization + proc_knl_map_dom = lp.preprocess_kernel(knl_map_dom) + lin_knl_map_dom = lp.get_one_linearized_kernel( + proc_knl_map_dom["loopy_kernel"], proc_knl_map_dom.callables_table) + + # }}} + + # {{{ Split iname and see if we get the same result + + knl_split_iname = ref_knl + knl_split_iname = lp.split_iname(knl_split_iname, "t", 32) + knl_split_iname = lp.prioritize_loops(knl_split_iname, "x, t_outer, t_inner") + proc_knl_split_iname = lp.preprocess_kernel(knl_split_iname) + lin_knl_split_iname = lp.get_one_linearized_kernel( + proc_knl_split_iname["loopy_kernel"], proc_knl_split_iname.callables_table) + + for d_map_domain, d_split_iname in zip( + knl_map_dom["loopy_kernel"].domains, + knl_split_iname["loopy_kernel"].domains): + d_map_domain_aligned = _ensure_dim_names_match_and_align( + d_map_domain, d_split_iname) + assert d_map_domain_aligned == d_split_iname + + for litem_map_domain, litem_split_iname in zip( + lin_knl_map_dom.linearization, lin_knl_split_iname.linearization): + assert litem_map_domain == litem_split_iname + + # Can't easily compare instructions because equivalent subscript + # expressions may have different orders + + # }}} + +# }}} + + +# {{{ test_map_domain_with_transform_map_missing_dims + +def test_map_domain_with_transform_map_missing_dims(): + # Make sure map_domain works correctly when the mapping doesn't include + # all the dims in the domain. + + # {{{ Make kernel + + knl = lp.make_kernel( + [ + "[nx,nt] -> {[x, y, z, t]: 0 <= x,y,z < nx and 0 <= t < nt}", + ], + """ + a[y,x,t,z] = b[y,x,t,z] {id=stmta} + """, + lang_version=(2018, 2), + ) + knl = lp.add_and_infer_dtypes(knl, {"b": np.float32}) + ref_knl = knl + + # }}} + + # {{{ Apply domain change mapping + + knl_map_dom = ref_knl # loop priority goes away, deps stay + + # Create map_domain mapping that only includes t and y + # (x and z should be unaffected) + import islpy as isl + transform_map = isl.BasicMap( + "[nx,nt] -> {[t, y] -> [t_outer, t_inner, y_new]: " + "0 <= t_inner < 32 and " + "32*t_outer + t_inner = t and " + "0 <= 32*t_outer + t_inner < nt and " + "y = y_new" + "}") + + # Call map_domain to transform kernel + knl_map_dom = lp.map_domain(knl_map_dom, transform_map) + + # Prioritize loops (prio should eventually be updated in map_domain?) + try: + # Use constrain_loop_nesting if it's available + desired_prio = "x, t_outer, t_inner, z, y_new" + knl_map_dom = lp.constrain_loop_nesting(knl_map_dom, desired_prio) + except AttributeError: + # For some reason, prioritize_loops can't handle the ordering above + # when linearizing knl_split_iname below + desired_prio = "z, y_new, x, t_outer, t_inner" + knl_map_dom = lp.prioritize_loops(knl_map_dom, desired_prio) + + # Get a linearization + proc_knl_map_dom = lp.preprocess_kernel(knl_map_dom) + lin_knl_map_dom = lp.get_one_linearized_kernel( + proc_knl_map_dom["loopy_kernel"], proc_knl_map_dom.callables_table) + + # }}} + + # {{{ Split iname and see if we get the same result + + knl_split_iname = ref_knl + knl_split_iname = lp.split_iname(knl_split_iname, "t", 32) + knl_split_iname = lp.rename_iname(knl_split_iname, "y", "y_new") + try: + # Use constrain_loop_nesting if it's available + knl_split_iname = lp.constrain_loop_nesting(knl_split_iname, desired_prio) + except AttributeError: + knl_split_iname = lp.prioritize_loops(knl_split_iname, desired_prio) + proc_knl_split_iname = lp.preprocess_kernel(knl_split_iname) + lin_knl_split_iname = lp.get_one_linearized_kernel( + proc_knl_split_iname["loopy_kernel"], proc_knl_split_iname.callables_table) + + for d_map_domain, d_split_iname in zip( + knl_map_dom["loopy_kernel"].domains, + knl_split_iname["loopy_kernel"].domains): + d_map_domain_aligned = _ensure_dim_names_match_and_align( + d_map_domain, d_split_iname) + assert d_map_domain_aligned == d_split_iname + + for litem_map_domain, litem_split_iname in zip( + lin_knl_map_dom.linearization, lin_knl_split_iname.linearization): + assert litem_map_domain == litem_split_iname + + # Can't easily compare instructions because equivalent subscript + # expressions may have different orders + + # }}} + +# }}} + + +def test_diamond_tiling(ctx_factory, interactive=False): + ctx = ctx_factory() + queue = cl.CommandQueue(ctx) + + ref_knl = lp.make_kernel( + "[nx,nt] -> {[ix, it]: 1<=ix {[ix, it] -> [tx, tt, tparity, itt, itx]: " + "16*(tx - tt) + itx - itt = ix - it and " + "16*(tx + tt + tparity) + itt + itx = ix + it and " + "0<=tparity<2 and 0 <= itx - itt < 16 and 0 <= itt+itx < 16}") + knl = lp.map_domain(knl_for_transform, m) + knl = lp.prioritize_loops(knl, "tt,tparity,tx,itt,itx") + + if interactive: + nx = 43 + u = np.zeros((nx, 200)) + x = np.linspace(-1, 1, nx) + dx = x[1] - x[0] + u[:, 0] = u[:, 1] = np.exp(-100*x**2) + + u_dev = cl.array.to_device(queue, u) + knl(queue, u=u_dev, dx=dx, dt=dx) + + u = u_dev.get() + import matplotlib.pyplot as plt + plt.imshow(u.T) + plt.show() + else: + types = {"dt,dx,u": np.float64} + knl = lp.add_and_infer_dtypes(knl, types) + ref_knl = lp.add_and_infer_dtypes(ref_knl, types) + + lp.auto_test_vs_ref(ref_knl, ctx, knl, + parameters={ + "nx": 200, "nt": 300, + "dx": 1, "dt": 1 + }) + + def test_extract_subst_with_iname_deps_in_templ(ctx_factory): knl = lp.make_kernel( "{[i, j, k]: 0<=i<100 and 0<=j,k<5}", From 34b7d8fd6a34ed91b432766cfd1b00d14f3acdc6 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 10 Aug 2021 13:06:54 -0500 Subject: [PATCH 05/41] remove FIXME now that dim_type usage is consistent --- loopy/transform/iname.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index 25cd6e4db..131056eab 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -1939,10 +1939,6 @@ def _find_aff_subst_from_map(iname, isl_map): raise LoopyError("no suitable equation for '%s' found" % iname) -# FIXME to match convention elsewhere, swap 'dt' and 'dim_type' identifiers -# (use dt to abbreviate islpy.dim_type, and use dim_type for variables -# containing a specific dim_type) - def _add_and_name_isl_dims(isl_map, dt, names): # (This function is also defined in independent, unmerged branch # statement-instance-order-and-lex-order-map, and used in child branches From 62d711b6c55cd6b135ad6c6d3eb7b36464ed601d Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 10 Aug 2021 16:59:59 -0500 Subject: [PATCH 06/41] add better encapsulation for case where domain map has fewer dims than domains (with _apply_identity_for_missing_map_dims func) --- loopy/transform/iname.py | 84 ++++++++++++++++++++++------------------ 1 file changed, 46 insertions(+), 38 deletions(-) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index 131056eab..f3369cea8 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -1977,6 +1977,40 @@ def _find_and_rename_dim(old_map, dim_types, old_name, new_name): return new_map +def _apply_identity_for_missing_map_dims(mapping, desired_dims): + + # If dims in s are missing from transform map, they need to be added + # so that, e.g, intersect_domain doesn't remove them. + # (assume ordering will be handled afterward) + + missing_dims = list( + set(desired_dims) - set(mapping.get_var_names(dim_type.in_))) + augmented_mapping = _add_and_name_isl_dims( + mapping, dim_type.in_, missing_dims) + + # We want these missing inames to map to themselves so that the map + # has no effect on them. Unfortunatley isl will break if the + # names of the out dims aren't unique, so we will temporariliy rename them + # (and then plan to change the names back afterward). + + # FIXME: need better way to make sure proxy dim names are unique within map + missing_dims_proxies = [d+"__prox" for d in missing_dims] + assert not set(missing_dims_proxies) & set( + augmented_mapping.get_var_dict().keys()) + + augmented_mapping = _add_and_name_isl_dims( + augmented_mapping, dim_type.out, missing_dims_proxies) + + proxy_name_pairs = list(zip(missing_dims, missing_dims_proxies)) + + # Set proxy iname equal to real iname with equality constraint + for real_iname, proxy_iname in proxy_name_pairs: + augmented_mapping = _add_eq_isl_constraint_from_names( + augmented_mapping, proxy_iname, real_iname) + + return augmented_mapping, proxy_name_pairs + + @for_each_kernel def map_domain(kernel, isl_map, within=None): # FIXME: Express _split_iname_backend in terms of this @@ -2049,46 +2083,22 @@ def process_set(s): # {{{ align dims of isl_map and s - # FIXME: Make this less gross - # FIXME: Make an exported/documented interface of this in islpy from islpy import _align_dim_type map_with_s_domain = isl.Map.from_domain(s) - # {{{ deal with dims missing from transform map (isl_map) - - # If dims in s are missing from transform map, they need to be added - # so that intersect_domain doesn't remove them. - # Order doesn't matter here because dims will be aligned in the next step. - dims_missing_from_transform_map = list( - set(s.get_var_names(dim_type.set)) - - set(isl_map.get_var_names(dim_type.in_))) - augmented_isl_map = _add_and_name_isl_dims( - isl_map, dim_type.in_, dims_missing_from_transform_map) - - # We want these missing inames to map to themselves so that the transform - # has no effect on them. Unfortunatley isl will break if the - # names of the out dims aren't unique, so we will temporariliy rename them - # and then change the names back afterward. - - # FIXME: need better way to make sure proxy dim names are unique - dims_missing_from_transform_map_proxies = [ - d+"__prox" for d in dims_missing_from_transform_map] - assert not set(dims_missing_from_transform_map_proxies) & set( - augmented_isl_map.get_var_dict().keys()) - - augmented_isl_map = _add_and_name_isl_dims( - augmented_isl_map, dim_type.out, dims_missing_from_transform_map_proxies) - - # Set proxy iname equal to real iname - for proxy_iname, real_iname in zip( - dims_missing_from_transform_map_proxies, - dims_missing_from_transform_map): - augmented_isl_map = _add_eq_isl_constraint_from_names( - augmented_isl_map, proxy_iname, real_iname) - - # }}} + # If there are dims in s that are not mapped by isl_map, add them + # to the in/out space of isl_map so that they remain unchanged. + # (temporary proxy dim names are needed in out space of transform + # map because isl won't allow any dim names to match, i.e., instead + # of just mapping {[unused_iname]->[unused_iname]}, we have to map + # {[unused_name]->[unused_name__prox] : unused_name__prox = unused_name}, + # and then rename unused_name__prox afterward.) + augmented_isl_map, proxy_name_pairs = _apply_identity_for_missing_map_dims( + isl_map, s.get_var_names(dim_type.set)) + # FIXME: Make this less gross + # FIXME: Make an exported/documented interface of this in islpy dim_types = [dim_type.param, dim_type.in_, dim_type.out] s_names = [ map_with_s_domain.get_dim_name(dt, i) @@ -2119,9 +2129,7 @@ def process_set(s): new_s = aligned_map.intersect_domain(s).range() # Now rename the proxy dims back to their original names - for proxy_iname, real_iname in zip( - dims_missing_from_transform_map_proxies, - dims_missing_from_transform_map): + for real_iname, proxy_iname in proxy_name_pairs: new_s = _find_and_rename_dim( new_s, [dim_type.set], proxy_iname, real_iname) From 05acfbee1f4626e16d458893cb09a9ee23ecc94c Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 10 Aug 2021 18:05:28 -0500 Subject: [PATCH 07/41] improve var names and error message related to mismatching of transform map and inames domain --- loopy/transform/iname.py | 41 ++++++++++++++++++++++------------------ 1 file changed, 23 insertions(+), 18 deletions(-) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index f3369cea8..0d5228275 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -2022,8 +2022,6 @@ def map_domain(kernel, isl_map, within=None): # FIXME: Document # FIXME: Support within - # FIXME: Right now, this requires all inames in a domain (or none) to - # be mapped. That makes this awkward to use. # {{{ within processing (disabled for now) if within is not None: @@ -2044,8 +2042,8 @@ def map_domain(kernel, isl_map, within=None): if not isl_map.is_bijective(): raise LoopyError("isl_map must be bijective") - new_inames = frozenset(isl_map.get_var_dict(dim_type.out)) - old_inames = frozenset(isl_map.get_var_dict(dim_type.in_)) + transform_map_out_dims = frozenset(isl_map.get_var_dict(dim_type.out)) + transform_map_in_dims = frozenset(isl_map.get_var_dict(dim_type.in_)) # {{{ solve for representation of old inames in terms of new @@ -2055,7 +2053,7 @@ def map_domain(kernel, isl_map, within=None): from loopy.symbolic import aff_to_expr from pymbolic import var - for iname in old_inames: + for iname in transform_map_in_dims: substitutions[iname] = aff_to_expr( _find_aff_subst_from_map(iname, isl_map)) var_substitutions[var(iname)] = aff_to_expr( @@ -2066,20 +2064,27 @@ def map_domain(kernel, isl_map, within=None): # }}} - def process_set(s): + def process_iname_dom(s): + # Make sure the inames we're transforming are all present in the domain + # if not transform_map_in_dims.issubset(frozenset(s.get_var_dict())): + # raise LoopyError("transform map %s attempts to map inames " + # "not present in domain %s. Transform map input inames " + # "must be a subset of the domain inames." + # % (isl_map, s)) + var_dict = s.get_var_dict() - overlap = old_inames & frozenset(var_dict) + overlap = transform_map_in_dims & frozenset(var_dict) if not overlap: # inames in s are not present in transform map, don't change s return s - if len(overlap) != len(old_inames): - raise LoopyError("loop domain '%s' involves a part " - "of the map domain inames. Domains must " - "either involve all or none of the map domain " - "inames." % s) + if len(overlap) != len(transform_map_in_dims): + raise LoopyError( + "Transform map %s attempts to map variables that are not present " + "in domain %s. This is only allowed if *none* of the mapped " + "variables are found in the domain." % (isl_map, s)) # {{{ align dims of isl_map and s @@ -2137,22 +2142,22 @@ def process_set(s): # FIXME: Revive _project_out_only_if_all_instructions_in_within - new_domains = [process_set(dom) for dom in kernel.domains] + new_domains = [process_iname_dom(dom) for dom in kernel.domains] # {{{ update within_inames new_insns = [] for insn in kernel.instructions: - overlap = old_inames & insn.within_inames + overlap = transform_map_in_dims & insn.within_inames if overlap and within(kernel, insn): - if len(overlap) != len(old_inames): + if len(overlap) != len(transform_map_in_dims): raise LoopyError("instruction '%s' is within only a part " "of the map domain inames. Instructions must " "either be within all or none of the map domain " "inames." % insn.id) - insn = insn.copy( - within_inames=(insn.within_inames - old_inames) | new_inames) + insn = insn.copy(within_inames=( + insn.within_inames - transform_map_in_dims) | transform_map_out_dims) else: # leave insn unmodified pass @@ -2169,7 +2174,7 @@ def process_set(s): rule_mapping_context = SubstitutionRuleMappingContext( kernel.substitutions, kernel.get_var_name_generator()) ins = _MapDomainMapper(rule_mapping_context, within, - new_inames, substitutions) + transform_map_out_dims, substitutions) kernel = ins.map_kernel(kernel) kernel = rule_mapping_context.finish_kernel(kernel) From 8b3195fb3809594ddbe190a9a4c983f1099bade9 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 10 Aug 2021 18:07:12 -0500 Subject: [PATCH 08/41] test map_domain cases where there is a mismatch between the inames in the transform map and the inames in a domain --- test/test_transform.py | 58 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 58 insertions(+) diff --git a/test/test_transform.py b/test/test_transform.py index 91db5ceeb..a07d31d81 100644 --- a/test/test_transform.py +++ b/test/test_transform.py @@ -780,6 +780,64 @@ def test_map_domain_with_transform_map_missing_dims(): # }}} + # {{{ Make sure there's an error if transform map contains *extra* dims + # Note, this is only okay if transform map 'in' dims don't match *any* + # domain inames, in which case nothing happens. + + # Not bijective + transform_map = isl.BasicMap( + "[nx,nt] -> {[t, y, rogue] -> [t_new, y_new]: " + "y = y_new and t = t_new" + "}") + + from loopy.diagnostic import LoopyError + knl_map_dom = ref_knl + try: + knl_map_dom = lp.map_domain(knl_map_dom, transform_map) + raise AssertionError() + except LoopyError as err: + assert "map must be bijective" in str(err) + + # Bijective and rogue dim + # (with some inames missing from in-dims, which would otherwise be okay) + transform_map = isl.BasicMap( + "[nx,nt] -> {[t, y, rogue] -> [t_new, y_new, rogue_new]: " + "y = y_new and t = t_new and rogue = rogue_new" + "}") + + try: + knl_map_dom = lp.map_domain(knl_map_dom, transform_map) + raise AssertionError() + except LoopyError as err: + assert ( + "attempts to map variables that are not present in domain" in str(err)) + + # Bijective and rogue dim + # (with all inames present in in-dims) + transform_map = isl.BasicMap( + "[nx,nt] -> {[t, y, x, z, rogue] -> [t_new, y_new, x_new, z_new, rogue_new]:" + "y = y_new and t = t_new and x = x_new and z = z_new and rogue = rogue_new" + "}") + + try: + knl_map_dom = lp.map_domain(knl_map_dom, transform_map) + raise AssertionError() + except LoopyError as err: + assert ( + "attempts to map variables that are not present in domain" in str(err)) + + # Bijective and rogue dim with *no* inames present in domain + # (allowed but does nothing) + transform_map = isl.BasicMap( + "[nx,nt] -> {[rogue] -> [rogue_new]: " + "rogue = rogue_new" + "}") + + # This should not raise an error + knl_map_dom = lp.map_domain(knl_map_dom, transform_map) + + # }}} + # }}} From 6ee3cf036e3b79080aa0c394438b76022e73ab62 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Wed, 11 Aug 2021 17:19:28 -0500 Subject: [PATCH 09/41] enforce that transform map in map_domain must apply to exactly one domain --- loopy/transform/iname.py | 41 +++++++++++++++++++--------------------- 1 file changed, 19 insertions(+), 22 deletions(-) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index 0d5228275..c11746799 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -2064,27 +2064,16 @@ def map_domain(kernel, isl_map, within=None): # }}} - def process_iname_dom(s): - # Make sure the inames we're transforming are all present in the domain - # if not transform_map_in_dims.issubset(frozenset(s.get_var_dict())): - # raise LoopyError("transform map %s attempts to map inames " - # "not present in domain %s. Transform map input inames " - # "must be a subset of the domain inames." - # % (isl_map, s)) + def process_set(s): + """Return the transformed set if transformation is possible, otherwise + return the original set. Also return an int representing + the number of sets that were transformed (0 or 1)""" - var_dict = s.get_var_dict() - - overlap = transform_map_in_dims & frozenset(var_dict) - - if not overlap: - # inames in s are not present in transform map, don't change s - return s - - if len(overlap) != len(transform_map_in_dims): - raise LoopyError( - "Transform map %s attempts to map variables that are not present " - "in domain %s. This is only allowed if *none* of the mapped " - "variables are found in the domain." % (isl_map, s)) + # Make sure the inames we're transforming are all present in the set + # (okay if map only transforms a *subset* of the inames in the set) + if not transform_map_in_dims.issubset(frozenset(s.get_var_dict())): + # Don't transform this set + return s, 0 # {{{ align dims of isl_map and s @@ -2138,11 +2127,19 @@ def process_iname_dom(s): new_s = _find_and_rename_dim( new_s, [dim_type.set], proxy_iname, real_iname) - return new_s + return new_s, 1 # FIXME: Revive _project_out_only_if_all_instructions_in_within - new_domains = [process_iname_dom(dom) for dom in kernel.domains] + new_doms_and_transform_ct = [process_set(dom) for dom in kernel.domains] + new_domains, transform_ct = zip(*new_doms_and_transform_ct) + if sum(transform_ct) != 1: + raise LoopyError( + "Transform map %s was applicable to %d domains. " + "Transform map must be applicable to exactly one domain. " + "A transform map is applicable to a domain if its input " + "inames are a subset of the domain inames." + % (isl_map, sum(transform_ct))) # {{{ update within_inames From f5b2257db8f0692ddb46e69787d57b938e4d9a5f Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Wed, 11 Aug 2021 17:19:53 -0500 Subject: [PATCH 10/41] test map validity checking/errors in map_domain --- test/test_transform.py | 90 +++++++++++++++++++++--------------------- 1 file changed, 46 insertions(+), 44 deletions(-) diff --git a/test/test_transform.py b/test/test_transform.py index a07d31d81..85911c028 100644 --- a/test/test_transform.py +++ b/test/test_transform.py @@ -692,20 +692,22 @@ def test_map_domain_vs_split_iname(): # }}} -# {{{ test_map_domain_with_transform_map_missing_dims +# {{{ test_map_domain_transform_map_validity_and_errors -def test_map_domain_with_transform_map_missing_dims(): - # Make sure map_domain works correctly when the mapping doesn't include - # all the dims in the domain. +def test_map_domain_transform_map_validity_and_errors(): # {{{ Make kernel knl = lp.make_kernel( [ "[nx,nt] -> {[x, y, z, t]: 0 <= x,y,z < nx and 0 <= t < nt}", + "[m] -> {[j]: 0 <= j < m}", ], """ a[y,x,t,z] = b[y,x,t,z] {id=stmta} + for j + <>temp = j {dep=stmta} + end """, lang_version=(2018, 2), ) @@ -714,6 +716,9 @@ def test_map_domain_with_transform_map_missing_dims(): # }}} + # Make sure map_domain works correctly when the mapping doesn't include + # all the dims in the domain. + # {{{ Apply domain change mapping knl_map_dom = ref_knl # loop priority goes away, deps stay @@ -729,10 +734,10 @@ def test_map_domain_with_transform_map_missing_dims(): "y = y_new" "}") - # Call map_domain to transform kernel + # Call map_domain to transform kernel; this should not produce an error knl_map_dom = lp.map_domain(knl_map_dom, transform_map) - # Prioritize loops (prio should eventually be updated in map_domain?) + # Prioritize loops (prio should eventually be updated in map_domain) try: # Use constrain_loop_nesting if it's available desired_prio = "x, t_outer, t_inner, z, y_new" @@ -780,9 +785,7 @@ def test_map_domain_with_transform_map_missing_dims(): # }}} - # {{{ Make sure there's an error if transform map contains *extra* dims - # Note, this is only okay if transform map 'in' dims don't match *any* - # domain inames, in which case nothing happens. + # {{{ Make sure we error on a map that is not bijective # Not bijective transform_map = isl.BasicMap( @@ -798,43 +801,42 @@ def test_map_domain_with_transform_map_missing_dims(): except LoopyError as err: assert "map must be bijective" in str(err) - # Bijective and rogue dim - # (with some inames missing from in-dims, which would otherwise be okay) - transform_map = isl.BasicMap( - "[nx,nt] -> {[t, y, rogue] -> [t_new, y_new, rogue_new]: " - "y = y_new and t = t_new and rogue = rogue_new" - "}") - - try: - knl_map_dom = lp.map_domain(knl_map_dom, transform_map) - raise AssertionError() - except LoopyError as err: - assert ( - "attempts to map variables that are not present in domain" in str(err)) - - # Bijective and rogue dim - # (with all inames present in in-dims) - transform_map = isl.BasicMap( - "[nx,nt] -> {[t, y, x, z, rogue] -> [t_new, y_new, x_new, z_new, rogue_new]:" - "y = y_new and t = t_new and x = x_new and z = z_new and rogue = rogue_new" - "}") - - try: - knl_map_dom = lp.map_domain(knl_map_dom, transform_map) - raise AssertionError() - except LoopyError as err: - assert ( - "attempts to map variables that are not present in domain" in str(err)) + # }}} - # Bijective and rogue dim with *no* inames present in domain - # (allowed but does nothing) - transform_map = isl.BasicMap( - "[nx,nt] -> {[rogue] -> [rogue_new]: " - "rogue = rogue_new" - "}") + # {{{ Make sure there's an error if transform map does not apply to + # exactly one domain. + + test_maps = [ + # Map where some inames match exactly one domain but there's also a + # rogue dim + isl.BasicMap( + "[nx,nt] -> {[t, y, rogue] -> [t_new, y_new, rogue_new]: " + "y = y_new and t = t_new and rogue = rogue_new" + "}"), + # Map where all inames match exactly one domain but there's also a + # rogue dim + isl.BasicMap( + "[nx,nt] -> {[t, y, x, z, rogue] -> " + "[t_new, y_new, x_new, z_new, rogue_new]: " + "y = y_new and t = t_new and x = x_new and z = z_new " + "and rogue = rogue_new" + "}"), + # Map where no inames match any domain + isl.BasicMap( + "[nx,nt] -> {[rogue] -> [rogue_new]: " + "rogue = rogue_new" + "}"), + ] - # This should not raise an error - knl_map_dom = lp.map_domain(knl_map_dom, transform_map) + for transform_map in test_maps: + try: + knl_map_dom = lp.map_domain(knl_map_dom, transform_map) + raise AssertionError() + except LoopyError as err: + assert ( + "was applicable to 0 domains. " + "Transform map must be applicable to exactly one domain." + in str(err)) # }}} From fddd305f1e0c0e2fda1423222194b03a693388b7 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Wed, 11 Aug 2021 18:05:09 -0500 Subject: [PATCH 11/41] remove map applicability logic from process_set --- loopy/transform/iname.py | 63 ++++++++++++++++++++++++++++------------ 1 file changed, 45 insertions(+), 18 deletions(-) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index c11746799..594ee0032 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -2065,15 +2065,8 @@ def map_domain(kernel, isl_map, within=None): # }}} def process_set(s): - """Return the transformed set if transformation is possible, otherwise - return the original set. Also return an int representing - the number of sets that were transformed (0 or 1)""" - - # Make sure the inames we're transforming are all present in the set - # (okay if map only transforms a *subset* of the inames in the set) - if not transform_map_in_dims.issubset(frozenset(s.get_var_dict())): - # Don't transform this set - return s, 0 + """Return the transformed set. Assume that map is applicable to this + set.""" # {{{ align dims of isl_map and s @@ -2127,19 +2120,53 @@ def process_set(s): new_s = _find_and_rename_dim( new_s, [dim_type.set], proxy_iname, real_iname) - return new_s, 1 + return new_s # FIXME: Revive _project_out_only_if_all_instructions_in_within - new_doms_and_transform_ct = [process_set(dom) for dom in kernel.domains] - new_domains, transform_ct = zip(*new_doms_and_transform_ct) - if sum(transform_ct) != 1: + # {{{ Apply the transform map to exactly one domain + + map_applied_to_one_dom = False + new_domains = [] + transform_map_rules = ( + "Transform map must be applicable to exactly one domain. " + "A transform map is applicable to a domain if its input " + "inames are a subset of the domain inames.") + + for old_domain in kernel.domains: + + # Make sure transform map is applicable to this set. Then transform. + + if not transform_map_in_dims.issubset( + frozenset(old_domain.get_var_dict())): + + # Map transforms inames that are not all present in the set. + # Don't transform. + new_domains.append(old_domain) + continue + + elif map_applied_to_one_dom: + + # Map is applicable to this domain, but this map was + # already applied. Error. + raise LoopyError( + "Transform map %s was applicable to more than one domain. %s" + % (isl_map, transform_map_rules)) + + else: + + # Map is applicable to this domain, and this map has not yet + # been applied. Transform. + new_domains.append(process_set(old_domain)) + map_applied_to_one_dom = True + + # If the map could not be applied to any domain, error. + if not map_applied_to_one_dom: raise LoopyError( - "Transform map %s was applicable to %d domains. " - "Transform map must be applicable to exactly one domain. " - "A transform map is applicable to a domain if its input " - "inames are a subset of the domain inames." - % (isl_map, sum(transform_ct))) + "Transform map %s was not applicable to any domain. %s" + % (isl_map, transform_map_rules)) + + # }}} # {{{ update within_inames From be524ff23e20d01885d930920601cfcb9e14b7b4 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Wed, 11 Aug 2021 18:05:29 -0500 Subject: [PATCH 12/41] update expected error string for map_domain test --- test/test_transform.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_transform.py b/test/test_transform.py index 85911c028..d72d35ba2 100644 --- a/test/test_transform.py +++ b/test/test_transform.py @@ -834,7 +834,7 @@ def test_map_domain_transform_map_validity_and_errors(): raise AssertionError() except LoopyError as err: assert ( - "was applicable to 0 domains. " + "was not applicable to any domain. " "Transform map must be applicable to exactly one domain." in str(err)) From e5b765c23e725ab690b829338be2164b16d24201 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Fri, 13 Aug 2021 14:19:43 -0500 Subject: [PATCH 13/41] remove arg from map_domain --- loopy/transform/iname.py | 38 ++++++-------------------------------- 1 file changed, 6 insertions(+), 32 deletions(-) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index 594ee0032..8d1f06d7b 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -1837,11 +1837,9 @@ def add_inames_to_insn(kernel, inames, insn_match): # {{{ map_domain class _MapDomainMapper(RuleAwareIdentityMapper): - def __init__(self, rule_mapping_context, within, new_inames, substitutions): + def __init__(self, rule_mapping_context, new_inames, substitutions): super(_MapDomainMapper, self).__init__(rule_mapping_context) - self.within = within - self.old_inames = frozenset(substitutions) self.new_inames = new_inames @@ -1850,10 +1848,7 @@ def __init__(self, rule_mapping_context, within, new_inames, substitutions): def map_reduction(self, expr, expn_state): red_overlap = frozenset(expr.inames) & self.old_inames arg_ctx_overlap = frozenset(expn_state.arg_context) & self.old_inames - if (red_overlap - and self.within( - expn_state.kernel, - expn_state.instruction)): + if red_overlap: if len(red_overlap) != len(self.old_inames): raise LoopyError("reduction '%s' involves a part " "of the map domain inames. Reductions must " @@ -1886,10 +1881,7 @@ def map_reduction(self, expr, expn_state): def map_variable(self, expr, expn_state): if (expr.name in self.old_inames - and expr.name not in expn_state.arg_context - and self.within( - expn_state.kernel, - expn_state.instruction)): + and expr.name not in expn_state.arg_context): return self.substitutions[expr.name] else: return super(_MapDomainMapper, self).map_variable(expr, expn_state) @@ -2012,7 +2004,7 @@ def _apply_identity_for_missing_map_dims(mapping, desired_dims): @for_each_kernel -def map_domain(kernel, isl_map, within=None): +def map_domain(kernel, isl_map): # FIXME: Express _split_iname_backend in terms of this # Missing/deleted for now: # - slab processing @@ -2021,24 +2013,6 @@ def map_domain(kernel, isl_map, within=None): # FIXME: Express affine_map_inames in terms of this, deprecate # FIXME: Document - # FIXME: Support within - - # {{{ within processing (disabled for now) - if within is not None: - raise NotImplementedError("within") - - from loopy.match import parse_match - within = parse_match(within) - - # {{{ return the same kernel if no kernel matches - - if not any(within(kernel, insn) for insn in kernel.instructions): - return kernel - - # }}} - - # }}} - if not isl_map.is_bijective(): raise LoopyError("isl_map must be bijective") @@ -2173,7 +2147,7 @@ def process_set(s): new_insns = [] for insn in kernel.instructions: overlap = transform_map_in_dims & insn.within_inames - if overlap and within(kernel, insn): + if overlap: if len(overlap) != len(transform_map_in_dims): raise LoopyError("instruction '%s' is within only a part " "of the map domain inames. Instructions must " @@ -2197,7 +2171,7 @@ def process_set(s): rule_mapping_context = SubstitutionRuleMappingContext( kernel.substitutions, kernel.get_var_name_generator()) - ins = _MapDomainMapper(rule_mapping_context, within, + ins = _MapDomainMapper(rule_mapping_context, transform_map_out_dims, substitutions) kernel = ins.map_kernel(kernel) From c0a613c5d1f0722cdaca71893387b05c97927b36 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Fri, 13 Aug 2021 15:13:53 -0500 Subject: [PATCH 14/41] in map_domain, error if there are any loop priorities/nest constraints involving the mapped inames --- loopy/transform/iname.py | 33 ++++++++++++++++++++++++++++++++- 1 file changed, 32 insertions(+), 1 deletion(-) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index 8d1f06d7b..a4e5fcee6 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -2003,22 +2003,53 @@ def _apply_identity_for_missing_map_dims(mapping, desired_dims): return augmented_mapping, proxy_name_pairs +def _error_if_any_iname_in_constraint( + inames, nest_constraints, + constraint_descriptor_str): + for constraint in nest_constraints: + for tier in constraint: + for iname in inames: + if tier.contains(iname): + raise ValueError( + "%s constraint %s contains iname(s) " + "transformed by map in map_domain." + % (constraint_descriptor_str, constraint)) + + @for_each_kernel def map_domain(kernel, isl_map): # FIXME: Express _split_iname_backend in terms of this # Missing/deleted for now: # - slab processing # - priorities processing - # FIXME: Process priorities # FIXME: Express affine_map_inames in terms of this, deprecate # FIXME: Document + # Make sure the map is bijective if not isl_map.is_bijective(): raise LoopyError("isl_map must be bijective") transform_map_out_dims = frozenset(isl_map.get_var_dict(dim_type.out)) transform_map_in_dims = frozenset(isl_map.get_var_dict(dim_type.in_)) + # {{{ Make sure that none of the mapped inames are involved in loop priorities + + if hasattr(kernel, "loop_priority") and kernel.loop_priority: + for prio in kernel.loop_priority: + if set(prio) & transform_map_in_dims: + raise ValueError( + "Loop priority %s contains iname(s) transformed by " + "map %s in map_domain." % (prio, isl_map)) + if hasattr(kernel, "loop_nest_constraints") and kernel.loop_nest_constraints: + _error_if_any_iname_in_constraint( + transform_map_in_dims, + kernel.loop_nest_constraints.must_nest, "Must-nest") + _error_if_any_iname_in_constraint( + transform_map_in_dims, + kernel.loop_nest_constraints.must_not_nest, "Must-not-nest") + + # }}} + # {{{ solve for representation of old inames in terms of new substitutions = {} From 88ebe1afa225840efaef2f06babebb96f005397e Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Fri, 13 Aug 2021 15:17:32 -0500 Subject: [PATCH 15/41] add test for map_domain error when there are any loop priorities/nest constraints involving the mapped inames --- test/test_transform.py | 24 +++++++++++++++++++++--- 1 file changed, 21 insertions(+), 3 deletions(-) diff --git a/test/test_transform.py b/test/test_transform.py index d72d35ba2..516d676c3 100644 --- a/test/test_transform.py +++ b/test/test_transform.py @@ -794,9 +794,9 @@ def test_map_domain_transform_map_validity_and_errors(): "}") from loopy.diagnostic import LoopyError - knl_map_dom = ref_knl + knl = ref_knl try: - knl_map_dom = lp.map_domain(knl_map_dom, transform_map) + knl = lp.map_domain(knl, transform_map) raise AssertionError() except LoopyError as err: assert "map must be bijective" in str(err) @@ -830,7 +830,7 @@ def test_map_domain_transform_map_validity_and_errors(): for transform_map in test_maps: try: - knl_map_dom = lp.map_domain(knl_map_dom, transform_map) + knl = lp.map_domain(knl, transform_map) raise AssertionError() except LoopyError as err: assert ( @@ -840,6 +840,24 @@ def test_map_domain_transform_map_validity_and_errors(): # }}} + # {{{ Make sure there's an error if we try to map inames in priorities + + knl = ref_knl + knl = lp.prioritize_loops(knl, "y, z") + knl = lp.prioritize_loops(knl, "x, z") + try: + transform_map = isl.BasicMap( + "[nx,nt] -> {[t, y] -> [t_new, y_new]: " + "y = y_new and t = t_new }") + knl = lp.map_domain(knl, transform_map) + raise AssertionError() + except ValueError as err: + assert ( + "Loop priority ('y', 'z') contains iname(s) " + "transformed by map" in str(err)) + + # }}} + # }}} From 13356df6dce015acc7bd48593649e1e59e2d588b Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Sat, 14 Aug 2021 16:47:36 -0500 Subject: [PATCH 16/41] in map_domain, rename isl_map->transform_map; add docstring --- loopy/transform/iname.py | 48 ++++++++++++++++++++++++---------------- 1 file changed, 29 insertions(+), 19 deletions(-) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index a4e5fcee6..0c6b27f1e 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -2017,20 +2017,29 @@ def _error_if_any_iname_in_constraint( @for_each_kernel -def map_domain(kernel, isl_map): +def map_domain(kernel, transform_map): + """Transform an iname domain by applying a mapping from existing inames to + new inames. + + :arg transform_map: A bijective :class:`islpy.Map` from existing inames to + new inames. To be applicable to a kernel domain, all input inames in + the map must be found in the domain. The map must be applicable to + exactly one domain found in *kernel.domains*. + + """ + # FIXME: Express _split_iname_backend in terms of this # Missing/deleted for now: # - slab processing # - priorities processing # FIXME: Express affine_map_inames in terms of this, deprecate - # FIXME: Document # Make sure the map is bijective - if not isl_map.is_bijective(): - raise LoopyError("isl_map must be bijective") + if not transform_map.is_bijective(): + raise LoopyError("transform_map must be bijective") - transform_map_out_dims = frozenset(isl_map.get_var_dict(dim_type.out)) - transform_map_in_dims = frozenset(isl_map.get_var_dict(dim_type.in_)) + transform_map_out_dims = frozenset(transform_map.get_var_dict(dim_type.out)) + transform_map_in_dims = frozenset(transform_map.get_var_dict(dim_type.in_)) # {{{ Make sure that none of the mapped inames are involved in loop priorities @@ -2039,7 +2048,7 @@ def map_domain(kernel, isl_map): if set(prio) & transform_map_in_dims: raise ValueError( "Loop priority %s contains iname(s) transformed by " - "map %s in map_domain." % (prio, isl_map)) + "map %s in map_domain." % (prio, transform_map)) if hasattr(kernel, "loop_nest_constraints") and kernel.loop_nest_constraints: _error_if_any_iname_in_constraint( transform_map_in_dims, @@ -2060,9 +2069,9 @@ def map_domain(kernel, isl_map): from pymbolic import var for iname in transform_map_in_dims: substitutions[iname] = aff_to_expr( - _find_aff_subst_from_map(iname, isl_map)) + _find_aff_subst_from_map(iname, transform_map)) var_substitutions[var(iname)] = aff_to_expr( - _find_aff_subst_from_map(iname, isl_map)) + _find_aff_subst_from_map(iname, transform_map)) applied_iname_rewrites.append(var_substitutions) del var_substitutions @@ -2073,21 +2082,22 @@ def process_set(s): """Return the transformed set. Assume that map is applicable to this set.""" - # {{{ align dims of isl_map and s + # {{{ align dims of transform_map and s from islpy import _align_dim_type map_with_s_domain = isl.Map.from_domain(s) - # If there are dims in s that are not mapped by isl_map, add them - # to the in/out space of isl_map so that they remain unchanged. + # If there are dims in s that are not mapped by transform_map, add them + # to the in/out space of transform_map so that they remain unchanged. # (temporary proxy dim names are needed in out space of transform # map because isl won't allow any dim names to match, i.e., instead # of just mapping {[unused_iname]->[unused_iname]}, we have to map # {[unused_name]->[unused_name__prox] : unused_name__prox = unused_name}, # and then rename unused_name__prox afterward.) - augmented_isl_map, proxy_name_pairs = _apply_identity_for_missing_map_dims( - isl_map, s.get_var_names(dim_type.set)) + augmented_transform_map, proxy_name_pairs = \ + _apply_identity_for_missing_map_dims( + transform_map, s.get_var_names(dim_type.set)) # FIXME: Make this less gross # FIXME: Make an exported/documented interface of this in islpy @@ -2098,9 +2108,9 @@ def process_set(s): for i in range(map_with_s_domain.dim(dt)) ] map_names = [ - augmented_isl_map.get_dim_name(dt, i) + augmented_transform_map.get_dim_name(dt, i) for dt in dim_types - for i in range(augmented_isl_map.dim(dt)) + for i in range(augmented_transform_map.dim(dt)) ] # (order doesn't matter in s_names/map_names, @@ -2109,7 +2119,7 @@ def process_set(s): # not sure why this isn't just handled inside _align_dim_type) aligned_map = _align_dim_type( dim_type.param, - augmented_isl_map, map_with_s_domain, False, + augmented_transform_map, map_with_s_domain, False, map_names, s_names) aligned_map = _align_dim_type( dim_type.in_, @@ -2156,7 +2166,7 @@ def process_set(s): # already applied. Error. raise LoopyError( "Transform map %s was applicable to more than one domain. %s" - % (isl_map, transform_map_rules)) + % (transform_map, transform_map_rules)) else: @@ -2169,7 +2179,7 @@ def process_set(s): if not map_applied_to_one_dom: raise LoopyError( "Transform map %s was not applicable to any domain. %s" - % (isl_map, transform_map_rules)) + % (transform_map, transform_map_rules)) # }}} From ab6076ad00bdea35a0d7a2935bdd299dc120df70 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Fri, 27 Aug 2021 09:37:56 -0500 Subject: [PATCH 17/41] clean up code for map_domain; clarify error message a bit --- loopy/transform/iname.py | 54 +++++++++++++++++++++++++--------------- 1 file changed, 34 insertions(+), 20 deletions(-) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index 0c6b27f1e..bdc94d8d5 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -2059,7 +2059,7 @@ def map_domain(kernel, transform_map): # }}} - # {{{ solve for representation of old inames in terms of new + # {{{ Solve for representation of old inames in terms of new substitutions = {} var_substitutions = {} @@ -2078,13 +2078,13 @@ def map_domain(kernel, transform_map): # }}} + # {{{ Function for applying mapping to one set + def process_set(s): """Return the transformed set. Assume that map is applicable to this set.""" - # {{{ align dims of transform_map and s - - from islpy import _align_dim_type + # {{{ Align dims of transform_map and s so that map can be applied map_with_s_domain = isl.Map.from_domain(s) @@ -2117,6 +2117,8 @@ def process_set(s): # _align_dim_type just converts these to sets # to determine which names are in both the obj and template, # not sure why this isn't just handled inside _align_dim_type) + + from islpy import _align_dim_type aligned_map = _align_dim_type( dim_type.param, augmented_transform_map, map_with_s_domain, False, @@ -2139,6 +2141,8 @@ def process_set(s): # FIXME: Revive _project_out_only_if_all_instructions_in_within + # }}} + # {{{ Apply the transform map to exactly one domain map_applied_to_one_dom = False @@ -2155,8 +2159,8 @@ def process_set(s): if not transform_map_in_dims.issubset( frozenset(old_domain.get_var_dict())): - # Map transforms inames that are not all present in the set. - # Don't transform. + # Map not applicable to this set because map transforms at least + # one iname that is not present in the set. Don't transform. new_domains.append(old_domain) continue @@ -2175,7 +2179,8 @@ def process_set(s): new_domains.append(process_set(old_domain)) map_applied_to_one_dom = True - # If the map could not be applied to any domain, error. + # If we get this far, either the map has been applied to 1 domain (good) + # or the map could not be applied to any domain, which should produce an error. if not map_applied_to_one_dom: raise LoopyError( "Transform map %s was not applicable to any domain. %s" @@ -2183,31 +2188,40 @@ def process_set(s): # }}} - # {{{ update within_inames + # {{{ Update within_inames for each statement - new_insns = [] - for insn in kernel.instructions: - overlap = transform_map_in_dims & insn.within_inames + # If we get this far, we know that the map was applied to exactly one domain, + # and that all the inames in transform_map_in_dims were transformed to + # inames in transform_map_out_dims. However, it's still possible that for some + # statements, stmt.within_inames will contain at least one but not all of the + # transformed inames (transform_map_in_dims). + # In this case, it's not clear what within_inames should be. Therefore, we + # require that if any transformed inames are found in stmt.within_inames, + # ALL transformed inames must be found in stmt.within_inames. + + new_stmts = [] + for stmt in kernel.instructions: + overlap = transform_map_in_dims & stmt.within_inames if overlap: if len(overlap) != len(transform_map_in_dims): - raise LoopyError("instruction '%s' is within only a part " - "of the map domain inames. Instructions must " - "either be within all or none of the map domain " - "inames." % insn.id) + raise LoopyError("Statement '%s' is within only a part " + "of the mapped inames in transformation map %s. " + "Statements must be within all or none of the mapped " + "inames." % (stmt.id, transform_map)) - insn = insn.copy(within_inames=( - insn.within_inames - transform_map_in_dims) | transform_map_out_dims) + stmt = stmt.copy(within_inames=( + stmt.within_inames - transform_map_in_dims) | transform_map_out_dims) else: - # leave insn unmodified + # Leave stmt unmodified pass - new_insns.append(insn) + new_stmts.append(stmt) # }}} kernel = kernel.copy( domains=new_domains, - instructions=new_insns, + instructions=new_stmts, applied_iname_rewrites=applied_iname_rewrites) rule_mapping_context = SubstitutionRuleMappingContext( From d710f8d3918d60debbd6c62bb8ac4776603bf762 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Fri, 27 Aug 2021 09:38:31 -0500 Subject: [PATCH 18/41] add map_domain test for handling of case where stmt.within_inames contains some but not all mapped inames --- test/test_transform.py | 46 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/test/test_transform.py b/test/test_transform.py index 516d676c3..152f8b5de 100644 --- a/test/test_transform.py +++ b/test/test_transform.py @@ -858,6 +858,52 @@ def test_map_domain_transform_map_validity_and_errors(): # }}} + # {{{ Make sure we error when stmt.within_inames contains at least one but + # not all mapped inames + + # {{{ Make kernel + + knl = lp.make_kernel( + [ + "[n, m] -> { [i, j]: 0 <= i < n and 0 <= j < m }", + "[ell] -> { [k]: 0 <= k < ell }", + ], + """ + for i + <>t0 = i {id=stmt0} + for j + <>t1 = j {id=stmt1, dep=stmt0} + end + <>t2 = i + 1 {id=stmt2, dep=stmt1} + end + for k + <>t3 = k {id=stmt3, dep=stmt2} + end + """, + lang_version=(2018, 2), + ) + + # }}} + + # This should fail: + try: + transform_map = isl.BasicMap( + "[n, m] -> {[i, j] -> [i_new, j_new]: " + "i_new = i + j and j_new = 2 + i }") + knl = lp.map_domain(knl, transform_map) + raise AssertionError() + except LoopyError as err: + assert ( + "Statements must be within all or none of the mapped inames" + in str(err)) + + # This should succeed: + transform_map = isl.BasicMap( + "[n, m] -> {[i] -> [i_new]: i_new = i + 2 }") + knl = lp.map_domain(knl, transform_map) + + # }}} + # }}} From 48e2c9a632c72edf4c8122aa434e7b85d5a4742a Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Fri, 27 Aug 2021 09:48:24 -0500 Subject: [PATCH 19/41] clarify comment in map_domain --- loopy/transform/iname.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index bdc94d8d5..1fca57fbc 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -2090,11 +2090,12 @@ def process_set(s): # If there are dims in s that are not mapped by transform_map, add them # to the in/out space of transform_map so that they remain unchanged. - # (temporary proxy dim names are needed in out space of transform - # map because isl won't allow any dim names to match, i.e., instead - # of just mapping {[unused_iname]->[unused_iname]}, we have to map + # We cannot just map {[unused_iname]->[unused_iname]} because isl won't + # allow any dim names to match, so temporary proxy dim names are needed + # in out space of transform map. I.e., we apply mapping # {[unused_name]->[unused_name__prox] : unused_name__prox = unused_name}, - # and then rename unused_name__prox afterward.) + # and then rename unused_name__prox afterward. + augmented_transform_map, proxy_name_pairs = \ _apply_identity_for_missing_map_dims( transform_map, s.get_var_names(dim_type.set)) From 6e7951a47a033ab24a4f19f6ffabe8b88421c352 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Sat, 28 Aug 2021 15:04:24 -0500 Subject: [PATCH 20/41] remove fixme --- test/test_transform.py | 1 - 1 file changed, 1 deletion(-) diff --git a/test/test_transform.py b/test/test_transform.py index 152f8b5de..86b7738a0 100644 --- a/test/test_transform.py +++ b/test/test_transform.py @@ -920,7 +920,6 @@ def test_diamond_tiling(ctx_factory, interactive=False): - u[ix, it]) """) - # FIXME: Handle priorities in map_domain knl_for_transform = ref_knl ref_knl = lp.prioritize_loops(ref_knl, "it, ix") From 061a8ecc8ec34747a7eca22ec51374337ef7bf36 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Sat, 28 Aug 2021 15:48:34 -0500 Subject: [PATCH 21/41] clean up and document _apply_identity_for_missing_map_dims --- loopy/transform/iname.py | 39 ++++++++++++++++++++++++++++----------- 1 file changed, 28 insertions(+), 11 deletions(-) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index 1fca57fbc..0cffb6edd 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -1970,23 +1970,35 @@ def _find_and_rename_dim(old_map, dim_types, old_name, new_name): def _apply_identity_for_missing_map_dims(mapping, desired_dims): + """For every variable v in *desired_dims* that is not found in the + input space for *mapping*, add input dimension v, output dimension + v_'proxy'_, and constraint v = v_'proxy'_ to the mapping. Also return a + list of the (v, v_'proxy'_) pairs. + """ + + # If the transform map in map_domain (below) does not contain all the + # inames in the iname domain (set) to which it is applied, the missing + # inames must be added to the transform map so that intersect_domain() + # doesn't remove them from the iname domain when the map is applied. + + # No two map dimension names can match, so we create a unique name for each + # new variable in the output dimension by appending _'proxy'_, and return a + # list of the (v, v_'proxy'_) pairs so that the proxy dims can be + # identified and replaced later. - # If dims in s are missing from transform map, they need to be added - # so that, e.g, intersect_domain doesn't remove them. - # (assume ordering will be handled afterward) + # (Apostrophes are not allowed in inames, so this suffix + # will not match any existing inames. This function is also used on + # dependency maps, which may contain variable names consisting of an iname + # suffixed with a single apostrophe.) + + # {{{ Find any missing vars and add them to the input and output space missing_dims = list( set(desired_dims) - set(mapping.get_var_names(dim_type.in_))) augmented_mapping = _add_and_name_isl_dims( mapping, dim_type.in_, missing_dims) - # We want these missing inames to map to themselves so that the map - # has no effect on them. Unfortunatley isl will break if the - # names of the out dims aren't unique, so we will temporariliy rename them - # (and then plan to change the names back afterward). - - # FIXME: need better way to make sure proxy dim names are unique within map - missing_dims_proxies = [d+"__prox" for d in missing_dims] + missing_dims_proxies = [d+"_'prox'_" for d in missing_dims] assert not set(missing_dims_proxies) & set( augmented_mapping.get_var_dict().keys()) @@ -1995,11 +2007,16 @@ def _apply_identity_for_missing_map_dims(mapping, desired_dims): proxy_name_pairs = list(zip(missing_dims, missing_dims_proxies)) - # Set proxy iname equal to real iname with equality constraint + # }}} + + # {{{ Add identity constraint (v = v_'proxy'_) for each new pair of dims + for real_iname, proxy_iname in proxy_name_pairs: augmented_mapping = _add_eq_isl_constraint_from_names( augmented_mapping, proxy_iname, real_iname) + # }}} + return augmented_mapping, proxy_name_pairs From fc0e265c79d55a3db175edaa3b649db2ac302ad4 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Sat, 28 Aug 2021 15:51:32 -0500 Subject: [PATCH 22/41] document _error_if_any_iname_in_constraint --- loopy/transform/iname.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index 0cffb6edd..b83543b40 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -2021,8 +2021,13 @@ def _apply_identity_for_missing_map_dims(mapping, desired_dims): def _error_if_any_iname_in_constraint( - inames, nest_constraints, - constraint_descriptor_str): + inames, nest_constraints, constraint_descriptor_str): + """Raise informative error if any iname in *inames* is constrained by any + nest constraint in *nest_constraints*. + """ + # (This function is only used when new machinery from + # new-loop-nest-constraints branch is detected.) + for constraint in nest_constraints: for tier in constraint: for iname in inames: From aee9c8313b7f6bd347e42b8b3d0f17083d285742 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Sat, 28 Aug 2021 15:57:18 -0500 Subject: [PATCH 23/41] minor changes to error messages in _MapDomainMapper and _find_aff_subst_from_map --- loopy/transform/iname.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index b83543b40..848d5d74a 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -1850,7 +1850,7 @@ def map_reduction(self, expr, expn_state): arg_ctx_overlap = frozenset(expn_state.arg_context) & self.old_inames if red_overlap: if len(red_overlap) != len(self.old_inames): - raise LoopyError("reduction '%s' involves a part " + raise LoopyError("Reduction '%s' involves a part " "of the map domain inames. Reductions must " "either involve all or none of the map domain " "inames." % str(expr)) @@ -1861,7 +1861,7 @@ def map_reduction(self, expr, expn_state): return super(_MapDomainMapper, self).map_reduction( expr, expn_state) else: - raise LoopyError("reduction '%s' has" + raise LoopyError("Reduction '%s' has" "some of the reduction variables affected " "by the map_domain shadowed by context. " "Either all or none must be shadowed." @@ -1928,7 +1928,7 @@ def _find_aff_subst_from_map(iname, isl_map): # not suitable, coefficient does not have unit coefficient continue - raise LoopyError("no suitable equation for '%s' found" % iname) + raise LoopyError("No suitable equation for '%s' found" % iname) def _add_and_name_isl_dims(isl_map, dt, names): From 468819d7e792e4160ac2527135b3d0f6ea510b1e Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Sat, 28 Aug 2021 16:02:52 -0500 Subject: [PATCH 24/41] vim fold markers for map_domain and associated functions --- loopy/transform/iname.py | 26 +++++++++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index 848d5d74a..78db2ded5 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -1834,7 +1834,9 @@ def add_inames_to_insn(kernel, inames, insn_match): # }}} -# {{{ map_domain +# {{{ map_domain and associated functions + +# {{{ _MapDomainMapper class _MapDomainMapper(RuleAwareIdentityMapper): def __init__(self, rule_mapping_context, new_inames, substitutions): @@ -1886,6 +1888,10 @@ def map_variable(self, expr, expn_state): else: return super(_MapDomainMapper, self).map_variable(expr, expn_state) +# }}} + + +# {{{ _find_aff_subst_from_map(iname, isl_map) def _find_aff_subst_from_map(iname, isl_map): if not isinstance(isl_map, isl.BasicMap): @@ -1930,6 +1936,10 @@ def _find_aff_subst_from_map(iname, isl_map): raise LoopyError("No suitable equation for '%s' found" % iname) +# }}} + + +# {{{ ISL map wrangling helper functions def _add_and_name_isl_dims(isl_map, dt, names): # (This function is also defined in independent, unmerged branch @@ -1968,6 +1978,10 @@ def _find_and_rename_dim(old_map, dim_types, old_name, new_name): dt, new_map.find_dim_by_name(dt, old_name), new_name) return new_map +# }}} + + +# {{{ _apply_identity_for_missing_map_dims(mapping, desired_dims) def _apply_identity_for_missing_map_dims(mapping, desired_dims): """For every variable v in *desired_dims* that is not found in the @@ -2019,6 +2033,10 @@ def _apply_identity_for_missing_map_dims(mapping, desired_dims): return augmented_mapping, proxy_name_pairs +# }}} + + +# {{{ _error_if_any_iname_in_constraint def _error_if_any_iname_in_constraint( inames, nest_constraints, constraint_descriptor_str): @@ -2037,6 +2055,10 @@ def _error_if_any_iname_in_constraint( "transformed by map in map_domain." % (constraint_descriptor_str, constraint)) +# }}} + + +# {{{ map_domain @for_each_kernel def map_domain(kernel, transform_map): @@ -2259,6 +2281,8 @@ def process_set(s): # }}} +# }}} + @for_each_kernel def add_inames_for_unused_hw_axes(kernel, within=None): From 1073f994a9d628a71b360ab67c644fc40d3c2738 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Sat, 28 Aug 2021 16:24:54 -0500 Subject: [PATCH 25/41] final cleanup of map_domain functions --- loopy/transform/iname.py | 46 +++++++++++++++++++++++----------------- 1 file changed, 27 insertions(+), 19 deletions(-) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index 78db2ded5..59b534f2d 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -2087,6 +2087,8 @@ def map_domain(kernel, transform_map): # {{{ Make sure that none of the mapped inames are involved in loop priorities + # kernel.loop_priority is being replaced with kernel.loop_nest_constraints, + # handle both attributes. if hasattr(kernel, "loop_priority") and kernel.loop_priority: for prio in kernel.loop_priority: if set(prio) & transform_map_in_dims: @@ -2122,7 +2124,7 @@ def map_domain(kernel, transform_map): # }}} - # {{{ Function for applying mapping to one set + # {{{ Function to apply mapping to one set def process_set(s): """Return the transformed set. Assume that map is applicable to this @@ -2130,38 +2132,41 @@ def process_set(s): # {{{ Align dims of transform_map and s so that map can be applied + # Create a map whose input space matches the set map_with_s_domain = isl.Map.from_domain(s) - # If there are dims in s that are not mapped by transform_map, add them - # to the in/out space of transform_map so that they remain unchanged. - # We cannot just map {[unused_iname]->[unused_iname]} because isl won't - # allow any dim names to match, so temporary proxy dim names are needed - # in out space of transform map. I.e., we apply mapping - # {[unused_name]->[unused_name__prox] : unused_name__prox = unused_name}, - # and then rename unused_name__prox afterward. + # {{{ Check for missing map dims and add them + + # For every iname v in the domain that is *not* found in the input + # space of the transform map, add input dimension v, output dimension + # v_'proxy'_, and constraint v = v_'proxy'_ to the transform map. + # Otherwise, v will be dropped from the domain when the map is applied. augmented_transform_map, proxy_name_pairs = \ _apply_identity_for_missing_map_dims( transform_map, s.get_var_names(dim_type.set)) - # FIXME: Make this less gross + # }}} + + # {{{ Align transform map input dims with set dims + # FIXME: Make an exported/documented interface of this in islpy + dim_types = [dim_type.param, dim_type.in_, dim_type.out] - s_names = [ + # Variables found in iname domain set + s_names = { map_with_s_domain.get_dim_name(dt, i) for dt in dim_types for i in range(map_with_s_domain.dim(dt)) - ] - map_names = [ + } + # Variables found in transform map + map_names = { augmented_transform_map.get_dim_name(dt, i) for dt in dim_types for i in range(augmented_transform_map.dim(dt)) - ] - - # (order doesn't matter in s_names/map_names, - # _align_dim_type just converts these to sets - # to determine which names are in both the obj and template, - # not sure why this isn't just handled inside _align_dim_type) + } + # (_align_dim_type uses these two sets to determine which names are in + # both the obj and template) from islpy import _align_dim_type aligned_map = _align_dim_type( @@ -2175,9 +2180,12 @@ def process_set(s): # }}} + # }}} + + # Apply the transform map to the domain new_s = aligned_map.intersect_domain(s).range() - # Now rename the proxy dims back to their original names + # Now rename any proxy dims back to their original names for real_iname, proxy_iname in proxy_name_pairs: new_s = _find_and_rename_dim( new_s, [dim_type.set], proxy_iname, real_iname) From e2ff759f0e7048ff207bcc9b20aec7a308d48637 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 31 Aug 2021 16:46:31 -0500 Subject: [PATCH 26/41] Update loopy/transform/iname.py MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit update arg without copy in _find_and_rename_dim Co-authored-by: Andreas Klöckner --- loopy/transform/iname.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index 59b534f2d..c8fcdcb32 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -1969,14 +1969,13 @@ def _add_eq_isl_constraint_from_names(isl_map, var1, var2): {1: 0, var1: 1, var2: -1})) -def _find_and_rename_dim(old_map, dim_types, old_name, new_name): +def _find_and_rename_dim(map, dim_types, old_name, new_name): # (This function is only used once here, but do not inline it; it is used many # times in child branch update-dependencies-during-transformations.) - new_map = old_map.copy() for dt in dim_types: - new_map = new_map.set_dim_name( - dt, new_map.find_dim_by_name(dt, old_name), new_name) - return new_map + map = map.set_dim_name( + dt, map.find_dim_by_name(dt, old_name), new_name) + return map # }}} From 4ce6c7a39111915eb8d93ef244c4e88eda6f6066 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 31 Aug 2021 16:51:47 -0500 Subject: [PATCH 27/41] Update loopy/transform/iname.py MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit in map_domain, use <= instead of issubset Co-authored-by: Andreas Klöckner --- loopy/transform/iname.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index c8fcdcb32..7b62f0bcc 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -2208,8 +2208,7 @@ def process_set(s): # Make sure transform map is applicable to this set. Then transform. - if not transform_map_in_dims.issubset( - frozenset(old_domain.get_var_dict())): + if not transform_map_in_dims <= frozenset(old_domain.get_var_dict()): # Map not applicable to this set because map transforms at least # one iname that is not present in the set. Don't transform. From a1534bd5dfab4ca9cdae3a6fd9696f9056f7f938 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 31 Aug 2021 17:24:14 -0500 Subject: [PATCH 28/41] move add_and_name_isl_dims, add_eq_isl_constraint_from_names, and find_and_rename_dim into isl_helpers --- loopy/isl_helpers.py | 50 +++++++++++++++++++++++++++++++++++++ loopy/transform/iname.py | 54 +++++++--------------------------------- 2 files changed, 59 insertions(+), 45 deletions(-) diff --git a/loopy/isl_helpers.py b/loopy/isl_helpers.py index d67df1154..448a62367 100644 --- a/loopy/isl_helpers.py +++ b/loopy/isl_helpers.py @@ -767,4 +767,54 @@ def subst_into_pwaff(new_space, pwaff, subst_dict): # }}} + +# {{{ add_and_name_isl_dims + +def add_and_name_isl_dims(isl_map, dt, names): + # (This function is also defined in independent, unmerged branch + # statement-instance-order-and-lex-order-map, and used in child branches + # thereof. Once these branches are all merged, it may make sense to move + # this function to a location for more general-purpose machinery. In the + # other branches, this function's name excludes the leading underscore.) + new_idx_start = isl_map.dim(dt) + new_map = isl_map.add_dims(dt, len(names)) + for i, name in enumerate(names): + new_map = new_map.set_dim_name(dt, new_idx_start+i, name) + return new_map + +# }}} + + +# {{{ add_eq_isl_constraint_from_names + +def add_eq_isl_constraint_from_names(isl_map, var1, var2): + # (This function is also defined in independent, unmerged branch + # statement-instance-order-and-lex-order-map, and used in child branches + # thereof. Once these branches are all merged, it may make sense to move + # this function to a location for more general-purpose machinery. In the + # other branches, this function's name excludes the leading underscore.) + + # add constraint var1 = var2 + + return isl_map.add_constraint( + isl.Constraint.eq_from_names( + isl_map.space, + {1: 0, var1: 1, var2: -1})) + +# }}} + + +# {{{ find_and_rename_dim + +def find_and_rename_dim(map, dim_types, old_name, new_name): + # (This function is only used once here, but do not inline it; it is used many + # times in child branch update-dependencies-during-transformations.) + for dt in dim_types: + map = map.set_dim_name( + dt, map.find_dim_by_name(dt, old_name), new_name) + return map + +# }}} + + # vim: foldmethod=marker diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index 7b62f0bcc..a94b2cb82 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -1939,47 +1939,6 @@ def _find_aff_subst_from_map(iname, isl_map): # }}} -# {{{ ISL map wrangling helper functions - -def _add_and_name_isl_dims(isl_map, dt, names): - # (This function is also defined in independent, unmerged branch - # statement-instance-order-and-lex-order-map, and used in child branches - # thereof. Once these branches are all merged, it may make sense to move - # this function to a location for more general-purpose machinery. In the - # other branches, this function's name excludes the leading underscore.) - new_idx_start = isl_map.dim(dt) - new_map = isl_map.add_dims(dt, len(names)) - for i, name in enumerate(names): - new_map = new_map.set_dim_name(dt, new_idx_start+i, name) - return new_map - - -def _add_eq_isl_constraint_from_names(isl_map, var1, var2): - # (This function is also defined in independent, unmerged branch - # statement-instance-order-and-lex-order-map, and used in child branches - # thereof. Once these branches are all merged, it may make sense to move - # this function to a location for more general-purpose machinery. In the - # other branches, this function's name excludes the leading underscore.) - - # add constraint var1 = var2 - - return isl_map.add_constraint( - isl.Constraint.eq_from_names( - isl_map.space, - {1: 0, var1: 1, var2: -1})) - - -def _find_and_rename_dim(map, dim_types, old_name, new_name): - # (This function is only used once here, but do not inline it; it is used many - # times in child branch update-dependencies-during-transformations.) - for dt in dim_types: - map = map.set_dim_name( - dt, map.find_dim_by_name(dt, old_name), new_name) - return map - -# }}} - - # {{{ _apply_identity_for_missing_map_dims(mapping, desired_dims) def _apply_identity_for_missing_map_dims(mapping, desired_dims): @@ -2004,18 +1963,21 @@ def _apply_identity_for_missing_map_dims(mapping, desired_dims): # dependency maps, which may contain variable names consisting of an iname # suffixed with a single apostrophe.) + from loopy.isl_helpers import ( + add_and_name_isl_dims, add_eq_isl_constraint_from_names) + # {{{ Find any missing vars and add them to the input and output space missing_dims = list( set(desired_dims) - set(mapping.get_var_names(dim_type.in_))) - augmented_mapping = _add_and_name_isl_dims( + augmented_mapping = add_and_name_isl_dims( mapping, dim_type.in_, missing_dims) missing_dims_proxies = [d+"_'prox'_" for d in missing_dims] assert not set(missing_dims_proxies) & set( augmented_mapping.get_var_dict().keys()) - augmented_mapping = _add_and_name_isl_dims( + augmented_mapping = add_and_name_isl_dims( augmented_mapping, dim_type.out, missing_dims_proxies) proxy_name_pairs = list(zip(missing_dims, missing_dims_proxies)) @@ -2025,7 +1987,7 @@ def _apply_identity_for_missing_map_dims(mapping, desired_dims): # {{{ Add identity constraint (v = v_'proxy'_) for each new pair of dims for real_iname, proxy_iname in proxy_name_pairs: - augmented_mapping = _add_eq_isl_constraint_from_names( + augmented_mapping = add_eq_isl_constraint_from_names( augmented_mapping, proxy_iname, real_iname) # }}} @@ -2185,8 +2147,10 @@ def process_set(s): new_s = aligned_map.intersect_domain(s).range() # Now rename any proxy dims back to their original names + + from loopy.isl_helpers import find_and_rename_dim for real_iname, proxy_iname in proxy_name_pairs: - new_s = _find_and_rename_dim( + new_s = find_and_rename_dim( new_s, [dim_type.set], proxy_iname, real_iname) return new_s From 3281b76ff2f8247b8cee83395b3a36a436161c21 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 31 Aug 2021 17:30:03 -0500 Subject: [PATCH 29/41] rename isl_map to isl_obj when sets are also allowed; remove comments about functions defined in other branches --- loopy/isl_helpers.py | 44 ++++++++++++++-------------------------- loopy/transform/iname.py | 8 ++++---- 2 files changed, 19 insertions(+), 33 deletions(-) diff --git a/loopy/isl_helpers.py b/loopy/isl_helpers.py index 448a62367..1be92d6f7 100644 --- a/loopy/isl_helpers.py +++ b/loopy/isl_helpers.py @@ -768,37 +768,25 @@ def subst_into_pwaff(new_space, pwaff, subst_dict): # }}} -# {{{ add_and_name_isl_dims - -def add_and_name_isl_dims(isl_map, dt, names): - # (This function is also defined in independent, unmerged branch - # statement-instance-order-and-lex-order-map, and used in child branches - # thereof. Once these branches are all merged, it may make sense to move - # this function to a location for more general-purpose machinery. In the - # other branches, this function's name excludes the leading underscore.) - new_idx_start = isl_map.dim(dt) - new_map = isl_map.add_dims(dt, len(names)) +# {{{ add_and_name_dims + +def add_and_name_dims(isl_obj, dt, names): + new_idx_start = isl_obj.dim(dt) + new_obj = isl_obj.add_dims(dt, len(names)) for i, name in enumerate(names): - new_map = new_map.set_dim_name(dt, new_idx_start+i, name) - return new_map + new_obj = new_obj.set_dim_name(dt, new_idx_start+i, name) + return new_obj # }}} -# {{{ add_eq_isl_constraint_from_names - -def add_eq_isl_constraint_from_names(isl_map, var1, var2): - # (This function is also defined in independent, unmerged branch - # statement-instance-order-and-lex-order-map, and used in child branches - # thereof. Once these branches are all merged, it may make sense to move - # this function to a location for more general-purpose machinery. In the - # other branches, this function's name excludes the leading underscore.) +# {{{ add_eq_constraint_from_names +def add_eq_constraint_from_names(isl_obj, var1, var2): # add constraint var1 = var2 - - return isl_map.add_constraint( + return isl_obj.add_constraint( isl.Constraint.eq_from_names( - isl_map.space, + isl_obj.space, {1: 0, var1: 1, var2: -1})) # }}} @@ -806,13 +794,11 @@ def add_eq_isl_constraint_from_names(isl_map, var1, var2): # {{{ find_and_rename_dim -def find_and_rename_dim(map, dim_types, old_name, new_name): - # (This function is only used once here, but do not inline it; it is used many - # times in child branch update-dependencies-during-transformations.) +def find_and_rename_dim(isl_obj, dim_types, old_name, new_name): for dt in dim_types: - map = map.set_dim_name( - dt, map.find_dim_by_name(dt, old_name), new_name) - return map + isl_obj = isl_obj.set_dim_name( + dt, isl_obj.find_dim_by_name(dt, old_name), new_name) + return isl_obj # }}} diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index a94b2cb82..d71b17dbb 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -1964,20 +1964,20 @@ def _apply_identity_for_missing_map_dims(mapping, desired_dims): # suffixed with a single apostrophe.) from loopy.isl_helpers import ( - add_and_name_isl_dims, add_eq_isl_constraint_from_names) + add_and_name_dims, add_eq_constraint_from_names) # {{{ Find any missing vars and add them to the input and output space missing_dims = list( set(desired_dims) - set(mapping.get_var_names(dim_type.in_))) - augmented_mapping = add_and_name_isl_dims( + augmented_mapping = add_and_name_dims( mapping, dim_type.in_, missing_dims) missing_dims_proxies = [d+"_'prox'_" for d in missing_dims] assert not set(missing_dims_proxies) & set( augmented_mapping.get_var_dict().keys()) - augmented_mapping = add_and_name_isl_dims( + augmented_mapping = add_and_name_dims( augmented_mapping, dim_type.out, missing_dims_proxies) proxy_name_pairs = list(zip(missing_dims, missing_dims_proxies)) @@ -1987,7 +1987,7 @@ def _apply_identity_for_missing_map_dims(mapping, desired_dims): # {{{ Add identity constraint (v = v_'proxy'_) for each new pair of dims for real_iname, proxy_iname in proxy_name_pairs: - augmented_mapping = add_eq_isl_constraint_from_names( + augmented_mapping = add_eq_constraint_from_names( augmented_mapping, proxy_iname, real_iname) # }}} From c7ac1f25c134983545ac7c63bc3173c1bf08d79f Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 31 Aug 2021 17:58:58 -0500 Subject: [PATCH 30/41] docstrings for add_and_name_isl_dims, add_eq_isl_constraint_from_names, and find_and_rename_dim --- loopy/isl_helpers.py | 52 +++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 51 insertions(+), 1 deletion(-) diff --git a/loopy/isl_helpers.py b/loopy/isl_helpers.py index 1be92d6f7..585bd4dcd 100644 --- a/loopy/isl_helpers.py +++ b/loopy/isl_helpers.py @@ -771,6 +771,23 @@ def subst_into_pwaff(new_space, pwaff, subst_dict): # {{{ add_and_name_dims def add_and_name_dims(isl_obj, dt, names): + """Append dimensions of the specified dimension type to the provided ISL + object, and set their names. + + :arg isl_obj: An :class:`islpy.Set` or :class:`islpy.Map` to which + new dimensions will be added. + + :arg dt: An :class:`islpy.dim_type`, i.e., an :class:`int`, specifying the + dimension type for the new dimensions. + + :arg names: An iterable of :class:`str` values specifying the names of the + new dimensions to be added. + + :returns: An object of the same type as *isl_obj* with the new dimensions + added and named. + + """ + new_idx_start = isl_obj.dim(dt) new_obj = isl_obj.add_dims(dt, len(names)) for i, name in enumerate(names): @@ -783,7 +800,21 @@ def add_and_name_dims(isl_obj, dt, names): # {{{ add_eq_constraint_from_names def add_eq_constraint_from_names(isl_obj, var1, var2): - # add constraint var1 = var2 + """Add constraint *var1* = *var2* to an ISL object. + + :arg isl_obj: An :class:`islpy.Set` or :class:`islpy.Map` to which + a new constraint will be added. + + :arg var1: A :class:`str` specifying the name of the first variable + involved in constraint *var1* = *var2*. + + :arg var2: A :class:`str` specifying the name of the second variable + involved in constraint *var1* = *var2*. + + :returns: An object of the same type as *isl_obj* with the constraint + *var1* = *var2*. + + """ return isl_obj.add_constraint( isl.Constraint.eq_from_names( isl_obj.space, @@ -795,6 +826,25 @@ def add_eq_constraint_from_names(isl_obj, var1, var2): # {{{ find_and_rename_dim def find_and_rename_dim(isl_obj, dim_types, old_name, new_name): + """Rename a dimension in an ISL object. + + :arg isl_obj: An :class:`islpy.Set` or :class:`islpy.Map` containing the + dimension to be renamed. + + :arg dim_types: An iterable of :class:`islpy.dim_type` values (i.e., + :class:`int` values) specifying the dimension types for any dimensions + to be renamed. + + :arg old_name: A :class:`str` specifying the name of the dimension to be + renamed. + + :arg new_name: A :class:`str` specifying the new name of the dimension to + be renamed. + + :returns: An object of the same type as *isl_obj* with the dimension + *old_name* renamed to *new_name*. + + """ for dt in dim_types: isl_obj = isl_obj.set_dim_name( dt, isl_obj.find_dim_by_name(dt, old_name), new_name) From 693ee86b63ec5639152e500cf0e0d87c188abe7b Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 31 Aug 2021 18:02:51 -0500 Subject: [PATCH 31/41] remove unnecessary variable renaming in add_and_name_dims --- loopy/isl_helpers.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/loopy/isl_helpers.py b/loopy/isl_helpers.py index 585bd4dcd..bcad428ba 100644 --- a/loopy/isl_helpers.py +++ b/loopy/isl_helpers.py @@ -789,10 +789,10 @@ def add_and_name_dims(isl_obj, dt, names): """ new_idx_start = isl_obj.dim(dt) - new_obj = isl_obj.add_dims(dt, len(names)) + isl_obj = isl_obj.add_dims(dt, len(names)) for i, name in enumerate(names): - new_obj = new_obj.set_dim_name(dt, new_idx_start+i, name) - return new_obj + isl_obj = isl_obj.set_dim_name(dt, new_idx_start+i, name) + return isl_obj # }}} From b8c3400250d79b49076e9185d8b6357700581726 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 31 Aug 2021 18:08:32 -0500 Subject: [PATCH 32/41] since find_and_rename_dim is only used for one dim type at a time, make it only accept a single dim type arg --- loopy/isl_helpers.py | 11 ++++------- loopy/transform/iname.py | 2 +- 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/loopy/isl_helpers.py b/loopy/isl_helpers.py index bcad428ba..57183109b 100644 --- a/loopy/isl_helpers.py +++ b/loopy/isl_helpers.py @@ -825,15 +825,14 @@ def add_eq_constraint_from_names(isl_obj, var1, var2): # {{{ find_and_rename_dim -def find_and_rename_dim(isl_obj, dim_types, old_name, new_name): +def find_and_rename_dim(isl_obj, dt, old_name, new_name): """Rename a dimension in an ISL object. :arg isl_obj: An :class:`islpy.Set` or :class:`islpy.Map` containing the dimension to be renamed. - :arg dim_types: An iterable of :class:`islpy.dim_type` values (i.e., - :class:`int` values) specifying the dimension types for any dimensions - to be renamed. + :arg dt: An :class:`islpy.dim_type` (i.e., :class:`int`) specifying the + dimension type containing the dimension to be renamed. :arg old_name: A :class:`str` specifying the name of the dimension to be renamed. @@ -845,10 +844,8 @@ def find_and_rename_dim(isl_obj, dim_types, old_name, new_name): *old_name* renamed to *new_name*. """ - for dt in dim_types: - isl_obj = isl_obj.set_dim_name( + return isl_obj.set_dim_name( dt, isl_obj.find_dim_by_name(dt, old_name), new_name) - return isl_obj # }}} diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index d71b17dbb..ccd0ea556 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -2151,7 +2151,7 @@ def process_set(s): from loopy.isl_helpers import find_and_rename_dim for real_iname, proxy_iname in proxy_name_pairs: new_s = find_and_rename_dim( - new_s, [dim_type.set], proxy_iname, real_iname) + new_s, dim_type.set, proxy_iname, real_iname) return new_s From 9801d1cffb4b20c61818f7716203986ca5d104d4 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 31 Aug 2021 18:19:27 -0500 Subject: [PATCH 33/41] more detailed documentation for _apply_identity_for_missing_map_dims --- loopy/transform/iname.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index ccd0ea556..5127d92c0 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -1946,6 +1946,16 @@ def _apply_identity_for_missing_map_dims(mapping, desired_dims): input space for *mapping*, add input dimension v, output dimension v_'proxy'_, and constraint v = v_'proxy'_ to the mapping. Also return a list of the (v, v_'proxy'_) pairs. + + :arg mapping: An :class:`islpy.Map`. + + :arg desired_dims: An iterable of :class:`str` specifying the names of the + desired map input dimensions. + + :returns: A two-tuple containing the mapping with the new dimensions and + constraints added, and a list of two-tuples of :class:`str` values + specifying the (v, v_'proxy'_) pairs. + """ # If the transform map in map_domain (below) does not contain all the From 09703eb041c7132fd8427085269ad14b7059aeb7 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 31 Aug 2021 18:27:12 -0500 Subject: [PATCH 34/41] in map_domain, don't recompute subst_from_map for no reason --- loopy/transform/iname.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index 5127d92c0..548f9ec01 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -2085,10 +2085,10 @@ def map_domain(kernel, transform_map): from loopy.symbolic import aff_to_expr from pymbolic import var for iname in transform_map_in_dims: - substitutions[iname] = aff_to_expr( - _find_aff_subst_from_map(iname, transform_map)) - var_substitutions[var(iname)] = aff_to_expr( - _find_aff_subst_from_map(iname, transform_map)) + subst_from_map = aff_to_expr( + _find_aff_subst_from_map(iname, transform_map)) + substitutions[iname] = subst_from_map + var_substitutions[var(iname)] = subst_from_map applied_iname_rewrites.append(var_substitutions) del var_substitutions From 3d24d437f6d291e4a40050164b1b33c1c7b0a23c Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 31 Aug 2021 18:36:56 -0500 Subject: [PATCH 35/41] use auto_test_vs_ref in test_map_domain_vs_split_iname --- test/test_transform.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/test/test_transform.py b/test/test_transform.py index 86b7738a0..38a15fe04 100644 --- a/test/test_transform.py +++ b/test/test_transform.py @@ -618,7 +618,7 @@ def _ensure_dim_names_match_and_align(obj_map, tgt_map): return align_spaces(obj_map, tgt_map) -def test_map_domain_vs_split_iname(): +def test_map_domain_vs_split_iname(ctx_factory): # {{{ Make kernel @@ -687,6 +687,9 @@ def test_map_domain_vs_split_iname(): # Can't easily compare instructions because equivalent subscript # expressions may have different orders + lp.auto_test_vs_ref(proc_knl_split_iname, ctx_factory(), proc_knl_map_dom, + parameters={"nx": 256, "nt": 256, "ni": 256}) + # }}} # }}} From 6594873494979511fd9c4aabcf547bc90f517ffd Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 31 Aug 2021 18:55:00 -0500 Subject: [PATCH 36/41] use get_attr instead of try/except to check for loop_nest_constraints attribute in kernel; also clarify purpose of some tests --- test/test_transform.py | 35 ++++++++++++++++++++++------------- 1 file changed, 22 insertions(+), 13 deletions(-) diff --git a/test/test_transform.py b/test/test_transform.py index 38a15fe04..a8c03d761 100644 --- a/test/test_transform.py +++ b/test/test_transform.py @@ -641,7 +641,7 @@ def test_map_domain_vs_split_iname(ctx_factory): # {{{ Apply domain change mapping - knl_map_dom = ref_knl # loop priority goes away, deps stay + knl_map_dom = ref_knl # Create map_domain mapping: import islpy as isl @@ -688,7 +688,7 @@ def test_map_domain_vs_split_iname(ctx_factory): # expressions may have different orders lp.auto_test_vs_ref(proc_knl_split_iname, ctx_factory(), proc_knl_map_dom, - parameters={"nx": 256, "nt": 256, "ni": 256}) + parameters={"nx": 128, "nt": 128, "ni": 128}) # }}} @@ -697,7 +697,7 @@ def test_map_domain_vs_split_iname(ctx_factory): # {{{ test_map_domain_transform_map_validity_and_errors -def test_map_domain_transform_map_validity_and_errors(): +def test_map_domain_transform_map_validity_and_errors(ctx_factory): # {{{ Make kernel @@ -719,12 +719,14 @@ def test_map_domain_transform_map_validity_and_errors(): # }}} - # Make sure map_domain works correctly when the mapping doesn't include - # all the dims in the domain. + # {{{ Make sure map_domain succeeds when we apply a map that includes 2 of 4 + # dims in the domain. - # {{{ Apply domain change mapping + # {{{ Apply domain change mapping that splits t and renames y; (similar to + # split_iname test above, but doesn't hurt to test this slightly different + # scenario) - knl_map_dom = ref_knl # loop priority goes away, deps stay + knl_map_dom = ref_knl # Create map_domain mapping that only includes t and y # (x and z should be unaffected) @@ -737,15 +739,17 @@ def test_map_domain_transform_map_validity_and_errors(): "y = y_new" "}") - # Call map_domain to transform kernel; this should not produce an error + # Call map_domain to transform kernel; this should *not* produce an error knl_map_dom = lp.map_domain(knl_map_dom, transform_map) - # Prioritize loops (prio should eventually be updated in map_domain) - try: - # Use constrain_loop_nesting if it's available + # Prioritize loops + + # Use constrain_loop_nesting if it's available + cln_attr = getattr(lp, "constrain_loop_nesting", None) + if cln_attr is not None: desired_prio = "x, t_outer, t_inner, z, y_new" knl_map_dom = lp.constrain_loop_nesting(knl_map_dom, desired_prio) - except AttributeError: + else: # For some reason, prioritize_loops can't handle the ordering above # when linearizing knl_split_iname below desired_prio = "z, y_new, x, t_outer, t_inner" @@ -758,7 +762,7 @@ def test_map_domain_transform_map_validity_and_errors(): # }}} - # {{{ Split iname and see if we get the same result + # {{{ Use split_iname, and rename_iname, and make sure we get the same result knl_split_iname = ref_knl knl_split_iname = lp.split_iname(knl_split_iname, "t", 32) @@ -786,6 +790,11 @@ def test_map_domain_transform_map_validity_and_errors(): # Can't easily compare instructions because equivalent subscript # expressions may have different orders + lp.auto_test_vs_ref(proc_knl_split_iname, ctx_factory(), proc_knl_map_dom, + parameters={"nx": 64, "nt": 64, "m": 64}) + + # }}} + # }}} # {{{ Make sure we error on a map that is not bijective From 94045941e8eaad77623b29ff5acacddd8b9e991b Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 31 Aug 2021 18:57:16 -0500 Subject: [PATCH 37/41] shrink problem size to make test_map_domain_transform_map_validity_and_errors run faster --- test/test_transform.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/test/test_transform.py b/test/test_transform.py index a8c03d761..29a6e2802 100644 --- a/test/test_transform.py +++ b/test/test_transform.py @@ -733,9 +733,9 @@ def test_map_domain_transform_map_validity_and_errors(ctx_factory): import islpy as isl transform_map = isl.BasicMap( "[nx,nt] -> {[t, y] -> [t_outer, t_inner, y_new]: " - "0 <= t_inner < 32 and " - "32*t_outer + t_inner = t and " - "0 <= 32*t_outer + t_inner < nt and " + "0 <= t_inner < 16 and " + "16*t_outer + t_inner = t and " + "0 <= 16*t_outer + t_inner < nt and " "y = y_new" "}") @@ -765,7 +765,7 @@ def test_map_domain_transform_map_validity_and_errors(ctx_factory): # {{{ Use split_iname, and rename_iname, and make sure we get the same result knl_split_iname = ref_knl - knl_split_iname = lp.split_iname(knl_split_iname, "t", 32) + knl_split_iname = lp.split_iname(knl_split_iname, "t", 16) knl_split_iname = lp.rename_iname(knl_split_iname, "y", "y_new") try: # Use constrain_loop_nesting if it's available @@ -791,7 +791,7 @@ def test_map_domain_transform_map_validity_and_errors(ctx_factory): # expressions may have different orders lp.auto_test_vs_ref(proc_knl_split_iname, ctx_factory(), proc_knl_map_dom, - parameters={"nx": 64, "nt": 64, "m": 64}) + parameters={"nx": 32, "nt": 32, "m": 32}) # }}} From 81886490f1d7c7464257f1f0b903f54a98b8db33 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 31 Aug 2021 19:00:17 -0500 Subject: [PATCH 38/41] use same loop prioritization regardless of whether we're using prioritize_loops() or constrain_loop_nesting() --- test/test_transform.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/test/test_transform.py b/test/test_transform.py index 29a6e2802..2eeea80ba 100644 --- a/test/test_transform.py +++ b/test/test_transform.py @@ -743,16 +743,13 @@ def test_map_domain_transform_map_validity_and_errors(ctx_factory): knl_map_dom = lp.map_domain(knl_map_dom, transform_map) # Prioritize loops + desired_prio = "x, t_outer, t_inner, z, y_new" # Use constrain_loop_nesting if it's available cln_attr = getattr(lp, "constrain_loop_nesting", None) if cln_attr is not None: - desired_prio = "x, t_outer, t_inner, z, y_new" knl_map_dom = lp.constrain_loop_nesting(knl_map_dom, desired_prio) else: - # For some reason, prioritize_loops can't handle the ordering above - # when linearizing knl_split_iname below - desired_prio = "z, y_new, x, t_outer, t_inner" knl_map_dom = lp.prioritize_loops(knl_map_dom, desired_prio) # Get a linearization From 434f04b13aabca268067ff9981d67b903258c6ac Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 31 Aug 2021 19:07:42 -0500 Subject: [PATCH 39/41] further clarification in comments about purpose of some map-domain tests --- test/test_transform.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/test/test_transform.py b/test/test_transform.py index 2eeea80ba..56bd4f77f 100644 --- a/test/test_transform.py +++ b/test/test_transform.py @@ -655,7 +655,8 @@ def test_map_domain_vs_split_iname(ctx_factory): knl_map_dom = lp.map_domain(knl_map_dom, transform_map) # Prioritize loops (prio should eventually be updated in map_domain?) - knl_map_dom = lp.prioritize_loops(knl_map_dom, "x, t_outer, t_inner") + loop_priority = "x, t_outer, t_inner" + knl_map_dom = lp.prioritize_loops(knl_map_dom, loop_priority) # Get a linearization proc_knl_map_dom = lp.preprocess_kernel(knl_map_dom) @@ -668,7 +669,7 @@ def test_map_domain_vs_split_iname(ctx_factory): knl_split_iname = ref_knl knl_split_iname = lp.split_iname(knl_split_iname, "t", 32) - knl_split_iname = lp.prioritize_loops(knl_split_iname, "x, t_outer, t_inner") + knl_split_iname = lp.prioritize_loops(knl_split_iname, loop_priority) proc_knl_split_iname = lp.preprocess_kernel(knl_split_iname) lin_knl_split_iname = lp.get_one_linearized_kernel( proc_knl_split_iname["loopy_kernel"], proc_knl_split_iname.callables_table) @@ -719,8 +720,8 @@ def test_map_domain_transform_map_validity_and_errors(ctx_factory): # }}} - # {{{ Make sure map_domain succeeds when we apply a map that includes 2 of 4 - # dims in the domain. + # {{{ Make sure map_domain *succeeds* when map includes 2 of 4 dims in one + # domain. # {{{ Apply domain change mapping that splits t and renames y; (similar to # split_iname test above, but doesn't hurt to test this slightly different @@ -759,7 +760,7 @@ def test_map_domain_transform_map_validity_and_errors(ctx_factory): # }}} - # {{{ Use split_iname, and rename_iname, and make sure we get the same result + # {{{ Use split_iname and rename_iname, and make sure we get the same result knl_split_iname = ref_knl knl_split_iname = lp.split_iname(knl_split_iname, "t", 16) @@ -870,7 +871,7 @@ def test_map_domain_transform_map_validity_and_errors(ctx_factory): # {{{ Make sure we error when stmt.within_inames contains at least one but # not all mapped inames - # {{{ Make kernel + # {{{ Make potentially problematic kernel knl = lp.make_kernel( [ From 332ebf55e839c909b19f939be8a24cc7e921d8a9 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 31 Aug 2021 19:25:19 -0500 Subject: [PATCH 40/41] tell pylint to chill out about a missing attribute --- test/test_transform.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_transform.py b/test/test_transform.py index 56bd4f77f..8c44a3f67 100644 --- a/test/test_transform.py +++ b/test/test_transform.py @@ -749,7 +749,7 @@ def test_map_domain_transform_map_validity_and_errors(ctx_factory): # Use constrain_loop_nesting if it's available cln_attr = getattr(lp, "constrain_loop_nesting", None) if cln_attr is not None: - knl_map_dom = lp.constrain_loop_nesting(knl_map_dom, desired_prio) + knl_map_dom = lp.constrain_loop_nesting(knl_map_dom, desired_prio) # noqa else: knl_map_dom = lp.prioritize_loops(knl_map_dom, desired_prio) From 6e1f1d9e7ce3fd56d53bd58841a1c04e7ffc0dba Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 31 Aug 2021 19:25:19 -0500 Subject: [PATCH 41/41] tell pylint to chill out about a missing attribute --- test/test_transform.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/test_transform.py b/test/test_transform.py index 56bd4f77f..3df2d6c2c 100644 --- a/test/test_transform.py +++ b/test/test_transform.py @@ -749,7 +749,8 @@ def test_map_domain_transform_map_validity_and_errors(ctx_factory): # Use constrain_loop_nesting if it's available cln_attr = getattr(lp, "constrain_loop_nesting", None) if cln_attr is not None: - knl_map_dom = lp.constrain_loop_nesting(knl_map_dom, desired_prio) + knl_map_dom = lp.constrain_loop_nesting( # noqa pylint:disable=no-member + knl_map_dom, desired_prio) else: knl_map_dom = lp.prioritize_loops(knl_map_dom, desired_prio)