[HLSL][DirectX] Add support for rootsig as a target environment (llvm#156373)

This pr implements support for a root signature as a target, as specified
[here](https://github.com/llvm/wg-hlsl/blob/main/proposals/0029-root-signature-driver-options.md#target-root-signature-version).

This is implemented in the following steps:
1. Add `rootsignature` as a shader model environment type and define
`rootsig` as a `target_profile`. Only valid as versions 1.0 and 1.1
2. Updates `HLSLFrontendAction` to invoke a special handling of
constructing the `ASTContext` if we are considering an `hlsl` file and
with a `rootsignature` target
3. Defines the special handling to minimally instantiate the `Parser`
and `Sema` to insert the `RootSignatureDecl`
4. Updates `CGHLSLRuntime` to emit the constructed root signature decl
as part of `dx.rootsignatures` with a `null` entry function
5. Updates `DXILRootSignature` to handle emitting a root signature
without an entry function
6. Updates `ToolChains/HLSL` to invoke `only-section=RTS0` to strip any
other generated information

Resolves: llvm#150286.

##### Implementation Considerations
Ideally we could invoke this as part of `clang-dxc` without the need of
a source file. However, the initialization of the `Parser` and `Lexer`
becomes quite complicated to handle this.

Technically, we could avoid generating any of the extra information that
is removed in step 6. However, it seems better to re-use the logic in
`llvm-objcopy` without any need for additional custom logic in
`DXILRootSignature`.

Author: inbelic

clang/include/clang/Basic/Attr.td

@@ -1195,6 +1195,7 @@ static llvm::Triple::EnvironmentType getEnvironmentType(llvm::StringRef Environm
              .Case("callable", llvm::Triple::Callable)
              .Case("mesh", llvm::Triple::Mesh)
              .Case("amplification", llvm::Triple::Amplification)
+             .Case("rootsignature", llvm::Triple::RootSignature)
              .Case("library", llvm::Triple::Library)
              .Default(llvm::Triple::UnknownEnvironment);
 }

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -13152,6 +13152,8 @@ def err_hlsl_attribute_needs_intangible_type: Error<"attribute %0 can be used on
 def err_hlsl_incorrect_num_initializers: Error<
   "too %select{few|many}0 initializers in list for type %1 "
   "(expected %2 but found %3)">;
+def err_hlsl_rootsignature_entry: Error<
+   "rootsignature specified as target environment but entry, %0, was not defined">;
 
 def err_hlsl_operator_unsupported : Error<
   "the '%select{&|*|->}0' operator is unsupported in HLSL">;

clang/include/clang/Driver/Options.td

@@ -9444,7 +9444,8 @@ def target_profile : DXCJoinedOrSeparate<"T">, MetaVarName<"<profile>">,
          "cs_6_0, cs_6_1, cs_6_2, cs_6_3, cs_6_4, cs_6_5, cs_6_6, cs_6_7,"
          "lib_6_3, lib_6_4, lib_6_5, lib_6_6, lib_6_7, lib_6_x,"
          "ms_6_5, ms_6_6, ms_6_7,"
-         "as_6_5, as_6_6, as_6_7">;
+         "as_6_5, as_6_6, as_6_7,"
+         "rootsig_1_0, rootsig_1_1">;
 def emit_pristine_llvm : DXCFlag<"emit-pristine-llvm">,
   HelpText<"Emit pristine LLVM IR from the frontend by not running any LLVM passes at all."
            "Same as -S + -emit-llvm + -disable-llvm-passes.">;

clang/include/clang/Parse/ParseHLSLRootSignature.h

@@ -240,6 +240,8 @@ IdentifierInfo *ParseHLSLRootSignature(Sema &Actions,
                                        llvm::dxbc::RootSignatureVersion Version,
                                        StringLiteral *Signature);
 
+void HandleRootSignatureTarget(Sema &S, StringRef EntryRootSig);
+
 } // namespace hlsl
 } // namespace clang
 

clang/include/clang/Sema/SemaHLSL.h

@@ -159,6 +159,8 @@ class SemaHLSL : public SemaBase {
     RootSigOverrideIdent = DeclIdent;
   }
 
+  HLSLRootSignatureDecl *lookupRootSignatureOverrideDecl(DeclContext *DC) const;
+
   // Returns true if any RootSignatureElement is invalid and a diagnostic was
   // produced
   bool

clang/lib/CodeGen/CGHLSLRuntime.cpp

@@ -70,18 +70,18 @@ void addDxilValVersion(StringRef ValVersionStr, llvm::Module &M) {
   DXILValMD->addOperand(Val);
 }
 
