Validate gfx1201 WMMA layout and compile first K kernel
This commit is contained in:
Codex committed
1 parent
64669c2c36
commit
6760294f60
10 files changed
+404
-10
No files matched your search
@@ -0,0 +1,127 @@
|
||||
#include "hip_api.h"
|
||||
#include <array>
|
||||
#include <cmath>
|
||||
#include <limits>
|
||||
|
||||
namespace {
|
||||
uint16_t half_integer(int value)
|
||||
{
|
||||
const unsigned n = static_cast<unsigned>(std::abs(value));
|
||||
if (n > 3u) throw std::runtime_error("test half integer out of range");
|
||||
const uint16_t bits[] = {0u, 0x3c00u, 0x4000u, 0x4200u};
|
||||
return static_cast<uint16_t>(bits[n] | (value < 0 ? 0x8000u : 0u));
|
||||
}
|
||||
struct Allocation {
|
||||
d4r::diag::HipApi& hip;
|
||||
void* pointer = nullptr;
|
||||
Allocation(d4r::diag::HipApi& api, size_t bytes) : hip(api) {
|
||||
hip.check(hip.hipMalloc(&pointer, bytes), "hipMalloc(WMMA)");
|
||||
}
|
||||
~Allocation() { if (pointer) (void)hip.hipFree(pointer); }
|
||||
};
|
||||
}
|
||||
|
||||
int main(int argc, char** argv)
|
||||
{
|
||||
using namespace d4r::diag;
|
||||
start();
|
||||
try {
|
||||
Args args(argc, argv);
|
||||
if (args.module.empty()) throw std::runtime_error("--module is required");
|
||||
HipApi hip(args.hip_root);
|
||||
hipDeviceProp_t properties{};
|
||||
hip.select_gfx1201(args.device, properties);
|
||||
hipModule_t module = nullptr;
|
||||
hip.check(hip.hipModuleLoad(&module, args.module.c_str()), "hipModuleLoad(WMMA)");
|
||||
struct ModuleCleanup {
|
||||
HipApi& hip; hipModule_t module;
|
||||
~ModuleCleanup() { if (module) (void)hip.hipModuleUnload(module); }
|
||||
} cleanup{hip, module};
|
||||
hipFunction_t raw_kernel = nullptr, adapter_kernel = nullptr, swap_kernel = nullptr;
|
||||
hip.check(hip.hipModuleGetFunction(&raw_kernel, module, "d4r_wmma_gfx1201"),
|
||||
"hipModuleGetFunction(WMMA)");
|
||||
hip.check(hip.hipModuleGetFunction(&adapter_kernel, module, "d4r_wmma_legacy_contract"),
|
||||
"hipModuleGetFunction(WMMA adapter)");
|
||||
hip.check(hip.hipModuleGetFunction(&swap_kernel, module, "d4r_wmma_swap_test"),
|
||||
"hipModuleGetFunction(WMMA swap)");
|
||||
constexpr size_t elements = 16u * 16u;
|
||||
std::array<uint16_t, elements> a{}, b{};
|
||||
std::array<float, elements> initial{}, output{};
|
||||
Allocation da(hip, sizeof(a)), db(hip, sizeof(b));
|
||||
Allocation dc(hip, sizeof(initial)), dout(hip, sizeof(output));
|
||||
void* swap_out = dout.pointer;
|
||||
void* swap_params[] = {&swap_out};
|
||||
hip.check(hip.hipModuleLaunchKernel(swap_kernel, 1, 1, 1, 32, 1, 1, 0,
|
||||
nullptr, swap_params, nullptr), "hipModuleLaunchKernel(WMMA swap)");
|
||||
hip.check(hip.hipDeviceSynchronize(), "hipDeviceSynchronize(WMMA swap)");
|
||||
hip.check(hip.hipMemcpy(output.data(), dout.pointer, sizeof(output), hipMemcpyDeviceToHost), "hipMemcpy(WMMA swap)");
|
||||
for (unsigned lane = 0; lane < 32u; ++lane) {
|
||||
if (output[lane] != static_cast<float>(lane ^ 16u)) {
|
||||
std::fprintf(stderr, "FAIL WMMA swap lane=%u expected=%u actual=%g\n",
|
||||
lane, lane ^ 16u, output[lane]);
|
||||
return 5;
|
||||
}
|
||||
if (output[32u + lane] != static_cast<float>(lane)) {
|
||||
std::fprintf(stderr, "FAIL WMMA lane-id lane=%u mbcnt=%g\n", lane, output[32u + lane]);
|
||||
return 5;
|
||||
}
|
||||
}
|
||||
float max_absolute_error = 0.0f;
|
||||
float max_relative_error = 0.0f;
|
||||
for (unsigned iteration = 0; iteration < args.iterations; ++iteration) {
|
||||
const unsigned steps = 1u + (iteration & 1u);
|
||||
for (unsigned row = 0; row < 16u; ++row) {
|
||||
for (unsigned k = 0; k < 16u; ++k) {
|
||||
a[row * 16u + k] = half_integer(static_cast<int>((row * 3u + k * 5u + iteration) % 7u) - 3);
|
||||
b[row * 16u + k] = half_integer(static_cast<int>((row * 5u + k * 2u + iteration * 3u) % 7u) - 3);
|
||||
initial[row * 16u + k] = static_cast<float>(static_cast<int>((row * 11u + k * 3u + iteration) % 17u) - 8) * 0.125f;
|
||||
}
|
||||
}
|
||||
hip.check(hip.hipMemcpy(da.pointer, a.data(), sizeof(a), hipMemcpyHostToDevice), "hipMemcpy(A)");
|
||||
hip.check(hip.hipMemcpy(db.pointer, b.data(), sizeof(b), hipMemcpyHostToDevice), "hipMemcpy(B)");
|
||||
hip.check(hip.hipMemcpy(dc.pointer, initial.data(), sizeof(initial), hipMemcpyHostToDevice), "hipMemcpy(C)");
|
||||
void* a_arg = da.pointer;
|
||||
void* b_arg = db.pointer;
|
||||
void* c_arg = dc.pointer;
|
||||
void* output_arg = dout.pointer;
|
||||
unsigned steps_arg = steps;
|
||||
void* params[] = {&a_arg, &b_arg, &c_arg, &output_arg, &steps_arg};
|
||||
for (unsigned mode = 0; mode < 2u; ++mode) {
|
||||
output.fill(std::numeric_limits<float>::quiet_NaN());
|
||||
hip.check(hip.hipMemcpy(dout.pointer, output.data(), sizeof(output), hipMemcpyHostToDevice), "hipMemcpy(output sentinel)");
|
||||
hip.check(hip.hipModuleLaunchKernel(mode ? adapter_kernel : raw_kernel, 1, 1, 1, 32, 1, 1, 0, nullptr,
|
||||
params, nullptr), mode ? "hipModuleLaunchKernel(WMMA adapter)" : "hipModuleLaunchKernel(WMMA raw)");
|
||||
hip.check(hip.hipDeviceSynchronize(), "hipDeviceSynchronize(WMMA)");
|
||||
hip.check(hip.hipMemcpy(output.data(), dout.pointer, sizeof(output), hipMemcpyDeviceToHost), "hipMemcpy(output)");
|
||||
for (unsigned row = 0; row < 16u; ++row) {
|
||||
for (unsigned col = 0; col < 16u; ++col) {
|
||||
int dot = 0;
|
||||
for (unsigned k = 0; k < 16u; ++k) {
|
||||
const int av = static_cast<int>((row * 3u + k * 5u + iteration) % 7u) - 3;
|
||||
const int bv = static_cast<int>((col * 5u + k * 2u + iteration * 3u) % 7u) - 3;
|
||||
dot += av * bv;
|
||||
}
|
||||
const float expected = initial[row * 16u + col] + static_cast<float>(static_cast<int>(steps) * dot);
|
||||
const float got = output[row * 16u + col];
|
||||
const float absolute = std::fabs(got - expected);
|
||||
const float relative = absolute / std::max(1.0f, std::fabs(expected));
|
||||
if (!std::isfinite(got) || absolute > 0.0001f) {
|
||||
std::fprintf(stderr, "FAIL WMMA mode=%s iteration=%u row=%u col=%u steps=%u expected=%.9g actual=%.9g abs=%.9g\n",
|
||||
mode ? "adapter" : "raw", iteration, row, col, steps, expected, got, absolute);
|
||||
return 5;
|
||||
}
|
||||
max_absolute_error = std::max(max_absolute_error, absolute);
|
||||
max_relative_error = std::max(max_relative_error, relative);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
std::printf("PASS GFX12_WMMA architecture=gfx1201 modes=raw,legacy_adapter iterations=%u max_abs=%.9g max_rel=%.9g\n",
|
||||
args.iterations, max_absolute_error, max_relative_error);
|
||||
return 0;
|
||||
} catch (const std::exception& error) {
|
||||
std::fprintf(stderr, "FAIL GFX12_WMMA %s\n", error.what());
|
||||
loaded_modules();
|
||||
return 4;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
#include "hip_api.h"
|
||||
|
||||
int main(int argc, char** argv)
|
||||
{
|
||||
using namespace d4r::diag;
|
||||
start();
|
||||
try {
|
||||
Args args(argc, argv);
|
||||
if (args.module.empty()) throw std::runtime_error("--module is required");
|
||||
HipApi hip(args.hip_root);
|
||||
hipDeviceProp_t properties{};
|
||||
hip.select_gfx1201(args.device, properties);
|
||||
hipModule_t module = nullptr;
|
||||
hip.check(hip.hipModuleLoad(&module, args.module.c_str()), "hipModuleLoad(K enc1)");
|
||||
struct Cleanup {
|
||||
HipApi& hip; hipModule_t module;
|
||||
~Cleanup() { if (module) (void)hip.hipModuleUnload(module); }
|
||||
} cleanup{hip, module};
|
||||
for (const char* name : {"dltss_pwin_enc1_layer_prep", "dltss_pwin_enc1_layer"}) {
|
||||
hipFunction_t function = nullptr;
|
||||
hip.check(hip.hipModuleGetFunction(&function, module, name), name);
|
||||
if (!function) throw std::runtime_error(std::string("Null K function: ") + name);
|
||||
}
|
||||
std::printf("PASS K_MODULE architecture=gfx1201 functions=prep,transformer transformer_executed=0\n");
|
||||
return 0;
|
||||
} catch (const std::exception& error) {
|
||||
std::fprintf(stderr, "FAIL K_MODULE %s\n", error.what());
|
||||
loaded_modules();
|
||||
return 4;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,67 @@
|
||||
// Independent gfx12 WMMA layout probe. This deliberately does not use d4r fragments.
|
||||
// A and B are row-major [logical row/column][K]; the intrinsic receives B,A
|
||||
// so its output is row=A, column=B, as in AMD's RDNA 4 WMMA guide.
|
||||
using h8 = _Float16 __attribute__((ext_vector_type(8)));
|
||||
using f8 = float __attribute__((ext_vector_type(8)));
|
||||
#include "../../kernels/common/wmma_backend.h"
|
||||
|
||||
extern "C" __attribute__((global)) void d4r_wmma_swap_test(float* output)
|
||||
{
|
||||
const unsigned int lane = __builtin_amdgcn_workitem_id_x();
|
||||
if (lane < 32u) {
|
||||
output[lane] = d4r_wmma_other_half(static_cast<float>(lane));
|
||||
output[32u + lane] = static_cast<float>(__builtin_amdgcn_mbcnt_lo(~0u, 0u));
|
||||
}
|
||||
}
|
||||
|
||||
extern "C" __attribute__((global)) void d4r_wmma_gfx1201(
|
||||
const unsigned short* a_bits, const unsigned short* b_bits,
|
||||
const float* initial_c, float* output, unsigned int steps)
|
||||
{
|
||||
const unsigned int lane = __builtin_amdgcn_workitem_id_x();
|
||||
if (lane >= 32u) return;
|
||||
const unsigned int row = lane & 15u;
|
||||
const unsigned int group = lane >> 4;
|
||||
const _Float16* a = reinterpret_cast<const _Float16*>(a_bits);
|
||||
const _Float16* b = reinterpret_cast<const _Float16*>(b_bits);
|
||||
h8 a_frag, b_frag;
|
||||
f8 c_frag;
|
||||
for (unsigned int i = 0; i < 8u; ++i) {
|
||||
a_frag[i] = a[row * 16u + group * 8u + i];
|
||||
b_frag[i] = b[row * 16u + group * 8u + i];
|
||||
c_frag[i] = initial_c[row * 16u + group * 8u + i];
|
||||
}
|
||||
f8 result = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12(
|
||||
b_frag, a_frag, c_frag);
|
||||
if (steps == 2u)
|
||||
result = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12(
|
||||
b_frag, a_frag, result);
|
||||
for (unsigned int i = 0; i < 8u; ++i)
|
||||
output[row * 16u + group * 8u + i] = result[i];
|
||||
}
|
||||
|
||||
// d4r's original gfx11 fragment contract, exercised independently of K/M.
|
||||
extern "C" __attribute__((global)) void d4r_wmma_legacy_contract(
|
||||
const unsigned short* a_bits, const unsigned short* b_bits,
|
||||
const float* initial_c, float* output, unsigned int steps)
|
||||
{
|
||||
const unsigned int lane = __builtin_amdgcn_workitem_id_x();
|
||||
if (lane >= 32u) return;
|
||||
const unsigned int token = lane & 15u;
|
||||
const unsigned int group = lane >> 4;
|
||||
const _Float16* a = reinterpret_cast<const _Float16*>(a_bits);
|
||||
const _Float16* b = reinterpret_cast<const _Float16*>(b_bits);
|
||||
d4r_wmma_h16 a_frag, b_frag;
|
||||
d4r_wmma_f8 c_frag;
|
||||
for (unsigned int i = 0; i < 16u; ++i) {
|
||||
a_frag[i] = a[token * 16u + i];
|
||||
b_frag[i] = b[token * 16u + i];
|
||||
}
|
||||
for (unsigned int i = 0; i < 8u; ++i)
|
||||
c_frag[i] = initial_c[(2u * i + group) * 16u + token];
|
||||
d4r_wmma_f8 result = d4r_wmma_legacy_layout(a_frag, b_frag, c_frag);
|
||||
if (steps == 2u)
|
||||
result = d4r_wmma_legacy_layout(a_frag, b_frag, result);
|
||||
for (unsigned int i = 0; i < 8u; ++i)
|
||||
output[(2u * i + group) * 16u + token] = result[i];
|
||||
}
|
||||
Reference in new issue
Block a user