[HLSL][DirectX] Add support for rootsig as a target environment

clang/include/clang/Basic/Attr.td

@@ -1185,6 +1185,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

@@ -13147,6 +13147,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

@@ -9432,7 +9432,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

@@ -157,6 +157,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

@@ -69,18 +69,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);
@@ -448,6 +448,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);
@@ -651,8 +664,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

@@ -62,6 +62,7 @@ class VarDecl;
 class ParmVarDecl;
 class InitListExpr;
 class HLSLBufferDecl;
+class HLSLRootSignatureDecl;
 class HLSLVkBindingAttr;
 class HLSLResourceBindingAttr;
 class Type;
@@ -151,6 +152,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

@@ -7537,7 +7537,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

@@ -1309,16 +1309,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(

clang/lib/Parse/ParseHLSLRootSignature.cpp

@@ -7,8 +7,9 @@
 //===----------------------------------------------------------------------===//
 
 #include "clang/Parse/ParseHLSLRootSignature.h"
-
+#include "clang/AST/ASTConsumer.h"
 #include "clang/Lex/LiteralSupport.h"
+#include "clang/Parse/Parser.h"
 #include "clang/Sema/Sema.h"
 
 using namespace llvm::hlsl::rootsig;
@@ -1472,5 +1473,38 @@ IdentifierInfo *ParseHLSLRootSignature(Sema &Actions,
   return DeclIdent;
 }
 
+void HandleRootSignatureTarget(Sema &S, StringRef EntryRootSig) {
+  ASTConsumer *Consumer = &S.getASTConsumer();
+
+  // Minimally initalize the parser. This does a couple things:
+  // - initializes Sema scope handling
+  // - invokes HLSLExternalSemaSource
+  // - invokes the preprocessor to lex the macros in the file
+  std::unique_ptr<Parser> P(new Parser(S.getPreprocessor(), S, true));
+  S.getPreprocessor().EnterMainSourceFile();
+
+  bool HaveLexer = S.getPreprocessor().getCurrentLexer();
+  if (HaveLexer) {
+    P->Initialize();
+    S.ActOnStartOfTranslationUnit();
+
+    // Skim through the file to parse to find the define
+    while (P->getCurToken().getKind() != tok::eof)
+      P->ConsumeAnyToken();
+
+    HLSLRootSignatureDecl *SignatureDecl =
+        S.HLSL().lookupRootSignatureOverrideDecl(
+            S.getASTContext().getTranslationUnitDecl());
+
+    if (SignatureDecl)
+      Consumer->HandleTopLevelDecl(DeclGroupRef(SignatureDecl));
+    else
+      S.getDiagnostics().Report(diag::err_hlsl_rootsignature_entry)
+          << EntryRootSig;
+  }
+
+  Consumer->HandleTranslationUnit(S.getASTContext());
+}
+
 } // namespace hlsl
 } // namespace clang

clang/lib/Sema/SemaHLSL.cpp

@@ -729,19 +729,15 @@ void SemaHLSL::ActOnTopLevelFunction(FunctionDecl *FD) {
 
   // If we have specified a root signature to override the entry function then
   // attach it now
-  if (RootSigOverrideIdent) {
-    LookupResult R(SemaRef, RootSigOverrideIdent, SourceLocation(),
-                   Sema::LookupOrdinaryName);
-    if (SemaRef.LookupQualifiedName(R, FD->getDeclContext()))
-      if (auto *SignatureDecl =
-              dyn_cast<HLSLRootSignatureDecl>(R.getFoundDecl())) {
-        FD->dropAttr<RootSignatureAttr>();
-        // We could look up the SourceRange of the macro here as well
-        AttributeCommonInfo AL(RootSigOverrideIdent, AttributeScopeInfo(),
-                               SourceRange(), ParsedAttr::Form::Microsoft());
-        FD->addAttr(::new (getASTContext()) RootSignatureAttr(
-            getASTContext(), AL, RootSigOverrideIdent, SignatureDecl));
-      }
+  HLSLRootSignatureDecl *SignatureDecl =
+      lookupRootSignatureOverrideDecl(FD->getDeclContext());
+  if (SignatureDecl) {
+    FD->dropAttr<RootSignatureAttr>();
+    // We could look up the SourceRange of the macro here as well
+    AttributeCommonInfo AL(RootSigOverrideIdent, AttributeScopeInfo(),
+                           SourceRange(), ParsedAttr::Form::Microsoft());
+    FD->addAttr(::new (getASTContext()) RootSignatureAttr(
+        getASTContext(), AL, RootSigOverrideIdent, SignatureDecl));
   }
 
   llvm::Triple::EnvironmentType Env = TargetInfo.getTriple().getEnvironment();
@@ -765,6 +761,8 @@ void SemaHLSL::ActOnTopLevelFunction(FunctionDecl *FD) {
     case llvm::Triple::UnknownEnvironment:
     case llvm::Triple::Library:
       break;
+    case llvm::Triple::RootSignature:
+      llvm_unreachable("rootsig environment has no functions");
     default:
       llvm_unreachable("Unhandled environment in triple");
     }
@@ -827,6 +825,8 @@ void SemaHLSL::CheckEntryPoint(FunctionDecl *FD) {
       }
     }
     break;
+  case llvm::Triple::RootSignature:
+    llvm_unreachable("rootsig environment has no function entry point");
   default:
     llvm_unreachable("Unhandled environment in triple");
   }
