File size: 5,178 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 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | // SPDX-License-Identifier: Apache-2.0
// TERMS_EPI compute (meteor-p150 OPT_PLAN #2, tt/epi.py): the TermsConv epilogue of tt/layers.py in one pass per
// tile, with the LLKs and the order of the stock programs it replaces, so the output is bit-identical (the SceneSeg
// EPI pattern, sceneseg-p150 OPT_REPORT R1.3 / R2.3):
// y = (p0 + p1) + p2 add_binary_tile<NearestEven> (binary_ng's fp32 SFPU add; operands unpacked to DEST as
// fp32: UnpackToDestFp32 on CBs 0..4)
// y = y + res add_binary_tile<NearestEven> (EPI_RES: ttaw residual_add(z, idt, "relu"))
// y = relu(y) relu_tile (EPI_RELU: binary_ng's fused post-activation of the last add, unary_op_utils)
// hi = bf16(y), lo = bf16(y - fp32(hi)) (EPI_SPLIT: tt.layers.split_terms = typecast, typecast back (exact),
// fp32 SFPU subtract, typecast), packed to CBs 17 / 18 as bf16
// DEST tiles: 0 = y, 1 = operand / hi, 2 = lo. EPI_MAIN 0 skips the fp32 main output (split only).
// NP = 1 with no residual / relu is the plain split of an fp32 map (split_terms as one program).
// CT args: [0] T, [1] NC, [2] GX, [3] NP. Defines: EPI_RES, EPI_RELU, EPI_SPLIT, EPI_MAIN.
#include <cstdint>
#include "api/compute/common.h"
#include "api/compute/compute_kernel_api.h"
#include "api/compute/copy_dest_values.h"
#include "api/compute/eltwise_binary_sfpu.h"
#include "api/compute/eltwise_unary/eltwise_unary.h"
#include "api/compute/eltwise_unary/relu.h"
#include "api/compute/eltwise_unary/typecast.h"
#include "api/compute/reconfig_data_format.h"
#include "api/compute/tile_move_copy.h"
#ifndef EPI_RES
#define EPI_RES 0
#endif
#ifndef EPI_RELU
#define EPI_RELU 0
#endif
#ifndef EPI_SPLIT
#define EPI_SPLIT 0
#endif
#ifndef EPI_MAIN
#define EPI_MAIN 1
#endif
using namespace ckernel;
constexpr uint32_t T = 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 NP = get_compile_time_arg_val(3);
constexpr uint32_t cb_res = 4, cb_out = 16, cb_hi = 17, cb_lo = 18;
void kernel_main() {
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
if (core >= NC) {
return;
}
constexpr uint32_t B = T / NC, E = T % NC;
const uint32_t t0 = core * B + (core < E ? core : E);
const uint32_t n = B + (core < E ? 1 : 0);
if (n == 0) {
return;
}
#if EPI_MAIN
compute_kernel_hw_startup(0, cb_out);
#else
compute_kernel_hw_startup(0, cb_hi);
#endif
for (uint32_t t = t0; t < t0 + n; ++t) {
cb_wait_front(0, 1);
if constexpr (NP > 1) {
cb_wait_front(1, 1);
}
if constexpr (NP > 2) {
cb_wait_front(2, 1);
}
#if EPI_RES
cb_wait_front(cb_res, 1);
#endif
tile_regs_acquire();
copy_init(0);
copy_tile(0, 0, 0);
if constexpr (NP > 1) {
copy_init(1);
copy_tile(1, 0, 1);
add_binary_tile_init();
add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
}
if constexpr (NP > 2) {
copy_init(2);
copy_tile(2, 0, 1);
add_binary_tile_init();
add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
}
#if EPI_RES
copy_init(cb_res);
copy_tile(cb_res, 0, 1);
add_binary_tile_init();
add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
#endif
#if EPI_RELU
relu_tile_init();
relu_tile(0);
#endif
#if EPI_SPLIT
// hi (tile 1) = bf16(y); lo (tile 2) = bf16(y - hi); y stays in tile 0
copy_dest_values_init();
copy_dest_values<DataFormat::Float32>(0, 1);
copy_dest_values<DataFormat::Float32>(0, 2);
typecast_tile_init<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>();
typecast_tile<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>(1);
sub_binary_tile_init();
sub_binary_tile<ckernel::DstRoundingMode::NearestEven>(2, 1, 2);
typecast_tile_init<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>();
typecast_tile<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>(2);
#endif
tile_regs_commit();
#if EPI_MAIN
cb_reserve_back(cb_out, 1);
#endif
#if EPI_SPLIT
cb_reserve_back(cb_hi, 1);
cb_reserve_back(cb_lo, 1);
#endif
tile_regs_wait();
#if EPI_MAIN
#if EPI_SPLIT
pack_reconfig_data_format(cb_out);
#endif
pack_tile(0, cb_out);
#endif
#if EPI_SPLIT
pack_reconfig_data_format(cb_hi);
pack_tile(1, cb_hi);
pack_tile(2, cb_lo);
#endif
tile_regs_release();
#if EPI_MAIN
cb_push_back(cb_out, 1);
#endif
#if EPI_SPLIT
cb_push_back(cb_hi, 1);
cb_push_back(cb_lo, 1);
#endif
cb_pop_front(0, 1);
if constexpr (NP > 1) {
cb_pop_front(1, 1);
}
if constexpr (NP > 2) {
cb_pop_front(2, 1);
}
#if EPI_RES
cb_pop_front(cb_res, 1);
#endif
}
}
|