[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 ↗

(On Diff #141232)

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 ↗

(On Diff #141232)

This could use a comma or something somewhere :)

[tra]

llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp

1244 ↗

(On Diff #141232)

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 ↗

(On Diff #141240)

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/trunk/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) {
+      assert(NumElts % 2 == 0 && "Vector must have even number of elements");
+      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/trunk/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. We consider it to be a legal type, which happens to be
+; loaded/stored 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) {
+; Larger f16 vectors may be split into individual f16 elements and multiple
+; loads/stores may be vectorized using f16 element type. Practically it's
+; limited to v4 variant only.
+; 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) {
+; Larger vectors are, again, loaded as v4i32. PTX has no v8 variants of loads/stores,
+; so load/store vectorizer has to convert v8f16 -> v4 x v2f16.
+; 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
+  %v4 = extractelement <8 x half> %a, i32 6
+  %sum1 = fadd half %v1, %v2
+  %sum2 = fadd half %v3, %v4
+  %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