TheRock/rocm-libraries/projects/hipblaslt/tensilelite/Tensile/Tests/unit/test_lraTileAssignment.py

#!/usr/bin/env python3
################################################################################
# GPU functional test for lraTileAssignment with parameterized tile configs
#
# Usage:
#   pytest test_lraTileAssignment.py -v -s
################################################################################

import os
import struct
import sys
import tempfile

import pytest
from types import SimpleNamespace

from gpu_test_helpers import (
    HAS_GFX950,
    GFX_TARGET,
    TileConfig,
    BPE, LOAD_WIDTH, WAVESIZE, NUM_THREADS, NUM_WAVES,
    generate_lra_asm,
    export_register,
    print_offset_grid,
)


# ---- Reference implementations ----

def compute_expected_lr_offset(thread_id, cfg, tileInfo, ldsStartOffsetB=None):
    """Python reference implementation matching lraTileAssignment kernel logic.

    Mirrors the kernel's exact steps:
      1. lane16, lane16Group from Serial
      2. Swizzling: rotation + VPermlane16SwapB32 on colOffset (always on)
      3. rowOffset = lane16 * depthUBytes (no splitOffset — commented out in kernel)
      4. _computeLROffset: colOffset stepping by numMFMACols
      5. Wave partitioning via _applyWavePartitionLROffset
      6. B matrix LDS base offset
    """
    depthUBytes = cfg.depth_u * BPE
    blockSize = depthUBytes // LOAD_WIDTH
    numRowsPerLDSBanks = (WAVESIZE * 4) // depthUBytes
    numMFMACols = tileInfo.mmaTileShape[1] * tileInfo.bpe // LOAD_WIDTH
    laneId = thread_id % WAVESIZE

    # --- Step 1: lane16, lane16Group ---
    lane16 = laneId & 15
    lane16Group = laneId >> 4

    # --- Step 2: Swizzling (always on in kernel) ---
    # Rotation
    lds_row_id = lane16 >> (numRowsPerLDSBanks.bit_length() - 1)
    rotation = (lds_row_id // 2) * 2
    colOffset = (rotation + lane16Group) % blockSize

    # VPermlane16SwapB32 with exec mask 0x33333333
    # Active lanes: those where (lane_pos % 4) < 2, i.e. lanes 0,1,4,5,8,9,12,13,...
    # For active lanes, swap colOffset value with lane (laneId XOR 16).
    # We need to compute colOffset for the partner lane too.
    partnerLaneId = laneId ^ 16
    partnerLane16 = partnerLaneId & 15  # same as lane16 (XOR 16 doesn't change lower 4 bits)
    partnerLane16Group = partnerLaneId >> 4
    partnerLdsRowId = partnerLane16 >> (numRowsPerLDSBanks.bit_length() - 1)
    partnerRotation = (partnerLdsRowId // 2) * 2
    partnerColOffset = (partnerRotation + partnerLane16Group) % blockSize

    # Check if this lane is active under exec mask 0x33333333
    isActive = (laneId % 4) < 2
    if isActive:
        colOffset = partnerColOffset  # swap: we get partner's value

    # --- Step 3: rowOffset (no splitOffset) ---
    rowOffset = lane16 * depthUBytes

    # --- Step 4: _computeLROffset ---
    offsets = []
    for lr_idx in range(tileInfo.numLRPerSubtile):
        if lr_idx == 0:
            newColOffset = colOffset
        else:
            newColOffset = (colOffset + numMFMACols * lr_idx) % blockSize
        offsets.append(newColOffset * LOAD_WIDTH + rowOffset)

    # --- Step 5: Wave partitioning (non-interleaved layout) ---
    waveId = thread_id // WAVESIZE
    partitionOffset = 0
    MT = tileInfo.globalMMATileGrid[0] * tileInfo.mmaTileShape[0]
    subIterKBytes = depthUBytes

    if tileInfo.loadRatioGR >= 2.0:
        pass  # no partition needed
    elif tileInfo.loadRatioGR == 1.0:
        # 2x2 config, non-interleaved: sInterval = MT * subIterKBytes // 2
        if tileInfo.tc == 'A':
            wavePartId = waveId & 1
        else:
            wavePartId = waveId >> 1
        partitionOffset = wavePartId * (MT * subIterKBytes // 2)
    elif tileInfo.loadRatioGR == 0.5:
        # 4x1 / 1x4 config, non-interleaved: sInterval = MT * subIterKBytes // 4
        partitionOffset = waveId * (MT * subIterKBytes // 4)
    else:
        raise NotImplementedError(f"Unsupported loadRatioGR: {tileInfo.loadRatioGR}")

    # --- Step 6: B matrix LDS base offset ---
    if tileInfo.tc == 'B':
        if ldsStartOffsetB is not None:
            partitionOffset += ldsStartOffsetB
        else:
            partitionOffset += cfg.mt_a * depthUBytes

    for i in range(len(offsets)):
        offsets[i] += partitionOffset

    return offsets


def compute_expected_lr_subtile(subtileId0, cfg, tileInfo):
    """Compute expected LR subtile register value.

    """
    subtile_rows = tileInfo.subtileShape[0] * tileInfo.mmaTileShape[0]
    depthU_bytes = cfg.depth_u * BPE
    return subtile_rows * depthU_bytes * subtileId0


# Tile configs to test
TILE_CONFIGS = [
    # 2x2 configs
    TileConfig(mt_a=256, mt_b=256, depth_u=64),
    TileConfig(mt_a=96, mt_b=256, depth_u=64),
    # 1x4 configs
    TileConfig(mt_a=80, mt_b=64, depth_u=64),
    # 4x1 configs
    TileConfig(mt_a=64, mt_b=80, depth_u=64),
    TileConfig(mt_a=128, mt_b=240, depth_u=64),
]


# ---- Pytest tests ----

@pytest.mark.skipif(not HAS_GFX950, reason=f"GPU tests require gfx950, found {GFX_TARGET}")
class TestLraTileAssignmentGPU:

    @pytest.fixture(params=TILE_CONFIGS, ids=lambda c: c.label)
    def lra_env(self, request, tmp_path):
        """Generate lraTileAssignment asm once per tile config."""
        cfg = request.param
        lra_asm, writer, tileInfoA, tileInfoB, kernel = generate_lra_asm(cfg)
        return SimpleNamespace(
            cfg=cfg,
            lra_asm=lra_asm,
            writer=writer,
            tileInfoA=tileInfoA,
            tileInfoB=tileInfoB,
            kernel=kernel,
            tmp_path=tmp_path,
        )

    def test_offset_a(self, lra_env):
        """Validate all sharedVgprLROffset vgprs for matrix A across all threads."""
        cfg = lra_env.cfg
        for idx, reg in enumerate(lra_env.tileInfoA.sharedVgprLROffset):
            results = export_register(lra_env.writer, lra_env.lra_asm, reg, False,
                                      cfg, lra_env.tmp_path, f"lr_offsetA_v{reg}_{cfg.label}")

            for tid in range(NUM_THREADS):
                expected = compute_expected_lr_offset(tid, cfg, lra_env.tileInfoA,
                                                      lra_env.writer.ldsStartOffsetB)
                assert results[tid] == expected[idx], \
                    f"[{cfg.label}] A LR offset[{idx}] v{reg} mismatch at tid={tid}: " \
                    f"got {results[tid]}, expected {expected[idx]}"

    def test_offset_b(self, lra_env):
        """Validate all sharedVgprLROffset vgprs for matrix B across all threads."""
        cfg = lra_env.cfg
        for idx, reg in enumerate(lra_env.tileInfoB.sharedVgprLROffset):
            results = export_register(lra_env.writer, lra_env.lra_asm, reg, False,
                                      cfg, lra_env.tmp_path, f"lr_offsetB_v{reg}_{cfg.label}")

            for tid in range(NUM_THREADS):
                expected = compute_expected_lr_offset(tid, cfg, lra_env.tileInfoB,
                                                      lra_env.writer.ldsStartOffsetB)
                assert results[tid] == expected[idx], \
                    f"[{cfg.label}] B LR offset[{idx}] v{reg} mismatch at tid={tid}: " \
                    f"got {results[tid]}, expected {expected[idx]}"




if __name__ == "__main__":
    """Run standalone without pytest."""
    import argparse
    parser = argparse.ArgumentParser(description="GPU test for lraTileAssignment")
    parser.add_argument("--grid", action="store_true",
                        help="Display offsets as 2D grid (waves x lanes) for A and B")
    parser.add_argument("--debug", action="store_true",
                        help="Display expected matrix in grid mode (implies --grid)")
    args = parser.parse_args()
    if args.debug:
        args.grid = True

    for cfg in TILE_CONFIGS:
        print(f"\n{'='*60}")
        print(f"  Tile Config: {cfg.label}")
        print(f"{'='*60}")

        lra_asm, writer, tileInfoA, tileInfoB, kernel = generate_lra_asm(cfg)

        with tempfile.TemporaryDirectory() as tmp_dir:
            tmp_path = type('P', (), {'__truediv__': lambda s, n: os.path.join(tmp_dir, n)})()

            # Print the generated assembly for inspection
            print("\n--- Generated Assembly (lraTileAssignment section) ---")
            for line in lra_asm.split('\n'):
                print(line)
            print("--- End ---\n")

            print(f"  TileInfoA: numLRPerSubtile={tileInfoA.numLRPerSubtile}, "
                  f"loadRatioLR={tileInfoA.loadRatioLR}, "
                  f"sharedVgprLROffset={tileInfoA.sharedVgprLROffset}")
            print(f"  TileInfoB: numLRPerSubtile={tileInfoB.numLRPerSubtile}, "
                  f"loadRatioLR={tileInfoB.loadRatioLR}, "
                  f"sharedVgprLROffset={tileInfoB.sharedVgprLROffset}")

            # Test all sharedVgprLROffset vgprs for both matrices
            for tc, tileInfo in [("A", tileInfoA), ("B", tileInfoB)]:
                for idx, reg in enumerate(tileInfo.sharedVgprLROffset):
                    results = export_register(writer, lra_asm, reg, False, cfg, tmp_path,
                                              f"lr_offset{tc}_v{reg}_{cfg.label}")

                    if args.grid:
                        print_offset_grid(f"Matrix {tc} LR GPU offset[{idx}] v{reg} ({cfg.label})",
                                          results, WAVESIZE, NUM_WAVES)

                        if args.debug:
                            expected = [compute_expected_lr_offset(tid, cfg, tileInfo, writer.ldsStartOffsetB)[idx]
                                        for tid in range(NUM_THREADS)]
                            print_offset_grid(f"Matrix {tc} LR EXPECTED offset[{idx}] ({cfg.label})",
                                              expected, WAVESIZE, NUM_WAVES)

                            mismatches = sum(1 for t in range(NUM_THREADS)
                                             if results[t] != expected[t])
                            if mismatches:
                                print(f"\n--- Matrix {tc} LR offset[{idx}] DIFF ({mismatches} mismatches) ---")
                                for w in range(NUM_WAVES):
                                    print(f"  w{w}: ", end="")
                                    for lane in range(WAVESIZE):
                                        tid = w * WAVESIZE + lane
                                        if results[tid] != expected[tid]:
                                            print(f" t{tid}:{results[tid]}!={expected[tid]}", end="")
                                    print()
                            else:
                                print(f"\n  Matrix {tc} LR offset[{idx}]: all match.")

                    errors = 0
                    for tid in range(NUM_THREADS):
                        exp = compute_expected_lr_offset(tid, cfg, tileInfo, writer.ldsStartOffsetB)[idx]
                        if results[tid] != exp:
                            errors += 1
                            if not args.grid:
                                print(f"  FAIL {tc} LR offset[{idx}] v{reg} tid={tid}: "
                                      f"got {results[tid]}, expected {exp}")
                        elif not args.grid and tid < 64:
                            print(f"  OK   {tc} LR offset[{idx}] v{reg} tid={tid}: {results[tid]}")

                    print(f"  Matrix {tc} LR offset[{idx}] v{reg}: "
                          f"{NUM_THREADS} threads, {errors} errors")