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

[Lucie Choi 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));
----------------
luciechoi wrote:

SG! modified

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