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

################################################################################
#
# Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved.
#
# SPDX-License-Identifier: MIT
#
################################################################################

from contextlib import contextmanager
from copy import deepcopy
from types import SimpleNamespace

import pytest
from rocisa.code import Module
from rocisa.enum import RegisterType
from rocisa.register import RegisterPool

import Tensile.KernelWriter as kw_module
from Tensile.KernelWriter import KernelWriter
import Tensile.KernelWriterAssembly as kwa_module
from Tensile.Components.StreamK import StreamKTwoTileDPFirst
from Tensile.Common.GlobalParameters import defaultSolution, globalParameters
from Tensile.Common.RequiredParameters import getRequiredParametersMin
from Tensile.Common.Types import IsaInfo, IsaVersion, SemanticVersion
from Tensile.Common.ValidParameters import validParameters
from Tensile.Contractions import SizeMapping
from Tensile.SolutionStructs.Solution import (
    Solution,
    _disableUnsupportedRuntimeStaggerU,
    validateParameterTypes,
)

pytestmark = pytest.mark.unit


# Tensile keeps process-global, module-level default dicts (`defaultSolution`,
# `globalParameters`) that `Solution.__init__` reads while constructing a solution.
# Some sibling unit tests mutate these in place -- e.g. test_MatrixInstructionConversion
# injects a "ProblemType" key into `defaultSolution`, which makes Solution.__init__'s
# `for key in defaultSolution` loop overwrite the already-converted ProblemType object
# with the raw config dict, leaving DataType a str and crashing
# assignProblemIndependentDerivedParameters. That manifested as order-dependent
# failures of the Solution-validation tests below under pytest-xdist. Snapshot the
# pristine defaults at import time (collection runs before any test executes, so they
# are clean here) and restore them around every test so Solution construction in this
# module is hermetic regardless of suite ordering.
_PRISTINE_DEFAULT_SOLUTION = deepcopy(defaultSolution)
_PRISTINE_GLOBAL_PARAMETERS = deepcopy(globalParameters)


@pytest.fixture(autouse=True)
def _isolate_global_solution_state():
    def _restore(target, pristine):
        target.clear()
        target.update(deepcopy(pristine))

    _restore(defaultSolution, _PRISTINE_DEFAULT_SOLUTION)
    _restore(globalParameters, _PRISTINE_GLOBAL_PARAMETERS)
    yield
    _restore(defaultSolution, _PRISTINE_DEFAULT_SOLUTION)
    _restore(globalParameters, _PRISTINE_GLOBAL_PARAMETERS)


def _module_with_comment(name, comment):
    module = Module(name)
    module.addComment0(comment)
    return module


def _tensor_module(name, comment, tensor_parameters):
    tc = tensor_parameters["tensorChar"]
    return _module_with_comment("%s_%s" % (name, tc), "%s %s" % (comment, tc))


class _ClassicPapWriter:
    def __init__(self, *, version=(9, 5, 0), use64b_shadow=False, use64b_shadow_mx=False):
        self.states = SimpleNamespace(
            a=SimpleNamespace(numVgprGlobalReadOffsets=2),
            b=SimpleNamespace(numVgprGlobalReadOffsets=2),
            ldsTensorTokenIdx=0,
            memTokenLdsBuffer0=0,
            memTokenLdsBuffer1=1,
            staggerUCode=False,
            unrollIdx=0,
            use64bShadowLimit=use64b_shadow,
            use64bShadowLimitMX=use64b_shadow_mx,
            version=version,
        )
        self._next_tmp_sgpr = 100
        self.vgprPool = _TrackingRegisterPool(RegisterType.Vgpr)
        self.startVgprGlobalReadOffsetA = 200
        self.startVgprGlobalReadOffsetB = 210

    @contextmanager
    def allocTmpSgpr(self, size, alignment=1, tag=""):
        base = self._next_tmp_sgpr
        self._next_tmp_sgpr += size + alignment
        yield SimpleNamespace(idx=base, size=size)

    def isSwapGlobalReadOrderForDtvOrDtl(self, kernel, prefetch1=False):
        return False

    def isPrefetchAcrossPersistentEnabled(self, kernel):
        return True

    def openSumAtLeastUnroll(self, kernel, prefetch=False, isOptNLL=True):
        return _module_with_comment("openSumAtLeastUnroll", "unit: open sum")

    def declareStaggerParms(self, kernel):
        return _module_with_comment("declareStaggerParms", "unit: declare stagger")

    def lwaTileAssignment(self, kernel, tensor_parameters):
        return _tensor_module("lwaTileAssignment", "unit: LWA tile", tensor_parameters)

    def graTileAssignment(self, kernel, tensor_parameters):
        return _tensor_module("graTileAssignment", "unit: tile assignment", tensor_parameters)

    def graUnrollAssignment(self, kernel, tensor_parameters):
        return _tensor_module("graUnrollAssignment", "unit: unroll assignment", tensor_parameters)

    def graTileOffsets(self, kernel, tensor_parameters):
        return _tensor_module("graTileOffsets", "unit: tile offsets", tensor_parameters)

    def graUnrollOffsets(self, kernel, tensor_parameters):
        return _tensor_module("graUnrollOffsets", "unit: unroll offsets", tensor_parameters)

    def graShift(self, kernel, tensor_parameters):
        return _tensor_module("graShift", "unit: shift", tensor_parameters)

    def graAddresses(self, kernel, tensor_parameters):
        return _tensor_module("graAddresses", "unit: GRA", tensor_parameters)

    def graFinalOffsets(self, kernel, tensor_parameters):
        return _tensor_module("graFinalOffsets", "unit: final offsets", tensor_parameters)

    def calculateStagger(self, kernel, tensor_parameters):
        return _tensor_module("calculateStagger", "unit: stagger", tensor_parameters)

    def directToLdsM0Update(self, kernel, offset, tensor_parameters, skipWait=False):
        return _tensor_module(
            "directToLdsM0Update",
            "unit: M0 %s skipWait=%s" % (tensor_parameters["tensorChar"], skipWait),
            tensor_parameters,
        )

    def globalReadDo(self, kernel, offset, tensor_parameters):
        return _tensor_module("globalReadDo", "unit: GR", tensor_parameters)

    def papDtlSaveLdsBank(self, kernel, tensor_parameters_a, tensor_parameters_b):
        return _module_with_comment("papDtlSaveLdsBank", "unit: save DTL LDS bank")


