[DirectX] Implement the ForwardHandleAccesses pass

llvm/include/llvm/Analysis/DXILResource.h

@@ -196,6 +196,24 @@ class SamplerExtType : public TargetExtType {
   }
 };
 
+class AnyResourceExtType : public TargetExtType {
+public:
+  AnyResourceExtType() = delete;
+  AnyResourceExtType(const AnyResourceExtType &) = delete;
+  AnyResourceExtType &operator=(const AnyResourceExtType &) = delete;
+
+  static bool classof(const TargetExtType *T) {
+    return isa<RawBufferExtType>(T) || isa<TypedBufferExtType>(T) ||
+           isa<TextureExtType>(T) || isa<MSTextureExtType>(T) ||
+           isa<FeedbackTextureExtType>(T) || isa<CBufferExtType>(T) ||
+           isa<SamplerExtType>(T);
+  }
+
+  static bool classof(const Type *T) {
+    return isa<TargetExtType>(T) && classof(cast<TargetExtType>(T));
+  }
+};
+
 /// The dx.Layout target extension type
 ///
 /// `target("dx.Layout", <Type>, <size>, [offsets...])`

llvm/lib/Target/DirectX/CMakeLists.txt

@@ -23,6 +23,7 @@ add_llvm_target(DirectXCodeGen
   DXILCBufferAccess.cpp
   DXILDataScalarization.cpp
   DXILFinalizeLinkage.cpp
+  DXILForwardHandleAccesses.cpp
   DXILFlattenArrays.cpp
   DXILIntrinsicExpansion.cpp
   DXILOpBuilder.cpp

llvm/lib/Target/DirectX/DXILForwardHandleAccesses.cpp

@@ -0,0 +1,165 @@
+//===- DXILForwardHandleAccesses.cpp - Cleanup Handles --------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "DXILForwardHandleAccesses.h"
+#include "DXILShaderFlags.h"
+#include "DirectX.h"
+#include "llvm/Analysis/DXILResource.h"
+#include "llvm/Analysis/Loads.h"
+#include "llvm/IR/DiagnosticInfo.h"
+#include "llvm/IR/Dominators.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/Intrinsics.h"
+#include "llvm/IR/IntrinsicsDirectX.h"
+#include "llvm/IR/Module.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Pass.h"
+#include "llvm/Transforms/Utils/Local.h"
+
+#define DEBUG_TYPE "dxil-forward-handle-accesses"
+
+using namespace llvm;
+
+static void diagnoseAmbiguousHandle(IntrinsicInst *NewII,
+                                    IntrinsicInst *PrevII) {
+  Function *F = NewII->getFunction();
+  LLVMContext &Context = F->getParent()->getContext();
+  Context.diagnose(DiagnosticInfoGeneric(
+      Twine("Handle at \"") + NewII->getName() + "\" overwrites handle at \"" +
+      PrevII->getName() + "\""));
+}
+
+static void diagnoseHandleNotFound(LoadInst *LI) {
+  Function *F = LI->getFunction();
+  LLVMContext &Context = F->getParent()->getContext();
+  Context.diagnose(DiagnosticInfoGeneric(
+      LI, Twine("Load of \"") + LI->getPointerOperand()->getName() +
+              "\" is not a global resource handle"));
+}
+
+static void diagnoseUndominatedLoad(LoadInst *LI, IntrinsicInst *Handle) {
+  Function *F = LI->getFunction();
+  LLVMContext &Context = F->getParent()->getContext();
+  Context.diagnose(DiagnosticInfoGeneric(
+      LI, Twine("Load at \"") + LI->getName() +
+              "\" is not dominated by handle creation at \"" +
+              Handle->getName() + "\""));
+}
+
+static void
+processHandle(IntrinsicInst *II,
+              DenseMap<GlobalVariable *, IntrinsicInst *> &HandleMap) {
+  for (User *U : II->users())
+    if (auto *SI = dyn_cast<StoreInst>(U))
+      if (auto *GV = dyn_cast<GlobalVariable>(SI->getPointerOperand())) {
+        auto Entry = HandleMap.try_emplace(GV, II);
+        if (Entry.second)
+          LLVM_DEBUG(dbgs() << "Added " << GV->getName() << " to handle map\n");
+        else
+          diagnoseAmbiguousHandle(II, Entry.first->second);
+      }
+}
+
+static bool forwardHandleAccesses(Function &F, DominatorTree &DT) {
+  bool Changed = false;
+
+  DenseMap<GlobalVariable *, IntrinsicInst *> HandleMap;
+  SmallVector<LoadInst *> LoadsToProcess;
+  for (BasicBlock &BB : F)
+    for (Instruction &Inst : BB)
+      if (auto *II = dyn_cast<IntrinsicInst>(&Inst)) {
+        switch (II->getIntrinsicID()) {
+        case Intrinsic::dx_resource_handlefrombinding:
+          processHandle(II, HandleMap);
+          break;
+        default:
+          continue;
+        }
+      } else if (auto *LI = dyn_cast<LoadInst>(&Inst))
+        if (isa<dxil::AnyResourceExtType>(LI->getType()))
+          LoadsToProcess.push_back(LI);
+
+  for (LoadInst *LI : LoadsToProcess) {
+    Value *V = LI->getPointerOperand();
+    auto *GV = dyn_cast<GlobalVariable>(LI->getPointerOperand());
+
+    // If we didn't find the global, we may need to walk through a level of
+    // indirection. This generally happens at -O0.
+    if (!GV)
+      if (auto *NestedLI = dyn_cast<LoadInst>(V)) {
+        BasicBlock::iterator BBI(NestedLI);
+        Value *Loaded = FindAvailableLoadedValue(
+            NestedLI, NestedLI->getParent(), BBI, 0, nullptr, nullptr);
+        GV = dyn_cast_or_null<GlobalVariable>(Loaded);
+      }
+
+    auto It = HandleMap.find(GV);
+    if (It == HandleMap.end()) {
+      diagnoseHandleNotFound(LI);
+      continue;
+    }
+    Changed = true;
+
+    if (!DT.dominates(It->second, LI)) {
+      diagnoseUndominatedLoad(LI, It->second);
+      continue;
+    }
+
+    LLVM_DEBUG(dbgs() << "Replacing uses of " << GV->getName() << " at "
+                      << LI->getName() << " with " << It->second->getName()
+                      << "\n");
+    LI->replaceAllUsesWith(It->second);
+    LI->eraseFromParent();
+  }
+
+  return Changed;
+}
+
+PreservedAnalyses DXILForwardHandleAccesses::run(Function &F,
+                                                 FunctionAnalysisManager &AM) {
+  PreservedAnalyses PA;
+
+  DominatorTree *DT = &AM.getResult<DominatorTreeAnalysis>(F);
+  bool Changed = forwardHandleAccesses(F, *DT);
+
+  if (!Changed)
+    return PreservedAnalyses::all();
+  return PA;
+}
+
+namespace {
+class DXILForwardHandleAccessesLegacy : public FunctionPass {
+public:
+  bool runOnFunction(Function &F) override {
+    DominatorTree *DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
+    return forwardHandleAccesses(F, *DT);
+  }
+  StringRef getPassName() const override {
+    return "DXIL Forward Handle Accesses";
+  }
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<DominatorTreeWrapperPass>();
+  }
+
+  DXILForwardHandleAccessesLegacy() : FunctionPass(ID) {}
+
+  static char ID; // Pass identification.
+};
+char DXILForwardHandleAccessesLegacy::ID = 0;
+} // end anonymous namespace
+
+INITIALIZE_PASS_BEGIN(DXILForwardHandleAccessesLegacy, DEBUG_TYPE,
+                      "DXIL Forward Handle Accesses", false, false)
+INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
+INITIALIZE_PASS_END(DXILForwardHandleAccessesLegacy, DEBUG_TYPE,
+                    "DXIL Forward Handle Accesses", false, false)
+
+FunctionPass *llvm::createDXILForwardHandleAccessesLegacyPass() {
+  return new DXILForwardHandleAccessesLegacy();
+}

