[HLSL][DirectX] Add the Qdx-rootsignature-strip driver option (#154454)

This pr adds the `Qstrip-rootsignature` as a `DXC` driver option.

To do so, this pr introduces the `BinaryModifyJobClass` as an `Action`
to modify a produced object file before its final output.

Further, it registers `llvm-objcopy` as the tool to modify a produced
`DXContainer` on the `HLSL` toolchain.

This allows us to specify the `Qstrip-rootsignature` option to
`clang-dxc` which will invoke `llvm-objcopy` with a
`--remove-section=RTS0` argument to implement its functionality.

Resolves: #150275.

Author: inbelic

clang/include/clang/Driver/Action.h

@@ -76,9 +76,10 @@ class Action {
     StaticLibJobClass,
     BinaryAnalyzeJobClass,
     BinaryTranslatorJobClass,
+    ObjcopyJobClass,
 
     JobClassFirst = PreprocessJobClass,
-    JobClassLast = BinaryTranslatorJobClass
+    JobClassLast = ObjcopyJobClass
   };
 
   // The offloading kind determines if this action is binded to a particular
@@ -687,6 +688,17 @@ class BinaryTranslatorJobAction : public JobAction {
   }
 };
 
+class ObjcopyJobAction : public JobAction {
+  void anchor() override;
+
+public:
+  ObjcopyJobAction(Action *Input, types::ID Type);
+
+  static bool classof(const Action *A) {
+    return A->getKind() == ObjcopyJobClass;
+  }
+};
+
 } // namespace driver
 } // namespace clang
 

clang/include/clang/Driver/Options.td

@@ -9404,6 +9404,11 @@ def res_may_alias : Option<["/", "-"], "res-may-alias", KIND_FLAG>,
   Visibility<[DXCOption, ClangOption, CC1Option]>,
   HelpText<"Assume that UAVs/SRVs may alias">,
   MarshallingInfoFlag<CodeGenOpts<"ResMayAlias">>;
+def dxc_strip_rootsignature :
+  Option<["/", "-"], "Qstrip-rootsignature", KIND_FLAG>,
+  Group<dxc_Group>,
+  Visibility<[DXCOption]>,
+  HelpText<"Omit the root signature from produced DXContainer">;
 def target_profile : DXCJoinedOrSeparate<"T">, MetaVarName<"<profile>">,
   HelpText<"Set target profile">,
   Values<"ps_6_0, ps_6_1, ps_6_2, ps_6_3, ps_6_4, ps_6_5, ps_6_6, ps_6_7,"

clang/lib/Driver/Action.cpp

@@ -52,6 +52,8 @@ const char *Action::getClassName(ActionClass AC) {
     return "binary-analyzer";
   case BinaryTranslatorJobClass:
     return "binary-translator";
+  case ObjcopyJobClass:
+    return "objcopy";
   }
 
   llvm_unreachable("invalid class");
@@ -467,3 +469,8 @@ void BinaryTranslatorJobAction::anchor() {}
 BinaryTranslatorJobAction::BinaryTranslatorJobAction(Action *Input,
                                                      types::ID Type)
     : JobAction(BinaryTranslatorJobClass, Input, Type) {}
+
+void ObjcopyJobAction::anchor() {}
+
+ObjcopyJobAction::ObjcopyJobAction(Action *Input, types::ID Type)
+    : JobAction(ObjcopyJobClass, Input, Type) {}

clang/lib/Driver/Driver.cpp

@@ -4642,16 +4642,28 @@ void Driver::BuildDefaultActions(Compilation &C, DerivedArgList &Args,
     }
   }
 
