onednn-src 0.1.13

/*******************************************************************************
* Copyright 2021 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 GRAPH_BACKEND_DNNL_PASSES_UTILS_HPP
#define GRAPH_BACKEND_DNNL_PASSES_UTILS_HPP

#include <algorithm>
#include <functional>
#include <iostream>
#include <map>
#include <memory>
#include <set>
#include <string>
#include <utility>
#include <vector>
#include <unordered_map>
#include <unordered_set>

#include "graph/interface/c_types_map.hpp"
#include "graph/interface/graph.hpp"
#include "graph/interface/op.hpp"
#include "graph/interface/value.hpp"

#include "graph/utils/utils.hpp"

#include "graph/backend/dnnl/subgraph.hpp"
#include "graph/backend/dnnl/utils.hpp"

#include "oneapi/dnnl/dnnl.hpp"

#define VCHECK_UTILS(cond, status, msg, ...) \
    VCONDCHECK(graph, create, check, utils, (cond), status, msg, ##__VA_ARGS__);

namespace dnnl {
namespace impl {
namespace graph {
namespace dnnl_impl {

// The pass_pipeline_t class is used to manage all transformation passes to run
// on a subgraph. Users should add passes need to run to the pipeline with a
// user defined order. And then call the run() method to run those added passes.
// After running each pass, the pipeline will choose to visualize the processed
// subgraph by using the visualizer.
class pass_pipeline_t {
public:
    using pass_signature
            = std::function<status_t(std::shared_ptr<subgraph_t> &)>;

    pass_pipeline_t() = default;

    pass_pipeline_t(const subgraph_visualizer_t &vis,
            bool enable_validator = true, bool enable_visualizer = true)
        : visualizer_(vis)
        , is_layout_sensitive_(false)
        , is_memory_sensitive_(false)
        , enable_validator_(enable_validator)
        , enable_visualizer_(enable_visualizer) {}

    // Reset the visualize arguments
    void reset_visualize_arg(
            bool is_layout_sensitive, bool is_memory_sensitive) {
        is_layout_sensitive_ = is_layout_sensitive;
        is_memory_sensitive_ = is_memory_sensitive;
    }

    // Add a pass to the pipeline. The current visualize arguments will be
    // recorded for the added pass and be used when visualize the subgraph
    // processed by this pass.
    void add_pass(const pass_signature &apass, const std::string &name) {
        passes_.emplace_back(apass);
        names_.emplace_back(name);
        is_layout_sensitives_.push_back(is_layout_sensitive_);
        is_memory_sensitives_.push_back(is_memory_sensitive_);
    }

    // Run all added passes
    status_t run(std::shared_ptr<subgraph_t> &sg) {
        status_t ret;
        for (size_t i = 0; i < passes_.size(); i++) {
            ret = passes_[i](sg);
            VCHECK_UTILS(ret == status::success, ret, "run pass %s failed",
                    names_[i].c_str());

            // Dump the subgraph to dot file
            if (enable_visualizer_) {
                visualizer_.run(sg, names_[i], is_layout_sensitives_[i],
                        is_memory_sensitives_[i]);
            }

            // Validate the subgraph after each pass
            if (enable_validator_) {
                ret = validator_.run(sg);
                VCHECK_UTILS(ret == status::success, ret,
                        "validation failed "
                        "after run pass %s",
                        names_[i].c_str());
            }
        }
        return status::success;
    }

private:
    // The added passes and their names
    std::vector<pass_signature> passes_;
    std::vector<std::string> names_;

    // The recorded visualize arguments for each pass
    std::vector<bool> is_layout_sensitives_;
    std::vector<bool> is_memory_sensitives_;

    subgraph_visualizer_t visualizer_;
    subgraph_validator_t validator_;

    // The current visualize arguments
    bool is_layout_sensitive_;
    bool is_memory_sensitive_;

    bool enable_validator_;
    bool enable_visualizer_;
};

#define BACKEND_DNNL_ADD_PASS(pipeline, pass) pipeline.add_pass(pass, #pass)

status_t set_given_inputs_outputs(std::shared_ptr<subgraph_t> &sg,
        const std::vector<logical_tensor_t> &inputs,
        const std::vector<logical_tensor_t> &outputs);

status_t set_given_inputs_outputs(std::vector<std::shared_ptr<op_t>> &subgraph,
        const std::vector<logical_tensor_t> &inputs,
        const std::vector<logical_tensor_t> &outputs);

void set_weight_bias_constant(std::shared_ptr<subgraph_t> &sg);

inline bool is_preprocess_op(op_t &op) {
    static const std::set<op_kind_t> preprocess_ops = {
            op_kind::_permute,
            op_kind::_to_group,
            op_kind::_from_group,
            op_kind::_unsqueeze,
            op_kind::_squeeze,
            op_kind::_reshape,
            op_kind::_transpose,
            op_kind::_identity,
    };
    return preprocess_ops.count(op.get_kind()) != 0;
}

void merge_common_eltwise_attrs(
        const std::shared_ptr<op_t> &org_op, std::shared_ptr<op_t> &new_op);

inline const std::map<op_kind_t, dnnl::algorithm> &get_eltwise_alg_map() {
    static const std::map<op_kind_t, dnnl::algorithm> &eltwise_alg_map = {
            {op_kind::Abs, dnnl::algorithm::eltwise_abs},
            {op_kind::Clamp, dnnl::algorithm::eltwise_clip_v2},
            {op_kind::Elu, dnnl::algorithm::eltwise_elu},
            {op_kind::Exp, dnnl::algorithm::eltwise_exp},
            {op_kind::GELU, dnnl::algorithm::eltwise_gelu_erf},
            {op_kind::HardSigmoid, dnnl::algorithm::eltwise_hardsigmoid},
            {op_kind::HardSwish, dnnl::algorithm::eltwise_hardswish},
            {op_kind::LeakyReLU, dnnl::algorithm::eltwise_relu},
            {op_kind::Log, dnnl::algorithm::eltwise_log},
            {op_kind::Mish, dnnl::algorithm::eltwise_mish},
            {op_kind::ReLU, dnnl::algorithm::eltwise_relu},
            {op_kind::Round, dnnl::algorithm::eltwise_round},
            {op_kind::Sigmoid, dnnl::algorithm::eltwise_logistic},
            {op_kind::Sqrt, dnnl::algorithm::eltwise_sqrt},
            {op_kind::Square, dnnl::algorithm::eltwise_square},
            {op_kind::Tanh, dnnl::algorithm::eltwise_tanh},
            {op_kind::AbsBackward, dnnl::algorithm::eltwise_abs},
            {op_kind::ClampBackward, dnnl::algorithm::eltwise_clip_v2},
            {op_kind::EluBackward, dnnl::algorithm::eltwise_elu},
            {op_kind::GELUBackward, dnnl::algorithm::eltwise_gelu_erf},
            {op_kind::HardSigmoidBackward,
                    dnnl::algorithm::eltwise_hardsigmoid},
            {op_kind::HardSwishBackward, dnnl::algorithm::eltwise_hardswish},
            {op_kind::MishBackward, dnnl::algorithm::eltwise_mish},
            {op_kind::ReLUBackward, dnnl::algorithm::eltwise_relu},
            {op_kind::SigmoidBackward, dnnl::algorithm::eltwise_logistic},
            {op_kind::ReLUBackward, dnnl::algorithm::eltwise_relu},
            {op_kind::SqrtBackward, dnnl::algorithm::eltwise_sqrt},
            {op_kind::TanhBackward, dnnl::algorithm::eltwise_tanh},
    };
    return eltwise_alg_map;
}

inline dnnl::algorithm get_eltwise_alg(
        const std::shared_ptr<op_t> &op, bool bwd) {
    using algo = dnnl::algorithm;

    const op_kind_t opk = op->get_kind();
    const auto &map = get_eltwise_alg_map();

    auto it = map.find(opk);
    if (it == map.end()) {
        assert(!"unexpected op kind");
        return algo::undef;
    } else {
        auto alg = it->second;
        // handle attributes
        if (opk == op_kind::GELU || opk == op_kind::GELUBackward) {
            if (op->has_attr(op_attr::mode)
                    && op->get_attr<std::string>(op_attr::mode)
                            == "gelu_tanh") {
                alg = algo::eltwise_gelu_tanh;
            }
        }

        if (bwd && op->has_attr(op_attr::use_dst)
                && op->get_attr<bool>(op_attr::use_dst)) {
            switch (opk) {
                case op_kind::ClampBackward:
                    alg = algo::eltwise_clip_v2_use_dst_for_bwd;
                    break;
                case op_kind::EluBackward:
                    alg = algo::eltwise_elu_use_dst_for_bwd;
                    break;
                case op_kind::ReLUBackward:
                    alg = algo::eltwise_relu_use_dst_for_bwd;
                    break;
                case op_kind::SigmoidBackward:
                    alg = algo::eltwise_logistic_use_dst_for_bwd;
                    break;
                case op_kind::SqrtBackward:
                    alg = algo::eltwise_sqrt_use_dst_for_bwd;
                    break;
                case op_kind::TanhBackward:
                    alg = algo::eltwise_tanh_use_dst_for_bwd;
                    break;
                default: break;
            }
        }

        return alg;
    }
}

inline const std::map<op_kind_t, dnnl::algorithm> &get_reduction_alg_map() {
    static const std::map<op_kind_t, dnnl::algorithm> &reduction_alg_map = {
            {op_kind::ReduceL1, dnnl::algorithm::reduction_norm_lp_power_p_sum},
            {op_kind::ReduceL2, dnnl::algorithm::reduction_norm_lp_sum},
            {op_kind::ReduceMax, dnnl::algorithm::reduction_max},
            {op_kind::ReduceMean, dnnl::algorithm::reduction_mean},
            {op_kind::ReduceMin, dnnl::algorithm::reduction_min},
            {op_kind::ReduceProd, dnnl::algorithm::reduction_mul},
            {op_kind::ReduceSum, dnnl::algorithm::reduction_sum},
    };
    return reduction_alg_map;
}

inline bool is_eltwise_kind(op_kind_t kind) {
    static const std::set<op_kind_t> eltwise_kinds {
            op_kind::Abs,
            op_kind::Clamp,
            op_kind::Elu,
            op_kind::Exp,
            op_kind::GELU,
            op_kind::HardSigmoid,
            op_kind::HardSwish,
            op_kind::LeakyReLU,
            op_kind::Log,
            op_kind::Mish,
            op_kind::ReLU,
            op_kind::Round,
            op_kind::Sigmoid,
            op_kind::SoftPlus,
            op_kind::Sqrt,
            op_kind::Square,
            op_kind::Tanh,
    };
    return eltwise_kinds.find(kind) != eltwise_kinds.end();
}

inline bool is_eltwise_bwd_kind(op_kind_t kind) {
    static const std::set<op_kind_t> eltwise_bwd_kinds {
            op_kind::AbsBackward,
            op_kind::ClampBackward,
            op_kind::EluBackward,
            op_kind::GELUBackward,
            op_kind::HardSigmoidBackward,
            op_kind::HardSwishBackward,
            op_kind::MishBackward,
            op_kind::ReLUBackward,
            op_kind::SigmoidBackward,
            op_kind::SqrtBackward,
            op_kind::TanhBackward,
    };
    return eltwise_bwd_kinds.find(kind) != eltwise_bwd_kinds.end();
}

inline bool is_binary_kind(op_kind_t kind) {
    static const std::set<op_kind_t> binary_kinds = {
            op_kind::Add,
            op_kind::Subtract,
            op_kind::Multiply,
            op_kind::Divide,
            op_kind::Minimum,
            op_kind::Maximum,
    };
    return binary_kinds.find(kind) != binary_kinds.end();
}

inline bool is_reduction_kind(op_kind_t kind) {
    static const std::set<op_kind_t> reduction_kinds = {
            op_kind::ReduceL1,
            op_kind::ReduceL2,
            op_kind::ReduceMax,
            op_kind::ReduceMean,
            op_kind::ReduceMin,
            op_kind::ReduceProd,
            op_kind::ReduceSum,
    };
    return reduction_kinds.find(kind) != reduction_kinds.end();
}

std::vector<value_t *> get_constant_block_output_values(
        const std::shared_ptr<subgraph_t> &sg);

status_t infer_shape(std::shared_ptr<subgraph_t> &sg);

const std::map<op_kind_t, dnnl::algorithm> &get_binary_alg_map();

// (3, 4) * (3, 4) is doable
// (1, 4) * (3, 4) is doable
// (3, 4, 5) * (4, 5) is doable
// (3, 4, 5) * (1, 5) is doable
// (3, 4, 5) * (2, 4, 5) is NOT doable
bool binary_doable(
        const std::vector<dim_t> &shape_0, const std::vector<dim_t> &shape_1);

bool prelu_doable(const std::vector<dim_t> &src_dims,
        const std::vector<dim_t> &wei_dims, const std::string &data_format,
        const bool per_channel_broadcast);

// Checks whether chain of Reshape, Transpose, Reshape is fusible
// to dnnl_shuffle. Returns following pair:
// (is_fusible, (axis, groups))
// axis and groups store relevant information only when 'is_fusible = true'.
std::pair<bool, std::pair<size_t, int64_t>> shuffle_fusible(
        const op_t *reshape0, op_t *reshape1, op_t *transpose);

// For some shapes, post binary will run into oneDNN's ref path and has poor
// performance. So, we check the shape in this function and only make
// per_tensor, per_channel, per_mb_w(MatMul) and full tensor broadcast
// binary able to be fused.
bool post_binary_fusible(const op_t *base_op, const op_t *bin_op,
        engine_kind_t ekind = engine_kind::cpu);

// binary + sqrt post-op fusion is unsupported on NVIDIA GPU
bool post_eltwise_fusible(
        const op_t *base_op, const op_t *elt_op, graph::engine_kind_t ekind);

// oneDNN support post depthwise conv fusion. This function is used to check if
// two conv ops can be fused as a conv + depthwise pattern.
bool post_depthwise_conv_fusible(
        const op_t *base_conv_op, const op_t *post_conv_op);

// Get the map between base op kind and fusible post ops kinds. The map is
// determined by oneDNN's fusion capability and may change. For example, a
// dnnl_eltwise op can't fuse dnnl_eltwise op, but dnnl_convolution can.
const std::unordered_map<op_kind_t, std::unordered_set<op_kind_t>> &
get_post_ops_fusible_map();

std::string kind2str(op_kind_t kind);

// This function is used to check if a dnnl_reorder op is converted from or act
// as a TypeCast op. This function will only return true for a dnnl_reorder op
// which only has different input/output data type.
bool is_typecast(const op_t *op);

bool with_runtime_scales(
        const std::shared_ptr<op_t> &op, bool is_input, size_t index);

bool with_runtime_zps(
        const std::shared_ptr<op_t> &op, bool is_input, size_t index);

// This function is used to check if a dnnl_reorder op is converted from or act
// as a Reorder op. This function will only return true for a dnnl_reorder op
// which is not for TypeCast or Quantization.
bool is_layout_reorder(const op_t *op);

// This function is used to clone a dnnl_mul_scales op
std::shared_ptr<op_t> clone_mul_scales(const std::shared_ptr<op_t> &scale_op);

// This function is used to inverse scales of a dnnl_mul_scales op
bool inverse_mul_scales(std::shared_ptr<op_t> &scale_op);

bool need_broadcast_for_inputs(
        const std::shared_ptr<op_t> &op, size_t index1, size_t index2);

} // namespace dnnl_impl
} // namespace graph
} // namespace impl
} // namespace dnnl

#endif