TLS loads opimization (hoist)

Reviewed By: craig.topper
Differential Revision: https://reviews.llvm.org/D120000

Author: xiangzh1

llvm/docs/LangRef.rst

@@ -2114,6 +2114,11 @@ example:
     function with a tail call. The prototype of a thunk should not be used for
     optimization purposes. The caller is expected to cast the thunk prototype to
     match the thunk target prototype.
+
+``"tls-load-hoist"``
+    This attribute indicates that the function will try to reduce redundant
+    tls address caculation by hoisting tls variable.
+
 ``uwtable[(sync|async)]``
     This attribute indicates that the ABI being targeted requires that
     an unwind table entry be produced for this function even if we can

llvm/include/llvm/CodeGen/MachinePassRegistry.def

@@ -47,6 +47,7 @@ FUNCTION_PASS("expand-reductions", ExpandReductionsPass, ())
 FUNCTION_PASS("expandvp", ExpandVectorPredicationPass, ())
 FUNCTION_PASS("lowerinvoke", LowerInvokePass, ())
 FUNCTION_PASS("scalarize-masked-mem-intrin", ScalarizeMaskedMemIntrinPass, ())
+FUNCTION_PASS("tlshoist", TLSVariableHoistPass, ())
 FUNCTION_PASS("verify", VerifierPass, ())
 #undef FUNCTION_PASS
 

llvm/include/llvm/InitializePasses.h

@@ -445,6 +445,7 @@ void initializeTargetLibraryInfoWrapperPassPass(PassRegistry&);
 void initializeTargetPassConfigPass(PassRegistry&);
 void initializeTargetTransformInfoWrapperPassPass(PassRegistry&);
 void initializeThreadSanitizerLegacyPassPass(PassRegistry&);
+void initializeTLSVariableHoistLegacyPassPass(PassRegistry &);
 void initializeTwoAddressInstructionPassPass(PassRegistry&);
 void initializeTypeBasedAAWrapperPassPass(PassRegistry&);
 void initializeTypePromotionPass(PassRegistry&);

llvm/include/llvm/LinkAllPasses.h