_ClassicPapWriter.setupPrefetchAcrossPersistentLoads = KernelWriter.setupPrefetchAcrossPersistentLoads


class _SetupNewTilePapTdmWriter:
    def __init__(self):
        self.states = SimpleNamespace(
            actualSummationLoops=1,
            doShadowInit=2,
            IncLdsBufSwitch=False,
            ldsTensorTokenIdx=0,
            memTokenLdsBuffer0=0,
            memTokenLdsBuffer1=1,
            staggerUCode=False,
            unrollIdx=0,
        )
        self.do = {"executeToInitEnd": False}
        self.dontAppendCode = False
        self.labels = _StubLabels()

    def _module(self, name):
        return _module_with_comment(name, "unit: %s" % name)

    def graWorkGroup(self, *args):
        return self._module("graWorkGroup")

    def localReadAddresses(self, *args):
        return self._module("localReadAddresses")

    def localWriteAddresses(self, *args):
        return self._module("localWriteAddresses")

    def removeGRSrdVariableSgprsFromPool(self, kernel):
        return self._module("removeGRSrdVariableSgprsFromPool")

    def initTDMDescriptorWaveSeparated(self, kernel, tpa, tpb):
        return self._module("initTDMDescriptorWaveSeparated_%s_%s" % (tpa["tensorChar"], tpb["tensorChar"]))

    def tdmGlobalOffsetWaveSeparated(self, kernel, tpa, tpb):
        return self._module("tdmGlobalOffsetWaveSeparated_%s_%s" % (tpa["tensorChar"], tpb["tensorChar"]))

    def removeGROffsetsVariableSgprsFromPool(self, kernel):
        return self._module("removeGROffsetsVariableSgprsFromPool")

    def tdmSetupIncrementWaveSeparated(self, kernel, tpa, tpb):
        return self._module("tdmSetupIncrementWaveSeparated_%s_%s" % (tpa["tensorChar"], tpb["tensorChar"]))

    def tdmApplyStreamKOffsetWaveSeparated(self, kernel, tpa, tpb):
        return self._module("tdmApplyStreamKOffsetWaveSeparated_%s_%s" % (tpa["tensorChar"], tpb["tensorChar"]))

    def releaseGlobalReadIncsSgprsAfterTdmWaveSep(self, kernel):
        return self._module("releaseGlobalReadIncsSgprsAfterTdmWaveSep")

    def undefineSgpr(self, name):
        return self._module("undefineSgpr_%s" % name)

    def initC(self, kernel):
        return self._module("initC")

    def calculateLoopNumIter(self, kernel, tpa, tpb, loop_idx):
        return self._module("calculateLoopNumIter")

    def localReadInitPointers(self, kernel, tpa, tpb):
        return self._module("localReadInitPointers_%s" % tpb["tensorChar"])

    def isPrefetchAcrossPersistentEnabled(self, kernel):
        return KernelWriter.isPrefetchAcrossPersistentEnabled(self, kernel)

    def papTdmRestoreLdsBank(self, kernel, tpa, tpb):
        return self._module("papTdmRestoreLdsBank")

    def isSwapGlobalReadOrderForDtvOrDtl(self, kernel, prefetch1=False):
        return False

    def openSumAtLeastUnroll(self, kernel, prefetch=False, isOptNLL=True):
        return self._module("openSumAtLeastUnroll")

    def directToLdsM0Update(self, kernel, offset, tensor_parameters, skipWait=False):
        return self._module("directToLdsM0Update_%s" % tensor_parameters["tensorChar"])

    def globalReadDo(self, kernel, offset, tensor_parameters, **kwargs):
        return self._module("globalReadDo_%s" % tensor_parameters["tensorChar"])

    def globalReadIncrementAB(self, kernel, tpa, tpb, loop_idx, prefetch_index):
        return self._module("globalReadIncrementAB")


