src/jit/trace_ops.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #include <brotensor/jit/trace.h> | ||
| 2 | #include "trace_dag.h" | ||
| 3 | #include <brotensor/ops.h> | ||
| 4 | #include "../cpu/cpu_jit.h" | ||
| 5 | |||
| 6 | #if BROTENSOR_HAS_CUDA | ||
| 7 | #include "../cuda/cuda_jit.h" | ||
| 8 | #endif | ||
| 9 | |||
| 10 | #include <cstdint> | ||
| 11 | #include <string> | ||
| 12 | #include <vector> | ||
| 13 | #include <stdexcept> | ||
| 14 | |||
| 15 | namespace brotensor::jit { | ||
| 16 | |||
| 17 | namespace { | ||
| 18 | |||
| 19 | // Result shape of a binary op whose operands may differ in rank: the operand | ||
| 20 | // with more elements wins, so `gate * x` and `x * gate` both produce x's | ||
| 21 | // shape when gate is a (1, D) row or a (1, 1) scalar. The trace records the | ||
| 22 | // broadcast as an addressing mode on that operand; nothing is materialised. | ||
| 23 | 20 | const Tensor& binary_shape(const Tensor& a, const Tensor& b) { | |
| 24 |
2/4✗ Branch 0 not taken.
✓ Branch 1 taken 20 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 20 times.
|
40 | return (static_cast<int64_t>(b.rows) * b.cols > |
| 25 | 20 | static_cast<int64_t>(a.rows) * a.cols) | |
| 26 | ✗ | ? b | |
| 27 | 20 | : a; | |
| 28 | } | ||
| 29 | |||
| 30 | 20 | void broadcast_shape(const Tensor& a, const Tensor& b, int& rows, int& cols) { | |
| 31 |
3/4✓ Branch 0 taken 3 times.
✓ Branch 1 taken 17 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 3 times.
|
20 | const bool ok_r = (a.rows == b.rows || a.rows == 1 || b.rows == 1); |
| 32 |
3/4✓ Branch 0 taken 2 times.
✓ Branch 1 taken 18 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
|
20 | const bool ok_c = (a.cols == b.cols || a.cols == 1 || b.cols == 1); |
| 33 |
2/4✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 20 times.
|
20 | if (ok_r && ok_c) { |
| 34 | 20 | rows = std::max(a.rows, b.rows); | |
| 35 | 20 | cols = std::max(a.cols, b.cols); | |
| 36 | 20 | } else { | |
| 37 | ✗ | const Tensor& r = binary_shape(a, b); | |
| 38 | ✗ | rows = r.rows; | |
| 39 | ✗ | cols = r.cols; | |
| 40 | } | ||
| 41 | 20 | } | |
| 42 | |||
| 43 | 1 | void broadcast_shape3(const Tensor& c, const Tensor& a, const Tensor& b, int& rows, int& cols) { | |
| 44 | 1 | int ab_r = 0, ab_c = 0; | |
| 45 | 1 | broadcast_shape(a, b, ab_r, ab_c); | |
| 46 |
1/6✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
|
2 | const bool ok_r = (c.rows == ab_r || c.rows == 1 || ab_r == 1) && |
| 47 |
1/4✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
1 | (a.rows == b.rows || a.rows == 1 || b.rows == 1); |
| 48 |
1/6✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
|
2 | const bool ok_c = (c.cols == ab_c || c.cols == 1 || ab_c == 1) && |
| 49 |
1/4✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
1 | (a.cols == b.cols || a.cols == 1 || b.cols == 1); |
| 50 |
2/4✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 1 time.
|
1 | if (ok_r && ok_c) { |
| 51 | 1 | rows = std::max(c.rows, ab_r); | |
| 52 | 1 | cols = std::max(c.cols, ab_c); | |
| 53 | 1 | } else { | |
| 54 | ✗ | const Tensor& r = binary_shape(c, binary_shape(a, b)); | |
| 55 | ✗ | rows = r.rows; | |
| 56 | ✗ | cols = r.cols; | |
| 57 | } | ||
| 58 | 1 | } | |
| 59 | |||
| 60 | ✗ | bool same_shape(const Tensor& a, const Tensor& b) { | |
| 61 | ✗ | return a.rows == b.rows && a.cols == b.cols; | |
| 62 | } | ||
| 63 | |||
| 64 | // An operand the caller did not supply — the default-constructed Tensor that | ||
| 65 | // means "no gamma" / "no beta". A symbolic intermediate also has a null | ||
| 66 | // `data`, so emptiness alone does not answer the question. | ||
| 67 | 3 | bool absent(const Tensor& t) { | |
| 68 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 3 times.
|
3 | return t.jit_slot < 0 && t.empty(); |
| 69 | } | ||
| 70 | |||
| 71 | ✗ | [[noreturn]] void no_eager_broadcast(const char* op) { | |
| 72 | ✗ | throw std::runtime_error( | |
| 73 | ✗ | std::string("brotensor::jit: eager ") + op + | |
| 74 | " does not broadcast; trace the expression (begin_trace/end_trace) or " | ||
| 75 | "call the shaped op directly"); | ||
| 76 | ✗ | } | |
| 77 | |||
| 78 | } // namespace | ||
| 79 | |||
| 80 | 8 | Tensor operator+(const Tensor& a, const Tensor& b) { | |
| 81 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
|
8 | if (is_tracing()) { |
| 82 | 8 | auto& ctx = TraceContext::current(); | |
| 83 | 8 | int slot_a = ctx.get_or_register_slot(a); | |
| 84 | 8 | int slot_b = ctx.get_or_register_slot(b); | |
| 85 | 8 | const Tensor& r = binary_shape(a, b); | |
| 86 | 8 | int rows = r.rows, cols = r.cols; | |
| 87 | 8 | broadcast_shape(a, b, rows, cols); | |
| 88 |
2/4✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
|
16 | return ctx.record_op(TraceOpKind::Add, {slot_a, slot_b}, 0.0f, |
| 89 | 8 | r.device, r.dtype, rows, cols); | |
| 90 | } | ||
| 91 | ✗ | if (!same_shape(a, b)) no_eager_broadcast("+"); | |
| 92 | ✗ | Tensor out = a.clone(); | |
| 93 | ✗ | brotensor::add_inplace(out, b); | |
| 94 | ✗ | return out; | |
| 95 | 8 | } | |
| 96 | |||
| 97 | 1 | Tensor operator+(const Tensor& a, float s) { | |
| 98 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
|
1 | if (is_tracing()) { |
| 99 | 1 | auto& ctx = TraceContext::current(); | |
| 100 | 1 | int slot_a = ctx.get_or_register_slot(a); | |
| 101 |
2/4✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
|
2 | return ctx.record_op(TraceOpKind::AddScalar, {slot_a}, s, |
| 102 | 1 | a.device, a.dtype, a.rows, a.cols); | |
| 103 | } | ||
| 104 | ✗ | Tensor out = a.clone(); | |
| 105 | ✗ | brotensor::add_scalar_inplace(out, s); | |
| 106 | ✗ | return out; | |
| 107 | 1 | } | |
| 108 | |||
| 109 | ✗ | Tensor operator-(const Tensor& a, const Tensor& b) { | |
| 110 | ✗ | if (is_tracing()) { | |
| 111 | ✗ | auto& ctx = TraceContext::current(); | |
| 112 | ✗ | int slot_a = ctx.get_or_register_slot(a); | |
| 113 | ✗ | int slot_b = ctx.get_or_register_slot(b); | |
| 114 | ✗ | const Tensor& r = binary_shape(a, b); | |
| 115 | ✗ | int rows = r.rows, cols = r.cols; | |
| 116 | ✗ | broadcast_shape(a, b, rows, cols); | |
| 117 | ✗ | return ctx.record_op(TraceOpKind::Sub, {slot_a, slot_b}, 0.0f, | |
| 118 | ✗ | r.device, r.dtype, rows, cols); | |
| 119 | } | ||
| 120 | ✗ | if (!same_shape(a, b)) no_eager_broadcast("-"); | |
| 121 | ✗ | Tensor out = a.clone(); | |
| 122 | ✗ | brotensor::axpby_inplace(out, b, 1.0f, -1.0f); | |
| 123 | ✗ | return out; | |
| 124 | ✗ | } | |
| 125 | |||
| 126 | ✗ | Tensor operator-(const Tensor& a, float s) { | |
| 127 | ✗ | return a + (-s); | |
| 128 | } | ||
| 129 | |||
| 130 | ✗ | Tensor operator-(float s, const Tensor& a) { | |
| 131 | ✗ | if (is_tracing()) { | |
| 132 | ✗ | auto& ctx = TraceContext::current(); | |
| 133 | ✗ | int slot_a = ctx.get_or_register_slot(a); | |
| 134 | ✗ | return ctx.record_op(TraceOpKind::ScalarSub, {slot_a}, s, | |
| 135 | ✗ | a.device, a.dtype, a.rows, a.cols); | |
| 136 | } | ||
| 137 | ✗ | Tensor out = a.clone(); | |
| 138 | ✗ | brotensor::scale_inplace(out, -1.0f); | |
| 139 | ✗ | brotensor::add_scalar_inplace(out, s); | |
| 140 | ✗ | return out; | |
| 141 | ✗ | } | |
| 142 | |||
| 143 | 11 | Tensor operator*(const Tensor& a, const Tensor& b) { | |
| 144 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 11 times.
|
11 | if (is_tracing()) { |
| 145 | 11 | auto& ctx = TraceContext::current(); | |
| 146 | 11 | int slot_a = ctx.get_or_register_slot(a); | |
| 147 | 11 | int slot_b = ctx.get_or_register_slot(b); | |
| 148 | 11 | const Tensor& r = binary_shape(a, b); | |
| 149 | 11 | int rows = r.rows, cols = r.cols; | |
| 150 | 11 | broadcast_shape(a, b, rows, cols); | |
| 151 |
2/4✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 11 times.
✗ Branch 3 not taken.
|
22 | return ctx.record_op(TraceOpKind::Mul, {slot_a, slot_b}, 0.0f, |
| 152 | 11 | r.device, r.dtype, rows, cols); | |
| 153 | } | ||
| 154 | ✗ | if (!same_shape(a, b)) { | |
| 155 | // The one eager broadcast brotensor already has a kernel for: a | ||
| 156 | // length-D row against (L, D). Anything else is an error. | ||
| 157 | ✗ | const Tensor& big = (static_cast<int64_t>(b.rows) * b.cols > | |
| 158 | ✗ | static_cast<int64_t>(a.rows) * a.cols) | |
| 159 | ✗ | ? b | |
| 160 | ✗ | : a; | |
| 161 | ✗ | const Tensor& small = (&big == &a) ? b : a; | |
| 162 | ✗ | if (small.rows == 1 && small.cols == big.cols) { | |
| 163 | ✗ | Tensor out; | |
| 164 | ✗ | brotensor::broadcast_mul(big, small, out); | |
| 165 | ✗ | return out; | |
| 166 | ✗ | } | |
| 167 | ✗ | no_eager_broadcast("*"); | |
| 168 | } | ||
| 169 | ✗ | Tensor out = a.clone(); | |
| 170 | ✗ | brotensor::mul_inplace(out, b); | |
| 171 | ✗ | return out; | |
| 172 | 11 | } | |
| 173 | |||
| 174 | ✗ | Tensor operator*(const Tensor& a, float s) { | |
| 175 | ✗ | if (is_tracing()) { | |
| 176 | ✗ | auto& ctx = TraceContext::current(); | |
| 177 | ✗ | int slot_a = ctx.get_or_register_slot(a); | |
| 178 | ✗ | return ctx.record_op(TraceOpKind::MulScalar, {slot_a}, s, | |
| 179 | ✗ | a.device, a.dtype, a.rows, a.cols); | |
| 180 | } | ||
| 181 | ✗ | Tensor out = a.clone(); | |
| 182 | ✗ | brotensor::scale_inplace(out, s); | |
| 183 | ✗ | return out; | |
| 184 | ✗ | } | |
| 185 | |||
| 186 | ✗ | Tensor operator/(const Tensor& a, const Tensor& b) { | |
| 187 | ✗ | if (is_tracing()) { | |
| 188 | ✗ | auto& ctx = TraceContext::current(); | |
| 189 | ✗ | int slot_a = ctx.get_or_register_slot(a); | |
| 190 | ✗ | int slot_b = ctx.get_or_register_slot(b); | |
| 191 | ✗ | const Tensor& r = binary_shape(a, b); | |
| 192 | ✗ | int rows = r.rows, cols = r.cols; | |
| 193 | ✗ | broadcast_shape(a, b, rows, cols); | |
| 194 | ✗ | return ctx.record_op(TraceOpKind::Div, {slot_a, slot_b}, 0.0f, | |
| 195 | ✗ | r.device, r.dtype, rows, cols); | |
| 196 | } | ||
| 197 | ✗ | if (!same_shape(a, b)) no_eager_broadcast("/"); | |
| 198 | ✗ | Tensor out = a.clone(); | |
| 199 | ✗ | if (a.is_host()) { | |
| 200 | ✗ | float* p_out = out.ptr(); | |
| 201 | ✗ | const float* p_b = b.ptr(); | |
| 202 | ✗ | int n = a.size(); | |
| 203 | ✗ | for (int i = 0; i < n; ++i) p_out[i] /= p_b[i]; | |
| 204 | ✗ | } else if (a.device.is_vulkan()) { | |
| 205 | ✗ | brotensor::div_inplace(out, b); | |
| 206 | ✗ | } else { | |
| 207 | #if BROTENSOR_HAS_CUDA | ||
| 208 | if (a.device.is_cuda()) { | ||
| 209 | detail::cuda::jit::launch_elementwise_div_ptx( | ||
| 210 | out.ptr(), | ||
| 211 | b.ptr(), | ||
| 212 | a.size() | ||
| 213 | ); | ||
| 214 | } else { | ||
| 215 | throw std::runtime_error("brotensor: elementwise eager div on GPU requires CUDA backend"); | ||
| 216 | } | ||
| 217 | #else | ||
| 218 | ✗ | throw std::runtime_error("brotensor: elementwise eager div on GPU requires CUDA backend"); | |
| 219 | #endif | ||
| 220 | } | ||
| 221 | ✗ | return out; | |
| 222 | ✗ | } | |
| 223 | |||
| 224 | ✗ | Tensor operator/(const Tensor& a, float s) { | |
| 225 | ✗ | return a * (1.0f / s); | |
| 226 | } | ||
| 227 | |||
| 228 | 2 | Tensor& operator+=(Tensor& a, const Tensor& b) { | |
| 229 |
1/2✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
|
2 | if (is_tracing()) { |
| 230 | 2 | auto& ctx = TraceContext::current(); | |
| 231 | 2 | int slot_a = ctx.get_or_register_slot(a); | |
| 232 | 2 | int slot_b = ctx.get_or_register_slot(b); | |
| 233 | 2 | ctx.record_inplace_op(TraceOpKind::Add, slot_a, slot_b, 0.0f, a); | |
| 234 | 2 | return a; | |
| 235 | } | ||
| 236 | ✗ | brotensor::add_inplace(a, b); | |
| 237 | ✗ | return a; | |
| 238 | 2 | } | |
| 239 | |||
| 240 | ✗ | Tensor& operator+=(Tensor& a, float s) { | |
| 241 | ✗ | if (is_tracing()) { | |
| 242 | ✗ | auto& ctx = TraceContext::current(); | |
| 243 | ✗ | int slot_a = ctx.get_or_register_slot(a); | |
| 244 | ✗ | ctx.record_inplace_op(TraceOpKind::AddScalar, slot_a, -1, s, a); | |
| 245 | ✗ | return a; | |
| 246 | } | ||
| 247 | ✗ | brotensor::add_scalar_inplace(a, s); | |
| 248 | ✗ | return a; | |
| 249 | ✗ | } | |
| 250 | |||
| 251 | ✗ | Tensor& operator*=(Tensor& a, const Tensor& b) { | |
| 252 | ✗ | if (is_tracing()) { | |
| 253 | ✗ | auto& ctx = TraceContext::current(); | |
| 254 | ✗ | int slot_a = ctx.get_or_register_slot(a); | |
| 255 | ✗ | int slot_b = ctx.get_or_register_slot(b); | |
| 256 | ✗ | ctx.record_inplace_op(TraceOpKind::Mul, slot_a, slot_b, 0.0f, a); | |
| 257 | ✗ | return a; | |
| 258 | } | ||
| 259 | ✗ | brotensor::mul_inplace(a, b); | |
| 260 | ✗ | return a; | |
| 261 | ✗ | } | |
| 262 | |||
| 263 | ✗ | Tensor& operator*=(Tensor& a, float s) { | |
| 264 | ✗ | if (is_tracing()) { | |
| 265 | ✗ | auto& ctx = TraceContext::current(); | |
| 266 | ✗ | int slot_a = ctx.get_or_register_slot(a); | |
| 267 | ✗ | ctx.record_inplace_op(TraceOpKind::MulScalar, slot_a, -1, s, a); | |
| 268 | ✗ | return a; | |
| 269 | } | ||
| 270 | ✗ | brotensor::scale_inplace(a, s); | |
| 271 | ✗ | return a; | |
| 272 | ✗ | } | |
| 273 | |||
| 274 | 7 | Tensor silu(const Tensor& a) { | |
| 275 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 7 times.
|
7 | if (is_tracing()) { |
| 276 | 7 | auto& ctx = TraceContext::current(); | |
| 277 | 7 | int slot_a = ctx.get_or_register_slot(a); | |
| 278 |
2/4✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 7 times.
✗ Branch 3 not taken.
|
14 | return ctx.record_op(TraceOpKind::SiLU, {slot_a}, 0.0f, |
| 279 | 7 | a.device, a.dtype, a.rows, a.cols); | |
| 280 | } | ||
| 281 | ✗ | Tensor out = Tensor::empty_on(a.device, a.rows, a.cols, a.dtype); | |
| 282 | ✗ | brotensor::silu_forward(a, out); | |
| 283 | ✗ | return out; | |
| 284 | 7 | } | |
| 285 | |||
| 286 | ✗ | Tensor gelu(const Tensor& a) { | |
| 287 | ✗ | if (is_tracing()) { | |
| 288 | ✗ | auto& ctx = TraceContext::current(); | |
| 289 | ✗ | int slot_a = ctx.get_or_register_slot(a); | |
| 290 | ✗ | return ctx.record_op(TraceOpKind::GELU, {slot_a}, 0.0f, | |
| 291 | ✗ | a.device, a.dtype, a.rows, a.cols); | |
| 292 | } | ||
| 293 | ✗ | Tensor out = Tensor::empty_on(a.device, a.rows, a.cols, a.dtype); | |
| 294 | ✗ | brotensor::gelu_forward(a, out); | |
| 295 | ✗ | return out; | |
| 296 | ✗ | } | |
| 297 | |||
| 298 | ✗ | Tensor relu(const Tensor& a) { | |
| 299 | ✗ | if (is_tracing()) { | |
| 300 | ✗ | auto& ctx = TraceContext::current(); | |
| 301 | ✗ | int slot_a = ctx.get_or_register_slot(a); | |
| 302 | ✗ | return ctx.record_op(TraceOpKind::ReLU, {slot_a}, 0.0f, | |
| 303 | ✗ | a.device, a.dtype, a.rows, a.cols); | |
| 304 | } | ||
| 305 | ✗ | Tensor out = Tensor::empty_on(a.device, a.rows, a.cols, a.dtype); | |
| 306 | ✗ | brotensor::relu_forward(a, out); | |
| 307 | ✗ | return out; | |
| 308 | ✗ | } | |
| 309 | |||
| 310 | 1 | void store(Tensor& dst, const Tensor& src) { | |
| 311 |
1/2✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
|
1 | if (is_tracing()) { |
| 312 | 1 | auto& ctx = TraceContext::current(); | |
| 313 | 1 | int slot_src = ctx.get_or_register_slot(src); | |
| 314 | 1 | ctx.record_store(slot_src, dst); | |
| 315 | 1 | return; | |
| 316 | } | ||
| 317 | ✗ | if (dst.rows != src.rows || dst.cols != src.cols || dst.dtype != src.dtype) { | |
| 318 | ✗ | throw std::runtime_error("brotensor::jit: store() needs dst and src to agree on " | |
| 319 | "shape and dtype"); | ||
| 320 | } | ||
| 321 | ✗ | brotensor::copy_d2d(src, 0, dst, 0, src.size()); | |
| 322 | 1 | } | |
| 323 | |||
| 324 | ✗ | Tensor tanh(const Tensor& a) { | |
| 325 | ✗ | if (is_tracing()) { | |
| 326 | ✗ | auto& ctx = TraceContext::current(); | |
| 327 | ✗ | int slot_a = ctx.get_or_register_slot(a); | |
| 328 | ✗ | return ctx.record_op(TraceOpKind::Tanh, {slot_a}, 0.0f, | |
| 329 | ✗ | a.device, a.dtype, a.rows, a.cols); | |
| 330 | } | ||
| 331 | ✗ | Tensor out = Tensor::empty_on(a.device, a.rows, a.cols, a.dtype); | |
| 332 | ✗ | brotensor::tanh_forward(a, out); | |
| 333 | ✗ | return out; | |
| 334 | ✗ | } | |
| 335 | |||
| 336 | ✗ | Tensor sigmoid(const Tensor& a) { | |
| 337 | ✗ | if (is_tracing()) { | |
| 338 | ✗ | auto& ctx = TraceContext::current(); | |
| 339 | ✗ | int slot_a = ctx.get_or_register_slot(a); | |
| 340 | ✗ | return ctx.record_op(TraceOpKind::Sigmoid, {slot_a}, 0.0f, | |
| 341 | ✗ | a.device, a.dtype, a.rows, a.cols); | |
| 342 | } | ||
| 343 | ✗ | Tensor out = Tensor::empty_on(a.device, a.rows, a.cols, a.dtype); | |
| 344 | ✗ | brotensor::sigmoid_forward(a, out); | |
| 345 | ✗ | return out; | |
| 346 | ✗ | } | |
| 347 | |||
| 348 | 1 | Tensor rms_norm(const Tensor& x, const Tensor& gamma, float eps) { | |
| 349 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
|
1 | if (is_tracing()) { |
| 350 | 1 | auto& ctx = TraceContext::current(); | |
| 351 | 1 | int slot_x = ctx.get_or_register_slot(x); | |
| 352 | 1 | std::vector<int> inputs = {slot_x}; | |
| 353 |
2/4✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
|
1 | if (!absent(gamma)) { |
| 354 |
2/4✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
|
1 | inputs.push_back(ctx.get_or_register_slot(gamma)); |
| 355 | 1 | } | |
| 356 |
2/4✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
|
2 | return ctx.record_op(TraceOpKind::RMSNorm, inputs, eps, |
| 357 | 1 | x.device, x.dtype, x.rows, x.cols); | |
| 358 | 1 | } | |
| 359 | |||
| 360 | ✗ | Tensor g = gamma; | |
| 361 | ✗ | if (g.empty()) { | |
| 362 | ✗ | std::vector<float> ones(x.cols, 1.0f); | |
| 363 | ✗ | g = Tensor::from_host_on(x.device, ones.data(), 1, x.cols); | |
| 364 | ✗ | } | |
| 365 | ✗ | Tensor out = Tensor::empty_on(x.device, x.rows, x.cols, x.dtype); | |
| 366 | ✗ | brotensor::rms_norm_forward(x, g, eps, out); | |
| 367 | ✗ | return out; | |
| 368 | 1 | } | |
| 369 | |||
| 370 | 1 | Tensor layernorm(const Tensor& x, const Tensor& gamma, const Tensor& beta, float eps) { | |
| 371 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
|
1 | if (is_tracing()) { |
| 372 | 1 | auto& ctx = TraceContext::current(); | |
| 373 | 1 | int slot_x = ctx.get_or_register_slot(x); | |
| 374 | 1 | std::vector<int> inputs = {slot_x}; | |
| 375 |
4/8✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
✓ Branch 4 taken 1 time.
✗ Branch 5 not taken.
✓ Branch 6 taken 1 time.
✗ Branch 7 not taken.
|
1 | if (!absent(gamma)) inputs.push_back(ctx.get_or_register_slot(gamma)); |
| 376 |
4/8✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
✓ Branch 4 taken 1 time.
✗ Branch 5 not taken.
✓ Branch 6 taken 1 time.
✗ Branch 7 not taken.
|
1 | if (!absent(beta)) inputs.push_back(ctx.get_or_register_slot(beta)); |
| 377 |
2/4✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
|
2 | return ctx.record_op(TraceOpKind::LayerNorm, inputs, eps, |
| 378 | 1 | x.device, x.dtype, x.rows, x.cols); | |
| 379 | 1 | } | |
| 380 | |||
| 381 | ✗ | Tensor g = gamma; | |
| 382 | ✗ | if (g.empty()) { | |
| 383 | ✗ | std::vector<float> ones(x.cols, 1.0f); | |
| 384 | ✗ | g = Tensor::from_host_on(x.device, ones.data(), 1, x.cols); | |
| 385 | ✗ | } | |
| 386 | ✗ | Tensor b = beta; | |
| 387 | ✗ | if (b.empty()) { | |
| 388 | ✗ | b = Tensor::zeros_on(x.device, 1, x.cols); | |
| 389 | ✗ | } | |
| 390 | ✗ | Tensor out = Tensor::empty_on(x.device, x.rows, x.cols, x.dtype); | |
| 391 | ✗ | brotensor::layernorm_forward_inference_batched(x, g, b, out, eps); | |
| 392 | ✗ | return out; | |
| 393 | 1 | } | |
| 394 | |||
| 395 | ✗ | Tensor modulate(const Tensor& x, const Tensor& scale, const Tensor& shift) { | |
| 396 | ✗ | if (is_tracing()) { | |
| 397 | ✗ | auto& ctx = TraceContext::current(); | |
| 398 | ✗ | int slot_x = ctx.get_or_register_slot(x); | |
| 399 | ✗ | int slot_s = ctx.get_or_register_slot(scale); | |
| 400 | ✗ | int slot_sh = ctx.get_or_register_slot(shift); | |
| 401 | ✗ | return ctx.record_op(TraceOpKind::Modulate, {slot_x, slot_s, slot_sh}, 0.0f, | |
| 402 | ✗ | x.device, x.dtype, x.rows, x.cols); | |
| 403 | } | ||
| 404 | |||
| 405 | ✗ | Tensor out = Tensor::empty_on(x.device, x.rows, x.cols, x.dtype); | |
| 406 | ✗ | if (x.is_host()) { | |
| 407 | ✗ | detail::cpu::jit::modulate(x.ptr(), scale.ptr(), shift.ptr(), out.ptr(), x.rows, x.cols); | |
| 408 | ✗ | } else if (x.device.is_vulkan()) { | |
| 409 | ✗ | brotensor::modulate(x, scale, shift, out); | |
| 410 | ✗ | } else { | |
| 411 | #if BROTENSOR_HAS_CUDA | ||
| 412 | detail::cuda::jit::launch_modulate_ptx( | ||
| 413 | static_cast<const float*>(x.data), | ||
| 414 | static_cast<const float*>(scale.data), | ||
| 415 | static_cast<const float*>(shift.data), | ||
| 416 | static_cast<float*>(out.data), | ||
| 417 | x.rows, x.cols | ||
| 418 | ); | ||
| 419 | #else | ||
| 420 | ✗ | throw std::runtime_error("brotensor: modulate on GPU requires CUDA backend"); | |
| 421 | #endif | ||
| 422 | } | ||
| 423 | ✗ | return out; | |
| 424 | ✗ | } | |
| 425 | |||
| 426 | 1 | Tensor where(const Tensor& cond, const Tensor& a, const Tensor& b) { | |
| 427 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
|
1 | if (is_tracing()) { |
| 428 | 1 | auto& ctx = TraceContext::current(); | |
| 429 | 1 | int slot_c = ctx.get_or_register_slot(cond); | |
| 430 | 1 | int slot_a = ctx.get_or_register_slot(a); | |
| 431 | 1 | int slot_b = ctx.get_or_register_slot(b); | |
| 432 | 1 | const Tensor& r = binary_shape(a, b); | |
| 433 | 1 | int rows = r.rows, cols = r.cols; | |
| 434 | 1 | broadcast_shape3(cond, a, b, rows, cols); | |
| 435 |
2/4✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
|
2 | return ctx.record_op(TraceOpKind::Where, {slot_c, slot_a, slot_b}, 0.0f, |
| 436 | 1 | r.device, r.dtype, rows, cols); | |
| 437 | } | ||
| 438 | |||
| 439 | ✗ | if (cond.device != a.device || a.device != b.device) { | |
| 440 | ✗ | throw std::runtime_error("brotensor::jit::where: device mismatch between operands"); | |
| 441 | } | ||
| 442 | ✗ | int rows = 0, cols = 0; | |
| 443 | ✗ | broadcast_shape3(cond, a, b, rows, cols); | |
| 444 | ✗ | auto valid_bcast = [&](const Tensor& t) { | |
| 445 | ✗ | if (t.rows == rows && t.cols == cols) return true; | |
| 446 | ✗ | const BroadcastKind bk = broadcast_of(t.rows, t.cols, rows, cols); | |
| 447 | ✗ | return bk != BroadcastKind::Full; | |
| 448 | ✗ | }; | |
| 449 | ✗ | if (!valid_bcast(cond) || !valid_bcast(a) || !valid_bcast(b)) { | |
| 450 | ✗ | throw std::runtime_error("brotensor::jit::where: incompatible operand shapes"); | |
| 451 | } | ||
| 452 | |||
| 453 | ✗ | const Tensor& r = binary_shape(a, b); | |
| 454 | ✗ | auto to_host_f32 = [](const Tensor& t) -> std::vector<float> { | |
| 455 | ✗ | Tensor h = t.is_host() ? Tensor::view(Device::CPU, const_cast<void*>(t.data), t.rows, t.cols, t.dtype) | |
| 456 | ✗ | : t.to(Device::CPU); | |
| 457 | ✗ | if (h.dtype != Dtype::FP32) { | |
| 458 | ✗ | Tensor c; | |
| 459 | ✗ | brotensor::cast(h, c, Dtype::FP32); | |
| 460 | ✗ | return c.to_host_vector(); | |
| 461 | ✗ | } | |
| 462 | ✗ | return h.to_host_vector(); | |
| 463 | ✗ | }; | |
| 464 | ✗ | const std::vector<float> vc = to_host_f32(cond); | |
| 465 | ✗ | const std::vector<float> va = to_host_f32(a); | |
| 466 | ✗ | const std::vector<float> vb = to_host_f32(b); | |
| 467 | |||
| 468 | ✗ | auto idx_of = [cols](const Tensor& t, int row, int col) -> std::size_t { | |
| 469 | ✗ | if (t.rows == 1 && t.cols == 1) return 0; | |
| 470 | ✗ | if (t.rows == 1) return static_cast<std::size_t>(col); | |
| 471 | ✗ | if (t.cols == 1) return static_cast<std::size_t>(row); | |
| 472 | ✗ | return static_cast<std::size_t>(row) * cols + col; | |
| 473 | ✗ | }; | |
| 474 | |||
| 475 | ✗ | std::vector<float> vout(static_cast<std::size_t>(rows) * cols); | |
| 476 | ✗ | for (int row = 0; row < rows; ++row) { | |
| 477 | ✗ | for (int col = 0; col < cols; ++col) { | |
| 478 | ✗ | const float c_val = vc[idx_of(cond, row, col)]; | |
| 479 | ✗ | const float a_val = va[idx_of(a, row, col)]; | |
| 480 | ✗ | const float b_val = vb[idx_of(b, row, col)]; | |
| 481 | ✗ | vout[static_cast<std::size_t>(row) * cols + col] = (c_val != 0.0f) ? a_val : b_val; | |
| 482 | ✗ | } | |
| 483 | ✗ | } | |
| 484 | |||
| 485 | ✗ | Tensor out_f32 = Tensor::from_host_on(r.device, vout.data(), rows, cols); | |
| 486 | ✗ | if (r.dtype == Dtype::FP32) return out_f32; | |
| 487 | ✗ | Tensor out; | |
| 488 | ✗ | brotensor::cast(out_f32, out, r.dtype); | |
| 489 | ✗ | return out; | |
| 490 | 1 | } | |
| 491 | |||
| 492 | 1 | Tensor clamp(const Tensor& x, float lo, float hi) { | |
| 493 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
|
1 | if (is_tracing()) { |
| 494 | 1 | auto& ctx = TraceContext::current(); | |
| 495 | 1 | int slot_x = ctx.get_or_register_slot(x); | |
| 496 |
2/4✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
|
2 | return ctx.record_op(TraceOpKind::Clamp, {slot_x}, lo, |
| 497 | 1 | x.device, x.dtype, x.rows, x.cols, hi); | |
| 498 | } | ||
| 499 | ✗ | Tensor out = x.clone(); | |
| 500 | ✗ | brotensor::clamp(out, lo, hi); | |
| 501 | ✗ | return out; | |
| 502 | 1 | } | |
| 503 | |||
| 504 | 1 | Tensor cast(const Tensor& x, Dtype out_dtype) { | |
| 505 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
|
1 | if (is_tracing()) { |
| 506 | 1 | auto& ctx = TraceContext::current(); | |
| 507 | 1 | int slot_x = ctx.get_or_register_slot(x); | |
| 508 |
2/4✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
|
2 | return ctx.record_op(TraceOpKind::Cast, {slot_x}, 0.0f, |
| 509 | 1 | x.device, out_dtype, x.rows, x.cols); | |
| 510 | } | ||
| 511 | ✗ | Tensor out; | |
| 512 | ✗ | brotensor::cast(x, out, out_dtype); | |
| 513 | ✗ | return out; | |
| 514 | 1 | } | |
| 515 | |||
| 516 | } // namespace brotensor::jit | ||
| 517 |