[llvm] [SPIR-V] Add NonSemantic.AuxData emission (-spirv-preserve-auxdata) (PR #200002)

[Arseniy Obolenskiy via llvm-commits]

================
@@ -0,0 +1,268 @@
+//===-- SPIRVAuxDataHandler.cpp - NonSemantic.AuxData emitter -*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "SPIRVAuxDataHandler.h"
+#include "MCTargetDesc/SPIRVMCTargetDesc.h"
+#include "SPIRVSubtarget.h"
+#include "SPIRVUtils.h"
+#include "llvm/CodeGen/AsmPrinter.h"
+#include "llvm/IR/Attributes.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/GlobalObject.h"
+#include "llvm/IR/GlobalVariable.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Metadata.h"
+#include "llvm/IR/Module.h"
+#include "llvm/MC/MCInst.h"
+#include "llvm/MC/MCStreamer.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/ErrorHandling.h"
+
+using namespace llvm;
+
+static cl::opt<bool> SPVPreserveAuxData(
+    "spirv-preserve-auxdata",
+    cl::desc("Preserve LLVM attributes and metadata as "
+             "NonSemantic.AuxData ExtInst annotations (requires "
+             "SPV_KHR_non_semantic_info)"),
+    cl::Optional, cl::Hidden, cl::init(false));
+
+namespace {
+enum AuxDataLinkageType : uint32_t {
+  AvailableExternally = 0,
+};
+
+constexpr unsigned NonSemanticAuxDataSet =
+    static_cast<unsigned>(SPIRV::InstructionSet::NonSemantic_AuxData);
+
+AttributeSet getGOAttrs(const GlobalObject *GO) {
+  if (const auto *F = dyn_cast<Function>(GO))
+    return F->getAttributes().getFnAttrs();
+  return cast<GlobalVariable>(GO)->getAttributes();
+}
+} // namespace
+
+SPIRVAuxDataHandler::SPIRVAuxDataHandler(AsmPrinter &AP, const Module &M)
+    : Asm(AP), Mod(M) {
+  for (const Function &F : M)
+    if (F.hasFnAttribute(SPIRV_WAS_AVAILABLE_EXTERNALLY_ATTR))
+      LinkagePreservedFns.push_back(&F);
+}
+
+bool SPIRVAuxDataHandler::hasWork() const { return SPVPreserveAuxData; }
+
+void SPIRVAuxDataHandler::prepareModuleOutput(const SPIRVSubtarget &ST,
+                                              SPIRV::ModuleAnalysisInfo &MAI) {
+  if (!hasWork())
+    return;
+  if (!ST.canUseExtension(SPIRV::Extension::SPV_KHR_non_semantic_info)) {
+    if (SPVPreserveAuxData)
+      report_fatal_error("-spirv-preserve-auxdata requires the "
+                         "SPV_KHR_non_semantic_info extension to be enabled.");
+    return;
+  }
+  MAI.Reqs.addExtension(SPIRV::Extension::SPV_KHR_non_semantic_info);
+  if (!MAI.ExtInstSetMap.count(NonSemanticAuxDataSet))
+    MAI.ExtInstSetMap[NonSemanticAuxDataSet] = MAI.getNextIDRegister();
+}
+
+MCRegister
+SPIRVAuxDataHandler::getOrEmitString(StringRef S,
+                                     SPIRV::ModuleAnalysisInfo &MAI) {
+  auto [It, Inserted] = StringRegs.try_emplace(S);
+  if (!Inserted)
+    return It->second;
+  MCRegister Reg = MAI.getNextIDRegister();
+  It->second = Reg;
+  MCInst Inst;
+  Inst.setOpcode(SPIRV::OpString);
+  Inst.addOperand(MCOperand::createReg(Reg));
+  addStringImm(S, Inst);
+  emitMCInst(Inst);
+  return Reg;
+}
+
+void SPIRVAuxDataHandler::collectAttributesFor(
+    const GlobalObject *GO, function_ref<MCRegister()> GetNameReg,
+    SPIRV::ModuleAnalysisInfo &MAI) {
+  AuxDataOpcode Opcode = isa<Function>(GO) ? FunctionAttributeOpcode
+                                           : GlobalVariableAttributeOpcode;
+  for (const Attribute &A : getGOAttrs(GO)) {
+    if (A.isStringAttribute() &&
+        A.getKindAsString() == SPIRV_WAS_AVAILABLE_EXTERNALLY_ATTR)
+      continue;
+    ExtInstRecord Rec;
+    Rec.Opcode = Opcode;
+    Rec.Operands.push_back(GetNameReg());
+    if (A.isStringAttribute()) {
+      Rec.Operands.push_back(getOrEmitString(A.getKindAsString(), MAI));
+      StringRef Val = A.getValueAsString();
+      if (!Val.empty())
+        Rec.Operands.push_back(getOrEmitString(Val, MAI));
+    } else {
+      Rec.Operands.push_back(
+          getOrEmitString(StringPool.save(A.getAsString()), MAI));
+    }
+    PendingRecords.push_back(std::move(Rec));
+  }
+}
+
+void SPIRVAuxDataHandler::collectMetadataFor(
+    const GlobalObject *GO, function_ref<MCRegister()> GetNameReg,
+    ArrayRef<StringRef> MDNames, SPIRV::ModuleAnalysisInfo &MAI) {
+  SmallVector<std::pair<unsigned, MDNode *>> AllMD;
+  GO->getAllMetadata(AllMD);
+  if (AllMD.empty())
+    return;
+  AuxDataOpcode Opcode =
+      isa<Function>(GO) ? FunctionMetadataOpcode : GlobalVariableMetadataOpcode;
+  // Skip non-MDString operands: emitting them would require a full value
+  // translation we can't safely drive from here.
+  auto CollectStrings =
+      [&](MDNode *MD) -> std::optional<SmallVector<MCRegister, 4>> {
+    SmallVector<MCRegister, 4> Out;
+    for (const MDOperand &MdOp : MD->operands()) {
+      auto *MDStr = dyn_cast_or_null<MDString>(MdOp.get());
+      if (!MDStr)
+        return std::nullopt;
+      Out.push_back(getOrEmitString(MDStr->getString(), MAI));
+    }
+    return Out;
+  };
+  for (const auto &MD : AllMD) {
+    if (MD.first == LLVMContext::MD_dbg)
+      continue;
+    StringRef MDName = MDNames[MD.first];
+    if (MDName == "spirv.Decorations" || MDName == "spirv.ParameterDecorations")
+      continue;
+    auto Operands = CollectStrings(MD.second);
+    if (!Operands)
+      continue;
+    ExtInstRecord Rec;
+    Rec.Opcode = Opcode;
+    Rec.Operands.push_back(GetNameReg());
+    Rec.Operands.push_back(getOrEmitString(MDName, MAI));
+    Rec.Operands.append(Operands->begin(), Operands->end());
+    PendingRecords.push_back(std::move(Rec));
+  }
+}
+
+void SPIRVAuxDataHandler::emitAuxDataStrings(SPIRV::ModuleAnalysisInfo &MAI) {
+  if (!SPVPreserveAuxData)
+    return;
+  if (!MAI.getExtInstSetReg(NonSemanticAuxDataSet).isValid())
+    return;
+  SmallVector<StringRef> MDNames;
+  Mod.getContext().getMDKindNames(MDNames);
+  for (const GlobalObject &GO : Mod.global_objects()) {
+    if (GO.isDeclaration())
+      continue;
+    // Defer the name OpString until the first record actually fires.
+    MCRegister NameReg;
+    auto GetNameReg = [&]() {
+      if (!NameReg.isValid())
+        NameReg = getOrEmitString(GO.getName(), MAI);
+      return NameReg;
+    };
+    collectAttributesFor(&GO, GetNameReg, MAI);
+    collectMetadataFor(&GO, GetNameReg, MDNames, MAI);
+  }
+}
+
+void SPIRVAuxDataHandler::emitAuxData(SPIRV::ModuleAnalysisInfo &MAI) {
+  MCRegister ExtSetReg = MAI.getExtInstSetReg(NonSemanticAuxDataSet);
+  if (!ExtSetReg.isValid())
+    return;
+
+  MCRegister VoidTypeReg = findOrEmitOpTypeVoid(MAI);
+
+  for (const ExtInstRecord &Rec : PendingRecords)
+    emitAuxDataExtInst(Rec.Opcode, VoidTypeReg, ExtSetReg, Rec.Operands, MAI);
+
+  if (LinkagePreservedFns.empty())
+    return;
+
+  MCRegister UInt32TypeReg = findOrEmitOpTypeUInt32(MAI);
+  MCRegister AEConstReg;
+  for (const Function *F : LinkagePreservedFns) {
+    MCRegister FnReg = MAI.getGlobalObjReg(F);
----------------
aobolensk wrote:

Fixed `<id>`s 

https://github.com/llvm/llvm-project/pull/200002