@@ -1107,6 +1107,18 @@ void SemaHLSL::ActOnFinishRootSignatureDecl(
   SemaRef.PushOnScopeChains(SignatureDecl, SemaRef.getCurScope());
 }
 
+HLSLRootSignatureDecl *
+SemaHLSL::lookupRootSignatureOverrideDecl(DeclContext *DC) const {
+  if (RootSigOverrideIdent) {
+    LookupResult R(SemaRef, RootSigOverrideIdent, SourceLocation(),
+                   Sema::LookupOrdinaryName);
+    if (SemaRef.LookupQualifiedName(R, DC))
+      return dyn_cast<HLSLRootSignatureDecl>(R.getFoundDecl());
+  }
+
+  return nullptr;
+}
+
 namespace {
 
 struct PerVisibilityBindingChecker {

clang/test/AST/HLSL/RootSignature-Target-AST.hlsl

@@ -0,0 +1,28 @@
+// RUN: %clang_cc1 -triple dxil-pc-shadermodel6.0-rootsignature -ast-dump \
+// RUN:  -hlsl-entry EntryRootSig -disable-llvm-passes -o - %s | FileCheck %s --check-prefixes=CHECK,CHECK-V1_1
+
+// RUN: %clang_cc1 -triple dxil-pc-shadermodel6.0-rootsignature -ast-dump \
+// RUN:  -fdx-rootsignature-version=rootsig_1_0 \
+// RUN:  -hlsl-entry EntryRootSig -disable-llvm-passes -o - %s | FileCheck %s --check-prefixes=CHECK,CHECK-V1_0
+
+// RUN: %clang_cc1 -triple dxil-pc-shadermodel6.0-rootsignature -ast-dump \
+// RUN:  -D CmdRS='"UAV(u0)"'\
+// RUN:  -hlsl-entry CmdRS -disable-llvm-passes -o - %s | FileCheck %s --check-prefix=CMD
+
+// CHECK: -HLSLRootSignatureDecl 0x{{.*}} {{.*}} implicit [[ENTRY_RS_DECL:__hlsl_rootsig_decl_\d*]]
+// CHECK-V1_0-SAME: version: 1.0,
+// CHECK-V1_1-SAME: version: 1.1,
+// CHECK-SAME: RootElements{
+// CHECK-SAME: RootCBV(b0,
+// CHECK-SAME:   space = 0, visibility = All,
+// CHECK-V1_0-SAME: flags = DataVolatile
+// CHECK-V1_1-SAME: flags = DataStaticWhileSetAtExecute
+// CHECK-SAME: )
+// CHECK-SAME: }
+#define EntryRootSig "CBV(b0)"
+
+// CMD: -HLSLRootSignatureDecl 0x{{.*}} {{.*}} implicit [[CMD_RS_DECL:__hlsl_rootsig_decl_\d*]]
+// CMD-SAME: version: 1.1,
+// CMD-SAME: RootElements{
+// CMD-SAME: RootUAV(u0, space = 0, visibility = All, flags = DataVolatile)
+// CMD-SAME: }

clang/test/CodeGenHLSL/RootSignature-Target.hlsl

@@ -0,0 +1,9 @@
+// RUN: %clang_cc1 -triple dxil-pc-shadermodel6.3-rootsignature \
+// RUN: -hlsl-entry EntryRS -emit-llvm -o - %s | FileCheck %s
+
+// CHECK: !dx.rootsignatures = !{![[#ENTRY:]]}
+// CHECK: ![[#ENTRY]] = !{null, ![[#ENTRY_RS:]], i32 2}
+// CHECK: ![[#ENTRY_RS]] = !{![[#ROOT_CBV:]]}
+// CHECK: ![[#ROOT_CBV]] = !{!"RootCBV", i32 0, i32 0, i32 0, i32 4}
+
+#define EntryRS "CBV(b0)"