// Copyright Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT

#include <rocRoller/AssemblyKernel.hpp>
#include <rocRoller/CodeGen/AddInstruction.hpp>
#include <rocRoller/CodeGen/Arithmetic/Add.hpp>
#include <rocRoller/CodeGen/CopyGenerator.hpp>
#include <rocRoller/Utilities/Component.hpp>

namespace rocRoller
{
    template <>
    std::shared_ptr<BinaryArithmeticGenerator<Expression::Add>>
        GetGenerator<Expression::Add>(Register::ValuePtr dst,
                                      Register::ValuePtr lhs,
                                      Register::ValuePtr rhs,
                                      Expression::Add const&)
    {
        // Choose the proper generator, based on the context, register type
        // and datatype.
        return Component::Get<BinaryArithmeticGenerator<Expression::Add>>(
            getContextFromValues(dst, lhs, rhs),
            promoteRegisterType(dst, lhs, rhs),
            promoteDataType(dst, lhs, rhs));
    }

    template <>
    Generator<Instruction>
        AddGenerator<Register::Type::Scalar, DataType::Int32>::generate(Register::ValuePtr dest,
                                                                        Register::ValuePtr lhs,
                                                                        Register::ValuePtr rhs,
                                                                        Expression::Add const&)
    {
        AssertFatal(lhs != nullptr);
        AssertFatal(rhs != nullptr);

        co_yield ScalarAddInt32(m_context, dest, lhs, rhs);
    }

    template <>
    Generator<Instruction>
        AddGenerator<Register::Type::Vector, DataType::Int32>::generate(Register::ValuePtr dest,
                                                                        Register::ValuePtr lhs,
                                                                        Register::ValuePtr rhs,
                                                                        Expression::Add const&)
    {
        AssertFatal(lhs != nullptr);
        AssertFatal(rhs != nullptr);

        co_yield m_context->copier()->ensureTypeCommutative(
            {Register::Type::Vector, Register::Type::Constant},
            lhs,
            {Register::Type::Vector, Register::Type::Constant},
            rhs);

        auto const& arch = m_context->targetArchitecture();
        auto const& gpu  = arch.target();

        if(arch.HasCapability(GPUCapability::HasExplicitNC))
        {
            co_yield_(Instruction("v_add_nc_i32", {dest}, {lhs, rhs}, {}, ""));
        }
        else if(gpu.isCDNAGPU())
        {
            co_yield_(Instruction("v_add_i32", {dest}, {lhs, rhs}, {}, ""));
        }
        else
        {
            Throw<FatalError>(fmt::format("AddGenerator<{}, {}> not implemented for {}.\n",
                                          toString(Register::Type::Vector),
                                          toString(DataType::Int32),
                                          gpu.toString()));
        }
    }

    template <>
    Generator<Instruction>
        AddGenerator<Register::Type::Scalar, DataType::UInt32>::generate(Register::ValuePtr dest,
                                                                         Register::ValuePtr lhs,
                                                                         Register::ValuePtr rhs,
                                                                         Expression::Add const&)
    {
        AssertFatal(lhs != nullptr);
        AssertFatal(rhs != nullptr);

        co_yield ScalarAddUInt32(m_context, dest, lhs, rhs);
    }

    template <>
    Generator<Instruction>
        AddGenerator<Register::Type::Vector, DataType::UInt32>::generate(Register::ValuePtr dest,
                                                                         Register::ValuePtr lhs,
                                                                         Register::ValuePtr rhs,
                                                                         Expression::Add const&)
    {
        AssertFatal(lhs != nullptr);
        AssertFatal(rhs != nullptr);

        co_yield m_context->copier()->ensureTypeCommutative(
            {Register::Type::Vector, Register::Type::Literal}, lhs, {Register::Type::Vector}, rhs);

        co_yield VectorAddUInt32(m_context, dest, lhs, rhs);
    }

    template <>
    Generator<Instruction>
        AddGenerator<Register::Type::M0, DataType::UInt32>::generate(Register::ValuePtr dest,
                                                                     Register::ValuePtr lhs,
                                                                     Register::ValuePtr rhs,
                                                                     Expression::Add const&)
    {
        AssertFatal(lhs != nullptr);
        AssertFatal(rhs != nullptr);

        co_yield ScalarAddUInt32(m_context, dest, lhs, rhs);
    }

    template <>
    Generator<Instruction>
        AddGenerator<Register::Type::Scalar, DataType::Int64>::generate(Register::ValuePtr dest,
                                                                        Register::ValuePtr lhs,
                                                                        Register::ValuePtr rhs,
                                                                        Expression::Add const&)
    {
        co_yield describeOpArgs("dest", dest, "lhs", lhs, "rhs", rhs);

        Register::ValuePtr l0, l1, r0, r1;
        co_yield get2DwordsScalar(l0, l1, lhs);
        co_yield get2DwordsScalar(r0, r1, rhs);

        co_yield ScalarAddUInt32(m_context, dest->subset({0}), l0, r0)
            .lock(Scheduling::Dependency::SCC, "Start of Int64 add, locking SCC");
        co_yield ScalarAddUInt32CarryInOut(m_context, dest->subset({1}), l1, r1)
            .unlock("End of Int64 add, unlocking SCC");
    }

