diff --git a/docs/spec/hir.md b/docs/spec/hir.md index 23e707a6..cfdc3dc4 100644 --- a/docs/spec/hir.md +++ b/docs/spec/hir.md @@ -619,11 +619,12 @@ their input when it states one. An input with `layout=None` produces a view with the registered relation. - `Slice` is normalized as `Slice(x, starts, sizes=..., strides=...)`. `starts` is a tuple of rank-0 integer operands; `sizes` and `strides` are - `ShapeDim` attributes. Its result shape is exactly `sizes` and MUST NOT contain - an induction `Var`. A start MAY be dim arithmetic over an induction `Var` — a - window moved off that loop's window by a compile-time offset. That start is an - address computed where it is read, not a value some op produces, so a walk over - compute ops MUST leave it alone. + `ShapeDim` attributes stored in the same IR normal form as every other dim. + Its result shape is exactly the normalized `sizes` and MUST NOT contain an + induction `Var`. A start MAY be dim arithmetic over an induction `Var` — a window moved off that + loop's window by a compile-time offset. That start is an address computed where + it is read, not a value some op produces, so a walk over compute ops MUST leave + it alone. - A plain-layout `Slice` with static starts MUST produce a `ComposedLayout`: its offset is the source offset plus the starts multiplied by the source strides, and its outer layout carries the sliced shape and retained strides (multiplied @@ -892,8 +893,16 @@ class Reshape(Op): - A plain C-order input reshapes to a C-order `Layout` over `new_shape`. An input with no assigned layout, or a non-contiguous plain input whose regroup cannot be expressed, has a `None` result layout. - - A fully-`Broadcast` `ShardLayout` input (every attr `Broadcast`, no genuine - sharding) reshapes to a plain (unsharded) output. + - A bare, fully-`Broadcast` `ShardLayout` input (every attr `Broadcast`, no + genuine sharding) carries that `ShardLayout` through `Reshape` when the + input layout positions can express `new_shape` by the view rules below. + When they cannot (including while either layout shape is symbolic), the + result layout is `None` and no error is raised because no genuine ownership + is discarded. This rule does not extend to a `ComposedLayout` whose outer + layout is a `ShardLayout`; that input follows the generic composed-layout + rule and may produce a `None` layout. + - A non-genuinely-sharded `UMAT` input reshaped to `()` remains the early + scalar case and produces a plain `UMAT` scalar. - A genuine `ShardLayout` input (at least one non-`Broadcast` attr) carries through `Reshape` when the reshape is expressible as a view over the input's layout positions (`layout.layout.shape`, [shard §7.1.1](./shard.md#711-layoutshape)): @@ -913,8 +922,9 @@ class Reshape(Op): block spans a boundary deeper than one divide, or two or more `Split`-bound positions interacting across the same regroup — is not yet supported and MUST fail closed. - - A reshape not expressible by the above MUST fail closed rather than - fabricate a layout. + - A reshape of a genuine `ShardLayout` not expressible by the above MUST fail + closed rather than fabricate or discard a layout. An unexpressible bare, + fully-`Broadcast` `ShardLayout` instead follows the `None`-layout rule above. ##### Cast ```python diff --git a/docs/spec/parser.md b/docs/spec/parser.md index 676ce07f..1a1fef5a 100644 --- a/docs/spec/parser.md +++ b/docs/spec/parser.md @@ -457,6 +457,10 @@ compile-time-expr ::= number-literal | identifier | compile-time-expr '.' identi rejected ([hir §1.3](./hir.md#13-op)). - Evaluation MUST NOT call anything reached from a speculative position: a value that is not statically reachable is parsed as IR instead. +- Dimension arithmetic has one canonical spelling after it enters IR, including + shape annotations and op attributes. Slice endpoints that contain the same + runtime scalar value MAY cancel to a static window size; an unrelated runtime + endpoint remains invalid under the ordinary `ShapeDim` rule. A **compile-time list** holds `Expr` elements and never reaches the IR: @@ -482,8 +486,10 @@ For a tensor slice, a run-time rank-0 integer is permitted as an endpoint only when the resulting window size is a compile-time dimension. The canonical spelling is `start:start + K`; the parser MUST reject an unrelated stop endpoint because `Slice.sizes` is a static attribute. The slice stride remains -compile-time. A tile window keeps its own length and MAY be **moved** by a -compile-time offset instead ([§1.7](#17-for-i-in-tile--for-i-in-range-hir-only)). +compile-time. An endpoint MAY also be a compile-time dimension, including the +axis's own symbolic extent; its window size is symbolic accordingly. A tile +window keeps its own length and MAY be **moved** by a compile-time offset instead +([§1.7](#17-for-i-in-tile--for-i-in-range-hir-only)). ## 2. DSL namespace surface diff --git a/docs/spec/types.md b/docs/spec/types.md index 91877246..b3441d4c 100644 --- a/docs/spec/types.md +++ b/docs/spec/types.md @@ -491,10 +491,22 @@ def ceildiv(a, b) -> Expr: - `DimVar` identity MUST be canonical per `(name, lo, hi)`. Same-name dimensions in one function signature MUST agree on bounds. - Producers of arithmetic dimension calls MUST route construction through - `simplify_dim`. - - `simplify_dim` MUST fold two integer-valued constant operands, except - division or modulo by zero; otherwise it MUST retain the canonical call. - It MUST NOT perform algebraic identity folding. + `simplify_dim`. A `DimVar` and the `Call` produced by dimension arithmetic + both support continued `+`, `-`, `*`, `//`, and `%` construction, including + the reflected forms. Other `Call` values do not support this arithmetic. + - `simplify_dim` MUST only construct dimension arithmetic: it wraps raw integer + operands as integer `Constant` values and rejects boolean operands, but MUST + NOT fold constants or apply algebraic identities. + - Every dimension stored in an IR type or op attribute MUST use the one + `normalize_dim` isl affine normal form. This covers function signatures, + inferred types, layouts, topology and mesh entries, and op attributes; it + does not cover dimension expressions used as `Expr` operands, such as a + `Slice` start address. All-constant expressions normalize to plain Python + integers. Runtime scalar `Var` leaves MUST be represented by object identity, + not name, so repeated uses of one value can cancel without conflating + distinct same-named values. Normalization MUST NOT apply `DimVar` envelope + bounds. Expressions outside the affine subset or not decodable as one + `ShapeDim` MUST remain unchanged. - `is_dim_expr` MUST accept non-boolean integers, `DimVar`, integer-valued `Constant`, and recursively valid calls to the seven dimension arithmetic operations, and MUST reject other values. diff --git a/src/tilefoundry/ir/core/expr.py b/src/tilefoundry/ir/core/expr.py index bc58a5e4..18e4dd3a 100644 --- a/src/tilefoundry/ir/core/expr.py +++ b/src/tilefoundry/ir/core/expr.py @@ -82,6 +82,44 @@ class Call(Expr): target: Op args: tuple[Expr, ...] + def _dim_binop(self, other, op_name: str, *, reverse: bool = False): + from tilefoundry.ir.types import dim # noqa: PLC0415 + + if not dim.is_dim_op_call(self): + return NotImplemented + operands = (other, self) if reverse else (self, other) + return dim._dim_binop(getattr(dim, op_name), *operands) + + def __add__(self, other): + return self._dim_binop(other, "DimAdd") + + def __radd__(self, other): + return self._dim_binop(other, "DimAdd", reverse=True) + + def __sub__(self, other): + return self._dim_binop(other, "DimSub") + + def __rsub__(self, other): + return self._dim_binop(other, "DimSub", reverse=True) + + def __mul__(self, other): + return self._dim_binop(other, "DimMul") + + def __rmul__(self, other): + return self._dim_binop(other, "DimMul", reverse=True) + + def __floordiv__(self, other): + return self._dim_binop(other, "DimFloorDiv") + + def __rfloordiv__(self, other): + return self._dim_binop(other, "DimFloorDiv", reverse=True) + + def __mod__(self, other): + return self._dim_binop(other, "DimMod") + + def __rmod__(self, other): + return self._dim_binop(other, "DimMod", reverse=True) + @dataclass(frozen=True) class Tuple(Expr): diff --git a/src/tilefoundry/ir/core/module.py b/src/tilefoundry/ir/core/module.py index 6d22607e..8bc7c280 100644 --- a/src/tilefoundry/ir/core/module.py +++ b/src/tilefoundry/ir/core/module.py @@ -16,6 +16,7 @@ from tilefoundry.ir.hir.function import Function as HirFunction from tilefoundry.ir.tir.prim_function import PrimFunction from tilefoundry.ir.types.shard.mesh import Topology +from tilefoundry.ir.types.substitute import canonicalize_topology_dims from tilefoundry.ir.types.tensor_type import TensorType from tilefoundry.target.base import Target, target_instance from tilefoundry.utils.spec_ref import spec_ref_render @@ -245,6 +246,9 @@ def __post_init__(self) -> None: f"method and a function/child module; names must be disjoint" ) if self.topologies is not None: + topologies = tuple(canonicalize_topology_dims(t) for t in self.topologies) + if topologies != self.topologies: + object.__setattr__(self, "topologies", topologies) names = [t.name for t in self.topologies] dupes = sorted({n for n in names if names.count(n) > 1}) if dupes: diff --git a/src/tilefoundry/ir/core/op.py b/src/tilefoundry/ir/core/op.py index 1a37a151..9ab47012 100644 --- a/src/tilefoundry/ir/core/op.py +++ b/src/tilefoundry/ir/core/op.py @@ -37,7 +37,9 @@ def _normalize_attr(name: str, value: Any) -> Any: """ if name == "storage": return resolve_storage(value) - return value + from tilefoundry.ir.types.dim_isl import normalize_dim_entries # noqa: PLC0415 + + return normalize_dim_entries(value) @dataclass(frozen=True) diff --git a/src/tilefoundry/ir/hir/function.py b/src/tilefoundry/ir/hir/function.py index e76f5935..951ea5b8 100644 --- a/src/tilefoundry/ir/hir/function.py +++ b/src/tilefoundry/ir/hir/function.py @@ -10,7 +10,7 @@ from tilefoundry.ir.hir.grid_region import GridRegionExpr from tilefoundry.ir.types import TensorType, Type, callable_type_for from tilefoundry.ir.types.dim import is_dim_expr -from tilefoundry.ir.types.substitute import substitute_shape_dim +from tilefoundry.ir.types.substitute import canonicalize_dims, substitute_shape_dim from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.contexts import FunctionScope, TypeInferContext @@ -46,7 +46,12 @@ def build( variants: tuple["Function", ...] = (), converters: tuple[tuple[str, "Function"], ...] = (), ) -> "Function": - """Construct a Function with the canonical CallableType.""" + """Construct a Function whose declarations and callable type are canonical.""" + for param in params: + canonical = canonicalize_dims(param.type) + if canonical is not param.type: + object.__setattr__(param, "type", canonical) + return_type = canonicalize_dims(return_type) return cls( name=name, params=params, diff --git a/src/tilefoundry/ir/hir/nn/conv2d.py b/src/tilefoundry/ir/hir/nn/conv2d.py index dcb6e78c..33060d24 100644 --- a/src/tilefoundry/ir/hir/nn/conv2d.py +++ b/src/tilefoundry/ir/hir/nn/conv2d.py @@ -13,6 +13,7 @@ from tilefoundry.ir.hir._shard_checks import check_multilinear_partials from tilefoundry.ir.types import TensorType from tilefoundry.ir.types.dim import DimAdd, DimFloorDiv, DimSub, simplify_dim +from tilefoundry.ir.types.dim_isl import normalize_dim from tilefoundry.ir.types.shape_helpers import i64_const, static_dim_value from tilefoundry.ir.types.shard import Layout, try_c_order_strides from tilefoundry.ir.types.shard.shard_layout import shard_layout_of, split_target_axes @@ -62,7 +63,7 @@ def _out_spatial(in_dim: Expr, k: int, s: int, p: int, d: int) -> Expr: sub_k = simplify_dim(DimSub, (add_pad, _i64(eff_k))) div_s = simplify_dim(DimFloorDiv, (sub_k, _i64(s))) plus_1 = simplify_dim(DimAdd, (div_s, _i64(1))) - return plus_1 + return normalize_dim(plus_1) def _pair(call, ctx, name: str, values: tuple, *, positive: bool) -> tuple[int, int]: diff --git a/src/tilefoundry/ir/hir/tensor/arange.py b/src/tilefoundry/ir/hir/tensor/arange.py index b88f2d0a..22fc2815 100644 --- a/src/tilefoundry/ir/hir/tensor/arange.py +++ b/src/tilefoundry/ir/hir/tensor/arange.py @@ -12,6 +12,7 @@ from tilefoundry.ir.core.register import register_op from tilefoundry.ir.types import DType, TensorType from tilefoundry.ir.types.dim import DimSub, ceildiv, is_dim_expr, simplify_dim +from tilefoundry.ir.types.dim_isl import normalize_dim from tilefoundry.ir.types.shape_dim import ShapeDim from tilefoundry.ir.types.shape_helpers import static_dim_value from tilefoundry.ir.types.storage import StorageKind @@ -46,7 +47,7 @@ def _(call: "Call", ctx: "TypeInferContext") -> TensorType: if op.dtype not in (DType.i32, DType.i64): ctx.error(call, f"dtype must be i32 or i64, got {op.dtype}") - difference = simplify_dim(DimSub, (op.end, op.start)) + difference = normalize_dim(simplify_dim(DimSub, (op.end, op.start))) static_difference = static_dim_value(difference) if static_difference is not None and static_difference < 0: ctx.error( diff --git a/src/tilefoundry/ir/hir/tensor/reshape.py b/src/tilefoundry/ir/hir/tensor/reshape.py index 5eecfde9..1eacbf7e 100644 --- a/src/tilefoundry/ir/hir/tensor/reshape.py +++ b/src/tilefoundry/ir/hir/tensor/reshape.py @@ -166,16 +166,13 @@ def _(call: "Call", ctx: "TypeInferContext") -> TensorType: new_layout = None if isinstance(x_ty.layout, ShardLayout): genuine = any(not isinstance(a, Broadcast) for a in x_ty.layout.attrs) - if not genuine: - new_layout = None - else: - new_layout = _carry_sharded_reshape(x_ty.layout, new_shape) - if new_layout is None: - ctx.error( - call, - "Reshape cannot express the sharded layout: new shape does " - "not align with the input layout factorization", - ) + new_layout = _carry_sharded_reshape(x_ty.layout, new_shape) + if new_layout is None and genuine: + ctx.error( + call, + "Reshape cannot express the sharded layout: new shape does " + "not align with the input layout factorization", + ) else: source = x_ty.layout if isinstance(source, Layout): diff --git a/src/tilefoundry/ir/hir/tensor/slice.py b/src/tilefoundry/ir/hir/tensor/slice.py index f774ce58..36bc9794 100644 --- a/src/tilefoundry/ir/hir/tensor/slice.py +++ b/src/tilefoundry/ir/hir/tensor/slice.py @@ -52,27 +52,13 @@ def _i64(value: int) -> Constant: def slice_size(begin: Expr, end: Expr, stride: Expr) -> Expr: - """Return ``max(0, ceil((end - begin) / stride))`` as a dimension Expr. - - Constant chains fold through ``simplify_dim``. A non-positive constant - stride returns zero explicitly because generic arithmetic folding does not - encode slice-domain semantics. - """ - if isinstance(begin, Constant) and isinstance(end, Constant) and isinstance(stride, Constant): - b, e, s = int(begin.value), int(end.value), int(stride.value) - if s <= 0: + """Return ``ceil((end - begin) / stride)`` as a dimension expression.""" + static = tuple(_constant_int(value) for value in (begin, end, stride)) + if all(value is not None for value in static): + start, stop, step = static + if step > 0 and stop < start: return _i64(0) - n = max(0, (e - b + s - 1) // s) - return _i64(n) - if isinstance(end, Call) and isinstance(end.target, DimAdd) and end.args[0] is begin: - window = end.args[1] - if isinstance(stride, Constant) and stride.value == 1: - return window - bump = simplify_dim( - DimAdd, - (window, simplify_dim(DimSub, (stride, _i64(1)))), - ) - return simplify_dim(DimFloorDiv, (bump, stride)) + diff = simplify_dim(DimSub, (end, begin)) bump = simplify_dim( diff --git a/src/tilefoundry/ir/types/dim.py b/src/tilefoundry/ir/types/dim.py index b765c6f5..283e12ee 100644 --- a/src/tilefoundry/ir/types/dim.py +++ b/src/tilefoundry/ir/types/dim.py @@ -136,23 +136,11 @@ class DimMax(Op): b = ParamDef(kind="input") -_DIM_FOLDERS: dict[type[Op], object] = { - DimAdd: lambda a, b: a + b, - DimSub: lambda a, b: a - b, - DimMul: lambda a, b: a * b, - DimFloorDiv: lambda a, b: a // b, - DimMod: lambda a, b: a % b, - DimMin: min, - DimMax: max, -} - - def simplify_dim(op_cls: type[Op], args: tuple) -> Expr: - """Fold dimension arithmetic when every operand is an integer constant. + """Build dimension arithmetic without applying algebraic rules. - Raw integers canonicalize to i64 constants. Unsupported operations and - division or modulo by zero remain calls for later verification; algebraic - identities are not folded. + Raw integers become i64 constants and bool remains invalid. The dimension + is normalized only when it enters IR. See [types §4](docs/spec/types.md#4-dim--symbolic-shape-dimensions). """ @@ -173,21 +161,6 @@ def _wrap(v): canon_args = tuple(_wrap(a) for a in args) - fold = _DIM_FOLDERS.get(op_cls) - if ( - fold is not None - and len(canon_args) == 2 - and all( - isinstance(a, Constant) and isinstance(a.value, int) and not isinstance(a.value, bool) - for a in canon_args - ) - ): - a_val = int(canon_args[0].value) - b_val = int(canon_args[1].value) - if op_cls in (DimFloorDiv, DimMod) and b_val == 0: - pass - else: - return Constant(type=ti64, value=fold(a_val, b_val)) return Call(type=ti64, target=op_cls(), args=canon_args) @@ -217,6 +190,21 @@ def is_dim_expr(value) -> bool: return False +def dim_expr(value) -> Expr: + """Convert a dimension value to an expression without simplifying it.""" + if isinstance(value, Expr): + return value + if isinstance(value, bool): + raise TypeError("dim_expr: bool is not a dimension") + if isinstance(value, int): + from .tensor_type import TensorType # noqa: PLC0415 + + return Constant(type=TensorType.umat_scalar(), value=value) + if isinstance(value, DimVar): + return simplify_dim(DimAdd, (value, 0)) + raise TypeError(f"dim_expr: expected int, DimVar, or Expr, got {type(value).__name__}") + + def is_dim_op_call(value) -> bool: """True iff *value* is a ``Call`` over one of this module's dim ops. @@ -230,11 +218,7 @@ def is_dim_op_call(value) -> bool: def dim_min(a, b) -> Expr: - """Symbolic ``min`` dim expression, folded to a ``Constant`` when both operands are static. - - Symbolic ``min(a, b)`` dim expression, folded to a ``Constant`` when both - operands are static. - """ + """Build a symbolic ``min(a, b)`` dimension expression.""" result = _dim_binop(DimMin, a, b) if result is NotImplemented: raise TypeError( @@ -260,8 +244,7 @@ def ceildiv(a, b) -> Expr: Composes existing dim-arithmetic ops — there is no dedicated ceil-div op. Operands may be ``int`` (non-bool), ``DimVar`` or ``Expr``; the - result is the same ``ShapeDim`` form produced by ``simplify_dim`` and - folds to a ``Constant`` when both operands are static. + result is the same ``ShapeDim`` form produced by ``simplify_dim``. """ num = simplify_dim(DimSub, (simplify_dim(DimAdd, (a, b)), 1)) return simplify_dim(DimFloorDiv, (num, b)) @@ -279,6 +262,7 @@ def ceildiv(a, b) -> Expr: "DimMax", "simplify_dim", "is_dim_expr", + "dim_expr", "is_dim_op_call", "dim_min", "dim_max", diff --git a/src/tilefoundry/ir/types/dim_isl.py b/src/tilefoundry/ir/types/dim_isl.py new file mode 100644 index 00000000..05718b66 --- /dev/null +++ b/src/tilefoundry/ir/types/dim_isl.py @@ -0,0 +1,302 @@ +"""The single ShapeDim <-> isl bridge and normalization authority.""" + +from __future__ import annotations + +import isl + +from tilefoundry.ir.core.expr import Call, Constant, Var + +from .dim import ( + _DIM_OP_TYPES, + DimAdd, + DimFloorDiv, + DimMax, + DimMin, + DimMod, + DimMul, + DimSub, + DimVar, +) +from .tensor_type import TensorType + + +def _is_const(node) -> bool: + if isinstance(node, bool): + return False + return isinstance(node, int) or isinstance(node, Constant) + + +def _bind_param( + value, + params: dict[str, tuple[int, int] | None], + param_map: dict[str, object] | None, + identities: dict[int, str] | None, +) -> str: + if isinstance(value, DimVar): + name = value.name + bound = (value.lo, value.hi) + previous = params.get(name) + if previous is not None and previous != bound: + raise ValueError( + f"DimVar {name!r} used with conflicting bounds {previous} vs {bound}" + ) + elif identities is not None: + key = id(value) + known = identities.get(key) + if known is not None: + return known + index = len(identities) + name = f"__tf_runtime_{index}" + while name in params: + index += 1 + name = f"__tf_runtime_{index}" + identities[key] = name + bound = None + else: + raise TypeError(f"unsupported ShapeDim {type(value).__name__}") + + params[name] = bound + if param_map is not None: + previous_value = param_map.get(name) + if previous_value is not None and previous_value is not value: + raise ValueError(f"isl parameter {name!r} maps to multiple dimension values") + param_map[name] = value + return name + + +def _range_expr( + dim, + params: dict[str, tuple[int, int] | None], + *, + param_map: dict[str, object] | None = None, + identities: dict[int, str] | None = None, +) -> str: + if isinstance(dim, bool): + raise TypeError("ShapeDim must not be bool") + if isinstance(dim, int): + return str(dim) + if isinstance(dim, Constant): + return str(int(dim.value)) + if isinstance(dim, (DimVar, Var)): + return _bind_param(dim, params, param_map, identities) + if isinstance(dim, Call): + op = type(dim.target) + if op not in _DIM_OP_TYPES: + return _bind_param(dim, params, param_map, identities) + a, b = dim.args + if op is DimMul and not (_is_const(a) or _is_const(b)): + name = _bind_param(dim, params, param_map, identities) + if params[name] is None: + params[name] = dim_range(dim) + return name + if op in (DimFloorDiv, DimMod) and not _is_const(b): + raise NotImplementedError( + f"{op.__name__} by a symbolic divisor has no isl representation" + ) + sa = _range_expr(a, params, param_map=param_map, identities=identities) + sb = _range_expr(b, params, param_map=param_map, identities=identities) + if op is DimAdd: + return f"({sa} + {sb})" + if op is DimSub: + return f"({sa} - {sb})" + if op is DimMul: + return f"({sa} * {sb})" + if op is DimFloorDiv: + return f"floor({sa}/{sb})" + if op is DimMod: + return f"({sa} mod {sb})" + if op is DimMax: + return f"max({sa}, {sb})" + if op is DimMin: + return f"min({sa}, {sb})" + raise TypeError(f"unsupported ShapeDim {type(dim).__name__}") + + +def _raw_dim_call(op_cls, args: tuple): + scalar = TensorType.umat_scalar() + + def wrap(value): + if isinstance(value, bool): + raise TypeError("bool is not a ShapeDim") + if isinstance(value, int): + return Constant(type=scalar, value=value) + return value + + return Call(type=scalar, target=op_cls(), args=tuple(wrap(arg) for arg in args)) + + +def _visit(expr, param_map: dict[str, object]): + if isinstance(expr, isl.ast_expr_int): + return int(expr.val().num_si()) + if isinstance(expr, isl.ast_expr_id): + name = expr.id().name() + if name not in param_map: + raise ValueError(f"isl identifier {name!r} has no known ShapeDim") + return param_map[name] + if isinstance(expr, isl.ast_expr_op): + op = expr.op_type() + Op = isl.ast_expr_op_type + if op == Op.MINUS: + return _raw_dim_call(DimSub, (0, _visit(expr.op_arg(0), param_map))) + a = _visit(expr.op_arg(0), param_map) + b = _visit(expr.op_arg(1), param_map) + if op == Op.ADD: + return _raw_dim_call(DimAdd, (a, b)) + if op == Op.SUB: + return _raw_dim_call(DimSub, (a, b)) + if op == Op.MUL: + return _raw_dim_call(DimMul, (a, b)) + if op in (Op.DIV, Op.PDIV_Q, Op.FDIV_Q): + return _raw_dim_call(DimFloorDiv, (a, b)) + if op == Op.PDIV_R: + return _raw_dim_call(DimMod, (a, b)) + if op == Op.MAX: + return _raw_dim_call(DimMax, (a, b)) + if op == Op.MIN: + return _raw_dim_call(DimMin, (a, b)) + raise NotImplementedError(f"ast_expr op {op!r} has no ShapeDim decoding") + raise NotImplementedError(f"unsupported ast_expr type {type(expr).__name__}") + + +def to_dim(pw_aff: "isl.pw_aff", param_map: dict[str, object]): + """Decode *pw_aff* into a ShapeDim using *param_map* for identifiers.""" + build = isl.ast_build.from_context(pw_aff.domain_space().universe_set()) + return _visit(build.expr_from(pw_aff), param_map) + + +def normalize_dim(value): + """Return the sole isl affine normal form for one dimension value.""" + if isinstance(value, int) and not isinstance(value, bool): + return value + if isinstance(value, DimVar): + return value + if not isinstance(value, (Constant, Var, Call)): + return value + try: + params: dict[str, tuple[int, int] | None] = {} + param_map: dict[str, object] = {} + expr = _range_expr( + value, + params, + param_map=param_map, + identities={}, + ) + prefix = f"[{', '.join(params)}] -> " if params else "" + normalized = to_dim(isl.pw_aff(prefix + f"{{ [{expr}] }}"), param_map) + return value if normalized == value else normalized + except (TypeError, ValueError, NotImplementedError, isl.Error): + return value + + +def normalize_dim_entries(value): + """Normalize dimension leaves in a tuple, preserving unchanged objects.""" + if isinstance(value, int) and not isinstance(value, bool): + return value + if isinstance(value, tuple): + entries = tuple(normalize_dim_entries(entry) for entry in value) + return value if all(a is b for a, b in zip(entries, value)) else entries + if isinstance(value, DimVar) or ( + isinstance(value, Constant) + and isinstance(value.value, int) + and not isinstance(value.value, bool) + ) or ( + isinstance(value, Call) and isinstance(value.target, _DIM_OP_TYPES) + ): + return normalize_dim(value) + return value + + +def dim_range(dim) -> tuple[int, int]: + """Return conservative half-open value bounds ``[lo, hi)`` for *dim*.""" + if isinstance(dim, bool): + raise TypeError("ShapeDim must not be bool") + if isinstance(dim, int): + return (dim, dim + 1) + if isinstance(dim, Constant): + value = int(dim.value) + return (value, value + 1) + if isinstance(dim, DimVar): + return (dim.lo, dim.hi) + if isinstance(dim, Call) and type(dim.target) is DimMul: + a, b = dim.args + if not (_is_const(a) or _is_const(b)): + alo, ahi = dim_range(a) + blo, bhi = dim_range(b) + corners = ( + alo * blo, + alo * (bhi - 1), + (ahi - 1) * blo, + (ahi - 1) * (bhi - 1), + ) + return (min(corners), max(corners) + 1) + params: dict[str, tuple[int, int] | None] = {} + expr = _range_expr(dim, params) + prefix = f"[{', '.join(params)}] -> " if params else "" + pw_aff = isl.pw_aff(prefix + f"{{ [{expr}] }}") + if params: + bounds = " and ".join( + f"{lo} <= {name} <= {hi - 1}" + for name, bound in params.items() + for lo, hi in (bound,) + ) + pw_aff = pw_aff.intersect_params(isl.set(prefix + f"{{ : {bounds} }}")) + return (int(pw_aff.min_val().num_si()), int(pw_aff.max_val().num_si()) + 1) + + +def to_domain(extents: tuple) -> tuple: + """Build a bounded iteration domain and its isl-parameter ShapeDim map.""" + param_map: dict[str, object] = {} + bounds: dict[str, tuple[int, int]] = {} + seen: dict = {} + names: list[str] = [] + + def bind(name: str, dim, lo: int, hi: int) -> None: + bound = (lo, hi) + previous = bounds.get(name) + if previous is not None and previous != bound: + raise ValueError( + f"isl parameter {name!r} used with conflicting bounds {previous} vs {bound}" + ) + if name not in bounds: + names.append(name) + bounds[name] = bound + param_map[name] = dim + + dims = [f"d{i}" for i in range(len(extents))] + constraints: list[str] = [] + for i, extent in enumerate(extents): + if isinstance(extent, bool): + raise TypeError("ShapeDim must not be bool") + if isinstance(extent, int): + constraints.append(f"0 <= d{i} < {extent}") + elif isinstance(extent, Constant): + constraints.append(f"0 <= d{i} < {int(extent.value)}") + elif isinstance(extent, DimVar): + bind(extent.name, extent, extent.lo, extent.hi) + constraints.append(f"0 <= d{i} < {extent.name}") + elif isinstance(extent, Call): + name = seen.get(extent) + if name is None: + name = f"D{i}" + seen[extent] = name + lo, hi = dim_range(extent) + bind(name, extent, lo, hi) + constraints.append(f"0 <= d{i} < {name}") + else: + raise TypeError(f"unsupported ShapeDim {type(extent).__name__}") + + constraints += [f"{bounds[name][0]} <= {name} < {bounds[name][1]}" for name in names] + prefix = f"[{', '.join(names)}] -> " if names else "" + if not dims: + return isl.set(prefix + "{ [] }"), param_map + body = f"{{ [{', '.join(dims)}] : {' and '.join(constraints)} }}" + return isl.set(prefix + body), param_map + + +__all__ = [ + "dim_range", + "normalize_dim", + "normalize_dim_entries", + "to_dim", + "to_domain", +] diff --git a/src/tilefoundry/ir/types/substitute.py b/src/tilefoundry/ir/types/substitute.py index 703225b4..ef6af0af 100644 --- a/src/tilefoundry/ir/types/substitute.py +++ b/src/tilefoundry/ir/types/substitute.py @@ -13,6 +13,7 @@ from tilefoundry.ir.core.expr import Call, Constant from .dim import _DIM_OP_TYPES, DimVar, simplify_dim +from .dim_isl import normalize_dim from .tensor_type import TensorType, TupleType, Type @@ -132,6 +133,85 @@ def substitute_dims(value: Type, bindings: Mapping[str, int]) -> Type: return value +def canonicalize_dims(value: Type) -> Type: + """Return *value* with every symbolic dimension in isl normal form.""" + if not has_symbolic_dims(value): + return value + if isinstance(value, TensorType): + shape = tuple(normalize_dim(entry) for entry in value.shape) + layout = _canonicalize_layout_dims(value.layout) + if shape == value.shape and layout is value.layout: + return value + return TensorType( + shape=shape, + dtype=value.dtype, + layout=layout, + storage=value.storage, + ) + if isinstance(value, TupleType): + fields = tuple(canonicalize_dims(field) for field in value.fields) + if fields == value.fields: + return value + return TupleType(fields=fields) + return value + + +def _canonicalize_layout_dims(layout: object) -> object: + if layout is None: + return layout + Layout, ComposedLayout, ShardLayout, _, _ = _shard_types() + if isinstance(layout, ShardLayout): + inner = _canonicalize_layout_dims(layout.layout) + mesh = _canonicalize_mesh_dims(layout.mesh) + if inner is layout.layout and mesh is layout.mesh: + return layout + return ShardLayout(layout=inner, attrs=layout.attrs, mesh=mesh) + if isinstance(layout, ComposedLayout): + outer = _canonicalize_layout_dims(layout.outer) + inner = _canonicalize_layout_dims(layout.inner) + offset = normalize_dim(layout.offset) + if outer is layout.outer and inner is layout.inner and offset == layout.offset: + return layout + return ComposedLayout(inner=inner, offset=offset, outer=outer) + if isinstance(layout, Layout): + shape = _canonicalize_nested(layout.shape) + strides = None if layout.strides is None else _canonicalize_nested(layout.strides) + if shape == layout.shape and strides == layout.strides: + return layout + return Layout(shape=shape, strides=strides) + return layout + + +def canonicalize_topology_dims(topology: object) -> object: + _, _, _, _, Topology = _shard_types() + if not isinstance(topology, Topology): + return topology + size = normalize_dim(topology.size) + if size == topology.size: + return topology + return Topology(topology.name, size) + + +def _canonicalize_mesh_dims(mesh: object) -> object: + _, _, _, Mesh, _ = _shard_types() + if not isinstance(mesh, Mesh): + return mesh + topologies = tuple(canonicalize_topology_dims(item) for item in mesh.topologies) + layout = _canonicalize_layout_dims(mesh.layout) + if topologies == mesh.topologies and layout is mesh.layout: + return mesh + return Mesh(topologies=topologies, layout=layout, names=mesh.names) + + +def _canonicalize_nested(entries: tuple) -> tuple: + return tuple( + _canonicalize_nested(entry) + if isinstance(entry, tuple) + else (None if entry is None else normalize_dim(entry)) + for entry in entries + ) + + def substitute_layout_dims(layout: object, bindings: Mapping[str, int]) -> object: """*layout* with its bound dimensions replaced. @@ -224,7 +304,7 @@ def substitute_shape_dim(entry: object, bindings: Mapping[str, int]) -> object: args = tuple(substitute_shape_dim(arg, bindings) for arg in entry.args) if args == tuple(entry.args): return entry - folded = simplify_dim(type(entry.target), args) + folded = normalize_dim(simplify_dim(type(entry.target), args)) if isinstance(folded, Constant) and isinstance(folded.value, int): return int(folded.value) @@ -281,6 +361,8 @@ def has_symbolic_dims(value: object) -> bool: __all__ = [ "DimSubstitutionError", + "canonicalize_dims", + "canonicalize_topology_dims", "dim_vars_by_name", "dim_vars_in", "has_symbolic_dims", diff --git a/src/tilefoundry/parser/base.py b/src/tilefoundry/parser/base.py index 01ca8763..dc058c54 100644 --- a/src/tilefoundry/parser/base.py +++ b/src/tilefoundry/parser/base.py @@ -33,14 +33,16 @@ from tilefoundry.ir.hir.tensor.slice import Slice from tilefoundry.ir.hir.tensor.tuple_get_item import TupleGetItem from tilefoundry.ir.types import DType, TensorType, TupleType -from tilefoundry.ir.types.dim import DimAdd, is_dim_expr, simplify_dim +from tilefoundry.ir.types.dim import DimAdd, dim_expr, is_dim_expr, simplify_dim +from tilefoundry.ir.types.dim_isl import normalize_dim from tilefoundry.ir.types.dtype import FloatDType from tilefoundry.ir.types.shape_helpers import i64_const from tilefoundry.ir.types.shard.layout import Layout from tilefoundry.ir.types.shard.mesh import Mesh from tilefoundry.ir.types.shard.shard_layout import ShardLayout, shard_layout_of from tilefoundry.ir.types.storage import StorageKind, resolve_storage -from tilefoundry.visitor_registry import typeinfer_registry +from tilefoundry.ir.types.substitute import canonicalize_dims +from tilefoundry.visitor_registry.visitors import TypeInferVisitor from .dispatch import ( Token, @@ -186,14 +188,14 @@ def _resolve_tensor_type(node: ast.AST, closure: dict[str, Any]) -> TensorType: """ result = try_parse_sugar_tensor_type(node, closure) if result is not None: - return result + return canonicalize_dims(result) try: code = compile(ast.Expression(body=node), "", "eval") val = eval(code, closure) # noqa: S307 — controlled internal eval except Exception as exc: raise VerifyError(f"failed to resolve type annotation: {exc}") if isinstance(val, TensorType): - return val + return canonicalize_dims(val) raise VerifyError(f"annotation did not resolve to TensorType, got {type(val).__name__}") @@ -613,13 +615,13 @@ def _lift_tensor_subscript(self, value, slc: ast.AST): collapsed.append(axis) continue b, e, s = self._slicer_for_dim(el, dim, axis) - b_expr = b if isinstance(b, Expr) else i64_const(int(b)) - e_expr = e if isinstance(e, Expr) else i64_const(int(e)) - s_expr = s if isinstance(s, Expr) else i64_const(int(s)) + b_expr = dim_expr(b) + e_expr = dim_expr(e) + s_expr = dim_expr(s) starts.append(b_expr) from tilefoundry.ir.hir.tensor.slice import slice_size # noqa: PLC0415 - size = slice_size(b_expr, e_expr, s_expr) + size = normalize_dim(slice_size(b_expr, e_expr, s_expr)) if not is_dim_expr(size): raise VerifyError( f"tensor subscript axis {axis}: a run-time start needs the " @@ -1296,10 +1298,7 @@ def _build_call( type=TensorType.scalar(DType.f32), target=op_inst, args=args, metadata=(*self._source_metadata(), *records), ) - fn = typeinfer_registry.lookup(type(op_inst)) - if fn is None: - raise VerifyError(f"no typeinfer registered for {type(op_inst).__name__}") - computed = fn(placeholder, self._ctx) + computed = TypeInferVisitor(self._ctx).visit(placeholder) return dataclasses.replace(placeholder, type=computed) def _eval_static_or_sugar( diff --git a/src/tilefoundry/parser/hir_parser.py b/src/tilefoundry/parser/hir_parser.py index b95ecf82..3dc9209c 100644 --- a/src/tilefoundry/parser/hir_parser.py +++ b/src/tilefoundry/parser/hir_parser.py @@ -33,6 +33,7 @@ is_dim_expr, simplify_dim, ) +from tilefoundry.ir.types.dim_isl import normalize_dim from tilefoundry.ir.types.shard import ( Broadcast, Layout, @@ -736,6 +737,9 @@ def _build_grid_for(self, node: ast.For) -> Expr: raise VerifyError( f"tile() takes 2 arguments (extent, step), got {len(loop_args)}" ) + start, extent, step = ( + normalize_dim(value) for value in (start, extent, step) + ) if not (is_dim_expr(start) and is_dim_expr(extent) and is_dim_expr(step)): offending = ", ".join( f"{label}={_dim_operand_str(value)}" diff --git a/src/tilefoundry/visitor_registry/isl_utility.py b/src/tilefoundry/visitor_registry/isl_utility.py index 95b3f123..f20bdbe0 100644 --- a/src/tilefoundry/visitor_registry/isl_utility.py +++ b/src/tilefoundry/visitor_registry/isl_utility.py @@ -1,216 +1,5 @@ -"""ShapeDim <-> isl bridge for relation-derived shape inference. +"""Compatibility exports for the ShapeDim <-> isl bridge owned by IR types.""" -A composite ``ShapeDim``'s arithmetic never enters the domain carried by -``AccessRelationResult`` — only its ``dim_range`` does, as the bound of a -freshly minted isl parameter (``to_domain``). Recovery reads that domain -back through isl's own ``ast_build`` (``to_dim``). -""" -from __future__ import annotations +from tilefoundry.ir.types.dim_isl import dim_range, to_dim, to_domain -import isl - -from tilefoundry.ir.core.expr import Call, Constant -from tilefoundry.ir.types.dim import ( - DimAdd, - DimFloorDiv, - DimMax, - DimMin, - DimMod, - DimMul, - DimSub, - DimVar, - simplify_dim, -) - - -def _is_const(node) -> bool: - if isinstance(node, bool): - return False - return isinstance(node, int) or isinstance(node, Constant) - - -def _range_expr(dim, params: dict) -> str: - if isinstance(dim, bool): - raise TypeError("ShapeDim must not be bool") - if isinstance(dim, int): - return str(dim) - if isinstance(dim, Constant): - return str(int(dim.value)) - if isinstance(dim, DimVar): - bound = (dim.lo, dim.hi) - prev = params.get(dim.name) - if prev is not None and prev != bound: - raise ValueError( - f"DimVar {dim.name!r} used with conflicting bounds {prev} vs {bound}" - ) - params[dim.name] = bound - return dim.name - if isinstance(dim, Call): - op = type(dim.target) - a, b = dim.args - if op is DimMul and not (_is_const(a) or _is_const(b)): - - - name = f"_t{len(params)}" - params[name] = dim_range(dim) - return name - if op in (DimFloorDiv, DimMod) and not _is_const(b): - raise NotImplementedError( - f"{op.__name__} by a symbolic divisor has no isl representation" - ) - sa, sb = _range_expr(a, params), _range_expr(b, params) - if op is DimAdd: - return f"({sa} + {sb})" - if op is DimSub: - return f"({sa} - {sb})" - if op is DimMul: - return f"({sa} * {sb})" - if op is DimFloorDiv: - return f"floor({sa}/{sb})" - if op is DimMod: - return f"({sa} mod {sb})" - if op is DimMax: - return f"max({sa}, {sb})" - if op is DimMin: - return f"min({sa}, {sb})" - raise NotImplementedError(f"dim op {op.__name__} has no isl representation") - raise TypeError(f"unsupported ShapeDim {type(dim).__name__}") - - -def dim_range(dim) -> tuple[int, int]: - """Half-open value bounds ``[lo, hi)`` of *dim*. - - Half-open value bounds ``[lo, hi)`` of *dim*: build its isl value - expression, bind every identifier to its own bound, and read the range - back from isl. The one case isl cannot express -- a product of two - non-constant terms -- falls back to interval arithmetic. - """ - if isinstance(dim, bool): - raise TypeError("ShapeDim must not be bool") - if isinstance(dim, int): - return (dim, dim + 1) - if isinstance(dim, Constant): - v = int(dim.value) - return (v, v + 1) - if isinstance(dim, DimVar): - return (dim.lo, dim.hi) - if isinstance(dim, Call) and type(dim.target) is DimMul: - a, b = dim.args - if not (_is_const(a) or _is_const(b)): - alo, ahi = dim_range(a) - blo, bhi = dim_range(b) - corners = (alo * blo, alo * (bhi - 1), (ahi - 1) * blo, (ahi - 1) * (bhi - 1)) - return (min(corners), max(corners) + 1) - params: dict = {} - expr = _range_expr(dim, params) - prefix = f"[{', '.join(params)}] -> " if params else "" - pa = isl.pw_aff(prefix + f"{{ [{expr}] }}") - if params: - bounds = " and ".join(f"{lo} <= {n} <= {hi - 1}" for n, (lo, hi) in params.items()) - pa = pa.intersect_params(isl.set(prefix + f"{{ : {bounds} }}")) - return (int(pa.min_val().num_si()), int(pa.max_val().num_si()) + 1) - - -def to_domain(extents: tuple) -> tuple: - """Bounded iteration domain ``{ [d0, ..., dn] : 0 <= di < extent_i }``. - - A static extent is an inline constraint; a bare ``DimVar`` is a - same-name isl parameter bound to its own ``[lo, hi)``; any other - ``ShapeDim`` mints an opaque parameter bound to ``dim_range(extent)``, - deduped by canonical expression across axes. Returns ``(domain, - param_map)`` where ``param_map`` resolves each isl parameter name back - to its ``ShapeDim`` -- this call's own data, not shared across calls. - """ - param_map: dict = {} - bounds: dict[str, tuple[int, int]] = {} - seen: dict = {} - names: list[str] = [] - - def _bind(name: str, dim, lo: int, hi: int) -> None: - bound = (lo, hi) - prev = bounds.get(name) - if prev is not None and prev != bound: - raise ValueError( - f"isl parameter {name!r} used with conflicting bounds {prev} vs {bound}" - ) - if name not in bounds: - names.append(name) - bounds[name] = bound - param_map[name] = dim - - dims = [f"d{i}" for i in range(len(extents))] - constraints: list[str] = [] - for i, ext in enumerate(extents): - if isinstance(ext, bool): - raise TypeError("ShapeDim must not be bool") - if isinstance(ext, int): - constraints.append(f"0 <= d{i} < {ext}") - elif isinstance(ext, Constant): - constraints.append(f"0 <= d{i} < {int(ext.value)}") - elif isinstance(ext, DimVar): - _bind(ext.name, ext, ext.lo, ext.hi) - constraints.append(f"0 <= d{i} < {ext.name}") - elif isinstance(ext, Call): - name = seen.get(ext) - if name is None: - name = f"D{i}" - seen[ext] = name - lo, hi = dim_range(ext) - _bind(name, ext, lo, hi) - constraints.append(f"0 <= d{i} < {name}") - else: - raise TypeError(f"unsupported ShapeDim {type(ext).__name__}") - - constraints += [f"{bounds[name][0]} <= {name} < {bounds[name][1]}" for name in names] - prefix = f"[{', '.join(names)}] -> " if names else "" - if not dims: - return isl.set(prefix + "{ [] }"), param_map - body = f"{{ [{', '.join(dims)}] : {' and '.join(constraints)} }}" - return isl.set(prefix + body), param_map - - -def _visit(expr, param_map: dict): - if isinstance(expr, isl.ast_expr_int): - return int(expr.val().num_si()) - if isinstance(expr, isl.ast_expr_id): - name = expr.id().name() - if name not in param_map: - raise ValueError(f"isl identifier {name!r} has no known ShapeDim") - return param_map[name] - if isinstance(expr, isl.ast_expr_op): - op = expr.op_type() - Op = isl.ast_expr_op_type - if op == Op.MINUS: - return simplify_dim(DimSub, (0, _visit(expr.op_arg(0), param_map))) - a = _visit(expr.op_arg(0), param_map) - b = _visit(expr.op_arg(1), param_map) - if op == Op.ADD: - return simplify_dim(DimAdd, (a, b)) - if op == Op.SUB: - return simplify_dim(DimSub, (a, b)) - if op == Op.MUL: - return simplify_dim(DimMul, (a, b)) - if op in (Op.DIV, Op.PDIV_Q, Op.FDIV_Q): - return simplify_dim(DimFloorDiv, (a, b)) - if op == Op.PDIV_R: - return simplify_dim(DimMod, (a, b)) - if op == Op.MAX: - return simplify_dim(DimMax, (a, b)) - if op == Op.MIN: - return simplify_dim(DimMin, (a, b)) - raise NotImplementedError(f"ast_expr op {op!r} has no ShapeDim decoding") - raise NotImplementedError(f"unsupported ast_expr type {type(expr).__name__}") - - -def to_dim(pw_aff: "isl.pw_aff", param_map: dict): - """To dim. - - Decode *pw_aff* into a ``ShapeDim`` via ``ast_build.expr_from`` plus - a generic ``ast_expr`` visitor. *param_map* resolves each isl - identifier the expression bottoms out on (from a prior ``to_domain``). - """ - build = isl.ast_build.from_context(pw_aff.domain_space().universe_set()) - return _visit(build.expr_from(pw_aff), param_map) - - -__all__ = ["dim_range", "to_domain", "to_dim"] +__all__ = ["dim_range", "to_dim", "to_domain"] diff --git a/src/tilefoundry/visitor_registry/visitors.py b/src/tilefoundry/visitor_registry/visitors.py index d247e52c..99c30fe0 100644 --- a/src/tilefoundry/visitor_registry/visitors.py +++ b/src/tilefoundry/visitor_registry/visitors.py @@ -15,6 +15,7 @@ from tilefoundry.ir.tir.shape import ShapeOf from tilefoundry.ir.tir.stmt import Stmt from tilefoundry.ir.tir.stmts import Evaluate, MeshScope +from tilefoundry.ir.types.substitute import canonicalize_dims from tilefoundry.ir.types.tensor_type import TupleType, Type, UnitType from tilefoundry.ir.visitor import ExprVisitor, StmtVisitor @@ -43,6 +44,9 @@ class TypeInferVisitor(ExprVisitor[Type]): def __init__(self, ctx: TypeInferContext) -> None: self.ctx = ctx + def visit(self, expr: Expr) -> Type: + return canonicalize_dims(super().visit(expr)) + def visit_Var(self, var: Var) -> Type: return var.type diff --git a/tests/analysis/test_isl_utility.py b/tests/analysis/test_isl_utility.py index fcdc3d9a..34ccfc67 100644 --- a/tests/analysis/test_isl_utility.py +++ b/tests/analysis/test_isl_utility.py @@ -5,6 +5,8 @@ import isl import pytest +from tilefoundry.ir.core.expr import Call, Var +from tilefoundry.ir.types import TensorType from tilefoundry.ir.types.dim import ( DimAdd, DimFloorDiv, @@ -16,12 +18,78 @@ DimVar, simplify_dim, ) -from tilefoundry.visitor_registry.isl_utility import dim_range, to_dim, to_domain +from tilefoundry.ir.types.dim_isl import dim_range, normalize_dim, to_dim, to_domain P = DimVar("P", 2048, 1_048_577) Q = DimVar("Q", 2, 33) +def test_normalize_dim_uses_isl_affine_normal_form(): + verbose = simplify_dim( + DimFloorDiv, + ( + simplify_dim( + DimAdd, + (simplify_dim(DimSub, (simplify_dim(DimAdd, (P, 0)), 0)), 0), + ), + 1, + ), + ) + quotient = simplify_dim(DimFloorDiv, (simplify_dim(DimAdd, (P, 8)), 4)) + expected = simplify_dim( + DimAdd, + (simplify_dim(DimFloorDiv, (P, 4)), 2), + ) + + constant = simplify_dim(DimMul, (simplify_dim(DimAdd, (4, 2)), 3)) + + assert normalize_dim(constant) == 18 + assert isinstance(normalize_dim(constant), int) + assert normalize_dim(verbose) is P + assert normalize_dim(quotient) == expected + + +def test_normalize_dim_leaves_unsupported_expressions_unchanged(): + symbolic_divisor = simplify_dim( + DimFloorDiv, + (simplify_dim(DimMul, (P, Q)), DimVar("G", 1, 65)), + ) + piecewise = simplify_dim(DimMin, (P, 8192)) + + assert normalize_dim(symbolic_divisor) is symbolic_divisor + assert normalize_dim(piecewise) is piecewise + + +def test_normalize_dim_keys_runtime_parameters_by_object_identity(): + scalar = TensorType.umat_scalar() + first = Var(type=scalar, name="start") + second = Var(type=scalar, name="start") + + def distance(left, right): + return simplify_dim( + DimSub, + ( + simplify_dim(DimAdd, (left, 9)), + simplify_dim(DimAdd, (right, 1)), + ), + ) + + assert normalize_dim(distance(first, first)) == 8 + distinct = normalize_dim(distance(first, second)) + assert isinstance(distinct, Call) + + def vars_in(value): + if isinstance(value, Var): + return [value] + if isinstance(value, Call): + return [leaf for arg in value.args for leaf in vars_in(arg)] + return [] + + leaves = vars_in(distinct) + assert any(leaf is first for leaf in leaves) + assert any(leaf is second for leaf in leaves) + + def test_dim_range_interval_arithmetic(): """Conservative half-open interval per dim kind, incl. nesting.""" assert dim_range(7) == (7, 8) diff --git a/tests/inspection/test_module_tree_roundtrip.py b/tests/inspection/test_module_tree_roundtrip.py index accfa63f..afc98c9f 100644 --- a/tests/inspection/test_module_tree_roundtrip.py +++ b/tests/inspection/test_module_tree_roundtrip.py @@ -47,7 +47,7 @@ def test_a_derived_topology_and_mesh_geometry_survive_the_round_trip() -> None: assert 'prefill_n = DimVar("prefill_n", 1, 65)' in source assert 'topology_only = DimVar("topology_only", 1, 1025)' in source - assert 'Topology("cta", ceildiv(prefill_n, 8))' in source + assert 'Topology("cta", ((prefill_n - 1) // 8) + 1)' in source assert 'Topology("thread", topology_only)' in source assert imported.topologies == DerivedPrefill.topologies assert imported.entry_function().params[0].type == ( @@ -61,6 +61,8 @@ def test_prefill_decode_specializations_survive_the_round_trip() -> None: imported = import_dsl(source) restored = import_dsl(as_script(imported)) + assert as_script(restored) == as_script(imported) + assert "arange(" in source assert "where(" in source for roundtripped in (imported, restored): @@ -81,6 +83,8 @@ def test_flash_split_k_decode_survives_the_round_trip() -> None: imported = import_dsl(source) restored = import_dsl(as_script(imported)) + assert as_script(restored) == as_script(imported) + for roundtripped in (imported, restored): assert roundtripped.topologies == FlashSplitKDecode.topologies slices = [ diff --git a/tests/ir/test_function_call_typeinfer.py b/tests/ir/test_function_call_typeinfer.py index 00fe77f6..3064bbf0 100644 --- a/tests/ir/test_function_call_typeinfer.py +++ b/tests/ir/test_function_call_typeinfer.py @@ -11,13 +11,16 @@ import pytest from tests.ops.typeinfer_utils import infer_call -from tilefoundry.ir.core import BindingMetadata, Call, Tuple, Var +from tilefoundry.ir.core import BindingMetadata, Call, Constant, Tuple, Var from tilefoundry.ir.core.errors import VerifyError from tilefoundry.ir.core.kinds import BinaryKind from tilefoundry.ir.hir.function import Function from tilefoundry.ir.hir.grid_region import GridRegionExpr from tilefoundry.ir.hir.math.binary import Binary +from tilefoundry.ir.hir.tensor.reshape import Reshape +from tilefoundry.ir.hir.tensor.slice import Slice from tilefoundry.ir.types import DType, TupleType, make_shard_tensor_type, make_tensor_type +from tilefoundry.ir.types.dim import DimMul, DimVar, simplify_dim from tilefoundry.ir.types.shard import make_mesh from tilefoundry.ir.types.shard.shard_layout import Broadcast, Partial, Split from tilefoundry.visitor_registry.contexts import TypeInferContext @@ -91,6 +94,49 @@ def test_explicit_sharded_formal_constrains_its_actual(): infer_call(f, make_shard_tensor_type((4, 8), mesh=_M, attrs=(Split(1),))) +def test_broadcast_formal_accepts_reshaped_runtime_slice(): + packed = make_shard_tensor_type((4, 8, 16), _F, mesh=_M, attrs=(Broadcast(),)) + sliced_type = make_shard_tensor_type((1, 8, 16), _F, mesh=_M, attrs=(Broadcast(),)) + formal = make_shard_tensor_type((8, 16), _F, mesh=_M, attrs=(Broadcast(),)) + packed_var = Var(type=packed, name="packed") + layer = Var(type=make_tensor_type((), DType.i64, storage="umat"), name="layer") + zero = Constant(type=make_tensor_type((), DType.i64), value=0) + starts = Tuple( + type=TupleType(fields=(layer.type, zero.type, zero.type)), + elements=(layer, zero, zero), + ) + sliced = Call( + type=sliced_type, + target=Slice(sizes=(1, 8, 16), strides=(1, 1, 1)), + args=(packed_var, starts), + ) + reshaped = Call(type=formal, target=Reshape(new_shape=(8, 16)), args=(sliced,)) + w = Var(type=formal, name="w") + callee = Function.build(name="consume", params=(w,), body=w, return_type=formal) + call = Call(type=formal, target=callee, args=(reshaped,)) + + assert TypeInferVisitor(TypeInferContext()).visit(call) == formal + + +def test_symbolic_arithmetic_signature_matches_inferred_argument(): + seq = DimVar("call_seq", 1, 4097) + authored_dim = simplify_dim(DimMul, (seq, 2)) + authored_type = make_tensor_type((authored_dim, 8), _F) + x = Var(type=authored_type, name="x") + stage = Function.build( + name="stage", + params=(x,), + body=x, + return_type=authored_type, + ) + y = Var(type=authored_type, name="y") + call = Call(type=authored_type, target=stage, args=(y,)) + expected_dim = simplify_dim(DimMul, (2, seq)) + + assert stage.params[0].type.shape == (expected_dim, 8) + assert TypeInferVisitor(TypeInferContext()).visit(call) == stage.return_type + + def test_plain_formal_rejects_shape_or_dtype_mismatch(): f = _add_callee(_PLAIN) diff --git a/tests/ir/test_simplify_dim.py b/tests/ir/test_simplify_dim.py index 2fc70660..715712a9 100644 --- a/tests/ir/test_simplify_dim.py +++ b/tests/ir/test_simplify_dim.py @@ -1,21 +1,19 @@ -"""Construction-time folding for dim arithmetic Calls. - -Folding is what makes a static shape one canonical value, so the boundaries are -where it must *not* fold (a symbolic operand, a division by zero, a bool) and -where the folded result must arrive as a plain ``int`` rather than a ``Constant`` -or a nested ``Call`` — two shapes that print and compare differently and produced -a real broadcast failure. -""" +"""Construction and IR-boundary normalization for dimension arithmetic.""" from __future__ import annotations import copy +import pytest + +import tilefoundry.ir.types.dim_isl as dim_isl +import tilefoundry.ir.types.substitute as dim_substitute from tilefoundry.ir.core import Tuple, TypeInferContext from tilefoundry.ir.core.expr import Call, Constant, Var from tilefoundry.ir.core.kinds import UnaryKind from tilefoundry.ir.hir._helpers import broadcast_shapes from tilefoundry.ir.hir.math.unary import Unary +from tilefoundry.ir.hir.tensor.reshape import Reshape from tilefoundry.ir.hir.tensor.slice import Slice from tilefoundry.ir.types import DType, TensorType, TupleType from tilefoundry.ir.types.dim import ( @@ -29,6 +27,9 @@ DimVar, simplify_dim, ) +from tilefoundry.ir.types.shard import ComposedLayout, Layout, Mesh, ShardLayout, Topology +from tilefoundry.ir.types.shard.shard_layout import Broadcast +from tilefoundry.visitor_registry.visitors import TypeInferVisitor def _i64(v: int) -> Constant: @@ -48,43 +49,26 @@ def _sym(name: str) -> Call: ) -def test_simplify_dim_folds_all_constant_args() -> None: - """Test simplify dim folds all constant args. - - When both args are int Constants, simplify_dim returns a folded - Constant with the canonical i64 dim type. Floor division follows Python - ``//`` (floor), not C truncation, which is the convention every tiling - expression is written against. - """ +def test_simplify_dim_only_constructs_constant_arithmetic() -> None: table = [ - (DimAdd, 3, 4, 7), - (DimSub, 10, 4, 6), - (DimMul, 3, 4, 12), - (DimFloorDiv, 17, 4, 4), - (DimMod, 17, 5, 2), - (DimMin, 7, 3, 3), - (DimMax, 7, 3, 7), - (DimFloorDiv, -7, 2, -4), + (DimAdd, 3, 4), + (DimSub, 10, 4), + (DimMul, 3, 4), + (DimFloorDiv, 17, 4), + (DimMod, 17, 5), + (DimMin, 7, 3), + (DimMax, 7, 3), + (DimFloorDiv, -7, 2), ] - for op_cls, a, b, expected in table: + for op_cls, a, b in table: result = simplify_dim(op_cls, (_i64(a), _i64(b))) - assert isinstance(result, Constant), ( - f"{op_cls.__name__}: expected Constant, got {type(result).__name__}" - ) - assert result.value == expected + assert isinstance(result, Call) + assert isinstance(result.target, op_cls) + assert result.args == (_i64(a), _i64(b)) assert result.type == TensorType.umat_scalar() -def test_simplify_dim_refuses_to_fold_outside_all_int_constants() -> None: - """Three non-foldable inputs, in either operand position. - - Three non-foldable inputs, in either operand position: - - a non-Constant arg (a symbolic DimVar Call): the Call survives with no - algebraic identity applied, so ``x + 0`` stays a Call; - - division by zero: folding to ``Constant(0)`` would mask a real bug, so the - Call survives for a later verify pass to flag; - - ``Constant(True)``: a bool is not an int dim value. - """ +def test_simplify_dim_constructs_symbolic_and_invalid_arithmetic() -> None: sym = _sym("M") for op_cls in (DimAdd, DimFloorDiv): for args in ((sym, _i64(0)), (_i64(0), sym)): @@ -104,6 +88,116 @@ def test_simplify_dim_refuses_to_fold_outside_all_int_constants() -> None: assert isinstance(bool_arg.target, DimAdd) +def test_dim_call_arithmetic_is_pure_construction_in_both_directions() -> None: + seq = DimVar("S_chain", 1, 1024) + base = seq - 1 + + expressions = ( + base + 8, + 8 + base, + base - 8, + 8 - base, + base * 8, + 8 * base, + base // 8, + 8 // base, + base % 8, + 8 % base, + ) + + assert all(isinstance(expr, Call) for expr in expressions) + assert [type(expr.target) for expr in expressions] == [ + DimAdd, + DimAdd, + DimSub, + DimSub, + DimMul, + DimMul, + DimFloorDiv, + DimFloorDiv, + DimMod, + DimMod, + ] + + +def test_non_dim_calls_do_not_gain_dimension_arithmetic() -> None: + tensor = TensorType.umat_tensor((8,), DType.f32) + value = Var(type=tensor, name="value") + ordinary = Call( + type=tensor, + target=Unary(kind=UnaryKind.NEG), + args=(value,), + ) + operations = ( + lambda: ordinary + 1, + lambda: 1 + ordinary, + lambda: ordinary - 1, + lambda: 1 - ordinary, + lambda: ordinary * 1, + lambda: 1 * ordinary, + lambda: ordinary // 1, + lambda: 1 // ordinary, + lambda: ordinary % 1, + lambda: 1 % ordinary, + ) + + for operation in operations: + with pytest.raises(TypeError): + operation() + + +def test_typeinfer_canonicalizes_equivalent_symbolic_shapes() -> None: + seq = DimVar("S_canonical", 1, 8193) + verbose = simplify_dim( + DimFloorDiv, + ( + simplify_dim( + DimAdd, + (simplify_dim(DimSub, (simplify_dim(DimAdd, (seq, 0)), 0)), 0), + ), + 1, + ), + ) + def layout(dim): + return ComposedLayout( + inner=Layout(shape=(dim,), strides=(dim,)), + offset=dim, + outer=ShardLayout( + layout=Layout(shape=(dim,), strides=(dim,)), + attrs=(Broadcast(),), + mesh=Mesh( + topologies=(Topology("cta", dim),), + layout=Layout(shape=(dim,), strides=(1,)), + ), + ), + ) + + verbose_type = TensorType( + shape=(verbose, 128), dtype=DType.f32, layout=layout(verbose), storage="gmem" + ) + direct_type = TensorType( + shape=(seq, 128), dtype=DType.f32, layout=layout(seq), storage="gmem" + ) + + inferred = TypeInferVisitor(TypeInferContext()).visit(Var(type=verbose_type, name="verbose")) + + assert inferred == direct_type + assert inferred.shape[0] is seq + + +def test_static_typeinfer_does_not_enter_dim_canonicalization( + monkeypatch: pytest.MonkeyPatch, +) -> None: + static = TensorType(shape=(8, 128), dtype=DType.f32, layout=None, storage="gmem") + + def fail_if_called(_): + raise AssertionError("static types must not enter isl canonicalization") + + monkeypatch.setattr(dim_substitute, "normalize_dim", fail_if_called) + + assert TypeInferVisitor(TypeInferContext()).visit(Var(type=static, name="static")) is static + + def test_a_fully_static_dim_has_one_canonical_int_representation() -> None: """Test a fully static dim has one canonical int representation. @@ -150,6 +244,26 @@ def test_a_fully_static_dim_has_one_canonical_int_representation() -> None: assert broadcast_shapes(sliced.shape, param.shape) == (1, 4, 32, 128) +def test_constant_dim_arithmetic_normalizes_when_stored_as_an_op_attribute() -> None: + constructed = simplify_dim(DimMul, (4, 2)) + assert isinstance(constructed, Call) + + reshape = Reshape(new_shape=(constructed, 16)) + + assert reshape.new_shape == (8, 16) + + +def test_static_op_attributes_do_not_enter_dim_normalization( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def fail_if_called(_): + raise AssertionError("static attributes must not enter isl normalization") + + monkeypatch.setattr(dim_isl, "normalize_dim", fail_if_called) + + assert Reshape(new_shape=(8, 16)).new_shape == (8, 16) + + def test_unary_propagates_dim_var_in_shape() -> None: """Test unary propagates dim var in shape. diff --git a/tests/ops/test_arange.py b/tests/ops/test_arange.py index 684bf3e0..1ee753c2 100644 --- a/tests/ops/test_arange.py +++ b/tests/ops/test_arange.py @@ -16,6 +16,7 @@ from tilefoundry.ir.hir.tensor.arange import Arange from tilefoundry.ir.types import DType, TensorType from tilefoundry.ir.types.dim import ceildiv +from tilefoundry.ir.types.dim_isl import normalize_dim from tilefoundry.ir.types.storage import StorageKind from tilefoundry.visitor_registry.contexts import TrafficBytes, TypeInferContext from tilefoundry.visitor_registry.visitors import TypeInferVisitor @@ -53,7 +54,7 @@ def test_arange_static_type_evaluation_and_cost(): def test_arange_symbolic_extent_resolves_from_runtime_shape(): call = _symbolic_arange.body assert isinstance(call, Call) and isinstance(call.target, Arange) - assert call.type.shape == (ceildiv(_N - 1, 2),) + assert call.type.shape == (normalize_dim(ceildiv(_N - 1, 2)),) assert call.type.dtype == DType.i32 actual = evaluate(_symbolic_arange, torch.zeros(8), device="cpu") diff --git a/tests/ops/test_reshape.py b/tests/ops/test_reshape.py index 33b01691..d45c9bc6 100644 --- a/tests/ops/test_reshape.py +++ b/tests/ops/test_reshape.py @@ -27,6 +27,7 @@ from tilefoundry.ir.types.dim import DimVar from tilefoundry.ir.types.shard import Layout, ShardLayout, make_mesh from tilefoundry.ir.types.shard.shard_layout import ( + Broadcast, Partial, Split, shard_layout_local_shape, @@ -178,6 +179,34 @@ def test_reshape_then_reshard_rmem_no_split_aliasing(): ) +def test_broadcast_reshape_and_reshard_order_agree_in_smem(): + source = make_tensor_type( + (1, 32, 128), + layout=Layout(shape=(1, 32, 128), strides=(4096, 128, 1)), + ) + packed_layout = ShardLayout( + layout=Layout(shape=(1, 32, 128), strides=None), + attrs=(Broadcast(),), + mesh=_M, + ) + reshaped_layout = ShardLayout( + layout=Layout(shape=(32, 128), strides=None), + attrs=(Broadcast(),), + mesh=_M, + ) + + resharded_then_reshaped = infer_call( + _reshape((32, 128)), + infer_call(Reshard(layout=packed_layout, storage=StorageKind.SMEM), source), + ) + reshaped_then_resharded = infer_call( + Reshard(layout=reshaped_layout, storage=StorageKind.SMEM), + infer_call(_reshape((32, 128)), source), + ) + + assert resharded_then_reshaped == reshaped_then_resharded + + _S = DimVar(name="seq_len", lo=1, hi=4096) diff --git a/tests/ops/test_slice.py b/tests/ops/test_slice.py index 9b820054..9015826f 100644 --- a/tests/ops/test_slice.py +++ b/tests/ops/test_slice.py @@ -17,9 +17,10 @@ from tilefoundry.ir.core.kinds import BinaryKind from tilefoundry.ir.hir.function import Function from tilefoundry.ir.hir.math.binary import Binary -from tilefoundry.ir.hir.tensor.slice import Slice +from tilefoundry.ir.hir.tensor.slice import Slice, slice_size from tilefoundry.ir.types import DType, TupleType, make_shard_tensor_type, make_tensor_type from tilefoundry.ir.types.dim import DimMul, DimVar, simplify_dim +from tilefoundry.ir.types.dim_isl import normalize_dim from tilefoundry.ir.types.shard import ComposedLayout, Layout, make_mesh from tilefoundry.ir.types.shard.shard_layout import ShardLayout, Split, shard_layout_of from tilefoundry.visitor_registry.contexts import CostContext, TrafficBytes @@ -62,6 +63,21 @@ def _windowed_shard(source, shape, strides) -> ShardLayout: ) +def test_reversed_static_window_is_an_empty_slice() -> None: + scalar = make_tensor_type((), DType.i64) + size = normalize_dim( + slice_size( + Constant(type=scalar, value=8), + Constant(type=scalar, value=4), + Constant(type=scalar, value=1), + ) + ) + + actual = _slice_type(make_tensor_type((8, 4), _F), (8, 0), (size, 4), (1, 1)) + + assert actual.shape == (0, 4) + + def test_slice_of_unbound_axis_preserves_the_shard_layout(): source = make_shard_tensor_type((16, 32), mesh=_M, attrs=(Split(0),)) actual = _slice_type( diff --git a/tests/parser/golden/hir_expressions.py b/tests/parser/golden/hir_expressions.py index 12a85d09..19aa6b8c 100644 --- a/tests/parser/golden/hir_expressions.py +++ b/tests/parser/golden/hir_expressions.py @@ -17,8 +17,8 @@ class HirExpressions: @func def dim_from_a_static_call( x: Tensor[(CTX_LEN,), "bf16"] - ) -> Tensor[(ceildiv(CTX_LEN, 128) * 128,), "bf16"]: - v0 = zeros(shape=(ceildiv(CTX_LEN, 128) * 128,), dtype="bf16", storage=gmem) + ) -> Tensor[((128 * ((CTX_LEN - 1) // 128)) + 128,), "bf16"]: + v0 = zeros(shape=((128 * ((CTX_LEN - 1) // 128)) + 128,), dtype="bf16", storage=gmem) return v0 @func @@ -142,6 +142,15 @@ def slice_strided_and_clamped( v4 = x[:, :, 1:10:3] return v4 + @func + def slice_to_symbolic_extents( + x: Tensor[(CTX_LEN, 128), "f32"] + ) -> Tensor[(CTX_LEN, 128), "f32"]: + v0 = 0 + v1 = 0 + v3 = x[:, :] + return v3 + @func def full_tile_window( x: Tensor[(8, 4), "f32"], diff --git a/tests/parser/programs.py b/tests/parser/programs.py index 715f979f..c2ef4be0 100644 --- a/tests/parser/programs.py +++ b/tests/parser/programs.py @@ -153,6 +153,12 @@ def index_counted_from_the_end(x: Tensor[(1, 4, 8), "f32"]) -> Tensor[(1, 4), "f def slice_strided_and_clamped(x: Tensor[(1, 4, 8), "f32"]) -> Tensor[(1, 4, 3), "f32"]: return x[:, :, 1:20:3] + @func + def slice_to_symbolic_extents( + x: Tensor[(CTX_LEN, _KD), "f32"], + ) -> Tensor[(CTX_LEN, _KD), "f32"]: + return x[0:CTX_LEN, 0:_KD] + @func def full_tile_window(x: Tensor[(8, 4), "f32"], seed: Tensor[(4, 4), "f32"]): out = add(seed, seed) # noqa: F405 @@ -484,6 +490,7 @@ class ParserProgram: "tuple-literal op input", "multi-output tuple unpack", "subscript indexing and slicing", + "slice endpoints at symbolic extents", "tile window", "window move by a compile-time offset", ), diff --git a/tests/parser/test_programs.py b/tests/parser/test_programs.py index ef3723f6..cc3d5f44 100644 --- a/tests/parser/test_programs.py +++ b/tests/parser/test_programs.py @@ -37,25 +37,29 @@ doubles_a_constant, returns_a_pair, ) +from tilefoundry.analysis.preflight import infer_authored_types from tilefoundry.analysis.walk import postorder from tilefoundry.dsl._stub_gen import regen_stubs from tilefoundry.evaluator import evaluate from tilefoundry.evaluator.dim import resolve_dim from tilefoundry.inspection import as_script from tilefoundry.ir.constraints import LayoutConstraint, constraint_metadata -from tilefoundry.ir.core import Call, Tuple, Var, get_metadata +from tilefoundry.ir.core import Call, Constant, Tuple, Var, get_metadata from tilefoundry.ir.hir.function import Function, elaborate from tilefoundry.ir.hir.grid_region import GridRegionExpr from tilefoundry.ir.hir.specialize import origin_of, specialize_function +from tilefoundry.ir.hir.tensor.slice import Slice from tilefoundry.ir.hir.verify import verify_function from tilefoundry.ir.types import DType, TupleType, make_shard_tensor_type -from tilefoundry.ir.types.dim import DimVar +from tilefoundry.ir.types.dim import DimAdd, DimVar from tilefoundry.ir.types.shard import Layout, Mesh, ShardLayout, Topology, make_mesh from tilefoundry.ir.types.shard.shard_layout import Broadcast, Partial, Split from tilefoundry.ir.types.storage import StorageKind from tilefoundry.parser import hir_parser from tilefoundry.parser.base import _ModuleCallee from tilefoundry.parser.sugar import parse_shard_layout_sugar +from tilefoundry.visitor_registry.contexts import TypeInferContext +from tilefoundry.visitor_registry.visitors import TypeInferVisitor @pytest.mark.parametrize("program", PROGRAMS, ids=[program.name for program in PROGRAMS]) @@ -414,6 +418,77 @@ def test_a_program_still_evaluates_to_what_torch_would_give() -> None: ) +def test_symbolic_slice_endpoints_preserve_shape_and_bind_across_a_call() -> None: + """Equivalent full windows retain the authored dimension at a call boundary.""" + prelude = ( + "from tilefoundry import func\n" + "from tilefoundry.dsl.tf import *\n" + "from tilefoundry.dsl import DimVar, Tensor\n" + ) + full_window = import_dsl( + prelude + + '\nCTX_LEN = DimVar("CTX_LEN", 1, 4097)\n' + '@func\ndef f(x: Tensor[(CTX_LEN, 128), "f32"]) ' + '-> Tensor[(CTX_LEN, 128), "f32"]:\n' + " return x[:, 0:128]\n" + ) + explicit_window = HirExpressions.lookup("slice_to_symbolic_extents") + assert isinstance(explicit_window.body, Call) + assert isinstance(explicit_window.body.target, Slice) + assert explicit_window.body.target.sizes == full_window.body.target.sizes + assert explicit_window.body.target.strides == full_window.body.target.strides + assert explicit_window.body.type == full_window.body.type + assert explicit_window.body.args[1] == full_window.body.args[1] + authored_extent = explicit_window.params[0].type.shape[0] + assert explicit_window.body.target.sizes[0] is authored_extent + assert explicit_window.body.type.shape[0] is authored_extent + + param = Var(type=explicit_window.return_type, name="window") + consumer = Function.build( + name="consume_window", + params=(param,), + body=param, + return_type=explicit_window.return_type, + ) + call = Call(type=consumer.return_type, target=consumer, args=(explicit_window.body,)) + assert TypeInferVisitor(TypeInferContext()).visit(call) == consumer.return_type + + dimension_start = import_dsl( + prelude + + '\nS = DimVar("m2_start_seq", 1, 4097)\n' + '@func\ndef f(x: Tensor[(S + 8, 128), "f32"]) -> Tensor[(8, 128), "f32"]:\n' + " return x[S:S + 8, 0:128]\n" + ) + start = dimension_start.body.args[1].elements[0] + assert isinstance(start, Call) and isinstance(start.target, DimAdd) + assert any(start.args[0] is arg for arg in dimension_start.params[0].type.shape[0].args) + assert isinstance(start.args[1], Constant) and start.args[1].value == 0 + + +def test_a_packed_cache_can_keep_a_symbolic_capacity_axis() -> None: + """A layer index and full symbolic windows can name every packed-cache axis.""" + packed = import_dsl( + "from tilefoundry import func, module\n" + "from tilefoundry.dsl.tf import *\n" + "from tilefoundry.dsl import DimVar, Tensor\n" + '\nCAP = DimVar("m2_capacity", 1, 4097)\n' + '@module(entry="run")\n' + "class PackedCache:\n" + " @func\n" + ' def run(kc: Tensor[(4, CAP, 8, 16), "f32"], ' + 'seed: Tensor[(CAP, 8, 16), "f32"]) -> Tensor[(CAP, 8, 16), "f32"]:\n' + " out = relu(seed)\n" + " for i in range(4):\n" + " out = add(out, kc[i, 0:CAP, 0:8, 0:16])\n" + " return out\n" + ) + entry = packed.entry_function() + infer_authored_types((entry,), packed) + capacity = entry.params[0].type.shape[1] + + assert entry.body.type.shape == (capacity, 8, 16) + + def _fused_reference(gu, seed): out = seed * 2 for lo in range(0, 4, 2): @@ -457,6 +532,15 @@ def test_a_range_scalar_and_a_runtime_endpoint_drive_a_slice_window() -> None: assert runtime_start.body.args[1].elements[0] is runtime_start.params[1] assert runtime_start.body.type.layout is None + shifted_start = import_dsl( + prelude + + '\n@func\ndef f(x: Tensor[(16, 4), "f32"], ' + 'start: Tensor[(), "i64"]) -> Tensor[(8, 4), "f32"]:\n' + " return x[start + 1:start + 9, :]\n" + ) + assert shifted_start.body.target.sizes == (8, 4) + assert shifted_start.body.type.shape == (8, 4) + x = torch.arange(32, dtype=torch.float32).reshape(8, 4) torch.testing.assert_close( evaluate(runtime_start, x, torch.tensor(2, dtype=torch.int64), device="cpu"), x[2:6, :]