[NVPTX] update how loads are optimized and disable on O0

Prince781 · Prince781 · commit 088391c7fd4c · 2025-05-22T03:00:58.000-07:00
diff --git a/llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp b/llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp
@@ -1189,11 +1189,25 @@ bool NVPTXDAGToDAGISel::tryLoadVector(SDNode *N) {
 
   unsigned FromTypeWidth = TotalWidth / getLoadStoreVectorNumElts(N);
 
+  LLVM_DEBUG({
+    dbgs() << "tryLoadVector on " << TLI->getTargetNodeName(N->getOpcode())
+           << ":\n";
+    dbgs() << "  load type: " << MemVT << "\n";
+    dbgs() << "  total load width: " << TotalWidth << " bits\n";
+    dbgs() << "  from type width: " << FromTypeWidth << " bits\n";
+    dbgs() << "  element type: " << EltVT << "\n";
+  });
+
   if (isSubVectorPackedInInteger(EltVT)) {
     assert(ExtensionType == ISD::NON_EXTLOAD);
     FromTypeWidth = EltVT.getSizeInBits();
     EltVT = MVT::getIntegerVT(FromTypeWidth);
     FromType = NVPTX::PTXLdStInstCode::Untyped;
+    LLVM_DEBUG({
+      dbgs() << "  packed integers detected:\n";
+      dbgs() << "    from type width: " << FromTypeWidth << " (new)\n";
+      dbgs() << "    element type: " << EltVT << " (new)\n";
+    });
   }
 
   assert(isPowerOf2_32(FromTypeWidth) && FromTypeWidth >= 8 &&
@@ -1501,9 +1515,23 @@ bool NVPTXDAGToDAGISel::tryStoreVector(SDNode *N) {
   SDValue N2 = N->getOperand(NumElts + 1);
   unsigned ToTypeWidth = TotalWidth / NumElts;
 
+  LLVM_DEBUG({
+    dbgs() << "tryStoreVector on " << TLI->getTargetNodeName(N->getOpcode())
+           << ":\n";
+    dbgs() << "  store type: " << StoreVT << "\n";
+    dbgs() << "  total store width: " << TotalWidth << " bits\n";
+    dbgs() << "  to type width: " << ToTypeWidth << " bits\n";
+    dbgs() << "  element type: " << EltVT << "\n";
+  });
+
   if (isSubVectorPackedInInteger(EltVT)) {
     ToTypeWidth = EltVT.getSizeInBits();
     EltVT = MVT::getIntegerVT(ToTypeWidth);
+    LLVM_DEBUG({
+      dbgs() << "  packed integers detected:\n";
+      dbgs() << "    to type width: " << ToTypeWidth << " (new)\n";
+      dbgs() << "    element type: " << EltVT << " (new)\n";
+    });
   }
 
   assert(isPowerOf2_32(ToTypeWidth) && ToTypeWidth >= 8 && ToTypeWidth <= 128 &&
diff --git a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
@@ -5323,9 +5323,20 @@ convertVectorLoad(SDNode *N, SelectionDAG &DAG, const NVPTXSubtarget &STI,
   return {{NewLD, LoadChain}};
 }
 
+static MachineMemOperand *
+getMachineMemOperandForType(const SelectionDAG &DAG,
+                            const MachineMemOperand *MMO,
+                            const MachinePointerInfo &PointerInfo, MVT VT) {
+  return DAG.getMachineFunction().getMachineMemOperand(MMO, PointerInfo,
+                                                       LLT(VT));
+}
+
 static SDValue PerformLoadCombine(SDNode *N,
                                   TargetLowering::DAGCombinerInfo &DCI,
                                   const NVPTXSubtarget &STI) {
+  if (DCI.DAG.getOptLevel() == CodeGenOptLevel::None)
+    return {};
+
   auto *MemN = cast<MemSDNode>(N);
   // only operate on vectors of f32s / i64s
   if (EVT MemVT = MemN->getMemoryVT();
@@ -5406,9 +5417,13 @@ static SDValue PerformLoadCombine(SDNode *N,
   // Do we have to tweak the opcode for an NVPTXISD::Load* or do we have to
   // rewrite an ISD::LOAD?
   std::optional<NVPTXISD::NodeType> NewOpcode;
+
+  // LoadV's are handled slightly different in ISelDAGToDAG.
+  bool IsLoadV = false;
   switch (N->getOpcode()) {
   case NVPTXISD::LoadV2:
     NewOpcode = NVPTXISD::LoadV4;
+    IsLoadV = true;
     break;
   case NVPTXISD::LoadParam:
     NewOpcode = NVPTXISD::LoadParamV2;
@@ -5449,9 +5464,22 @@ static SDValue PerformLoadCombine(SDNode *N,
       }
     }
 
+    MVT LoadVT = MVT::f32;
+    MachineMemOperand *MMO = MemN->getMemOperand();
+
+    if (IsLoadV) {
+      // Some loads must have an operand type that matches the number of results
+      // and the type of each result. Because we changed a vNi64 to v(N*2)f32 we
+      // have to update it here. Note that LoadParam is not handled the same way
+      // in NVPXISelDAGToDAG so we only do this for LoadV*.
+      LoadVT = MVT::getVectorVT(MVT::f32, NumElts);
+      MMO = getMachineMemOperandForType(DCI.DAG, MMO, MemN->getPointerInfo(),
+                                        LoadVT);
+    }
+
     NewLoad = DCI.DAG.getMemIntrinsicNode(
         *NewOpcode, SDLoc(N), DCI.DAG.getVTList(VTs),
-        SmallVector<SDValue>(N->ops()), MVT::f32, MemN->getMemOperand());
+        SmallVector<SDValue>(N->ops()), LoadVT, MMO);
     NewChain = NewLoad.getValue(*NewChainIdx);
     if (NewGlueIdx)
       NewGlue = NewLoad.getValue(*NewGlueIdx);
@@ -5550,6 +5578,9 @@ static SDValue PerformStoreCombineHelper(SDNode *N,
     // as the previous value will become unused and eliminated later.
     return N->getOperand(0);
 
+  if (DCI.DAG.getOptLevel() == CodeGenOptLevel::None)
+    return {};
+
   auto *MemN = cast<MemSDNode>(N);
   if (MemN->getMemoryVT() == MVT::v2f32) {
     // try to fold, and expand:
@@ -5581,13 +5612,15 @@ static SDValue PerformStoreCombineHelper(SDNode *N,
     if (NewOpcode) {
       // copy chain, offset from existing store
       SmallVector<SDValue> NewOps = {N->getOperand(0), N->getOperand(1)};
+      unsigned NumElts = 0;
       // gather all operands to expand
       for (unsigned I = 2, E = N->getNumOperands(); I < E; ++I) {
         SDValue CurrentOp = N->getOperand(I);
         if (CurrentOp->getOpcode() == ISD::BUILD_VECTOR) {
           assert(CurrentOp.getValueType() == MVT::v2f32);
           NewOps.push_back(CurrentOp.getOperand(0));
           NewOps.push_back(CurrentOp.getOperand(1));
+          NumElts += 2;
         } else {
           NewOps.clear();
           break;
diff --git a/llvm/test/CodeGen/NVPTX/f16x2-instructions.ll b/llvm/test/CodeGen/NVPTX/f16x2-instructions.ll
@@ -614,7 +614,7 @@ define <2 x float> @test_select_cc_f32_f16(<2 x float> %a, <2 x float> %b,
 ; CHECK-F16:       {
 ; CHECK-F16-NEXT:    .reg .pred %p<3>;
 ; CHECK-F16-NEXT:    .reg .b32 %r<9>;
-; CHECK-F16-NEXT:    .reg .b64 %rd<3>;
+; CHECK-F16-NEXT:    .reg .b64 %rd<4>;
 ; CHECK-F16-EMPTY:
 ; CHECK-F16-NEXT:  // %bb.0:
 ; CHECK-F16-NEXT:    ld.param.b32 %r2, [test_select_cc_f32_f16_param_3];
@@ -626,15 +626,16 @@ define <2 x float> @test_select_cc_f32_f16(<2 x float> %a, <2 x float> %b,
 ; CHECK-F16-NEXT:    mov.b64 {%r5, %r6}, %rd1;
 ; CHECK-F16-NEXT:    selp.f32 %r7, %r6, %r4, %p2;
 ; CHECK-F16-NEXT:    selp.f32 %r8, %r5, %r3, %p1;
-; CHECK-F16-NEXT:    st.param.v2.b32 [func_retval0], {%r8, %r7};
+; CHECK-F16-NEXT:    mov.b64 %rd3, {%r8, %r7};
+; CHECK-F16-NEXT:    st.param.b64 [func_retval0], %rd3;
 ; CHECK-F16-NEXT:    ret;
 ;
 ; CHECK-NOF16-LABEL: test_select_cc_f32_f16(
 ; CHECK-NOF16:       {
 ; CHECK-NOF16-NEXT:    .reg .pred %p<3>;
 ; CHECK-NOF16-NEXT:    .reg .b16 %rs<5>;
 ; CHECK-NOF16-NEXT:    .reg .b32 %r<13>;
-; CHECK-NOF16-NEXT:    .reg .b64 %rd<3>;
+; CHECK-NOF16-NEXT:    .reg .b64 %rd<4>;
 ; CHECK-NOF16-EMPTY:
 ; CHECK-NOF16-NEXT:  // %bb.0:
 ; CHECK-NOF16-NEXT:    ld.param.b32 %r2, [test_select_cc_f32_f16_param_3];
@@ -653,7 +654,8 @@ define <2 x float> @test_select_cc_f32_f16(<2 x float> %a, <2 x float> %b,
 ; CHECK-NOF16-NEXT:    mov.b64 {%r9, %r10}, %rd1;
 ; CHECK-NOF16-NEXT:    selp.f32 %r11, %r10, %r8, %p2;
 ; CHECK-NOF16-NEXT:    selp.f32 %r12, %r9, %r7, %p1;
-; CHECK-NOF16-NEXT:    st.param.v2.b32 [func_retval0], {%r12, %r11};
+; CHECK-NOF16-NEXT:    mov.b64 %rd3, {%r12, %r11};
+; CHECK-NOF16-NEXT:    st.param.b64 [func_retval0], %rd3;
 ; CHECK-NOF16-NEXT:    ret;
                                            <2 x half> %c, <2 x half> %d) #0 {
   %cc = fcmp une <2 x half> %c, %d
@@ -1563,13 +1565,15 @@ define <2 x float> @test_fpext_2xfloat(<2 x half> %a) #0 {
 ; CHECK:       {
 ; CHECK-NEXT:    .reg .b16 %rs<3>;
 ; CHECK-NEXT:    .reg .b32 %r<4>;
+; CHECK-NEXT:    .reg .b64 %rd<2>;
 ; CHECK-EMPTY:
 ; CHECK-NEXT:  // %bb.0:
 ; CHECK-NEXT:    ld.param.b32 %r1, [test_fpext_2xfloat_param_0];
 ; CHECK-NEXT:    mov.b32 {%rs1, %rs2}, %r1;
 ; CHECK-NEXT:    cvt.f32.f16 %r2, %rs2;
 ; CHECK-NEXT:    cvt.f32.f16 %r3, %rs1;
-; CHECK-NEXT:    st.param.v2.b32 [func_retval0], {%r3, %r2};
+; CHECK-NEXT:    mov.b64 %rd1, {%r3, %r2};
+; CHECK-NEXT:    st.param.b64 [func_retval0], %rd1;
 ; CHECK-NEXT:    ret;
   %r = fpext <2 x half> %a to <2 x float>
   ret <2 x float> %r
@@ -2054,6 +2058,7 @@ define <2 x float> @test_copysign_extended(<2 x half> %a, <2 x half> %b) #0 {
 ; CHECK-F16:       {
 ; CHECK-F16-NEXT:    .reg .b16 %rs<3>;
 ; CHECK-F16-NEXT:    .reg .b32 %r<8>;
+; CHECK-F16-NEXT:    .reg .b64 %rd<2>;
 ; CHECK-F16-EMPTY:
 ; CHECK-F16-NEXT:  // %bb.0:
 ; CHECK-F16-NEXT:    ld.param.b32 %r2, [test_copysign_extended_param_1];
@@ -2064,13 +2069,15 @@ define <2 x float> @test_copysign_extended(<2 x half> %a, <2 x half> %b) #0 {
 ; CHECK-F16-NEXT:    mov.b32 {%rs1, %rs2}, %r5;
 ; CHECK-F16-NEXT:    cvt.f32.f16 %r6, %rs2;
 ; CHECK-F16-NEXT:    cvt.f32.f16 %r7, %rs1;
-; CHECK-F16-NEXT:    st.param.v2.b32 [func_retval0], {%r7, %r6};
+; CHECK-F16-NEXT:    mov.b64 %rd1, {%r7, %r6};
+; CHECK-F16-NEXT:    st.param.b64 [func_retval0], %rd1;
 ; CHECK-F16-NEXT:    ret;
 ;
 ; CHECK-NOF16-LABEL: test_copysign_extended(
 ; CHECK-NOF16:       {
 ; CHECK-NOF16-NEXT:    .reg .b16 %rs<11>;
 ; CHECK-NOF16-NEXT:    .reg .b32 %r<5>;
+; CHECK-NOF16-NEXT:    .reg .b64 %rd<2>;
 ; CHECK-NOF16-EMPTY:
 ; CHECK-NOF16-NEXT:  // %bb.0:
 ; CHECK-NOF16-NEXT:    ld.param.b32 %r2, [test_copysign_extended_param_1];
@@ -2085,7 +2092,8 @@ define <2 x float> @test_copysign_extended(<2 x half> %a, <2 x half> %b) #0 {
 ; CHECK-NOF16-NEXT:    or.b16 %rs10, %rs9, %rs8;
 ; CHECK-NOF16-NEXT:    cvt.f32.f16 %r3, %rs10;
 ; CHECK-NOF16-NEXT:    cvt.f32.f16 %r4, %rs7;
-; CHECK-NOF16-NEXT:    st.param.v2.b32 [func_retval0], {%r4, %r3};
+; CHECK-NOF16-NEXT:    mov.b64 %rd1, {%r4, %r3};
+; CHECK-NOF16-NEXT:    st.param.b64 [func_retval0], %rd1;
 ; CHECK-NOF16-NEXT:    ret;
   %r = call <2 x half> @llvm.copysign.f16(<2 x half> %a, <2 x half> %b)
   %xr = fpext <2 x half> %r to <2 x float>
diff --git a/llvm/test/CodeGen/NVPTX/f32x2-instructions.ll b/llvm/test/CodeGen/NVPTX/f32x2-instructions.ll
