[K/N] Fallback in case turned off liveness analysis

#KT-65153 Fixed
2024-01-22 17:01:47 +02:00
parent 6c05219485
commit 27fa3ebfb3
3 changed files with 114 additions and 16 deletions
@@ -18,6 +18,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.ir.expressions.IrSuspensionPoint

 internal class InlineFunctionOriginInfo(val irFunction: IrFunction, val irFile: IrFile, val startOffset: Int, val endOffset: Int)

@@ -73,6 +74,8 @@ internal class NativeGenerationState(
    val calledFromExportedInlineFunctions = mutableSetOf<IrFunction>()
    val constructedFromExportedInlineFunctions = mutableSetOf<IrClass>()
    val inlineFunctionOrigins = mutableMapOf<IrFunction, InlineFunctionOriginInfo>()
+    val liveVariablesAtSuspensionPoints = mutableMapOf<IrSuspensionPoint, List<IrVariable>>()
+    val visibleVariablesAtSuspensionPoints = mutableMapOf<IrSuspensionPoint, List<IrVariable>>()

    private val localClassNames = mutableMapOf<IrAttributeContainer, String>()
    fun getLocalClassName(container: IrAttributeContainer): String? = localClassNames[container.attributeOwnerId]
@@ -5,6 +5,7 @@

 package org.jetbrains.kotlin.backend.konan.driver.phases

+import org.jetbrains.kotlin.backend.common.BodyLoweringPass
 import org.jetbrains.kotlin.backend.common.FileLoweringPass
 import org.jetbrains.kotlin.backend.common.lower.*
 import org.jetbrains.kotlin.backend.common.lower.coroutines.AddContinuationToNonLocalSuspendFunctionsLowering
@@ -13,14 +14,13 @@ import org.jetbrains.kotlin.backend.common.lower.inline.LocalClassesExtractionFr
 import org.jetbrains.kotlin.backend.common.lower.inline.LocalClassesInInlineFunctionsLowering
 import org.jetbrains.kotlin.backend.common.lower.inline.LocalClassesInInlineLambdasLowering
 import org.jetbrains.kotlin.backend.common.lower.loops.ForLoopsLowering
+import org.jetbrains.kotlin.backend.common.lower.optimizations.LivenessAnalysis
 import org.jetbrains.kotlin.backend.common.lower.optimizations.PropertyAccessorInlineLowering
 import org.jetbrains.kotlin.backend.common.phaser.*
 import org.jetbrains.kotlin.backend.common.runOnFilePostfix
 import org.jetbrains.kotlin.backend.konan.Context
 import org.jetbrains.kotlin.backend.konan.NativeGenerationState
 import org.jetbrains.kotlin.backend.konan.driver.PhaseEngine
-import org.jetbrains.kotlin.ir.declarations.IrFile
-import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 import org.jetbrains.kotlin.backend.konan.*
 import org.jetbrains.kotlin.backend.konan.driver.utilities.getDefaultIrActions
 import org.jetbrains.kotlin.backend.konan.ir.FunctionsWithoutBoundCheckGenerator
@@ -33,7 +33,9 @@ import org.jetbrains.kotlin.backend.konan.lower.ReturnsInsertionLowering
 import org.jetbrains.kotlin.backend.konan.lower.UnboxInlineLowering
 import org.jetbrains.kotlin.backend.konan.optimizations.KonanBCEForLoopBodyTransformer
 import org.jetbrains.kotlin.ir.IrElement
-import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.ir.declarations.*
+import org.jetbrains.kotlin.ir.expressions.IrBody
+import org.jetbrains.kotlin.ir.expressions.IrSuspensionPoint
 import org.jetbrains.kotlin.ir.interpreter.IrInterpreterConfiguration
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
 import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
@@ -385,8 +387,32 @@ private val coroutinesPhase = createFileLoweringPhase(
        prerequisite = setOf(localFunctionsPhase, finallyBlocksPhase, kotlinNothingValueExceptionPhase)
 )

+private val coroutinesLivenessAnalysisFallbackPhase = createFileLoweringPhase(
+        lowering = ::CoroutinesLivenessAnalysisFallback,
+        name = "CoroutinesLivenessAnalysisFallback",
+        description = "Compute visible variables at suspension points",
+        prerequisite = setOf(coroutinesPhase)
+)
+
+private val coroutinesLivenessAnalysisPhase = createFileLoweringPhase(
+        lowering = { context: NativeGenerationState ->
+            object : BodyLoweringPass {
+                override fun lower(irBody: IrBody, container: IrDeclaration) {
+                    val liveVariablesAtSuspensionPoints = context.liveVariablesAtSuspensionPoints
+                    LivenessAnalysis.run(irBody) { it is IrSuspensionPoint }
+                            .forEach { (irElement, liveVariables) ->
+                                liveVariablesAtSuspensionPoints[irElement as IrSuspensionPoint] = liveVariables
+                            }
+                }
+            }
+        },
+        name = "CoroutinesLivenessAnalysis",
+        description = "Run liveness analysis for coroutines",
+        prerequisite = setOf(coroutinesPhase)
+)
+
 private val coroutinesVarSpillingPhase = createFileLoweringPhase(
-        ::CoroutinesVarSpillingLowering,
+        lowering = ::CoroutinesVarSpillingLowering,
        name = "CoroutinesVarSpilling",
        description = "Save/restore coroutines variables before/after suspension",
        prerequisite = setOf(coroutinesPhase)
@@ -554,6 +580,9 @@ private fun PhaseEngine<NativeGenerationState>.getAllLowerings() = listOfNotNull
        varargPhase,
        kotlinNothingValueExceptionPhase,
        coroutinesPhase,
+        // Either of these could be turned off without losing correctness.
+        coroutinesLivenessAnalysisPhase, // This is more optimal
+        coroutinesLivenessAnalysisFallbackPhase, // While this is simple
        coroutinesVarSpillingPhase,
        typeOperatorPhase,
        expressionBodyTransformPhase,
@@ -5,43 +5,45 @@

 package org.jetbrains.kotlin.backend.konan.lower

-import org.jetbrains.kotlin.backend.common.BodyLoweringPass
+import org.jetbrains.kotlin.backend.common.*
 import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
 import org.jetbrains.kotlin.backend.common.lower.irBlock
-import org.jetbrains.kotlin.backend.common.lower.optimizations.LivenessAnalysis
-import org.jetbrains.kotlin.backend.konan.Context
+import org.jetbrains.kotlin.backend.konan.NativeGenerationState
+import org.jetbrains.kotlin.backend.konan.ir.KonanSymbols
 import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
+import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.builders.declarations.buildField
 import org.jetbrains.kotlin.ir.builders.irGet
 import org.jetbrains.kotlin.ir.builders.irGetField
 import org.jetbrains.kotlin.ir.builders.irSet
 import org.jetbrains.kotlin.ir.builders.irSetField
 import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.expressions.IrBody
-import org.jetbrains.kotlin.ir.expressions.IrCall
-import org.jetbrains.kotlin.ir.expressions.IrExpression
-import org.jetbrains.kotlin.ir.expressions.IrSuspensionPoint
+import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.symbols.IrVariableSymbol
 import org.jetbrains.kotlin.ir.util.addChild
 import org.jetbrains.kotlin.ir.util.overrides
 import org.jetbrains.kotlin.ir.util.parentAsClass
+import org.jetbrains.kotlin.ir.util.render
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformer
+import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
+import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid

 internal val DECLARATION_ORIGIN_COROUTINE_VAR_SPILLING = IrDeclarationOriginImpl("COROUTINE_VAR_SPILLING")

-internal class CoroutinesVarSpillingLowering(val context: Context) : BodyLoweringPass {
+internal class CoroutinesVarSpillingLowering(val generationState: NativeGenerationState) : BodyLoweringPass {
+    private val context = generationState.context
    private val irFactory = context.irFactory
    private val symbols = context.ir.symbols
+    private val invokeSuspendFunction = symbols.invokeSuspendFunction
    private val saveCoroutineState = symbols.saveCoroutineState
    private val restoreCoroutineState = symbols.restoreCoroutineState

    override fun lower(irBody: IrBody, container: IrDeclaration) {
        val thisReceiver = (container as? IrSimpleFunction)?.dispatchReceiverParameter
-        if (thisReceiver == null || !container.overrides(context.ir.symbols.invokeSuspendFunction.owner))
+        if (thisReceiver == null || !container.overrides(invokeSuspendFunction.owner))
            return

        val coroutineClass = container.parentAsClass
-        val liveLocals = LivenessAnalysis.run(irBody) { it is IrSuspensionPoint }

        // TODO: optimize by using the same property for different locals.
        val localToPropertyMap = mutableMapOf<IrVariableSymbol, IrField>()
@@ -64,7 +66,10 @@ internal class CoroutinesVarSpillingLowering(val context: Context) : BodyLowerin
        val irBuilder = context.createIrBuilder(container.symbol, container.startOffset, container.endOffset)
        irBody.transformChildren(object : IrElementTransformer<List<IrVariable>> {
            override fun visitSuspensionPoint(expression: IrSuspensionPoint, data: List<IrVariable>): IrExpression {
-                expression.transformChildren(this, liveLocals[expression]!!)
+                val liveVariables = generationState.liveVariablesAtSuspensionPoints[expression]
+                        ?: generationState.visibleVariablesAtSuspensionPoints[expression]
+                        ?: error("No live variables for ${container.render()} at ${expression.suspensionPointIdParameter.name}")
+                expression.transformChildren(this, liveVariables)

                return expression
            }
@@ -94,4 +99,65 @@ internal class CoroutinesVarSpillingLowering(val context: Context) : BodyLowerin
            }
        }, data = emptyList())
    }
-}
+}
+
+internal class CoroutinesLivenessAnalysisFallback(val generationState: NativeGenerationState) : FileLoweringPass, IrElementVisitorVoid {
+    private val invokeSuspendFunction = generationState.context.ir.symbols.invokeSuspendFunction
+
+    override fun lower(irFile: IrFile) {
+        if (generationState.liveVariablesAtSuspensionPoints.isEmpty())
+            irFile.acceptChildrenVoid(this)
+    }
+
+    override fun visitElement(element: IrElement) {
+        element.acceptChildrenVoid(this)
+    }
+
+    override fun visitFunction(declaration: IrFunction) {
+        val body = declaration.body
+        if (body != null && declaration.dispatchReceiverParameter != null
+                && (declaration as? IrSimpleFunction)?.overrides(invokeSuspendFunction.owner) == true
+        ) {
+            computeVisibleVariablesAtSuspensionPoints(body)
+        }
+    }
+
+    private fun computeVisibleVariablesAtSuspensionPoints(body: IrBody) {
+        body.acceptChildrenVoid(object : IrElementVisitorVoid {
+            val scopeStack = mutableListOf<MutableSet<IrVariable>>(mutableSetOf())
+
+            override fun visitElement(element: IrElement) {
+                element.acceptChildrenVoid(this)
+            }
+
+            override fun visitContainerExpression(expression: IrContainerExpression) {
+                if (!expression.isTransparentScope)
+                    scopeStack.push(mutableSetOf())
+                super.visitContainerExpression(expression)
+                if (!expression.isTransparentScope)
+                    scopeStack.pop()
+            }
+
+            override fun visitCatch(aCatch: IrCatch) {
+                scopeStack.push(mutableSetOf())
+                super.visitCatch(aCatch)
+                scopeStack.pop()
+            }
+
+            override fun visitVariable(declaration: IrVariable) {
+                super.visitVariable(declaration)
+                scopeStack.peek()!!.add(declaration)
+            }
+
+            override fun visitSuspensionPoint(expression: IrSuspensionPoint) {
+                // Skip suspensionPointIdParameter, because we don't want to save it.
+                expression.result.acceptChildrenVoid(this)
+                expression.resumeResult.acceptChildrenVoid(this)
+
+                val visibleVariables = mutableListOf<IrVariable>()
+                scopeStack.forEach { visibleVariables += it }
+                generationState.visibleVariablesAtSuspensionPoints[expression] = visibleVariables
+            }
+        })
+    }
+}