GCC Code Coverage Report

Directory: brotensor/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 87.6% 535 / 0 / 611
Functions: 67.7% 21 / 0 / 31
Branches: 48.8% 493 / 0 / 1010
src/cpu/cpu_jit.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 rms_norm(const float* X, 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 ✗ const float* xr = X + static_cast<size_t>(b) * D;
15 ✗ float* yr = Y + static_cast<size_t>(b) * D;
16 ✗ float sum = 0.0f;
17 ✗ for (int i = 0; i < D; ++i) sum += xr[i] * xr[i];
18 ✗ float rrms = 1.0f / std::sqrt(sum * inv_D + eps);
19 ✗ for (int i = 0; i < D; ++i) yr[i] = xr[i] * gamma[i] * rrms;
20 ✗ };
21 ✗ if (B > 1 && static_cast<int64_t>(B) * D >= 262144) {
22 ✗ detail::cpu::parallel_for(static_cast<size_t>(B), [&](size_t bi) {
23 ✗ row_fn(static_cast<int>(bi));
24 ✗ });
25 ✗ } else {
26 ✗ for (int b = 0; b < B; ++b) row_fn(b);
27 }
28 ✗ }
29
30 247 inline void swiglu(const float* X, float* Y, int B, int D) {
31
2/4
✓ Branch 0 taken 247 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 247 times.
 247 if (B <= 0 || D <= 0) return;
32 3364 auto row_fn = [&](int b) {
33 3117 const float* gate = X + static_cast<size_t>(b) * 2 * D;
34 3117 const float* up = gate + D;
35 3117 float* out = Y + static_cast<size_t>(b) * D;
36
2/2
✓ Branch 0 taken 3525511 times.
✓ Branch 1 taken 3117 times.
 3528628 for (int i = 0; i < D; ++i) {
37 3525511 float g = gate[i];
38 3525511 float silu = g / (1.0f + std::exp(-g));
39 3525511 out[i] = silu * up[i];
40 3525511 }
41 3117 };
42
4/4
✓ Branch 0 taken 186 times.
✓ Branch 1 taken 61 times.
✓ Branch 2 taken 124 times.
✓ Branch 3 taken 62 times.
 247 if (B > 1 && static_cast<int64_t>(B) * D >= 16384) {
43
1/2
✓ Branch 0 taken 62 times.
✗ Branch 1 not taken.
 2377 detail::cpu::parallel_for(static_cast<size_t>(B), [&](size_t bi) {
44 2315 row_fn(static_cast<int>(bi));
45 2315 });
46 62 } else {
47
2/2
✓ Branch 0 taken 807 times.
✓ Branch 1 taken 185 times.
 992 for (int b = 0; b < B; ++b) row_fn(b);
48 }
49 247 }
50
51 246 inline void modulate(const float* X, const float* scale, const float* shift, float* Y, int L, int D) {
52
2/4
✓ Branch 0 taken 246 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 246 times.
 246 if (L <= 0 || D <= 0) return;
53 2975 auto row_fn = [&](int l) {
54 2729 const float* xr = X + static_cast<size_t>(l) * D;
55 2729 float* yr = Y + static_cast<size_t>(l) * D;
56
2/2
✓ Branch 0 taken 8453035 times.
✓ Branch 1 taken 2729 times.
 8455764 for (int d = 0; d < D; ++d) {
57 8453035 yr[d] = xr[d] * (1.0f + scale[d]) + shift[d];
58 8453035 }
59 2729 };
60
3/4
✓ Branch 0 taken 185 times.
✓ Branch 1 taken 61 times.
✓ Branch 2 taken 185 times.
✗ Branch 3 not taken.
 246 if (L > 1 && static_cast<int64_t>(L) * D >= 262144) {
61 ✗ detail::cpu::parallel_for(static_cast<size_t>(L), [&](size_t li) {
62 ✗ row_fn(static_cast<int>(li));
63 ✗ });
64 ✗ } else {
65
2/2
✓ Branch 0 taken 2729 times.
✓ Branch 1 taken 246 times.
 2975 for (int l = 0; l < L; ++l) row_fn(l);
66 }
67 246 }
68
69 245 inline void broadcast_mul(const float* X, const float* v, float* Y, int L, int D) {
70
2/4
✓ Branch 0 taken 245 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 245 times.
 245 if (L <= 0 || D <= 0) return;
71 2973 auto row_fn = [&](int l) {
72 2728 const float* xr = X + static_cast<size_t>(l) * D;
73 2728 float* yr = Y + static_cast<size_t>(l) * D;
74
2/2
✓ Branch 0 taken 8453034 times.
✓ Branch 1 taken 2728 times.
 8455762 for (int d = 0; d < D; ++d) {
75 8453034 yr[d] = xr[d] * v[d];
76 8453034 }
77 2728 };
78
3/4
✓ Branch 0 taken 185 times.
✓ Branch 1 taken 60 times.
✓ Branch 2 taken 185 times.
✗ Branch 3 not taken.
 245 if (L > 1 && static_cast<int64_t>(L) * D >= 262144) {
79 ✗ detail::cpu::parallel_for(static_cast<size_t>(L), [&](size_t li) {
80 ✗ row_fn(static_cast<int>(li));
81 ✗ });
82 ✗ } else {
83
2/2
✓ Branch 0 taken 2728 times.
✓ Branch 1 taken 245 times.
 2973 for (int l = 0; l < L; ++l) row_fn(l);
84 }
85 245 }
86
87 ✗ inline void layernorm_forward_inference_batched(const float* X, const float* gamma, const float* beta,
88 float eps, float* Y, int R, int D) {
89 ✗ if (R <= 0 || D <= 0) return;
90 ✗ const float inv_D = 1.0f / static_cast<float>(D);
91 ✗ auto row_fn = [&](int r) {
92 ✗ const float* xr = X + static_cast<size_t>(r) * D;
93 ✗ float* yr = Y + static_cast<size_t>(r) * D;
94 ✗ float sum = 0.0f;
95 ✗ for (int i = 0; i < D; ++i) sum += xr[i];
96 ✗ float mean = sum * inv_D;
97 ✗ float sumsq = 0.0f;
98 ✗ for (int i = 0; i < D; ++i) {
99 ✗ float diff = xr[i] - mean;
100 ✗ sumsq += diff * diff;
101 ✗ }
102 ✗ float var = sumsq * inv_D;
103 ✗ float rstd = 1.0f / std::sqrt(var + eps);
104 ✗ for (int i = 0; i < D; ++i) {
105 ✗ float g = gamma ? gamma[i] : 1.0f;
106 ✗ float b = beta ? beta[i] : 0.0f;
107 ✗ yr[i] = g * ((xr[i] - mean) * rstd) + b;
108 ✗ }
109 ✗ };
110 ✗ if (R > 1 && static_cast<int64_t>(R) * D >= 262144) {
111 ✗ detail::cpu::parallel_for(static_cast<size_t>(R), [&](size_t ri) {
112 ✗ row_fn(static_cast<int>(ri));
113 ✗ });
114 ✗ } else {
115 ✗ for (int r = 0; r < R; ++r) row_fn(r);
116 }
117 ✗ }
118
119 } // namespace ref
120 } // namespace brotensor::detail::cpu::jit
121
122 #if BROTENSOR_HAS_BRASS_JIT
123
124 #include <brass/codegen/kernel_jit.hpp>
125 #include <brass/codegen/ml_fusion.hpp>
126 #include <brass/mir/builder.hpp>
127 #include <brass/mir/module.hpp>
128 #include <brass/mir/function.hpp>
129 #include <cmath>
130 #include <mutex>
131 #include <vector>
132 #include <memory>
133 #include <algorithm>
134 #include <iostream>
135
136 namespace brotensor::detail::cpu::jit {
137
138 namespace {
139
140 using namespace brass;
141 using namespace brass::codegen;
142
143 struct JitKernels {
144 bool available = false;
145
146 using AdaLNModulateFn = void (*)(const float* x, const float* scale, const float* shift, float* out, uint64_t n);
147 AdaLNModulateFn adaln_fn = nullptr;
148
149 using BroadcastMulFn = void (*)(const float* x, const float* v, float* out, uint64_t n);
150 BroadcastMulFn broadcast_mul_fn = nullptr;
151
152 using RmsNormRowFn = void (*)(const float* x, const float* gamma, float* y, uint64_t d, float eps, float inv_d);
153 RmsNormRowFn rms_norm_row_fn = nullptr;
154
155 using LayerNormRowFn = void (*)(const float* x, const float* gamma, const float* beta, float* y, uint64_t d, float eps, float inv_d);
156 LayerNormRowFn layernorm_row_fn = nullptr;
157
158 std::vector<std::shared_ptr<JitExecutionEngine>> engines;
159 };
160
161 // Horizontal sum of an 8-wide float vector using destination scratch buffer
162 24 static Value* reduce_vsum8(KernelBuilder& kb, Value* scratch, Value* vsum) {
163 24 kb.vstore_f32x8(scratch, vsum);
164 24 Value* s0 = kb.load_f32(scratch, 0); Value* s1 = kb.load_f32(scratch, 4);
165 24 Value* s2 = kb.load_f32(scratch, 8); Value* s3 = kb.load_f32(scratch, 12);
166 24 Value* s4 = kb.load_f32(scratch, 16); Value* s5 = kb.load_f32(scratch, 20);
167 24 Value* s6 = kb.load_f32(scratch, 24); Value* s7 = kb.load_f32(scratch, 28);
168 24 Value* sum0123 = kb.add(kb.add(s0, s1), kb.add(s2, s3));
169 24 Value* sum4567 = kb.add(kb.add(s4, s5), kb.add(s6, s7));
170 24 return kb.add(sum0123, sum4567);
171 }
172
173 // ── 1. Vectorized Broadcast Mul (AVX2 8-wide + scalar remainder) ────────────
174 8 static KernelFunction build_broadcast_mul(KernelJit& jit) {
175 8 Module mod("mod_broadcast_mul");
176
6/12
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 8 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 8 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 8 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 8 times.
✗ Branch 11 not taken.
 8 Function* fn = mod.create_function("broadcast_mul", Type::void_type(), {
177 32 Type::ptr(), Type::ptr(), Type::ptr(), Type::i64()
178 });
179
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 KernelBuilder kb(mod, fn);
180
181
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* entry = kb.builder().append_block("entry");
182
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* x = kb.builder().add_block_param(entry, Type::ptr());
183
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* v = kb.builder().add_block_param(entry, Type::ptr());
184
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* out = kb.builder().add_block_param(entry, Type::ptr());
185
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* n = kb.builder().add_block_param(entry, Type::i64());
186
187
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* v_head = kb.builder().create_block("v_head");
188
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* v_body = kb.builder().create_block("v_body");
189
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* s_head = kb.builder().create_block("s_head");
190
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* s_body = kb.builder().create_block("s_body");
191
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* exit = kb.builder().create_block("exit");
192
193 8 kb.position_at_end(entry);
194
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* zero_i64 = kb.builder().build_iconst_i64(0);
195
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* one_i64 = kb.builder().build_iconst_i64(1);
196
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* two_i64 = kb.builder().build_iconst_i64(2);
197
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* eight_i64 = kb.builder().build_iconst_i64(8);
198
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* mask_eight = kb.builder().build_iconst_i64(~int64_t(7));
199
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vec_n = kb.builder().build_and(n, mask_eight);
200
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br(v_head, {zero_i64});
201
202 // Vector loop: 8 elements per iteration
203
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(v_head);
204
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* iv_v = kb.builder().add_block_param(v_head, Type::i64());
205 8 kb.position_at_end(v_head);
206
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* v_cond = kb.builder().build_slt(iv_v, vec_n);
207
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br_if(v_cond, v_body, {}, s_head, {iv_v});
208
209
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(v_body);
210 8 kb.position_at_end(v_body);
211
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* byte_off = kb.builder().build_shl(iv_v, two_i64);
212
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* px = kb.builder().build_add(x, byte_off);
213
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* pv = kb.builder().build_add(v, byte_off);
214
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* py = kb.builder().build_add(out, byte_off);
215
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vx = kb.vload_f32x8(px);
216
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vv = kb.vload_f32x8(pv);
217
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vy = kb.vmul(vx, vv);
218
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.vstore_f32x8(py, vy);
219
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* next_iv_v = kb.builder().build_add(iv_v, eight_i64);
220
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br(v_head, {next_iv_v});
221
222 // Scalar remainder loop
223
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(s_head);
224
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* iv_s = kb.builder().add_block_param(s_head, Type::i64());
225 8 kb.position_at_end(s_head);
226
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* s_cond = kb.builder().build_slt(iv_s, n);
227
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br_if(s_cond, s_body, exit);
228
229
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(s_body);
230 8 kb.position_at_end(s_body);
231
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* xi = kb.load_f32_indexed(x, iv_s, 4, 0);
232
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vi = kb.load_f32_indexed(v, iv_s, 4, 0);
233
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* yi = kb.mul(xi, vi);
234
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.store_f32_indexed(out, iv_s, yi, 4, 0);
235
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* next_iv_s = kb.builder().build_add(iv_s, one_i64);
236
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br(s_head, {next_iv_s});
237
238
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(exit);
239 8 kb.position_at_end(exit);
240
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_ret_void();
241
242
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 return jit.compile(mod, "broadcast_mul");
243 8 }
244
245 // ── 2. Vectorized AdaLN Modulate (AVX2 8-wide + FMA + scalar remainder) ─────
246 8 static KernelFunction build_adaln_modulate(KernelJit& jit) {
247 8 Module mod("mod_adaln_modulate");
248
7/14
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 8 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 8 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 8 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 8 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 8 times.
✗ Branch 13 not taken.
 8 Function* fn = mod.create_function("fused_adaln_modulate", Type::void_type(), {
249 40 Type::ptr(), Type::ptr(), Type::ptr(), Type::ptr(), Type::i64()
250 });
251
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 KernelBuilder kb(mod, fn);
252
253
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* entry = kb.builder().append_block("entry");
254
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* x = kb.builder().add_block_param(entry, Type::ptr());
255
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* scale = kb.builder().add_block_param(entry, Type::ptr());
256
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* shift = kb.builder().add_block_param(entry, Type::ptr());
257
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* out = kb.builder().add_block_param(entry, Type::ptr());
258
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* n = kb.builder().add_block_param(entry, Type::i64());
259
260
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* v_head = kb.builder().create_block("v_head");
261
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* v_body = kb.builder().create_block("v_body");
262
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* s_head = kb.builder().create_block("s_head");
263
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* s_body = kb.builder().create_block("s_body");
264
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* exit = kb.builder().create_block("exit");
265
266 8 kb.position_at_end(entry);
267
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* zero_i64 = kb.builder().build_iconst_i64(0);
268
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* one_i64 = kb.builder().build_iconst_i64(1);
269
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* two_i64 = kb.builder().build_iconst_i64(2);
270
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* eight_i64 = kb.builder().build_iconst_i64(8);
271
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* one_f32 = kb.builder().build_fconst_f32(1.0f);
272
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* v_one = kb.vbroadcast(Type::f32x8(), one_f32);
273
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* mask_eight = kb.builder().build_iconst_i64(~int64_t(7));
274
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vec_n = kb.builder().build_and(n, mask_eight);
275
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br(v_head, {zero_i64});
276
277 // Vector loop: modulated = x * (1.0f + scale) + shift via AVX2 FMA
278
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(v_head);
279
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* iv_v = kb.builder().add_block_param(v_head, Type::i64());
280 8 kb.position_at_end(v_head);
281
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* v_cond = kb.builder().build_slt(iv_v, vec_n);
282
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br_if(v_cond, v_body, {}, s_head, {iv_v});
283
284
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(v_body);
285 8 kb.position_at_end(v_body);
286
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* byte_off = kb.builder().build_shl(iv_v, two_i64);
287
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* px = kb.builder().build_add(x, byte_off);
288
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* pscale = kb.builder().build_add(scale, byte_off);
289
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* pshift = kb.builder().build_add(shift, byte_off);
290
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* py = kb.builder().build_add(out, byte_off);
291
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vx = kb.vload_f32x8(px);
292
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vscale = kb.vload_f32x8(pscale);
293
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vshift = kb.vload_f32x8(pshift);
294
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* v_1_plus_scale = kb.vadd(v_one, vscale);
295
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vy = kb.vfma(vx, v_1_plus_scale, vshift);
296
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.vstore_f32x8(py, vy);
297
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* next_iv_v = kb.builder().build_add(iv_v, eight_i64);
298
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br(v_head, {next_iv_v});
299
300 // Scalar remainder loop
301
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(s_head);
302
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* iv_s = kb.builder().add_block_param(s_head, Type::i64());
303 8 kb.position_at_end(s_head);
304
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* s_cond = kb.builder().build_slt(iv_s, n);
305
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br_if(s_cond, s_body, exit);
306
307
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(s_body);
308 8 kb.position_at_end(s_body);
309
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* xi = kb.load_f32_indexed(x, iv_s, 4, 0);
310
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* scale_i = kb.load_f32_indexed(scale, iv_s, 4, 0);
311
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* shift_i = kb.load_f32_indexed(shift, iv_s, 4, 0);
312
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* one_plus_scale = kb.add(one_f32, scale_i);
313
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* yi = kb.builder().build_fma_f32(xi, one_plus_scale, shift_i);
314
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.store_f32_indexed(out, iv_s, yi, 4, 0);
315
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* next_iv_s = kb.builder().build_add(iv_s, one_i64);
316
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br(s_head, {next_iv_s});
317
318
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(exit);
319 8 kb.position_at_end(exit);
320
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_ret_void();
321
322
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 return jit.compile(mod, "fused_adaln_modulate");
323 8 }
324
325 // ── 3. Vectorized RMSNorm Row (AVX2 FMA reduction + AVX2 affine scale) ─────
326 8 static KernelFunction build_rms_norm_row(KernelJit& jit) {
327 8 Module mod("mod_rms_norm_row");
328
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 mod.add_external_symbol("rsqrtf");
329
8/16
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 8 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 8 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 8 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 8 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 8 times.
✗ Branch 13 not taken.
✓ Branch 14 taken 8 times.
✗ Branch 15 not taken.
 8 Function* fn = mod.create_function("rms_norm_row", Type::void_type(), {
330 48 Type::ptr(), Type::ptr(), Type::ptr(), Type::i64(), Type::f32(), Type::f32()
331 });
332
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 KernelBuilder kb(mod, fn);
333
334
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* entry = kb.builder().append_block("entry");
335
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* x = kb.builder().add_block_param(entry, Type::ptr());
336
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* gamma = kb.builder().add_block_param(entry, Type::ptr());
337
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* y = kb.builder().add_block_param(entry, Type::ptr());
338
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* d = kb.builder().add_block_param(entry, Type::i64());
339
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* eps = kb.builder().add_block_param(entry, Type::f32());
340
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* inv_d = kb.builder().add_block_param(entry, Type::f32());
341
342
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* l1_vhead = kb.builder().create_block("l1_vhead");
343
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* l1_vbody = kb.builder().create_block("l1_vbody");
344
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* l1_vexit = kb.builder().create_block("l1_vexit");
345
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* l1_shead = kb.builder().create_block("l1_shead");
346
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* l1_sbody = kb.builder().create_block("l1_sbody");
347
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* l1_exit = kb.builder().create_block("l1_exit");
348
349
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* l2_vhead = kb.builder().create_block("l2_vhead");
350
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* l2_vbody = kb.builder().create_block("l2_vbody");
351
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* l2_shead = kb.builder().create_block("l2_shead");
352
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* l2_sbody = kb.builder().create_block("l2_sbody");
353
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* l2_exit = kb.builder().create_block("l2_exit");
354
355 8 kb.position_at_end(entry);
356
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* zero_i64 = kb.builder().build_iconst_i64(0);
357
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* one_i64 = kb.builder().build_iconst_i64(1);
358
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* two_i64 = kb.builder().build_iconst_i64(2);
359
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* eight_i64 = kb.builder().build_iconst_i64(8);
360
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* zero_f32 = kb.builder().build_fconst_f32(0.0f);
361
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vzero = kb.vzero(Type::f32x8());
362
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* mask_eight = kb.builder().build_iconst_i64(~int64_t(7));
363
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vec_d = kb.builder().build_and(d, mask_eight);
364
365
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* has_vec = kb.builder().build_sge(d, eight_i64);
366
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br_if(has_vec, l1_vhead, {zero_i64, vzero}, l1_shead, {zero_i64, zero_f32});
367
368 // Pass 1: Vector loop accumulating sum_sq via vfma
369
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(l1_vhead);
370
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* iv1_v = kb.builder().add_block_param(l1_vhead, Type::i64());
371
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vsum = kb.builder().add_block_param(l1_vhead, Type::f32x8());
372 8 kb.position_at_end(l1_vhead);
373
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* cond1_v = kb.builder().build_slt(iv1_v, vec_d);
374
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br_if(cond1_v, l1_vbody, {}, l1_vexit, {iv1_v, vsum});
375
376
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(l1_vbody);
377 8 kb.position_at_end(l1_vbody);
378
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* byte_off1 = kb.builder().build_shl(iv1_v, two_i64);
379
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* px1 = kb.builder().build_add(x, byte_off1);
380
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vx1 = kb.vload_f32x8(px1);
381
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* next_vsum = kb.vfma(vx1, vx1, vsum);
382
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* next_iv1_v = kb.builder().build_add(iv1_v, eight_i64);
383
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br(l1_vhead, {next_iv1_v, next_vsum});
384
385 // Pass 1: Vector exit -> horizontal vector reduction via y scratch buffer
386
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(l1_vexit);
387
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* exit_iv1 = kb.builder().add_block_param(l1_vexit, Type::i64());
388
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* exit_vsum = kb.builder().add_block_param(l1_vexit, Type::f32x8());
389 8 kb.position_at_end(l1_vexit);
390
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* init_scalar_sum = reduce_vsum8(kb, y, exit_vsum);
391
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br(l1_shead, {exit_iv1, init_scalar_sum});
392
393 // Pass 1: Scalar remainder loop
394
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(l1_shead);
395
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* iv1_s = kb.builder().add_block_param(l1_shead, Type::i64());
396
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* sum_sq = kb.builder().add_block_param(l1_shead, Type::f32());
397 8 kb.position_at_end(l1_shead);
398
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* cond1_s = kb.builder().build_slt(iv1_s, d);
399
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br_if(cond1_s, l1_sbody, {}, l1_exit, {sum_sq});
400
401
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(l1_sbody);
402 8 kb.position_at_end(l1_sbody);
403
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* xi = kb.load_f32_indexed(x, iv1_s, 4, 0);
404
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* xi_sq = kb.mul(xi, xi);
405
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* next_sum = kb.add(sum_sq, xi_sq);
406
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* next_iv1_s = kb.builder().build_add(iv1_s, one_i64);
407
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br(l1_shead, {next_iv1_s, next_sum});
408
409 // Pass 1 exit -> compute rrms = rsqrtf(mean_sq + eps)
410
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(l1_exit);
411
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* final_sum = kb.builder().add_block_param(l1_exit, Type::f32());
412 8 kb.position_at_end(l1_exit);
413
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* mean_sq = kb.mul(final_sum, inv_d);
414
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* denom = kb.add(mean_sq, eps);
415
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 Value* rrms = kb.builder().build_call("rsqrtf", Type::f32(), {denom});
416
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vrrms = kb.vbroadcast(Type::f32x8(), rrms);
417
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br(l2_vhead, {zero_i64, rrms, vrrms});
418
419 // Pass 2: Vector loop y[i] = x[i] * gamma[i] * rrms
420
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(l2_vhead);
421
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* iv2_v = kb.builder().add_block_param(l2_vhead, Type::i64());
422
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* rrms_val = kb.builder().add_block_param(l2_vhead, Type::f32());
423
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vrrms_val = kb.builder().add_block_param(l2_vhead, Type::f32x8());
424 8 kb.position_at_end(l2_vhead);
425
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* cond2_v = kb.builder().build_slt(iv2_v, vec_d);
426
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br_if(cond2_v, l2_vbody, {}, l2_shead, {iv2_v, rrms_val});
427
428
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(l2_vbody);
429 8 kb.position_at_end(l2_vbody);
430
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* byte_off2 = kb.builder().build_shl(iv2_v, two_i64);
431
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* px2 = kb.builder().build_add(x, byte_off2);
432
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* pg2 = kb.builder().build_add(gamma, byte_off2);
433
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* py2 = kb.builder().build_add(y, byte_off2);
434
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vx2 = kb.vload_f32x8(px2);
435
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vg2 = kb.vload_f32x8(pg2);
436
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vxg = kb.vmul(vx2, vg2);
437
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vy = kb.vmul(vxg, vrrms_val);
438
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.vstore_f32x8(py2, vy);
439
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* next_iv2_v = kb.builder().build_add(iv2_v, eight_i64);
440
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br(l2_vhead, {next_iv2_v, rrms_val, vrrms_val});
441
442 // Pass 2: Scalar remainder loop
443
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(l2_shead);
444
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* iv2_s = kb.builder().add_block_param(l2_shead, Type::i64());
445
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* rrms_scalar = kb.builder().add_block_param(l2_shead, Type::f32());
446 8 kb.position_at_end(l2_shead);
447
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* cond2_s = kb.builder().build_slt(iv2_s, d);
448
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br_if(cond2_s, l2_sbody, l2_exit);
449
450
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(l2_sbody);
451 8 kb.position_at_end(l2_sbody);
452
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* x2 = kb.load_f32_indexed(x, iv2_s, 4, 0);
453
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* g2 = kb.load_f32_indexed(gamma, iv2_s, 4, 0);
454
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* xg2 = kb.mul(x2, g2);
455
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* out_val = kb.mul(xg2, rrms_scalar);
456
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.store_f32_indexed(y, iv2_s, out_val, 4, 0);
457
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* next_iv2_s = kb.builder().build_add(iv2_s, one_i64);
458
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br(l2_shead, {next_iv2_s, rrms_scalar});
459
460
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(l2_exit);
461 8 kb.position_at_end(l2_exit);
462
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_ret_void();
463
464
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 return jit.compile(mod, "rms_norm_row");
465 8 }
466
467 // ── 4. Vectorized LayerNorm Row (AVX2 mean + var + affine) ──────────────────
468 8 static KernelFunction build_layernorm_row(KernelJit& jit) {
469 8 Module mod("mod_layernorm_row");
470
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 mod.add_external_symbol("rsqrtf");
471
9/18
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 8 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 8 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 8 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 8 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 8 times.
✗ Branch 13 not taken.
✓ Branch 14 taken 8 times.
✗ Branch 15 not taken.
✓ Branch 16 taken 8 times.
✗ Branch 17 not taken.
 8 Function* fn = mod.create_function("layernorm_row", Type::void_type(), {
472 56 Type::ptr(), Type::ptr(), Type::ptr(), Type::ptr(), Type::i64(), Type::f32(), Type::f32()
473 });
474
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 KernelBuilder kb(mod, fn);
475
476
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* entry = kb.builder().append_block("entry");
477
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* x = kb.builder().add_block_param(entry, Type::ptr());
478
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* gamma = kb.builder().add_block_param(entry, Type::ptr());
479
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* beta = kb.builder().add_block_param(entry, Type::ptr());
480
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* y = kb.builder().add_block_param(entry, Type::ptr());
481
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* d = kb.builder().add_block_param(entry, Type::i64());
482
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* eps = kb.builder().add_block_param(entry, Type::f32());
483
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* inv_d = kb.builder().add_block_param(entry, Type::f32());
484
485 // Pass 1 blocks (mean)
486
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* l1_vhead = kb.builder().create_block("ln_l1_vhead");
487
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* l1_vbody = kb.builder().create_block("ln_l1_vbody");
488
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* l1_vexit = kb.builder().create_block("ln_l1_vexit");
489
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* l1_shead = kb.builder().create_block("ln_l1_shead");
490
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* l1_sbody = kb.builder().create_block("ln_l1_sbody");
491
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* l1_exit = kb.builder().create_block("ln_l1_exit");
492
493 // Pass 2 blocks (variance)
494
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* l2_vhead = kb.builder().create_block("ln_l2_vhead");
495
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* l2_vbody = kb.builder().create_block("ln_l2_vbody");
496
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* l2_vexit = kb.builder().create_block("ln_l2_vexit");
497
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* l2_shead = kb.builder().create_block("ln_l2_shead");
498
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* l2_sbody = kb.builder().create_block("ln_l2_sbody");
499
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* l2_exit = kb.builder().create_block("ln_l2_exit");
500
501 // Pass 3 blocks (affine)
502
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* l3_vhead = kb.builder().create_block("ln_l3_vhead");
503
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* l3_vbody = kb.builder().create_block("ln_l3_vbody");
504
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* l3_shead = kb.builder().create_block("ln_l3_shead");
505
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* l3_sbody = kb.builder().create_block("ln_l3_sbody");
506
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 BasicBlock* l3_exit = kb.builder().create_block("ln_l3_exit");
507
508 8 kb.position_at_end(entry);
509
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* zero_i64 = kb.builder().build_iconst_i64(0);
510
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* one_i64 = kb.builder().build_iconst_i64(1);
511
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* two_i64 = kb.builder().build_iconst_i64(2);
512
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* eight_i64 = kb.builder().build_iconst_i64(8);
513
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* zero_f32 = kb.builder().build_fconst_f32(0.0f);
514
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vzero = kb.vzero(Type::f32x8());
515
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* mask_eight = kb.builder().build_iconst_i64(~int64_t(7));
516
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vec_d = kb.builder().build_and(d, mask_eight);
517
518
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* has_vec = kb.builder().build_sge(d, eight_i64);
519
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br_if(has_vec, l1_vhead, {zero_i64, vzero}, l1_shead, {zero_i64, zero_f32});
520
521 // Pass 1: Vector sum
522
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(l1_vhead);
523
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* iv1_v = kb.builder().add_block_param(l1_vhead, Type::i64());
524
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vsum1 = kb.builder().add_block_param(l1_vhead, Type::f32x8());
525 8 kb.position_at_end(l1_vhead);
526
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* cond1_v = kb.builder().build_slt(iv1_v, vec_d);
527
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br_if(cond1_v, l1_vbody, {}, l1_vexit, {iv1_v, vsum1});
528
529
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(l1_vbody);
530 8 kb.position_at_end(l1_vbody);
531
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* byte_off1 = kb.builder().build_shl(iv1_v, two_i64);
532
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* px1 = kb.builder().build_add(x, byte_off1);
533
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vx1 = kb.vload_f32x8(px1);
534
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* next_vsum1 = kb.vadd(vsum1, vx1);
535
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* next_iv1_v = kb.builder().build_add(iv1_v, eight_i64);
536
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br(l1_vhead, {next_iv1_v, next_vsum1});
537
538 // Pass 1: Vector exit -> horizontal sum
539
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(l1_vexit);
540
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* exit_iv1 = kb.builder().add_block_param(l1_vexit, Type::i64());
541
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* exit_vsum1 = kb.builder().add_block_param(l1_vexit, Type::f32x8());
542 8 kb.position_at_end(l1_vexit);
543
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* init_sum1 = reduce_vsum8(kb, y, exit_vsum1);
544
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br(l1_shead, {exit_iv1, init_sum1});
545
546 // Pass 1: Scalar remainder
547
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(l1_shead);
548
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* iv1_s = kb.builder().add_block_param(l1_shead, Type::i64());
549
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* sum1 = kb.builder().add_block_param(l1_shead, Type::f32());
550 8 kb.position_at_end(l1_shead);
551
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* cond1_s = kb.builder().build_slt(iv1_s, d);
552
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br_if(cond1_s, l1_sbody, {}, l1_exit, {sum1});
553
554
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(l1_sbody);
555 8 kb.position_at_end(l1_sbody);
556
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* xi1 = kb.load_f32_indexed(x, iv1_s, 4, 0);
557
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* next_sum1 = kb.add(sum1, xi1);
558
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* next_iv1_s = kb.builder().build_add(iv1_s, one_i64);
559
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br(l1_shead, {next_iv1_s, next_sum1});
560
561 // Pass 1: Mean computed -> enter Pass 2
562
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(l1_exit);
563
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* final_sum1 = kb.builder().add_block_param(l1_exit, Type::f32());
564 8 kb.position_at_end(l1_exit);
565
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* mean = kb.mul(final_sum1, inv_d);
566
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vmean = kb.vbroadcast(Type::f32x8(), mean);
567
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br_if(has_vec, l2_vhead, {zero_i64, vzero, mean, vmean}, l2_shead, {zero_i64, zero_f32, mean});
568
569 // Pass 2: Vector sum of squared deviations
570
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(l2_vhead);
571
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* iv2_v = kb.builder().add_block_param(l2_vhead, Type::i64());
572
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vsum2 = kb.builder().add_block_param(l2_vhead, Type::f32x8());
573
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* m2_scalar = kb.builder().add_block_param(l2_vhead, Type::f32());
574
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vm2 = kb.builder().add_block_param(l2_vhead, Type::f32x8());
575 8 kb.position_at_end(l2_vhead);
576
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* cond2_v = kb.builder().build_slt(iv2_v, vec_d);
577
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br_if(cond2_v, l2_vbody, {}, l2_vexit, {iv2_v, vsum2, m2_scalar});
578
579
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(l2_vbody);
580 8 kb.position_at_end(l2_vbody);
581
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* byte_off2 = kb.builder().build_shl(iv2_v, two_i64);
582
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* px2 = kb.builder().build_add(x, byte_off2);
583
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vx2 = kb.vload_f32x8(px2);
584
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vdev = kb.vsub(vx2, vm2);
585
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* next_vsum2 = kb.vfma(vdev, vdev, vsum2);
586
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* next_iv2_v = kb.builder().build_add(iv2_v, eight_i64);
587
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br(l2_vhead, {next_iv2_v, next_vsum2, m2_scalar, vm2});
588
589 // Pass 2: Vector exit -> horizontal sum
590
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(l2_vexit);
591
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* exit_iv2 = kb.builder().add_block_param(l2_vexit, Type::i64());
592
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* exit_vsum2 = kb.builder().add_block_param(l2_vexit, Type::f32x8());
593
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* m2_out = kb.builder().add_block_param(l2_vexit, Type::f32());
594 8 kb.position_at_end(l2_vexit);
595
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* init_sumsq = reduce_vsum8(kb, y, exit_vsum2);
596
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br(l2_shead, {exit_iv2, init_sumsq, m2_out});
597
598 // Pass 2: Scalar remainder
599
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(l2_shead);
600
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* iv2_s = kb.builder().add_block_param(l2_shead, Type::i64());
601
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* sumsq = kb.builder().add_block_param(l2_shead, Type::f32());
602
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* m2_s = kb.builder().add_block_param(l2_shead, Type::f32());
603 8 kb.position_at_end(l2_shead);
604
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* cond2_s = kb.builder().build_slt(iv2_s, d);
605
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br_if(cond2_s, l2_sbody, {}, l2_exit, {sumsq, m2_s});
606
607
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(l2_sbody);
608 8 kb.position_at_end(l2_sbody);
609
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* xi2 = kb.load_f32_indexed(x, iv2_s, 4, 0);
610
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* diff2 = kb.sub(xi2, m2_s);
611
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* diff_sq2 = kb.mul(diff2, diff2);
612
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* next_sumsq = kb.add(sumsq, diff_sq2);
613
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* next_iv2_s = kb.builder().build_add(iv2_s, one_i64);
614
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br(l2_shead, {next_iv2_s, next_sumsq, m2_s});
615
616 // Pass 2 exit -> compute rstd
617
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(l2_exit);
618
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* final_sumsq = kb.builder().add_block_param(l2_exit, Type::f32());
619
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* final_mean = kb.builder().add_block_param(l2_exit, Type::f32());
620 8 kb.position_at_end(l2_exit);
621
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* var = kb.mul(final_sumsq, inv_d);
622
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* var_eps = kb.add(var, eps);
623
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 Value* rstd = kb.builder().build_call("rsqrtf", Type::f32(), {var_eps});
624
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vmean3 = kb.vbroadcast(Type::f32x8(), final_mean);
625
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vrstd3 = kb.vbroadcast(Type::f32x8(), rstd);
626
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br(l3_vhead, {zero_i64, final_mean, rstd, vmean3, vrstd3});
627
628 // Pass 3: Vector affine y[i] = gamma[i] * ((x[i] - mean) * rstd) + beta[i]
629
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(l3_vhead);
630
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* iv3_v = kb.builder().add_block_param(l3_vhead, Type::i64());
631
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* m3_s = kb.builder().add_block_param(l3_vhead, Type::f32());
632
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* r3_s = kb.builder().add_block_param(l3_vhead, Type::f32());
633
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vm3 = kb.builder().add_block_param(l3_vhead, Type::f32x8());
634
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vr3 = kb.builder().add_block_param(l3_vhead, Type::f32x8());
635 8 kb.position_at_end(l3_vhead);
636
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* cond3_v = kb.builder().build_slt(iv3_v, vec_d);
637
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br_if(cond3_v, l3_vbody, {}, l3_shead, {iv3_v, m3_s, r3_s});
638
639
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(l3_vbody);
640 8 kb.position_at_end(l3_vbody);
641
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* byte_off3 = kb.builder().build_shl(iv3_v, two_i64);
642
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* px3 = kb.builder().build_add(x, byte_off3);
643
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* pg3 = kb.builder().build_add(gamma, byte_off3);
644
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* pb3 = kb.builder().build_add(beta, byte_off3);
645
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* py3 = kb.builder().build_add(y, byte_off3);
646
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vx3 = kb.vload_f32x8(px3);
647
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vg3 = kb.vload_f32x8(pg3);
648
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vb3 = kb.vload_f32x8(pb3);
649
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vdiff3 = kb.vsub(vx3, vm3);
650
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vxhat3 = kb.vmul(vdiff3, vr3);
651
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* vy3 = kb.vfma(vg3, vxhat3, vb3);
652
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.vstore_f32x8(py3, vy3);
653
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* next_iv3_v = kb.builder().build_add(iv3_v, eight_i64);
654
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br(l3_vhead, {next_iv3_v, m3_s, r3_s, vm3, vr3});
655
656 // Pass 3: Scalar remainder
657
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(l3_shead);
658
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* iv3_s = kb.builder().add_block_param(l3_shead, Type::i64());
659
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* m3_sc = kb.builder().add_block_param(l3_shead, Type::f32());
660
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* r3_sc = kb.builder().add_block_param(l3_shead, Type::f32());
661 8 kb.position_at_end(l3_shead);
662
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* cond3_s = kb.builder().build_slt(iv3_s, d);
663
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br_if(cond3_s, l3_sbody, l3_exit);
664
665
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(l3_sbody);
666 8 kb.position_at_end(l3_sbody);
667
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* x3 = kb.load_f32_indexed(x, iv3_s, 4, 0);
668
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* g3 = kb.load_f32_indexed(gamma, iv3_s, 4, 0);
669
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* b3 = kb.load_f32_indexed(beta, iv3_s, 4, 0);
670
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* diff3 = kb.sub(x3, m3_sc);
671
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* xhat3 = kb.mul(diff3, r3_sc);
672
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* y3 = kb.builder().build_fma_f32(g3, xhat3, b3);
673
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.store_f32_indexed(y, iv3_s, y3, 4, 0);
674
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 Value* next_iv3_s = kb.builder().build_add(iv3_s, one_i64);
675
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_br(l3_shead, {next_iv3_s, m3_sc, r3_sc});
676
677
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 fn->append_block(l3_exit);
678 8 kb.position_at_end(l3_exit);
679
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kb.builder().build_ret_void();
680
681
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 return jit.compile(mod, "layernorm_row");
682 8 }
683
684 // ── Kernel Initialization & JIT Engine Management ───────────────────────────
685 2273 static const JitKernels& get_kernels() {
686
3/4
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 2265 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 8 times.
 2273 static JitKernels kernels;
687 static std::once_flag init_flag;
688 2281 std::call_once(init_flag, [] {
689 try {
690 8 KernelOptions opts;
691 8 opts.enable_optimizations = true;
692 8 opts.enable_avx2 = true;
693 8 opts.enable_fma = true;
694 8 opts.enable_vectorize = true;
695 8 opts.enable_unroll = true;
696 8 opts.unroll_factor = 4;
697 8 opts.enable_fp_reassociation = true;
698
699
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 KernelJit jit(opts);
700
701 // (1. SwiGLU runs the reference loop; see swiglu_forward below.)
702
703 // 2. Vectorized AdaLN Modulate
704
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 KernelFunction k_adaln = build_adaln_modulate(jit);
705 8 kernels.adaln_fn = k_adaln.as<JitKernels::AdaLNModulateFn>();
706
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kernels.engines.push_back(k_adaln.engine());
707
708 // 3. Vectorized Broadcast Mul
709
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 KernelFunction k_bmul = build_broadcast_mul(jit);
710 8 kernels.broadcast_mul_fn = k_bmul.as<JitKernels::BroadcastMulFn>();
711
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kernels.engines.push_back(k_bmul.engine());
712
713 // 4. Vectorized RMSNorm row
714
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 KernelFunction k_rmsnorm = build_rms_norm_row(jit);
715 8 kernels.rms_norm_row_fn = k_rmsnorm.as<JitKernels::RmsNormRowFn>();
716
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kernels.engines.push_back(k_rmsnorm.engine());
717
718 // 5. Vectorized LayerNorm row
719
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 KernelFunction k_ln = build_layernorm_row(jit);
720 8 kernels.layernorm_row_fn = k_ln.as<JitKernels::LayerNormRowFn>();
721
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kernels.engines.push_back(k_ln.engine());
722
723
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 16 kernels.available = (kernels.adaln_fn != nullptr &&
724
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 kernels.broadcast_mul_fn != nullptr &&
725
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 kernels.rms_norm_row_fn != nullptr &&
726 8 kernels.layernorm_row_fn != nullptr);
727
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 8 } catch (const std::exception& e) {
728 ✗ std::cerr << "brotensor: Brass JIT initialization failed: " << e.what() << "\n";
729 ✗ kernels.available = false;
730 ✗ }
731 8 });
732 2273 return kernels;
733 }
734
735 } // namespace
736
737 1262 bool is_jit_available() {
738 1262 return get_kernels().available;
739 }
740
741 254 void rms_norm_forward(const float* X, const float* gamma, float eps, float* Y, int B, int D) {
742 254 const auto& k = get_kernels();
743
2/4
✓ Branch 0 taken 254 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 254 times.
 254 if (!k.available || k.rms_norm_row_fn == nullptr) {
744 ✗ ref::rms_norm(X, gamma, eps, Y, B, D);
745 ✗ return;
746 }
747
748 254 const float inv_D = 1.0f / static_cast<float>(D);
749
750 // Single-thread vectorized loop executes in sub-microsecond latency.
751 // Multi-threaded pool dispatch only when workload amortizes thread wake-up overhead.
752
3/4
✓ Branch 0 taken 192 times.
✓ Branch 1 taken 62 times.
✓ Branch 2 taken 192 times.
✗ Branch 3 not taken.
 254 if (B > 1 && static_cast<int64_t>(B) * D >= 262144) {
753 ✗ detail::cpu::parallel_for(static_cast<std::size_t>(B), [&](std::size_t bi) {
754 ✗ const int b = static_cast<int>(bi);
755 ✗ const float* xr = X + static_cast<std::size_t>(b) * D;
756 ✗ float* yr = Y + static_cast<std::size_t>(b) * D;
757 ✗ k.rms_norm_row_fn(xr, gamma, yr, static_cast<uint64_t>(D), eps, inv_D);
758 ✗ });
759 ✗ } else {
760
2/2
✓ Branch 0 taken 2769 times.
✓ Branch 1 taken 254 times.
 3023 for (int b = 0; b < B; ++b) {
761 2769 const float* xr = X + static_cast<std::size_t>(b) * D;
762 2769 float* yr = Y + static_cast<std::size_t>(b) * D;
763 2769 k.rms_norm_row_fn(xr, gamma, yr, static_cast<uint64_t>(D), eps, inv_D);
764 2769 }
765 }
766 254 }
767
768 // SwiGLU has no JIT kernel. brass's MlFusionCompiler::compile_swiglu evaluates
769 // exp(-g) as a degree-5 Taylor polynomial of -g/16 raised to the 16th power;
770 // the polynomial turns negative once -g/16 passes about -3.5, so every gate
771 // above ~55 came out with silu(g) ~ 0 instead of g, and ViT-H SwiGLU FFNs
772 // (DINOv3) put gates well past that. A JIT kernel calling the exact expf per
773 // element measured ~4x slower than this loop (MSVC vectorizes std::exp), so
774 // the row-parallel reference is the kernel.
775 247 void swiglu_forward(const float* X, float* Y, int B, int D) {
776 247 ref::swiglu(X, Y, B, D);
777 247 }
778
779 // modulate and broadcast_mul are one load-op-store stream per element, which
780 // the compiler already vectorizes: on AArch64 clang's NEON loop is unrolled
781 // and runs these about 2x faster than the JIT's f32x8 (a q-register pair per
782 // iteration, not unrolled), so there the row-parallel reference is the kernel.
783 // The norms keep the JIT everywhere: their reductions are where it wins.
784 #if defined(__aarch64__) || defined(_M_ARM64)
785 constexpr bool kJitStreamingKernels = false;
786 #else
787 constexpr bool kJitStreamingKernels = true;
788 #endif
789
790 246 void modulate(const float* X, const float* scale, const float* shift, float* Y, int L, int D) {
791 246 const auto& k = get_kernels();
792 if (!kJitStreamingKernels || !k.available || k.adaln_fn == nullptr) {
793 246 ref::modulate(X, scale, shift, Y, L, D);
794 246 return;
795 }
796
797 if (L > 1 && static_cast<int64_t>(L) * D >= 262144) {
798 detail::cpu::parallel_for(static_cast<std::size_t>(L), [&](std::size_t li) {
799 const int l = static_cast<int>(li);
800 const float* xr = X + static_cast<std::size_t>(l) * D;
801 float* yr = Y + static_cast<std::size_t>(l) * D;
802 k.adaln_fn(xr, scale, shift, yr, static_cast<uint64_t>(D));
803 });
804 } else {
805 for (int l = 0; l < L; ++l) {
806 const float* xr = X + static_cast<std::size_t>(l) * D;
807 float* yr = Y + static_cast<std::size_t>(l) * D;
808 k.adaln_fn(xr, scale, shift, yr, static_cast<uint64_t>(D));
809 }
810 }
811 }
812
813 245 void broadcast_mul(const float* X, const float* v, float* Y, int L, int D) {
814 245 const auto& k = get_kernels();
815 if (!kJitStreamingKernels || !k.available || k.broadcast_mul_fn == nullptr) {
816 245 ref::broadcast_mul(X, v, Y, L, D);
817 245 return;
818 }
819
820 if (L > 1 && static_cast<int64_t>(L) * D >= 262144) {
821 detail::cpu::parallel_for(static_cast<std::size_t>(L), [&](std::size_t li) {
822 const int l = static_cast<int>(li);
823 const float* xr = X + static_cast<std::size_t>(l) * D;
824 float* yr = Y + static_cast<std::size_t>(l) * D;
825 k.broadcast_mul_fn(xr, v, yr, static_cast<uint64_t>(D));
826 });
827 } else {
828 for (int l = 0; l < L; ++l) {
829 const float* xr = X + static_cast<std::size_t>(l) * D;
830 float* yr = Y + static_cast<std::size_t>(l) * D;
831 k.broadcast_mul_fn(xr, v, yr, static_cast<uint64_t>(D));
832 }
833 }
834 }
835
836 266 void layernorm_forward_inference_batched(const float* X, const float* gamma, const float* beta, float eps, float* Y, int R, int D) {
837 266 const auto& k = get_kernels();
838
2/4
✓ Branch 0 taken 266 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 266 times.
 266 if (!k.available || k.layernorm_row_fn == nullptr) {
839 ✗ ref::layernorm_forward_inference_batched(X, gamma, beta, eps, Y, R, D);
840 ✗ return;
841 }
842
843 266 const float inv_D = 1.0f / static_cast<float>(D);
844
845
3/4
✓ Branch 0 taken 188 times.
✓ Branch 1 taken 78 times.
✓ Branch 2 taken 188 times.
✗ Branch 3 not taken.
 266 if (R > 1 && static_cast<int64_t>(R) * D >= 262144) {
846 ✗ detail::cpu::parallel_for(static_cast<std::size_t>(R), [&](std::size_t ri) {
847 ✗ const int r = static_cast<int>(ri);
848 ✗ const float* xr = X + static_cast<std::size_t>(r) * D;
849 ✗ float* yr = Y + static_cast<std::size_t>(r) * D;
850 ✗ k.layernorm_row_fn(xr, gamma, beta, yr, static_cast<uint64_t>(D), eps, inv_D);
851 ✗ });
852 ✗ } else {
853
2/2
✓ Branch 0 taken 3183 times.
✓ Branch 1 taken 266 times.
 3449 for (int r = 0; r < R; ++r) {
854 3183 const float* xr = X + static_cast<std::size_t>(r) * D;
855 3183 float* yr = Y + static_cast<std::size_t>(r) * D;
856 3183 k.layernorm_row_fn(xr, gamma, beta, yr, static_cast<uint64_t>(D), eps, inv_D);
857 3183 }
858 }
859 266 }
860
861 } // namespace brotensor::detail::cpu::jit
862
863 #else // !BROTENSOR_HAS_BRASS_JIT
864
865 namespace brotensor::detail::cpu::jit {
866
867 bool is_jit_available() { return false; }
868 void rms_norm_forward(const float* X, const float* gamma, float eps, float* Y, int B, int D) {
869 ref::rms_norm(X, gamma, eps, Y, B, D);
870 }
871 void swiglu_forward(const float* X, float* Y, int B, int D) {
872 ref::swiglu(X, Y, B, D);
873 }
874 void modulate(const float* X, const float* scale, const float* shift, float* Y, int L, int D) {
875 ref::modulate(X, scale, shift, Y, L, D);
876 }
877 void broadcast_mul(const float* X, const float* v, float* Y, int L, int D) {
878 ref::broadcast_mul(X, v, Y, L, D);
879 }
880 void layernorm_forward_inference_batched(const float* X, const float* gamma, const float* beta, float eps, float* Y, int R, int D) {
881 ref::layernorm_forward_inference_batched(X, gamma, beta, eps, Y, R, D);
882 }
883
884 } // namespace brotensor::detail::cpu::jit
885
886 #endif // BROTENSOR_HAS_BRASS_JIT
887