asgardius/virtualx-engine
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions
virtualx-engine/thirdparty/oidn/mkl-dnn/src/cpu/simple_reorder.hpp
Juan Linietsky 1bea8e1eac New lightmapper 
-Added LocalVector (needed it)
-Added stb_rect_pack (It's pretty cool, we could probably use it for other stuff too)
-Fixes and changes all around the place
-Added library for 128 bits fixed point (required for Delaunay3D)
2020-05-10 15:59:09 -03:00

1022 lines
40 KiB
C++
Raw Blame History

/*******************************************************************************
* Copyright 2016-2018 Intel Corporation
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*******************************************************************************/
#ifndef CPU_SIMPLE_REORDER_HPP
#define CPU_SIMPLE_REORDER_HPP
#include <assert.h>
#include "c_types_map.hpp"
#include "type_helpers.hpp"
#include "math_utils.hpp"
#include "mkldnn_thread.hpp"
#include "utils.hpp"
#include "tag_traits.hpp"
#include "cpu_reorder_pd.hpp"
#include "cpu_primitive.hpp"
#include "simple_q10n.hpp"
#include "cpu_isa_traits.hpp"
namespace mkldnn {
namespace impl {
namespace cpu {
using namespace mkldnn::impl::status;
using namespace mkldnn::impl::format_tag;
using namespace mkldnn::impl::data_type;
using bd = block_dim_t;
using ib = inner_blk_t;
using namespace mkldnn::impl::utils;
using math::saturate;
template<impl::data_type_t type>
using data_t = typename prec_traits<type>::type;
template<impl::data_type_t type_i, impl::data_type_t type_o>
using _qz_a1b0 = qz_a1b0<data_t<type_i>, data_t<type_o>>;
template<impl::data_type_t type_i, impl::data_type_t type_o>
using _qz = qz<data_t<type_i>, data_t<type_o>>;
namespace fmt_order {
const bool keep = true;
const bool reverse = false;
const bool any = keep;
}
namespace spec {
struct direct_copy {};
struct direct_copy_except_dim_0 {};
struct reference {};
struct conv_s8s8 {};
}
#define SIMPLE_REORDER_TEMPL_DECL \
impl::data_type_t type_i, impl::format_tag_t tag_i, \
impl::data_type_t type_o, impl::format_tag_t tag_o, bool order_keep
#define SIMPLE_REORDER_TEMPL_CALL \
type_i, tag_i, type_o, tag_o, order_keep
#define DECLARE_COMMON_PARAMS() \
const memory_desc_wrapper &input_d = pd->src_md(); \
const memory_desc_wrapper &output_d = pd->dst_md(); \
const float alpha = pd->alpha(); MAYBE_UNUSED(alpha); \
const float beta = pd->beta(); MAYBE_UNUSED(beta);
/* specific reorders: common template */
template <SIMPLE_REORDER_TEMPL_DECL, typename spec = void>
struct simple_reorder_impl {};
namespace {
inline bool simple_fmt_check(bool order_keep, impl::format_tag_t tag_i,
impl::format_tag_t tag_o, const memory_desc_wrapper &input_d,
const memory_desc_wrapper &output_d) {
return input_d.matches_tag(order_keep ? tag_i : tag_o)
&& output_d.matches_tag(order_keep ? tag_o : tag_i);
}
inline bool simple_attr_check(const primitive_attr_t *attr, bool many_scales_support) {
if (many_scales_support)
return true;
return IMPLICATION(attr, attr->output_scales_.mask_ == 0);
}
}
/* specific reorders: implementation */
template <SIMPLE_REORDER_TEMPL_DECL>
struct simple_reorder_impl<SIMPLE_REORDER_TEMPL_CALL,
typename utils::enable_if<tag_i == any && (false
|| tag_o == hwio
|| tag_o == hwigo)
, spec::conv_s8s8>::type>
{
static bool is_applicable(const memory_desc_wrapper &input_d,
const memory_desc_wrapper &output_d, const primitive_attr_t *attr)
{
const size_t D_mask = utils::array_product(input_d.dims(),
math::ilog2q(attr->output_scales_.mask_ + 1));
const int oc = (input_d.dims()[tag_o == hwigo + 0]);
const int g = (tag_o == hwigo) ? (input_d.dims()[0]) : 1;
return output_d.matches_tag(tag_o)
&& (output_d.extra().flags & memory_extra_flags::compensation_conv_s8s8)
&& (input_d.data_type() == f32 || input_d.data_type() == s8)
&& output_d.data_type() == s8
&& (D_mask == 1 || D_mask == (size_t)g * oc);
}
static status_t execute(const cpu_reorder_pd_t *pd,
const data_t<type_i> *input, data_t<type_o> *output) {
DECLARE_COMMON_PARAMS();
static constexpr bool w_groups = tag_o == hwigo;
const auto &dims = input_d.dims();
const auto &pdims = output_d.padded_dims();
const int G = w_groups ? dims[0] : 1;
const int OC = dims[w_groups + 0];
const int IC = dims[w_groups + 1];
const int H = dims[w_groups + 2];
const int W = dims[w_groups + 3];
const float *scales = pd->attr()->output_scales_.scales_;
const size_t D_mask = utils::array_product(input_d.dims(),
math::ilog2q(pd->attr()->output_scales_.mask_ + 1));
assert(output_d.extra().flags
& memory_extra_flags::compensation_conv_s8s8);
float adj_scale =
(output_d.extra().flags & memory_extra_flags::scale_adjust)
? output_d.extra().scale_adjust : 1.f;
size_t offset = G * pdims[w_groups + 0] * pdims[w_groups + 1] * H * W;
int32_t *cp = reinterpret_cast<int32_t *>(output + offset);
parallel_nd(G, OC, [&](int g, int oc) {
cp[g * OC + oc] = 0;
for (int ic = 0; ic < IC; ic++)
for (int h = 0; h < H; h++)
for (int w = 0; w < W; w++) {
auto i = input[input_d.blk_off<!w_groups>(g, oc, ic, h, w)];
auto &o = output[output_d.blk_off<!w_groups>(g, oc, ic, h, w)];
const float s = scales[(D_mask == 1) ? 0 : g * OC + oc];
o = qz_b0<data_t<type_i>, data_t<type_o>>()(
i, s * adj_scale);
cp[g * OC + oc] -= (int32_t)o;
}
cp [g * OC + oc] *= 128;
});
return success;
}
};
template <SIMPLE_REORDER_TEMPL_DECL>
struct simple_reorder_impl<SIMPLE_REORDER_TEMPL_CALL,
typename utils::enable_if<
(tag_i == goiw && tag_o == gOIw4i16o4i)
|| (tag_i == oiw && tag_o == OIw4i16o4i)
|| (tag_i == goihw && tag_o == gOIhw4i16o4i)
|| (tag_i == oihw && tag_o == OIhw4i16o4i)
|| (tag_i == goihw && tag_o == gOIhw2i8o4i)
|| (tag_i == goihw && tag_o == gOIhw4o4i)
, spec::conv_s8s8>::type>
{
static bool is_applicable(const memory_desc_wrapper &input_d,
const memory_desc_wrapper &output_d, const primitive_attr_t *attr)
{
const size_t D_mask = utils::array_product(input_d.dims(),
math::ilog2q(attr->output_scales_.mask_ + 1));
const bool w_groups = !utils::one_of(tag_o, OIw4i16o4i, OIhw4i16o4i);
const int oc = (input_d.dims()[w_groups ? 1 : 0]);
const int g = w_groups ? input_d.dims()[0] : 1;
return input_d.matches_tag(tag_i)
&& output_d.matches_tag(tag_o)
&& (output_d.extra().flags & memory_extra_flags::compensation_conv_s8s8)
&& (input_d.data_type() == f32 || input_d.data_type() == s8)
&& output_d.data_type() == s8
&& (D_mask == 1 || D_mask == (size_t)g * oc);
}
static status_t execute(const cpu_reorder_pd_t *pd,
const data_t<type_i> *input, data_t<type_o> *output) {
DECLARE_COMMON_PARAMS();
static constexpr bool w_groups =
!utils::one_of(tag_o, OIw4i16o4i, OIhw4i16o4i);
constexpr int is_1d =
utils::one_of(tag_o, gOIw4i16o4i, OIw4i16o4i);
constexpr int blksize = tag_traits<tag_o>::inner_blks == ib::_4b4c
? 4
: tag_traits<tag_o>::inner_blks == ib::_2c8b4c
? 8
: 16;
const auto &_g_oihw_d = order_keep ? input_d : output_d;
const auto &dims = input_d.dims();
const auto &pdims = order_keep
? output_d.padded_dims()
: input_d.padded_dims();
const int G = w_groups ? dims[0] : 1;
const int OC = dims[w_groups + 0];
const int NB_OC = pdims[w_groups + 0] / blksize;
const int IC = dims[w_groups + 1];
const int NB_IC = pdims[w_groups + 1] / blksize;
const int H = is_1d ? 1 : dims[w_groups + 2];
const int W = dims[w_groups + 3 - is_1d];
const float *scales = pd->attr()->output_scales_.scales_;
const size_t D_mask = utils::array_product(input_d.dims(),
math::ilog2q(pd->attr()->output_scales_.mask_ + 1));
assert(output_d.extra().flags
& memory_extra_flags::compensation_conv_s8s8);
float adj_scale =
(output_d.extra().flags & memory_extra_flags::scale_adjust)
? output_d.extra().scale_adjust : 1.f;
auto ker = [&](const data_t<type_i> *inp, data_t<type_o> *out,
int32_t *c, const float *s, const int oc_block, const int ic_block) {
# define index AB_or_BC_blk_off<tag_traits<tag_o>::inner_blks>
for (int ic = 0; ic < ic_block; ++ic) {
for (int oc = 0; oc < oc_block; ++oc) {
const auto _g_oihw_off =
oc * _g_oihw_d.blocking_desc().strides[w_groups + 0]
+ ic * _g_oihw_d.blocking_desc().strides[w_groups + 1];
out[index(oc, ic)]
= qz_b0<data_t<type_i>, data_t<type_o>>()(
inp[_g_oihw_off], s[oc] * adj_scale);
c[oc] -= (128 * (int32_t)(out[index(oc, ic)]));
}
}
# undef index
};
constexpr int i_mult = blksize;
constexpr int o_mult = 1;
size_t offset = G * pdims[w_groups+0] * pdims[w_groups+1] * H * W;
int32_t *cp = reinterpret_cast<int32_t *>(output + offset);
parallel_nd(G * NB_OC * blksize, [&](int i) {
cp[i] = 0;
});
# define wei_blk_off(md, g, o, i, h, w) \
(is_1d ? (md).blk_off<!w_groups>(g, o, i, w) \
: (md).blk_off<!w_groups>(g, o, i, h, w))
parallel_nd(G, NB_OC, [&](int g, int O) {
for (int I = 0; I < NB_IC; I++)
for (int h = 0; h < H; h++)
for (int w = 0; w < W; w++) {
auto i = &input[wei_blk_off(
input_d, g, i_mult * O, i_mult * I, h, w)];
auto o = &output[wei_blk_off(
output_d, g, o_mult * O, o_mult * I, h, w)];
const int oc_block = nstl::min(blksize, OC - O * blksize);
const int ic_block = nstl::min(blksize, IC - I * blksize);
int _offset = (g * NB_OC + O) * blksize;
ker(i, o, (order_keep) ? &cp[_offset] : nullptr,
&scales[(D_mask == 1) ? 0 : _offset],
oc_block, ic_block);
}
});
# undef wei_blk_off
return success;
}
};
template <SIMPLE_REORDER_TEMPL_DECL>
struct simple_reorder_impl<SIMPLE_REORDER_TEMPL_CALL,
typename utils::enable_if<false
||(tag_i == goiw && tag_o == Goiw16g)
||(tag_i == goihw && tag_o == Goihw16g)
, spec::conv_s8s8>::type>
{
static bool is_applicable(const memory_desc_wrapper &input_d,
const memory_desc_wrapper &output_d, const primitive_attr_t *attr) {
const size_t D_mask = utils::array_product(input_d.dims(),
math::ilog2q(attr->output_scales_.mask_ + 1));
const int oc = input_d.dims()[1];
const int g = input_d.dims()[0];
return true
&& order_keep
&& input_d.matches_tag(tag_i)
&& output_d.matches_tag(tag_o)
&& (output_d.extra().flags & memory_extra_flags::compensation_conv_s8s8)
&& (input_d.data_type() == f32 || input_d.data_type() == s8)
&& output_d.data_type() == s8
&& (D_mask == 1 || D_mask == (size_t)g * oc);
}
static status_t execute(const cpu_reorder_pd_t *pd,
const data_t<type_i> *input, data_t<type_o> *output) {
DECLARE_COMMON_PARAMS();
constexpr bool is_1d = tag_i == goiw;
constexpr int blksize = 16;
const auto &dims = input_d.dims();
const auto &pdims = output_d.padded_dims();
const int G = dims[0];
const int Gp = pdims[0];
const int OC = dims[1];
const int IC = dims[2];
const int H = is_1d ? 1 : dims[3];
const int W = dims[4 - is_1d];
const size_t D_mask = utils::array_product(input_d.dims(),
math::ilog2q(pd->attr()->output_scales_.mask_ + 1));
const float *scales = pd->attr()->output_scales_.scales_;
assert(output_d.extra().flags
& memory_extra_flags::compensation_conv_s8s8);
float adj_scale =
(output_d.extra().flags & memory_extra_flags::scale_adjust)
? output_d.extra().scale_adjust : 1.f;
auto ker = [&](const data_t<type_i> *inp, data_t<type_o> *out,
int32_t *cp, const float *s, const int g_block) {
PRAGMA_OMP_SIMD()
for (int g = 0; g < g_block; g++) {
const auto i_off = g * input_d.blocking_desc().strides[0];
out[g] = qz_b0<data_t<type_i>, data_t<type_o>>()(
inp[i_off], s[g * OC] * adj_scale);
cp[g * OC] -= 128 * (int32_t)(out[g]);
}
};
size_t cp_offset = output_d.size() - output_d.additional_buffer_size();
int32_t *cp = reinterpret_cast<int32_t *>(output + cp_offset);
parallel_nd((Gp/blksize) * OC, [&](int ib) {
PRAGMA_OMP_SIMD()
for (int i = 0; i < blksize; i++)
cp[ib * blksize + i] = 0;
});
# define wei_blk_off(md, g, o, i, h, w) \
(is_1d ? (md).blk_off(g, o, i, w) : (md).blk_off(g, o, i, h, w))
parallel_nd(Gp/blksize, OC, [&](int gb, int O) {
for (int I = 0; I < IC; I++) {
for (int h = 0; h < H; h++)
for (int w = 0; w < W; w++)
{
const int g_block = nstl::min(G - gb * blksize, blksize);
const auto inp = &input[wei_blk_off(
input_d, gb * blksize, O, I, h, w)];
const auto out = &output[wei_blk_off(
output_d, gb, O, I, h, w)];
int offset = gb * blksize + O;
ker(inp, out, &cp[offset],
&scales[(D_mask == 1) ? 0 : offset], g_block);
}
}
});
# undef wei_blk_off
return success;
}
};
/* reorders with tail support */
template <SIMPLE_REORDER_TEMPL_DECL>
struct simple_reorder_impl<SIMPLE_REORDER_TEMPL_CALL,
typename utils::enable_if<false
|| (tag_i == nCdhw8c && tag_o == nCdhw16c)
|| (tag_i == nChw8c && tag_o == nChw16c)
|| (tag_i == nCw8c && tag_o == nCw16c)
>::type>
{
static bool is_applicable(const memory_desc_wrapper &input_d,
const memory_desc_wrapper &output_d, const primitive_attr_t *attr)
{
return simple_fmt_check(order_keep, tag_i, tag_o, input_d, output_d)
&& simple_attr_check(attr, false);
}
static status_t execute(const cpu_reorder_pd_t *pd,
const data_t<type_i> *input, data_t<type_o> *output) {
DECLARE_COMMON_PARAMS();
constexpr int is_1d = tag_i == nCw8c;
constexpr int is_3d = tag_i == nCdhw8c;
constexpr int blksize_16 = 16;
constexpr int blksize_8 = 8;
constexpr int ic_mult = order_keep ? 2 : 1;
constexpr int oc_mult = order_keep ? 1 : 2;
const auto &dims = input_d.dims();
const auto &pdims = order_keep ? output_d.padded_dims()
: input_d.padded_dims();
const int C = dims[1];
const int D = is_3d ? dims[2] : 1;
const int H = is_1d ? 1 : dims[2 + is_3d];
const int W = dims[3 + is_3d - is_1d];
auto ker = [&](const data_t<type_i> *i, data_t<type_o> *o,
const int block_16) {
const int nb = (block_16 - 1) / blksize_8 + 1;
if (alpha == 1.0 && beta == 0.0) {
for (int b = 0; b < nb; ++b) {
const ptrdiff_t i_off = order_keep ? b : b * blksize_8;
const ptrdiff_t o_off = order_keep ? b * blksize_8 : b;
const int block_8 = nstl::min(blksize_8,
block_16 - b * blksize_8);
for (int c = 0; c < block_8; ++c) {
o[o_off + c] = _qz_a1b0<type_i, type_o>()(
i[i_off + c]);
}
}
} else {
for (int b = 0; b < nb; ++b) {
const ptrdiff_t i_off = order_keep ? b : b * blksize_8;
const ptrdiff_t o_off = order_keep ? b * blksize_8 : b;
const int block_8 = nstl::min(blksize_8,
block_16 - b * blksize_8);
for (int c = 0; c < block_8; ++c) {
o[o_off + c] = _qz<type_i, type_o>()(i[i_off + c],
o[o_off + c], alpha, beta);
}
}
}
};
# define data_blk_off(md, n, c, d, h, w) \
( is_1d ? (md).blk_off(n, c, w) \
: is_3d ? (md).blk_off(n, c, d, h, w) : (md).blk_off(n, c, h, w))
parallel_nd(dims[0], pdims[1] / blksize_16, D, H, W,
[&](int n, int nb_c, int d, int h, int w) {
auto i = &input[data_blk_off(input_d, n, ic_mult * nb_c, d, h, w)];
auto o = &output[data_blk_off(output_d, n, oc_mult * nb_c, d, h, w)];
const int block_16 = nstl::min(blksize_16, C - nb_c * blksize_16);
ker(i, o, block_16);
});
# undef data_blk_off
return success;
}
};
#define PLAIN_TO_BLOCKED_IS_APPLICABLE() \
static bool is_applicable(const memory_desc_wrapper &input_d, \
const memory_desc_wrapper &output_d, const primitive_attr_t *attr) { \
return simple_attr_check(attr, false) && (order_keep \
? output_d.matches_tag(tag_o) && input_d.is_plain() \
: input_d.matches_tag(tag_o) && output_d.is_plain()); \
}
template <SIMPLE_REORDER_TEMPL_DECL>
struct simple_reorder_impl<SIMPLE_REORDER_TEMPL_CALL,
typename utils::enable_if<tag_i == any
&& (tag_traits<tag_o>::block_dims == bd::_A
|| tag_traits<tag_o>::block_dims == bd::_B)
&& tag_traits<tag_o>::ndims >= 3
&& tag_traits<tag_o>::ndims <= 6
>::type>
{
PLAIN_TO_BLOCKED_IS_APPLICABLE();
static status_t execute(const cpu_reorder_pd_t *pd,
const data_t<type_i> *input, data_t<type_o> *output) {
DECLARE_COMMON_PARAMS();
const auto &flat_d = order_keep ? input_d : output_d;
const auto &block_d = order_keep ? output_d : input_d;
const auto &dims = input_d.dims();
const auto &pdims = block_d.padded_dims();
constexpr int ndims = tag_traits<tag_o>::ndims;
constexpr int blk_idx = tag_traits<tag_o>::block_dims == bd::_A ? 0 : 1;
const dim_t H0 = dims[0];
const dim_t H1 = dims[1];
const dim_t M0 = ndims >= 6 ? dims[ndims - 4] : 1;
const dim_t M1 = ndims >= 5 ? dims[ndims - 3] : 1;
const dim_t M2 = ndims >= 4 ? dims[ndims - 2] : 1;
const dim_t L = dims[ndims - 1];
const dim_t l_blk_stride = block_d.blocking_desc().strides[ndims - 1];
constexpr int blksize = false ? 0
: utils::one_of(tag_traits<tag_o>::inner_blks, ib::_4a, ib::_4b) ? 4
: utils::one_of(tag_traits<tag_o>::inner_blks, ib::_8a, ib::_8b) ? 8
: 16;
auto ker = [&](const data_t<type_i> *i, data_t<type_o> *o, int block) {
if (alpha == 1.0 && beta == 0.0) {
for (int l = 0; l < L; ++l)
for (int blk = 0; blk < block; ++blk) {
const dim_t flat_off = 0
+ blk * flat_d.blocking_desc().strides[blk_idx]
+ l * flat_d.blocking_desc().strides[ndims - 1];
if (order_keep) {
o[l * l_blk_stride + blk] = _qz_a1b0<type_i, type_o>()(
i[flat_off]);
} else {
o[flat_off] = _qz_a1b0<type_i, type_o>()(
i[l * l_blk_stride + blk]);
}
}
} else {
for (int l = 0; l < L; ++l)
for (int blk = 0; blk < block; ++blk) {
const dim_t flat_off = 0
+ blk * flat_d.blocking_desc().strides[blk_idx]
+ l * flat_d.blocking_desc().strides[ndims - 1];
if (order_keep) {
o[l * l_blk_stride + blk] = _qz<type_i, type_o>()(
i[flat_off], o[l * blksize + blk],
alpha, beta);
} else {
o[flat_off] = _qz<type_i, type_o>()(
i[l * l_blk_stride + blk], o[flat_off],
alpha, beta);
}
}
}
};
# define off(md, h0, h1, m0, m1, m2) \
(ndims >= 6 ? (md).blk_off(h0, h1, m0, m1, m2) \
: ndims >= 5 ? (md).blk_off(h0, h1, m1, m2) \
: ndims >= 4 ? (md).blk_off(h0, h1, m2) \
: /* ndims >= 3 ? */ (md).blk_off(h0, h1))
constexpr int i_mult = order_keep ? blksize : 1;
constexpr int o_mult = order_keep ? 1 : blksize;
if (blk_idx == 0) {
const dim_t BH0 = pdims[0] / blksize;
parallel_nd(BH0, H1, M0, M1, M2,
[&](dim_t bh0, dim_t h1, dim_t m0, dim_t m1, dim_t m2) {
auto i = &input[off(input_d, bh0 * i_mult, h1, m0, m1, m2)];
auto o = &output[off(output_d, bh0 * o_mult, h1, m0, m1, m2)];
const int block = nstl::min<int>(blksize, H0 - bh0 * blksize);
ker(i, o, block);
});
} else if (blk_idx == 1) {
const dim_t BH1 = pdims[1] / blksize;
parallel_nd(H0, BH1, M0, M1, M2,
[&](dim_t h0, dim_t bh1, dim_t m0, dim_t m1, dim_t m2) {
auto i = &input[off(input_d, h0, bh1 * i_mult, m0, m1, m2)];
auto o = &output[off(output_d, h0, bh1 * o_mult, m0, m1, m2)];
const int block = nstl::min<int>(blksize, H1 - bh1 * blksize);
ker(i, o, block);
});
} else {
assert(!"unimplemented");
}
# undef off
return success;
}
};
template <SIMPLE_REORDER_TEMPL_DECL>
struct simple_reorder_impl<SIMPLE_REORDER_TEMPL_CALL,
typename utils::enable_if<tag_i == any
&& (tag_traits<tag_o>::block_dims == bd::_AB
|| tag_traits<tag_o>::block_dims == bd::_BC)
&& IMPLICATION(tag_traits<tag_o>::block_dims == bd::_AB,
tag_traits<tag_o>::ndims >= 3 && tag_traits<tag_o>::ndims <= 5)
&& IMPLICATION(tag_traits<tag_o>::block_dims == bd::_BC,
tag_traits<tag_o>::ndims >= 4 && tag_traits<tag_o>::ndims <= 6)
>::type>
{
PLAIN_TO_BLOCKED_IS_APPLICABLE();
static status_t execute(const cpu_reorder_pd_t *pd,
const data_t<type_i> *input, data_t<type_o> *output) {
DECLARE_COMMON_PARAMS();
const auto &flat_d = order_keep ? input_d : output_d;
const auto &dims = input_d.dims();
const auto &pdims = order_keep
? output_d.padded_dims()
: input_d.padded_dims();
constexpr int ndims = tag_traits<tag_o>::ndims;
static constexpr bool with_g = tag_traits<tag_o>::block_dims == bd::_BC;
const dim_t G = with_g ? dims[0] : 1;
const dim_t H0 = dims[0 + with_g];
const dim_t H1 = dims[1 + with_g];
const dim_t M0 = ndims >= 5 + with_g ? dims[ndims - 3] : 1;
const dim_t M1 = ndims >= 4 + with_g ? dims[ndims - 2] : 1;
const dim_t M2 = ndims >= 3 + with_g ? dims[ndims - 1] : 1;
constexpr int blksize_0 = false ? 0
: utils::one_of(tag_traits<tag_o>::inner_blks,
ib::_4b4a, ib::_4b4c, ib::_4c4b)
? 4
: utils::one_of(tag_traits<tag_o>::inner_blks,
ib::_8a8b, ib::_8b8a, ib::_8b8c, ib::_8c8b, ib::_2c8b4c)
? 8
: utils::one_of(tag_traits<tag_o>::inner_blks,
ib::_16a16b, ib::_16a4b, ib::_16b16a, ib::_16b4c,
ib::_16b16c, ib::_16c16b, ib::_8a16b2a, ib::_4b16a4b,
ib::_8b16a2b, ib::_8b16c2b, ib::_4c16b4c, ib::_8c16b2c)
? 16 : INT_MIN;
constexpr int blksize_1 = utils::one_of(tag_traits<tag_o>::inner_blks,
ib::_8a8b, ib::_8b8a, ib::_8b8c, ib::_8c8b, ib::_2c8b4c)
? 8
: utils::one_of(tag_traits<tag_o>::inner_blks,
ib::_16a16b, ib::_16b16a, ib::_16b16c, ib::_16c16b,
ib::_8a16b2a, ib::_4b16a4b, ib::_8b16a2b, ib::_8b16c2b,
ib::_4c16b4c, ib::_8c16b2c)
? 16
: utils::one_of(tag_traits<tag_o>::inner_blks,
ib::_4b4a, ib::_4b4c, ib::_4c4b,
ib::_16a4b, ib::_16b4c)
? 4
: INT_MIN;
const dim_t NB_H0 = pdims[0 + with_g] / blksize_0;
const dim_t NB_H1 = pdims[1 + with_g] / blksize_1;
auto ker = [&](const data_t<type_i> *i, data_t<type_o> *o,
const int block_h0, const int block_h1) {
# define blk_off AB_or_BC_blk_off<tag_traits<tag_o>::inner_blks>
if (alpha == 1.0 && beta == 0.0) {
for (int h0 = 0; h0 < block_h0; ++h0)
for (int h1 = 0; h1 < block_h1; ++h1) {
const dim_t flat_off = 0
+ h0 * flat_d.blocking_desc().strides[with_g + 0]
+ h1 * flat_d.blocking_desc().strides[with_g + 1];
if (order_keep) {
o[blk_off(h0, h1)] = _qz_a1b0<type_i, type_o>()(
i[flat_off]);
} else {
o[flat_off] = _qz_a1b0<type_i, type_o>()(
i[blk_off(h0, h1)]);
}
}
} else {
for (int h0 = 0; h0 < block_h0; ++h0)
for (int h1 = 0; h1 < block_h1; ++h1) {
const dim_t flat_off = 0
+ h0 * flat_d.blocking_desc().strides[with_g + 0]
+ h1 * flat_d.blocking_desc().strides[with_g + 1];
if (order_keep) {
o[blk_off(h0, h1)] = _qz<type_i, type_o>()(i[flat_off],
o[blk_off(h0, h1)], alpha, beta);
} else {
o[flat_off] = _qz<type_i, type_o>()(i[blk_off(h0, h1)],
o[flat_off], alpha, beta);
}
}
}
# undef blk_off
};
constexpr int i_mult_0 = order_keep ? blksize_0 : 1;
constexpr int o_mult_0 = order_keep ? 1 : blksize_0;
constexpr int i_mult_1 = order_keep ? blksize_1 : 1;
constexpr int o_mult_1 = order_keep ? 1 : blksize_1;
# define off(md, g, h0, h1, m0, m1, m2) \
(ndims >= 5 + with_g ? (md).blk_off<!with_g>(g, h0, h1, m0, m1, m2) \
: ndims >= 4 + with_g ? (md).blk_off<!with_g>(g, h0, h1, m1, m2) \
: /* ndims >= 3 + with_g ? */ (md).blk_off<!with_g>(g, h0, h1, m2))
parallel_nd(G, NB_H0, NB_H1, M0, M1, M2,
[&](dim_t g, dim_t nb_h0, dim_t nb_h1, dim_t m0, dim_t m1, dim_t m2) {
auto i = &input[off(input_d,
g, i_mult_0 * nb_h0, i_mult_1 * nb_h1, m0, m1, m2)];
auto o = &output[off(output_d,
g, o_mult_0 * nb_h0, o_mult_1 * nb_h1, m0, m1, m2)];
const int block_h0 = nstl::min<int>(blksize_0, H0 - nb_h0 * blksize_0);
const int block_h1 = nstl::min<int>(blksize_1, H1 - nb_h1 * blksize_1);
ker(i, o, block_h0, block_h1);
});
# undef off
return success;
}
};
/* generic and direct-copy reorders */
template <SIMPLE_REORDER_TEMPL_DECL>
struct simple_reorder_impl<SIMPLE_REORDER_TEMPL_CALL,
typename utils::enable_if<
tag_i == any && tag_o == any && order_keep == fmt_order::any,
spec::direct_copy>::type>
{
static bool is_applicable(const memory_desc_wrapper &input_d,
const memory_desc_wrapper &output_d, const primitive_attr_t *attr) {
/* FIXME: is the formula correct? */
return input_d.similar_to(output_d, true, false, 0)
&& input_d.is_dense() && output_d.is_dense()
&& simple_attr_check(attr, false);
}
static status_t execute(const cpu_reorder_pd_t *pd,
const data_t<type_i> *input, data_t<type_o> *output) {
DECLARE_COMMON_PARAMS();
assert(input_d.is_dense());
input += input_d.blk_off(0);
output += output_d.blk_off(0);
const size_t nelems = input_d.nelems();
constexpr int block_size = 16;
const auto num_blocks = nelems / block_size;
const auto rem_elems = nelems % block_size;
parallel(0, [&](const int ithr, const int nthr) {
size_t start{0}, end{0};
balance211(num_blocks, nthr, ithr, start, end);
start = start * block_size;
end = end * block_size;
if (alpha == 1.0 && beta == 0.0) {
PRAGMA_OMP_SIMD()
for (size_t e = start; e < end; ++e) {
output[e] = qz_a1b0<data_t<type_i>, data_t<type_o>>()
(input[e]);
}
} else if (alpha == 1.0) {
PRAGMA_OMP_SIMD()
for (size_t e = start; e < end; ++e) {
output[e] = qz_a1<data_t<type_i>, data_t<type_o>>()
(input[e], output[e], beta);
}
} else if (beta == 0.0) {
PRAGMA_OMP_SIMD()
for (size_t e = start; e < end; ++e) {
output[e] = qz_b0<data_t<type_i>, data_t<type_o>>()
(input[e], alpha);
}
} else {
PRAGMA_OMP_SIMD()
for (size_t e = start; e < end; ++e) {
output[e] = qz<data_t<type_i>, data_t<type_o>>()
(input[e], output[e], alpha, beta);
}
}
if (rem_elems != 0 && ithr == nthr - 1){
if (alpha == 1.0 && beta == 0.0) {
PRAGMA_OMP_SIMD()
for (size_t e = nelems - rem_elems; e < nelems; ++e) {
output[e] = qz_a1b0<data_t<type_i>,
data_t<type_o>>()(input[e]);
}
} else if (alpha == 1.0) {
PRAGMA_OMP_SIMD()
for (size_t e = nelems - rem_elems; e < nelems; ++e) {
output[e] = qz_a1<data_t<type_i>,
data_t<type_o>>()(input[e], output[e], beta);
}
} else if (beta == 0.0) {
PRAGMA_OMP_SIMD()
for (size_t e = nelems - rem_elems; e < nelems; ++e) {
output[e] = qz_b0<data_t<type_i>,
data_t<type_o>>()(input[e], alpha);
}
} else {
PRAGMA_OMP_SIMD()
for (size_t e = nelems - rem_elems; e < nelems; ++e) {
output[e] = qz<data_t<type_i>, data_t<type_o>>()
(input[e], output[e], alpha, beta);
}
}
}
});
return success;
}
};
template <SIMPLE_REORDER_TEMPL_DECL>
struct simple_reorder_impl<SIMPLE_REORDER_TEMPL_CALL,
typename utils::enable_if<
tag_i == any && tag_o == any && order_keep == fmt_order::any,
spec::direct_copy_except_dim_0>::type>
{
static bool is_applicable(const memory_desc_wrapper &input_d,
const memory_desc_wrapper &output_d, const primitive_attr_t *attr) {
auto is_dense_no_0 = [](const memory_desc_wrapper &data_d) {
return nelems_no_dim_0(data_d) == _size_no_dim_0(data_d);
};
/* FIXME: is the formula correct? */
return input_d.similar_to(output_d, true, false, 1)
&& is_dense_no_0(input_d) && is_dense_no_0(output_d)
&& simple_attr_check(attr, false);
}
static status_t execute(const cpu_reorder_pd_t *pd,
const data_t<type_i> *input, data_t<type_o> *output) {
DECLARE_COMMON_PARAMS();
input += input_d.blk_off(0);
output += output_d.blk_off(0);
const int N = input_d.dims()[0];
const dim_t is = input_d.blocking_desc().strides[0];
const dim_t os = output_d.blocking_desc().strides[0];
const dim_t nelems_no_d0 = nelems_no_dim_0(input_d);
const dim_t work_amount = N * nelems_no_d0;
if (alpha == 1.0 && beta == 0.0) {
parallel(0, [&](const int ithr, const int nthr) {
dim_t n{0}, dim1_s{0};
dim_t start{0}, end{0};
balance211(work_amount, nthr, ithr, start, end);
nd_iterator_init(start, n, N, dim1_s, nelems_no_d0);
while(start < end) {
dim_t work_rem = end - start;
dim_t dim1_e = dim1_s + work_rem > nelems_no_d0
? nelems_no_d0 : dim1_s + work_rem;
PRAGMA_OMP_SIMD()
for (dim_t e = dim1_s; e < dim1_e; ++e) {
output[os * n + e] = _qz_a1b0<type_i, type_o>()(
input[is * n + e]);
}
nd_iterator_jump(start, end, n, N, dim1_s, nelems_no_d0);
}
});
} else {
parallel(0, [&](const int ithr, const int nthr) {
dim_t n{0}, dim1_s{0};
dim_t start{0}, end{0};
balance211(work_amount, nthr, ithr, start, end);
nd_iterator_init(start, n, N, dim1_s, nelems_no_d0);
while(start < end) {
dim_t work_rem = end - start;
dim_t dim1_e =
dim1_s + work_rem > nelems_no_d0 ? nelems_no_d0
: dim1_s + work_rem;
PRAGMA_OMP_SIMD()
for (dim_t e = dim1_s; e < dim1_e; ++e){
output[os * n + e] = _qz<type_i, type_o>()(
input[is * n + e], output[os * n + e], alpha,
beta);
}
nd_iterator_jump(start, end, n, N, dim1_s, nelems_no_d0);
}
});
}
return success;
}
private:
static dim_t nelems_no_dim_0(const memory_desc_wrapper &data_d) {
const int ndims = data_d.ndims();
if (ndims <= 1) return 1;
return utils::array_product(data_d.dims() + 1, data_d.ndims() - 1);
}
static dim_t _size_no_dim_0(const memory_desc_wrapper &data_d) {
dims_t blocks;
data_d.compute_blocks(blocks);
const auto &blk = data_d.blocking_desc();
dim_t blk_size = 1;
for (int iblk = 0; iblk < blk.inner_nblks; ++iblk)
blk_size *= blk.inner_blks[iblk];
dim_t max_size = blk_size;
for (int d = 1; d < data_d.ndims(); ++d) {
max_size = nstl::max(max_size,
data_d.padded_dims()[d] / blocks[d] * blk.strides[d]);
}
return max_size;
}
};
template <SIMPLE_REORDER_TEMPL_DECL>
struct simple_reorder_impl<SIMPLE_REORDER_TEMPL_CALL,
typename utils::enable_if<
tag_i == any && tag_o == any && order_keep == fmt_order::any,
spec::reference>::type>
{
static bool is_applicable(const memory_desc_wrapper &input_d,
const memory_desc_wrapper &output_d, const primitive_attr_t *attr) {
/* supported smask: 0x0...011..10...0,
* i.e. 1 should be contiguous */
int smask = attr ? attr->output_scales_.mask_ : 0;
for (; smask > 0 && !(smask & 0x1); smask >>= 1);
for (; smask > 0 && smask & 0x1; smask >>= 1);
return true
&& input_d.is_blocking_desc()
&& output_d.is_blocking_desc()
&& !output_d.is_additional_buffer()
&& !input_d.is_additional_buffer()
&& smask == 0;
}
static status_t execute(const cpu_reorder_pd_t *pd,
const data_t<type_i> *input, data_t<type_o> *output) {
DECLARE_COMMON_PARAMS();
const size_t nelems = input_d.nelems();
int ndims_start = 0, ndims_mask = 0;
int smask = pd->attr()->output_scales_.mask_;
for (; smask > 0 && !(smask & 0x1); smask >>= 1) ++ndims_start;
for (; smask > 0 && smask & 0x1; smask >>= 1) ++ndims_mask;
assert(smask == 0);
const ptrdiff_t D_start
= utils::array_product(input_d.dims(), ndims_start);
const ptrdiff_t D_mask
= utils::array_product(input_d.dims() + ndims_start, ndims_mask);
const ptrdiff_t D_rest = nelems / D_start / D_mask;
const float *scales = pd->attr()->output_scales_.scales_;
parallel_nd(D_start, D_mask, D_rest,
[&](ptrdiff_t ds, ptrdiff_t dm, ptrdiff_t dr) {
const float scale = scales[dm];
const size_t e = (ds * D_mask + dm) * D_rest + dr;
const auto &i = input[input_d.off_l(e)];
auto &o = output[output_d.off_l(e)];
o = _qz<type_i, type_o>()(i, o, scale, beta);
});
return success;
}
};
/* high level class declaration */
template <SIMPLE_REORDER_TEMPL_DECL, typename spec = void>
struct simple_reorder_t: public cpu_primitive_t {
struct pd_t: public cpu_reorder_pd_t {
using cpu_reorder_pd_t::cpu_reorder_pd_t;
DECLARE_COMMON_PD_T("simple:any", simple_reorder_t);
static status_t create(reorder_pd_t **reorder_pd,
engine_t *engine, const primitive_attr_t *attr,
engine_t *src_engine, const memory_desc_t *src_md,
engine_t *dst_engine, const memory_desc_t *dst_md) {
bool args_ok = true
&& src_md->data_type == type_i
&& dst_md->data_type == type_o
&& simple_reorder_impl<SIMPLE_REORDER_TEMPL_CALL, spec>::
is_applicable(src_md, dst_md, attr);
if (!args_ok)
return status::invalid_arguments;
auto _pd = new pd_t(engine, attr, src_engine, src_md, dst_engine,
dst_md);
if (_pd == nullptr) return status::out_of_memory;
if (_pd->init() != status::success) {
delete _pd;
return status::unimplemented;
}
return safe_ptr_assign<reorder_pd_t>(*reorder_pd, _pd);
}
};
simple_reorder_t(const pd_t *apd): cpu_primitive_t(apd) {}
virtual status_t execute(const exec_ctx_t &ctx) const override {
auto input = CTX_IN_MEM(const data_t<type_i> *, MKLDNN_ARG_FROM);
auto output = CTX_OUT_MEM(data_t<type_o> *, MKLDNN_ARG_TO);
simple_reorder_impl<SIMPLE_REORDER_TEMPL_CALL, spec>::execute(
pd(), input, output);
return status::success;
}
private:
const pd_t *pd() const { return (const pd_t *)primitive_t::pd(); }
};
#undef SIMPLE_REORDER_TEMPL_DECL
#undef SIMPLE_REORDER_TEMPL_CALL
}
}
}
#endif
// vim: et ts=4 sw=4 cindent cino^=l0,\:0,N-s
Reference in a new issue View git blame Copy permalink