kernels/rr: staged-image Swin layer engine prototypes
Per-layer staged-image engine for the eleven translated denoiser Swin layers: descriptors (verified offsets per phase), cut recipes that splice an isolated native phase into NVIDIA's PTX, wave templates, epilogue notes and GPU probes. The isolated phase-8 splice produces a byte-identical final 4K frame but is slower than the translated body (14.90 vs 6.59 ms) - a correctness experiment, not a shipping improvement; docs/rr-engine.md records the design and the next direct-converter win (B operands built from transposed accumulators).
This commit is contained in:
woutermans committed
1 parent
798e25611c
commit
0fe3a28cb7
58 files changed
+28330
-5
No files matched your search
@@ -0,0 +1,785 @@
|
||||
// Steps 1 and 2 of the engine's validation: e4m3 GEMMs in real shared memory, every HIP call
|
||||
// checked, compared against an independently computed host model.
|
||||
//
|
||||
// step 1 one k16 GEMM (the second k16 half is zero) -- operand form, lane/half indexing,
|
||||
// and the D layout the engine assumes
|
||||
// step 2 a 64-wide k chain as two k32 steps with the accumulator rounded to f16 between
|
||||
// them, which is what kernels/tex/enc0_tail.hip does (`acc[i] = (half_t)cf[i];
|
||||
// cf[i] = (float)acc[i];`)
|
||||
//
|
||||
// The host model is e4m3 decode plus a double-precision sum, written from the layout facts in
|
||||
// kernels/common/wmma_layout.h and the usage in kernels/tex/enc0_tail.hip; it shares no code
|
||||
// with the device side.
|
||||
//
|
||||
// Build and run:
|
||||
// hipcc --offload-arch=gfx1201 -DD4R_TEX_FP8 -DD4R_WMMA_LAYOUT=12 -O2 \
|
||||
// -o /tmp/engine_test kernels/tools/rrswin_engine_k16_test.cpp && /tmp/engine_test
|
||||
#include <hip/hip_runtime.h>
|
||||
#include <cstdio>
|
||||
#include <cstdint>
|
||||
#include <cmath>
|
||||
#include <cstring>
|
||||
#include <cstdlib>
|
||||
#include <vector>
|
||||
|
||||
#include "../rr/rrswin_engine.h"
|
||||
|
||||
#define CHECK(expr) \
|
||||
do \
|
||||
{ \
|
||||
hipError_t e_ = (expr); \
|
||||
if (e_ != hipSuccess) \
|
||||
{ \
|
||||
std::printf("HIP error at %s:%d: %s -> %s\n", __FILE__, __LINE__, #expr, \
|
||||
hipGetErrorString(e_)); \
|
||||
std::exit(2); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
// the documented enc0 weight layout, written here independently of the device header
|
||||
static int rrswin_waddr_host(int k, int n)
|
||||
{
|
||||
return 512 * (n >> 4) + 64 * (n & 7) + 16 * ((k >> 1) & 3) + 8 * ((n >> 3) & 1)
|
||||
+ 4 * ((k >> 3) & 1) + 2 * ((k >> 4) & 1) + (k & 1);
|
||||
}
|
||||
|
||||
static const int kRows = 16;
|
||||
static const int kChunk = 16; // bytes per k16 chunk
|
||||
static const int kChunks = 4; // k = 64 in four k16 chunks (= two k32 steps)
|
||||
static const int kRowBytes = kChunk * kChunks;
|
||||
static const int kSlots = 128; // prepared weight slots: 32 per k32 step (16 rows x 2 k16), two steps per GEMM
|
||||
|
||||
static uint8_t enc(double v)
|
||||
{
|
||||
const int sign = v < 0;
|
||||
v = std::fabs(v);
|
||||
if (v < 0.001953125)
|
||||
return (uint8_t)(sign ? 0x80 : 0x00);
|
||||
int e = 0;
|
||||
while (v >= 2.0) { v /= 2.0; ++e; }
|
||||
while (v < 1.0) { v *= 2.0; --e; }
|
||||
const int code = ((e + 7) << 3) | (int)std::lround((v - 1.0) * 8.0);
|
||||
return (uint8_t)(code | (sign ? 0x80 : 0));
|
||||
}
|
||||
|
||||
static double dec(uint8_t c)
|
||||
{
|
||||
const int sign = c & 0x80, e = (c >> 3) & 0xf, m = c & 7;
|
||||
const double mag = e == 0 ? (m / 8.0) * std::pow(2.0, -6.0) : (1.0 + m / 8.0) * std::pow(2.0, e - 7.0);
|
||||
return sign ? -mag : mag;
|
||||
}
|
||||
|
||||
// D VGPR i of half h is row i + 8h and lane l is column l & 15, so a value's (row, column)
|
||||
// fixes the lane and slot that must hold it.
|
||||
static void check(const char* what, const std::vector<float>& got, const std::vector<uint8_t>& w,
|
||||
const std::vector<uint8_t>& x, int kk, double tol)
|
||||
{
|
||||
int bad = 0, checked = 0;
|
||||
for (int l = 0; l < 32; ++l)
|
||||
{
|
||||
const int h = l >> 4, rcol = l & 15;
|
||||
for (int i = 0; i < 8; ++i)
|
||||
{
|
||||
const int row = i + 8 * h;
|
||||
// weight row `row`'s k16 chunk h of k32 step s is stored at slot 32*s + 16*h + row,
|
||||
// which is byte 16*slot -- NOT inside the same 16 bytes as the row's first chunk.
|
||||
auto wslot = [&](int s, int chunk, int r) { return (size_t)(32 * s + 16 * chunk + r); };
|
||||
auto dot = [&](int from, int to) {
|
||||
double s = 0.0;
|
||||
for (int c = from; c < to; ++c)
|
||||
s += dec(w[wslot(c / 32, (c % 32) / 16, row) * kChunk + (c % 16)]) *
|
||||
dec(x[(size_t)rcol * kRowBytes + c]);
|
||||
return s;
|
||||
};
|
||||
double ref = dot(0, kk);
|
||||
if (kk > 32)
|
||||
{
|
||||
// the caller's f16 rounding of the accumulator between the two k32 steps
|
||||
const double first = (double)(_Float16)dot(0, 32);
|
||||
ref = first + dot(32, kk);
|
||||
}
|
||||
const float g = got[i * 32 + l];
|
||||
++checked;
|
||||
if (std::fabs((double)g - ref) > tol * (1.0 + std::fabs(ref)))
|
||||
{
|
||||
if (bad < 6)
|
||||
std::printf(" lane %2d vgpr %d: D[row %2d][col %2d] got %+.6f ref %+.6f\n",
|
||||
l, i, row, rcol, g, ref);
|
||||
++bad;
|
||||
}
|
||||
}
|
||||
}
|
||||
std::printf("%s: checked %d, mismatched %d\n", what, checked, bad);
|
||||
if (bad)
|
||||
{
|
||||
// report which (row, column) each mismatching value *does* correspond to, which is
|
||||
// the mapping the engine would have to use
|
||||
int shown = 0;
|
||||
for (int l = 0; l < 32 && shown < 4; ++l)
|
||||
{
|
||||
for (int i = 0; i < 8 && shown < 4; ++i)
|
||||
{
|
||||
const float g = got[i * 32 + l];
|
||||
int best_row = -1, best_col = -1;
|
||||
double best = 1e30;
|
||||
for (int row = 0; row < 16; ++row)
|
||||
for (int col = 0; col < 16; ++col)
|
||||
{
|
||||
double ref = 0.0;
|
||||
for (int c = 0; c < kk; ++c)
|
||||
ref += dec(w[(size_t)(32 * (c / 32) + 16 * ((c % 32) / 16) + row) * kChunk + (c % 16)]) *
|
||||
dec(x[(size_t)col * kRowBytes + c]);
|
||||
if (std::fabs((double)g - ref) < best) { best = std::fabs((double)g - ref); best_row = row; best_col = col; }
|
||||
}
|
||||
if (best < tol)
|
||||
{
|
||||
std::printf(" lane %2d vgpr %d = D[row %2d][col %2d] (engine assumed row %2d col %2d)\n",
|
||||
l, i, best_row, best_col, i + 8 * (l >> 4), l & 15);
|
||||
++shown;
|
||||
}
|
||||
}
|
||||
}
|
||||
std::exit(1);
|
||||
}
|
||||
}
|
||||
|
||||
// The requantisation the device performs in swin_stage_operand, written independently here:
|
||||
// RNE satfinite e4m3 of an f16 pair, low half to the low byte.
|
||||
static uint8_t requant_code(uint16_t h)
|
||||
{
|
||||
const uint32_t mag = h & 0x7fffu;
|
||||
uint32_t normal = ((mag - 0x2000u + 0x3fu + ((mag >> 7) & 1u)) & 0xffffu) >> 7;
|
||||
const uint32_t e = mag >> 10;
|
||||
const uint32_t clamped = e < 1u ? 1u : (e > 8u ? 8u : e);
|
||||
const uint32_t shift = 16u - clamped;
|
||||
const uint32_t sig = (mag & 0x3ffu) | (e != 0 ? 0x400u : 0u);
|
||||
const uint32_t sub = (sig + (1u << (shift - 1u)) - 1u + ((sig >> shift) & 1u)) >> shift;
|
||||
uint32_t code = mag >= 0x2400u ? normal : sub;
|
||||
if (mag > 0x5f00u)
|
||||
code = 0x7eu;
|
||||
if (mag > 0x7c00u)
|
||||
code = 0x7fu;
|
||||
return (uint8_t)(((h >> 8) & 0x80u) | code);
|
||||
}
|
||||
|
||||
// the f16 bits the device would round to, then the e4m3 code
|
||||
static uint8_t requant_code_bytes(double v)
|
||||
{
|
||||
const _Float16 half = (_Float16)v;
|
||||
uint16_t bits;
|
||||
std::memcpy(&bits, &half, sizeof(bits));
|
||||
return requant_code(bits);
|
||||
}
|
||||
|
||||
static double requant(double v)
|
||||
{
|
||||
const _Float16 half = (_Float16)v;
|
||||
uint16_t bits;
|
||||
std::memcpy(&bits, &half, sizeof(bits));
|
||||
return dec(requant_code(bits));
|
||||
}
|
||||
|
||||
__global__ void step1_kernel(const uint8_t* x_img, const uint8_t* w_prep, float* out)
|
||||
{
|
||||
__shared__ uint8_t img[kRows * kRowBytes];
|
||||
for (int i = threadIdx.x; i < kRows * kRowBytes; i += blockDim.x)
|
||||
img[i] = x_img[i];
|
||||
__syncthreads();
|
||||
f8v c = {0, 0, 0, 0, 0, 0, 0, 0};
|
||||
c = rrswin::swin_gemm_tile(c, (const kslot*)w_prep, 0, 16, img, kRowBytes, 0);
|
||||
for (int i = 0; i < 8; ++i)
|
||||
out[i * 32 + threadIdx.x] = c[i];
|
||||
}
|
||||
|
||||
__global__ void step2_kernel(const uint8_t* x_img, const uint8_t* w_prep, float* out)
|
||||
{
|
||||
__shared__ uint8_t img[kRows * kRowBytes];
|
||||
for (int i = threadIdx.x; i < kRows * kRowBytes; i += blockDim.x)
|
||||
img[i] = x_img[i];
|
||||
__syncthreads();
|
||||
f8v c = {0, 0, 0, 0, 0, 0, 0, 0};
|
||||
c = rrswin::swin_gemm_tile(c, (const kslot*)w_prep, 0, 16, img, kRowBytes, 0);
|
||||
half_t acc[8];
|
||||
for (int i = 0; i < 8; ++i)
|
||||
acc[i] = (half_t)c[i];
|
||||
for (int i = 0; i < 8; ++i)
|
||||
c[i] = (float)acc[i];
|
||||
// the second k32 step: slots 32..63 and the image's chunks 2 and 3
|
||||
c = rrswin::swin_gemm_tile(c, (const kslot*)w_prep, 32, 16, img + 2 * kChunk, kRowBytes, 0);
|
||||
for (int i = 0; i < 8; ++i)
|
||||
out[i * 32 + threadIdx.x] = c[i];
|
||||
}
|
||||
|
||||
// Step 3: the accumulator-to-operand mapping. A first k32 GEMM produces D, D is staged as
|
||||
// the next phase's activation image with swin_stage_operand, and a second k32 GEMM consumes
|
||||
// it against a second weight slot. The host model re-does exactly that, requantising with
|
||||
// its own e4m3 encoder, so a wrong mapping (row/column/k-half) shows up as a mismatch.
|
||||
__global__ void step3_kernel(const uint8_t* x_img, const uint8_t* w_prep, uint8_t* staged, float* out)
|
||||
{
|
||||
__shared__ uint8_t img[kRows * kRowBytes];
|
||||
for (int i = threadIdx.x; i < kRows * kRowBytes; i += blockDim.x)
|
||||
img[i] = x_img[i];
|
||||
__syncthreads();
|
||||
f8v c = {0, 0, 0, 0, 0, 0, 0, 0};
|
||||
c = rrswin::swin_gemm_tile(c, (const kslot*)w_prep, 0, 16, img, kRowBytes, 0);
|
||||
__syncthreads();
|
||||
// the staging writes 16 rows, i.e. k 0..15; the image's second k16 chunk must not keep
|
||||
// the values the first GEMM's operand left there, or the second GEMM sums them too
|
||||
for (int i = threadIdx.x; i < kRows * kRowBytes; i += blockDim.x)
|
||||
if ((i % kRowBytes) >= 16)
|
||||
img[i] = 0;
|
||||
__syncthreads();
|
||||
rrswin::swin_stage_operand(img, kRowBytes, 0, 0, c);
|
||||
__syncthreads();
|
||||
for (int i = threadIdx.x; i < kRows * kRowBytes; i += blockDim.x)
|
||||
staged[i] = img[i];
|
||||
__syncthreads();
|
||||
f8v d = {0, 0, 0, 0, 0, 0, 0, 0};
|
||||
d = rrswin::swin_gemm_tile(d, (const kslot*)w_prep, 64, 16, img, kRowBytes, 0);
|
||||
for (int i = 0; i < 8; ++i)
|
||||
out[i * 32 + threadIdx.x] = d[i];
|
||||
}
|
||||
|
||||
// Step 4: the weight preparation. Lane t prepares weight row t's two k16 slots (32 bytes), then
|
||||
// the phase runs with a known activation image and the result is compared against a host model
|
||||
// that reads the weight image through the documented waddr. This is what proves the prep's slot
|
||||
// addressing, which the GEMM alone cannot see.
|
||||
__global__ void step4_kernel(const uint8_t* x_img, const uint8_t* wimg, uint8_t* slots, float* out)
|
||||
{
|
||||
__shared__ uint8_t img[kRows * kRowBytes];
|
||||
for (int i = threadIdx.x; i < kRows * kRowBytes; i += blockDim.x)
|
||||
img[i] = x_img[i];
|
||||
__syncthreads();
|
||||
const int t = threadIdx.x & 15;
|
||||
if ((threadIdx.x >> 4) == 0)
|
||||
rrswin::swin_prep_phase(slots, 0, 16, 1, wimg, rrswin::swin_waddr_enc0);
|
||||
__syncthreads();
|
||||
f8v c = {0, 0, 0, 0, 0, 0, 0, 0};
|
||||
c = rrswin::swin_gemm_tile(c, (const kslot*)slots, 0, 16, img, kRowBytes, 0);
|
||||
(void)t;
|
||||
for (int i = 0; i < 8; ++i)
|
||||
out[i * 32 + threadIdx.x] = c[i];
|
||||
}
|
||||
|
||||
// Step 4c: a 32-row phase (two m-tiles of 16 weight rows). slot_step is n_rows = 32, so tile 1's
|
||||
// rows sit past tile 0's second k16 slot; with a fixed step of 16 it would read them as tile 0's
|
||||
// second k16 operand.
|
||||
__global__ void step4c_kernel(const uint8_t* x_img, const uint8_t* wimg, uint8_t* slots, float* out)
|
||||
{
|
||||
__shared__ uint8_t img[kRows * kRowBytes];
|
||||
for (int i = threadIdx.x; i < kRows * kRowBytes; i += blockDim.x)
|
||||
img[i] = x_img[i];
|
||||
__syncthreads();
|
||||
if (threadIdx.x == 0)
|
||||
rrswin::swin_prep_phase(slots, 0, 32, 1, wimg, rrswin::swin_waddr_enc0);
|
||||
__syncthreads();
|
||||
f8v c = {0, 0, 0, 0, 0, 0, 0, 0};
|
||||
c = rrswin::swin_gemm_tile(c, (const kslot*)slots, 0, 32, img, kRowBytes, 0);
|
||||
for (int i = 0; i < 8; ++i)
|
||||
out[i * 32 + threadIdx.x] = c[i];
|
||||
f8v d = {0, 0, 0, 0, 0, 0, 0, 0};
|
||||
d = rrswin::swin_gemm_tile(d, (const kslot*)slots, 16, 32, img, kRowBytes, 0);
|
||||
for (int i = 0; i < 8; ++i)
|
||||
out[256 + i * 32 + threadIdx.x] = d[i];
|
||||
}
|
||||
|
||||
// Step 5: the phase loop. Phase 1 is a 32-row (two m-tiles) x 16-token x k64 GEMM that stages
|
||||
// its accumulator; phase 2 is 16x16x32 reading that staged image. The host model redoes the
|
||||
// whole chain, including the e4m3 requantisation between the phases.
|
||||
__global__ void step5_kernel(const uint8_t* x_img, const uint8_t* wimg, uint8_t* slots,
|
||||
uint8_t* staged, float* out)
|
||||
{
|
||||
__shared__ uint8_t img[16 * 64];
|
||||
__shared__ uint8_t mid[16 * 32];
|
||||
for (int i = threadIdx.x; i < 16 * 64; i += blockDim.x)
|
||||
img[i] = x_img[i];
|
||||
for (int i = threadIdx.x; i < 16 * 32; i += blockDim.x)
|
||||
mid[i] = 0;
|
||||
__syncthreads();
|
||||
if (threadIdx.x == 0)
|
||||
{
|
||||
rrswin::swin_prep_phase(slots + 0, 0, 32, 2, wimg, rrswin::swin_waddr_enc0); // phase 1
|
||||
rrswin::swin_prep_phase(slots + 32 * 32 * 2, 0, 16, 1, wimg + 8192, rrswin::swin_waddr_enc0);
|
||||
}
|
||||
__syncthreads();
|
||||
rrswin::Phase p1 = {2, 1, 2, 0, 0, 0};
|
||||
rrswin::swin_phase(p1, (const kslot*)slots, 32, img, 64, mid, 32, nullptr);
|
||||
__syncthreads();
|
||||
for (int i = threadIdx.x; i < 16 * 32; i += blockDim.x)
|
||||
staged[i] = mid[i];
|
||||
rrswin::Phase p2 = {1, 1, 1, 128, 1, 0}; // slot 128 = byte 2048, past phase 1's slots
|
||||
f8v acc[1];
|
||||
rrswin::swin_phase(p2, (const kslot*)slots, 16, mid, 32, nullptr, 0, acc);
|
||||
for (int i = 0; i < 8; ++i)
|
||||
out[i * 32 + threadIdx.x] = acc[0][i];
|
||||
}
|
||||
|
||||
// Step 6: the clamped cubic activation. The host model below is written from the specification
|
||||
// in kernels/rr/rrswin_enc0.hip's comment (y = clamp(x, +-2), out = x * (0.5 + y * (0.412109375
|
||||
// - 0.0810546875*|y|)), f16 throughout, constants 0x3698/0x2d30/0x3800/0x4000) rather than from
|
||||
// the device code, so agreement means the port is faithful.
|
||||
static uint16_t act_model(uint16_t bits)
|
||||
{
|
||||
const auto as_half = [](uint16_t b) { _Float16 h; std::memcpy(&h, &b, sizeof(h)); return h; };
|
||||
const auto as_bits = [](const _Float16& h) { uint16_t b; std::memcpy(&b, &h, sizeof(b)); return b; };
|
||||
const _Float16 x = as_half(bits);
|
||||
const _Float16 two = as_half(0x4000), slope = as_half(0x2d30), mid = as_half(0x3698), half0 = as_half(0x3800);
|
||||
_Float16 y = x > -two ? x : -two;
|
||||
y = y < two ? y : two;
|
||||
y = y > -y ? y : -y;
|
||||
const _Float16 g = (_Float16)(half0 + y * (mid - slope * y));
|
||||
return as_bits((_Float16)(x * g));
|
||||
}
|
||||
|
||||
__global__ void step6_kernel(const uint32_t* pairs, uint32_t* out)
|
||||
{
|
||||
const uint32_t p = pairs[threadIdx.x];
|
||||
out[threadIdx.x] = rrswin::swin_act_pair(p);
|
||||
}
|
||||
|
||||
// Step 7: fragment-to-image staging. A host model of the documented layout builds the fragment
|
||||
// registers from a known image, the device stages them back, and the result must equal the
|
||||
// original image -- a round trip that fails if any (row, k) mapping is wrong.
|
||||
static void fragment_model(const std::vector<uint8_t>& img, int stride, int m_tiles,
|
||||
std::vector<uint32_t>& frag, int lane)
|
||||
{
|
||||
const int g = lane >> 2, t = lane & 3;
|
||||
uint32_t* mine = frag.data() + (size_t)4 * m_tiles * lane; // each lane's own registers
|
||||
for (int mi = 0; mi < m_tiles; ++mi)
|
||||
for (int j = 0; j < 4; ++j)
|
||||
{
|
||||
const uint32_t k0 = 2 * t + (j & 1) + 16u * ((j >> 1) & 1);
|
||||
const uint32_t k2 = 8u + k0;
|
||||
mine[4 * mi + 0] |= (uint32_t)img[(size_t)(16 * mi + g) * stride + k0] << (8 * j);
|
||||
mine[4 * mi + 1] |= (uint32_t)img[(size_t)(16 * mi + g + 8) * stride + k0] << (8 * j);
|
||||
mine[4 * mi + 2] |= (uint32_t)img[(size_t)(16 * mi + g) * stride + k2] << (8 * j);
|
||||
mine[4 * mi + 3] |= (uint32_t)img[(size_t)(16 * mi + g + 8) * stride + k2] << (8 * j);
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void step7_kernel(const uint32_t* frag, uint8_t* staged, int m_tiles)
|
||||
{
|
||||
__shared__ uint8_t img[64 * 32];
|
||||
for (int i = threadIdx.x; i < 64 * 32; i += blockDim.x)
|
||||
img[i] = 0;
|
||||
__syncthreads();
|
||||
rrswin::swin_stage_fragment(img, 32, frag + 4 * m_tiles * threadIdx.x, m_tiles);
|
||||
__syncthreads();
|
||||
for (int i = threadIdx.x; i < 64 * 32; i += blockDim.x)
|
||||
staged[i] = img[i];
|
||||
}
|
||||
|
||||
int main()
|
||||
{
|
||||
std::vector<uint8_t> x(kRows * kRowBytes, 0), w(kSlots * kChunk, 0);
|
||||
// Well-spread e4m3 codes: small/huge and equal sums must not collide, or the mapping
|
||||
// search below cannot tell one (row, column) from another.
|
||||
uint32_t lcg = 0x12345678u;
|
||||
auto next_code = [&lcg]() -> uint8_t {
|
||||
lcg = lcg * 1664525u + 1013904223u;
|
||||
uint32_t mag = (lcg >> 13) & 0x3fu; // exponents 0..7 and mantissas
|
||||
if (mag == 0)
|
||||
mag = 0x21;
|
||||
const uint32_t sign = (lcg >> 7) & 1u;
|
||||
return (uint8_t)((sign << 7) | mag);
|
||||
};
|
||||
for (int r = 0; r < kRows; ++r)
|
||||
for (int c = 0; c < kChunks * kChunk; ++c)
|
||||
{
|
||||
// row r's k16 chunk h of k32 step `step` lives at slot 32*step + 16*h + r
|
||||
const int step = c / 32, h = (c % 32) / 16, kk = c % 16;
|
||||
w[(size_t)(32 * step + 16 * h + r) * kChunk + kk] = next_code();
|
||||
x[(size_t)r * kRowBytes + c] = next_code();
|
||||
}
|
||||
|
||||
uint8_t *dx, *dw;
|
||||
float* dout;
|
||||
CHECK(hipMalloc(&dx, x.size()));
|
||||
CHECK(hipMalloc(&dw, w.size()));
|
||||
CHECK(hipMalloc(&dout, 8 * 32 * sizeof(float)));
|
||||
std::vector<float> got(8 * 32);
|
||||
|
||||
CHECK(hipMemcpy(dx, x.data(), x.size(), hipMemcpyHostToDevice));
|
||||
CHECK(hipMemcpy(dw, w.data(), w.size(), hipMemcpyHostToDevice));
|
||||
step1_kernel<<<1, 32>>>(dx, dw, dout);
|
||||
CHECK(hipGetLastError());
|
||||
CHECK(hipDeviceSynchronize());
|
||||
CHECK(hipMemcpy(got.data(), dout, got.size() * sizeof(float), hipMemcpyDeviceToHost));
|
||||
check("step 1 one k32 step (k 0..31)", got, w, x, 32, 2e-3);
|
||||
|
||||
step2_kernel<<<1, 32>>>(dx, dw, dout);
|
||||
CHECK(hipGetLastError());
|
||||
CHECK(hipDeviceSynchronize());
|
||||
CHECK(hipMemcpy(got.data(), dout, got.size() * sizeof(float), hipMemcpyDeviceToHost));
|
||||
check("step 2 k64 as two k32 with f16 rounding", got, w, x, 64, 2e-3);
|
||||
|
||||
// step 3: two chained GEMMs with the accumulator staged as the next operand
|
||||
{
|
||||
std::vector<uint8_t> w2 = w; // slots 64..95 are zero so far
|
||||
for (int r = 0; r < kRows; ++r)
|
||||
for (int c = 0; c < 32; ++c)
|
||||
{
|
||||
const int h = c / 16, kk = c % 16;
|
||||
w2[(size_t)(64 + 16 * h + r) * kChunk + kk] = next_code();
|
||||
}
|
||||
CHECK(hipMemcpy(dw, w2.data(), w2.size(), hipMemcpyHostToDevice));
|
||||
uint8_t* dstaged;
|
||||
CHECK(hipMalloc(&dstaged, kRows * kRowBytes));
|
||||
step3_kernel<<<1, 32>>>(dx, dw, dstaged, dout);
|
||||
CHECK(hipGetLastError());
|
||||
CHECK(hipDeviceSynchronize());
|
||||
CHECK(hipMemcpy(got.data(), dout, got.size() * sizeof(float), hipMemcpyDeviceToHost));
|
||||
|
||||
std::vector<uint8_t> staged(kRows * kRowBytes, 0);
|
||||
CHECK(hipMemcpy(staged.data(), dstaged, staged.size(), hipMemcpyDeviceToHost));
|
||||
auto w1_at = [&](int m, int k) { return w[(size_t)(32 * (k / 32) + 16 * ((k % 32) / 16) + m) * kChunk + (k % 16)]; };
|
||||
// where did each expected staged byte land, and what is at its expected slot?
|
||||
int pl = 0;
|
||||
for (int n = 0; n < 16 && pl < 6; ++n)
|
||||
for (int m = 0; m < 16 && pl < 6; ++m)
|
||||
{
|
||||
double d1 = 0.0;
|
||||
for (int j = 0; j < 32; ++j)
|
||||
d1 += dec(w1_at(m, j)) * dec(x[(size_t)n * kRowBytes + j]);
|
||||
const uint8_t want = requant_code(((_Float16)(float)d1) ? [&] {
|
||||
const _Float16 h = (_Float16)d1;
|
||||
uint16_t b;
|
||||
std::memcpy(&b, &h, sizeof(b));
|
||||
return b;
|
||||
}() : 0);
|
||||
const uint8_t have = staged[(size_t)n * kRowBytes + m];
|
||||
std::printf(" staged[col %2d][k %2d] = 0x%02x, expected 0x%02x%s\n", n, m, have, want,
|
||||
have == want ? "" : " <-- differs");
|
||||
++pl;
|
||||
}
|
||||
auto w2_at = [&](int m, int k) { return w2[(size_t)(64 + 32 * (k / 32) + 16 * ((k % 32) / 16) + m) * kChunk + (k % 16)]; };
|
||||
int bad = 0, checked = 0;
|
||||
for (int l = 0; l < 32; ++l)
|
||||
{
|
||||
const int h = l >> 4, col = l & 15;
|
||||
for (int i = 0; i < 8; ++i)
|
||||
{
|
||||
const int m = i + 8 * h;
|
||||
// The staging covers 16 rows (the first GEMM's m), so the second GEMM's
|
||||
// k 16..31 operand is zero; the model sums the same 16 k.
|
||||
double ref = 0.0;
|
||||
for (int k = 0; k < 16; ++k)
|
||||
{
|
||||
double d1 = 0.0;
|
||||
for (int j = 0; j < 32; ++j)
|
||||
d1 += dec(w1_at(k, j)) * dec(x[(size_t)col * kRowBytes + j]);
|
||||
ref += dec(w2_at(m, k)) * requant(d1);
|
||||
}
|
||||
const float g = got[i * 32 + l];
|
||||
++checked;
|
||||
if (std::fabs((double)g - ref) > 2e-3 * (1.0 + std::fabs(ref)))
|
||||
{
|
||||
if (bad < 4)
|
||||
std::printf(" step3 lane %2d vgpr %d: got %+.6f ref %+.6f\n", l, i, g, ref);
|
||||
++bad;
|
||||
}
|
||||
}
|
||||
}
|
||||
std::printf("step 3 accumulator staged as the next operand: checked %d, mismatched %d\n",
|
||||
checked, bad);
|
||||
CHECK(hipFree(dstaged));
|
||||
if (bad)
|
||||
std::exit(1);
|
||||
}
|
||||
// step 4: weight preparation, checked through the GEMM against the documented layout
|
||||
{
|
||||
std::vector<uint8_t> wimg(65536, 0);
|
||||
for (int k = 0; k < 32; ++k)
|
||||
for (int n = 0; n < 16; ++n)
|
||||
wimg[rrswin_waddr_host(k, n)] = next_code();
|
||||
|
||||
uint8_t* dwimg;
|
||||
uint8_t* dslots;
|
||||
CHECK(hipMalloc(&dwimg, wimg.size()));
|
||||
CHECK(hipMalloc(&dslots, 16 * 16 * 2));
|
||||
CHECK(hipMemset(dslots, 0, 16 * 16 * 2));
|
||||
CHECK(hipMemcpy(dwimg, wimg.data(), wimg.size(), hipMemcpyHostToDevice));
|
||||
step4_kernel<<<1, 32>>>(dx, dwimg, dslots, dout);
|
||||
CHECK(hipGetLastError());
|
||||
CHECK(hipDeviceSynchronize());
|
||||
CHECK(hipMemcpy(got.data(), dout, got.size() * sizeof(float), hipMemcpyDeviceToHost));
|
||||
|
||||
std::vector<uint8_t> slotsHost(16 * 16 * 2, 0);
|
||||
CHECK(hipMemcpy(slotsHost.data(), dslots, slotsHost.size(), hipMemcpyDeviceToHost));
|
||||
int bad = 0, checked = 0;
|
||||
for (int hus = 0; hus < 2; ++hus)
|
||||
for (int n = 0; n < 16; ++n)
|
||||
for (int j = 0; j < 16; ++j)
|
||||
{
|
||||
const uint8_t have = slotsHost[16 * (16 * hus + n) + j];
|
||||
const uint8_t want = wimg[rrswin_waddr_host(16 * hus + j, n)];
|
||||
++checked;
|
||||
if (have != want)
|
||||
{
|
||||
if (bad < 4)
|
||||
std::printf(" slot[h%u n%2d j%2d] = 0x%02x, expected 0x%02x\n", hus, n, j, have, want);
|
||||
++bad;
|
||||
}
|
||||
}
|
||||
std::printf("step 4a prep slots: checked %d, mismatched %d\n", checked, bad);
|
||||
|
||||
int gbad = 0, gchecked = 0;
|
||||
for (int l = 0; l < 32; ++l)
|
||||
{
|
||||
const int h = l >> 4, col = l & 15;
|
||||
for (int i = 0; i < 8; ++i)
|
||||
{
|
||||
const int n = i + 8 * h;
|
||||
double ref = 0.0;
|
||||
for (int k = 0; k < 32; ++k)
|
||||
ref += dec(wimg[rrswin_waddr_host(k, n)]) * dec(x[(size_t)col * kRowBytes + k]);
|
||||
const float g = got[i * 32 + l];
|
||||
++gchecked;
|
||||
if (std::fabs((double)g - ref) > 2e-3 * (1.0 + std::fabs(ref)))
|
||||
{
|
||||
if (gbad < 4)
|
||||
std::printf(" step4 GEMM lane %2d vgpr %d: got %+.6f ref %+.6f\n", l, i, g, ref);
|
||||
++gbad;
|
||||
}
|
||||
}
|
||||
}
|
||||
std::printf("step 4b prep through the GEMM: checked %d, mismatched %d\n", gchecked, gbad);
|
||||
|
||||
|
||||
CHECK(hipFree(dwimg));
|
||||
CHECK(hipFree(dslots));
|
||||
if (bad || gbad)
|
||||
std::exit(1);
|
||||
}
|
||||
// step 4c: 32 rows, two m-tiles, slot_step = n_rows
|
||||
{
|
||||
std::vector<uint8_t> wimg(131072, 0);
|
||||
for (int k = 0; k < 32; ++k)
|
||||
for (int n = 0; n < 32; ++n)
|
||||
wimg[rrswin_waddr_host(k, n)] = next_code();
|
||||
|
||||
uint8_t* dwimg;
|
||||
uint8_t* dslots;
|
||||
float* dout2;
|
||||
CHECK(hipMalloc(&dwimg, wimg.size()));
|
||||
CHECK(hipMalloc(&dslots, 32 * 16 * 2));
|
||||
CHECK(hipMalloc(&dout2, 512 * sizeof(float)));
|
||||
CHECK(hipMemset(dslots, 0, 32 * 16 * 2));
|
||||
CHECK(hipMemcpy(dwimg, wimg.data(), wimg.size(), hipMemcpyHostToDevice));
|
||||
step4c_kernel<<<1, 32>>>(dx, dwimg, dslots, dout2);
|
||||
CHECK(hipGetLastError());
|
||||
CHECK(hipDeviceSynchronize());
|
||||
std::vector<float> got2(512);
|
||||
CHECK(hipMemcpy(got2.data(), dout2, got2.size() * sizeof(float), hipMemcpyDeviceToHost));
|
||||
|
||||
int bad = 0, checked = 0;
|
||||
for (int tile = 0; tile < 2; ++tile)
|
||||
for (int l = 0; l < 32; ++l)
|
||||
{
|
||||
const int h = l >> 4, col = l & 15;
|
||||
for (int i = 0; i < 8; ++i)
|
||||
{
|
||||
const int n = 16 * tile + i + 8 * h;
|
||||
double ref = 0.0;
|
||||
for (int k = 0; k < 32; ++k)
|
||||
ref += dec(wimg[rrswin_waddr_host(k, n)]) * dec(x[(size_t)col * kRowBytes + k]);
|
||||
const float g = got2[tile * 256 + i * 32 + l];
|
||||
++checked;
|
||||
if (std::fabs((double)g - ref) > 2e-3 * (1.0 + std::fabs(ref)))
|
||||
{
|
||||
if (bad < 4)
|
||||
std::printf(" step4c tile %d lane %2d vgpr %d: got %+.6f ref %+.6f\n",
|
||||
tile, l, i, g, ref);
|
||||
++bad;
|
||||
}
|
||||
}
|
||||
}
|
||||
std::printf("step 4c 32-row phase, two m-tiles: checked %d, mismatched %d\n", checked, bad);
|
||||
CHECK(hipFree(dwimg));
|
||||
CHECK(hipFree(dslots));
|
||||
CHECK(hipFree(dout2));
|
||||
if (bad)
|
||||
std::exit(1);
|
||||
}
|
||||
// step 5: the phase loop, two phases chained through a staged image
|
||||
{
|
||||
std::vector<uint8_t> wimg(32768, 0);
|
||||
for (int k = 0; k < 32; ++k)
|
||||
for (int n = 0; n < 48; ++n)
|
||||
wimg[(n < 32 ? rrswin_waddr_host(k, n) : 8192 + rrswin_waddr_host(k, n - 32))] = next_code();
|
||||
|
||||
uint8_t *dwimg, *dslots, *dstaged;
|
||||
float* dout2;
|
||||
CHECK(hipMalloc(&dwimg, wimg.size()));
|
||||
CHECK(hipMalloc(&dslots, (32 * 32 * 2 + 16 * 16 * 2)));
|
||||
CHECK(hipMalloc(&dstaged, 16 * 32));
|
||||
CHECK(hipMalloc(&dout2, 256 * sizeof(float)));
|
||||
CHECK(hipMemset(dslots, 0, 32 * 32 * 2 + 16 * 16 * 2));
|
||||
CHECK(hipMemcpy(dwimg, wimg.data(), wimg.size(), hipMemcpyHostToDevice));
|
||||
step5_kernel<<<1, 32>>>(dx, dwimg, dslots, dstaged, dout2);
|
||||
CHECK(hipGetLastError());
|
||||
CHECK(hipDeviceSynchronize());
|
||||
std::vector<float> got2(256);
|
||||
CHECK(hipMemcpy(got2.data(), dout2, got2.size() * sizeof(float), hipMemcpyDeviceToHost));
|
||||
|
||||
std::vector<uint8_t> staged(16 * 32, 0);
|
||||
CHECK(hipMemcpy(staged.data(), dstaged, staged.size(), hipMemcpyDeviceToHost));
|
||||
auto w1 = [&](int m, int k) { return wimg[rrswin_waddr_host(k, m)]; };
|
||||
{
|
||||
int sbad = 0, schecked = 0;
|
||||
for (int col = 0; col < 16 && sbad < 4; ++col)
|
||||
for (int m = 0; m < 32; ++m)
|
||||
{
|
||||
double d1 = 0.0;
|
||||
for (int j = 0; j < 64; ++j)
|
||||
d1 += dec(w1(m, j)) * dec(x[(size_t)col * kRowBytes + j]);
|
||||
const uint8_t want = requant_code_bytes(d1);
|
||||
const uint8_t have = staged[(size_t)col * 32 + m];
|
||||
++schecked;
|
||||
if (have != want)
|
||||
{
|
||||
if (sbad < 4)
|
||||
std::printf(" staged5[col %2d][k %2d] = 0x%02x, expected 0x%02x\n", col, m, have, want);
|
||||
++sbad;
|
||||
}
|
||||
}
|
||||
std::printf("step 5a staged image: checked %d, mismatched %d\n", schecked, sbad);
|
||||
}
|
||||
auto w2 = [&](int m, int k) { return wimg[8192 + rrswin_waddr_host(k, m)]; };
|
||||
int bad = 0, checked = 0;
|
||||
for (int l = 0; l < 32; ++l)
|
||||
{
|
||||
const int h = l >> 4, col = l & 15;
|
||||
for (int i = 0; i < 8; ++i)
|
||||
{
|
||||
const int n = i + 8 * h;
|
||||
// phase 1's staged image: row = token col, k = phase 1's m
|
||||
// phase 2's k spans 32: phase 1's two m-tiles land in the staged image's two
|
||||
// k16 chunks
|
||||
double ref = 0.0;
|
||||
for (int k = 0; k < 32; ++k)
|
||||
{
|
||||
double d1 = 0.0;
|
||||
for (int j = 0; j < 64; ++j)
|
||||
d1 += dec(w1(k, j)) * dec(x[(size_t)col * kRowBytes + j]);
|
||||
ref += dec(w2(n, k)) * requant(d1);
|
||||
}
|
||||
const float g = got2[i * 32 + l];
|
||||
++checked;
|
||||
if (std::fabs((double)g - ref) > 3e-3 * (1.0 + std::fabs(ref)))
|
||||
{
|
||||
if (bad < 4)
|
||||
std::printf(" step5 lane %2d vgpr %d: got %+.6f ref %+.6f\n", l, i, g, ref);
|
||||
++bad;
|
||||
}
|
||||
}
|
||||
}
|
||||
std::printf("step 5 phase loop, two phases: checked %d, mismatched %d\n", checked, bad);
|
||||
CHECK(hipFree(dwimg));
|
||||
CHECK(hipFree(dslots));
|
||||
CHECK(hipFree(dstaged));
|
||||
CHECK(hipFree(dout2));
|
||||
if (bad)
|
||||
std::exit(1);
|
||||
}
|
||||
// step 6: the clamped cubic activation, over f16 pairs spanning and crossing the clamps
|
||||
{
|
||||
std::vector<uint32_t> pairs(64);
|
||||
for (int i = 0; i < 64; ++i)
|
||||
{
|
||||
const _Float16 lo = (_Float16)((i - 32) * 0.21f);
|
||||
const _Float16 hi = (_Float16)((i - 16) * -0.13f + 1.7f);
|
||||
uint16_t bl, bh;
|
||||
std::memcpy(&bl, &lo, sizeof(bl));
|
||||
std::memcpy(&bh, &hi, sizeof(bh));
|
||||
pairs[i] = (uint32_t)bl | ((uint32_t)bh << 16);
|
||||
}
|
||||
uint32_t *dp, *dop;
|
||||
CHECK(hipMalloc(&dp, pairs.size() * 4));
|
||||
CHECK(hipMalloc(&dop, pairs.size() * 4));
|
||||
CHECK(hipMemcpy(dp, pairs.data(), pairs.size() * 4, hipMemcpyHostToDevice));
|
||||
step6_kernel<<<1, 64>>>(dp, dop);
|
||||
CHECK(hipGetLastError());
|
||||
CHECK(hipDeviceSynchronize());
|
||||
std::vector<uint32_t> got6(pairs.size());
|
||||
CHECK(hipMemcpy(got6.data(), dop, got6.size() * 4, hipMemcpyDeviceToHost));
|
||||
int bad = 0, checked = 0;
|
||||
for (int i = 0; i < 64; ++i)
|
||||
{
|
||||
const uint32_t want = (uint32_t)act_model((uint16_t)(pairs[i] & 0xffff))
|
||||
| ((uint32_t)act_model((uint16_t)(pairs[i] >> 16)) << 16);
|
||||
++checked;
|
||||
// The device applies the formula with packed f16x2 operations, which is what the PTX
|
||||
// does; this scalar _Float16 model can land one ulp away. A logic error moves whole
|
||||
// clamps, so one ulp is the right bar here; the oracle comparison is the bit-exact
|
||||
// check, later.
|
||||
auto within_ulp = [](uint16_t a, uint16_t b) {
|
||||
// compare as f16 by ordering the sign-magnitude representation
|
||||
const auto key = [](uint16_t v) -> int {
|
||||
return (v & 0x8000) ? (int)(0x8000 - (v & 0x7fff)) : (int)(v & 0x7fff);
|
||||
};
|
||||
const int d = key(a) - key(b);
|
||||
return d >= -1 && d <= 1;
|
||||
};
|
||||
if (!within_ulp((uint16_t)(got6[i] & 0xffff), (uint16_t)(want & 0xffff)) ||
|
||||
!within_ulp((uint16_t)(got6[i] >> 16), (uint16_t)(want >> 16)))
|
||||
{
|
||||
if (bad < 4)
|
||||
std::printf(" act pair %2d: got %08x want %08x\n", i, got6[i], want);
|
||||
++bad;
|
||||
}
|
||||
}
|
||||
std::printf("step 6 clamped cubic activation (1 ulp): checked %d, mismatched %d\n", checked, bad);
|
||||
CHECK(hipFree(dp));
|
||||
CHECK(hipFree(dop));
|
||||
if (bad)
|
||||
std::exit(1);
|
||||
}
|
||||
// step 7: fragment -> image staging round trip
|
||||
{
|
||||
const int m_tiles = 4;
|
||||
std::vector<uint8_t> src(64 * 32, 0);
|
||||
uint32_t lcg7 = 0x2468ace0u;
|
||||
for (auto& v : src)
|
||||
{
|
||||
lcg7 = lcg7 * 1664525u + 1013904223u;
|
||||
v = (uint8_t)((lcg7 >> 21) & 0x7fu);
|
||||
}
|
||||
std::vector<uint32_t> frag(32 * 4 * m_tiles, 0);
|
||||
for (int lane = 0; lane < 32; ++lane)
|
||||
fragment_model(src, 32, m_tiles, frag, lane);
|
||||
|
||||
uint32_t* dfrag;
|
||||
uint8_t* dstaged7;
|
||||
CHECK(hipMalloc(&dfrag, frag.size() * 4));
|
||||
CHECK(hipMalloc(&dstaged7, src.size()));
|
||||
CHECK(hipMemcpy(dfrag, frag.data(), frag.size() * 4, hipMemcpyHostToDevice));
|
||||
CHECK(hipMemset(dstaged7, 0, src.size()));
|
||||
step7_kernel<<<1, 32>>>(dfrag, dstaged7, m_tiles);
|
||||
CHECK(hipGetLastError());
|
||||
CHECK(hipDeviceSynchronize());
|
||||
std::vector<uint8_t> got7(src.size(), 0);
|
||||
CHECK(hipMemcpy(got7.data(), dstaged7, src.size(), hipMemcpyDeviceToHost));
|
||||
int bad = 0, checked = 0;
|
||||
for (size_t i = 0; i < src.size(); ++i)
|
||||
{
|
||||
++checked;
|
||||
if (got7[i] != src[i])
|
||||
{
|
||||
if (bad < 4)
|
||||
std::printf(" stage byte %zu (row %zu k %zu): got 0x%02x want 0x%02x\n", i, i / 32, i % 32,
|
||||
got7[i], src[i]);
|
||||
++bad;
|
||||
}
|
||||
}
|
||||
std::printf("step 7 fragment staging round trip: checked %d, mismatched %d\n", checked, bad);
|
||||
CHECK(hipFree(dfrag));
|
||||
CHECK(hipFree(dstaged7));
|
||||
if (bad)
|
||||
std::exit(1);
|
||||
}
|
||||
CHECK(hipFree(dx));
|
||||
CHECK(hipFree(dw));
|
||||
CHECK(hipFree(dout));
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,93 @@
|
||||
// Test a translated PTX MMA, not a direct call to its lowering helpers.
|
||||
// The first warp catches masked-EXEC padding (A=1, B=0.5, C=1 => D=17).
|
||||
// The second checks every row/column with A[r,k]=(r+1)/8, B[k,c]=(c+1)/8.
|
||||
// Its nonuniform bias C[r,c]=8*r+c+0.25 also checks accumulator layout conversion.
|
||||
// A second MMA checks partial accumulator writes separated by a dependent read:
|
||||
// double the low-row bias through the shadow, negate the high-row bias.
|
||||
// All expected values are exactly representable in f16.
|
||||
//
|
||||
// Build: hipcc -O2 kernels/tools/rrswin_fp8_unpaired_test.cpp -o /tmp/fp8_unpaired_test
|
||||
// Emit with the patched ZLUDA d4r_emit example:
|
||||
// env D4R_ZLUDA_WMMA=1 D4R_ZLUDA_WMMA_FP8=1 D4R_ZLUDA_WMMA_FP8_NATIVE=1 \
|
||||
// d4r_emit kernels/tools/rrswin_fp8_unpaired.ptx /tmp/fp8_unpaired gfx1201
|
||||
// Run: /tmp/fp8_unpaired_test /tmp/fp8_unpaired/module.hsaco
|
||||
#include <hip/hip_runtime.h>
|
||||
#include <cstdint>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
|
||||
#define CHECK(expr) do { \
|
||||
const hipError_t error = (expr); \
|
||||
if (error != hipSuccess) { \
|
||||
std::fprintf(stderr, "%s: %s\n", #expr, hipGetErrorString(error)); \
|
||||
std::exit(2); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
static uint16_t half_bits(float value)
|
||||
{
|
||||
const _Float16 half = static_cast<_Float16>(value);
|
||||
uint16_t bits;
|
||||
std::memcpy(&bits, &half, sizeof(bits));
|
||||
return bits;
|
||||
}
|
||||
|
||||
int main(int argc, char **argv)
|
||||
{
|
||||
if (argc != 2) {
|
||||
std::fprintf(stderr, "usage: %s module.hsaco\n", argv[0]);
|
||||
return 2;
|
||||
}
|
||||
constexpr unsigned lanes = 64;
|
||||
constexpr uint8_t eighths[16] = {
|
||||
0x20, 0x28, 0x2c, 0x30, 0x32, 0x34, 0x36, 0x38,
|
||||
0x39, 0x3a, 0x3b, 0x3c, 0x3d, 0x3e, 0x3f, 0x40
|
||||
};
|
||||
uint32_t input[lanes * 8], output[lanes * 4];
|
||||
uint16_t expected[lanes * 8];
|
||||
for (unsigned i = 0; i < lanes; ++i) {
|
||||
const unsigned g = (i % 32) / 4, t = i % 4;
|
||||
const uint32_t low = i < 32 ? 0x38383838u : uint32_t(eighths[g]) * 0x01010101u;
|
||||
const uint32_t high = i < 32 ? low : uint32_t(eighths[g + 8]) * 0x01010101u;
|
||||
const uint32_t b = i < 32 ? 0x30303030u : uint32_t(eighths[g]) * 0x01010101u;
|
||||
uint32_t fragment[8] = {low, high, low, high, b, b, 0, 0};
|
||||
for (unsigned row = 0; row < 2; ++row)
|
||||
for (unsigned column = 0; column < 2; ++column) {
|
||||
const float bias = i < 32 ? 1.f : float(8 * (g + 8 * row) + 2 * t + column) + 0.25f;
|
||||
fragment[6 + row] |= uint32_t(half_bits(bias)) << (16 * column);
|
||||
const float product = i < 32 ? 16.f : float((g + 8 * row + 1) * (2 * t + column + 1)) / 2.f;
|
||||
expected[i * 8 + row * 2 + column] = half_bits(bias + product);
|
||||
expected[i * 8 + 4 + row * 2 + column] = half_bits((row == 0 ? 2.f : -1.f) * bias + product);
|
||||
}
|
||||
std::memcpy(input + i * 8, fragment, sizeof(fragment));
|
||||
}
|
||||
hipModule_t module;
|
||||
hipFunction_t function;
|
||||
CHECK(hipModuleLoad(&module, argv[1]));
|
||||
CHECK(hipModuleGetFunction(&function, module, "fp8_unpaired"));
|
||||
uint32_t *device_input, *device_output;
|
||||
CHECK(hipMalloc(&device_input, sizeof(input)));
|
||||
CHECK(hipMalloc(&device_output, sizeof(output)));
|
||||
CHECK(hipMemcpy(device_input, input, sizeof(input), hipMemcpyHostToDevice));
|
||||
CHECK(hipMemset(device_output, 0, sizeof(output)));
|
||||
void *args[] = {&device_input, &device_output};
|
||||
CHECK(hipModuleLaunchKernel(function, 2, 1, 1, 32, 1, 1, 0, nullptr, args, nullptr));
|
||||
CHECK(hipDeviceSynchronize());
|
||||
CHECK(hipMemcpy(output, device_output, sizeof(output), hipMemcpyDeviceToHost));
|
||||
uint16_t observed[lanes * 8];
|
||||
std::memcpy(observed, output, sizeof(observed));
|
||||
unsigned bad = 0;
|
||||
for (unsigned i = 0; i < lanes * 8; ++i)
|
||||
if (observed[i] != expected[i]) {
|
||||
if (bad < 8)
|
||||
std::fprintf(stderr, "warp=%u lane=%u case=%u half=%u: got %04x, expected %04x\n",
|
||||
i / 256, (i / 8) % 32, (i / 4) % 2, i % 4, observed[i], expected[i]);
|
||||
++bad;
|
||||
}
|
||||
CHECK(hipFree(device_output));
|
||||
CHECK(hipFree(device_input));
|
||||
CHECK(hipModuleUnload(module));
|
||||
std::printf("unpaired FP8 MMA and partial accumulators: %u/512 differing half values\n", bad);
|
||||
return bad ? 1 : 0;
|
||||
}
|
||||
@@ -0,0 +1,184 @@
|
||||
// Fit the engine's D mapping by measurement.
|
||||
//
|
||||
// The WMMA computes D[m][n] = sum_k A[m][k] B[k][n]. With A (the prepared weight slot) filled
|
||||
// as value_for((b + c * r) % 16) in byte b of row r, and B (the activation image) an identity
|
||||
// selector (row n holds 1.0 at byte n, zero elsewhere), the value that comes out at a given
|
||||
// (vgpr, lane) is value_for((n + c * m) % 16): one equation in the two indices the engine's
|
||||
// indexing must report. Two runs with different c (1 and 3, both coprime with 16) solve for
|
||||
// both. Nothing here assumes which (r, n) is which -- the point is to measure it.
|
||||
//
|
||||
// Build and run:
|
||||
// hipcc --offload-arch=gfx1201 -DD4R_TEX_FP8 -DD4R_WMMA_LAYOUT=12 -O2 \
|
||||
// -o /tmp/fit kernels/tools/rrswin_operand_fit.cpp && /tmp/fit
|
||||
#include <hip/hip_runtime.h>
|
||||
#include <cstdio>
|
||||
#include <cstdint>
|
||||
#include <cmath>
|
||||
#include <cstdlib>
|
||||
#include <vector>
|
||||
|
||||
#include "../rr/rrswin_engine.h"
|
||||
|
||||
#define CHECK(expr) \
|
||||
do \
|
||||
{ \
|
||||
hipError_t e_ = (expr); \
|
||||
if (e_ != hipSuccess) \
|
||||
{ \
|
||||
std::printf("HIP error at %s:%d: %s -> %s\n", __FILE__, __LINE__, #expr, \
|
||||
hipGetErrorString(e_)); \
|
||||
std::exit(2); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
static const int kRows = 16, kChunk = 16, kRowBytes = 32;
|
||||
|
||||
static uint8_t value_for(int k)
|
||||
{
|
||||
const int e = 1 + (k & 7), m = (k >> 3) & 1;
|
||||
return (uint8_t)((e << 3) | (m ? 4 : 0));
|
||||
}
|
||||
|
||||
static double dec(uint8_t c)
|
||||
{
|
||||
const int sign = c & 0x80, e = (c >> 3) & 0xf, m = c & 7;
|
||||
const double mag = e == 0 ? (m / 8.0) * std::pow(2.0, -6.0) : (1.0 + m / 8.0) * std::pow(2.0, e - 7.0);
|
||||
return sign ? -mag : mag;
|
||||
}
|
||||
|
||||
static int value_index(double v)
|
||||
{
|
||||
for (int q = 0; q < 16; ++q)
|
||||
if (std::fabs(v - dec(value_for(q))) < 1e-3)
|
||||
return q;
|
||||
return -1;
|
||||
}
|
||||
|
||||
__global__ void fit_kernel(const uint8_t* x_img, const uint8_t* w_prep, float* out)
|
||||
{
|
||||
__shared__ uint8_t img[kRows * kRowBytes];
|
||||
for (int i = threadIdx.x; i < kRows * kRowBytes; i += blockDim.x)
|
||||
img[i] = x_img[i];
|
||||
__syncthreads();
|
||||
f8v c = {0, 0, 0, 0, 0, 0, 0, 0};
|
||||
c = rrswin::swin_gemm_tile(c, (const kslot*)w_prep, 0, img, kRowBytes, 0);
|
||||
for (int i = 0; i < 8; ++i)
|
||||
out[i * 32 + threadIdx.x] = c[i];
|
||||
}
|
||||
|
||||
// run one pattern and return the value_for index at every (vgpr i, lane l)
|
||||
static std::vector<int> run(int coefficient, const std::vector<uint8_t>& x)
|
||||
{
|
||||
std::vector<uint8_t> w(32 * kChunk, 0);
|
||||
for (int r = 0; r < kRows; ++r)
|
||||
for (int b = 0; b < 16; ++b)
|
||||
w[(size_t)r * kChunk + b] = value_for((b + coefficient * r) % 16);
|
||||
|
||||
uint8_t *dx, *dw;
|
||||
float* dout;
|
||||
CHECK(hipMalloc(&dx, x.size()));
|
||||
CHECK(hipMalloc(&dw, w.size()));
|
||||
CHECK(hipMalloc(&dout, 8 * 32 * sizeof(float)));
|
||||
CHECK(hipMemcpy(dx, x.data(), x.size(), hipMemcpyHostToDevice));
|
||||
CHECK(hipMemcpy(dw, w.data(), w.size(), hipMemcpyHostToDevice));
|
||||
fit_kernel<<<1, 32>>>(dx, dw, dout);
|
||||
CHECK(hipGetLastError());
|
||||
CHECK(hipDeviceSynchronize());
|
||||
std::vector<float> got(8 * 32);
|
||||
CHECK(hipMemcpy(got.data(), dout, got.size() * sizeof(float), hipMemcpyDeviceToHost));
|
||||
std::vector<int> idx(8 * 32);
|
||||
for (int i = 0; i < 8 * 32; ++i)
|
||||
idx[i] = value_index(got[i]);
|
||||
CHECK(hipFree(dx));
|
||||
CHECK(hipFree(dw));
|
||||
CHECK(hipFree(dout));
|
||||
return idx;
|
||||
}
|
||||
|
||||
// Experiment 3: A row r carries a distinct value at byte 0 only, and every B row is 1.0 at
|
||||
// byte 0, so D[i][l] reads back the A row that element was actually computed from. This is
|
||||
// the operand row selection, which experiment 2 cannot separate from the D indexing.
|
||||
static std::vector<int> run_rows(const std::vector<uint8_t>& x)
|
||||
{
|
||||
std::vector<uint8_t> w(32 * kChunk, 0);
|
||||
for (int r = 0; r < kRows; ++r)
|
||||
w[(size_t)r * kChunk + 0] = value_for(r);
|
||||
uint8_t *dx, *dw;
|
||||
float* dout;
|
||||
CHECK(hipMalloc(&dx, x.size()));
|
||||
CHECK(hipMalloc(&dw, w.size()));
|
||||
CHECK(hipMalloc(&dout, 8 * 32 * sizeof(float)));
|
||||
CHECK(hipMemcpy(dx, x.data(), x.size(), hipMemcpyHostToDevice));
|
||||
CHECK(hipMemcpy(dw, w.data(), w.size(), hipMemcpyHostToDevice));
|
||||
fit_kernel<<<1, 32>>>(dx, dw, dout);
|
||||
CHECK(hipGetLastError());
|
||||
CHECK(hipDeviceSynchronize());
|
||||
std::vector<float> got(8 * 32);
|
||||
CHECK(hipMemcpy(got.data(), dout, got.size() * sizeof(float), hipMemcpyDeviceToHost));
|
||||
std::vector<int> idx(8 * 32);
|
||||
for (int i = 0; i < 8 * 32; ++i)
|
||||
idx[i] = value_index(got[i]);
|
||||
CHECK(hipFree(dx));
|
||||
CHECK(hipFree(dw));
|
||||
CHECK(hipFree(dout));
|
||||
return idx;
|
||||
}
|
||||
|
||||
int main()
|
||||
{
|
||||
// B: row n holds 1.0 at byte n only.
|
||||
const uint8_t one = 0x38;
|
||||
std::vector<uint8_t> x(kRows * kRowBytes, 0);
|
||||
for (int n = 0; n < kRows; ++n)
|
||||
x[(size_t)n * kRowBytes + n] = one;
|
||||
|
||||
const std::vector<int> a = run(1, x);
|
||||
const std::vector<int> b = run(3, x);
|
||||
const std::vector<int> c = run(5, x);
|
||||
|
||||
std::printf("(vgpr i, lane l) -> (sum1, sum3) = ((n+m)%%16, (n+3m)%%16), solved m,n:\n");
|
||||
int shown = 0;
|
||||
for (int l = 0; l < 32 && shown < 14; l += 4)
|
||||
for (int i = 0; i < 8 && shown < 14; ++i, ++shown)
|
||||
{
|
||||
const int s1 = a[i * 32 + l], s3 = b[i * 32 + l], s5 = c[i * 32 + l];
|
||||
if (s1 < 0 || s3 < 0 || s5 < 0)
|
||||
{
|
||||
std::printf(" i%d l%-3d unreadable (sums %d %d)\n", i, l, s1, s3);
|
||||
continue;
|
||||
}
|
||||
// Three equations n + c*m = s_c (mod 16) for c = 1, 3, 5 pin both indices down.
|
||||
int m = -1, n = -1, hits = 0;
|
||||
for (int mm = 0; mm < 16; ++mm)
|
||||
for (int nn = 0; nn < 16; ++nn)
|
||||
if ((nn + mm) % 16 == s1 && (nn + 3 * mm) % 16 == s3 && (nn + 5 * mm) % 16 == s5)
|
||||
{
|
||||
if (m < 0) { m = mm; n = nn; }
|
||||
++hits;
|
||||
}
|
||||
std::printf(" i%d l%-3d sums %2d %2d %2d -> m=%2d n=%2d%s\n", i, l, s1, s3, s5, m, n,
|
||||
hits == 1 ? "" : " (AMBIGUOUS)");
|
||||
}
|
||||
// experiment 3
|
||||
{
|
||||
std::vector<uint8_t> xb(kRows * kRowBytes, 0);
|
||||
for (int n = 0; n < kRows; ++n)
|
||||
xb[(size_t)n * kRowBytes + 0] = one;
|
||||
const std::vector<int> rows = run_rows(xb);
|
||||
std::printf("\nA row used at (vgpr i, lane l):\n ");
|
||||
for (int i = 0; i < 8; ++i) std::printf(" i%-2d", i);
|
||||
std::printf("\n");
|
||||
for (int l = 0; l < 32; ++l)
|
||||
{
|
||||
std::printf("l%-4d ", l);
|
||||
for (int i = 0; i < 8; ++i)
|
||||
{
|
||||
const int v = rows[i * 32 + l];
|
||||
if (v < 0) std::printf(" ? ");
|
||||
else std::printf(" %-2d ", v);
|
||||
}
|
||||
std::printf("\n");
|
||||
}
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,120 @@
|
||||
// Layout probe: which k does each byte of the A (weight) and B (activation) operand hold?
|
||||
//
|
||||
// Method, with no assumption about byte order:
|
||||
// * A (the prepared weight slot) row r holds a distinct e4m3 value per byte of its k16 row.
|
||||
// * B (the staged activation) is an identity selector: column c is 1.0 at k = c, 0 elsewhere.
|
||||
// * The WMMA computes D[r][c] = sum_k A[r][k] * B[k][c] = A[r][k(c)], so the value that
|
||||
// comes back names the k that column c's byte was paired with. Both operands are filled
|
||||
// byte-by-byte in memory order, so the printout is a direct read of the hardware's slot
|
||||
// order for each side.
|
||||
//
|
||||
// Build and run:
|
||||
// hipcc --offload-arch=gfx1201 -DD4R_TEX_FP8 -DD4R_WMMA_LAYOUT=12 -O2 \
|
||||
// -o /tmp/probe kernels/tools/rrswin_operand_probe.cpp && /tmp/probe
|
||||
#include <hip/hip_runtime.h>
|
||||
#include <cstdio>
|
||||
#include <cstdint>
|
||||
#include <cmath>
|
||||
#include <cstdlib>
|
||||
#include <vector>
|
||||
|
||||
#include "../rr/rrswin_engine.h"
|
||||
|
||||
#define CHECK(expr) \
|
||||
do \
|
||||
{ \
|
||||
hipError_t e_ = (expr); \
|
||||
if (e_ != hipSuccess) \
|
||||
{ \
|
||||
std::printf("HIP error at %s:%d: %s -> %s\n", __FILE__, __LINE__, #expr, \
|
||||
hipGetErrorString(e_)); \
|
||||
std::exit(2); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
static const int kRows = 16, kChunk = 16, kRowBytes = 32;
|
||||
|
||||
// 16 distinct, exactly representable e4m3 values indexed by k
|
||||
static uint8_t value_for(int k)
|
||||
{
|
||||
const int e = 1 + (k & 7); // 1..8
|
||||
const int m = (k >> 3) & 1; // 0 or 1
|
||||
return (uint8_t)((e << 3) | (m ? 4 : 0));
|
||||
}
|
||||
|
||||
static double dec(uint8_t c)
|
||||
{
|
||||
const int sign = c & 0x80, e = (c >> 3) & 0xf, m = c & 7;
|
||||
const double mag = e == 0 ? (m / 8.0) * std::pow(2.0, -6.0) : (1.0 + m / 8.0) * std::pow(2.0, e - 7.0);
|
||||
return sign ? -mag : mag;
|
||||
}
|
||||
|
||||
__global__ void probe_kernel(const uint8_t* x_img, const uint8_t* w_prep, float* out)
|
||||
{
|
||||
__shared__ uint8_t img[kRows * kRowBytes];
|
||||
for (int i = threadIdx.x; i < kRows * kRowBytes; i += blockDim.x)
|
||||
img[i] = x_img[i];
|
||||
__syncthreads();
|
||||
f8v c = {0, 0, 0, 0, 0, 0, 0, 0};
|
||||
c = rrswin::swin_gemm_tile(c, (const kslot*)w_prep, 0, img, kRowBytes, 0);
|
||||
for (int i = 0; i < 8; ++i)
|
||||
out[i * 32 + threadIdx.x] = c[i];
|
||||
}
|
||||
|
||||
int main()
|
||||
{
|
||||
// A: weight row r, byte b of chunk 0 = value_for(b); the other chunk zero.
|
||||
// B: activation column c = 1.0 at byte c of chunk 0, zero elsewhere.
|
||||
const uint8_t one = 0x38; // e4m3 1.0
|
||||
std::vector<uint8_t> w(32 * kChunk, 0), x(kRows * kRowBytes, 0);
|
||||
// A row r carries only 2^(r-8) at byte 0, so the value read back names the row the
|
||||
// engine actually selected; B column c is 1.0 at k = c.
|
||||
// A row r = 1.0 at byte 0 only; B column c holds a distinct value per byte of its k16 row.
|
||||
for (int r = 0; r < kRows; ++r)
|
||||
w[(size_t)(r) * kChunk + 0] = one;
|
||||
// B row r must be distinguishable by its byte 0 as well, or "which row" is invisible:
|
||||
// byte b of row r carries value_for((b + r) % 16).
|
||||
for (int c = 0; c < kRows; ++c)
|
||||
for (int b = 0; b < 16; ++b)
|
||||
x[(size_t)c * kRowBytes + b] = value_for((b + c) % 16);
|
||||
|
||||
uint8_t *dx, *dw;
|
||||
float* dout;
|
||||
CHECK(hipMalloc(&dx, x.size()));
|
||||
CHECK(hipMalloc(&dw, w.size()));
|
||||
CHECK(hipMalloc(&dout, 8 * 32 * sizeof(float)));
|
||||
CHECK(hipMemcpy(dx, x.data(), x.size(), hipMemcpyHostToDevice));
|
||||
CHECK(hipMemcpy(dw, w.data(), w.size(), hipMemcpyHostToDevice));
|
||||
probe_kernel<<<1, 32>>>(dx, dw, dout);
|
||||
CHECK(hipGetLastError());
|
||||
CHECK(hipDeviceSynchronize());
|
||||
std::vector<float> got(8 * 32);
|
||||
CHECK(hipMemcpy(got.data(), dout, got.size() * sizeof(float), hipMemcpyDeviceToHost));
|
||||
|
||||
// The engine's assumed indexing: value at D[row = i + 8h][col = lane & 15].
|
||||
// A is the identity at byte 0, so D(row, col) = B[k0][col] = value_for((k0 + row(col)) % 16).
|
||||
// Print the value_for index of every (vgpr, lane) so the true (row, column) assignment can
|
||||
// be read off the table instead of assumed.
|
||||
std::printf("value_for index at got[vgpr i][lane l] (rows = l 0..31, cols = i 0..7):\n");
|
||||
std::printf(" ");
|
||||
for (int i = 0; i < 8; ++i) std::printf(" i%-3d", i);
|
||||
std::printf("\n");
|
||||
for (int l = 0; l < 32; ++l)
|
||||
{
|
||||
std::printf("l%-4d ", l);
|
||||
for (int i = 0; i < 8; ++i)
|
||||
{
|
||||
const double v = got[i * 32 + l];
|
||||
int n = -1;
|
||||
for (int q = 0; q < 16; ++q)
|
||||
if (std::fabs(v - dec(value_for(q))) < 1e-3) n = q;
|
||||
if (n < 0) std::printf(" ? ");
|
||||
else std::printf(" %-3d ", n);
|
||||
}
|
||||
std::printf("\n");
|
||||
}
|
||||
CHECK(hipFree(dx));
|
||||
CHECK(hipFree(dw));
|
||||
CHECK(hipFree(dout));
|
||||
return 0;
|
||||
}
|
||||
Executable
+421
@@ -0,0 +1,421 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Solve phase 8's activation row from the oracle, to derive the fragment map instead of assuming it.
|
||||
|
||||
For one token row, the oracle's accumulator gives 32 equations (one per weight column):
|
||||
|
||||
D(tok, wt) = sum_k weight(k, wt) * a(k)
|
||||
|
||||
with `weight` read from the weight image through the corroborated waddr (the PTX's own loads put
|
||||
phase 8's B at 21696/22208). That is a 32x32 linear system in the 32 unknown activation values
|
||||
a(k) of that token row, so it can be solved directly -- no assumption about the fragment layout.
|
||||
The solution is then snapped to e4m3 codes and compared with the bytes the lanes actually hold,
|
||||
which names the byte that belongs at each k: the fragment map, derived.
|
||||
|
||||
usage: phase8_fit.py <oracle-p8c-dir> [token]
|
||||
"""
|
||||
import pathlib
|
||||
import struct
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
|
||||
BASE = 46000000
|
||||
SLOT = 16384
|
||||
ORE = 32
|
||||
ACC_OFF = 32 * 4
|
||||
WEIGHT_BASE = 21696
|
||||
TOKEN_TILES = 6 # the activations' 96 rows
|
||||
WEIGHT_TILES = 4 # the weights' 32 rows
|
||||
ROWS = 96
|
||||
COLS = 32
|
||||
|
||||
|
||||
def find(dirname, needle):
|
||||
for p in sorted(pathlib.Path(dirname).iterdir()):
|
||||
if "swin_enc0_kernel" in p.name and needle in p.name:
|
||||
return p
|
||||
raise SystemExit(f"no dump with {needle} in {dirname}")
|
||||
|
||||
|
||||
def e4m3_codes():
|
||||
"""Every e4m3 code and its value."""
|
||||
out = []
|
||||
for code in range(256):
|
||||
mag = code & 0x7F
|
||||
if (code & 0x7F) == 0x7F:
|
||||
continue # NaN
|
||||
if mag >= 8:
|
||||
bits = (((mag >> 3) + 120) << 23) | ((mag & 7) << 20)
|
||||
else:
|
||||
bits = struct.unpack('<I', struct.pack('<f', float(mag) * 2.0 ** -9))[0]
|
||||
value = struct.unpack('<f', struct.pack('<I', bits | ((code & 0x80) << 24)))[0]
|
||||
out.append((code, value))
|
||||
return out
|
||||
|
||||
|
||||
def waddr(k, n):
|
||||
return (512 * (n >> 4) + 64 * (n & 7) + 16 * ((k >> 1) & 3) + 8 * ((n >> 3) & 1)
|
||||
+ 4 * ((k >> 3) & 1) + 2 * ((k >> 4) & 1) + (k & 1))
|
||||
|
||||
|
||||
def f16(bits):
|
||||
return struct.unpack('<e', struct.pack('<H', bits & 0xFFFF))[0]
|
||||
|
||||
|
||||
def snap_error(W, D):
|
||||
"""Max distance from the solved activation values to the nearest e4m3 code.
|
||||
|
||||
The true activation values are e4m3, so a correct (weight map, D map) pair makes this small;
|
||||
a wrong one leaves values scattered, as the 15.8 measured for the assumed pair. This is the
|
||||
discriminating test -- the least-squares residual is not, because the system is square.
|
||||
"""
|
||||
a, *_ = np.linalg.lstsq(W, D, rcond=None)
|
||||
codes = e4m3_codes()
|
||||
errs = [min(abs(cv[1] - a[k]) for cv in codes) for k in range(32)]
|
||||
# the mean is the useful statistic: the max is the same for many wrong permutations, because
|
||||
# they differ only in how many values land close, not in the worst one
|
||||
return float(np.mean(errs))
|
||||
|
||||
|
||||
def test_zluda_map(dump, wimg, tok, D, W):
|
||||
"""Test the map ZLUDA's own lowering documents, rather than the body's comment.
|
||||
|
||||
zluda_ptx_impl.cpp's e4m3_a_k16_fragment says: "upper holds row groupID (4 columns from
|
||||
4*threadID), lower holds row groupID+8", and the k16 halves go to registers {x,y} and {z,w}.
|
||||
That makes an A byte j of register x the k = 4*(lane&3) + j of row (lane>>2), not the
|
||||
k = 2*t + (j&1) + 16*((j>>1)&1) rrswin_enc0.hip claims. This assembles the row from those
|
||||
bytes and compares it with the value solved out of the oracle.
|
||||
"""
|
||||
a, *_ = np.linalg.lstsq(W, D, rcond=None)
|
||||
codes = e4m3_codes()
|
||||
solved = [min(codes, key=lambda cv: abs(cv[1] - a[k]))[0] for k in range(32)]
|
||||
held = {}
|
||||
for lane in range(32):
|
||||
if (lane >> 2) != (tok % 8):
|
||||
continue
|
||||
high = 1 if tok >= 8 else 0
|
||||
ops = struct.unpack(f'<{ORE}I', dump[BASE + lane * ORE * 4:BASE + lane * ORE * 4 + ORE * 4])
|
||||
mi = tok // 16
|
||||
rx = 4 * mi + (2 if high else 0)
|
||||
rz = 4 * mi + (3 if high else 1)
|
||||
for j in range(4):
|
||||
held[4 * (lane & 3) + j] = (ops[rx] >> (8 * j)) & 0xFF
|
||||
held[16 + 4 * (lane & 3) + j] = (ops[rz] >> (8 * j)) & 0xFF
|
||||
hit = sum(1 for k in range(32) if held.get(k) == solved[k])
|
||||
return hit, solved, held
|
||||
|
||||
|
||||
def main():
|
||||
d = sys.argv[1] if len(sys.argv) > 1 else "/tmp/d4r-rr-perf/d-p8c"
|
||||
tok = int(sys.argv[2]) if len(sys.argv) > 2 else 0
|
||||
dump = find(d, "arg048").read_bytes()
|
||||
wimg = find(d, "arg040").read_bytes()
|
||||
|
||||
# the oracle's D for this token: rows 16*mi + 8*rh + g live in lane M with g = M>>2, t = M&3
|
||||
mi, h, g = tok // 16, (tok % 16) // 8, (tok % 16) % 8
|
||||
rh = h # the register pair's row half is the lane's half
|
||||
D = np.zeros(COLS)
|
||||
lane_bytes = {}
|
||||
for widx in range(WEIGHT_TILES):
|
||||
for t in range(4):
|
||||
M = 16 * h + t # derived: lane = 16*((tok%16)/8) + (wt%8)/2
|
||||
off = BASE + M * ORE * 4 + ACC_OFF
|
||||
acc = struct.unpack('<48I', dump[off:off + 48 * 4])
|
||||
for c in range(2):
|
||||
wt = 8 * widx + 2 * t + c
|
||||
reg = (mi * 4 + widx) * 2 + rh
|
||||
D[wt] = f16((acc[reg] >> 16) if c else (acc[reg] & 0xFFFF))
|
||||
# this lane's A bytes, for naming the solved values
|
||||
ops = struct.unpack(f'<{ORE}I', dump[BASE + M * ORE * 4:BASE + M * ORE * 4 + ORE * 4])
|
||||
lane_bytes[M] = [((ops[4 * mi + r] >> (8 * j)) & 0xFF) for r in range(4) for j in range(4)]
|
||||
|
||||
W = np.zeros((COLS, 32))
|
||||
for wt in range(COLS):
|
||||
for k in range(32):
|
||||
W[wt, k] = e4m3_codes_value(wimg[WEIGHT_BASE + waddr(k, wt)])
|
||||
|
||||
a, *_ = np.linalg.lstsq(W, D, rcond=None)
|
||||
residual = float(np.abs(W @ a - D).max())
|
||||
# The system is square -- 32 weight columns for 32 unknown k -- so a small residual proves
|
||||
# nothing by itself: it can be solved exactly whether or not the model is right. The real
|
||||
# test is whether the *solution* lands on e4m3 codes, since the true activation values are
|
||||
# e4m3. Snap error is that measure.
|
||||
codes_all = e4m3_codes()
|
||||
_ = snap_error # used by the search below
|
||||
snap_err = max(min(abs(cv[1] - a[k]) for cv in codes_all) for k in range(32))
|
||||
print(f"token {tok}: solved 32 activation values, residual {residual:.3e} "
|
||||
f"(square system, so not evidence), max snap-to-e4m3 error {snap_err:.3e} "
|
||||
f"({'consistent with e4m3 inputs' if snap_err < 1e-3 else 'NOT e4m3: the weights or the D mapping is wrong'})")
|
||||
|
||||
codes = e4m3_codes()
|
||||
snaps = []
|
||||
for k in range(32):
|
||||
best = min(codes, key=lambda cv: abs(cv[1] - a[k]))
|
||||
snaps.append(best[0])
|
||||
print("solved and snapped activation row (k: code):")
|
||||
print(" ".join(f"{c:02x}" for c in snaps))
|
||||
# Derive the token -> lane assignment: for each candidate token row, solve and see whether
|
||||
# the resulting codes are exactly the bytes some lane group holds.
|
||||
print("token -> lane-group derivation (which solved row is a permutation of held bytes):")
|
||||
for cand in range(96):
|
||||
cmi, crh, cg = cand // 16, (cand % 16) // 8, (cand % 16) % 8
|
||||
Dc = np.zeros(COLS)
|
||||
held = []
|
||||
for widx in range(WEIGHT_TILES):
|
||||
for t in range(4):
|
||||
M = 16 * crh + t
|
||||
off = BASE + M * ORE * 4 + ACC_OFF
|
||||
acc = struct.unpack('<48I', dump[off:off + 48 * 4])
|
||||
for c in range(2):
|
||||
wt = 8 * widx + 2 * t + c
|
||||
reg = (cmi * 4 + widx) * 2 + crh
|
||||
Dc[wt] = f16((acc[reg] >> 16) if c else (acc[reg] & 0xFFFF))
|
||||
ops = struct.unpack(f'<{ORE}I', dump[BASE + M * ORE * 4:BASE + M * ORE * 4 + ORE * 4])
|
||||
held += [((ops[4 * cmi + (0 if crh == 0 else 1)] >> (8 * j)) & 0xFF) for j in range(4)]
|
||||
held += [((ops[4 * cmi + (2 if crh == 0 else 3)] >> (8 * j)) & 0xFF) for j in range(4)]
|
||||
ac, *_ = np.linalg.lstsq(W, Dc, rcond=None)
|
||||
res = float(np.abs(W @ ac - Dc).max())
|
||||
snaps_c = [min(codes, key=lambda cv: abs(cv[1] - ac[k]))[0] for k in range(32)]
|
||||
if sorted(snaps_c) == sorted(held) and res < 1e-9:
|
||||
print(f" token row {cand:2d} (mi {cmi}, rh {crh}, g {cg}) solves to exactly the 32 bytes "
|
||||
f"held by lanes {4*cg}..{4*cg+3} of m-tile {cmi} row half {crh}")
|
||||
break
|
||||
else:
|
||||
print(" no token row's solution matched the bytes held for its group")
|
||||
|
||||
# Derive the register-to-row assignment: solve every row exactly, then ask which
|
||||
# (lane, m-tile, row half) register set holds that row's values.
|
||||
print("register-to-row derivation:")
|
||||
solved = {}
|
||||
for cand in range(96):
|
||||
cmi, crh, cg = cand // 16, (cand % 16) // 8, (cand % 16) % 8
|
||||
Dc = np.zeros(COLS)
|
||||
for widx in range(WEIGHT_TILES):
|
||||
for t in range(4):
|
||||
M = 4 * cg + t
|
||||
off = BASE + M * ORE * 4 + ACC_OFF
|
||||
acc = struct.unpack('<48I', dump[off:off + 48 * 4])
|
||||
for c in range(2):
|
||||
Dc[8 * widx + 2 * t + c] = f16((acc[((cmi * 4 + widx) * 2 + crh)] >> (16 if c else 0)) & 0xFFFF)
|
||||
ac, *_ = np.linalg.lstsq(W, Dc, rcond=None)
|
||||
solved[cand] = [min(codes, key=lambda cv: abs(cv[1] - ac[k]))[0] for k in range(32)]
|
||||
|
||||
hits = {}
|
||||
for lane in range(32):
|
||||
ops = struct.unpack(f'<{ORE}I', dump[BASE + lane * ORE * 4:BASE + lane * ORE * 4 + ORE * 4])
|
||||
for mi in range(6):
|
||||
for rh in (0, 1):
|
||||
regs = (1, 3) if rh else (0, 2)
|
||||
held = [((ops[4 * mi + r] >> (8 * j)) & 0xFF) for r in regs for j in range(4)]
|
||||
for cand in range(96):
|
||||
want = list(solved[cand])
|
||||
ok = all(want.count(b) >= held.count(b) for b in set(held))
|
||||
if ok:
|
||||
hits.setdefault(cand, []).append((lane, mi, rh))
|
||||
for cand in sorted(hits)[:6]:
|
||||
(lane, mi, rh) = hits[cand][0]
|
||||
print(f" row {cand:2d} (mi {cand//16}, rh {(cand%16)//8}, g {(cand%16)%8}): held by "
|
||||
f"lane {lane:2d} m-tile {mi} row-half {rh} ({len(hits[cand])} candidates)")
|
||||
|
||||
# Search the plausible variants of the two maps and report the snap error for each. The
|
||||
# correct pair makes the solution land on e4m3 codes.
|
||||
print("map candidates (max snap-to-e4m3 error of the solved activation row):")
|
||||
import itertools
|
||||
|
||||
def wmap(k, wt, variant):
|
||||
if variant == 0:
|
||||
return waddr(k, wt)
|
||||
if variant == 1: # k16 halves swapped
|
||||
return waddr((k + 8) % 32 if (k % 16) < 8 else (k - 8), wt)
|
||||
if variant == 2: # the two (k>>1)&3 / (k>>3)&1 terms exchanged
|
||||
n_ = wt
|
||||
return (512 * (n_ >> 4) + 64 * (n_ & 7) + 16 * ((k >> 3) & 1) + 8 * ((n_ >> 3) & 1)
|
||||
+ 4 * ((k >> 1) & 3) + 2 * ((k >> 4) & 1) + (k & 1))
|
||||
if variant == 3: # the k&1 / (k>>1)&1 terms exchanged
|
||||
n_ = wt
|
||||
return (512 * (n_ >> 4) + 64 * (n_ & 7) + 16 * ((k >> 1) & 3) + 8 * ((n_ >> 3) & 1)
|
||||
+ 4 * ((k >> 3) & 1) + 2 * ((k >> 4) & 1) + ((k >> 1) & 1) * 1)
|
||||
return waddr(k, wt)
|
||||
|
||||
def dmap(tok, wt, variant):
|
||||
mi, rh, g = tok // 16, (tok % 16) // 8, (tok % 16) % 8
|
||||
M = 4 * g + ((wt % 8) // 2)
|
||||
off = BASE + M * ORE * 4 + ACC_OFF
|
||||
acc = struct.unpack('<48I', dump[off:off + 48 * 4])
|
||||
if variant == 0:
|
||||
widx, c = wt // 8, wt % 2
|
||||
reg = (mi * 4 + widx) * 2 + rh
|
||||
elif variant == 1: # n-tile major instead of m-tile major
|
||||
widx, c = wt // 8, wt % 2
|
||||
reg = (widx * 6 + mi) * 2 + rh
|
||||
elif variant == 2: # row halves swapped
|
||||
widx, c = wt // 8, wt % 2
|
||||
reg = (mi * 4 + widx) * 2 + (1 - rh)
|
||||
else:
|
||||
return None
|
||||
return f16((acc[reg] >> 16) if c else (acc[reg] & 0xFFFF))
|
||||
|
||||
for wv, dv in itertools.product(range(4), range(3)):
|
||||
Wv = np.zeros((COLS, 32))
|
||||
for wt in range(COLS):
|
||||
for k in range(32):
|
||||
Wv[wt, k] = e4m3_codes_value(wimg[WEIGHT_BASE + wmap(k, wt, wv)])
|
||||
Dv = np.zeros(COLS)
|
||||
ok = True
|
||||
for wt in range(COLS):
|
||||
v = dmap(tok, wt, dv)
|
||||
if v is None:
|
||||
ok = False
|
||||
break
|
||||
Dv[wt] = v
|
||||
if not ok:
|
||||
continue
|
||||
print(f" weight map {wv}, D map {dv}: {snap_error(Wv, Dv):.3e}")
|
||||
|
||||
# The body's waddr is a hand-derived inverse, i.e. a bit permutation of the k index onto the
|
||||
# byte-offset bits. Search the permutation instead of trusting it: 5 k bits onto the offset
|
||||
# bits {0,1,2,4,5} (offset bit 3 belongs to n), with the n terms fixed. The correct
|
||||
# permutation makes the solved activation row land exactly on e4m3 codes.
|
||||
# NOTE: this search's encoding of the permutation is not yet right -- candidates cluster
|
||||
# around 0.43 rather than reaching ~0, so the winner found here is not trustworthy. It does
|
||||
# establish the thing that matters at this stage: the maps in rrswin_enc0.hip give a snap
|
||||
# error of 15.8, i.e. they are *wrong*, and the e4m3-snap error is the test that shows it.
|
||||
print("k-bit permutation search (offset bits {0,1,2,4,5} <- k bits):")
|
||||
# Both maps are searched together: the k-bit permutation (5! of them) and the D register
|
||||
# order, since a wrong D map makes every permutation look wrong.
|
||||
def dval(order, wt):
|
||||
mi, rh, g = tok // 16, (tok % 16) // 8, (tok % 16) % 8
|
||||
M = 16 * rh + ((wt % 8) // 2) # derived lane mapping
|
||||
off = BASE + M * ORE * 4 + ACC_OFF
|
||||
acc = struct.unpack('<48I', dump[off:off + 48 * 4])
|
||||
widx, c = wt // 8, wt % 2
|
||||
if order == 0:
|
||||
reg = (mi * 4 + widx) * 2 + rh
|
||||
elif order == 1:
|
||||
reg = (widx * 6 + mi) * 2 + rh
|
||||
elif order == 2:
|
||||
reg = (mi * 4 + widx) * 2 + (1 - rh)
|
||||
else:
|
||||
reg = (mi * 4 + (3 - widx)) * 2 + rh
|
||||
return f16((acc[reg] >> 16) if c else (acc[reg] & 0xFFFF))
|
||||
|
||||
# A k-bit permutation only relabels which k holds which value, so it cannot change the
|
||||
# *set* of solved values and the snap statistic is blind to it. The test that does work is
|
||||
# permutation-invariant and uses the registers themselves: do the solved values equal the
|
||||
# bytes the lanes hold? Only the D map can be searched this way.
|
||||
print("D-map search by multiset match with the held bytes:")
|
||||
Wbase = np.zeros((COLS, 32))
|
||||
for wt in range(COLS):
|
||||
for k in range(32):
|
||||
Wbase[wt, k] = e4m3_codes_value(wimg[WEIGHT_BASE + waddr(k, wt)])
|
||||
held_all = set()
|
||||
for lane in range(32):
|
||||
ops = struct.unpack(f'<{ORE}I', dump[BASE + lane * ORE * 4:BASE + lane * ORE * 4 + ORE * 4])
|
||||
for mi in range(6):
|
||||
for r in range(4):
|
||||
for j in range(4):
|
||||
held_all.add((ops[4 * mi + r] >> (8 * j)) & 0xFF)
|
||||
for order in range(4):
|
||||
Dv = np.array([dval(order, wt) for wt in range(COLS)])
|
||||
a, *_ = np.linalg.lstsq(Wbase, Dv, rcond=None)
|
||||
codes_all = e4m3_codes()
|
||||
snaps = [min(codes_all, key=lambda cv: abs(cv[1] - a[k]))[0] for k in range(32)]
|
||||
inside = sum(1 for c in snaps if c in held_all)
|
||||
# strict: the solved row must equal the union of the bytes the four lanes of this token's
|
||||
# group hold for this row half, as a multiset
|
||||
g_ = (tok % 16) % 8
|
||||
strict = []
|
||||
for lane in range(4 * g_, 4 * g_ + 4):
|
||||
ops = struct.unpack(f'<{ORE}I', dump[BASE + lane * ORE * 4:BASE + lane * ORE * 4 + ORE * 4])
|
||||
regs = (1, 3) if order == 2 else (0, 2)
|
||||
for r in regs:
|
||||
for j in range(4):
|
||||
strict.append((ops[4 * (tok // 16) + r] >> (8 * j)) & 0xFF)
|
||||
exact = sorted(snaps) == sorted(strict)
|
||||
print(f" D order {order}: {inside:2d} of 32 solved values are values the payload holds"
|
||||
f"{' EXACT multiset match with the group\'s held bytes' if exact else ''}")
|
||||
# Derive the placement: which byte position holds each k. The solved row gives the value for
|
||||
# every k in the *weight* labelling (the one waddr uses, which the exact-w-values test just
|
||||
# corroborated). The payload gives the bytes and their (lane, register, byte) positions. A
|
||||
# position is the k label whose solved value it holds, so matching them derives the map the
|
||||
# engine must use for its activation image to be consistent with the weights.
|
||||
hit, solved, held = test_zluda_map(dump, wimg, tok, D, W)
|
||||
print(f"ZLUDA's documented map (k = 4*(lane&3)+j, row = lane>>2): {hit} of 32 k match exactly")
|
||||
if hit < 32:
|
||||
print(" k : solved held")
|
||||
for k in range(32):
|
||||
mark = "" if held.get(k) == solved[k] else " <-- differs"
|
||||
print(f" {k:2d}: {solved[k]:02x} {held.get(k, 0):02x}{mark}")
|
||||
|
||||
print("placement: which byte position holds each k label")
|
||||
a_all, *_ = np.linalg.lstsq(Wbase, D, rcond=None)
|
||||
codes_all = e4m3_codes()
|
||||
label_of = {}
|
||||
for k in range(32):
|
||||
snapped = min(codes_all, key=lambda cv: abs(cv[1] - a_all[k]))[0]
|
||||
label_of.setdefault(snapped, []).append(k)
|
||||
# The solved row's values appear somewhere in the payload but not in the group I assumed, so
|
||||
# the row-to-group mapping is what to derive: try every (g, row half) and report which group's
|
||||
# bytes are exactly the solved row's values.
|
||||
# wmma_layout.h says "A operand lane l = row r" with r = l & 15, while rrswin_enc0.hip's
|
||||
# header says row g = lane >> 2. They are different conventions, so test the other one: a
|
||||
# row's 16 bytes per lane over two lanes (m and m+16) rather than four lanes of a group.
|
||||
print(" row-to-pair search (row = lane & 15, 16 bytes per lane):")
|
||||
solved_multi = sorted(label_of.keys())
|
||||
for row in range(16):
|
||||
held = []
|
||||
for lane in (row, row + 16):
|
||||
ops = struct.unpack(f'<{ORE}I', dump[BASE + lane * ORE * 4:BASE + lane * ORE * 4 + ORE * 4])
|
||||
for rr in range(4):
|
||||
for j in range(4):
|
||||
held.append((ops[4 * (tok // 16) + rr] >> (8 * j)) & 0xFF)
|
||||
hit = sum(1 for b in solved_multi if b in held)
|
||||
miss = [b for b in solved_multi if b not in held]
|
||||
if hit >= 28:
|
||||
print(f" row {row:2d}: {hit} of 32 solved values held by lanes {row},{row+16}"
|
||||
f"{' EXACT' if not miss else ''}")
|
||||
|
||||
print(" row-to-group search:")
|
||||
solved_multi = sorted(label_of.keys())
|
||||
for cand_g in range(8):
|
||||
for cand_r in (0, 1):
|
||||
held = []
|
||||
for lane in range(4 * cand_g, 4 * cand_g + 4):
|
||||
ops = struct.unpack(f'<{ORE}I', dump[BASE + lane * ORE * 4:BASE + lane * ORE * 4 + ORE * 4])
|
||||
regs = (1, 3) if cand_r else (0, 2)
|
||||
for rr in regs:
|
||||
for j in range(4):
|
||||
held.append((ops[4 * (tok // 16) + rr] >> (8 * j)) & 0xFF)
|
||||
hit = sum(1 for b in solved_multi if b in held)
|
||||
if hit >= 28:
|
||||
print(f" g {cand_g}, row half {cand_r}: {hit} of {len(solved_multi)} solved values held")
|
||||
|
||||
g_ = (tok % 16) % 8
|
||||
print(" group lanes " + " ".join(str(4 * g_ + t) for t in range(4)) + ":")
|
||||
for lane in range(4 * g_, 4 * g_ + 4):
|
||||
ops = struct.unpack(f'<{ORE}I', dump[BASE + lane * ORE * 4:BASE + lane * ORE * 4 + ORE * 4])
|
||||
for r in range(4):
|
||||
row = "[row g]" if r in (0, 2) else "[row g+8]"
|
||||
ks = []
|
||||
for j in range(4):
|
||||
b = (ops[4 * (tok // 16) + r] >> (8 * j)) & 0xFF
|
||||
cand = label_of.get(b, [])
|
||||
ks.append(f"{b:02x}->k{cand[0]}" if len(cand) == 1 else f"{b:02x}->?")
|
||||
print(f" lane {lane:2d} reg {r} {row}: " + " ".join(ks))
|
||||
|
||||
print("bytes held by the lanes of this token's group:")
|
||||
for M in sorted(lane_bytes):
|
||||
print(f" lane {M:2d}: " + " ".join(f"{b:02x}" for b in lane_bytes[M]))
|
||||
|
||||
|
||||
def e4m3_codes_value(code):
|
||||
mag = code & 0x7F
|
||||
if mag >= 8:
|
||||
bits = (((mag >> 3) + 120) << 23) | ((mag & 7) << 20)
|
||||
else:
|
||||
bits = struct.unpack('<I', struct.pack('<f', float(mag) * 2.0 ** -9))[0]
|
||||
return struct.unpack('<f', struct.pack('<I', bits | ((code & 0x80) << 24)))[0]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Executable
+114
@@ -0,0 +1,114 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Host model of enc0 phase 8 from the oracle's own operands, to localise the pilot's residual.
|
||||
|
||||
Reads the combined p8c dump (32 operand registers then 48 accumulator registers per lane per
|
||||
block) and does everything the pilot does except on the host, in f64, from the *documented*
|
||||
layouts:
|
||||
- the A fragments (the first 24 registers) staged into an activation image with the layout the
|
||||
enc0 body documents,
|
||||
- the weights taken from the weight image through waddr (the same lookup the engine's prep uses),
|
||||
- D = the sum over k of weight * activation, compared with the oracle's own D.
|
||||
If this reproduces the oracle's D, the operand model is right and any disagreement the pilot shows
|
||||
is the engine or its glue; if it does not, the operand model is what is wrong.
|
||||
|
||||
usage: phase8_model.py <oracle-p8c-dir>
|
||||
"""
|
||||
import bisect
|
||||
import pathlib
|
||||
import struct
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, "/home/wouter/git/d4r/kernels/rr")
|
||||
import rr_layer_spec as R # noqa: E402
|
||||
|
||||
BASE = 46000000
|
||||
SLOT = 16384
|
||||
ORE = 32 # operand registers per lane
|
||||
AREGS = 24 # A fragments (6 m-tiles of 4)
|
||||
WEIGHT_BASE = 21696
|
||||
|
||||
|
||||
def find(dirname, needle):
|
||||
for p in sorted(pathlib.Path(dirname).iterdir()):
|
||||
if "swin_enc0_kernel" in p.name and needle in p.name:
|
||||
return p
|
||||
raise SystemExit(f"no dump with {needle} in {dirname}")
|
||||
|
||||
|
||||
def e4m3_to_f32(code):
|
||||
mag = code & 0x7F
|
||||
if mag >= 8:
|
||||
bits = (((mag >> 3) + 120) << 23) | ((mag & 7) << 20)
|
||||
else:
|
||||
bits = struct.unpack('<I', struct.pack('<f', float(mag) * 2.0 ** -9))[0]
|
||||
return struct.unpack('<f', struct.pack('<I', bits | ((code & 0x80) << 24)))[0]
|
||||
|
||||
|
||||
def waddr(k, n):
|
||||
return (512 * (n >> 4) + 64 * (n & 7) + 16 * ((k >> 1) & 3) + 8 * ((n >> 3) & 1)
|
||||
+ 4 * ((k >> 3) & 1) + 2 * ((k >> 4) & 1) + (k & 1))
|
||||
|
||||
|
||||
def f16(bits):
|
||||
return struct.unpack('<e', struct.pack('<H', bits & 0xFFFF))[0]
|
||||
|
||||
|
||||
def main():
|
||||
d = sys.argv[1] if len(sys.argv) > 1 else "/tmp/d4r-rr-perf/d-p8c"
|
||||
dump = find(d, "arg048").read_bytes()
|
||||
wimg = find(d, "arg040").read_bytes()
|
||||
|
||||
block, lane = 0, 0
|
||||
off = BASE + block * SLOT + lane * (ORE + 48) * 4
|
||||
|
||||
# Each lane supplies only eight K positions for each token row. Assemble
|
||||
# all 32 lanes before taking a dot product; the combined payload has 80
|
||||
# registers per lane, not 32.
|
||||
img = [[0] * 32 for _ in range(96)]
|
||||
for source_lane in range(32):
|
||||
source_off = BASE + block * SLOT + source_lane * (ORE + 48) * 4
|
||||
ops = struct.unpack(f'<{ORE}I', dump[source_off:source_off + ORE * 4])
|
||||
for mi in range(6):
|
||||
f = ops[4 * mi:4 * mi + 4]
|
||||
for j in range(4):
|
||||
k0 = 2 * (source_lane & 3) + (j & 1) + 16 * ((j >> 1) & 1)
|
||||
k2 = 8 + k0
|
||||
img[16 * mi + (source_lane >> 2)][k0] = (f[0] >> (8 * j)) & 0xFF
|
||||
img[16 * mi + (source_lane >> 2) + 8][k0] = (f[1] >> (8 * j)) & 0xFF
|
||||
img[16 * mi + (source_lane >> 2)][k2] = (f[2] >> (8 * j)) & 0xFF
|
||||
img[16 * mi + (source_lane >> 2) + 8][k2] = (f[3] >> (8 * j)) & 0xFF
|
||||
|
||||
# the oracle's D, this lane
|
||||
doff = off + ORE * 4
|
||||
acc = struct.unpack('<48I', dump[doff:doff + 48 * 4])
|
||||
g, t = lane >> 2, lane & 3
|
||||
|
||||
# This lane holds only the columns 2t and 2t+1 of each 8-wide n-tile, and only the token rows
|
||||
# 16*mi + 8*rh + g -- reading any other (weight, token) from its registers reads another
|
||||
# lane's data.
|
||||
print("weight token model oracle delta")
|
||||
for wt in (2 * t, 2 * t + 1):
|
||||
for mi in (0, 1): # this lane holds one token row per row half
|
||||
for rh in (0, 1):
|
||||
tok = 16 * mi + 8 * rh + g
|
||||
total = 0.0
|
||||
for k in range(32):
|
||||
total += e4m3_to_f32(wimg[WEIGHT_BASE + waddr(k, wt)]) * e4m3_to_f32(img[tok][k])
|
||||
total = struct.unpack('<e', struct.pack('<e', total))[0]
|
||||
widx = wt // 8
|
||||
reg = (mi * 4 + widx) * 2 + rh
|
||||
pair = acc[reg]
|
||||
c = wt % 2
|
||||
want = f16((pair >> 16) if c else (pair & 0xFFFF))
|
||||
# which (register, half) of this lane does the model's value actually equal?
|
||||
hits = []
|
||||
for rb in range(48):
|
||||
for hh in (0, 1):
|
||||
v = f16((acc[rb] >> 16) if hh else (acc[rb] & 0xFFFF))
|
||||
if abs(v - total) < 1e-4 * (1 + abs(total)):
|
||||
hits.append(f"r{rb}.{hh}")
|
||||
print(f"{wt:6d} {tok:6d} {total:+.6f} {want:+.6f} delta {total - want:+.6f} matches {hits[:4]}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,193 @@
|
||||
// Compare enc0 phase 8's native staged-image GEMM with a combined p8c oracle
|
||||
// capture: 24 A fragments, 8 B fragments and 48 D registers per lane.
|
||||
// The engine produces D[weight][token]; the oracle stores D[token][weight].
|
||||
// Every HIP call is checked, and a mismatch produces a nonzero exit status.
|
||||
//
|
||||
// hipcc --offload-arch=gfx1201 -DD4R_TEX_FP8 -DD4R_WMMA_LAYOUT=12 -O2 \
|
||||
// -o /tmp/pilot kernels/tools/rrswin_pilot_enc0_p8.cpp
|
||||
// /tmp/pilot <oracle-p8c-dir>
|
||||
#include <hip/hip_runtime.h>
|
||||
#include <cstdio>
|
||||
#include <cstdint>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include <dirent.h>
|
||||
|
||||
#include "../rr/rrswin_engine.h"
|
||||
|
||||
#define CHECK(expr) \
|
||||
do \
|
||||
{ \
|
||||
hipError_t e_ = (expr); \
|
||||
if (e_ != hipSuccess) \
|
||||
{ \
|
||||
std::printf("HIP error at %s:%d: %s -> %s\n", __FILE__, __LINE__, #expr, \
|
||||
hipGetErrorString(e_)); \
|
||||
std::exit(2); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
static const int kDebugSlot = 16384; // swin_debug_slot(80) for the combined stage
|
||||
static const int kDebugBase = 46000000;
|
||||
|
||||
static std::vector<uint8_t> slurp(const std::string& dir, const char* what, size_t* size_out)
|
||||
{
|
||||
// the dump names are launch-<n>-<kernel>-<arg>-<kind>-<base>-fmt<n>.bin; matching the pieces
|
||||
// is more robust than a glob over a name that contains '+0x0'
|
||||
DIR* d = opendir(dir.c_str());
|
||||
if (d == nullptr)
|
||||
{
|
||||
std::printf("cannot open %s\n", dir.c_str());
|
||||
std::exit(2);
|
||||
}
|
||||
std::string found;
|
||||
for (dirent* e = readdir(d); e != nullptr; e = readdir(d))
|
||||
{
|
||||
const std::string name = e->d_name;
|
||||
if (name.find("swin_enc0_kernel") != std::string::npos &&
|
||||
name.find(std::string("-") + what + "-") != std::string::npos)
|
||||
{
|
||||
found = dir + "/" + name;
|
||||
break;
|
||||
}
|
||||
}
|
||||
closedir(d);
|
||||
if (found.empty())
|
||||
{
|
||||
std::printf("no dump with %s in %s\n", what, dir.c_str());
|
||||
std::exit(2);
|
||||
}
|
||||
FILE* f = std::fopen(found.c_str(), "rb");
|
||||
std::fseek(f, 0, SEEK_END);
|
||||
const long n = std::ftell(f);
|
||||
std::fseek(f, 0, SEEK_SET);
|
||||
std::vector<uint8_t> data((size_t)n);
|
||||
if (std::fread(data.data(), 1, data.size(), f) != data.size())
|
||||
{
|
||||
std::printf("short read of %s\n", found.c_str());
|
||||
std::exit(2);
|
||||
}
|
||||
std::fclose(f);
|
||||
if (size_out)
|
||||
*size_out = data.size();
|
||||
return data;
|
||||
}
|
||||
|
||||
// One phase over one debug slot: stage six m-tiles of A fragments, prep the weights, run the
|
||||
// phase, write D[weight][token] out. The weight preparation reads the image through the
|
||||
// documented enc0 layout, offset by the phase's byte base.
|
||||
__global__ void pilot_kernel(const uint32_t* frag, const uint8_t* wimg, int weight_base,
|
||||
float* out)
|
||||
{
|
||||
__shared__ uint8_t img[96 * 32];
|
||||
__shared__ uint8_t weight_slots[32 * 2 * 16];
|
||||
for (int i = threadIdx.x; i < 96 * 32; i += blockDim.x)
|
||||
img[i] = 0;
|
||||
__syncthreads();
|
||||
rrswin::swin_stage_fragment(img, 32, frag + (size_t)24 * (blockIdx.x * 32u + threadIdx.x), 6);
|
||||
if (threadIdx.x == 0)
|
||||
rrswin::swin_prep_phase(weight_slots, 0, 32, 1, wimg + weight_base, rrswin::swin_waddr_enc0);
|
||||
__syncthreads();
|
||||
const rrswin::Phase phase = {2, 6, 1, 0, 1, 0};
|
||||
f8v acc[12];
|
||||
rrswin::swin_phase(phase, (const kslot*)weight_slots, 32, img, 32, nullptr, 0, acc);
|
||||
for (int tile = 0; tile < 12; ++tile)
|
||||
for (int i = 0; i < 8; ++i)
|
||||
out[(size_t)blockIdx.x * 3072 + (tile * 8 + i) * 32 + threadIdx.x] = acc[tile][i];
|
||||
}
|
||||
|
||||
int main(int argc, char** argv)
|
||||
{
|
||||
if (argc < 2)
|
||||
{
|
||||
std::printf("usage: %s <oracle-p8c-dir>\n", argv[0]);
|
||||
return 2;
|
||||
}
|
||||
size_t no = 0, nw = 0;
|
||||
const std::vector<uint8_t> p8o = slurp(argv[1], "arg048", &no); // combined: 32 operands + 48 accumulators
|
||||
const std::vector<uint8_t> p8 = p8o;
|
||||
const std::vector<uint8_t> wimg = slurp(argv[1], "arg040", &nw);
|
||||
std::printf("arg048 (combined p8c) %zu bytes, arg040 %zu bytes\n", no, nw);
|
||||
if (no < (size_t)kDebugBase + 64 * kDebugSlot)
|
||||
{
|
||||
std::printf("dumps are too small for the payload window\n");
|
||||
return 2;
|
||||
}
|
||||
|
||||
// the phase's A fragments: 24 registers per lane, 32 lanes, per block
|
||||
std::vector<uint32_t> frag(64 * 32 * 24);
|
||||
for (int b = 0; b < 64; ++b)
|
||||
for (int lane = 0; lane < 32; ++lane)
|
||||
{
|
||||
const uint8_t* p = p8o.data() + kDebugBase + (size_t)b * kDebugSlot + (size_t)lane * 80 * 4;
|
||||
for (int r = 0; r < 24; ++r)
|
||||
std::memcpy(&frag[((size_t)b * 32 + lane) * 24 + r], p + 4 * r, 4);
|
||||
}
|
||||
|
||||
uint32_t* dfrag;
|
||||
uint8_t* dwimg;
|
||||
float* dout;
|
||||
CHECK(hipMalloc(&dfrag, frag.size() * 4));
|
||||
CHECK(hipMalloc(&dwimg, wimg.size()));
|
||||
CHECK(hipMalloc(&dout, 64 * 3072 * sizeof(float)));
|
||||
CHECK(hipMemcpy(dfrag, frag.data(), frag.size() * 4, hipMemcpyHostToDevice));
|
||||
CHECK(hipMemcpy(dwimg, wimg.data(), wimg.size(), hipMemcpyHostToDevice));
|
||||
CHECK(hipMemset(dout, 0, 64 * 3072 * sizeof(float)));
|
||||
pilot_kernel<<<64, 32>>>(dfrag, dwimg, 21696, dout);
|
||||
CHECK(hipGetLastError());
|
||||
CHECK(hipDeviceSynchronize());
|
||||
|
||||
std::vector<float> got(64 * 3072);
|
||||
CHECK(hipMemcpy(got.data(), dout, got.size() * 4, hipMemcpyDeviceToHost));
|
||||
|
||||
// Compare the engine's D[weight][token] with the oracle's D[token][weight].
|
||||
//
|
||||
// The two sides do not use the same lane for the same element, which is what earlier attempts
|
||||
// got wrong by reusing one lane for both:
|
||||
// engine: acc[mt][nt][i] of lane L holds D[row = 16*mt + i + 8*(L>>4)][col = 16*nt + (L&15)]
|
||||
// oracle: the p8 payload's register (mi, widx, rh) of lane M holds D[row = 16*mi + 8*rh + g]
|
||||
// [col = 8*widx + 2*t + c], with g = M>>2 and t = M&3
|
||||
// so an element identified by (weight w, token k) lives in lane 16*((w%16)/8) + (k%16) on the
|
||||
// engine side and in lane 4*((k%16)%8) + (w%8)/2 on the oracle side.
|
||||
int bad = 0, checked = 0;
|
||||
for (int b = 0; b < 64; ++b)
|
||||
{
|
||||
const uint8_t* d = p8.data() + kDebugBase + (size_t)b * kDebugSlot;
|
||||
for (int wt = 0; wt < 32; ++wt)
|
||||
for (int tok = 0; tok < 96; ++tok)
|
||||
{
|
||||
const int mt = wt / 16, i = (wt % 16) & 7, h = (wt % 16) / 8, nt = tok / 16;
|
||||
const int eLane = 16 * h + (tok % 16);
|
||||
const float gotv = got[(size_t)b * 3072 + ((mt * 6 + nt) * 8 + i) * 32 + eLane];
|
||||
|
||||
const int mi = tok / 16, rh = (tok % 16) / 8, g = (tok % 16) % 8;
|
||||
const int widx = wt / 8, t = (wt % 8) / 2, c = wt % 2;
|
||||
const int oLane = 4 * g + t;
|
||||
const uint16_t* regs = (const uint16_t*)(d + (size_t)oLane * 80 * 4 + 32 * 4);
|
||||
const int reg = (mi * 4 + widx) * 2 + rh;
|
||||
const uint32_t u = (uint32_t)regs[2 * reg] | ((uint32_t)regs[2 * reg + 1] << 16);
|
||||
const uint16_t bits = (uint16_t)(c ? (u >> 16) : (u & 0xffff));
|
||||
const _Float16 hb = *reinterpret_cast<const _Float16*>(&bits);
|
||||
const float oracle_f = (float)hb;
|
||||
// the engine's accumulator is f32 (the WMMA's), the oracle's mma accumulates in
|
||||
// f16, so the comparison is at f16: round the engine's value to f16 first
|
||||
const _Float16 got_h = (_Float16)gotv;
|
||||
const float got_r = (float)got_h;
|
||||
++checked;
|
||||
if (got_r != oracle_f)
|
||||
{
|
||||
if (bad < 4)
|
||||
std::printf(" block %2d weight %2d token %2d: engine %+.6f oracle %+.6f\n",
|
||||
b, wt, tok, got_r, oracle_f);
|
||||
++bad;
|
||||
}
|
||||
}
|
||||
}
|
||||
std::printf("phase-8 pilot: checked %d, mismatched %d\n", checked, bad);
|
||||
CHECK(hipFree(dfrag));
|
||||
CHECK(hipFree(dwimg));
|
||||
CHECK(hipFree(dout));
|
||||
return bad != 0;
|
||||
}
|
||||
Executable
+30
@@ -0,0 +1,30 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Dump the mma statements of a swin entry's phase range with their A/B/C/D registers.
|
||||
|
||||
Usage: swin_mma.py LAYER LO_STMT [HI_STMT]
|
||||
LAYER is a rr_layer_spec alias (enc0, ...). One line per mma, with the
|
||||
statement index, the physical PTX line and the register operands in order.
|
||||
"""
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "rr"))
|
||||
import rr_layer_spec as R # noqa: E402
|
||||
|
||||
|
||||
def main():
|
||||
layer = sys.argv[1] if len(sys.argv) > 1 else "enc0"
|
||||
lo = int(sys.argv[2]) if len(sys.argv) > 2 else 0
|
||||
hi = int(sys.argv[3]) if len(sys.argv) > 3 else 10 ** 9
|
||||
path, entry = R.resolve(layer, None, None)
|
||||
lines = open(path).read().split("\n")
|
||||
first, last = next((a, b) for (nm, a, b) in R.find_entries(lines) if nm == entry)
|
||||
stmts, _ = R.linearise(lines, first, last)
|
||||
for s in stmts:
|
||||
if s.is_mma and lo <= s.idx <= hi:
|
||||
print(f"s{s.idx} L{s.line}: " + " ".join(re.findall(r"%r\d+", s.text)))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Executable
+165
@@ -0,0 +1,165 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Emit the oracle debug stages for cuda_dldn_engine_swin_enc0_kernel, phases 8..22.
|
||||
|
||||
For each phase it reads NVIDIA's PTX, keeps the mma whose D pair is not another mma's
|
||||
C operand inside the same phase (the last k-step of an accumulating chain), and lists
|
||||
those tiles in mma order, which is m-tile major and n-tile minor -- the order the
|
||||
native body's stage_regs() writes. Operand stages list the phase's A fragments, B
|
||||
fragment pairs and C seeds as the mma read them.
|
||||
|
||||
Writes kernels/rr/swin_enc0_debug_stages.py, which make_ptx.py folds into its own
|
||||
SWIN_DEBUG_STAGES.
|
||||
"""
|
||||
import re
|
||||
import sys
|
||||
import os
|
||||
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "rr"))
|
||||
import rr_layer_spec as R # noqa: E402
|
||||
|
||||
ENTRY = "cuda_dldn_engine_swin_enc0_kernel"
|
||||
ZERO_C = "%r10854" # mov.u32 %r10854, 0 (PTX line 3920): the unfused phases' C
|
||||
OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "rr",
|
||||
"swin_enc0_debug_stages.py")
|
||||
|
||||
PHASES = {
|
||||
8: (8057, 8080), 9: (8088, 8111), 10: (8283, 8330), 11: (12169, 12216),
|
||||
12: (12272, 12287), 13: (13175, 13190), 14: (13900, 13915), 15: (13980, 13995),
|
||||
16: (14596, 14611), 17: (14676, 14691), 18: (15292, 15307), 19: (15372, 15387),
|
||||
20: (15988, 16003), 21: (16418, 16449), 22: (16607, 16614),
|
||||
}
|
||||
|
||||
HEADER = '''#!/usr/bin/env python3
|
||||
"""Oracle debug stages for cuda_dldn_engine_swin_enc0_kernel's phases 8..22.
|
||||
|
||||
Each stage is a (file line after which the store block goes, registers in order) tuple for
|
||||
kernels/tex/make_ptx.py's SWIN_DEBUG_STAGES, so D4R_RRSWIN_DEBUG=p8 (p8o, p9, p10, ...)
|
||||
builds NVIDIA's own body with that stage's registers stored into the plane arena, which
|
||||
kernels/rr/rrswin_enc0.hip's D4R_RRSWIN_DEBUG_STAGE reproduces from its own registers.
|
||||
|
||||
GENERATED by kernels/tools/swin_stages.py from the corpus PTX -- do not edit. An
|
||||
accumulator stage keeps the mma whose D pair no other mma of the phase reads as its C (the
|
||||
last k-step of an accumulating chain) and lists the tiles in mma order: m-tile major,
|
||||
n-tile minor, the row pair (row g, row g+8) at each base. An operand stage lists the A
|
||||
fragments, the B fragment pairs and the C seeds as the mma read them.
|
||||
"""
|
||||
|
||||
|
||||
def swin_pairs(bases):
|
||||
"""One register pair per base: a tile's (row g, row g+8) D registers."""
|
||||
return sum(([f"%r{b + rh}" for rh in (0, 1)] for b in bases), [])
|
||||
|
||||
|
||||
def swin_quads(bases):
|
||||
"""One A fragment per base: four b32."""
|
||||
return sum(([f"%r{b + o}" for o in range(4)] for b in bases), [])
|
||||
|
||||
|
||||
def stages():
|
||||
"""The stage table: name -> (splice line, registers)."""
|
||||
return {
|
||||
'''
|
||||
|
||||
|
||||
def phase_mmas(stmts, p):
|
||||
lo, hi = PHASES[p]
|
||||
out = []
|
||||
for s in stmts:
|
||||
if not s.is_mma or not (lo <= s.idx <= hi):
|
||||
continue
|
||||
r = re.findall(r"%r\d+", s.text)
|
||||
out.append(dict(D=(r[0], r[1]), A=tuple(r[2:6]), B=(r[6], r[7]), C=(r[8], r[9]),
|
||||
line=s.line))
|
||||
return out
|
||||
|
||||
|
||||
def outputs(mmas):
|
||||
"""The last k-step of every chain: the mma whose D pair no other mma reads as its C."""
|
||||
cs = set()
|
||||
for m in mmas:
|
||||
cs.update(m["C"])
|
||||
return [m for m in mmas if m["D"][0] not in cs and m["D"][1] not in cs]
|
||||
|
||||
|
||||
def dbase(mmas):
|
||||
return [int(m["D"][0][2:]) for m in outputs(mmas)]
|
||||
|
||||
|
||||
def abase(mmas):
|
||||
seen, out = set(), []
|
||||
for m in mmas:
|
||||
if m["A"][0] in seen:
|
||||
continue
|
||||
seen.add(m["A"][0])
|
||||
out.append(int(m["A"][0][2:]))
|
||||
return out
|
||||
|
||||
|
||||
def bbase(mmas):
|
||||
seen, out = set(), []
|
||||
for m in mmas:
|
||||
if m["B"][0] in seen:
|
||||
continue
|
||||
seen.add(m["B"][0])
|
||||
out.append(int(m["B"][0][2:]))
|
||||
return out
|
||||
|
||||
|
||||
def cseed(mmas):
|
||||
"""The distinct C registers that are not another mma's D inside the same phase (an
|
||||
accumulating chain's previous step is not a seed), in first-appearance order."""
|
||||
ds = set()
|
||||
for m in mmas:
|
||||
ds.update(m["D"])
|
||||
seen, out = set(), []
|
||||
for m in mmas:
|
||||
for r in m["C"]:
|
||||
if r in seen or r in ds or r == ZERO_C:
|
||||
continue
|
||||
seen.add(r)
|
||||
out.append(r)
|
||||
return out
|
||||
|
||||
|
||||
def nums(xs):
|
||||
return ", ".join(str(x) for x in xs)
|
||||
|
||||
|
||||
def entries(p, mmas, after):
|
||||
"""The two dict entries of one phase: its output tiles, and its mma operands."""
|
||||
d, a, b, cs = dbase(mmas), abase(mmas), bbase(mmas), cseed(mmas)
|
||||
parts = [f"swin_quads([{nums(a)}])", f"swin_pairs([{nums(b)}])"]
|
||||
n = 4 * len(a) + 2 * len(b)
|
||||
kinds = "A fragments, B fragment pairs"
|
||||
if cs:
|
||||
parts.append("[" + ", ".join(f'"{r}"' for r in cs) + "]")
|
||||
n += len(cs)
|
||||
kinds += ", C seeds"
|
||||
pad = -n % 4 # swin_debug_block stores four registers per v4, so pad with the zero
|
||||
if pad:
|
||||
parts.append("[" + ", ".join([f'"{ZERO_C}"'] * pad) + "]")
|
||||
n += pad
|
||||
return [f" # phase {p}: {len(mmas)} mma, {len(d)} output tiles, "
|
||||
f"splice after file line {after}",
|
||||
f' "p{p}": ({after}, swin_pairs([{nums(d)}])),',
|
||||
f' "p{p}o": ({after}, {" + ".join(parts)}),'
|
||||
f" # {kinds}; {n} registers"]
|
||||
|
||||
|
||||
def main():
|
||||
path, entry_name = R.resolve("enc0", None, None)
|
||||
lines = open(path).read().split("\n")
|
||||
first, last = next((a, b) for (nm, a, b) in R.find_entries(lines) if nm == entry_name)
|
||||
stmts, _ = R.linearise(lines, first, last)
|
||||
body = []
|
||||
for p in sorted(PHASES):
|
||||
mmas = phase_mmas(stmts, p)
|
||||
after = max(m["line"] for m in mmas) + 3
|
||||
assert lines[after - 1].rstrip().endswith(";"), lines[after - 1]
|
||||
body.extend(entries(p, mmas, after))
|
||||
with open(OUT, "w") as f:
|
||||
f.write(HEADER + "\n".join(body) + "\n }\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in new issue
Block a user