[DirectX backend] Translate cbuffer into DXIL.

[python3kgae]

Summary

cbuffer is a global ptr in address space 4 with annotation in named metadata "hlsl.cbufs".
A new pass DXILCBufferLoweringLegacy is added to collect users of cbuffer
and generate createHandle and cbufferLoad from the users of the cbuffer.

[python3kgae]

Rebase

[python3kgae]

Code Cleanup

[python3kgae]

Refactoring to share code for cbuffer and RWBuffer.

[python3kgae]

Refactoring collectZeroOffsetMemAccess for buffer lower.

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
+  DXILCBufferLowering.cpp
   DXILMetadata.cpp
   DXILOpBuilder.cpp
   DXILOpLowering.cpp
   DXILPrepare.cpp
   DXILTranslateMetadata.cpp
+  MemAccessLowerHelper.cpp
   PointerTypeAnalysis.cpp
 
   LINK_COMPONENTS
   AsmPrinter
   Core
   SelectionDAG
   Support
   DirectXInfo

llvm/lib/Target/DirectX/DXIL.td

 }
 
 def Unary : dxil_class<"Unary">;
 def Binary : dxil_class<"Binary">;
 def FlattenedThreadIdInGroupClass : dxil_class<"FlattenedThreadIdInGroup">;
 def ThreadIdInGroupClass : dxil_class<"ThreadIdInGroup">;
 def ThreadIdClass : dxil_class<"ThreadId">;
 def GroupIdClass : dxil_class<"GroupId">;
+def CBufferLoadClass : dxil_class<"CBufferLoad">;
+def CreateHandleClass : dxil_class<"CreateHandle">;
 
 def binary_uint : dxil_category<"Binary uint">;
 def unary_float : dxil_category<"Unary float">;
 def ComputeID : dxil_category<"Compute/Mesh/Amplification shader">;