class _StubGsu:
    def setupNewTile(self, writer, kernel, tpa, tpb, tpm):
        return _module_with_comment("gsuSetupNewTile", "unit: GSU setup")


class _StubTdmComp:
    def getLdsAddrSgprName(self, group_name):
        return "%s+1" % group_name


class _PapTdmDescriptorRefreshWriter:
    def __init__(self):
        self._next_tmp_sgpr = 300
        self.recomputed_waveidx = []
        self.init_waveidx = []
        self.global_offset_waveidx = []

    @contextmanager
    def allocTmpSgpr(self, size, alignment=1, tag=""):
        base = self._next_tmp_sgpr
        self._next_tmp_sgpr += size + alignment
        yield SimpleNamespace(idx=base, size=size)

    def papTdmRecomputeWaveIdx(self, kernel, wave_idx_sgpr):
        self.recomputed_waveidx.append(wave_idx_sgpr)
        return _module_with_comment("papTdmRecomputeWaveIdx", "unit: recompute WaveIdx")

    def initTDMDescriptorWaveSeparated(self, kernel, tpa, tpb, wave_idx_sgpr="WaveIdx"):
        self.init_waveidx.append(wave_idx_sgpr)
        return _module_with_comment("initTDMDescriptorWaveSeparated", "unit: init TDM descriptor")

    def tdmGlobalOffsetWaveSeparated(self, kernel, tpa, tpb, wave_idx_sgpr="WaveIdx"):
        self.global_offset_waveidx.append(wave_idx_sgpr)
        return _module_with_comment("tdmGlobalOffsetWaveSeparated", "unit: global offset")

    def tdmApplyStreamKOffsetWaveSeparated(self, kernel, tpa, tpb):
        return _module_with_comment("tdmApplyStreamKOffsetWaveSeparated", "unit: StreamK offset")


class _TrackingRegisterPool(RegisterPool):
    def __init__(self, register_type):
        super().__init__(0, register_type, defaultPreventOverflow=False, printRP=False)
        self.checked_out = []
        self.checked_in = []

    def checkOutAligned(self, size, alignment=1, tag="", *args, **kwargs):
        base = super().checkOutAligned(size, alignment, tag, *args, **kwargs)
        self.checked_out.append((base, size, tag))
        return base

    def checkIn(self, vgpr):
        super().checkIn(vgpr)
        self.checked_in.append(vgpr)


class _StubLabels:
    def __init__(self):
        self._count = 0

    def getNameInc(self, name):
        self._count += 1
        return "%s_%u" % (name, self._count)


class _ClassicPapWrapperWriter:
    def __init__(self):
        self.labels = _StubLabels()
        self.states = SimpleNamespace(unrollIdx=0)
        self.vgprPool = _TrackingRegisterPool(RegisterType.Vgpr)

    def isPrefetchAcrossPersistentEnabled(self, kernel):
        return True

    @contextmanager
    def allocPapTileIdentitySgprs(self, kernel):
        yield {
            "WorkGroup0": 100,
            "WorkGroup1": 101,
            "WorkGroup2": 102,
            "StreamKLocalStart": 103,
            "StreamKLocalEnd": 104,
        }

    def papCheckpointCurrentTileIdentity(self, kernel, prev_tile):
        return _module_with_comment("papCheckpointCurrentTileIdentity", "unit: checkpoint tile")

    def loopCounterName(self, kernel, loop_idx):
        return "LoopCounterL"

    def calculateLoopNumIter(self, kernel, tpa, tpb, loop_idx):
        return _module_with_comment("calculateLoopNumIter", "unit: calculate loop num iter")

    def setupPrefetchAcrossPersistentLoads(self, kernel, tpa, tpb, isOptNLL=True):
        return _module_with_comment("setupPrefetchAcrossPersistentLoads", "unit: setup PAP loads")

    def papRestoreCurrentTileIdentity(self, kernel, prev_tile):
        return _module_with_comment("papRestoreCurrentTileIdentity", "unit: restore tile")


