src/api/native_tensor_extra_util.h
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #pragma once | ||
| 2 | |||
| 3 | // native_tensor_extra_util.h — argument checks shared by the native_tensor_extra* | ||
| 4 | // files, on top of api_internal.h's needElems / needSameSize / needSameDtype. | ||
| 5 | // | ||
| 6 | // Two things the extra group needs beyond those: | ||
| 7 | // - dims products that cannot overflow (a JS caller can pass any int32, and | ||
| 8 | // five of them multiply past INT64_MAX), and output extents that must fit | ||
| 9 | // the int the ops size tensors with; | ||
| 10 | // - INT32 index operands whose VALUES address memory (scatter rows, RoPE | ||
| 11 | // positions, sequence bounds, segment offsets). The kernels do not | ||
| 12 | // bounds-check those, so the native reads them on the host and checks | ||
| 13 | // them before the launch. The same read lets an FP32 tensor of whole | ||
| 14 | // numbers (what GpuTensor.prototype.upload always produces) stand in for | ||
| 15 | // an INT32 one, converted to an INT32 copy on the operand's device. | ||
| 16 | |||
| 17 | #include "api_internal.h" | ||
| 18 | |||
| 19 | #include <climits> | ||
| 20 | #include <cstdint> | ||
| 21 | #include <initializer_list> | ||
| 22 | #include <string> | ||
| 23 | #include <vector> | ||
| 24 | |||
| 25 | namespace brotensor::api::extra { | ||
| 26 | |||
| 27 | // Product of dims in 64 bits: -1 when any dim is negative, INT64_MAX when the | ||
| 28 | // product would overflow (needElems then reports it as too large). | ||
| 29 | 326 | inline int64_t prod(std::initializer_list<int64_t> dims) { | |
| 30 | 326 | int64_t n = 1; | |
| 31 |
3/4✓ Branch 0 taken 960 times.
✓ Branch 1 taken 326 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 960 times.
|
1286 | for (int64_t d : dims) if (d < 0) return -1; |
| 32 |
2/2✓ Branch 0 taken 948 times.
✓ Branch 1 taken 312 times.
|
1260 | for (int64_t d : dims) { |
| 33 |
2/2✓ Branch 0 taken 14 times.
✓ Branch 1 taken 934 times.
|
948 | if (d == 0) return 0; |
| 34 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 934 times.
|
934 | if (n > INT64_MAX / d) return INT64_MAX; |
| 35 | 934 | n *= d; | |
| 36 | } | ||
| 37 | 312 | return n; | |
| 38 | 326 | } | |
| 39 | |||
| 40 | // An output extent (or element count) the op will allocate: positive and | ||
| 41 | // within the int the ops size tensors with. | ||
| 42 | 150 | inline bool needExtent(const char* label, const char* what, int64_t n) { | |
| 43 |
1/2✓ Branch 0 taken 150 times.
✗ Branch 1 not taken.
|
150 | if (n <= 0) { |
| 44 | ✗ | setError(std::string(label) + ": " + what + " is " + std::to_string(n) + ", must be positive"); | |
| 45 | ✗ | return false; | |
| 46 | } | ||
| 47 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 150 times.
|
150 | if (n > INT_MAX) { |
| 48 | ✗ | setError(std::string(label) + ": " + what + " is too large"); | |
| 49 | ✗ | return false; | |
| 50 | } | ||
| 51 | 150 | return true; | |
| 52 | 150 | } | |
| 53 | |||
| 54 | // A scalar argument in [lo, hi]. | ||
| 55 | 36 | inline bool needRange(const char* label, const char* what, int64_t v, int64_t lo, int64_t hi) { | |
| 56 |
2/4✓ Branch 0 taken 36 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 36 times.
|
36 | if (v < lo || v > hi) { |
| 57 | ✗ | setError(std::string(label) + ": " + what + " is " + std::to_string(v) + ", must be in [" + | |
| 58 | ✗ | std::to_string(lo) + ", " + std::to_string(hi) + "]"); | |
| 59 | ✗ | return false; | |
| 60 | } | ||
| 61 | 36 | return true; | |
| 62 | 36 | } | |
| 63 | |||
| 64 | 184 | inline bool needDtype(const char* label, const char* name, const brotensor::Tensor* t, brotensor::Dtype dt) { | |
| 65 |
1/2✓ Branch 0 taken 184 times.
✗ Branch 1 not taken.
|
184 | if (!t) { |
| 66 | ✗ | setError(std::string(label) + ": " + name + " must be a GpuTensor"); | |
| 67 | ✗ | return false; | |
| 68 | } | ||
| 69 |
2/2✓ Branch 0 taken 2 times.
✓ Branch 1 taken 182 times.
|
184 | if (t->dtype != dt) { |
| 70 |
7/14✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 2 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 2 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 2 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 2 times.
✗ Branch 13 not taken.
|
2 | setError(std::string(label) + ": " + name + " is " + dtypeToString(t->dtype) + ", expected " + |
| 71 | 2 | dtypeToString(dt)); | |
| 72 | 2 | return false; | |
| 73 | } | ||
| 74 | 182 | return true; | |
| 75 | 184 | } | |
| 76 | |||
| 77 | // FP32 / FP16 / BF16 — the float storage the dtype-dispatched kernels read. | ||
| 78 | 134 | inline bool needFloat(const char* label, const char* name, const brotensor::Tensor* t) { | |
| 79 |
1/2✓ Branch 0 taken 134 times.
✗ Branch 1 not taken.
|
134 | if (!t) { |
| 80 | ✗ | setError(std::string(label) + ": " + name + " must be a GpuTensor"); | |
| 81 | ✗ | return false; | |
| 82 | } | ||
| 83 |
5/6✓ Branch 0 taken 10 times.
✓ Branch 1 taken 124 times.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 8 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 2 times.
|
134 | if (t->dtype != brotensor::Dtype::FP32 && t->dtype != brotensor::Dtype::FP16 && |
| 84 | 2 | t->dtype != brotensor::Dtype::BF16) { | |
| 85 |
6/14✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 2 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 2 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 2 times.
✗ Branch 11 not taken.
✗ Branch 12 not taken.
✗ Branch 13 not taken.
|
2 | setError(std::string(label) + ": " + name + " is " + dtypeToString(t->dtype) + |
| 86 | ", expected fp32, fp16 or bf16"); | ||
| 87 | 2 | return false; | |
| 88 | } | ||
| 89 | 132 | return true; | |
| 90 | 134 | } | |
| 91 | |||
| 92 | // An optional operand: null passes; otherwise it must hold `n` elements of | ||
| 93 | // `ref`'s dtype (ref null = any dtype). | ||
| 94 | 48 | inline bool optElems(const char* label, const char* name, const brotensor::Tensor* t, int64_t n, | |
| 95 | const brotensor::Tensor* ref) { | ||
| 96 |
2/2✓ Branch 0 taken 32 times.
✓ Branch 1 taken 16 times.
|
48 | if (!t) return true; |
| 97 |
2/2✓ Branch 0 taken 22 times.
✓ Branch 1 taken 10 times.
|
32 | if (!needElems(label, name, t, n)) return false; |
| 98 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 22 times.
|
22 | return !ref || needSameDtype(label, name, t, ref); |
| 99 | 48 | } | |
| 100 | |||
| 101 | // An optional FP32 operand (the `const float* d_mask` slots): null passes. | ||
| 102 | 44 | inline bool optF32Elems(const char* label, const char* name, const brotensor::Tensor* t, int64_t n) { | |
| 103 |
2/2✓ Branch 0 taken 18 times.
✓ Branch 1 taken 26 times.
|
44 | if (!t) return true; |
| 104 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 18 times.
|
18 | return needDtype(label, name, t, brotensor::Dtype::FP32) && needElems(label, name, t, n); |
| 105 | 44 | } | |
| 106 | |||
| 107 | // The values of an index operand on the host: INT32 read as-is, FP32 whole | ||
| 108 | // numbers converted. Anything else (or fewer than `n` elements) is an error. | ||
| 109 | bool hostInt32(const char* label, const char* name, const brotensor::Tensor* t, int64_t n, | ||
| 110 | std::vector<int32_t>& out); | ||
| 111 | |||
| 112 | // `t` itself when it is already INT32, else an INT32 copy of `vals` shaped | ||
| 113 | // like `t` on `t`'s device, held in `scratch`. | ||
| 114 | const brotensor::Tensor& int32Operand(const brotensor::Tensor& t, const std::vector<int32_t>& vals, | ||
| 115 | brotensor::Tensor& scratch); | ||
| 116 | |||
| 117 | // Every one of the first `n` values lies in [lo, hi). | ||
| 118 | bool needIndicesIn(const char* label, const char* name, const std::vector<int32_t>& vals, int64_t n, | ||
| 119 | int64_t lo, int64_t hi); | ||
| 120 | |||
| 121 | } // namespace brotensor::api::extra | ||
| 122 |