-
+def Resources : dxil_category<"Resources">;
 
 // The parameter description for a DXIL instruction
 class dxil_param<int _pos, string type, string _name, string _doc,
                  bit _is_const = 0, string _enum_name = "",
                  int _max_value = 0> {
   int pos = _pos;           // position in parameter list
   string llvm_type = type; // llvm type name, $o for overload, $r for resource
                            // type, $cb for legacy cbuffer, $u4 for u4 struct

 
 def FlattenedThreadIdInGroup :dxil_op< "FlattenedThreadIdInGroup", 96,  FlattenedThreadIdInGroupClass, ComputeID,
    "provides a flattened index for a given thread within a given group (SV_GroupIndex)", "i32;",  "rn",
   [
     dxil_param<0,  "i32",  "",  "result">,
     dxil_param<1,  "i32",  "opcode",  "DXIL opcode">
   ]>,
   dxil_map_intrinsic<int_dx_flattened_thread_id_in_group>;
+
+
+def CBufferLoad : dxil_op< "CBufferLoad", 58, CBufferLoadClass,Resources, "loads a value from a constant buffer resource", "half;float;double;i8;i16;i32;i64;", "ro",
+  [
+    dxil_param<0, "$o", "", "the value for the constant buffer variable">,
+    dxil_param<1, "i32",  "opcode",  "DXIL opcode">,
+    dxil_param<2, "dx.types.Handle", "srv", "cbuffer handle">,
+    dxil_param<3, "i32", "byteOffset", "linear byte offset of value">,
+    dxil_param<4, "i32", "alignment", "load access alignment", 1>
+  ]>;
+
+def CreateHandle : dxil_op< "CreateHandle", 57, CreateHandleClass, Resources, "creates the handle to a resource",
+  "void;", "ro",
+  [
+    dxil_param<0, "dx.types.Handle", "", "the handle to the resource">,
+    dxil_param<1, "i32",  "opcode",  "DXIL opcode">,
+    dxil_param<2, "i8", "resourceClass", "the class of resource to create (SRV, UAV, CBuffer, Sampler)", 1>, // maps to DxilResourceBase::Class
+    dxil_param<3, "i32", "rangeId", "range identifier for resource", 1>,
+    dxil_param<4, "i32", "index", "zero-based index into range">,
+    dxil_param<5, "i1", "nonUniformIndex", "non-uniform resource index", 1>
+  ]>;

llvm/lib/Target/DirectX/DXILCBufferLowering.cpp

+//===- DXILCBufferLower.cpp - Lowering CBuffer to DIXL --------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file This file contains passes to lower cbuffer load to DXIL.
+//===----------------------------------------------------------------------===//
+
+#include "DXILConstants.h"
+#include "DXILOpBuilder.h"
+#include "DirectX.h"
+#include "MemAccessLowerHelper.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Analysis/Utils/Local.h"
+#include "llvm/CodeGen/Passes.h"
+#include "llvm/IR/GetElementPtrTypeIterator.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instruction.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/Intrinsics.h"
+#include "llvm/IR/IntrinsicsDirectX.h"
+#include "llvm/IR/Metadata.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/PassManager.h"
+#include "llvm/Pass.h"
+#include "llvm/Support/ErrorHandling.h"
+
+#define DEBUG_TYPE "dxil-cbuf-lower"
+
+using namespace llvm;
+using namespace llvm::dxil;
+
+static bool lowerCBufferAccess(Module &M) {
+  // FIXME: Allocate resource binding first.
+  const StringRef Name = "hlsl.cbufs";
+  auto *ResTable = M.getNamedMetadata(Name);
+  if (!ResTable)
+    return false;
+  const DataLayout &DL = M.getDataLayout();
+
+  for (auto *Res : ResTable->operands()) {
+    assert(Res->getNumOperands() == 5 && "invalid resource metadata");
+    auto *GVMD = cast<ValueAsMetadata>(Res->getOperand(0).get());
+    auto *GV = cast<GlobalVariable>(GVMD->getValue());
+    std::vector<MemAccessLowerHelper::MemAccess> AccessList;
+    MemAccessLowerHelper::collectZeroOffsetMemAccess(GV, AccessList, DL);
+
+    SmallDenseMap<Function *, CallInst *> HandleMap;
+    uint64_t RangeID =
+        mdconst::extract<ConstantInt>(Res->getOperand(2))->getLimitedValue();
+    ConstantInt *CBIndex = mdconst::extract<ConstantInt>(Res->getOperand(3));
+    assert(CBIndex->getLimitedValue() != -1ULL && "unallocated binding");
+    for (auto &Access : AccessList) {
+      Value *Index = Access.Index;
+      Instruction *User = Access.User;
+      // make handle in the Function.
+      Function *F = User->getParent()->getParent();
+      auto it = HandleMap.find(F);
+      CallInst *Hdl = nullptr;
+      if (it == HandleMap.end()) {
+        IRBuilder<> B(&*F->getEntryBlock().getFirstInsertionPt());
+        DXILOpBuilder DXILB(M, B);
+        Hdl = DXILB.createCreateHandle(
+            static_cast<unsigned>(dxil::ResourceClass::CBuffer), RangeID,
+            CBIndex, false);
+        HandleMap[F] = Hdl;
+      } else
+        Hdl = it->second;
+
+      LoadInst *LI = cast<LoadInst>(User);
+      Type *Ty = LI->getType();
+      Value *CBLd = nullptr;
+
+      IRBuilder<> B(LI);
+      DXILOpBuilder DXILB(M, B);
+      if (Ty->isIntegerTy() || Ty->isFloatingPointTy()) {
+        CBLd = DXILB.createCBufferLoad(Ty, Hdl, Index,
+                                       DL.getPrefTypeAlign(Ty).value());
+      } else if (isa<VectorType>(Ty)) {
+        // Only support fixed vectory type.
+        auto *VT = cast<FixedVectorType>(Ty);
+        Value *Result = PoisonValue::get(VT);
+        Type *EltTy = VT->getElementType();
+        uint64_t Align = DL.getPrefTypeAlign(EltTy).value();
+        for (unsigned i = 0; i < VT->getNumElements(); ++i) {
+          Value *Offset =
+              B.CreateAdd(Index, B.getInt32(i * DL.getTypeAllocSize(EltTy)));
+          Value *Elt = DXILB.createCBufferLoad(EltTy, Hdl, Offset, Align);
+          Result = B.CreateInsertElement(Result, Elt, i);
+        }
+        CBLd = Result;
+      } else {
+        assert(0 && "invalid type for cbuffer load.");
+      }
+      LI->replaceAllUsesWith(CBLd);
+      LI->eraseFromParent();
+    }
+  }
+  return true;
+}
+
+namespace {
+
+/// A pass that lowering cbuffer access into DXIL.
+class DXILCBufLowering : public PassInfoMixin<DXILCBufLowering> {
+public:
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &) {
+    if (lowerCBufferAccess(M))
+      return PreservedAnalyses::none();
+    return PreservedAnalyses::all();
+  }
+};
+} // namespace
+
+namespace {
+class DXILCBufferLoweringLegacy : public ModulePass {
+public:
+  bool runOnModule(Module &M) override { return lowerCBufferAccess(M); }
+  StringRef getPassName() const override { return "DXIL CBuffer lowering"; }
+  DXILCBufferLoweringLegacy() : ModulePass(ID) {}
+
+  static char ID; // Pass identification.
+};
+char DXILCBufferLoweringLegacy::ID = 0;
+
+} // end anonymous namespace
+
+INITIALIZE_PASS_BEGIN(DXILCBufferLoweringLegacy, DEBUG_TYPE,
+                      "DXIL CBuffer lowering", false, false)
+INITIALIZE_PASS_END(DXILCBufferLoweringLegacy, DEBUG_TYPE,
+                    "DXIL CBuffer lowering", false, false)
+
+ModulePass *llvm::createDXILCBufferLoweringLegacyPass() {
+  return new DXILCBufferLoweringLegacy();
+}

