[K/N] Custom bitcode optimization pass to avoid multiple loads of TLS

^KT-50075

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/CheckExternalCalls.kt

@@ -9,12 +9,6 @@ import kotlinx.cinterop.toCValues
 import llvm.*
 import org.jetbrains.kotlin.backend.konan.llvm.*
 
-private fun getBasicBlocks(function: LLVMValueRef) =
-        generateSequence(LLVMGetFirstBasicBlock(function)) { LLVMGetNextBasicBlock(it) }
-
-private fun getInstructions(function: LLVMBasicBlockRef) =
-        generateSequence(LLVMGetFirstInstruction(function)) { LLVMGetNextInstruction(it) }
-
 private fun LLVMValueRef.isFunctionCall() = LLVMIsACallInst(this) != null || LLVMIsAInvokeInst(this) != null
 
 private fun LLVMValueRef.isExternalFunction() = LLVMGetFirstBasicBlock(this) == null
@@ -172,18 +166,8 @@ private const val functionListSizeGlobal = "Kotlin_callsCheckerKnownFunctionsCou
 internal fun checkLlvmModuleExternalCalls(context: Context) {
     val staticData = context.llvm.staticData
 
-    val annotations = staticData.getGlobal("llvm.global.annotations")?.getInitializer()
 
-    val ignoredFunctions = annotations?.run {
+    val ignoredFunctions = (context.llvm.runtimeAnnotationMap["no_external_calls_check"] ?: emptyList())
-        getOperands(this).mapNotNull {
-            val annotationName = LLVMGetInitializer(LLVMGetOperand(LLVMGetOperand(it, 1), 0))?.getAsCString()
-            if (annotationName == "no_external_calls_check") {
-                LLVMGetOperand(LLVMGetOperand(it, 0), 0)!!.name
-            } else {
-                null
-            }
-        }.toSet()
-    } ?: emptySet()
 
     val goodFunctions = staticData.getGlobal("Kotlin_callsCheckerGoodFunctionNames")?.getInitializer()?.run {
         getOperands(this).map {
@@ -193,7 +177,7 @@ internal fun checkLlvmModuleExternalCalls(context: Context) {
 
     val checker = CallsChecker(context, goodFunctions)
     getFunctions(context.llvmModule!!)
-            .filter { !it.isExternalFunction() && it.name !in ignoredFunctions }
+            .filter { !it.isExternalFunction() && it !in ignoredFunctions }
             .forEach(checker::processFunction)
     // otherwise optimiser can inline it
     staticData.getGlobal(functionListGlobal)?.setExternallyInitialized(true);

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/OptimizationPipeline.kt

@@ -5,7 +5,7 @@ import kotlinx.cinterop.memScoped
 import kotlinx.cinterop.ptr
 import kotlinx.cinterop.value
 import llvm.*
-import org.jetbrains.kotlin.backend.konan.llvm.makeVisibilityHiddenLikeLlvmInternalizePass
+import org.jetbrains.kotlin.backend.konan.llvm.*
 import org.jetbrains.kotlin.konan.target.*
 
 private fun initializeLlvmGlobalPassRegistry() {
@@ -25,17 +25,6 @@ private fun initializeLlvmGlobalPassRegistry() {
     LLVMInitializeObjCARCOpts(passRegistry)
 }
 
-internal fun shouldRunLateBitcodePasses(context: Context): Boolean {
-    return context.coverage.enabled
-}
-
-internal fun runLateBitcodePasses(context: Context, llvmModule: LLVMModuleRef) {
-    val passManager = LLVMCreatePassManager()!!
-    LLVMKotlinAddTargetLibraryInfoWrapperPass(passManager, context.llvm.targetTriple)
-    context.coverage.addLateLlvmPasses(passManager)
-    LLVMRunPassManager(passManager, llvmModule)
-    LLVMDisposePassManager(passManager)
-}
 
 private class LlvmPipelineConfiguration(context: Context) {
 
@@ -183,9 +172,6 @@ internal fun runLlvmOptimizationPipeline(context: Context) {
         LLVMDisposeTargetMachine(targetMachine)
         LLVMDisposePassManager(modulePasses)
     }
-    if (shouldRunLateBitcodePasses(context)) {
-        runLateBitcodePasses(context, llvmModule)
-    }
 }
 
 internal fun RelocationModeFlags.currentRelocationMode(context: Context): RelocationModeFlags.Mode =

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/ToplevelPhases.kt

@@ -512,6 +512,16 @@ internal val bitcodePhase = NamedCompilerPhase(
                 cStubsPhase
 )
 
+private val bitcodePostprocessingPhase = NamedCompilerPhase(
+        name = "BitcodePostprocessing",
+        description = "Optimize and rewrite bitcode",
+        lower = checkExternalCallsPhase then
+                bitcodeOptimizationPhase then
+                coveragePhase then
+                optimizeTLSDataLoadsPhase then
+                rewriteExternalCallsCheckerGlobals
+)
+
 private val backendCodegen = namedUnitPhase(
         name = "Backend codegen",
         description = "Backend code generation",
@@ -527,9 +537,7 @@ private val backendCodegen = namedUnitPhase(
                 verifyBitcodePhase then
                 printBitcodePhase then
                 linkBitcodeDependenciesPhase then
-                checkExternalCallsPhase then
+                bitcodePostprocessingPhase then
-                bitcodeOptimizationPhase then
-                rewriteExternalCallsCheckerGlobals then
                 unitSink()
 )
 
@@ -578,10 +586,11 @@ internal fun PhaseConfig.konanPhasesConfig(config: KonanConfig) {
         disableUnless(buildAdditionalCacheInfoPhase, config.produce.isCache)
         disableUnless(exportInternalAbiPhase, config.produce.isCache)
         disableIf(backendCodegen, config.produce == CompilerOutputKind.LIBRARY)
-        disableUnless(bitcodeOptimizationPhase, config.produce.involvesLinkStage)
+        disableUnless(bitcodePostprocessingPhase, config.produce.involvesLinkStage)
         disableUnless(linkBitcodeDependenciesPhase, config.produce.involvesLinkStage)
-        disableUnless(checkExternalCallsPhase, config.produce.involvesLinkStage && getBoolean(KonanConfigKeys.CHECK_EXTERNAL_CALLS))
+        disableUnless(checkExternalCallsPhase, getBoolean(KonanConfigKeys.CHECK_EXTERNAL_CALLS))
-        disableUnless(rewriteExternalCallsCheckerGlobals, config.produce.involvesLinkStage && getBoolean(KonanConfigKeys.CHECK_EXTERNAL_CALLS))
+        disableUnless(rewriteExternalCallsCheckerGlobals, getBoolean(KonanConfigKeys.CHECK_EXTERNAL_CALLS))
+        disableUnless(optimizeTLSDataLoadsPhase, getBoolean(KonanConfigKeys.OPTIMIZATION))
         disableUnless(objectFilesPhase, config.produce.involvesLinkStage)
         disableUnless(linkerPhase, config.produce.involvesLinkStage)
         disableIf(testProcessorPhase, getNotNull(KonanConfigKeys.GENERATE_TEST_RUNNER) == TestRunnerKind.NONE)

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/BitcodePhases.kt

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.backend.common.phaser.PhaserState
 import org.jetbrains.kotlin.backend.common.phaser.namedUnitPhase
 import org.jetbrains.kotlin.backend.konan.*
 import org.jetbrains.kotlin.backend.konan.descriptors.GlobalHierarchyAnalysis
+import org.jetbrains.kotlin.backend.konan.llvm.coverage.runCoveragePass
 import org.jetbrains.kotlin.backend.konan.lower.InlineClassPropertyAccessorsLowering
 import org.jetbrains.kotlin.backend.konan.lower.RedundantCoercionsCleaner
 import org.jetbrains.kotlin.backend.konan.lower.ReturnsInsertionLowering
@@ -368,6 +369,18 @@ internal val bitcodeOptimizationPhase = makeKonanModuleOpPhase(
         op = { context, _ -> runLlvmOptimizationPipeline(context) }
 )
 
+internal val coveragePhase = makeKonanModuleOpPhase(
+        name = "Coverage",
+        description = "Produce coverage information",
+        op = { context, _ -> runCoveragePass(context) }
+)
+
+internal val optimizeTLSDataLoadsPhase = makeKonanModuleOpPhase(
+        name = "OptimizeTLSDataLoads",
+        description = "Optimize multiple loads of thread data",
+        op = { context, _ -> removeMultipleThreadDataLoads(context) }
+)
+
 internal val produceOutputPhase = namedUnitPhase(
         name = "ProduceOutput",
         description = "Produce output",

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/ContextUtils.kt

@@ -590,6 +590,18 @@ internal class Llvm(val context: Context, val llvmModule: LLVMModuleRef) : Runti
     val initializersGenerationState = InitializersGenerationState()
     val boxCacheGlobals = mutableMapOf<BoxCache, StaticData.Global>()
 
+    val runtimeAnnotationMap by lazy {
+        context.llvm.staticData.getGlobal("llvm.global.annotations")
+                ?.getInitializer()
+                ?.let { getOperands(it) }
+                ?.groupBy(
+                        { LLVMGetInitializer(LLVMGetOperand(LLVMGetOperand(it, 1), 0))?.getAsCString() ?: "" },
+                        { LLVMGetOperand(LLVMGetOperand(it, 0), 0)!! }
+                )
+                ?.filterKeys { it != "" }
+                ?: emptyMap()
+    }
+
 
     private object lazyRtFunction {
         operator fun provideDelegate(

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/LlvmUtils.kt

@@ -436,6 +436,13 @@ internal fun getGlobalAliases(module: LLVMModuleRef) =
 internal fun getFunctions(module: LLVMModuleRef) =
         generateSequence(LLVMGetFirstFunction(module), { LLVMGetNextFunction(it) })
 
+internal fun getBasicBlocks(function: LLVMValueRef) =
+        generateSequence(LLVMGetFirstBasicBlock(function)) { LLVMGetNextBasicBlock(it) }
+
+internal fun getInstructions(block: LLVMBasicBlockRef) =
+        generateSequence(LLVMGetFirstInstruction(block)) { LLVMGetNextInstruction(it) }
+
+
 internal fun getGlobals(module: LLVMModuleRef) =
         generateSequence(LLVMGetFirstGlobal(module), { LLVMGetNextGlobal(it) })
 

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/coverage/CoverageManager.kt

@@ -119,4 +119,13 @@ internal class CoverageManager(val context: Context) {
         } else {
             emptyList()
         }
+}
+
+internal fun runCoveragePass(context: Context) {
+    if (!context.coverage.enabled) return
+    val passManager = LLVMCreatePassManager()!!
+    LLVMKotlinAddTargetLibraryInfoWrapperPass(passManager, context.llvm.targetTriple)
+    context.coverage.addLateLlvmPasses(passManager)
+    LLVMRunPassManager(passManager, context.llvmModule)
+    LLVMDisposePassManager(passManager)
 }

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/optimizations/RemoveMultipleThreadDataLoads.kt

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.backend.konan.optimizations
+
+import llvm.*
+import org.jetbrains.kotlin.backend.konan.Context
+import org.jetbrains.kotlin.backend.konan.llvm.getBasicBlocks
+import org.jetbrains.kotlin.backend.konan.llvm.getFunctions
+import org.jetbrains.kotlin.backend.konan.llvm.getInstructions
+import org.jetbrains.kotlin.backend.konan.llvm.name
+
+private fun filterLoads(block: LLVMBasicBlockRef, variable: LLVMValueRef) = getInstructions(block)
+        .mapNotNull { LLVMIsALoadInst(it) }
+        .filter { inst ->
+            LLVMGetOperand(inst, 0)?.let { LLVMIsAGlobalVariable(it) } == variable
+        }
+
+private fun process(function: LLVMValueRef, currentThreadTLV: LLVMValueRef) {
+    val entry = LLVMGetEntryBasicBlock(function) ?: return
+    val load = filterLoads(entry, currentThreadTLV).firstOrNull() ?: return
+    getBasicBlocks(function)
+            .flatMap { filterLoads(it, currentThreadTLV) }
+            .filter { it != load }
+            .toList() // to force evaluating of all sequences above, because removing something during iteration is bad idea
+            .forEach {
+                LLVMReplaceAllUsesWith(it, load)
+                LLVMInstructionEraseFromParent(it)
+            }
+}
+
+internal fun removeMultipleThreadDataLoads(context: Context) {
+    val currentThreadTLV = context.llvm.runtimeAnnotationMap["current_thread_tlv"]?.singleOrNull() ?: return
+
+    getFunctions(context.llvmModule!!)
+            .filter { it.name?.startsWith("kfun:") == true }
+            .filterNot { LLVMIsDeclaration(it) == 1 }
+            .forEach { process(it, currentThreadTLV) }
+}

kotlin-native/runtime/src/mm/cpp/ThreadRegistry.hpp

@@ -55,7 +55,7 @@ private:
     ThreadRegistry();
     ~ThreadRegistry();
 
-    static THREAD_LOCAL_VARIABLE Node* currentThreadDataNode_;
+    static THREAD_LOCAL_VARIABLE Node* currentThreadDataNode_ __attribute__((annotate("current_thread_tlv")));
 
     SingleLockList<ThreadData, Mutex> list_;
 };