[NVPTX] Correctly lower extending loads for fp16 vectors.

[tra]

Summary

Fixes https://github.com/llvm/llvm-project/issues/63436

Improves lowering of extending FP vector loads.

[tra]

Reproducer: https://godbolt.org/z/nzsd16MrP

[nirvedhmeshram]

llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp

3605

I would like to understand how this is safe, since for target nodes we will always get IsSigned = false . Could there be cases where we would actually need a signed convert?

[tra]

llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp

3605

Good point.
I'm pretty sure we have not run into any case like these, because in such a case we'd skip extending the types and would end up with the same error we're dealing with now.
Integers appear to follow a somewhat different code path and do not try to do extending load in one instruction.

However, I do not know it for sure, so I'll update the patch to keep behavior for integers unchanged.

[nirvedhmeshram]

llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp

3605

I did confirm same example works if it is i16/i32 instead of f16/f32 . So I am okay either way.

[tra]

Preserve behavior for lowering integer vectors.

[tra]

Fixed an error in lowering extending vector loads.

We were bailing out due to comparing the alignment of *extended* type, instead
of the type we're actually loading and that resulted in unnecessary vector
splitting.

[tra]

Added more tests for vector loads w/ small alignment.

llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.h

bool SelectADDRsi_imp(SDNode *OpNode, SDValue Addr, SDValue &Base,
                         SDValue &Offset, MVT mvt);
   bool SelectADDRsi(SDNode *OpNode, SDValue Addr, SDValue &Base,
                     SDValue &Offset);
   bool SelectADDRsi64(SDNode *OpNode, SDValue Addr, SDValue &Base,
                       SDValue &Offset);
 
   bool ChkMemSDNodeAddressSpace(SDNode *N, unsigned int spN) const;
 
-  static unsigned GetConvertOpcode(MVT DestTy, MVT SrcTy, bool IsSigned);
+  static unsigned GetConvertOpcode(MVT DestTy, MVT SrcTy, LoadSDNode *N);
 };
 } // end namespace llvm
 
 #endif

llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp

bool NVPTXDAGToDAGISel::tryLDGLDU(SDNode *N) {
   // memory type (in this case, i8) and our types will not match (the node needs
   // to return an i32 in this case). Our LDG/LDU nodes do not support the
   // concept of sign-/zero-extension, so emulate it here by adding an explicit
   // CVT instruction. Ptxas should clean up any redundancies here.
 
   EVT OrigType = N->getValueType(0);
   LoadSDNode *LdNode = dyn_cast<LoadSDNode>(N);
 
-  if (OrigType != EltVT && LdNode) {
+  if (OrigType != EltVT &&
+      (LdNode || (OrigType.isFloatingPoint() && EltVT.isFloatingPoint()))) {
     // We have an extending-load. The instruction we selected operates on the
     // smaller type, but the SDNode we are replacing has the larger type. We
     // need to emit a CVT to make the types match.
-    bool IsSigned = LdNode->getExtensionType() == ISD::SEXTLOAD;
-    unsigned CvtOpc = GetConvertOpcode(OrigType.getSimpleVT(),
-                                       EltVT.getSimpleVT(), IsSigned);
+    unsigned CvtOpc =
+        GetConvertOpcode(OrigType.getSimpleVT(), EltVT.getSimpleVT(), LdNode);
 
     // For each output value, apply the manual sign/zero-extension and make sure
     // all users of the load go through that CVT.
     for (unsigned i = 0; i != NumElts; ++i) {
       SDValue Res(LD, i);
       SDValue OrigVal(N, i);
 
       SDNode *CvtNode =
case InlineAsm::Constraint_m: // memory
     break;
   }
   return true;
 }
 
 /// GetConvertOpcode - Returns the CVT_ instruction opcode that implements a
 /// conversion from \p SrcTy to \p DestTy.
 unsigned NVPTXDAGToDAGISel::GetConvertOpcode(MVT DestTy, MVT SrcTy,
-                                             bool IsSigned) {
+                                             LoadSDNode *LdNode) {
+  bool IsSigned = LdNode && LdNode->getExtensionType() == ISD::SEXTLOAD;
   switch (SrcTy.SimpleTy) {
   default:
     llvm_unreachable("Unhandled source type");
   case MVT::i8:
     switch (DestTy.SimpleTy) {
     default:
       llvm_unreachable("Unhandled dest type");
     case MVT::i16:
default:
       llvm_unreachable("Unhandled dest type");
     case MVT::i8:
       return IsSigned ? NVPTX::CVT_s8_s64 : NVPTX::CVT_u8_u64;
     case MVT::i16:
       return IsSigned ? NVPTX::CVT_s16_s64 : NVPTX::CVT_u16_u64;
     case MVT::i32:
       return IsSigned ? NVPTX::CVT_s32_s64 : NVPTX::CVT_u32_u64;
     }
+  case MVT::f16:
+    switch (DestTy.SimpleTy) {
+    default:
+      llvm_unreachable("Unhandled dest type");
+    case MVT::f32:
+      return NVPTX::CVT_f32_f16;
+    case MVT::f64:
+      return NVPTX::CVT_f64_f16;
+    }
   }
 }

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

case MVT::v8f16: // <4 x f16x2>
     // This is a "native" vector type
     break;
   }
 
   LoadSDNode *LD = cast<LoadSDNode>(N);
 
   Align Alignment = LD->getAlign();
   auto &TD = DAG.getDataLayout();
-  Align PrefAlign = TD.getPrefTypeAlign(ResVT.getTypeForEVT(*DAG.getContext()));
+  Align PrefAlign =
+      TD.getPrefTypeAlign(LD->getMemoryVT().getTypeForEVT(*DAG.getContext()));
   if (Alignment < PrefAlign) {
     // This load is not sufficiently aligned, so bail out and let this vector
     // load be scalarized.  Note that we may still be able to emit smaller
     // vector loads.  For example, if we are loading a <4 x float> with an
     // alignment of 8, this check will fail but the legalizer will try again
     // with 2 x <2 x float>, which will succeed with an alignment of 8.
     return;
   }

llvm/test/CodeGen/NVPTX/vector-loads.ll

; CHECK: ld.v2.u8
   %t32 = load i8, ptr %t31, align 1
   %t33 = icmp ult i8 %t28, %t32
   %t34 = select i1 %t33, i8 %t32, i8 %t28
   store i8 %t34, ptr %t31
 ; CHECK: ret
   ret void
 }
 
