diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/optimizations/LivenessAnalysis.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/optimizations/LivenessAnalysis.kt index 97681569237..d25ca14b575 100644 --- a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/optimizations/LivenessAnalysis.kt +++ b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/optimizations/LivenessAnalysis.kt @@ -49,6 +49,7 @@ object LivenessAnalysis { private val loopEndsLV = mutableMapOf() private val loopStartsLV = mutableMapOf() private var catchesLV = BitSet() + private val suspensionPointIdParameters = BitSet() fun run(body: IrBody): Map> { 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()}") } diff --git a/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/driver/phases/NativeLoweringPhases.kt b/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/driver/phases/NativeLoweringPhases.kt index 30f1fe6aeab..e75338faca3 100644 --- a/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/driver/phases/NativeLoweringPhases.kt +++ b/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/driver/phases/NativeLoweringPhases.kt @@ -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.getAllLowerings() = listOfNotNull varargPhase, kotlinNothingValueExceptionPhase, coroutinesPhase, + coroutinesVarSpillingPhase, typeOperatorPhase, expressionBodyTransformPhase, objectClassesPhase, diff --git a/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/ir/Ir.kt b/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/ir/Ir.kt index 9c1e8e3e09f..795854302cb 100644 --- a/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/ir/Ir.kt +++ b/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/ir/Ir.kt @@ -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"))!! diff --git a/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/IntrinsicGenerator.kt b/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/IntrinsicGenerator.kt index f67df8610ce..9578041a6bd 100644 --- a/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/IntrinsicGenerator.kt +++ b/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/IntrinsicGenerator.kt @@ -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, diff --git a/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/CoroutinesVarSpillingLowering.kt b/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/CoroutinesVarSpillingLowering.kt new file mode 100644 index 00000000000..18ba1b96ded --- /dev/null +++ b/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/CoroutinesVarSpillingLowering.kt @@ -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() + 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> { + override fun visitSuspensionPoint(expression: IrSuspensionPoint, data: List): IrExpression { + expression.transformChildren(this, liveLocals[expression]!!) + + return expression + } + + override fun visitCall(expression: IrCall, data: List): 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()) + } +} \ No newline at end of file diff --git a/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/NativeSuspendFunctionLowering.kt b/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/NativeSuspendFunctionLowering.kt index e79845cae32..d4825378030 100644 --- a/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/NativeSuspendFunctionLowering.kt +++ b/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/NativeSuspendFunctionLowering.kt @@ -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(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() - // 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) = diff --git a/kotlin-native/runtime/src/main/kotlin/kotlin/native/internal/Coroutines.kt b/kotlin-native/runtime/src/main/kotlin/kotlin/native/internal/Coroutines.kt index c1af7acd4c6..776cced94da 100644 --- a/kotlin-native/runtime/src/main/kotlin/kotlin/native/internal/Coroutines.kt +++ b/kotlin-native/runtime/src/main/kotlin/kotlin/native/internal/Coroutines.kt @@ -25,3 +25,12 @@ internal inline suspend fun getCoroutineContext(): CoroutineContext = @TypedIntrinsic(IntrinsicType.RETURN_IF_SUSPENDED) @PublishedApi internal external suspend fun 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() + diff --git a/kotlin-native/runtime/src/main/kotlin/kotlin/native/internal/IntrinsicType.kt b/kotlin-native/runtime/src/main/kotlin/kotlin/native/internal/IntrinsicType.kt index 5e6cba7b151..8ad74d86578 100644 --- a/kotlin-native/runtime/src/main/kotlin/kotlin/native/internal/IntrinsicType.kt +++ b/kotlin-native/runtime/src/main/kotlin/kotlin/native/internal/IntrinsicType.kt @@ -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"