code/tt_meteor/tt/kernels/argmax_rows.cpp

// 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();
}