DXILWriterPass.cpp

Go to the documentation of this file.

//===- DXILWriterPass.cpp - Bitcode writing 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
//
//===----------------------------------------------------------------------===//
//
// DXILWriterPass implementation.
//
//===----------------------------------------------------------------------===//
 
#include "DXILWriterPass.h"
#include "DXILBitcodeWriter.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Analysis/ModuleSummaryAnalysis.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/DerivedTypes.h"
#include "llvm/IR/GlobalVariable.h"
#include "llvm/IR/IntrinsicInst.h"
#include "llvm/IR/Intrinsics.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/PassManager.h"
#include "llvm/InitializePasses.h"
#include "llvm/Pass.h"
#include "llvm/Support/Alignment.h"
#include "llvm/Transforms/Utils/ModuleUtils.h"
 
using namespace llvm;
using namespace llvm::dxil;
 
namespace {
class WriteDXILPass : public llvm::ModulePass {
  raw_ostream &OS; // raw_ostream to print on
 
public:
  static char ID; // Pass identification, replacement for typeid
  WriteDXILPass() : ModulePass(ID), OS(dbgs()) {
    initializeWriteDXILPassPass(*PassRegistry::getPassRegistry());
  }
 
  explicit WriteDXILPass(raw_ostream &o) : ModulePass(ID), OS(o) {
    initializeWriteDXILPassPass(*PassRegistry::getPassRegistry());
  }
 
  StringRef getPassName() const override { return "Bitcode Writer"; }
 
  bool runOnModule(Module &M) override {
    WriteDXILToFile(M, OS);
    return false;
  }
  void getAnalysisUsage(AnalysisUsage &AU) const override {
    AU.setPreservesAll();
  }
};
 
static void legalizeLifetimeIntrinsics(Module &M) {
  LLVMContext &Ctx = M.getContext();
  Type *I64Ty = IntegerType::get(Ctx, 64);
  Type *PtrTy = PointerType::get(Ctx, 0);
  Intrinsic::ID LifetimeIIDs[2] = {Intrinsic::lifetime_start,
                                   Intrinsic::lifetime_end};
  for (Intrinsic::ID &IID : LifetimeIIDs) {
    Function *F = M.getFunction(Intrinsic::getName(IID, {PtrTy}, &M));
    if (!F)
      continue;
 
    // Get or insert an LLVM 3.7-compliant lifetime intrinsic function of the
    // form `void @llvm.lifetime.[start/end](i64, ptr)` with the NoUnwind
    // attribute
    AttributeList Attr;
    Attr = Attr.addFnAttribute(Ctx, Attribute::NoUnwind);
    FunctionCallee LifetimeCallee = M.getOrInsertFunction(
        Intrinsic::getBaseName(IID), Attr, Type::getVoidTy(Ctx), I64Ty, PtrTy);
 
    // Replace all calls to lifetime intrinsics with calls to the
    // LLVM 3.7-compliant version of the lifetime intrinsic
    for (User *U : make_early_inc_range(F->users())) {
      CallInst *CI = dyn_cast<CallInst>(U);
      assert(CI &&
             "Expected user of a lifetime intrinsic function to be a CallInst");
 
      // LLVM 3.7 lifetime intrinics require an i8* operand, so we insert
      // a bitcast to ensure that is the case
      Value *PtrOperand = CI->getArgOperand(0);
      PointerType *PtrOpPtrTy = cast<PointerType>(PtrOperand->getType());
      Value *NoOpBitCast = CastInst::Create(Instruction::BitCast, PtrOperand,
                                            PtrOpPtrTy, "", CI->getIterator());
 
      // LLVM 3.7 lifetime intrinsics have an explicit size operand, whose value
      // we can obtain from the pointer operand which must be an AllocaInst (as
      // of https://github.com/llvm/llvm-project/pull/149310)
      AllocaInst *AI = dyn_cast<AllocaInst>(PtrOperand);
      assert(AI &&
             "The pointer operand of a lifetime intrinsic call must be an "
             "AllocaInst");
      std::optional<TypeSize> AllocSize =
          AI->getAllocationSize(CI->getDataLayout());
      assert(AllocSize.has_value() &&
             "Expected the allocation size of AllocaInst to be known");
      CallInst *NewCI = CallInst::Create(
          LifetimeCallee,
          {ConstantInt::get(I64Ty, AllocSize.value().getFixedValue()),
           NoOpBitCast},
          "", CI->getIterator());
      for (Attribute ParamAttr : CI->getParamAttributes(0))
        NewCI->addParamAttr(1, ParamAttr);
 
      CI->eraseFromParent();
    }
 
    F->eraseFromParent();
  }
}
 
static void removeLifetimeIntrinsics(Module &M) {
  Intrinsic::ID LifetimeIIDs[2] = {Intrinsic::lifetime_start,
                                   Intrinsic::lifetime_end};
  for (Intrinsic::ID &IID : LifetimeIIDs) {
    Function *F = M.getFunction(Intrinsic::getBaseName(IID));
    if (!F)
      continue;
 
    for (User *U : make_early_inc_range(F->users())) {
      CallInst *CI = dyn_cast<CallInst>(U);
      assert(CI && "Expected user of lifetime function to be a CallInst");
      BitCastInst *BCI = dyn_cast<BitCastInst>(CI->getArgOperand(1));
      assert(BCI && "Expected pointer operand of CallInst to be a BitCastInst");
      CI->eraseFromParent();
      BCI->eraseFromParent();
    }
    F->eraseFromParent();
  }
}
 
class EmbedDXILPass : public llvm::ModulePass {
public:
  static char ID; // Pass identification, replacement for typeid
  EmbedDXILPass() : ModulePass(ID) {
    initializeEmbedDXILPassPass(*PassRegistry::getPassRegistry());
  }
 
