diff --git a/src/xtc/schedules/descript.py b/src/xtc/schedules/descript.py index d1100508..abde85da 100644 --- a/src/xtc/schedules/descript.py +++ b/src/xtc/schedules/descript.py @@ -2,7 +2,8 @@ # SPDX-License-Identifier: BSD-3-Clause # Copyright (c) 2024-2026 The XTC Project Authors # -from typing import Any, cast +from math import factorial +from typing import Any from dataclasses import dataclass, field from copy import deepcopy @@ -15,16 +16,20 @@ ParameterSplitOrigin, ) -from .exceptions import ScheduleInterpretError +from .exceptions import ScheduleInterpretError, ScheduleParseError from .parsing import ( ScheduleParser, ScheduleSpec, SplitDecl, TileDecl, AxisDecl, + PRTDecl, Annotations, YAMLParser, literal, + ansor_tile, + tup_list, + pre_parse, ) @@ -66,15 +71,24 @@ 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) + _levels: dict[str, dict[str, literal | None]] = field(default_factory=dict) def interpret(self, spec: ScheduleSpec, root: str) -> ParameterLoopNest: """Interpret a schedule specification into a LoopNest.""" 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( @@ -86,17 +100,19 @@ 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} - interchange: list[str] = list(head) + # Track finak split and corresponding missing size last_split: list[tuple[literal, literal | None]] = [] - sizes: dict[str, literal] = {} + # Name of the last loop nest (declared here for typing/liveness) loop_name: str = "" + + node.interchange = list(head) + 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: @@ -108,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, @@ -117,18 +132,37 @@ 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) + self._apply_annotations(item, loop_name, axes, 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) + self._apply_annotations(item, loop_name, axes, node) + elif isinstance(item, PRTDecl): + loop_name = self._interpret_prt( + item=item, + loop_name=loop_name, + node=node, + axes=axes, + ) - # Reaplace the placeholder of the last split with its size + # 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(): + 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))}") + + # 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): @@ -139,23 +173,16 @@ 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})." ) - 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]], @@ -187,23 +214,23 @@ 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 # 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 + # i[x:] => i[x:end] 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: + + # 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 = ( inner_size.replace(_SPLIT_LEFT, "_") @@ -214,14 +241,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: - inner_size = z - x = cut + # Split from x, of size z, that fits in y y = current_size - 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) + + # 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): @@ -252,6 +282,10 @@ 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) + # Recursively interpret the nested schedule into the child node self._interpret_spec_into_node( spec=item.body, @@ -265,29 +299,44 @@ def _interpret_tile( self, item: TileDecl, node: ParameterLoopNestNode, - interchange: list[str], - sizes: dict[str, literal], axes: dict[str, list[literal]], ) -> 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 - sizes[loop_name] = item.size - interchange.append(loop_name) + + # 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: + 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) + + # Alias for axes[item.axis] + list_axis: list[int | str] if item.axis in axes: list_axis = axes[item.axis] 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( @@ -295,31 +344,35 @@ 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( 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) + + # Update axes list_axis.append(item.size) + self._update_levels(item.axis, axes[item.axis][-1]) return loop_name 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 + + # 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 @@ -329,9 +382,197 @@ def _interpret_axis( f"Axis {axis_name} is scheduled twice (or more)." ) - interchange.append(axis_name) + # 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: + 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) + + 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, + loop_name: str, + 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 + unroll_factor = item.annotations.unroll_factor + vectorize = item.annotations.vectorize + parallelize = item.annotations.parallelize + 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) + + 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, + level=level + "_" + axis if level else "", + ) + + 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: + 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( + "P scheme used, but P axes are not specified" + ) + 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)) + + 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" + ) + if not self.abstract_r_dims: + raise ScheduleInterpretError( + "U scheme used, but R axes are not specified" + ) + 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)) + + case "W": + # Add an output write buffer at this level + node.buffer_at[loop_name] = None + return loop_name + + 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 + def _check_axis_existence(self, axis: str) -> None: """Check that an axis is defined.""" if axis not in self.abstract_dims: @@ -341,30 +582,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: @@ -381,6 +626,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 @@ -388,12 +639,36 @@ 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): + node.buffer_parameters[loop_name] = annotations.pack_specified + 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): + 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 _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, @@ -443,25 +718,38 @@ 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 loop_nest(self, node_name: str, spec: dict[str, dict[str, Any]] | str): + 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]] | tup_list | 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", [])) + 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_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/parameter_loop_nest.py b/src/xtc/schedules/parameter_loop_nest.py index ee1eba5b..45454117 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 @@ -11,8 +12,6 @@ NodeT = TypeVar("NodeT", bound="Node[Any]") -from .exceptions import ScheduleValidationError - literal = int | str @@ -114,20 +113,25 @@ 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) 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) + _levels: dict[str, dict[str, literal]] = field(default_factory=dict) 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 +139,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 +147,50 @@ 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(): 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, @@ -426,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. @@ -448,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. - """ - ) diff --git a/src/xtc/schedules/parsing.py b/src/xtc/schedules/parsing.py index bbea72af..13ad9418 100644 --- a/src/xtc/schedules/parsing.py +++ b/src/xtc/schedules/parsing.py @@ -7,12 +7,16 @@ 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]] +CONSTRAINTS = "constraints" +SCHEDULE = "schedule" def toliteral(s: str) -> literal: @@ -22,6 +26,36 @@ def toliteral(s: str) -> literal: return s +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()] + 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. @@ -45,9 +79,11 @@ 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 + interchange: str = "" + level: str = "" partial: bool = False full: bool = False @@ -93,7 +129,20 @@ 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): + if self.shape not in ansor_tile: + # Should be unreachable + raise ScheduleParseError(f"Invalid tile declaration {self.shape}") + + +ScheduleItem = SplitDecl | TileDecl | AxisDecl | PRTDecl @dataclass(frozen=True) @@ -110,19 +159,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) @@ -131,17 +182,20 @@ def _parse_declaration(self, declaration: str, value: Any) -> ScheduleItem: if "#" in declaration: return self._parse_tile(declaration, value) + if len(declaration) == 1 and declaration in ansor_tile: + return self._parse_ansor_tile(declaration, value) + # 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: @@ -156,13 +210,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_axis_ref(self, declaration: str, value: dict) -> AxisDecl: + 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: 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]], declaration: str + ) -> Annotations: """Parse annotation dict into Annotations object.""" unroll_factor: literal | None = None @@ -170,13 +232,15 @@ def _parse_annotations(self, value: dict[str, Any], context: str) -> 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 + interchange: str = "" + level: str = "" 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: @@ -215,17 +279,43 @@ def _parse_annotations(self, value: dict[str, Any], context: str) -> 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 "interchange": + if param is None: + 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": 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, @@ -236,38 +326,40 @@ def _parse_annotations(self, value: dict[str, Any], context: str) -> Annotations buffer_specified=buffer_specified, pack=pack, pack_specified=pack_specified, + interchange=interchange, + level=level, partial=partial, full=full, ) 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 @@ -302,51 +394,54 @@ 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) -> 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(): + 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 diff --git a/src/xtc/search/strategies.py b/src/xtc/search/strategies.py index e32caa54..a89b9893 100644 --- a/src/xtc/search/strategies.py +++ b/src/xtc/search/strategies.py @@ -4,11 +4,12 @@ # 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 import numpy as np +import re from xtc.itf.graph import Graph from xtc.itf.schd import Scheduler @@ -1029,25 +1030,55 @@ 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, spec: dict[str, dict[str, Any]] | str, - constraints: list[str] = [], + 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 = list(self._op.dims) + 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 = [ + 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, + 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, ) @@ -1056,11 +1087,48 @@ 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._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 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: + raise ValueError(f"Level {level} is not defined in the spec.") + level = levels[level] + 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, 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)} return sizes @@ -1072,7 +1140,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 ) @@ -1158,14 +1226,25 @@ def __init__( self, graph: Graph, spec: dict[str, dict[str, Any]] | str, - constraints: list[str] = [], + 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._sample_shape: list[str] = [] super().__init__( - graph, spec, constraints, partial_tiles, partial_unrolls, initialize + graph, + spec, + constraints, + axes, + tensors, + functions, + partial_tiles, + partial_unrolls, + initialize, ) @override @@ -1179,7 +1258,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) 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/schedules/test_matmul_descript_extend_prp.py b/tests/filecheck/schedules/test_matmul_descript_extend_prp.py new file mode 100644 index 00000000..d2af45ac --- /dev/null +++ b/tests/filecheck/schedules/test_matmul_descript_extend_prp.py @@ -0,0 +1,111 @@ +# 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, 'prt_u_k_k_0': 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: 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 new file mode 100644 index 00000000..0bac4571 --- /dev/null +++ b/tests/filecheck/schedules/test_matmul_descript_extend_twuop.py @@ -0,0 +1,214 @@ +# 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 = """ +T: parallelize +W: +U: +O: +P: 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_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 +# 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: #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> +# 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.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: ?>> +# 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: %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: %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"} +# 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_conv_descript_prp.py b/tests/filecheck/search/test_conv_descript_prp.py new file mode 100644 index 00000000..e848f68c --- /dev/null +++ b/tests/filecheck/search/test_conv_descript_prp.py @@ -0,0 +1,31 @@ +# 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) + +for x in sorted(strategy._constraints): + print(x) +print(sum(1 for _ in strategy.sample(100))) + +# 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 new file mode 100644 index 00000000..dc7d1f00 --- /dev/null +++ b/tests/filecheck/search/test_matmul_descript_interchange.py @@ -0,0 +1,31 @@ +# 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) + +for x in sorted(strategy._constraints): + print(x) +print(sum(1 for _ in strategy.sample(100))) + +# 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 new file mode 100644 index 00000000..45b60a9b --- /dev/null +++ b/tests/filecheck/search/test_matmul_descript_pprprp.py @@ -0,0 +1,33 @@ +# 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 = """ +- 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) + +for x in sorted(strategy._constraints): + print(x) +print(sum(1 for _ in strategy.sample(100))) + +# 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 new file mode 100644 index 00000000..c417a550 --- /dev/null +++ b/tests/filecheck/search/test_matmul_descript_pprprp_interchange.py @@ -0,0 +1,32 @@ +# 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 = """ +- P: +- P: +- U: +- R: +- P:""" +strategy = Strategy(graph, spec, initialize=False,) + +for x in sorted(strategy._constraints): + print(x) +print(sum(1 for _ in strategy.sample(100))) + +# 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 new file mode 100644 index 00000000..e383621e --- /dev/null +++ b/tests/filecheck/search/test_matmul_descript_twuop.py @@ -0,0 +1,34 @@ +# 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 = """ +T: parallelize +W: +U: +O: +P: unroll vectorize +""" +strategy = Strategy(graph, spec, initialize=False, partial_tiles=True, partial_unrolls=True) + +for x in sorted(strategy._constraints): + print(x) +print(sum(1 for _ in strategy.sample(100))) + +# 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 d1ce888d..fc240710 100644 --- a/tests/filecheck/search/test_matmul_descript_yaml_goto.py +++ b/tests/filecheck/search/test_matmul_descript_yaml_goto.py @@ -23,40 +23,70 @@ 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} - - nvr * mr * kr <= {reorder_buffer} - - kc * nr <= {nb_words_L1} - - kc * mc <= {nb_words_L2} - - kc * nc <= {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) + footprint(A, L1) <= {nb_words_L1} + - footprint(A, L2) <= {nb_words_L2} + - footprint(B, L3) <= {nb_words_L3} +schedule: j: k: - B: pack - i: - A: pack - j#nc: - i#mc: + B: pack=pack_B pad + i: level=L3 + A: pack pad=pad_A + j#nc: level=L2 + i#mc: level=L1 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)) +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, +) + +for x in sorted(strategy._constraints): + print(x) 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-NEXT: 100 +# 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