Fix #1629 - crash when using NonUniformResourceIndex (#1646)

Fix crash caused by NonUniformSet being invalidated before use
Instead of using NonUniformSet, defer lowering of
NonUniformResourceIndex until last, then mark all GEPs that use
the incoming value.

This will mark all uses of the value used in NonUniformResourceIndex
as non-uniform, including indexing that was not marked non-uniform,
if it ends up being the same index.  This matches the prior behavior,
and avoids loss of non-unifom metadata when merging GEPs, but could
be considered not quite correct, although the difference in indexing
would probably not have been intentional.

include/dxc/HLSL/HLOperationLower.h

@@ -14,7 +14,6 @@
 
 namespace llvm {
 class Instruction;
-class Value;
 class LoadInst;
 class Function;
 }
@@ -26,6 +25,5 @@ class HLSLExtensionsCodegenHelper;
 
 void TranslateBuiltinOperations(
     HLModule &HLM, HLSLExtensionsCodegenHelper *extCodegenHelper,
-    std::unordered_set<llvm::LoadInst *> &UpdateCounterSet,
-    std::unordered_set<llvm::Value *> &NonUniformSet);
+    std::unordered_set<llvm::LoadInst *> &UpdateCounterSet);
 }

lib/HLSL/DxilGenerationPass.cpp

@@ -280,14 +280,12 @@ public:
     }
 
     std::unordered_set<LoadInst *> UpdateCounterSet;
-    std::unordered_set<Value *> NonUniformSet;
 
-    GenerateDxilOperations(M, UpdateCounterSet, NonUniformSet);
+    GenerateDxilOperations(M, UpdateCounterSet);
 
-    GenerateDxilCBufferHandles(NonUniformSet);
+    GenerateDxilCBufferHandles();
     MarkUpdateCounter(UpdateCounterSet);
     LowerHLCreateHandle();
-    MarkNonUniform(NonUniformSet);
 
     // LowerHLCreateHandle() should have translated HLCreateHandle to CreateHandleForLib.
     // Clean up HLCreateHandle functions.
@@ -333,14 +331,12 @@ private:
   void MarkUpdateCounter(std::unordered_set<LoadInst *> &UpdateCounterSet);
   // Generate DXIL cbuffer handles.
   void
-  GenerateDxilCBufferHandles(std::unordered_set<Value *> &NonUniformSet);
+  GenerateDxilCBufferHandles();
 
   // change built-in funtion into DXIL operations
   void GenerateDxilOperations(Module &M,
-                              std::unordered_set<LoadInst *> &UpdateCounterSet,
-                              std::unordered_set<Value *> &NonUniformSet);
+                              std::unordered_set<LoadInst *> &UpdateCounterSet);
   void LowerHLCreateHandle();
-  void MarkNonUniform(std::unordered_set<Value *> &NonUniformSet);
 
   // Translate precise attribute into HL function call.
   void TranslatePreciseAttribute();
@@ -398,17 +394,6 @@ void DxilGenerationPass::LowerHLCreateHandle() {
   }
 }
 
-void DxilGenerationPass::MarkNonUniform(
-    std::unordered_set<Value *> &NonUniformSet) {
-  for (Value *V : NonUniformSet) {
-    for (User *U : V->users()) {
-      if (GetElementPtrInst *I = dyn_cast<GetElementPtrInst>(U)) {
-        DxilMDHelper::MarkNonUniform(I);
-      }
-    }
-  }
-}
-
 static void
 MarkUavUpdateCounter(Value* LoadOrGEP,
                      DxilResource &res,
@@ -449,8 +434,7 @@ void DxilGenerationPass::MarkUpdateCounter(
   }
 }
 
-void DxilGenerationPass::GenerateDxilCBufferHandles(
-    std::unordered_set<Value *> &NonUniformSet) {
+void DxilGenerationPass::GenerateDxilCBufferHandles() {
   // For CBuffer, handle are mapped to HLCreateHandle.
   OP *hlslOP = m_pHLModule->GetOP();
   Value *opArg = hlslOP->GetU32Const((unsigned)OP::OpCode::CreateHandleForLib);
@@ -512,14 +496,6 @@ void DxilGenerationPass::GenerateDxilCBufferHandles(
         }
         // Add GEP for cbv array use.
         Value *GEP = Builder.CreateGEP(GV, {zeroIdx, CBIndex});
-        /*
-        if (!NonUniformSet.count(CBIndex))
-          args[DXIL::OperandIndex::kCreateHandleIsUniformOpIdx] =
-              hlslOP->GetI1Const(0);
-        else
-          args[DXIL::OperandIndex::kCreateHandleIsUniformOpIdx] =
-              hlslOP->GetI1Const(1);*/
-
         Value *V = Builder.CreateLoad(GEP);
         CallInst *handle = Builder.CreateCall(createHandle, {opArg, V}, handleName);
         CI->replaceAllUsesWith(handle);
@@ -530,8 +506,7 @@ void DxilGenerationPass::GenerateDxilCBufferHandles(
 }
 
 void DxilGenerationPass::GenerateDxilOperations(
-    Module &M, std::unordered_set<LoadInst *> &UpdateCounterSet,
-    std::unordered_set<Value *> &NonUniformSet) {
+    Module &M, std::unordered_set<LoadInst *> &UpdateCounterSet) {
   // remove all functions except entry function
   Function *entry = m_pHLModule->GetEntryFunction();
   const ShaderModel *pSM = m_pHLModule->GetShaderModel();
@@ -557,7 +532,7 @@ void DxilGenerationPass::GenerateDxilOperations(
   }
 
   TranslateBuiltinOperations(*m_pHLModule, m_extensionsCodegenHelper,
-                             UpdateCounterSet, NonUniformSet);
+                             UpdateCounterSet);
 
   // Remove unused HL Operation functions.
   std::vector<Function *> deadList;

lib/HLSL/HLOperationLower.cpp

@@ -75,7 +75,6 @@ private:
   };
   std::unordered_map<Value *, ResAttribute> HandleMetaMap;
   std::unordered_set<LoadInst *> &UpdateCounterSet;
-  std::unordered_set<Value *> &NonUniformSet;
   // Map from pointer of cbuffer to pointer of resource.
   // For cbuffer like this:
   //   cbuffer A {
@@ -87,9 +86,8 @@ private:
 
 public:
   HLObjectOperationLowerHelper(HLModule &HLM,
-                               std::unordered_set<LoadInst *> &UpdateCounter,
-                               std::unordered_set<Value *> &NonUniform)
-      : HLM(HLM), UpdateCounterSet(UpdateCounter), NonUniformSet(NonUniform) {}
+                               std::unordered_set<LoadInst *> &UpdateCounter)
+      : HLM(HLM), UpdateCounterSet(UpdateCounter) {}
   DXIL::ResourceClass GetRC(Value *Handle) {
     ResAttribute &Res = FindCreateHandleResourceBase(Handle);
     return Res.RC;
@@ -110,7 +108,6 @@ public:
     std::unordered_set<Value *> resSet;
     MarkHasCounterOnCreateHandle(handle, resSet);
   }
