[IR][K/N] Extracted explicit var spilling phase from coroutines lowering

Not only does it give some flexibility but also it could be turned off in the future (by providing trivial replacement) as a workaround #KT-65153
2024-01-18 16:55:35 +02:00
parent 954eade78c
commit 6c05219485
8 changed files with 168 additions and 110 deletions
@@ -49,6 +49,7 @@ object LivenessAnalysis {
        private val loopEndsLV = mutableMapOf<IrLoop, BitSet>()
        private val loopStartsLV = mutableMapOf<IrLoop, BitSet>()
        private var catchesLV = BitSet()
+        private val suspensionPointIdParameters = BitSet()

        fun run(body: IrBody): Map<IrElement, List<IrVariable>> {
            body.accept(this, BitSet() /* No variable is live at the end */)
@@ -81,6 +82,8 @@ object LivenessAnalysis {
                val elementLV = filteredElementEndsLV.getOrPut(element) { BitSet() }
                elementLV.or(liveVariables)
                elementLV.or(catchesLV)
+                // Suspension points id parameters are never live since they are provided at codegen.
+                elementLV.andNot(suspensionPointIdParameters)
            }
            return compute()
        }
@@ -184,6 +187,14 @@ object LivenessAnalysis {
            currentCatchesLV
        }

+        override fun visitSuspensionPoint(expression: IrSuspensionPoint, data: BitSet): BitSet {
+            val variableId = getVariableId(expression.suspensionPointIdParameter)
+            suspensionPointIdParameters.set(variableId)
+            return super.visitSuspensionPoint(expression, data).also {
+                suspensionPointIdParameters.clear(variableId)
+            }
+        }
+
        override fun visitBreak(jump: IrBreak, data: BitSet) = saveAndCompute(jump, data) {
            loopEndsLV[jump.loop] ?: error("Break from an unknown loop ${jump.loop.render()}")
        }
@@ -385,6 +385,13 @@ private val coroutinesPhase = createFileLoweringPhase(
        prerequisite = setOf(localFunctionsPhase, finallyBlocksPhase, kotlinNothingValueExceptionPhase)
 )

+private val coroutinesVarSpillingPhase = createFileLoweringPhase(
+        ::CoroutinesVarSpillingLowering,
+        name = "CoroutinesVarSpilling",
+        description = "Save/restore coroutines variables before/after suspension",
+        prerequisite = setOf(coroutinesPhase)
+)
+
 private val typeOperatorPhase = createFileLoweringPhase(
        ::TypeOperatorLowering,
        name = "TypeOperators",
@@ -547,6 +554,7 @@ private fun PhaseEngine<NativeGenerationState>.getAllLowerings() = listOfNotNull
        varargPhase,
        kotlinNothingValueExceptionPhase,
        coroutinesPhase,
+        coroutinesVarSpillingPhase,
        typeOperatorPhase,
        expressionBodyTransformPhase,
        objectClassesPhase,
@@ -385,6 +385,9 @@ internal class KonanSymbols(
    override val coroutineSuspendedGetter =
            findTopLevelPropertyGetter(StandardNames.COROUTINES_INTRINSICS_PACKAGE_FQ_NAME, COROUTINE_SUSPENDED_NAME, null)

+    val saveCoroutineState = internalFunction("saveCoroutineState")
+    val restoreCoroutineState = internalFunction("restoreCoroutineState")
+
    val cancellationException = topLevelClass(KonanFqNames.cancellationException)

    val kotlinResult = irBuiltIns.findClass(Name.identifier("Result"))!!
@@ -74,6 +74,8 @@ internal enum class IntrinsicType {
    // Coroutines
    GET_CONTINUATION,
    RETURN_IF_SUSPENDED,
+    SAVE_COROUTINE_STATE,
+    RESTORE_COROUTINE_STATE,
    // Interop
    INTEROP_READ_BITS,
    INTEROP_WRITE_BITS,
@@ -279,6 +281,8 @@ internal class IntrinsicGenerator(private val environment: IntrinsicGeneratorEnv
                IntrinsicType.GET_AND_ADD_ARRAY_ELEMENT -> emitGetAndAddArrayElement(callSite, args)
                IntrinsicType.GET_CONTINUATION,
                IntrinsicType.RETURN_IF_SUSPENDED,
+                IntrinsicType.SAVE_COROUTINE_STATE,
+                IntrinsicType.RESTORE_COROUTINE_STATE,
                IntrinsicType.INTEROP_BITS_TO_FLOAT,
                IntrinsicType.INTEROP_BITS_TO_DOUBLE,
                IntrinsicType.INTEROP_SIGN_EXTEND,
@@ -0,0 +1,97 @@
+/*
+ * Copyright 2010-2024 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.lower
+
+import org.jetbrains.kotlin.backend.common.BodyLoweringPass
+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.descriptors.DescriptorVisibilities
+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.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.visitors.IrElementTransformer
+
+internal val DECLARATION_ORIGIN_COROUTINE_VAR_SPILLING = IrDeclarationOriginImpl("COROUTINE_VAR_SPILLING")
+
+internal class CoroutinesVarSpillingLowering(val context: Context) : BodyLoweringPass {
+    private val irFactory = context.irFactory
+    private val symbols = context.ir.symbols
+    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))
+            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>()
+        fun getFieldForSpilling(variable: IrVariable) = localToPropertyMap.getOrPut(variable.symbol) {
+            variable.isVar = true // Make variables mutable in order to save/restore them.
+            irFactory.buildField {
+                startOffset = coroutineClass.startOffset
+                endOffset = coroutineClass.endOffset
+                origin = DECLARATION_ORIGIN_COROUTINE_VAR_SPILLING
+                name = variable.name
+                type = variable.type
+                visibility = DescriptorVisibilities.PRIVATE
+                isFinal = false
+            }.apply {
+                coroutineClass.addChild(this)
+            }
+        }
+
+        // Save/restore state at suspension points.
+        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]!!)
+
+                return expression
+            }
+
+            override fun visitCall(expression: IrCall, data: List<IrVariable>): IrExpression {
+                expression.transformChildren(this, data)
+
+                return when (expression.symbol) {
+                    saveCoroutineState -> irBuilder.run {
+                        irBlock(expression) {
+                            for (variable in data) {
+                                val field = getFieldForSpilling(variable)
+                                +irSetField(irGet(thisReceiver), field, irGet(variable))
+                            }
+                        }
+                    }
+                    restoreCoroutineState -> irBuilder.run {
+                        irBlock(expression) {
+                            for (variable in data) {
+                                val field = getFieldForSpilling(variable)
+                                +irSet(variable, irGetField(irGet(thisReceiver), field))
+                            }
+                        }
+                    }
+                    else -> expression
+                }
+            }
+        }, data = emptyList())
+    }
+}
@@ -2,11 +2,8 @@ package org.jetbrains.kotlin.backend.konan.lower

 import org.jetbrains.kotlin.backend.common.descriptors.synthesizedName
 import org.jetbrains.kotlin.backend.common.lower.*
-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.descriptors.DescriptorVisibilities
-import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.ObsoleteDescriptorBasedAPI
@@ -17,8 +14,6 @@ import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.IrSuspendableExpressionImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrSuspensionPointImpl
 import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
-import org.jetbrains.kotlin.ir.symbols.IrVariableSymbol
-import org.jetbrains.kotlin.ir.symbols.impl.IrSimpleFunctionSymbolImpl
 import org.jetbrains.kotlin.ir.symbols.impl.IrVariableSymbolImpl
 import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.types.isUnit
@@ -33,6 +28,8 @@ internal class NativeSuspendFunctionsLowering(
 ) : AbstractSuspendFunctionsLowering<Context>(generationState.context) {
    private val symbols = context.ir.symbols
    private val fileLowerState = generationState.fileLowerState
+    private val saveCoroutineState = symbols.saveCoroutineState
+    private val restoreCoroutineState = symbols.restoreCoroutineState

    override val stateMachineMethodName = Name.identifier("invokeSuspend")

@@ -83,23 +80,10 @@ internal class NativeSuspendFunctionsLowering(
            val suspendResult = irVar("suspendResult".synthesizedName, context.irBuiltIns.anyNType, true)

            // Extract all suspend calls to temporaries in order to make correct jumps to them.
-            originalBody.transformChildrenVoid(ExpressionSlicer(labelField.type, transformingFunction, saveState, restoreState, suspendResult, resultArgument))
-
-            val liveLocals = LivenessAnalysis.run(originalBody) { it is IrSuspensionPoint }
-
-            val localToPropertyMap = mutableMapOf<IrVariableSymbol, IrField>()
-            // TODO: optimize by using the same property for different locals.
-            liveLocals.values.forEach { scope ->
-                scope.forEach {
-                    it.isVar = true // Make variables mutable in order to save/restore them.
-                    localToPropertyMap.getOrPut(it.symbol) {
-                        coroutineClass.addField(it.name, it.type, true)
-                    }
-                }
-            }
+            originalBody.transformChildrenVoid(
+                    ExpressionSlicer(labelField.type, transformingFunction, suspendResult, resultArgument, thisReceiver, labelField))

            originalBody.transformChildrenVoid(object : IrElementTransformerVoid() {
-
                // Replace returns to refer to the new function.
                override fun visitReturn(expression: IrReturn): IrExpression {
                    expression.transformChildrenVoid(this)
@@ -125,48 +109,10 @@ internal class NativeSuspendFunctionsLowering(
                            ?: return expression
                    return irSetField(irGet(thisReceiver), capturedValue, expression.value)
                }
-
-                // Save/restore state at suspension points.
-                override fun visitExpression(expression: IrExpression): IrExpression {
-                    expression.transformChildrenVoid(this)
-
-                    val suspensionPoint = expression as? IrSuspensionPoint
-                            ?: return expression
-
-                    suspensionPoint.transformChildrenVoid(object : IrElementTransformerVoid() {
-                        override fun visitCall(expression: IrCall): IrExpression {
-                            expression.transformChildrenVoid(this)
-
-                            when (expression.symbol) {
-                                saveState.symbol -> {
-                                    val scope = liveLocals[suspensionPoint]!!
-                                    return irBlock(expression) {
-                                        for (variable in scope)
-                                            +irSetField(irGet(thisReceiver), localToPropertyMap[variable.symbol]!!, irGet(variable))
-                                        +irSetField(
-                                                irGet(thisReceiver),
-                                                labelField,
-                                                irGet(suspensionPoint.suspensionPointIdParameter)
-                                        )
-                                    }
-                                }
-                                restoreState.symbol -> {
-                                    val scope = liveLocals[suspensionPoint]!!
-                                    return irBlock(expression) {
-                                        for (variable in scope)
-                                            +irSet(variable, irGetField(irGet(thisReceiver), localToPropertyMap[variable.symbol]!!))
-                                    }
-                                }
-                            }
-                            return expression
-                        }
-                    })
-
-                    return suspensionPoint
-                }
            })
+
            originalBody.setDeclarationsParent(stateMachineFunction)
-            val statements = (originalBody as IrBlockBody).statements
+
            +suspendResult
            +IrSuspendableExpressionImpl(
                    startOffset, endOffset,
@@ -174,7 +120,7 @@ internal class NativeSuspendFunctionsLowering(
                    suspensionPointId = irGetField(irGet(stateMachineFunction.dispatchReceiverParameter!!), labelField),
                    result = irBlock(startOffset, endOffset) {
                        +irThrowIfNotNull(irExceptionOrNull(irGet(resultArgument))) // Coroutine might start with an exception.
-                        statements.forEach { +it }
+                        (originalBody as IrBlockBody).statements.forEach { +it }
                    })
            if (transformingFunction.returnType.isUnit())
                +irReturn(irGetObject(symbols.unit))                             // Insert explicit return for Unit functions.
@@ -184,9 +130,14 @@ internal class NativeSuspendFunctionsLowering(
    private var tempIndex = 0
    private var suspensionPointIdIndex = 0

-    private inner class ExpressionSlicer(val suspensionPointIdType: IrType, val irFunction: IrFunction,
-                                         val saveState: IrFunction, val restoreState: IrFunction,
-                                         val suspendResult: IrVariable, val resultArgument: IrValueParameter): IrElementTransformerVoid() {
+    private inner class ExpressionSlicer(
+            val suspensionPointIdType: IrType,
+            val irFunction: IrFunction,
+            val suspendResult: IrVariable,
+            val resultArgument: IrValueParameter,
+            val thisReceiver: IrValueParameter,
+            val labelField: IrField,
+    ) : IrElementTransformerVoid() {
        // TODO: optimize - it has square complexity.

        override fun visitSetField(expression: IrSetField): IrExpression {
@@ -250,6 +201,19 @@ internal class NativeSuspendFunctionsLowering(
                    }
                }

+                val suspensionPointIdParameter by lazy {
+                    irVar("suspensionPointId${suspensionPointIdIndex++}".synthesizedName, suspensionPointIdType)
+                }
+
+                fun saveState() = irBlock {
+                    +irCall(saveCoroutineState)
+                    +irSetField(
+                            irGet(thisReceiver),
+                            labelField,
+                            irGet(suspensionPointIdParameter)
+                    )
+                }
+
                var calledSaveState = false
                var suspendCall: IrExpression? = null
                when {
@@ -257,7 +221,7 @@ internal class NativeSuspendFunctionsLowering(
                        calledSaveState = true
                        val firstArgument = newChildren[2]!!
                        newChildren[2] = irBlock(firstArgument) {
-                            +irCall(saveState)
+                            +saveState()
                            +firstArgument
                        }
                        suspendCall = newChildren[2]
@@ -271,12 +235,12 @@ internal class NativeSuspendFunctionsLowering(
                            newChildren[lastChildIndex] =
                                    irBlock(lastChild) {
                                        if (lastChild.isPure()) {
-                                            +irCall(saveState)
+                                            +saveState()
                                            +lastChild
                                        } else {
                                            val tmp = irVar(lastChild)
                                            +tmp
-                                            +irCall(saveState)
+                                            +saveState()
                                            +irGet(tmp)
                                        }
                                    }
@@ -306,19 +270,16 @@ internal class NativeSuspendFunctionsLowering(
                        startOffset                = startOffset,
                        endOffset                  = endOffset,
                        type                       = context.irBuiltIns.anyNType,
-                        suspensionPointIdParameter = irVar(
-                                "suspensionPointId${suspensionPointIdIndex++}".synthesizedName,
-                                suspensionPointIdType
-                        ),
+                        suspensionPointIdParameter = suspensionPointIdParameter,
                        result                     = irBlock(startOffset, endOffset) {
                            if (!calledSaveState)
-                                +irCall(saveState)
+                                +saveState()
                            +irSet(suspendResult.symbol, suspendCall)
                            +irReturnIfSuspended(suspendResult)
                            +irGet(suspendResult)
                        },
                        resumeResult               = irBlock(startOffset, endOffset) {
-                            +irCall(restoreState)
+                            +irCall(restoreCoroutineState)
                            +irGetOrThrow(irGet(resultArgument))
                        })
                val expressionResult = when {
@@ -399,43 +360,6 @@ internal class NativeSuspendFunctionsLowering(
        }
    }

-    // These are marker functions to split up the lowering on two parts.
-    private val saveState =
-        context.irFactory.createSimpleFunction(
-                SYNTHETIC_OFFSET,
-                SYNTHETIC_OFFSET,
-                IrDeclarationOrigin.DEFINED,
-                "saveState".synthesizedName,
-                DescriptorVisibilities.PRIVATE,
-                isInline = false,
-                isExpect = false,
-                context.irBuiltIns.unitType,
-                Modality.ABSTRACT,
-                IrSimpleFunctionSymbolImpl(),
-                isTailrec = false,
-                isSuspend = false,
-                isOperator = false,
-                isInfix = false,
-        )
-
-    private val restoreState =
-            context.irFactory.createSimpleFunction(
-                    SYNTHETIC_OFFSET,
-                    SYNTHETIC_OFFSET,
-                    IrDeclarationOrigin.DEFINED,
-                    "restoreState".synthesizedName,
-                    DescriptorVisibilities.PRIVATE,
-                    isInline = false,
-                    isExpect = false,
-                    context.irBuiltIns.unitType,
-                    Modality.ABSTRACT,
-                    IrSimpleFunctionSymbolImpl(),
-                    isTailrec = false,
-                    isSuspend = false,
-                    isOperator = false,
-                    isInfix = false,
-            )
-
    private fun IrBuilderWithScope.irVar(name: Name, type: IrType,
                                         isMutable: Boolean = false,
                                         initializer: IrExpression? = null) =
@@ -25,3 +25,12 @@ internal inline suspend fun getCoroutineContext(): CoroutineContext =
@TypedIntrinsic(IntrinsicType.RETURN_IF_SUSPENDED)
@PublishedApi
 internal external suspend fun <T> returnIfSuspended(@Suppress("UNUSED_PARAMETER") argument: Any?): T
+
+@TypedIntrinsic(IntrinsicType.SAVE_COROUTINE_STATE)
+@PublishedApi
+internal external fun saveCoroutineState()
+
+@TypedIntrinsic(IntrinsicType.RESTORE_COROUTINE_STATE)
+@PublishedApi
+internal external fun restoreCoroutineState()
+
@@ -64,6 +64,8 @@ internal class IntrinsicType {
        // Coroutines
        const val GET_CONTINUATION              = "GET_CONTINUATION"
        const val RETURN_IF_SUSPENDED           = "RETURN_IF_SUSPENDED"
+        const val SAVE_COROUTINE_STATE          = "SAVE_COROUTINE_STATE"
+        const val RESTORE_COROUTINE_STATE       = "RESTORE_COROUTINE_STATE"

        // Interop
        const val INTEROP_READ_PRIMITIVE        = "INTEROP_READ_PRIMITIVE"