-  // Call validator for dxil when -Vd not in Args.
   if (C.getDefaultToolChain().getTriple().isDXIL()) {
-    // Only add action when needValidation.
     const auto &TC =
         static_cast<const toolchains::HLSLToolChain &>(C.getDefaultToolChain());
+
+    // Call objcopy for manipulation of the unvalidated DXContainer when an
+    // option in Args requires it.
+    if (TC.requiresObjcopy(Args)) {
+      Action *LastAction = Actions.back();
+      // llvm-objcopy expects an unvalidated DXIL container (TY_OBJECT).
+      if (LastAction->getType() == types::TY_Object)
+        Actions.push_back(
+            C.MakeAction<ObjcopyJobAction>(LastAction, types::TY_Object));
+    }
+
+    // Call validator for dxil when -Vd not in Args.
     if (TC.requiresValidation(Args)) {
       Action *LastAction = Actions.back();
       Actions.push_back(C.MakeAction<BinaryAnalyzeJobAction>(
           LastAction, types::TY_DX_CONTAINER));
     }
+
+    // Call metal-shaderconverter when targeting metal.
     if (TC.requiresBinaryTranslation(Args)) {
       Action *LastAction = Actions.back();
       // Metal shader converter runs on DXIL containers, which can either be
@@ -6253,8 +6265,9 @@ const char *Driver::GetNamedOutputPath(Compilation &C, const JobAction &JA,
        C.getArgs().hasArg(options::OPT_dxc_Fo)) ||
       JA.getType() == types::TY_DX_CONTAINER) {
     StringRef FoValue = C.getArgs().getLastArgValue(options::OPT_dxc_Fo);
-    // If we are targeting DXIL and not validating or translating, we should set
-    // the final result file. Otherwise we should emit to a temporary.
+    // If we are targeting DXIL and not validating/translating/objcopying, we
+    // should set the final result file. Otherwise we should emit to a
+    // temporary.
     if (C.getDefaultToolChain().getTriple().isDXIL()) {
       const auto &TC = static_cast<const toolchains::HLSLToolChain &>(
           C.getDefaultToolChain());

clang/lib/Driver/ToolChain.cpp

@@ -652,6 +652,7 @@ Tool *ToolChain::getTool(Action::ActionClass AC) const {
   case Action::VerifyDebugInfoJobClass:
   case Action::BinaryAnalyzeJobClass:
   case Action::BinaryTranslatorJobClass:
+  case Action::ObjcopyJobClass:
     llvm_unreachable("Invalid tool kind.");
 
   case Action::CompileJobClass:

clang/lib/Driver/ToolChains/HLSL.cpp

@@ -294,6 +294,32 @@ void tools::hlsl::MetalConverter::ConstructJob(
                                          Exec, CmdArgs, Inputs, Input));
 }
 
+void tools::hlsl::LLVMObjcopy::ConstructJob(Compilation &C, const JobAction &JA,
+                                            const InputInfo &Output,
+                                            const InputInfoList &Inputs,
+                                            const ArgList &Args,
+                                            const char *LinkingOutput) const {
+
+  std::string ObjcopyPath = getToolChain().GetProgramPath("llvm-objcopy");
+  const char *Exec = Args.MakeArgString(ObjcopyPath);
+
+  ArgStringList CmdArgs;
+  assert(Inputs.size() == 1 && "Unable to handle multiple inputs.");
+  const InputInfo &Input = Inputs[0];
+  CmdArgs.push_back(Input.getFilename());
+  CmdArgs.push_back(Output.getFilename());
+
+  if (Args.hasArg(options::OPT_dxc_strip_rootsignature)) {
+    const char *Frs = Args.MakeArgString("--remove-section=RTS0");
+    CmdArgs.push_back(Frs);
+  }
+
+  assert(CmdArgs.size() > 2 && "Unnecessary invocation of objcopy.");
+
+  C.addCommand(std::make_unique<Command>(JA, *this, ResponseFileSupport::None(),
+                                         Exec, CmdArgs, Inputs, Input));
+}
+
 /// DirectX Toolchain
 HLSLToolChain::HLSLToolChain(const Driver &D, const llvm::Triple &Triple,
                              const ArgList &Args)
@@ -314,6 +340,10 @@ Tool *clang::driver::toolchains::HLSLToolChain::getTool(
     if (!MetalConverter)
       MetalConverter.reset(new tools::hlsl::MetalConverter(*this));
     return MetalConverter.get();
+  case Action::ObjcopyJobClass:
+    if (!LLVMObjcopy)
+      LLVMObjcopy.reset(new tools::hlsl::LLVMObjcopy(*this));
+    return LLVMObjcopy.get();
   default:
     return ToolChain::getTool(AC);
   }
@@ -458,16 +488,22 @@ bool HLSLToolChain::requiresBinaryTranslation(DerivedArgList &Args) const {
   return Args.hasArg(options::OPT_metal) && Args.hasArg(options::OPT_dxc_Fo);
 }
 
+bool HLSLToolChain::requiresObjcopy(DerivedArgList &Args) const {
+  return Args.hasArg(options::OPT_dxc_Fo) &&
+         Args.hasArg(options::OPT_dxc_strip_rootsignature);
+}
+
 bool HLSLToolChain::isLastJob(DerivedArgList &Args,
                               Action::ActionClass AC) const {
-  bool HasTranslation = requiresBinaryTranslation(Args);
-  bool HasValidation = requiresValidation(Args);
-  // If translation and validation are not required, we should only have one
-  // action.
-  if (!HasTranslation && !HasValidation)
-    return true;
-  if ((HasTranslation && AC == Action::BinaryTranslatorJobClass) ||
-      (!HasTranslation && HasValidation && AC == Action::BinaryAnalyzeJobClass))
-    return true;
-  return false;
+  // Note: we check in the reverse order of execution
+  if (requiresBinaryTranslation(Args))
+    return AC == Action::Action::BinaryTranslatorJobClass;
+  if (requiresValidation(Args))
+    return AC == Action::Action::BinaryAnalyzeJobClass;
+  if (requiresObjcopy(Args))
+    return AC == Action::Action::ObjcopyJobClass;
+
+  // No translation, validation, or objcopy are required, so this action must
+  // output to the result file.
+  return true;
 }

clang/lib/Driver/ToolChains/HLSL.h

@@ -42,6 +42,20 @@ class LLVM_LIBRARY_VISIBILITY MetalConverter : public Tool {
                     const llvm::opt::ArgList &TCArgs,
                     const char *LinkingOutput) const override;
 };
+
+class LLVM_LIBRARY_VISIBILITY LLVMObjcopy : public Tool {
+public:
+  LLVMObjcopy(const ToolChain &TC)
+      : Tool("hlsl::LLVMObjcopy", "llvm-objcopy", TC) {}
+
+  bool hasIntegratedCPP() const override { return false; }
+
+  void ConstructJob(Compilation &C, const JobAction &JA,
+                    const InputInfo &Output, const InputInfoList &Inputs,
+                    const llvm::opt::ArgList &TCArgs,
+                    const char *LinkingOutput) const override;
+};
+
 } // namespace hlsl
 } // namespace tools
 
