[SPIRV] Handle aggregate arguments to spv_store

[AlexVlx]

This patch handles the special case where an extract value yields an aggregate result, which then is used as an argument to a store. The SPIRV BE uses special intrinsics (spv_extractv and spv_store) to represent these through IRTranslator, however this creates a problem: spv_store is called as a function, and IRTranslator cannot handle arguments that take more than a vreg. For other functions, the aggregate argument replacement pass would have solved things, but it does not apply here. Hence, we apply the same mutate-into-Int32 solution here when dealing with stores, and restore the extract value's type (which we have available as a ValueAttr) during instruction selection.

[llvmbot]

@llvm/pr-subscribers-backend-spir-v

Author: Alex Voicu (AlexVlx)

Changes

This patch handles the special case where an extract value yields an aggregate result, which then is used as an argument to a store. The SPIRV BE uses special intrinsics (spv_extractv and spv_store) to represent these through IRTranslator, however this creates a problem: spv_store is called as a function, and IRTranslator cannot handle arguments that take more than a vreg. For other functions, the aggregate argument replacement pass would have solved things, but it does not apply here. Hence, we apply the same mutate-into-Int32 solution here when dealing with stores, and restore the extract value's type (which we have available as a ValueAttr) during instruction selection.

---

Full diff: https://github.com/llvm/llvm-project/pull/172348.diff

3 Files Affected:

• (modified) llvm/lib/Target/SPIRV/SPIRVEmitIntrinsics.cpp (+13)
• (modified) llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp (+10)
• (modified) llvm/test/CodeGen/SPIRV/instructions/nested-composites.ll (+30)

