okooooooo/EVEngine
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
42d9be4eedc75c47cebcfc3e12525425083fcbb2
EVEngine/src/modules/tensor/KernelGenWgslReduction.cpp
T

216 lines
9.2 KiB
C++
Raw Blame History

#include <algorithm>
#include <cmath>
#include <sstream>
#include "common/Exception.h"
#include "tensor/KernelGenWgslInternal.h"
#include "tensor/Quant.h"
namespace eve::tensor::wgsl_detail {
void genSoftmax(const Graph &g, const FusedGroup &grp, KernelSpec &out) {
const GraphNode &sn = g.node(grp.outputNode);
const GraphNode &X = g.node(sn.in0);
const int axis = sn.i0;
int outer = 1, reduce = 1, inner = 1;
for (int k = 0; k < axis; ++k) outer *= X.dims[k];
reduce = X.dims[axis];
for (int k = axis + 1; k < X.rank; ++k) inner *= X.dims[k];
const int rows = outer * inner;
const bool logMode = grp.logMode;
std::ostringstream os1, os2;
os1 << header(kLocalSize);
os1 << bufferDecl(0, "in_");
os1 << bufferDecl(1, "mx");
os1 << bufferDecl(2, "sm");
os1 << pushConstant();
os1 << "@compute @workgroup_size(workgroupX, workgroupY)\nfn main(@builtin(global_invocation_id) globalId: "
"vec3<u32>, @builtin(local_invocation_id) localId: vec3<u32>, @builtin(workgroup_id) groupId: vec3<u32>) "
"{\n";
os1 << " var i_: u32 = globalId.x;\n";
os1 << " if (i_ >= " << rows << "u) { return; }\n";
os1 << " var o_: u32 = i_ / " << inner << "u;\n";
os1 << " var ii: u32 = i_ % " << inner << "u;\n";
os1 << " var m: f32 = -3.402823e38;\n";
os1 << " for (var j: u32 = 0u; j < " << reduce << "u; j++) {\n";
os1 << " m = max(m, in_[(o_ * " << reduce << "u + j) * " << inner << "u + ii]);\n";
os1 << " }\n";
os1 << " var s: f32 = 0.0;\n";
os1 << " for (var j: u32 = 0u; j < " << reduce << "u; j++) {\n";
os1 << " s += exp(in_[(o_ * " << reduce << "u + j) * " << inner << "u + ii] - m);\n";
os1 << " }\n";
os1 << " mx[i_] = m;\n";
os1 << " sm[i_] = s;\n";
os1 << "}\n";
os2 << header(kLocalSize);
os2 << bufferDecl(0, "in_");
os2 << bufferDecl(1, "mx");
os2 << bufferDecl(2, "sm");
os2 << bufferDecl(3, "o");
os2 << pushConstant();
os2 << "@compute @workgroup_size(workgroupX, workgroupY)\nfn main(@builtin(global_invocation_id) globalId: "
"vec3<u32>, @builtin(local_invocation_id) localId: vec3<u32>, @builtin(workgroup_id) groupId: vec3<u32>) "
"{\n";
os2 << " var i_: u32 = globalId.x;\n";
os2 << " if (i_ >= " << sn.size << "u) { return; }\n";
os2 << " var o_: u32 = (i_ / (" << reduce * inner << "u)) * " << inner << "u + (i_ % " << inner << "u);\n";
os2 << " var m: f32 = mx[o_];\n";
os2 << " var s: f32 = sm[o_];\n";
os2 << " var x: f32 = in_[i_];\n";
if (logMode) {
os2 << " o[i_] = (x - m) - log(s);\n";
} else {
os2 << " o[i_] = exp(x - m) / s;\n";
}
os2 << "}\n";
out.pass1 = os1.str();
out.pass2 = os2.str();
out.groupsX1 = groupsFor(rows);
out.groupsX2 = groupsFor(sn.size);
out.inputCount = 1;
out.inputsReadPass1 = 1;
out.statsCount = 2;
out.statsSize = rows;
out.twoPass = true;
return;
}
void genNorm(const Graph &g, const FusedGroup &grp, bool rms, KernelSpec &out) {
const GraphNode &nn = g.node(grp.outputNode);
const GraphNode &X = g.node(nn.in0);
const int cols = X.dims[X.rank - 1];
const int rows = X.size / cols;
const float eps = nn.s0;
const bool hasScale = grp.hasScale;
const bool hasBias = !rms && grp.hasBias;
const int inputCount = 1 + (hasScale ? 1 : 0) + (hasBias ? 1 : 0);
const int statsCount = rms ? 1 : 2;
const int outBinding = inputCount + statsCount;
std::ostringstream os1, os2;
os1 << header(kLocalSize);
os1 << bufferDecl(0, "in_");
os1 << bufferDecl(1, "st0");
if (statsCount > 1) os1 << bufferDecl(2, "st1");
os1 << pushConstant();
os1 << "@compute @workgroup_size(workgroupX, workgroupY)\nfn main(@builtin(global_invocation_id) globalId: "
"vec3<u32>, @builtin(local_invocation_id) localId: vec3<u32>, @builtin(workgroup_id) groupId: vec3<u32>) "
"{\n";
os1 << " var i_: u32 = globalId.x;\n";
os1 << " if (i_ >= " << rows << "u) { return; }\n";
os1 << " var s0: f32 = 0.0;\n";
if (statsCount > 1) os1 << " var s1: f32 = 0.0;\n";
os1 << " for (var j: u32 = 0u; j < " << cols << "u; j++) {\n";
os1 << " var v: f32 = in_[i_ * " << cols << "u + j];\n";
if (rms) {
os1 << " s0 += v * v;\n";
} else {
os1 << " s0 += v; s1 += v * v;\n";
}
os1 << " }\n";
os1 << " st0[i_] = s0;\n";
if (statsCount > 1) os1 << " st1[i_] = s1;\n";
os1 << "}\n";
os2 << header(kLocalSize);
for (int k = 0; k < inputCount; ++k) os2 << bufferDecl(k, ("a" + std::to_string(k)).c_str());
os2 << bufferDecl(inputCount, "st0");
if (statsCount > 1) os2 << bufferDecl(inputCount + 1, "st1");
os2 << bufferDecl(outBinding, "o");
os2 << pushConstant();
os2 << "@compute @workgroup_size(workgroupX, workgroupY)\nfn main(@builtin(global_invocation_id) globalId: "
"vec3<u32>, @builtin(local_invocation_id) localId: vec3<u32>, @builtin(workgroup_id) groupId: vec3<u32>) "
"{\n";
os2 << " var i_: u32 = globalId.x;\n";
os2 << " if (i_ >= " << nn.size << "u) { return; }\n";
os2 << " var r: u32 = i_ / " << cols << "u;\n";
os2 << " var c: u32 = i_ % " << cols << "u;\n";
if (rms) {
os2 << " var inv: f32 = 1.0 / sqrt(st0[r] / " << cols << ".0 + " << scalarStr(eps) << ");\n";
os2 << " var y: f32 = a0[i_] * inv;\n";
if (hasScale) os2 << " y *= a1[c];\n";
} else {
os2 << " var mean: f32 = st0[r] / " << cols << ".0;\n";
os2 << " var variance: f32 = st1[r] / " << cols << ".0 - mean * mean;\n";
os2 << " variance = max(variance, 0.0);\n";
os2 << " var inv: f32 = 1.0 / sqrt(variance + " << scalarStr(eps) << ");\n";
os2 << " var y: f32 = (a0[i_] - mean) * inv;\n";
if (hasScale) os2 << " y *= a1[c];\n";
if (hasBias) os2 << " y += a" << (hasScale ? 2 : 1) << "[c];\n";
}
os2 << " o[i_] = y;\n";
os2 << "}\n";
out.pass1 = os1.str();
out.pass2 = os2.str();
out.groupsX1 = groupsFor(rows);
out.groupsX2 = groupsFor(nn.size);
out.inputCount = inputCount;
out.inputsReadPass1 = 1;
out.statsCount = statsCount;
out.statsSize = rows;
out.twoPass = true;
return;
}
void genReduceOrArgmax(const Graph &g, const FusedGroup &grp, bool argmax, KernelSpec &out) {
const GraphNode &rn = g.node(grp.outputNode);
const GraphNode &X = g.node(rn.in0);
const int axis = rn.i0;
int outer = 1, reduce = 1, inner = 1;
for (int k = 0; k < axis; ++k) outer *= X.dims[k];
reduce = X.dims[axis];
for (int k = axis + 1; k < X.rank; ++k) inner *= X.dims[k];
const int outSize = outer * inner;
std::ostringstream os;
os << header(kLocalSize);
os << bufferDecl(0, "in_");
os << bufferDecl(1, "o");
os << pushConstant();
os << "@compute @workgroup_size(workgroupX, workgroupY)\nfn main(@builtin(global_invocation_id) globalId: "
"vec3<u32>, @builtin(local_invocation_id) localId: vec3<u32>, @builtin(workgroup_id) groupId: vec3<u32>) {\n";
os << " var i_: u32 = globalId.x;\n";
os << " if (i_ >= " << outSize << "u) { return; }\n";
os << " var o_: u32 = i_ / " << inner << "u;\n";
os << " var ii: u32 = i_ % " << inner << "u;\n";
if (argmax) {
os << " var best: f32 = -3.402823e38;\n";
os << " var bestJ: f32 = 0.0;\n";
os << " for (var j: u32 = 0u; j < " << reduce << "u; j++) {\n";
os << " var v: f32 = in_[(o_ * " << reduce << "u + j) * " << inner << "u + ii];\n";
os << " if (v > best) { best = v; bestJ = f32(j); }\n";
os << " }\n";
os << " o[i_] = bestJ;\n";
} else {
switch (grp.op) {
case OpType::ReduceSum:
case OpType::ReduceMean:
os << " var acc: f32 = 0.0;\n";
os << " for (var j: u32 = 0u; j < " << reduce << "u; j++) { acc += in_[(o_ * " << reduce << "u + j) * "
<< inner << "u + ii]; }\n";
if (grp.op == OpType::ReduceMean) os << " acc /= " << scalarStr(float(reduce)) << ";\n";
break;
case OpType::ReduceMin:
os << " var acc: f32 = 3.402823e38;\n";
os << " for (var j: u32 = 0u; j < " << reduce << "u; j++) { acc = min(acc, in_[(o_ * " << reduce
<< "u + j) * " << inner << "u + ii]); }\n";
break;
case OpType::ReduceMax:
os << " var acc: f32 = -3.402823e38;\n";
os << " for (var j: u32 = 0u; j < " << reduce << "u; j++) { acc = max(acc, in_[(o_ * " << reduce
<< "u + j) * " << inner << "u + ii]); }\n";
break;
default: throw eve::Exception("Tensor WGSL: unsupported kernel variant or binding count");
}
os << " o[i_] = acc;\n";
}
os << "}\n";
out.pass1.clear();
out.pass2 = os.str();
out.groupsX2 = groupsFor(outSize);
out.inputCount = 1;
return;
}
} // namespace eve::tensor::wgsl_detail
Reference in new issue View Git Blame Copy Permalink