//===-- SPIRVInstPrinter.cpp - Output SPIR-V MCInsts as ASM -----*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This class prints a SPIR-V MCInst to a .s file.
//
//===----------------------------------------------------------------------===//

#include "SPIRVInstPrinter.h"
#include "SPIRV.h"
#include "SPIRVBaseInfo.h"
#include "llvm/ADT/APFloat.h"
#include "llvm/MC/MCAsmInfo.h"
#include "llvm/MC/MCExpr.h"
#include "llvm/MC/MCInst.h"
#include "llvm/MC/MCInstrInfo.h"
#include "llvm/MC/MCSymbol.h"
#include "llvm/Support/ErrorHandling.h"

using namespace llvm;
using namespace llvm::SPIRV;

#define DEBUG_TYPE "asm-printer"

// Include the auto-generated portion of the assembly writer.
#include "SPIRVGenAsmWriter.inc"

void SPIRVInstPrinter::printRemainingVariableOps(const MCInst *MI,
                                                 unsigned StartIndex,
                                                 raw_ostream &O,
                                                 bool SkipFirstSpace,
                                                 bool SkipImmediates) {
  const unsigned NumOps = MI->getNumOperands();
  for (unsigned i = StartIndex; i < NumOps; ++i) {
    if (!SkipImmediates || !MI->getOperand(i).isImm()) {
      if (!SkipFirstSpace || i != StartIndex)
        O << ' ';
      printOperand(MI, i, O);
    }
  }
}

void SPIRVInstPrinter::printOpConstantVarOps(const MCInst *MI,
                                             unsigned StartIndex,
                                             raw_ostream &O) {
  unsigned IsBitwidth16 = MI->getFlags() & SPIRV::INST_PRINTER_WIDTH16;
  const unsigned NumVarOps = MI->getNumOperands() - StartIndex;

  if (MI->getOpcode() == SPIRV::OpConstantI && NumVarOps > 2) {
    // Look up the bitwidth of this int type register from
    // IntTypeBitwidths map.
    MCRegister IntTypeReg = MI->getOperand(1).getReg();
    unsigned Bitwidth = IntTypeBitwidths.at(IntTypeReg);

    // SPV_ALTERA_arbitrary_precision_integers allows for integer widths greater
    // than 64, which will be encoded via multiple operands.
    const unsigned TotalBits = NumVarOps * 32;
    APInt Val(TotalBits, 0);
    for (unsigned i = 0; i < NumVarOps; ++i) {
      uint64_t Word = MI->getOperand(StartIndex + i).getImm();
      Val |= APInt(TotalBits, Word) << (i * 32);
    }
    APInt ActualVal = Val.trunc(Bitwidth);
    O << ' ';
    ActualVal.print(O, /*isSigned=*/false);
    return;
  }

  assert((NumVarOps == 1 || NumVarOps == 2) &&
         "Unsupported number of bits for literal variable");

  O << ' ';

  uint64_t Imm = MI->getOperand(StartIndex).getImm();

  // Handle 64 bit literals.
  if (NumVarOps == 2) {
    Imm |= (MI->getOperand(StartIndex + 1).getImm() << 32);
  }

  // Format and print float values.
  if (MI->getOpcode() == SPIRV::OpConstantF && IsBitwidth16 == 0) {
    APFloat FP = NumVarOps == 1 ? APFloat(APInt(32, Imm).bitsToFloat())
                                : APFloat(APInt(64, Imm).bitsToDouble());

    // Print infinity and NaN as hex floats.
    // TODO: Make sure subnormal numbers are handled correctly as they may also
    // require hex float notation.
    if (FP.isInfinity()) {
      if (FP.isNegative())
        O << '-';
      O << "0x1p+128";
      return;
    }
    if (FP.isNaN()) {
      O << "0x1.8p+128";
      return;
    }

    // Format val as a decimal floating point or scientific notation (whichever
    // is shorter), with enough digits of precision to produce the exact value.
    O << format("%.*g", std::numeric_limits<double>::max_digits10,
                FP.convertToDouble());

    return;
  }

  // Print integer values directly.
  O << Imm;
}

