#include "graph/backend/dnnl/executables/concat.hpp"
namespace dnnl {
namespace impl {
namespace graph {
namespace dnnl_impl {
#define VCHECK_OP_EXECUTABLE(cond, msg, ...) \
if (!(cond)) { VERROR(graph, op_executable, msg, ##__VA_ARGS__); }
concat_executable_t::desc_t concat_executable_t::create_desc(
std::shared_ptr<op_t> &op, const dnnl::engine &p_engine,
pd_cache_t &pd_cache, const fpmath_t &fpmath, bool use_block_layout) {
if (pd_cache.find(op.get()) != pd_cache.end()) {
return {graph::utils::any_cast<dnnl::concat::primitive_desc>(
pd_cache.at(op.get())),
false};
}
auto get_forced_format_tag = [](const dims &in_dims) -> format_tag {
if (in_dims.size() == 4)
return format_tag::acdb;
else
return get_ncx_format(in_dims);
};
const auto rank = op->get_output_logical_tensor(0).ndims;
const auto res = utils::try_reverse_axis(
op->get_attr<int64_t>(op_attr::axis), rank);
VCHECK_OP_EXECUTABLE(res.first, "invalid axis for concat");
const auto axis = res.second;
dnnl::primitive_attr prm_attr;
if (op->has_attr(op_attr::fusion_info)) {
const fusion_info_t &fusion_info
= op->get_attr<fusion_info_t>(op_attr::fusion_info);
prm_attr = make_dnnl_primitive_attr(op, fusion_info);
}
prm_attr.set_scratchpad_mode(dnnl::scratchpad_mode::user);
std::vector<memory::desc> src_mds;
src_mds.reserve(op->num_inputs());
for (const auto &in_val : op->get_input_values()) {
const auto tmp_desc
= make_dnnl_memory_desc(in_val->get_logical_tensor());
src_mds.emplace_back(tmp_desc.get_dims(), tmp_desc.get_data_type(),
get_forced_format_tag(tmp_desc.get_dims()));
}
const auto tmp_desc
= make_dnnl_memory_desc(op->get_output_logical_tensor(0));
auto dst = memory::desc {tmp_desc.get_dims(), tmp_desc.get_data_type(),
get_forced_format_tag(tmp_desc.get_dims())};
dnnl::concat::primitive_desc pd(
p_engine, dst, static_cast<int>(axis), src_mds, prm_attr);
pd_cache.insert({op.get(), pd});
return {pd, false};
}
arg_indices_t concat_executable_t::get_arg_indices(const op_t *op) {
return get_arg_indices_for_miso_op(op);
}
} } } }