| |
| |
| |
| |
| |
| #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.h" |
|
|
| using namespace ckernel; |
|
|
| 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 cb_a = 0, cb_b = 2, cb_out = 16; |
|
|
| 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 n_units = B + (core < E ? 1 : 0); |
| if (n_units == 0) { |
| return; |
| } |
| compute_kernel_hw_startup(cb_b, cb_a, cb_out); |
| add_init(cb_b, cb_a); |
| for (uint32_t u = 0; u < n_units; ++u) { |
| cb_wait_front(cb_b, 3); |
| cb_wait_front(cb_a, 3); |
| cb_reserve_back(cb_out, 3); |
| tile_regs_acquire(); |
| for (uint32_t k = 0; k < 3; ++k) { |
| add_tiles(cb_b, cb_a, k, k, k); |
| } |
| tile_regs_commit(); |
| tile_regs_wait(); |
| for (uint32_t k = 0; k < 3; ++k) { |
| pack_tile(k, cb_out); |
| } |
| tile_regs_release(); |
| cb_push_back(cb_out, 3); |
| cb_pop_front(cb_a, 3); |
| cb_pop_front(cb_b, 3); |
| } |
| } |
|
|
