[DX] Create globals for DXContainer parts

[beanz]

Summary

DXContainer files have a handful of sections that need to be written.
This adds a pass to write the section data into IR globals, and writes
the shader flag data into a global.

The test cases here verify that the shader flags are correctly written
from the IR into the global and emitted to the DXContainer.

This change also fixes a bug in the MCDXContainerWriter, where the size
of the dxbc::ProgramHeader was not being included in the part offset
calcuations. This is verified to be working by the new testcases where
obj2yaml can properly dump part data for parts after the DXIL part.

Resolves issue #57742 (https://github.com/llvm/llvm-project/issues/57742)

[python3kgae]

llvm/lib/Target/DirectX/DirectX.h

47

Do we need to call initializeShaderFlagsAnalysisWrapperPass somewhere?

[beanz]

Add call to initializeShaderFlagsAnalysisWrapperPass in LLVMInitializeDirectXTarget.

llvm/lib/MC/MCDXContainerWriter.cpp

if (SectionSize == 0)
       continue;
 
     assert(SectionSize < std::numeric_limits<uint32_t>::max() &&
            "Section size too large for DXContainer");
 
     PartOffsets.push_back(PartOffset);
     PartOffset += sizeof(dxbc::PartHeader) + SectionSize;
     PartOffset = alignTo(PartOffset, Align(4ul));
+    // The DXIL part also writes a program header, so we need to include its
+    // size when computing the offset for a part after the DXIL part.
+    if (Sec.getName() == "DXIL")
+      PartOffset += sizeof(dxbc::ProgramHeader);
   }
   assert(PartOffset < std::numeric_limits<uint32_t>::max() &&
          "Part data too large for DXContainer");
 
   uint64_t PartStart =
       sizeof(dxbc::Header) + (PartOffsets.size() * sizeof(uint32_t));
   uint64_t FileSize = PartStart + PartOffset;
   assert(FileSize < std::numeric_limits<uint32_t>::max() &&

llvm/lib/Target/DirectX/CMakeLists.txt

 add_public_tablegen_target(DirectXCommonTableGen)
 
 add_llvm_target(DirectXCodeGen
   DirectXAsmPrinter.cpp
   DirectXInstrInfo.cpp
   DirectXRegisterInfo.cpp
   DirectXSubtarget.cpp
   DirectXTargetMachine.cpp
+  DXContainerGlobals.cpp
   DXILMetadata.cpp
   DXILOpBuilder.cpp
   DXILOpLowering.cpp
   DXILPrepare.cpp
   DXILPrettyPrinter.cpp
   DXILResource.cpp
   DXILResourceAnalysis.cpp
   DXILShaderFlags.cpp

llvm/lib/Target/DirectX/DXContainerGlobals.cpp

+//===- DXContainerGlobals.cpp - DXContainer global generator pass ---------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// DXContainerGlobalsPass implementation.
+//
+//===----------------------------------------------------------------------===//
+
+#include "DXILShaderFlags.h"
+#include "DirectX.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/CodeGen/Passes.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Pass.h"
+#include "llvm/Transforms/Utils/ModuleUtils.h"
+
+using namespace llvm;
+using namespace llvm::dxil;
+
+namespace {
+class DXContainerGlobals : public llvm::ModulePass {
+
+public:
+  static char ID; // Pass identification, replacement for typeid
+  DXContainerGlobals() : ModulePass(ID) {
+    initializeDXContainerGlobalsPass(*PassRegistry::getPassRegistry());
+  }
+
+  StringRef getPassName() const override {
+    return "DXContainer Global Emitter";
+  }
+
+  bool runOnModule(Module &M) override;
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.setPreservesAll();
+    AU.addRequired<ShaderFlagsAnalysisWrapper>();
+  }
+};
+
+} // namespace
+
+bool DXContainerGlobals::runOnModule(Module &M) {
+  const uint64_t Flags =
+      (uint64_t)(getAnalysis<ShaderFlagsAnalysisWrapper>().getShaderFlags());
+
+  Constant *FlagsConstant = ConstantInt::get(M.getContext(), APInt(64, Flags));
+  auto *GV = new llvm::GlobalVariable(M, FlagsConstant->getType(), true,
+                                      GlobalValue::PrivateLinkage,
+                                      FlagsConstant, "dx.sfi0");
+  GV->setSection("SFI0");
+  GV->setAlignment(Align(4));
+  appendToCompilerUsed(M, {GV});
+  return true;
+}
+
+char DXContainerGlobals::ID = 0;
+INITIALIZE_PASS_BEGIN(DXContainerGlobals, "dxil-globals",
+                      "DXContainer Global Emitter", false, true)
+INITIALIZE_PASS_DEPENDENCY(ShaderFlagsAnalysisWrapper)
+INITIALIZE_PASS_END(DXContainerGlobals, "dxil-globals",
+                    "DXContainer Global Emitter", false, true)
+
+ModulePass *llvm::createDXContainerGlobalsPass() {
+  return new DXContainerGlobals();
+}

llvm/lib/Target/DirectX/DXILShaderFlags.h

class ShaderFlagsAnalysisPrinter
     : public PassInfoMixin<ShaderFlagsAnalysisPrinter> {
   raw_ostream &OS;
 
 public:
   explicit ShaderFlagsAnalysisPrinter(raw_ostream &OS) : OS(OS) {}
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
 };
 
+/// Wrapper pass for the legacy pass manager.
+///
+/// This is required because the passes that will depend on this are codegen
+/// passes which run through the legacy pass manager.
+class ShaderFlagsAnalysisWrapper : public ModulePass {
+  ComputedShaderFlags Flags;
+
+public:
+  static char ID;
+
+  ShaderFlagsAnalysisWrapper() : ModulePass(ID) {}
+
+  const ComputedShaderFlags &getShaderFlags() { return Flags; }
+
+  bool runOnModule(Module &M) override {
+    Flags = ComputedShaderFlags::computeFlags(M);
+    return false;
+  }
+};
+
 } // namespace dxil
 } // namespace llvm
 
 #endif // LLVM_TARGET_DIRECTX_DXILSHADERFLAGS_H