-  void MarkNonUniform(Value *V) { NonUniformSet.insert(V); }
 
   Value *GetOrCreateResourceForCbPtr(GetElementPtrInst *CbPtr,
                                      GlobalVariable *CbGV, MDNode *MD) {
@@ -536,14 +533,19 @@ Value *TrivialIsSpecialFloat(CallInst *CI, IntrinsicOp IOP, OP::OpCode opcode,
 
 Value *TranslateNonUniformResourceIndex(CallInst *CI, IntrinsicOp IOP, OP::OpCode opcode,
                       HLOperationLowerHelper &helper,  HLObjectOperationLowerHelper *pObjHelper, bool &Translated) {
-  for (User *U : CI->users()) {
-    if (CastInst *I = dyn_cast<CastInst>(U)) {
-      pObjHelper->MarkNonUniform(I);
+  Value *V = CI->getArgOperand(HLOperandIndex::kUnaryOpSrc0Idx);
+  CI->replaceAllUsesWith(V);
+  for (User *U : V->users()) {
+    if (GetElementPtrInst *I = dyn_cast<GetElementPtrInst>(U)) {
+      DxilMDHelper::MarkNonUniform(I);
+    } else if (CastInst *castI = dyn_cast<CastInst>(U)) {
+      for (User *castU : castI->users()) {
+        if (GetElementPtrInst *I = dyn_cast<GetElementPtrInst>(castU)) {
+          DxilMDHelper::MarkNonUniform(I);
+        }
+      }
     }
   }
-  Value *V = CI->getArgOperand(HLOperandIndex::kUnaryOpSrc0Idx);
-  pObjHelper->MarkNonUniform(V);
-  CI->replaceAllUsesWith(V);
   return nullptr;
 }
 
@@ -7089,15 +7091,15 @@ namespace hlsl {
 
 void TranslateBuiltinOperations(
     HLModule &HLM, HLSLExtensionsCodegenHelper *extCodegenHelper,
-    std::unordered_set<LoadInst *> &UpdateCounterSet,
-    std::unordered_set<Value *> &NonUniformSet) {
+    std::unordered_set<LoadInst *> &UpdateCounterSet) {
   HLOperationLowerHelper helper(HLM);
 
-  HLObjectOperationLowerHelper objHelper = {HLM, UpdateCounterSet,
-                                            NonUniformSet};
+  HLObjectOperationLowerHelper objHelper = {HLM, UpdateCounterSet};
 
   Module *M = HLM.GetModule();
 
+  SmallVector<Function *, 4> NonUniformResourceIndexIntrinsics;
+
   // generate dxil operation
   for (iplist<Function>::iterator F : M->getFunctionList()) {
     if (F->user_empty())
@@ -7114,8 +7116,23 @@ void TranslateBuiltinOperations(
       TranslateHLExtension(F, extCodegenHelper, helper.hlslOP);
       continue;
     }
+    if (group == HLOpcodeGroup::HLIntrinsic) {
+      CallInst *CI = cast<CallInst>(*F->user_begin()); // must be call inst
+      unsigned opcode = hlsl::GetHLOpcode(CI);
+      if (opcode == (unsigned)IntrinsicOp::IOP_NonUniformResourceIndex) {
+        NonUniformResourceIndexIntrinsics.push_back(F);
+        continue;
+      }
+    }
     TranslateHLBuiltinOperation(F, helper, group, &objHelper);
   }
+
+  // Translate last so value placed in NonUniformSet is still valid.
+  if (!NonUniformResourceIndexIntrinsics.empty()) {
+    for (auto F : NonUniformResourceIndexIntrinsics) {
+      TranslateHLBuiltinOperation(F, helper, HLOpcodeGroup::HLIntrinsic, &objHelper);
+    }
+  }
 }
 
 }