Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
135 lines
7.7 KiB
C++
135 lines
7.7 KiB
C++
#include "engine/framework/core/backend.h"
|
|
#include "engine/framework/core/execution_context.h"
|
|
#include "engine/framework/core/module.h"
|
|
#include "engine/framework/modules/recurrent_modules.h"
|
|
#include "engine/framework/modules/structural_modules.h"
|
|
#include "test_assert.h"
|
|
#include "engine/framework/debug/profiler.h"
|
|
#include "ggml-alloc.h"
|
|
|
|
#include <cmath>
|
|
#include <iostream>
|
|
#include <memory>
|
|
#include <vector>
|
|
|
|
namespace {
|
|
using namespace engine;
|
|
|
|
std::vector<float> check(int64_t frames, int64_t batch, bool reverse,
|
|
core::BackendType backend = core::BackendType::Cpu, int64_t features = 3, int64_t width = 4, bool reuse = false) {
|
|
core::ExecutionContext execution({backend, 0, 1});
|
|
std::unique_ptr<ggml_context, decltype(&ggml_free)> context(
|
|
ggml_init({32 * 1024 * 1024, nullptr, true}), ggml_free);
|
|
core::ModuleBuildContext ctx{context.get(), "lstm_sequence_test", backend};
|
|
using Shape = core::TensorShape;
|
|
std::vector<core::TensorValue> inputs;
|
|
auto tensor = [&](Shape shape) {
|
|
auto value = core::make_tensor(ctx, GGML_TYPE_F32, shape);
|
|
inputs.push_back(value);
|
|
ggml_set_input(value.tensor);
|
|
return value;
|
|
};
|
|
const auto input = tensor(Shape::from_dims({frames, batch, features}));
|
|
const auto hidden = tensor(Shape::from_dims({batch, width}));
|
|
const auto cell = tensor(Shape::from_dims({batch, width}));
|
|
const modules::LSTMSequenceWeights weights{{
|
|
tensor(Shape::from_dims({4 * width, features})), tensor(Shape::from_dims({4 * width, width})),
|
|
tensor(Shape::from_dims({4 * width})), tensor(Shape::from_dims({4 * width}))}};
|
|
const auto batched = modules::LSTMSequenceModule({features, width, reverse, true}).build(ctx, input, hidden, cell, weights);
|
|
const auto fusion_friendly = modules::LSTMSequenceModule({features, width, reverse, true, true})
|
|
.build(ctx, input, hidden, cell, weights);
|
|
const auto bidirectional = modules::BidirectionalLSTMModule({features, width, false, true}).build(
|
|
ctx, input, hidden, cell, hidden, cell, {weights.cell, weights.cell});
|
|
std::vector<modules::LSTMSequenceOutputs> single;
|
|
auto * graph = ggml_new_graph_custom(context.get(), 8192, false);
|
|
ggml_set_output(bidirectional.sequence.tensor);
|
|
ggml_build_forward_expand(graph, bidirectional.sequence.tensor);
|
|
for (const auto & value : {batched.sequence, batched.hidden, batched.cell}) {
|
|
ggml_set_output(value.tensor);
|
|
ggml_build_forward_expand(graph, value.tensor);
|
|
}
|
|
for (const auto & value : {fusion_friendly.sequence, fusion_friendly.hidden, fusion_friendly.cell}) {
|
|
ggml_set_output(value.tensor);
|
|
ggml_build_forward_expand(graph, value.tensor);
|
|
}
|
|
for (int64_t b = 0; b < batch; ++b) {
|
|
auto x = modules::SliceModule({1, b, 1}).build(ctx, input);
|
|
x = core::ensure_backend_addressable_layout(ctx, x);
|
|
x = core::reshape_tensor(ctx, x, Shape::from_dims({frames, features}));
|
|
single.push_back(modules::LSTMSequenceModule({features, width, reverse}).build(ctx, x,
|
|
modules::SliceModule({0, b, 1}).build(ctx, hidden),
|
|
modules::SliceModule({0, b, 1}).build(ctx, cell), weights));
|
|
for (const auto & value : {single.back().sequence, single.back().hidden, single.back().cell}) {
|
|
ggml_set_output(value.tensor);
|
|
ggml_build_forward_expand(graph, value.tensor);
|
|
}
|
|
}
|
|
std::unique_ptr<ggml_backend_buffer, decltype(&ggml_backend_buffer_free)> buffer(
|
|
reuse ? nullptr : ggml_backend_alloc_ctx_tensors(context.get(), execution.backend()), ggml_backend_buffer_free);
|
|
std::unique_ptr<ggml_gallocr, decltype(&ggml_gallocr_free)> alloc(
|
|
reuse ? ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution.backend())) : nullptr, ggml_gallocr_free);
|
|
if (reuse) test::require(alloc && ggml_gallocr_reserve(alloc.get(), graph) && ggml_gallocr_alloc_graph(alloc.get(), graph), "LSTM graph allocation failed");
|
|
else test::require(buffer != nullptr, "LSTM allocation failed");
|
|
for (size_t i = 0; i < inputs.size(); ++i) {
|
|
std::vector<float> values(ggml_nelements(inputs[i].tensor));
|
|
for (size_t j = 0; j < values.size(); ++j) values[j] = 0.2f * std::sin(static_cast<float>(j + 7 * i));
|
|
core::write_tensor_f32(inputs[i], values);
|
|
}
|
|
test::require(core::compute_backend_graph(execution.backend(), graph) == GGML_STATUS_SUCCESS, "LSTM compute failed");
|
|
const auto sequence = core::read_tensor_f32(batched.sequence.tensor);
|
|
const auto both_directions = core::read_tensor_f32(bidirectional.sequence.tensor);
|
|
const auto final_hidden = core::read_tensor_f32(batched.hidden.tensor);
|
|
const auto final_cell = core::read_tensor_f32(batched.cell.tensor);
|
|
const auto fusion_sequence = core::read_tensor_f32(fusion_friendly.sequence.tensor);
|
|
const auto fusion_hidden = core::read_tensor_f32(fusion_friendly.hidden.tensor);
|
|
const auto fusion_cell = core::read_tensor_f32(fusion_friendly.cell.tensor);
|
|
for (size_t i = 0; i < sequence.size(); ++i)
|
|
test::require_close(fusion_sequence[i], sequence[i], 0.0f, "fusion-friendly sequence");
|
|
for (size_t i = 0; i < final_hidden.size(); ++i) {
|
|
test::require_close(fusion_hidden[i], final_hidden[i], 0.0f, "fusion-friendly hidden");
|
|
test::require_close(fusion_cell[i], final_cell[i], 0.0f, "fusion-friendly cell");
|
|
}
|
|
for (int64_t b = 0; b < batch; ++b) {
|
|
const auto expected = core::read_tensor_f32(single[b].sequence.tensor);
|
|
const auto expected_hidden = core::read_tensor_f32(single[b].hidden.tensor);
|
|
const auto expected_cell = core::read_tensor_f32(single[b].cell.tensor);
|
|
for (int64_t t = 0; t < frames; ++t)
|
|
for (int64_t c = 0; c < width; ++c) {
|
|
test::require_close(sequence[(t * batch + b) * width + c], expected[t * width + c], 1e-6f, "batched sequence");
|
|
test::require_close(both_directions[(t * batch + b) * (2 * width) + (reverse ? width : 0) + c],
|
|
expected[t * width + c], 1e-6f, "bidirectional sequence");
|
|
}
|
|
for (int64_t c = 0; c < width; ++c) {
|
|
test::require_close(final_hidden[b * width + c], expected_hidden[c], 1e-6f, "batched hidden");
|
|
test::require_close(final_cell[b * width + c], expected_cell[c], 1e-6f, "batched cell");
|
|
}
|
|
}
|
|
core::release_backend_graph_resources(execution.backend(), graph, true);
|
|
return sequence;
|
|
}
|
|
} // namespace
|
|
|
|
int main(int argc, char ** argv) {
|
|
try {
|
|
engine::debug::configure_logging({true});
|
|
const auto started = std::chrono::steady_clock::now();
|
|
if (argc > 2 && std::string(argv[1]) == "--backend" && std::string(argv[2]) == "webgpu") {
|
|
const auto expected = check(1, 1, false, core::BackendType::Cpu, 640, 640, true);
|
|
const auto actual = check(1, 1, false, core::BackendType::WebGPU, 640, 640, true);
|
|
for (size_t i = 0; i < expected.size(); ++i)
|
|
test::require_close(actual[i], expected[i], 1e-5f, "WebGPU LSTM CPU reference");
|
|
debug::timing_log_scalar("session.wall_ms", debug::elapsed_ms(started));
|
|
std::cout << "WebGPU 640-wide LSTM matches CPU with graph buffer reuse\n";
|
|
return 0;
|
|
}
|
|
for (const auto frames : {1, 5, 6})
|
|
for (const auto batch : {1, 3})
|
|
for (const bool reverse : {false, true}) check(frames, batch, reverse);
|
|
debug::timing_log_scalar("session.wall_ms", debug::elapsed_ms(started));
|
|
std::cout << "Batched LSTM matches independent sequences in all 12 cases\n";
|
|
return 0;
|
|
} catch (const std::exception & error) {
|
|
std::cerr << error.what() << '\n';
|
|
return 1;
|
|
}
|
|
}
|