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
101 lines
3.5 KiB
C++
101 lines
3.5 KiB
C++
#pragma once
|
|
|
|
#include "engine/framework/core/execution_context.h"
|
|
#include "engine/framework/core/module.h"
|
|
#include "engine/community_models/minimax_h3/types.h"
|
|
|
|
#include "ggml-alloc.h"
|
|
#include "ggml-backend.h"
|
|
|
|
#include <cstddef>
|
|
#include <cstdint>
|
|
#include <vector>
|
|
|
|
namespace engine::models::minimax_h3 {
|
|
|
|
struct MiniMaxH3SamplerTargetSpec {
|
|
engine::core::TensorShape packed_state_shape;
|
|
engine::core::TensorShape prediction_shape;
|
|
int64_t active_row_start = 0;
|
|
};
|
|
|
|
struct MiniMaxH3SamplerGraphSpec {
|
|
MiniMaxH3SamplerTargetSpec primary;
|
|
MiniMaxH3SamplerTargetSpec secondary;
|
|
};
|
|
|
|
struct MiniMaxH3SamplerInput {
|
|
ggml_tensor * primary_state = nullptr;
|
|
ggml_tensor * secondary_state = nullptr;
|
|
ggml_tensor * primary_velocity = nullptr;
|
|
ggml_tensor * secondary_velocity = nullptr;
|
|
float primary_sigma = 0.0F;
|
|
float secondary_sigma = 0.0F;
|
|
float primary_sigma_delta = 0.0F;
|
|
float secondary_sigma_delta = 0.0F;
|
|
};
|
|
|
|
struct MiniMaxH3SamplerOutput {
|
|
std::vector<float> next_primary;
|
|
std::vector<float> next_secondary;
|
|
};
|
|
|
|
class MiniMaxH3SamplerGraph {
|
|
public:
|
|
MiniMaxH3SamplerGraph(
|
|
engine::core::ExecutionContext & execution,
|
|
const MiniMaxH3SamplerGraphSpec & spec,
|
|
MiniMaxH3SamplerMode mode);
|
|
~MiniMaxH3SamplerGraph();
|
|
|
|
void run(const MiniMaxH3SamplerInput & input, MiniMaxH3SamplerOutput & output);
|
|
|
|
double device_copy_ms() const;
|
|
double graph_compute_ms() const;
|
|
double output_read_ms() const;
|
|
double history_update_ms() const;
|
|
|
|
private:
|
|
engine::core::ExecutionContext & execution_;
|
|
const MiniMaxH3SamplerGraphSpec spec_;
|
|
const MiniMaxH3SamplerMode mode_;
|
|
ggml_context * ctx_ = nullptr;
|
|
ggml_context * input_ctx_ = nullptr;
|
|
ggml_cgraph * graph_ = nullptr;
|
|
ggml_gallocr_t gallocr_ = nullptr;
|
|
ggml_backend_buffer_t input_buffer_ = nullptr;
|
|
engine::core::HostGraphPlan plan_;
|
|
engine::core::TensorValue primary_state_t_;
|
|
engine::core::TensorValue secondary_state_t_;
|
|
engine::core::TensorValue primary_velocity_t_;
|
|
engine::core::TensorValue secondary_velocity_t_;
|
|
engine::core::TensorValue previous_primary_denoised_t_;
|
|
engine::core::TensorValue previous_secondary_denoised_t_;
|
|
engine::core::TensorValue previous_primary_velocity_t_;
|
|
engine::core::TensorValue previous_secondary_velocity_t_;
|
|
engine::core::TensorValue primary_sigma_t_;
|
|
engine::core::TensorValue secondary_sigma_t_;
|
|
engine::core::TensorValue primary_state_coeff_t_;
|
|
engine::core::TensorValue primary_velocity_coeff_t_;
|
|
engine::core::TensorValue primary_denoised_coeff_t_;
|
|
engine::core::TensorValue primary_previous_denoised_coeff_t_;
|
|
engine::core::TensorValue primary_previous_velocity_coeff_t_;
|
|
engine::core::TensorValue secondary_state_coeff_t_;
|
|
engine::core::TensorValue secondary_velocity_coeff_t_;
|
|
engine::core::TensorValue secondary_denoised_coeff_t_;
|
|
engine::core::TensorValue secondary_previous_denoised_coeff_t_;
|
|
engine::core::TensorValue secondary_previous_velocity_coeff_t_;
|
|
engine::core::TensorValue denoised_primary_t_;
|
|
engine::core::TensorValue denoised_secondary_t_;
|
|
engine::core::TensorValue next_primary_t_;
|
|
engine::core::TensorValue next_secondary_t_;
|
|
float previous_primary_sigma_ = 0.0F;
|
|
float previous_secondary_sigma_ = 0.0F;
|
|
bool has_history_ = false;
|
|
double device_copy_ms_ = 0.0;
|
|
double graph_compute_ms_ = 0.0;
|
|
double output_read_ms_ = 0.0;
|
|
double history_update_ms_ = 0.0;
|
|
};
|
|
|
|
} // namespace engine::models::minimax_h3
|