forked from 0xShug0/audio.cpp
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathlinear_module.cpp
More file actions
114 lines (89 loc) · 3.48 KB
/
Copy pathlinear_module.cpp
File metadata and controls
114 lines (89 loc) · 3.48 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
#include "engine/framework/modules/linear_module.h"
#include <stdexcept>
namespace engine::modules {
namespace {
const core::ModulePortSpec kLinearInputs[] = {
{"input", core::PortKind::Activation, false},
{"weight", core::PortKind::Parameter, false},
{"bias", core::PortKind::Parameter, true},
};
const core::ModulePortSpec kLinearOutputs[] = {
{"output", core::PortKind::Activation, false},
};
const core::ModuleSchema kLinearSchema = {
"Linear",
"nn.primitive",
kLinearInputs,
3,
kLinearOutputs,
1,
"Applies an affine projection over the last logical dimension.",
};
core::TensorShape flatten_to_matrix_shape(const core::TensorShape & shape) {
if (shape.rank == 1) {
return core::TensorShape::from_dims({1, shape.last_dim()});
}
return core::TensorShape::from_dims({shape.prefix_elements(), shape.last_dim()});
}
void validate_weight_shape(const LinearConfig & config, const LinearWeights & weights) {
core::validate_shape(
weights.weight,
core::TensorShape::from_dims({config.out_features, config.in_features}),
"weight");
if (!config.use_bias) {
return;
}
if (!weights.bias.has_value()) {
throw std::runtime_error("bias is required when LinearConfig.use_bias is true");
}
core::validate_shape(
*weights.bias,
core::TensorShape::from_dims({config.out_features}),
"bias");
}
} // namespace
LinearModule::LinearModule(LinearConfig config) : config_(config) {
if (config_.in_features <= 0 || config_.out_features <= 0) {
throw std::runtime_error("LinearConfig features must be positive");
}
}
const LinearConfig & LinearModule::config() const noexcept {
return config_;
}
const core::ModuleSchema & LinearModule::schema() const noexcept {
return static_schema();
}
core::TensorValue LinearModule::build(
core::ModuleBuildContext & ctx,
const core::TensorValue & input,
const LinearWeights & weights) const {
if (ctx.ggml == nullptr) {
throw std::runtime_error("ModuleBuildContext.ggml is null");
}
core::validate_rank_between(input, 1, core::kMaxTensorRank, "input");
core::validate_last_dim(input, config_.in_features, "input");
validate_weight_shape(config_, weights);
const core::TensorValue contiguous_input = core::ensure_backend_addressable_layout(ctx, input);
const core::TensorShape matrix_input_shape = flatten_to_matrix_shape(contiguous_input.shape);
core::TensorValue matrix_input = core::reshape_tensor(ctx, contiguous_input, matrix_input_shape);
ggml_tensor * projected_raw = ggml_mul_mat(ctx.ggml, weights.weight.tensor, matrix_input.tensor);
if (config_.precision != GGML_PREC_DEFAULT) {
ggml_mul_mat_set_prec(projected_raw, config_.precision);
}
core::TensorValue projected = core::wrap_tensor(
projected_raw,
core::TensorShape::from_dims({matrix_input_shape.at(0), config_.out_features}),
GGML_TYPE_F32);
if (config_.use_bias) {
ggml_tensor * biased_raw = ggml_add(ctx.ggml, projected.tensor, weights.bias->tensor);
projected = core::wrap_tensor(
biased_raw,
core::TensorShape::from_dims({matrix_input_shape.at(0), config_.out_features}),
GGML_TYPE_F32);
}
return core::reshape_tensor(ctx, projected, input.shape.with_last_dim(config_.out_features));
}
const core::ModuleSchema & LinearModule::static_schema() noexcept {
return kLinearSchema;
}
} // namespace engine::modules