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,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;
|
||||
}
|
||||
Reference in new issue
Block a user