[DirectX] Add extract-section to llvm-objcopy and implement it for DXContainer

llvm/docs/CommandGuide/llvm-objcopy.rst

@@ -79,6 +79,15 @@ multiple file formats.
  Enable deterministic mode when copying archives, i.e. use 0 for archive member
  header UIDs, GIDs and timestamp fields. On by default.
 
+.. option:: --extract-section <section>=<file>
+
+ Extract the specified section ``<section>`` into the file ``<file>`` as a
+ seperate object. Can be specified multiple times to extract multiple sections.
+ ``<file>`` is unrelated to the input and output files provided to
+ :program:`llvm-objcopy` and as such the normal copying and editing
+ operations will still be performed. No operations are performed on the sections
+ prior to dumping them.
+
 .. option:: --globalize-symbol <symbol>
 
  Mark any defined symbols named ``<symbol>`` as global symbols in the output.

llvm/include/llvm/ObjCopy/CommonConfig.h

@@ -233,6 +233,7 @@ struct CommonConfig {
   SmallVector<StringRef, 0> DumpSection;
   SmallVector<NewSectionInfo, 0> UpdateSection;
   SmallVector<SectionPatternAddressUpdate, 0> ChangeSectionAddress;
+  SmallVector<StringRef, 0> ExtractSection;
 
   // Section matchers
   NameMatcher KeepSection;

llvm/include/llvm/ObjCopy/ConfigManager.h

@@ -27,7 +27,7 @@ struct LLVM_ABI ConfigManager : public MultiFormatConfig {
 
   const CommonConfig &getCommonConfig() const override { return Common; }
 
-  Expected<const ELFConfig &> getELFConfig() const override { return ELF; }
+  Expected<const ELFConfig &> getELFConfig() const override;
 
   Expected<const COFFConfig &> getCOFFConfig() const override;
 

llvm/lib/ObjCopy/ConfigManager.cpp

@@ -13,6 +13,13 @@
 using namespace llvm;
 using namespace llvm::objcopy;
 
+Expected<const ELFConfig &> ConfigManager::getELFConfig() const {
+  if (!Common.ExtractSection.empty())
+    return createStringError(llvm::errc::invalid_argument,
+                             "option is not supported for ELF");
+  return ELF;
+}
+
 Expected<const COFFConfig &> ConfigManager::getCOFFConfig() const {
   if (!Common.SplitDWO.empty() || !Common.SymbolsPrefix.empty() ||
       !Common.SymbolsPrefixRemove.empty() || !Common.SymbolsToSkip.empty() ||
@@ -27,7 +34,7 @@ Expected<const COFFConfig &> ConfigManager::getCOFFConfig() const {
       Common.DiscardMode == DiscardType::Locals ||
       !Common.SymbolsToAdd.empty() || Common.GapFill != 0 ||
       Common.PadTo != 0 || Common.ChangeSectionLMAValAll != 0 ||
-      !Common.ChangeSectionAddress.empty())
+      !Common.ChangeSectionAddress.empty() || !Common.ExtractSection.empty())
     return createStringError(llvm::errc::invalid_argument,
                              "option is not supported for COFF");
 
@@ -48,7 +55,7 @@ Expected<const MachOConfig &> ConfigManager::getMachOConfig() const {
       Common.DiscardMode == DiscardType::Locals ||
       !Common.SymbolsToAdd.empty() || Common.GapFill != 0 ||
       Common.PadTo != 0 || Common.ChangeSectionLMAValAll != 0 ||
-      !Common.ChangeSectionAddress.empty())
+      !Common.ChangeSectionAddress.empty() || !Common.ExtractSection.empty())
     return createStringError(llvm::errc::invalid_argument,
                              "option is not supported for MachO");
 
@@ -69,7 +76,7 @@ Expected<const WasmConfig &> ConfigManager::getWasmConfig() const {
       !Common.SetSectionFlags.empty() || !Common.SetSectionType.empty() ||
       !Common.SymbolsToRename.empty() || Common.GapFill != 0 ||
       Common.PadTo != 0 || Common.ChangeSectionLMAValAll != 0 ||
-      !Common.ChangeSectionAddress.empty())
+      !Common.ChangeSectionAddress.empty() || !Common.ExtractSection.empty())
     return createStringError(llvm::errc::invalid_argument,
                              "only flags for section dumping, removal, and "
                              "addition are supported");
@@ -99,7 +106,7 @@ Expected<const XCOFFConfig &> ConfigManager::getXCOFFConfig() const {
       Common.Weaken || Common.StripUnneeded || Common.DecompressDebugSections ||
       Common.GapFill != 0 || Common.PadTo != 0 ||
       Common.ChangeSectionLMAValAll != 0 ||
-      !Common.ChangeSectionAddress.empty()) {
+      !Common.ChangeSectionAddress.empty() || !Common.ExtractSection.empty()) {
     return createStringError(
         llvm::errc::invalid_argument,
         "no flags are supported yet, only basic copying is allowed");
@@ -124,9 +131,8 @@ ConfigManager::getDXContainerConfig() const {
       Common.DecompressDebugSections || Common.GapFill != 0 ||
       Common.PadTo != 0 || Common.ChangeSectionLMAValAll != 0 ||
       !Common.ChangeSectionAddress.empty()) {
-    return createStringError(
-        llvm::errc::invalid_argument,
-        "no flags are supported yet, only basic copying is allowed");
+    return createStringError(llvm::errc::invalid_argument,
+                             "option is not supported for DXContainer");
   }
   return DXContainer;
 }

llvm/lib/ObjCopy/DXContainer/DXContainerObjcopy.cpp

@@ -11,14 +11,48 @@
 #include "DXContainerWriter.h"
 #include "llvm/ObjCopy/CommonConfig.h"
 #include "llvm/ObjCopy/DXContainer/DXContainerConfig.h"
+#include "llvm/Support/raw_ostream.h"
 
 namespace llvm {
 namespace objcopy {
 namespace dxbc {
 
 using namespace object;
 
+static Error extractPartAsObject(StringRef PartName, StringRef OutFilename,
+                                 StringRef InputFilename, const Object &Obj) {
+  for (const Part &P : Obj.Parts)
+    if (P.Name == PartName) {
+      Object PartObj;
+      PartObj.Header = Obj.Header;
+      PartObj.Parts.push_back({P.Name, P.Data});
+      PartObj.recomputeHeader();
+
+      auto Write = [&OutFilename, &PartObj](raw_ostream &Out) -> Error {
+        DXContainerWriter Writer(PartObj, Out);
+        if (Error E = Writer.write())
+          return createFileError(OutFilename, std::move(E));
+        return Error::success();
+      };
+
+      return writeToOutput(OutFilename, Write);
+    }
+
+  return createFileError(InputFilename, object_error::parse_failed,
+                         "part '%s' not found", PartName.str().c_str());
+}
+
 static Error handleArgs(const CommonConfig &Config, Object &Obj) {
+  // Extract all sections before any modifications.
+  for (StringRef Flag : Config.ExtractSection) {
+    StringRef SectionName;
+    StringRef FileName;
+    std::tie(SectionName, FileName) = Flag.split('=');
+    if (Error E = extractPartAsObject(SectionName, FileName,
+                                      Config.InputFilename, Obj))
+      return E;
+  }
+
   std::function<bool(const Part &)> RemovePred = [](const Part &) {
     return false;
   };

llvm/test/tools/llvm-objcopy/DXContainer/extract-section-basic.test

@@ -0,0 +1,111 @@
+## Tests that a separate DXContainer is created for the RTS0 (root signature)
+## part, when--extract-section is specified.
+
+# RUN: yaml2obj %s -o %t
+# RUN: llvm-objcopy %t --extract-section=RTS0=%t.rts0.out
+# RUN: obj2yaml %t.rts0.out | FileCheck %s --implicit-check-not=Name:
+
+## The DXContainer described below was generated with:
+
+## `clang-dxc -T cs_6_7 test.hlsl /Fo temp.dxo`
+## `obj2yaml temp.dxo`
+
+## and has the DXIL section trimmed for readability.
+
+## ``` test.hlsl
+## [RootSignature("")]
+## [numthreads(1,1,1)]
+## void main() {}
+## ```
+
+# CHECK:      Header:
+# CHECK-NEXT:   Hash:
+# CHECK:        Version:
+# CHECK-NEXT:      Major:           1
+# CHECK-NEXT:      Minor:           0
+# CHECK-NEXT:   FileSize:       68
+# CHECK-NEXT:   PartCount:     1
+# CHECK-NEXT:   PartOffsets:   [ 36 ]
+# CHECK-NEXT:   Parts:
+# CHECK-NEXT:   Name:            RTS0
+# CHECK-NEXT    Size:            24
+# CHECK-NEXT    RootSignature:
+# CHECK-NEXT      Version:         2
+# CHECK-NEXT      NumRootParameters: 0
+# CHECK-NEXT      RootParametersOffset: 24
+# CHECK-NEXT      NumStaticSamplers: 0
+# CHECK-NEXT      StaticSamplersOffset: 24
+# CHECK-NEXT      Parameters:      []
+
+--- !dxcontainer
+Header:
+  Hash:            [ 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
+                     0x0, 0x0, 0x0, 0x0, 0x0, 0x0 ]
+  Version:
+    Major:           1
+    Minor:           0
+  FileSize:        1984
+  PartCount:       7
+  PartOffsets:     [ 60, 1792, 1808, 1836, 1852, 1868, 1900 ]
+Parts:
+  - Name:            DXIL
+    Size:            1724
+    Program:
+      MajorVersion:    6
+      MinorVersion:    7
+      ShaderKind:      5
+      Size:            28
+      DXILMajorVersion: 1
+      DXILMinorVersion: 7
+      DXILSize:        4
+      DXIL:            [ 0x42, 0x43, 0xC0, 0xDE, ]
+  - Name:            SFI0
+    Size:            8
+  - Name:            HASH
+    Size:            20
+    Hash:
+      IncludesSource:  false
+      Digest:          [ 0x9F, 0xD1, 0xD9, 0xE2, 0x49, 0xFB, 0x3A, 0x6C,
+                         0x8C, 0x14, 0x8A, 0x96, 0x1C, 0x7D, 0x85, 0xA9 ]
+  - Name:            ISG1
+    Size:            8
+    Signature:
+      Parameters:      []
+  - Name:            OSG1
+    Size:            8
+    Signature:
+      Parameters:      []
+  - Name:            RTS0
+    Size:            24
+    RootSignature:
+      Version:         2
+      NumRootParameters: 0
+      RootParametersOffset: 24
+      NumStaticSamplers: 0
+      StaticSamplersOffset: 24
+      Parameters:      []
+  - Name:            PSV0
+    Size:            76
+    PSVInfo:
+      Version:         3
+      ShaderStage:     5
+      MinimumWaveLaneCount: 0
+      MaximumWaveLaneCount: 4294967295
+      UsesViewID:      0
+      SigInputVectors: 0
+      SigOutputVectors: [ 0, 0, 0, 0 ]
+      NumThreadsX:     1
+      NumThreadsY:     1
+      NumThreadsZ:     1
+      EntryName:       main
+      ResourceStride:  24
+      Resources:       []
+      SigInputElements: []
+      SigOutputElements: []
+      SigPatchOrPrimElements: []
+      InputOutputMap:
+        - [  ]
+        - [  ]
+        - [  ]
+        - [  ]
+...

llvm/test/tools/llvm-objcopy/DXContainer/extract-section-errs.test

@@ -0,0 +1,40 @@
+## Check that llvm-objcopy reports a suitable error when it can't find the
+## section to extract.
+
+## We can't extract a part that doesn't exist.
+# RUN: yaml2obj %s --docnum=1 -o %t1
+# RUN: not llvm-objcopy %t1 --extract-section=UNKNOWN=%t.unknown.out 2>&1 | FileCheck %s -DFILE=%t1 --check-prefix=ERROR1
+
+# ERROR1: error: '[[FILE]]': part 'UNKNOWN' not found
+
+--- !dxcontainer
+Header:
+  Hash:      [ 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
+               0x0, 0x0, 0x0, 0x0, 0x0, 0x0 ]
+  Version:
+    Major:   1
+    Minor:   0
+  PartCount: 1
+Parts:
+  - Name:    FKE0
+    Size:    8
+...
+
+## We can't extract a part that is specified incorrectly.
+# RUN: yaml2obj %s --docnum=2 -o %t2
+# RUN: not llvm-objcopy %t2 --extract-section=FKE0,%t.fke0.out 2>&1 | FileCheck %s -DFILE=%t2 --check-prefix=ERROR2
+
+# ERROR2: error: bad format for --extract-section, expected section=file
+
+--- !dxcontainer
+Header:
+  Hash:      [ 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
+               0x0, 0x0, 0x0, 0x0, 0x0, 0x0 ]
+  Version:
+    Major:   1
+    Minor:   0
+  PartCount: 1
+Parts:
+  - Name:    FKE0
+    Size:    8
+...

llvm/test/tools/llvm-objcopy/DXContainer/extract-section-headers.test

@@ -0,0 +1,53 @@
+## Tests that a separate DXContainer is created with only the specified section
+## for each --extract-section specified.
+
+# RUN: yaml2obj %s -o %t
+# RUN: llvm-objcopy %t --extract-section=FKE1=%t.fke1.out --extract-section=FKE4=%t.fke4.out
+# RUN: obj2yaml %t.fke1.out | FileCheck %s --check-prefixes=CHECK,FKE1 --implicit-check-not=Name:
+# RUN: obj2yaml %t.fke4.out | FileCheck %s --check-prefixes=CHECK,FKE4 --implicit-check-not=Name:
+
+# FKE1:       FileSize:       52
+# FKE4:       FileSize:       1732
+# CHECK-NEXT: PartCount:     1
+# CHECK-NEXT: PartOffsets:   [ 36 ]
+# CHECK-NEXT: Parts:
+# FKE1-NEXT:    Name: FKE1
+# FKE1-NEXT:    Size: 8
+# FKE4-NEXT:    Name: FKE4
+# FKE4-NEXT:    Size: 1688
+
+--- !dxcontainer
+Header:
+  Hash:        [ 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
+                 0x0, 0x0, 0x0, 0x0, 0x0, 0x0 ]
+  Version:
+    Major:     1
+    Minor:     0
+  FileSize:    1996
+  PartCount:   7
+  PartOffsets: [ 60, 76, 92, 108, 236, 1932, 1960 ]
+Parts:
+  - Name: FKE0
+    Size: 8
+  - Name: FKE1
+    Size: 8
+  - Name: FKE2
+    Size: 8
+  - Name: FKE3
+    Size: 120
+  - Name: FKE4
+    Size: 1688
+  - Name: FKE5
+    Size: 20
+  - Name: DXIL
+    Size: 28
+    Program:
+      MajorVersion:     6
+      MinorVersion:     5
+      ShaderKind:       5
+      Size:             8
+      DXILMajorVersion: 1
+      DXILMinorVersion: 5
+      DXILSize:         4
+      DXIL:             [ 0x42, 0x43, 0xC0, 0xDE, ]
+...

llvm/tools/llvm-objcopy/CommonOpts.td

@@ -23,6 +23,12 @@ def enable_deterministic_archives
     : Flag<["--"], "enable-deterministic-archives">,
       HelpText<"Enable deterministic mode when operating on archives (use "
                "zero for UIDs, GIDs, and timestamps).">;
+
+defm extract_section
+    : Eq<"extract-section",
+         "Extract section named <section> into standalone object in file <file>">,
+      MetaVarName<"section=file">;
+
 def D : Flag<["-"], "D">,
         Alias<enable_deterministic_archives>,
         HelpText<"Alias for --enable-deterministic-archives">;

llvm/tools/llvm-objcopy/ObjcopyOptions.cpp

@@ -1082,6 +1082,14 @@ objcopy::parseObjcopyOptions(ArrayRef<const char *> ArgsArr,
           "bad format for --dump-section, expected section=file");
     Config.DumpSection.push_back(Value);
   }
+  for (auto *Arg : InputArgs.filtered(OBJCOPY_extract_section)) {
+    StringRef Value(Arg->getValue());
+    if (Value.split('=').second.empty())
+      return createStringError(
+          errc::invalid_argument,
+          "bad format for --extract-section, expected section=file");
+    Config.ExtractSection.push_back(Value);
+  }
   Config.StripAll = InputArgs.hasArg(OBJCOPY_strip_all);
   Config.StripAllGNU = InputArgs.hasArg(OBJCOPY_strip_all_gnu);
   Config.StripDebug = InputArgs.hasArg(OBJCOPY_strip_debug);