[interop] Optimized unmanaged memory alloc/free

kotlin-native/Interop/Runtime/src/jvm/kotlin/kotlinx/cinterop/JvmNativeMem.kt

@@ -16,6 +16,7 @@
 
 package kotlinx.cinterop
 
+import org.jetbrains.kotlin.konan.util.nativeMemoryAllocator
 import sun.misc.Unsafe
 
 private val NativePointed.address: Long
@@ -95,19 +96,20 @@ internal object nativeMemUtils {
         return unsafe.allocateInstance(T::class.java) as T
     }
 
-    fun alloc(size: Long, align: Int): NativePointed {
+    internal fun allocRaw(size: Long, align: Int): NativePtr {
-        val address = unsafe.allocateMemory(
+        val address = unsafe.allocateMemory(size)
-                if (size == 0L) 1L else size // It is a hack: `sun.misc.Unsafe` can't allocate zero bytes
-        )
-
         if (address % align != 0L) TODO(align.toString())
-        return interpretOpaquePointed(address)
+        return address
     }
 
-    fun free(mem: NativePtr) {
+    internal fun freeRaw(mem: NativePtr) {
         unsafe.freeMemory(mem)
     }
 
+    fun alloc(size: Long, align: Int) = interpretOpaquePointed(nativeMemoryAllocator.alloc(size, align))
+
+    fun free(mem: NativePtr) = nativeMemoryAllocator.free(mem)
+
     private val unsafe = with(Unsafe::class.java.getDeclaredField("theUnsafe")) {
         isAccessible = true
         return@with this.get(null) as Unsafe

kotlin-native/Interop/Runtime/src/main/kotlin/kotlinx/cinterop/Utils.kt

@@ -395,7 +395,7 @@ private object EmptyCString: CValues<ByteVar>() {
     override val align get() = 1
 
     private val placement =
-            interpretCPointer<ByteVar>(nativeMemUtils.alloc(1, 1).rawPtr)!!.also {
+            interpretCPointer<ByteVar>(nativeMemUtils.allocRaw(1, 1))!!.also {
                 it[0] = 0.toByte()
             }
 

kotlin-native/Interop/Runtime/src/native/kotlin/kotlinx/cinterop/NativeMem.kt

@@ -118,14 +118,22 @@ internal object nativeMemUtils {
     }
 
     fun alloc(size: Long, align: Int): NativePointed {
+        return interpretOpaquePointed(allocRaw(size, align))
+    }
+
+    fun free(mem: NativePtr) {
+        freeRaw(mem)
+    }
+
+    internal fun allocRaw(size: Long, align: Int): NativePtr {
         val ptr = malloc(size, align)
         if (ptr == nativeNullPtr) {
             throw OutOfMemoryError("unable to allocate native memory")
         }
-        return interpretOpaquePointed(ptr)
+        return ptr
     }
 
-    fun free(mem: NativePtr) {
+    internal fun freeRaw(mem: NativePtr) {
         cfree(mem)
     }
 }

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

@@ -46,9 +46,9 @@ import org.jetbrains.kotlin.ir.symbols.IrSymbol
 import org.jetbrains.kotlin.ir.symbols.impl.IrFieldSymbolImpl
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.konan.library.KonanLibraryLayout
+import org.jetbrains.kotlin.konan.util.disposeNativeMemoryAllocator
 import org.jetbrains.kotlin.library.SerializedIrModule
 import org.jetbrains.kotlin.resolve.descriptorUtil.isEffectivelyExternal
-import org.jetbrains.kotlin.utils.addToStdlib.getOrPut
 
 /**
  * Offset for synthetic elements created by lowerings and not attributable to other places in the source code.
@@ -347,6 +347,14 @@ internal class Context(config: KonanConfig) : KonanBackendContext(config) {
         llvmDisposed = true
     }
 
+    private var nativeMemFreed = false
+
+    fun freeNativeMem() {
+        if (nativeMemFreed) return
+        disposeNativeMemoryAllocator()
+        nativeMemFreed = true
+    }
+
     val cStubsManager = CStubsManager(config.target)
 
     val coverage = CoverageManager(this)

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

@@ -28,7 +28,11 @@ fun runTopLevelPhases(konanConfig: KonanConfig, environment: KotlinCoreEnvironme
     try {
         toplevelPhase.cast<CompilerPhase<Context, Unit, Unit>>().invokeToplevel(context.phaseConfig, context, Unit)
     } finally {
-        context.disposeLlvm()
+        try {
+            context.disposeLlvm()
+        } finally {
+            context.freeNativeMem()
+        }
     }
 }
 

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

@@ -438,7 +438,8 @@ val toplevelPhase: CompilerPhase<*, Unit, Unit> = namedUnitPhase(
                                 unitSink()
                 ) then
                 objectFilesPhase then
-                linkerPhase
+                linkerPhase then
+                freeNativeMemPhase
 )
 
 internal fun PhaseConfig.disableIf(phase: AnyNamedPhase, condition: Boolean) {

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

@@ -74,6 +74,16 @@ internal val disposeLLVMPhase = namedUnitPhase(
         }
 )
 
+internal val freeNativeMemPhase = namedUnitPhase(
+        name = "FreeNativeMem",
+        description = "Free native memory used by interop",
+        lower = object : CompilerPhase<Context, Unit, Unit> {
+            override fun invoke(phaseConfig: PhaseConfig, phaserState: PhaserState<Unit>, context: Context, input: Unit) {
+                context.freeNativeMem()
+            }
+        }
+)
+
 internal val RTTIPhase = makeKonanModuleOpPhase(
         name = "RTTI",
         description = "RTTI generation",

kotlin-native/utilities/basic-utils/src/main/kotlin/NativeMemoryAllocator.kt

@@ -0,0 +1,120 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+package org.jetbrains.kotlin.konan.util
+
+import sun.misc.Unsafe
+
+private val allocatorHolder = ThreadLocal<NativeMemoryAllocator>()
+val nativeMemoryAllocator: NativeMemoryAllocator
+    get() = allocatorHolder.get() ?: NativeMemoryAllocator().also { allocatorHolder.set(it) }
+
+fun disposeNativeMemoryAllocator() {
+    allocatorHolder.get()?.freeAll()
+    allocatorHolder.remove()
+}
+
+// 256 buckets for sizes <= 2048 padded to 8
+// 256 buckets for sizes <= 64KB padded to 256
+// 256 buckets for sizes <= 1MB padded to 4096
+private const val ChunkBucketSize = 256
+// Alignments are such that overhead is approx 10%.
+private const val SmallChunksSizeAlignment = 8
+private const val MediumChunksSizeAlignment = 256
+private const val BigChunksSizeAlignment = 4096
+private const val MaxSmallSize = ChunkBucketSize * SmallChunksSizeAlignment
+private const val MaxMediumSize = ChunkBucketSize * MediumChunksSizeAlignment
+private const val MaxBigSize = ChunkBucketSize * BigChunksSizeAlignment
+private const val ChunkHeaderSize = 2 * Int.SIZE_BYTES // chunk size + alignment hop size.
+
+private const val RawChunkSize: Long = 4 * 1024 * 1024
+
+class NativeMemoryAllocator {
+    private fun alignUp(x: Long, align: Int) = (x + align - 1) and (align - 1).toLong().inv()
+    private fun alignUp(x: Int, align: Int) = (x + align - 1) and (align - 1).inv()
+
+    private val smallChunks = LongArray(ChunkBucketSize)
+    private val mediumChunks = LongArray(ChunkBucketSize)
+    private val bigChunks = LongArray(ChunkBucketSize)
+
+    // Chunk layout: [chunk size,...padding...,diff to start,aligned data start,.....data.....]
+    fun alloc(size: Long, align: Int): Long {
+        val totalChunkSize = ChunkHeaderSize + size + align
+        val ptr = ChunkHeaderSize + when {
+            totalChunkSize <= MaxSmallSize -> allocFromFreeList(totalChunkSize.toInt(), SmallChunksSizeAlignment, smallChunks)
+            totalChunkSize <= MaxMediumSize -> allocFromFreeList(totalChunkSize.toInt(), MediumChunksSizeAlignment, mediumChunks)
+            totalChunkSize <= MaxBigSize -> allocFromFreeList(totalChunkSize.toInt(), BigChunksSizeAlignment, bigChunks)
+            else -> unsafe.allocateMemory(totalChunkSize).also {
+                // The actual size is not used. Just put value bigger than the biggest threshold.
+                unsafe.putInt(it, Int.MAX_VALUE)
+            }
+        }
+        val alignedPtr = alignUp(ptr, align)
+        unsafe.putInt(alignedPtr - Int.SIZE_BYTES, (alignedPtr - ptr).toInt())
+        return alignedPtr
+    }
+
+    private fun allocFromFreeList(size: Int, align: Int, freeList: LongArray): Long {
+        val paddedSize = alignUp(size, align)
+        val index = paddedSize / align - 1
+        val chunk = freeList[index]
+        val ptr = if (chunk == 0L)
+            allocRaw(paddedSize)
+        else {
+            val nextChunk = unsafe.getLong(chunk)
+            freeList[index] = nextChunk
+            chunk
+        }
+        unsafe.putInt(ptr, paddedSize)
+        return ptr
+    }
+
+    private fun freeToFreeList(paddedSize: Int, align: Int, freeList: LongArray, chunk: Long) {
+        require(paddedSize > 0 && paddedSize % align == 0)
+        val index = paddedSize / align - 1
+        unsafe.putLong(chunk, freeList[index])
+        freeList[index] = chunk
+    }
+
+    private val rawChunks = mutableListOf<Long>()
+    private var rawOffset = 0
+
+    private fun allocRaw(size: Int): Long {
+        if (rawChunks.isEmpty() || rawOffset + size > RawChunkSize) {
+            val newRawChunk = unsafe.allocateMemory(RawChunkSize)
+            rawChunks.add(newRawChunk)
+            rawOffset = size
+            return newRawChunk
+        }
+        return (rawChunks.last() + rawOffset).also { rawOffset += size }
+    }
+
+    fun free(mem: Long) {
+        val chunkStart = mem - ChunkHeaderSize - unsafe.getInt(mem - Int.SIZE_BYTES)
+        val chunkSize = unsafe.getInt(chunkStart)
+        when {
+            chunkSize <= MaxSmallSize -> freeToFreeList(chunkSize, SmallChunksSizeAlignment, smallChunks, chunkStart)
+            chunkSize <= MaxMediumSize -> freeToFreeList(chunkSize, MediumChunksSizeAlignment, mediumChunks, chunkStart)
+            chunkSize <= MaxBigSize -> freeToFreeList(chunkSize, BigChunksSizeAlignment, bigChunks, chunkStart)
+            else -> unsafe.freeMemory(chunkStart)
+        }
+    }
+
+    fun freeAll() {
+        for (i in 0 until ChunkBucketSize) {
+            smallChunks[i] = 0L
+            mediumChunks[i] = 0L
+            bigChunks[i] = 0L
+        }
+        for (chunk in rawChunks)
+            unsafe.freeMemory(chunk)
+        rawChunks.clear()
+    }
+
+    private val unsafe = with(Unsafe::class.java.getDeclaredField("theUnsafe")) {
+        isAccessible = true
+        return@with this.get(null) as Unsafe
+    }
+}