@@ -65,6 +79,13 @@ class LLVM_LIBRARY_VISIBILITY HLSLToolChain : public ToolChain {
   static std::optional<std::string> parseTargetProfile(StringRef TargetProfile);
   bool requiresValidation(llvm::opt::DerivedArgList &Args) const;
   bool requiresBinaryTranslation(llvm::opt::DerivedArgList &Args) const;
+  bool requiresObjcopy(llvm::opt::DerivedArgList &Args) const;
+
+  /// If we are targeting DXIL then the last job should output the DXContainer
+  /// to the specified output file with /Fo. Otherwise, we will emit to a
+  /// temporary file for the next job to use.
+  ///
+  /// Returns true if we should output to the final result file.
   bool isLastJob(llvm::opt::DerivedArgList &Args, Action::ActionClass AC) const;
 
   // Set default DWARF version to 4 for DXIL uses version 4.
@@ -73,6 +94,7 @@ class LLVM_LIBRARY_VISIBILITY HLSLToolChain : public ToolChain {
 private:
   mutable std::unique_ptr<tools::hlsl::Validator> Validator;
   mutable std::unique_ptr<tools::hlsl::MetalConverter> MetalConverter;
+  mutable std::unique_ptr<tools::hlsl::LLVMObjcopy> LLVMObjcopy;
 };
 
 } // end namespace toolchains

clang/test/Driver/dxc_strip_rootsignature.hlsl

@@ -0,0 +1,15 @@
+// Create a dummy dxv to run
+// RUN: mkdir -p %t.dir
+// RUN: echo "dxv" > %t.dir/dxv && chmod 754 %t.dir/dxv
+
+// RUN: %clang_dxc -Qstrip-rootsignature --dxv-path=%t.dir -T cs_6_0 /Fo %t.dxo -### %s 2>&1 | FileCheck %s
+
+// Test to demonstrate that we specify to the root signature with the
+// -Qstrip-rootsignature option and that it occurs before DXV
+
+// CHECK: "{{.*}}llvm-objcopy{{(.exe)?}}" "{{.*}}.obj" "{{.*}}.obj" "--remove-section=RTS0"
+// CHECK: "{{.*}}dxv{{(.exe)?}}" "{{.*}}.obj" "-o" "{{.*}}.dxo"
+
+[shader("compute"), RootSignature("")]
+[numthreads(1,1,1)]
+void EmptyEntry() {}