File size: 4,981 Bytes
e91e989 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | // SPDX-License-Identifier: Apache-2.0
// ARGMAX v2 (meteor-p150 OPT2 round 4 R4.2, tt/argmax.py, METEOR_ARGMAX=fast): argmax_rows.cpp with the same data
// movement and the same compare (ArgMaxNC semantics: fp32 values, strict `>` in channel order, first maximum on ties,
// -0.0 == +0.0), only the per-row reduction rewritten. v1 spent ~32 RISC cycles per element (the per-element face index
// arithmetic and a branchy key transform); the reads are not the limit (the kernel is RISC-bound: 42 MB in 1.9 ms).
// - the key transform is branchless: z = (b << 1) ? b : 0 (-0.0 -> +0.0), key = z ^ ((int32(z) >> 31) | 0x80000000)
// (sign set: ~z; sign clear: z | 0x80000000): the monotone unsigned order of the floats, as v1's key_of;
// - the channels of one row are walked face by face (16 contiguous words), the loop bounds are compile-time (C);
// - two rows (r and r + 16: the same word offsets in face rows 0 and 1) are reduced together, so their loads and
// compares interleave.
// Work unit, work split, CT args, common RT args and CBs: as argmax_rows.cpp (no per-core RT args, probe P3; reads
// and writes on one RISC / NOC, SYS-1419).
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t NU = get_compile_time_arg_val(0);
constexpr uint32_t NC = get_compile_time_arg_val(1);
constexpr uint32_t GX = get_compile_time_arg_val(2);
constexpr uint32_t C = get_compile_time_arg_val(3);
constexpr uint32_t CT = get_compile_time_arg_val(4);
constexpr auto in_args = TensorAccessorArgs<5>();
constexpr auto out_args = TensorAccessorArgs<in_args.next_compile_time_args_offset()>();
constexpr uint32_t TB = 4096; // fp32 tile bytes
constexpr uint32_t SLOT = CT * TB;
constexpr uint32_t OB = 128; // one output row: 32 uint32
constexpr uint32_t NF = (C + 15) / 16; // faces (16-channel groups) per row
FORCE_INLINE uint32_t key_of(uint32_t b) {
const uint32_t z = (b << 1) ? b : 0u; // -0.0 == +0.0
return z ^ (static_cast<uint32_t>(static_cast<int32_t>(z) >> 31) | 0x80000000u);
}
// channels [16 g, 16 g + n) of rows r (word offset `a` of face row 0) and r + 16 (a + 512)
template <uint32_t G>
FORCE_INLINE void face(const uint32_t* src, uint32_t a, uint32_t& k0, uint32_t& b0, uint32_t& k1, uint32_t& b1) {
constexpr uint32_t n = (C - 16 * G < 16) ? C - 16 * G : 16;
constexpr uint32_t off = (G >> 1) * 1024 + (G & 1) * 256; // tile G / 2, face column G % 2
const uint32_t* p = src + off + a;
#pragma GCC unroll 16
for (uint32_t j = (G == 0 ? 1 : 0); j < n; ++j) {
const uint32_t x0 = key_of(p[j]);
const uint32_t x1 = key_of(p[j + 512]);
if (x0 > k0) {
k0 = x0;
b0 = 16 * G + j;
}
if (x1 > k1) {
k1 = x1;
b1 = 16 * G + j;
}
}
}
template <uint32_t G>
FORCE_INLINE void faces(const uint32_t* src, uint32_t a, uint32_t& k0, uint32_t& b0, uint32_t& k1, uint32_t& b1) {
if constexpr (G < NF) {
face<G>(src, a, k0, b0, k1, b1);
faces<G + 1>(src, a, k0, b0, k1, b1);
}
}
void kernel_main() {
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
if (core >= NC) {
return;
}
constexpr uint32_t B = NU / NC, E = NU % NC;
const uint32_t u0 = core * B + (core < E ? core : E);
const uint32_t nu = B + (core < E ? 1 : 0);
if (nu == 0) {
return;
}
const auto in = TensorAccessor(in_args, get_common_arg_val<uint32_t>(0), TB);
const auto out = TensorAccessor(out_args, get_common_arg_val<uint32_t>(1), OB);
const uint32_t ibuf = get_write_ptr(0);
const uint32_t obuf = get_write_ptr(1);
auto issue_reads = [&](uint32_t u, uint32_t buf) {
for (uint32_t k = 0; k < CT; ++k) {
noc_async_read(in.get_noc_addr(u * CT + k), buf + k * TB, TB);
}
};
issue_reads(u0, ibuf);
for (uint32_t i = 0; i < nu; ++i) {
const uint32_t ib = ibuf + (i & 1) * SLOT;
const uint32_t ob = obuf + (i & 1) * OB;
noc_async_read_barrier();
if (i + 1 < nu) {
issue_reads(u0 + i + 1, ibuf + ((i + 1) & 1) * SLOT);
}
if (i >= 2) {
noc_async_writes_flushed(); // the write of unit i - 2 (same slot) is out
}
const uint32_t* src = reinterpret_cast<const uint32_t*>(ib);
uint32_t* dst = reinterpret_cast<uint32_t*>(ob);
for (uint32_t r = 0; r < 16; ++r) {
const uint32_t a = r * 16; // row r of face row 0; row r + 16 at a + 512
uint32_t k0 = key_of(src[a]), b0 = 0;
uint32_t k1 = key_of(src[a + 512]), b1 = 0;
faces<0>(src, a, k0, b0, k1, b1);
dst[r] = b0;
dst[r + 16] = b1;
}
noc_async_write(ob, out.get_noc_addr(u0 + i), OB);
}
noc_async_write_barrier();
}
|