CodeGenPrepare: Remove threadlocal_address intrinsic when cheap to recompute.

The `threadlocal_address` intrinsic is currently ignored/removed for
instruction selection by the `SelectionDAGBuilder` (see also
https://reviews.llvm.org/D125291 ).

However being an Instruction means `SelectionDAG` will assign a register
to it and share the value across basic blocks. This sharing is
suboptimal in the "LocalExec" TLS model on x86 where it is cheaper to
just recompute the address. We saw a 0.5% regression in a codebase with
a lot of TLS usage (HHVM).

This introduces a new `cheapToRecomputeTLSAddress` target lowering
callback and removes the `threadlocal_address` intrinsic in
`CodeGenPrepare` to restore the efficient behavior from before the
introduction of the `threadlocal_address` intrinsic.

This fixes #87437

Author: MatzeB

llvm/include/llvm/CodeGen/TargetLowering.h

@@ -3393,6 +3393,13 @@ class TargetLoweringBase {
     return nullptr;
   }
 
+  /// Returns true if thread local storage (TLS) addresses are so cheap to
+  /// re-compute that it is not worth keeping them in a register between basic
+  /// blocks.
+  virtual bool cheapToRecomputeTLSAddress(const GlobalVariable &) const {
+    return false;
+  }
+
   //===--------------------------------------------------------------------===//
   // Runtime Library hooks
   //

llvm/lib/CodeGen/CodeGenPrepare.cpp

@@ -2528,6 +2528,18 @@ bool CodeGenPrepare::optimizeCallInst(CallInst *CI, ModifyDT &ModifiedDT) {
       return optimizeGatherScatterInst(II, II->getArgOperand(0));
     case Intrinsic::masked_scatter:
       return optimizeGatherScatterInst(II, II->getArgOperand(1));
+    case Intrinsic::threadlocal_address:
+      // SelectionDAGBuilder currently skips this intrinsic anyway; but removing
+      // it earlier means the addresses will not be kept in registers accross
+      // basic blocks but recomputed. This is preferable on architectures where
+      // TLS is part of normal addressing modes.
+      GlobalVariable &GV = cast<GlobalVariable>(*II->getArgOperand(0));
+      if (TLI->cheapToRecomputeTLSAddress(GV)) {
+        replaceAllUsesWith(II, &GV, FreshBBs, IsHugeFunc);
+        II->eraseFromParent();
+        return true;
+      }
+      break;
     }
 
     SmallVector<Value *, 2> PtrOps;

llvm/lib/Target/X86/X86ISelLowering.cpp

@@ -18913,6 +18913,16 @@ X86TargetLowering::LowerGlobalTLSAddress(SDValue Op, SelectionDAG &DAG) const {
   llvm_unreachable("TLS not implemented for this target.");
 }
 
+bool X86TargetLowering::cheapToRecomputeTLSAddress(
+    const GlobalVariable &GV) const {
+  if (Subtarget.is64Bit() && Subtarget.isTargetELF()) {
+    const TargetMachine &TM = getTargetMachine();
+    TLSModel::Model Model = TM.getTLSModel(&GV);
+    return Model == TLSModel::LocalExec;
+  }
+  return false;
+}
+
 /// Lower SRA_PARTS and friends, which return two i32 values
 /// and take a 2 x i32 value to shift plus a shift amount.
 /// TODO: Can this be moved to general expansion code?

llvm/lib/Target/X86/X86ISelLowering.h

@@ -1661,6 +1661,8 @@ namespace llvm {
     SDValue LowerGlobalTLSAddress(SDValue Op, SelectionDAG &DAG) const;
     SDValue LowerExternalSymbol(SDValue Op, SelectionDAG &DAG) const;
 
+    bool cheapToRecomputeTLSAddress(const GlobalVariable &GV) const override;
+
     /// Creates target global address or external symbol nodes for calls or
     /// other uses.
     SDValue LowerGlobalOrExternal(SDValue Op, SelectionDAG &DAG,

llvm/test/CodeGen/X86/tls-multi-use.ll

@@ -11,31 +11,20 @@ declare nonnull ptr @llvm.threadlocal.address.p0(ptr nonnull)
 define i32 @tls_multi_use(i32 %arg) {
 ; CHECK-LABEL: tls_multi_use:
 ; CHECK:       # %bb.0: # %entry
-; CHECK-NEXT:    pushq %r14
-; CHECK-NEXT:    .cfi_def_cfa_offset 16
 ; CHECK-NEXT:    pushq %rbx
-; CHECK-NEXT:    .cfi_def_cfa_offset 24
-; CHECK-NEXT:    pushq %rax
-; CHECK-NEXT:    .cfi_def_cfa_offset 32
-; CHECK-NEXT:    .cfi_offset %rbx, -24
-; CHECK-NEXT:    .cfi_offset %r14, -16
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    .cfi_offset %rbx, -16
 ; CHECK-NEXT:    movl %fs:foo@TPOFF, %ebx
 ; CHECK-NEXT:    testl %edi, %edi
 ; CHECK-NEXT:    movl %ebx, %eax
 ; CHECK-NEXT:    jne .LBB0_2
 ; CHECK-NEXT:  # %bb.1: # %if.then
-; CHECK-NEXT:    movq %fs:0, %rax
-; CHECK-NEXT:    leaq foo@TPOFF(%rax), %r14
 ; CHECK-NEXT:    callq rand@PLT
-; CHECK-NEXT:    movl (%r14), %eax
+; CHECK-NEXT:    movl %fs:foo@TPOFF, %eax
 ; CHECK-NEXT:  .LBB0_2: # %if.end
 ; CHECK-NEXT:    addl %eax, %ebx
 ; CHECK-NEXT:    movl %ebx, %eax
-; CHECK-NEXT:    addq $8, %rsp
-; CHECK-NEXT:    .cfi_def_cfa_offset 24
 ; CHECK-NEXT:    popq %rbx
-; CHECK-NEXT:    .cfi_def_cfa_offset 16
-; CHECK-NEXT:    popq %r14
 ; CHECK-NEXT:    .cfi_def_cfa_offset 8
 ; CHECK-NEXT:    retq
 ;