[WASM] Optimize varargs without spreads

2022-10-26 18:44:15 +02:00
parent 913ce9d817
commit 346b2f162c
19 changed files with 211 additions and 37 deletions
@@ -77,7 +77,7 @@ abstract class AbstractValueUsageTransformer(
    protected open fun IrExpression.useAsResult(enclosing: IrExpression): IrExpression =
        this.useAs(enclosing.type)

-    protected open fun IrExpression.useAsVarargElement(expression: IrVararg): IrExpression = this
+    protected open fun useAsVarargElement(element: IrExpression, expression: IrVararg): IrExpression = element

    override fun visitPropertyReference(expression: IrPropertyReference): IrExpression {
        TODO()
@@ -233,7 +233,7 @@ abstract class AbstractValueUsageTransformer(
                is IrSpreadElement ->
                    element.expression = element.expression.useAs(expression.type)
                is IrExpression -> {
-                    expression.putElement(i, element.useAsVarargElement(expression))
+                    expression.putElement(i, useAsVarargElement(element, expression))
                }
            }
        }
@@ -108,16 +108,14 @@ abstract class AbstractValueUsageLowering(val context: JsCommonBackendContext) :
        return this.useAsArgument(expression.target.valueParameters[parameter.index])
    }

-
-    override fun IrExpression.useAsVarargElement(expression: IrVararg): IrExpression {
-        return this.useAs(
+    override fun useAsVarargElement(element: IrExpression, expression: IrVararg): IrExpression =
+        element.useAs(
            // Do not box primitive inline classes
-            if (icUtils.isTypeInlined(type) && !icUtils.isTypeInlined(expression.type) && !expression.type.isPrimitiveArray())
+            if (icUtils.isTypeInlined(element.type) && !icUtils.isTypeInlined(expression.type) && !expression.type.isPrimitiveArray())
                irBuiltIns.anyNType
            else
                expression.varargElementType
        )
-    }
 }

 class AutoboxingTransformer(context: JsCommonBackendContext) : AbstractValueUsageLowering(context) {
@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.ir.declarations.impl.IrExternalPackageFragmentImpl
 import org.jetbrains.kotlin.ir.declarations.impl.IrFileImpl
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 import org.jetbrains.kotlin.ir.symbols.impl.DescriptorlessExternalPackageFragmentSymbol
+import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.types.IrTypeSystemContext
 import org.jetbrains.kotlin.ir.types.IrTypeSystemContextImpl
 import org.jetbrains.kotlin.ir.util.SymbolTable
@@ -37,6 +37,14 @@ internal class WasmUsefulDeclarationProcessor(
            super.visitVariable(declaration, data)
        }

+        override fun visitVararg(expression: IrVararg, data: IrDeclaration) {
+            expression.type.getClass()!!
+                .constructors
+                .firstOrNull { it.hasWasmPrimitiveConstructorAnnotation() }
+                ?.enqueue(data, "implicit vararg constructor")
+            super.visitVararg(expression, data)
+        }
+
        private fun tryToProcessIntrinsicCall(from: IrDeclaration, call: IrCall): Boolean = when (call.symbol) {
            context.wasmSymbols.unboxIntrinsic -> {
                val fromType = call.getTypeArgument(0)
@@ -73,6 +73,59 @@ class BodyGenerator(
        error("Unexpected element of type ${element::class}")
    }

+    private fun tryGenerateConstVarargArray(irVararg: IrVararg, wasmArrayType: WasmImmediate.GcType): Boolean {
+        if (irVararg.elements.isEmpty()) return false
+
+        val kind = (irVararg.elements[0] as? IrConst<*>)?.kind ?: return false
+        if (kind == IrConstKind.String || kind == IrConstKind.Null) return false
+        if (irVararg.elements.any { it !is IrConst<*> || it.kind != kind }) return false
+
+        val elementConstValues = irVararg.elements.map { (it as IrConst<*>).value!! }
+
+        val resource = when (irVararg.varargElementType) {
+            irBuiltIns.byteType -> elementConstValues.map { (it as Byte).toLong() } to WasmI8
+            irBuiltIns.booleanType -> elementConstValues.map { if (it as Boolean) 1L else 0L } to WasmI8
+            irBuiltIns.intType -> elementConstValues.map { (it as Int).toLong() } to WasmI32
+            irBuiltIns.shortType -> elementConstValues.map { (it as Short).toLong() } to WasmI16
+            irBuiltIns.longType -> elementConstValues.map { it as Long } to WasmI64
+            else -> return false
+        }
+
+        val constantArrayId = context.referenceConstantArray(resource)
+
+        body.buildConstI32(0)
+        body.buildConstI32(irVararg.elements.size)
+        body.buildInstr(WasmOp.ARRAY_NEW_DATA, wasmArrayType, WasmImmediate.DataIdx(constantArrayId))
+        return true
+    }
+
+    private fun tryGenerateVarargArray(irVararg: IrVararg, wasmArrayType: WasmImmediate.GcType) {
+        irVararg.elements.forEach {
+            check(it is IrExpression)
+            generateExpression(it)
+        }
+
+        val length = WasmImmediate.ConstI32(irVararg.elements.size)
+        body.buildInstr(WasmOp.ARRAY_NEW_FIXED, wasmArrayType, length)
+    }
+
+    override fun visitVararg(expression: IrVararg) {
+        val arrayClass = expression.type.getClass()!!
+
+        val wasmArrayType = arrayClass.constructors
+            .mapNotNull { it.valueParameters.singleOrNull()?.type }
+            .firstOrNull { it.getClass()?.getWasmArrayAnnotation() != null }
+            ?.getRuntimeClass(irBuiltIns)?.symbol
+            ?.let(context::referenceGcType)
+            ?.let(WasmImmediate::GcType)
+
+        check(wasmArrayType != null)
+
+        generateAnyParameters(arrayClass.symbol)
+        if (!tryGenerateConstVarargArray(expression, wasmArrayType)) tryGenerateVarargArray(expression, wasmArrayType)
+        body.buildStructNew(context.referenceGcType(expression.type.getRuntimeClass(irBuiltIns).symbol))
+    }
+
    override fun visitThrow(expression: IrThrow) {
        generateExpression(expression.value)
        body.buildThrow(context.tagIdx)
@@ -256,14 +309,18 @@ class BodyGenerator(
        generateCall(expression)
    }

+    private fun generateBox(expression: IrExpression, type: IrType) {
+        val klassSymbol = type.getRuntimeClass(irBuiltIns).symbol
+        generateAnyParameters(klassSymbol)
+        generateExpression(expression)
+        body.buildStructNew(context.referenceGcType(klassSymbol))
+    }
+
    private fun generateCall(call: IrFunctionAccessExpression) {
        // Box intrinsic has an additional klass ID argument.
        // Processing it separately
        if (call.symbol == wasmSymbols.boxIntrinsic) {
-            val klassSymbol = call.getTypeArgument(0)!!.getRuntimeClass(irBuiltIns).symbol
-            generateAnyParameters(klassSymbol)
-            generateExpression(call.getValueArgument(0)!!)
-            body.buildStructNew(context.referenceGcType(klassSymbol))
+            generateBox(call.getValueArgument(0)!!, call.getTypeArgument(0)!!)
            return
        }

@@ -491,17 +548,17 @@ class BodyGenerator(
                body.buildInstr(WasmOp.ARRAY_COPY, immediate, immediate)
            }

-            wasmSymbols.wasmArrayNewData0 -> {
-                val immediate = WasmImmediate.GcType(
-                    context.referenceGcType(call.getTypeArgument(0)!!.getRuntimeClass(irBuiltIns).symbol)
-                )
-                body.buildInstr(WasmOp.ARRAY_NEW_DATA, immediate, WasmImmediate.DataIdx(0))
-            }
-
            wasmSymbols.stringGetPoolSize -> {
                body.buildConstI32Symbol(context.stringPoolSize)
            }

+            wasmSymbols.wasmArrayNewData0 -> {
+                val arrayGcType = WasmImmediate.GcType(
+                    context.referenceGcType(call.getTypeArgument(0)!!.getRuntimeClass(irBuiltIns).symbol)
+                )
+                body.buildInstr(WasmOp.ARRAY_NEW_DATA, arrayGcType, WasmImmediate.DataIdx(0))
+            }
+
            else -> {
                return false
            }
@@ -10,7 +10,10 @@ import org.jetbrains.kotlin.wasm.ir.WasmSymbol
 // Representation of constant data in Wasm memory

 internal const val CHAR_SIZE_BYTES = 2
+internal const val BYTE_SIZE_BYTES = 1
+internal const val SHORT_SIZE_BYTES = 2
 internal const val INT_SIZE_BYTES = 4
+internal const val LONG_SIZE_BYTES = 8

 sealed class ConstantDataElement {
    abstract val sizeInBytes: Int
@@ -45,6 +48,23 @@ class ConstantDataIntField(val name: String, val value: WasmSymbol<Int>) : Const
    override val sizeInBytes: Int = INT_SIZE_BYTES
 }

+class ConstantDataIntegerArray(val name: String, val value: List<Long>, private val integerSize: Int) : ConstantDataElement() {
+    override fun toBytes(): ByteArray {
+        val array = ByteArray(value.size * integerSize)
+        repeat(value.size) { i ->
+            value[i].toLittleEndianBytesTo(array, i * integerSize, integerSize)
+        }
+        return array
+    }
+
+    override fun dump(indent: String, startAddress: Int): String {
+        if (value.isEmpty()) return ""
+        return "${addressToString(startAddress)}: $indent i${integerSize * 8}[] : ${toBytes().contentToString()}   ;; $name\n"
+    }
+
+    override val sizeInBytes: Int = value.size * integerSize
+}
+
 class ConstantDataIntArray(val name: String, val value: List<WasmSymbol<Int>>) : ConstantDataElement() {
    override fun toBytes(): ByteArray {
        return value.fold(byteArrayOf()) { acc, el -> acc + el.owner.toLittleEndianBytes() }
@@ -95,6 +115,13 @@ class ConstantDataStruct(val name: String, val elements: List<ConstantDataElemen
    override val sizeInBytes: Int = elements.map { it.sizeInBytes }.sum()
 }

+fun Long.toLittleEndianBytesTo(to: ByteArray, offset: Int, size: Int) {
+    for (i in 0 until size) {
+        to[offset + i] = (this ushr (i * 8)).toByte()
+    }
+}
+
+
 fun Int.toLittleEndianBytes(): ByteArray {
    return ByteArray(4) {
        (this ushr (it * 8)).toByte()
@@ -38,6 +38,8 @@ interface WasmBaseCodegenContext {

    fun referenceStringLiteralAddressAndId(string: String): Pair<WasmSymbol<Int>, WasmSymbol<Int>>

+    fun referenceConstantArray(resource: Pair<List<Long>, WasmType>): WasmSymbol<Int>
+
    fun transformType(irType: IrType): WasmType
    fun transformFieldType(irType: IrType): WasmType

@@ -40,6 +40,8 @@ class WasmCompiledModuleFragment(val irBuiltIns: IrBuiltIns) {
        ReferencableElements<String, Int>()
    val stringLiteralPoolId =
        ReferencableElements<String, Int>()
+    val constantArrayDataSegmentId =
+        ReferencableElements<Pair<List<Long>, WasmType>, Int>()

    private val tagFuncType = WasmFunctionType(
        listOf(
@@ -148,6 +150,18 @@ class WasmCompiledModuleFragment(val irBuiltIns: IrBuiltIns) {

        val data = mutableListOf<WasmData>()
        data.add(WasmData(WasmDataMode.Passive, stringDataSectionBytes.toByteArray()))
+        constantArrayDataSegmentId.unbound.forEach { (constantArraySegment, symbol) ->
+            symbol.bind(data.size)
+            val integerSize = when (constantArraySegment.second) {
+                WasmI8 -> BYTE_SIZE_BYTES
+                WasmI16 -> SHORT_SIZE_BYTES
+                WasmI32 -> INT_SIZE_BYTES
+                WasmI64 -> LONG_SIZE_BYTES
+                else -> TODO("type ${constantArraySegment.second} is not implemented")
+            }
+            val constData = ConstantDataIntegerArray("constant_array", constantArraySegment.first, integerSize)
+            data.add(WasmData(WasmDataMode.Passive, constData.toBytes()))
+        }
        typeInfo.buildData(data, address = { klassIds.getValue(it) })

        val masterInitFunctionType = WasmFunctionType(emptyList(), emptyList())
@@ -203,6 +217,7 @@ class WasmCompiledModuleFragment(val irBuiltIns: IrBuiltIns) {
            startFunction = null,  // Module is initialized via export call
            elements = emptyList(),
            data = data,
+            dataCount = true,
            tags = listOf(tag)
        )
        module.calculateIds()
@@ -67,6 +67,9 @@ class WasmModuleCodegenContextImpl(
        return address to id
    }

+    override fun referenceConstantArray(resource: Pair<List<Long>, WasmType>): WasmSymbol<Int> =
+        wasmFragment.constantArrayDataSegmentId.reference(resource)
+
    override fun generateTypeInfo(irClass: IrClassSymbol, typeInfo: ConstantDataElement) {
        wasmFragment.typeInfo.define(irClass, typeInfo)
    }
@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.backend.wasm.ir2wasm.isExternalType
 import org.jetbrains.kotlin.ir.backend.js.ir.JsIrBuilder
 import org.jetbrains.kotlin.ir.backend.js.lower.AbstractValueUsageLowering
 import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.IrVararg
 import org.jetbrains.kotlin.ir.types.IrType

 /**
@@ -26,4 +27,10 @@ class ExplicitlyCastExternalTypesLowering(wasmContext: WasmBackendContext) : Abs

        return this
    }
+
+    override fun useAsVarargElement(element: IrExpression, expression: IrVararg): IrExpression =
+        if (isExternalType(element.type))
+            element.useAs(irBuiltIns.anyNType)
+        else
+            super.useAsVarargElement(element, expression)
 }
@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.backend.wasm.lower

 import org.jetbrains.kotlin.backend.common.FileLoweringPass
 import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
+import org.jetbrains.kotlin.backend.common.lower.DeclarationIrBuilder
 import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
 import org.jetbrains.kotlin.backend.common.lower.irComposite
 import org.jetbrains.kotlin.backend.wasm.WasmBackendContext
@@ -20,7 +21,6 @@ import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.util.OperatorNameConventions
-import java.lang.IllegalArgumentException

 internal class WasmVarargExpressionLowering(
    private val context: WasmBackendContext
@@ -44,11 +44,11 @@ internal class WasmVarargExpressionLowering(
            get() = arrayClass.symbol in context.wasmSymbols.unsignedTypesToUnsignedArrays.values

        val primaryConstructor: IrConstructor
-            get() {
+            get() =
                if (isUnsigned)
-                    return arrayClass.constructors.find { it.valueParameters.singleOrNull()?.type == context.irBuiltIns.intType }!!
-                return arrayClass.primaryConstructor!!
-            }
+                    arrayClass.constructors.find { it.valueParameters.singleOrNull()?.type == context.irBuiltIns.intType }!!
+                else arrayClass.primaryConstructor!!
+
        val constructors
            get() = arrayClass.constructors

@@ -74,7 +74,6 @@ internal class WasmVarargExpressionLowering(

                return func?.owner ?: throw IllegalArgumentException("copyInto is not found for ${arrayType.render()}")
            }
-
    }

    private fun IrBlockBuilder.irCreateArray(size: IrExpression, arrDescr: ArrayDescr) =
@@ -164,16 +163,37 @@ internal class WasmVarargExpressionLowering(
            this@irSize.irSize()
        }

+    private fun tryVisitWithNoSpread(irVararg: IrVararg, builder: DeclarationIrBuilder): IrExpression {
+        val irVarargType = irVararg.type
+        if (!irVarargType.isUnsignedArray()) return irVararg
+
+        val unsignedConstructor = irVarargType.getClass()!!.primaryConstructor!!
+        val constructorParameterType = unsignedConstructor.valueParameters[0].type
+        val signedElementType = constructorParameterType.getArrayElementType(context.irBuiltIns)
+
+        irVararg.type = constructorParameterType
+        irVararg.varargElementType = signedElementType
+        return builder.irCall(unsignedConstructor.symbol, irVarargType).also {
+            it.putValueArgument(0, irVararg)
+        }
+    }
+
    override fun visitVararg(expression: IrVararg): IrExpression {
        // Optimization in case if we have a single spread element
-        if (expression.elements.size == 1 && expression.elements.first() is IrSpreadElement) {
-            val spreadExpr = (expression.elements.first() as IrSpreadElement).expression
+        val singleSpreadElement = expression.elements.singleOrNull() as? IrSpreadElement
+        if (singleSpreadElement != null) {
+            val spreadExpr = singleSpreadElement.expression
            if (isImmediatelyCreatedArray(spreadExpr))
                return spreadExpr.transform(this, null)
        }

        // Lower nested varargs
        val irVararg = super.visitVararg(expression) as IrVararg
+        val builder = context.createIrBuilder(currentScope!!.scope.scopeOwnerSymbol)
+
+        if (irVararg.elements.none { it is IrSpreadElement }) {
+            return tryVisitWithNoSpread(irVararg, builder)
+        }

        // Create temporary variable for each element and emit them all at once to preserve
        // argument evaluation order as per kotlin language spec.
@@ -202,9 +222,7 @@ internal class WasmVarargExpressionLowering(
                yield(VarargSegmentBuilder.Plain(currentElements.toList(), context))
        }.toList()

-        val builder = context.createIrBuilder(currentScope!!.scope.scopeOwnerSymbol)
        val destArrayDescr = ArrayDescr(irVararg.type, context)
-
        return builder.irComposite(irVararg) {
            // Emit all of the variables first so that all vararg expressions
            // are evaluated only once and in order of their appearance.