diff --git a/doc/ref_kernel.rst b/doc/ref_kernel.rst index 922315685..09eceb1d1 100644 --- a/doc/ref_kernel.rst +++ b/doc/ref_kernel.rst @@ -220,6 +220,8 @@ Tag Meaning Identifiers ----------- +.. _reserved-identifiers: + Reserved Identifiers ^^^^^^^^^^^^^^^^^^^^ diff --git a/loopy/schedule/checker/__init__.py b/loopy/schedule/checker/__init__.py new file mode 100644 index 000000000..51100bc92 --- /dev/null +++ b/loopy/schedule/checker/__init__.py @@ -0,0 +1,144 @@ +__copyright__ = "Copyright (C) 2019 James Stevens" + +__license__ = """ +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +""" + + +# {{{ create a pairwise schedules for statement pairs + +def get_schedules_for_statement_pairs( + knl, + linearization_items, + insn_id_pairs, + ): + r"""For each statement pair in a subset of all statement pairs found in a + linearized kernel, determine the (relative) order in which the statement + instances are executed. For each pair, describe this relative ordering with + a pair of mappings from statement instances to points in a single + lexicographic ordering (a ``pairwise schedule''). When determining the + relative ordering, ignore concurrent inames. + + :arg knl: A preprocessed :class:`loopy.kernel.LoopKernel` containing the + linearization items that will be used to create a schedule. + + :arg linearization_items: A list of :class:`loopy.schedule.ScheduleItem` + (to be renamed to `loopy.schedule.LinearizationItem`) containing + all linearization items for which pairwise schedules will be + created. To allow usage of this routine during linearization, a + truncated (i.e. partial) linearization may be passed through this + argument. + + :arg insn_id_pairs: A list containing pairs of instruction + identifiers. + + :returns: A dictionary mapping each two-tuple of instruction identifiers + provided in `insn_id_pairs` to a corresponding two-tuple containing two + :class:`islpy.Map`\ s representing a pairwise schedule as two + mappings from statement instances to lexicographic time, one for + each of the two statements. + + .. doctest: + + >>> import loopy as lp + >>> import numpy as np + >>> # Make kernel ----------------------------------------------------------- + >>> knl = lp.make_kernel( + ... "{[i,j,k]: 0<=i>> knl = lp.add_and_infer_dtypes(knl, {"a": np.float32, "b": np.float32}) + >>> knl = lp.prioritize_loops(knl, "i,j") + >>> knl = lp.prioritize_loops(knl, "i,k") + >>> # Preprocess + >>> knl = lp.preprocess_kernel(knl) + >>> # Get a linearization + >>> knl = lp.get_one_linearized_kernel( + ... knl["loopy_kernel"], knl.callables_table) + >>> # Get a pairwise schedule ----------------------------------------------- + >>> from loopy.schedule.checker import get_schedules_for_statement_pairs + >>> # Get two maps ---------------------------------------------------------- + >>> schedules = get_schedules_for_statement_pairs( + ... knl, + ... knl.linearization, + ... [("insn_a", "insn_b")], + ... ) + >>> # Print maps + >>> print("\n".join( + ... str(m).replace("{ ", "{\n").replace(" :", "\n:") + ... for m in schedules[("insn_a", "insn_b")] + ... )) + [pi, pj, pk] -> { + [_lp_linchk_stmt = 0, i, j, k] -> [_lp_linchk_l0 = i, _lp_linchk_l1 = 0] + : 0 <= i < pi and 0 <= j < pj and 0 <= k < pk } + [pi, pj, pk] -> { + [_lp_linchk_stmt = 1, i, j, k] -> [_lp_linchk_l0 = i, _lp_linchk_l1 = 1] + : 0 <= i < pi and 0 <= j < pj and 0 <= k < pk } + + """ + + # {{{ make sure kernel has been preprocessed + + from loopy.kernel import KernelState + assert knl.state in [ + KernelState.PREPROCESSED, + KernelState.LINEARIZED] + + # }}} + + # {{{ Find any EnterLoop inames that are tagged as concurrent + # so that generate_pairwise_schedule() knows to ignore them + # (In the future, this shouldn't be necessary because there + # won't be any inames with ConcurrentTags in EnterLoop linearization items. + # Test which exercises this: test_linearization_checker_with_stroud_bernstein()) + from loopy.schedule.checker.utils import ( + partition_inames_by_concurrency, + get_EnterLoop_inames, + ) + conc_inames, _ = partition_inames_by_concurrency(knl) + enterloop_inames = get_EnterLoop_inames(linearization_items) + conc_loop_inames = conc_inames & enterloop_inames + + # The only concurrent EnterLoop inames should be Vec and ILP + from loopy.kernel.data import (VectorizeTag, IlpBaseTag) + for conc_iname in conc_loop_inames: + # Assert that there exists an ilp or vectorize tag (out of the + # potentially multiple other tags on this concurrent iname). + assert any( + isinstance(tag, (VectorizeTag, IlpBaseTag)) + for tag in knl.iname_to_tags[conc_iname]) + + # }}} + + # {{{ Create two mappings from {statement instance: lex point} + + # include only instructions involved in this dependency + from loopy.schedule.checker.schedule import generate_pairwise_schedules + return generate_pairwise_schedules( + knl, + linearization_items, + insn_id_pairs, + loops_to_ignore=conc_loop_inames, + ) + + # }}} + +# }}} diff --git a/loopy/schedule/checker/schedule.py b/loopy/schedule/checker/schedule.py new file mode 100644 index 000000000..64dd377f2 --- /dev/null +++ b/loopy/schedule/checker/schedule.py @@ -0,0 +1,315 @@ +__copyright__ = "Copyright (C) 2019 James Stevens" + +__license__ = """ +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +""" + +import islpy as isl + +__doc__ = """ + +.. data:: LIN_CHECK_IDENTIFIER_PREFIX + + The prefix for identifiers involved in linearization checking. + +.. data:: LEX_VAR_PREFIX + + E.g., a prefix of ``_lp_linchk_lex`` might yield lexicographic dimension + variables ``_lp_linchk_lex0``, ``_lp_linchk_lex1``, ``_lp_linchk_lex2``. Cf. + :ref:`reserved-identifiers`. + +.. data:: STATEMENT_VAR_NAME + + Set the :class:`str` specifying the prefix to be used for the variables + representing the dimensions in the lexicographic ordering used in a + pairwise schedule. + +""" + +LIN_CHECK_IDENTIFIER_PREFIX = "_lp_linchk_" +LEX_VAR_PREFIX = "%sl" % (LIN_CHECK_IDENTIFIER_PREFIX) +STATEMENT_VAR_NAME = "%sstmt" % (LIN_CHECK_IDENTIFIER_PREFIX) + + +def _pad_tuple_with_zeros(tup, desired_length): + return tup[:] + tuple([0]*(desired_length-len(tup))) + + +def _simplify_lex_dims(tup0, tup1): + """Simplify a pair of lex tuples in order to reduce the complexity of + resulting maps. Remove lex tuple dimensions with matching integer values + since these do not provide information on relative ordering. Once a + dimension is found where both tuples have non-matching integer values, + remove any faster-updating lex dimensions since they are not necessary + to specify a relative ordering. + """ + + new_tup0 = [] + new_tup1 = [] + + # Loop over dims from slowest updating to fastest + for d0, d1 in zip(tup0, tup1): + if isinstance(d0, int) and isinstance(d1, int): + + # Both vals are ints for this dim + if d0 == d1: + # Do not keep this dim + continue + elif d0 > d1: + # These ints inform us about the relative ordering of + # two statements. While their values may be larger than 1 in + # the lexicographic ordering describing a larger set of + # statements, in a pairwise schedule, only ints 0 and 1 are + # necessary to specify relative order. To keep the pairwise + # schedules as simple and comprehensible as possible, use only + # integers 0 and 1 to specify this relative ordering. + # (doesn't take much extra time since we are already going + # through these to remove unnecessary lex tuple dims) + new_tup0.append(1) + new_tup1.append(0) + + # No further dims needed to fully specify ordering + break + else: # d1 > d0 + new_tup0.append(0) + new_tup1.append(1) + + # No further dims needed to fully specify ordering + break + else: + # Keep this dim without modifying + new_tup0.append(d0) + new_tup1.append(d1) + + if len(new_tup0) == 0: + # Statements map to the exact same point(s) in the lex ordering, + # which is okay, but to represent this, our lex tuple cannot be empty. + return (0, ), (0, ) + else: + return tuple(new_tup0), tuple(new_tup1) + + +def generate_pairwise_schedules( + knl, + linearization_items, + insn_id_pairs, + loops_to_ignore=frozenset(), + ): + r"""For each statement pair in a subset of all statement pairs found in a + linearized kernel, determine the (relative) order in which the statement + instances are executed. For each pair, describe this relative ordering with + a pair of mappings from statement instances to points in a single + lexicographic ordering (a ``pairwise schedule''). + + :arg knl: A preprocessed :class:`loopy.kernel.LoopKernel` containing the + linearization items that will be used to create a schedule. This + kernel will be used to get the domains associated with the inames + used in the statements. + + :arg linearization_items: A list of :class:`loopy.schedule.ScheduleItem` + (to be renamed to `loopy.schedule.LinearizationItem`) containing + all linearization items for which pairwise schedules will be + created. To allow usage of this routine during linearization, a + truncated (i.e. partial) linearization may be passed through this + argument. + + :arg insn_id_pairs: A list containing pairs of instruction identifiers. + + :arg loops_to_ignore: A set of inames that will be ignored when + determining the relative ordering of statements. This will typically + contain concurrent inames tagged with the ``vec`` or ``ilp`` array + access tags. + + :returns: A dictionary mapping each two-tuple of instruction identifiers + provided in `insn_id_pairs` to a corresponding two-tuple containing two + :class:`islpy.Map`\ s representing a pairwise schedule as two + mappings from statement instances to lexicographic time, one for + each of the two statements. + """ + + from loopy.schedule import (EnterLoop, LeaveLoop, Barrier, RunInstruction) + + all_insn_ids = set().union(*insn_id_pairs) + + # First, use one pass through linearization_items to generate a lexicographic + # ordering describing the relative order of *all* statements represented by + # all_insn_ids + + # For each statement, map the insn_id to a tuple representing points + # in the lexicographic ordering containing items of :class:`int` or + # :class:`str` :mod:`loopy` inames. + stmt_instances = {} + + # Keep track of the next tuple of points in our lexicographic + # ordering, initially this as a 1-d point with value 0 + next_insn_lex_tuple = [0] + + for linearization_item in linearization_items: + if isinstance(linearization_item, EnterLoop): + iname = linearization_item.iname + if iname in loops_to_ignore: + continue + + # Increment next_insn_lex_tuple[-1] for statements in the section + # of code after this EnterLoop. + # (not technically necessary if no statement was added in the + # previous section; gratuitous incrementing is counteracted + # in the simplification step below) + next_insn_lex_tuple[-1] += 1 + + # Upon entering a loop, add one lex dimension for the loop variable, + # add second lex dim to enumerate sections of code within new loop + next_insn_lex_tuple.append(iname) + next_insn_lex_tuple.append(0) + + elif isinstance(linearization_item, LeaveLoop): + if linearization_item.iname in loops_to_ignore: + continue + + # Upon leaving a loop, + # pop lex dimension for enumerating code sections within this loop, and + # pop lex dimension for the loop variable, and + # increment lex dim val enumerating items in current section of code + next_insn_lex_tuple.pop() + next_insn_lex_tuple.pop() + + # Increment next_insn_lex_tuple[-1] for statements in the section + # of code after this LeaveLoop. + # (not technically necessary if no statement was added in the + # previous section; gratuitous incrementing is counteracted + # in the simplification step below) + next_insn_lex_tuple[-1] += 1 + + elif isinstance(linearization_item, (RunInstruction, Barrier)): + from loopy.schedule.checker.utils import ( + get_insn_id_from_linearization_item, + ) + lp_insn_id = get_insn_id_from_linearization_item(linearization_item) + + if lp_insn_id is None: + assert isinstance(linearization_item, Barrier) + + # Barriers without insn ids were inserted as a result of a + # dependency. They don't themselves have dependencies. Ignore them. + + # FIXME: It's possible that we could record metadata about them + # (e.g. what dependency produced them) and verify that they're + # adequately protecting all statement instance pairs. + + continue + + # Only process listed insns, otherwise ignore + if lp_insn_id in all_insn_ids: + # Add item to stmt_instances + stmt_instances[lp_insn_id] = tuple(next_insn_lex_tuple) + + # Increment lex dim val enumerating items in current section of code + next_insn_lex_tuple[-1] += 1 + + else: + from loopy.schedule import (CallKernel, ReturnFromKernel) + # No action needed for these types of linearization item + assert isinstance( + linearization_item, (CallKernel, ReturnFromKernel)) + pass + + # To save time, stop when we've found all statements + if len(stmt_instances.keys()) == len(all_insn_ids): + break + + from loopy.schedule.checker.utils import ( + sorted_union_of_names_in_isl_sets, + create_symbolic_map_from_tuples, + add_dims_to_isl_set, + ) + + def _get_map_for_stmt_inst(insn_id, lex_points, int_sid, out_names_sched): + + # Get inames domain for statement instance (a BasicSet) + dom = knl.get_inames_domain( + knl.id_to_insn[insn_id].within_inames) + + # Create map space (an isl space in current implementation) + # {('statement', ) -> + # (lexicographic ordering dims)} + dom_inames_ordered = sorted_union_of_names_in_isl_sets([dom]) + + in_names_sched = [STATEMENT_VAR_NAME] + dom_inames_ordered[:] + sched_space = isl.Space.create_from_names( + isl.DEFAULT_CONTEXT, + in_=in_names_sched, out=out_names_sched, params=[]) + + # Insert 'statement' dim into domain so that its space allows + # for intersection with sched map later + dom_to_intersect = [ + add_dims_to_isl_set( + dom, isl.dim_type.set, [STATEMENT_VAR_NAME], 0), ] + + # Each map will map statement instances -> lex time. + # Right now, statement instance tuples consist of single int. + # Add all inames from domains to each map domain tuple. + tuple_pair = [( + (int_sid, ) + tuple(dom_inames_ordered), + lex_points + )] + + # Create map + return create_symbolic_map_from_tuples( + tuple_pairs_with_domains=zip(tuple_pair, dom_to_intersect), + space=sched_space, + ) + + # Second, create pairwise schedules for each individual pair of insns + + pairwise_schedules = {} + for insn_ids in insn_id_pairs: + lex_tuples = [stmt_instances[insn_id] for insn_id in insn_ids] + + # Simplify tuples to the extent possible ------------------------------------ + + # At this point, one of the lex tuples may have more dimensions than another; + # the missing dims are the fastest-updating dims, and their values should + # be zero. Add them. + max_lex_dims = max([len(lex_tuple) for lex_tuple in lex_tuples]) + lex_tuples_padded = [ + _pad_tuple_with_zeros(lex_tuple, max_lex_dims) + for lex_tuple in lex_tuples] + + lex_tuples_simplified = _simplify_lex_dims(*lex_tuples_padded) + + # Now generate maps from the blueprint -------------------------------------- + + # Create names for the output dimensions + out_names_sched = [ + LEX_VAR_PREFIX+str(i) for i in range(len(lex_tuples_simplified[0]))] + + # Determine integer IDs that will represent each statement in mapping + # (dependency map creation assumes sid_before=0 and sid_after=1, unless + # before and after refer to same stmt, in which case sid_before=sid_after=0) + int_sids = [0, 0] if insn_ids[0] == insn_ids[1] else [0, 1] + + sched_maps = [ + _get_map_for_stmt_inst(insn_id, lex_tuple, int_sid, out_names_sched) + for insn_id, lex_tuple, int_sid + in zip(insn_ids, lex_tuples_simplified, int_sids) + ] + + pairwise_schedules[tuple(insn_ids)] = tuple(sched_maps) + + return pairwise_schedules diff --git a/loopy/schedule/checker/utils.py b/loopy/schedule/checker/utils.py new file mode 100644 index 000000000..8e2a82a01 --- /dev/null +++ b/loopy/schedule/checker/utils.py @@ -0,0 +1,223 @@ +__copyright__ = "Copyright (C) 2019 James Stevens" + +__license__ = """ +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +""" + +import islpy as isl + + +def prettier_map_string(map_obj): + return str(map_obj + ).replace("{ ", "{\n").replace(" }", "\n}").replace("; ", ";\n") + + +def add_dims_to_isl_set(isl_set, dim_type, names, new_idx_start): + new_set = isl_set.insert_dims( + dim_type, new_idx_start, len(names) + ).set_dim_name(dim_type, new_idx_start, names[0]) + for i, name in enumerate(names[1:]): + new_set = new_set.set_dim_name(dim_type, new_idx_start+1+i, name) + return new_set + + +def reorder_dims_by_name( + isl_set, dim_type, desired_dims_ordered): + """Return an isl_set with the dimensions of the specified dim_type + in the specified order. + + :arg isl_set: A :class:`islpy.Set` whose dimensions are + to be reordered. + + :arg dim_type: A :class:`islpy.dim_type`, i.e., an :class:`int`, + specifying the dimension to be reordered. + + :arg desired_dims_ordered: A :class:`list` of :class:`str` elements + representing the desired dimensions in order by dimension name. + + :returns: An :class:`islpy.Set` matching `isl_set` with the + dimension order matching `desired_dims_ordered`. + + """ + + assert dim_type != isl.dim_type.param + assert set(isl_set.get_var_names(dim_type)) == set(desired_dims_ordered) + + other_dim_type = isl.dim_type.param + other_dim_len = len(isl_set.get_var_names(other_dim_type)) + + new_set = isl_set.copy() + for desired_idx, name in enumerate(desired_dims_ordered): + + current_idx = new_set.find_dim_by_name(dim_type, name) + if current_idx != desired_idx: + # First move to other dim because isl is stupid + new_set = new_set.move_dims( + other_dim_type, other_dim_len, dim_type, current_idx, 1) + # Now move it where we actually want it + new_set = new_set.move_dims( + dim_type, desired_idx, other_dim_type, other_dim_len, 1) + + return new_set + + +def ensure_dim_names_match_and_align(obj_map, tgt_map): + + # first make sure names match + assert all( + set(obj_map.get_var_names(dt)) == set(tgt_map.get_var_names(dt)) + for dt in [isl.dim_type.in_, isl.dim_type.out, isl.dim_type.param]) + + return isl.align_spaces(obj_map, tgt_map) + + +def sorted_union_of_names_in_isl_sets( + isl_sets, + set_dim=isl.dim_type.set): + r"""Return a sorted list of the union of all variable names found in + the provided :class:`islpy.Set`\ s. + """ + + inames = set().union(*[isl_set.get_var_names(set_dim) for isl_set in isl_sets]) + + # Sorting is not necessary, but keeps results consistent between runs + return sorted(inames) + + +def create_symbolic_map_from_tuples( + tuple_pairs_with_domains, + space, + ): + """Return an :class:`islpy.Map` constructed using the provided space, + mapping input->output tuples provided in `tuple_pairs_with_domains`, + with each set of tuple variables constrained by the domains provided. + + :arg tuple_pairs_with_domains: A :class:`list` with each element being + a tuple of the form `((tup_in, tup_out), domain)`. + `tup_in` and `tup_out` are tuples containing elements of type + :class:`int` and :class:`str` representing values for the + input and output dimensions in `space`, and `domain` is a + :class:`islpy.Set` constraining variable bounds. + + :arg space: A :class:`islpy.Space` to be used to create the map. + + :returns: A :class:`islpy.Map` constructed using the provided space + as follows. For each `((tup_in, tup_out), domain)` in + `tuple_pairs_with_domains`, map + `(tup_in)->(tup_out) : domain`, where `tup_in` and `tup_out` are + numeric or symbolic values assigned to the input and output + dimension variables in `space`, and `domain` specifies conditions + on these values. + + """ + # TODO allow None for domains + + dim_type = isl.dim_type + + space_out_names = space.get_var_names(dim_type.out) + space_in_names = space.get_var_names(isl.dim_type.in_) + + # Get islvars from space + islvars = isl.affs_from_space( + space.move_dims( + isl.dim_type.out, 0, + isl.dim_type.in_, 0, + len(space_in_names), + ).range() + ) + + def _conjunction_of_dim_eq_conditions(dim_names, values, islvars): + condition = islvars[0].eq_set(islvars[0]) + for dim_name, val in zip(dim_names, values): + if isinstance(val, int): + condition = condition \ + & islvars[dim_name].eq_set(islvars[0]+val) + else: + condition = condition \ + & islvars[dim_name].eq_set(islvars[val]) + return condition + + # Initialize union of maps to empty + union_of_maps = isl.Map.from_domain( + islvars[0].eq_set(islvars[0]+1) # 0 == 1 (false) + ).move_dims( + dim_type.out, 0, dim_type.in_, len(space_in_names), len(space_out_names)) + + # Loop through tuple pairs + for (tup_in, tup_out), dom in tuple_pairs_with_domains: + + # Set values for 'in' dimension using tuple vals + condition = _conjunction_of_dim_eq_conditions( + space_in_names, tup_in, islvars) + + # Set values for 'out' dimension using tuple vals + condition = condition & _conjunction_of_dim_eq_conditions( + space_out_names, tup_out, islvars) + + # Convert set to map by moving dimensions around + map_from_set = isl.Map.from_domain(condition) + map_from_set = map_from_set.move_dims( + dim_type.out, 0, dim_type.in_, + len(space_in_names), len(space_out_names)) + + # Align the *out* dims of dom with the space *in_* dims + # in preparation for intersection + dom_with_set_dim_aligned = reorder_dims_by_name( + dom, isl.dim_type.set, + space_in_names, + ) + + # Intersect domain with this map + union_of_maps = union_of_maps.union( + map_from_set.intersect_domain(dom_with_set_dim_aligned)) + + return union_of_maps + + +def partition_inames_by_concurrency(knl): + from loopy.kernel.data import ConcurrentTag + conc_inames = set() + non_conc_inames = set() + + all_inames = knl.all_inames() + for iname in all_inames: + if knl.iname_tags_of_type(iname, ConcurrentTag): + conc_inames.add(iname) + else: + non_conc_inames.add(iname) + + return conc_inames, all_inames-conc_inames + + +def get_insn_id_from_linearization_item(linearization_item): + from loopy.schedule import Barrier + if isinstance(linearization_item, Barrier): + return linearization_item.originating_insn_id + else: + return linearization_item.insn_id + + +def get_EnterLoop_inames(linearization_items): + from loopy.schedule import EnterLoop + + # Note: each iname must live in len-1 list to avoid char separation + return set().union(*[ + [item.iname, ] for item in linearization_items + if isinstance(item, EnterLoop) + ]) diff --git a/test/test_linearization_checker.py b/test/test_linearization_checker.py new file mode 100644 index 000000000..dff64764c --- /dev/null +++ b/test/test_linearization_checker.py @@ -0,0 +1,306 @@ +from __future__ import division, print_function + +__copyright__ = "Copyright (C) 2019 James Stevens" + +__license__ = """ +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +""" + +import six # noqa: F401 +import sys +import numpy as np +import loopy as lp +from pyopencl.tools import ( # noqa + pytest_generate_tests_for_pyopencl + as pytest_generate_tests) +from loopy.version import LOOPY_USE_LANGUAGE_VERSION_2018_2 # noqa +import logging +from loopy import ( + preprocess_kernel, + get_one_linearized_kernel, +) +from loopy.schedule.checker.schedule import ( + LEX_VAR_PREFIX, + STATEMENT_VAR_NAME, +) + +logger = logging.getLogger(__name__) + + +def test_lexschedule_creation(): + import islpy as isl + from loopy.schedule.checker import ( + get_schedules_for_statement_pairs, + ) + from loopy.schedule.checker.utils import ( + ensure_dim_names_match_and_align, + ) + + # example kernel + # insn_c depends on insn_b only to create deterministic order + # insn_d depends on insn_c only to create deterministic order + knl = lp.make_kernel( + [ + "{[i]: 0<=itemp = b[i,k] {id=insn_a} + end + for j + a[i,j] = temp + 1 {id=insn_b,dep=insn_a} + c[i,j] = d[i,j] {id=insn_c,dep=insn_b} + end + end + for t + e[t] = f[t] {id=insn_d, dep=insn_c} + end + """, + assumptions="pi,pj,pk,pt >= 1", + ) + knl = lp.add_and_infer_dtypes( + knl, + {"b": np.float32, "d": np.float32, "f": np.float32}) + knl = lp.prioritize_loops(knl, "i,k") + knl = lp.prioritize_loops(knl, "i,j") + + # get a linearization + knl = preprocess_kernel(knl) + knl = get_one_linearized_kernel(knl["loopy_kernel"], knl.callables_table) + linearization_items = knl.linearization + + def _lex_space_string(dim_vals): + # Return a string describing lex space dimension assignments + # (used to create maps below) + return ", ".join( + ["%s%d=%s" % (LEX_VAR_PREFIX, idx, str(val)) + for idx, val in enumerate(dim_vals)]) + + insn_id_pairs = [ + ("insn_a", "insn_b"), + ("insn_a", "insn_c"), + ("insn_a", "insn_d"), + ("insn_b", "insn_c"), + ("insn_b", "insn_d"), + ("insn_c", "insn_d"), + ] + sched_maps = get_schedules_for_statement_pairs( + knl, + linearization_items, + insn_id_pairs, + ) + + # Relationship between insn_a and insn_b --------------------------------------- + + # Get two maps + sched_map_before, sched_map_after = sched_maps[("insn_a", "insn_b")] + + # Create expected maps, align, compare + + sched_map_before_expected = isl.Map( + "[pi, pk] -> { [%s=0, i, k] -> [%s] : 0 <= i < pi and 0 <= k < pk }" + % ( + STATEMENT_VAR_NAME, + _lex_space_string(["i", "0"]), + ) + ) + sched_map_before_expected = ensure_dim_names_match_and_align( + sched_map_before_expected, sched_map_before) + + sched_map_after_expected = isl.Map( + "[pi, pj] -> { [%s=1, i, j] -> [%s] : 0 <= i < pi and 0 <= j < pj }" + % ( + STATEMENT_VAR_NAME, + _lex_space_string(["i", "1"]), + ) + ) + sched_map_after_expected = ensure_dim_names_match_and_align( + sched_map_after_expected, sched_map_after) + + assert sched_map_before == sched_map_before_expected + assert sched_map_after == sched_map_after_expected + + # ------------------------------------------------------------------------------ + # Relationship between insn_a and insn_c --------------------------------------- + + # Get two maps + sched_map_before, sched_map_after = sched_maps[("insn_a", "insn_c")] + + # Create expected maps, align, compare + + sched_map_before_expected = isl.Map( + "[pi, pk] -> { [%s=0, i, k] -> [%s] : 0 <= i < pi and 0 <= k < pk }" + % ( + STATEMENT_VAR_NAME, + _lex_space_string(["i", "0"]), + ) + ) + sched_map_before_expected = ensure_dim_names_match_and_align( + sched_map_before_expected, sched_map_before) + + sched_map_after_expected = isl.Map( + "[pi, pj] -> { [%s=1, i, j] -> [%s] : 0 <= i < pi and 0 <= j < pj }" + % ( + STATEMENT_VAR_NAME, + _lex_space_string(["i", "1"]), + ) + ) + sched_map_after_expected = ensure_dim_names_match_and_align( + sched_map_after_expected, sched_map_after) + + assert sched_map_before == sched_map_before_expected + assert sched_map_after == sched_map_after_expected + + # ------------------------------------------------------------------------------ + # Relationship between insn_a and insn_d --------------------------------------- + + # Get two maps + sched_map_before, sched_map_after = sched_maps[("insn_a", "insn_d")] + + # Create expected maps, align, compare + + sched_map_before_expected = isl.Map( + "[pi, pk] -> { [%s=0, i, k] -> [%s] : 0 <= i < pi and 0 <= k < pk }" + % ( + STATEMENT_VAR_NAME, + _lex_space_string([0, ]), + ) + ) + sched_map_before_expected = ensure_dim_names_match_and_align( + sched_map_before_expected, sched_map_before) + + sched_map_after_expected = isl.Map( + "[pt] -> { [%s=1, t] -> [%s] : 0 <= t < pt }" + % ( + STATEMENT_VAR_NAME, + _lex_space_string([1, ]), + ) + ) + sched_map_after_expected = ensure_dim_names_match_and_align( + sched_map_after_expected, sched_map_after) + + assert sched_map_before == sched_map_before_expected + assert sched_map_after == sched_map_after_expected + + # ------------------------------------------------------------------------------ + # Relationship between insn_b and insn_c --------------------------------------- + + # Get two maps + sched_map_before, sched_map_after = sched_maps[("insn_b", "insn_c")] + + # Create expected maps, align, compare + + sched_map_before_expected = isl.Map( + "[pi, pj] -> { [%s=0, i, j] -> [%s] : 0 <= i < pi and 0 <= j < pj }" + % ( + STATEMENT_VAR_NAME, + _lex_space_string(["i", "j", 0]), + ) + ) + sched_map_before_expected = ensure_dim_names_match_and_align( + sched_map_before_expected, sched_map_before) + + sched_map_after_expected = isl.Map( + "[pi, pj] -> { [%s=1, i, j] -> [%s] : 0 <= i < pi and 0 <= j < pj }" + % ( + STATEMENT_VAR_NAME, + _lex_space_string(["i", "j", 1]), + ) + ) + sched_map_after_expected = ensure_dim_names_match_and_align( + sched_map_after_expected, sched_map_after) + + assert sched_map_before == sched_map_before_expected + assert sched_map_after == sched_map_after_expected + + # ------------------------------------------------------------------------------ + # Relationship between insn_b and insn_d --------------------------------------- + + # Get two maps + sched_map_before, sched_map_after = sched_maps[("insn_b", "insn_d")] + + # Create expected maps, align, compare + + sched_map_before_expected = isl.Map( + "[pi, pj] -> { [%s=0, i, j] -> [%s] : 0 <= i < pi and 0 <= j < pj }" + % ( + STATEMENT_VAR_NAME, + _lex_space_string([0, ]), + ) + ) + sched_map_before_expected = ensure_dim_names_match_and_align( + sched_map_before_expected, sched_map_before) + + sched_map_after_expected = isl.Map( + "[pt] -> { [%s=1, t] -> [%s] : 0 <= t < pt }" + % ( + STATEMENT_VAR_NAME, + _lex_space_string([1, ]), + ) + ) + sched_map_after_expected = ensure_dim_names_match_and_align( + sched_map_after_expected, sched_map_after) + + assert sched_map_before == sched_map_before_expected + assert sched_map_after == sched_map_after_expected + + # ------------------------------------------------------------------------------ + # Relationship between insn_c and insn_d --------------------------------------- + + # Get two maps + sched_map_before, sched_map_after = sched_maps[("insn_c", "insn_d")] + + # Create expected maps, align, compare + + sched_map_before_expected = isl.Map( + "[pi, pj] -> { [%s=0, i, j] -> [%s] : 0 <= i < pi and 0 <= j < pj }" + % ( + STATEMENT_VAR_NAME, + _lex_space_string([0, ]), + ) + ) + sched_map_before_expected = ensure_dim_names_match_and_align( + sched_map_before_expected, sched_map_before) + + sched_map_after_expected = isl.Map( + "[pt] -> { [%s=1, t] -> [%s] : 0 <= t < pt }" + % ( + STATEMENT_VAR_NAME, + _lex_space_string([1, ]), + ) + ) + sched_map_after_expected = ensure_dim_names_match_and_align( + sched_map_after_expected, sched_map_after) + + assert sched_map_before == sched_map_before_expected + assert sched_map_after == sched_map_after_expected + + +if __name__ == "__main__": + if len(sys.argv) > 1: + exec(sys.argv[1]) + else: + from pytest import main + main([__file__]) + +# vim: foldmethod=marker