/*
// Copyright (c) 2019 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.
*/

#include "jitter.h"
#include "tensor_type.h"
#include <string>
#include <memory>
#include <utility>

#include <quantize/quantize_kernel_params.h>
#include <eltwise/eltwise_kernel_base.h>
#include <activation/activation_kernel_base.h>

namespace {
class JitTerm {
public:
    explicit JitTerm(std::string text)
        : text(std::move(text)) {}

    std::string str() const { return text; }

    JitTerm gt(const JitTerm& rhs) const {
        JitTerm jit_term { "(" + text + ">" + rhs.str() + ")" };
        return jit_term;
    }

    JitTerm ge(const JitTerm& rhs) const {
        JitTerm jit_term {"(" + text + ">=" + rhs.str() + ")"};
        return jit_term;
    }

    JitTerm le(const JitTerm& rhs) const {
        JitTerm jit_term {"(" + text + "<=" + rhs.str() + ")"};
        return jit_term;
    }

    JitTerm eq(const JitTerm& rhs) const {
        JitTerm jit_term {"(" + text + "==" + rhs.str() + ")"};
        return jit_term;
    }

private:
    std::string text;
};

JitTerm operator+(const JitTerm& lhs, const JitTerm& rhs) {
    JitTerm jit_term{"(" + lhs.str() + " + " + rhs.str() + ")"};
    return jit_term;
}

JitTerm operator-(const JitTerm& lhs, const JitTerm& rhs) {
    JitTerm jit_term{"(" + lhs.str() + " - " + rhs.str() + ")"};
    return jit_term;
}

JitTerm operator*(const JitTerm& lhs, const JitTerm& rhs) {
    JitTerm jit_term{"(" + lhs.str() + " * " + rhs.str() + ")"};
    return jit_term;
}

JitTerm operator/(const JitTerm& lhs, const JitTerm& rhs) {
    JitTerm jit_term{"(" + lhs.str() + " / " + rhs.str() + ")"};
    return jit_term;
}

JitTerm neg(const JitTerm& arg) {
    JitTerm jit_term{"(-" + arg.str() + ")"};
    return jit_term;
}

JitTerm ternary(const JitTerm& condition,
                const JitTerm& true_expr,
                const JitTerm& false_expr) {
    JitTerm jit_term{"(" + condition.str() + " ? " + true_expr.str() + " : " + false_expr.str() + ")"};
    return jit_term;
}
JitTerm isinf(const JitTerm& arg) {
    JitTerm jit_term{"(isinf(" + arg.str() + "))"};
    return jit_term;
}

JitTerm exp(const JitTerm& arg) {
    JitTerm jit_term{"(exp(" + arg.str() + "))"};
    return jit_term;
}

JitTerm log(const JitTerm& arg) {
    JitTerm jit_term{"(log(" + arg.str() + "))"};
    return jit_term;
}

JitTerm operator"" _jit(const char* str, size_t) {
    JitTerm jit_term{str};
    return jit_term;
}
}  // namespace

namespace kernel_selector {

std::string toCLType(WeightsType wType) {
    switch (wType) {
        case WeightsType::BINARY:
            return GetTypeName<uint32_t>();
        case WeightsType::INT8:
            return GetTypeName<int8_t>();
        case WeightsType::UINT8:
            return GetTypeName<uint8_t>();
        case WeightsType::F16:
            return "half";
        case WeightsType::F32:
            return GetTypeName<float>();
        default:
            return "";
    }
}

std::string toCLType(Datatype dType) {
    switch (dType) {
        case Datatype::BINARY:
            return GetTypeName<uint32_t>();
        case Datatype::INT8:
            return GetTypeName<int8_t>();
        case Datatype::UINT8:
            return GetTypeName<uint8_t>();
        case Datatype::INT16:
            return GetTypeName<int16_t>();
        case Datatype::UINT16:
            return GetTypeName<uint16_t>();
        case Datatype::INT32:
            return GetTypeName<int32_t>();
        case Datatype::UINT32:
            return GetTypeName<uint32_t>();
        case Datatype::INT64:
            return GetTypeName<int64_t>();
        case Datatype::F16:
            return "half";
        case Datatype::F32:
            return GetTypeName<float>();
        default:
            return "";
    }
}

std::string getMeanOpString(MeanOp op) {
    switch (op) {
        case MeanOp::NONE:
            return "val";
        case MeanOp::DIV:
            return "val/mean_val";
        case MeanOp::MUL:
            return "val*mean_val";
        case MeanOp::SUB:
            return "val-mean_val";
        default:
            return "";
    }
}

std::string toCodeString(float val) {
    if (std::isinf(val))
        return std::signbit(val) ? "-INFINITY" : "INFINITY";
    std::stringstream ss;
    // Workaround GCC compiler/STL bug
    ss << "as_float(0x" << std::hex << *reinterpret_cast<uint32_t*>(&val) << ")";

    ss << " /*" << std::scientific << val << "*/";
    return ss.str();
}

std::string toCodeString(double val) {
    if (std::isinf(val))
        return std::signbit(val) ? "-INFINITY" : "INFINITY";
    std::stringstream ss;
    // Workaround GCC compiler/STL bug
    ss << "as_double(0x" << std::hex << *reinterpret_cast<uint64_t*>(&val) << ")";

    ss << " /*" << std::scientific << val << "*/";
    return ss.str();
}

JitDefinitions JitConstants::GetDefinitions() const {
    JitDefinitions definitons;
    definitons.reserve(_constants.size() * 6);  // assuming max 6 pairs per jit_constant

    for (auto& constant : _constants) {
        auto def = constant->GetDefinitions();
        definitons.insert(definitons.end(), def.begin(), def.end());
    }
    return definitons;
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
// TensorBaseTJitConstant
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <typename DType, typename Layout>
class TensorBaseTJitConstant : public JitConstant {
protected:
    explicit TensorBaseTJitConstant(const std::string& name) : JitConstant(name) {}

public:
    JitDefinitions GetDefinitions(const Tensor::TensorBaseT<DType, Layout>& t) const {
        JitDefinitions definitions{
            {_name + "_OFFSET", toCodeString(t.GetFirstElementOffset())},
            {_name + "_VIEW_OFFSET", toCodeString(t.GetViewOffset())},
            {_name + "_LENGTH", toCodeString(t.LogicalSize())},
            {_name + "_DIMS", toCodeString(t.GetDims().size())},
            {_name + "_SIMPLE", toCodeString(t.SimpleLayout())},
            {_name + "_GROUPED", toCodeString(t.GroupedLayout())},
            {_name + "_LAYOUT_" + toString(t.GetLayout()), "1"},
        };

        auto type_defs = MakeTypeJitConstants(t.GetDType(), _name).GetDefinitions();
        definitions.insert(definitions.end(), type_defs.begin(), type_defs.end());

        definitions.push_back({_name + "_SIZE", toCodeString(t.GetDims().size())});
        definitions.push_back(
            {_name + "_SIZES",
             toVectorString(t.GetDims(), "size_t", KERNEL_SELECTOR_TENSOR_DIM_MAX, 1, [](const Tensor::Dim& d) { return d.v; })});
        definitions.push_back(
            {_name + "_PITCHES",
             toVectorString(t.GetDims(), "size_t", KERNEL_SELECTOR_TENSOR_DIM_MAX, 1, [](const Tensor::Dim& d) { return d.pitch; })});
        definitions.push_back(
            {_name + "_PAD_BEFORE",
             toVectorString(t.GetDims(), "size_t", KERNEL_SELECTOR_TENSOR_DIM_MAX, 0, [](const Tensor::Dim& d) { return d.pad.before; })});
        definitions.push_back(
            {_name + "_PAD_AFTER",
             toVectorString(t.GetDims(), "size_t", KERNEL_SELECTOR_TENSOR_DIM_MAX, 0, [](const Tensor::Dim& d) { return d.pad.after; })});

        return definitions;
    }
};

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
// DataTensorJitConstant
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
class DataTensorJitConstant : public TensorBaseTJitConstant<Datatype, DataLayout> {
    const DataTensor _tensor;

public:
    DataTensorJitConstant(const std::string& name, const DataTensor& t) : TensorBaseTJitConstant(name), _tensor(t) {}

    JitDefinitions GetDefinitions() const override;
};

JitDefinitions DataTensorJitConstant::GetDefinitions() const {
    JitDefinitions baseDefinitions = TensorBaseTJitConstant::GetDefinitions(_tensor);

    JitDefinitions definitions{
        {_name + "_SIZE_X", toCodeString(_tensor.X().v)},
        {_name + "_SIZE_Y", toCodeString(_tensor.Y().v)},
        {_name + "_SIZE_Z", toCodeString(_tensor.Z().v)},
        {_name + "_SIZE_W", toCodeString(_tensor.W().v)},
        {_name + "_FEATURE_NUM", toCodeString(_tensor.Feature().v)},
        {_name + "_BATCH_NUM", toCodeString(_tensor.Batch().v)},
        {_name + "_X_PITCH", toCodeString(_tensor.X().pitch)},
        {_name + "_Y_PITCH", toCodeString(_tensor.Y().pitch)},
        {_name + "_Z_PITCH", toCodeString(_tensor.Z().pitch)},
        {_name + "_W_PITCH", toCodeString(_tensor.W().pitch)},
        {_name + "_FEATURE_PITCH", toCodeString(_tensor.Feature().pitch)},
        {_name + "_BATCH_PITCH", toCodeString(_tensor.Batch().pitch)},
        {_name + "_PAD_BEFORE_SIZE_X", toCodeString(_tensor.X().pad.before)},
        {_name + "_PAD_BEFORE_SIZE_Y", toCodeString(_tensor.Y().pad.before)},
        {_name + "_PAD_BEFORE_SIZE_Z", toCodeString(_tensor.Z().pad.before)},
        {_name + "_PAD_BEFORE_SIZE_W", toCodeString(_tensor.W().pad.before)},
        {_name + "_PAD_BEFORE_FEATURE_NUM", toCodeString(_tensor.Feature().pad.before)},
        {_name + "_PAD_BEFORE_BATCH_NUM", toCodeString(_tensor.Batch().pad.before)},
        {_name + "_PAD_AFTER_SIZE_X", toCodeString(_tensor.X().pad.after)},
        {_name + "_PAD_AFTER_SIZE_Y", toCodeString(_tensor.Y().pad.after)},
        {_name + "_PAD_AFTER_SIZE_Z", toCodeString(_tensor.Z().pad.after)},
        {_name + "_PAD_AFTER_SIZE_W", toCodeString(_tensor.W().pad.after)},
        {_name + "_PAD_AFTER_FEATURE_NUM", toCodeString(_tensor.Feature().pad.after)},
        {_name + "_PAD_AFTER_BATCH_NUM", toCodeString(_tensor.Batch().pad.after)},
    };

    auto is_common_nd_layout = [](std::vector<Tensor::DataChannelName> common_channels, DataLayout l) -> bool {
        for (size_t c = 0; c < static_cast<size_t>(Tensor::DataChannelName::COUNT); c++) {
            auto channel = static_cast<Tensor::DataChannelName>(c);
            if (DataTensor::Channelndex(l, channel) != -1) {
                if (std::find(common_channels.begin(), common_channels.end(), channel) == common_channels.end()) {
                    return false;
                }
            }
        }
        return true;
    };

    std::string index_func_name;
    std::string safe_index_func_name;
    std::string raw_index_func_name;
    std::string index_func_val;
    std::string safe_index_func_val;
    std::string raw_index_func_val;

    // TODO: add support for other layouts
    auto layout = _tensor.GetLayout();
    if (DataTensor::ChannelsCount(layout) <= 4) {
        std::vector<Tensor::DataChannelName> base_4d_channels = {
                Tensor::DataChannelName::BATCH,
                Tensor::DataChannelName::FEATURE,
                Tensor::DataChannelName::Y,
                Tensor::DataChannelName::X,
        };
        bool is_common_4d_layout = is_common_nd_layout(base_4d_channels, layout);
        if (is_common_4d_layout) {
            index_func_name = _name + "_GET_INDEX(b, f, y, x)";
            safe_index_func_name = _name + "_GET_INDEX_SAFE(b, f, y, x)";
            raw_index_func_name = _name + "_GET_INDEX_RAW(b, f, y, x)";

            if (_tensor.SimpleLayout()) {
                index_func_val = "GET_DATA_INDEX(" + _name + ", b, f, y, x)";
                safe_index_func_val = "GET_DATA_INDEX_SAFE(" + _name + ", b, f, y, x)";
                raw_index_func_val = "GET_DATA_INDEX_RAW(" + _name + ", b, f, y, x)";
            } else if (layout == DataLayout::b_fs_yx_fsv16 ||
                       layout == DataLayout::b_fs_yx_fsv32 ||
                       layout == DataLayout::byxf_af32 ||
                       layout == DataLayout::fs_bs_yx_bsv4_fsv32 ||
                       layout == DataLayout::b_fs_yx_fsv4 ||
                       layout == DataLayout::fs_b_yx_fsv32 ||
                       layout == DataLayout::bs_fs_yx_bsv16_fsv16) {
                auto layout_str = toString(layout);
                index_func_val = "GET_DATA_" + layout_str + "_INDEX(" + _name + ", b, f, y, x)";
                raw_index_func_val = "GET_DATA_" + layout_str + "_INDEX(" + _name + ", b, f, y, x)";
                if (layout == DataLayout::b_fs_yx_fsv16 || layout == DataLayout::b_fs_yx_fsv32 || layout == DataLayout::bs_fs_yx_bsv16_fsv16)
                    safe_index_func_val = "GET_DATA_" + layout_str + "_INDEX_SAFE(" + _name + ", b, f, y, x)";
                else
                    safe_index_func_val = "GET_DATA_" + layout_str + "_INDEX(" + _name + ", b, f, y, x)";
            } else if (layout == DataLayout::bs_f_bsv8__af8 ||
                       layout == DataLayout::bs_f_bsv16__af8) {
                size_t sub_group_size = layout == DataLayout::bs_f_bsv16__af8 ? 16 : 8;
                index_func_val = "GET_DATA_BS_FYX_BSV8_INDEX(" + _name + ", b, f, y, x," + std::to_string(sub_group_size) + ")";
                safe_index_func_val = "GET_DATA_BS_FYX_BSV8_INDEX(" + _name + ", b, f, y, x," + std::to_string(sub_group_size) + ")";
                raw_index_func_val = "GET_DATA_BS_FYX_BSV8_INDEX(" + _name + ", b, f, y, x," + std::to_string(sub_group_size) + ")";
            } else {
                index_func_val =  "GET_DATA_INDEX_RAW(" + _name + ", b, f, y, x)";
                safe_index_func_val = "GET_DATA_INDEX_RAW(" + _name + ", b, f, y, x)";
                raw_index_func_val = "GET_DATA_INDEX_RAW(" + _name + ", b, f, y, x)";
            }
        } else {
            // TODO: implement support of non-default layouts with 4 channels
            assert(0);
        }
    } else if (DataTensor::ChannelsCount(layout) == 5) {
        std::vector<Tensor::DataChannelName> base_5d_channels = {
                Tensor::DataChannelName::BATCH,
                Tensor::DataChannelName::FEATURE,
                Tensor::DataChannelName::Z,
                Tensor::DataChannelName::Y,
                Tensor::DataChannelName::X,
        };
        bool is_common_5d_layout = is_common_nd_layout(base_5d_channels, layout);
        if (is_common_5d_layout) {
            index_func_name = _name + "_GET_INDEX(b, f, z, y, x)";
            safe_index_func_name = _name + "_GET_INDEX_SAFE(b, f, z, y, x)";
            raw_index_func_name = _name + "_GET_INDEX_RAW(b, f, z, y, x)";

            if (_tensor.SimpleLayout()) {
                index_func_val = "GET_DATA_INDEX_5D(" + _name + ", b, f, z, y, x)";
                safe_index_func_val = "GET_DATA_INDEX_5D_SAFE("+ _name + ", b, f, z, y, x)";
                raw_index_func_val = "GET_DATA_INDEX_5D_RAW("+ _name + ", b, f, z, y, x)";
            } else if (layout == DataLayout::b_fs_zyx_fsv16) {
                index_func_val = "GET_DATA_B_FS_ZYX_FSV16_INDEX(" + _name + ", b, f, z, y, x)";
                raw_index_func_val = "GET_DATA_B_FS_ZYX_FSV16_INDEX(" + _name + ", b, f, z, y, x)";
                safe_index_func_val = "GET_DATA_B_FS_ZYX_FSV16_INDEX_SAFE(" + _name + ", b, f, z, y, x)";
            } else if (layout == DataLayout::bs_fs_zyx_bsv16_fsv16) {
                index_func_val = "GET_DATA_BS_FS_ZYX_BSV16_FSV16_INDEX(" + _name + ", b, f, z, y, x)";
                raw_index_func_val = "GET_DATA_BS_FS_ZYX_BSV16_FSV16_INDEX(" + _name + ", b, f, z, y, x)";
                safe_index_func_val = "GET_DATA_BS_FS_ZYX_BSV16_FSV16_INDEX_SAFE(" + _name + ", b, f, z, y, x)";
            } else if (layout == DataLayout::b_fs_zyx_fsv32) {
                index_func_val = "GET_DATA_B_FS_ZYX_FSV32_INDEX(" + _name + ", b, f, z, y, x)";
                raw_index_func_val = "GET_DATA_B_FS_ZYX_FSV32_INDEX(" + _name + ", b, f, z, y, x)";
                safe_index_func_val = "GET_DATA_B_FS_ZYX_FSV32_INDEX_SAFE(" + _name + ", b, f, z, y, x)";
            } else {
                index_func_val = "GET_DATA_INDEX_5D_RAW(" + _name + ", b, f, z, y, x)";
                safe_index_func_val = "GET_DATA_INDEX_5D_RAW(" + _name + ", b, f, z, y, x)";
                raw_index_func_val = "GET_DATA_INDEX_5D_RAW(" + _name + ", b, f, z, y, x)";
            }
        } else {
            // TODO: implement support of non-default layouts with 5 channels
            assert(0);
        }
    } else if (DataTensor::ChannelsCount(layout) == 6) {
        std::vector<Tensor::DataChannelName> base_6d_channels = {
                Tensor::DataChannelName::BATCH,
                Tensor::DataChannelName::FEATURE,
                Tensor::DataChannelName::W,
                Tensor::DataChannelName::Z,
                Tensor::DataChannelName::Y,
                Tensor::DataChannelName::X,
        };
        bool is_common_6d_layout = is_common_nd_layout(base_6d_channels, layout);
        if (is_common_6d_layout) {
            index_func_name = _name + "_GET_INDEX(b, f, w, z, y, x)";
            safe_index_func_name = _name + "_GET_INDEX_SAFE(b, f, w, z, y, x)";
            raw_index_func_name = _name + "_GET_INDEX_RAW(b, f, w, z, y, x)";

            index_func_val = "GET_DATA_INDEX_6D(" + _name + ", b, f, w, z, y, x)";
            safe_index_func_val = "GET_DATA_INDEX_6D_SAFE(" + _name + ", b, f, w, z, y, x)";
            raw_index_func_val = "GET_DATA_INDEX_6D_RAW(" + _name + ", b, f, w, z, y, x)";
        } else {
            // TODO: implement support of non-default layouts with 6 channels
            assert(0);
        }
    } else {
        throw std::runtime_error("Unsupported channels count(" + std::to_string(DataTensor::ChannelsCount(layout)) +
                                 ") in layout: " + toString(layout));
    }

    std::string offset = std::to_string(_tensor.GetFirstElementOffset());
    if (_tensor.LogicalSize() == 1) {
        // if tensor contains single element we can always return 0 for safe function
        if (_tensor.PitchesDifferFromLogicalDims()) {
            definitions.push_back({ safe_index_func_name, offset });
            definitions.push_back({ index_func_name, offset });
        } else {
            definitions.push_back({ safe_index_func_name, "0" });
            definitions.push_back({ index_func_name, "0" });
        }
    } else if (_tensor.LogicalSize() == _tensor.Feature().v) {
        // We support broadcast only if corresponding dimension is equal to 1.
        // Otherwise, dimensions should be equal and using "f" should be safe.
        if (_tensor.PitchesDifferFromLogicalDims()) {
            std::string f_pitch = std::to_string(_tensor.Feature().pitch);
            definitions.push_back({ safe_index_func_name, "(" + offset + " + (f) * " + f_pitch + ")" });
            definitions.push_back({ index_func_name, "(" + offset + " + (f) * " + f_pitch + ")" });
        } else {
            definitions.push_back({ safe_index_func_name, "f" });
            definitions.push_back({ index_func_name, "f" });
        }
    } else {
        definitions.push_back({ safe_index_func_name, safe_index_func_val });
        definitions.push_back({ index_func_name, index_func_val });
    }
    definitions.push_back({ raw_index_func_name, raw_index_func_val });

    definitions.insert(definitions.end(), baseDefinitions.begin(), baseDefinitions.end());

    return definitions;
}

std::shared_ptr<JitConstant> MakeJitConstant(const std::string& name, const DataTensor& value) {
    return std::static_pointer_cast<JitConstant>(std::make_shared<DataTensorJitConstant>(name, value));
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
// WeightTensorJitConstant
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
class WeightTensorJitConstant : public TensorBaseTJitConstant<WeightsType, WeightsLayout> {
    const WeightsTensor _tensor;

public:
    WeightTensorJitConstant(const std::string& name, const WeightsTensor& t) : TensorBaseTJitConstant(name), _tensor(t) {}

    JitDefinitions GetDefinitions() const override;
};

JitDefinitions WeightTensorJitConstant::GetDefinitions() const {
    JitDefinitions baseDefinitions = TensorBaseTJitConstant::GetDefinitions(_tensor);

    JitDefinitions definitions{
        {_name + "_SIZE_X", toCodeString(_tensor.X().v)},
        {_name + "_SIZE_Y", toCodeString(_tensor.Y().v)},
        {_name + "_SIZE_Z", toCodeString(_tensor.Z().v)},
        {_name + "_IFM_NUM", toCodeString(_tensor.IFM().v)},
        {_name + "_OFM_NUM", toCodeString(_tensor.OFM().v)},
        {_name + "_GROUPS_NUM", toCodeString(_tensor.G().v)},
        {_name + "_X_PITCH", toCodeString(_tensor.X().pitch)},
        {_name + "_Y_PITCH", toCodeString(_tensor.Y().pitch)},
        {_name + "_Z_PITCH", toCodeString(_tensor.Z().pitch)},
        {_name + "_IFM_PITCH", toCodeString(_tensor.IFM().pitch)},
        {_name + "_OFM_PITCH", toCodeString(_tensor.OFM().pitch)},
        {_name + "_GROUPS_PITCH", toCodeString(_tensor.G().pitch)},
    };

    definitions.insert(definitions.end(), baseDefinitions.begin(), baseDefinitions.end());

    return definitions;
}

std::shared_ptr<JitConstant> MakeJitConstant(const std::string& name, const WeightsTensor& value) {
    return std::static_pointer_cast<JitConstant>(std::make_shared<WeightTensorJitConstant>(name, value));
}

JitConstants MakeActivationJitConstants(ActivationFunction activation_function,
                                        Datatype out_dt,
                                        const std::string& suffix,
                                        bool use_type_parameter,
                                        bool disable_type_conversion) {
    std::string name = "ACTIVATION_FUNC" + suffix;
    JitConstants jitConstants = {};

    jitConstants.Merge(MakeTypeJitConstants(out_dt, name));
    // See the comment in the jitter.h regarding `use_type_parameter`.
    // The "CAT" macro is expected to be defined through the inlcusion of
    // 'common.cl' in the kernel.
    auto type_handler =
        [use_type_parameter, name](const std::string& prefix,
                                   const std::string& suffix) -> std::string {
        if (!use_type_parameter)
            return prefix + name + suffix;

        std::string result = "jit_type";

        // Process the prefix first, otherwise when doing "CAT(TO_,
        // CAT(NAME, _TYPE))" the second concatenation will be expanded
        // fully first resulting in something like "TO_float".
        if (!prefix.empty())
            result = "CAT(" + prefix + ", " + result + ")";

        if (!suffix.empty())
            result = "CAT(" + result + ", " + suffix + ")";

        return result;
    };

    const JitTerm one{type_handler("", "_VAL_ONE")};
    const JitTerm zero{type_handler("", "_VAL_ZERO")};
    const JitTerm input{"input"};
    auto max_func = [type_handler](const JitTerm& lhs,
                                   const JitTerm& rhs) -> JitTerm {
        JitTerm jit_term{"(" + type_handler("", "_MAX_FUNC") + "(" + lhs.str() + ", " + rhs.str() + "))"};
        return jit_term;
    };
    auto min_func = [type_handler](const JitTerm& lhs,
                                   const JitTerm& rhs) -> JitTerm {
        JitTerm jit_term{"(" + type_handler("", "_MIN_FUNC") + "(" + lhs.str() + ", " + rhs.str() + "))"};
        return jit_term;
    };
    auto abs_func = [type_handler](const JitTerm& hs) -> JitTerm {
        JitTerm jit_term{"(" + type_handler("", "_ABS_FUNC") + "(" + hs.str() + "))"};
        return jit_term;
    };
    auto to_type = [type_handler](const JitTerm& arg) -> JitTerm {
        JitTerm jit_term{type_handler("TO_", "_TYPE") + "(" + arg.str() + ")"};
        return jit_term;
    };

    std::string macro_def = name + (use_type_parameter ? "(jit_type, input, m, n)" : "(input, m, n)");
    std::string macro_def_grad = name + (use_type_parameter ? "(jit_type, input_grad, input, m, n)"
                                                            : "(input_grad, input, m, n)");
    std::string macro_def_params = use_type_parameter ? "(jit_type, input, params)" : "(input, params)";

    jitConstants.AddConstant(MakeJitConstant("ACTIVATION_PARAMS" + suffix, "NL_M" + suffix + ", NL_N" + suffix));

    // TODO: use native_exp and use cast for APL
    switch (activation_function) {
        case ActivationFunction::LOGISTIC:
            jitConstants.AddConstant(MakeJitConstant(macro_def, (one / (one + exp(neg(input)))).str()));
            break;
        case ActivationFunction::HYPERBOLIC_TAN:
            jitConstants.AddConstant(MakeJitConstant(macro_def, "(tanh(input))"));
            break;
        case ActivationFunction::RELU:
            jitConstants.AddConstant(MakeJitConstant(macro_def, max_func(zero, input).str()));
            break;
        case ActivationFunction::RELU_NEGATIVE_SLOPE: {
            const JitTerm slope = disable_type_conversion ? "m"_jit : to_type("m"_jit);
            jitConstants.AddConstant(MakeJitConstant(
                macro_def,
                ternary(isinf(slope),
                        ternary(input.ge(zero), input, neg(slope)),
                        max_func(input, zero) + (slope * min_func(input, zero)))
                    .str()));
            break;
        }
        case ActivationFunction::ELU: {
            auto alpha = disable_type_conversion ? "m"_jit : to_type("m"_jit);
            jitConstants.AddConstant(MakeJitConstant(
                macro_def,
                (max_func(input, zero) + (alpha * (exp(min_func(input, zero)) - one)))
                    .str()));
            break;
        }
        case ActivationFunction::CLAMP: {
            const JitTerm m = disable_type_conversion ? "m"_jit : to_type("m"_jit);
            const JitTerm n = disable_type_conversion ? "n"_jit : to_type("n"_jit);
            jitConstants.AddConstant(MakeJitConstant(
                 macro_def,
                 max_func(m, min_func(n, input)).str()));
            break;
        }
        case ActivationFunction::SOFTRELU:
            jitConstants.AddConstant(MakeJitConstant(macro_def, log(one + exp(input)).str()));
            break;
        case ActivationFunction::ABS:
            jitConstants.AddConstant(MakeJitConstant(macro_def, "(fabs(input))"));
            break;
        case ActivationFunction::LINEAR: {
            const JitTerm m = disable_type_conversion ? "m"_jit : to_type("m"_jit);
            const JitTerm n = disable_type_conversion ? "n"_jit : to_type("n"_jit);
            jitConstants.AddConstant(MakeJitConstant(macro_def, (m * input + n).str()));
            break;
        }
        case ActivationFunction::SQUARE:
            jitConstants.AddConstant(MakeJitConstant(macro_def, "(input*input)"));
            break;
        case ActivationFunction::SQRT:
            jitConstants.AddConstant(MakeJitConstant(macro_def, "(sqrt(input))"));
            break;
        case ActivationFunction::SIN:
            jitConstants.AddConstant(MakeJitConstant(macro_def, "(sin(input))"));
            break;
        case ActivationFunction::ASIN:
            jitConstants.AddConstant(MakeJitConstant(macro_def, "(asin(input))"));
            break;
        case ActivationFunction::SINH:
            jitConstants.AddConstant(MakeJitConstant(macro_def, "(sinh(input))"));
            break;
        case ActivationFunction::ASINH:
            jitConstants.AddConstant(MakeJitConstant(macro_def, "(asinh(input))"));
            break;
        case ActivationFunction::COS:
            jitConstants.AddConstant(MakeJitConstant(macro_def, "(cos(input))"));
            break;
        case ActivationFunction::ACOS:
            jitConstants.AddConstant(MakeJitConstant(macro_def, "(acos(input))"));
            break;
        case ActivationFunction::COSH:
            jitConstants.AddConstant(MakeJitConstant(macro_def, "(cosh(input))"));
            break;
        case ActivationFunction::ACOSH:
            jitConstants.AddConstant(MakeJitConstant(macro_def, "(acosh(input))"));
            break;
        case ActivationFunction::LOG:
            jitConstants.AddConstant(MakeJitConstant(macro_def, "(log(input))"));
            break;
        case ActivationFunction::LOG2:
            jitConstants.AddConstant(MakeJitConstant(macro_def, "(log2(input))"));
            break;
        case ActivationFunction::EXP:
            jitConstants.AddConstant(MakeJitConstant(macro_def, "(exp(input))"));
            break;
        case ActivationFunction::POW: {
            const JitTerm m = disable_type_conversion ? "m"_jit : to_type("m"_jit);
            jitConstants.AddConstant(MakeJitConstant(macro_def, "(pow(input," + m.str() + "))"));
            break;
        }
        case ActivationFunction::RELU_GRAD:
            jitConstants.AddConstant(MakeJitConstant(
                macro_def_grad,
                ("input_grad"_jit * ternary(input.gt(zero), one, zero)).str()));
            macro_def_params = use_type_parameter ? "(jit_type, input_grad, input, params)" : "(input_grad, input, params)";
            break;
        case ActivationFunction::RELU_NEGATIVE_SLOPE_GRAD: {
            const JitTerm slope = disable_type_conversion ? "m"_jit : to_type("m"_jit);
            jitConstants.AddConstant(MakeJitConstant(
                macro_def_grad,
                ("input_grad"_jit * (ternary(input.gt(zero), one, zero) + (to_type(slope) * ternary(input.le(zero), one, zero))))
                    .str()));
            macro_def_params = use_type_parameter ? "(jit_type, input_grad, input, params)" : "(input_grad, input, params)";
            break;
        }
        case ActivationFunction::NONE_GRAD:
            jitConstants.AddConstant(MakeJitConstant(macro_def_grad, "input_grad"));
            macro_def_params = use_type_parameter ? "(jit_type, input_grad, input, params)" : "(input_grad, input, params)";
            break;
        case ActivationFunction::TAN:
            jitConstants.AddConstant(MakeJitConstant(macro_def, "(tan(input))"));
            break;
        case ActivationFunction::ATAN:
            jitConstants.AddConstant(MakeJitConstant(macro_def, "(atan(input))"));
            break;
        case ActivationFunction::ATANH:
            jitConstants.AddConstant(MakeJitConstant(macro_def, "(atanh(input))"));
            break;
        case ActivationFunction::FLOOR:
            jitConstants.AddConstant(MakeJitConstant(macro_def, "(floor(input))"));
            break;
        case ActivationFunction::CEIL:
            jitConstants.AddConstant(MakeJitConstant(macro_def, "(ceil(input))"));
            break;
        case ActivationFunction::NEGATIVE:
            jitConstants.AddConstant(MakeJitConstant(macro_def, "(-input)"));
            break;
        case ActivationFunction::ERF:
            jitConstants.AddConstant(MakeJitConstant(macro_def, "erf(input)"));
            break;
        case ActivationFunction::HARD_SIGMOID: {
            auto alpha = disable_type_conversion ? "m"_jit : to_type("m"_jit);
            auto beta =  disable_type_conversion ? "n"_jit : to_type("n"_jit);
            jitConstants.AddConstant(MakeJitConstant(
                    macro_def,
                    max_func(zero, min_func(one, (JitTerm)((alpha * input + beta).str()))).str()));
            break;
        }
        case ActivationFunction::SIGN:
            jitConstants.AddConstant(MakeJitConstant(
                    macro_def,
                    ternary(input.gt(zero), one, ternary(input.eq(zero), zero, neg(one))).str()));
            break;
        case ActivationFunction::RECIPROCAL:
            jitConstants.AddConstant(MakeJitConstant(macro_def, (one / input).str()));
            break;
        case ActivationFunction::SELU: {
            auto alpha = disable_type_conversion ? "m"_jit : to_type("m"_jit);
            auto gamma = disable_type_conversion ? "n"_jit : to_type("n"_jit);
            jitConstants.AddConstant(MakeJitConstant(
                    macro_def,
                    ternary(input.le(zero), gamma * (alpha * exp(input) - alpha), gamma * input).str()));
            break;
        }
        case ActivationFunction::SOFTPLUS: {
            jitConstants.AddConstant(MakeJitConstant(
                    macro_def,
                    log(exp(input) + one).str()));
            break;
        }
        case ActivationFunction::SOFTSIGN: {
            jitConstants.AddConstant(MakeJitConstant(
                    macro_def,
                    (input / (one + abs_func(input))).str()));
            break;
        }
        case ActivationFunction::SWISH: {
            jitConstants.AddConstant(MakeJitConstant(
                    macro_def,
                    (input / (one + exp(neg(input)))).str()));
            break;
        }
        case ActivationFunction::NOT:
            jitConstants.AddConstant(MakeJitConstant(
                macro_def,
                ternary(input.eq(zero), one, zero)
                    .str()));  // the workaround for OpenCL's vector type result (!input)
            break;
        case ActivationFunction::NONE:
        default:
            jitConstants.AddConstant(MakeJitConstant(macro_def, "input"));
            break;
    }

    jitConstants.AddConstant(MakeJitConstant("ACTIVATION" + suffix + macro_def_params, name + macro_def_params));

    return jitConstants;
}

JitConstants MakeTypeJitConstants(Datatype dataType, const std::string& macroName) {
    std::string type;
    std::string max_val;
    std::string min_val;
    std::string val_one;
    std::string val_zero;
    std::string to_type;
    std::string to_type_sat;
    std::string as_type;
    std::string max_func;
    std::string min_func;
    std::string abs_func;
    std::string type_size;
    bool is_fp;
    switch (dataType) {
        case Datatype::INT8:
            type = "char";
            max_val = "CHAR_MAX";
            min_val = "CHAR_MIN";
            val_one = "(char) 1";
            val_zero = "(char) 0";
            to_type = "convert_char(v)";
            to_type_sat = "convert_char_sat(v)";
            as_type = "as_char(v)";
            max_func = "max";
            min_func = "min";
            abs_func = "abs";
            type_size = "1";
            is_fp = false;
            break;
        case Datatype::UINT8:
            type = "uchar";
            max_val = "UCHAR_MAX";
            min_val = "0";
            val_one = "(uchar) 1";
            val_zero = "(uchar) 0";
            to_type = "convert_uchar(v)";
            to_type_sat = "convert_uchar_sat(v)";
            as_type = "as_uchar(v)";
            max_func = "max";
            min_func = "min";
            abs_func = "abs";
            type_size = "1";
            is_fp = false;
            break;
        case Datatype::INT16:
            type = "short";
            max_val = "SHRT_MAX";
            min_val = "SHRT_MIN";
            val_one = "(short) 1";
            val_zero = "(short) 0";
            to_type = "convert_short(v)";
            to_type_sat = "convert_short_sat(v)";
            as_type = "as_short(v)";
            max_func = "max";
            min_func = "min";
            abs_func = "abs";
            type_size = "2";
            is_fp = false;
            break;
        case Datatype::UINT16:
            type = "ushort";
            max_val = "USHRT_MAX";
            min_val = "0";
            val_one = "(ushort) 1";
            val_zero = "(ushort) 0";
            to_type = "convert_ushort(v)";
            to_type_sat = "convert_ushort_sat(v)";
            as_type = "as_ushort(v)";
            max_func = "max";
            min_func = "min";
            abs_func = "abs";
            type_size = "2";
            is_fp = false;
            break;
        case Datatype::INT32:
            type = "int";
            max_val = "INT_MAX";
            min_val = "INT_MIN";
            val_one = "(int) 1";
            val_zero = "(int) 0";
            to_type = "convert_int(v)";
            to_type_sat = "convert_int_sat(v)";
            as_type = "as_int(v)";
            max_func = "max";
            min_func = "min";
            abs_func = "abs";
            type_size = "4";
            is_fp = false;
            break;
        case Datatype::UINT32:
        case Datatype::BINARY:
            type = "uint";
            max_val = "UINT_MAX";
            min_val = "0";
            val_one = "(uint) 1";
            val_zero = "(uint) 0";
            to_type = "convert_uint(v)";
            to_type_sat = "convert_uint_sat(v)";
            as_type = "as_uint(v)";
            max_func = "max";
            min_func = "min";
            abs_func = "abs";
            type_size = "4";
            is_fp = false;
            break;
        case Datatype::INT64:
            type = "long";
            max_val = "LONG_MAX";
            min_val = "LONG_MIN";
            val_one = "(long) 1";
            val_zero = "(long) 0";
            to_type = "convert_long(v)";
            to_type_sat = "convert_long_sat(v)";
            as_type = "as_long(v)";
            max_func = "max";
            min_func = "min";
            abs_func = "abs";
            type_size = "8";
            is_fp = false;
            break;
        case Datatype::F16:
            type = "half";
            max_val = "HALF_MAX";
            min_val = "-" + macroName + "_VAL_MAX";
            val_one = "1.0h";
            val_zero = "0.0h";
            to_type = "convert_half(v)";
            to_type_sat = "convert_half(v)";
            as_type = "as_half(v)";
            max_func = "fmax";
            min_func = "fmin";
            abs_func = "fabs";
            type_size = "2";
            is_fp = true;
            break;
        default:
            type = "float";
            max_val = "FLT_MAX";
            min_val = "-" + macroName + "_VAL_MAX";
            val_one = "1.0f";
            val_zero = "0.0f";
            to_type = "convert_float(v)";
            to_type_sat = "convert_float(v)";
            as_type = "as_float(v)";
            max_func = "fmax";
            min_func = "fmin";
            abs_func = "fabs";
            type_size = "4";
            is_fp = true;
            break;
    }

    return JitConstants{
        MakeJitConstant(macroName + "_TYPE", type),
        MakeJitConstant(macroName + "_VAL_MAX", max_val),
        MakeJitConstant(macroName + "_VAL_MIN", min_val),
        MakeJitConstant(macroName + "_VAL_ONE", val_one),
        MakeJitConstant(macroName + "_VAL_ZERO", val_zero),
        MakeJitConstant("TO_" + macroName + "_TYPE(v)", to_type),
        MakeJitConstant("TO_" + macroName + "_TYPE_SAT(v)", to_type_sat),
        MakeJitConstant("AS_" + macroName + "_TYPE(v)", as_type),
        MakeJitConstant(macroName + "_MAX_FUNC", max_func),
        MakeJitConstant(macroName + "_MIN_FUNC", min_func),
        MakeJitConstant(macroName + "_ABS_FUNC", abs_func),
        MakeJitConstant(macroName + "_TYPE_SIZE", type_size),
        MakeJitConstant(macroName + "_IS_FP", is_fp),
    };
}
JitConstants MakeTypeJitConstants(WeightsType weightsType, const std::string& macroName) {
    switch (weightsType) {
        case WeightsType::UNSUPPORTED:
            return MakeTypeJitConstants(Datatype::UNSUPPORTED, macroName);
        case WeightsType::F16:
            return MakeTypeJitConstants(Datatype::F16, macroName);
        case WeightsType::F32:
            return MakeTypeJitConstants(Datatype::F32, macroName);
        case WeightsType::INT8:
            return MakeTypeJitConstants(Datatype::INT8, macroName);
        case WeightsType::UINT8:
            return MakeTypeJitConstants(Datatype::UINT8, macroName);
        case WeightsType::BINARY:
            return MakeTypeJitConstants(Datatype::UINT32, macroName);
    }
    assert(false || "Unreachable!");
    // FIXME: Is there some builtin_unreachable available?
    return MakeTypeJitConstants(Datatype::UNSUPPORTED, macroName);
}

JitConstants MakeActivationJitConstants(const base_activation_params& params,
                                        Datatype out_dt,
                                        const std::string& suffix,
                                        bool use_type_parameter,
                                        bool disable_type_conversion) {
    auto jitConstants = JitConstants{MakeJitConstant("NL_M" + suffix, params.m),
                                     MakeJitConstant("NL_N" + suffix, params.n)};
    jitConstants.Merge(MakeActivationJitConstants(
        params.function, out_dt, suffix, use_type_parameter, disable_type_conversion));
    return jitConstants;
}

JitConstants MakeActivationJitConstants(std::vector<kernel_selector::base_activation_params> params,
                                        Datatype out_dt,
                                        const std::string& suffix,
                                        bool use_type_parameter,
                                        bool disable_type_conversion) {
    JitConstants res = {};
    if (params.empty()) {
        return MakeActivationJitConstants({ActivationFunction::NONE, 0.f, 0.f}, out_dt,
                                          suffix, use_type_parameter, disable_type_conversion);
    }
    std::string res_activation = "";
    std::string activation_params = "";
    for (size_t i = 0; i < params.size(); i++) {
        std::string activation_suffix = suffix + "_" + std::to_string(i);
        auto jitConstants = JitConstants{MakeJitConstant("NL_M" + activation_suffix, params[i].m),
                                         MakeJitConstant("NL_N" + activation_suffix, params[i].n)};
        jitConstants.Merge(MakeActivationJitConstants(
                params[i].function, out_dt, activation_suffix, use_type_parameter, disable_type_conversion));
        res.Merge(jitConstants);

        if (i == 0) {
            if (params[i].gradient) {
                activation_params = use_type_parameter ? "(jit_type, input_grad, input, params)" : "(input_grad, input, params)";
            } else {
                activation_params = use_type_parameter ? "(jit_type, input, params)" : "(input, params)";
            }
            res_activation = "ACTIVATION_FUNC" + activation_suffix + activation_params;
        } else {
            res_activation = "ACTIVATION" + activation_suffix + "(" + (use_type_parameter ? "jit_type, " : "") +
                             (params[i].gradient ? "input_grad, " : "") +
                             res_activation + ", ACTIVATION_PARAMS" + activation_suffix + ")";
        }
    }
    if (params[params.size() - 1].gradient) {
        activation_params = use_type_parameter ? "(jit_type, input_grad, input, params)" : "(input_grad, input, params)";
    } else {
        activation_params = use_type_parameter ? "(jit_type, input, params)" : "(input, params)";
    }
    res.AddConstant(MakeJitConstant("ACTIVATION_PARAMS" + suffix, "ACTIVATION_PARAMS" + suffix + "_0"));
    res.AddConstant(MakeJitConstant("ACTIVATION" + suffix + activation_params, res_activation));
    return res;
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
// MakeLoopUnrollParamsJitConstants
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
JitConstants MakeLoopUnrollParamsJitConstants(uint32_t loopCount) {
    JitConstants jit{
        MakeJitConstant("LOOP0(VAR, STMT)", ""),
        MakeJitConstant("LOOP1(VAR, STMT)", "(STMT); (VAR)++;"),
    };

    for (uint32_t i = 2; i <= loopCount + 1; i++) {
        jit.AddConstant({
            MakeJitConstant("LOOP" + toCodeString(i) + "(VAR, STMT)", "LOOP" + toCodeString(i - 1) + "(VAR, STMT); (STMT); (VAR)++;"),
        });
    }

    jit.AddConstant({
        MakeJitConstant("LOOP(N, VAR, STMT)", "CAT(LOOP, N)((VAR), (STMT))"),
    });

    return jit;
}

bool FusedOpsCodeGenerator::CanPreloadData(const FusedOpsConfiguration& conf) const {
    if (conf.loop_axes.empty())
        return true;

    bool can_preload = true;
    // Check that tensor offset doesn't have dependency from the loop dimensions
    for (auto& d : conf.loop_axes) {
        for (auto& t : desc.tensors) {
            auto idx = idx_desc{conf.bfzyx_idx_order, t};
            switch (d) {
                case Tensor::DataChannelName::BATCH:   can_preload &= idx.b == "0"; break;
                case Tensor::DataChannelName::FEATURE: can_preload &= idx.f == "0"; break;
                case Tensor::DataChannelName::Z:       can_preload &= idx.z == "0"; break;
                case Tensor::DataChannelName::Y:       can_preload &= idx.y == "0"; break;
                case Tensor::DataChannelName::X:       can_preload &= idx.x == "0"; break;
                default: return false;
            }
        }
    }

    return can_preload;
}

std::string FusedOpsCodeGenerator::GetTypeStr() const {
    switch (desc.GetType()) {
        case KernelType::ELTWISE: return "eltwise";
        case KernelType::SCALE: return "scale";
        case KernelType::QUANTIZE: return "quantize";
        case KernelType::ACTIVATION: return "activation";
        case KernelType::UNKNOWN: throw std::runtime_error("Invalid type of fused operation. Fused op can't have type UNKNOWN");
        default: return "";
    }
}

JitConstants FusedOpsCodeGenerator::MakeFusedTensorJitConstants(const FusedOpsConfiguration& /*conf*/) const {
    JitConstants jit{};
    for (size_t op_input_id = 0; op_input_id < desc.tensors.size(); op_input_id++) {
        std::string name = GetInputTensorName(op_input_id);
        jit.AddConstant(MakeJitConstant(name, desc.tensors[op_input_id]));
    }
    jit.AddConstant(MakeJitConstant(GetOutputTensorName(), desc.output_tensor));
    return jit;
}

JitConstants FusedOpsCodeGenerator::MakeInputDeclsJitConstants(const FusedOpsConfiguration& /*conf*/) const {
    JitConstants jit = {};

    std::string input_decls = "";
    for (size_t op_input_id = 0; op_input_id < desc.tensors.size(); op_input_id++) {
        std::string ptr_name = GetInputPtrName(op_input_id);
        input_decls += "\\\n\tconst __global " + toCLType(desc.tensors[op_input_id].GetDType()) +
                       "* " + ptr_name + (op_input_id == desc.tensors.size() - 1 ? "" : ",");
    }

    jit.AddConstant(MakeJitConstant("FUSED_OP"+std::to_string(desc.op_id)+"_DECLS", input_decls));
    return jit;
}

JitConstants FusedOpsCodeGenerator::MakeLoadJitConstants(const FusedOpsConfiguration& conf, const DataTensor prim_output) const {
    JitConstants jit = {};

    auto vec_size = conf.vec_size;
    auto idx = conf.bfzyx_idx_order;

    std::string load_decls = "";
    static int i = 0;
    // TODO: check if there is a use case for index reuse or it can be removed
    bool reuse_index = false;
    bool safe_load = conf.boundary_check == FusedOpsConfiguration::BoundaryCheck::ENABLED;
    std::string reused_idx = "reused_idx_" + std::to_string(i++);
    if (reuse_index) {
        load_decls += "\\\n\tint " + reused_idx + " = " +  GetIdx(0, idx_desc{idx, desc.tensors[0]}, safe_load) + ";";
    }

    for (auto op_input_id : GetRequiredInputs()) {
        load_decls += "\\\n\t" + GetInputTypeName(op_input_id, vec_size) + " " + GetInputVarName(op_input_id) + " = " +
                      GetJitLoad(conf, op_input_id, prim_output, reuse_index, reused_idx) + ";";
    }
    jit.AddConstant(MakeJitConstant("FUSED_OP"+std::to_string(desc.op_id)+"_LOAD" + conf.suffix, load_decls));

    return jit;
}

JitConstants FusedOpsCodeGenerator::MakeOpJitConstants(const FusedOpsConfiguration& conf,
                                                       const std::string in_var, const Datatype in_type,
                                                       std::string& out_var, Datatype& out_type) const {
    JitConstants jit = {};

    std::string op_decls = "";
    auto vec_size = conf.vec_size;
    auto idx = conf.bfzyx_idx_order;

    out_var = GetOutputVarName(in_var);
    out_type = desc.output_tensor.GetDType();

    std::vector<std::string> in_vars_converted;
    for (size_t i = 0; i < desc.tensors.size(); i++) {
        auto in_name = GetInputVarName(i);
        if (desc.tensors[0].GetDType() != desc.output_tensor.GetDType()) {
            in_name = ConvertToOutputType(in_name, vec_size);
        }
        in_vars_converted.push_back(in_name);
    }

    switch (desc.GetType()) {
        case KernelType::SCALE: {
            op_decls += "\\\n\t" + GetOutputType(vec_size) + " " + out_var + " = " +
                        in_vars_converted[0] + " * " + ConvertToOutputType(in_var, vec_size) + ";";
            if (desc.tensors.size() > 1) {
                op_decls += "\\\n\t" + out_var + " += " + in_vars_converted[1] + ";";
            }
            break;
        }
        case KernelType::ELTWISE: {
            op_decls += "\\\n\t" + GetOutputType(vec_size) + " " + out_var + " = " + in_vars_converted[0] +
                        " + " + ConvertToOutputType(in_var, vec_size) + ";";
            break;
        }
        case KernelType::QUANTIZE: {
            auto p = desc.GetOpParams<quantize_fuse_params>();
            if (!p)
                throw std::runtime_error("[clDNN] Quantize fuse params can't be nullptr");

            // We can't convert inputs to output data type, because it might be equal to UINT8 or INT8, so we convert the data
            // to the zero tensor's (input_lo) type
            std::string tmp_var = in_var;
            std::string tmp_type;
            std::string in_converted = in_var;
            if (in_type != desc.tensors[0].GetDType()) {
                tmp_type = GetType(desc.tensors[0].GetDType(), vec_size);
                tmp_var = out_var + "_tmp";
                in_converted = ConvertToType(in_var, desc.tensors[0].GetDType(), vec_size);
            }

            auto post_scale = p->per_tensor_output_scale ? Broadcast(std::to_string(p->out_scale), desc.tensors[0].GetDType(), vec_size) : GetInputVarName(6);
            auto post_shift = p->per_tensor_output_shift ? Broadcast(std::to_string(p->out_shift), desc.tensors[0].GetDType(), vec_size) : GetInputVarName(7);
            auto pre_scale = p->per_tensor_input_scale ? Broadcast(std::to_string(p->in_scale), desc.tensors[0].GetDType(), vec_size) : GetInputVarName(4);
            auto pre_shift = p->per_tensor_input_shift ? Broadcast(std::to_string(p->in_shift), desc.tensors[0].GetDType(), vec_size) : GetInputVarName(5);
            auto in_lo = p->per_tensor_input_range ? Broadcast(std::to_string(p->in_lo), desc.tensors[0].GetDType(), vec_size) : GetInputVarName(0);
            auto in_hi = p->per_tensor_input_range ? Broadcast(std::to_string(p->in_hi), desc.tensors[0].GetDType(), vec_size) : GetInputVarName(1);

            op_decls += "\\\n\t" + tmp_type + " " + tmp_var + " = min(max(" + in_lo + ", " + in_converted + "), " + in_hi + ");";
            op_decls += "\\\n\t" + tmp_var + " = " + tmp_var + "*" + pre_scale + ";";
            if (p->has_pre_shift)
                op_decls += "\\\n\t" + tmp_var + " = " + tmp_var + " + " + pre_shift + ";";
            op_decls += "\\\n\t" + tmp_var + " = round(" + tmp_var + ");";

            bool need_round = (p->has_post_scale || p->has_post_shift) &&
                              (desc.output_tensor.GetDType() == Datatype::UINT8 || desc.output_tensor.GetDType() == Datatype::INT8);
            if (p->has_post_scale)
                op_decls += "\\\n\t" + tmp_var + " = (" + tmp_var + "*" + post_scale + ");";
            if (p->has_post_shift)
                op_decls += "\\\n\t" + tmp_var + " = (" + tmp_var + " + " + post_shift + ");";
            if (need_round)
                op_decls += "\\\n\t" + tmp_var + " = round(" + tmp_var + ");";

            op_decls += "\\\n\t" + GetOutputType(vec_size) + " " + out_var + " = " + ConvertToOutputTypeSat(tmp_var, vec_size) +";";
            break;
        }
        case KernelType::ACTIVATION: {
            auto p = desc.GetOpParams<activation_fuse_params>();
            base_activation_params activation_p = p->param;
            op_decls += "\\\n\t" + GetOutputType(vec_size) + " " + out_var + " = " + in_var + ";";
            if (activation_p.function != ActivationFunction::NONE) {
                auto suffix = "_FUSED_OP"+std::to_string(desc.op_id) + conf.suffix;
                std::string nl_m = std::to_string(activation_p.m);
                std::string nl_n = std::to_string(activation_p.n);

                if (desc.tensors.size() == 1) {
                    if (desc.tensors[0].GetDType() != out_type) {
                        nl_m = ConvertToOutputType(GetInputVarName(0), vec_size);
                    } else {
                        nl_m = GetInputVarName(0);
                    }
                } else {
                    nl_m = Broadcast(nl_m, out_type, vec_size);
                }

                nl_n = Broadcast(nl_n, out_type, vec_size);

                // Disable type casts in activation, since current jit generator for activation don't respect vector size of parameters.
                // So conversion is explicitly done in params declaration
                jit.Merge(MakeActivationJitConstants(activation_p.function, out_type, suffix, false, true));
                std::string params = nl_m + ","+ nl_n;
                op_decls += "\\\n\t" + out_var + " = ACTIVATION_FUNC" + suffix + "(" + out_var + ", " + params + ");";
            }
            break;
        }
        default: break;
    }

    jit.AddConstant(MakeJitConstant("FUSED_OP"+std::to_string(desc.op_id)+"_ACTION" + conf.suffix, op_decls));

    return jit;
}

std::string FusedOpsCodeGenerator::GetInputTensorName(size_t input_id) const {
    return "FUSED_OP_" + std::to_string(desc.op_id) + "_INPUT" + std::to_string(input_id);
}

std::string FusedOpsCodeGenerator::GetOutputTensorName() const {
    return "FUSED_OP_" + std::to_string(desc.op_id) + "_OUTPUT";
}

std::string FusedOpsCodeGenerator::GetInputTypeName(size_t input_id, size_t vec_size) const {
    if (vec_size == 0 || vec_size > 8)
        throw std::invalid_argument("Invalid vector size in jit definitions: " + std::to_string(vec_size));
    std::string scalar_type = GetInputTensorName(input_id) + "_TYPE";
    if (vec_size > 1)
        return "MAKE_VECTOR_TYPE(" + scalar_type + "," + std::to_string(vec_size) + ")";
    else
        return scalar_type;
}

std::string FusedOpsCodeGenerator::GetIdx(size_t input_id, idx_desc idx, bool should_be_safe) const {
    std::string idx_order = "";

    if (DataTensor::ChannelsCount(desc.tensors[input_id].GetLayout()) <= 4) {
        idx_order = idx.b + "," + idx.f + "," + idx.y + "," + idx.x;
    } else if (DataTensor::ChannelsCount(desc.tensors[input_id].GetLayout()) == 5) {
        idx_order = idx.b + "," + idx.f + "," + idx.z + "," + idx.y + "," + idx.x;
    }

    if (should_be_safe) {
        return GetInputTensorName(input_id) + "_GET_INDEX_SAFE(" + idx_order +")";
    } else {
        return GetInputTensorName(input_id) + "_GET_INDEX(" + idx_order +")";
    }
}

std::string FusedOpsCodeGenerator::GetJitLoad(const FusedOpsConfiguration& conf, size_t input_id, const DataTensor prim_output,
                                                          bool reuse_index, std::string reused_idx) const {
    auto& input_tensor = desc.tensors[input_id];
    size_t vec_size = 1;
    auto input_dt = input_tensor.GetDType();
    if (desc.GetType() == KernelType::ELTWISE) {
        if (input_tensor.LogicalSize() == prim_output.LogicalSize() &&
            input_tensor.GetLayout() != prim_output.GetLayout() && conf.vec_size > 1) {
            throw std::runtime_error("[clDNN] Mixed layouts of input tensors are not supported in fused eltwise");
        }
        vec_size = conf.vec_size;
    }

    if (conf.vec_axis == Tensor::DataChannelName::FEATURE &&
        DataTensor::Extract(input_tensor.GetLayout(), conf.vec_axis, input_tensor.GetDims()).v != 1) {
        vec_size = conf.vec_size;
    }

    auto idx = conf.bfzyx_idx_order;
    if (vec_size == 0 || vec_size > 8)
        throw std::invalid_argument("Invalid vector size in jit definitions: " + std::to_string(vec_size));

    bool safe_load = conf.boundary_check == FusedOpsConfiguration::BoundaryCheck::ENABLED;

    std::string index_func_call_vec = reuse_index ? reused_idx : GetIdx(input_id, idx_desc{idx, desc.tensors[input_id]}, safe_load);
    std::string index_func_call = reuse_index ? reused_idx : GetIdx(input_id, idx_desc{idx, desc.tensors[input_id]}, safe_load);
    if (conf.index_type == FusedOpsConfiguration::IndexType::LINEAR_OFFSET) {
        std::string offset = conf.bfzyx_idx_order[0];
        if (safe_load)
            offset = "(" + offset + " % " + std::to_string(input_tensor.LogicalSize()) + ")";
        if (vec_size > 1)
            return "((const __global " + toCLType(input_dt) + std::to_string(vec_size) + "*)(" +
                   GetInputPtrName(input_id) + " + " + offset + "))[0]";
        else
            return GetInputPtrName(input_id) + "[" + offset + "]";
    } else {
        // TODO: Need to add smarter vectors handling:
        // 1. Boundary checks for safe load
        // 2. If in given configuration data can't be loaded by a simple UNIT_BLOCK_READx call or load from casted ptr,
        //    we can gather the data to vector
        if (conf.load_type == FusedOpsConfiguration::LoadType::LT_ALIGNED_READ) {
            std::string vs = vec_size > 1 ? std::to_string(vec_size)  : "";
            std::string block_read;

            if (input_dt == Datatype::F32) {
                block_read = CastToType(" intel_sub_group_block_read" + vs + "("
                                        + "(const __global uint*)(" + GetInputPtrName(input_id) + " + " + index_func_call_vec + "))",
                                        input_dt, vec_size);
            } else if (input_dt == Datatype::F16) {
                block_read = CastToType(" intel_sub_group_block_read_us" + vs + "("
                                        + "(const __global ushort*)(" + GetInputPtrName(input_id) + " + " + index_func_call_vec + "))",
                                        input_dt, vec_size);
            } else {
                throw std::runtime_error("Aligned load is not supported yet for " + toCLType(input_dt) + " data type");
            }

            if (vec_size > 1) {
                return block_read;
            } else if (input_tensor.LogicalSize() > 1) {
                // Currently we assume that in such scenario we can safely load sub_group_size elements from the pointer
                return Broadcast(block_read, input_dt, conf.vec_size);
            } else {
                // Input has only one element, so broadcast it for the whole vector size
                return Broadcast(GetInputPtrName(input_id) + "[" + index_func_call + "]", input_dt, conf.vec_size);
            }
        } else {
            if (vec_size > 1) {
                return "((const __global " + toCLType(input_dt) + std::to_string(vec_size) + "*)(" +
                       GetInputPtrName(input_id) + " + " + index_func_call_vec + "))[0]";
            } else {
                return GetInputPtrName(input_id) + "[" + index_func_call + "]";
            }
        }
    }
}

std::string FusedOpsCodeGenerator::GetInputPtrName(size_t input_id) const {
    return GetTypeStr() + std::to_string(desc.op_id) + "_input" + std::to_string(input_id);
}

std::string FusedOpsCodeGenerator::GetInputVarName(size_t input_id) const {
    return GetTypeStr() + std::to_string(desc.op_id) + "_data" + std::to_string(input_id);
}

std::string FusedOpsCodeGenerator::GetOutputVarName(std::string input_var) const {
    static int i = 0;
    std::replace(input_var.begin(), input_var.end(), '[', '_');
    std::replace(input_var.begin(), input_var.end(), ']', '_');
    std::replace(input_var.begin(), input_var.end(), ' ', '_');
    return input_var + "_" + std::to_string(i++);
}

std::string FusedOpsCodeGenerator::GetType(Datatype dt, size_t vec_size) const {
    if (vec_size > 1)
        return toCLType(dt) + std::to_string(vec_size);
    else
        return toCLType(dt);
}

std::string FusedOpsCodeGenerator::GetOutputType(size_t vec_size) const {
    return GetType(desc.output_tensor.GetDType(), vec_size);
}

std::string FusedOpsCodeGenerator::ConvertToType(std::string var, Datatype dt, size_t vec_size) const {
    return "convert_" + GetType(dt, vec_size) + "(" + var + ")";
}

std::string FusedOpsCodeGenerator::CastToType(std::string var, Datatype dt, size_t vec_size) const {
    return "as_" + GetType(dt, vec_size) + "(" + var + ")";
}

std::string FusedOpsCodeGenerator::ConvertToOutputType(std::string var, size_t vec_size) const {
    return ConvertToType(var, desc.output_tensor.GetDType(), vec_size);
}

std::string FusedOpsCodeGenerator::Broadcast(std::string var, Datatype dt, size_t vec_size) const {
    return "(" + GetType(dt, vec_size) + ")(" + var + ")";
}

std::string FusedOpsCodeGenerator::ConvertToOutputTypeSat(std::string var, size_t vec_size) const {
    if (desc.output_tensor.GetDType() == Datatype::F32 || desc.output_tensor.GetDType() == Datatype::F16)
        return "convert_" + GetOutputType(vec_size) + "(" + var + ")";
    else
        return "convert_" + GetOutputType(vec_size) + "_sat(" + var + ")";
}

std::vector<size_t> FusedOpsCodeGenerator::GetRequiredInputs() const {
    switch (desc.GetType()) {
        case KernelType::QUANTIZE: {
            auto p = std::dynamic_pointer_cast<quantize_fuse_params>(desc.op_params);
            if (p) {
                std::vector<size_t> res = {};
                if (!p->per_tensor_input_range) {
                    res.push_back(0);
                    res.push_back(1);
                }
                if (!p->per_tensor_input_scale)
                    res.push_back(4);
                if (p->has_pre_shift && !p->per_tensor_input_shift)
                    res.push_back(5);
                if (p->has_post_scale && !p->per_tensor_output_scale)
                    res.push_back(6);
                if (p->has_post_shift && !p->per_tensor_output_shift)
                    res.push_back(7);

                return res;
            }
            return {};
        }
        default: {
            std::vector<size_t> res;
            for (size_t i = 0; i < desc.tensors.size(); i++) {
                res.push_back(i);
            }
            return res;
        }
    }
}

}  // namespace kernel_selector