llvm/lib/Target/DirectX/DXILShaderFlags.cpp

 }
 
 PreservedAnalyses ShaderFlagsAnalysisPrinter::run(Module &M,
                                                   ModuleAnalysisManager &AM) {
   ComputedShaderFlags Flags = AM.getResult<ShaderFlagsAnalysis>(M);
   Flags.print(OS);
   return PreservedAnalyses::all();
 }
+
+char ShaderFlagsAnalysisWrapper::ID = 0;
+
+INITIALIZE_PASS(ShaderFlagsAnalysisWrapper, "dx-shader-flag-analysis",
+                "DXIL Shader Flag Analysis", true, true)

llvm/lib/Target/DirectX/DirectX.h

 
 /// Pass to emit metadata for DXIL.
 ModulePass *createDXILTranslateMetadataPass();
 
 /// Initializer for DXILTranslateMetadata.
 void initializeDXILResourceWrapperPass(PassRegistry &);
 
 /// Pass to pretty print DXIL metadata.
 ModulePass *createDXILPrettyPrinterPass(raw_ostream &OS);
 
 /// Initializer for DXILPrettyPrinter.
 void initializeDXILPrettyPrinterPass(PassRegistry &);
 
+/// Initializer for dxil::ShaderFlagsAnalysisWrapper pass.
+void initializeShaderFlagsAnalysisWrapperPass(PassRegistry &);
+
+/// Initializer for DXContainerGlobals pass.
+void initializeDXContainerGlobalsPass(PassRegistry &);
+
+/// Pass for generating DXContainer part globals.
+ModulePass *createDXContainerGlobalsPass();
 } // namespace llvm
 
 #endif // LLVM_LIB_TARGET_DIRECTX_DIRECTX_H

