西风逍遥游
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).
1060 lines
41 KiB
C++
1060 lines
41 KiB
C++
#include "tensor/TF.h"
|
|
#include "tensor/CpuKernels.h"
|
|
#include "tensor/GpuBackend.h"
|
|
#include "tensor/Graph.h"
|
|
#include "tensor/Quant.h"
|
|
#include "tensor/Tensor.h"
|
|
|
|
#include "common/Exception.h"
|
|
|
|
#include <simplesquirrel/simplesquirrel.hpp>
|
|
|
|
#include <cmath>
|
|
#include <cstdint>
|
|
#include <cstring>
|
|
|
|
namespace eve::tensor {
|
|
namespace {
|
|
// Below this element count the GPU dispatch/readback round-trip costs more
|
|
// than the CPU loop saves; only worth trying the compute-shader path above it.
|
|
constexpr int kGpuReduceMinSize = 1 << 14; // 16384 floats
|
|
|
|
int convOutSize(int inSize, int kernel, int stride, int pad) {
|
|
const int out = (inSize + 2 * pad - kernel) / stride + 1;
|
|
if (out <= 0) throw Exception("TF.conv: output size must be > 0");
|
|
return out;
|
|
}
|
|
|
|
int normalizeAxis(int axis, int rank) {
|
|
if (axis < 0) axis += rank;
|
|
if (axis < 0 || axis >= rank) throw Exception("TF: axis out of range");
|
|
return axis;
|
|
}
|
|
|
|
void reduceOutDims(const int *dims, int rank, int axis, int keepDims, int *out, int &outRank) {
|
|
outRank = rank;
|
|
for (int k = 0; k < rank; ++k) out[k] = dims[k];
|
|
if (keepDims) {
|
|
out[axis] = 1;
|
|
} else if (rank == 1) {
|
|
out[0] = 1;
|
|
} else {
|
|
for (int k = axis; k < rank - 1; ++k) out[k] = out[k + 1];
|
|
out[rank - 1] = 0;
|
|
--outRank;
|
|
}
|
|
}
|
|
|
|
} // namespace
|
|
|
|
Module_IMPL(TF, new TF());
|
|
|
|
TF::TF() : seed_(1), rngState_(1) {}
|
|
|
|
float TF::nextUniform() const {
|
|
rngState_ = rngState_ * 1664525u + 1013904223u;
|
|
return float(rngState_ >> 8) * (1.f / 16777216.f);
|
|
}
|
|
|
|
float TF::nextGaussian() const {
|
|
float u1 = nextUniform();
|
|
float u2 = nextUniform();
|
|
if (u1 < 1e-7f) u1 = 1e-7f;
|
|
return std::sqrt(-2.f * std::log(u1)) * std::cos(6.28318530718f * u2);
|
|
}
|
|
|
|
void TF::setRandomSeed(uint32_t seed) {
|
|
seed_ = seed == 0 ? 1u : seed;
|
|
rngState_ = seed_;
|
|
}
|
|
|
|
uint32_t TF::getRandomSeed() const { return seed_; }
|
|
|
|
void TF::pushTrace(Func *f) {
|
|
if (f) traceStack_.push_back(f);
|
|
}
|
|
|
|
void TF::popTrace(Func *f) {
|
|
if (traceStack_.empty()) return;
|
|
if (traceStack_.back() == f) {
|
|
traceStack_.pop_back();
|
|
return;
|
|
}
|
|
for (auto it = traceStack_.begin(); it != traceStack_.end(); ++it) {
|
|
if (*it == f) {
|
|
traceStack_.erase(it);
|
|
return;
|
|
}
|
|
}
|
|
}
|
|
|
|
Func *TF::tracing() const {
|
|
return traceStack_.empty() ? nullptr : traceStack_.back();
|
|
}
|
|
|
|
Func *TF::func() { return new Func(this); }
|
|
|
|
Tensor *TF::filled(const int *dims, int rank, float value) {
|
|
if (Func *f = tracing()) return f->emitFill(dims, rank, value);
|
|
auto *t = new Tensor(dims, rank);
|
|
t->fill(value);
|
|
return t;
|
|
}
|
|
|
|
Tensor *TF::zeros1(int d0) {
|
|
int d[] = {d0};
|
|
return filled(d, 1, 0.f);
|
|
}
|
|
Tensor *TF::zeros2(int d0, int d1) {
|
|
int d[] = {d0, d1};
|
|
return filled(d, 2, 0.f);
|
|
}
|
|
Tensor *TF::zeros3(int d0, int d1, int d2) {
|
|
int d[] = {d0, d1, d2};
|
|
return filled(d, 3, 0.f);
|
|
}
|
|
Tensor *TF::zeros4(int d0, int d1, int d2, int d3) {
|
|
int d[] = {d0, d1, d2, d3};
|
|
return filled(d, 4, 0.f);
|
|
}
|
|
Tensor *TF::zeros5(int d0, int d1, int d2, int d3, int d4) {
|
|
int d[] = {d0, d1, d2, d3, d4};
|
|
return filled(d, 5, 0.f);
|
|
}
|
|
Tensor *TF::zeros6(int d0, int d1, int d2, int d3, int d4, int d5) {
|
|
int d[] = {d0, d1, d2, d3, d4, d5};
|
|
return filled(d, 6, 0.f);
|
|
}
|
|
|
|
Tensor *TF::ones1(int d0) {
|
|
int d[] = {d0};
|
|
return filled(d, 1, 1.f);
|
|
}
|
|
Tensor *TF::ones2(int d0, int d1) {
|
|
int d[] = {d0, d1};
|
|
return filled(d, 2, 1.f);
|
|
}
|
|
Tensor *TF::ones3(int d0, int d1, int d2) {
|
|
int d[] = {d0, d1, d2};
|
|
return filled(d, 3, 1.f);
|
|
}
|
|
Tensor *TF::ones4(int d0, int d1, int d2, int d3) {
|
|
int d[] = {d0, d1, d2, d3};
|
|
return filled(d, 4, 1.f);
|
|
}
|
|
Tensor *TF::ones5(int d0, int d1, int d2, int d3, int d4) {
|
|
int d[] = {d0, d1, d2, d3, d4};
|
|
return filled(d, 5, 1.f);
|
|
}
|
|
Tensor *TF::ones6(int d0, int d1, int d2, int d3, int d4, int d5) {
|
|
int d[] = {d0, d1, d2, d3, d4, d5};
|
|
return filled(d, 6, 1.f);
|
|
}
|
|
|
|
Tensor *TF::fill1(int d0, float value) {
|
|
int d[] = {d0};
|
|
return filled(d, 1, value);
|
|
}
|
|
Tensor *TF::fill2(int d0, int d1, float value) {
|
|
int d[] = {d0, d1};
|
|
return filled(d, 2, value);
|
|
}
|
|
Tensor *TF::fill3(int d0, int d1, int d2, float value) {
|
|
int d[] = {d0, d1, d2};
|
|
return filled(d, 3, value);
|
|
}
|
|
Tensor *TF::fill4(int d0, int d1, int d2, int d3, float value) {
|
|
int d[] = {d0, d1, d2, d3};
|
|
return filled(d, 4, value);
|
|
}
|
|
|
|
Tensor *TF::constantScalar(float value) {
|
|
int d[] = {1};
|
|
return filled(d, 1, value);
|
|
}
|
|
|
|
Tensor *TF::arange(int n) {
|
|
if (n <= 0) throw Exception("TF.arange: n must be > 0");
|
|
if (tracing()) throw Exception("TF.arange: not supported while tracing (use inputs/constants)");
|
|
auto *t = new Tensor(n);
|
|
for (int i = 0; i < n; ++i) t->set1(i, float(i));
|
|
return t;
|
|
}
|
|
|
|
Tensor *TF::linspace(float start, float end, int n) {
|
|
if (n <= 0) throw Exception("TF.linspace: n must be > 0");
|
|
if (tracing()) throw Exception("TF.linspace: not supported while tracing");
|
|
auto *t = new Tensor(n);
|
|
if (n == 1) {
|
|
t->set1(0, start);
|
|
return t;
|
|
}
|
|
float step = (end - start) / float(n - 1);
|
|
for (int i = 0; i < n; ++i) t->set1(i, start + step * float(i));
|
|
return t;
|
|
}
|
|
|
|
Tensor *TF::eye(int n) {
|
|
if (n <= 0) throw Exception("TF.eye: n must be > 0");
|
|
if (Func *f = tracing()) {
|
|
auto *eager = new Tensor(n, n);
|
|
for (int i = 0; i < n; ++i) eager->set2(i, i, 1.f);
|
|
int id = f->ensureNode(eager);
|
|
delete eager;
|
|
return Tensor::makeSymbolic(&f->graph(), id, f->graph().node(id).dims,
|
|
f->graph().node(id).rank);
|
|
}
|
|
auto *t = new Tensor(n, n);
|
|
for (int i = 0; i < n; ++i) t->set2(i, i, 1.f);
|
|
return t;
|
|
}
|
|
|
|
Tensor *TF::randomUniform1(int d0) {
|
|
if (tracing()) throw Exception("TF.randomUniform: not supported while tracing");
|
|
auto *t = zeros1(d0);
|
|
for (int i = 0; i < t->getSize(); ++i) t->set(i, nextUniform());
|
|
return t;
|
|
}
|
|
|
|
Tensor *TF::randomUniform2(int d0, int d1) {
|
|
if (tracing()) throw Exception("TF.randomUniform: not supported while tracing");
|
|
auto *t = zeros2(d0, d1);
|
|
for (int i = 0; i < t->getSize(); ++i) t->set(i, nextUniform());
|
|
return t;
|
|
}
|
|
|
|
Tensor *TF::randomUniform3(int d0, int d1, int d2) {
|
|
if (tracing()) throw Exception("TF.randomUniform: not supported while tracing");
|
|
auto *t = zeros3(d0, d1, d2);
|
|
for (int i = 0; i < t->getSize(); ++i) t->set(i, nextUniform());
|
|
return t;
|
|
}
|
|
|
|
Tensor *TF::randomUniform4(int d0, int d1, int d2, int d3) {
|
|
if (tracing()) throw Exception("TF.randomUniform: not supported while tracing");
|
|
auto *t = zeros4(d0, d1, d2, d3);
|
|
for (int i = 0; i < t->getSize(); ++i) t->set(i, nextUniform());
|
|
return t;
|
|
}
|
|
|
|
Tensor *TF::randomNormal1(int d0) {
|
|
if (tracing()) throw Exception("TF.randomNormal: not supported while tracing");
|
|
auto *t = zeros1(d0);
|
|
for (int i = 0; i < t->getSize(); ++i) t->set(i, nextGaussian());
|
|
return t;
|
|
}
|
|
|
|
Tensor *TF::randomNormal2(int d0, int d1) {
|
|
if (tracing()) throw Exception("TF.randomNormal: not supported while tracing");
|
|
auto *t = zeros2(d0, d1);
|
|
for (int i = 0; i < t->getSize(); ++i) t->set(i, nextGaussian());
|
|
return t;
|
|
}
|
|
|
|
Tensor *TF::randomNormal3(int d0, int d1, int d2) {
|
|
if (tracing()) throw Exception("TF.randomNormal: not supported while tracing");
|
|
auto *t = zeros3(d0, d1, d2);
|
|
for (int i = 0; i < t->getSize(); ++i) t->set(i, nextGaussian());
|
|
return t;
|
|
}
|
|
|
|
Tensor *TF::randomNormal4(int d0, int d1, int d2, int d3) {
|
|
if (tracing()) throw Exception("TF.randomNormal: not supported while tracing");
|
|
auto *t = zeros4(d0, d1, d2, d3);
|
|
for (int i = 0; i < t->getSize(); ++i) t->set(i, nextGaussian());
|
|
return t;
|
|
}
|
|
|
|
#define EVE_TF_UNARY(name, method, opType) \
|
|
Tensor *TF::name(Tensor *a) { \
|
|
if (!a) throw Exception("TF." #name ": null"); \
|
|
if (Func *f = tracing()) return f->emitUnary(opType, a); \
|
|
a->ensureEager(#name); \
|
|
return a->method(); \
|
|
}
|
|
|
|
EVE_TF_UNARY(neg, neg, OpType::Neg)
|
|
EVE_TF_UNARY(abs, abs, OpType::Abs)
|
|
EVE_TF_UNARY(sqrt, sqrt, OpType::Sqrt)
|
|
EVE_TF_UNARY(exp, exp, OpType::Exp)
|
|
EVE_TF_UNARY(log, log, OpType::Log)
|
|
EVE_TF_UNARY(sin, sin, OpType::Sin)
|
|
EVE_TF_UNARY(cos, cos, OpType::Cos)
|
|
EVE_TF_UNARY(tanh, tanh, OpType::Tanh)
|
|
EVE_TF_UNARY(relu, relu, OpType::Relu)
|
|
EVE_TF_UNARY(sigmoid, sigmoid, OpType::Sigmoid)
|
|
EVE_TF_UNARY(gelu, gelu, OpType::Gelu)
|
|
EVE_TF_UNARY(silu, silu, OpType::Silu)
|
|
|
|
#undef EVE_TF_UNARY
|
|
|
|
#define EVE_TF_BINARY(name, method, opType) \
|
|
Tensor *TF::name(Tensor *a, Tensor *b) { \
|
|
if (!a || !b) throw Exception("TF." #name ": null"); \
|
|
if (Func *f = tracing()) return f->emitBinary(opType, a, b); \
|
|
a->ensureEager(#name); \
|
|
b->ensureEager(#name); \
|
|
return a->method(b); \
|
|
}
|
|
|
|
EVE_TF_BINARY(add, add, OpType::Add)
|
|
EVE_TF_BINARY(sub, sub, OpType::Sub)
|
|
EVE_TF_BINARY(multiply, multiply, OpType::Multiply)
|
|
EVE_TF_BINARY(div, div, OpType::Divide)
|
|
|
|
#undef EVE_TF_BINARY
|
|
|
|
Tensor *TF::addScalar(Tensor *a, float s) {
|
|
if (!a) throw Exception("TF.addScalar: null");
|
|
if (Func *f = tracing()) return f->emitUnaryScalar(OpType::AddScalar, a, s);
|
|
return a->addScalar(s);
|
|
}
|
|
Tensor *TF::subScalar(Tensor *a, float s) {
|
|
if (!a) throw Exception("TF.subScalar: null");
|
|
if (Func *f = tracing()) return f->emitUnaryScalar(OpType::SubScalar, a, s);
|
|
return a->subScalar(s);
|
|
}
|
|
Tensor *TF::mulScalar(Tensor *a, float s) {
|
|
if (!a) throw Exception("TF.mulScalar: null");
|
|
if (Func *f = tracing()) return f->emitUnaryScalar(OpType::MulScalar, a, s);
|
|
return a->mulScalar(s);
|
|
}
|
|
Tensor *TF::divScalar(Tensor *a, float s) {
|
|
if (!a) throw Exception("TF.divScalar: null");
|
|
if (Func *f = tracing()) return f->emitUnaryScalar(OpType::DivScalar, a, s);
|
|
return a->divScalar(s);
|
|
}
|
|
Tensor *TF::powScalar(Tensor *a, float exp) {
|
|
if (!a) throw Exception("TF.powScalar: null");
|
|
if (Func *f = tracing()) return f->emitUnaryScalar(OpType::PowScalar, a, exp);
|
|
return a->powScalar(exp);
|
|
}
|
|
Tensor *TF::clamp(Tensor *a, float lo, float hi) {
|
|
if (!a) throw Exception("TF.clamp: null");
|
|
if (Func *f = tracing()) return f->emitUnaryScalar(OpType::Clamp, a, lo, hi);
|
|
return a->clamp(lo, hi);
|
|
}
|
|
Tensor *TF::maximumScalar(Tensor *a, float s) {
|
|
if (!a) throw Exception("TF.maximumScalar: null");
|
|
if (Func *f = tracing()) return f->emitUnaryScalar(OpType::MaximumScalar, a, s);
|
|
return a->maximumScalar(s);
|
|
}
|
|
Tensor *TF::minimumScalar(Tensor *a, float s) {
|
|
if (!a) throw Exception("TF.minimumScalar: null");
|
|
if (Func *f = tracing()) return f->emitUnaryScalar(OpType::MinimumScalar, a, s);
|
|
return a->minimumScalar(s);
|
|
}
|
|
|
|
Tensor *TF::matmul(Tensor *a, Tensor *b) {
|
|
if (!a || !b) throw Exception("TF.matmul: null");
|
|
if (Func *f = tracing()) return f->emitMatMul(a, b);
|
|
return a->matmul(b);
|
|
}
|
|
|
|
Tensor *TF::transpose(Tensor *a) {
|
|
if (!a) throw Exception("TF.transpose: null");
|
|
if (Func *f = tracing()) return f->emitTranspose(a);
|
|
return a->transpose();
|
|
}
|
|
|
|
Tensor *TF::permute2(Tensor *a, int a0, int a1) {
|
|
if (!a) throw Exception("TF.permute2: null");
|
|
int order[] = {a0, a1};
|
|
if (Func *f = tracing()) return f->emitPermute(a, order, 2);
|
|
return a->permute(order, 2);
|
|
}
|
|
Tensor *TF::permute3(Tensor *a, int a0, int a1, int a2) {
|
|
if (!a) throw Exception("TF.permute3: null");
|
|
int order[] = {a0, a1, a2};
|
|
if (Func *f = tracing()) return f->emitPermute(a, order, 3);
|
|
return a->permute(order, 3);
|
|
}
|
|
Tensor *TF::permute4(Tensor *a, int a0, int a1, int a2, int a3) {
|
|
if (!a) throw Exception("TF.permute4: null");
|
|
int order[] = {a0, a1, a2, a3};
|
|
if (Func *f = tracing()) return f->emitPermute(a, order, 4);
|
|
return a->permute(order, 4);
|
|
}
|
|
Tensor *TF::permute5(Tensor *a, int a0, int a1, int a2, int a3, int a4) {
|
|
if (!a) throw Exception("TF.permute5: null");
|
|
int order[] = {a0, a1, a2, a3, a4};
|
|
if (Func *f = tracing()) return f->emitPermute(a, order, 5);
|
|
return a->permute(order, 5);
|
|
}
|
|
Tensor *TF::permute6(Tensor *a, int a0, int a1, int a2, int a3, int a4, int a5) {
|
|
if (!a) throw Exception("TF.permute6: null");
|
|
int order[] = {a0, a1, a2, a3, a4, a5};
|
|
if (Func *f = tracing()) return f->emitPermute(a, order, 6);
|
|
return a->permute(order, 6);
|
|
}
|
|
|
|
Tensor *TF::reshape1(Tensor *a, int d0) {
|
|
if (!a) throw Exception("TF.reshape1: null");
|
|
int d[] = {d0};
|
|
if (Func *f = tracing()) return f->emitReshape(a, d, 1);
|
|
return a->reshape1(d0);
|
|
}
|
|
Tensor *TF::reshape2(Tensor *a, int d0, int d1) {
|
|
if (!a) throw Exception("TF.reshape2: null");
|
|
int d[] = {d0, d1};
|
|
if (Func *f = tracing()) return f->emitReshape(a, d, 2);
|
|
return a->reshape2(d0, d1);
|
|
}
|
|
Tensor *TF::reshape3(Tensor *a, int d0, int d1, int d2) {
|
|
if (!a) throw Exception("TF.reshape3: null");
|
|
int d[] = {d0, d1, d2};
|
|
if (Func *f = tracing()) return f->emitReshape(a, d, 3);
|
|
return a->reshape3(d0, d1, d2);
|
|
}
|
|
Tensor *TF::reshape4(Tensor *a, int d0, int d1, int d2, int d3) {
|
|
if (!a) throw Exception("TF.reshape4: null");
|
|
int d[] = {d0, d1, d2, d3};
|
|
if (Func *f = tracing()) return f->emitReshape(a, d, 4);
|
|
return a->reshape4(d0, d1, d2, d3);
|
|
}
|
|
Tensor *TF::reshape5(Tensor *a, int d0, int d1, int d2, int d3, int d4) {
|
|
if (!a) throw Exception("TF.reshape5: null");
|
|
int d[] = {d0, d1, d2, d3, d4};
|
|
if (Func *f = tracing()) return f->emitReshape(a, d, 5);
|
|
return a->reshape5(d0, d1, d2, d3, d4);
|
|
}
|
|
Tensor *TF::reshape6(Tensor *a, int d0, int d1, int d2, int d3, int d4, int d5) {
|
|
if (!a) throw Exception("TF.reshape6: null");
|
|
int d[] = {d0, d1, d2, d3, d4, d5};
|
|
if (Func *f = tracing()) return f->emitReshape(a, d, 6);
|
|
return a->reshape6(d0, d1, d2, d3, d4, d5);
|
|
}
|
|
|
|
Tensor *TF::flatten(Tensor *a) {
|
|
if (!a) throw Exception("TF.flatten: null");
|
|
int d[] = {a->getSize()};
|
|
if (Func *f = tracing()) return f->emitReshape(a, d, 1);
|
|
return a->flatten();
|
|
}
|
|
|
|
Tensor *TF::where(Tensor *cond, Tensor *a, Tensor *b) {
|
|
if (!cond || !a || !b) throw Exception("TF.where: null");
|
|
if (Func *f = tracing()) return f->emitTernary(OpType::Where, cond, a, b);
|
|
if (cond->getRank() != a->getRank() || a->getRank() != b->getRank() ||
|
|
cond->getSize() != a->getSize() || a->getSize() != b->getSize())
|
|
throw Exception("TF.where: shape mismatch");
|
|
for (int i = 0; i < cond->getRank(); ++i) {
|
|
if (cond->getDim(i) != a->getDim(i) || a->getDim(i) != b->getDim(i))
|
|
throw Exception("TF.where: shape mismatch");
|
|
}
|
|
int dims[Tensor::kMaxRank] = {};
|
|
for (int i = 0; i < a->getRank(); ++i) dims[i] = a->getDim(i);
|
|
auto *out = new Tensor(dims, a->getRank());
|
|
for (int i = 0; i < a->getSize(); ++i)
|
|
out->set(i, cond->get(i) > 0.5f ? a->get(i) : b->get(i));
|
|
return out;
|
|
}
|
|
|
|
// --- neural / speech / terrain ops ---------------------------------------
|
|
|
|
Tensor *TF::softmax(Tensor *a, int axis) {
|
|
if (!a) throw Exception("TF.softmax: null");
|
|
if (Func *f = tracing()) return f->emitSoftmax(a, axis, false);
|
|
a->ensureEager("softmax");
|
|
axis = normalizeAxis(axis, a->getRank());
|
|
int dims[Tensor::kMaxRank] = {};
|
|
for (int k = 0; k < a->getRank(); ++k) dims[k] = a->getDim(k);
|
|
auto *out = new Tensor(dims, a->getRank());
|
|
kernels::softmax(a->data(), dims, a->getRank(), axis, false, out->data());
|
|
return out;
|
|
}
|
|
|
|
Tensor *TF::logSoftmax(Tensor *a, int axis) {
|
|
if (!a) throw Exception("TF.logSoftmax: null");
|
|
if (Func *f = tracing()) return f->emitSoftmax(a, axis, true);
|
|
a->ensureEager("logSoftmax");
|
|
axis = normalizeAxis(axis, a->getRank());
|
|
int dims[Tensor::kMaxRank] = {};
|
|
for (int k = 0; k < a->getRank(); ++k) dims[k] = a->getDim(k);
|
|
auto *out = new Tensor(dims, a->getRank());
|
|
kernels::softmax(a->data(), dims, a->getRank(), axis, true, out->data());
|
|
return out;
|
|
}
|
|
|
|
Tensor *TF::layernorm(Tensor *a, float eps) {
|
|
if (!a) throw Exception("TF.layernorm: null");
|
|
if (Func *f = tracing()) return f->emitLayerNorm(a, nullptr, nullptr, eps);
|
|
a->ensureEager("layernorm");
|
|
const int cols = a->getDim(a->getRank() - 1);
|
|
const int rows = a->getSize() / cols;
|
|
int dims[Tensor::kMaxRank] = {};
|
|
for (int k = 0; k < a->getRank(); ++k) dims[k] = a->getDim(k);
|
|
auto *out = new Tensor(dims, a->getRank());
|
|
kernels::layernorm(a->data(), rows, cols, nullptr, nullptr, eps, out->data());
|
|
return out;
|
|
}
|
|
|
|
Tensor *TF::layernormWB(Tensor *a, Tensor *scale, Tensor *bias, float eps) {
|
|
if (!a || !scale || !bias) throw Exception("TF.layernormWB: null");
|
|
if (Func *f = tracing()) return f->emitLayerNorm(a, scale, bias, eps);
|
|
a->ensureEager("layernormWB");
|
|
const int cols = a->getDim(a->getRank() - 1);
|
|
if (scale->getSize() != cols || bias->getSize() != cols)
|
|
throw Exception("TF.layernormWB: scale/bias must match last dim");
|
|
const int rows = a->getSize() / cols;
|
|
int dims[Tensor::kMaxRank] = {};
|
|
for (int k = 0; k < a->getRank(); ++k) dims[k] = a->getDim(k);
|
|
auto *out = new Tensor(dims, a->getRank());
|
|
kernels::layernorm(a->data(), rows, cols, scale->data(), bias->data(), eps, out->data());
|
|
return out;
|
|
}
|
|
|
|
Tensor *TF::rmsnorm(Tensor *a, float eps) {
|
|
if (!a) throw Exception("TF.rmsnorm: null");
|
|
if (Func *f = tracing()) return f->emitRMSNorm(a, nullptr, eps);
|
|
a->ensureEager("rmsnorm");
|
|
const int cols = a->getDim(a->getRank() - 1);
|
|
const int rows = a->getSize() / cols;
|
|
int dims[Tensor::kMaxRank] = {};
|
|
for (int k = 0; k < a->getRank(); ++k) dims[k] = a->getDim(k);
|
|
auto *out = new Tensor(dims, a->getRank());
|
|
kernels::rmsnorm(a->data(), rows, cols, nullptr, eps, out->data());
|
|
return out;
|
|
}
|
|
|
|
Tensor *TF::rmsnormW(Tensor *a, Tensor *scale, float eps) {
|
|
if (!a || !scale) throw Exception("TF.rmsnormW: null");
|
|
if (Func *f = tracing()) return f->emitRMSNorm(a, scale, eps);
|
|
a->ensureEager("rmsnormW");
|
|
const int cols = a->getDim(a->getRank() - 1);
|
|
if (scale->getSize() != cols) throw Exception("TF.rmsnormW: scale must match last dim");
|
|
const int rows = a->getSize() / cols;
|
|
int dims[Tensor::kMaxRank] = {};
|
|
for (int k = 0; k < a->getRank(); ++k) dims[k] = a->getDim(k);
|
|
auto *out = new Tensor(dims, a->getRank());
|
|
kernels::rmsnorm(a->data(), rows, cols, scale->data(), eps, out->data());
|
|
return out;
|
|
}
|
|
|
|
Tensor *TF::conv1d(Tensor *x, Tensor *w, int stride, int pad) {
|
|
return conv1dBias(x, w, nullptr, stride, pad);
|
|
}
|
|
|
|
Tensor *TF::conv1dBias(Tensor *x, Tensor *w, Tensor *bias, int stride, int pad) {
|
|
if (!x || !w) throw Exception("TF.conv1d: null");
|
|
if (Func *f = tracing()) return f->emitConv1d(x, w, bias, stride, pad);
|
|
if (x->getRank() != 3 || w->getRank() != 3) throw Exception("TF.conv1d: rank 3 required");
|
|
if (x->getDim(1) != w->getDim(1)) throw Exception("TF.conv1d: channel mismatch");
|
|
if (bias && bias->getSize() != w->getDim(0)) throw Exception("TF.conv1d: bias mismatch");
|
|
const int OL = convOutSize(x->getDim(2), w->getDim(2), stride, pad);
|
|
auto *out = new Tensor(x->getDim(0), w->getDim(0), OL);
|
|
int xd[3] = {x->getDim(0), x->getDim(1), x->getDim(2)};
|
|
int wd[3] = {w->getDim(0), w->getDim(1), w->getDim(2)};
|
|
kernels::conv1d(x->data(), xd, w->data(), wd, bias ? bias->data() : nullptr, stride, pad,
|
|
out->data());
|
|
return out;
|
|
}
|
|
|
|
Tensor *TF::conv2d(Tensor *x, Tensor *w, int stride, int pad) {
|
|
return conv2dBias(x, w, nullptr, stride, pad);
|
|
}
|
|
|
|
Tensor *TF::conv2dBias(Tensor *x, Tensor *w, Tensor *bias, int stride, int pad) {
|
|
if (!x || !w) throw Exception("TF.conv2d: null");
|
|
if (Func *f = tracing()) return f->emitConv2d(x, w, bias, stride, pad);
|
|
if (x->getRank() != 4 || w->getRank() != 4) throw Exception("TF.conv2d: rank 4 required");
|
|
if (x->getDim(1) != w->getDim(1)) throw Exception("TF.conv2d: channel mismatch");
|
|
if (bias && bias->getSize() != w->getDim(0)) throw Exception("TF.conv2d: bias mismatch");
|
|
const int OH = convOutSize(x->getDim(2), w->getDim(2), stride, pad);
|
|
const int OW = convOutSize(x->getDim(3), w->getDim(3), stride, pad);
|
|
auto *out = new Tensor(x->getDim(0), w->getDim(0), OH, OW);
|
|
int xd[4] = {x->getDim(0), x->getDim(1), x->getDim(2), x->getDim(3)};
|
|
int wd[4] = {w->getDim(0), w->getDim(1), w->getDim(2), w->getDim(3)};
|
|
kernels::conv2d(x->data(), xd, w->data(), wd, bias ? bias->data() : nullptr, stride, pad,
|
|
out->data());
|
|
return out;
|
|
}
|
|
|
|
Tensor *TF::maxpool2d(Tensor *x, int ksize, int stride, int pad) {
|
|
if (!x) throw Exception("TF.maxpool2d: null");
|
|
if (Func *f = tracing()) return f->emitPool(OpType::MaxPool2d, x, ksize, stride, pad);
|
|
if (x->getRank() != 4) throw Exception("TF.maxpool2d: rank 4 required");
|
|
const int OH = convOutSize(x->getDim(2), ksize, stride, pad);
|
|
const int OW = convOutSize(x->getDim(3), ksize, stride, pad);
|
|
auto *out = new Tensor(x->getDim(0), x->getDim(1), OH, OW);
|
|
int xd[4] = {x->getDim(0), x->getDim(1), x->getDim(2), x->getDim(3)};
|
|
kernels::maxpool2d(x->data(), xd, ksize, stride, pad, out->data());
|
|
return out;
|
|
}
|
|
|
|
Tensor *TF::avgpool2d(Tensor *x, int ksize, int stride, int pad) {
|
|
if (!x) throw Exception("TF.avgpool2d: null");
|
|
if (Func *f = tracing()) return f->emitPool(OpType::AvgPool2d, x, ksize, stride, pad);
|
|
if (x->getRank() != 4) throw Exception("TF.avgpool2d: rank 4 required");
|
|
const int OH = convOutSize(x->getDim(2), ksize, stride, pad);
|
|
const int OW = convOutSize(x->getDim(3), ksize, stride, pad);
|
|
auto *out = new Tensor(x->getDim(0), x->getDim(1), OH, OW);
|
|
int xd[4] = {x->getDim(0), x->getDim(1), x->getDim(2), x->getDim(3)};
|
|
kernels::avgpool2d(x->data(), xd, ksize, stride, pad, out->data());
|
|
return out;
|
|
}
|
|
|
|
Tensor *TF::embedding(Tensor *table, Tensor *indices) {
|
|
if (!table || !indices) throw Exception("TF.embedding: null");
|
|
if (Func *f = tracing()) return f->emitEmbedding(table, indices);
|
|
if (table->getRank() != 2) throw Exception("TF.embedding: table rank 2 required");
|
|
int dims[Tensor::kMaxRank] = {};
|
|
for (int k = 0; k < indices->getRank(); ++k) dims[k] = indices->getDim(k);
|
|
dims[indices->getRank()] = table->getDim(1);
|
|
auto *out = new Tensor(dims, indices->getRank() + 1);
|
|
if (table->isQuantized()) {
|
|
const std::vector<float> tf32 = table->dequantized();
|
|
kernels::embedding(tf32.data(), table->getDim(0), table->getDim(1), indices->data(),
|
|
indices->getSize(), out->data());
|
|
} else {
|
|
kernels::embedding(table->data(), table->getDim(0), table->getDim(1), indices->data(),
|
|
indices->getSize(), out->data());
|
|
}
|
|
return out;
|
|
}
|
|
|
|
Tensor *TF::quantizeWeight(Tensor *a, const std::string &dtype, int group) {
|
|
if (!a) throw Exception("TF.quantizeWeight: null");
|
|
if (tracing()) throw Exception("TF.quantizeWeight: quantize eager weights before tracing");
|
|
DType dt = DType::Float32;
|
|
if (!parseDType(dtype, dt) || !q::isQuantDType(dt))
|
|
throw Exception("TF.quantizeWeight: expected fp16/fp8/fp4/int8/int4, got '%s'",
|
|
dtype.c_str());
|
|
a->ensureEager("quantizeWeight");
|
|
if (a->isQuantized()) throw Exception("TF.quantizeWeight: input is already quantized");
|
|
q::QuantPayload p = q::quantize(a->data(), a->getSize(), dt, group);
|
|
auto *out = new Tensor(dt, a->getRank() > 0 ? a->dims_ : nullptr, a->getRank());
|
|
out->bytes_ = std::move(p.bytes);
|
|
out->qScales_ = std::move(p.scales);
|
|
out->qGroup_ = p.group;
|
|
return out;
|
|
}
|
|
|
|
Tensor *TF::concat2(Tensor *a, Tensor *b, int axis) {
|
|
Tensor *ins[] = {a, b};
|
|
return concatN(ins, 2, axis);
|
|
}
|
|
|
|
Tensor *TF::concat3(Tensor *a, Tensor *b, Tensor *c, int axis) {
|
|
Tensor *ins[] = {a, b, c};
|
|
return concatN(ins, 3, axis);
|
|
}
|
|
|
|
Tensor *TF::concat4(Tensor *a, Tensor *b, Tensor *c, Tensor *d, int axis) {
|
|
Tensor *ins[] = {a, b, c, d};
|
|
return concatN(ins, 4, axis);
|
|
}
|
|
|
|
namespace {
|
|
|
|
Tensor *concatEager(Tensor *const *ins, int n, int axis) {
|
|
for (int k = 0; k < n; ++k)
|
|
if (!ins[k]) throw Exception("TF.concat: null");
|
|
const int rank = ins[0]->getRank();
|
|
axis = normalizeAxis(axis, rank);
|
|
int dims[Tensor::kMaxRank] = {};
|
|
for (int k = 0; k < rank; ++k) {
|
|
if (k == axis) {
|
|
int total = 0;
|
|
for (int t = 0; t < n; ++t) total += ins[t]->getDim(k);
|
|
dims[k] = total;
|
|
} else {
|
|
dims[k] = ins[0]->getDim(k);
|
|
for (int t = 1; t < n; ++t)
|
|
if (ins[t]->getDim(k) != dims[k]) throw Exception("TF.concat: dims mismatch");
|
|
}
|
|
}
|
|
auto *out = new Tensor(dims, rank);
|
|
const float *ptrs[4] = {};
|
|
int inDims[4][Tensor::kMaxRank] = {};
|
|
const int *dimsPtr[4] = {};
|
|
int inRanks[4] = {};
|
|
for (int k = 0; k < n; ++k) {
|
|
ptrs[k] = ins[k]->data();
|
|
for (int d = 0; d < rank; ++d) inDims[k][d] = ins[k]->getDim(d);
|
|
dimsPtr[k] = inDims[k];
|
|
inRanks[k] = rank;
|
|
}
|
|
kernels::concat(ptrs, dimsPtr, inRanks, n, axis, out->data(), dims, rank);
|
|
return out;
|
|
}
|
|
|
|
} // namespace
|
|
|
|
Tensor *TF::concatN(Tensor *const *ins, int n, int axis) {
|
|
if (Func *f = tracing()) return f->emitConcat(const_cast<const Tensor *const *>(ins), n, axis);
|
|
return concatEager(ins, n, axis);
|
|
}
|
|
|
|
Tensor *TF::slice(Tensor *a, int axis, int begin, int end) {
|
|
if (!a) throw Exception("TF.slice: null");
|
|
if (Func *f = tracing()) return f->emitSlice(a, axis, begin, end);
|
|
a->ensureEager("slice");
|
|
axis = normalizeAxis(axis, a->getRank());
|
|
if (begin < 0 || end < begin || end > a->getDim(axis))
|
|
throw Exception("TF.slice: range out of bounds");
|
|
int srcDims[Tensor::kMaxRank] = {};
|
|
int dims[Tensor::kMaxRank] = {};
|
|
for (int k = 0; k < a->getRank(); ++k) {
|
|
srcDims[k] = a->getDim(k);
|
|
dims[k] = a->getDim(k);
|
|
}
|
|
dims[axis] = end - begin;
|
|
auto *out = new Tensor(dims, a->getRank());
|
|
kernels::sliceOp(a->data(), srcDims, a->getRank(), axis, begin, end, out->data(), dims,
|
|
a->getRank());
|
|
return out;
|
|
}
|
|
|
|
#define EVE_TF_AXIS_REDUCE(name, opType) \
|
|
Tensor *TF::name(Tensor *a, int axis, int keepDims) { \
|
|
if (!a) throw Exception("TF." #name ": null"); \
|
|
if (Func *f = tracing()) return f->emitReduce(opType, a, axis, keepDims != 0); \
|
|
a->ensureEager(#name); \
|
|
axis = normalizeAxis(axis, a->getRank()); \
|
|
int srcDims[Tensor::kMaxRank] = {}; \
|
|
for (int k = 0; k < a->getRank(); ++k) srcDims[k] = a->getDim(k); \
|
|
int od[Tensor::kMaxRank] = {}; \
|
|
int outRank = 0; \
|
|
reduceOutDims(srcDims, a->getRank(), axis, keepDims != 0, od, outRank); \
|
|
auto *out = new Tensor(od, outRank); \
|
|
kernels::reduceAxis(opType, a->data(), srcDims, a->getRank(), axis, out->data(), od, \
|
|
outRank); \
|
|
return out; \
|
|
}
|
|
|
|
EVE_TF_AXIS_REDUCE(sumAxis, OpType::ReduceSum)
|
|
EVE_TF_AXIS_REDUCE(meanAxis, OpType::ReduceMean)
|
|
EVE_TF_AXIS_REDUCE(minAxis, OpType::ReduceMin)
|
|
EVE_TF_AXIS_REDUCE(maxAxis, OpType::ReduceMax)
|
|
|
|
#undef EVE_TF_AXIS_REDUCE
|
|
|
|
Tensor *TF::argmax(Tensor *a, int axis, int keepDims) {
|
|
if (!a) throw Exception("TF.argmax: null");
|
|
if (Func *f = tracing()) return f->emitArgMax(a, axis, keepDims != 0);
|
|
a->ensureEager("argmax");
|
|
axis = normalizeAxis(axis, a->getRank());
|
|
int srcDims[Tensor::kMaxRank] = {};
|
|
for (int k = 0; k < a->getRank(); ++k) srcDims[k] = a->getDim(k);
|
|
int od[Tensor::kMaxRank] = {};
|
|
int outRank = 0;
|
|
reduceOutDims(srcDims, a->getRank(), axis, keepDims != 0, od, outRank);
|
|
auto *out = new Tensor(DType::Int32, od, outRank);
|
|
kernels::argmax(a->data(), srcDims, a->getRank(), axis, out->data(), od, outRank);
|
|
return out;
|
|
}
|
|
|
|
Tensor *TF::cast(Tensor *a, const std::string &dtype) {
|
|
if (!a) throw Exception("TF.cast: null");
|
|
DType dt = DType::Float32;
|
|
if (!parseDType(dtype, dt)) throw Exception("TF.cast: unknown dtype '%s'", dtype.c_str());
|
|
if (Func *f = tracing()) return f->emitCast(a, dt);
|
|
a->ensureEager("cast");
|
|
int dims[Tensor::kMaxRank] = {};
|
|
for (int k = 0; k < a->getRank(); ++k) dims[k] = a->getDim(k);
|
|
auto *out = new Tensor(dt, dims, a->getRank());
|
|
std::memcpy(out->data(), a->data(), sizeof(float) * static_cast<size_t>(a->getSize()));
|
|
return out;
|
|
}
|
|
|
|
Tensor *TF::sdpa(Tensor *q, Tensor *k, Tensor *v, float scale) {
|
|
return sdpaMasked(q, k, v, nullptr, scale);
|
|
}
|
|
|
|
Tensor *TF::sdpaMasked(Tensor *q, Tensor *k, Tensor *v, Tensor *mask, float scale) {
|
|
if (!q || !k || !v) throw Exception("TF.sdpa: null");
|
|
if (Func *f = tracing()) return f->emitSdpa(q, k, v, mask, scale);
|
|
if (q->getRank() != 4 || k->getRank() != 4 || v->getRank() != 4)
|
|
throw Exception("TF.sdpa: rank 4 required");
|
|
if (q->getDim(0) != k->getDim(0) || q->getDim(1) != k->getDim(1) ||
|
|
q->getDim(3) != k->getDim(3) || k->getDim(2) != v->getDim(2) ||
|
|
q->getDim(3) != v->getDim(3))
|
|
throw Exception("TF.sdpa: q/k/v shape mismatch");
|
|
if (mask && (mask->getRank() != 4 || mask->getDim(0) != q->getDim(0) ||
|
|
mask->getDim(1) != q->getDim(1) || mask->getDim(2) != q->getDim(2) ||
|
|
mask->getDim(3) != k->getDim(2)))
|
|
throw Exception("TF.sdpa: mask shape mismatch");
|
|
auto *out = new Tensor(q->getDim(0), q->getDim(1), q->getDim(2), q->getDim(3));
|
|
kernels::sdpa(q->data(), k->data(), v->data(), mask ? mask->data() : nullptr, q->getDim(0),
|
|
q->getDim(1), q->getDim(2), k->getDim(2), q->getDim(3), scale, out->data());
|
|
return out;
|
|
}
|
|
|
|
Tensor *TF::resize2d(Tensor *a, int outW, int outH, int mode) {
|
|
if (!a) throw Exception("TF.resize2d: null");
|
|
if (Func *f = tracing()) return f->emitResize2d(a, outH, outW, mode);
|
|
if (a->getRank() != 4) throw Exception("TF.resize2d: rank 4 required");
|
|
if (outW <= 0 || outH <= 0) throw Exception("TF.resize2d: bad output size");
|
|
auto *out = new Tensor(a->getDim(0), a->getDim(1), outH, outW);
|
|
int xd[4] = {a->getDim(0), a->getDim(1), a->getDim(2), a->getDim(3)};
|
|
kernels::resize2d(a->data(), xd, outW, outH, mode, out->data());
|
|
return out;
|
|
}
|
|
|
|
float TF::reduceSum(Tensor *a) {
|
|
if (!a) throw Exception("TF.reduceSum: null");
|
|
if (tracing()) throw Exception("TF.reduceSum: not supported while tracing");
|
|
a->ensureEager("reduceSum");
|
|
float gpuResult = 0.f;
|
|
if (a->getSize() >= kGpuReduceMinSize && gpuReduce(a->data(), a->getSize(), 0, gpuResult))
|
|
return gpuResult;
|
|
return a->reduceSum();
|
|
}
|
|
float TF::reduceMean(Tensor *a) {
|
|
if (!a) throw Exception("TF.reduceMean: null");
|
|
if (tracing()) throw Exception("TF.reduceMean: not supported while tracing");
|
|
a->ensureEager("reduceMean");
|
|
if (a->getSize() <= 0) return 0.f;
|
|
return reduceSum(a) / float(a->getSize());
|
|
}
|
|
float TF::reduceMin(Tensor *a) {
|
|
if (!a) throw Exception("TF.reduceMin: null");
|
|
if (tracing()) throw Exception("TF.reduceMin: not supported while tracing");
|
|
a->ensureEager("reduceMin");
|
|
float gpuResult = 0.f;
|
|
if (a->getSize() >= kGpuReduceMinSize && gpuReduce(a->data(), a->getSize(), 1, gpuResult))
|
|
return gpuResult;
|
|
return a->reduceMin();
|
|
}
|
|
float TF::reduceMax(Tensor *a) {
|
|
if (!a) throw Exception("TF.reduceMax: null");
|
|
if (tracing()) throw Exception("TF.reduceMax: not supported while tracing");
|
|
a->ensureEager("reduceMax");
|
|
float gpuResult = 0.f;
|
|
if (a->getSize() >= kGpuReduceMinSize && gpuReduce(a->data(), a->getSize(), 2, gpuResult))
|
|
return gpuResult;
|
|
return a->reduceMax();
|
|
}
|
|
|
|
void TF::expose(ssq::Table &table) {
|
|
auto cls = table.addClass(name, TF::create, false);
|
|
expose(cls);
|
|
|
|
auto ten = table.addClass<Tensor>(
|
|
"Tensor", std::function<Tensor *()>([]() -> Tensor * { return nullptr; }), true);
|
|
|
|
ten.addFunc("isSymbolic", &Tensor::isSymbolic);
|
|
ten.addFunc("isEager", &Tensor::isEager);
|
|
ten.addFunc("getRank", &Tensor::getRank);
|
|
ten.addFunc("getSize", &Tensor::getSize);
|
|
ten.addFunc("getDim", &Tensor::getDim);
|
|
ten.addFunc("getDim0", &Tensor::getDim0);
|
|
ten.addFunc("getDim1", &Tensor::getDim1);
|
|
ten.addFunc("getDim2", &Tensor::getDim2);
|
|
ten.addFunc("getDim3", &Tensor::getDim3);
|
|
ten.addFunc("getDim4", &Tensor::getDim4);
|
|
ten.addFunc("getDim5", &Tensor::getDim5);
|
|
ten.addFunc("getDevice", &Tensor::getDevice);
|
|
ten.addFunc("getDtype", &Tensor::getDtype);
|
|
ten.addFunc("isQuantized", &Tensor::isQuantized);
|
|
|
|
ten.addFunc("get", &Tensor::get);
|
|
ten.addFunc("set", &Tensor::set);
|
|
ten.addFunc("get1", &Tensor::get1);
|
|
ten.addFunc("set1", &Tensor::set1);
|
|
ten.addFunc("get2", &Tensor::get2);
|
|
ten.addFunc("set2", &Tensor::set2);
|
|
ten.addFunc("get3", &Tensor::get3);
|
|
ten.addFunc("set3", &Tensor::set3);
|
|
ten.addFunc("get4", &Tensor::get4);
|
|
ten.addFunc("set4", &Tensor::set4);
|
|
ten.addFunc("get5", &Tensor::get5);
|
|
ten.addFunc("set5", &Tensor::set5);
|
|
ten.addFunc("get6", &Tensor::get6);
|
|
ten.addFunc("set6", &Tensor::set6);
|
|
|
|
ten.addFunc("fill", &Tensor::fill);
|
|
ten.addFunc("copyFrom", &Tensor::copyFrom);
|
|
ten.addFunc("clone", &Tensor::clone);
|
|
|
|
ten.addFunc("add", &Tensor::add);
|
|
ten.addFunc("sub", &Tensor::sub);
|
|
ten.addFunc("multiply", &Tensor::multiply);
|
|
ten.addFunc("div", &Tensor::div);
|
|
ten.addFunc("addScalar", &Tensor::addScalar);
|
|
ten.addFunc("subScalar", &Tensor::subScalar);
|
|
ten.addFunc("mulScalar", &Tensor::mulScalar);
|
|
ten.addFunc("divScalar", &Tensor::divScalar);
|
|
ten.addFunc("neg", &Tensor::neg);
|
|
ten.addFunc("abs", &Tensor::abs);
|
|
ten.addFunc("sqrt", &Tensor::sqrt);
|
|
ten.addFunc("exp", &Tensor::exp);
|
|
ten.addFunc("log", &Tensor::log);
|
|
ten.addFunc("sin", &Tensor::sin);
|
|
ten.addFunc("cos", &Tensor::cos);
|
|
ten.addFunc("tanh", &Tensor::tanh);
|
|
ten.addFunc("relu", &Tensor::relu);
|
|
ten.addFunc("sigmoid", &Tensor::sigmoid);
|
|
ten.addFunc("gelu", &Tensor::gelu);
|
|
ten.addFunc("silu", &Tensor::silu);
|
|
ten.addFunc("powScalar", &Tensor::powScalar);
|
|
ten.addFunc("clamp", &Tensor::clamp);
|
|
ten.addFunc("maximumScalar", &Tensor::maximumScalar);
|
|
ten.addFunc("minimumScalar", &Tensor::minimumScalar);
|
|
|
|
ten.addFunc("addInPlace", &Tensor::addInPlace);
|
|
ten.addFunc("multiplyInPlace", &Tensor::multiplyInPlace);
|
|
ten.addFunc("addScalarInPlace", &Tensor::addScalarInPlace);
|
|
ten.addFunc("mulScalarInPlace", &Tensor::mulScalarInPlace);
|
|
ten.addFunc("reluInPlace", &Tensor::reluInPlace);
|
|
|
|
ten.addFunc("reduceSum", &Tensor::reduceSum);
|
|
ten.addFunc("reduceMean", &Tensor::reduceMean);
|
|
ten.addFunc("reduceMin", &Tensor::reduceMin);
|
|
ten.addFunc("reduceMax", &Tensor::reduceMax);
|
|
ten.addFunc("dot", &Tensor::dot);
|
|
|
|
ten.addFunc("matmul", &Tensor::matmul);
|
|
ten.addFunc("transpose", &Tensor::transpose);
|
|
ten.addFunc("reshape1", &Tensor::reshape1);
|
|
ten.addFunc("reshape2", &Tensor::reshape2);
|
|
ten.addFunc("reshape3", &Tensor::reshape3);
|
|
ten.addFunc("reshape4", &Tensor::reshape4);
|
|
ten.addFunc("reshape5", &Tensor::reshape5);
|
|
ten.addFunc("reshape6", &Tensor::reshape6);
|
|
ten.addFunc("flatten", &Tensor::flatten);
|
|
|
|
auto fn = table.addClass<Func>(
|
|
"Func", std::function<Func *()>([]() -> Func * { return nullptr; }), true);
|
|
fn.addFunc("input1", &Func::input1);
|
|
fn.addFunc("input2", &Func::input2);
|
|
fn.addFunc("input3", &Func::input3);
|
|
fn.addFunc("input4", &Func::input4);
|
|
fn.addFunc("input5", &Func::input5);
|
|
fn.addFunc("input6", &Func::input6);
|
|
fn.addFunc("setOutput", &Func::setOutput);
|
|
fn.addFunc("compile", &Func::compile);
|
|
|
|
auto cf = table.addClass<CompiledFunction>(
|
|
"CompiledFunction",
|
|
std::function<CompiledFunction *()>([]() -> CompiledFunction * { return nullptr; }), true);
|
|
cf.addFunc("run0", &CompiledFunction::run0);
|
|
cf.addFunc("run1", &CompiledFunction::run1);
|
|
cf.addFunc("run2", &CompiledFunction::run2);
|
|
cf.addFunc("run3", &CompiledFunction::run3);
|
|
cf.addFunc("run4", &CompiledFunction::run4);
|
|
cf.addFunc("run5", &CompiledFunction::run5);
|
|
cf.addFunc("run6", &CompiledFunction::run6);
|
|
cf.addFunc("getPlaceholderCount", &CompiledFunction::getPlaceholderCount);
|
|
cf.addFunc("getDevice", &CompiledFunction::getDevice);
|
|
}
|
|
|
|
void TF::expose(ssq::Class &cls) {
|
|
cls.addFunc("getName", &TF::getName);
|
|
cls.addFunc("func", &TF::func);
|
|
cls.addFunc("quantizeWeight", &TF::quantizeWeight);
|
|
|
|
cls.addFunc("zeros1", &TF::zeros1);
|
|
cls.addFunc("zeros2", &TF::zeros2);
|
|
cls.addFunc("zeros3", &TF::zeros3);
|
|
cls.addFunc("zeros4", &TF::zeros4);
|
|
cls.addFunc("zeros5", &TF::zeros5);
|
|
cls.addFunc("zeros6", &TF::zeros6);
|
|
cls.addFunc("ones1", &TF::ones1);
|
|
cls.addFunc("ones2", &TF::ones2);
|
|
cls.addFunc("ones3", &TF::ones3);
|
|
cls.addFunc("ones4", &TF::ones4);
|
|
cls.addFunc("ones5", &TF::ones5);
|
|
cls.addFunc("ones6", &TF::ones6);
|
|
cls.addFunc("fill1", &TF::fill1);
|
|
cls.addFunc("fill2", &TF::fill2);
|
|
cls.addFunc("fill3", &TF::fill3);
|
|
cls.addFunc("fill4", &TF::fill4);
|
|
cls.addFunc("constantScalar", &TF::constantScalar);
|
|
cls.addFunc("arange", &TF::arange);
|
|
cls.addFunc("linspace", &TF::linspace);
|
|
cls.addFunc("eye", &TF::eye);
|
|
cls.addFunc("randomUniform1", &TF::randomUniform1);
|
|
cls.addFunc("randomUniform2", &TF::randomUniform2);
|
|
cls.addFunc("randomUniform3", &TF::randomUniform3);
|
|
cls.addFunc("randomUniform4", &TF::randomUniform4);
|
|
cls.addFunc("randomNormal1", &TF::randomNormal1);
|
|
cls.addFunc("randomNormal2", &TF::randomNormal2);
|
|
cls.addFunc("randomNormal3", &TF::randomNormal3);
|
|
cls.addFunc("randomNormal4", &TF::randomNormal4);
|
|
cls.addFunc("rand1", &TF::rand1);
|
|
cls.addFunc("rand2", &TF::rand2);
|
|
cls.addFunc("rand3", &TF::rand3);
|
|
cls.addFunc("rand4", &TF::rand4);
|
|
cls.addFunc("randn1", &TF::randn1);
|
|
cls.addFunc("randn2", &TF::randn2);
|
|
cls.addFunc("randn3", &TF::randn3);
|
|
cls.addFunc("randn4", &TF::randn4);
|
|
cls.addFunc("setRandomSeed", &TF::setRandomSeed);
|
|
cls.addFunc("getRandomSeed", &TF::getRandomSeed);
|
|
|
|
cls.addFunc("add", &TF::add);
|
|
cls.addFunc("sub", &TF::sub);
|
|
cls.addFunc("multiply", &TF::multiply);
|
|
cls.addFunc("div", &TF::div);
|
|
cls.addFunc("addScalar", &TF::addScalar);
|
|
cls.addFunc("subScalar", &TF::subScalar);
|
|
cls.addFunc("mulScalar", &TF::mulScalar);
|
|
cls.addFunc("divScalar", &TF::divScalar);
|
|
cls.addFunc("neg", &TF::neg);
|
|
cls.addFunc("abs", &TF::abs);
|
|
cls.addFunc("sqrt", &TF::sqrt);
|
|
cls.addFunc("exp", &TF::exp);
|
|
cls.addFunc("log", &TF::log);
|
|
cls.addFunc("sin", &TF::sin);
|
|
cls.addFunc("cos", &TF::cos);
|
|
cls.addFunc("tanh", &TF::tanh);
|
|
cls.addFunc("relu", &TF::relu);
|
|
cls.addFunc("sigmoid", &TF::sigmoid);
|
|
cls.addFunc("gelu", &TF::gelu);
|
|
cls.addFunc("silu", &TF::silu);
|
|
cls.addFunc("powScalar", &TF::powScalar);
|
|
cls.addFunc("clamp", &TF::clamp);
|
|
cls.addFunc("maximumScalar", &TF::maximumScalar);
|
|
cls.addFunc("minimumScalar", &TF::minimumScalar);
|
|
cls.addFunc("matmul", &TF::matmul);
|
|
cls.addFunc("transpose", &TF::transpose);
|
|
cls.addFunc("permute2", &TF::permute2);
|
|
cls.addFunc("permute3", &TF::permute3);
|
|
cls.addFunc("permute4", &TF::permute4);
|
|
cls.addFunc("permute5", &TF::permute5);
|
|
cls.addFunc("permute6", &TF::permute6);
|
|
cls.addFunc("reshape1", &TF::reshape1);
|
|
cls.addFunc("reshape2", &TF::reshape2);
|
|
cls.addFunc("reshape3", &TF::reshape3);
|
|
cls.addFunc("reshape4", &TF::reshape4);
|
|
cls.addFunc("reshape5", &TF::reshape5);
|
|
cls.addFunc("reshape6", &TF::reshape6);
|
|
cls.addFunc("flatten", &TF::flatten);
|
|
cls.addFunc("where", &TF::where);
|
|
cls.addFunc("reduceSum", &TF::reduceSum);
|
|
cls.addFunc("reduceMean", &TF::reduceMean);
|
|
cls.addFunc("reduceMin", &TF::reduceMin);
|
|
cls.addFunc("reduceMax", &TF::reduceMax);
|
|
|
|
cls.addFunc("softmax", &TF::softmax);
|
|
cls.addFunc("logSoftmax", &TF::logSoftmax);
|
|
cls.addFunc("layernorm", &TF::layernorm);
|
|
cls.addFunc("layernormWB", &TF::layernormWB);
|
|
cls.addFunc("rmsnorm", &TF::rmsnorm);
|
|
cls.addFunc("rmsnormW", &TF::rmsnormW);
|
|
cls.addFunc("conv1d", &TF::conv1d);
|
|
cls.addFunc("conv1dBias", &TF::conv1dBias);
|
|
cls.addFunc("conv2d", &TF::conv2d);
|
|
cls.addFunc("conv2dBias", &TF::conv2dBias);
|
|
cls.addFunc("maxpool2d", &TF::maxpool2d);
|
|
cls.addFunc("avgpool2d", &TF::avgpool2d);
|
|
cls.addFunc("embedding", &TF::embedding);
|
|
cls.addFunc("concat2", &TF::concat2);
|
|
cls.addFunc("concat3", &TF::concat3);
|
|
cls.addFunc("concat4", &TF::concat4);
|
|
cls.addFunc("slice", &TF::slice);
|
|
cls.addFunc("sumAxis", &TF::sumAxis);
|
|
cls.addFunc("meanAxis", &TF::meanAxis);
|
|
cls.addFunc("minAxis", &TF::minAxis);
|
|
cls.addFunc("maxAxis", &TF::maxAxis);
|
|
cls.addFunc("argmax", &TF::argmax);
|
|
cls.addFunc("cast", &TF::cast);
|
|
cls.addFunc("sdpa", &TF::sdpa);
|
|
cls.addFunc("sdpaMasked", &TF::sdpaMasked);
|
|
cls.addFunc("resize2d", &TF::resize2d);
|
|
}
|
|
|
|
} // namespace eve::tensor
|