From 7fbc7af440e37205dad8d7400134acc730f53912 Mon Sep 17 00:00:00 2001 From: "[6~" Date: Tue, 12 Nov 2019 22:05:06 -0600 Subject: [PATCH 1/8] Add wave equation solver on the way to diamond tiling demo --- test/test_transform.py | 32 +++++++++++++++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/test/test_transform.py b/test/test_transform.py index cdc0c14b8..0906c9b34 100644 --- a/test/test_transform.py +++ b/test/test_transform.py @@ -552,7 +552,6 @@ def test_split_iname_only_if_in_within(): def test_nested_substs_in_insns(ctx_factory): ctx = ctx_factory() - import loopy as lp ref_knl = lp.make_kernel( "{[i]: 0<=i<10}", @@ -570,6 +569,37 @@ def test_nested_substs_in_insns(ctx_factory): lp.auto_test_vs_ref(ref_knl, ctx, knl) +def test_diamond_tiling(ctx_factory): + ctx = ctx_factory() + queue = cl.CommandQueue(ctx) + + ref_knl = lp.make_kernel( + "[nx,nt] -> {[ix, it]: 1<=ix 1: exec(sys.argv[1]) -- GitLab From af822259a1be28cf30f710fe0bbd4bf731ef54b5 Mon Sep 17 00:00:00 2001 From: "[6~" Date: Wed, 13 Nov 2019 00:56:16 -0600 Subject: [PATCH 2/8] Add initial prototype of map_domain and (currently broken) test of diamond tiling --- loopy/__init__.py | 4 +- loopy/transform/iname.py | 234 +++++++++++++++++++++++++++++++++++++++ test/test_transform.py | 30 ++++- 3 files changed, 263 insertions(+), 5 deletions(-) diff --git a/loopy/__init__.py b/loopy/__init__.py index b60de6e2d..f8a43df6d 100644 --- a/loopy/__init__.py +++ b/loopy/__init__.py @@ -78,7 +78,7 @@ from loopy.transform.iname import ( affine_map_inames, find_unused_axis_tag, make_reduction_inames_unique, has_schedulable_iname_nesting, get_iname_duplication_options, - add_inames_to_insn) + add_inames_to_insn, map_domain) from loopy.transform.instruction import ( find_instructions, map_instructions, @@ -195,7 +195,7 @@ __all__ = [ "affine_map_inames", "find_unused_axis_tag", "make_reduction_inames_unique", "has_schedulable_iname_nesting", "get_iname_duplication_options", - "add_inames_to_insn", + "add_inames_to_insn", "map_domain", "add_prefetch", "change_arg_to_image", "tag_array_axes", "tag_data_axes", diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index 96c8252ef..e627604c1 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -1759,4 +1759,238 @@ def add_inames_to_insn(knl, inames, insn_match): # }}} +# {{{ map_domain + +class _MapDomainMapper(RuleAwareIdentityMapper): + def __init__(self, rule_mapping_context, within, new_inames, substitutions): + super(_MapDomainMapper, self).__init__(rule_mapping_context) + + self.within = within + + self.old_inames = frozenset(substitutions) + self.new_inames = new_inames + + self.substitutions = substitutions + + def map_reduction(self, expr, expn_state): + overlap = frozenset(expr.inames) & self.old_inames + if (overlap + and self.split_iname not in expn_state.arg_context + and self.within( + expn_state.kernel, + expn_state.instruction)): + # FIXME + if len(overlap) != len(self.old_inames): + raise LoopyError(...) + + raise NotImplementedError("reductions") + new_inames = list(expr.inames) + new_inames.remove(self.split_iname) + new_inames.extend([self.outer_iname, self.inner_iname]) + + from loopy.symbolic import Reduction + return Reduction(expr.operation, tuple(new_inames), + self.rec(expr.expr, expn_state), + expr.allow_simultaneous) + else: + return super(_MapDomainMapper, self).map_reduction(expr, expn_state) + + def map_variable(self, expr, expn_state): + if (expr.name in self.old_inames + and expr.name not in expn_state.arg_context + and self.within( + expn_state.kernel, + expn_state.instruction)): + return self.substitutions[expr.name] + else: + return super(_MapDomainMapper, self).map_variable(expr, expn_state) + + +def _find_aff_subst_from_map(iname, isl_map): + if not isinstance(isl_map, isl.BasicMap): + raise RuntimeError("isl_map must be a BasicMap") + + dt, dim_idx = isl_map.get_var_dict()[iname] + + assert dt == dim_type.in_ + + # Force isl to solve for only this iname on its side of the map, by + # projecting out all other "in" variables. + isl_map = isl_map.project_out(dt, dim_idx+1, isl_map.dim(dt)-(dim_idx+1)) + isl_map = isl_map.project_out(dt, 0, dim_idx) + dim_idx = 0 + + # Convert map to set to avoid "domain of affine expression should be a set". + # The old "in" variable will be the last of the out_dims. + new_dim_idx = isl_map.dim(dim_type.out) + isl_map = isl_map.move_dims( + dim_type.out, isl_map.dim(dim_type.out), + dt, dim_idx, 1) + isl_map = isl_map.range() # now a set + dt = dim_type.set + dim_idx = new_dim_idx + del new_dim_idx + + for cns in isl_map.get_constraints(): + if cns.is_equality() and cns.involves_dims(dt, dim_idx, 1): + coeff = cns.get_coefficient_val(dt, dim_idx) + cns_zeroed = cns.set_coefficient_val(dt, dim_idx, 0) + if cns_zeroed.involves_dims(dt, dim_idx, 1): + # not suitable, constraint still involves dim, perhaps in a div + continue + + if coeff.is_one(): + return -cns_zeroed.get_aff() + elif coeff.is_negone(): + return cns_zeroed.get_aff() + else: + # not suitable, coefficient does not have unit coefficient + continue + + raise LoopyError("no suitable equation for '%s' found" % iname) + + +def map_domain(kernel, isl_map, within=None): + # FIXME: Express _split_iname_backend in terms of this + # Missing/deleted for now: + # - slab processing + # - priorities processing + # FIXME: Process priorities + # FIXME: Express affine_map_inames in terms of this, deprecate + # FIXME: Document + + # FIXME: Support within + + # {{{ within processing (disabled for now) + if within is not None: + raise NotImplementedError("within") + + from loopy.match import parse_match + within = parse_match(within) + + # {{{ return the same kernel if no kernel matches + + def _do_not_transform_if_no_within_matches(): + for insn in kernel.instructions: + if within(kernel, insn): + return + + return kernel + + _do_not_transform_if_no_within_matches() + + # }}} + + # }}} + + if not isl_map.is_bijective(): + raise LoopyError("isl_map must be bijective") + + new_inames = frozenset(isl_map.get_var_dict(dim_type.out)) + old_inames = frozenset(isl_map.get_var_dict(dim_type.in_)) + + # {{{ solve for representation of old inames in terms of new + + substitutions = {} + var_substitutions = {} + applied_iname_rewrites = kernel.applied_iname_rewrites[:] + + from loopy.symbolic import aff_to_expr + from pymbolic import var + for iname in old_inames: + substitutions[iname] = aff_to_expr( + _find_aff_subst_from_map(iname, isl_map)) + var_substitutions[var(iname)] = aff_to_expr( + _find_aff_subst_from_map(iname, isl_map)) + + applied_iname_rewrites.append(var_substitutions) + del var_substitutions + + # }}} + + def process_set(s): + var_dict = s.get_var_dict() + + overlap = old_inames & frozenset(var_dict) + if overlap and len(overlap) != len(old_inames): + raise LoopyError("loop domain '%s' involves a part " + "of the map domain inames. Domains must " + "either involve all or none of the map domain " + "inames." % s) + + # {{{ align dims of isl_map and s + + # FIXME: Make this less gross + # FIXME: Make an exported/documented interface of this in islpy + from islpy import _align_dim_type + + map_with_s_domain = isl.Map.from_domain(s) + + dim_types = [dim_type.param, dim_type.in_, dim_type.out] + s_names = [ + map_with_s_domain.get_dim_name(dt, i) + for dt in dim_types + for i in range(map_with_s_domain.dim(dt)) + ] + map_names = [ + isl_map.get_dim_name(dt, i) + for dt in dim_types + for i in range(isl_map.dim(dt)) + ] + aligned_map = _align_dim_type( + dim_type.param, + isl_map, map_with_s_domain, obj_bigger_ok=False, + obj_names=map_names, tgt_names=s_names) + aligned_map = _align_dim_type( + dim_type.in_, + isl_map, map_with_s_domain, obj_bigger_ok=False, + obj_names=map_names, tgt_names=s_names) + + # }}} + + return aligned_map.intersect_domain(s).range() + + # FIXME: Revive _project_out_only_if_all_instructions_in_within + + new_domains = [process_set(dom) for dom in kernel.domains] + + # {{{ update within_inames + + new_insns = [] + for insn in kernel.instructions: + overlap = old_inames & insn.within_inames + if overlap and within(kernel, insn): + if len(overlap) != len(old_inames): + raise LoopyError("instruction '%s' is within only a part " + "of the map domain inames. Instructions must " + "either be within all or none of the map domain " + "inames." % insn.id) + + insn = insn.copy( + within_inames=(insn.within_inames - old_inames) | new_inames) + else: + # leave insn unmodified + pass + + new_insns.append(insn) + + # }}} + + kernel = kernel.copy( + domains=new_domains, + instructions=new_insns, + applied_iname_rewrites=applied_iname_rewrites) + + rule_mapping_context = SubstitutionRuleMappingContext( + kernel.substitutions, kernel.get_var_name_generator()) + ins = _MapDomainMapper(rule_mapping_context, within, + new_inames, substitutions) + + kernel = ins.map_kernel(kernel) + kernel = rule_mapping_context.finish_kernel(kernel) + + return kernel + +# }}} + # vim: foldmethod=marker diff --git a/test/test_transform.py b/test/test_transform.py index 0906c9b34..9b184f4b0 100644 --- a/test/test_transform.py +++ b/test/test_transform.py @@ -578,21 +578,45 @@ def test_diamond_tiling(ctx_factory): """ u[ix, it+2] = ( 2*u[ix, it+1] - + dt**2/dx**2 * (u[ix+1, it+1] - 2*u[ix, it+1] + u[ix-1, it+1]) + + dt**2/dx**2 * (u[ix+1, it+1] - 2*u[ix, it+1] + u[ix-1, it+1]) - u[ix, it]) """) + ref_knl = lp.set_options(ref_knl, write_cl=True) + + # FIXME: Handle priorities in map_domain + knl_for_transform = ref_knl ref_knl = lp.prioritize_loops(ref_knl, "it, ix") - ref_knl = lp.set_options(ref_knl, write_cl=True) nx = 43 u = np.zeros((nx, 200)) x = np.linspace(-1, 1, nx) dx = x[1] - x[0] u[:, 0] = u[:, 1] = np.exp(-100*x**2) + import islpy as isl + if 1: + m = isl.BasicMap( + "[nx,nt] -> {[ix, it] -> [tx, tt, tparity, itt, itx]: " + "16*(tx - tt + tparity) + itx - itt = ix - it and " + "16*(tx + tt) + itt + itx = ix + it and " + "0<=tparity<2 and 0 <= itx - itt < 16 and 0 <= itt+itx < 16}") + knl = lp.map_domain(knl_for_transform, m) + knl = lp.prioritize_loops(knl, "tt,tparity,tx,itt,itx") + else: + # This is more like what split_iname does, but it is *not* + # a correct tiling for the stencil operator. + m = isl.BasicMap( + "[nx,nt] -> {[ix, it] -> [tx, tt, itt, itx]: " + "16*tx + itx = ix and " + "16*tt + itt = it and " + "0 <= itx < 16 and 0 <= itt< 16}") + + knl = lp.map_domain(knl_for_transform, m) + knl = lp.prioritize_loops(knl, "tt,tx,itt,itx") + u_dev = cl.array.to_device(queue, u) - ref_knl(queue, u=u_dev, dx=dx, dt=dx) + knl(queue, u=u_dev, dx=dx, dt=dx) u = u_dev.get() import matplotlib.pyplot as plt -- GitLab From 70ee84dfda8139f8826aada1d49f4a6c9a0d7f24 Mon Sep 17 00:00:00 2001 From: "[6~" Date: Wed, 13 Nov 2019 18:07:41 -0600 Subject: [PATCH 3/8] Make placeholder error message Py2-compatible --- 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 e627604c1..6236ace72 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -1781,7 +1781,7 @@ class _MapDomainMapper(RuleAwareIdentityMapper): expn_state.instruction)): # FIXME if len(overlap) != len(self.old_inames): - raise LoopyError(...) + raise LoopyError("...") raise NotImplementedError("reductions") new_inames = list(expr.inames) -- GitLab From 3c4c1ff2e46a8feecf7ac1053cb2a90fc7aebb25 Mon Sep 17 00:00:00 2001 From: "[6~" Date: Wed, 13 Nov 2019 18:07:59 -0600 Subject: [PATCH 4/8] test_diamond_tiling: Remove split_iname analog --- test/test_transform.py | 26 +++++++------------------- 1 file changed, 7 insertions(+), 19 deletions(-) diff --git a/test/test_transform.py b/test/test_transform.py index 9b184f4b0..6c15b91ec 100644 --- a/test/test_transform.py +++ b/test/test_transform.py @@ -595,25 +595,13 @@ def test_diamond_tiling(ctx_factory): u[:, 0] = u[:, 1] = np.exp(-100*x**2) import islpy as isl - if 1: - m = isl.BasicMap( - "[nx,nt] -> {[ix, it] -> [tx, tt, tparity, itt, itx]: " - "16*(tx - tt + tparity) + itx - itt = ix - it and " - "16*(tx + tt) + itt + itx = ix + it and " - "0<=tparity<2 and 0 <= itx - itt < 16 and 0 <= itt+itx < 16}") - knl = lp.map_domain(knl_for_transform, m) - knl = lp.prioritize_loops(knl, "tt,tparity,tx,itt,itx") - else: - # This is more like what split_iname does, but it is *not* - # a correct tiling for the stencil operator. - m = isl.BasicMap( - "[nx,nt] -> {[ix, it] -> [tx, tt, itt, itx]: " - "16*tx + itx = ix and " - "16*tt + itt = it and " - "0 <= itx < 16 and 0 <= itt< 16}") - - knl = lp.map_domain(knl_for_transform, m) - knl = lp.prioritize_loops(knl, "tt,tx,itt,itx") + m = isl.BasicMap( + "[nx,nt] -> {[ix, it] -> [tx, tt, tparity, itt, itx]: " + "16*(tx - tt + tparity) + itx - itt = ix - it and " + "16*(tx + tt) + itt + itx = ix + it and " + "0<=tparity<2 and 0 <= itx - itt < 16 and 0 <= itt+itx < 16}") + knl = lp.map_domain(knl_for_transform, m) + knl = lp.prioritize_loops(knl, "tt,tparity,tx,itt,itx") u_dev = cl.array.to_device(queue, u) knl(queue, u=u_dev, dx=dx, dt=dx) -- GitLab From b61dbcf7aeb705e7071d2129fc20bb40facbf076 Mon Sep 17 00:00:00 2001 From: "[6~" Date: Wed, 13 Nov 2019 18:08:47 -0600 Subject: [PATCH 5/8] Fix diamond tile mapping --- test/test_transform.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/test_transform.py b/test/test_transform.py index 6c15b91ec..9cb1af4ab 100644 --- a/test/test_transform.py +++ b/test/test_transform.py @@ -597,8 +597,8 @@ def test_diamond_tiling(ctx_factory): import islpy as isl m = isl.BasicMap( "[nx,nt] -> {[ix, it] -> [tx, tt, tparity, itt, itx]: " - "16*(tx - tt + tparity) + itx - itt = ix - it and " - "16*(tx + tt) + itt + itx = ix + it and " + "16*(tx - tt) + itx - itt = ix - it and " + "16*(tx + tt + tparity) + itt + itx = ix + it and " "0<=tparity<2 and 0 <= itx - itt < 16 and 0 <= itt+itx < 16}") knl = lp.map_domain(knl_for_transform, m) knl = lp.prioritize_loops(knl, "tt,tparity,tx,itt,itx") -- GitLab From c6610312666c6f8bb38acbb439bf7fc1199aa529 Mon Sep 17 00:00:00 2001 From: "[6~" Date: Thu, 14 Nov 2019 16:37:20 -0600 Subject: [PATCH 6/8] Make test_diamond_tiling an actual test --- test/test_transform.py | 38 ++++++++++++++++++++++++-------------- 1 file changed, 24 insertions(+), 14 deletions(-) diff --git a/test/test_transform.py b/test/test_transform.py index 9cb1af4ab..d8030358b 100644 --- a/test/test_transform.py +++ b/test/test_transform.py @@ -569,7 +569,7 @@ def test_nested_substs_in_insns(ctx_factory): lp.auto_test_vs_ref(ref_knl, ctx, knl) -def test_diamond_tiling(ctx_factory): +def test_diamond_tiling(ctx_factory, interactive=False): ctx = ctx_factory() queue = cl.CommandQueue(ctx) @@ -581,19 +581,12 @@ def test_diamond_tiling(ctx_factory): + dt**2/dx**2 * (u[ix+1, it+1] - 2*u[ix, it+1] + u[ix-1, it+1]) - u[ix, it]) """) - ref_knl = lp.set_options(ref_knl, write_cl=True) # FIXME: Handle priorities in map_domain knl_for_transform = ref_knl ref_knl = lp.prioritize_loops(ref_knl, "it, ix") - nx = 43 - u = np.zeros((nx, 200)) - x = np.linspace(-1, 1, nx) - dx = x[1] - x[0] - u[:, 0] = u[:, 1] = np.exp(-100*x**2) - import islpy as isl m = isl.BasicMap( "[nx,nt] -> {[ix, it] -> [tx, tt, tparity, itt, itx]: " @@ -603,13 +596,30 @@ def test_diamond_tiling(ctx_factory): knl = lp.map_domain(knl_for_transform, m) knl = lp.prioritize_loops(knl, "tt,tparity,tx,itt,itx") - u_dev = cl.array.to_device(queue, u) - knl(queue, u=u_dev, dx=dx, dt=dx) + if interactive: + nx = 43 + u = np.zeros((nx, 200)) + x = np.linspace(-1, 1, nx) + dx = x[1] - x[0] + u[:, 0] = u[:, 1] = np.exp(-100*x**2) + + u_dev = cl.array.to_device(queue, u) + knl(queue, u=u_dev, dx=dx, dt=dx) - u = u_dev.get() - import matplotlib.pyplot as plt - plt.imshow(u.T) - plt.show() + u = u_dev.get() + import matplotlib.pyplot as plt + plt.imshow(u.T) + plt.show() + else: + types = {"dt,dx,u": np.float64} + knl = lp.add_and_infer_dtypes(knl, types) + ref_knl = lp.add_and_infer_dtypes(ref_knl, types) + + lp.auto_test_vs_ref(ref_knl, ctx, knl, + parameters={ + "nx": 200, "nt": 300, + "dx": 1, "dt": 1 + }) if __name__ == "__main__": -- GitLab From e1ae8cf8080f93f8cb1252a194414c2cf60f1f78 Mon Sep 17 00:00:00 2001 From: "[6~" Date: Thu, 14 Nov 2019 16:37:43 -0600 Subject: [PATCH 7/8] Fix handling of reductions in map_domain --- loopy/transform/iname.py | 32 +++++++++++++++++++++++--------- 1 file changed, 23 insertions(+), 9 deletions(-) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index 6236ace72..c49f26137 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -1773,20 +1773,34 @@ class _MapDomainMapper(RuleAwareIdentityMapper): self.substitutions = substitutions def map_reduction(self, expr, expn_state): - overlap = frozenset(expr.inames) & self.old_inames - if (overlap - and self.split_iname not in expn_state.arg_context + red_overlap = frozenset(expr.inames) & self.old_inames + arg_ctx_overlap = frozenset(expn_state.arg_context) & self.old_inames + if (red_overlap and self.within( expn_state.kernel, expn_state.instruction)): - # FIXME - if len(overlap) != len(self.old_inames): - raise LoopyError("...") + if len(red_overlap) != len(self.old_inames): + raise LoopyError("reduction '%s' involves a part " + "of the map domain inames. Reductions must " + "either involve all or none of the map domain " + "inames." % str(expr)) + + if arg_ctx_overlap: + if arg_ctx_overlap == red_overlap: + # All variables are shadowed by context, that's OK. + return super(_MapDomainMapper, self).map_reduction( + expr, expn_state) + else: + raise LoopyError("reduction '%s' has" + "some of the reduction variables affected " + "by the map_domain shadowed by context. " + "Either all or none must be shadowed." + % str(expr)) - raise NotImplementedError("reductions") new_inames = list(expr.inames) - new_inames.remove(self.split_iname) - new_inames.extend([self.outer_iname, self.inner_iname]) + for old_iname in self.old_inames: + new_inames.remove(old_iname) + new_inames.extend(self.new_inames) from loopy.symbolic import Reduction return Reduction(expr.operation, tuple(new_inames), -- GitLab From da35b59d0188c0e838178235bf3d742df2b4c6a4 Mon Sep 17 00:00:00 2001 From: Nicholas Christensen Date: Mon, 1 Feb 2021 11:51:08 -0600 Subject: [PATCH 8/8] Update to use new islpy _laign_dim_type parameters --- loopy/transform/iname.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index fbd0d57bc..d832adbd7 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -1975,6 +1975,16 @@ def map_domain(kernel, isl_map, within=None): for dt in dim_types for i in range(isl_map.dim(dt)) ] + aligned_map = _align_dim_type( + dim_type.param, + isl_map, map_with_s_domain, False, + map_names, s_names) + aligned_map = _align_dim_type( + dim_type.in_, + isl_map, map_with_s_domain, False, + map_names, s_names) + # Old code + """ aligned_map = _align_dim_type( dim_type.param, isl_map, map_with_s_domain, obj_bigger_ok=False, @@ -1983,7 +1993,7 @@ def map_domain(kernel, isl_map, within=None): dim_type.in_, isl_map, map_with_s_domain, obj_bigger_ok=False, obj_names=map_names, tgt_names=s_names) - + """ # }}} return aligned_map.intersect_domain(s).range() -- GitLab