[RISCV] Fix evaluating %pcrel_lo against global and weak symbols

Summary:
Previously, we would erroneously turn %pcrel_lo(label), where label has
a %pcrel_hi against a weak symbol, into %pcrel_lo(label + offset), as
evaluatePCRelLo would believe the target independent logic was going to
fold it. Moreover, even if that were fixed, shouldForceRelocation lacks
an MCAsmLayout and thus cannot evaluate the %pcrel_hi fixup to a value
and check the symbol, so we would then erroneously constant-fold the
%pcrel_lo whilst leaving the %pcrel_hi intact. After D72197, this same
sequence also occurs for symbols with global binding, which is triggered
in real-world code.

Instead, as discussed in D71978, we introduce a new FKF_IsTarget flag to
avoid these kinds of issues. All the resolution logic happens in one
place, with no coordination required between RISCAsmBackend and
RISCVMCExpr to ensure they implement the same logic twice. Although the
implementation of %pcrel_hi can be left as target independent, we make
it target dependent to ensure that they are handled identically to
%pcrel_lo, otherwise we risk one of them being constant folded but the
other being preserved. This also allows us to properly support fixup
pairs where the instructions are in different fragments.

Reviewers: asb, lenary, efriedma

Reviewed By: efriedma

Subscribers: arichardson, hiraditya, rbar, johnrusso, simoncook, sabuasal, niosHD, kito-cheng, shiva0217, MaskRay, zzheng, edward-jones, rogfer01, MartinMosbeck, brucehoult, the_o, rkruppe, PkmX, jocewei, psnobl, benna, Jim, s.egerton, pzheng, sameer.abuasal, apazos, luismarques, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D73211

(cherry picked from commit 3f5976c97dbfefb4669abcf968bd79a9a64c18e0)

4 files changed, 79 insertions, 108 deletions

diff --git a/llvm/lib/Target/RISCV/MCTargetDesc/RISCVAsmBackend.cpp b/llvm/lib/Target/RISCV/MCTargetDesc/RISCVAsmBackend.cpp
index 5881a0a86ef..373d0ccb185 100644
--- a/llvm/lib/Target/RISCV/MCTargetDesc/RISCVAsmBackend.cpp
+++ b/llvm/lib/Target/RISCV/MCTargetDesc/RISCVAsmBackend.cpp
@@ -9,6 +9,7 @@
 #include "RISCVAsmBackend.h"
 #include "RISCVMCExpr.h"
 #include "llvm/ADT/APInt.h"
+#include "llvm/MC/MCAsmLayout.h"
 #include "llvm/MC/MCAssembler.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCDirectives.h"
