gemv_quantv1.cpp

#define CL_TARGET_OPENCL_VERSION 120

#include <CL/cl.h>

#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <limits>
#include <random>
#include <string>
#include <vector>

#include "half.hpp"

using half_float::half;
using std::cerr;
using std::cout;
using std::endl;
using std::string;
using std::vector;

namespace {

constexpr int kBlockSize = 32;
constexpr float kQMax = 127.0f;
constexpr float kScaleEpsilon = 1.0e-6f;

struct Options {
    int m = 122753;
    int k = 2304;
    int warmup = 100;
    int iterations = 1000;
    int platform = 0;
    int device = 0;
};

struct ErrorStats {
    float max_abs = 0.0f;
    float max_rel = 0.0f;
};

void checkErr(cl_int err, const char *operation) {
    if (err != CL_SUCCESS) {
        cerr << "ERROR: " << operation << " (OpenCL error " << err << ")\n";
        std::exit(EXIT_FAILURE);
    }
}

int parsePositive(const char *value, const char *name) {
    char *end = nullptr;
    const long parsed = std::strtol(value, &end, 10);
    if (*value == '\0' || *end != '\0' || parsed <= 0 ||
        parsed > std::numeric_limits<int>::max()) {
        cerr << "Invalid value for " << name << ": " << value << endl;
        std::exit(EXIT_FAILURE);
    }
    return static_cast<int>(parsed);
}

Options parseOptions(int argc, char **argv) {
    Options options;
    for (int i = 1; i < argc; ++i) {
        const string arg = argv[i];
        if (arg == "--help") {
            cout << "Usage: gemv_quantv1 [--m rows] [--k columns] [--warmup count] "
                    "[--iters count] [--platform index] [--device index]\n";
            std::exit(EXIT_SUCCESS);
        }
        if (i + 1 >= argc) {
            cerr << "Missing value after " << arg << endl;
            std::exit(EXIT_FAILURE);
        }
        if (arg == "--m") options.m = parsePositive(argv[++i], "--m");
        else if (arg == "--k") options.k = parsePositive(argv[++i], "--k");
        else if (arg == "--warmup") options.warmup = parsePositive(argv[++i], "--warmup");
        else if (arg == "--iters") options.iterations = parsePositive(argv[++i], "--iters");
        else if (arg == "--platform") options.platform = parsePositive(argv[++i], "--platform") - 1;
        else if (arg == "--device") options.device = parsePositive(argv[++i], "--device") - 1;
        else {
            cerr << "Unknown option: " << arg << endl;
            std::exit(EXIT_FAILURE);
        }
    }
    if (options.k % kBlockSize != 0) {
        cerr << "--k must be a multiple of " << kBlockSize << endl;
        std::exit(EXIT_FAILURE);
    }
    return options;
}

vector<half> generateF16(size_t count, uint32_t seed) {
    std::mt19937 generator(seed);
    std::uniform_real_distribution<float> distribution(-0.5f, 0.5f);
    vector<half> values(count);
    for (half &value : values) value = half(distribution(generator));
    return values;
}

int8_t clampInt8(int value) {
    return static_cast<int8_t>(std::max(-127, std::min(127, value)));
}

// Split layout: q[row, k] and scale[row, block]. It avoids the 34-byte
// interleaved structure and lets a 32-lane group read one full block together.
void quantizeQ8(const vector<half> &source, vector<int8_t> &quantized,
                vector<float> &scales, int m, int k) {
    const int blocks = k / kBlockSize;
    quantized.resize(static_cast<size_t>(m) * k);
    scales.resize(static_cast<size_t>(m) * blocks);
    for (int row = 0; row < m; ++row) {
        for (int block = 0; block < blocks; ++block) {
            const size_t offset = static_cast<size_t>(row) * k + block * kBlockSize;
            float maxAbs = 0.0f;
            for (int lane = 0; lane < kBlockSize; ++lane)
                maxAbs = std::max(maxAbs, std::fabs(static_cast<float>(source[offset + lane])));

            const float scale = std::max(maxAbs / kQMax, kScaleEpsilon);
            scales[static_cast<size_t>(row) * blocks + block] = scale;
            for (int lane = 0; lane < kBlockSize; ++lane) {
                const float value = static_cast<float>(source[offset + lane]) / scale;
                quantized[offset + lane] = clampInt8(static_cast<int>(std::nearbyint(value)));
            }
        }
    }
}

void computeReferences(const vector<half> &original, const vector<int8_t> &quantized,
                       const vector<float> &scales, const vector<half> &x,
                       vector<float> &originalRef, vector<float> &quantizedRef,
                       int m, int k, float alpha) {
    const int blocks = k / kBlockSize;
    originalRef.resize(m);
    quantizedRef.resize(m);
    for (int row = 0; row < m; ++row) {
        float originalSum = 0.0f;
        float quantizedSum = 0.0f;
        for (int block = 0; block < blocks; ++block) {
            float blockSum = 0.0f;
            const float scale = scales[static_cast<size_t>(row) * blocks + block];
            const size_t offset = static_cast<size_t>(row) * k + block * kBlockSize;
            for (int lane = 0; lane < kBlockSize; ++lane) {
                const float xv = static_cast<float>(x[block * kBlockSize + lane]);
                originalSum += static_cast<float>(original[offset + lane]) * xv;
                blockSum += static_cast<float>(quantized[offset + lane]) * xv;
            }
            quantizedSum += blockSum * scale;
        }
        originalRef[row] = alpha * originalSum;
        quantizedRef[row] = alpha * quantizedSum;
    }
}

ErrorStats compare(const vector<float> &expected, const vector<float> &actual) {
    ErrorStats stats;
    for (size_t i = 0; i < expected.size(); ++i) {
        const float absError = std::fabs(expected[i] - actual[i]);
        const float denominator = std::max(std::fabs(expected[i]), 1.0e-6f);
        stats.max_abs = std::max(stats.max_abs, absError);
        stats.max_rel = std::max(stats.max_rel, absError / denominator);
    }
    return stats;
}

const char *kernelSource = R"CLC(
// NVIDIA's OpenCL C implementation does not expose cl_khr_fp16. Activations
// therefore remain binary16 in memory and are converted explicitly to float.
inline float load_f16(__global const ushort *data, int index) {
    const uint h = data[index];
    const uint sign = (h & 0x8000u) << 16;
    uint exponent = (h >> 10) & 0x1fu;
    uint mantissa = h & 0x3ffu;
    uint bits;
    if (exponent == 0u) {
        if (mantissa == 0u) return as_float(sign);
        int shift = 0;
        while ((mantissa & 0x400u) == 0u) { mantissa <<= 1; ++shift; }
        mantissa &= 0x3ffu;
        bits = sign | ((uint)(113 - shift) << 23) | (mantissa << 13);
    } else if (exponent == 31u) {
        bits = sign | 0x7f800000u | (mantissa << 13);
    } else {
        bits = sign | ((exponent + 112u) << 23) | (mantissa << 13);
    }
    return as_float(bits);
}

__kernel void gemv_q8_scalar(__global const char *q,
                             __global const float *scales,
                             __global const ushort *x,
                             __global float *y,
                             int M, int K, float alpha) {
    const int row = (int)get_global_id(0);
    if (row >= M) return;
    const int blocks = K / 32;
    float sum = 0.0f;
    for (int block = 0; block < blocks; ++block) {
        float block_sum = 0.0f;
        const int offset = row * K + block * 32;
        for (int lane = 0; lane < 32; ++lane)
            block_sum = fma((float)q[offset + lane], load_f16(x, block * 32 + lane), block_sum);
        sum = fma(block_sum, scales[row * blocks + block], sum);
    }
    y[row] = alpha * sum;
}

__kernel void gemv_q8_group32(__global const char *q,
                              __global const float *scales,
                              __global const float *x,
                              __global float *y,
                              int M, int K, float alpha,
                              __local float *scratch) {
    const int row = (int)get_group_id(0);
    const int lane = (int)get_local_id(0);
    if (row >= M) return;

    const int blocks = K / 32;
    float partial = 0.0f;
    for (int block = 0; block < blocks; ++block) {
        const int index = row * K + block * 32 + lane;
        const float product = (float)q[index] * x[block * 32 + lane];
        partial = fma(product, scales[row * blocks + block], partial);
    }
    scratch[lane] = partial;
    barrier(CLK_LOCAL_MEM_FENCE);

    // A portable tree reduction. On NVIDIA this maps naturally to one warp;
    // barriers keep it legal on OpenCL 1.2 devices without subgroup support.
    for (int offset = 16; offset > 0; offset >>= 1) {
        if (lane < offset) scratch[lane] += scratch[lane + offset];
        barrier(CLK_LOCAL_MEM_FENCE);
    }
    if (lane == 0) y[row] = alpha * scratch[0];
}
)CLC";