class _StubStreamK:
    def prefetchAcrossPersistentSetupNextTile(self, writer, kernel, tpa, tpb, skipLroReset=False):
        return _module_with_comment("prefetchAcrossPersistentSetupNextTile", "unit: setup next tile")


def _problem_type(**overrides):
    problem_type = {
        "BiasSrc": "D",
        "Gradient": False,
        "MXBlockA": 0,
        "MXBlockB": 0,
        "NumIndicesSummation": 1,
        "Sparse": 0,
        "UseBias": False,
    }
    problem_type.update(overrides)
    return problem_type


def _kernel_from(base, **overrides):
    kernel = deepcopy(base)
    problem_type_overrides = overrides.pop("ProblemType", None)
    kernel.update(overrides)
    if problem_type_overrides is not None:
        kernel["ProblemType"].update(problem_type_overrides)
    return kernel


_CLASSIC_KERNEL_BASE = {
    "BufferLoad": True,
    "DirectToLdsA": False,
    "DirectToLdsB": False,
    "DirectToVgprA": False,
    "DirectToVgprB": False,
    "EdgeType": "None",
    "GuaranteeNoPartialA": False,
    "GuaranteeNoPartialB": False,
    "NoTailLoop": False,
    "PrefetchGlobalRead": 2,
    "PrefetchGL2": 0,
    "ProblemType": _problem_type(),
    "StreamKForceDPOnly": 0,
    "UseGeneralizedNLCOneA": False,
    "UseGeneralizedNLCOneB": False,
    "_UseSgprForGRO": False,
    "enableTDMA": False,
    "enableTDMB": False,
}


_SETUP_NEW_TILE_TDM_BASE = _kernel_from(
    _CLASSIC_KERNEL_BASE,
    ClusterBarrier=False,
    DirectToVgprSparseMetadata=False,
    enableTDMA=True,
    enableTDMB=True,
    enableTDMMetadata=False,
    GuaranteeNoPartialA=True,
    GuaranteeNoPartialB=True,
    GuaranteeNoPartialMetadata=False,
    MIWaveGroup=[2, 1],
    Multicast=False,
    NumWaves=2,
    PrefetchAcrossPersistent=1,
    PrefetchGlobalRead=1,
    StreamK=3,
    SuppressNoLoadLoop=False,
    TDMInst=3,
    UseCustomMainLoopSchedule=0,
    UseSubtileImpl=False,
)


def _classic_kernel(**overrides):
    return _kernel_from(_CLASSIC_KERNEL_BASE, **overrides)


def _setup_new_tile_tdm_kernel(prefetch_across_persistent=1, **overrides):
    return _kernel_from(
        _SETUP_NEW_TILE_TDM_BASE,
        PrefetchAcrossPersistent=prefetch_across_persistent,
        **overrides,
    )


def _pap_wrapper_kernel(**overrides):
    return _classic_kernel(
        PrefetchAcrossPersistent=1,
        StreamK=3,
        SpaceFillingAlgo=[],
        **overrides,
    )


def _tensor_parameters(*, with_mx=False, with_metadata=False):
    tpa = {"tensorChar": "A", "isSwizzled": False}
    tpb = {"tensorChar": "B", "isSwizzled": False}
    if with_mx:
        tpa["MX"] = {"tensorChar": "MXSA", "isSwizzled": False}
        tpb["MX"] = {"tensorChar": "MXSB", "isSwizzled": False}
    if with_metadata:
        metadata = {"tensorChar": "Metadata", "isSwizzled": False}
        tpa.update({"is_sparse": False, "tpsMetadata": metadata})
        tpb.update({"is_sparse": False, "tpsMetadata": metadata})
    return tpa, tpb


_RUNTIME_STAGGER_BASE = {
    "PrefetchAcrossPersistent": 0,
    "TDMInst": 0,
    "StaggerU": 32,
    "StaggerUMapping": 2,
    "StaggerUStride": 256,
    "InternalSupportParams": {"SupportCustomStaggerU": True},
    "ProblemType": {"MXBlockA": 0, "MXBlockB": 0},
    "enableTDMA": False,
    "enableTDMB": False,
    "DirectToLdsA": False,
    "DirectToLdsB": False,
}


def _stagger_runtime_state(*, pap=False, tdm=False, **overrides):
    state = deepcopy(_RUNTIME_STAGGER_BASE)
    if pap:
        state["PrefetchAcrossPersistent"] = 1
    if tdm:
        state.update({"TDMInst": 3, "enableTDMA": True, "enableTDMB": True})
    state.update(overrides)
    return state


class _DefaultFalseDict(dict):
    def __missing__(self, key):
        return False