llvm/lib/Target/DirectX/DXILForwardHandleAccesses.h

@@ -0,0 +1,28 @@
+//===- DXILForwardHandleAccesses.h - Cleanup Handles ------------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// \file Eliminate redundant stores and loads from handle globals.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIB_TARGET_DIRECTX_DXILFORWARDHANDLEACCESS_H
+#define LLVM_LIB_TARGET_DIRECTX_DXILFORWARDHANDLEACCESS_H
+
+#include "llvm/IR/PassManager.h"
+
+namespace llvm {
+
+class DXILForwardHandleAccesses
+    : public PassInfoMixin<DXILForwardHandleAccesses> {
+public:
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+};
+
+} // namespace llvm
+
+#endif // LLVM_LIB_TARGET_DIRECTX_DXILFORWARDHANDLEACCESS_H

llvm/lib/Target/DirectX/DirectX.h

@@ -53,6 +53,12 @@ void initializeDXILFlattenArraysLegacyPass(PassRegistry &);
 /// Pass to flatten arrays into a one dimensional DXIL legal form
 ModulePass *createDXILFlattenArraysLegacyPass();
 
+/// Initializer for DXIL Forward Handle Accesses Pass
+void initializeDXILForwardHandleAccessesLegacyPass(PassRegistry &);
+
+/// Pass to eliminate redundant stores and loads from handle globals.
+FunctionPass *createDXILForwardHandleAccessesLegacyPass();
+
 /// Initializer DXIL legalizationPass
 void initializeDXILLegalizeLegacyPass(PassRegistry &);
 

