[CHERIoT] Use captable relocations to optimize code size

lld/ELF/Arch/Cheri.cpp

@@ -341,7 +341,8 @@ void CheriCapRelocsSection::addCapReloc(CheriCapRelocLocation loc,
             "\n>>> Target: " + target.verboseToString(ctx));
 
   bool canWriteLoc = (loc.section->flags & SHF_WRITE) || !ctx.arg.zText;
-  if (!canWriteLoc) {
+  if (!canWriteLoc &&
+      !(ctx.arg.isCheriot && loc.section->name == ".captable")) {
     readOnlyCapRelocsError(ctx, *target.sym(),
                            "\n>>> referenced by " + sourceMsg());
     return;
@@ -595,7 +596,9 @@ void CheriCapRelocsSection::writeTo(uint8_t *buf) {
 CheriCapTableSection::CheriCapTableSection(Ctx &ctx)
     : SyntheticSection(
           ctx, ".captable", SHT_PROGBITS,
-          SHF_ALLOC | SHF_WRITE, /* XXX: actually RELRO for BIND_NOW*/
+          SHF_ALLOC | (ctx.arg.isCheriot
+                           ? 0
+                           : SHF_WRITE), /* XXX: actually RELRO for BIND_NOW*/
           ctx.arg.capabilitySize) {
   assert(ctx.arg.capabilitySize > 0);
   this->entsize = ctx.arg.capabilitySize;

lld/ELF/Arch/RISCV.cpp

@@ -124,6 +124,19 @@ static uint32_t setLO12_S(uint32_t insn, uint32_t imm) {
          (extractBits(imm, 4, 0) << 7);
 }
 
+static bool isCheriotHighPartReloc(RelType ty) {
+  switch (ty) {
+  case R_RISCV_CHERIOT_COMPARTMENT_HI:
+  case R_RISCV_CHERI_CAPTAB_PCREL_HI20:
+    // FIXME: This can be enabled once they handle cheriot 11-bit AUIPC shifts.
+    // case R_RISCV_CHERI_TLS_IE_CAPTAB_PCREL_HI20:
+    // case R_RISCV_CHERI_TLS_GD_CAPTAB_PCREL_HI20:
+    return true;
+  default:
+    return false;
+  }
+}
+
 RISCV::RISCV(Ctx &ctx) : TargetInfo(ctx) {
   copyRel = R_RISCV_COPY;
   pltRel = R_RISCV_JUMP_SLOT;
@@ -540,6 +553,17 @@ void RISCV::relocate(uint8_t *loc, const Relocation &rel, uint64_t val) const {
   case R_RISCV_TLS_GOT_HI20:
   case R_RISCV_TPREL_HI20:
   case R_RISCV_HI20: {
+    if (ctx.arg.isCheriot && isCheriotHighPartReloc(rel.type)) {
+      // Similar to R_RISCV_CHERIOT_COMPARTMENT_HI, except that these
+      // always use AUIPCC.
+      if (int64_t(val) < 0)
+        val = (val + 0x7ff) & ~0x7ff;
+      val = int64_t(val) >> 11;
+      checkInt(ctx, loc, SignExtend64(val + 0x800, bits) >> 12, 20, rel);
+      write32le(loc, (read32le(loc) & 0xFFF) | (val << 12));
+      return;
+    }
+
     uint64_t hi = val + 0x800;
     checkInt(ctx, loc, SignExtend64(hi, bits) >> 12, 20, rel);
     write32le(loc, (read32le(loc) & 0xFFF) | (hi & 0xFFFFF000));
@@ -1096,19 +1120,24 @@ static bool rewriteCheriotLowRelocs(Ctx &ctx, InputSection &sec) {
 
         const Relocation *target = nullptr;
         for (auto it = range.first; it != range.second; ++it)
-          if (it->type == R_RISCV_CHERIOT_COMPARTMENT_HI) {
+          if (isCheriotHighPartReloc(it->type)) {
             target = &*it;
             break;
           }
         if (!target) {
-          error(
-              "Could not find R_RISCV_CHERIOT_COMPARTMENT_HI relocation for " +
-              toStr(ctx, *r.sym));
+          error("Could not find high-part relocation for " +
+                toStr(ctx, *r.sym) + " at 0x" +
+                llvm::utohexstr(r.offset, false, 8) + "@" + toStr(ctx, &sec));
         }
         // If the target is PCC-relative then the auipcc can't be erased and so
         // skip the rewriting.
         if (isPCCRelative(ctx, nullptr, target->sym))
           continue;
+
+        // captab relocations similarly are never relaxed, so we can skip them.
+        if (target->type == R_RISCV_CHERI_CAPTAB_PCREL_HI20)
+          continue;
+
         // Update our relocation to point to the target thing.
         r.sym = target->sym;
         r.addend = target->addend;

lld/ELF/InputSection.cpp

@@ -1056,7 +1056,9 @@ uint64_t InputSectionBase::getRelocTargetVA(Ctx &ctx, const Relocation &r,
   case R_CHERIOT_COMPARTMENT_CGPREL_LO_I: {
     if (isPCCRelative(ctx, nullptr, r.sym)) {
       if (const Relocation *hiRel = getRISCVPCRelHi20(ctx, this, r)) {
-        if (isPCCRelative(ctx, nullptr, hiRel->sym))
+        if (isPCCRelative(ctx, nullptr, hiRel->sym) ||
+            // captab is always PC-relative on CHERIoT
+            hiRel->type == R_RISCV_CHERI_CAPTAB_PCREL_HI20)
           return getRelocTargetVA(ctx, *hiRel, r.sym->getVA(ctx));
         return getBiasedCGPOffsetLo12(ctx, *hiRel->sym);
       }

llvm/lib/Target/RISCV/RISCVExpandPseudoInsts.cpp

@@ -707,7 +707,8 @@ bool RISCVExpandPseudo::expandAuipccInstPair(
       MF->getSubtarget<RISCVSubtarget>().isRV32E() && Symbol.isGlobal() &&
       isa<GlobalVariable>(Symbol.getGlobal()) &&
       (cast<GlobalVariable>(Symbol.getGlobal())->getSection() !=
-       ".compartment_imports"))
+       ".compartment_imports") &&
+      FlagsHi != RISCVII::MO_CAPTAB_PCREL_HI)
     BuildMI(NewMBB, DL, TII->get(RISCV::CSetBoundsImm), DestReg)
         .addReg(DestReg)
         .addDisp(Symbol, 0, RISCVII::MO_CHERIOT_COMPARTMENT_SIZE);

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

@@ -94,6 +94,11 @@ static cl::opt<bool>
                                "be combined with a shift"),
                       cl::init(true));
 
+static cl::opt<bool> CheriotDisableCapTable(
+    "cheriot-disable-cap-table", cl::Hidden,
+    cl::desc("Disable cap tables for references to globals on CHERIoT"),
+    cl::init(false));
+
 RISCVTargetLowering::RISCVTargetLowering(const TargetMachine &TM,
                                          const RISCVSubtarget &STI)
     : TargetLowering(TM), Subtarget(STI) {
@@ -8977,14 +8982,17 @@ static SDValue getLargeExternalSymbol(ExternalSymbolSDNode *N, const SDLoc &DL,
 template <class NodeTy>
 SDValue RISCVTargetLowering::getAddr(NodeTy *N, EVT Ty, SelectionDAG &DAG,
                                      bool IsLocal, bool CanDeriveFromPcc,
-                                     bool IsExternWeak) const {
+                                     bool IsExternWeak,
+                                     const GlobalValue *GV) const {
   SDLoc DL(N);
 
   if (RISCVABI::isCheriPureCapABI(Subtarget.getTargetABI())) {
     bool IsCheriot = Subtarget.getTargetABI() == RISCVABI::ABI_CHERIOT ||
                      Subtarget.getTargetABI() == RISCVABI::ABI_CHERIOT_BAREMETAL;
     SDValue Addr = getTargetNode(N, DL, Ty, DAG, 0);
-    if ((IsLocal && CanDeriveFromPcc) || IsCheriot) {
+    bool NotHighUseGV = !GV || GV->getNumUses() < 4;
+    if ((IsLocal && CanDeriveFromPcc) ||
+        (IsCheriot && (CheriotDisableCapTable || NotHighUseGV))) {
       // Use PC-relative addressing to access the symbol. This generates the
       // pattern (PseudoCLLC sym), which expands to
       // (cincoffsetimm (auipcc %pcrel_hi(sym)) %pcrel_lo(auipc)).
@@ -9121,7 +9129,7 @@ SDValue RISCVTargetLowering::lowerGlobalAddress(SDValue Op,
   }
 
   return getAddr(N, Ty, DAG, GV->isDSOLocal(), /*CanDeriveFromPcc=*/false,
-                 GV->hasExternalWeakLinkage());
+                 GV->hasExternalWeakLinkage(), GV);
 }
 
 SDValue RISCVTargetLowering::lowerBlockAddress(SDValue Op,

llvm/lib/Target/RISCV/RISCVISelLowering.h

@@ -500,7 +500,8 @@ class RISCVTargetLowering : public TargetLowering {
 
   template <class NodeTy>
   SDValue getAddr(NodeTy *N, EVT Ty, SelectionDAG &DAG, bool IsLocal,
-                  bool CanDeriveFromPcc, bool IsExternWeak = false) const;
+                  bool CanDeriveFromPcc, bool IsExternWeak = false,
+                  const GlobalValue *GV = nullptr) const;
   SDValue getStaticTLSAddr(GlobalAddressSDNode *N, EVT Ty, SelectionDAG &DAG,
                            bool NotLocal) const;
   SDValue getDynamicTLSAddr(GlobalAddressSDNode *N, EVT Ty,