GCC Code Coverage Report

Directory: brotensor/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 85.9% 122 / 0 / 142
Functions: 87.5% 7 / 0 / 8
Branches: 61.2% 101 / 0 / 165
src/jit/trace_plan.cpp
Line Branch Exec Source
1 // Device-neutral fusion plans (trace_plan.h). Moved out of the CUDA trace
2 // compiler so the Vulkan compiler builds exactly the same plans.
3
4 #include "trace_plan.h"
5
6 #include <algorithm>
7 #include <cstddef>
8 #include <cstdint>
9
10 namespace brotensor::jit::plan {
11
12 340 int elem_bytes(Dtype d) {
13
1/3
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✓ Branch 2 taken 340 times.
 340 switch (d) {
14 340 case Dtype::FP32: return 4;
15 case Dtype::FP16:
16 ✗ case Dtype::BF16: return 2;
17 ✗ default: return 0;
18 }
19 340 }
20
21 36 bool dtype_supported(Dtype d) { return elem_bytes(d) != 0; }
22
23 36 int vec_width_for(Dtype d) {
24 36 const int e = elem_bytes(d);
25
1/2
✓ Branch 0 taken 36 times.
✗ Branch 1 not taken.
 36 return e ? 16 / e : 1;
26 }
27
28 32 void row_norm_geometry(int cols, int rows, int lanes, int& tpr, int& rpb) {
29
1/2
✓ Branch 0 taken 32 times.
✗ Branch 1 not taken.
 32 if (lanes < 1) lanes = 1;
30 32 const int per_row = (cols + lanes - 1) / lanes; // threads to cover a row
31 32 int t = 32;
32
4/4
✓ Branch 0 taken 16 times.
✓ Branch 1 taken 72 times.
✓ Branch 2 taken 56 times.
✓ Branch 3 taken 32 times.
 88 while (t < per_row && t < kRowThreads) t <<= 1;
33 32 int r = kRowThreads / t;
34 // The out-of-range exit has to be block-uniform, or the threads that stay
35 // hit a barrier the ones that left never will.
36
3/4
✓ Branch 0 taken 16 times.
✓ Branch 1 taken 16 times.
✓ Branch 2 taken 32 times.
✗ Branch 3 not taken.
 32 while (r > 1 && rows % r != 0) r >>= 1;
37 32 tpr = t;
38 32 rpb = r;
39 32 }
40
41 8 bool build_plan(const TraceDAG& dag, ElementwisePlan& plan) {
42 8 int rows = 0, cols = 0;
43 8 dag.dominant_shape(rows, cols);
44
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 8 times.
 8 if (rows <= 0 || cols <= 0) return false;
45
46 8 plan.rows = rows;
47 8 plan.cols = cols;
48 8 plan.numel = static_cast<std::int64_t>(rows) * cols;
49
50 8 int lanes = 8;
51
2/2
✓ Branch 0 taken 38 times.
✓ Branch 1 taken 8 times.
 46 for (const auto& n : dag.nodes()) {
52 38 const bool is_in = (n.op == TraceOpKind::Input);
53 38 const bool is_out = n.is_live_at_end;
54
4/4
✓ Branch 0 taken 12 times.
✓ Branch 1 taken 26 times.
✓ Branch 2 taken 10 times.
✓ Branch 3 taken 2 times.
 38 if (!is_in && !is_out) continue;
55
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 36 times.
 36 if (!dtype_supported(n.dtype)) return false;
56
57 36 const BroadcastKind bk = broadcast_of(n.rows, n.cols, rows, cols);
58
3/4
✓ Branch 0 taken 22 times.
✓ Branch 1 taken 14 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 22 times.
 36 if (bk == BroadcastKind::Full &&
59 22 static_cast<std::int64_t>(n.rows) * n.cols != plan.numel) {
60 ✗ return false;
61 }
62
3/4
✓ Branch 0 taken 10 times.
✓ Branch 1 taken 26 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 10 times.
 36 if (is_out && bk != BroadcastKind::Full) return false;
63
64 36 lanes = std::min(lanes, vec_width_for(n.dtype));
65
2/2
✓ Branch 0 taken 22 times.
✓ Branch 1 taken 14 times.
 36 if (bk == BroadcastKind::Row) plan.any_row_bcast = true;
66
1/2
✓ Branch 0 taken 36 times.
✗ Branch 1 not taken.
 36 if (bk == BroadcastKind::Col) plan.any_col_bcast = true;
67
68 36 const BufferSpec spec{n.id, n.dtype, bk};
69
2/2
✓ Branch 0 taken 10 times.
✓ Branch 1 taken 26 times.
 36 if (is_in) plan.inputs.push_back(spec);
70
2/2
✓ Branch 0 taken 26 times.
✓ Branch 1 taken 10 times.
 36 if (is_out) plan.outputs.push_back(spec);
71 }
72
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 if (plan.outputs.empty()) return false;
73
74
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (plan.numel % lanes != 0) lanes = 1;
75
2/4
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 if ((plan.any_row_bcast || plan.any_col_bcast) && cols % lanes != 0) lanes = 1;
76 8 plan.vec = lanes;
77 8 return true;
78 8 }
79
80 8 bool build_row_norm_plan(const TraceDAG& dag, RowNormPlan& plan) {
81 8 const auto& nodes = dag.nodes();
82 8 const TraceNode* norm = nullptr;
83
2/2
✓ Branch 0 taken 38 times.
✓ Branch 1 taken 8 times.
 46 for (const auto& n : nodes) {
84
4/4
✓ Branch 0 taken 36 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 6 times.
✓ Branch 3 taken 30 times.
 38 if (n.op == TraceOpKind::RMSNorm || n.op == TraceOpKind::LayerNorm) {
85
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 if (norm) return false; // more than one reduction
86 8 norm = &n;
87 8 }
88 }
89
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 8 times.
 8 if (!norm || norm->inputs.empty()) return false;
90
91
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!build_plan(dag, plan.ew)) return false;
92
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 if (plan.ew.cols != norm->cols) return false;
93
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 if (plan.ew.rows != norm->rows) return false;
94
95 22 auto input_index = [&](int node_id) {
96
1/2
✓ Branch 0 taken 42 times.
✗ Branch 1 not taken.
 42 for (std::size_t i = 0; i < plan.ew.inputs.size(); ++i) {
97
2/2
✓ Branch 0 taken 14 times.
✓ Branch 1 taken 28 times.
 42 if (plan.ew.inputs[i].node_id == node_id) return static_cast<int>(i);
98 28 }
99 ✗ return -1;
100 14 };
101
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 if (norm->inputs.size() > 1) {
102 8 plan.gamma_input = input_index(norm->inputs[1]);
103
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (plan.gamma_input < 0) return false;
104 8 }
105
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 6 times.
 8 if (norm->inputs.size() > 2) {
106 6 plan.beta_input = input_index(norm->inputs[2]);
107
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (plan.beta_input < 0) return false;
108 6 }
109
110 8 plan.x_node = norm->inputs[0];
111 8 plan.norm_node = norm->id;
112 8 plan.reduce = (norm->op == TraceOpKind::LayerNorm) ? RowReduce::Mean : RowReduce::RMS;
113
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 plan.eps = norm->scalar > 0.0f ? norm->scalar : 1e-5f;
114
115 // Reverse reachability from the reduction's data input. Node ids are
116 // assigned in creation order, which is topological, so one backward sweep
117 // is enough.
118 8 std::vector<char> in_pre(nodes.size(), 0);
119
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 if (plan.x_node >= 0) in_pre[static_cast<std::size_t>(plan.x_node)] = 1;
120
2/2
✓ Branch 0 taken 38 times.
✓ Branch 1 taken 8 times.
 46 for (std::size_t i = nodes.size(); i-- > 0;) {
121
2/2
✓ Branch 0 taken 22 times.
✓ Branch 1 taken 16 times.
 38 if (!in_pre[i]) continue;
122
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 16 times.
 24 for (int src : nodes[i].inputs) {
123
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 if (src >= 0) in_pre[static_cast<std::size_t>(src)] = 1;
124 }
125 }
126 8 in_pre[static_cast<std::size_t>(norm->id)] = 0;
127
128
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 38 times.
 46 for (const auto& n : nodes) {
129
2/2
✓ Branch 0 taken 12 times.
✓ Branch 1 taken 26 times.
 38 if (n.op == TraceOpKind::Input) continue;
130
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 4 times.
 12 if (n.id == norm->id) continue;
131
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
 4 if (in_pre[static_cast<std::size_t>(n.id)]) {
132
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 plan.pre_ids.push_back(n.id);
133 4 } else {
134 ✗ plan.post_ids.push_back(n.id);
135 }
136 }
137
138
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 plan.pre_input.assign(plan.ew.inputs.size(), 0);
139
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 26 times.
 34 for (std::size_t i = 0; i < plan.ew.inputs.size(); ++i) {
140 26 const int id = plan.ew.inputs[i].node_id;
141 26 plan.pre_input[i] = in_pre[static_cast<std::size_t>(id)] ? 1 : 0;
142 26 }
143
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 plan.pre_output.assign(plan.ew.outputs.size(), 0);
144
2/2
✓ Branch 0 taken 10 times.
✓ Branch 1 taken 8 times.
 18 for (std::size_t i = 0; i < plan.ew.outputs.size(); ++i) {
145 10 const int id = plan.ew.outputs[i].node_id;
146 10 plan.pre_output[i] = in_pre[static_cast<std::size_t>(id)] ? 1 : 0;
147 10 }
148
149 // Pass two can skip the pre chain entirely when its result is one of the
150 // trace's own outputs — pass one writes it, pass two reads it back — and
151 // when nothing downstream of the reduction needs any *other* pre value.
152 8 plan.x_output = -1;
153
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 4 times.
 8 if (!plan.pre_ids.empty()) {
154 4 int x_out = -1;
155
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 4 times.
 10 for (std::size_t i = 0; i < plan.ew.outputs.size(); ++i) {
156
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 2 times.
 6 if (plan.ew.outputs[i].node_id == plan.x_node) x_out = static_cast<int>(i);
157 6 }
158 4 bool post_only_needs_x = true;
159
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
 4 for (int pid : plan.post_ids) {
160 ✗ for (int src : nodes[static_cast<std::size_t>(pid)].inputs) {
161 ✗ if (src < 0) continue;
162 ✗ if (src == plan.x_node) continue;
163 ✗ if (nodes[static_cast<std::size_t>(src)].op == TraceOpKind::Input) continue;
164 ✗ if (in_pre[static_cast<std::size_t>(src)]) post_only_needs_x = false;
165 }
166 }
167
3/4
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
 4 if (x_out >= 0 && post_only_needs_x) plan.x_output = x_out;
168 4 }
169
170 // Without the reload, pass two replays the pre chain from its inputs —
171 // and pass one has already overwritten whatever it stored. An in-place
172 // pre chain (`x += p` with x a pass-one output) would replay as
173 // x + p + p. Only the reload form is correct once pass one stores.
174
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 6 times.
 8 if (plan.x_output < 0) {
175
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 6 times.
 12 for (char w : plan.pre_output) {
176
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
 6 if (w) return false;
177 }
178 6 }
179
180 // What pass two loads: with the reload it is the reduction's gamma/beta
181 // plus whatever the post chain reads; without it, everything.
182
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 plan.post_input.assign(plan.ew.inputs.size(), plan.x_output >= 0 ? 0 : 1);
183
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 2 times.
 8 if (plan.x_output >= 0) {
184
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (plan.gamma_input >= 0) plan.post_input[static_cast<std::size_t>(plan.gamma_input)] = 1;
185
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (plan.beta_input >= 0) plan.post_input[static_cast<std::size_t>(plan.beta_input)] = 1;
186
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 for (int pid : plan.post_ids) {
187 ✗ for (int src : nodes[static_cast<std::size_t>(pid)].inputs) {
188 ✗ if (src < 0) continue;
189 ✗ for (std::size_t i = 0; i < plan.ew.inputs.size(); ++i) {
190 ✗ if (plan.ew.inputs[i].node_id == src) plan.post_input[i] = 1;
191 ✗ }
192 }
193 }
194 2 }
195
196 // The row kernel's stride loop takes `vec` columns per thread per step.
197 // build_plan already reduced `vec` to the narrowest dtype in play; the row
198 // form additionally needs D to divide evenly so a vector never straddles a
199 // row boundary.
200
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (plan.ew.cols % plan.ew.vec != 0) plan.ew.vec = 1;
201 8 return true;
202 8 }
203
204 ✗ bool all_fp32(const TraceDAG& dag) {
205 ✗ for (const auto& n : dag.nodes()) {
206 ✗ if (n.dtype != Dtype::FP32) return false;
207 }
208 ✗ return true;
209 ✗ }
210
211 } // namespace brotensor::jit::plan
212