From ea6bb44bbe36f126ec47a47d7fcf1528bf0ba7ba Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Sun, 26 Jun 2022 00:00:27 -0500 Subject: [PATCH 1/4] Implements pw_qpolynomial_to_expr --- loopy/symbolic.py | 28 +++++++++++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/loopy/symbolic.py b/loopy/symbolic.py index fd6013416..584b0dc10 100644 --- a/loopy/symbolic.py +++ b/loopy/symbolic.py @@ -24,7 +24,7 @@ """ -from typing import ClassVar, Tuple +from typing import ClassVar, Tuple, Union from functools import reduce, cached_property from sys import intern import re @@ -69,6 +69,7 @@ from loopy.diagnostic import LoopyError from loopy.diagnostic import (ExpressionToAffineConversionError, UnableToDetermineAccessRangeError) +from loopy.typing import ExpressionT __doc__ = """ @@ -2792,4 +2793,29 @@ def is_tuple_of_expressions_equal(a, b): # }}} + +def _is_isl_set_universe(isl_set: Union[isl.BasicSet, isl.Set]): + if isinstance(isl_set, isl.BasicSet): + return isl_set.is_universe() + else: + assert isinstance(isl_set, isl.Set) + return isl_set.complement().is_empty() + + +def pw_qpolynomial_to_expr(pw_qpoly: isl.PwQPolynomial + ) -> ExpressionT: + from pymbolic.primitives import If + + result = 0 + + for bset, qpoly in reversed(pw_qpoly.get_pieces()): + if _is_isl_set_universe(bset): + result = qpolynomial_to_expr(qpoly) + else: + result = If(set_to_cond_expr(bset), + qpolynomial_to_expr(qpoly), + result) + + return result + # vim: foldmethod=marker From 9c9c41d567c0abfd19698cb27c3aee4bfab0fd81 Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Sat, 25 Jun 2022 20:39:37 -0500 Subject: [PATCH 2/4] Implements Loechner Reindexing --- doc/misc.rst | 7 + doc/ref_transform.rst | 2 + loopy/__init__.py | 4 + loopy/transform/reindex.py | 325 +++++++++++++++++++++++++++++++++++++ 4 files changed, 338 insertions(+) create mode 100644 loopy/transform/reindex.py diff --git a/doc/misc.rst b/doc/misc.rst index 3fea6fdd4..6821c3cef 100644 --- a/doc/misc.rst +++ b/doc/misc.rst @@ -456,6 +456,13 @@ Here's a Bibtex entry for your convenience:: doi = "{10.1145/2627373.2627387}", } +References +========== + +.. [Seghir_2006] Seghir and Loechner, Proceedings of the 2006 International + Conference on Compilers, Architecture and Synthesis for Embedded systems, + `(DOI) `__ + Getting help ============ diff --git a/doc/ref_transform.rst b/doc/ref_transform.rst index 9ef012d66..54117ee10 100644 --- a/doc/ref_transform.rst +++ b/doc/ref_transform.rst @@ -54,6 +54,8 @@ Influencing data access .. autofunction:: allocate_temporaries_for_base_storage +.. automodule:: loopy.transform.reindex + Padding Data ------------ diff --git a/loopy/__init__.py b/loopy/__init__.py index 4796c1f59..14372304a 100644 --- a/loopy/__init__.py +++ b/loopy/__init__.py @@ -108,6 +108,8 @@ from loopy.transform.fusion import fuse_kernels from loopy.transform.concatenate import concatenate_arrays +from loopy.transform.reindex import reindex_temporary_using_seghir_loechner_scheme + from loopy.transform.arithmetic import ( fold_constants, collect_common_factors_on_increment) @@ -239,6 +241,8 @@ "fold_constants", "collect_common_factors_on_increment", + "reindex_temporary_using_seghir_loechner_scheme", + "split_array_axis", "split_array_dim", "split_arg_axis", "find_padding_multiple", "add_padding", diff --git a/loopy/transform/reindex.py b/loopy/transform/reindex.py new file mode 100644 index 000000000..50f80eba1 --- /dev/null +++ b/loopy/transform/reindex.py @@ -0,0 +1,325 @@ +""" +.. currentmodule:: loopy + +.. autofunction:: reindex_temporary_using_seghir_loechner_scheme +""" + +__copyright__ = "Copyright (C) 2022 Kaushik Kulkarni" + +__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 +from typing import Union, Iterable, Tuple +from loopy.typing import ExpressionT +from loopy.kernel import LoopKernel +from loopy.diagnostic import LoopyError +from loopy.symbolic import CombineMapper +from loopy.kernel.instruction import (MultiAssignmentBase, + CInstruction, BarrierInstruction) +from loopy.symbolic import RuleAwareIdentityMapper + + +ISLMapT = Union[isl.BasicMap, isl.Map] +ISLSetT = Union[isl.BasicSet, isl.Set] + + +def _add_prime_to_dim_names(isl_map: ISLMapT, + dts: Iterable[isl.dim_type]) -> ISLMapT: + """ + Returns a copy of *isl_map* with dims of types *dts* having their names + suffixed with an apostrophe (``'``). + + .. testsetup:: + + >>> import islpy as isl + >>> from loopy.transform.reindex import _add_prime_to_dim_names + + .. doctest:: + + >>> amap = isl.Map("{[i]->[j=2i]}") + >>> _add_prime_to_dim_names(amap, [isl.dim_type.in_, isl.dim_type.out]) + Map("{ [i'] -> [j' = 2i'] }") + """ + for dt in dts: + for idim in range(isl_map.dim(dt)): + old_name = isl_map.get_dim_name(dt, idim) + new_name = f"{old_name}'" + isl_map = isl_map.set_dim_name(dt, idim, new_name) + + return isl_map + + +def _get_seghir_loechner_reindexing_from_range(access_range: ISLSetT + ) -> Tuple[isl.PwQPolynomial, + isl.PwQPolynomial]: + """ + Returns ``(reindex_map, new_shape)``, where, + + * ``reindex_map`` is a quasi-polynomial of the form ``[i1, .., in] -> {f(i1, + .., in)}`` representing that an array indexed via the subscripts + ``[i1, ..,in]`` should be re-indexed into a 1-dimensional array as + ``f(i1, .., in)``. + * ``new_shape`` is a quasi-polynomial corresponding to the shape of the + re-indexed 1-dimensional array. + """ + + # {{{ create amap: an ISL map which is an identity map from access_map's range + + amap = isl.BasicMap.identity( + access_range + .space + .add_dims(isl.dim_type.in_, access_range.dim(isl.dim_type.out))) + + # set amap's dim names + for idim in range(amap.dim(isl.dim_type.in_)): + amap = amap.set_dim_name(isl.dim_type.in_, idim, + f"_lpy_in_{idim}") + amap = amap.set_dim_name(isl.dim_type.out, idim, + f"_lpy_out_{idim}") + + amap = amap.intersect_domain(access_range) + + # }}} + + n_in = amap.dim(isl.dim_type.out) + n_out = amap.dim(isl.dim_type.out) + + amap_lexmin = amap.lexmin() + primed_amap_lexmin = _add_prime_to_dim_names(amap_lexmin, [isl.dim_type.in_, + isl.dim_type.out]) + + lex_lt_map = isl.Map.lex_lt_map(primed_amap_lexmin, amap_lexmin) + + # make the lexmin map parametric in terms of it's previous access expressions. + lex_lt_set = (lex_lt_map + .move_dims(isl.dim_type.param, 0, isl.dim_type.out, 0, n_in) + .domain()) + + # {{{ initialize amap_to_count + + amap_to_count = _add_prime_to_dim_names(amap, [isl.dim_type.in_]) + amap_to_count = amap_to_count.insert_dims(isl.dim_type.param, 0, n_in) + + for idim in range(n_in): + amap_to_count = amap_to_count.set_dim_name( + isl.dim_type.param, idim, + amap.get_dim_name(isl.dim_type.in_, idim)) + + amap_to_count = amap_to_count.intersect_domain(lex_lt_set) + + # }}} + + result = amap_to_count.range().card() + + # {{{ simplify 'result' by gisting with 'access_range' + + aligned_access_range = access_range.move_dims(isl.dim_type.param, 0, + isl.dim_type.set, 0, n_out) + + for idim in range(result.dim(isl.dim_type.param)): + aligned_access_range = ( + aligned_access_range + .set_dim_name(isl.dim_type.param, idim, + result.space.get_dim_name(isl.dim_type.param, + idim))) + + result = result.gist_params(aligned_access_range.params()) + + # }}} + + return result, access_range.card() + + +class _IndexCollector(CombineMapper): + """ + A mapper that collects all instances of + :class:`pymbolic.primitives.Subscript` accessing :attr:`var_name`. + """ + def __init__(self, var_name): + super().__init__() + self.var_name = var_name + + def combine(self, values): + from functools import reduce + return reduce(frozenset.union, values, frozenset()) + + def map_subscript(self, expr): + if expr.aggregate.name == self.var_name: + return frozenset([expr]) | super().map_subscript(expr) + else: + return super().map_subscript(expr) + + def map_constant(self, expr): + return frozenset() + + map_variable = map_constant + map_function_symbol = map_constant + map_tagged_variable = map_constant + map_type_cast = map_constant + map_nan = map_constant + + +class ReindexingApplier(RuleAwareIdentityMapper): + def __init__(self, rule_mapping_context, + var_to_reindex, + reindexed_var_name, + new_index_expr, + index_names): + + super().__init__(rule_mapping_context) + + self.var_to_reindex = var_to_reindex + self.reindexed_var_name = reindexed_var_name + self.new_index_expr = new_index_expr + self.index_names = index_names + + def map_subscript(self, expr, expn_state): + if expr.aggregate.name != self.var_to_reindex: + return super().map_subscript(expr, expn_state) + + from loopy.symbolic import SubstitutionMapper + from pymbolic.mapper.substitutor import make_subst_func + from pymbolic.primitives import Subscript, Variable + + rec_indices = tuple(self.rec(idx, expn_state) for idx in expr.index_tuple) + + assert len(self.index_names) == len(rec_indices) + subst_func = make_subst_func(dict(zip(self.index_names, rec_indices))) + + return SubstitutionMapper(subst_func)( + Subscript(Variable(self.reindexed_var_name), + self.new_index_expr) + ) + + +def reindex_temporary_using_seghir_loechner_scheme(kernel: LoopKernel, + var_name: str, + ) -> LoopKernel: + """ + Returns a kernel with expressions of the form ``var_name[i1, .., in]`` + replaced with ``var_name_reindexed[f(i1, .., in)]`` where ``f`` is a + quasi-polynomial as outlined in [Seghir_2006]_. + """ + from loopy.transform.subst import expand_subst + from loopy.symbolic import (BatchedAccessMapMapper, pw_qpolynomial_to_expr, + SubstitutionRuleMappingContext) + + if var_name not in kernel.temporary_variables: + raise LoopyError(f"'{var_name}' not in temporary variable in kernel" + f" '{kernel.name}'.") + + # {{{ compute the access_range of *var_name* in *kernel* + + subst_kernel = expand_subst(kernel) + access_map_recorder = BatchedAccessMapMapper( + subst_kernel, + frozenset([var_name])) + + access_exprs: Tuple[ExpressionT, ...] + + for insn in subst_kernel.instructions: + if var_name in insn.dependency_names(): + if isinstance(insn, MultiAssignmentBase): + access_exprs = (insn.assignees, + insn.expression, + tuple(insn.predicates)) + elif isinstance(insn, (CInstruction, BarrierInstruction)): + access_exprs = tuple(insn.predicates) + else: + raise NotImplementedError(type(insn)) + + access_map_recorder(access_exprs, insn.within_inames) + + vng = kernel.get_var_name_generator() + new_var_name = vng(var_name+"_reindexed") + + access_range = access_map_recorder.get_access_range(var_name) + + del subst_kernel + del access_map_recorder + + # }}} + + subst, new_shape = _get_seghir_loechner_reindexing_from_range( + access_range) + + # {{{ simplify new_shape with the assumptions from kernel + + new_shape = new_shape.gist_params(kernel.assumptions) + + # }}} + + # {{{ update kernel.temporary_variables + + new_shape = new_shape.drop_unused_params() + + new_temps = dict(kernel.temporary_variables).copy() + new_temps[new_var_name] = new_temps.pop(var_name).copy( + name=new_var_name, + shape=pw_qpolynomial_to_expr(new_shape), + strides=None, + dim_tags=None, + dim_names=None, + ) + + kernel = kernel.copy(temporary_variables=new_temps) + + # }}} + + # {{{ perform the substitution i.e. reindex the accesses + + subst_expr = pw_qpolynomial_to_expr(subst) + subst_dim_names = tuple( + subst.space.get_dim_name(isl.dim_type.param, idim) + for idim in range(access_range.dim(isl.dim_type.out))) + assert not (set(subst_dim_names) & kernel.all_variable_names()) + + rule_mapping_context = SubstitutionRuleMappingContext(kernel.substitutions, + vng) + reindexing_mapper = ReindexingApplier(rule_mapping_context, + var_name, new_var_name, + subst_expr, subst_dim_names) + + def _does_access_var_name(kernel, insn, *args): + return var_name in insn.dependency_names() + + kernel = reindexing_mapper.map_kernel(kernel, + within=_does_access_var_name, + map_args=False, + map_tvs=False) + kernel = rule_mapping_context.finish_kernel(kernel) + + # }}} + + # Note: Distributing a piece of code that depends on loopy and distributes + # code that conditionally/unconditionally calls this routine does *NOT* + # become a derivative of GPLv2. Since, as per point (0) of GPLV2 a + # derivative is defined as: "a work containing the Program or a portion of + # it, either verbatim or with modifications and/or translated into another + # language." + # + # Loopy does *NOT* contain any portion of the barvinok library in it's + # source code. + + return kernel + +# vim: fdm=marker From d94668d44050e70e89fbb3e37f610dbf14e84cc4 Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Sun, 26 Jun 2022 22:23:28 -0500 Subject: [PATCH 3/4] test reindex_using_seghir_loechner_scheme --- test/test_transform.py | 89 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 89 insertions(+) diff --git a/test/test_transform.py b/test/test_transform.py index 1b56344e7..92c270597 100644 --- a/test/test_transform.py +++ b/test/test_transform.py @@ -1668,6 +1668,95 @@ def test_remove_predicates_from_insn(): assert t_unit == ref_t_unit +def test_reindexing_strided_access(ctx_factory): + import islpy as isl + + if not hasattr(isl.Set, "card"): + pytest.skip("No barvinok support") + + ctx = ctx_factory() + + tunit = lp.make_kernel( + "{[i, j]: 0<=j,i<10}", + """ + <> tmp[2*i, 2*j] = a[i, j] + out[i, j] = tmp[2*i, 2*j]**2 + """) + + tunit = lp.add_dtypes(tunit, {"a": "float64"}) + ref_tunit = tunit + + knl = lp.reindex_temporary_using_seghir_loechner_scheme(tunit.default_entrypoint, + "tmp") + tunit = tunit.with_kernel(knl) + + tv, = tunit.default_entrypoint.temporary_variables.values() + assert tv.shape == (100,) + + lp.auto_test_vs_ref(ref_tunit, ctx, tunit) + + +def test_reindexing_figurate(ctx_factory): + import islpy as isl + + if not hasattr(isl.Set, "card"): + pytest.skip("No barvinok support") + + ctx = ctx_factory() + + tunit = lp.make_kernel( + "{[i, j]: 0<=j<=i<10}", + """ + <> tmp[2*i, 2*j] = a[i, j] + out[i, j] = tmp[2*i, 2*j]**2 + """) + + tunit = lp.add_dtypes(tunit, {"a": "float64"}) + ref_tunit = tunit + + knl = lp.reindex_temporary_using_seghir_loechner_scheme(tunit.default_entrypoint, + "tmp") + tunit = tunit.with_kernel(knl) + + tv, = tunit.default_entrypoint.temporary_variables.values() + assert tv.shape == (55,) + + lp.auto_test_vs_ref(ref_tunit, ctx, tunit) + + +def test_reindexing_figurate_parametric_shape(ctx_factory): + import islpy as isl + from loopy.symbolic import parse + + if not hasattr(isl.Set, "card"): + pytest.skip("No barvinok support") + + ctx = ctx_factory() + + tunit = lp.make_kernel( + "{[i, j]: 0<=j<=i tmp[i, j] = a[i, j] + out[i, j] = tmp[i, j]**2 + """, + assumptions="n > 0", + ) + + tunit = lp.add_dtypes(tunit, {"a": "float64"}) + tunit = lp.set_temporary_address_space(tunit, "tmp", + lp.AddressSpace.GLOBAL) + ref_tunit = tunit + + knl = lp.reindex_temporary_using_seghir_loechner_scheme(tunit.default_entrypoint, + "tmp") + tunit = tunit.with_kernel(knl) + + tv, = tunit.default_entrypoint.temporary_variables.values() + assert tv.shape == (parse("(n + n**2) // 2"),) + + lp.auto_test_vs_ref(ref_tunit, ctx, tunit, parameters={"n": 20}) + + if __name__ == "__main__": if len(sys.argv) > 1: exec(sys.argv[1]) From 1a0d21b9544d4c114c9704994efcd5f890d3d377 Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Wed, 6 Jul 2022 11:03:28 -0500 Subject: [PATCH 4/4] adds a CI job with barvinok --- .github/workflows/ci.yml | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2017f93cf..f6b7ef104 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -104,6 +104,20 @@ jobs: ( test_py_project ) ( test_py_project ) + pytest_with_barvinok: + name: Conda Pytest with Barvinok + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v2 + - name: "Main Script" + run: | + CONDA_ENVIRONMENT=.test-conda-env-py3.yml + echo "- barvinok" >> "$CONDA_ENVIRONMENT" + curl -L -O https://tiker.net/ci-support-v0 + . ./ci-support-v0 + build_py_project_in_conda_env + test_py_project + examples: name: Conda Examples runs-on: ubuntu-latest