From dd870ea116eba400495ec1d794ed4aa2b7e4a5c7 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 27 Jul 2021 16:24:33 -0500 Subject: [PATCH 01/18] add three functions that will be used for updating loop nest constraints during transformation (initially copied in from old loop-nest-constraints-v2 branch with minor changes) --- loopy/transform/iname.py | 280 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 280 insertions(+) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index bff63102d..67d5f0346 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -799,6 +799,286 @@ def get_graph_sources(graph): # }}} +# }}} + + +# {{{ Updating constraints during transformation + +# {{{ replace_inames_in_nest_constraints + +def replace_inames_in_nest_constraints( + inames_to_replace, replacement_inames, old_constraints, + coalesce_new_iname_duplicates=False, + ): + """ + :arg inames_to_replace: A set of inames that may exist in + `old_constraints`, each of which is to be replaced with all inames + in `replacement_inames`. + :arg replacement_inames: A set of inames, all of which will repalce each + iname in `inames_to_replace` in `old_constraints`. + :arg old_constraints: An iterable of tuples containing one or more + :class:`UnexpandedInameSet` objects. + """ + + # replace each iname in inames_to_replace + # with *all* inames in replacement_inames + + # loop through old_constraints and handle each nesting independently + new_constraints = set() + for old_nesting in old_constraints: + # loop through each iname_set in this nesting and perform replacement + new_nesting = [] + for iname_set in old_nesting: + + # find inames to be replaced + inames_found = inames_to_replace & iname_set.inames + + # create the new set of inames with the replacements + if inames_found: + new_inames = iname_set.inames - inames_found + new_inames.update(replacement_inames) + else: + new_inames = iname_set.inames.copy() + + new_nesting.append( + UnexpandedInameSet(new_inames, iname_set.complement)) + + # if we've removed things, new_nesting might only contain 1 item, + # in which case it's meaningless and we should just remove it + if len(new_nesting) > 1: + new_constraints.add(tuple(new_nesting)) + + # When joining inames, we may need to coalesce: + # e.g., if we join `i` and `j` into `ij`, and old_nesting was + # [{i, k}, {j, h}], at this point we have [{ij, k}, {ij, h}] + # which contains a cycle. If coalescing is enabled, change this + # to [{k}, ij, {h}] to remove the cycle. + if coalesce_new_iname_duplicates: + + def coalesce_duplicate_inames_in_nesting(nesting, coalesce_candidates): + # TODO would like this to be fully generic, but for now, assumes + # all UnexpandedInameSets have complement=False, which works if + # we're only using this for must_nest constraints since they cannot + # have complements + for iname_set in nesting: + assert not iname_set.complement + + import copy + # copy and convert nesting to list so we can modify + coalesced_nesting = list(copy.deepcopy(nesting)) + + # repeat coalescing step until we don't find any adjacent pairs + # containing duplicates (among coalesce_candidates) + found_duplicates = True + while found_duplicates: + found_duplicates = False + # loop through each iname_set in nesting and coalesce + # (assume new_nesting has at least 2 items) + i = 0 + while i < len(coalesced_nesting)-1: + iname_set_before = coalesced_nesting[i] + iname_set_after = coalesced_nesting[i+1] + # coalesce for each iname candidate + for iname in coalesce_candidates: + if (iname_set_before.inames == set([iname, ]) and + iname_set_after.inames == set([iname, ])): + # before/after contain single iname to be coalesced, + # -> remove iname_set_after + del coalesced_nesting[i+1] + found_duplicates = True + elif (iname_set_before.inames == set([iname, ]) and + iname in iname_set_after.inames): + # before contains single iname to be coalesced, + # after contains iname along with others, + # -> remove iname from iname_set_after.inames + coalesced_nesting[i+1] = UnexpandedInameSet( + inames=iname_set_after.inames - set([iname, ]), + complement=iname_set_after.complement, + ) + found_duplicates = True + elif (iname in iname_set_before.inames and + iname_set_after.inames == set([iname, ])): + # after contains single iname to be coalesced, + # before contains iname along with others, + # -> remove iname from iname_set_before.inames + coalesced_nesting[i] = UnexpandedInameSet( + inames=iname_set_before.inames - set([iname, ]), + complement=iname_set_before.complement, + ) + found_duplicates = True + elif (iname in iname_set_before.inames and + iname in iname_set_after.inames): + # before and after contain iname along with others, + # -> remove iname from iname_set_{before,after}.inames + # and insert it in between them + coalesced_nesting[i] = UnexpandedInameSet( + inames=iname_set_before.inames - set([iname, ]), + complement=iname_set_before.complement, + ) + coalesced_nesting[i+1] = UnexpandedInameSet( + inames=iname_set_after.inames - set([iname, ]), + complement=iname_set_after.complement, + ) + coalesced_nesting.insert(i+1, UnexpandedInameSet( + inames=set([iname, ]), + complement=False, + )) + found_duplicates = True + # else, iname was not found in both sets, so do nothing + i = i + 1 + + return tuple(coalesced_nesting) + + # loop through new_constraints; handle each nesting independently + coalesced_constraints = set() + for new_nesting in new_constraints: + coalesced_constraints.add( + coalesce_duplicate_inames_in_nesting( + new_nesting, replacement_inames)) + + return coalesced_constraints + else: + return new_constraints + +# }}} + + +# {{{ replace_inames_in_graph + +def replace_inames_in_graph( + inames_to_replace, replacement_inames, old_graph): + # replace each iname in inames_to_replace with all inames in replacement_inames + + new_graph = {} + iname_to_replace_found_as_key = False + union_of_inames_after_for_replaced_keys = set() + for iname, inames_after in old_graph.items(): + # create new inames_after + new_inames_after = inames_after.copy() + inames_found = inames_to_replace & new_inames_after + + if inames_found: + new_inames_after -= inames_found + new_inames_after.update(replacement_inames) + + # update dict + if iname in inames_to_replace: + iname_to_replace_found_as_key = True + union_of_inames_after_for_replaced_keys = \ + union_of_inames_after_for_replaced_keys | new_inames_after + # don't add this iname as a key in new graph, + # its replacements will be added below + else: + new_graph[iname] = new_inames_after + + # add replacement iname keys + if iname_to_replace_found_as_key: + for new_key in replacement_inames: + new_graph[new_key] = union_of_inames_after_for_replaced_keys.copy() + + # check for cycle + from pytools.graph import contains_cycle + if contains_cycle(new_graph): + raise ValueError( + "replace_inames_in_graph: Loop priority cycle detected. " + "Cannot replace inames %s with inames %s." + % (inames_to_replace, replacement_inames)) + + return new_graph + +# }}} + + +# {{{ replace_inames_in_all_nest_constraints + +def replace_inames_in_all_nest_constraints( + kernel, old_inames, new_inames, + coalesce_new_iname_duplicates=False, + pairs_that_must_not_voilate_constraints=set(), + ): + # replace each iname in old_inames with all inames in new_inames + + # get old must_nest and must_not_nest + # (must_nest_graph will be rebuilt) + if kernel.loop_nest_constraints: + old_must_nest = kernel.loop_nest_constraints.must_nest + old_must_not_nest = kernel.loop_nest_constraints.must_not_nest + # (these could still be None) + else: + old_must_nest = None + old_must_not_nest = None + + if old_must_nest: + # check to make sure special pairs don't conflict with constraints + for iname_before, iname_after in pairs_that_must_not_voilate_constraints: + if iname_before in kernel.loop_nest_constraints.must_nest_graph[ + iname_after]: + raise ValueError( + "Implied nestings violate existing must-nest constraints." + "\nimplied nestings: %s\nmust-nest constraints: %s" + % (pairs_that_must_not_voilate_constraints, old_must_nest)) + + new_must_nest = replace_inames_in_nest_constraints( + old_inames, new_inames, old_must_nest, + coalesce_new_iname_duplicates=coalesce_new_iname_duplicates, + ) + else: + new_must_nest = None + + if old_must_not_nest: + # check to make sure special pairs don't conflict with constraints + if not loop_nest_constraints_satisfied( + pairs_that_must_not_voilate_constraints, + must_not_nest_constraints=old_must_not_nest): + raise ValueError( + "Implied nestings violate existing must-not-nest constraints." + "\nimplied nestings: %s\nmust-not-nest constraints: %s" + % (pairs_that_must_not_voilate_constraints, old_must_not_nest)) + + new_must_not_nest = replace_inames_in_nest_constraints( + old_inames, new_inames, old_must_not_nest, + coalesce_new_iname_duplicates=False, + # (for now, never coalesce must-not-nest constraints) + ) + # each must not nest constraint may only contain two tiers + # TODO coalesce_new_iname_duplicates? + else: + new_must_not_nest = None + + # Rebuild must_nest graph + if new_must_nest: + new_must_nest_graph = {} + new_all_inames = ( + kernel.all_inames() - set(old_inames)) | set(new_inames) + from pytools.graph import CycleError + for must_nest_tuple in new_must_nest: + try: + new_must_nest_graph = add_to_must_nest_graph( + new_must_nest_graph, must_nest_tuple, new_all_inames) + except CycleError: + raise ValueError( + "Loop priority cycle detected when replacing inames %s " + "with inames %s. Previous must_nest constraints: %s" + % (old_inames, new_inames, old_must_nest)) + + # make sure none of the must_nest constraints violate must_not_nest + # this may not catch all problems + check_must_not_nest_against_must_nest_graph( + new_must_not_nest, new_must_nest_graph) + else: + new_must_nest_graph = None + + return kernel.copy( + loop_nest_constraints=LoopNestConstraints( + must_nest=new_must_nest, + must_not_nest=new_must_not_nest, + must_nest_graph=new_must_nest_graph, + ) + ) + +# }}} + + # }}} # }}} From 9826e73b6c10c65b873b08282303a1a31406b8d3 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 27 Jul 2021 17:21:11 -0500 Subject: [PATCH 02/18] refactor replace_inames_in_nest_constraints: move _coalesce_duplicate_inames_in_one_nesting to top level and document; clarify code+comments --- loopy/transform/iname.py | 245 ++++++++++++++++++++++++--------------- 1 file changed, 153 insertions(+), 92 deletions(-) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index 67d5f0346..abe58300e 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -806,34 +806,157 @@ def get_graph_sources(graph): # {{{ replace_inames_in_nest_constraints +# {{{ (helper) _coalesce_duplicate_inames_in_one_nesting + +def _coalesce_duplicate_inames_in_one_nesting(nesting, coalesce_candidates): + r"""Coalesce inames that are duplicated in adjacent loop tiers. + + :arg nesting: A tuple containing at least two + :class:`UnexpandedInameSet`\ s, neither of which may have + ``complement`` attribute set to ``True``. + + :arg coalesce_candidates: A set of inames to be coalesced. + + :returns: A tuple containing :class:`UnexpandedInameSet`\ s with iname + candidates that were duplicated in adjacent loop tiers coalesced. + For example, if ``nesting=({ij, k}, {ij, h})`` and ``ij`` is in + ``coalesce_candidates``, return ``({k}, {ij}, {h})``. + """ + + # For now, assume all UnexpandedInameSets have + # complement=False, which works if we're only using this for + # must_nest constraints since they cannot have complements + for iname_set in nesting: + assert not iname_set.complement + + # Copy and convert nesting to list so we can modify + import copy + coalesced_nesting = list(copy.deepcopy(nesting)) + + # Repeat coalescing step until we don't find any adjacent pairs + # containing duplicates (among coalesce_candidates) + + found_duplicates = True + while found_duplicates: + + # Loop through each adjacent pair of iname_sets (adjacent + # tiers) in nesting and coalesce in place (assume nesting has + # at least 2 items) + found_duplicates = False + i = 0 + while i < len(coalesced_nesting)-1: + + # Get pair of adjacent nesting tiers (iname sets) + iname_set_before = coalesced_nesting[i] + iname_set_after = coalesced_nesting[i+1] + + # {{{ For each iname candidate, coalesce if necessary + + for iname in coalesce_candidates: + + # Four possible cases depending on where iname is found + # and whether other inames exist in the tiers: + + if (iname_set_before.inames == set([iname, ]) and + iname_set_after.inames == set([iname, ])): + # Before/after both contain a lone iname to be coalesced, + # -> remove iname_set_after + del coalesced_nesting[i+1] + found_duplicates = True + elif (iname_set_before.inames == set([iname, ]) and + iname in iname_set_after.inames): + # Before contains a lone iname to be coalesced, + # after contains iname along with others, + # -> remove iname from iname_set_after.inames + coalesced_nesting[i+1] = UnexpandedInameSet( + inames=iname_set_after.inames - set([iname, ]), + complement=iname_set_after.complement, + ) + found_duplicates = True + elif (iname in iname_set_before.inames and + iname_set_after.inames == set([iname, ])): + # After contains a lone iname to be coalesced, + # before contains iname along with others, + # -> remove iname from iname_set_before.inames + coalesced_nesting[i] = UnexpandedInameSet( + inames=iname_set_before.inames - set([iname, ]), + complement=iname_set_before.complement, + ) + found_duplicates = True + elif (iname in iname_set_before.inames and + iname in iname_set_after.inames): + # Before and after both contain iname along with others, + # -> remove iname from iname_set_{before,after}.inames + # and insert it in a new tier between them + coalesced_nesting[i] = UnexpandedInameSet( + inames=iname_set_before.inames - set([iname, ]), + complement=iname_set_before.complement, + ) + coalesced_nesting[i+1] = UnexpandedInameSet( + inames=iname_set_after.inames - set([iname, ]), + complement=iname_set_after.complement, + ) + coalesced_nesting.insert(i+1, UnexpandedInameSet( + inames=set([iname, ]), + complement=False, + )) + found_duplicates = True + + # Else, iname was not found in both sets, so do nothing + + # }}} + + # Increment i to move to next pair of tiers + i = i + 1 + + return tuple(coalesced_nesting) + +# }}} + + def replace_inames_in_nest_constraints( inames_to_replace, replacement_inames, old_constraints, coalesce_new_iname_duplicates=False, ): - """ + """Replace each iname in inames_to_replace with *all* inames in + replacement_inames + :arg inames_to_replace: A set of inames that may exist in - `old_constraints`, each of which is to be replaced with all inames + ``old_constraints``, each of which is to be replaced with all inames in `replacement_inames`. - :arg replacement_inames: A set of inames, all of which will repalce each - iname in `inames_to_replace` in `old_constraints`. + + :arg replacement_inames: A set of inames, all of which will replace each + iname in ``inames_to_replace`` in ``old_constraints``. + :arg old_constraints: An iterable of tuples containing one or more :class:`UnexpandedInameSet` objects. + + :arg coalesce_new_iname_duplicates: A :class:`bool` determining whether to + coalesce replaced inames that are duplicated in adjacent loop tiers as + a result of the replacement. For example, if must-nest constraint + ``({i, k}, {j, h})`` exists, and then a transformation joins ``i`` and + ``j`` into ``ij`` replacing ``i`` and ``j`` with ``ij`` would yield + constraint ``({ij, k}, {ij, h})``, which contains a cycle. If + coalescing is enabled, this new constraint will be + changed to to ``({k}, {ij}, {h})``, which removes the cycle. + """ - # replace each iname in inames_to_replace - # with *all* inames in replacement_inames + # {{{ Replace each iname in inames_to_replace w/ all inames in replacement_inames - # loop through old_constraints and handle each nesting independently new_constraints = set() + for old_nesting in old_constraints: - # loop through each iname_set in this nesting and perform replacement + + # {{{ Perform replacement on each iname_set in this nesting + new_nesting = [] for iname_set in old_nesting: - # find inames to be replaced + # Find inames to be replaced inames_found = inames_to_replace & iname_set.inames - # create the new set of inames with the replacements + # Create the new set of inames with the replacements if inames_found: new_inames = iname_set.inames - inames_found new_inames.update(replacement_inames) @@ -843,103 +966,41 @@ def replace_inames_in_nest_constraints( new_nesting.append( UnexpandedInameSet(new_inames, iname_set.complement)) - # if we've removed things, new_nesting might only contain 1 item, - # in which case it's meaningless and we should just remove it + # If we've removed things, new_nesting might only contain 1 item, + # in which case it's meaningless and we should just exclude it if len(new_nesting) > 1: new_constraints.add(tuple(new_nesting)) - # When joining inames, we may need to coalesce: - # e.g., if we join `i` and `j` into `ij`, and old_nesting was - # [{i, k}, {j, h}], at this point we have [{ij, k}, {ij, h}] - # which contains a cycle. If coalescing is enabled, change this - # to [{k}, ij, {h}] to remove the cycle. - if coalesce_new_iname_duplicates: + # }}} - def coalesce_duplicate_inames_in_nesting(nesting, coalesce_candidates): - # TODO would like this to be fully generic, but for now, assumes - # all UnexpandedInameSets have complement=False, which works if - # we're only using this for must_nest constraints since they cannot - # have complements - for iname_set in nesting: - assert not iname_set.complement - - import copy - # copy and convert nesting to list so we can modify - coalesced_nesting = list(copy.deepcopy(nesting)) - - # repeat coalescing step until we don't find any adjacent pairs - # containing duplicates (among coalesce_candidates) - found_duplicates = True - while found_duplicates: - found_duplicates = False - # loop through each iname_set in nesting and coalesce - # (assume new_nesting has at least 2 items) - i = 0 - while i < len(coalesced_nesting)-1: - iname_set_before = coalesced_nesting[i] - iname_set_after = coalesced_nesting[i+1] - # coalesce for each iname candidate - for iname in coalesce_candidates: - if (iname_set_before.inames == set([iname, ]) and - iname_set_after.inames == set([iname, ])): - # before/after contain single iname to be coalesced, - # -> remove iname_set_after - del coalesced_nesting[i+1] - found_duplicates = True - elif (iname_set_before.inames == set([iname, ]) and - iname in iname_set_after.inames): - # before contains single iname to be coalesced, - # after contains iname along with others, - # -> remove iname from iname_set_after.inames - coalesced_nesting[i+1] = UnexpandedInameSet( - inames=iname_set_after.inames - set([iname, ]), - complement=iname_set_after.complement, - ) - found_duplicates = True - elif (iname in iname_set_before.inames and - iname_set_after.inames == set([iname, ])): - # after contains single iname to be coalesced, - # before contains iname along with others, - # -> remove iname from iname_set_before.inames - coalesced_nesting[i] = UnexpandedInameSet( - inames=iname_set_before.inames - set([iname, ]), - complement=iname_set_before.complement, - ) - found_duplicates = True - elif (iname in iname_set_before.inames and - iname in iname_set_after.inames): - # before and after contain iname along with others, - # -> remove iname from iname_set_{before,after}.inames - # and insert it in between them - coalesced_nesting[i] = UnexpandedInameSet( - inames=iname_set_before.inames - set([iname, ]), - complement=iname_set_before.complement, - ) - coalesced_nesting[i+1] = UnexpandedInameSet( - inames=iname_set_after.inames - set([iname, ]), - complement=iname_set_after.complement, - ) - coalesced_nesting.insert(i+1, UnexpandedInameSet( - inames=set([iname, ]), - complement=False, - )) - found_duplicates = True - # else, iname was not found in both sets, so do nothing - i = i + 1 - - return tuple(coalesced_nesting) + # }}} + + # {{{ Coalesce new iname duplicates + + # If coalesce_new_iname_duplicates=True, + # coalesce replaced inames that are duplicated in adjacent loop tiers as + # a result of the replacement. For example, if must-nest constraint + # ``({i, k}, {j, h})`` exists, and then a transformation joins ``i`` and + # ``j`` into ``ij`` replacing ``i`` and ``j`` with ``ij`` would yield + # constraint ``({ij, k}, {ij, h})``, which contains a cycle. If + # coalescing is enabled, this new constraint will be + # changed to to ``({k}, {ij}, {h})``, which removes the cycle. + + if coalesce_new_iname_duplicates: # loop through new_constraints; handle each nesting independently coalesced_constraints = set() for new_nesting in new_constraints: coalesced_constraints.add( - coalesce_duplicate_inames_in_nesting( + _coalesce_duplicate_inames_in_one_nesting( new_nesting, replacement_inames)) return coalesced_constraints else: return new_constraints + # }}} + # }}} From 2624adee9750bbd1562260138bc53e9d2e6a5592 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 27 Jul 2021 17:38:54 -0500 Subject: [PATCH 03/18] clearer comments in replace_inames_in_graph --- loopy/transform/iname.py | 38 +++++++++++++++++++++++++++----------- 1 file changed, 27 insertions(+), 11 deletions(-) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index abe58300e..c5b6658d3 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -1008,40 +1008,56 @@ def replace_inames_in_nest_constraints( def replace_inames_in_graph( inames_to_replace, replacement_inames, old_graph): - # replace each iname in inames_to_replace with all inames in replacement_inames - - new_graph = {} + """Replace each iname in inames_to_replace with all inames in + replacement_inames""" + # TODO more thorough documentation after initial code review + + # For any iname to replace is found as a key, their children will + # need to become children of the replacement inames. Keep track of all + # these children during initial replacement pass for graph values; then add + # them as children to the new keys afterward. iname_to_replace_found_as_key = False union_of_inames_after_for_replaced_keys = set() + + new_graph = {} + + # {{{ Replace graph values + for iname, inames_after in old_graph.items(): - # create new inames_after + + # Create new inames_after (graph children) with replacements new_inames_after = inames_after.copy() inames_found = inames_to_replace & new_inames_after - if inames_found: new_inames_after -= inames_found new_inames_after.update(replacement_inames) - # update dict + # If graph key is also one of the inames to be replaced, track its + # children but don't insert it into the new graph (replacement key will + # be added below). + # Otherwise, add original iname as key with new children to new graph if iname in inames_to_replace: iname_to_replace_found_as_key = True union_of_inames_after_for_replaced_keys = \ union_of_inames_after_for_replaced_keys | new_inames_after - # don't add this iname as a key in new graph, - # its replacements will be added below else: new_graph[iname] = new_inames_after - # add replacement iname keys + # }}} + + # {{{ Add replacement inames as keys, and add appropriate children + if iname_to_replace_found_as_key: for new_key in replacement_inames: new_graph[new_key] = union_of_inames_after_for_replaced_keys.copy() - # check for cycle + # }}} + + # Check for cycle from pytools.graph import contains_cycle if contains_cycle(new_graph): raise ValueError( - "replace_inames_in_graph: Loop priority cycle detected. " + "replace_inames_in_graph: Loop nesting constraint cycle detected. " "Cannot replace inames %s with inames %s." % (inames_to_replace, replacement_inames)) From 3e57099937b8b18c5bc3f3dc767903b85437a01a Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 27 Jul 2021 17:46:37 -0500 Subject: [PATCH 04/18] clean up code/comments in replace_inames_in_all_nest_constraints --- loopy/transform/iname.py | 31 +++++++++++++++++++++---------- 1 file changed, 21 insertions(+), 10 deletions(-) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index c5b6658d3..7d87ee7c0 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -1010,7 +1010,7 @@ def replace_inames_in_graph( inames_to_replace, replacement_inames, old_graph): """Replace each iname in inames_to_replace with all inames in replacement_inames""" - # TODO more thorough documentation after initial code review + # TODO more thorough docstring after initial code review # For any iname to replace is found as a key, their children will # need to become children of the replacement inames. Keep track of all @@ -1073,9 +1073,10 @@ def replace_inames_in_all_nest_constraints( coalesce_new_iname_duplicates=False, pairs_that_must_not_voilate_constraints=set(), ): - # replace each iname in old_inames with all inames in new_inames + """Replace each iname in old_inames with all inames in new_inames""" + # TODO more thorough docstring after initial code review - # get old must_nest and must_not_nest + # Get old must_nest and must_not_nest, if they exist # (must_nest_graph will be rebuilt) if kernel.loop_nest_constraints: old_must_nest = kernel.loop_nest_constraints.must_nest @@ -1085,8 +1086,10 @@ def replace_inames_in_all_nest_constraints( old_must_nest = None old_must_not_nest = None + # {{{ Perform replacement on must-nest constraints + if old_must_nest: - # check to make sure special pairs don't conflict with constraints + # Check to make sure special pairs don't conflict with constraints for iname_before, iname_after in pairs_that_must_not_voilate_constraints: if iname_before in kernel.loop_nest_constraints.must_nest_graph[ iname_after]: @@ -1095,6 +1098,7 @@ def replace_inames_in_all_nest_constraints( "\nimplied nestings: %s\nmust-nest constraints: %s" % (pairs_that_must_not_voilate_constraints, old_must_nest)) + # Compute new must nest constraints with replacements new_must_nest = replace_inames_in_nest_constraints( old_inames, new_inames, old_must_nest, coalesce_new_iname_duplicates=coalesce_new_iname_duplicates, @@ -1102,6 +1106,10 @@ def replace_inames_in_all_nest_constraints( else: new_must_nest = None + # }}} + + # {{{ Perform replacement on must-not-nest constraints + if old_must_not_nest: # check to make sure special pairs don't conflict with constraints if not loop_nest_constraints_satisfied( @@ -1112,17 +1120,19 @@ def replace_inames_in_all_nest_constraints( "\nimplied nestings: %s\nmust-not-nest constraints: %s" % (pairs_that_must_not_voilate_constraints, old_must_not_nest)) + # Compute new must not nest constraints with replacements new_must_not_nest = replace_inames_in_nest_constraints( old_inames, new_inames, old_must_not_nest, coalesce_new_iname_duplicates=False, # (for now, never coalesce must-not-nest constraints) ) - # each must not nest constraint may only contain two tiers - # TODO coalesce_new_iname_duplicates? else: new_must_not_nest = None - # Rebuild must_nest graph + # }}} + + # {{{ Rebuild must_nest graph + if new_must_nest: new_must_nest_graph = {} new_all_inames = ( @@ -1138,13 +1148,15 @@ def replace_inames_in_all_nest_constraints( "with inames %s. Previous must_nest constraints: %s" % (old_inames, new_inames, old_must_nest)) - # make sure none of the must_nest constraints violate must_not_nest - # this may not catch all problems + # Make sure none of the must_nest constraints violate must_not_nest + # (this may not catch all problems) check_must_not_nest_against_must_nest_graph( new_must_not_nest, new_must_nest_graph) else: new_must_nest_graph = None + # }}} + return kernel.copy( loop_nest_constraints=LoopNestConstraints( must_nest=new_must_nest, @@ -1155,7 +1167,6 @@ def replace_inames_in_all_nest_constraints( # }}} - # }}} # }}} From 462afeda060ac7eb8f277b9cde27bb904e6d0abf Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 27 Jul 2021 17:54:27 -0500 Subject: [PATCH 05/18] change default arg value from set to frozenset in replace_inames_in_all_nest_constraints --- loopy/transform/iname.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index 7d87ee7c0..52a269915 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -1071,7 +1071,7 @@ def replace_inames_in_graph( def replace_inames_in_all_nest_constraints( kernel, old_inames, new_inames, coalesce_new_iname_duplicates=False, - pairs_that_must_not_voilate_constraints=set(), + pairs_that_must_not_voilate_constraints=frozenset(), ): """Replace each iname in old_inames with all inames in new_inames""" # TODO more thorough docstring after initial code review From 375a6a2410932af7e2d1a014ec3afa84fb97f97e Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 27 Jul 2021 18:15:01 -0500 Subject: [PATCH 06/18] handle nest constraints during split_iname --- loopy/transform/iname.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index 52a269915..f2b9bbc66 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -1359,6 +1359,19 @@ def _split_iname_backend(kernel, iname_to_split, new_prio = new_prio + (prio_iname,) new_priorities.append(new_prio) + # {{{ Update nest constraints + + # Add {inner,outer} wherever iname_to_split is found in constraints. + # Let remove_unused_inames handle removal of the old iname if necessary + + # update must_nest, must_not_nest, and must_nest_graph + kernel = replace_inames_in_all_nest_constraints( + kernel, + set([iname_to_split, ]), [iname_to_split, inner_iname, outer_iname], + ) + + # }}} + kernel = kernel.copy( domains=new_domains, iname_slab_increments=iname_slab_increments, From 1a722738b3ab80d7b4478a03aa116db6b705e4a2 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 27 Jul 2021 18:18:33 -0500 Subject: [PATCH 07/18] handle nest constraints during remove_unused_inames --- loopy/transform/iname.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index f2b9bbc66..7a29a3bf2 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -2335,6 +2335,15 @@ def remove_unused_inames(kernel, inames=None): # }}} + # {{{ Remove inames from loop nest constraints + + kernel = replace_inames_in_all_nest_constraints( + kernel, old_inames=unused_inames, new_inames=[], + coalesce_new_iname_duplicates=False, + ) + + # }}} + return kernel From 73917af8b7e76925ea9055c2544fe630d0db4ca2 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 27 Jul 2021 18:25:45 -0500 Subject: [PATCH 08/18] test constraint updating during split_iname --- test/test_nest_constraints.py | 110 ++++++++++++++++++++++++++++++++++ 1 file changed, 110 insertions(+) diff --git a/test/test_nest_constraints.py b/test/test_nest_constraints.py index 70f821715..38bdd0a4a 100644 --- a/test/test_nest_constraints.py +++ b/test/test_nest_constraints.py @@ -58,6 +58,12 @@ def _process_and_linearize(prog, knl_name="loopy_kernel"): proc_prog[knl_name], proc_prog.callables_table) return lin_prog + +def _linearize_and_get_nestings(prog, knl_name="loopy_kernel"): + from loopy.transform.iname import get_iname_nestings + lin_knl = _process_and_linearize(prog, knl_name) + return get_iname_nestings(lin_knl.linearization) + # }}} @@ -643,6 +649,110 @@ def order_of_enterloops(lin_items): # }}} +# {{{ Test constraint updating during transformation + +# {{{ test_constraint_updating_split_iname + +def test_constraint_updating_split_iname(): + + ref_knl = lp.make_kernel( + "{ [g,h,i,j,k]: 0<=g,h,i,j,k 1: exec(sys.argv[1]) From 7ba7a3672772667333c65d6d3ffcc1222a135801 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 27 Jul 2021 18:33:45 -0500 Subject: [PATCH 09/18] update nest constraints during duplicate_inames --- loopy/transform/iname.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index 7a29a3bf2..76dd67fc3 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -1981,6 +1981,14 @@ def duplicate_inames(kernel, inames, within, new_inames=None, suffix=None, from loopy.kernel.tools import DomainChanger domch = DomainChanger(kernel, frozenset([old_iname])) + # {{{ Update nest constraints + + # (don't remove any unused inames yet, that happens later) + kernel = replace_inames_in_all_nest_constraints( + kernel, set([old_iname, ]), [old_iname, new_iname]) + + # }}} + from loopy.isl_helpers import duplicate_axes kernel = kernel.copy( domains=domch.get_domains_with( @@ -2010,6 +2018,18 @@ def duplicate_inames(kernel, inames, within, new_inames=None, suffix=None, # }}} + # Why isn't remove_unused_inames called on kernel here? + + # {{{ Remove any newly unused inames from nest constraints + + now_unused_inames = (set(inames) - get_used_inames(kernel)) & set(inames) + kernel = replace_inames_in_all_nest_constraints( + kernel, old_inames=now_unused_inames, new_inames=[], + coalesce_new_iname_duplicates=False, + ) + + # }}} + return kernel # }}} From 24231cb1b7c44654f8161c8e4e7e5f2385a3b890 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 27 Jul 2021 18:34:04 -0500 Subject: [PATCH 10/18] test updating of nest constraints during duplicate_inames --- test/test_nest_constraints.py | 103 ++++++++++++++++++++++++++++++++++ 1 file changed, 103 insertions(+) diff --git a/test/test_nest_constraints.py b/test/test_nest_constraints.py index 38bdd0a4a..6cb1999e9 100644 --- a/test/test_nest_constraints.py +++ b/test/test_nest_constraints.py @@ -750,6 +750,109 @@ def test_constraint_updating_split_iname(): # }}} + +# {{{ test_constraint_updating_duplicate_inames + +def test_constraint_updating_duplicate_inames(): + + ref_knl = lp.make_kernel( + "{ [g,h,i,j,k]: 0<=g,h,i,j,k Date: Tue, 27 Jul 2021 18:40:43 -0500 Subject: [PATCH 11/18] raise NotImplementedError in rename_iname when there are nest_constraints and does_exist=True --- loopy/transform/iname.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index 76dd67fc3..4e4becdce 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -2223,6 +2223,17 @@ def rename_iname(kernel, old_iname, new_iname, existing_ok=False, within=None): "--cannot rename" % new_iname) if does_exist: + + # FIXME implement loop nest constraint updating when does_exist=True + # (otherwise constraints handled in duplicate_iname and remove_unused_inames) + if kernel.loop_nest_constraints and ( + kernel.loop_nest_constraints.must_nest or + kernel.loop_nest_constraints.must_not_nest or + kernel.loop_nest_constraints.must_nest_graph): + raise NotImplementedError( + "rename_iname() does not yet handle new loop nest " + "constraints when does_exist=True.") + # {{{ check that the domains match up dom = kernel.get_inames_domain(frozenset((old_iname, new_iname))) From 335cbd9a2eba51ba0eb37b8f306b897aee541ce1 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 27 Jul 2021 18:46:45 -0500 Subject: [PATCH 12/18] test updating of nest constraints during rename_iname --- test/test_nest_constraints.py | 37 +++++++++++++++++++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/test/test_nest_constraints.py b/test/test_nest_constraints.py index 6cb1999e9..655c3c397 100644 --- a/test/test_nest_constraints.py +++ b/test/test_nest_constraints.py @@ -851,6 +851,43 @@ def test_constraint_updating_duplicate_inames(): # }}} + +# {{{ test_constraint_updating_rename_iname + +def test_constraint_updating_rename_iname(): + + ref_knl = lp.make_kernel( + "{ [g,h,i,j,k]: 0<=g,h,i,j,k Date: Tue, 27 Jul 2021 18:54:41 -0500 Subject: [PATCH 13/18] check for nest constraint conflicts in tag_inames --- loopy/transform/iname.py | 39 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index 4e4becdce..c318d7ed7 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -1855,6 +1855,7 @@ def parse_tag(tag): # }}} + from loopy.kernel.data import ConcurrentTag, VectorizeTag knl_inames = kernel.inames.copy() for name, new_tag in iname_to_tag.items(): if not new_tag: @@ -1865,6 +1866,44 @@ def parse_tag(tag): knl_inames[name] = knl_inames[name].tagged(new_tag) + # {{{ Check for incompatible loop nest constraints + + if isinstance(new_tag, VectorizeTag): + + # {{{ vec_inames will be nested innermost, check whether this + # conflicts with must-nest constraints + + must_nest_graph = (kernel.loop_nest_constraints.must_nest_graph + if kernel.loop_nest_constraints else None) + + if must_nest_graph and must_nest_graph.get(iname, set()): + # iname is not a leaf + raise ValueError( + "Loop priorities provided specify that iname %s nest " + "outside of inames %s, but vectorized inames " + "must nest innermost. Cannot tag %s with 'vec' tag." + % (iname, must_nest_graph.get(iname, set()), iname)) + + # }}} + + elif isinstance(new_tag, ConcurrentTag) and kernel.loop_nest_constraints: + + # {{{ Don't allow tagging of must_nest iname as concurrent + + must_nest = kernel.loop_nest_constraints.must_nest + + if must_nest: + for nesting in must_nest: + for iname_set in nesting: + if iname in iname_set.inames: + raise ValueError("cannot tag '%s' as concurrent--" + "iname involved in must-nest constraint %s." + % (iname, nesting)) + + # }}} + + # }}} + return kernel.copy(inames=knl_inames) # }}} From bf2cb81e83f700192a308137549372261c4591b0 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 27 Jul 2021 18:58:11 -0500 Subject: [PATCH 14/18] test checking of nest constraint conflicts in tag_inames --- test/test_nest_constraints.py | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/test/test_nest_constraints.py b/test/test_nest_constraints.py index 655c3c397..d684949ea 100644 --- a/test/test_nest_constraints.py +++ b/test/test_nest_constraints.py @@ -888,6 +888,39 @@ def test_constraint_updating_rename_iname(): # }}} + +# {{{ test_constraint_handling_tag_inames + +def test_constraint_handling_tag_inames(): + + ref_knl = lp.make_kernel( + "{ [g,h,i,j,k]: 0<=g,h,i,j,k Date: Tue, 27 Jul 2021 19:05:02 -0500 Subject: [PATCH 15/18] test adding of vec tag that conflicts with a must_nest constraint --- test/test_nest_constraints.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/test/test_nest_constraints.py b/test/test_nest_constraints.py index d684949ea..7a5a10cc0 100644 --- a/test/test_nest_constraints.py +++ b/test/test_nest_constraints.py @@ -458,6 +458,9 @@ def is_innermost(iname, lin_items): "iname h tagged with ConcurrentTag, " "cannot use iname in must-nest constraint" in str(e)) + # (testing these two operations in opposite order in + # test_constraint_handling_tag_inames) + # Add a must_not_nest constraint that conflicts with a vec tag # and attempt to linearize knl = ref_knl @@ -917,7 +920,15 @@ def test_constraint_handling_tag_inames(): "cannot tag 'i' as concurrent--iname involved in must-nest constraint" in str(e)) - # Need to test anything else here...? + # Add a vec tag that conflicts with a must_nest constraint + knl = ref_knl + knl = lp.constrain_loop_nesting(knl, must_nest=("{g,h,i,j}", "{k}")) + try: + knl = lp.tag_inames(knl, {"h": "vec"}) + raise AssertionError() + except ValueError as e: + assert ( + "vectorized inames must nest innermost" in str(e)) # }}} From a9d1d38ee235dbb2eb5be099746500a64fb2e31f Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 27 Jul 2021 19:26:32 -0500 Subject: [PATCH 16/18] update loop nest constraints during join_iname; remove pairs_that_must_not_voilate_constraints arg and processing from replace_inames_in_graph() --- loopy/transform/iname.py | 79 ++++++++++++++++++++++++++++++---------- 1 file changed, 59 insertions(+), 20 deletions(-) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index c318d7ed7..33c81bd03 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -1071,7 +1071,6 @@ def replace_inames_in_graph( def replace_inames_in_all_nest_constraints( kernel, old_inames, new_inames, coalesce_new_iname_duplicates=False, - pairs_that_must_not_voilate_constraints=frozenset(), ): """Replace each iname in old_inames with all inames in new_inames""" # TODO more thorough docstring after initial code review @@ -1089,15 +1088,6 @@ def replace_inames_in_all_nest_constraints( # {{{ Perform replacement on must-nest constraints if old_must_nest: - # Check to make sure special pairs don't conflict with constraints - for iname_before, iname_after in pairs_that_must_not_voilate_constraints: - if iname_before in kernel.loop_nest_constraints.must_nest_graph[ - iname_after]: - raise ValueError( - "Implied nestings violate existing must-nest constraints." - "\nimplied nestings: %s\nmust-nest constraints: %s" - % (pairs_that_must_not_voilate_constraints, old_must_nest)) - # Compute new must nest constraints with replacements new_must_nest = replace_inames_in_nest_constraints( old_inames, new_inames, old_must_nest, @@ -1111,15 +1101,6 @@ def replace_inames_in_all_nest_constraints( # {{{ Perform replacement on must-not-nest constraints if old_must_not_nest: - # check to make sure special pairs don't conflict with constraints - if not loop_nest_constraints_satisfied( - pairs_that_must_not_voilate_constraints, - must_not_nest_constraints=old_must_not_nest): - raise ValueError( - "Implied nestings violate existing must-not-nest constraints." - "\nimplied nestings: %s\nmust-not-nest constraints: %s" - % (pairs_that_must_not_voilate_constraints, old_must_not_nest)) - # Compute new must not nest constraints with replacements new_must_not_nest = replace_inames_in_nest_constraints( old_inames, new_inames, old_must_not_nest, @@ -1609,7 +1590,7 @@ def join_inames(kernel, inames, new_iname=None, tag=None, within=None): from loopy.match import parse_match within = parse_match(within) - # {{{ return the same kernel if no kernel matches + # {{{ return the same kernel if no insn matches if not any(within(kernel, insn) for insn in kernel.instructions): return kernel @@ -1704,6 +1685,64 @@ def subst_within_inames(fid): applied_iname_rewrites=kernel.applied_iname_rewrites + [subst_dict] )) + # {{{ Update loop nest constraints + + if kernel.loop_nest_constraints and ( + kernel.loop_nest_constraints.must_nest or + kernel.loop_nest_constraints.must_not_nest or + kernel.loop_nest_constraints.must_nest_graph): + + if within != parse_match(None): + raise NotImplementedError( + "join_inames() does not yet handle new loop nest " + "constraints when within is not None.") + + # When joining inames, we create *implied* loop nestings. + + # {{{ Make sure that these implied nestings don't violate existing + # constraints. + + implied_nestings = set() + inames_orig_order = inames[::-1] # This was reversed above + for i, iname_before in enumerate(inames_orig_order[:-1]): + for iname_after in inames_orig_order[i+1:]: + implied_nestings.add((iname_before, iname_after)) + + old_must_nest = kernel.loop_nest_constraints.must_nest + old_must_not_nest = kernel.loop_nest_constraints.must_not_nest + # (these could still be None) + + # Check to make sure special pairs don't conflict with constraints + if old_must_nest: + for iname_before, iname_after in implied_nestings: + if iname_before in kernel.loop_nest_constraints.must_nest_graph[ + iname_after]: + raise ValueError( + "Implied nestings violate existing must-nest constraints." + "\nimplied nestings: %s\nmust-nest constraints: %s" + % (implied_nestings, old_must_nest)) + + if old_must_not_nest: + if not loop_nest_constraints_satisfied( + implied_nestings, + must_not_nest_constraints=old_must_not_nest): + raise ValueError( + "Implied nestings violate existing must-not-nest constraints." + "\nimplied nestings: %s\nmust-not-nest constraints: %s" + % (implied_nestings, old_must_not_nest)) + + # }}} + + # When replacing inames, coalesce newly created iname duplicates + # (see docstring for replace_inames_in_all_nest_constraints) + kernel = replace_inames_in_all_nest_constraints( + kernel, set(inames), [new_iname], + coalesce_new_iname_duplicates=True, + ) + # (will fail if cycle is created in must-nest graph) + + # }}} + from loopy.match import parse_stack_match within = parse_stack_match(within) From 7f438e62508a2faabc22da1c2ef8dd17dd0f2981 Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 27 Jul 2021 19:27:48 -0500 Subject: [PATCH 17/18] test updating of nest constraints during join_inames --- test/test_nest_constraints.py | 83 +++++++++++++++++++++++++++++++++++ 1 file changed, 83 insertions(+) diff --git a/test/test_nest_constraints.py b/test/test_nest_constraints.py index 7a5a10cc0..6ad96c2a9 100644 --- a/test/test_nest_constraints.py +++ b/test/test_nest_constraints.py @@ -932,6 +932,89 @@ def test_constraint_handling_tag_inames(): # }}} + +# {{{ test_constraint_updating_join_inames + +def test_constraint_updating_join_inames(): + + ref_knl = lp.make_kernel( + "{ [g,h,i,j,k]: 0<=g,h,i,j,k<1024 }", + """ + out[g,h,i,j,k] = 2*a[g,h,i,j,k] {id=insn} + """, + ) + ref_knl = lp.add_and_infer_dtypes(ref_knl, {"a": np.dtype(np.float32)}) + + knl = ref_knl + knl = lp.constrain_loop_nesting( + knl, + must_nest=("i", "{g, h, j, k}"), + must_not_nest=("h", "g"), + ) + knl = lp.constrain_loop_nesting( + knl, + must_nest=("{g, h, i}", "{j, k}"), + ) + knl = lp.join_inames(knl, inames=["g", "h"], new_iname="gh") + loop_nesting = _linearize_and_get_nestings(knl)[0] # only one nesting + assert loop_nesting[0] == "i" + assert loop_nesting[1] == "gh" + assert set(loop_nesting[2:]) == set(["j", "k"]) + + knl = ref_knl + knl = lp.constrain_loop_nesting( + knl, + must_nest=("i", "{g, h, j, k}"), + must_not_nest=("h", "g"), + ) + knl = lp.constrain_loop_nesting( + knl, + must_nest=("{g, h, i}", "{j, k}"), + ) + knl = lp.join_inames(knl, inames=["j", "k"], new_iname="jk") + + loop_nesting = _linearize_and_get_nestings(knl)[0] # only one nesting + assert loop_nesting[0] == "i" + assert loop_nesting[1:3] == ("g", "h") + assert loop_nesting[3] == "jk" + + knl = ref_knl + knl = lp.constrain_loop_nesting( + knl, + must_nest=("h", "i", "g", "{j, k}"), + ) + knl = lp.join_inames(knl, inames=["i", "g"], new_iname="ig") + loop_nesting = _linearize_and_get_nestings(knl)[0] # only one nesting + assert loop_nesting[0] == "h" + assert loop_nesting[1] == "ig" + assert set(loop_nesting[2:4]) == set(["j", "k"]) + + # Test cycle detection + knl = ref_knl + knl = lp.constrain_loop_nesting( + knl, + must_nest=("i", "{g, h}", "{j, k}"), + ) + try: + lp.join_inames(knl, inames=["i", "k"], new_iname="ik") + raise AssertionError() + except ValueError as e: + assert "cycle" in str(e) + + # Test implied nesting that creates constraint violation + knl = ref_knl + knl = lp.constrain_loop_nesting( + knl, + must_not_nest=("i", "k"), + ) + try: + lp.join_inames(knl, inames=["i", "k"], new_iname="ik") + raise AssertionError() + except ValueError as e: + assert "Implied nestings violate existing must-not-nest" in str(e) + +# }}} + # }}} # }}} From 751ab73eb4d8bd48e8ac5fcc509098429dac115b Mon Sep 17 00:00:00 2001 From: jdsteve2 Date: Tue, 27 Jul 2021 19:29:54 -0500 Subject: [PATCH 18/18] test coalescing of newly created inames (during join_iname) --- test/test_nest_constraints.py | 147 ++++++++++++++++++++++++++++++++++ 1 file changed, 147 insertions(+) diff --git a/test/test_nest_constraints.py b/test/test_nest_constraints.py index 6ad96c2a9..1dbe8567d 100644 --- a/test/test_nest_constraints.py +++ b/test/test_nest_constraints.py @@ -1015,6 +1015,153 @@ def test_constraint_updating_join_inames(): # }}} + +# {{{ test_iname_coalescing_in_loop_nest_constraints + +def test_iname_coalescing_in_loop_nest_constraints(): + + def get_sets_of_inames(iname_sets_tuple, iname_universe): + # convert UnexpandedInameSets to sets + sets_of_inames = [] + for iname_set in iname_sets_tuple: + sets_of_inames.append( + iname_set.get_inames_represented(iname_universe)) + return sets_of_inames + + ref_knl = lp.make_kernel( + "{ [g,h,i,j,k]: 0<=g,h,i,j,k<1024 }", + """ + out[g,h,i,j,k] = 2*a[g,h,i,j,k] {id=insn} + """, + ) + # (join_inames errors if domain bound is variable) + + ref_knl = lp.add_and_infer_dtypes(ref_knl, {"a": np.dtype(np.float32)}) + + knl = ref_knl + knl = lp.constrain_loop_nesting( + knl, + must_nest=("i", "g", "h", "j", "k"), + ) + knl = lp.join_inames(knl, inames=["g", "h"], new_iname="gh") + knl = knl["loopy_kernel"] + new_must_nest = get_sets_of_inames( + list(knl.loop_nest_constraints.must_nest)[0], knl.all_inames()) + expected_must_nest = [ + set(["i", ]), set(["gh", ]), set(["j", ]), set(["k", ])] + assert new_must_nest == expected_must_nest + + knl = ref_knl + knl = lp.constrain_loop_nesting( + knl, + must_nest=("{i, g}", "h", "j", "k"), + ) + knl = lp.join_inames(knl, inames=["g", "h"], new_iname="gh") + knl = knl["loopy_kernel"] + new_must_nest = get_sets_of_inames( + list(knl.loop_nest_constraints.must_nest)[0], knl.all_inames()) + expected_must_nest = [ + set(["i", ]), set(["gh", ]), set(["j", ]), set(["k", ])] + assert new_must_nest == expected_must_nest + + knl = ref_knl + knl = lp.constrain_loop_nesting( + knl, + must_nest=("i", "g", "{h, j}", "k"), + ) + knl = lp.join_inames(knl, inames=["g", "h"], new_iname="gh") + knl = knl["loopy_kernel"] + new_must_nest = get_sets_of_inames( + list(knl.loop_nest_constraints.must_nest)[0], knl.all_inames()) + expected_must_nest = [ + set(["i", ]), set(["gh", ]), set(["j", ]), set(["k", ])] + assert new_must_nest == expected_must_nest + + knl = ref_knl + knl = lp.constrain_loop_nesting( + knl, + must_nest=("i", "g", "{h, j, k}"), + ) + knl = lp.join_inames(knl, inames=["g", "h"], new_iname="gh") + knl = knl["loopy_kernel"] + new_must_nest = get_sets_of_inames( + list(knl.loop_nest_constraints.must_nest)[0], knl.all_inames()) + expected_must_nest = [ + set(["i", ]), set(["gh", ]), set(["j", "k"])] + assert new_must_nest == expected_must_nest + + knl = ref_knl + knl = lp.constrain_loop_nesting( + knl, + must_nest=("i", "{g, h}", "j", "k"), + ) + knl = lp.join_inames(knl, inames=["g", "h"], new_iname="gh") + knl = knl["loopy_kernel"] + new_must_nest = get_sets_of_inames( + list(knl.loop_nest_constraints.must_nest)[0], knl.all_inames()) + expected_must_nest = [ + set(["i", ]), set(["gh", ]), set(["j", ]), set(["k", ])] + assert new_must_nest == expected_must_nest + + knl = ref_knl + knl = lp.constrain_loop_nesting( + knl, + must_nest=("{i, g}", "{h, j, k}"), + ) + knl = lp.join_inames(knl, inames=["g", "h"], new_iname="gh") + knl = knl["loopy_kernel"] + new_must_nest = get_sets_of_inames( + list(knl.loop_nest_constraints.must_nest)[0], knl.all_inames()) + expected_must_nest = [ + set(["i", ]), set(["gh", ]), set(["j", "k"])] + assert new_must_nest == expected_must_nest + + knl = ref_knl + knl = lp.constrain_loop_nesting( + knl, + must_nest=("i", "g", "j", "h", "k"), + ) + try: + knl = lp.join_inames(knl, inames=["g", "h"], new_iname="gh") + raise AssertionError() + except ValueError as e: + assert "contains cycle" in str(e) + + knl = ref_knl + knl = lp.constrain_loop_nesting( + knl, + must_nest=("{i, g}", "j", "{h, k}"), + ) + try: + knl = lp.join_inames(knl, inames=["g", "h"], new_iname="gh") + raise AssertionError() + except ValueError as e: + assert "contains cycle" in str(e) + + knl = ref_knl + knl = lp.constrain_loop_nesting( + knl, + must_nest=("{i, h}", "j", "{g, k}"), + ) + try: + knl = lp.join_inames(knl, inames=["g", "h"], new_iname="gh") + raise AssertionError() + except ValueError as e: + assert "nestings violate existing must-nest" in str(e) + + knl = ref_knl + knl = lp.constrain_loop_nesting( + knl, + must_not_nest=("g", "h"), + ) + try: + knl = lp.join_inames(knl, inames=["g", "h"], new_iname="gh") + raise AssertionError() + except ValueError as e: + assert "nestings violate existing must-not-nest" in str(e) + +# }}} + # }}} # }}}