Logo ROOT  
Reference Guide
 
Loading...
Searching...
No Matches
SOFIE_common_helpers.cxx
Go to the documentation of this file.
1/// \file SOFIE_common_helpers.cxx
2/// Standalone definitions of the SOFIE inference helpers (Im2col, Gemm_Call, ...).
3/// RModel records which helpers a model needs (RModel::AddNeededHelperFunction)
4/// and dumps only those into the generated namespace, so the emitted header is
5/// self-contained (no TMVA/SOFIE_common.hxx include). These are dependency-free
6/// copies of the SOFIE_common.hxx originals: keep them in sync when those change.
7
9
10namespace TMVA {
11namespace Experimental {
12namespace SOFIE {
13
14namespace {
15
16// Standalone helper snippets, each emitted verbatim inside the generated model
17// namespace. They depend only on the C++ stdlib (headers collected in
18// GenerateHelperFunctionsCode) and, for Gemm, on the local BLAS::sgemm_ below.
19
20// extern "C" declaration of the BLAS routine Gemm_Call needs, in a nested BLAS
21// namespace. Skipped when the caller already declared sgemm_ (see the
22// sgemmAlreadyDeclared parameter of GenerateHelperFunctionsCode).
23constexpr const char *kBlasSgemm = R"SOFIE(
24namespace BLAS {
25extern "C" void sgemm_(const char *transa, const char *transb, const int *m, const int *n, const int *k,
26 const float *alpha, const float *A, const int *lda, const float *B, const int *ldb,
27 const float *beta, float *C, const int *ldc);
28} // namespace BLAS
29)SOFIE";
30
31constexpr const char *kConvertShapeToLength = R"SOFIE(
32inline std::size_t ConvertShapeToLength(const std::vector<std::size_t> &shape)
33{
34 std::size_t length = 1;
35 for (auto &dim : shape)
36 length *= dim;
37 return length;
38}
39)SOFIE";
40
41constexpr const char *kConvertShapeToString = R"SOFIE(
42inline std::string ConvertShapeToString(const std::vector<std::size_t> &shape)
43{
44 std::stringstream out;
45 out << "{ ";
46 for (std::size_t i = 0; i < shape.size(); i++) {
47 out << shape[i];
48 if (i < shape.size() - 1)
49 out << " , ";
50 }
51 out << " }";
52 return out.str();
53}
54)SOFIE";
55
56// Branchless bounds check `0 <= a < b`: casting to unsigned collapses it to a
57// single comparison (a negative `a` wraps to a large value that fails `< b`).
58constexpr const char *kIsAGeZero = R"SOFIE(
59inline bool is_a_ge_zero_and_a_lt_b(int a, int b)
60{
61 return static_cast<unsigned>(a) < static_cast<unsigned>(b);
62}
63)SOFIE";
64
65// im2col: re-arrange convolution input into a matrix usable directly by BLAS.
66// It loops over each element of the filtered region first, following the input
67// layout, so reads/writes stay consecutive in memory; the result is already
68// transposed -- a (channels*kernel_h*kernel_w , output_h*output_w) matrix.
69// Example: input a1 a2 a3
70// b1 b2 b3 with a 2x2 kernel (k1,k2,k3,k4) and padding 1
71// c1 c2 c3
72// gives a 4x16 matrix, output-ordered (all elements for k1, then k2, ...):
73// ( 0 0 0 0 0 a1 a2 a3 0 b1 b2 b3 0 c1 c2 c3 ) k1
74// ( 0 0 0 0 a1 a2 a3 0 b1 b2 b3 0 c1 c2 c3 0 ) k2
75// ( 0 a1 a2 a3 0 b1 b2 b3 0 c1 c2 c3 0 0 0 0 ) k3
76// ( a1 a2 a3 0 b1 b2 b3 0 c1 c2 c3 0 0 0 0 0 ) k4
77// Per-axis begin/end padding can differ (ONNX "pads" attribute, and odd total
78// padding from the SAME_UPPER / SAME_LOWER autopad modes).
79constexpr const char *kIm2col = R"SOFIE(
80template <typename T>
81void Im2col(const T *data_im, const int channels, const int height, const int width, const int kernel_h,
82 const int kernel_w, const int pad_h_begin, const int pad_h_end, const int pad_w_begin,
83 const int pad_w_end, const int stride_h, const int stride_w,
84 const int dilation_h, const int dilation_w, T *data_col)
85{
86 const int output_h = (height + pad_h_begin + pad_h_end - (dilation_h * (kernel_h - 1) + 1)) / stride_h + 1;
87 const int output_w = (width + pad_w_begin + pad_w_end - (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1;
88 const int channel_size = height * width;
89 for (int channel = channels; channel--; data_im += channel_size) {
90 for (int kernel_row = 0; kernel_row < kernel_h; kernel_row++) {
91 for (int kernel_col = 0; kernel_col < kernel_w; kernel_col++) {
92 int input_row = -pad_h_begin + kernel_row * dilation_h;
93 for (int output_rows = output_h; output_rows; output_rows--) {
94 if (!is_a_ge_zero_and_a_lt_b(input_row, height)) {
95 for (int output_cols = output_w; output_cols; output_cols--) {
96 *(data_col++) = 0;
97 }
98 } else {
99 int input_col = -pad_w_begin + kernel_col * dilation_w;
100 for (int output_col = output_w; output_col; output_col--) {
101 if (is_a_ge_zero_and_a_lt_b(input_col, width)) {
102 *(data_col++) = data_im[input_row * width + input_col];
103 } else {
104 *(data_col++) = 0;
105 }
106 input_col += stride_w;
107 }
108 }
109 input_row += stride_h;
110 }
111 }
112 }
113 }
114}
115)SOFIE";
116
117constexpr const char *kIm2col3d = R"SOFIE(
118template <typename T>
119void Im2col_3d(const T *data_im, const int channels,
120 const int depth, const int height, const int width,
121 const int kernel_d, const int kernel_h, const int kernel_w,
122 const int pad_d_begin, const int pad_d_end, const int pad_h_begin, const int pad_h_end,
123 const int pad_w_begin, const int pad_w_end,
124 const int stride_d, const int stride_h, const int stride_w,
125 const int dilation_d, const int dilation_h, const int dilation_w, T *data_col)
126{
127 const int output_h = (height + pad_h_begin + pad_h_end - (dilation_h * (kernel_h - 1) + 1)) / stride_h + 1;
128 const int output_w = (width + pad_w_begin + pad_w_end - (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1;
129 const int output_d = (depth + pad_d_begin + pad_d_end - (dilation_d * (kernel_d - 1) + 1)) / stride_d + 1;
130 const int channel_size = height * width * depth;
131 for (int channel = channels; channel--; data_im += channel_size) {
132 for (int kernel_depth = 0; kernel_depth < kernel_d; kernel_depth++) {
133 for (int kernel_row = 0; kernel_row < kernel_h; kernel_row++) {
134 for (int kernel_col = 0; kernel_col < kernel_w; kernel_col++) {
135 int input_dep = -pad_d_begin + kernel_depth * dilation_d;
136 for (int output_dep = output_d; output_dep; output_dep--) {
137 if (!is_a_ge_zero_and_a_lt_b(input_dep, depth)) {
138 for (int output_rows = output_h; output_rows; output_rows--) {
139 for (int output_cols = output_w; output_cols; output_cols--) {
140 *(data_col++) = 0;
141 }
142 }
143 } else {
144 int input_row = -pad_h_begin + kernel_row * dilation_h;
145 for (int output_rows = output_h; output_rows; output_rows--) {
146 if (!is_a_ge_zero_and_a_lt_b(input_row, height)) {
147 for (int output_cols = output_w; output_cols; output_cols--) {
148 *(data_col++) = 0;
149 }
150 } else {
151 int input_col = -pad_w_begin + kernel_col * dilation_w;
152 for (int output_col = output_w; output_col; output_col--) {
153 if (is_a_ge_zero_and_a_lt_b(input_col, width)) {
154 *(data_col++) = data_im[input_dep * width * height + input_row * width + input_col];
155 } else {
156 *(data_col++) = 0;
157 }
158 input_col += stride_w;
159 }
160 }
161 input_row += stride_h;
162 }
163 }
164 input_dep += stride_d;
165 }
166 }
167 }
168 }
169 }
170}
171)SOFIE";
172
173constexpr const char *kCol2im = R"SOFIE(
174template <typename Dtype>
175void col2im(const Dtype *data_col, const int channels,
176 const int height, const int width, const int kernel_h, const int kernel_w,
177 const int pad_h_begin, const int pad_h_end, const int pad_w_begin, const int pad_w_end,
178 const int stride_h, const int stride_w,
179 const int dilation_h, const int dilation_w,
180 Dtype *data_im)
181{
182 // output must start zeroed: col2im scatters with += so overlapping columns accumulate
183 std::fill(data_im, data_im + height * width * channels, 0.);
184 const int output_h = (height + pad_h_begin + pad_h_end - (dilation_h * (kernel_h - 1) + 1)) / stride_h + 1;
185 const int output_w = (width + pad_w_begin + pad_w_end - (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1;
186 const int channel_size = height * width;
187 for (int channel = channels; channel--; data_im += channel_size) {
188 for (int kernel_row = 0; kernel_row < kernel_h; kernel_row++) {
189 for (int kernel_col = 0; kernel_col < kernel_w; kernel_col++) {
190 int input_row = -pad_h_begin + kernel_row * dilation_h;
191 for (int output_rows = output_h; output_rows; output_rows--) {
192 if (!is_a_ge_zero_and_a_lt_b(input_row, height)) {
193 data_col += output_w;
194 } else {
195 int input_col = -pad_w_begin + kernel_col * dilation_w;
196 for (int output_col = output_w; output_col; output_col--) {
197 if (is_a_ge_zero_and_a_lt_b(input_col, width)) {
198 data_im[input_row * width + input_col] += *data_col;
199 }
200 data_col++;
201 input_col += stride_w;
202 }
203 }
204 input_row += stride_h;
205 }
206 }
207 }
208 }
209}
210)SOFIE";
211
212// Broadcast helpers implementing numpy broadcasting rules (see
213// https://numpy.org/doc/stable/user/basics.broadcasting.html and
214// https://github.com/onnx/onnx/blob/main/docs/Broadcasting.md). Unidirectional
215// broadcast: only the input shape is stretched to targetShape, not vice versa.
216// These are rewritten with respect to TMVA/SOFIE_common.hxx to avoid the
217// dependency on std::span (which would require C++20 or ROOT's RSpan.hxx in the
218// generated code): the input is passed as a raw pointer plus a length instead.
219constexpr const char *kBroadcastTensor = R"SOFIE(
220template <typename T>
221void BroadcastTensor(const T *data, std::size_t curLength, const std::vector<std::size_t> &shape,
222 const std::vector<std::size_t> &targetShape, T *broadcastedData)
223{
224 std::size_t size = shape.size();
225 if (size > 1 && shape.front() == targetShape.front() && shape.back() == 1) {
226 std::size_t bsize = targetShape.back();
227 for (int k = int(size) - 2; k >= 0; k--) {
228 if (shape[k] != 1)
229 break;
230 bsize *= targetShape[k];
231 }
232 for (std::size_t i = 0; i < curLength; i++) {
233 std::fill(broadcastedData + i * bsize, broadcastedData + (i + 1) * bsize, data[i]);
234 }
235 return;
236 }
237
238 std::copy(data, data + curLength, broadcastedData);
239 std::size_t arrayNum = 1;
240 std::vector<T> newData(ConvertShapeToLength(targetShape));
241
242 for (std::size_t idx = 0; idx < size; idx++) {
243 std::size_t dim = shape[idx];
244 std::size_t targetDim = targetShape[idx];
245 if (dim == 1 && targetDim > 1) {
246 std::size_t newLength = curLength * targetDim;
247 std::size_t arrayLength = curLength / arrayNum;
248 if (arrayLength > 1) {
249 for (std::size_t arrayIdx = 0; arrayIdx < arrayNum; arrayIdx++) {
250 for (std::size_t targetIdx = 0; targetIdx < targetDim; targetIdx++) {
251 std::size_t offset = arrayIdx * arrayLength * targetDim + targetIdx * arrayLength;
252 std::copy(broadcastedData + arrayIdx * arrayLength,
253 broadcastedData + (arrayIdx + 1) * arrayLength,
254 newData.begin() + offset);
255 }
256 }
257 } else {
258 for (std::size_t arrayIdx = 0; arrayIdx < arrayNum; arrayIdx++) {
259 std::fill(newData.begin() + arrayIdx * targetDim,
260 newData.begin() + (arrayIdx + 1) * targetDim, broadcastedData[arrayIdx]);
261 }
262 }
263 curLength = newLength;
264 std::copy(newData.begin(), newData.begin() + newLength, broadcastedData);
265 }
266 arrayNum *= targetDim;
267 }
268}
269
270template <typename T>
271T *CreateBroadcastTensor(const T *data, const std::vector<std::size_t> &shape,
272 const std::vector<std::size_t> &targetShape, std::size_t targetLength)
273{
274 T *broadcastedData = new T[targetLength];
275 std::size_t curLength = ConvertShapeToLength(shape);
276 BroadcastTensor<T>(data, curLength, shape, targetShape, broadcastedData);
277 return broadcastedData;
278}
279
280template <typename T>
281T *UnidirectionalBroadcast(const T *data, const std::vector<std::size_t> &shape,
282 const std::vector<std::size_t> &targetShape)
283{
284 if (shape.size() < targetShape.size()) {
285 std::size_t targetSize = targetShape.size();
286 std::vector<std::size_t> newShape(targetSize, 1);
287 std::size_t offset = targetSize - shape.size();
288 std::copy(shape.begin(), shape.end(), newShape.begin() + offset);
289 return CreateBroadcastTensor(data, newShape, targetShape, ConvertShapeToLength(targetShape));
290 }
291 return CreateBroadcastTensor(data, shape, targetShape, ConvertShapeToLength(targetShape));
292}
293
294template <typename T>
295void UnidirectionalBroadcast(const T *data, const std::vector<std::size_t> &shape,
296 const std::vector<std::size_t> &targetShape, T *broadcastedData)
297{
298 std::size_t curLength = ConvertShapeToLength(shape);
299 if (shape.size() < targetShape.size()) {
300 std::size_t targetSize = targetShape.size();
301 std::vector<std::size_t> newShape(targetSize, 1);
302 std::size_t offset = targetSize - shape.size();
303 std::copy(shape.begin(), shape.end(), newShape.begin() + offset);
304 BroadcastTensor(data, curLength, newShape, targetShape, broadcastedData);
305 return;
306 }
307 BroadcastTensor(data, curLength, shape, targetShape, broadcastedData);
308}
309)SOFIE";
310
311constexpr const char *kBroadcastConvBias = R"SOFIE(
312template <typename T>
313T *BroadcastConvBias(const T *data, const std::size_t channel, const std::vector<std::size_t> &targetShape)
314{
315 std::size_t size = targetShape.size();
316 if (targetShape[1] != channel) {
317 std::stringstream ss;
318 ss << "TMVA::SOFIE - Error broadcasting Conv Bias of shape {";
319 ss << std::to_string(channel);
320 ss << "} to ";
321 ss << ConvertShapeToString(targetShape);
322 throw std::runtime_error(ss.str());
323 }
324
325 std::size_t targetLength = ConvertShapeToLength(targetShape);
326 T *newData = new T[targetLength];
327
328 if (targetLength == channel) {
329 std::copy(data, data + channel, newData);
330 return newData;
331 }
332
333 std::size_t cStride = 1;
334 for (std::size_t i = 2; i < size; i++)
335 cStride *= targetShape[i];
336 for (std::size_t i = 0; i < channel; i++) {
337 std::fill(newData + i * cStride, newData + (i + 1) * cStride, data[i]);
338 }
339 std::size_t batch = targetShape[0];
340 std::size_t bStride = channel * cStride;
341 for (std::size_t i = 1; i < batch; i++) {
342 std::copy(newData, newData + bStride, newData + i * bStride);
343 }
344 return newData;
345}
346)SOFIE";
347
348constexpr const char *kGemmCall = R"SOFIE(
349inline void Gemm_Call(float *output, bool transa, bool transb, int m, int n, int k, float alpha, const float *A,
350 const float *B, float beta, const float *C)
351{
352 char ct = 't';
353 char cn = 'n';
354 const int *lda = transa ? &k : &m;
355 const int *ldb = transb ? &n : &k;
356 const int *ldc = &m;
357 if (C != nullptr) {
358 std::copy(C, C + m * n, output);
359 }
360 BLAS::sgemm_(transa ? &ct : &cn, transb ? &ct : &cn, &m, &n, &k, &alpha, A, lda, B, ldb, &beta, output, ldc);
361}
362)SOFIE";
363
364// Custom Clad reverse-mode pullbacks for the helpers (they used to live in
365// Math/CladDerivator.h). Gemm_Call and Copy need a hand-written pullback because
366// their bodies call bodyless routines (sgemm_, std::copy -> memmove) that Clad
367// cannot differentiate; Fill and Relu are included for completeness so a
368// differentiated model needs neither SOFIE_common.hxx nor CladDerivator.h. How
369// they are placed and found is explained at the emission site
370// (GenerateHelperFunctionsCode).
371constexpr const char *kGemmCallPullback = R"SOFIE(
372inline void Gemm_Call_pullback(float *output, bool transa, bool transb, int m, int n, int k, float alpha,
373 const float *A, const float *B, float beta, const float *C, float *_d_output, bool *,
374 bool *, int *, int *, int *, float *_d_alpha, float *_d_A, float *_d_B, float *_d_beta,
375 float *_d_C)
376{
377 // TODO:
378 // - fix and test the implementation for alpha != 1.0
379 if (alpha != 1.0f) {
380 return;
381 }
382
383 // beta needs to be one because we want to add to _d_A and _d_B instead of
384 // overwriting it.
385 float one = 1.;
386
387 // ---- dA ----
388 if (!transa) {
389 // dA += dY * op(B)^T
390 Gemm_Call(_d_A, false, !transb, m, k, n, one, _d_output, B, one, _d_A);
391 } else {
392 // dA += op(B) * dY^T
393 Gemm_Call(_d_A, transb, true, k, m, n, one, B, _d_output, one, _d_A);
394 }
395
396 // ---- dB ----
397 if (!transb) {
398 // dB += op(A)^T * dY
399 Gemm_Call(_d_B, !transa, false, k, n, m, one, A, _d_output, one, _d_B);
400 } else {
401 // dB += dY^T * op(A)
402 Gemm_Call(_d_B, true, transa, n, k, m, one, _d_output, A, one, _d_B);
403 }
404
405 int sizeC = n * m;
406
407 for (int i = 0; i < sizeC; ++i) {
408 if (C) {
409 *_d_alpha += _d_output[i] * (output[i] - beta * C[i]);
410 *_d_beta += _d_output[i] * C[i];
411 } else {
412 *_d_alpha += _d_output[i] * output[i];
413 }
414 if (_d_C)
415 _d_C[i] += _d_output[i] * beta;
416 }
417}
418)SOFIE";
419
420// Pullback for the Copy helper. Copy's body uses std::copy, which lowers to the
421// bodyless __builtin_memmove that Clad cannot differentiate, so a hand-written
422// pullback is required. The generated Copy is a template; this float overload
423// matches its float instantiation (the only one used by inference code).
424constexpr const char *kCopyPullback = R"SOFIE(
425inline void Copy_pullback(float *output, const float *input, int size, float *_d_output, float *_d_input, int *)
426{
427 for (int i = 0; i < size; i++) {
428 output[i] = input[i];
429 _d_input[i] += _d_output[i];
430 _d_output[i] = 0.F;
431 }
432}
433)SOFIE";
434
435// Pullback for the Fill helper (std::fill -> bodyless builtin).
436constexpr const char *kFillPullback = R"SOFIE(
437inline void Fill_pullback(float *output, float value, int size, float *_d_output, float *_d_value, int *)
438{
439 for (int i = 0; i < size; i++) {
440 output[i] = value;
441 *_d_value += _d_output[i];
442 _d_output[i] = 0.F;
443 }
444}
445)SOFIE";
446
447// Pullback for the Relu helper. Relu's body is differentiable by Clad on its
448// own, but providing the pullback keeps the derivative identical to the one
449// previously supplied by Math/CladDerivator.h.
450constexpr const char *kReluPullback = R"SOFIE(
451inline void Relu_pullback(float *output, const float *input, int size, float *_d_output, float *_d_input, int *)
452{
453 for (int i = 0; i < size; i++) {
454 output[i] = input[i] > 0.F ? input[i] : 0.F;
455 float _r_d0 = _d_output[i];
456 _d_output[i] = 0.F;
457 if (input[i] > 0.F)
458 _d_input[i] += _r_d0;
459 }
460}
461)SOFIE";
462
463// Custom Clad forward-mode pushforwards for the same helpers, following
464// clad's convention for void functions: the original parameters followed by
465// same-type tangent clones of every parameter. Clad computes Hessians as
466// reverse-mode derivatives of forward-mode code, so the bodies below are
467// themselves reverse-differentiated and must consist of plain loops and of
468// calls with custom pullbacks only (for the namespace alias on the
469// Gemm_Call calls, see the comment at the emission site). Note that
470// Gemm_Call_pullback silently bails out for alpha != 1 (see the TODO there),
471// so Hessians inherit that restriction.
472constexpr const char *kGemmCallPushforward = R"SOFIE(
473inline void Gemm_Call_pushforward(float *output, bool transa, bool transb, int m, int n, int k, float alpha,
474 const float *A, const float *B, float beta, const float *C, float *_d_output, bool,
475 bool, int, int, int, float _d_alpha, const float *_d_A, const float *_d_B,
476 float _d_beta, const float *_d_C)
477{
478 // Primal: output = alpha * op(A) op(B) + beta * C, where a null C means
479 // "accumulate onto the existing output" (see Gemm_Call).
480 // Tangent:
481 // d_output = alpha * (op(dA) op(B) + op(A) op(dB)) + beta * dC
482 // + d_alpha * op(A) op(B) + d_beta * C
483 // with (C, dC) read as (output, d_output) in accumulate mode. The tangent
484 // is therefore computed first, while output and d_output still hold the
485 // values the primal call overwrites. The pointer null-checks are hoisted
486 // into bools: with a pointer-typed ternary condition, Clad's reverse pass
487 // (as of v2.4) hoists and tapes the condition with the const qualifier
488 // dropped, generating code that does not compile.
489 const bool hasC = C != nullptr;
490 const bool hasdC = _d_C != nullptr;
491 SOFIE_MODEL_NS::Gemm_Call(_d_output, transa, transb, m, n, k, alpha, _d_A, B, (hasC && !hasdC) ? 0.0f : beta, _d_C);
492 SOFIE_MODEL_NS::Gemm_Call(_d_output, transa, transb, m, n, k, alpha, A, _d_B, 1.0f, nullptr);
493 if (_d_alpha != 0.0f) {
494 SOFIE_MODEL_NS::Gemm_Call(_d_output, transa, transb, m, n, k, _d_alpha, A, B, 1.0f, nullptr);
495 }
496 if (_d_beta != 0.0f) {
497 if (hasC) {
498 for (int i = 0; i < m * n; ++i) {
499 _d_output[i] += _d_beta * C[i];
500 }
501 } else {
502 for (int i = 0; i < m * n; ++i) {
503 _d_output[i] += _d_beta * output[i];
504 }
505 }
506 }
507 SOFIE_MODEL_NS::Gemm_Call(output, transa, transb, m, n, k, alpha, A, B, beta, C);
508}
509)SOFIE";
510
511constexpr const char *kCopyPushforward = R"SOFIE(
512inline void Copy_pushforward(float *output, const float *input, int size, float *_d_output, const float *_d_input,
513 int)
514{
515 for (int i = 0; i < size; i++) {
516 output[i] = input[i];
517 _d_output[i] = _d_input[i];
518 }
519}
520)SOFIE";
521
522constexpr const char *kFillPushforward = R"SOFIE(
523inline void Fill_pushforward(float *output, float value, int size, float *_d_output, float _d_value, int)
524{
525 for (int i = 0; i < size; i++) {
526 output[i] = value;
527 _d_output[i] = _d_value;
528 }
529}
530)SOFIE";
531
532constexpr const char *kReluPushforward = R"SOFIE(
533inline void Relu_pushforward(float *output, float const *input, int size, float *_d_output, float const *_d_input,
534 int)
535{
536 // Tangent first: the generated code applies Relu in place, so input/output
537 // (and their tangents) may alias.
538 for (int i = 0; i < size; i++) {
539 _d_output[i] = (input[i] > 0.0f) ? _d_input[i] : 0.0f;
540 output[i] = (input[i] > 0.0f) ? input[i] : 0.0f;
541 }
542}
543)SOFIE";
544
545constexpr const char *kRelu = R"SOFIE(
546inline void Relu(float *output, float const *input, int size)
547{
548 for (int i = 0; i < size; i++) {
549 output[i] = (input[i] > 0.0f) ? input[i] : 0.0f;
550 }
551}
552)SOFIE";
553
554constexpr const char *kFill = R"SOFIE(
555inline void Fill(float *output, float value, int size)
556{
557 std::fill(output, output + size, value);
558}
559)SOFIE";
560
561constexpr const char *kCopy = R"SOFIE(
562template <class T>
563inline void Copy(T *output, T const *input, int size)
564{
565 std::copy(input, input + size, output);
566}
567)SOFIE";
568
569constexpr const char *kReadTensorFromStream = R"SOFIE(
570inline float ParseFloatToken(const std::string &s)
571{
572 if (s == "inf")
573 return std::numeric_limits<float>::infinity();
574 if (s == "-inf")
575 return -std::numeric_limits<float>::infinity();
576 if (s == "nan")
577 return std::numeric_limits<float>::quiet_NaN();
578 return std::stof(s);
579}
580
581template <class T>
582void ReadTensorFromStream(std::istream &is, T &target, std::string const &expectedName, std::size_t expectedLength)
583{
584 std::string name;
585 std::size_t length;
586 is >> name >> length;
587 if (name != expectedName) {
588 std::string err_msg =
589 "TMVA-SOFIE failed to read the correct tensor name; expected name is " + expectedName + " , read " + name;
590 throw std::runtime_error(err_msg);
591 }
592 if (length != expectedLength) {
593 std::string err_msg = "TMVA-SOFIE failed to read the correct tensor size; expected size is " +
594 std::to_string(expectedLength) + " , read " + std::to_string(length);
595 throw std::runtime_error(err_msg);
596 }
597 std::string token;
598 for (std::size_t i = 0; i < length; ++i) {
599 is >> token;
600 target[i] = ParseFloatToken(token);
601 }
602 if (is.fail()) {
603 throw std::runtime_error("TMVA-SOFIE failed to read the values for tensor " + expectedName);
604 }
605}
606)SOFIE";
607
608// Writer and reader support for the safetensors weight format
609// (https://huggingface.co/docs/safetensors): an 8-byte little-endian JSON
610// header size, the JSON header mapping tensor names to
611// {"dtype", "shape", "data_offsets"}, then the concatenated raw little-endian
612// payloads. data_offsets are relative to the start of the payload area, i.e.
613// just after the JSON header.
614// The header is scanned with a small self-contained streaming parser (no JSON
615// value tree is built), keeping the generated header free of any third-party
616// dependency.
617
618// Non-owning view over a safetensors weight payload, either loaded from a
619// file or provided directly from memory.
620constexpr const char *kSafetensorsBlob = R"SOFIE(
621// Non-owning view over a safetensors weight payload, either loaded from a
622// file or provided directly from memory. The payload is little-endian per the
623// safetensors specification.
624struct SafetensorsBlob {
625 const char *data = nullptr;
626 std::size_t size = 0;
627 constexpr SafetensorsBlob() = default;
628 constexpr SafetensorsBlob(const char *d, std::size_t n) : data(d), size(n) {}
629 SafetensorsBlob(const std::string &s) : data(s.data()), size(s.size()) {}
630};
631)SOFIE";
632
633constexpr const char *kSafetensorsReader = R"SOFIE(
634// Bookkeeping of one tensor inside a safetensors blob
635struct SafetensorsTensorInfo {
636 std::uint64_t fBegin = 0; // payload offset of the tensor in the blob
637 std::uint64_t fEnd = 0;
638 std::string fDtype;
639};
640
641// Scanner for the JSON header of a safetensors blob. The header structure (a
642// flat object of {"dtype", "shape", "data_offsets"} tensor entries plus an
643// optional "__metadata__" member) is fixed, so the header is streamed straight
644// into the tensor index instead of being parsed into a JSON value tree.
645// Content that is not needed ("shape", "__metadata__", unknown members) is
646// still scanned completely, so malformed headers are rejected rather than
647// silently skipped.
648class SafetensorsHeaderScanner {
649public:
650 SafetensorsHeaderScanner(std::string_view text, std::uint64_t dataBegin, std::uint64_t payloadBytes)
651 : fText(text), fDataBegin(dataBegin), fPayloadBytes(payloadBytes)
652 {
653 }
654
655 void Scan(std::unordered_map<std::string, SafetensorsTensorInfo> &tensors)
656 {
657 SkipWhitespace();
658 Expect('{');
659 SkipWhitespace();
660 if (!TryConsume('}')) {
661 while (true) {
662 SkipWhitespace();
663 if (Peek() != '"')
664 Fail("expected tensor name");
665 const std::string name = ParseString(true);
666 SkipWhitespace();
667 Expect(':');
668 SkipWhitespace();
669 if (name == "__metadata__") {
670 SkipValue();
671 } else {
672 ParseTensorEntry(name, tensors);
673 }
674 SkipWhitespace();
675 if (TryConsume('}'))
676 break;
677 Expect(',');
678 }
679 }
680 SkipWhitespace();
681 if (fPos != fText.size())
682 Fail("unexpected trailing content");
683 }
684
685private:
686 std::string_view fText;
687 std::size_t fPos = 0;
688 int fDepth = 0; // recursion budget for skipped content
689 std::uint64_t fDataBegin; // absolute offset of the payload area in the blob
690 std::uint64_t fPayloadBytes; // size of the payload area
691
692 [[noreturn]] void Fail(const std::string &message)
693 {
694 throw std::runtime_error("tmva-sofie: malformed JSON in safetensors header at offset " +
695 std::to_string(fPos) + " : " + message);
696 }
697
698 void SkipWhitespace()
699 {
700 while (fPos < fText.size()) {
701 const char c = fText[fPos];
702 if (c == ' ' || c == '\t' || c == '\n' || c == '\r')
703 ++fPos;
704 else
705 break;
706 }
707 }
708
709 char Peek()
710 {
711 if (fPos >= fText.size())
712 Fail("unexpected end of input");
713 return fText[fPos];
714 }
715
716 void Expect(char c)
717 {
718 if (Peek() != c) {
719 std::string msg = "expected '";
720 msg += c;
721 msg += "'";
722 Fail(msg);
723 }
724 ++fPos;
725 }
726
727 bool TryConsume(char c)
728 {
729 if (fPos < fText.size() && fText[fPos] == c) {
730 ++fPos;
731 return true;
732 }
733 return false;
734 }
735
736 void ExpectLiteral(const char *literal)
737 {
738 const std::size_t len = std::string_view(literal).size();
739 if (fText.compare(fPos, len, literal) != 0)
740 Fail(std::string("expected '") + literal + "'");
741 fPos += len;
742 }
743
744 // JSON string. Decoding only matters for the strings that are kept (tensor
745 // names and dtypes); those are ASCII in practice, so \u escapes above 0x7F
746 // are rejected rather than decoded (such a name could never match the
747 // tensor names the generated code looks up). In skip mode only the scan
748 // position matters.
749 std::string ParseString(bool keep)
750 {
751 Expect('"');
752 std::string out;
753 while (true) {
754 if (fPos >= fText.size())
755 Fail("unterminated string");
756 const char c = fText[fPos++];
757 if (c == '"')
758 return out;
759 if (c != '\\') {
760 if (static_cast<unsigned char>(c) < 0x20)
761 Fail("unescaped control character in string");
762 if (keep)
763 out.push_back(c);
764 continue;
765 }
766 if (fPos >= fText.size())
767 Fail("truncated escape sequence");
768 const char esc = fText[fPos++];
769 std::uint32_t cp;
770 switch (esc) {
771 case '"': cp = '"'; break;
772 case '\\': cp = '\\'; break;
773 case '/': cp = '/'; break;
774 case 'b': cp = '\b'; break;
775 case 'f': cp = '\f'; break;
776 case 'n': cp = '\n'; break;
777 case 'r': cp = '\r'; break;
778 case 't': cp = '\t'; break;
779 case 'u': cp = ParseHexEscape(); break;
780 default: Fail("invalid escape sequence");
781 }
782 if (keep) {
783 if (cp >= 0x80)
784 Fail("unsupported non-ASCII escape in tensor name");
785 out.push_back(static_cast<char>(cp));
786 }
787 }
788 }
789
790 std::uint32_t ParseHexEscape()
791 {
792 if (fPos + 4 > fText.size())
793 Fail("truncated \\u escape");
794 std::uint32_t cp = 0;
795 for (int i = 0; i < 4; ++i) {
796 const char c = fText[fPos++];
797 cp <<= 4;
798 if (c >= '0' && c <= '9')
799 cp |= static_cast<std::uint32_t>(c - '0');
800 else if (c >= 'a' && c <= 'f')
801 cp |= static_cast<std::uint32_t>(c - 'a' + 10);
802 else if (c >= 'A' && c <= 'F')
803 cp |= static_cast<std::uint32_t>(c - 'A' + 10);
804 else
805 Fail("invalid hex digit in \\u escape");
806 }
807 return cp;
808 }
809
810 // Non-negative decimal integer; the safetensors specification fixes the
811 // data_offsets grammar to plain integers.
812 std::uint64_t ParseUInt()
813 {
814 if (fPos >= fText.size() || fText[fPos] < '0' || fText[fPos] > '9')
815 Fail("expected a non-negative integer");
816 std::uint64_t v = 0;
817 while (fPos < fText.size() && fText[fPos] >= '0' && fText[fPos] <= '9') {
818 const std::uint64_t d = std::uint64_t(fText[fPos] - '0');
819 if (v > (~std::uint64_t(0) - d) / 10)
820 Fail("integer out of range");
821 v = 10 * v + d;
822 ++fPos;
823 }
824 if (fPos < fText.size() && (fText[fPos] == '.' || fText[fPos] == 'e' || fText[fPos] == 'E'))
825 Fail("expected a non-negative integer");
826 return v;
827 }
828
829 void ParseTensorEntry(const std::string &name, std::unordered_map<std::string, SafetensorsTensorInfo> &tensors)
830 {
831 Expect('{');
832 SkipWhitespace();
833 std::string dtype;
834 std::uint64_t begin = 0;
835 std::uint64_t end = 0;
836 bool haveDtype = false;
837 bool haveOffsets = false;
838 if (!TryConsume('}')) {
839 while (true) {
840 SkipWhitespace();
841 if (Peek() != '"')
842 Fail("expected member name in tensor entry");
843 const std::string member = ParseString(true);
844 SkipWhitespace();
845 Expect(':');
846 SkipWhitespace();
847 if (member == "dtype") {
848 dtype = ParseString(true);
849 haveDtype = true;
850 } else if (member == "data_offsets") {
851 Expect('[');
852 SkipWhitespace();
853 begin = ParseUInt();
854 SkipWhitespace();
855 Expect(',');
856 SkipWhitespace();
857 end = ParseUInt();
858 SkipWhitespace();
859 Expect(']');
860 haveOffsets = true;
861 } else {
862 SkipValue();
863 }
864 SkipWhitespace();
865 if (TryConsume('}'))
866 break;
867 Expect(',');
868 }
869 }
870 if (!haveDtype || !haveOffsets)
871 throw std::runtime_error("tmva-sofie: invalid entry for tensor " + name + " in safetensors blob");
872 // Validate against the payload area before computing blob offsets, so
873 // that the additions below cannot wrap around.
874 if (begin > end || end > fPayloadBytes)
875 throw std::runtime_error("tmva-sofie: invalid data offsets for tensor " + name + " in safetensors blob");
876 tensors.emplace(name, SafetensorsTensorInfo{fDataBegin + begin, fDataBegin + end, dtype});
877 }
878
879 void SkipNumber()
880 {
881 // the value of skipped content does not matter, only the scan position
882 const std::size_t begin = fPos;
883 while (fPos < fText.size()) {
884 const char c = fText[fPos];
885 if ((c >= '0' && c <= '9') || c == '-' || c == '+' || c == '.' || c == 'e' || c == 'E')
886 ++fPos;
887 else
888 break;
889 }
890 if (fPos == begin)
891 Fail("unexpected character");
892 }
893
894 void SkipValue()
895 {
896 // A real safetensors header nests at most ~4 levels deep; the depth
897 // budget stops deeply nested malformed headers from exhausting the stack
898 if (++fDepth > 64)
899 Fail("JSON nested too deeply");
900 SkipWhitespace();
901 const char c = Peek();
902 if (c == '{') {
903 ++fPos;
904 SkipWhitespace();
905 if (!TryConsume('}')) {
906 while (true) {
907 SkipWhitespace();
908 if (Peek() != '"')
909 Fail("expected object key");
910 ParseString(false);
911 SkipWhitespace();
912 Expect(':');
913 SkipValue();
914 SkipWhitespace();
915 if (TryConsume('}'))
916 break;
917 Expect(',');
918 }
919 }
920 } else if (c == '[') {
921 ++fPos;
922 SkipWhitespace();
923 if (!TryConsume(']')) {
924 while (true) {
925 SkipValue();
926 SkipWhitespace();
927 if (TryConsume(']'))
928 break;
929 Expect(',');
930 }
931 }
932 } else if (c == '"') {
933 ParseString(false);
934 } else if (c == 't') {
935 ExpectLiteral("true");
936 } else if (c == 'f') {
937 ExpectLiteral("false");
938 } else if (c == 'n') {
939 ExpectLiteral("null");
940 } else {
941 SkipNumber();
942 }
943 --fDepth;
944 }
945};
946
947class SafetensorsReader {
948public:
949 explicit SafetensorsReader(SafetensorsBlob blob) : fData(blob.data), fSize(blob.size)
950 {
951 // safetensors payloads are little-endian by specification
952 const std::uint16_t one = 1;
953 if (!*reinterpret_cast<const std::uint8_t *>(&one))
954 throw std::runtime_error("tmva-sofie: safetensors weights can only be read on a little-endian host");
955
956 if (fSize < 8)
957 throw std::runtime_error("tmva-sofie: truncated safetensors blob: missing JSON header size");
958 std::uint64_t headerSize = 0;
959 for (int i = 0; i < 8; ++i)
960 headerSize |= std::uint64_t(static_cast<unsigned char>(fData[i])) << (8 * i);
961 // safeguard against absurd allocations on a corrupted payload
962 if (headerSize > (std::uint64_t(1) << 30))
963 throw std::runtime_error("tmva-sofie: invalid JSON header size in safetensors blob");
964 if (headerSize > fSize - 8)
965 throw std::runtime_error("tmva-sofie: truncated JSON header in safetensors blob");
966 const std::uint64_t dataBegin = 8 + headerSize;
967 SafetensorsHeaderScanner(std::string_view(fData + 8, headerSize), dataBegin, fSize - dataBegin).Scan(fTensors);
968 }
969
970 // Copy the payload of tensor `name` into `target`, checking data type and
971 // size against what the generated code expects
972 template <class T>
973 void Read(const std::string &name, std::vector<T> &target, std::size_t expectedLength,
974 const std::string &expectedDtype)
975 {
976 auto it = fTensors.find(name);
977 if (it == fTensors.end())
978 throw std::runtime_error("tmva-sofie: tensor " + name + " not found in safetensors blob");
979 const SafetensorsTensorInfo &info = it->second;
980 if (info.fDtype != expectedDtype)
981 throw std::runtime_error("tmva-sofie: tensor " + name + " in safetensors blob has dtype " + info.fDtype +
982 " , expected " + expectedDtype);
983 const std::uint64_t nbytes = info.fEnd - info.fBegin;
984 if (nbytes != expectedLength * sizeof(T))
985 throw std::runtime_error("tmva-sofie: tensor " + name + " in safetensors blob has " +
986 std::to_string(nbytes) + " bytes , expected " +
987 std::to_string(expectedLength * sizeof(T)));
988 target.resize(expectedLength);
989 std::memcpy(target.data(), fData + info.fBegin, nbytes);
990 }
991
992private:
993 const char *fData;
994 std::size_t fSize;
995 std::unordered_map<std::string, SafetensorsTensorInfo> fTensors;
996};
997)SOFIE";
998
999// Constexpr helpers carrying static/symbolic shape metadata for the model's
1000// input tensors into the emitted code.
1001constexpr const char *kInputTensorDims = R"SOFIE(
1002struct SingleDim {
1003 enum class Kind { Static, Symbolic };
1004 Kind kind;
1005 std::size_t dim;
1006 std::string_view name;
1007 constexpr SingleDim(std::size_t v) : kind(Kind::Static), dim(v), name() {}
1008 constexpr SingleDim(const char *v) : kind(Kind::Symbolic), dim(0), name(v) {}
1009};
1010
1011struct TensorDims {
1012 const SingleDim *data;
1013 std::size_t size;
1014 constexpr std::size_t total_size() const
1015 {
1016 std::size_t result = 1;
1017 for (std::size_t i = 0; i < size; ++i) {
1018 result *= data[i].dim;
1019 }
1020 return result;
1021 }
1022};
1023
1024template <class Arr>
1025constexpr TensorDims makeDims(Arr const &arr)
1026{
1027 return TensorDims{arr.data(), arr.size()};
1028}
1029)SOFIE";
1030
1031constexpr const char *kDynamicMemory = R"SOFIE(
1032struct TensorLifeInfo {
1033 int begin; // start time (operator index) of the tensor's lifetime
1034 int end; // end time (operator index)
1035 std::size_t size; // size in bytes
1036};
1037
1038struct MemoryResult {
1039 std::size_t total_bytes = 0; // total memory needed
1040 std::vector<std::size_t> offsets; // resulting offset for each tensor
1041};
1042
1043namespace memory_detail {
1044struct FreeBlock {
1045 std::size_t offset;
1046 std::size_t size;
1047 // order by offset for deterministic coalescing
1048 bool operator<(const FreeBlock &other) const { return offset < other.offset; }
1049};
1050struct MemoryEvent {
1051 int t; // time (operator index)
1052 int type; // 0 = END, 1 = START
1053 int idx; // tensor index
1054 bool operator<(const MemoryEvent &o) const
1055 {
1056 if (t != o.t)
1057 return t < o.t;
1058 return type < o.type; // END before START at the same time
1059 }
1060};
1061} // namespace memory_detail
1062
1063// Greedy best-fit planner with a coalescing free list.
1064inline MemoryResult OrganizeMemory(const std::vector<TensorLifeInfo> &tensorsInfo)
1065{
1066 using memory_detail::FreeBlock;
1067 using memory_detail::MemoryEvent;
1068 for (const auto &t : tensorsInfo) {
1069 if (!(t.end > t.begin)) {
1070 throw std::runtime_error("Each tensor must have end > begin.");
1071 }
1072 }
1073
1074 std::vector<MemoryEvent> events;
1075 events.reserve(tensorsInfo.size() * 2);
1076 for (int i = 0; i < (int)tensorsInfo.size(); ++i) {
1077 events.push_back({tensorsInfo[i].end, 0, i});
1078 events.push_back({tensorsInfo[i].begin, 1, i});
1079 }
1080 std::sort(events.begin(), events.end());
1081
1082 std::vector<std::size_t> tensorsOffset(tensorsInfo.size());
1083 std::set<FreeBlock> free_list;
1084 std::unordered_map<int, std::size_t> live_size;
1085 std::unordered_map<int, std::size_t> live_offset;
1086 std::size_t total_bytes = 0;
1087
1088 auto allocate_best_fit = [&](std::size_t need) -> std::size_t {
1089 // Smallest free block with size >= need. free_list is ordered by offset, so
1090 // this scans linearly; for very large tensor sets a size-keyed multimap
1091 // would avoid the O(n) scan.
1092 auto best = free_list.end();
1093 for (auto it = free_list.begin(); it != free_list.end(); ++it) {
1094 if (it->size >= need) {
1095 if (best == free_list.end() || it->size < best->size)
1096 best = it;
1097 }
1098 }
1099 if (best != free_list.end()) {
1100 std::size_t off = best->offset;
1101 if (best->size == need) {
1102 free_list.erase(best);
1103 } else {
1104 FreeBlock updated{best->offset + need, best->size - need};
1105 free_list.erase(best);
1106 free_list.insert(updated);
1107 }
1108 return off;
1109 }
1110 std::size_t off = total_bytes;
1111 total_bytes += need;
1112 return off;
1113 };
1114
1115 auto try_coalesce = [&](std::set<FreeBlock>::iterator it) {
1116 if (it != free_list.begin()) {
1117 auto prev = std::prev(it);
1118 if (prev->offset + prev->size == it->offset) {
1119 FreeBlock merged{prev->offset, prev->size + it->size};
1120 free_list.erase(prev);
1121 it = free_list.erase(it);
1122 it = free_list.insert(merged).first;
1123 }
1124 }
1125 auto next = std::next(it);
1126 if (next != free_list.end() && it->offset + it->size == next->offset) {
1127 FreeBlock merged{it->offset, it->size + next->size};
1128 free_list.erase(next);
1129 it = free_list.erase(it);
1130 free_list.insert(merged);
1131 }
1132 };
1133
1134 for (const auto &e : events) {
1135 if (e.type == 0) {
1136 auto it_sz = live_size.find(e.idx);
1137 auto it_off = live_offset.find(e.idx);
1138 if (it_sz != live_size.end() && it_off != live_offset.end()) {
1139 FreeBlock fb{it_off->second, it_sz->second};
1140 auto it = free_list.insert(fb).first;
1141 try_coalesce(it);
1142 live_size.erase(it_sz);
1143 live_offset.erase(it_off);
1144 }
1145 } else {
1146 auto &t = tensorsInfo[e.idx];
1147 std::size_t off = allocate_best_fit(t.size);
1148 tensorsOffset[e.idx] = off;
1149 live_size[e.idx] = t.size;
1150 live_offset[e.idx] = off;
1151 }
1152 }
1153
1154 return MemoryResult{total_bytes, std::move(tensorsOffset)};
1155}
1156)SOFIE";
1157
1158} // anonymous namespace
1159
1160HelperFunctionsCode GenerateHelperFunctionsCode(const std::set<std::string> &neededHelpers,
1161 const std::string &modelNamespace, bool sgemmAlreadyDeclared)
1162{
1163 auto need = [&](const char *key) { return neededHelpers.count(key) > 0; };
1164
1165 const bool im2col = need("Im2col");
1166 const bool im2col3d = need("Im2col_3d");
1167 const bool col2im = need("col2im");
1168 const bool uniBroadcast = need("UnidirectionalBroadcast");
1169 const bool convBias = need("BroadcastConvBias");
1170 const bool gemm = need("Gemm_Call");
1171 const bool relu = need("Relu");
1172 const bool fill = need("Fill");
1173 const bool copy = need("Copy");
1174 const bool readTensor = need("ReadTensorFromStream");
1175 const bool safetensorsBlob = need("SafetensorsBlob");
1176 const bool readSafetensors = need("SafetensorsReader");
1177 const bool inputDims = need("InputTensorDims");
1178 const bool dynMemory = need("DynamicMemory");
1179
1180 const bool im2colFamily = im2col || im2col3d || col2im;
1181 const bool needConvertLength = uniBroadcast || convBias;
1182 const bool needConvertString = convBias;
1183
1184 // ---- collect the required standard headers -----------------------------
1185 std::set<std::string> stdHeaders;
1186 auto addStd = [&](std::initializer_list<const char *> hs) {
1187 for (auto h : hs)
1188 stdHeaders.insert(h);
1189 };
1190
1191 if (im2colFamily || uniBroadcast || convBias || gemm || fill || copy || dynMemory)
1192 addStd({"algorithm"});
1193 if (needConvertLength || uniBroadcast || dynMemory)
1194 addStd({"vector", "cstddef"});
1195 if (needConvertString || convBias)
1196 addStd({"sstream", "string", "stdexcept"});
1197 if (readTensor)
1198 addStd({"string", "istream", "stdexcept", "limits"});
1199 if (safetensorsBlob)
1200 addStd({"cstddef", "string"});
1201 if (readSafetensors)
1202 addStd({"cstdint", "cstring", "string", "string_view", "unordered_map", "stdexcept", "vector"});
1203 if (inputDims)
1204 addStd({"array", "string_view", "cstddef"});
1205 if (dynMemory)
1206 addStd({"set", "unordered_map", "stdexcept", "iterator"});
1207
1208 std::string includes;
1209 for (auto const &h : stdHeaders)
1210 includes += "#include <" + h + ">\n";
1211
1212 // ---- assemble the definitions ------------------------------------------
1213 // The order matters: a definition must precede any non-dependent use of it.
1214 std::string defs;
1215 defs += "\n// --- Standalone SOFIE inference helper functions ---\n";
1216
1217 // sgemm_ declaration for Gemm_Call, unless the caller already emitted one.
1218 if (gemm && !sgemmAlreadyDeclared)
1219 defs += kBlasSgemm;
1220
1221 if (needConvertLength)
1222 defs += kConvertShapeToLength;
1223 if (needConvertString)
1224 defs += kConvertShapeToString;
1225
1226 if (im2colFamily || uniBroadcast || convBias) {
1227 defs += "\nnamespace UTILITY {\n";
1228 if (im2colFamily)
1229 defs += kIsAGeZero;
1230 if (im2col)
1231 defs += kIm2col;
1232 if (im2col3d)
1233 defs += kIm2col3d;
1234 if (col2im)
1235 defs += kCol2im;
1236 if (uniBroadcast)
1237 defs += kBroadcastTensor;
1238 if (convBias)
1239 defs += kBroadcastConvBias;
1240 defs += "} // namespace UTILITY\n";
1241 }
1242
1243 if (gemm)
1244 defs += kGemmCall;
1245 if (relu)
1246 defs += kRelu;
1247 if (fill)
1248 defs += kFill;
1249 if (copy)
1250 defs += kCopy;
1251 if (readTensor)
1252 defs += kReadTensorFromStream;
1253 if (safetensorsBlob)
1254 defs += kSafetensorsBlob;
1255 if (readSafetensors)
1256 defs += kSafetensorsReader;
1257 if (inputDims)
1258 defs += kInputTensorDims;
1259 if (dynMemory)
1260 defs += kDynamicMemory;
1261
1262 defs += "// --- End of SOFIE inference helper functions ---\n\n";
1263
1264 // ---- Clad custom derivatives (pullbacks and pushforwards) --------------
1265 // Some helpers cannot be differentiated automatically by Clad (Gemm_Call
1266 // calls the bodyless BLAS routine sgemm_). For those we emit a hand-written
1267 // pullback (reverse mode) and pushforward (forward mode). Clad looks up
1268 // custom derivatives in
1269 // clad::custom_derivatives::<function-namespace>, so they are placed
1270 // there (mirroring the generated model namespace) rather than next to the
1271 // function. The definitions reference the model's own helpers via a using
1272 // declaration / namespace alias, so no Clad header is pulled in and a user
1273 // who never differentiates the model simply carries unused inline
1274 // functions.
1275 std::string cladDefs;
1276 if (gemm || copy || fill || relu) {
1277 cladDefs += "\nnamespace clad {\nnamespace custom_derivatives {\nnamespace " + modelNamespace + " {\n";
1278 if (gemm) {
1279 // Gemm_Call_pullback calls Gemm_Call, so bring it into scope. The
1280 // pushforward uses the SOFIE_MODEL_NS alias instead: its body gets
1281 // reverse-differentiated by Clad for Hessians, and Clad resolves the
1282 // custom Gemm_Call_pullback for the inner calls only when they do
1283 // not go through a using-declaration.
1284 cladDefs += "using ::" + modelNamespace + "::Gemm_Call;\n";
1285 cladDefs += "namespace SOFIE_MODEL_NS = ::" + modelNamespace + ";\n";
1286 cladDefs += kGemmCallPullback;
1287 cladDefs += kGemmCallPushforward;
1288 }
1289 if (copy) {
1290 cladDefs += kCopyPullback;
1291 cladDefs += kCopyPushforward;
1292 }
1293 if (fill) {
1294 cladDefs += kFillPullback;
1295 cladDefs += kFillPushforward;
1296 }
1297 if (relu) {
1298 cladDefs += kReluPullback;
1299 cladDefs += kReluPushforward;
1300 }
1301 cladDefs += "} // namespace " + modelNamespace + "\n} // namespace custom_derivatives\n} // namespace clad\n";
1302 }
1303
1304 return HelperFunctionsCode{std::move(includes), std::move(defs), std::move(cladDefs)};
1305}
1306
1307} // namespace SOFIE
1308} // namespace Experimental
1309} // namespace TMVA
#define h(i)
Definition RSha256.hxx:106
constexpr const char * kInputTensorDims
constexpr const char * kDynamicMemory
HelperFunctionsCode GenerateHelperFunctionsCode(const std::set< std::string > &neededHelpers, const std::string &modelNamespace, bool sgemmAlreadyDeclared=false)
Return the standalone C++ source of the inference helper functions requested in neededHelpers (see RM...
create variable transformations