[lld][LoongArch] Support TLSDESC GD/LD to IE/LE (#123715)

Support TLSDESC to initial-exec or local-exec optimizations. Introduce a
new hook RE_LOONGARCH_RELAX_TLS_GD_TO_IE_PAGE_PC and use existing
R_RELAX_TLS_GD_TO_IE_ABS to support TLSDESC => IE, while use existing
R_RELAX_TLS_GD_TO_LE to support TLSDESC => LE.
    
In normal or medium code model, there are two forms of code sequences:
* pcalau12i  $a0, %desc_pc_hi20(sym_desc)
* addi.d     $a0, $a0, %desc_pc_lo12(sym_desc)
* ld.d       $ra, $a0, %desc_ld(sym_desc)
* jirl       $ra, $ra, %desc_call(sym_desc)
------
* pcaddi     $a0, %desc_pcrel_20(sym_desc)
* ld.d       $ra, $a0, %desc_ld(sym_desc)
* jirl       $ra, $ra, %desc_call(sym_desc)
    
Convert to IE:
* pcalau12i $a0, %ie_pc_hi20(sym_ie)
* ld.[wd]   $a0, $a0, %ie_pc_lo12(sym_ie)

Convert to LE:
* lu12i.w $a0, %le_hi20(sym_le) # le_hi20 != 0, otherwise NOP
* ori $a0 src, %le_lo12(sym_le) # le_hi20 != 0, src = $a0, otherwise src = $zero

Simplicity, whether tlsdescToIe or tlsdescToLe, we always tend to
convert the preceding instructions to NOPs, due to both forms of code
sequence (corresponding to relocation combinations:
R_LARCH_TLS_DESC_PC_HI20+R_LARCH_TLS_DESC_PC_LO12 and
R_LARCH_TLS_DESC_PCREL20_S2) have same process.
    
TODO: When relaxation enables, redundant NOPs can be removed. It will be
implemented in a future patch.
    
Note: All forms of TLSDESC code sequences should not appear interleaved
in the normal, medium or extreme code model, which compilers do not
generate and lld is unsupported. This is thanks to the guard in
PostRASchedulerList.cpp in llvm.
```
Calls are not scheduling boundaries before register allocation,
but post-ra we don't gain anything by scheduling across calls
since we don't need to worry about register pressure.
```

Author: ylzsx

lld/ELF/Arch/LoongArch.cpp

@@ -39,8 +39,13 @@ class LoongArch final : public TargetInfo {
   void relocate(uint8_t *loc, const Relocation &rel,
                 uint64_t val) const override;
   bool relaxOnce(int pass) const override;
+  RelExpr adjustTlsExpr(RelType type, RelExpr expr) const override;
   void relocateAlloc(InputSectionBase &sec, uint8_t *buf) const override;
   void finalizeRelax(int passes) const override;
+
+private:
+  void tlsdescToIe(uint8_t *loc, const Relocation &rel, uint64_t val) const;
+  void tlsdescToLe(uint8_t *loc, const Relocation &rel, uint64_t val) const;
 };
 } // end anonymous namespace
 