string getDeviceString(cl_device_id device, cl_device_info key) {
    size_t size = 0;
    checkErr(clGetDeviceInfo(device, key, 0, nullptr, &size), "clGetDeviceInfo(size)");
    string result(size, '\0');
    checkErr(clGetDeviceInfo(device, key, size, &result[0], nullptr), "clGetDeviceInfo(value)");
    if (!result.empty() && result.back() == '\0') result.pop_back();
    return result;
}

cl_device_id chooseDevice(int platformIndex, int deviceIndex, cl_platform_id &chosenPlatform) {
    cl_uint platformCount = 0;
    checkErr(clGetPlatformIDs(0, nullptr, &platformCount), "clGetPlatformIDs(count)");
    vector<cl_platform_id> platforms(platformCount);
    checkErr(clGetPlatformIDs(platformCount, platforms.data(), nullptr), "clGetPlatformIDs");
    if (platformIndex < 0 || platformIndex >= static_cast<int>(platformCount)) {
        cerr << "Platform index is out of range (available: " << platformCount << ")\n";
        std::exit(EXIT_FAILURE);
    }
    chosenPlatform = platforms[platformIndex];
    cl_uint deviceCount = 0;
    checkErr(clGetDeviceIDs(chosenPlatform, CL_DEVICE_TYPE_GPU, 0, nullptr, &deviceCount),
             "clGetDeviceIDs(count)");
    vector<cl_device_id> devices(deviceCount);
    checkErr(clGetDeviceIDs(chosenPlatform, CL_DEVICE_TYPE_GPU, deviceCount, devices.data(), nullptr),
             "clGetDeviceIDs");
    if (deviceIndex < 0 || deviceIndex >= static_cast<int>(deviceCount)) {
        cerr << "Device index is out of range (available GPUs: " << deviceCount << ")\n";
        std::exit(EXIT_FAILURE);
    }
    return devices[deviceIndex];
}

cl_program buildProgram(cl_context context, cl_device_id device) {
    cl_int err = CL_SUCCESS;
    cl_program program = clCreateProgramWithSource(context, 1, &kernelSource, nullptr, &err);
    checkErr(err, "clCreateProgramWithSource");
    err = clBuildProgram(program, 1, &device, "-cl-std=CL1.2", nullptr, nullptr);
    if (err != CL_SUCCESS) {
        size_t size = 0;
        clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, 0, nullptr, &size);
        string log(size, '\0');
        clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, size, &log[0], nullptr);
        cerr << "OpenCL build failed:\n" << log << endl;
        clReleaseProgram(program);
        std::exit(EXIT_FAILURE);
    }
    return program;
}

double benchmark(cl_command_queue queue, cl_kernel kernel, size_t globalSize,
                 size_t localSize, int warmup, int iterations) {
    for (int i = 0; i < warmup; ++i)
        checkErr(clEnqueueNDRangeKernel(queue, kernel, 1, nullptr, &globalSize, &localSize,
                                       0, nullptr, nullptr), "warmup enqueue");
    checkErr(clFinish(queue), "warmup finish");

    vector<cl_event> events(iterations);
    for (int i = 0; i < iterations; ++i)
        checkErr(clEnqueueNDRangeKernel(queue, kernel, 1, nullptr, &globalSize, &localSize,
                                       0, nullptr, &events[i]), "timed enqueue");
    checkErr(clWaitForEvents(iterations, events.data()), "clWaitForEvents");

    double totalNanoseconds = 0.0;
    for (cl_event event : events) {
        cl_ulong start = 0, end = 0;
        checkErr(clGetEventProfilingInfo(event, CL_PROFILING_COMMAND_START, sizeof(start), &start,
                                        nullptr), "profiling start");
        checkErr(clGetEventProfilingInfo(event, CL_PROFILING_COMMAND_END, sizeof(end), &end,
                                        nullptr), "profiling end");
        totalNanoseconds += static_cast<double>(end - start);
        clReleaseEvent(event);
    }
    return totalNanoseconds / iterations / 1.0e6;
}

void setCommonArgs(cl_kernel kernel, cl_mem q, cl_mem scales, cl_mem x, cl_mem y,
                   int m, int k, float alpha, bool localScratch) {
    int index = 0;
    checkErr(clSetKernelArg(kernel, index++, sizeof(q), &q), "clSetKernelArg(q)");
    checkErr(clSetKernelArg(kernel, index++, sizeof(scales), &scales), "clSetKernelArg(scales)");
    checkErr(clSetKernelArg(kernel, index++, sizeof(x), &x), "clSetKernelArg(x)");
    checkErr(clSetKernelArg(kernel, index++, sizeof(y), &y), "clSetKernelArg(y)");
    checkErr(clSetKernelArg(kernel, index++, sizeof(m), &m), "clSetKernelArg(M)");
    checkErr(clSetKernelArg(kernel, index++, sizeof(k), &k), "clSetKernelArg(K)");
    checkErr(clSetKernelArg(kernel, index++, sizeof(alpha), &alpha), "clSetKernelArg(alpha)");
    if (localScratch)
        checkErr(clSetKernelArg(kernel, index++, kBlockSize * sizeof(float), nullptr),
                 "clSetKernelArg(scratch)");
}

}  // namespace

int main(int argc, char **argv) {
    static_assert(sizeof(half) == sizeof(uint16_t), "half must use IEEE binary16 storage");
    const Options options = parseOptions(argc, argv);
    const float alpha = 0.5f;

    cout << "Shape: M=" << options.m << ", N=1, K=" << options.k << endl;
    vector<half> weights = generateF16(static_cast<size_t>(options.m) * options.k, 66);
    vector<half> activation = generateF16(options.k, 67);
    vector<float> activationF32(options.k);
    std::transform(activation.begin(), activation.end(), activationF32.begin(),
                   [](half value) { return static_cast<float>(value); });
    vector<int8_t> quantized;
    vector<float> scales;
    quantizeQ8(weights, quantized, scales, options.m, options.k);

    vector<float> originalReference, quantizedReference;
    computeReferences(weights, quantized, scales, activation, originalReference,
                      quantizedReference, options.m, options.k, alpha);
    const ErrorStats quantizationError = compare(originalReference, quantizedReference);
    weights.clear();
    weights.shrink_to_fit();

    cl_platform_id platform = nullptr;
    cl_device_id device = chooseDevice(options.platform, options.device, platform);
    cout << "Device: " << getDeviceString(device, CL_DEVICE_NAME) << '\n'
         << "Driver: " << getDeviceString(device, CL_DRIVER_VERSION) << '\n'
         << "OpenCL C: " << getDeviceString(device, CL_DEVICE_OPENCL_C_VERSION) << endl;

    cl_int err = CL_SUCCESS;
    cl_context context = clCreateContext(nullptr, 1, &device, nullptr, nullptr, &err);
    checkErr(err, "clCreateContext");
    cl_command_queue queue = clCreateCommandQueue(context, device, CL_QUEUE_PROFILING_ENABLE, &err);
    checkErr(err, "clCreateCommandQueue");
    cl_program program = buildProgram(context, device);

    cl_mem qBuffer = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR,
                                    quantized.size() * sizeof(int8_t), quantized.data(), &err);
    checkErr(err, "clCreateBuffer(q)");
    cl_mem scaleBuffer = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR,
                                        scales.size() * sizeof(float), scales.data(), &err);
    checkErr(err, "clCreateBuffer(scales)");
    cl_mem xBuffer = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR,
                                    activation.size() * sizeof(half), activation.data(), &err);
    checkErr(err, "clCreateBuffer(x)");
    cl_mem xF32Buffer = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR,
                                       activationF32.size() * sizeof(float), activationF32.data(), &err);
    checkErr(err, "clCreateBuffer(xF32)");
    cl_mem yBuffer = clCreateBuffer(context, CL_MEM_WRITE_ONLY,
                                    static_cast<size_t>(options.m) * sizeof(float), nullptr, &err);
    checkErr(err, "clCreateBuffer(y)");

    cl_kernel scalar = clCreateKernel(program, "gemv_q8_scalar", &err);
    checkErr(err, "clCreateKernel(scalar)");
    cl_kernel group32 = clCreateKernel(program, "gemv_q8_group32", &err);
    checkErr(err, "clCreateKernel(group32)");
    setCommonArgs(scalar, qBuffer, scaleBuffer, xBuffer, yBuffer,
                  options.m, options.k, alpha, false);
    setCommonArgs(group32, qBuffer, scaleBuffer, xF32Buffer, yBuffer,
                  options.m, options.k, alpha, true);

    constexpr size_t scalarLocalSize = 128;
    const size_t scalarGlobalSize =
        (static_cast<size_t>(options.m) + scalarLocalSize - 1) / scalarLocalSize * scalarLocalSize;
    const double scalarMs = benchmark(queue, scalar, scalarGlobalSize, scalarLocalSize,
                                      options.warmup, options.iterations);
    vector<float> scalarOutput(options.m);
    checkErr(clEnqueueReadBuffer(queue, yBuffer, CL_TRUE, 0,
                                 scalarOutput.size() * sizeof(float), scalarOutput.data(),
                                 0, nullptr, nullptr), "read scalar output");

    constexpr size_t groupLocalSize = kBlockSize;
    const size_t groupGlobalSize = static_cast<size_t>(options.m) * groupLocalSize;
    const double groupMs = benchmark(queue, group32, groupGlobalSize, groupLocalSize,
                                     options.warmup, options.iterations);
    vector<float> groupOutput(options.m);
    checkErr(clEnqueueReadBuffer(queue, yBuffer, CL_TRUE, 0,
                                 groupOutput.size() * sizeof(float), groupOutput.data(),
                                 0, nullptr, nullptr), "read optimized output");

    const ErrorStats scalarError = compare(quantizedReference, scalarOutput);
    const ErrorStats groupError = compare(quantizedReference, groupOutput);
    cout << std::fixed << std::setprecision(6)
         << "Quantization vs FP16 weights: max_abs=" << quantizationError.max_abs
         << ", max_rel=" << quantizationError.max_rel << '\n'
         << "gemv_q8_scalar:  avg_ms=" << scalarMs
         << ", max_abs=" << scalarError.max_abs << ", max_rel=" << scalarError.max_rel << '\n'
         << "gemv_q8_group32: avg_ms=" << groupMs
         << ", max_abs=" << groupError.max_abs << ", max_rel=" << groupError.max_rel << '\n'
         << "Speedup (scalar/group32): " << scalarMs / groupMs << "x\n";

    const bool correct = scalarError.max_abs <= 1.0e-3f && groupError.max_abs <= 1.0e-3f;
    cout << "Correctness: " << (correct ? "PASS" : "FAIL") << endl;

    clReleaseKernel(group32);
    clReleaseKernel(scalar);
    clReleaseMemObject(yBuffer);
    clReleaseMemObject(xF32Buffer);
    clReleaseMemObject(xBuffer);
    clReleaseMemObject(scaleBuffer);
    clReleaseMemObject(qBuffer);
    clReleaseProgram(program);
    clReleaseCommandQueue(queue);
    clReleaseContext(context);
    return correct ? EXIT_SUCCESS : EXIT_FAILURE;
}