[K/N] Better removing of extra safepoints

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.backend.konan
 
 import llvm.*
+import org.jetbrains.kotlin.backend.common.LoggingContext
 import org.jetbrains.kotlin.backend.konan.descriptors.*
 import org.jetbrains.kotlin.backend.konan.ir.KonanIr
 import org.jetbrains.kotlin.library.SerializedMetadata
@@ -457,6 +458,10 @@ internal class Context(config: KonanConfig) : KonanBackendContext(config) {
     fun shouldUseDebugInfoFromNativeLibs() = shouldContainAnyDebugInfo() && config.useDebugInfoInNativeLibs
 
     fun shouldOptimize() = config.optimizationsEnabled
+    fun shouldInlineSafepoints() = when {
+        config.target.family.isAppleFamily -> config.target.architecture != Architecture.ARM32 // disable for watchos_arm32 and similar
+        else -> true
+    }
     fun ghaEnabled() = ::globalHierarchyAnalysisResult.isInitialized
     fun useLazyFileInitializers() = config.propertyLazyInitialization
 
@@ -530,11 +535,11 @@ private fun MemberScope.getContributedClassifier(name: String) =
 private fun MemberScope.getContributedFunctions(name: String) =
         this.getContributedFunctions(Name.identifier(name), NoLookupLocation.FROM_BUILTINS)
 
-internal class ContextLogger(val context: Context) {
+internal class ContextLogger(val context: LoggingContext) {
     operator fun String.unaryPlus() = context.log { this }
 }
 
-internal fun Context.logMultiple(messageBuilder: ContextLogger.() -> Unit) {
+internal fun LoggingContext.logMultiple(messageBuilder: ContextLogger.() -> Unit) {
     if (!inVerbosePhase) return
     with(ContextLogger(this)) { messageBuilder() }
 }

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

@@ -390,11 +390,12 @@ internal val produceOutputPhase = namedUnitPhase(
 
 internal val removeRedundantSafepointsPhase = makeKonanModuleOpPhase(
         name = "RemoveRedundantSafepoints",
-        description = "Leave only one safepoint in a basic block",
+        description = "Remove function prologue safepoints inlined to another function",
         op = { context, _ ->
-            if (context.config.target.architecture == Architecture.ARM32 && context.config.target.family.isAppleFamily) {
-                RemoveRedundantSafepointsPass(context as LoggingContext).runOnModule(context.llvmModule!!)
-            }
+            RemoveRedundantSafepointsPass(context).runOnModule(
+                    module = context.llvmModule!!,
+                    isSafepointInliningAllowed = context.shouldInlineSafepoints()
+            )
         }
 )
 

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

@@ -7,34 +7,50 @@ package org.jetbrains.kotlin.backend.konan.optimizations
 
 import llvm.*
 import org.jetbrains.kotlin.backend.common.LoggingContext
+import org.jetbrains.kotlin.backend.konan.llvm.*
 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
+import org.jetbrains.kotlin.backend.konan.logMultiple
 
 /**
  * Removes all Kotlin_mm_safePointFunctionPrologue from basic block except the first one.
- *
- * Currently, this pass is useful only for watchos_arm32, ios_arm32 targets because Kotlin_mm_* functions are marked
- * as noinline there.
+ * Also, if first basic block in function contains call to Kotlin_mm_safePointFunctionPrologue, all other calls would be removed.
+ * Also, calls, which are not removed are inlined (except arm32 apple targets)
  */
-class RemoveRedundantSafepointsPass(
+internal class RemoveRedundantSafepointsPass(
         private val loggingContext: LoggingContext
 ) {
     var totalPrologueSafepointsCount = 0
     var removedPrologueSafepointsCount = 0
 
-    fun runOnFunction(function: LLVMValueRef) {
+    fun runOnFunction(function: LLVMValueRef, isSafepointInliningAllowed: Boolean) {
+        val firstBlock = LLVMGetFirstBasicBlock(function) ?: return
+        val firstBlockHasSafepoint = getInstructions(firstBlock).any { isPrologueSafepointCallsite(it) }
         getBasicBlocks(function).forEach { bb ->
-            val unnecessaryPrologueSafepointCallsites = getInstructions(bb)
+            val removeFirst = firstBlockHasSafepoint && bb != firstBlock
+            val prologueSafepointCallsites = getInstructions(bb)
                     .filter { isPrologueSafepointCallsite(it) }
-                    .onEach { totalPrologueSafepointsCount += 1 }
-                    .drop(1)
                     .toList()
-            unnecessaryPrologueSafepointCallsites.forEach {
+            totalPrologueSafepointsCount += prologueSafepointCallsites.size
+            prologueSafepointCallsites.drop(if (removeFirst) 0 else 1).forEach {
                 LLVMInstructionEraseFromParent(it)
                 removedPrologueSafepointsCount += 1
             }
+            if (!removeFirst && isSafepointInliningAllowed) {
+                prologueSafepointCallsites
+                        .firstOrNull()
+                        ?.apply {
+                            if (LLVMIsDeclaration(LLVMGetCalledValue(this)) == 0) {
+                                if (LLVMInlineCall(this) == 0) {
+                                    loggingContext.logMultiple {
+                                        +"Failed to Inline safepoint to ${function.name}"
+                                        +llvm2string(function)
+                                    }
+                                }
+                            }
+                        }
+            }
         }
     }
 
@@ -42,18 +58,15 @@ class RemoveRedundantSafepointsPass(
             (LLVMIsACallInst(insn) != null || LLVMIsAInvokeInst(insn) != null)
                     && LLVMGetCalledValue(insn)?.name == prologueSafepointFunctionName
 
-    fun runOnModule(module: LLVMModuleRef) {
+    fun runOnModule(module: LLVMModuleRef, isSafepointInliningAllowed: Boolean) {
         totalPrologueSafepointsCount = 0
         removedPrologueSafepointsCount = 0
         getFunctions(module)
-                .filter { it.name?.startsWith("kfun:") == true }
                 .filterNot { LLVMIsDeclaration(it) == 1 }
-                .forEach(this::runOnFunction)
-        loggingContext.log {
-            """
-               Total prologue safepoints: $totalPrologueSafepointsCount
-               Removed prologue safepoints: $removedPrologueSafepointsCount
-            """.trimIndent()
+                .forEach { runOnFunction(it, isSafepointInliningAllowed) }
+        loggingContext.logMultiple {
+               +"Total prologue safepoints: $totalPrologueSafepointsCount"
+               +"Removed prologue safepoints: $removedPrologueSafepointsCount"
         }
     }
 

kotlin-native/libllvmext/src/main/cpp/CAPIExtensions.cpp

@@ -9,6 +9,7 @@
 #include <llvm/IR/Instructions.h>
 #include <llvm/IR/LegacyPassManager.h>
 #include <llvm/Transforms/ObjCARC.h>
+#include <llvm/Transforms/Utils/Cloning.h>
 
 using namespace llvm;
 
@@ -47,3 +48,8 @@ void LLVMKotlinInitializeTargets() {
 void LLVMSetNoTailCall(LLVMValueRef Call) {
     unwrap<CallInst>(Call)->setTailCallKind(CallInst::TCK_NoTail);
 }
+
+int LLVMInlineCall(LLVMValueRef call) {
+  InlineFunctionInfo IFI;
+  return InlineFunction(*unwrap<CallBase>(call), IFI).isSuccess();
+}

kotlin-native/libllvmext/src/main/include/CAPIExtensions.h

@@ -23,6 +23,8 @@ void LLVMKotlinInitializeTargets();
 
 void LLVMSetNoTailCall(LLVMValueRef Call);
 
+int LLVMInlineCall(LLVMValueRef call);
+
 # ifdef __cplusplus
 }
 # endif

kotlin-native/runtime/src/mm/cpp/Memory.cpp

@@ -586,7 +586,8 @@ extern "C" void CheckGlobalsAccessible() {
     // Always accessible
 }
 
-extern "C" RUNTIME_NOTHROW CODEGEN_INLINE_POLICY void Kotlin_mm_safePointFunctionPrologue() {
+// it would be inlined manually in RemoveRedundantSafepointsPass
+extern "C" RUNTIME_NOTHROW NO_INLINE void Kotlin_mm_safePointFunctionPrologue() {
     auto* threadData = mm::ThreadRegistry::Instance().CurrentThreadData();
     AssertThreadState(threadData, ThreadState::kRunnable);
     threadData->gc().SafePointFunctionPrologue();