Fix bug when lower NonUniformResourceIndex. (#3417)

Should only mark nonUniform on user of CI.

lib/DXIL/DxilMetadataHelper.cpp

@@ -2405,17 +2405,22 @@ bool DxilMDHelper::IsKnownMetadataID(LLVMContext &Ctx, unsigned ID)
 
 void DxilMDHelper::GetKnownMetadataIDs(LLVMContext &Ctx, SmallVectorImpl<unsigned> *pIDs)
 {
-    auto AddIdIfExists = [&Ctx, &pIDs](StringRef Name) {
-        unsigned ID = 0;
-        if (Ctx.findMDKindID(hlsl::DxilMDHelper::kDxilPreciseAttributeMDName, &ID))
-        {
-            pIDs->push_back(ID);
-        }
-    };
+  auto AddIdIfExists = [&Ctx, &pIDs](StringRef Name) {
+    unsigned ID = 0;
+    if (Ctx.findMDKindID(hlsl::DxilMDHelper::kDxilPreciseAttributeMDName,
+                         &ID)) {
+      pIDs->push_back(ID);
+    }
+    if (Ctx.findMDKindID(hlsl::DxilMDHelper::kDxilNonUniformAttributeMDName,
+                         &ID)) {
+      pIDs->push_back(ID);
+    }
+  };
 
-    AddIdIfExists(hlsl::DxilMDHelper::kDxilPreciseAttributeMDName);
-    AddIdIfExists(hlsl::DxilMDHelper::kDxilNonUniformAttributeMDName);
+  AddIdIfExists(hlsl::DxilMDHelper::kDxilPreciseAttributeMDName);
+  AddIdIfExists(hlsl::DxilMDHelper::kDxilNonUniformAttributeMDName);
 }
+
 void DxilMDHelper::combineDxilMetadata(llvm::Instruction *K,
                                        const llvm::Instruction *J) {
   if (IsMarkedNonUniform(J))

lib/HLSL/HLOperationLower.cpp

@@ -532,9 +532,8 @@ bool IsResourceGEP(GetElementPtrInst *I) {
 Value *TranslateNonUniformResourceIndex(CallInst *CI, IntrinsicOp IOP, OP::OpCode opcode,
                       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()) {
+  for (User *U : CI->users()) {
     if (GetElementPtrInst *I = dyn_cast<GetElementPtrInst>(U)) {
       // Only mark on GEP which point to resource.
       if (IsResourceGEP(I))
@@ -552,6 +551,7 @@ Value *TranslateNonUniformResourceIndex(CallInst *CI, IntrinsicOp IOP, OP::OpCod
         DxilMDHelper::MarkNonUniform(CI);
     }
   }
+  CI->replaceAllUsesWith(V);
   return nullptr;
 }
 

tools/clang/test/HLSLFileCheck/hlsl/intrinsics/NonUniformResourceIndex/nonUniformIndex_propagation.hlsl

@@ -0,0 +1,20 @@
+// RUN: %dxc -T ps_6_0  %s | FileCheck %s
+
+// Make sure only set nonUniformIndex when mark it.
+// CHECK:call %dx.types.Handle @dx.op.createHandle(i32 57, i8 0, i32 0, i32 %{{.*}}, i1 false)
+// CHECK:call %dx.types.Handle @dx.op.createHandle(i32 57, i8 0, i32 1, i32 %{{.*}}, i1 true)
+
+Buffer<float> Tuniform[16];
+Buffer<float> Tnonuniform[16];
+
+float main(int i : IN) : SV_Target
+{
+    if (i >= 8)
+    {
+        return Tuniform[i].Load(0);
+    }
+    else
+    {
+        return Tnonuniform[NonUniformResourceIndex(i)].Load(0);
+    }
+}