def _pap_solution_isa_info_map():
    isa = IsaVersion(9, 5, 0)
    asm_caps = _DefaultFalseDict(
        {
            "SupportedISA": True,
            "HasMFMA": True,
            "HasTDM": True,
            "HasDirectToLds": True,
            "HasDirectToLdsx4": True,
            "HasNTModifier": True,
            "HasMFMA_f64": True,
            "HasGLTr8B64": True,
            "HasGLTr16B128": True,
            "HasLDSTr": True,
        }
    )
    arch_caps = _DefaultFalseDict(
        {
            "DeviceLDS": 65536,
            "HasEccHalf": True,
            "HasSchedMode": True,
            "HasAccCD": True,
            "HasMXScaleSwizzle": True,
        }
    )
    return {isa: IsaInfo(asm_caps, arch_caps, _DefaultFalseDict(), _DefaultFalseDict())}


def _pap_solution_config(**overrides):
    config = {
        "ISA": [9, 5, 0],
        "EnableMatrixInstruction": True,
        "ProblemType": {
            "OperationType": "GEMM",
            "DataType": "s",
            "DestDataType": "s",
            "ComputeDataType": "s",
            "TransposeA": True,
            "TransposeB": False,
            "UseBeta": True,
            "Batched": True,
            "StridedBatched": True,
        },
        "MatrixInstruction": [32, 32, 8, 1],
        "MIBlock": [32, 32, 8, 1, 1, 1],
        "MIWaveGroup": [2, 1],
        "MIWaveTile": [1, 1],
        "MIInputPerThread": 1,
        "WorkGroup": [32, 2, 1],
        "StreamK": 3,
        "PrefetchAcrossPersistent": 1,
        "PrefetchGlobalRead": 1,
        "ScheduleIterAlg": 0,
        "TDMInst": 3,
        "StaggerU": 0,
        "1LDSBuffer": 0,
        "BufferLoad": True,
        "BufferStore": True,
        "StoreRemapVectorWidth": 0,
        "SuppressNoLoadLoop": False,
        "DirectToVgprA": False,
        "DirectToVgprB": False,
        "UseSubtileImpl": False,
        "DepthU": 32,
        "LocalSplitU": 1,
    }
    problem_type_overrides = overrides.pop("ProblemType", None)
    config.update(overrides)
    if problem_type_overrides is not None:
        config["ProblemType"].update(problem_type_overrides)
    return config


def _pap_solution(**overrides):
    assembler = SimpleNamespace(
        code_object_version="default",
        rocm_version=SemanticVersion(6, 4, 0),
    )
    return Solution(
        _pap_solution_config(**overrides),
        False,
        True,
        False,
        assembler,
        _pap_solution_isa_info_map(),
    )


def _module_items(module):
    return [module.getItem(i) for i in range(module.itemsSize())]


def _module_index(items, name):
    return next(i for i, item in enumerate(items) if isinstance(item, Module) and item.name == name)


def _module_names(module_or_items):
    items = _module_items(module_or_items) if isinstance(module_or_items, Module) else module_or_items
    return [item.name for item in items if isinstance(item, Module)]


def _instruction_indices(items, instruction_type, *, dst_contains=None, src_contains=None):
    indices = []
    for i, item in enumerate(items):
        if not isinstance(item, instruction_type):
            continue
        dst = str(getattr(item, "dst", ""))
        srcs = [str(item_src) for item_src in getattr(item, "srcs", [])]
        if dst_contains is not None and dst_contains not in dst:
            continue
        if src_contains is not None and not any(src_contains in src for src in srcs):
            continue
        indices.append(i)
    return indices


def _instruction_index(items, instruction_type, dst, src):
    return next(
        i
        for i, item in enumerate(items)
        if isinstance(item, instruction_type)
        and str(item.dst) == dst
        and [str(item_src) for item_src in item.srcs] == [src]
    )


def _setup_new_tile_module_names(prefetch_across_persistent):
    tpa, tpb = _tensor_parameters(with_metadata=True)
    module = KernelWriter.setupNewTile(
        _SetupNewTilePapTdmWriter(),
        _setup_new_tile_tdm_kernel(prefetch_across_persistent=prefetch_across_persistent),
        tpa,
        tpb,
    )
    return _module_names(module)


def _prefetch_across_persistent(monkeypatch, *, skip_barrier=False, **kernel_overrides):
    monkeypatch.setattr(kwa_module.Component.StreamK, "find", lambda writer: _StubStreamK())
    writer = _ClassicPapWrapperWriter()
    module = kwa_module.KernelWriterAssembly.prefetchAcrossPersistent(
        writer,
        _pap_wrapper_kernel(**kernel_overrides),
        *_tensor_parameters(),
        skipBarrier=skip_barrier,
    )
    return writer, _module_items(module)


