From cd34d2f099a32eab5806981cffbee00804df9247 Mon Sep 17 00:00:00 2001 From: The-Monk Date: Thu, 23 Jul 2026 23:02:34 -0400 Subject: [PATCH] ck_tile: fix 2:4-sparse SWMMAC correctness on gfx1201/RDNA4 (3 bugs) + repro Fixes three bugs in the sparse MmaOpFamily compress path that produce wrong results on gfx1201 (v_swmmac_*_iu4): phantom fill on <2-nonzero 2:4 groups, pk_int4_t group-of-4 byte miscount (OOB), and idx-metadata ordering mismatch. Adds a standalone repro that fails on unpatched CK (max_abs_err=352) and passes with the fix (max_abs_err=0). Refs #3753. --- .../arch/mma/sparse/sparse_mma_pipeline.hpp | 17 +- .../arch/mma/sparse/sparse_transforms.hpp | 236 ++++++++++-- .../sparse_swmmac_correctness_repro.cpp | 336 ++++++++++++++++++ 3 files changed, 551 insertions(+), 38 deletions(-) create mode 100644 test/ck_tile/gfx1201_sparse_swmmac/sparse_swmmac_correctness_repro.cpp diff --git a/include/ck_tile/core/arch/mma/sparse/sparse_mma_pipeline.hpp b/include/ck_tile/core/arch/mma/sparse/sparse_mma_pipeline.hpp index b7a1ecb31f..b19284033b 100644 --- a/include/ck_tile/core/arch/mma/sparse/sparse_mma_pipeline.hpp +++ b/include/ck_tile/core/arch/mma/sparse/sparse_mma_pipeline.hpp @@ -203,7 +203,18 @@ struct SparseMmaPipeline : public MmaPipelineBase::vector_size; static constexpr index_t ExternalAVecSize = InternalAVecSize * MmaOp::kCompressionRatio; static constexpr index_t TotalUncompressedElems = FragsM * FragsK * ExternalAVecSize; - static constexpr index_t TotalCompressedElems = FragsM * FragsK * InternalAVecSize; + // Stage-17c fix (local, not upstreamed): the sparsity idx metadata is one + // 2-bit field per LOGICAL compressed element, not per physical storage + // element -- these differ for packed sub-byte A types (pk_int4_t, + // APackedSize=2). Confirmed against the RDNA4 ISA (Sec 7.12.3 Structured + // Sparse Matrices): "Wave32: each lane has 8 index values per lane" for a + // K=32 iu4 tile, where InternalAVecSize (physical pk_int4_t BYTE count + // per lane) is only 4 -- the ISA's own count is exactly + // InternalAVecSize * APackedSize, confirming idx tracks NIBBLES, not + // BYTES, for packed A. For APackedSize==1 (int8/fp8/fp16/...) this is a + // no-op (matches the original formula exactly). + static constexpr index_t TotalCompressedElems = + FragsM * FragsK * InternalAVecSize * MmaOp::APackedSize; // Variable-length idx type for the whole wave-tile (spans multiple int32_t words if needed) static constexpr index_t IdxNumWords = sparse::detail::idx_words_needed; @@ -257,7 +268,7 @@ struct SparseMmaPipeline : public MmaPipelineBase( + sparse::detail::extract_fragment_idx( idx, bm, bk)); } } @@ -275,7 +286,7 @@ struct SparseMmaPipeline : public MmaPipelineBase( + sparse::detail::extract_fragment_idx( idx, bm, bk)); } } diff --git a/include/ck_tile/core/arch/mma/sparse/sparse_transforms.hpp b/include/ck_tile/core/arch/mma/sparse/sparse_transforms.hpp index a0d55ce468..e996793a31 100644 --- a/include/ck_tile/core/arch/mma/sparse/sparse_transforms.hpp +++ b/include/ck_tile/core/arch/mma/sparse/sparse_transforms.hpp @@ -8,6 +8,7 @@ #include "ck_tile/core/config.hpp" #include "ck_tile/core/numeric/ext_vector_base.hpp" #include "ck_tile/core/numeric/integer.hpp" +#include "ck_tile/core/numeric/pk_int4.hpp" #include "ck_tile/core/numeric/vector_type.hpp" #include "ck_tile/core/utility/bit_cast.hpp" #include "ck_tile/core/utility/functional.hpp" @@ -55,46 +56,202 @@ struct SparseIdxPack * If fewer than CompressedSize non-zeros are found, remaining fields * default to 2 (see below). */ -template +// Stage-17c fix (local, not upstreamed): packed-sub-byte awareness. +// +// CompressedSize here is measured in PHYSICAL storage elements of ADataType +// (e.g. whole pk_int4_t BYTES, each holding 2 packed 4-bit values). The +// original code indexed a_vec[] and ran its "group of 4 / keep <=2" scan at +// that PHYSICAL granularity -- correct for PackedSize==1 (int8/fp8/fp16/...), +// but WRONG for PackedSize>1 (pk_int4_t): structured 2:4 sparsity is defined +// over LOGICAL elements (4-bit values), not physical storage words. Verified +// against the RDNA4 ISA (Sec 7.12.3, "Wave32: each lane has 8 index values +// per lane" for a K=32 iu4 tile whose physical AVecType is only 4 BYTES -- +// i.e. one idx field per LOGICAL nibble, not per physical byte) and via a +// device-side dump: real per-nibble 2:4 data makes EVERY physical byte in a +// group of 4 bytes "nonzero" (since a byte is nonzero whenever either of its +// two nibbles survived), systematically exceeding the old code's "<=2 +// nonzero per group of 4" assumption and writing out of bounds into +// nonzero_elems[2]/[3] (undefined behavior -- observed as register aliasing +// that silently corrupted the first two survivors on real hardware). +// Stage-17c fix (local, not upstreamed): LogicalADataType defaults to +// ADataType (preserves old behavior for every existing/non-packed caller), +// but callers that know the TRUE logical element type of a packed vector +// (e.g. pk_int4_t, whose physical ext_vector storage is represented as plain +// int8_t -- see vector_traits>::scalar_type in +// vector_type.hpp, which deliberately erases pk_int4_t to int8_t because +// clang's ext_vector_type attribute needs a true scalar) can pass it +// explicitly so PackedSize is read from the REAL logical type, not the +// erased physical one. Without this, `ADataType` here is ALWAYS int8_t for +// pk_int4_t vectors (confirmed: the packedness tag is lost one layer up, in +// SparseCompressTransform::execExtVec, well before reaching this function -- +// an earlier version of this fix that only checked +// numeric_traits::PackedSize here had NO EFFECT for exactly this +// reason, verified by a before/after diff showing byte-identical output). +template static CK_TILE_DEVICE auto compress_a_impl(AVec& a_vec) { - static constexpr index_t NumIdxWords = idx_words_needed; - // idx holds one 2-bit index per output element (total CompressedSize entries), - // packed across NumIdxWords int32_t words. - // It is initialized to the pattern 0b10 for every field. This matches - // what the hardware expects when there are fewer than two non-zero values - // in a 4-element group - the unused output is treated as coming from slot 2. - // The loop below will clear and set each field as real non-zeros are seen. + static constexpr index_t PackedSize = numeric_traits::PackedSize; + // LOGICAL field count: equals CompressedSize when PackedSize==1 (fully + // backward compatible with the original formula), doubles it for + // pk_int4_t-style 2-packed types. + static constexpr index_t LogicalCompressedSize = CompressedSize * PackedSize; + static constexpr index_t NumIdxWords = idx_words_needed; + + // idx holds one 2-bit index per LOGICAL output element (total + // LogicalCompressedSize entries), packed across NumIdxWords int32_t + // words. Initialized to the pattern 0b10 for every field -- see below. SparseIdxPack idx{}; - static_for<0, CompressedSize, 1>{}([&](auto k) { + static_for<0, LogicalCompressedSize, 1>{}([&](auto k) { constexpr uint32_t bit_pos = static_cast(k) * 2u; constexpr uint32_t word = bit_pos / 32u; constexpr uint32_t shift = bit_pos % 32u; idx.words[word] |= static_cast(2u << shift); }); - static_for<0, CompressedSize / 2, 1>{}([&](auto i) { - ADataType nonzero_elems[2] = {a_vec[i * 4 + 2], a_vec[i * 4 + 3]}; - int32_t non_zero_pos = 0; - - static_for<0, 4, 1>{}([&](auto j) { - if(static_cast(a_vec[i * 4 + j]) != 0.0f) - { - nonzero_elems[non_zero_pos] = a_vec[i * 4 + j]; - // clear the two-bit field for this output and insert j - const uint32_t field_idx = - static_cast(i) * 2u + static_cast(non_zero_pos); - const uint32_t bit_pos = field_idx * 2u; - const uint32_t word = bit_pos / 32u; - const uint32_t shift = bit_pos % 32u; - idx.words[word] &= ~static_cast(0b11u << shift); - idx.words[word] |= static_cast(static_cast(j) << shift); - ++non_zero_pos; - } + // Shared idx-field writer (used by both the scalar and packed paths). + auto set_idx_field = [&](uint32_t field_idx, uint32_t j_value) { + const uint32_t bit_pos = field_idx * 2u; + const uint32_t word = bit_pos / 32u; + const uint32_t shift = bit_pos % 32u; + idx.words[word] &= ~static_cast(0b11u << shift); + idx.words[word] |= static_cast(j_value << shift); + }; + + if constexpr(PackedSize == 1) + { + // Original scalar path, byte-granular groups of 4 -- unchanged. + static_for<0, CompressedSize / 2, 1>{}([&](auto i) { + // Stage-17b fix (local, not upstreamed): a group with fewer than + // 2 real non-zeros leaves one (or both) output slots unassigned + // by the scan below; its idx stays at the default 2 (see + // comment above). For that default idx to be safe REGARDLESS of + // what value the other slot claims for position 2, an + // unassigned slot's VALUE must be a true zero -- not copied + // from some fixed input position, because that position can + // itself be the group's real (only) survivor. The original + // code read defaults from {a[2], a[3]}: this is only safe when + // position 3 is *guaranteed* zero (true for the canonical + // "keep slots 0,2" synthetic pattern) -- for genuinely + // arbitrary-position 2:4 data (e.g. groups whose sole survivor + // sits at position 3, or, subtly, at position 2 -- both occur + // in real Quark-quantized weights, which are only 97.7% + // exactly-2:4), that default silently duplicates a real value + // onto a mismatched idx (measured: {0,0,0,44} reconstructed as + // {0,0,44,44}, or in the slot0-collides-with-slot1 direction, a + // legitimate lone survivor at position 2 got double-counted + // against B, producing small K-independent errors even on the + // "canonical" synthetic pattern whenever its random fill + // happened to leave slot 0 zero and slot 2 as the true lone + // survivor). An always-zero default is correct in every case: + // 0/1/2-nonzero groups all still reconstruct exactly, and any + // idx collision on an unassigned slot contributes nothing. + ADataType nonzero_elems[2] = {static_cast(0), static_cast(0)}; + int32_t non_zero_pos = 0; + + static_for<0, 4, 1>{}([&](auto j) { + if(static_cast(a_vec[i * 4 + j]) != 0.0f) + { + nonzero_elems[non_zero_pos] = a_vec[i * 4 + j]; + set_idx_field(static_cast(i) * 2u + static_cast(non_zero_pos), + static_cast(j)); + ++non_zero_pos; + } + }); + a_vec[i * 2] = nonzero_elems[0]; + a_vec[i * 2 + 1] = nonzero_elems[1]; }); - a_vec[i * 2] = nonzero_elems[0]; - a_vec[i * 2 + 1] = nonzero_elems[1]; - }); + } + else + { + static_assert(PackedSize == 2 && std::is_same_v, + "compress_a_impl packed path only implements pk_int4_t (2x4-bit) currently"); + // Packed-nibble path (pk_int4_t): each iteration consumes 2 + // PHYSICAL input bytes (= 4 LOGICAL nibbles, one real 2:4 group) and + // produces 1 PHYSICAL output byte (the <=2 survivor nibbles, + // packed). CompressedSize physical output bytes -> CompressedSize + // iterations (unlike the scalar path's CompressedSize/2, because + // here one physical byte carries 2 logical elements). + // + // Nibble ordering matches pk_int4_t's own CK_TILE_USE_PK4_LAYOUT_SHUFFLE + // convention (on by default, see pk_int4.hpp): "logical element 0" + // of a packed byte is the HIGH nibble, "element 1" is the LOW + // nibble -- so within a group of 4 logical positions (j=0..3) built + // from 2 consecutive physical bytes (byte0=a_vec[2i], byte1=a_vec[2i+1]): + // j=0 -> byte0 high, j=1 -> byte0 low, j=2 -> byte1 high, j=3 -> byte1 low. + // Raw byte-pointer view of a_vec: clang's native ext_vector_type + // operator[] does not support a WRITE with a single (non-scaled) + // compile-time-constant index in this context (only the scalar + // path's `a_vec[i*2]=...` pattern compiles) -- sidestep entirely by + // reinterpreting a_vec as a flat uint8_t* for both the packed reads + // and the packed write below (same technique already validated in + // this Stage-17c audit's own probe fill code). + uint8_t * a_bytes = reinterpret_cast(&a_vec); + static_for<0, CompressedSize, 1>{}([&](auto i) { + const uint8_t byte0 = a_bytes[2 * static_cast(i) + 0]; + const uint8_t byte1 = a_bytes[2 * static_cast(i) + 1]; + auto get_nibble = [](uint8_t byte, bool high) -> int8_t { + uint8_t nib = high ? ((byte >> 4) & 0x0Fu) : (byte & 0x0Fu); + int8_t val = static_cast(nib); + if(val & 0x08) val = static_cast(val | 0xF0); // sign-extend 4-bit + return val; + }; + const int8_t nib[4] = { + get_nibble(byte0, true), get_nibble(byte0, false), + get_nibble(byte1, true), get_nibble(byte1, false)}; + + // Same true-zero-default fix as the scalar path (Stage-17b), + // applied at nibble granularity. + // + // Stage-17c CLOSURE fix: the swmmac IU4 hardware reads the two + // idx fields of a group SWAPPED relative to which compressed + // value they govern, AND applies a "XOR 1" transform to the + // field's raw 2-bit value to get the real reconstructed + // position. Empirically determined (device measurement, not + // derived from the ISA doc alone -- the general Sec 7.12.3 + // pseudocode does not capture this IU4-specific detail) via a + // well-formed (always-exactly-2-survivor, no default-idx + // ambiguity) position-pair sweep on real hardware: for a group + // whose HIGH-nibble survivor (my slot 0, found first in scan + // order) sat at true position j0 and LOW-nibble survivor (slot + // 1, found second) sat at true position j1, the hardware + // reconstructs correctly ONLY when field(i*2+1) [normally + // "idx1"] is written as (j0 XOR 1) and field(i*2+0) [normally + // "idx0"] is written as (j1 XOR 1) -- i.e. slot0's position goes + // into the OTHER field, XORed with 1, and vice versa. Verified + // against all 6 position-pair combinations (C(4,2)) with exact + // arithmetic match (predicted C01 == measured C01 in every + // case) before being applied here. This ALSO fully explains an + // earlier single-survivor sweep's "always reconstructs at + // position 3" finding: the untouched default field (value 2, + // unused-slot placeholder) XORed with 1 gives 3, constant, + // regardless of the real survivor's true position -- exactly + // what was observed. + int8_t survivor[2] = {0, 0}; + int32_t non_zero_pos = 0; + static_for<0, 4, 1>{}([&](auto j) { + if(nib[static_cast(j)] != 0) + { + survivor[non_zero_pos] = nib[static_cast(j)]; + // slot 0 (HIGH nibble) -> field i*2+1; slot 1 (LOW nibble) -> field i*2+0. + const uint32_t target_field = + static_cast(i) * 2u + (1u - static_cast(non_zero_pos)); + set_idx_field(target_field, static_cast(j) ^ 1u); + ++non_zero_pos; + } + }); + // Re-pack: element0 (survivor[0]) -> HIGH nibble, element1 + // (survivor[1]) -> LOW nibble, matching the same shuffle + // convention used to unpack above (and used by the register-map + // fill on the host/caller side). + const uint8_t out_byte = static_cast( + ((static_cast(survivor[0]) & 0x0Fu) << 4) | + (static_cast(survivor[1]) & 0x0Fu)); + // Output bytes are always written at index <= input bytes read + // (i <= 2i), so writing through the same raw-byte view in + // increasing i order never clobbers not-yet-read input data. + a_bytes[static_cast(i)] = out_byte; + }); + } return idx; } @@ -146,7 +303,14 @@ struct SparseCompressTransform { // This function takes A in uncompressed form as a big ext_vector, and returns a // compressed ext_vector. - template + // Stage-17c fix (local, not upstreamed): added LogicalADataType, defaulted + // to vector_traits::scalar_type so every EXISTING call site + // (including direct execExtVec() callers that don't know about packing) + // behaves exactly as before. exec() below (the real Pipeline::exec() call + // path) passes the TRUE logical element type explicitly, which is the + // only way compress_a_impl can see PackedSize > 1 (see its own comment). + template >::scalar_type> CK_TILE_DEVICE static decltype(auto) execExtVec(VecType& v) { using VecTraits = vector_traits>; @@ -155,12 +319,13 @@ struct SparseCompressTransform static constexpr index_t CompressedSize = VecN / CompressionRatio; using VecCompressed = ext_vector_t; using IdxType = - sparse::detail::SparseIdxPack>; + sparse::detail::SparseIdxPack::PackedSize>>; static_assert(VecN % CompressionRatio == 0, "VecN must be divisible by CompressionRatio"); static_assert(CompressedSize > 0, "CompressedSize must be > 0"); - auto idx = sparse::detail::compress_a_impl(v); + auto idx = sparse::detail::compress_a_impl(v); return std::tuple(*ck_tile::bit_cast(&v), idx); } @@ -174,11 +339,12 @@ struct SparseCompressTransform CK_TILE_DEVICE static decltype(auto) exec(ATensor& a_tensor) { // Properties of ATensor as a big ext vector. - using ADataType = typename ATensor::DataType; + using ADataType = typename ATensor::DataType; // TRUE logical type (e.g. pk_int4_t) constexpr index_t VecN = ATensor::get_thread_buffer_size(); using VecType = ext_vector_t; - return execExtVec(a_tensor.get_thread_buffer().template get_as().template at<0>()); + return execExtVec( + a_tensor.get_thread_buffer().template get_as().template at<0>()); } }; diff --git a/test/ck_tile/gfx1201_sparse_swmmac/sparse_swmmac_correctness_repro.cpp b/test/ck_tile/gfx1201_sparse_swmmac/sparse_swmmac_correctness_repro.cpp new file mode 100644 index 0000000000..29cfed1ae7 --- /dev/null +++ b/test/ck_tile/gfx1201_sparse_swmmac/sparse_swmmac_correctness_repro.cpp @@ -0,0 +1,336 @@ +// Stage-1 (2:4-sparse CK library-grade GEMM feasibility, 2026-07-20): +// warp-tile correctness proof for CK's ck_tile::SparseMmaPipeline +// (int8_t x int8_t -> int32_t, 16x16x32 WMMA SWMMAC, gfx1201), extending +// Stage-8's "raw atom executes" finding to "a validated PIPELINE produces +// correct 2:4-sparse GEMM results, checked against a CPU reference". +// +// This is a non-gtest standalone port of CK's OWN test methodology +// (test/ck_tile/core/arch/mma/pipeline/test_amdgcn_sparse_mma.cpp + +// pipeline_tests_helper.hpp), NOT independently re-derived -- reusing +// AMD's validated register-mapping logic (TileDistrEncRegMap) is much +// lower-risk than hand-deriving the per-lane sparse layout from scratch, +// and is exactly the "library value" this Stage-1 gate is asking about: +// does CK's own machinery work correctly for the SPARSE MmaOpFamily on +// gfx1201, not just compile. +// +// Data flow tested: dense A (int8, with 2:4 structured zeros already +// applied per-4-group -- CK's own convention: A is provided UNCOMPRESSED, +// SparseCompressTransform runs ON THE FLY inside Pipeline::exec()) x dense +// B -> int32 C, ONE warp, ONE 16x16x{32,64,128} tile. +#include "ck_tile/core/arch/arch.hpp" +#include "ck_tile/core/arch/mma/mma.hpp" +#include "ck_tile/core/arch/mma/sparse/sparse_mma_pipeline.hpp" +#include "ck_tile/core/arch/mma/utility/tile_distribution_encoding_calculator.hpp" +#include "ck_tile/core/arch/mma/utility/tile_distribution_encoding_register_mapper.hpp" +#include "ck_tile/core/numeric/integer.hpp" +#include "ck_tile/core/numeric/type_convert.hpp" +#include "ck_tile/core/numeric/vector_type.hpp" +#include "ck_tile/host/hip_check_error.hpp" +#include "ck_tile/host/kernel_launch.hpp" +#include "ck_tile/host/stream_config.hpp" + +#include + +#include "real_a_tile.h" // REAL_A_16x128: Quark int8 2:4 weight tile (arbitrary-position zeros) + +#include +#include +#include +#include +#include +#include +#include + +using namespace ck_tile; +using namespace ck_tile::core::arch; +using namespace ck_tile::core::arch::mma; + +// CRITICAL fix (found the hard way, matches this fork's own already-known +// __GFX12__/RDNA4 macro-visibility lesson): SparseMmaPipeline's DEFAULT +// CompilerTarget template param (getCMakeCompilerTarget()) resolves via +// the __gfx1201__ preprocessor macro, which clang's HIP frontend only +// predefines during the DEVICE compilation pass of a TU -- NEVER during +// the HOST pass. Any HOST-side code (like this file's run_test(), which +// computes sizeof(AWarpTensor) etc. to size/fill host buffers before +// upload) that instantiates SparseMmaPipeline<...> with the DEFAULTED +// target silently gets HOST-context resolution (__gfx1201__ undefined -> +// CK_TILE_ARCH_GFX1201=false -> falls through to an UNSUPPORTED dummy +// MmaOp with kM=kN=kK=1, kCompressionRatio=1) -- a COMPLETELY DIFFERENT, +// much smaller, wrong-shaped type than what the DEVICE kernel body (same +// source, device pass, __gfx1201__ defined) resolves. Host buffer sizing +// mismatched against device tensor layout -> heap corruption (confirmed: +// this exact bug crashed the first version of this probe with a glibc +// malloc assertion failure). FIX: specify CompilerTarget EXPLICITLY +// (a pure compile-time type, not macro-dependent) everywhere -- exactly +// what CK's OWN test file does (CompilerTargetGfx950 = decltype(...)), +// which I initially skipped by relying on the convenience default. Not a +// CK bug -- a usage requirement whenever a Pipeline type crosses the +// host/device boundary in one TU. +using Gfx1201Target = decltype(make_amdgcn_gfx12_target()); + +// --- Reference (CPU, int64 accumulate -- exact for int8 x int8 up to K=~2^47, no overflow risk) --- +static void reference_matmul_i8(std::vector & C, const std::vector & A, + const std::vector & B, uint32_t M, uint32_t N, uint32_t K) { + for (uint32_t m = 0; m < M; ++m) { + for (uint32_t n = 0; n < N; ++n) { + int64_t acc = 0; + for (uint32_t k = 0; k < K; ++k) { + acc += (int64_t) A[m * K + k] * (int64_t) B[k * N + n]; + } + C[m * N + n] = (int32_t) acc; + } + } +} + +// Apply 2:4 sparsity: every group of 4 consecutive K elements keeps slots +// 0,2, zeros slots 1,3 (matches CK's own test convention exactly). +static void apply_sparse_pattern(std::vector & A, uint32_t M, uint32_t K) { + for (uint32_t m = 0; m < M; ++m) { + for (uint32_t k = 0; k < K; k += 4) { + if (k + 1 < K) A[m * K + k + 1] = 0; + if (k + 3 < K) A[m * K + k + 3] = 0; + } + } +} + +template +struct SparseGemmKernel { + static constexpr int kBlockSize = 32; // wave32 on gfx1201 + + __device__ void operator()(const void * a_per_lane, const void * b_per_lane, void * c_per_lane) const { + using ATensor = typename Pipeline::AWarpTensor; + using BTensor = typename Pipeline::BWarpTensor; + using CTensor = typename Pipeline::CWarpTensor; + + const uint32_t lane = threadIdx.x; + + ATensor a; + BTensor b; + CTensor c; + __builtin_memcpy(&a, static_cast(a_per_lane) + lane * sizeof(ATensor), sizeof(ATensor)); + __builtin_memcpy(&b, static_cast(b_per_lane) + lane * sizeof(BTensor), sizeof(BTensor)); + __builtin_memset(&c, 0, sizeof(CTensor)); + + if constexpr (MmaOpTraits::IsSupported) { +#ifdef DEBUG_DEVICE + if (lane == 0) { + ATensor a_dbg = a; + constexpr index_t VecN = ATensor::get_thread_buffer_size(); + using RawVec = ext_vector_t; + auto & raw = a_dbg.get_thread_buffer().template get_as().template at<0>(); + printf(" [DEV lane0] raw uncompressed a_vec (VecN=%d):", (int)VecN); + for (index_t i = 0; i < VecN; ++i) printf(" %d", (int)raw[i]); + printf("\n"); + auto ab_pair = Pipeline::ATransform::execExtVec(raw); + auto & compressed = std::get<0>(ab_pair); + auto & idxpk = std::get<1>(ab_pair); + using CompVecTraits = vector_traits>; + printf(" [DEV lane0] compressed a_vec (size=%d):", (int)CompVecTraits::vector_size); + for (index_t i = 0; i < CompVecTraits::vector_size; ++i) printf(" %d", (int)compressed[i]); + printf("\n [DEV lane0] idx words:"); + for (auto w : idxpk.words) printf(" 0x%08x", (unsigned)w); + printf("\n"); + } +#endif + Pipeline::exec(a, b, c); + __builtin_memcpy(static_cast(c_per_lane) + lane * sizeof(CTensor), &c, sizeof(CTensor)); + } + } +}; + +// Per-lane layout construction, ported from pipeline_tests_helper.hpp's +// fill_a_fragments/fill_b_fragments/extract_c_matrix (AMD's own validated +// register-map logic -- TileDistrEncRegMap -- reused verbatim, not +// re-derived). +// FIX (Stage-17b): the old version rebuilt the register map from a *bare* +// TileDistrEncCalc (all defaults: UncompressedA=false, kIter=1), which +// yields the COMPRESSED-A encoding (K dimension already halved by +// kCompressionRatio) -- the layout the *intrinsic* consumes post-compression, +// not the layout the *pipeline* expects the caller to supply pre-compression. +// It then hand-expanded that compressed coordinate back out via +// `k_local = k_compressed * kCompressionRatio + sub_pos`, which silently +// assumes the compressed-K stride advances in lockstep, whole-4-group chunks +// -- true only for the "zeros at slots 1,3" synthetic pattern, not for +// arbitrary-position 2:4 data (where compress_a_impl's own group-of-4 scan +// must see the TRUE contiguous [k*4 .. k*4+3] slice). +// +// SparseMmaPipeline already computes and PUBLICLY EXPOSES exactly the right +// encoding for this: `Pipeline::AWarpDstrEncoding` is built internally via +// `EncCalc = TileDistrEncCalc` (see +// sparse_mma_pipeline.hpp). With UncompressedA=true the K sub-dimension is +// NOT divided by kCompressionRatio, so `calc_matrix_indices_from_lane_vector` +// returns the GLOBAL, uncompressed (m, k) coordinate directly -- already in +// the exact per-lane flat vector order that SparseCompressTransform::exec() +// (and hence compress_a_impl's real contiguous-4 grouping) expects. Using +// Pipeline's own type instead of re-deriving it removes the bug at the root: +// no more manual compression-ratio math, no bm/bk frag loop needed (the K +// sub-dim already spans the WHOLE WaveTileK once kIter=FragsK is baked in). +template +static void fill_a_fragments(typename Pipeline::AWarpTensor * a_per_lane, + const std::vector & A_matrix, uint32_t K, uint32_t waveSize) { + using ARegMap = TileDistrEncRegMap; + using AFragScalar = typename Pipeline::ADataType; + constexpr index_t a_vec_size = ARegMap::num_vector_items; // full uncompressed per-lane count + + for (uint32_t lane = 0; lane < waveSize; ++lane) { + auto * lane_a = reinterpret_cast(&a_per_lane[lane]); + for (index_t v = 0; v < a_vec_size; ++v) { + auto coords = ARegMap::calc_matrix_indices_from_lane_vector(lane, v); + uint32_t m_global = coords[0]; + uint32_t k_global = coords[1]; + lane_a[v] = static_cast(A_matrix[m_global * K + k_global]); +#ifdef DEBUG_FILL_A + if (m_global == 0 && K == 32) { + printf(" [DBG m=0] lane=%2u v=%2d -> k=%2u val=%d\n", lane, (int)v, k_global, (int)lane_a[v]); + } +#endif + } + } +} + +template +static void fill_b_fragments(typename Pipeline::BWarpTensor * b_per_lane, + const std::vector & B_matrix, uint32_t N, uint32_t waveSize) { + // Same root-cause class as A (see fill_a_fragments comment): a bare + // TileDistrEncCalc defaults to kIter=1, describing ONE MmaOp::kK + // fragment's worth of B. SparseMmaPipeline's real EncCalc uses kIter=FragsK + // (baked into Pipeline::BWarpDstrEncoding), so its K sub-dim already spans + // the WHOLE WaveTileK -- the old per-fragment bn/bk loop with frag_offset + // happened to reproduce the right flat order only in combination with the + // old (also-mismatched) A fill; once A is fixed to match Pipeline's true + // encoding, B must match it too or the two mismatches stop canceling out. + using BRegMap = TileDistrEncRegMap; + using BFragScalar = typename Pipeline::BDataType; + constexpr index_t b_vec_size = BRegMap::num_vector_items; // spans the full WaveTileK already + + for (uint32_t lane = 0; lane < waveSize; ++lane) { + auto * lane_b = reinterpret_cast(&b_per_lane[lane]); + for (index_t v = 0; v < b_vec_size; ++v) { + auto coords = BRegMap::calc_matrix_indices_from_lane_vector(lane, v); + uint32_t n_global = coords[0]; + uint32_t k_global = coords[1]; + lane_b[v] = static_cast(B_matrix[k_global * N + n_global]); + } + } +} + +template +static void extract_c_matrix(const typename Pipeline::CWarpTensor * c_per_lane, + std::vector & C_matrix, uint32_t N, uint32_t waveSize) { + // C's encoding doesn't depend on kIter (see get_cwarp_dstr_encoding()), so this + // was not actually part of the bug -- switched to Pipeline's own alias anyway + // for consistency/robustness with A and B above. + using CRegMap = TileDistrEncRegMap; + using CFragScalar = typename Pipeline::CDataType; + + constexpr uint32_t FragM = Pipeline::MmaOp::kM; + constexpr uint32_t FragN = Pipeline::MmaOp::kN; + constexpr uint32_t FragsM = Pipeline::FragsM; + constexpr uint32_t FragsN = Pipeline::FragsN; + constexpr index_t c_vec_size = CRegMap::num_vector_items; + + for (uint32_t lane = 0; lane < waveSize; ++lane) { + auto * lane_c = reinterpret_cast(&c_per_lane[lane]); + for (uint32_t bm = 0; bm < FragsM; ++bm) { + for (uint32_t bn = 0; bn < FragsN; ++bn) { + uint32_t frag_offset = (bm * FragsN + bn) * c_vec_size; + for (index_t v = 0; v < c_vec_size; ++v) { + auto coords = CRegMap::calc_matrix_indices_from_lane_vector(lane, v); + uint32_t m_local = coords[0]; + uint32_t n_local = coords[1]; + uint32_t m_global = bm * FragM + m_local; + uint32_t n_global = bn * FragN + n_local; + C_matrix[m_global * N + n_global] = static_cast(lane_c[frag_offset + v]); + } + } + } + } +} + +template +static bool run_test(const char * label) { + using Pipeline = SparseMmaPipeline; + using AWarpTensor = typename Pipeline::AWarpTensor; + using BWarpTensor = typename Pipeline::BWarpTensor; + using CWarpTensor = typename Pipeline::CWarpTensor; + + const uint32_t M = WaveTileM, N = WaveTileN, K = WaveTileK, waveSize = 32; + + std::mt19937 rng(42); + std::uniform_int_distribution dist(-8, 8); // small range, avoids int8 overflow noise + std::vector A(M * K), B(K * N); +#ifdef USE_REAL_TILE + // A = REAL Quark int8 2:4 weight tile (q_proj L0, arbitrary-position zeros). + // Tests whether CK's on-the-fly SparseCompressTransform handles genuine 2:4 + // (zeros anywhere in each 4-group). REAL_A_16x128 is row-major [16][128]. + for (uint32_t m = 0; m < M; ++m) + for (uint32_t k = 0; k < K; ++k) + A[m * K + k] = REAL_A_16x128[m * 128 + k]; + for (auto & v : B) v = (int8_t) dist(rng); +#else + // Synthetic control: zeros in slots 1,3 (CK's assumed kept-pattern 0,2). + for (auto & v : A) v = (int8_t) dist(rng); + for (auto & v : B) v = (int8_t) dist(rng); + apply_sparse_pattern(A, M, K); +#endif + + std::vector C_expected(M * N, 0), C_actual(M * N, 0); + reference_matmul_i8(C_expected, A, B, M, N, K); +#ifdef DEBUG_DEVICE + if (K == 32) { + printf(" [HOST] B column 0, k=0..31:"); + for (uint32_t k = 0; k < K; ++k) printf(" %d", (int)B[k * N + 0]); + printf("\n"); + } +#endif + + const size_t a_buf_size = waveSize * sizeof(AWarpTensor); + const size_t b_buf_size = waveSize * sizeof(BWarpTensor); + const size_t c_buf_size = waveSize * sizeof(CWarpTensor); + std::vector h_a(a_buf_size, 0), h_b(b_buf_size, 0), h_c(c_buf_size, 0); + + fill_a_fragments(reinterpret_cast(h_a.data()), A, K, waveSize); + fill_b_fragments(reinterpret_cast(h_b.data()), B, N, waveSize); + + void *d_a, *d_b, *d_c; + HIP_CHECK_ERROR(hipMalloc(&d_a, a_buf_size)); + HIP_CHECK_ERROR(hipMalloc(&d_b, b_buf_size)); + HIP_CHECK_ERROR(hipMalloc(&d_c, c_buf_size)); + HIP_CHECK_ERROR(hipMemcpy(d_a, h_a.data(), a_buf_size, hipMemcpyHostToDevice)); + HIP_CHECK_ERROR(hipMemcpy(d_b, h_b.data(), b_buf_size, hipMemcpyHostToDevice)); + HIP_CHECK_ERROR(hipMemset(d_c, 0, c_buf_size)); + + ck_tile::launch_kernel(ck_tile::stream_config{}, + ck_tile::make_kernel(SparseGemmKernel{}, dim3(1), dim3(waveSize), 0, d_a, d_b, d_c)); + HIP_CHECK_ERROR(hipDeviceSynchronize()); + + HIP_CHECK_ERROR(hipMemcpy(h_c.data(), d_c, c_buf_size, hipMemcpyDeviceToHost)); + extract_c_matrix(reinterpret_cast(h_c.data()), C_actual, N, waveSize); + + HIP_CHECK_ERROR(hipFree(d_a)); + HIP_CHECK_ERROR(hipFree(d_b)); + HIP_CHECK_ERROR(hipFree(d_c)); + + long max_abs_err = 0; + for (size_t i = 0; i < C_actual.size(); ++i) { + max_abs_err = std::max(max_abs_err, std::labs((long) C_actual[i] - (long) C_expected[i])); + } + const bool pass = (max_abs_err == 0); + printf("[%s] M=%u N=%u K=%u -> max_abs_err=%ld %s\n", label, M, N, K, max_abs_err, pass ? "PASS" : "FAIL"); + if (!pass) { + printf(" sample: C_expected[0]=%d C_actual[0]=%d\n", C_expected[0], C_actual[0]); + } + return pass; +} + +int main() { + printf("=== Stage-1: ck_tile::SparseMmaPipeline warp-tile correctness (gfx1201 WMMA SWMMAC) ===\n"); + bool all_pass = true; + all_pass &= run_test<16, 16, 32>("K32_single_frag"); + all_pass &= run_test<16, 16, 64>("K64_2frag"); + all_pass &= run_test<16, 16, 128>("K128_4frag"); + printf("=== %s ===\n", all_pass ? "ALL PASS" : "SOME FAILED"); + return all_pass ? 0 : 1; +}