[NVPTX] Fix crash caused by ComputePTXValueVTs (#104524)

When [lowering return
values](https://github.com/llvm/llvm-project/blob/99a10f1fe8a7e4b0fdb4c6dd5e7f24f87e0d3695/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp#L3422)
from LLVM IR to SelectionDAG, we check that [the number of values
`SelectionDAG` tells us to return is equal to the number of values that
`ComputePTXValueVTs()` tells us to
return](https://github.com/llvm/llvm-project/blob/99a10f1fe8a7e4b0fdb4c6dd5e7f24f87e0d3695/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp#L3441).
However, this check can fail on valid IR. For example:

```
define <6 x half> @foo() {
  ret <6 x half> zeroinitializer
}
```

`ComputePTXValueVTs()` tells us to return ***3*** `v2f16` values, while
`SelectionDAG` tells us to return ***6*** `f16` values. Thus, the
compiler will crash.

`ComputePTXValueVTs()` [supports all `half` element vectors with an even
number of
elements](https://github.com/llvm/llvm-project/blob/99a10f1fe8a7e4b0fdb4c6dd5e7f24f87e0d3695/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp#L213).
Whereas `SelectionDAG` [only supports power-of-2 sized
vectors](https://github.com/llvm/llvm-project/blob/4e078e3797098daa40d254447c499bcf61415308/llvm/lib/CodeGen/TargetLoweringBase.cpp#L1580).
This is the root of the discrepancy.

Assuming that the developers who added the code to
`ComputePTXValueVTs()` overlooked this, I've restricted
`ComputePTXValueVTs()` to compute the same number of return values as
`SelectionDAG`, instead of extending `SelectionDAG` to support
non-power-of-2 sized vectors.

Author: justinfargnoli

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

@@ -207,10 +207,15 @@ static void ComputePTXValueVTs(const TargetLowering &TLI, const DataLayout &DL,
     if (VT.isVector()) {
       unsigned NumElts = VT.getVectorNumElements();
       EVT EltVT = VT.getVectorElementType();
-      // Vectors with an even number of f16 elements will be passed to
-      // us as an array of v2f16/v2bf16 elements. We must match this so we
-      // stay in sync with Ins/Outs.
-      if ((Is16bitsType(EltVT.getSimpleVT())) && NumElts % 2 == 0) {
+      // We require power-of-2 sized vectors becuase
+      // TargetLoweringBase::getVectorTypeBreakdown() which is invoked in
+      // ComputePTXValueVTs() cannot currently break down non-power-of-2 sized
+      // vectors.
+      if ((Is16bitsType(EltVT.getSimpleVT())) && NumElts % 2 == 0 &&
+          isPowerOf2_32(NumElts)) {
+        // Vectors with an even number of f16 elements will be passed to
+        // us as an array of v2f16/v2bf16 elements. We must match this so we
+        // stay in sync with Ins/Outs.
         switch (EltVT.getSimpleVT().SimpleTy) {
         case MVT::f16:
           EltVT = MVT::v2f16;
@@ -226,7 +231,8 @@ static void ComputePTXValueVTs(const TargetLowering &TLI, const DataLayout &DL,
         }
         NumElts /= 2;
       } else if (EltVT.getSimpleVT() == MVT::i8 &&
-                 (NumElts % 4 == 0 || NumElts == 3)) {
+                 ((NumElts % 4 == 0 && isPowerOf2_32(NumElts)) ||
+                  NumElts == 3)) {
         // v*i8 are formally lowered as v4i8
         EltVT = MVT::v4i8;
         NumElts = (NumElts + 3) / 4;

llvm/test/CodeGen/NVPTX/compute-ptx-value-vts.ll

@@ -0,0 +1,61 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc < %s -march=nvptx64 -mcpu=sm_20 | FileCheck %s
+
+target triple = "nvptx-nvidia-cuda"
+
+define <6 x half> @half6() {
+; CHECK-LABEL: half6(
+; CHECK:       {
+; CHECK-NEXT:    .reg .b16 %rs<2>;
+; CHECK-EMPTY:
+; CHECK-NEXT:  // %bb.0:
+; CHECK-NEXT:    mov.b16 %rs1, 0x0000;
+; CHECK-NEXT:    st.param.v4.b16 [func_retval0+0], {%rs1, %rs1, %rs1, %rs1};
+; CHECK-NEXT:    st.param.v2.b16 [func_retval0+8], {%rs1, %rs1};
+; CHECK-NEXT:    ret;
+  ret <6 x half> zeroinitializer
+}
+
+define <10 x half> @half10() {
+; CHECK-LABEL: half10(
+; CHECK:       {
+; CHECK-NEXT:    .reg .b16 %rs<2>;
+; CHECK-EMPTY:
+; CHECK-NEXT:  // %bb.0:
+; CHECK-NEXT:    mov.b16 %rs1, 0x0000;
+; CHECK-NEXT:    st.param.v4.b16 [func_retval0+0], {%rs1, %rs1, %rs1, %rs1};
+; CHECK-NEXT:    st.param.v4.b16 [func_retval0+8], {%rs1, %rs1, %rs1, %rs1};
+; CHECK-NEXT:    st.param.v2.b16 [func_retval0+16], {%rs1, %rs1};
+; CHECK-NEXT:    ret;
+  ret <10 x half> zeroinitializer
+}
+
+define <12 x i8> @byte12() {
+; CHECK-LABEL: byte12(
+; CHECK:       {
+; CHECK-NEXT:    .reg .b16 %rs<2>;
+; CHECK-EMPTY:
+; CHECK-NEXT:  // %bb.0:
+; CHECK-NEXT:    mov.u16 %rs1, 0;
+; CHECK-NEXT:    st.param.v4.b8 [func_retval0+0], {%rs1, %rs1, %rs1, %rs1};
+; CHECK-NEXT:    st.param.v4.b8 [func_retval0+4], {%rs1, %rs1, %rs1, %rs1};
+; CHECK-NEXT:    st.param.v4.b8 [func_retval0+8], {%rs1, %rs1, %rs1, %rs1};
+; CHECK-NEXT:    ret;
+  ret <12 x i8> zeroinitializer
+}
+
+define <20 x i8> @byte20() {
+; CHECK-LABEL: byte20(
+; CHECK:       {
+; CHECK-NEXT:    .reg .b16 %rs<2>;
+; CHECK-EMPTY:
+; CHECK-NEXT:  // %bb.0:
+; CHECK-NEXT:    mov.u16 %rs1, 0;
+; CHECK-NEXT:    st.param.v4.b8 [func_retval0+0], {%rs1, %rs1, %rs1, %rs1};
+; CHECK-NEXT:    st.param.v4.b8 [func_retval0+4], {%rs1, %rs1, %rs1, %rs1};
+; CHECK-NEXT:    st.param.v4.b8 [func_retval0+8], {%rs1, %rs1, %rs1, %rs1};
+; CHECK-NEXT:    st.param.v4.b8 [func_retval0+12], {%rs1, %rs1, %rs1, %rs1};
+; CHECK-NEXT:    st.param.v4.b8 [func_retval0+16], {%rs1, %rs1, %rs1, %rs1};
+; CHECK-NEXT:    ret;
+  ret <20 x i8> zeroinitializer
+}