#include <algorithm>
#include <functional>
#include <limits>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "graph/interface/allocator.hpp"
#include "graph/interface/backend.hpp"
#include "graph/interface/shape_infer.hpp"
#include "graph/utils/utils.hpp"
#include "graph/backend/dnnl/common.hpp"
#include "graph/backend/dnnl/dnnl_backend.hpp"
#include "graph/backend/dnnl/scratchpad.hpp"
#include "common/primitive_desc_iface.hpp"
#include "common/primitive_iface.hpp"
#include "common/stream.hpp"
#if DNNL_CPU_RUNTIME != DNNL_RUNTIME_SYCL
const size_t DNNL_CPU_MEMALIGNMENT = 64;
#endif
#ifdef DNNL_WITH_SYCL
#include "oneapi/dnnl/dnnl_sycl.hpp"
const size_t DNNL_SYCL_MEMALIGNMENT = 64;
#endif
#if DNNL_GPU_RUNTIME == DNNL_RUNTIME_OCL
#include "oneapi/dnnl/dnnl_ocl.hpp"
const size_t DNNL_OCL_MEMALIGNMENT = 0;
#endif
#if DNNL_CPU_RUNTIME == DNNL_RUNTIME_THREADPOOL
#include "cpu/cpu_stream.hpp"
#include "oneapi/dnnl/dnnl_threadpool.hpp"
#endif
namespace dnnl {
namespace impl {
namespace graph {
namespace dnnl_impl {
void *dnnl_allocator_t::malloc(size_t size, const dnnl::engine &p_engine,
const graph::allocator_t *alc, allocator_t::mem_type_t type) {
if (p_engine.get_kind() == dnnl::engine::kind::cpu) {
#if DNNL_CPU_RUNTIME == DNNL_RUNTIME_SYCL
return alc->allocate(size, dnnl::sycl_interop::get_device(p_engine),
dnnl::sycl_interop::get_context(p_engine),
{type, DNNL_SYCL_MEMALIGNMENT});
#else
return alc->allocate(size, {type, DNNL_CPU_MEMALIGNMENT});
#endif
} else if (p_engine.get_kind() == dnnl::engine::kind::gpu) {
#if DNNL_GPU_RUNTIME == DNNL_RUNTIME_SYCL
return alc->allocate(size, dnnl::sycl_interop::get_device(p_engine),
dnnl::sycl_interop::get_context(p_engine),
{type, DNNL_SYCL_MEMALIGNMENT});
#elif DNNL_GPU_RUNTIME == DNNL_RUNTIME_OCL
return alc->allocate(size, dnnl::ocl_interop::get_device(p_engine),
dnnl::ocl_interop::get_context(p_engine),
{type, DNNL_OCL_MEMALIGNMENT});
#else
return nullptr;
#endif
} else {
return nullptr;
}
}
void dnnl_allocator_t::free(
void *p, const dnnl::engine &p_engine, const allocator_t *alc) {
if (p_engine.get_kind() == dnnl::engine::kind::cpu) {
#if DNNL_CPU_RUNTIME == DNNL_RUNTIME_SYCL
assert(!"use event based free");
#else
return alc->deallocate(p);
#endif
} else if (p_engine.get_kind() == dnnl::engine::kind::gpu) {
#if DNNL_GPU_RUNTIME == DNNL_RUNTIME_SYCL \
|| DNNL_GPU_RUNTIME == DNNL_RUNTIME_OCL
assert(!"use event based free");
#endif
}
}
#ifdef DNNL_WITH_SYCL
void dnnl_allocator_t::free(void *p, const dnnl::engine &p_engine,
const allocator_t *alc, const ::sycl::event &deps) {
if (p_engine.get_kind() == dnnl::engine::kind::cpu) {
#if DNNL_CPU_RUNTIME == DNNL_RUNTIME_SYCL
alc->deallocate(p, dnnl::sycl_interop::get_device(p_engine),
dnnl::sycl_interop::get_context(p_engine), deps);
#else
alc->deallocate(p);
#endif
} else if (p_engine.get_kind() == dnnl::engine::kind::gpu) {
#if DNNL_GPU_RUNTIME == DNNL_RUNTIME_SYCL
alc->deallocate(p, dnnl::sycl_interop::get_device(p_engine),
dnnl::sycl_interop::get_context(p_engine), deps);
#endif
}
}
#endif
#if DNNL_GPU_RUNTIME == DNNL_RUNTIME_OCL
void dnnl_allocator_t::free(void *p, const dnnl::engine &p_engine,
const allocator_t *alc, const cl_event &deps) {
if (p_engine.get_kind() != dnnl::engine::kind::gpu) {
assert(!"the engine kind should be gpu");
return;
}
alc->deallocate(p, dnnl::ocl_interop::get_device(p_engine),
dnnl::ocl_interop::get_context(p_engine), deps);
}
#endif
format_tag get_ncx_format(size_t ndim) {
switch (ndim) {
case 1: return format_tag::a;
case 2: return format_tag::ab;
case 3: return format_tag::abc;
case 4: return format_tag::abcd;
case 5: return format_tag::abcde;
case 6: return format_tag::abcdef;
default: return format_tag::undef;
}
}
format_tag get_ncx_format(const dims &adims) {
const auto size = adims.size();
return get_ncx_format(size);
}
dims get_compatible_dilates(const dims &dilates, size_t input_size) {
if (!dilates.empty() && !graph::utils::any_le(dilates, static_cast<dim>(0)))
return utils::fmap(dilates, [](dim x) { return x - 1; });
return dims(input_size - 2, 0);
}
dims group_dims(const dims &adims, dim groups) {
auto new_dims = adims;
new_dims.insert(new_dims.begin(), groups);
new_dims[1] /= groups;
return new_dims;
}
dnnl::engine make_dnnl_engine(const engine_t &g_engine) {
dnnl::engine engine;
engine.reset(const_cast<engine_t *>(&g_engine), true); return engine;
}
dnnl::engine make_host_engine() {
return dnnl::engine(dnnl::engine::kind::cpu, 0);
}
dnnl::stream make_dnnl_stream(
const dnnl::engine &p_engine, const stream_t &g_stream) {
UNUSED(p_engine);
dnnl::stream strm;
strm.reset(const_cast<stream_t *>(&g_stream), true); return strm;
}
dnnl::memory::desc make_dnnl_memory_desc(const logical_tensor_t <) {
const logical_tensor_wrapper_t ltw(lt);
const auto dtype = ltw.data_type() == graph::data_type::boolean
? data_type::u8
: static_cast<data_type>(ltw.data_type());
if (ltw.is_host_scalar()) { return dnnl::memory::desc::host_scalar(dtype); }
if (ltw.is_opaque()) {
#ifdef DNNL_GRAPH_LAYOUT_DEBUG
const auto format_tag
= static_cast<dnnl::memory::format_tag>(ltw.layout_id());
if (format_tag < dnnl::memory::format_tag::format_tag_last
&& format_tag > dnnl::memory::format_tag::any) {
return {ltw.vdims(), dtype, format_tag};
}
#endif
const auto &td = dnnl_backend_t::get_singleton().get_mem_desc(
static_cast<size_t>(ltw.layout_id()));
return graph::utils::any_cast<memory::desc>(td.value());
} else if (ltw.is_any()) {
if (ltw.ndims() > 0) {
return {ltw.vdims(), dtype, dnnl::memory::format_tag::any};
} else if (ltw.ndims() == 0) {
return {dims {1}, dtype, dnnl::memory::format_tag::any};
} else {
return {dims {}, dtype, dnnl::memory::format_tag::any};
}
} else if (ltw.is_strided()) {
if (ltw.ndims() > 0) {
return {ltw.vdims(), dtype, ltw.vstrides()};
} else if (ltw.ndims() == 0) {
return {dims {1}, dtype, dims {1}};
} else {
assertm(false,
"An empty strided logical tensor can't be convert to "
"memory desc");
return {dims {}, dtype, dims {}};
}
} else {
return {};
}
}
dnnl::memory make_dnnl_memory(const dnnl::memory::desc &md,
const dnnl::engine &p_engine, void *handle) {
if (p_engine.get_kind() == dnnl::engine::kind::cpu) {
#if DNNL_CPU_RUNTIME == DNNL_RUNTIME_SYCL
return dnnl::sycl_interop::make_memory(
md, p_engine, dnnl::sycl_interop::memory_kind::usm, handle);
#else
return dnnl::memory(md, p_engine, handle);
#endif
} else if (p_engine.get_kind() == dnnl::engine::kind::gpu) {
#if DNNL_GPU_RUNTIME == DNNL_RUNTIME_SYCL
return dnnl::sycl_interop::make_memory(
md, p_engine, dnnl::sycl_interop::memory_kind::usm, handle);
#elif DNNL_GPU_RUNTIME == DNNL_RUNTIME_OCL
return dnnl::ocl_interop::make_memory(
md, p_engine, dnnl::ocl_interop::memory_kind::usm, handle);
#else
return dnnl::memory(md, p_engine, handle);
#endif
} else {
assert(!"only cpu and gpu memory are valid");
return {};
}
}
dnnl::memory make_dnnl_memory(
const tensor_t &atensor, const dnnl::engine &p_engine) {
const logical_tensor_t lt = atensor.get_logical_tensor();
const dnnl::memory::desc md = make_dnnl_memory_desc(lt);
const logical_tensor_wrapper_t ltw(lt);
if (ltw.is_host_scalar()) {
assert(!"make_dnnl_memory: got host scalar tensor.");
return {};
} else {
return make_dnnl_memory(md, p_engine, atensor.get_data_handle());
}
}
memory::desc expand(const memory::desc &adesc, int tgt_ndims) {
int64_t org_ndims = adesc.get_ndims();
dnnl::memory::dims expanded_dims = adesc.get_dims();
expanded_dims.insert(expanded_dims.begin(), tgt_ndims - org_ndims, 1);
return adesc.reshape(expanded_dims);
}
std::vector<int64_t> get_last_two_dims_permutation(int ndims) {
assert(ndims > 1);
std::vector<int64_t> axes(ndims);
std::iota(axes.begin(), axes.end(), 0);
const auto last_dim = static_cast<dims::size_type>(ndims - 1);
std::swap(axes[last_dim], axes[last_dim - 1]);
return axes;
}
std::vector<int64_t> get_permutation(int ndims, const std::string &from_format,
const std::string &to_format) {
assert(ndims > 2);
assert(from_format == "NCX" || from_format == "NXC" || from_format == "IOX"
|| from_format == "OIX" || from_format == "XIO"
|| from_format == "XOI");
assert(to_format == "NCX" || to_format == "NXC" || to_format == "IOX"
|| to_format == "OIX" || to_format == "XIO" || to_format == "XOI");
size_t spatial_dims = static_cast<size_t>(ndims - 2);
std::vector<int64_t> axes(ndims);
size_t axes_idx = 0;
for (const auto &ch : from_format) {
size_t to_idx = to_format.find(ch);
if (ch == 'X') {
for (size_t spatial_idx = 0; spatial_idx < spatial_dims;
++spatial_idx) {
axes[axes_idx++] = to_idx + spatial_idx;
}
} else {
if (to_idx > to_format.find('X')) {
axes[axes_idx++] = to_idx + spatial_dims - 1;
} else {
axes[axes_idx++] = to_idx;
}
}
}
return axes;
}
memory::desc transpose(const memory::desc &adesc, dim dim0, dim dim1) {
std::vector<int> axes(static_cast<std::size_t>(adesc.get_ndims()));
std::iota(axes.begin(), axes.end(), 0);
axes[static_cast<std::size_t>(dim0)] = dim1;
axes[static_cast<std::size_t>(dim1)] = dim0;
return adesc.permute_axes(axes);
}
memory::desc to_grouped(const memory::desc &adesc, dim groups) {
auto grouped_shape = group_dims(adesc.get_dims(), groups);
return adesc.reshape(grouped_shape);
}
memory::desc from_grouped(const memory::desc &adesc) {
auto new_dims = adesc.get_dims();
const dim groups = new_dims.front();
new_dims.erase(new_dims.begin());
new_dims[0] *= groups;
return adesc.reshape(new_dims, true);
}
memory::desc to_format_any(const memory::desc &adesc) {
return memory::desc(
adesc.get_dims(), adesc.get_data_type(), memory::format_tag::any);
}
dims get_ncx_strides(const dims &shape) {
auto _shape = shape;
for (size_t i = 0; i < _shape.size(); i++) {
if (_shape[i] == 0) _shape[i] = 1;
}
dims strides(_shape.size());
for (auto it = _shape.begin(); it < _shape.end(); ++it) {
const auto val = std::accumulate(
std::next(it), _shape.end(), 1, std::multiplies<dim_t>());
const auto dist = std::distance(_shape.begin(), it);
strides[static_cast<size_t>(dist)] = val;
}
return strides;
}
dims get_nxc_strides(const dims &shape) {
auto _shape = shape;
for (size_t i = 0; i < _shape.size(); i++) {
if (_shape[i] == 0) _shape[i] = 1;
}
dims strides(_shape.size());
dim tmp, tmp1, tmp2;
switch (_shape.size()) {
case 3:
strides[0] = _shape[1] * _shape[2];
strides[1] = 1;
strides[2] = _shape[1];
break;
case 4:
tmp = _shape[1] * _shape[3];
strides[0] = tmp * _shape[2];
strides[1] = 1;
strides[2] = tmp;
strides[3] = _shape[1];
break;
case 5:
tmp1 = _shape[1] * _shape[4];
tmp2 = tmp1 * _shape[3];
strides[0] = tmp2 * _shape[2];
strides[1] = 1;
strides[2] = tmp2;
strides[3] = tmp1;
strides[4] = _shape[1];
break;
case 6:
tmp1 = _shape[1] * _shape[5];
tmp2 = tmp1 * _shape[3] * _shape[4];
strides[0] = tmp2 * _shape[2];
strides[1] = 1;
strides[2] = tmp2;
strides[3] = tmp1 * _shape[4];
strides[4] = tmp1;
strides[5] = _shape[1];
break;
default: strides = get_ncx_strides(_shape);
}
return strides;
}
memory::desc to_nxc_format(const memory::desc &adesc) {
if (is_format(adesc, "nxc")) return adesc;
const auto ori_dims = adesc.get_dims();
dims strides = get_nxc_strides(ori_dims);
return {ori_dims, adesc.get_data_type(), strides};
}
bool is_format(const memory::desc &adesc, memory::format_tag tag) {
return adesc == memory::desc(adesc.get_dims(), adesc.get_data_type(), tag);
}
bool is_format(const memory::desc &adesc, const std::string &tag) {
if (!graph::utils::one_of(tag, "ncx", "nxc")) {
assertm(false, "wrong tag to check memory format");
return false;
}
if (adesc.get_format_kind() != format_kind::blocked
|| adesc.get_inner_nblks() != 0)
return false;
const auto &strides = adesc.get_strides();
const auto &shape = adesc.get_dims();
if ("ncx" == tag) { return strides == get_ncx_strides(shape); }
return strides == get_nxc_strides(shape);
}
bool is_4c_blocked(const memory::desc &adesc) {
if (adesc.get_format_kind() != format_kind::blocked) return false;
return adesc.get_inner_nblks() == 1 && adesc.get_inner_idxs()[0] == 1
&& adesc.get_inner_blks()[0] == 4;
}
bool is_plain(const memory::desc &adesc) {
if (adesc.get_format_kind() != format_kind::blocked) return false;
return adesc.get_inner_nblks() == 0;
}
dims get_dense_strides(const dims &shape) {
dims strides(shape.size());
for (auto it = shape.begin(); it < shape.end(); ++it) {
const auto val = std::accumulate(
std::next(it), shape.end(), 1, [](dim_t x, dim_t y) {
return std::max<dim_t>(x, 1) * std::max<dim_t>(y, 1);
});
const auto dist = std::distance(shape.begin(), it);
strides[static_cast<size_t>(dist)] = val;
}
return strides;
}
memory::desc to_ncx_format(const memory::desc &adesc) {
return memory::desc(adesc.get_dims(), adesc.get_data_type(),
get_ncx_format(adesc.get_ndims()));
}
inline bool maybe_reorder_value(const value_t *val) {
for (const auto &consumer : val->get_consumers()) {
if (consumer.get_op().get_kind() == graph::op_kind::Reorder) {
return true;
}
}
bool is_out_value = val->has_producer()
&& val->get_producer().get_kind() == graph::op_kind::Reorder;
return is_out_value;
}
void set_all_layout_to_any(std::vector<std::shared_ptr<op_t>> &subgraph) {
for (auto &cur_op : subgraph) {
for (const auto &val : cur_op->get_input_values()) {
if (maybe_reorder_value(val.get())) continue;
val->set_layout_type(layout_type::any);
}
for (const auto &val : cur_op->get_output_values()) {
if (maybe_reorder_value(val.get())) continue;
val->set_layout_type(layout_type::any);
}
}
}
status_t fill_layout_info(logical_tensor_t *lt, const memory::desc &md) {
const logical_tensor_wrapper_t ltw(lt);
if (ltw.is_any()) { const int lt_ndims = ltw.ndims();
const int md_ndims = md.get_ndims();
if (md_ndims == 0) {
if (lt_ndims < 0) {
lt->layout_type = layout_type::undef;
return status::success;
} else {
assertm(false,
"The logical tensor should be also empty when the "
"memory desc is empty");
return status::invalid_arguments;
}
}
if (lt_ndims < 0 && md_ndims > 0) { lt->ndims = md_ndims;
const auto &dims = md.get_dims();
std::copy(dims.data(), dims.data() + md_ndims, lt->dims);
lt->data_type = static_cast<data_type_t>(md.get_data_type());
}
if (lt_ndims == 0 && graph::utils::prod(md.get_dims()) == 1) { lt->layout_type = layout_type::strided;
}
if (lt->id != std::numeric_limits<size_t>::max() && is_plain(md)) {
lt->layout_type = layout_type::strided;
graph::utils::array_copy(lt->layout.strides,
md.get_strides().data(), md.get_ndims());
} else {
graph::utils::optional_t<size_t> layout_id
= dnnl_backend_t::get_singleton().set_mem_desc(md);
lt->layout.layout_id = layout_id.value();
lt->layout_type = layout_type::opaque;
}
}
return status::success;
}
status_t fill_layout_info(
const std::shared_ptr<value_t> &val, const memory::desc &md) {
logical_tensor_t lt = val->get_logical_tensor();
const logical_tensor_wrapper_t ltw(lt);
if (ltw.is_any()) { const int lt_ndims = ltw.ndims();
const int md_ndims = md.get_ndims();
if (md_ndims == 0) {
if (lt_ndims < 0) {
val->set_layout_type(layout_type::undef);
return status::success;
} else {
assertm(false,
"The logical tensor should be also empty when the "
"memory desc is empty");
return status::invalid_arguments;
}
}
if (lt_ndims < 0 && md_ndims > 0) { val->set_dims(md.get_dims());
val->set_data_type(
static_cast<impl::data_type_t>(md.get_data_type()));
}
if (lt_ndims == 0 && graph::utils::prod(md.get_dims()) == 1) { val->set_strides({});
}
if (ltw.id() != std::numeric_limits<size_t>::max() && is_plain(md)) {
val->set_strides(md.get_strides());
} else {
val->set_layout_id(
dnnl_backend_t::get_singleton().set_mem_desc(md).value());
}
}
return status::success;
}
std::shared_ptr<value_t> insert_empty_scratchpad(std::shared_ptr<op_t> &op) {
logical_tensor_t lt = empty_logical_tensor_with_default_id();
auto scratchpad_val = std::make_shared<value_t>(*op, op->num_outputs(), lt);
op->add_output(scratchpad_val);
scratchpad_val->set_data_type(graph::data_type::u8);
return scratchpad_val;
}
std::shared_ptr<value_t> insert_empty_workspace(std::shared_ptr<op_t> &op) {
logical_tensor_t lt = empty_logical_tensor_with_default_id();
auto workspace_val = std::make_shared<value_t>(*op, op->num_outputs(), lt);
op->add_output(workspace_val);
return workspace_val;
}
dnnl::memory::format_tag get_format_tag(const dnnl::memory::desc &md) {
std::string blk_tag = md2fmt_tag_str(md.get());
dnnl::memory::format_tag format_tag = dnnl::memory::format_tag::undef;
for (size_t tag = 0; tag < dnnl_format_tag_last; ++tag) {
if (dnnl_fmt_tag2str((dnnl_format_tag_t)tag) == blk_tag) {
format_tag = static_cast<dnnl::memory::format_tag>(tag);
break;
}
}
return format_tag;
}
size_t generate_constant_md_hash(
size_t part_id, const std::vector<dnnl::memory::desc> &const_mds) {
size_t key = 0;
key = hash_combine(key, part_id);
for (auto &md : const_mds) {
auto md_hash = impl::primitive_hashing::get_md_hash(*md.get());
key = hash_combine(key, md_hash);
}
return key;
}
dnnl::accumulation_mode str2accumulation_mode(
const std::string &accumulation_mode_str) {
if (accumulation_mode_str == "strict") {
return dnnl::accumulation_mode::strict;
} else if (accumulation_mode_str == "relaxed") {
return dnnl::accumulation_mode::relaxed;
} else if (accumulation_mode_str == "any") {
return dnnl::accumulation_mode::any;
} else if (accumulation_mode_str == "s32") {
return dnnl::accumulation_mode::s32;
} else if (accumulation_mode_str == "f32") {
return dnnl::accumulation_mode::f32;
} else if (accumulation_mode_str == "f16") {
return dnnl::accumulation_mode::f16;
} else {
assert(!"unknown accumulation mode");
return dnnl::accumulation_mode::strict;
}
}
void prolong_temporary_scratchpad_lifetime(const stream_t *g_stream,
const std::shared_ptr<temporary_scratchpad_t> &scratchpad) {
#if DNNL_CPU_RUNTIME == DNNL_RUNTIME_THREADPOOL
auto *tp_stream
= dnnl::impl::utils::downcast<dnnl::impl::cpu::cpu_stream_t *>(
const_cast<stream_t *>(g_stream));
tp_stream->before_exec_hook();
parallel(1, [=](int, int) { UNUSED(scratchpad); });
tp_stream->after_exec_hook();
#endif
}
status_t dnnl_primitive_execute_without_tp_hook(const primitive &prim,
const stream &astream,
const std::unordered_map<int, memory> &exec_args) {
std::vector<dnnl_exec_arg_t> vec_args;
vec_args.reserve(exec_args.size());
for (const auto &a : exec_args)
vec_args.push_back({a.first, a.second.get(true)});
const primitive_iface_t *primitive_iface = prim.get();
stream_t *stream = astream.get();
int nargs = static_cast<int>(vec_args.size());
const dnnl_exec_arg_t *c_args = vec_args.data();
bool ok = true && !dnnl::impl::utils::any_null(primitive_iface, stream)
&& primitive_iface->engine() == stream->engine()
&& IMPLICATION(nargs > 0, c_args != nullptr);
if (!ok) return status::invalid_arguments;
exec_args_t args;
status_t status = cvt_primitive_args(
primitive_iface->pd()->impl().get(), nargs, c_args, args);
if (status != status::success) return status;
exec_ctx_t ctx(stream, std::move(args));
status = dnnl::impl::primitive_execute(primitive_iface, ctx);
return status;
}
} } } }