[DirectX] Make dx.RawBuffer an op that can't be replaced (llvm#154620)

fixes llvm#152348

SimplifyCFG collapses raw buffer store from a if\else load into a
select.

This change prevents the TargetExtType dx.Rawbuffer from being replace
thus preserving the if\else blocks.

A further change was needed to eliminate the phi node before we process
Intrinsic::dx_resource_getpointer in DXILResourceAccess.cpp

Author: farzonl

clang/test/CodeGenHLSL/builtins/hlsl_resource_t.hlsl

@@ -2,23 +2,28 @@
 
 using handle_float_t = __hlsl_resource_t [[hlsl::resource_class(UAV)]] [[hlsl::contained_type(float)]];
 
+struct CustomResource {
+  handle_float_t h;
+};
+
+// CHECK: %struct.CustomResource = type { target("dx.TypedBuffer", float, 1, 0, 0) }
 // CHECK: %"class.hlsl::RWBuffer" = type { target("dx.TypedBuffer", <4 x float>, 1, 0, 0)
 // CHECK: %"class.hlsl::StructuredBuffer" = type { target("dx.RawBuffer", %struct.MyStruct, 0, 0)
 // CHECK: %struct.MyStruct = type <{ <4 x float>, <2 x i32> }>
 
-// CHECK: define hidden void @_Z2faU9_Res_u_CTfu17__hlsl_resource_t(target("dx.TypedBuffer", float, 1, 0, 0) %a)
-// CHECK: call void @_Z4foo1U9_Res_u_CTfu17__hlsl_resource_t(target("dx.TypedBuffer", float, 1, 0, 0) %0)
-// CHECK: declare hidden void @_Z4foo1U9_Res_u_CTfu17__hlsl_resource_t(target("dx.TypedBuffer", float, 1, 0, 0))
+// CHECK: define hidden void @_Z2fa14CustomResource(ptr noundef byval(%struct.CustomResource) align 1 %a)
+// CHECK: call void @_Z4foo114CustomResource(ptr noundef byval(%struct.CustomResource) align 1 %agg.tmp)
+// CHECK: declare hidden void @_Z4foo114CustomResource(ptr noundef byval(%struct.CustomResource) align 1)
 
-void foo1(handle_float_t res);
+void foo1(CustomResource res);
 