void SPIRVInstPrinter::recordIntType(const MCInst *MI) {
  MCRegister IntTypeReg = MI->getOperand(0).getReg();
  unsigned Bitwidth = MI->getOperand(1).getImm();
  IntTypeBitwidths[IntTypeReg] = Bitwidth;
}

void SPIRVInstPrinter::recordOpExtInstImport(const MCInst *MI) {
  MCRegister Reg = MI->getOperand(0).getReg();
  auto Name = getSPIRVStringOperand(*MI, 1);
  auto Set = getExtInstSetFromString(std::move(Name));
  ExtInstSetIDs.insert({Reg, Set});
}

void SPIRVInstPrinter::printInst(const MCInst *MI, uint64_t Address,
                                 StringRef Annot, const MCSubtargetInfo &STI,
                                 raw_ostream &OS) {
  const unsigned OpCode = MI->getOpcode();
  printInstruction(MI, Address, OS);
  if (OpCode == SPIRV::OpTypeInt) {
    recordIntType(MI);
  }

  if (OpCode == SPIRV::OpDecorate || OpCode == SPIRV::OpDecorateId) {
    printOpDecorate(MI, OS);
  } else if (OpCode == SPIRV::OpExtInstImport) {
    recordOpExtInstImport(MI);
  } else if (OpCode == SPIRV::OpExtInst) {
    printOpExtInst(MI, OS);
  } else if (OpCode == SPIRV::UNKNOWN_type) {
    printUnknownType(MI, OS);
  } else {
    // Print any extra operands for variadic instructions.
    const MCInstrDesc &MCDesc = MII.get(OpCode);
    if (MCDesc.isVariadic()) {
      const unsigned NumFixedOps = MCDesc.getNumOperands();
      const unsigned LastFixedIndex = NumFixedOps - 1;
      const int FirstVariableIndex = NumFixedOps;
      if (NumFixedOps > 0 && MCDesc.operands()[LastFixedIndex].OperandType ==
                                 MCOI::OPERAND_UNKNOWN) {
        // For instructions where a custom type (not reg or immediate) comes as
        // the last operand before the variable_ops. This is usually a StringImm
        // operand, but there are a few other cases.
        switch (OpCode) {
        case SPIRV::OpTypeImage:
          OS << ' ';
          printSymbolicOperand<OperandCategory::AccessQualifierOperand>(
              MI, FirstVariableIndex, OS);
          break;
        case SPIRV::OpVariable:
          OS << ' ';
          printOperand(MI, FirstVariableIndex, OS);
          break;
        case SPIRV::OpEntryPoint: {
          // Print the interface ID operands, skipping the name's string
          // literal.
          printRemainingVariableOps(MI, NumFixedOps, OS, false, true);
          break;
        }
        case SPIRV::OpMemberDecorate:
          printRemainingVariableOps(MI, NumFixedOps, OS);
          break;
        case SPIRV::OpExecutionMode:
        case SPIRV::OpExecutionModeId:
        case SPIRV::OpLoopMerge:
        case SPIRV::OpLoopControlINTEL: {
          // Print any literals after the OPERAND_UNKNOWN argument normally.
          printRemainingVariableOps(MI, NumFixedOps, OS);
          break;
        }
        default:
          break; // printStringImm has already been handled.
        }
      } else {
        // For instructions with no fixed ops or a reg/immediate as the final
        // fixed operand, we can usually print the rest with "printOperand", but
        // check for a few cases with custom types first.
        switch (OpCode) {
        case SPIRV::OpLoad:
        case SPIRV::OpStore:
          OS << ' ';
          printSymbolicOperand<OperandCategory::MemoryOperandOperand>(
              MI, FirstVariableIndex, OS);
          printRemainingVariableOps(MI, FirstVariableIndex + 1, OS);
          break;
        case SPIRV::OpSwitch:
          if (MI->getFlags() & SPIRV::INST_PRINTER_WIDTH64) {
            // In binary format 64-bit types are split into two 32-bit operands,
            // but in text format combine these into a single 64-bit value as
            // this is what tools such as spirv-as require.
            const unsigned NumOps = MI->getNumOperands();
            for (unsigned OpIdx = NumFixedOps; OpIdx < NumOps;) {
              if (OpIdx + 1 >= NumOps || !MI->getOperand(OpIdx).isImm() ||
                  !MI->getOperand(OpIdx + 1).isImm()) {
                llvm_unreachable("Unexpected OpSwitch operands");
                continue;
              }
              OS << ' ';
              uint64_t LowBits = MI->getOperand(OpIdx).getImm();
              uint64_t HighBits = MI->getOperand(OpIdx + 1).getImm();
              uint64_t CombinedValue = (HighBits << 32) | LowBits;
              OS << formatImm(CombinedValue);
              OpIdx += 2;

              // Next should be the label
              if (OpIdx < NumOps) {
                OS << ' ';
                printOperand(MI, OpIdx, OS);
                OpIdx++;
              }
            }
          } else {
            printRemainingVariableOps(MI, NumFixedOps, OS);
          }
          break;
        case SPIRV::OpImageSampleImplicitLod:
        case SPIRV::OpImageSampleDrefImplicitLod:
        case SPIRV::OpImageSampleProjImplicitLod:
        case SPIRV::OpImageSampleProjDrefImplicitLod:
        case SPIRV::OpImageFetch:
        case SPIRV::OpImageGather:
        case SPIRV::OpImageDrefGather:
        case SPIRV::OpImageRead:
        case SPIRV::OpImageWrite:
        case SPIRV::OpImageSparseSampleImplicitLod:
        case SPIRV::OpImageSparseSampleDrefImplicitLod:
        case SPIRV::OpImageSparseSampleProjImplicitLod:
        case SPIRV::OpImageSparseSampleProjDrefImplicitLod:
        case SPIRV::OpImageSparseFetch:
        case SPIRV::OpImageSparseGather:
        case SPIRV::OpImageSparseDrefGather:
        case SPIRV::OpImageSparseRead:
        case SPIRV::OpImageSampleFootprintNV:
          OS << ' ';
          printSymbolicOperand<OperandCategory::ImageOperandOperand>(
              MI, FirstVariableIndex, OS);
          printRemainingVariableOps(MI, NumFixedOps + 1, OS);
          break;
        case SPIRV::OpCopyMemory:
        case SPIRV::OpCopyMemorySized: {
          const unsigned NumOps = MI->getNumOperands();
          for (unsigned i = NumFixedOps; i < NumOps; ++i) {
            OS << ' ';
            printSymbolicOperand<OperandCategory::MemoryOperandOperand>(MI, i,
                                                                        OS);
            if (MI->getOperand(i).getImm() & MemoryOperand::Aligned) {
              assert(i + 1 < NumOps && "Missing alignment operand");
              OS << ' ';
              printOperand(MI, i + 1, OS);
              i += 1;
            }
          }
          break;
        }
        case SPIRV::OpConstantI:
        case SPIRV::OpConstantF:
          // The last fixed operand along with any variadic operands that follow
          // are part of the variable value.
          assert(NumFixedOps > 0 && "Expected at least one fixed operand");
          printOpConstantVarOps(MI, NumFixedOps - 1, OS);
          break;
        case SPIRV::OpCooperativeMatrixMulAddKHR: {
          const unsigned NumOps = MI->getNumOperands();
          if (NumFixedOps == NumOps)
            break;

          OS << ' ';
          const unsigned MulAddOp = MI->getOperand(FirstVariableIndex).getImm();
          if (MulAddOp == 0) {
            printSymbolicOperand<
                OperandCategory::CooperativeMatrixOperandsOperand>(
                MI, FirstVariableIndex, OS);
          } else {
            std::string Buffer;
            for (unsigned Mask = 0x1;
                 Mask != SPIRV::CooperativeMatrixOperands::
                             MatrixResultBFloat16ComponentsINTEL;
                 Mask <<= 1) {
              if (MulAddOp & Mask) {
                if (!Buffer.empty())
                  Buffer += '|';
                Buffer += getSymbolicOperandMnemonic(
                    OperandCategory::CooperativeMatrixOperandsOperand, Mask);
              }
            }
            OS << Buffer;
          }
          break;
        }
        case SPIRV::OpSubgroupMatrixMultiplyAccumulateINTEL: {
          const unsigned NumOps = MI->getNumOperands();
          if (NumFixedOps >= NumOps)
            break;
          OS << ' ';
          const unsigned Flags = MI->getOperand(NumOps - 1).getImm();
          if (Flags == 0) {
            printSymbolicOperand<
                OperandCategory::MatrixMultiplyAccumulateOperandsOperand>(
                MI, NumOps - 1, OS);
          } else {
            std::string Buffer;
            for (unsigned Mask = 0x1;
                 Mask <= SPIRV::MatrixMultiplyAccumulateOperands::
                             MatrixBPackedBFloat16INTEL;
                 Mask <<= 1) {
              if (Flags & Mask) {
                if (!Buffer.empty())
                  Buffer += '|';
                Buffer += getSymbolicOperandMnemonic(
                    OperandCategory::MatrixMultiplyAccumulateOperandsOperand,
                    Mask);
              }
            }
            OS << Buffer;
          }
          break;
        }
        case SPIRV::OpSDot:
        case SPIRV::OpUDot:
        case SPIRV::OpSUDot:
        case SPIRV::OpSDotAccSat:
        case SPIRV::OpUDotAccSat:
        case SPIRV::OpSUDotAccSat: {
          const unsigned NumOps = MI->getNumOperands();
          if (NumOps > NumFixedOps) {
            OS << ' ';
            printSymbolicOperand<OperandCategory::PackedVectorFormatsOperand>(
                MI, NumOps - 1, OS);
            break;
          }
          break;
        }
        case SPIRV::OpPredicatedLoadINTEL:
        case SPIRV::OpPredicatedStoreINTEL: {
          const unsigned NumOps = MI->getNumOperands();
          if (NumOps > NumFixedOps) {
            OS << ' ';
            printSymbolicOperand<OperandCategory::MemoryOperandOperand>(
                MI, NumOps - 1, OS);
            break;
          }
          break;
        }
        default:
          printRemainingVariableOps(MI, NumFixedOps, OS);
          break;
        }
      }
    }
  }

  printAnnotation(OS, Annot);
}

