[HLSL][RootSignature] Introduce HLSLFrontendAction to implement rootsig-define (#154639)

This pr implements the functionality of `rootsig-define` as described
[here](https://github.com/llvm/wg-hlsl/blob/main/proposals/0029-root-signature-driver-options.md#option--rootsig-define).

This is accomplished by:
- Defining the `fdx-rootsignature-define`, and `rootsig-define` alias,
driver options. It simply specifies the name of a macro that will expand
to a `LiteralString` to be interpreted as a root signature.
- Introduces a new general frontend action wrapper,
`HLSLFrontendAction`. This class allows us to introduce `HLSL` specific
behaviour on the underlying action (primarily `ASTFrontendAction`).
Which will be further extended, or modularly wrapped, when considering
future DXC options.
- Using `HLSLFrontendAction` we can add a new `PPCallback` that will
eagerly parse the root signature specified with `rootsig-define` and
push it as a `TopLevelDecl` to `Sema`. This occurs when the macro has
been lexed.
- Since the root signature is parsed early, before any function
declarations, we can then simply attach it to the entry function once it
is encountered. Overwriting any applicable root signature attrs.

Resolves #150274

##### Implementation considerations

To implement this feature, note that:
1. We need access to all defined macros. These are created as part of
the first `Lex` in `Parser::Initialize` after `PP->EnterMainSourceFile`
2. `RootSignatureDecl` must be added to `Sema` before
`Consumer->HandleTranslationUnit` is invoked in `ParseAST`

Therefore, we can't handle the root signature in
`HLSLFrontendAction::ExecuteAction` before (from 1.) or after (from 2.)
invoking the underlying `ASTFrontendAction`.

This means we could alternatively:
- Manually handle this case
[here](https://github.com/llvm/llvm-project/blob/ac8f0bb070c9071742b6f6ce03bebc9d87217830/clang/lib/Parse/ParseAST.cpp#L168)
before parsing the first top level decl.
- Hook into when we [return the entry function
decl](https://github.com/llvm/llvm-project/blob/ac8f0bb070c9071742b6f6ce03bebc9d87217830/clang/lib/Parse/Parser.cpp#L1190)
and then parse the root signature and override its `RootSignatureAttr`.

The proposed solution handles this in the most modular way which should
work on any `FrontendAction` that might use the `Parser` without
invoking `ParseAST`, and, is not subject to needing to call the hook in
multiple different places of function declarators.

Author: inbelic

clang/include/clang/Basic/LangOptions.h

@@ -552,6 +552,10 @@ class LangOptions : public LangOptionsBase {
   llvm::dxbc::RootSignatureVersion HLSLRootSigVer =
       llvm::dxbc::RootSignatureVersion::V1_1;
 
+  /// The HLSL root signature that will be used to overide the root signature
+  /// used for the shader entry point.
+  std::string HLSLRootSigOverride;
+
   // Indicates if the wasm-opt binary must be ignored in the case of a
   // WebAssembly target.
   bool NoWasmOpt = false;

clang/include/clang/Driver/Options.td

@@ -9436,6 +9436,18 @@ def dxc_rootsig_ver :
   Alias<fdx_rootsignature_version>,
   Group<dxc_Group>,
   Visibility<[DXCOption]>;
+def fdx_rootsignature_define :
+  Joined<["-"], "fdx-rootsignature-define=">,
+  Group<dxc_Group>,
+  Visibility<[ClangOption, CC1Option]>,
+  MarshallingInfoString<LangOpts<"HLSLRootSigOverride">, "\"\"">,
+  HelpText<"Override entry function root signature with root signature at "
+           "given macro name.">;
+def dxc_rootsig_define :
+  Separate<["-"], "rootsig-define">,
+  Alias<fdx_rootsignature_define>,
+  Group<dxc_Group>,
+  Visibility<[DXCOption]>;
 def hlsl_entrypoint : Option<["-"], "hlsl-entry", KIND_SEPARATE>,
                       Group<dxc_Group>,
                       Visibility<[ClangOption, CC1Option]>,

clang/include/clang/Frontend/FrontendActions.h

@@ -329,6 +329,18 @@ class GetDependenciesByModuleNameAction : public PreprocessOnlyAction {
       : ModuleName(ModuleName) {}
 };
 
+//===----------------------------------------------------------------------===//
+// HLSL Specific Actions
+//===----------------------------------------------------------------------===//
+
+class HLSLFrontendAction : public WrapperFrontendAction {
+protected:
+  void ExecuteAction() override;
+
+public:
+  HLSLFrontendAction(std::unique_ptr<FrontendAction> WrappedAction);
+};
+
 }  // end namespace clang
 
 #endif

clang/include/clang/Parse/ParseHLSLRootSignature.h

@@ -236,6 +236,10 @@ class RootSignatureParser {
   RootSignatureToken CurToken;
 };
 
+IdentifierInfo *ParseHLSLRootSignature(Sema &Actions,
+                                       llvm::dxbc::RootSignatureVersion Version,
+                                       StringLiteral *Signature);
+
 } // namespace hlsl
 } // namespace clang
 

clang/include/clang/Sema/SemaHLSL.h

@@ -153,6 +153,10 @@ class SemaHLSL : public SemaBase {
   ActOnFinishRootSignatureDecl(SourceLocation Loc, IdentifierInfo *DeclIdent,
                                ArrayRef<hlsl::RootSignatureElement> Elements);
 
+  void SetRootSignatureOverride(IdentifierInfo *DeclIdent) {
+    RootSigOverrideIdent = DeclIdent;
+  }
+
   // Returns true if any RootSignatureElement is invalid and a diagnostic was
   // produced
   bool
@@ -221,6 +225,8 @@ class SemaHLSL : public SemaBase {
 
   uint32_t ImplicitBindingNextOrderID = 0;
 
+  IdentifierInfo *RootSigOverrideIdent = nullptr;
+
 private:
   void collectResourceBindingsOnVarDecl(VarDecl *D);
   void collectResourceBindingsOnUserRecordDecl(const VarDecl *VD,

clang/lib/CodeGen/CodeGenModule.cpp

@@ -7534,6 +7534,9 @@ void CodeGenModule::EmitTopLevelDecl(Decl *D) {
             getContext().getCanonicalTagType(cast<EnumDecl>(D)));
     break;
 
+  case Decl::HLSLRootSignature:
+    // Will be handled by attached function
+    break;
   case Decl::HLSLBuffer:
     getHLSLRuntime().addBuffer(cast<HLSLBufferDecl>(D));
     break;

clang/lib/Driver/ToolChains/Clang.cpp

@@ -3801,6 +3801,7 @@ static void RenderHLSLOptions(const ArgList &Args, ArgStringList &CmdArgs,
       options::OPT_disable_llvm_passes,
       options::OPT_fnative_half_type,
       options::OPT_hlsl_entrypoint,
+      options::OPT_fdx_rootsignature_define,
       options::OPT_fdx_rootsignature_version};
   if (!types::isHLSL(InputType))
     return;

clang/lib/Driver/ToolChains/HLSL.cpp

@@ -351,6 +351,13 @@ HLSLToolChain::TranslateArgs(const DerivedArgList &Args, StringRef BoundArch,
       A->claim();
       continue;
     }
+    if (A->getOption().getID() == options::OPT_dxc_rootsig_define) {
+      DAL->AddJoinedArg(nullptr,
+                        Opts.getOption(options::OPT_fdx_rootsignature_define),
+                        A->getValue());
+      A->claim();
+      continue;
+    }
     if (A->getOption().getID() == options::OPT__SLASH_O) {
       StringRef OStr = A->getValue();
       if (OStr == "d") {

clang/lib/Frontend/CompilerInvocation.cpp

@@ -640,6 +640,10 @@ static bool FixupInvocation(CompilerInvocation &Invocation,
     Diags.Report(diag::err_drv_argument_not_allowed_with)
         << "-fdx-rootsignature-version" << GetInputKindName(IK);
 
+  if (Args.hasArg(OPT_fdx_rootsignature_define) && !LangOpts.HLSL)
+    Diags.Report(diag::err_drv_argument_not_allowed_with)
+        << "-fdx-rootsignature-define" << GetInputKindName(IK);
+
   if (Args.hasArg(OPT_fgpu_allow_device_init) && !LangOpts.HIP)
     Diags.Report(diag::warn_ignored_hip_only_option)
         << Args.getLastArg(OPT_fgpu_allow_device_init)->getAsString(Args);

clang/lib/Frontend/FrontendActions.cpp

@@ -22,6 +22,7 @@
 #include "clang/Lex/HeaderSearch.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Lex/PreprocessorOptions.h"
+#include "clang/Parse/ParseHLSLRootSignature.h"
 #include "clang/Sema/TemplateInstCallback.h"
 #include "clang/Serialization/ASTReader.h"
 #include "clang/Serialization/ASTWriter.h"
@@ -1241,3 +1242,85 @@ void GetDependenciesByModuleNameAction::ExecuteAction() {
   PPCallbacks *CB = PP.getPPCallbacks();
   CB->moduleImport(SourceLocation(), Path, ModResult);
 }
+
+//===----------------------------------------------------------------------===//
+// HLSL Specific Actions
+//===----------------------------------------------------------------------===//
+
+class InjectRootSignatureCallback : public PPCallbacks {
+private:
+  Sema &Actions;
+  StringRef RootSigName;
+  llvm::dxbc::RootSignatureVersion Version;
+
+  std::optional<StringLiteral *> processStringLiteral(ArrayRef<Token> Tokens) {
+    for (Token Tok : Tokens)
+      if (!tok::isStringLiteral(Tok.getKind()))
+        return std::nullopt;
+
+    ExprResult StringResult = Actions.ActOnUnevaluatedStringLiteral(Tokens);
+    if (StringResult.isInvalid())
+      return std::nullopt;
+
+    if (auto Signature = dyn_cast<StringLiteral>(StringResult.get()))
+      return Signature;
+
+    return std::nullopt;
+  }
+
+public:
+  void MacroDefined(const Token &MacroNameTok,
+                    const MacroDirective *MD) override {
+    if (RootSigName != MacroNameTok.getIdentifierInfo()->getName())
+      return;
+
+    const MacroInfo *MI = MD->getMacroInfo();
+    auto Signature = processStringLiteral(MI->tokens());
+    if (!Signature.has_value()) {
+      Actions.getDiagnostics().Report(MI->getDefinitionLoc(),
+                                      diag::err_expected_string_literal)
+          << /*in attributes...*/ 4 << "RootSignature";
+      return;
+    }
+
+    IdentifierInfo *DeclIdent =
+        hlsl::ParseHLSLRootSignature(Actions, Version, *Signature);
+    Actions.HLSL().SetRootSignatureOverride(DeclIdent);
+  }
+
+  InjectRootSignatureCallback(Sema &Actions, StringRef RootSigName,
+                              llvm::dxbc::RootSignatureVersion Version)
+      : PPCallbacks(), Actions(Actions), RootSigName(RootSigName),
+        Version(Version) {}
+};
+
+void HLSLFrontendAction::ExecuteAction() {
+  // Pre-requisites to invoke
+  CompilerInstance &CI = getCompilerInstance();
+  if (!CI.hasASTContext() || !CI.hasPreprocessor())
+    return WrapperFrontendAction::ExecuteAction();
+
+  // InjectRootSignatureCallback requires access to invoke Sema to lookup/
+  // register a root signature declaration. The wrapped action is required to
+  // account for this by only creating a Sema if one doesn't already exist
+  // (like we have done, and, ASTFrontendAction::ExecuteAction)
+  if (!CI.hasSema())
+    CI.createSema(getTranslationUnitKind(),
+                  /*CodeCompleteConsumer=*/nullptr);
+  Sema &S = CI.getSema();
+
+  // Register HLSL specific callbacks
+  auto LangOpts = CI.getLangOpts();
+  auto MacroCallback = std::make_unique<InjectRootSignatureCallback>(
+      S, LangOpts.HLSLRootSigOverride, LangOpts.HLSLRootSigVer);
+
+  Preprocessor &PP = CI.getPreprocessor();
+  PP.addPPCallbacks(std::move(MacroCallback));
+
+  // Invoke as normal
+  WrapperFrontendAction::ExecuteAction();
+}
+
+HLSLFrontendAction::HLSLFrontendAction(
+    std::unique_ptr<FrontendAction> WrappedAction)
+    : WrapperFrontendAction(std::move(WrappedAction)) {}