@@ -28,8 +29,6 @@ using namespace llvm;
 bool RISCVAsmBackend::shouldForceRelocation(const MCAssembler &Asm,
                                             const MCFixup &Fixup,
                                             const MCValue &Target) {
-  bool ShouldForce = false;
-
   switch (Fixup.getTargetKind()) {
   default:
     break;
@@ -44,40 +43,9 @@ bool RISCVAsmBackend::shouldForceRelocation(const MCAssembler &Asm,
   case RISCV::fixup_riscv_tls_got_hi20:
   case RISCV::fixup_riscv_tls_gd_hi20:
     return true;
-  case RISCV::fixup_riscv_pcrel_lo12_i:
-  case RISCV::fixup_riscv_pcrel_lo12_s:
-    // For pcrel_lo12, force a relocation if the target of the corresponding
-    // pcrel_hi20 is not in the same fragment.
-    const MCFixup *T = cast<RISCVMCExpr>(Fixup.getValue())->getPCRelHiFixup();
-    if (!T) {
-      Asm.getContext().reportError(Fixup.getLoc(),
-                                   "could not find corresponding %pcrel_hi");
-      return false;
-    }
-
-    switch (T->getTargetKind()) {
-    default:
-      llvm_unreachable("Unexpected fixup kind for pcrel_lo12");
-      break;
-    case RISCV::fixup_riscv_got_hi20:
-    case RISCV::fixup_riscv_tls_got_hi20:
-    case RISCV::fixup_riscv_tls_gd_hi20:
-      ShouldForce = true;
-      break;
-    case RISCV::fixup_riscv_pcrel_hi20: {
-      MCFragment *TFragment = T->getValue()->findAssociatedFragment();
-      MCFragment *FixupFragment = Fixup.getValue()->findAssociatedFragment();
-      assert(FixupFragment && "We should have a fragment for this fixup");
-      ShouldForce =
-          !TFragment || TFragment->getParent() != FixupFragment->getParent();
-      break;
-    }
-    }
-    break;
   }
 
-  return ShouldForce || STI.getFeatureBits()[RISCV::FeatureRelax] ||
-         ForceRelocs;
+  return STI.getFeatureBits()[RISCV::FeatureRelax] || ForceRelocs;
 }
 
 bool RISCVAsmBackend::fixupNeedsRelaxationAdvanced(const MCFixup &Fixup,
@@ -284,6 +252,67 @@ static uint64_t adjustFixupValue(const MCFixup &Fixup, uint64_t Value,
   }
 }
 
+bool RISCVAsmBackend::evaluateTargetFixup(
+    const MCAssembler &Asm, const MCAsmLayout &Layout, const MCFixup &Fixup,
+    const MCFragment *DF, const MCValue &Target, uint64_t &Value,
+    bool &WasForced) {
+  const MCFixup *AUIPCFixup;
+  const MCFragment *AUIPCDF;
+  MCValue AUIPCTarget;
+  switch (Fixup.getTargetKind()) {
+  default:
+    llvm_unreachable("Unexpected fixup kind!");
+  case RISCV::fixup_riscv_pcrel_hi20:
+    AUIPCFixup = &Fixup;
+    AUIPCDF = DF;
+    AUIPCTarget = Target;
+    break;
+  case RISCV::fixup_riscv_pcrel_lo12_i:
+  case RISCV::fixup_riscv_pcrel_lo12_s: {
+    AUIPCFixup = cast<RISCVMCExpr>(Fixup.getValue())->getPCRelHiFixup(&AUIPCDF);
+    if (!AUIPCFixup) {
+      Asm.getContext().reportError(Fixup.getLoc(),
+                                   "could not find corresponding %pcrel_hi");
+      return true;
+    }
+
+    // MCAssembler::evaluateFixup will emit an error for this case when it sees
+    // the %pcrel_hi, so don't duplicate it when also seeing the %pcrel_lo.
+    const MCExpr *AUIPCExpr = AUIPCFixup->getValue();
+    if (!AUIPCExpr->evaluateAsRelocatable(AUIPCTarget, &Layout, AUIPCFixup))
+      return true;
+    break;
+  }
+  }
+
+  if (!AUIPCTarget.getSymA() || AUIPCTarget.getSymB())
+    return false;
+
+  const MCSymbolRefExpr *A = AUIPCTarget.getSymA();
+  const MCSymbol &SA = A->getSymbol();
+  if (A->getKind() != MCSymbolRefExpr::VK_None || SA.isUndefined())
+    return false;
+
+  auto *Writer = Asm.getWriterPtr();
+  if (!Writer)
+    return false;
+
+  bool IsResolved = Writer->isSymbolRefDifferenceFullyResolvedImpl(
+      Asm, SA, *AUIPCDF, false, true);
+  if (!IsResolved)
+    return false;
+
+  Value = Layout.getSymbolOffset(SA) + AUIPCTarget.getConstant();
+  Value -= Layout.getFragmentOffset(AUIPCDF) + AUIPCFixup->getOffset();
+
+  if (shouldForceRelocation(Asm, *AUIPCFixup, AUIPCTarget)) {
+    WasForced = true;
+    return false;
+  }
+
+  return true;
+}
+
 void RISCVAsmBackend::applyFixup(const MCAssembler &Asm, const MCFixup &Fixup,
                                  const MCValue &Target,
                                  MutableArrayRef<char> Data, uint64_t Value,
diff --git a/llvm/lib/Target/RISCV/MCTargetDesc/RISCVAsmBackend.h b/llvm/lib/Target/RISCV/MCTargetDesc/RISCVAsmBackend.h
index 254249c87dc..1c3c587355a 100644
--- a/llvm/lib/Target/RISCV/MCTargetDesc/RISCVAsmBackend.h
+++ b/llvm/lib/Target/RISCV/MCTargetDesc/RISCVAsmBackend.h
@@ -65,6 +65,11 @@ public:
                                      const MCAsmLayout &Layout,
                                      MCAlignFragment &AF) override;
 
+  bool evaluateTargetFixup(const MCAssembler &Asm, const MCAsmLayout &Layout,
+                           const MCFixup &Fixup, const MCFragment *DF,
+                           const MCValue &Target, uint64_t &Value,
+                           bool &WasForced) override;
+
   void applyFixup(const MCAssembler &Asm, const MCFixup &Fixup,
                   const MCValue &Target, MutableArrayRef<char> Data,
                   uint64_t Value, bool IsResolved,
@@ -101,9 +106,12 @@ public:
       { "fixup_riscv_hi20",         12,     20,  0 },
       { "fixup_riscv_lo12_i",       20,     12,  0 },
       { "fixup_riscv_lo12_s",        0,     32,  0 },
-      { "fixup_riscv_pcrel_hi20",   12,     20,  MCFixupKindInfo::FKF_IsPCRel },
-      { "fixup_riscv_pcrel_lo12_i", 20,     12,  MCFixupKindInfo::FKF_IsPCRel },
-      { "fixup_riscv_pcrel_lo12_s",  0,     32,  MCFixupKindInfo::FKF_IsPCRel },
+      { "fixup_riscv_pcrel_hi20",   12,     20,
+        MCFixupKindInfo::FKF_IsPCRel | MCFixupKindInfo::FKF_IsTarget },
+      { "fixup_riscv_pcrel_lo12_i", 20,     12,
+        MCFixupKindInfo::FKF_IsPCRel | MCFixupKindInfo::FKF_IsTarget },
+      { "fixup_riscv_pcrel_lo12_s",  0,     32,
+        MCFixupKindInfo::FKF_IsPCRel | MCFixupKindInfo::FKF_IsTarget },
       { "fixup_riscv_got_hi20",     12,     20,  MCFixupKindInfo::FKF_IsPCRel },
       { "fixup_riscv_tprel_hi20",   12,     20,  0 },
       { "fixup_riscv_tprel_lo12_i", 20,     12,  0 },
diff --git a/llvm/lib/Target/RISCV/MCTargetDesc/RISCVMCExpr.cpp b/llvm/lib/Target/RISCV/MCTargetDesc/RISCVMCExpr.cpp
index 7aa9b5e7d68..2a6f372e50b 100644
--- a/llvm/lib/Target/RISCV/MCTargetDesc/RISCVMCExpr.cpp
+++ b/llvm/lib/Target/RISCV/MCTargetDesc/RISCVMCExpr.cpp
@@ -47,7 +47,7 @@ void RISCVMCExpr::printImpl(raw_ostream &OS, const MCAsmInfo *MAI) const {
     OS << ')';
 }
 
-const MCFixup *RISCVMCExpr::getPCRelHiFixup() const {
+const MCFixup *RISCVMCExpr::getPCRelHiFixup(const MCFragment **DFOut) const {
   MCValue AUIPCLoc;
   if (!getSubExpr()->evaluateAsRelocatable(AUIPCLoc, nullptr, nullptr))
     return nullptr;
@@ -81,6 +81,8 @@ const MCFixup *RISCVMCExpr::getPCRelHiFixup() const {
     case RISCV::fixup_riscv_tls_got_hi20:
     case RISCV::fixup_riscv_tls_gd_hi20:
     case RISCV::fixup_riscv_pcrel_hi20:
+      if (DFOut)
+        *DFOut = DF;
       return &F;
     }
   }
@@ -88,74 +90,9 @@ const MCFixup *RISCVMCExpr::getPCRelHiFixup() const {
   return nullptr;
 }
 
-bool RISCVMCExpr::evaluatePCRelLo(MCValue &Res, const MCAsmLayout *Layout,
-                                  const MCFixup *Fixup) const {
-  // VK_RISCV_PCREL_LO has to be handled specially.  The MCExpr inside is
-  // actually the location of a auipc instruction with a VK_RISCV_PCREL_HI fixup
-  // pointing to the real target.  We need to generate an MCValue in the form of
-  // (<real target> + <offset from this fixup to the auipc fixup>).  The Fixup
-  // is pcrel relative to the VK_RISCV_PCREL_LO fixup, so we need to add the
-  // offset to the VK_RISCV_PCREL_HI Fixup from VK_RISCV_PCREL_LO to correct.
-
-  // Don't try to evaluate if the fixup will be forced as a relocation (e.g.
-  // as linker relaxation is enabled). If we evaluated pcrel_lo in this case,
-  // the modified fixup will be converted into a relocation that no longer
-  // points to the pcrel_hi as the linker requires.
-  auto &RAB =
-      static_cast<RISCVAsmBackend &>(Layout->getAssembler().getBackend());
-  if (RAB.willForceRelocations())
-    return false;
-
-  MCValue AUIPCLoc;
-  if (!getSubExpr()->evaluateAsValue(AUIPCLoc, *Layout))
-    return false;
-
-  const MCSymbolRefExpr *AUIPCSRE = AUIPCLoc.getSymA();
-  // Don't try to evaluate %pcrel_hi/%pcrel_lo pairs that cross fragment
-  // boundries.
-  if (!AUIPCSRE ||
-      findAssociatedFragment() != AUIPCSRE->findAssociatedFragment())
-    return false;
-
-  const MCSymbol *AUIPCSymbol = &AUIPCSRE->getSymbol();
-  if (!AUIPCSymbol)
-    return false;
-
-  const MCFixup *TargetFixup = getPCRelHiFixup();
-  if (!TargetFixup)
-    return false;
-
-  if ((unsigned)TargetFixup->getKind() != RISCV::fixup_riscv_pcrel_hi20)
-    return false;
-
-  MCValue Target;
-  if (!TargetFixup->getValue()->evaluateAsValue(Target, *Layout))
-    return false;
-
-  if (!Target.getSymA() || !Target.getSymA()->getSymbol().isInSection())
-    return false;
-
-  if (&Target.getSymA()->getSymbol().getSection() !=
-      findAssociatedFragment()->getParent())
-    return false;
-
-  // We must use TargetFixup rather than AUIPCSymbol here. They will almost
-  // always have the same offset, except for the case when AUIPCSymbol is at
-  // the end of a fragment and the fixup comes from offset 0 in the next
-  // fragment.
-  uint64_t AUIPCOffset = TargetFixup->getOffset();
-
-  Res = MCValue::get(Target.getSymA(), nullptr,
-                     Target.getConstant() + (Fixup->getOffset() - AUIPCOffset));
-  return true;
-}
-
 bool RISCVMCExpr::evaluateAsRelocatableImpl(MCValue &Res,
                                             const MCAsmLayout *Layout,
                                             const MCFixup *Fixup) const {
-  if (Kind == VK_RISCV_PCREL_LO && evaluatePCRelLo(Res, Layout, Fixup))
-    return true;
-
   if (!getSubExpr()->evaluateAsRelocatable(Res, Layout, Fixup))
     return false;
 
diff --git a/llvm/lib/Target/RISCV/MCTargetDesc/RISCVMCExpr.h b/llvm/lib/Target/RISCV/MCTargetDesc/RISCVMCExpr.h
index 921df376f3d..167e7d553e7 100644
--- a/llvm/lib/Target/RISCV/MCTargetDesc/RISCVMCExpr.h
+++ b/llvm/lib/Target/RISCV/MCTargetDesc/RISCVMCExpr.h
@@ -46,9 +46,6 @@ private:
 
   int64_t evaluateAsInt64(int64_t Value) const;
 
-  bool evaluatePCRelLo(MCValue &Res, const MCAsmLayout *Layout,
-                       const MCFixup *Fixup) const;
-
   explicit RISCVMCExpr(const MCExpr *Expr, VariantKind Kind)
       : Expr(Expr), Kind(Kind) {}
 
@@ -61,11 +58,11 @@ public:
   const MCExpr *getSubExpr() const { return Expr; }
 
   /// Get the corresponding PC-relative HI fixup that a VK_RISCV_PCREL_LO
-  /// points to.
+  /// points to, and optionally the fragment containing it.
   ///
   /// \returns nullptr if this isn't a VK_RISCV_PCREL_LO pointing to a
   /// known PC-relative HI fixup.
-  const MCFixup *getPCRelHiFixup() const;
+  const MCFixup *getPCRelHiFixup(const MCFragment **DFOut) const;
 
   void printImpl(raw_ostream &OS, const MCAsmInfo *MAI) const override;
   bool evaluateAsRelocatableImpl(MCValue &Res, const MCAsmLayout *Layout,