GCC Code Coverage Report

Directory: brotensor/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 82.8% 571 / 0 / 690
Functions: 45.5% 10 / 0 / 22
Branches: 47.4% 556 / 0 / 1174
src/cpu/cpu_jit_fused.cpp
Line Branch Exec Source
1 #include "cpu_jit.h"
2 #include <brotensor/detail/cpu/thread_pool.h>
3 #include <algorithm>
4 #include <cmath>
5 #include <cstdint>
6
7 namespace brotensor::detail::cpu::jit {
8 namespace ref {
9
10 ✗ inline void fused_residual_rmsnorm(float* X, const float* res, const float* gamma, float eps, float* Y, int B, int D) {
11 ✗ if (B <= 0 || D <= 0) return;
12 ✗ const float inv_D = 1.0f / static_cast<float>(D);
13 ✗ auto row_fn = [&](int b) {
14 ✗ float* xr = X + static_cast<size_t>(b) * D;
15 ✗ const float* rr = res + static_cast<size_t>(b) * D;
16 ✗ float* yr = Y + static_cast<size_t>(b) * D;
17 ✗ float sum = 0.0f;
18 ✗ for (int i = 0; i < D; ++i) {
19 ✗ xr[i] += rr[i];
20 ✗ sum += xr[i] * xr[i];
21 ✗ }
22 ✗ float rrms = 1.0f / std::sqrt(sum * inv_D + eps);
23 ✗ for (int i = 0; i < D; ++i) {
24 ✗ yr[i] = xr[i] * gamma[i] * rrms;
25 ✗ }
26 ✗ };
27 ✗ if (B > 1 && static_cast<int64_t>(B) * D >= 262144) {
28 ✗ detail::cpu::parallel_for(static_cast<size_t>(B), [&](size_t bi) {
29 ✗ row_fn(static_cast<int>(bi));
30 ✗ });
31 ✗ } else {
32 ✗ for (int b = 0; b < B; ++b) row_fn(b);
33 }
34 ✗ }
35
36 ✗ inline void fused_residual_layernorm(float* X, const float* res, const float* gamma, const float* beta,
37 float eps, float* Y, int B, int D) {
38 ✗ if (B <= 0 || D <= 0) return;
39 ✗ const float inv_D = 1.0f / static_cast<float>(D);
40 ✗ auto row_fn = [&](int b) {
41 ✗ float* xr = X + static_cast<size_t>(b) * D;
42 ✗ const float* rr = res + static_cast<size_t>(b) * D;
43 ✗ float* yr = Y + static_cast<size_t>(b) * D;
44 ✗ float sum = 0.0f;
45 ✗ for (int i = 0; i < D; ++i) {
46 ✗ xr[i] += rr[i];
47 ✗ sum += xr[i];
48 ✗ }
49 ✗ float mean = sum * inv_D;
50 ✗ float sumsq = 0.0f;
51 ✗ for (int i = 0; i < D; ++i) {
52 ✗ float diff = xr[i] - mean;
53 ✗ sumsq += diff * diff;
54 ✗ }
55 ✗ float var = sumsq * inv_D;
56 ✗ float rstd = 1.0f / std::sqrt(var + eps);
57 ✗ for (int i = 0; i < D; ++i) {
58 ✗ float g = gamma ? gamma[i] : 1.0f;
59 ✗ float bv = beta ? beta[i] : 0.0f;
60 ✗ yr[i] = g * ((xr[i] - mean) * rstd) + bv;
61 ✗ }
62 ✗ };
63 ✗ if (B > 1 && static_cast<int64_t>(B) * D >= 262144) {
64 ✗ detail::cpu::parallel_for(static_cast<size_t>(B), [&](size_t bi) {
65 ✗ row_fn(static_cast<int>(bi));
66 ✗ });
67 ✗ } else {
68 ✗ for (int b = 0; b < B; ++b) row_fn(b);
69 }
70 ✗ }
71
72 ✗ inline void fused_layernorm_modulate(const float* X, const float* gamma, const float* beta,
73 const float* scale, const float* shift, float eps, float* Y, int R, int D) {
74 ✗ if (R <= 0 || D <= 0) return;
75 ✗ const float inv_D = 1.0f / static_cast<float>(D);
76 ✗ auto row_fn = [&](int r) {
77 ✗ const float* xr = X + static_cast<size_t>(r) * D;
78 ✗ float* yr = Y + static_cast<size_t>(r) * D;
79 ✗ float sum = 0.0f;
80 ✗ for (int i = 0; i < D; ++i) sum += xr[i];
81 ✗ float mean = sum * inv_D;
82 ✗ float sumsq = 0.0f;
83 ✗ for (int i = 0; i < D; ++i) {
84 ✗ float diff = xr[i] - mean;
85 ✗ sumsq += diff * diff;
86 ✗ }
87 ✗ float var = sumsq * inv_D;
88 ✗ float rstd = 1.0f / std::sqrt(var + eps);
89 ✗ for (int i = 0; i < D; ++i) {
90 ✗ float g = gamma ? gamma[i] : 1.0f;
91 ✗ float bv = beta ? beta[i] : 0.0f;
92 ✗ float norm = g * ((xr[i] - mean) * rstd) + bv;
93 ✗ yr[i] = norm * (1.0f + scale[i]) + shift[i];
94 ✗ }
95 ✗ };
96 ✗ if (R > 1 && static_cast<int64_t>(R) * D >= 262144) {
97 ✗ detail::cpu::parallel_for(static_cast<size_t>(R), [&](size_t ri) {
98 ✗ row_fn(static_cast<int>(ri));
99 ✗ });
100 ✗ } else {
101 ✗ for (int r = 0; r < R; ++r) row_fn(r);
102 }
103 ✗ }
104
105 } // namespace ref
106 } // namespace brotensor::detail::cpu::jit
107
108 #if BROTENSOR_HAS_BRASS_JIT
109
110 #include <brass/codegen/kernel_jit.hpp>
111 #include <brass/mir/builder.hpp>
112 #include <brass/mir/module.hpp>
113 #include <brass/mir/function.hpp>
114 #include <cmath>
115 #include <mutex>
116 #include <vector>
117 #include <memory>
118 #include <algorithm>
119 #include <iostream>
120
121 namespace brotensor::detail::cpu::jit {
122
123 namespace {
124
125 using namespace brass;
126 using namespace brass::codegen;
127
128 struct JitFusedKernels {
129 bool available = false;
130
131 using FusedResidualRmsNormRowFn = void (*)(float* x, const float* res, const float* gamma, float* y, uint64_t d, float eps, float inv_d);
132 FusedResidualRmsNormRowFn fused_residual_rmsnorm_row_fn = nullptr;
133
134 using FusedResidualLayerNormRowFn = void (*)(float* x, const float* res, const float* gamma, const float* beta, float* y, uint64_t d, float eps, float inv_d);
135 FusedResidualLayerNormRowFn fused_residual_layernorm_row_fn = nullptr;
136
137 using FusedLayerNormModulateRowFn = void (*)(const float* x, const float* gamma, const float* beta, const float* scale, const float* shift, float* y, uint64_t d, float eps, float inv_d);
138 FusedLayerNormModulateRowFn fused_layernorm_modulate_row_fn = nullptr;
139
140 std::vector<std::shared_ptr<JitExecutionEngine>> engines;
141 };
142
143 // Horizontal sum of an 8-wide float vector using destination scratch buffer
144 15 static Value* reduce_vsum8(KernelBuilder& kb, Value* scratch, Value* vsum) {
145 15 kb.vstore_f32x8(scratch, vsum);
146 15 Value* s0 = kb.load_f32(scratch, 0); Value* s1 = kb.load_f32(scratch, 4);
147 15 Value* s2 = kb.load_f32(scratch, 8); Value* s3 = kb.load_f32(scratch, 12);
148 15 Value* s4 = kb.load_f32(scratch, 16); Value* s5 = kb.load_f32(scratch, 20);
149 15 Value* s6 = kb.load_f32(scratch, 24); Value* s7 = kb.load_f32(scratch, 28);
150 15 Value* sum0123 = kb.add(kb.add(s0, s1), kb.add(s2, s3));
151 15 Value* sum4567 = kb.add(kb.add(s4, s5), kb.add(s6, s7));
152 15 return kb.add(sum0123, sum4567);
153 }
154
155 // ── 1. Fused Residual RMSNorm Row (Option B: X += res in-place + RMSNorm) ───
156 //
157 // Computes in-place residual addition X += res, writes X, and accumulates sum_sq
158 // in the EXACT same vector register pass, eliminating a full memory read/write pass.
159 3 static KernelFunction build_fused_residual_rmsnorm_row(KernelJit& jit) {
160 3 Module mod("mod_fused_residual_rmsnorm_row");
161
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 mod.add_external_symbol("rsqrtf");
162
9/18
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 3 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 3 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 3 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 3 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 3 times.
✗ Branch 13 not taken.
✓ Branch 14 taken 3 times.
✗ Branch 15 not taken.
✓ Branch 16 taken 3 times.
✗ Branch 17 not taken.
 3 Function* fn = mod.create_function("fused_residual_rmsnorm_row", Type::void_type(), {
163 21 Type::ptr(), Type::ptr(), Type::ptr(), Type::ptr(), Type::i64(), Type::f32(), Type::f32()
164 });
165
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 KernelBuilder kb(mod, fn);
166
167
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* entry = kb.builder().append_block("entry");
168
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* x = kb.builder().add_block_param(entry, Type::ptr());
169
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* res = kb.builder().add_block_param(entry, Type::ptr());
170
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* gamma = kb.builder().add_block_param(entry, Type::ptr());
171
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* y = kb.builder().add_block_param(entry, Type::ptr());
172
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* d = kb.builder().add_block_param(entry, Type::i64());
173
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* eps = kb.builder().add_block_param(entry, Type::f32());
174
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* inv_d = kb.builder().add_block_param(entry, Type::f32());
175
176
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l1_vhead = kb.builder().create_block("res_rms_l1_vhead");
177
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l1_vbody = kb.builder().create_block("res_rms_l1_vbody");
178
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l1_vexit = kb.builder().create_block("res_rms_l1_vexit");
179
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l1_shead = kb.builder().create_block("res_rms_l1_shead");
180
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l1_sbody = kb.builder().create_block("res_rms_l1_sbody");
181
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l1_exit = kb.builder().create_block("res_rms_l1_exit");
182
183
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l2_vhead = kb.builder().create_block("res_rms_l2_vhead");
184
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l2_vbody = kb.builder().create_block("res_rms_l2_vbody");
185
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l2_shead = kb.builder().create_block("res_rms_l2_shead");
186
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l2_sbody = kb.builder().create_block("res_rms_l2_sbody");
187
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l2_exit = kb.builder().create_block("res_rms_l2_exit");
188
189 3 kb.position_at_end(entry);
190
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* zero_i64 = kb.builder().build_iconst_i64(0);
191
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* one_i64 = kb.builder().build_iconst_i64(1);
192
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* two_i64 = kb.builder().build_iconst_i64(2);
193
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* eight_i64 = kb.builder().build_iconst_i64(8);
194
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* zero_f32 = kb.builder().build_fconst_f32(0.0f);
195
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vzero = kb.vzero(Type::f32x8());
196
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* mask_eight = kb.builder().build_iconst_i64(~int64_t(7));
197
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vec_d = kb.builder().build_and(d, mask_eight);
198
199
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* has_vec = kb.builder().build_sge(d, eight_i64);
200
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br_if(has_vec, l1_vhead, {zero_i64, vzero}, l1_shead, {zero_i64, zero_f32});
201
202 // Pass 1: Vector loop X += res in-place AND accumulates sum_sq in same register pass
203
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l1_vhead);
204
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* iv1_v = kb.builder().add_block_param(l1_vhead, Type::i64());
205
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vsum = kb.builder().add_block_param(l1_vhead, Type::f32x8());
206 3 kb.position_at_end(l1_vhead);
207
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* cond1_v = kb.builder().build_slt(iv1_v, vec_d);
208
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br_if(cond1_v, l1_vbody, {}, l1_vexit, {iv1_v, vsum});
209
210
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l1_vbody);
211 3 kb.position_at_end(l1_vbody);
212
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* byte_off1 = kb.builder().build_shl(iv1_v, two_i64);
213
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* px1 = kb.builder().build_add(x, byte_off1);
214
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* pres1 = kb.builder().build_add(res, byte_off1);
215
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vx1 = kb.vload_f32x8(px1);
216
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vr1 = kb.vload_f32x8(pres1);
217
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vx_new = kb.vadd(vx1, vr1);
218
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.vstore_f32x8(px1, vx_new); // X += res in-place
219
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* next_vsum = kb.vfma(vx_new, vx_new, vsum); // sum_sq accumulation in exact same pass
220
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* next_iv1_v = kb.builder().build_add(iv1_v, eight_i64);
221
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br(l1_vhead, {next_iv1_v, next_vsum});
222
223 // Pass 1: Vector exit -> horizontal reduction
224
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l1_vexit);
225
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* exit_iv1 = kb.builder().add_block_param(l1_vexit, Type::i64());
226
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* exit_vsum = kb.builder().add_block_param(l1_vexit, Type::f32x8());
227 3 kb.position_at_end(l1_vexit);
228
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* init_scalar_sum = reduce_vsum8(kb, y, exit_vsum);
229
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br(l1_shead, {exit_iv1, init_scalar_sum});
230
231 // Pass 1: Scalar remainder
232
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l1_shead);
233
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* iv1_s = kb.builder().add_block_param(l1_shead, Type::i64());
234
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* sum_sq = kb.builder().add_block_param(l1_shead, Type::f32());
235 3 kb.position_at_end(l1_shead);
236
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* cond1_s = kb.builder().build_slt(iv1_s, d);
237
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br_if(cond1_s, l1_sbody, {}, l1_exit, {sum_sq});
238
239
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l1_sbody);
240 3 kb.position_at_end(l1_sbody);
241
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* xi = kb.load_f32_indexed(x, iv1_s, 4, 0);
242
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* ri = kb.load_f32_indexed(res, iv1_s, 4, 0);
243
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* xi_new = kb.add(xi, ri);
244
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.store_f32_indexed(x, iv1_s, xi_new, 4, 0);
245
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* xi_sq = kb.mul(xi_new, xi_new);
246
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* next_sum = kb.add(sum_sq, xi_sq);
247
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* next_iv1_s = kb.builder().build_add(iv1_s, one_i64);
248
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br(l1_shead, {next_iv1_s, next_sum});
249
250 // Pass 1 exit -> compute rrms
251
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l1_exit);
252
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* final_sum = kb.builder().add_block_param(l1_exit, Type::f32());
253 3 kb.position_at_end(l1_exit);
254
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* mean_sq = kb.mul(final_sum, inv_d);
255
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* denom = kb.add(mean_sq, eps);
256
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 Value* rrms = kb.builder().build_call("rsqrtf", Type::f32(), {denom});
257
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vrrms = kb.vbroadcast(Type::f32x8(), rrms);
258
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br(l2_vhead, {zero_i64, rrms, vrrms});
259
260 // Pass 2: Vector loop y = x * gamma * rrms
261
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l2_vhead);
262
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* iv2_v = kb.builder().add_block_param(l2_vhead, Type::i64());
263
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* rrms_val = kb.builder().add_block_param(l2_vhead, Type::f32());
264
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vrrms_val = kb.builder().add_block_param(l2_vhead, Type::f32x8());
265 3 kb.position_at_end(l2_vhead);
266
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* cond2_v = kb.builder().build_slt(iv2_v, vec_d);
267
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br_if(cond2_v, l2_vbody, {}, l2_shead, {iv2_v, rrms_val});
268
269
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l2_vbody);
270 3 kb.position_at_end(l2_vbody);
271
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* byte_off2 = kb.builder().build_shl(iv2_v, two_i64);
272
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* px2 = kb.builder().build_add(x, byte_off2);
273
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* pg2 = kb.builder().build_add(gamma, byte_off2);
274
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* py2 = kb.builder().build_add(y, byte_off2);
275
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vx2 = kb.vload_f32x8(px2);
276
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vg2 = kb.vload_f32x8(pg2);
277
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vxg = kb.vmul(vx2, vg2);
278
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vy = kb.vmul(vxg, vrrms_val);
279
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.vstore_f32x8(py2, vy);
280
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* next_iv2_v = kb.builder().build_add(iv2_v, eight_i64);
281
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br(l2_vhead, {next_iv2_v, rrms_val, vrrms_val});
282
283 // Pass 2: Scalar remainder
284
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l2_shead);
285
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* iv2_s = kb.builder().add_block_param(l2_shead, Type::i64());
286
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* rrms_scalar = kb.builder().add_block_param(l2_shead, Type::f32());
287 3 kb.position_at_end(l2_shead);
288
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* cond2_s = kb.builder().build_slt(iv2_s, d);
289
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br_if(cond2_s, l2_sbody, l2_exit);
290
291
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l2_sbody);
292 3 kb.position_at_end(l2_sbody);
293
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* x2 = kb.load_f32_indexed(x, iv2_s, 4, 0);
294
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* g2 = kb.load_f32_indexed(gamma, iv2_s, 4, 0);
295
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* xg2 = kb.mul(x2, g2);
296
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* out_val = kb.mul(xg2, rrms_scalar);
297
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.store_f32_indexed(y, iv2_s, out_val, 4, 0);
298
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* next_iv2_s = kb.builder().build_add(iv2_s, one_i64);
299
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br(l2_shead, {next_iv2_s, rrms_scalar});
300
301
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l2_exit);
302 3 kb.position_at_end(l2_exit);
303
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_ret_void();
304
305
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 return jit.compile(mod, "fused_residual_rmsnorm_row");
306 3 }
307
308 // ── 2. Fused Residual LayerNorm Row (X += res in-place + LayerNorm) ─────────
309 //
310 // Computes in-place residual addition X += res, writes X, and computes mean and
311 // variance across vector registers, writing Y = gamma * (X - mean) * rstd + beta
312 // directly in registers without intermediate buffers.
313 3 static KernelFunction build_fused_residual_layernorm_row(KernelJit& jit) {
314 3 Module mod("mod_fused_residual_layernorm_row");
315
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 mod.add_external_symbol("rsqrtf");
316
10/20
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 3 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 3 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 3 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 3 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 3 times.
✗ Branch 13 not taken.
✓ Branch 14 taken 3 times.
✗ Branch 15 not taken.
✓ Branch 16 taken 3 times.
✗ Branch 17 not taken.
✓ Branch 18 taken 3 times.
✗ Branch 19 not taken.
 3 Function* fn = mod.create_function("fused_residual_layernorm_row", Type::void_type(), {
317 24 Type::ptr(), Type::ptr(), Type::ptr(), Type::ptr(), Type::ptr(), Type::i64(), Type::f32(), Type::f32()
318 });
319
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 KernelBuilder kb(mod, fn);
320
321
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* entry = kb.builder().append_block("entry");
322
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* x = kb.builder().add_block_param(entry, Type::ptr());
323
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* res = kb.builder().add_block_param(entry, Type::ptr());
324
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* gamma = kb.builder().add_block_param(entry, Type::ptr());
325
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* beta = kb.builder().add_block_param(entry, Type::ptr());
326
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* y = kb.builder().add_block_param(entry, Type::ptr());
327
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* d = kb.builder().add_block_param(entry, Type::i64());
328
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* eps = kb.builder().add_block_param(entry, Type::f32());
329
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* inv_d = kb.builder().add_block_param(entry, Type::f32());
330
331 // Pass 1 blocks (residual add + mean)
332
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l1_vhead = kb.builder().create_block("res_ln_l1_vhead");
333
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l1_vbody = kb.builder().create_block("res_ln_l1_vbody");
334
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l1_vexit = kb.builder().create_block("res_ln_l1_vexit");
335
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l1_shead = kb.builder().create_block("res_ln_l1_shead");
336
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l1_sbody = kb.builder().create_block("res_ln_l1_sbody");
337
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l1_exit = kb.builder().create_block("res_ln_l1_exit");
338
339 // Pass 2 blocks (variance)
340
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l2_vhead = kb.builder().create_block("res_ln_l2_vhead");
341
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l2_vbody = kb.builder().create_block("res_ln_l2_vbody");
342
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l2_vexit = kb.builder().create_block("res_ln_l2_vexit");
343
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l2_shead = kb.builder().create_block("res_ln_l2_shead");
344
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l2_sbody = kb.builder().create_block("res_ln_l2_sbody");
345
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l2_exit = kb.builder().create_block("res_ln_l2_exit");
346
347 // Pass 3 blocks (gamma * xhat + beta)
348
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l3_vhead = kb.builder().create_block("res_ln_l3_vhead");
349
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l3_vbody = kb.builder().create_block("res_ln_l3_vbody");
350
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l3_shead = kb.builder().create_block("res_ln_l3_shead");
351
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l3_sbody = kb.builder().create_block("res_ln_l3_sbody");
352
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l3_exit = kb.builder().create_block("res_ln_l3_exit");
353
354 3 kb.position_at_end(entry);
355
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* zero_i64 = kb.builder().build_iconst_i64(0);
356
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* one_i64 = kb.builder().build_iconst_i64(1);
357
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* two_i64 = kb.builder().build_iconst_i64(2);
358
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* eight_i64 = kb.builder().build_iconst_i64(8);
359
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* zero_f32 = kb.builder().build_fconst_f32(0.0f);
360
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vzero = kb.vzero(Type::f32x8());
361
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* mask_eight = kb.builder().build_iconst_i64(~int64_t(7));
362
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vec_d = kb.builder().build_and(d, mask_eight);
363
364
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* has_vec = kb.builder().build_sge(d, eight_i64);
365
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br_if(has_vec, l1_vhead, {zero_i64, vzero}, l1_shead, {zero_i64, zero_f32});
366
367 // Pass 1: Vector sum + in-place residual addition X += res
368
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l1_vhead);
369
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* iv1_v = kb.builder().add_block_param(l1_vhead, Type::i64());
370
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vsum1 = kb.builder().add_block_param(l1_vhead, Type::f32x8());
371 3 kb.position_at_end(l1_vhead);
372
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* cond1_v = kb.builder().build_slt(iv1_v, vec_d);
373
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br_if(cond1_v, l1_vbody, {}, l1_vexit, {iv1_v, vsum1});
374
375
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l1_vbody);
376 3 kb.position_at_end(l1_vbody);
377
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* byte_off1 = kb.builder().build_shl(iv1_v, two_i64);
378
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* px1 = kb.builder().build_add(x, byte_off1);
379
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* pres1 = kb.builder().build_add(res, byte_off1);
380
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vx1 = kb.vload_f32x8(px1);
381
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vr1 = kb.vload_f32x8(pres1);
382
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vx_new = kb.vadd(vx1, vr1);
383
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.vstore_f32x8(px1, vx_new); // X += res in-place!
384
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* next_vsum1 = kb.vadd(vsum1, vx_new);
385
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* next_iv1_v = kb.builder().build_add(iv1_v, eight_i64);
386
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br(l1_vhead, {next_iv1_v, next_vsum1});
387
388 // Pass 1: Vector exit -> horizontal sum
389
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l1_vexit);
390
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* exit_iv1 = kb.builder().add_block_param(l1_vexit, Type::i64());
391
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* exit_vsum1 = kb.builder().add_block_param(l1_vexit, Type::f32x8());
392 3 kb.position_at_end(l1_vexit);
393
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* init_sum1 = reduce_vsum8(kb, y, exit_vsum1);
394
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br(l1_shead, {exit_iv1, init_sum1});
395
396 // Pass 1: Scalar remainder
397
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l1_shead);
398
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* iv1_s = kb.builder().add_block_param(l1_shead, Type::i64());
399
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* sum1 = kb.builder().add_block_param(l1_shead, Type::f32());
400 3 kb.position_at_end(l1_shead);
401
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* cond1_s = kb.builder().build_slt(iv1_s, d);
402
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br_if(cond1_s, l1_sbody, {}, l1_exit, {sum1});
403
404
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l1_sbody);
405 3 kb.position_at_end(l1_sbody);
406
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* xi1 = kb.load_f32_indexed(x, iv1_s, 4, 0);
407
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* ri1 = kb.load_f32_indexed(res, iv1_s, 4, 0);
408
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* xi_new = kb.add(xi1, ri1);
409
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.store_f32_indexed(x, iv1_s, xi_new, 4, 0); // X += res in-place!
410
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* next_sum1 = kb.add(sum1, xi_new);
411
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* next_iv1_s = kb.builder().build_add(iv1_s, one_i64);
412
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br(l1_shead, {next_iv1_s, next_sum1});
413
414 // Pass 1: Mean computed -> enter Pass 2
415
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l1_exit);
416
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* final_sum1 = kb.builder().add_block_param(l1_exit, Type::f32());
417 3 kb.position_at_end(l1_exit);
418
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* mean = kb.mul(final_sum1, inv_d);
419
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vmean = kb.vbroadcast(Type::f32x8(), mean);
420
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br_if(has_vec, l2_vhead, {zero_i64, vzero, mean, vmean}, l2_shead, {zero_i64, zero_f32, mean});
421
422 // Pass 2: Vector sum of squared deviations
423
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l2_vhead);
424
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* iv2_v = kb.builder().add_block_param(l2_vhead, Type::i64());
425
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vsum2 = kb.builder().add_block_param(l2_vhead, Type::f32x8());
426
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* m2_scalar = kb.builder().add_block_param(l2_vhead, Type::f32());
427
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vm2 = kb.builder().add_block_param(l2_vhead, Type::f32x8());
428 3 kb.position_at_end(l2_vhead);
429
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* cond2_v = kb.builder().build_slt(iv2_v, vec_d);
430
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br_if(cond2_v, l2_vbody, {}, l2_vexit, {iv2_v, vsum2, m2_scalar});
431
432
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l2_vbody);
433 3 kb.position_at_end(l2_vbody);
434
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* byte_off2 = kb.builder().build_shl(iv2_v, two_i64);
435
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* px2 = kb.builder().build_add(x, byte_off2);
436
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vx2 = kb.vload_f32x8(px2);
437
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vdev = kb.vsub(vx2, vm2);
438
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* next_vsum2 = kb.vfma(vdev, vdev, vsum2);
439
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* next_iv2_v = kb.builder().build_add(iv2_v, eight_i64);
440
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br(l2_vhead, {next_iv2_v, next_vsum2, m2_scalar, vm2});
441
442 // Pass 2: Vector exit -> horizontal sum
443
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l2_vexit);
444
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* exit_iv2 = kb.builder().add_block_param(l2_vexit, Type::i64());
445
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* exit_vsum2 = kb.builder().add_block_param(l2_vexit, Type::f32x8());
446
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* m2_out = kb.builder().add_block_param(l2_vexit, Type::f32());
447 3 kb.position_at_end(l2_vexit);
448
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* init_sumsq = reduce_vsum8(kb, y, exit_vsum2);
449
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br(l2_shead, {exit_iv2, init_sumsq, m2_out});
450
451 // Pass 2: Scalar remainder
452
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l2_shead);
453
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* iv2_s = kb.builder().add_block_param(l2_shead, Type::i64());
454
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* sumsq = kb.builder().add_block_param(l2_shead, Type::f32());
455
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* m2_s = kb.builder().add_block_param(l2_shead, Type::f32());
456 3 kb.position_at_end(l2_shead);
457
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* cond2_s = kb.builder().build_slt(iv2_s, d);
458
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br_if(cond2_s, l2_sbody, {}, l2_exit, {sumsq, m2_s});
459
460
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l2_sbody);
461 3 kb.position_at_end(l2_sbody);
462
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* xi2 = kb.load_f32_indexed(x, iv2_s, 4, 0);
463
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* diff2 = kb.sub(xi2, m2_s);
464
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* diff_sq2 = kb.mul(diff2, diff2);
465
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* next_sumsq = kb.add(sumsq, diff_sq2);
466
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* next_iv2_s = kb.builder().build_add(iv2_s, one_i64);
467
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br(l2_shead, {next_iv2_s, next_sumsq, m2_s});
468
469 // Pass 2 exit -> compute rstd
470
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l2_exit);
471
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* final_sumsq = kb.builder().add_block_param(l2_exit, Type::f32());
472
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* final_mean = kb.builder().add_block_param(l2_exit, Type::f32());
473 3 kb.position_at_end(l2_exit);
474
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* var = kb.mul(final_sumsq, inv_d);
475
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* var_eps = kb.add(var, eps);
476
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 Value* rstd = kb.builder().build_call("rsqrtf", Type::f32(), {var_eps});
477
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vmean3 = kb.vbroadcast(Type::f32x8(), final_mean);
478
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vrstd3 = kb.vbroadcast(Type::f32x8(), rstd);
479
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br(l3_vhead, {zero_i64, final_mean, rstd, vmean3, vrstd3});
480
481 // Pass 3: Vector LayerNorm Y = gamma * (X - mean) * rstd + beta
482
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l3_vhead);
483
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* iv3_v = kb.builder().add_block_param(l3_vhead, Type::i64());
484
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* m3_s = kb.builder().add_block_param(l3_vhead, Type::f32());
485
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* r3_s = kb.builder().add_block_param(l3_vhead, Type::f32());
486
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vm3 = kb.builder().add_block_param(l3_vhead, Type::f32x8());
487
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vr3 = kb.builder().add_block_param(l3_vhead, Type::f32x8());
488 3 kb.position_at_end(l3_vhead);
489
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* cond3_v = kb.builder().build_slt(iv3_v, vec_d);
490
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br_if(cond3_v, l3_vbody, {}, l3_shead, {iv3_v, m3_s, r3_s});
491
492
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l3_vbody);
493 3 kb.position_at_end(l3_vbody);
494
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* byte_off3 = kb.builder().build_shl(iv3_v, two_i64);
495
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* px3 = kb.builder().build_add(x, byte_off3);
496
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* pg3 = kb.builder().build_add(gamma, byte_off3);
497
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* pb3 = kb.builder().build_add(beta, byte_off3);
498
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* py3 = kb.builder().build_add(y, byte_off3);
499
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vx3 = kb.vload_f32x8(px3);
500
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vg3 = kb.vload_f32x8(pg3);
501
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vb3 = kb.vload_f32x8(pb3);
502
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vdiff3 = kb.vsub(vx3, vm3);
503
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vxhat3 = kb.vmul(vdiff3, vr3);
504
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vy3 = kb.vfma(vg3, vxhat3, vb3); // y = gamma * xhat + beta
505
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.vstore_f32x8(py3, vy3);
506
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* next_iv3_v = kb.builder().build_add(iv3_v, eight_i64);
507
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br(l3_vhead, {next_iv3_v, m3_s, r3_s, vm3, vr3});
508
509 // Pass 3: Scalar remainder
510
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l3_shead);
511
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* iv3_s = kb.builder().add_block_param(l3_shead, Type::i64());
512
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* m3_sc = kb.builder().add_block_param(l3_shead, Type::f32());
513
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* r3_sc = kb.builder().add_block_param(l3_shead, Type::f32());
514 3 kb.position_at_end(l3_shead);
515
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* cond3_s = kb.builder().build_slt(iv3_s, d);
516
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br_if(cond3_s, l3_sbody, l3_exit);
517
518
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l3_sbody);
519 3 kb.position_at_end(l3_sbody);
520
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* x3 = kb.load_f32_indexed(x, iv3_s, 4, 0);
521
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* g3 = kb.load_f32_indexed(gamma, iv3_s, 4, 0);
522
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* b3 = kb.load_f32_indexed(beta, iv3_s, 4, 0);
523
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* diff3 = kb.sub(x3, m3_sc);
524
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* xhat3 = kb.mul(diff3, r3_sc);
525
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* y3 = kb.builder().build_fma_f32(g3, xhat3, b3);
526
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.store_f32_indexed(y, iv3_s, y3, 4, 0);
527
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* next_iv3_s = kb.builder().build_add(iv3_s, one_i64);
528
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br(l3_shead, {next_iv3_s, m3_sc, r3_sc});
529
530
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l3_exit);
531 3 kb.position_at_end(l3_exit);
532
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_ret_void();
533
534
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 return jit.compile(mod, "fused_residual_layernorm_row");
535 3 }
536
537 // ── 3. Fused LayerNorm + Modulate Row (Option B: Zero Intermediate Memory) ──
538 //
539 // Computes LayerNorm directly into registers and immediately applies AdaLN
540 // modulation by (1 + scale) + shift into output Y, writing zero intermediate
541 // tensors to memory.
542 3 static KernelFunction build_fused_layernorm_modulate_row(KernelJit& jit) {
543 3 Module mod("mod_fused_layernorm_modulate_row");
544
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 mod.add_external_symbol("rsqrtf");
545
11/22
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 3 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 3 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 3 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 3 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 3 times.
✗ Branch 13 not taken.
✓ Branch 14 taken 3 times.
✗ Branch 15 not taken.
✓ Branch 16 taken 3 times.
✗ Branch 17 not taken.
✓ Branch 18 taken 3 times.
✗ Branch 19 not taken.
✓ Branch 20 taken 3 times.
✗ Branch 21 not taken.
 3 Function* fn = mod.create_function("fused_layernorm_modulate_row", Type::void_type(), {
546 27 Type::ptr(), Type::ptr(), Type::ptr(), Type::ptr(), Type::ptr(), Type::ptr(), Type::i64(), Type::f32(), Type::f32()
547 });
548
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 KernelBuilder kb(mod, fn);
549
550
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* entry = kb.builder().append_block("entry");
551
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* x = kb.builder().add_block_param(entry, Type::ptr());
552
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* gamma = kb.builder().add_block_param(entry, Type::ptr());
553
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* beta = kb.builder().add_block_param(entry, Type::ptr());
554
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* scale = kb.builder().add_block_param(entry, Type::ptr());
555
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* shift = kb.builder().add_block_param(entry, Type::ptr());
556
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* y = kb.builder().add_block_param(entry, Type::ptr());
557
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* d = kb.builder().add_block_param(entry, Type::i64());
558
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* eps = kb.builder().add_block_param(entry, Type::f32());
559
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* inv_d = kb.builder().add_block_param(entry, Type::f32());
560
561 // Pass 1 blocks (mean)
562
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l1_vhead = kb.builder().create_block("flm_l1_vhead");
563
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l1_vbody = kb.builder().create_block("flm_l1_vbody");
564
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l1_vexit = kb.builder().create_block("flm_l1_vexit");
565
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l1_shead = kb.builder().create_block("flm_l1_shead");
566
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l1_sbody = kb.builder().create_block("flm_l1_sbody");
567
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l1_exit = kb.builder().create_block("flm_l1_exit");
568
569 // Pass 2 blocks (variance)
570
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l2_vhead = kb.builder().create_block("flm_l2_vhead");
571
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l2_vbody = kb.builder().create_block("flm_l2_vbody");
572
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l2_vexit = kb.builder().create_block("flm_l2_vexit");
573
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l2_shead = kb.builder().create_block("flm_l2_shead");
574
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l2_sbody = kb.builder().create_block("flm_l2_sbody");
575
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l2_exit = kb.builder().create_block("flm_l2_exit");
576
577 // Pass 3 blocks (fused LayerNorm + Modulate directly to Y)
578
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l3_vhead = kb.builder().create_block("flm_l3_vhead");
579
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l3_vbody = kb.builder().create_block("flm_l3_vbody");
580
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l3_shead = kb.builder().create_block("flm_l3_shead");
581
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l3_sbody = kb.builder().create_block("flm_l3_sbody");
582
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 BasicBlock* l3_exit = kb.builder().create_block("flm_l3_exit");
583
584 3 kb.position_at_end(entry);
585
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* zero_i64 = kb.builder().build_iconst_i64(0);
586
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* one_i64 = kb.builder().build_iconst_i64(1);
587
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* two_i64 = kb.builder().build_iconst_i64(2);
588
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* eight_i64 = kb.builder().build_iconst_i64(8);
589
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* zero_f32 = kb.builder().build_fconst_f32(0.0f);
590
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* one_f32 = kb.builder().build_fconst_f32(1.0f);
591
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vzero = kb.vzero(Type::f32x8());
592
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vone = kb.vbroadcast(Type::f32x8(), one_f32);
593
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* mask_eight = kb.builder().build_iconst_i64(~int64_t(7));
594
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vec_d = kb.builder().build_and(d, mask_eight);
595
596
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* has_vec = kb.builder().build_sge(d, eight_i64);
597
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br_if(has_vec, l1_vhead, {zero_i64, vzero}, l1_shead, {zero_i64, zero_f32});
598
599 // Pass 1: Vector sum
600
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l1_vhead);
601
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* iv1_v = kb.builder().add_block_param(l1_vhead, Type::i64());
602
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vsum1 = kb.builder().add_block_param(l1_vhead, Type::f32x8());
603 3 kb.position_at_end(l1_vhead);
604
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* cond1_v = kb.builder().build_slt(iv1_v, vec_d);
605
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br_if(cond1_v, l1_vbody, {}, l1_vexit, {iv1_v, vsum1});
606
607
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l1_vbody);
608 3 kb.position_at_end(l1_vbody);
609
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* byte_off1 = kb.builder().build_shl(iv1_v, two_i64);
610
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* px1 = kb.builder().build_add(x, byte_off1);
611
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vx1 = kb.vload_f32x8(px1);
612
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* next_vsum1 = kb.vadd(vsum1, vx1);
613
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* next_iv1_v = kb.builder().build_add(iv1_v, eight_i64);
614
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br(l1_vhead, {next_iv1_v, next_vsum1});
615
616 // Pass 1: Vector exit -> horizontal sum
617
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l1_vexit);
618
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* exit_iv1 = kb.builder().add_block_param(l1_vexit, Type::i64());
619
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* exit_vsum1 = kb.builder().add_block_param(l1_vexit, Type::f32x8());
620 3 kb.position_at_end(l1_vexit);
621
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* init_sum1 = reduce_vsum8(kb, y, exit_vsum1);
622
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br(l1_shead, {exit_iv1, init_sum1});
623
624 // Pass 1: Scalar remainder
625
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l1_shead);
626
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* iv1_s = kb.builder().add_block_param(l1_shead, Type::i64());
627
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* sum1 = kb.builder().add_block_param(l1_shead, Type::f32());
628 3 kb.position_at_end(l1_shead);
629
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* cond1_s = kb.builder().build_slt(iv1_s, d);
630
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br_if(cond1_s, l1_sbody, {}, l1_exit, {sum1});
631
632
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l1_sbody);
633 3 kb.position_at_end(l1_sbody);
634
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* xi1 = kb.load_f32_indexed(x, iv1_s, 4, 0);
635
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* next_sum1 = kb.add(sum1, xi1);
636
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* next_iv1_s = kb.builder().build_add(iv1_s, one_i64);
637
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br(l1_shead, {next_iv1_s, next_sum1});
638
639 // Pass 1: Mean computed -> enter Pass 2
640
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l1_exit);
641
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* final_sum1 = kb.builder().add_block_param(l1_exit, Type::f32());
642 3 kb.position_at_end(l1_exit);
643
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* mean = kb.mul(final_sum1, inv_d);
644
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vmean = kb.vbroadcast(Type::f32x8(), mean);
645
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br_if(has_vec, l2_vhead, {zero_i64, vzero, mean, vmean}, l2_shead, {zero_i64, zero_f32, mean});
646
647 // Pass 2: Vector sum of squared deviations
648
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l2_vhead);
649
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* iv2_v = kb.builder().add_block_param(l2_vhead, Type::i64());
650
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vsum2 = kb.builder().add_block_param(l2_vhead, Type::f32x8());
651
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* m2_scalar = kb.builder().add_block_param(l2_vhead, Type::f32());
652
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vm2 = kb.builder().add_block_param(l2_vhead, Type::f32x8());
653 3 kb.position_at_end(l2_vhead);
654
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* cond2_v = kb.builder().build_slt(iv2_v, vec_d);
655
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br_if(cond2_v, l2_vbody, {}, l2_vexit, {iv2_v, vsum2, m2_scalar});
656
657
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l2_vbody);
658 3 kb.position_at_end(l2_vbody);
659
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* byte_off2 = kb.builder().build_shl(iv2_v, two_i64);
660
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* px2 = kb.builder().build_add(x, byte_off2);
661
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vx2 = kb.vload_f32x8(px2);
662
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vdev = kb.vsub(vx2, vm2);
663
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* next_vsum2 = kb.vfma(vdev, vdev, vsum2);
664
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* next_iv2_v = kb.builder().build_add(iv2_v, eight_i64);
665
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br(l2_vhead, {next_iv2_v, next_vsum2, m2_scalar, vm2});
666
667 // Pass 2: Vector exit -> horizontal sum
668
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l2_vexit);
669
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* exit_iv2 = kb.builder().add_block_param(l2_vexit, Type::i64());
670
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* exit_vsum2 = kb.builder().add_block_param(l2_vexit, Type::f32x8());
671
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* m2_out = kb.builder().add_block_param(l2_vexit, Type::f32());
672 3 kb.position_at_end(l2_vexit);
673
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* init_sumsq = reduce_vsum8(kb, y, exit_vsum2);
674
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br(l2_shead, {exit_iv2, init_sumsq, m2_out});
675
676 // Pass 2: Scalar remainder
677
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l2_shead);
678
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* iv2_s = kb.builder().add_block_param(l2_shead, Type::i64());
679
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* sumsq = kb.builder().add_block_param(l2_shead, Type::f32());
680
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* m2_s = kb.builder().add_block_param(l2_shead, Type::f32());
681 3 kb.position_at_end(l2_shead);
682
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* cond2_s = kb.builder().build_slt(iv2_s, d);
683
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br_if(cond2_s, l2_sbody, {}, l2_exit, {sumsq, m2_s});
684
685
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l2_sbody);
686 3 kb.position_at_end(l2_sbody);
687
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* xi2 = kb.load_f32_indexed(x, iv2_s, 4, 0);
688
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* diff2 = kb.sub(xi2, m2_s);
689
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* diff_sq2 = kb.mul(diff2, diff2);
690
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* next_sumsq = kb.add(sumsq, diff_sq2);
691
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* next_iv2_s = kb.builder().build_add(iv2_s, one_i64);
692
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br(l2_shead, {next_iv2_s, next_sumsq, m2_s});
693
694 // Pass 2 exit -> compute rstd
695
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l2_exit);
696
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* final_sumsq = kb.builder().add_block_param(l2_exit, Type::f32());
697
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* final_mean = kb.builder().add_block_param(l2_exit, Type::f32());
698 3 kb.position_at_end(l2_exit);
699
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* var = kb.mul(final_sumsq, inv_d);
700
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* var_eps = kb.add(var, eps);
701
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 Value* rstd = kb.builder().build_call("rsqrtf", Type::f32(), {var_eps});
702
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vmean3 = kb.vbroadcast(Type::f32x8(), final_mean);
703
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vrstd3 = kb.vbroadcast(Type::f32x8(), rstd);
704
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br(l3_vhead, {zero_i64, final_mean, rstd, vmean3, vrstd3});
705
706 // Pass 3: Vector fused LayerNorm + AdaLN Modulate directly into registers, zero intermediate RAM
707
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l3_vhead);
708
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* iv3_v = kb.builder().add_block_param(l3_vhead, Type::i64());
709
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* m3_s = kb.builder().add_block_param(l3_vhead, Type::f32());
710
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* r3_s = kb.builder().add_block_param(l3_vhead, Type::f32());
711
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vm3 = kb.builder().add_block_param(l3_vhead, Type::f32x8());
712
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vr3 = kb.builder().add_block_param(l3_vhead, Type::f32x8());
713 3 kb.position_at_end(l3_vhead);
714
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* cond3_v = kb.builder().build_slt(iv3_v, vec_d);
715
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br_if(cond3_v, l3_vbody, {}, l3_shead, {iv3_v, m3_s, r3_s});
716
717
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l3_vbody);
718 3 kb.position_at_end(l3_vbody);
719
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* byte_off3 = kb.builder().build_shl(iv3_v, two_i64);
720
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* px3 = kb.builder().build_add(x, byte_off3);
721
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* pg3 = kb.builder().build_add(gamma, byte_off3);
722
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* pb3 = kb.builder().build_add(beta, byte_off3);
723
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* pscale3 = kb.builder().build_add(scale, byte_off3);
724
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* pshift3 = kb.builder().build_add(shift, byte_off3);
725
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* py3 = kb.builder().build_add(y, byte_off3);
726
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vx3 = kb.vload_f32x8(px3);
727
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vg3 = kb.vload_f32x8(pg3);
728
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vb3 = kb.vload_f32x8(pb3);
729
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vscale3 = kb.vload_f32x8(pscale3);
730
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vshift3 = kb.vload_f32x8(pshift3);
731
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vdiff3 = kb.vsub(vx3, vm3);
732
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vxhat3 = kb.vmul(vdiff3, vr3);
733
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* v_ln3 = kb.vfma(vg3, vxhat3, vb3); // LayerNorm in registers
734
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* v_1_plus_scale = kb.vadd(vone, vscale3);
735
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* vy3 = kb.vfma(v_ln3, v_1_plus_scale, vshift3); // Modulated in registers
736
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.vstore_f32x8(py3, vy3); // Stored directly to destination!
737
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* next_iv3_v = kb.builder().build_add(iv3_v, eight_i64);
738
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br(l3_vhead, {next_iv3_v, m3_s, r3_s, vm3, vr3});
739
740 // Pass 3: Scalar remainder
741
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l3_shead);
742
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* iv3_s = kb.builder().add_block_param(l3_shead, Type::i64());
743
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* m3_sc = kb.builder().add_block_param(l3_shead, Type::f32());
744
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* r3_sc = kb.builder().add_block_param(l3_shead, Type::f32());
745 3 kb.position_at_end(l3_shead);
746
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* cond3_s = kb.builder().build_slt(iv3_s, d);
747
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br_if(cond3_s, l3_sbody, l3_exit);
748
749
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l3_sbody);
750 3 kb.position_at_end(l3_sbody);
751
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* x3 = kb.load_f32_indexed(x, iv3_s, 4, 0);
752
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* g3 = kb.load_f32_indexed(gamma, iv3_s, 4, 0);
753
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* b3 = kb.load_f32_indexed(beta, iv3_s, 4, 0);
754
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* s3 = kb.load_f32_indexed(scale, iv3_s, 4, 0);
755
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* sh3 = kb.load_f32_indexed(shift, iv3_s, 4, 0);
756
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* diff3 = kb.sub(x3, m3_sc);
757
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* xhat3 = kb.mul(diff3, r3_sc);
758
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* ln_val3 = kb.builder().build_fma_f32(g3, xhat3, b3);
759
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* one_plus_s3 = kb.add(one_f32, s3);
760
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* y3 = kb.builder().build_fma_f32(ln_val3, one_plus_s3, sh3);
761
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.store_f32_indexed(y, iv3_s, y3, 4, 0);
762
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 Value* next_iv3_s = kb.builder().build_add(iv3_s, one_i64);
763
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_br(l3_shead, {next_iv3_s, m3_sc, r3_sc});
764
765
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 fn->append_block(l3_exit);
766 3 kb.position_at_end(l3_exit);
767
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kb.builder().build_ret_void();
768
769
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 3 return jit.compile(mod, "fused_layernorm_modulate_row");
770 3 }
771
772 // ── Kernel Initialization & JIT Engine Management ───────────────────────────
773 490 static const JitFusedKernels& get_fused_kernels() {
774
3/4
✓ Branch 0 taken 3 times.
✓ Branch 1 taken 487 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 3 times.
 490 static JitFusedKernels kernels;
775 static std::once_flag init_flag;
776 493 std::call_once(init_flag, [] {
777 try {
778 3 KernelOptions opts;
779 3 opts.enable_optimizations = true;
780 3 opts.enable_avx2 = true;
781 3 opts.enable_fma = true;
782 3 opts.enable_vectorize = true;
783 3 opts.enable_unroll = true;
784 3 opts.unroll_factor = 4;
785 3 opts.enable_fp_reassociation = true;
786
787
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 KernelJit jit(opts);
788
789 // 1. Fused Residual RMSNorm row (Option B)
790
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 KernelFunction k_fused_res_rms = build_fused_residual_rmsnorm_row(jit);
791 3 kernels.fused_residual_rmsnorm_row_fn = k_fused_res_rms.as<JitFusedKernels::FusedResidualRmsNormRowFn>();
792
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kernels.engines.push_back(k_fused_res_rms.engine());
793
794 // 2. Fused Residual LayerNorm row
795
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 KernelFunction k_fused_res_ln = build_fused_residual_layernorm_row(jit);
796 3 kernels.fused_residual_layernorm_row_fn = k_fused_res_ln.as<JitFusedKernels::FusedResidualLayerNormRowFn>();
797
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kernels.engines.push_back(k_fused_res_ln.engine());
798
799 // 3. Fused LayerNorm Modulate row (Option B)
800
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 KernelFunction k_fused_ln_mod = build_fused_layernorm_modulate_row(jit);
801 3 kernels.fused_layernorm_modulate_row_fn = k_fused_ln_mod.as<JitFusedKernels::FusedLayerNormModulateRowFn>();
802
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 kernels.engines.push_back(k_fused_ln_mod.engine());
803
804
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 6 kernels.available = (kernels.fused_residual_rmsnorm_row_fn != nullptr &&
805
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 3 times.
 3 kernels.fused_residual_layernorm_row_fn != nullptr &&
806 3 kernels.fused_layernorm_modulate_row_fn != nullptr);
807
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 3 } catch (const std::exception& e) {
808 ✗ std::cerr << "brotensor: Brass JIT fused kernels initialization failed: " << e.what() << "\n";
809 ✗ kernels.available = false;
810 ✗ }
811 3 });
812 490 return kernels;
813 }
814
815 } // namespace
816
817 246 void fused_residual_rmsnorm(float* X, const float* res, const float* gamma, float eps, float* Y, int B, int D) {
818 246 const auto& k = get_fused_kernels();
819
2/4
✓ Branch 0 taken 246 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 246 times.
 246 if (!k.available || k.fused_residual_rmsnorm_row_fn == nullptr) {
820 ✗ ref::fused_residual_rmsnorm(X, res, gamma, eps, Y, B, D);
821 ✗ return;
822 }
823
824 246 const float inv_D = 1.0f / static_cast<float>(D);
825
826
3/4
✓ Branch 0 taken 185 times.
✓ Branch 1 taken 61 times.
✓ Branch 2 taken 185 times.
✗ Branch 3 not taken.
 246 if (B > 1 && static_cast<int64_t>(B) * D >= 262144) {
827 ✗ detail::cpu::parallel_for(static_cast<std::size_t>(B), [&](std::size_t bi) {
828 ✗ const int b = static_cast<int>(bi);
829 ✗ float* xr = X + static_cast<std::size_t>(b) * D;
830 ✗ const float* rr = res + static_cast<std::size_t>(b) * D;
831 ✗ float* yr = Y + static_cast<std::size_t>(b) * D;
832 ✗ k.fused_residual_rmsnorm_row_fn(xr, rr, gamma, yr, static_cast<uint64_t>(D), eps, inv_D);
833 ✗ });
834 ✗ } else {
835
2/2
✓ Branch 0 taken 2753 times.
✓ Branch 1 taken 246 times.
 2999 for (int b = 0; b < B; ++b) {
836 2753 float* xr = X + static_cast<std::size_t>(b) * D;
837 2753 const float* rr = res + static_cast<std::size_t>(b) * D;
838 2753 float* yr = Y + static_cast<std::size_t>(b) * D;
839 2753 k.fused_residual_rmsnorm_row_fn(xr, rr, gamma, yr, static_cast<uint64_t>(D), eps, inv_D);
840 2753 }
841 }
842 246 }
843
844 2 void fused_residual_layernorm(float* X, const float* res, const float* gamma, const float* beta,
845 float eps, float* Y, int B, int D) {
846 2 const auto& k = get_fused_kernels();
847
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
 2 if (!k.available || k.fused_residual_layernorm_row_fn == nullptr) {
848 ✗ ref::fused_residual_layernorm(X, res, gamma, beta, eps, Y, B, D);
849 ✗ return;
850 }
851
852 2 const float inv_D = 1.0f / static_cast<float>(D);
853
854
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 if (B > 1 && static_cast<int64_t>(B) * D >= 262144) {
855 ✗ detail::cpu::parallel_for(static_cast<std::size_t>(B), [&](std::size_t bi) {
856 ✗ const int b = static_cast<int>(bi);
857 ✗ float* xr = X + static_cast<std::size_t>(b) * D;
858 ✗ const float* rr = res + static_cast<std::size_t>(b) * D;
859 ✗ float* yr = Y + static_cast<std::size_t>(b) * D;
860 ✗ k.fused_residual_layernorm_row_fn(xr, rr, gamma, beta, yr, static_cast<uint64_t>(D), eps, inv_D);
861 ✗ });
862 ✗ } else {
863
2/2
✓ Branch 0 taken 397 times.
✓ Branch 1 taken 2 times.
 399 for (int b = 0; b < B; ++b) {
864 397 float* xr = X + static_cast<std::size_t>(b) * D;
865 397 const float* rr = res + static_cast<std::size_t>(b) * D;
866 397 float* yr = Y + static_cast<std::size_t>(b) * D;
867 397 k.fused_residual_layernorm_row_fn(xr, rr, gamma, beta, yr, static_cast<uint64_t>(D), eps, inv_D);
868 397 }
869 }
870 2 }
871
872 242 void fused_layernorm_modulate(const float* X, const float* gamma, const float* beta,
873 const float* scale, const float* shift, float eps, float* Y, int R, int D) {
874 242 const auto& k = get_fused_kernels();
875
2/4
✓ Branch 0 taken 242 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 242 times.
 242 if (!k.available || k.fused_layernorm_modulate_row_fn == nullptr) {
876 ✗ ref::fused_layernorm_modulate(X, gamma, beta, scale, shift, eps, Y, R, D);
877 ✗ return;
878 }
879
880 242 const float inv_D = 1.0f / static_cast<float>(D);
881
882
3/4
✓ Branch 0 taken 182 times.
✓ Branch 1 taken 60 times.
✓ Branch 2 taken 182 times.
✗ Branch 3 not taken.
 242 if (R > 1 && static_cast<int64_t>(R) * D >= 262144) {
883 ✗ detail::cpu::parallel_for(static_cast<std::size_t>(R), [&](std::size_t ri) {
884 ✗ const int r = static_cast<int>(ri);
885 ✗ const float* xr = X + static_cast<std::size_t>(r) * D;
886 ✗ float* yr = Y + static_cast<std::size_t>(r) * D;
887 ✗ k.fused_layernorm_modulate_row_fn(xr, gamma, beta, scale, shift, yr, static_cast<uint64_t>(D), eps, inv_D);
888 ✗ });
889 ✗ } else {
890
2/2
✓ Branch 0 taken 2716 times.
✓ Branch 1 taken 242 times.
 2958 for (int r = 0; r < R; ++r) {
891 2716 const float* xr = X + static_cast<std::size_t>(r) * D;
892 2716 float* yr = Y + static_cast<std::size_t>(r) * D;
893 2716 k.fused_layernorm_modulate_row_fn(xr, gamma, beta, scale, shift, yr, static_cast<uint64_t>(D), eps, inv_D);
894 2716 }
895 }
896 242 }
897
898 } // namespace brotensor::detail::cpu::jit
899
900 #else // !BROTENSOR_HAS_BRASS_JIT
901
902 namespace brotensor::detail::cpu::jit {
903
904 void fused_residual_rmsnorm(float* X, const float* res, const float* gamma, float eps, float* Y, int B, int D) {
905 ref::fused_residual_rmsnorm(X, res, gamma, eps, Y, B, D);
906 }
907 void fused_residual_layernorm(float* X, const float* res, const float* gamma, const float* beta, float eps, float* Y, int B, int D) {
908 ref::fused_residual_layernorm(X, res, gamma, beta, eps, Y, B, D);
909 }
910 void fused_layernorm_modulate(const float* X, const float* gamma, const float* beta,
911 const float* scale, const float* shift, float eps, float* Y, int R, int D) {
912 ref::fused_layernorm_modulate(X, gamma, beta, scale, shift, eps, Y, R, D);
913 }
914
915 } // namespace brotensor::detail::cpu::jit
916
917 #endif // BROTENSOR_HAS_BRASS_JIT
918