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

#include <sstream>
#include <string>

#include "gpu/intel/jit/ir/tensor.hpp"
#include "gpu/intel/jit/ngen/ngen.hpp"

namespace dnnl {
namespace impl {
namespace gpu {
namespace intel {
namespace jit {

// Possible backend instruction sets
enum class fma_kind_t {
    undef = 0,
    mad,
    dp4a,
    dpas,
    dpasw,
};

static auto fma_kind_names = nstl::to_array({
        make_enum_name(fma_kind_t::undef, "undef"),
        make_enum_name(fma_kind_t::mad, "mad"),
        make_enum_name(fma_kind_t::dp4a, "dp4a"),
        make_enum_name(fma_kind_t::dpas, "dpas"),
        make_enum_name(fma_kind_t::dpasw, "dpasw"),
});
GPU_DEFINE_PARSE_ENUM(fma_kind_t, fma_kind_names)

inline bool is_dp_fma(fma_kind_t kind) {
    switch (kind) {
        case fma_kind_t::dp4a:
        case fma_kind_t::dpas:
        case fma_kind_t::dpasw: return true;
        default: return false;
    }
}

fma_kind_t get_supported_fma_kind(
        const hw_t &hw, const type_t &a, const type_t &b, const type_t &c);

int get_simd_size(const hw_t &hw, fma_kind_t kind, const type_t &a,
        const type_t &b, const type_t &c);

class multiply_desc_t {
public:
    multiply_desc_t() = default;

    multiply_desc_t(const layout_t &a_layout, const layout_t &b_layout,
            bool force_c_upconvert)
        : a_layout_(a_layout), b_layout_(b_layout) {
        ir_assert(a_layout.ndims() == dim_idx_t(2)
                && b_layout.ndims() == dim_idx_t(2))
                << "Expected 2D layouts, A layout: " << a_layout
                << " B layout: " << b_layout;

        c_type_ = get_c_type(a_type(), b_type(), force_c_upconvert);
    }

    const layout_t &a_layout() const { return a_layout_; }
    const layout_t &b_layout() const { return b_layout_; }

    const type_t &a_type() const { return a_layout_.type(); }
    const type_t &b_type() const { return b_layout_.type(); }
    const type_t &c_type() const { return c_type_; }

    dim_t m() const { return a_layout_.dims()[0]; }
    dim_t n() const { return b_layout_.dims()[1]; }
    dim_t k() const { return a_layout_.dims()[1]; }

    static type_t get_c_type(
            const type_t &a, const type_t &b, bool force_c_upconvert);

private:
    layout_t a_layout_;
    layout_t b_layout_;
    type_t c_type_;
};

// Function representing DPAS instruction.
class dpas_t : public func_impl_t {
public:
    IR_DECL_DERIVED_TYPE_ID(dpas_t, func_impl_t)

    static func_t make(bool is_dpasw, int exec_size, uint8_t sdepth,
            uint8_t rcount, const type_t &dst_type, const type_t &src1_type,
            const type_t &src2_type) {
        return func_t(new dpas_t(is_dpasw, exec_size, sdepth, rcount, dst_type,
                src1_type, src2_type));
    }

    static func_t make_dpasw(const dpas_t &dpas) {
        return func_t(new dpas_t(true, dpas.exec_size, dpas.sdepth, dpas.rcount,
                dpas.dst_type, dpas.src1_type, dpas.src2_type));
    }

    static func_t make_dp4a(int exec_size, const type_t &dst_type,
            const type_t &src1_type, const type_t &src2_type) {
        return make(/*is_dpasw=*/false, exec_size, /*sdepth=*/1, /*rcount=*/1,
                dst_type, src1_type, src2_type);
    }

    static bool is_dp4a_call(const stmt_t &s) {
        auto call = s.as_ptr<func_call_t>();
        return call && call->func.as<dpas_t>().is_dp4a();
    }

    bool is_dp4a() const { return rcount == 1 && sdepth == 1; }

    bool is_equal(const object_impl_t &obj) const override {
        if (!obj.is<self_type>()) return false;
        auto &other = obj.as<self_type>();

        return (is_dpasw == other.is_dpasw) && (exec_size == other.exec_size)
                && (sdepth == other.sdepth) && (rcount == other.rcount)
                && (dst_type == other.dst_type)
                && (src1_type == other.src1_type)
                && (src2_type == other.src2_type);
    }
    std::string str() const override {
        std::ostringstream oss;
        oss << (is_dpasw ? "dpasw" : is_dp4a() ? "dp4a" : "dpas");
        if (!is_dp4a()) {
            oss << std::to_string(sdepth) << "x" << std::to_string(rcount);
        }
        oss << ".x" << std::to_string(exec_size);
        return oss.str();
    }

    IR_DEFINE_ARG_GET(dst, 0)
    IR_DEFINE_ARG_GET(src0, 1)
    IR_DEFINE_ARG_GET(src1, 2)
    IR_DEFINE_ARG_GET(src2, 3)