llvm/lib/Target/DirectX/DXILConstants.h

 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_LIB_TARGET_DIRECTX_DXILCONSTANTS_H
 #define LLVM_LIB_TARGET_DIRECTX_DXILCONSTANTS_H
 
 namespace llvm {
 namespace dxil {
 
+enum class ResourceClass { SRV = 0, UAV, CBuffer, Sampler, Invalid };
+
 #define DXIL_OP_ENUM
 #include "DXILOperation.inc"
 #undef DXIL_OP_ENUM
 
 } // namespace dxil
 } // namespace llvm
 
 #endif

llvm/lib/Target/DirectX/DXILOpBuilder.h

 public:
   DXILOpBuilder(Module &M, IRBuilderBase &B) : M(M), B(B) {}
   CallInst *createDXILOpCall(dxil::OpCode OpCode, Type *OverloadTy,
                              llvm::iterator_range<Use *> Args);
   Type *getOverloadTy(dxil::OpCode OpCode, FunctionType *FT,
                       bool NoOpCodeParam);
   static const char *getOpCodeName(dxil::OpCode DXILOp);
 
+  CallInst *createCreateHandle(int8_t ResClass, int RangeID, Value *Index,
+                               bool NonUniformIndex);
+  CallInst *createCBufferLoad(Type *OverloadTy, Value *Hdl, Value *ByteOffset,
+                              uint32_t Alignment);
+
 private:
   Module &M;
   IRBuilderBase &B;
 };
 
 } // namespace dxil
 } // namespace llvm
 
 #endif

llvm/lib/Target/DirectX/DXILOpBuilder.cpp

if (Kind == ParameterKind::CBUFFER_RET ||
     OverloadType = ST->getElementType(0);
   }
   return OverloadType;
 }
 
 const char *DXILOpBuilder::getOpCodeName(dxil::OpCode DXILOp) {
   return ::getOpCodeName(DXILOp);
 }
+
+CallInst *DXILOpBuilder::createCreateHandle(int8_t ResClass, int RangeID,
+                                            Value *Index,
+                                            bool NonUniformIndex) {
+  auto Fn =
+      getOrCreateDXILOpFunction(dxil::OpCode::CreateHandle, B.getVoidTy(), M);
+  return B.CreateCall(Fn, {B.getInt32((int32_t)dxil::OpCode::CreateHandle),
+                           B.getInt8(ResClass), B.getInt32(RangeID), Index,
+                           B.getInt1(NonUniformIndex)});
+}
+
+CallInst *DXILOpBuilder::createCBufferLoad(Type *OverloadTy, Value *Hdl,
+                                           Value *ByteOffset,
+                                           uint32_t Alignment) {
+  auto Fn = getOrCreateDXILOpFunction(dxil::OpCode::CBufferLoad, OverloadTy, M);
+  return B.CreateCall(Fn, {B.getInt32((int32_t)dxil::OpCode::CBufferLoad), Hdl,
+                           ByteOffset, B.getInt32(Alignment)});
+}
+
 } // namespace dxil
 } // namespace llvm