-void fa(handle_float_t a) {
+void fa(CustomResource a) {
     foo1(a);
 }
 
-// CHECK: define hidden void @_Z2fbU9_Res_u_CTfu17__hlsl_resource_t(target("dx.TypedBuffer", float, 1, 0, 0) %a)
-void fb(handle_float_t a) {
-    handle_float_t b = a;
+// CHECK: define hidden void @_Z2fb14CustomResource(ptr noundef byval(%struct.CustomResource) align 1 %a)
+void fb(CustomResource a) {
+    CustomResource b = a;
 }
 
 // CHECK: define hidden void @_Z2fcN4hlsl8RWBufferIDv4_fEE(ptr noundef byval(%"class.hlsl::RWBuffer") align 4 %a)

llvm/include/llvm/IR/DerivedTypes.h

@@ -845,8 +845,11 @@ class TargetExtType : public Type {
     /// This type may be allocated on the stack, either as the allocated type
     /// of an alloca instruction or as a byval function parameter.
     CanBeLocal = 1U << 2,
-    // This type may be used as an element in a vector.
+    /// This type may be used as an element in a vector.
     CanBeVectorElement = 1U << 3,
+    // This type can only be used in intrinsic arguments and return values.
+    /// In particular, it cannot be used in select and phi instructions.
+    IsTokenLike = 1U << 4,
   };
 
   /// Returns true if the target extension type contains the given property.

llvm/include/llvm/IR/Type.h

@@ -233,6 +233,9 @@ class Type {
   /// Return true if this is 'token'.
   bool isTokenTy() const { return getTypeID() == TokenTyID; }
 
+  /// Returns true if this is 'token' or a token-like target type.s
+  bool isTokenLikeTy() const;
+
   /// True if this is an instance of IntegerType.
   bool isIntegerTy() const { return getTypeID() == IntegerTyID; }
 

llvm/lib/IR/Type.cpp

@@ -1036,7 +1036,8 @@ static TargetTypeInfo getTargetTypeInfo(const TargetExtType *Ty) {
   // DirectX resources
   if (Name.starts_with("dx."))
     return TargetTypeInfo(PointerType::get(C, 0), TargetExtType::CanBeGlobal,
-                          TargetExtType::CanBeLocal);
+                          TargetExtType::CanBeLocal,
+                          TargetExtType::IsTokenLike);
 
   // Opaque types in the AMDGPU name space.
   if (Name == "amdgcn.named.barrier") {
@@ -1054,6 +1055,14 @@ static TargetTypeInfo getTargetTypeInfo(const TargetExtType *Ty) {
   return TargetTypeInfo(Type::getVoidTy(C));
 }
 
+bool Type::isTokenLikeTy() const {
+  if (isTokenTy())
+    return true;
+  if (auto *TT = dyn_cast<TargetExtType>(this))
+    return TT->hasProperty(TargetExtType::Property::IsTokenLike);
+  return false;
+}
+
 Type *TargetExtType::getLayoutType() const {
   return getTargetTypeInfo(this).LayoutType;
 }

llvm/lib/IR/Verifier.cpp

@@ -3008,7 +3008,7 @@ void Verifier::visitFunction(const Function &F) {
     if (!IsIntrinsic) {
       Check(!Arg.getType()->isMetadataTy(),
             "Function takes metadata but isn't an intrinsic", &Arg, &F);
-      Check(!Arg.getType()->isTokenTy(),
+      Check(!Arg.getType()->isTokenLikeTy(),
             "Function takes token but isn't an intrinsic", &Arg, &F);
       Check(!Arg.getType()->isX86_AMXTy(),
             "Function takes x86_amx but isn't an intrinsic", &Arg, &F);
@@ -3022,7 +3022,7 @@ void Verifier::visitFunction(const Function &F) {
   }
 
   if (!IsIntrinsic) {
-    Check(!F.getReturnType()->isTokenTy(),
+    Check(!F.getReturnType()->isTokenLikeTy(),
           "Function returns a token but isn't an intrinsic", &F);
     Check(!F.getReturnType()->isX86_AMXTy(),
           "Function returns a x86_amx but isn't an intrinsic", &F);
@@ -3636,7 +3636,7 @@ void Verifier::visitPHINode(PHINode &PN) {
         "PHI nodes not grouped at top of basic block!", &PN, PN.getParent());
 
   // Check that a PHI doesn't yield a Token.
-  Check(!PN.getType()->isTokenTy(), "PHI nodes cannot have token type!");
+  Check(!PN.getType()->isTokenLikeTy(), "PHI nodes cannot have token type!");
 
   // Check that all of the values of the PHI node have the same type as the
   // result.
@@ -3841,14 +3841,14 @@ void Verifier::visitCallBase(CallBase &Call) {
     for (Type *ParamTy : FTy->params()) {
       Check(!ParamTy->isMetadataTy(),
             "Function has metadata parameter but isn't an intrinsic", Call);
-      Check(!ParamTy->isTokenTy(),
+      Check(!ParamTy->isTokenLikeTy(),
             "Function has token parameter but isn't an intrinsic", Call);
     }
   }
 
   // Verify that indirect calls don't return tokens.
   if (!Call.getCalledFunction()) {
-    Check(!FTy->getReturnType()->isTokenTy(),
+    Check(!FTy->getReturnType()->isTokenLikeTy(),
           "Return type cannot be token for indirect call!");
     Check(!FTy->getReturnType()->isX86_AMXTy(),
           "Return type cannot be x86_amx for indirect call!");

llvm/lib/Target/DirectX/DXILResourceAccess.cpp

@@ -8,14 +8,19 @@
 
 #include "DXILResourceAccess.h"
 #include "DirectX.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/Analysis/DXILResource.h"
+#include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/IntrinsicsDirectX.h"
+#include "llvm/IR/User.h"
 #include "llvm/InitializePasses.h"
+#include "llvm/Transforms/Utils/ValueMapper.h"
 
 #define DEBUG_TYPE "dxil-resource-access"
 
@@ -198,6 +203,112 @@ static void createLoadIntrinsic(IntrinsicInst *II, LoadInst *LI, Value *Offset,
   llvm_unreachable("Unhandled case in switch");
 }
 
+static SmallVector<Instruction *> collectBlockUseDef(Instruction *Start) {
+  SmallPtrSet<Instruction *, 32> Visited;
+  SmallVector<Instruction *, 32> Worklist;
+  SmallVector<Instruction *> Out;
+  auto *BB = Start->getParent();
+
+  // Seed with direct users in this block.
+  for (User *U : Start->users()) {
+    if (auto *I = dyn_cast<Instruction>(U)) {
+      if (I->getParent() == BB)
+        Worklist.push_back(I);
+    }
+  }
+
+  // BFS over transitive users, constrained to the same block.
+  while (!Worklist.empty()) {
+    Instruction *I = Worklist.pop_back_val();
+    if (!Visited.insert(I).second)
+      continue;
+    Out.push_back(I);
+
+    for (User *U : I->users()) {
+      if (auto *J = dyn_cast<Instruction>(U)) {
+        if (J->getParent() == BB)
+          Worklist.push_back(J);
+      }
+    }
+    for (Use &V : I->operands()) {
+      if (auto *J = dyn_cast<Instruction>(V)) {
+        if (J->getParent() == BB && V != Start)
+          Worklist.push_back(J);
+      }
+    }
+  }
+
+  // Order results in program order.
+  DenseMap<const Instruction *, unsigned> Ord;
+  unsigned Idx = 0;
+  for (Instruction &I : *BB)
+    Ord[&I] = Idx++;
+
+  llvm::sort(Out, [&](Instruction *A, Instruction *B) {
+    return Ord.lookup(A) < Ord.lookup(B);
+  });
+
+  return Out;
+}
+
+static void phiNodeRemapHelper(PHINode *Phi, BasicBlock *BB,
+                               IRBuilder<> &Builder,
+                               SmallVector<Instruction *> &UsesInBlock) {
+
+  ValueToValueMapTy VMap;
+  Value *Val = Phi->getIncomingValueForBlock(BB);
+  VMap[Phi] = Val;
+  Builder.SetInsertPoint(&BB->back());
+  for (Instruction *I : UsesInBlock) {
+    // don't clone over the Phi just remap them
+    if (auto *PhiNested = dyn_cast<PHINode>(I)) {
+      VMap[PhiNested] = PhiNested->getIncomingValueForBlock(BB);
+      continue;
+    }
+    Instruction *Clone = I->clone();
+    RemapInstruction(Clone, VMap,
+                     RF_NoModuleLevelChanges | RF_IgnoreMissingLocals);
+    Builder.Insert(Clone);
+    VMap[I] = Clone;
+  }
+}
+
+static void phiNodeReplacement(IntrinsicInst *II,
+                               SmallVectorImpl<Instruction *> &PrevBBDeadInsts,
+                               SetVector<BasicBlock *> &DeadBB) {
+  SmallVector<Instruction *> CurrBBDeadInsts;
+  for (User *U : II->users()) {
+    auto *Phi = dyn_cast<PHINode>(U);
+    if (!Phi)
+      continue;
+
+    IRBuilder<> Builder(Phi);
+    SmallVector<Instruction *> UsesInBlock = collectBlockUseDef(Phi);
+    bool HasReturnUse = isa<ReturnInst>(UsesInBlock.back());
+
+    for (unsigned I = 0, E = Phi->getNumIncomingValues(); I < E; I++) {
+      auto *CurrIncomingBB = Phi->getIncomingBlock(I);
+      phiNodeRemapHelper(Phi, CurrIncomingBB, Builder, UsesInBlock);
+      if (HasReturnUse)
+        PrevBBDeadInsts.push_back(&CurrIncomingBB->back());
+    }
+
+    CurrBBDeadInsts.push_back(Phi);
+
+    for (Instruction *I : UsesInBlock) {
+      CurrBBDeadInsts.push_back(I);
+    }
+    if (HasReturnUse) {
+      BasicBlock *PhiBB = Phi->getParent();
+      DeadBB.insert(PhiBB);
+    }
+  }
+  // Traverse the now-dead instructions in RPO and remove them.
+  for (Instruction *Dead : llvm::reverse(CurrBBDeadInsts))
+    Dead->eraseFromParent();
+  CurrBBDeadInsts.clear();
+}
+
 static void replaceAccess(IntrinsicInst *II, dxil::ResourceTypeInfo &RTI) {
   // Process users keeping track of indexing accumulated from GEPs.
   struct AccessAndOffset {
@@ -229,7 +340,6 @@ static void replaceAccess(IntrinsicInst *II, dxil::ResourceTypeInfo &RTI) {
     } else if (auto *LI = dyn_cast<LoadInst>(Current.Access)) {
       createLoadIntrinsic(II, LI, Current.Offset, RTI);
       DeadInsts.push_back(LI);
-
     } else
       llvm_unreachable("Unhandled instruction - pointer escaped?");
   }
@@ -242,13 +352,27 @@ static void replaceAccess(IntrinsicInst *II, dxil::ResourceTypeInfo &RTI) {
 
 static bool transformResourcePointers(Function &F, DXILResourceTypeMap &DRTM) {
   SmallVector<std::pair<IntrinsicInst *, dxil::ResourceTypeInfo>> Resources;
-  for (BasicBlock &BB : F)
+  SetVector<BasicBlock *> DeadBB;
+  SmallVector<Instruction *> PrevBBDeadInsts;
+  for (BasicBlock &BB : make_early_inc_range(F)) {
+    for (Instruction &I : make_early_inc_range(BB))
+      if (auto *II = dyn_cast<IntrinsicInst>(&I))
+        if (II->getIntrinsicID() == Intrinsic::dx_resource_getpointer)
+          phiNodeReplacement(II, PrevBBDeadInsts, DeadBB);
+
     for (Instruction &I : BB)
       if (auto *II = dyn_cast<IntrinsicInst>(&I))
         if (II->getIntrinsicID() == Intrinsic::dx_resource_getpointer) {
           auto *HandleTy = cast<TargetExtType>(II->getArgOperand(0)->getType());
           Resources.emplace_back(II, DRTM[HandleTy]);
         }
+  }
+  for (auto *Dead : PrevBBDeadInsts)
+    Dead->eraseFromParent();
+  PrevBBDeadInsts.clear();
+  for (auto *Dead : DeadBB)
+    Dead->eraseFromParent();
+  DeadBB.clear();
 
   for (auto &[II, RI] : Resources)
     replaceAccess(II, RI);
@@ -279,7 +403,6 @@ class DXILResourceAccessLegacy : public FunctionPass {
   bool runOnFunction(Function &F) override {
     DXILResourceTypeMap &DRTM =
         getAnalysis<DXILResourceTypeWrapperPass>().getResourceTypeMap();
-
     return transformResourcePointers(F, DRTM);
   }
   StringRef getPassName() const override { return "DXIL Resource Access"; }

llvm/lib/Transforms/Utils/Local.cpp

@@ -3839,7 +3839,7 @@ void llvm::maybeMarkSanitizerLibraryCallNoBuiltin(
 bool llvm::canReplaceOperandWithVariable(const Instruction *I, unsigned OpIdx) {
   const auto *Op = I->getOperand(OpIdx);
   // We can't have a PHI with a metadata or token type.
-  if (Op->getType()->isMetadataTy() || Op->getType()->isTokenTy())
+  if (Op->getType()->isMetadataTy() || Op->getType()->isTokenLikeTy())
     return false;
 
   // swifterror pointers can only be used by a load, store, or as a swifterror