#include "graph/backend/dnnl/kernels/shuffle.hpp"
#include "graph/backend/dnnl/patterns/fusions.hpp"
#include "graph/backend/dnnl/patterns/pattern_matcher_pass.hpp"
namespace dnnl {
namespace impl {
namespace graph {
namespace dnnl_impl {
namespace pattern {
namespace pm = graph::utils::pm;
using in_edges_t = pm::in_edges_t;
using pb_graph = pm::pb_graph_t;
using FCreatePattern = graph::pass::FCreatePattern;
DNNL_BACKEND_REGISTER_PATTERN_DEF_BEGIN(shuffle_fusion)
DNNL_BACKEND_REGISTER_PATTERN_MATCHER_PASS(dnnl, shuffle_fusion)
.set_priority(8.2f)
.set_kind(partition_kind_t::misc_post_ops)
.set_attr<FCreatePattern>("FCreatePattern",
[](const std::shared_ptr<pb_graph> &pgraph) -> void {
pm::pb_op_t *reshape0
= pgraph->append_op(graph::op_kind::StaticReshape);
pm::pb_op_t *transpose
= pgraph->append_op(graph::op_kind::StaticTranspose,
in_edges_t {in_edge(0, reshape0, 0)});
pgraph->append_op(graph::op_kind::StaticReshape,
in_edges_t {in_edge(0, transpose, 0)});
})
.set_attr<FCreateKernel>("FCreateKernel", []() -> kernel_ptr {
return std::make_shared<shuffle_fwd_t>();
});
DNNL_BACKEND_REGISTER_PATTERN_DEF_END
} } } } }