void SPIRVInstPrinter::printOpExtInst(const MCInst *MI, raw_ostream &O) {
  // The fixed operands have already been printed, so just need to decide what
  // type of ExtInst operands to print based on the instruction set and number.
  const MCInstrDesc &MCDesc = MII.get(MI->getOpcode());
  unsigned NumFixedOps = MCDesc.getNumOperands();
  const auto NumOps = MI->getNumOperands();
  if (NumOps == NumFixedOps)
    return;

  O << ' ';

  // TODO: implement special printing for OpenCLExtInst::vstor*.
  printRemainingVariableOps(MI, NumFixedOps, O, true);
}

void SPIRVInstPrinter::printOpDecorate(const MCInst *MI, raw_ostream &O) {
  // The fixed operands have already been printed, so just need to decide what
  // type of decoration operands to print based on the Decoration type.
  const MCInstrDesc &MCDesc = MII.get(MI->getOpcode());
  unsigned NumFixedOps = MCDesc.getNumOperands();

  if (NumFixedOps != MI->getNumOperands()) {
    auto DecOp = MI->getOperand(NumFixedOps - 1);
    auto Dec = static_cast<Decoration::Decoration>(DecOp.getImm());

    O << ' ';

    switch (Dec) {
    case Decoration::BuiltIn:
      printSymbolicOperand<OperandCategory::BuiltInOperand>(MI, NumFixedOps, O);
      break;
    case Decoration::UniformId:
      printSymbolicOperand<OperandCategory::ScopeOperand>(MI, NumFixedOps, O);
      break;
    case Decoration::FuncParamAttr:
      printSymbolicOperand<OperandCategory::FunctionParameterAttributeOperand>(
          MI, NumFixedOps, O);
      break;
    case Decoration::FPRoundingMode:
      printSymbolicOperand<OperandCategory::FPRoundingModeOperand>(
          MI, NumFixedOps, O);
      break;
    case Decoration::FPFastMathMode:
      printSymbolicOperand<OperandCategory::FPFastMathModeOperand>(
          MI, NumFixedOps, O);
      break;
    case Decoration::LinkageAttributes:
    case Decoration::UserSemantic:
      printStringImm(MI, NumFixedOps, O);
      break;
    case Decoration::HostAccessINTEL:
      printOperand(MI, NumFixedOps, O);
      if (NumFixedOps + 1 < MI->getNumOperands()) {
        O << ' ';
        printStringImm(MI, NumFixedOps + 1, O);
      }
      break;
    default:
      printRemainingVariableOps(MI, NumFixedOps, O, true);
      break;
    }
  }
}

