69 lines
3.1 KiB
C++
69 lines
3.1 KiB
C++
#include <cmath>
|
|
#include <memory>
|
|
#include <vector>
|
|
#include "tensor/Graph.h"
|
|
#include "tensor/KernelGen.h"
|
|
#include "tensor/Optimizer.h"
|
|
#include "tensor/TF.h"
|
|
#include "tensor/Tensor.h"
|
|
#include "zeroerr/assert.h"
|
|
#include "zeroerr/unittest.h"
|
|
using namespace eve::tensor;
|
|
|
|
TEST_CASE("tensor.optimizer.residualAttentionDependenciesAndSlots") {
|
|
auto *tf = TF::create();
|
|
const int T = 4, D = 8;
|
|
std::unique_ptr<Tensor> W(tf->fill2(D, D, -0.1f));
|
|
for (int i = 0; i < D; ++i) W->set2(i, i, 0.7f + float(i) * 0.1f);
|
|
std::unique_ptr<Func> fn(tf->func());
|
|
Tensor *in = fn->input3(1, T, D);
|
|
Tensor *l = tf->layernorm(in, 1e-5f);
|
|
Tensor *m = tf->reshape2(l, T, D);
|
|
Tensor *mm = tf->matmul(m, W.get());
|
|
Tensor *q = tf->reshape4(mm, 1, 1, T, D);
|
|
Tensor *a = tf->sdpa(q, q, q, 0.353553f);
|
|
Tensor *a3 = tf->reshape3(a, 1, T, D);
|
|
fn->setOutput(tf->add(in, a3));
|
|
|
|
OptimizedGraph opt = optimizeGraph(fn->graph(), fn->outputNode());
|
|
REQUIRE(!opt.order.empty());
|
|
REQUIRE(opt.outputNode == fn->outputNode());
|
|
REQUIRE(opt.nodeSlot.size() == size_t(fn->graph().nodeCount()));
|
|
std::vector<bool> visited(size_t(fn->graph().nodeCount()), false);
|
|
for (int id : opt.order) {
|
|
REQUIRE(id >= 0);
|
|
REQUIRE(id < fn->graph().nodeCount());
|
|
REQUIRE(!visited[size_t(id)]);
|
|
const auto &node = fn->graph().node(id);
|
|
for (int input : {node.in0, node.in1, node.in2, node.in3, node.in4}) {
|
|
if (input >= 0) REQUIRE(visited[size_t(input)]);
|
|
}
|
|
visited[size_t(id)] = true;
|
|
}
|
|
for (const auto &group : opt.groups) {
|
|
REQUIRE(group.outputNode >= 0);
|
|
const int slot = opt.nodeSlot[size_t(group.outputNode)];
|
|
REQUIRE(slot >= 0);
|
|
REQUIRE(size_t(slot) < opt.slotSize.size());
|
|
REQUIRE(opt.slotSize[size_t(slot)] >= fn->graph().node(group.outputNode).size);
|
|
if (group.kind == GroupKind::Elementwise) {
|
|
KernelSpec spec;
|
|
generateKernel(fn->graph(), group, spec);
|
|
REQUIRE(!spec.pass2.empty());
|
|
}
|
|
}
|
|
std::unique_ptr<CompiledFunction> compiled(fn->compile());
|
|
std::unique_ptr<Tensor> input(tf->zeros3(1, T, D));
|
|
for (int i = 0; i < input->getSize(); ++i) input->set(i, float(i % 7) * 0.2f);
|
|
std::unique_ptr<Tensor> normalized(tf->layernorm(input.get(), 1e-5f));
|
|
std::unique_ptr<Tensor> matrix(tf->reshape2(normalized.get(), T, D));
|
|
std::unique_ptr<Tensor> multiplied(tf->matmul(matrix.get(), W.get()));
|
|
std::unique_ptr<Tensor> query(tf->reshape4(multiplied.get(), 1, 1, T, D));
|
|
std::unique_ptr<Tensor> attended(tf->sdpa(query.get(), query.get(), query.get(), 0.353553f));
|
|
std::unique_ptr<Tensor> reshaped(tf->reshape3(attended.get(), 1, T, D));
|
|
std::unique_ptr<Tensor> expected(tf->add(input.get(), reshaped.get()));
|
|
std::unique_ptr<Tensor> actual(compiled->run1(input.get()));
|
|
REQUIRE(actual->getSize() == expected->getSize());
|
|
for (int i = 0; i < actual->getSize(); ++i) REQUIRE(std::fabs(actual->get(i) - expected->get(i)) < 1e-4f);
|
|
}
|