915 lines
34 KiB
C++
915 lines
34 KiB
C++
#include "tensor/Graph.h"
|
|
#include "tensor/CpuKernels.h"
|
|
#include "tensor/GpuBackend.h"
|
|
#include "tensor/Optimizer.h"
|
|
#include "tensor/Quant.h"
|
|
#include "tensor/TF.h"
|
|
|
|
#include "common/Exception.h"
|
|
|
|
#include <algorithm>
|
|
#include <cmath>
|
|
#include <cstring>
|
|
|
|
namespace eve::tensor {
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Func: tracing graph builder
|
|
// ---------------------------------------------------------------------------
|
|
|
|
Func::Func(TF *owner) : owner_(owner) {
|
|
if (!owner_) throw eve::Exception("Func: null owner");
|
|
owner_->pushTrace(this);
|
|
}
|
|
|
|
Func::~Func() {
|
|
if (tracing_ && owner_) owner_->popTrace(this);
|
|
}
|
|
|
|
GraphNode Func::makeShapeNode(OpType type, const int *dims, int rank) {
|
|
GraphNode n;
|
|
n.type = type;
|
|
n.rank = rank;
|
|
for (int i = 0; i < Tensor::kMaxRank; ++i) n.dims[i] = 0;
|
|
for (int i = 0; i < rank; ++i) n.dims[i] = dims[i];
|
|
n.size = Graph::product(dims, rank);
|
|
return n;
|
|
}
|
|
|
|
Tensor *Func::makeSymbolicFromNode(int nodeId) {
|
|
const auto &n = graph_.node(nodeId);
|
|
auto *t = Tensor::makeSymbolic(&graph_, nodeId, n.dims, n.rank);
|
|
t->setDtype(static_cast<DType>(n.dtype));
|
|
return t;
|
|
}
|
|
|
|
namespace {
|
|
|
|
GraphNode makeInputNode(Graph &graph, OpType type, const int *dims, int rank, int slot) {
|
|
GraphNode n;
|
|
n.type = type;
|
|
n.rank = rank;
|
|
for (int i = 0; i < Tensor::kMaxRank; ++i) n.dims[i] = 0;
|
|
for (int i = 0; i < rank; ++i) n.dims[i] = dims[i];
|
|
n.size = Graph::product(dims, rank);
|
|
n.placeholderSlot = slot;
|
|
return n;
|
|
}
|
|
|
|
int normalizeAxisChecked(int axis, int rank) {
|
|
if (axis < 0) axis += rank;
|
|
if (axis < 0 || axis >= rank) throw eve::Exception("Func: axis out of range");
|
|
return axis;
|
|
}
|
|
|
|
int convOutSize(int inSize, int kernel, int stride, int pad) {
|
|
const int out = (inSize + 2 * pad - kernel) / stride + 1;
|
|
if (out <= 0) throw eve::Exception("Func: conv output size must be > 0");
|
|
return out;
|
|
}
|
|
|
|
} // namespace
|
|
|
|
Tensor *Func::input1(int d0) {
|
|
int d[] = {d0};
|
|
int id = graph_.addNode(makeInputNode(graph_, OpType::Placeholder, d, 1, placeholderCount_++));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
|
|
Tensor *Func::input2(int d0, int d1) {
|
|
int d[] = {d0, d1};
|
|
int id = graph_.addNode(makeInputNode(graph_, OpType::Placeholder, d, 2, placeholderCount_++));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
|
|
Tensor *Func::input3(int d0, int d1, int d2) {
|
|
int d[] = {d0, d1, d2};
|
|
int id = graph_.addNode(makeInputNode(graph_, OpType::Placeholder, d, 3, placeholderCount_++));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
|
|
Tensor *Func::input4(int d0, int d1, int d2, int d3) {
|
|
int d[] = {d0, d1, d2, d3};
|
|
int id = graph_.addNode(makeInputNode(graph_, OpType::Placeholder, d, 4, placeholderCount_++));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
|
|
Tensor *Func::input5(int d0, int d1, int d2, int d3, int d4) {
|
|
int d[] = {d0, d1, d2, d3, d4};
|
|
int id = graph_.addNode(makeInputNode(graph_, OpType::Placeholder, d, 5, placeholderCount_++));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
|
|
Tensor *Func::input6(int d0, int d1, int d2, int d3, int d4, int d5) {
|
|
int d[] = {d0, d1, d2, d3, d4, d5};
|
|
int id = graph_.addNode(makeInputNode(graph_, OpType::Placeholder, d, 6, placeholderCount_++));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
|
|
void Func::setOutput(Tensor *t) {
|
|
if (!t) throw eve::Exception("Func.setOutput: null");
|
|
outputNode_ = ensureNode(t);
|
|
}
|
|
|
|
int Func::ensureNode(const Tensor *t) {
|
|
if (!t) throw eve::Exception("Func.ensureNode: null");
|
|
if (t->isSymbolic()) {
|
|
if (t->graph() != &graph_) throw eve::Exception("Func: tensor from another graph");
|
|
return t->nodeId();
|
|
}
|
|
t->ensureEager("capture");
|
|
auto n = makeShapeNode(OpType::Const, t->dims_, t->rank_);
|
|
if (t->isQuantized()) {
|
|
n.constBytes = t->qBytes();
|
|
n.constScales = t->qScales();
|
|
n.qGroup = t->qGroup();
|
|
} else {
|
|
n.constData.assign(t->data(), t->data() + t->getSize());
|
|
}
|
|
n.dtype = static_cast<int>(t->dtype_);
|
|
return graph_.addNode(std::move(n));
|
|
}
|
|
|
|
Tensor *Func::emitFill(const int *dims, int rank, float value) {
|
|
auto n = makeShapeNode(OpType::Const, dims, rank);
|
|
n.constData.assign(static_cast<size_t>(n.size), value);
|
|
int id = graph_.addNode(std::move(n));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
|
|
Tensor *Func::emitUnary(OpType type, const Tensor *x) {
|
|
int ix = ensureNode(x);
|
|
const auto &src = graph_.node(ix);
|
|
auto n = makeShapeNode(type, src.dims, src.rank);
|
|
n.in0 = ix;
|
|
n.dtype = src.dtype;
|
|
int id = graph_.addNode(std::move(n));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
|
|
Tensor *Func::emitUnaryScalar(OpType type, const Tensor *x, float s0, float s1) {
|
|
int ix = ensureNode(x);
|
|
const auto &src = graph_.node(ix);
|
|
auto n = makeShapeNode(type, src.dims, src.rank);
|
|
n.in0 = ix;
|
|
n.s0 = s0;
|
|
n.s1 = s1;
|
|
n.dtype = src.dtype;
|
|
int id = graph_.addNode(std::move(n));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
|
|
Tensor *Func::emitBinary(OpType type, const Tensor *a, const Tensor *b) {
|
|
int ia = ensureNode(a);
|
|
int ib = ensureNode(b);
|
|
const auto &na = graph_.node(ia);
|
|
const auto &nb = graph_.node(ib);
|
|
int od[Tensor::kMaxRank] = {};
|
|
int orank = 0;
|
|
if (!kernels::broadcastShape(na.dims, na.rank, nb.dims, nb.rank, od, orank))
|
|
throw eve::Exception("Func binary: broadcast shape mismatch");
|
|
auto n = makeShapeNode(type, od, orank);
|
|
n.in0 = ia;
|
|
n.in1 = ib;
|
|
int id = graph_.addNode(std::move(n));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
|
|
Tensor *Func::emitTernary(OpType type, const Tensor *a, const Tensor *b, const Tensor *c) {
|
|
int ia = ensureNode(a);
|
|
int ib = ensureNode(b);
|
|
int ic = ensureNode(c);
|
|
const auto &na = graph_.node(ia);
|
|
auto n = makeShapeNode(type, na.dims, na.rank);
|
|
n.in0 = ia;
|
|
n.in1 = ib;
|
|
n.in2 = ic;
|
|
int id = graph_.addNode(std::move(n));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
|
|
Tensor *Func::emitMatMul(const Tensor *a, const Tensor *b) {
|
|
int ia = ensureNode(a);
|
|
int ib = ensureNode(b);
|
|
const auto &na = graph_.node(ia);
|
|
const auto &nb = graph_.node(ib);
|
|
if (na.rank == 2 && nb.rank == 2) {
|
|
if (na.dims[1] != nb.dims[0]) throw eve::Exception("Func.matmul: inner dims mismatch");
|
|
int od[] = {na.dims[0], nb.dims[1]};
|
|
auto n = makeShapeNode(OpType::MatMul, od, 2);
|
|
n.in0 = ia;
|
|
n.in1 = ib;
|
|
int id = graph_.addNode(std::move(n));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
if (na.rank == 3 && nb.rank == 3) {
|
|
if (na.dims[0] != nb.dims[0] || na.dims[2] != nb.dims[1])
|
|
throw eve::Exception("Func.matmul: batched dims mismatch");
|
|
int od[] = {na.dims[0], na.dims[1], nb.dims[2]};
|
|
auto n = makeShapeNode(OpType::MatMul, od, 3);
|
|
n.in0 = ia;
|
|
n.in1 = ib;
|
|
int id = graph_.addNode(std::move(n));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
throw eve::Exception("Func.matmul: rank 2x2 or 3x3 required");
|
|
}
|
|
|
|
Tensor *Func::emitTranspose(const Tensor *x) {
|
|
int ix = ensureNode(x);
|
|
const auto &src = graph_.node(ix);
|
|
if (src.rank != 2) throw eve::Exception("Func.transpose: rank 2 required");
|
|
int order[] = {1, 0};
|
|
return emitPermute(x, order, 2);
|
|
}
|
|
|
|
Tensor *Func::emitPermute(const Tensor *x, const int *order, int rank) {
|
|
int ix = ensureNode(x);
|
|
const auto &src = graph_.node(ix);
|
|
if (rank != src.rank) throw eve::Exception("Func.permute: rank mismatch");
|
|
int od[Tensor::kMaxRank] = {};
|
|
for (int k = 0; k < rank; ++k) {
|
|
if (order[k] < 0 || order[k] >= rank)
|
|
throw eve::Exception("Func.permute: order out of range");
|
|
od[k] = src.dims[order[k]];
|
|
}
|
|
auto n = makeShapeNode(OpType::Permute, od, rank);
|
|
n.in0 = ix;
|
|
for (int k = 0; k < rank; ++k) n.perm[k] = order[k];
|
|
n.permRank = rank;
|
|
int id = graph_.addNode(std::move(n));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
|
|
Tensor *Func::emitReshape(const Tensor *x, const int *dims, int rank) {
|
|
int ix = ensureNode(x);
|
|
const auto &src = graph_.node(ix);
|
|
int newSize = Graph::product(dims, rank);
|
|
if (newSize != src.size) throw eve::Exception("Func.reshape: size mismatch");
|
|
auto n = makeShapeNode(OpType::Reshape, dims, rank);
|
|
n.in0 = ix;
|
|
n.dtype = src.dtype;
|
|
int id = graph_.addNode(std::move(n));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
|
|
Tensor *Func::emitSoftmax(const Tensor *x, int axis, bool logMode) {
|
|
int ix = ensureNode(x);
|
|
const auto &src = graph_.node(ix);
|
|
axis = normalizeAxisChecked(axis, src.rank);
|
|
auto n = makeShapeNode(logMode ? OpType::LogSoftmax : OpType::Softmax, src.dims, src.rank);
|
|
n.in0 = ix;
|
|
n.i0 = axis;
|
|
int id = graph_.addNode(std::move(n));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
|
|
namespace {
|
|
|
|
void checkParamVector(const GraphNode &src, int cols, const char *what) {
|
|
if (src.rank == 1 && src.dims[0] == cols) return;
|
|
if (src.rank == 2 && src.dims[0] == 1 && src.dims[1] == cols) return;
|
|
throw eve::Exception("Func.%s: expected shape [%d]", what, cols);
|
|
}
|
|
|
|
} // namespace
|
|
|
|
Tensor *Func::emitLayerNorm(const Tensor *x, const Tensor *scale, const Tensor *bias, float eps) {
|
|
int ix = ensureNode(x);
|
|
const auto &src = graph_.node(ix);
|
|
if (src.rank < 1) throw eve::Exception("Func.layernorm: rank >= 1 required");
|
|
const int cols = src.dims[src.rank - 1];
|
|
auto n = makeShapeNode(OpType::LayerNorm, src.dims, src.rank);
|
|
n.in0 = ix;
|
|
n.s0 = eps;
|
|
if (scale) {
|
|
int is = ensureNode(scale);
|
|
checkParamVector(graph_.node(is), cols, "layernorm scale");
|
|
n.in1 = is;
|
|
}
|
|
if (bias) {
|
|
int ib = ensureNode(bias);
|
|
checkParamVector(graph_.node(ib), cols, "layernorm bias");
|
|
n.in2 = ib;
|
|
}
|
|
int id = graph_.addNode(std::move(n));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
|
|
Tensor *Func::emitRMSNorm(const Tensor *x, const Tensor *scale, float eps) {
|
|
int ix = ensureNode(x);
|
|
const auto &src = graph_.node(ix);
|
|
if (src.rank < 1) throw eve::Exception("Func.rmsnorm: rank >= 1 required");
|
|
const int cols = src.dims[src.rank - 1];
|
|
auto n = makeShapeNode(OpType::RMSNorm, src.dims, src.rank);
|
|
n.in0 = ix;
|
|
n.s0 = eps;
|
|
if (scale) {
|
|
int is = ensureNode(scale);
|
|
checkParamVector(graph_.node(is), cols, "rmsnorm scale");
|
|
n.in1 = is;
|
|
}
|
|
int id = graph_.addNode(std::move(n));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
|
|
Tensor *Func::emitConv1d(const Tensor *x, const Tensor *w, const Tensor *bias, int stride,
|
|
int pad) {
|
|
int ix = ensureNode(x);
|
|
int iw = ensureNode(w);
|
|
const auto &nx = graph_.node(ix);
|
|
const auto &nw = graph_.node(iw);
|
|
if (nx.rank != 3 || nw.rank != 3) throw eve::Exception("Func.conv1d: rank 3 required");
|
|
if (nx.dims[1] != nw.dims[1]) throw eve::Exception("Func.conv1d: channel mismatch");
|
|
const int OL = convOutSize(nx.dims[2], nw.dims[2], stride, pad);
|
|
int od[] = {nx.dims[0], nw.dims[0], OL};
|
|
auto n = makeShapeNode(OpType::Conv1d, od, 3);
|
|
n.in0 = ix;
|
|
n.in1 = iw;
|
|
n.i0 = stride;
|
|
n.i1 = pad;
|
|
if (bias) {
|
|
int ib = ensureNode(bias);
|
|
const auto &nb = graph_.node(ib);
|
|
if (nb.rank != 1 || nb.dims[0] != n.dims[1]) throw eve::Exception("Func.conv1d: bias shape mismatch");
|
|
n.in2 = ib;
|
|
}
|
|
int id = graph_.addNode(std::move(n));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
|
|
Tensor *Func::emitConv2d(const Tensor *x, const Tensor *w, const Tensor *bias, int stride,
|
|
int pad) {
|
|
int ix = ensureNode(x);
|
|
int iw = ensureNode(w);
|
|
const auto &nx = graph_.node(ix);
|
|
const auto &nw = graph_.node(iw);
|
|
if (nx.rank != 4 || nw.rank != 4) throw eve::Exception("Func.conv2d: rank 4 required");
|
|
if (nx.dims[1] != nw.dims[1]) throw eve::Exception("Func.conv2d: channel mismatch");
|
|
const int OH = convOutSize(nx.dims[2], nw.dims[2], stride, pad);
|
|
const int OW = convOutSize(nx.dims[3], nw.dims[3], stride, pad);
|
|
int od[] = {nx.dims[0], nw.dims[0], OH, OW};
|
|
auto n = makeShapeNode(OpType::Conv2d, od, 4);
|
|
n.in0 = ix;
|
|
n.in1 = iw;
|
|
n.i0 = stride;
|
|
n.i1 = pad;
|
|
if (bias) {
|
|
int ib = ensureNode(bias);
|
|
const auto &nb = graph_.node(ib);
|
|
if (nb.rank != 1 || nb.dims[0] != n.dims[1]) throw eve::Exception("Func.conv2d: bias shape mismatch");
|
|
n.in2 = ib;
|
|
}
|
|
int id = graph_.addNode(std::move(n));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
|
|
Tensor *Func::emitPool(OpType type, const Tensor *x, int ksize, int stride, int pad) {
|
|
int ix = ensureNode(x);
|
|
const auto &src = graph_.node(ix);
|
|
if (src.rank != 4) throw eve::Exception("Func.pool: rank 4 required");
|
|
const int OH = convOutSize(src.dims[2], ksize, stride, pad);
|
|
const int OW = convOutSize(src.dims[3], ksize, stride, pad);
|
|
int od[] = {src.dims[0], src.dims[1], OH, OW};
|
|
auto n = makeShapeNode(type, od, 4);
|
|
n.in0 = ix;
|
|
n.i0 = ksize;
|
|
n.i1 = stride;
|
|
n.i2 = pad;
|
|
int id = graph_.addNode(std::move(n));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
|
|
Tensor *Func::emitEmbedding(const Tensor *table, const Tensor *indices) {
|
|
int it = ensureNode(table);
|
|
int ii = ensureNode(indices);
|
|
const auto &nt = graph_.node(it);
|
|
const auto &ni = graph_.node(ii);
|
|
if (nt.rank != 2) throw eve::Exception("Func.embedding: table rank 2 required");
|
|
if (ni.rank < 1 || ni.rank + 1 > Tensor::kMaxRank)
|
|
throw eve::Exception("Func.embedding: index rank out of range");
|
|
int od[Tensor::kMaxRank] = {};
|
|
for (int k = 0; k < ni.rank; ++k) od[k] = ni.dims[k];
|
|
od[ni.rank] = nt.dims[1];
|
|
auto n = makeShapeNode(OpType::Embedding, od, ni.rank + 1);
|
|
n.in0 = it;
|
|
n.in1 = ii;
|
|
int id = graph_.addNode(std::move(n));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
|
|
Tensor *Func::emitConcat(const Tensor *const *ins, int n, int axis) {
|
|
if (!ins || n < 2 || n > 4) throw eve::Exception("Func.concat: 2..4 inputs required");
|
|
int ids[4] = {};
|
|
const GraphNode *ns[4] = {};
|
|
// Importing an eager input may reallocate the graph's node array.
|
|
for (int k = 0; k < n; ++k) ids[k] = ensureNode(ins[k]);
|
|
for (int k = 0; k < n; ++k) {
|
|
ns[k] = &graph_.node(ids[k]);
|
|
if (ns[k]->rank != ns[0]->rank)
|
|
throw eve::Exception("Func.concat: rank mismatch");
|
|
}
|
|
axis = normalizeAxisChecked(axis, ns[0]->rank);
|
|
int od[Tensor::kMaxRank] = {};
|
|
for (int k = 0; k < ns[0]->rank; ++k) {
|
|
if (k == axis) {
|
|
int total = 0;
|
|
for (int t = 0; t < n; ++t) total += ns[t]->dims[k];
|
|
od[k] = total;
|
|
} else {
|
|
od[k] = ns[0]->dims[k];
|
|
for (int t = 1; t < n; ++t)
|
|
if (ns[t]->dims[k] != od[k])
|
|
throw eve::Exception("Func.concat: dims mismatch on axis %d", k);
|
|
}
|
|
}
|
|
auto g = makeShapeNode(OpType::Concat, od, ns[0]->rank);
|
|
g.in0 = ids[0];
|
|
g.in1 = ids[1];
|
|
if (n > 2) g.in2 = ids[2];
|
|
if (n > 3) g.in3 = ids[3];
|
|
g.i0 = axis;
|
|
int id = graph_.addNode(std::move(g));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
|
|
Tensor *Func::emitSlice(const Tensor *x, int axis, int begin, int end) {
|
|
int ix = ensureNode(x);
|
|
const auto &src = graph_.node(ix);
|
|
axis = normalizeAxisChecked(axis, src.rank);
|
|
if (begin < 0 || end < begin || end > src.dims[axis])
|
|
throw eve::Exception("Func.slice: range out of bounds");
|
|
int od[Tensor::kMaxRank] = {};
|
|
for (int k = 0; k < src.rank; ++k) od[k] = src.dims[k];
|
|
od[axis] = end - begin;
|
|
auto n = makeShapeNode(OpType::Slice, od, src.rank);
|
|
n.in0 = ix;
|
|
n.i0 = axis;
|
|
n.i1 = begin;
|
|
n.i2 = end;
|
|
int id = graph_.addNode(std::move(n));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
|
|
Tensor *Func::emitReduce(OpType type, const Tensor *x, int axis, bool keepDims) {
|
|
int ix = ensureNode(x);
|
|
const auto &src = graph_.node(ix);
|
|
axis = normalizeAxisChecked(axis, src.rank);
|
|
int od[Tensor::kMaxRank] = {};
|
|
int orank = src.rank;
|
|
for (int k = 0; k < src.rank; ++k) od[k] = src.dims[k];
|
|
if (keepDims) {
|
|
od[axis] = 1;
|
|
} else {
|
|
if (orank == 1) {
|
|
od[0] = 1; // Tensor rank must stay >= 1
|
|
} else {
|
|
for (int k = axis; k < orank - 1; ++k) od[k] = od[k + 1];
|
|
od[orank - 1] = 0;
|
|
--orank;
|
|
}
|
|
}
|
|
auto n = makeShapeNode(type, od, orank);
|
|
n.in0 = ix;
|
|
n.i0 = axis;
|
|
n.i1 = keepDims ? 1 : 0;
|
|
int id = graph_.addNode(std::move(n));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
|
|
Tensor *Func::emitArgMax(const Tensor *x, int axis, bool keepDims) {
|
|
int ix = ensureNode(x);
|
|
const auto &src = graph_.node(ix);
|
|
axis = normalizeAxisChecked(axis, src.rank);
|
|
int od[Tensor::kMaxRank] = {};
|
|
int orank = src.rank;
|
|
for (int k = 0; k < src.rank; ++k) od[k] = src.dims[k];
|
|
if (keepDims) {
|
|
od[axis] = 1;
|
|
} else {
|
|
if (orank == 1) {
|
|
od[0] = 1;
|
|
} else {
|
|
for (int k = axis; k < orank - 1; ++k) od[k] = od[k + 1];
|
|
od[orank - 1] = 0;
|
|
--orank;
|
|
}
|
|
}
|
|
auto n = makeShapeNode(OpType::ArgMax, od, orank);
|
|
n.in0 = ix;
|
|
n.i0 = axis;
|
|
n.i1 = keepDims ? 1 : 0;
|
|
n.dtype = static_cast<int>(DType::Int32);
|
|
int id = graph_.addNode(std::move(n));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
|
|
Tensor *Func::emitCast(const Tensor *x, DType dtype) {
|
|
int ix = ensureNode(x);
|
|
const auto &src = graph_.node(ix);
|
|
auto n = makeShapeNode(OpType::Cast, src.dims, src.rank);
|
|
n.in0 = ix;
|
|
n.dtype = static_cast<int>(dtype);
|
|
int id = graph_.addNode(std::move(n));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
|
|
Tensor *Func::emitSdpa(const Tensor *q, const Tensor *k, const Tensor *v, const Tensor *mask,
|
|
float scale) {
|
|
int iq = ensureNode(q);
|
|
int ik = ensureNode(k);
|
|
int iv = ensureNode(v);
|
|
const int im = mask ? ensureNode(mask) : -1;
|
|
const auto &nq = graph_.node(iq);
|
|
const auto &nk = graph_.node(ik);
|
|
const auto &nv = graph_.node(iv);
|
|
if (nq.rank != 4 || nk.rank != 4 || nv.rank != 4)
|
|
throw eve::Exception("Func.sdpa: rank 4 required");
|
|
if (nq.dims[0] != nk.dims[0] || nq.dims[1] != nk.dims[1] ||
|
|
nq.dims[3] != nk.dims[3] || nk.dims[2] != nv.dims[2] || nq.dims[3] != nv.dims[3])
|
|
throw eve::Exception("Func.sdpa: q/k/v shape mismatch");
|
|
auto n = makeShapeNode(OpType::ScaledDotProductAttention, nq.dims, 4);
|
|
n.in0 = iq;
|
|
n.in1 = ik;
|
|
n.in2 = iv;
|
|
n.s0 = scale;
|
|
if (mask) {
|
|
const auto &nm = graph_.node(im);
|
|
if (nm.rank != 4 || nm.dims[0] != nq.dims[0] || nm.dims[1] != nq.dims[1] ||
|
|
nm.dims[2] != nq.dims[2] || nm.dims[3] != nk.dims[2])
|
|
throw eve::Exception("Func.sdpa: mask shape mismatch");
|
|
n.in3 = im;
|
|
}
|
|
int id = graph_.addNode(std::move(n));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
|
|
Tensor *Func::emitResize2d(const Tensor *x, int outH, int outW, int mode) {
|
|
int ix = ensureNode(x);
|
|
const auto &src = graph_.node(ix);
|
|
if (src.rank != 4) throw eve::Exception("Func.resize2d: rank 4 required");
|
|
if (outH <= 0 || outW <= 0) throw eve::Exception("Func.resize2d: bad output size");
|
|
int od[] = {src.dims[0], src.dims[1], outH, outW};
|
|
auto n = makeShapeNode(OpType::Resize2d, od, 4);
|
|
n.in0 = ix;
|
|
n.i0 = mode;
|
|
int id = graph_.addNode(std::move(n));
|
|
return makeSymbolicFromNode(id);
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// CompiledFunction
|
|
// ---------------------------------------------------------------------------
|
|
|
|
CompiledFunction::CompiledFunction() = default;
|
|
CompiledFunction::~CompiledFunction() = default;
|
|
|
|
CompiledFunction *Func::compile() {
|
|
if (outputNode_ < 0) throw eve::Exception("Func.compile: setOutput required");
|
|
tracing_ = false;
|
|
if (owner_) owner_->popTrace(this);
|
|
return CompiledFunction::fromFunc(this);
|
|
}
|
|
|
|
CompiledFunction *CompiledFunction::fromFunc(Func *fn) {
|
|
if (!fn) throw eve::Exception("CompiledFunction: null func");
|
|
auto *cf = new CompiledFunction();
|
|
cf->graph_ = fn->graph(); // copy
|
|
cf->outputNode_ = fn->outputNode();
|
|
cf->placeholderCount_ = fn->placeholderCount();
|
|
cf->device_ = "cpu";
|
|
cf->optimized_ = std::make_unique<OptimizedGraph>(optimizeGraph(cf->graph_, cf->outputNode_));
|
|
cf->order_ = cf->optimized_->order;
|
|
|
|
// Best-effort: run on GPU via eve::gpgpu compute shaders when a Vulkan
|
|
// device is available. Falls back to the CPU interpreter below otherwise.
|
|
try {
|
|
cf->gpuProgram_.reset(GpuProgram::tryBuild(cf->graph_, *cf->optimized_, cf->outputNode_));
|
|
} catch (...) {
|
|
cf->gpuProgram_.reset();
|
|
}
|
|
if (cf->gpuProgram_) cf->device_ = "gpu";
|
|
return cf;
|
|
}
|
|
|
|
Tensor *CompiledFunction::run0() { return runWithFeeds(nullptr, 0); }
|
|
|
|
Tensor *CompiledFunction::run1(Tensor *in0) {
|
|
Tensor *feeds[] = {in0};
|
|
return runWithFeeds(feeds, 1);
|
|
}
|
|
|
|
Tensor *CompiledFunction::run2(Tensor *in0, Tensor *in1) {
|
|
Tensor *feeds[] = {in0, in1};
|
|
return runWithFeeds(feeds, 2);
|
|
}
|
|
|
|
Tensor *CompiledFunction::run3(Tensor *in0, Tensor *in1, Tensor *in2) {
|
|
Tensor *feeds[] = {in0, in1, in2};
|
|
return runWithFeeds(feeds, 3);
|
|
}
|
|
|
|
Tensor *CompiledFunction::run4(Tensor *in0, Tensor *in1, Tensor *in2, Tensor *in3) {
|
|
Tensor *feeds[] = {in0, in1, in2, in3};
|
|
return runWithFeeds(feeds, 4);
|
|
}
|
|
|
|
Tensor *CompiledFunction::run5(Tensor *in0, Tensor *in1, Tensor *in2, Tensor *in3, Tensor *in4) {
|
|
Tensor *feeds[] = {in0, in1, in2, in3, in4};
|
|
return runWithFeeds(feeds, 5);
|
|
}
|
|
|
|
Tensor *CompiledFunction::run6(Tensor *in0, Tensor *in1, Tensor *in2, Tensor *in3, Tensor *in4,
|
|
Tensor *in5) {
|
|
Tensor *feeds[] = {in0, in1, in2, in3, in4, in5};
|
|
return runWithFeeds(feeds, 6);
|
|
}
|
|
|
|
Tensor *CompiledFunction::runWithFeeds(Tensor *const *feeds, int nFeeds) {
|
|
if (nFeeds != placeholderCount_)
|
|
throw eve::Exception("CompiledFunction.run: expected %d feeds, got %d", placeholderCount_,
|
|
nFeeds);
|
|
for (int i = 0; i < nFeeds; ++i) {
|
|
if (!feeds[i]) throw eve::Exception("CompiledFunction.run: null feed");
|
|
feeds[i]->ensureEager("run");
|
|
}
|
|
|
|
const int n = graph_.nodeCount();
|
|
|
|
// Validate placeholder shapes (shared by the GPU and CPU execution paths).
|
|
for (int i = 0; i < n; ++i) {
|
|
const auto &nd = graph_.node(i);
|
|
if (nd.type != OpType::Placeholder) continue;
|
|
const int slot = nd.placeholderSlot;
|
|
if (slot < 0 || slot >= nFeeds) throw eve::Exception("CompiledFunction: bad placeholder slot");
|
|
Tensor *feed = feeds[slot];
|
|
if (feed->getRank() != nd.rank || feed->getSize() != nd.size)
|
|
throw eve::Exception("CompiledFunction: feed shape mismatch");
|
|
for (int a = 0; a < nd.rank; ++a)
|
|
if (feed->getDim(a) != nd.dims[a])
|
|
throw eve::Exception("CompiledFunction: feed shape mismatch");
|
|
}
|
|
|
|
const auto &outN = graph_.node(outputNode_);
|
|
|
|
if (gpuProgram_) {
|
|
std::vector<const float *> ptrs(static_cast<size_t>(nFeeds));
|
|
for (int i = 0; i < nFeeds; ++i) ptrs[static_cast<size_t>(i)] = feeds[i]->data();
|
|
std::vector<float> result = gpuProgram_->run(ptrs);
|
|
auto *out = new Tensor(static_cast<DType>(outN.dtype), outN.dims, outN.rank);
|
|
if (int(result.size()) != out->getSize())
|
|
throw eve::Exception("CompiledFunction: gpu output size mismatch");
|
|
std::memcpy(out->data(), result.data(), sizeof(float) * static_cast<size_t>(out->getSize()));
|
|
return out;
|
|
}
|
|
|
|
std::vector<std::vector<float>> bufs(static_cast<size_t>(n));
|
|
for (int i = 0; i < n; ++i) {
|
|
const auto &nd = graph_.node(i);
|
|
if (nd.type != OpType::Placeholder) continue;
|
|
Tensor *feed = feeds[nd.placeholderSlot];
|
|
bufs[static_cast<size_t>(i)].assign(feed->data(), feed->data() + feed->getSize());
|
|
}
|
|
|
|
for (int nodeId : order_) executeNode(nodeId, bufs);
|
|
|
|
auto *out = new Tensor(static_cast<DType>(outN.dtype), outN.dims, outN.rank);
|
|
const auto &src = bufs[static_cast<size_t>(outputNode_)];
|
|
if (int(src.size()) != out->getSize())
|
|
throw eve::Exception("CompiledFunction: output size mismatch");
|
|
std::memcpy(out->data(), src.data(), sizeof(float) * static_cast<size_t>(out->getSize()));
|
|
return out;
|
|
}
|
|
|
|
void CompiledFunction::executeNode(int nodeId, std::vector<std::vector<float>> &bufs) const {
|
|
const auto &nd = graph_.node(nodeId);
|
|
auto &out = bufs[static_cast<size_t>(nodeId)];
|
|
const auto in = [&](int id) -> const std::vector<float> & {
|
|
return bufs[static_cast<size_t>(id)];
|
|
};
|
|
|
|
switch (nd.type) {
|
|
case OpType::Placeholder:
|
|
return; // already filled
|
|
case OpType::Const:
|
|
if (!nd.constBytes.empty()) {
|
|
out.assign(static_cast<size_t>(nd.size), 0.f);
|
|
q::dequantizeAll(static_cast<DType>(nd.dtype), nd.constBytes.data(),
|
|
nd.constScales.data(), nd.qGroup, nd.size, out.data());
|
|
} else {
|
|
out = nd.constData;
|
|
}
|
|
return;
|
|
case OpType::Add:
|
|
case OpType::Sub:
|
|
case OpType::Multiply:
|
|
case OpType::Divide: {
|
|
const auto &na = graph_.node(nd.in0);
|
|
const auto &nb = graph_.node(nd.in1);
|
|
out.resize(static_cast<size_t>(nd.size));
|
|
kernels::binaryOp(nd.type, in(nd.in0).data(), na.dims, na.rank, in(nd.in1).data(),
|
|
nb.dims, nb.rank, out.data(), nd.dims, nd.rank);
|
|
return;
|
|
}
|
|
case OpType::Neg:
|
|
case OpType::Abs:
|
|
case OpType::Sqrt:
|
|
case OpType::Exp:
|
|
case OpType::Log:
|
|
case OpType::Sin:
|
|
case OpType::Cos:
|
|
case OpType::Tanh:
|
|
case OpType::Relu:
|
|
case OpType::Sigmoid:
|
|
case OpType::Gelu:
|
|
case OpType::Silu:
|
|
case OpType::AddScalar:
|
|
case OpType::SubScalar:
|
|
case OpType::MulScalar:
|
|
case OpType::DivScalar:
|
|
case OpType::PowScalar:
|
|
case OpType::Clamp:
|
|
case OpType::MaximumScalar:
|
|
case OpType::MinimumScalar: {
|
|
out.resize(static_cast<size_t>(nd.size));
|
|
kernels::unaryOp(nd.type, in(nd.in0).data(), nd.size, out.data(), nd.s0, nd.s1);
|
|
return;
|
|
}
|
|
case OpType::Where: {
|
|
const auto &c = in(nd.in0);
|
|
const auto &a = in(nd.in1);
|
|
const auto &b = in(nd.in2);
|
|
out.resize(static_cast<size_t>(nd.size));
|
|
for (int i = 0; i < nd.size; ++i)
|
|
out[static_cast<size_t>(i)] =
|
|
c[static_cast<size_t>(i)] > 0.5f ? a[static_cast<size_t>(i)]
|
|
: b[static_cast<size_t>(i)];
|
|
return;
|
|
}
|
|
case OpType::MatMul: {
|
|
const auto &A = graph_.node(nd.in0);
|
|
const auto &B = graph_.node(nd.in1);
|
|
const auto &a = in(nd.in0);
|
|
const auto &b = in(nd.in1);
|
|
out.assign(static_cast<size_t>(nd.size), 0.f);
|
|
if (nd.rank == 2) {
|
|
const int m = A.dims[0], k = A.dims[1], n = B.dims[1];
|
|
for (int i = 0; i < m; ++i) {
|
|
for (int j = 0; j < n; ++j) {
|
|
double acc = 0.0;
|
|
for (int t = 0; t < k; ++t)
|
|
acc += double(a[static_cast<size_t>(i * k + t)]) *
|
|
double(b[static_cast<size_t>(t * n + j)]);
|
|
out[static_cast<size_t>(i * n + j)] = float(acc);
|
|
}
|
|
}
|
|
} else {
|
|
const int batch = A.dims[0], m = A.dims[1], k = A.dims[2], n = B.dims[2];
|
|
for (int bb = 0; bb < batch; ++bb) {
|
|
const float *ap = a.data() + size_t(bb) * m * k;
|
|
const float *bp = b.data() + size_t(bb) * k * n;
|
|
float *cp = out.data() + size_t(bb) * m * n;
|
|
for (int i = 0; i < m; ++i)
|
|
for (int j = 0; j < n; ++j) {
|
|
double acc = 0.0;
|
|
for (int t = 0; t < k; ++t) acc += double(ap[i * k + t]) * double(bp[t * n + j]);
|
|
cp[i * n + j] = float(acc);
|
|
}
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
case OpType::Transpose:
|
|
case OpType::Permute: {
|
|
const auto &X = graph_.node(nd.in0);
|
|
out.resize(static_cast<size_t>(nd.size));
|
|
int order[Tensor::kMaxRank] = {};
|
|
for (int k = 0; k < nd.rank; ++k)
|
|
order[k] = nd.type == OpType::Transpose ? (k == 0 ? 1 : 0) : nd.perm[k];
|
|
kernels::permute(in(nd.in0).data(), X.dims, nd.rank, order, out.data(), nd.dims);
|
|
return;
|
|
}
|
|
case OpType::Reshape:
|
|
case OpType::Flatten:
|
|
case OpType::Cast:
|
|
out = in(nd.in0);
|
|
return;
|
|
case OpType::Softmax:
|
|
case OpType::LogSoftmax: {
|
|
out.resize(static_cast<size_t>(nd.size));
|
|
kernels::softmax(in(nd.in0).data(), nd.dims, nd.rank, nd.i0,
|
|
nd.type == OpType::LogSoftmax, out.data());
|
|
return;
|
|
}
|
|
case OpType::LayerNorm: {
|
|
out.resize(static_cast<size_t>(nd.size));
|
|
const int cols = nd.dims[nd.rank - 1];
|
|
const int rows = nd.size / cols;
|
|
kernels::layernorm(in(nd.in0).data(), rows, cols,
|
|
nd.in1 >= 0 ? in(nd.in1).data() : nullptr,
|
|
nd.in2 >= 0 ? in(nd.in2).data() : nullptr, nd.s0, out.data());
|
|
return;
|
|
}
|
|
case OpType::RMSNorm: {
|
|
out.resize(static_cast<size_t>(nd.size));
|
|
const int cols = nd.dims[nd.rank - 1];
|
|
const int rows = nd.size / cols;
|
|
kernels::rmsnorm(in(nd.in0).data(), rows, cols,
|
|
nd.in1 >= 0 ? in(nd.in1).data() : nullptr, nd.s0, out.data());
|
|
return;
|
|
}
|
|
case OpType::Conv1d:
|
|
case OpType::Conv2d: {
|
|
out.resize(static_cast<size_t>(nd.size));
|
|
const auto &X = graph_.node(nd.in0);
|
|
const auto &W = graph_.node(nd.in1);
|
|
if (nd.type == OpType::Conv1d)
|
|
kernels::conv1d(in(nd.in0).data(), X.dims, in(nd.in1).data(), W.dims,
|
|
nd.in2 >= 0 ? in(nd.in2).data() : nullptr, nd.i0, nd.i1, out.data());
|
|
else
|
|
kernels::conv2d(in(nd.in0).data(), X.dims, in(nd.in1).data(), W.dims,
|
|
nd.in2 >= 0 ? in(nd.in2).data() : nullptr, nd.i0, nd.i1, out.data());
|
|
return;
|
|
}
|
|
case OpType::MaxPool2d:
|
|
case OpType::AvgPool2d: {
|
|
out.resize(static_cast<size_t>(nd.size));
|
|
const auto &X = graph_.node(nd.in0);
|
|
if (nd.type == OpType::MaxPool2d)
|
|
kernels::maxpool2d(in(nd.in0).data(), X.dims, nd.i0, nd.i1, nd.i2, out.data());
|
|
else
|
|
kernels::avgpool2d(in(nd.in0).data(), X.dims, nd.i0, nd.i1, nd.i2, out.data());
|
|
return;
|
|
}
|
|
case OpType::Embedding: {
|
|
out.resize(static_cast<size_t>(nd.size));
|
|
const auto &T = graph_.node(nd.in0);
|
|
const auto &I = graph_.node(nd.in1);
|
|
kernels::embedding(in(nd.in0).data(), T.dims[0], T.dims[1], in(nd.in1).data(),
|
|
I.size, out.data());
|
|
return;
|
|
}
|
|
case OpType::Concat: {
|
|
out.resize(static_cast<size_t>(nd.size));
|
|
const float *ins[4] = {};
|
|
const int *dims[4] = {};
|
|
int ranks[4] = {};
|
|
int n = 2;
|
|
if (nd.in3 >= 0) n = 4;
|
|
else if (nd.in2 >= 0) n = 3;
|
|
const int ids[4] = {nd.in0, nd.in1, nd.in2, nd.in3};
|
|
for (int k = 0; k < n; ++k) {
|
|
ins[k] = in(ids[k]).data();
|
|
dims[k] = graph_.node(ids[k]).dims;
|
|
ranks[k] = graph_.node(ids[k]).rank;
|
|
}
|
|
kernels::concat(ins, dims, ranks, n, nd.i0, out.data(), nd.dims, nd.rank);
|
|
return;
|
|
}
|
|
case OpType::Slice: {
|
|
out.resize(static_cast<size_t>(nd.size));
|
|
const auto &X = graph_.node(nd.in0);
|
|
kernels::sliceOp(in(nd.in0).data(), X.dims, X.rank, nd.i0, nd.i1, nd.i2, out.data(),
|
|
nd.dims, nd.rank);
|
|
return;
|
|
}
|
|
case OpType::ReduceSum:
|
|
case OpType::ReduceMean:
|
|
case OpType::ReduceMin:
|
|
case OpType::ReduceMax: {
|
|
out.resize(static_cast<size_t>(nd.size));
|
|
const auto &X = graph_.node(nd.in0);
|
|
kernels::reduceAxis(nd.type, in(nd.in0).data(), X.dims, X.rank, nd.i0, out.data(),
|
|
nd.dims, nd.rank);
|
|
return;
|
|
}
|
|
case OpType::ArgMax: {
|
|
out.resize(static_cast<size_t>(nd.size));
|
|
const auto &X = graph_.node(nd.in0);
|
|
kernels::argmax(in(nd.in0).data(), X.dims, X.rank, nd.i0, out.data(), nd.dims,
|
|
nd.rank);
|
|
return;
|
|
}
|
|
case OpType::ScaledDotProductAttention: {
|
|
out.resize(static_cast<size_t>(nd.size));
|
|
const auto &Q = graph_.node(nd.in0);
|
|
const auto &K = graph_.node(nd.in1);
|
|
const int B = Q.dims[0], H = Q.dims[1], T = Q.dims[2], D = Q.dims[3];
|
|
const int S = K.dims[2];
|
|
kernels::sdpa(in(nd.in0).data(), in(nd.in1).data(), in(nd.in2).data(),
|
|
nd.in3 >= 0 ? in(nd.in3).data() : nullptr, B, H, T, S, D, nd.s0,
|
|
out.data());
|
|
return;
|
|
}
|
|
case OpType::Resize2d: {
|
|
out.resize(static_cast<size_t>(nd.size));
|
|
const auto &X = graph_.node(nd.in0);
|
|
kernels::resize2d(in(nd.in0).data(), X.dims, nd.dims[3], nd.dims[2], nd.i0,
|
|
out.data());
|
|
return;
|
|
}
|
|
}
|
|
}
|
|
|
|
} // namespace eve::tensor
|