def _streamk_with_stubbed_tile_indexing():
    streamk = StreamKTwoTileDPFirst()
    streamk.skTileIndex = lambda writer, kernel, s_tmp, tpa, tpb, skipLroReset=False: (
        _module_with_comment("skTileIndex", "unit: tile index")
    )
    streamk.skIndexToWG = lambda writer, kernel, s_tmp: _module_with_comment(
        "skIndexToWG", "unit: index to WG"
    )
    return streamk


def _streamk_wgm_writer():
    writer = SimpleNamespace(
        sgprPool=RegisterPool(0, RegisterType.Sgpr, defaultPreventOverflow=False, printRP=False),
        states=SimpleNamespace(WGMTransformLevels=-1),
    )

    # prefetchAcrossPersistentSetupNextTile now takes its SKPrefetchTemp through the
    # growable allocTmpSgpr path (see StreamK.py); provide a counter-based stub
    # matching the other PAP mock writers (skTileIndex/skIndexToWG are stubbed, so
    # the actual register index is irrelevant here).
    next_tmp = [100]

    @contextmanager
    def _alloc_tmp_sgpr(size, alignment=1, tag=""):
        base = next_tmp[0]
        next_tmp[0] += size + alignment
        yield SimpleNamespace(idx=base, size=size)

    writer.allocTmpSgpr = _alloc_tmp_sgpr
    return writer


def test_pap_is_valid_solution_parameter():
    assert validParameters["PrefetchAcrossPersistent"] == [0, 1]
    assert defaultSolution["PrefetchAcrossPersistent"] == 0
    assert "PrefetchAcrossPersistent" in getRequiredParametersMin()
    assert "prefetchAcrossPersistent" in SizeMapping.StateKeys
    validateParameterTypes({"PrefetchAcrossPersistent": 1})


def test_solution_validation_accepts_minimal_pap_tdm_contract():
    assert _pap_solution()["Valid"] is True


@pytest.mark.parametrize(
    "overrides, reason",
    [
        pytest.param(
            {"1LDSBuffer": 1},
            "PrefetchAcrossPersistent requires 1LDSBuffer != 1",
            id="rejects_single_lds_buffer",
        ),
        pytest.param(
            {"StaggerU": 32},
            "TDM + PrefetchAcrossPersistent with StaggerU is not implemented",
            id="rejects_nonzero_staggeru",
        ),
        pytest.param(
            # ScheduleIterAlg 3 so the SIA0+StreamK rule (which already caps PGR at 2)
            # doesn't reject first; this exercises the PAP PGR<=2 contract directly.
            {"PrefetchGlobalRead": 3, "ScheduleIterAlg": 3},
            "PrefetchAcrossPersistent requires PrefetchGlobalRead in [1, 2]",
            id="rejects_pgr_above_two",
        ),
    ],
)
def test_solution_validation_rejects_unsupported_pap_tdm_contracts(capsys, overrides, reason):
    solution = _pap_solution(**overrides)

    assert solution["Valid"] is False
    assert reason in capsys.readouterr().out


@pytest.mark.parametrize(
    "pap, tdm, expected",
    [
        pytest.param(
            True,
            True,
            (0, 0, 0, False),
            id="pap_tdm_disables_runtime_custom_staggeru",
        ),
        pytest.param(
            False,
            True,
            (32, 2, 256, True),
            id="non_pap_tdm_keeps_runtime_custom_staggeru",
        ),
    ],
)
def test_runtime_staggeru_controls_for_tdm_pap(pap, tdm, expected):
    state = _stagger_runtime_state(pap=pap, tdm=tdm, StaggerU=0 if pap else 32)

    _disableUnsupportedRuntimeStaggerU(state)

    stagger_u, mapping, stride, support_custom = expected
    assert state["StaggerU"] == stagger_u
    assert state["StaggerUMapping"] == mapping
    assert state["StaggerUStride"] == stride
    assert state["InternalSupportParams"]["SupportCustomStaggerU"] is support_custom


def test_setup_new_tile_releases_waveidx_for_pap_wave_separated_tdm(monkeypatch):
    monkeypatch.setattr(kw_module.Component.GSU, "find", lambda writer: _StubGsu())

    pap_module_names = _setup_new_tile_module_names(prefetch_across_persistent=1)
    non_pap_module_names = _setup_new_tile_module_names(prefetch_across_persistent=0)

    assert "papTdmRestoreLdsBank" in pap_module_names
    assert "undefineSgpr_WaveIdx" in pap_module_names
    assert "undefineSgpr_WaveIdx" in non_pap_module_names
    assert pap_module_names.index("undefineSgpr_WaveIdx") < pap_module_names.index("papTdmRestoreLdsBank")


