Skip to content

fix(parser): fold a body-local constant expression before the shard divisibility check #119

Description

@zhen8838

Why

  • A body-local name for a constant expression is not folded. width = A + B with A = B = 4096 stays an IR Binary Call, so 8192 @ m.u is not recognised as divisible by a 128-position mesh and the program is refused.
  • The refusal prints an internal repr rather than the expression the author wrote:
tilefoundry: error: shard_layout_local_shape: layout dim 2 (Call(type=TensorType(shape=(),
dtype=IntegerDType(name='i64', bit_width=64, signed=True), layout=None,
storage=<StorageKind.UMAT: 6>), target=Binary(kind=<BinaryKind.ADD: 'add'>),
args=(Constant(...value=4096), Constant(...value=4096)))) and mesh axis 0 extent 128
do not have a decidable divisibility relation; bind symbolic dimensions before local projection

What

  • The same arithmetic written at module level (W = A + B outside the class) is folded and accepted.
  • The same arithmetic written inline in the layout tuple ((1, 16, (A + B) @ m.u)) is also accepted.
  • Only naming it inside the body fails, so the constant reaches the shard check unfolded.

Contract

  • Two spellings of one extent should be one extent. Today a body-local name changes whether a legal shard layout is legal.
  • The message tells an author to "bind symbolic dimensions" for a dimension that is not symbolic; 4096 + 4096 is fully static.

Risk

  • Two separable fixes: fold static integer Binary in the body (removes the refusal), and print the authored expression instead of the node repr (removes the misleading advice). The second is worth doing even if the first is not, because the same message is reachable from a genuinely open dimension.

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, A, B = 128, 4096, 4096
_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:
            width = A + B                                  # <-- refused
            return tf.reshard(x, (1, 16, width @ m.u), "smem")
tilefoundry analyze case.py:M /tmp/r.txt --compute-cost

Replacing the marked line with (1, 16, (A + B) @ m.u) inline, or naming W = A + B at module level, is accepted.

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

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions