西风逍遥游
cff9bb4d0c
Refactor the tensor module around AITemplate-style static compilation: - Optimizer: DCE, constant folding, elementwise chain fusion, matmul/conv bias+activation epilogue fusion, liveness-based static memory planning, topological group execution order. - KernelGen: shape-baked GLSL per fused group; two-pass softmax/layernorm; shared-memory SDPA; naive + 16x16 tiled matmul templates. - GpuBackend: arena buffer allocation per memory plan, matmul autotuning, single command-buffer submission via the new gpgpu::Sequence (uploads, dispatches and downloads recorded once per run, Kompute-style). - shaderc in-process GLSL->SPIR-V on Windows (no glslc.exe dependency), with glslc spawn fallback. - CPU reference kernels shared by eager and the compiled interpreter; no-Vulkan fallback keeps compiled graphs runnable headless. Neural-network coverage: softmax/logsoftmax, layernorm/rmsnorm (with weight/bias), gelu/silu, conv1d/conv2d, pool, embedding, concat/slice, axis reduces, argmax, cast, fused SDPA (+masked), resize2d, batched matmul, seeded random. Weight-only quantization (activations stay fp32): - fp16 (IEEE half), fp8 (e4m3), fp4 (e2m1), int8, int4 with per-group block scales (MXFP-style: scale = maxAbs / format max magnitude). - Quantized weights capture into graph Const nodes; CPU interpreter dequantizes at materialization; GPU matmul/embedding kernels dequantize inline from packed uint buffers (naive matmul variant). - TF::quantizeWeight API + script binding. Mini-LLM validation (TinyStories-gpt2-3M, 8 layers d=64): - test/tensor_llm.cpp builds the full GPT-2 graph from engine ops and matches numpy reference: probs within 3e-7, greedy generation 24/24 identical (~100 ms/token CPU), compiled path exact. - tensor.llm.quantizedDtypes: fp16/int8/fp8/int4 all keep 16/16 greedy top-1 (fp4 16/16 with proper block scaling), weights 27.6 MB -> 3.3-13 MB. - scripts/export_tiny_gpt2.py converts a PyTorch .bin checkpoint (no torch dependency) plus GPT-2 tokenizer into engine assets + numpy reference. Tests: gpgpu 8/8, tensor 31/31 (original suite) + tensor.llm.* + tensor.quant.* Assets in test/assets/mini_llm are gitignored (regenerate via the script).
64 lines
2.4 KiB
C++
64 lines
2.4 KiB
C++
#ifndef EVE_TENSOR_KERNELGEN_H
|
|
#define EVE_TENSOR_KERNELGEN_H
|
|
|
|
#include "tensor/Optimizer.h"
|
|
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
namespace eve::tensor {
|
|
|
|
/**
|
|
* A generated, fully specialized compute kernel (AITemplate-style codegen).
|
|
*
|
|
* Shapes are baked into the GLSL as constants; the only runtime inputs are the
|
|
* storage buffers themselves, so a CompiledFunction is a static executable
|
|
* program with no per-dispatch shape computation.
|
|
*
|
|
* Binding layout:
|
|
* - single-pass kernels: inputs 0..k-1, output binding k;
|
|
* - two-pass kernels (softmax / layernorm / rmsnorm): pass1 binds the first
|
|
* `inputsReadPass1` inputs + `statsCount` stats buffers; pass2 binds all
|
|
* inputs + stats + the output as the last binding.
|
|
*/
|
|
struct KernelSpec {
|
|
std::string pass1; // empty when single-pass
|
|
std::string pass2; // always set
|
|
int groupsX1 = 0, groupsY1 = 1, groupsZ1 = 1;
|
|
int groupsX2 = 0, groupsY2 = 1, groupsZ2 = 1;
|
|
int inputCount = 0; // total group input buffers
|
|
int inputsReadPass1 = 0; // leading inputs bound in pass1
|
|
int statsCount = 0; // per-row stats buffers (allocated by the runtime)
|
|
int statsSize = 0; // elements per stats buffer (rows)
|
|
bool twoPass = false;
|
|
|
|
// Weight-only quantization of a matmul B / embedding table input.
|
|
// qDtype == 0 means fp32 storage. int8/int4 add a per-group scales buffer.
|
|
int qDtype = 0; // static_cast<int>(DType)
|
|
int qGroup = 0; // elements per scale group
|
|
int scalesBinding = -1; // storage binding of the scales buffer
|
|
int outputBinding = -1; // storage binding of the output (override)
|
|
};
|
|
|
|
/**
|
|
* Generate the specialized GLSL kernel(s) for a fused group.
|
|
* Returns false when the group cannot be lowered (e.g. too many inputs for the
|
|
* fixed 8-binding descriptor layout) — callers fall back to the CPU interpreter.
|
|
*/
|
|
bool generateKernel(const Graph &graph, const FusedGroup &group, KernelSpec &out);
|
|
|
|
/**
|
|
* Generate a specific matmul variant for autotuning (tiled=false: thread-per-
|
|
* output naive; tiled=true: shared-memory 16x16 tiles). Only rank-2 matmuls
|
|
* support the tiled variant.
|
|
*/
|
|
bool generateMatMulVariant(const Graph &graph, const FusedGroup &group, bool tiled,
|
|
KernelSpec &out);
|
|
|
|
/** Max storage bindings available per generated kernel. */
|
|
constexpr int kMaxKernelBindings = 8;
|
|
|
|
} // namespace eve::tensor
|
|
|
|
#endif // EVE_TENSOR_KERNELGEN_H
|