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

@@ -2798,6 +2798,12 @@ class TargetLoweringBase {
                                      Type *Ty, unsigned AddrSpace,
                                      Instruction *I = nullptr) const;
 
+  /// Returns true if the targets addressing mode can target thread local
+  /// storage (TLS).
+  virtual bool addressingModeSupportsTLS(const GlobalValue &) const {
+    return false;
+  }
+
   /// Return the prefered common base offset.
   virtual int64_t getPreferredLargeGEPBaseOffset(int64_t MinOffset,
                                                  int64_t MaxOffset) const {

llvm/lib/CodeGen/CodeGenPrepare.cpp

@@ -5082,6 +5082,15 @@ bool AddressingModeMatcher::matchOperationAddr(User *AddrInst, unsigned Opcode,
     }
     return true;
   }
+  case Instruction::Call:
+    if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(AddrInst)) {
+      if (II->getIntrinsicID() == Intrinsic::threadlocal_address) {
+        GlobalValue &GV = cast<GlobalValue>(*II->getArgOperand(0));
+        if (TLI.addressingModeSupportsTLS(GV))
+          return matchAddr(AddrInst->getOperand(0), Depth);
+      }
+    }
+    break;
   }
   return false;
 }
@@ -5620,11 +5629,16 @@ bool CodeGenPrepare::optimizeMemoryInst(Instruction *MemoryInst, Value *Addr,
         return Modified;
     }
 
