fix(analysis): count M once in MatMul per-unit flops

[zhen8838]

Why

• MatMul per-unit flops come out global / 2 — whatever the grid is — when the sharded axis of the output is its last one. The same gemm with the sharded axis leading reports global / 128, which is right for a 128-CTA mesh.
• performance is issued against that column, so two spellings of one gemm read 14x apart on predicted-ns while moving the same bytes and reporting the same ideal-ns.

LastAxis     flops bf16 global=25769803776 per-unit=12884901888  ratio   2.0   1820708 predicted-ns
StripMajor   flops bf16 global=25769803776 per-unit=  201326592  ratio 128.0    128676 predicted-ns

What

• visitor_registry/op_cost.py::_matmul reads shapes off the local types:

m, k, n = lhs.shape[-2], lhs.shape[-1], rhs.shape[-1]
batch = math.prod(output.shape[:-2])

• Sharding an axis replaces it with (units, per-unit). Sharded last, the new axis lands at -2: the local output is (1, 64, 1, 96), so batch = prod((1, 64)) = 64 and m = 64 too. M is counted twice and the units divisor never appears. Sharded leading, the new axis lands inside the batch region at extent 1 and the arithmetic is right.
• cast on the same value in the same program reports global / 128, so this is the shape arithmetic in _matmul, not the local projection.

Contract

• The per-unit column of compute-cost, and every predicted-ns derived from it, is currently a function of which axis a weight's layout puts the grid on. --topology cta is documented as changing flops_per_unit for a program that shards; a sharded program reporting global / 2 is reporting that it barely sharded.
• flops (global) and ideal-ns are unaffected.

Risk

• Fixing it changes reported flops_per_unit/predicted-ns for existing last-axis-sharded programs — in the direction of the grid divisor.
• Suggested direction: take m, k, n and the batch from operand/result tensor axes rather than local layout ranks. layout_axis_to_tensor_axis is already in that file and _conv2d beside it uses it for the same reason.

Reproducer

Two modules, identical arithmetic and traffic; the only difference is the two marked lines.

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

H, GRID, PT, BK, NOUT, NPER = 2048, 128, 64, 128, 12288, 96
S = DimVar("seq_len", 0, 8193)
_H200, _CTA = CudaTarget("nvidia.h200_sxm"), Topology("cta", GRID)


@module(entry="proj", target=_H200, topologies=(_CTA,))
class LastAxis:
    @func
    def proj(x: Tensor[(1, S, H), "bf16"],
             w: ConstTensor[(H, NOUT), "bf16"],                 # (K, N)
             out: Tensor[(1, S, NOUT), "bf16"]) -> Tensor[(1, S, NOUT), "bf16"]:
        with Mesh(("cta",), layout=(GRID,), names=("u",)) as m:
            acc = out
            for t in tile(S, PT):
                base = t + 0
                xs = tf.reshard(x[:, base:base + PT, 0:BK], (1, PT, BK), "smem")
                ws = tf.reshard(w[0:BK, :], (BK, NOUT @ m.u), "smem")   # <-- last axis
                acc = tf.insert_slice(acc, tf.matmul(xs, ws), (0, base, 0))
            return acc


@module(entry="proj", target=_H200, topologies=(_CTA,))
class StripMajor:
    @func
    def proj(x: Tensor[(1, S, H), "bf16"],
             w: ConstTensor[(GRID, H, NPER), "bf16"],           # strip-major
             out: Tensor[(GRID, S, NPER), "bf16"]) -> Tensor[(GRID, S, NPER), "bf16"]:
        with Mesh(("cta",), layout=(GRID,), names=("u",)) as m:
            acc = out
            for t in tile(S, PT):
                base = t + 0
                xs = tf.reshard(x[:, base:base + PT, 0:BK], (1, PT, BK), "smem")
                ws = tf.reshard(w[:, 0:BK, :], (GRID @ m.u, BK, NPER), "smem")  # <-- leading
                acc = tf.insert_slice(acc, tf.matmul(xs, ws), (0, base, 0))
            return acc

for M in LastAxis StripMajor; do
  tilefoundry analyze case.py:$M /tmp/$M.txt \
    --compute-cost --memory --roofline --performance --dim seq_len=8192
  grep -E '^# (compute-cost|roofline|performance) ' /tmp/$M.txt
done

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

[zhen8838]

Fixed in #123, with a regression test that fails without the change.