#include "src/naive/warp_affine/warp_affine_cv.h"
#include "src/naive/handle.h"
#include "src/common/cv/common.h"
#include "src/common/cv/helper.h"
#include "src/common/cv/interp_helper.h"
#include "src/common/utils.h"
#include "src/common/warp_common.h"
#include "src/naive/handle.h"
#include <climits>
#include <cstring>
using namespace megdnn;
using namespace naive;
using namespace megcv;
using namespace warp;
namespace {
constexpr size_t BLOCK_SZ = 64_z;
template <typename T, InterpolationMode imode, BorderMode bmode, size_t CH>
void warp_affine_cv(
const Mat<T>& src, Mat<T>& dst, const float* trans, const float border_value,
size_t task_id) {
double M[6];
rep(i, 6) M[i] = trans[i];
T bvalue[3] = {(T)border_value, (T)border_value, (T)border_value};
std::vector<int> _adelta(dst.cols() * 2);
int *adelta = _adelta.data(), *bdelta = adelta + dst.cols();
constexpr int AB_BITS = 10 > INTER_BITS ? 10 : INTER_BITS;
constexpr int AB_SCALE = 1 << AB_BITS;
size_t dstcols = dst.cols();
for (size_t x = 0; x < dstcols; ++x) {
adelta[x] = saturate_cast<int>(M[0] * x * AB_SCALE);
bdelta[x] = saturate_cast<int>(M[3] * x * AB_SCALE);
}
size_t x1, y1, dstrows = dst.rows();
size_t BLOCK_SZ_H = std::min(BLOCK_SZ / 2, dstrows);
size_t BLOCK_SZ_W = std::min(BLOCK_SZ * BLOCK_SZ / BLOCK_SZ_H, dstcols);
BLOCK_SZ_H = std::min(BLOCK_SZ * BLOCK_SZ / BLOCK_SZ_W, dstrows);
size_t width_block_size = div_ceil<size_t>(dstcols, BLOCK_SZ_W);
size_t y = (task_id / width_block_size) * BLOCK_SZ_H;
size_t x = (task_id % width_block_size) * BLOCK_SZ_W;
short XY[BLOCK_SZ * BLOCK_SZ * 2], A[BLOCK_SZ * BLOCK_SZ];
int round_delta =
(imode == IMode::INTER_NEAREST ? AB_SCALE / 2
: AB_SCALE / INTER_TAB_SIZE / 2);
size_t bw = std::min(BLOCK_SZ_W, dstcols - x);
size_t bh = std::min(BLOCK_SZ_H, dstrows - y);
Mat<short> _XY(bh, bw, 2, XY);
Mat<T> dpart(dst, y, bh, x, bw);
for (y1 = 0; y1 < bh; ++y1) {
short* xy = XY + y1 * bw * 2;
int X0 = saturate_cast<int>((M[1] * (y + y1) + M[2]) * AB_SCALE) + round_delta;
int Y0 = saturate_cast<int>((M[4] * (y + y1) + M[5]) * AB_SCALE) + round_delta;
if (imode == IMode::INTER_NEAREST) {
x1 = 0;
for (; x1 < bw; x1++) {
int X = (X0 + adelta[x + x1]) >> AB_BITS;
int Y = (Y0 + bdelta[x + x1]) >> AB_BITS;
xy[x1 * 2] = saturate_cast<short>(X);
xy[x1 * 2 + 1] = saturate_cast<short>(Y);
}
} else {
short* alpha = A + y1 * bw;
x1 = 0;
for (; x1 < bw; x1++) {
int X = (X0 + adelta[x + x1]) >> (AB_BITS - INTER_BITS);
int Y = (Y0 + bdelta[x + x1]) >> (AB_BITS - INTER_BITS);
xy[x1 * 2] = saturate_cast<short>(X >> INTER_BITS);
xy[x1 * 2 + 1] = saturate_cast<short>(Y >> INTER_BITS);
alpha[x1] =
(short)((Y & (INTER_TAB_SIZE - 1)) * INTER_TAB_SIZE +
(X & (INTER_TAB_SIZE - 1)));
}
}
}
Mat<ushort> _matA(bh, bw, 1, (ushort*)(A));
remap<T, imode, bmode, CH, RemapVec<T, CH>>(src, dpart, _XY, _matA, bvalue);
}
}
void megdnn::naive::warp_affine_cv_exec(
_megdnn_tensor_in src, _megdnn_tensor_in trans, _megdnn_tensor_in dst,
float border_value, BorderMode bmode, InterpolationMode imode, Handle* handle) {
size_t ch = dst.layout[3];
size_t width = dst.layout[2];
size_t height = dst.layout[1];
const size_t batch = dst.layout.shape[0];
size_t BLOCK_SZ_H = std::min(BLOCK_SZ / 2, height);
size_t BLOCK_SZ_W = std::min(BLOCK_SZ * BLOCK_SZ / BLOCK_SZ_H, width);
BLOCK_SZ_H = std::min(BLOCK_SZ * BLOCK_SZ / BLOCK_SZ_W, height);
size_t parallelism_batch =
div_ceil<size_t>(height, BLOCK_SZ_H) * div_ceil<size_t>(width, BLOCK_SZ_W);
megdnn_assert(
ch == 1 || ch == 3 || ch == 2,
"unsupported src channel: %zu, avaiable channel size: 1/2/3", ch);
if (dst.layout.dtype.enumv() == DTypeEnum::Float32) {
#define cb(_imode, _bmode, _ch) \
auto task = [src, trans, dst, border_value, parallelism_batch]( \
size_t index, size_t) { \
size_t batch_id = index / parallelism_batch; \
size_t task_id = index % parallelism_batch; \
Mat<float> src_mat = TensorND2Mat<float>(src, batch_id); \
Mat<float> dst_mat = TensorND2Mat<float>(dst, batch_id); \
auto task_trans_ptr = trans.ptr<float>() + batch_id * 2 * 3; \
warp_affine_cv< \
float MEGDNN_COMMA _imode MEGDNN_COMMA _bmode MEGDNN_COMMA _ch>( \
src_mat MEGDNN_COMMA const_cast<Mat<float>&>(dst_mat) \
MEGDNN_COMMA task_trans_ptr MEGDNN_COMMA border_value, \
task_id); \
}; \
MEGDNN_DISPATCH_MULTI_THREAD_CPU_KERN( \
static_cast<naive::HandleImpl*>(handle), batch* parallelism_batch, task);
DISPATCH_IMODE(imode, bmode, ch, cb)
} else if (dst.layout.dtype.enumv() == DTypeEnum::Uint8) {
#undef cb
#define cb(_imode, _bmode, _ch) \
auto task = [src, trans, dst, border_value, parallelism_batch]( \
size_t index, size_t) { \
size_t batch_id = index / parallelism_batch; \
size_t task_id = index % parallelism_batch; \
Mat<uchar> src_mat = TensorND2Mat<uchar>(src, batch_id); \
Mat<uchar> dst_mat = TensorND2Mat<uchar>(dst, batch_id); \
auto task_trans_ptr = trans.ptr<float>() + batch_id * 2 * 3; \
warp_affine_cv< \
uchar MEGDNN_COMMA _imode MEGDNN_COMMA _bmode MEGDNN_COMMA _ch>( \
src_mat MEGDNN_COMMA const_cast<Mat<uchar>&>(dst_mat) \
MEGDNN_COMMA task_trans_ptr MEGDNN_COMMA border_value, \
task_id); \
}; \
MEGDNN_DISPATCH_MULTI_THREAD_CPU_KERN( \
static_cast<naive::HandleImpl*>(handle), batch* parallelism_batch, task);
DISPATCH_IMODE(imode, bmode, ch, cb)
#undef cb
} else {
megdnn_throw("Unsupported datatype of WarpAffine optr.");
}
}