[SPIRV] Add OpAccessChain instruction support (#66253)

This commit adds 2 new instructions in the selector:
 - OpAccessChain
 - OpInBoundsAccessChain.

The choice between the two relies on the `inbounds` marker.

Those instruction are not used for OpenCL, to maintain the same
behavior as previously. They are only added when building for logical
SPIR-V, as it doesn't support the pointer equivalent.

Because logical SPIR-V doesn't support pointer cast either, the
assign_ptr_type intrinsic need to be generated so OpAccessChain gets
lowered with the correct pointer type, instead of i8*.

Fixes #66107

---------

Signed-off-by: Nathan Gauër <brioche@google.com>

llvm/lib/Target/SPIRV/SPIRVEmitIntrinsics.cpp

@@ -413,7 +413,7 @@ void SPIRVEmitIntrinsics::insertAssignPtrTypeIntrs(Instruction *I) {
     EltTyConst = Constant::getNullValue(AI->getAllocatedType());
     AddressSpace = AI->getAddressSpace();
   } else if (auto *GEP = dyn_cast<GetElementPtrInst>(I)) {
-    EltTyConst = Constant::getNullValue(GEP->getSourceElementType());
+    EltTyConst = Constant::getNullValue(GEP->getResultElementType());
     AddressSpace = GEP->getPointerAddressSpace();
   } else {
     llvm_unreachable("Unexpected instruction!");

llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp

@@ -1299,18 +1299,28 @@ bool SPIRVInstructionSelector::selectExtractElt(Register ResVReg,
 bool SPIRVInstructionSelector::selectGEP(Register ResVReg,
                                          const SPIRVType *ResType,
                                          MachineInstr &I) const {
-  // In general we should also support OpAccessChain instrs here (i.e. not
-  // PtrAccessChain) but SPIRV-LLVM Translator doesn't emit them at all and so
-  // do we to stay compliant with its test and more importantly consumers.
-  unsigned Opcode = I.getOperand(2).getImm() ? SPIRV::OpInBoundsPtrAccessChain
-                                             : SPIRV::OpPtrAccessChain;
+  const bool isGEPInBounds = I.getOperand(2).getImm();
+
+  // OpAccessChain could be used for OpenCL, but the SPIRV-LLVM Translator only
+  // relies on PtrAccessChain, so we'll try not to deviate. For Vulkan however,
+  // we have to use Op[InBounds]AccessChain.
+  const unsigned Opcode = STI.isVulkanEnv()
+                              ? (isGEPInBounds ? SPIRV::OpInBoundsAccessChain
+                                               : SPIRV::OpAccessChain)
+                              : (isGEPInBounds ? SPIRV::OpInBoundsPtrAccessChain
+                                               : SPIRV::OpPtrAccessChain);
+
   auto Res = BuildMI(*I.getParent(), I, I.getDebugLoc(), TII.get(Opcode))
                  .addDef(ResVReg)
                  .addUse(GR.getSPIRVTypeID(ResType))
                  // Object to get a pointer to.
                  .addUse(I.getOperand(3).getReg());
   // Adding indices.
-  for (unsigned i = 4; i < I.getNumExplicitOperands(); ++i)
+  const unsigned StartingIndex =
+      (Opcode == SPIRV::OpAccessChain || Opcode == SPIRV::OpInBoundsAccessChain)
+          ? 5
+          : 4;
+  for (unsigned i = StartingIndex; i < I.getNumExplicitOperands(); ++i)
     Res.addUse(I.getOperand(i).getReg());
   return Res.constrainAllUses(TII, TRI, RBI);
 }

llvm/test/CodeGen/SPIRV/logical-access-chain.ll

@@ -0,0 +1,20 @@
+; RUN: llc -O0 -mtriple=spirv-unknown-unknown %s -o - | FileCheck %s
+
+; CHECK:      [[uint:%[0-9]+]] = OpTypeInt 32 0
+; CHECK:     [[uint2:%[0-9]+]] = OpTypeVector [[uint]] 2
+; CHECK:    [[uint_1:%[0-9]+]] = OpConstant [[uint]] 1
+; CHECK:  [[ptr_uint:%[0-9]+]] = OpTypePointer Function [[uint]]
+; CHECK: [[ptr_uint2:%[0-9]+]] = OpTypePointer Function [[uint2]]
+
+define void @main() #1 {
+entry:
+  %0 = alloca <2 x i32>, align 4
+; CHECK: [[var:%[0-9]+]] = OpVariable [[ptr_uint2]] Function
+
+  %1 = getelementptr <2 x i32>, ptr %0, i32 0, i32 1
+; CHECK: {{%[0-9]+}} = OpAccessChain [[ptr_uint]] [[var]] [[uint_1]]
+
+  ret void
+}
+
+attributes #1 = { "hlsl.numthreads"="4,8,16" "hlsl.shader"="compute" }

llvm/test/CodeGen/SPIRV/logical-struct-access.ll

@@ -0,0 +1,59 @@
+; RUN: llc -O0 -mtriple=spirv-unknown-unknown %s -o - | FileCheck %s
+
+; CHECK: [[uint:%[0-9]+]] = OpTypeInt 32 0
+
+%A = type {
+  i32,
+  i32
+}
+; CHECK:    [[A:%[0-9]+]] = OpTypeStruct [[uint]] [[uint]]
+
+%B = type {
+  %A,
+  i32,
+  %A
+}
+; CHECK:    [[B:%[0-9]+]] = OpTypeStruct [[A]] [[uint]] [[A]]
+
+; CHECK: [[uint_0:%[0-9]+]] = OpConstant [[uint]] 0
+; CHECK: [[uint_1:%[0-9]+]] = OpConstant [[uint]] 1
+; CHECK: [[uint_2:%[0-9]+]] = OpConstant [[uint]] 2
+
+; CHECK: [[ptr_uint:%[0-9]+]] = OpTypePointer Function [[uint]]
+; CHECK:    [[ptr_A:%[0-9]+]] = OpTypePointer Function [[A]]
+; CHECK:    [[ptr_B:%[0-9]+]] = OpTypePointer Function [[B]]
+
+define void @main() #1 {
+entry:
+  %0 = alloca %B, align 4
+; CHECK: [[tmp:%[0-9]+]] = OpVariable [[ptr_B]] Function
+
+  %1 = getelementptr %B, ptr %0, i32 0, i32 0
+; CHECK: {{%[0-9]+}} = OpAccessChain [[ptr_A]] [[tmp]] [[uint_0]]
+  %2 = getelementptr inbounds %B, ptr %0, i32 0, i32 0
+; CHECK: {{%[0-9]+}} = OpInBoundsAccessChain [[ptr_A]] [[tmp]] [[uint_0]]
+
+  %3 = getelementptr %B, ptr %0, i32 0, i32 1
+; CHECK: {{%[0-9]+}} = OpAccessChain [[ptr_uint]] [[tmp]] [[uint_1]]
+  %4 = getelementptr inbounds %B, ptr %0, i32 0, i32 1
+; CHECK: {{%[0-9]+}} = OpInBoundsAccessChain [[ptr_uint]] [[tmp]] [[uint_1]]
+
+  %5 = getelementptr %B, ptr %0, i32 0, i32 2
+; CHECK: {{%[0-9]+}} = OpAccessChain [[ptr_A]] [[tmp]] [[uint_2]]
+  %6 = getelementptr inbounds %B, ptr %0, i32 0, i32 2
+; CHECK: {{%[0-9]+}} = OpInBoundsAccessChain [[ptr_A]] [[tmp]] [[uint_2]]
+
+  %7 = getelementptr %B, ptr %0, i32 0, i32 2, i32 1
+; CHECK: {{%[0-9]+}} = OpAccessChain [[ptr_uint]] [[tmp]] [[uint_2]] [[uint_1]]
+  %8 = getelementptr inbounds %B, ptr %0, i32 0, i32 2, i32 1
+; CHECK: {{%[0-9]+}} = OpInBoundsAccessChain [[ptr_uint]] [[tmp]] [[uint_2]] [[uint_1]]
+
+  %9 = getelementptr %B, ptr %0, i32 0, i32 2
+  %10 = getelementptr %A, ptr %9, i32 0, i32 1
+; CHECK: [[x:%[0-9]+]] = OpAccessChain [[ptr_A]] [[tmp]] [[uint_2]]
+; CHECK:   {{%[0-9]+}} = OpAccessChain [[ptr_uint]] [[x]] [[uint_1]]
+
+  ret void
+}
+
+attributes #1 = { "hlsl.numthreads"="4,8,16" "hlsl.shader"="compute" }