AMDGPU: Fix crash on v3i15 kernel arguments

[rampitec]

LGTM

[arsenm]

f734ce0488d45c8073892abb6805c9fd9d99fd46

llvm/lib/Target/AMDGPU/AMDGPUISelLowering.cpp

for (unsigned Value = 0, NumValues = ValueVTs.size();
       if (MemVT.isVector() && MemVT.getVectorNumElements() == 1)
         MemVT = MemVT.getScalarType();
 
       // Round up vec3/vec5 argument.
       if (MemVT.isVector() && !MemVT.isPow2VectorType()) {
         assert(MemVT.getVectorNumElements() == 3 ||
                MemVT.getVectorNumElements() == 5);
         MemVT = MemVT.getPow2VectorType(State.getContext());
+      } else if (!MemVT.isSimple() && !MemVT.isVector()) {
+        MemVT = MemVT.getRoundIntegerType(State.getContext());
       }
 
       unsigned PartOffset = 0;
       for (unsigned i = 0; i != NumRegs; ++i) {
         State.addLoc(CCValAssign::getCustomMem(InIndex++, RegisterVT,
                                                BasePartOffset + PartOffset,
                                                MemVT.getSimpleVT(),
                                                CCValAssign::Full));

llvm/test/CodeGen/AMDGPU/kernel-argument-dag-lowering.ll

 ; GCN-LABEL: {{^}}array_3xi16:
 ; HSA-VI: s_load_dword s{{[0-9]+}}, s[4:5], 0x0
 ; HSA-VI: s_load_dword s{{[0-9]+}}, s[4:5], 0x4
 define amdgpu_kernel void @array_3xi16(i8 %arg0, [3 x i16] %arg1) {
   store volatile i8 %arg0, i8 addrspace(1)* undef
   store volatile [3 x i16] %arg1, [3 x i16] addrspace(1)* undef
   ret void
 }
+
+; GCN-LABEL: {{^}}v2i15_arg:
+; GCN: s_load_dword [[DWORD:s[0-9]+]]
+; GCN-DAG: s_bfe_u32 [[BFE:s[0-9]+]], [[DWORD]], 0x100010{{$}}
+; GCN-DAG: s_and_b32 [[AND:s[0-9]+]], [[DWORD]], 0x7fff{{$}}
+define amdgpu_kernel void @v2i15_arg(<2 x i15> addrspace(1)* nocapture %out, <2 x i15> %in) {
+entry:
+  store <2 x i15> %in, <2 x i15> addrspace(1)* %out, align 4
+  ret void
+}
+
+; GCN-LABEL: {{^}}v3i15_arg:
+; GCN: s_load_dword [[DWORD:s[0-9]+]]
+; GCN: s_lshl_b64
+; GCN: s_and_b32
+; GCN: s_and_b32
+; GCN: s_or_b32
+define amdgpu_kernel void @v3i15_arg(<3 x i15> addrspace(1)* nocapture %out, <3 x i15> %in) {
+entry:
+  store <3 x i15> %in, <3 x i15> addrspace(1)* %out, align 4
+  ret void
+}