fix(ir): preserve layouts and symbolic slice dimensions

docs/spec/hir.md

@@ -619,11 +619,12 @@ their input when it states one. An input with `layout=None` produces a view with
   the registered relation.
 - `Slice` is normalized as `Slice(x, starts, sizes=..., strides=...)`.
   `starts` is a tuple of rank-0 integer operands; `sizes` and `strides` are
-  `ShapeDim` attributes. Its result shape is exactly `sizes` and MUST NOT contain
-  an induction `Var`. A start MAY be dim arithmetic over an induction `Var` — a
-  window moved off that loop's window by a compile-time offset. That start is an
-  address computed where it is read, not a value some op produces, so a walk over
-  compute ops MUST leave it alone.
+  `ShapeDim` attributes stored in the same IR normal form as every other dim.
+  Its result shape is exactly the normalized `sizes` and MUST NOT contain an
+  induction `Var`. A start MAY be dim arithmetic over an induction `Var` — a window moved off that
+  loop's window by a compile-time offset. That start is an address computed where
+  it is read, not a value some op produces, so a walk over compute ops MUST leave
+  it alone.
 - A plain-layout `Slice` with static starts MUST produce a `ComposedLayout`: its
   offset is the source offset plus the starts multiplied by the source strides,
   and its outer layout carries the sliced shape and retained strides (multiplied
@@ -892,8 +893,16 @@ class Reshape(Op):
   - A plain C-order input reshapes to a C-order `Layout` over `new_shape`. An
     input with no assigned layout, or a non-contiguous plain input whose regroup
     cannot be expressed, has a `None` result layout.
-  - A fully-`Broadcast` `ShardLayout` input (every attr `Broadcast`, no genuine
-    sharding) reshapes to a plain (unsharded) output.
+  - A bare, fully-`Broadcast` `ShardLayout` input (every attr `Broadcast`, no
+    genuine sharding) carries that `ShardLayout` through `Reshape` when the
+    input layout positions can express `new_shape` by the view rules below.
+    When they cannot (including while either layout shape is symbolic), the
+    result layout is `None` and no error is raised because no genuine ownership
+    is discarded. This rule does not extend to a `ComposedLayout` whose outer
+    layout is a `ShardLayout`; that input follows the generic composed-layout
+    rule and may produce a `None` layout.
+  - A non-genuinely-sharded `UMAT` input reshaped to `()` remains the early
+    scalar case and produces a plain `UMAT` scalar.
   - A genuine `ShardLayout` input (at least one non-`Broadcast` attr) carries
     through `Reshape` when the reshape is expressible as a view over the
     input's layout positions (`layout.layout.shape`, [shard §7.1.1](./shard.md#711-layoutshape)):
@@ -913,8 +922,9 @@ class Reshape(Op):
     block spans a boundary deeper than one divide, or two or more
     `Split`-bound positions interacting across the same regroup — is not yet
     supported and MUST fail closed.
-  - A reshape not expressible by the above MUST fail closed rather than
-    fabricate a layout.
+  - A reshape of a genuine `ShardLayout` not expressible by the above MUST fail
+    closed rather than fabricate or discard a layout. An unexpressible bare,
+    fully-`Broadcast` `ShardLayout` instead follows the `None`-layout rule above.
 
 ##### Cast
 ```python

docs/spec/parser.md

@@ -457,6 +457,10 @@ compile-time-expr ::= number-literal | identifier | compile-time-expr '.' identi
   rejected ([hir §1.3](./hir.md#13-op)).
 - Evaluation MUST NOT call anything reached from a speculative position: a value
   that is not statically reachable is parsed as IR instead.
+- Dimension arithmetic has one canonical spelling after it enters IR, including
+  shape annotations and op attributes. Slice endpoints that contain the same
+  runtime scalar value MAY cancel to a static window size; an unrelated runtime
+  endpoint remains invalid under the ordinary `ShapeDim` rule.
 
 A **compile-time list** holds `Expr` elements and never reaches the IR:
 
@@ -482,8 +486,10 @@ For a tensor slice, a run-time rank-0 integer is permitted as an endpoint only
 when the resulting window size is a compile-time dimension. The canonical
 spelling is `start:start + K`; the parser MUST reject an unrelated stop
 endpoint because `Slice.sizes` is a static attribute. The slice stride remains
-compile-time. A tile window keeps its own length and MAY be **moved** by a
-compile-time offset instead ([§1.7](#17-for-i-in-tile--for-i-in-range-hir-only)).
+compile-time. An endpoint MAY also be a compile-time dimension, including the
+axis's own symbolic extent; its window size is symbolic accordingly. A tile
+window keeps its own length and MAY be **moved** by a compile-time offset instead
+([§1.7](#17-for-i-in-tile--for-i-in-range-hir-only)).
 
 ## 2. DSL namespace surface
 

docs/spec/types.md

@@ -491,10 +491,22 @@ def ceildiv(a, b) -> Expr:
   - `DimVar` identity MUST be canonical per `(name, lo, hi)`. Same-name
     dimensions in one function signature MUST agree on bounds.
   - Producers of arithmetic dimension calls MUST route construction through
-    `simplify_dim`.
-  - `simplify_dim` MUST fold two integer-valued constant operands, except
-    division or modulo by zero; otherwise it MUST retain the canonical call.
-    It MUST NOT perform algebraic identity folding.
+    `simplify_dim`. A `DimVar` and the `Call` produced by dimension arithmetic
+    both support continued `+`, `-`, `*`, `//`, and `%` construction, including
+    the reflected forms. Other `Call` values do not support this arithmetic.
+  - `simplify_dim` MUST only construct dimension arithmetic: it wraps raw integer
+    operands as integer `Constant` values and rejects boolean operands, but MUST
+    NOT fold constants or apply algebraic identities.
+  - Every dimension stored in an IR type or op attribute MUST use the one
+    `normalize_dim` isl affine normal form. This covers function signatures,
+    inferred types, layouts, topology and mesh entries, and op attributes; it
+    does not cover dimension expressions used as `Expr` operands, such as a
+    `Slice` start address. All-constant expressions normalize to plain Python
+    integers. Runtime scalar `Var` leaves MUST be represented by object identity,
+    not name, so repeated uses of one value can cancel without conflating
+    distinct same-named values. Normalization MUST NOT apply `DimVar` envelope
+    bounds. Expressions outside the affine subset or not decodable as one
+    `ShapeDim` MUST remain unchanged.
   - `is_dim_expr` MUST accept non-boolean integers, `DimVar`, integer-valued
     `Constant`, and recursively valid calls to the seven dimension arithmetic
     operations, and MUST reject other values.

src/tilefoundry/ir/core/expr.py

@@ -82,6 +82,44 @@ class Call(Expr):
     target: Op
     args: tuple[Expr, ...]
 
+    def _dim_binop(self, other, op_name: str, *, reverse: bool = False):
+        from tilefoundry.ir.types import dim  # noqa: PLC0415
+
+        if not dim.is_dim_op_call(self):
+            return NotImplemented
+        operands = (other, self) if reverse else (self, other)
+        return dim._dim_binop(getattr(dim, op_name), *operands)
+
+    def __add__(self, other):
+        return self._dim_binop(other, "DimAdd")
+
+    def __radd__(self, other):
+        return self._dim_binop(other, "DimAdd", reverse=True)
+
+    def __sub__(self, other):
+        return self._dim_binop(other, "DimSub")
+
+    def __rsub__(self, other):
+        return self._dim_binop(other, "DimSub", reverse=True)
+
+    def __mul__(self, other):
+        return self._dim_binop(other, "DimMul")
+
+    def __rmul__(self, other):
+        return self._dim_binop(other, "DimMul", reverse=True)
+
+    def __floordiv__(self, other):
+        return self._dim_binop(other, "DimFloorDiv")
+
+    def __rfloordiv__(self, other):
+        return self._dim_binop(other, "DimFloorDiv", reverse=True)
+
+    def __mod__(self, other):
+        return self._dim_binop(other, "DimMod")
+
+    def __rmod__(self, other):
+        return self._dim_binop(other, "DimMod", reverse=True)
+
 
 @dataclass(frozen=True)
 class Tuple(Expr):

src/tilefoundry/ir/core/module.py

@@ -16,6 +16,7 @@
 from tilefoundry.ir.hir.function import Function as HirFunction
 from tilefoundry.ir.tir.prim_function import PrimFunction
 from tilefoundry.ir.types.shard.mesh import Topology
+from tilefoundry.ir.types.substitute import canonicalize_topology_dims
 from tilefoundry.ir.types.tensor_type import TensorType
 from tilefoundry.target.base import Target, target_instance
 from tilefoundry.utils.spec_ref import spec_ref_render
@@ -245,6 +246,9 @@ def __post_init__(self) -> None:
                 f"method and a function/child module; names must be disjoint"
             )
         if self.topologies is not None:
+            topologies = tuple(canonicalize_topology_dims(t) for t in self.topologies)
+            if topologies != self.topologies:
+                object.__setattr__(self, "topologies", topologies)
             names = [t.name for t in self.topologies]
             dupes = sorted({n for n in names if names.count(n) > 1})
             if dupes:

src/tilefoundry/ir/core/op.py

@@ -37,7 +37,9 @@ def _normalize_attr(name: str, value: Any) -> Any:
     """
     if name == "storage":
         return resolve_storage(value)
-    return value
+    from tilefoundry.ir.types.dim_isl import normalize_dim_entries  # noqa: PLC0415
+
+    return normalize_dim_entries(value)
 
 
 @dataclass(frozen=True)

src/tilefoundry/ir/hir/function.py

@@ -10,7 +10,7 @@
 from tilefoundry.ir.hir.grid_region import GridRegionExpr
 from tilefoundry.ir.types import TensorType, Type, callable_type_for
 from tilefoundry.ir.types.dim import is_dim_expr
-from tilefoundry.ir.types.substitute import substitute_shape_dim
+from tilefoundry.ir.types.substitute import canonicalize_dims, substitute_shape_dim
 from tilefoundry.visitor_registry import register_typeinfer
 from tilefoundry.visitor_registry.contexts import FunctionScope, TypeInferContext
 
@@ -46,7 +46,12 @@ def build(
         variants: tuple["Function", ...] = (),
         converters: tuple[tuple[str, "Function"], ...] = (),
     ) -> "Function":
-        """Construct a Function with the canonical CallableType."""
+        """Construct a Function whose declarations and callable type are canonical."""
+        for param in params:
+            canonical = canonicalize_dims(param.type)
+            if canonical is not param.type:
+                object.__setattr__(param, "type", canonical)
+        return_type = canonicalize_dims(return_type)
         return cls(
             name=name,
             params=params,

src/tilefoundry/ir/hir/nn/conv2d.py

@@ -13,6 +13,7 @@
 from tilefoundry.ir.hir._shard_checks import check_multilinear_partials
 from tilefoundry.ir.types import TensorType
 from tilefoundry.ir.types.dim import DimAdd, DimFloorDiv, DimSub, simplify_dim
+from tilefoundry.ir.types.dim_isl import normalize_dim
 from tilefoundry.ir.types.shape_helpers import i64_const, static_dim_value
 from tilefoundry.ir.types.shard import Layout, try_c_order_strides
 from tilefoundry.ir.types.shard.shard_layout import shard_layout_of, split_target_axes
@@ -62,7 +63,7 @@ def _out_spatial(in_dim: Expr, k: int, s: int, p: int, d: int) -> Expr:
     sub_k = simplify_dim(DimSub, (add_pad, _i64(eff_k)))
     div_s = simplify_dim(DimFloorDiv, (sub_k, _i64(s)))
     plus_1 = simplify_dim(DimAdd, (div_s, _i64(1)))
-    return plus_1
+    return normalize_dim(plus_1)
 
 
 def _pair(call, ctx, name: str, values: tuple, *, positive: bool) -> tuple[int, int]:

src/tilefoundry/ir/hir/tensor/arange.py

@@ -12,6 +12,7 @@
 from tilefoundry.ir.core.register import register_op
 from tilefoundry.ir.types import DType, TensorType
 from tilefoundry.ir.types.dim import DimSub, ceildiv, is_dim_expr, simplify_dim
+from tilefoundry.ir.types.dim_isl import normalize_dim
 from tilefoundry.ir.types.shape_dim import ShapeDim
 from tilefoundry.ir.types.shape_helpers import static_dim_value
 from tilefoundry.ir.types.storage import StorageKind
@@ -46,7 +47,7 @@ def _(call: "Call", ctx: "TypeInferContext") -> TensorType:
     if op.dtype not in (DType.i32, DType.i64):
         ctx.error(call, f"dtype must be i32 or i64, got {op.dtype}")
 
-    difference = simplify_dim(DimSub, (op.end, op.start))
+    difference = normalize_dim(simplify_dim(DimSub, (op.end, op.start)))
     static_difference = static_dim_value(difference)
     if static_difference is not None and static_difference < 0:
         ctx.error(

src/tilefoundry/ir/hir/tensor/reshape.py

@@ -166,16 +166,13 @@ def _(call: "Call", ctx: "TypeInferContext") -> TensorType:
     new_layout = None
     if isinstance(x_ty.layout, ShardLayout):
         genuine = any(not isinstance(a, Broadcast) for a in x_ty.layout.attrs)
-        if not genuine:
-            new_layout = None
-        else:
-            new_layout = _carry_sharded_reshape(x_ty.layout, new_shape)
-            if new_layout is None:
-                ctx.error(
-                    call,
-                    "Reshape cannot express the sharded layout: new shape does "
-                    "not align with the input layout factorization",
-                )
+        new_layout = _carry_sharded_reshape(x_ty.layout, new_shape)
+        if new_layout is None and genuine:
+            ctx.error(
+                call,
+                "Reshape cannot express the sharded layout: new shape does "
+                "not align with the input layout factorization",
+            )
     else:
         source = x_ty.layout
         if isinstance(source, Layout):

src/tilefoundry/ir/hir/tensor/slice.py

@@ -52,27 +52,13 @@ def _i64(value: int) -> Constant:
 
 
 def slice_size(begin: Expr, end: Expr, stride: Expr) -> Expr:
-    """Return ``max(0, ceil((end - begin) / stride))`` as a dimension Expr.
-
-    Constant chains fold through ``simplify_dim``. A non-positive constant
-    stride returns zero explicitly because generic arithmetic folding does not
-    encode slice-domain semantics.
-    """
-    if isinstance(begin, Constant) and isinstance(end, Constant) and isinstance(stride, Constant):
-        b, e, s = int(begin.value), int(end.value), int(stride.value)
-        if s <= 0:
+    """Return ``ceil((end - begin) / stride)`` as a dimension expression."""
+    static = tuple(_constant_int(value) for value in (begin, end, stride))
+    if all(value is not None for value in static):
+        start, stop, step = static
+        if step > 0 and stop < start:
             return _i64(0)
-        n = max(0, (e - b + s - 1) // s)
-        return _i64(n)
-    if isinstance(end, Call) and isinstance(end.target, DimAdd) and end.args[0] is begin:
-        window = end.args[1]
-        if isinstance(stride, Constant) and stride.value == 1:
-            return window
-        bump = simplify_dim(
-            DimAdd,
-            (window, simplify_dim(DimSub, (stride, _i64(1)))),
-        )
-        return simplify_dim(DimFloorDiv, (bump, stride))
+
     diff = simplify_dim(DimSub, (end, begin))
 
     bump = simplify_dim(

src/tilefoundry/ir/types/dim.py

@@ -136,23 +136,11 @@ class DimMax(Op):
     b = ParamDef(kind="input")
 
 
-_DIM_FOLDERS: dict[type[Op], object] = {
-    DimAdd: lambda a, b: a + b,
-    DimSub: lambda a, b: a - b,
-    DimMul: lambda a, b: a * b,
-    DimFloorDiv: lambda a, b: a // b,
-    DimMod: lambda a, b: a % b,
-    DimMin: min,
-    DimMax: max,
-}
-
-
 def simplify_dim(op_cls: type[Op], args: tuple) -> Expr:
-    """Fold dimension arithmetic when every operand is an integer constant.
+    """Build dimension arithmetic without applying algebraic rules.
 
-    Raw integers canonicalize to i64 constants. Unsupported operations and
-    division or modulo by zero remain calls for later verification; algebraic
-    identities are not folded.
+    Raw integers become i64 constants and bool remains invalid. The dimension
+    is normalized only when it enters IR.
 
     See [types §4](docs/spec/types.md#4-dim--symbolic-shape-dimensions).
     """