  StringRef getPassName() const override { return "DXIL Embedder"; }
 
  bool runOnModule(Module &M) override {
    std::string Data;
    llvm::raw_string_ostream OS(Data);
 
    Triple OriginalTriple = M.getTargetTriple();
    // Set to DXIL triple when write to bitcode.
    // Only the output bitcode need to be DXIL triple.
    M.setTargetTriple(Triple("dxil-ms-dx"));
 
    // Perform late legalization of lifetime intrinsics that would otherwise
    // fail the Module Verifier if performed in an earlier pass
    legalizeLifetimeIntrinsics(M);
 
    WriteDXILToFile(M, OS);
 
    // We no longer need lifetime intrinsics after bitcode serialization, so we
    // simply remove them to keep the Module Verifier happy after our
    // not-so-legal legalizations
    removeLifetimeIntrinsics(M);
 
    // Recover triple.
    M.setTargetTriple(OriginalTriple);
 
    Constant *ModuleConstant =
        ConstantDataArray::get(M.getContext(), arrayRefFromStringRef(Data));
    auto *GV = new llvm::GlobalVariable(M, ModuleConstant->getType(), true,
                                        GlobalValue::PrivateLinkage,
                                        ModuleConstant, "dx.dxil");
    GV->setSection("DXIL");
    GV->setAlignment(Align(4));
    appendToCompilerUsed(M, {GV});
    return true;
  }
 
  void getAnalysisUsage(AnalysisUsage &AU) const override {
    AU.setPreservesAll();
  }
};
} // namespace
 
char WriteDXILPass::ID = 0;
INITIALIZE_PASS_BEGIN(WriteDXILPass, "dxil-write-bitcode", "Write Bitcode",
                      false, true)
INITIALIZE_PASS_DEPENDENCY(ModuleSummaryIndexWrapperPass)
INITIALIZE_PASS_END(WriteDXILPass, "dxil-write-bitcode", "Write Bitcode", false,
                    true)
 
ModulePass *llvm::createDXILWriterPass(raw_ostream &Str) {
  return new WriteDXILPass(Str);
}
 
char EmbedDXILPass::ID = 0;
INITIALIZE_PASS(EmbedDXILPass, "dxil-embed", "Embed DXIL", false, true)
 
ModulePass *llvm::createDXILEmbedderPass() { return new EmbedDXILPass(); }