From b41f1651df5be48b7a8ea7266f4d863c7de32898 Mon Sep 17 00:00:00 2001 From: Addison Date: Fri, 17 Jul 2026 16:31:39 -0500 Subject: [PATCH 01/30] update HappensAfter contract to reflect v2 --- loopy/kernel/instruction.py | 28 ++++++++++------------------ 1 file changed, 10 insertions(+), 18 deletions(-) diff --git a/loopy/kernel/instruction.py b/loopy/kernel/instruction.py index 01dc59146..61d2a8d4f 100644 --- a/loopy/kernel/instruction.py +++ b/loopy/kernel/instruction.py @@ -47,6 +47,7 @@ from typing_extensions import Self, override import islpy as isl +import namedisl as nisl import pymbolic.primitives as p from pytools import ImmutableRecord, memoize_method from pytools.tag import Tag, Taggable, tag_dataclass @@ -119,35 +120,28 @@ 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. - - .. 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`. + 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:: 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: isl.Map | None + instances_rel: nisl.Map | None # }}} @@ -356,14 +350,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): From 7d1ddbf53974dacf52a4a88ef44cc74dccf61332 Mon Sep 17 00:00:00 2001 From: Addison Date: Fri, 17 Jul 2026 22:00:04 -0500 Subject: [PATCH 02/30] implement add_lexicographic_happens_after + tests --- loopy/kernel/dependency.py | 95 ++++++++++++++++++++++++ loopy/kernel/instruction.py | 2 +- test/test_dependency.py | 143 ++++++++++++++++++++++++++++++++++++ 3 files changed, 239 insertions(+), 1 deletion(-) create mode 100644 loopy/kernel/dependency.py create mode 100644 test/test_dependency.py diff --git a/loopy/kernel/dependency.py b/loopy/kernel/dependency.py new file mode 100644 index 000000000..161bdc3f4 --- /dev/null +++ b/loopy/kernel/dependency.py @@ -0,0 +1,95 @@ +from constantdict import constantdict +import namedisl as nisl +from namedisl import DimType + +from loopy import for_each_kernel +from loopy.kernel import LoopKernel +from loopy.kernel.instruction import HappensAfter + +from pytools.graph import compute_topological_order + + +def _prefix_names(obj: nisl.Set, prefix: str, dim_type: DimType) -> nisl.Set: + return obj.rename_dims( + ((name, name + prefix) for name in obj.space.dimtype_to_names[dim_type]) + ) + + +@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_insns = [] + for i, insn in enumerate(kernel.instructions): + new_happens_after = {} + + preds = (insn,) if i == 0 else (insn, kernel.instructions[i - 1]) + + # FIXME: yuck + after_domain = nisl.make_set( + kernel.get_inames_domain(insn.within_inames).to_set() + ) + + after_inames = after_domain.space.dimtype_to_names[DimType.out] + after_domain = _prefix_names(after_domain, "_after", DimType.out) + for pred in preds: + before_domain = nisl.make_set( + kernel.get_inames_domain(pred.within_inames).to_set() + ) + + before_inames = before_domain.space.dimtype_to_names[DimType.out] + before_domain = _prefix_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.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 pred.id == insn.id: + ordered_instances = strict_lex + else: + ordered_instances = strict_lex | equal_prefix + + instances_rel = ordered_instances.as_map( + in_names=tuple(f"{name}_after" for name in after_inames) + ) + + new_happens_after[pred.id] = HappensAfter( + instances_rel=instances_rel + ) + + new_insns.append(insn.copy(happens_after=new_happens_after)) + + return kernel.copy(instructions=new_insns) diff --git a/loopy/kernel/instruction.py b/loopy/kernel/instruction.py index 82a6ceebc..813fe3071 100644 --- a/loopy/kernel/instruction.py +++ b/loopy/kernel/instruction.py @@ -141,7 +141,7 @@ class HappensAfter: statement-level dependencies of prior versions of :mod:`loopy`. """ - instances_rel: nisl.Map | None + instances_rel: nisl.Map | nisl.BasicMap | None # }}} diff --git a/test/test_dependency.py b/test/test_dependency.py new file mode 100644 index 000000000..1080459f0 --- /dev/null +++ b/test/test_dependency.py @@ -0,0 +1,143 @@ +import namedisl as nisl + +import loopy as lp +import loopy.kernel.dependency as dep + + +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_distinct_loop_nests() -> None: + t_unit = lp.make_kernel( + [ + "{ [i] : 0 <= i < N }", + "{ [j] : 0 <= j < M }", + ], + """ + a[i] = 2 * a[i] {id=S} + b[j] = 2 * b[j] {id=T} + """, + ) + + kernel = dep.add_lexicographic_happens_after(t_unit).default_entrypoint + cross_relation = kernel.id_to_insn["T"].happens_after["S"].instances_rel + + assert cross_relation is not None + assert cross_relation.equals( + nisl.make_map(""" + [N, M] -> { + [j_after] -> [i_before] : + 0 <= j_after < M and + 0 <= i_before < N + } + """) + ) + + +def test_add_lexicographic_happens_after_with_five_inames() -> None: + t_unit = lp.make_kernel( + """ + { [q, z, a, m, b] : + 0 <= q < 2 and + 0 <= z < 2 and + 0 <= a < 2 and + 0 <= m < 2 and + 0 <= b < 2 + } + """, + "out[q, z, a, m, b] = q + z + a + m + b {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(""" + { + [q_after, z_after, a_after, m_after, b_after] -> + [q_before, z_before, a_before, m_before, b_before] : + 0 <= q_after, q_before < 2 and + 0 <= z_after, z_before < 2 and + 0 <= a_after, a_before < 2 and + 0 <= m_after, m_before < 2 and + 0 <= b_after, b_before < 2 and + (q_before < q_after or + (q_before = q_after and z_before < z_after) or + (q_before = q_after and z_before = z_after and + a_before < a_after) or + (q_before = q_after and z_before = z_after and + a_before = a_after and m_before < m_after) or + (q_before = q_after and z_before = z_after and + a_before = a_after and m_before = m_after and + b_before < b_after)) + } + """) + ) + + +if __name__ == "__main__": + import sys + + if len(sys.argv) > 1: + exec(sys.argv[1]) + else: + from pytest import main + + main([__file__]) From 5c82ac4455450a53fd7db166feb7b94fb58fa19c Mon Sep 17 00:00:00 2001 From: Addison Date: Sun, 19 Jul 2026 01:23:41 -0500 Subject: [PATCH 03/30] add access reln finding machinery + v2 chaser --- loopy/kernel/dependency.py | 426 ++++++++++++++++++++++++++++++++++- loopy/kernel/instruction.py | 13 +- test/test_dependency.py | 433 ++++++++++++++++++++++++++++++++++++ 3 files changed, 857 insertions(+), 15 deletions(-) diff --git a/loopy/kernel/dependency.py b/loopy/kernel/dependency.py index 161bdc3f4..aa6dc5690 100644 --- a/loopy/kernel/dependency.py +++ b/loopy/kernel/dependency.py @@ -1,17 +1,205 @@ -from constantdict import constantdict +from __future__ import annotations + +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.kernel import LoopKernel -from loopy.kernel.instruction import HappensAfter +from loopy.kernel.instruction import ( + HappensAfter, + InstructionBase, + MultiAssignmentBase, +) +from loopy.symbolic import ( + LinearSubscript, + Reduction, + SubArrayRef, + WalkMapper, + aff_from_expr, +) + + +if TYPE_CHECKING: + from collections.abc import Mapping + + from namedisl.core import NamedIslObjectT + + from pymbolic.typing import Expression + + from loopy.kernel import LoopKernel + from loopy.typing import ShapeType + + +@final +class AccessType(Enum): + read = 0 + write = 1 -from pytools.graph import compute_topological_order +class AccessRelationFinder(WalkMapper[[str, AccessType]]): + kernel: LoopKernel + variables: frozenset[str] + _additional_inames: frozenset[str] + _read_relations: dict[str, dict[str, nisl.Map]] + _write_relations: dict[str, dict[str, nisl.Map]] + + def __init__(self, kernel: LoopKernel): + self.kernel = kernel + self.variables = frozenset(kernel.all_variable_names()) + self._additional_inames = frozenset() + self._read_relations = {insn.id: {} for insn in kernel.instructions} + self._write_relations = {insn.id: {} for insn in kernel.instructions} + + super().__init__() + + def _get_access_relation( + self, + domain: nisl.Set, + subscript: tuple[Expression, ...], + assumptions: nisl.BasicSet | None = None, + shape: ShapeType | None = None, + allowed_constant_names: frozenset[str] | None = None, + ) -> nisl.Map: + instance_names = domain.space.dimtype_to_names[DimType.out] + cell_names = tuple(f"ax_{axis}" for axis in range(len(subscript))) + + access_set = domain.add_dims(DimType.out, cell_names) + coordinates = access_set.pw_affs + for cell_name, index_expr in zip(cell_names, subscript, strict=True): + index_aff = nisl.make_aff( + aff_from_expr( + access_set.space.as_isl_set_space(), + index_expr, + ) + ).as_pw_aff() + + access_set = access_set & coordinates[cell_name].eq_set(index_aff) + + return access_set.as_map(in_names=instance_names) + + def _insn_writes_var(self, insn_id: str, var: str) -> bool: + return ( + var in self.kernel.writer_map() + and insn_id in self.kernel.writer_map()[var] + ) + + def _insn_reads_var(self, insn_id: str, var: str) -> bool: + return ( + var in self.kernel.reader_map() + and insn_id in self.kernel.reader_map()[var] + ) + + def _insn_accesses_var(self, insn_id: str, var: str) -> bool: + return self._insn_reads_var(insn_id, var) | self._insn_writes_var( + insn_id, var + ) -def _prefix_names(obj: nisl.Set, prefix: str, dim_type: DimType) -> nisl.Set: + def _record_access( + self, + insn_id: str, + var: str, + subscript: tuple[Expression, ...], + access_type: AccessType, + ) -> None: + if not self._insn_accesses_var(insn_id, var): + return + + insn = self.kernel.id_to_insn[insn_id] + domain_inames = insn.within_inames | self._additional_inames + inames_domain = nisl.make_set( + self.kernel.get_inames_domain(domain_inames).to_set() + ) + access_rel = self._get_access_relation(inames_domain, subscript) + + additional_inames = self._additional_inames - insn.within_inames + if additional_inames: + access_rel = access_rel.project_out(additional_inames) + + match access_type: + case AccessType.read: + previous = self._read_relations[insn_id].get(var) + self._read_relations[insn_id][var] = ( + access_rel if previous is None else previous | access_rel + ) + case AccessType.write: + previous = self._write_relations[insn_id].get(var) + self._write_relations[insn_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({ + insn_id: constantdict(self._read_relations[insn_id]) + for insn_id in self._read_relations + }) + + @cached_property + def write_relations(self) -> Mapping[str, Mapping[str, nisl.Map]]: + return constantdict({ + insn_id: constantdict(self._write_relations[insn_id]) + for insn_id in self._write_relations + }) + + @override + def map_subscript( + self, expr: prim.Subscript, /, insn_id: str, access_type: AccessType + ) -> None: + assert isinstance(expr.aggregate, prim.Variable) + self._record_access( + insn_id, expr.aggregate.name, expr.index_tuple, access_type + ) + + @override + def map_linear_subscript( + self, expr: LinearSubscript, /, insn_id: str, access_type: AccessType + ) -> None: + self.rec(expr.index, insn_id, AccessType.read) + + assert isinstance(expr.aggregate, prim.Variable) + self._record_access( + insn_id, expr.aggregate.name, (expr.index,), access_type + ) + + @override + def map_reduction( + self, expr: Reduction, /, insn_id: str, access_type: AccessType + ) -> None: + previous_inames = self._additional_inames + self._additional_inames |= frozenset(expr.inames) + try: + WalkMapper.map_reduction(self, expr, insn_id, access_type) + finally: + self._additional_inames = previous_inames + + @override + def map_sub_array_ref( + self, expr: SubArrayRef, /, insn_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, insn_id, access_type) + finally: + self._additional_inames = previous_inames + + +def _suffix_names( + obj: NamedIslObjectT, prefix: str, dim_type: DimType +) -> NamedIslObjectT: return obj.rename_dims( - ((name, name + prefix) for name in obj.space.dimtype_to_names[dim_type]) + (name, name + prefix) for name in obj.space.dimtype_to_names[dim_type] ) @@ -29,7 +217,7 @@ def add_lexicographic_happens_after(kernel: LoopKernel) -> LoopKernel: immediately preceding statement. """ - new_insns = [] + new_insns: list[InstructionBase] = [] for i, insn in enumerate(kernel.instructions): new_happens_after = {} @@ -41,14 +229,14 @@ def add_lexicographic_happens_after(kernel: LoopKernel) -> LoopKernel: ) after_inames = after_domain.space.dimtype_to_names[DimType.out] - after_domain = _prefix_names(after_domain, "_after", DimType.out) + after_domain = _suffix_names(after_domain, "_after", DimType.out) for pred in preds: before_domain = nisl.make_set( kernel.get_inames_domain(pred.within_inames).to_set() ) before_inames = before_domain.space.dimtype_to_names[DimType.out] - before_domain = _prefix_names(before_domain, "_before", 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 @@ -93,3 +281,223 @@ def add_lexicographic_happens_after(kernel: LoopKernel) -> LoopKernel: new_insns.append(insn.copy(happens_after=new_happens_after)) return kernel.copy(instructions=new_insns) + + +def _relax_strict_happens_after_inner( + kernel: LoopKernel, + sink_id: str, + source_id: str, + var: str, + sink_access_type: AccessType, + incoming_instances_rel: nisl.Map, + live_access_rel: nisl.Map, + rel_finder: AccessRelationFinder, + happens_after: dict[str, HappensAfter], +) -> Mapping[str, HappensAfter]: + """ + Recursively finds conflicting accesses to *var* by *sink* and *source* to + determine the minimal required execution order between statement instances + of *source* and *sink*. + + :arg sink_id: The ID of the statement whose instances will be in the domain + of the resulting :class:`namedisl.Map`. + + :arg source_id: The ID of the statement whose instances will be in the range + of the resuling :class:`namedisl.Map`. + + :arg var: The variable for which we are performing data dependence analysis. + + :arg sink_access_type: A :class:`AccessType` describing whether *sink_id* + reads or writes *var*. This determines how live instances are removed from + *live_access_rel*. + + :arg incoming_instances_rel: The incoming :class:`namedisl.Map` describing + how each sink and source instance are related. + + :arg live_access_rel: A :class:`namedisl.Map` describing the set of live + accesses by *sink_id* to *var*. When conflicts are found, the conflicting + relation is used to remove elements from this relation. + + :arg rel_finder: A :class:`AccessRelationFinder` with access relations + constructed before entering this routine. + + :arg happens_after: A mapping from statement IDs to + :class:`loopy.HappensAfter` recording the dependencies from *sink* to all + statements in *happens_after*. + + :returns: The updated precise dependencies for *source*. + """ + + def record_conflicts(source_relation: nisl.Map) -> nisl.Map: + source_relation = _suffix_names(source_relation, "_before", DimType.in_) + conflicts = live_access_rel.apply_range(source_relation.reverse()) + + req_order = incoming_instances_rel & conflicts + previous = happens_after.get(source_id) + if not req_order.is_empty(): + happens_after[source_id] = ( + HappensAfter(req_order) + if previous is None + else HappensAfter(req_order | previous.instances_rel) # pyright: ignore[reportOperatorIssue] + ) + + return live_access_rel & req_order.apply_range(source_relation) + + def normalize_interface_and_compose( + sink_map: nisl.Map, source_map: nisl.Map + ) -> nisl.Map: + + sink_map = sink_map.rename_dims( + + (name, name[: len(name) - len("_before")]) + for name in sink_map.space.out_names + + ) + + source_map = source_map.rename_dims( + + (name, name[: len(name) - len("_after")]) + for name in source_map.space.in_names + + ) + + return sink_map.apply_range(source_map) + + match sink_access_type: + # compute raw + case AccessType.read: + if var in rel_finder.write_relations[source_id]: + source_relation = rel_finder.write_relations[source_id][var] + + caught_instances = record_conflicts(source_relation) + live_access_rel = live_access_rel - caught_instances + + # compute waw, war + case AccessType.write: + if var in rel_finder.write_relations[source_id]: + source_relation = rel_finder.write_relations[source_id][var] + + caught_instances = record_conflicts(source_relation) + live_access_rel = live_access_rel - caught_instances + + # don't update live_access_rel; does not find a "most recent writer" + if var in rel_finder.read_relations[source_id]: + source_relation = rel_finder.read_relations[source_id][var] + _ = record_conflicts(source_relation) + + case _: + raise ValueError("unknown access type") + + # recurse + if not live_access_rel.is_empty() and (sink_id != source_id): + source_insn = kernel.id_to_insn[source_id] + for src_dep_id, src_happens_after in source_insn.happens_after.items(): + if src_dep_id == source_id: + continue + + if src_happens_after.instances_rel is None: + raise ValueError( + "All `HappensAfter`s must have precise dependencies " + "defined to use precise dependency finding machinery." + ) + + outgoing_instances_rel = normalize_interface_and_compose( + incoming_instances_rel, src_happens_after.instances_rel + ).coalesce() + + _relax_strict_happens_after_inner( + kernel, + sink_id, + src_dep_id, + var, + sink_access_type, + outgoing_instances_rel, + live_access_rel, + rel_finder, + happens_after, + ) + + return happens_after + + +@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. + """ + + coarse_dependency_graph: dict[str, frozenset[str]] = {} + for insn in kernel.instructions: + coarse_dependency_graph[insn.id] = frozenset({ + dep for dep in insn.happens_after if dep != insn.id + }) + + topo_sort = compute_topological_order(coarse_dependency_graph) + + rel_finder = AccessRelationFinder(kernel) + for insn in kernel.instructions: + if isinstance(insn, MultiAssignmentBase): + for assignee in insn.assignees: + rel_finder(assignee, insn.id, AccessType.write) + rel_finder(insn.expression, insn.id, AccessType.read) + for pred in insn.predicates: + rel_finder(pred, insn.id, AccessType.read) + + # FIXME: clean up. kind of gross + new_insns: list[InstructionBase] = [] + for sink_id in topo_sort: + new_happens_after: dict[str, HappensAfter] = {} + old_happens_after = kernel.id_to_insn[sink_id].happens_after + for var, read_rel in rel_finder.read_relations[sink_id].items(): + read_rel = _suffix_names(read_rel, "_after", DimType.in_) + for source_id, happens_after in old_happens_after.items(): + if happens_after.instances_rel is None: + raise ValueError( + "All `HappensAfter`s must have precise dependencies " + "defined to use precise dependency finding machinery." + ) + + _relax_strict_happens_after_inner( + kernel, + sink_id, + source_id, + var, + AccessType.read, + happens_after.instances_rel, + read_rel, + rel_finder, + new_happens_after, + ) + + for var, write_rel in rel_finder.write_relations[sink_id].items(): + write_rel = _suffix_names(write_rel, "_after", DimType.in_) + for source_id, happens_after in old_happens_after.items(): + if happens_after.instances_rel is None: + raise ValueError( + "All `HappensAfter`s must have precise dependencies " + "defined to use precise dependency finding machinery." + ) + + _relax_strict_happens_after_inner( + kernel, + sink_id, + source_id, + var, + AccessType.write, + happens_after.instances_rel, + write_rel, + rel_finder, + new_happens_after, + ) + + new_insns.append( + kernel.id_to_insn[sink_id].copy( + happens_after=constantdict(new_happens_after) + ) + ) + + return kernel.copy(instructions=new_insns) diff --git a/loopy/kernel/instruction.py b/loopy/kernel/instruction.py index 813fe3071..435ae4971 100644 --- a/loopy/kernel/instruction.py +++ b/loopy/kernel/instruction.py @@ -47,7 +47,6 @@ from typing_extensions import Self, override import islpy as isl -import namedisl as nisl import pymbolic.primitives as p from pytools import ImmutableRecord, memoize_method from pytools.tag import Tag, Taggable, tag_dataclass @@ -59,6 +58,8 @@ if TYPE_CHECKING: + import namedisl as nisl + from pymbolic import Expression from loopy.kernel import LoopKernel @@ -120,8 +121,8 @@ 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 analyze and verify - data dependencies are respected after a :class:`loopy.LoopKernel` has been + 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:: instances_rel @@ -132,8 +133,8 @@ class HappensAfter: the range. Map dimensions are named according to the order of appearance of the - 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 + 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*, @@ -141,7 +142,7 @@ class HappensAfter: statement-level dependencies of prior versions of :mod:`loopy`. """ - instances_rel: nisl.Map | nisl.BasicMap | None + instances_rel: nisl.Map | None # }}} diff --git a/test/test_dependency.py b/test/test_dependency.py index 1080459f0..001d8b9b3 100644 --- a/test/test_dependency.py +++ b/test/test_dependency.py @@ -1,7 +1,10 @@ import namedisl as nisl +from pymbolic import var + import loopy as lp import loopy.kernel.dependency as dep +from loopy.symbolic import SubArrayRef def test_add_lexicographic_happens_after_is_strict_for_self() -> None: @@ -132,6 +135,436 @@ def test_add_lexicographic_happens_after_with_five_inames() -> None: ) +def test_access_relation_finder_keeps_instruction_maps_separate() -> None: + t_unit = lp.make_kernel( + "{ [i] : 0 <= i < N }", + """ + a[i] = 1 {id=S} + b[i] = 2 {id=T} + """, + ) + + kernel = t_unit.default_entrypoint + rel_find = dep.AccessRelationFinder(kernel) + rel_find(kernel.id_to_insn["S"].assignee, "S", dep.AccessType.write) + rel_find(kernel.id_to_insn["T"].assignee, "T", dep.AccessType.write) + + assert rel_find.write_relations["S"].keys() == {"a"} + assert rel_find.write_relations["T"].keys() == {"b"} + + +def test_access_relation_finder_distinguishes_reads_and_writes() -> None: + t_unit = lp.make_kernel( + "{ [i] : 1 <= i < N }", + "a[i] = a[i - 1] {id=S}", + ) + + kernel = t_unit.default_entrypoint + insn = kernel.id_to_insn["S"] + rel_find = dep.AccessRelationFinder(kernel) + rel_find(insn.assignee, insn.id, dep.AccessType.write) + rel_find(insn.expression, insn.id, dep.AccessType.read) + + assert rel_find.read_relations["S"]["a"].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_linear_subscript() -> None: + t_unit = lp.make_kernel( + "{ [i, j] : 0 <= i < NI and 0 <= j < NJ }", + "out[i, j] = a[[2*i + j]] {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] -> { + [i, j] -> [ax_0 = 2*i + j] : + 0 <= i < NI and 0 <= j < NJ + } + """) + ) + + +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) + 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_finds_direct_raw() -> None: + kernel = _relax_strict_happens_after( + """ + a[i, j] = 1 {id=S} + b[i, j] = a[i, j] {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_relax_strict_happens_after_finds_direct_waw() -> None: + kernel = _relax_strict_happens_after( + """ + a[i, j] = 1 {id=S} + a[i, j] = 2 {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_relax_strict_happens_after_finds_direct_war() -> None: + kernel = _relax_strict_happens_after( + """ + b[i, j] = a[i, j] {id=S} + a[i, j] = 2 {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_relax_strict_happens_after_finds_self_raw() -> None: + kernel = _relax_strict_happens_after( + "a[i] = a[i - 1] {id=S}", + "{ [i] : 1 <= i < N }", + ) + + required_order = kernel.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 + } + """) + ) + + +def test_relax_strict_happens_after_drops_conflict_free_edge() -> None: + kernel = _relax_strict_happens_after(""" + a[i] = 1 {id=S} + b[i] = 2 {id=T} + """) + + assert "S" not in kernel.id_to_insn["T"].happens_after + + +def test_relax_strict_happens_after_finds_recursive_raw() -> None: + kernel = _relax_strict_happens_after( + """ + a[i, j] = 1 {id=S} + b[i, j] = 2 {id=T} + c[i, j] = a[i, j] {id=U} + """, + "{ [i, j] : 0 <= i < NI and 0 <= j < NJ }", + ) + + 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(""" + [NI, NJ] -> { + [i_after, j_after] -> + [i_before = i_after, j_before = j_after] : + 0 <= i_after < NI and 0 <= j_after < NJ + } + """) + ) + assert "T" not in kernel.id_to_insn["U"].happens_after + + +def test_relax_strict_happens_after_stops_at_most_recent_writer() -> None: + kernel = _relax_strict_happens_after( + """ + a[i, j] = 1 {id=S} + a[i, j] = 2 {id=T} + c[i, j] = a[i, j] {id=U} + """, + "{ [i, j] : 0 <= i < NI and 0 <= j < NJ }", + ) + + required_order = kernel.id_to_insn["U"].happens_after["T"].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 + } + """) + ) + assert "S" not in kernel.id_to_insn["U"].happens_after + + +def test_relax_strict_happens_after_partitions_partial_writer_footprint() -> None: + kernel = _relax_strict_happens_after( + """ + a[i, j] = 1 {id=S} + a[2*i, j] = 2 {id=T} + c[i, j] = a[i, j] {id=U} + """, + "{ [i, j] : 0 <= i < NI and 0 <= j < NJ }", + ) + + recent_order = kernel.id_to_insn["U"].happens_after["T"].instances_rel + fallback_order = kernel.id_to_insn["U"].happens_after["S"].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 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 and + i_after mod 2 = 1 + } + """) + ) + + +def test_relax_strict_happens_after_composes_distinct_loop_nests() -> None: + t_unit = lp.make_kernel( + [ + "{ [i] : 0 <= i < 4 }", + "{ [j] : 0 <= j < 3 }", + "{ [k] : 0 <= k < 4 }", + ], + """ + a[i] = 1 {id=S} + b[j] = 2 {id=T} + c[k] = a[k] {id=U} + """, + ) + 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["U"].happens_after["S"].instances_rel + assert required_order is not None + assert required_order.equals( + nisl.make_map(""" + { + [k_after] -> [i_before = k_after] : 0 <= k_after < 4 + } + """) + ) + + +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_relax_strict_happens_after_drops_empty_same_variable_edge() -> None: + kernel = _relax_strict_happens_after(""" + a[i] = 1 {id=S} + b[i] = a[i + N] {id=T} + """) + + assert "S" not in kernel.id_to_insn["T"].happens_after + + +def test_relax_strict_happens_after_inner_uses_live_sink_accesses() -> None: + t_unit = lp.make_kernel( + "{ [i] : 0 <= i < N }", + """ + a[i] = 1 {id=S} + b[i] = a[i] {id=T} + """, + ) + kernel = dep.add_lexicographic_happens_after(t_unit).default_entrypoint + + rel_finder = dep.AccessRelationFinder(kernel) + for insn in kernel.instructions: + assert isinstance(insn, lp.MultiAssignmentBase) + rel_finder(insn.assignee, insn.id, dep.AccessType.write) + rel_finder(insn.expression, insn.id, dep.AccessType.read) + + incoming_relation = kernel.id_to_insn["T"].happens_after["S"].instances_rel + assert incoming_relation is not None + + live_access_relation = rel_finder.read_relations["T"]["a"].rename_dims(( + ("i", "i_after"), + )) + live_access_relation = live_access_relation & nisl.make_map(""" + [N] -> { + [i_after] -> [ax_0] : i_after = 0 + } + """) + + happens_after = dep._relax_strict_happens_after_inner( + kernel, + "T", + "S", + "a", + dep.AccessType.read, + incoming_relation, + live_access_relation, + rel_finder, + {}, + ) + + required_order = happens_after["S"].instances_rel + assert required_order is not None + assert required_order.equals( + nisl.make_map(""" + [N] -> { + [i_after = 0] -> [i_before = 0] : N > 0 + } + """) + ) + + if __name__ == "__main__": import sys From 459a6c5c92dd44214d0edf36acba6f920c5510ff Mon Sep 17 00:00:00 2001 From: Addison Date: Sun, 19 Jul 2026 01:31:18 -0500 Subject: [PATCH 04/30] formatting --- loopy/kernel/dependency.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/loopy/kernel/dependency.py b/loopy/kernel/dependency.py index aa6dc5690..64027259d 100644 --- a/loopy/kernel/dependency.py +++ b/loopy/kernel/dependency.py @@ -348,17 +348,13 @@ def normalize_interface_and_compose( ) -> nisl.Map: sink_map = sink_map.rename_dims( - - (name, name[: len(name) - len("_before")]) - for name in sink_map.space.out_names - + (name, name[: len(name) - len("_before")]) + for name in sink_map.space.out_names ) source_map = source_map.rename_dims( - - (name, name[: len(name) - len("_after")]) - for name in source_map.space.in_names - + (name, name[: len(name) - len("_after")]) + for name in source_map.space.in_names ) return sink_map.apply_range(source_map) From 97c74cfaae7325b7ea846e466cc75db933448feb Mon Sep 17 00:00:00 2001 From: Addison Date: Sun, 19 Jul 2026 01:37:46 -0500 Subject: [PATCH 05/30] clean up pass --- loopy/kernel/dependency.py | 156 +++++++++++++++++-------------------- test/test_dependency.py | 5 ++ 2 files changed, 75 insertions(+), 86 deletions(-) diff --git a/loopy/kernel/dependency.py b/loopy/kernel/dependency.py index 64027259d..10e50f8de 100644 --- a/loopy/kernel/dependency.py +++ b/loopy/kernel/dependency.py @@ -34,7 +34,6 @@ from pymbolic.typing import Expression from loopy.kernel import LoopKernel - from loopy.typing import ShapeType @final @@ -44,15 +43,15 @@ class AccessType(Enum): class AccessRelationFinder(WalkMapper[[str, AccessType]]): + """Collect per-instruction statement-instance-to-cell access relations.""" + kernel: LoopKernel - variables: frozenset[str] _additional_inames: frozenset[str] _read_relations: dict[str, dict[str, nisl.Map]] _write_relations: dict[str, dict[str, nisl.Map]] def __init__(self, kernel: LoopKernel): self.kernel = kernel - self.variables = frozenset(kernel.all_variable_names()) self._additional_inames = frozenset() self._read_relations = {insn.id: {} for insn in kernel.instructions} self._write_relations = {insn.id: {} for insn in kernel.instructions} @@ -63,9 +62,6 @@ def _get_access_relation( self, domain: nisl.Set, subscript: tuple[Expression, ...], - assumptions: nisl.BasicSet | None = None, - shape: ShapeType | None = None, - allowed_constant_names: frozenset[str] | None = None, ) -> nisl.Map: instance_names = domain.space.dimtype_to_names[DimType.out] cell_names = tuple(f"ax_{axis}" for axis in range(len(subscript))) @@ -97,7 +93,7 @@ def _insn_reads_var(self, insn_id: str, var: str) -> bool: ) def _insn_accesses_var(self, insn_id: str, var: str) -> bool: - return self._insn_reads_var(insn_id, var) | self._insn_writes_var( + return self._insn_reads_var(insn_id, var) or self._insn_writes_var( insn_id, var ) @@ -139,15 +135,15 @@ def _record_access( @cached_property def read_relations(self) -> Mapping[str, Mapping[str, nisl.Map]]: return constantdict({ - insn_id: constantdict(self._read_relations[insn_id]) - for insn_id in self._read_relations + insn_id: constantdict(relations) + for insn_id, relations in self._read_relations.items() }) @cached_property def write_relations(self) -> Mapping[str, Mapping[str, nisl.Map]]: return constantdict({ - insn_id: constantdict(self._write_relations[insn_id]) - for insn_id in self._write_relations + insn_id: constantdict(relations) + for insn_id, relations in self._write_relations.items() }) @override @@ -196,10 +192,10 @@ def map_sub_array_ref( def _suffix_names( - obj: NamedIslObjectT, prefix: str, dim_type: DimType + obj: NamedIslObjectT, suffix: str, dim_type: DimType ) -> NamedIslObjectT: return obj.rename_dims( - (name, name + prefix) for name in obj.space.dimtype_to_names[dim_type] + (name, name + suffix) for name in obj.space.dimtype_to_names[dim_type] ) @@ -219,20 +215,19 @@ def add_lexicographic_happens_after(kernel: LoopKernel) -> LoopKernel: new_insns: list[InstructionBase] = [] for i, insn in enumerate(kernel.instructions): - new_happens_after = {} + new_happens_after: dict[str, HappensAfter] = {} - preds = (insn,) if i == 0 else (insn, kernel.instructions[i - 1]) + sources = (insn,) if i == 0 else (insn, kernel.instructions[i - 1]) - # FIXME: yuck after_domain = nisl.make_set( kernel.get_inames_domain(insn.within_inames).to_set() ) after_inames = after_domain.space.dimtype_to_names[DimType.out] after_domain = _suffix_names(after_domain, "_after", DimType.out) - for pred in preds: + for source in sources: before_domain = nisl.make_set( - kernel.get_inames_domain(pred.within_inames).to_set() + kernel.get_inames_domain(source.within_inames).to_set() ) before_inames = before_domain.space.dimtype_to_names[DimType.out] @@ -265,7 +260,7 @@ def add_lexicographic_happens_after(kernel: LoopKernel) -> LoopKernel: ) equal_prefix = equal_prefix & after_aff.eq_set(before_aff) - if pred.id == insn.id: + if source.id == insn.id: ordered_instances = strict_lex else: ordered_instances = strict_lex | equal_prefix @@ -274,7 +269,7 @@ def add_lexicographic_happens_after(kernel: LoopKernel) -> LoopKernel: in_names=tuple(f"{name}_after" for name in after_inames) ) - new_happens_after[pred.id] = HappensAfter( + new_happens_after[source.id] = HappensAfter( instances_rel=instances_rel ) @@ -303,7 +298,7 @@ def _relax_strict_happens_after_inner( of the resulting :class:`namedisl.Map`. :arg source_id: The ID of the statement whose instances will be in the range - of the resuling :class:`namedisl.Map`. + of the resulting :class:`namedisl.Map`. :arg var: The variable for which we are performing data dependence analysis. @@ -325,56 +320,61 @@ def _relax_strict_happens_after_inner( :class:`loopy.HappensAfter` recording the dependencies from *sink* to all statements in *happens_after*. - :returns: The updated precise dependencies for *source*. + :returns: The updated precise dependencies for *sink_id*. """ def record_conflicts(source_relation: nisl.Map) -> nisl.Map: source_relation = _suffix_names(source_relation, "_before", DimType.in_) + + # live_access_rel; source_relation^-1 conflicts = live_access_rel.apply_range(source_relation.reverse()) - req_order = incoming_instances_rel & conflicts + # Only conflicts ordered along this graph path are required. + required_order = incoming_instances_rel & conflicts previous = happens_after.get(source_id) - if not req_order.is_empty(): - happens_after[source_id] = ( - HappensAfter(req_order) - if previous is None - else HappensAfter(req_order | previous.instances_rel) # pyright: ignore[reportOperatorIssue] - ) + if not required_order.is_empty(): + if previous is None: + combined_order = required_order + else: + assert previous.instances_rel is not None + combined_order = required_order | previous.instances_rel + + happens_after[source_id] = HappensAfter(combined_order) - return live_access_rel & req_order.apply_range(source_relation) + # Retire only live accesses supplied by an ordered source instance. + return live_access_rel & required_order.apply_range(source_relation) def normalize_interface_and_compose( - sink_map: nisl.Map, source_map: nisl.Map + incoming_relation: nisl.Map, next_edge_relation: nisl.Map ) -> nisl.Map: - - sink_map = sink_map.rename_dims( + incoming_relation = incoming_relation.rename_dims( (name, name[: len(name) - len("_before")]) - for name in sink_map.space.out_names + for name in incoming_relation.space.out_names ) - source_map = source_map.rename_dims( + next_edge_relation = next_edge_relation.rename_dims( (name, name[: len(name) - len("_after")]) - for name in source_map.space.in_names + for name in next_edge_relation.space.in_names ) - return sink_map.apply_range(source_map) + return incoming_relation.apply_range(next_edge_relation) match sink_access_type: - # compute raw + # Read-after-write case AccessType.read: if var in rel_finder.write_relations[source_id]: source_relation = rel_finder.write_relations[source_id][var] - caught_instances = record_conflicts(source_relation) - live_access_rel = live_access_rel - caught_instances + caught_accesses = record_conflicts(source_relation) + live_access_rel = live_access_rel - caught_accesses - # compute waw, war + # Write-after-write and write-after-read case AccessType.write: if var in rel_finder.write_relations[source_id]: source_relation = rel_finder.write_relations[source_id][var] - caught_instances = record_conflicts(source_relation) - live_access_rel = live_access_rel - caught_instances + caught_accesses = record_conflicts(source_relation) + live_access_rel = live_access_rel - caught_accesses # don't update live_access_rel; does not find a "most recent writer" if var in rel_finder.read_relations[source_id]: @@ -384,7 +384,7 @@ def normalize_interface_and_compose( case _: raise ValueError("unknown access type") - # recurse + # Continue backward through the strict-order graph. if not live_access_rel.is_empty() and (sink_id != source_id): source_insn = kernel.id_to_insn[source_id] for src_dep_id, src_happens_after in source_insn.happens_after.items(): @@ -432,7 +432,7 @@ def relax_strict_happens_after(kernel: LoopKernel) -> LoopKernel: dep for dep in insn.happens_after if dep != insn.id }) - topo_sort = compute_topological_order(coarse_dependency_graph) + topological_order = compute_topological_order(coarse_dependency_graph) rel_finder = AccessRelationFinder(kernel) for insn in kernel.instructions: @@ -443,53 +443,37 @@ def relax_strict_happens_after(kernel: LoopKernel) -> LoopKernel: for pred in insn.predicates: rel_finder(pred, insn.id, AccessType.read) - # FIXME: clean up. kind of gross new_insns: list[InstructionBase] = [] - for sink_id in topo_sort: + for sink_id in topological_order: new_happens_after: dict[str, HappensAfter] = {} old_happens_after = kernel.id_to_insn[sink_id].happens_after - for var, read_rel in rel_finder.read_relations[sink_id].items(): - read_rel = _suffix_names(read_rel, "_after", DimType.in_) - for source_id, happens_after in old_happens_after.items(): - if happens_after.instances_rel is None: - raise ValueError( - "All `HappensAfter`s must have precise dependencies " - "defined to use precise dependency finding machinery." - ) - - _relax_strict_happens_after_inner( - kernel, - sink_id, - source_id, - var, - AccessType.read, - happens_after.instances_rel, - read_rel, - rel_finder, - new_happens_after, + 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_ ) - - for var, write_rel in rel_finder.write_relations[sink_id].items(): - write_rel = _suffix_names(write_rel, "_after", DimType.in_) - for source_id, happens_after in old_happens_after.items(): - if happens_after.instances_rel is None: - raise ValueError( - "All `HappensAfter`s must have precise dependencies " - "defined to use precise dependency finding machinery." + for source_id, happens_after in old_happens_after.items(): + if happens_after.instances_rel is None: + raise ValueError( + "All `HappensAfter`s must have precise dependencies " + "defined to use precise dependency finding machinery." + ) + + _relax_strict_happens_after_inner( + kernel, + sink_id, + source_id, + var, + sink_access_type, + happens_after.instances_rel, + access_relation, + rel_finder, + new_happens_after, ) - _relax_strict_happens_after_inner( - kernel, - sink_id, - source_id, - var, - AccessType.write, - happens_after.instances_rel, - write_rel, - rel_finder, - new_happens_after, - ) - new_insns.append( kernel.id_to_insn[sink_id].copy( happens_after=constantdict(new_happens_after) diff --git a/test/test_dependency.py b/test/test_dependency.py index 001d8b9b3..9c495c370 100644 --- a/test/test_dependency.py +++ b/test/test_dependency.py @@ -5,6 +5,9 @@ import loopy as lp import loopy.kernel.dependency as dep 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: @@ -233,6 +236,8 @@ def test_access_relation_finder_handles_sub_array_ref() -> None: 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")], From 5f96a3ec8f9eaf3a9d9be187beca0dba2d8b5865 Mon Sep 17 00:00:00 2001 From: Addison Date: Sun, 19 Jul 2026 11:14:02 -0500 Subject: [PATCH 06/30] hook kernel/dependency.py into docs; add namedisl to docs config --- doc/conf.py | 1 + doc/ref_internals.rst | 4 ++++ 2 files changed, 5 insertions(+) diff --git a/doc/conf.py b/doc/conf.py index 5ac143233..ed66683c2 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -17,6 +17,7 @@ "constantdict": ("https://matthiasdiener.github.io/constantdict/", None), "genpy": ("https://documen.tician.de/genpy", None), "islpy": ("https://documen.tician.de/islpy", None), + "namedisl": ("https://documen.tician.de/namedisl", None), "numpy": ("https://numpy.org/doc/stable/", None), "pymbolic": ("https://documen.tician.de/pymbolic", None), "pyopencl": ("https://documen.tician.de/pyopencl", None), diff --git a/doc/ref_internals.rst b/doc/ref_internals.rst index 02edceca6..9e817a22f 100644 --- a/doc/ref_internals.rst +++ b/doc/ref_internals.rst @@ -60,3 +60,7 @@ Schedule .. automodule:: loopy.schedule .. automodule:: loopy.schedule.tools .. automodule:: loopy.schedule.tree + +Dependence Analysis +------------------- +.. automodule:: loopy.kernel.dependency From a009c5be35c61e46c146366b3cac16c8488a5c10 Mon Sep 17 00:00:00 2001 From: Addison Date: Sun, 19 Jul 2026 14:01:55 -0500 Subject: [PATCH 07/30] temporarily use unnamed HappensAfter.instances_rel --- loopy/kernel/dependency.py | 15 ++++++---- loopy/kernel/instruction.py | 2 +- test/test_dependency.py | 56 +++++++++++++++++++++++++++++++------ 3 files changed, 59 insertions(+), 14 deletions(-) diff --git a/loopy/kernel/dependency.py b/loopy/kernel/dependency.py index 10e50f8de..a4a184da0 100644 --- a/loopy/kernel/dependency.py +++ b/loopy/kernel/dependency.py @@ -270,7 +270,7 @@ def add_lexicographic_happens_after(kernel: LoopKernel) -> LoopKernel: ) new_happens_after[source.id] = HappensAfter( - instances_rel=instances_rel + instances_rel=instances_rel.as_isl() ) new_insns.append(insn.copy(happens_after=new_happens_after)) @@ -337,9 +337,12 @@ def record_conflicts(source_relation: nisl.Map) -> nisl.Map: combined_order = required_order else: assert previous.instances_rel is not None - combined_order = required_order | previous.instances_rel + # FIXME: remove named conversion + previous_instances_rel = nisl.make_map(previous.instances_rel) + combined_order = required_order | previous_instances_rel - happens_after[source_id] = HappensAfter(combined_order) + # FIXME: remove unnamed conversion + happens_after[source_id] = HappensAfter(combined_order.as_isl()) # Retire only live accesses supplied by an ordered source instance. return live_access_rel & required_order.apply_range(source_relation) @@ -397,8 +400,10 @@ def normalize_interface_and_compose( "defined to use precise dependency finding machinery." ) + # FIXME: removed named conversion + src_instances_rel = nisl.make_map(src_happens_after.instances_rel) outgoing_instances_rel = normalize_interface_and_compose( - incoming_instances_rel, src_happens_after.instances_rel + incoming_instances_rel, src_instances_rel ).coalesce() _relax_strict_happens_after_inner( @@ -468,7 +473,7 @@ def relax_strict_happens_after(kernel: LoopKernel) -> LoopKernel: source_id, var, sink_access_type, - happens_after.instances_rel, + nisl.make_map(happens_after.instances_rel), access_relation, rel_finder, new_happens_after, diff --git a/loopy/kernel/instruction.py b/loopy/kernel/instruction.py index 435ae4971..28b3c013f 100644 --- a/loopy/kernel/instruction.py +++ b/loopy/kernel/instruction.py @@ -142,7 +142,7 @@ class HappensAfter: statement-level dependencies of prior versions of :mod:`loopy`. """ - instances_rel: nisl.Map | None + instances_rel: isl.Map | None # }}} diff --git a/test/test_dependency.py b/test/test_dependency.py index 9c495c370..4dcd7f0ac 100644 --- a/test/test_dependency.py +++ b/test/test_dependency.py @@ -26,6 +26,10 @@ def test_add_lexicographic_happens_after_is_strict_for_self() -> None: assert self_relation is not None assert previous_relation is not None + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + self_relation = nisl.make_map(self_relation) + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + previous_relation = nisl.make_map(previous_relation) assert self_relation.equals( nisl.make_map(""" [N] -> { @@ -54,6 +58,8 @@ def test_add_lexicographic_happens_after_uses_domain_dimension_order() -> None: self_relation = kernel.id_to_insn["S"].happens_after["S"].instances_rel assert self_relation is not None + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + self_relation = nisl.make_map(self_relation) assert self_relation.equals( nisl.make_map(""" [NZ, NA] -> { @@ -85,6 +91,8 @@ def test_add_lexicographic_happens_after_orders_distinct_loop_nests() -> None: cross_relation = kernel.id_to_insn["T"].happens_after["S"].instances_rel assert cross_relation is not None + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + cross_relation = nisl.make_map(cross_relation) assert cross_relation.equals( nisl.make_map(""" [N, M] -> { @@ -114,6 +122,8 @@ def test_add_lexicographic_happens_after_with_five_inames() -> None: self_relation = kernel.id_to_insn["S"].happens_after["S"].instances_rel assert self_relation is not None + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + self_relation = nisl.make_map(self_relation) assert self_relation.equals( nisl.make_map(""" { @@ -276,6 +286,8 @@ def test_relax_strict_happens_after_finds_direct_raw() -> None: required_order = kernel.id_to_insn["T"].happens_after["S"].instances_rel assert required_order is not None + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + required_order = nisl.make_map(required_order) assert required_order.equals( nisl.make_map(""" [NI, NJ] -> { @@ -298,6 +310,8 @@ def test_relax_strict_happens_after_finds_direct_waw() -> None: required_order = kernel.id_to_insn["T"].happens_after["S"].instances_rel assert required_order is not None + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + required_order = nisl.make_map(required_order) assert required_order.equals( nisl.make_map(""" [NI, NJ] -> { @@ -320,6 +334,8 @@ def test_relax_strict_happens_after_finds_direct_war() -> None: required_order = kernel.id_to_insn["T"].happens_after["S"].instances_rel assert required_order is not None + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + required_order = nisl.make_map(required_order) assert required_order.equals( nisl.make_map(""" [NI, NJ] -> { @@ -339,6 +355,8 @@ def test_relax_strict_happens_after_finds_self_raw() -> None: required_order = kernel.id_to_insn["S"].happens_after["S"].instances_rel assert required_order is not None + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + required_order = nisl.make_map(required_order) assert required_order.equals( nisl.make_map(""" [N] -> { @@ -369,6 +387,8 @@ def test_relax_strict_happens_after_finds_recursive_raw() -> None: required_order = kernel.id_to_insn["U"].happens_after["S"].instances_rel assert required_order is not None + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + required_order = nisl.make_map(required_order) assert required_order.equals( nisl.make_map(""" [NI, NJ] -> { @@ -393,6 +413,8 @@ def test_relax_strict_happens_after_stops_at_most_recent_writer() -> None: required_order = kernel.id_to_insn["U"].happens_after["T"].instances_rel assert required_order is not None + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + required_order = nisl.make_map(required_order) assert required_order.equals( nisl.make_map(""" [NI, NJ] -> { @@ -405,7 +427,9 @@ def test_relax_strict_happens_after_stops_at_most_recent_writer() -> None: assert "S" not in kernel.id_to_insn["U"].happens_after -def test_relax_strict_happens_after_partitions_partial_writer_footprint() -> None: +def test_relax_strict_happens_after_partitions_partial_writer_footprint() -> ( + None +): kernel = _relax_strict_happens_after( """ a[i, j] = 1 {id=S} @@ -419,6 +443,10 @@ def test_relax_strict_happens_after_partitions_partial_writer_footprint() -> Non fallback_order = kernel.id_to_insn["U"].happens_after["S"].instances_rel assert recent_order is not None assert fallback_order is not None + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + recent_order = nisl.make_map(recent_order) + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + fallback_order = nisl.make_map(fallback_order) assert recent_order.equals( nisl.make_map(""" [NI, NJ] -> { @@ -462,6 +490,8 @@ def test_relax_strict_happens_after_composes_distinct_loop_nests() -> None: required_order = kernel.id_to_insn["U"].happens_after["S"].instances_rel assert required_order is not None + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + required_order = nisl.make_map(required_order) assert required_order.equals( nisl.make_map(""" { @@ -489,18 +519,24 @@ def test_relax_strict_happens_after_unions_branched_paths() -> None: "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 = 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 + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + required_order = nisl.make_map(required_order) assert required_order.equals( nisl.make_map(""" [N] -> { @@ -537,6 +573,8 @@ def test_relax_strict_happens_after_inner_uses_live_sink_accesses() -> None: incoming_relation = kernel.id_to_insn["T"].happens_after["S"].instances_rel assert incoming_relation is not None + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + incoming_relation = nisl.make_map(incoming_relation) live_access_relation = rel_finder.read_relations["T"]["a"].rename_dims(( ("i", "i_after"), @@ -561,6 +599,8 @@ def test_relax_strict_happens_after_inner_uses_live_sink_accesses() -> None: required_order = happens_after["S"].instances_rel assert required_order is not None + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + required_order = nisl.make_map(required_order) assert required_order.equals( nisl.make_map(""" [N] -> { From 10affd5fff7862b5ca27369f2e76ba400ae912a5 Mon Sep 17 00:00:00 2001 From: Addison Date: Mon, 20 Jul 2026 22:48:18 -0500 Subject: [PATCH 08/30] add precise dependency verification --- loopy/codegen/__init__.py | 7 + loopy/kernel/dependency.py | 20 ++ loopy/schedule/verification.py | 523 +++++++++++++++++++++++++++++++++ test/test_dependency.py | 513 ++++++++++++++++++++++++++++++++ 4 files changed, 1063 insertions(+) create mode 100644 loopy/schedule/verification.py diff --git a/loopy/codegen/__init__.py b/loopy/codegen/__init__.py index 45ecb64e9..f988b1544 100644 --- a/loopy/codegen/__init__.py +++ b/loopy/codegen/__init__.py @@ -332,6 +332,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 index a4a184da0..9b84911fd 100644 --- a/loopy/kernel/dependency.py +++ b/loopy/kernel/dependency.py @@ -12,6 +12,7 @@ from pytools.graph import compute_topological_order from loopy import for_each_kernel +from loopy.diagnostic import LoopyError from loopy.kernel.instruction import ( HappensAfter, InstructionBase, @@ -42,6 +43,25 @@ class AccessType(Enum): write = 1 +def has_precise_dependencies(kernel: LoopKernel) -> bool: + has_precise = False + has_legacy = False + for insn in kernel.instructions: + for happens_after in insn.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 + + class AccessRelationFinder(WalkMapper[[str, AccessType]]): """Collect per-instruction statement-instance-to-cell access relations.""" diff --git a/loopy/schedule/verification.py b/loopy/schedule/verification.py new file mode 100644 index 000000000..4fa9c827c --- /dev/null +++ b/loopy/schedule/verification.py @@ -0,0 +1,523 @@ +from constantdict import constantdict +import islpy as isl +import namedisl as nisl +from namedisl import DimType + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass + +from typing_extensions import override + +from loopy import KernelState, LoopKernel, for_each_kernel +from loopy.diagnostic import LoopyError +from loopy.kernel.instruction import BarrierKind, HappensAfter +from loopy.schedule import ( + Barrier, + CallKernel, + EnterLoop, + LeaveLoop, + ReturnFromKernel, + RunInstruction, + ScheduleItem, +) + + +@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." + ) + + 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)): + 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) + + if dom_str: + full_str = dom_str + ", " + ran_str + else: + full_str = ran_str + + # FIXME: isl -> named conversion + domain = nisl.make_set( + kernel.get_inames_domain(stmt.within_inames).as_set() + ).project_out_except([*stmt.within_inames, *kernel.all_params()]) + + 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)}] : " + f"{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 + + # FIXME: isl -> named conversion + domain = nisl.make_set( + kernel.get_inames_domain(frozenset(inames)).as_set() + ).project_out_except([*inames, *kernel.all_params()]) + + 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( + timestamp_names: Sequence[str], +) -> nisl.Map: + later_names = tuple(f"{name}_later" for name in timestamp_names) + earlier_names = tuple(f"{name}_earlier" for name in timestamp_names) + + joint = nisl.make_set( + f"{{ [{', '.join([*later_names, *earlier_names])}] }}" + ) + affs = joint.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, +) -> tuple[Mapping[str, nisl.Map], Sequence[nisl.Map], nisl.Map]: + max_stmt_tstamp_len = max( + len(record.timestamp) for _, record in prec_sched.statements.items() + ) + + 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 = [f"__ts_{i}" for i in range(max_tstamp_len)] + + 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(timestamp_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: str, +) -> nisl.Map: + return _suffix_dim_names(relation, DimType.out, f"_{role}") + + +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, +) -> nisl.Map: + axis_inames = _hardware_axis_inames(kernel, stmt_id, include_local_axes) + input_names = instance_domain.space.dimtype_to_names[DimType.out] + hardware_names = tuple( + f"__{kind}_{axis}" for kind, axis in sorted(axis_inames) + ) + constraints = " and ".join( + f"__{kind}_{axis} = {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)}]" + 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, +) -> 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" + ) + + sink_hardware = _build_hardware_id_relation( + kernel, sink_id, sink_domain, "_after", include_local_axes + ) + source_hardware = _build_hardware_id_relation( + kernel, source_id, source_domain, "_before", include_local_axes + ) + 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, +) -> nisl.Map: + sink = _suffix_dim_names(stmt_relns[sink_id], DimType.in_, "_after") + source = _suffix_dim_names(stmt_relns[source_id], DimType.in_, "_before") + + enforced = ( + _timestamp_relation_for_role(sink, "later") + .apply_range(timestamp_lex) + .apply_range(_timestamp_relation_for_role(source, "earlier").reverse()) + ) + enforced = enforced & _build_same_hardware_scope_relation( + kernel, + sink_id, + source_id, + sink.domain(), + source.domain(), + include_local_axes=True, + ) + + 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") + .apply_range(timestamp_lex) + .apply_range( + _timestamp_relation_for_role(barrier, "earlier").reverse() + ) + ) + barrier_to_source = ( + _timestamp_relation_for_role(barrier, "later") + .apply_range(timestamp_lex) + .apply_range( + _timestamp_relation_for_role(source, "earlier").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, + ) + ) + + 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) + stmt_relns, barrier_relns, timestamp_lex = _build_timestamp_relations( + kernel, prec_sched + ) + + 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 = nisl.make_map( + happens_after.instances_rel + .reset_tuple_id(isl.dim_type.in_) + .reset_tuple_id(isl.dim_type.out) + ) + enforced = _build_enforced_order( + kernel, + sink.id, + source_id, + prec_sched, + stmt_relns, + barrier_relns, + timestamp_lex, + ) + 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/test/test_dependency.py b/test/test_dependency.py index 4dcd7f0ac..7c44224e1 100644 --- a/test/test_dependency.py +++ b/test/test_dependency.py @@ -1,9 +1,31 @@ import namedisl as nisl +import pytest from pymbolic import var import loopy as lp import loopy.kernel.dependency as dep +from loopy.diagnostic import LoopyError +from loopy.kernel.instruction import HappensAfter + +from loopy.schedule.verification import ( + _BarrierRecord, + _PreciseSchedule, + _StatementRecord, + _build_enforced_order, + _build_strict_lexicographic_order, + _build_timestamp_relations, + _get_timestamp_points_from_linearization, + verify_happens_after_is_enforced, +) +from loopy.schedule import ( + Barrier, + CallKernel, + EnterLoop, + LeaveLoop, + ReturnFromKernel, + RunInstruction, +) from loopy.symbolic import SubArrayRef from loopy.version import ( LOOPY_USE_LANGUAGE_VERSION_2018_2, # ruff:ignore[unused-import] @@ -299,6 +321,37 @@ def test_relax_strict_happens_after_finds_direct_raw() -> None: ) +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_finds_direct_waw() -> None: kernel = _relax_strict_happens_after( """ @@ -610,6 +663,466 @@ def test_relax_strict_happens_after_inner_uses_live_sink_accesses() -> None: ) +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: + kernel = 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} + """, + ).default_entrypoint.copy( + state=lp.KernelState.LINEARIZED, + linearization=( + CallKernel("device_program"), + RunInstruction("before"), + EnterLoop("i"), + RunInstruction("outer"), + EnterLoop("j"), + EnterLoop("k"), + RunInstruction("deep"), + LeaveLoop("k"), + LeaveLoop("j"), + RunInstruction("after_inner"), + LeaveLoop("i"), + RunInstruction("after"), + ReturnFromKernel("device_program"), + ), + ) + + 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"}) + kernel = t_unit.default_entrypoint.copy( + state=lp.KernelState.LINEARIZED, + linearization=( + CallKernel("device_program"), + EnterLoop("i"), + RunInstruction("S"), + LeaveLoop("i"), + ReturnFromKernel("device_program"), + ), + ) + + 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 + ) + 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_strict_lexicographic_timestamp_order() -> None: + order = _build_strict_lexicographic_order(("t0", "t1", "t2")) + + 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 + ) + 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: + local_barrier = Barrier( + comment="local synchronization", + synchronization_kind="local", + mem_kind="local", + originating_insn_id=None, + ) + t_unit = lp.make_kernel( + """ + { + [g, l, i] : + 0 <= g < 2 and 0 <= l < 4 and 0 <= i < 2 + } + """, + """ + a[g, l, i] = g + l + i {id=source} + b[g, l, i] = a[g, l, i] {id=sink} + """, + ) + t_unit = lp.tag_inames(t_unit, {"g": "g.0", "l": "l.0"}) + kernel = t_unit.default_entrypoint.copy( + state=lp.KernelState.LINEARIZED, + linearization=( + CallKernel("device_program"), + EnterLoop("i"), + RunInstruction("source"), + local_barrier, + RunInstruction("sink"), + LeaveLoop("i"), + ReturnFromKernel("device_program"), + ), + ) + + precise_schedule = _get_timestamp_points_from_linearization(kernel) + stmt_relations, barrier_relations, timestamp_order = ( + _build_timestamp_relations(kernel, precise_schedule) + ) + enforced = _build_enforced_order( + kernel, + "sink", + "source", + precise_schedule, + stmt_relations, + barrier_relations, + timestamp_order, + ) + + 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: + global_barrier = Barrier( + comment="global synchronization", + synchronization_kind="global", + mem_kind="global", + originating_insn_id=None, + ) + t_unit = lp.make_kernel( + "{ [g, l] : 0 <= g < 2 and 0 <= l < 4 }", + """ + a[g, l] = g + l {id=source} + b[g, l] = a[g, l] {id=sink} + """, + ) + t_unit = lp.tag_inames(t_unit, {"g": "g.0", "l": "l.0"}) + kernel = t_unit.default_entrypoint.copy( + state=lp.KernelState.LINEARIZED, + linearization=( + CallKernel("phase0"), + RunInstruction("source"), + ReturnFromKernel("phase0"), + global_barrier, + CallKernel("phase1"), + RunInstruction("sink"), + ReturnFromKernel("phase1"), + ), + ) + + precise_schedule = _get_timestamp_points_from_linearization(kernel) + stmt_relations, barrier_relations, timestamp_order = ( + _build_timestamp_relations(kernel, precise_schedule) + ) + enforced = _build_enforced_order( + kernel, + "sink", + "source", + precise_schedule, + stmt_relations, + barrier_relations, + timestamp_order, + ) + + 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_verification() -> None: + t_unit = lp.make_kernel( + "{ [i] : 0 <= i < N }", + """ + a[i] = 2 * a[i] {id=S} + b[i] = 2 * b[i] {id=T} + c[i] = a[i] + b[i] {id=U} + """, + ) + + t_unit = dep.add_lexicographic_happens_after(t_unit) + t_unit = dep.relax_strict_happens_after(t_unit) + t_unit = lp.preprocess_program(t_unit) + t_unit = lp.linearize(t_unit) + t_unit = verify_happens_after_is_enforced(t_unit) + lp.generate_code_v2(t_unit) + + +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_handles_explicit_barrier_instruction() -> None: + t_unit = lp.make_kernel( + "{ [i] : 0 <= i < N }", + """ + a[i] = i {id=S} + ... gbarrier {id=B} + b[i] = a[i] {id=T} + """, + ) + t_unit = dep.add_lexicographic_happens_after(t_unit) + + kernel = t_unit.default_entrypoint + barrier_insn = kernel.id_to_insn["B"] + barrier = Barrier( + comment="explicit global barrier", + synchronization_kind=barrier_insn.synchronization_kind, + mem_kind=barrier_insn.mem_kind, + originating_insn_id="B", + ) + kernel = kernel.copy( + state=lp.KernelState.LINEARIZED, + linearization=( + CallKernel("phase0"), + EnterLoop("i"), + RunInstruction("S"), + LeaveLoop("i"), + ReturnFromKernel("phase0"), + barrier, + CallKernel("phase1"), + EnterLoop("i"), + RunInstruction("T"), + LeaveLoop("i"), + ReturnFromKernel("phase1"), + ), + ) + t_unit = t_unit.with_kernel(kernel) + + verify_happens_after_is_enforced(t_unit) + + if __name__ == "__main__": import sys From f88b459dc5afe7cde66a910b89b7da9afd5e71f3 Mon Sep 17 00:00:00 2001 From: Addison Date: Mon, 20 Jul 2026 23:02:31 -0500 Subject: [PATCH 09/30] add namedisl git url to pyproject.toml temporarily --- pyproject.toml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 2b502eb5e..fe3cb56a4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,6 +44,8 @@ dependencies = [ "typing-extensions>=4", "strenum>=0.4.15", + + "namedisl @ git+https://github.com/inducer/namedisl.git@main" ] [project.optional-dependencies] pyopencl = [ From 007bca968bddf140b82f4650fe031fe728e0aa22 Mon Sep 17 00:00:00 2001 From: Addison Date: Mon, 20 Jul 2026 23:03:49 -0500 Subject: [PATCH 10/30] fix ruff complaints --- loopy/kernel/instruction.py | 1 - loopy/schedule/verification.py | 25 +++++++++++++------------ test/test_dependency.py | 21 ++++++++++----------- 3 files changed, 23 insertions(+), 24 deletions(-) diff --git a/loopy/kernel/instruction.py b/loopy/kernel/instruction.py index 28b3c013f..0c0e2dd01 100644 --- a/loopy/kernel/instruction.py +++ b/loopy/kernel/instruction.py @@ -58,7 +58,6 @@ if TYPE_CHECKING: - import namedisl as nisl from pymbolic import Expression diff --git a/loopy/schedule/verification.py b/loopy/schedule/verification.py index 4fa9c827c..9ca7da156 100644 --- a/loopy/schedule/verification.py +++ b/loopy/schedule/verification.py @@ -1,16 +1,16 @@ -from constantdict import constantdict -import islpy as isl -import namedisl as nisl -from namedisl import DimType +from __future__ import annotations -from collections.abc import Mapping, Sequence from dataclasses import dataclass +from typing import TYPE_CHECKING + +import namedisl as nisl +from constantdict import constantdict +from namedisl import DimType -from typing_extensions import override +import islpy as isl from loopy import KernelState, LoopKernel, for_each_kernel from loopy.diagnostic import LoopyError -from loopy.kernel.instruction import BarrierKind, HappensAfter from loopy.schedule import ( Barrier, CallKernel, @@ -22,6 +22,10 @@ ) +if TYPE_CHECKING: + from collections.abc import Mapping, Sequence + + @dataclass(frozen=True) class _PreciseScheduleRecord: timestamp: Sequence[int | str] @@ -162,10 +166,7 @@ def _build_statement_timestamp_relations( dom_str = ", ".join(name for name in stmt.within_inames) - if dom_str: - full_str = dom_str + ", " + ran_str - else: - full_str = ran_str + full_str = dom_str + ", " + ran_str if dom_str else ran_str # FIXME: isl -> named conversion domain = nisl.make_set( @@ -269,7 +270,7 @@ def _build_timestamp_relations( prec_sched: _PreciseSchedule, ) -> tuple[Mapping[str, nisl.Map], Sequence[nisl.Map], nisl.Map]: max_stmt_tstamp_len = max( - len(record.timestamp) for _, record in prec_sched.statements.items() + len(record.timestamp) for record in prec_sched.statements.values() ) max_bar_tstamp_len = -1 diff --git a/test/test_dependency.py b/test/test_dependency.py index 7c44224e1..8020bed03 100644 --- a/test/test_dependency.py +++ b/test/test_dependency.py @@ -7,17 +7,6 @@ import loopy.kernel.dependency as dep from loopy.diagnostic import LoopyError from loopy.kernel.instruction import HappensAfter - -from loopy.schedule.verification import ( - _BarrierRecord, - _PreciseSchedule, - _StatementRecord, - _build_enforced_order, - _build_strict_lexicographic_order, - _build_timestamp_relations, - _get_timestamp_points_from_linearization, - verify_happens_after_is_enforced, -) from loopy.schedule import ( Barrier, CallKernel, @@ -26,6 +15,16 @@ 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] From 4f72d9f9d57bef799b31c38b38fe3b230c804929 Mon Sep 17 00:00:00 2001 From: Addison Date: Mon, 20 Jul 2026 23:05:48 -0500 Subject: [PATCH 11/30] formatting + add verification to docs --- doc/ref_internals.rst | 1 + loopy/schedule/verification.py | 17 ++++++----------- 2 files changed, 7 insertions(+), 11 deletions(-) diff --git a/doc/ref_internals.rst b/doc/ref_internals.rst index 9e817a22f..7d3442bcd 100644 --- a/doc/ref_internals.rst +++ b/doc/ref_internals.rst @@ -64,3 +64,4 @@ Schedule Dependence Analysis ------------------- .. automodule:: loopy.kernel.dependency +.. automodule:: loopy.schedule.verification diff --git a/loopy/schedule/verification.py b/loopy/schedule/verification.py index 9ca7da156..5761f022f 100644 --- a/loopy/schedule/verification.py +++ b/loopy/schedule/verification.py @@ -174,19 +174,14 @@ def _build_statement_timestamp_relations( ).project_out_except([*stmt.within_inames, *kernel.all_params()]) if stmt.within_inames: - reln = nisl.make_set(f"{{[{full_str}]}}").as_map( - 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 - ) + for name, pos in zip(timestamp_names, timestamp, strict=True) ) reln = nisl.make_map( - f"{{ [] -> [{', '.join(timestamp_names)}] : " - f"{constraints} }}" + f"{{ [] -> [{', '.join(timestamp_names)}] : {constraints} }}" ) reln = reln.intersect_domain(domain) @@ -501,9 +496,9 @@ def verify_happens_after_is_enforced(kernel: LoopKernel) -> LoopKernel: ) required = nisl.make_map( - happens_after.instances_rel - .reset_tuple_id(isl.dim_type.in_) - .reset_tuple_id(isl.dim_type.out) + happens_after.instances_rel.reset_tuple_id( + isl.dim_type.in_ + ).reset_tuple_id(isl.dim_type.out) ) enforced = _build_enforced_order( kernel, From 0ec2850a57f5714e9dd503a6577d82b3013ba504 Mon Sep 17 00:00:00 2001 From: Addison Date: Mon, 20 Jul 2026 23:07:59 -0500 Subject: [PATCH 12/30] add file headers --- loopy/kernel/dependency.py | 25 ++++++++++++++ loopy/schedule/verification.py | 25 ++++++++++++++ test/test_dependency.py | 63 ++++++++++++++++++++++++++-------- 3 files changed, 99 insertions(+), 14 deletions(-) diff --git a/loopy/kernel/dependency.py b/loopy/kernel/dependency.py index 9b84911fd..926a24d6d 100644 --- a/loopy/kernel/dependency.py +++ b/loopy/kernel/dependency.py @@ -1,5 +1,30 @@ 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 diff --git a/loopy/schedule/verification.py b/loopy/schedule/verification.py index 5761f022f..edc67d9d3 100644 --- a/loopy/schedule/verification.py +++ b/loopy/schedule/verification.py @@ -1,5 +1,30 @@ 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 diff --git a/test/test_dependency.py b/test/test_dependency.py index 8020bed03..91a585bd9 100644 --- a/test/test_dependency.py +++ b/test/test_dependency.py @@ -1,3 +1,28 @@ +__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 pytest @@ -330,9 +355,7 @@ def test_has_precise_dependencies() -> None: ).default_entrypoint assert not dep.has_precise_dependencies(legacy_kernel) - precise_kernel = dep.add_lexicographic_happens_after( - 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"] @@ -746,7 +769,9 @@ def test_statement_timestamps_with_imperfect_loop_nesting() -> None: 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), + "outer": _StatementRecord( + timestamp=(0, 2, "i", 3), subkernel_idx=0 + ), "deep": _StatementRecord( timestamp=(0, 2, "i", 4, "j", 5, "k", 6), subkernel_idx=0, @@ -802,7 +827,8 @@ def test_timestamp_relation_keeps_parallel_inames_in_instance_domain() -> None: def test_strict_lexicographic_timestamp_order() -> None: order = _build_strict_lexicographic_order(("t0", "t1", "t2")) - assert order.equals(nisl.make_map(""" + assert order.equals( + nisl.make_map(""" { [t0_later, t1_later, t2_later] -> [t0_earlier, t1_earlier, t2_earlier] : @@ -817,7 +843,8 @@ def test_strict_lexicographic_timestamp_order() -> None: t1_later = t1_earlier and t2_later > t2_earlier } - """)) + """) + ) def test_statement_timestamps_record_local_and_global_barriers() -> None: @@ -895,7 +922,8 @@ def test_statement_timestamps_record_local_and_global_barriers() -> None: kernel, precise_schedule ) assert len(barrier_relations) == 2 - assert barrier_relations[0].equals(nisl.make_map(""" + assert barrier_relations[0].equals( + nisl.make_map(""" { [batch, i] -> [__ts_0, __ts_1, __ts_2, __ts_3, __ts_4, __ts_5] : @@ -904,8 +932,10 @@ def test_statement_timestamps_record_local_and_global_barriers() -> None: __ts_2 = 1 and __ts_3 = 2 and __ts_4 = i and __ts_5 = 4 } - """)) - assert barrier_relations[1].equals(nisl.make_map(""" + """) + ) + assert barrier_relations[1].equals( + nisl.make_map(""" { [batch] -> [__ts_0, __ts_1, __ts_2, __ts_3, __ts_4, __ts_5] : @@ -914,7 +944,8 @@ def test_statement_timestamps_record_local_and_global_barriers() -> None: __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: @@ -964,7 +995,8 @@ def test_local_barrier_orders_work_items_in_the_same_group() -> None: timestamp_order, ) - assert enforced.equals(nisl.make_map(""" + assert enforced.equals( + nisl.make_map(""" { [g_after, l_after, i_after] -> [g_before, l_before, i_before] : @@ -974,7 +1006,8 @@ def test_local_barrier_orders_work_items_in_the_same_group() -> None: 0 <= l_before < 4 and 0 <= i_before <= i_after } - """)) + """) + ) def test_global_barrier_orders_all_work_items() -> None: @@ -1019,13 +1052,15 @@ def test_global_barrier_orders_all_work_items() -> None: timestamp_order, ) - assert enforced.equals(nisl.make_map(""" + 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_verification() -> None: From b6393af2bdc44c60d0c4a5402382ac8120dc16e3 Mon Sep 17 00:00:00 2001 From: Addison Date: Mon, 20 Jul 2026 23:11:54 -0500 Subject: [PATCH 13/30] align namedisl usage with upstream namedisl api --- loopy/kernel/dependency.py | 4 ++-- loopy/schedule/verification.py | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/loopy/kernel/dependency.py b/loopy/kernel/dependency.py index 926a24d6d..64b5bba53 100644 --- a/loopy/kernel/dependency.py +++ b/loopy/kernel/dependency.py @@ -112,7 +112,7 @@ def _get_access_relation( cell_names = tuple(f"ax_{axis}" for axis in range(len(subscript))) access_set = domain.add_dims(DimType.out, cell_names) - coordinates = access_set.pw_affs + coordinates = access_set.var_pw_affs for cell_name, index_expr in zip(cell_names, subscript, strict=True): index_aff = nisl.make_aff( aff_from_expr( @@ -292,7 +292,7 @@ def add_lexicographic_happens_after(kernel: LoopKernel) -> LoopKernel: shared_order = after_order joint_domain = after_domain & before_domain - affs = joint_domain.pw_affs + affs = joint_domain.var_pw_affs strict_lex = joint_domain - joint_domain equal_prefix = joint_domain diff --git a/loopy/schedule/verification.py b/loopy/schedule/verification.py index edc67d9d3..fb644c685 100644 --- a/loopy/schedule/verification.py +++ b/loopy/schedule/verification.py @@ -195,7 +195,7 @@ def _build_statement_timestamp_relations( # FIXME: isl -> named conversion domain = nisl.make_set( - kernel.get_inames_domain(stmt.within_inames).as_set() + kernel.get_inames_domain(stmt.within_inames).to_set() ).project_out_except([*stmt.within_inames, *kernel.all_params()]) if stmt.within_inames: @@ -239,7 +239,7 @@ def _build_barrier_timestamp_relations( # FIXME: isl -> named conversion domain = nisl.make_set( - kernel.get_inames_domain(frozenset(inames)).as_set() + kernel.get_inames_domain(frozenset(inames)).to_set() ).project_out_except([*inames, *kernel.all_params()]) if inames: @@ -267,7 +267,7 @@ def _build_strict_lexicographic_order( joint = nisl.make_set( f"{{ [{', '.join([*later_names, *earlier_names])}] }}" ) - affs = joint.pw_affs + affs = joint.var_pw_affs strict_lex = joint - joint equal_prefix = joint From eec7919531bd60efc63a425e77d1bc67b594e360 Mon Sep 17 00:00:00 2001 From: Addison Date: Mon, 20 Jul 2026 23:13:02 -0500 Subject: [PATCH 14/30] fix ruff complaints (again) --- loopy/kernel/dependency.py | 1 + loopy/schedule/verification.py | 1 + 2 files changed, 2 insertions(+) diff --git a/loopy/kernel/dependency.py b/loopy/kernel/dependency.py index 64b5bba53..5fded4789 100644 --- a/loopy/kernel/dependency.py +++ b/loopy/kernel/dependency.py @@ -1,5 +1,6 @@ from __future__ import annotations + __copyright__ = """ Copyright (C) 2026 Addison Alvey-Blanco """ diff --git a/loopy/schedule/verification.py b/loopy/schedule/verification.py index fb644c685..b2502e646 100644 --- a/loopy/schedule/verification.py +++ b/loopy/schedule/verification.py @@ -1,5 +1,6 @@ from __future__ import annotations + __copyright__ = """ Copyright (C) 2026 Addison Alvey-Blanco """ From f4ac7c8af757a74e273213adced3cd2ff32541d2 Mon Sep 17 00:00:00 2001 From: Addison Date: Tue, 21 Jul 2026 08:20:01 -0500 Subject: [PATCH 15/30] minor: prefer statement over instruction --- loopy/kernel/dependency.py | 112 ++++++++++++++++++------------------- 1 file changed, 56 insertions(+), 56 deletions(-) diff --git a/loopy/kernel/dependency.py b/loopy/kernel/dependency.py index 5fded4789..deddb9ee1 100644 --- a/loopy/kernel/dependency.py +++ b/loopy/kernel/dependency.py @@ -72,8 +72,8 @@ class AccessType(Enum): def has_precise_dependencies(kernel: LoopKernel) -> bool: has_precise = False has_legacy = False - for insn in kernel.instructions: - for happens_after in insn.happens_after.values(): + for stmt in kernel.instructions: + for happens_after in stmt.happens_after.values(): if happens_after.instances_rel is None: has_legacy = True else: @@ -99,8 +99,8 @@ class AccessRelationFinder(WalkMapper[[str, AccessType]]): def __init__(self, kernel: LoopKernel): self.kernel = kernel self._additional_inames = frozenset() - self._read_relations = {insn.id: {} for insn in kernel.instructions} - self._write_relations = {insn.id: {} for insn in kernel.instructions} + self._read_relations = {stmt.id: {} for stmt in kernel.instructions} + self._write_relations = {stmt.id: {} for stmt in kernel.instructions} super().__init__() @@ -126,53 +126,53 @@ def _get_access_relation( return access_set.as_map(in_names=instance_names) - def _insn_writes_var(self, insn_id: str, var: str) -> bool: + def _stmt_writes_var(self, stmt_id: str, var: str) -> bool: return ( var in self.kernel.writer_map() - and insn_id in self.kernel.writer_map()[var] + and stmt_id in self.kernel.writer_map()[var] ) - def _insn_reads_var(self, insn_id: str, var: str) -> bool: + def _stmt_reads_var(self, stmt_id: str, var: str) -> bool: return ( var in self.kernel.reader_map() - and insn_id in self.kernel.reader_map()[var] + and stmt_id in self.kernel.reader_map()[var] ) - def _insn_accesses_var(self, insn_id: str, var: str) -> bool: - return self._insn_reads_var(insn_id, var) or self._insn_writes_var( - insn_id, var + def _stmt_accesses_var(self, stmt_id: str, var: str) -> bool: + return self._stmt_reads_var(stmt_id, var) or self._stmt_writes_var( + stmt_id, var ) def _record_access( self, - insn_id: str, + stmt_id: str, var: str, subscript: tuple[Expression, ...], access_type: AccessType, ) -> None: - if not self._insn_accesses_var(insn_id, var): + if not self._stmt_accesses_var(stmt_id, var): return - insn = self.kernel.id_to_insn[insn_id] - domain_inames = insn.within_inames | self._additional_inames + stmt = self.kernel.id_to_insn[stmt_id] + domain_inames = stmt.within_inames | self._additional_inames inames_domain = nisl.make_set( self.kernel.get_inames_domain(domain_inames).to_set() ) access_rel = self._get_access_relation(inames_domain, subscript) - additional_inames = self._additional_inames - insn.within_inames + 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[insn_id].get(var) - self._read_relations[insn_id][var] = ( + 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[insn_id].get(var) - self._write_relations[insn_id][var] = ( + 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 _: @@ -181,58 +181,58 @@ def _record_access( @cached_property def read_relations(self) -> Mapping[str, Mapping[str, nisl.Map]]: return constantdict({ - insn_id: constantdict(relations) - for insn_id, relations in self._read_relations.items() + 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({ - insn_id: constantdict(relations) - for insn_id, relations in self._write_relations.items() + stmt_id: constantdict(relations) + for stmt_id, relations in self._write_relations.items() }) @override def map_subscript( - self, expr: prim.Subscript, /, insn_id: str, access_type: AccessType + self, expr: prim.Subscript, /, stmt_id: str, access_type: AccessType ) -> None: assert isinstance(expr.aggregate, prim.Variable) self._record_access( - insn_id, expr.aggregate.name, expr.index_tuple, access_type + stmt_id, expr.aggregate.name, expr.index_tuple, access_type ) @override def map_linear_subscript( - self, expr: LinearSubscript, /, insn_id: str, access_type: AccessType + self, expr: LinearSubscript, /, stmt_id: str, access_type: AccessType ) -> None: - self.rec(expr.index, insn_id, AccessType.read) + self.rec(expr.index, stmt_id, AccessType.read) assert isinstance(expr.aggregate, prim.Variable) self._record_access( - insn_id, expr.aggregate.name, (expr.index,), access_type + stmt_id, expr.aggregate.name, (expr.index,), access_type ) @override def map_reduction( - self, expr: Reduction, /, insn_id: str, access_type: AccessType + 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, insn_id, access_type) + WalkMapper.map_reduction(self, expr, stmt_id, access_type) finally: self._additional_inames = previous_inames @override def map_sub_array_ref( - self, expr: SubArrayRef, /, insn_id: str, access_type: AccessType + 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, insn_id, access_type) + self.rec(expr.subscript, stmt_id, access_type) finally: self._additional_inames = previous_inames @@ -259,14 +259,14 @@ def add_lexicographic_happens_after(kernel: LoopKernel) -> LoopKernel: immediately preceding statement. """ - new_insns: list[InstructionBase] = [] - for i, insn in enumerate(kernel.instructions): + new_stmts: list[InstructionBase] = [] + for i, stmt in enumerate(kernel.instructions): new_happens_after: dict[str, HappensAfter] = {} - sources = (insn,) if i == 0 else (insn, kernel.instructions[i - 1]) + sources = (stmt,) if i == 0 else (stmt, kernel.instructions[i - 1]) after_domain = nisl.make_set( - kernel.get_inames_domain(insn.within_inames).to_set() + kernel.get_inames_domain(stmt.within_inames).to_set() ) after_inames = after_domain.space.dimtype_to_names[DimType.out] @@ -306,7 +306,7 @@ def add_lexicographic_happens_after(kernel: LoopKernel) -> LoopKernel: ) equal_prefix = equal_prefix & after_aff.eq_set(before_aff) - if source.id == insn.id: + if source.id == stmt.id: ordered_instances = strict_lex else: ordered_instances = strict_lex | equal_prefix @@ -319,9 +319,9 @@ def add_lexicographic_happens_after(kernel: LoopKernel) -> LoopKernel: instances_rel=instances_rel.as_isl() ) - new_insns.append(insn.copy(happens_after=new_happens_after)) + new_stmts.append(stmt.copy(happens_after=new_happens_after)) - return kernel.copy(instructions=new_insns) + return kernel.copy(instructions=new_stmts) def _relax_strict_happens_after_inner( @@ -435,8 +435,8 @@ def normalize_interface_and_compose( # Continue backward through the strict-order graph. if not live_access_rel.is_empty() and (sink_id != source_id): - source_insn = kernel.id_to_insn[source_id] - for src_dep_id, src_happens_after in source_insn.happens_after.items(): + source_stmt = kernel.id_to_insn[source_id] + for src_dep_id, src_happens_after in source_stmt.happens_after.items(): if src_dep_id == source_id: continue @@ -478,23 +478,23 @@ def relax_strict_happens_after(kernel: LoopKernel) -> LoopKernel: """ coarse_dependency_graph: dict[str, frozenset[str]] = {} - for insn in kernel.instructions: - coarse_dependency_graph[insn.id] = frozenset({ - dep for dep in insn.happens_after if dep != insn.id + for stmt in kernel.instructions: + coarse_dependency_graph[stmt.id] = frozenset({ + dep for dep in stmt.happens_after if dep != stmt.id }) topological_order = compute_topological_order(coarse_dependency_graph) rel_finder = AccessRelationFinder(kernel) - for insn in kernel.instructions: - if isinstance(insn, MultiAssignmentBase): - for assignee in insn.assignees: - rel_finder(assignee, insn.id, AccessType.write) - rel_finder(insn.expression, insn.id, AccessType.read) - for pred in insn.predicates: - rel_finder(pred, insn.id, AccessType.read) - - new_insns: list[InstructionBase] = [] + 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] = {} old_happens_after = kernel.id_to_insn[sink_id].happens_after @@ -525,10 +525,10 @@ def relax_strict_happens_after(kernel: LoopKernel) -> LoopKernel: new_happens_after, ) - new_insns.append( + new_stmts.append( kernel.id_to_insn[sink_id].copy( happens_after=constantdict(new_happens_after) ) ) - return kernel.copy(instructions=new_insns) + return kernel.copy(instructions=new_stmts) From 33b226a40356dcc004e0b76c908aba9a7ae18b21 Mon Sep 17 00:00:00 2001 From: Addison Date: Tue, 21 Jul 2026 08:42:18 -0500 Subject: [PATCH 16/30] ignore precise self edges for scheduling --- loopy/schedule/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/loopy/schedule/__init__.py b/loopy/schedule/__init__.py index ec5be4f19..b8ef1ce50 100644 --- a/loopy/schedule/__init__.py +++ b/loopy/schedule/__init__.py @@ -981,6 +981,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: From 229330354734010432bff175de0a0913c8f400f6 Mon Sep 17 00:00:00 2001 From: Addison Date: Tue, 21 Jul 2026 08:44:19 -0500 Subject: [PATCH 17/30] handle scalar accesses during access relation gathering --- loopy/kernel/dependency.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/loopy/kernel/dependency.py b/loopy/kernel/dependency.py index deddb9ee1..808d671df 100644 --- a/loopy/kernel/dependency.py +++ b/loopy/kernel/dependency.py @@ -192,6 +192,12 @@ def write_relations(self) -> Mapping[str, Mapping[str, nisl.Map]]: 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 From 18910b2387eec020e4a30457810d9259859ad897 Mon Sep 17 00:00:00 2001 From: Addison Date: Tue, 21 Jul 2026 09:03:46 -0500 Subject: [PATCH 18/30] improve dependency testing suite --- test/test_dependency.py | 602 ++++++++++++++-------------------------- 1 file changed, 211 insertions(+), 391 deletions(-) diff --git a/test/test_dependency.py b/test/test_dependency.py index 91a585bd9..db2ead6cc 100644 --- a/test/test_dependency.py +++ b/test/test_dependency.py @@ -24,8 +24,14 @@ import namedisl as nisl +import numpy as np import pytest +import pyopencl as cl +from pyopencl.tools import ( # ruff:ignore[unused-import] + pytest_generate_tests_for_pyopencl as pytest_generate_tests, +) + from pymbolic import var import loopy as lp @@ -121,110 +127,81 @@ def test_add_lexicographic_happens_after_uses_domain_dimension_order() -> None: ) -def test_add_lexicographic_happens_after_orders_distinct_loop_nests() -> None: +def test_add_lexicographic_happens_after_orders_mixed_loop_nests() -> None: t_unit = lp.make_kernel( [ - "{ [i] : 0 <= i < N }", - "{ [j] : 0 <= j < M }", + "[NI] -> { [i] : 0 <= i < NI }", + "[i, NJ] -> { [j] : 0 <= j < NJ }", + "[i, NK] -> { [k] : 0 <= k < NK }", + "[NQ] -> { [q] : 0 <= q < NQ }", ], """ - a[i] = 2 * a[i] {id=S} - b[j] = 2 * b[j] {id=T} + 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 - cross_relation = kernel.id_to_insn["T"].happens_after["S"].instances_rel + 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 cross_relation is not None + 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 + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + shared_nest_relation = nisl.make_map(shared_nest_relation) # FIXME: Remove conversion once HappensAfter stores namedisl.Map. - cross_relation = nisl.make_map(cross_relation) - assert cross_relation.equals( + disjoint_nest_relation = nisl.make_map(disjoint_nest_relation) + assert shared_nest_relation.equals( nisl.make_map(""" - [N, M] -> { - [j_after] -> [i_before] : - 0 <= j_after < M and - 0 <= i_before < N + [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 } """) ) - - -def test_add_lexicographic_happens_after_with_five_inames() -> None: - t_unit = lp.make_kernel( - """ - { [q, z, a, m, b] : - 0 <= q < 2 and - 0 <= z < 2 and - 0 <= a < 2 and - 0 <= m < 2 and - 0 <= b < 2 - } - """, - "out[q, z, a, m, b] = q + z + a + m + b {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 - # FIXME: Remove conversion once HappensAfter stores namedisl.Map. - self_relation = nisl.make_map(self_relation) - assert self_relation.equals( + assert disjoint_nest_relation.equals( nisl.make_map(""" - { - [q_after, z_after, a_after, m_after, b_after] -> - [q_before, z_before, a_before, m_before, b_before] : - 0 <= q_after, q_before < 2 and - 0 <= z_after, z_before < 2 and - 0 <= a_after, a_before < 2 and - 0 <= m_after, m_before < 2 and - 0 <= b_after, b_before < 2 and - (q_before < q_after or - (q_before = q_after and z_before < z_after) or - (q_before = q_after and z_before = z_after and - a_before < a_after) or - (q_before = q_after and z_before = z_after and - a_before = a_after and m_before < m_after) or - (q_before = q_after and z_before = z_after and - a_before = a_after and m_before = m_after and - b_before < b_after)) + [NI, NK, NQ] -> { + [q_after] -> [i_before, k_before] : + 0 <= q_after < NQ and + 0 <= i_before < NI and + 0 <= k_before < NK } """) ) -def test_access_relation_finder_keeps_instruction_maps_separate() -> None: +def test_access_relation_finder_tracks_reads_and_writes_per_statement() -> None: t_unit = lp.make_kernel( - "{ [i] : 0 <= i < N }", + "{ [i] : 1 <= i < N }", """ - a[i] = 1 {id=S} - b[i] = 2 {id=T} + 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) - rel_find(kernel.id_to_insn["S"].assignee, "S", dep.AccessType.write) - rel_find(kernel.id_to_insn["T"].assignee, "T", dep.AccessType.write) + 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"} - - -def test_access_relation_finder_distinguishes_reads_and_writes() -> None: - t_unit = lp.make_kernel( - "{ [i] : 1 <= i < N }", - "a[i] = a[i - 1] {id=S}", - ) - - kernel = t_unit.default_entrypoint - insn = kernel.id_to_insn["S"] - rel_find = dep.AccessRelationFinder(kernel) - rel_find(insn.assignee, insn.id, dep.AccessType.write) - rel_find(insn.expression, insn.id, dep.AccessType.read) - - assert rel_find.read_relations["S"]["a"].equals( + 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( @@ -232,27 +209,6 @@ def test_access_relation_finder_distinguishes_reads_and_writes() -> None: ) -def test_access_relation_finder_handles_linear_subscript() -> None: - t_unit = lp.make_kernel( - "{ [i, j] : 0 <= i < NI and 0 <= j < NJ }", - "out[i, j] = a[[2*i + j]] {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] -> { - [i, j] -> [ax_0 = 2*i + j] : - 0 <= i < NI and 0 <= j < NJ - } - """) - ) - - def test_access_relation_finder_handles_reduction() -> None: t_unit = lp.make_kernel( """ @@ -321,11 +277,22 @@ def _relax_strict_happens_after( return dep.relax_strict_happens_after(t_unit).default_entrypoint -def test_relax_strict_happens_after_finds_direct_raw() -> None: +@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( - """ - a[i, j] = 1 {id=S} - b[i, j] = a[i, j] {id=T} + f""" + {source} {{id=S}} + {sink} {{id=T}} """, "{ [i, j] : 0 <= i < NI and 0 <= j < NJ }", ) @@ -374,148 +341,66 @@ def test_has_precise_dependencies() -> None: dep.has_precise_dependencies(mixed_kernel) -def test_relax_strict_happens_after_finds_direct_waw() -> None: - kernel = _relax_strict_happens_after( - """ - a[i, j] = 1 {id=S} - a[i, j] = 2 {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 - # FIXME: Remove conversion once HappensAfter stores namedisl.Map. - required_order = nisl.make_map(required_order) - 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_relax_strict_happens_after_finds_direct_war() -> None: +def test_relax_strict_happens_after_tracks_scalar_accesses() -> None: kernel = _relax_strict_happens_after( """ - b[i, j] = a[i, j] {id=S} - a[i, j] = 2 {id=T} + <> tmp = i {id=S} + out[i] = tmp {id=T} """, - "{ [i, j] : 0 <= i < NI and 0 <= j < NJ }", + "{ [i] : 0 <= i < N }", ) required_order = kernel.id_to_insn["T"].happens_after["S"].instances_rel assert required_order is not None # FIXME: Remove conversion once HappensAfter stores namedisl.Map. required_order = nisl.make_map(required_order) - 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_relax_strict_happens_after_finds_self_raw() -> None: - kernel = _relax_strict_happens_after( - "a[i] = a[i - 1] {id=S}", - "{ [i] : 1 <= i < N }", - ) - - required_order = kernel.id_to_insn["S"].happens_after["S"].instances_rel - assert required_order is not None - # FIXME: Remove conversion once HappensAfter stores namedisl.Map. - required_order = nisl.make_map(required_order) assert required_order.equals( nisl.make_map(""" [N] -> { - [i_after] -> [i_before = i_after - 1] : 2 <= i_after < N + [i_after] -> [i_before] : + 0 <= i_before <= i_after < N } """) ) -def test_relax_strict_happens_after_drops_conflict_free_edge() -> None: - kernel = _relax_strict_happens_after(""" +@pytest.mark.parametrize( + "instructions", + ( + """ a[i] = 1 {id=S} b[i] = 2 {id=T} - """) - - assert "S" not in kernel.id_to_insn["T"].happens_after - - -def test_relax_strict_happens_after_finds_recursive_raw() -> None: - kernel = _relax_strict_happens_after( - """ - a[i, j] = 1 {id=S} - b[i, j] = 2 {id=T} - c[i, j] = a[i, j] {id=U} """, - "{ [i, j] : 0 <= i < NI and 0 <= j < NJ }", - ) - - required_order = kernel.id_to_insn["U"].happens_after["S"].instances_rel - assert required_order is not None - # FIXME: Remove conversion once HappensAfter stores namedisl.Map. - required_order = nisl.make_map(required_order) - 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 - } - """) - ) - assert "T" not in kernel.id_to_insn["U"].happens_after - - -def test_relax_strict_happens_after_stops_at_most_recent_writer() -> None: - kernel = _relax_strict_happens_after( """ - a[i, j] = 1 {id=S} - a[i, j] = 2 {id=T} - c[i, j] = a[i, j] {id=U} + a[i] = 1 {id=S} + b[i] = a[i + N] {id=T} """, - "{ [i, j] : 0 <= i < NI and 0 <= j < NJ }", - ) + ), + ids=("different-variables", "disjoint-footprints"), +) +def test_relax_strict_happens_after_drops_nonconflicting_edges( + instructions: str, +) -> None: + kernel = _relax_strict_happens_after(instructions) - required_order = kernel.id_to_insn["U"].happens_after["T"].instances_rel - assert required_order is not None - # FIXME: Remove conversion once HappensAfter stores namedisl.Map. - required_order = nisl.make_map(required_order) - 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 - } - """) - ) - assert "S" not in kernel.id_to_insn["U"].happens_after + assert "S" not in kernel.id_to_insn["T"].happens_after -def test_relax_strict_happens_after_partitions_partial_writer_footprint() -> ( +def test_relax_strict_happens_after_tracks_live_footprints_through_a_chain() -> ( None ): kernel = _relax_strict_happens_after( """ - a[i, j] = 1 {id=S} - a[2*i, j] = 2 {id=T} - c[i, j] = a[i, j] {id=U} + 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["U"].happens_after["T"].instances_rel - fallback_order = kernel.id_to_insn["U"].happens_after["S"].instances_rel + 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 # FIXME: Remove conversion once HappensAfter stores namedisl.Map. @@ -534,6 +419,7 @@ def test_relax_strict_happens_after_partitions_partial_writer_footprint() -> ( } """) ) + assert "B" not in kernel.id_to_insn["D"].happens_after assert fallback_order.equals( nisl.make_map(""" [NI, NJ] -> { @@ -547,21 +433,49 @@ def test_relax_strict_happens_after_partitions_partial_writer_footprint() -> ( ) -def test_relax_strict_happens_after_composes_distinct_loop_nests() -> None: +def test_relax_strict_happens_after_composes_user_supplied_relations() -> None: t_unit = lp.make_kernel( [ - "{ [i] : 0 <= i < 4 }", - "{ [j] : 0 <= j < 3 }", - "{ [k] : 0 <= k < 4 }", + "[N] -> { [i] : 0 <= i < 2*N }", + "[N] -> { [j] : 0 <= j < N }", + "[N] -> { [k] : 1 <= k < N }", ], """ a[i] = 1 {id=S} - b[j] = 2 {id=T} - c[k] = a[k] {id=U} + tmp[j] = 0 {id=T} + out[k] = a[2*k - 2] {id=U} """, ) - t_unit = dep.add_lexicographic_happens_after(t_unit) - kernel = dep.relax_strict_happens_after(t_unit).default_entrypoint + 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 + } + """).as_isl() + ) + u_after_t = HappensAfter( + instances_rel=nisl.make_map(""" + [N] -> { + [k_after] -> [j_before = k_after - 1] : + 1 <= k_after < N + } + """).as_isl() + ) + 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 @@ -569,8 +483,9 @@ def test_relax_strict_happens_after_composes_distinct_loop_nests() -> None: required_order = nisl.make_map(required_order) assert required_order.equals( nisl.make_map(""" - { - [k_after] -> [i_before = k_after] : 0 <= k_after < 4 + [N] -> { + [k_after] -> [i_before = 2*k_after - 2] : + 1 <= k_after < N } """) ) @@ -621,70 +536,6 @@ def test_relax_strict_happens_after_unions_branched_paths() -> None: ) -def test_relax_strict_happens_after_drops_empty_same_variable_edge() -> None: - kernel = _relax_strict_happens_after(""" - a[i] = 1 {id=S} - b[i] = a[i + N] {id=T} - """) - - assert "S" not in kernel.id_to_insn["T"].happens_after - - -def test_relax_strict_happens_after_inner_uses_live_sink_accesses() -> None: - t_unit = lp.make_kernel( - "{ [i] : 0 <= i < N }", - """ - a[i] = 1 {id=S} - b[i] = a[i] {id=T} - """, - ) - kernel = dep.add_lexicographic_happens_after(t_unit).default_entrypoint - - rel_finder = dep.AccessRelationFinder(kernel) - for insn in kernel.instructions: - assert isinstance(insn, lp.MultiAssignmentBase) - rel_finder(insn.assignee, insn.id, dep.AccessType.write) - rel_finder(insn.expression, insn.id, dep.AccessType.read) - - incoming_relation = kernel.id_to_insn["T"].happens_after["S"].instances_rel - assert incoming_relation is not None - # FIXME: Remove conversion once HappensAfter stores namedisl.Map. - incoming_relation = nisl.make_map(incoming_relation) - - live_access_relation = rel_finder.read_relations["T"]["a"].rename_dims(( - ("i", "i_after"), - )) - live_access_relation = live_access_relation & nisl.make_map(""" - [N] -> { - [i_after] -> [ax_0] : i_after = 0 - } - """) - - happens_after = dep._relax_strict_happens_after_inner( - kernel, - "T", - "S", - "a", - dep.AccessType.read, - incoming_relation, - live_access_relation, - rel_finder, - {}, - ) - - required_order = happens_after["S"].instances_rel - assert required_order is not None - # FIXME: Remove conversion once HappensAfter stores namedisl.Map. - required_order = nisl.make_map(required_order) - assert required_order.equals( - nisl.make_map(""" - [N] -> { - [i_after = 0] -> [i_before = 0] : N > 0 - } - """) - ) - - def test_statement_timestamps_with_calls_inside_outer_loop() -> None: kernel = lp.make_kernel( """{ @@ -738,7 +589,7 @@ def test_statement_timestamps_with_calls_inside_outer_loop() -> None: def test_statement_timestamps_with_imperfect_loop_nesting() -> None: - kernel = lp.make_kernel( + t_unit = lp.make_kernel( "{ [i, j, k] : 0 <= i, j, k < 8 }", """ <> before = 0 {id=before} @@ -747,24 +598,10 @@ def test_statement_timestamps_with_imperfect_loop_nesting() -> None: <> after_inner = i {id=after_inner} <> after = 0 {id=after} """, - ).default_entrypoint.copy( - state=lp.KernelState.LINEARIZED, - linearization=( - CallKernel("device_program"), - RunInstruction("before"), - EnterLoop("i"), - RunInstruction("outer"), - EnterLoop("j"), - EnterLoop("k"), - RunInstruction("deep"), - LeaveLoop("k"), - LeaveLoop("j"), - RunInstruction("after_inner"), - LeaveLoop("i"), - RunInstruction("after"), - ReturnFromKernel("device_program"), - ), ) + 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={ @@ -791,16 +628,8 @@ def test_timestamp_relation_keeps_parallel_inames_in_instance_domain() -> None: "out[g, l, i] = g + l + i {id=S}", ) t_unit = lp.tag_inames(t_unit, {"g": "g.0", "l": "l.0"}) - kernel = t_unit.default_entrypoint.copy( - state=lp.KernelState.LINEARIZED, - linearization=( - CallKernel("device_program"), - EnterLoop("i"), - RunInstruction("S"), - LeaveLoop("i"), - ReturnFromKernel("device_program"), - ), - ) + 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( @@ -949,12 +778,6 @@ def test_statement_timestamps_record_local_and_global_barriers() -> None: def test_local_barrier_orders_work_items_in_the_same_group() -> None: - local_barrier = Barrier( - comment="local synchronization", - synchronization_kind="local", - mem_kind="local", - originating_insn_id=None, - ) t_unit = lp.make_kernel( """ { @@ -963,23 +786,18 @@ def test_local_barrier_orders_work_items_in_the_same_group() -> None: } """, """ - a[g, l, i] = g + l + i {id=source} - b[g, l, i] = a[g, l, i] {id=sink} + <> tmp[g, l, i] = g + l + i {id=source} + out[g, l, i] = tmp[g, (l + 1) % 4, i] {id=sink} """, ) - t_unit = lp.tag_inames(t_unit, {"g": "g.0", "l": "l.0"}) - kernel = t_unit.default_entrypoint.copy( - state=lp.KernelState.LINEARIZED, - linearization=( - CallKernel("device_program"), - EnterLoop("i"), - RunInstruction("source"), - local_barrier, - RunInstruction("sink"), - LeaveLoop("i"), - ReturnFromKernel("device_program"), - ), + 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) stmt_relations, barrier_relations, timestamp_order = ( @@ -1011,32 +829,19 @@ def test_local_barrier_orders_work_items_in_the_same_group() -> None: def test_global_barrier_orders_all_work_items() -> None: - global_barrier = Barrier( - comment="global synchronization", - synchronization_kind="global", - mem_kind="global", - originating_insn_id=None, - ) t_unit = lp.make_kernel( "{ [g, l] : 0 <= g < 2 and 0 <= l < 4 }", """ a[g, l] = g + l {id=source} - b[g, l] = a[g, l] {id=sink} + 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"}) - kernel = t_unit.default_entrypoint.copy( - state=lp.KernelState.LINEARIZED, - linearization=( - CallKernel("phase0"), - RunInstruction("source"), - ReturnFromKernel("phase0"), - global_barrier, - CallKernel("phase1"), - RunInstruction("sink"), - ReturnFromKernel("phase1"), - ), - ) + 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) stmt_relations, barrier_relations, timestamp_order = ( @@ -1063,22 +868,60 @@ def test_global_barrier_orders_all_work_items() -> None: ) -def test_verification() -> None: - t_unit = lp.make_kernel( - "{ [i] : 0 <= i < N }", - """ - a[i] = 2 * a[i] {id=S} - b[i] = 2 * b[i] {id=T} - c[i] = a[i] + b[i] {id=U} - """, +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) - t_unit = lp.preprocess_program(t_unit) - t_unit = lp.linearize(t_unit) - t_unit = verify_happens_after_is_enforced(t_unit) - lp.generate_code_v2(t_unit) + + required_order = ( + t_unit.default_entrypoint + .id_to_insn["S"] + .happens_after["S"] + .instances_rel + ) + assert required_order is not None + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + assert nisl.make_map(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: @@ -1119,40 +962,17 @@ def test_verification_rejects_unenforced_order() -> None: def test_verification_handles_explicit_barrier_instruction() -> None: t_unit = lp.make_kernel( - "{ [i] : 0 <= i < N }", + "{ : }", """ - a[i] = i {id=S} + a[0] = 1 {id=S} ... gbarrier {id=B} - b[i] = a[i] {id=T} + b[0] = a[0] {id=T} """, + seq_dependencies=True, ) t_unit = dep.add_lexicographic_happens_after(t_unit) - - kernel = t_unit.default_entrypoint - barrier_insn = kernel.id_to_insn["B"] - barrier = Barrier( - comment="explicit global barrier", - synchronization_kind=barrier_insn.synchronization_kind, - mem_kind=barrier_insn.mem_kind, - originating_insn_id="B", - ) - kernel = kernel.copy( - state=lp.KernelState.LINEARIZED, - linearization=( - CallKernel("phase0"), - EnterLoop("i"), - RunInstruction("S"), - LeaveLoop("i"), - ReturnFromKernel("phase0"), - barrier, - CallKernel("phase1"), - EnterLoop("i"), - RunInstruction("T"), - LeaveLoop("i"), - ReturnFromKernel("phase1"), - ), - ) - t_unit = t_unit.with_kernel(kernel) + t_unit = lp.preprocess_program(t_unit) + t_unit = lp.linearize(t_unit) verify_happens_after_is_enforced(t_unit) From b46c17af6d888eba54dd023f3b0a1ddf2a8b82a0 Mon Sep 17 00:00:00 2001 From: Addison Date: Tue, 21 Jul 2026 09:11:02 -0500 Subject: [PATCH 19/30] use name generator to avoid conflicts in relation construction --- loopy/kernel/dependency.py | 11 ++++- loopy/schedule/verification.py | 85 +++++++++++++++++++++++++--------- test/test_dependency.py | 83 +++++++++++++++++++++++++++++++-- 3 files changed, 151 insertions(+), 28 deletions(-) diff --git a/loopy/kernel/dependency.py b/loopy/kernel/dependency.py index 808d671df..b720a7a37 100644 --- a/loopy/kernel/dependency.py +++ b/loopy/kernel/dependency.py @@ -59,6 +59,7 @@ from namedisl.core import NamedIslObjectT from pymbolic.typing import Expression + from pytools import UniqueNameGenerator from loopy.kernel import LoopKernel @@ -95,12 +96,16 @@ class AccessRelationFinder(WalkMapper[[str, AccessType]]): _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] 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 = [] super().__init__() @@ -110,7 +115,11 @@ def _get_access_relation( subscript: tuple[Expression, ...], ) -> nisl.Map: instance_names = domain.space.dimtype_to_names[DimType.out] - cell_names = tuple(f"ax_{axis}" for axis in range(len(subscript))) + 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 diff --git a/loopy/schedule/verification.py b/loopy/schedule/verification.py index b2502e646..c027334c3 100644 --- a/loopy/schedule/verification.py +++ b/loopy/schedule/verification.py @@ -51,6 +51,8 @@ if TYPE_CHECKING: from collections.abc import Mapping, Sequence + from pytools import UniqueNameGenerator + @dataclass(frozen=True) class _PreciseScheduleRecord: @@ -260,11 +262,9 @@ def _build_barrier_timestamp_relations( def _build_strict_lexicographic_order( - timestamp_names: Sequence[str], + later_names: Sequence[str], + earlier_names: Sequence[str], ) -> nisl.Map: - later_names = tuple(f"{name}_later" for name in timestamp_names) - earlier_names = tuple(f"{name}_earlier" for name in timestamp_names) - joint = nisl.make_set( f"{{ [{', '.join([*later_names, *earlier_names])}] }}" ) @@ -289,6 +289,7 @@ def _build_strict_lexicographic_order( 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() @@ -301,7 +302,15 @@ def _build_timestamp_relations( ) max_tstamp_len = max(max_stmt_tstamp_len, max_bar_tstamp_len) - timestamp_names = [f"__ts_{i}" for i in range(max_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 @@ -311,7 +320,9 @@ def _build_timestamp_relations( kernel, prec_sched.barriers, timestamp_names ) - timestamp_lex = _build_strict_lexicographic_order(timestamp_names) + timestamp_lex = _build_strict_lexicographic_order( + later_names, earlier_names + ) return constantdict(stmt_relns), bar_relns, timestamp_lex @@ -329,9 +340,15 @@ def _suffix_dim_names( def _timestamp_relation_for_role( relation: nisl.Map, - role: str, + role_names: Sequence[str], ) -> nisl.Map: - return _suffix_dim_names(relation, DimType.out, f"_{role}") + return relation.rename_dims( + zip( + relation.space.dimtype_to_names[DimType.out], + role_names, + strict=True, + ) + ) def _hardware_axis_inames( @@ -373,14 +390,13 @@ def _build_hardware_id_relation( 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] - hardware_names = tuple( - f"__{kind}_{axis}" for kind, axis in sorted(axis_inames) - ) constraints = " and ".join( - f"__{kind}_{axis} = {axis_inames[kind, axis]}{instance_suffix}" + 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 "" @@ -388,7 +404,7 @@ def _build_hardware_id_relation( relation = nisl.make_map( "{ " f"[{', '.join(input_names)}] -> " - f"[{', '.join(hardware_names)}]" + f"[{', '.join(hardware_names[key] for key in sorted(axis_inames))}]" f"{constraint_str} " "}" ) @@ -402,6 +418,7 @@ def _build_same_hardware_scope_relation( 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) @@ -412,11 +429,26 @@ def _build_same_hardware_scope_relation( 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 + 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 + kernel, + source_id, + source_domain, + "_before", + include_local_axes, + hardware_names, ) return sink_hardware.apply_range(source_hardware.reverse()) @@ -429,14 +461,19 @@ def _build_enforced_order( 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") + _timestamp_relation_for_role(sink, later_names) .apply_range(timestamp_lex) - .apply_range(_timestamp_relation_for_role(source, "earlier").reverse()) + .apply_range( + _timestamp_relation_for_role(source, earlier_names).reverse() + ) ) enforced = enforced & _build_same_hardware_scope_relation( kernel, @@ -445,6 +482,7 @@ def _build_enforced_order( sink.domain(), source.domain(), include_local_axes=True, + name_generator=name_generator, ) sink_record = prec_sched.statements[sink_id] @@ -462,17 +500,17 @@ def _build_enforced_order( barrier_reln, DimType.in_, f"_barrier_{barrier_idx}" ) sink_to_barrier = ( - _timestamp_relation_for_role(sink, "later") + _timestamp_relation_for_role(sink, later_names) .apply_range(timestamp_lex) .apply_range( - _timestamp_relation_for_role(barrier, "earlier").reverse() + _timestamp_relation_for_role(barrier, earlier_names).reverse() ) ) barrier_to_source = ( - _timestamp_relation_for_role(barrier, "later") + _timestamp_relation_for_role(barrier, later_names) .apply_range(timestamp_lex) .apply_range( - _timestamp_relation_for_role(source, "earlier").reverse() + _timestamp_relation_for_role(source, earlier_names).reverse() ) ) through_barrier = sink_to_barrier.apply_range(barrier_to_source) @@ -487,6 +525,7 @@ def _build_enforced_order( sink.domain(), source.domain(), include_local_axes=False, + name_generator=name_generator, ) ) @@ -508,8 +547,9 @@ def verify_happens_after_is_enforced(kernel: LoopKernel) -> LoopKernel: 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 + kernel, prec_sched, name_generator ) for sink in kernel.instructions: @@ -534,6 +574,7 @@ def verify_happens_after_is_enforced(kernel: LoopKernel) -> LoopKernel: stmt_relns, barrier_relns, timestamp_lex, + name_generator, ) missing = required - enforced if not missing.is_empty(): diff --git a/test/test_dependency.py b/test/test_dependency.py index db2ead6cc..5b91a7770 100644 --- a/test/test_dependency.py +++ b/test/test_dependency.py @@ -209,6 +209,31 @@ def test_access_relation_finder_tracks_reads_and_writes_per_statement() -> None: ) +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( """ @@ -638,7 +663,7 @@ def test_timestamp_relation_keeps_parallel_inames_in_instance_domain() -> None: ) timestamp_relations, _, _ = _build_timestamp_relations( - kernel, precise_schedule + kernel, precise_schedule, kernel.get_var_name_generator() ) timestamp_relation = timestamp_relations["S"] assert timestamp_relation.equals( @@ -653,8 +678,48 @@ def test_timestamp_relation_keeps_parallel_inames_in_instance_domain() -> None: ) +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", "t1", "t2")) + order = _build_strict_lexicographic_order( + ("t0_later", "t1_later", "t2_later"), + ("t0_earlier", "t1_earlier", "t2_earlier"), + ) assert order.equals( nisl.make_map(""" @@ -748,7 +813,7 @@ def test_statement_timestamps_record_local_and_global_barriers() -> None: ) _, barrier_relations, _ = _build_timestamp_relations( - kernel, precise_schedule + kernel, precise_schedule, kernel.get_var_name_generator() ) assert len(barrier_relations) == 2 assert barrier_relations[0].equals( @@ -800,8 +865,11 @@ def test_local_barrier_orders_work_items_in_the_same_group() -> None: 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) + _build_timestamp_relations( + kernel, precise_schedule, name_generator + ) ) enforced = _build_enforced_order( kernel, @@ -811,6 +879,7 @@ def test_local_barrier_orders_work_items_in_the_same_group() -> None: stmt_relations, barrier_relations, timestamp_order, + name_generator, ) assert enforced.equals( @@ -844,8 +913,11 @@ def test_global_barrier_orders_all_work_items() -> None: 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) + _build_timestamp_relations( + kernel, precise_schedule, name_generator + ) ) enforced = _build_enforced_order( kernel, @@ -855,6 +927,7 @@ def test_global_barrier_orders_all_work_items() -> None: stmt_relations, barrier_relations, timestamp_order, + name_generator, ) assert enforced.equals( From d535428ca543d7486cc6e1c81e2215c4ad149138 Mon Sep 17 00:00:00 2001 From: Addison Date: Tue, 21 Jul 2026 09:11:24 -0500 Subject: [PATCH 20/30] fix ruff complaints --- test/test_dependency.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/test/test_dependency.py b/test/test_dependency.py index 5b91a7770..6988f0c1e 100644 --- a/test/test_dependency.py +++ b/test/test_dependency.py @@ -28,12 +28,11 @@ 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, ) -from pymbolic import var - import loopy as lp import loopy.kernel.dependency as dep from loopy.diagnostic import LoopyError From cf0200d43840aea11ebaeb5a1360825833fc32f6 Mon Sep 17 00:00:00 2001 From: Addison Date: Tue, 21 Jul 2026 09:19:01 -0500 Subject: [PATCH 21/30] update baseline to include InstructionBase.copy() usage in dependency.py --- .basedpyright/baseline.json | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index 29853be78..f5763171a 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -26144,6 +26144,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", @@ -99951,4 +99969,4 @@ } ] } -} \ No newline at end of file +} From e4169c809915589164c6820b1c5c5a48099b869f Mon Sep 17 00:00:00 2001 From: Addison Date: Tue, 21 Jul 2026 09:35:43 -0500 Subject: [PATCH 22/30] expand substitution rules when finding access relations --- loopy/kernel/dependency.py | 30 ++++++++++++------------------ test/test_dependency.py | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 44 insertions(+), 18 deletions(-) diff --git a/loopy/kernel/dependency.py b/loopy/kernel/dependency.py index b720a7a37..7e250625e 100644 --- a/loopy/kernel/dependency.py +++ b/loopy/kernel/dependency.py @@ -48,6 +48,7 @@ LinearSubscript, Reduction, SubArrayRef, + SubstitutionRuleExpander, WalkMapper, aff_from_expr, ) @@ -98,6 +99,8 @@ class AccessRelationFinder(WalkMapper[[str, AccessType]]): _write_relations: dict[str, dict[str, nisl.Map]] _name_generator: UniqueNameGenerator _cell_names: list[str] + _storage_variables: frozenset[str] + _subst_expander: SubstitutionRuleExpander def __init__(self, kernel: LoopKernel): self.kernel = kernel @@ -106,9 +109,17 @@ def __init__(self, kernel: LoopKernel): self._write_relations = {stmt.id: {} for stmt in kernel.instructions} self._name_generator = kernel.get_var_name_generator() self._cell_names = [] + 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, @@ -135,23 +146,6 @@ def _get_access_relation( return access_set.as_map(in_names=instance_names) - def _stmt_writes_var(self, stmt_id: str, var: str) -> bool: - return ( - var in self.kernel.writer_map() - and stmt_id in self.kernel.writer_map()[var] - ) - - def _stmt_reads_var(self, stmt_id: str, var: str) -> bool: - return ( - var in self.kernel.reader_map() - and stmt_id in self.kernel.reader_map()[var] - ) - - def _stmt_accesses_var(self, stmt_id: str, var: str) -> bool: - return self._stmt_reads_var(stmt_id, var) or self._stmt_writes_var( - stmt_id, var - ) - def _record_access( self, stmt_id: str, @@ -159,7 +153,7 @@ def _record_access( subscript: tuple[Expression, ...], access_type: AccessType, ) -> None: - if not self._stmt_accesses_var(stmt_id, var): + if var not in self._storage_variables: return stmt = self.kernel.id_to_insn[stmt_id] diff --git a/test/test_dependency.py b/test/test_dependency.py index 6988f0c1e..f1423b14f 100644 --- a/test/test_dependency.py +++ b/test/test_dependency.py @@ -301,6 +301,38 @@ def _relax_strict_happens_after( 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 + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + required_order = nisl.make_map(required_order) + assert required_order.equals( + nisl.make_map(""" + [N] -> { + [i_after] -> [i_before = i_after] : 0 <= i_after < N + } + """) + ) + + @pytest.mark.parametrize( ("source", "sink"), ( From 483d3bfc3c25d82be41503b21f14a4b1fc932c4e Mon Sep 17 00:00:00 2001 From: Addison Date: Tue, 21 Jul 2026 13:19:09 -0500 Subject: [PATCH 23/30] minor formatting --- loopy/kernel/dependency.py | 40 +++++++++++++++++++------------------- 1 file changed, 20 insertions(+), 20 deletions(-) diff --git a/loopy/kernel/dependency.py b/loopy/kernel/dependency.py index 7e250625e..e42db9d1e 100644 --- a/loopy/kernel/dependency.py +++ b/loopy/kernel/dependency.py @@ -71,25 +71,6 @@ class AccessType(Enum): write = 1 -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 - - class AccessRelationFinder(WalkMapper[[str, AccessType]]): """Collect per-instruction statement-instance-to-cell access relations.""" @@ -130,7 +111,7 @@ def _get_access_relation( axis = len(self._cell_names) self._cell_names.append(self._name_generator(f"ax_{axis}")) - cell_names = tuple(self._cell_names[:len(subscript)]) + cell_names = tuple(self._cell_names[: len(subscript)]) access_set = domain.add_dims(DimType.out, cell_names) coordinates = access_set.var_pw_affs @@ -246,6 +227,25 @@ def map_sub_array_ref( self._additional_inames = previous_inames +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: From f9d95df5c055122c15e793f1d95c9cec82dcc7a9 Mon Sep 17 00:00:00 2001 From: Addison Date: Tue, 21 Jul 2026 16:04:39 -0500 Subject: [PATCH 24/30] add primitive for transforming happens afters when inames are transformed --- loopy/kernel/dependency.py | 104 +++++++++++++ loopy/transform/iname.py | 11 ++ test/test_dependency.py | 294 ++++++++++++++++++++++++++++++++++++- 3 files changed, 402 insertions(+), 7 deletions(-) diff --git a/loopy/kernel/dependency.py b/loopy/kernel/dependency.py index e42db9d1e..536d65511 100644 --- a/loopy/kernel/dependency.py +++ b/loopy/kernel/dependency.py @@ -227,6 +227,110 @@ def map_sub_array_ref( self._additional_inames = previous_inames +def apply_affine_transform_to_happens_afters( + kernel: LoopKernel, affine_reln: nisl.Map +) -> LoopKernel: + """ + Applies an affine transformation to all relevant happens-after relations. + """ + + 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" + ) + + # FIXME: Remove conversion once LoopKernel domains use namedisl.Set. + stmt_domain = nisl.make_set( + kernel.get_inames_domain(stmt.within_inames).to_set() + ) + 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, 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 + + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + instances_rel = nisl.make_map(happens_after.instances_rel) + 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 + ) + 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) + proxy_renames.extend(src_proxy_renames) + + instances_rel = instances_rel.rename_dims(proxy_renames).coalesce() + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + new_happens_after[src_id] = HappensAfter(instances_rel.as_isl()) + + 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 diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index 891066fa9..8bb8408a9 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -2290,6 +2290,17 @@ def process_set(s: isl.BasicSet) -> isl.BasicSet: # }}} + from loopy.kernel.dependency import ( + apply_affine_transform_to_happens_afters, + has_precise_dependencies, + ) + if has_precise_dependencies(kernel): + import namedisl as nisl + # FIXME: Remove conversion once map_domain accepts namedisl.Map. + kernel = apply_affine_transform_to_happens_afters( + kernel, nisl.make_map(transform_map.to_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/test/test_dependency.py b/test/test_dependency.py index f1423b14f..b9cc29f63 100644 --- a/test/test_dependency.py +++ b/test/test_dependency.py @@ -27,6 +27,7 @@ import numpy as np import pytest +import islpy as isl import pyopencl as cl from pymbolic import var from pyopencl.tools import ( # ruff:ignore[unused-import] @@ -179,6 +180,289 @@ def test_add_lexicographic_happens_after_orders_mixed_loop_nests() -> None: ) +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 + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + return nisl.make_map(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 = isl.BasicMap(""" + [NI] -> { + [i] -> [io, ii] : + i = 4*io + ii and 0 <= ii < 4 + } + """) + + expected = dep.apply_affine_transform_to_happens_afters( + t_unit.default_entrypoint, + nisl.make_map(transform_map.to_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_access_relation_finder_tracks_reads_and_writes_per_statement() -> None: t_unit = lp.make_kernel( "{ [i] : 1 <= i < N }", @@ -310,7 +594,7 @@ def test_relax_strict_happens_after_expands_substitution_rules() -> None: """, substitutions={ "inner": lp.SubstitutionRule( - "inner", ("j",), var("a")[2*var("j") + 1] + "inner", ("j",), var("a")[2 * var("j") + 1] ), "outer": lp.SubstitutionRule( "outer", ("k",), var("inner")(var("k") + 1) @@ -898,9 +1182,7 @@ def test_local_barrier_orders_work_items_in_the_same_group() -> None: 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 - ) + _build_timestamp_relations(kernel, precise_schedule, name_generator) ) enforced = _build_enforced_order( kernel, @@ -946,9 +1228,7 @@ def test_global_barrier_orders_all_work_items() -> None: 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 - ) + _build_timestamp_relations(kernel, precise_schedule, name_generator) ) enforced = _build_enforced_order( kernel, From 7c480384a185ab2405335cd76fdfb283df599336 Mon Sep 17 00:00:00 2001 From: Addison Date: Tue, 21 Jul 2026 17:05:14 -0500 Subject: [PATCH 25/30] add splicing primitives for consumer/producer insertion; hook split_iname into precise deps --- loopy/kernel/dependency.py | 327 +++++++++++++++++- loopy/transform/iname.py | 27 ++ test/test_dependency.py | 690 +++++++++++++++++++++++++++++++++++++ 3 files changed, 1040 insertions(+), 4 deletions(-) diff --git a/loopy/kernel/dependency.py b/loopy/kernel/dependency.py index 536d65511..3d0c9f4ea 100644 --- a/loopy/kernel/dependency.py +++ b/loopy/kernel/dependency.py @@ -234,6 +234,7 @@ def apply_affine_transform_to_happens_afters( 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(): @@ -277,7 +278,7 @@ def build_xform_reln( nonxformed_names, output_proxy_names, strict=True ) ) - return xform_reln, proxy_renames + return xform_reln.coalesce(), proxy_renames new_stmts: list[InstructionBase] = [] for sink_stmt in kernel.instructions: @@ -297,14 +298,16 @@ def build_xform_reln( continue # FIXME: Remove conversion once HappensAfter stores namedisl.Map. - instances_rel = nisl.make_map(happens_after.instances_rel) + instances_rel = nisl.make_map( + 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: @@ -317,7 +320,9 @@ def build_xform_reln( ) src_xform_reln, src_proxy_renames = src_xform - instances_rel = instances_rel.apply_range(src_xform_reln) + 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() @@ -358,6 +363,320 @@ def _suffix_names( ) +def _statement_instance_set( + kernel: LoopKernel, stmt: InstructionBase, suffix: str +) -> nisl.Set: + # FIXME: Remove conversion once LoopKernel domains use namedisl.Set. + instance_set = nisl.make_set( + kernel.get_inames_domain(stmt.within_inames).to_set() + ) + 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 _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, +) -> 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}'" + ) + + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + previous_rel = nisl.make_map(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() + + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + happens_after[source_id] = HappensAfter(instances_rel.as_isl()) + + +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}'" + ) + + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + inherited = _compose_happens_after_relations( + consumer_to_anchor, + nisl.make_map(happens_after.instances_rel), + ) + _add_or_union_happens_after( + new_happens_after, consumer_id, source_id, inherited + ) + + 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}'" + ) + + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + anchor_order = nisl.make_map(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: + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + new_happens_after[anchor_id] = HappensAfter( + remaining_anchor_order.as_isl() + ) + + _add_or_union_happens_after( + new_happens_after, sink.id, producer_id, redirected_order + ) + + 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: """ diff --git a/loopy/transform/iname.py b/loopy/transform/iname.py index 8bb8408a9..8206d83a5 100644 --- a/loopy/transform/iname.py +++ b/loopy/transform/iname.py @@ -305,6 +305,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) @@ -339,6 +341,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) diff --git a/test/test_dependency.py b/test/test_dependency.py index b9cc29f63..c08192a28 100644 --- a/test/test_dependency.py +++ b/test/test_dependency.py @@ -463,6 +463,624 @@ def test_map_domain_transforms_precise_happens_after() -> None: ) +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 + } + """).as_isl()) + 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 + } + """).as_isl()) + 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 + } + """).as_isl()) + + 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 + } + """).as_isl()) + c_after_p = HappensAfter(nisl.make_map(""" + [N] -> { + [ic_after, lane_after] -> [i_before = ic_after] : + 0 <= ic_after < N and lane_after = 0 + } + """).as_isl()) + 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("{ [] -> [] }").as_isl()), + }, + "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 + } + """).as_isl()) + 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 + } + """).as_isl()) + 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 + } + """).as_isl()) + 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 + } + """).as_isl()) + 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 + } + """).as_isl()) + 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 + } + """).as_isl()) + s_after_g = HappensAfter(nisl.make_map(""" + [N] -> { + [i_after] -> [ip_before = i_after] : + 0 <= i_after < N and 2*i_after >= N + } + """).as_isl()) + 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 + } + """).as_isl()), + }, + }[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.as_isl())}, + "G": {}, + "S": {"A": HappensAfter(same_instance.as_isl())}, + }[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 + } + """).as_isl()) + 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 }", @@ -1252,6 +1870,78 @@ def test_global_barrier_orders_all_work_items() -> None: ) +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 + } + """).as_isl()) + 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: From f896110811ae39660f043c4b807d41b479df3cc4 Mon Sep 17 00:00:00 2001 From: Addison Date: Fri, 24 Jul 2026 21:24:43 -0500 Subject: [PATCH 26/30] bugfixes causing overapproximation of required ordering in some cases --- loopy/kernel/dependency.py | 84 ++++++++++++++++++ loopy/schedule/verification.py | 5 ++ test/test_dependency.py | 153 +++++++++++++++++++++++++++++++++ 3 files changed, 242 insertions(+) diff --git a/loopy/kernel/dependency.py b/loopy/kernel/dependency.py index 3d0c9f4ea..3a10d99c3 100644 --- a/loopy/kernel/dependency.py +++ b/loopy/kernel/dependency.py @@ -40,6 +40,7 @@ from loopy import for_each_kernel from loopy.diagnostic import LoopyError from loopy.kernel.instruction import ( + CInstruction, HappensAfter, InstructionBase, MultiAssignmentBase, @@ -405,6 +406,73 @@ def _compose_happens_after_relations( 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" + ) + + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + self_relation = nisl.make_map(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" + ) + + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + relation = nisl.make_map(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() + + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + new_happens_after[source_id] = HappensAfter(relation.as_isl()) + + 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, @@ -909,6 +977,22 @@ def relax_strict_happens_after(kernel: LoopKernel) -> LoopKernel: 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: coarse_dependency_graph[stmt.id] = frozenset({ diff --git a/loopy/schedule/verification.py b/loopy/schedule/verification.py index c027334c3..9bdd272b3 100644 --- a/loopy/schedule/verification.py +++ b/loopy/schedule/verification.py @@ -84,6 +84,8 @@ def _get_timestamp_points_from_linearization( "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]: @@ -91,6 +93,9 @@ def build_timestamp_from_stack( 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) diff --git a/test/test_dependency.py b/test/test_dependency.py index c08192a28..bb588b268 100644 --- a/test/test_dependency.py +++ b/test/test_dependency.py @@ -1299,6 +1299,54 @@ def test_has_precise_dependencies() -> None: 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( """ @@ -1449,6 +1497,71 @@ def test_relax_strict_happens_after_composes_user_supplied_relations() -> None: ) +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 + } + """).as_isl()) + c_after_p = HappensAfter(nisl.make_map(""" + { + [i_after] -> [i_before = i_after] : + 0 <= i_after < 8 + } + """).as_isl()) + 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 + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + required_order = nisl.make_map(required_order) + 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 + } + """).as_isl()) + 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 }", @@ -2034,6 +2147,46 @@ def test_verification_rejects_unenforced_order() -> None: 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 + } + """).as_isl() + ) + 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( "{ : }", From 8a6a2f496865f5ea7c7b46b6ffeb432e96ae2fd2 Mon Sep 17 00:00:00 2001 From: Addison Date: Fri, 24 Jul 2026 22:43:14 -0500 Subject: [PATCH 27/30] bugfix: actually find most-recent-writer in scalar access case --- loopy/kernel/dependency.py | 66 +++++++++++++++++++++++++++++++------- test/test_dependency.py | 36 +++++++++++++++++++-- 2 files changed, 89 insertions(+), 13 deletions(-) diff --git a/loopy/kernel/dependency.py b/loopy/kernel/dependency.py index 3a10d99c3..b9cb65985 100644 --- a/loopy/kernel/dependency.py +++ b/loopy/kernel/dependency.py @@ -869,14 +869,48 @@ def _relax_strict_happens_after_inner( :returns: The updated precise dependencies for *sink_id*. """ - def record_conflicts(source_relation: nisl.Map) -> nisl.Map: + def record_conflicts( + source_relation: nisl.Map, + *, + select_most_recent_writer: bool, + ) -> nisl.Map: source_relation = _suffix_names(source_relation, "_before", DimType.in_) - # live_access_rel; source_relation^-1 - conflicts = live_access_rel.apply_range(source_relation.reverse()) + sink_names = incoming_instances_rel.space.in_names + cell_names = live_access_rel.space.out_names + + candidate_set = ( + live_access_rel.as_set() + & incoming_instances_rel.as_set() + & source_relation.as_set() + ) + candidates = candidate_set.as_map( + in_names=(*sink_names, *cell_names) + ) - # Only conflicts ordered along this graph path are required. - required_order = incoming_instances_rel & conflicts + if select_most_recent_writer: + self_happens_after = kernel.id_to_insn[ + source_id + ].happens_after.get(source_id) + if self_happens_after is not None: + assert self_happens_after.instances_rel is not None + # FIXME: Remove conversion once HappensAfter stores + # namedisl.Map. + self_relation = nisl.make_map( + self_happens_after.instances_rel + ) + dominated = candidates & _compose_happens_after_relations( + candidates, self_relation + ) + candidates = (candidates - dominated).coalesce() + + required_order = ( + candidates + .as_set() + .project_out(cell_names) + .as_map(in_names=sink_names) + .coalesce() + ) previous = happens_after.get(source_id) if not required_order.is_empty(): if previous is None: @@ -885,13 +919,14 @@ def record_conflicts(source_relation: nisl.Map) -> nisl.Map: assert previous.instances_rel is not None # FIXME: remove named conversion previous_instances_rel = nisl.make_map(previous.instances_rel) - combined_order = required_order | previous_instances_rel + combined_order = ( + required_order | previous_instances_rel + ).coalesce() # FIXME: remove unnamed conversion happens_after[source_id] = HappensAfter(combined_order.as_isl()) - # Retire only live accesses supplied by an ordered source instance. - return live_access_rel & required_order.apply_range(source_relation) + return candidates.domain().as_map(in_names=sink_names).coalesce() def normalize_interface_and_compose( incoming_relation: nisl.Map, next_edge_relation: nisl.Map @@ -914,7 +949,10 @@ def normalize_interface_and_compose( if var in rel_finder.write_relations[source_id]: source_relation = rel_finder.write_relations[source_id][var] - caught_accesses = record_conflicts(source_relation) + caught_accesses = record_conflicts( + source_relation, + select_most_recent_writer=True, + ) live_access_rel = live_access_rel - caught_accesses # Write-after-write and write-after-read @@ -922,13 +960,19 @@ def normalize_interface_and_compose( if var in rel_finder.write_relations[source_id]: source_relation = rel_finder.write_relations[source_id][var] - caught_accesses = record_conflicts(source_relation) + caught_accesses = record_conflicts( + source_relation, + select_most_recent_writer=True, + ) live_access_rel = live_access_rel - caught_accesses # don't update live_access_rel; does not find a "most recent writer" if var in rel_finder.read_relations[source_id]: source_relation = rel_finder.read_relations[source_id][var] - _ = record_conflicts(source_relation) + _ = record_conflicts( + source_relation, + select_most_recent_writer=False, + ) case _: raise ValueError("unknown access type") diff --git a/test/test_dependency.py b/test/test_dependency.py index bb588b268..d42befd18 100644 --- a/test/test_dependency.py +++ b/test/test_dependency.py @@ -1363,8 +1363,40 @@ def test_relax_strict_happens_after_tracks_scalar_accesses() -> None: assert required_order.equals( nisl.make_map(""" [N] -> { - [i_after] -> [i_before] : - 0 <= i_before <= i_after < N + [i_after] -> [i_before = i_after] : + 0 <= i_after < N + } + """) + ) + + +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 + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + required_order = nisl.make_map(required_order) + 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 } """) ) From c48e84158a876d6c164773193e4d37a461fe7d36 Mon Sep 17 00:00:00 2001 From: Addison Date: Sun, 26 Jul 2026 22:36:43 -0500 Subject: [PATCH 28/30] more underapproximation bugfixes --- loopy/kernel/dependency.py | 65 ++++++++++++++++++++++--------- test/test_dependency.py | 80 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 127 insertions(+), 18 deletions(-) diff --git a/loopy/kernel/dependency.py b/loopy/kernel/dependency.py index b9cb65985..6d1c35969 100644 --- a/loopy/kernel/dependency.py +++ b/loopy/kernel/dependency.py @@ -56,7 +56,7 @@ if TYPE_CHECKING: - from collections.abc import Mapping + from collections.abc import Collection, Mapping from namedisl.core import NamedIslObjectT @@ -126,7 +126,7 @@ def _get_access_relation( access_set = access_set & coordinates[cell_name].eq_set(index_aff) - return access_set.as_map(in_names=instance_names) + return _set_as_map(access_set, in_names=instance_names) def _record_access( self, @@ -228,6 +228,20 @@ def map_sub_array_ref( self._additional_inames = previous_inames +def _set_as_map( + set_: nisl.Set, in_names: Collection[str] +) -> nisl.Map: + if in_names: + return set_.as_map(in_names) + + # FIXME: Use Set.as_map directly once namedisl handles an empty input + # dimension space. + domain = set_.project_out(set_.space.out_names) + result = nisl.make_map_from_domain_and_range(domain, set_) + assert isinstance(result, nisl.Map) + return result + + def apply_affine_transform_to_happens_afters( kernel: LoopKernel, affine_reln: nisl.Map ) -> LoopKernel: @@ -811,7 +825,8 @@ def add_lexicographic_happens_after(kernel: LoopKernel) -> LoopKernel: else: ordered_instances = strict_lex | equal_prefix - instances_rel = ordered_instances.as_map( + instances_rel = _set_as_map( + ordered_instances, in_names=tuple(f"{name}_after" for name in after_inames) ) @@ -884,7 +899,8 @@ def record_conflicts( & incoming_instances_rel.as_set() & source_relation.as_set() ) - candidates = candidate_set.as_map( + candidates = _set_as_map( + candidate_set, in_names=(*sink_names, *cell_names) ) @@ -908,9 +924,10 @@ def record_conflicts( candidates .as_set() .project_out(cell_names) - .as_map(in_names=sink_names) - .coalesce() ) + required_order = _set_as_map( + required_order, in_names=sink_names + ).coalesce() previous = happens_after.get(source_id) if not required_order.is_empty(): if previous is None: @@ -926,7 +943,9 @@ def record_conflicts( # FIXME: remove unnamed conversion happens_after[source_id] = HappensAfter(combined_order.as_isl()) - return candidates.domain().as_map(in_names=sink_names).coalesce() + return _set_as_map( + candidates.domain(), in_names=sink_names + ).coalesce() def normalize_interface_and_compose( incoming_relation: nisl.Map, next_edge_relation: nisl.Map @@ -957,6 +976,15 @@ def normalize_interface_and_compose( # Write-after-write and write-after-read case AccessType.write: + # Readers must be recorded before a writer retires the live + # sink-cell relation. + if var in rel_finder.read_relations[source_id]: + source_relation = rel_finder.read_relations[source_id][var] + _ = record_conflicts( + source_relation, + select_most_recent_writer=False, + ) + if var in rel_finder.write_relations[source_id]: source_relation = rel_finder.write_relations[source_id][var] @@ -966,14 +994,6 @@ def normalize_interface_and_compose( ) live_access_rel = live_access_rel - caught_accesses - # don't update live_access_rel; does not find a "most recent writer" - if var in rel_finder.read_relations[source_id]: - source_relation = rel_finder.read_relations[source_id][var] - _ = record_conflicts( - source_relation, - select_most_recent_writer=False, - ) - case _: raise ValueError("unknown access type") @@ -1039,9 +1059,18 @@ def relax_strict_happens_after(kernel: LoopKernel) -> LoopKernel: coarse_dependency_graph: dict[str, frozenset[str]] = {} for stmt in kernel.instructions: - coarse_dependency_graph[stmt.id] = frozenset({ - dep for dep in stmt.happens_after if dep != stmt.id - }) + 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 + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + relation = nisl.make_map(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) diff --git a/test/test_dependency.py b/test/test_dependency.py index d42befd18..12d632cd7 100644 --- a/test/test_dependency.py +++ b/test/test_dependency.py @@ -1370,6 +1370,22 @@ def test_relax_strict_happens_after_tracks_scalar_accesses() -> None: ) +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 + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + required_order = nisl.make_map(required_order) + assert required_order.equals(nisl.make_map("{ [] -> [] }")) + + def test_relax_strict_happens_after_selects_writers_per_cell() -> None: kernel = _relax_strict_happens_after( """ @@ -1402,6 +1418,36 @@ def test_relax_strict_happens_after_selects_writers_per_cell() -> None: ) +def test_relax_strict_happens_after_records_readers_before_writer() -> 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 + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + required_order = nisl.make_map(required_order) + assert required_order.equals( + nisl.make_map(""" + [N] -> { + [i_after] -> [i_before] : + 0 <= i_before <= 1 and + i_before <= i_after < N + } + """) + ) + + @pytest.mark.parametrize( "instructions", ( @@ -1424,6 +1470,40 @@ def test_relax_strict_happens_after_drops_nonconflicting_edges( 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 + # FIXME: Remove conversion once HappensAfter stores namedisl.Map. + cross_order = nisl.make_map(cross_order) + 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.as_isl()), + }), + b_insn.copy(happens_after={ + "A": HappensAfter(empty_order.as_isl()), + "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 ): From 06027081616f63fa4c4086ea3fdf38cc925ff106 Mon Sep 17 00:00:00 2001 From: Addison Date: Tue, 18 Aug 2026 11:00:13 -0500 Subject: [PATCH 29/30] adapt precise verification to namedisl API --- loopy/schedule/verification.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/loopy/schedule/verification.py b/loopy/schedule/verification.py index 9bdd272b3..3dfb84c3b 100644 --- a/loopy/schedule/verification.py +++ b/loopy/schedule/verification.py @@ -204,7 +204,9 @@ def _build_statement_timestamp_relations( # FIXME: isl -> named conversion domain = nisl.make_set( kernel.get_inames_domain(stmt.within_inames).to_set() - ).project_out_except([*stmt.within_inames, *kernel.all_params()]) + ).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) @@ -248,7 +250,9 @@ def _build_barrier_timestamp_relations( # FIXME: isl -> named conversion domain = nisl.make_set( kernel.get_inames_domain(frozenset(inames)).to_set() - ).project_out_except([*inames, *kernel.all_params()]) + ).project_out_except( + [*inames, *kernel.all_params()], dim_type="all" + ) if inames: relation = nisl.make_set(f"{{[{full_str}]}}").as_map(inames) From 87508cef398e705657ef30b06d4305eb22e47236 Mon Sep 17 00:00:00 2001 From: Addison Date: Fri, 21 Aug 2026 15:29:50 -0500 Subject: [PATCH 30/30] fix global most-recent-writer selection --- loopy/kernel/dependency.py | 372 ++++++++++++++++++------------------- test/test_dependency.py | 115 +++++++++++- test/test_fusion.py | 14 +- test/test_loop_fusion.py | 13 +- test/test_transform.py | 13 +- 5 files changed, 329 insertions(+), 198 deletions(-) diff --git a/loopy/kernel/dependency.py b/loopy/kernel/dependency.py index 8097c0f87..3e68cf182 100644 --- a/loopy/kernel/dependency.py +++ b/loopy/kernel/dependency.py @@ -52,6 +52,7 @@ SubstitutionRuleExpander, WalkMapper, aff_from_expr, + get_dependencies, ) @@ -81,6 +82,7 @@ class AccessRelationFinder(WalkMapper[[str, AccessType]]): _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 @@ -91,6 +93,12 @@ def __init__(self, kernel: LoopKernel): 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) @@ -107,6 +115,16 @@ def _get_access_relation( 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) @@ -826,198 +844,140 @@ def add_lexicographic_happens_after(kernel: LoopKernel) -> LoopKernel: return kernel.copy(instructions=new_stmts) -def _relax_strict_happens_after_inner( +def _compute_reachable_happens_after( kernel: LoopKernel, - sink_id: str, - source_id: str, - var: str, - sink_access_type: AccessType, - incoming_instances_rel: nisl.Map, - live_access_rel: nisl.Map, - rel_finder: AccessRelationFinder, - happens_after: dict[str, HappensAfter], -) -> Mapping[str, HappensAfter]: - """ - Recursively finds conflicting accesses to *var* by *sink* and *source* to - determine the minimal required execution order between statement instances - of *source* and *sink*. + 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]] = {} - :arg sink_id: The ID of the statement whose instances will be in the domain - of the resulting :class:`namedisl.Map`. - - :arg source_id: The ID of the statement whose instances will be in the range - of the resulting :class:`namedisl.Map`. - - :arg var: The variable for which we are performing data dependence analysis. - - :arg sink_access_type: A :class:`AccessType` describing whether *sink_id* - reads or writes *var*. This determines how live instances are removed from - *live_access_rel*. - - :arg incoming_instances_rel: The incoming :class:`namedisl.Map` describing - how each sink and source instance are related. + 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() + ) - :arg live_access_rel: A :class:`namedisl.Map` describing the set of live - accesses by *sink_id* to *var*. When conflicts are found, the conflicting - relation is used to remove elements from this relation. + result[sink_id] = reachable - :arg rel_finder: A :class:`AccessRelationFinder` with access relations - constructed before entering this routine. + return result - :arg happens_after: A mapping from statement IDs to - :class:`loopy.HappensAfter` recording the dependencies from *sink* to all - statements in *happens_after*. - :returns: The updated precise dependencies for *sink_id*. - """ +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 - def record_conflicts( - source_relation: nisl.Map, - *, - select_most_recent_writer: bool, - ) -> nisl.Map: - source_relation = _suffix_names(source_relation, "_before", DimType.in_) + dominated = dominated | ( + candidate_relation + & _compose_happens_after_relations( + writer_relation, writer_to_candidate + ) + ) - sink_names = incoming_instances_rel.space.in_names - cell_names = live_access_rel.space.out_names + remaining = (candidate_relation - dominated).coalesce() + if not remaining.is_empty(): + result[candidate_id] = remaining - candidate_set = ( - live_access_rel.as_set() - & incoming_instances_rel.as_set() - & source_relation.as_set() - ) - candidates = _set_as_map( - candidate_set, - in_names=(*sink_names, *cell_names) - ) + return result - if select_most_recent_writer: - self_happens_after = kernel.id_to_insn[ - source_id - ].happens_after.get(source_id) - if self_happens_after is not None: - assert self_happens_after.instances_rel is not None - self_relation = self_happens_after.instances_rel - dominated = candidates & _compose_happens_after_relations( - candidates, self_relation - ) - candidates = (candidates - dominated).coalesce() - required_order = ( - candidates - .as_set() - .project_out(cell_names) - ) +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( - required_order, in_names=sink_names + candidate_relation.as_set().project_out(cell_names), + in_names=sink_names, ).coalesce() previous = happens_after.get(source_id) - if not required_order.is_empty(): - if previous is None: - combined_order = required_order - else: - assert previous.instances_rel is not None - previous_instances_rel = previous.instances_rel - 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, - ) - - return _set_as_map( - candidates.domain(), in_names=sink_names - ).coalesce() + if previous is None: + combined_order = required_order + else: + assert previous.instances_rel is not None + combined_order = ( + required_order | previous.instances_rel + ).coalesce() - def normalize_interface_and_compose( - incoming_relation: nisl.Map, next_edge_relation: nisl.Map - ) -> nisl.Map: - incoming_relation = incoming_relation.rename_dims( - (name, name[: len(name) - len("_before")]) - for name in incoming_relation.space.out_names + variable_name = ( + var + if previous is None or previous.variable_name == var + else None ) - - next_edge_relation = next_edge_relation.rename_dims( - (name, name[: len(name) - len("_after")]) - for name in next_edge_relation.space.in_names + happens_after[source_id] = HappensAfter( + combined_order, + variable_name=variable_name, ) - return incoming_relation.apply_range(next_edge_relation) - - match sink_access_type: - # Read-after-write - case AccessType.read: - if var in rel_finder.write_relations[source_id]: - source_relation = rel_finder.write_relations[source_id][var] - - caught_accesses = record_conflicts( - source_relation, - select_most_recent_writer=True, - ) - live_access_rel = live_access_rel - caught_accesses - - # Write-after-write and write-after-read - case AccessType.write: - # Readers must be recorded before a writer retires the live - # sink-cell relation. - if var in rel_finder.read_relations[source_id]: - source_relation = rel_finder.read_relations[source_id][var] - _ = record_conflicts( - source_relation, - select_most_recent_writer=False, - ) - - if var in rel_finder.write_relations[source_id]: - source_relation = rel_finder.write_relations[source_id][var] - - caught_accesses = record_conflicts( - source_relation, - select_most_recent_writer=True, - ) - live_access_rel = live_access_rel - caught_accesses - - case _: - raise ValueError("unknown access type") - - # Continue backward through the strict-order graph. - if not live_access_rel.is_empty() and (sink_id != source_id): - source_stmt = kernel.id_to_insn[source_id] - for src_dep_id, src_happens_after in source_stmt.happens_after.items(): - if src_dep_id == source_id: - continue - - if src_happens_after.instances_rel is None: - raise ValueError( - "All `HappensAfter`s must have precise dependencies " - "defined to use precise dependency finding machinery." - ) - - src_instances_rel = src_happens_after.instances_rel - outgoing_instances_rel = normalize_interface_and_compose( - incoming_instances_rel, src_instances_rel - ).coalesce() - - _relax_strict_happens_after_inner( - kernel, - sink_id, - src_dep_id, - var, - sink_access_type, - outgoing_instances_rel, - live_access_rel, - rel_finder, - happens_after, - ) - - return happens_after - @for_each_kernel def relax_strict_happens_after(kernel: LoopKernel) -> LoopKernel: @@ -1060,6 +1020,9 @@ def relax_strict_happens_after(kernel: LoopKernel) -> LoopKernel: 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: @@ -1073,7 +1036,6 @@ def relax_strict_happens_after(kernel: LoopKernel) -> LoopKernel: new_stmts: list[InstructionBase] = [] for sink_id in topological_order: new_happens_after: dict[str, HappensAfter] = {} - old_happens_after = kernel.id_to_insn[sink_id].happens_after for sink_access_type, access_relations in ( (AccessType.read, rel_finder.read_relations[sink_id]), (AccessType.write, rel_finder.write_relations[sink_id]), @@ -1082,22 +1044,56 @@ def relax_strict_happens_after(kernel: LoopKernel) -> LoopKernel: access_relation = _suffix_names( access_relation, "_after", DimType.in_ ) - for source_id, happens_after in old_happens_after.items(): - if happens_after.instances_rel is None: - raise ValueError( - "All `HappensAfter`s must have precise dependencies " - "defined to use precise dependency finding machinery." + 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, + ) - _relax_strict_happens_after_inner( - kernel, - sink_id, - source_id, + 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, - sink_access_type, - happens_after.instances_rel, - access_relation, - rel_finder, new_happens_after, ) diff --git a/test/test_dependency.py b/test/test_dependency.py index 758d26109..00e8651d8 100644 --- a/test/test_dependency.py +++ b/test/test_dependency.py @@ -1098,6 +1098,29 @@ def test_access_relation_finder_tracks_reads_and_writes_per_statement() -> None: ) +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( """ @@ -1396,7 +1419,7 @@ def test_relax_strict_happens_after_selects_writers_per_cell() -> None: ) -def test_relax_strict_happens_after_records_readers_before_writer() -> None: +def test_relax_strict_happens_after_selects_live_readers() -> None: t_unit = lp.make_kernel( "{ [i] : 0 <= i < N }", """ @@ -1416,12 +1439,19 @@ def test_relax_strict_happens_after_records_readers_before_writer() -> None: assert required_order.equals( nisl.make_map(""" [N] -> { - [i_after] -> [i_before] : - 0 <= i_before <= 1 and - i_before <= i_after < 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( @@ -1513,14 +1543,85 @@ def test_relax_strict_happens_after_tracks_live_footprints_through_a_chain() -> [NI, NJ] -> { [i_after, j_after] -> [i_before = i_after, j_before = j_after] : - 0 <= i_after < NI and - 0 <= j_after < NJ and - i_after mod 2 = 1 + 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( [ 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))