|
| 1 | +__copyright__ = "Copyright (C) 2019 James Stevens" |
| 2 | + |
| 3 | +__license__ = """ |
| 4 | +Permission is hereby granted, free of charge, to any person obtaining a copy |
| 5 | +of this software and associated documentation files (the "Software"), to deal |
| 6 | +in the Software without restriction, including without limitation the rights |
| 7 | +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell |
| 8 | +copies of the Software, and to permit persons to whom the Software is |
| 9 | +furnished to do so, subject to the following conditions: |
| 10 | +
|
| 11 | +The above copyright notice and this permission notice shall be included in |
| 12 | +all copies or substantial portions of the Software. |
| 13 | +
|
| 14 | +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR |
| 15 | +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, |
| 16 | +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE |
| 17 | +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER |
| 18 | +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, |
| 19 | +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN |
| 20 | +THE SOFTWARE. |
| 21 | +""" |
| 22 | + |
| 23 | + |
| 24 | +# {{{ create a pairwise schedules for statement pairs |
| 25 | + |
| 26 | +def get_schedules_for_statement_pairs( |
| 27 | + knl, |
| 28 | + linearization_items, |
| 29 | + insn_id_pairs, |
| 30 | + ): |
| 31 | + r"""For each statement pair in a subset of all statement pairs found in a |
| 32 | + linearized kernel, determine the (relative) order in which the statement |
| 33 | + instances are executed. For each pair, describe this relative ordering with |
| 34 | + a pair of mappings from statement instances to points in a single |
| 35 | + lexicographic ordering (a ``pairwise schedule''). When determining the |
| 36 | + relative ordering, ignore concurrent inames. |
| 37 | +
|
| 38 | + :arg knl: A preprocessed :class:`loopy.kernel.LoopKernel` containing the |
| 39 | + linearization items that will be used to create a schedule. |
| 40 | +
|
| 41 | + :arg linearization_items: A list of :class:`loopy.schedule.ScheduleItem` |
| 42 | + (to be renamed to `loopy.schedule.LinearizationItem`) containing |
| 43 | + all linearization items for which pairwise schedules will be |
| 44 | + created. To allow usage of this routine during linearization, a |
| 45 | + truncated (i.e. partial) linearization may be passed through this |
| 46 | + argument. |
| 47 | +
|
| 48 | + :arg insn_id_pairs: A list containing pairs of instruction |
| 49 | + identifiers. |
| 50 | +
|
| 51 | + :returns: A dictionary mapping each two-tuple of instruction identifiers |
| 52 | + provided in `insn_id_pairs` to a corresponding two-tuple containing two |
| 53 | + :class:`islpy.Map`\ s representing a pairwise schedule as two |
| 54 | + mappings from statement instances to lexicographic time, one for |
| 55 | + each of the two statements. |
| 56 | +
|
| 57 | + .. doctest: |
| 58 | +
|
| 59 | + >>> import loopy as lp |
| 60 | + >>> import numpy as np |
| 61 | + >>> # Make kernel ----------------------------------------------------------- |
| 62 | + >>> knl = lp.make_kernel( |
| 63 | + ... "{[i,j,k]: 0<=i<pi and 0<=j<pj and 0<=k<pk}", |
| 64 | + ... [ |
| 65 | + ... "a[i,j] = j {id=insn_a}", |
| 66 | + ... "b[i,k] = k+a[i,0] {id=insn_b,dep=insn_a}", |
| 67 | + ... ]) |
| 68 | + >>> knl = lp.add_and_infer_dtypes(knl, {"a": np.float32, "b": np.float32}) |
| 69 | + >>> knl = lp.prioritize_loops(knl, "i,j") |
| 70 | + >>> knl = lp.prioritize_loops(knl, "i,k") |
| 71 | + >>> # Get a linearization |
| 72 | + >>> knl = lp.get_one_linearized_kernel(lp.preprocess_kernel(knl)) |
| 73 | + >>> # Get a pairwise schedule ----------------------------------------------- |
| 74 | + >>> from loopy.schedule.checker import get_schedules_for_statement_pairs |
| 75 | + >>> # Get two maps ---------------------------------------------------------- |
| 76 | + >>> schedules = get_schedules_for_statement_pairs( |
| 77 | + ... knl, |
| 78 | + ... knl.linearization, |
| 79 | + ... [("insn_a", "insn_b")], |
| 80 | + ... ) |
| 81 | + >>> # Print maps |
| 82 | + >>> print("\n".join( |
| 83 | + ... str(m).replace("{ ", "{\n").replace(" :", "\n:") |
| 84 | + ... for m in schedules[("insn_a", "insn_b")] |
| 85 | + ... )) |
| 86 | + [pi, pj, pk] -> { |
| 87 | + [_lp_linchk_stmt = 0, i, j, k] -> [_lp_linchk_l0 = i, _lp_linchk_l1 = 0] |
| 88 | + : 0 <= i < pi and 0 <= j < pj and 0 <= k < pk } |
| 89 | + [pi, pj, pk] -> { |
| 90 | + [_lp_linchk_stmt = 1, i, j, k] -> [_lp_linchk_l0 = i, _lp_linchk_l1 = 1] |
| 91 | + : 0 <= i < pi and 0 <= j < pj and 0 <= k < pk } |
| 92 | +
|
| 93 | + """ |
| 94 | + |
| 95 | + # {{{ make sure kernel has been preprocessed |
| 96 | + |
| 97 | + from loopy.kernel import KernelState |
| 98 | + assert knl.state in [ |
| 99 | + KernelState.PREPROCESSED, |
| 100 | + KernelState.LINEARIZED] |
| 101 | + |
| 102 | + # }}} |
| 103 | + |
| 104 | + # {{{ Find any EnterLoop inames that are tagged as concurrent |
| 105 | + # so that generate_pairwise_schedule() knows to ignore them |
| 106 | + # (In the future, this shouldn't be necessary because there |
| 107 | + # won't be any inames with ConcurrentTags in EnterLoop linearization items. |
| 108 | + # Test which exercises this: test_linearization_checker_with_stroud_bernstein()) |
| 109 | + from loopy.schedule.checker.utils import ( |
| 110 | + partition_inames_by_concurrency, |
| 111 | + get_EnterLoop_inames, |
| 112 | + ) |
| 113 | + conc_inames, _ = partition_inames_by_concurrency(knl) |
| 114 | + enterloop_inames = get_EnterLoop_inames(linearization_items) |
| 115 | + conc_loop_inames = conc_inames & enterloop_inames |
| 116 | + |
| 117 | + # The only concurrent EnterLoop inames should be Vec and ILP |
| 118 | + from loopy.kernel.data import (VectorizeTag, IlpBaseTag) |
| 119 | + for conc_iname in conc_loop_inames: |
| 120 | + # Assert that there exists an ilp or vectorize tag (out of the |
| 121 | + # potentially multiple other tags on this concurrent iname). |
| 122 | + assert any( |
| 123 | + isinstance(tag, (VectorizeTag, IlpBaseTag)) |
| 124 | + for tag in knl.iname_to_tags[conc_iname]) |
| 125 | + |
| 126 | + # }}} |
| 127 | + |
| 128 | + # {{{ Create two mappings from {statement instance: lex point} |
| 129 | + |
| 130 | + # include only instructions involved in this dependency |
| 131 | + from loopy.schedule.checker.schedule import generate_pairwise_schedules |
| 132 | + return generate_pairwise_schedules( |
| 133 | + knl, |
| 134 | + linearization_items, |
| 135 | + insn_id_pairs, |
| 136 | + loops_to_ignore=conc_loop_inames, |
| 137 | + ) |
| 138 | + |
| 139 | + # }}} |
| 140 | + |
| 141 | +# }}} |
0 commit comments