#include "graph/backend/dnnl/executables/shuffle.hpp"
namespace dnnl {
namespace impl {
namespace graph {
namespace dnnl_impl {
shuffle_executable_t::desc_t shuffle_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()) {
auto pd = graph::utils::any_cast<dnnl::shuffle_forward::primitive_desc>(
pd_cache.at(op.get()));
return {pd, true};
}
const int group = static_cast<int>(op->get_attr<int64_t>(op_attr::groups));
const int axis = static_cast<int>(op->get_attr<int64_t>(op_attr::axis));
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);
auto src = make_dnnl_memory_desc(op->get_input_logical_tensor(0));
auto dst = make_dnnl_memory_desc(op->get_output_logical_tensor(0));
dst = to_format_any(dst);
dnnl::shuffle_forward::primitive_desc pd(p_engine,
prop_kind::forward_inference, src, dst, axis, group, prm_attr);
pd_cache.insert({op.get(), pd});
return {pd, false};
}
arg_indices_t shuffle_executable_t::get_arg_indices(const op_t *op) {
return get_arg_indices_for_siso_op(op);
}
} } } }