257 lines
11 KiB
C++
257 lines
11 KiB
C++
#pragma once
|
|
#include "../utils.h"
|
|
#include <runtime_op_utility.h>
|
|
#include <targets/k210/k210_runtime_op_utility.h>
|
|
|
|
namespace nncase
|
|
{
|
|
namespace kernels
|
|
{
|
|
namespace k210
|
|
{
|
|
inline void kpu_upload(const uint8_t *src, uint8_t *dest, const runtime_shape_t &in_shape)
|
|
{
|
|
if (in_shape[3] % 64 == 0)
|
|
{
|
|
std::copy(src, src + kernels::details::compute_size(in_shape), dest);
|
|
}
|
|
else
|
|
{
|
|
auto layout = targets::k210::get_kpu_row_layout(in_shape[3]);
|
|
auto fmap_size = targets::k210::get_kpu_bytes(in_shape[3], in_shape[2], in_shape[1]);
|
|
|
|
for (int32_t batch = 0; batch < in_shape[0]; batch++)
|
|
{
|
|
auto batch_origin = dest + (size_t)batch * fmap_size;
|
|
for (int32_t oc = 0; oc < in_shape[1]; oc++)
|
|
{
|
|
auto channel_origin = batch_origin + (size_t)oc / layout.groups * layout.row_len * in_shape[2] * 64 + (size_t)oc % layout.groups * layout.row_pitch;
|
|
for (int32_t y = 0; y < in_shape[2]; y++)
|
|
{
|
|
auto y_origin = channel_origin + (size_t)y * layout.row_len * 64;
|
|
std::copy(src, src + in_shape[3], y_origin);
|
|
src += in_shape[3];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#if NNCASE_TARGET_K210_SIMULATOR
|
|
|
|
inline void kpu_download(const uint8_t *src, uint8_t *dest, const runtime_shape_t &in_shape)
|
|
{
|
|
if (in_shape[3] % 64 == 0)
|
|
{
|
|
std::copy(src, src + kernels::details::compute_size(in_shape), dest);
|
|
}
|
|
else
|
|
{
|
|
auto layout = targets::k210::get_kpu_row_layout(in_shape[3]);
|
|
auto fmap_size = targets::k210::get_kpu_bytes(in_shape[3], in_shape[2], in_shape[1]);
|
|
|
|
for (int32_t batch = 0; batch < in_shape[0]; batch++)
|
|
{
|
|
auto batch_origin = src + (size_t)batch * fmap_size;
|
|
for (int32_t oc = 0; oc < in_shape[1]; oc++)
|
|
{
|
|
auto channel_origin = batch_origin + (size_t)oc / layout.groups * layout.row_len * in_shape[2] * 64 + (size_t)oc % layout.groups * layout.row_pitch;
|
|
for (int32_t y = 0; y < in_shape[2]; y++)
|
|
{
|
|
auto y_origin = channel_origin + (size_t)y * layout.row_len * 64;
|
|
for (int32_t x = 0; x < in_shape[3]; x++)
|
|
*dest++ = y_origin[x];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <bool IsDepthwise, int32_t FilterSize>
|
|
void kpu_conv2d(const uint8_t *input, int64_t *workspace, uint8_t *output, const uint8_t *weights, int32_t in_h, int32_t in_w, int32_t in_channels, int32_t out_channels, uint8_t pad_value, int32_t arg_x,
|
|
int32_t shift_x, int32_t arg_w, int32_t shift_w, int64_t arg_add, const targets::k210::kpu_batchnorm_segment *batchnorm, const targets::k210::kpu_activation_table_t &activation)
|
|
{
|
|
const auto channel_size = size_t(in_h) * in_w;
|
|
// conv
|
|
{
|
|
auto out_it = workspace;
|
|
const auto pad = FilterSize == 1 ? 0 : 1;
|
|
const auto groups = IsDepthwise ? out_channels : 1;
|
|
const auto g_ic = IsDepthwise ? 1 : in_channels / groups;
|
|
const auto g_oc = IsDepthwise ? 1 : out_channels;
|
|
|
|
for (int32_t og = 0; og < groups; og++)
|
|
{
|
|
const uint8_t *w_group_p = weights + (size_t)og * g_oc * g_ic * FilterSize * FilterSize;
|
|
|
|
for (int32_t oc = 0; oc < g_oc; oc++)
|
|
{
|
|
const uint8_t *w_oc_p = w_group_p + (size_t)oc * g_ic * FilterSize * FilterSize;
|
|
|
|
for (int32_t oy = 0; oy < in_h; oy++)
|
|
{
|
|
for (int32_t ox = 0; ox < in_w; ox++)
|
|
{
|
|
const int32_t in_y_origin = oy - pad;
|
|
const int32_t in_x_origin = ox - pad;
|
|
int64_t value = 0;
|
|
int64_t sum_x = 0, sum_w = 0;
|
|
|
|
for (int32_t ic = 0; ic < g_ic; ic++)
|
|
{
|
|
const uint8_t *in_c_p = input + ((size_t)og * g_ic + ic) * in_h * in_w;
|
|
const uint8_t *w_ic_p = w_oc_p + (size_t)ic * FilterSize * FilterSize;
|
|
|
|
for (int32_t ky = 0; ky < FilterSize; ky++)
|
|
{
|
|
for (int32_t kx = 0; kx < FilterSize; kx++)
|
|
{
|
|
const int32_t in_y = in_y_origin + ky;
|
|
const int32_t in_x = in_x_origin + kx;
|
|
|
|
uint8_t x;
|
|
if (in_x < 0 || in_x >= in_w
|
|
|| in_y < 0 || in_y >= in_h)
|
|
x = pad_value;
|
|
else
|
|
x = in_c_p[in_y * in_w + in_x];
|
|
|
|
uint8_t w = w_ic_p[ky * FilterSize + kx];
|
|
|
|
sum_x += x;
|
|
sum_w += w;
|
|
value += (int32_t)x * w;
|
|
}
|
|
}
|
|
}
|
|
|
|
*out_it++ = value + (arg_x * sum_x >> shift_x) + (arg_w * sum_w >> shift_w) + arg_add * g_ic;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// bn act
|
|
{
|
|
auto src_it = workspace;
|
|
auto out_it = output;
|
|
for (int32_t oc = 0; oc < out_channels; oc++)
|
|
{
|
|
const auto &bn = batchnorm[oc];
|
|
for (size_t i = 0; i < channel_size; i++)
|
|
{
|
|
auto value = (*src_it++ * bn.mul >> bn.shift) + bn.add;
|
|
auto &seg = *std::find_if(activation.rbegin(), activation.rend(), [value](const targets::k210::kpu_activation_segment &seg) {
|
|
return value > seg.start_x;
|
|
});
|
|
value = runtime::carry_shift((value - seg.start_x) * seg.mul, seg.shift);
|
|
*out_it++ = (uint8_t)std::clamp(value, int64_t(0), int64_t(255));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
inline void kpu_pool2d(const uint8_t *input, uint8_t *output, int32_t in_h, int32_t in_w, int32_t in_channels, targets::k210::kpu_pool_type_t pool_type)
|
|
{
|
|
using namespace targets::k210;
|
|
|
|
const auto filter = get_kpu_filter_size(pool_type);
|
|
const auto stride = get_kpu_filter_stride(pool_type);
|
|
const auto out_h = get_kpu_pool_output_size(in_h, pool_type);
|
|
const auto out_w = get_kpu_pool_output_size(in_w, pool_type);
|
|
|
|
for (int32_t oc = 0; oc < in_channels; oc++)
|
|
{
|
|
auto in_c_p = input + (size_t)oc * in_h * in_w;
|
|
|
|
for (int32_t oy = 0; oy < out_h; oy++)
|
|
{
|
|
for (int32_t ox = 0; ox < out_w; ox++)
|
|
{
|
|
const int32_t in_y_origin = oy * stride;
|
|
const int32_t in_x_origin = ox * stride;
|
|
int32_t value = 0;
|
|
|
|
switch (pool_type)
|
|
{
|
|
case kpu_pool_bypass:
|
|
{
|
|
const int32_t in_y = in_y_origin;
|
|
const int32_t in_x = in_x_origin;
|
|
|
|
value = in_c_p[in_y * in_w + in_x];
|
|
break;
|
|
}
|
|
case kpu_pool_max_2_s2:
|
|
case kpu_pool_max_2_s1:
|
|
case kpu_pool_max_4_s4:
|
|
{
|
|
for (int32_t ky = 0; ky < filter; ky++)
|
|
{
|
|
for (int32_t kx = 0; kx < filter; kx++)
|
|
{
|
|
const int32_t in_y = in_y_origin + ky;
|
|
const int32_t in_x = in_x_origin + kx;
|
|
int32_t in_v;
|
|
|
|
if (in_y < 0 || in_y >= in_h || in_x < 0 || in_x >= in_w)
|
|
in_v = 0;
|
|
else
|
|
in_v = in_c_p[in_y * in_w + in_x];
|
|
|
|
value = std::max(value, in_v);
|
|
}
|
|
}
|
|
|
|
break;
|
|
}
|
|
case kpu_pool_mean_2_s2:
|
|
case kpu_pool_mean_2_s1:
|
|
case kpu_pool_mean_4_s4:
|
|
{
|
|
for (int32_t ky = 0; ky < filter; ky++)
|
|
{
|
|
for (int32_t kx = 0; kx < filter; kx++)
|
|
{
|
|
const int32_t in_y = std::clamp(in_y_origin + ky, 0, in_h - 1);
|
|
const int32_t in_x = std::clamp(in_x_origin + kx, 0, in_w - 1);
|
|
const int32_t in_v = in_c_p[in_y * in_w + in_x];
|
|
|
|
value += in_v;
|
|
}
|
|
}
|
|
|
|
value /= filter * filter;
|
|
break;
|
|
}
|
|
case kpu_pool_left_top_2_s2:
|
|
case kpu_pool_left_top_4_s4:
|
|
case kpu_pool_right_top_2_s2:
|
|
{
|
|
auto k_off = get_kpu_select_pool_offset(pool_type);
|
|
const int32_t in_y = in_y_origin + k_off[0];
|
|
const int32_t in_x = in_x_origin + k_off[1];
|
|
int32_t in_v;
|
|
|
|
if (in_y < 0 || in_y >= in_h || in_x < 0 || in_x >= in_w)
|
|
in_v = 0;
|
|
else
|
|
in_v = in_c_p[in_y * in_w + in_x];
|
|
|
|
value = in_v;
|
|
break;
|
|
}
|
|
}
|
|
|
|
*output++ = (uint8_t)value;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#endif
|
|
}
|
|
}
|
|
}
|