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

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> SplitSection;
 
   // Section matchers
   NameMatcher KeepSection;

llvm/lib/ObjCopy/ConfigManager.cpp

@@ -27,7 +27,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.SplitSection.empty())
     return createStringError(llvm::errc::invalid_argument,
                              "option is not supported for COFF");
 
@@ -48,7 +48,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.SplitSection.empty())
     return createStringError(llvm::errc::invalid_argument,
                              "option is not supported for MachO");
 
@@ -69,7 +69,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.SplitSection.empty())
     return createStringError(llvm::errc::invalid_argument,
                              "only flags for section dumping, removal, and "
                              "addition are supported");
@@ -99,7 +99,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.SplitSection.empty()) {
     return createStringError(
         llvm::errc::invalid_argument,
         "no flags are supported yet, only basic copying is allowed");
@@ -126,9 +126,8 @@ ConfigManager::getDXContainerConfig() const {
       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");
   }
 
   // If a flag is listed here, then it has support for DXContainer:

llvm/lib/ObjCopy/DXContainer/DXContainerObjcopy.cpp

@@ -11,13 +11,37 @@
 #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 splitPartAsObject(StringRef PartName, StringRef OutFilename,
+                               StringRef InputFilename, const Object &Obj) {
+  for (const Part &P : Obj.Parts)
+    if (P.Name == PartName) {
+      auto PartObj = std::make_unique<Object>();
+      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) {
   std::function<bool(const Part &)> RemovePred = [](const Part &) {
     return false;
@@ -28,6 +52,32 @@ static Error handleArgs(const CommonConfig &Config, Object &Obj) {
       return Config.ToRemove.matches(P.Name);
     };
 
+  if (!Config.SplitSection.empty()) {
+    for (StringRef Flag : Config.SplitSection) {
+      StringRef SectionName;
+      StringRef FileName;
+      std::tie(SectionName, FileName) = Flag.split('=');
+
+      if (Error E = splitPartAsObject(SectionName, FileName,
+                                      Config.InputFilename, Obj))
+        return E;
+    }
+
+    RemovePred = [&Config, RemovePred](const Part &P) {
+      if (RemovePred(P))
+        return true;
+
+      for (StringRef Flag : Config.SplitSection) {
+        bool CanContain = Flag.size() > P.Name.size();
+        if (CanContain && Flag.starts_with(P.Name) &&
+            Flag[P.Name.size()] == '=')
+          return true;
+      }
+
+      return false;
+    };
+  }
+
   if (auto E = Obj.removeParts(RemovePred))
     return E;
 

llvm/lib/ObjCopy/DXContainer/DXContainerObject.h

@@ -42,8 +42,6 @@ struct Object {
   }
 
   Error removeParts(PartPred ToRemove);
-
-private:
   void recomputeHeader();
 };
 

llvm/test/tools/llvm-objcopy/DXContainer/errs-split.test

@@ -0,0 +1,21 @@
+## Check that llvm-objcopy reports a suitable error when it
+## can't find the section to split
+
+## We can't have multiple DXIL parts.
+# RUN: yaml2obj %s --docnum=1 -o %t1
+# RUN: not llvm-objcopy %t1 --split-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
+...

llvm/test/tools/llvm-objcopy/DXContainer/split-headers.test

@@ -0,0 +1,58 @@
+## Tests that a separate DXContainer is created with only the specified section
+## for each --split-section specified. Also ensure that the part is removed
+## from the original object.
+
+# RUN: yaml2obj %s -o %t
+# RUN: llvm-objcopy %t --split-section=FKE1=%t.fke1.out --split-section=FKE4=%t.fke4.out
+# RUN: obj2yaml %t | FileCheck %s --check-prefix=ORIG
+# RUN: obj2yaml %t.fke1.out | FileCheck %s --check-prefixes=CHECK,FKE1
+# RUN: obj2yaml %t.fke4.out | FileCheck %s --check-prefixes=CHECK,FKE4
+
+--- !dxcontainer
+Header:
+  Hash:            [ 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
+                     0x0, 0x0, 0x0, 0x0, 0x0, 0x0 ]
+  Version:
+    Major:           1
+    Minor:           0
+# ORIG: FileSize:       276
+# FKE1: FileSize:       52
+# FKE4: FileSize:       1732
+  FileSize:        1996
+# CHECK: PartCount:     1
+# CHECK: PartOffsets:   [ 36 ]
+  PartCount:       7
+  PartOffsets:     [ 60, 76, 92, 108, 236, 1932, 1960 ]
+Parts:
+# FKE1: Name:            FKE1
+# FKE4: Name:            FKE4
+# ORIG-NOT: FKE1
+# ORIG-NOT: FKE4
+# CHECK-NOT: FKE0
+# CHECK-NOT: FKE2
+# CHECK-NOT: FKE3
+# CHECK-NOT: FKE5
+  - 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, ]
+...