diff --git a/loopy/schedule/checker/utils.py b/loopy/schedule/checker/utils.py index 12d3480e1..2ba5f37e8 100644 --- a/loopy/schedule/checker/utils.py +++ b/loopy/schedule/checker/utils.py @@ -91,6 +91,48 @@ def reorder_dims_by_name( return new_set +def move_dim_to_index( + isl_map, dim_name, dt, destination_idx): + """Return an isl map with the specified dimension moved to + the specified index. + + :arg isl_map: A :class:`islpy.Map`. + + :arg dim_name: A :class:`str` specifying the name of the dimension + to be moved. + + :arg dt: A :class:`islpy.dim_type`, i.e., an :class:`int`, + specifying the type of dimension to be reordered. + + :arg destination_idx: A :class:`int` specifying the desired dimension + index of the dimention to be moved. + + :returns: An :class:`islpy.Map` matching `isl_map` with the + specified dimension moved to the specified index. + + """ + + assert dt != dim_type.param + + layover_dt = dim_type.param + layover_dim_len = len(isl_map.get_var_names(layover_dt)) + + current_idx = isl_map.find_dim_by_name(dt, dim_name) + if current_idx == -1: + raise ValueError("Dimension name %s not found in dim type %s of %s" + % (dim_name, dt, isl_map)) + + if current_idx != destination_idx: + # First move to other dim because isl is stupid + new_map = isl_map.move_dims( + layover_dt, layover_dim_len, dt, current_idx, 1) + # Now move it where we actually want it + new_map = new_map.move_dims( + dt, destination_idx, layover_dt, layover_dim_len, 1) + + return new_map + + def move_dims_by_name( isl_obj, dst_type, dst_pos_start, src_type, dim_names): dst_pos = dst_pos_start @@ -197,7 +239,6 @@ def append_mark_to_isl_map_var_names(old_isl_map, dt, mark): def append_mark_to_strings(strings, mark): - assert isinstance(strings, list) return [s+mark for s in strings] @@ -340,6 +381,14 @@ def sorted_union_of_names_in_isl_sets( return sorted(inames) +def convert_map_to_set(isl_map): + n_in_dims = len(isl_map.get_var_names(dim_type.in_)) + n_out_dims = len(isl_map.get_var_names(dim_type.out)) + return isl_map.move_dims( + dim_type.in_, n_in_dims, dim_type.out, 0, n_out_dims + ).domain(), n_in_dims, n_out_dims + + def create_symbolic_map_from_tuples( tuple_pairs_with_domains, space, diff --git a/loopy/transform/add_barrier.py b/loopy/transform/add_barrier.py index 7a220418f..e738c7130 100644 --- a/loopy/transform/add_barrier.py +++ b/loopy/transform/add_barrier.py @@ -90,10 +90,37 @@ def add_barrier(kernel, insn_before="", insn_after="", id_based_on=None, new_kernel = kernel.copy(instructions=kernel.instructions + [barrier_to_add]) if insn_after is not None: + # TODO this should be a new dependency new_kernel = add_dependency(kernel=new_kernel, insn_match=insn_after, depends_on="id:"+id) + for insn_before_id in insns_before: + # make v2 dep: + from loopy.schedule.checker.utils import ( + append_mark_to_strings, + make_dep_map, + ) + from loopy.schedule.checker.schedule import BEFORE_MARK + inames_before = new_kernel.id_to_insn[insn_before_id].within_inames + inames_before_marked = append_mark_to_strings( + inames_before, BEFORE_MARK) + + inames_after = set(within_inames) if within_inames else set() + + shared_inames = inames_after & inames_before + + in_space_str = ", ".join(inames_before_marked) + out_space_str = ", ".join(inames_after) + constraint_str = " and ".join([ + "{0}{1} = {0}".format(iname, BEFORE_MARK) for iname in shared_inames]) + + dep_v2 = make_dep_map( + f"{{ [{in_space_str}] -> [{out_space_str}] : {constraint_str} }}", + knl_with_domains=new_kernel) + new_kernel = add_dependency( + new_kernel, "id:"+id, ("id:"+insn_before_id, dep_v2)) + return new_kernel # }}} diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index 85a548896..cd6568a8d 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -272,6 +272,84 @@ def _split_iname_backend(kernel, iname_to_split, fixed_length, fixed_length_is_inner) for dom in kernel.domains] + # {{{ Split iname in dependencies + + from loopy.transform.instruction import map_dependency_maps + from loopy.schedule.checker.schedule import BEFORE_MARK + from loopy.schedule.checker.utils import ( + convert_map_to_set, + remove_dims_by_name, + ) + + def _split_iname_in_depender(dep): + + # If iname is not present in dep, return unmodified dep + if iname_to_split not in dep.get_var_names(dim_type.out): + return dep + + # Temporarily convert map to set for processing + set_from_map, n_in_dims, n_out_dims = convert_map_to_set(dep) + + # Split iname + set_from_map = _split_iname_in_set( + set_from_map, iname_to_split, inner_iname, outer_iname, + fixed_length, fixed_length_is_inner) + + # Dim order: [old_inames' ..., old_inames ..., i_outer, i_inner] + + # Convert set back to map + map_from_set = isl.Map.from_domain(set_from_map) + # Move original out dims + 2 new dims: + map_from_set = map_from_set.move_dims( + dim_type.out, 0, dim_type.in_, n_in_dims, n_out_dims+2) + + # Remove iname that was split: + map_from_set = remove_dims_by_name( + map_from_set, dim_type.out, [iname_to_split]) + + return map_from_set + + def _split_iname_in_dependee(dep): + + iname_to_split_marked = iname_to_split+BEFORE_MARK + + # If iname is not present in dep, return unmodified dep + if iname_to_split_marked not in dep.get_var_names(dim_type.in_): + return dep + + # Temporarily convert map to set for processing + set_from_map, n_in_dims, n_out_dims = convert_map_to_set(dep) + + # Split iname' + set_from_map = _split_iname_in_set( + set_from_map, iname_to_split_marked, + inner_iname+BEFORE_MARK, outer_iname+BEFORE_MARK, + fixed_length, fixed_length_is_inner) + + # Dim order: [old_inames' ..., old_inames ..., i_outer', i_inner'] + + # Convert set back to map + map_from_set = isl.Map.from_domain(set_from_map) + # Move original out dims new dims: + map_from_set = map_from_set.move_dims( + dim_type.out, 0, dim_type.in_, n_in_dims, n_out_dims) + + # Remove iname that was split: + map_from_set = remove_dims_by_name( + map_from_set, dim_type.in_, [iname_to_split_marked]) + + return map_from_set + + false_id_match = "not id:*" + kernel = map_dependency_maps( + kernel, _split_iname_in_depender, + stmt_match_depender=within, stmt_match_dependee=false_id_match) + kernel = map_dependency_maps( + kernel, _split_iname_in_dependee, + stmt_match_depender=false_id_match, stmt_match_dependee=within) + + # }}} + from pymbolic import var inner = var(inner_iname) outer = var(outer_iname) @@ -881,7 +959,7 @@ def duplicate_inames(kernel, inames, within, new_inames=None, suffix=None, new_inames = [iname.strip() for iname in new_inames.split(",")] from loopy.match import parse_stack_match - within = parse_stack_match(within) + within_sm = parse_stack_match(within) if new_inames is None: new_inames = [None] * len(inames) @@ -891,6 +969,7 @@ def duplicate_inames(kernel, inames, within, new_inames=None, suffix=None, name_gen = kernel.get_var_name_generator() + # Generate new iname names for i, iname in enumerate(inames): new_iname = new_inames[i] @@ -913,7 +992,7 @@ def duplicate_inames(kernel, inames, within, new_inames=None, suffix=None, # }}} - # {{{ duplicate the inames + # {{{ duplicate the inames in domains for old_iname, new_iname in zip(inames, new_inames): from loopy.kernel.tools import DomainChanger @@ -924,6 +1003,40 @@ def duplicate_inames(kernel, inames, within, new_inames=None, suffix=None, domains=domch.get_domains_with( duplicate_axes(domch.domain, [old_iname], [new_iname]))) + # {{{ *Rename* iname in dependencies + + # TODO use find_and_rename_dims for simpler code + # (see example in rename_iname) + from loopy.transform.instruction import map_dependency_maps + from loopy.schedule.checker.schedule import BEFORE_MARK + old_iname_p = old_iname+BEFORE_MARK + new_iname_p = new_iname+BEFORE_MARK + + def _rename_iname_in_dim_out(dep): + # update iname in out-dim + out_idx = dep.find_dim_by_name(dim_type.out, old_iname) + if out_idx != -1: + dep = dep.set_dim_name(dim_type.out, out_idx, new_iname) + return dep + + def _rename_iname_in_dim_in(dep): + # update iname in in-dim + in_idx = dep.find_dim_by_name(dim_type.in_, old_iname_p) + if in_idx != -1: + dep = dep.set_dim_name(dim_type.in_, in_idx, new_iname_p) + return dep + + # TODO figure out match vs stack_match + false_id_match = "not id:*" + kernel = map_dependency_maps( + kernel, _rename_iname_in_dim_out, + stmt_match_depender=within, stmt_match_dependee=false_id_match) + kernel = map_dependency_maps( + kernel, _rename_iname_in_dim_in, + stmt_match_depender=false_id_match, stmt_match_dependee=within) + + # }}} + # }}} # {{{ change the inames in the code @@ -932,10 +1045,10 @@ def duplicate_inames(kernel, inames, within, new_inames=None, suffix=None, kernel.substitutions, name_gen) indup = _InameDuplicator(rule_mapping_context, old_to_new=dict(list(zip(inames, new_inames))), - within=within) + within=within_sm) kernel = rule_mapping_context.finish_kernel( - indup.map_kernel(kernel, within=within)) + indup.map_kernel(kernel, within=within_sm)) # }}} @@ -1200,6 +1313,41 @@ def rename_iname(kernel, old_iname, new_iname, existing_ok=False, within=None): kernel = kernel.copy(instructions=new_instructions) + # {{{ Rename iname in dependencies + + from loopy.transform.instruction import map_dependency_maps + from loopy.schedule.checker.schedule import BEFORE_MARK + from loopy.isl_helpers import find_and_rename_dims + old_iname_p = old_iname+BEFORE_MARK + new_iname_p = new_iname+BEFORE_MARK + + def _rename_iname_in_dim_out(dep): + # Update iname in out-dim (depender dim). + + # For now, dim should definitely exist because this will only + # be called on dependers + return find_and_rename_dims( + dep, dim_type.out, {old_iname: new_iname}) + + def _rename_iname_in_dim_in(dep): + # Update iname in in-dim (dependee dim). + + # For now, dim should definitely exist because this will only + # be called on dependees + return find_and_rename_dims( + dep, dim_type.in_, {old_iname_p: new_iname_p}) + + # TODO figure out match vs stack_match + false_id_match = "not id:*" + kernel = map_dependency_maps( + kernel, _rename_iname_in_dim_out, + stmt_match_depender=within, stmt_match_dependee=false_id_match) + kernel = map_dependency_maps( + kernel, _rename_iname_in_dim_in, + stmt_match_depender=false_id_match, stmt_match_dependee=within) + + # }}} + else: kernel = duplicate_inames( kernel, [old_iname], within=within, new_inames=[new_iname]) @@ -1227,6 +1375,19 @@ def get_used_inames(kernel): return used_inames +def remove_vars_from_set(s, remove_vars): + from copy import deepcopy + new_s = deepcopy(s) + for var in remove_vars: + try: + dt, idx = s.get_var_dict()[var] + except KeyError: + continue + else: + new_s = new_s.project_out(dt, idx, 1) + return new_s + + @for_each_kernel def remove_unused_inames(kernel, inames=None): """Delete those among *inames* that are unused, i.e. project them @@ -1254,22 +1415,26 @@ def remove_unused_inames(kernel, inames=None): # {{{ remove them - domains = kernel.domains - for iname in unused_inames: - new_domains = [] + new_domains = [] + for dom in kernel.domains: + new_domains.append(remove_vars_from_set(dom, unused_inames)) + + kernel = kernel.copy(domains=new_domains) - for dom in domains: - try: - dt, idx = dom.get_var_dict()[iname] - except KeyError: - pass - else: - dom = dom.project_out(dt, idx, 1) - new_domains.append(dom) + # }}} - domains = new_domains + # {{{ Remove inames from deps - kernel = kernel.copy(domains=domains) + from loopy.transform.instruction import map_dependency_maps + from loopy.schedule.checker.schedule import BEFORE_MARK + from loopy.schedule.checker.utils import append_mark_to_strings + unused_inames_marked = append_mark_to_strings(unused_inames, BEFORE_MARK) + + def _remove_iname_from_dep(dep): + return remove_vars_from_set( + remove_vars_from_set(dep, unused_inames), unused_inames_marked) + + kernel = map_dependency_maps(kernel, _remove_iname_from_dep) # }}} @@ -1903,7 +2068,8 @@ def _find_aff_subst_from_map(iname, isl_map): # 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, dim_idx+1, isl_map.dim(dt)-(dim_idx+1)) isl_map = isl_map.project_out(dt, 0, dim_idx) dim_idx = 0 @@ -2097,10 +2263,16 @@ def map_domain(kernel, transform_map): # {{{ Function to apply mapping to one set + from loopy.isl_helpers import find_and_rename_dims + def process_set(s): """Return the transformed set. Assume that map is applicable to this set.""" + # At this point, overlap condition check guarantees that the + # in-dims of the transform map are a subset of the dims we're + # about to change. + # {{{ Align dims of transform_map and s so that map can be applied # Create a map whose input space matches the set @@ -2122,18 +2294,17 @@ def process_set(s): # {{{ Align transform map input dims with set dims # FIXME: Make an exported/documented interface of this in islpy - - dim_types = [dim_type.param, dim_type.in_, dim_type.out] + dts = [dim_type.param, dim_type.in_, dim_type.out] # Variables found in iname domain set s_names = { map_with_s_domain.get_dim_name(dt, i) - for dt in dim_types + for dt in dts for i in range(map_with_s_domain.dim(dt)) } # Variables found in transform map map_names = { augmented_transform_map.get_dim_name(dt, i) - for dt in dim_types + for dt in dts for i in range(augmented_transform_map.dim(dt)) } # (_align_dim_type uses these two sets to determine which names are in @@ -2157,8 +2328,6 @@ def process_set(s): new_s = aligned_map.intersect_domain(s).range() # Now rename any proxy dims back to their original names - - from loopy.isl_helpers import find_and_rename_dims new_s = find_and_rename_dims( new_s, dim_type.set, dict([pair[::-1] for pair in proxy_name_pairs])) # (reverse pair order) @@ -2213,6 +2382,156 @@ def process_set(s): # }}} + # {{{ update dependencies + + # Prep transform map to be applied to dependency + from loopy.transform.instruction import map_dependency_maps + from loopy.schedule.checker.utils import ( + append_mark_to_isl_map_var_names, + move_dim_to_index, + ) + from loopy.schedule.checker.schedule import ( + BEFORE_MARK, + STATEMENT_VAR_NAME, + ) + + names_to_ignore = set([STATEMENT_VAR_NAME, STATEMENT_VAR_NAME+BEFORE_MARK]) + + # Create version of transform map with before marks + # (for aligning when applying map to dependee portion of deps) + transform_map_marked = append_mark_to_isl_map_var_names( + append_mark_to_isl_map_var_names(transform_map, dim_type.in_, BEFORE_MARK), + dim_type.out, BEFORE_MARK) + + def _apply_transform_map_to_depender(dep_map): + # (since 'out' dim of dep is unmarked, use unmarked transform map) + + # Find overlap between transform map inames and *depender* inames + # (ignore the statement var name, which may have been + # added to a transform map or s) + mapped_inames = set( + transform_map.get_var_names(dim_type.in_)) - names_to_ignore + overlap = set(dep_map.range().get_var_dict()) & mapped_inames + + if overlap: + # All of the transform map inames must be in the overlap. + if len(overlap) != len(mapped_inames): + raise LoopyError( + "Attempted to transform range of dependency map %s " + "(which corresponds to the depender statement) using " + "transform map %s, but at least one dependency iname is not " + "present in the transform map." % (dep_map, transform_map)) + + # At this point, overlap condition check guarantees that the + # in-dims of the transform map are a subset of the dims we're + # about to change. + + # If there are any out-dims (depender dims) in dep_map that are not + # mapped by the transform map, add them to the in/out space of the + # transform map so that they remain unchanged. + # (temporary proxy dim names are needed in out space of transform + # map because isl won't allow any dim names to match, i.e., instead + # of just mapping {[unused_name]->[unused_name]}, we have to map + # {[unused_name]->[unused_name__prox] : unused_name__prox = unused_name}, + # and then rename unused_name__prox afterward.) + ( + augmented_trans_map, proxy_name_pairs + ) = _apply_identity_for_missing_map_dims( + transform_map, dep_map.get_var_names(dim_type.out)) + + # Align 'in_' dim of transform map with 'out' dim of dep + from loopy.schedule.checker.utils import reorder_dims_by_name + augmented_trans_map_aligned = reorder_dims_by_name( + augmented_trans_map, dim_type.in_, + dep_map.get_var_names(dim_type.out)) + + # Apply transform map to dep output dims + new_dep_map = dep_map.apply_range(augmented_trans_map_aligned) + + # Now rename the proxy dims back to their original names + new_dep_map = find_and_rename_dims( + new_dep_map, dim_type.out, + dict([pair[::-1] for pair in proxy_name_pairs])) # (reverse order) + + # Statement var may have moved, so put it back at the beginning + new_dep_map = move_dim_to_index( + new_dep_map, STATEMENT_VAR_NAME, dim_type.out, 0) + + return new_dep_map + else: # No overlap + # Transform map inames are not present in depender, don't change dep_map + return dep_map + + def _apply_transform_map_to_dependee(dep_map): + # (since 'in_' dim of dep is marked, use transform_map_marked) + + # Find overlap between transform map inames and *dependee* inames + # (ignore the statement var name, which may have been + # added to a transform map or s) + mapped_inames = set( + transform_map_marked.get_var_names(dim_type.in_)) - names_to_ignore + overlap = set(dep_map.domain().get_var_dict()) & mapped_inames + + if overlap: + # All of the transform map inames must be in the overlap. + if len(overlap) != len(mapped_inames): + raise LoopyError( + "Attempted to transform domain of dependency map %s " + "(which corresponds to the dependee statement) using " + "transform map %s, but at least one dependency iname is not " + "present in the transform map." % (dep_map, transform_map)) + + # At this point, overlap condition check guarantees that the + # in-dims of the transform map are a subset of the dims we're + # about to change. + + # If there are any in-dims (dependee dims) in dep_map that are not + # mapped by the transform map, add them to the in/out space of the + # transform map so that they remain unchanged. + # (temporary proxy dim names are needed in out space of transform + # map because isl won't allow any dim names to match, i.e., instead + # of just mapping {[unused_name]->[unused_name]}, we have to map + # {[unused_name]->[unused_name__prox] : unused_name__prox = unused_name}, + # and then rename unused_name__prox afterward.) + ( + augmented_trans_map_marked, proxy_name_pairs + ) = _apply_identity_for_missing_map_dims( + transform_map_marked, dep_map.get_var_names(dim_type.in_)) + + # Align 'in_' dim of transform map with 'in_' dim of dep + from loopy.schedule.checker.utils import reorder_dims_by_name + augmented_trans_map_aligned = reorder_dims_by_name( + augmented_trans_map_marked, dim_type.in_, + dep_map.get_var_names(dim_type.in_)) + + # Apply transform map to dep input dims + new_dep_map = dep_map.apply_domain(augmented_trans_map_aligned) + + # Now rename the proxy dims back to their original names + new_dep_map = find_and_rename_dims( + new_dep_map, dim_type.in_, + dict([pair[::-1] for pair in proxy_name_pairs])) # (reverse order) + + # Statement var may have moved, so put it back at the beginning + new_dep_map = move_dim_to_index( + new_dep_map, STATEMENT_VAR_NAME+BEFORE_MARK, dim_type.in_, 0) + + return new_dep_map + else: # No overlap + # Transform map inames are not present in dependee, don't change dep_map + return dep_map + + false_id_match = "not id:*" + true_id_match = "id:*" + kernel = map_dependency_maps( + kernel, _apply_transform_map_to_depender, + stmt_match_depender=true_id_match, stmt_match_dependee=false_id_match) + kernel = map_dependency_maps( + kernel, _apply_transform_map_to_dependee, + stmt_match_depender=false_id_match, stmt_match_dependee=true_id_match) + + # }}} + # {{{ Update within_inames for each statement # If we get this far, we know that the map was applied to exactly one domain, @@ -2264,6 +2583,8 @@ def process_set(s): # }}} +# {{{ add_inames_for_unused_hw_axes + @for_each_kernel def add_inames_for_unused_hw_axes(kernel, within=None): """ @@ -2374,4 +2695,6 @@ def add_inames_for_unused_hw_axes(kernel, within=None): return kernel.copy(instructions=new_insns) +# }}} + # vim: foldmethod=marker diff --git a/loopy/transform/instruction.py b/loopy/transform/instruction.py index 3c3c16274..9548aa11e 100644 --- a/loopy/transform/instruction.py +++ b/loopy/transform/instruction.py @@ -173,6 +173,131 @@ def add_dep(stmt): # }}} +# {{{ map dependencies v2 + +# Terminiology: +# stmtX.dependencies: # <- "stmt dependencies" = full dict of deps +# {stmt0: [dep_map00, dep_map01, ...], # <- "one dependency" +# stmt1: [dep_map10, dep_map11, ...], +# ...} +# one dependency includes one "dependency list", which contains "dep maps" + + +# {{{ _parse_match_if_necessary + +def _parse_match_if_necessary(match_candidate): + from loopy.match import ( + MatchExpressionBase, + StackMatch, + ) + if not isinstance( + match_candidate, (MatchExpressionBase, StackMatch)): + from loopy.match import parse_match + # TODO assumes StackMatches are already parsed + # TODO determine when to use parse_stack_match (AKQ) + return parse_match(match_candidate) + else: + return match_candidate + +# }}} + + +# {{{ map_dependency_lists + +def map_dependency_lists( + kernel, f, stmt_match_depender="id:*", stmt_match_dependee="id:*"): + """Replace dependency_list with f(dependency_list) if conditions match. + + Let dependency_list = depender_stmt.dependencies[dependee_stmt]; + + Set dependency_list = f(dependency_list) if: + + - stmt_match_depender matches depender_stmt OR + - stmt_match_dependee matches dependee_stmt + - (but don't call f() twice if both depender and dependee match) + + """ + # TODO refactor this so that separate depender and dependee updating + # functions can be provided and used in a single pass + from loopy.match import ( + StackMatch, + ) + + match_depender = _parse_match_if_necessary(stmt_match_depender) + match_dependee = _parse_match_if_necessary(stmt_match_dependee) + + # TODO figure out right way to simultaneously handle + # both MatchExpressionBase and StackMatch + if isinstance(match_depender, StackMatch): + extra_match_depender_args = [()] + else: + extra_match_depender_args = [] + if isinstance(match_dependee, StackMatch): + extra_match_dependee_args = [()] + else: + extra_match_dependee_args = [] + + new_stmts = [] + for stmt in kernel.instructions: + new_deps = {} + if match_depender(kernel, stmt, *extra_match_depender_args): + # Stmt matches as depender + # Replace all deps + for dep_id, dep_maps in stmt.dependencies.items(): + new_deps[dep_id] = f(dep_maps) + else: + # Stmt didn't match as a depender + # Replace deps matching dependees + for dep_id, dep_maps in stmt.dependencies.items(): + if match_dependee( + kernel, kernel.id_to_insn[dep_id], + *extra_match_dependee_args): + new_deps[dep_id] = f(dep_maps) + else: + new_deps[dep_id] = dep_maps + new_stmts.append(stmt.copy(dependencies=new_deps)) + + return kernel.copy(instructions=new_stmts) + +# }}} + + +# {{{ map_dependency_maps + +def map_dependency_maps( + kernel, f, stmt_match_depender="id:*", stmt_match_dependee="id:*"): + # Set dep_map = f(dep_map) for dep_map in: + # All dependencies of stmts matching stmt_match_depender + # All dependencies ON stmts matching stmt_match_dependee + + def _update_dep_maps(dep_maps): + return [f(dep_map) for dep_map in dep_maps] + + return map_dependency_lists( + kernel, _update_dep_maps, stmt_match_depender, stmt_match_dependee) + +# }}} + + +# {{{ map_stmt_dependencies + +def map_stmt_dependencies(kernel, stmt_match, f): + """Set stmt.dependences = f(stmt.dependencies) for stmts matching stmt_match""" + # (Only modifies dependencies for depender; + # does not search for matching dependee statements) + + def _update_deps(stmt): + # pass stmt to f because might need info + new_deps = f(stmt.dependencies, stmt) + return stmt.copy(dependencies=new_deps) + + return map_instructions(kernel, stmt_match, _update_deps) + +# }}} + +# }}} + + # {{{ remove_instructions def _toposort_of_subset_of_insns(kernel, subset_insns): @@ -263,13 +388,24 @@ def remove_instructions(kernel, insn_ids): assert (new_deps & insn_ids) == frozenset() + # {{{ Remove any new-world stmt inst dependencies on removed stmts + + new_dependencies = insn.dependencies + for removed_id in insn_ids: + # TODO propagate these intelligently? + new_dependencies.pop(removed_id, None) + + # }}} + # update no_sync_with new_no_sync_with = frozenset((insn_id, scope) for insn_id, scope in insn.no_sync_with if insn_id not in insn_ids) - new_insns.append( - insn.copy(depends_on=new_deps, no_sync_with=new_no_sync_with)) + new_insns.append(insn.copy( + depends_on=new_deps, + dependencies=new_dependencies, + no_sync_with=new_no_sync_with)) return kernel.copy( instructions=new_insns) diff --git a/loopy/transform/parameter.py b/loopy/transform/parameter.py index 4916dd4e7..ab9cb9eb2 100644 --- a/loopy/transform/parameter.py +++ b/loopy/transform/parameter.py @@ -91,6 +91,34 @@ def process_set(s): new_domains = [process_set(dom) for dom in kernel.domains] + # {{{ Fix parameter in deps + + from loopy.transform.instruction import map_dependency_maps + from loopy.schedule.checker.utils import convert_map_to_set + + def _fix_parameter_in_dep(dep): + # For efficiency: could check for param presence first + dim_type = isl.dim_type + + # Temporarily convert map to set for processing + set_from_map, n_in_dims, n_out_dims = convert_map_to_set(dep) + + # Fix param + set_from_map = process_set(set_from_map) + + # Now set dims look like [inames' ..., inames ...] + # Convert set back to map + map_from_set = isl.Map.from_domain(set_from_map) + # Move original out dims back + map_from_set = map_from_set.move_dims( + dim_type.out, 0, dim_type.in_, n_in_dims, n_out_dims) + + return map_from_set + + kernel = map_dependency_maps(kernel, _fix_parameter_in_dep) + + # }}} + from pymbolic.mapper.substitutor import make_subst_func subst_func = make_subst_func({name: value}) diff --git a/loopy/transform/precompute.py b/loopy/transform/precompute.py index 730f21542..ad61c00c7 100644 --- a/loopy/transform/precompute.py +++ b/loopy/transform/precompute.py @@ -156,6 +156,9 @@ def __init__(self, rule_mapping_context, subst_name, subst_tag, within, self.compute_read_variables = compute_read_variables self.compute_insn_depends_on = set() + # TODO determine whether there's a better strategy for this + self.things_replaced = set() + def map_substitution(self, name, tag, arguments, expn_state): if not ( name == self.subst_name @@ -236,6 +239,9 @@ def map_kernel(self, kernel): insn.depends_on | frozenset([self.compute_dep_id]))) + if hasattr(insn, "id"): + self.things_replaced.add(insn.id) + for dep in insn.depends_on: if dep in excluded_insn_ids: continue @@ -1058,6 +1064,121 @@ def add_assumptions(d): from loopy.kernel.tools import assign_automatic_axes kernel = assign_automatic_axes(kernel, callables_table) + # {{{ update dependencies + # FIXME Handle deps in precompute + """ + # Get some values that will be useful later + fetch_stmt_id = compute_insn_id + fetch_stmt = kernel.id_to_insn[compute_insn_id] + fetch_inames = fetch_stmt.within_inames + + # Go through all stmts that now use the fetch stuff + for usage_stmt_id in invr.things_replaced: + from loopy.schedule.checker.utils import ( + make_dep_map, + append_mark_to_strings, + remove_dims_by_name, + add_and_name_isl_dims, + insert_and_name_isl_dims, + ) + from loopy.schedule.checker.schedule import ( + BEFORE_MARK, + STATEMENT_VAR_NAME, + ) + # Get some values that will be useful later + usage_stmt = kernel.id_to_insn[usage_stmt_id] + usage_inames = usage_stmt.within_inames + shared_inames = fetch_inames & usage_inames + # TODO understand why this isn't true: + # assert shared_inames == usage_stmt.within_inames - set(sweep_inames) + fetch_inames_not_shared = fetch_inames - shared_inames + + # {{{ create dep fetch_stmt->usage_stmt : SAME(shared_inames) + + dep_in_names = list(fetch_inames) # want a copy anyway + dep_in_names_marked = append_mark_to_strings(dep_in_names, BEFORE_MARK) + dep_out_names = usage_inames + + in_space_str = ", ".join(dep_in_names_marked) + out_space_str = ", ".join(dep_out_names) + constraint_str = " and ".join([ + "{0}{1} = {0}".format(iname, BEFORE_MARK) for iname in shared_inames]) + dep_usage_on_fetch = make_dep_map( + f"{{ [{in_space_str}] -> [{out_space_str}] : {constraint_str} }}", + knl_with_domains=kernel) + # (add this dep below after next step) + + # }}} + + from islpy import dim_type + for dependee_id, old_deps in usage_stmt.dependencies.items(): + for old_dep in old_deps: + # old dep: dependee->usage_stmt + # {{{ create dep dependee->fetch_stmt + + new_dep = old_dep.copy() + + old_out_inames = old_dep.get_var_names(dim_type.out) + assert ( + set(old_out_inames) - set([STATEMENT_VAR_NAME, ]) == + set(usage_inames)) + + non_shared_inames = set(usage_inames) - shared_inames + # Remove inames from old out dims that won't appear in new out dims + for non_shared_iname in non_shared_inames: + new_dep = remove_dims_by_name( + new_dep, dim_type.out, [non_shared_iname]) + + # These new out inames will take on full domain values + assert ( + (set(usage_inames) - non_shared_inames) | fetch_inames_not_shared + == fetch_inames) + + # Add new_unconstrained_out_names to out dims + new_dep = add_and_name_isl_dims( + new_dep, dim_type.out, fetch_inames_not_shared) + + # Intersect dom for fetch_inames_not_shared + dom_to_intersect = kernel.get_inames_domain( + fetch_inames_not_shared + ).project_out_except(fetch_inames_not_shared, [dim_type.set]) + + dom_to_intersect_aligned = isl.align_spaces( + dom_to_intersect, new_dep.range(), + obj_bigger_ok=True) # e.g., params might exist? + + new_dep = new_dep.intersect_range(dom_to_intersect_aligned) + + # {{{ Old dep might have been self-dep, set stmt var correctly + + # add and remove stmt dim + new_dep = remove_dims_by_name( + new_dep, dim_type.out, [STATEMENT_VAR_NAME]) + new_dep = insert_and_name_isl_dims( + new_dep, dim_type.out, [STATEMENT_VAR_NAME], 0) + # set stmt dim value + sid_out = 0 if fetch_stmt_id == dependee_id else 1 + new_dep = new_dep.add_constraint( + isl.Constraint.eq_from_names( + new_dep.space, + {1: sid_out, STATEMENT_VAR_NAME: -1})) + # }}} + + # Add this dep: dependee->fetch : dep + kernel = lp.add_dependency( + kernel, "id:%s" % (fetch_stmt_id), + ("id:%s" % (dependee_id), new_dep)) + + # }}} + + # Add other new dep from above: fetch->usage + kernel = lp.add_dependency( + kernel, "id:%s" % (usage_stmt_id), + ("id:%s" % (fetch_stmt_id), dep_usage_on_fetch)) + + """ + # }}} + return kernel diff --git a/loopy/transform/subst.py b/loopy/transform/subst.py index fd6d93f09..adfec78be 100644 --- a/loopy/transform/subst.py +++ b/loopy/transform/subst.py @@ -234,9 +234,9 @@ def __init__(self, rule_mapping_context, lhs_name, definition_insn_ids, self.definition_insn_id_to_subst_name = {} - self.saw_unmatched_usage_sites = {} + self.unmatched_usage_sites_found = {} for def_id in self.definition_insn_ids: - self.saw_unmatched_usage_sites[def_id] = False + self.unmatched_usage_sites_found[def_id] = set() def get_subst_name(self, def_insn_id): try: @@ -278,7 +278,7 @@ def transform_access(self, index, expn_state): expn_state.kernel, expn_state.instruction, expn_state.stack): - self.saw_unmatched_usage_sites[my_def_id] = True + self.unmatched_usage_sites_found[my_def_id].add(my_insn_id) return None subst_name = self.get_subst_name(my_def_id) @@ -362,6 +362,7 @@ def get_relevant_definition_insn_id(usage_insn_id): return def_id usage_to_definition = {} + definition_to_usage_ids = {} for insn in dep_kernel.instructions: if lhs_name not in insn.read_dependency_names(): @@ -374,11 +375,26 @@ def get_relevant_definition_insn_id(usage_insn_id): % (lhs_name, insn.id)) usage_to_definition[insn.id] = def_id + definition_to_usage_ids.setdefault(def_id, set()).add(insn.id) + # these insns may be removed so can't get within_inames later + definition_to_within_inames = {} + for def_id in definition_to_usage_ids.keys(): + definition_to_within_inames[def_id] = kernel.id_to_insn[def_id].within_inames + + # Get deps for subst_def statements before any of them get removed + definition_id_to_deps = {} + from copy import deepcopy definition_insn_ids = set() for insn in kernel.instructions: if lhs_name in insn.write_dependency_names(): definition_insn_ids.add(insn.id) + definition_id_to_deps[insn.id] = deepcopy(insn.dependencies) + + # usage_to_definition maps each usage to the most recent assignment to the var, + # (most recent "definition"), + # so set(usage_to_definition.values()) is a subset of definition_insn_ids, + # which contains ALL the insns where the var is assigned # }}} @@ -443,7 +459,7 @@ def get_relevant_definition_insn_id(usage_insn_id): new_args = kernel.args if lhs_name in kernel.temporary_variables: - if not any(tts.saw_unmatched_usage_sites.values()): + if not any(tts.unmatched_usage_sites_found.values()): # All usage sites matched--they're now substitution rules. # We can get rid of the variable. @@ -451,7 +467,7 @@ def get_relevant_definition_insn_id(usage_insn_id): del new_temp_vars[lhs_name] if lhs_name in kernel.arg_dict and not force_retain_argument: - if not any(tts.saw_unmatched_usage_sites.values()): + if not any(tts.unmatched_usage_sites_found.values()): # All usage sites matched--they're now substitution rules. # We can get rid of the argument @@ -464,13 +480,90 @@ def get_relevant_definition_insn_id(usage_insn_id): # }}} import loopy as lp + # Remove defs if the subst expression is not still used anywhere kernel = lp.remove_instructions( kernel, { insn_id - for insn_id, still_used in tts.saw_unmatched_usage_sites.items() + for insn_id, still_used in tts.unmatched_usage_sites_found.items() if not still_used}) + # {{{ update dependencies + + from loopy.transform.instruction import map_stmt_dependencies + + # Add dependencies from each subst_def to any statement where its + # LHS was found and the subst was performed + for subst_def_id, subst_usage_ids in definition_to_usage_ids.items(): + + unmatched_usage_ids = tts.unmatched_usage_sites_found[subst_def_id] + matched_usage_ids = subst_usage_ids - unmatched_usage_ids + if matched_usage_ids: + import islpy as isl + dim_type = isl.dim_type + # Create match condition string: + match_any_matched_usage_id = " or ".join( + ["id:%s" % (usage_id) for usage_id in matched_usage_ids]) + + subst_def_deps_dict = definition_id_to_deps[subst_def_id] + old_dep_out_inames = definition_to_within_inames[subst_def_id] + + def _add_deps_to_stmt(old_dep_dict, stmt): + # old_dep_dict: prev dep dict for this stmt + + # want to add old dep from def stmt to usage stmt, + # but if inames of def stmt don't match inames of usage stmt, + # need to get rid of unwanted inames in old dep out dims and add + # any missing inames (inames from usage stmt not present in def stmt) + new_dep_out_inames = stmt.within_inames + out_inames_to_project_out = old_dep_out_inames - new_dep_out_inames + out_inames_to_add = new_dep_out_inames - old_dep_out_inames + # inames_domain for new inames to add + dom_for_new_inames = kernel.get_inames_domain( + out_inames_to_add + ).project_out_except(out_inames_to_add, [dim_type.set]) + + # process and add the old deps + for depends_on_id, old_dep_list in subst_def_deps_dict.items(): + # pu.db + + new_dep_list = [] + for old_dep in old_dep_list: + # TODO figure out when copies are necessary + new_dep = deepcopy(old_dep) + + # project out inames from old dep (out dim) that don't apply + # to this statement + for old_iname in out_inames_to_project_out: + idx_of_old_iname = old_dep.find_dim_by_name( + dim_type.out, old_iname) + assert idx_of_old_iname != -1 + new_dep = new_dep.project_out( + dim_type.out, idx_of_old_iname, 1) + + # add inames from this stmt that were not present in old dep + from loopy.schedule.checker.utils import ( + add_and_name_isl_dims, + ) + new_dep = add_and_name_isl_dims( + new_dep, dim_type.out, out_inames_to_add) + + # add inames domain for new inames + dom_aligned = isl.align_spaces( + dom_for_new_inames, new_dep.range()) + + # Intersect domain with dep + new_dep = new_dep.intersect_range(dom_aligned) + new_dep_list.append(new_dep) + + old_dep_dict.setdefault(depends_on_id, []).extend(new_dep_list) + return old_dep_dict + + kernel = map_stmt_dependencies( + kernel, match_any_matched_usage_id, _add_deps_to_stmt) + + # }}} + return kernel.copy( substitutions=new_substs, temporary_variables=new_temp_vars, diff --git a/test/test_linearization_checker.py b/test/test_linearization_checker.py index 4258a2c52..ded53bd98 100644 --- a/test/test_linearization_checker.py +++ b/test/test_linearization_checker.py @@ -51,6 +51,7 @@ from loopy.schedule.checker import ( get_pairwise_statement_orderings, ) +dim_type = isl.dim_type logger = logging.getLogger(__name__) @@ -93,12 +94,11 @@ def _isl_map_with_marked_dims(s, placeholder_mark="'"): from loopy.schedule.checker.utils import ( append_mark_to_isl_map_var_names, ) - dt = isl.dim_type if BEFORE_MARK == "'": # ISL will ignore the apostrophe; manually name the in_ vars return append_mark_to_isl_map_var_names( isl.Map(s.replace(placeholder_mark, BEFORE_MARK)), - dt.in_, + dim_type.in_, BEFORE_MARK) else: return isl.Map(s.replace(placeholder_mark, BEFORE_MARK)) @@ -143,13 +143,59 @@ def _check_orderings_for_stmt_pair( _align_and_compare_maps(maps_to_compare) -def _process_and_linearize(knl, knl_name="loopy_kernel"): +def _process_and_linearize(prog, knl_name="loopy_kernel"): # Return linearization items along with the preprocessed kernel and # linearized kernel - proc_knl = preprocess_kernel(knl) - lin_knl = get_one_linearized_kernel( - proc_knl[knl_name], proc_knl.callables_table) - return lin_knl.linearization, proc_knl[knl_name], lin_knl + proc_prog = preprocess_kernel(prog) + lin_prog = get_one_linearized_kernel( + proc_prog[knl_name], proc_prog.callables_table) + return lin_prog.linearization, proc_prog[knl_name], lin_prog + +# }}} + + +# {{{ Helper functions for dependency tests + + +def _compare_dependencies( + prog, deps_expected, return_unsatisfied=False, knl_name="loopy_kernel"): + + deps_found = {} + for stmt in prog[knl_name].instructions: + if hasattr(stmt, "dependencies") and stmt.dependencies: + deps_found[stmt.id] = stmt.dependencies + + assert deps_found.keys() == deps_expected.keys() + + for stmt_id_after, dep_dict_found in deps_found.items(): + + dep_dict_expected = deps_expected[stmt_id_after] + + # Ensure deps for stmt_id_after match + assert dep_dict_found.keys() == dep_dict_expected.keys() + + for stmt_id_before, dep_list_found in dep_dict_found.items(): + + # Ensure deps from (stmt_id_before -> stmt_id_after) match + dep_list_expected = dep_dict_expected[stmt_id_before] + print("comparing deps %s->%s" % (stmt_id_before, stmt_id_after)) + assert len(dep_list_found) == len(dep_list_expected) + _align_and_compare_maps(zip(dep_list_found, dep_list_expected)) + + if not return_unsatisfied: + return + + # Get unsatisfied deps + lin_items, proc_prog, lin_prog = _process_and_linearize(prog, knl_name) + unsatisfied_deps = lp.find_unsatisfied_dependencies( + proc_prog, lin_items, stop_on_first_violation=False) + + # Make sure dep checking also works with just linearized kernel + unsatisfied_deps_2 = lp.find_unsatisfied_dependencies( + lin_prog, stop_on_first_violation=False) + assert len(unsatisfied_deps) == len(unsatisfied_deps_2) + + return unsatisfied_deps # }}} @@ -1778,6 +1824,8 @@ def test_blex_map_transitivity_with_duplicate_conc_inames(): # {{{ Dependency tests +# {{{ Dependency creation and checking (without transformations) + # {{{ test_add_dependency_with_new_deps def test_add_dependency_with_new_deps(): @@ -1813,20 +1861,18 @@ def test_add_dependency_with_new_deps(): knl_with_domains=knl["loopy_kernel"]) knl = lp.add_dependency(knl, "id:stmt_b", ("id:stmt_a", dep_b_on_a)) - # Make sure knl instructions all have the expected deps - for stmt in knl["loopy_kernel"].instructions: - if stmt.id == "stmt_b": - assert stmt.dependencies == { - "stmt_a": [dep_b_on_a, ], - } - else: - assert not stmt.dependencies + # Compare deps to expected deps (but don't check for satisfaction yet) + _compare_dependencies( + knl, + {"stmt_b": { + "stmt_a": [dep_b_on_a, ]}}) # {{{ Test make_dep_map while we're here dep_b_on_a_test = _isl_map_with_marked_dims( "[pi] -> {{ [{3}'=0, i'] -> [{3}=1, i] : i > i' " - "and {0} and {1} and {2} }}".format( + "and {0} and {1} and {2} " + "}}".format( i_range_str, i_range_str_p, assumptions_str, @@ -1844,14 +1890,11 @@ def test_add_dependency_with_new_deps(): knl_with_domains=knl["loopy_kernel"]) knl = lp.add_dependency(knl, "id:stmt_b", ("id:stmt_a", dep_b_on_a_2)) - # Make sure knl instructions all have the expected deps - for stmt in knl["loopy_kernel"].instructions: - if stmt.id == "stmt_b": - assert stmt.dependencies == { - "stmt_a": [dep_b_on_a, dep_b_on_a_2], - } - else: - assert not stmt.dependencies + # Compare deps to expected deps (but don't check for satisfaction yet) + _compare_dependencies( + knl, + {"stmt_b": { + "stmt_a": [dep_b_on_a, dep_b_on_a_2]}}) # {{{ Test make_dep_map while we're here @@ -1884,37 +1927,19 @@ def test_add_dependency_with_new_deps(): knl = lp.add_dependency(knl, "id:stmt_c", ("id:stmt_a", dep_c_on_a)) knl = lp.add_dependency(knl, "id:stmt_c", ("id:stmt_b", dep_c_on_b)) - # Make sure knl instructions all have the expected deps - for stmt in knl["loopy_kernel"].instructions: - if stmt.id == "stmt_b": - assert stmt.dependencies == { - "stmt_a": [dep_b_on_a, dep_b_on_a_2], - } - elif stmt.id == "stmt_c": - assert stmt.dependencies == { - "stmt_a": [dep_c_on_a, ], - "stmt_b": [dep_c_on_b, ], - } - else: - assert not stmt.dependencies - - # {{{ Now make sure deps can be checked. These should be satisfied. - - lin_items, proc_knl, lin_knl = _process_and_linearize(knl) - - unsatisfied_deps = lp.find_unsatisfied_dependencies( - proc_knl, lin_items) - - assert not unsatisfied_deps - - # Make sure dep checking also works when only the linearized kernel is - # provided to find_unsatisfied_dependencies() - unsatisfied_deps = lp.find_unsatisfied_dependencies(lin_knl) + # Compare deps and make sure they are satisfied + unsatisfied_deps = _compare_dependencies( + knl, + { + "stmt_b": { + "stmt_a": [dep_b_on_a, dep_b_on_a_2]}, + "stmt_c": { + "stmt_a": [dep_c_on_a, ], "stmt_b": [dep_c_on_b, ]}, + }, + return_unsatisfied=True) assert not unsatisfied_deps - # }}} - # }}} @@ -2009,16 +2034,11 @@ def test_new_dependencies_finite_diff(): # Prioritize loops correctly knl = lp.prioritize_loops(knl, "t,x") - # Make sure deps are satisfied - lin_items, proc_knl, lin_knl = _process_and_linearize(knl) - - unsatisfied_deps = lp.find_unsatisfied_dependencies( - proc_knl, lin_items) - - assert not unsatisfied_deps - - # Make sure dep checking also works with just linearized kernel - unsatisfied_deps = lp.find_unsatisfied_dependencies(lin_knl) + # Compare deps and make sure they are satisfied + unsatisfied_deps = _compare_dependencies( + knl, + {"stmt": {"stmt": [dep, ]}, }, + return_unsatisfied=True) assert not unsatisfied_deps @@ -2030,11 +2050,11 @@ def test_new_dependencies_finite_diff(): knl = ref_knl knl = lp.prioritize_loops(knl, "x,t") - # Make sure unsatisfied deps are caught - lin_items, proc_knl, lin_knl = _process_and_linearize(knl) - - unsatisfied_deps = lp.find_unsatisfied_dependencies( - proc_knl, lin_items, stop_on_first_violation=False) + # Compare deps and make sure unsatisfied deps are caught + unsatisfied_deps = _compare_dependencies( + knl, + {"stmt": {"stmt": [dep, ]}, }, + return_unsatisfied=True) assert len(unsatisfied_deps) == 1 @@ -2078,7 +2098,531 @@ def test_new_dependencies_finite_diff(): knl = lp.add_dtypes( knl, {"u": np.float32, "dx": np.float32, "dt": np.float32}) - # Make sure deps are satisfied + knl = lp.add_dependency(knl, "id:stmt", ("id:stmt", dep)) + + knl = lp.prioritize_loops(knl, "t,x") + knl = lp.tag_inames(knl, "x:l.0") + + # Compare deps and make sure they are satisfied + unsatisfied_deps = _compare_dependencies( + knl, + {"stmt": {"stmt": [dep, ]}, }, + return_unsatisfied=True) + + assert not unsatisfied_deps + + # }}} + +# }}} + +# }}} + + +# {{{ Dependency handling during transformations + +# {{{ test_fix_parameters_with_dependencies + +def test_fix_parameters_with_dependencies(): + knl = lp.make_kernel( + "{[i,j]: 0 <= i < n and 0 <= j < m}", + """ + <>temp0 = 0.1*i+j {id=stmt0} + <>tsq = temp0**2+i+j {id=stmt1,dep=stmt0} + a[i,j] = 23*tsq + 25*tsq+j {id=stmt2,dep=stmt1} + """) + + knl = lp.add_and_infer_dtypes(knl, {"a": np.float32}) + + dep_orig = _isl_map_with_marked_dims( + "[n,m] -> {{ [{0}'=0, i', j']->[{0}=1, i, j] : " + "0 <= i,i' < n and 0 <= j,j' < m " + "and i' = i and j' = j" + "}}".format(STATEMENT_VAR_NAME)) + + from copy import deepcopy + knl = lp.add_dependency(knl, "id:stmt1", ("id:stmt0", deepcopy(dep_orig))) + knl = lp.add_dependency(knl, "id:stmt2", ("id:stmt1", deepcopy(dep_orig))) + + fix_val = 64 + knl = lp.fix_parameters(knl, m=fix_val) + + dep_exp = _isl_map_with_marked_dims( + "[n] -> {{ [{0}'=0, i', j']->[{0}=1, i, j] : " + "0 <= i,i' < n and 0 <= j,j' < {1} " + "and i' = i and j' = j" + "}}".format(STATEMENT_VAR_NAME, fix_val)) + + # Compare deps and make sure they are satisfied + unsatisfied_deps = _compare_dependencies( + knl, + { + "stmt1": {"stmt0": [dep_exp, ]}, + "stmt2": {"stmt1": [dep_exp, ]}, + }, + return_unsatisfied=True) + + assert not unsatisfied_deps + +# }}} + + +# {{{ test_assignment_to_subst_with_dependencies + +def test_assignment_to_subst_with_dependencies(): + knl = lp.make_kernel( + "{[i]: 0 <= i < n}", + """ + <>temp0 = 0.1*i {id=stmt0} + <>tsq = temp0**2 {id=stmt1,dep=stmt0} + a[i] = 23*tsq + 25*tsq {id=stmt2,dep=stmt1} + <>temp3 = 3*tsq {id=stmt3,dep=stmt1} + <>temp4 = 5.5*i {id=stmt4,dep=stmt1} + """) + + # TODO test with multiple subst definition sites + knl = lp.add_and_infer_dtypes(knl, {"a": np.float32}) + + dep_eq = _isl_map_with_marked_dims( + "[n] -> {{ [{0}'=0, i']->[{0}=1, i] : " + "0 <= i,i' < n and i' = i" + "}}".format(STATEMENT_VAR_NAME)) + dep_le = _isl_map_with_marked_dims( + "[n] -> {{ [{0}'=0, i']->[{0}=1, i] : " + "0 <= i,i' < n and i' <= i" + "}}".format(STATEMENT_VAR_NAME)) + + from copy import deepcopy + knl = lp.add_dependency(knl, "id:stmt1", ("id:stmt0", deepcopy(dep_le))) + knl = lp.add_dependency( + knl, "id:stmt2 or id:stmt3 or id:stmt4", + ("id:stmt1", deepcopy(dep_eq))) + + knl = lp.assignment_to_subst(knl, "tsq") + + # Compare deps and make sure they are satisfied + unsatisfied_deps = _compare_dependencies( + knl, + { + "stmt2": {"stmt0": [dep_le, ]}, + "stmt3": {"stmt0": [dep_le, ]}, + }, + return_unsatisfied=True) + # (stmt4 dep was removed because dependee was removed, but dependee's + # deps were not added to stmt4 because the substitution was not made + # in stmt4) TODO this behavior will change when we propagate deps properly + + assert not unsatisfied_deps + + # Test using 'within' -------------------------------------------------- + + knl = lp.make_kernel( + "{[i]: 0 <= i < n}", + """ + <>temp0 = 0.1*i {id=stmt0} + <>tsq = temp0**2 {id=stmt1,dep=stmt0} + a[i] = 23*tsq + 25*tsq {id=stmt2,dep=stmt1} + <>temp3 = 3*tsq {id=stmt3,dep=stmt1} + <>temp4 = 5.5*i {id=stmt4,dep=stmt1} + <>temp5 = 5.6*tsq*i {id=stmt5,dep=stmt1} + """) + + knl = lp.add_and_infer_dtypes(knl, {"a": np.float32}) + + knl = lp.add_dependency(knl, "id:stmt1", ("id:stmt0", deepcopy(dep_le))) + knl = lp.add_dependency( + knl, "id:stmt2 or id:stmt3 or id:stmt4 or id:stmt5", + ("id:stmt1", deepcopy(dep_eq))) + + knl = lp.assignment_to_subst(knl, "tsq", within="id:stmt2 or id:stmt3") + + # Replacement will not be made in stmt5, so stmt1 will not be removed, + # which means no deps will be removed, and the statements where the replacement + # *was* made (stmt2 and stmt3) will still receive the deps from stmt1 + # TODO this behavior may change when we propagate deps properly + + # Compare deps and make sure they are satisfied + unsatisfied_deps = _compare_dependencies( + knl, + { + "stmt1": {"stmt0": [dep_le, ]}, + "stmt2": { + "stmt0": [dep_le, ], "stmt1": [dep_eq, ]}, + "stmt3": { + "stmt0": [dep_le, ], "stmt1": [dep_eq, ]}, + "stmt4": {"stmt1": [dep_eq, ]}, + "stmt5": {"stmt1": [dep_eq, ]}, + }, + return_unsatisfied=True) + + assert not unsatisfied_deps + + # test case where subst def is removed, has deps, and + # inames of subst_def don't match subst usage + + knl = lp.make_kernel( + "{[i,j,k,m]: 0 <= i,j,k,m < n}", + """ + for i,j + <>temp0 = 0.1*i {id=stmt0} + end + for k + <>tsq = temp0**2 {id=stmt1,dep=stmt0} + end + for m + <>res = 23*tsq + 25*tsq {id=stmt2,dep=stmt1} + end + """) + knl = lp.add_and_infer_dtypes(knl, {"temp0,tsq,res": np.float32}) + + dep_1_on_0 = make_dep_map( + "[n] -> { [i', j']->[k] : 0 <= i',j',k < n }", self_dep=False) + dep_2_on_1 = make_dep_map( + "[n] -> { [k']->[m] : 0 <= k',m < n }", self_dep=False) + + from copy import deepcopy + knl = lp.add_dependency(knl, "id:stmt1", ("id:stmt0", deepcopy(dep_1_on_0))) + knl = lp.add_dependency(knl, "id:stmt2", ("id:stmt1", deepcopy(dep_2_on_1))) + + knl = lp.assignment_to_subst(knl, "tsq") + + dep_exp = make_dep_map( + "[n] -> { [i', j']->[m] : 0 <= i',j',m < n }", self_dep=False) + + # Compare deps and make sure they are satisfied + unsatisfied_deps = _compare_dependencies( + knl, + { + "stmt2": {"stmt0": [dep_exp, ]}, + }, + return_unsatisfied=True) + + assert not unsatisfied_deps + +# }}} + + +# {{{ test_duplicate_inames_with_dependencies + +def test_duplicate_inames_with_dependencies(): + + knl = lp.make_kernel( + "{[i,j]: 0 <= i,j < n}", + """ + b[i,j] = a[i,j] {id=stmtb} + c[i,j] = a[i,j] {id=stmtc,dep=stmtb} + """) + knl = lp.add_and_infer_dtypes(knl, {"a": np.float32}) + + dep_eq = _isl_map_with_marked_dims( + "[n] -> {{ [{0}'=0, i', j']->[{0}=1, i, j] : " + "0 <= i,i',j,j' < n and i' = i and j' = j" + "}}".format(STATEMENT_VAR_NAME)) + + # Create dep stmtb->stmtc + knl = lp.add_dependency(knl, "id:stmtc", ("id:stmtb", dep_eq)) + + ref_knl = knl + + # {{{ Duplicate j within stmtc + + knl = lp.duplicate_inames(knl, ["j"], within="id:stmtc", new_inames=["j_new"]) + + dep_exp = _isl_map_with_marked_dims( + "[n] -> {{ [{0}'=0, i', j']->[{0}=1, i, j_new] : " + "0 <= i,i',j_new,j' < n and i' = i and j' = j_new" + "}}".format(STATEMENT_VAR_NAME)) + + # Compare deps and make sure they are satisfied + unsatisfied_deps = _compare_dependencies( + knl, + {"stmtc": {"stmtb": [dep_exp, ]}}, + return_unsatisfied=True) + + assert not unsatisfied_deps + + # }}} + + # {{{ Duplicate j within stmtb + + knl = ref_knl + knl = lp.duplicate_inames(knl, ["j"], within="id:stmtb", new_inames=["j_new"]) + + dep_exp = _isl_map_with_marked_dims( + "[n] -> {{ [{0}'=0, i', j_new']->[{0}=1, i, j] : " + "0 <= i,i',j,j_new' < n and i' = i and j_new' = j" + "}}".format(STATEMENT_VAR_NAME)) + + # Compare deps and make sure they are satisfied + unsatisfied_deps = _compare_dependencies( + knl, + {"stmtc": {"stmtb": [dep_exp, ]}}, + return_unsatisfied=True) + + assert not unsatisfied_deps + + # }}} + + # {{{ Duplicate j within stmtb and stmtc + + knl = ref_knl + knl = lp.duplicate_inames( + knl, ["j"], within="id:stmtb or id:stmtc", new_inames=["j_new"]) + + dep_exp = _isl_map_with_marked_dims( + "[n] -> {{ [{0}'=0, i', j_new']->[{0}=1, i, j_new] : " + "0 <= i,i',j_new,j_new' < n and i' = i and j_new' = j_new" + "}}".format(STATEMENT_VAR_NAME)) + + # Compare deps and make sure they are satisfied + unsatisfied_deps = _compare_dependencies( + knl, + {"stmtc": {"stmtb": [dep_exp, ]}}, + return_unsatisfied=True) + + assert not unsatisfied_deps + + # }}} + +# }}} + + +# {{{ test_rename_inames_with_dependencies + +def test_rename_inames_with_dependencies(): + # When rename_iname is called and the new iname + # *doesn't* already exist, then duplicate_inames is called, + # and we test that elsewhere. Here we test the case where + # rename_iname is called and the new iname already exists. + + knl = lp.make_kernel( + "{[i,j,m,j_new]: 0 <= i,j,m,j_new < n}", + """ + b[i,j] = a[i,j] {id=stmtb} + c[i,j] = a[i,j] {id=stmtc,dep=stmtb} + e[i,j_new] = 1.1 + d[m] = 5.5 {id=stmtd,dep=stmtc} + """) + knl = lp.add_and_infer_dtypes(knl, {"a,d": np.float32}) + + dep_c_on_b = make_dep_map( + "[n] -> { [i', j']->[i, j] : 0 <= i,i',j,j' < n and i' = i and j' = j }", + self_dep=False) + dep_c_on_c = make_dep_map( + "[n] -> { [i', j']->[i, j] : 0 <= i,i',j,j' < n and i' < i and j' < j }", + self_dep=True) + dep_d_on_c = make_dep_map( + "[n] -> { [i', j']->[m] : 0 <= m,i',j' < n }", + self_dep=False) + + # Create dep stmtb->stmtc + knl = lp.add_dependency(knl, "id:stmtc", ("id:stmtb", dep_c_on_b)) + knl = lp.add_dependency(knl, "id:stmtc", ("id:stmtc", dep_c_on_c)) + knl = lp.add_dependency(knl, "id:stmtd", ("id:stmtc", dep_d_on_c)) + + # Rename j within stmtc + + knl = lp.rename_iname( + knl, "j", "j_new", within="id:stmtc", existing_ok=True) + + dep_c_on_b_exp = make_dep_map( + "[n] -> { [i', j']->[i, j_new] : " + "0 <= i,i',j_new,j' < n and i' = i and j' = j_new}", + self_dep=False) + dep_c_on_c_exp = make_dep_map( + "[n] -> { [i', j_new']->[i, j_new] : " + "0 <= i,i',j_new,j_new' < n and i' < i and j_new' < j_new }", + self_dep=True) + dep_d_on_c_exp = make_dep_map( + "[n] -> { [i', j_new']->[m] : 0 <= m,i',j_new' < n }", + self_dep=False) + + # Compare deps and make sure they are satisfied + unsatisfied_deps = _compare_dependencies( + knl, + { + "stmtc": {"stmtb": [dep_c_on_b_exp, ], "stmtc": [dep_c_on_c_exp, ]}, + "stmtd": {"stmtc": [dep_d_on_c_exp, ]}, + }, + return_unsatisfied=True) + + assert not unsatisfied_deps + +# }}} + + +# {{{ test_split_iname_with_dependencies + +def test_split_iname_with_dependencies(): + knl = lp.make_kernel( + "{[i]: 0<=i { %s : 0 <= i < p and i' = i }" + % (dep_inout_space_str)) + + knl = lp.add_dependency(knl, "id:stmt1", ("id:stmt0", dep_satisfied)) + knl = lp.split_iname(knl, "i", 32) + + dep_exp = _isl_map_with_marked_dims( + "[p] -> {{ [{0}'=0, i_outer', i_inner'] -> [{0}=1, i_outer, i_inner] : " + "0 <= i_inner, i_inner' < 32" # new bounds + " and 0 <= 32*i_outer + i_inner < p" # transformed bounds (0 <= i < p) + " and 0 <= 32*i_outer' + i_inner' < p" # transformed bounds (0 <= i' < p) + " and i_inner + 32*i_outer = 32*i_outer' + i_inner'" # i = i' + "}}".format(STATEMENT_VAR_NAME)) + + # Compare deps and make sure they are satisfied + unsatisfied_deps = _compare_dependencies( + knl, + {"stmt1": {"stmt0": [dep_exp, ]}}, + return_unsatisfied=True) + + assert not unsatisfied_deps + + # }}} + + # {{{ Split iname within stmt1 and make sure dep is correct + + knl = deepcopy(ref_knl) + + knl = lp.add_dependency(knl, "id:stmt1", ("id:stmt0", dep_satisfied)) + knl = lp.split_iname(knl, "i", 32, within="id:stmt1") + + dep_exp = _isl_map_with_marked_dims( + "[p] -> {{ [{0}'=0, i'] -> [{0}=1, i_outer, i_inner] : " + "0 <= i_inner < 32" # new bounds + " and 0 <= 32*i_outer + i_inner < p" # transformed bounds (0 <= i < p) + " and 0 <= i' < p" # original bounds + " and i_inner + 32*i_outer = i'" # transform {i = i'} + "}}".format(STATEMENT_VAR_NAME)) + + # Compare deps and make sure they are satisfied + unsatisfied_deps = _compare_dependencies( + knl, + {"stmt1": {"stmt0": [dep_exp, ]}}, + return_unsatisfied=True) + + assert not unsatisfied_deps + + # }}} + + # {{{ Split iname within stmt0 and make sure dep is correct + + knl = deepcopy(ref_knl) + + knl = lp.add_dependency(knl, "id:stmt1", ("id:stmt0", dep_satisfied)) + knl = lp.split_iname(knl, "i", 32, within="id:stmt0") + + dep_exp = _isl_map_with_marked_dims( + "[p] -> {{ [{0}'=0, i_outer', i_inner'] -> [{0}=1, i] : " + "0 <= i_inner' < 32" # new bounds + " and 0 <= i < p" # original bounds + " and 0 <= 32*i_outer' + i_inner' < p" # transformed bounds (0 <= i' < p) + " and i = 32*i_outer' + i_inner'" # transform {i = i'} + "}}".format(STATEMENT_VAR_NAME)) + + # Compare deps and make sure they are satisfied + unsatisfied_deps = _compare_dependencies( + knl, + {"stmt1": {"stmt0": [dep_exp, ]}}, + return_unsatisfied=True) + + assert not unsatisfied_deps + + # }}} + + # {{{ Check dep that should not be satisfied + + knl = deepcopy(ref_knl) + + dep_unsatisfied = _isl_map_with_marked_dims( + "[p] -> { %s : 0 <= i < p and i' = i + 1 }" + % (dep_inout_space_str)) + + knl = lp.add_dependency(knl, "id:stmt1", ("id:stmt0", dep_unsatisfied)) + knl = lp.split_iname(knl, "i", 32) + + dep_exp = _isl_map_with_marked_dims( + "[p] -> {{ [{0}'=0, i_outer', i_inner'] -> [{0}=1, i_outer, i_inner] : " + "0 <= i_inner, i_inner' < 32" # new bounds + " and 0 <= 32*i_outer + i_inner < p" # transformed bounds (0 <= i < p) + " and 0 <= 32*i_outer' + i_inner' - 1 < p" # trans. bounds (0 <= i'-1 < p) + " and i_inner + 32*i_outer + 1 = 32*i_outer' + i_inner'" # i' = i + 1 + "}}".format(STATEMENT_VAR_NAME)) + + # Compare deps and make sure they are satisfied + unsatisfied_deps = _compare_dependencies( + knl, + {"stmt1": {"stmt0": [dep_exp, ]}}, + return_unsatisfied=True) + + assert len(unsatisfied_deps) == 1 + + # }}} + + # {{{ Deps that should be satisfied after gratuitous splitting + + knl = lp.make_kernel( + "{[i,j,k,m]: 0<=i,j,k,m { %s : %s and i' = i and k' = k}" + % (dep_ik_space_str, ik_bounds_str)) + dep_stmt1_on_stmt0_lt = _isl_map_with_marked_dims( + "[p] -> { %s : %s and i' < i and k' < k}" + % (dep_ik_space_str, ik_bounds_str)) + dep_stmt3_on_stmt2_eq = _isl_map_with_marked_dims( + "[p] -> { %s : %s and i' = i and k' = k and j' = j and m' = m}" + % (dep_ijkm_space_str, ijkm_bounds_str)) + + knl = lp.add_dependency(knl, "id:stmt1", ("id:stmt0", dep_stmt1_on_stmt0_eq)) + knl = lp.add_dependency(knl, "id:stmt1", ("id:stmt0", dep_stmt1_on_stmt0_lt)) + knl = lp.add_dependency(knl, "id:stmt3", ("id:stmt2", dep_stmt3_on_stmt2_eq)) + + # Gratuitous splitting + knl = lp.split_iname(knl, "i", 64) + knl = lp.split_iname(knl, "j", 64) + knl = lp.split_iname(knl, "k", 64) + knl = lp.split_iname(knl, "m", 64) + knl = lp.split_iname(knl, "i_inner", 8) + knl = lp.split_iname(knl, "j_inner", 8) + knl = lp.split_iname(knl, "k_inner", 8) + knl = lp.split_iname(knl, "m_inner", 8) + knl = lp.split_iname(knl, "i_outer", 4) + knl = lp.split_iname(knl, "j_outer", 4) + knl = lp.split_iname(knl, "k_outer", 4) + knl = lp.split_iname(knl, "m_outer", 4) + + # Get a linearization lin_items, proc_knl, lin_knl = _process_and_linearize(knl) unsatisfied_deps = lp.find_unsatisfied_dependencies( @@ -2090,6 +2634,389 @@ def test_new_dependencies_finite_diff(): # }}} + +# {{{ test map domain with dependencies + +# {{{ test_map_domain_with_only_partial_dep_pair_affected + +def test_map_domain_with_only_partial_dep_pair_affected(): + + # Split an iname using map_domain, and have (misaligned) deps + # where only the dependee uses the split iname + + # {{{ Make kernel + + knl = lp.make_kernel( + [ + "[nx,nt] -> {[x, t]: 0 <= x < nx and 0 <= t < nt}", + "[ni] -> {[i]: 0 <= i < ni}", + ], + """ + a[x,t] = b[x,t] {id=stmta} + c[x,t] = d[x,t] {id=stmtc,dep=stmta} + e[i] = f[i] {id=stmte,dep=stmtc} + """, + lang_version=(2018, 2), + ) + knl = lp.add_and_infer_dtypes(knl, {"b,d,f": np.float32}) + + # }}} + + # {{{ Add dependencies + + dep_c_on_a = _isl_map_with_marked_dims( + "[nx, nt] -> {{" + "[{0}' = 0, x', t'] -> [{0} = 1, x, t] : " + "0 <= x,x' < nx and 0 <= t,t' < nt and " + "t' <= t and x' <= x" + "}}".format(STATEMENT_VAR_NAME)) + + knl = lp.add_dependency( + knl, "id:stmtc", ("id:stmta", dep_c_on_a)) + + # Intentionally make order of x and t different from transform_map below + # to test alignment steps in map_domain + dep_e_on_c = _isl_map_with_marked_dims( + "[nx, nt, ni] -> {{" + "[{0}' = 0, t', x'] -> [{0} = 1, i] : " + "0 <= x' < nx and 0 <= t' < nt and 0 <= i < ni" + "}}".format(STATEMENT_VAR_NAME)) + + knl = lp.add_dependency( + knl, "id:stmte", ("id:stmtc", dep_e_on_c)) + + # }}} + + # {{{ Apply domain change mapping + + # Create map_domain mapping: + import islpy as isl + transform_map = isl.BasicMap( + "[nt] -> {[t] -> [t_outer, t_inner]: " + "0 <= t_inner < 32 and " + "32*t_outer + t_inner = t and " + "0 <= 32*t_outer + t_inner < nt}") + + # Call map_domain to transform kernel + knl = lp.map_domain(knl, transform_map) + + # Prioritize loops (prio should eventually be updated in map_domain?) + knl = lp.prioritize_loops(knl, "x, t_outer, t_inner") + + # }}} + + # {{{ Create expected dependencies + + dep_c_on_a_exp = _isl_map_with_marked_dims( + "[nx, nt] -> {{" + "[{0}' = 0, x', t_outer', t_inner'] -> [{0} = 1, x, t_outer, t_inner] : " + "0 <= x,x' < nx and " # old bounds + "0 <= t_inner,t_inner' < 32 and " # new bounds + "0 <= 32*t_outer + t_inner < nt and " # new bounds + "0 <= 32*t_outer' + t_inner' < nt and " # new bounds + "32*t_outer' + t_inner' <= 32*t_outer + t_inner and " # new constraint t'<=t + "x' <= x" # old constraint + "}}".format(STATEMENT_VAR_NAME)) + + dep_e_on_c_exp = _isl_map_with_marked_dims( + "[nx, nt, ni] -> {{" + "[{0}' = 0, x', t_outer', t_inner'] -> [{0} = 1, i] : " + "0 <= x' < nx and 0 <= i < ni and " # old bounds + "0 <= t_inner' < 32 and " # new bounds + "0 <= 32*t_outer' + t_inner' < nt" # new bounds + "}}".format(STATEMENT_VAR_NAME)) + + # }}} + + # {{{ Make sure deps are correct and satisfied + + # Compare deps and make sure they are satisfied + unsatisfied_deps = _compare_dependencies( + knl, + { + "stmtc": { + "stmta": [dep_c_on_a_exp, ]}, + "stmte": { + "stmtc": [dep_e_on_c_exp, ]}, + }, + return_unsatisfied=True) + + assert not unsatisfied_deps + + # }}} + +# }}} + + +# {{{ test_map_domain_with_inames_missing_in_transform_map + +def test_map_domain_with_inames_missing_in_transform_map(): + + # Make sure map_domain updates deps correctly when the mapping doesn't + # include all the dims in the domain. + + # {{{ Make kernel + + knl = lp.make_kernel( + "[nx,nt] -> {[x, y, z, t]: 0 <= x,y,z < nx and 0 <= t < nt}", + """ + a[y,x,t,z] = b[y,x,t,z] {id=stmta} + """, + lang_version=(2018, 2), + ) + knl = lp.add_and_infer_dtypes(knl, {"b": np.float32}) + + # }}} + + # {{{ Create dependency + + dep = _isl_map_with_marked_dims( + "[nx, nt] -> {{" + "[{0}' = 0, x', y', z', t'] -> [{0} = 0, x, y, z, t] : " + "0 <= x,y,z,x',y',z' < nx and 0 <= t,t' < nt and " + "t' < t and x' < x and y' < y and z' < z" + "}}".format(STATEMENT_VAR_NAME)) + + knl = lp.add_dependency(knl, "id:stmta", ("id:stmta", dep)) + + # }}} + + # {{{ Apply domain change mapping + + # Create map_domain mapping that only includes t and y + # (x and z should be unaffected) + transform_map = isl.BasicMap( + "[nx,nt] -> {[t, y] -> [t_outer, t_inner, y_new]: " + "0 <= t_inner < 32 and " + "32*t_outer + t_inner = t and " + "0 <= 32*t_outer + t_inner < nt and " + "y = y_new" + "}") + + # Call map_domain to transform kernel + knl = lp.map_domain(knl, transform_map) + + # }}} + + # {{{ Create expected dependency after transformation + + dep_exp = _isl_map_with_marked_dims( + "[nx, nt] -> {{" + "[{0}' = 0, x', y_new', z', t_outer', t_inner'] -> " + "[{0} = 0, x, y_new, z, t_outer, t_inner] : " + "0 <= x,z,x',z' < nx " # old bounds + "and 0 <= t_inner,t_inner' < 32 and 0 <= y_new,y_new' < nx " # new bounds + "and 0 <= 32*t_outer + t_inner < nt " # new bounds + "and 0 <= 32*t_outer' + t_inner' < nt " # new bounds + "and x' < x and z' < z " # old constraints + "and y_new' < y_new " # new constraint + "and 32*t_outer' + t_inner' < 32*t_outer + t_inner" # new constraint + "}}".format(STATEMENT_VAR_NAME)) + + # }}} + + # {{{ Make sure deps are correct and satisfied + + # Compare deps and make sure they are satisfied + unsatisfied_deps = _compare_dependencies( + knl, + {"stmta": {"stmta": [dep_exp, ]}}, + return_unsatisfied=True) + + assert not unsatisfied_deps + + # }}} + + +# }}} + + +# {{{ test_map_domain_with_stencil_dependencies + +def test_map_domain_with_stencil_dependencies(): + + # {{{ Make kernel + + knl = lp.make_kernel( + "[nx,nt] -> {[ix, it]: 1<=ix {{" + "[{0}' = 0, ix', it'] -> [{0} = 0, ix, it = 1 + it'] : " + "0 < ix' <= -2 + nx and 0 <= it' <= -2 + nt and ix >= -1 + ix' and " + "0 < ix <= 1 + ix' and ix <= -2 + nx; " + "[statement' = 0, ix', it'] -> [statement = 0, ix = ix', it = 2 + it'] : " + "0 < ix' <= -2 + nx and 0 <= it' <= -3 + nt" + "}}".format(STATEMENT_VAR_NAME)) + + knl = lp.add_dependency( + knl, "id:"+stmt_after, ("id:"+stmt_before, dep_map)) + + # }}} + + # {{{ Check deps *without* map_domain transformation + + ref_knl = knl + + # Prioritize loops + knl = lp.prioritize_loops(knl, ("it", "ix")) # valid + #knl = lp.prioritize_loops(knl, ("ix", "it")) # invalid + + # Compare deps and make sure they are satisfied + unsatisfied_deps = _compare_dependencies( + knl, + {"stmt": {"stmt": [dep_map, ]}}, + return_unsatisfied=True) + + assert not unsatisfied_deps + + # }}} + + # {{{ Check dependency after domain change mapping + + knl = ref_knl # loop priority goes away, deps stay + + # Create map_domain mapping: + transform_map = isl.BasicMap( + "[nx,nt] -> {[ix, it] -> [tx, tt, tparity, itt, itx]: " + "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}") + + # Call map_domain to transform kernel + knl = lp.map_domain(knl, transform_map) + + # Prioritize loops (prio should eventually be updated in map_domain?) + knl = lp.prioritize_loops(knl, "tt,tparity,tx,itt,itx") + + # {{{ Create expected dependency + + # Prep transform map to be applied to dependency + from loopy.schedule.checker.utils import ( + insert_and_name_isl_dims, + add_eq_isl_constraint_from_names, + append_mark_to_isl_map_var_names, + ) + # Insert 'statement' dim into transform map + transform_map = insert_and_name_isl_dims( + transform_map, dim_type.in_, [STATEMENT_VAR_NAME+BEFORE_MARK], 0) + transform_map = insert_and_name_isl_dims( + transform_map, dim_type.out, [STATEMENT_VAR_NAME], 0) + # Add stmt = stmt' constraint + transform_map = add_eq_isl_constraint_from_names( + transform_map, STATEMENT_VAR_NAME, STATEMENT_VAR_NAME+BEFORE_MARK) + + # Apply transform map to dependency + mapped_dep_map = dep_map.apply_range(transform_map).apply_domain(transform_map) + mapped_dep_map = append_mark_to_isl_map_var_names( + mapped_dep_map, dim_type.in_, BEFORE_MARK) + + # }}} + + # Compare deps and make sure they are satisfied + unsatisfied_deps = _compare_dependencies( + knl, + {"stmt": {"stmt": [mapped_dep_map, ]}}, + return_unsatisfied=True) + + assert not unsatisfied_deps + + # }}} + +# }}} + +# }}} + + +# {{{ test_add_prefetch_with_dependencies + +# FIXME handle deps during prefetch + +''' + +def test_add_prefetch_with_dependencies(): + + lp.set_caching_enabled(False) + knl = lp.make_kernel( + "[p] -> { [i,j,k,m] : 0 <= i,j < p and 0 <= k,m < 16}", + """ + for i,j,k,m + a[i+1,j+1,k+1,m+1] = a[i,j,k,m] {id=stmt} + end + """, + assumptions="p >= 1", + lang_version=(2018, 2) + ) + knl = lp.add_and_infer_dtypes(knl, {"a": np.float32}) + + dep_init = make_dep_map( + "{ [i',j',k',m'] -> [i,j,k,m] : " + "i' + 1 = i and j' + 1 = j and k' + 1 = k and m' + 1 = m }", + self_dep=True, knl_with_domains=knl) + knl = lp.add_dependency(knl, "id:stmt", ("id:stmt", dep_init)) + + # Compare deps and make sure they are satisfied + unsatisfied_deps = _compare_dependencies( + knl, + {"stmt": {"stmt": [dep_init, ]}}, + return_unsatisfied=True) + + assert not unsatisfied_deps + + knl = lp.add_prefetch( + knl, "a", sweep_inames=["k", "m"], + fetch_outer_inames=frozenset({"i", "j"}), + # dim_arg_names=["k_fetch", "m_fetch"], # TODO not sure why these don't work + ) + + # create expected deps + dep_stmt_on_fetch_exp = make_dep_map( + "{ [i',j',a_dim_2',a_dim_3'] -> [i,j,k,m] : " + "i' = i and j' = j }", + knl_with_domains=knl) + dep_fetch_on_stmt_exp = make_dep_map( + "{ [i',j',k',m'] -> [i,j,a_dim_2,a_dim_3] : " + "i' + 1 = i and j' + 1 = j " + "and 0 <= k',m' < 15 " + "}", + knl_with_domains=knl) + # (make_dep_map will set k',m' upper bound to 16, so add manually^) + + # Why is this necessary to avoid dependency cycle? + knl.id_to_insn["a_fetch_rule"].depends_on_is_final = True + + # Compare deps and make sure they are satisfied + unsatisfied_deps = _compare_dependencies( + knl, + { + "stmt": {"stmt": [dep_init], "a_fetch_rule": [dep_stmt_on_fetch_exp]}, + "a_fetch_rule": {"stmt": [dep_fetch_on_stmt_exp]}, + }, + return_unsatisfied=True) + + assert not unsatisfied_deps + +''' + +# }}} + +# }}} + # }}} diff --git a/test/test_transform.py b/test/test_transform.py index 3df2d6c2c..3bc5456ba 100644 --- a/test/test_transform.py +++ b/test/test_transform.py @@ -320,6 +320,7 @@ def test_join_inames(ctx_factory): knl = lp.add_prefetch(knl, "a", sweep_inames=["i", "j"], default_tag="l.auto") knl = lp.join_inames(knl, ["a_dim_0", "a_dim_1"]) + # TODO why does this lead to inames key error for 'a_dim_1_and_a_dim_0' ?? lp.auto_test_vs_ref(ref_knl, ctx, knl, print_ref_code=True) @@ -574,7 +575,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_prg = lp.make_kernel( "{[i]: 0<=i<10}", @@ -674,10 +674,13 @@ def test_map_domain_vs_split_iname(ctx_factory): lin_knl_split_iname = lp.get_one_linearized_kernel( proc_knl_split_iname["loopy_kernel"], proc_knl_split_iname.callables_table) + from loopy.schedule.checker.utils import ( + ensure_dim_names_match_and_align, + ) for d_map_domain, d_split_iname in zip( knl_map_dom["loopy_kernel"].domains, knl_split_iname["loopy_kernel"].domains): - d_map_domain_aligned = _ensure_dim_names_match_and_align( + d_map_domain_aligned = ensure_dim_names_match_and_align( d_map_domain, d_split_iname) assert d_map_domain_aligned == d_split_iname @@ -775,10 +778,13 @@ def test_map_domain_transform_map_validity_and_errors(ctx_factory): lin_knl_split_iname = lp.get_one_linearized_kernel( proc_knl_split_iname["loopy_kernel"], proc_knl_split_iname.callables_table) + from loopy.schedule.checker.utils import ( + ensure_dim_names_match_and_align, + ) for d_map_domain, d_split_iname in zip( knl_map_dom["loopy_kernel"].domains, knl_split_iname["loopy_kernel"].domains): - d_map_domain_aligned = _ensure_dim_names_match_and_align( + d_map_domain_aligned = ensure_dim_names_match_and_align( d_map_domain, d_split_iname) assert d_map_domain_aligned == d_split_iname