File size: 3,323 Bytes
e0adbab | 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 | // SPDX-License-Identifier: Apache-2.0
// S1 compute (sceneseg-p150 OPT_PLAN S1, tt/subpixel.py): per output tile of the sub-pixel up-block, with the EPI
// LLKs (epi_compute.cpp):
// y = Zhi + Zlo add_binary_tile<NearestEven> (fp32 SFPU, operands unpacked to DEST as fp32)
// y = (y + Shi) + Slo (S1_SKIP: the composed skip conv's two products)
// y = y + bias the reader's bias tile (interior row or the border pattern of this row tile)
// y = gelu(y) gelu_tile<false> (S1_ACT 2: ttnn.gelu's exact erf form)
// out = bf16(y) typecast_tile<Float32, Float16_b> (S1_OUT_BF16)
// CT args: [0] R, [1] CT, [2] NC, [3] GX (work split as the reader: units u = c * R + r, contiguous per core).
// Defines: S1_ACT, S1_OUT_BF16, S1_SKIP, S1_ZLO (the Z lo product present: two weight terms), S1_SLO (likewise for
// the skip conv; round 4 P2: 0 = one weight term, no add, no CB).
#include <cstdint>
#include "api/compute/common.h"
#include "api/compute/compute_kernel_api.h"
#include "api/compute/eltwise_binary_sfpu.h"
#include "api/compute/eltwise_unary/eltwise_unary.h"
#include "api/compute/eltwise_unary/gelu.h"
#include "api/compute/eltwise_unary/typecast.h"
#include "api/compute/tile_move_copy.h"
using namespace ckernel;
constexpr uint32_t R = get_compile_time_arg_val(0);
constexpr uint32_t CT = get_compile_time_arg_val(1);
constexpr uint32_t NC = get_compile_time_arg_val(2);
constexpr uint32_t GX = get_compile_time_arg_val(3);
constexpr uint32_t cb_zh = 0, cb_zl = 1, cb_sh = 2, cb_b = 3, cb_sl = 6, cb_out = 16;
FORCE_INLINE void add_from(uint32_t cb) {
copy_init(cb);
copy_tile(cb, 0, 1);
add_binary_tile_init();
add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
}
void kernel_main() {
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
if (core >= NC) {
return;
}
constexpr uint32_t U = R * CT;
constexpr uint32_t B = U / NC, E = U % NC;
const uint32_t n = B + (core < E ? 1 : 0);
if (n == 0) {
return;
}
compute_kernel_hw_startup(cb_zh, cb_out);
for (uint32_t t = 0; t < n; ++t) {
cb_wait_front(cb_zh, 1);
#if S1_ZLO
cb_wait_front(cb_zl, 1);
#endif
cb_wait_front(cb_b, 1);
#if S1_SKIP
cb_wait_front(cb_sh, 1);
#if S1_SLO
cb_wait_front(cb_sl, 1);
#endif
#endif
tile_regs_acquire();
copy_init(cb_zh);
copy_tile(cb_zh, 0, 0);
#if S1_ZLO
add_from(cb_zl);
#endif
#if S1_SKIP
add_from(cb_sh);
#if S1_SLO
add_from(cb_sl);
#endif
#endif
add_from(cb_b);
#if S1_ACT == 2
gelu_tile_init<false>();
gelu_tile<false>(0);
#endif
#if S1_OUT_BF16
typecast_tile_init<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>();
typecast_tile<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>(0);
#endif
tile_regs_commit();
cb_reserve_back(cb_out, 1);
tile_regs_wait();
pack_tile(0, cb_out);
tile_regs_release();
cb_push_back(cb_out, 1);
cb_pop_front(cb_zh, 1);
#if S1_ZLO
cb_pop_front(cb_zl, 1);
#endif
cb_pop_front(cb_b, 1);
#if S1_SKIP
cb_pop_front(cb_sh, 1);
#if S1_SLO
cb_pop_front(cb_sl, 1);
#endif
#endif
}
}
|