Function _validate_mma_alloc_shape

python/tvm/backend/cuda/lang/alloc_pool.py:102–173

Validate that *shape* / *dtype* / *swizzle_mode* are mutually compatible. ``mma_shared_layout`` tiles a swizzle atom of shape ``[8, swizzle_bytes / dtype_bytes]`` over the last two logical dimensions of *shape*. If the row width or row count of the request is smaller than (or not a mult

(shape, dtype, swizzle_mode)

def _validate_mma_alloc_shape(shape, dtype, swizzle_mode):
    """Validate that *shape* / *dtype* / *swizzle_mode* are mutually compatible.

    ``mma_shared_layout`` tiles a swizzle atom of shape ``[8, swizzle_bytes / dtype_bytes]``
    over the last two logical dimensions of *shape*. If the row width or row count of
    the request is smaller than (or not a multiple of) the atom, the underlying
    ``Layout.tile_to`` lowers to a ``floordiv``/``floormod`` by zero and raises an
    opaque internal "Divide by zero" diagnostic from ``tile_tile_ops.cc``. Catch the
    misconfiguration here so callers see *what* is wrong and *how* to fix it.

    Validation skipped when *swizzle_mode* is ``SWIZZLE_NONE`` (no atom).
    """
    from tvm.backend.cuda.operator.tile_primitive.tma_utils import SwizzleMode

    if swizzle_mode == SwizzleMode.SWIZZLE_NONE:
        return

    if len(shape) < 2:
        raise ValueError(
            f"alloc_mma shape={tuple(shape)} has fewer than 2 dimensions; "
            f"swizzled MMA layouts tile over the last two dims (rows, cols). "
            f"Use swizzle_mode='none' for 1-D allocations."
        )

    # Only validate concrete int dims; symbolic dims fall through (the analyzer
    # in C++ will still ICHECK on them, but at least we don't false-positive).
    rows = shape[-2]
    cols = shape[-1]
    if not (isinstance(rows, int) and isinstance(cols, int)):
        return

    dtype_bytes = DataType(dtype).bits // 8
    if dtype_bytes == 0:
        # Sub-byte dtype (e.g. float4); ``cols`` is already in element units, so
        # use a fractional check expressed via bits.
        col_bits = cols * DataType(dtype).bits
        atom_bits = _swizzle_atom_bytes(swizzle_mode) * 8
        if col_bits < atom_bits or col_bits % atom_bits != 0:
            row_bytes = col_bits // 8 if col_bits % 8 == 0 else col_bits / 8
            atom_bytes = _swizzle_atom_bytes(swizzle_mode)
            suggestion = _suggest_swizzle_for_row_bytes(col_bits // 8 if col_bits >= 8 else 0)
            raise ValueError(
                f"alloc_mma shape={tuple(shape)} with dtype={dtype!r} produces "
                f"{row_bytes}B rows, which is incompatible with the {atom_bytes}B "
                f"swizzle atom selected by {swizzle_mode.name}. "
                f"Use swizzle_mode=SwizzleMode.{suggestion}, or widen shape[-1] "
                f"to a multiple of "
                f"{(atom_bits + DataType(dtype).bits - 1) // DataType(dtype).bits} elements."
            )
    else:
        row_bytes = cols * dtype_bytes
        atom_bytes = _swizzle_atom_bytes(swizzle_mode)
        if row_bytes < atom_bytes or row_bytes % atom_bytes != 0:
            suggestion = _suggest_swizzle_for_row_bytes(row_bytes)
            min_cols = atom_bytes // dtype_bytes
            raise ValueError(
                f"alloc_mma shape={tuple(shape)} with dtype={dtype!r} produces "
                f"{row_bytes}B rows, which is incompatible with the {atom_bytes}B "