+; CHECK-LABEL: extv8f16_global_a16(
+define void @extv8f16_global_a16(ptr addrspace(1) noalias readonly align 16 %dst, ptr addrspace(1) noalias readonly align 16 %src) #0 {
+; CHECK: ld.global.v4.b16 {%f
+; CHECK: ld.global.v4.b16 {%f
+  %v = load <8 x half>, ptr addrspace(1) %src, align 16
+  %ext = fpext <8 x half> %v to <8 x float>
+; CHECK: st.global.v4.f32
+; CHECK: st.global.v4.f32
+  store <8 x float> %ext, ptr addrspace(1) %dst, align 16
+  ret void
+}
+
+; CHECK-LABEL: extv8f16_global_a4(
+define void @extv8f16_global_a4(ptr addrspace(1) noalias readonly align 16 %dst, ptr addrspace(1) noalias readonly align 16 %src) #0 {
+; CHECK: ld.global.v2.b16 {%f
+; CHECK: ld.global.v2.b16 {%f
+; CHECK: ld.global.v2.b16 {%f
+; CHECK: ld.global.v2.b16 {%f
+  %v = load <8 x half>, ptr addrspace(1) %src, align 4
+  %ext = fpext <8 x half> %v to <8 x float>
+; CHECK: st.global.v4.f32
+; CHECK: st.global.v4.f32
+  store <8 x float> %ext, ptr addrspace(1) %dst, align 16
+  ret void
+}
+
+
+; CHECK-LABEL: extv8f16_generic_a16(
+define void @extv8f16_generic_a16(ptr noalias readonly align 16 %dst, ptr noalias readonly align 16 %src) #0 {
+; CHECK: ld.v4.b16 {%f
+; CHECK: ld.v4.b16 {%f
+  %v = load <8 x half>, ptr %src, align 16
+  %ext = fpext <8 x half> %v to <8 x float>
+; CHECK: st.v4.f32
+; CHECK: st.v4.f32
+  store <8 x float> %ext, ptr %dst, align 16
+  ret void
+}
+
+; CHECK-LABEL: extv8f16_generic_a4(
+define void @extv8f16_generic_a4(ptr noalias readonly align 16 %dst, ptr noalias readonly align 16 %src) #0 {
+; CHECK: ld.v2.b16 {%f
+; CHECK: ld.v2.b16 {%f
+; CHECK: ld.v2.b16 {%f
+; CHECK: ld.v2.b16 {%f
+  %v = load <8 x half>, ptr %src, align 4
+  %ext = fpext <8 x half> %v to <8 x float>
+; CHECK: st.v4.f32
+; CHECK: st.v4.f32
+  store <8 x float> %ext, ptr %dst, align 16
+  ret void
+}
+
 
 !1 = !{i32 0, i32 64}