[AMDGPU] Check for num elts in SelectVOP3PMods

[piotr]

Summary

The rest of the code section assumes there are exactly two elements
in the vector (Lo, Hi), so add the check before entering the section.

[arsenm]

Does this solve a real problem?

[piotr]

In D133852#3789211, @arsenm wrote:

Does this solve a real problem?

Yes, the wrong path can be entered by one of the v_wmma instructions.

[arsenm]

In D133852#3789325, @piotr wrote:

In D133852#3789211, @arsenm wrote:

Does this solve a real problem?

Yes, the wrong path can be entered by one of the v_wmma instructions.

Should get a test then

[piotr]

Added test.

llvm/lib/Target/AMDGPU/AMDGPUISelDAGToDAG.cpp

bool AMDGPUDAGToDAGISel::SelectVOP3PMods(SDValue In, SDValue &Src,
   unsigned Mods = 0;
   Src = In;
 
   if (Src.getOpcode() == ISD::FNEG) {
     Mods ^= (SISrcMods::NEG | SISrcMods::NEG_HI);
     Src = Src.getOperand(0);
   }
 
-  if (Src.getOpcode() == ISD::BUILD_VECTOR &&
+  if (Src.getOpcode() == ISD::BUILD_VECTOR && Src.getNumOperands() == 2 &&
       (!IsDOT || !Subtarget->hasDOTOpSelHazard())) {
     unsigned VecMods = Mods;
 
     SDValue Lo = stripBitcast(Src.getOperand(0));
     SDValue Hi = stripBitcast(Src.getOperand(1));
 
     if (Lo.getOpcode() == ISD::FNEG) {
       Lo = stripBitcast(Lo.getOperand(0));

llvm/test/CodeGen/AMDGPU/wmma_modifiers.ll

+; RUN: llc -march=amdgcn -mcpu=gfx1100 -mattr=-wavefrontsize32,+wavefrontsize64 -verify-machineinstrs < %s | FileCheck %s
+
+declare <4 x float> @llvm.amdgcn.wmma.f32.16x16x16.f16(<16 x half>, <16 x half>, <4 x float>)
+
+; Make sure we don't crash when trying to select modifiers in SelectVOP3PMods.
+
+define amdgpu_cs void @xyz () {
+; CHECK-LABEL: xyz:
+; CHECK: v_wmma_f32_16x16x16_f16 v[0:3], v[4:11], v[4:11], v[0:3]
+
+.entry:
+  br label %loop
+loop:
+  %ld = load <8 x float>, <8 x float> addrspace(5)* null, align 32
+  %in_shuffle = shufflevector <8 x float> %ld, <8 x float> undef, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+  %wmma = call <4 x float> @llvm.amdgcn.wmma.f32.16x16x16.f16(<16 x half> undef, <16 x half> undef, <4 x float> %in_shuffle)
+  %out_shuffle = shufflevector <4 x float> %wmma, <4 x float> undef, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 undef, i32 undef, i32 undef, i32 undef>
+  store <8 x float> %out_shuffle, <8 x float> addrspace(5)* null, align 32
+  br i1 false, label %.exit, label %loop
+.exit:
+  ret void
+}