Support NonUniformResourceIndex for dynamic resource. (#3333)

lib/HLSL/DxilCondenseResources.cpp

@@ -658,8 +658,10 @@ public:
     FailOnPoisonResources();
 
     bool bChanged = false;
-    if (DM.GetShaderModel()->IsSM66Plus())
+    if (DM.GetShaderModel()->IsSM66Plus()) {
       bChanged = PatchDynamicTBuffers(DM);
+      SetNonUniformIndexForDynamicResource(DM);
+    }
 
     unsigned numResources = DM.GetCBuffers().size() + DM.GetUAVs().size() +
                             DM.GetSRVs().size() + DM.GetSamplers().size();
@@ -751,6 +753,7 @@ private:
   bool PatchTBuffers(DxilModule &DM);
   void PatchTBufferUse(Value *V, DxilModule &DM, DenseSet<Value *> &patchedSet);
   void UpdateCBufferUsage();
+  void SetNonUniformIndexForDynamicResource(DxilModule &DM);
 };
 
 } // namespace
@@ -2717,6 +2720,26 @@ void DxilLowerCreateHandleForLib::UpdateCBufferUsage() {
  }
 }
 
+void DxilLowerCreateHandleForLib::SetNonUniformIndexForDynamicResource(
+    DxilModule &DM) {
+  hlsl::OP *hlslOP = DM.GetOP();
+  Value *TrueVal = hlslOP->GetI1Const(true);
+  for (auto it : hlslOP->GetOpFuncList(DXIL::OpCode::CreateHandleFromHeap)) {
+    Function *F = it.second;
+    if (!F)
+      continue;
+    for (User *U : F->users()) {
+      CallInst *CI = cast<CallInst>(U);
+      if (!DxilMDHelper::IsMarkedNonUniform(CI))
+        continue;
+      // Set NonUniform to be true.
+      CI->setOperand(DxilInst_CreateHandleFromHeap::arg_nonUniformIndex,
+                     TrueVal);
+      // Clear nonUniform metadata.
+      CI->setMetadata(DxilMDHelper::kDxilNonUniformAttributeMDName, nullptr);
+    }
+  }
+}
 
 char DxilLowerCreateHandleForLib::ID = 0;
 

lib/HLSL/HLOperationLower.cpp

@@ -533,7 +533,7 @@ Value *TranslateNonUniformResourceIndex(CallInst *CI, IntrinsicOp IOP, OP::OpCod
                       HLOperationLowerHelper &helper,  HLObjectOperationLowerHelper *pObjHelper, bool &Translated) {
   Value *V = CI->getArgOperand(HLOperandIndex::kUnaryOpSrc0Idx);
   CI->replaceAllUsesWith(V);
-
+  Type *hdlTy = helper.hlslOP.GetHandleType();
   for (User *U : V->users()) {
     if (GetElementPtrInst *I = dyn_cast<GetElementPtrInst>(U)) {
       // Only mark on GEP which point to resource.
@@ -548,7 +548,7 @@ Value *TranslateNonUniformResourceIndex(CallInst *CI, IntrinsicOp IOP, OP::OpCod
         }
       }
     } else if (CallInst *CI = dyn_cast<CallInst>(U)) {
-      if (hlsl::GetHLOpcodeGroup(CI->getCalledFunction()) == hlsl::HLOpcodeGroup::HLCreateHandle)
+      if (CI->getType() == hdlTy)
         DxilMDHelper::MarkNonUniform(CI);
     }
   }

tools/clang/test/HLSLFileCheck/hlsl/intrinsics/createHandleFromHeap/NonUniformDynamic.hlsl

@@ -0,0 +1,21 @@
+// RUN: %dxc -T cs_6_6 %s | %FileCheck %s
+
+// Make sure nonUniformIndex is true.
+// CHECK:call %dx.types.Handle @dx.op.createHandleFromHeap(i32 218, i32 %{{[0-9]+}}, i1 false, i1 true)
+// CHECK:call %dx.types.Handle @dx.op.createHandleFromHeap(i32 218, i32 %{{[0-9]+}}, i1 false, i1 true)
+
+float read(uint ID) {
+  Buffer<float> buf = ResourceDescriptorHeap[NonUniformResourceIndex(ID)];
+  return buf[0];
+}
+
+void write(uint ID, float f) {
+  RWBuffer<float> buf = ResourceDescriptorHeap[NonUniformResourceIndex(ID)];
+  buf[0] = f;
+}
+
+[numthreads(8, 8, 1)]
+void main( uint2 ID : SV_DispatchThreadID) {
+  float v = read(ID.x);
+  write(ID.y, v);
+}