GCC Code Coverage Report

Directory: brotensor/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 80.6% 429 / 0 / 532
Functions: 73.3% 22 / 0 / 30
Branches: 39.8% 379 / 0 / 952
src/jit/trace_compiler_cpu.cpp
Line Branch Exec Source
1 #include "trace_compiler.h"
2 #include "trace_dag.h"
3 #include "trace_cache.h"
4 #include "../cpu/cpu_jit.h"
5 #include <brotensor/detail/cpu/thread_pool.h>
6
7 #if BROTENSOR_HAS_BRASS_JIT
8 #include <brass/codegen/kernel_jit.hpp>
9 #include <brass/codegen/trace_builder.hpp>
10 #include <brass/mir/builder.hpp>
11 #include <brass/mir/module.hpp>
12 #include <brass/mir/function.hpp>
13 #include <brass/mir/types.hpp>
14 #endif
15
16 #include <chrono>
17 #include <cmath>
18 #include <string>
19 #include <vector>
20 #include <algorithm>
21 #include <stdexcept>
22 #include <iostream>
23
24 namespace brotensor::jit {
25
26 10 FusionPattern TraceCompiler::classify_dag(const TraceDAG& dag) {
27 10 const auto& nodes = dag.nodes();
28
29 // Check for Fused Residual + RMSNorm: in-place Add followed by RMSNorm consuming it
30
2/2
✓ Branch 0 taken 60 times.
✓ Branch 1 taken 9 times.
 69 for (const auto& n : nodes) {
31
3/4
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 59 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 1 time.
 60 if (n.op == TraceOpKind::RMSNorm && !n.inputs.empty()) {
32 1 int in_id = n.inputs[0];
33
2/4
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 1 time.
 1 if (in_id >= 0 && in_id < static_cast<int>(nodes.size())) {
34 1 const auto& prev = nodes[in_id];
35
2/4
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 1 time.
 1 if (prev.op == TraceOpKind::Add && prev.is_inplace) {
36 1 return FusionPattern::ResidualRMSNorm;
37 }
38 ✗ }
39 ✗ }
40 }
41
42 // Check for Fused LayerNorm + Modulate: LayerNorm followed by Modulate / Affine
43
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 50 times.
 58 for (const auto& n : nodes) {
44
2/2
✓ Branch 0 taken 49 times.
✓ Branch 1 taken 1 time.
 50 if (n.op == TraceOpKind::LayerNorm) {
45
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
 1 for (int out_id : n.outputs) {
46
2/4
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 1 time.
 1 if (out_id >= 0 && out_id < static_cast<int>(nodes.size())) {
47
2/4
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
 1 if (nodes[out_id].op == TraceOpKind::Modulate || nodes[out_id].op == TraceOpKind::Mul) {
48 1 return FusionPattern::LayerNormModulate;
49 }
50 ✗ }
51 }
52 ✗ }
53 }
54
55 8 return FusionPattern::Elementwise;
56 10 }
57
58 static CudaTraceCompilerFn s_cuda_trace_compiler_fn = nullptr;
59
60 ✗ CudaTraceCompilerFn get_cuda_trace_compiler_hook() {
61 ✗ return s_cuda_trace_compiler_fn;
62 }
63
64 ✗ void register_cuda_trace_compiler(CudaTraceCompilerFn fn) {
65 ✗ s_cuda_trace_compiler_fn = fn;
66 ✗ }
67
68 static VulkanTraceCompilerFn s_vulkan_trace_compiler_fn = nullptr;
69
70 ✗ VulkanTraceCompilerFn get_vulkan_trace_compiler_hook() {
71 ✗ return s_vulkan_trace_compiler_fn;
72 }
73
74 ✗ void register_vulkan_trace_compiler(VulkanTraceCompilerFn fn) {
75 ✗ s_vulkan_trace_compiler_fn = fn;
76 ✗ }
77
78 11 TraceHandle TraceCompiler::compile_and_cache(const TraceDAG& dag) {
79 11 uint64_t hash = dag.compute_hash();
80
81 // Check trace cache for previously compiled kernel
82 11 auto cached = TraceCache::instance().lookup(hash);
83
3/4
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 1 time.
 11 if (cached && cached->rebind_fn) {
84 1 std::vector<void*> input_ptrs;
85 1 std::vector<void*> output_ptrs;
86
2/2
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 8 times.
 9 for (const auto& node : dag.nodes()) {
87
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 4 times.
 8 if (node.op == TraceOpKind::Input) {
88
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 input_ptrs.push_back(node.buffer);
89 4 }
90
2/2
✓ Branch 0 taken 7 times.
✓ Branch 1 taken 1 time.
 8 if (node.is_live_at_end) {
91
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 output_ptrs.push_back(node.buffer);
92 1 }
93 }
94
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 auto re_bound = cached->rebind_fn(input_ptrs, output_ptrs);
95 1 re_bound->signature_hash = hash;
96 1 re_bound->is_cache_hit = true;
97 1 re_bound->node_count = dag.node_count();
98
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 re_bound->execute();
99
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 return TraceHandle(re_bound);
100 1 }
101
102 10 FusionPattern pattern = classify_dag(dag);
103
104 // Which compiler gets the DAG. CUDA has the PTX compiler and Vulkan the
105 // SPIR-V one (brass MIR); a Vulkan DAG must never reach the CPU compiler
106 // (host code over device addresses, unordered against the device
107 // stream), so what the Vulkan compiler cannot fuse is replayed op by op
108 // instead. Everything else is the CPU compiler's.
109 10 bool is_cuda = false;
110 10 bool is_vulkan = false;
111
2/2
✓ Branch 0 taken 60 times.
✓ Branch 1 taken 10 times.
 70 for (const auto& n : dag.nodes()) {
112
2/4
✓ Branch 0 taken 60 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 60 times.
✗ Branch 3 not taken.
 60 if (n.device.is_cuda()) is_cuda = true;
113
2/4
✓ Branch 0 taken 60 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 60 times.
✗ Branch 3 not taken.
 60 if (n.device.is_vulkan()) is_vulkan = true;
114 }
115
116 10 const auto compile_t0 = std::chrono::steady_clock::now();
117 10 std::shared_ptr<TraceHandleImpl> handle;
118
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 10 times.
 10 if (is_cuda) {
119 ✗ auto cuda_fn = get_cuda_trace_compiler_hook();
120 ✗ if (!cuda_fn) {
121 ✗ throw std::runtime_error("brotensor::jit: CUDA JIT compiler hook is not registered or CUDA is unavailable");
122 }
123 ✗ handle = cuda_fn(dag, pattern);
124
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 } else if (is_vulkan) {
125 ✗ if (auto vk_fn = get_vulkan_trace_compiler_hook()) handle = vk_fn(dag, pattern);
126 ✗ if (!handle) handle = eager::compile_eager(dag);
127 ✗ } else {
128
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 handle = cpu::compile_cpu(dag, pattern);
129 }
130
131 10 handle->signature_hash = hash;
132 10 handle->is_cuda = is_cuda;
133 10 handle->node_count = dag.node_count();
134 10 handle->is_cache_hit = false;
135 // The CUDA and Vulkan compilers name their own fusion and time their own
136 // builds, the eager replay counts its own launches; the CPU path is always
137 // one fused call through brass's host codegen.
138
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 8 times.
 10 if (handle->launch_count == 0) handle->launch_count = 1;
139
2/4
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 10 times.
✗ Branch 3 not taken.
 10 if (!is_cuda && !is_vulkan) {
140
2/4
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 10 times.
✗ Branch 3 not taken.
 10 if (handle->fusion_name == nullptr || handle->fusion_name[0] == '\0') {
141 ✗ handle->fusion_name = "cpu-avx2-fused";
142 ✗ }
143
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 20 handle->compile_us = std::chrono::duration<double, std::micro>(
144
2/4
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 10 times.
✗ Branch 3 not taken.
 10 std::chrono::steady_clock::now() - compile_t0).count();
145 10 }
146
147
2/4
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 10 times.
✗ Branch 3 not taken.
 10 TraceCache::instance().insert(hash, handle);
148
149 // Initial execution to populate outputs
150
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 handle->execute();
151
152
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 return TraceHandle(handle);
153 11 }
154
155 11 TraceHandle end_trace() {
156 11 TraceDAG dag = TraceContext::current().end();
157
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 return TraceCompiler::compile_and_cache(dag);
158 11 }
159
160 namespace cpu {
161
162 #if BROTENSOR_HAS_BRASS_JIT
163
164 using ElementwiseFn = void (*)(float* const* out_ptrs, const float* const* in_ptrs, int64_t n);
165
166 // What the host codegen below can actually be handed.
167 //
168 // It walks every operand as a flat run of `total_numel` FP32 values. A (1, D)
169 // broadcast row or a half-precision buffer would be read past the end of its
170 // allocation — a crash, or worse, silently wrong numbers — so anything that is
171 // not the trace's full shape in FP32 is refused here. end_trace() then throws
172 // and the caller falls back to its eager path, which is the contract the CUDA
173 // compiler already follows for expressions it cannot fuse.
174 6 static void require_cpu_fusable(const TraceDAG& dag) {
175 6 int rows = 0, cols = 0;
176 6 dag.dominant_shape(rows, cols);
177 6 const int64_t dom_numel = static_cast<int64_t>(rows) * cols;
178
2/2
✓ Branch 0 taken 40 times.
✓ Branch 1 taken 6 times.
 46 for (const TraceNode& n : dag.nodes()) {
179
1/2
✓ Branch 0 taken 40 times.
✗ Branch 1 not taken.
 40 if (n.dtype != Dtype::FP32) {
180 ✗ throw std::runtime_error(
181 "brotensor::jit: the CPU trace compiler is FP32 only");
182 }
183 40 const BroadcastKind bk = broadcast_of(n.rows, n.cols, rows, cols);
184 40 const int64_t numel = static_cast<int64_t>(n.rows) * n.cols;
185
3/4
✓ Branch 0 taken 37 times.
✓ Branch 1 taken 3 times.
✓ Branch 2 taken 37 times.
✗ Branch 3 not taken.
 40 if (bk == BroadcastKind::Full && numel != dom_numel) {
186 ✗ throw std::runtime_error(
187 ✗ "brotensor::jit: operand shape (" + std::to_string(n.rows) + ", " +
188 ✗ std::to_string(n.cols) + ") cannot broadcast to dominant shape (" +
189 ✗ std::to_string(rows) + ", " + std::to_string(cols) + ")");
190 }
191
3/4
✓ Branch 0 taken 7 times.
✓ Branch 1 taken 33 times.
✓ Branch 2 taken 7 times.
✗ Branch 3 not taken.
 40 if (n.is_live_at_end && bk != BroadcastKind::Full) {
192 ✗ throw std::runtime_error(
193 "brotensor::jit: a traced output must have the trace's full (rows, cols) shape");
194 }
195 }
196 6 }
197
198 // The emitted code dereferences every buffer on the host. CUDA device
199 // pointers and Vulkan buffer device addresses are not host addresses (and
200 // even where one happened to be readable, the host code would be unordered
201 // against the device stream that produces its inputs). Metal storage is
202 // MTLResourceStorageModeShared, whose contents pointer is a host address.
203 10 static void require_host_addressable(const TraceDAG& dag) {
204
2/2
✓ Branch 0 taken 60 times.
✓ Branch 1 taken 10 times.
 70 for (const TraceNode& n : dag.nodes()) {
205
1/2
✓ Branch 0 taken 60 times.
✗ Branch 1 not taken.
 60 if (n.device.is_cuda() || n.device.is_vulkan()) {
206 ✗ throw std::runtime_error(
207 ✗ "brotensor::jit: the CPU trace compiler was handed a " +
208 ✗ to_string(n.device) + " tensor; it only runs host-addressable memory");
209 }
210 }
211 10 }
212
213 6 static std::shared_ptr<TraceHandleImpl> compile_cpu_elementwise(const TraceDAG& dag) {
214 6 require_cpu_fusable(dag);
215 6 const auto& nodes = dag.nodes();
216
217 6 int dom_rows = 0, dom_cols = 0;
218 6 dag.dominant_shape(dom_rows, dom_cols);
219
220 6 std::vector<int> input_node_ids;
221 6 std::vector<BroadcastKind> input_bcast;
222 6 std::vector<int> output_node_ids;
223 6 std::unordered_map<int, int> node_to_in_idx;
224 6 std::unordered_map<int, int> node_to_out_idx;
225
226 6 bool any_bcast = false;
227 6 bool any_row_bcast = false;
228 6 bool any_col_bcast = false;
229 6 bool has_where = false;
230
231
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 40 times.
 46 for (const auto& n : nodes) {
232
2/2
✓ Branch 0 taken 19 times.
✓ Branch 1 taken 21 times.
 40 if (n.op == TraceOpKind::Input) {
233
1/2
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
 19 node_to_in_idx[n.id] = static_cast<int>(input_node_ids.size());
234
1/2
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
 19 input_node_ids.push_back(n.id);
235
1/2
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
 19 const BroadcastKind bk = broadcast_of(n.rows, n.cols, dom_rows, dom_cols);
236
1/2
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
 19 input_bcast.push_back(bk);
237
2/2
✓ Branch 0 taken 16 times.
✓ Branch 1 taken 3 times.
 19 if (bk != BroadcastKind::Full) any_bcast = true;
238
2/2
✓ Branch 0 taken 18 times.
✓ Branch 1 taken 1 time.
 19 if (bk == BroadcastKind::Row) any_row_bcast = true;
239
2/2
✓ Branch 0 taken 17 times.
✓ Branch 1 taken 2 times.
 19 if (bk == BroadcastKind::Col) any_col_bcast = true;
240 19 }
241
2/2
✓ Branch 0 taken 39 times.
✓ Branch 1 taken 1 time.
 40 if (n.op == TraceOpKind::Where) {
242 1 has_where = true;
243 1 }
244
2/2
✓ Branch 0 taken 33 times.
✓ Branch 1 taken 7 times.
 40 if (n.is_live_at_end) {
245
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 node_to_out_idx[n.id] = static_cast<int>(output_node_ids.size());
246
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 output_node_ids.push_back(n.id);
247 7 }
248 }
249
250 6 const int64_t total_numel = static_cast<int64_t>(dom_rows) * static_cast<int64_t>(dom_cols);
251
4/4
✓ Branch 0 taken 5 times.
✓ Branch 1 taken 1 time.
✓ Branch 2 taken 1 time.
✓ Branch 3 taken 4 times.
 6 const bool use_vec_loop = !any_row_bcast && !any_col_bcast && !has_where;
252
253 using namespace brass;
254 using namespace brass::codegen;
255
256 static uint64_t kernel_id_counter = 0;
257
2/4
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
 6 std::string mod_name = "mod_trace_cpu_" + std::to_string(++kernel_id_counter);
258
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 Module mod(mod_name);
259
260
5/10
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 6 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 6 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 6 times.
✗ Branch 9 not taken.
 6 Function* fn = mod.create_function("trace_elementwise_kernel", Type::void_type(), {
261 6 Type::ptr(), // float* const* out_ptrs
262 6 Type::ptr(), // const float* const* in_ptrs
263 6 Type::i64() // int64_t n
264 });
265
266
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 TraceBuilder tb(mod, fn);
267
268
2/4
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
 6 BasicBlock* entry = tb.builder().append_block("entry");
269
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 Value* p_out_ptrs = tb.builder().add_block_param(entry, Type::ptr());
270
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 Value* p_in_ptrs = tb.builder().add_block_param(entry, Type::ptr());
271
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 Value* n_val = tb.builder().add_block_param(entry, Type::i64());
272
273
2/4
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
 6 BasicBlock* shead = tb.builder().create_block("shead");
274
2/4
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
 6 BasicBlock* sbody = tb.builder().create_block("sbody");
275
2/4
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
 6 BasicBlock* exit_bb = tb.builder().create_block("exit");
276
277 6 tb.position_at_end(entry);
278
279 // Load input and output base pointers into local SSA values
280 6 std::vector<Value*> in_ptrs_val;
281
2/2
✓ Branch 0 taken 19 times.
✓ Branch 1 taken 6 times.
 25 for (size_t k = 0; k < input_node_ids.size(); ++k) {
282
1/2
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
 19 Value* p = tb.builder().build_load(Type::ptr(), p_in_ptrs, static_cast<int32_t>(k * 8));
283
1/2
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
 19 in_ptrs_val.push_back(p);
284 19 }
285 6 std::vector<Value*> out_ptrs_val;
286
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 7 times.
 13 for (size_t m = 0; m < output_node_ids.size(); ++m) {
287
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 Value* p = tb.builder().build_load(Type::ptr(), p_out_ptrs, static_cast<int32_t>(m * 8));
288
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 out_ptrs_val.push_back(p);
289 7 }
290
291
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 Value* zero_i64 = tb.builder().build_iconst_i64(0);
292
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 Value* one_i64 = tb.builder().build_iconst_i64(1);
293
294
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 2 times.
 6 if (use_vec_loop) {
295
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 4 BasicBlock* vhead = tb.builder().create_block("vhead");
296
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 4 BasicBlock* vbody = tb.builder().create_block("vbody");
297
298
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 Value* two_i64 = tb.builder().build_iconst_i64(2);
299
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 Value* eight_i64 = tb.builder().build_iconst_i64(8);
300
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 Value* mask_eight = tb.builder().build_iconst_i64(~int64_t(7));
301
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 Value* vec_n = tb.builder().build_and(n_val, mask_eight);
302
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 Value* has_vec = tb.builder().build_sge(n_val, eight_i64);
303
304
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 tb.builder().build_br_if(has_vec, vhead, {zero_i64}, shead, {zero_i64});
305
306 // ── Vector AVX2 Loop (8-wide) ─────────────────────────────────────────
307
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 fn->append_block(vhead);
308
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 Value* iv_v = tb.builder().add_block_param(vhead, Type::i64());
309 4 tb.position_at_end(vhead);
310
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 Value* cond_v = tb.builder().build_slt(iv_v, vec_n);
311
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 tb.builder().build_br_if(cond_v, vbody, {}, shead, {iv_v});
312
313
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 fn->append_block(vbody);
314 4 tb.position_at_end(vbody);
315
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 Value* byte_off_v = tb.builder().build_shl(iv_v, two_i64);
316
317 4 std::unordered_map<int, Value*> v_node_vals;
318
2/2
✓ Branch 0 taken 12 times.
✓ Branch 1 taken 4 times.
 16 for (size_t k = 0; k < input_node_ids.size(); ++k) {
319 12 int nid = input_node_ids[k];
320
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (input_bcast[k] == BroadcastKind::Scalar) {
321 ✗ Value* sv = tb.load_f32(in_ptrs_val[k], 0);
322 ✗ v_node_vals[nid] = tb.vbroadcast(Type::f32x8(), sv);
323 ✗ } else {
324
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 Value* addr = tb.builder().build_add(in_ptrs_val[k], byte_off_v);
325
2/4
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
 12 v_node_vals[nid] = tb.vload_f32x8(addr);
326 }
327 12 }
328
329
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 27 times.
 31 for (const auto& n : nodes) {
330
2/2
✓ Branch 0 taken 15 times.
✓ Branch 1 taken 12 times.
 27 if (n.op == TraceOpKind::Input) continue;
331
332 15 Value* res = nullptr;
333
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 11 times.
 15 if (n.op == TraceOpKind::Add) {
334
3/6
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 4 times.
✗ Branch 5 not taken.
 4 res = tb.vadd(v_node_vals[n.inputs[0]], v_node_vals[n.inputs[1]]);
335
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 11 times.
 15 } else if (n.op == TraceOpKind::Sub) {
336 ✗ res = tb.vsub(v_node_vals[n.inputs[0]], v_node_vals[n.inputs[1]]);
337
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 5 times.
 11 } else if (n.op == TraceOpKind::Mul) {
338
3/6
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 6 times.
✗ Branch 5 not taken.
 6 res = tb.vmul(v_node_vals[n.inputs[0]], v_node_vals[n.inputs[1]]);
339
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 5 times.
 11 } else if (n.op == TraceOpKind::Div) {
340 ✗ res = tb.vdiv(v_node_vals[n.inputs[0]], v_node_vals[n.inputs[1]]);
341
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 5 times.
 5 } else if (n.op == TraceOpKind::AddScalar) {
342 ✗ Value* c = tb.vbroadcast(Type::f32x8(), tb.builder().build_fconst_f32(n.scalar));
343 ✗ res = tb.vadd(v_node_vals[n.inputs[0]], c);
344
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 5 times.
 5 } else if (n.op == TraceOpKind::SubScalar) {
345 ✗ Value* c = tb.vbroadcast(Type::f32x8(), tb.builder().build_fconst_f32(n.scalar));
346 ✗ res = tb.vsub(v_node_vals[n.inputs[0]], c);
347
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 5 times.
 5 } else if (n.op == TraceOpKind::MulScalar) {
348 ✗ Value* c = tb.vbroadcast(Type::f32x8(), tb.builder().build_fconst_f32(n.scalar));
349 ✗ res = tb.vmul(v_node_vals[n.inputs[0]], c);
350
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 5 times.
 5 } else if (n.op == TraceOpKind::DivScalar) {
351 ✗ Value* c = tb.vbroadcast(Type::f32x8(), tb.builder().build_fconst_f32(n.scalar));
352 ✗ res = tb.vdiv(v_node_vals[n.inputs[0]], c);
353
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 5 times.
 5 } else if (n.op == TraceOpKind::ScalarSub) {
354 ✗ Value* c = tb.vbroadcast(Type::f32x8(), tb.builder().build_fconst_f32(n.scalar));
355 ✗ res = tb.vsub(c, v_node_vals[n.inputs[0]]);
356
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 5 times.
 5 } else if (n.op == TraceOpKind::ScalarDiv) {
357 ✗ Value* c = tb.vbroadcast(Type::f32x8(), tb.builder().build_fconst_f32(n.scalar));
358 ✗ res = tb.vdiv(c, v_node_vals[n.inputs[0]]);
359
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 5 times.
 5 } else if (n.op == TraceOpKind::FMA) {
360 ✗ res = tb.vfma(v_node_vals[n.inputs[0]], v_node_vals[n.inputs[1]], v_node_vals[n.inputs[2]]);
361
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 5 times.
 5 } else if (n.op == TraceOpKind::Modulate) {
362 ✗ Value* c1 = tb.vbroadcast(Type::f32x8(), tb.builder().build_fconst_f32(1.0f));
363 ✗ res = tb.vfma(v_node_vals[n.inputs[0]], tb.vadd(v_node_vals[n.inputs[1]], c1), v_node_vals[n.inputs[2]]);
364
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 5 times.
 5 } else if (n.op == TraceOpKind::Clamp) {
365 ✗ Value* lo_v = tb.vbroadcast(Type::f32x8(), tb.builder().build_fconst_f32(n.scalar));
366 ✗ Value* hi_v = tb.vbroadcast(Type::f32x8(), tb.builder().build_fconst_f32(n.scalar2));
367 ✗ res = tb.vmin(tb.vmax(v_node_vals[n.inputs[0]], lo_v), hi_v);
368
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 1 time.
 5 } else if (n.op == TraceOpKind::SiLU) {
369
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 4 res = tb.vsilu_f32x8(v_node_vals[n.inputs[0]]);
370
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
 5 } else if (n.op == TraceOpKind::GELU) {
371 ✗ res = tb.vgelu_f32x8(v_node_vals[n.inputs[0]]);
372
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
 1 } else if (n.op == TraceOpKind::ReLU) {
373 ✗ res = tb.vrelu_f32x8(v_node_vals[n.inputs[0]]);
374
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
 1 } else if (n.op == TraceOpKind::Tanh) {
375 ✗ res = tb.vtanh_f32x8(v_node_vals[n.inputs[0]]);
376
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
 1 } else if (n.op == TraceOpKind::Sigmoid) {
377 ✗ res = tb.vsigmoid_f32x8(v_node_vals[n.inputs[0]]);
378
1/4
✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 1 } else if (n.op == TraceOpKind::Copy || n.op == TraceOpKind::Cast) {
379
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 res = v_node_vals[n.inputs[0]];
380 1 } else {
381 ✗ throw std::runtime_error(
382 ✗ std::string("brotensor::jit: op '") + op_kind_name(n.op) +
383 "' has no CPU elementwise form");
384 }
385
1/2
✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
 15 v_node_vals[n.id] = res;
386
387
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 11 times.
 15 if (n.is_live_at_end) {
388
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 int out_idx = node_to_out_idx[n.id];
389
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 Value* addr = tb.builder().build_add(out_ptrs_val[out_idx], byte_off_v);
390
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 tb.vstore_f32x8(addr, res);
391 4 }
392 }
393
394
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 Value* next_iv_v = tb.builder().build_add(iv_v, eight_i64);
395
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 tb.builder().build_br(vhead, {next_iv_v});
396 4 } else {
397
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 tb.builder().build_br(shead, {zero_i64});
398 }
399
400 // ── Scalar Loop ──────────────────────────────────────────────────────────
401
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 fn->append_block(shead);
402
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 Value* iv_s = tb.builder().add_block_param(shead, Type::i64());
403 6 tb.position_at_end(shead);
404
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 Value* cond_s = tb.builder().build_slt(iv_s, n_val);
405
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 tb.builder().build_br_if(cond_s, sbody, {}, exit_bb, {});
406
407
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 fn->append_block(sbody);
408 6 tb.position_at_end(sbody);
409
410
4/4
✓ Branch 0 taken 5 times.
✓ Branch 1 taken 1 time.
✓ Branch 2 taken 1 time.
✓ Branch 3 taken 4 times.
 8 Value* cols_i64 = (any_row_bcast || any_col_bcast)
411
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 ? tb.builder().build_iconst_i64(dom_cols)
412 : nullptr;
413
3/4
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 5 times.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
 6 Value* col_s = any_row_bcast ? tb.builder().build_umod(iv_s, cols_i64) : nullptr;
414
3/4
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 4 times.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 6 Value* row_s = any_col_bcast ? tb.builder().build_udiv(iv_s, cols_i64) : nullptr;
415
416 6 std::unordered_map<int, Value*> s_node_vals;
417
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 19 times.
 25 for (size_t k = 0; k < input_node_ids.size(); ++k) {
418 19 int nid = input_node_ids[k];
419
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 19 times.
 19 if (input_bcast[k] == BroadcastKind::Scalar) {
420 ✗ s_node_vals[nid] = tb.load_f32(in_ptrs_val[k], 0);
421
2/2
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 18 times.
 19 } else if (input_bcast[k] == BroadcastKind::Row) {
422
2/4
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
 1 s_node_vals[nid] = tb.load_f32_indexed(in_ptrs_val[k], col_s, 4, 0);
423
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 16 times.
 19 } else if (input_bcast[k] == BroadcastKind::Col) {
424
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 s_node_vals[nid] = tb.load_f32_indexed(in_ptrs_val[k], row_s, 4, 0);
425 2 } else {
426
2/4
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
 16 s_node_vals[nid] = tb.load_f32_indexed(in_ptrs_val[k], iv_s, 4, 0);
427 }
428 19 }
429
430
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 40 times.
 46 for (const auto& n : nodes) {
431
2/2
✓ Branch 0 taken 21 times.
✓ Branch 1 taken 19 times.
 40 if (n.op == TraceOpKind::Input) continue;
432
433 21 Value* res = nullptr;
434
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 15 times.
 21 if (n.op == TraceOpKind::Add) {
435
3/6
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 6 times.
✗ Branch 5 not taken.
 6 res = tb.add(s_node_vals[n.inputs[0]], s_node_vals[n.inputs[1]]);
436
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 15 times.
 21 } else if (n.op == TraceOpKind::Sub) {
437 ✗ res = tb.sub(s_node_vals[n.inputs[0]], s_node_vals[n.inputs[1]]);
438
2/2
✓ Branch 0 taken 7 times.
✓ Branch 1 taken 8 times.
 15 } else if (n.op == TraceOpKind::Mul) {
439
3/6
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 7 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 7 times.
✗ Branch 5 not taken.
 7 res = tb.mul(s_node_vals[n.inputs[0]], s_node_vals[n.inputs[1]]);
440
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 15 } else if (n.op == TraceOpKind::Div) {
441 ✗ res = tb.div(s_node_vals[n.inputs[0]], s_node_vals[n.inputs[1]]);
442
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 } else if (n.op == TraceOpKind::AddScalar) {
443 ✗ Value* c = tb.builder().build_fconst_f32(n.scalar);
444 ✗ res = tb.add(s_node_vals[n.inputs[0]], c);
445
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 } else if (n.op == TraceOpKind::SubScalar) {
446 ✗ Value* c = tb.builder().build_fconst_f32(n.scalar);
447 ✗ res = tb.sub(s_node_vals[n.inputs[0]], c);
448
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 } else if (n.op == TraceOpKind::MulScalar) {
449 ✗ Value* c = tb.builder().build_fconst_f32(n.scalar);
450 ✗ res = tb.mul(s_node_vals[n.inputs[0]], c);
451
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 } else if (n.op == TraceOpKind::DivScalar) {
452 ✗ Value* c = tb.builder().build_fconst_f32(n.scalar);
453 ✗ res = tb.div(s_node_vals[n.inputs[0]], c);
454
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 } else if (n.op == TraceOpKind::ScalarSub) {
455 ✗ Value* c = tb.builder().build_fconst_f32(n.scalar);
456 ✗ res = tb.sub(c, s_node_vals[n.inputs[0]]);
457
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 } else if (n.op == TraceOpKind::ScalarDiv) {
458 ✗ Value* c = tb.builder().build_fconst_f32(n.scalar);
459 ✗ res = tb.div(c, s_node_vals[n.inputs[0]]);
460
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 } else if (n.op == TraceOpKind::FMA) {
461 ✗ res = tb.fma(s_node_vals[n.inputs[0]], s_node_vals[n.inputs[1]], s_node_vals[n.inputs[2]]);
462
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 } else if (n.op == TraceOpKind::Modulate) {
463 ✗ Value* c1 = tb.builder().build_fconst_f32(1.0f);
464 ✗ res = tb.fma(s_node_vals[n.inputs[0]], tb.add(s_node_vals[n.inputs[1]], c1), s_node_vals[n.inputs[2]]);
465
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 6 times.
 8 } else if (n.op == TraceOpKind::Clamp) {
466
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 Value* lo_c = tb.builder().build_fconst_f32(n.scalar);
467
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 Value* hi_c = tb.builder().build_fconst_f32(n.scalar2);
468
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 4 res = tb.builder().build_fmin_f32(
469
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 tb.builder().build_fmax_f32(s_node_vals[n.inputs[0]], lo_c), hi_c);
470
2/2
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 5 times.
 8 } else if (n.op == TraceOpKind::Where) {
471
2/4
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
 1 Value* c_bits = tb.builder().build_bitcast_i32_f32(s_node_vals[n.inputs[0]]);
472
2/4
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
 1 Value* mag_bits = tb.builder().build_and(c_bits, tb.builder().build_iconst_i32(0x7FFFFFFF));
473
2/4
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
 1 Value* nz = tb.builder().build_ne(mag_bits, tb.builder().build_iconst_i32(0));
474
3/6
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
✓ Branch 4 taken 1 time.
✗ Branch 5 not taken.
 1 res = tb.builder().build_select(nz, s_node_vals[n.inputs[1]], s_node_vals[n.inputs[2]]);
475
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 1 time.
 6 } else if (n.op == TraceOpKind::SiLU) {
476
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 4 res = tb.silu_f32(s_node_vals[n.inputs[0]]);
477
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
 5 } else if (n.op == TraceOpKind::GELU) {
478 ✗ res = tb.gelu_f32(s_node_vals[n.inputs[0]]);
479
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
 1 } else if (n.op == TraceOpKind::ReLU) {
480 ✗ res = tb.relu_f32(s_node_vals[n.inputs[0]]);
481
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
 1 } else if (n.op == TraceOpKind::Tanh) {
482 ✗ res = tb.tanh_f32(s_node_vals[n.inputs[0]]);
483
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
 1 } else if (n.op == TraceOpKind::Sigmoid) {
484 ✗ res = tb.sigmoid_f32(s_node_vals[n.inputs[0]]);
485
1/4
✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 1 } else if (n.op == TraceOpKind::Copy || n.op == TraceOpKind::Cast) {
486
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 res = s_node_vals[n.inputs[0]];
487 1 } else {
488 ✗ throw std::runtime_error(
489 ✗ std::string("brotensor::jit: op '") + op_kind_name(n.op) +
490 "' has no CPU elementwise form");
491 }
492
1/2
✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
 21 s_node_vals[n.id] = res;
493
494
2/2
✓ Branch 0 taken 7 times.
✓ Branch 1 taken 14 times.
 21 if (n.is_live_at_end) {
495
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 int out_idx = node_to_out_idx[n.id];
496
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 tb.store_f32_indexed(out_ptrs_val[out_idx], iv_s, res, 4, 0);
497 7 }
498 }
499
500
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 Value* next_iv_s = tb.builder().build_add(iv_s, one_i64);
501
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 tb.builder().build_br(shead, {next_iv_s});
502
503
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 fn->append_block(exit_bb);
504 6 tb.position_at_end(exit_bb);
505
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 tb.builder().build_ret_void();
506
507 6 KernelOptions opts;
508 6 opts.enable_optimizations = true;
509 6 opts.enable_avx2 = true;
510 6 opts.enable_fma = true;
511 6 opts.enable_vectorize = true;
512
513
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 auto jit = std::make_shared<KernelJit>(opts);
514
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 KernelFunction kfn = jit->compile(*fn);
515
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!kfn.is_valid()) {
516 ✗ throw std::runtime_error("brotensor::jit: Failed to JIT compile CPU elementwise trace kernel");
517 }
518
519 6 auto kernel_ptr = kfn.as<ElementwiseFn>();
520 6 auto engine_ref = kfn.engine();
521
522
3/6
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 6 times.
✗ Branch 5 not taken.
 19 auto make_bound_handle = [kernel_ptr, engine_ref, total_numel, any_bcast, input_node_ids, output_node_ids]
523 (const std::vector<void*>& in_buffers, const std::vector<void*>& out_buffers) {
524 7 auto h = std::make_shared<TraceHandleImpl>();
525
526
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 std::vector<const float*> in_f(in_buffers.size());
527
2/2
✓ Branch 0 taken 23 times.
✓ Branch 1 taken 7 times.
 30 for (size_t k = 0; k < in_buffers.size(); ++k) in_f[k] = static_cast<const float*>(in_buffers[k]);
528
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 std::vector<float*> out_f(out_buffers.size());
529
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 7 times.
 15 for (size_t m = 0; m < out_buffers.size(); ++m) out_f[m] = static_cast<float*>(out_buffers[m]);
530
531
3/8
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 7 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 7 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✗ Branch 7 not taken.
 565 h->execute_fn = [kernel_ptr, engine_ref, total_numel, any_bcast, in_f, out_f]() {
532
4/4
✓ Branch 0 taken 556 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 554 times.
✓ Branch 3 taken 2 times.
 558 if (!any_bcast && total_numel >= 262144) {
533 2 auto& pool = detail::cpu::ThreadPool::instance();
534 2 int n_threads = pool.num_threads();
535 2 int64_t chunk = (total_numel + n_threads - 1) / n_threads;
536 2 chunk = (chunk + 7) & ~int64_t(7); // Align to 8-wide AVX2
537
538
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 8 detail::cpu::parallel_for(n_threads, [&](std::size_t tid_sz) {
539 6 int64_t tid = static_cast<int64_t>(tid_sz);
540 6 int64_t start = tid * chunk;
541
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
 6 if (start >= total_numel) return;
542 6 int64_t count = std::min(chunk, total_numel - start);
543
544 6 std::vector<const float*> chunk_in(in_f.size());
545
2/2
✓ Branch 0 taken 12 times.
✓ Branch 1 taken 6 times.
 18 for (size_t k = 0; k < in_f.size(); ++k) {
546 12 chunk_in[k] = in_f[k] + start;
547 12 }
548
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 std::vector<float*> chunk_out(out_f.size());
549
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 6 times.
 12 for (size_t m = 0; m < out_f.size(); ++m) {
550 6 chunk_out[m] = out_f[m] + start;
551 6 }
552
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 kernel_ptr(chunk_out.data(), chunk_in.data(), count);
553 6 });
554 2 } else {
555 556 kernel_ptr(out_f.data(), in_f.data(), total_numel);
556 }
557 558 };
558 7 return h;
559
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 };
560
561 6 std::vector<void*> initial_in_buffers;
562
3/4
✓ Branch 0 taken 19 times.
✓ Branch 1 taken 6 times.
✓ Branch 2 taken 19 times.
✗ Branch 3 not taken.
 25 for (int nid : input_node_ids) initial_in_buffers.push_back(nodes[nid].buffer);
563 6 std::vector<void*> initial_out_buffers;
564
3/4
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 7 times.
✓ Branch 2 taken 7 times.
✗ Branch 3 not taken.
 13 for (int nid : output_node_ids) initial_out_buffers.push_back(nodes[nid].buffer);
565
566
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 auto handle = make_bound_handle(initial_in_buffers, initial_out_buffers);
567
2/4
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
 7 handle->rebind_fn = [make_bound_handle](const std::vector<void*>& in_b, const std::vector<void*>& out_b) {
568 1 return make_bound_handle(in_b, out_b);
569 };
570
571 6 return handle;
572 6 }
573
574 10 std::shared_ptr<TraceHandleImpl> compile_cpu(const TraceDAG& dag, FusionPattern pattern) {
575 10 const auto& nodes = dag.nodes();
576 10 require_host_addressable(dag);
577
578
2/2
✓ Branch 0 taken 59 times.
✓ Branch 1 taken 8 times.
 67 for (const auto& n : nodes) {
579
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 57 times.
 59 if (n.dtype != Dtype::FP32) {
580 2 return eager::compile_eager(dag);
581 }
582 }
583
584
2/2
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 7 times.
 8 if (pattern == FusionPattern::ResidualRMSNorm) {
585 1 const TraceNode* node_rms = nullptr;
586 1 const TraceNode* node_add = nullptr;
587 1 const TraceNode* node_res = nullptr;
588 1 const TraceNode* node_gamma = nullptr;
589
590
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 5 times.
 5 for (const auto& n : nodes) {
591
3/4
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 1 time.
 5 if (n.op != TraceOpKind::RMSNorm || n.inputs.empty()) continue;
592 1 const TraceNode& add = nodes[n.inputs[0]];
593 // The hand kernel below is exactly `x += res; rms_norm(x, gamma)`.
594 // Anything else — no residual, no gain, or downstream ops — is not
595 // this pattern.
596
2/4
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 1 time.
 1 if (add.op != TraceOpKind::Add || add.inputs.size() < 2) continue;
597
2/4
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 1 time.
 1 if (n.inputs.size() < 2 || !n.outputs.empty()) continue;
598 1 node_rms = &n;
599 1 node_add = &add;
600 1 node_res = &nodes[add.inputs[1]];
601 1 node_gamma = &nodes[n.inputs[1]];
602 1 break;
603 }
604
605
4/8
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
✓ Branch 4 taken 1 time.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 1 time.
 1 if (node_rms && node_add && node_res && node_gamma) {
606 1 int B = node_add->rows;
607 1 int D = node_add->cols;
608
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 float eps = node_rms->scalar > 0.0f ? node_rms->scalar : 1e-5f;
609
610 2 auto make_handle = [B, D, eps](void* x_buf, const void* res_buf, const void* gamma_buf, void* y_buf) {
611 1 auto h = std::make_shared<TraceHandleImpl>();
612
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 3 h->execute_fn = [x_buf, res_buf, gamma_buf, y_buf, B, D, eps]() {
613 2 detail::cpu::jit::fused_residual_rmsnorm(
614 2 static_cast<float*>(x_buf),
615 2 static_cast<const float*>(res_buf),
616 2 static_cast<const float*>(gamma_buf),
617 2 eps,
618 2 static_cast<float*>(y_buf),
619 2 B, D
620 );
621 2 };
622 1 return h;
623
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 };
624
625 1 void* x_buf = node_add->buffer;
626 1 const void* res_buf = node_res->buffer;
627
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 const void* gamma_buf = node_gamma ? node_gamma->buffer : nullptr;
628 1 void* y_buf = node_rms->buffer;
629
630 1 auto handle = make_handle(x_buf, res_buf, gamma_buf, y_buf);
631
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 handle->rebind_fn = [make_handle](const std::vector<void*>& in_b, const std::vector<void*>& out_b) {
632 ✗ const void* res_p = in_b.size() > 1 ? in_b[1] : nullptr;
633 ✗ const void* gamma_p = in_b.size() > 2 ? in_b[2] : nullptr;
634 ✗ void* x_p = out_b.size() > 0 ? out_b[0] : in_b[0];
635 ✗ void* y_p = out_b.size() > 1 ? out_b[1] : (out_b.size() > 0 ? out_b[0] : nullptr);
636 ✗ return make_handle(x_p, res_p, gamma_p, y_p);
637 };
638 1 return handle;
639 1 }
640
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 1 time.
 7 } else if (pattern == FusionPattern::LayerNormModulate) {
641 1 const TraceNode* node_ln = nullptr;
642 1 const TraceNode* node_gamma = nullptr;
643 1 const TraceNode* node_beta = nullptr;
644 1 const TraceNode* node_scale = nullptr;
645 1 const TraceNode* node_shift = nullptr;
646 1 const TraceNode* node_out = nullptr;
647
648
2/2
✓ Branch 0 taken 9 times.
✓ Branch 1 taken 1 time.
 10 for (const auto& n : nodes) {
649
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 1 time.
 9 if (n.op == TraceOpKind::LayerNorm) {
650 1 node_ln = &n;
651
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
 1 if (n.inputs.size() > 1) node_gamma = &nodes[n.inputs[1]];
652
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
 1 if (n.inputs.size() > 2) node_beta = &nodes[n.inputs[2]];
653 1 }
654
3/4
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 1 time.
 9 if (n.is_live_at_end && n.op != TraceOpKind::LayerNorm) {
655 1 node_out = &n;
656 1 }
657
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (n.op == TraceOpKind::Modulate) {
658 ✗ if (n.inputs.size() > 1) node_scale = &nodes[n.inputs[1]];
659 ✗ if (n.inputs.size() > 2) node_shift = &nodes[n.inputs[2]];
660 ✗ }
661 }
662
663 // Find scale/shift from expression ln * (1 + scale) + shift if not directly Modulate op
664
1/4
✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 1 if (!node_scale || !node_shift) {
665
2/2
✓ Branch 0 taken 9 times.
✓ Branch 1 taken 1 time.
 10 for (const auto& n : nodes) {
666
3/4
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 1 time.
 9 if (n.op == TraceOpKind::AddScalar && n.inputs.size() > 0) {
667 1 node_scale = &nodes[n.inputs[0]];
668 1 }
669
3/4
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 1 time.
 9 if (n.op == TraceOpKind::Add && n.inputs.size() > 1) {
670 1 int rhs = n.inputs[1];
671
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
 1 if (nodes[rhs].op == TraceOpKind::Input) {
672 1 node_shift = &nodes[rhs];
673 1 }
674 1 }
675 }
676 1 }
677
678 // The hand kernel dereferences all five operands, so a trace missing
679 // any of them (a non-affine LayerNorm, a modulate with no shift) is
680 // not this pattern and falls through to the generic path.
681
6/12
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
✓ Branch 4 taken 1 time.
✗ Branch 5 not taken.
✓ Branch 6 taken 1 time.
✗ Branch 7 not taken.
✓ Branch 8 taken 1 time.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✓ Branch 11 taken 1 time.
 1 if (node_ln && node_out && node_gamma && node_beta && node_scale && node_shift) {
682 1 int R = node_ln->rows;
683 1 int D = node_ln->cols;
684
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 float eps = node_ln->scalar > 0.0f ? node_ln->scalar : 1e-5f;
685
686 2 auto make_handle = [R, D, eps](const void* x, const void* g, const void* b,
687 const void* scale, const void* shift, void* y) {
688 1 auto h = std::make_shared<TraceHandleImpl>();
689
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 3 h->execute_fn = [x, g, b, scale, shift, y, R, D, eps]() {
690 2 detail::cpu::jit::fused_layernorm_modulate(
691 2 static_cast<const float*>(x),
692 2 static_cast<const float*>(g),
693 2 static_cast<const float*>(b),
694 2 static_cast<const float*>(scale),
695 2 static_cast<const float*>(shift),
696 2 eps,
697 2 static_cast<float*>(y),
698 2 R, D
699 );
700 2 };
701 1 return h;
702
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 };
703
704 1 const void* x_buf = nodes[node_ln->inputs[0]].buffer;
705
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
 1 const void* g_buf = node_gamma ? node_gamma->buffer : nullptr;
706
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
 1 const void* b_buf = node_beta ? node_beta->buffer : nullptr;
707
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
 1 const void* s_buf = node_scale ? node_scale->buffer : nullptr;
708
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
 1 const void* sh_buf = node_shift ? node_shift->buffer : nullptr;
709 1 void* y_buf = node_out->buffer;
710
711 1 auto handle = make_handle(x_buf, g_buf, b_buf, s_buf, sh_buf, y_buf);
712
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 handle->rebind_fn = [make_handle](const std::vector<void*>& in_b, const std::vector<void*>& out_b) {
713 ✗ const void* x = in_b.size() > 0 ? in_b[0] : nullptr;
714 ✗ const void* g = in_b.size() > 1 ? in_b[1] : nullptr;
715 ✗ const void* b = in_b.size() > 2 ? in_b[2] : nullptr;
716 ✗ const void* sc = in_b.size() > 3 ? in_b[3] : nullptr;
717 ✗ const void* sh = in_b.size() > 4 ? in_b[4] : nullptr;
718 ✗ void* y = out_b.size() > 0 ? out_b[0] : nullptr;
719 ✗ return make_handle(x, g, b, sc, sh, y);
720 };
721 1 return handle;
722 1 }
723 ✗ }
724
725
2/2
✓ Branch 0 taken 40 times.
✓ Branch 1 taken 6 times.
 46 for (const auto& n : nodes) {
726
2/4
✓ Branch 0 taken 40 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 40 times.
 40 if (n.op == TraceOpKind::RMSNorm || n.op == TraceOpKind::LayerNorm) {
727 ✗ return eager::compile_eager(dag);
728 }
729 }
730
731 6 return compile_cpu_elementwise(dag);
732 10 }
733
734 #else // !BROTENSOR_HAS_BRASS_JIT
735
736 std::shared_ptr<TraceHandleImpl> compile_cpu(const TraceDAG& /*dag*/, FusionPattern /*pattern*/) {
737 throw std::runtime_error("brotensor::jit: Brass CPU JIT compiler is not enabled in this build");
738 }
739
740 #endif // BROTENSOR_HAS_BRASS_JIT
741
742 } // namespace cpu
743 } // namespace brotensor::jit
744