llvm/lib/Target/DirectX/DirectXPassRegistry.def

@@ -38,6 +38,7 @@ MODULE_PASS("print<dxil-root-signature>", dxil::RootSignatureAnalysisPrinter(dbg
 #ifndef FUNCTION_PASS
 #define FUNCTION_PASS(NAME, CREATE_PASS)
 #endif
+FUNCTION_PASS("dxil-forward-handle-accesses", DXILForwardHandleAccesses())
 FUNCTION_PASS("dxil-resource-access", DXILResourceAccess())
 FUNCTION_PASS("dxil-legalize", DXILLegalizePass())
 #undef FUNCTION_PASS

llvm/lib/Target/DirectX/DirectXTargetMachine.cpp

@@ -15,6 +15,7 @@
 #include "DXILCBufferAccess.h"
 #include "DXILDataScalarization.h"
 #include "DXILFlattenArrays.h"
+#include "DXILForwardHandleAccesses.h"
 #include "DXILIntrinsicExpansion.h"
 #include "DXILLegalizePass.h"
 #include "DXILOpLowering.h"
@@ -66,6 +67,7 @@ extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeDirectXTarget() {
   initializeRootSignatureAnalysisWrapperPass(*PR);
   initializeDXILFinalizeLinkageLegacyPass(*PR);
   initializeDXILPrettyPrinterLegacyPass(*PR);
+  initializeDXILForwardHandleAccessesLegacyPass(*PR);
   initializeDXILCBufferAccessLegacyPass(*PR);
 }
 
@@ -105,6 +107,7 @@ class DirectXPassConfig : public TargetPassConfig {
     ScalarizerPassOptions DxilScalarOptions;
     DxilScalarOptions.ScalarizeLoadStore = true;
     addPass(createScalarizerPass(DxilScalarOptions));
+    addPass(createDXILForwardHandleAccessesLegacyPass());
     addPass(createDXILLegalizeLegacyPass());
     addPass(createDXILTranslateMetadataLegacyPass());
     addPass(createDXILOpLoweringLegacyPass());

llvm/test/CodeGen/DirectX/ForwardHandleAccesses/alloca.ll

@@ -0,0 +1,20 @@
+; RUN: not opt -S -dxil-forward-handle-accesses -mtriple=dxil--shadermodel6.3-library %s 2>&1 | FileCheck %s
+
+; CHECK: error: Load of "buf" is not a global resource handle
+
+%"class.hlsl::RWStructuredBuffer" = type { target("dx.RawBuffer", <4 x float>, 1, 0) }
+@Buf = internal global %"class.hlsl::RWStructuredBuffer" poison, align 4
+
+define float @f() {
+entry:
+  %buf = alloca target("dx.RawBuffer", <4 x float>, 1, 0), align 4
+  %h = call target("dx.RawBuffer", <4 x float>, 1, 0) @llvm.dx.resource.handlefrombinding(i32 0, i32 0, i32 1, i32 0, i1 false)
+  store target("dx.RawBuffer", <4 x float>, 1, 0) %h, ptr %buf, align 4
+
+  %b = load target("dx.RawBuffer", <4 x float>, 1, 0), ptr %buf, align 4
+  %l = call { <4 x float>, i1 } @llvm.dx.resource.load.rawbuffer(target("dx.RawBuffer", <4 x float>, 1, 0) %b, i32 0, i32 0)
+  %x = extractvalue { <4 x float>, i1 } %l, 0
+  %v = extractelement <4 x float> %x, i32 0
+
+  ret float %v
+}

llvm/test/CodeGen/DirectX/ForwardHandleAccesses/ambiguous.ll

@@ -0,0 +1,21 @@
+; RUN: not opt -S -dxil-forward-handle-accesses -mtriple=dxil--shadermodel6.3-library %s 2>&1 | FileCheck %s
+
+; CHECK: error: Handle at "h2" overwrites handle at "h1"
+
+%"class.hlsl::RWStructuredBuffer" = type { target("dx.RawBuffer", <4 x float>, 1, 0) }
+@Buf = internal global %"class.hlsl::RWStructuredBuffer" poison, align 4
+
+define float @f() {
+entry:
+  %h1 = call target("dx.RawBuffer", <4 x float>, 1, 0) @llvm.dx.resource.handlefrombinding(i32 0, i32 0, i32 1, i32 0, i1 false)
+  store target("dx.RawBuffer", <4 x float>, 1, 0) %h1, ptr @Buf, align 4
+  %h2 = call target("dx.RawBuffer", <4 x float>, 1, 0) @llvm.dx.resource.handlefrombinding(i32 0, i32 1, i32 1, i32 0, i1 false)
+  store target("dx.RawBuffer", <4 x float>, 1, 0) %h2, ptr @Buf, align 4
+
+  %b = load target("dx.RawBuffer", <4 x float>, 1, 0), ptr @Buf, align 4
+  %l = call { <4 x float>, i1 } @llvm.dx.resource.load.rawbuffer(target("dx.RawBuffer", <4 x float>, 1, 0) %b, i32 0, i32 0)
+  %x = extractvalue { <4 x float>, i1 } %l, 0
+  %v = extractelement <4 x float> %x, i32 0
+
+  ret float %v
+}

llvm/test/CodeGen/DirectX/ForwardHandleAccesses/buffer-O0.ll

@@ -0,0 +1,44 @@
+; RUN: opt -S -dxil-forward-handle-accesses -mtriple=dxil--shadermodel6.3-library %s | FileCheck %s
+
+%"class.hlsl::RWStructuredBuffer" = type { target("dx.RawBuffer", <4 x float>, 1, 0) }
+
+@_ZL2In = internal global %"class.hlsl::RWStructuredBuffer" poison, align 4
+@_ZL3Out = internal global %"class.hlsl::RWStructuredBuffer" poison, align 4
+
+define void @main() #1 {
+entry:
+  %this.addr.i.i.i = alloca ptr, align 4
+  %this.addr.i.i = alloca ptr, align 4
+  %this.addr.i1 = alloca ptr, align 4
+  %Index.addr.i2 = alloca i32, align 4
+  %this.addr.i = alloca ptr, align 4
+  %Index.addr.i = alloca i32, align 4
+  ; CHECK: [[IN:%.*]] = call target("dx.RawBuffer", <4 x float>, 1, 0) @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_v4f32_1_0t(i32 0, i32 0, i32 1, i32 0, i1 false)
+  %_ZL2In_h.i.i = call target("dx.RawBuffer", <4 x float>, 1, 0) @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_v4f32_1_0t(i32 0, i32 0, i32 1, i32 0, i1 false)
+  store target("dx.RawBuffer", <4 x float>, 1, 0) %_ZL2In_h.i.i, ptr @_ZL2In, align 4
+  store ptr @_ZL2In, ptr %this.addr.i.i, align 4
+  %this1.i.i = load ptr, ptr %this.addr.i.i, align 4
+  ; CHECK: [[OUT:%.*]] = call target("dx.RawBuffer", <4 x float>, 1, 0) @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_v4f32_1_0t(i32 0, i32 1, i32 1, i32 0, i1 false)
+  %_ZL3Out_h.i.i = call target("dx.RawBuffer", <4 x float>, 1, 0) @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_v4f32_1_0t(i32 0, i32 1, i32 1, i32 0, i1 false)
+  store target("dx.RawBuffer", <4 x float>, 1, 0) %_ZL3Out_h.i.i, ptr @_ZL3Out, align 4
+  store ptr @_ZL3Out, ptr %this.addr.i.i.i, align 4
+  %this1.i.i.i = load ptr, ptr %this.addr.i.i.i, align 4
+  store ptr @_ZL2In, ptr %this.addr.i1, align 4
+  store i32 0, ptr %Index.addr.i2, align 4
+  %this1.i3 = load ptr, ptr %this.addr.i1, align 4
+  ; CHECK-NOT: load target("dx.RawBuffer", <4 x float>, 1, 0)
+  %0 = load target("dx.RawBuffer", <4 x float>, 1, 0), ptr %this1.i3, align 4
+  %1 = load i32, ptr %Index.addr.i2, align 4
+  ; CHECK: call { <4 x float>, i1 } @llvm.dx.resource.load.rawbuffer.v4f32.tdx.RawBuffer_v4f32_1_0t(target("dx.RawBuffer", <4 x float>, 1, 0) [[IN]],
+  %2 = call { <4 x float>, i1 } @llvm.dx.resource.load.rawbuffer.v4f32.tdx.RawBuffer_v4f32_1_0t(target("dx.RawBuffer", <4 x float>, 1, 0) %0, i32 %1, i32 0)
+  %3 = extractvalue { <4 x float>, i1 } %2, 0
+  store ptr @_ZL3Out, ptr %this.addr.i, align 4
+  store i32 0, ptr %Index.addr.i, align 4
+  %this1.i = load ptr, ptr %this.addr.i, align 4
+  ; CHECK-NOT: load target("dx.RawBuffer", <4 x float>, 1, 0)
+  %4 = load target("dx.RawBuffer", <4 x float>, 1, 0), ptr %this1.i, align 4
+  %5 = load i32, ptr %Index.addr.i, align 4
+  ; CHECK: call void @llvm.dx.resource.store.rawbuffer.tdx.RawBuffer_v4f32_1_0t.v4f32(target("dx.RawBuffer", <4 x float>, 1, 0) [[OUT]],
+  call void @llvm.dx.resource.store.rawbuffer.tdx.RawBuffer_v4f32_1_0t.v4f32(target("dx.RawBuffer", <4 x float>, 1, 0) %4, i32 %5, i32 0, <4 x float> %3)
+  ret void
+}

llvm/test/CodeGen/DirectX/ForwardHandleAccesses/cbuffer-access.ll

@@ -0,0 +1,23 @@
+; RUN: opt -S -dxil-forward-handle-accesses -mtriple=dxil--shadermodel6.3-library %s | FileCheck %s
+
+%__cblayout_CB = type <{ float, i32, i32 }>
+%struct.Scalars = type { float, i32, i32 }
+
+@CB.cb = local_unnamed_addr global target("dx.CBuffer", target("dx.Layout", %__cblayout_CB, 12, 0, 4, 8)) poison
+
+define void @main() local_unnamed_addr #1 {
+entry:
+  ; CHECK: [[CB:%.*]] = tail call target({{.*}}) @llvm.dx.resource.handlefrombinding
+  %h = tail call target("dx.CBuffer", target("dx.Layout", %__cblayout_CB, 12, 0, 4, 8)) @llvm.dx.resource.handlefrombinding(i32 0, i32 0, i32 1, i32 0, i1 false)
+  store target("dx.CBuffer", target("dx.Layout", %__cblayout_CB, 12, 0, 4, 8)) %h, ptr @CB.cb, align 4
+  %_ZL3Out_h.i.i = tail call target("dx.RawBuffer", %struct.Scalars, 1, 0) @llvm.dx.resource.handlefrombinding(i32 0, i32 0, i32 1, i32 0, i1 false)
+  ; CHECK-NOT: load target({{.*}}), ptr @CB.cb
+  %cb = load target("dx.CBuffer", target("dx.Layout", %__cblayout_CB, 12, 0, 4, 8)), ptr @CB.cb, align 4
+  ; CHECK: call { float, float, float, float } @llvm.dx.resource.load.cbufferrow.4.{{.*}}(target({{.*}}) [[CB]], i32 0)
+  %0 = call { float, float, float, float } @llvm.dx.resource.load.cbufferrow.4(target("dx.CBuffer", target("dx.Layout", %__cblayout_CB, 12, 0, 4, 8)) %cb, i32 0)
+  %1 = extractvalue { float, float, float, float } %0, 0
+  call void @llvm.dx.resource.store.rawbuffer(target("dx.RawBuffer", %struct.Scalars, 1, 0) %_ZL3Out_h.i.i, i32 0, i32 0, float %1)
+  ret void
+}
+
+attributes #0 = { mustprogress nofree noinline norecurse nosync nounwind willreturn memory(readwrite, argmem: write, inaccessiblemem: none) "approx-func-fp-math"="false" "frame-pointer"="all" "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" "no-infs-fp-math"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="true" "stack-protector-buffer-size"="8" }