-    if (AddrMode.BaseGV) {
+    GlobalValue *BaseGV = AddrMode.BaseGV;
+    if (BaseGV != nullptr) {
       if (ResultPtr)
         return Modified;
 
-      ResultPtr = AddrMode.BaseGV;
+      if (BaseGV->isThreadLocal()) {
+        ResultPtr = Builder.CreateThreadLocalAddress(BaseGV);
+      } else {
+        ResultPtr = BaseGV;
+      }
     }
 
     // If the real base value actually came from an inttoptr, then the matcher

llvm/lib/Target/X86/X86ISelLowering.cpp

@@ -18928,6 +18928,30 @@ X86TargetLowering::LowerGlobalTLSAddress(SDValue Op, SelectionDAG &DAG) const {
   llvm_unreachable("TLS not implemented for this target.");
 }
 
+bool X86TargetLowering::addressingModeSupportsTLS(const GlobalValue &GV) const {
+  if (Subtarget.is64Bit() && Subtarget.isTargetELF()) {
+    const TargetMachine &TM = getTargetMachine();
+    TLSModel::Model Model = TM.getTLSModel(&GV);
+    switch (Model) {
+    case TLSModel::LocalExec:
+    case TLSModel::InitialExec:
+      // We can include the %fs segment register in addressing modes.
+      return true;
+    case TLSModel::LocalDynamic:
+    case TLSModel::GeneralDynamic:
+      // These models do not result in %fs relative addresses unless
+      // TLS descriptior are used.
+      //
+      // Even in the case of TLS descriptors we currently have no way to model
+      // the difference between %fs access and the computations needed for the
+      // offset and returning `true` for TLS-desc currently duplicates both
+      // which is detrimental :-/
+      return false;
+    }
+  }
+  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

@@ -1323,6 +1323,8 @@ namespace llvm {
                                Type *Ty, unsigned AS,
                                Instruction *I = nullptr) const override;
 
+    bool addressingModeSupportsTLS(const GlobalValue &GV) const override;
+
     /// Return true if the specified immediate is legal
     /// icmp immediate, that is the target has icmp instructions which can
     /// compare a register against the immediate without having to materialize

llvm/test/CodeGen/X86/codegen-prepare-addrmode-tls.ll

@@ -13,65 +13,47 @@ declare nonnull ptr @llvm.threadlocal.address.p0(ptr nonnull)
 define i32 @func_local_tls(i32 %arg0, i32 %arg1) nounwind {
 ; NOPIC-LABEL: func_local_tls:
 ; NOPIC:       # %bb.0: # %entry
-; NOPIC-NEXT:    pushq %r14
 ; NOPIC-NEXT:    pushq %rbx
-; NOPIC-NEXT:    pushq %rax
 ; NOPIC-NEXT:    movl %fs:foo_local@TPOFF, %ebx
 ; NOPIC-NEXT:    testl %edi, %edi
 ; NOPIC-NEXT:    movl %ebx, %eax
 ; NOPIC-NEXT:    jne .LBB0_2
 ; NOPIC-NEXT:  # %bb.1: # %if.then
-; NOPIC-NEXT:    movq %fs:0, %rax
-; NOPIC-NEXT:    leaq foo_local@TPOFF(%rax), %r14
 ; NOPIC-NEXT:    callq effect@PLT
-; NOPIC-NEXT:    movl (%r14), %eax
+; NOPIC-NEXT:    movl %fs:foo_local@TPOFF, %eax
 ; NOPIC-NEXT:  .LBB0_2: # %if.end
 ; NOPIC-NEXT:    addl %ebx, %eax
-; NOPIC-NEXT:    addq $8, %rsp
 ; NOPIC-NEXT:    popq %rbx
-; NOPIC-NEXT:    popq %r14
 ; NOPIC-NEXT:    retq
 ;
 ; PIC-LABEL: func_local_tls:
 ; PIC:       # %bb.0: # %entry
-; PIC-NEXT:    pushq %r14
 ; PIC-NEXT:    pushq %rbx
-; PIC-NEXT:    pushq %rax
 ; PIC-NEXT:    movl %fs:.Lfoo_local$local@TPOFF, %ebx
 ; PIC-NEXT:    testl %edi, %edi
 ; PIC-NEXT:    movl %ebx, %eax
 ; PIC-NEXT:    jne .LBB0_2
 ; PIC-NEXT:  # %bb.1: # %if.then
-; PIC-NEXT:    movq %fs:0, %rax
-; PIC-NEXT:    leaq .Lfoo_local$local@TPOFF(%rax), %r14
 ; PIC-NEXT:    callq effect@PLT
-; PIC-NEXT:    movl (%r14), %eax
+; PIC-NEXT:    movl %fs:.Lfoo_local$local@TPOFF, %eax
 ; PIC-NEXT:  .LBB0_2: # %if.end
 ; PIC-NEXT:    addl %ebx, %eax
-; PIC-NEXT:    addq $8, %rsp
 ; PIC-NEXT:    popq %rbx
-; PIC-NEXT:    popq %r14
 ; PIC-NEXT:    retq
 ;
 ; TLSDESC-LABEL: func_local_tls:
 ; TLSDESC:       # %bb.0: # %entry
-; TLSDESC-NEXT:    pushq %r14
 ; TLSDESC-NEXT:    pushq %rbx
-; TLSDESC-NEXT:    pushq %rax
 ; TLSDESC-NEXT:    movl %fs:.Lfoo_local$local@TPOFF, %ebx
 ; TLSDESC-NEXT:    testl %edi, %edi
 ; TLSDESC-NEXT:    movl %ebx, %eax
 ; TLSDESC-NEXT:    jne .LBB0_2
 ; TLSDESC-NEXT:  # %bb.1: # %if.then
-; TLSDESC-NEXT:    movq %fs:0, %rax
-; TLSDESC-NEXT:    leaq .Lfoo_local$local@TPOFF(%rax), %r14
 ; TLSDESC-NEXT:    callq effect@PLT
-; TLSDESC-NEXT:    movl (%r14), %eax
+; TLSDESC-NEXT:    movl %fs:.Lfoo_local$local@TPOFF, %eax
 ; TLSDESC-NEXT:  .LBB0_2: # %if.end
 ; TLSDESC-NEXT:    addl %ebx, %eax
-; TLSDESC-NEXT:    addq $8, %rsp
 ; TLSDESC-NEXT:    popq %rbx
-; TLSDESC-NEXT:    popq %r14
 ; TLSDESC-NEXT:    retq
 entry:
   %addr = tail call ptr @llvm.threadlocal.address.p0(ptr @foo_local)
@@ -106,9 +88,8 @@ define i32 @func_nonlocal_tls(i32 %arg0, i32 %arg1) nounwind {
 ; NOPIC-NEXT:    movl %ebx, %eax
 ; NOPIC-NEXT:    jne .LBB1_2
 ; NOPIC-NEXT:  # %bb.1: # %if.then
-; NOPIC-NEXT:    addq %fs:0, %r14
 ; NOPIC-NEXT:    callq effect@PLT
-; NOPIC-NEXT:    movl (%r14), %eax
+; NOPIC-NEXT:    movl %fs:(%r14), %eax
 ; NOPIC-NEXT:  .LBB1_2: # %if.end
 ; NOPIC-NEXT:    addl %ebx, %eax
 ; NOPIC-NEXT:    addq $8, %rsp

llvm/test/Transforms/CodeGenPrepare/X86/sink-addrmode-tls.ll

@@ -0,0 +1,83 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4
+; RUN: opt -S -passes='require<profile-summary>,function(codegenprepare)' %s | FileCheck %s
+
+target triple = "x86_64--linux-gnu"
+
+@foo = dso_local thread_local(localexec) global i32 0, align 4
+
+declare void @effect()
+declare nonnull ptr @llvm.threadlocal.address.p0(ptr nonnull)
+
+define i32 @func0(i32 %arg) {
+; CHECK-LABEL: define i32 @func0(
+; CHECK-SAME: i32 [[ARG:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[ADDR:%.*]] = tail call ptr @llvm.threadlocal.address.p0(ptr @foo)
+; CHECK-NEXT:    [[LOAD0:%.*]] = load i32, ptr [[ADDR]], align 4
+; CHECK-NEXT:    [[COND:%.*]] = icmp eq i32 [[ARG]], 0
+; CHECK-NEXT:    br i1 [[COND]], label [[IF_THEN:%.*]], label [[IF_END:%.*]]
+; CHECK:       if.then:
+; CHECK-NEXT:    tail call void @effect()
+; CHECK-NEXT:    [[TMP0:%.*]] = call align 4 ptr @llvm.threadlocal.address.p0(ptr align 4 @foo)
+; CHECK-NEXT:    [[LOAD1:%.*]] = load i32, ptr [[TMP0]], align 4
+; CHECK-NEXT:    br label [[IF_END]]
+; CHECK:       if.end:
+; CHECK-NEXT:    [[PHI:%.*]] = phi i32 [ [[LOAD1]], [[IF_THEN]] ], [ [[LOAD0]], [[ENTRY:%.*]] ]
+; CHECK-NEXT:    [[RET:%.*]] = add i32 [[PHI]], [[LOAD0]]
+; CHECK-NEXT:    ret i32 [[RET]]
+;
+entry:
+  %addr = tail call ptr @llvm.threadlocal.address.p0(ptr @foo)
+  %load0 = load i32, ptr %addr, align 4
+  %cond = icmp eq i32 %arg, 0
+  br i1 %cond, label %if.then, label %if.end
+
+if.then:
+  tail call void @effect()
+  %load1 = load i32, ptr %addr, align 4
+  br label %if.end
+
+if.end:
+  %phi = phi i32 [ %load1, %if.then ], [ %load0, %entry ]
+  %ret = add i32 %phi, %load0
+  ret i32 %ret
+}
+
+define i32 @func1(i32 %arg0, i32 %arg1) {
+; CHECK-LABEL: define i32 @func1(
+; CHECK-SAME: i32 [[ARG0:%.*]], i32 [[ARG1:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[ADDR:%.*]] = tail call ptr @llvm.threadlocal.address.p0(ptr @foo)
+; CHECK-NEXT:    [[LOAD0:%.*]] = load i32, ptr [[ADDR]], align 4
+; CHECK-NEXT:    [[COND:%.*]] = icmp eq i32 [[ARG0]], 0
+; CHECK-NEXT:    br i1 [[COND]], label [[IF_THEN:%.*]], label [[IF_END:%.*]]
+; CHECK:       if.then:
+; CHECK-NEXT:    tail call void @effect()
+; CHECK-NEXT:    [[X:%.*]] = add i32 [[ARG1]], 42
+; CHECK-NEXT:    [[ADDR1:%.*]] = getelementptr inbounds i32, ptr [[ADDR]], i32 [[X]]
+; CHECK-NEXT:    [[TMP0:%.*]] = call align 4 ptr @llvm.threadlocal.address.p0(ptr align 4 @foo)
+; CHECK-NEXT:    [[LOAD1:%.*]] = load i32, ptr [[TMP0]], align 4
+; CHECK-NEXT:    br label [[IF_END]]
+; CHECK:       if.end:
+; CHECK-NEXT:    [[PHI:%.*]] = phi i32 [ [[LOAD1]], [[IF_THEN]] ], [ [[LOAD0]], [[ENTRY:%.*]] ]
+; CHECK-NEXT:    [[RET:%.*]] = add i32 [[PHI]], [[LOAD0]]
+; CHECK-NEXT:    ret i32 [[RET]]
+;
+entry:
+  %addr = tail call ptr @llvm.threadlocal.address.p0(ptr @foo)
+  %load0 = load i32, ptr %addr, align 4
+  %cond = icmp eq i32 %arg0, 0
+  br i1 %cond, label %if.then, label %if.end
+
+if.then:
+  tail call void @effect()
+  %x = add i32 %arg1, 42
+  %addr1 = getelementptr inbounds i32, ptr %addr, i32 %x
+  %load1 = load i32, ptr %addr, align 4
+  br label %if.end
+
+if.end:
+  %phi = phi i32 [ %load1, %if.then ], [ %load0, %entry ]
+  %ret = add i32 %phi, %load0
+  ret i32 %ret
+}