    stmt_t operator()(const expr_t &dst, const expr_t &src0, const expr_t &src1,
            const expr_t &src2) const {
        return call({dst, src0, src1, src2});
    }

    int dst_size() const {
        return exec_size * (int)(rcount * sizeof(uint32_t));
    }
    int src0_size() const { return dst_size(); }
    int src1_size() const {
        return exec_size * (int)(sdepth * sizeof(uint32_t));
    }
    int src2_size() const {
        const int dpas_size = sdepth * rcount * sizeof(uint32_t);
        return is_dpasw ? dpas_size / 2 : dpas_size;
    }

    layout_t a_layout() const;
    layout_t b_layout() const;
    layout_t c_layout() const;

    bool matches(const multiply_desc_t &desc) const;

    static bool matches_types(
            const hw_t &hw, const type_t &a, const type_t &b, const type_t &c);
    static bool is_src_type(type_t type);

    bool is_dpasw;

    int exec_size;
    uint8_t sdepth;
    uint8_t rcount;

    type_t dst_type; // src0 type is same as dst_type.
    type_t src1_type;
    type_t src2_type;

private:
    dpas_t(bool is_dpasw, int exec_size, uint8_t sdepth, uint8_t rcount,
            const type_t &dst_type, const type_t &src1_type,
            const type_t &src2_type)
        : func_impl_t(_type_info())
        , is_dpasw(is_dpasw)
        , exec_size(exec_size)
        , sdepth(sdepth)
        , rcount(rcount)
        , dst_type(dst_type)
        , src1_type(src1_type)
        , src2_type(src2_type) {}
};

// Function representing MAD instruction.
class mad_t : public func_impl_t {
public:
    IR_DECL_DERIVED_TYPE_ID(mad_t, func_impl_t)

    static func_t make(const hw_t &hw, const type_t &dst_type, int exec_size,
            const type_t &src1_type, int src1_stride, const type_t src2_type,
            int src2_stride) {
        return func_t(new mad_t(hw, dst_type, exec_size, src1_type, src1_stride,
                src2_type, src2_stride));
    }

    std::string str() const override {
        std::ostringstream oss;
        oss << "madx" << exec_size;
        return oss.str();
    }

    IR_DEFINE_ARG_GET(dst, 0)
    IR_DEFINE_ARG_GET(src0, 1)
    IR_DEFINE_ARG_GET(src1, 2)
    IR_DEFINE_ARG_GET(src2, 3)

    stmt_t operator()(const expr_t &dst, const expr_t &src0, const expr_t &src1,
            const expr_t &src2) const {
        return call({dst, src0, src1, src2});
    }

    int dst_size() const { return exec_size * dst_type.size(); }
    int src0_size() const { return dst_size(); }
    int src1_size() const {
        return std::max(
                src1_type.size(), src1_stride * src1_type.size() * exec_size);
    }
    int src2_size() const {
        return std::max(
                src2_type.size(), src2_stride * src2_type.size() * exec_size);
    }

    static bool matches_types(
            const hw_t &hw, const type_t &a, const type_t &b, const type_t &c);

    static const int max_exec_size = 32;
    static int get_max_exec_size_bytes(const hw_t &hw) {
        return hw >= ngen::HW::XeHPC ? 128 : 64;
    }
    static int get_simd_size(
            const hw_t &hw, const type_t &a, const type_t &b, const type_t &c) {
        int max_size = max_exec_size;
        int max_exec_size_bytes = get_max_exec_size_bytes(hw);
        int max_type_size = (utils::one_of(type_t::bf8(), a, b, c)
                                    || utils::one_of(type_t::hf8(), a, b, c))
                ? 2
                : std::max(a.size(), std::max(b.size(), c.size()));
        return std::min(max_size, max_exec_size_bytes / max_type_size);
    }
    int get_exec_size() const { return exec_size; }

    type_t dst_type;
    type_t src1_type;
    type_t src2_type;

    int exec_size;
    int src1_stride;
    int src2_stride;

private:
    mad_t(const hw_t &hw, const type_t &dst_type, int exec_size,
            const type_t &src1_type, int src1_stride, const type_t &src2_type,
            int src2_stride)
        : func_impl_t(_type_info())
        , dst_type(dst_type)
        , src1_type(src1_type)
        , src2_type(src2_type)
        , exec_size(exec_size)
        , src1_stride(src1_stride)
        , src2_stride(src2_stride) {
        int max_exec_size_bytes = get_max_exec_size_bytes(hw);
        ir_assert(math::is_pow2(exec_size));

        ir_assert(exec_size <= max_exec_size);
        ir_assert(dst_size() <= max_exec_size_bytes);
        ir_assert(src1_size() <= max_exec_size_bytes);
        ir_assert(src2_size() <= max_exec_size_bytes);
    }
};

} // namespace jit
} // namespace intel
} // namespace gpu
} // namespace impl
} // namespace dnnl

#endif