-void addRootSignature(llvm::dxbc::RootSignatureVersion RootSigVer,
-                      ArrayRef<llvm::hlsl::rootsig::RootElement> Elements,
-                      llvm::Function *Fn, llvm::Module &M) {
+void addRootSignatureMD(llvm::dxbc::RootSignatureVersion RootSigVer,
+                        ArrayRef<llvm::hlsl::rootsig::RootElement> Elements,
+                        llvm::Function *Fn, llvm::Module &M) {
   auto &Ctx = M.getContext();
 
   llvm::hlsl::rootsig::MetadataBuilder RSBuilder(Ctx, Elements);
   MDNode *RootSignature = RSBuilder.BuildRootSignature();
 
   ConstantAsMetadata *Version = ConstantAsMetadata::get(ConstantInt::get(
       llvm::Type::getInt32Ty(Ctx), llvm::to_underlying(RootSigVer)));
-  MDNode *MDVals =
-      MDNode::get(Ctx, {ValueAsMetadata::get(Fn), RootSignature, Version});
+  ValueAsMetadata *EntryFunc = Fn ? ValueAsMetadata::get(Fn) : nullptr;
+  MDNode *MDVals = MDNode::get(Ctx, {EntryFunc, RootSignature, Version});
 
   StringRef RootSignatureValKey = "dx.rootsignatures";
   auto *RootSignatureValMD = M.getOrInsertNamedMetadata(RootSignatureValKey);
@@ -449,6 +449,19 @@ void CGHLSLRuntime::addBuffer(const HLSLBufferDecl *BufDecl) {
   }
 }
 
+void CGHLSLRuntime::addRootSignature(
+    const HLSLRootSignatureDecl *SignatureDecl) {
+  llvm::Module &M = CGM.getModule();
+  Triple T(M.getTargetTriple());
+
+  // Generated later with the function decl if not targeting root signature
+  if (T.getEnvironment() != Triple::EnvironmentType::RootSignature)
+    return;
+
+  addRootSignatureMD(SignatureDecl->getVersion(),
+                     SignatureDecl->getRootElements(), nullptr, M);
+}
+
 llvm::TargetExtType *
 CGHLSLRuntime::getHLSLBufferLayoutType(const RecordType *StructType) {
   const auto Entry = LayoutTypes.find(StructType);
@@ -685,8 +698,8 @@ void CGHLSLRuntime::emitEntryFunction(const FunctionDecl *FD,
   for (const Attr *Attr : FD->getAttrs()) {
     if (const auto *RSAttr = dyn_cast<RootSignatureAttr>(Attr)) {
       auto *RSDecl = RSAttr->getSignatureDecl();
-      addRootSignature(RSDecl->getVersion(), RSDecl->getRootElements(), EntryFn,
-                       M);
+      addRootSignatureMD(RSDecl->getVersion(), RSDecl->getRootElements(),
+                         EntryFn, M);
     }
   }
 }

clang/lib/CodeGen/CGHLSLRuntime.h

@@ -64,6 +64,7 @@ class VarDecl;
 class ParmVarDecl;
 class InitListExpr;
 class HLSLBufferDecl;
+class HLSLRootSignatureDecl;
 class HLSLVkBindingAttr;
 class HLSLResourceBindingAttr;
 class Type;
@@ -171,6 +172,7 @@ class CGHLSLRuntime {
   void generateGlobalCtorDtorCalls();
 
   void addBuffer(const HLSLBufferDecl *D);
+  void addRootSignature(const HLSLRootSignatureDecl *D);
   void finishCodeGen();
 
   void setHLSLEntryAttributes(const FunctionDecl *FD, llvm::Function *Fn);

clang/lib/CodeGen/CodeGenModule.cpp

@@ -7545,7 +7545,7 @@ void CodeGenModule::EmitTopLevelDecl(Decl *D) {
     break;
 
   case Decl::HLSLRootSignature:
-    // Will be handled by attached function
+    getHLSLRuntime().addRootSignature(cast<HLSLRootSignatureDecl>(D));
     break;
   case Decl::HLSLBuffer:
     getHLSLRuntime().addBuffer(cast<HLSLBufferDecl>(D));

clang/lib/Driver/ToolChains/HLSL.cpp

@@ -62,11 +62,15 @@ bool isLegalShaderModel(Triple &T) {
     VersionTuple MinVer(6, 5);
     return MinVer <= Version;
   } break;
+  case Triple::EnvironmentType::RootSignature:
+    VersionTuple MinVer(1, 0);
+    VersionTuple MaxVer(1, 1);
+    return MinVer <= Version && Version <= MaxVer;
   }
   return false;
 }
 
-std::optional<std::string> tryParseProfile(StringRef Profile) {
+std::optional<llvm::Triple> tryParseTriple(StringRef Profile) {
   // [ps|vs|gs|hs|ds|cs|ms|as]_[major]_[minor]
   SmallVector<StringRef, 3> Parts;
   Profile.split(Parts, "_");
@@ -84,6 +88,7 @@ std::optional<std::string> tryParseProfile(StringRef Profile) {
           .Case("lib", Triple::EnvironmentType::Library)
           .Case("ms", Triple::EnvironmentType::Mesh)
           .Case("as", Triple::EnvironmentType::Amplification)
+          .Case("rootsig", Triple::EnvironmentType::RootSignature)
           .Default(Triple::EnvironmentType::UnknownEnvironment);
   if (Kind == Triple::EnvironmentType::UnknownEnvironment)
     return std::nullopt;
@@ -147,8 +152,14 @@ std::optional<std::string> tryParseProfile(StringRef Profile) {
   T.setOSName(Triple::getOSTypeName(Triple::OSType::ShaderModel).str() +
               VersionTuple(Major, Minor).getAsString());
   T.setEnvironment(Kind);
-  if (isLegalShaderModel(T))
-    return T.getTriple();
+
+  return T;
+}
+
+std::optional<std::string> tryParseProfile(StringRef Profile) {
+  std::optional<llvm::Triple> MaybeT = tryParseTriple(Profile);
+  if (MaybeT && isLegalShaderModel(*MaybeT))
+    return MaybeT->getTriple();
   else
     return std::nullopt;
 }
@@ -258,6 +269,19 @@ bool checkExtensionArgsAreValid(ArrayRef<std::string> SpvExtensionArgs,
   }
   return AllValid;
 }
+
+bool isRootSignatureTarget(StringRef Profile) {
+  if (std::optional<llvm::Triple> T = tryParseTriple(Profile))
+    return T->getEnvironment() == Triple::EnvironmentType::RootSignature;
+  return false;
+}
+
+bool isRootSignatureTarget(DerivedArgList &Args) {
+  if (const Arg *A = Args.getLastArg(options::OPT_target_profile))
+    return isRootSignatureTarget(A->getValue());
+  return false;
+}
+
 } // namespace
 
 void tools::hlsl::Validator::ConstructJob(Compilation &C, const JobAction &JA,
@@ -317,6 +341,12 @@ void tools::hlsl::LLVMObjcopy::ConstructJob(Compilation &C, const JobAction &JA,
     CmdArgs.push_back(Frs);
   }
 
+  if (const Arg *A = Args.getLastArg(options::OPT_target_profile))
+    if (isRootSignatureTarget(A->getValue())) {
+      const char *Fos = Args.MakeArgString("--only-section=RTS0");
+      CmdArgs.push_back(Fos);
+    }
+
   assert(CmdArgs.size() > 2 && "Unnecessary invocation of objcopy.");
 
   C.addCommand(std::make_unique<Command>(JA, *this, ResponseFileSupport::None(),
@@ -493,7 +523,8 @@ bool HLSLToolChain::requiresBinaryTranslation(DerivedArgList &Args) const {
 
 bool HLSLToolChain::requiresObjcopy(DerivedArgList &Args) const {
   return Args.hasArg(options::OPT_dxc_Fo) &&
-         Args.hasArg(options::OPT_dxc_strip_rootsignature);
+         (Args.hasArg(options::OPT_dxc_strip_rootsignature) ||
+          isRootSignatureTarget(Args));
 }
 
 bool HLSLToolChain::isLastJob(DerivedArgList &Args,

clang/lib/Frontend/FrontendActions.cpp

@@ -1310,16 +1310,27 @@ void HLSLFrontendAction::ExecuteAction() {
                   /*CodeCompleteConsumer=*/nullptr);
   Sema &S = CI.getSema();
 
+  auto &TargetInfo = CI.getASTContext().getTargetInfo();
+  bool IsRootSignatureTarget =
+      TargetInfo.getTriple().getEnvironment() == llvm::Triple::RootSignature;
+  StringRef HLSLEntry = TargetInfo.getTargetOpts().HLSLEntry;
+
   // Register HLSL specific callbacks
   auto LangOpts = CI.getLangOpts();
+  StringRef RootSigName =
+      IsRootSignatureTarget ? HLSLEntry : LangOpts.HLSLRootSigOverride;
+
   auto MacroCallback = std::make_unique<InjectRootSignatureCallback>(
-      S, LangOpts.HLSLRootSigOverride, LangOpts.HLSLRootSigVer);
+      S, RootSigName, LangOpts.HLSLRootSigVer);
 
   Preprocessor &PP = CI.getPreprocessor();
   PP.addPPCallbacks(std::move(MacroCallback));
 
-  // Invoke as normal
-  WrapperFrontendAction::ExecuteAction();
+  // If we are targeting a root signature, invoke custom handling
+  if (IsRootSignatureTarget)
+    return hlsl::HandleRootSignatureTarget(S, HLSLEntry);
+  else // otherwise, invoke as normal
+    return WrapperFrontendAction::ExecuteAction();
 }
 
 HLSLFrontendAction::HLSLFrontendAction(