[llvm] [SPIR-V] Support `nonuniformindex` intrsinsic in SPIRV CodeGen. (PR #162540)

[Steven Perron via llvm-commits]

================
@@ -3713,6 +3721,58 @@ bool SPIRVInstructionSelector::selectResourceGetPointer(
       .constrainAllUses(TII, TRI, RBI);
 }
 
+bool SPIRVInstructionSelector::selectResourceNonUniformIndex(
+    Register &ResVReg, const SPIRVType *ResType, MachineInstr &I) const {
+  Register ObjReg = I.getOperand(2).getReg();
+  if (!BuildCOPY(ResVReg, ObjReg, I))
+    return false;
+
+  buildOpDecorate(ResVReg, I, TII, SPIRV::Decoration::NonUniformEXT, {});
+  // Check for the registers that use the index marked as non-uniform
+  // and recursively mark them as non-uniform.
+  // Per the spec, it's necessary that the final argument used for
+  // load/store/sample/atomic must be decorated, so we need to propagate the
+  // decoration through access chains and copies.
+  // https://docs.vulkan.org/samples/latest/samples/extensions/descriptor_indexing/README.html#_when_to_use_non_uniform_indexing_qualifier
+  recursivelyDecorateChildAsNonUniform(ResVReg, ResType, I);
+  return true;
+}
+
+void SPIRVInstructionSelector::recursivelyDecorateChildAsNonUniform(
+    Register &NonUniformReg, const SPIRVType *RegType, MachineInstr &I) const {
+  std::vector<std::tuple<Register, const SPIRVType *, MachineInstr *>> WorkList;
+  bool isDecorated = false;
+  for (MachineInstr &Use :
+       RegType->getMF()->getRegInfo().use_instructions(NonUniformReg)) {
+    if (Use.getOpcode() != SPIRV::OpDecorate &&
+        Use.getOpcode() != SPIRV::OpAccessChain &&
+        Use.getOpcode() != SPIRV::OpCopyObject &&
+        Use.getOpcode() != SPIRV::OpLoad)
+      continue;
+
+    if (Use.getOpcode() == SPIRV::OpDecorate &&
+        Use.getOperand(1).getImm() == SPIRV::Decoration::NonUniformEXT) {
+      isDecorated = true;
+      continue;
+    }
+
+    Register ResultReg = Use.getOperand(0).getReg();
+    SPIRVType *ResultType = GR.getResultType(ResultReg);
+    WorkList.push_back(std::make_tuple(ResultReg, ResultType, &Use));
+  }
+
+  if (!isDecorated) {
+    buildOpDecorate(NonUniformReg, I, TII, SPIRV::Decoration::NonUniformEXT,
+                    {});
+  }
+
+  for (auto &Item : WorkList) {
+    recursivelyDecorateChildAsNonUniform(std::get<0>(Item), std::get<1>(Item),
+                                         *std::get<2>(Item));
----------------
s-perron wrote:

I won't happen generally, but we sometimes get corner cases where the recursion goes so deep that the compiler crashes. Can you change this to avoid the recursion? You can add the inputs to the work list at the start. Add a loop runs while the worklist is not empty, take the first element in the list, ...

https://github.com/llvm/llvm-project/pull/162540