From 08f102df2cff8fd89bf3247bc7fbe96a470aed20 Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Sun, 6 Mar 2022 16:40:34 -0600 Subject: [PATCH] Move simplify_via_aff to be near simplify_using_aff, for discoverability --- loopy/check.py | 2 +- loopy/isl_helpers.py | 22 ++++++++-------------- loopy/symbolic.py | 16 +++++++++++++++- loopy/transform/buffer.py | 2 +- loopy/transform/pack_and_unpack_args.py | 3 ++- loopy/transform/precompute.py | 2 +- test/test_isl.py | 3 +-- 7 files changed, 29 insertions(+), 21 deletions(-) diff --git a/loopy/check.py b/loopy/check.py index 9634ba6a0..736ffaf00 100644 --- a/loopy/check.py +++ b/loopy/check.py @@ -1453,7 +1453,7 @@ def _are_sub_array_refs_equivalent(sar1, sar2, caller): from loopy.symbolic import SubstitutionMapper from pymbolic.mapper.substitutor import make_subst_func - from loopy.isl_helpers import simplify_via_aff + from loopy.symbolic import simplify_via_aff subst_func = make_subst_func({iname1.name: iname2 for iname1, iname2 in zip(sar1.swept_inames, sar2.swept_inames) diff --git a/loopy/isl_helpers.py b/loopy/isl_helpers.py index 45f74d70a..5eef5f393 100644 --- a/loopy/isl_helpers.py +++ b/loopy/isl_helpers.py @@ -1,6 +1,5 @@ """isl helpers""" - __copyright__ = "Copyright (C) 2012 Andreas Kloeckner" __license__ = """ @@ -28,6 +27,7 @@ from loopy.diagnostic import StaticValueFindingError, LoopyError import islpy as isl from islpy import dim_type +from warnings import warn def pw_aff_to_aff(pw_aff): @@ -448,19 +448,6 @@ def boxify(cache_manager, domain, box_inames, context): # }}} -def simplify_via_aff(expr): - from loopy.symbolic import aff_to_expr, guarded_aff_from_expr, get_dependencies - from loopy.diagnostic import ExpressionToAffineConversionError - - deps = sorted(get_dependencies(expr)) - try: - return aff_to_expr(guarded_aff_from_expr( - isl.Space.create_from_names(isl.DEFAULT_CONTEXT, list(deps)), - expr)) - except ExpressionToAffineConversionError: - return expr - - def project_out(set, inames): for iname in inames: var_dict = set.get_var_dict() @@ -849,4 +836,11 @@ def find_and_rename_dim(isl_obj, dt, old_name, new_name): # }}} +def simplify_via_aff(expr): + warn("simplify_via_aff has moved to loopy.symbolic. " + "Importing it from loopy.isl_helpers will stop working in July 2022.", + DeprecationWarning, stacklevel=2) + from loopy.symbolic import simplify_via_aff + return simplify_via_aff(expr) + # vim: foldmethod=marker diff --git a/loopy/symbolic.py b/loopy/symbolic.py index af5965593..39af9ee7d 100644 --- a/loopy/symbolic.py +++ b/loopy/symbolic.py @@ -1962,6 +1962,19 @@ def qpolynomial_from_expr(space, expr): # {{{ simplify using aff +def simplify_via_aff(expr): + from loopy.symbolic import aff_to_expr, guarded_aff_from_expr, get_dependencies + from loopy.diagnostic import ExpressionToAffineConversionError + + deps = sorted(get_dependencies(expr)) + try: + return aff_to_expr(guarded_aff_from_expr( + isl.Space.create_from_names(isl.DEFAULT_CONTEXT, list(deps)), + expr)) + except ExpressionToAffineConversionError: + return expr + + @memoize_on_first_arg def simplify_using_aff(kernel, expr): """ @@ -1970,6 +1983,7 @@ def simplify_using_aff(kernel, expr): :arg expr: An instance of :class:`pymbolic.primitives.Expression`. """ deps = get_dependencies(expr) + inames = deps & kernel.all_inames() # FIXME: Ideally, we should find out what inames are usable and allow @@ -1981,7 +1995,7 @@ def simplify_using_aff(kernel, expr): .project_out_except(inames, [dim_type.set])) non_inames = deps - set(domain.get_var_dict().keys()) - non_inames = set([name for name in set(non_inames) if name.isidentifier()]) + non_inames = {name for name in set(non_inames) if name.isidentifier()} if non_inames: cur_dim = domain.dim(isl.dim_type.set) domain = domain.insert_dims(isl.dim_type.set, cur_dim, len(non_inames)) diff --git a/loopy/transform/buffer.py b/loopy/transform/buffer.py index 5af437c05..cd7ee0d06 100644 --- a/loopy/transform/buffer.py +++ b/loopy/transform/buffer.py @@ -102,7 +102,7 @@ class ArrayAccessReplacer(RuleAwareIdentityMapper): continue ax_index = index[i] - from loopy.isl_helpers import simplify_via_aff + from loopy.symbolic import simplify_via_aff ax_index = simplify_via_aff( ax_index - abm.storage_base_indices[i]) diff --git a/loopy/transform/pack_and_unpack_args.py b/loopy/transform/pack_and_unpack_args.py index 9ced3fdc6..daf9316cd 100644 --- a/loopy/transform/pack_and_unpack_args.py +++ b/loopy/transform/pack_and_unpack_args.py @@ -202,7 +202,8 @@ def pack_and_unpack_args_for_call_for_single_kernel(kernel, # {{{ getting the lhs for packing and rhs for unpacking - from loopy.isl_helpers import simplify_via_aff, make_slab + from loopy.symbolic import simplify_via_aff + from loopy.isl_helpers import make_slab flatten_index = simplify_via_aff( sum(dim_tag.stride*idx for dim_tag, idx in diff --git a/loopy/transform/precompute.py b/loopy/transform/precompute.py index 52ee88c20..8b755cca9 100644 --- a/loopy/transform/precompute.py +++ b/loopy/transform/precompute.py @@ -206,7 +206,7 @@ class RuleInvocationReplacer(RuleAwareIdentityMapper): # an iname ax_index = var(sax_source) - from loopy.isl_helpers import simplify_via_aff + from loopy.symbolic import simplify_via_aff ax_index = simplify_via_aff(ax_index - sax_base_idx) stor_subscript.append(ax_index) diff --git a/test/test_isl.py b/test/test_isl.py index 5ddb344e1..720bf1a5f 100644 --- a/test/test_isl.py +++ b/test/test_isl.py @@ -89,8 +89,7 @@ def test_subst_into_pwaff(): def test_simplify_via_aff_reproducibility(): # See https://github.com/inducer/loopy/pull/349 - from loopy.symbolic import parse - from loopy.isl_helpers import simplify_via_aff + from loopy.symbolic import parse, simplify_via_aff expr = parse("i+i_0") -- GitLab