GCC Code Coverage Report

Directory: brotensor/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 100.0% 21 / 0 / 21
Functions: 100.0% 14 / 0 / 14
Branches: -% 0 / 0 / 0
src/jit/trace_dag.h
Line Branch Exec Source
1 #pragma once
2
3 #include <brotensor/tensor.h>
4 #include <cstdint>
5 #include <cstddef>
6 #include <vector>
7 #include <unordered_map>
8 #include <string>
9
10 namespace brotensor::jit {
11
12 enum class TraceOpKind : uint8_t {
13 Input = 0,
14 Add,
15 Sub,
16 Mul,
17 Div,
18 AddScalar,
19 SubScalar,
20 MulScalar,
21 DivScalar,
22 ScalarSub,
23 ScalarDiv,
24 FMA,
25 SiLU,
26 GELU,
27 ReLU,
28 RMSNorm,
29 LayerNorm,
30 Modulate,
31 Tanh,
32 Sigmoid,
33 Copy,
34 Where,
35 Clamp,
36 Cast
37 };
38
39 const char* op_kind_name(TraceOpKind op);
40
41 // How a node's buffer maps onto the trace's dominant (rows, cols) shape.
42 //
43 // Row, Col and Scalar are the broadcast forms: a (1, D) row, a (R, 1) column,
44 // or a (1, 1) scalar combined against an (R, D) activation. The emitters turn
45 // these into a different address computation for that operand, not into a
46 // materialised expansion — a broadcast operand contributes D, R, or 1 elements
47 // of traffic, not R*D.
48 enum class BroadcastKind : uint8_t { Full = 0, Row = 1, Scalar = 2, Col = 3 };
49
50 BroadcastKind broadcast_of(int node_rows, int node_cols, int rows, int cols);
51
52 230 struct TraceNode {
53 112 int id = -1;
54 112 TraceOpKind op = TraceOpKind::Input;
55 std::vector<int> inputs;
56 std::vector<int> outputs; // consumer node indices
57 112 float scalar = 0.0f;
58 112 float scalar2 = 0.0f;
59 // Destination or input data pointer — and null for every node the fused
60 // kernel evaluates in registers and never spills.
61 //
62 // A buffer is present on exactly three kinds of node: an Input and an
63 // in-place target, whose storage the caller owns and whose address is
64 // therefore stable for the whole trace; a store() destination, likewise
65 // caller-owned; and a live-at-end intermediate, which end_trace()
66 // allocates and installs into the Tensor the caller is still holding.
67 // Every other node — every dead intermediate — carries a null buffer and
68 // is addressed by `id` alone, which is the reason the tracer no longer
69 // allocates a full-size tensor per traced op. Nothing hashes the buffer,
70 // so a null one cannot collide two different traces.
71 112 void* buffer = nullptr;
72 112 Device device = Device::CPU;
73 112 Dtype dtype = Dtype::FP32;
74 112 int rows = 0;
75 112 int cols = 0;
76 112 bool is_inplace = false;
77 112 int inplace_target_slot = -1;
78 112 bool is_live_at_end = false;
79 };
80
81 class TraceDAG {
82 public:
83 18 TraceDAG() = default;
84
85 int register_input(const Tensor& t);
86 int add_node(TraceOpKind op, const std::vector<int>& inputs, float scalar = 0.0f,
87 void* buffer = nullptr, const Device& dev = Device::CPU,
88 Dtype dt = Dtype::FP32, int rows = 0, int cols = 0,
89 float scalar2 = 0.0f);
90 int add_inplace_node(TraceOpKind op, int target_slot, int arg_slot,
91 void* buffer, float scalar = 0.0f, float scalar2 = 0.0f);
92
93 void analyze_liveness();
94 uint64_t compute_hash() const;
95
96 // The trace's dominant shape: the largest (rows * cols) any node carries.
97 // Every non-broadcast node must match it; (1, cols), (rows, 1), and (1, 1)
98 // nodes are broadcast operands.
99 void dominant_shape(int& rows, int& cols) const;
100
101 73 const std::vector<TraceNode>& nodes() const noexcept { return nodes_; }
102 32 std::vector<TraceNode>& nodes() noexcept { return nodes_; }
103 16 const TraceNode& node(int id) const { return nodes_.at(id); }
104 11 size_t node_count() const noexcept { return nodes_.size(); }
105
106 // True when `slot` still describes `t` (same buffer, shape, dtype, device).
107 bool slot_matches(int slot, const Tensor& t) const;
108 void reset();
109
110 private:
111 std::vector<TraceNode> nodes_;
112 std::unordered_map<const void*, int> input_ptr_to_slot_;
113 };
114
115 // ─── TraceContext ───────────────────────────────────────────────────────────
116 //
117 // The thread's trace in progress, and the symbolic tensors it has handed out.
118 //
119 // A traced op does not produce a tensor; it produces a *name* for a node of
120 // the DAG. record_op() returns that name as a Tensor with a null `data`, a
121 // real shape/dtype/device and `jit_slot` set — so the expression the caller
122 // wrote goes on type-checking and composing exactly as before while allocating
123 // nothing. Node identity follows suit: a symbolic operand is identified by its
124 // slot, and only a caller-owned buffer (an input, an in-place target, a store()
125 // destination) is still identified by its pointer, which is safe precisely
126 // because the caller owns it and it cannot be recycled mid-trace.
127 //
128 // The catch a tracked-pointer table solves: a value the caller still holds
129 // when the trace closes has to end up owning a real buffer it can read from.
130 // Every symbolic Tensor handed out is recorded here as slot -> Tensor*, and
131 // Tensor's move ctor / move assignment / destructor keep that table pointing
132 // at wherever the caller moved the value (see detail::jit_slot_retarget in
133 // tensor.h). end() then allocates for exactly those slots.
134 class TraceContext {
135 public:
136 static TraceContext& current();
137
138 void begin();
139 // Runs liveness, gives every still-held live-at-end intermediate a real
140 // buffer, and returns the DAG. Any symbolic Tensor the caller still holds
141 // whose value the fused kernel only consumes internally is neutralised to
142 // an empty tensor — its contents were never computed to memory.
143 TraceDAG end();
144 // Throws away a trace in progress and neutralises every symbolic Tensor
145 // the caller still holds to an empty tensor. Leaves the context idle.
146 void discard();
147 39 bool is_active() const noexcept { return active_; }
148
149 int get_or_register_slot(const Tensor& t);
150
151 // Records `op` and returns the symbolic Tensor standing for its result.
152 Tensor record_op(TraceOpKind op, const std::vector<int>& in_slots, float scalar,
153 const Device& dev, Dtype dt, int rows, int cols,
154 float scalar2 = 0.0f);
155
156 // Records a Copy into `dst`, a buffer the caller already owns.
157 void record_store(int src_slot, const Tensor& dst);
158
159 // Records a 3-arg cast(src, dst, out_dtype) into `dst`.
160 void record_cast_out(int src_slot, const Tensor& src, Tensor& dst, Dtype out_dtype);
161
162 int record_inplace_op(TraceOpKind op, int target_slot, int arg_slot, float scalar,
163 Tensor& target_tensor, float scalar2 = 0.0f);
164
165 // Tensor lifetime callback: the symbolic Tensor for `slot` moved from
166 // `from` to `to`, or was destroyed (`to` == nullptr). Never throws and
167 // never allocates — Tensor's move and destructor are noexcept.
168 void retarget(int slot, const Tensor* from, Tensor* to) noexcept;
169
170 private:
171 3 TraceContext() = default;
172
173 // Allocates and installs the buffers for live-at-end nodes that do not
174 // already have one; demotes a live-at-end node nobody can observe.
175 void materialize_();
176 // Turns every still-tracked symbolic Tensor into an ordinary empty one.
177 void release_tracked_() noexcept;
178
179 1 bool active_ = false;
180 TraceDAG dag_;
181 std::unordered_map<const void*, int> active_ptr_to_slot_;
182 std::unordered_map<int, Tensor*> slot_to_tensor_;
183 };
184
185 } // namespace brotensor::jit
186