src/jit/ptx_emit_rownorm.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #include "ptx_emit.h" | ||
| 2 | #include "ptx_emit_detail.h" | ||
| 3 | |||
| 4 | #include <sstream> | ||
| 5 | #include <string> | ||
| 6 | #include <unordered_map> | ||
| 7 | #include <vector> | ||
| 8 | |||
| 9 | namespace brotensor::jit::ptx { | ||
| 10 | |||
| 11 | using namespace brotensor::jit::ptx::detail; | ||
| 12 | |||
| 13 | namespace { | ||
| 14 | |||
| 15 | // ─── row-norm kernel ──────────────────────────────────────────────────────── | ||
| 16 | // | ||
| 17 | // A block holds `rpb` rows of `tpr` threads; each group of tpr threads walks | ||
| 18 | // its row in strides of tpr * lanes columns, so every access moves 16 bytes. | ||
| 19 | // row_norm_geometry picks the pair — wide rows get all 256 threads and one | ||
| 20 | // row per block, narrow ones pack several rows into a block instead of idling | ||
| 21 | // most of it. | ||
| 22 | // | ||
| 23 | // RMS pass 1 pre chain (the nodes the reduction depends on), sum of | ||
| 24 | // squares, pass-one outputs stored | ||
| 25 | // pass 2 x (read back, or the pre chain again — holding a whole | ||
| 26 | // row in registers would cap D), * rstd, gamma, post | ||
| 27 | // chain, stores | ||
| 28 | // Mean pass 1 pre chain, sum, pass-one outputs stored | ||
| 29 | // pass 1b x again, sum of squared deviations from the mean | ||
| 30 | // pass 2 as above with (x - mean) * rstd, gamma, beta | ||
| 31 | // | ||
| 32 | // Each reduction is a warp butterfly plus, when a row spans several warps, | ||
| 33 | // one shared-memory round across that row's warps. | ||
| 34 | // | ||
| 35 | // LayerNorm's variance is the two-pass sum of squared deviations, never | ||
| 36 | // E[x^2] - E[x]^2 (which cancels catastrophically once a row's mean dwarfs | ||
| 37 | // its spread — the rule in CLAUDE.md), with the same structure as the Vulkan | ||
| 38 | // emitter (spirv_emit_rownorm.cpp). The extra pass re-reads a row the block | ||
| 39 | // just touched, out of L2. | ||
| 40 | // | ||
| 41 | // The pre chain is what lets `x += p; rms_norm(x)` fuse: the reduction needs | ||
| 42 | // the updated x that this same kernel writes. It also means the row is read | ||
| 43 | // more than once (twice for RMSNorm, three times for LayerNorm), but every | ||
| 44 | // read after the first is the same 8-64 KB the block just touched, so it | ||
| 45 | // comes out of L2 rather than HBM. | ||
| 46 | // | ||
| 47 | // D, the row count and eps are baked in as immediates. The trace signature | ||
| 48 | // already carries the shape, so a different shape is a different trace and a | ||
| 49 | // different compile — and the kernel loses every parameter decode and every | ||
| 50 | // runtime bound it would otherwise carry. | ||
| 51 | |||
| 52 | 16 | std::string emit_rn_entry(const TraceDAG& dag, const RowNormPlan& plan, int lanes, | |
| 53 | const std::string& name) { | ||
| 54 | 16 | const ElementwisePlan& ew = plan.ew; | |
| 55 | 16 | const int D = ew.cols; | |
| 56 | 16 | const int R = ew.rows; | |
| 57 | 16 | const bool mean = (plan.reduce == RowReduce::Mean); | |
| 58 | 16 | int tpr = kRowThreads, rpb = 1; | |
| 59 | 16 | row_norm_geometry(D, R, lanes, tpr, rpb); | |
| 60 | // tpr is never below 32, so a row's threads always fill whole warps and | ||
| 61 | // the butterfly below is always a full-warp one. | ||
| 62 | 16 | const int warps = tpr / 32; // warps cooperating on one row | |
| 63 | 16 | const int step = tpr * lanes; | |
| 64 | |||
| 65 | 16 | Alloc a; | |
| 66 | 16 | std::stringstream b; | |
| 67 | |||
| 68 | 16 | const std::size_t n_out = ew.outputs.size(); | |
| 69 | 16 | const std::size_t n_in = ew.inputs.size(); | |
| 70 | |||
| 71 |
2/4✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
|
16 | std::vector<std::string> out_ptr(n_out), in_ptr(n_in); |
| 72 |
2/2✓ Branch 0 taken 16 times.
✓ Branch 1 taken 20 times.
|
36 | for (std::size_t m = 0; m < n_out; ++m) { |
| 73 |
1/2✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
|
20 | out_ptr[m] = a.RD(); |
| 74 |
7/14✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 20 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 20 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 20 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 20 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 20 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 20 times.
✗ Branch 13 not taken.
|
20 | b << " ld.param.u64 " << out_ptr[m] << ", [" << name << "_out_" << m << "];\n"; |
| 75 | 20 | } | |
| 76 |
2/2✓ Branch 0 taken 16 times.
✓ Branch 1 taken 52 times.
|
68 | for (std::size_t k = 0; k < n_in; ++k) { |
| 77 |
1/2✓ Branch 0 taken 52 times.
✗ Branch 1 not taken.
|
52 | in_ptr[k] = a.RD(); |
| 78 |
7/14✓ Branch 0 taken 52 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 52 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 52 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 52 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 52 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 52 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 52 times.
✗ Branch 13 not taken.
|
52 | b << " ld.param.u64 " << in_ptr[k] << ", [" << name << "_in_" << k << "];\n"; |
| 79 | 52 | } | |
| 80 | |||
| 81 | // tid splits into (which row of this block, which thread within the row). | ||
| 82 |
3/6✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 16 times.
✗ Branch 5 not taken.
|
16 | const std::string r_tid = a.R(), r_row = a.R(), r_c0 = a.R(); |
| 83 |
2/4✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
|
16 | const std::string r_sub = a.R(), r_local = a.R(); |
| 84 |
3/6✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 16 times.
✗ Branch 5 not taken.
|
16 | b << " mov.u32 " << r_tid << ", %tid.x;\n"; |
| 85 |
3/6✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 16 times.
✗ Branch 5 not taken.
|
16 | b << " mov.u32 " << r_row << ", %ctaid.x;\n"; |
| 86 |
2/2✓ Branch 0 taken 8 times.
✓ Branch 1 taken 8 times.
|
16 | if (rpb == 1) { |
| 87 |
3/6✓ 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.
|
8 | b << " mov.u32 " << r_local << ", 0;\n"; |
| 88 |
5/10✓ 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.
|
8 | b << " mov.u32 " << r_sub << ", " << r_tid << ";\n"; |
| 89 | 8 | } else { | |
| 90 | // tpr is a power of two, so both halves are a shift and a mask. | ||
| 91 | 8 | int log_tpr = 0; | |
| 92 |
2/2✓ Branch 0 taken 44 times.
✓ Branch 1 taken 8 times.
|
52 | while ((1 << log_tpr) < tpr) ++log_tpr; |
| 93 |
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 | b << " shr.u32 " << r_local << ", " << r_tid << ", " << log_tpr << ";\n"; |
| 94 |
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 | b << " and.b32 " << r_sub << ", " << r_tid << ", " << (tpr - 1) << ";\n"; |
| 95 |
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 | b << " mad.lo.u32 " << r_row << ", " << r_row << ", " << rpb << ", " |
| 96 |
2/4✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
|
8 | << r_local << ";\n"; |
| 97 | } | ||
| 98 | { | ||
| 99 | // rpb divides R (row_norm_geometry guarantees it), so this exit is | ||
| 100 | // uniform across the block and the barrier below stays safe. | ||
| 101 |
1/2✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
|
16 | const std::string p = a.P(); |
| 102 |
7/14✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 16 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 16 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 16 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 16 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 16 times.
✗ Branch 13 not taken.
|
16 | b << " setp.ge.u32 " << p << ", " << r_row << ", " << R << ";\n"; |
| 103 |
5/10✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 16 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 16 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 16 times.
✗ Branch 9 not taken.
|
16 | b << " @" << p << " bra $L_" << name << "_done;\n"; |
| 104 | 16 | } | |
| 105 | // This thread's first column. | ||
| 106 |
2/2✓ Branch 0 taken 8 times.
✓ Branch 1 taken 8 times.
|
16 | if (lanes == 1) { |
| 107 |
5/10✓ 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.
|
8 | b << " mov.u32 " << r_c0 << ", " << r_sub << ";\n"; |
| 108 | 8 | } else { | |
| 109 |
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 | b << " mul.lo.u32 " << r_c0 << ", " << r_sub << ", " << lanes << ";\n"; |
| 110 | } | ||
| 111 | |||
| 112 | // Element index of column 0 of this row, and its byte offset per dtype. | ||
| 113 |
1/2✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
|
16 | const std::string r_rowbase = a.R(); |
| 114 |
7/14✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 16 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 16 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 16 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 16 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 16 times.
✗ Branch 13 not taken.
|
16 | b << " mul.lo.u32 " << r_rowbase << ", " << r_row << ", " << D << ";\n"; |
| 115 | |||
| 116 | 16 | std::unordered_map<int, std::string> rowbyte, colbcastbyte; | |
| 117 | 132 | auto row_bytes = [&](Dtype dt) -> std::string { | |
| 118 | 116 | const int e = elem_bytes(dt); | |
| 119 | 116 | auto it = rowbyte.find(e); | |
| 120 |
2/2✓ Branch 0 taken 100 times.
✓ Branch 1 taken 16 times.
|
116 | if (it != rowbyte.end()) return it->second; |
| 121 | 16 | const std::string rd = a.RD(); | |
| 122 |
7/14✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 16 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 16 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 16 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 16 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 16 times.
✗ Branch 13 not taken.
|
16 | b << " mul.wide.u32 " << rd << ", " << r_rowbase << ", " << e << ";\n"; |
| 123 |
2/4✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
|
16 | rowbyte[e] = rd; |
| 124 | 16 | return rd; | |
| 125 |
1/2✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
|
132 | }; |
| 126 | 16 | auto colbcast_bytes = [&](Dtype dt) -> std::string { | |
| 127 | ✗ | const int e = elem_bytes(dt); | |
| 128 | ✗ | auto it = colbcastbyte.find(e); | |
| 129 | ✗ | if (it != colbcastbyte.end()) return it->second; | |
| 130 | ✗ | const std::string rd = a.RD(); | |
| 131 | ✗ | b << " mul.wide.u32 " << rd << ", " << r_row << ", " << e << ";\n"; | |
| 132 | ✗ | colbcastbyte[e] = rd; | |
| 133 | ✗ | return rd; | |
| 134 | ✗ | }; | |
| 135 | // Every element size the trace touches, resolved before either loop so the | ||
| 136 | // stride loops do not recompute a loop-invariant address. | ||
| 137 |
2/2✓ Branch 0 taken 52 times.
✓ Branch 1 taken 16 times.
|
68 | for (const BufferSpec& s : ew.inputs) { |
| 138 |
1/2✓ Branch 0 taken 52 times.
✗ Branch 1 not taken.
|
52 | row_bytes(s.dtype); |
| 139 |
1/4✗ Branch 0 not taken.
✓ Branch 1 taken 52 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
52 | if (s.bcast == BroadcastKind::Col) colbcast_bytes(s.dtype); |
| 140 | } | ||
| 141 |
3/4✓ Branch 0 taken 16 times.
✓ Branch 1 taken 20 times.
✓ Branch 2 taken 20 times.
✗ Branch 3 not taken.
|
36 | for (const BufferSpec& s : ew.outputs) row_bytes(s.dtype); |
| 142 | |||
| 143 | // Per-pass cache of the byte offsets materialised for the column the | ||
| 144 | // stride loop is on. The two passes are separate PTX loops, so each gets | ||
| 145 | // its own. | ||
| 146 | struct Offsets { | ||
| 147 | std::unordered_map<int, std::string> col, full; | ||
| 148 | }; | ||
| 149 | |||
| 150 | 88 | auto col_off = [&](Offsets& o, const std::string& r_c, Dtype dt) -> std::string { | |
| 151 | 72 | const int e = elem_bytes(dt); | |
| 152 | 72 | auto it = o.col.find(e); | |
| 153 |
2/2✓ Branch 0 taken 28 times.
✓ Branch 1 taken 44 times.
|
72 | if (it != o.col.end()) return it->second; |
| 154 | 44 | const std::string rd = a.RD(); | |
| 155 |
7/14✓ Branch 0 taken 44 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 44 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 44 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 44 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 44 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 44 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 44 times.
✗ Branch 13 not taken.
|
44 | b << " mul.wide.u32 " << rd << ", " << r_c << ", " << e << ";\n"; |
| 156 |
2/4✓ Branch 0 taken 44 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 44 times.
✗ Branch 3 not taken.
|
44 | o.col[e] = rd; |
| 157 | 44 | return rd; | |
| 158 |
1/2✓ Branch 0 taken 44 times.
✗ Branch 1 not taken.
|
116 | }; |
| 159 | 96 | auto full_off = [&](Offsets& o, const std::string& r_c, Dtype dt) -> std::string { | |
| 160 | 80 | const int e = elem_bytes(dt); | |
| 161 | 80 | auto it = o.full.find(e); | |
| 162 |
2/2✓ Branch 0 taken 36 times.
✓ Branch 1 taken 44 times.
|
80 | if (it != o.full.end()) return it->second; |
| 163 | 44 | const std::string rb = row_bytes(dt); | |
| 164 |
1/2✓ Branch 0 taken 44 times.
✗ Branch 1 not taken.
|
44 | const std::string co = col_off(o, r_c, dt); |
| 165 |
1/2✓ Branch 0 taken 44 times.
✗ Branch 1 not taken.
|
44 | const std::string rd = a.RD(); |
| 166 |
7/14✓ Branch 0 taken 44 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 44 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 44 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 44 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 44 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 44 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 44 times.
✗ Branch 13 not taken.
|
44 | b << " add.u64 " << rd << ", " << rb << ", " << co << ";\n"; |
| 167 |
2/4✓ Branch 0 taken 44 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 44 times.
✗ Branch 3 not taken.
|
44 | o.full[e] = rd; |
| 168 | 44 | return rd; | |
| 169 |
1/2✓ Branch 0 taken 44 times.
✗ Branch 1 not taken.
|
80 | }; |
| 170 | |||
| 171 | // Loads the inputs `want` selects (empty = all) at the column in `r_c`. | ||
| 172 | 56 | auto emit_inputs = [&](Offsets& o, const std::string& r_c, | |
| 173 | const std::vector<char>& want) -> ValueMap { | ||
| 174 | 40 | ValueMap vals; | |
| 175 |
2/2✓ Branch 0 taken 132 times.
✓ Branch 1 taken 40 times.
|
172 | for (std::size_t k = 0; k < n_in; ++k) { |
| 176 |
3/4✓ Branch 0 taken 132 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 80 times.
✓ Branch 3 taken 52 times.
|
132 | if (!want.empty() && !want[k]) continue; |
| 177 | 80 | const BufferSpec& s = ew.inputs[k]; | |
| 178 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 80 times.
|
80 | const bool single_lane = (s.bcast == BroadcastKind::Scalar || s.bcast == BroadcastKind::Col); |
| 179 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 80 times.
|
80 | const int ln = single_lane ? 1 : lanes; |
| 180 | 80 | std::string addr; | |
| 181 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 80 times.
|
80 | if (s.bcast == BroadcastKind::Scalar) { |
| 182 | ✗ | addr = in_ptr[k]; | |
| 183 | ✗ | } else { | |
| 184 |
2/2✓ Branch 0 taken 28 times.
✓ Branch 1 taken 52 times.
|
160 | const std::string off = (s.bcast == BroadcastKind::Row) |
| 185 |
1/2✓ Branch 0 taken 28 times.
✗ Branch 1 not taken.
|
28 | ? col_off(o, r_c, s.dtype) |
| 186 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 52 times.
|
104 | : (s.bcast == BroadcastKind::Col) |
| 187 | ✗ | ? colbcast_bytes(s.dtype) | |
| 188 |
1/2✓ Branch 0 taken 52 times.
✗ Branch 1 not taken.
|
52 | : full_off(o, r_c, s.dtype); |
| 189 |
1/2✓ Branch 0 taken 80 times.
✗ Branch 1 not taken.
|
80 | addr = a.RD(); |
| 190 |
7/14✓ Branch 0 taken 80 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 80 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 80 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 80 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 80 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 80 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 80 times.
✗ Branch 13 not taken.
|
80 | b << " add.u64 " << addr << ", " << in_ptr[k] << ", " << off << ";\n"; |
| 191 | 80 | } | |
| 192 |
1/2✓ Branch 0 taken 80 times.
✗ Branch 1 not taken.
|
80 | std::vector<std::string> v = emit_load(b, a, s.dtype, ln, addr); |
| 193 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 80 times.
|
80 | if (single_lane) { |
| 194 | ✗ | const std::string only = v[0]; | |
| 195 | ✗ | v.assign(static_cast<std::size_t>(lanes), only); | |
| 196 | ✗ | } | |
| 197 |
1/2✓ Branch 0 taken 80 times.
✗ Branch 1 not taken.
|
80 | vals[s.node_id] = std::move(v); |
| 198 | 80 | } | |
| 199 | 40 | return vals; | |
| 200 |
1/2✓ Branch 0 taken 40 times.
✗ Branch 1 not taken.
|
40 | }; |
| 201 | |||
| 202 | // x at the column in `r_c` after pass one: read back from the output | ||
| 203 | // pass one wrote, or (no pre output) the pre chain replayed from the | ||
| 204 | // inputs. With `vals`, also loads what pass two needs and leaves it there. | ||
| 205 | 44 | auto reload_x = [&](Offsets& o, const std::string& r_c, | |
| 206 | ValueMap* vals) -> std::vector<std::string> { | ||
| 207 | 28 | ValueMap local; | |
| 208 |
2/2✓ Branch 0 taken 16 times.
✓ Branch 1 taken 12 times.
|
28 | ValueMap& v = vals ? *vals : local; |
| 209 |
2/2✓ Branch 0 taken 8 times.
✓ Branch 1 taken 20 times.
|
28 | if (plan.x_output >= 0) { |
| 210 |
3/4✓ Branch 0 taken 4 times.
✓ Branch 1 taken 4 times.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
|
8 | if (vals) v = emit_inputs(o, r_c, plan.post_input); |
| 211 | 8 | const BufferSpec& xs = ew.outputs[static_cast<std::size_t>(plan.x_output)]; | |
| 212 |
1/2✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
|
8 | const std::string off = full_off(o, r_c, xs.dtype); |
| 213 |
1/2✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
|
8 | const std::string addr = a.RD(); |
| 214 |
4/8✓ 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.
|
8 | b << " add.u64 " << addr << ", " << out_ptr[static_cast<std::size_t>(plan.x_output)] |
| 215 |
3/6✓ 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.
|
8 | << ", " << off << ";\n"; |
| 216 |
1/2✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
|
8 | std::vector<std::string> x = emit_load(b, a, xs.dtype, lanes, addr); |
| 217 |
2/4✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
|
8 | v[plan.x_node] = x; |
| 218 | 8 | return x; | |
| 219 | 8 | } | |
| 220 | // Without the reload the plan guarantees pass one stored nothing, so | ||
| 221 | // the inputs are still what pass one read. | ||
| 222 |
3/4✓ Branch 0 taken 12 times.
✓ Branch 1 taken 8 times.
✓ Branch 2 taken 20 times.
✗ Branch 3 not taken.
|
20 | v = emit_inputs(o, r_c, vals ? plan.post_input : plan.pre_input); |
| 223 |
1/2✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
|
20 | v = emit_chain(b, a, dag, lanes, std::move(v), &plan.pre_ids); |
| 224 |
2/4✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 20 times.
✗ Branch 3 not taken.
|
20 | return v.at(plan.x_node); |
| 225 | 28 | }; | |
| 226 | |||
| 227 | // The row total of `acc` in every thread of the row: a warp butterfly, | ||
| 228 | // then, when the row spans several warps, one shared-memory round across | ||
| 229 | // ITS warps. `region` picks the slot set (one entry per warp of the | ||
| 230 | // block, at most eight): squares at 0, sums after them, so the two | ||
| 231 | // LayerNorm rounds touch different words and need no barrier between | ||
| 232 | // them. | ||
| 233 | 44 | auto block_reduce = [&](const std::string& acc, int region, const char* tag) { | |
| 234 | 28 | const std::string rb = a.R(); | |
| 235 |
2/2✓ Branch 0 taken 140 times.
✓ Branch 1 taken 28 times.
|
168 | for (int m = 16; m >= 1; m >>= 1) { |
| 236 |
1/2✓ Branch 0 taken 140 times.
✗ Branch 1 not taken.
|
140 | const std::string t = a.F(); |
| 237 |
5/10✓ Branch 0 taken 140 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 140 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 140 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 140 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 140 times.
✗ Branch 9 not taken.
|
140 | b << " mov.b32 " << rb << ", " << acc << ";\n"; |
| 238 |
7/14✓ Branch 0 taken 140 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 140 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 140 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 140 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 140 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 140 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 140 times.
✗ Branch 13 not taken.
|
140 | b << " shfl.sync.bfly.b32 " << rb << ", " << rb << ", " << m << ", 31, -1;\n"; |
| 239 |
5/10✓ Branch 0 taken 140 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 140 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 140 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 140 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 140 times.
✗ Branch 9 not taken.
|
140 | b << " mov.b32 " << t << ", " << rb << ";\n"; |
| 240 |
7/14✓ Branch 0 taken 140 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 140 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 140 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 140 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 140 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 140 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 140 times.
✗ Branch 13 not taken.
|
140 | b << " add.f32 " << acc << ", " << acc << ", " << t << ";\n"; |
| 241 | 140 | } | |
| 242 |
2/2✓ Branch 0 taken 21 times.
✓ Branch 1 taken 7 times.
|
28 | if (warps == 1) return; |
| 243 | 21 | const int off = region * (kRowThreads / 32) * 4; | |
| 244 |
4/8✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 21 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 21 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 21 times.
✗ Branch 7 not taken.
|
21 | const std::string r_warp = a.R(), r_lane = a.R(), r_sb = a.R(), r_sa = a.R(); |
| 245 |
2/4✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 21 times.
✗ Branch 3 not taken.
|
21 | const std::string p_lane0 = a.P(), rt = a.R(); |
| 246 |
5/10✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 21 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 21 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 21 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 21 times.
✗ Branch 9 not taken.
|
21 | b << " shr.u32 " << r_warp << ", " << r_tid << ", 5;\n"; |
| 247 |
5/10✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 21 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 21 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 21 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 21 times.
✗ Branch 9 not taken.
|
21 | b << " and.b32 " << r_lane << ", " << r_tid << ", 31;\n"; |
| 248 |
3/6✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 21 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 21 times.
✗ Branch 5 not taken.
|
21 | b << " mov.u32 " << r_sb << ", rn_red;\n"; |
| 249 |
9/16✓ Branch 0 taken 9 times.
✓ Branch 1 taken 12 times.
✓ Branch 2 taken 9 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 9 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 9 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 9 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 9 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 9 times.
✗ Branch 13 not taken.
✓ Branch 14 taken 9 times.
✗ Branch 15 not taken.
|
21 | if (off) b << " add.u32 " << r_sb << ", " << r_sb << ", " << off << ";\n"; |
| 250 |
5/10✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 21 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 21 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 21 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 21 times.
✗ Branch 9 not taken.
|
21 | b << " setp.eq.u32 " << p_lane0 << ", " << r_lane << ", 0;\n"; |
| 251 |
7/14✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 21 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 21 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 21 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 21 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 21 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 21 times.
✗ Branch 13 not taken.
|
21 | b << " @!" << p_lane0 << " bra $L_" << name << "_" << tag << "_nostore;\n"; |
| 252 |
7/14✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 21 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 21 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 21 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 21 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 21 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 21 times.
✗ Branch 13 not taken.
|
21 | b << " mad.lo.u32 " << r_sa << ", " << r_warp << ", 4, " << r_sb << ";\n"; |
| 253 |
5/10✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 21 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 21 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 21 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 21 times.
✗ Branch 9 not taken.
|
21 | b << " mov.b32 " << rt << ", " << acc << ";\n"; |
| 254 |
5/10✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 21 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 21 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 21 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 21 times.
✗ Branch 9 not taken.
|
21 | b << " st.shared.b32 [" << r_sa << "], " << rt << ";\n"; |
| 255 |
5/10✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 21 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 21 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 21 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 21 times.
✗ Branch 9 not taken.
|
21 | b << "$L_" << name << "_" << tag << "_nostore:\n"; |
| 256 |
1/2✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
|
21 | b << " bar.sync 0;\n"; |
| 257 | |||
| 258 | // This row's warps start at (row within block) * warps. Every thread | ||
| 259 | // folds them itself, so the result needs no second broadcast. | ||
| 260 |
1/2✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
|
21 | const std::string r_base = a.R(); |
| 261 |
2/2✓ Branch 0 taken 14 times.
✓ Branch 1 taken 7 times.
|
21 | if (rpb == 1) { |
| 262 |
5/10✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 14 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 14 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 14 times.
✗ Branch 9 not taken.
|
14 | b << " mov.u32 " << r_base << ", " << r_sb << ";\n"; |
| 263 | 14 | } else { | |
| 264 |
6/12✓ 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 taken 7 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 7 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 7 times.
✗ Branch 11 not taken.
|
7 | b << " mad.lo.u32 " << r_base << ", " << r_local << ", " << (warps * 4) |
| 265 |
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 | << ", " << r_sb << ";\n"; |
| 266 | } | ||
| 267 |
3/6✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 21 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 21 times.
✗ Branch 5 not taken.
|
21 | b << " mov.f32 " << acc << ", 0f00000000;\n"; |
| 268 |
2/2✓ Branch 0 taken 126 times.
✓ Branch 1 taken 21 times.
|
147 | for (int w = 0; w < warps; ++w) { |
| 269 |
3/6✓ Branch 0 taken 126 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 126 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 126 times.
|
126 | const std::string ra = a.R(), rv = a.R(), fv = a.F(); |
| 270 |
7/14✓ Branch 0 taken 126 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 126 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 126 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 126 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 126 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 126 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 126 times.
✗ Branch 13 not taken.
|
126 | b << " add.u32 " << ra << ", " << r_base << ", " << (w * 4) << ";\n"; |
| 271 |
5/10✓ Branch 0 taken 126 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 126 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 126 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 126 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 126 times.
✗ Branch 9 not taken.
|
126 | b << " ld.shared.b32 " << rv << ", [" << ra << "];\n"; |
| 272 |
5/10✓ Branch 0 taken 126 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 126 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 126 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 126 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 126 times.
✗ Branch 9 not taken.
|
126 | b << " mov.b32 " << fv << ", " << rv << ";\n"; |
| 273 |
7/14✓ Branch 0 taken 126 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 126 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 126 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 126 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 126 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 126 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 126 times.
✗ Branch 13 not taken.
|
126 | b << " add.f32 " << acc << ", " << acc << ", " << fv << ";\n"; |
| 274 | 126 | } | |
| 275 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 28 times.
|
28 | }; |
| 276 | |||
| 277 | // ── pass 1: the pre chain, then the sum (Mean) or sum of squares (RMS) ── | ||
| 278 |
1/2✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
|
16 | const std::string f_acc = a.F(); |
| 279 |
3/6✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 16 times.
✗ Branch 5 not taken.
|
16 | b << " mov.f32 " << f_acc << ", 0f00000000;\n"; |
| 280 | { | ||
| 281 |
2/4✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
|
16 | const std::string r_c = a.R(), p = a.P(); |
| 282 |
5/10✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 16 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 16 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 16 times.
✗ Branch 9 not taken.
|
16 | b << " mov.u32 " << r_c << ", " << r_c0 << ";\n"; |
| 283 |
3/6✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 16 times.
✗ Branch 5 not taken.
|
16 | b << "$L_" << name << "_sum:\n"; |
| 284 |
7/14✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 16 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 16 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 16 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 16 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 16 times.
✗ Branch 13 not taken.
|
16 | b << " setp.ge.u32 " << p << ", " << r_c << ", " << D << ";\n"; |
| 285 |
5/10✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 16 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 16 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 16 times.
✗ Branch 9 not taken.
|
16 | b << " @" << p << " bra $L_" << name << "_sum_end;\n"; |
| 286 | |||
| 287 | 16 | Offsets o1; | |
| 288 | ValueMap vals = | ||
| 289 |
2/4✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
|
16 | emit_chain(b, a, dag, lanes, emit_inputs(o1, r_c, plan.pre_input), &plan.pre_ids); |
| 290 |
2/4✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
|
16 | const std::vector<std::string> x = vals.at(plan.x_node); |
| 291 |
2/2✓ Branch 0 taken 16 times.
✓ Branch 1 taken 40 times.
|
56 | for (int j = 0; j < lanes; ++j) { |
| 292 | 40 | const std::string& xj = x[static_cast<std::size_t>(j)]; | |
| 293 |
2/2✓ Branch 0 taken 10 times.
✓ Branch 1 taken 30 times.
|
40 | if (mean) { |
| 294 |
7/14✓ Branch 0 taken 30 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 30 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 30 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 30 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 30 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 30 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 30 times.
✗ Branch 13 not taken.
|
30 | b << " add.f32 " << f_acc << ", " << f_acc << ", " << xj << ";\n"; |
| 295 | 30 | } else { | |
| 296 |
9/18✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 10 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 10 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 10 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 10 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 10 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 10 times.
✗ Branch 13 not taken.
✓ Branch 14 taken 10 times.
✗ Branch 15 not taken.
✓ Branch 16 taken 10 times.
✗ Branch 17 not taken.
|
10 | b << " fma.rn.f32 " << f_acc << ", " << xj << ", " << xj << ", " << f_acc << ";\n"; |
| 297 | } | ||
| 298 | 40 | } | |
| 299 | // Outputs the pre chain produced are final here — writing them now is | ||
| 300 | // what lets pass two read the value back instead of recomputing it. | ||
| 301 |
2/2✓ Branch 0 taken 16 times.
✓ Branch 1 taken 20 times.
|
36 | for (std::size_t m = 0; m < n_out; ++m) { |
| 302 |
3/4✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✓ Branch 3 taken 4 times.
|
20 | if (m >= plan.pre_output.size() || !plan.pre_output[m]) continue; |
| 303 | 4 | const BufferSpec& s = ew.outputs[m]; | |
| 304 |
1/2✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
|
4 | const std::string off = full_off(o1, r_c, s.dtype); |
| 305 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
|
4 | const std::string addr = a.RD(); |
| 306 |
7/14✓ 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.
✓ Branch 6 taken 4 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 4 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 4 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 4 times.
✗ Branch 13 not taken.
|
4 | b << " add.u64 " << addr << ", " << out_ptr[m] << ", " << off << ";\n"; |
| 307 |
2/4✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
|
4 | emit_store(b, a, s.dtype, lanes, addr, vals.at(s.node_id)); |
| 308 | 4 | } | |
| 309 | |||
| 310 |
7/14✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 16 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 16 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 16 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 16 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 16 times.
✗ Branch 13 not taken.
|
16 | b << " add.u32 " << r_c << ", " << r_c << ", " << step << ";\n"; |
| 311 |
3/6✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 16 times.
✗ Branch 5 not taken.
|
16 | b << " bra $L_" << name << "_sum;\n"; |
| 312 |
3/6✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 16 times.
✗ Branch 5 not taken.
|
16 | b << "$L_" << name << "_sum_end:\n"; |
| 313 | 16 | } | |
| 314 | |||
| 315 | // ── reductions; for LayerNorm, pass 1b over the deviations ────────────── | ||
| 316 |
2/4✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
|
16 | const std::string f_mean = a.F(), f_rstd = a.F(); |
| 317 | 16 | const float inv_d = 1.0f / static_cast<float>(D); | |
| 318 |
2/2✓ Branch 0 taken 12 times.
✓ Branch 1 taken 4 times.
|
16 | if (mean) { |
| 319 |
1/2✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
|
12 | block_reduce(f_acc, 1, "sum"); |
| 320 |
8/16✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 12 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 12 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 12 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 12 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 12 times.
✗ Branch 13 not taken.
✓ Branch 14 taken 12 times.
✗ Branch 15 not taken.
|
12 | b << " mul.f32 " << f_mean << ", " << f_acc << ", " << fhex(inv_d) << ";\n"; |
| 321 | |||
| 322 |
1/2✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
|
12 | const std::string f_dev = a.F(); |
| 323 |
3/6✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 12 times.
✗ Branch 5 not taken.
|
12 | b << " mov.f32 " << f_dev << ", 0f00000000;\n"; |
| 324 |
2/4✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
|
12 | const std::string r_c = a.R(), p = a.P(); |
| 325 |
5/10✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 12 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 12 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 12 times.
✗ Branch 9 not taken.
|
12 | b << " mov.u32 " << r_c << ", " << r_c0 << ";\n"; |
| 326 |
3/6✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 12 times.
✗ Branch 5 not taken.
|
12 | b << "$L_" << name << "_dev:\n"; |
| 327 |
7/14✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 12 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 12 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 12 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 12 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 12 times.
✗ Branch 13 not taken.
|
12 | b << " setp.ge.u32 " << p << ", " << r_c << ", " << D << ";\n"; |
| 328 |
5/10✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 12 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 12 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 12 times.
✗ Branch 9 not taken.
|
12 | b << " @" << p << " bra $L_" << name << "_dev_end;\n"; |
| 329 | 12 | Offsets o1b; | |
| 330 |
1/2✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
|
12 | const std::vector<std::string> x = reload_x(o1b, r_c, nullptr); |
| 331 |
2/2✓ Branch 0 taken 12 times.
✓ Branch 1 taken 30 times.
|
42 | for (int j = 0; j < lanes; ++j) { |
| 332 |
1/2✓ Branch 0 taken 30 times.
✗ Branch 1 not taken.
|
30 | const std::string d = a.F(); |
| 333 |
7/14✓ Branch 0 taken 30 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 30 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 30 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 30 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 30 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 30 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 30 times.
✗ Branch 13 not taken.
|
30 | b << " sub.f32 " << d << ", " << x[static_cast<std::size_t>(j)] << ", " << f_mean << ";\n"; |
| 334 |
9/18✓ Branch 0 taken 30 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 30 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 30 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 30 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 30 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 30 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 30 times.
✗ Branch 13 not taken.
✓ Branch 14 taken 30 times.
✗ Branch 15 not taken.
✓ Branch 16 taken 30 times.
✗ Branch 17 not taken.
|
30 | b << " fma.rn.f32 " << f_dev << ", " << d << ", " << d << ", " << f_dev << ";\n"; |
| 335 | 30 | } | |
| 336 |
7/14✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 12 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 12 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 12 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 12 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 12 times.
✗ Branch 13 not taken.
|
12 | b << " add.u32 " << r_c << ", " << r_c << ", " << step << ";\n"; |
| 337 |
3/6✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 12 times.
✗ Branch 5 not taken.
|
12 | b << " bra $L_" << name << "_dev;\n"; |
| 338 |
3/6✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 12 times.
✗ Branch 5 not taken.
|
12 | b << "$L_" << name << "_dev_end:\n"; |
| 339 | |||
| 340 |
1/2✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
|
12 | block_reduce(f_dev, 0, "dev"); |
| 341 |
1/2✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
|
12 | const std::string var = a.F(); |
| 342 |
8/16✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 12 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 12 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 12 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 12 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 12 times.
✗ Branch 13 not taken.
✓ Branch 14 taken 12 times.
✗ Branch 15 not taken.
|
12 | b << " mul.f32 " << var << ", " << f_dev << ", " << fhex(inv_d) << ";\n"; |
| 343 |
8/16✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 12 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 12 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 12 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 12 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 12 times.
✗ Branch 13 not taken.
✓ Branch 14 taken 12 times.
✗ Branch 15 not taken.
|
12 | b << " add.f32 " << var << ", " << var << ", " << fhex(plan.eps) << ";\n"; |
| 344 |
5/10✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 12 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 12 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 12 times.
✗ Branch 9 not taken.
|
12 | b << " rsqrt.approx.f32 " << f_rstd << ", " << var << ";\n"; |
| 345 | 12 | } else { | |
| 346 |
1/2✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
|
4 | block_reduce(f_acc, 0, "sq"); |
| 347 |
1/2✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
|
4 | const std::string ms = a.F(); |
| 348 |
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 | b << " mov.f32 " << f_mean << ", 0f00000000;\n"; |
| 349 |
8/16✓ 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.
✓ Branch 6 taken 4 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 4 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 4 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 4 times.
✗ Branch 13 not taken.
✓ Branch 14 taken 4 times.
✗ Branch 15 not taken.
|
4 | b << " mul.f32 " << ms << ", " << f_acc << ", " << fhex(inv_d) << ";\n"; |
| 350 |
8/16✓ 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.
✓ Branch 6 taken 4 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 4 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 4 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 4 times.
✗ Branch 13 not taken.
✓ Branch 14 taken 4 times.
✗ Branch 15 not taken.
|
4 | b << " add.f32 " << ms << ", " << ms << ", " << fhex(plan.eps) << ";\n"; |
| 351 |
5/10✓ 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.
✓ Branch 6 taken 4 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 4 times.
✗ Branch 9 not taken.
|
4 | b << " rsqrt.approx.f32 " << f_rstd << ", " << ms << ";\n"; |
| 352 | 4 | } | |
| 353 | |||
| 354 | // ── pass 2: pre chain, normalise, post chain, store ───────────────────── | ||
| 355 | { | ||
| 356 |
2/4✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
|
16 | const std::string r_c = a.R(), p = a.P(); |
| 357 |
5/10✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 16 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 16 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 16 times.
✗ Branch 9 not taken.
|
16 | b << " mov.u32 " << r_c << ", " << r_c0 << ";\n"; |
| 358 |
3/6✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 16 times.
✗ Branch 5 not taken.
|
16 | b << "$L_" << name << "_apply:\n"; |
| 359 |
7/14✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 16 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 16 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 16 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 16 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 16 times.
✗ Branch 13 not taken.
|
16 | b << " setp.ge.u32 " << p << ", " << r_c << ", " << D << ";\n"; |
| 360 |
5/10✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 16 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 16 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 16 times.
✗ Branch 9 not taken.
|
16 | b << " @" << p << " bra $L_" << name << "_apply_end;\n"; |
| 361 | |||
| 362 | 16 | Offsets o2; | |
| 363 | 16 | ValueMap vals; | |
| 364 |
1/2✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
|
16 | const std::vector<std::string> x = reload_x(o2, r_c, &vals); |
| 365 | 16 | std::vector<std::string> gam, bet; | |
| 366 |
1/2✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
|
16 | if (plan.gamma_input >= 0) { |
| 367 |
2/4✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
|
16 | gam = vals.at(ew.inputs[static_cast<std::size_t>(plan.gamma_input)].node_id); |
| 368 | 16 | } | |
| 369 |
2/2✓ Branch 0 taken 12 times.
✓ Branch 1 taken 4 times.
|
16 | if (plan.beta_input >= 0) { |
| 370 |
2/4✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
|
12 | bet = vals.at(ew.inputs[static_cast<std::size_t>(plan.beta_input)].node_id); |
| 371 | 12 | } | |
| 372 | |||
| 373 |
1/2✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
|
16 | std::vector<std::string> norm(static_cast<std::size_t>(lanes)); |
| 374 |
2/2✓ Branch 0 taken 16 times.
✓ Branch 1 taken 40 times.
|
56 | for (int j = 0; j < lanes; ++j) { |
| 375 |
1/2✓ Branch 0 taken 40 times.
✗ Branch 1 not taken.
|
40 | const std::string f_n = a.F(); |
| 376 | 40 | const std::string& xj = x[static_cast<std::size_t>(j)]; | |
| 377 |
2/2✓ Branch 0 taken 30 times.
✓ Branch 1 taken 10 times.
|
40 | if (mean) { |
| 378 |
7/14✓ Branch 0 taken 30 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 30 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 30 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 30 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 30 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 30 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 30 times.
✗ Branch 13 not taken.
|
30 | b << " sub.f32 " << f_n << ", " << xj << ", " << f_mean << ";\n"; |
| 379 |
7/14✓ Branch 0 taken 30 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 30 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 30 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 30 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 30 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 30 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 30 times.
✗ Branch 13 not taken.
|
30 | b << " mul.f32 " << f_n << ", " << f_n << ", " << f_rstd << ";\n"; |
| 380 | 30 | } else { | |
| 381 |
7/14✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 10 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 10 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 10 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 10 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 10 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 10 times.
✗ Branch 13 not taken.
|
10 | b << " mul.f32 " << f_n << ", " << xj << ", " << f_rstd << ";\n"; |
| 382 | } | ||
| 383 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 40 times.
|
40 | if (!gam.empty()) { |
| 384 |
2/2✓ Branch 0 taken 30 times.
✓ Branch 1 taken 10 times.
|
40 | if (!bet.empty()) { |
| 385 |
5/10✓ Branch 0 taken 30 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 30 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 30 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 30 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 30 times.
✗ Branch 9 not taken.
|
30 | b << " fma.rn.f32 " << f_n << ", " << f_n << ", " |
| 386 |
2/4✓ Branch 0 taken 30 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 30 times.
✗ Branch 3 not taken.
|
30 | << gam[static_cast<std::size_t>(j)] << ", " |
| 387 |
2/4✓ Branch 0 taken 30 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 30 times.
✗ Branch 3 not taken.
|
30 | << bet[static_cast<std::size_t>(j)] << ";\n"; |
| 388 | 30 | } else { | |
| 389 |
5/10✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 10 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 10 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 10 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 10 times.
✗ Branch 9 not taken.
|
10 | b << " mul.f32 " << f_n << ", " << f_n << ", " |
| 390 |
2/4✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 10 times.
✗ Branch 3 not taken.
|
10 | << gam[static_cast<std::size_t>(j)] << ";\n"; |
| 391 | } | ||
| 392 | 40 | } | |
| 393 |
1/2✓ Branch 0 taken 40 times.
✗ Branch 1 not taken.
|
40 | norm[static_cast<std::size_t>(j)] = f_n; |
| 394 | 40 | } | |
| 395 |
1/2✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
|
16 | vals[plan.norm_node] = std::move(norm); |
| 396 | |||
| 397 |
1/2✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
|
16 | vals = emit_chain(b, a, dag, lanes, std::move(vals), &plan.post_ids); |
| 398 | |||
| 399 |
2/2✓ Branch 0 taken 16 times.
✓ Branch 1 taken 20 times.
|
36 | for (std::size_t m = 0; m < n_out; ++m) { |
| 400 |
3/4✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✓ Branch 3 taken 4 times.
|
20 | if (m < plan.pre_output.size() && plan.pre_output[m]) continue; // written in pass one |
| 401 | 16 | const BufferSpec& s = ew.outputs[m]; | |
| 402 |
1/2✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
|
16 | const std::string off = full_off(o2, r_c, s.dtype); |
| 403 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 16 times.
|
16 | const std::string addr = a.RD(); |
| 404 |
7/14✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 16 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 16 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 16 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 16 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 16 times.
✗ Branch 13 not taken.
|
16 | b << " add.u64 " << addr << ", " << out_ptr[m] << ", " << off << ";\n"; |
| 405 |
2/4✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
|
16 | emit_store(b, a, s.dtype, lanes, addr, vals.at(s.node_id)); |
| 406 | 16 | } | |
| 407 | |||
| 408 |
7/14✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 16 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 16 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 16 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 16 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 16 times.
✗ Branch 13 not taken.
|
16 | b << " add.u32 " << r_c << ", " << r_c << ", " << step << ";\n"; |
| 409 |
3/6✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 16 times.
✗ Branch 5 not taken.
|
16 | b << " bra $L_" << name << "_apply;\n"; |
| 410 |
3/6✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 16 times.
✗ Branch 5 not taken.
|
16 | b << "$L_" << name << "_apply_end:\n"; |
| 411 | 16 | } | |
| 412 | |||
| 413 |
3/6✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 16 times.
✗ Branch 5 not taken.
|
16 | b << "$L_" << name << "_done:\n"; |
| 414 |
1/2✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
|
16 | b << " ret;\n"; |
| 415 | |||
| 416 |
1/2✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
|
16 | std::stringstream ss; |
| 417 |
3/6✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 16 times.
✗ Branch 5 not taken.
|
16 | ss << ".visible .entry " << name << "(\n"; |
| 418 |
2/2✓ Branch 0 taken 20 times.
✓ Branch 1 taken 16 times.
|
36 | for (std::size_t m = 0; m < n_out; ++m) { |
| 419 |
5/10✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 20 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 20 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 20 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 20 times.
✗ Branch 9 not taken.
|
20 | ss << " .param .u64 " << name << "_out_" << m << ",\n"; |
| 420 | 20 | } | |
| 421 |
2/2✓ Branch 0 taken 16 times.
✓ Branch 1 taken 36 times.
|
52 | for (std::size_t k = 0; k + 1 < n_in; ++k) { |
| 422 |
5/10✓ Branch 0 taken 36 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 36 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 36 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 36 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 36 times.
✗ Branch 9 not taken.
|
36 | ss << " .param .u64 " << name << "_in_" << k << ",\n"; |
| 423 | 36 | } | |
| 424 |
5/10✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 16 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 16 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 16 times.
✗ Branch 9 not taken.
|
16 | ss << " .param .u64 " << name << "_in_" << (n_in - 1) << "\n"; |
| 425 |
1/2✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
|
16 | ss << ") {\n"; |
| 426 |
2/4✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
|
16 | ss << declarations(a); |
| 427 |
2/4✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
|
16 | ss << b.str(); |
| 428 |
1/2✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
|
16 | ss << "}\n\n"; |
| 429 |
1/2✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
|
16 | return ss.str(); |
| 430 | 16 | } | |
| 431 | |||
| 432 | } // namespace | ||
| 433 | |||
| 434 | 8 | std::string emit_row_norm(const TraceDAG& dag, const RowNormPlan& plan, | |
| 435 | const std::string& arch) { | ||
| 436 | 8 | std::stringstream ss; | |
| 437 |
1/2✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
|
8 | ss << ".version 7.8\n"; |
| 438 |
3/6✓ 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.
|
8 | ss << ".target " << arch << "\n"; |
| 439 |
1/2✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
|
8 | ss << ".address_size 64\n\n"; |
| 440 | // Per-CTA partials: one slot per warp for the sum of squares (RMS) or of | ||
| 441 | // squared deviations (LayerNorm), and for LayerNorm one more set for the | ||
| 442 | // plain sum. Shared by both entries. | ||
| 443 |
3/6✓ 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.
|
8 | ss << ".shared .align 4 .b32 rn_red[" << (2 * (kRowThreads / 32)) << "];\n\n"; |
| 444 |
3/6✓ 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.
|
8 | ss << emit_rn_entry(dag, plan, plan.ew.vec, kEntryRowNorm); |
| 445 |
3/6✓ 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.
|
8 | ss << emit_rn_entry(dag, plan, 1, kEntryRowNormScalar); |
| 446 |
1/2✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
|
8 | return ss.str(); |
| 447 | 8 | } | |
| 448 | |||
| 449 | } // namespace brotensor::jit::ptx | ||
| 450 |