242 lines
10 KiB
C++
242 lines
10 KiB
C++
#include "zeroerr/assert.h"
|
|
#include "zeroerr/unittest.h"
|
|
|
|
#include "common/Exception.h"
|
|
#include "tensor/Graph.h"
|
|
#include "tensor/KernelGen.h"
|
|
#include "tensor/Optimizer.h"
|
|
#include "tensor/TF.h"
|
|
#include "tensor/Tensor.h"
|
|
|
|
#include "graphics/AmbientOcclusion.h"
|
|
#include "graphics/AntiAliasing.h"
|
|
#include "graphics/Canvas.h"
|
|
#include "graphics/DrawItem2D.h"
|
|
#include "graphics/Font.h"
|
|
#include "graphics/GBuffer.h"
|
|
#include "graphics/GlobalIllumination.h"
|
|
#include "graphics/Graphics.h"
|
|
#include "graphics/Grass.h"
|
|
#include "graphics/Light.h"
|
|
#include "graphics/Material.h"
|
|
#include "graphics/Mesh.h"
|
|
#include "graphics/Outline.h"
|
|
#include "graphics/Quad.h"
|
|
#include "graphics/RenderControl.h"
|
|
#include "graphics/ScreenSpaceReflection.h"
|
|
#include "graphics/Shader.h"
|
|
#include "graphics/Texture.h"
|
|
#include "graphics/Volumetric.h"
|
|
#include "graphics/Water.h"
|
|
#include "graphics/Waterfall.h"
|
|
#include "window/Window.h"
|
|
|
|
#include <algorithm>
|
|
#include <cmath>
|
|
#include <functional>
|
|
#include <memory>
|
|
|
|
using namespace eve::tensor;
|
|
|
|
namespace {
|
|
|
|
/** GPU tensor tests need a live device; headless init is enough. */
|
|
void initHeadlessGfx(zeroerr::TestContext *_ZEROERR_TEST_CONTEXT) {
|
|
auto *gfx = eve::graphics::Graphics::create();
|
|
REQUIRE(gfx != nullptr);
|
|
gfx->initHeadless(320, 240);
|
|
}
|
|
|
|
} // namespace
|
|
|
|
TEST_CASE("tensor.gpu.fusedElementwiseChain") {
|
|
initHeadlessGfx(_ZEROERR_TEST_CONTEXT);
|
|
auto *tf = TF::create();
|
|
std::unique_ptr<Func> fn(tf->func());
|
|
Tensor *x = fn->input2(4, 4);
|
|
Tensor *y = tf->mulScalar(x, 2.f);
|
|
y = tf->addScalar(y, 1.f);
|
|
y = tf->relu(y);
|
|
y = tf->mulScalar(y, 0.5f);
|
|
fn->setOutput(y);
|
|
std::unique_ptr<CompiledFunction> compiled(fn->compile());
|
|
REQUIRE_EQ(compiled->getDevice(), std::string("gpu"));
|
|
|
|
std::unique_ptr<Tensor> xd(tf->randomNormal2(4, 4));
|
|
std::unique_ptr<Tensor> out(compiled->run1(xd.get()));
|
|
std::unique_ptr<Tensor> ref(xd->mulScalar(2.f));
|
|
ref.reset(ref->addScalar(1.f));
|
|
ref.reset(ref->relu());
|
|
ref.reset(ref->mulScalar(0.5f));
|
|
for (int i = 0; i < out->getSize(); ++i) REQUIRE(std::fabs(out->get(i) - ref->get(i)) < 1e-4f);
|
|
}
|
|
|
|
TEST_CASE("tensor.gpu.softmaxLayernormConv") {
|
|
initHeadlessGfx(_ZEROERR_TEST_CONTEXT);
|
|
auto *tf = TF::create();
|
|
std::unique_ptr<Tensor> w(tf->fill4(2, 1, 3, 3, 0.1f));
|
|
|
|
std::unique_ptr<Func> fn(tf->func());
|
|
Tensor *in = fn->input4(1, 1, 6, 6);
|
|
Tensor *y = tf->conv2d(in, w.get(), 1, 0);
|
|
y = tf->relu(y);
|
|
y = tf->flatten(y);
|
|
y = tf->layernorm(y, 1e-5f);
|
|
fn->setOutput(y);
|
|
std::unique_ptr<CompiledFunction> compiled(fn->compile());
|
|
REQUIRE_EQ(compiled->getDevice(), std::string("gpu"));
|
|
|
|
std::unique_ptr<Tensor> xd(tf->zeros4(1, 1, 6, 6));
|
|
for (int i = 0; i < xd->getSize(); ++i) xd->set(i, float((i % 7) - 3));
|
|
std::unique_ptr<Tensor> out(compiled->run1(xd.get()));
|
|
|
|
std::unique_ptr<Tensor> ref(tf->conv2d(xd.get(), w.get(), 1, 0));
|
|
ref.reset(ref->relu());
|
|
ref.reset(ref->flatten());
|
|
ref.reset(tf->layernorm(ref.get(), 1e-5f));
|
|
REQUIRE_EQ(out->getSize(), ref->getSize());
|
|
for (int i = 0; i < out->getSize(); ++i) REQUIRE(std::fabs(out->get(i) - ref->get(i)) < 1e-3f);
|
|
}
|
|
|
|
TEST_CASE("tensor.gpu.transformerBlockInference") {
|
|
initHeadlessGfx(_ZEROERR_TEST_CONTEXT);
|
|
auto *tf = TF::create();
|
|
const int B = 1, T = 4, D = 8, C = 3;
|
|
const float scale = 1.f / std::sqrt(float(D));
|
|
const float eps = 1e-5f;
|
|
|
|
auto rand = [&](int d0, int d1, float s) {
|
|
std::unique_ptr<Tensor> w(tf->randomNormal2(d0, d1));
|
|
w.reset(w->mulScalar(s));
|
|
return w;
|
|
};
|
|
std::unique_ptr<Tensor> Wq(rand(D, D, 0.15f));
|
|
std::unique_ptr<Tensor> Wk(rand(D, D, 0.15f));
|
|
std::unique_ptr<Tensor> Wv(rand(D, D, 0.15f));
|
|
std::unique_ptr<Tensor> W1(rand(D, 2 * D, 0.15f));
|
|
std::unique_ptr<Tensor> W2(rand(2 * D, D, 0.15f));
|
|
std::unique_ptr<Tensor> Wc(rand(D, C, 0.15f));
|
|
|
|
std::unique_ptr<Tensor> xEager(tf->randomNormal3(B, T, D));
|
|
|
|
std::unique_ptr<Func> fn(tf->func());
|
|
Tensor *x = fn->input3(B, T, D);
|
|
Tensor *l1 = tf->layernorm(x, eps);
|
|
Tensor *l1m = tf->reshape2(l1, B * T, D);
|
|
Tensor *qq = tf->reshape4(tf->matmul(l1m, Wq.get()), 1, 1, T, D);
|
|
Tensor *kk = tf->reshape4(tf->matmul(l1m, Wk.get()), 1, 1, T, D);
|
|
Tensor *vv = tf->reshape4(tf->matmul(l1m, Wv.get()), 1, 1, T, D);
|
|
Tensor *a1 = tf->reshape3(tf->sdpa(qq, kk, vv, scale), B, T, D);
|
|
Tensor *x1s = tf->add(x, a1);
|
|
Tensor *l2 = tf->layernorm(x1s, eps);
|
|
Tensor *l2m = tf->reshape2(l2, B * T, D);
|
|
Tensor *hh = tf->matmul(l2m, W1.get());
|
|
hh = tf->gelu(hh);
|
|
hh = tf->matmul(hh, W2.get());
|
|
Tensor *x2s = tf->add(x1s, tf->reshape3(hh, B, T, D));
|
|
Tensor *lg = tf->matmul(tf->reshape2(x2s, B * T, D), Wc.get());
|
|
fn->setOutput(tf->softmax(lg, 1));
|
|
std::unique_ptr<CompiledFunction> compiled(fn->compile());
|
|
REQUIRE_EQ(compiled->getDevice(), std::string("gpu"));
|
|
|
|
std::unique_ptr<Tensor> out(compiled->run1(xEager.get()));
|
|
REQUIRE(std::fabs(out->reduceSum() - float(B * T)) < 1e-3f);
|
|
std::unique_ptr<Tensor> action(tf->argmax(out.get(), 1, 0));
|
|
REQUIRE_EQ(action->getSize(), B * T);
|
|
|
|
// cross-check the generated GPU kernels against the eager CPU reference
|
|
std::unique_ptr<Tensor> ln1(tf->layernorm(xEager.get(), eps));
|
|
std::unique_ptr<Tensor> ln1m(tf->reshape2(ln1.get(), B * T, D));
|
|
std::unique_ptr<Tensor> q(tf->reshape4(tf->matmul(ln1m.get(), Wq.get()), 1, 1, T, D));
|
|
std::unique_ptr<Tensor> k(tf->reshape4(tf->matmul(ln1m.get(), Wk.get()), 1, 1, T, D));
|
|
std::unique_ptr<Tensor> v(tf->reshape4(tf->matmul(ln1m.get(), Wv.get()), 1, 1, T, D));
|
|
std::unique_ptr<Tensor> att(tf->reshape3(tf->sdpa(q.get(), k.get(), v.get(), scale), B, T, D));
|
|
std::unique_ptr<Tensor> x1(tf->add(xEager.get(), att.get()));
|
|
std::unique_ptr<Tensor> ln2(tf->layernorm(x1.get(), eps));
|
|
std::unique_ptr<Tensor> ln2m(tf->reshape2(ln2.get(), B * T, D));
|
|
std::unique_ptr<Tensor> h(tf->matmul(ln2m.get(), W1.get()));
|
|
h.reset(tf->gelu(h.get()));
|
|
h.reset(tf->matmul(h.get(), W2.get()));
|
|
std::unique_ptr<Tensor> x2(tf->add(x1.get(), tf->reshape3(h.get(), B, T, D)));
|
|
std::unique_ptr<Tensor> logits(tf->matmul(tf->reshape2(x2.get(), B * T, D), Wc.get()));
|
|
std::unique_ptr<Tensor> ref(tf->softmax(logits.get(), 1));
|
|
for (int i = 0; i < out->getSize(); ++i) REQUIRE(std::fabs(out->get(i) - ref->get(i)) < 1e-2f);
|
|
}
|
|
|
|
TEST_CASE("tensor.gpu.sdpaMatchesCpu") {
|
|
initHeadlessGfx(_ZEROERR_TEST_CONTEXT);
|
|
auto *tf = TF::create();
|
|
const int B = 1, H = 1, T = 4, S = 4, D = 8;
|
|
std::unique_ptr<Tensor> q(tf->randomNormal4(B, H, T, D));
|
|
std::unique_ptr<Tensor> k(tf->randomNormal4(B, H, S, D));
|
|
std::unique_ptr<Tensor> v(tf->randomNormal4(B, H, S, D));
|
|
std::unique_ptr<Tensor> ref(tf->sdpa(q.get(), k.get(), v.get(), 0.353553f));
|
|
|
|
std::unique_ptr<Func> fn(tf->func());
|
|
Tensor *qq = fn->input4(B, H, T, D);
|
|
Tensor *kk = fn->input4(B, H, S, D);
|
|
Tensor *vv = fn->input4(B, H, S, D);
|
|
fn->setOutput(tf->sdpa(qq, kk, vv, 0.353553f));
|
|
std::unique_ptr<CompiledFunction> compiled(fn->compile());
|
|
REQUIRE_EQ(compiled->getDevice(), std::string("gpu"));
|
|
std::unique_ptr<Tensor> out(compiled->run3(q.get(), k.get(), v.get()));
|
|
for (int i = 0; i < out->getSize(); ++i) REQUIRE(std::fabs(out->get(i) - ref->get(i)) < 1e-4f);
|
|
}
|
|
|
|
namespace {
|
|
|
|
void checkGpuGraphMatchesEager(zeroerr::TestContext *_ZEROERR_TEST_CONTEXT, TF *tf,
|
|
const std::function<Tensor *(Tensor *)> &build, Tensor *feed, const char *name) {
|
|
std::unique_ptr<Func> fn(tf->func());
|
|
Tensor *in = fn->input3(1, 4, 8);
|
|
Tensor *outSym = build(in);
|
|
fn->setOutput(outSym);
|
|
std::unique_ptr<CompiledFunction> compiled(fn->compile());
|
|
REQUIRE_EQ(compiled->getDevice(), std::string("gpu"));
|
|
std::unique_ptr<Tensor> out(compiled->run1(feed));
|
|
std::unique_ptr<Tensor> ref(build(feed));
|
|
REQUIRE_EQ(out->getRank(), ref->getRank());
|
|
REQUIRE_EQ(out->getSize(), ref->getSize());
|
|
float worst = 0.f;
|
|
for (int i = 0; i < out->getSize(); ++i) {
|
|
const float error = std::fabs(out->get(i) - ref->get(i));
|
|
REQUIRE(std::isfinite(error));
|
|
worst = std::max(worst, error);
|
|
}
|
|
if (worst >= 1e-3f) fprintf(stderr, "stage %s worst=%.6g\n", name, double(worst));
|
|
REQUIRE_LT(worst, 1e-3f);
|
|
}
|
|
|
|
} // namespace
|
|
|
|
TEST_CASE("tensor.gpu.stageBisect") {
|
|
initHeadlessGfx(_ZEROERR_TEST_CONTEXT);
|
|
auto *tf = TF::create();
|
|
const int T = 4, D = 8;
|
|
std::unique_ptr<Tensor> x(tf->randomNormal3(1, T, D));
|
|
std::unique_ptr<Tensor> W(tf->randomNormal2(D, D));
|
|
|
|
auto stage1 = [&](Tensor *in) {
|
|
Tensor *l = tf->layernorm(in, 1e-5f);
|
|
Tensor *m = tf->reshape2(l, T, D);
|
|
return tf->reshape4(tf->matmul(m, W.get()), 1, 1, T, D);
|
|
};
|
|
auto stage2 = [&](Tensor *in) {
|
|
Tensor *q = stage1(in);
|
|
return tf->sdpa(q, q, q, 0.353553f);
|
|
};
|
|
auto stage3 = [&](Tensor *in) {
|
|
Tensor *a = stage2(in);
|
|
return tf->add(in, tf->reshape3(a, 1, T, D));
|
|
};
|
|
auto stage4 = [&](Tensor *in) {
|
|
Tensor *r = stage3(in);
|
|
return tf->layernorm(r, 1e-5f);
|
|
};
|
|
|
|
checkGpuGraphMatchesEager(_ZEROERR_TEST_CONTEXT, tf, stage1, x.get(), "stage1");
|
|
checkGpuGraphMatchesEager(_ZEROERR_TEST_CONTEXT, tf, stage2, x.get(), "stage2");
|
|
checkGpuGraphMatchesEager(_ZEROERR_TEST_CONTEXT, tf, stage3, x.get(), "stage3");
|
|
checkGpuGraphMatchesEager(_ZEROERR_TEST_CONTEXT, tf, stage4, x.get(), "stage4");
|
|
}
|