[K/N] Implement pass profiling in LLVMKotlinRunPasses

^KT-79381

Author: projedi

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

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.backend.konan.serialization.SerializedClassFields
 import org.jetbrains.kotlin.backend.konan.serialization.SerializedEagerInitializedFile
 import org.jetbrains.kotlin.backend.konan.serialization.SerializedInlineFunctionReference
 import org.jetbrains.kotlin.ir.declarations.*
+import org.jetbrains.kotlin.util.PerformanceManager
 
 internal class FileLowerState {
     private var functionReferenceCount = 0
@@ -59,6 +60,7 @@ internal class NativeGenerationState(
     val llvmModuleSpecification: LlvmModuleSpecification,
     val outputFiles: OutputFiles,
     val llvmModuleName: String,
+    override val performanceManager: PerformanceManager?,
 ) : BasicPhaseContext(config), BackendContextHolder, LlvmIrHolder, BitcodePostProcessingContext {
     val outputFile = outputFiles.mainFileName
 

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

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.backend.konan.driver.PhaseContext
 import org.jetbrains.kotlin.backend.konan.llvm.*
 import org.jetbrains.kotlin.config.nativeBinaryOptions.StackProtectorMode
 import org.jetbrains.kotlin.konan.target.*
+import org.jetbrains.kotlin.util.PerformanceManager
 import java.io.Closeable
 
 enum class LlvmOptimizationLevel(val value: Int) {
@@ -182,7 +183,8 @@ internal fun createLTOFinalPipelineConfig(
 
 abstract class LlvmOptimizationPipeline(
         private val config: LlvmPipelineConfig,
-        private val logger: LoggingContext? = null
+        private val performanceManager: PerformanceManager?,
+        private val logger: LoggingContext? = null,
 ) : Closeable {
     open fun executeCustomPreprocessing(config: LlvmPipelineConfig, module: LLVMModuleRef) {}
 
@@ -220,7 +222,7 @@ abstract class LlvmOptimizationPipeline(
         val passDescription = passes.joinToString(",")
         logger?.log {
             """
-                Running ${pipelineName} with the following parameters:
+                Running $pipelineName with the following parameters:
                 target_triple: ${config.targetTriple}
                 cpu_model: ${config.cpuModel}
                 cpu_features: ${config.cpuFeatures}
@@ -231,15 +233,21 @@ abstract class LlvmOptimizationPipeline(
             """.trimIndent()
         }
         if (passDescription.isEmpty()) return
-        val errorCode = LLVMKotlinRunPasses(
-                llvmModule,
-                passDescription,
-                targetMachine,
-                InlinerThreshold = config.inlineThreshold ?: -1,
-        )
+        val (errorCode, profile) = withLLVMPassesProfile(performanceManager != null, pipelineName) {
+            LLVMKotlinRunPasses(
+                    llvmModule,
+                    passDescription,
+                    targetMachine,
+                    InlinerThreshold = config.inlineThreshold ?: -1,
+                    Profile = it,
+            )
+        }
         require(errorCode == null) {
             LLVMGetErrorMessage(errorCode)!!.toKString()
         }
+        profile?.let {
+            performanceManager?.addLlvmPassesProfile(it)
+        }
         if (config.timePasses) {
             LLVMPrintAllTimersToStdOut()
             LLVMClearAllTimers()
@@ -272,9 +280,9 @@ abstract class LlvmOptimizationPipeline(
     }
 }
 
-class MandatoryOptimizationPipeline(config: LlvmPipelineConfig, logger: LoggingContext? = null) :
-        LlvmOptimizationPipeline(config, logger) {
-    override val pipelineName = "New PM Mandatory llvm optimizations"
+class MandatoryOptimizationPipeline(config: LlvmPipelineConfig, performanceManager: PerformanceManager?, logger: LoggingContext? = null) :
+        LlvmOptimizationPipeline(config, performanceManager, logger) {
+    override val pipelineName = "llvm-mandatory"
     override val passes = buildList {
         if (config.objCPasses) {
             // Lower ObjC ARC intrinsics (e.g. `@llvm.objc.clang.arc.use(...)`).
@@ -293,15 +301,15 @@ class MandatoryOptimizationPipeline(config: LlvmPipelineConfig, logger: LoggingC
     }
 }
 
-class ModuleOptimizationPipeline(config: LlvmPipelineConfig, logger: LoggingContext? = null) :
-        LlvmOptimizationPipeline(config, logger) {
-    override val pipelineName = "New PM Module LLVM optimizations"
+class ModuleOptimizationPipeline(config: LlvmPipelineConfig, performanceManager: PerformanceManager?, logger: LoggingContext? = null) :
+        LlvmOptimizationPipeline(config, performanceManager, logger) {
+    override val pipelineName = "llvm-default"
     override val passes = listOf(config.modulePasses ?: "default<$optimizationFlag>")
 }
 
-class LTOOptimizationPipeline(config: LlvmPipelineConfig, logger: LoggingContext? = null) :
-        LlvmOptimizationPipeline(config, logger) {
-    override val pipelineName = "New PM LTO LLVM optimizations"
+class LTOOptimizationPipeline(config: LlvmPipelineConfig, performanceManager: PerformanceManager?, logger: LoggingContext? = null) :
+        LlvmOptimizationPipeline(config, performanceManager, logger) {
+    override val pipelineName = "llvm-lto"
     override val passes =
             if (config.ltoPasses != null) listOf(config.ltoPasses)
             else buildList {
@@ -318,9 +326,9 @@ class LTOOptimizationPipeline(config: LlvmPipelineConfig, logger: LoggingContext
             }
 }
 
-class ThreadSanitizerPipeline(config: LlvmPipelineConfig, logger: LoggingContext? = null) :
-        LlvmOptimizationPipeline(config, logger) {
-    override val pipelineName = "New PM thread sanitizer"
+class ThreadSanitizerPipeline(config: LlvmPipelineConfig, performanceManager: PerformanceManager?, logger: LoggingContext? = null) :
+        LlvmOptimizationPipeline(config, performanceManager, logger) {
+    override val pipelineName = "llvm-tsan"
     override val passes = listOf("tsan-module,function(tsan)")
 
     override fun executeCustomPreprocessing(config: LlvmPipelineConfig, module: LLVMModuleRef) {

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

@@ -15,7 +15,13 @@ import org.jetbrains.kotlin.config.phaser.NamedCompilerPhase
 import org.jetbrains.kotlin.backend.konan.ConfigChecks
 import org.jetbrains.kotlin.backend.konan.KonanConfig
 import org.jetbrains.kotlin.cli.common.messages.MessageCollector
+import org.jetbrains.kotlin.cli.common.perfManager
 import org.jetbrains.kotlin.config.CommonConfigurationKeys
+import org.jetbrains.kotlin.util.PerformanceManager
+
+internal interface PerformanceManagerContext {
+    val performanceManager: PerformanceManager?
+}
 
 /**
  * Context is a set of resources that is shared between different phases. PhaseContext is a "minimal context",
@@ -35,7 +41,7 @@ import org.jetbrains.kotlin.config.CommonConfigurationKeys
  * * On the other hand, middle- and bitcode phases are hard to decouple due to the way the code was written many years ago.
  * It will take some time to rewrite it properly.
  */
-internal interface PhaseContext : LoggingContext, ConfigChecks, ErrorReportingContext, DisposableContext
+internal interface PhaseContext : LoggingContext, ConfigChecks, ErrorReportingContext, DisposableContext, PerformanceManagerContext
 
 internal open class BasicPhaseContext(
         override val config: KonanConfig,
@@ -48,6 +54,9 @@ internal open class BasicPhaseContext(
     override fun dispose() {
 
     }
+
+    override val performanceManager: PerformanceManager?
+        get() = config.configuration.perfManager
 }
 
 internal fun PhaseEngine.Companion.startTopLevel(config: KonanConfig, body: (PhaseEngine<PhaseContext>) -> Unit) {

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/driver/phases/Bitcode.kt

@@ -29,6 +29,7 @@ import org.jetbrains.kotlin.backend.konan.llvm.verifyModule
 import org.jetbrains.kotlin.backend.konan.optimizations.RemoveRedundantSafepointsPass
 import org.jetbrains.kotlin.backend.konan.optimizations.removeMultipleThreadDataLoads
 import org.jetbrains.kotlin.config.nativeBinaryOptions.SanitizerKind
+import org.jetbrains.kotlin.util.PerformanceManager
 import java.io.File
 import kotlin.sequences.forEach
 
@@ -68,15 +69,16 @@ internal val RewriteExternalCallsCheckerGlobals = createSimpleNamedCompilerPhase
 
 internal class OptimizationState(
         konanConfig: KonanConfig,
-        val llvmConfig: LlvmPipelineConfig
+        val llvmConfig: LlvmPipelineConfig,
+        override val performanceManager: PerformanceManager?,
 ) : BasicPhaseContext(konanConfig)
 
-internal fun optimizationPipelinePass(name: String, pipeline: (LlvmPipelineConfig, LoggingContext) -> LlvmOptimizationPipeline) =
+internal fun optimizationPipelinePass(name: String, pipeline: (LlvmPipelineConfig, PerformanceManager?, LoggingContext) -> LlvmOptimizationPipeline) =
         createSimpleNamedCompilerPhase<OptimizationState, LLVMModuleRef>(
                 name = name,
                 postactions = getDefaultLlvmModuleActions(),
         ) { context, module ->
-            pipeline(context.llvmConfig, context).use {
+            pipeline(context.llvmConfig, context.performanceManager, context).use {
                 it.execute(module)
             }
         }
@@ -165,7 +167,7 @@ internal fun <T : BitcodePostProcessingContext> PhaseEngine<T>.runBitcodePostPro
             closedWorld = context.config.isFinalBinary,
             timePasses = context.config.phaseConfig.needProfiling,
     )
-    useContext(OptimizationState(context.config, optimizationConfig)) {
+    useContext(OptimizationState(context.config, optimizationConfig, context.performanceManager)) {
         val module = this@runBitcodePostProcessing.context.llvmModule
         it.runPhase(StackProtectorPhase, module)
         it.runPhase(MandatoryBitcodeLLVMPostprocessingPhase, module)

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/driver/phases/TopLevelPhases.kt

@@ -95,7 +95,8 @@ internal fun <C : PhaseContext> PhaseEngine<C>.runBackend(backendContext: Contex
             val outputFiles = OutputFiles(outputPath, config.target, config.produce)
             val generationState = NativeGenerationState(context.config, backendContext,
                     fragment.cacheDeserializationStrategy, fragment.dependenciesTracker, fragment.llvmModuleSpecification, outputFiles,
-                    llvmModuleName = "out" // TODO: Currently, all llvm modules are named as "out" which might lead to collisions.
+                    llvmModuleName = "out", // TODO: Currently, all llvm modules are named as "out" which might lead to collisions.
+                    fragment.performanceManager
             )
             try {
                 val module = fragment.irModule

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

@@ -0,0 +1,90 @@
+/*
+ * Copyright 2010-2025 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.llvm
+
+import kotlinx.cinterop.*
+import llvm.LLVMKotlinDisposePassesProfile
+import llvm.LLVMKotlinPassesProfileAsString
+import llvm.LLVMKotlinPassesProfileRef
+import llvm.LLVMKotlinPassesProfileRefVar
+import org.jetbrains.kotlin.util.DynamicStats
+import org.jetbrains.kotlin.util.PerformanceManager
+import org.jetbrains.kotlin.util.PhaseType
+import org.jetbrains.kotlin.util.PlatformType
+import org.jetbrains.kotlin.util.Time
+import org.jetbrains.kotlin.util.UnitStats
+
+internal data class LlvmPassProfile(val pass: String, val duration: Time)
+
+private fun LLVMKotlinPassesProfileRef.parse(pipelineName: String) = LLVMKotlinPassesProfileAsString(this)!!
+        .toKString()
+        .lineSequence()
+        .filter { it.isNotEmpty() }
+        .map {
+            val line = it.split("\t")
+            check(line.size == 2) {
+                "Expected line `$it` to have exactly 2 tab-separated columns"
+            }
+            val pass = "$pipelineName.${line[0]}"
+            val nanos = line[1].toLong()
+            // Computing user and system times for each llvm phase gives a big overhead.
+            val userNanos = 0L
+            val cpuNanos = 0L
+            LlvmPassProfile(pass, Time(nanos, userNanos, cpuNanos))
+        }
+
+@JvmInline
+internal value class LlvmPassesProfile private constructor(val entries: List<LlvmPassProfile>) {
+    constructor(profile: LLVMKotlinPassesProfileRef, pipelineName: String) : this(profile.parse(pipelineName).toList())
+}
+
+/**
+ * Run [block] that expects [LLVMKotlinPassesProfileRef] as output variable and return its result augmented by the parsed profile.
+ *
+ * @param enabled when false, no profile will be computed
+ * @param pipelineName will be prepended to each phase name
+ */
+internal inline fun <T> withLLVMPassesProfile(
+        enabled: Boolean,
+        pipelineName: String,
+        block: (CValuesRef<LLVMKotlinPassesProfileRefVar>?) -> T,
+): Pair<T, LlvmPassesProfile?> {
+    if (!enabled) {
+        return block(null) to null
+    }
+    return memScoped {
+        val profileRef = alloc<LLVMKotlinPassesProfileRefVar>().also {
+            it.value = null
+        }
+        val result = block(profileRef.ptr)
+        val profile = profileRef.value
+        defer {
+            LLVMKotlinDisposePassesProfile(profile)
+        }
+        result to LlvmPassesProfile(profile!!, pipelineName)
+    }
+}
+
+internal fun PerformanceManager.addLlvmPassesProfile(profile: LlvmPassesProfile) {
+    addOtherUnitStats(UnitStats(
+            name = "$presentableName (llvm)",
+            outputKind = null,
+            platform = PlatformType.Native,
+            filesCount = 0,
+            linesCount = 0,
+            initStats = null,
+            analysisStats = null,
+            translationToIrStats = null,
+            irPreLoweringStats = null,
+            irSerializationStats = null,
+            klibWritingStats = null,
+            irLoweringStats = null,
+            backendStats = null,
+            dynamicStats = profile.entries.map { (pass, duration) ->
+                DynamicStats(PhaseType.Backend, pass, duration)
+            }
+    ))
+}

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

@@ -40,10 +40,10 @@ internal fun linkRuntimeModules(generationState: NativeGenerationState, runtimeN
     }
     val config = createLTOPipelineConfigForRuntime(generationState)
 
-    MandatoryOptimizationPipeline(config, generationState).use {
+    MandatoryOptimizationPipeline(config, generationState.performanceManager, generationState).use {
         it.execute(runtimeModule)
     }
-    ModuleOptimizationPipeline(config, generationState).use {
+    ModuleOptimizationPipeline(config, generationState.performanceManager, generationState).use {
         it.execute(runtimeModule)
     }
 

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

@@ -9,6 +9,8 @@
 #include <llvm/Support/Timer.h>
 #include <llvm/Transforms/Utils/Cloning.h>
 
+#include "PassesProfileHandler.h"
+
 using namespace llvm;
 
 namespace {
@@ -57,8 +59,10 @@ extern "C" LLVMErrorRef LLVMKotlinRunPasses(
         LLVMModuleRef M,
         const char *Passes,
         LLVMTargetMachineRef TM,
-        int InlinerThreshold
+        int InlinerThreshold,
+        LLVMKotlinPassesProfileRef* Profile
 ) {
+    // Implementation is taken from https://github.com/Kotlin/llvm-project/blob/0fa53d5183ec3c0654631d719dd6dfa7a270ca98/llvm/lib/Passes/PassBuilderBindings.cpp#L47
     TargetMachine *Machine = unwrap(TM);
     Module *Mod = unwrap(M);
 
@@ -81,10 +85,20 @@ extern "C" LLVMErrorRef LLVMKotlinRunPasses(
     StandardInstrumentations SI(Mod->getContext(), false, false);
     SI.registerCallbacks(PIC, &MAM);
 
+    PassesProfileHandler PPH(Profile != nullptr);
+    // Putting last to make this the last callback for before* events;
+    // the handler will additionally make sure its after* events are handled before anything else.
+    // This makes it so the profile tracks phases only, ignoring other callbacks.
+    PPH.registerCallbacks(PIC);
+
     ModulePassManager MPM;
     if (auto Err = PB.parsePassPipeline(MPM, Passes))
         return wrap(std::move(Err));
     MPM.run(*Mod, MAM);
 
+    if (Profile != nullptr) {
+        *Profile = wrap(new PassesProfile(PPH.serialize()));
+    }
+
     return LLVMErrorSuccess;
 }