From f94613951e74731526def8ea4e66966e6f6a37d8 Mon Sep 17 00:00:00 2001 From: Leon Frenot Date: Mon, 6 Jul 2026 14:30:03 +0200 Subject: [PATCH 01/23] Fix strategy default values --- src/xtc/search/strategies.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/xtc/search/strategies.py b/src/xtc/search/strategies.py index e32caa54..8213067d 100644 --- a/src/xtc/search/strategies.py +++ b/src/xtc/search/strategies.py @@ -1033,7 +1033,7 @@ def __init__( self, graph: Graph, spec: dict[str, dict[str, Any]] | str, - constraints: list[str] = [], + constraints: list[str] | None = None, partial_tiles: bool = False, partial_unrolls: bool = False, initialize: bool = True, @@ -1056,7 +1056,9 @@ def __init__( self._loop_nest = loop_nest input_constraints = loop_nest.collect_constraints() self._orders: dict[str, list] = {} - self._constraints = constraints + input_constraints + self._constraints = ( + constraints + input_constraints if constraints else input_constraints + ) self._initialized = False if initialize: self._initialize() @@ -1158,7 +1160,7 @@ def __init__( self, graph: Graph, spec: dict[str, dict[str, Any]] | str, - constraints: list[str] = [], + constraints: list[str] | None = None, partial_tiles: bool = False, partial_unrolls: bool = False, initialize: bool = True, @@ -1179,7 +1181,6 @@ def sample(self, num: int, seed: int | None = 0) -> Iterator[Sample]: @override def generate(self, scheduler: Scheduler, sample: Sample) -> None: - assert isinstance(self._sample_shape, list) sample = dict(zip(self._sample_shape, sample)) super().generate(scheduler, sample) From 5950318bcc718dd7b2c9a2bfa87c489424c4b4e6 Mon Sep 17 00:00:00 2001 From: Leon Frenot Date: Mon, 6 Jul 2026 14:31:33 +0200 Subject: [PATCH 02/23] Fix tensor name and allow nondefault axis names --- src/xtc/search/strategies.py | 24 +++++++++++++++++++++--- 1 file changed, 21 insertions(+), 3 deletions(-) diff --git a/src/xtc/search/strategies.py b/src/xtc/search/strategies.py index 8213067d..13afb993 100644 --- a/src/xtc/search/strategies.py +++ b/src/xtc/search/strategies.py @@ -1034,6 +1034,8 @@ def __init__( graph: Graph, spec: dict[str, dict[str, Any]] | str, constraints: list[str] | None = None, + axes: list[str] | None = None, + tensors: list[str] | None = None, partial_tiles: bool = False, partial_unrolls: bool = False, initialize: bool = True, @@ -1041,13 +1043,20 @@ def __init__( self._graph = graph self._op = graph.outputs_nodes[0].operation self._stats: dict[str, int] = {} - self._axes = list(self._op.dims) + self._axes = axes if axes else list(self._op.dims) self._sizes = self._constant_sizes() self._sample_names: list[str] = [] + tensors = [ + chr(ord("A") + i) + for i in range( + len(graph.outputs_nodes[0].inputs) + + len(graph.outputs_nodes[0].outputs) + ) + ] descript = Descript( abstract_dims=self._axes, abstract_dim_sizes=dict(self._sizes), - abstract_matrix=["A", "B", "C"], + abstract_matrix=tensors, partial_tiles=partial_tiles, partial_unrolls=partial_unrolls, ) @@ -1161,13 +1170,22 @@ def __init__( graph: Graph, spec: dict[str, dict[str, Any]] | str, constraints: list[str] | None = None, + axes: list[str] | None = None, + tensors: list[str] | None = None, partial_tiles: bool = False, partial_unrolls: bool = False, initialize: bool = True, ) -> None: self._sample_shape: list[str] = [] super().__init__( - graph, spec, constraints, partial_tiles, partial_unrolls, initialize + graph, + spec, + constraints, + axes, + tensors, + partial_tiles, + partial_unrolls, + initialize, ) @override From aea7a856360d757733ce538c7ed9805739a9ea6c Mon Sep 17 00:00:00 2001 From: Leon Frenot Date: Mon, 6 Jul 2026 16:42:41 +0200 Subject: [PATCH 03/23] Starting to implement Ansor-style PRT for descript --- src/xtc/schedules/descript.py | 112 ++++++++++++++++++++++++++++++++-- src/xtc/schedules/parsing.py | 28 ++++++++- 2 files changed, 134 insertions(+), 6 deletions(-) diff --git a/src/xtc/schedules/descript.py b/src/xtc/schedules/descript.py index d1100508..e9ecac7f 100644 --- a/src/xtc/schedules/descript.py +++ b/src/xtc/schedules/descript.py @@ -15,16 +15,18 @@ ParameterSplitOrigin, ) -from .exceptions import ScheduleInterpretError +from .exceptions import ScheduleInterpretError, ScheduleParseError from .parsing import ( ScheduleParser, ScheduleSpec, SplitDecl, TileDecl, AxisDecl, + PRTDecl, Annotations, YAMLParser, literal, + ansor_tile, ) @@ -66,9 +68,19 @@ class ScheduleInterpreter: abstract_dims: list[str] abstract_dim_sizes: dict[str, int] = field(default_factory=dict) + abstract_p_dims: list[str] = field(default_factory=list) + abstract_r_dims: list[str] = field(default_factory=list) abstract_matrix: list[str] = field(default_factory=list) partial_tiles: bool = False partial_unrolls: bool = False + _abstract_dim_fresh: dict[str, int] = field(default_factory=dict) + + def __post_init__(self): + self._abstract_dim_fresh = {i: -1 for i in self.abstract_dims} + + def _fresh(self, axis: str): + self._abstract_dim_fresh[axis] += 1 + return f"prt_{axis}_{self._abstract_dim_fresh[axis]}" def interpret(self, spec: ScheduleSpec, root: str) -> ParameterLoopNest: """Interpret a schedule specification into a LoopNest.""" @@ -127,6 +139,14 @@ def _interpret_spec_into_node( if _loop_name: loop_name = _loop_name self._apply_annotations(item.annotations, loop_name, sizes, node) + elif isinstance(item, PRTDecl): + self._interpret_prt( + item=item, + node=node, + interchange=interchange, + sizes=sizes, + axes=axes, + ) # Reaplace the placeholder of the last split with its size if len(last_split) > 0: @@ -332,6 +352,82 @@ def _interpret_axis( interchange.append(axis_name) return axis_name + def _interpret_prt( + self, + item: PRTDecl, + node: ParameterLoopNestNode, + interchange: list[str], + sizes: dict[str, literal], + axes: dict[str, list[literal]], + ): + for t in item.shape: + match t: + case "P": + for a in self.abstract_p_dims: + self._interpret_tile( + TileDecl(a, self._fresh(a), Annotations()), + node, + interchange, + sizes, + axes, + ) + # TODO: annotations + case "R": + for a in self.abstract_r_dims: + self._interpret_tile( + TileDecl(a, self._fresh(a), Annotations()), + node, + interchange, + sizes, + axes, + ) + # TODO: annotations + case "T": + for a in self.abstract_dims: + self._interpret_tile( + TileDecl(a, self._fresh(a), Annotations()), + node, + interchange, + sizes, + axes, + ) + # TODO: annotations + case "U": + raise ScheduleInterpretError( + "TODO: Ansor-style U tile unimplemented (random order)" + ) + case "O": + a = self.abstract_p_dims[0] + self._interpret_tile( + TileDecl(a, self._fresh(a), Annotations()), + node, + interchange, + sizes, + axes, + ) + for a in self.abstract_r_dims: + self._interpret_tile( + TileDecl(a, self._fresh(a), Annotations()), + node, + interchange, + sizes, + axes, + ) + for a in self.abstract_p_dims[1:]: + self._interpret_tile( + TileDecl(a, self._fresh(a), Annotations()), + node, + interchange, + sizes, + axes, + ) + case "W": + raise ScheduleInterpretError( + "TODO: Write buffer in Ansor-styke unimplemented" + ) + case "F": + raise ScheduleInterpretError("TODO: Fusion unimplemented") + def _check_axis_existence(self, axis: str) -> None: """Check that an axis is defined.""" if axis not in self.abstract_dims: @@ -443,25 +539,33 @@ class Descript: abstract_dims: list[str] abstract_dim_sizes: dict[str, int] = field(default_factory=dict) + abstract_p_dims: list[str] = field(default_factory=list) + abstract_r_dims: list[str] = field(default_factory=list) abstract_matrix: list[str] = field(default_factory=list) partial_tiles: bool = False partial_unrolls: bool = False + def __post_init__(self): + for m in self.abstract_matrix: + if m in ansor_tile: + raise ScheduleParseError(f"Forbidden abstract matrix name: {m}") + def loop_nest(self, node_name: str, spec: dict[str, dict[str, Any]] | str): if isinstance(spec, str): yaml_parser = YAMLParser() spec = yaml_parser.parse(spec) - spec_dict = cast(dict[str, dict[str, Any]], spec) - constraints = cast(list[str], spec_dict.pop("constraints", [])) + constraints = cast(list[str], spec.pop("constraints", [])) # Parse the specification into an AST parser = ScheduleParser() - ast = parser.parse(spec_dict) + ast = parser.parse(spec) # Interpret the AST into a LoopNest interpreter = ScheduleInterpreter( abstract_dims=self.abstract_dims, abstract_dim_sizes=self.abstract_dim_sizes, + abstract_p_dims=self.abstract_p_dims, + abstract_r_dims=self.abstract_r_dims, abstract_matrix=self.abstract_matrix, partial_tiles=self.partial_tiles, partial_unrolls=self.partial_unrolls, diff --git a/src/xtc/schedules/parsing.py b/src/xtc/schedules/parsing.py index bbea72af..00a109c2 100644 --- a/src/xtc/schedules/parsing.py +++ b/src/xtc/schedules/parsing.py @@ -13,6 +13,7 @@ from .exceptions import ScheduleParseError literal = int | str +ansor_tile = "PRTUOWF" def toliteral(s: str) -> literal: @@ -93,7 +94,21 @@ class AxisDecl: annotations: Annotations -ScheduleItem = SplitDecl | TileDecl | AxisDecl +@dataclass(frozen=True) +class PRTDecl: + """AST Type: Ansor-style tile declarations""" + + shape: str + annotations: Annotations + + def __post__init__(self): + for s in self.shape: + if s not in ansor_tile: + # Should be unreachable + raise ScheduleParseError(f"Invalid tile declaration {s}") + + +ScheduleItem = SplitDecl | TileDecl | AxisDecl | PRTDecl @dataclass(frozen=True) @@ -131,6 +146,9 @@ def _parse_declaration(self, declaration: str, value: Any) -> ScheduleItem: if "#" in declaration: return self._parse_tile(declaration, value) + if declaration[0] in ansor_tile: + return self._parse_ansor_tile(declaration, value) + # Must be a direct axis reference return self._parse_axis_ref(declaration, value) @@ -156,6 +174,12 @@ def _parse_tile(self, declaration: str, value: dict) -> TileDecl: annotations = self._parse_annotations(value, declaration) return TileDecl(axis=axis_name, size=size, annotations=annotations) + def _parse_ansor_tile(self, declaration: str, value: dict) -> PRTDecl: + """Parse an ansor-style tile reference.""" + + annotations = self._parse_annotations(value, declaration) + return PRTDecl(shape=declaration, annotations=annotations) + def _parse_axis_ref(self, declaration: str, value: dict) -> AxisDecl: """Parse a direct axis reference.""" @@ -302,7 +326,7 @@ def _parse_split_syntax( class YAMLParser: """Parses a YAML specification into a dict-based schedule specification.""" - def parse(self, spec: str): + def parse(self, spec: str) -> dict[str, dict[str, Any]]: """Parses a YAML specification into a dict-based schedule specification.""" descript_spec = strictyaml.load(spec).data if not isinstance(descript_spec, dict): From d1296f92888f60def072d227da357e82a1a26464 Mon Sep 17 00:00:00 2001 From: Leon Frenot Date: Tue, 7 Jul 2026 14:24:36 +0200 Subject: [PATCH 04/23] Annotations for descript PRT TODO: change dict representation to list of tuples to support separating every PRT --- src/xtc/schedules/descript.py | 121 ++++++++++++++++-- src/xtc/search/strategies.py | 2 + .../search/test_matmul_descript_pprprp.py | 22 ++++ 3 files changed, 131 insertions(+), 14 deletions(-) create mode 100644 tests/filecheck/search/test_matmul_descript_pprprp.py diff --git a/src/xtc/schedules/descript.py b/src/xtc/schedules/descript.py index e9ecac7f..0a6832af 100644 --- a/src/xtc/schedules/descript.py +++ b/src/xtc/schedules/descript.py @@ -140,7 +140,7 @@ def _interpret_spec_into_node( loop_name = _loop_name self._apply_annotations(item.annotations, loop_name, sizes, node) elif isinstance(item, PRTDecl): - self._interpret_prt( + loop_name = self._interpret_prt( item=item, node=node, interchange=interchange, @@ -359,44 +359,107 @@ def _interpret_prt( interchange: list[str], sizes: dict[str, literal], axes: dict[str, list[literal]], - ): - for t in item.shape: + ) -> str: + loop_name = "" + + parallelize = item.annotations.parallelize + unroll = item.annotations.unroll_specified + + def _parallelize(): + if isinstance(parallelize, str): + node.parallelize_parameters[loop_name] = parallelize + node.constraints.append(f"{parallelize} in {{0, 1}}") + else: + node.parallelize.append(loop_name) + + def _unroll(): + unroll_factor = item.annotations.unroll_factor + if unroll_factor is None: + # None means "unroll fully" - use the loop size + if loop_name not in sizes: + raise ScheduleInterpretError( + f"{loop_name}'s size being unknown, an unroll factor is needed." + ) + unroll_factor = sizes[loop_name] + elif isinstance(unroll_factor, int) and unroll_factor <= 0: + raise ScheduleInterpretError( + f'`{{"unroll" = {unroll_factor}}}`: unroll parameter should be strictly positive.' + ) + elif isinstance(unroll_factor, str): + unroll_factor += "_prt_" + loop_name + if self.partial_unrolls: + node.constraints.append(f"{unroll_factor} <= {sizes[loop_name]}") + else: + node.constraints.append(f"{unroll_factor} || {sizes[loop_name]}") + node.unroll[loop_name] = unroll_factor + + for idx, t in enumerate(item.shape): match t: case "P": + if not self.abstract_p_dims: + raise ScheduleInterpretError( + "P scheme used, but P axes are not specified" + ) for a in self.abstract_p_dims: - self._interpret_tile( - TileDecl(a, self._fresh(a), Annotations()), + fresh_a = self._fresh(a) + loop_name = self._interpret_tile( + TileDecl(a, fresh_a, Annotations()), node, interchange, sizes, axes, ) - # TODO: annotations + if parallelize: + _parallelize() + parallelize = False + if unroll and idx + 1 == len(item.shape): + _unroll() case "R": + if not self.abstract_r_dims: + raise ScheduleInterpretError( + "R scheme used, but R axes are not specified" + ) for a in self.abstract_r_dims: - self._interpret_tile( + loop_name = self._interpret_tile( TileDecl(a, self._fresh(a), Annotations()), node, interchange, sizes, axes, ) - # TODO: annotations + if parallelize: + raise ScheduleInterpretError( + "Cannot parallelize a reduction axis" + ) + if unroll and idx + 1 == len(item.shape): + _unroll() case "T": for a in self.abstract_dims: - self._interpret_tile( + loop_name = self._interpret_tile( TileDecl(a, self._fresh(a), Annotations()), node, interchange, sizes, axes, ) - # TODO: annotations + if parallelize: + _parallelize() + parallelize = False + if unroll and idx + 1 == len(item.shape): + _unroll() case "U": raise ScheduleInterpretError( "TODO: Ansor-style U tile unimplemented (random order)" ) case "O": + if not self.abstract_p_dims: + raise ScheduleInterpretError( + "O scheme used, but P axes are not specified" + ) + if not self.abstract_r_dims: + raise ScheduleInterpretError( + "O scheme used, but R axes are not specified" + ) a = self.abstract_p_dims[0] self._interpret_tile( TileDecl(a, self._fresh(a), Annotations()), @@ -405,6 +468,11 @@ def _interpret_prt( sizes, axes, ) + if parallelize: + _parallelize() + parallelize = False + if unroll and idx + 1 == len(item.shape): + _unroll() for a in self.abstract_r_dims: self._interpret_tile( TileDecl(a, self._fresh(a), Annotations()), @@ -413,21 +481,46 @@ def _interpret_prt( sizes, axes, ) + if unroll and idx + 1 == len(item.shape): + _unroll() for a in self.abstract_p_dims[1:]: - self._interpret_tile( + loop_name = self._interpret_tile( TileDecl(a, self._fresh(a), Annotations()), node, interchange, sizes, axes, ) + if unroll and idx + 1 == len(item.shape): + _unroll() case "W": - raise ScheduleInterpretError( - "TODO: Write buffer in Ansor-styke unimplemented" - ) + node.buffer_at[loop_name] = None case "F": raise ScheduleInterpretError("TODO: Fusion unimplemented") + if item.annotations.vectorize: + if isinstance(item.annotations.vectorize, str): + node.vectorize_parameters[loop_name] = item.annotations.vectorize + node.constraints.append(f"{item.annotations.vectorize} in {{0, 1}}") + else: + node.vectorize.append(loop_name) + + if item.annotations.buffer_specified: + node.buffer_at[loop_name] = item.annotations.buffer + + if item.annotations.pack_specified and item.annotations.pack is not None: + input_matrix, mtype, pad = item.annotations.pack + if isinstance(input_matrix, str): + idx = self.abstract_matrix.index(input_matrix) + if idx == len(self.abstract_matrix) - 1: + node.buffer_at[loop_name] = mtype + else: + node.pack_at[loop_name] = (idx, mtype, pad) + else: + node.pack_at[loop_name] = (input_matrix, mtype, pad) + + return loop_name + def _check_axis_existence(self, axis: str) -> None: """Check that an axis is defined.""" if axis not in self.abstract_dims: diff --git a/src/xtc/search/strategies.py b/src/xtc/search/strategies.py index 13afb993..2e934648 100644 --- a/src/xtc/search/strategies.py +++ b/src/xtc/search/strategies.py @@ -1057,6 +1057,8 @@ def __init__( abstract_dims=self._axes, abstract_dim_sizes=dict(self._sizes), abstract_matrix=tensors, + abstract_p_dims=list(self._op.dims_kind("P")), + abstract_r_dims=list(self._op.dims_kind("R")), partial_tiles=partial_tiles, partial_unrolls=partial_unrolls, ) diff --git a/tests/filecheck/search/test_matmul_descript_pprprp.py b/tests/filecheck/search/test_matmul_descript_pprprp.py new file mode 100644 index 00000000..d93ef6c6 --- /dev/null +++ b/tests/filecheck/search/test_matmul_descript_pprprp.py @@ -0,0 +1,22 @@ +# RUN: python %s 2>&1 | filecheck %s +# REQUIRES: module_xvs +""" +Test strategy Goto on matmul +""" + +import utils +from xtc.search.strategies import Strategy_Descript as Strategy + +graph = utils.get_graph_matmul() +backend = utils.get_backend(graph) +spec = """ +PPRP: parallelize +R: unroll=u_k +P: unroll vectorize""" +strategy = Strategy(graph, spec, initialize=False) + +print(sorted(strategy._constraints)) +print(sum(1 for _ in strategy.sample(100))) + +# CHECK: ['prt_i_0 || {21}', 'prt_i_1 || {21, prt_i_0}', 'prt_i_2 || {21, prt_i_0, prt_i_1}', 'prt_i_3 || {21, prt_i_0, prt_i_1, prt_i_2}', 'prt_j_0 || {32}', 'prt_j_1 || {32, prt_j_0}', 'prt_j_2 || {32, prt_j_0, prt_j_1}', 'prt_j_3 || {32, prt_j_0, prt_j_1, prt_j_2}', 'prt_k_0 || {12}', 'prt_k_1 || {12, prt_k_0}', 'u_k_prt_k1 || prt_k_1'] +# CHECK-NEXT: 100 From 19714ca7a6f470b4ea4ab80052a22e6a2f7b645e Mon Sep 17 00:00:00 2001 From: Leon Frenot Date: Tue, 7 Jul 2026 17:38:07 +0200 Subject: [PATCH 05/23] Change parsing to support lists of tuples for mutiple entries with the same key --- src/xtc/schedules/descript.py | 15 ++++-- src/xtc/schedules/parsing.py | 91 +++++++++++++++++++++++------------ 2 files changed, 72 insertions(+), 34 deletions(-) diff --git a/src/xtc/schedules/descript.py b/src/xtc/schedules/descript.py index 0a6832af..22a40b16 100644 --- a/src/xtc/schedules/descript.py +++ b/src/xtc/schedules/descript.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: BSD-3-Clause # Copyright (c) 2024-2026 The XTC Project Authors # -from typing import Any, cast +from typing import Any from dataclasses import dataclass, field from copy import deepcopy @@ -27,6 +27,8 @@ YAMLParser, literal, ansor_tile, + tup_list, + pre_parse, ) @@ -643,12 +645,17 @@ def __post_init__(self): if m in ansor_tile: raise ScheduleParseError(f"Forbidden abstract matrix name: {m}") - def loop_nest(self, node_name: str, spec: dict[str, dict[str, Any]] | str): + def loop_nest( + self, node_name: str, spec: dict[str, dict[str, Any]] | tup_list | str + ): if isinstance(spec, str): yaml_parser = YAMLParser() spec = yaml_parser.parse(spec) - - constraints = cast(list[str], spec.pop("constraints", [])) + else: + spec = pre_parse(spec) + constraints: Any = [v for k, v in spec if k == "constraints"] + constraints = constraints[0] if constraints else None + spec = [(k, v) for k, v in spec if k != "constraints"] # Parse the specification into an AST parser = ScheduleParser() ast = parser.parse(spec) diff --git a/src/xtc/schedules/parsing.py b/src/xtc/schedules/parsing.py index 00a109c2..2907067e 100644 --- a/src/xtc/schedules/parsing.py +++ b/src/xtc/schedules/parsing.py @@ -7,13 +7,14 @@ from typing import Any from dataclasses import dataclass import re -import strictyaml +import yaml from typing_extensions import override from .exceptions import ScheduleParseError literal = int | str ansor_tile = "PRTUOWF" +tup_list = list[tuple[str, Any]] def toliteral(s: str) -> literal: @@ -23,6 +24,28 @@ def toliteral(s: str) -> literal: return s +def pre_parse(s: dict[str, Any] | tup_list | list[dict[str, Any]]) -> tup_list: + if isinstance(s, dict): + return [(k, _pre_parse(v)) for k, v in s.items()] + out: tup_list = [] + for s_ in s: + if isinstance(s_, dict): + out += [(k, _pre_parse(v)) for k, v in s_.items()] + else: + k, v = s_ + out.append((k, _pre_parse(v))) + return out + + +def _pre_parse(v: Any) -> Any: + if isinstance(v, dict): + return pre_parse(v) + elif isinstance(v, list) and isinstance(v[0], dict | list): + return pre_parse(v) + else: + return v + + @dataclass(frozen=True) class Annotations: """AST Type : annotations that can be applied to a loop. @@ -125,19 +148,21 @@ class ScheduleParser: _SPLIT_PATTERN = re.compile(r"^(.*)\[(-\w+|\w*)?:(-\w+|\w*)?\]$") _SPLIT_MIDDLE_PATTERN = re.compile(r"^(.*)\[:(\w*):\]$") - def parse(self, spec: dict[str, Any]) -> ScheduleSpec: + def parse(self, spec: list[tuple[str, list]] | dict[str, Any]) -> ScheduleSpec: """Parse a schedule specification dict into an AST.""" + if isinstance(spec, dict): + spec = pre_parse(spec) + items: list[ScheduleItem] = [] - for declaration, value in spec.items(): + for declaration, value in spec: item = self._parse_declaration(declaration, value) items.append(item) return ScheduleSpec(items=tuple(items)) - def _parse_declaration(self, declaration: str, value: Any) -> ScheduleItem: + def _parse_declaration(self, declaration: str, value: list) -> ScheduleItem: """Parse a single declaration into a ScheduleItem.""" - assert isinstance(value, dict) # Try split declaration first (e.g., "axis[0:10]") if ":" in declaration: return self._parse_split(declaration, value) @@ -152,14 +177,14 @@ def _parse_declaration(self, declaration: str, value: Any) -> ScheduleItem: # Must be a direct axis reference return self._parse_axis_ref(declaration, value) - def _parse_split(self, declaration: str, value: dict) -> SplitDecl: + def _parse_split(self, declaration: str, value: list) -> SplitDecl: """Parse a split declaration like 'axis[start:end]' or 'axis[:size:]'.""" axis_name, start, end, size = self._parse_split_syntax(declaration) body = self.parse(value) return SplitDecl(axis=axis_name, start=start, end=end, body=body, size=size) - def _parse_tile(self, declaration: str, value: dict) -> TileDecl: + def _parse_tile(self, declaration: str, value: list) -> TileDecl: """Parse a tile declaration like 'axis#size'.""" parts = declaration.split("#") if len(parts) != 2: @@ -174,19 +199,21 @@ def _parse_tile(self, declaration: str, value: dict) -> TileDecl: annotations = self._parse_annotations(value, declaration) return TileDecl(axis=axis_name, size=size, annotations=annotations) - def _parse_ansor_tile(self, declaration: str, value: dict) -> PRTDecl: + def _parse_ansor_tile(self, declaration: str, value: list) -> PRTDecl: """Parse an ansor-style tile reference.""" annotations = self._parse_annotations(value, declaration) return PRTDecl(shape=declaration, annotations=annotations) - def _parse_axis_ref(self, declaration: str, value: dict) -> AxisDecl: + def _parse_axis_ref(self, declaration: str, value: list) -> AxisDecl: """Parse a direct axis reference.""" annotations = self._parse_annotations(value, declaration) return AxisDecl(axis=declaration, annotations=annotations) - def _parse_annotations(self, value: dict[str, Any], context: str) -> Annotations: + def _parse_annotations( + self, value: list[tuple[str, Any]], context: str + ) -> Annotations: """Parse annotation dict into Annotations object.""" unroll_factor: literal | None = None @@ -200,7 +227,7 @@ def _parse_annotations(self, value: dict[str, Any], context: str) -> Annotations partial = False full = False - for key, param in value.items(): + for key, param in value: match key: case "unroll": if param is True or param is None: @@ -326,51 +353,55 @@ def _parse_split_syntax( class YAMLParser: """Parses a YAML specification into a dict-based schedule specification.""" - def parse(self, spec: str) -> dict[str, dict[str, Any]]: + def parse(self, spec: str) -> tup_list: """Parses a YAML specification into a dict-based schedule specification.""" - descript_spec = strictyaml.load(spec).data - if not isinstance(descript_spec, dict): + descript_spec = yaml.safe_load(spec) + if not isinstance(descript_spec, dict | list): raise ScheduleParseError( f"Wrong format: YAML input parses to {type(descript_spec)}." ) + descript_spec = pre_parse(descript_spec) return self._parse(descript_spec) - def _parse(self, spec: dict[str, Any]) -> dict[str, dict]: + def _parse(self, spec: tup_list) -> tup_list: """Parses a dict YAML specification into a schedule specification.""" - constraints = spec.pop("constraints", []) - descript_spec: dict[str, dict] = {} - for a, v in spec.items(): + # constraints = spec.pop("constraints", []) + descript_spec: tup_list = [] + for a, v in spec: + if a == "constraints": + descript_spec.append((a, v)) + continue if isinstance(v, str): d = self._split(v) - elif isinstance(v, dict): + elif isinstance(v, list): d = v + elif v is None: + d = [] else: raise ScheduleParseError( f"Value {v} of key {a} is neither a string nor a dict." ) - size = d.get("size", None) + size = [v for k, v in d if k == "size"] if size: - d.pop("size") - a = f"{a}#{size}" + d = [(k, v) for k, v in d if k != "size"] + a = f"{a}#{size[0]}" if ":" in a: - descript_spec[a] = self._parse(d) + descript_spec.append((a, self._parse(d))) else: - descript_spec[a] = d - if constraints: - descript_spec["constraints"] = constraints + descript_spec.append((a, d)) return descript_spec - def _split(self, s: str) -> dict[str, Any]: + def _split(self, s: str) -> tup_list: """Splits a string of 'keyword's and 'keyword=value's separated by spaces into a dict.""" - d: dict[str, Any] = {} + d: tup_list = [] for s in s.split(): if "=" not in s: - d[s] = None + d.append((s, None)) else: x, y = s.split("=") try: tmp = eval(y) except (NameError, SyntaxError): tmp = y - d[x] = tmp + d.append((x, tmp)) return d From da72a912eb8f1146e89b5bb3478fbeec442fc555 Mon Sep 17 00:00:00 2001 From: Leon Frenot Date: Wed, 8 Jul 2026 15:07:39 +0200 Subject: [PATCH 06/23] Parameters for packing and padding --- src/xtc/schedules/descript.py | 17 +++++++ src/xtc/schedules/parameter_loop_nest.py | 21 ++++++++- src/xtc/schedules/parsing.py | 45 ++++++++++++------- .../search/test_matmul_descript_yaml_goto.py | 7 ++- 4 files changed, 70 insertions(+), 20 deletions(-) diff --git a/src/xtc/schedules/descript.py b/src/xtc/schedules/descript.py index 22a40b16..0f85bfb1 100644 --- a/src/xtc/schedules/descript.py +++ b/src/xtc/schedules/descript.py @@ -581,10 +581,27 @@ def _apply_annotations( idx = self.abstract_matrix.index(input_matrix) if idx == len(self.abstract_matrix) - 1: node.buffer_at[loop_name] = mtype + if isinstance(annotations.pack_specified, str): + node.buffer_parameters[loop_name] = annotations.pack_specified + node.constraints.append( + f"{annotations.buffer_specified} in {{0, 1}}" + ) else: node.pack_at[loop_name] = (idx, mtype, pad) + if isinstance(annotations.pack_specified, str): + node.pack_parameters[loop_name] = annotations.pack_specified + node.constraints.append( + f"{annotations.pack_specified} in {{0, 1}}" + ) + if isinstance(pad, str): + node.constraints.append(f"{pad} in {{0,1}}") else: node.pack_at[loop_name] = (input_matrix, mtype, pad) + if isinstance(annotations.pack_specified, str): + node.pack_parameters[loop_name] = annotations.pack_specified + node.constraints.append(f"{annotations.pack_specified} in {{0, 1}}") + if isinstance(pad, str): + node.constraints.append(f"{pad} in {{0,1}}") def _check_splitting_intervals( self, diff --git a/src/xtc/schedules/parameter_loop_nest.py b/src/xtc/schedules/parameter_loop_nest.py index ee1eba5b..b1b76ca7 100644 --- a/src/xtc/schedules/parameter_loop_nest.py +++ b/src/xtc/schedules/parameter_loop_nest.py @@ -120,7 +120,9 @@ class ParameterLoopNestNode(Node["ParameterLoopNestNode"]): parallelize_parameters: dict[str, str] = field(default_factory=dict) unroll: dict[str, literal] = field(default_factory=dict) buffer_at: dict[str, str | None] = field(default_factory=dict) + buffer_parameters: dict[str, str] = field(default_factory=dict) pack_at: dict[str, tuple[int, str | None, bool | str]] = field(default_factory=dict) + pack_parameters: dict[str, str] = field(default_factory=dict) constraints: list[str] = field(default_factory=list) def apply_sample(self, sample: dict[str, int]) -> LoopNestNode: @@ -128,6 +130,7 @@ def apply_sample(self, sample: dict[str, int]) -> LoopNestNode: Replaces the string parameters with the correspondings values in sample. And builds the corresponding LoopNestNode.""" root = self.root + tiles = { a: { b: sample[v_b] if isinstance(v_b, str) else v_b @@ -135,6 +138,7 @@ def apply_sample(self, sample: dict[str, int]) -> LoopNestNode: } for a, v_a in self.tiles.items() } + splits = { a: { b: sample[v_b] if isinstance(v_b, str) else v_b @@ -142,30 +146,45 @@ def apply_sample(self, sample: dict[str, int]) -> LoopNestNode: } for a, v_a in self.splits.items() } + interchange = self.interchange + vectorize = self.vectorize for a, v in self.vectorize_parameters.items(): if sample.get(v, False): vectorize.append(a) + parallelize = self.parallelize for a, v in self.parallelize_parameters.items(): if sample.get(v, False): parallelize.append(a) + unroll = { a: sample[v_a] if isinstance(v_a, str) else v_a for a, v_a in self.unroll.items() } - buffer_at = self.buffer_at + + buffer_at = self.buffer_at.copy() + for a, v in self.buffer_parameters.items(): + if not sample.get(v, True): + buffer_at.pop(a) + pack_at = { a: (a1, a2, bool(sample[a3]) if isinstance(a3, str) else a3) for a, (a1, a2, a3) in self.pack_at.items() } + for a, v in self.pack_parameters.items(): + if not sample.get(v, True): + pack_at.pop(a) + children = [child.apply_sample(sample) for child in self.children] + split_origin = ( self.split_origin.apply_sample(sample) if self.split_origin is not None else None ) + return LoopNestNode( root=root, tiles=tiles, diff --git a/src/xtc/schedules/parsing.py b/src/xtc/schedules/parsing.py index 2907067e..773bccbd 100644 --- a/src/xtc/schedules/parsing.py +++ b/src/xtc/schedules/parsing.py @@ -69,9 +69,9 @@ class Annotations: vectorize: bool | str = False parallelize: bool | str = False buffer: str | None = None - buffer_specified: bool = False + buffer_specified: bool | str = False pack: tuple[literal, str | None, bool | str] | None = None - pack_specified: bool = False + pack_specified: bool | str = False partial: bool = False full: bool = False @@ -212,7 +212,7 @@ def _parse_axis_ref(self, declaration: str, value: list) -> AxisDecl: return AxisDecl(axis=declaration, annotations=annotations) def _parse_annotations( - self, value: list[tuple[str, Any]], context: str + self, value: list[tuple[str, Any]], declaration: str ) -> Annotations: """Parse annotation dict into Annotations object.""" @@ -221,9 +221,9 @@ def _parse_annotations( vectorize: bool | str = False parallelize: bool | str = False buffer: str | None = None - buffer_specified = False + buffer_specified: bool | str = False pack: tuple[literal, str | None, bool | str] | None = None - pack_specified = False + pack_specified: bool | str = False partial = False full = False @@ -266,17 +266,32 @@ def _parse_annotations( buffer = None if param == "default" else param buffer_specified = True case "pack": - pack = self._parse_pack_param(param, context) - pack_specified = True + if isinstance(param, str): + pack = (declaration, None, False) + pack_specified = param + else: + pack = self._parse_pack_param(param, declaration) + pack_specified = True + case "pad": + if pack is None: + raise ScheduleParseError( + f"pad annotation before/without pack on {declaration}: {key}" + ) + declaration_, mtype_, pad_ = pack + pack = (declaration_, mtype_, param) case "partial": partial = True case "full": full = True case _: - raise ScheduleParseError(f"Unknown annotation on {context}: {key}") + raise ScheduleParseError( + f"Unknown annotation on {declaration}: {key}" + ) if partial and full: - raise ScheduleParseError(f"{context} has both annotations full and partial") + raise ScheduleParseError( + f"{declaration} has both annotations full and partial" + ) return Annotations( unroll_factor=unroll_factor, @@ -292,33 +307,33 @@ def _parse_annotations( ) def _parse_pack_param( - self, param: Any, context: str + self, param: Any, declaration: str ) -> tuple[literal, str | None, bool | str] | None: """Parse pack parameter into (input_idx, mtype, pad) tuple.""" if param is None: - return (context, None, False) + return (declaration, None, False) if not isinstance(param, (list, tuple)) or len(param) != 3: raise ScheduleParseError( - f'`{{"pack" = {param}}}` on {context}: pack parameter should be a tuple (input_idx, mtype, pad).' + f'`{{"pack" = {param}}}` on {declaration}: pack parameter should be a tuple (input_idx, mtype, pad).' ) input_idx, mtype, pad = param if not isinstance(input_idx, literal): raise ScheduleParseError( - f'`{{"pack" = {param}}}` on {context}: input_idx should be an integer.' + f'`{{"pack" = {param}}}` on {declaration}: input_idx should be an integer.' ) if not isinstance(mtype, str | None): raise ScheduleParseError( - f'`{{"pack" = {param}}}` on {context}: mtype should be a string or None.' + f'`{{"pack" = {param}}}` on {declaration}: mtype should be a string or None.' ) if not isinstance(pad, bool | str): raise ScheduleParseError( - f'`{{"pack" = {param}}}` on {context}: pad should be a boolean.' + f'`{{"pack" = {param}}}` on {declaration}: pad should be a boolean.' ) # Convert "default" to None for mtype diff --git a/tests/filecheck/search/test_matmul_descript_yaml_goto.py b/tests/filecheck/search/test_matmul_descript_yaml_goto.py index d1ce888d..041884c7 100644 --- a/tests/filecheck/search/test_matmul_descript_yaml_goto.py +++ b/tests/filecheck/search/test_matmul_descript_yaml_goto.py @@ -42,21 +42,20 @@ - kc * nc <= {nb_words_L3} j: k: - B: pack + B: pack=pack_B pad i: - A: pack + A: pack pad=pad_A j#nc: i#mc: k#kc: unroll=kr i#mr: unroll full j#nr: vectorize full """ -print(spec) strategy = Strategy(graph, spec, partial_tiles=True, partial_unrolls=True, initialize=False) print(sorted(strategy._constraints)) print(sum(1 for _ in strategy.sample(100))) -# CHECK: ['1 + nvr + nvr * mr <= 32', 'kc * mc <= 262144', 'kc * nc <= 9437184', 'kc * nr <= 8192', 'kc <= 1024', 'kr <= kc', 'mc <= 1024', 'mr || {1024, mc}', 'nc <= 1024', 'nr == 16 * nvr', 'nr || {1024, nc}', 'nvr * mr * kr <= 256', 'nvr * mr >= 8'] +# CHECK: ['1 + nvr + nvr * mr <= 32', 'kc * mc <= 262144', 'kc * nc <= 9437184', 'kc * nr <= 8192', 'kc <= 1024', 'kr <= kc', 'mc <= 1024', 'mr || {1024, mc}', 'nc <= 1024', 'nr == 16 * nvr', 'nr || {1024, nc}', 'nvr * mr * kr <= 256', 'nvr * mr >= 8', 'pack_B in {0, 1}', 'pad_A in {0,1}'] #CHECK-NEXT: 100 From 500ab90929ddf91a4661514b92f07017a5a216ff Mon Sep 17 00:00:00 2001 From: Leon Frenot Date: Thu, 9 Jul 2026 10:43:34 +0200 Subject: [PATCH 07/23] Partial and full annotations for descript PRT --- src/xtc/schedules/descript.py | 16 ++++++++++------ .../search/test_matmul_descript_pprprp.py | 6 +++--- 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/src/xtc/schedules/descript.py b/src/xtc/schedules/descript.py index 0f85bfb1..e094fb95 100644 --- a/src/xtc/schedules/descript.py +++ b/src/xtc/schedules/descript.py @@ -367,6 +367,10 @@ def _interpret_prt( parallelize = item.annotations.parallelize unroll = item.annotations.unroll_specified + annotations = Annotations( + partial=item.annotations.partial, full=item.annotations.full + ) + def _parallelize(): if isinstance(parallelize, str): node.parallelize_parameters[loop_name] = parallelize @@ -405,7 +409,7 @@ def _unroll(): for a in self.abstract_p_dims: fresh_a = self._fresh(a) loop_name = self._interpret_tile( - TileDecl(a, fresh_a, Annotations()), + TileDecl(a, fresh_a, annotations), node, interchange, sizes, @@ -423,7 +427,7 @@ def _unroll(): ) for a in self.abstract_r_dims: loop_name = self._interpret_tile( - TileDecl(a, self._fresh(a), Annotations()), + TileDecl(a, self._fresh(a), annotations), node, interchange, sizes, @@ -438,7 +442,7 @@ def _unroll(): case "T": for a in self.abstract_dims: loop_name = self._interpret_tile( - TileDecl(a, self._fresh(a), Annotations()), + TileDecl(a, self._fresh(a), annotations), node, interchange, sizes, @@ -464,7 +468,7 @@ def _unroll(): ) a = self.abstract_p_dims[0] self._interpret_tile( - TileDecl(a, self._fresh(a), Annotations()), + TileDecl(a, self._fresh(a), annotations), node, interchange, sizes, @@ -477,7 +481,7 @@ def _unroll(): _unroll() for a in self.abstract_r_dims: self._interpret_tile( - TileDecl(a, self._fresh(a), Annotations()), + TileDecl(a, self._fresh(a), annotations), node, interchange, sizes, @@ -487,7 +491,7 @@ def _unroll(): _unroll() for a in self.abstract_p_dims[1:]: loop_name = self._interpret_tile( - TileDecl(a, self._fresh(a), Annotations()), + TileDecl(a, self._fresh(a), annotations), node, interchange, sizes, diff --git a/tests/filecheck/search/test_matmul_descript_pprprp.py b/tests/filecheck/search/test_matmul_descript_pprprp.py index d93ef6c6..21e0f32c 100644 --- a/tests/filecheck/search/test_matmul_descript_pprprp.py +++ b/tests/filecheck/search/test_matmul_descript_pprprp.py @@ -12,11 +12,11 @@ spec = """ PPRP: parallelize R: unroll=u_k -P: unroll vectorize""" -strategy = Strategy(graph, spec, initialize=False) +P: unroll vectorize full""" +strategy = Strategy(graph, spec, initialize=False, partial_tiles=True, partial_unrolls=True) print(sorted(strategy._constraints)) print(sum(1 for _ in strategy.sample(100))) -# CHECK: ['prt_i_0 || {21}', 'prt_i_1 || {21, prt_i_0}', 'prt_i_2 || {21, prt_i_0, prt_i_1}', 'prt_i_3 || {21, prt_i_0, prt_i_1, prt_i_2}', 'prt_j_0 || {32}', 'prt_j_1 || {32, prt_j_0}', 'prt_j_2 || {32, prt_j_0, prt_j_1}', 'prt_j_3 || {32, prt_j_0, prt_j_1, prt_j_2}', 'prt_k_0 || {12}', 'prt_k_1 || {12, prt_k_0}', 'u_k_prt_k1 || prt_k_1'] +# CHECK: ['prt_i_0 <= 21', 'prt_i_1 <= prt_i_0', 'prt_i_2 <= prt_i_1', 'prt_i_3 || {21, prt_i_0, prt_i_1, prt_i_2}', 'prt_j_0 <= 32', 'prt_j_1 <= prt_j_0', 'prt_j_2 <= prt_j_1', 'prt_j_3 || {32, prt_j_0, prt_j_1, prt_j_2}', 'prt_k_0 <= 12', 'prt_k_1 <= prt_k_0', 'u_k_prt_k1 <= prt_k_1'] # CHECK-NEXT: 100 From e9b0afedf73786232a890450b137588c437d0ce2 Mon Sep 17 00:00:00 2001 From: Leon Frenot Date: Thu, 9 Jul 2026 16:30:56 +0200 Subject: [PATCH 08/23] Parameters for loop interchange in descript --- src/xtc/schedules/descript.py | 61 +++++++----- src/xtc/schedules/parameter_loop_nest.py | 9 +- src/xtc/schedules/parsing.py | 8 ++ ...test_matmul_descript_extend_interchange.py | 99 +++++++++++++++++++ .../test_matmul_descript_interchange.py | 27 +++++ ...test_matmul_descript_pprprp_interchange.py | 22 +++++ 6 files changed, 203 insertions(+), 23 deletions(-) create mode 100644 tests/filecheck/schedules/test_matmul_descript_extend_interchange.py create mode 100644 tests/filecheck/search/test_matmul_descript_interchange.py create mode 100644 tests/filecheck/search/test_matmul_descript_pprprp_interchange.py diff --git a/src/xtc/schedules/descript.py b/src/xtc/schedules/descript.py index e094fb95..d4410c3c 100644 --- a/src/xtc/schedules/descript.py +++ b/src/xtc/schedules/descript.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: BSD-3-Clause # Copyright (c) 2024-2026 The XTC Project Authors # +from math import factorial from typing import Any from dataclasses import dataclass, field from copy import deepcopy @@ -79,6 +80,7 @@ class ScheduleInterpreter: def __post_init__(self): self._abstract_dim_fresh = {i: -1 for i in self.abstract_dims} + self._abstract_dim_fresh["interchange"] = -1 def _fresh(self, axis: str): self._abstract_dim_fresh[axis] += 1 @@ -102,7 +104,7 @@ def _interpret_spec_into_node( """Interpret a schedule spec into an existing node (mutates node).""" # Track state during interpretation previous_cut: dict[str, literal | None] = {a: 0 for a in self.abstract_dims} - interchange: list[str] = list(head) + node.interchange = list(head) last_split: list[tuple[literal, literal | None]] = [] sizes: dict[str, literal] = {} loop_name: str = "" @@ -122,7 +124,6 @@ def _interpret_spec_into_node( item=item, node=node, root=root, - interchange=interchange, previous_cut=previous_cut, axes=axes, last_split=last_split, @@ -131,13 +132,12 @@ def _interpret_spec_into_node( loop_name = self._interpret_tile( item=item, node=node, - interchange=interchange, sizes=sizes, axes=axes, ) self._apply_annotations(item.annotations, loop_name, sizes, node) elif isinstance(item, AxisDecl): - _loop_name = self._interpret_axis(item, interchange) + _loop_name = self._interpret_axis(item, node) if _loop_name: loop_name = _loop_name self._apply_annotations(item.annotations, loop_name, sizes, node) @@ -145,11 +145,13 @@ def _interpret_spec_into_node( loop_name = self._interpret_prt( item=item, node=node, - interchange=interchange, sizes=sizes, axes=axes, ) + for k, v_d in node.interchange_groups.items(): + node.constraints.append(f"1 <= {k} <= {factorial(len(v_d))}") + # Reaplace the placeholder of the last split with its size if len(last_split) > 0: a0, b0 = last_split[0] @@ -170,14 +172,11 @@ def _interpret_spec_into_node( f"Splitting of {axis} unachieved (stops at {cut})." ) - node.interchange = interchange - def _interpret_split( self, item: SplitDecl, node: ParameterLoopNestNode, root: str, - interchange: list[str], previous_cut: dict[str, literal | None], axes: dict[str, list[literal]], last_split: list[tuple[literal, literal | None]], @@ -215,7 +214,7 @@ def _interpret_split( # Save the cutting points of the new dimensions node.splits[axis_name][new_dim_name] = x - interchange.append(new_dim_name) + node.interchange.append(new_dim_name) inner_size = None if y is None: @@ -243,7 +242,7 @@ def _interpret_split( assert x is not None # Save the cutting points of the new dimensions node.splits[axis_name][new_dim_name] = x - interchange.append(new_dim_name) + node.interchange.append(new_dim_name) if isinstance(z, int) and isinstance(x, int): previous_cut[axis_name] = x + z if not isinstance(y, int): @@ -287,7 +286,6 @@ def _interpret_tile( self, item: TileDecl, node: ParameterLoopNestNode, - interchange: list[str], sizes: dict[str, literal], axes: dict[str, list[literal]], ) -> str: @@ -301,7 +299,18 @@ def _interpret_tile( ) node.tiles[item.axis][loop_name] = item.size sizes[loop_name] = item.size - interchange.append(loop_name) + + inter_group = item.annotations.interchange + if inter_group: + if inter_group in node.interchange_groups: + node.interchange_groups[inter_group][len(node.interchange)] = loop_name + else: + node.interchange_groups[inter_group] = { + len(node.interchange): loop_name + } + node.interchange.append(loop_name) + + list_axis: list[int | str] if item.axis in axes: list_axis = axes[item.axis] else: @@ -336,10 +345,11 @@ def _interpret_tile( def _interpret_axis( self, item: AxisDecl, - interchange: list[str], + node: ParameterLoopNestNode, ) -> str: """Interpret a direct axis reference. Returns the loop name.""" axis_name = item.axis + interchange = node.interchange if axis_name in self.abstract_matrix: return "" self._check_axis_existence(axis_name) @@ -351,14 +361,22 @@ def _interpret_axis( f"Axis {axis_name} is scheduled twice (or more)." ) - interchange.append(axis_name) + inter_group = item.annotations.interchange + if inter_group: + if inter_group in node.interchange_groups: + node.interchange_groups[inter_group][len(node.interchange)] = axis_name + else: + node.interchange_groups[inter_group] = { + len(node.interchange): axis_name + } + node.interchange.append(axis_name) + return axis_name def _interpret_prt( self, item: PRTDecl, node: ParameterLoopNestNode, - interchange: list[str], sizes: dict[str, literal], axes: dict[str, list[literal]], ) -> str: @@ -367,8 +385,13 @@ def _interpret_prt( parallelize = item.annotations.parallelize unroll = item.annotations.unroll_specified + interchange = item.annotations.interchange + if interchange == "interchange": + interchange = self._fresh("interchange") annotations = Annotations( - partial=item.annotations.partial, full=item.annotations.full + partial=item.annotations.partial, + full=item.annotations.full, + interchange=interchange, ) def _parallelize(): @@ -411,7 +434,6 @@ def _unroll(): loop_name = self._interpret_tile( TileDecl(a, fresh_a, annotations), node, - interchange, sizes, axes, ) @@ -429,7 +451,6 @@ def _unroll(): loop_name = self._interpret_tile( TileDecl(a, self._fresh(a), annotations), node, - interchange, sizes, axes, ) @@ -444,7 +465,6 @@ def _unroll(): loop_name = self._interpret_tile( TileDecl(a, self._fresh(a), annotations), node, - interchange, sizes, axes, ) @@ -470,7 +490,6 @@ def _unroll(): self._interpret_tile( TileDecl(a, self._fresh(a), annotations), node, - interchange, sizes, axes, ) @@ -483,7 +502,6 @@ def _unroll(): self._interpret_tile( TileDecl(a, self._fresh(a), annotations), node, - interchange, sizes, axes, ) @@ -493,7 +511,6 @@ def _unroll(): loop_name = self._interpret_tile( TileDecl(a, self._fresh(a), annotations), node, - interchange, sizes, axes, ) diff --git a/src/xtc/schedules/parameter_loop_nest.py b/src/xtc/schedules/parameter_loop_nest.py index b1b76ca7..7a2293cc 100644 --- a/src/xtc/schedules/parameter_loop_nest.py +++ b/src/xtc/schedules/parameter_loop_nest.py @@ -3,6 +3,7 @@ # Copyright (c) 2024-2026 The XTC Project Authors # from __future__ import annotations +import itertools from dataclasses import dataclass, field from typing import Generic, TypeVar, Any @@ -114,6 +115,7 @@ class ParameterLoopNestNode(Node["ParameterLoopNestNode"]): tiles: dict[str, dict[str, literal]] splits: dict[str, dict[str, literal]] = field(default_factory=dict) interchange: list[str] = field(default_factory=list) + interchange_groups: dict[str, dict[int, str]] = field(default_factory=dict) vectorize: list[str] = field(default_factory=list) vectorize_parameters: dict[str, str] = field(default_factory=dict) parallelize: list[str] = field(default_factory=list) @@ -147,7 +149,12 @@ def apply_sample(self, sample: dict[str, int]) -> LoopNestNode: for a, v_a in self.splits.items() } - interchange = self.interchange + interchange = self.interchange.copy() + for a, v_d in self.interchange_groups.items(): + if a in sample: + v_ = iter(list(itertools.permutations(v_d.values()))[sample[a] - 1]) + for i in v_d.keys(): + interchange[i] = next(v_) vectorize = self.vectorize for a, v in self.vectorize_parameters.items(): diff --git a/src/xtc/schedules/parsing.py b/src/xtc/schedules/parsing.py index 773bccbd..4f26105e 100644 --- a/src/xtc/schedules/parsing.py +++ b/src/xtc/schedules/parsing.py @@ -72,6 +72,7 @@ class Annotations: buffer_specified: bool | str = False pack: tuple[literal, str | None, bool | str] | None = None pack_specified: bool | str = False + interchange: str = "" partial: bool = False full: bool = False @@ -224,6 +225,7 @@ def _parse_annotations( buffer_specified: bool | str = False pack: tuple[literal, str | None, bool | str] | None = None pack_specified: bool | str = False + interchange: str = "" partial = False full = False @@ -279,6 +281,11 @@ def _parse_annotations( ) declaration_, mtype_, pad_ = pack pack = (declaration_, mtype_, param) + case "interchange": + if param is None: + interchange = "interchange" + else: + interchange = param case "partial": partial = True case "full": @@ -302,6 +309,7 @@ def _parse_annotations( buffer_specified=buffer_specified, pack=pack, pack_specified=pack_specified, + interchange=interchange, partial=partial, full=full, ) diff --git a/tests/filecheck/schedules/test_matmul_descript_extend_interchange.py b/tests/filecheck/schedules/test_matmul_descript_extend_interchange.py new file mode 100644 index 00000000..83370521 --- /dev/null +++ b/tests/filecheck/schedules/test_matmul_descript_extend_interchange.py @@ -0,0 +1,99 @@ +# RUN: python -O %s 2>&1 | filecheck %s +# REQUIRES: module_tvm +# REQUIRES: module_xvs + +import xtc.graphs.xtc.op as O +from xtc.backends.tvm import Backend +from xtc.search.strategies import Strategy_Descript as Strategy + +I, J, K, dtype = 4, 32, 512, "float32" +a = O.tensor((I, K), dtype, name="A") +b = O.tensor((K, J), dtype, name="B") + +with O.graph(name="matmul") as gb: + O.matmul(a, b, name="C") + +graph = gb.graph +print(graph) + +impl = Backend(graph, always_vectorize=False, no_alias=True) + +sch = impl.get_scheduler() + +spec = """ +k: interchange +i: interchange +j: interchange +i#2: +j#16: +""" + +sample = {"interchange": 2} + +strategy = Strategy(graph, spec) + +strategy.generate(sch, sample) + +sched = sch.schedule() + +comp = impl.get_compiler( + shared_lib=True, + dump_file="matmul_descript_extend_interchange", + print_source_ir=True, + print_transformed_ir=True, +) +module = comp.compile(sched) +executor = module.get_executor(validate=True) +res = executor.execute() +print(f"CODE: {res}") +#CHECK: graph: +#CHECK-NEXT: name: matmul +#CHECK-NEXT: inputs: +#CHECK-NEXT: - %0 : 4x512xfloat32 +#CHECK-NEXT: - %1 : 512x32xfloat32 +#CHECK-NEXT: outputs: +#CHECK-NEXT: - %2 : 4x32xfloat32 +#CHECK-NEXT: nodes: +#CHECK-NEXT: - %2: matmul(%0, %1) {name = 'C'} : [4x512xfloat32, 512x32xfloat32] -> [4x32xfloat32] +#CHECK-EMPTY: +#CHECK-NEXT: # from tvm.script import ir as I +#CHECK-NEXT: # from tvm.script import tir as T +#CHECK-EMPTY: +#CHECK-NEXT: @I.ir_module +#CHECK-NEXT: class Module: +#CHECK-NEXT: @T.prim_func +#CHECK-NEXT: def main(_0: T.Buffer((4, 512), "float32"), _1: T.Buffer((512, 32), "float32"), C: T.Buffer((4, 32), "float32")): +#CHECK-NEXT: T.func_attr({"from_legacy_te_schedule": T.bool(True), "tir.noalias": T.bool(True)}) +#CHECK-NEXT: for i, j in T.grid(4, 32): +#CHECK-NEXT: C_1 = T.Buffer((128,), data=C.data) +#CHECK-NEXT: C_1[i * 32 + j] = T.float32(0.0) +#CHECK-NEXT: for k in range(512): +#CHECK-NEXT: cse_var_1: T.int32 = i * 32 + j +#CHECK-NEXT: _0_1 = T.Buffer((2048,), data=_0.data) +#CHECK-NEXT: _1_1 = T.Buffer((16384,), data=_1.data) +#CHECK-NEXT: C_1[cse_var_1] = C_1[cse_var_1] + _0_1[i * 512 + k] * _1_1[k * 32 + j] +#CHECK-NEXT: O = obj['C'] +#CHECK-NEXT: i, j, = O.op.axis +#CHECK-NEXT: k, = O.op.reduce_axis +#CHECK-NEXT: i, i0 = sch[O].split(i, factor=2) +#CHECK-NEXT: j, j0 = sch[O].split(j, factor=16) +#CHECK-NEXT: sch[O].reorder(k, j, i, i0, j0) +#CHECK-EMPTY: +#CHECK-NEXT: # from tvm.script import ir as I +#CHECK-NEXT: # from tvm.script import tir as T +#CHECK-EMPTY: +#CHECK-NEXT: @I.ir_module +#CHECK-NEXT: class Module: +#CHECK-NEXT: @T.prim_func +#CHECK-NEXT: def main(_0: T.Buffer((4, 512), "float32"), _1: T.Buffer((512, 32), "float32"), C: T.Buffer((4, 32), "float32")): +#CHECK-NEXT: T.func_attr({"from_legacy_te_schedule": T.bool(True), "tir.noalias": T.bool(True)}) +#CHECK-NEXT: C_1 = T.Buffer((128,), data=C.data) +#CHECK-NEXT: for j_outer_init, i_outer_init, i_inner_init, j_inner_init in T.grid(2, 2, 2, 16): +#CHECK-NEXT: C_1[i_outer_init * 64 + i_inner_init * 32 + j_outer_init * 16 + j_inner_init] = T.float32(0.0) +#CHECK-NEXT: for k, j_outer, i_outer, i_inner, j_inner in T.grid(512, 2, 2, 2, 16): +#CHECK-NEXT: cse_var_2: T.int32 = j_outer * 16 +#CHECK-NEXT: cse_var_1: T.int32 = i_outer * 64 + i_inner * 32 + cse_var_2 + j_inner +#CHECK-NEXT: _0_1 = T.Buffer((2048,), data=_0.data) +#CHECK-NEXT: _1_1 = T.Buffer((16384,), data=_1.data) +#CHECK-NEXT: C_1[cse_var_1] = C_1[cse_var_1] + _0_1[i_outer * 1024 + i_inner * 512 + k] * _1_1[k * 32 + cse_var_2 + j_inner] +#CHECK-NEXT: CODE: 0 diff --git a/tests/filecheck/search/test_matmul_descript_interchange.py b/tests/filecheck/search/test_matmul_descript_interchange.py new file mode 100644 index 00000000..c240538d --- /dev/null +++ b/tests/filecheck/search/test_matmul_descript_interchange.py @@ -0,0 +1,27 @@ +# RUN: python %s 2>&1 | filecheck %s +# REQUIRES: module_xvs +""" +Test simple schedule on matmul +""" + +import utils +from xtc.search.strategies import Strategy_Descript as Strategy + +graph = utils.get_graph_matmul() +backend = utils.get_backend(graph) +spec = { + "k": {}, + "i": {}, + "j": {}, + "i#i1": {"interchange": "int"}, + "j#j1": {"interchange": "int"}, + "j#j2": {} +} + +strategy = Strategy(graph, spec, initialize=False) + +print(sorted(strategy._constraints)) +print(sum(1 for _ in strategy.sample(100))) + +# CHECK: ['1 <= int <= 2', 'i1 || {21}', 'j1 || {32}', 'j2 || {32, j1}'] +# CHECK-NEXT: 100 diff --git a/tests/filecheck/search/test_matmul_descript_pprprp_interchange.py b/tests/filecheck/search/test_matmul_descript_pprprp_interchange.py new file mode 100644 index 00000000..3e03e728 --- /dev/null +++ b/tests/filecheck/search/test_matmul_descript_pprprp_interchange.py @@ -0,0 +1,22 @@ +# RUN: python %s 2>&1 | filecheck %s +# REQUIRES: module_xvs +""" +Test strategy PPRPRP with interchange on matmul +""" + +import utils +from xtc.search.strategies import Strategy_Descript as Strategy + +graph = utils.get_graph_matmul() +backend = utils.get_backend(graph) +spec = """ +- PP: +- RP: interchange +- RP:""" +strategy = Strategy(graph, spec, initialize=False,) + +print(sorted(strategy._constraints)) +print(sum(1 for _ in strategy.sample(100))) + +# CHECK: ['1 <= prt_interchange_0 <= 6', 'prt_i_0 || {21}', 'prt_i_1 || {21, prt_i_0}', 'prt_i_2 || {21, prt_i_0, prt_i_1}', 'prt_i_3 || {21, prt_i_0, prt_i_1, prt_i_2}', 'prt_j_0 || {32}', 'prt_j_1 || {32, prt_j_0}', 'prt_j_2 || {32, prt_j_0, prt_j_1}', 'prt_j_3 || {32, prt_j_0, prt_j_1, prt_j_2}', 'prt_k_0 || {12}', 'prt_k_1 || {12, prt_k_0}'] +# CHECK-NEXT: 100 From 987a3e4b3a463e086dee2e37816d850db220680b Mon Sep 17 00:00:00 2001 From: Leon Frenot Date: Fri, 10 Jul 2026 16:56:06 +0200 Subject: [PATCH 09/23] U tile in descript PRT --- src/xtc/schedules/descript.py | 26 +++++++++++++++++++------- 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/src/xtc/schedules/descript.py b/src/xtc/schedules/descript.py index d4410c3c..627e4364 100644 --- a/src/xtc/schedules/descript.py +++ b/src/xtc/schedules/descript.py @@ -78,12 +78,11 @@ class ScheduleInterpreter: partial_unrolls: bool = False _abstract_dim_fresh: dict[str, int] = field(default_factory=dict) - def __post_init__(self): - self._abstract_dim_fresh = {i: -1 for i in self.abstract_dims} - self._abstract_dim_fresh["interchange"] = -1 - def _fresh(self, axis: str): - self._abstract_dim_fresh[axis] += 1 + if axis not in self._abstract_dim_fresh: + self._abstract_dim_fresh[axis] = 0 + else: + self._abstract_dim_fresh[axis] += 1 return f"prt_{axis}_{self._abstract_dim_fresh[axis]}" def interpret(self, spec: ScheduleSpec, root: str) -> ParameterLoopNest: @@ -474,9 +473,22 @@ def _unroll(): if unroll and idx + 1 == len(item.shape): _unroll() case "U": - raise ScheduleInterpretError( - "TODO: Ansor-style U tile unimplemented (random order)" + annotations_u = Annotations( + partial=item.annotations.partial, + full=item.annotations.full, + interchange=self._fresh("interchange_u"), ) + for a in self.abstract_dims: + loop_name = self._interpret_tile( + TileDecl(a, self._fresh(a), annotations_u), + node, + sizes, + axes, + ) + if parallelize: + raise ScheduleInterpretError("Cannot parallelize a U tile.") + if unroll and idx + 1 == len(item.shape): + _unroll() case "O": if not self.abstract_p_dims: raise ScheduleInterpretError( From 52ea84f64c009fac0f917f07f7c715b285e7425a Mon Sep 17 00:00:00 2001 From: Leon Frenot Date: Wed, 15 Jul 2026 10:53:59 +0200 Subject: [PATCH 10/23] Check that each axis is used at most once per interchange --- src/xtc/schedules/descript.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/xtc/schedules/descript.py b/src/xtc/schedules/descript.py index 627e4364..96c48af7 100644 --- a/src/xtc/schedules/descript.py +++ b/src/xtc/schedules/descript.py @@ -149,6 +149,13 @@ def _interpret_spec_into_node( ) for k, v_d in node.interchange_groups.items(): + read: set[str] = set() + for k_, v_d_ in v_d.items(): + for a in self.abstract_dims: + if a in v_d_: + if a in read: + raise ScheduleInterpretError(f"Axis {a} is used twice in interchange {k}.") + read.add(a) node.constraints.append(f"1 <= {k} <= {factorial(len(v_d))}") # Reaplace the placeholder of the last split with its size From f88436df5f5e72acaa75be1fbbd2c21a1c090b9a Mon Sep 17 00:00:00 2001 From: Leon Frenot Date: Wed, 15 Jul 2026 15:03:50 +0200 Subject: [PATCH 11/23] Support for footprint in descript constraints --- src/xtc/schedules/descript.py | 65 +++++++++++++++---- src/xtc/schedules/parameter_loop_nest.py | 1 + src/xtc/schedules/parsing.py | 9 +++ src/xtc/search/strategies.py | 38 +++++++++++ .../search/test_matmul_descript_yaml_goto.py | 14 ++-- 5 files changed, 108 insertions(+), 19 deletions(-) diff --git a/src/xtc/schedules/descript.py b/src/xtc/schedules/descript.py index 96c48af7..e874d100 100644 --- a/src/xtc/schedules/descript.py +++ b/src/xtc/schedules/descript.py @@ -77,6 +77,7 @@ class ScheduleInterpreter: partial_tiles: bool = False partial_unrolls: bool = False _abstract_dim_fresh: dict[str, int] = field(default_factory=dict) + _levels: dict[str, dict[str, literal | None]] = field(default_factory=dict) def _fresh(self, axis: str): if axis not in self._abstract_dim_fresh: @@ -90,6 +91,11 @@ def interpret(self, spec: ScheduleSpec, root: str) -> ParameterLoopNest: loop_nest = ParameterLoopNest(abstract_dims=self.abstract_dims) root_node = loop_nest.build_root_node(root) self._interpret_spec_into_node(spec, root_node, root, head=[]) + levels = { + k: {k_: v_ if v_ is not None else 1 for k_, v_ in level.items()} + for k, level in self._levels.items() + } + root_node._levels = levels return loop_nest def _interpret_spec_into_node( @@ -154,7 +160,9 @@ def _interpret_spec_into_node( for a in self.abstract_dims: if a in v_d_: if a in read: - raise ScheduleInterpretError(f"Axis {a} is used twice in interchange {k}.") + raise ScheduleInterpretError( + f"Axis {a} is used twice in interchange {k}." + ) read.add(a) node.constraints.append(f"1 <= {k} <= {factorial(len(v_d))}") @@ -279,6 +287,9 @@ def _interpret_split( ) node.add_child(child_node) + if current_size is not None: + self._update_levels(item.axis, current_size) + # Recursively interpret the nested schedule into the child node self._interpret_spec_into_node( spec=item.body, @@ -346,6 +357,7 @@ def _interpret_tile( node.constraints.append(s) list_axis.append(item.size) + self._update_levels(item.axis, axes[item.axis][-1]) return loop_name def _interpret_axis( @@ -377,6 +389,8 @@ def _interpret_axis( } node.interchange.append(axis_name) + if self.abstract_dim_sizes: + self._update_levels(axis_name, self.abstract_dim_sizes[axis_name]) return axis_name def _interpret_prt( @@ -390,15 +404,25 @@ def _interpret_prt( parallelize = item.annotations.parallelize unroll = item.annotations.unroll_specified + level = item.annotations.level interchange = item.annotations.interchange if interchange == "interchange": interchange = self._fresh("interchange") - annotations = Annotations( - partial=item.annotations.partial, - full=item.annotations.full, - interchange=interchange, - ) + + def _annotations(axis: str): + if not level: + return Annotations( + partial=item.annotations.partial, + full=item.annotations.full, + interchange=interchange, + ) + return Annotations( + partial=item.annotations.partial, + full=item.annotations.full, + interchange=interchange, + level=level + axis, + ) def _parallelize(): if isinstance(parallelize, str): @@ -428,6 +452,10 @@ def _unroll(): node.constraints.append(f"{unroll_factor} || {sizes[loop_name]}") node.unroll[loop_name] = unroll_factor + if level and len(item.shape) > 1: + raise ScheduleInterpretError( + "Cannot use the same level on mulitple PRT tiles." + ) for idx, t in enumerate(item.shape): match t: case "P": @@ -438,7 +466,7 @@ def _unroll(): for a in self.abstract_p_dims: fresh_a = self._fresh(a) loop_name = self._interpret_tile( - TileDecl(a, fresh_a, annotations), + TileDecl(a, fresh_a, _annotations(a)), node, sizes, axes, @@ -455,7 +483,7 @@ def _unroll(): ) for a in self.abstract_r_dims: loop_name = self._interpret_tile( - TileDecl(a, self._fresh(a), annotations), + TileDecl(a, self._fresh(a), _annotations(a)), node, sizes, axes, @@ -469,7 +497,7 @@ def _unroll(): case "T": for a in self.abstract_dims: loop_name = self._interpret_tile( - TileDecl(a, self._fresh(a), annotations), + TileDecl(a, self._fresh(a), _annotations(a)), node, sizes, axes, @@ -480,6 +508,8 @@ def _unroll(): if unroll and idx + 1 == len(item.shape): _unroll() case "U": + if level: + raise ScheduleInterpretError("Cannot use level on a U tile.") annotations_u = Annotations( partial=item.annotations.partial, full=item.annotations.full, @@ -507,7 +537,7 @@ def _unroll(): ) a = self.abstract_p_dims[0] self._interpret_tile( - TileDecl(a, self._fresh(a), annotations), + TileDecl(a, self._fresh(a), _annotations(a)), node, sizes, axes, @@ -519,7 +549,7 @@ def _unroll(): _unroll() for a in self.abstract_r_dims: self._interpret_tile( - TileDecl(a, self._fresh(a), annotations), + TileDecl(a, self._fresh(a), _annotations(a)), node, sizes, axes, @@ -528,7 +558,7 @@ def _unroll(): _unroll() for a in self.abstract_p_dims[1:]: loop_name = self._interpret_tile( - TileDecl(a, self._fresh(a), annotations), + TileDecl(a, self._fresh(a), _annotations(a)), node, sizes, axes, @@ -612,6 +642,12 @@ def _apply_annotations( else: node.parallelize.append(loop_name) + if annotations.level: + level = annotations.level + if level in self._levels: + raise ScheduleInterpretError(f"Level {level} used multiple times.") + self._levels[level] = {k: None for k in self.abstract_dims} + if annotations.buffer_specified: node.buffer_at[loop_name] = annotations.buffer @@ -643,6 +679,11 @@ def _apply_annotations( if isinstance(pad, str): node.constraints.append(f"{pad} in {{0,1}}") + def _update_levels(self, axis: str, size: literal): + for name, level in self._levels.items(): + if level[axis] is None: + level[axis] = size + def _check_splitting_intervals( self, item: SplitDecl, diff --git a/src/xtc/schedules/parameter_loop_nest.py b/src/xtc/schedules/parameter_loop_nest.py index 7a2293cc..f37aa5c5 100644 --- a/src/xtc/schedules/parameter_loop_nest.py +++ b/src/xtc/schedules/parameter_loop_nest.py @@ -126,6 +126,7 @@ class ParameterLoopNestNode(Node["ParameterLoopNestNode"]): pack_at: dict[str, tuple[int, str | None, bool | str]] = field(default_factory=dict) pack_parameters: dict[str, str] = field(default_factory=dict) constraints: list[str] = field(default_factory=list) + _levels: dict[str, dict[str, literal]] = field(default_factory=dict) def apply_sample(self, sample: dict[str, int]) -> LoopNestNode: """ diff --git a/src/xtc/schedules/parsing.py b/src/xtc/schedules/parsing.py index 4f26105e..750d7aec 100644 --- a/src/xtc/schedules/parsing.py +++ b/src/xtc/schedules/parsing.py @@ -73,6 +73,7 @@ class Annotations: pack: tuple[literal, str | None, bool | str] | None = None pack_specified: bool | str = False interchange: str = "" + level: str = "" partial: bool = False full: bool = False @@ -226,6 +227,7 @@ def _parse_annotations( pack: tuple[literal, str | None, bool | str] | None = None pack_specified: bool | str = False interchange: str = "" + level: str = "" partial = False full = False @@ -286,6 +288,12 @@ def _parse_annotations( interchange = "interchange" else: interchange = param + case "level": + if param is None: + raise ScheduleParseError( + f"Level annotation without name on {declaration}: {key}" + ) + level = param case "partial": partial = True case "full": @@ -310,6 +318,7 @@ def _parse_annotations( pack=pack, pack_specified=pack_specified, interchange=interchange, + level=level, partial=partial, full=full, ) diff --git a/src/xtc/search/strategies.py b/src/xtc/search/strategies.py index 2e934648..8490e451 100644 --- a/src/xtc/search/strategies.py +++ b/src/xtc/search/strategies.py @@ -9,6 +9,7 @@ from collections.abc import Sequence, Mapping, Iterator, Generator import itertools import numpy as np +import re from xtc.itf.graph import Graph from xtc.itf.schd import Scheduler @@ -1029,6 +1030,8 @@ def _filter(self, samples: Iterator[VecSample]) -> Iterator[VecSample]: from xvs.execute import execute_static, execute_dynamic class Strategy_Descript(Strategy): + _FP_PATTERN = re.compile(r"footprint\((\w*),\s*(\w*)\)") + def __init__( self, graph: Graph, @@ -1044,6 +1047,21 @@ def __init__( self._op = graph.outputs_nodes[0].operation self._stats: dict[str, int] = {} self._axes = axes if axes else list(self._op.dims) + accesses: tuple[tuple[str, ...], ...] = ( + self._op.accesses_maps[1] + self._op.accesses_maps[2] + ) + if axes: + if len(axes) != len(self._op.dims): + raise ValueError( + f"Number of user axes: {axes} is different from operation's axes: {self._op.dims}" + ) + self._axes = axes + access_swap = { + list(self._op.dims)[i]: axes[i] for i in range(len(axes)) + } + accesses = tuple((tuple((access_swap[x] for x in d)) for d in accesses)) + else: + self._axes = list(self._op.dims) self._sizes = self._constant_sizes() self._sample_names: list[str] = [] tensors = [ @@ -1070,10 +1088,30 @@ def __init__( self._constraints = ( constraints + input_constraints if constraints else input_constraints ) + self._footprints(tensors, accesses) self._initialized = False if initialize: self._initialize() + def _footprints( + self, tensors: list[str], accesses: tuple[tuple[str, ...], ...] + ): + if not self._constraints: + return + if not self._loop_nest.root_node: + return + levels = self._loop_nest.root_node._levels + for i, constraint in enumerate(self._constraints): + match = self._FP_PATTERN.match(constraint) + if match: + tensor, level = match.groups() + tensor = tensors.index(tensor) + if level not in levels: + raise ValueError(f"Level {level} is not defined in the spec.") + level = levels[level] + footprint = " * ".join(str(level[x]) for x in accesses[tensor]) + self._constraints[i] = self._FP_PATTERN.sub(footprint, constraint) + def _constant_sizes(self) -> Mapping[str, int]: sizes = {a: v for a, v in self._op.dims.items() if isinstance(v, int)} return sizes diff --git a/tests/filecheck/search/test_matmul_descript_yaml_goto.py b/tests/filecheck/search/test_matmul_descript_yaml_goto.py index 041884c7..591568ce 100644 --- a/tests/filecheck/search/test_matmul_descript_yaml_goto.py +++ b/tests/filecheck/search/test_matmul_descript_yaml_goto.py @@ -37,16 +37,16 @@ - nr == {vector_size} * nvr - nvr * mr >= {ilp} - nvr * mr * kr <= {reorder_buffer} - - kc * nr <= {nb_words_L1} - - kc * mc <= {nb_words_L2} - - kc * nc <= {nb_words_L3} + - footprint(B, L1) <= {nb_words_L1} + - footprint(A, L2) <= {nb_words_L2} + - footprint(B, L3) <= {nb_words_L3} j: k: B: pack=pack_B pad - i: + i: level=L3 A: pack pad=pad_A - j#nc: - i#mc: + j#nc: level=L2 + i#mc: level=L1 k#kc: unroll=kr i#mr: unroll full j#nr: vectorize full @@ -57,5 +57,5 @@ print(sorted(strategy._constraints)) print(sum(1 for _ in strategy.sample(100))) -# CHECK: ['1 + nvr + nvr * mr <= 32', 'kc * mc <= 262144', 'kc * nc <= 9437184', 'kc * nr <= 8192', 'kc <= 1024', 'kr <= kc', 'mc <= 1024', 'mr || {1024, mc}', 'nc <= 1024', 'nr == 16 * nvr', 'nr || {1024, nc}', 'nvr * mr * kr <= 256', 'nvr * mr >= 8', 'pack_B in {0, 1}', 'pad_A in {0,1}'] +# CHECK: ['1 + nvr + nvr * mr <= 32', 'kc * nc <= 9437184', 'kc * nr <= 8192', 'kc <= 1024', 'kr <= kc', 'mc * kc <= 262144', 'mc <= 1024', 'mr || {1024, mc}', 'nc <= 1024', 'nr == 16 * nvr', 'nr || {1024, nc}', 'nvr * mr * kr <= 256', 'nvr * mr >= 8', 'pack_B in {0, 1}', 'pad_A in {0,1}'] #CHECK-NEXT: 100 From 00c408a856f006156ec86adb2554032774e20cd7 Mon Sep 17 00:00:00 2001 From: Leon Frenot Date: Wed, 15 Jul 2026 15:29:04 +0200 Subject: [PATCH 12/23] Fix PRT tiles if first occurence of axis --- src/xtc/schedules/descript.py | 64 ++++------ .../test_matmul_descript_extend_prp.py | 112 ++++++++++++++++++ .../search/test_matmul_descript_pprprp.py | 2 +- ...test_matmul_descript_pprprp_interchange.py | 2 +- 4 files changed, 134 insertions(+), 46 deletions(-) create mode 100644 tests/filecheck/schedules/test_matmul_descript_extend_prp.py diff --git a/src/xtc/schedules/descript.py b/src/xtc/schedules/descript.py index e874d100..deca4181 100644 --- a/src/xtc/schedules/descript.py +++ b/src/xtc/schedules/descript.py @@ -410,7 +410,7 @@ def _interpret_prt( if interchange == "interchange": interchange = self._fresh("interchange") - def _annotations(axis: str): + def _annotations(axis: str) -> Annotations: if not level: return Annotations( partial=item.annotations.partial, @@ -452,6 +452,18 @@ def _unroll(): node.constraints.append(f"{unroll_factor} || {sizes[loop_name]}") node.unroll[loop_name] = unroll_factor + def _interpret(axis: str, annotations: Annotations | None = None): + if axis not in node.interchange: + if not annotations: + annotations = Annotations() + return self._interpret_axis(AxisDecl(axis, annotations), node) + else: + if not annotations: + annotations = _annotations(a) + return self._interpret_tile( + TileDecl(a, self._fresh(a), annotations), node, sizes, axes + ) + if level and len(item.shape) > 1: raise ScheduleInterpretError( "Cannot use the same level on mulitple PRT tiles." @@ -464,13 +476,7 @@ def _unroll(): "P scheme used, but P axes are not specified" ) for a in self.abstract_p_dims: - fresh_a = self._fresh(a) - loop_name = self._interpret_tile( - TileDecl(a, fresh_a, _annotations(a)), - node, - sizes, - axes, - ) + loop_name = _interpret(a) if parallelize: _parallelize() parallelize = False @@ -482,12 +488,7 @@ def _unroll(): "R scheme used, but R axes are not specified" ) for a in self.abstract_r_dims: - loop_name = self._interpret_tile( - TileDecl(a, self._fresh(a), _annotations(a)), - node, - sizes, - axes, - ) + loop_name = _interpret(a) if parallelize: raise ScheduleInterpretError( "Cannot parallelize a reduction axis" @@ -496,12 +497,7 @@ def _unroll(): _unroll() case "T": for a in self.abstract_dims: - loop_name = self._interpret_tile( - TileDecl(a, self._fresh(a), _annotations(a)), - node, - sizes, - axes, - ) + loop_name = _interpret(a) if parallelize: _parallelize() parallelize = False @@ -516,12 +512,7 @@ def _unroll(): interchange=self._fresh("interchange_u"), ) for a in self.abstract_dims: - loop_name = self._interpret_tile( - TileDecl(a, self._fresh(a), annotations_u), - node, - sizes, - axes, - ) + loop_name = _interpret(a, annotations_u) if parallelize: raise ScheduleInterpretError("Cannot parallelize a U tile.") if unroll and idx + 1 == len(item.shape): @@ -536,33 +527,18 @@ def _unroll(): "O scheme used, but R axes are not specified" ) a = self.abstract_p_dims[0] - self._interpret_tile( - TileDecl(a, self._fresh(a), _annotations(a)), - node, - sizes, - axes, - ) + _interpret(a) if parallelize: _parallelize() parallelize = False if unroll and idx + 1 == len(item.shape): _unroll() for a in self.abstract_r_dims: - self._interpret_tile( - TileDecl(a, self._fresh(a), _annotations(a)), - node, - sizes, - axes, - ) + _interpret(a) if unroll and idx + 1 == len(item.shape): _unroll() for a in self.abstract_p_dims[1:]: - loop_name = self._interpret_tile( - TileDecl(a, self._fresh(a), _annotations(a)), - node, - sizes, - axes, - ) + loop_name = _interpret(a) if unroll and idx + 1 == len(item.shape): _unroll() case "W": diff --git a/tests/filecheck/schedules/test_matmul_descript_extend_prp.py b/tests/filecheck/schedules/test_matmul_descript_extend_prp.py new file mode 100644 index 00000000..9894632d --- /dev/null +++ b/tests/filecheck/schedules/test_matmul_descript_extend_prp.py @@ -0,0 +1,112 @@ +# RUN: python -O %s 2>&1 | filecheck %s +# REQUIRES: module_tvm +# REQUIRES: module_xvs + +import xtc.graphs.xtc.op as O +from xtc.backends.tvm import Backend +from xtc.search.strategies import Strategy_Descript as Strategy + +I, J, K, dtype = 4, 32, 512, "float32" +a = O.tensor((I, K), dtype, name="A") +b = O.tensor((K, J), dtype, name="B") + +with O.graph(name="matmul") as gb: + O.matmul(a, b, name="C") + +graph = gb.graph +print(graph) + +impl = Backend(graph, always_vectorize=False, no_alias=True) + +sch = impl.get_scheduler() + +spec = """ +- P: parallelize +- R: unroll=u_k +- P: unroll vectorize""" + +strategy = Strategy(graph, spec) + +sample = {'prt_i_0': 1, 'u_k_prt_k': 8, 'prt_j_0': 2} + +strategy.generate(sch, sample) + +sched = sch.schedule() + +comp = impl.get_compiler( + shared_lib=True, + dump_file="matmul_descript_extend_prp", + print_source_ir=True, + print_transformed_ir=True, +) +module = comp.compile(sched) +executor = module.get_executor(validate=True) +res = executor.execute() +print(f"CODE: {res}") + +#CHECK:graph: +#CHECK-NEXT: name: matmul +#CHECK-NEXT: inputs: +#CHECK-NEXT: - %0 : 4x512xfloat32 +#CHECK-NEXT: - %1 : 512x32xfloat32 +#CHECK-NEXT: outputs: +#CHECK-NEXT: - %2 : 4x32xfloat32 +#CHECK-NEXT: nodes: +#CHECK-NEXT: - %2: matmul(%0, %1) {name = 'C'} : [4x512xfloat32, 512x32xfloat32] -> [4x32xfloat32] +#CHECK-EMPTY: +#CHECK-NEXT:# from tvm.script import ir as I +#CHECK-NEXT:# from tvm.script import tir as T +#CHECK-EMPTY: +#CHECK-NEXT:@I.ir_module +#CHECK-NEXT:class Module: +#CHECK-NEXT: @T.prim_func +#CHECK-NEXT: def main(_0: T.Buffer((4, 512), "float32"), _1: T.Buffer((512, 32), "float32"), C: T.Buffer((4, 32), "float32")): +#CHECK-NEXT: T.func_attr({"from_legacy_te_schedule": T.bool(True), "tir.noalias": T.bool(True)}) +#CHECK-NEXT: for i, j in T.grid(4, 32): +#CHECK-NEXT: C_1 = T.Buffer((128,), data=C.data) +#CHECK-NEXT: C_1[i * 32 + j] = T.float32(0.0) +#CHECK-NEXT: for k in range(512): +#CHECK-NEXT: cse_var_1: T.int32 = i * 32 + j +#CHECK-NEXT: _0_1 = T.Buffer((2048,), data=_0.data) +#CHECK-NEXT: _1_1 = T.Buffer((16384,), data=_1.data) +#CHECK-NEXT: C_1[cse_var_1] = C_1[cse_var_1] + _0_1[i * 512 + k] * _1_1[k * 32 + j] +#CHECK-NEXT:O = obj['C'] +#CHECK-NEXT:i, j, = O.op.axis +#CHECK-NEXT:k, = O.op.reduce_axis +#CHECK-NEXT:k, __u_k = sch[O].split(k, factor=8) +#CHECK-NEXT:i, i0 = sch[O].split(i, factor=1) +#CHECK-NEXT:j, j0 = sch[O].split(j, factor=2) +#CHECK-NEXT:sch[O].reorder(i, j, k, __u_k, i0, j0) +#CHECK-NEXT:sch[O].unroll(__u_k) +#CHECK-NEXT:sch[O].unroll(i0) +#CHECK-NEXT:sch[O].vectorize(j0) +#CHECK-NEXT:sch[O].parallel(i) +#CHECK-EMPTY: +#CHECK-NEXT:# from tvm.script import ir as I +#CHECK-NEXT:# from tvm.script import tir as T +#CHECK-EMPTY: +#CHECK-NEXT:@I.ir_module +#CHECK-NEXT:class Module: +#CHECK-NEXT: @T.prim_func +#CHECK-NEXT: def main(_0: T.Buffer((4, 512), "float32"), _1: T.Buffer((512, 32), "float32"), C: T.Buffer((4, 32), "float32")): +#CHECK-NEXT: T.func_attr({"from_legacy_te_schedule": T.bool(True), "tir.noalias": T.bool(True)}) +#CHECK-NEXT: for i_outer in T.parallel(4): +#CHECK-NEXT: for j_outer in range(16): +#CHECK-NEXT: C_1 = T.Buffer((128,), data=C.data) +#CHECK-NEXT: C_1[i_outer * 32 + j_outer * 2:i_outer * 32 + j_outer * 2 + 2] = T.Broadcast(T.float32(0.0), 2) +#CHECK-NEXT: for k_outer in range(64): +#CHECK-NEXT: cse_var_4: T.int32 = j_outer * 2 +#CHECK-NEXT: cse_var_3: T.int32 = k_outer * 256 + cse_var_4 +#CHECK-NEXT: cse_var_2: T.int32 = i_outer * 512 + k_outer * 8 +#CHECK-NEXT: cse_var_1: T.int32 = i_outer * 32 + cse_var_4 +#CHECK-NEXT: _0_1 = T.Buffer((2048,), data=_0.data) +#CHECK-NEXT: _1_1 = T.Buffer((16384,), data=_1.data) +#CHECK-NEXT: C_1[cse_var_1:cse_var_1 + 2] = C_1[cse_var_1:cse_var_1 + 2] + T.Broadcast(_0_1[cse_var_2], 2) * _1_1[cse_var_3:cse_var_3 + 2] +#CHECK-NEXT: C_1[cse_var_1:cse_var_1 + 2] = C_1[cse_var_1:cse_var_1 + 2] + T.Broadcast(_0_1[cse_var_2 + 1], 2) * _1_1[cse_var_3 + 32:cse_var_3 + 32 + 2] +#CHECK-NEXT: C_1[cse_var_1:cse_var_1 + 2] = C_1[cse_var_1:cse_var_1 + 2] + T.Broadcast(_0_1[cse_var_2 + 2], 2) * _1_1[cse_var_3 + 64:cse_var_3 + 64 + 2] +#CHECK-NEXT: C_1[cse_var_1:cse_var_1 + 2] = C_1[cse_var_1:cse_var_1 + 2] + T.Broadcast(_0_1[cse_var_2 + 3], 2) * _1_1[cse_var_3 + 96:cse_var_3 + 96 + 2] +#CHECK-NEXT: C_1[cse_var_1:cse_var_1 + 2] = C_1[cse_var_1:cse_var_1 + 2] + T.Broadcast(_0_1[cse_var_2 + 4], 2) * _1_1[cse_var_3 + 128:cse_var_3 + 128 + 2] +#CHECK-NEXT: C_1[cse_var_1:cse_var_1 + 2] = C_1[cse_var_1:cse_var_1 + 2] + T.Broadcast(_0_1[cse_var_2 + 5], 2) * _1_1[cse_var_3 + 160:cse_var_3 + 160 + 2] +#CHECK-NEXT: C_1[cse_var_1:cse_var_1 + 2] = C_1[cse_var_1:cse_var_1 + 2] + T.Broadcast(_0_1[cse_var_2 + 6], 2) * _1_1[cse_var_3 + 192:cse_var_3 + 192 + 2] +#CHECK-NEXT: C_1[cse_var_1:cse_var_1 + 2] = C_1[cse_var_1:cse_var_1 + 2] + T.Broadcast(_0_1[cse_var_2 + 7], 2) * _1_1[cse_var_3 + 224:cse_var_3 + 224 + 2] +#CHECK-NEXT:CODE: 0 diff --git a/tests/filecheck/search/test_matmul_descript_pprprp.py b/tests/filecheck/search/test_matmul_descript_pprprp.py index 21e0f32c..cdfa6279 100644 --- a/tests/filecheck/search/test_matmul_descript_pprprp.py +++ b/tests/filecheck/search/test_matmul_descript_pprprp.py @@ -18,5 +18,5 @@ print(sorted(strategy._constraints)) print(sum(1 for _ in strategy.sample(100))) -# CHECK: ['prt_i_0 <= 21', 'prt_i_1 <= prt_i_0', 'prt_i_2 <= prt_i_1', 'prt_i_3 || {21, prt_i_0, prt_i_1, prt_i_2}', 'prt_j_0 <= 32', 'prt_j_1 <= prt_j_0', 'prt_j_2 <= prt_j_1', 'prt_j_3 || {32, prt_j_0, prt_j_1, prt_j_2}', 'prt_k_0 <= 12', 'prt_k_1 <= prt_k_0', 'u_k_prt_k1 <= prt_k_1'] +# CHECK: ['prt_i_0 <= 21', 'prt_i_1 <= prt_i_0', 'prt_i_2 || {21, prt_i_0, prt_i_1}', 'prt_j_0 <= 32', 'prt_j_1 <= prt_j_0', 'prt_j_2 || {32, prt_j_0, prt_j_1}', 'prt_k_0 <= 12', 'u_k_prt_k0 <= prt_k_0'] # CHECK-NEXT: 100 diff --git a/tests/filecheck/search/test_matmul_descript_pprprp_interchange.py b/tests/filecheck/search/test_matmul_descript_pprprp_interchange.py index 3e03e728..f8909e47 100644 --- a/tests/filecheck/search/test_matmul_descript_pprprp_interchange.py +++ b/tests/filecheck/search/test_matmul_descript_pprprp_interchange.py @@ -18,5 +18,5 @@ print(sorted(strategy._constraints)) print(sum(1 for _ in strategy.sample(100))) -# CHECK: ['1 <= prt_interchange_0 <= 6', 'prt_i_0 || {21}', 'prt_i_1 || {21, prt_i_0}', 'prt_i_2 || {21, prt_i_0, prt_i_1}', 'prt_i_3 || {21, prt_i_0, prt_i_1, prt_i_2}', 'prt_j_0 || {32}', 'prt_j_1 || {32, prt_j_0}', 'prt_j_2 || {32, prt_j_0, prt_j_1}', 'prt_j_3 || {32, prt_j_0, prt_j_1, prt_j_2}', 'prt_k_0 || {12}', 'prt_k_1 || {12, prt_k_0}'] +# CHECK: ['1 <= prt_interchange_0 <= 6', 'prt_i_0 || {21}', 'prt_i_1 || {21, prt_i_0}', 'prt_i_2 || {21, prt_i_0, prt_i_1}', 'prt_j_0 || {32}', 'prt_j_1 || {32, prt_j_0}', 'prt_j_2 || {32, prt_j_0, prt_j_1}', 'prt_k_0 || {12}'] # CHECK-NEXT: 100 From b2a8bfcf35c73c22c9ed2674e0ee72d90bdbdece Mon Sep 17 00:00:00 2001 From: Leon Frenot Date: Thu, 16 Jul 2026 11:30:34 +0200 Subject: [PATCH 13/23] Remove useless check functions from parameter loop nest --- src/xtc/schedules/parameter_loop_nest.py | 128 ----------------------- 1 file changed, 128 deletions(-) diff --git a/src/xtc/schedules/parameter_loop_nest.py b/src/xtc/schedules/parameter_loop_nest.py index f37aa5c5..45454117 100644 --- a/src/xtc/schedules/parameter_loop_nest.py +++ b/src/xtc/schedules/parameter_loop_nest.py @@ -12,8 +12,6 @@ NodeT = TypeVar("NodeT", bound="Node[Any]") -from .exceptions import ScheduleValidationError - literal = int | str @@ -453,14 +451,6 @@ def build_root_node(self, root: str) -> ParameterLoopNestNode: self.root_node = node return node - def check(self): - assert self.root_node is not None - info = ParameterLoopInfo.build_from_node(self.root_node) - self._check_use_defined_dims(info) - self._check_vectorization_consistency() - self._check_tiling_consistency(info) - self._check_sizes(info) - def apply_sample(self, sample: dict[str, int]) -> LoopNest: """ Replaces the string parameters with the correspondings values in sample. @@ -475,121 +465,3 @@ def collect_constraints(self) -> list[str]: for node in self.nodes: constraints += node.constraints return constraints - - def _check_use_defined_dims(self, info: ParameterLoopInfo): - for dim in self.abstract_dims: - if dim not in info.dims: - raise ScheduleValidationError(f"{dim} defined but never used") - - def _check_vectorization_consistency(self): - for sched in self.nodes: - vect_above = False - for loop_name in sched.interchange: - if loop_name in sched.vectorize: - vect_above = True - elif vect_above: - raise ScheduleValidationError( - f"Inner loop {loop_name} isn't vectorized but an outer one is." - ) - - def _check_tiling_consistency(self, info: ParameterLoopInfo) -> None: - seen_axes: dict[str, literal | None] = {} - for sched in self.nodes: - for loop_name in sched.interchange: - if loop_name in info.dims: - seen_axes[loop_name] = None - elif loop_name in info.splits_to_axis: - axis = info.splits_to_axis[loop_name] - seen_axes[axis] = sched.splits[axis][loop_name] - elif loop_name in info.tiles_to_axis: - axis = info.tiles_to_axis[loop_name] - size = sched.tiles[axis][loop_name] - if axis not in seen_axes: - raise ScheduleValidationError( - f""" - `{axis}#{size}`: {axis} has not been materialized yet. - """ - ) - seen_axes[axis] = size - - def _check_sizes(self, info: ParameterLoopInfo): - current_size_of_split: dict[str, literal | None] = {} - for sched in self.nodes: - current_size_of_tile: dict[str, literal] = {} - if sched.split_origin is not None: - axis = sched.split_origin.axis - start = sched.split_origin.start - end = sched.split_origin.end - if end is not None and start is not None: - current_size_of_split[axis] = ( - end - start - if isinstance(end, int) and isinstance(start, int) - else f"{end} - {start}" - ) - else: - current_size_of_split[axis] = None - - for loop_name in sched.interchange: - axis = info.loops_to_axis[loop_name] - current_sizes = ( - {d: None for d in info.dims} - | current_size_of_split - | current_size_of_tile - ) - loop_size = None - if loop_name in info.dims: - if loop_name not in current_size_of_split: - current_size_of_split[loop_name] = None - elif loop_name in info.tiles_to_axis: - loop_size = sched.tiles[axis][loop_name] - ParameterLoopNest._must_be_smaller_routine( - new_size=loop_size, - current_sizes=current_sizes, - loop_name=loop_name, - axis=axis, - ) - current_size_of_tile[axis] = loop_size - elif ( - loop_name in info.splits_to_axis - and loop_name in info.splits_to_sizes - ): - loop_size = info.splits_to_sizes[loop_name] - ParameterLoopNest._must_be_smaller_routine( - new_size=loop_size, - current_sizes=current_sizes, - loop_name=loop_name, - axis=axis, - ) - current_size_of_split[axis] = loop_size - - if loop_name in sched.unroll: - unroll_factor = sched.unroll[loop_name] - if ( - loop_size - and isinstance(loop_size, int) - and isinstance(unroll_factor, int) - and loop_size < unroll_factor - ): - raise ScheduleValidationError( - f'`{{"unroll" = {unroll_factor}}}`: unroll factor should be smaller than {loop_size}.' - ) - - @staticmethod - def _must_be_smaller_routine( - new_size: literal, - current_sizes: dict[str, literal | None], - loop_name: str, - axis: str, - ): - old_size = current_sizes[axis] - if ( - old_size is not None - and isinstance(new_size, int) - and isinstance(old_size, int) - and new_size > old_size - ): - raise ScheduleValidationError( - f""" - Inner loop {loop_name} on axis {axis} must be smaller than outer loop. - """ - ) From a1837646bb54bafc5eea83b40f645532abcc8816 Mon Sep 17 00:00:00 2001 From: Leon Frenot Date: Thu, 16 Jul 2026 12:06:37 +0200 Subject: [PATCH 14/23] Missing tests for descript PRT tiles --- .../test_matmul_descript_extend_twuop.py | 208 ++++++++++++++++++ .../search/test_matmul_descript_twuop.py | 21 ++ 2 files changed, 229 insertions(+) create mode 100644 tests/filecheck/schedules/test_matmul_descript_extend_twuop.py create mode 100644 tests/filecheck/search/test_matmul_descript_twuop.py diff --git a/tests/filecheck/schedules/test_matmul_descript_extend_twuop.py b/tests/filecheck/schedules/test_matmul_descript_extend_twuop.py new file mode 100644 index 00000000..b4c64337 --- /dev/null +++ b/tests/filecheck/schedules/test_matmul_descript_extend_twuop.py @@ -0,0 +1,208 @@ +# RUN: python -O %s 2>&1 | filecheck %s +# REQUIRES: module_tvm +# REQUIRES: module_xvs + +import xtc.graphs.xtc.op as O +from xtc.backends.mlir import Backend +from xtc.search.strategies import Strategy_Descript as Strategy + +I, J, K, dtype = 8, 32, 512, "float32" +a = O.tensor((I, K), dtype, name="A") +b = O.tensor((K, J), dtype, name="B") + +with O.graph(name="matmul") as gb: + O.matmul(a, b, name="C") + +graph = gb.graph +print(graph) + +impl = Backend(graph, always_vectorize=False, no_alias=True) + +sch = impl.get_scheduler() + +spec = """ +TWUOP: parallelize unroll vectorize +""" + +strategy = Strategy(graph, spec) + +sample = { + "prt_j_0": 32, + "prt_j_1": 16, + "prt_j_2": 2, + "prt_i_0": 8, + "prt_i_1": 4, + "prt_i_2": 2, + "prt_k_0": 8, + "prt_k_1": 2, + "prt_interchange_u_0": 2, +} + +strategy.generate(sch, sample) + +sched = sch.schedule() + +comp = impl.get_compiler( + shared_lib=True, + dump_file="matmul_descript_extend_twuop", + print_source_ir=True, + print_transformed_ir=True, +) +module = comp.compile(sched) +executor = module.get_executor(validate=True) +res = executor.execute() +print(f"CODE: {res}") + +# CHECK: // -----// IR Dump Before transform //----- // +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @matmul(%arg0: memref<8x512xf32> {llvm.noalias}, %arg1: memref<512x32xf32> {llvm.noalias}, %arg2: memref<8x32xf32> {llvm.noalias}) { +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%arg2 : memref<8x32xf32>) +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%arg0, %arg1 : memref<8x512xf32>, memref<512x32xf32>) outs(%arg2 : memref<8x32xf32>) +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @_vecto(%arg0: !transform.any_op {transform.consumed}) { +# CHECK-NEXT: transform.structured.vectorize %arg0 : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { +# CHECK-NEXT: %0 = transform.structured.match attributes {__xtc_id_C_0_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_linalg_op, %loops = transform.structured.tile_using_for %0 tile_sizes [1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops "./i" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_0, %loops_1 = transform.structured.tile_using_for %tiled_linalg_op tile_sizes [0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_1 "./j" : !transform.any_op +# CHECK-NEXT: %1 = transform.structured.match attributes {__xtc_id_C_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_op, %forall_op = transform.structured.tile_using_forall %1 tile_sizes [8, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %forall_op "./i" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_2, %loops_3 = transform.structured.tile_using_for %tiled_op tile_sizes [0, 32, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_3 "./j" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_4, %loops_5 = transform.structured.tile_using_for %tiled_linalg_op_2 tile_sizes [0, 0, 8] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_5 "./k" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_6, %loops_7 = transform.structured.tile_using_for %tiled_linalg_op_4 tile_sizes [4, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_7 "./i0" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_8, %loops_9 = transform.structured.tile_using_for %tiled_linalg_op_6 tile_sizes [0, 0, 2] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_9 "./k0" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_10, %loops_11 = transform.structured.tile_using_for %tiled_linalg_op_8 tile_sizes [0, 16, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_11 "./j0" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_12, %loops_13 = transform.structured.tile_using_for %tiled_linalg_op_10 tile_sizes [2, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_13 "./i1" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_14, %loops_15 = transform.structured.tile_using_for %tiled_linalg_op_12 tile_sizes [0, 0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_15 "./k1" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_16, %loops_17 = transform.structured.tile_using_for %tiled_linalg_op_14 tile_sizes [0, 2, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_17 "./j1" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_18, %loops_19 = transform.structured.tile_using_for %tiled_linalg_op_16 tile_sizes [1, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_19 "./i2" : !transform.any_op +# CHECK-NEXT: transform.include @_vecto failures(suppress) (%tiled_linalg_op_18) : (!transform.any_op) -> () +# CHECK-NEXT: transform.loop.unroll %loops_19 {factor = 2 : i64} : !transform.any_op +# CHECK-NEXT: %2 = transform.get_parent_op %forall_op {isolated_from_above} : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: transform.apply_patterns to %2 { +# CHECK-NEXT: transform.apply_patterns.vector.reduction_to_contract +# CHECK-NEXT: transform.apply_patterns.vector.transfer_permutation_patterns +# CHECK-NEXT: } : !transform.any_op +# CHECK-NEXT: transform.apply_patterns to %2 { +# CHECK-NEXT: transform.apply_patterns.vector.lower_outerproduct +# CHECK-NEXT: transform.apply_patterns.vector.lower_contraction +# CHECK-NEXT: } : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-EMPTY: +# CHECK-NEXT: // -----// IR Dump After transform //----- // +# CHECK-NEXT: #map = affine_map<(d0) -> (d0 * 8)> +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @matmul(%arg0: memref<8x512xf32> {llvm.noalias}, %arg1: memref<512x32xf32> {llvm.noalias}, %arg2: memref<8x32xf32> {llvm.noalias}) { +# CHECK-NEXT: %cst = arith.constant dense<0.000000e+00> : vector<1x2xf32> +# CHECK-NEXT: %0 = ub.poison : f32 +# CHECK-NEXT: %c16 = arith.constant 16 : index +# CHECK-NEXT: %c2 = arith.constant 2 : index +# CHECK-NEXT: %c4 = arith.constant 4 : index +# CHECK-NEXT: %c512 = arith.constant 512 : index +# CHECK-NEXT: %c32 = arith.constant 32 : index +# CHECK-NEXT: %cst_0 = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: %c0 = arith.constant 0 : index +# CHECK-NEXT: %c8 = arith.constant 8 : index +# CHECK-NEXT: %c1 = arith.constant 1 : index +# CHECK-NEXT: scf.for %arg3 = %c0 to %c8 step %c1 { +# CHECK-NEXT: %subview = memref.subview %arg2[%arg3, 0] [1, 32] [1, 1] : memref<8x32xf32> to memref<1x32xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: scf.for %arg4 = %c0 to %c32 step %c1 { +# CHECK-NEXT: %subview_1 = memref.subview %subview[0, %arg4] [1, 1] [1, 1] : memref<1x32xf32, strided<[32, 1], offset: ?>> to memref<1x1xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst_0 : f32) outs(%subview_1 : memref<1x1xf32, strided<[32, 1], offset: ?>>) +# CHECK-NEXT: } {"./j"} +# CHECK-NEXT: } {"./i"} +# CHECK-NEXT: scf.forall (%arg3) in (1) { +# CHECK-NEXT: %1 = affine.apply #map(%arg3) +# CHECK-NEXT: %subview = memref.subview %arg0[%1, 0] [8, 512] [1, 1] : memref<8x512xf32> to memref<8x512xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_1 = memref.subview %arg1[0, 0] [512, 32] [1, 1] : memref<512x32xf32> to memref<512x32xf32, strided<[32, 1]>> +# CHECK-NEXT: %subview_2 = memref.subview %arg2[%1, 0] [8, 32] [1, 1] : memref<8x32xf32> to memref<8x32xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: scf.for %arg4 = %c0 to %c32 step %c32 { +# CHECK-NEXT: %subview_3 = memref.subview %subview_1[0, %arg4] [512, 32] [1, 1] : memref<512x32xf32, strided<[32, 1]>> to memref<512x32xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: %subview_4 = memref.subview %subview_2[0, %arg4] [8, 32] [1, 1] : memref<8x32xf32, strided<[32, 1], offset: ?>> to memref<8x32xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: scf.for %arg5 = %c0 to %c512 step %c8 { +# CHECK-NEXT: %subview_5 = memref.subview %subview[0, %arg5] [8, 8] [1, 1] : memref<8x512xf32, strided<[512, 1], offset: ?>> to memref<8x8xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_6 = memref.subview %subview_3[%arg5, 0] [8, 32] [1, 1] : memref<512x32xf32, strided<[32, 1], offset: ?>> to memref<8x32xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: scf.for %arg6 = %c0 to %c8 step %c4 { +# CHECK-NEXT: %subview_7 = memref.subview %subview_5[%arg6, 0] [4, 8] [1, 1] : memref<8x8xf32, strided<[512, 1], offset: ?>> to memref<4x8xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_8 = memref.subview %subview_4[%arg6, 0] [4, 32] [1, 1] : memref<8x32xf32, strided<[32, 1], offset: ?>> to memref<4x32xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: scf.for %arg7 = %c0 to %c8 step %c2 { +# CHECK-NEXT: %subview_9 = memref.subview %subview_7[0, %arg7] [4, 2] [1, 1] : memref<4x8xf32, strided<[512, 1], offset: ?>> to memref<4x2xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_10 = memref.subview %subview_6[%arg7, 0] [2, 32] [1, 1] : memref<8x32xf32, strided<[32, 1], offset: ?>> to memref<2x32xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: scf.for %arg8 = %c0 to %c32 step %c16 { +# CHECK-NEXT: %subview_11 = memref.subview %subview_10[0, %arg8] [2, 16] [1, 1] : memref<2x32xf32, strided<[32, 1], offset: ?>> to memref<2x16xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: %subview_12 = memref.subview %subview_8[0, %arg8] [4, 16] [1, 1] : memref<4x32xf32, strided<[32, 1], offset: ?>> to memref<4x16xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: scf.for %arg9 = %c0 to %c4 step %c2 { +# CHECK-NEXT: %subview_13 = memref.subview %subview_9[%arg9, 0] [2, 2] [1, 1] : memref<4x2xf32, strided<[512, 1], offset: ?>> to memref<2x2xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_14 = memref.subview %subview_12[%arg9, 0] [2, 16] [1, 1] : memref<4x16xf32, strided<[32, 1], offset: ?>> to memref<2x16xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: scf.for %arg10 = %c0 to %c2 step %c1 { +# CHECK-NEXT: %subview_15 = memref.subview %subview_13[0, %arg10] [2, 1] [1, 1] : memref<2x2xf32, strided<[512, 1], offset: ?>> to memref<2x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_16 = memref.subview %subview_11[%arg10, 0] [1, 16] [1, 1] : memref<2x16xf32, strided<[32, 1], offset: ?>> to memref<1x16xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: scf.for %arg11 = %c0 to %c16 step %c2 { +# CHECK-NEXT: %subview_17 = memref.subview %subview_16[0, %arg11] [1, 2] [1, 1] : memref<1x16xf32, strided<[32, 1], offset: ?>> to memref<1x2xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: %subview_18 = memref.subview %subview_14[0, %arg11] [2, 2] [1, 1] : memref<2x16xf32, strided<[32, 1], offset: ?>> to memref<2x2xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: %subview_19 = memref.subview %subview_15[%c0, 0] [1, 1] [1, 1] : memref<2x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_20 = memref.subview %subview_18[%c0, 0] [1, 2] [1, 1] : memref<2x2xf32, strided<[32, 1], offset: ?>> to memref<1x2xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: %2 = vector.transfer_read %subview_19[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x1xf32, strided<[512, 1], offset: ?>>, vector<1x1xf32> +# CHECK-NEXT: %3 = vector.transfer_read %subview_17[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x2xf32, strided<[32, 1], offset: ?>>, vector<1x2xf32> +# CHECK-NEXT: %4 = vector.transfer_read %subview_20[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x2xf32, strided<[32, 1], offset: ?>>, vector<1x2xf32> +# CHECK-NEXT: %5 = vector.extract %3[0] : vector<2xf32> from vector<1x2xf32> +# CHECK-NEXT: %6 = vector.extract %2[0, 0] : f32 from vector<1x1xf32> +# CHECK-NEXT: %7 = vector.broadcast %6 : f32 to vector<2xf32> +# CHECK-NEXT: %8 = vector.extract %4[0] : vector<2xf32> from vector<1x2xf32> +# CHECK-NEXT: %9 = vector.fma %7, %5, %8 : vector<2xf32> +# CHECK-NEXT: %10 = vector.insert %9, %cst [0] : vector<2xf32> into vector<1x2xf32> +# CHECK-NEXT: vector.transfer_write %10, %subview_20[%c0, %c0] {in_bounds = [true, true]} : vector<1x2xf32>, memref<1x2xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: %subview_21 = memref.subview %subview_15[%c1, 0] [1, 1] [1, 1] : memref<2x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> +# CHECK-NEXT: %subview_22 = memref.subview %subview_18[%c1, 0] [1, 2] [1, 1] : memref<2x2xf32, strided<[32, 1], offset: ?>> to memref<1x2xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: %11 = vector.transfer_read %subview_21[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x1xf32, strided<[512, 1], offset: ?>>, vector<1x1xf32> +# CHECK-NEXT: %12 = vector.transfer_read %subview_17[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x2xf32, strided<[32, 1], offset: ?>>, vector<1x2xf32> +# CHECK-NEXT: %13 = vector.transfer_read %subview_22[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x2xf32, strided<[32, 1], offset: ?>>, vector<1x2xf32> +# CHECK-NEXT: %14 = vector.extract %12[0] : vector<2xf32> from vector<1x2xf32> +# CHECK-NEXT: %15 = vector.extract %11[0, 0] : f32 from vector<1x1xf32> +# CHECK-NEXT: %16 = vector.broadcast %15 : f32 to vector<2xf32> +# CHECK-NEXT: %17 = vector.extract %13[0] : vector<2xf32> from vector<1x2xf32> +# CHECK-NEXT: %18 = vector.fma %16, %14, %17 : vector<2xf32> +# CHECK-NEXT: %19 = vector.insert %18, %cst [0] : vector<2xf32> into vector<1x2xf32> +# CHECK-NEXT: vector.transfer_write %19, %subview_22[%c0, %c0] {in_bounds = [true, true]} : vector<1x2xf32>, memref<1x2xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: } {"./j1"} +# CHECK-NEXT: } {"./k1"} +# CHECK-NEXT: } {"./i1"} +# CHECK-NEXT: } {"./j0"} +# CHECK-NEXT: } {"./k0"} +# CHECK-NEXT: } {"./i0"} +# CHECK-NEXT: } {"./k"} +# CHECK-NEXT: } {"./j"} +# CHECK-NEXT: } {"./i"} +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-EMPTY: +# CHECK-NEXT: graph: +# CHECK-NEXT: name: matmul +# CHECK-NEXT: inputs: +# CHECK-NEXT: - %0 : 8x512xfloat32 +# CHECK-NEXT: - %1 : 512x32xfloat32 +# CHECK-NEXT: outputs: +# CHECK-NEXT: - %2 : 8x32xfloat32 +# CHECK-NEXT: nodes: +# CHECK-NEXT: - %2: matmul(%0, %1) {name = 'C'} : [8x512xfloat32, 512x32xfloat32] -> [8x32xfloat32] +# CHECK-EMPTY: +# CHECK-NEXT: CODE: 0 diff --git a/tests/filecheck/search/test_matmul_descript_twuop.py b/tests/filecheck/search/test_matmul_descript_twuop.py new file mode 100644 index 00000000..7719b756 --- /dev/null +++ b/tests/filecheck/search/test_matmul_descript_twuop.py @@ -0,0 +1,21 @@ +# RUN: python %s 2>&1 | filecheck %s +# REQUIRES: module_xvs +""" +Test strategy Goto on matmul +""" + +import utils +from xtc.search.strategies import Strategy_Descript as Strategy + +graph = utils.get_graph_matmul() +backend = utils.get_backend(graph) +spec = """ +TWUOP: parallelize unroll vectorize +""" +strategy = Strategy(graph, spec, initialize=False, partial_tiles=True, partial_unrolls=True) + +print(sorted(strategy._constraints)) +print(sum(1 for _ in strategy.sample(100))) + +# CHECK: ['1 <= prt_interchange_u_0 <= 6', 'prt_i_0 <= 21', 'prt_i_1 <= prt_i_0', 'prt_i_2 <= prt_i_1', 'prt_j_0 <= 32', 'prt_j_1 <= prt_j_0', 'prt_j_2 <= prt_j_1', 'prt_k_0 <= 12', 'prt_k_1 <= prt_k_0'] +# CHECK-NEXT: 100 From 0746ffdd3ec90d3eca48e44e5f651f8527f31c59 Mon Sep 17 00:00:00 2001 From: Leon Frenot Date: Thu, 16 Jul 2026 14:04:36 +0200 Subject: [PATCH 15/23] Remove "sizes" from descript interpreter Replaced with "axes" when needed --- src/xtc/schedules/descript.py | 97 +++++++++++++++++------------------ 1 file changed, 46 insertions(+), 51 deletions(-) diff --git a/src/xtc/schedules/descript.py b/src/xtc/schedules/descript.py index deca4181..32937e26 100644 --- a/src/xtc/schedules/descript.py +++ b/src/xtc/schedules/descript.py @@ -111,13 +111,11 @@ def _interpret_spec_into_node( previous_cut: dict[str, literal | None] = {a: 0 for a in self.abstract_dims} node.interchange = list(head) last_split: list[tuple[literal, literal | None]] = [] - sizes: dict[str, literal] = {} loop_name: str = "" if not axes: axes = dict() if self.abstract_dim_sizes: for a, v in self.abstract_dim_sizes.items(): - sizes[a] = v if a not in axes: axes[a] = [v] if not axes: @@ -137,20 +135,18 @@ def _interpret_spec_into_node( loop_name = self._interpret_tile( item=item, node=node, - sizes=sizes, axes=axes, ) - self._apply_annotations(item.annotations, loop_name, sizes, node) + self._apply_annotations(item, loop_name, axes, node) elif isinstance(item, AxisDecl): _loop_name = self._interpret_axis(item, node) if _loop_name: loop_name = _loop_name - self._apply_annotations(item.annotations, loop_name, sizes, node) + self._apply_annotations(item, loop_name, axes, node) elif isinstance(item, PRTDecl): loop_name = self._interpret_prt( item=item, node=node, - sizes=sizes, axes=axes, ) @@ -177,11 +173,7 @@ def _interpret_spec_into_node( # Check that all splits are complete for axis, cut in previous_cut.items(): - if ( - cut is not None - and isinstance(cut, int) - and cut not in [0, sizes.get(axis, 0)] - ): + if cut is not None and isinstance(cut, int) and cut not in [0] + axes[axis]: raise ScheduleInterpretError( f"Splitting of {axis} unachieved (stops at {cut})." ) @@ -303,7 +295,6 @@ def _interpret_tile( self, item: TileDecl, node: ParameterLoopNestNode, - sizes: dict[str, literal], axes: dict[str, list[literal]], ) -> str: """Interpret a tile declaration. Returns the loop name.""" @@ -315,7 +306,6 @@ def _interpret_tile( f"`{item}`: tile sizes should be strictly positive." ) node.tiles[item.axis][loop_name] = item.size - sizes[loop_name] = item.size inter_group = item.annotations.interchange if inter_group: @@ -397,7 +387,6 @@ def _interpret_prt( self, item: PRTDecl, node: ParameterLoopNestNode, - sizes: dict[str, literal], axes: dict[str, list[literal]], ) -> str: loop_name = "" @@ -431,25 +420,27 @@ def _parallelize(): else: node.parallelize.append(loop_name) - def _unroll(): + def _unroll(axis: str): unroll_factor = item.annotations.unroll_factor - if unroll_factor is None: - # None means "unroll fully" - use the loop size - if loop_name not in sizes: - raise ScheduleInterpretError( - f"{loop_name}'s size being unknown, an unroll factor is needed." - ) - unroll_factor = sizes[loop_name] - elif isinstance(unroll_factor, int) and unroll_factor <= 0: + if isinstance(unroll_factor, int) and unroll_factor <= 0: raise ScheduleInterpretError( f'`{{"unroll" = {unroll_factor}}}`: unroll parameter should be strictly positive.' ) - elif isinstance(unroll_factor, str): - unroll_factor += "_prt_" + loop_name - if self.partial_unrolls: - node.constraints.append(f"{unroll_factor} <= {sizes[loop_name]}") - else: - node.constraints.append(f"{unroll_factor} || {sizes[loop_name]}") + if unroll_factor is None or isinstance(unroll_factor, str): + if not axes[axis]: + raise ScheduleInterpretError( + f"{loop_name}'s size being unknown, an unroll factor is needed." + ) + axis_size = axes[axis][-1] + if unroll_factor is None: + # None means "unroll fully" - use the loop size + unroll_factor = axis_size + elif isinstance(unroll_factor, str): + unroll_factor += "_prt_" + loop_name + if self.partial_unrolls: + node.constraints.append(f"{unroll_factor} <= {axis_size}") + else: + node.constraints.append(f"{unroll_factor} || {axis_size}") node.unroll[loop_name] = unroll_factor def _interpret(axis: str, annotations: Annotations | None = None): @@ -461,7 +452,7 @@ def _interpret(axis: str, annotations: Annotations | None = None): if not annotations: annotations = _annotations(a) return self._interpret_tile( - TileDecl(a, self._fresh(a), annotations), node, sizes, axes + TileDecl(a, self._fresh(a), annotations), node, axes ) if level and len(item.shape) > 1: @@ -481,7 +472,7 @@ def _interpret(axis: str, annotations: Annotations | None = None): _parallelize() parallelize = False if unroll and idx + 1 == len(item.shape): - _unroll() + _unroll(a) case "R": if not self.abstract_r_dims: raise ScheduleInterpretError( @@ -494,7 +485,7 @@ def _interpret(axis: str, annotations: Annotations | None = None): "Cannot parallelize a reduction axis" ) if unroll and idx + 1 == len(item.shape): - _unroll() + _unroll(a) case "T": for a in self.abstract_dims: loop_name = _interpret(a) @@ -502,7 +493,7 @@ def _interpret(axis: str, annotations: Annotations | None = None): _parallelize() parallelize = False if unroll and idx + 1 == len(item.shape): - _unroll() + _unroll(a) case "U": if level: raise ScheduleInterpretError("Cannot use level on a U tile.") @@ -516,7 +507,7 @@ def _interpret(axis: str, annotations: Annotations | None = None): if parallelize: raise ScheduleInterpretError("Cannot parallelize a U tile.") if unroll and idx + 1 == len(item.shape): - _unroll() + _unroll(a) case "O": if not self.abstract_p_dims: raise ScheduleInterpretError( @@ -532,15 +523,15 @@ def _interpret(axis: str, annotations: Annotations | None = None): _parallelize() parallelize = False if unroll and idx + 1 == len(item.shape): - _unroll() + _unroll(a) for a in self.abstract_r_dims: _interpret(a) if unroll and idx + 1 == len(item.shape): - _unroll() + _unroll(a) for a in self.abstract_p_dims[1:]: loop_name = _interpret(a) if unroll and idx + 1 == len(item.shape): - _unroll() + _unroll(a) case "W": node.buffer_at[loop_name] = None case "F": @@ -578,30 +569,34 @@ def _check_axis_existence(self, axis: str) -> None: def _apply_annotations( self, - annotations: Annotations, + item: TileDecl | AxisDecl, loop_name: str, - sizes: dict[str, literal], + axes: dict[str, list[literal]], node: ParameterLoopNestNode, ) -> None: """Apply annotations to a loop in the node.""" + annotations = item.annotations + if annotations.unroll_specified: unroll_factor = annotations.unroll_factor - if unroll_factor is None: - # None means "unroll fully" - use the loop size - if loop_name not in sizes: - raise ScheduleInterpretError( - f"{loop_name}'s size being unknown, an unroll factor is needed." - ) - unroll_factor = sizes[loop_name] - elif isinstance(unroll_factor, int) and unroll_factor <= 0: + if isinstance(unroll_factor, int) and unroll_factor <= 0: raise ScheduleInterpretError( f'`{{"unroll" = {unroll_factor}}}`: unroll parameter should be strictly positive.' ) - elif isinstance(unroll_factor, str): - if self.partial_unrolls: - node.constraints.append(f"{unroll_factor} <= {sizes[loop_name]}") + if unroll_factor is None or isinstance(unroll_factor, str): + if not axes[item.axis]: + raise ScheduleInterpretError( + f"{loop_name}'s size being unknown, an unroll factor is needed." + ) + axis_size = axes[item.axis][-1] + if unroll_factor is None: + # None means "unroll fully" - use the loop size + unroll_factor = axis_size else: - node.constraints.append(f"{unroll_factor} || {sizes[loop_name]}") + if self.partial_unrolls: + node.constraints.append(f"{unroll_factor} <= {axis_size}") + else: + node.constraints.append(f"{unroll_factor} || {axis_size}") node.unroll[loop_name] = unroll_factor if annotations.vectorize: From a5e045f221799b0db6a885e54ecf18dcb3b79b2b Mon Sep 17 00:00:00 2001 From: Leon Frenot Date: Thu, 16 Jul 2026 15:30:00 +0200 Subject: [PATCH 16/23] Minor refactors --- src/xtc/schedules/descript.py | 92 ++++++++++++++++------------------- 1 file changed, 43 insertions(+), 49 deletions(-) diff --git a/src/xtc/schedules/descript.py b/src/xtc/schedules/descript.py index 32937e26..0cc1f282 100644 --- a/src/xtc/schedules/descript.py +++ b/src/xtc/schedules/descript.py @@ -79,13 +79,6 @@ class ScheduleInterpreter: _abstract_dim_fresh: dict[str, int] = field(default_factory=dict) _levels: dict[str, dict[str, literal | None]] = field(default_factory=dict) - def _fresh(self, axis: str): - if axis not in self._abstract_dim_fresh: - self._abstract_dim_fresh[axis] = 0 - else: - self._abstract_dim_fresh[axis] += 1 - return f"prt_{axis}_{self._abstract_dim_fresh[axis]}" - def interpret(self, spec: ScheduleSpec, root: str) -> ParameterLoopNest: """Interpret a schedule specification into a LoopNest.""" loop_nest = ParameterLoopNest(abstract_dims=self.abstract_dims) @@ -109,9 +102,11 @@ def _interpret_spec_into_node( """Interpret a schedule spec into an existing node (mutates node).""" # Track state during interpretation previous_cut: dict[str, literal | None] = {a: 0 for a in self.abstract_dims} - node.interchange = list(head) last_split: list[tuple[literal, literal | None]] = [] loop_name: str = "" + + node.interchange = list(head) + if not axes: axes = dict() if self.abstract_dim_sizes: @@ -222,15 +217,10 @@ def _interpret_split( node.splits[axis_name][new_dim_name] = x node.interchange.append(new_dim_name) - inner_size = None if y is None: y = current_size - if isinstance(x, int): - if x == 0: - inner_size = y - elif isinstance(y, int): - inner_size = y - x - if inner_size is None: + + if isinstance(x, str) or (x != 0 and isinstance(y, str)): inner_size = root[1:] + new_dim_name inner_size = ( inner_size.replace(_SPLIT_LEFT, "_") @@ -242,10 +232,8 @@ def _interpret_split( node.constraints.append(f"{x} <= {y}") node.constraints.append(f"{inner_size} + {x} == {y}") else: - inner_size = z - x = cut y = current_size - assert x is not None + # Save the cutting points of the new dimensions node.splits[axis_name][new_dim_name] = x node.interchange.append(new_dim_name) @@ -299,12 +287,13 @@ def _interpret_tile( ) -> str: """Interpret a tile declaration. Returns the loop name.""" self._check_axis_existence(item.axis) - tile_num = len(node.tiles[item.axis]) - loop_name = f"{item.axis}{tile_num}" if isinstance(item.size, int) and item.size <= 0: raise ScheduleInterpretError( f"`{item}`: tile sizes should be strictly positive." ) + + tile_num = len(node.tiles[item.axis]) + loop_name = f"{item.axis}{tile_num}" node.tiles[item.axis][loop_name] = item.size inter_group = item.annotations.interchange @@ -315,8 +304,10 @@ def _interpret_tile( node.interchange_groups[inter_group] = { len(node.interchange): loop_name } + node.interchange.append(loop_name) + # Alias for axes[item.axis] list_axis: list[int | str] if item.axis in axes: list_axis = axes[item.axis] @@ -338,11 +329,7 @@ def _interpret_tile( raise ScheduleInterpretError( f"`{item}` is a full tile, but the axis sizes are unknown." ) - s = ( - ", ".join(map(str, list_axis)) - if len(list_axis) > 1 - else str(list_axis[0]) - ) + s = ", ".join(map(str, list_axis)) s = f"{item.size} || {{{s}}}" node.constraints.append(s) list_axis.append(item.size) @@ -369,6 +356,7 @@ def _interpret_axis( f"Axis {axis_name} is scheduled twice (or more)." ) + # If the axis is in an interchange group, update that group inter_group = item.annotations.interchange if inter_group: if inter_group in node.interchange_groups: @@ -377,19 +365,26 @@ def _interpret_axis( node.interchange_groups[inter_group] = { len(node.interchange): axis_name } + node.interchange.append(axis_name) if self.abstract_dim_sizes: self._update_levels(axis_name, self.abstract_dim_sizes[axis_name]) return axis_name + def _fresh(self, axis: str): + if axis not in self._abstract_dim_fresh: + self._abstract_dim_fresh[axis] = 0 + else: + self._abstract_dim_fresh[axis] += 1 + return f"prt_{axis}_{self._abstract_dim_fresh[axis]}" + def _interpret_prt( self, item: PRTDecl, node: ParameterLoopNestNode, axes: dict[str, list[literal]], ) -> str: - loop_name = "" parallelize = item.annotations.parallelize unroll = item.annotations.unroll_specified @@ -399,18 +394,25 @@ def _interpret_prt( if interchange == "interchange": interchange = self._fresh("interchange") + loop_name: str = "" + + if level and len(item.shape) > 1: + raise ScheduleInterpretError( + "Cannot use the same level on mulitple PRT tiles." + ) + def _annotations(axis: str) -> Annotations: - if not level: + if level: return Annotations( partial=item.annotations.partial, full=item.annotations.full, interchange=interchange, + level=level + axis, ) return Annotations( partial=item.annotations.partial, full=item.annotations.full, interchange=interchange, - level=level + axis, ) def _parallelize(): @@ -444,21 +446,16 @@ def _unroll(axis: str): node.unroll[loop_name] = unroll_factor def _interpret(axis: str, annotations: Annotations | None = None): + if not annotations: + annotations = _annotations(axis) + if axis not in node.interchange: - if not annotations: - annotations = Annotations() return self._interpret_axis(AxisDecl(axis, annotations), node) else: - if not annotations: - annotations = _annotations(a) return self._interpret_tile( - TileDecl(a, self._fresh(a), annotations), node, axes + TileDecl(axis, self._fresh(axis), annotations), node, axes ) - if level and len(item.shape) > 1: - raise ScheduleInterpretError( - "Cannot use the same level on mulitple PRT tiles." - ) for idx, t in enumerate(item.shape): match t: case "P": @@ -517,19 +514,16 @@ def _interpret(axis: str, annotations: Annotations | None = None): raise ScheduleInterpretError( "O scheme used, but R axes are not specified" ) - a = self.abstract_p_dims[0] - _interpret(a) - if parallelize: - _parallelize() - parallelize = False - if unroll and idx + 1 == len(item.shape): - _unroll(a) - for a in self.abstract_r_dims: + dims = ( + [self.abstract_p_dims[0]] + + self.abstract_r_dims + + self.abstract_p_dims[1:] + ) + for a in dims: _interpret(a) - if unroll and idx + 1 == len(item.shape): - _unroll(a) - for a in self.abstract_p_dims[1:]: - loop_name = _interpret(a) + if parallelize: + _parallelize() + parallelize = False if unroll and idx + 1 == len(item.shape): _unroll(a) case "W": From 5178622001cf9d6439eaf68c5a30f462b6dabee5 Mon Sep 17 00:00:00 2001 From: Leon Frenot Date: Thu, 16 Jul 2026 15:30:06 +0200 Subject: [PATCH 17/23] Comments --- src/xtc/schedules/descript.py | 50 ++++++++++++++++++++++++++++++++--- 1 file changed, 47 insertions(+), 3 deletions(-) diff --git a/src/xtc/schedules/descript.py b/src/xtc/schedules/descript.py index 0cc1f282..94750b17 100644 --- a/src/xtc/schedules/descript.py +++ b/src/xtc/schedules/descript.py @@ -100,9 +100,11 @@ def _interpret_spec_into_node( axes: dict[str, list[literal]] | None = None, ) -> None: """Interpret a schedule spec into an existing node (mutates node).""" - # Track state during interpretation + # Track previous split for each axis previous_cut: dict[str, literal | None] = {a: 0 for a in self.abstract_dims} + # Track finak split and corresponding missing size last_split: list[tuple[literal, literal | None]] = [] + # Name of the last loop nest (declared here for typing/liveness) loop_name: str = "" node.interchange = list(head) @@ -145,6 +147,8 @@ def _interpret_spec_into_node( axes=axes, ) + # Check interchange groups only use each axis at most once + # And create the sampling constraint for k, v_d in node.interchange_groups.items(): read: set[str] = set() for k_, v_d_ in v_d.items(): @@ -157,7 +161,7 @@ def _interpret_spec_into_node( read.add(a) node.constraints.append(f"1 <= {k} <= {factorial(len(v_d))}") - # Reaplace the placeholder of the last split with its size + # Replace the placeholder of the last split with its size if len(last_split) > 0: a0, b0 = last_split[0] if isinstance(a0, int) and not isinstance(b0, int): @@ -209,7 +213,10 @@ def _interpret_split( new_dim_name = f"{axis_name}{_SPLIT_LEFT}{new_dim_index}{_SPLIT_RIGHT}" new_root_name = f"{root}{_NODE_SEP}{new_dim_name}" + # i[x:y] if z is None: + # Split from x to y + # Update the previous cut previous_cut[axis_name] = y @@ -217,9 +224,11 @@ def _interpret_split( node.splits[axis_name][new_dim_name] = x node.interchange.append(new_dim_name) + # i[x:] => i[x:end] if y is None: y = current_size + # Add constraints if the size of the split, or the size of the tile being split, is unknown if isinstance(x, str) or (x != 0 and isinstance(y, str)): inner_size = root[1:] + new_dim_name inner_size = ( @@ -231,12 +240,17 @@ def _interpret_split( if isinstance(x, str): node.constraints.append(f"{x} <= {y}") node.constraints.append(f"{inner_size} + {x} == {y}") + + # i[:z:] else: + # Split from x, of size z, that fits in y y = current_size # Save the cutting points of the new dimensions node.splits[axis_name][new_dim_name] = x node.interchange.append(new_dim_name) + + # Update the previous cut, and add constraints if needed if isinstance(z, int) and isinstance(x, int): previous_cut[axis_name] = x + z if not isinstance(y, int): @@ -267,6 +281,7 @@ def _interpret_split( ) node.add_child(child_node) + # TODO: is this the correct footprint? if current_size is not None: self._update_levels(item.axis, current_size) @@ -296,6 +311,7 @@ def _interpret_tile( loop_name = f"{item.axis}{tile_num}" node.tiles[item.axis][loop_name] = item.size + # If the tile is in an interchange group, update that group inter_group = item.annotations.interchange if inter_group: if inter_group in node.interchange_groups: @@ -314,9 +330,12 @@ def _interpret_tile( else: list_axis = [] axes[item.axis] = list_axis + + # Add the contraints if isinstance(item.size, str): partial = item.annotations.partial full = item.annotations.full + # Partial tile if partial or (not full and self.partial_tiles): if not list_axis: raise ScheduleInterpretError( @@ -324,6 +343,7 @@ def _interpret_tile( ) old_size = list_axis[-1] node.constraints.append(f"{item.size} <= {old_size}") + # Full tile else: if not self.abstract_dim_sizes: raise ScheduleInterpretError( @@ -332,6 +352,8 @@ def _interpret_tile( s = ", ".join(map(str, list_axis)) s = f"{item.size} || {{{s}}}" node.constraints.append(s) + + # Update axes list_axis.append(item.size) self._update_levels(item.axis, axes[item.axis][-1]) @@ -345,8 +367,11 @@ def _interpret_axis( """Interpret a direct axis reference. Returns the loop name.""" axis_name = item.axis interchange = node.interchange + + # AxisDecl is actually a bufferization line if axis_name in self.abstract_matrix: return "" + self._check_axis_existence(axis_name) # Unreachable when built from a Python dict (because keys @@ -385,11 +410,18 @@ def _interpret_prt( node: ParameterLoopNestNode, axes: dict[str, list[literal]], ) -> str: + """Interpret a sequence of PRT tiles. Returns the last loop name. + The comportement of annotations is derived from BaseStrategyPRTScheme. + Parallelize is applied on the first tile only. + Unroll is applied on the last tile only. + Vectorize is applied on the last axis only. + Pack/Buffer are applied after the last axis. + """ + # Collect annotations parallelize = item.annotations.parallelize unroll = item.annotations.unroll_specified level = item.annotations.level - interchange = item.annotations.interchange if interchange == "interchange": interchange = self._fresh("interchange") @@ -459,6 +491,7 @@ def _interpret(axis: str, annotations: Annotations | None = None): for idx, t in enumerate(item.shape): match t: case "P": + # All P-axes in order if not self.abstract_p_dims: raise ScheduleInterpretError( "P scheme used, but P axes are not specified" @@ -471,6 +504,7 @@ def _interpret(axis: str, annotations: Annotations | None = None): if unroll and idx + 1 == len(item.shape): _unroll(a) case "R": + # All R-axes in order if not self.abstract_r_dims: raise ScheduleInterpretError( "R scheme used, but R axes are not specified" @@ -484,6 +518,7 @@ def _interpret(axis: str, annotations: Annotations | None = None): if unroll and idx + 1 == len(item.shape): _unroll(a) case "T": + # All axes in order for a in self.abstract_dims: loop_name = _interpret(a) if parallelize: @@ -492,6 +527,7 @@ def _interpret(axis: str, annotations: Annotations | None = None): if unroll and idx + 1 == len(item.shape): _unroll(a) case "U": + # All axes in a random order (creates an interchange group) if level: raise ScheduleInterpretError("Cannot use level on a U tile.") annotations_u = Annotations( @@ -506,6 +542,7 @@ def _interpret(axis: str, annotations: Annotations | None = None): if unroll and idx + 1 == len(item.shape): _unroll(a) case "O": + # first P-axis then R-axes, then remaining P-axis if not self.abstract_p_dims: raise ScheduleInterpretError( "O scheme used, but P axes are not specified" @@ -527,10 +564,13 @@ def _interpret(axis: str, annotations: Annotations | None = None): if unroll and idx + 1 == len(item.shape): _unroll(a) case "W": + # Add an output write buffer at this level node.buffer_at[loop_name] = None case "F": + # Fuse producer at this level raise ScheduleInterpretError("TODO: Fusion unimplemented") + # Vectorize the last axis if item.annotations.vectorize: if isinstance(item.annotations.vectorize, str): node.vectorize_parameters[loop_name] = item.annotations.vectorize @@ -538,9 +578,11 @@ def _interpret(axis: str, annotations: Annotations | None = None): else: node.vectorize.append(loop_name) + # Write buffer after the last axis if item.annotations.buffer_specified: node.buffer_at[loop_name] = item.annotations.buffer + # Read buffer after the last axis if item.annotations.pack_specified and item.annotations.pack is not None: input_matrix, mtype, pad = item.annotations.pack if isinstance(input_matrix, str): @@ -620,6 +662,7 @@ def _apply_annotations( input_matrix, mtype, pad = annotations.pack if isinstance(input_matrix, str): idx = self.abstract_matrix.index(input_matrix) + # Write buffer if idx == len(self.abstract_matrix) - 1: node.buffer_at[loop_name] = mtype if isinstance(annotations.pack_specified, str): @@ -627,6 +670,7 @@ def _apply_annotations( node.constraints.append( f"{annotations.buffer_specified} in {{0, 1}}" ) + # Read buffer else: node.pack_at[loop_name] = (idx, mtype, pad) if isinstance(annotations.pack_specified, str): From 0c84fb564279b3ef75b3b19bde5b6700a6dd1b44 Mon Sep 17 00:00:00 2001 From: Leon Frenot Date: Thu, 16 Jul 2026 16:00:36 +0200 Subject: [PATCH 18/23] User defined functions in constraints --- src/xtc/search/strategies.py | 8 +++-- .../search/test_matmul_descript_yaml_goto.py | 31 ++++++++++++++----- 2 files changed, 29 insertions(+), 10 deletions(-) diff --git a/src/xtc/search/strategies.py b/src/xtc/search/strategies.py index 8490e451..8b993e21 100644 --- a/src/xtc/search/strategies.py +++ b/src/xtc/search/strategies.py @@ -4,7 +4,7 @@ # from abc import abstractmethod from dataclasses import dataclass, field -from typing import TypeAlias, Any +from typing import TypeAlias, Any, Callable from typing_extensions import override from collections.abc import Sequence, Mapping, Iterator, Generator import itertools @@ -1039,11 +1039,13 @@ def __init__( constraints: list[str] | None = None, axes: list[str] | None = None, tensors: list[str] | None = None, + functions: dict[str, Callable] | None = None, partial_tiles: bool = False, partial_unrolls: bool = False, initialize: bool = True, ) -> None: self._graph = graph + self._functions = functions self._op = graph.outputs_nodes[0].operation self._stats: dict[str, int] = {} self._axes = axes if axes else list(self._op.dims) @@ -1123,7 +1125,7 @@ def _initialize(self): self._initialized = True return max_enum = int(1 + np.log2(max(self._sizes.values()))) - context = Context() + context = Context(functions=self._functions) constraints, self.constants = constraints_from_str( list(self._constraints), context=context ) @@ -1212,6 +1214,7 @@ def __init__( constraints: list[str] | None = None, axes: list[str] | None = None, tensors: list[str] | None = None, + functions: dict[str, Callable] | None = None, partial_tiles: bool = False, partial_unrolls: bool = False, initialize: bool = True, @@ -1223,6 +1226,7 @@ def __init__( constraints, axes, tensors, + functions, partial_tiles, partial_unrolls, initialize, diff --git a/tests/filecheck/search/test_matmul_descript_yaml_goto.py b/tests/filecheck/search/test_matmul_descript_yaml_goto.py index 591568ce..62a29f8a 100644 --- a/tests/filecheck/search/test_matmul_descript_yaml_goto.py +++ b/tests/filecheck/search/test_matmul_descript_yaml_goto.py @@ -23,19 +23,20 @@ nb_registers = 32 nb_fma = 2 fma_latency = 4 -ilp = nb_fma*fma_latency +ilp = nb_fma * fma_latency vector_size = 16 elt_size = 4 reorder_buffer = 256 -nb_words_L1 = 32*1024//elt_size -nb_words_L2 = 1024*1024//elt_size -nb_words_L3 = 36*1024*1024//elt_size +nb_words_L1 = 32 * 1024 // elt_size +nb_words_L2 = 1024 * 1024 // elt_size +nb_words_L3 = 36 * 1024 * 1024 // elt_size +RoB = 200 spec = f""" constraints: - 1 + nvr + nvr * mr <= {nb_registers} - nr == {vector_size} * nvr - - nvr * mr >= {ilp} + - ilp(nvr, mr, mc, nc, kc) >= {ilp} - nvr * mr * kr <= {reorder_buffer} - footprint(B, L1) <= {nb_words_L1} - footprint(A, L2) <= {nb_words_L2} @@ -52,10 +53,24 @@ j#nr: vectorize full """ -strategy = Strategy(graph, spec, partial_tiles=True, partial_unrolls=True, initialize=False) + +def fn_ilp(nvr, mr, mc, nc, kc): + if nvr * mr * kc >= RoB: + return nvr * mr + return min(nvr * mr * mc * nc, RoB / kc) + + +strategy = Strategy( + graph, + spec, + functions={"ilp": fn_ilp}, + partial_tiles=True, + partial_unrolls=True, + initialize=False, +) print(sorted(strategy._constraints)) print(sum(1 for _ in strategy.sample(100))) -# CHECK: ['1 + nvr + nvr * mr <= 32', 'kc * nc <= 9437184', 'kc * nr <= 8192', 'kc <= 1024', 'kr <= kc', 'mc * kc <= 262144', 'mc <= 1024', 'mr || {1024, mc}', 'nc <= 1024', 'nr == 16 * nvr', 'nr || {1024, nc}', 'nvr * mr * kr <= 256', 'nvr * mr >= 8', 'pack_B in {0, 1}', 'pad_A in {0,1}'] -#CHECK-NEXT: 100 +# CHECK: ['1 + nvr + nvr * mr <= 32', 'ilp(nvr, mr, mc, nc, kc) >= 8', 'kc * nc <= 9437184', 'kc * nr <= 8192', 'kc <= 1024', 'kr <= kc', 'mc * kc <= 262144', 'mc <= 1024', 'mr || {1024, mc}', 'nc <= 1024', 'nr == 16 * nvr', 'nr || {1024, nc}', 'nvr * mr * kr <= 256', 'pack_B in {0, 1}', 'pad_A in {0,1}'] +# CHECK-NEXT: 100 From 002fcf9c4ce157ab9f2e89fbb7d1afb064fe22f9 Mon Sep 17 00:00:00 2001 From: Leon Frenot Date: Mon, 20 Jul 2026 13:42:06 +0200 Subject: [PATCH 19/23] "schedule" entry in yaml descript spec --- src/xtc/schedules/parsing.py | 15 ++++++++++++--- .../search/test_matmul_descript_yaml_goto.py | 17 +++++++++-------- 2 files changed, 21 insertions(+), 11 deletions(-) diff --git a/src/xtc/schedules/parsing.py b/src/xtc/schedules/parsing.py index 750d7aec..ad049f68 100644 --- a/src/xtc/schedules/parsing.py +++ b/src/xtc/schedules/parsing.py @@ -15,6 +15,8 @@ literal = int | str ansor_tile = "PRTUOWF" tup_list = list[tuple[str, Any]] +CONSTRAINTS = "constraints" +SCHEDULE = "schedule" def toliteral(s: str) -> literal: @@ -25,9 +27,17 @@ def toliteral(s: str) -> literal: def pre_parse(s: dict[str, Any] | tup_list | list[dict[str, Any]]) -> tup_list: + out: tup_list = [] if isinstance(s, dict): + if SCHEDULE in s: + if CONSTRAINTS in s: + out = [(CONSTRAINTS, s[CONSTRAINTS])] + if not set(s.keys()) <= {CONSTRAINTS, SCHEDULE}: + raise ScheduleParseError( + f"Parsed entries other than schedule and constraints in spec:\n{s}" + ) + return out + pre_parse(s[SCHEDULE]) return [(k, _pre_parse(v)) for k, v in s.items()] - out: tup_list = [] for s_ in s: if isinstance(s_, dict): out += [(k, _pre_parse(v)) for k, v in s_.items()] @@ -397,10 +407,9 @@ def parse(self, spec: str) -> tup_list: def _parse(self, spec: tup_list) -> tup_list: """Parses a dict YAML specification into a schedule specification.""" - # constraints = spec.pop("constraints", []) descript_spec: tup_list = [] for a, v in spec: - if a == "constraints": + if a == CONSTRAINTS: descript_spec.append((a, v)) continue if isinstance(v, str): diff --git a/tests/filecheck/search/test_matmul_descript_yaml_goto.py b/tests/filecheck/search/test_matmul_descript_yaml_goto.py index 62a29f8a..fc777d73 100644 --- a/tests/filecheck/search/test_matmul_descript_yaml_goto.py +++ b/tests/filecheck/search/test_matmul_descript_yaml_goto.py @@ -33,14 +33,15 @@ RoB = 200 spec = f""" - constraints: - - 1 + nvr + nvr * mr <= {nb_registers} - - nr == {vector_size} * nvr - - ilp(nvr, mr, mc, nc, kc) >= {ilp} - - nvr * mr * kr <= {reorder_buffer} - - footprint(B, L1) <= {nb_words_L1} - - footprint(A, L2) <= {nb_words_L2} - - footprint(B, L3) <= {nb_words_L3} +constraints: + - 1 + nvr + nvr * mr <= {nb_registers} + - nr == {vector_size} * nvr + - ilp(nvr, mr, mc, nc, kc) >= {ilp} + - nvr * mr * kr <= {reorder_buffer} + - footprint(B, L1) <= {nb_words_L1} + - footprint(A, L2) <= {nb_words_L2} + - footprint(B, L3) <= {nb_words_L3} +schedule: j: k: B: pack=pack_B pad From e8146f3964ae20908cca59195ea9135b47da749e Mon Sep 17 00:00:00 2001 From: Leon Frenot Date: Tue, 21 Jul 2026 16:16:50 +0200 Subject: [PATCH 20/23] Only allow one PRT tile per line --- src/xtc/schedules/descript.py | 285 ++++++++---------- src/xtc/schedules/parsing.py | 9 +- .../test_matmul_descript_extend_prp.py | 133 ++++---- .../test_matmul_descript_extend_twuop.py | 60 ++-- .../search/test_matmul_descript_pprprp.py | 11 +- ...test_matmul_descript_pprprp_interchange.py | 10 +- .../search/test_matmul_descript_twuop.py | 6 +- 7 files changed, 252 insertions(+), 262 deletions(-) diff --git a/src/xtc/schedules/descript.py b/src/xtc/schedules/descript.py index 94750b17..abde85da 100644 --- a/src/xtc/schedules/descript.py +++ b/src/xtc/schedules/descript.py @@ -143,6 +143,7 @@ def _interpret_spec_into_node( elif isinstance(item, PRTDecl): loop_name = self._interpret_prt( item=item, + loop_name=loop_name, node=node, axes=axes, ) @@ -407,6 +408,7 @@ def _fresh(self, axis: str): def _interpret_prt( self, item: PRTDecl, + loop_name: str, node: ParameterLoopNestNode, axes: dict[str, list[literal]], ) -> str: @@ -419,180 +421,155 @@ def _interpret_prt( """ # Collect annotations + unroll_factor = item.annotations.unroll_factor + vectorize = item.annotations.vectorize parallelize = item.annotations.parallelize - unroll = item.annotations.unroll_specified + buffer = item.annotations.buffer + buffer_specified = item.annotations.buffer_specified + pack = item.annotations.pack + pack_specified = item.annotations.pack_specified + level = item.annotations.level interchange = item.annotations.interchange if interchange == "interchange": interchange = self._fresh("interchange") + if not interchange and item.shape == "U": + interchange = self._fresh("interchange_u") + + def _annotations( + axis: str, + first_p_axes: bool, + last_axis: bool, + ): + unroll_factor_ = unroll_factor + if isinstance(unroll_factor_, str): + unroll_factor_ = self._fresh(unroll_factor_ + "_" + axis) - loop_name: str = "" - - if level and len(item.shape) > 1: - raise ScheduleInterpretError( - "Cannot use the same level on mulitple PRT tiles." - ) - - def _annotations(axis: str) -> Annotations: - if level: - return Annotations( - partial=item.annotations.partial, - full=item.annotations.full, - interchange=interchange, - level=level + axis, - ) return Annotations( + unroll_factor=unroll_factor_, + unroll_specified=item.annotations.unroll_specified, + vectorize=vectorize if last_axis else False, + parallelize=parallelize if first_p_axes else False, + buffer=buffer, + buffer_specified=buffer_specified if last_axis else False, + pack=pack, + pack_specified=pack_specified if last_axis else False, + interchange=interchange, partial=item.annotations.partial, full=item.annotations.full, - interchange=interchange, + level=level + "_" + axis if level else "", ) - def _parallelize(): - if isinstance(parallelize, str): - node.parallelize_parameters[loop_name] = parallelize - node.constraints.append(f"{parallelize} in {{0, 1}}") + def _interpret(axis: str, annotations: Annotations): + decl: AxisDecl | TileDecl + if axis not in node.interchange: + decl = AxisDecl(axis, annotations) + loop_name = self._interpret_axis(decl, node) else: - node.parallelize.append(loop_name) - - def _unroll(axis: str): - unroll_factor = item.annotations.unroll_factor - if isinstance(unroll_factor, int) and unroll_factor <= 0: - raise ScheduleInterpretError( - f'`{{"unroll" = {unroll_factor}}}`: unroll parameter should be strictly positive.' - ) - if unroll_factor is None or isinstance(unroll_factor, str): - if not axes[axis]: + decl = TileDecl(axis, self._fresh(axis), annotations) + loop_name = self._interpret_tile(decl, node, axes) + self._apply_annotations(decl, loop_name, axes, node) + return loop_name + + tile: list[tuple[str, Annotations]] = [] + match item.shape: + case "P": + # All P-axes in order + if not self.abstract_p_dims: raise ScheduleInterpretError( - f"{loop_name}'s size being unknown, an unroll factor is needed." + "P scheme used, but P axes are not specified" ) - axis_size = axes[axis][-1] - if unroll_factor is None: - # None means "unroll fully" - use the loop size - unroll_factor = axis_size - elif isinstance(unroll_factor, str): - unroll_factor += "_prt_" + loop_name - if self.partial_unrolls: - node.constraints.append(f"{unroll_factor} <= {axis_size}") - else: - node.constraints.append(f"{unroll_factor} || {axis_size}") - node.unroll[loop_name] = unroll_factor - - def _interpret(axis: str, annotations: Annotations | None = None): - if not annotations: - annotations = _annotations(axis) - - if axis not in node.interchange: - return self._interpret_axis(AxisDecl(axis, annotations), node) - else: - return self._interpret_tile( - TileDecl(axis, self._fresh(axis), annotations), node, axes + for axis in self.abstract_p_dims[:-1]: + annotations = _annotations(axis, first_p_axes=True, last_axis=False) + tile.append((axis, annotations)) + axis = self.abstract_p_dims[-1] + annotations = _annotations(axis, first_p_axes=True, last_axis=True) + tile.append((axis, annotations)) + + case "R": + # All R-axes in order + if not self.abstract_r_dims: + raise ScheduleInterpretError( + "R scheme used, but R axes are not specified" + ) + for axis in self.abstract_r_dims[:-1]: + annotations = _annotations( + axis, first_p_axes=False, last_axis=False + ) + tile.append((axis, annotations)) + axis = self.abstract_r_dims[-1] + annotations = _annotations(axis, first_p_axes=False, last_axis=True) + tile.append((axis, annotations)) + + case "T" | "U": + # All axes in order/a random order + # Warning, annotations with a randomized order can have unintended comportments for now + # TODO: Fix those + first_p_axes = True + for axis in self.abstract_dims[:-1]: + if ( + first_p_axes + and self.abstract_p_dims + and axis not in self.abstract_p_dims + ): + first_p_axes = False + annotations = _annotations( + axis, first_p_axes=first_p_axes, last_axis=False + ) + tile.append((axis, annotations)) + axis = self.abstract_dims[-1] + if ( + first_p_axes + and self.abstract_p_dims + and axis not in self.abstract_p_dims + ): + first_p_axes = False + annotations = _annotations( + axis, first_p_axes=first_p_axes, last_axis=True ) + tile.append((axis, annotations)) - for idx, t in enumerate(item.shape): - match t: - case "P": - # All P-axes in order - if not self.abstract_p_dims: - raise ScheduleInterpretError( - "P scheme used, but P axes are not specified" - ) - for a in self.abstract_p_dims: - loop_name = _interpret(a) - if parallelize: - _parallelize() - parallelize = False - if unroll and idx + 1 == len(item.shape): - _unroll(a) - case "R": - # All R-axes in order - if not self.abstract_r_dims: - raise ScheduleInterpretError( - "R scheme used, but R axes are not specified" - ) - for a in self.abstract_r_dims: - loop_name = _interpret(a) - if parallelize: - raise ScheduleInterpretError( - "Cannot parallelize a reduction axis" - ) - if unroll and idx + 1 == len(item.shape): - _unroll(a) - case "T": - # All axes in order - for a in self.abstract_dims: - loop_name = _interpret(a) - if parallelize: - _parallelize() - parallelize = False - if unroll and idx + 1 == len(item.shape): - _unroll(a) - case "U": - # All axes in a random order (creates an interchange group) - if level: - raise ScheduleInterpretError("Cannot use level on a U tile.") - annotations_u = Annotations( - partial=item.annotations.partial, - full=item.annotations.full, - interchange=self._fresh("interchange_u"), + case "O": + # first P-axis then R-axes, then remaining P-axis + if not self.abstract_p_dims: + raise ScheduleInterpretError( + "U scheme used, but P axes are not specified" ) - for a in self.abstract_dims: - loop_name = _interpret(a, annotations_u) - if parallelize: - raise ScheduleInterpretError("Cannot parallelize a U tile.") - if unroll and idx + 1 == len(item.shape): - _unroll(a) - case "O": - # first P-axis then R-axes, then remaining P-axis - if not self.abstract_p_dims: - raise ScheduleInterpretError( - "O scheme used, but P axes are not specified" - ) - if not self.abstract_r_dims: - raise ScheduleInterpretError( - "O scheme used, but R axes are not specified" - ) - dims = ( - [self.abstract_p_dims[0]] - + self.abstract_r_dims - + self.abstract_p_dims[1:] + if not self.abstract_r_dims: + raise ScheduleInterpretError( + "U scheme used, but R axes are not specified" ) - for a in dims: - _interpret(a) - if parallelize: - _parallelize() - parallelize = False - if unroll and idx + 1 == len(item.shape): - _unroll(a) - case "W": - # Add an output write buffer at this level - node.buffer_at[loop_name] = None - case "F": - # Fuse producer at this level - raise ScheduleInterpretError("TODO: Fusion unimplemented") - - # Vectorize the last axis - if item.annotations.vectorize: - if isinstance(item.annotations.vectorize, str): - node.vectorize_parameters[loop_name] = item.annotations.vectorize - node.constraints.append(f"{item.annotations.vectorize} in {{0, 1}}") - else: - node.vectorize.append(loop_name) + axis = self.abstract_p_dims[0] + annotations = _annotations( + axis, first_p_axes=True, last_axis=len(self.abstract_dims) > 1 + ) + tile = [(axis, annotations)] + axes_ = self.abstract_r_dims + self.abstract_p_dims[1:] + for axis in axes_[:-1]: + annotations = _annotations( + axis, first_p_axes=False, last_axis=False + ) + tile.append((axis, annotations)) + annotations = _annotations(axis, first_p_axes=False, last_axis=True) + tile.append((axes_[-1], annotations)) - # Write buffer after the last axis - if item.annotations.buffer_specified: - node.buffer_at[loop_name] = item.annotations.buffer + case "W": + # Add an output write buffer at this level + node.buffer_at[loop_name] = None + return loop_name - # Read buffer after the last axis - if item.annotations.pack_specified and item.annotations.pack is not None: - input_matrix, mtype, pad = item.annotations.pack - if isinstance(input_matrix, str): - idx = self.abstract_matrix.index(input_matrix) - if idx == len(self.abstract_matrix) - 1: - node.buffer_at[loop_name] = mtype - else: - node.pack_at[loop_name] = (idx, mtype, pad) - else: - node.pack_at[loop_name] = (input_matrix, mtype, pad) + case "F": + # Fuse producer at this level + raise ScheduleInterpretError("TODO: Fusion unimplemented") + + case _: + # Should be unreachable + raise ScheduleInterpretError( + f"PRTDecl with an invalid name: {item.shape}" + ) + + for axis, annotations in tile: + loop_name = _interpret(axis, annotations) return loop_name diff --git a/src/xtc/schedules/parsing.py b/src/xtc/schedules/parsing.py index ad049f68..13ad9418 100644 --- a/src/xtc/schedules/parsing.py +++ b/src/xtc/schedules/parsing.py @@ -137,10 +137,9 @@ class PRTDecl: annotations: Annotations def __post__init__(self): - for s in self.shape: - if s not in ansor_tile: - # Should be unreachable - raise ScheduleParseError(f"Invalid tile declaration {s}") + if self.shape not in ansor_tile: + # Should be unreachable + raise ScheduleParseError(f"Invalid tile declaration {self.shape}") ScheduleItem = SplitDecl | TileDecl | AxisDecl | PRTDecl @@ -183,7 +182,7 @@ def _parse_declaration(self, declaration: str, value: list) -> ScheduleItem: if "#" in declaration: return self._parse_tile(declaration, value) - if declaration[0] in ansor_tile: + if len(declaration) == 1 and declaration in ansor_tile: return self._parse_ansor_tile(declaration, value) # Must be a direct axis reference diff --git a/tests/filecheck/schedules/test_matmul_descript_extend_prp.py b/tests/filecheck/schedules/test_matmul_descript_extend_prp.py index 9894632d..d2af45ac 100644 --- a/tests/filecheck/schedules/test_matmul_descript_extend_prp.py +++ b/tests/filecheck/schedules/test_matmul_descript_extend_prp.py @@ -27,7 +27,7 @@ strategy = Strategy(graph, spec) -sample = {'prt_i_0': 1, 'u_k_prt_k': 8, 'prt_j_0': 2} +sample = {'prt_i_0': 1, 'prt_u_k_k_0': 8, 'prt_j_0': 2} strategy.generate(sch, sample) @@ -44,69 +44,68 @@ res = executor.execute() print(f"CODE: {res}") -#CHECK:graph: -#CHECK-NEXT: name: matmul -#CHECK-NEXT: inputs: -#CHECK-NEXT: - %0 : 4x512xfloat32 -#CHECK-NEXT: - %1 : 512x32xfloat32 -#CHECK-NEXT: outputs: -#CHECK-NEXT: - %2 : 4x32xfloat32 -#CHECK-NEXT: nodes: -#CHECK-NEXT: - %2: matmul(%0, %1) {name = 'C'} : [4x512xfloat32, 512x32xfloat32] -> [4x32xfloat32] -#CHECK-EMPTY: -#CHECK-NEXT:# from tvm.script import ir as I -#CHECK-NEXT:# from tvm.script import tir as T -#CHECK-EMPTY: -#CHECK-NEXT:@I.ir_module -#CHECK-NEXT:class Module: -#CHECK-NEXT: @T.prim_func -#CHECK-NEXT: def main(_0: T.Buffer((4, 512), "float32"), _1: T.Buffer((512, 32), "float32"), C: T.Buffer((4, 32), "float32")): -#CHECK-NEXT: T.func_attr({"from_legacy_te_schedule": T.bool(True), "tir.noalias": T.bool(True)}) -#CHECK-NEXT: for i, j in T.grid(4, 32): -#CHECK-NEXT: C_1 = T.Buffer((128,), data=C.data) -#CHECK-NEXT: C_1[i * 32 + j] = T.float32(0.0) -#CHECK-NEXT: for k in range(512): -#CHECK-NEXT: cse_var_1: T.int32 = i * 32 + j -#CHECK-NEXT: _0_1 = T.Buffer((2048,), data=_0.data) -#CHECK-NEXT: _1_1 = T.Buffer((16384,), data=_1.data) -#CHECK-NEXT: C_1[cse_var_1] = C_1[cse_var_1] + _0_1[i * 512 + k] * _1_1[k * 32 + j] -#CHECK-NEXT:O = obj['C'] -#CHECK-NEXT:i, j, = O.op.axis -#CHECK-NEXT:k, = O.op.reduce_axis -#CHECK-NEXT:k, __u_k = sch[O].split(k, factor=8) -#CHECK-NEXT:i, i0 = sch[O].split(i, factor=1) -#CHECK-NEXT:j, j0 = sch[O].split(j, factor=2) -#CHECK-NEXT:sch[O].reorder(i, j, k, __u_k, i0, j0) -#CHECK-NEXT:sch[O].unroll(__u_k) -#CHECK-NEXT:sch[O].unroll(i0) -#CHECK-NEXT:sch[O].vectorize(j0) -#CHECK-NEXT:sch[O].parallel(i) -#CHECK-EMPTY: -#CHECK-NEXT:# from tvm.script import ir as I -#CHECK-NEXT:# from tvm.script import tir as T -#CHECK-EMPTY: -#CHECK-NEXT:@I.ir_module -#CHECK-NEXT:class Module: -#CHECK-NEXT: @T.prim_func -#CHECK-NEXT: def main(_0: T.Buffer((4, 512), "float32"), _1: T.Buffer((512, 32), "float32"), C: T.Buffer((4, 32), "float32")): -#CHECK-NEXT: T.func_attr({"from_legacy_te_schedule": T.bool(True), "tir.noalias": T.bool(True)}) -#CHECK-NEXT: for i_outer in T.parallel(4): -#CHECK-NEXT: for j_outer in range(16): -#CHECK-NEXT: C_1 = T.Buffer((128,), data=C.data) -#CHECK-NEXT: C_1[i_outer * 32 + j_outer * 2:i_outer * 32 + j_outer * 2 + 2] = T.Broadcast(T.float32(0.0), 2) -#CHECK-NEXT: for k_outer in range(64): -#CHECK-NEXT: cse_var_4: T.int32 = j_outer * 2 -#CHECK-NEXT: cse_var_3: T.int32 = k_outer * 256 + cse_var_4 -#CHECK-NEXT: cse_var_2: T.int32 = i_outer * 512 + k_outer * 8 -#CHECK-NEXT: cse_var_1: T.int32 = i_outer * 32 + cse_var_4 -#CHECK-NEXT: _0_1 = T.Buffer((2048,), data=_0.data) -#CHECK-NEXT: _1_1 = T.Buffer((16384,), data=_1.data) -#CHECK-NEXT: C_1[cse_var_1:cse_var_1 + 2] = C_1[cse_var_1:cse_var_1 + 2] + T.Broadcast(_0_1[cse_var_2], 2) * _1_1[cse_var_3:cse_var_3 + 2] -#CHECK-NEXT: C_1[cse_var_1:cse_var_1 + 2] = C_1[cse_var_1:cse_var_1 + 2] + T.Broadcast(_0_1[cse_var_2 + 1], 2) * _1_1[cse_var_3 + 32:cse_var_3 + 32 + 2] -#CHECK-NEXT: C_1[cse_var_1:cse_var_1 + 2] = C_1[cse_var_1:cse_var_1 + 2] + T.Broadcast(_0_1[cse_var_2 + 2], 2) * _1_1[cse_var_3 + 64:cse_var_3 + 64 + 2] -#CHECK-NEXT: C_1[cse_var_1:cse_var_1 + 2] = C_1[cse_var_1:cse_var_1 + 2] + T.Broadcast(_0_1[cse_var_2 + 3], 2) * _1_1[cse_var_3 + 96:cse_var_3 + 96 + 2] -#CHECK-NEXT: C_1[cse_var_1:cse_var_1 + 2] = C_1[cse_var_1:cse_var_1 + 2] + T.Broadcast(_0_1[cse_var_2 + 4], 2) * _1_1[cse_var_3 + 128:cse_var_3 + 128 + 2] -#CHECK-NEXT: C_1[cse_var_1:cse_var_1 + 2] = C_1[cse_var_1:cse_var_1 + 2] + T.Broadcast(_0_1[cse_var_2 + 5], 2) * _1_1[cse_var_3 + 160:cse_var_3 + 160 + 2] -#CHECK-NEXT: C_1[cse_var_1:cse_var_1 + 2] = C_1[cse_var_1:cse_var_1 + 2] + T.Broadcast(_0_1[cse_var_2 + 6], 2) * _1_1[cse_var_3 + 192:cse_var_3 + 192 + 2] -#CHECK-NEXT: C_1[cse_var_1:cse_var_1 + 2] = C_1[cse_var_1:cse_var_1 + 2] + T.Broadcast(_0_1[cse_var_2 + 7], 2) * _1_1[cse_var_3 + 224:cse_var_3 + 224 + 2] -#CHECK-NEXT:CODE: 0 +# CHECK: graph: +# CHECK-NEXT: name: matmul +# CHECK-NEXT: inputs: +# CHECK-NEXT: - %0 : 4x512xfloat32 +# CHECK-NEXT: - %1 : 512x32xfloat32 +# CHECK-NEXT: outputs: +# CHECK-NEXT: - %2 : 4x32xfloat32 +# CHECK-NEXT: nodes: +# CHECK-NEXT: - %2: matmul(%0, %1) {name = 'C'} : [4x512xfloat32, 512x32xfloat32] -> [4x32xfloat32] +# CHECK-EMPTY: +# CHECK-NEXT: # from tvm.script import ir as I +# CHECK-NEXT: # from tvm.script import tir as T +# CHECK-EMPTY: +# CHECK-NEXT: @I.ir_module +# CHECK-NEXT: class Module: +# CHECK-NEXT: @T.prim_func +# CHECK-NEXT: def main(_0: T.Buffer((4, 512), "float32"), _1: T.Buffer((512, 32), "float32"), C: T.Buffer((4, 32), "float32")): +# CHECK-NEXT: T.func_attr({"from_legacy_te_schedule": T.bool(True), "tir.noalias": T.bool(True)}) +# CHECK-NEXT: for i, j in T.grid(4, 32): +# CHECK-NEXT: C_1 = T.Buffer((128,), data=C.data) +# CHECK-NEXT: C_1[i * 32 + j] = T.float32(0.0) +# CHECK-NEXT: for k in range(512): +# CHECK-NEXT: cse_var_1: T.int32 = i * 32 + j +# CHECK-NEXT: _0_1 = T.Buffer((2048,), data=_0.data) +# CHECK-NEXT: _1_1 = T.Buffer((16384,), data=_1.data) +# CHECK-NEXT: C_1[cse_var_1] = C_1[cse_var_1] + _0_1[i * 512 + k] * _1_1[k * 32 + j] +# CHECK-NEXT: O = obj['C'] +# CHECK-NEXT: i, j, = O.op.axis +# CHECK-NEXT: k, = O.op.reduce_axis +# CHECK-NEXT: k, __u_k = sch[O].split(k, factor=8) +# CHECK-NEXT: i, i0 = sch[O].split(i, factor=1) +# CHECK-NEXT: j, j0 = sch[O].split(j, factor=2) +# CHECK-NEXT: sch[O].reorder(i, j, k, __u_k, i0, j0) +# CHECK-NEXT: sch[O].unroll(__u_k) +# CHECK-NEXT: sch[O].unroll(i0) +# CHECK-NEXT: sch[O].vectorize(j0) +# CHECK-NEXT: j = sch[O].fuse(i, j) +# CHECK-NEXT: sch[O].parallel(j) +# CHECK-EMPTY: +# CHECK-NEXT: # from tvm.script import ir as I +# CHECK-NEXT: # from tvm.script import tir as T +# CHECK-EMPTY: +# CHECK-NEXT: @I.ir_module +# CHECK-NEXT: class Module: +# CHECK-NEXT: @T.prim_func +# CHECK-NEXT: def main(_0: T.Buffer((4, 512), "float32"), _1: T.Buffer((512, 32), "float32"), C: T.Buffer((4, 32), "float32")): +# CHECK-NEXT: T.func_attr({"from_legacy_te_schedule": T.bool(True), "tir.noalias": T.bool(True)}) +# CHECK-NEXT: for i_outer_j_outer_fused in T.parallel(64): +# CHECK-NEXT: C_1 = T.Buffer((128,), data=C.data) +# CHECK-NEXT: C_1[i_outer_j_outer_fused * 2:i_outer_j_outer_fused * 2 + 2] = T.Broadcast(T.float32(0.0), 2) +# CHECK-NEXT: for k_outer in range(64): +# CHECK-NEXT: cse_var_3: T.int32 = i_outer_j_outer_fused * 2 +# CHECK-NEXT: cse_var_2: T.int32 = k_outer * 256 + i_outer_j_outer_fused % 16 * 2 +# CHECK-NEXT: cse_var_1: T.int32 = i_outer_j_outer_fused // 16 * 512 + k_outer * 8 +# CHECK-NEXT: _0_1 = T.Buffer((2048,), data=_0.data) +# CHECK-NEXT: _1_1 = T.Buffer((16384,), data=_1.data) +# CHECK-NEXT: C_1[cse_var_3:cse_var_3 + 2] = C_1[cse_var_3:cse_var_3 + 2] + T.Broadcast(_0_1[cse_var_1], 2) * _1_1[cse_var_2:cse_var_2 + 2] +# CHECK-NEXT: C_1[cse_var_3:cse_var_3 + 2] = C_1[cse_var_3:cse_var_3 + 2] + T.Broadcast(_0_1[cse_var_1 + 1], 2) * _1_1[cse_var_2 + 32:cse_var_2 + 32 + 2] +# CHECK-NEXT: C_1[cse_var_3:cse_var_3 + 2] = C_1[cse_var_3:cse_var_3 + 2] + T.Broadcast(_0_1[cse_var_1 + 2], 2) * _1_1[cse_var_2 + 64:cse_var_2 + 64 + 2] +# CHECK-NEXT: C_1[cse_var_3:cse_var_3 + 2] = C_1[cse_var_3:cse_var_3 + 2] + T.Broadcast(_0_1[cse_var_1 + 3], 2) * _1_1[cse_var_2 + 96:cse_var_2 + 96 + 2] +# CHECK-NEXT: C_1[cse_var_3:cse_var_3 + 2] = C_1[cse_var_3:cse_var_3 + 2] + T.Broadcast(_0_1[cse_var_1 + 4], 2) * _1_1[cse_var_2 + 128:cse_var_2 + 128 + 2] +# CHECK-NEXT: C_1[cse_var_3:cse_var_3 + 2] = C_1[cse_var_3:cse_var_3 + 2] + T.Broadcast(_0_1[cse_var_1 + 5], 2) * _1_1[cse_var_2 + 160:cse_var_2 + 160 + 2] +# CHECK-NEXT: C_1[cse_var_3:cse_var_3 + 2] = C_1[cse_var_3:cse_var_3 + 2] + T.Broadcast(_0_1[cse_var_1 + 6], 2) * _1_1[cse_var_2 + 192:cse_var_2 + 192 + 2] +# CHECK-NEXT: C_1[cse_var_3:cse_var_3 + 2] = C_1[cse_var_3:cse_var_3 + 2] + T.Broadcast(_0_1[cse_var_1 + 7], 2) * _1_1[cse_var_2 + 224:cse_var_2 + 224 + 2] +# CHECK-NEXT: CODE: 0 diff --git a/tests/filecheck/schedules/test_matmul_descript_extend_twuop.py b/tests/filecheck/schedules/test_matmul_descript_extend_twuop.py index b4c64337..0bac4571 100644 --- a/tests/filecheck/schedules/test_matmul_descript_extend_twuop.py +++ b/tests/filecheck/schedules/test_matmul_descript_extend_twuop.py @@ -21,7 +21,11 @@ sch = impl.get_scheduler() spec = """ -TWUOP: parallelize unroll vectorize +T: parallelize +W: +U: +O: +P: unroll vectorize """ strategy = Strategy(graph, spec) @@ -74,9 +78,9 @@ # CHECK-NEXT: %1 = transform.structured.match attributes {__xtc_id_C_} in %arg0 : (!transform.any_op) -> !transform.any_op # CHECK-NEXT: %tiled_op, %forall_op = transform.structured.tile_using_forall %1 tile_sizes [8, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) # CHECK-NEXT: transform.annotate %forall_op "./i" : !transform.any_op -# CHECK-NEXT: %tiled_linalg_op_2, %loops_3 = transform.structured.tile_using_for %tiled_op tile_sizes [0, 32, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) -# CHECK-NEXT: transform.annotate %loops_3 "./j" : !transform.any_op -# CHECK-NEXT: %tiled_linalg_op_4, %loops_5 = transform.structured.tile_using_for %tiled_linalg_op_2 tile_sizes [0, 0, 8] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: %tiled_op_2, %forall_op_3 = transform.structured.tile_using_forall %tiled_op tile_sizes [0, 32, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %forall_op_3 "./j" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_4, %loops_5 = transform.structured.tile_using_for %tiled_op_2 tile_sizes [0, 0, 8] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) # CHECK-NEXT: transform.annotate %loops_5 "./k" : !transform.any_op # CHECK-NEXT: %tiled_linalg_op_6, %loops_7 = transform.structured.tile_using_for %tiled_linalg_op_4 tile_sizes [4, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) # CHECK-NEXT: transform.annotate %loops_7 "./i0" : !transform.any_op @@ -109,6 +113,7 @@ # CHECK-EMPTY: # CHECK-NEXT: // -----// IR Dump After transform //----- // # CHECK-NEXT: #map = affine_map<(d0) -> (d0 * 8)> +# CHECK-NEXT: #map1 = affine_map<(d0) -> (d0 * 32)> # CHECK-NEXT: module attributes {transform.with_named_sequence} { # CHECK-NEXT: func.func @matmul(%arg0: memref<8x512xf32> {llvm.noalias}, %arg1: memref<512x32xf32> {llvm.noalias}, %arg2: memref<8x32xf32> {llvm.noalias}) { # CHECK-NEXT: %cst = arith.constant dense<0.000000e+00> : vector<1x2xf32> @@ -134,9 +139,10 @@ # CHECK-NEXT: %subview = memref.subview %arg0[%1, 0] [8, 512] [1, 1] : memref<8x512xf32> to memref<8x512xf32, strided<[512, 1], offset: ?>> # CHECK-NEXT: %subview_1 = memref.subview %arg1[0, 0] [512, 32] [1, 1] : memref<512x32xf32> to memref<512x32xf32, strided<[32, 1]>> # CHECK-NEXT: %subview_2 = memref.subview %arg2[%1, 0] [8, 32] [1, 1] : memref<8x32xf32> to memref<8x32xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: scf.for %arg4 = %c0 to %c32 step %c32 { -# CHECK-NEXT: %subview_3 = memref.subview %subview_1[0, %arg4] [512, 32] [1, 1] : memref<512x32xf32, strided<[32, 1]>> to memref<512x32xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: %subview_4 = memref.subview %subview_2[0, %arg4] [8, 32] [1, 1] : memref<8x32xf32, strided<[32, 1], offset: ?>> to memref<8x32xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: scf.forall (%arg4) in (1) { +# CHECK-NEXT: %2 = affine.apply #map1(%arg4) +# CHECK-NEXT: %subview_3 = memref.subview %subview_1[0, %2] [512, 32] [1, 1] : memref<512x32xf32, strided<[32, 1]>> to memref<512x32xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: %subview_4 = memref.subview %subview_2[0, %2] [8, 32] [1, 1] : memref<8x32xf32, strided<[32, 1], offset: ?>> to memref<8x32xf32, strided<[32, 1], offset: ?>> # CHECK-NEXT: scf.for %arg5 = %c0 to %c512 step %c8 { # CHECK-NEXT: %subview_5 = memref.subview %subview[0, %arg5] [8, 8] [1, 1] : memref<8x512xf32, strided<[512, 1], offset: ?>> to memref<8x8xf32, strided<[512, 1], offset: ?>> # CHECK-NEXT: %subview_6 = memref.subview %subview_3[%arg5, 0] [8, 32] [1, 1] : memref<512x32xf32, strided<[32, 1], offset: ?>> to memref<8x32xf32, strided<[32, 1], offset: ?>> @@ -160,28 +166,28 @@ # CHECK-NEXT: %subview_18 = memref.subview %subview_14[0, %arg11] [2, 2] [1, 1] : memref<2x16xf32, strided<[32, 1], offset: ?>> to memref<2x2xf32, strided<[32, 1], offset: ?>> # CHECK-NEXT: %subview_19 = memref.subview %subview_15[%c0, 0] [1, 1] [1, 1] : memref<2x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> # CHECK-NEXT: %subview_20 = memref.subview %subview_18[%c0, 0] [1, 2] [1, 1] : memref<2x2xf32, strided<[32, 1], offset: ?>> to memref<1x2xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: %2 = vector.transfer_read %subview_19[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x1xf32, strided<[512, 1], offset: ?>>, vector<1x1xf32> -# CHECK-NEXT: %3 = vector.transfer_read %subview_17[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x2xf32, strided<[32, 1], offset: ?>>, vector<1x2xf32> -# CHECK-NEXT: %4 = vector.transfer_read %subview_20[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x2xf32, strided<[32, 1], offset: ?>>, vector<1x2xf32> -# CHECK-NEXT: %5 = vector.extract %3[0] : vector<2xf32> from vector<1x2xf32> -# CHECK-NEXT: %6 = vector.extract %2[0, 0] : f32 from vector<1x1xf32> -# CHECK-NEXT: %7 = vector.broadcast %6 : f32 to vector<2xf32> -# CHECK-NEXT: %8 = vector.extract %4[0] : vector<2xf32> from vector<1x2xf32> -# CHECK-NEXT: %9 = vector.fma %7, %5, %8 : vector<2xf32> -# CHECK-NEXT: %10 = vector.insert %9, %cst [0] : vector<2xf32> into vector<1x2xf32> -# CHECK-NEXT: vector.transfer_write %10, %subview_20[%c0, %c0] {in_bounds = [true, true]} : vector<1x2xf32>, memref<1x2xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: %3 = vector.transfer_read %subview_19[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x1xf32, strided<[512, 1], offset: ?>>, vector<1x1xf32> +# CHECK-NEXT: %4 = vector.transfer_read %subview_17[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x2xf32, strided<[32, 1], offset: ?>>, vector<1x2xf32> +# CHECK-NEXT: %5 = vector.transfer_read %subview_20[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x2xf32, strided<[32, 1], offset: ?>>, vector<1x2xf32> +# CHECK-NEXT: %6 = vector.extract %4[0] : vector<2xf32> from vector<1x2xf32> +# CHECK-NEXT: %7 = vector.extract %3[0, 0] : f32 from vector<1x1xf32> +# CHECK-NEXT: %8 = vector.broadcast %7 : f32 to vector<2xf32> +# CHECK-NEXT: %9 = vector.extract %5[0] : vector<2xf32> from vector<1x2xf32> +# CHECK-NEXT: %10 = vector.fma %8, %6, %9 : vector<2xf32> +# CHECK-NEXT: %11 = vector.insert %10, %cst [0] : vector<2xf32> into vector<1x2xf32> +# CHECK-NEXT: vector.transfer_write %11, %subview_20[%c0, %c0] {in_bounds = [true, true]} : vector<1x2xf32>, memref<1x2xf32, strided<[32, 1], offset: ?>> # CHECK-NEXT: %subview_21 = memref.subview %subview_15[%c1, 0] [1, 1] [1, 1] : memref<2x1xf32, strided<[512, 1], offset: ?>> to memref<1x1xf32, strided<[512, 1], offset: ?>> # CHECK-NEXT: %subview_22 = memref.subview %subview_18[%c1, 0] [1, 2] [1, 1] : memref<2x2xf32, strided<[32, 1], offset: ?>> to memref<1x2xf32, strided<[32, 1], offset: ?>> -# CHECK-NEXT: %11 = vector.transfer_read %subview_21[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x1xf32, strided<[512, 1], offset: ?>>, vector<1x1xf32> -# CHECK-NEXT: %12 = vector.transfer_read %subview_17[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x2xf32, strided<[32, 1], offset: ?>>, vector<1x2xf32> -# CHECK-NEXT: %13 = vector.transfer_read %subview_22[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x2xf32, strided<[32, 1], offset: ?>>, vector<1x2xf32> -# CHECK-NEXT: %14 = vector.extract %12[0] : vector<2xf32> from vector<1x2xf32> -# CHECK-NEXT: %15 = vector.extract %11[0, 0] : f32 from vector<1x1xf32> -# CHECK-NEXT: %16 = vector.broadcast %15 : f32 to vector<2xf32> -# CHECK-NEXT: %17 = vector.extract %13[0] : vector<2xf32> from vector<1x2xf32> -# CHECK-NEXT: %18 = vector.fma %16, %14, %17 : vector<2xf32> -# CHECK-NEXT: %19 = vector.insert %18, %cst [0] : vector<2xf32> into vector<1x2xf32> -# CHECK-NEXT: vector.transfer_write %19, %subview_22[%c0, %c0] {in_bounds = [true, true]} : vector<1x2xf32>, memref<1x2xf32, strided<[32, 1], offset: ?>> +# CHECK-NEXT: %12 = vector.transfer_read %subview_21[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x1xf32, strided<[512, 1], offset: ?>>, vector<1x1xf32> +# CHECK-NEXT: %13 = vector.transfer_read %subview_17[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x2xf32, strided<[32, 1], offset: ?>>, vector<1x2xf32> +# CHECK-NEXT: %14 = vector.transfer_read %subview_22[%c0, %c0], %0 {in_bounds = [true, true]} : memref<1x2xf32, strided<[32, 1], offset: ?>>, vector<1x2xf32> +# CHECK-NEXT: %15 = vector.extract %13[0] : vector<2xf32> from vector<1x2xf32> +# CHECK-NEXT: %16 = vector.extract %12[0, 0] : f32 from vector<1x1xf32> +# CHECK-NEXT: %17 = vector.broadcast %16 : f32 to vector<2xf32> +# CHECK-NEXT: %18 = vector.extract %14[0] : vector<2xf32> from vector<1x2xf32> +# CHECK-NEXT: %19 = vector.fma %17, %15, %18 : vector<2xf32> +# CHECK-NEXT: %20 = vector.insert %19, %cst [0] : vector<2xf32> into vector<1x2xf32> +# CHECK-NEXT: vector.transfer_write %20, %subview_22[%c0, %c0] {in_bounds = [true, true]} : vector<1x2xf32>, memref<1x2xf32, strided<[32, 1], offset: ?>> # CHECK-NEXT: } {"./j1"} # CHECK-NEXT: } {"./k1"} # CHECK-NEXT: } {"./i1"} diff --git a/tests/filecheck/search/test_matmul_descript_pprprp.py b/tests/filecheck/search/test_matmul_descript_pprprp.py index cdfa6279..746d88f5 100644 --- a/tests/filecheck/search/test_matmul_descript_pprprp.py +++ b/tests/filecheck/search/test_matmul_descript_pprprp.py @@ -10,13 +10,16 @@ graph = utils.get_graph_matmul() backend = utils.get_backend(graph) spec = """ -PPRP: parallelize -R: unroll=u_k -P: unroll vectorize full""" +- P: parallelize +- P: +- R: +- P: +- R: unroll=u_k +- P: unroll vectorize full""" strategy = Strategy(graph, spec, initialize=False, partial_tiles=True, partial_unrolls=True) print(sorted(strategy._constraints)) print(sum(1 for _ in strategy.sample(100))) -# CHECK: ['prt_i_0 <= 21', 'prt_i_1 <= prt_i_0', 'prt_i_2 || {21, prt_i_0, prt_i_1}', 'prt_j_0 <= 32', 'prt_j_1 <= prt_j_0', 'prt_j_2 || {32, prt_j_0, prt_j_1}', 'prt_k_0 <= 12', 'u_k_prt_k0 <= prt_k_0'] +# CHECK: ['prt_i_0 <= 21', 'prt_i_1 <= prt_i_0', 'prt_i_2 || {21, prt_i_0, prt_i_1}', 'prt_j_0 <= 32', 'prt_j_1 <= prt_j_0', 'prt_j_2 || {32, prt_j_0, prt_j_1}', 'prt_k_0 <= 12', 'prt_u_k_k_0 <= prt_k_0'] # CHECK-NEXT: 100 diff --git a/tests/filecheck/search/test_matmul_descript_pprprp_interchange.py b/tests/filecheck/search/test_matmul_descript_pprprp_interchange.py index f8909e47..332732cf 100644 --- a/tests/filecheck/search/test_matmul_descript_pprprp_interchange.py +++ b/tests/filecheck/search/test_matmul_descript_pprprp_interchange.py @@ -10,13 +10,15 @@ graph = utils.get_graph_matmul() backend = utils.get_backend(graph) spec = """ -- PP: -- RP: interchange -- RP:""" +- P: +- P: +- U: +- R: +- P:""" strategy = Strategy(graph, spec, initialize=False,) print(sorted(strategy._constraints)) print(sum(1 for _ in strategy.sample(100))) -# CHECK: ['1 <= prt_interchange_0 <= 6', 'prt_i_0 || {21}', 'prt_i_1 || {21, prt_i_0}', 'prt_i_2 || {21, prt_i_0, prt_i_1}', 'prt_j_0 || {32}', 'prt_j_1 || {32, prt_j_0}', 'prt_j_2 || {32, prt_j_0, prt_j_1}', 'prt_k_0 || {12}'] +# CHECK: ['1 <= prt_interchange_u_0 <= 6', 'prt_i_0 || {21}', 'prt_i_1 || {21, prt_i_0}', 'prt_i_2 || {21, prt_i_0, prt_i_1}', 'prt_j_0 || {32}', 'prt_j_1 || {32, prt_j_0}', 'prt_j_2 || {32, prt_j_0, prt_j_1}', 'prt_k_0 || {12}'] # CHECK-NEXT: 100 diff --git a/tests/filecheck/search/test_matmul_descript_twuop.py b/tests/filecheck/search/test_matmul_descript_twuop.py index 7719b756..16fd2411 100644 --- a/tests/filecheck/search/test_matmul_descript_twuop.py +++ b/tests/filecheck/search/test_matmul_descript_twuop.py @@ -10,7 +10,11 @@ graph = utils.get_graph_matmul() backend = utils.get_backend(graph) spec = """ -TWUOP: parallelize unroll vectorize +T: parallelize +W: +U: +O: +P: unroll vectorize """ strategy = Strategy(graph, spec, initialize=False, partial_tiles=True, partial_unrolls=True) From 3e44cce6bc42b40e70e4ae25c017b1de6148f76d Mon Sep 17 00:00:00 2001 From: Leon Frenot Date: Tue, 21 Jul 2026 17:04:27 +0200 Subject: [PATCH 21/23] Fix footprint for non-linear access functions --- src/xtc/search/strategies.py | 12 ++++++++- .../search/test_conv_descript_prp.py | 26 +++++++++++++++++++ .../search/test_matmul_descript_yaml_goto.py | 2 +- 3 files changed, 38 insertions(+), 2 deletions(-) create mode 100644 tests/filecheck/search/test_conv_descript_prp.py diff --git a/src/xtc/search/strategies.py b/src/xtc/search/strategies.py index 8b993e21..9e4cde32 100644 --- a/src/xtc/search/strategies.py +++ b/src/xtc/search/strategies.py @@ -1111,7 +1111,17 @@ def _footprints( if level not in levels: raise ValueError(f"Level {level} is not defined in the spec.") level = levels[level] - footprint = " * ".join(str(level[x]) for x in accesses[tensor]) + level_accesses = [] + for access in accesses[tensor]: + if access in level: + level_accesses.append(str(level[access])) + else: + for k, v in level.items(): + access = re.sub(rf"\b{k}\b", f"({v}-1)", access) + access += "+1" + level_accesses.append(access) + + footprint = "*".join(level_accesses) self._constraints[i] = self._FP_PATTERN.sub(footprint, constraint) def _constant_sizes(self) -> Mapping[str, int]: diff --git a/tests/filecheck/search/test_conv_descript_prp.py b/tests/filecheck/search/test_conv_descript_prp.py new file mode 100644 index 00000000..2bc1542a --- /dev/null +++ b/tests/filecheck/search/test_conv_descript_prp.py @@ -0,0 +1,26 @@ +# RUN: python %s 2>&1 | filecheck %s +# REQUIRES: module_xvs +""" +Test strategy Goto on matmul +""" + +import utils +from xtc.search.strategies import Strategy_Descript as Strategy + +graph = utils.get_graph_conv2d() +backend = utils.get_backend(graph) +spec = """ +constraints: + - footprint(A, L_r) > 1 +schedule: + - P: + - R: level=L + - P: +""" +strategy = Strategy(graph, spec, initialize=False) + +print(sorted(strategy._constraints)) +print(sum(1 for _ in strategy.sample(100))) + +# CHECK: ['prt_b_0 || {2}', 'prt_b_0*(prt_h_0-1)*2+(1-1)+1*(prt_w_0-1)*2+(7-1)+1*3 > 1', 'prt_f_0 || {32}', 'prt_h_0 || {2}', 'prt_w_0 || {2}'] +# CHECK-NEXT: 48 diff --git a/tests/filecheck/search/test_matmul_descript_yaml_goto.py b/tests/filecheck/search/test_matmul_descript_yaml_goto.py index fc777d73..ac06fd6b 100644 --- a/tests/filecheck/search/test_matmul_descript_yaml_goto.py +++ b/tests/filecheck/search/test_matmul_descript_yaml_goto.py @@ -73,5 +73,5 @@ def fn_ilp(nvr, mr, mc, nc, kc): print(sorted(strategy._constraints)) print(sum(1 for _ in strategy.sample(100))) -# CHECK: ['1 + nvr + nvr * mr <= 32', 'ilp(nvr, mr, mc, nc, kc) >= 8', 'kc * nc <= 9437184', 'kc * nr <= 8192', 'kc <= 1024', 'kr <= kc', 'mc * kc <= 262144', 'mc <= 1024', 'mr || {1024, mc}', 'nc <= 1024', 'nr == 16 * nvr', 'nr || {1024, nc}', 'nvr * mr * kr <= 256', 'pack_B in {0, 1}', 'pad_A in {0,1}'] +# CHECK: ['1 + nvr + nvr * mr <= 32', 'ilp(nvr, mr, mc, nc, kc) >= 8', 'kc <= 1024', 'kc*nc <= 9437184', 'kc*nr <= 8192', 'kr <= kc', 'mc <= 1024', 'mc*kc <= 262144', 'mr || {1024, mc}', 'nc <= 1024', 'nr == 16 * nvr', 'nr || {1024, nc}', 'nvr * mr * kr <= 256', 'pack_B in {0, 1}', 'pad_A in {0,1}'] # CHECK-NEXT: 100 From 4032812c460b8cf1c4fc7671478f4571fcdaf511 Mon Sep 17 00:00:00 2001 From: Leon Frenot Date: Wed, 22 Jul 2026 11:34:16 +0200 Subject: [PATCH 22/23] Fix footprint regex --- src/xtc/search/strategies.py | 13 +++++++++---- .../search/test_matmul_descript_yaml_goto.py | 4 ++-- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/src/xtc/search/strategies.py b/src/xtc/search/strategies.py index 9e4cde32..a89b9893 100644 --- a/src/xtc/search/strategies.py +++ b/src/xtc/search/strategies.py @@ -1103,9 +1103,10 @@ def _footprints( if not self._loop_nest.root_node: return levels = self._loop_nest.root_node._levels - for i, constraint in enumerate(self._constraints): - match = self._FP_PATTERN.match(constraint) - if match: + for i in range(len(self._constraints)): + constraint = self._constraints[i] + match = self._FP_PATTERN.search(constraint) + while match: tensor, level = match.groups() tensor = tensors.index(tensor) if level not in levels: @@ -1122,7 +1123,11 @@ def _footprints( level_accesses.append(access) footprint = "*".join(level_accesses) - self._constraints[i] = self._FP_PATTERN.sub(footprint, constraint) + self._constraints[i] = self._FP_PATTERN.sub( + footprint, constraint, count=1 + ) + constraint = self._constraints[i] + match = self._FP_PATTERN.search(constraint) def _constant_sizes(self) -> Mapping[str, int]: sizes = {a: v for a, v in self._op.dims.items() if isinstance(v, int)} diff --git a/tests/filecheck/search/test_matmul_descript_yaml_goto.py b/tests/filecheck/search/test_matmul_descript_yaml_goto.py index ac06fd6b..cd3b36ad 100644 --- a/tests/filecheck/search/test_matmul_descript_yaml_goto.py +++ b/tests/filecheck/search/test_matmul_descript_yaml_goto.py @@ -38,7 +38,7 @@ - nr == {vector_size} * nvr - ilp(nvr, mr, mc, nc, kc) >= {ilp} - nvr * mr * kr <= {reorder_buffer} - - footprint(B, L1) <= {nb_words_L1} + - footprint(B, L1) + footprint(A, L1) <= {nb_words_L1} - footprint(A, L2) <= {nb_words_L2} - footprint(B, L3) <= {nb_words_L3} schedule: @@ -73,5 +73,5 @@ def fn_ilp(nvr, mr, mc, nc, kc): print(sorted(strategy._constraints)) print(sum(1 for _ in strategy.sample(100))) -# CHECK: ['1 + nvr + nvr * mr <= 32', 'ilp(nvr, mr, mc, nc, kc) >= 8', 'kc <= 1024', 'kc*nc <= 9437184', 'kc*nr <= 8192', 'kr <= kc', 'mc <= 1024', 'mc*kc <= 262144', 'mr || {1024, mc}', 'nc <= 1024', 'nr == 16 * nvr', 'nr || {1024, nc}', 'nvr * mr * kr <= 256', 'pack_B in {0, 1}', 'pad_A in {0,1}'] +# CHECK: ['1 + nvr + nvr * mr <= 32', 'ilp(nvr, mr, mc, nc, kc) >= 8', 'kc <= 1024', 'kc*nc <= 9437184', 'kc*nr + mr*kc <= 8192', 'kr <= kc', 'mc <= 1024', 'mc*kc <= 262144', 'mr || {1024, mc}', 'nc <= 1024', 'nr == 16 * nvr', 'nr || {1024, nc}', 'nvr * mr * kr <= 256', 'pack_B in {0, 1}', 'pad_A in {0,1}'] # CHECK-NEXT: 100 From e5f722da6de147bfc6e2f59ddc1200ae6bf34df9 Mon Sep 17 00:00:00 2001 From: Leon Frenot Date: Wed, 22 Jul 2026 14:44:28 +0200 Subject: [PATCH 23/23] Change descript tests to have one result per line --- .../search/test_conv_descript_prp.py | 9 +++++++-- .../search/test_matmul_descript_3axes.py | 7 +++++-- .../search/test_matmul_descript_goto.py | 13 +++++++++++-- .../test_matmul_descript_interchange.py | 8 ++++++-- .../search/test_matmul_descript_pprprp.py | 12 ++++++++++-- ...test_matmul_descript_pprprp_interchange.py | 12 ++++++++++-- .../search/test_matmul_descript_simple.py | 7 +++++-- .../search/test_matmul_descript_split.py | 10 ++++++++-- .../test_matmul_descript_split_in_split.py | 15 +++++++++++++-- .../search/test_matmul_descript_twuop.py | 13 +++++++++++-- .../search/test_matmul_descript_yaml_goto.py | 19 +++++++++++++++++-- .../test_matmul_descript_yaml_simple.py | 7 +++++-- .../search/test_matmul_descript_yaml_split.py | 15 +++++++++++++-- 13 files changed, 121 insertions(+), 26 deletions(-) diff --git a/tests/filecheck/search/test_conv_descript_prp.py b/tests/filecheck/search/test_conv_descript_prp.py index 2bc1542a..e848f68c 100644 --- a/tests/filecheck/search/test_conv_descript_prp.py +++ b/tests/filecheck/search/test_conv_descript_prp.py @@ -19,8 +19,13 @@ """ strategy = Strategy(graph, spec, initialize=False) -print(sorted(strategy._constraints)) +for x in sorted(strategy._constraints): + print(x) print(sum(1 for _ in strategy.sample(100))) -# CHECK: ['prt_b_0 || {2}', 'prt_b_0*(prt_h_0-1)*2+(1-1)+1*(prt_w_0-1)*2+(7-1)+1*3 > 1', 'prt_f_0 || {32}', 'prt_h_0 || {2}', 'prt_w_0 || {2}'] +# CHECK: prt_b_0 || {2} +# CHECK-NEXT: prt_b_0*(prt_h_0-1)*2+(1-1)+1*(prt_w_0-1)*2+(7-1)+1*3 > 1 +# CHECK-NEXT: prt_f_0 || {32} +# CHECK-NEXT: prt_h_0 || {2} +# CHECK-NEXT: prt_w_0 || {2} # CHECK-NEXT: 48 diff --git a/tests/filecheck/search/test_matmul_descript_3axes.py b/tests/filecheck/search/test_matmul_descript_3axes.py index 6e484343..0356dcb1 100644 --- a/tests/filecheck/search/test_matmul_descript_3axes.py +++ b/tests/filecheck/search/test_matmul_descript_3axes.py @@ -19,8 +19,11 @@ } strategy = Strategy(graph, spec, initialize=False) -print(sorted(strategy._constraints)) +for x in sorted(strategy._constraints): + print(x) print(sum(1 for _ in strategy.sample(100))) -# CHECK:['iR || {21}', 'jR || {32}', 'kR || {12}'] +# CHECK: iR || {21} +# CHECK-NEXT: jR || {32} +# CHECK-NEXT: kR || {12} # CHECK-NEXT:100 diff --git a/tests/filecheck/search/test_matmul_descript_goto.py b/tests/filecheck/search/test_matmul_descript_goto.py index 6e6597f8..998cda60 100644 --- a/tests/filecheck/search/test_matmul_descript_goto.py +++ b/tests/filecheck/search/test_matmul_descript_goto.py @@ -21,8 +21,17 @@ constraint = ["iR * jR <= 56"] strategy = Strategy(graph, spec, constraints=constraint, initialize=False) -print(sorted(strategy._constraints)) +for x in sorted(strategy._constraints): + print(x) print(sum(1 for _ in strategy.sample(100))) -# CHECK: ['iL2 || {21}', 'iR * jR <= 56', 'iR || {21, iL2}', 'jL3 || {32}', 'jR || {32, jL3}', 'j_parallel in {0, 1}', 'j_vectorise in {0, 1}', 'kL1 || {12}', 'k_unroll || kL1'] +# CHECK: iL2 || {21} +# CHECK-NEXT: iR * jR <= 56 +# CHECK-NEXT: iR || {21, iL2} +# CHECK-NEXT: jL3 || {32} +# CHECK-NEXT: jR || {32, jL3} +# CHECK-NEXT: j_parallel in {0, 1} +# CHECK-NEXT: j_vectorise in {0, 1} +# CHECK-NEXT: kL1 || {12} +# CHECK-NEXT: k_unroll || kL1 # CHECK-NEXT: 100 diff --git a/tests/filecheck/search/test_matmul_descript_interchange.py b/tests/filecheck/search/test_matmul_descript_interchange.py index c240538d..dc7d1f00 100644 --- a/tests/filecheck/search/test_matmul_descript_interchange.py +++ b/tests/filecheck/search/test_matmul_descript_interchange.py @@ -20,8 +20,12 @@ strategy = Strategy(graph, spec, initialize=False) -print(sorted(strategy._constraints)) +for x in sorted(strategy._constraints): + print(x) print(sum(1 for _ in strategy.sample(100))) -# CHECK: ['1 <= int <= 2', 'i1 || {21}', 'j1 || {32}', 'j2 || {32, j1}'] +# CHECK: 1 <= int <= 2 +# CHECK-NEXT: i1 || {21} +# CHECK-NEXT: j1 || {32} +# CHECK-NEXT: j2 || {32, j1} # CHECK-NEXT: 100 diff --git a/tests/filecheck/search/test_matmul_descript_pprprp.py b/tests/filecheck/search/test_matmul_descript_pprprp.py index 746d88f5..45b60a9b 100644 --- a/tests/filecheck/search/test_matmul_descript_pprprp.py +++ b/tests/filecheck/search/test_matmul_descript_pprprp.py @@ -18,8 +18,16 @@ - P: unroll vectorize full""" strategy = Strategy(graph, spec, initialize=False, partial_tiles=True, partial_unrolls=True) -print(sorted(strategy._constraints)) +for x in sorted(strategy._constraints): + print(x) print(sum(1 for _ in strategy.sample(100))) -# CHECK: ['prt_i_0 <= 21', 'prt_i_1 <= prt_i_0', 'prt_i_2 || {21, prt_i_0, prt_i_1}', 'prt_j_0 <= 32', 'prt_j_1 <= prt_j_0', 'prt_j_2 || {32, prt_j_0, prt_j_1}', 'prt_k_0 <= 12', 'prt_u_k_k_0 <= prt_k_0'] +# CHECK: prt_i_0 <= 21 +# CHECK-NEXT: prt_i_1 <= prt_i_0 +# CHECK-NEXT: prt_i_2 || {21, prt_i_0, prt_i_1} +# CHECK-NEXT: prt_j_0 <= 32 +# CHECK-NEXT: prt_j_1 <= prt_j_0 +# CHECK-NEXT: prt_j_2 || {32, prt_j_0, prt_j_1} +# CHECK-NEXT: prt_k_0 <= 12 +# CHECK-NEXT: prt_u_k_k_0 <= prt_k_0 # CHECK-NEXT: 100 diff --git a/tests/filecheck/search/test_matmul_descript_pprprp_interchange.py b/tests/filecheck/search/test_matmul_descript_pprprp_interchange.py index 332732cf..c417a550 100644 --- a/tests/filecheck/search/test_matmul_descript_pprprp_interchange.py +++ b/tests/filecheck/search/test_matmul_descript_pprprp_interchange.py @@ -17,8 +17,16 @@ - P:""" strategy = Strategy(graph, spec, initialize=False,) -print(sorted(strategy._constraints)) +for x in sorted(strategy._constraints): + print(x) print(sum(1 for _ in strategy.sample(100))) -# CHECK: ['1 <= prt_interchange_u_0 <= 6', 'prt_i_0 || {21}', 'prt_i_1 || {21, prt_i_0}', 'prt_i_2 || {21, prt_i_0, prt_i_1}', 'prt_j_0 || {32}', 'prt_j_1 || {32, prt_j_0}', 'prt_j_2 || {32, prt_j_0, prt_j_1}', 'prt_k_0 || {12}'] +# CHECK: 1 <= prt_interchange_u_0 <= 6 +# CHECK-NEXT: prt_i_0 || {21} +# CHECK-NEXT: prt_i_1 || {21, prt_i_0} +# CHECK-NEXT: prt_i_2 || {21, prt_i_0, prt_i_1} +# CHECK-NEXT: prt_j_0 || {32} +# CHECK-NEXT: prt_j_1 || {32, prt_j_0} +# CHECK-NEXT: prt_j_2 || {32, prt_j_0, prt_j_1} +# CHECK-NEXT: prt_k_0 || {12} # CHECK-NEXT: 100 diff --git a/tests/filecheck/search/test_matmul_descript_simple.py b/tests/filecheck/search/test_matmul_descript_simple.py index c67b7835..30dc07a5 100644 --- a/tests/filecheck/search/test_matmul_descript_simple.py +++ b/tests/filecheck/search/test_matmul_descript_simple.py @@ -20,8 +20,11 @@ strategy = Strategy(graph, spec, initialize=False) -print(sorted(strategy._constraints)) +for x in sorted(strategy._constraints): + print(x) print(sum(1 for _ in strategy.sample(100))) -# CHECK: ['i1 || {21}', 'j1 || {32}', 'j2 || {32, j1}'] +# CHECK: i1 || {21} +# CHECK-NEXT: j1 || {32} +# CHECK-NEXT: j2 || {32, j1} # CHECK-NEXT: 84 diff --git a/tests/filecheck/search/test_matmul_descript_split.py b/tests/filecheck/search/test_matmul_descript_split.py index 72b6fe55..be16e417 100644 --- a/tests/filecheck/search/test_matmul_descript_split.py +++ b/tests/filecheck/search/test_matmul_descript_split.py @@ -27,8 +27,14 @@ } strategy = Strategy(graph, spec, initialize=False) -print(sorted(strategy._constraints)) +for x in sorted(strategy._constraints): + print(x) print(sum(1 for _ in strategy.sample(100))) -# CHECK: ['iL3 || {21}', 'iR1 || {21, iL3, 7}', 'iR2 || {21, iL3, 7}', 'jDDR || {32}', 'jR1 || {32, jDDR}', 'jR2 || {32, jDDR}'] +# CHECK: iL3 || {21} +# CHECK-NEXT: iR1 || {21, iL3, 7} +# CHECK-NEXT: iR2 || {21, iL3, 7} +# CHECK-NEXT: jDDR || {32} +# CHECK-NEXT: jR1 || {32, jDDR} +# CHECK-NEXT: jR2 || {32, jDDR} # CHECK-NEXT: 100 diff --git a/tests/filecheck/search/test_matmul_descript_split_in_split.py b/tests/filecheck/search/test_matmul_descript_split_in_split.py index bc0ec0c1..fb0d640b 100644 --- a/tests/filecheck/search/test_matmul_descript_split_in_split.py +++ b/tests/filecheck/search/test_matmul_descript_split_in_split.py @@ -30,8 +30,19 @@ } strategy = Strategy(graph, spec, initialize=False) -print(sorted(strategy._constraints)) +for x in sorted(strategy._constraints): + print(x) print(sum(1 for _ in strategy.sample(100))) -# CHECK: ['iL2 || {21}', 'iR1 || {21, iL2, 3}', 'iR2 || {21, iL2}', 'iR3 || {21, iL2, 3}', 'iS + 2 == 3', 'i_1_ + 6 == iL2', 'i_1_ <= iL2', 'jDDR || {32}', 'jR1 || {32, jDDR}', 'jR2 || {32, jDDR}', 'jR3 || {32, jDDR}'] +# CHECK: iL2 || {21} +# CHECK-NEXT: iR1 || {21, iL2, 3} +# CHECK-NEXT: iR2 || {21, iL2} +# CHECK-NEXT: iR3 || {21, iL2, 3} +# CHECK-NEXT: iS + 2 == 3 +# CHECK-NEXT: i_1_ + 6 == iL2 +# CHECK-NEXT: i_1_ <= iL2 +# CHECK-NEXT: jDDR || {32} +# CHECK-NEXT: jR1 || {32, jDDR} +# CHECK-NEXT: jR2 || {32, jDDR} +# CHECK-NEXT: jR3 || {32, jDDR} # CHECK-NEXT: 100 diff --git a/tests/filecheck/search/test_matmul_descript_twuop.py b/tests/filecheck/search/test_matmul_descript_twuop.py index 16fd2411..e383621e 100644 --- a/tests/filecheck/search/test_matmul_descript_twuop.py +++ b/tests/filecheck/search/test_matmul_descript_twuop.py @@ -18,8 +18,17 @@ """ strategy = Strategy(graph, spec, initialize=False, partial_tiles=True, partial_unrolls=True) -print(sorted(strategy._constraints)) +for x in sorted(strategy._constraints): + print(x) print(sum(1 for _ in strategy.sample(100))) -# CHECK: ['1 <= prt_interchange_u_0 <= 6', 'prt_i_0 <= 21', 'prt_i_1 <= prt_i_0', 'prt_i_2 <= prt_i_1', 'prt_j_0 <= 32', 'prt_j_1 <= prt_j_0', 'prt_j_2 <= prt_j_1', 'prt_k_0 <= 12', 'prt_k_1 <= prt_k_0'] +# CHECK: 1 <= prt_interchange_u_0 <= 6 +# CHECK-NEXT: prt_i_0 <= 21 +# CHECK-NEXT: prt_i_1 <= prt_i_0 +# CHECK-NEXT: prt_i_2 <= prt_i_1 +# CHECK-NEXT: prt_j_0 <= 32 +# CHECK-NEXT: prt_j_1 <= prt_j_0 +# CHECK-NEXT: prt_j_2 <= prt_j_1 +# CHECK-NEXT: prt_k_0 <= 12 +# CHECK-NEXT: prt_k_1 <= prt_k_0 # CHECK-NEXT: 100 diff --git a/tests/filecheck/search/test_matmul_descript_yaml_goto.py b/tests/filecheck/search/test_matmul_descript_yaml_goto.py index cd3b36ad..fc240710 100644 --- a/tests/filecheck/search/test_matmul_descript_yaml_goto.py +++ b/tests/filecheck/search/test_matmul_descript_yaml_goto.py @@ -70,8 +70,23 @@ def fn_ilp(nvr, mr, mc, nc, kc): initialize=False, ) -print(sorted(strategy._constraints)) +for x in sorted(strategy._constraints): + print(x) print(sum(1 for _ in strategy.sample(100))) -# CHECK: ['1 + nvr + nvr * mr <= 32', 'ilp(nvr, mr, mc, nc, kc) >= 8', 'kc <= 1024', 'kc*nc <= 9437184', 'kc*nr + mr*kc <= 8192', 'kr <= kc', 'mc <= 1024', 'mc*kc <= 262144', 'mr || {1024, mc}', 'nc <= 1024', 'nr == 16 * nvr', 'nr || {1024, nc}', 'nvr * mr * kr <= 256', 'pack_B in {0, 1}', 'pad_A in {0,1}'] +# CHECK: 1 + nvr + nvr * mr <= 32 +# CHECK-NEXT: ilp(nvr, mr, mc, nc, kc) >= 8 +# CHECK-NEXT: kc <= 1024 +# CHECK-NEXT: kc*nc <= 9437184 +# CHECK-NEXT: kc*nr + mr*kc <= 8192 +# CHECK-NEXT: kr <= kc +# CHECK-NEXT: mc <= 1024 +# CHECK-NEXT: mc*kc <= 262144 +# CHECK-NEXT: mr || {1024, mc} +# CHECK-NEXT: nc <= 1024 +# CHECK-NEXT: nr == 16 * nvr +# CHECK-NEXT: nr || {1024, nc} +# CHECK-NEXT: nvr * mr * kr <= 256 +# CHECK-NEXT: pack_B in {0, 1} +# CHECK-NEXT: pad_A in {0,1} # CHECK-NEXT: 100 diff --git a/tests/filecheck/search/test_matmul_descript_yaml_simple.py b/tests/filecheck/search/test_matmul_descript_yaml_simple.py index 031f3aa5..c704dda8 100644 --- a/tests/filecheck/search/test_matmul_descript_yaml_simple.py +++ b/tests/filecheck/search/test_matmul_descript_yaml_simple.py @@ -19,8 +19,11 @@ """ strategy = Strategy(graph, spec) -print(sorted(strategy._constraints)) +for x in sorted(strategy._constraints): + print(x) print(sum(1 for _ in strategy.sample(100))) -# CHECK: ['i1 || {21}', 'j1 || {32}', 'j2 || {32, j1}'] +# CHECK: i1 || {21} +# CHECK-NEXT: j1 || {32} +# CHECK-NEXT: j2 || {32, j1} # CHECK-NEXT: 84 diff --git a/tests/filecheck/search/test_matmul_descript_yaml_split.py b/tests/filecheck/search/test_matmul_descript_yaml_split.py index b493bbf1..f8dec437 100644 --- a/tests/filecheck/search/test_matmul_descript_yaml_split.py +++ b/tests/filecheck/search/test_matmul_descript_yaml_split.py @@ -26,8 +26,19 @@ """ strategy = Strategy(graph, spec) -print(sorted(strategy._constraints)) +for x in sorted(strategy._constraints): + print(x) print(sum(1 for _ in strategy.sample(100))) -# CHECK: ['SR || {12}', 'iL2 || {21, iL3}', 'iL3 || {21}', 'iR1 || {21, iL3, iL2}', 'iR2 || {21, iL3, iL2}', 'iS <= iL2', 'i_1_ + iS == iL2', 'i_1_ <= iL2', 'jDDR || {32}', 'jR1 || {32, jDDR}', 'jR2 || {32, jDDR}'] +# CHECK: SR || {12} +# CHECK-NEXT: iL2 || {21, iL3} +# CHECK-NEXT: iL3 || {21} +# CHECK-NEXT: iR1 || {21, iL3, iL2} +# CHECK-NEXT: iR2 || {21, iL3, iL2} +# CHECK-NEXT: iS <= iL2 +# CHECK-NEXT: i_1_ + iS == iL2 +# CHECK-NEXT: i_1_ <= iL2 +# CHECK-NEXT: jDDR || {32} +# CHECK-NEXT: jR1 || {32, jDDR} +# CHECK-NEXT: jR2 || {32, jDDR} # CHECK-NEXT: 100