From 5cedb0d5daf4ea0b628f92a165eaf0dd4b5715e9 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Mon, 21 Sep 2026 17:17:37 +0800 Subject: [PATCH 01/25] feat(hir): add value range metadata --- docs/spec/hir.md | 4 ++ docs/spec/passes.md | 3 +- docs/spec/types.md | 7 +++ docs/spec/visitor-registry.md | 21 ++++--- src/tilefoundry/analysis/check.py | 8 ++- src/tilefoundry/ir/core/__init__.py | 2 + src/tilefoundry/ir/core/metadata.py | 9 +++ src/tilefoundry/ir/hir/sharding/mesh_coord.py | 20 +++++- src/tilefoundry/ir/types/dim_isl.py | 60 +++++++++++------- src/tilefoundry/parser/ast_pattern.py | 10 ++- src/tilefoundry/parser/pattern_nodes.py | 2 + src/tilefoundry/passes/pass_manager.py | 10 ++- src/tilefoundry/visitor_registry/contexts.py | 8 +++ .../visitor_registry/registries.py | 17 +++++- src/tilefoundry/visitor_registry/visitors.py | 61 ++++++++++++++++--- 15 files changed, 197 insertions(+), 45 deletions(-) diff --git a/docs/spec/hir.md b/docs/spec/hir.md index de420615..d579e13b 100644 --- a/docs/spec/hir.md +++ b/docs/spec/hir.md @@ -63,6 +63,10 @@ class Function(Expr): rules are stated below. - mutable during the compiler's authorised typing, metadata, and specialization updates; fields that are not updated retain structural equality and hashing semantics. + - value-range metadata is refreshed by a non-type-writing whole-function + inference walk. Parser construction MAY attach an already-proved range to a leaf + whose value is fixed by lexical geometry; `MeshCoord` uses its concrete mesh + axis extent for this purpose. - a `Function` MUST NOT declare or override execution context. The `Module` that owns it declares the `Target` and the ordered `Topology` hierarchy its body runs against ([core-ir §1](./core-ir.md#1-module)). diff --git a/docs/spec/passes.md b/docs/spec/passes.md index c07ed349..89439333 100644 --- a/docs/spec/passes.md +++ b/docs/spec/passes.md @@ -208,7 +208,8 @@ parser already runs eager typeinfer ([parser](./parser.md)), so a pass runs, `PassManager` re-runs the relevant analysis on that pass's **dirty scope**: -- HIR-side: changed `Function`s rerun `typeinfer`. +- HIR-side: changed `Function`s rerun the existing checks and refresh available + `RangeMetadata` over the complete function without rewriting `.type`. - TIR-side: changed `PrimFunction`s rerun `verify`, which recursively retriggers `typeinfer` on the embedded Expr fields and refreshes their `.type`. diff --git a/docs/spec/types.md b/docs/spec/types.md index 9b0f82e9..35d288a7 100644 --- a/docs/spec/types.md +++ b/docs/spec/types.md @@ -527,6 +527,13 @@ def ceildiv(a, b) -> Expr: operations, and MUST reject other values. - `ceildiv(a, b)` MUST compose the existing add, subtract, and floor-divide operations; it does not introduce a distinct Op. + - `dim_range(value)` MUST return conservative half-open bounds from + `RangeMetadata` before attempting structural dimension arithmetic. A value + with neither stored nor structurally derivable bounds returns `None`. + Unsupported symbolic divisors remain an error rather than an unknown range. + - A bounded non-dimension `Expr` leaf in dimension arithmetic MUST become one + identity-deduplicated isl parameter carrying its stored bounds. An unbounded + leaf remains an unconstrained parameter for consumers that permit one. --- diff --git a/docs/spec/visitor-registry.md b/docs/spec/visitor-registry.md index 20132e98..1ada6b38 100644 --- a/docs/spec/visitor-registry.md +++ b/docs/spec/visitor-registry.md @@ -238,7 +238,9 @@ def register_typeinfer(op_cls: type[Op]): ... # decorator: register a typein ``` - constraints: - - handler signature is `(call: Call, ctx: TypeInferContext) -> Type`. + - handler signature is `(call: Call, ctx: TypeInferContext) -> Type | TypeInferResults`. + A bare `Type` states no value range; the decorator normalizes it to + `TypeInferResults(type)` without changing the other registries. ```python # example @@ -251,7 +253,7 @@ Visitor: ```python class TypeInferVisitor(ExprVisitor[Type]): - def __init__(self, *, memo=None, owns_body=True): ... + def __init__(self, *, memo=None, owns_body=True, ranges=False): ... def visit(self, expr: Expr, ctx: TypeInferContext) -> Type: ... def visit_leaf_Var(self, var: Var, operands, ctx): ... def visit_leaf_Constant(self, c: Constant, operands, ctx): ... @@ -261,7 +263,7 @@ class TypeInferVisitor(ExprVisitor[Type]): def visit_MeshRegion(self, region, ctx): ... def visit_leaf_ShapeOf(self, shape_of: ShapeOf, operands, ctx) -> Type: ... -def inference_type(expr: Expr, ctx: TypeInferContext | None = None) -> Type: ... +def inference_type(expr: Expr, ctx: TypeInferContext | None = None, *, ranges=False) -> Type: ... ``` - constraints: @@ -286,13 +288,16 @@ def inference_type(expr: Expr, ctx: TypeInferContext | None = None) -> Type: ... replaced child context. The region result type is the body's type. - `visit_leaf_ShapeOf` returns the node's declared rank-0 i32 type. - `inference_type` creates a fresh non-owning visitor and returns the inferred - type without writing it to `expr.type`; traversal-wide inference continues - to use its own shared visitor and memo. + type without writing it to `expr.type` by default. With `ranges=True`, it + replaces or removes each `RangeMetadata` over the complete tree without + rewriting stored types. Partial parser inference never writes ranges. Lifecycle: parser builds a `TypeInferContext` and infers each newly built call -at parse time (see [parser](./parser.md)). The analysis preflight then walks -each authored body and writes every `Expr.type` in place before -consumers run. The visitor and its context share the current scope's memo; +at parse time (see [parser](./parser.md)). Once a complete HIR Function is +formed, and after a pass replaces one, whole-function inference refreshes its +available ranges. Analysis does the same after cloning its complete view. This +range-only refresh deliberately leaves the existing type lifecycle unchanged. +The visitor and its context share the current scope's memo; `type_of` is a constant-time lookup used to expose bindings to handlers. ### 4.1 Access relation service — `access_relation` diff --git a/src/tilefoundry/analysis/check.py b/src/tilefoundry/analysis/check.py index 05b4b40e..13107804 100644 --- a/src/tilefoundry/analysis/check.py +++ b/src/tilefoundry/analysis/check.py @@ -14,6 +14,7 @@ BindingMetadata, Call, Expr, + RangeMetadata, Var, ) from tilefoundry.ir.core.module import ( @@ -67,6 +68,7 @@ MemoryMetadata, PerformanceMetadata, PerformanceSummaryMetadata, + RangeMetadata, RooflineMetadata, TrafficMetadata, } @@ -595,7 +597,11 @@ def check_program( f"got {budget!r}" ) derived = InlineCloner(module, function, budget).clone() - inference_type(derived.body, TypeInferContext(scope=FunctionScope(module, derived))) + inference_type( + derived, + TypeInferContext(scope=FunctionScope(module, derived)), + ranges=True, + ) _require_concrete_geometry(module, derived, error_type=AnalysisError) target = module.resolve_target() for topology in module.effective_topologies(): diff --git a/src/tilefoundry/ir/core/__init__.py b/src/tilefoundry/ir/core/__init__.py index dd5c8912..e8cd5a19 100644 --- a/src/tilefoundry/ir/core/__init__.py +++ b/src/tilefoundry/ir/core/__init__.py @@ -7,6 +7,7 @@ from .metadata import ( BindingMetadata, IRMetadata, + RangeMetadata, SourceSpanMetadata, attach_metadata, binding_name, @@ -39,6 +40,7 @@ "Call", "Tuple", "IRMetadata", + "RangeMetadata", "BindingMetadata", "SourceSpanMetadata", "attach_metadata", diff --git a/src/tilefoundry/ir/core/metadata.py b/src/tilefoundry/ir/core/metadata.py index 5d6340e7..e2cf2d79 100644 --- a/src/tilefoundry/ir/core/metadata.py +++ b/src/tilefoundry/ir/core/metadata.py @@ -28,6 +28,14 @@ class SourceSpanMetadata(IRMetadata): end_column: int | None = None +@dataclass(frozen=True) +class RangeMetadata(IRMetadata): + """Conservative half-open value bounds ``[lo, hi)`` for one expression.""" + + lo: int + hi: int + + def get_metadata[T: IRMetadata](expr: "Expr", cls: type[T]) -> T | None: """Return the metadata whose concrete class is ``cls``, if present.""" return next((value for value in expr.metadata if type(value) is cls), None) @@ -123,6 +131,7 @@ def remove_metadata(expr: "Expr", cls: type[IRMetadata]) -> "Expr": __all__ = [ "IRMetadata", "BindingMetadata", + "RangeMetadata", "SourceSpanMetadata", "binding_name", "describe_expr", diff --git a/src/tilefoundry/ir/hir/sharding/mesh_coord.py b/src/tilefoundry/ir/hir/sharding/mesh_coord.py index b5be907a..d3340971 100644 --- a/src/tilefoundry/ir/hir/sharding/mesh_coord.py +++ b/src/tilefoundry/ir/hir/sharding/mesh_coord.py @@ -9,6 +9,7 @@ from tilefoundry.ir.core.pattern import Scalar from tilefoundry.ir.core.register import register_op from tilefoundry.ir.types import DType, TensorType +from tilefoundry.ir.types.shape_helpers import static_dim_value from tilefoundry.ir.types.shard.mesh import Mesh from tilefoundry.ir.types.shard.scope_match import covered_by_scope from tilefoundry.ir.types.storage import StorageKind @@ -17,6 +18,7 @@ measures_without_reading, register_access_relation, ) +from tilefoundry.visitor_registry.contexts import TypeInferResults @register_op @@ -37,7 +39,7 @@ class MeshCoord(Op): @register_typeinfer(MeshCoord) -def _(call: "Call", ctx: "TypeInferContext") -> TensorType: +def _(call: "Call", ctx: "TypeInferContext") -> TypeInferResults: """A coordinate is one number about this unit, so it carries no placement.""" if not isinstance(call.target.mesh, Mesh): ctx.error(call, "MeshCoord.mesh must be a Mesh") @@ -45,7 +47,21 @@ def _(call: "Call", ctx: "TypeInferContext") -> TensorType: ctx.error(call, "MeshCoord.mesh must be bound by the current mesh scope") if not call.args: ctx.error(call, "missing required input 'axis'") - return TensorType(shape=(), dtype=DType.i64, layout=None, storage=StorageKind.RMEM) + shape = call.target.mesh.layout.shape + axis = static_dim_value(call.args[0]) + if axis is not None and not 0 <= axis < len(shape): + ctx.error(call, f"axis {axis} is out of range for rank-{len(shape)} mesh") + extents = (shape[axis],) if axis is not None else shape + bounds = ( + (0, max(extents)) + if extents + and all(isinstance(extent, int) and not isinstance(extent, bool) for extent in extents) + else None + ) + return TypeInferResults( + TensorType(shape=(), dtype=DType.i64, layout=None, storage=StorageKind.RMEM), + value_range=bounds, + ) @register_eval(MeshCoord) diff --git a/src/tilefoundry/ir/types/dim_isl.py b/src/tilefoundry/ir/types/dim_isl.py index 41e7728d..d73e3774 100644 --- a/src/tilefoundry/ir/types/dim_isl.py +++ b/src/tilefoundry/ir/types/dim_isl.py @@ -4,8 +4,9 @@ import isl -from tilefoundry.ir.core.expr import Call, Constant, Var +from tilefoundry.ir.core.expr import Call, Constant, Expr, Var from tilefoundry.ir.core.kinds import BinaryKind +from tilefoundry.ir.core.metadata import RangeMetadata, get_metadata from .dim import ( _DIM_OP_TYPES, @@ -52,20 +53,22 @@ def _bind_param( raise ValueError( f"DimVar {name!r} used with conflicting bounds {previous} vs {bound}" ) - elif identities is not None: + else: + stored = get_metadata(value, RangeMetadata) if isinstance(value, Expr) else None + if stored is None and identities is None: + raise TypeError(f"unsupported ShapeDim {type(value).__name__}") key = id(value) - known = identities.get(key) + known = identities.get(key) if identities is not None else None if known is not None: return known - index = len(identities) + index = len(identities) if identities is not None else len(params) 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__}") + if identities is not None: + identities[key] = name + bound = (stored.lo, stored.hi) if stored is not None else None params[name] = bound if param_map is not None: @@ -160,7 +163,7 @@ def _dim_range_visitor_type(): if _DIM_RANGE_VISITOR_TYPE is None: from tilefoundry.ir.visitor import ExprVisitor # noqa: PLC0415 - class _DimRangeVisitor(ExprVisitor[tuple[int, int]]): + class _DimRangeVisitor(ExprVisitor[tuple[int, int] | None]): def visit_Constant(self, value: Constant, ctx=None) -> tuple[int, int]: number = int(value.value) return number, number + 1 @@ -168,12 +171,16 @@ def visit_Constant(self, value: Constant, ctx=None) -> tuple[int, int]: def visit_DimVar(self, value: DimVar, ctx=None) -> tuple[int, int]: return value.lo, value.hi - def visit_Call(self, value: Call, ctx=None) -> tuple[int, int]: + def visit_Call(self, value: Call, ctx=None) -> tuple[int, int] | None: if type(value.target) is DimMul: a, b = value.args if not (_is_const(a) or _is_const(b)): - alo, ahi = self.visit(a, ctx) - blo, bhi = self.visit(b, ctx) + a_bounds = self.visit(a, ctx) + b_bounds = self.visit(b, ctx) + if a_bounds is None or b_bounds is None: + return None + alo, ahi = a_bounds + blo, bhi = b_bounds corners = ( alo * blo, alo * (bhi - 1), @@ -182,7 +189,9 @@ def visit_Call(self, value: Call, ctx=None) -> tuple[int, int]: ) return min(corners), max(corners) + 1 params: dict[str, tuple[int, int] | None] = {} - expr = _range_expr(value, params) + expr = _range_expr(value, params, identities={}) + if any(bound is None for bound in params.values()): + return None prefix = f"[{', '.join(params)}] -> " if params else "" pw_aff = isl.pw_aff(prefix + f"{{ [{expr}] }}") if params: @@ -194,12 +203,12 @@ def visit_Call(self, value: Call, ctx=None) -> tuple[int, int]: 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 default_visit(self, value, ctx=None) -> tuple[int, int]: + def default_visit(self, value, ctx=None) -> tuple[int, int] | None: if isinstance(value, bool): raise TypeError("ShapeDim must not be bool") if isinstance(value, int): return value, value + 1 - raise TypeError(f"unsupported ShapeDim {type(value).__name__}") + return None _DIM_RANGE_VISITOR_TYPE = _DimRangeVisitor return _DIM_RANGE_VISITOR_TYPE @@ -309,20 +318,22 @@ def normalize_dim_entries(value): return value -def dim_range(dim) -> tuple[int, int]: +def dim_range(dim) -> tuple[int, int] | None: """Return conservative half-open value bounds ``[lo, hi)`` for *dim*.""" + stored = get_metadata(dim, RangeMetadata) if isinstance(dim, Expr) else None + if stored is not None: + return stored.lo, stored.hi return _dim_range_visitor_type()().visit(dim) 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]] = {} + bounds: dict[str, tuple[int, int] | None] = {} seen: dict = {} names: list[str] = [] - def bind(name: str, dim, lo: int, hi: int) -> None: - bound = (lo, hi) + def bind(name: str, dim, bound: tuple[int, int] | None) -> None: previous = bounds.get(name) if previous is not None and previous != bound: raise ValueError( @@ -343,20 +354,23 @@ def bind(name: str, dim, lo: int, hi: int) -> None: 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) + 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) + bind(name, extent, dim_range(extent)) 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] + constraints += [ + f"{bound[0]} <= {name} < {bound[1]}" + for name in names + if (bound := bounds[name]) is not None + ] prefix = f"[{', '.join(names)}] -> " if names else "" if not dims: return isl.set(prefix + "{ [] }"), param_map diff --git a/src/tilefoundry/parser/ast_pattern.py b/src/tilefoundry/parser/ast_pattern.py index 4925ef44..12a3825f 100644 --- a/src/tilefoundry/parser/ast_pattern.py +++ b/src/tilefoundry/parser/ast_pattern.py @@ -91,7 +91,7 @@ from tilefoundry.ir.visitor import BindingSubstitutionCloner from tilefoundry.target import MemoryHierarchyFacts, Target, UnsupportedCapabilityError from tilefoundry.visitor_registry.contexts import FunctionScope, TypeInferContext -from tilefoundry.visitor_registry.visitors import TypeInferVisitor +from tilefoundry.visitor_registry.visitors import TypeInferVisitor, inference_type T = TypeVar("T") _TYPE_INFER_CONTEXT = "" @@ -286,6 +286,7 @@ def attach_authored_metadata(value: object, node: ast.AST, context: "MatchContex TupleType=TupleType, TypeInferContext=TypeInferContext, FunctionScope=FunctionScope, + inference_type=inference_type, TypeInferVisitor=TypeInferVisitor, TupleGetItem=TupleGetItem, Unary=Unary, @@ -1192,6 +1193,13 @@ def _finalize(self, cls: type) -> object: for function in functions: if isinstance(function, runtime.Function): verify_function(function, module=result) + runtime.inference_type( + function, + runtime.TypeInferContext( + scope=runtime.FunctionScope(result, function) + ), + ranges=True, + ) elif isinstance(function, runtime.PrimFunction): verify_prim_function(function, module_fns=result) return result diff --git a/src/tilefoundry/parser/pattern_nodes.py b/src/tilefoundry/parser/pattern_nodes.py index 9cdec03f..b6b5e0e3 100644 --- a/src/tilefoundry/parser/pattern_nodes.py +++ b/src/tilefoundry/parser/pattern_nodes.py @@ -25,6 +25,7 @@ from tilefoundry.ir.constraints.layout import _LAYOUT_WILDCARD from tilefoundry.ir.core import ( BindingMetadata, + RangeMetadata, attach_metadata, get_metadata, ) @@ -3162,6 +3163,7 @@ def construct(match, children, context): return cached index = _constant(axis) coordinate = _infer_call(runtime.MeshCoord(mesh=mesh), (index,), context) + attach_metadata(coordinate, RangeMetadata(0, extent)) context.function.state.mesh_coordinates[cache_key] = coordinate return coordinate diff --git a/src/tilefoundry/passes/pass_manager.py b/src/tilefoundry/passes/pass_manager.py index ceb74f26..6e2413bb 100644 --- a/src/tilefoundry/passes/pass_manager.py +++ b/src/tilefoundry/passes/pass_manager.py @@ -16,6 +16,8 @@ from tilefoundry.ir.hir.verify import verify_function as verify_hir_function from tilefoundry.ir.tir.prim_function import PrimFunction from tilefoundry.ir.tir.verify import verify_prim_function +from tilefoundry.visitor_registry.contexts import FunctionScope, TypeInferContext +from tilefoundry.visitor_registry.visitors import inference_type from .pass_base import Pass @@ -51,9 +53,15 @@ def _post_pass_recheck(self, prev: Module, curr: Module) -> None: if prev_by_name.get(fn.name) is fn: continue if isinstance(fn, HirFunction): - verify_hir_function(fn) + verify_hir_function(fn, module=curr) + inference_type( + fn, + TypeInferContext(scope=FunctionScope(curr, fn)), + ranges=True, + ) elif isinstance(fn, PrimFunction): verify_prim_function(fn, module_fns=curr) + def _check_requires(self) -> None: seen: set[str] = set() for p in self.passes: diff --git a/src/tilefoundry/visitor_registry/contexts.py b/src/tilefoundry/visitor_registry/contexts.py index f51c5fc8..8d737110 100644 --- a/src/tilefoundry/visitor_registry/contexts.py +++ b/src/tilefoundry/visitor_registry/contexts.py @@ -19,6 +19,14 @@ from tilefoundry.ir.types.utils import local_type_of +@dataclass(frozen=True) +class TypeInferResults: + """The type and optional value range produced by one inference rule.""" + + type: Type + value_range: tuple[int, int] | None = None + + @dataclass(frozen=True) class FunctionScope: """Where a walk is reading: one Module tree, and whose body it is in. diff --git a/src/tilefoundry/visitor_registry/registries.py b/src/tilefoundry/visitor_registry/registries.py index f486fe6d..881589d8 100644 --- a/src/tilefoundry/visitor_registry/registries.py +++ b/src/tilefoundry/visitor_registry/registries.py @@ -77,11 +77,26 @@ class Role(Enum): """Keyed by ``(target, role, class)``; see :func:`register_codegen`.""" -register_typeinfer = typeinfer_registry.decorator() register_verify_stmt = verify_stmt_registry.decorator() register_cost_evaluator = cost_evaluator_registry.decorator() +def register_typeinfer(cls: type) -> Callable[[Callable], Callable]: + """Register one type rule, normalizing its result without affecting peers.""" + + def decorator(fn: Callable) -> Callable: + def wrapped(*args, **kwargs): + from .contexts import TypeInferResults # noqa: PLC0415 + + result = fn(*args, **kwargs) + return result if isinstance(result, TypeInferResults) else TypeInferResults(result) + + typeinfer_registry.register(cls, wrapped) + return fn + + return decorator + + def register_codegen( target: type["Target"], role: Role, cls: type ) -> Callable[[Callable], Callable]: diff --git a/src/tilefoundry/visitor_registry/visitors.py b/src/tilefoundry/visitor_registry/visitors.py index 8c763140..9eecd1be 100644 --- a/src/tilefoundry/visitor_registry/visitors.py +++ b/src/tilefoundry/visitor_registry/visitors.py @@ -14,6 +14,11 @@ from dataclasses import replace from tilefoundry.ir.core.expr import Call, Constant, Expr, Tuple, Var +from tilefoundry.ir.core.metadata import ( + RangeMetadata, + attach_metadata, + detach_metadata, +) from tilefoundry.ir.hir.function import Function from tilefoundry.ir.hir.loop_region import LoopRegion from tilefoundry.ir.hir.mesh_region import MeshRegion @@ -21,6 +26,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.callable_type import callable_type_for from tilefoundry.ir.types.shard.mesh import composed from tilefoundry.ir.types.shard.scope_match import covered_by_scope, storage_reaches from tilefoundry.ir.types.shard.shard_layout import ShardLayout @@ -32,7 +38,9 @@ from .contexts import ( Cost, CostContext, + FunctionScope, TypeInferContext, + TypeInferResults, VerifyContext, ) from .registries import ( @@ -56,11 +64,14 @@ class TypeInferVisitor(ExprVisitor[Type]): stale ``expr.type``. """ - def __init__(self, *, memo=None, owns_body: bool = True) -> None: + def __init__( + self, *, memo=None, owns_body: bool = True, ranges: bool = False + ) -> None: super().__init__(memo=memo) self._memo_supplied = memo is not None self._visit_depth = 0 self._owns_body = owns_body + self._ranges = ranges def visit(self, expr: Expr, ctx: TypeInferContext) -> Type: """Derive one type while preserving the active execution domain.""" @@ -72,9 +83,18 @@ def visit(self, expr: Expr, ctx: TypeInferContext) -> Type: self._memo = ctx.memo self._visit_depth += 1 try: - result = canonicalize_dims(super().visit(expr, ctx)) + results = super().visit(expr, ctx) + if not isinstance(results, TypeInferResults): + results = TypeInferResults(results) + result = canonicalize_dims(results.type) + self._memo[id(expr)] = (expr, result) if self._owns_body: expr.type = result + if self._ranges: + if results.value_range is None: + detach_metadata(expr, RangeMetadata) + else: + attach_metadata(expr, RangeMetadata(*results.value_range)) return result finally: self._visit_depth -= 1 @@ -155,7 +175,7 @@ def _call_function( cached = ctx.instantiated_memo.get(key) if cached is not None: return cached - result = TypeInferVisitor(memo=memo, owns_body=False).visit( + result = TypeInferVisitor(memo=memo, owns_body=False, ranges=False).visit( callee.body, ctx.for_callee(callee) ) ctx.instantiated_memo[key] = result @@ -177,7 +197,11 @@ def visit_LoopRegion(self, region: LoopRegion, ctx: TypeInferContext) -> Type: id(region.induction_var): (region.induction_var, region.induction_var.annotation), **{id(phi): (phi, type_) for phi, type_ in zip(region.carried_args, inits)}, } - inner = TypeInferVisitor(memo=memo, owns_body=self._owns_body) + inner = TypeInferVisitor( + memo=memo, + owns_body=self._owns_body, + ranges=self._ranges, + ) body_type = inner.visit(region.body, ctx) for y in region.yield_values: inner.visit(y, ctx) @@ -223,6 +247,24 @@ def visit_MeshRegion(self, expr: MeshRegion, ctx: TypeInferContext) -> Type: mesh = composed((ctx.current_mesh, expr.mesh)) if ctx.current_mesh else expr.mesh return self.visit(expr.body, replace(ctx, current_mesh=mesh, memo=memo)) + def visit_Function(self, fn: Function, ctx: TypeInferContext) -> Type: + """Refresh one complete function after binding its parameter types.""" + if ctx.scope is not None and ctx.scope.function is not fn: + ctx = replace(ctx, scope=FunctionScope(ctx.scope.module, fn)) + memo = {id(param): (param, param.annotation) for param in fn.params} + if fn.body is not None: + TypeInferVisitor( + memo=memo, + owns_body=self._owns_body, + ranges=self._ranges, + ).visit(fn.body, replace(ctx, memo=memo)) + for nested in (*fn.variants, *(converter for _, converter in fn.converters)): + TypeInferVisitor( + owns_body=self._owns_body, + ranges=self._ranges, + ).visit(nested, ctx) + return callable_type_for(fn.params, fn.return_type) + def visit_leaf_ShapeOf( self, shape_of: ShapeOf, _operands, ctx: TypeInferContext ) -> Type: @@ -237,9 +279,14 @@ def default_visit_leaf(self, expr: Expr, _operands, ctx: TypeInferContext) -> Ty ctx.error(expr, f"no typeinfer rule for Expr subclass {type(expr).__name__}") -def inference_type(expr: Expr, ctx: TypeInferContext | None = None) -> Type: - """Infer and return *expr*'s type without writing it back to the IR.""" - return TypeInferVisitor(owns_body=False).visit( +def inference_type( + expr: Expr, + ctx: TypeInferContext | None = None, + *, + ranges: bool = False, +) -> Type: + """Infer *expr*; ``ranges=True`` refreshes value ranges but not stored types.""" + return TypeInferVisitor(owns_body=False, ranges=ranges).visit( expr, ctx if ctx is not None else TypeInferContext() ) From 307c5978f75877cd4de9b3e08d8a771ff634ccb4 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Mon, 21 Sep 2026 17:27:38 +0800 Subject: [PATCH 02/25] fix(hir): preserve cached value ranges --- src/tilefoundry/ir/types/dim_isl.py | 15 +++++++++------ src/tilefoundry/passes/pass_manager.py | 2 +- src/tilefoundry/visitor_registry/visitors.py | 3 ++- 3 files changed, 12 insertions(+), 8 deletions(-) diff --git a/src/tilefoundry/ir/types/dim_isl.py b/src/tilefoundry/ir/types/dim_isl.py index d73e3774..661e9a28 100644 --- a/src/tilefoundry/ir/types/dim_isl.py +++ b/src/tilefoundry/ir/types/dim_isl.py @@ -55,19 +55,18 @@ def _bind_param( ) else: stored = get_metadata(value, RangeMetadata) if isinstance(value, Expr) else None - if stored is None and identities is None: + if identities is None: raise TypeError(f"unsupported ShapeDim {type(value).__name__}") key = id(value) - known = identities.get(key) if identities is not None else None + known = identities.get(key) if known is not None: return known - index = len(identities) if identities is not None else len(params) + index = len(identities) name = f"__tf_runtime_{index}" while name in params: index += 1 name = f"__tf_runtime_{index}" - if identities is not None: - identities[key] = name + identities[key] = name bound = (stored.lo, stored.hi) if stored is not None else None params[name] = bound @@ -327,7 +326,11 @@ def dim_range(dim) -> tuple[int, int] | None: def to_domain(extents: tuple) -> tuple: - """Build a bounded iteration domain and its isl-parameter ShapeDim map.""" + """Build an iteration domain and its isl-parameter ShapeDim map. + + A ``Call`` without a value range becomes an unconstrained parameter. Consumers + that require a bounded domain must reject that parameter explicitly. + """ param_map: dict[str, object] = {} bounds: dict[str, tuple[int, int] | None] = {} seen: dict = {} diff --git a/src/tilefoundry/passes/pass_manager.py b/src/tilefoundry/passes/pass_manager.py index 6e2413bb..7f9581a6 100644 --- a/src/tilefoundry/passes/pass_manager.py +++ b/src/tilefoundry/passes/pass_manager.py @@ -53,7 +53,7 @@ def _post_pass_recheck(self, prev: Module, curr: Module) -> None: if prev_by_name.get(fn.name) is fn: continue if isinstance(fn, HirFunction): - verify_hir_function(fn, module=curr) + verify_hir_function(fn) inference_type( fn, TypeInferContext(scope=FunctionScope(curr, fn)), diff --git a/src/tilefoundry/visitor_registry/visitors.py b/src/tilefoundry/visitor_registry/visitors.py index 9eecd1be..510ba073 100644 --- a/src/tilefoundry/visitor_registry/visitors.py +++ b/src/tilefoundry/visitor_registry/visitors.py @@ -75,6 +75,7 @@ def __init__( def visit(self, expr: Expr, ctx: TypeInferContext) -> Type: """Derive one type while preserving the active execution domain.""" + cached = id(expr) in self._memo outermost = self._visit_depth == 0 if outermost: if self._memo_supplied: @@ -90,7 +91,7 @@ def visit(self, expr: Expr, ctx: TypeInferContext) -> Type: self._memo[id(expr)] = (expr, result) if self._owns_body: expr.type = result - if self._ranges: + if self._ranges and not cached: if results.value_range is None: detach_metadata(expr, RangeMetadata) else: From 0dae3300cc5fcffd2f294d83c28513f04f2798e3 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Mon, 21 Sep 2026 17:32:47 +0800 Subject: [PATCH 03/25] fix(hir): narrow symbolic dimension detection --- docs/spec/hir.md | 4 +++- docs/spec/types.md | 4 ++++ src/tilefoundry/ir/types/substitute.py | 2 +- 3 files changed, 8 insertions(+), 2 deletions(-) diff --git a/docs/spec/hir.md b/docs/spec/hir.md index d579e13b..44e5d5eb 100644 --- a/docs/spec/hir.md +++ b/docs/spec/hir.md @@ -1690,4 +1690,6 @@ def is_concrete(fn: Function) -> bool: re-pointed instead would answer both with whichever was written last. - `residual_dims` and `dim_vars_reached` MUST inspect the whole function graph, including signatures, bodies, Op attributes, loop bounds, variants, - and called functions. `is_concrete` additionally checks the return type. + and called functions. `is_concrete` additionally checks the return type and + is false exactly when a reachable required extent still contains a `DimVar`; + runtime values without a `DimVar` are rejected later by their consumer. diff --git a/docs/spec/types.md b/docs/spec/types.md index 35d288a7..a8a9b61c 100644 --- a/docs/spec/types.md +++ b/docs/spec/types.md @@ -525,6 +525,10 @@ def ceildiv(a, b) -> Expr: - `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. + - `has_symbolic_dims(value)` MUST report whether a `DimVar` is reachable, not + whether dimension arithmetic contains a runtime `Expr`. A dimension call + recursively checks its operands; non-dimension expressions do not become + symbolic merely by participating in that arithmetic. - `ceildiv(a, b)` MUST compose the existing add, subtract, and floor-divide operations; it does not introduce a distinct Op. - `dim_range(value)` MUST return conservative half-open bounds from diff --git a/src/tilefoundry/ir/types/substitute.py b/src/tilefoundry/ir/types/substitute.py index 403628bd..62d1afb2 100644 --- a/src/tilefoundry/ir/types/substitute.py +++ b/src/tilefoundry/ir/types/substitute.py @@ -388,7 +388,7 @@ def has_symbolic_dims(value: object) -> bool: if isinstance(value, DimVar): return True if isinstance(value, Call) and isinstance(value.target, _DIM_OP_TYPES): - return True + return any(has_symbolic_dims(arg) for arg in value.args) if isinstance(value, TensorType): return has_symbolic_dims(value.shape) or has_symbolic_dims(value.layout) if isinstance(value, TupleType): From 1989df1d38846770611872c204a143e3f3ccf1d0 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Mon, 21 Sep 2026 17:46:50 +0800 Subject: [PATCH 04/25] feat(analysis): support bounded loop endpoints --- docs/spec/analysis.md | 13 +++ src/tilefoundry/analysis/scope.py | 105 ++++++++++++------- src/tilefoundry/utils/isl_utils.py | 70 ++++++++++--- src/tilefoundry/visitor_registry/visitors.py | 2 +- 4 files changed, 140 insertions(+), 50 deletions(-) diff --git a/docs/spec/analysis.md b/docs/spec/analysis.md index 83396a85..9484ed77 100644 --- a/docs/spec/analysis.md +++ b/docs/spec/analysis.md @@ -1218,6 +1218,19 @@ call site, source expressions shared by identity remain one shared expression in the clone; sharing never aliases the independently cloned body of another call site. +A loop `start` or `extent` MAY depend on the unit running it when every runtime +leaf carries a half-open value range. `Scope.domain` MUST retain the complete +affine expression and represent each such leaf as one identity-deduplicated isl +parameter constrained by that range. A leaf without a range MUST be rejected as +runtime-computed. `step` MUST remain a literal: a parametric stride requires a +product or modulus by two unknowns and has no isl Presburger representation. + +`cardinality` with bounded free parameters is a conservative upper bound: it +counts every feasible corner of the parameter box and returns the maximum. +`Scope.trips()` MUST instead fix child and parent domains to the same corner, +divide their counts there, and then take the maximum ratio; dividing two +independently maximized counts is not a valid trip bound. + ### 2.2 Target-selected Analyzers ```python diff --git a/src/tilefoundry/analysis/scope.py b/src/tilefoundry/analysis/scope.py index 8c9607ed..6f83baba 100644 --- a/src/tilefoundry/analysis/scope.py +++ b/src/tilefoundry/analysis/scope.py @@ -14,11 +14,11 @@ from tilefoundry.ir.hir.tensor.reshape import Reshape from tilefoundry.ir.hir.tensor.slice import Slice from tilefoundry.ir.types import TensorType -from tilefoundry.ir.types.dim import DimVar +from tilefoundry.ir.types.dim_isl import _range_expr from tilefoundry.ir.types.shape_helpers import static_dim_value from tilefoundry.ir.types.utils import local_type_of from tilefoundry.ir.visitor import expr_children -from tilefoundry.utils.isl_utils import cardinality +from tilefoundry.utils.isl_utils import _count, _param_corners, cardinality from tilefoundry.visitor_registry.access_relation import ( AccessRelations, access_relation_registry, @@ -91,17 +91,18 @@ def trips(self) -> int: return cached if self.parent is None: return 1 - if isinstance(self.owner, LoopRegion): - start, extent, step = self.owner.start, self.owner.extent, self.owner.step - if all(isinstance(value, int) for value in (start, extent, step)): - result = 1 if step <= 0 or extent <= start else -(-(extent - start) // step) - self._trips_cache = result - return result - count = cardinality(self.domain) - parent_count = cardinality(self.parent.domain) - if count is None or not parent_count: - return 1 - result = max(1, count // parent_count) + domain = self.domain + parent = self.parent.domain.align_params(domain.get_space()) + domain = domain.align_params(parent.get_space()) + corners = _param_corners(domain.params().intersect(parent.params())) + ratios = [] + for corner in corners or (): + count = _count(domain.intersect_params(corner)) + parent_count = _count(parent.intersect_params(corner)) + if count is None or not parent_count: + continue + ratios.append(max(1, count // parent_count)) + result = max(ratios, default=1) self._trips_cache = result return result @@ -222,6 +223,48 @@ def _induction_of(loop: LoopRegion) -> str: return getattr(loop.induction_var, "name", None) or "" +def _unbounded_loop_bound(loop: LoopRegion, which: str) -> None: + value = getattr(loop, which) + if which == "extent": + raise AnalysisError( + f"loop {_induction_of(loop)!r} has a trip count the program computes " + f"at run time from {value_label(value) or 'a value'!r}, so no " + f"per-occurrence total can be scaled by it; bind the extent to a " + f"literal, or state it as an open dimension" + ) + raise AnalysisError( + f"loop {_induction_of(loop)!r} takes its {which} from " + f"{value_label(value) or 'a value'!r}, which the program computes at run time; " + f"analysis needs a literal {which} or a stated value range" + ) + + +def _bound_or_param( + loop: LoopRegion, + which: str, + params: dict[str, tuple[int, int] | None], + param_map: dict[str, object], + identities: dict[int, str], +) -> str: + """Render one start/extent from bounded leaves, or reject an unknown value.""" + value = getattr(loop, which) + number = static_dim_value(value) + if number is not None: + return str(number) + try: + rendered = _range_expr( + value, + params, + param_map=param_map, + identities=identities, + ) + except (TypeError, ValueError, NotImplementedError, isl.Error): + _unbounded_loop_bound(loop, which) + if any(bound is None for bound in params.values()): + _unbounded_loop_bound(loop, which) + return rendered + + def _domain_for(owner: Function | LoopRegion, parent: Scope | None) -> isl.set: if isinstance(owner, Function): return isl.set("{ [] }") @@ -232,38 +275,26 @@ def _domain_for(owner: Function | LoopRegion, parent: Scope | None) -> isl.set: loops.append(cursor.owner) cursor = cursor.parent loops.reverse() - params: dict[str, DimVar] = {} + params: dict[str, tuple[int, int] | None] = {} + param_map: dict[str, object] = {} + identities: dict[int, str] = {} bounds: list[str] = [] for index, loop in enumerate(loops + [owner]): - start = static_dim_value(loop.start) + start = _bound_or_param(loop, "start", params, param_map, identities) + stop = _bound_or_param(loop, "extent", params, param_map, identities) step = static_dim_value(loop.step) - if start is None or step is None: - which = "start" if start is None else "step" - culprit = value_label(loop.start if start is None else loop.step) + if step is None: raise AnalysisError( - f"loop {_induction_of(loop)!r} takes its {which} from " - f"{culprit or 'a value'!r}, which the program computes at run time; " - f"analysis needs a literal {which}" + f"loop {_induction_of(loop)!r} takes its step from " + f"{value_label(loop.step) or 'a value'!r}; analysis needs a literal " + "step, because a parametric stride has no isl representation" ) - extent = loop.extent - if isinstance(extent, DimVar): - params[extent.name] = extent - stop = extent.name - else: - value = static_dim_value(extent) - if value is None: - raise AnalysisError( - f"loop {_induction_of(loop)!r} has a trip count the program computes " - f"at run time from {value_label(extent) or 'a value'!r}, so no " - f"per-occurrence total can be scaled by it; bind the extent to a " - f"literal, or state it as an open dimension" - ) - stop = str(value) bounds.append(f"{start} <= p{index} < {stop}") if step != 1: bounds.append(f"(p{index} - {start}) mod {step} = 0") - for name, dim in params.items(): - bounds.append(f"{dim.lo} <= {name} < {dim.hi}") + for name, bound in params.items(): + assert bound is not None + bounds.append(f"{bound[0]} <= {name} < {bound[1]}") names = ", ".join(f"p{index}" for index in range(len(loops) + 1)) prefix = f"[{', '.join(params)}] -> " if params else "" return isl.set(f"{prefix}{{ [{names}] : {' and '.join(bounds)} }}") diff --git a/src/tilefoundry/utils/isl_utils.py b/src/tilefoundry/utils/isl_utils.py index f57aaafb..63f94a7d 100644 --- a/src/tilefoundry/utils/isl_utils.py +++ b/src/tilefoundry/utils/isl_utils.py @@ -2,36 +2,82 @@ from __future__ import annotations +from itertools import product + import isl __all__ = ["as_multi_aff", "cardinality", "equates", "involved_dims"] -def cardinality(image: "isl.set") -> int | None: - """How many points *image* holds, or None when that is not a finite number. - - A box is counted from its bounds: every axis is an independent bounded - interval, so their lengths multiply, at a cost set by rank alone. ISL's own - count has a closed form for simple sets and loses it as rank grows: on one - 462M-point image, 4.0 ms at two dimensions and 2.3 s at four. The cost - tracks rank, not points. Anything that is not a box falls back to it. - """ +def _count(image: "isl.set") -> int | None: + """Count *image* after every parameter has been fixed by its caller.""" image = image.coalesce() if image.is_box(): - product = 1 + amount = 1 for axis in range(image.tuple_dim()): if not image.dim_is_bounded(isl.dim_type.SET, axis): break low, high = image.dim_min_val(axis), image.dim_max_val(axis) if not (low.is_int() and high.is_int()): break - product *= high.get_num_si() - low.get_num_si() + 1 + amount *= high.get_num_si() - low.get_num_si() + 1 else: - return product + return amount amount = image.count_val() return amount.get_num_si() if amount.is_int() else None +def _param_corners(image: "isl.set") -> tuple["isl.set", ...] | None: + """Fix every parameter to each feasible corner of its bounded box.""" + context = image.params() + count = context.dim(isl.dim_type.PARAM) + if not count: + return (context,) + axes = context.move_dims( + isl.dim_type.SET, + 0, + isl.dim_type.PARAM, + 0, + count, + ) + values: list[tuple[int, ...]] = [] + for axis in range(count): + if not axes.dim_is_bounded(isl.dim_type.SET, axis): + return None + low = axes.dim_min_val(axis) + high = axes.dim_max_val(axis) + if not (low.is_int() and high.is_int()): + return None + lo, hi = low.get_num_si(), high.get_num_si() + values.append((lo,) if lo == hi else (lo, hi)) + corners = [] + for point in product(*values): + corner = context + for axis, value in enumerate(point): + corner = corner.fix_si(isl.dim_type.PARAM, axis, value) + if not corner.is_empty(): + corners.append(corner) + return tuple(corners) + + +def cardinality(image: "isl.set") -> int | None: + """Return a finite point count, or a corner-wise upper bound with params. + + A box is counted from its bounds: every axis is an independent bounded + interval, so their lengths multiply, at a cost set by rank alone. ISL's own + count has a closed form for simple sets and loses it as rank grows: on one + 462M-point image, 4.0 ms at two dimensions and 2.3 s at four. The cost + tracks rank, not points. Anything that is not a box falls back to it. When + bounded parameters remain free, each feasible box corner is counted and the + maximum is returned; an unbounded parameter has no finite answer. + """ + corners = _param_corners(image) + if corners is None: + return None + counts = tuple(_count(image.intersect_params(corner)) for corner in corners) + return None if not counts or any(count is None for count in counts) else max(counts) + + def equates(relation: "isl.map", out_axis: int, in_dim: int) -> bool: """Whether *relation* everywhere sends domain dim *in_dim* to *out_axis*. diff --git a/src/tilefoundry/visitor_registry/visitors.py b/src/tilefoundry/visitor_registry/visitors.py index 510ba073..1506e493 100644 --- a/src/tilefoundry/visitor_registry/visitors.py +++ b/src/tilefoundry/visitor_registry/visitors.py @@ -75,13 +75,13 @@ def __init__( def visit(self, expr: Expr, ctx: TypeInferContext) -> Type: """Derive one type while preserving the active execution domain.""" - cached = id(expr) in self._memo outermost = self._visit_depth == 0 if outermost: if self._memo_supplied: ctx = replace(ctx, memo=self._memo) else: self._memo = ctx.memo + cached = id(expr) in self._memo self._visit_depth += 1 try: results = super().visit(expr, ctx) From 2433684cfba49afaa3dac2b563411129853c6c99 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Mon, 21 Sep 2026 17:55:24 +0800 Subject: [PATCH 05/25] fix(analysis): enumerate bounded parameter domains --- docs/spec/analysis.md | 11 ++--- src/tilefoundry/analysis/scope.py | 23 +++++----- src/tilefoundry/ir/types/dim_isl.py | 8 ++-- src/tilefoundry/utils/isl_utils.py | 66 ++++++++++++++++++----------- 4 files changed, 65 insertions(+), 43 deletions(-) diff --git a/docs/spec/analysis.md b/docs/spec/analysis.md index 9484ed77..8f20dac5 100644 --- a/docs/spec/analysis.md +++ b/docs/spec/analysis.md @@ -1225,11 +1225,12 @@ parameter constrained by that range. A leaf without a range MUST be rejected as runtime-computed. `step` MUST remain a literal: a parametric stride requires a product or modulus by two unknowns and has no isl Presburger representation. -`cardinality` with bounded free parameters is a conservative upper bound: it -counts every feasible corner of the parameter box and returns the maximum. -`Scope.trips()` MUST instead fix child and parent domains to the same corner, -divide their counts there, and then take the maximum ratio; dividing two -independently maximized counts is not a valid trip bound. +`cardinality` with bounded free parameters MUST enumerate every feasible integer +point and return the true maximum when the parameter box contains at most 4096 +points; for a larger box it MUST return unknown. `Scope.trips()` MUST likewise +fix child and parent domains to the same parameter point, divide their counts +there, and then take the maximum ratio; dividing two independently maximized +counts is not a valid trip bound. ### 2.2 Target-selected Analyzers diff --git a/src/tilefoundry/analysis/scope.py b/src/tilefoundry/analysis/scope.py index 6f83baba..2036ff91 100644 --- a/src/tilefoundry/analysis/scope.py +++ b/src/tilefoundry/analysis/scope.py @@ -14,11 +14,11 @@ from tilefoundry.ir.hir.tensor.reshape import Reshape from tilefoundry.ir.hir.tensor.slice import Slice from tilefoundry.ir.types import TensorType -from tilefoundry.ir.types.dim_isl import _range_expr +from tilefoundry.ir.types.dim_isl import range_expr from tilefoundry.ir.types.shape_helpers import static_dim_value from tilefoundry.ir.types.utils import local_type_of from tilefoundry.ir.visitor import expr_children -from tilefoundry.utils.isl_utils import _count, _param_corners, cardinality +from tilefoundry.utils.isl_utils import cardinality, count, param_points from tilefoundry.visitor_registry.access_relation import ( AccessRelations, access_relation_registry, @@ -94,14 +94,14 @@ def trips(self) -> int: domain = self.domain parent = self.parent.domain.align_params(domain.get_space()) domain = domain.align_params(parent.get_space()) - corners = _param_corners(domain.params().intersect(parent.params())) + points = param_points(domain.params().intersect(parent.params())) ratios = [] - for corner in corners or (): - count = _count(domain.intersect_params(corner)) - parent_count = _count(parent.intersect_params(corner)) - if count is None or not parent_count: + for point in points or (): + amount = count(domain.intersect_params(point)) + parent_count = count(parent.intersect_params(point)) + if amount is None or not parent_count: continue - ratios.append(max(1, count // parent_count)) + ratios.append(max(1, amount // parent_count)) result = max(ratios, default=1) self._trips_cache = result return result @@ -252,7 +252,7 @@ def _bound_or_param( if number is not None: return str(number) try: - rendered = _range_expr( + rendered = range_expr( value, params, param_map=param_map, @@ -293,7 +293,10 @@ def _domain_for(owner: Function | LoopRegion, parent: Scope | None) -> isl.set: if step != 1: bounds.append(f"(p{index} - {start}) mod {step} = 0") for name, bound in params.items(): - assert bound is not None + if bound is None: + raise AnalysisError( + f"loop domain parameter {name!r} has no stated value range" + ) bounds.append(f"{bound[0]} <= {name} < {bound[1]}") names = ", ".join(f"p{index}" for index in range(len(loops) + 1)) prefix = f"[{', '.join(params)}] -> " if params else "" diff --git a/src/tilefoundry/ir/types/dim_isl.py b/src/tilefoundry/ir/types/dim_isl.py index 661e9a28..9c42c219 100644 --- a/src/tilefoundry/ir/types/dim_isl.py +++ b/src/tilefoundry/ir/types/dim_isl.py @@ -188,7 +188,7 @@ def visit_Call(self, value: Call, ctx=None) -> tuple[int, int] | None: ) return min(corners), max(corners) + 1 params: dict[str, tuple[int, int] | None] = {} - expr = _range_expr(value, params, identities={}) + expr = range_expr(value, params, identities={}) if any(bound is None for bound in params.values()): return None prefix = f"[{', '.join(params)}] -> " if params else "" @@ -213,13 +213,14 @@ def default_visit(self, value, ctx=None) -> tuple[int, int] | None: return _DIM_RANGE_VISITOR_TYPE -def _range_expr( +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: + """Render *dim* as an isl expression and register its leaf parameters.""" return _range_expr_visitor_type()(params, param_map, identities).visit(dim) @@ -286,7 +287,7 @@ def normalize_dim(value): try: params: dict[str, tuple[int, int] | None] = {} param_map: dict[str, object] = {} - expr = _range_expr( + expr = range_expr( value, params, param_map=param_map, @@ -385,6 +386,7 @@ def bind(name: str, dim, bound: tuple[int, int] | None) -> None: "dim_range", "normalize_dim", "normalize_dim_entries", + "range_expr", "to_dim", "to_domain", ] diff --git a/src/tilefoundry/utils/isl_utils.py b/src/tilefoundry/utils/isl_utils.py index 63f94a7d..68b8b278 100644 --- a/src/tilefoundry/utils/isl_utils.py +++ b/src/tilefoundry/utils/isl_utils.py @@ -6,10 +6,19 @@ import isl -__all__ = ["as_multi_aff", "cardinality", "equates", "involved_dims"] +__all__ = [ + "as_multi_aff", + "cardinality", + "count", + "equates", + "involved_dims", + "param_points", +] +_PARAM_POINT_LIMIT = 4096 -def _count(image: "isl.set") -> int | None: + +def count(image: "isl.set") -> int | None: """Count *image* after every parameter has been fixed by its caller.""" image = image.coalesce() if image.is_box(): @@ -27,8 +36,8 @@ def _count(image: "isl.set") -> int | None: return amount.get_num_si() if amount.is_int() else None -def _param_corners(image: "isl.set") -> tuple["isl.set", ...] | None: - """Fix every parameter to each feasible corner of its bounded box.""" +def param_points(image: "isl.set") -> tuple["isl.set", ...] | None: + """Fix parameters to every feasible point of a small bounded integer box.""" context = image.params() count = context.dim(isl.dim_type.PARAM) if not count: @@ -40,7 +49,8 @@ def _param_corners(image: "isl.set") -> tuple["isl.set", ...] | None: 0, count, ) - values: list[tuple[int, ...]] = [] + values: list[range] = [] + box_points = 1 for axis in range(count): if not axes.dim_is_bounded(isl.dim_type.SET, axis): return None @@ -49,33 +59,39 @@ def _param_corners(image: "isl.set") -> tuple["isl.set", ...] | None: if not (low.is_int() and high.is_int()): return None lo, hi = low.get_num_si(), high.get_num_si() - values.append((lo,) if lo == hi else (lo, hi)) - corners = [] + choices = range(lo, hi + 1) + box_points *= len(choices) + if box_points > _PARAM_POINT_LIMIT: + return None + values.append(choices) + points = [] for point in product(*values): - corner = context + fixed = context for axis, value in enumerate(point): - corner = corner.fix_si(isl.dim_type.PARAM, axis, value) - if not corner.is_empty(): - corners.append(corner) - return tuple(corners) + fixed = fixed.fix_si(isl.dim_type.PARAM, axis, value) + if not fixed.is_empty(): + points.append(fixed) + return tuple(points) def cardinality(image: "isl.set") -> int | None: - """Return a finite point count, or a corner-wise upper bound with params. - - A box is counted from its bounds: every axis is an independent bounded - interval, so their lengths multiply, at a cost set by rank alone. ISL's own - count has a closed form for simple sets and loses it as rank grows: on one - 462M-point image, 4.0 ms at two dimensions and 2.3 s at four. The cost - tracks rank, not points. Anything that is not a box falls back to it. When - bounded parameters remain free, each feasible box corner is counted and the - maximum is returned; an unbounded parameter has no finite answer. + """Return a finite point count, maximizing over a small parameter box. + + A box is counted by multiplying its bounded axis lengths, at a cost set by + rank alone; anything else falls back to ISL's count. With bounded free + parameters, every feasible integer point in a box of at most 4096 points is + counted and the true maximum is returned. A larger or unbounded parameter + box has no answer. """ - corners = _param_corners(image) - if corners is None: + if not image.dim(isl.dim_type.PARAM): + return count(image) + points = param_points(image) + if points is None: + return None + counts = tuple(count(image.intersect_params(point)) for point in points) + if any(amount is None for amount in counts): return None - counts = tuple(_count(image.intersect_params(corner)) for corner in corners) - return None if not counts or any(count is None for count in counts) else max(counts) + return max(counts, default=0) def equates(relation: "isl.map", out_axis: int, in_dim: int) -> bool: From 67273b60bae0fc16d2857d49d37b2b1f5b900e31 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Mon, 21 Sep 2026 17:58:23 +0800 Subject: [PATCH 06/25] fix(analysis): reject unknown trip counts --- src/tilefoundry/analysis/scope.py | 8 +++++++- src/tilefoundry/utils/isl_utils.py | 8 ++++---- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/src/tilefoundry/analysis/scope.py b/src/tilefoundry/analysis/scope.py index 2036ff91..1ba6485d 100644 --- a/src/tilefoundry/analysis/scope.py +++ b/src/tilefoundry/analysis/scope.py @@ -95,8 +95,14 @@ def trips(self) -> int: parent = self.parent.domain.align_params(domain.get_space()) domain = domain.align_params(parent.get_space()) points = param_points(domain.params().intersect(parent.params())) + if points is None: + raise AnalysisError( + f"loop {_induction_of(self.owner)!r} has a parameter box exceeding " + "the 4096-point analysis limit, so its trip count cannot be " + "determined" + ) ratios = [] - for point in points or (): + for point in points: amount = count(domain.intersect_params(point)) parent_count = count(parent.intersect_params(point)) if amount is None or not parent_count: diff --git a/src/tilefoundry/utils/isl_utils.py b/src/tilefoundry/utils/isl_utils.py index 68b8b278..411f8ebd 100644 --- a/src/tilefoundry/utils/isl_utils.py +++ b/src/tilefoundry/utils/isl_utils.py @@ -39,19 +39,19 @@ def count(image: "isl.set") -> int | None: def param_points(image: "isl.set") -> tuple["isl.set", ...] | None: """Fix parameters to every feasible point of a small bounded integer box.""" context = image.params() - count = context.dim(isl.dim_type.PARAM) - if not count: + param_count = context.dim(isl.dim_type.PARAM) + if not param_count: return (context,) axes = context.move_dims( isl.dim_type.SET, 0, isl.dim_type.PARAM, 0, - count, + param_count, ) values: list[range] = [] box_points = 1 - for axis in range(count): + for axis in range(param_count): if not axes.dim_is_bounded(isl.dim_type.SET, axis): return None low = axes.dim_min_val(axis) From 4416a682cb1c2e8508811eeac276e15c7a665adf Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Mon, 21 Sep 2026 18:00:43 +0800 Subject: [PATCH 07/25] fix(analysis): bind access parameters by name --- src/tilefoundry/analysis/scope.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/src/tilefoundry/analysis/scope.py b/src/tilefoundry/analysis/scope.py index 1ba6485d..4761550a 100644 --- a/src/tilefoundry/analysis/scope.py +++ b/src/tilefoundry/analysis/scope.py @@ -334,6 +334,11 @@ def _bind_access( relation = relation.intersect_domain(scope_domain) params = dict(getattr(boundary.pattern, "parameters", ()) or ()) for name in list(params): + param_index = relation.find_dim_by_name(isl.dim_type.PARAM, name) + if param_index < 0: + raise AnalysisError( + f"access pattern parameter {name!r} is missing from its relation" + ) value = params[name] number = static_dim_value(value) if number is None: @@ -346,7 +351,7 @@ def _bind_access( term = _widest_allowed(relation, name, operand.type) exact = False if term is None: - relation = relation.project_out(isl.dim_type.PARAM, 0, 1) + relation = relation.project_out(isl.dim_type.PARAM, param_index, 1) continue else: term = type( @@ -356,7 +361,9 @@ def _bind_access( def placed(kind: str, sign: int, constant: int) -> isl.constraint: constraint = getattr(isl.constraint, f"alloc_{kind}")(local) - constraint = constraint.set_coefficient_si(isl.dim_type.PARAM, 0, sign) + constraint = constraint.set_coefficient_si( + isl.dim_type.PARAM, param_index, sign + ) if term.loop_axis is not None: constraint = constraint.set_coefficient_si( isl.dim_type.IN, term.loop_axis, -sign * term.stride @@ -368,7 +375,7 @@ def placed(kind: str, sign: int, constant: int) -> isl.constraint: else: relation = relation.add_constraint(placed("inequality", 1, -term.low)) relation = relation.add_constraint(placed("inequality", -1, term.high)) - relation = relation.project_out(isl.dim_type.PARAM, 0, 1) + relation = relation.project_out(isl.dim_type.PARAM, param_index, 1) try: held = local_type_of(operand.type) if narrow else operand.type except (TypeError, ValueError, NotImplementedError): From 9325b5543a68dbb87b27369243143687aac46084 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Mon, 21 Sep 2026 18:06:27 +0800 Subject: [PATCH 08/25] feat(analysis): distinguish bounded access parameters --- docs/spec/analysis.md | 5 ++ src/tilefoundry/analysis/allocation.py | 10 ++-- src/tilefoundry/analysis/scope.py | 69 ++++++++++++++++++-------- src/tilefoundry/utils/isl_utils.py | 10 ++++ 4 files changed, 71 insertions(+), 23 deletions(-) diff --git a/docs/spec/analysis.md b/docs/spec/analysis.md index 8f20dac5..a6d2e8ea 100644 --- a/docs/spec/analysis.md +++ b/docs/spec/analysis.md @@ -387,6 +387,11 @@ concrete arrangement is not reported: no address or per-value buffer identity is a conclusion of this analysis. `rmem` is not address-solved and reports only the largest single projected logical value. +An access relation that retains a parameter with a stated finite range remains +exact: it is neither widened nor unknown, and overlap MAY be proved from it. +Widened relations and relations with an unbounded parameter MUST NOT prove +overlap. + - constraints: - Placement MUST be settled for the addressable levels `gmem` and `smem` only, once per capacity domain that holds a buffer -- the whole target for a level diff --git a/src/tilefoundry/analysis/allocation.py b/src/tilefoundry/analysis/allocation.py index f2b3083c..9967a65b 100644 --- a/src/tilefoundry/analysis/allocation.py +++ b/src/tilefoundry/analysis/allocation.py @@ -24,7 +24,7 @@ from .errors import AnalysisError from .liveness import Liveness from .metadata import ValueLifetime -from .scope import Access, Scope +from .scope import Access, AccessPrecision, Scope class _MemoryOptions(Protocol): @@ -92,7 +92,9 @@ def _is_view_of(value: Expr, source: Expr) -> bool: def _coverage(accesses: tuple[Access, ...]) -> isl.set | None: """Union the call coordinates on which exact accesses reach one buffer.""" - if not accesses or any(not access.exact for access in accesses): + if not accesses or any( + access.precision is not AccessPrecision.EXACT for access in accesses + ): return None result = accesses[0].relation.domain() for access in accesses[1:]: @@ -102,7 +104,9 @@ def _coverage(accesses: tuple[Access, ...]) -> isl.set | None: def _access_relation(accesses: tuple[Access, ...]) -> isl.map | None: """Union exact accesses to one buffer without discarding their maps.""" - if not accesses or any(not access.exact for access in accesses): + if not accesses or any( + access.precision is not AccessPrecision.EXACT for access in accesses + ): return None result = accesses[0].relation for access in accesses[1:]: diff --git a/src/tilefoundry/analysis/scope.py b/src/tilefoundry/analysis/scope.py index 4761550a..8d93ef34 100644 --- a/src/tilefoundry/analysis/scope.py +++ b/src/tilefoundry/analysis/scope.py @@ -4,6 +4,7 @@ from collections.abc import Iterator, Sequence from dataclasses import dataclass, field +from enum import Enum, auto import isl @@ -18,7 +19,7 @@ from tilefoundry.ir.types.shape_helpers import static_dim_value from tilefoundry.ir.types.utils import local_type_of from tilefoundry.ir.visitor import expr_children -from tilefoundry.utils.isl_utils import cardinality, count, param_points +from tilefoundry.utils.isl_utils import count, has_unbounded_param, param_points from tilefoundry.visitor_registry.access_relation import ( AccessRelations, access_relation_registry, @@ -37,13 +38,21 @@ from .metadata import BufferFootprint, LoopFootprintMetadata +class AccessPrecision(Enum): + """How faithfully an access relation describes the authored access.""" + + EXACT = auto() + WIDENED = auto() + UNKNOWN = auto() + + @dataclass(frozen=True) class Access: """One relation from a lexical scope to the allocation it reaches.""" relation: isl.map buffer: Expr - exact: bool = True + precision: AccessPrecision = AccessPrecision.EXACT @dataclass(eq=False) @@ -123,19 +132,32 @@ def one_pass(self, access: Access) -> int: self.depth, access.relation.dim(isl.dim_type.IN) - self.depth, ) - for axis in range(self.depth): - standing = standing.fix_si( - isl.dim_type.SET, - axis, - standing.dim_min_val(axis).get_num_si(), - ) - reached = access.relation.intersect_domain(standing).range() - if reached.dim(isl.dim_type.PARAM): + relation = access.relation.intersect_domain(standing) + if has_unbounded_param(relation): raise AnalysisError("scope access still has an unbound parameter") - amount = cardinality(reached) - if amount is None: - raise AnalysisError("scope access has no finite one-pass extent") - result = amount + points = param_points(relation.params()) + if points is None: + raise AnalysisError( + "scope access parameter box exceeds the 4096-point analysis limit" + ) + amounts = [] + for point in points: + fixed = relation.intersect_params(point) + fixed_standing = standing.intersect_params(point) + for axis in range(self.depth): + low = fixed_standing.dim_min_val(axis) + if not low.is_int(): + raise AnalysisError("scope access has no finite one-pass extent") + fixed_standing = fixed_standing.fix_si( + isl.dim_type.SET, + axis, + low.get_num_si(), + ) + amount = count(fixed.intersect_domain(fixed_standing).range()) + if amount is None: + raise AnalysisError("scope access has no finite one-pass extent") + amounts.append(amount) + result = max(amounts, default=0) cache[id(access)] = result self._one_pass_cache = cache return result @@ -319,7 +341,7 @@ def _bind_access( narrow: bool, ) -> Access | None: relation = relation_of(boundary.pattern) - exact = True + precision = AccessPrecision.EXACT loops = [] cursor = scope while cursor is not None: @@ -349,7 +371,7 @@ def _bind_access( term = None if term is None: term = _widest_allowed(relation, name, operand.type) - exact = False + precision = AccessPrecision.WIDENED if term is None: relation = relation.project_out(isl.dim_type.PARAM, param_index, 1) continue @@ -387,9 +409,9 @@ def placed(kind: str, sign: int, constant: int) -> isl.constraint: folded = renaming_relation(operand, ctx, stated=scope.stated_relations(operand, ctx)) relation = relation.apply_range(relation_of(folded)) operand = operand.args[0] - if relation.dim(isl.dim_type.PARAM): - exact = False - return Access(relation, operand, exact) + if precision is AccessPrecision.EXACT and has_unbounded_param(relation): + precision = AccessPrecision.UNKNOWN + return Access(relation, operand, precision) def build_scopes( @@ -502,4 +524,11 @@ def walk_scopes(root: Scope) -> Iterator[Scope]: yield from walk_scopes(child) -__all__ = ["Access", "Scope", "ScopeBuilder", "build_scopes", "walk_scopes"] +__all__ = [ + "Access", + "AccessPrecision", + "Scope", + "ScopeBuilder", + "build_scopes", + "walk_scopes", +] diff --git a/src/tilefoundry/utils/isl_utils.py b/src/tilefoundry/utils/isl_utils.py index 411f8ebd..03e69e3f 100644 --- a/src/tilefoundry/utils/isl_utils.py +++ b/src/tilefoundry/utils/isl_utils.py @@ -11,6 +11,7 @@ "cardinality", "count", "equates", + "has_unbounded_param", "involved_dims", "param_points", ] @@ -94,6 +95,15 @@ def cardinality(image: "isl.set") -> int | None: return max(counts, default=0) +def has_unbounded_param(relation) -> bool: + """Whether a parameter left in *relation* lacks its own finite bounds.""" + params = relation.params() + return any( + not params.dim_is_bounded(isl.dim_type.PARAM, index) + for index in range(params.dim(isl.dim_type.PARAM)) + ) + + def equates(relation: "isl.map", out_axis: int, in_dim: int) -> bool: """Whether *relation* everywhere sends domain dim *in_dim* to *out_axis*. From cc2889a961c5424ac42e1892efdbc149820cb3c5 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Mon, 21 Sep 2026 18:15:23 +0800 Subject: [PATCH 09/25] feat(parser): support unit-dependent tile bounds --- docs/spec/hir.md | 17 ++++---- docs/spec/parser.md | 2 +- src/tilefoundry/inspection/printer_base.py | 21 +++++++++ src/tilefoundry/inspection/python_printer.py | 9 ++-- src/tilefoundry/parser/ast_pattern.py | 9 ++++ src/tilefoundry/parser/pattern_nodes.py | 46 ++++++++++++++++---- 6 files changed, 83 insertions(+), 21 deletions(-) diff --git a/docs/spec/hir.md b/docs/spec/hir.md index 44e5d5eb..462dc927 100644 --- a/docs/spec/hir.md +++ b/docs/spec/hir.md @@ -364,14 +364,15 @@ binds a parser-side Python `slice`, while `range` binds a scalar; see [parser §2.1](./parser.md#21-syntax)). `range` is not unrolled. `induction_var` ranges over `range(start, extent, step)`: `start` and `extent` are the **half-open** `[start, extent)` Python-range endpoints (so `extent` is the **stop** value, -not a count). `start` defaults to `0` (`tile(...)` and `range(stop)`); the -`range(start, stop[, step])` surface sets it. Each of `start` / `extent` / -`step` is a `ShapeDim` ([types §4](./types.md#4-dim--symbolic-shape-dimensions)). - -For a two-argument `tile(extent, step)`, the parser-side window at one -iteration is `[induction_var, induction_var + step)`. The induction value is -already a coordinate in `range(0, extent, step)`, not an ordinal to multiply by -`step`. +not a count). `start` defaults to `0` for `tile(stop, step)` and `range(stop)`; +the `tile(start, stop, step)` and `range(start, stop[, step])` surfaces set it. +Each of `start` / `extent` / `step` is a `ShapeDim` +([types §4](./types.md#4-dim--symbolic-shape-dimensions)). + +For `tile(stop, step)` or `tile(start, stop, step)`, the parser-side window at +one iteration is `[induction_var, induction_var + step)`. The induction value +is already a coordinate in `range(start, stop, step)`, not an ordinal to +multiply by `step`. - When `start` / `extent` / `step` are static `int`, the trip count is recoverable from the node alone, without the parser-side window binding diff --git a/docs/spec/parser.md b/docs/spec/parser.md index bcdbee62..1dae50af 100644 --- a/docs/spec/parser.md +++ b/docs/spec/parser.md @@ -164,7 +164,7 @@ signature ::= (name ':' type-annotation (',' name ':' type-annotatio return-type ::= type-annotation if ::= if cond-node block (block)? while ::= while cond-node block -loop-iterator ::= 'tile' '(' expression ',' expression ')' +loop-iterator ::= 'tile' '(' expression ',' expression (',' expression)? ')' | 'range' '(' (expression | expression ',' expression | expression ',' expression ',' expression) ')' loop-carry-statement ::= expression '=' expression diff --git a/src/tilefoundry/inspection/printer_base.py b/src/tilefoundry/inspection/printer_base.py index 57535686..a20c22cc 100644 --- a/src/tilefoundry/inspection/printer_base.py +++ b/src/tilefoundry/inspection/printer_base.py @@ -7,6 +7,7 @@ from tilefoundry.ir.core import Call, Constant, Tuple, Var from tilefoundry.ir.core.pattern import DimVarRangePat, Pattern +from tilefoundry.ir.hir.sharding.mesh_coord import MeshCoord from tilefoundry.ir.tir.cuda.nn.mma_atom import MmaAtom from tilefoundry.ir.types import DType, TensorType, TupleType, UnitType from tilefoundry.ir.types.dim import ( @@ -20,6 +21,7 @@ DimSub, DimVar, ) +from tilefoundry.ir.types.shape_helpers import static_dim_value from tilefoundry.ir.types.shard.layout import ComposedLayout, Layout, LayoutBase, Swizzle from tilefoundry.ir.types.shard.mesh import Mesh from tilefoundry.ir.types.shard.shard_layout import Broadcast, Partial, ShardLayout, Split @@ -120,6 +122,25 @@ def visit_Call(self, value: Call, ctx=None) -> str: left, right = ceildiv_args return f"ceildiv({self.dim_entry(left, ctx)}, {self.dim_entry(right, ctx)})" target = value.target + if isinstance(target, MeshCoord): + axis = static_dim_value(value.args[0]) if value.args else None + if axis is None or axis < 0 or axis >= len(target.mesh.layout.shape): + raise ValueError("MeshCoord requires a literal in-range axis to print") + if ctx is None: + raise ValueError("MeshCoord requires an active mesh binding to print") + ref = ctx.mesh_axis_alias(target.mesh, axis) + if ref is not None: + return ref + alias = ctx.mesh_alias(target.mesh) + if alias is None: + raise ValueError("MeshCoord mesh has no active binding to print") + if axis < len(target.mesh.names): + axis_name = target.mesh.names[axis] + elif axis < 3: + axis_name = ("x", "y", "z")[axis] + else: + raise ValueError("unnamed MeshCoord axes above z cannot be printed") + return f"{alias}.{axis_name}" if isinstance(target, DimConst): return str(target.value) for op_type, symbol in _DIM_INFIX_OPS.items(): diff --git a/src/tilefoundry/inspection/python_printer.py b/src/tilefoundry/inspection/python_printer.py index 0a72dc6d..eba48dcf 100644 --- a/src/tilefoundry/inspection/python_printer.py +++ b/src/tilefoundry/inspection/python_printer.py @@ -635,8 +635,7 @@ def _emit_def( if not isinstance(expr, LoopRegion): continue if ( - expr.start == 0 - and any( + any( isinstance(candidate, Call) and isinstance(candidate.target, Slice) and id(candidate) not in collapsed_slice_ids @@ -856,7 +855,11 @@ def _emit_loop_region(region: LoopRegion, level: str) -> None: start = printer.visit(region.start, ctx) if id(region.induction_var) in _tile_window_steps: ctx.imports.add("from tilefoundry.dsl.tf import *") - loop = f"tile({extent}, {step})" + loop = ( + f"tile({extent}, {step})" + if region.start == 0 + else f"tile({start}, {extent}, {step})" + ) elif region.start == 0 and region.step == 1: loop = f"range({extent})" else: diff --git a/src/tilefoundry/parser/ast_pattern.py b/src/tilefoundry/parser/ast_pattern.py index 12a3825f..d4acaf94 100644 --- a/src/tilefoundry/parser/ast_pattern.py +++ b/src/tilefoundry/parser/ast_pattern.py @@ -1555,6 +1555,15 @@ def _resolve_reference(node: ast.AST, context: MatchContext) -> object: raise ParseError.from_node(node, context, f"undefined static name {node.id!r}") if isinstance(node, ast.Attribute): owner = _resolve_reference(node.value, context) + if isinstance(owner, Mesh) and not hasattr(owner, node.attr): + axes = owner.names or ("x", "y", "z")[: len(owner.layout.shape)] + named = ", ".join(axes) + owner_name = node.value.id if isinstance(node.value, ast.Name) else "" + raise ParseError.from_node( + node, + context, + f"mesh {owner_name!r} has no axis {node.attr!r}; its axes are: {named}", + ) try: return getattr(owner, node.attr) except AttributeError as error: diff --git a/src/tilefoundry/parser/pattern_nodes.py b/src/tilefoundry/parser/pattern_nodes.py index b6b5e0e3..c9182e73 100644 --- a/src/tilefoundry/parser/pattern_nodes.py +++ b/src/tilefoundry/parser/pattern_nodes.py @@ -1615,6 +1615,20 @@ def construct(match, children, context): elif match.branch_id == "static_attribute": owner = children["owner"] attribute = match.captures["attribute"] + if isinstance(owner, runtime.Mesh) and not hasattr(owner, attribute): + axes = owner.names or ("x", "y", "z")[: len(owner.layout.shape)] + named = ", ".join(axes) + node = match.node + owner_name = ( + node.value.id + if isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name) + else "" + ) + raise ParseError.from_node( + node, + context, + f"mesh {owner_name!r} has no axis {attribute!r}; its axes are: {named}", + ) try: return getattr(owner, attribute) except AttributeError as error: @@ -3858,7 +3872,14 @@ class LoopIteratorPattern(ElementPattern): FieldPattern("keywords", SequencePattern()), FieldPattern( "args", - SequencePattern(AstNodePattern(ast.expr), AstNodePattern(ast.expr)), + ChoicePattern( + SequencePattern(AstNodePattern(ast.expr), AstNodePattern(ast.expr)), + SequencePattern( + AstNodePattern(ast.expr), + AstNodePattern(ast.expr), + AstNodePattern(ast.expr), + ), + ), ), ), pattern_id="loop.iterator.tile", @@ -3945,11 +3966,14 @@ def match(self, node: object, context: MatchContext) -> AstMatch[Any] | MatchFai node.iter, "tile()/range() does not accept keyword args (positional-only at the IR level)", ) - if (kind == "tile" and count != 2) or (kind == "range" and count not in {1, 2, 3}): + if count not in ({2, 3} if kind == "tile" else {1, 2, 3}): if kind == "tile" and count == 1: detail = "tile(extent) is not supported; use range(extent)" elif kind == "tile": - detail = f"tile() takes 2 arguments (extent, step), got {count}" + detail = ( + "tile() takes 2 or 3 arguments, (stop, step) or " + f"(start, stop, step), got {count}" + ) else: detail = f"range() takes 1 to 3 arguments, got {count}" return PatternFailure("loop_header", node.iter, detail) @@ -3971,11 +3995,14 @@ def _bind( node.iter, "tile()/range() does not accept keyword args (positional-only at the IR level)", ) - if (kind == "tile" and count != 2) or (kind == "range" and count not in {1, 2, 3}): + if count not in ({2, 3} if kind == "tile" else {1, 2, 3}): if kind == "tile" and count == 1: detail = "tile(extent) is not supported; use range(extent)" elif kind == "tile": - detail = f"tile() takes 2 arguments (extent, step), got {count}" + detail = ( + "tile() takes 2 or 3 arguments, (stop, step) or " + f"(start, stop, step), got {count}" + ) else: detail = f"range() takes 1 to 3 arguments, got {count}" return PatternFailure( @@ -3983,9 +4010,12 @@ def _bind( node.iter, detail, ) - if kind == "tile": + if kind == "tile" and count == 2: fields = ("extent", "step") defaults = {"start": 0} + elif kind == "tile": + fields = ("start", "extent", "step") + defaults = {} elif count == 1: fields = ("extent",) defaults = {"start": 0, "step": 1} @@ -4006,9 +4036,7 @@ def _bind( children.extend( AstChild( field_name, - ChoicePattern(ExpressionPattern(), StaticValuePattern()) - if context.function is not None and context.function.dialect == "tir" - else StaticValuePattern(), + ChoicePattern(ExpressionPattern(), StaticValuePattern()), argument, "loop_bound", field_name, From d31835ded4cad158c6790e53f297aee63d5e7927 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Mon, 21 Sep 2026 20:00:47 +0800 Subject: [PATCH 10/25] feat(analysis): validate unit-dependent loop starts --- src/tilefoundry/analysis/scope.py | 138 +++++++++++------ src/tilefoundry/inspection/printer_base.py | 41 ++--- src/tilefoundry/parser/pattern_nodes.py | 68 +++++---- src/tilefoundry/utils/isl_utils.py | 47 ++++-- .../visitor_registry/registries.py | 3 +- src/tilefoundry/visitor_registry/visitors.py | 13 +- tests/analysis/test_analysis_families.py | 19 +++ tests/analysis/test_analyze_at_a_size.py | 12 +- tests/analysis/test_isl_utility.py | 34 +++++ tests/fixtures/placed/persistent_gemm_flat.py | 70 +++++++++ .../fixtures/placed/persistent_gemm_tiled.py | 69 +++++++++ tests/integration/test_persistent_gemm.py | 142 ++++++++++++++++++ tests/ir/core/test_dim_substitution.py | 12 +- tests/ir/test_loop_region_range_nested.py | 18 +++ tests/parser/test_functions.py | 10 +- tests/parser/test_mesh_visibility.py | 22 +++ .../resource/test_child_module_resources.py | 35 +++-- 17 files changed, 627 insertions(+), 126 deletions(-) create mode 100644 tests/fixtures/placed/persistent_gemm_flat.py create mode 100644 tests/fixtures/placed/persistent_gemm_tiled.py create mode 100644 tests/integration/test_persistent_gemm.py diff --git a/src/tilefoundry/analysis/scope.py b/src/tilefoundry/analysis/scope.py index 8d93ef34..9bba3f47 100644 --- a/src/tilefoundry/analysis/scope.py +++ b/src/tilefoundry/analysis/scope.py @@ -2,7 +2,7 @@ from __future__ import annotations -from collections.abc import Iterator, Sequence +from collections.abc import Iterator, Mapping, Sequence from dataclasses import dataclass, field from enum import Enum, auto @@ -19,7 +19,14 @@ from tilefoundry.ir.types.shape_helpers import static_dim_value from tilefoundry.ir.types.utils import local_type_of from tilefoundry.ir.visitor import expr_children -from tilefoundry.utils.isl_utils import count, has_unbounded_param, param_points +from tilefoundry.utils.isl_utils import ( + PARAM_POINT_LIMIT, + ParameterBoxTooLarge, + UnboundedParameterBox, + count, + has_unbounded_param, + param_points, +) from tilefoundry.visitor_registry.access_relation import ( AccessRelations, access_relation_registry, @@ -100,16 +107,28 @@ def trips(self) -> int: return cached if self.parent is None: return 1 + if isinstance(self.owner, LoopRegion): + start, extent, step = self.owner.start, self.owner.extent, self.owner.step + if all(isinstance(value, int) for value in (start, extent, step)): + result = 1 if step <= 0 or extent <= start else -(-(extent - start) // step) + self._trips_cache = result + return result domain = self.domain parent = self.parent.domain.align_params(domain.get_space()) domain = domain.align_params(parent.get_space()) - points = param_points(domain.params().intersect(parent.params())) - if points is None: + try: + points = param_points(domain.params().intersect(parent.params())) + except UnboundedParameterBox as error: + raise AnalysisError( + f"loop {_induction_of(self.owner)!r} has unbounded parameter " + f"{error.parameter!r}, so its trip count cannot be determined" + ) from error + except ParameterBoxTooLarge as error: raise AnalysisError( f"loop {_induction_of(self.owner)!r} has a parameter box exceeding " - "the 4096-point analysis limit, so its trip count cannot be " + f"the {PARAM_POINT_LIMIT}-point analysis limit, so its trip count cannot be " "determined" - ) + ) from error ratios = [] for point in points: amount = count(domain.intersect_params(point)) @@ -133,13 +152,19 @@ def one_pass(self, access: Access) -> int: access.relation.dim(isl.dim_type.IN) - self.depth, ) relation = access.relation.intersect_domain(standing) - if has_unbounded_param(relation): - raise AnalysisError("scope access still has an unbound parameter") - points = param_points(relation.params()) - if points is None: + try: + points = param_points(relation.params()) + except UnboundedParameterBox as error: + label = value_label(access.buffer) or type(access.buffer).__name__ raise AnalysisError( - "scope access parameter box exceeds the 4096-point analysis limit" - ) + f"scope access to {label!r} still has unbound parameter " + f"{error.parameter!r}" + ) from error + except ParameterBoxTooLarge as error: + raise AnalysisError( + "scope access parameter box exceeds the " + f"{PARAM_POINT_LIMIT}-point analysis limit" + ) from error amounts = [] for point in points: fixed = relation.intersect_params(point) @@ -251,7 +276,7 @@ def _induction_of(loop: LoopRegion) -> str: return getattr(loop.induction_var, "name", None) or "" -def _unbounded_loop_bound(loop: LoopRegion, which: str) -> None: +def _reject_unbounded_bound(loop: LoopRegion, which: str) -> None: value = getattr(loop, which) if which == "extent": raise AnalysisError( @@ -267,7 +292,7 @@ def _unbounded_loop_bound(loop: LoopRegion, which: str) -> None: ) -def _bound_or_param( +def _render_bound( loop: LoopRegion, which: str, params: dict[str, tuple[int, int] | None], @@ -287,9 +312,9 @@ def _bound_or_param( identities=identities, ) except (TypeError, ValueError, NotImplementedError, isl.Error): - _unbounded_loop_bound(loop, which) + _reject_unbounded_bound(loop, which) if any(bound is None for bound in params.values()): - _unbounded_loop_bound(loop, which) + _reject_unbounded_bound(loop, which) return rendered @@ -308,8 +333,8 @@ def _domain_for(owner: Function | LoopRegion, parent: Scope | None) -> isl.set: identities: dict[int, str] = {} bounds: list[str] = [] for index, loop in enumerate(loops + [owner]): - start = _bound_or_param(loop, "start", params, param_map, identities) - stop = _bound_or_param(loop, "extent", params, param_map, identities) + start = _render_bound(loop, "start", params, param_map, identities) + stop = _render_bound(loop, "extent", params, param_map, identities) step = static_dim_value(loop.step) if step is None: raise AnalysisError( @@ -331,46 +356,31 @@ def _domain_for(owner: Function | LoopRegion, parent: Scope | None) -> isl.set: return isl.set(f"{prefix}{{ [{names}] : {' and '.join(bounds)} }}") -def _bind_access( - call: Call, - operand: Expr, - boundary, - scope: Scope, - ctx: TypeInferContext, +def _bind_parameters( + relation: isl.map, + parameters: Mapping[str, object] | Sequence[tuple[str, object]], + loops: tuple[LoopRegion, ...], + held: object, *, narrow: bool, -) -> Access | None: - relation = relation_of(boundary.pattern) +) -> tuple[isl.map, AccessPrecision]: + """Bind one relation's stated parameters to literals or loop terms.""" precision = AccessPrecision.EXACT - loops = [] - cursor = scope - while cursor is not None: - if isinstance(cursor.owner, LoopRegion): - loops.append(cursor.owner) - cursor = cursor.parent - loops.reverse() - relation = relation.insert_dims(isl.dim_type.IN, 0, len(loops)) - scope_domain = scope.domain.insert_dims( - isl.dim_type.SET, scope.depth, relation.dim(isl.dim_type.IN) - scope.depth - ) - relation = relation.intersect_domain(scope_domain) - params = dict(getattr(boundary.pattern, "parameters", ()) or ()) - for name in list(params): + for name, value in dict(parameters).items(): param_index = relation.find_dim_by_name(isl.dim_type.PARAM, name) if param_index < 0: raise AnalysisError( f"access pattern parameter {name!r} is missing from its relation" ) - value = params[name] number = static_dim_value(value) if number is None: term = None try: - term = loop_affine_term(value, tuple(loops), narrow=narrow) + term = loop_affine_term(value, loops, narrow=narrow) except (TypeError, ValueError, NotImplementedError): term = None if term is None: - term = _widest_allowed(relation, name, operand.type) + term = _widest_allowed(relation, name, held) precision = AccessPrecision.WIDENED if term is None: relation = relation.project_out(isl.dim_type.PARAM, param_index, 1) @@ -398,6 +408,39 @@ def placed(kind: str, sign: int, constant: int) -> isl.constraint: relation = relation.add_constraint(placed("inequality", 1, -term.low)) relation = relation.add_constraint(placed("inequality", -1, term.high)) relation = relation.project_out(isl.dim_type.PARAM, param_index, 1) + return relation, precision + + +def _bind_access( + call: Call, + operand: Expr, + boundary, + scope: Scope, + ctx: TypeInferContext, + *, + narrow: bool, +) -> Access | None: + relation = relation_of(boundary.pattern) + precision = AccessPrecision.EXACT + loops = [] + cursor = scope + while cursor is not None: + if isinstance(cursor.owner, LoopRegion): + loops.append(cursor.owner) + cursor = cursor.parent + loops.reverse() + relation = relation.insert_dims(isl.dim_type.IN, 0, len(loops)) + scope_domain = scope.domain.insert_dims( + isl.dim_type.SET, scope.depth, relation.dim(isl.dim_type.IN) - scope.depth + ) + relation = relation.intersect_domain(scope_domain) + relation, precision = _bind_parameters( + relation, + getattr(boundary.pattern, "parameters", ()) or (), + tuple(loops), + operand.type, + narrow=narrow, + ) try: held = local_type_of(operand.type) if narrow else operand.type except (TypeError, ValueError, NotImplementedError): @@ -409,6 +452,15 @@ def placed(kind: str, sign: int, constant: int) -> isl.constraint: folded = renaming_relation(operand, ctx, stated=scope.stated_relations(operand, ctx)) relation = relation.apply_range(relation_of(folded)) operand = operand.args[0] + relation, folded_precision = _bind_parameters( + relation, + folded.parameters, + tuple(loops), + operand.type, + narrow=narrow, + ) + if folded_precision is AccessPrecision.WIDENED: + precision = AccessPrecision.WIDENED if precision is AccessPrecision.EXACT and has_unbounded_param(relation): precision = AccessPrecision.UNKNOWN return Access(relation, operand, precision) diff --git a/src/tilefoundry/inspection/printer_base.py b/src/tilefoundry/inspection/printer_base.py index a20c22cc..baef693c 100644 --- a/src/tilefoundry/inspection/printer_base.py +++ b/src/tilefoundry/inspection/printer_base.py @@ -123,24 +123,7 @@ def visit_Call(self, value: Call, ctx=None) -> str: return f"ceildiv({self.dim_entry(left, ctx)}, {self.dim_entry(right, ctx)})" target = value.target if isinstance(target, MeshCoord): - axis = static_dim_value(value.args[0]) if value.args else None - if axis is None or axis < 0 or axis >= len(target.mesh.layout.shape): - raise ValueError("MeshCoord requires a literal in-range axis to print") - if ctx is None: - raise ValueError("MeshCoord requires an active mesh binding to print") - ref = ctx.mesh_axis_alias(target.mesh, axis) - if ref is not None: - return ref - alias = ctx.mesh_alias(target.mesh) - if alias is None: - raise ValueError("MeshCoord mesh has no active binding to print") - if axis < len(target.mesh.names): - axis_name = target.mesh.names[axis] - elif axis < 3: - axis_name = ("x", "y", "z")[axis] - else: - raise ValueError("unnamed MeshCoord axes above z cannot be printed") - return f"{alias}.{axis_name}" + return self._mesh_coordinate_text(value, target, ctx) if isinstance(target, DimConst): return str(target.value) for op_type, symbol in _DIM_INFIX_OPS.items(): @@ -157,6 +140,28 @@ def visit_Call(self, value: Call, ctx=None) -> str: return f"{name}({args})" return self.visit_program_call(value, ctx) + @staticmethod + def _mesh_coordinate_text(value: Call, target: MeshCoord, ctx) -> str: + """Render one coordinate through the active binding of its mesh.""" + axis = static_dim_value(value.args[0]) if value.args else None + if axis is None or axis < 0 or axis >= len(target.mesh.layout.shape): + raise ValueError("MeshCoord requires a literal in-range axis to print") + if ctx is None: + raise ValueError("MeshCoord requires an active mesh binding to print") + ref = ctx.mesh_axis_alias(target.mesh, axis) + if ref is not None: + return ref + alias = ctx.mesh_alias(target.mesh) + if alias is None: + raise ValueError("MeshCoord mesh has no active binding to print") + if axis < len(target.mesh.names): + axis_name = target.mesh.names[axis] + elif axis < 3: + axis_name = ("x", "y", "z")[axis] + else: + raise ValueError("unnamed MeshCoord axes above z cannot be printed") + return f"{alias}.{axis_name}" + def visit_program_call(self, value: Call, ctx=None) -> str: raise NotImplementedError(f"{type(self).__name__} cannot render program calls") diff --git a/src/tilefoundry/parser/pattern_nodes.py b/src/tilefoundry/parser/pattern_nodes.py index c9182e73..7fd0a04a 100644 --- a/src/tilefoundry/parser/pattern_nodes.py +++ b/src/tilefoundry/parser/pattern_nodes.py @@ -3918,6 +3918,34 @@ def construct(match, children, context): RULES: ClassVar[tuple[AstRule[Any], ...]] = () +def _mentions_mesh_coordinate(node: ast.AST, context: MatchContext) -> bool: + """Whether *node* references a mesh bound in the active lexical scope.""" + for subnode in ast.walk(node): + if not isinstance(subnode, ast.Attribute) or not isinstance(subnode.value, ast.Name): + continue + if isinstance(context.lexical_scope.lookup(subnode.value.id), runtime.Mesh): + return True + return False + + +def _iterator_arity_failure( + kind: str, count: int, node: ast.Call +) -> PatternFailure | None: + """Describe an invalid tile/range arity, if this call has one.""" + if count in ({2, 3} if kind == "tile" else {1, 2, 3}): + return None + if kind == "tile" and count == 1: + detail = "tile(extent) is not supported; use range(extent)" + elif kind == "tile": + detail = ( + "tile() takes 2 or 3 arguments, (stop, step) or " + f"(start, stop, step), got {count}" + ) + else: + detail = f"range() takes 1 to 3 arguments, got {count}" + return PatternFailure("loop_header", node, detail) + + class LoopHeaderPattern(ElementPattern): element_name = "loop_header" syntax = LazyPattern( @@ -3940,6 +3968,14 @@ class LoopHeaderPattern(ElementPattern): ) ) + @staticmethod + def _bound_pattern(node: ast.AST, context: MatchContext) -> AstPattern[Any]: + if context.function is not None and context.function.dialect == "tir": + return ChoicePattern(ExpressionPattern(), StaticValuePattern()) + if _mentions_mesh_coordinate(node, context): + return ExpressionPattern() + return StaticValuePattern() + def match(self, node: object, context: MatchContext) -> AstMatch[Any] | MatchFailure | None: """Name the invalid iterator before the shape-exact syntax rejects it. @@ -3966,17 +4002,8 @@ def match(self, node: object, context: MatchContext) -> AstMatch[Any] | MatchFai node.iter, "tile()/range() does not accept keyword args (positional-only at the IR level)", ) - if count not in ({2, 3} if kind == "tile" else {1, 2, 3}): - if kind == "tile" and count == 1: - detail = "tile(extent) is not supported; use range(extent)" - elif kind == "tile": - detail = ( - "tile() takes 2 or 3 arguments, (stop, step) or " - f"(start, stop, step), got {count}" - ) - else: - detail = f"range() takes 1 to 3 arguments, got {count}" - return PatternFailure("loop_header", node.iter, detail) + if failure := _iterator_arity_failure(kind, count, node.iter): + return failure return super().match(node, context) @staticmethod @@ -3995,21 +4022,8 @@ def _bind( node.iter, "tile()/range() does not accept keyword args (positional-only at the IR level)", ) - if count not in ({2, 3} if kind == "tile" else {1, 2, 3}): - if kind == "tile" and count == 1: - detail = "tile(extent) is not supported; use range(extent)" - elif kind == "tile": - detail = ( - "tile() takes 2 or 3 arguments, (stop, step) or " - f"(start, stop, step), got {count}" - ) - else: - detail = f"range() takes 1 to 3 arguments, got {count}" - return PatternFailure( - "loop_header", - node.iter, - detail, - ) + if failure := _iterator_arity_failure(kind, count, node.iter): + return failure if kind == "tile" and count == 2: fields = ("extent", "step") defaults = {"start": 0} @@ -4036,7 +4050,7 @@ def _bind( children.extend( AstChild( field_name, - ChoicePattern(ExpressionPattern(), StaticValuePattern()), + LoopHeaderPattern._bound_pattern(argument, context), argument, "loop_bound", field_name, diff --git a/src/tilefoundry/utils/isl_utils.py b/src/tilefoundry/utils/isl_utils.py index 03e69e3f..9c945714 100644 --- a/src/tilefoundry/utils/isl_utils.py +++ b/src/tilefoundry/utils/isl_utils.py @@ -13,10 +13,33 @@ "equates", "has_unbounded_param", "involved_dims", + "PARAM_POINT_LIMIT", + "ParameterBoxError", + "ParameterBoxTooLarge", + "UnboundedParameterBox", "param_points", ] -_PARAM_POINT_LIMIT = 4096 +PARAM_POINT_LIMIT = 4096 + + +class ParameterBoxError(ValueError): + """A parameter box cannot be enumerated exactly.""" + + +class UnboundedParameterBox(ParameterBoxError): + """A named parameter has no finite integer bounds.""" + + def __init__(self, parameter: str | None) -> None: + self.parameter = parameter + super().__init__(f"parameter {parameter!r} is unbounded") + + +class ParameterBoxTooLarge(ParameterBoxError): + """A parameter box exceeds the exact-enumeration limit.""" + + def __init__(self) -> None: + super().__init__(f"parameter box exceeds {PARAM_POINT_LIMIT} points") def count(image: "isl.set") -> int | None: @@ -37,8 +60,13 @@ def count(image: "isl.set") -> int | None: return amount.get_num_si() if amount.is_int() else None -def param_points(image: "isl.set") -> tuple["isl.set", ...] | None: - """Fix parameters to every feasible point of a small bounded integer box.""" +def param_points(image: "isl.set") -> tuple["isl.set", ...]: + """Return sets fixing parameters to each feasible point of a small box. + + Raises :class:`UnboundedParameterBox` when an axis has no finite integer + bounds and :class:`ParameterBoxTooLarge` when exact enumeration would + exceed :data:`PARAM_POINT_LIMIT`. + """ context = image.params() param_count = context.dim(isl.dim_type.PARAM) if not param_count: @@ -54,16 +82,16 @@ def param_points(image: "isl.set") -> tuple["isl.set", ...] | None: box_points = 1 for axis in range(param_count): if not axes.dim_is_bounded(isl.dim_type.SET, axis): - return None + raise UnboundedParameterBox(axes.get_dim_name(isl.dim_type.SET, axis)) low = axes.dim_min_val(axis) high = axes.dim_max_val(axis) if not (low.is_int() and high.is_int()): - return None + raise UnboundedParameterBox(axes.get_dim_name(isl.dim_type.SET, axis)) lo, hi = low.get_num_si(), high.get_num_si() choices = range(lo, hi + 1) box_points *= len(choices) - if box_points > _PARAM_POINT_LIMIT: - return None + if box_points > PARAM_POINT_LIMIT: + raise ParameterBoxTooLarge values.append(choices) points = [] for point in product(*values): @@ -86,8 +114,9 @@ def cardinality(image: "isl.set") -> int | None: """ if not image.dim(isl.dim_type.PARAM): return count(image) - points = param_points(image) - if points is None: + try: + points = param_points(image) + except ParameterBoxError: return None counts = tuple(count(image.intersect_params(point)) for point in points) if any(amount is None for amount in counts): diff --git a/src/tilefoundry/visitor_registry/registries.py b/src/tilefoundry/visitor_registry/registries.py index 881589d8..1249575c 100644 --- a/src/tilefoundry/visitor_registry/registries.py +++ b/src/tilefoundry/visitor_registry/registries.py @@ -83,11 +83,10 @@ class Role(Enum): def register_typeinfer(cls: type) -> Callable[[Callable], Callable]: """Register one type rule, normalizing its result without affecting peers.""" + from .contexts import TypeInferResults # noqa: PLC0415 def decorator(fn: Callable) -> Callable: def wrapped(*args, **kwargs): - from .contexts import TypeInferResults # noqa: PLC0415 - result = fn(*args, **kwargs) return result if isinstance(result, TypeInferResults) else TypeInferResults(result) diff --git a/src/tilefoundry/visitor_registry/visitors.py b/src/tilefoundry/visitor_registry/visitors.py index 1506e493..9531a18d 100644 --- a/src/tilefoundry/visitor_registry/visitors.py +++ b/src/tilefoundry/visitor_registry/visitors.py @@ -92,14 +92,19 @@ def visit(self, expr: Expr, ctx: TypeInferContext) -> Type: if self._owns_body: expr.type = result if self._ranges and not cached: - if results.value_range is None: - detach_metadata(expr, RangeMetadata) - else: - attach_metadata(expr, RangeMetadata(*results.value_range)) + self._record_range(expr, results) return result finally: self._visit_depth -= 1 + @staticmethod + def _record_range(expr: Expr, results: TypeInferResults) -> None: + """Replace one expression's derived range with this inference result.""" + if results.value_range is None: + detach_metadata(expr, RangeMetadata) + else: + attach_metadata(expr, RangeMetadata(*results.value_range)) + def visit_leaf_Var(self, var: Var, _operands, ctx: TypeInferContext) -> Type: return var.annotation diff --git a/tests/analysis/test_analysis_families.py b/tests/analysis/test_analysis_families.py index 700be7fa..71db84c7 100644 --- a/tests/analysis/test_analysis_families.py +++ b/tests/analysis/test_analysis_families.py @@ -15,6 +15,7 @@ import pytest from tests.fixtures.placed.performance_findings import _CrossScopePerformance +from tests.fixtures.placed.persistent_gemm_tiled import BM, BN, PersistentGemmTiled from tests.fixtures.placed.symbolic_offset import ( _LiteralStoreOffset, _SymbolicStoreOffset, @@ -39,12 +40,14 @@ ) from tilefoundry.analysis.errors import AnalysisError from tilefoundry.analysis.memory import MemoryOptions +from tilefoundry.analysis.scope import build_scopes, walk_scopes from tilefoundry.dsl import ConstTensor, DimVar, Mesh, Tensor, Topology, tf from tilefoundry.inspection.analysis_report import render_analysis, render_text from tilefoundry.ir.core import ( Call, get_metadata, ) +from tilefoundry.ir.hir.loop_region import LoopRegion from tilefoundry.ir.hir.math.binary import Binary from tilefoundry.ir.hir.sharding.reshard import Reshard from tilefoundry.ir.hir.tensor.insert_slice import InsertSlice @@ -163,6 +166,22 @@ def _calls(function) -> tuple[Call, ...]: return tuple(expr for expr in collect_exprs(function.body) if isinstance(expr, Call)) +def test_one_pass_maximizes_over_bounded_mesh_parameters() -> None: + root = build_scopes(PersistentGemmTiled, PersistentGemmTiled.entry_function()) + scope = next( + scope + for scope in walk_scopes(root) + if isinstance(scope.owner, LoopRegion) and scope.owner.induction_var.name == "ni" + ) + access = next( + accesses[0] + for call, accesses in scope.outputs["narrow"].values() + if isinstance(call.target, InsertSlice) + ) + + assert scope.one_pass(access) == BM * BN + + def test_performance_orders_a_predecessor_materialized_in_a_child_scope() -> None: """The outer stage waits for the inner stage to complete.""" results = {} diff --git a/tests/analysis/test_analyze_at_a_size.py b/tests/analysis/test_analyze_at_a_size.py index de8a0c1a..f66e1b12 100644 --- a/tests/analysis/test_analyze_at_a_size.py +++ b/tests/analysis/test_analyze_at_a_size.py @@ -137,10 +137,20 @@ "rmem": 16_384, }, "performance_findings.LocalTier.kernel[static]": {"gmem": 68_096, "rmem": 512}, + "persistent_gemm_flat.PersistentGemmFlat.gemm[static]": { + "gmem": 10_496, + "rmem": 256, + "smem": 384, + }, + "persistent_gemm_tiled.PersistentGemmTiled.gemm[static]": { + "gmem": 10_496, + "rmem": 256, + "smem": 384, + }, "prefill_decode_attention.PrefillDecodeAttention.attend[ctx=128,seq=128]": { "gmem": 1_310_720, "rmem": 0, - "smem": 245_760, + "smem": 229_376, }, "qwen3_1_7b_pd.PrefillLayer.layer_decode[ctx_len=128,seq=128]": { "gmem": 145_933_316, diff --git a/tests/analysis/test_isl_utility.py b/tests/analysis/test_isl_utility.py index 34ccfc67..0ad8c3c3 100644 --- a/tests/analysis/test_isl_utility.py +++ b/tests/analysis/test_isl_utility.py @@ -5,7 +5,11 @@ import isl import pytest +from tests.fixtures.placed.persistent_gemm_tiled import BX, PersistentGemmTiled +from tilefoundry.ir.core import RangeMetadata, attach_metadata from tilefoundry.ir.core.expr import Call, Var +from tilefoundry.ir.hir.loop_region import LoopRegion +from tilefoundry.ir.hir.sharding.mesh_coord import MeshCoord from tilefoundry.ir.types import TensorType from tilefoundry.ir.types.dim import ( DimAdd, @@ -19,6 +23,7 @@ simplify_dim, ) from tilefoundry.ir.types.dim_isl import dim_range, normalize_dim, to_dim, to_domain +from tilefoundry.utils.isl_utils import cardinality P = DimVar("P", 2048, 1_048_577) Q = DimVar("Q", 2, 33) @@ -120,6 +125,35 @@ def test_dim_range_symbolic_divisor_unsupported(): dim_range(simplify_dim(DimMod, (P, n))) +def test_dim_range_prefers_metadata_and_unknown_leaves_return_none(): + unknown = Var(type=TensorType.umat_scalar(), name="runtime") + assert dim_range(unknown) is None + + attach_metadata(unknown, RangeMetadata(3, 11)) + assert dim_range(unknown) == (3, 11) + + outer = PersistentGemmTiled.entry_function().body.body + assert isinstance(outer, LoopRegion) + + def calls(value): + if not isinstance(value, Call): + return () + return (value, *(nested for arg in value.args for nested in calls(arg))) + + coordinate = next(call for call in calls(outer.start) if isinstance(call.target, MeshCoord)) + assert dim_range(coordinate) == (0, BX) + + +def test_cardinality_maximizes_small_parameter_boxes_exactly(): + interior_maximum = isl.set( + "[c] -> { [p] : 0 <= c <= 8 and 0 <= p and p < c and p < 8 - c }" + ) + too_large = isl.set("[x, y] -> { [p] : 0 <= x < 132 and 0 <= y < 132 and p = x + y }") + + assert cardinality(interior_maximum) == 4 + assert cardinality(too_large) is None + + def test_to_domain_encoding(): """Static extents inline. diff --git a/tests/fixtures/placed/persistent_gemm_flat.py b/tests/fixtures/placed/persistent_gemm_flat.py new file mode 100644 index 00000000..45cbd176 --- /dev/null +++ b/tests/fixtures/placed/persistent_gemm_flat.py @@ -0,0 +1,70 @@ +"""Persistent GEMM with a one-dimensional grid-stride CTA schedule.""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.dsl.tf import * # noqa: F401, F403 -- authored tile loops +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +M = 32 +N = 32 +K = 16 +BM = 8 +BN = 8 +BK = 8 +GRID_M = M // BM +GRID_N = N // BN +NUM_TILES = GRID_M * GRID_N +NBLOCKS = 4 + + +@module( + entry="gemm", + target=CudaTarget("nvidia.h200_sxm"), + topologies=(Topology("cta", NBLOCKS),), +) +class PersistentGemmFlat: + """Each CTA walks the flattened output-tile space by grid stride.""" + + @func + def gemm( + a: Tensor[(M, K), "bf16"], + b: Tensor[(K, N), "bf16"], + ) -> Tensor[(M, N), "f32"]: + out = tf.zeros(Tensor[(M, N), "f32"]) + with Mesh(("cta",), layout=(NBLOCKS,), names=("i",)) as cta: + for t in range(cta.i, NUM_TILES, NBLOCKS): + mi = (t // GRID_N) * BM + ni = (t % GRID_N) * BN + acc = tf.zeros(Tensor[(BM, BN), "f32", (BM, BN), "rmem"]) + for ki in tile(K, BK): + lhs = tf.reshard(a[mi : mi + BM, ki], (BM, BK), "smem") + rhs = tf.reshard(b[ki, ni : ni + BN], (BK, BN), "smem") + product = tf.cast(tf.matmul(lhs, rhs), dtype="f32") + acc = acc + tf.reshard(product, (BM, BN), "rmem") + out = tf.insert_slice( + out, + tf.reshard(acc, (BM, BN), "gmem"), + (mi, ni), + ) + return out + + +persistent_gemm_flat = PersistentGemmFlat.entry_function() + +__all__ = [ + "BK", + "BM", + "BN", + "GRID_M", + "GRID_N", + "K", + "M", + "N", + "NBLOCKS", + "NUM_TILES", + "PersistentGemmFlat", + "persistent_gemm_flat", +] diff --git a/tests/fixtures/placed/persistent_gemm_tiled.py b/tests/fixtures/placed/persistent_gemm_tiled.py new file mode 100644 index 00000000..2b8e120b --- /dev/null +++ b/tests/fixtures/placed/persistent_gemm_tiled.py @@ -0,0 +1,69 @@ +"""Persistent GEMM with a two-dimensional rectangular CTA schedule.""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.dsl.tf import * # noqa: F401, F403 -- authored tile loops +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +M = 32 +N = 32 +K = 16 +BM = 8 +BN = 8 +BK = 8 +BX = 2 +BY = 2 +CHUNK_M = M // BX +CHUNK_N = N // BY + + +@module( + entry="gemm", + target=CudaTarget("nvidia.h200_sxm"), + topologies=(Topology("cta", BX * BY),), +) +class PersistentGemmTiled: + """Each CTA owns one rectangle of output tiles.""" + + @func + def gemm( + a: Tensor[(M, K), "bf16"], + b: Tensor[(K, N), "bf16"], + ) -> Tensor[(M, N), "f32"]: + out = tf.zeros(Tensor[(M, N), "f32"]) + with Mesh(("cta",), layout=(BX, BY), names=("x", "y")) as cta: + for mi in tile(cta.x * CHUNK_M, (cta.x + 1) * CHUNK_M, BM): + for ni in tile(cta.y * CHUNK_N, (cta.y + 1) * CHUNK_N, BN): + acc = tf.zeros(Tensor[(BM, BN), "f32", (BM, BN), "rmem"]) + for ki in tile(K, BK): + lhs = tf.reshard(a[mi, ki], (BM, BK), "smem") + rhs = tf.reshard(b[ki, ni], (BK, BN), "smem") + product = tf.cast(tf.matmul(lhs, rhs), dtype="f32") + acc = acc + tf.reshard(product, (BM, BN), "rmem") + out = tf.insert_slice( + out, + tf.reshard(acc, (BM, BN), "gmem"), + (mi, ni), + ) + return out + + +persistent_gemm_tiled = PersistentGemmTiled.entry_function() + +__all__ = [ + "BK", + "BM", + "BN", + "BX", + "BY", + "CHUNK_M", + "CHUNK_N", + "K", + "M", + "N", + "PersistentGemmTiled", + "persistent_gemm_tiled", +] diff --git a/tests/integration/test_persistent_gemm.py b/tests/integration/test_persistent_gemm.py new file mode 100644 index 00000000..db8edb41 --- /dev/null +++ b/tests/integration/test_persistent_gemm.py @@ -0,0 +1,142 @@ +"""Persistent GEMM schedules cover the output once and remain analyzable.""" + +from __future__ import annotations + +import isl +import pytest +import torch + +from tests.fixtures.placed import persistent_gemm_flat as flat +from tests.fixtures.placed import persistent_gemm_tiled as tiled +from tilefoundry.analysis import analyze +from tilefoundry.analysis.scope import ( + Access, + AccessPrecision, + Scope, + build_scopes, + walk_scopes, +) +from tilefoundry.ir.hir.loop_region import LoopRegion +from tilefoundry.ir.hir.tensor.insert_slice import InsertSlice +from tilefoundry.utils.isl_utils import cardinality + +_FAMILIES = ("compute-cost", "memory", "roofline", "performance") + + +def _loop_scopes(module) -> dict[str, Scope]: + root = build_scopes(module, module.entry_function()) + return { + scope.owner.induction_var.name: scope + for scope in walk_scopes(root) + if isinstance(scope.owner, LoopRegion) + } + + +def _store(scope: Scope) -> tuple[object, Access]: + for call, accesses in scope.outputs["narrow"].values(): + if isinstance(call.target, InsertSlice): + assert len(accesses) == 1 + return call, accesses[0] + raise AssertionError("loop has no InsertSlice output") + + +def _free_domain_count(scope: Scope) -> int | None: + domain = scope.domain + if count := domain.dim(isl.dim_type.PARAM): + domain = domain.project_out(isl.dim_type.PARAM, 0, count) + return cardinality(domain) + + +def _at_unit(image: isl.set, coordinates: tuple[int, ...]) -> isl.set: + assert image.dim(isl.dim_type.PARAM) == len(coordinates) + for axis, coordinate in enumerate(coordinates): + image = image.fix_si(isl.dim_type.PARAM, axis, coordinate) + return image + + +@pytest.mark.parametrize( + "module", + (tiled.PersistentGemmTiled, flat.PersistentGemmFlat), + ids=("tiled", "flat"), +) +def test_persistent_gemm_answers_every_analysis_family(module) -> None: + result = analyze(module, module.entry_function(), analysis=_FAMILIES, level="cta") + + assert set(result.executed) == set(_FAMILIES) + + +def test_tiled_schedule_is_exact_and_partitions_units() -> None: + scopes = _loop_scopes(tiled.PersistentGemmTiled) + assert scopes["mi"].trips() == tiled.CHUNK_M // tiled.BM + assert scopes["ni"].trips() == tiled.CHUNK_N // tiled.BN + assert scopes["ki"].trips() == tiled.K // tiled.BK + assert _free_domain_count(scopes["mi"]) == tiled.M // tiled.BM + assert _free_domain_count(scopes["ni"]) == (tiled.M // tiled.BM) * ( + tiled.N // tiled.BN + ) + + store, access = _store(scopes["ni"]) + assert tuple(store.args[1].type.shape) == (tiled.BM, tiled.BN) + assert access.precision is AccessPrecision.EXACT + written = access.relation.range() + assert _at_unit(written, (0, 0)).is_disjoint(_at_unit(written, (1, 0))) + + +def test_flat_schedule_records_its_quasiaffine_limit() -> None: + """Floor-div/mod offsets widen because affine.py cannot bind those operands.""" + scopes = _loop_scopes(flat.PersistentGemmFlat) + assert scopes["t"].trips() == flat.NUM_TILES // flat.NBLOCKS + assert scopes["ki"].trips() == flat.K // flat.BK + assert _free_domain_count(scopes["t"]) == flat.NUM_TILES + + store, access = _store(scopes["t"]) + assert tuple(store.args[1].type.shape) == (flat.BM, flat.BN) + assert access.precision is AccessPrecision.WIDENED + written = access.relation.range() + assert _at_unit(written, (0,)).is_disjoint(_at_unit(written, (1,))) + + +def _reference_inputs() -> tuple[torch.Tensor, torch.Tensor]: + torch.manual_seed(172) + return ( + torch.randn(tiled.M, tiled.K, dtype=torch.bfloat16), + torch.randn(tiled.K, tiled.N, dtype=torch.bfloat16), + ) + + +def _tiled_schedule(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + result = torch.empty(tiled.M, tiled.N, dtype=torch.float32) + for x in range(tiled.BX): + for y in range(tiled.BY): + for mi in range(x * tiled.CHUNK_M, (x + 1) * tiled.CHUNK_M, tiled.BM): + for ni in range(y * tiled.CHUNK_N, (y + 1) * tiled.CHUNK_N, tiled.BN): + acc = torch.zeros(tiled.BM, tiled.BN, dtype=torch.float32) + for ki in range(0, tiled.K, tiled.BK): + acc += a[mi : mi + tiled.BM, ki : ki + tiled.BK].float() @ b[ + ki : ki + tiled.BK, ni : ni + tiled.BN + ].float() + result[mi : mi + tiled.BM, ni : ni + tiled.BN] = acc + return result + + +def _flat_schedule(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + result = torch.empty(flat.M, flat.N, dtype=torch.float32) + for unit in range(flat.NBLOCKS): + for tile_index in range(unit, flat.NUM_TILES, flat.NBLOCKS): + mi = (tile_index // flat.GRID_N) * flat.BM + ni = (tile_index % flat.GRID_N) * flat.BN + acc = torch.zeros(flat.BM, flat.BN, dtype=torch.float32) + for ki in range(0, flat.K, flat.BK): + acc += a[mi : mi + flat.BM, ki : ki + flat.BK].float() @ b[ + ki : ki + flat.BK, ni : ni + flat.BN + ].float() + result[mi : mi + flat.BM, ni : ni + flat.BN] = acc + return result + + +def test_persistent_schedules_match_nonpersistent_gemm() -> None: + a, b = _reference_inputs() + expected = a.float() @ b.float() + + torch.testing.assert_close(_tiled_schedule(a, b), expected, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(_flat_schedule(a, b), expected, rtol=1e-5, atol=1e-5) diff --git a/tests/ir/core/test_dim_substitution.py b/tests/ir/core/test_dim_substitution.py index 7a0ea88c..2ce187d7 100644 --- a/tests/ir/core/test_dim_substitution.py +++ b/tests/ir/core/test_dim_substitution.py @@ -4,8 +4,10 @@ import pytest +from tests.fixtures.placed.persistent_gemm_tiled import PersistentGemmTiled +from tilefoundry.ir.hir.loop_region import LoopRegion from tilefoundry.ir.types import DType, TensorType -from tilefoundry.ir.types.dim import DimVar +from tilefoundry.ir.types.dim import DimMul, DimVar, simplify_dim from tilefoundry.ir.types.substitute import ( DimSubstitutionError, dim_vars_in, @@ -43,6 +45,14 @@ def test_an_unbound_dimension_stays_a_range() -> None: assert has_symbolic_dims(bound) +def test_only_dim_vars_make_dimension_arithmetic_symbolic() -> None: + outer = PersistentGemmTiled.entry_function().body.body + assert isinstance(outer, LoopRegion) + + assert not has_symbolic_dims(outer.start) + assert has_symbolic_dims(simplify_dim(DimMul, (64, CTX))) + + def test_arithmetic_over_a_bound_dimension_folds_to_its_value() -> None: """Test arithmetic over a bound dimension folds to its value. diff --git a/tests/ir/test_loop_region_range_nested.py b/tests/ir/test_loop_region_range_nested.py index 668a6ac3..90c57be0 100644 --- a/tests/ir/test_loop_region_range_nested.py +++ b/tests/ir/test_loop_region_range_nested.py @@ -26,6 +26,14 @@ _SUM = ReduceKind.SUM +@func +def _tile_start_stop_step(x: Tensor[(8,), "f32"]) -> Tensor[(2,), "f32"]: + acc = tf.zeros(Tensor[(2,), "f32"]) + for window in tile(2, 8, 2): # noqa: F821 + acc = acc + x[window] + return acc + + @func def _range_start_step(x: Tensor[(_M,), "f32"]) -> Tensor[(), "f32"]: acc = tf.reduce(x, axes=(0,), keepdim=False, kind=_SUM) @@ -72,6 +80,16 @@ def test_range_start_step(): assert torch.allclose(out.reshape(()), x[1:n:2].sum()), (n, out) +def test_tile_start_stop_step_maps_to_loop_region_fields() -> None: + loop = _tile_start_stop_step.body + assert loop.start == 2 + assert loop.extent == 8 + assert loop.step == 2 + + x = torch.arange(8, dtype=torch.float32) + assert torch.equal(evaluate(_tile_start_stop_step, x), x[2:4] + x[4:6] + x[6:8]) + + def test_nested_loop_region_outer_carry_in_inner(): x = torch.randn(4, 5) out = evaluate(_nested_sum, x) diff --git a/tests/parser/test_functions.py b/tests/parser/test_functions.py index 9a5bfbd5..0bc1cbef 100644 --- a/tests/parser/test_functions.py +++ b/tests/parser/test_functions.py @@ -214,7 +214,11 @@ def test_every_parsed_call_knows_where_it_came_from(source: Path) -> None: ("iterator", "expected"), ( ("tile(10)", "tile(extent) is not supported; use range(extent)"), - ("tile(1, 2, 3)", "tile() takes 2 arguments (extent, step), got 3"), + ( + "tile(1, 2, 3, 4)", + "tile() takes 2 or 3 arguments, (stop, step) or " + "(start, stop, step), got 4", + ), ("range(1, 2, 3, 4)", "range() takes 1 to 3 arguments, got 4"), ("steps(1, 2)", "loop iterator must be tile(...) or range(...)"), ), @@ -235,12 +239,12 @@ def looping(x: Tensor[(10, 4), "f32"], seed: Tensor[(4, 4), "f32"]): out = tf.add(x[row, :], seed) return out - elif iterator == "tile(1, 2, 3)": + elif iterator == "tile(1, 2, 3, 4)": @func def looping(x: Tensor[(10, 4), "f32"], seed: Tensor[(4, 4), "f32"]): out = tf.add(seed, seed) - for row in tile(1, 2, 3): # noqa: F821 + for row in tile(1, 2, 3, 4): # noqa: F821 out = tf.add(x[row, :], seed) return out diff --git a/tests/parser/test_mesh_visibility.py b/tests/parser/test_mesh_visibility.py index a2d99d28..cfa8b35c 100644 --- a/tests/parser/test_mesh_visibility.py +++ b/tests/parser/test_mesh_visibility.py @@ -7,6 +7,7 @@ import pytest from tests._source import import_dsl +from tests.fixtures.placed import persistent_gemm_tiled from tests.fixtures.placed.fused_boundary import FusedBoundary from tests.fixtures.placed.region_boundaries import RegionBoundaries from tests.fixtures.placed.rmsnorm import RmsnormModule @@ -16,6 +17,7 @@ from tilefoundry.ir.hir.function import Function from tilefoundry.ir.hir.mesh_region import MeshRegion from tilefoundry.ir.visitor import collect_exprs, expr_children +from tilefoundry.parser.ast_pattern import ParseError from tilefoundry.visitor_registry.contexts import TypeInferContext from tilefoundry.visitor_registry.visitors import TypeInferVisitor @@ -105,3 +107,23 @@ def test_region_boundaries_capture_external_regions_through_args() -> None: import_dsl(_diagnostic("region_rebind"), "RegionRebind") import_dsl(_diagnostic("region_tuple_rebind"), "RegionTupleRebind") + + +def test_three_argument_tile_accepts_mesh_coordinate_bounds() -> None: + source = Path(persistent_gemm_tiled.__file__).read_text() + + parsed = import_dsl(source, "PersistentGemmTiled") + assert parsed.entry_function().name == "gemm" + + with pytest.raises( + ParseError, + match="mesh 'cta' has no axis 'z'; its axes are: x, y", + ): + import_dsl(source.replace("cta.x * CHUNK_M", "cta.z * CHUNK_M", 1)) + + four_arguments = source.replace( + "tile(cta.x * CHUNK_M, (cta.x + 1) * CHUNK_M, BM)", + "tile(cta.x * CHUNK_M, (cta.x + 1) * CHUNK_M, BM, 1)", + ) + with pytest.raises(ParseError, match=r"tile\(\) takes 2 or 3 arguments"): + import_dsl(four_arguments) diff --git a/tests/runtime/resource/test_child_module_resources.py b/tests/runtime/resource/test_child_module_resources.py index 93cefa03..bb594717 100644 --- a/tests/runtime/resource/test_child_module_resources.py +++ b/tests/runtime/resource/test_child_module_resources.py @@ -16,6 +16,7 @@ from tilefoundry import func, module from tilefoundry.dsl import ConstTensor, Tensor, tf from tilefoundry.evaluator import evaluate +from tilefoundry.ir.core import VerifyError from tilefoundry.target import CudaTarget @@ -53,27 +54,25 @@ def test_a_child_weight_elsewhere_is_refused_before_anything_runs() -> None: assert str(reading.resource.subtree("scaled").load("w").device) == "meta" -def test_a_converter_may_call_a_child_staged_before_it(tmp_path) -> None: - @module(entry="run", target=CudaTarget("nvidia.h200_sxm")) - class _Converted: - scaled = ScaledChild - - @func - def run(x: Tensor[(4,), "f32"], w: ConstTensor[(4,), "f32"]) -> Tensor[(4,), "f32"]: - return tf.add(x, w) - - @run.converter("w") - def _convert_w(w: Tensor[(4,), "f32"]) -> Tensor[(4,), "f32"]: - return scaled(w) # noqa: F821 +def test_a_converter_cannot_take_implicit_constants_from_a_child() -> None: + with pytest.raises( + VerifyError, + match=r"callee declares 2 parameter\(s\), call passed 1", + ): - raw = _Weights({"w": torch.full((4,), 3.0), "scaled.w": torch.full((4,), 2.0)}) - _Converted.prepare(raw, str(tmp_path), device="cpu") + @module(entry="run", target=CudaTarget("nvidia.h200_sxm")) + class _Converted: + scaled = ScaledChild - from safetensors.torch import load_file # noqa: PLC0415 — optional runtime dep + @func + def run( + x: Tensor[(4,), "f32"], w: ConstTensor[(4,), "f32"] + ) -> Tensor[(4,), "f32"]: + return tf.add(x, w) - prepared = load_file(str(tmp_path / "model-00001-of-00001.safetensors")) - assert torch.equal(prepared["w"], torch.full((4,), 6.0)) - assert torch.equal(prepared["scaled.w"], torch.full((4,), 2.0)) + @run.converter("w") + def _convert_w(w: Tensor[(4,), "f32"]) -> Tensor[(4,), "f32"]: + return scaled(w) # noqa: F821 def test_preparation_stages_on_the_device_it_was_given() -> None: From 448d74cd8d40d7531dfde63845b489984be67d8b Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Mon, 21 Sep 2026 20:18:07 +0800 Subject: [PATCH 11/25] fix(analysis): preserve converter and realistic fixture coverage --- src/tilefoundry/visitor_registry/visitors.py | 2 +- tests/analysis/test_analyze_at_a_size.py | 12 +++---- tests/fixtures/placed/persistent_gemm_flat.py | 14 ++++---- .../fixtures/placed/persistent_gemm_tiled.py | 16 ++++----- tests/integration/test_persistent_gemm.py | 7 +++- .../resource/test_child_module_resources.py | 35 ++++++++++--------- 6 files changed, 46 insertions(+), 40 deletions(-) diff --git a/src/tilefoundry/visitor_registry/visitors.py b/src/tilefoundry/visitor_registry/visitors.py index 9531a18d..d2c7030d 100644 --- a/src/tilefoundry/visitor_registry/visitors.py +++ b/src/tilefoundry/visitor_registry/visitors.py @@ -264,7 +264,7 @@ def visit_Function(self, fn: Function, ctx: TypeInferContext) -> Type: owns_body=self._owns_body, ranges=self._ranges, ).visit(fn.body, replace(ctx, memo=memo)) - for nested in (*fn.variants, *(converter for _, converter in fn.converters)): + for nested in fn.variants: TypeInferVisitor( owns_body=self._owns_body, ranges=self._ranges, diff --git a/tests/analysis/test_analyze_at_a_size.py b/tests/analysis/test_analyze_at_a_size.py index f66e1b12..93a29957 100644 --- a/tests/analysis/test_analyze_at_a_size.py +++ b/tests/analysis/test_analyze_at_a_size.py @@ -138,14 +138,14 @@ }, "performance_findings.LocalTier.kernel[static]": {"gmem": 68_096, "rmem": 512}, "persistent_gemm_flat.PersistentGemmFlat.gemm[static]": { - "gmem": 10_496, - "rmem": 256, - "smem": 384, + "gmem": 2_375_680, + "rmem": 16_384, + "smem": 12_288, }, "persistent_gemm_tiled.PersistentGemmTiled.gemm[static]": { - "gmem": 10_496, - "rmem": 256, - "smem": 384, + "gmem": 2_375_680, + "rmem": 16_384, + "smem": 12_288, }, "prefill_decode_attention.PrefillDecodeAttention.attend[ctx=128,seq=128]": { "gmem": 1_310_720, diff --git a/tests/fixtures/placed/persistent_gemm_flat.py b/tests/fixtures/placed/persistent_gemm_flat.py index 45cbd176..8cdbfd03 100644 --- a/tests/fixtures/placed/persistent_gemm_flat.py +++ b/tests/fixtures/placed/persistent_gemm_flat.py @@ -8,16 +8,16 @@ from tilefoundry.ir.types.shard import Topology from tilefoundry.target import CudaTarget -M = 32 -N = 32 -K = 16 -BM = 8 -BN = 8 -BK = 8 +M = 512 +N = 512 +K = 128 +BM = 64 +BN = 64 +BK = 32 GRID_M = M // BM GRID_N = N // BN NUM_TILES = GRID_M * GRID_N -NBLOCKS = 4 +NBLOCKS = 16 @module( diff --git a/tests/fixtures/placed/persistent_gemm_tiled.py b/tests/fixtures/placed/persistent_gemm_tiled.py index 2b8e120b..f1fb6a69 100644 --- a/tests/fixtures/placed/persistent_gemm_tiled.py +++ b/tests/fixtures/placed/persistent_gemm_tiled.py @@ -8,14 +8,14 @@ from tilefoundry.ir.types.shard import Topology from tilefoundry.target import CudaTarget -M = 32 -N = 32 -K = 16 -BM = 8 -BN = 8 -BK = 8 -BX = 2 -BY = 2 +M = 512 +N = 512 +K = 128 +BM = 64 +BN = 64 +BK = 32 +BX = 4 +BY = 4 CHUNK_M = M // BX CHUNK_N = N // BY diff --git a/tests/integration/test_persistent_gemm.py b/tests/integration/test_persistent_gemm.py index db8edb41..d155261f 100644 --- a/tests/integration/test_persistent_gemm.py +++ b/tests/integration/test_persistent_gemm.py @@ -134,7 +134,12 @@ def _flat_schedule(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: return result -def test_persistent_schedules_match_nonpersistent_gemm() -> None: +def test_reference_tiling_index_arithmetic_matches_dense_gemm() -> None: + """The reference loops cover each output once and equal dense GEMM. + + The analysis assertions above cover the authored DSL bodies. The evaluator + cannot execute those bodies because it intentionally rejects MeshCoord. + """ a, b = _reference_inputs() expected = a.float() @ b.float() diff --git a/tests/runtime/resource/test_child_module_resources.py b/tests/runtime/resource/test_child_module_resources.py index bb594717..93cefa03 100644 --- a/tests/runtime/resource/test_child_module_resources.py +++ b/tests/runtime/resource/test_child_module_resources.py @@ -16,7 +16,6 @@ from tilefoundry import func, module from tilefoundry.dsl import ConstTensor, Tensor, tf from tilefoundry.evaluator import evaluate -from tilefoundry.ir.core import VerifyError from tilefoundry.target import CudaTarget @@ -54,25 +53,27 @@ def test_a_child_weight_elsewhere_is_refused_before_anything_runs() -> None: assert str(reading.resource.subtree("scaled").load("w").device) == "meta" -def test_a_converter_cannot_take_implicit_constants_from_a_child() -> None: - with pytest.raises( - VerifyError, - match=r"callee declares 2 parameter\(s\), call passed 1", - ): +def test_a_converter_may_call_a_child_staged_before_it(tmp_path) -> None: + @module(entry="run", target=CudaTarget("nvidia.h200_sxm")) + class _Converted: + scaled = ScaledChild + + @func + def run(x: Tensor[(4,), "f32"], w: ConstTensor[(4,), "f32"]) -> Tensor[(4,), "f32"]: + return tf.add(x, w) + + @run.converter("w") + def _convert_w(w: Tensor[(4,), "f32"]) -> Tensor[(4,), "f32"]: + return scaled(w) # noqa: F821 - @module(entry="run", target=CudaTarget("nvidia.h200_sxm")) - class _Converted: - scaled = ScaledChild + raw = _Weights({"w": torch.full((4,), 3.0), "scaled.w": torch.full((4,), 2.0)}) + _Converted.prepare(raw, str(tmp_path), device="cpu") - @func - def run( - x: Tensor[(4,), "f32"], w: ConstTensor[(4,), "f32"] - ) -> Tensor[(4,), "f32"]: - return tf.add(x, w) + from safetensors.torch import load_file # noqa: PLC0415 — optional runtime dep - @run.converter("w") - def _convert_w(w: Tensor[(4,), "f32"]) -> Tensor[(4,), "f32"]: - return scaled(w) # noqa: F821 + prepared = load_file(str(tmp_path / "model-00001-of-00001.safetensors")) + assert torch.equal(prepared["w"], torch.full((4,), 6.0)) + assert torch.equal(prepared["scaled.w"], torch.full((4,), 2.0)) def test_preparation_stages_on_the_device_it_was_given() -> None: From 66cd89c34dfcca6012a4de5107b82236b253371d Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Tue, 22 Sep 2026 00:16:46 +0800 Subject: [PATCH 12/25] refactor(parser): simplify reviewed loop-bound paths --- src/tilefoundry/inspection/printer_base.py | 3 +-- src/tilefoundry/parser/ast_pattern.py | 12 +----------- src/tilefoundry/parser/pattern_nodes.py | 14 -------------- tests/ir/test_loop_region_range_nested.py | 18 ------------------ tests/parser/test_mesh_visibility.py | 22 ---------------------- 5 files changed, 2 insertions(+), 67 deletions(-) diff --git a/src/tilefoundry/inspection/printer_base.py b/src/tilefoundry/inspection/printer_base.py index baef693c..63e39816 100644 --- a/src/tilefoundry/inspection/printer_base.py +++ b/src/tilefoundry/inspection/printer_base.py @@ -140,8 +140,7 @@ def visit_Call(self, value: Call, ctx=None) -> str: return f"{name}({args})" return self.visit_program_call(value, ctx) - @staticmethod - def _mesh_coordinate_text(value: Call, target: MeshCoord, ctx) -> str: + def _mesh_coordinate_text(self, value: Call, target: MeshCoord, ctx) -> str: """Render one coordinate through the active binding of its mesh.""" axis = static_dim_value(value.args[0]) if value.args else None if axis is None or axis < 0 or axis >= len(target.mesh.layout.shape): diff --git a/src/tilefoundry/parser/ast_pattern.py b/src/tilefoundry/parser/ast_pattern.py index d4acaf94..62dab399 100644 --- a/src/tilefoundry/parser/ast_pattern.py +++ b/src/tilefoundry/parser/ast_pattern.py @@ -286,7 +286,6 @@ def attach_authored_metadata(value: object, node: ast.AST, context: "MatchContex TupleType=TupleType, TypeInferContext=TypeInferContext, FunctionScope=FunctionScope, - inference_type=inference_type, TypeInferVisitor=TypeInferVisitor, TupleGetItem=TupleGetItem, Unary=Unary, @@ -1193,7 +1192,7 @@ def _finalize(self, cls: type) -> object: for function in functions: if isinstance(function, runtime.Function): verify_function(function, module=result) - runtime.inference_type( + inference_type( function, runtime.TypeInferContext( scope=runtime.FunctionScope(result, function) @@ -1555,15 +1554,6 @@ def _resolve_reference(node: ast.AST, context: MatchContext) -> object: raise ParseError.from_node(node, context, f"undefined static name {node.id!r}") if isinstance(node, ast.Attribute): owner = _resolve_reference(node.value, context) - if isinstance(owner, Mesh) and not hasattr(owner, node.attr): - axes = owner.names or ("x", "y", "z")[: len(owner.layout.shape)] - named = ", ".join(axes) - owner_name = node.value.id if isinstance(node.value, ast.Name) else "" - raise ParseError.from_node( - node, - context, - f"mesh {owner_name!r} has no axis {node.attr!r}; its axes are: {named}", - ) try: return getattr(owner, node.attr) except AttributeError as error: diff --git a/src/tilefoundry/parser/pattern_nodes.py b/src/tilefoundry/parser/pattern_nodes.py index 7fd0a04a..f3737cf4 100644 --- a/src/tilefoundry/parser/pattern_nodes.py +++ b/src/tilefoundry/parser/pattern_nodes.py @@ -1615,20 +1615,6 @@ def construct(match, children, context): elif match.branch_id == "static_attribute": owner = children["owner"] attribute = match.captures["attribute"] - if isinstance(owner, runtime.Mesh) and not hasattr(owner, attribute): - axes = owner.names or ("x", "y", "z")[: len(owner.layout.shape)] - named = ", ".join(axes) - node = match.node - owner_name = ( - node.value.id - if isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name) - else "" - ) - raise ParseError.from_node( - node, - context, - f"mesh {owner_name!r} has no axis {attribute!r}; its axes are: {named}", - ) try: return getattr(owner, attribute) except AttributeError as error: diff --git a/tests/ir/test_loop_region_range_nested.py b/tests/ir/test_loop_region_range_nested.py index 90c57be0..668a6ac3 100644 --- a/tests/ir/test_loop_region_range_nested.py +++ b/tests/ir/test_loop_region_range_nested.py @@ -26,14 +26,6 @@ _SUM = ReduceKind.SUM -@func -def _tile_start_stop_step(x: Tensor[(8,), "f32"]) -> Tensor[(2,), "f32"]: - acc = tf.zeros(Tensor[(2,), "f32"]) - for window in tile(2, 8, 2): # noqa: F821 - acc = acc + x[window] - return acc - - @func def _range_start_step(x: Tensor[(_M,), "f32"]) -> Tensor[(), "f32"]: acc = tf.reduce(x, axes=(0,), keepdim=False, kind=_SUM) @@ -80,16 +72,6 @@ def test_range_start_step(): assert torch.allclose(out.reshape(()), x[1:n:2].sum()), (n, out) -def test_tile_start_stop_step_maps_to_loop_region_fields() -> None: - loop = _tile_start_stop_step.body - assert loop.start == 2 - assert loop.extent == 8 - assert loop.step == 2 - - x = torch.arange(8, dtype=torch.float32) - assert torch.equal(evaluate(_tile_start_stop_step, x), x[2:4] + x[4:6] + x[6:8]) - - def test_nested_loop_region_outer_carry_in_inner(): x = torch.randn(4, 5) out = evaluate(_nested_sum, x) diff --git a/tests/parser/test_mesh_visibility.py b/tests/parser/test_mesh_visibility.py index cfa8b35c..a2d99d28 100644 --- a/tests/parser/test_mesh_visibility.py +++ b/tests/parser/test_mesh_visibility.py @@ -7,7 +7,6 @@ import pytest from tests._source import import_dsl -from tests.fixtures.placed import persistent_gemm_tiled from tests.fixtures.placed.fused_boundary import FusedBoundary from tests.fixtures.placed.region_boundaries import RegionBoundaries from tests.fixtures.placed.rmsnorm import RmsnormModule @@ -17,7 +16,6 @@ from tilefoundry.ir.hir.function import Function from tilefoundry.ir.hir.mesh_region import MeshRegion from tilefoundry.ir.visitor import collect_exprs, expr_children -from tilefoundry.parser.ast_pattern import ParseError from tilefoundry.visitor_registry.contexts import TypeInferContext from tilefoundry.visitor_registry.visitors import TypeInferVisitor @@ -107,23 +105,3 @@ def test_region_boundaries_capture_external_regions_through_args() -> None: import_dsl(_diagnostic("region_rebind"), "RegionRebind") import_dsl(_diagnostic("region_tuple_rebind"), "RegionTupleRebind") - - -def test_three_argument_tile_accepts_mesh_coordinate_bounds() -> None: - source = Path(persistent_gemm_tiled.__file__).read_text() - - parsed = import_dsl(source, "PersistentGemmTiled") - assert parsed.entry_function().name == "gemm" - - with pytest.raises( - ParseError, - match="mesh 'cta' has no axis 'z'; its axes are: x, y", - ): - import_dsl(source.replace("cta.x * CHUNK_M", "cta.z * CHUNK_M", 1)) - - four_arguments = source.replace( - "tile(cta.x * CHUNK_M, (cta.x + 1) * CHUNK_M, BM)", - "tile(cta.x * CHUNK_M, (cta.x + 1) * CHUNK_M, BM, 1)", - ) - with pytest.raises(ParseError, match=r"tile\(\) takes 2 or 3 arguments"): - import_dsl(four_arguments) From 9bc10d16c7fc10d83d1ad5c534d78846b6761634 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Tue, 22 Sep 2026 00:16:52 +0800 Subject: [PATCH 13/25] docs(analysis): structure parameter constraints --- docs/spec/analysis.md | 33 +++++++++++++++------------------ 1 file changed, 15 insertions(+), 18 deletions(-) diff --git a/docs/spec/analysis.md b/docs/spec/analysis.md index a6d2e8ea..fb00467c 100644 --- a/docs/spec/analysis.md +++ b/docs/spec/analysis.md @@ -387,12 +387,10 @@ concrete arrangement is not reported: no address or per-value buffer identity is a conclusion of this analysis. `rmem` is not address-solved and reports only the largest single projected logical value. -An access relation that retains a parameter with a stated finite range remains -exact: it is neither widened nor unknown, and overlap MAY be proved from it. -Widened relations and relations with an unbounded parameter MUST NOT prove -overlap. - - constraints: + - An access relation that keeps a parameter with a stated finite range is + exact and MAY prove overlap. A widened relation, and one with an unbounded + parameter, MUST NOT. - Placement MUST be settled for the addressable levels `gmem` and `smem` only, once per capacity domain that holds a buffer -- the whole target for a level owned target-wide, one per owning position otherwise -- with two buffers in @@ -1223,19 +1221,18 @@ call site, source expressions shared by identity remain one shared expression in the clone; sharing never aliases the independently cloned body of another call site. -A loop `start` or `extent` MAY depend on the unit running it when every runtime -leaf carries a half-open value range. `Scope.domain` MUST retain the complete -affine expression and represent each such leaf as one identity-deduplicated isl -parameter constrained by that range. A leaf without a range MUST be rejected as -runtime-computed. `step` MUST remain a literal: a parametric stride requires a -product or modulus by two unknowns and has no isl Presburger representation. - -`cardinality` with bounded free parameters MUST enumerate every feasible integer -point and return the true maximum when the parameter box contains at most 4096 -points; for a larger box it MUST return unknown. `Scope.trips()` MUST likewise -fix child and parent domains to the same parameter point, divide their counts -there, and then take the maximum ratio; dividing two independently maximized -counts is not a valid trip bound. +- constraints: + - A loop `start` or `extent` MAY be unit-dependent. Every runtime leaf in one + MUST carry a half-open value range, and `Scope.domain` MUST keep the whole + affine expression with each such leaf as one identity-deduplicated isl + parameter constrained by that range. A leaf without a range MUST be refused. + - A loop `step` MUST be a literal; a parametric stride has no isl + representation. + - `cardinality` MUST enumerate every feasible integer point of a parameter box + of at most `PARAM_POINT_LIMIT` points and return the maximum, and MUST + report unknown for a larger box. + - `Scope.trips()` MUST fix child and parent domains to the same parameter + point before dividing, and take the maximum of those ratios. ### 2.2 Target-selected Analyzers From 94681a9344071cf7918957b681854e241fc60c78 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Tue, 22 Sep 2026 00:16:58 +0800 Subject: [PATCH 14/25] test(analysis): remove superseded persistent coverage --- tests/analysis/test_analysis_families.py | 19 --- tests/fixtures/placed/persistent_gemm_flat.py | 3 - .../fixtures/placed/persistent_gemm_tiled.py | 3 - tests/integration/test_persistent_gemm.py | 147 ------------------ 4 files changed, 172 deletions(-) delete mode 100644 tests/integration/test_persistent_gemm.py diff --git a/tests/analysis/test_analysis_families.py b/tests/analysis/test_analysis_families.py index 71db84c7..700be7fa 100644 --- a/tests/analysis/test_analysis_families.py +++ b/tests/analysis/test_analysis_families.py @@ -15,7 +15,6 @@ import pytest from tests.fixtures.placed.performance_findings import _CrossScopePerformance -from tests.fixtures.placed.persistent_gemm_tiled import BM, BN, PersistentGemmTiled from tests.fixtures.placed.symbolic_offset import ( _LiteralStoreOffset, _SymbolicStoreOffset, @@ -40,14 +39,12 @@ ) from tilefoundry.analysis.errors import AnalysisError from tilefoundry.analysis.memory import MemoryOptions -from tilefoundry.analysis.scope import build_scopes, walk_scopes from tilefoundry.dsl import ConstTensor, DimVar, Mesh, Tensor, Topology, tf from tilefoundry.inspection.analysis_report import render_analysis, render_text from tilefoundry.ir.core import ( Call, get_metadata, ) -from tilefoundry.ir.hir.loop_region import LoopRegion from tilefoundry.ir.hir.math.binary import Binary from tilefoundry.ir.hir.sharding.reshard import Reshard from tilefoundry.ir.hir.tensor.insert_slice import InsertSlice @@ -166,22 +163,6 @@ def _calls(function) -> tuple[Call, ...]: return tuple(expr for expr in collect_exprs(function.body) if isinstance(expr, Call)) -def test_one_pass_maximizes_over_bounded_mesh_parameters() -> None: - root = build_scopes(PersistentGemmTiled, PersistentGemmTiled.entry_function()) - scope = next( - scope - for scope in walk_scopes(root) - if isinstance(scope.owner, LoopRegion) and scope.owner.induction_var.name == "ni" - ) - access = next( - accesses[0] - for call, accesses in scope.outputs["narrow"].values() - if isinstance(call.target, InsertSlice) - ) - - assert scope.one_pass(access) == BM * BN - - def test_performance_orders_a_predecessor_materialized_in_a_child_scope() -> None: """The outer stage waits for the inner stage to complete.""" results = {} diff --git a/tests/fixtures/placed/persistent_gemm_flat.py b/tests/fixtures/placed/persistent_gemm_flat.py index 8cdbfd03..39818ce3 100644 --- a/tests/fixtures/placed/persistent_gemm_flat.py +++ b/tests/fixtures/placed/persistent_gemm_flat.py @@ -52,8 +52,6 @@ def gemm( return out -persistent_gemm_flat = PersistentGemmFlat.entry_function() - __all__ = [ "BK", "BM", @@ -66,5 +64,4 @@ def gemm( "NBLOCKS", "NUM_TILES", "PersistentGemmFlat", - "persistent_gemm_flat", ] diff --git a/tests/fixtures/placed/persistent_gemm_tiled.py b/tests/fixtures/placed/persistent_gemm_tiled.py index f1fb6a69..415e2c23 100644 --- a/tests/fixtures/placed/persistent_gemm_tiled.py +++ b/tests/fixtures/placed/persistent_gemm_tiled.py @@ -51,8 +51,6 @@ def gemm( return out -persistent_gemm_tiled = PersistentGemmTiled.entry_function() - __all__ = [ "BK", "BM", @@ -65,5 +63,4 @@ def gemm( "M", "N", "PersistentGemmTiled", - "persistent_gemm_tiled", ] diff --git a/tests/integration/test_persistent_gemm.py b/tests/integration/test_persistent_gemm.py deleted file mode 100644 index d155261f..00000000 --- a/tests/integration/test_persistent_gemm.py +++ /dev/null @@ -1,147 +0,0 @@ -"""Persistent GEMM schedules cover the output once and remain analyzable.""" - -from __future__ import annotations - -import isl -import pytest -import torch - -from tests.fixtures.placed import persistent_gemm_flat as flat -from tests.fixtures.placed import persistent_gemm_tiled as tiled -from tilefoundry.analysis import analyze -from tilefoundry.analysis.scope import ( - Access, - AccessPrecision, - Scope, - build_scopes, - walk_scopes, -) -from tilefoundry.ir.hir.loop_region import LoopRegion -from tilefoundry.ir.hir.tensor.insert_slice import InsertSlice -from tilefoundry.utils.isl_utils import cardinality - -_FAMILIES = ("compute-cost", "memory", "roofline", "performance") - - -def _loop_scopes(module) -> dict[str, Scope]: - root = build_scopes(module, module.entry_function()) - return { - scope.owner.induction_var.name: scope - for scope in walk_scopes(root) - if isinstance(scope.owner, LoopRegion) - } - - -def _store(scope: Scope) -> tuple[object, Access]: - for call, accesses in scope.outputs["narrow"].values(): - if isinstance(call.target, InsertSlice): - assert len(accesses) == 1 - return call, accesses[0] - raise AssertionError("loop has no InsertSlice output") - - -def _free_domain_count(scope: Scope) -> int | None: - domain = scope.domain - if count := domain.dim(isl.dim_type.PARAM): - domain = domain.project_out(isl.dim_type.PARAM, 0, count) - return cardinality(domain) - - -def _at_unit(image: isl.set, coordinates: tuple[int, ...]) -> isl.set: - assert image.dim(isl.dim_type.PARAM) == len(coordinates) - for axis, coordinate in enumerate(coordinates): - image = image.fix_si(isl.dim_type.PARAM, axis, coordinate) - return image - - -@pytest.mark.parametrize( - "module", - (tiled.PersistentGemmTiled, flat.PersistentGemmFlat), - ids=("tiled", "flat"), -) -def test_persistent_gemm_answers_every_analysis_family(module) -> None: - result = analyze(module, module.entry_function(), analysis=_FAMILIES, level="cta") - - assert set(result.executed) == set(_FAMILIES) - - -def test_tiled_schedule_is_exact_and_partitions_units() -> None: - scopes = _loop_scopes(tiled.PersistentGemmTiled) - assert scopes["mi"].trips() == tiled.CHUNK_M // tiled.BM - assert scopes["ni"].trips() == tiled.CHUNK_N // tiled.BN - assert scopes["ki"].trips() == tiled.K // tiled.BK - assert _free_domain_count(scopes["mi"]) == tiled.M // tiled.BM - assert _free_domain_count(scopes["ni"]) == (tiled.M // tiled.BM) * ( - tiled.N // tiled.BN - ) - - store, access = _store(scopes["ni"]) - assert tuple(store.args[1].type.shape) == (tiled.BM, tiled.BN) - assert access.precision is AccessPrecision.EXACT - written = access.relation.range() - assert _at_unit(written, (0, 0)).is_disjoint(_at_unit(written, (1, 0))) - - -def test_flat_schedule_records_its_quasiaffine_limit() -> None: - """Floor-div/mod offsets widen because affine.py cannot bind those operands.""" - scopes = _loop_scopes(flat.PersistentGemmFlat) - assert scopes["t"].trips() == flat.NUM_TILES // flat.NBLOCKS - assert scopes["ki"].trips() == flat.K // flat.BK - assert _free_domain_count(scopes["t"]) == flat.NUM_TILES - - store, access = _store(scopes["t"]) - assert tuple(store.args[1].type.shape) == (flat.BM, flat.BN) - assert access.precision is AccessPrecision.WIDENED - written = access.relation.range() - assert _at_unit(written, (0,)).is_disjoint(_at_unit(written, (1,))) - - -def _reference_inputs() -> tuple[torch.Tensor, torch.Tensor]: - torch.manual_seed(172) - return ( - torch.randn(tiled.M, tiled.K, dtype=torch.bfloat16), - torch.randn(tiled.K, tiled.N, dtype=torch.bfloat16), - ) - - -def _tiled_schedule(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: - result = torch.empty(tiled.M, tiled.N, dtype=torch.float32) - for x in range(tiled.BX): - for y in range(tiled.BY): - for mi in range(x * tiled.CHUNK_M, (x + 1) * tiled.CHUNK_M, tiled.BM): - for ni in range(y * tiled.CHUNK_N, (y + 1) * tiled.CHUNK_N, tiled.BN): - acc = torch.zeros(tiled.BM, tiled.BN, dtype=torch.float32) - for ki in range(0, tiled.K, tiled.BK): - acc += a[mi : mi + tiled.BM, ki : ki + tiled.BK].float() @ b[ - ki : ki + tiled.BK, ni : ni + tiled.BN - ].float() - result[mi : mi + tiled.BM, ni : ni + tiled.BN] = acc - return result - - -def _flat_schedule(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: - result = torch.empty(flat.M, flat.N, dtype=torch.float32) - for unit in range(flat.NBLOCKS): - for tile_index in range(unit, flat.NUM_TILES, flat.NBLOCKS): - mi = (tile_index // flat.GRID_N) * flat.BM - ni = (tile_index % flat.GRID_N) * flat.BN - acc = torch.zeros(flat.BM, flat.BN, dtype=torch.float32) - for ki in range(0, flat.K, flat.BK): - acc += a[mi : mi + flat.BM, ki : ki + flat.BK].float() @ b[ - ki : ki + flat.BK, ni : ni + flat.BN - ].float() - result[mi : mi + flat.BM, ni : ni + flat.BN] = acc - return result - - -def test_reference_tiling_index_arithmetic_matches_dense_gemm() -> None: - """The reference loops cover each output once and equal dense GEMM. - - The analysis assertions above cover the authored DSL bodies. The evaluator - cannot execute those bodies because it intentionally rejects MeshCoord. - """ - a, b = _reference_inputs() - expected = a.float() @ b.float() - - torch.testing.assert_close(_tiled_schedule(a, b), expected, rtol=1e-5, atol=1e-5) - torch.testing.assert_close(_flat_schedule(a, b), expected, rtol=1e-5, atol=1e-5) From 67ebcf4d7223ae0aa91ebc2b959f1c91de6384dd Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Tue, 22 Sep 2026 00:17:24 +0800 Subject: [PATCH 15/25] refactor(types): colocate ShapeDim with tensor types --- src/tilefoundry/ir/hir/loop_region.py | 2 +- src/tilefoundry/ir/hir/tensor/arange.py | 2 +- src/tilefoundry/ir/hir/tensor/topk.py | 2 +- src/tilefoundry/ir/types/shape_dim.py | 11 ----------- src/tilefoundry/ir/types/shard/int_tuple.py | 2 +- src/tilefoundry/ir/types/shard/mesh.py | 2 +- src/tilefoundry/ir/types/tensor_type.py | 8 +++++++- 7 files changed, 12 insertions(+), 17 deletions(-) delete mode 100644 src/tilefoundry/ir/types/shape_dim.py diff --git a/src/tilefoundry/ir/hir/loop_region.py b/src/tilefoundry/ir/hir/loop_region.py index e281e88f..e8354550 100644 --- a/src/tilefoundry/ir/hir/loop_region.py +++ b/src/tilefoundry/ir/hir/loop_region.py @@ -3,7 +3,7 @@ from dataclasses import dataclass from tilefoundry.ir.core import Expr, Var -from tilefoundry.ir.types.shape_dim import ShapeDim +from tilefoundry.ir.types.tensor_type import ShapeDim @dataclass(unsafe_hash=True) diff --git a/src/tilefoundry/ir/hir/tensor/arange.py b/src/tilefoundry/ir/hir/tensor/arange.py index de377238..9b3ef3c9 100644 --- a/src/tilefoundry/ir/hir/tensor/arange.py +++ b/src/tilefoundry/ir/hir/tensor/arange.py @@ -12,7 +12,7 @@ from tilefoundry.ir.core.register import register_op from tilefoundry.ir.types import DType, TensorType from tilefoundry.ir.types.dim import is_dim_expr -from tilefoundry.ir.types.shape_dim import ShapeDim +from tilefoundry.ir.types.tensor_type import ShapeDim from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( AccessRelations, diff --git a/src/tilefoundry/ir/hir/tensor/topk.py b/src/tilefoundry/ir/hir/tensor/topk.py index 0bdccfec..51c91bce 100644 --- a/src/tilefoundry/ir/hir/tensor/topk.py +++ b/src/tilefoundry/ir/hir/tensor/topk.py @@ -30,7 +30,6 @@ DimVar, is_dim_expr, ) -from tilefoundry.ir.types.shape_dim import ShapeDim from tilefoundry.ir.types.shard import Layout, try_c_order_strides from tilefoundry.ir.types.shard.shard_layout import ( ShardLayout, @@ -38,6 +37,7 @@ layout_axis_to_tensor_axis, shard_layout_of, ) +from tilefoundry.ir.types.tensor_type import ShapeDim from tilefoundry.ir.visitor import ExprVisitor from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( diff --git a/src/tilefoundry/ir/types/shape_dim.py b/src/tilefoundry/ir/types/shape_dim.py deleted file mode 100644 index cdfb2779..00000000 --- a/src/tilefoundry/ir/types/shape_dim.py +++ /dev/null @@ -1,11 +0,0 @@ -"""Type alias for static, symbolic, or expression-valued shape entries. - -The string forward reference avoids an import cycle between core expressions -and tensor types; annotations never require evaluating it at runtime. - -See [types §4](docs/spec/types.md#4-dim--symbolic-shape-dimensions). -""" - -from __future__ import annotations - -type ShapeDim = "int | DimVar | Expr" diff --git a/src/tilefoundry/ir/types/shard/int_tuple.py b/src/tilefoundry/ir/types/shard/int_tuple.py index 681a7db8..05b726f1 100644 --- a/src/tilefoundry/ir/types/shard/int_tuple.py +++ b/src/tilefoundry/ir/types/shard/int_tuple.py @@ -4,7 +4,7 @@ from typing import Union, overload -from tilefoundry.ir.types.shape_dim import ShapeDim +from tilefoundry.ir.types.tensor_type import ShapeDim IntTuple = Union[int, tuple["IntTuple", ...]] diff --git a/src/tilefoundry/ir/types/shard/mesh.py b/src/tilefoundry/ir/types/shard/mesh.py index 510d2439..6c6e5071 100644 --- a/src/tilefoundry/ir/types/shard/mesh.py +++ b/src/tilefoundry/ir/types/shard/mesh.py @@ -3,10 +3,10 @@ from dataclasses import dataclass from functools import lru_cache -from tilefoundry.ir.types.shape_dim import ShapeDim from tilefoundry.ir.types.shard.int_tuple import flatten from tilefoundry.ir.types.shard.layout import ComposedLayout, Layout from tilefoundry.ir.types.shard.layout_algebra import c_order_strides, unflatten +from tilefoundry.ir.types.tensor_type import ShapeDim @dataclass(frozen=True) diff --git a/src/tilefoundry/ir/types/tensor_type.py b/src/tilefoundry/ir/types/tensor_type.py index ebefecbb..cd5c92e4 100644 --- a/src/tilefoundry/ir/types/tensor_type.py +++ b/src/tilefoundry/ir/types/tensor_type.py @@ -6,7 +6,13 @@ from tilefoundry.ir.types.storage import StorageKind, resolve_storage from .dtype import DType -from .shape_dim import ShapeDim + +type ShapeDim = "int | DimVar | Expr" +"""A shape entry: static, symbolic, or expression-valued. + +The string forward reference keeps this module free of any import of core +expressions; annotations never require evaluating it at runtime. +""" def _canonicalize_static_dims(shape: tuple) -> tuple: From daa33960bbff36c5a52ee960503d1de3650549a2 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Tue, 22 Sep 2026 00:38:56 +0800 Subject: [PATCH 16/25] refactor(analysis): separate iteration scope responsibilities --- docs/spec/analysis.md | 24 +- docs/spec/code-organization.md | 8 +- src/tilefoundry/analysis/access.py | 182 ++++++ src/tilefoundry/analysis/allocation.py | 25 +- src/tilefoundry/analysis/api.py | 2 +- src/tilefoundry/analysis/iteration_scope.py | 352 ++++++++++++ src/tilefoundry/analysis/loop_domain.py | 106 ++++ src/tilefoundry/analysis/memory.py | 4 +- src/tilefoundry/analysis/performance.py | 6 +- src/tilefoundry/analysis/scope.py | 586 -------------------- src/tilefoundry/analysis/visitor.py | 4 +- tests/analysis/test_analyze_at_a_size.py | 2 +- 12 files changed, 682 insertions(+), 619 deletions(-) create mode 100644 src/tilefoundry/analysis/access.py create mode 100644 src/tilefoundry/analysis/iteration_scope.py create mode 100644 src/tilefoundry/analysis/loop_domain.py delete mode 100644 src/tilefoundry/analysis/scope.py diff --git a/docs/spec/analysis.md b/docs/spec/analysis.md index fb00467c..60ec77c7 100644 --- a/docs/spec/analysis.md +++ b/docs/spec/analysis.md @@ -1206,12 +1206,13 @@ def analyze( renderings of it and of the Metadata on the IR, and MUST NOT be fields of it. -### 2.1 Shared Scope and Access +### 2.1 Shared IterationScope and Access The normalized HIR is visited once per `analyze()` call. That visit produces a -`Scope` tree parallel to Function/LoopRegion nesting and `Access` relations -for the narrow and device views. `Scope.domain` is the accumulated authored -loop domain; `Scope.accesses` and `Scope.refused` are the only family inputs for +`IterationScope` tree parallel to Function/LoopRegion nesting and `Access` +relations for the narrow and device views. `IterationScope.domain` is the +accumulated authored loop domain; `IterationScope.accesses` and +`IterationScope.refused` are the only family inputs for loop footprints, movement, and placement. An `Access` stores only its relation and allocation expression; storage level and element width are read from the allocation type. A refused descendant makes its owning scope unknown for that @@ -1223,15 +1224,16 @@ call site. - constraints: - A loop `start` or `extent` MAY be unit-dependent. Every runtime leaf in one - MUST carry a half-open value range, and `Scope.domain` MUST keep the whole - affine expression with each such leaf as one identity-deduplicated isl - parameter constrained by that range. A leaf without a range MUST be refused. + MUST carry a half-open value range, and `IterationScope.domain` MUST keep the + whole affine expression with each such leaf as one identity-deduplicated + isl parameter constrained by that range. A leaf without a range MUST be + refused. - A loop `step` MUST be a literal; a parametric stride has no isl representation. - `cardinality` MUST enumerate every feasible integer point of a parameter box of at most `PARAM_POINT_LIMIT` points and return the maximum, and MUST report unknown for a larger box. - - `Scope.trips()` MUST fix child and parent domains to the same parameter + - `IterationScope.trips()` MUST fix child and parent domains to the same parameter point before dividing, and take the maximum of those ratios. ### 2.2 Target-selected Analyzers @@ -1242,8 +1244,8 @@ class AnalyzeContext: target: Target topology_level: str | None options: object | None - root: Scope - current: Scope + root: IterationScope + current: IterationScope AnalysisCallable = Callable[ @@ -1287,7 +1289,7 @@ class Target: - constraints: - `AnalysisCallable` MUST receive the normalized Function graph and one `AnalyzeContext` carrying the exact Module, Target, resolved topology level, - caller options, and the shared root/current `Scope` view. The + caller options, and the shared root/current `IterationScope` view. The `topology_level` MAY be `None` only when the Module declares no topology; options MAY be `None`. - Analyze MUST obtain every root and dependency from the same exact Target diff --git a/docs/spec/code-organization.md b/docs/spec/code-organization.md index 57f701b3..37072bf2 100644 --- a/docs/spec/code-organization.md +++ b/docs/spec/code-organization.md @@ -34,8 +34,10 @@ truth for the directory's structure and invariants. | `analysis/api.py` | [analysis](./analysis.md) | The public composed Analyze operation: shared authored-program check and normalization, one per-call `AnalyzeContext`, dependency closure, ordering, single execution per member, Metadata-ownership enforcement, and semantic result assembly. | | `analysis/registry.py` | [analysis](./analysis.md) | The built-in Analyzer declarations. It re-exports the immutable `Analyzer` descriptor from `target/services.py` and holds no Target dispatch table. | | `analysis/errors.py` | [analysis](./analysis.md) | `AnalysisError`, the one diagnostic the whole analysis layer raises, so catching an analysis failure catches every analysis failure rather than the subset the caller happened to import. | -| `analysis/visitor.py` | [analysis](./analysis.md) | The per-call `AnalyzeContext`, carrying the shared root/current lexical `Scope` while a family traverses its work. | -| `analysis/scope.py` | [analysis](./analysis.md) | The shared `Scope` tree and `Access` relations built once from normalized HIR; families query these views instead of constructing parallel structure. | +| `analysis/visitor.py` | [analysis](./analysis.md) | The per-call `AnalyzeContext`, carrying the shared root/current `IterationScope` while a family traverses its work. | +| `analysis/iteration_scope.py` | [analysis](./analysis.md) | The shared `IterationScope` tree built once from normalized HIR; families query it instead of constructing parallel structure. | +| `analysis/access.py` | [analysis](./analysis.md) | Access relations resolved against the authored iteration scopes. | +| `analysis/loop_domain.py` | [analysis](./analysis.md) | isl iteration domains built from authored loop bounds. | | `analysis/affine.py` | [analysis](./analysis.md) | The shared loop-affine term parser used by scope binding and authored-loop footprint binding, including constant loop strides and bounded invariant offsets. It does not introduce a second affine graph representation. | | `analysis/footprint.py` | [analysis](./analysis.md) | Target-independent authored-loop access images, buffer-view folding, and deduplicated versus repeated byte readings. Requires no separate time map. | | `analysis/report.py` | [analysis](./analysis.md) | Structured analysis report data, including record-family registration, field serialization, and target-aware report-only projections. It depends only on analysis/core modules; inspection consumes it to produce text and source annotations. | @@ -45,7 +47,7 @@ truth for the directory's structure and invariants. | `analysis/compute_cost.py` | [analysis](./analysis.md) | The `compute-cost` family: logical flops per DType and bytes per storage level, from the authored program alone. | | `analysis/memory.py` | [analysis](./analysis.md) | The `memory` family: value lifetimes, per-level peaks, and the capacity comparisons against a target's hierarchy — failing on an over-full addressable level and advising on an over-full cache. | | `analysis/roofline.py` | [analysis](./analysis.md) | The `roofline` family: the recorded work divided by the target's published rates, per Call and aggregated per Function. Adds no count of its own. | -| `analysis/performance.py` | [analysis](./analysis.md) | The `performance` family: occurrences projected from the shared `Scope` tree into flat timeline records and one function envelope, scaled by parallel capacity. It introduces no second scope tree. | +| `analysis/performance.py` | [analysis](./analysis.md) | The `performance` family: occurrences projected from the shared `IterationScope` tree into flat timeline records and one function envelope, scaled by parallel capacity. It introduces no second scope tree. | | `visitor_registry/` | [visitor-registry](./visitor-registry.md) | Shared registry instances and derived visitors: access-relation construction, contexts, ISL helpers, relation building, shard propagation, type inference, verification, code generation, and cost evaluation. | | `visitor_registry/op_cost.py` | [analysis](./analysis.md) | Each operation's per-instance flops and bytes, registered into the shared cost-evaluator registry. Owned here rather than by any target package, because the work an operation asks for follows from its own semantics and operand types on every backend. | | `inspection/analysis_report.py` | [inspection](./inspection.md) | Presentation of analysis-owned report data as text and annotated source. Analysis owns the structured report data and JSON dump; inspection owns how a human reads it. | diff --git a/src/tilefoundry/analysis/access.py b/src/tilefoundry/analysis/access.py new file mode 100644 index 00000000..2651782c --- /dev/null +++ b/src/tilefoundry/analysis/access.py @@ -0,0 +1,182 @@ +"""Access relations resolved against authored iteration scopes.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from enum import Enum, auto + +import isl + +from tilefoundry.ir.core import Call, Expr +from tilefoundry.ir.hir.loop_region import LoopRegion +from tilefoundry.ir.hir.tensor.reshape import Reshape +from tilefoundry.ir.hir.tensor.slice import Slice +from tilefoundry.ir.types import TensorType +from tilefoundry.ir.types.shape_helpers import static_dim_value +from tilefoundry.ir.types.utils import local_type_of +from tilefoundry.utils.isl_utils import has_unbounded_param +from tilefoundry.visitor_registry.access_relation import ( + BoundaryRelation, + index_set, + relation_of, + renaming_relation, +) +from tilefoundry.visitor_registry.contexts import TypeInferContext + +from .affine import LoopAffineTerm, loop_affine_term +from .errors import AnalysisError +from .footprint import _widest_allowed + + +class AccessPrecision(Enum): + """How faithfully an access relation describes the authored access.""" + + EXACT = auto() + WIDENED = auto() + UNKNOWN = auto() + + +@dataclass(frozen=True) +class Access: + """One relation from an iteration scope to the allocation it reaches.""" + + relation: isl.map + buffer: Expr + precision: AccessPrecision = AccessPrecision.EXACT + + +def _parameter_term( + value: object, + relation: isl.map, + name: str, + loops: tuple[LoopRegion, ...], + held: object, + *, + narrow: bool, +) -> tuple[LoopAffineTerm | None, AccessPrecision]: + number = static_dim_value(value) + if number is not None: + return LoopAffineTerm(None, 0, number, number), AccessPrecision.EXACT + try: + term = loop_affine_term(value, loops, narrow=narrow) + except (TypeError, ValueError, NotImplementedError): + term = None + if term is not None: + return term, AccessPrecision.EXACT + return _widest_allowed(relation, name, held), AccessPrecision.WIDENED + + +def _constrain_parameter( + relation: isl.map, + param_index: int, + term: LoopAffineTerm, +) -> isl.map: + local = isl.local_space.from_space(relation.get_space()) + + def placed(kind: str, sign: int, constant: int) -> isl.constraint: + constraint = getattr(isl.constraint, f"alloc_{kind}")(local) + constraint = constraint.set_coefficient_si(isl.dim_type.PARAM, param_index, sign) + if term.loop_axis is not None: + constraint = constraint.set_coefficient_si( + isl.dim_type.IN, term.loop_axis, -sign * term.stride + ) + return constraint.set_constant_si(constant) + + if term.low == term.high: + return relation.add_constraint(placed("equality", 1, -term.low)) + relation = relation.add_constraint(placed("inequality", 1, -term.low)) + return relation.add_constraint(placed("inequality", -1, term.high)) + + +def eliminate_parameters( + relation: isl.map, + parameters: Mapping[str, object] | Sequence[tuple[str, object]], + loops: tuple[LoopRegion, ...], + held: object, + *, + narrow: bool, +) -> tuple[isl.map, AccessPrecision]: + """Eliminate declared parameters using literals, loop terms, or widening.""" + precision = AccessPrecision.EXACT + for name, value in dict(parameters).items(): + param_index = relation.find_dim_by_name(isl.dim_type.PARAM, name) + if param_index < 0: + raise AnalysisError(f"access pattern parameter {name!r} is missing from its relation") + term, resolved_precision = _parameter_term( + value, relation, name, loops, held, narrow=narrow + ) + if resolved_precision is AccessPrecision.WIDENED: + precision = AccessPrecision.WIDENED + if term is not None: + relation = _constrain_parameter(relation, param_index, term) + relation = relation.project_out(isl.dim_type.PARAM, param_index, 1) + return relation, precision + + +def _enclosing_loops(scope: "IterationScope") -> tuple[LoopRegion, ...]: + loops = [] + cursor = scope + while cursor is not None: + if isinstance(cursor.owner, LoopRegion): + loops.append(cursor.owner) + cursor = cursor.parent + loops.reverse() + return tuple(loops) + + +def resolve_access( + operand: Expr, + boundary: BoundaryRelation, + scope: "IterationScope", + ctx: TypeInferContext, + *, + narrow: bool, +) -> Access | None: + """Resolve one declared boundary into an access from its iteration scope.""" + relation = relation_of(boundary.pattern) + precision = AccessPrecision.EXACT + loops = _enclosing_loops(scope) + relation = relation.insert_dims(isl.dim_type.IN, 0, len(loops)) + scope_domain = scope.domain.insert_dims( + isl.dim_type.SET, scope.depth, relation.dim(isl.dim_type.IN) - scope.depth + ) + relation = relation.intersect_domain(scope_domain) + relation, precision = eliminate_parameters( + relation, + getattr(boundary.pattern, "parameters", ()) or (), + loops, + operand.type, + narrow=narrow, + ) + try: + held = local_type_of(operand.type) if narrow else operand.type + except (TypeError, ValueError, NotImplementedError): + return None + box = index_set(tuple(held.shape)) if isinstance(held, TensorType) else None + if box is not None: + relation = relation.intersect_range(box) + while isinstance(operand, Call) and isinstance(operand.target, (Slice, Reshape)): + folded = renaming_relation(operand, ctx, stated=scope.stated_relations(operand, ctx)) + relation = relation.apply_range(relation_of(folded)) + operand = operand.args[0] + relation, folded_precision = eliminate_parameters( + relation, + folded.parameters, + loops, + operand.type, + narrow=narrow, + ) + if folded_precision is AccessPrecision.WIDENED: + precision = AccessPrecision.WIDENED + if precision is AccessPrecision.EXACT and has_unbounded_param(relation): + precision = AccessPrecision.UNKNOWN + return Access(relation, operand, precision) + + +__all__ = [ + "Access", + "AccessPrecision", + "eliminate_parameters", + "resolve_access", +] diff --git a/src/tilefoundry/analysis/allocation.py b/src/tilefoundry/analysis/allocation.py index 9967a65b..edc8dfdb 100644 --- a/src/tilefoundry/analysis/allocation.py +++ b/src/tilefoundry/analysis/allocation.py @@ -21,10 +21,11 @@ from tilefoundry.utils.isl_utils import equates from tilefoundry.visitor_registry.access_relation import index_set +from .access import Access, AccessPrecision from .errors import AnalysisError +from .iteration_scope import IterationScope from .liveness import Liveness from .metadata import ValueLifetime -from .scope import Access, AccessPrecision, Scope class _MemoryOptions(Protocol): @@ -61,7 +62,7 @@ class _OperandConstraint: class _ConstraintContext: """One memory-level model while the HIR visitor applies logical relations.""" - current: Scope + current: IterationScope liveness: Liveness values: tuple[AllocationValue, ...] boxes_by_expr: dict[int, int] @@ -114,7 +115,7 @@ def _access_relation(accesses: tuple[Access, ...]) -> isl.map | None: return result.coalesce() -def _value_domain(value: Expr, scope: Scope) -> isl.set | None: +def _value_domain(value: Expr, scope: IterationScope) -> isl.set | None: """The complete loop-aware coordinate domain of one material value.""" try: held = local_type_of(value.type) @@ -163,7 +164,7 @@ def _complete_operand_relation( operand: Expr, inputs: isl.map | None, outputs: isl.map | None, - scope: Scope, + scope: IterationScope, ) -> isl.map | None: """Return a composed relation only when it covers the whole operand.""" if inputs is None or outputs is None: @@ -178,7 +179,7 @@ def _complete_operand_relation( return None -def _identity_result_relation(node: Call, scope: Scope) -> isl.map | None: +def _identity_result_relation(node: Call, scope: IterationScope) -> isl.map | None: """Map a same-shaped operand to the result while retaining loop axes.""" domain = _value_domain(node, scope) if domain is None: @@ -201,9 +202,11 @@ def _intervals_by_expr(liveness: Liveness) -> dict[int, tuple[int, int]]: } -def _corresponding_carry(source: Expr, operand: Expr, scope: Scope) -> LoopRegion | None: +def _corresponding_carry( + source: Expr, operand: Expr, scope: IterationScope +) -> LoopRegion | None: """Find the carry whose own yield is ``source``, without walking its body.""" - cursor: Scope | None = scope + cursor: IterationScope | None = scope while cursor is not None: loop = cursor.owner if isinstance(loop, LoopRegion): @@ -214,7 +217,9 @@ def _corresponding_carry(source: Expr, operand: Expr, scope: Scope) -> LoopRegio return None -def _tie_is_live(source: Expr, operand: Expr, scope: Scope, liveness: Liveness) -> bool: +def _tie_is_live( + source: Expr, operand: Expr, scope: IterationScope, liveness: Liveness +) -> bool: """Prove that reusing ``operand`` cannot clobber a later ordinary use.""" intervals = _intervals_by_expr(liveness) source_interval = intervals.get(id(source)) @@ -236,7 +241,7 @@ def _tie_is_live(source: Expr, operand: Expr, scope: Scope, liveness: Liveness) def _analyze_operand_constraints( - node: Call, scope: Scope, liveness: Liveness + node: Call, scope: IterationScope, liveness: Liveness ) -> tuple[_OperandConstraint, ...]: """Prove exact logical relations between one result and its operands.""" recorded_inputs = scope.accesses.get("narrow", {}).get(id(node)) @@ -516,7 +521,7 @@ def solve_allocation( memory_level: str, values: tuple[AllocationValue, ...], liveness: Liveness, - root: Scope, + root: IterationScope, *, options: _MemoryOptions, ) -> AllocationResult: diff --git a/src/tilefoundry/analysis/api.py b/src/tilefoundry/analysis/api.py index e1ef0ea2..5282cd1a 100644 --- a/src/tilefoundry/analysis/api.py +++ b/src/tilefoundry/analysis/api.py @@ -15,9 +15,9 @@ from tilefoundry.analysis.check import _resolve_program_geometry, check_program from tilefoundry.analysis.errors import AnalysisError +from tilefoundry.analysis.iteration_scope import ScopeBuilder from tilefoundry.analysis.registry import Analyzer from tilefoundry.analysis.report import render_json, report_data -from tilefoundry.analysis.scope import ScopeBuilder from tilefoundry.analysis.visitor import AnalyzeContext from tilefoundry.dump import DumpFlags, dump from tilefoundry.ir.core import IRMetadata diff --git a/src/tilefoundry/analysis/iteration_scope.py b/src/tilefoundry/analysis/iteration_scope.py new file mode 100644 index 00000000..0af688e9 --- /dev/null +++ b/src/tilefoundry/analysis/iteration_scope.py @@ -0,0 +1,352 @@ +"""Shared authored iteration scopes used by analysis families.""" + +from __future__ import annotations + +from collections.abc import Iterator, Sequence +from dataclasses import dataclass, field + +import isl + +from tilefoundry.ir.core import Call, Expr, value_label, value_labels +from tilefoundry.ir.core.module import Module +from tilefoundry.ir.hir.function import Function +from tilefoundry.ir.hir.loop_region import LoopRegion +from tilefoundry.ir.visitor import expr_children +from tilefoundry.utils.isl_utils import ( + PARAM_POINT_LIMIT, + ParameterBoxTooLarge, + UnboundedParameterBox, + count, + param_points, +) +from tilefoundry.visitor_registry.access_relation import ( + AccessRelations, + access_relation_registry, + projected, + relations_of, + static_bytes, +) +from tilefoundry.visitor_registry.contexts import FunctionScope, TypeInferContext + +from .access import Access, resolve_access +from .errors import AnalysisError +from .loop_domain import induction_name, iteration_domain +from .metadata import BufferFootprint, LoopFootprintMetadata + + +@dataclass(eq=False) +class IterationScope: + """One Function or authored loop, with all accesses below it.""" + + owner: Function | LoopRegion + parent: IterationScope | None + children: tuple[IterationScope, ...] + depth: int + domain: isl.set + accesses: dict[str, dict[int, tuple[Call, tuple[Access, ...]]]] = field(default_factory=dict) + outputs: dict[str, dict[int, tuple[Call, tuple[Access, ...]]]] = field(default_factory=dict) + relations: dict[int, tuple[Call, AccessRelations]] = field(default_factory=dict) + refused: dict[str, frozenset[Call]] = field(default_factory=dict) + _variance: dict[int, frozenset[int]] = field(default_factory=dict, repr=False) + + def stated_relations(self, call: Call, ctx: TypeInferContext) -> AccessRelations: + """Return the Op-declared relations recorded in this scope chain.""" + cursor: IterationScope | None = self + while cursor is not None: + stored = cursor.relations.get(id(call)) + if stored is not None and stored[0] is call: + return stored[1] + cursor = cursor.parent + return relations_of(call, ctx) + + def is_variant(self, value: Expr) -> bool: + """Whether *value* depends on this loop's induction or carry values.""" + if not isinstance(self.owner, LoopRegion): + return False + root = self + while root.parent is not None: + root = root.parent + return id(self) in root._variance.get(id(value), frozenset()) + + def trips(self) -> int: + """Return this scope's iteration count relative to its parent.""" + cached = getattr(self, "_trips_cache", None) + if cached is not None: + return cached + if self.parent is None: + return 1 + if isinstance(self.owner, LoopRegion): + start, extent, step = self.owner.start, self.owner.extent, self.owner.step + if all(isinstance(value, int) for value in (start, extent, step)): + result = 1 if step <= 0 or extent <= start else -(-(extent - start) // step) + self._trips_cache = result + return result + domain = self.domain + parent = self.parent.domain.align_params(domain.get_space()) + domain = domain.align_params(parent.get_space()) + try: + points = param_points(domain.params().intersect(parent.params())) + except UnboundedParameterBox as error: + raise AnalysisError( + f"loop {induction_name(self.owner)!r} has unbounded parameter " + f"{error.parameter!r}, so its trip count cannot be determined" + ) from error + except ParameterBoxTooLarge as error: + raise AnalysisError( + f"loop {induction_name(self.owner)!r} has a parameter box exceeding " + f"the {PARAM_POINT_LIMIT}-point analysis limit, so its trip count cannot be " + "determined" + ) from error + ratios = [] + for point in points: + amount = count(domain.intersect_params(point)) + parent_count = count(parent.intersect_params(point)) + if amount is None or not parent_count: + continue + ratios.append(max(1, amount // parent_count)) + result = max(ratios, default=1) + self._trips_cache = result + return result + + def elements_per_trip(self, access: Access) -> int: + """Count elements reached while this scope's loop axes are held still.""" + cache = getattr(self, "_elements_per_trip_cache", {}) + cached = cache.get(id(access)) + if cached is not None: + return cached + standing = self.domain.insert_dims( + isl.dim_type.SET, + self.depth, + access.relation.dim(isl.dim_type.IN) - self.depth, + ) + relation = access.relation.intersect_domain(standing) + try: + points = param_points(relation.params()) + except UnboundedParameterBox as error: + label = value_label(access.buffer) or type(access.buffer).__name__ + raise AnalysisError( + f"scope access to {label!r} still has unbound parameter {error.parameter!r}" + ) from error + except ParameterBoxTooLarge as error: + raise AnalysisError( + f"scope access parameter box exceeds the {PARAM_POINT_LIMIT}-point analysis limit" + ) from error + amounts = [] + for point in points: + fixed = relation.intersect_params(point) + fixed_standing = standing.intersect_params(point) + for axis in range(self.depth): + low = fixed_standing.dim_min_val(axis) + if not low.is_int(): + raise AnalysisError("scope access has no finite one-pass extent") + fixed_standing = fixed_standing.fix_si( + isl.dim_type.SET, + axis, + low.get_num_si(), + ) + amount = count(fixed.intersect_domain(fixed_standing).range()) + if amount is None: + raise AnalysisError("scope access has no finite one-pass extent") + amounts.append(amount) + result = max(amounts, default=0) + cache[id(access)] = result + self._elements_per_trip_cache = cache + return result + + def accesses_in(self, view: str) -> Iterator[Access]: + """Yield accesses owned by this scope and all descendant scopes.""" + for _call, values in self.accesses.get(view, {}).values(): + yield from values + for child in self.children: + yield from child.accesses_in(view) + + def is_complete(self, view: str) -> bool: + """Whether this scope and every descendant answered every access.""" + if self.refused.get(view): + return False + return all(child.is_complete(view) for child in self.children) + + def footprint(self) -> LoopFootprintMetadata: + """Summarize device and per-unit access bytes for this scope. + + Two structurally equal buffers are distinct allocations, so identity + groups the rows. It does not order or name them: an address is whatever + the allocator handed out this run, and a report exists to be compared + against another run. + """ + rows: dict[tuple[int, str], tuple[Expr, int, int, int]] = {} + for view, scale in (("narrow", "bytes"), ("device", "device_bytes")): + for access in self.accesses_in(view): + try: + amount = self.elements_per_trip(access) + except AnalysisError: + continue + size = static_bytes(access.buffer.type) + if size is None: + continue + device_amount = amount * max(1, self.trips()) + key = (id(access.buffer), str(getattr(access.buffer.type, "storage", "unknown"))) + current = rows.get(key, (access.buffer, len(rows), 0, 0)) + rows[key] = ( + current[0], + current[1], + current[2] + (amount * size if scale == "bytes" else 0), + current[3] + (device_amount * size if scale == "device_bytes" else 0), + ) + entries = list(rows.items()) + labels = value_labels(buffer for _, (buffer, _, _, _) in entries) + ordered = sorted( + (label, memory_level, local, device) + for label, ((_, memory_level), (_, _, local, device)) in zip(labels, entries) + ) + footprints = tuple( + BufferFootprint( + buffer=label, + memory_level=memory_level, + bytes=local, + device_bytes=device, + repeated_bytes=local * self.trips(), + ) + for label, memory_level, local, device in ordered + ) + return LoopFootprintMetadata( + footprints=footprints, + known=self.is_complete("narrow") and self.is_complete("device"), + ) + + +class ScopeBuilder: + """Build one IterationScope tree and its access views for a Function.""" + + def __init__( + self, + module: Module, + graph: Function, + *, + views: Sequence[str] = ("narrow", "device"), + ) -> None: + self.graph = graph + self.views = tuple(views) + self.type_ctx = TypeInferContext(scope=FunctionScope(module, graph)) + self.seeds: dict[int, IterationScope] + self.variance: dict[int, frozenset[int]] + self.seen: set[int] + + def _empty_accesses(self) -> dict[str, dict[int, tuple[Call, tuple[Access, ...]]]]: + return {view: {} for view in self.views} + + def _record_accesses(self, expr: Call, scope: IterationScope) -> None: + if ( + isinstance(expr.target, Function) + or access_relation_registry.lookup(type(expr.target)) is None + ): + return + try: + stated = relations_of(expr, self.type_ctx) + scope.relations[id(expr)] = (expr, stated) + local_relations = projected(stated, expr, self.type_ctx) + except (NotImplementedError, TypeError, ValueError, isl.Error): + for view in self.views: + scope.refused[view] = scope.refused.get(view, frozenset()) | {expr} + return + for view in self.views: + narrow = view == "narrow" + built: list[Access] = [] + for index, boundary in enumerate(local_relations.inputs): + if index >= len(expr.args): + continue + access = resolve_access( + expr.args[index], boundary, scope, self.type_ctx, narrow=narrow + ) + if access is not None: + built.append(access) + scope.accesses.setdefault(view, {})[id(expr)] = (expr, tuple(built)) + written: list[Access] = [] + for boundary in local_relations.outputs: + access = resolve_access(expr, boundary, scope, self.type_ctx, narrow=narrow) + if access is not None: + written.append(access) + scope.outputs.setdefault(view, {})[id(expr)] = (expr, tuple(written)) + + def _record_variance(self, expr: Expr, operands: tuple[Expr, ...]) -> None: + changing: set[int] = set() + for operand in operands: + changing.update(self.variance.get(id(operand), frozenset())) + if (loop := self.seeds.get(id(expr))) is not None: + changing.add(id(loop)) + self.variance[id(expr)] = frozenset(changing) + + def _visit(self, expr: Expr, scope: IterationScope) -> None: + if id(expr) in self.seen: + return + self.seen.add(id(expr)) + if isinstance(expr, LoopRegion): + for operand in expr.init_args: + self._visit(operand, scope) + child = IterationScope( + expr, + scope, + (), + scope.depth + 1, + iteration_domain(expr, scope), + self._empty_accesses(), + ) + scope.children = (*scope.children, child) + self.seeds[id(expr.induction_var)] = child + for carried in expr.carried_args: + self.seeds[id(carried)] = child + self._visit(expr.body, child) + for operand in expr.yield_values: + self._visit(operand, child) + self._record_variance(expr, expr_children(expr)) + return + operands = expr_children(expr) + for operand in operands: + self._visit(operand, scope) + if isinstance(expr, Call): + self._record_accesses(expr, scope) + self._record_variance(expr, operands) + + def build(self) -> IterationScope: + self.seeds = {} + self.variance = {} + self.seen = set() + root = IterationScope( + self.graph, + None, + (), + 0, + iteration_domain(self.graph, None), + self._empty_accesses(), + ) + for param in self.graph.params: + self._visit(param, root) + if self.graph.body is not None: + self._visit(self.graph.body, root) + root._variance = self.variance + return root + + +def build_scopes( + module: Module, + graph: Function, + *, + views: Sequence[str] = ("narrow", "device"), +) -> IterationScope: + """Build the iteration-scope tree and access views in one normalized walk.""" + return ScopeBuilder(module, graph, views=views).build() + + +def walk_scopes(root: IterationScope) -> Iterator[IterationScope]: + """Yield an iteration scope and its descendants in lexical order.""" + yield root + for child in root.children: + yield from walk_scopes(child) + + +__all__ = [ + "IterationScope", + "ScopeBuilder", + "build_scopes", + "walk_scopes", +] diff --git a/src/tilefoundry/analysis/loop_domain.py b/src/tilefoundry/analysis/loop_domain.py new file mode 100644 index 00000000..2feccb63 --- /dev/null +++ b/src/tilefoundry/analysis/loop_domain.py @@ -0,0 +1,106 @@ +"""Construction of isl iteration domains from authored loop bounds.""" + +from __future__ import annotations + +import isl + +from tilefoundry.ir.core import value_label +from tilefoundry.ir.hir.function import Function +from tilefoundry.ir.hir.loop_region import LoopRegion +from tilefoundry.ir.types.dim_isl import range_expr +from tilefoundry.ir.types.shape_helpers import static_dim_value + +from .errors import AnalysisError + + +def induction_name(loop: LoopRegion) -> str: + """Name a loop by the induction variable authored for it.""" + return getattr(loop.induction_var, "name", None) or "" + + +def refuse_runtime_bound(loop: LoopRegion, which: str) -> None: + """Reject a loop bound that has no literal value or stated range.""" + value = getattr(loop, which) + if which == "extent": + raise AnalysisError( + f"loop {induction_name(loop)!r} has a trip count the program computes " + f"at run time from {value_label(value) or 'a value'!r}, so no " + f"per-occurrence total can be scaled by it; bind the extent to a " + f"literal, or state it as an open dimension" + ) + raise AnalysisError( + f"loop {induction_name(loop)!r} takes its {which} from " + f"{value_label(value) or 'a value'!r}, which the program computes at run time; " + f"analysis needs a literal {which} or a stated value range" + ) + + +def bound_to_isl_expr( + loop: LoopRegion, + which: str, + params: dict[str, tuple[int, int] | None], + param_map: dict[str, object], + identities: dict[int, str], +) -> str: + """Render one start or extent from bounded leaves into isl syntax.""" + value = getattr(loop, which) + number = static_dim_value(value) + if number is not None: + return str(number) + try: + rendered = range_expr( + value, + params, + param_map=param_map, + identities=identities, + ) + except (TypeError, ValueError, NotImplementedError, isl.Error): + refuse_runtime_bound(loop, which) + if any(bound is None for bound in params.values()): + refuse_runtime_bound(loop, which) + return rendered + + +def iteration_domain(owner: Function | LoopRegion, parent: "IterationScope | None") -> isl.set: + """Build the accumulated authored iteration domain for one scope owner.""" + if isinstance(owner, Function): + return isl.set("{ [] }") + loops: list[LoopRegion] = [] + cursor = parent + while cursor is not None: + if isinstance(cursor.owner, LoopRegion): + loops.append(cursor.owner) + cursor = cursor.parent + loops.reverse() + params: dict[str, tuple[int, int] | None] = {} + param_map: dict[str, object] = {} + identities: dict[int, str] = {} + bounds: list[str] = [] + for index, loop in enumerate(loops + [owner]): + start = bound_to_isl_expr(loop, "start", params, param_map, identities) + stop = bound_to_isl_expr(loop, "extent", params, param_map, identities) + step = static_dim_value(loop.step) + if step is None: + raise AnalysisError( + f"loop {induction_name(loop)!r} takes its step from " + f"{value_label(loop.step) or 'a value'!r}; analysis needs a literal " + "step, because a parametric stride has no isl representation" + ) + bounds.append(f"{start} <= p{index} < {stop}") + if step != 1: + bounds.append(f"(p{index} - {start}) mod {step} = 0") + for name, bound in params.items(): + if bound is None: + raise AnalysisError(f"loop domain parameter {name!r} has no stated value range") + bounds.append(f"{bound[0]} <= {name} < {bound[1]}") + names = ", ".join(f"p{index}" for index in range(len(loops) + 1)) + prefix = f"[{', '.join(params)}] -> " if params else "" + return isl.set(f"{prefix}{{ [{names}] : {' and '.join(bounds)} }}") + + +__all__ = [ + "bound_to_isl_expr", + "induction_name", + "iteration_domain", + "refuse_runtime_bound", +] diff --git a/src/tilefoundry/analysis/memory.py b/src/tilefoundry/analysis/memory.py index a6f3ae3b..4c4e2ceb 100644 --- a/src/tilefoundry/analysis/memory.py +++ b/src/tilefoundry/analysis/memory.py @@ -1,4 +1,4 @@ -"""Memory-family projection from shared Scope and Access records.""" +"""Memory-family projection from shared IterationScope and Access records.""" from __future__ import annotations @@ -437,7 +437,7 @@ def default_visit_leaf( def analyze_memory(function: Function, context: AnalyzeContext) -> None: - """Attach traffic and per-loop footprints from the shared Scope tree.""" + """Attach traffic and per-loop footprints from the shared IterationScope tree.""" module = context.module topology_level = context.topology_level facts = context.target.get_facts(MemoryHierarchyFacts) diff --git a/src/tilefoundry/analysis/performance.py b/src/tilefoundry/analysis/performance.py index 993d48a5..da0b96b9 100644 --- a/src/tilefoundry/analysis/performance.py +++ b/src/tilefoundry/analysis/performance.py @@ -1,4 +1,4 @@ -"""Place modeled work by querying the shared lexical Scope tree.""" +"""Place modeled work by querying the shared IterationScope tree.""" from __future__ import annotations @@ -16,6 +16,7 @@ from .compute_cost import _local_duration_ns from .errors import AnalysisError from .facts import ParallelCapacityFacts, PerformanceServiceFacts, ThroughputFacts +from .iteration_scope import IterationScope from .metadata import ( ComputeCostMetadata, PerformanceMetadata, @@ -24,7 +25,6 @@ TimelineMetadata, TrafficMetadata, ) -from .scope import Scope from .visitor import AnalyzeContext SELECTOR = "performance" @@ -36,7 +36,7 @@ class PerformanceContext(AnalyzeContext): facts: ThroughputFacts | None = None services: PerformanceServiceFacts | None = None - occurrences: list[tuple[Scope, Call, int]] = field(default_factory=list) + occurrences: list[tuple[IterationScope, Call, int]] = field(default_factory=list) class PerformanceVisitor(ExprVisitor[None]): diff --git a/src/tilefoundry/analysis/scope.py b/src/tilefoundry/analysis/scope.py deleted file mode 100644 index 9bba3f47..00000000 --- a/src/tilefoundry/analysis/scope.py +++ /dev/null @@ -1,586 +0,0 @@ -"""The shared lexical scopes and access relations used by analysis families.""" - -from __future__ import annotations - -from collections.abc import Iterator, Mapping, Sequence -from dataclasses import dataclass, field -from enum import Enum, auto - -import isl - -from tilefoundry.ir.core import Call, Expr, value_label, value_labels -from tilefoundry.ir.core.module import Module -from tilefoundry.ir.hir.function import Function -from tilefoundry.ir.hir.loop_region import LoopRegion -from tilefoundry.ir.hir.tensor.reshape import Reshape -from tilefoundry.ir.hir.tensor.slice import Slice -from tilefoundry.ir.types import TensorType -from tilefoundry.ir.types.dim_isl import range_expr -from tilefoundry.ir.types.shape_helpers import static_dim_value -from tilefoundry.ir.types.utils import local_type_of -from tilefoundry.ir.visitor import expr_children -from tilefoundry.utils.isl_utils import ( - PARAM_POINT_LIMIT, - ParameterBoxTooLarge, - UnboundedParameterBox, - count, - has_unbounded_param, - param_points, -) -from tilefoundry.visitor_registry.access_relation import ( - AccessRelations, - access_relation_registry, - index_set, - projected, - relation_of, - relations_of, - renaming_relation, - static_bytes, -) -from tilefoundry.visitor_registry.contexts import FunctionScope, TypeInferContext - -from .affine import loop_affine_term -from .errors import AnalysisError -from .footprint import _widest_allowed -from .metadata import BufferFootprint, LoopFootprintMetadata - - -class AccessPrecision(Enum): - """How faithfully an access relation describes the authored access.""" - - EXACT = auto() - WIDENED = auto() - UNKNOWN = auto() - - -@dataclass(frozen=True) -class Access: - """One relation from a lexical scope to the allocation it reaches.""" - - relation: isl.map - buffer: Expr - precision: AccessPrecision = AccessPrecision.EXACT - - -@dataclass(eq=False) -class Scope: - """One Function or authored loop, with all accesses below it.""" - - owner: Function | LoopRegion - parent: "Scope | None" - children: tuple["Scope", ...] - depth: int - domain: isl.set - accesses: dict[str, dict[int, tuple[Call, tuple[Access, ...]]]] = field(default_factory=dict) - outputs: dict[str, dict[int, tuple[Call, tuple[Access, ...]]]] = field(default_factory=dict) - relations: dict[int, tuple[Call, AccessRelations]] = field(default_factory=dict) - refused: dict[str, frozenset[Call]] = field(default_factory=dict) - _variance: dict[int, frozenset[int]] = field(default_factory=dict, repr=False) - - def stated_relations(self, call: Call, ctx: TypeInferContext) -> AccessRelations: - """Return the Op-declared relations recorded in this scope chain.""" - cursor: Scope | None = self - while cursor is not None: - stored = cursor.relations.get(id(call)) - if stored is not None and stored[0] is call: - return stored[1] - cursor = cursor.parent - return relations_of(call, ctx) - - def is_variant(self, value: Expr) -> bool: - """Whether *value* depends on this loop's induction or carry values.""" - if not isinstance(self.owner, LoopRegion): - return False - root = self - while root.parent is not None: - root = root.parent - return id(self) in root._variance.get(id(value), frozenset()) - - def is_invariant(self, value: Expr) -> bool: - """Whether *value* is independent of this loop's induction values.""" - return not self.is_variant(value) - - def trips(self) -> int: - """Return this scope's iteration count relative to its parent.""" - cached = getattr(self, "_trips_cache", None) - if cached is not None: - return cached - if self.parent is None: - return 1 - if isinstance(self.owner, LoopRegion): - start, extent, step = self.owner.start, self.owner.extent, self.owner.step - if all(isinstance(value, int) for value in (start, extent, step)): - result = 1 if step <= 0 or extent <= start else -(-(extent - start) // step) - self._trips_cache = result - return result - domain = self.domain - parent = self.parent.domain.align_params(domain.get_space()) - domain = domain.align_params(parent.get_space()) - try: - points = param_points(domain.params().intersect(parent.params())) - except UnboundedParameterBox as error: - raise AnalysisError( - f"loop {_induction_of(self.owner)!r} has unbounded parameter " - f"{error.parameter!r}, so its trip count cannot be determined" - ) from error - except ParameterBoxTooLarge as error: - raise AnalysisError( - f"loop {_induction_of(self.owner)!r} has a parameter box exceeding " - f"the {PARAM_POINT_LIMIT}-point analysis limit, so its trip count cannot be " - "determined" - ) from error - ratios = [] - for point in points: - amount = count(domain.intersect_params(point)) - parent_count = count(parent.intersect_params(point)) - if amount is None or not parent_count: - continue - ratios.append(max(1, amount // parent_count)) - result = max(ratios, default=1) - self._trips_cache = result - return result - - def one_pass(self, access: Access) -> int: - """Count one pass of this scope's relation with loop axes held still.""" - cache = getattr(self, "_one_pass_cache", {}) - cached = cache.get(id(access)) - if cached is not None: - return cached - standing = self.domain.insert_dims( - isl.dim_type.SET, - self.depth, - access.relation.dim(isl.dim_type.IN) - self.depth, - ) - relation = access.relation.intersect_domain(standing) - try: - points = param_points(relation.params()) - except UnboundedParameterBox as error: - label = value_label(access.buffer) or type(access.buffer).__name__ - raise AnalysisError( - f"scope access to {label!r} still has unbound parameter " - f"{error.parameter!r}" - ) from error - except ParameterBoxTooLarge as error: - raise AnalysisError( - "scope access parameter box exceeds the " - f"{PARAM_POINT_LIMIT}-point analysis limit" - ) from error - amounts = [] - for point in points: - fixed = relation.intersect_params(point) - fixed_standing = standing.intersect_params(point) - for axis in range(self.depth): - low = fixed_standing.dim_min_val(axis) - if not low.is_int(): - raise AnalysisError("scope access has no finite one-pass extent") - fixed_standing = fixed_standing.fix_si( - isl.dim_type.SET, - axis, - low.get_num_si(), - ) - amount = count(fixed.intersect_domain(fixed_standing).range()) - if amount is None: - raise AnalysisError("scope access has no finite one-pass extent") - amounts.append(amount) - result = max(amounts, default=0) - cache[id(access)] = result - self._one_pass_cache = cache - return result - - def over(self, access: Access) -> isl.set: - """Return the source elements reached while this scope varies.""" - cache = getattr(self, "_over_cache", {}) - cached = cache.get(id(access)) - if cached is not None: - return cached - domain = self.domain.insert_dims( - isl.dim_type.SET, - self.depth, - access.relation.dim(isl.dim_type.IN) - self.depth, - ) - for axis in range(self.depth): - domain = domain.fix_si( - isl.dim_type.SET, - axis, - domain.dim_min_val(axis).get_num_si(), - ) - result = access.relation.intersect_domain(domain).range() - cache[id(access)] = result - self._over_cache = cache - return result - - def reaching(self, view: str) -> Iterator[Access]: - """Yield accesses owned by this scope and all descendant scopes.""" - for _call, values in self.accesses.get(view, {}).values(): - yield from values - for child in self.children: - yield from child.reaching(view) - - def known(self, view: str) -> bool: - """Whether this scope and every descendant answered every access.""" - if self.refused.get(view): - return False - return all(child.known(view) for child in self.children) - - def footprint(self) -> LoopFootprintMetadata: - """Summarize device and per-unit access bytes for this scope. - - Two structurally equal buffers are distinct allocations, so identity - groups the rows. It does not order or name them: an address is whatever - the allocator handed out this run, and a report exists to be compared - against another run. - """ - rows: dict[tuple[int, str], tuple[Expr, int, int, int]] = {} - for view, scale in (("narrow", "bytes"), ("device", "device_bytes")): - for access in self.reaching(view): - try: - amount = self.one_pass(access) - except AnalysisError: - continue - size = static_bytes(access.buffer.type) - if size is None: - continue - device_amount = amount * max(1, self.trips()) - key = (id(access.buffer), str(getattr(access.buffer.type, "storage", "unknown"))) - current = rows.get(key, (access.buffer, len(rows), 0, 0)) - rows[key] = ( - current[0], - current[1], - current[2] + (amount * size if scale == "bytes" else 0), - current[3] + (device_amount * size if scale == "device_bytes" else 0), - ) - entries = list(rows.items()) - labels = value_labels(buffer for _, (buffer, _, _, _) in entries) - ordered = sorted( - (label, memory_level, local, device) - for label, ((_, memory_level), (_, _, local, device)) in zip(labels, entries) - ) - footprints = tuple( - BufferFootprint( - buffer=label, - memory_level=memory_level, - bytes=local, - device_bytes=device, - repeated_bytes=local * self.trips(), - ) - for label, memory_level, local, device in ordered - ) - return LoopFootprintMetadata( - footprints=footprints, - known=self.known("narrow") and self.known("device"), - ) - - -def _induction_of(loop: LoopRegion) -> str: - """How a diagnostic names one loop: by the variable the author bound it to.""" - return getattr(loop.induction_var, "name", None) or "" - - -def _reject_unbounded_bound(loop: LoopRegion, which: str) -> None: - value = getattr(loop, which) - if which == "extent": - raise AnalysisError( - f"loop {_induction_of(loop)!r} has a trip count the program computes " - f"at run time from {value_label(value) or 'a value'!r}, so no " - f"per-occurrence total can be scaled by it; bind the extent to a " - f"literal, or state it as an open dimension" - ) - raise AnalysisError( - f"loop {_induction_of(loop)!r} takes its {which} from " - f"{value_label(value) or 'a value'!r}, which the program computes at run time; " - f"analysis needs a literal {which} or a stated value range" - ) - - -def _render_bound( - loop: LoopRegion, - which: str, - params: dict[str, tuple[int, int] | None], - param_map: dict[str, object], - identities: dict[int, str], -) -> str: - """Render one start/extent from bounded leaves, or reject an unknown value.""" - value = getattr(loop, which) - number = static_dim_value(value) - if number is not None: - return str(number) - try: - rendered = range_expr( - value, - params, - param_map=param_map, - identities=identities, - ) - except (TypeError, ValueError, NotImplementedError, isl.Error): - _reject_unbounded_bound(loop, which) - if any(bound is None for bound in params.values()): - _reject_unbounded_bound(loop, which) - return rendered - - -def _domain_for(owner: Function | LoopRegion, parent: Scope | None) -> isl.set: - if isinstance(owner, Function): - return isl.set("{ [] }") - loops: list[LoopRegion] = [] - cursor = parent - while cursor is not None: - if isinstance(cursor.owner, LoopRegion): - loops.append(cursor.owner) - cursor = cursor.parent - loops.reverse() - params: dict[str, tuple[int, int] | None] = {} - param_map: dict[str, object] = {} - identities: dict[int, str] = {} - bounds: list[str] = [] - for index, loop in enumerate(loops + [owner]): - start = _render_bound(loop, "start", params, param_map, identities) - stop = _render_bound(loop, "extent", params, param_map, identities) - step = static_dim_value(loop.step) - if step is None: - raise AnalysisError( - f"loop {_induction_of(loop)!r} takes its step from " - f"{value_label(loop.step) or 'a value'!r}; analysis needs a literal " - "step, because a parametric stride has no isl representation" - ) - bounds.append(f"{start} <= p{index} < {stop}") - if step != 1: - bounds.append(f"(p{index} - {start}) mod {step} = 0") - for name, bound in params.items(): - if bound is None: - raise AnalysisError( - f"loop domain parameter {name!r} has no stated value range" - ) - bounds.append(f"{bound[0]} <= {name} < {bound[1]}") - names = ", ".join(f"p{index}" for index in range(len(loops) + 1)) - prefix = f"[{', '.join(params)}] -> " if params else "" - return isl.set(f"{prefix}{{ [{names}] : {' and '.join(bounds)} }}") - - -def _bind_parameters( - relation: isl.map, - parameters: Mapping[str, object] | Sequence[tuple[str, object]], - loops: tuple[LoopRegion, ...], - held: object, - *, - narrow: bool, -) -> tuple[isl.map, AccessPrecision]: - """Bind one relation's stated parameters to literals or loop terms.""" - precision = AccessPrecision.EXACT - for name, value in dict(parameters).items(): - param_index = relation.find_dim_by_name(isl.dim_type.PARAM, name) - if param_index < 0: - raise AnalysisError( - f"access pattern parameter {name!r} is missing from its relation" - ) - number = static_dim_value(value) - if number is None: - term = None - try: - term = loop_affine_term(value, loops, narrow=narrow) - except (TypeError, ValueError, NotImplementedError): - term = None - if term is None: - term = _widest_allowed(relation, name, held) - precision = AccessPrecision.WIDENED - if term is None: - relation = relation.project_out(isl.dim_type.PARAM, param_index, 1) - continue - else: - term = type( - "Term", (), {"loop_axis": None, "stride": 0, "low": number, "high": number} - )() - local = isl.local_space.from_space(relation.get_space()) - - def placed(kind: str, sign: int, constant: int) -> isl.constraint: - constraint = getattr(isl.constraint, f"alloc_{kind}")(local) - constraint = constraint.set_coefficient_si( - isl.dim_type.PARAM, param_index, sign - ) - if term.loop_axis is not None: - constraint = constraint.set_coefficient_si( - isl.dim_type.IN, term.loop_axis, -sign * term.stride - ) - return constraint.set_constant_si(constant) - - if term.low == term.high: - relation = relation.add_constraint(placed("equality", 1, -term.low)) - else: - relation = relation.add_constraint(placed("inequality", 1, -term.low)) - relation = relation.add_constraint(placed("inequality", -1, term.high)) - relation = relation.project_out(isl.dim_type.PARAM, param_index, 1) - return relation, precision - - -def _bind_access( - call: Call, - operand: Expr, - boundary, - scope: Scope, - ctx: TypeInferContext, - *, - narrow: bool, -) -> Access | None: - relation = relation_of(boundary.pattern) - precision = AccessPrecision.EXACT - loops = [] - cursor = scope - while cursor is not None: - if isinstance(cursor.owner, LoopRegion): - loops.append(cursor.owner) - cursor = cursor.parent - loops.reverse() - relation = relation.insert_dims(isl.dim_type.IN, 0, len(loops)) - scope_domain = scope.domain.insert_dims( - isl.dim_type.SET, scope.depth, relation.dim(isl.dim_type.IN) - scope.depth - ) - relation = relation.intersect_domain(scope_domain) - relation, precision = _bind_parameters( - relation, - getattr(boundary.pattern, "parameters", ()) or (), - tuple(loops), - operand.type, - narrow=narrow, - ) - try: - held = local_type_of(operand.type) if narrow else operand.type - except (TypeError, ValueError, NotImplementedError): - return None - box = index_set(tuple(held.shape)) if isinstance(held, TensorType) else None - if box is not None: - relation = relation.intersect_range(box) - while isinstance(operand, Call) and isinstance(operand.target, (Slice, Reshape)): - folded = renaming_relation(operand, ctx, stated=scope.stated_relations(operand, ctx)) - relation = relation.apply_range(relation_of(folded)) - operand = operand.args[0] - relation, folded_precision = _bind_parameters( - relation, - folded.parameters, - tuple(loops), - operand.type, - narrow=narrow, - ) - if folded_precision is AccessPrecision.WIDENED: - precision = AccessPrecision.WIDENED - if precision is AccessPrecision.EXACT and has_unbounded_param(relation): - precision = AccessPrecision.UNKNOWN - return Access(relation, operand, precision) - - -def build_scopes( - module: Module, - graph: Function, - *, - views: Sequence[str] = ("narrow", "device"), -) -> Scope: - """Build the scope tree and both access views in one normalized walk.""" - - def empty_accesses() -> dict[str, dict[int, tuple[Call, tuple[Access, ...]]]]: - return {view: {} for view in views} - - type_ctx = TypeInferContext(scope=FunctionScope(module, graph)) - seeds: dict[int, Scope] = {} - variance: dict[int, frozenset[int]] = {} - seen: set[int] = set() - - def record_accesses(expr: Call, scope: Scope) -> None: - if ( - isinstance(expr.target, Function) - or access_relation_registry.lookup(type(expr.target)) is None - ): - return - try: - stated = relations_of(expr, type_ctx) - scope.relations[id(expr)] = (expr, stated) - local_relations = projected(stated, expr, type_ctx) - except (NotImplementedError, TypeError, ValueError, isl.Error): - for view in views: - scope.refused[view] = scope.refused.get(view, frozenset()) | {expr} - return - for view in views: - narrow = view == "narrow" - built: list[Access] = [] - for index, boundary in enumerate(local_relations.inputs): - if index >= len(expr.args): - continue - access = _bind_access( - expr, expr.args[index], boundary, scope, type_ctx, narrow=narrow - ) - if access is not None: - built.append(access) - scope.accesses.setdefault(view, {})[id(expr)] = (expr, tuple(built)) - written: list[Access] = [] - for boundary in local_relations.outputs: - access = _bind_access(expr, expr, boundary, scope, type_ctx, narrow=narrow) - if access is not None: - written.append(access) - scope.outputs.setdefault(view, {})[id(expr)] = (expr, tuple(written)) - - def record_variance(expr: Expr, operands: tuple[Expr, ...]) -> None: - changing: set[int] = set() - for operand in operands: - changing.update(variance.get(id(operand), frozenset())) - if (loop := seeds.get(id(expr))) is not None: - changing.add(id(loop)) - variance[id(expr)] = frozenset(changing) - - def visit(expr: Expr, scope: Scope) -> None: - if id(expr) in seen: - return - seen.add(id(expr)) - if isinstance(expr, LoopRegion): - for operand in expr.init_args: - visit(operand, scope) - child = Scope( - expr, scope, (), scope.depth + 1, _domain_for(expr, scope), empty_accesses() - ) - scope.children = (*scope.children, child) - seeds[id(expr.induction_var)] = child - for carried in expr.carried_args: - seeds[id(carried)] = child - visit(expr.body, child) - for operand in expr.yield_values: - visit(operand, child) - record_variance(expr, expr_children(expr)) - return - operands = expr_children(expr) - for operand in operands: - visit(operand, scope) - if isinstance(expr, Call): - record_accesses(expr, scope) - record_variance(expr, operands) - - root = Scope(graph, None, (), 0, _domain_for(graph, None), empty_accesses()) - for param in graph.params: - visit(param, root) - if graph.body is not None: - visit(graph.body, root) - root._variance = variance - return root - - -class ScopeBuilder: - """Build one lexical Scope tree and its access views for a derived Function.""" - - def __init__(self, module: Module, graph: Function) -> None: - self.module = module - self.graph = graph - - def build(self) -> Scope: - return build_scopes(self.module, self.graph) - - -def walk_scopes(root: Scope) -> Iterator[Scope]: - """Yield a scope and its descendants in lexical order.""" - yield root - for child in root.children: - yield from walk_scopes(child) - - -__all__ = [ - "Access", - "AccessPrecision", - "Scope", - "ScopeBuilder", - "build_scopes", - "walk_scopes", -] diff --git a/src/tilefoundry/analysis/visitor.py b/src/tilefoundry/analysis/visitor.py index a785d70b..4812b587 100644 --- a/src/tilefoundry/analysis/visitor.py +++ b/src/tilefoundry/analysis/visitor.py @@ -16,8 +16,8 @@ class AnalyzeContext: target: Target topology_level: str | None options: object | None - root: "Scope" - current: "Scope" + root: "IterationScope" + current: "IterationScope" __all__ = ["AnalyzeContext"] diff --git a/tests/analysis/test_analyze_at_a_size.py b/tests/analysis/test_analyze_at_a_size.py index 93a29957..6df4a8c5 100644 --- a/tests/analysis/test_analyze_at_a_size.py +++ b/tests/analysis/test_analyze_at_a_size.py @@ -33,7 +33,7 @@ ) from tilefoundry.analysis.compute_cost import _local_duration_ns from tilefoundry.analysis.errors import AnalysisError -from tilefoundry.analysis.scope import build_scopes, walk_scopes +from tilefoundry.analysis.iteration_scope import build_scopes, walk_scopes from tilefoundry.ir.core import Call, describe_expr, get_metadata from tilefoundry.ir.hir.function import Function from tilefoundry.ir.hir.loop_region import LoopRegion From 279643b5e401820ba75346acc480ce68ee8255c6 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Tue, 22 Sep 2026 00:47:36 +0800 Subject: [PATCH 17/25] refactor(analysis): centralize enclosing loop traversal --- src/tilefoundry/analysis/access.py | 14 +------------- src/tilefoundry/analysis/iteration_scope.py | 11 +++++++++++ src/tilefoundry/analysis/loop_domain.py | 10 ++-------- 3 files changed, 14 insertions(+), 21 deletions(-) diff --git a/src/tilefoundry/analysis/access.py b/src/tilefoundry/analysis/access.py index 2651782c..8ffe6536 100644 --- a/src/tilefoundry/analysis/access.py +++ b/src/tilefoundry/analysis/access.py @@ -114,17 +114,6 @@ def eliminate_parameters( return relation, precision -def _enclosing_loops(scope: "IterationScope") -> tuple[LoopRegion, ...]: - loops = [] - cursor = scope - while cursor is not None: - if isinstance(cursor.owner, LoopRegion): - loops.append(cursor.owner) - cursor = cursor.parent - loops.reverse() - return tuple(loops) - - def resolve_access( operand: Expr, boundary: BoundaryRelation, @@ -135,8 +124,7 @@ def resolve_access( ) -> Access | None: """Resolve one declared boundary into an access from its iteration scope.""" relation = relation_of(boundary.pattern) - precision = AccessPrecision.EXACT - loops = _enclosing_loops(scope) + loops = scope.enclosing_loops() relation = relation.insert_dims(isl.dim_type.IN, 0, len(loops)) scope_domain = scope.domain.insert_dims( isl.dim_type.SET, scope.depth, relation.dim(isl.dim_type.IN) - scope.depth diff --git a/src/tilefoundry/analysis/iteration_scope.py b/src/tilefoundry/analysis/iteration_scope.py index 0af688e9..f69b4053 100644 --- a/src/tilefoundry/analysis/iteration_scope.py +++ b/src/tilefoundry/analysis/iteration_scope.py @@ -68,6 +68,17 @@ def is_variant(self, value: Expr) -> bool: root = root.parent return id(self) in root._variance.get(id(value), frozenset()) + def enclosing_loops(self) -> tuple[LoopRegion, ...]: + """Return this scope's loop owners in outer-to-inner order.""" + loops = [] + cursor: IterationScope | None = self + while cursor is not None: + if isinstance(cursor.owner, LoopRegion): + loops.append(cursor.owner) + cursor = cursor.parent + loops.reverse() + return tuple(loops) + def trips(self) -> int: """Return this scope's iteration count relative to its parent.""" cached = getattr(self, "_trips_cache", None) diff --git a/src/tilefoundry/analysis/loop_domain.py b/src/tilefoundry/analysis/loop_domain.py index 2feccb63..49c50683 100644 --- a/src/tilefoundry/analysis/loop_domain.py +++ b/src/tilefoundry/analysis/loop_domain.py @@ -65,18 +65,12 @@ def iteration_domain(owner: Function | LoopRegion, parent: "IterationScope | Non """Build the accumulated authored iteration domain for one scope owner.""" if isinstance(owner, Function): return isl.set("{ [] }") - loops: list[LoopRegion] = [] - cursor = parent - while cursor is not None: - if isinstance(cursor.owner, LoopRegion): - loops.append(cursor.owner) - cursor = cursor.parent - loops.reverse() + loops = () if parent is None else parent.enclosing_loops() params: dict[str, tuple[int, int] | None] = {} param_map: dict[str, object] = {} identities: dict[int, str] = {} bounds: list[str] = [] - for index, loop in enumerate(loops + [owner]): + for index, loop in enumerate((*loops, owner)): start = bound_to_isl_expr(loop, "start", params, param_map, identities) stop = bound_to_isl_expr(loop, "extent", params, param_map, identities) step = static_dim_value(loop.step) From 890e9c1469b7d797bc2c361ec96f76b1e761429d Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Tue, 22 Sep 2026 00:56:17 +0800 Subject: [PATCH 18/25] refactor(analysis): unify cardinality counting --- docs/spec/analysis.md | 3 ++- src/tilefoundry/analysis/iteration_scope.py | 8 +++--- src/tilefoundry/utils/isl_utils.py | 29 ++++++++++++++------- 3 files changed, 26 insertions(+), 14 deletions(-) diff --git a/docs/spec/analysis.md b/docs/spec/analysis.md index 60ec77c7..4ecfd32c 100644 --- a/docs/spec/analysis.md +++ b/docs/spec/analysis.md @@ -1232,7 +1232,8 @@ call site. representation. - `cardinality` MUST enumerate every feasible integer point of a parameter box of at most `PARAM_POINT_LIMIT` points and return the maximum, and MUST - report unknown for a larger box. + report unknown for a larger box. It MUST count directly when every retained + parameter is already fixed to one integer point. - `IterationScope.trips()` MUST fix child and parent domains to the same parameter point before dividing, and take the maximum of those ratios. diff --git a/src/tilefoundry/analysis/iteration_scope.py b/src/tilefoundry/analysis/iteration_scope.py index f69b4053..7a8d1f9e 100644 --- a/src/tilefoundry/analysis/iteration_scope.py +++ b/src/tilefoundry/analysis/iteration_scope.py @@ -16,7 +16,7 @@ PARAM_POINT_LIMIT, ParameterBoxTooLarge, UnboundedParameterBox, - count, + cardinality, param_points, ) from tilefoundry.visitor_registry.access_relation import ( @@ -110,8 +110,8 @@ def trips(self) -> int: ) from error ratios = [] for point in points: - amount = count(domain.intersect_params(point)) - parent_count = count(parent.intersect_params(point)) + amount = cardinality(domain.intersect_params(point)) + parent_count = cardinality(parent.intersect_params(point)) if amount is None or not parent_count: continue ratios.append(max(1, amount // parent_count)) @@ -155,7 +155,7 @@ def elements_per_trip(self, access: Access) -> int: axis, low.get_num_si(), ) - amount = count(fixed.intersect_domain(fixed_standing).range()) + amount = cardinality(fixed.intersect_domain(fixed_standing).range()) if amount is None: raise AnalysisError("scope access has no finite one-pass extent") amounts.append(amount) diff --git a/src/tilefoundry/utils/isl_utils.py b/src/tilefoundry/utils/isl_utils.py index 9c945714..606654ba 100644 --- a/src/tilefoundry/utils/isl_utils.py +++ b/src/tilefoundry/utils/isl_utils.py @@ -9,7 +9,6 @@ __all__ = [ "as_multi_aff", "cardinality", - "count", "equates", "has_unbounded_param", "involved_dims", @@ -42,8 +41,8 @@ def __init__(self) -> None: super().__init__(f"parameter box exceeds {PARAM_POINT_LIMIT} points") -def count(image: "isl.set") -> int | None: - """Count *image* after every parameter has been fixed by its caller.""" +def _count(image: "isl.set") -> int | None: + """Count *image* without enumerating its parameter space.""" image = image.coalesce() if image.is_box(): amount = 1 @@ -60,6 +59,14 @@ def count(image: "isl.set") -> int | None: return amount.get_num_si() if amount.is_int() else None +def _parameters_are_fixed(context: "isl.set") -> bool: + """Whether every retained parameter has one explicit integer value.""" + return all( + context.plain_get_val_if_fixed(isl.dim_type.PARAM, axis).is_int() + for axis in range(context.dim(isl.dim_type.PARAM)) + ) + + def param_points(image: "isl.set") -> tuple["isl.set", ...]: """Return sets fixing parameters to each feasible point of a small box. @@ -107,18 +114,22 @@ def cardinality(image: "isl.set") -> int | None: """Return a finite point count, maximizing over a small parameter box. A box is counted by multiplying its bounded axis lengths, at a cost set by - rank alone; anything else falls back to ISL's count. With bounded free - parameters, every feasible integer point in a box of at most 4096 points is - counted and the true maximum is returned. A larger or unbounded parameter - box has no answer. + rank alone; anything else falls back to ISL's count. Retained parameters + already fixed to one integer point are counted directly. With bounded free + parameters, every feasible integer point in a box of at most 4096 points + is counted and the true maximum is returned. A larger or unbounded + parameter box has no answer. """ if not image.dim(isl.dim_type.PARAM): - return count(image) + return _count(image) + context = image.params() + if context.is_empty() or _parameters_are_fixed(context): + return _count(image) try: points = param_points(image) except ParameterBoxError: return None - counts = tuple(count(image.intersect_params(point)) for point in points) + counts = tuple(_count(image.intersect_params(point)) for point in points) if any(amount is None for amount in counts): return None return max(counts, default=0) From 49319704b12c74ca56ebfdf644713e5a2342b311 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Tue, 22 Sep 2026 00:59:47 +0800 Subject: [PATCH 19/25] test(analysis): define empty parameter cardinality --- docs/spec/analysis.md | 4 +++- tests/analysis/test_isl_utility.py | 8 ++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/docs/spec/analysis.md b/docs/spec/analysis.md index 4ecfd32c..902d0933 100644 --- a/docs/spec/analysis.md +++ b/docs/spec/analysis.md @@ -1233,7 +1233,9 @@ call site. - `cardinality` MUST enumerate every feasible integer point of a parameter box of at most `PARAM_POINT_LIMIT` points and return the maximum, and MUST report unknown for a larger box. It MUST count directly when every retained - parameter is already fixed to one integer point. + parameter is already fixed to one integer point. An empty parameter context + MUST count as zero; a non-empty context with an unbounded parameter MUST + report unknown. - `IterationScope.trips()` MUST fix child and parent domains to the same parameter point before dividing, and take the maximum of those ratios. diff --git a/tests/analysis/test_isl_utility.py b/tests/analysis/test_isl_utility.py index 0ad8c3c3..b9b9fd5f 100644 --- a/tests/analysis/test_isl_utility.py +++ b/tests/analysis/test_isl_utility.py @@ -154,6 +154,14 @@ def test_cardinality_maximizes_small_parameter_boxes_exactly(): assert cardinality(too_large) is None +def test_cardinality_distinguishes_empty_and_unbounded_parameter_contexts(): + empty = isl.set("[c] -> { [i] : 0 <= i < 4 and c >= 3 and c <= 1 }") + unbounded = isl.set("[c] -> { [i] : i = 0 }") + + assert cardinality(empty) == 0 + assert cardinality(unbounded) is None + + def test_to_domain_encoding(): """Static extents inline. From 2a657b5622c528b1d2536256584028610824968c Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Tue, 22 Sep 2026 01:17:07 +0800 Subject: [PATCH 20/25] refactor(types): clarify dim isl boundaries --- docs/spec/types.md | 11 + src/tilefoundry/analysis/access.py | 2 +- src/tilefoundry/analysis/allocation.py | 2 +- src/tilefoundry/analysis/footprint.py | 2 +- src/tilefoundry/analysis/loop_domain.py | 4 +- src/tilefoundry/ir/hir/nn/rope.py | 2 +- src/tilefoundry/ir/hir/tensor/where.py | 4 +- src/tilefoundry/ir/types/dim.py | 6 + src/tilefoundry/ir/types/dim_isl.py | 191 ++++++++++-------- .../visitor_registry/access_relation.py | 20 +- .../visitor_registry/isl_utility.py | 6 +- tests/analysis/test_analysis_invariants.py | 2 +- tests/analysis/test_isl_utility.py | 49 +++-- 13 files changed, 167 insertions(+), 134 deletions(-) diff --git a/docs/spec/types.md b/docs/spec/types.md index a8a9b61c..5fbdb2e1 100644 --- a/docs/spec/types.md +++ b/docs/spec/types.md @@ -531,6 +531,17 @@ def ceildiv(a, b) -> Expr: symbolic merely by participating in that arithmetic. - `ceildiv(a, b)` MUST compose the existing add, subtract, and floor-divide operations; it does not introduce a distinct Op. + - `ir.types.dim` MUST own dimension IR definitions, construction, and + structural predicates without depending on isl. `ir.types.dim_isl` MUST own + conversion between dimension IR and isl, affine normalization, shape-domain + construction, and conservative value-range queries. + - `dim_to_isl_expr` MUST render one dimension expression while registering + its leaf parameters; `isl_to_dim` MUST decode an isl affine expression using + that parameter map. `shape_to_isl_domain` MUST return one shape's iteration + domain and parameter map. + - `index_set` MUST be the non-negative, all-literal shape specialization of + `shape_to_isl_domain`. It MUST return `None` for a negative, boolean, or + non-literal extent rather than constructing a symbolic or empty domain. - `dim_range(value)` MUST return conservative half-open bounds from `RangeMetadata` before attempting structural dimension arithmetic. A value with neither stored nor structurally derivable bounds returns `None`. diff --git a/src/tilefoundry/analysis/access.py b/src/tilefoundry/analysis/access.py index 8ffe6536..959bfe85 100644 --- a/src/tilefoundry/analysis/access.py +++ b/src/tilefoundry/analysis/access.py @@ -13,12 +13,12 @@ from tilefoundry.ir.hir.tensor.reshape import Reshape from tilefoundry.ir.hir.tensor.slice import Slice from tilefoundry.ir.types import TensorType +from tilefoundry.ir.types.dim_isl import index_set from tilefoundry.ir.types.shape_helpers import static_dim_value from tilefoundry.ir.types.utils import local_type_of from tilefoundry.utils.isl_utils import has_unbounded_param from tilefoundry.visitor_registry.access_relation import ( BoundaryRelation, - index_set, relation_of, renaming_relation, ) diff --git a/src/tilefoundry/analysis/allocation.py b/src/tilefoundry/analysis/allocation.py index edc8dfdb..9251ce6a 100644 --- a/src/tilefoundry/analysis/allocation.py +++ b/src/tilefoundry/analysis/allocation.py @@ -16,10 +16,10 @@ from tilefoundry.ir.hir.tensor.reshape import Reshape from tilefoundry.ir.hir.tensor.slice import Slice from tilefoundry.ir.types import TensorType +from tilefoundry.ir.types.dim_isl import index_set from tilefoundry.ir.types.utils import local_type_of from tilefoundry.ir.visitor import ExprVisitor from tilefoundry.utils.isl_utils import equates -from tilefoundry.visitor_registry.access_relation import index_set from .access import Access, AccessPrecision from .errors import AnalysisError diff --git a/src/tilefoundry/analysis/footprint.py b/src/tilefoundry/analysis/footprint.py index 499fbf65..2a3ea9e2 100644 --- a/src/tilefoundry/analysis/footprint.py +++ b/src/tilefoundry/analysis/footprint.py @@ -7,8 +7,8 @@ import isl from tilefoundry.ir.types import TensorType +from tilefoundry.ir.types.dim_isl import index_set from tilefoundry.utils.isl_utils import cardinality -from tilefoundry.visitor_registry.access_relation import index_set from .affine import LoopAffineTerm diff --git a/src/tilefoundry/analysis/loop_domain.py b/src/tilefoundry/analysis/loop_domain.py index 49c50683..dab3d0fd 100644 --- a/src/tilefoundry/analysis/loop_domain.py +++ b/src/tilefoundry/analysis/loop_domain.py @@ -7,7 +7,7 @@ from tilefoundry.ir.core import value_label from tilefoundry.ir.hir.function import Function from tilefoundry.ir.hir.loop_region import LoopRegion -from tilefoundry.ir.types.dim_isl import range_expr +from tilefoundry.ir.types.dim_isl import dim_to_isl_expr from tilefoundry.ir.types.shape_helpers import static_dim_value from .errors import AnalysisError @@ -48,7 +48,7 @@ def bound_to_isl_expr( if number is not None: return str(number) try: - rendered = range_expr( + rendered = dim_to_isl_expr( value, params, param_map=param_map, diff --git a/src/tilefoundry/ir/hir/nn/rope.py b/src/tilefoundry/ir/hir/nn/rope.py index fc6fc82e..75b36c76 100644 --- a/src/tilefoundry/ir/hir/nn/rope.py +++ b/src/tilefoundry/ir/hir/nn/rope.py @@ -22,12 +22,12 @@ from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._shard_checks import check_multilinear_partials, reject_partials from tilefoundry.ir.types import TupleType +from tilefoundry.ir.types.dim_isl import index_set from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( AccessRelations, AffineAccess, BoundaryRelation, - index_set, iterating, logical_coordinates, reached_at, diff --git a/src/tilefoundry/ir/hir/tensor/where.py b/src/tilefoundry/ir/hir/tensor/where.py index da6d31d4..cfe40499 100644 --- a/src/tilefoundry/ir/hir/tensor/where.py +++ b/src/tilefoundry/ir/hir/tensor/where.py @@ -27,7 +27,7 @@ relations_of, shape_from_relation, ) -from tilefoundry.visitor_registry.isl_utility import to_domain +from tilefoundry.visitor_registry.isl_utility import shape_to_isl_domain from tilefoundry.visitor_registry.shard_propagate import derive_output_shard_layout @@ -50,7 +50,7 @@ def _broadcast_all(shapes: tuple[tuple, ...]) -> tuple: def _maps(shapes: tuple[tuple, ...]) -> tuple[object, tuple[AffineAccess, ...], dict]: out_shape = _broadcast_all(shapes) rank = len(out_shape) - domain, param_map = to_domain(out_shape) + domain, param_map = shape_to_isl_domain(out_shape) dims = [f"d{i}" for i in range(rank)] source = "[" + ", ".join(dims) + "]" maps = [] diff --git a/src/tilefoundry/ir/types/dim.py b/src/tilefoundry/ir/types/dim.py index 9fdab94b..d8092484 100644 --- a/src/tilefoundry/ir/types/dim.py +++ b/src/tilefoundry/ir/types/dim.py @@ -1,3 +1,9 @@ +"""Dimension IR definitions, construction, and structural predicates. + +This module does not depend on isl. Conversion to and from isl, affine +normalization, and value-range queries belong to :mod:`dim_isl`. +""" + from __future__ import annotations from ..core.expr import Call, Constant, Expr diff --git a/src/tilefoundry/ir/types/dim_isl.py b/src/tilefoundry/ir/types/dim_isl.py index 9c42c219..82186d37 100644 --- a/src/tilefoundry/ir/types/dim_isl.py +++ b/src/tilefoundry/ir/types/dim_isl.py @@ -1,4 +1,9 @@ -"""The single ShapeDim <-> isl bridge and normalization authority.""" +"""Conversion between dimension IR and isl, plus affine range queries. + +Dimension definitions and construction stay in :mod:`dim`; this module owns +the dim-to-isl rendering, isl-to-dim decoding, shape-domain construction, +affine normalization, and conservative value-range calculation. +""" from __future__ import annotations @@ -32,7 +37,6 @@ BinaryKind.MAX: DimMax, } - def _is_const(node) -> bool: if isinstance(node, bool): return False @@ -50,9 +54,7 @@ def _bind_param( 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}" - ) + raise ValueError(f"DimVar {name!r} used with conflicting bounds {previous} vs {bound}") else: stored = get_metadata(value, RangeMetadata) if isinstance(value, Expr) else None if identities is None: @@ -78,16 +80,16 @@ def _bind_param( return name -_RANGE_EXPR_VISITOR_TYPE = None +_DIM_TO_ISL_EXPR_VISITOR_TYPE = None _DIM_RANGE_VISITOR_TYPE = None -def _range_expr_visitor_type(): - global _RANGE_EXPR_VISITOR_TYPE - if _RANGE_EXPR_VISITOR_TYPE is None: +def _dim_to_isl_expr_visitor_type(): + global _DIM_TO_ISL_EXPR_VISITOR_TYPE + if _DIM_TO_ISL_EXPR_VISITOR_TYPE is None: from tilefoundry.ir.visitor import ExprVisitor # noqa: PLC0415 - class _RangeExprVisitor(ExprVisitor[str]): + class _DimToIslExprVisitor(ExprVisitor[str]): def __init__(self, params, param_map, identities) -> None: super().__init__() self.params = params @@ -153,67 +155,11 @@ def default_visit(self, value, ctx=None) -> str: raise TypeError(f"unsupported ShapeDim {type(value).__name__}") return _bind_param(value, self.params, self.param_map, self.identities) - _RANGE_EXPR_VISITOR_TYPE = _RangeExprVisitor - return _RANGE_EXPR_VISITOR_TYPE - - -def _dim_range_visitor_type(): - global _DIM_RANGE_VISITOR_TYPE - if _DIM_RANGE_VISITOR_TYPE is None: - from tilefoundry.ir.visitor import ExprVisitor # noqa: PLC0415 - - class _DimRangeVisitor(ExprVisitor[tuple[int, int] | None]): - def visit_Constant(self, value: Constant, ctx=None) -> tuple[int, int]: - number = int(value.value) - return number, number + 1 - - def visit_DimVar(self, value: DimVar, ctx=None) -> tuple[int, int]: - return value.lo, value.hi - - def visit_Call(self, value: Call, ctx=None) -> tuple[int, int] | None: - if type(value.target) is DimMul: - a, b = value.args - if not (_is_const(a) or _is_const(b)): - a_bounds = self.visit(a, ctx) - b_bounds = self.visit(b, ctx) - if a_bounds is None or b_bounds is None: - return None - alo, ahi = a_bounds - blo, bhi = b_bounds - 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(value, params, identities={}) - if any(bound is None for bound in params.values()): - return None - 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 default_visit(self, value, ctx=None) -> tuple[int, int] | None: - if isinstance(value, bool): - raise TypeError("ShapeDim must not be bool") - if isinstance(value, int): - return value, value + 1 - return None - - _DIM_RANGE_VISITOR_TYPE = _DimRangeVisitor - return _DIM_RANGE_VISITOR_TYPE + _DIM_TO_ISL_EXPR_VISITOR_TYPE = _DimToIslExprVisitor + return _DIM_TO_ISL_EXPR_VISITOR_TYPE -def range_expr( +def dim_to_isl_expr( dim, params: dict[str, tuple[int, int] | None], *, @@ -221,7 +167,7 @@ def range_expr( identities: dict[int, str] | None = None, ) -> str: """Render *dim* as an isl expression and register its leaf parameters.""" - return _range_expr_visitor_type()(params, param_map, identities).visit(dim) + return _dim_to_isl_expr_visitor_type()(params, param_map, identities).visit(dim) def _raw_dim_call(op_cls, args: tuple): @@ -237,7 +183,7 @@ def wrap(value): return Call(type=scalar, target=op_cls(), args=tuple(wrap(arg) for arg in args)) -def _visit(expr, param_map: dict[str, object]): +def _visit_isl_expr(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): @@ -249,9 +195,9 @@ def _visit(expr, param_map: dict[str, object]): 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) + return _raw_dim_call(DimSub, (0, _visit_isl_expr(expr.op_arg(0), param_map))) + a = _visit_isl_expr(expr.op_arg(0), param_map) + b = _visit_isl_expr(expr.op_arg(1), param_map) if op == Op.ADD: return _raw_dim_call(DimAdd, (a, b)) if op == Op.SUB: @@ -270,10 +216,10 @@ def _visit(expr, param_map: dict[str, object]): raise NotImplementedError(f"unsupported ast_expr type {type(expr).__name__}") -def to_dim(pw_aff: "isl.pw_aff", param_map: dict[str, object]): +def isl_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) + return _visit_isl_expr(build.expr_from(pw_aff), param_map) def normalize_dim(value): @@ -287,14 +233,14 @@ def normalize_dim(value): try: params: dict[str, tuple[int, int] | None] = {} param_map: dict[str, object] = {} - expr = range_expr( + expr = dim_to_isl_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) + normalized = isl_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 @@ -307,17 +253,75 @@ def normalize_dim_entries(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) + 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_visitor_type(): + global _DIM_RANGE_VISITOR_TYPE + if _DIM_RANGE_VISITOR_TYPE is None: + from tilefoundry.ir.visitor import ExprVisitor # noqa: PLC0415 + + class _DimRangeVisitor(ExprVisitor[tuple[int, int] | None]): + def visit_Constant(self, value: Constant, ctx=None) -> tuple[int, int]: + number = int(value.value) + return number, number + 1 + + def visit_DimVar(self, value: DimVar, ctx=None) -> tuple[int, int]: + return value.lo, value.hi + + def visit_Call(self, value: Call, ctx=None) -> tuple[int, int] | None: + if type(value.target) is DimMul: + a, b = value.args + if not (_is_const(a) or _is_const(b)): + a_bounds = self.visit(a, ctx) + b_bounds = self.visit(b, ctx) + if a_bounds is None or b_bounds is None: + return None + alo, ahi = a_bounds + blo, bhi = b_bounds + 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 = dim_to_isl_expr(value, params, identities={}) + if any(bound is None for bound in params.values()): + return None + 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 default_visit(self, value, ctx=None) -> tuple[int, int] | None: + if isinstance(value, bool): + raise TypeError("ShapeDim must not be bool") + if isinstance(value, int): + return value, value + 1 + return None + + _DIM_RANGE_VISITOR_TYPE = _DimRangeVisitor + return _DIM_RANGE_VISITOR_TYPE + + def dim_range(dim) -> tuple[int, int] | None: """Return conservative half-open value bounds ``[lo, hi)`` for *dim*.""" stored = get_metadata(dim, RangeMetadata) if isinstance(dim, Expr) else None @@ -326,7 +330,7 @@ def dim_range(dim) -> tuple[int, int] | None: return _dim_range_visitor_type()().visit(dim) -def to_domain(extents: tuple) -> tuple: +def shape_to_isl_domain(extents: tuple) -> tuple[isl.set, dict[str, object]]: """Build an iteration domain and its isl-parameter ShapeDim map. A ``Call`` without a value range becomes an unconstrained parameter. Consumers @@ -382,11 +386,22 @@ def bind(name: str, dim, bound: tuple[int, int] | None) -> None: return isl.set(prefix + body), param_map +def index_set(shape: tuple) -> isl.set | None: + """Return the coordinate set for a non-negative literal shape.""" + if any( + not isinstance(extent, int) or isinstance(extent, bool) or extent < 0 for extent in shape + ): + return None + domain, _ = shape_to_isl_domain(shape) + return domain + + __all__ = [ + "dim_to_isl_expr", "dim_range", + "index_set", + "isl_to_dim", "normalize_dim", "normalize_dim_entries", - "range_expr", - "to_dim", - "to_domain", + "shape_to_isl_domain", ] diff --git a/src/tilefoundry/visitor_registry/access_relation.py b/src/tilefoundry/visitor_registry/access_relation.py index 16530088..dfff412d 100644 --- a/src/tilefoundry/visitor_registry/access_relation.py +++ b/src/tilefoundry/visitor_registry/access_relation.py @@ -16,7 +16,7 @@ from tilefoundry.ir.hir._helpers import is_one from tilefoundry.ir.types import TensorType, TupleType, Type, tensor_bytes -from tilefoundry.ir.types.dim_isl import to_dim, to_domain +from tilefoundry.ir.types.dim_isl import index_set, isl_to_dim, shape_to_isl_domain from tilefoundry.ir.types.shape_helpers import static_dim_value from tilefoundry.ir.types.shard.shard_layout import layout_axis_to_tensor_axis from tilefoundry.utils.isl_utils import cardinality @@ -396,7 +396,7 @@ def shape_from_relation( return tuple(extents[axis] for axis in range(rank)) image = reached.range() bindings = parameters_of(relations) - return tuple(to_dim(image.dim_max(axis).add_constant(1), bindings) for axis in range(rank)) + return tuple(isl_to_dim(image.dim_max(axis).add_constant(1), bindings) for axis in range(rank)) def boundary_maps(relations: AccessRelations) -> tuple["isl.map", ...]: @@ -664,7 +664,7 @@ def iterating(extents: "Sequence", relations: "AccessRelations") -> "AccessRelat that space, which is one relation empty somewhere, not a second space. """ try: - domain, named = to_domain(tuple(extents)) + domain, named = shape_to_isl_domain(tuple(extents)) except (TypeError, ValueError, isl.Error) as error: raise ValueError( f"an Op states it iterates {tuple(extents)}, which is no space to walk: {error}" @@ -708,19 +708,6 @@ def relation_of(pattern: "AffineAccess") -> "isl.map": return pattern.relation -def index_set(shape) -> "isl.set | None": - """The coordinates one value legally has, or nothing when its shape is not numbers.""" - if any( - not isinstance(extent, int) or isinstance(extent, bool) or extent < 0 for extent in shape - ): - return None - if not shape: - return isl.set("{ [] }") - names = ", ".join(f"d{index}" for index in range(len(shape))) - guards = " and ".join(f"0 <= d{index} < {extent}" for index, extent in enumerate(shape)) - return isl.set(f"{{ [{names}] : {guards} }}") - - def _as_number(value) -> int | None: """The number a bound parameter's value is, when it is one.""" number = static_dim_value(value) @@ -1226,7 +1213,6 @@ def static_bytes(type_: "Type") -> int | None: "AffineAccess", "BoundaryRelation", "access_relation_registry", - "index_set", "iterating", "identity_access", "identity_relations", diff --git a/src/tilefoundry/visitor_registry/isl_utility.py b/src/tilefoundry/visitor_registry/isl_utility.py index f20bdbe0..fc471a64 100644 --- a/src/tilefoundry/visitor_registry/isl_utility.py +++ b/src/tilefoundry/visitor_registry/isl_utility.py @@ -1,5 +1,5 @@ -"""Compatibility exports for the ShapeDim <-> isl bridge owned by IR types.""" +"""Compatibility exports for the dim <-> isl bridge owned by IR types.""" -from tilefoundry.ir.types.dim_isl import dim_range, to_dim, to_domain +from tilefoundry.ir.types.dim_isl import dim_range, isl_to_dim, shape_to_isl_domain -__all__ = ["dim_range", "to_dim", "to_domain"] +__all__ = ["dim_range", "isl_to_dim", "shape_to_isl_domain"] diff --git a/tests/analysis/test_analysis_invariants.py b/tests/analysis/test_analysis_invariants.py index 1111b4f2..02c846ca 100644 --- a/tests/analysis/test_analysis_invariants.py +++ b/tests/analysis/test_analysis_invariants.py @@ -53,6 +53,7 @@ make_tensor_type, tensor_bytes, ) +from tilefoundry.ir.types.dim_isl import index_set from tilefoundry.ir.types.shard import Topology, make_mesh from tilefoundry.ir.types.shard.shard_layout import Split as ShardSplit from tilefoundry.ir.types.storage import StorageKind @@ -61,7 +62,6 @@ AffineAccess, BoundaryRelation, access_relation_registry, - index_set, local_relations_of, relation_of, relations_of, diff --git a/tests/analysis/test_isl_utility.py b/tests/analysis/test_isl_utility.py index b9b9fd5f..a14a962c 100644 --- a/tests/analysis/test_isl_utility.py +++ b/tests/analysis/test_isl_utility.py @@ -1,4 +1,4 @@ -"""isl_utility — dim_range, to_domain encode, to_dim decode.""" +"""dim_isl — range queries and conversion between dimensions and isl.""" from __future__ import annotations @@ -22,7 +22,13 @@ DimVar, simplify_dim, ) -from tilefoundry.ir.types.dim_isl import dim_range, normalize_dim, to_dim, to_domain +from tilefoundry.ir.types.dim_isl import ( + dim_range, + index_set, + isl_to_dim, + normalize_dim, + shape_to_isl_domain, +) from tilefoundry.utils.isl_utils import cardinality P = DimVar("P", 2048, 1_048_577) @@ -145,9 +151,7 @@ def calls(value): def test_cardinality_maximizes_small_parameter_boxes_exactly(): - interior_maximum = isl.set( - "[c] -> { [p] : 0 <= c <= 8 and 0 <= p and p < c and p < 8 - c }" - ) + interior_maximum = isl.set("[c] -> { [p] : 0 <= c <= 8 and 0 <= p and p < c and p < 8 - c }") too_large = isl.set("[x, y] -> { [p] : 0 <= x < 132 and 0 <= y < 132 and p = x + y }") assert cardinality(interior_maximum) == 4 @@ -162,38 +166,49 @@ def test_cardinality_distinguishes_empty_and_unbounded_parameter_contexts(): assert cardinality(unbounded) is None -def test_to_domain_encoding(): +def test_shape_to_isl_domain_encoding(): """Static extents inline. Static extents inline; a bare DimVar keeps its own param name; a composite mints one opaque param bounded by ``dim_range`` and dedups across axes on the canonical expression. """ - dom, param_map = to_domain((8, 4)) + dom, param_map = shape_to_isl_domain((8, 4)) assert dom.dim(isl.dim_type.PARAM) == 0 assert dom.dim(isl.dim_type.SET) == 2 assert param_map == {} - dom, param_map = to_domain((P,)) + dom, param_map = shape_to_isl_domain((P,)) assert dom.get_dim_name(isl.dim_type.PARAM, 0) == "P" assert param_map == {"P": P} d = simplify_dim(DimFloorDiv, (P, 4)) - dom, param_map = to_domain((d, 128, d)) + dom, param_map = shape_to_isl_domain((d, 128, d)) assert dom.dim(isl.dim_type.PARAM) == 1 name = dom.get_dim_name(isl.dim_type.PARAM, 0) lo, hi = dim_range(d) assert f"{lo} <= {name} <= {hi - 1}" in str(dom) assert param_map[name] is d - dom, param_map = to_domain(()) + dom, param_map = shape_to_isl_domain(()) assert dom.dim(isl.dim_type.SET) == 0 assert param_map == {} -def test_to_domain_same_name_conflicting_bounds_raises(): +def test_shape_to_isl_domain_same_name_conflicting_bounds_raises(): with pytest.raises(ValueError, match="conflicting bounds"): - to_domain((DimVar("S", 1, 8), DimVar("S", 1, 16))) + shape_to_isl_domain((DimVar("S", 1, 8), DimVar("S", 1, 16))) + + +def test_index_set_is_the_nonnegative_literal_shape_special_case(): + for shape in ((8, 4), (), (0, 3)): + domain, param_map = shape_to_isl_domain(shape) + assert param_map == {} + assert index_set(shape).is_equal(domain) + + assert index_set((-1, 3)) is None + assert index_set((P,)) is None + assert index_set((True,)) is None def test_round_trip_lossless_for_every_dim_kind(): @@ -205,11 +220,11 @@ def test_round_trip_lossless_for_every_dim_kind(): the map is refused instead of being invented as an opaque dim nothing can resolve later. """ - assert to_dim(isl.pw_aff("{ [42] }"), {}) == 42 + assert isl_to_dim(isl.pw_aff("{ [42] }"), {}) == 42 named = isl.pw_aff("[P] -> { [P] }") - assert to_dim(named, {"P": P}) is P + assert isl_to_dim(named, {"P": P}) is P with pytest.raises(ValueError, match="no known ShapeDim"): - to_dim(named, {}) + isl_to_dim(named, {}) dims = ( 128, @@ -222,8 +237,8 @@ def test_round_trip_lossless_for_every_dim_kind(): simplify_dim(DimMod, (P, 128)), simplify_dim(DimAdd, (128, simplify_dim(DimFloorDiv, (P, 4)))), ) - domain, param_map = to_domain(dims) + domain, param_map = shape_to_isl_domain(dims) recovered = tuple( - to_dim(domain.dim_max(i).add_constant(1), param_map) for i in range(len(dims)) + isl_to_dim(domain.dim_max(i).add_constant(1), param_map) for i in range(len(dims)) ) assert recovered == dims From bedef606932a7c27d067e4cde95e28e2fb86f492 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Tue, 22 Sep 2026 01:34:27 +0800 Subject: [PATCH 21/25] refactor(visitor-registry): isolate type inference visitor --- docs/spec/visitor-registry.md | 4 + src/tilefoundry/analysis/check.py | 2 +- src/tilefoundry/ir/hir/specialize.py | 2 +- src/tilefoundry/ir/types/dim_isl.py | 1 + src/tilefoundry/parser/ast_pattern.py | 2 +- src/tilefoundry/passes/pass_manager.py | 2 +- src/tilefoundry/visitor_registry/__init__.py | 6 +- src/tilefoundry/visitor_registry/typeinfer.py | 269 +++++++++++++++ src/tilefoundry/visitor_registry/visitors.py | 307 +----------------- tests/analysis/test_mesh_region_cost.py | 2 +- tests/evaluator/eval_utils.py | 2 +- tests/ir/test_function_call_typeinfer.py | 2 +- tests/ir/test_simplify_dim.py | 2 +- tests/ir/types/test_mma_fragment_layouts.py | 2 +- tests/ops/ir/cost_utils.py | 3 +- tests/ops/ir/test_arange.py | 2 +- tests/ops/ir/test_cache_update.py | 3 +- tests/ops/ir/test_quant.py | 2 +- tests/ops/ir/test_slice.py | 3 +- tests/ops/ir/test_topk.py | 2 +- tests/ops/ir/typeinfer_utils.py | 2 +- tests/parser/test_mesh_visibility.py | 2 +- 22 files changed, 306 insertions(+), 318 deletions(-) create mode 100644 src/tilefoundry/visitor_registry/typeinfer.py diff --git a/docs/spec/visitor-registry.md b/docs/spec/visitor-registry.md index 1ada6b38..8598a11d 100644 --- a/docs/spec/visitor-registry.md +++ b/docs/spec/visitor-registry.md @@ -251,6 +251,10 @@ def _(call: Call, ctx: TypeInferContext) -> TensorType: ... Visitor: +`TypeInferVisitor` and `inference_type` live in +`tilefoundry.visitor_registry.typeinfer`. Verification, code-generation, and +cost-evaluation visitors remain in `tilefoundry.visitor_registry.visitors`. + ```python class TypeInferVisitor(ExprVisitor[Type]): def __init__(self, *, memo=None, owns_body=True, ranges=False): ... diff --git a/src/tilefoundry/analysis/check.py b/src/tilefoundry/analysis/check.py index 13107804..565f7156 100644 --- a/src/tilefoundry/analysis/check.py +++ b/src/tilefoundry/analysis/check.py @@ -49,7 +49,7 @@ from tilefoundry.ir.visitor import BindingSubstitutionCloner, collect_exprs from tilefoundry.target import UnsupportedCapabilityError from tilefoundry.visitor_registry.contexts import FunctionScope, TypeInferContext -from tilefoundry.visitor_registry.visitors import inference_type +from tilefoundry.visitor_registry.typeinfer import inference_type from .errors import AnalysisError from .facts import ParallelCapacityFacts, PerformanceServiceFacts diff --git a/src/tilefoundry/ir/hir/specialize.py b/src/tilefoundry/ir/hir/specialize.py index 36393407..f7be3a84 100644 --- a/src/tilefoundry/ir/hir/specialize.py +++ b/src/tilefoundry/ir/hir/specialize.py @@ -29,7 +29,7 @@ from tilefoundry.ir.types.tensor_type import TensorType, Type from tilefoundry.ir.visitor import ExprCloner, ExprVisitor, ExprWalker from tilefoundry.visitor_registry.contexts import TypeInferContext -from tilefoundry.visitor_registry.visitors import TypeInferVisitor +from tilefoundry.visitor_registry.typeinfer import TypeInferVisitor from .function import Function diff --git a/src/tilefoundry/ir/types/dim_isl.py b/src/tilefoundry/ir/types/dim_isl.py index 82186d37..106b0a74 100644 --- a/src/tilefoundry/ir/types/dim_isl.py +++ b/src/tilefoundry/ir/types/dim_isl.py @@ -37,6 +37,7 @@ BinaryKind.MAX: DimMax, } + def _is_const(node) -> bool: if isinstance(node, bool): return False diff --git a/src/tilefoundry/parser/ast_pattern.py b/src/tilefoundry/parser/ast_pattern.py index 62dab399..4b5161e7 100644 --- a/src/tilefoundry/parser/ast_pattern.py +++ b/src/tilefoundry/parser/ast_pattern.py @@ -91,7 +91,7 @@ from tilefoundry.ir.visitor import BindingSubstitutionCloner from tilefoundry.target import MemoryHierarchyFacts, Target, UnsupportedCapabilityError from tilefoundry.visitor_registry.contexts import FunctionScope, TypeInferContext -from tilefoundry.visitor_registry.visitors import TypeInferVisitor, inference_type +from tilefoundry.visitor_registry.typeinfer import TypeInferVisitor, inference_type T = TypeVar("T") _TYPE_INFER_CONTEXT = "" diff --git a/src/tilefoundry/passes/pass_manager.py b/src/tilefoundry/passes/pass_manager.py index 7f9581a6..008edb2b 100644 --- a/src/tilefoundry/passes/pass_manager.py +++ b/src/tilefoundry/passes/pass_manager.py @@ -17,7 +17,7 @@ from tilefoundry.ir.tir.prim_function import PrimFunction from tilefoundry.ir.tir.verify import verify_prim_function from tilefoundry.visitor_registry.contexts import FunctionScope, TypeInferContext -from tilefoundry.visitor_registry.visitors import inference_type +from tilefoundry.visitor_registry.typeinfer import inference_type from .pass_base import Pass diff --git a/src/tilefoundry/visitor_registry/__init__.py b/src/tilefoundry/visitor_registry/__init__.py index 1d949d8d..4277e002 100644 --- a/src/tilefoundry/visitor_registry/__init__.py +++ b/src/tilefoundry/visitor_registry/__init__.py @@ -2,9 +2,11 @@ To keep this package importable early (ir.core imports from here during its own __init__), the package __init__ re-exports **only** the -lightweight registry bits. Contexts and Visitors live in submodules -``contexts`` and ``visitors`` and should be imported from there. +lightweight registry bits. Contexts live in ``contexts``, type inference in +``typeinfer``, and the remaining derived visitors in ``visitors``; import each +from its owning submodule. """ + from __future__ import annotations from .registries import ( diff --git a/src/tilefoundry/visitor_registry/typeinfer.py b/src/tilefoundry/visitor_registry/typeinfer.py new file mode 100644 index 00000000..68761213 --- /dev/null +++ b/src/tilefoundry/visitor_registry/typeinfer.py @@ -0,0 +1,269 @@ +"""Type inference over expressions, functions, and authored regions.""" + +from __future__ import annotations + +from dataclasses import replace + +from tilefoundry.ir.core.expr import Call, Constant, Expr, Tuple, Var +from tilefoundry.ir.core.metadata import ( + RangeMetadata, + attach_metadata, + detach_metadata, +) +from tilefoundry.ir.hir.function import Function +from tilefoundry.ir.hir.loop_region import LoopRegion +from tilefoundry.ir.hir.mesh_region import MeshRegion +from tilefoundry.ir.hir.sharding.reshard import Reshard as HirReshard +from tilefoundry.ir.tir.shape import ShapeOf +from tilefoundry.ir.types.callable_type import callable_type_for +from tilefoundry.ir.types.shard.mesh import composed +from tilefoundry.ir.types.shard.scope_match import covered_by_scope, storage_reaches +from tilefoundry.ir.types.shard.shard_layout import ShardLayout +from tilefoundry.ir.types.substitute import canonicalize_dims +from tilefoundry.ir.types.tensor_type import TupleType, Type +from tilefoundry.ir.types.utils import types_compatible +from tilefoundry.ir.visitor import ExprVisitor + +from .contexts import FunctionScope, TypeInferContext, TypeInferResults +from .registries import typeinfer_registry + + +class TypeInferVisitor(ExprVisitor[Type]): + """Derive one type for each ``Expr`` kind. + + See [hir §1.1](docs/spec/hir.md#11-function) and + [visitor-registry §4](docs/spec/visitor-registry.md#4-instance-1--typeinfer). + + The context owns one identity memo for the current inference scope. A + missing leaf raises through ``default_visit_leaf`` rather than trusting a + stale ``expr.type``. + """ + + def __init__(self, *, memo=None, owns_body: bool = True, ranges: bool = False) -> None: + super().__init__(memo=memo) + self._memo_supplied = memo is not None + self._visit_depth = 0 + self._owns_body = owns_body + self._ranges = ranges + + def visit(self, expr: Expr, ctx: TypeInferContext) -> Type: + """Derive one type while preserving the active execution domain.""" + outermost = self._visit_depth == 0 + if outermost: + if self._memo_supplied: + ctx = replace(ctx, memo=self._memo) + else: + self._memo = ctx.memo + cached = id(expr) in self._memo + self._visit_depth += 1 + try: + results = super().visit(expr, ctx) + if not isinstance(results, TypeInferResults): + results = TypeInferResults(results) + result = canonicalize_dims(results.type) + self._memo[id(expr)] = (expr, result) + if self._owns_body: + expr.type = result + if self._ranges and not cached: + self._record_range(expr, results) + return result + finally: + self._visit_depth -= 1 + + @staticmethod + def _record_range(expr: Expr, results: TypeInferResults) -> None: + """Replace one expression's derived range with this inference result.""" + if results.value_range is None: + detach_metadata(expr, RangeMetadata) + else: + attach_metadata(expr, RangeMetadata(*results.value_range)) + + def visit_leaf_Var(self, var: Var, _operands, ctx: TypeInferContext) -> Type: + return var.annotation + + def visit_leaf_Constant(self, c: Constant, _operands, ctx: TypeInferContext) -> Type: + return c.type + + def visit_leaf_Call(self, call: Call, arg_types, ctx: TypeInferContext) -> Type: + target = call.target + if ctx.current_mesh is not None and not isinstance(target, HirReshard): + for index, arg_type in enumerate(arg_types): + layout = getattr(arg_type, "layout", None) + if not isinstance(layout, ShardLayout): + continue + if not covered_by_scope(layout.mesh, ctx.current_mesh): + ctx.error( + call, + f"input {index} is laid out more finely than the scope it " + "runs in; write it inside that scope, or reshard it back first", + ) + if not storage_reaches(arg_type.storage, layout.mesh, ctx.current_mesh): + ctx.error( + call, + f"input {index} is laid out more coarsely and kept in " + f"{arg_type.storage.name.lower()}, which does not reach the units " + "this runs on; reshard it to smem or gmem first", + ) + if isinstance(target, Function): + return self._call_function(call, target, arg_types, ctx) + op_cls = type(target) + fn = typeinfer_registry.lookup(op_cls) + if fn is None: + ctx.error(call, f"no typeinfer registered for {op_cls.__name__}") + return fn(call, ctx) + + def _call_function( + self, + call: Call, + callee: Function, + arg_types: tuple[Type, ...], + ctx: TypeInferContext, + ) -> Type: + child = ctx.child_for(callee) + supplied = tuple( + param for param in callee.params if not (child is not None and param.is_const) + ) + if len(arg_types) != len(supplied): + kind = "activation(s)" if child is not None else "parameter(s)" + ctx.error( + call, + f"hir Function call {callee.name!r}: arity mismatch — " + f"callee declares {len(supplied)} {kind}, call passed {len(arg_types)}", + ) + + given = iter(enumerate(arg_types)) + memo = {} + for param in callee.params: + if child is not None and param.is_const: + memo[id(param)] = (param, param.annotation) + continue + index, arg_type = next(given) + declared = param.annotation + if not types_compatible(declared, arg_type): + ctx.error( + call, + f"hir Function call {callee.name!r}: arg {index} type mismatch — " + f"callee param {param.name!r} expects {declared!r}, got {arg_type!r}", + ) + memo[id(param)] = (param, arg_type) + + if callee.body is None or callee.variants: + return callee.return_type + + key = (id(callee), arg_types) + cached = ctx.instantiated_memo.get(key) + if cached is not None: + return cached + result = TypeInferVisitor(memo=memo, owns_body=False, ranges=False).visit( + callee.body, ctx.for_callee(callee) + ) + ctx.instantiated_memo[key] = result + return result + + def visit_leaf_Tuple(self, tup: Tuple, operands, ctx: TypeInferContext) -> Type: + """Visit Tuple. + + Structural: the field types of the (possibly just-elaborated) + elements, never the node's own stamped ``.type`` ([hir §1.1](docs/spec/hir.md#11-function)). + """ + return TupleType(fields=operands) + + def visit_LoopRegion(self, region: LoopRegion, ctx: TypeInferContext) -> Type: + """Infer a loop after binding its induction and carried variables.""" + inits = tuple(self.visit(arg, ctx) for arg in region.init_args) + memo = { + **self._memo, + id(region.induction_var): (region.induction_var, region.induction_var.annotation), + **{id(phi): (phi, type_) for phi, type_ in zip(region.carried_args, inits)}, + } + inner = TypeInferVisitor( + memo=memo, + owns_body=self._owns_body, + ranges=self._ranges, + ) + body_type = inner.visit(region.body, ctx) + for y in region.yield_values: + inner.visit(y, ctx) + if not region.carried_args: + return body_type + if len(region.carried_args) == 1: + return inner.visit(region.carried_args[0], ctx) + return TupleType(fields=tuple(inner.visit(phi, ctx) for phi in region.carried_args)) + + def visit_MeshRegion(self, expr: MeshRegion, ctx: TypeInferContext) -> Type: + """Type a region against the participants in force inside it. + + Entering a scope composes it onto the mesh in force, so the body reads + the whole nesting. The resulting mesh goes down on a child context, so + the caller's own scope survives the recursion. The region types as its + body does: who runs the work is not a fact about the shape of what it + produced, and what one unit costs is cost's question. + """ + arg_types = tuple(self.visit(arg, ctx) for arg in expr.args) + if len(arg_types) != len(expr.params): + ctx.error( + expr, + f"mesh scope expects {len(expr.params)} argument(s), got {len(arg_types)}", + ) + for index, (param, arg_type) in enumerate(zip(expr.params, arg_types, strict=True)): + if not types_compatible(param.annotation, arg_type): + ctx.error( + expr, + f"mesh scope arg {index} type mismatch for param {param.name!r}", + ) + memo = { + **ctx.memo, + **{ + id(param): (param, arg_type) + for param, arg_type in zip(expr.params, arg_types, strict=True) + }, + } + from tilefoundry.ir.hir.verify import _verify_isolated # noqa: PLC0415 + + _verify_isolated(expr, ctx) + mesh = composed((ctx.current_mesh, expr.mesh)) if ctx.current_mesh else expr.mesh + return self.visit(expr.body, replace(ctx, current_mesh=mesh, memo=memo)) + + def visit_Function(self, fn: Function, ctx: TypeInferContext) -> Type: + """Refresh one complete function after binding its parameter types.""" + if ctx.scope is not None and ctx.scope.function is not fn: + ctx = replace(ctx, scope=FunctionScope(ctx.scope.module, fn)) + memo = {id(param): (param, param.annotation) for param in fn.params} + if fn.body is not None: + TypeInferVisitor( + memo=memo, + owns_body=self._owns_body, + ranges=self._ranges, + ).visit(fn.body, replace(ctx, memo=memo)) + for nested in fn.variants: + TypeInferVisitor( + owns_body=self._owns_body, + ranges=self._ranges, + ).visit(nested, ctx) + return callable_type_for(fn.params, fn.return_type) + + def visit_leaf_ShapeOf(self, shape_of: ShapeOf, _operands, ctx: TypeInferContext) -> Type: + """A ``tir.ShapeOf`` always carries its own concrete (rank-0 i32) type at construction. + + A ``tir.ShapeOf`` always carries its own concrete (rank-0 i32) + type at construction; it has no children to derive from. + """ + return shape_of.type + + def default_visit_leaf(self, expr: Expr, _operands, ctx: TypeInferContext) -> Type: + ctx.error(expr, f"no typeinfer rule for Expr subclass {type(expr).__name__}") + + +def inference_type( + expr: Expr, + ctx: TypeInferContext | None = None, + *, + ranges: bool = False, +) -> Type: + """Infer *expr*; ``ranges=True`` refreshes value ranges but not stored types.""" + return TypeInferVisitor(owns_body=False, ranges=ranges).visit( + expr, ctx if ctx is not None else TypeInferContext() + ) + + +__all__ = ["TypeInferVisitor", "inference_type"] diff --git a/src/tilefoundry/visitor_registry/visitors.py b/src/tilefoundry/visitor_registry/visitors.py index d2c7030d..cfbfb4f0 100644 --- a/src/tilefoundry/visitor_registry/visitors.py +++ b/src/tilefoundry/visitor_registry/visitors.py @@ -1,302 +1,25 @@ -"""Derived Visitors — TypeInferVisitor / VerifyVisitor / CodegenVisitor / CostEvaluator. +"""Verification, code generation, and cost-evaluation visitors.""" -`DispatchRegistry` instance with a traversal skeleton from -tilefoundry.ir.visitor. - -The `registry` is exposed as an advanced constructor param (default: the -canonical module-level registry for that analysis). Default path uses the -module-level registry directly; passing a custom one is an advanced -extension point for sandbox tests or grouped dispatch. -""" from __future__ import annotations from collections.abc import Callable -from dataclasses import replace -from tilefoundry.ir.core.expr import Call, Constant, Expr, Tuple, Var -from tilefoundry.ir.core.metadata import ( - RangeMetadata, - attach_metadata, - detach_metadata, -) -from tilefoundry.ir.hir.function import Function -from tilefoundry.ir.hir.loop_region import LoopRegion -from tilefoundry.ir.hir.mesh_region import MeshRegion -from tilefoundry.ir.hir.sharding.reshard import Reshard as HirReshard -from tilefoundry.ir.tir.shape import ShapeOf +from tilefoundry.ir.core.expr import Call, Expr from tilefoundry.ir.tir.stmt import Stmt from tilefoundry.ir.tir.stmts import Evaluate, MeshScope -from tilefoundry.ir.types.callable_type import callable_type_for -from tilefoundry.ir.types.shard.mesh import composed -from tilefoundry.ir.types.shard.scope_match import covered_by_scope, storage_reaches -from tilefoundry.ir.types.shard.shard_layout import ShardLayout -from tilefoundry.ir.types.substitute import canonicalize_dims -from tilefoundry.ir.types.tensor_type import TupleType, Type, UnitType -from tilefoundry.ir.types.utils import types_compatible -from tilefoundry.ir.visitor import ExprVisitor, ExprWalker, StmtVisitor +from tilefoundry.ir.types.tensor_type import UnitType +from tilefoundry.ir.visitor import ExprWalker, StmtVisitor -from .contexts import ( - Cost, - CostContext, - FunctionScope, - TypeInferContext, - TypeInferResults, - VerifyContext, -) +from .contexts import Cost, CostContext, VerifyContext from .registries import ( DispatchRegistry, Role, cost_evaluator_registry, spelled, - typeinfer_registry, verify_stmt_registry, ) -class TypeInferVisitor(ExprVisitor[Type]): - """Derive one type for each ``Expr`` kind. - - See [hir §1.1](docs/spec/hir.md#11-function) and - [visitor-registry §4](docs/spec/visitor-registry.md#4-instance-1--typeinfer). - - The context owns one identity memo for the current inference scope. A - missing leaf raises through ``default_visit_leaf`` rather than trusting a - stale ``expr.type``. - """ - - def __init__( - self, *, memo=None, owns_body: bool = True, ranges: bool = False - ) -> None: - super().__init__(memo=memo) - self._memo_supplied = memo is not None - self._visit_depth = 0 - self._owns_body = owns_body - self._ranges = ranges - - def visit(self, expr: Expr, ctx: TypeInferContext) -> Type: - """Derive one type while preserving the active execution domain.""" - outermost = self._visit_depth == 0 - if outermost: - if self._memo_supplied: - ctx = replace(ctx, memo=self._memo) - else: - self._memo = ctx.memo - cached = id(expr) in self._memo - self._visit_depth += 1 - try: - results = super().visit(expr, ctx) - if not isinstance(results, TypeInferResults): - results = TypeInferResults(results) - result = canonicalize_dims(results.type) - self._memo[id(expr)] = (expr, result) - if self._owns_body: - expr.type = result - if self._ranges and not cached: - self._record_range(expr, results) - return result - finally: - self._visit_depth -= 1 - - @staticmethod - def _record_range(expr: Expr, results: TypeInferResults) -> None: - """Replace one expression's derived range with this inference result.""" - if results.value_range is None: - detach_metadata(expr, RangeMetadata) - else: - attach_metadata(expr, RangeMetadata(*results.value_range)) - - def visit_leaf_Var(self, var: Var, _operands, ctx: TypeInferContext) -> Type: - return var.annotation - - def visit_leaf_Constant(self, c: Constant, _operands, ctx: TypeInferContext) -> Type: - return c.type - - def visit_leaf_Call(self, call: Call, arg_types, ctx: TypeInferContext) -> Type: - target = call.target - if ctx.current_mesh is not None and not isinstance(target, HirReshard): - for index, arg_type in enumerate(arg_types): - layout = getattr(arg_type, "layout", None) - if not isinstance(layout, ShardLayout): - continue - if not covered_by_scope(layout.mesh, ctx.current_mesh): - ctx.error( - call, - f"input {index} is laid out more finely than the scope it " - "runs in; write it inside that scope, or reshard it back first", - ) - if not storage_reaches(arg_type.storage, layout.mesh, ctx.current_mesh): - ctx.error( - call, - f"input {index} is laid out more coarsely and kept in " - f"{arg_type.storage.name.lower()}, which does not reach the units " - "this runs on; reshard it to smem or gmem first", - ) - if isinstance(target, Function): - return self._call_function(call, target, arg_types, ctx) - op_cls = type(target) - fn = typeinfer_registry.lookup(op_cls) - if fn is None: - ctx.error(call, f"no typeinfer registered for {op_cls.__name__}") - return fn(call, ctx) - - def _call_function( - self, - call: Call, - callee: Function, - arg_types: tuple[Type, ...], - ctx: TypeInferContext, - ) -> Type: - child = ctx.child_for(callee) - supplied = tuple( - param for param in callee.params if not (child is not None and param.is_const) - ) - if len(arg_types) != len(supplied): - kind = "activation(s)" if child is not None else "parameter(s)" - ctx.error( - call, - f"hir Function call {callee.name!r}: arity mismatch — " - f"callee declares {len(supplied)} {kind}, call passed {len(arg_types)}", - ) - - given = iter(enumerate(arg_types)) - memo = {} - for param in callee.params: - if child is not None and param.is_const: - memo[id(param)] = (param, param.annotation) - continue - index, arg_type = next(given) - declared = param.annotation - if not types_compatible(declared, arg_type): - ctx.error( - call, - f"hir Function call {callee.name!r}: arg {index} type mismatch — " - f"callee param {param.name!r} expects {declared!r}, got {arg_type!r}", - ) - memo[id(param)] = (param, arg_type) - - if callee.body is None or callee.variants: - return callee.return_type - - key = (id(callee), arg_types) - cached = ctx.instantiated_memo.get(key) - if cached is not None: - return cached - result = TypeInferVisitor(memo=memo, owns_body=False, ranges=False).visit( - callee.body, ctx.for_callee(callee) - ) - ctx.instantiated_memo[key] = result - return result - - def visit_leaf_Tuple(self, tup: Tuple, operands, ctx: TypeInferContext) -> Type: - """Visit Tuple. - - Structural: the field types of the (possibly just-elaborated) - elements, never the node's own stamped ``.type`` ([hir §1.1](docs/spec/hir.md#11-function)). - """ - return TupleType(fields=operands) - - def visit_LoopRegion(self, region: LoopRegion, ctx: TypeInferContext) -> Type: - """Infer a loop after binding its induction and carried variables.""" - inits = tuple(self.visit(arg, ctx) for arg in region.init_args) - memo = { - **self._memo, - id(region.induction_var): (region.induction_var, region.induction_var.annotation), - **{id(phi): (phi, type_) for phi, type_ in zip(region.carried_args, inits)}, - } - inner = TypeInferVisitor( - memo=memo, - owns_body=self._owns_body, - ranges=self._ranges, - ) - body_type = inner.visit(region.body, ctx) - for y in region.yield_values: - inner.visit(y, ctx) - if not region.carried_args: - return body_type - if len(region.carried_args) == 1: - return inner.visit(region.carried_args[0], ctx) - return TupleType(fields=tuple(inner.visit(phi, ctx) for phi in region.carried_args)) - - def visit_MeshRegion(self, expr: MeshRegion, ctx: TypeInferContext) -> Type: - """Type a region against the participants in force inside it. - - Entering a scope composes it onto the mesh in force, so the body reads - the whole nesting. The resulting mesh goes down on a child context, so - the caller's own scope survives the recursion. The region types as its - body does: who runs the work is not a fact about the shape of what it - produced, and what one unit costs is cost's question. - """ - arg_types = tuple(self.visit(arg, ctx) for arg in expr.args) - if len(arg_types) != len(expr.params): - ctx.error( - expr, - f"mesh scope expects {len(expr.params)} argument(s), got {len(arg_types)}", - ) - for index, (param, arg_type) in enumerate( - zip(expr.params, arg_types, strict=True) - ): - if not types_compatible(param.annotation, arg_type): - ctx.error( - expr, - f"mesh scope arg {index} type mismatch for param {param.name!r}", - ) - memo = { - **ctx.memo, - **{ - id(param): (param, arg_type) - for param, arg_type in zip(expr.params, arg_types, strict=True) - }, - } - from tilefoundry.ir.hir.verify import _verify_isolated # noqa: PLC0415 - - _verify_isolated(expr, ctx) - mesh = composed((ctx.current_mesh, expr.mesh)) if ctx.current_mesh else expr.mesh - return self.visit(expr.body, replace(ctx, current_mesh=mesh, memo=memo)) - - def visit_Function(self, fn: Function, ctx: TypeInferContext) -> Type: - """Refresh one complete function after binding its parameter types.""" - if ctx.scope is not None and ctx.scope.function is not fn: - ctx = replace(ctx, scope=FunctionScope(ctx.scope.module, fn)) - memo = {id(param): (param, param.annotation) for param in fn.params} - if fn.body is not None: - TypeInferVisitor( - memo=memo, - owns_body=self._owns_body, - ranges=self._ranges, - ).visit(fn.body, replace(ctx, memo=memo)) - for nested in fn.variants: - TypeInferVisitor( - owns_body=self._owns_body, - ranges=self._ranges, - ).visit(nested, ctx) - return callable_type_for(fn.params, fn.return_type) - - def visit_leaf_ShapeOf( - self, shape_of: ShapeOf, _operands, ctx: TypeInferContext - ) -> Type: - """A ``tir.ShapeOf`` always carries its own concrete (rank-0 i32) type at construction. - - A ``tir.ShapeOf`` always carries its own concrete (rank-0 i32) - type at construction; it has no children to derive from. - """ - return shape_of.type - - def default_visit_leaf(self, expr: Expr, _operands, ctx: TypeInferContext) -> Type: - ctx.error(expr, f"no typeinfer rule for Expr subclass {type(expr).__name__}") - - -def inference_type( - expr: Expr, - ctx: TypeInferContext | None = None, - *, - ranges: bool = False, -) -> Type: - """Infer *expr*; ``ranges=True`` refreshes value ranges but not stored types.""" - return TypeInferVisitor(owns_body=False, ranges=ranges).visit( - expr, ctx if ctx is not None else TypeInferContext() - ) - - class VerifyVisitor(StmtVisitor[None]): """Dispatch verify_stmt per Stmt subclass. @@ -312,18 +35,11 @@ def __init__( ctx: VerifyContext, registry: DispatchRegistry = verify_stmt_registry, ) -> None: - - - - self.ctx = ctx self.registry = registry def generic_visit(self, stmt: Stmt) -> None: if isinstance(stmt, Evaluate): - - - op = stmt.callable fn = self.registry.lookup(type(op)) if fn is not None: @@ -401,15 +117,8 @@ def __init__( def visit_Call(self, call: Call, ctx: CostContext) -> Cost: fn = self.registry.lookup(type(call.target)) if fn is None: - ctx.error( - call, f"no cost evaluator registered for {type(call.target).__name__}" - ) + ctx.error(call, f"no cost evaluator registered for {type(call.target).__name__}") return fn(call, ctx) -__all__ = [ - "TypeInferVisitor", - "inference_type", - "VerifyVisitor", - "CodegenVisitor", - "CostEvaluator", -] + +__all__ = ["VerifyVisitor", "CodegenVisitor", "CostEvaluator"] diff --git a/tests/analysis/test_mesh_region_cost.py b/tests/analysis/test_mesh_region_cost.py index 1a848f61..c2469b2c 100644 --- a/tests/analysis/test_mesh_region_cost.py +++ b/tests/analysis/test_mesh_region_cost.py @@ -18,7 +18,7 @@ from tilefoundry.ir.visitor import collect_exprs from tilefoundry.target import CudaTarget from tilefoundry.visitor_registry.contexts import TypeInferContext -from tilefoundry.visitor_registry.visitors import TypeInferVisitor +from tilefoundry.visitor_registry.typeinfer import TypeInferVisitor _TARGET = CudaTarget("nvidia.h200_sxm") _TOPOLOGIES = (Topology("cta", 1), Topology("thread", 4)) diff --git a/tests/evaluator/eval_utils.py b/tests/evaluator/eval_utils.py index 3b20d387..ddf98a8a 100644 --- a/tests/evaluator/eval_utils.py +++ b/tests/evaluator/eval_utils.py @@ -17,7 +17,7 @@ from tilefoundry.ir.core import Call, Var from tilefoundry.ir.types import DType, TensorType, make_tensor_type from tilefoundry.visitor_registry.contexts import TypeInferContext -from tilefoundry.visitor_registry.visitors import TypeInferVisitor +from tilefoundry.visitor_registry.typeinfer import TypeInferVisitor _DTYPE_OF = { torch.float32: DType.f32, diff --git a/tests/ir/test_function_call_typeinfer.py b/tests/ir/test_function_call_typeinfer.py index 943756b9..be78fa59 100644 --- a/tests/ir/test_function_call_typeinfer.py +++ b/tests/ir/test_function_call_typeinfer.py @@ -24,7 +24,7 @@ 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 -from tilefoundry.visitor_registry.visitors import TypeInferVisitor +from tilefoundry.visitor_registry.typeinfer import TypeInferVisitor _F = DType.f32 _M = make_mesh((4,)) diff --git a/tests/ir/test_simplify_dim.py b/tests/ir/test_simplify_dim.py index 0bb6919e..afa55943 100644 --- a/tests/ir/test_simplify_dim.py +++ b/tests/ir/test_simplify_dim.py @@ -29,7 +29,7 @@ ) 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 +from tilefoundry.visitor_registry.typeinfer import TypeInferVisitor def _i64(v: int) -> Constant: diff --git a/tests/ir/types/test_mma_fragment_layouts.py b/tests/ir/types/test_mma_fragment_layouts.py index d42a20c0..5934c3b5 100644 --- a/tests/ir/types/test_mma_fragment_layouts.py +++ b/tests/ir/types/test_mma_fragment_layouts.py @@ -15,7 +15,7 @@ from tilefoundry.ir.types import DType, TensorType from tilefoundry.ir.types.shard import ShardLayout, Split, product from tilefoundry.ir.types.storage import StorageKind -from tilefoundry.visitor_registry.visitors import inference_type +from tilefoundry.visitor_registry.typeinfer import inference_type _ATOM = make_atom(SM80_16x8x16_F32BF16BF16F32_TN) A_FRAG_SHARD = _ATOM.A diff --git a/tests/ops/ir/cost_utils.py b/tests/ops/ir/cost_utils.py index ca3f7a7c..1001e96e 100644 --- a/tests/ops/ir/cost_utils.py +++ b/tests/ops/ir/cost_utils.py @@ -9,7 +9,8 @@ from tilefoundry.ir.types import DType, TensorType from tilefoundry.ir.types.shard import Topology from tilefoundry.visitor_registry.contexts import CostContext, TrafficBytes, TypeInferContext -from tilefoundry.visitor_registry.visitors import CostEvaluator, TypeInferVisitor +from tilefoundry.visitor_registry.typeinfer import TypeInferVisitor +from tilefoundry.visitor_registry.visitors import CostEvaluator @dataclass(frozen=True) diff --git a/tests/ops/ir/test_arange.py b/tests/ops/ir/test_arange.py index 3e4a823a..f3d77c4f 100644 --- a/tests/ops/ir/test_arange.py +++ b/tests/ops/ir/test_arange.py @@ -23,7 +23,7 @@ from tilefoundry.ir.types.shard import Layout, Mesh, Topology, composed from tilefoundry.ir.types.storage import StorageKind from tilefoundry.visitor_registry.contexts import TrafficBytes, TypeInferContext -from tilefoundry.visitor_registry.visitors import TypeInferVisitor +from tilefoundry.visitor_registry.typeinfer import TypeInferVisitor _N = DimVar("arange_n", 1, 17) diff --git a/tests/ops/ir/test_cache_update.py b/tests/ops/ir/test_cache_update.py index 6350afd5..6c36c768 100644 --- a/tests/ops/ir/test_cache_update.py +++ b/tests/ops/ir/test_cache_update.py @@ -39,7 +39,8 @@ from tilefoundry.ir.visitor import collect_exprs from tilefoundry.target import CudaTarget from tilefoundry.visitor_registry.contexts import CostContext, TrafficBytes, TypeInferContext -from tilefoundry.visitor_registry.visitors import CostEvaluator, TypeInferVisitor +from tilefoundry.visitor_registry.typeinfer import TypeInferVisitor +from tilefoundry.visitor_registry.visitors import CostEvaluator def _ref(cache, cur_pos, s, new): diff --git a/tests/ops/ir/test_quant.py b/tests/ops/ir/test_quant.py index 06e8705b..0aabdb9c 100644 --- a/tests/ops/ir/test_quant.py +++ b/tests/ops/ir/test_quant.py @@ -35,7 +35,7 @@ split_target_axes, ) from tilefoundry.visitor_registry.contexts import TypeInferContext -from tilefoundry.visitor_registry.visitors import TypeInferVisitor +from tilefoundry.visitor_registry.typeinfer import TypeInferVisitor _BF = DType.bf16 _FP8 = DType.fp8e4m3 diff --git a/tests/ops/ir/test_slice.py b/tests/ops/ir/test_slice.py index d53e4f5a..fced03e9 100644 --- a/tests/ops/ir/test_slice.py +++ b/tests/ops/ir/test_slice.py @@ -22,7 +22,8 @@ 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 -from tilefoundry.visitor_registry.visitors import CostEvaluator, TypeInferVisitor +from tilefoundry.visitor_registry.typeinfer import TypeInferVisitor +from tilefoundry.visitor_registry.visitors import CostEvaluator _F = DType.f32 _M = make_mesh((4,)) diff --git a/tests/ops/ir/test_topk.py b/tests/ops/ir/test_topk.py index beb28416..f7f23e54 100644 --- a/tests/ops/ir/test_topk.py +++ b/tests/ops/ir/test_topk.py @@ -40,7 +40,7 @@ from tilefoundry.ir.types.shard import Layout, make_mesh from tilefoundry.ir.types.shard.shard_layout import Broadcast, Partial, ShardLayout, Split from tilefoundry.visitor_registry.contexts import TypeInferContext -from tilefoundry.visitor_registry.visitors import TypeInferVisitor +from tilefoundry.visitor_registry.typeinfer import TypeInferVisitor _F32 = DType.f32 _I64 = DType.i64 diff --git a/tests/ops/ir/typeinfer_utils.py b/tests/ops/ir/typeinfer_utils.py index 3094d42d..0c6b2146 100644 --- a/tests/ops/ir/typeinfer_utils.py +++ b/tests/ops/ir/typeinfer_utils.py @@ -18,7 +18,7 @@ from tilefoundry.ir.types.shard.layout import Layout from tilefoundry.ir.types.shard.shard_layout import ShardLayout, Split, shard_layout_local_shape from tilefoundry.visitor_registry.contexts import TypeInferContext -from tilefoundry.visitor_registry.visitors import TypeInferVisitor +from tilefoundry.visitor_registry.typeinfer import TypeInferVisitor def raw_shard_tensor_type( diff --git a/tests/parser/test_mesh_visibility.py b/tests/parser/test_mesh_visibility.py index a2d99d28..d78064d3 100644 --- a/tests/parser/test_mesh_visibility.py +++ b/tests/parser/test_mesh_visibility.py @@ -17,7 +17,7 @@ from tilefoundry.ir.hir.mesh_region import MeshRegion from tilefoundry.ir.visitor import collect_exprs, expr_children from tilefoundry.visitor_registry.contexts import TypeInferContext -from tilefoundry.visitor_registry.visitors import TypeInferVisitor +from tilefoundry.visitor_registry.typeinfer import TypeInferVisitor _DIAGNOSTICS = Path(__file__).parents[1] / "fixtures" / "diagnostics" From de334e5fc70fa7fd5e4b8a515fe4b11b11904f63 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Tue, 22 Sep 2026 01:45:07 +0800 Subject: [PATCH 22/25] test(analysis): cover persistent schedule invariants --- tests/analysis/test_analyze_at_a_size.py | 73 +++++++++++++++++++++++- 1 file changed, 71 insertions(+), 2 deletions(-) diff --git a/tests/analysis/test_analyze_at_a_size.py b/tests/analysis/test_analyze_at_a_size.py index 6df4a8c5..8423df05 100644 --- a/tests/analysis/test_analyze_at_a_size.py +++ b/tests/analysis/test_analyze_at_a_size.py @@ -11,8 +11,9 @@ from __future__ import annotations -from dataclasses import replace +from dataclasses import dataclass, replace +import isl import pytest from tests.fixtures.placed.gqa_decode import GqaOnline @@ -31,9 +32,10 @@ TrafficMetadata, analyze, ) +from tilefoundry.analysis.access import Access, AccessPrecision from tilefoundry.analysis.compute_cost import _local_duration_ns from tilefoundry.analysis.errors import AnalysisError -from tilefoundry.analysis.iteration_scope import build_scopes, walk_scopes +from tilefoundry.analysis.iteration_scope import IterationScope, build_scopes, walk_scopes from tilefoundry.ir.core import Call, describe_expr, get_metadata from tilefoundry.ir.hir.function import Function from tilefoundry.ir.hir.loop_region import LoopRegion @@ -42,6 +44,7 @@ residual_dims, variant_for, ) +from tilefoundry.ir.hir.tensor.insert_slice import InsertSlice from tilefoundry.ir.types.shard import ( Topology, ) @@ -53,6 +56,16 @@ FAMILIES = ("compute-cost", "memory", "roofline", "performance") CASES = placed_cases() INVENTORY = [pytest.param(case, id=case.id) for case in CASES] + + +@dataclass(frozen=True) +class _PersistentScheduleExpectation: + loop_trips: tuple[tuple[str, int], ...] + store_loop: str + store_precision: AccessPrecision + compared_units: tuple[tuple[int, ...], tuple[int, ...]] + + EXPECTED_MEMORY_PEAKS = { "derived_prefill.DerivedPrefill.prefill[prefill_n=64,topology_only=128]": { "gmem": 288, @@ -209,8 +222,61 @@ "tp_all_to_all.TransposeShard.transpose_shard[static]": {"gmem": 256}, "weighted_twin.Weighted.scaled[static]": {"gmem": 1_348, "rmem": 4}, } +EXPECTED_PERSISTENT_SCHEDULES = { + "persistent_gemm_flat.PersistentGemmFlat.gemm[static]": _PersistentScheduleExpectation( + loop_trips=(("t", 4), ("ki", 4)), + store_loop="t", + store_precision=AccessPrecision.WIDENED, + compared_units=((0,), (1,)), + ), + "persistent_gemm_tiled.PersistentGemmTiled.gemm[static]": _PersistentScheduleExpectation( + loop_trips=(("mi", 2), ("ni", 2), ("ki", 4)), + store_loop="ni", + store_precision=AccessPrecision.EXACT, + compared_units=((0, 0), (1, 0)), + ), +} assert set(EXPECTED_MEMORY_PEAKS) == {case.id for case in CASES} +assert set(EXPECTED_PERSISTENT_SCHEDULES) <= {case.id for case in CASES} + + +def _loop_scopes(result: AnalysisResult) -> dict[str, IterationScope]: + return { + scope.owner.induction_var.name: scope + for scope in walk_scopes(build_scopes(result.module, result.function)) + if isinstance(scope.owner, LoopRegion) + } + + +def _insert_slice_output(scope: IterationScope) -> Access: + for call, accesses in scope.outputs.get("narrow", {}).values(): + if isinstance(call.target, InsertSlice): + assert len(accesses) == 1 + return accesses[0] + raise AssertionError("loop has no InsertSlice output") + + +def _at_unit(image: isl.set, coordinates: tuple[int, ...]) -> isl.set: + assert image.dim(isl.dim_type.PARAM) == len(coordinates) + for axis, coordinate in enumerate(coordinates): + image = image.fix_si(isl.dim_type.PARAM, axis, coordinate) + return image + + +def _assert_persistent_schedule( + result: AnalysisResult, + expected: _PersistentScheduleExpectation, +) -> None: + scopes = _loop_scopes(result) + for name, trips in expected.loop_trips: + assert scopes[name].trips() == trips + + store = _insert_slice_output(scopes[expected.store_loop]) + assert store.precision is expected.store_precision + written = store.relation.range() + first, second = (_at_unit(written, unit) for unit in expected.compared_units) + assert first.is_disjoint(second) def _aimed(): @@ -418,6 +484,9 @@ def test_every_concrete_program_predicts_coherently(case: ConcreteCase) -> None: assert placement is not None observed = {item.level: item.peak_bytes for item in placement.footprint} assert observed == EXPECTED_MEMORY_PEAKS[case.id] + expected_schedule = EXPECTED_PERSISTENT_SCHEDULES.get(case.id) + if expected_schedule is not None: + _assert_persistent_schedule(result, expected_schedule) @pytest.mark.parametrize("family", FAMILIES) From f18b3f4b31e75ee716cb546e39dce0b0d3efee77 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Tue, 22 Sep 2026 11:01:23 +0800 Subject: [PATCH 23/25] refactor(ir): move isl interop out of types --- docs/spec/code-organization.md | 1 + docs/spec/types.md | 7 ++++--- src/tilefoundry/analysis/access.py | 2 +- src/tilefoundry/analysis/allocation.py | 2 +- src/tilefoundry/analysis/footprint.py | 2 +- src/tilefoundry/analysis/loop_domain.py | 2 +- src/tilefoundry/ir/core/op.py | 2 +- src/tilefoundry/ir/hir/nn/conv2d.py | 2 +- src/tilefoundry/ir/hir/nn/rope.py | 2 +- src/tilefoundry/ir/hir/tensor/concat.py | 2 +- src/tilefoundry/ir/hir/tensor/slice.py | 2 +- .../ir/{types/dim_isl.py => isl_interop.py} | 15 ++++++++------- src/tilefoundry/ir/types/dim.py | 3 ++- src/tilefoundry/ir/types/substitute.py | 2 +- src/tilefoundry/parser/ast_pattern.py | 2 +- .../visitor_registry/access_relation.py | 2 +- src/tilefoundry/visitor_registry/isl_utility.py | 4 ++-- tests/analysis/test_analysis_invariants.py | 2 +- tests/analysis/test_isl_utility.py | 16 ++++++++-------- tests/ir/test_simplify_dim.py | 4 ++-- tests/ops/ir/test_arange.py | 2 +- tests/ops/ir/test_slice.py | 2 +- 22 files changed, 42 insertions(+), 38 deletions(-) rename src/tilefoundry/ir/{types/dim_isl.py => isl_interop.py} (97%) diff --git a/docs/spec/code-organization.md b/docs/spec/code-organization.md index 37072bf2..4bf90aff 100644 --- a/docs/spec/code-organization.md +++ b/docs/spec/code-organization.md @@ -22,6 +22,7 @@ truth for the directory's structure and invariants. | `ir/types/shard/` | [shard](./shard.md) | Shard / layout sublayer: `Topology` / `Mesh` / `Layout` / `ComposedLayout` / `ShardLayout` / `ShardAttr` (`Split` / `Broadcast` / `Dynamic` / `Partial`). The physical nesting reflects the spec's "sublayer" relationship. | | `ir/constraints/` | [parser](./parser.md) | Authored `where(layout=..., mesh=..., storage=...)` constraint records: the shared base plus layout, mesh, and storage constraints, attached by the parser and read back by the Python printer. | | `ir/visitor.py` | [visitor-mutator](./visitor-mutator.md) | `ExprFunctor` / `ExprVisitor` / `ExprWalker` / `ExprCollector` / `ExprCloner` / `BindingSubstitutionCloner` / `StmtVisitor` / `StmtMutator` / `StmtExprMutator`, plus `collect_exprs`, value-operand/function-value queries, and the canonical `PrimFunction` walk and rewrite entries. | +| `ir/isl_interop.py` | [types](./types.md) | Interoperation between dimension and shape IR values and isl: expression rendering and decoding, normalization, value ranges, and shape-domain construction. Pure isl operations remain in `utils/isl_utils.py`. | | `ir/hir/` | [hir](./hir.md) | HIR Op layer; one subdirectory per category (`math/` / `tensor/` / `nn/` / `shape/` / `sharding/`). One real Op per `.py` ([§2](#2-file-naming-and-content-rules) rule 1); surface-alias schemas have no per-name file and live in each category's `aliases.py` ([§2](#2-file-naming-and-content-rules) rule 5). | | `ir/tir/` | [tir](./tir.md) | TIR layer: `stmt.py` re-exports the `Stmt` base from `ir/core/stmt.py`; `stmts.py` hosts the general TIR `Stmt` subclasses (`LetStmt` / `Evaluate` / `Sequential` / `MeshScope` / …); `prim_function.py`; effect Ops and TIR-owned Expr Ops by category (`memory/` / `nn/` / …); `launch.py` owns `Launch` and its authored launch-attribute descriptors; `arith.py` / `reduce.py` for tag-dispatched `Binary` / `Unary` / `Reduce`; `intrinsic.py` for the `@intrinsic` decorator. Target-specific nodes nest under `ir/tir///` (e.g. `ir/tir/cuda/nn/mma.py`) per [§2](#2-file-naming-and-content-rules) Rule 1c. | | `parser/` | [parser](./parser.md) | DSL → IR parsing: `base.py` (shared visitor base + dispatch), `hir_parser.py` (`@func` body), `tir_parser.py` (`@prim_func` body), layout sugar / range-slice / dispatch modules. **Not under `ir/`** — the parser is a producer of IR, not an IR sublayer. | diff --git a/docs/spec/types.md b/docs/spec/types.md index 5fbdb2e1..6b7d8c13 100644 --- a/docs/spec/types.md +++ b/docs/spec/types.md @@ -532,9 +532,10 @@ def ceildiv(a, b) -> Expr: - `ceildiv(a, b)` MUST compose the existing add, subtract, and floor-divide operations; it does not introduce a distinct Op. - `ir.types.dim` MUST own dimension IR definitions, construction, and - structural predicates without depending on isl. `ir.types.dim_isl` MUST own - conversion between dimension IR and isl, affine normalization, shape-domain - construction, and conservative value-range queries. + structural predicates without depending on isl. `ir.isl_interop` MUST own + conversion between dimension and shape IR values and isl, affine + normalization, shape-domain construction, and conservative value-range + queries. - `dim_to_isl_expr` MUST render one dimension expression while registering its leaf parameters; `isl_to_dim` MUST decode an isl affine expression using that parameter map. `shape_to_isl_domain` MUST return one shape's iteration diff --git a/src/tilefoundry/analysis/access.py b/src/tilefoundry/analysis/access.py index 959bfe85..9bb7eb85 100644 --- a/src/tilefoundry/analysis/access.py +++ b/src/tilefoundry/analysis/access.py @@ -12,8 +12,8 @@ from tilefoundry.ir.hir.loop_region import LoopRegion from tilefoundry.ir.hir.tensor.reshape import Reshape from tilefoundry.ir.hir.tensor.slice import Slice +from tilefoundry.ir.isl_interop import index_set from tilefoundry.ir.types import TensorType -from tilefoundry.ir.types.dim_isl import index_set from tilefoundry.ir.types.shape_helpers import static_dim_value from tilefoundry.ir.types.utils import local_type_of from tilefoundry.utils.isl_utils import has_unbounded_param diff --git a/src/tilefoundry/analysis/allocation.py b/src/tilefoundry/analysis/allocation.py index 9251ce6a..8d4484a9 100644 --- a/src/tilefoundry/analysis/allocation.py +++ b/src/tilefoundry/analysis/allocation.py @@ -15,8 +15,8 @@ from tilefoundry.ir.hir.tensor.insert_slice import InsertSlice from tilefoundry.ir.hir.tensor.reshape import Reshape from tilefoundry.ir.hir.tensor.slice import Slice +from tilefoundry.ir.isl_interop import index_set from tilefoundry.ir.types import TensorType -from tilefoundry.ir.types.dim_isl import index_set from tilefoundry.ir.types.utils import local_type_of from tilefoundry.ir.visitor import ExprVisitor from tilefoundry.utils.isl_utils import equates diff --git a/src/tilefoundry/analysis/footprint.py b/src/tilefoundry/analysis/footprint.py index 2a3ea9e2..c874bd88 100644 --- a/src/tilefoundry/analysis/footprint.py +++ b/src/tilefoundry/analysis/footprint.py @@ -6,8 +6,8 @@ import isl +from tilefoundry.ir.isl_interop import index_set from tilefoundry.ir.types import TensorType -from tilefoundry.ir.types.dim_isl import index_set from tilefoundry.utils.isl_utils import cardinality from .affine import LoopAffineTerm diff --git a/src/tilefoundry/analysis/loop_domain.py b/src/tilefoundry/analysis/loop_domain.py index dab3d0fd..8985b4fc 100644 --- a/src/tilefoundry/analysis/loop_domain.py +++ b/src/tilefoundry/analysis/loop_domain.py @@ -7,7 +7,7 @@ from tilefoundry.ir.core import value_label from tilefoundry.ir.hir.function import Function from tilefoundry.ir.hir.loop_region import LoopRegion -from tilefoundry.ir.types.dim_isl import dim_to_isl_expr +from tilefoundry.ir.isl_interop import dim_to_isl_expr from tilefoundry.ir.types.shape_helpers import static_dim_value from .errors import AnalysisError diff --git a/src/tilefoundry/ir/core/op.py b/src/tilefoundry/ir/core/op.py index 84b895cc..92ce17fc 100644 --- a/src/tilefoundry/ir/core/op.py +++ b/src/tilefoundry/ir/core/op.py @@ -37,7 +37,7 @@ def _normalize_attr(name: str, value: Any) -> Any: """ if name == "storage": return resolve_storage(value) - from tilefoundry.ir.types.dim_isl import normalize_dim_entries # noqa: PLC0415 + from tilefoundry.ir.isl_interop import normalize_dim_entries # noqa: PLC0415 return normalize_dim_entries(value) diff --git a/src/tilefoundry/ir/hir/nn/conv2d.py b/src/tilefoundry/ir/hir/nn/conv2d.py index 66d4f272..b9526e73 100644 --- a/src/tilefoundry/ir/hir/nn/conv2d.py +++ b/src/tilefoundry/ir/hir/nn/conv2d.py @@ -11,9 +11,9 @@ from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._shard_checks import check_multilinear_partials +from tilefoundry.ir.isl_interop import normalize_dim 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 Split, shard_layout_of, split_target_axes diff --git a/src/tilefoundry/ir/hir/nn/rope.py b/src/tilefoundry/ir/hir/nn/rope.py index 75b36c76..484f4e09 100644 --- a/src/tilefoundry/ir/hir/nn/rope.py +++ b/src/tilefoundry/ir/hir/nn/rope.py @@ -21,8 +21,8 @@ from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op from tilefoundry.ir.hir._shard_checks import check_multilinear_partials, reject_partials +from tilefoundry.ir.isl_interop import index_set from tilefoundry.ir.types import TupleType -from tilefoundry.ir.types.dim_isl import index_set from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( AccessRelations, diff --git a/src/tilefoundry/ir/hir/tensor/concat.py b/src/tilefoundry/ir/hir/tensor/concat.py index 064c9e01..ab4695e6 100644 --- a/src/tilefoundry/ir/hir/tensor/concat.py +++ b/src/tilefoundry/ir/hir/tensor/concat.py @@ -17,9 +17,9 @@ reject_dynamic_shards, require_uniform_partial_slices, ) +from tilefoundry.ir.isl_interop import normalize_dim_entries from tilefoundry.ir.types import TensorType from tilefoundry.ir.types.dim import DimAdd, simplify_dim -from tilefoundry.ir.types.dim_isl import normalize_dim_entries from tilefoundry.ir.types.shard import ( Layout, Split, diff --git a/src/tilefoundry/ir/hir/tensor/slice.py b/src/tilefoundry/ir/hir/tensor/slice.py index 8d6ecddb..7220b0de 100644 --- a/src/tilefoundry/ir/hir/tensor/slice.py +++ b/src/tilefoundry/ir/hir/tensor/slice.py @@ -10,6 +10,7 @@ from tilefoundry.ir.core.param_def import ParamDef from tilefoundry.ir.core.pattern import Tensor from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.isl_interop import dim_range from tilefoundry.ir.types import TensorType from tilefoundry.ir.types.dim import ( DimAdd, @@ -18,7 +19,6 @@ DimSub, simplify_dim, ) -from tilefoundry.ir.types.dim_isl import dim_range from tilefoundry.ir.types.shape_helpers import i64_const from tilefoundry.ir.types.shard import ( ComposedLayout, diff --git a/src/tilefoundry/ir/types/dim_isl.py b/src/tilefoundry/ir/isl_interop.py similarity index 97% rename from src/tilefoundry/ir/types/dim_isl.py rename to src/tilefoundry/ir/isl_interop.py index 106b0a74..1a9c0388 100644 --- a/src/tilefoundry/ir/types/dim_isl.py +++ b/src/tilefoundry/ir/isl_interop.py @@ -1,8 +1,9 @@ -"""Conversion between dimension IR and isl, plus affine range queries. +"""Interoperation between dimension and shape IR values and isl. -Dimension definitions and construction stay in :mod:`dim`; this module owns -the dim-to-isl rendering, isl-to-dim decoding, shape-domain construction, -affine normalization, and conservative value-range calculation. +Dimension definitions and construction stay in :mod:`tilefoundry.ir.types.dim`; +pure isl operations stay in :mod:`tilefoundry.utils.isl_utils`. This module owns +the boundary between those layers: rendering, decoding, normalization, value +ranges, and shape domains. """ from __future__ import annotations @@ -13,7 +14,7 @@ from tilefoundry.ir.core.kinds import BinaryKind from tilefoundry.ir.core.metadata import RangeMetadata, get_metadata -from .dim import ( +from .types.dim import ( _DIM_OP_TYPES, DimAdd, DimFloorDiv, @@ -24,8 +25,8 @@ DimSub, DimVar, ) -from .dtype import IntegerDType -from .tensor_type import TensorType +from .types.dtype import IntegerDType +from .types.tensor_type import TensorType _INTEGER_BINARY_DIM_OP = { BinaryKind.ADD: DimAdd, diff --git a/src/tilefoundry/ir/types/dim.py b/src/tilefoundry/ir/types/dim.py index d8092484..207d452d 100644 --- a/src/tilefoundry/ir/types/dim.py +++ b/src/tilefoundry/ir/types/dim.py @@ -1,7 +1,8 @@ """Dimension IR definitions, construction, and structural predicates. This module does not depend on isl. Conversion to and from isl, affine -normalization, and value-range queries belong to :mod:`dim_isl`. +normalization, and value-range queries belong to +:mod:`tilefoundry.ir.isl_interop`. """ from __future__ import annotations diff --git a/src/tilefoundry/ir/types/substitute.py b/src/tilefoundry/ir/types/substitute.py index 62d1afb2..595a6fc2 100644 --- a/src/tilefoundry/ir/types/substitute.py +++ b/src/tilefoundry/ir/types/substitute.py @@ -11,9 +11,9 @@ from collections.abc import Mapping from tilefoundry.ir.core.expr import Call, Constant +from tilefoundry.ir.isl_interop import normalize_dim from .dim import _DIM_OP_TYPES, DimVar, simplify_dim -from .dim_isl import normalize_dim from .tensor_type import TensorType, TupleType, Type diff --git a/src/tilefoundry/parser/ast_pattern.py b/src/tilefoundry/parser/ast_pattern.py index 4b5161e7..08715ff2 100644 --- a/src/tilefoundry/parser/ast_pattern.py +++ b/src/tilefoundry/parser/ast_pattern.py @@ -48,6 +48,7 @@ from tilefoundry.ir.hir.tensor.reshape import Reshape from tilefoundry.ir.hir.tensor.slice import Slice, slice_size from tilefoundry.ir.hir.tensor.tuple_get_item import TupleGetItem +from tilefoundry.ir.isl_interop import normalize_dim from tilefoundry.ir.tir.prim_function import PrimFunction from tilefoundry.ir.tir.shape import ShapeOf from tilefoundry.ir.tir.stmts import ( @@ -73,7 +74,6 @@ dim_expr, simplify_dim, ) -from tilefoundry.ir.types.dim_isl import normalize_dim from tilefoundry.ir.types.shard import ( Broadcast, Layout, diff --git a/src/tilefoundry/visitor_registry/access_relation.py b/src/tilefoundry/visitor_registry/access_relation.py index dfff412d..aac1c5b0 100644 --- a/src/tilefoundry/visitor_registry/access_relation.py +++ b/src/tilefoundry/visitor_registry/access_relation.py @@ -15,8 +15,8 @@ import isl from tilefoundry.ir.hir._helpers import is_one +from tilefoundry.ir.isl_interop import index_set, isl_to_dim, shape_to_isl_domain from tilefoundry.ir.types import TensorType, TupleType, Type, tensor_bytes -from tilefoundry.ir.types.dim_isl import index_set, isl_to_dim, shape_to_isl_domain from tilefoundry.ir.types.shape_helpers import static_dim_value from tilefoundry.ir.types.shard.shard_layout import layout_axis_to_tensor_axis from tilefoundry.utils.isl_utils import cardinality diff --git a/src/tilefoundry/visitor_registry/isl_utility.py b/src/tilefoundry/visitor_registry/isl_utility.py index fc471a64..c7191386 100644 --- a/src/tilefoundry/visitor_registry/isl_utility.py +++ b/src/tilefoundry/visitor_registry/isl_utility.py @@ -1,5 +1,5 @@ -"""Compatibility exports for the dim <-> isl bridge owned by IR types.""" +"""Compatibility exports for the dim <-> isl bridge owned by IR.""" -from tilefoundry.ir.types.dim_isl import dim_range, isl_to_dim, shape_to_isl_domain +from tilefoundry.ir.isl_interop import dim_range, isl_to_dim, shape_to_isl_domain __all__ = ["dim_range", "isl_to_dim", "shape_to_isl_domain"] diff --git a/tests/analysis/test_analysis_invariants.py b/tests/analysis/test_analysis_invariants.py index 02c846ca..3c4e62fd 100644 --- a/tests/analysis/test_analysis_invariants.py +++ b/tests/analysis/test_analysis_invariants.py @@ -45,6 +45,7 @@ from tilefoundry.ir.core.pattern import Tensor as TensorPattern from tilefoundry.ir.hir.tensor.insert_slice import InsertSlice from tilefoundry.ir.hir.tensor.slice import Slice as SliceOp +from tilefoundry.ir.isl_interop import index_set from tilefoundry.ir.types import ( DType, TensorType, @@ -53,7 +54,6 @@ make_tensor_type, tensor_bytes, ) -from tilefoundry.ir.types.dim_isl import index_set from tilefoundry.ir.types.shard import Topology, make_mesh from tilefoundry.ir.types.shard.shard_layout import Split as ShardSplit from tilefoundry.ir.types.storage import StorageKind diff --git a/tests/analysis/test_isl_utility.py b/tests/analysis/test_isl_utility.py index a14a962c..1cd073cd 100644 --- a/tests/analysis/test_isl_utility.py +++ b/tests/analysis/test_isl_utility.py @@ -1,4 +1,4 @@ -"""dim_isl — range queries and conversion between dimensions and isl.""" +"""isl_interop — range queries and conversion between IR dimensions and isl.""" from __future__ import annotations @@ -10,6 +10,13 @@ from tilefoundry.ir.core.expr import Call, Var from tilefoundry.ir.hir.loop_region import LoopRegion from tilefoundry.ir.hir.sharding.mesh_coord import MeshCoord +from tilefoundry.ir.isl_interop import ( + dim_range, + index_set, + isl_to_dim, + normalize_dim, + shape_to_isl_domain, +) from tilefoundry.ir.types import TensorType from tilefoundry.ir.types.dim import ( DimAdd, @@ -22,13 +29,6 @@ DimVar, simplify_dim, ) -from tilefoundry.ir.types.dim_isl import ( - dim_range, - index_set, - isl_to_dim, - normalize_dim, - shape_to_isl_domain, -) from tilefoundry.utils.isl_utils import cardinality P = DimVar("P", 2048, 1_048_577) diff --git a/tests/ir/test_simplify_dim.py b/tests/ir/test_simplify_dim.py index afa55943..939aee0a 100644 --- a/tests/ir/test_simplify_dim.py +++ b/tests/ir/test_simplify_dim.py @@ -6,7 +6,7 @@ import pytest -import tilefoundry.ir.types.dim_isl as dim_isl +import tilefoundry.ir.isl_interop as isl_interop import tilefoundry.ir.types.substitute as dim_substitute from tilefoundry.ir.core import Tuple, TypeInferContext from tilefoundry.ir.core.expr import Call, Constant, Var @@ -265,7 +265,7 @@ def test_static_op_attributes_do_not_enter_dim_normalization( def fail_if_called(_): raise AssertionError("static attributes must not enter isl normalization") - monkeypatch.setattr(dim_isl, "normalize_dim", fail_if_called) + monkeypatch.setattr(isl_interop, "normalize_dim", fail_if_called) assert Reshape(new_shape=(8, 16)).new_shape == (8, 16) diff --git a/tests/ops/ir/test_arange.py b/tests/ops/ir/test_arange.py index f3d77c4f..4ad8c6dc 100644 --- a/tests/ops/ir/test_arange.py +++ b/tests/ops/ir/test_arange.py @@ -17,9 +17,9 @@ from tilefoundry.ir.hir.sharding.mesh_coord import MeshCoord from tilefoundry.ir.hir.specialize import residual_dims, specialize_concretely from tilefoundry.ir.hir.tensor.arange import Arange +from tilefoundry.ir.isl_interop import normalize_dim 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.shard import Layout, Mesh, Topology, composed from tilefoundry.ir.types.storage import StorageKind from tilefoundry.visitor_registry.contexts import TrafficBytes, TypeInferContext diff --git a/tests/ops/ir/test_slice.py b/tests/ops/ir/test_slice.py index fced03e9..58a8f0e8 100644 --- a/tests/ops/ir/test_slice.py +++ b/tests/ops/ir/test_slice.py @@ -16,9 +16,9 @@ from tilefoundry.ir.core.kinds import BinaryKind from tilefoundry.ir.hir.math.binary import Binary from tilefoundry.ir.hir.tensor.slice import Slice, slice_size +from tilefoundry.ir.isl_interop import normalize_dim 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 From 1942765deab686e3319a3242eb7f12b8a8c4ac1c Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Tue, 22 Sep 2026 11:20:13 +0800 Subject: [PATCH 24/25] refactor(analysis): clarify loop term ownership --- docs/spec/analysis.md | 4 ++ docs/spec/code-organization.md | 4 +- src/tilefoundry/analysis/access.py | 50 ++++++++++++-- src/tilefoundry/analysis/footprint.py | 68 ++----------------- .../analysis/{affine.py => loop_terms.py} | 22 +++--- 5 files changed, 63 insertions(+), 85 deletions(-) rename src/tilefoundry/analysis/{affine.py => loop_terms.py} (89%) diff --git a/docs/spec/analysis.md b/docs/spec/analysis.md index 902d0933..98724821 100644 --- a/docs/spec/analysis.md +++ b/docs/spec/analysis.md @@ -1222,6 +1222,10 @@ call site, source expressions shared by identity remain one shared expression in the clone; sharing never aliases the independently cloned body of another call site. +`analysis.loop_terms` resolves HIR values to `LoopTerm` without depending on +isl. `analysis.access` turns those terms into isl constraints and owns access +widening; the two modules do not define a second affine graph representation. + - constraints: - A loop `start` or `extent` MAY be unit-dependent. Every runtime leaf in one MUST carry a half-open value range, and `IterationScope.domain` MUST keep the diff --git a/docs/spec/code-organization.md b/docs/spec/code-organization.md index 4bf90aff..5844ef96 100644 --- a/docs/spec/code-organization.md +++ b/docs/spec/code-organization.md @@ -39,8 +39,8 @@ truth for the directory's structure and invariants. | `analysis/iteration_scope.py` | [analysis](./analysis.md) | The shared `IterationScope` tree built once from normalized HIR; families query it instead of constructing parallel structure. | | `analysis/access.py` | [analysis](./analysis.md) | Access relations resolved against the authored iteration scopes. | | `analysis/loop_domain.py` | [analysis](./analysis.md) | isl iteration domains built from authored loop bounds. | -| `analysis/affine.py` | [analysis](./analysis.md) | The shared loop-affine term parser used by scope binding and authored-loop footprint binding, including constant loop strides and bounded invariant offsets. It does not introduce a second affine graph representation. | -| `analysis/footprint.py` | [analysis](./analysis.md) | Target-independent authored-loop access images, buffer-view folding, and deduplicated versus repeated byte readings. Requires no separate time map. | +| `analysis/loop_terms.py` | [analysis](./analysis.md) | IR-only resolution of constants, authored loop axes, strides, and bounded invariant offsets into `LoopTerm`; it has no isl dependency. | +| `analysis/footprint.py` | [analysis](./analysis.md) | Reserved for future target-independent authored-loop footprint policy; access-relation construction stays in `analysis/access.py`. | | `analysis/report.py` | [analysis](./analysis.md) | Structured analysis report data, including record-family registration, field serialization, and target-aware report-only projections. It depends only on analysis/core modules; inspection consumes it to produce text and source annotations. | | `analysis/check.py` | [analysis](./analysis.md) | The shared authored-program gate for analysis: authored-type re-derivation, authored validation, call-context validation, and checker-specific input checks. Established once per public call rather than per family. | | `analysis/facts.py` | [analysis](./analysis.md) | The narrow Facts aggregates the analysis families declare — the memory hierarchy graph, the throughput rates, and the parallel capacity. It is the record of how much hardware each measurement rests on, and names no backend; a Fact shared across consumer families belongs in `target/facts.py`. | diff --git a/src/tilefoundry/analysis/access.py b/src/tilefoundry/analysis/access.py index 9bb7eb85..1ad05fe1 100644 --- a/src/tilefoundry/analysis/access.py +++ b/src/tilefoundry/analysis/access.py @@ -16,7 +16,7 @@ from tilefoundry.ir.types import TensorType from tilefoundry.ir.types.shape_helpers import static_dim_value from tilefoundry.ir.types.utils import local_type_of -from tilefoundry.utils.isl_utils import has_unbounded_param +from tilefoundry.utils.isl_utils import cardinality, has_unbounded_param from tilefoundry.visitor_registry.access_relation import ( BoundaryRelation, relation_of, @@ -24,9 +24,8 @@ ) from tilefoundry.visitor_registry.contexts import TypeInferContext -from .affine import LoopAffineTerm, loop_affine_term from .errors import AnalysisError -from .footprint import _widest_allowed +from .loop_terms import LoopTerm, loop_affine_term class AccessPrecision(Enum): @@ -46,6 +45,42 @@ class Access: precision: AccessPrecision = AccessPrecision.EXACT +def widest_allowed(access: isl.map, name: str, held: object) -> LoopTerm | None: + """The value a parameter may take that reaches the most of its operand. + + A footprint is an upper bound, so a parameter nobody here can place takes + whichever end of its legal range touches more: where a window sits does not + change how much of it there is, but how long it is does. Both ends are the + Op's own contract, read off the relation rather than guessed. + """ + names = [ + access.get_dim_name(isl.dim_type.PARAM, index) + for index in range(access.dim(isl.dim_type.PARAM)) + ] + space = f"[{', '.join(names)}] -> " + probe = isl.set(f"{space}{{ [x] : x = {name} }}").intersect_params(access.params()) + ends = (probe.dim_min_val(0), probe.dim_max_val(0)) + if not all(end.is_int() for end in ends): + return None + box = index_set(tuple(held.shape)) if isinstance(held, TensorType) else None + if box is None or box.dim(isl.dim_type.SET) != access.dim(isl.dim_type.OUT): + least = ends[0].get_num_si() + return LoopTerm(None, 0, least, least) + best: tuple[int, int] | None = None + for value in sorted({end.get_num_si() for end in ends}): + reach = ( + access.intersect_params(isl.set(f"{space}{{ : {name} = {value} }}")) + .range() + .intersect(box) + ) + amount = cardinality(reach) + if amount is None: + return None + if best is None or amount > best[0]: + best = (amount, value) + return None if best is None else LoopTerm(None, 0, best[1], best[1]) + + def _parameter_term( value: object, relation: isl.map, @@ -54,23 +89,23 @@ def _parameter_term( held: object, *, narrow: bool, -) -> tuple[LoopAffineTerm | None, AccessPrecision]: +) -> tuple[LoopTerm | None, AccessPrecision]: number = static_dim_value(value) if number is not None: - return LoopAffineTerm(None, 0, number, number), AccessPrecision.EXACT + return LoopTerm(None, 0, number, number), AccessPrecision.EXACT try: term = loop_affine_term(value, loops, narrow=narrow) except (TypeError, ValueError, NotImplementedError): term = None if term is not None: return term, AccessPrecision.EXACT - return _widest_allowed(relation, name, held), AccessPrecision.WIDENED + return widest_allowed(relation, name, held), AccessPrecision.WIDENED def _constrain_parameter( relation: isl.map, param_index: int, - term: LoopAffineTerm, + term: LoopTerm, ) -> isl.map: local = isl.local_space.from_space(relation.get_space()) @@ -167,4 +202,5 @@ def resolve_access( "AccessPrecision", "eliminate_parameters", "resolve_access", + "widest_allowed", ] diff --git a/src/tilefoundry/analysis/footprint.py b/src/tilefoundry/analysis/footprint.py index c874bd88..15ec3a75 100644 --- a/src/tilefoundry/analysis/footprint.py +++ b/src/tilefoundry/analysis/footprint.py @@ -1,65 +1,5 @@ -"""Measure authored-loop buffer access without requiring a schedule.""" +"""Reserved for future target-independent authored-loop footprint analysis. -from __future__ import annotations - -import math - -import isl - -from tilefoundry.ir.isl_interop import index_set -from tilefoundry.ir.types import TensorType -from tilefoundry.utils.isl_utils import cardinality - -from .affine import LoopAffineTerm - - -class _Unavailable(Exception): - """An access that cannot be represented by this authored-loop model.""" - - -def _static_loop_bound(value: object) -> int: - if isinstance(value, int) and not isinstance(value, bool): - return value - raise _Unavailable - - -def _widest_allowed( - access: isl.map, name: str, held: object -) -> LoopAffineTerm | None: - """The value a parameter may take that reaches the most of its operand. - - A footprint is an upper bound, so a parameter nobody here can place takes - whichever end of its legal range touches more: where a window sits does not - change how much of it there is, but how long it is does. Both ends are the - Op's own contract, read off the relation rather than guessed. - """ - names = [ - access.get_dim_name(isl.dim_type.PARAM, index) - for index in range(access.dim(isl.dim_type.PARAM)) - ] - space = f"[{', '.join(names)}] -> " - probe = isl.set(f"{space}{{ [x] : x = {name} }}").intersect_params(access.params()) - ends = (probe.dim_min_val(0), probe.dim_max_val(0)) - if not all(end.is_int() for end in ends): - return None - box = index_set(tuple(held.shape)) if isinstance(held, TensorType) else None - if box is None or box.dim(isl.dim_type.SET) != access.dim(isl.dim_type.OUT): - least = ends[0].get_num_si() - return LoopAffineTerm(None, 0, least, least) - best: tuple[int, int] | None = None - for value in sorted({end.get_num_si() for end in ends}): - reach = access.intersect_params( - isl.set(f"{space}{{ : {name} = {value} }}") - ).range().intersect(box) - amount = cardinality(reach) - if amount is None: - return None - if best is None or amount > best[0]: - best = (amount, value) - return None if best is None else LoopAffineTerm(None, 0, best[1], best[1]) - - - - -def _packed_bytes(elements: int, bit_width: int) -> int: - return math.ceil(elements * bit_width / 8) +Access relation construction lives in :mod:`access`; a follow-up will add +footprint-specific policy here. +""" diff --git a/src/tilefoundry/analysis/affine.py b/src/tilefoundry/analysis/loop_terms.py similarity index 89% rename from src/tilefoundry/analysis/affine.py rename to src/tilefoundry/analysis/loop_terms.py index 1c6f17ec..b31932f8 100644 --- a/src/tilefoundry/analysis/affine.py +++ b/src/tilefoundry/analysis/loop_terms.py @@ -16,7 +16,7 @@ @dataclass(frozen=True) -class LoopAffineTerm: +class LoopTerm: """One authored loop coefficient plus a compile-time offset interval.""" loop_axis: int | None @@ -69,12 +69,10 @@ def static_range(expr: Expr, *, narrow: bool) -> tuple[int, int] | None: return None -def _loop_term( - value: Expr, loops: tuple[LoopRegion, ...], *, narrow: bool -) -> LoopAffineTerm | None: +def _loop_term(value: Expr, loops: tuple[LoopRegion, ...], *, narrow: bool) -> LoopTerm | None: for index, loop in enumerate(loops): if loop.induction_var is value: - return LoopAffineTerm(index, 1, 0, 0) + return LoopTerm(index, 1, 0, 0) if not isinstance(value, Call) or not isinstance(value.target, Binary): return None if value.target.kind is BinaryKind.ADD: @@ -85,7 +83,7 @@ def _loop_term( term = _loop_term(candidate, loops, narrow=narrow) bounds = static_range(invariant, narrow=narrow) if term is not None and bounds is not None: - return LoopAffineTerm( + return LoopTerm( term.loop_axis, term.stride, term.low + bounds[0], @@ -101,7 +99,7 @@ def _loop_term( if term is not None and bounds is not None and bounds[0] == bounds[1]: factor = bounds[0] offsets = (term.low * factor, term.high * factor) - return LoopAffineTerm( + return LoopTerm( term.loop_axis, term.stride * factor, min(offsets), @@ -112,14 +110,14 @@ def _loop_term( def loop_affine_term( value: Expr, loops: tuple[LoopRegion, ...], *, narrow: bool -) -> LoopAffineTerm | None: +) -> LoopTerm | None: """Resolve a constant, loop variable, affine offset, or constant stride.""" base, offset = window_base(value) if base is None: - return LoopAffineTerm(None, 0, offset, offset) + return LoopTerm(None, 0, offset, offset) term = _loop_term(base, loops, narrow=narrow) if term is not None: - return LoopAffineTerm( + return LoopTerm( term.loop_axis, term.stride, term.low + offset, @@ -128,7 +126,7 @@ def loop_affine_term( bounds = static_range(base, narrow=narrow) if bounds is None: return None - return LoopAffineTerm(None, 0, bounds[0] + offset, bounds[1] + offset) + return LoopTerm(None, 0, bounds[0] + offset, bounds[1] + offset) -__all__ = ["LoopAffineTerm", "loop_affine_term"] +__all__ = ["LoopTerm", "loop_affine_term"] From 713390c7764970ee2e31dc19e57e5be6c06b9f52 Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Tue, 22 Sep 2026 11:29:34 +0800 Subject: [PATCH 25/25] refactor(analysis): publish shared analysis helpers --- src/tilefoundry/analysis/api.py | 4 ++-- src/tilefoundry/analysis/check.py | 2 +- src/tilefoundry/analysis/compute_cost.py | 2 +- src/tilefoundry/analysis/performance.py | 4 ++-- src/tilefoundry/cli/analyze.py | 4 ++-- tests/analysis/test_analysis_families.py | 8 ++++---- tests/analysis/test_analyze_at_a_size.py | 4 ++-- 7 files changed, 14 insertions(+), 14 deletions(-) diff --git a/src/tilefoundry/analysis/api.py b/src/tilefoundry/analysis/api.py index 5282cd1a..c436232f 100644 --- a/src/tilefoundry/analysis/api.py +++ b/src/tilefoundry/analysis/api.py @@ -13,7 +13,7 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass -from tilefoundry.analysis.check import _resolve_program_geometry, check_program +from tilefoundry.analysis.check import check_program, resolve_program_geometry from tilefoundry.analysis.errors import AnalysisError from tilefoundry.analysis.iteration_scope import ScopeBuilder from tilefoundry.analysis.registry import Analyzer @@ -141,7 +141,7 @@ def analyze( roots = _roots(analysis) result_module = module try: - module, function = _resolve_program_geometry( + module, function = resolve_program_geometry( module, function, dims, diff --git a/src/tilefoundry/analysis/check.py b/src/tilefoundry/analysis/check.py index 565f7156..a3874ee2 100644 --- a/src/tilefoundry/analysis/check.py +++ b/src/tilefoundry/analysis/check.py @@ -162,7 +162,7 @@ def _program_dim_vars(module: Module, function: Function) -> dict[str, object]: return found -def _resolve_program_geometry( +def resolve_program_geometry( module: Module, function: Function, dims: Mapping[str, int] | None, diff --git a/src/tilefoundry/analysis/compute_cost.py b/src/tilefoundry/analysis/compute_cost.py index c8907c33..eeb32427 100644 --- a/src/tilefoundry/analysis/compute_cost.py +++ b/src/tilefoundry/analysis/compute_cost.py @@ -51,7 +51,7 @@ def _is_structural_occurrence( ) -def _local_duration_ns( +def local_duration_ns( cost: ComputeCostMetadata, facts: ThroughputFacts, services: PerformanceServiceFacts, diff --git a/src/tilefoundry/analysis/performance.py b/src/tilefoundry/analysis/performance.py index da0b96b9..ad1f1cfb 100644 --- a/src/tilefoundry/analysis/performance.py +++ b/src/tilefoundry/analysis/performance.py @@ -13,7 +13,7 @@ from tilefoundry.ir.types.shape_helpers import static_dim_value from tilefoundry.ir.visitor import ExprVisitor -from .compute_cost import _local_duration_ns +from .compute_cost import local_duration_ns from .errors import AnalysisError from .facts import ParallelCapacityFacts, PerformanceServiceFacts, ThroughputFacts from .iteration_scope import IterationScope @@ -63,7 +63,7 @@ def default_visit_leaf( raise AnalysisError(f"performance: missing compute/memory record for {expr!r}") if ctx.facts is None or ctx.services is None: raise AnalysisError("performance: visitor context is missing target facts") - duration = _local_duration_ns( + duration = local_duration_ns( cost, ctx.facts, ctx.services, diff --git a/src/tilefoundry/cli/analyze.py b/src/tilefoundry/cli/analyze.py index 144ac3b5..d964e5c0 100644 --- a/src/tilefoundry/cli/analyze.py +++ b/src/tilefoundry/cli/analyze.py @@ -13,7 +13,7 @@ from tilefoundry.analysis import analyze, check_program from tilefoundry.analysis.check import ( _program_dim_vars, - _resolve_program_geometry, + resolve_program_geometry, ) from tilefoundry.cli.source import load_authored_ir, suggested_extents from tilefoundry.inspection import PythonPrintOptions, as_script @@ -142,7 +142,7 @@ def run_authored_analysis( raise ValueError(f"analyze needs one EXTENT for every open dimension: {guidance}") if not analyses: try: - checked_module, checked = _resolve_program_geometry( + checked_module, checked = resolve_program_geometry( module, function, dims, diff --git a/tests/analysis/test_analysis_families.py b/tests/analysis/test_analysis_families.py index 700be7fa..9d023f9c 100644 --- a/tests/analysis/test_analysis_families.py +++ b/tests/analysis/test_analysis_families.py @@ -35,7 +35,7 @@ ) from tilefoundry.analysis.api import analyze from tilefoundry.analysis.compute_cost import ( - _local_duration_ns, + local_duration_ns, ) from tilefoundry.analysis.errors import AnalysisError from tilefoundry.analysis.memory import MemoryOptions @@ -429,10 +429,10 @@ def test_a_price_is_refused_where_the_machine_states_no_rate_to_pay_it_at() -> N match=r"^performance: selected topology level 'thread', but the target's " r"one-unit throughputs are stated for 'cta'$", ): - _local_duration_ns(work, throughput, services, level="thread") + local_duration_ns(work, throughput, services, level="thread") with pytest.raises(AnalysisError, match=r"unknown compute dtype 'f9e9m9'"): - _local_duration_ns( + local_duration_ns( replace( work, flops=Breakdown((*work.flops.kinds, ("f9e9m9", Spread(8, 8, (8,))))), @@ -461,7 +461,7 @@ def test_a_price_is_refused_where_the_machine_states_no_rate_to_pay_it_at() -> N AnalysisError, match=rf"no one-unit throughput for level '{throughput.bandwidth_level}' at 'cta'", ): - _local_duration_ns( + local_duration_ns( ComputeCostMetadata(), throughput, replace(services, unit_bandwidth=()), diff --git a/tests/analysis/test_analyze_at_a_size.py b/tests/analysis/test_analyze_at_a_size.py index 8423df05..1d19c498 100644 --- a/tests/analysis/test_analyze_at_a_size.py +++ b/tests/analysis/test_analyze_at_a_size.py @@ -33,7 +33,7 @@ analyze, ) from tilefoundry.analysis.access import Access, AccessPrecision -from tilefoundry.analysis.compute_cost import _local_duration_ns +from tilefoundry.analysis.compute_cost import local_duration_ns from tilefoundry.analysis.errors import AnalysisError from tilefoundry.analysis.iteration_scope import IterationScope, build_scopes, walk_scopes from tilefoundry.ir.core import Call, describe_expr, get_metadata @@ -325,7 +325,7 @@ def assert_performance_contract(result: AnalysisResult) -> None: continue cost = get_metadata(expr, ComputeCostMetadata) assert cost is not None - duration = _local_duration_ns( + duration = local_duration_ns( cost, throughput, services,