diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index 409cb2c85..264fde4af 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -14016,6 +14016,24 @@ } } ], + "./loopy/kernel/dependency.py": [ + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 25, + "endColumn": 34, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 12, + "endColumn": 43, + "lineCount": 1 + } + } + ], "./loopy/kernel/function_interface.py": [ { "code": "reportAny", @@ -74225,4 +74243,4 @@ } ] } -} \ No newline at end of file +} diff --git a/doc/ref_internals.rst b/doc/ref_internals.rst index 02edceca6..7d3442bcd 100644 --- a/doc/ref_internals.rst +++ b/doc/ref_internals.rst @@ -60,3 +60,8 @@ Schedule .. automodule:: loopy.schedule .. automodule:: loopy.schedule.tools .. automodule:: loopy.schedule.tree + +Dependence Analysis +------------------- +.. automodule:: loopy.kernel.dependency +.. automodule:: loopy.schedule.verification diff --git a/loopy/codegen/__init__.py b/loopy/codegen/__init__.py index 6cad3620f..027546be6 100644 --- a/loopy/codegen/__init__.py +++ b/loopy/codegen/__init__.py @@ -316,6 +316,13 @@ def generate_code_for_a_single_kernel( raise LoopyError("cannot generate code for a kernel that has not been " "scheduled") + from loopy.kernel.dependency import has_precise_dependencies + if has_precise_dependencies(kernel): + from loopy.schedule.verification import ( + verify_happens_after_is_enforced, + ) + kernel = verify_happens_after_is_enforced(kernel) + codegen_plog = ProcessLogger(logger, f"{kernel.name}: generate code") # {{{ examine arg list diff --git a/loopy/kernel/dependency.py b/loopy/kernel/dependency.py new file mode 100644 index 000000000..3e68cf182 --- /dev/null +++ b/loopy/kernel/dependency.py @@ -0,0 +1,1106 @@ +from __future__ import annotations + + +__copyright__ = """ +Copyright (C) 2026 Addison Alvey-Blanco +""" + +__license__ = """ +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +""" + + +from enum import Enum +from functools import cached_property +from typing import TYPE_CHECKING, final, override + +import namedisl as nisl +from constantdict import constantdict +from namedisl import DimType + +from pymbolic import primitives as prim +from pytools.graph import compute_topological_order + +from loopy import for_each_kernel +from loopy.diagnostic import LoopyError +from loopy.kernel.instruction import ( + CInstruction, + HappensAfter, + InstructionBase, + MultiAssignmentBase, +) +from loopy.symbolic import ( + LinearSubscript, + Reduction, + SubArrayRef, + SubstitutionRuleExpander, + WalkMapper, + aff_from_expr, + get_dependencies, +) + + +if TYPE_CHECKING: + from collections.abc import Collection, Mapping + + from namedisl.core import NamedIslObjectT + + from pymbolic.typing import Expression + from pytools import UniqueNameGenerator + + from loopy.kernel import LoopKernel + + +@final +class AccessType(Enum): + read = 0 + write = 1 + + +class AccessRelationFinder(WalkMapper[[str, AccessType]]): + """Collect per-instruction statement-instance-to-cell access relations.""" + + kernel: LoopKernel + _additional_inames: frozenset[str] + _read_relations: dict[str, dict[str, nisl.Map]] + _write_relations: dict[str, dict[str, nisl.Map]] + _name_generator: UniqueNameGenerator + _cell_names: list[str] + _constant_names: frozenset[str] + _storage_variables: frozenset[str] + _subst_expander: SubstitutionRuleExpander + + def __init__(self, kernel: LoopKernel): + self.kernel = kernel + self._additional_inames = frozenset() + self._read_relations = {stmt.id: {} for stmt in kernel.instructions} + self._write_relations = {stmt.id: {} for stmt in kernel.instructions} + self._name_generator = kernel.get_var_name_generator() + self._cell_names = [] + from loopy.kernel.data import ValueArg + self._constant_names = frozenset( + arg.name for arg in kernel.args + if isinstance(arg, ValueArg) + and arg.name not in kernel.get_written_variables() + ) + self._storage_variables = frozenset(kernel.non_iname_variable_names()) + self._subst_expander = SubstitutionRuleExpander(kernel.substitutions) + + super().__init__() + + @override + def __call__( + self, expr: Expression, stmt_id: str, access_type: AccessType + ) -> None: + self.rec(self._subst_expander(expr), stmt_id, access_type) + + def _get_access_relation( + self, + domain: nisl.Set, + subscript: tuple[Expression, ...], + ) -> nisl.Map: + subscript_dependencies = frozenset( + dependency + for index in subscript + for dependency in get_dependencies(index) + ) + domain = domain.add_dims( + DimType.param, + (subscript_dependencies & self._constant_names) + - domain.space.param_names, + ) + instance_names = domain.space.dimtype_to_names[DimType.out] + while len(self._cell_names) < len(subscript): + axis = len(self._cell_names) + self._cell_names.append(self._name_generator(f"ax_{axis}")) + + cell_names = tuple(self._cell_names[: len(subscript)]) + + access_set = domain.add_dims(DimType.out, cell_names) + coordinates = access_set.var_pw_affs + for cell_name, index_expr in zip(cell_names, subscript, strict=True): + index_aff = aff_from_expr( + access_set.var_affs, + index_expr, + ).as_pw_aff() + + access_set = access_set & coordinates[cell_name].eq_set(index_aff) + + return _set_as_map(access_set, in_names=instance_names) + + def _record_access( + self, + stmt_id: str, + var: str, + subscript: tuple[Expression, ...], + access_type: AccessType, + ) -> None: + if var not in self._storage_variables: + return + + stmt = self.kernel.id_to_insn[stmt_id] + domain_inames = stmt.within_inames | self._additional_inames + inames_domain = self.kernel.get_inames_domain(domain_inames) + access_rel = self._get_access_relation(inames_domain, subscript) + + additional_inames = self._additional_inames - stmt.within_inames + if additional_inames: + access_rel = access_rel.project_out(additional_inames) + + match access_type: + case AccessType.read: + previous = self._read_relations[stmt_id].get(var) + self._read_relations[stmt_id][var] = ( + access_rel if previous is None else previous | access_rel + ) + case AccessType.write: + previous = self._write_relations[stmt_id].get(var) + self._write_relations[stmt_id][var] = ( + access_rel if previous is None else previous | access_rel + ) + case _: + raise ValueError("unknown AccessType") + + @cached_property + def read_relations(self) -> Mapping[str, Mapping[str, nisl.Map]]: + return constantdict({ + stmt_id: constantdict(relations) + for stmt_id, relations in self._read_relations.items() + }) + + @cached_property + def write_relations(self) -> Mapping[str, Mapping[str, nisl.Map]]: + return constantdict({ + stmt_id: constantdict(relations) + for stmt_id, relations in self._write_relations.items() + }) + + @override + def map_variable( + self, expr: prim.Variable, /, stmt_id: str, access_type: AccessType + ) -> None: + self._record_access(stmt_id, expr.name, (), access_type) + + @override + def map_subscript( + self, expr: prim.Subscript, /, stmt_id: str, access_type: AccessType + ) -> None: + assert isinstance(expr.aggregate, prim.Variable) + self._record_access( + stmt_id, expr.aggregate.name, expr.index_tuple, access_type + ) + + @override + def map_linear_subscript( + self, expr: LinearSubscript, /, stmt_id: str, access_type: AccessType + ) -> None: + self.rec(expr.index, stmt_id, AccessType.read) + + assert isinstance(expr.aggregate, prim.Variable) + self._record_access( + stmt_id, expr.aggregate.name, (expr.index,), access_type + ) + + @override + def map_reduction( + self, expr: Reduction, /, stmt_id: str, access_type: AccessType + ) -> None: + previous_inames = self._additional_inames + self._additional_inames |= frozenset(expr.inames) + try: + WalkMapper.map_reduction(self, expr, stmt_id, access_type) + finally: + self._additional_inames = previous_inames + + @override + def map_sub_array_ref( + self, expr: SubArrayRef, /, stmt_id: str, access_type: AccessType + ) -> None: + previous_inames = self._additional_inames + self._additional_inames |= frozenset( + iname.name for iname in expr.swept_inames + ) + try: + self.rec(expr.subscript, stmt_id, access_type) + finally: + self._additional_inames = previous_inames + + +def _set_as_map( + set_: nisl.Set, in_names: Collection[str] +) -> nisl.Map: + return set_.as_map(in_names) + + +def apply_affine_transform_to_happens_afters( + kernel: LoopKernel, affine_reln: nisl.Map +) -> LoopKernel: + """ + Applies an affine transformation to all relevant happens-after relations. + """ + + affine_reln = affine_reln.coalesce() + transformed_inames = frozenset(affine_reln.space.in_names) + name_generator = kernel.get_var_name_generator() + for names in affine_reln.space.dimtype_to_names.values(): + name_generator.add_names(names, conflicting_ok=True) + + def build_xform_reln( + stmt: InstructionBase, suffix: str + ) -> tuple[nisl.Map, tuple[tuple[str, str], ...]] | None: + overlap = stmt.within_inames & transformed_inames + if not overlap: + return None + if overlap != transformed_inames: + raise LoopyError( + f"statement '{stmt.id}' is within only part of the affine " + "transformation's input inames" + ) + + stmt_domain = kernel.get_inames_domain(stmt.within_inames) + stmt_inames = stmt_domain.space.dimtype_to_names[DimType.out] + nonxformed_names = tuple( + name for name in stmt_inames if name not in transformed_inames + ) + output_proxy_names = tuple( + name_generator(f"{name}_new_") for name in nonxformed_names + ) + + xform_reln = affine_reln.add_dims(DimType.in_, nonxformed_names) + xform_reln = xform_reln.add_dims(DimType.out, output_proxy_names) + xform_reln = xform_reln.equate_dims(tuple(zip( + nonxformed_names, output_proxy_names, strict=True + ))) + xform_reln = _suffix_names(xform_reln, suffix, DimType.in_) + xform_reln = _suffix_names(xform_reln, suffix, DimType.out) + + proxy_renames = tuple( + (f"{proxy}{suffix}", f"{name}{suffix}") + for name, proxy in zip( + nonxformed_names, output_proxy_names, strict=True + ) + ) + return xform_reln.coalesce(), proxy_renames + + new_stmts: list[InstructionBase] = [] + for sink_stmt in kernel.instructions: + sink_xform = build_xform_reln(sink_stmt, "_after") + new_happens_after: dict[str, HappensAfter] = {} + + for src_id, happens_after in sink_stmt.happens_after.items(): + if happens_after.instances_rel is None: + raise LoopyError( + "cannot determine precise happens-after information" + ) + + src_stmt = kernel.id_to_insn[src_id] + src_xform = build_xform_reln(src_stmt, "_before") + if sink_xform is None and src_xform is None: + new_happens_after[src_id] = happens_after + continue + + instances_rel = happens_after.instances_rel.coalesce() + proxy_renames: list[tuple[str, str]] = [] + + if sink_xform is not None: + sink_xform_reln, sink_proxy_renames = sink_xform + instances_rel = sink_xform_reln.reverse().apply_range( + instances_rel + ).coalesce() + proxy_renames.extend(sink_proxy_renames) + + if src_xform is not None: + if instances_rel.space.dim(DimType.in_) == 0: + dummy_name = name_generator("happens_after_dummy") + instances_rel = ( + instances_rel + .add_dims(DimType.in_, (dummy_name,)) + .project_out((dummy_name,)) + ) + + src_xform_reln, src_proxy_renames = src_xform + instances_rel = instances_rel.apply_range( + src_xform_reln + ).coalesce() + proxy_renames.extend(src_proxy_renames) + + instances_rel = instances_rel.rename_dims(proxy_renames).coalesce() + new_happens_after[src_id] = HappensAfter( + instances_rel, + variable_name=happens_after.variable_name, + ) + + new_stmts.append( + sink_stmt.copy(happens_after=constantdict(new_happens_after)) + ) + + return kernel.copy(instructions=new_stmts) + + +def has_precise_dependencies(kernel: LoopKernel) -> bool: + has_precise = False + has_legacy = False + for stmt in kernel.instructions: + for happens_after in stmt.happens_after.values(): + if happens_after.instances_rel is None: + has_legacy = True + else: + has_precise = True + + if has_precise and has_legacy: + raise LoopyError( + f"kernel '{kernel.name}' mixes precise and legacy " + "happens-after dependencies" + ) + + return has_precise + + +def _suffix_names( + obj: NamedIslObjectT, suffix: str, dim_type: DimType +) -> NamedIslObjectT: + return obj.rename_dims( + (name, name + suffix) for name in obj.space.dimtype_to_names[dim_type] + ) + + +def _statement_instance_set( + kernel: LoopKernel, stmt: InstructionBase, suffix: str +) -> nisl.Set: + instance_set = kernel.get_inames_domain(stmt.within_inames) + unused_inames = instance_set.space.out_names - stmt.within_inames + if unused_inames: + instance_set = instance_set.project_out(unused_inames) + + return _suffix_names(instance_set, suffix, DimType.out).coalesce() + + +def _compose_happens_after_relations( + first: nisl.Map, second: nisl.Map +) -> nisl.Map: + first = first.coalesce() + second = second.coalesce() + first_interface = tuple( + name.removesuffix("_before") + for name in first.space.dimtype_to_names[DimType.out] + ) + second_interface = tuple( + name.removesuffix("_after") + for name in second.space.dimtype_to_names[DimType.in_] + ) + if frozenset(first_interface) != frozenset(second_interface): + raise LoopyError( + "cannot compose happens-after relations with different " + "intermediate instance spaces" + ) + + first = first.rename_dims(zip( + first.space.dimtype_to_names[DimType.out], first_interface, strict=True + )) + second = second.rename_dims(zip( + second.space.dimtype_to_names[DimType.in_], second_interface, strict=True + )) + return first.apply_range(second).coalesce() + + +def _saturate_cross_relations_with_self_relations( + kernel: LoopKernel, +) -> LoopKernel: + self_relations: dict[str, nisl.Map] = {} + for stmt in kernel.instructions: + happens_after = stmt.happens_after.get(stmt.id) + if happens_after is None: + continue + if happens_after.instances_rel is None: + raise LoopyError( + "self-relation saturation requires precise dependencies" + ) + + self_relation = happens_after.instances_rel.coalesce() + if not ( + _compose_happens_after_relations( + self_relation, self_relation + ) - self_relation + ).is_empty(): + raise LoopyError( + f"self happens-after relation for '{stmt.id}' is not " + "transitive" + ) + + self_relations[stmt.id] = self_relation + + new_stmts: list[InstructionBase] = [] + for stmt in kernel.instructions: + new_happens_after: dict[str, HappensAfter] = {} + for source_id, happens_after in stmt.happens_after.items(): + if happens_after.instances_rel is None: + raise LoopyError( + "self-relation saturation requires precise dependencies" + ) + + relation = happens_after.instances_rel.coalesce() + if source_id != stmt.id: + self_relation = self_relations.get(stmt.id) + if self_relation is not None: + relation = ( + relation + | _compose_happens_after_relations( + self_relation, relation + ) + ).coalesce() + + self_relation = self_relations.get(source_id) + if self_relation is not None: + relation = ( + relation + | _compose_happens_after_relations( + relation, self_relation + ) + ).coalesce() + + new_happens_after[source_id] = HappensAfter( + relation, + variable_name=happens_after.variable_name, + ) + + new_stmts.append( + stmt.copy(happens_after=constantdict(new_happens_after)) + ) + + return kernel.copy(instructions=tuple(new_stmts)) + + +def _validate_instance_mapping( + relation: nisl.Map, + domain_instances: nisl.Set, + range_instances: nisl.Set, + *, + relation_name: str, + domain_name: str, + range_name: str, +) -> nisl.Map: + relation = relation.coalesce() + if ( + relation.space.dimtype_to_names[DimType.in_] + != domain_instances.space.dimtype_to_names[DimType.out] + ): + raise LoopyError( + f"{relation_name} relation has the wrong {domain_name} " + "instance space" + ) + if ( + relation.space.dimtype_to_names[DimType.out] + != range_instances.space.dimtype_to_names[DimType.out] + ): + raise LoopyError( + f"{relation_name} relation has the wrong {range_name} " + "instance space" + ) + if not (relation.domain() - domain_instances).is_empty(): + raise LoopyError( + f"{relation_name} relation contains instances outside the " + f"{domain_name} domain" + ) + if not (relation.range() - range_instances).is_empty(): + raise LoopyError( + f"{relation_name} relation contains instances outside the " + f"{range_name} domain" + ) + + return relation + + +def _add_or_union_happens_after( + happens_after: dict[str, HappensAfter], + sink_id: str, + source_id: str, + instances_rel: nisl.Map, + variable_name: str | None = None, +) -> None: + instances_rel = instances_rel.coalesce() + if instances_rel.is_empty(): + return + + previous = happens_after.get(source_id) + if previous is not None: + if previous.instances_rel is None: + raise LoopyError( + "cannot combine precise and imprecise happens-after " + f"relations for '{sink_id}' and '{source_id}'" + ) + + previous_rel = previous.instances_rel + if ( + previous_rel.space.dimtype_to_names[DimType.in_] + != instances_rel.space.dimtype_to_names[DimType.in_] + or previous_rel.space.dimtype_to_names[DimType.out] + != instances_rel.space.dimtype_to_names[DimType.out] + ): + raise LoopyError( + "cannot union happens-after relations with different " + "statement instance spaces" + ) + instances_rel = (previous_rel | instances_rel).coalesce() + if variable_name != previous.variable_name: + variable_name = None + + happens_after[source_id] = HappensAfter( + instances_rel, + variable_name=variable_name, + ) + + +def splice_happens_after_as_consumer( + kernel: LoopKernel, + consumer_id: str, + anchor_id: str, + consumer_to_anchor: nisl.Map, +) -> LoopKernel: + """Give *consumer_id* the incoming dependencies of *anchor_id*. + + *consumer_to_anchor* maps consumer instances to the anchor instances whose + incoming dependencies they inherit. Its input dimensions use the consumer + inames suffixed with ``"_after"`` and its output dimensions use the anchor + inames suffixed with ``"_before"``. + + The anchor's self-edge is not inherited. Existing dependencies of the + consumer are preserved and unioned with inherited dependencies to the same + source. + """ + + if consumer_id == anchor_id: + raise LoopyError("consumer and anchor instructions must be distinct") + if not has_precise_dependencies(kernel): + raise LoopyError("consumer splicing requires precise dependencies") + + try: + consumer = kernel.id_to_insn[consumer_id] + anchor = kernel.id_to_insn[anchor_id] + except KeyError as err: + raise LoopyError(f"unknown instruction ID '{err.args[0]}'") from err + + consumer_instances = _statement_instance_set(kernel, consumer, "_after") + anchor_instances = _statement_instance_set(kernel, anchor, "_before") + consumer_to_anchor = _validate_instance_mapping( + consumer_to_anchor, + consumer_instances, + anchor_instances, + relation_name="consumer-to-anchor", + domain_name="consumer", + range_name="anchor", + ) + + new_happens_after = dict(consumer.happens_after) + for source_id, happens_after in anchor.happens_after.items(): + if source_id == anchor_id: + continue + if happens_after.instances_rel is None: + raise LoopyError( + "cannot inherit an imprecise happens-after relation from " + f"'{anchor_id}' to '{source_id}'" + ) + + inherited = _compose_happens_after_relations( + consumer_to_anchor, + happens_after.instances_rel, + ) + _add_or_union_happens_after( + new_happens_after, + consumer_id, + source_id, + inherited, + variable_name=happens_after.variable_name, + ) + + return kernel.copy(instructions=tuple( + stmt.copy(happens_after=constantdict(new_happens_after)) + if stmt.id == consumer_id else stmt + for stmt in kernel.instructions + )) + + +def splice_happens_after_as_producer( + kernel: LoopKernel, + producer_id: str, + anchor_id: str, + anchor_to_producer: nisl.Map, +) -> LoopKernel: + """Redirect consumers of *anchor_id* to *producer_id*. + + *anchor_to_producer* maps the anchor instances being replaced to the + producer instances that replace them. Its input dimensions use the anchor + inames suffixed with ``"_after"`` and its output dimensions use the producer + inames suffixed with ``"_before"``. + + If the map covers only part of the anchor instance space, dependencies on + the remaining anchor instances are preserved. Existing dependencies on the + producer are unioned with the redirected dependencies. + """ + + if producer_id == anchor_id: + raise LoopyError("producer and anchor instructions must be distinct") + if not has_precise_dependencies(kernel): + raise LoopyError("producer splicing requires precise dependencies") + + try: + producer = kernel.id_to_insn[producer_id] + anchor = kernel.id_to_insn[anchor_id] + except KeyError as err: + raise LoopyError(f"unknown instruction ID '{err.args[0]}'") from err + + anchor_instances = _statement_instance_set(kernel, anchor, "_after") + producer_instances = _statement_instance_set(kernel, producer, "_before") + anchor_to_producer = _validate_instance_mapping( + anchor_to_producer, + anchor_instances, + producer_instances, + relation_name="anchor-to-producer", + domain_name="anchor", + range_name="producer", + ) + + mapped_anchor_instances = ( + anchor_to_producer.domain() + .coalesce() + .rename_dims(zip( + anchor_to_producer.space.dimtype_to_names[DimType.in_], + tuple( + name.removesuffix("_after") + "_before" + for name in anchor_to_producer.space.dimtype_to_names[DimType.in_] + ), + strict=True, + )) + ) + + new_stmts: list[InstructionBase] = [] + for sink in kernel.instructions: + if sink.id in {anchor_id, producer_id}: + new_stmts.append(sink) + continue + + happens_after = sink.happens_after.get(anchor_id) + if happens_after is None: + new_stmts.append(sink) + continue + if happens_after.instances_rel is None: + raise LoopyError( + "cannot redirect an imprecise happens-after relation from " + f"'{sink.id}' to '{anchor_id}'" + ) + + anchor_order = happens_after.instances_rel.coalesce() + redirected_anchor_order = anchor_order.intersect_range( + mapped_anchor_instances + ).coalesce() + remaining_anchor_order = ( + anchor_order - redirected_anchor_order + ).coalesce() + redirected_order = _compose_happens_after_relations( + redirected_anchor_order, anchor_to_producer + ) + + new_happens_after = dict(sink.happens_after) + if remaining_anchor_order.is_empty(): + del new_happens_after[anchor_id] + else: + new_happens_after[anchor_id] = HappensAfter( + remaining_anchor_order, + variable_name=happens_after.variable_name, + ) + + _add_or_union_happens_after( + new_happens_after, + sink.id, + producer_id, + redirected_order, + variable_name=happens_after.variable_name, + ) + + new_stmts.append( + sink.copy(happens_after=constantdict(new_happens_after)) + ) + + return kernel.copy(instructions=tuple(new_stmts)) + + +def splice_happens_after_as_consumer_and_producer( + kernel: LoopKernel, + instruction_id: str, + anchor_id: str, + instruction_to_anchor: nisl.Map, + anchor_to_instruction: nisl.Map, +) -> LoopKernel: + """Splice *instruction_id* across both sides of *anchor_id*. + + The new instruction inherits the anchor's incoming dependencies according + to *instruction_to_anchor*. Dependencies on the mapped anchor instances are + redirected to the new instruction according to *anchor_to_instruction*. + The two relations are supplied independently and need not be inverses. + """ + + kernel = splice_happens_after_as_consumer( + kernel, + instruction_id, + anchor_id, + instruction_to_anchor, + ) + return splice_happens_after_as_producer( + kernel, + instruction_id, + anchor_id, + anchor_to_instruction, + ) + + +@for_each_kernel +def add_lexicographic_happens_after(kernel: LoopKernel) -> LoopKernel: + """ + Imposes a strict lexicographic order on all statements in *kernel*. The + order of statements as they appear in the kernel is used to impose the + dependence relations. + + The following two conditions are true of the order imposed by this routine: + 1. All statements will have a self-dependence relation defined + 2. All statements except the first statement (as dictated by kernel order) + will have a dependence relation defined between itself and the + immediately preceding statement. + """ + + new_stmts: list[InstructionBase] = [] + for i, stmt in enumerate(kernel.instructions): + new_happens_after: dict[str, HappensAfter] = {} + + sources = (stmt,) if i == 0 else (stmt, kernel.instructions[i - 1]) + + after_domain = kernel.get_inames_domain(stmt.within_inames) + + after_inames = after_domain.space.dimtype_to_names[DimType.out] + after_domain = _suffix_names(after_domain, "_after", DimType.out) + for source in sources: + before_domain = kernel.get_inames_domain(source.within_inames) + + before_inames = before_domain.space.dimtype_to_names[DimType.out] + before_domain = _suffix_names(before_domain, "_before", DimType.out) + + # lexicographic order necessitates agreement between before and + # after on the order of shared inames + shared_inames = frozenset(before_inames) & frozenset(after_inames) + before_order = tuple( + iname for iname in before_inames if iname in shared_inames + ) + after_order = tuple( + iname for iname in after_inames if iname in shared_inames + ) + + assert before_order == after_order + shared_order = after_order + + joint_domain = after_domain & before_domain + affs = joint_domain.var_pw_affs + + strict_lex = joint_domain - joint_domain + equal_prefix = joint_domain + for iname in shared_order: + after_aff = affs[f"{iname}_after"] + before_aff = affs[f"{iname}_before"] + + strict_lex = strict_lex | ( + equal_prefix & after_aff.gt_set(before_aff) + ) + equal_prefix = equal_prefix & after_aff.eq_set(before_aff) + + if source.id == stmt.id: + ordered_instances = strict_lex + else: + ordered_instances = strict_lex | equal_prefix + + instances_rel = _set_as_map( + ordered_instances, + in_names=tuple(f"{name}_after" for name in after_inames) + ) + + new_happens_after[source.id] = HappensAfter( + instances_rel=instances_rel + ) + + new_stmts.append(stmt.copy(happens_after=new_happens_after)) + + return kernel.copy(instructions=new_stmts) + + +def _compute_reachable_happens_after( + kernel: LoopKernel, + topological_order: list[str], +) -> dict[str, dict[str, nisl.Map]]: + """Return every nonempty fine-grained branch of the supplied order.""" + result: dict[str, dict[str, nisl.Map]] = {} + + for sink_id in topological_order: + reachable: dict[str, nisl.Map] = {} + self_happens_after = kernel.id_to_insn[sink_id].happens_after.get(sink_id) + if self_happens_after is not None: + assert self_happens_after.instances_rel is not None + if not self_happens_after.instances_rel.is_empty(): + reachable[sink_id] = self_happens_after.instances_rel + + for intermediate_id in topological_order: + if intermediate_id == sink_id: + sink_to_intermediate = None + else: + sink_to_intermediate = reachable.get(intermediate_id) + if sink_to_intermediate is None: + continue + + intermediate = kernel.id_to_insn[intermediate_id] + for source_id, happens_after in intermediate.happens_after.items(): + if source_id == intermediate_id: + continue + + assert happens_after.instances_rel is not None + edge_relation = happens_after.instances_rel + if edge_relation.is_empty(): + continue + + if sink_to_intermediate is None: + path_relation = edge_relation + else: + path_relation = _compose_happens_after_relations( + sink_to_intermediate, edge_relation + ) + if path_relation.is_empty(): + continue + + previous = reachable.get(source_id) + reachable[source_id] = ( + path_relation + if previous is None + else (previous | path_relation).coalesce() + ) + + result[sink_id] = reachable + + return result + + +def _find_conflicting_access_candidates( + sink_access_relation: nisl.Map, + sink_to_source: nisl.Map, + source_access_relation: nisl.Map, +) -> nisl.Map: + """Return ``(sink instance, cell) -> source instance`` conflicts.""" + source_access_relation = _suffix_names( + source_access_relation, "_before", DimType.in_ + ) + sink_names = sink_to_source.space.in_names + cell_names = sink_access_relation.space.out_names + candidate_set = ( + sink_access_relation.as_set() + & sink_to_source.as_set() + & source_access_relation.as_set() + ) + return _set_as_map( + candidate_set, in_names=(*sink_names, *cell_names) + ).coalesce() + + +def _discard_candidates_preceding_writers( + candidates: Mapping[str, nisl.Map], + writer_candidates: Mapping[str, nisl.Map], + reachable_order: Mapping[str, Mapping[str, nisl.Map]], +) -> dict[str, nisl.Map]: + """Remove candidates with a same-cell writer ordered after them.""" + result: dict[str, nisl.Map] = {} + for candidate_id, candidate_relation in candidates.items(): + dominated = candidate_relation - candidate_relation + for writer_id, writer_relation in writer_candidates.items(): + writer_to_candidate = reachable_order[writer_id].get(candidate_id) + if writer_to_candidate is None: + continue + + dominated = dominated | ( + candidate_relation + & _compose_happens_after_relations( + writer_relation, writer_to_candidate + ) + ) + + remaining = (candidate_relation - dominated).coalesce() + if not remaining.is_empty(): + result[candidate_id] = remaining + + return result + + +def _record_candidate_order( + candidates: Mapping[str, nisl.Map], + sink_names: Collection[str], + cell_names: Collection[str], + var: str, + happens_after: dict[str, HappensAfter], +) -> None: + for source_id, candidate_relation in candidates.items(): + required_order = _set_as_map( + candidate_relation.as_set().project_out(cell_names), + in_names=sink_names, + ).coalesce() + previous = happens_after.get(source_id) + if previous is None: + combined_order = required_order + else: + assert previous.instances_rel is not None + combined_order = ( + required_order | previous.instances_rel + ).coalesce() + + variable_name = ( + var + if previous is None or previous.variable_name == var + else None + ) + happens_after[source_id] = HappensAfter( + combined_order, + variable_name=variable_name, + ) + + +@for_each_kernel +def relax_strict_happens_after(kernel: LoopKernel) -> LoopKernel: + """ + Relaxes an incoming strict execution order imposed on statements in *kernel* + through dependence analysis. + + :returns: *kernel* with the minimally required execution order on statement + instances in a program needed to satisfy data dependencies. + """ + + for stmt in kernel.instructions: + if isinstance(stmt, CInstruction): + raise LoopyError( + "precise dependency relaxation does not support " + f"CInstruction '{stmt.id}'" + ) + + for temporary in kernel.temporary_variables.values(): + if temporary.base_storage is not None: + raise LoopyError( + "precise dependency relaxation does not support temporary " + f"'{temporary.name}' with base_storage" + ) + + kernel = _saturate_cross_relations_with_self_relations(kernel) + + coarse_dependency_graph: dict[str, frozenset[str]] = {} + for stmt in kernel.instructions: + dependencies: set[str] = set() + for source_id, happens_after in stmt.happens_after.items(): + if source_id == stmt.id: + continue + + assert happens_after.instances_rel is not None + relation = happens_after.instances_rel + if not relation.is_empty(): + dependencies.add(source_id) + + coarse_dependency_graph[stmt.id] = frozenset(dependencies) + + topological_order = compute_topological_order(coarse_dependency_graph) + reachable_order = _compute_reachable_happens_after( + kernel, topological_order + ) + + rel_finder = AccessRelationFinder(kernel) + for stmt in kernel.instructions: + if isinstance(stmt, MultiAssignmentBase): + for assignee in stmt.assignees: + rel_finder(assignee, stmt.id, AccessType.write) + rel_finder(stmt.expression, stmt.id, AccessType.read) + for pred in stmt.predicates: + rel_finder(pred, stmt.id, AccessType.read) + + new_stmts: list[InstructionBase] = [] + for sink_id in topological_order: + new_happens_after: dict[str, HappensAfter] = {} + for sink_access_type, access_relations in ( + (AccessType.read, rel_finder.read_relations[sink_id]), + (AccessType.write, rel_finder.write_relations[sink_id]), + ): + for var, access_relation in access_relations.items(): + access_relation = _suffix_names( + access_relation, "_after", DimType.in_ + ) + writer_candidates: dict[str, nisl.Map] = {} + reader_candidates: dict[str, nisl.Map] = {} + for source_id, sink_to_source in reachable_order[sink_id].items(): + source_writes = rel_finder.write_relations[source_id] + if var in source_writes: + candidates = _find_conflicting_access_candidates( + access_relation, + sink_to_source, + source_writes[var], + ) + if not candidates.is_empty(): + writer_candidates[source_id] = candidates + + if sink_access_type == AccessType.write: + source_reads = rel_finder.read_relations[source_id] + if var in source_reads: + candidates = _find_conflicting_access_candidates( + access_relation, + sink_to_source, + source_reads[var], + ) + if not candidates.is_empty(): + reader_candidates[source_id] = candidates + + most_recent_writers = _discard_candidates_preceding_writers( + writer_candidates, + writer_candidates, + reachable_order, + ) + sink_names = access_relation.space.in_names + cell_names = access_relation.space.out_names + _record_candidate_order( + most_recent_writers, + sink_names, + cell_names, + var, + new_happens_after, + ) + + if sink_access_type == AccessType.write: + live_readers = _discard_candidates_preceding_writers( + reader_candidates, + writer_candidates, + reachable_order, + ) + _record_candidate_order( + live_readers, + sink_names, + cell_names, + var, + new_happens_after, + ) + + new_stmts.append( + kernel.id_to_insn[sink_id].copy( + happens_after=constantdict(new_happens_after) + ) + ) + + return kernel.copy(instructions=new_stmts) diff --git a/loopy/kernel/instruction.py b/loopy/kernel/instruction.py index 028e8b73a..6aeecf008 100644 --- a/loopy/kernel/instruction.py +++ b/loopy/kernel/instruction.py @@ -58,6 +58,7 @@ if TYPE_CHECKING: + from pymbolic import Expression from loopy.kernel import LoopKernel @@ -119,35 +120,36 @@ class UseStreamingStoreTag(Tag): @dataclass(frozen=True) class HappensAfter: """A class representing a "happens-after" relationship between two - statements found in a :class:`loopy.LoopKernel`. Used to validate that a - given kernel transformation respects the data dependencies in a given - program. + statements found in a :class:`loopy.LoopKernel`. Used to analyze and verify + data dependencies are respected after a :class:`loopy.LoopKernel` has been + scheduled. .. attribute:: variable_name - The name of the variable responsible for the dependency. For - backward compatibility purposes, this may be *None*. In this case, the - dependency semantics revert to the deprecated, statement-level - dependencies of prior versions of :mod:`loopy`. + The name of the variable responsible for the dependency, if the + dependency records one variable explicitly. This may be *None* when the + metadata is unavailable or one exact relation combines several + variables. .. attribute:: instances_rel - An :class:`islpy.Map` representing the precise happens-after + An :class:`namedisl.Map` representing the precise happens-after relationship. The domain and range are sets of statement instances. The - instances in the domain are required to execute before the instances in + instances in the domain are required to execute after the instances in the range. Map dimensions are named according to the order of appearance of the - inames in a :mod:`loopy` program. The dimension names in the range are - appended with a prime to signify that the mapped instances are distinct. + inames in a :mod:`loopy` program. The names in the domain are suffixed + with "_after" and the names in the range are suffixed with "_before" to + signify that the instances are distinct. As a (deprecated) matter of backward compatibility, this may be *None*, in which case the semantics revert to the (underspecified) statement-level dependencies of prior versions of :mod:`loopy`. """ - variable_name: str | None instances_rel: nisl.Map | None + variable_name: str | None = None # }}} @@ -356,14 +358,12 @@ def __init__(self, happens_after = constantdict({ after_id.strip(): HappensAfter( - variable_name=None, instances_rel=None) for after_id in happens_after.split(",") if after_id.strip()}) elif isinstance(happens_after, frozenset): happens_after = constantdict({ after_id: HappensAfter( - variable_name=None, instances_rel=None) for after_id in happens_after}) elif isinstance(happens_after, dict): diff --git a/loopy/schedule/__init__.py b/loopy/schedule/__init__.py index 86d779414..28691e67d 100644 --- a/loopy/schedule/__init__.py +++ b/loopy/schedule/__init__.py @@ -980,6 +980,8 @@ def _generate_loop_schedules_v2(kernel: LoopKernel) -> Sequence[ScheduleItem]: insn_loop_inames = insn.within_inames & loop_inames for dep_id in insn.depends_on: + if dep_id == insn.id: + continue dep = kernel.id_to_insn[dep_id] dep_loop_inames = dep.within_inames & loop_inames # Enforce instruction dep: diff --git a/loopy/schedule/verification.py b/loopy/schedule/verification.py new file mode 100644 index 000000000..2270c0cf7 --- /dev/null +++ b/loopy/schedule/verification.py @@ -0,0 +1,587 @@ +from __future__ import annotations + + +__copyright__ = """ +Copyright (C) 2026 Addison Alvey-Blanco +""" + +__license__ = """ +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +""" + + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import namedisl as nisl +from constantdict import constantdict +from namedisl import DimType + +from loopy import KernelState, LoopKernel, for_each_kernel +from loopy.diagnostic import LoopyError +from loopy.schedule import ( + Barrier, + CallKernel, + EnterLoop, + LeaveLoop, + ReturnFromKernel, + RunInstruction, + ScheduleItem, +) + + +if TYPE_CHECKING: + from collections.abc import Mapping, Sequence + + from pytools import UniqueNameGenerator + + +@dataclass(frozen=True) +class _PreciseScheduleRecord: + timestamp: Sequence[int | str] + + +@dataclass(frozen=True) +class _StatementRecord(_PreciseScheduleRecord): + subkernel_idx: int | None + + +@dataclass(frozen=True) +class _BarrierRecord(_PreciseScheduleRecord): + barrier: Barrier + subkernel_idx: int | None = None + + +@dataclass(frozen=True) +class _PreciseSchedule: + statements: Mapping[str, _StatementRecord] + barriers: Sequence[_BarrierRecord] + + +def _get_timestamp_points_from_linearization( + kernel: LoopKernel, +) -> _PreciseSchedule: + if kernel.linearization is None or kernel.state != KernelState.LINEARIZED: + raise LoopyError( + "Kernel must be linearized before instance-level analysis." + ) + + from loopy.kernel.data import VectorizeTag + + def build_timestamp_from_stack( + stack: list[tuple[int, ScheduleItem]], + ) -> Sequence[int | str]: + tstamp: Sequence[int | str] = [] + for frame in stack: + match frame: + case (x, EnterLoop(iname=iname)): + if kernel.iname_tags_of_type(iname, VectorizeTag): + continue + + tstamp.append(x) + tstamp.append(iname) + + case (x, CallKernel(_)): + tstamp.append(x) + + case _: + pass + + return tuple(tstamp) + + def find_most_recent_subkernel_idx( + stack: list[tuple[int, ScheduleItem]], + ) -> int | None: + subkernel_idx = None + for frame in stack: + pos, sched_item = frame + if isinstance(sched_item, CallKernel): + subkernel_idx = pos + break + return subkernel_idx + + stack: list[tuple[int, ScheduleItem]] = [] + stmt_records: dict[str, _StatementRecord] = {} + bar_records: list[_BarrierRecord] = [] + for i, node in enumerate(kernel.linearization): + match node: + case CallKernel(_): + stack.append((i, node)) + + case ReturnFromKernel(_): + stack.pop() + + case EnterLoop(_): + stack.append((i, node)) + + case LeaveLoop(_): + stack.pop() + + case RunInstruction(insn_id=stmt_id): + tstamp = build_timestamp_from_stack(stack) + tstamp = (*tstamp, i) + + subkernel_idx = find_most_recent_subkernel_idx(stack) + + if subkernel_idx is None: + raise LoopyError( + f"could not determine subkernel index for {stmt_id}" + ) + + stmt_records[stmt_id] = _StatementRecord( + timestamp=tstamp, subkernel_idx=subkernel_idx + ) + + case Barrier(_): + tstamp = build_timestamp_from_stack(stack) + tstamp = (*tstamp, i) + + subkernel_idx = find_most_recent_subkernel_idx(stack) + + bar_records.append( + _BarrierRecord( + timestamp=tstamp, + barrier=node, + subkernel_idx=subkernel_idx, + ) + ) + + if node.originating_insn_id is not None: + stmt_records[node.originating_insn_id] = _StatementRecord( + timestamp=tstamp, + subkernel_idx=subkernel_idx, + ) + + case _: + pass + + return _PreciseSchedule( + statements=constantdict(stmt_records), barriers=tuple(bar_records) + ) + + +def _build_statement_timestamp_relations( + kernel: LoopKernel, + stmt_records: Mapping[str, _StatementRecord], + timestamp_names: Sequence[str], +) -> Mapping[str, nisl.Map]: + stmt_relns: dict[str, nisl.Map] = {} + for stmt_id, record in stmt_records.items(): + stmt = kernel.id_to_insn[stmt_id] + + # pad so that composition interface matches across all statements + pad = len(timestamp_names) - len(record.timestamp) + timestamp = list(record.timestamp) + timestamp.extend([0 for _ in range(pad)]) + + ran_str = ", ".join( + f"{name} = {pos}" + for name, pos in zip(timestamp_names, timestamp, strict=True) + ) + + dom_str = ", ".join(name for name in stmt.within_inames) + + full_str = dom_str + ", " + ran_str if dom_str else ran_str + + domain = kernel.get_inames_domain( + stmt.within_inames + ).project_out_except( + [*stmt.within_inames, *kernel.all_params()], dim_type="all" + ) + + if stmt.within_inames: + reln = nisl.make_set(f"{{[{full_str}]}}").as_map(stmt.within_inames) + else: + constraints = " and ".join( + f"{name} = {pos}" + for name, pos in zip(timestamp_names, timestamp, strict=True) + ) + reln = nisl.make_map( + f"{{ [] -> [{', '.join(timestamp_names)}] : {constraints} }}" + ) + + reln = reln.intersect_domain(domain) + + stmt_relns[stmt_id] = reln + + return stmt_relns + + +def _build_barrier_timestamp_relations( + kernel: LoopKernel, + barrier_records: Sequence[_BarrierRecord], + timestamp_names: Sequence[str], +) -> Sequence[nisl.Map]: + barrier_relns: list[nisl.Map] = [] + for record in barrier_records: + inames = tuple( + value for value in record.timestamp if isinstance(value, str) + ) + + pad = len(timestamp_names) - len(record.timestamp) + timestamp = [*record.timestamp, *(0 for _ in range(pad))] + + ran_str = ", ".join( + f"{name} = {pos}" + for name, pos in zip(timestamp_names, timestamp, strict=True) + ) + dom_str = ", ".join(inames) + full_str = f"{dom_str}, {ran_str}" if dom_str else ran_str + + domain = kernel.get_inames_domain( + frozenset(inames) + ).project_out_except( + [*inames, *kernel.all_params()], dim_type="all" + ) + + if inames: + relation = nisl.make_set(f"{{[{full_str}]}}").as_map(inames) + else: + constraints = " and ".join( + f"{name} = {pos}" + for name, pos in zip(timestamp_names, timestamp, strict=True) + ) + relation = nisl.make_map( + f"{{ [] -> [{', '.join(timestamp_names)}] : {constraints} }}" + ) + + barrier_relns.append(relation.intersect_domain(domain)) + + return tuple(barrier_relns) + + +def _build_strict_lexicographic_order( + later_names: Sequence[str], + earlier_names: Sequence[str], +) -> nisl.Map: + joint = nisl.make_set( + f"{{ [{', '.join([*later_names, *earlier_names])}] }}" + ) + affs = joint.var_pw_affs + + strict_lex = joint - joint + equal_prefix = joint + + for later_name, earlier_name in zip( + later_names, earlier_names, strict=True + ): + strict_lex = strict_lex | ( + equal_prefix & affs[later_name].gt_set(affs[earlier_name]) + ) + equal_prefix = equal_prefix & affs[later_name].eq_set( + affs[earlier_name] + ) + + return strict_lex.as_map(later_names) + + +def _build_timestamp_relations( + kernel: LoopKernel, + prec_sched: _PreciseSchedule, + name_generator: UniqueNameGenerator, +) -> tuple[Mapping[str, nisl.Map], Sequence[nisl.Map], nisl.Map]: + max_stmt_tstamp_len = max( + len(record.timestamp) for record in prec_sched.statements.values() + ) + + max_bar_tstamp_len = -1 + if prec_sched.barriers: + max_bar_tstamp_len = max( + len(record.timestamp) for record in prec_sched.barriers + ) + + max_tstamp_len = max(max_stmt_tstamp_len, max_bar_tstamp_len) + timestamp_names = tuple( + name_generator(f"__ts_{i}") for i in range(max_tstamp_len) + ) + later_names = tuple( + name_generator(f"{name}_later") for name in timestamp_names + ) + earlier_names = tuple( + name_generator(f"{name}_earlier") for name in timestamp_names + ) + + stmt_relns = _build_statement_timestamp_relations( + kernel, prec_sched.statements, timestamp_names + ) + + bar_relns = _build_barrier_timestamp_relations( + kernel, prec_sched.barriers, timestamp_names + ) + + timestamp_lex = _build_strict_lexicographic_order( + later_names, earlier_names + ) + + return constantdict(stmt_relns), bar_relns, timestamp_lex + + +def _suffix_dim_names( + relation: nisl.Map, + dim_type: DimType, + suffix: str, +) -> nisl.Map: + return relation.rename_dims( + (name, f"{name}{suffix}") + for name in relation.space.dimtype_to_names[dim_type] + ) + + +def _timestamp_relation_for_role( + relation: nisl.Map, + role_names: Sequence[str], +) -> nisl.Map: + return relation.rename_dims( + zip( + relation.space.dimtype_to_names[DimType.out], + role_names, + strict=True, + ) + ) + + +def _hardware_axis_inames( + kernel: LoopKernel, + stmt_id: str, + include_local_axes: bool, +) -> Mapping[tuple[str, int], str]: + from loopy.kernel.data import GroupInameTag, LocalInameTag + + result: dict[tuple[str, int], str] = {} + for iname in kernel.id_to_insn[stmt_id].within_inames: + tags = kernel.iname_tags_of_type( + iname, (GroupInameTag, LocalInameTag), max_num=1 + ) + if not tags: + continue + + (tag,) = tags + if isinstance(tag, GroupInameTag): + key = ("group", tag.axis) + elif include_local_axes: + key = ("local", tag.axis) + else: + continue + + if key in result: + raise LoopyError( + f"instruction '{stmt_id}' uses multiple inames for " + f"hardware axis '{key[0]}.{key[1]}'" + ) + result[key] = iname + + return constantdict(result) + + +def _build_hardware_id_relation( + kernel: LoopKernel, + stmt_id: str, + instance_domain: nisl.Set, + instance_suffix: str, + include_local_axes: bool, + hardware_names: Mapping[tuple[str, int], str], +) -> nisl.Map: + axis_inames = _hardware_axis_inames(kernel, stmt_id, include_local_axes) + input_names = instance_domain.space.dimtype_to_names[DimType.out] + constraints = " and ".join( + f"{hardware_names[kind, axis]} = " + f"{axis_inames[kind, axis]}{instance_suffix}" + for kind, axis in sorted(axis_inames) + ) + constraint_str = f" : {constraints}" if constraints else "" + + relation = nisl.make_map( + "{ " + f"[{', '.join(input_names)}] -> " + f"[{', '.join(hardware_names[key] for key in sorted(axis_inames))}]" + f"{constraint_str} " + "}" + ) + return relation.intersect_domain(instance_domain) + + +def _build_same_hardware_scope_relation( + kernel: LoopKernel, + sink_id: str, + source_id: str, + sink_domain: nisl.Set, + source_domain: nisl.Set, + include_local_axes: bool, + name_generator: UniqueNameGenerator, +) -> nisl.Map: + sink_axes = _hardware_axis_inames(kernel, sink_id, include_local_axes) + source_axes = _hardware_axis_inames(kernel, source_id, include_local_axes) + if sink_axes.keys() != source_axes.keys(): + scope = "work-item" if include_local_axes else "work-group" + raise LoopyError( + f"cannot compare the {scope} instances of '{sink_id}' and " + f"'{source_id}': their hardware axes differ" + ) + + hardware_names = { + key: name_generator(f"__{key[0]}_{key[1]}") + for key in sorted(sink_axes) + } + + sink_hardware = _build_hardware_id_relation( + kernel, + sink_id, + sink_domain, + "_after", + include_local_axes, + hardware_names, + ) + source_hardware = _build_hardware_id_relation( + kernel, + source_id, + source_domain, + "_before", + include_local_axes, + hardware_names, + ) + return sink_hardware.apply_range(source_hardware.reverse()) + + +def _build_enforced_order( + kernel: LoopKernel, + sink_id: str, + source_id: str, + prec_sched: _PreciseSchedule, + stmt_relns: Mapping[str, nisl.Map], + barrier_relns: Sequence[nisl.Map], + timestamp_lex: nisl.Map, + name_generator: UniqueNameGenerator, +) -> nisl.Map: + sink = _suffix_dim_names(stmt_relns[sink_id], DimType.in_, "_after") + source = _suffix_dim_names(stmt_relns[source_id], DimType.in_, "_before") + later_names = timestamp_lex.space.dimtype_to_names[DimType.in_] + earlier_names = timestamp_lex.space.dimtype_to_names[DimType.out] + + enforced = ( + _timestamp_relation_for_role(sink, later_names) + .apply_range(timestamp_lex) + .apply_range( + _timestamp_relation_for_role(source, earlier_names).reverse() + ) + ) + enforced = enforced & _build_same_hardware_scope_relation( + kernel, + sink_id, + source_id, + sink.domain(), + source.domain(), + include_local_axes=True, + name_generator=name_generator, + ) + + sink_record = prec_sched.statements[sink_id] + source_record = prec_sched.statements[source_id] + for barrier_idx, (barrier_record, barrier_reln) in enumerate( + zip(prec_sched.barriers, barrier_relns, strict=True) + ): + if barrier_record.barrier.synchronization_kind == "local" and ( + sink_record.subkernel_idx != barrier_record.subkernel_idx + or source_record.subkernel_idx != barrier_record.subkernel_idx + ): + continue + + barrier = _suffix_dim_names( + barrier_reln, DimType.in_, f"_barrier_{barrier_idx}" + ) + sink_to_barrier = ( + _timestamp_relation_for_role(sink, later_names) + .apply_range(timestamp_lex) + .apply_range( + _timestamp_relation_for_role(barrier, earlier_names).reverse() + ) + ) + barrier_to_source = ( + _timestamp_relation_for_role(barrier, later_names) + .apply_range(timestamp_lex) + .apply_range( + _timestamp_relation_for_role(source, earlier_names).reverse() + ) + ) + through_barrier = sink_to_barrier.apply_range(barrier_to_source) + + if barrier_record.barrier.synchronization_kind == "local": + through_barrier = ( + through_barrier + & _build_same_hardware_scope_relation( + kernel, + sink_id, + source_id, + sink.domain(), + source.domain(), + include_local_axes=False, + name_generator=name_generator, + ) + ) + + enforced = enforced | through_barrier + + return enforced.coalesce() + + +@for_each_kernel +def verify_happens_after_is_enforced(kernel: LoopKernel) -> LoopKernel: + """ + Verifies that the linearization of *kernel* enforces the instance-level order + required by the :class:`loopy.HappensAfter`s of each + statement in *kernel*. + + Assumes *kernel* possesses a valid linearization. + """ + if kernel.linearization is None or kernel.state != KernelState.LINEARIZED: + raise LoopyError("Kernel must be linearized before verification.") + + prec_sched = _get_timestamp_points_from_linearization(kernel) + name_generator = kernel.get_var_name_generator() + stmt_relns, barrier_relns, timestamp_lex = _build_timestamp_relations( + kernel, prec_sched, name_generator + ) + + for sink in kernel.instructions: + for source_id, happens_after in sink.happens_after.items(): + if happens_after.instances_rel is None: + raise LoopyError( + "precise happens-after verification requires an " + f"instance relation for '{sink.id}' after " + f"'{source_id}'" + ) + + required = happens_after.instances_rel + enforced = _build_enforced_order( + kernel, + sink.id, + source_id, + prec_sched, + stmt_relns, + barrier_relns, + timestamp_lex, + name_generator, + ) + missing = required - enforced + if not missing.is_empty(): + raise LoopyError( + f"schedule does not enforce '{sink.id}' after " + f"'{source_id}': missing order {missing}" + ) + + return kernel diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index c19953f2a..0618ea5fb 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -291,6 +291,8 @@ def _split_iname_backend( matching contexts. See :func:`loopy.match.parse_match` for syntax. """ + within_was_specified = within is not None + from loopy.match import parse_match within = parse_match(within) @@ -325,6 +327,31 @@ def _split_iname_backend( if inner_iname is None: inner_iname = vng(f"{iname_to_split}_inner") + if fixed_length_is_inner: + from loopy.kernel.dependency import ( + apply_affine_transform_to_happens_afters, + has_precise_dependencies, + ) + if has_precise_dependencies(kernel): + # FIXME: Support statement-filtered affine happens-after updates. + if within_was_specified: + raise LoopyError( + "split_iname does not support 'within' when the kernel " + "has precise dependencies" + ) + + import namedisl as nisl + split_reln = nisl.make_map(f""" + {{ [{iname_to_split}] -> [{outer_iname}, {inner_iname}] : + {iname_to_split} = ( + {fixed_length}*{outer_iname} + {inner_iname}) and + 0 <= {inner_iname} < {fixed_length} + }} + """) + kernel = apply_affine_transform_to_happens_afters( + kernel, split_reln + ) + new_domains = [ _split_iname_in_set(dom, iname_to_split, inner_iname, outer_iname, fixed_length, fixed_length_is_inner) @@ -2022,6 +2049,15 @@ def process_set(s: nisl.Set) -> nisl.Set: # }}} + from loopy.kernel.dependency import ( + apply_affine_transform_to_happens_afters, + has_precise_dependencies, + ) + if has_precise_dependencies(kernel): + kernel = apply_affine_transform_to_happens_afters( + kernel, transform_map + ) + # {{{ Update within_inames for each statement # If we get this far, we know that the map was applied to exactly one domain, diff --git a/pyproject.toml b/pyproject.toml index 0f6bd6f2f..c0da5a7dd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,6 +45,8 @@ dependencies = [ "typing-extensions>=4", "strenum>=0.4.15", + + "namedisl @ git+https://github.com/inducer/namedisl.git@main" ] [project.optional-dependencies] pyopencl = [ diff --git a/test/test_dependency.py b/test/test_dependency.py new file mode 100644 index 000000000..00e8651d8 --- /dev/null +++ b/test/test_dependency.py @@ -0,0 +1,2391 @@ +__copyright__ = """ +Copyright (C) 2026 Addison Alvey-Blanco +""" + +__license__ = """ +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +""" + + +import namedisl as nisl +import numpy as np +import pytest + +import pyopencl as cl +from pymbolic import var +from pyopencl.tools import ( # ruff:ignore[unused-import] + pytest_generate_tests_for_pyopencl as pytest_generate_tests, +) + +import loopy as lp +import loopy.kernel.dependency as dep +from loopy.diagnostic import LoopyError +from loopy.kernel.instruction import HappensAfter +from loopy.schedule import ( + Barrier, + CallKernel, + EnterLoop, + LeaveLoop, + ReturnFromKernel, + RunInstruction, +) +from loopy.schedule.verification import ( + _BarrierRecord, + _build_enforced_order, + _build_strict_lexicographic_order, + _build_timestamp_relations, + _get_timestamp_points_from_linearization, + _PreciseSchedule, + _StatementRecord, + verify_happens_after_is_enforced, +) +from loopy.symbolic import SubArrayRef +from loopy.version import ( + LOOPY_USE_LANGUAGE_VERSION_2018_2, # ruff:ignore[unused-import] +) + + +def test_add_lexicographic_happens_after_is_strict_for_self() -> None: + t_unit = lp.make_kernel( + "{ [i] : 0 <= i < N }", + """ + a[i] = 2 * a[i] {id=S} + b[i] = a[i] {id=T} + """, + ) + + kernel = dep.add_lexicographic_happens_after(t_unit).default_entrypoint + + self_relation = kernel.id_to_insn["T"].happens_after["T"].instances_rel + previous_relation = kernel.id_to_insn["T"].happens_after["S"].instances_rel + + assert self_relation is not None + assert previous_relation is not None + assert self_relation.equals( + nisl.make_map(""" + [N] -> { + [i_after] -> [i_before] : + 0 <= i_before < i_after < N + } + """) + ) + assert previous_relation.equals( + nisl.make_map(""" + [N] -> { + [i_after] -> [i_before] : + 0 <= i_before <= i_after < N + } + """) + ) + + +def test_add_lexicographic_happens_after_uses_domain_dimension_order() -> None: + t_unit = lp.make_kernel( + "{ [z, a] : 0 <= z < NZ and 0 <= a < NA }", + "out[z, a] = z + a {id=S}", + ) + + kernel = dep.add_lexicographic_happens_after(t_unit).default_entrypoint + self_relation = kernel.id_to_insn["S"].happens_after["S"].instances_rel + + assert self_relation is not None + assert self_relation.equals( + nisl.make_map(""" + [NZ, NA] -> { + [z_after, a_after] -> [z_before, a_before] : + 0 <= z_after < NZ and + 0 <= z_before < NZ and + 0 <= a_after < NA and + 0 <= a_before < NA and + (z_before < z_after or + (z_before = z_after and a_before < a_after)) + } + """) + ) + + +def test_add_lexicographic_happens_after_orders_mixed_loop_nests() -> None: + t_unit = lp.make_kernel( + [ + "[NI] -> { [i] : 0 <= i < NI }", + "[i, NJ] -> { [j] : 0 <= j < NJ }", + "[i, NK] -> { [k] : 0 <= k < NK }", + "[NQ] -> { [q] : 0 <= q < NQ }", + ], + """ + a[i, j] = i + j {id=S} + b[i, k] = i + k {id=T} + c[q] = q {id=U} + """, + ) + + kernel = dep.add_lexicographic_happens_after(t_unit).default_entrypoint + shared_nest_relation = ( + kernel.id_to_insn["T"].happens_after["S"].instances_rel + ) + disjoint_nest_relation = ( + kernel.id_to_insn["U"].happens_after["T"].instances_rel + ) + + assert kernel.id_to_insn["T"].happens_after.keys() == {"S", "T"} + assert kernel.id_to_insn["U"].happens_after.keys() == {"T", "U"} + assert shared_nest_relation is not None + assert disjoint_nest_relation is not None + assert shared_nest_relation.equals( + nisl.make_map(""" + [NI, NJ, NK] -> { + [i_after, k_after] -> [i_before, j_before] : + 0 <= i_before <= i_after < NI and + 0 <= j_before < NJ and + 0 <= k_after < NK + } + """) + ) + assert disjoint_nest_relation.equals( + nisl.make_map(""" + [NI, NK, NQ] -> { + [q_after] -> [i_before, k_before] : + 0 <= q_after < NQ and + 0 <= i_before < NI and + 0 <= k_before < NK + } + """) + ) + + +def _get_precise_relation( + kernel: lp.LoopKernel, sink_id: str, source_id: str +) -> nisl.Map: + relation = kernel.id_to_insn[sink_id].happens_after[source_id].instances_rel + assert relation is not None + return relation + + +def test_affine_happens_after_transform_tiling() -> None: + t_unit = lp.make_kernel( + [ + "[N] -> { [i] : 0 <= i < N }", + "[M] -> { [j] : 0 <= j < M }", + ], + """ + a[j] = 1 {id=A} + b[i, j] = 2 {id=B} + c[i, j] = 3 {id=C} + d[j] = 4 {id=D} + """, + [ + lp.GlobalArg("a,d", shape=(var("M"),)), + lp.GlobalArg("b,c", shape=(var("N"), var("M"))), + "...", + ], + ) + kernel = dep.add_lexicographic_happens_after(t_unit).default_entrypoint + a_self = _get_precise_relation(kernel, "A", "A") + d_self = _get_precise_relation(kernel, "D", "D") + + kernel = dep.apply_affine_transform_to_happens_afters( + kernel, + nisl.make_map(""" + [N] -> { + [i] -> [io, ii] : + i = 4*io + ii and 0 <= ii < 4 + } + """), + ) + + assert _get_precise_relation(kernel, "B", "B").equals( + nisl.make_map(""" + [N, M] -> { + [io_after, ii_after, j_after] -> + [io_before, ii_before, j_before] : + 0 <= 4*io_after + ii_after < N and + 0 <= ii_after < 4 and + 0 <= j_after < M and + 0 <= 4*io_before + ii_before < N and + 0 <= ii_before < 4 and + 0 <= j_before < M and + (4*io_before + ii_before < 4*io_after + ii_after or + (4*io_before + ii_before = 4*io_after + ii_after and + j_before < j_after)) + } + """) + ) + assert _get_precise_relation(kernel, "B", "A").equals( + nisl.make_map(""" + [N, M] -> { + [io_after, ii_after, j_after] -> [j_before] : + 0 <= 4*io_after + ii_after < N and + 0 <= ii_after < 4 and + 0 <= j_before <= j_after < M + } + """) + ) + assert _get_precise_relation(kernel, "C", "B").equals( + nisl.make_map(""" + [N, M] -> { + [io_after, ii_after, j_after] -> + [io_before, ii_before, j_before] : + 0 <= 4*io_after + ii_after < N and + 0 <= ii_after < 4 and + 0 <= j_after < M and + 0 <= 4*io_before + ii_before < N and + 0 <= ii_before < 4 and + 0 <= j_before < M and + (4*io_before + ii_before < 4*io_after + ii_after or + (4*io_before + ii_before = 4*io_after + ii_after and + j_before <= j_after)) + } + """) + ) + assert _get_precise_relation(kernel, "D", "C").equals( + nisl.make_map(""" + [N, M] -> { + [j_after] -> [io_before, ii_before, j_before] : + 0 <= 4*io_before + ii_before < N and + 0 <= ii_before < 4 and + 0 <= j_before <= j_after < M + } + """) + ) + assert _get_precise_relation(kernel, "A", "A").equals(a_self) + assert _get_precise_relation(kernel, "D", "D").equals(d_self) + + +def test_affine_happens_after_transform_skew() -> None: + t_unit = lp.make_kernel( + """ + [NI, NJ] -> { + [i, j] : 0 <= i < NI and 0 <= j < NJ + } + """, + "out[i, j] = i + j {id=S}", + ) + kernel = dep.add_lexicographic_happens_after(t_unit).default_entrypoint + kernel = dep.apply_affine_transform_to_happens_afters( + kernel, + nisl.make_map(""" + { [i, j] -> [is, js] : is = i and js = i + j } + """), + ) + + assert _get_precise_relation(kernel, "S", "S").equals( + nisl.make_map(""" + [NI, NJ] -> { + [is_after, js_after] -> [is_before, js_before] : + 0 <= is_after < NI and + 0 <= js_after - is_after < NJ and + 0 <= is_before < NI and + 0 <= js_before - is_before < NJ and + (is_before < is_after or + (is_before = is_after and + js_before - is_before < js_after - is_after)) + } + """) + ) + + +def test_affine_happens_after_transform_permuted_axes() -> None: + t_unit = lp.make_kernel( + """ + [NI, NJ, NK] -> { + [i, j, k] : + 0 <= i < NI and 0 <= j < NJ and 0 <= k < NK + } + """, + "out[i, j, k] = i + j + k {id=S}", + ) + kernel = dep.add_lexicographic_happens_after(t_unit).default_entrypoint + kernel = dep.apply_affine_transform_to_happens_afters( + kernel, + nisl.make_map(""" + { + [i, j, k] -> [j_new, i_new, k_new] : + j_new = j and i_new = i and k_new = k + } + """), + ) + + assert _get_precise_relation(kernel, "S", "S").equals( + nisl.make_map(""" + [NI, NJ, NK] -> { + [j_new_after, i_new_after, k_new_after] -> + [j_new_before, i_new_before, k_new_before] : + 0 <= i_new_after < NI and + 0 <= j_new_after < NJ and + 0 <= k_new_after < NK and + 0 <= i_new_before < NI and + 0 <= j_new_before < NJ and + 0 <= k_new_before < NK and + (i_new_before < i_new_after or + (i_new_before = i_new_after and + j_new_before < j_new_after) or + (i_new_before = i_new_after and + j_new_before = j_new_after and + k_new_before < k_new_after)) + } + """) + ) + + +def test_affine_happens_after_transform_scalar_endpoints() -> None: + t_unit = lp.make_kernel( + "[N] -> { [i] : 0 <= i < N }", + """ + a[0] = 1 {id=A} + b[i] = 2 {id=B} + c[0] = 3 {id=C} + """, + [ + lp.GlobalArg("a,c", shape=(1,)), + lp.GlobalArg("b", shape=(var("N"),)), + "...", + ], + ) + kernel = dep.add_lexicographic_happens_after(t_unit).default_entrypoint + kernel = dep.apply_affine_transform_to_happens_afters( + kernel, + nisl.make_map("{ [i] -> [ip = i + 2] }"), + ) + + assert _get_precise_relation(kernel, "B", "A").equals( + nisl.make_map(""" + [N] -> { [ip_after] -> [] : 2 <= ip_after < N + 2 } + """) + ) + assert _get_precise_relation(kernel, "C", "B").equals( + nisl.make_map(""" + [N] -> { [] -> [ip_before] : 2 <= ip_before < N + 2 } + """) + ) + + +def test_affine_happens_after_transform_rejects_partial_nest() -> None: + t_unit = lp.make_kernel( + [ + "[NI] -> { [i] : 0 <= i < NI }", + "[NJ] -> { [j] : 0 <= j < NJ }", + ], + "out[i] = i {id=S}", + ) + kernel = dep.add_lexicographic_happens_after(t_unit).default_entrypoint + + with pytest.raises(LoopyError): + dep.apply_affine_transform_to_happens_afters( + kernel, + nisl.make_map(""" + { + [i, j] -> [i_new, j_new] : + i_new = i and j_new = j + } + """), + ) + + +def test_affine_happens_after_transform_avoids_proxy_name_collisions() -> None: + t_unit = lp.make_kernel( + "[NI, NJ] -> { [i, j] : 0 <= i < NI and 0 <= j < NJ }", + "out[i, j] = i + j {id=S}", + ) + kernel = dep.add_lexicographic_happens_after(t_unit).default_entrypoint + kernel = dep.apply_affine_transform_to_happens_afters( + kernel, + nisl.make_map("{ [i] -> [j_new_] : j_new_ = i }"), + ) + + assert _get_precise_relation(kernel, "S", "S").equals( + nisl.make_map(""" + [NI, NJ] -> { + [j_new__after, j_after] -> [j_new__before, j_before] : + 0 <= j_new__after < NI and 0 <= j_after < NJ and + 0 <= j_new__before < NI and 0 <= j_before < NJ and + (j_new__before < j_new__after or + (j_new__before = j_new__after and j_before < j_after)) + } + """) + ) + + +def test_map_domain_transforms_precise_happens_after() -> None: + t_unit = lp.make_kernel( + "[NI, NJ] -> { [i, j] : 0 <= i < NI and 0 <= j < NJ }", + """ + a[i, j] = i + j {id=A} + b[i, j] = a[i, j] {id=B} + """, + ) + t_unit = dep.add_lexicographic_happens_after(t_unit) + transform_map = nisl.make_map(""" + [NI] -> { + [i] -> [io, ii] : + i = 4*io + ii and 0 <= ii < 4 + } + """) + + expected = dep.apply_affine_transform_to_happens_afters( + t_unit.default_entrypoint, + transform_map, + ) + kernel = lp.map_domain(t_unit, transform_map).default_entrypoint + + assert kernel.id_to_insn["A"].within_inames == frozenset({"io", "ii", "j"}) + assert kernel.id_to_insn["B"].within_inames == frozenset({"io", "ii", "j"}) + for sink_id, source_id in [("A", "A"), ("B", "A"), ("B", "B")]: + assert _get_precise_relation(kernel, sink_id, source_id).equals( + _get_precise_relation(expected, sink_id, source_id) + ) + + +def test_split_iname_transforms_precise_happens_after() -> None: + t_unit = lp.make_kernel( + "[NI, NJ] -> { [i, j] : 0 <= i < NI and 0 <= j < NJ }", + """ + a[i, j] = i + j {id=A} + b[i, j] = a[i, j] {id=B} + """, + ) + t_unit = dep.add_lexicographic_happens_after(t_unit) + split_reln = nisl.make_map(""" + { [i] -> [io, ii] : + i = 4*io + ii and 0 <= ii < 4 + } + """) + + expected = dep.apply_affine_transform_to_happens_afters( + t_unit.default_entrypoint, split_reln + ) + kernel = lp.split_iname( + t_unit, "i", 4, outer_iname="io", inner_iname="ii" + ).default_entrypoint + + assert kernel.id_to_insn["A"].within_inames == frozenset({"io", "ii", "j"}) + assert kernel.id_to_insn["B"].within_inames == frozenset({"io", "ii", "j"}) + for sink_id, source_id in [("A", "A"), ("B", "A"), ("B", "B")]: + assert _get_precise_relation(kernel, sink_id, source_id).equals( + _get_precise_relation(expected, sink_id, source_id) + ) + + +def test_split_iname_rejects_within_for_precise_happens_after() -> None: + t_unit = lp.make_kernel( + "{ [i] : 0 <= i < 16 }", + """ + a[i] = i {id=A} + b[i] = a[i] {id=B} + """, + ) + t_unit = dep.add_lexicographic_happens_after(t_unit) + + with pytest.raises(LoopyError, match="does not support 'within'"): + lp.split_iname(t_unit, "i", 4, within="id:A") + + +def test_splice_happens_after_as_consumer_inherits_branched_order() -> None: + t_unit = lp.make_kernel( + [ + "[NI, NJ] -> { [i, j] : 0 <= i < NI and 0 <= j < NJ }", + """ + [NI, NJ] -> { + [tile, lane, jc] : + 0 <= 2*tile + lane < NI and + 0 <= lane < 2 and 0 <= jc < NJ + } + """, + ], + """ + p[i, j] = i + j {id=P} + q[i] = i {id=Q} + a[i, j] = p[i, j] + q[i] {id=A} + c[tile, lane, jc] = 0 {id=C} + s[i, j] = a[i, j] {id=S} + """, + ) + kernel = dep.add_lexicographic_happens_after(t_unit).default_entrypoint + + a_after_p = HappensAfter(nisl.make_map(""" + [NI, NJ] -> { + [i_after, j_after] -> [i_before, j_before] : + i_before = i_after and j_before = j_after and + 0 <= i_after < NI and 0 <= j_after < NJ + } + """)) + a_after_q = HappensAfter(nisl.make_map(""" + [NI, NJ] -> { + [i_after, j_after] -> [i_before] : + i_before = i_after and + 0 <= i_after < NI and 0 <= j_after < NJ + } + """)) + s_after_a = HappensAfter(nisl.make_map(""" + [NI, NJ] -> { + [i_after, j_after] -> [i_before, j_before] : + i_before = i_after and j_before = j_after and + 0 <= i_after < NI and 0 <= j_after < NJ + } + """)) + + new_happens_after = { + "P": {"P": kernel.id_to_insn["P"].happens_after["P"]}, + "Q": {"Q": kernel.id_to_insn["Q"].happens_after["Q"]}, + "A": { + "A": kernel.id_to_insn["A"].happens_after["A"], + "P": a_after_p, + "Q": a_after_q, + }, + "C": {"C": kernel.id_to_insn["C"].happens_after["C"]}, + "S": { + "S": kernel.id_to_insn["S"].happens_after["S"], + "A": s_after_a, + }, + } + kernel = kernel.copy(instructions=tuple( + stmt.copy(happens_after=new_happens_after[stmt.id]) + for stmt in kernel.instructions + )) + original_anchor_order = kernel.id_to_insn["A"].happens_after + original_successor_order = kernel.id_to_insn["S"].happens_after + original_consumer_self_order = kernel.id_to_insn["C"].happens_after["C"] + + kernel = dep.splice_happens_after_as_consumer( + kernel, + "C", + "A", + nisl.make_map(""" + [NI, NJ] -> { + [tile_after, lane_after, jc_after] -> [i_before, j_before] : + i_before = 2*tile_after + lane_after and + j_before = jc_after and + 0 <= 2*tile_after + lane_after < NI and + 0 <= lane_after < 2 and 0 <= jc_after < NJ + } + """), + ) + + assert kernel.id_to_insn["C"].happens_after.keys() == {"C", "P", "Q"} + assert kernel.id_to_insn["C"].happens_after["C"] == ( + original_consumer_self_order + ) + assert kernel.id_to_insn["A"].happens_after == original_anchor_order + assert kernel.id_to_insn["S"].happens_after == original_successor_order + assert _get_precise_relation(kernel, "C", "P").equals( + nisl.make_map(""" + [NI, NJ] -> { + [tile_after, lane_after, jc_after] -> [i_before, j_before] : + i_before = 2*tile_after + lane_after and + j_before = jc_after and + 0 <= 2*tile_after + lane_after < NI and + 0 <= lane_after < 2 and 0 <= jc_after < NJ + } + """) + ) + assert _get_precise_relation(kernel, "C", "Q").equals( + nisl.make_map(""" + [NI, NJ] -> { + [tile_after, lane_after, jc_after] -> [i_before] : + i_before = 2*tile_after + lane_after and + 0 <= 2*tile_after + lane_after < NI and + 0 <= lane_after < 2 and 0 <= jc_after < NJ + } + """) + ) + + +def test_splice_happens_after_as_consumer_unions_existing_order() -> None: + t_unit = lp.make_kernel( + [ + "[N] -> { [i] : 0 <= i < N }", + "[N] -> { [ic, lane] : 0 <= ic < N and 0 <= lane < 4 }", + ], + """ + p[i] = i {id=P} + a[i] = p[i] {id=A} + c[ic, lane] = 0 {id=C} + """, + ) + kernel = dep.add_lexicographic_happens_after(t_unit).default_entrypoint + a_after_p = HappensAfter(nisl.make_map(""" + [N] -> { + [i_after] -> [i_before = i_after] : 0 <= i_after < N + } + """)) + c_after_p = HappensAfter(nisl.make_map(""" + [N] -> { + [ic_after, lane_after] -> [i_before = ic_after] : + 0 <= ic_after < N and lane_after = 0 + } + """)) + new_happens_after = { + "P": {"P": kernel.id_to_insn["P"].happens_after["P"]}, + "A": { + "A": kernel.id_to_insn["A"].happens_after["A"], + "P": a_after_p, + }, + "C": { + "C": kernel.id_to_insn["C"].happens_after["C"], + "P": c_after_p, + }, + } + kernel = kernel.copy(instructions=tuple( + stmt.copy(happens_after=new_happens_after[stmt.id]) + for stmt in kernel.instructions + )) + + kernel = dep.splice_happens_after_as_consumer( + kernel, + "C", + "A", + nisl.make_map(""" + [N] -> { + [ic_after, lane_after] -> [i_before = ic_after] : + 0 <= ic_after < N and 1 <= lane_after < 4 + } + """), + ) + + assert _get_precise_relation(kernel, "C", "P").equals( + nisl.make_map(""" + [N] -> { + [ic_after, lane_after] -> [i_before = ic_after] : + 0 <= ic_after < N and 0 <= lane_after < 4 + } + """) + ) + + +def test_splice_happens_after_as_consumer_handles_scalar_anchor() -> None: + t_unit = lp.make_kernel( + "[N] -> { [i] : 0 <= i < N }", + """ + p[0] = 1 {id=P} + a[0] = p[0] {id=A} + c[i] = 0 {id=C} + """, + ) + kernel = dep.add_lexicographic_happens_after(t_unit).default_entrypoint + kernel = kernel.copy(instructions=tuple( + stmt.copy(happens_after={ + "P": {"P": kernel.id_to_insn["P"].happens_after["P"]}, + "A": { + "A": kernel.id_to_insn["A"].happens_after["A"], + "P": HappensAfter(nisl.make_map("{ [] -> [] }")), + }, + "C": {"C": kernel.id_to_insn["C"].happens_after["C"]}, + }[stmt.id]) + for stmt in kernel.instructions + )) + + kernel = dep.splice_happens_after_as_consumer( + kernel, + "C", + "A", + nisl.make_map("[N] -> { [i_after] -> [] : 0 <= i_after < N }"), + ) + + assert _get_precise_relation(kernel, "C", "P").equals( + nisl.make_map("[N] -> { [i_after] -> [] : 0 <= i_after < N }") + ) + + +def test_splice_happens_after_as_producer_redirects_branched_order() -> None: + t_unit = lp.make_kernel( + [ + "[NI, NJ] -> { [i, j] : 0 <= i < NI and 0 <= j < NJ }", + """ + [NI, NJ] -> { + [tile, lane, jp] : + 0 <= 2*tile + lane < NI and + 0 <= lane < 2 and 0 <= jp < NJ + } + """, + ], + """ + q[i, j] = i + j {id=Q} + a[i, j] = q[i, j] {id=A} + g[tile, lane, jp] = 0 {id=G} + s[i, j] = a[i, j] + q[i, j] {id=S} + t[i] = sum(j, a[i, j]) {id=T} + """, + ) + kernel = t_unit.default_entrypoint + + a_after_q = HappensAfter(nisl.make_map(""" + [NI, NJ] -> { + [i_after, j_after] -> [i_before, j_before] : + i_before = i_after and j_before = j_after and + 0 <= i_after < NI and 0 <= j_after < NJ + } + """)) + g_after_q = HappensAfter(nisl.make_map(""" + [NI, NJ] -> { + [tile_after, lane_after, jp_after] -> [i_before, j_before] : + i_before = 2*tile_after + lane_after and + j_before = jp_after and + 0 <= 2*tile_after + lane_after < NI and + 0 <= lane_after < 2 and 0 <= jp_after < NJ + } + """)) + s_after_a = HappensAfter(nisl.make_map(""" + [NI, NJ] -> { + [i_after, j_after] -> [i_before, j_before] : + i_before = i_after and j_before = j_after and + 0 <= i_after < NI and 0 <= j_after < NJ + } + """)) + s_after_q = HappensAfter(nisl.make_map(""" + [NI, NJ] -> { + [i_after, j_after] -> [i_before, j_before] : + i_before = i_after and j_before = j_after and + 0 <= i_after < NI and 0 <= j_after < NJ + } + """)) + t_after_a = HappensAfter(nisl.make_map(""" + [NI, NJ] -> { + [i_after] -> [i_before, j_before] : + i_before = i_after and + 0 <= i_after < NI and 0 <= j_before < NJ + } + """)) + happens_after = { + "Q": {}, + "A": {"Q": a_after_q}, + "G": {"Q": g_after_q}, + "S": {"A": s_after_a, "Q": s_after_q}, + "T": {"A": t_after_a}, + } + kernel = kernel.copy(instructions=tuple( + stmt.copy(happens_after=happens_after[stmt.id]) + for stmt in kernel.instructions + )) + original_anchor_order = kernel.id_to_insn["A"].happens_after + original_producer_order = kernel.id_to_insn["G"].happens_after + original_s_after_q = kernel.id_to_insn["S"].happens_after["Q"] + + kernel = dep.splice_happens_after_as_producer( + kernel, + "G", + "A", + nisl.make_map(""" + [NI, NJ] -> { + [i_after, j_after] -> [tile_before, lane_before, jp_before] : + i_after = 2*tile_before + lane_before and + j_after = jp_before and + 0 <= i_after < NI and 0 <= j_after < NJ and + 0 <= lane_before < 2 + } + """), + ) + + assert kernel.id_to_insn["A"].happens_after == original_anchor_order + assert kernel.id_to_insn["G"].happens_after == original_producer_order + assert kernel.id_to_insn["S"].happens_after.keys() == {"G", "Q"} + assert kernel.id_to_insn["S"].happens_after["Q"] == original_s_after_q + assert kernel.id_to_insn["T"].happens_after.keys() == {"G"} + assert _get_precise_relation(kernel, "S", "G").equals( + nisl.make_map(""" + [NI, NJ] -> { + [i_after, j_after] -> [tile_before, lane_before, jp_before] : + i_after = 2*tile_before + lane_before and + j_after = jp_before and + 0 <= i_after < NI and 0 <= j_after < NJ and + 0 <= lane_before < 2 + } + """) + ) + assert _get_precise_relation(kernel, "T", "G").equals( + nisl.make_map(""" + [NI, NJ] -> { + [i_after] -> [tile_before, lane_before, jp_before] : + i_after = 2*tile_before + lane_before and + 0 <= i_after < NI and 0 <= lane_before < 2 and + 0 <= jp_before < NJ + } + """) + ) + + +def test_splice_happens_after_as_producer_preserves_unmapped_order() -> None: + t_unit = lp.make_kernel( + [ + "[N] -> { [i] : 0 <= i < N }", + "[N] -> { [ip] : 0 <= ip < N }", + ], + """ + a[i] = i {id=A} + g[ip] = 0 {id=G} + s[i] = a[i] {id=S} + """, + ) + kernel = t_unit.default_entrypoint + s_after_a = HappensAfter(nisl.make_map(""" + [N] -> { + [i_after] -> [i_before = i_after] : 0 <= i_after < N + } + """)) + s_after_g = HappensAfter(nisl.make_map(""" + [N] -> { + [i_after] -> [ip_before = i_after] : + 0 <= i_after < N and 2*i_after >= N + } + """)) + kernel = kernel.copy(instructions=tuple( + stmt.copy(happens_after={ + "A": {}, + "G": {}, + "S": {"A": s_after_a, "G": s_after_g}, + }[stmt.id]) + for stmt in kernel.instructions + )) + + kernel = dep.splice_happens_after_as_producer( + kernel, + "G", + "A", + nisl.make_map(""" + [N] -> { + [i_after] -> [ip_before = 0] : + 0 <= i_after < N and 2*i_after < N + } + """), + ) + + assert _get_precise_relation(kernel, "S", "A").equals( + nisl.make_map(""" + [N] -> { + [i_after] -> [i_before = i_after] : + 0 <= i_after < N and 2*i_after >= N + } + """) + ) + assert _get_precise_relation(kernel, "S", "G").equals( + nisl.make_map(""" + [N] -> { + [i_after] -> [ip_before = 0] : + 0 <= i_after < N and 2*i_after < N; + [i_after] -> [ip_before = i_after] : + 0 <= i_after < N and 2*i_after >= N + } + """) + ) + + +def test_splice_happens_after_as_producer_handles_scalar_producer() -> None: + t_unit = lp.make_kernel( + "[N] -> { [i] : 0 <= i < N }", + """ + a[i] = i {id=A} + g[0] = 0 {id=G} + s[i] = a[i] {id=S} + """, + ) + kernel = t_unit.default_entrypoint + kernel = kernel.copy(instructions=tuple( + stmt.copy(happens_after={ + "A": {}, + "G": {}, + "S": { + "A": HappensAfter(nisl.make_map(""" + [N] -> { + [i_after] -> [i_before = i_after] : + 0 <= i_after < N + } + """)), + }, + }[stmt.id]) + for stmt in kernel.instructions + )) + + kernel = dep.splice_happens_after_as_producer( + kernel, + "G", + "A", + nisl.make_map("[N] -> { [i_after] -> [] : 0 <= i_after < N }"), + ) + + assert kernel.id_to_insn["S"].happens_after.keys() == {"G"} + assert _get_precise_relation(kernel, "S", "G").equals( + nisl.make_map("[N] -> { [i_after] -> [] : 0 <= i_after < N }") + ) + + +def test_splice_happens_after_as_consumer_and_producer() -> None: + t_unit = lp.make_kernel( + [ + "[NI, NJ] -> { [i, j] : 0 <= i < NI and 0 <= j < NJ }", + """ + [NI, NJ] -> { + [tile, lane, jg] : + 0 <= 2*tile + lane < NI and + 0 <= lane < 2 and 0 <= jg < NJ + } + """, + ], + """ + d[i, j] = i + j {id=D} + a[i, j] = d[i, j] {id=A} + g[tile, lane, jg] = 0 {id=G} + s[i, j] = a[i, j] {id=S} + """, + ) + kernel = t_unit.default_entrypoint + same_instance = nisl.make_map(""" + [NI, NJ] -> { + [i_after, j_after] -> [i_before, j_before] : + i_before = i_after and j_before = j_after and + 0 <= i_after < NI and 0 <= j_after < NJ + } + """) + kernel = kernel.copy(instructions=tuple( + stmt.copy(happens_after={ + "D": {}, + "A": {"D": HappensAfter(same_instance)}, + "G": {}, + "S": {"A": HappensAfter(same_instance)}, + }[stmt.id]) + for stmt in kernel.instructions + )) + + instruction_to_anchor = nisl.make_map(""" + [NI, NJ] -> { + [tile_after, lane_after, jg_after] -> [i_before, j_before] : + i_before = 2*tile_after + lane_after and + j_before = jg_after and + 0 <= 2*tile_after + lane_after < NI and + 0 <= lane_after < 2 and 0 <= jg_after < NJ + } + """) + anchor_to_instruction = nisl.make_map(""" + [NI, NJ] -> { + [i_after, j_after] -> [tile_before, lane_before, jg_before] : + i_after = 2*tile_before + lane_before and + j_after = jg_before and + 0 <= i_after < NI and 0 <= j_after < NJ and + 0 <= lane_before < 2 + } + """) + kernel = dep.splice_happens_after_as_consumer_and_producer( + kernel, + "G", + "A", + instruction_to_anchor, + anchor_to_instruction, + ) + + assert kernel.id_to_insn["A"].happens_after.keys() == {"D"} + assert kernel.id_to_insn["G"].happens_after.keys() == {"D"} + assert kernel.id_to_insn["S"].happens_after.keys() == {"G"} + assert _get_precise_relation(kernel, "G", "D").equals( + instruction_to_anchor + ) + assert _get_precise_relation(kernel, "S", "G").equals( + anchor_to_instruction + ) + + +def test_combined_splice_accepts_independent_partial_maps() -> None: + t_unit = lp.make_kernel( + [ + "[N] -> { [i] : 0 <= i < N }", + "[N] -> { [b] : 0 <= b < N }", + ], + """ + d[i] = i {id=D} + a[i] = d[i] {id=A} + g[b] = 0 {id=G} + s[i] = a[i] {id=S} + """, + ) + kernel = t_unit.default_entrypoint + same_instance = HappensAfter(nisl.make_map(""" + [N] -> { + [i_after] -> [i_before = i_after] : 0 <= i_after < N + } + """)) + kernel = kernel.copy(instructions=tuple( + stmt.copy(happens_after={ + "D": {}, + "A": {"D": same_instance}, + "G": {}, + "S": {"A": same_instance}, + }[stmt.id]) + for stmt in kernel.instructions + )) + + kernel = dep.splice_happens_after_as_consumer_and_producer( + kernel, + "G", + "A", + nisl.make_map(""" + [N] -> { + [b_after] -> [i_before = 2*b_after] : + 0 <= 2*b_after < N + } + """), + nisl.make_map(""" + [N] -> { + [i_after] -> [b_before = i_after] : + 0 <= i_after < N and 2*i_after >= N + } + """), + ) + + assert _get_precise_relation(kernel, "G", "D").equals( + nisl.make_map(""" + [N] -> { + [b_after] -> [i_before = 2*b_after] : + 0 <= 2*b_after < N + } + """) + ) + assert _get_precise_relation(kernel, "S", "A").equals( + nisl.make_map(""" + [N] -> { + [i_after] -> [i_before = i_after] : + 0 <= i_after < N and 2*i_after < N + } + """) + ) + assert _get_precise_relation(kernel, "S", "G").equals( + nisl.make_map(""" + [N] -> { + [i_after] -> [b_before = i_after] : + 0 <= i_after < N and 2*i_after >= N + } + """) + ) + + +def test_access_relation_finder_tracks_reads_and_writes_per_statement() -> None: + t_unit = lp.make_kernel( + "{ [i] : 1 <= i < N }", + """ + a[i] = b[i - 1] + c[i] {id=S} + b[i] = a[i] {id=T} + """, + ) + + kernel = t_unit.default_entrypoint + rel_find = dep.AccessRelationFinder(kernel) + for stmt in kernel.instructions: + assert isinstance(stmt, lp.MultiAssignmentBase) + for assignee in stmt.assignees: + rel_find(assignee, stmt.id, dep.AccessType.write) + rel_find(stmt.expression, stmt.id, dep.AccessType.read) + + assert rel_find.read_relations["S"].keys() == {"b", "c"} + assert rel_find.write_relations["S"].keys() == {"a"} + assert rel_find.read_relations["T"].keys() == {"a"} + assert rel_find.write_relations["T"].keys() == {"b"} + assert rel_find.read_relations["S"]["b"].equals( + nisl.make_map("[N] -> { [i] -> [ax_0 = i - 1] : 1 <= i < N }") + ) + assert rel_find.write_relations["S"]["a"].equals( + nisl.make_map("[N] -> { [i] -> [ax_0 = i] : 1 <= i < N }") + ) + + +def test_access_relation_finder_handles_value_arg_subscript() -> None: + t_unit = lp.make_kernel( + "{ [i] : 0 <= i < N }", + "out[i] = a[offset + i] {id=S}", + [ + lp.GlobalArg("a", shape=("N + offset",)), + lp.ValueArg("offset", np.int32), + "...", + ], + ) + + kernel = t_unit.default_entrypoint + insn = kernel.id_to_insn["S"] + rel_find = dep.AccessRelationFinder(kernel) + rel_find(insn.expression, insn.id, dep.AccessType.read) + + assert rel_find.read_relations["S"]["a"].equals(nisl.make_map(""" + [N, offset] -> { + [i] -> [ax_0 = offset + i] : 0 <= i < N + } + """)) + + +def test_access_relation_names_do_not_clash_with_inames() -> None: + t_unit = lp.make_kernel( + """ + { [ax_0, ax_1] : + 0 <= ax_0 < N and 0 <= ax_1 < M + } + """, + "out[ax_0, ax_1] = a[ax_0] + b[ax_1] {id=S}", + ) + + kernel = t_unit.default_entrypoint + insn = kernel.id_to_insn["S"] + rel_find = dep.AccessRelationFinder(kernel) + rel_find(insn.expression, insn.id, dep.AccessType.read) + + a_cell_names = rel_find.read_relations["S"]["a"].space.dimtype_to_names[ + nisl.DimType.out + ] + b_cell_names = rel_find.read_relations["S"]["b"].space.dimtype_to_names[ + nisl.DimType.out + ] + assert a_cell_names == b_cell_names + assert set(a_cell_names).isdisjoint(kernel.all_variable_names()) + + +def test_access_relation_finder_handles_reduction() -> None: + t_unit = lp.make_kernel( + """ + { [i, j, k] : + 0 <= i < NI and 0 <= j < NJ and 0 <= k < NK + } + """, + "out[i, j] = sum(k, a[i, j, k]) {id=S}", + ) + + kernel = t_unit.default_entrypoint + insn = kernel.id_to_insn["S"] + rel_find = dep.AccessRelationFinder(kernel) + rel_find(insn.expression, insn.id, dep.AccessType.read) + + assert rel_find.read_relations["S"]["a"].equals( + nisl.make_map(""" + [NI, NJ, NK] -> { + [i, j] -> [ax_0 = i, ax_1 = j, ax_2] : + 0 <= i < NI and + 0 <= j < NJ and + 0 <= ax_2 < NK + } + """) + ) + + +def test_access_relation_finder_handles_sub_array_ref() -> None: + t_unit = lp.make_kernel( + """ + { [i, j, k] : + 0 <= i < NI and 0 <= j < NJ and 0 <= k < NK + } + """, + "out[i, j] = a[i, j, 0] {id=S}", + ) + + kernel = t_unit.default_entrypoint + rel_find = dep.AccessRelationFinder(kernel) + # Build the swept access directly so this remains a mapper unit test and + # does not require setting up an array-valued callable. + sub_array_ref = SubArrayRef( + (var("k"),), + var("a")[var("i"), var("j"), var("k")], + ) + rel_find(sub_array_ref, "S", dep.AccessType.read) + + assert rel_find.read_relations["S"]["a"].equals( + nisl.make_map(""" + [NI, NJ, NK] -> { + [i, j] -> [ax_0 = i, ax_1 = j, ax_2] : + 0 <= i < NI and + 0 <= j < NJ and + 0 <= ax_2 < NK + } + """) + ) + + +def _relax_strict_happens_after( + instructions: str, + domain: str = "{ [i] : 0 <= i < N }", +) -> lp.LoopKernel: + t_unit = lp.make_kernel(domain, instructions) + t_unit = dep.add_lexicographic_happens_after(t_unit) + return dep.relax_strict_happens_after(t_unit).default_entrypoint + + +def test_relax_strict_happens_after_expands_substitution_rules() -> None: + t_unit = lp.make_kernel( + "{ [i] : 0 <= i < N }", + """ + a[2*i + 3] = i {id=S} + out[i] = outer(i) {id=T} + """, + substitutions={ + "inner": lp.SubstitutionRule( + "inner", ("j",), var("a")[2 * var("j") + 1] + ), + "outer": lp.SubstitutionRule( + "outer", ("k",), var("inner")(var("k") + 1) + ), + }, + ) + t_unit = dep.add_lexicographic_happens_after(t_unit) + kernel = dep.relax_strict_happens_after(t_unit).default_entrypoint + + required_order = kernel.id_to_insn["T"].happens_after["S"].instances_rel + assert required_order is not None + assert required_order.equals( + nisl.make_map(""" + [N] -> { + [i_after] -> [i_before = i_after] : 0 <= i_after < N + } + """) + ) + + +@pytest.mark.parametrize( + ("source", "sink"), + ( + ("a[i, j] = 1", "b[i, j] = a[i, j]"), + ("a[i, j] = 1", "a[i, j] = 2"), + ("b[i, j] = a[i, j]", "a[i, j] = 2"), + ), + ids=("read-after-write", "write-after-write", "write-after-read"), +) +def test_relax_strict_happens_after_finds_direct_hazards( + source: str, sink: str +) -> None: + kernel = _relax_strict_happens_after( + f""" + {source} {{id=S}} + {sink} {{id=T}} + """, + "{ [i, j] : 0 <= i < NI and 0 <= j < NJ }", + ) + + required_order = kernel.id_to_insn["T"].happens_after["S"].instances_rel + assert required_order is not None + assert required_order.equals( + nisl.make_map(""" + [NI, NJ] -> { + [i_after, j_after] -> + [i_before = i_after, j_before = j_after] : + 0 <= i_after < NI and 0 <= j_after < NJ + } + """) + ) + + +def test_has_precise_dependencies() -> None: + legacy_kernel = lp.make_kernel( + "{ [i] : 0 <= i < N }", + """ + a[i] = i {id=S} + b[i] = a[i] {id=T, dep=S} + """, + ).default_entrypoint + assert not dep.has_precise_dependencies(legacy_kernel) + + precise_kernel = dep.add_lexicographic_happens_after(legacy_kernel) + assert dep.has_precise_dependencies(precise_kernel) + + t_insn = precise_kernel.id_to_insn["T"] + mixed_happens_after = dict(t_insn.happens_after) + mixed_happens_after["S"] = HappensAfter(instances_rel=None) + mixed_kernel = precise_kernel.copy( + instructions=tuple( + insn.copy(happens_after=mixed_happens_after) + if insn.id == "T" + else insn + for insn in precise_kernel.instructions + ) + ) + + with pytest.raises(LoopyError, match="mixes precise and legacy"): + dep.has_precise_dependencies(mixed_kernel) + + +def test_relax_strict_happens_after_rejects_c_instruction() -> None: + t_unit = lp.make_kernel( + "{ [i] : 0 <= i < N }", + [ + lp.CInstruction( + "i", + "tmp[i] = 1;", + assignees=("tmp[i]",), + id="C", + ), + ], + temporary_variables={ + "tmp": lp.TemporaryVariable( + "tmp", dtype=np.float32, shape=("N",) + ), + }, + ) + t_unit = dep.add_lexicographic_happens_after(t_unit) + + with pytest.raises(LoopyError, match="does not support CInstruction 'C'"): + dep.relax_strict_happens_after(t_unit) + + +def test_relax_strict_happens_after_rejects_base_storage() -> None: + t_unit = lp.make_kernel( + "{ [i] : 0 <= i < N }", + """ + a[i] = i {id=A} + b[i] = a[i] {id=B} + """, + temporary_variables={ + name: lp.TemporaryVariable( + name, + dtype=np.float32, + shape=("N",), + base_storage="base", + ) + for name in ("a", "b") + }, + ) + t_unit = dep.add_lexicographic_happens_after(t_unit) + + with pytest.raises( + LoopyError, match="temporary 'a' with base_storage" + ): + dep.relax_strict_happens_after(t_unit) + + +def test_relax_strict_happens_after_tracks_scalar_accesses() -> None: + kernel = _relax_strict_happens_after( + """ + <> tmp = i {id=S} + out[i] = tmp {id=T} + """, + "{ [i] : 0 <= i < N }", + ) + + required_order = kernel.id_to_insn["T"].happens_after["S"].instances_rel + assert required_order is not None + assert required_order.equals( + nisl.make_map(""" + [N] -> { + [i_after] -> [i_before = i_after] : + 0 <= i_after < N + } + """) + ) + + +def test_relax_strict_happens_after_tracks_scalar_statements() -> None: + kernel = _relax_strict_happens_after( + """ + <> tmp = 1 {id=S} + out = tmp {id=T} + """, + "{ : }", + ) + + required_order = kernel.id_to_insn["T"].happens_after["S"].instances_rel + assert required_order is not None + assert required_order.equals(nisl.make_map("{ [] -> [] }")) + + +def test_relax_strict_happens_after_selects_writers_per_cell() -> None: + kernel = _relax_strict_happens_after( + """ + x[i, j] = i + j {id=S} + out[i, j] = x[i, j] + x[i, j - 1] {id=T} + """, + """ + { + [i, j] : + 0 <= i < NI and 0 <= j < NJ + } + """, + ) + + required_order = kernel.id_to_insn["T"].happens_after["S"].instances_rel + assert required_order is not None + assert required_order.equals( + nisl.make_map(""" + [NI, NJ] -> { + [j_after, i_after] -> + [i_before = i_after, j_before] : + 0 <= i_after < NI and + 0 <= j_after < NJ and + 0 <= j_before < NJ and + j_after - 1 <= j_before <= j_after + } + """) + ) + + +def test_relax_strict_happens_after_selects_live_readers() -> None: + t_unit = lp.make_kernel( + "{ [i] : 0 <= i < N }", + """ + a[i + 1] = a[i] + 1 {id=S} + a[1] = i {id=T} + """, + [ + lp.GlobalArg("a", dtype=np.int32, shape=("N + 1",)), + "...", + ], + ) + t_unit = dep.add_lexicographic_happens_after(t_unit) + kernel = dep.relax_strict_happens_after(t_unit).default_entrypoint + + required_order = kernel.id_to_insn["T"].happens_after["S"].instances_rel + assert required_order is not None + assert required_order.equals( + nisl.make_map(""" + [N] -> { + [i_after] -> [i_before = i_after] : + 0 <= i_after <= 1 and i_after < N + } + """) + ) + self_order = kernel.id_to_insn["T"].happens_after["T"].instances_rel + assert self_order is not None + assert self_order.equals(nisl.make_map(""" + [N] -> { + [i_after] -> [i_before = i_after - 1] : + 0 < i_after < N + } + """)) + + +@pytest.mark.parametrize( + "instructions", + ( + """ + a[i] = 1 {id=S} + b[i] = 2 {id=T} + """, + """ + a[i] = 1 {id=S} + b[i] = a[i + N] {id=T} + """, + ), + ids=("different-variables", "disjoint-footprints"), +) +def test_relax_strict_happens_after_drops_nonconflicting_edges( + instructions: str, +) -> None: + kernel = _relax_strict_happens_after(instructions) + + assert "S" not in kernel.id_to_insn["T"].happens_after + + +def test_relax_strict_happens_after_ignores_empty_coarse_edges() -> None: + t_unit = lp.make_kernel( + "{ [i] : 0 <= i < N }", + """ + a[i] = i {id=A} + b[i] = i {id=B} + """, + ) + kernel = dep.add_lexicographic_happens_after(t_unit).default_entrypoint + + cross_order = kernel.id_to_insn["B"].happens_after["A"].instances_rel + assert cross_order is not None + empty_order = cross_order - cross_order + + a_insn = kernel.id_to_insn["A"] + b_insn = kernel.id_to_insn["B"] + kernel = kernel.copy(instructions=( + a_insn.copy(happens_after={ + "A": a_insn.happens_after["A"], + "B": HappensAfter(empty_order), + }), + b_insn.copy(happens_after={ + "A": HappensAfter(empty_order), + "B": b_insn.happens_after["B"], + }), + )) + + kernel = dep.relax_strict_happens_after(kernel) + + assert all(not stmt.happens_after for stmt in kernel.instructions) + + +def test_relax_strict_happens_after_tracks_live_footprints_through_a_chain() -> ( + None +): + kernel = _relax_strict_happens_after( + """ + a[i, j] = 1 {id=A} + b[i, j] = 2 {id=B} + a[2*i, j] = 3 {id=C} + out[i, j] = a[i, j] {id=D} + """, + "{ [i, j] : 0 <= i < NI and 0 <= j < NJ }", + ) + + recent_order = kernel.id_to_insn["D"].happens_after["C"].instances_rel + fallback_order = kernel.id_to_insn["D"].happens_after["A"].instances_rel + assert recent_order is not None + assert fallback_order is not None + assert recent_order.equals( + nisl.make_map(""" + [NI, NJ] -> { + [i_after, j_after] -> [i_before, j_before] : + i_after = 2*i_before and + j_after = j_before and + 0 <= i_after < NI and + 0 <= i_before < NI and + 0 <= j_after < NJ + } + """) + ) + assert "B" not in kernel.id_to_insn["D"].happens_after + assert fallback_order.equals( + nisl.make_map(""" + [NI, NJ] -> { + [i_after, j_after] -> + [i_before = i_after, j_before = j_after] : + 0 < i_after < NI and + 0 <= j_after < NJ + } + """) + ) + + +def test_relax_strict_happens_after_selects_writers_globally() -> None: + t_unit = lp.make_kernel( + "{ : }", + """ + a = 0 {id=W0} + a = 1 {id=W1} + out = a {id=R} + """, + ) + kernel = dep.add_lexicographic_happens_after(t_unit).default_entrypoint + r_insn = kernel.id_to_insn["R"] + kernel = kernel.copy(instructions=tuple( + insn.copy(happens_after={ + **insn.happens_after, + "W0": r_insn.happens_after["W1"], + }) + if insn.id == "R" else insn + for insn in kernel.instructions + )) + + kernel = dep.relax_strict_happens_after(kernel) + + assert "W0" not in kernel.id_to_insn["R"].happens_after + required_order = kernel.id_to_insn["R"].happens_after["W1"].instances_rel + assert required_order is not None + assert required_order.equals(nisl.make_map("{ [] -> [] }")) + + +def test_relax_strict_happens_after_keeps_incomparable_writers() -> None: + t_unit = lp.make_kernel( + [ + "{ [q] : 0 <= q < 2 }", + "{ [p] : p = 0 }", + "{ [a_idx] : a_idx = 0 }", + ], + """ + x[0] = p {id=P} + x[q] = q {id=Q} + out[a_idx] = x[0] {id=A} + """, + ) + kernel = t_unit.default_entrypoint + q_after_p = HappensAfter(nisl.make_map(""" + { [q_after = 1] -> [p_before = 0] } + """)) + a_after_q = HappensAfter(nisl.make_map(""" + { [a_idx_after = 0] -> [q_before] : 0 <= q_before < 2 } + """)) + kernel = kernel.copy(instructions=tuple( + insn.copy(happens_after={ + "P": {}, + "Q": {"P": q_after_p}, + "A": {"Q": a_after_q}, + }[insn.id]) + for insn in kernel.instructions + )) + + kernel = dep.relax_strict_happens_after(kernel) + + a_happens_after = kernel.id_to_insn["A"].happens_after + q_order = a_happens_after["Q"].instances_rel + p_order = a_happens_after["P"].instances_rel + assert q_order is not None + assert p_order is not None + assert q_order.equals(nisl.make_map(""" + { [a_idx_after = 0] -> [q_before = 0] } + """)) + assert p_order.equals(nisl.make_map(""" + { [a_idx_after = 0] -> [p_before = 0] } + """)) + + +def test_relax_strict_happens_after_composes_user_supplied_relations() -> None: + t_unit = lp.make_kernel( + [ + "[N] -> { [i] : 0 <= i < 2*N }", + "[N] -> { [j] : 0 <= j < N }", + "[N] -> { [k] : 1 <= k < N }", + ], + """ + a[i] = 1 {id=S} + tmp[j] = 0 {id=T} + out[k] = a[2*k - 2] {id=U} + """, + ) + kernel = dep.add_lexicographic_happens_after(t_unit).default_entrypoint + t_after_s = HappensAfter( + instances_rel=nisl.make_map(""" + [N] -> { + [j_after] -> [i_before = 2*j_after] : + 0 <= j_after < N + } + """) + ) + u_after_t = HappensAfter( + instances_rel=nisl.make_map(""" + [N] -> { + [k_after] -> [j_before = k_after - 1] : + 1 <= k_after < N + } + """) + ) + kernel = kernel.copy( + instructions=tuple( + stmt.copy( + happens_after={ + stmt.id: stmt.happens_after[stmt.id], + **({"S": t_after_s} if stmt.id == "T" else {}), + **({"T": u_after_t} if stmt.id == "U" else {}), + } + ) + for stmt in kernel.instructions + ) + ) + kernel = dep.relax_strict_happens_after(kernel) + + required_order = kernel.id_to_insn["U"].happens_after["S"].instances_rel + assert required_order is not None + assert required_order.equals( + nisl.make_map(""" + [N] -> { + [k_after] -> [i_before = 2*k_after - 2] : + 1 <= k_after < N + } + """) + ) + + +def test_relax_strict_happens_after_traverses_self_edges() -> None: + t_unit = lp.make_kernel( + "{ [i] : 0 <= i < 8 }", + """ + x[i + 1] = inp[i] {id=P} + out[i] = x[i] {id=C} + """, + ) + kernel = t_unit.default_entrypoint + c_after_c = HappensAfter(nisl.make_map(""" + { + [i_after] -> [i_before] : + 0 <= i_before < i_after < 8 + } + """)) + c_after_p = HappensAfter(nisl.make_map(""" + { + [i_after] -> [i_before = i_after] : + 0 <= i_after < 8 + } + """)) + kernel = kernel.copy(instructions=tuple( + stmt.copy(happens_after={ + "P": {}, + "C": {"C": c_after_c, "P": c_after_p}, + }[stmt.id]) + for stmt in kernel.instructions + )) + + kernel = dep.relax_strict_happens_after(kernel) + + required_order = kernel.id_to_insn["C"].happens_after["P"].instances_rel + assert required_order is not None + assert required_order.equals(nisl.make_map(""" + { + [i_after] -> [i_before = i_after - 1] : + 1 <= i_after < 8 + } + """)) + + +def test_relax_strict_happens_after_rejects_nontransitive_self_order() -> None: + t_unit = lp.make_kernel( + "{ [i] : 0 <= i < 8 }", + "out[i] = i {id=S}", + ) + kernel = t_unit.default_entrypoint + immediate_predecessor = HappensAfter(nisl.make_map(""" + { + [i_after] -> [i_before = i_after - 1] : + 1 <= i_after < 8 + } + """)) + kernel = kernel.copy(instructions=( + kernel.id_to_insn["S"].copy( + happens_after={"S": immediate_predecessor} + ), + )) + + with pytest.raises(LoopyError, match="self.*'S'.*not transitive"): + dep.relax_strict_happens_after(kernel) + + +def test_relax_strict_happens_after_unions_branched_paths() -> None: + t_unit = lp.make_kernel( + "{ [i] : 0 <= i < N }", + """ + a[i] = 1 {id=A} + b[i] = 2 {id=B} + c[i] = 3 {id=C} + d[i] = a[i] {id=D} + """, + ) + kernel = dep.add_lexicographic_happens_after(t_unit).default_entrypoint + cross_order = kernel.id_to_insn["B"].happens_after["A"] + predecessors = { + "A": (), + "B": ("A",), + "C": ("A",), + "D": ("B", "C"), + } + kernel = kernel.copy( + instructions=[ + insn.copy( + happens_after={ + insn.id: insn.happens_after[insn.id], + **dict.fromkeys(predecessors[insn.id], cross_order), + } + ) + for insn in kernel.instructions + ] + ) + + kernel = dep.relax_strict_happens_after(kernel) + + required_order = kernel.id_to_insn["D"].happens_after["A"].instances_rel + assert required_order is not None + assert required_order.equals( + nisl.make_map(""" + [N] -> { + [i_after] -> [i_before = i_after] : 0 <= i_after < N + } + """) + ) + + +def test_statement_timestamps_with_calls_inside_outer_loop() -> None: + kernel = lp.make_kernel( + """{ + [batch, i, j, k] : + 0 <= batch < 4 and 0 <= i < 8 and + 0 <= j < 16 and 0 <= k < 32 + }""", + """ + a[batch, i, j] = i + j {id=A} + b[batch, i] = i {id=B} + c[batch, k] = k {id=C} + """, + ).default_entrypoint.copy( + state=lp.KernelState.LINEARIZED, + linearization=( + EnterLoop("batch"), + CallKernel("phase0"), + EnterLoop("i"), + EnterLoop("j"), + RunInstruction("A"), + LeaveLoop("j"), + RunInstruction("B"), + LeaveLoop("i"), + ReturnFromKernel("phase0"), + CallKernel("phase1"), + EnterLoop("k"), + RunInstruction("C"), + LeaveLoop("k"), + ReturnFromKernel("phase1"), + LeaveLoop("batch"), + ), + ) + + assert _get_timestamp_points_from_linearization(kernel) == _PreciseSchedule( + statements={ + "A": _StatementRecord( + timestamp=(0, "batch", 1, 2, "i", 3, "j", 4), + subkernel_idx=1, + ), + "B": _StatementRecord( + timestamp=(0, "batch", 1, 2, "i", 6), + subkernel_idx=1, + ), + "C": _StatementRecord( + timestamp=(0, "batch", 9, 10, "k", 11), + subkernel_idx=9, + ), + }, + barriers=(), + ) + + +def test_statement_timestamps_with_imperfect_loop_nesting() -> None: + t_unit = lp.make_kernel( + "{ [i, j, k] : 0 <= i, j, k < 8 }", + """ + <> before = 0 {id=before} + <> outer = i {id=outer} + <> deep = i + j + k {id=deep} + <> after_inner = i {id=after_inner} + <> after = 0 {id=after} + """, + ) + t_unit = dep.add_lexicographic_happens_after(t_unit) + t_unit = lp.preprocess_program(t_unit) + kernel = lp.linearize(t_unit).default_entrypoint + + assert _get_timestamp_points_from_linearization(kernel) == _PreciseSchedule( + statements={ + "before": _StatementRecord(timestamp=(0, 1), subkernel_idx=0), + "outer": _StatementRecord( + timestamp=(0, 2, "i", 3), subkernel_idx=0 + ), + "deep": _StatementRecord( + timestamp=(0, 2, "i", 4, "j", 5, "k", 6), + subkernel_idx=0, + ), + "after_inner": _StatementRecord( + timestamp=(0, 2, "i", 9), subkernel_idx=0 + ), + "after": _StatementRecord(timestamp=(0, 11), subkernel_idx=0), + }, + barriers=(), + ) + + +def test_timestamp_relation_keeps_parallel_inames_in_instance_domain() -> None: + t_unit = lp.make_kernel( + "{ [g, l, i] : 0 <= g < 4 and 0 <= l < 8 and 0 <= i < 16 }", + "out[g, l, i] = g + l + i {id=S}", + ) + t_unit = lp.tag_inames(t_unit, {"g": "g.0", "l": "l.0"}) + t_unit = lp.preprocess_program(t_unit) + kernel = lp.linearize(t_unit).default_entrypoint + + precise_schedule = _get_timestamp_points_from_linearization(kernel) + assert precise_schedule.statements["S"] == _StatementRecord( + timestamp=(0, 1, "i", 2), + subkernel_idx=0, + ) + + timestamp_relations, _, _ = _build_timestamp_relations( + kernel, precise_schedule, kernel.get_var_name_generator() + ) + timestamp_relation = timestamp_relations["S"] + assert timestamp_relation.equals( + nisl.make_map(""" + { + [g, l, i] -> [__ts_0, __ts_1, __ts_2, __ts_3] : + 0 <= g < 4 and 0 <= l < 8 and 0 <= i < 16 and + __ts_0 = 0 and __ts_1 = 1 and + __ts_2 = i and __ts_3 = 2 + } + """) + ) + + +def test_analysis_names_do_not_clash_with_kernel_names() -> None: + t_unit = lp.make_kernel( + """ + [__group_0] -> { + [g, __ts_0, __ts_0_later] : + 0 <= g < __group_0 and + 0 <= __ts_0, __ts_0_later < 2 + } + """, + """ + a[g, __ts_0, __ts_0_later] = g {id=S} + out[g, __ts_0, __ts_0_later] = a[g, __ts_0, __ts_0_later] {id=T} + """, + ) + t_unit = dep.add_lexicographic_happens_after(t_unit) + t_unit = dep.relax_strict_happens_after(t_unit) + t_unit = lp.tag_inames(t_unit, {"g": "g.0"}) + t_unit = lp.preprocess_program(t_unit) + kernel = lp.linearize(t_unit).default_entrypoint + + precise_schedule = _get_timestamp_points_from_linearization(kernel) + stmt_relations, _, timestamp_order = _build_timestamp_relations( + kernel, precise_schedule, kernel.get_var_name_generator() + ) + + timestamp_names = stmt_relations["S"].space.dimtype_to_names[ + nisl.DimType.out + ] + later_names = timestamp_order.space.dimtype_to_names[nisl.DimType.in_] + earlier_names = timestamp_order.space.dimtype_to_names[nisl.DimType.out] + analysis_names = {*timestamp_names, *later_names, *earlier_names} + + assert len(analysis_names) == 3 * len(timestamp_names) + assert analysis_names.isdisjoint(kernel.all_variable_names()) + verify_happens_after_is_enforced(kernel) + + +def test_strict_lexicographic_timestamp_order() -> None: + order = _build_strict_lexicographic_order( + ("t0_later", "t1_later", "t2_later"), + ("t0_earlier", "t1_earlier", "t2_earlier"), + ) + + assert order.equals( + nisl.make_map(""" + { + [t0_later, t1_later, t2_later] -> + [t0_earlier, t1_earlier, t2_earlier] : + t0_later > t0_earlier; + [t0_later, t1_later, t2_later] -> + [t0_earlier, t1_earlier, t2_earlier] : + t0_later = t0_earlier and + t1_later > t1_earlier; + [t0_later, t1_later, t2_later] -> + [t0_earlier, t1_earlier, t2_earlier] : + t0_later = t0_earlier and + t1_later = t1_earlier and + t2_later > t2_earlier + } + """) + ) + + +def test_statement_timestamps_record_local_and_global_barriers() -> None: + local_barrier = Barrier( + comment="local synchronization", + synchronization_kind="local", + mem_kind="local", + originating_insn_id=None, + ) + global_barrier = Barrier( + comment="global synchronization", + synchronization_kind="global", + mem_kind="global", + originating_insn_id=None, + ) + kernel = lp.make_kernel( + "{ [batch, i, j] : 0 <= batch, i, j < 8 }", + """ + a[batch, i] = i {id=producer} + b[batch, i] = a[batch, i] {id=after_local} + c[batch, j] = j {id=consumer} + """, + ).default_entrypoint.copy( + state=lp.KernelState.LINEARIZED, + linearization=( + EnterLoop("batch"), + CallKernel("phase0"), + EnterLoop("i"), + RunInstruction("producer"), + local_barrier, + RunInstruction("after_local"), + LeaveLoop("i"), + ReturnFromKernel("phase0"), + global_barrier, + CallKernel("phase1"), + EnterLoop("j"), + RunInstruction("consumer"), + LeaveLoop("j"), + ReturnFromKernel("phase1"), + LeaveLoop("batch"), + ), + ) + + precise_schedule = _get_timestamp_points_from_linearization(kernel) + assert precise_schedule == _PreciseSchedule( + statements={ + "producer": _StatementRecord( + timestamp=(0, "batch", 1, 2, "i", 3), + subkernel_idx=1, + ), + "after_local": _StatementRecord( + timestamp=(0, "batch", 1, 2, "i", 5), + subkernel_idx=1, + ), + "consumer": _StatementRecord( + timestamp=(0, "batch", 9, 10, "j", 11), + subkernel_idx=9, + ), + }, + barriers=( + _BarrierRecord( + timestamp=(0, "batch", 1, 2, "i", 4), + barrier=local_barrier, + subkernel_idx=1, + ), + _BarrierRecord( + timestamp=(0, "batch", 8), + barrier=global_barrier, + subkernel_idx=None, + ), + ), + ) + + _, barrier_relations, _ = _build_timestamp_relations( + kernel, precise_schedule, kernel.get_var_name_generator() + ) + assert len(barrier_relations) == 2 + assert barrier_relations[0].equals( + nisl.make_map(""" + { + [batch, i] -> + [__ts_0, __ts_1, __ts_2, __ts_3, __ts_4, __ts_5] : + 0 <= batch < 8 and 0 <= i < 8 and + __ts_0 = 0 and __ts_1 = batch and + __ts_2 = 1 and __ts_3 = 2 and + __ts_4 = i and __ts_5 = 4 + } + """) + ) + assert barrier_relations[1].equals( + nisl.make_map(""" + { + [batch] -> + [__ts_0, __ts_1, __ts_2, __ts_3, __ts_4, __ts_5] : + 0 <= batch < 8 and + __ts_0 = 0 and __ts_1 = batch and + __ts_2 = 8 and __ts_3 = 0 and + __ts_4 = 0 and __ts_5 = 0 + } + """) + ) + + +def test_local_barrier_orders_work_items_in_the_same_group() -> None: + t_unit = lp.make_kernel( + """ + { + [g, l, i] : + 0 <= g < 2 and 0 <= l < 4 and 0 <= i < 2 + } + """, + """ + <> tmp[g, l, i] = g + l + i {id=source} + out[g, l, i] = tmp[g, (l + 1) % 4, i] {id=sink} + """, + ) + t_unit = dep.add_lexicographic_happens_after(t_unit) + t_unit = dep.relax_strict_happens_after(t_unit) + t_unit = lp.set_temporary_address_space( + t_unit, "tmp", lp.AddressSpace.LOCAL + ) + t_unit = lp.tag_inames(t_unit, {"g": "g.0", "l": "l.0"}) + t_unit = lp.preprocess_program(t_unit) + kernel = lp.linearize(t_unit).default_entrypoint + + precise_schedule = _get_timestamp_points_from_linearization(kernel) + name_generator = kernel.get_var_name_generator() + stmt_relations, barrier_relations, timestamp_order = ( + _build_timestamp_relations(kernel, precise_schedule, name_generator) + ) + enforced = _build_enforced_order( + kernel, + "sink", + "source", + precise_schedule, + stmt_relations, + barrier_relations, + timestamp_order, + name_generator, + ) + + assert enforced.equals( + nisl.make_map(""" + { + [g_after, l_after, i_after] -> + [g_before, l_before, i_before] : + 0 <= g_after < 2 and 0 <= l_after < 4 and + 0 <= i_after < 2 and + g_before = g_after and + 0 <= l_before < 4 and + 0 <= i_before <= i_after + } + """) + ) + + +def test_global_barrier_orders_all_work_items() -> None: + t_unit = lp.make_kernel( + "{ [g, l] : 0 <= g < 2 and 0 <= l < 4 }", + """ + a[g, l] = g + l {id=source} + out[g, l] = a[(g + 1) % 2, l] {id=sink} + """, + ) + t_unit = dep.add_lexicographic_happens_after(t_unit) + t_unit = dep.relax_strict_happens_after(t_unit) + t_unit = lp.tag_inames(t_unit, {"g": "g.0", "l": "l.0"}) + t_unit = lp.set_options(t_unit, insert_gbarriers=True) + t_unit = lp.preprocess_program(t_unit) + kernel = lp.linearize(t_unit).default_entrypoint + + precise_schedule = _get_timestamp_points_from_linearization(kernel) + name_generator = kernel.get_var_name_generator() + stmt_relations, barrier_relations, timestamp_order = ( + _build_timestamp_relations(kernel, precise_schedule, name_generator) + ) + enforced = _build_enforced_order( + kernel, + "sink", + "source", + precise_schedule, + stmt_relations, + barrier_relations, + timestamp_order, + name_generator, + ) + + assert enforced.equals( + nisl.make_map(""" + { + [g_after, l_after] -> [g_before, l_before] : + 0 <= g_after < 2 and 0 <= l_after < 4 and + 0 <= g_before < 2 and 0 <= l_before < 4 + } + """) + ) + + +def test_numerical_affine_and_splicing_integration( + ctx_factory: cl.CtxFactory, +) -> None: + args = [ + lp.GlobalArg("x,out", dtype=np.int32, shape=(64,)), + "...", + ] + ref_t_unit = lp.make_kernel( + "{ [i] : 0 <= i < 64 }", + "out[i] = 2*(x[i] + 1) + 3", + args, + ) + t_unit = lp.make_kernel( + "{ [i] : 0 <= i < 64 }", + """ + out[i] = g[i] + 3 {id=S} + <> g[i] = 2*d[i] {id=G} + <> a[i] = 2*d[i] {id=A} + <> d[i] = x[i] + 1 {id=D} + """, + args, + ) + same_instance = HappensAfter(nisl.make_map(""" + { + [i_after] -> [i_before = i_after] : 0 <= i_after < 64 + } + """)) + kernel = t_unit.default_entrypoint + kernel = kernel.copy(instructions=tuple( + stmt.copy(happens_after={ + "S": {"A": same_instance}, + "G": {}, + "A": {"D": same_instance}, + "D": {}, + }[stmt.id]) + for stmt in kernel.instructions + )) + t_unit = t_unit.with_kernel(kernel) + + t_unit = lp.split_iname( + t_unit, + "i", + 4, + outer_iname="io", + inner_iname="ii", + ) + same_tiled_instance = nisl.make_map(""" + { + [io_after, ii_after] -> [io_before, ii_before] : + io_before = io_after and ii_before = ii_after and + 0 <= 4*io_after + ii_after < 64 and 0 <= ii_after < 4 + } + """) + kernel = dep.splice_happens_after_as_consumer_and_producer( + t_unit.default_entrypoint, + "G", + "A", + same_tiled_instance, + same_tiled_instance, + ) + t_unit = t_unit.with_kernel(kernel) + t_unit = lp.prioritize_loops(t_unit, "io,ii") + + lp.auto_test_vs_ref( + ref_t_unit, + ctx_factory(), + t_unit, + print_code=False, + quiet=True, + ) + + +def test_verification_enforces_self_recurrence( + ctx_factory: cl.CtxFactory, +) -> None: + ref_t_unit = lp.make_kernel( + "{ [i] : 1 <= i < N }", + "a[i] = a[i - 1] + x[i] {id=S}", + [lp.GlobalArg("a,x", dtype=np.int32, shape=lp.auto), "..."], + ) + t_unit = ref_t_unit + t_unit = dep.add_lexicographic_happens_after(t_unit) + t_unit = dep.relax_strict_happens_after(t_unit) + + required_order = ( + t_unit.default_entrypoint + .id_to_insn["S"] + .happens_after["S"] + .instances_rel + ) + assert required_order is not None + assert required_order.equals( + nisl.make_map(""" + [N] -> { + [i_after] -> [i_before = i_after - 1] : + 2 <= i_after < N + } + """) + ) + + lp.auto_test_vs_ref( + ref_t_unit, + ctx_factory(), + t_unit, + parameters={"N": 64}, + print_code=False, + quiet=True, + ) + + parallel_t_unit = lp.tag_inames(t_unit, {"i": "l.0"}) + parallel_kernel = parallel_t_unit.default_entrypoint.copy( + state=lp.KernelState.LINEARIZED, + linearization=( + CallKernel("device_program"), + RunInstruction("S"), + ReturnFromKernel("device_program"), + ), + ) + with pytest.raises( + LoopyError, + match="schedule does not enforce 'S' after 'S'", + ): + verify_happens_after_is_enforced( + parallel_t_unit.with_kernel(parallel_kernel) + ) + + +def test_verification_rejects_unenforced_order() -> None: + t_unit = lp.make_kernel( + "{ [i] : 0 <= i < N }", + """ + a[i] = i {id=S} + b[i] = a[i] {id=T} + """, + ) + t_unit = dep.add_lexicographic_happens_after(t_unit) + + kernel = t_unit.default_entrypoint.copy( + state=lp.KernelState.LINEARIZED, + linearization=( + CallKernel("device_program"), + EnterLoop("i"), + RunInstruction("T"), + RunInstruction("S"), + LeaveLoop("i"), + ReturnFromKernel("device_program"), + ), + ) + t_unit = t_unit.with_kernel(kernel) + + with pytest.raises( + LoopyError, + match="schedule does not enforce 'T' after 'S'", + ): + verify_happens_after_is_enforced(t_unit) + + with pytest.raises( + LoopyError, + match="schedule does not enforce 'T' after 'S'", + ): + lp.generate_code_v2(t_unit) + + +def test_verification_does_not_order_vector_lanes() -> None: + t_unit = lp.make_kernel( + "{ [v] : 0 <= v < 4 }", + "out[v] = v {id=S}", + ) + kernel = t_unit.default_entrypoint + required_order = HappensAfter( + nisl.make_map(""" + { + [v_after] -> [v_before = v_after - 1] : + 1 <= v_after < 4 + } + """) + ) + kernel = kernel.copy( + instructions=( + kernel.id_to_insn["S"].copy( + happens_after={"S": required_order} + ), + ) + ) + t_unit = lp.tag_inames(t_unit.with_kernel(kernel), {"v": "vec"}) + kernel = t_unit.default_entrypoint.copy( + state=lp.KernelState.LINEARIZED, + linearization=( + CallKernel("device_program"), + EnterLoop("v"), + RunInstruction("S"), + LeaveLoop("v"), + ReturnFromKernel("device_program"), + ), + ) + + with pytest.raises( + LoopyError, + match="schedule does not enforce 'S' after 'S'", + ): + verify_happens_after_is_enforced(t_unit.with_kernel(kernel)) + + +def test_verification_handles_explicit_barrier_instruction() -> None: + t_unit = lp.make_kernel( + "{ : }", + """ + a[0] = 1 {id=S} + ... gbarrier {id=B} + b[0] = a[0] {id=T} + """, + seq_dependencies=True, + ) + t_unit = dep.add_lexicographic_happens_after(t_unit) + t_unit = lp.preprocess_program(t_unit) + t_unit = lp.linearize(t_unit) + + verify_happens_after_is_enforced(t_unit) + + +if __name__ == "__main__": + import sys + + if len(sys.argv) > 1: + exec(sys.argv[1]) + else: + from pytest import main + + main([__file__]) diff --git a/test/test_fusion.py b/test/test_fusion.py index c7ded797b..a006881b1 100644 --- a/test/test_fusion.py +++ b/test/test_fusion.py @@ -22,6 +22,7 @@ import numpy as np +import pytest import namedisl as nisl import pyopencl as cl @@ -30,6 +31,7 @@ ) import loopy as lp +import loopy.kernel.dependency as dep def test_two_kernel_fusion(ctx_factory: cl.CtxFactory): @@ -57,7 +59,13 @@ def test_two_kernel_fusion(ctx_factory: cl.CtxFactory): np.testing.assert_allclose(out.get(), np.arange(100, 110)) -def test_write_block_matrix_fusion(ctx_factory: cl.CtxFactory): +@pytest.mark.parametrize( + "precise_dependencies", (False, True), ids=("legacy", "precise") +) +def test_write_block_matrix_fusion( + ctx_factory: cl.CtxFactory, + precise_dependencies: bool, +): """ A slightly more complicated fusion test, where all sub-kernels write into the same global matrix, but @@ -170,6 +178,10 @@ def write_into_mat_prg(): bidirectional=True, force=True ) + if precise_dependencies: + fused_knl = dep.add_lexicographic_happens_after(fused_knl) + fused_knl = dep.relax_strict_happens_after(fused_knl) + _evt, result = fused_knl(queue, **kwargs) result = result["result"] np.testing.assert_allclose(result, answer) diff --git a/test/test_loop_fusion.py b/test/test_loop_fusion.py index 2a1dd46ff..903576c7f 100644 --- a/test/test_loop_fusion.py +++ b/test/test_loop_fusion.py @@ -29,6 +29,7 @@ import pyopencl as cl import loopy as lp +import loopy.kernel.dependency as dep logger = logging.getLogger(__name__) @@ -195,7 +196,13 @@ def test_loop_fusion_constrained_by_outer_loop_deps(ctx_factory: cl.CtxFactory): lp.auto_test_vs_ref(ref_knl, ctx, knl) -def test_loop_fusion_with_loop_carried_deps1(ctx_factory: cl.CtxFactory): +@pytest.mark.parametrize( + "precise_dependencies", (False, True), ids=("legacy", "precise") +) +def test_loop_fusion_with_loop_carried_deps1( + ctx_factory: cl.CtxFactory, + precise_dependencies: bool, +): ctx = ctx_factory() knl = lp.make_kernel( @@ -227,6 +234,10 @@ def test_loop_fusion_with_loop_carried_deps1(ctx_factory: cl.CtxFactory): == 1 ) + if precise_dependencies: + knl = dep.add_lexicographic_happens_after(knl) + knl = dep.relax_strict_happens_after(knl) + lp.auto_test_vs_ref(ref_knl, ctx, knl) diff --git a/test/test_transform.py b/test/test_transform.py index 2d4d0b542..5594c7503 100644 --- a/test/test_transform.py +++ b/test/test_transform.py @@ -34,6 +34,7 @@ from pytools.tag import Tag import loopy as lp +import loopy.kernel.dependency as dep from loopy.version import ( LOOPY_USE_LANGUAGE_VERSION_2018_2, # ruff:ignore[unused-import] ) @@ -1345,7 +1346,13 @@ def test_prefetch_with_within(ctx_factory: cl.CtxFactory): lp.auto_test_vs_ref(ref_t_unit, ctx_factory(), t_unit) -def test_privatize_with_nonzero_lbound(ctx_factory: cl.CtxFactory): +@pytest.mark.parametrize( + "precise_dependencies", (False, True), ids=("legacy", "precise") +) +def test_privatize_with_nonzero_lbound( + ctx_factory: cl.CtxFactory, + precise_dependencies: bool, +): ctx = ctx_factory() queue = cl.CommandQueue(ctx) @@ -1362,6 +1369,10 @@ def test_privatize_with_nonzero_lbound(ctx_factory: cl.CtxFactory): knl = lp.privatize_temporaries_with_inames(knl, {"j"}) assert knl["arange_10_to_14"].temporary_variables["tmp"].shape == (4,) + if precise_dependencies: + knl = dep.add_lexicographic_happens_after(knl) + knl = dep.relax_strict_happens_after(knl) + _, (out, ) = knl(queue) np.testing.assert_allclose(out.get()[10:14], np.arange(10, 14))