diff --git a/llvm/lib/Target/SPIRV/SPIRVEmitIntrinsics.cpp b/llvm/lib/Target/SPIRV/SPIRVEmitIntrinsics.cpp
index fe3ba8b7c114f..b9d26cd6a95f5 100644
--- a/llvm/lib/Target/SPIRV/SPIRVEmitIntrinsics.cpp
+++ b/llvm/lib/Target/SPIRV/SPIRVEmitIntrinsics.cpp
@@ -2032,6 +2032,19 @@ Instruction *SPIRVEmitIntrinsics::visitStoreInst(StoreInst &I) {
   MachineMemOperand::Flags Flags =
       TLI->getStoreMemOperandFlags(I, CurrF->getDataLayout());
   auto *PtrOp = I.getPointerOperand();
+
+  if (I.getValueOperand()->getType()->isAggregateType()) {
+    // It is possible that what used to be an ExtractValueInst has been replaced
+    // with a call to the spv_extractv intrinsic, and that said call hasn't
+    // had its return type replaced with i32 during the dedicated pass (because
+    // it was emitted later); we have to handle this here, because IRTranslator
+    // cannot deal with multi-register types at the moment.
+    CallBase *CB = dyn_cast<CallBase>(I.getValueOperand());
+    assert(CB && CB->getIntrinsicID() == Intrinsic::spv_extractv &&
+           "Unexpected argument of aggregate type, should be spv_extractv!");
+    CB->mutateType(B.getInt32Ty());
+  }
+
   auto *NewI = B.CreateIntrinsic(
       Intrinsic::spv_store, {I.getValueOperand()->getType(), PtrOp->getType()},
       {I.getValueOperand(), PtrOp, B.getInt16(Flags),
diff --git a/llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp b/llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp
index 2e4563795e8f0..dd2e3f1951850 100644
--- a/llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp
+++ b/llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp
@@ -3297,6 +3297,16 @@ bool SPIRVInstructionSelector::selectInsertVal(Register ResVReg,
 bool SPIRVInstructionSelector::selectExtractVal(Register ResVReg,
                                                 const SPIRVType *ResType,
                                                 MachineInstr &I) const {
+  Type *MaybeResTy = nullptr;
+  StringRef ResName;
+  if (GR.findValueAttrs(&I, MaybeResTy, ResName) &&
+      MaybeResTy != GR.getTypeForSPIRVType(ResType)) {
+    assert(!MaybeResTy || MaybeResTy->isAggregateType() &&
+           "Expected aggregate type for extractv instruction");
+    ResType = GR.getOrCreateSPIRVType(MaybeResTy, I,
+                                      SPIRV::AccessQualifier::ReadWrite, false);
+    GR.assignSPIRVTypeToVReg(ResType, ResVReg, *I.getMF());
+  }
   MachineBasicBlock &BB = *I.getParent();
   auto MIB = BuildMI(BB, I, I.getDebugLoc(), TII.get(SPIRV::OpCompositeExtract))
                  .addDef(ResVReg)
diff --git a/llvm/test/CodeGen/SPIRV/instructions/nested-composites.ll b/llvm/test/CodeGen/SPIRV/instructions/nested-composites.ll
index 88e992f183452..7b49002766a26 100644
--- a/llvm/test/CodeGen/SPIRV/instructions/nested-composites.ll
+++ b/llvm/test/CodeGen/SPIRV/instructions/nested-composites.ll
@@ -1,13 +1,19 @@
 ; RUN: llc -verify-machineinstrs -O0 -mtriple=spirv32-unknown-unknown %s -o - | FileCheck %s
+; RUN: llc -verify-machineinstrs -O0 -mtriple=spirv64-unknown-unknown %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv32-unknown-unknown %s -o - -filetype=obj | spirv-val %}
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv64-unknown-unknown %s -o - -filetype=obj | spirv-val %}
 
 ; CHECK-DAG: OpName [[FOOBAR:%.+]] "foobar"
 ; CHECK-DAG: OpName [[PRODUCER:%.+]] "producer"
 ; CHECK-DAG: OpName [[CONSUMER:%.+]] "consumer"
+; CHECK-DAG: OpName [[SNEAKY:%.+]] "sneaky"
+; CHECK-DAG: OpName [[ARR_RET:%.+]] "arr_ret"
 
 ; CHECK-NOT: DAG-FENCE
 
 %ty1 = type {i16, i32}
 %ty2 = type {%ty1, i64}
+%ty3 = type {ptr addrspace(4), ptr addrspace(4), [8 x i8]}
 
 ; CHECK-DAG: [[I16:%.+]] = OpTypeInt 16
 ; CHECK-DAG: [[I32:%.+]] = OpTypeInt 32
@@ -18,6 +24,12 @@
 ; CHECK-DAG: [[UNDEF_I64:%.+]] = OpUndef [[I64]]
 ; CHECK-DAG: [[UNDEF_TY2:%.+]] = OpUndef [[TY2]]
 ; CHECK-DAG: [[CST_42:%.+]] = OpConstant [[I32]] 42
+; CHECK-DAG: [[I8:%.+]] = OpTypeInt 8 0
+; CHECK-DAG: [[PTR:%.+]] = OpTypePointer Generic [[I8]]
+; CHECK-DAG: [[CST_8:%.+]] = OpConstant [[I32]] 8
+; CHECK-DAG: [[ARR:%.+]] = OpTypeArray [[I8]] [[CST_8]]
+; CHECK-DAG: [[TY3:%.+]] = OpTypeStruct [[PTR]] [[PTR]] [[ARR]]
+; CHECK-DAG: [[PTR_ARR:%.+]] = OpTypePointer Generic [[ARR]]
 
 ; CHECK-NOT: DAG-FENCE
 
@@ -62,3 +74,21 @@ define i32 @consumer(%ty2 %agg) {
 ; CHECK: [[RET:%.+]] = OpCompositeExtract [[I32]] [[AGG]] 0 1
 ; CHECK: OpReturnValue [[RET]]
 ; CHECK: OpFunctionEnd
+
+declare %ty3 @arr_ret()
+
+define void @sneaky(ptr addrspace(4) %p, [8 x i8] %a) {
+  %1 = call spir_func %ty3 @arr_ret()
+  %2 = getelementptr inbounds nuw %ty3, ptr addrspace(4) %p, i32 0, i32 2
+  %3 = extractvalue %ty3 %1, 2
+  store [8 x i8] %3, ptr addrspace(4) %2, align 8
+  ret void
+}
+
+; CHECK: [[SNEAKY]] = OpFunction
+; CHECK: [[PTR_AGG:%.+]] = OpFunctionParameter
+; CHECK: [[A:%.+]] = OpFunctionParameter [[ARR]]
+; CHECK: [[AGG_RET:%.+]] = OpFunctionCall [[TY3]]
+; CHECK: [[GEP:%.+]] = OpInBoundsPtrAccessChain [[PTR_ARR]] [[PTR_AGG]]
+; CHECK: [[EXTRACTV:%.+]] = OpCompositeExtract [[ARR]] [[AGG_RET]] 2
+; CHECK: OpStore [[GEP]] [[EXTRACTV]]

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.