void SPIRVInstPrinter::printUnknownType(const MCInst *MI, raw_ostream &O) {
  const auto EnumOperand = MI->getOperand(1);
  assert(EnumOperand.isImm() &&
         "second operand of UNKNOWN_type must be opcode!");

  const auto Enumerant = EnumOperand.getImm();
  const auto NumOps = MI->getNumOperands();

  // Print the opcode using the spirv-as unknown opcode syntax
  O << "OpUnknown(" << Enumerant << ", " << NumOps << ") ";

  // The result ID must be printed after the opcode when using this syntax
  printOperand(MI, 0, O);

  O << " ";

  const MCInstrDesc &MCDesc = MII.get(MI->getOpcode());
  unsigned NumFixedOps = MCDesc.getNumOperands();
  if (NumOps == NumFixedOps)
    return;

  // Print the rest of the operands
  printRemainingVariableOps(MI, NumFixedOps, O, true);
}

void SPIRVInstPrinter::printOperand(const MCInst *MI, unsigned OpNo,
                                    raw_ostream &O) {
  if (OpNo < MI->getNumOperands()) {
    const MCOperand &Op = MI->getOperand(OpNo);
    if (Op.isReg())
      O << '%' << (getIDFromRegister(Op.getReg().id()) + 1);
    else if (Op.isImm()) {
      int64_t Imm = Op.getImm();
      // For OpVectorShuffle:
      // A Component literal may also be FFFFFFFF, which means the corresponding
      // result component has no source and is undefined.
      // LLVM representation of poison/undef becomes -1 when lowered to MI.
      if (MI->getOpcode() == SPIRV::OpVectorShuffle && Imm == -1)
        O << "0xFFFFFFFF";
      else
        O << formatImm(Imm);
    } else if (Op.isDFPImm())
      O << formatImm((double)Op.getDFPImm());
    else if (Op.isExpr())
      MAI.printExpr(O, *Op.getExpr());
    else
      llvm_unreachable("Unexpected operand type");
  }
}

