[Backport to 16] Backport spirv-preserve-auxdata changes (#2187)

Backport the below changes to 16:

f729c49
89d658c
9823690
d498f48
d24b9c6

I made the following changes that are not in the original changes:

Port tests to not use opaque pointers and related flags except for one test
Fix patch fail due to missing unrelated function in this branch

Signed-off-by: Sarnie, Nick <nick.sarnie@intel.com>

include/LLVMSPIRVOpts.h

@@ -90,7 +90,9 @@ enum class DebugInfoEIS : uint32_t {
 /// \brief Helper class to manage SPIR-V translation
 class TranslatorOpts {
 public:
-  using ExtensionsStatusMap = std::map<ExtensionID, bool>;
+  // Unset optional means not directly specified by user
+  using ExtensionsStatusMap = std::map<ExtensionID, std::optional<bool>>;
+
   using ArgList = llvm::SmallVector<llvm::StringRef, 4>;
 
   TranslatorOpts() = default;
@@ -107,11 +109,14 @@ class TranslatorOpts {
     if (ExtStatusMap.end() == I)
       return false;
 
-    return I->second;
+    return I->second && *I->second;
   }
 
   void setAllowedToUseExtension(ExtensionID Extension, bool Allow = true) {
-    ExtStatusMap[Extension] = Allow;
+    // Only allow using the extension if it has not already been disabled
+    auto I = ExtStatusMap.find(Extension);
+    if (I == ExtStatusMap.end() || !I->second || (*I->second) == true)
+      ExtStatusMap[Extension] = Allow;
   }
 
   VersionNumber getMaxVersion() const { return MaxVersion; }
@@ -122,6 +127,10 @@ class TranslatorOpts {
 
   void setMemToRegEnabled(bool Mem2Reg) { SPIRVMemToReg = Mem2Reg; }
 
+  bool preserveAuxData() const { return PreserveAuxData; }
+
+  void setPreserveAuxData(bool ArgValue) { PreserveAuxData = ArgValue; }
+
   void setGenKernelArgNameMDEnabled(bool ArgNameMD) {
     GenKernelArgNameMD = ArgNameMD;
   }
@@ -230,6 +239,8 @@ class TranslatorOpts {
   // Add a workaround to preserve OpenCL kernel_arg_type and
   // kernel_arg_type_qual metadata through OpString
   bool PreserveOCLKernelArgTypeMetadataThroughString = false;
+
+  bool PreserveAuxData = false;
 };
 
 } // namespace SPIRV

lib/SPIRV/SPIRVReader.cpp

@@ -3302,10 +3302,16 @@ bool SPIRVToLLVM::translate() {
     if (!postProcessBuiltinsReturningStruct(M, IsCpp))
       return false;
   }
+
+  for (SPIRVExtInst *EI : BM->getAuxDataInstVec()) {
+    transAuxDataInst(EI);
+  }
+
   eraseUselessFunctions(M);
 
   DbgTran->addDbgInfoVersion();
   DbgTran->finalize();
+
   return true;
 }
 
@@ -4407,6 +4413,67 @@ Instruction *SPIRVToLLVM::transOCLBuiltinFromExtInst(SPIRVExtInst *BC,
   return CI;
 }
 
+void SPIRVToLLVM::transAuxDataInst(SPIRVExtInst *BC) {
+  assert(BC->getExtSetKind() == SPIRV::SPIRVEIS_NonSemantic_AuxData);
+  if (!BC->getModule()->preserveAuxData())
+    return;
+  auto Args = BC->getArguments();
+  // Args 0 and 1 are common between attributes and metadata.
+  // 0 is the function, 1 is the name of the attribute/metadata as a string
+  auto *SpvFcn = BC->getModule()->getValue(Args[0]);
+  auto *F = static_cast<Function *>(getTranslatedValue(SpvFcn));
+  assert(F && "Function should already have been translated!");
+  auto AttrOrMDName = BC->getModule()->get<SPIRVString>(Args[1])->getStr();
+  switch (BC->getExtOp()) {
+  case NonSemanticAuxData::FunctionAttribute: {
+    assert(Args.size() < 4 && "Unexpected FunctionAttribute Args");
+    // If this attr was specially handled and added elsewhere, skip it.
+    const Attribute::AttrKind AsKind =
+        Attribute::getAttrKindFromName(AttrOrMDName);
+    if (AsKind != Attribute::None && F->hasFnAttribute(AsKind))
+      return;
+    if (AsKind == Attribute::None && F->hasFnAttribute(AttrOrMDName))
+      return;
+    // For attributes, arg 2 is the attribute value as a string, which may not
+    // exist.
+    if (Args.size() == 3) {
+      auto AttrValue = BC->getModule()->get<SPIRVString>(Args[2])->getStr();
+      F->addFnAttr(AttrOrMDName, AttrValue);
+    } else {
+      if (AsKind != Attribute::None)
+        F->addFnAttr(AsKind);
+      else
+        F->addFnAttr(AttrOrMDName);
+    }
+    break;
+  }
+  case NonSemanticAuxData::FunctionMetadata: {
+    // If this metadata was specially handled and added elsewhere, skip it.
+    if (F->hasMetadata(AttrOrMDName))
+      return;
+    SmallVector<Metadata *> MetadataArgs;
+    // Process the metadata values.
+    for (size_t CurArg = 2; CurArg < Args.size(); CurArg++) {
+      auto *Arg = BC->getModule()->get<SPIRVEntry>(Args[CurArg]);
+      // For metadata, the metadata values can be either values or strings.
+      if (Arg->getOpCode() == OpString) {
+        auto *ArgAsStr = static_cast<SPIRVString *>(Arg);
+        MetadataArgs.push_back(
+            MDString::get(F->getContext(), ArgAsStr->getStr()));
+      } else {
+        auto *ArgAsVal = static_cast<SPIRVValue *>(Arg);
+        auto *TranslatedMD = transValue(ArgAsVal, F, nullptr);
+        MetadataArgs.push_back(ValueAsMetadata::get(TranslatedMD));
+      }
+    }
+    F->setMetadata(AttrOrMDName, MDNode::get(*Context, MetadataArgs));
+    break;
+  }
+  default:
+    llvm_unreachable("Invalid op");
+  }
+}
+
 // SPIR-V only contains language version. Use OpenCL language version as
 // SPIR version.
 void SPIRVToLLVM::transSourceLanguage() {

lib/SPIRV/SPIRVReader.h

@@ -100,6 +100,7 @@ class SPIRVToLLVM : private BuiltinCallHelper {
   bool transDecoration(SPIRVValue *, Value *);
   bool transAlign(SPIRVValue *, Value *);
   Instruction *transOCLBuiltinFromExtInst(SPIRVExtInst *BC, BasicBlock *BB);
+  void transAuxDataInst(SPIRVExtInst *BC);
   std::vector<Value *> transValue(const std::vector<SPIRVValue *> &,
                                   Function *F, BasicBlock *);
   Function *transFunction(SPIRVFunction *F);

lib/SPIRV/SPIRVWriter.cpp

@@ -903,6 +903,8 @@ SPIRVFunction *LLVMToSPIRVBase::transFunctionDecl(Function *F) {
 
   transFunctionMetadataAsUserSemanticDecoration(BF, F);
 
+  transAuxDataInst(BF, F);
+
   SPIRVDBG(dbgs() << "[transFunction] " << *F << " => ";
            spvdbgs() << *BF << '\n';)
   return BF;
@@ -1082,6 +1084,82 @@ SPIRVValue *LLVMToSPIRVBase::transConstantUse(Constant *C) {
   return BM->addUnaryInst(OpBitcast, ExpectedType, Trans, nullptr);
 }
 
+void LLVMToSPIRVBase::transAuxDataInst(SPIRVFunction *BF, Function *F) {
+  auto *BM = BF->getModule();
+  if (!BM->preserveAuxData())
+    return;
+  BM->addExtension(SPIRV::ExtensionID::SPV_KHR_non_semantic_info);
+  const auto &FnAttrs = F->getAttributes().getFnAttrs();
+  for (const auto &Attr : FnAttrs) {
+    std::vector<SPIRVWord> Ops;
+    Ops.push_back(BF->getId());
+    if (Attr.isStringAttribute()) {
+      // Format for String attributes is:
+      // NonSemanticAuxDataFunctionAttribute Fcn AttrName AttrValue
+      // or, if no value:
+      // NonSemanticAuxDataFunctionAttribute Fcn AttrName
+      //
+      // AttrName and AttrValue are always Strings
+      const StringRef AttrKind = Attr.getKindAsString();
+      const StringRef AttrValue = Attr.getValueAsString();
+      auto *KindSpvString = BM->getString(AttrKind.str());
+      Ops.push_back(KindSpvString->getId());
+      if (!AttrValue.empty()) {
+        auto *ValueSpvString = BM->getString(AttrValue.str());
+        Ops.push_back(ValueSpvString->getId());
+      }
+    } else {
+      // Format for other types is:
+      // NonSemanticAuxDataFunctionAttribute Fcn AttrStr
+      // AttrStr is always a String.
+      const std::string AttrStr = Attr.getAsString();
+      auto *AttrSpvString = BM->getString(AttrStr);
+      Ops.push_back(AttrSpvString->getId());
+    }
+    BM->addAuxData(NonSemanticAuxData::FunctionAttribute,
+                   transType(Type::getVoidTy(F->getContext())), Ops);
+  }
+  SmallVector<std::pair<unsigned, MDNode *>> AllMD;
+  SmallVector<StringRef> MDNames;
+  F->getContext().getMDKindNames(MDNames);
+  F->getAllMetadata(AllMD);
+  for (auto MD : AllMD) {
+    const std::string MDName = MDNames[MD.first].str();
+
+    // spirv.Decorations, spirv.ParameterDecorations and debug information are
+    // handled elsewhere for both forward and reverse translation and are
+    // complicated to support here, so just skip them.
+    if (MDName == SPIRV_MD_DECORATIONS ||
+        MDName == SPIRV_MD_PARAMETER_DECORATIONS ||
+        MD.first == LLVMContext::MD_dbg)
+      continue;
+
+    // Format for metadata is:
+    // NonSemanticAuxDataFunctionMetadata Fcn MDName MDVals...
+    // MDName is always a String, MDVals have different types as explained
+    // below. Also note this instruction has a variable number of operands
+    std::vector<SPIRVWord> Ops;
+    Ops.push_back(BF->getId());
+    Ops.push_back(BM->getString(MDName)->getId());
+    for (unsigned int OpIdx = 0; OpIdx < MD.second->getNumOperands(); OpIdx++) {
+      const auto &CurOp = MD.second->getOperand(OpIdx);
+      if (auto *MDStr = dyn_cast<MDString>(CurOp)) {
+        // For MDString, MDVal is String
+        auto *SPIRVStr = BM->getString(MDStr->getString().str());
+        Ops.push_back(SPIRVStr->getId());
+      } else if (auto *ValueAsMeta = dyn_cast<ValueAsMetadata>(CurOp)) {
+        // For Value metadata, MDVal is a SPIRVValue
+        auto *SPIRVVal = transValue(ValueAsMeta->getValue(), nullptr);
+        Ops.push_back(SPIRVVal->getId());
+      } else {
+        assert(false && "Unsupported metadata type");
+      }
+    }
+    BM->addAuxData(NonSemanticAuxData::FunctionMetadata,
+                   transType(Type::getVoidTy(F->getContext())), Ops);
+  }
+}
+
 SPIRVValue *LLVMToSPIRVBase::transConstant(Value *V) {
   if (auto CPNull = dyn_cast<ConstantPointerNull>(V))
     return BM->addNullConstant(
@@ -2739,6 +2817,11 @@ bool LLVMToSPIRVBase::transBuiltinSet() {
             SPIRVBuiltinSetNameMap::map(BM->getDebugInfoEIS()), &EISId))
       return false;
   }
+  if (BM->preserveAuxData()) {
+    if (!BM->importBuiltinSet(
+            SPIRVBuiltinSetNameMap::map(SPIRVEIS_NonSemantic_AuxData), &EISId))
+      return false;
+  }
   return true;
 }
 

lib/SPIRV/SPIRVWriter.h

@@ -122,6 +122,7 @@ class LLVMToSPIRVBase : protected BuiltinCallHelper {
   SPIRVFunction *transFunctionDecl(Function *F);
   void transVectorComputeMetadata(Function *F);
   void transFPGAFunctionMetadata(SPIRVFunction *BF, Function *F);
+  void transAuxDataInst(SPIRVFunction *BF, Function *F);
   void transFunctionMetadataAsUserSemanticDecoration(SPIRVFunction *BF,
                                                      Function *F);
   bool transGlobalVariables();

lib/SPIRV/libSPIRV/NonSemantic.AuxData.h

@@ -0,0 +1,33 @@
+/*
+** Copyright (c) 2023 The Khronos Group Inc.
+**
+** Permission is hereby granted, free of charge, to any person obtaining a copy
+** of this software and/or associated documentation files (the "Materials"),
+** to deal in the Materials without restriction, including without limitation
+** the rights to use, copy, modify, merge, publish, distribute, sublicense,
+** and/or sell copies of the Materials, and to permit persons to whom the
+** Materials are furnished to do so, subject to the following conditions:
+**
+** The above copyright notice and this permission notice shall be included in
+** all copies or substantial portions of the Materials.
+**
+** MODIFICATIONS TO THIS FILE MAY MEAN IT NO LONGER ACCURATELY REFLECTS KHRONOS
+** STANDARDS. THE UNMODIFIED, NORMATIVE VERSIONS OF KHRONOS SPECIFICATIONS AND
+** HEADER INFORMATION ARE LOCATED AT https://www.khronos.org/registry/
+**
+** THE MATERIALS ARE PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+** OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+** FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+** THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+** LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+** FROM,OUT OF OR IN CONNECTION WITH THE MATERIALS OR THE USE OR OTHER DEALINGS
+** IN THE MATERIALS.
+*/
+
+namespace NonSemanticAuxData {
+enum Instruction {
+  FunctionMetadata = 0,
+  FunctionAttribute = 1,
+  PreserveCount = 2
+};
+} // namespace NonSemanticAuxData

lib/SPIRV/libSPIRV/SPIRVEnum.h

@@ -80,6 +80,7 @@ enum SPIRVExtInstSetKind {
   SPIRVEIS_OpenCL_DebugInfo_100,
   SPIRVEIS_NonSemantic_Shader_DebugInfo_100,
   SPIRVEIS_NonSemantic_Shader_DebugInfo_200,
+  SPIRVEIS_NonSemantic_AuxData,
   SPIRVEIS_Count,
 };
 
@@ -135,6 +136,7 @@ template <> inline void SPIRVMap<SPIRVExtInstSetKind, std::string>::init() {
       "NonSemantic.Shader.DebugInfo.100");
   add(SPIRVEIS_NonSemantic_Shader_DebugInfo_200,
       "NonSemantic.Shader.DebugInfo.200");
+  add(SPIRVEIS_NonSemantic_AuxData, "NonSemantic.AuxData");
 }
 typedef SPIRVMap<SPIRVExtInstSetKind, std::string> SPIRVBuiltinSetNameMap;
 

lib/SPIRV/libSPIRV/SPIRVExtInst.h

@@ -40,6 +40,7 @@
 #ifndef SPIRV_LIBSPIRV_SPIRVEXTINST_H
 #define SPIRV_LIBSPIRV_SPIRVEXTINST_H
 
+#include "NonSemantic.AuxData.h"
 #include "OpenCL.std.h"
 #include "SPIRV.debug.h"
 #include "SPIRVEnum.h"
@@ -263,6 +264,16 @@ template <> inline void SPIRVMap<SPIRVDebugExtOpKind, std::string>::init() {
 }
 SPIRV_DEF_NAMEMAP(SPIRVDebugExtOpKind, SPIRVDebugExtOpMap)
 
+typedef NonSemanticAuxData::Instruction NonSemanticAuxDataOpKind;
+template <>
+inline void SPIRVMap<NonSemanticAuxDataOpKind, std::string>::init() {
+  add(NonSemanticAuxData::FunctionMetadata,
+      "NonSemanticAuxDataFunctionMetadata");
+  add(NonSemanticAuxData::FunctionAttribute,
+      "NonSemanticAuxDataFunctionAttribute");
+}
+SPIRV_DEF_NAMEMAP(NonSemanticAuxDataOpKind, NonSemanticAuxDataOpMap)
+
 } // namespace SPIRV
 
 #endif // SPIRV_LIBSPIRV_SPIRVEXTINST_H

lib/SPIRV/libSPIRV/SPIRVInstruction.h

@@ -42,6 +42,7 @@
 
 #include "SPIRVBasicBlock.h"
 #include "SPIRVEnum.h"
+#include "SPIRVFunction.h"
 #include "SPIRVIsValidEnum.h"
 #include "SPIRVOpCode.h"
 #include "SPIRVStream.h"
@@ -50,6 +51,7 @@
 #include <cassert>
 #include <functional>
 #include <iostream>
+#include <optional>
 #include <unordered_set>
 #include <utility>
 #include <vector>
@@ -1763,7 +1765,8 @@ class SPIRVExtInst : public SPIRVFunctionCallGeneric<OpExtInst, 5> {
     assert((ExtSetKind == SPIRVEIS_OpenCL || ExtSetKind == SPIRVEIS_Debug ||
             ExtSetKind == SPIRVEIS_OpenCL_DebugInfo_100 ||
             ExtSetKind == SPIRVEIS_NonSemantic_Shader_DebugInfo_100 ||
-            ExtSetKind == SPIRVEIS_NonSemantic_Shader_DebugInfo_200) &&
+            ExtSetKind == SPIRVEIS_NonSemantic_Shader_DebugInfo_200 ||
+            ExtSetKind == SPIRVEIS_NonSemantic_AuxData) &&
            "not supported");
   }
   void encode(spv_ostream &O) const override {
@@ -1778,6 +1781,9 @@ class SPIRVExtInst : public SPIRVFunctionCallGeneric<OpExtInst, 5> {
     case SPIRVEIS_NonSemantic_Shader_DebugInfo_200:
       getEncoder(O) << ExtOpDebug;
       break;
+    case SPIRVEIS_NonSemantic_AuxData:
+      getEncoder(O) << ExtOpNonSemanticAuxData;
+      break;
     default:
       assert(0 && "not supported");
       getEncoder(O) << ExtOp;
@@ -1797,6 +1803,9 @@ class SPIRVExtInst : public SPIRVFunctionCallGeneric<OpExtInst, 5> {
     case SPIRVEIS_NonSemantic_Shader_DebugInfo_200:
       getDecoder(I) >> ExtOpDebug;
       break;
+    case SPIRVEIS_NonSemantic_AuxData:
+      getDecoder(I) >> ExtOpNonSemanticAuxData;
+      break;
     default:
       assert(0 && "not supported");
       getDecoder(I) >> ExtOp;
@@ -1842,13 +1851,20 @@ class SPIRVExtInst : public SPIRVFunctionCallGeneric<OpExtInst, 5> {
     return ArgTypes;
   }
 
+  std::optional<ExtensionID> getRequiredExtension() const override {
+    if (SPIRVBuiltinSetNameMap::map(ExtSetKind).find("NonSemantic.") == 0)
+      return ExtensionID::SPV_KHR_non_semantic_info;
+    return {};
+  }
+
 protected:
   SPIRVExtInstSetKind ExtSetKind;
   SPIRVId ExtSetId;
   union {
     SPIRVWord ExtOp;
     OCLExtOpKind ExtOpOCL;
     SPIRVDebugExtOpKind ExtOpDebug;
+    NonSemanticAuxDataOpKind ExtOpNonSemanticAuxData;
   };
 };