def test_pap_tdm_descriptor_refresh_threads_temporary_waveidx(monkeypatch):
    monkeypatch.setattr(kwa_module.TensorDataMoverLoad, "find", lambda writer: _StubTdmComp())
    writer = _PapTdmDescriptorRefreshWriter()
    kernel = {"LdsOffsetA_Blk": 0, "StreamK": 3}
    tpa, tpb = {"tensorChar": "A"}, {"tensorChar": "B"}

    module = kwa_module.KernelWriterAssembly.papTdmUpdateDescriptor(writer, kernel, tpa, tpb)

    assert _module_names(module) == [
        "papTdmRecomputeWaveIdx",
        "initTDMDescriptorWaveSeparated",
        "tdmGlobalOffsetWaveSeparated",
        "tdmApplyStreamKOffsetWaveSeparated",
    ]
    assert writer.recomputed_waveidx == [300]
    assert writer.init_waveidx == [300]
    assert writer.global_offset_waveidx == [300]


def test_classic_pap_primes_mx_first_pgr_group_before_marking_primed():
    writer = _ClassicPapWriter(version=(9, 5, 0))
    kernel = _classic_kernel(ProblemType={"MXBlockA": 32, "MXBlockB": 32, "Sparse": 0})
    tpa, tpb = _tensor_parameters(with_mx=True)

    module = writer.setupPrefetchAcrossPersistentLoads(kernel, tpa, tpb)
    items = _module_items(module)

    gr_a = _module_index(items, "globalReadDo_A")
    gr_mxsa = _module_index(items, "globalReadDo_MXSA")
    gr_mxsb = _module_index(items, "globalReadDo_MXSB")
    gr_b = _module_index(items, "globalReadDo_B")
    primed = _instruction_index(items, kwa_module.SMovB32, "s[sgprSkPrefetchPrimed]", "1")

    assert gr_a < gr_mxsa
    assert gr_mxsa < gr_mxsb
    assert gr_mxsb < gr_b
    assert gr_b < primed


def test_classic_pap_restores_gfx1250_shadow_limit_descriptor_encoding():
    writer = _ClassicPapWriter(version=(12, 5, 0), use64b_shadow=True)
    kernel = _classic_kernel()
    tpa, tpb = _tensor_parameters()

    module = writer.setupPrefetchAcrossPersistentLoads(kernel, tpa, tpb)
    items = _module_items(module)

    assert _instruction_indices(items, kwa_module.SMovB64, src_contains="ShadowLimitA+0")
    assert _instruction_indices(items, kwa_module.SMovB64, dst_contains="ShadowLimitA+0")
    assert len(_instruction_indices(items, kwa_module.SLShiftRightB32, dst_contains="Srd")) == 2


def test_classic_pap_shiftptr_refreshes_and_restores_gro_for_next_tile_loads():
    writer = _ClassicPapWriter()
    kernel = _classic_kernel(EdgeType="ShiftPtr")
    tpa, tpb = _tensor_parameters()

    module = writer.setupPrefetchAcrossPersistentLoads(kernel, tpa, tpb)
    items = _module_items(module)
    gro_snapshot_bases = {
        tag: base for base, _, tag in writer.vgprPool.checked_out if tag.endswith("GROSnapshot")
    }

    for tc in ("A", "B"):
        snapshot_base = gro_snapshot_bases["PAP%sGROSnapshot" % tc]
        checkpoint = _instruction_index(
            items,
            kwa_module.VMovB32,
            "v%u" % snapshot_base,
            "v[vgprGlobalReadOffset%s+0]" % tc,
        )
        refresh = _module_index(items, "lwaTileAssignment_%s" % tc)
        first_load = _module_index(items, "globalReadDo_%s" % tc)
        restore = _instruction_index(
            items,
            kwa_module.VMovB32,
            "v[vgprGlobalReadOffset%s+0]" % tc,
            "v%u" % snapshot_base,
        )

        assert checkpoint < refresh
        assert refresh < first_load
        assert first_load < restore

    gro_snapshots = [tag for _, _, tag in writer.vgprPool.checked_out if tag.endswith("GROSnapshot")]
    assert gro_snapshots == ["PAPAGROSnapshot", "PAPBGROSnapshot"]
    assert len(writer.vgprPool.checked_in) == len(gro_snapshots)


def test_classic_pap_saves_direct_to_lds_bank_state_after_priming():
    writer = _ClassicPapWriter(version=(9, 5, 0))
    kernel = _classic_kernel(DirectToLdsA=True)
    tpa, tpb = _tensor_parameters()

    module = writer.setupPrefetchAcrossPersistentLoads(kernel, tpa, tpb)
    items = _module_items(module)
    primed = _instruction_index(items, kwa_module.SMovB32, "s[sgprSkPrefetchPrimed]", "1")
    save_lds_bank = _module_index(items, "papDtlSaveLdsBank")

    assert primed < save_lds_bank
    assert writer.states.ldsTensorTokenIdx == writer.states.memTokenLdsBuffer1


