[SelectionDAG][RISCV] Avoid store merging across function calls

llvm/include/llvm/CodeGen/TargetLowering.h

@@ -3520,6 +3520,10 @@ class TargetLoweringBase {
   /// The default implementation just freezes the set of reserved registers.
   virtual void finalizeLowering(MachineFunction &MF) const;
 
+  /// Returns true if it's profitable to allow merging store of loads when there
+  /// are functions calls between the load and the store.
+  virtual bool shouldMergeStoreOfLoadsOverCall(EVT, EVT) const { return true; }
+
   //===----------------------------------------------------------------------===//
   //  GlobalISel Hooks
   //===----------------------------------------------------------------------===//

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

@@ -792,6 +792,12 @@ namespace {
         SmallVectorImpl<MemOpLink> &StoreNodes, unsigned NumStores,
         SDNode *RootNode);
 
+    /// Helper function for tryStoreMergeOfLoads. Checks if the load/store
+    /// chain has a call in it. \return True if a call is found.
+    bool hasCallInLdStChain(SmallVectorImpl<MemOpLink> &StoreNodes,
+                            SmallVectorImpl<MemOpLink> &LoadNodes,
+                            unsigned NumStores);
+
     /// This is a helper function for mergeConsecutiveStores. Given a list of
     /// store candidates, find the first N that are consecutive in memory.
     /// Returns 0 if there are not at least 2 consecutive stores to try merging.
@@ -21152,6 +21158,48 @@ bool DAGCombiner::checkMergeStoreCandidatesForDependencies(
   return true;
 }
 
+bool DAGCombiner::hasCallInLdStChain(SmallVectorImpl<MemOpLink> &StoreNodes,
+                                     SmallVectorImpl<MemOpLink> &LoadNodes,
+                                     unsigned NumStores) {
+  for (unsigned i = 0; i < NumStores; ++i) {
+    StoreSDNode *St = cast<StoreSDNode>(StoreNodes[i].MemNode);
+    LoadSDNode *Ld = cast<LoadSDNode>(LoadNodes[i].MemNode);
+    assert(Ld == cast<LoadSDNode>(peekThroughBitcasts(St->getValue())) &&
+           "Load and store mismatch");
+
+    SmallPtrSet<const SDNode *, 32> Visited;
+    SmallVector<std::pair<const SDNode *, bool>, 8> Worklist;
+    Worklist.emplace_back(St->getChain().getNode(), false);
+
+    while (!Worklist.empty()) {
+      auto [Node, FoundCall] = Worklist.pop_back_val();
+      if (!Visited.insert(Node).second || Node->getNumOperands() == 0)
+        continue;
+
+      switch (Node->getOpcode()) {
+      case ISD::CALLSEQ_END:
+        Worklist.emplace_back(Node->getOperand(0).getNode(), true);
+        break;
+      case ISD::TokenFactor:
+        for (SDValue Op : Node->ops())
+          Worklist.emplace_back(Op.getNode(), FoundCall);
+        break;
+      case ISD::LOAD:
+        if (Node == Ld)
+          return FoundCall;
+        [[fallthrough]];
+      default:
+        assert(Node->getOperand(0).getValueType() == MVT::Other &&
+               "Invalid chain type");
+        Worklist.emplace_back(Node->getOperand(0).getNode(), FoundCall);
+        break;
+      }
+    }
+    return false;
+  }
+  return false;
+}
+
 unsigned
 DAGCombiner::getConsecutiveStores(SmallVectorImpl<MemOpLink> &StoreNodes,
                                   int64_t ElementSizeBytes) const {
@@ -21598,6 +21646,15 @@ bool DAGCombiner::tryStoreMergeOfLoads(SmallVectorImpl<MemOpLink> &StoreNodes,
       JointMemOpVT = EVT::getIntegerVT(Context, SizeInBits);
     }
 
+    // Check if there is a call in the load/store chain.
+    if (!TLI.shouldMergeStoreOfLoadsOverCall(MemVT, JointMemOpVT) &&
+        hasCallInLdStChain(StoreNodes, LoadNodes, NumElem)) {
+      StoreNodes.erase(StoreNodes.begin(), StoreNodes.begin() + NumElem);
+      LoadNodes.erase(LoadNodes.begin(), LoadNodes.begin() + NumElem);
+      NumConsecutiveStores -= NumElem;
+      continue;
+    }
+
     SDLoc LoadDL(LoadNodes[0].MemNode);
     SDLoc StoreDL(StoreNodes[0].MemNode);
 

llvm/lib/Target/RISCV/RISCVISelLowering.h

@@ -1070,6 +1070,13 @@ class RISCVTargetLowering : public TargetLowering {
     return false;
   }
 
+  /// Disables storing and loading vectors by default when there are function
+  /// calls between the load and store, since these are more expensive than just
+  /// using scalars
+  bool shouldMergeStoreOfLoadsOverCall(EVT SrcVT, EVT MergedVT) const override {
+    return SrcVT.isScalarInteger() == MergedVT.isScalarInteger();
+  }
+
   /// For available scheduling models FDIV + two independent FMULs are much
   /// faster than two FDIVs.
   unsigned combineRepeatedFPDivisors() const override;

llvm/test/CodeGen/RISCV/stores-of-loads-merging.ll

@@ -1,7 +1,6 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc < %s -mtriple=riscv64 -mattr=+v | FileCheck %s
 
-declare i32 @llvm.experimental.constrained.fptosi.i32.f64(double, metadata)
 declare void @g()
 
 ; TODO: Merging scalars into vectors is unprofitable because we have no
@@ -15,40 +14,40 @@ define void @f(ptr %m, ptr %n, ptr %p, ptr %q, ptr %r, ptr %s, double %t) {
 ; CHECK-NEXT:    sd s0, 32(sp) # 8-byte Folded Spill
 ; CHECK-NEXT:    sd s1, 24(sp) # 8-byte Folded Spill
 ; CHECK-NEXT:    sd s2, 16(sp) # 8-byte Folded Spill
+; CHECK-NEXT:    sd s3, 8(sp) # 8-byte Folded Spill
+; CHECK-NEXT:    sd s4, 0(sp) # 8-byte Folded Spill
 ; CHECK-NEXT:    .cfi_offset ra, -8
 ; CHECK-NEXT:    .cfi_offset s0, -16
 ; CHECK-NEXT:    .cfi_offset s1, -24
 ; CHECK-NEXT:    .cfi_offset s2, -32
-; CHECK-NEXT:    csrr a6, vlenb
-; CHECK-NEXT:    sub sp, sp, a6
-; CHECK-NEXT:    .cfi_escape 0x0f, 0x0d, 0x72, 0x00, 0x11, 0x30, 0x22, 0x11, 0x01, 0x92, 0xa2, 0x38, 0x00, 0x1e, 0x22 # sp + 48 + 1 * vlenb
+; CHECK-NEXT:    .cfi_offset s3, -40
+; CHECK-NEXT:    .cfi_offset s4, -48
 ; CHECK-NEXT:    mv s0, a5
 ; CHECK-NEXT:    mv s1, a4
 ; CHECK-NEXT:    vsetivli zero, 2, e64, m1, ta, ma
 ; CHECK-NEXT:    vle64.v v8, (a0)
 ; CHECK-NEXT:    vse64.v v8, (a1)
-; CHECK-NEXT:    vle64.v v8, (a2)
-; CHECK-NEXT:    addi a0, sp, 16
-; CHECK-NEXT:    vs1r.v v8, (a0) # Unknown-size Folded Spill
+; CHECK-NEXT:    ld s3, 0(a2)
+; CHECK-NEXT:    ld s4, 8(a2)
 ; CHECK-NEXT:    mv s2, a3
 ; CHECK-NEXT:    call g
-; CHECK-NEXT:    addi a0, sp, 16
-; CHECK-NEXT:    vl1r.v v8, (a0) # Unknown-size Folded Reload
+; CHECK-NEXT:    sd s3, 0(s2)
+; CHECK-NEXT:    sd s4, 8(s2)
 ; CHECK-NEXT:    vsetivli zero, 2, e64, m1, ta, ma
-; CHECK-NEXT:    vse64.v v8, (s2)
 ; CHECK-NEXT:    vle64.v v8, (s1)
 ; CHECK-NEXT:    vse64.v v8, (s0)
-; CHECK-NEXT:    csrr a0, vlenb
-; CHECK-NEXT:    add sp, sp, a0
-; CHECK-NEXT:    .cfi_def_cfa sp, 48
 ; CHECK-NEXT:    ld ra, 40(sp) # 8-byte Folded Reload
 ; CHECK-NEXT:    ld s0, 32(sp) # 8-byte Folded Reload
 ; CHECK-NEXT:    ld s1, 24(sp) # 8-byte Folded Reload
 ; CHECK-NEXT:    ld s2, 16(sp) # 8-byte Folded Reload
+; CHECK-NEXT:    ld s3, 8(sp) # 8-byte Folded Reload
+; CHECK-NEXT:    ld s4, 0(sp) # 8-byte Folded Reload
 ; CHECK-NEXT:    .cfi_restore ra
 ; CHECK-NEXT:    .cfi_restore s0
 ; CHECK-NEXT:    .cfi_restore s1
 ; CHECK-NEXT:    .cfi_restore s2
+; CHECK-NEXT:    .cfi_restore s3
+; CHECK-NEXT:    .cfi_restore s4
 ; CHECK-NEXT:    addi sp, sp, 48
 ; CHECK-NEXT:    .cfi_def_cfa_offset 0
 ; CHECK-NEXT:    ret
@@ -77,13 +76,13 @@ define void @f(ptr %m, ptr %n, ptr %p, ptr %q, ptr %r, ptr %s, double %t) {
   ret void
 }
 
-define void @f1(ptr %m, ptr %n, ptr %p, ptr %q, ptr %r, ptr %s, double %t) {
+define void @f1(ptr %p, ptr %q, double %t) {
 ; CHECK-LABEL: f1:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetivli zero, 2, e64, m1, ta, ma
-; CHECK-NEXT:    vle64.v v8, (a2)
+; CHECK-NEXT:    vle64.v v8, (a0)
 ; CHECK-NEXT:    fcvt.wu.d a0, fa0, rtz
-; CHECK-NEXT:    vse64.v v8, (a3)
+; CHECK-NEXT:    vse64.v v8, (a1)
 ; CHECK-NEXT:    ret
   %x0 = load i64, ptr %p
   %p.1 = getelementptr i64, ptr %p, i64 1
@@ -92,7 +91,6 @@ define void @f1(ptr %m, ptr %n, ptr %p, ptr %q, ptr %r, ptr %s, double %t) {
   store i64 %x0, ptr %q
   %q.1 = getelementptr i64, ptr %q, i64 1
   store i64 %x1, ptr %q.1
-
   ret void
 }