llvm/lib/Target/DirectX/DirectXTargetMachine.cpp

 extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeDirectXTarget() {
   RegisterTargetMachine<DirectXTargetMachine> X(getTheDirectXTarget());
   auto *PR = PassRegistry::getPassRegistry();
   initializeDXILPrepareModulePass(*PR);
   initializeEmbedDXILPassPass(*PR);
   initializeDXILOpLoweringLegacyPass(*PR);
   initializeDXILTranslateMetadataPass(*PR);
   initializeDXILResourceWrapperPass(*PR);
+  initializeShaderFlagsAnalysisWrapperPass(*PR);
 }
 
 class DXILTargetObjectFile : public TargetLoweringObjectFile {
 public:
   DXILTargetObjectFile() = default;
 
   MCSection *getExplicitSectionGlobal(const GlobalObject *GO, SectionKind Kind,
                                       const TargetMachine &TM) const override {
bool DirectXTargetMachine::addPassesToEmitFile(
     PassManagerBase &PM, raw_pwrite_stream &Out, raw_pwrite_stream *DwoOut,
     CodeGenFileType FileType, bool DisableVerify,
     MachineModuleInfoWrapperPass *MMIWP) {
   TargetPassConfig *PassConfig = createPassConfig(PM);
   PassConfig->addCodeGenPrepare();
 
   if (TargetPassConfig::willCompleteCodeGenPipeline()) {
     PM.add(createDXILEmbedderPass());
+    // We embed the other DXContainer globals after embedding DXIL so that the
+    // globals don't pollute the DXIL.
+    PM.add(createDXContainerGlobalsPass());
   }
   switch (FileType) {
   case CGFT_AssemblyFile:
     PM.add(createDXILPrettyPrinterPass(Out));
     PM.add(createPrintModulePass(Out, "", true));
     break;
   case CGFT_ObjectFile:
     if (TargetPassConfig::willCompleteCodeGenPipeline()) {

llvm/test/CodeGen/DirectX/ShaderFlags/double-extensions.ll

 ; RUN: opt -S --passes="print-dx-shader-flags" 2>&1 %s | FileCheck %s
+; RUN: llc %s --filetype=obj -o - | obj2yaml | FileCheck %s --check-prefix=DXC
 
 target triple = "dxil-pc-shadermodel6.7-library"
 
 ; CHECK: ; Shader Flags Value: 0x00000021
 ; CHECK: ; Note: shader requires additional functionality:
 ; CHECK-NEXT: ;       Double-precision floating point
 ; CHECK-NEXT: ;       Double-precision extensions for 11.1
 ; CHECK-NEXT: {{^;$}}
 define double @div(double %a, double %b) {
   %res = fdiv double %a, %b
   ret double %res
 }
+
+
+; DXC: - Name:            SFI0
+; DXC-NEXT:     Size:            16
+; DXC-NEXT:     Flags:
+; DXC-NEXT:       Doubles:         true
+; DXC-NOT:   {{[A-Za-z]+: +true}}
+; DXC:            DX11_1_DoubleExtensions:         true
+; DXC-NOT:   {{[A-Za-z]+: +true}}
+; DXC:       NextUnusedBit:   false
+; DXC: ...

llvm/test/CodeGen/DirectX/ShaderFlags/doubles.ll

 ; RUN: opt -S --passes="print-dx-shader-flags" 2>&1 %s | FileCheck %s
+; RUN: llc %s --filetype=obj -o - | obj2yaml | FileCheck %s --check-prefix=DXC
 
 target triple = "dxil-pc-shadermodel6.7-library"
 
 ; CHECK: ; Shader Flags Value: 0x00000001
 ; CHECK: ; Note: shader requires additional functionality:
 ; CHECK-NEXT: ;       Double-precision floating point
 ; CHECK-NEXT: {{^;$}}
 define double @add(double %a, double %b) {
   %sum = fadd double %a, %b
   ret double %sum
 }
+
+; DXC: - Name:            SFI0
+; DXC-NEXT:     Size:            16
+; DXC-NEXT:     Flags:
+; DXC-NEXT:       Doubles:         true
+; DXC-NOT:   {{[A-Za-z]+: +true}}
+; DXC:       NextUnusedBit:   false
+; DXC: ...

llvm/test/CodeGen/DirectX/embed-dxil.ll

 ; RUN: llc %s --filetype=asm -o - | FileCheck %s
 ; RUN: opt %s -dxil-embed -S -o - | FileCheck %s
 ; RUN: llc %s --filetype=obj -o - | obj2yaml | FileCheck %s --check-prefix=DXC
 target triple = "dxil-unknown-shadermodel6.5-library"
 
 define i32 @add(i32 %a, i32 %b) {
   %sum = add i32 %a, %b
   ret i32 %sum
 }
 
 ; CHECK: @dx.dxil = private constant [[BC_TYPE:\[[0-9]+ x i8\]]] c"BC\C0\DE{{[^"]+}}", section "DXIL", align 4
-; CHECK: @llvm.compiler.used = appending global [1 x ptr] [ptr @dx.dxil], section "llvm.metadata"
+
+; The dxil global should be the first here because we generate it before the
+; other globals. If it isn't the first here, that's probably a bug.
+; CHECK: @llvm.compiler.used = appending global {{\[[0-9]+ x ptr\]}} [ptr @dx.dxil
 
 ; This is using regex matches on some sizes, offsets and fields. These are all
 ; going to change as the DirectX backend continues to evolve and implement more
 ; features. Rather than extending this test to cover those future features, this
 ; test's matches are extremely fuzzy so that it won't break.
 
 ; DXC: --- !dxcontainer
 ; DXC-NEXT: Header:

 ; DXC-NEXT:       MajorVersion:    6
 ; DXC-NEXT:       MinorVersion:    5
 ; DXC-NEXT:       ShaderKind:      6
 ; DXC-NEXT:       Size:            [[#div(SIZE,4) - 2]]
 ; DXC-NEXT:       DXILMajorVersion: [[#]]
 ; DXC-NEXT:       DXILMinorVersion: [[#]]
 ; DXC-NEXT:       DXILSize:        [[#SIZE - 32]]
 ; DXC-NEXT:       DXIL:            [ 0x42, 0x43, 0xC0, 0xDE,
+; DXC:      - Name:            SFI0
+; DXC-NEXT:   Size:            16
+; DXC-NOT:    Flags: