GCC Code Coverage Report

Directory: brotensor/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 48.9% 137 / 0 / 280
Functions: 68.3% 28 / 0 / 41
Branches: 24.6% 96 / 0 / 390
src/jit/trace_eager.cpp
Line Branch Exec Source
1 // ─── Eager trace execution — Vulkan DAGs with no fusion ─────────────────────
2 //
3 // The trace JIT compiles a DAG to host code (brass, CPU), to PTX (CUDA) or to
4 // SPIR-V (Vulkan, trace_compiler_vulkan.cpp). A Vulkan DAG the Vulkan compiler
5 // has no fusion for (two reductions, an op with no elementwise form, ...), or
6 // a build without that compiler, must still never reach the CPU compiler: its
7 // host code would walk device addresses outside the device's ordering (a
8 // Vulkan address is not even a host pointer). Such a trace is therefore
9 // executed by replaying the DAG one node at a time through the ordinary
10 // device-dispatched ops, on the current stream, in program order.
11 //
12 // That keeps every trace correct on Vulkan — same results as the eager op
13 // sequence the trace replaces, same TraceHandle replay/rebind contract — but
14 // it is not fused: one launch per traced op, and an intermediate the fused
15 // kernel would keep in registers is a scratch tensor here (allocated per
16 // execute(), released after its last consumer). launch_count() and
17 // fusion_name() ("eager-unfused") say so.
18 //
19 // Operand rules mirror the fused compilers: every value carries its node's
20 // (rows, cols); a binary operand may be the full shape, a (1, cols) row or a
21 // (1, 1) scalar; a norm/modulate gain is a length-cols vector; an operand of
22 // another dtype is cast to the node's dtype first (math is FP32 inside each
23 // op either way). Anything else is rejected when the trace is compiled, so
24 // end_trace() throws and the caller takes its eager path.
25
26 #include "trace_compiler.h"
27 #include "trace_dag.h"
28 #include "trace_cache.h"
29
30 #include <brotensor/ops.h>
31
32 #include <chrono>
33 #include <cstddef>
34 #include <cstring>
35 #include <memory>
36 #include <stdexcept>
37 #include <string>
38 #include <vector>
39
40 namespace brotensor::jit::eager {
41
42 namespace {
43
44 ✗ [[noreturn]] void reject(const TraceNode& n, const char* why) {
45 ✗ throw std::runtime_error(std::string("brotensor::jit: eager trace execution cannot run ") +
46 ✗ op_kind_name(n.op) + " on " + to_string(n.device) + ": " + why);
47 ✗ }
48
49 3 bool is_unary(TraceOpKind op) {
50
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 1 time.
 3 switch (op) {
51 case TraceOpKind::AddScalar: case TraceOpKind::SubScalar:
52 case TraceOpKind::MulScalar: case TraceOpKind::DivScalar:
53 case TraceOpKind::ScalarSub: case TraceOpKind::SiLU:
54 case TraceOpKind::GELU: case TraceOpKind::ReLU:
55 case TraceOpKind::Tanh: case TraceOpKind::Sigmoid:
56 case TraceOpKind::Copy: case TraceOpKind::Clamp:
57 case TraceOpKind::Cast:
58 2 return true;
59 default:
60 1 return false;
61 }
62 3 }
63
64 1 bool is_binary(TraceOpKind op) {
65
1/4
✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 1 return op == TraceOpKind::Add || op == TraceOpKind::Sub ||
66 ✗ op == TraceOpKind::Mul || op == TraceOpKind::Div;
67 }
68
69 14 int64_t numel(int rows, int cols) { return static_cast<int64_t>(rows) * cols; }
70
71 // Shape and op checks, done once at compile time so a trace this executor
72 // cannot run fails in end_trace() rather than on some later replay.
73 2 void validate(const TraceDAG& dag) {
74 2 const auto& nodes = dag.nodes();
75
1/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 2 if (nodes.empty()) throw std::runtime_error("brotensor::jit: empty trace");
76 2 const Device dev = nodes.front().device;
77
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 2 times.
 8 for (const TraceNode& n : nodes) {
78
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
 6 if (n.device != dev) reject(n, "the trace spans more than one device");
79
4/6
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 4 times.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 2 times.
 6 if (n.dtype != Dtype::FP32 && n.dtype != Dtype::FP16 && n.dtype != Dtype::BF16) {
80 ✗ reject(n, "only FP32 / FP16 / BF16 values are supported");
81 }
82
2/2
✓ Branch 0 taken 3 times.
✓ Branch 1 taken 3 times.
 6 if (n.op == TraceOpKind::Input) continue;
83 7 auto operand = [&](std::size_t k) -> const TraceNode& {
84
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (k >= n.inputs.size() || n.inputs[k] < 0 ||
85 4 n.inputs[k] >= static_cast<int>(nodes.size())) {
86 ✗ reject(n, "missing operand");
87 }
88 4 return nodes[static_cast<std::size_t>(n.inputs[k])];
89 };
90 3 auto vector_like = [&](const TraceNode& v) {
91 ✗ return numel(v.rows, v.cols) == n.cols && (v.rows == 1 || v.cols == 1);
92 };
93
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 1 time.
 3 if (is_unary(n.op)) {
94 2 const TraceNode& a = operand(0);
95
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (numel(a.rows, a.cols) != numel(n.rows, n.cols)) reject(n, "operand shape mismatch");
96
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 3 } else if (is_binary(n.op)) {
97
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 1 time.
 3 for (std::size_t k = 0; k < 2; ++k) {
98 2 const TraceNode& a = operand(k);
99 2 const BroadcastKind bk = broadcast_of(a.rows, a.cols, n.rows, n.cols);
100
3/4
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 1 time.
✗ Branch 2 not taken.
✓ Branch 3 taken 1 time.
 2 if (bk == BroadcastKind::Full && numel(a.rows, a.cols) != numel(n.rows, n.cols)) {
101 ✗ reject(n, "operand is neither the full shape, a row, a column nor a scalar");
102 }
103 2 }
104
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 1 } else if (n.op == TraceOpKind::Where) {
105 ✗ for (std::size_t k = 0; k < 3; ++k) {
106 ✗ const TraceNode& a = operand(k);
107 ✗ const BroadcastKind bk = broadcast_of(a.rows, a.cols, n.rows, n.cols);
108 ✗ if (bk == BroadcastKind::Full && numel(a.rows, a.cols) != numel(n.rows, n.cols)) {
109 ✗ reject(n, "Where operand is neither the full shape, a row, a column nor a scalar");
110 }
111 ✗ }
112 ✗ } else if (n.op == TraceOpKind::FMA) {
113 ✗ for (std::size_t k = 0; k < 3; ++k) {
114 ✗ if (numel(operand(k).rows, operand(k).cols) != numel(n.rows, n.cols)) {
115 ✗ reject(n, "FMA operands must be the full shape");
116 }
117 ✗ }
118 ✗ } else if (n.op == TraceOpKind::RMSNorm || n.op == TraceOpKind::LayerNorm ||
119 ✗ n.op == TraceOpKind::Modulate) {
120 ✗ if (numel(operand(0).rows, operand(0).cols) != numel(n.rows, n.cols)) {
121 ✗ reject(n, "operand shape mismatch");
122 }
123 ✗ const std::size_t want = n.op == TraceOpKind::Modulate ? 3 : n.inputs.size();
124 ✗ for (std::size_t k = 1; k < want; ++k) {
125 ✗ if (!vector_like(operand(k))) reject(n, "gain / shift must be a length-cols vector");
126 ✗ }
127 ✗ } else {
128 ✗ reject(n, "no eager form");
129 }
130 }
131 2 }
132
133 5 Tensor view_of(const TraceNode& n, void* buffer) {
134 5 return Tensor::view(n.device, buffer, n.rows, n.cols, n.dtype);
135 }
136
137 // `t` reshaped to (rows, cols) without copying (same element count).
138 4 Tensor reshaped(const Tensor& t, int rows, int cols) {
139 4 return Tensor::view(t.device, t.data, rows, cols, t.dtype);
140 }
141
142 4 Tensor as_dtype(const Tensor& t, Dtype dt, int rows, int cols) {
143
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 2 times.
 4 if (t.dtype == dt) return reshaped(t, rows, cols);
144 2 Tensor out;
145
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 brotensor::cast(t, out, dt);
146 2 out.rows = rows;
147 2 out.cols = cols;
148 2 return out;
149 6 }
150
151 1 Tensor as_dtype(const Tensor& t, Dtype dt) {
152 1 return as_dtype(t, dt, t.rows, t.cols);
153 }
154
155 // Copies `src` into `dst` unless they already are the same storage.
156 4 void assign(Tensor& dst, const Tensor& src) {
157
2/2
✓ Branch 0 taken 3 times.
✓ Branch 1 taken 1 time.
 4 if (dst.data == src.data) return;
158
3/6
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 3 times.
 3 if (dst.device.is_cpu() && src.device.is_cpu() && dst.dtype == src.dtype) {
159 3 std::memcpy(dst.host_raw_mut(), src.host_raw(), dst.bytes());
160 3 return;
161 }
162 ✗ brotensor::copy_d2d(src, 0, dst, 0, static_cast<int>(dst.size()));
163 4 }
164
165 // A (1, cols) row, (rows, 1) column, or (1, 1) scalar operand spread to the full (rows, cols).
166 1 Tensor expand(const Tensor& small, int rows, int cols, Dtype dt) {
167
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 const Dtype work_dt = small.device.is_cpu() ? Dtype::FP32 : dt;
168 1 Tensor full = Tensor::zeros_on(small.device, rows, cols, work_dt);
169
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 Tensor s = as_dtype(small, work_dt);
170
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 if (small.size() == 1) {
171 ✗ Tensor col = reshaped(full, rows * cols, 1);
172 ✗ brotensor::add_row_bias_inplace(col, reshaped(s, 1, 1));
173
2/4
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
 1 } else if (small.rows == rows && small.cols == 1) {
174
2/4
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
 1 brotensor::add_channel_bias_inplace(full, reshaped(s, rows, 1), rows, cols);
175 1 } else {
176 ✗ brotensor::add_row_bias_inplace(full, reshaped(s, 1, cols));
177 }
178
2/6
✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 1 return work_dt == dt ? full : as_dtype(full, dt);
179 1 }
180
181 ✗ class Program {
182 public:
183 4 explicit Program(const TraceDAG& dag) : nodes_(dag.nodes()) {
184
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 last_use_.assign(nodes_.size(), -1);
185
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 2 times.
 8 for (const TraceNode& n : nodes_) {
186
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 6 times.
 10 for (int in : n.inputs) {
187
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
 4 if (in >= 0) last_use_[static_cast<std::size_t>(in)] = n.id;
188 }
189 }
190
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 2 times.
 8 for (const TraceNode& n : nodes_) {
191
2/2
✓ Branch 0 taken 3 times.
✓ Branch 1 taken 3 times.
 6 if (n.op == TraceOpKind::Input) continue;
192
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 1 time.
 3 const bool needed = n.buffer != nullptr || last_use_[static_cast<std::size_t>(n.id)] >= 0;
193
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 3 times.
 3 if (needed) ++launches_;
194 }
195 4 }
196
197 2 std::size_t launches() const { return launches_; }
198
199 // Input buffers in node order, then the buffers of live-at-end nodes in
200 // node order — the order TraceCompiler::compile_and_cache rebinds in.
201 ✗ std::shared_ptr<Program> rebound(const std::vector<void*>& in,
202 const std::vector<void*>& out) const {
203 ✗ auto p = std::make_shared<Program>(*this);
204 ✗ std::size_t i = 0, o = 0;
205 ✗ for (TraceNode& n : p->nodes_) {
206 ✗ if (n.op == TraceOpKind::Input) {
207 ✗ if (i < in.size()) n.buffer = in[i++];
208 ✗ }
209 ✗ if (n.is_live_at_end) {
210 ✗ if (o < out.size()) n.buffer = out[o++];
211 ✗ }
212 }
213 ✗ return p;
214 ✗ }
215
216 2 void run() const {
217 2 std::vector<Tensor> vals(nodes_.size());
218
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 2 times.
 8 for (const TraceNode& n : nodes_) {
219 6 const std::size_t id = static_cast<std::size_t>(n.id);
220
2/2
✓ Branch 0 taken 3 times.
✓ Branch 1 taken 3 times.
 6 if (n.op == TraceOpKind::Input) {
221
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 vals[id] = view_of(n, n.buffer);
222 3 continue;
223 }
224 // A value nothing reads and nothing stores is dead code.
225
3/4
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
 3 if (n.buffer == nullptr && last_use_[id] < 0) continue;
226
227
3/4
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 1 time.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 4 Tensor dst = n.buffer ? view_of(n, n.buffer)
228
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 : Tensor::empty_on(n.device, n.rows, n.cols, n.dtype);
229
5/8
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 1 time.
✓ Branch 5 taken 2 times.
✗ Branch 6 not taken.
✓ Branch 7 taken 1 time.
 4 if (n.device.is_cpu() && n.dtype != Dtype::FP32 &&
230
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 n.op != TraceOpKind::Copy && n.op != TraceOpKind::Cast) {
231 ✗ TraceNode fp32_node = n;
232 ✗ fp32_node.dtype = Dtype::FP32;
233 ✗ Tensor tmp = Tensor::empty_on(n.device, n.rows, n.cols, Dtype::FP32);
234 ✗ evaluate(fp32_node, vals, tmp);
235 ✗ brotensor::cast(tmp, dst, n.dtype);
236 ✗ } else {
237
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 evaluate(n, vals, dst);
238 }
239 3 vals[id] = std::move(dst);
240
241
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 3 times.
 7 for (int in : n.inputs) {
242
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 4 times.
 4 if (in >= 0 && last_use_[static_cast<std::size_t>(in)] == n.id) {
243 4 vals[static_cast<std::size_t>(in)] = Tensor();
244 4 }
245 }
246 3 }
247 2 }
248
249 private:
250 6 const Tensor& arg(const TraceNode& n, std::size_t k, const std::vector<Tensor>& vals) const {
251 6 return vals[static_cast<std::size_t>(n.inputs[k])];
252 }
253
254 // Operand k as a full-shape tensor of the node's dtype.
255 4 Tensor full_arg(const TraceNode& n, std::size_t k, const std::vector<Tensor>& vals) const {
256 4 const Tensor& a = arg(n, k, vals);
257
2/2
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 3 times.
 4 if (numel(a.rows, a.cols) != numel(n.rows, n.cols)) return expand(a, n.rows, n.cols, n.dtype);
258 3 return as_dtype(a, n.dtype, n.rows, n.cols);
259 4 }
260
261 // A gain / shift vector as (1, cols) in the node's dtype.
262 ✗ Tensor row_arg(const TraceNode& n, std::size_t k, const std::vector<Tensor>& vals) const {
263 ✗ return as_dtype(arg(n, k, vals), n.dtype, 1, n.cols);
264 }
265
266 3 void evaluate(const TraceNode& n, const std::vector<Tensor>& vals, Tensor& dst) const {
267
2/17
✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✗ Branch 7 not taken.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
✗ Branch 12 not taken.
✗ Branch 13 not taken.
✗ Branch 14 not taken.
✗ Branch 15 not taken.
✗ Branch 16 not taken.
 3 switch (n.op) {
268 case TraceOpKind::Copy:
269 case TraceOpKind::Cast:
270
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 assign(dst, full_arg(n, 0, vals));
271 2 return;
272 case TraceOpKind::Clamp:
273 ✗ assign(dst, full_arg(n, 0, vals));
274 ✗ brotensor::clamp(dst, n.scalar, n.scalar2);
275 ✗ return;
276 case TraceOpKind::Where: {
277 ✗ Tensor res = brotensor::jit::where(arg(n, 0, vals),
278 ✗ full_arg(n, 1, vals),
279 ✗ full_arg(n, 2, vals));
280 ✗ assign(dst, res);
281 return;
282 ✗ }
283 case TraceOpKind::AddScalar:
284 case TraceOpKind::SubScalar:
285 ✗ assign(dst, full_arg(n, 0, vals));
286 ✗ brotensor::add_scalar_inplace(dst, n.op == TraceOpKind::AddScalar ? n.scalar : -n.scalar);
287 ✗ return;
288 case TraceOpKind::MulScalar:
289 case TraceOpKind::DivScalar:
290 ✗ assign(dst, full_arg(n, 0, vals));
291 ✗ brotensor::scale_inplace(dst, n.op == TraceOpKind::MulScalar ? n.scalar : 1.0f / n.scalar);
292 ✗ return;
293 case TraceOpKind::ScalarSub:
294 ✗ assign(dst, full_arg(n, 0, vals));
295 ✗ brotensor::scale_inplace(dst, -1.0f);
296 ✗ brotensor::add_scalar_inplace(dst, n.scalar);
297 ✗ return;
298 ✗ case TraceOpKind::SiLU: brotensor::silu_forward(full_arg(n, 0, vals), dst); return;
299 ✗ case TraceOpKind::GELU: brotensor::gelu_forward(full_arg(n, 0, vals), dst); return;
300 ✗ case TraceOpKind::ReLU: brotensor::relu_forward(full_arg(n, 0, vals), dst); return;
301 ✗ case TraceOpKind::Tanh: brotensor::tanh_forward(full_arg(n, 0, vals), dst); return;
302 ✗ case TraceOpKind::Sigmoid: brotensor::sigmoid_forward(full_arg(n, 0, vals), dst); return;
303 case TraceOpKind::Add:
304 case TraceOpKind::Sub:
305 case TraceOpKind::Mul:
306 case TraceOpKind::Div:
307 1 binary(n, vals, dst);
308 1 return;
309 case TraceOpKind::FMA: {
310 ✗ Tensor t = full_arg(n, 0, vals);
311 ✗ Tensor prod = Tensor::empty_on(n.device, n.rows, n.cols, n.dtype);
312 ✗ assign(prod, t);
313 ✗ brotensor::mul_inplace(prod, full_arg(n, 1, vals));
314 ✗ brotensor::add_inplace(prod, full_arg(n, 2, vals));
315 ✗ assign(dst, prod);
316 return;
317 ✗ }
318 case TraceOpKind::RMSNorm: {
319 ✗ Tensor g = n.inputs.size() > 1 ? row_arg(n, 1, vals) : ones(n);
320 ✗ into(dst, full_arg(n, 0, vals), [&](const Tensor& x, Tensor& y) {
321 ✗ brotensor::rms_norm_forward(x, g, n.scalar, y);
322 ✗ });
323 return;
324 ✗ }
325 case TraceOpKind::LayerNorm: {
326 ✗ Tensor g = n.inputs.size() > 1 ? row_arg(n, 1, vals) : ones(n);
327 ✗ Tensor b = n.inputs.size() > 2 ? row_arg(n, 2, vals)
328 ✗ : Tensor::zeros_on(n.device, 1, n.cols, n.dtype);
329 ✗ into(dst, full_arg(n, 0, vals), [&](const Tensor& x, Tensor& y) {
330 ✗ brotensor::layernorm_forward_inference_batched(x, g, b, y, n.scalar);
331 ✗ });
332 return;
333 ✗ }
334 case TraceOpKind::Modulate: {
335 ✗ Tensor sc = row_arg(n, 1, vals);
336 ✗ Tensor sh = row_arg(n, 2, vals);
337 ✗ into(dst, full_arg(n, 0, vals), [&](const Tensor& x, Tensor& y) {
338 ✗ brotensor::modulate(x, sc, sh, y);
339 ✗ });
340 return;
341 ✗ }
342 default:
343 ✗ reject(n, "no eager form");
344 }
345 3 }
346
347 ✗ Tensor ones(const TraceNode& n) const {
348 ✗ Tensor t = Tensor::zeros_on(n.device, 1, n.cols, n.dtype);
349 ✗ brotensor::add_scalar_inplace(t, 1.0f);
350 ✗ return t;
351 ✗ }
352
353 // Runs a row op whose output may not alias its input: through a scratch
354 // tensor when dst is the input's storage (an in-place target).
355 template <class F>
356 ✗ void into(Tensor& dst, const Tensor& x, F&& op) const {
357 ✗ if (dst.data != x.data) {
358 ✗ op(x, dst);
359 ✗ return;
360 }
361 ✗ Tensor tmp = Tensor::empty_on(dst.device, dst.rows, dst.cols, dst.dtype);
362 ✗ op(x, tmp);
363 ✗ assign(dst, tmp);
364 ✗ }
365
366 1 void binary(const TraceNode& n, const std::vector<Tensor>& vals, Tensor& dst) const {
367 1 const Tensor& ra = arg(n, 0, vals);
368 1 const Tensor& rb = arg(n, 1, vals);
369 // Multiply by a row / scalar has a kernel of its own; no expansion.
370
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 if (n.op == TraceOpKind::Mul) {
371 ✗ const bool a_small = numel(ra.rows, ra.cols) != numel(n.rows, n.cols);
372 ✗ const bool b_small = numel(rb.rows, rb.cols) != numel(n.rows, n.cols);
373 ✗ if (a_small != b_small) {
374 ✗ const Tensor& small = a_small ? ra : rb;
375 ✗ if (small.size() == 1 || (small.rows == 1 && small.cols == n.cols)) {
376 ✗ const Tensor big = full_arg(n, a_small ? 1 : 0, vals);
377 ✗ const Tensor v = as_dtype(small, n.dtype);
378 ✗ Tensor tmp = dst.data == big.data
379 ✗ ? Tensor::empty_on(n.device, n.rows, n.cols, n.dtype)
380 ✗ : reshaped(dst, n.rows, n.cols);
381 ✗ if (small.size() == 1) {
382 ✗ Tensor y = reshaped(tmp, n.rows * n.cols, 1);
383 ✗ brotensor::broadcast_mul(reshaped(big, n.rows * n.cols, 1), reshaped(v, 1, 1), y);
384 ✗ } else {
385 ✗ brotensor::broadcast_mul(big, reshaped(v, 1, n.cols), tmp);
386 }
387 ✗ assign(dst, tmp);
388 return;
389 ✗ }
390 ✗ }
391 ✗ }
392 1 Tensor a = full_arg(n, 0, vals);
393
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 Tensor b = full_arg(n, 1, vals);
394 // dst = a op b, with dst possibly the storage of a (an in-place
395 // target) or of b. Work in dst when it is a's storage or neither.
396
1/4
✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 2 Tensor acc = dst.data == b.data && dst.data != a.data
397 ✗ ? Tensor::empty_on(n.device, n.rows, n.cols, n.dtype)
398
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 : reshaped(dst, n.rows, n.cols);
399
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 assign(acc, a);
400
1/5
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
 1 switch (n.op) {
401
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 case TraceOpKind::Add: brotensor::add_inplace(acc, b); break;
402 ✗ case TraceOpKind::Sub: brotensor::axpby_inplace(acc, b, 1.0f, -1.0f); break;
403 ✗ case TraceOpKind::Mul: brotensor::mul_inplace(acc, b); break;
404 ✗ case TraceOpKind::Div: brotensor::div_inplace(acc, b); break;
405 ✗ default: reject(n, "not a binary op");
406 }
407
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 assign(dst, acc);
408 1 }
409
410 std::vector<TraceNode> nodes_;
411 std::vector<int> last_use_;
412 2 std::size_t launches_ = 0;
413 };
414
415 2 std::shared_ptr<TraceHandleImpl> make_handle(std::shared_ptr<const Program> prog) {
416 2 auto h = std::make_shared<TraceHandleImpl>();
417
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 4 h->execute_fn = [prog] { prog->run(); };
418
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 h->rebind_fn = [prog](const std::vector<void*>& in, const std::vector<void*>& out) {
419 ✗ return make_handle(prog->rebound(in, out));
420 ✗ };
421
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 h->launch_count = prog->launches();
422 2 h->fusion_name = "eager-unfused";
423 2 return h;
424 2 }
425
426 } // namespace
427
428 2 std::shared_ptr<TraceHandleImpl> compile_eager(const TraceDAG& dag) {
429 2 const auto t0 = std::chrono::steady_clock::now();
430 2 validate(dag);
431
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 auto h = make_handle(std::make_shared<const Program>(dag));
432
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 4 h->compile_us = std::chrono::duration<double, std::micro>(
433
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 std::chrono::steady_clock::now() - t0).count();
434 2 return h;
435 2 }
436
437 } // namespace brotensor::jit::eager
438