| |
| |
| |
| |
| |
| |
| |
| #include <cstdint> |
|
|
| #include "api/compute/common.h" |
| #include "api/compute/compute_kernel_api.h" |
| #include "api/compute/compute_kernel_hw_startup.h" |
| #include "api/compute/eltwise_binary_sfpu.h" |
| #include "api/compute/eltwise_unary/eltwise_unary.h" |
| #include "api/compute/eltwise_unary/typecast.h" |
| #include "api/compute/pack.h" |
| #include "api/compute/reconfig_data_format.h" |
| #include "api/compute/tile_move_copy.h" |
|
|
| constexpr uint32_t MODE = get_compile_time_arg_val(0); |
| constexpr uint32_t NT = 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 NCOLS = get_compile_time_arg_val(4); |
| constexpr uint32_t HAS_C = get_compile_time_arg_val(5); |
| constexpr uint32_t cb0 = 0, cb1 = 1, cb2 = 2, cbo0 = 16, cbo1 = 17; |
| constexpr uint32_t F32 = static_cast<uint32_t>(DataFormat::Float32); |
| constexpr uint32_t I32 = static_cast<uint32_t>(DataFormat::Int32); |
|
|
| void kernel_main() { |
| const uint32_t ci = get_absolute_logical_y() * GX + get_absolute_logical_x(); |
| if (ci >= NC) { |
| return; |
| } |
| const uint32_t t0 = ci * NT / NC; |
| const uint32_t t1 = (ci + 1) * NT / NC; |
| if (t0 >= t1) { |
| return; |
| } |
| compute_kernel_hw_startup(cb0, cbo0); |
| if constexpr (MODE == 0 && HAS_C) { |
| cb_wait_front(cb2, NCOLS); |
| } |
| for (uint32_t t = t0; t < t1; ++t) { |
| cb_wait_front(cb0, 1); |
| cb_wait_front(cb1, 1); |
| if constexpr (MODE == 1) { |
| cb_wait_front(cb2, 1); |
| } |
| tile_regs_acquire(); |
| reconfig_data_format_srca(cb0); |
| copy_tile_to_dst_init_short(cb0); |
| copy_tile(cb0, 0, 0); |
| reconfig_data_format_srca(cb1); |
| copy_tile_to_dst_init_short(cb1); |
| copy_tile(cb1, 0, 1); |
| if constexpr (MODE == 0) { |
| add_binary_tile_init(); |
| add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0); |
| if constexpr (HAS_C) { |
| reconfig_data_format_srca(cb2); |
| copy_tile_to_dst_init_short(cb2); |
| copy_tile(cb2, t % NCOLS, 1); |
| add_binary_tile_init(); |
| add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0); |
| } |
| } else { |
| mul_binary_tile_init(); |
| mul_binary_tile(0, 1, 0); |
| reconfig_data_format_srca(cb2); |
| copy_tile_to_dst_init_short(cb2); |
| copy_tile(cb2, 0, 1); |
| add_binary_tile_init(); |
| add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0); |
| reconfig_data_format_srca(cb0); |
| copy_tile_to_dst_init_short(cb0); |
| copy_tile(cb0, 0, 2); |
| sigmoid_tile_init<false>(); |
| sigmoid_tile<VectorMode::RC, false>(2); |
| } |
| typecast_tile_init<F32, I32>(); |
| typecast_tile<F32, I32>(0); |
| tile_regs_commit(); |
| cb_reserve_back(cbo0, 1); |
| if constexpr (MODE == 1) { |
| cb_reserve_back(cbo1, 1); |
| } |
| tile_regs_wait(); |
| pack_reconfig_data_format(cbo0); |
| pack_tile(0, cbo0); |
| if constexpr (MODE == 1) { |
| pack_reconfig_data_format(cbo1); |
| pack_tile(2, cbo1); |
| } |
| tile_regs_release(); |
| cb_push_back(cbo0, 1); |
| if constexpr (MODE == 1) { |
| cb_push_back(cbo1, 1); |
| } |
| cb_pop_front(cb0, 1); |
| cb_pop_front(cb1, 1); |
| if constexpr (MODE == 1) { |
| cb_pop_front(cb2, 1); |
| } |
| } |
| } |
|
|