def test_classic_pap_checkpoints_loop_counters_in_vgprs_around_next_tile_recount(monkeypatch):
    writer, items = _prefetch_across_persistent(monkeypatch)

    loop_vgpr = next(base for base, _, tag in writer.vgprPool.checked_out if tag == "PAP loop counters")
    orig_loop_vgpr = loop_vgpr + 1
    loop_checkpoint = _instruction_index(items, kwa_module.VMovB32, "v%u" % loop_vgpr, "s[sgprLoopCounterL]")
    orig_loop_checkpoint = _instruction_index(items, kwa_module.VMovB32, "v%u" % orig_loop_vgpr, "s[sgprOrigLoopCounter]")
    calculate_loop_num_iter = _module_index(items, "calculateLoopNumIter")
    setup_pap_loads = _module_index(items, "setupPrefetchAcrossPersistentLoads")
    loop_restore = _instruction_index(items, kwa_module.VReadfirstlaneB32, "s[sgprLoopCounterL]", "v%u" % loop_vgpr)
    orig_loop_restore = _instruction_index(items, kwa_module.VReadfirstlaneB32, "s[sgprOrigLoopCounter]", "v%u" % orig_loop_vgpr)

    assert loop_checkpoint < calculate_loop_num_iter
    assert orig_loop_checkpoint < calculate_loop_num_iter
    assert calculate_loop_num_iter < setup_pap_loads
    assert setup_pap_loads < loop_restore
    assert loop_restore < orig_loop_restore
    assert writer.vgprPool.checked_in == [loop_vgpr]


def test_dp_only_pap_skips_loop_counter_checkpoint(monkeypatch):
    # DP-only StreamK keeps LoopCounter/OrigLoopCounter constant (idempotent
    # recompute, PAP never runs on the last tile), so prefetchAcrossPersistent
    # skips the 2-VGPR checkpoint/restore entirely
    # (KernelWriterAssembly: snapshotLoopCounter = not StreamKForceDPOnly).
    writer, items = _prefetch_across_persistent(monkeypatch, StreamKForceDPOnly=1)

    assert not any(tag == "PAP loop counters" for _, _, tag in writer.vgprPool.checked_out)
    assert not _instruction_indices(items, kwa_module.VReadfirstlaneB32, dst_contains="sgprLoopCounterL")
    assert not _instruction_indices(items, kwa_module.VReadfirstlaneB32, dst_contains="sgprOrigLoopCounter")
    # the next-tile prefetch setup still runs; only the loop-counter snapshot is skipped
    _module_index(items, "setupPrefetchAcrossPersistentLoads")


def test_classic_pap_can_skip_internal_barrier_after_caller_sync(monkeypatch):
    _, with_barrier = _prefetch_across_persistent(monkeypatch, skip_barrier=False)
    _, without_barrier = _prefetch_across_persistent(monkeypatch, skip_barrier=True)

    assert _instruction_indices(with_barrier, kwa_module.SBarrier)
    assert not _instruction_indices(without_barrier, kwa_module.SBarrier)


@pytest.mark.parametrize(
    "space_filling_algo, remap_name, expected_transform_levels",
    [
        pytest.param([], "DefaultWGM", -1, id="default_wgm"),
        pytest.param([{"foo": "bar"}], "SpaceFillingCurveWalk", 1, id="space_filling_wgm"),
    ],
)
def test_streamk_pap_next_tile_setup_applies_wgm_remap(
    monkeypatch,
    space_filling_algo,
    remap_name,
    expected_transform_levels,
):
    import Tensile.Components.WorkGroupMappingAlgos as wgm_algos

    monkeypatch.setattr(
        wgm_algos,
        remap_name,
        lambda writer, kernel, sgpr_wgm: _module_with_comment(
            remap_name, "unit: %s remap" % remap_name
        ),
    )

    writer = _streamk_wgm_writer()
    module = _streamk_with_stubbed_tile_indexing().prefetchAcrossPersistentSetupNextTile(
        writer,
        {"SpaceFillingAlgo": space_filling_algo},
        {"tensorChar": "A"},
        {"tensorChar": "B"},
    )
    items = _module_items(module)

    tile_index = _module_index(items, "skTileIndex")
    index_to_wg = _module_index(items, "skIndexToWG")
    wgm_remap = _module_index(items, remap_name)
    assert tile_index < index_to_wg
    assert index_to_wg < wgm_remap
    assert writer.states.WGMTransformLevels == expected_transform_levels