[llvm] abd927e - [AMDGPU] Check for num elts in SelectVOP3PMods

[Piotr Sobczak via llvm-commits]

Author: Piotr Sobczak
Date: 2022-09-14T20:00:19+02:00
New Revision: abd927e5a849f0dbb2f8a8fe232519424fc84c52

URL: https://github.com/llvm/llvm-project/commit/abd927e5a849f0dbb2f8a8fe232519424fc84c52
DIFF: https://github.com/llvm/llvm-project/commit/abd927e5a849f0dbb2f8a8fe232519424fc84c52.diff

LOG: [AMDGPU] Check for num elts in SelectVOP3PMods

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.

Differential Revision: https://reviews.llvm.org/D133852

Added: 
    llvm/test/CodeGen/AMDGPU/wmma_modifiers.ll

Modified: 
    llvm/lib/Target/AMDGPU/AMDGPUISelDAGToDAG.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Target/AMDGPU/AMDGPUISelDAGToDAG.cpp b/llvm/lib/Target/AMDGPU/AMDGPUISelDAGToDAG.cpp
index 8425948be92b9..2e1d6b27b3ae6 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUISelDAGToDAG.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUISelDAGToDAG.cpp
@@ -2704,7 +2704,7 @@ bool AMDGPUDAGToDAGISel::SelectVOP3PMods(SDValue In, SDValue &Src,
     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;
 

diff  --git a/llvm/test/CodeGen/AMDGPU/wmma_modifiers.ll b/llvm/test/CodeGen/AMDGPU/wmma_modifiers.ll
new file mode 100644
index 0000000000000..c8312c0f039dc
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/wmma_modifiers.ll
@@ -0,0 +1,22 @@
+; 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
+}