File size: 4,629 Bytes
8eeb550 | 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 | // SPDX-License-Identifier: Apache-2.0
// ARGMAX (meteor-p150 OPT round 5 R5.3, tt/argmax.py): per-row argmax over the C channels of a [1, 1, R, Cpad] fp32
// TILE tensor (R % 32 == 0) -> uint32 ROW_MAJOR [1, R / 32, 32] (row r at element r: the stock channel_argmax output
// after its untilize), in ONE data-movement program. Replaces ttnn's transpose + TILE reshape (W = C -> 32-row
// planes: 1.1-3.3 ms at METEOR's maps) + ArgMaxNC + untilize.
// Semantics of ArgMaxNC (argmax_nc_compute.cpp): fp32 values (unpacked to DEST as fp32), running max with a strict
// `>` in channel order (the first maximum wins ties). Here the compare is done on the fp32 bits as an unsigned key
// (sign-magnitude -> monotone order; -0.0 is mapped to +0.0 first, so the two zeros tie as in a float compare).
// NaN inputs are not specified (ArgMaxNC's NaN handling is not PyTorch's either).
// Work unit = 32 rows (one tile row: Ct tiles of 4 KB); element (r, c) of a tile row sits in tile c / 32, face
// (r / 16) * 2 + (c % 32) / 16, at (r % 16) * 16 + c % 16 within the face. One RISC / NOC does the reads and the
// 128 B output writes (no two-NOC small DRAM writes, SYS-1419); the reads of unit u + 1 are in flight while unit u is
// reduced.
// Work split (no per-core RT args, probe P3): core i = y * GX + x of the first NC cores owns units
// [i * B + min(i, E), + B + (i < E)), B = NU / NC, E = NU % NC.
// CT args: [0] NU (= R / 32), [1] NC, [2] GX, [3] C, [4] Ct, then the TensorAccessorArgs of the input and the output.
// Common RT args: [in_addr, out_addr]. CB 0: 2 * Ct tiles (input slots); CB 1: 2 * 128 B (output slots); both used as
// plain L1 scratch.
#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
FORCE_INLINE uint32_t key_of(uint32_t b) {
b = (b == 0x80000000u) ? 0u : b; // -0.0 == +0.0
return (b & 0x80000000u) ? ~b : (b | 0x80000000u); // monotone unsigned order of the floats
}
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 < 32; ++r) {
const uint32_t rbase = (r >> 4) * 512 + (r & 15) * 16; // face row offset (faces 0 / 2 for c % 32 < 16)
uint32_t best = 0;
uint32_t bkey = key_of(src[rbase]);
for (uint32_t c = 1; c < C; ++c) {
const uint32_t cc = c & 31;
const uint32_t idx = (c >> 5) * 1024 + rbase + (cc >> 4) * 256 + (cc & 15);
const uint32_t k = key_of(src[idx]);
if (k > bkey) {
bkey = k;
best = c;
}
}
dst[r] = best;
}
noc_async_write(ob, out.get_noc_addr(u0 + i), OB);
}
noc_async_write_barrier();
}
|