llvm/lib/Target/DirectX/DirectX.h

 ModulePass *createDXILPrepareModulePass();
 
 /// Initializer for DXILOpLowering
 void initializeDXILOpLoweringLegacyPass(PassRegistry &);
 
 /// Pass to lowering LLVM intrinsic call to DXIL op function call.
 ModulePass *createDXILOpLoweringLegacyPass();
 
+/// Initializer for DXILCBufferLowering
+void initializeDXILCBufferLoweringLegacyPass(PassRegistry &);
+
+/// Pass to lowering CBuffer to DXIL op function call.
+ModulePass *createDXILCBufferLoweringLegacyPass();
+
 /// Initializer for DXILTranslateMetadata.
 void initializeDXILTranslateMetadataPass(PassRegistry &);
 
 /// Pass to emit metadata for DXIL.
 ModulePass *createDXILTranslateMetadataPass();
 } // 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);
+  initializeDXILCBufferLoweringLegacyPass(*PR);
 }
 
 class DXILTargetObjectFile : public TargetLoweringObjectFile {
 public:
   DXILTargetObjectFile() = default;
 
   MCSection *getExplicitSectionGlobal(const GlobalObject *GO, SectionKind Kind,
                                       const TargetMachine &TM) const override {
DirectXPassConfig(DirectXTargetMachine &TM, PassManagerBase &PM)
       : TargetPassConfig(TM, PM) {}
 
   DirectXTargetMachine &getDirectXTargetMachine() const {
     return getTM<DirectXTargetMachine>();
   }
 
   FunctionPass *createTargetRegisterAllocator(bool) override { return nullptr; }
   void addCodeGenPrepare() override {
+    addPass(createDXILCBufferLoweringLegacyPass());
     addPass(createDXILOpLoweringLegacyPass());
     addPass(createDXILPrepareModulePass());
     addPass(createDXILTranslateMetadataPass());
   }
 };
 
 DirectXTargetMachine::DirectXTargetMachine(const Target &T, const Triple &TT,
                                            StringRef CPU, StringRef FS,

llvm/lib/Target/DirectX/MemAccessLowerHelper.h

+//===- Target/DirectX/MemAccessLowerHelper.h - Mem access helper --------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Utils to help lowering memory access for resources.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TARGET_DIRECTX_MEMACCESSLOWERHELPER_H
+#define LLVM_TARGET_DIRECTX_MEMACCESSLOWERHELPER_H
+
+#include <vector>
+
+namespace llvm {
+
+class Value;
+class Instruction;
+class GlobalVariable;
+class DataLayout;
+
+namespace dxil {
+
+namespace MemAccessLowerHelper {
+
+struct MemAccess {
+  Value *Index;    // Address when access memory.
+  uint32_t Offset; // Immediate Offset when access memory for legacy cbuffer and
+                   // structured buffer.
+  Instruction *User; // The instruction which access memory like Load/Store.
+};
+
+/// Collect all memory access for non-legacy cbuffer/ typed buffer global
+/// variable \c GV.
+void collectZeroOffsetMemAccess(Value *Ptr, std::vector<MemAccess> &AccessList,
+                                const DataLayout &DL);
+} // namespace MemAccessLowerHelper
+
+} // namespace dxil
+
+} // namespace llvm
+
+#endif // LLVM_TARGET_DIRECTX_MEMACCESSLOWERHELPER_H

llvm/lib/Target/DirectX/MemAccessLowerHelper.cpp

+//===- MemAccessLowerHelper.cpp - Mem access helper -------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+/// Utils to help lowering memory access for resources.
+//===----------------------------------------------------------------------===//
+
+#include "MemAccessLowerHelper.h"
+#include "llvm/Analysis/Utils/Local.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/GlobalVariable.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/Operator.h"
+
+using namespace llvm;
+using namespace llvm::dxil::MemAccessLowerHelper;
+
+static void collectZeroOffsetAccess(User *U, Value *Addr,
+                                    std::vector<MemAccess> &AccessList,
+                                    const DataLayout &DL) {
+  if (auto *GEP = dyn_cast<GEPOperator>(U)) {
+    // Calculate new Addr.
+    Value *NewAddr = Addr;
+    IRBuilder<> B(GEP->getContext());
+    if (auto *Inst = dyn_cast<Instruction>(GEP))
+      B.SetInsertPoint(Inst);
+    if (GEP->hasAllConstantIndices()) {
+      SmallVector<Value *> IdxList(GEP->idx_begin(), GEP->idx_end());
+      NewAddr = B.CreateAdd(Addr, B.getInt32(DL.getIndexedOffsetInType(
+                                      GEP->getSourceElementType(), IdxList)));
+    } else {
+      Value *Offset = EmitGEPOffset(&B, DL, GEP, /*NoAssumptions=*/true);
+      NewAddr = B.CreateAdd(Addr, Offset);
+    }
+
+    for (User *GEPU : GEP->users())
+      collectZeroOffsetAccess(GEPU, NewAddr, AccessList, DL);
+  } else if (isa<BitCastOperator>(U) || isa<AddrSpaceCastOperator>(U)) {
+    for (User *AU : U->users())
+      collectZeroOffsetAccess(AU, Addr, AccessList, DL);
+  } else if (auto *LI = dyn_cast<LoadInst>(U)) {
+    MemAccess Access = {Addr, 0, LI};
+    AccessList.emplace_back(Access);
+  } else if (auto *SI = dyn_cast<StoreInst>(U)) {
+    MemAccess Access = {Addr, 0, SI};
+    AccessList.emplace_back(Access);
+  } else
+    llvm_unreachable("unsupported user");
+}
+
+void llvm::dxil::MemAccessLowerHelper::collectZeroOffsetMemAccess(
+    Value *Ptr, std::vector<MemAccess> &AccessList, const DataLayout &DL) {
+  auto &Ctx = Ptr->getContext();
+  Value *OffsetZero = ConstantInt::get(Type::getInt32Ty(Ctx), 0);
+  for (User *U : Ptr->users())
+    collectZeroOffsetAccess(U, OffsetZero, AccessList, DL);
+}

llvm/test/CodeGen/DirectX/cbuf.ll

+; RUN: opt -S -dxil-cbuf-lower < %s | FileCheck %s
+target datalayout = "e-m:e-p:32:32-i1:32-i8:8-i16:16-i32:32-i64:64-f16:16-f32:32-f64:64-n8:16:32:64"
+target triple = "dxil-unknown-shadermodel6.7-library"
+
+; Make sure generate create handle.
+; CHECK:%[[HDL:.+]] = call %dx.types.Handle @dx.op.createHandle(i32 57, i8 2, i32 0, i32 2, i1 false)
+
+; Make sure load at offset 0/8/20 for @A.cb.float/double/<2 x i32>.y
+; CHECK:call float @dx.op.cbufferLoad.f32(i32 58, %dx.types.Handle %[[HDL]], i32 0, i32 4)
+; CHECK:call double @dx.op.cbufferLoad.f64(i32 58, %dx.types.Handle %[[HDL]], i32 8, i32 8)
+; CHECK:call i32 @dx.op.cbufferLoad.i32(i32 58, %dx.types.Handle %[[HDL]], i32 20, i32 4)
[e-mail]. = external constant { float, i32, double, <2 x i32> }
+
+; Function Attrs: noinline nounwind optnone
+define noundef float @"?foo@@YAMXZ"() #0 {
+entry:
+  %0 = load float, ptr @A.cb., align 4
+  %conv = fpext float %0 to double
+  %1 = load double, ptr getelementptr inbounds ({ float, i32, double, <2 x i32> }, ptr @A.cb., i32 0, i32 2), align 8
+  %2 = load <2 x i32>, ptr getelementptr inbounds ({ float, i32, double, <2 x i32> }, ptr @A.cb., i32 0, i32 3), align 8
+  %3 = extractelement <2 x i32> %2, i32 1
+  %conv1 = sitofp i32 %3 to double
+  %4 = call double @llvm.fmuladd.f64(double %1, double %conv1, double %conv)
+  %conv2 = fptrunc double %4 to float
+  ret float %conv2
+}
+
+; Function Attrs: nocallback nofree nosync nounwind readnone speculatable willreturn
+declare double @llvm.fmuladd.f64(double, double, double) #1
+
+attributes #0 = { noinline nounwind }
+attributes #1 = { nocallback nofree nosync nounwind readnone speculatable willreturn }
+
+!hlsl.cbufs = !{!1}
+
+!1 = !{ptr @A.cb., !"A.cb.ty", i32 0, i32 2, i32 1}