    template <>
    Generator<Instruction>
        AddGenerator<Register::Type::Vector, DataType::Int64>::generate(Register::ValuePtr dest,
                                                                        Register::ValuePtr lhs,
                                                                        Register::ValuePtr rhs,
                                                                        Expression::Add const&)
    {
        co_yield describeOpArgs("dest", dest, "lhs", lhs, "rhs", rhs);

        Register::ValuePtr l0, l1, r0, r1;
        co_yield get2DwordsVector(l0, l1, lhs);
        co_yield get2DwordsVector(r0, r1, rhs);

        if(r0->regType() == Register::Type::Literal)
        {
            std::swap(l0, r0);
        }

        if(r1->regType() == Register::Type::Literal)
        {
            std::swap(l1, r1);
        }

        if(r0->regType() == Register::Type::Scalar || r0->regType() == Register::Type::Literal)
        {
            co_yield moveToVGPR(r0);
        }

        if(r1->regType() == Register::Type::Scalar || r1->regType() == Register::Type::Literal)
        {
            co_yield moveToVGPR(r1);
        }

        if(l0->regType() == Register::Type::Scalar)
        {
            co_yield moveToVGPR(l0);
        }

        // v_addc_co_u32 can only accept inline constant literals
        if(l1->regType() == Register::Type::Scalar
           || (l1->regType() == Register::Type::Literal
               && !m_context->targetArchitecture().isSupportedConstantValue(l1)))
        {
            co_yield moveToVGPR(l1);
        }

        bool useVCC = (l0->regType() == Register::Type::Literal
                       && !m_context->targetArchitecture().isSupportedConstantValue(l0))
                      || (l1->regType() == Register::Type::Literal
                          && !m_context->targetArchitecture().isSupportedConstantValue(l1));

        auto dependency = useVCC ? Scheduling::Dependency::VCC : Scheduling::Dependency::Count;

        auto carry
            = useVCC ? m_context->getVCC() : Register::Value::WavefrontPlaceholder(m_context, 1);

        co_yield VectorAddUInt32CarryOut(
            m_context, dest->subset({0}), carry, l0, r0, "least significant half")
            .lock(dependency, "Start of Int64 add, locking VCC");
        co_yield VectorAddUInt32CarryInOut(
            m_context, dest->subset({1}), carry, carry, l1, r1, "most significant half")
            .unlock(dependency, "End of Int64 add, Unlocking VCC.");
    }

    template <>
    Generator<Instruction>
        AddGenerator<Register::Type::Vector, DataType::Halfx2>::generate(Register::ValuePtr dest,
                                                                         Register::ValuePtr lhs,
                                                                         Register::ValuePtr rhs,
                                                                         Expression::Add const&)
    {
        AssertFatal(lhs != nullptr);
        AssertFatal(rhs != nullptr);

        co_yield_(Instruction("v_pk_add_f16", {dest}, {lhs, rhs}, {}, ""));
    }

    template <>
    Generator<Instruction>
        AddGenerator<Register::Type::Vector, DataType::Half>::generate(Register::ValuePtr dest,
                                                                       Register::ValuePtr lhs,
                                                                       Register::ValuePtr rhs,
                                                                       Expression::Add const&)
    {
        AssertFatal(lhs != nullptr);
        AssertFatal(rhs != nullptr);

        co_yield_(Instruction("v_add_f16", {dest}, {lhs, rhs}, {}, ""));
    }

    template <>
    Generator<Instruction>
        AddGenerator<Register::Type::Vector, DataType::BFloat16>::generate(Register::ValuePtr dest,
                                                                           Register::ValuePtr lhs,
                                                                           Register::ValuePtr rhs,
                                                                           Expression::Add const&)
    {
        AssertFatal(lhs != nullptr);
        AssertFatal(rhs != nullptr);

        auto lit16 = Register::Value::Literal(16u);

        co_yield generateOp<Expression::ShiftL>(lhs, lhs, lit16);
        co_yield generateOp<Expression::ShiftL>(rhs, rhs, lit16);

        co_yield_(Instruction("v_add_f32", {dest}, {lhs, rhs}, {}, ""));
        co_yield generateOp<Expression::LogicalShiftR>(dest, dest, lit16);

        co_yield generateOp<Expression::LogicalShiftR>(lhs, lhs, lit16);
        co_yield generateOp<Expression::LogicalShiftR>(rhs, rhs, lit16);
    }

    template <>
    Generator<Instruction>
        AddGenerator<Register::Type::Vector, DataType::Float>::generate(Register::ValuePtr dest,
                                                                        Register::ValuePtr lhs,
                                                                        Register::ValuePtr rhs,
                                                                        Expression::Add const&)
    {
        AssertFatal(lhs != nullptr);
        AssertFatal(rhs != nullptr);

        co_yield m_context->copier()->ensureTypeCommutative(
            {Register::Type::Vector, Register::Type::Constant}, lhs, {Register::Type::Vector}, rhs);

        co_yield_(Instruction("v_add_f32", {dest}, {lhs, rhs}, {}, ""));
    }

    template <>
    Generator<Instruction>
        AddGenerator<Register::Type::Vector, DataType::Double>::generate(Register::ValuePtr dest,
                                                                         Register::ValuePtr lhs,
                                                                         Register::ValuePtr rhs,
                                                                         Expression::Add const&)
    {
        AssertFatal(lhs != nullptr);
        AssertFatal(rhs != nullptr);

        co_yield m_context->copier()->ensureTypeCommutative(
            {Register::Type::Vector, Register::Type::Constant},
            lhs,
            {Register::Type::Vector, Register::Type::Constant},
            rhs);

        co_yield_(Instruction("v_add_f64", {dest}, {lhs, rhs}, {}, ""));
    }

}