[NVPTX] handle more cases for loads and stores

Split unaligned stores and loads of v2f32.

Add DAGCombiner rules for:
 - target-independent stores that store a v2f32 BUILD_VECTOR. We
   scalarize the value and rewrite the store

Fix test cases.

Author: Prince781

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

@@ -865,7 +865,7 @@ NVPTXTargetLowering::NVPTXTargetLowering(const NVPTXTargetMachine &TM,
   setTargetDAGCombine({ISD::ADD, ISD::AND, ISD::EXTRACT_VECTOR_ELT, ISD::FADD,
                        ISD::MUL, ISD::SHL, ISD::SREM, ISD::UREM, ISD::VSELECT,
                        ISD::BUILD_VECTOR, ISD::ADDRSPACECAST, ISD::FP_ROUND,
-                       ISD::TRUNCATE, ISD::LOAD, ISD::BITCAST});
+                       ISD::TRUNCATE, ISD::LOAD, ISD::STORE, ISD::BITCAST});
 
   // setcc for f16x2 and bf16x2 needs special handling to prevent
   // legalizer's attempt to scalarize it due to v2i1 not being legal.
@@ -3242,10 +3242,10 @@ SDValue NVPTXTargetLowering::LowerLOAD(SDValue Op, SelectionDAG &DAG) const {
   if (Op.getValueType() == MVT::i1)
     return LowerLOADi1(Op, DAG);
 
-  // v2f16/v2bf16/v2i16/v4i8 are legal, so we can't rely on legalizer to handle
-  // unaligned loads and have to handle it here.
+  // v2f16/v2bf16/v2i16/v4i8/v2f32 are legal, so we can't rely on legalizer to
+  // handle unaligned loads and have to handle it here.
   EVT VT = Op.getValueType();
-  if (Isv2x16VT(VT) || VT == MVT::v4i8) {
+  if (Isv2x16VT(VT) || VT == MVT::v4i8 || VT == MVT::v2f32) {
     LoadSDNode *Load = cast<LoadSDNode>(Op);
     EVT MemVT = Load->getMemoryVT();
     if (!allowsMemoryAccessForAlignment(*DAG.getContext(), DAG.getDataLayout(),
@@ -3289,22 +3289,23 @@ SDValue NVPTXTargetLowering::LowerSTORE(SDValue Op, SelectionDAG &DAG) const {
   if (VT == MVT::i1)
     return LowerSTOREi1(Op, DAG);
 
-  // v2f16 is legal, so we can't rely on legalizer to handle unaligned
-  // stores and have to handle it here.
-  if ((Isv2x16VT(VT) || VT == MVT::v4i8) &&
+  // v2f16/v2bf16/v2i16/v4i8/v2f32 are legal, so we can't rely on legalizer to
+  // handle unaligned stores and have to handle it here.
+  if ((Isv2x16VT(VT) || VT == MVT::v4i8 || VT == MVT::v2f32) &&
       !allowsMemoryAccessForAlignment(*DAG.getContext(), DAG.getDataLayout(),
                                       VT, *Store->getMemOperand()))
     return expandUnalignedStore(Store, DAG);
 
-  // v2f16, v2bf16 and v2i16 don't need special handling.
-  if (Isv2x16VT(VT) || VT == MVT::v4i8)
+  // v2f16/v2bf16/v2i16/v4i8/v2f32 don't need special handling.
+  if (Isv2x16VT(VT) || VT == MVT::v4i8 || VT == MVT::v2f32)
     return SDValue();
 
   return LowerSTOREVector(Op, DAG);
 }
 
-SDValue
-NVPTXTargetLowering::LowerSTOREVector(SDValue Op, SelectionDAG &DAG) const {
+static SDValue convertVectorStore(SDValue Op, SelectionDAG &DAG,
+                                  const SmallVectorImpl<SDValue> &Elements,
+                                  const NVPTXSubtarget &STI) {
   MemSDNode *N = cast<MemSDNode>(Op.getNode());
   SDValue Val = N->getOperand(1);
   SDLoc DL(N);
@@ -3369,6 +3370,8 @@ NVPTXTargetLowering::LowerSTOREVector(SDValue Op, SelectionDAG &DAG) const {
                                 NumEltsPerSubVector);
       Ops.push_back(DAG.getBuildVector(EltVT, DL, SubVectorElts));
     }
+  } else if (!Elements.empty()) {
+    Ops.insert(Ops.end(), Elements.begin(), Elements.end());
   } else {
     SDValue V = DAG.getBitcast(MVT::getVectorVT(EltVT, NumElts), Val);
     for (const unsigned I : llvm::seq(NumElts)) {
@@ -3392,10 +3395,20 @@ NVPTXTargetLowering::LowerSTOREVector(SDValue Op, SelectionDAG &DAG) const {
       DAG.getMemIntrinsicNode(Opcode, DL, DAG.getVTList(MVT::Other), Ops,
                               N->getMemoryVT(), N->getMemOperand());
 
-  // return DCI.CombineTo(N, NewSt, true);
   return NewSt;
 }
 
+// Default variant where we don't pass in elements.
+static SDValue convertVectorStore(SDValue Op, SelectionDAG &DAG,
+                                  const NVPTXSubtarget &STI) {
+  return convertVectorStore(Op, DAG, SmallVector<SDValue>{}, STI);
+}
+
+SDValue NVPTXTargetLowering::LowerSTOREVector(SDValue Op,
+                                              SelectionDAG &DAG) const {
+  return convertVectorStore(Op, DAG, STI);
+}
+
 // st i1 v, addr
 //    =>
 // v1 = zxt v to i16
@@ -5539,6 +5552,9 @@ static SDValue PerformStoreCombineHelper(SDNode *N,
     // -->
     //   StoreRetvalV2 {a, b}
     // likewise for V2 -> V4 case
+    //
+    // We also handle target-independent stores, which require us to first
+    // convert to StoreV2.
 
     std::optional<NVPTXISD::NodeType> NewOpcode;
     switch (N->getOpcode()) {
@@ -5564,8 +5580,8 @@ static SDValue PerformStoreCombineHelper(SDNode *N,
         SDValue CurrentOp = N->getOperand(I);
         if (CurrentOp->getOpcode() == ISD::BUILD_VECTOR) {
           assert(CurrentOp.getValueType() == MVT::v2f32);
-          NewOps.push_back(CurrentOp.getNode()->getOperand(0));
-          NewOps.push_back(CurrentOp.getNode()->getOperand(1));
+          NewOps.push_back(CurrentOp.getOperand(0));
+          NewOps.push_back(CurrentOp.getOperand(1));
         } else {
           NewOps.clear();
           break;
@@ -6342,6 +6358,19 @@ static SDValue PerformBITCASTCombine(SDNode *N,
   return SDValue();
 }
 
+static SDValue PerformStoreCombine(SDNode *N,
+                                   TargetLowering::DAGCombinerInfo &DCI,
+                                   const NVPTXSubtarget &STI) {
+  // check if the store'd value can be scalarized
+  SDValue StoredVal = N->getOperand(1);
+  if (StoredVal.getValueType() == MVT::v2f32 &&
+      StoredVal.getOpcode() == ISD::BUILD_VECTOR) {
+    SmallVector<SDValue> Elements(StoredVal->op_values());
+    return convertVectorStore(SDValue(N, 0), DCI.DAG, Elements, STI);
+  }
+  return SDValue();
+}
+
 SDValue NVPTXTargetLowering::PerformDAGCombine(SDNode *N,
                                                DAGCombinerInfo &DCI) const {
   CodeGenOptLevel OptLevel = getTargetMachine().getOptLevel();
@@ -6371,6 +6400,8 @@ SDValue NVPTXTargetLowering::PerformDAGCombine(SDNode *N,
     case NVPTXISD::LoadParam:
     case NVPTXISD::LoadParamV2:
       return PerformLoadCombine(N, DCI, STI);
+    case ISD::STORE:
+      return PerformStoreCombine(N, DCI, STI);
     case NVPTXISD::StoreParam:
     case NVPTXISD::StoreParamV2:
     case NVPTXISD::StoreParamV4:

llvm/test/CodeGen/NVPTX/aggregate-return.ll

@@ -10,9 +10,9 @@ define void @test_v2f32(<2 x float> %input, ptr %output) {
 ; CHECK-LABEL: @test_v2f32
   %call = tail call <2 x float> @barv(<2 x float> %input)
 ; CHECK: .param .align 8 .b8 retval0[8];
-; CHECK: ld.param.v2.b32 {[[E0:%r[0-9]+]], [[E1:%r[0-9]+]]}, [retval0];
+; CHECK: ld.param.b64 [[E0_1:%rd[0-9]+]], [retval0];
   store <2 x float> %call, ptr %output, align 8
-; CHECK: st.v2.b32 [{{%rd[0-9]+}}], {[[E0]], [[E1]]}
+; CHECK: st.b64 [{{%rd[0-9]+}}], [[E0_1]]
   ret void
 }
 

llvm/test/CodeGen/NVPTX/f32x2-instructions.ll

@@ -512,14 +512,13 @@ define <2 x float> @test_frem_ftz(<2 x float> %a, <2 x float> %b) #2 {
 define void @test_ldst_v2f32(ptr %a, ptr %b) #0 {
 ; CHECK-LABEL: test_ldst_v2f32(
 ; CHECK:       {
-; CHECK-NEXT:    .reg .b32 %r<3>;
-; CHECK-NEXT:    .reg .b64 %rd<3>;
+; CHECK-NEXT:    .reg .b64 %rd<4>;
 ; CHECK-EMPTY:
 ; CHECK-NEXT:  // %bb.0:
 ; CHECK-NEXT:    ld.param.b64 %rd2, [test_ldst_v2f32_param_1];
 ; CHECK-NEXT:    ld.param.b64 %rd1, [test_ldst_v2f32_param_0];
-; CHECK-NEXT:    ld.v2.b32 {%r1, %r2}, [%rd1];
-; CHECK-NEXT:    st.v2.b32 [%rd2], {%r1, %r2};
+; CHECK-NEXT:    ld.b64 %rd3, [%rd1];
+; CHECK-NEXT:    st.b64 [%rd2], %rd3;
 ; CHECK-NEXT:    ret;
   %t1 = load <2 x float>, ptr %a
   store <2 x float> %t1, ptr %b, align 32

llvm/test/CodeGen/NVPTX/load-with-non-coherent-cache.ll

@@ -108,9 +108,9 @@ define ptx_kernel void @foo10(ptr noalias readonly %from, ptr %to) {
 }
 
 ; SM20-LABEL: .visible .entry foo11(
-; SM20: ld.global.v2.b32
+; SM20: ld.global.b64
 ; SM35-LABEL: .visible .entry foo11(
-; SM35: ld.global.nc.v2.b32
+; SM35: ld.global.nc.b64
 define ptx_kernel void @foo11(ptr noalias readonly %from, ptr %to) {
   %1 = load <2 x float>, ptr %from
   store <2 x float> %1, ptr %to

llvm/test/CodeGen/NVPTX/misaligned-vector-ldst.ll

@@ -26,7 +26,7 @@ define <4 x float> @t2(ptr %p1) {
 ; CHECK-LABEL: t3
 define <4 x float> @t3(ptr %p1) {
 ; CHECK-NOT: ld.v4
-; CHECK: ld.v2
+; CHECK: ld.b64
   %r = load <4 x float>, ptr %p1, align 8
   ret <4 x float> %r
 }