void SPIRVInstPrinter::printStringImm(const MCInst *MI, unsigned OpNo,
                                      raw_ostream &O) {
  const unsigned NumOps = MI->getNumOperands();
  unsigned StrStartIndex = OpNo;
  while (StrStartIndex < NumOps) {
    if (MI->getOperand(StrStartIndex).isReg())
      break;

    std::string Str = getSPIRVStringOperand(*MI, StrStartIndex);
    if (StrStartIndex != OpNo)
      O << ' '; // Add a space if we're starting a new string/argument.
    O << '"';
    for (char c : Str) {
      // Escape ", \n characters (might break for complex UTF-8).
      if (c == '\n') {
        O.write("\\n", 2);
      } else {
        if (c == '"')
          O.write('\\');
        O.write(c);
      }
    }
    O << '"';

    unsigned numOpsInString = (Str.size() / 4) + 1;
    StrStartIndex += numOpsInString;

    // Check for final Op of "OpDecorate %x %stringImm %linkageAttribute".
    if (MI->getOpcode() == SPIRV::OpDecorate &&
        MI->getOperand(1).getImm() ==
            static_cast<unsigned>(Decoration::LinkageAttributes)) {
      O << ' ';
      printSymbolicOperand<OperandCategory::LinkageTypeOperand>(
          MI, StrStartIndex, O);
      break;
    }
  }
}

void SPIRVInstPrinter::printExtension(const MCInst *MI, unsigned OpNo,
                                      raw_ostream &O) {
  auto SetReg = MI->getOperand(2).getReg();
  auto Set = ExtInstSetIDs[SetReg];
  auto Op = MI->getOperand(OpNo).getImm();
  O << getExtInstName(Set, Op);
}

template <OperandCategory::OperandCategory category>
void SPIRVInstPrinter::printSymbolicOperand(const MCInst *MI, unsigned OpNo,
                                            raw_ostream &O) {
  if (OpNo < MI->getNumOperands()) {
    O << getSymbolicOperandMnemonic(category, MI->getOperand(OpNo).getImm());
  }
}