@@ -177,6 +177,7 @@ namespace {
       (void) llvm::createStripDeadDebugInfoPass();
       (void) llvm::createStripDeadPrototypesPass();
       (void) llvm::createTailCallEliminationPass();
+      (void)llvm::createTLSVariableHoistPass();
       (void) llvm::createJumpThreadingPass();
       (void) llvm::createDFAJumpThreadingPass();
       (void) llvm::createUnifyFunctionExitNodesPass();

llvm/include/llvm/Transforms/Scalar.h

@@ -427,6 +427,12 @@ extern char &InferAddressSpacesID;
 // "block_weights" metadata.
 FunctionPass *createLowerExpectIntrinsicPass();
 
+//===----------------------------------------------------------------------===//
+//
+// TLSVariableHoist - This pass reduce duplicated TLS address call.
+//
+FunctionPass *createTLSVariableHoistPass();
+
 //===----------------------------------------------------------------------===//
 //
 // LowerConstantIntrinsicss - Expand any remaining llvm.objectsize and

llvm/include/llvm/Transforms/Scalar/TLSVariableHoist.h

@@ -0,0 +1,131 @@
+//==- TLSVariableHoist.h ------ Remove Redundant TLS Loads -------*- C++ -*-==//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass identifies/eliminates Redundant TLS Loads if related option is set.
+// For example:
+// static __thread int x;
+// int g();
+// int f(int c) {
+//   int *px = &x;
+//   while (c--)
+//     *px += g();
+//   return *px;
+// }
+//
+// will generate Redundant TLS Loads by compiling it with
+// clang++ -fPIC -ftls-model=global-dynamic -O2 -S
+//
+// .LBB0_2:                                # %while.body
+//                                         # =>This Inner Loop Header: Depth=1
+//         callq   _Z1gv@PLT
+//         movl    %eax, %ebp
+//         leaq    _ZL1x@TLSLD(%rip), %rdi
+//         callq   __tls_get_addr@PLT
+//         addl    _ZL1x@DTPOFF(%rax), %ebp
+//         movl    %ebp, _ZL1x@DTPOFF(%rax)
+//         addl    $-1, %ebx
+//         jne     .LBB0_2
+//         jmp     .LBB0_3
+// .LBB0_4:                                # %entry.while.end_crit_edge
+//         leaq    _ZL1x@TLSLD(%rip), %rdi
+//         callq   __tls_get_addr@PLT
+//         movl    _ZL1x@DTPOFF(%rax), %ebp
+//
+// The Redundant TLS Loads will hurt the performance, especially in loops.
+// So we try to eliminate/move them if required by customers, let it be:
+//
+// # %bb.0:                                # %entry
+//         ...
+//         movl    %edi, %ebx
+//         leaq    _ZL1x@TLSLD(%rip), %rdi
+//         callq   __tls_get_addr@PLT
+//         leaq    _ZL1x@DTPOFF(%rax), %r14
+//         testl   %ebx, %ebx
+//         je      .LBB0_1
+// .LBB0_2:                                # %while.body
+//                                         # =>This Inner Loop Header: Depth=1
+//         callq   _Z1gv@PLT
+//         addl    (%r14), %eax
+//         movl    %eax, (%r14)
+//         addl    $-1, %ebx
+//         jne     .LBB0_2
+//         jmp     .LBB0_3
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TRANSFORMS_SCALAR_TLSVARIABLEHOIST_H
+#define LLVM_TRANSFORMS_SCALAR_TLSVARIABLEHOIST_H
+
+#include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/IR/PassManager.h"
+
+namespace llvm {
+
+class BasicBlock;
+class DominatorTree;
+class Function;
+class GlobalVariable;
+class Instruction;
+
+/// A private "module" namespace for types and utilities used by
+/// TLSVariableHoist. These are implementation details and should
+/// not be used by clients.
+namespace tlshoist {
+
+/// Keeps track of the user of a TLS variable and the operand index
+/// where the variable is used.
+struct TLSUser {
+  Instruction *Inst;
+  unsigned OpndIdx;
+
+  TLSUser(Instruction *Inst, unsigned Idx) : Inst(Inst), OpndIdx(Idx) {}
+};
+
+/// Keeps track of a TLS variable candidate and its users.
+struct TLSCandidate {
+  SmallVector<TLSUser, 8> Users;
+
+  /// Add the user to the use list and update the cost.
+  void addUser(Instruction *Inst, unsigned Idx) {
+    Users.push_back(TLSUser(Inst, Idx));
+  }
+};
+
+} // end namespace tlshoist
+
+class TLSVariableHoistPass : public PassInfoMixin<TLSVariableHoistPass> {
+public:
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+
+  // Glue for old PM.
+  bool runImpl(Function &F, DominatorTree &DT, LoopInfo &LI);
+
+private:
+  DominatorTree *DT;
+  LoopInfo *LI;
+
+  /// Keeps track of TLS variable candidates found in the function.
+  using TLSCandMapType = MapVector<GlobalVariable *, tlshoist::TLSCandidate>;
+  TLSCandMapType TLSCandMap;
+
+  void collectTLSCandidates(Function &Fn);
+  void collectTLSCandidate(Instruction *Inst);
+  Instruction *getNearestLoopDomInst(BasicBlock *BB, Loop *L);
+  Instruction *getDomInst(Instruction *I1, Instruction *I2);
+  BasicBlock::iterator findInsertPos(Function &Fn, GlobalVariable *GV,
+                                     BasicBlock *&PosBB);
+  Instruction *genBitCastInst(Function &Fn, GlobalVariable *GV);
+  bool tryReplaceTLSCandidates(Function &Fn);
+  bool tryReplaceTLSCandidate(Function &Fn, GlobalVariable *GV);
+};
+
+} // end namespace llvm
+
+#endif // LLVM_TRANSFORMS_SCALAR_TLSVARIABLEHOIST_H

llvm/lib/CodeGen/TargetPassConfig.cpp

@@ -922,6 +922,9 @@ void TargetPassConfig::addIRPasses() {
   // Allow disabling it for testing purposes.
   if (!DisableExpandReductions)
     addPass(createExpandReductionsPass());
+
+  if (getOptLevel() != CodeGenOpt::None)
+    addPass(createTLSVariableHoistPass());
 }
 
 /// Turn exception handling constructs into something the code generators can

llvm/lib/Passes/PassBuilder.cpp

@@ -212,6 +212,7 @@
 #include "llvm/Transforms/Scalar/SpeculativeExecution.h"
 #include "llvm/Transforms/Scalar/StraightLineStrengthReduce.h"
 #include "llvm/Transforms/Scalar/StructurizeCFG.h"
+#include "llvm/Transforms/Scalar/TLSVariableHoist.h"
 #include "llvm/Transforms/Scalar/TailRecursionElimination.h"
 #include "llvm/Transforms/Scalar/WarnMissedTransforms.h"
 #include "llvm/Transforms/Utils/AddDiscriminators.h"

llvm/lib/Passes/PassRegistry.def

@@ -361,6 +361,7 @@ FUNCTION_PASS("verify<safepoint-ir>", SafepointIRVerifierPass())
 FUNCTION_PASS("verify<scalar-evolution>", ScalarEvolutionVerifierPass())
 FUNCTION_PASS("view-cfg", CFGViewerPass())
 FUNCTION_PASS("view-cfg-only", CFGOnlyViewerPass())
+FUNCTION_PASS("tlshoist", TLSVariableHoistPass())
 FUNCTION_PASS("transform-warning", WarnMissedTransformationsPass())
 FUNCTION_PASS("tsan", ThreadSanitizerPass())
 FUNCTION_PASS("memprof", MemProfilerPass())

llvm/lib/Transforms/Scalar/CMakeLists.txt

@@ -77,6 +77,7 @@ add_llvm_component_library(LLVMScalarOpts
   StraightLineStrengthReduce.cpp
   StructurizeCFG.cpp
   TailRecursionElimination.cpp
+  TLSVariableHoist.cpp
   WarnMissedTransforms.cpp
 
   ADDITIONAL_HEADER_DIRS

llvm/lib/Transforms/Scalar/Scalar.cpp

@@ -101,6 +101,7 @@ void llvm::initializeScalarOpts(PassRegistry &Registry) {
   initializeSimpleLoopUnswitchLegacyPassPass(Registry);
   initializeSinkingLegacyPassPass(Registry);
   initializeTailCallElimPass(Registry);
+  initializeTLSVariableHoistLegacyPassPass(Registry);
   initializeSeparateConstOffsetFromGEPLegacyPassPass(Registry);
   initializeSpeculativeExecutionLegacyPassPass(Registry);
   initializeStraightLineStrengthReduceLegacyPassPass(Registry);