@@ -58,6 +63,7 @@ enum Op {
   LU12I_W = 0x14000000,
   PCADDI = 0x18000000,
   PCADDU12I = 0x1c000000,
+  PCALAU12I = 0x1a000000,
   LD_W = 0x28800000,
   LD_D = 0x28c00000,
   JIRL = 0x4c000000,
@@ -69,6 +75,7 @@ enum Reg {
   R_ZERO = 0,
   R_RA = 1,
   R_TP = 2,
+  R_A0 = 4,
   R_T0 = 12,
   R_T1 = 13,
   R_T2 = 14,
@@ -961,7 +968,8 @@ static bool relax(Ctx &ctx, InputSection &sec) {
     case R_LARCH_TLS_LD_PC_HI20:
     case R_LARCH_TLS_DESC_PC_HI20:
       // The overflow check for i+2 will be carried out in isPairRelaxable.
-      if (isPairRelaxable(relocs, i))
+      if (r.expr != RE_LOONGARCH_RELAX_TLS_GD_TO_IE_PAGE_PC &&
+          r.expr != R_RELAX_TLS_GD_TO_LE && isPairRelaxable(relocs, i))
         relaxPCHi20Lo12(ctx, sec, i, loc, r, relocs[i + 2], remove);
       break;
     case R_LARCH_CALL36:
@@ -1046,6 +1054,104 @@ static void tlsIeToLe(uint8_t *loc, const Relocation &rel, uint64_t val) {
   }
 }
 
+// Convert TLSDESC GD/LD to IE.
+// In normal or medium code model, there are two forms of code sequences:
+//  * pcalau12i  $a0, %desc_pc_hi20(sym_desc)
+//  * addi.d     $a0, $a0, %desc_pc_lo12(sym_desc)
+//  * ld.d       $ra, $a0, %desc_ld(sym_desc)
+//  * jirl       $ra, $ra, %desc_call(sym_desc)
+//  ------
+//  * pcaddi $a0, %desc_pcrel_20(a)
+//  * load $ra, $a0, %desc_ld(a)
+//  * jirl $ra, $ra, %desc_call(a)
+//
+// The code sequence obtained is as follows:
+//  * pcalau12i $a0, %ie_pc_hi20(sym_ie)
+//  * ld.[wd]   $a0, $a0, %ie_pc_lo12(sym_ie)
+//
+// Simplicity, whether tlsdescToIe or tlsdescToLe, we always tend to convert the
+// preceding instructions to NOPs, due to both forms of code sequence
+// (corresponding to relocation combinations:
+// R_LARCH_TLS_DESC_PC_HI20+R_LARCH_TLS_DESC_PC_LO12 and
+// R_LARCH_TLS_DESC_PCREL20_S2) have same process.
+//
+// When relaxation enables, redundant NOPs can be removed.
+void LoongArch::tlsdescToIe(uint8_t *loc, const Relocation &rel,
+                            uint64_t val) const {
+  switch (rel.type) {
+  case R_LARCH_TLS_DESC_PC_HI20:
+  case R_LARCH_TLS_DESC_PC_LO12:
+  case R_LARCH_TLS_DESC_PCREL20_S2:
+    write32le(loc, insn(ANDI, R_ZERO, R_ZERO, 0)); // nop
+    break;
+  case R_LARCH_TLS_DESC_LD:
+    write32le(loc, insn(PCALAU12I, R_A0, 0, 0)); // pcalau12i $a0, %ie_pc_hi20
+    relocateNoSym(loc, R_LARCH_TLS_IE_PC_HI20, val);
+    break;
+  case R_LARCH_TLS_DESC_CALL:
+    write32le(loc, insn(ctx.arg.is64 ? LD_D : LD_W, R_A0, R_A0,
+                        0)); // ld.[wd] $a0, $a0, %ie_pc_lo12
+    relocateNoSym(loc, R_LARCH_TLS_IE_PC_LO12, val);
+    break;
+  default:
+    llvm_unreachable("unsupported relocation for TLSDESC to IE");
+  }
+}
+
+// Convert TLSDESC GD/LD to LE.
+// The code sequence obtained in the normal or medium code model is as follows:
+//  * lu12i.w   $a0, %le_hi20(sym)      # le_hi20 != 0, otherwise NOP
+//  * ori       $a0, src, %le_lo12(sym) # le_hi20 != 0, src = $a0,
+//                                      # otherwise,    src = $zero
+// See the comment in tlsdescToIe for detailed information.
+void LoongArch::tlsdescToLe(uint8_t *loc, const Relocation &rel,
+                            uint64_t val) const {
+  assert(isInt<32>(val) &&
+         "val exceeds the range of medium code model in tlsdescToLe");
+
+  bool isUInt12 = isUInt<12>(val);
+  switch (rel.type) {
+  case R_LARCH_TLS_DESC_PC_HI20:
+  case R_LARCH_TLS_DESC_PC_LO12:
+  case R_LARCH_TLS_DESC_PCREL20_S2:
+    write32le(loc, insn(ANDI, R_ZERO, R_ZERO, 0)); // nop
+    break;
+  case R_LARCH_TLS_DESC_LD:
+    if (isUInt12)
+      write32le(loc, insn(ANDI, R_ZERO, R_ZERO, 0)); // nop
+    else
+      write32le(loc, insn(LU12I_W, R_A0, extractBits(val, 31, 12),
+                          0)); // lu12i.w $a0, %le_hi20
+    break;
+  case R_LARCH_TLS_DESC_CALL:
+    if (isUInt12)
+      write32le(loc, insn(ORI, R_A0, R_ZERO, val)); // ori $a0, $zero, %le_lo12
+    else
+      write32le(loc,
+                insn(ORI, R_A0, R_A0, lo12(val))); // ori $a0, $a0, %le_lo12
+    break;
+  default:
+    llvm_unreachable("unsupported relocation for TLSDESC to LE");
+  }
+}
+
+// During TLSDESC GD_TO_IE, the converted code sequence always includes an
+// instruction related to the Lo12 relocation (ld.[wd]). To obtain correct val
+// in `getRelocTargetVA`, expr of this instruction should be adjusted to
+// R_RELAX_TLS_GD_TO_IE_ABS, while expr of other instructions related to the
+// Hi20 relocation (pcalau12i) should be adjusted to
+// RE_LOONGARCH_RELAX_TLS_GD_TO_IE_PAGE_PC. Specifically, in the normal or
+// medium code model, the instruction with relocation R_LARCH_TLS_DESC_CALL is
+// the candidate of Lo12 relocation.
+RelExpr LoongArch::adjustTlsExpr(RelType type, RelExpr expr) const {
+  if (expr == R_RELAX_TLS_GD_TO_IE) {
+    if (type != R_LARCH_TLS_DESC_CALL)
+      return RE_LOONGARCH_RELAX_TLS_GD_TO_IE_PAGE_PC;
+    return R_RELAX_TLS_GD_TO_IE_ABS;
+  }
+  return expr;
+}
+
 void LoongArch::relocateAlloc(InputSectionBase &sec, uint8_t *buf) const {
   const unsigned bits = ctx.arg.is64 ? 64 : 32;
   uint64_t secAddr = sec.getOutputSection()->addr;
@@ -1074,7 +1180,7 @@ void LoongArch::relocateAlloc(InputSectionBase &sec, uint8_t *buf) const {
         //  * i+2 -- R_LARCH_TLS_IE64_PC_LO20
         //  * i+3 -- R_LARCH_TLS_IE64_PC_HI12
         isExtreme =
-            (i + 2 < size && relocs[i + 2].type == R_LARCH_TLS_IE64_PC_LO20);
+            i + 2 < size && relocs[i + 2].type == R_LARCH_TLS_IE64_PC_LO20;
       }
       if (isExtreme) {
         rel.expr = getRelExpr(rel.type, *rel.sym, loc);
@@ -1088,6 +1194,47 @@ void LoongArch::relocateAlloc(InputSectionBase &sec, uint8_t *buf) const {
         tlsIeToLe(loc, rel, val);
       }
       continue;
+    case RE_LOONGARCH_RELAX_TLS_GD_TO_IE_PAGE_PC:
+      if (rel.type == R_LARCH_TLS_DESC_PC_HI20) {
+        // LoongArch does not support TLSDESC GD/LD to LE/IE optimization in the
+        // extreme code model. In these cases, the relocs are as follows:
+        //
+        //  * i   -- R_LARCH_TLS_DESC_PC_HI20
+        //  * i+1 -- R_LARCH_TLS_DESC_PC_LO12
+        //  * i+2 -- R_LARCH_TLS_DESC64_PC_LO20
+        //  * i+3 -- R_LARCH_TLS_DESC64_PC_HI12
+        isExtreme =
+            i + 2 < size && relocs[i + 2].type == R_LARCH_TLS_DESC64_PC_LO20;
+      }
+      [[fallthrough]];
+    case R_RELAX_TLS_GD_TO_IE_ABS:
+      if (isExtreme) {
+        if (rel.type == R_LARCH_TLS_DESC_CALL)
+          continue;
+        rel.expr = getRelExpr(rel.type, *rel.sym, loc);
+        val = SignExtend64(sec.getRelocTargetVA(ctx, rel, secAddr + rel.offset),
+                           bits);
+        relocateNoSym(loc, rel.type, val);
+      } else {
+        tlsdescToIe(loc, rel, val);
+      }
+      continue;
+    case R_RELAX_TLS_GD_TO_LE:
+      if (rel.type == R_LARCH_TLS_DESC_PC_HI20) {
+        isExtreme =
+            i + 2 < size && relocs[i + 2].type == R_LARCH_TLS_DESC64_PC_LO20;
+      }
+      if (isExtreme) {
+        if (rel.type == R_LARCH_TLS_DESC_CALL)
+          continue;
+        rel.expr = getRelExpr(rel.type, *rel.sym, loc);
+        val = SignExtend64(sec.getRelocTargetVA(ctx, rel, secAddr + rel.offset),
+                           bits);
+        relocateNoSym(loc, rel.type, val);
+      } else {
+        tlsdescToLe(loc, rel, val);
+      }
+      continue;
     default:
       break;
     }

lld/ELF/InputSection.cpp

@@ -830,6 +830,7 @@ uint64_t InputSectionBase::getRelocTargetVA(Ctx &ctx, const Relocation &r,
   case R_GOTPLT_PC:
     return r.sym->getGotPltVA(ctx) + a - p;
   case RE_LOONGARCH_GOT_PAGE_PC:
+  case RE_LOONGARCH_RELAX_TLS_GD_TO_IE_PAGE_PC:
     if (r.sym->hasFlag(NEEDS_TLSGD))
       return getLoongArchPageDelta(ctx.in.got->getGlobalDynAddr(*r.sym) + a, p,
                                    r.type);

lld/ELF/Relocations.cpp

@@ -1340,22 +1340,10 @@ unsigned RelocationScanner::handleTlsRelocation(RelExpr expr, RelType type,
   if (ctx.arg.emachine == EM_MIPS)
     return handleMipsTlsRelocation(ctx, type, sym, *sec, offset, addend, expr);
 
-  // LoongArch does not yet implement transition from TLSDESC to LE/IE, so
-  // generate TLSDESC dynamic relocation for the dynamic linker to handle.
-  if (ctx.arg.emachine == EM_LOONGARCH &&
-      oneof<RE_LOONGARCH_TLSDESC_PAGE_PC, R_TLSDESC, R_TLSDESC_PC,
-            R_TLSDESC_CALL>(expr)) {
-    if (expr != R_TLSDESC_CALL) {
-      sym.setFlags(NEEDS_TLSDESC);
-      sec->addReloc({expr, type, offset, addend, &sym});
-    }
-    return 1;
-  }
-
   bool isRISCV = ctx.arg.emachine == EM_RISCV;
 
   if (oneof<RE_AARCH64_TLSDESC_PAGE, R_TLSDESC, R_TLSDESC_CALL, R_TLSDESC_PC,
-            R_TLSDESC_GOTPLT>(expr) &&
+            R_TLSDESC_GOTPLT, RE_LOONGARCH_TLSDESC_PAGE_PC>(expr) &&
       ctx.arg.shared) {
     // R_RISCV_TLSDESC_{LOAD_LO12,ADD_LO12_I,CALL} reference a label. Do not
     // set NEEDS_TLSDESC on the label.
@@ -1369,10 +1357,14 @@ unsigned RelocationScanner::handleTlsRelocation(RelExpr expr, RelType type,
     return 1;
   }
 
-  // LoongArch supports IE to LE optimization in non-extreme code model.
+  // LoongArch supports IE to LE, DESC GD/LD to IE/LE optimizations in
+  // non-extreme code model.
   bool execOptimizeInLoongArch =
       ctx.arg.emachine == EM_LOONGARCH &&
-      (type == R_LARCH_TLS_IE_PC_HI20 || type == R_LARCH_TLS_IE_PC_LO12);
+      (type == R_LARCH_TLS_IE_PC_HI20 || type == R_LARCH_TLS_IE_PC_LO12 ||
+       type == R_LARCH_TLS_DESC_PC_HI20 || type == R_LARCH_TLS_DESC_PC_LO12 ||
+       type == R_LARCH_TLS_DESC_LD || type == R_LARCH_TLS_DESC_CALL ||
+       type == R_LARCH_TLS_DESC_PCREL20_S2);
 
   // ARM, Hexagon, LoongArch and RISC-V do not support GD/LD to IE/LE
   // optimizations.
@@ -1431,9 +1423,23 @@ unsigned RelocationScanner::handleTlsRelocation(RelExpr expr, RelType type,
     return 1;
   }
 
+  // LoongArch does not support transition from TLSDESC to LE/IE in the extreme
+  // code model, in which NEEDS_TLSDESC should set, rather than NEEDS_TLSGD. So
+  // we check independently.
+  if (ctx.arg.emachine == EM_LOONGARCH &&
+      oneof<RE_LOONGARCH_TLSDESC_PAGE_PC, R_TLSDESC, R_TLSDESC_PC,
+            R_TLSDESC_CALL>(expr) &&
+      !execOptimize) {
+    if (expr != R_TLSDESC_CALL) {
+      sym.setFlags(NEEDS_TLSDESC);
+      sec->addReloc({expr, type, offset, addend, &sym});
+    }
+    return 1;
+  }
+
   if (oneof<RE_AARCH64_TLSDESC_PAGE, R_TLSDESC, R_TLSDESC_CALL, R_TLSDESC_PC,
             R_TLSDESC_GOTPLT, R_TLSGD_GOT, R_TLSGD_GOTPLT, R_TLSGD_PC,
-            RE_LOONGARCH_TLSGD_PAGE_PC>(expr)) {
+            RE_LOONGARCH_TLSGD_PAGE_PC, RE_LOONGARCH_TLSDESC_PAGE_PC>(expr)) {
     if (!execOptimize) {
       sym.setFlags(NEEDS_TLSGD);
       sec->addReloc({expr, type, offset, addend, &sym});

lld/ELF/Relocations.h

@@ -131,6 +131,7 @@ enum RelExpr {
   RE_LOONGARCH_GOT_PAGE_PC,
   RE_LOONGARCH_TLSGD_PAGE_PC,
   RE_LOONGARCH_TLSDESC_PAGE_PC,
+  RE_LOONGARCH_RELAX_TLS_GD_TO_IE_PAGE_PC,
 };
 
 // Architecture-neutral representation of relocation.