fix(dsl): make X @ mesh.axis mean one thing in reshard layouts and Tensor types

[zhen8838]

Why

• X @ mesh.axis means two different things in the two places an author writes it. In a reshard layout it keeps the tensor's rank; in a Tensor[...] type it splits the axis into two. So the two cannot meet:

tilefoundry: error: Binary: Binary ADD: cannot broadcast shapes (1, 16, 8192) and (1, 16, 128, 64)

• Both operands were written with the same sugar over the same mesh in the same body.

What

• tf.reshard(x, (1, 16, 8192 @ m.u), "smem") has type Tensor[(1, 16, 8192), ..., ((1, 16, 128 @ cta.u, 64), ...)] — logical rank 3, the split recorded in the layout.
• tf.zeros(Tensor[(1, 16, 8192 @ m.u), "f32", "smem"]) has logical shape (1, 16, 128, 64) — the split is in the shape.
• Consequence: an accumulator seeded with tf.zeros cannot be added to, or accumulated into by, anything produced by a resharded op, which is the ordinary way to open a reduction. docs/tutorial/authoring.md's Stage3_Fused works only because every operand there is tf.zeros/tf.full_like of the same seed, so the two forms never meet.

Contract

• One notation, one meaning. Whichever way it is resolved, reshard and Tensor[...] should agree on whether @ changes rank.

Risk

• If Tensor[...] is changed to match reshard, programs that index the split axes of a tf.zeros result break. If reshard is changed to match Tensor[...], most sharded programs break. A third option — reject the sugar in one of the two positions — is the smallest change and at least turns a broadcast error into a statement about the notation.
• Workaround today: open the accumulator with the first iteration's matmul instead of tf.zeros, which is what we ended up doing everywhere.

Reproducer

from __future__ import annotations
from tilefoundry import func, module
from tilefoundry.dsl import Mesh, Tensor, tf
from tilefoundry.dsl.tf import *  # noqa: F401, F403
from tilefoundry.ir.types.shard import Topology
from tilefoundry.target import CudaTarget

N = 128
_H200, _CTA = CudaTarget("nvidia.h200_sxm"), Topology("cta", N)


@module(entry="f", target=_H200, topologies=(_CTA,))
class M:
    @func
    def f(x: Tensor[(1, 16, 8192), "f32"]) -> Tensor[(1, 16, 8192), "f32"]:
        with Mesh(("cta",), layout=(N,), names=("u",)) as m:
            xs = tf.reshard(x, (1, 16, 8192 @ m.u), "smem")          # rank kept
            zs = tf.zeros(Tensor[(1, 16, 8192 @ m.u), "f32", "smem"])  # rank grown
            return xs + zs

tilefoundry analyze case.py:M /tmp/r.txt --compute-cost

tilefoundry 0.0.2.dev43+g45032b967 (45032b9), target nvidia.h200_sxm.

[zhen8838]

Not fixed in #123 — it wants a maintainer's call, and the PR's Risk section says why. Referenced there so it does not get lost.

[zhen8838]

Fixed by #150, landed on main as 18edc35.

A Tensor shape slot now keeps the extents as written, so both spellings agree:

tf.reshard(x, (1, 16, 8192 @ m.u), "smem")   shape=(1,16,8192)
Tensor[(1, 16, 8192 @ m.u), "f32", "smem"]   shape=(1,16,8192)

xs + zs parses, so the matmul-seeded accumulator workaround your Risk section describes can be undone where it was used. I have not undone it anywhere.

The cause was narrower than the issue's three options suggested. Both slots parse the sugar with the same PlacedLayoutPattern, whose contract -- stated by its own LayoutShapeRule -- is to produce a layout. That is the whole answer for a layout slot, which only asks where a value goes. A shape slot has no operand to read a shape from, so it recovered one from the layout, and a layout's shape is the divided one; the shape as written was computed while parsing and then dropped. So neither reshard nor Tensor[...] needed its meaning changed, and ShardLayout was not touched.

Two consequences worth knowing:

• A layout may now be stated once, not twice. Tensor[(8 @ m.b, 8), "f32", (8 @ m.b, 8)] is refused, where previously the layout slot silently replaced what the shape slot implied -- and could replace it with None.
• All 27 authored uses of shape-slot sugar still parse, 18 of them in the shipped nemotron example.