[NVPTX] Fixed vectorized LDG for f16.

[tra]

Summary

v2f16 is a special case in NVPTX. v4f16 may be loaded as a pair of v2f16
and that was not previously handled correctly by tryLDGLDU()

[bixia]

llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp

1244

I played with this and found that the only f16 vector that can reach here is v2f16, due to the way that the legalizer is set up for the target. If we can somehow make v4f16 reach here, I believe that we will still see an error in routine GetConvertOpcode, similar to what we see for v2f16 before this fix. For this reason, I prefer an assertion, something like this:

if (EltVT == MVT::v2f16) {

assert(EltVT.getVectorNumElements() == 2 && "missed legalizing vector-of-f16");

} else {

NumElts = EltVT.getVectorNumElements();
 EltVT = EltVT.getVectorElementType();

}

[jlebar]

lgtm modulo Bixia's comment.

llvm/test/CodeGen/NVPTX/ldg-invariant.ll

17

This could use a comma or something somewhere :)

[tra]

llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp

1244

Please take a look at the test cases in ldg-invariants.ll in this patch. v8f16 variant also gets here.

Legalizing node: t6: v8f16,ch = load<(invariant load 16 from %ir.ptr, addrspace 1)> t0, t23, undef:i64
...
Legalizing node: t24: v2f16,v2f16,v2f16,v2f16,ch = NVPTXISD::LoadV4<(invariant load 16 from %ir.ptr, addrspace 1)> t0, t23, undef:i64, Constant:i64<0>

Potentially v4f16 could also get here as a load of a pair of v2f16 elements. If you insist on assertion, it should accept all three cases.

[tra]

Updated test comments.

[bixia]

llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp

1244

What do we do if NumElts % 2 != 0? Can we assert the situation, assuming that the legalizer can guarantee this?

[tra]

Made the check for even number of elements an assertion.
Cosmetic typo fix in the tests.

[tra]

@bixia PTAL.

llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp

bool NVPTXDAGToDAGISel::tryLDGLDU(SDNode *N) {
   SDNode *LD;
   SDValue Base, Offset, Addr;
 
   EVT EltVT = Mem->getMemoryVT();
   unsigned NumElts = 1;
   if (EltVT.isVector()) {
     NumElts = EltVT.getVectorNumElements();
     EltVT = EltVT.getVectorElementType();
+    // vectors of f16 are loaded/stored as multiples of v2f16 elements.
+    if (EltVT == MVT::f16 && N->getValueType(0) == MVT::v2f16 &&
+        NumElts % 2 == 0) {
+      EltVT = MVT::v2f16;
+      NumElts /= 2;
+    }
   }
 
   // Build the "promoted" result VTList for the load. If we are really loading
   // i8s, then the return type will be promoted to i16 since we do not expose
   // 8-bit registers in NVPTX.
   EVT NodeVT = (EltVT == MVT::i8) ? MVT::i16 : EltVT;
   SmallVector<EVT, 5> InstVTs;
   for (unsigned i = 0; i != NumElts; ++i) {

llvm/test/CodeGen/NVPTX/ldg-invariant.ll

 ; RUN: llc < %s -march=nvptx64 -mcpu=sm_35 | FileCheck %s
 
 ; Check that invariant loads from the global addrspace are lowered to
 ; ld.global.nc.
 
 ; CHECK-LABEL: @ld_global
 define i32 @ld_global(i32 addrspace(1)* %ptr) {
 ; CHECK: ld.global.nc.{{[a-z]}}32
   %a = load i32, i32 addrspace(1)* %ptr, !invariant.load !0
   ret i32 %a
 }
 
+; CHECK-LABEL: @ld_global_v2f16
+define half @ld_global_v2f16(<2 x half> addrspace(1)* %ptr) {
+; load of v2f16 is weird as we consider it to be a legal type and don't split it
+; further into individual elements with a side effect of loading and storing it
+; as a 32-bit scalar.
+; CHECK: ld.global.nc.b32
+  %a = load <2 x half>, <2 x half> addrspace(1)* %ptr, !invariant.load !0
+  %v1 = extractelement <2 x half> %a, i32 0
+  %v2 = extractelement <2 x half> %a, i32 1
+  %sum = fadd half %v1, %v2
+  ret half %sum
+}
+
+; CHECK-LABEL: @ld_global_v4f16
+define half @ld_global_v4f16(<4 x half> addrspace(1)* %ptr) {
+; All other variants of f16 vectors get normal vectorized loads.stores.
+; CHECK: ld.global.nc.v4.b16
+  %a = load <4 x half>, <4 x half> addrspace(1)* %ptr, !invariant.load !0
+  %v1 = extractelement <4 x half> %a, i32 0
+  %v2 = extractelement <4 x half> %a, i32 1
+  %v3 = extractelement <4 x half> %a, i32 2
+  %v4 = extractelement <4 x half> %a, i32 3
+  %sum1 = fadd half %v1, %v2
+  %sum2 = fadd half %v3, %v4
+  %sum = fadd half %sum1, %sum2
+  ret half %sum
+}
+
+; CHECK-LABEL: @ld_global_v8f16
+define half @ld_global_v8f16(<8 x half> addrspace(1)* %ptr) {
+; But v8f16 is again loaded as v4i32, because PTX has no v8 variants of loads/stores.
+; CHECK: ld.global.nc.v4.b32
+  %a = load <8 x half>, <8 x half> addrspace(1)* %ptr, !invariant.load !0
+  %v1 = extractelement <8 x half> %a, i32 0
+  %v2 = extractelement <8 x half> %a, i32 2
+  %v3 = extractelement <8 x half> %a, i32 4
+  %v8 = extractelement <8 x half> %a, i32 6
+  %sum1 = fadd half %v1, %v2
+  %sum2 = fadd half %v3, %v8
+  %sum = fadd half %sum1, %sum2
+  ret half %sum
+}
+
+
 ; CHECK-LABEL: @ld_global_v2i32
 define i32 @ld_global_v2i32(<2 x i32> addrspace(1)* %ptr) {
 ; CHECK: ld.global.nc.v2.{{[a-z]}}32
   %a = load <2 x i32>, <2 x i32> addrspace(1)* %ptr, !invariant.load !0
   %v1 = extractelement <2 x i32> %a, i32 0
   %v2 = extractelement <2 x i32> %a, i32 1
   %sum = add i32 %v1, %v2
   ret i32 %sum