@@ -173,21 +161,6 @@ def _wrap(v):
 
     canon_args = tuple(_wrap(a) for a in args)
 
-    fold = _DIM_FOLDERS.get(op_cls)
-    if (
-        fold is not None
-        and len(canon_args) == 2
-        and all(
-            isinstance(a, Constant) and isinstance(a.value, int) and not isinstance(a.value, bool)
-            for a in canon_args
-        )
-    ):
-        a_val = int(canon_args[0].value)
-        b_val = int(canon_args[1].value)
-        if op_cls in (DimFloorDiv, DimMod) and b_val == 0:
-            pass
-        else:
-            return Constant(type=ti64, value=fold(a_val, b_val))
     return Call(type=ti64, target=op_cls(), args=canon_args)
 
 
@@ -217,6 +190,21 @@ def is_dim_expr(value) -> bool:
     return False
 
 
+def dim_expr(value) -> Expr:
+    """Convert a dimension value to an expression without simplifying it."""
+    if isinstance(value, Expr):
+        return value
+    if isinstance(value, bool):
+        raise TypeError("dim_expr: bool is not a dimension")
+    if isinstance(value, int):
+        from .tensor_type import TensorType  # noqa: PLC0415
+
+        return Constant(type=TensorType.umat_scalar(), value=value)
+    if isinstance(value, DimVar):
+        return simplify_dim(DimAdd, (value, 0))
+    raise TypeError(f"dim_expr: expected int, DimVar, or Expr, got {type(value).__name__}")
+
+
 def is_dim_op_call(value) -> bool:
     """True iff *value* is a ``Call`` over one of this module's dim ops.
 
@@ -230,11 +218,7 @@ def is_dim_op_call(value) -> bool:
 
 
 def dim_min(a, b) -> Expr:
-    """Symbolic ``min`` dim expression, folded to a ``Constant`` when both operands are static.
-
-    Symbolic ``min(a, b)`` dim expression, folded to a ``Constant`` when both
-    operands are static.
-    """
+    """Build a symbolic ``min(a, b)`` dimension expression."""
     result = _dim_binop(DimMin, a, b)
     if result is NotImplemented:
         raise TypeError(
@@ -260,8 +244,7 @@ def ceildiv(a, b) -> Expr:
 
     Composes existing dim-arithmetic ops — there is no dedicated ceil-div
     op. Operands may be ``int`` (non-bool), ``DimVar`` or ``Expr``; the
-    result is the same ``ShapeDim`` form produced by ``simplify_dim`` and
-    folds to a ``Constant`` when both operands are static.
+    result is the same ``ShapeDim`` form produced by ``simplify_dim``.
     """
     num = simplify_dim(DimSub, (simplify_dim(DimAdd, (a, b)), 1))
     return simplify_dim(DimFloorDiv, (num, b))
@@ -279,6 +262,7 @@ def ceildiv(a, b) -> Expr:
     "DimMax",
     "simplify_dim",
     "is_dim_expr",
+    "dim_expr",
     "is_dim_op_call",
     "dim_min",
     "dim_max",