Support JEP-370 foreign memory

.github/workflows/build.yml

@@ -13,16 +13,13 @@ jobs:
     runs-on: ${{matrix.os}}
     timeout-minutes: 40
     steps:
-      - name: Checkout the repo
-        uses: actions/checkout@v2
-      - name: Set up JDK 11
-        uses: DeLaGuardo/setup-graalvm@4.0
+      - uses: actions/checkout@v2
+      - uses: DeLaGuardo/setup-graalvm@4.0
         with:
           graalvm: 21.2.0
-          java: java11
+          java: java16
           arch: amd64
-      - name: Cache gradle
-        uses: actions/cache@v2
+      - uses: actions/cache@v2
         with:
           path: |
             ~/.gradle/caches
@@ -30,12 +27,10 @@ jobs:
           key: ${{ runner.os }}-gradle-${{ hashFiles('*.gradle.kts') }}
           restore-keys: |
             ${{ runner.os }}-gradle-
-      - name: Cache konan
-        uses: actions/cache@v2
+      - uses: actions/cache@v2
         with:
           path: ~/.konan
           key: ${{ runner.os }}-gradle-${{ hashFiles('*.gradle.kts') }}
           restore-keys: |
             ${{ runner.os }}-gradle-
-      - name: Build
-        run: ./gradlew build --build-cache --no-daemon --stacktrace
+      - run: ./gradlew build --build-cache --no-daemon --stacktrace

.github/workflows/pages.yml

@@ -13,11 +13,19 @@ jobs:
       - uses: DeLaGuardo/setup-graalvm@4.0
         with:
           graalvm: 21.2.0
-          java: java11
+          java: java16
           arch: amd64
       - uses: actions/cache@v2
         with:
-          path: ~/.gradle/caches
+          path: |
+            ~/.gradle/caches
+            ~/.gradle/wrapper
+          key: ${{ runner.os }}-gradle-${{ hashFiles('*.gradle.kts') }}
+          restore-keys: |
+            ${{ runner.os }}-gradle-
+      - uses: actions/cache@v2
+        with:
+          path: ~/.konan
           key: ${{ runner.os }}-gradle-${{ hashFiles('*.gradle.kts') }}
           restore-keys: |
             ${{ runner.os }}-gradle-

.github/workflows/publish.yml

@@ -14,16 +14,13 @@ jobs:
         os: [ macOS-latest, windows-latest ]
     runs-on: ${{matrix.os}}
     steps:
-      - name: Checkout the repo
-        uses: actions/checkout@v2
-      - name: Set up JDK 11
-        uses: DeLaGuardo/setup-graalvm@4.0
+      - uses: actions/checkout@v2
+      - uses: DeLaGuardo/setup-graalvm@4.0
         with:
           graalvm: 21.2.0
-          java: java11
+          java: java16
           arch: amd64
-      - name: Cache gradle
-        uses: actions/cache@v2
+      - uses: actions/cache@v2
         with:
           path: |
             ~/.gradle/caches
@@ -31,22 +28,19 @@ jobs:
           key: ${{ runner.os }}-gradle-${{ hashFiles('*.gradle.kts') }}
           restore-keys: |
             ${{ runner.os }}-gradle-
-      - name: Cache konan
-        uses: actions/cache@v2
+      - uses: actions/cache@v2
         with:
           path: ~/.konan
           key: ${{ runner.os }}-gradle-${{ hashFiles('*.gradle.kts') }}
           restore-keys: |
             ${{ runner.os }}-gradle-
-      - name: Publish Windows Artifacts
-        if: matrix.os == 'windows-latest'
+      - if: matrix.os == 'windows-latest'
         shell: cmd
         run: >
           ./gradlew release --no-daemon --build-cache -Ppublishing.enabled=true
           -Ppublishing.space.user=${{ secrets.SPACE_APP_ID }}
           -Ppublishing.space.token=${{ secrets.SPACE_APP_SECRET }}
-      - name: Publish Mac Artifacts
-        if: matrix.os == 'macOS-latest'
+      - if: matrix.os == 'macOS-latest'
         run: >
           ./gradlew release --no-daemon --build-cache -Ppublishing.enabled=true -Ppublishing.platform=macosX64
           -Ppublishing.space.user=${{ secrets.SPACE_APP_ID }}

gradle.properties

@@ -4,6 +4,7 @@
 #
 
 kotlin.code.style=official
+kotlin.jupyter.add.scanner=false
 kotlin.mpp.enableGranularSourceSetsMetadata=true
 kotlin.mpp.stability.nowarn=true
 kotlin.native.enableDependencyPropagation=false

kmath-jupyter/src/main/kotlin/space/kscience/kmath/jupyter/KMathJupyter.kt

@@ -11,7 +11,6 @@ import kotlinx.html.stream.createHTML
 import kotlinx.html.unsafe
 import org.jetbrains.kotlinx.jupyter.api.DisplayResult
 import org.jetbrains.kotlinx.jupyter.api.HTML
-import org.jetbrains.kotlinx.jupyter.api.annotations.JupyterLibrary
 import org.jetbrains.kotlinx.jupyter.api.libraries.JupyterIntegration
 import space.kscience.kmath.ast.rendering.FeaturedMathRendererWithPostProcess
 import space.kscience.kmath.ast.rendering.MathMLSyntaxRenderer

kmath-memory/api/kmath-memory.api

@@ -70,3 +70,8 @@ public abstract interface class space/kscience/kmath/memory/MemoryWriter {
 	public abstract fun writeShort (IS)V
 }
 
+public final class space/kscience/kmath/memory/foreign/ForeignMemoryKt {
+	public static final fun allocateAsForeign (Lspace/kscience/kmath/memory/Memory$Companion;I)Lspace/kscience/kmath/memory/Memory;
+	public static final fun wrapAsForeign (Lspace/kscience/kmath/memory/Memory$Companion;[B)Lspace/kscience/kmath/memory/Memory;
+}
+

kmath-memory/build.gradle.kts

@@ -10,3 +10,7 @@ readme {
         An API and basic implementation for arranging objects in a continuous memory block.
     """.trimIndent()
 }
+
+tasks.jvmTest {
+    jvmArgs("--add-modules", "jdk.incubator.foreign")
+}

kmath-memory/src/jvmMain/kotlin/space/kscience/kmath/memory/foreign/ForeignMemory.kt

@@ -0,0 +1,54 @@
+/*
+ * Copyright 2018-2021 KMath contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+ */
+
+package space.kscience.kmath.memory.foreign
+
+import jdk.incubator.foreign.MemorySegment
+import space.kscience.kmath.memory.Memory
+import space.kscience.kmath.memory.MemoryReader
+import space.kscience.kmath.memory.MemoryWriter
+import java.lang.ref.Cleaner
+
+/**
+ * Allocates memory using JDK Foreign Memory API. It should be even faster than default ByteBuffer memory provided by
+ * [space.kscience.kmath.memory.allocate].
+ */
+public fun Memory.Companion.allocateAsForeign(length: Int): Memory =
+    ForeignMemory(MemorySegment.allocateNative(length.toLong()))
+
+/**
+ * Wraps a [Memory] around existing [ByteArray]. This operation is unsafe since the array is not copied
+ * and could be mutated independently from the resulting [Memory].
+ *
+ * The memory is wrapped to JDK Foreign Memory segment.
+ */
+public fun Memory.Companion.wrapAsForeign(array: ByteArray): Memory = ForeignMemory(MemorySegment.ofArray(array))
+
+private val cleaner: Cleaner by lazy { Cleaner.create() }
+
+private fun cleaningRunnable(scope: MemorySegment): Runnable = Runnable { scope.close() }
+
+internal class ForeignMemory(val scope: MemorySegment) : Memory, AutoCloseable {
+    private val cleanable: Cleaner.Cleanable = cleaner.register(this, cleaningRunnable(scope))
+
+    override val size: Int
+        get() = Math.toIntExact(scope.byteSize())
+
+    private val writer: MemoryWriter = ForeignWriter(this)
+    private val reader: MemoryReader = ForeignReader(this)
+
+    override fun view(offset: Int, length: Int): ForeignMemory =
+        ForeignMemory(scope.asSlice(offset.toLong(), length.toLong()))
+
+    override fun copy(): Memory {
+        val newScope = MemorySegment.allocateNative(scope.byteSize())
+        newScope.copyFrom(scope)
+        return ForeignMemory(newScope)
+    }
+
+    override fun reader(): MemoryReader = reader
+    override fun writer(): MemoryWriter = writer
+    override fun close(): Unit = cleanable.clean()
+}

kmath-memory/src/jvmMain/kotlin/space/kscience/kmath/memory/foreign/ForeignReader.kt

@@ -0,0 +1,23 @@
+/*
+ * Copyright 2018-2021 KMath contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+ */
+
+package space.kscience.kmath.memory.foreign
+
+import jdk.incubator.foreign.MemoryAccess
+import jdk.incubator.foreign.MemorySegment
+import space.kscience.kmath.memory.MemoryReader
+
+internal class ForeignReader(override val memory: ForeignMemory) : MemoryReader {
+    private val scope: MemorySegment
+        get() = memory.scope
+
+    override fun readDouble(offset: Int): Double = MemoryAccess.getDoubleAtOffset(scope, offset.toLong())
+    override fun readFloat(offset: Int): Float = MemoryAccess.getFloatAtOffset(scope, offset.toLong())
+    override fun readByte(offset: Int): Byte = MemoryAccess.getByteAtOffset(scope, offset.toLong())
+    override fun readShort(offset: Int): Short = MemoryAccess.getShortAtOffset(scope, offset.toLong())
+    override fun readInt(offset: Int): Int = MemoryAccess.getIntAtOffset(scope, offset.toLong())
+    override fun readLong(offset: Int): Long = MemoryAccess.getLongAtOffset(scope, offset.toLong())
+    override fun release(): Unit = Unit
+}

kmath-memory/src/jvmMain/kotlin/space/kscience/kmath/memory/foreign/ForeignWriter.kt

@@ -0,0 +1,30 @@
+/*
+ * Copyright 2018-2021 KMath contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+ */
+
+package space.kscience.kmath.memory.foreign
+
+import jdk.incubator.foreign.MemoryAccess
+import jdk.incubator.foreign.MemorySegment
+import space.kscience.kmath.memory.MemoryWriter
+
+internal class ForeignWriter(override val memory: ForeignMemory) : MemoryWriter {
+    private val scope: MemorySegment
+        get() = memory.scope
+
+    override fun writeDouble(offset: Int, value: Double): Unit =
+        MemoryAccess.setDoubleAtOffset(scope, offset.toLong(), value)
+
+    override fun writeFloat(offset: Int, value: Float): Unit =
+        MemoryAccess.setFloatAtOffset(scope, offset.toLong(), value)
+
+    override fun writeByte(offset: Int, value: Byte): Unit = MemoryAccess.setByteAtOffset(scope, offset.toLong(), value)
+
+    override fun writeShort(offset: Int, value: Short): Unit =
+        MemoryAccess.setShortAtOffset(scope, offset.toLong(), value)
+
+    override fun writeInt(offset: Int, value: Int): Unit = MemoryAccess.setIntAtOffset(scope, offset.toLong(), value)
+    override fun writeLong(offset: Int, value: Long): Unit = MemoryAccess.setLongAtOffset(scope, offset.toLong(), value)
+    override fun release(): Unit = Unit
+}

kmath-memory/src/jvmTest/kotlin/space/kscience/kmath/memory/ByteBufferMemoryTest.kt

@@ -0,0 +1,56 @@
+/*
+ * Copyright 2018-2021 KMath contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+ */
+
+package space.kscience.kmath.memory
+
+import kotlin.test.Test
+import kotlin.test.assertEquals
+
+internal class ByteBufferMemoryTest {
+    private fun getMemory(int: Int) = Memory.allocate(int)
+
+    @Test
+    fun size() {
+        val mem = getMemory(66666)
+        assertEquals(66666, mem.size)
+    }
+
+    @Test
+    fun view() {
+        val mem = getMemory(4242)
+        val sub = mem.view(10, 10)
+        sub.write { writeInt(0, 1000000) }
+        assertEquals(10, sub.size)
+        assertEquals(1000000, mem.read { readInt(10) })
+        assertEquals(1000000, sub.read { readInt(0) })
+    }
+
+    @Test
+    fun copy() {
+        val mem = getMemory(8)
+        mem.write { writeDouble(0, 12.0) }
+        val copy = mem.copy()
+        assertEquals(12.0, copy.read { readDouble(0) })
+    }
+
+    @Test
+    fun reader() {
+        val mem = getMemory(8)
+        val rd = mem.reader()
+        assertEquals(0, rd.readLong(0))
+        rd.release()
+    }
+
+    @Test
+    fun writer() {
+        val mem = getMemory(4)
+        val wr = mem.writer()
+        wr.writeFloat(0, 6f)
+        val rd = mem.reader()
+        assertEquals(6f, rd.readFloat(0))
+        rd.release()
+        wr.release()
+    }
+}

kmath-memory/src/jvmTest/kotlin/space/kscience/kmath/memory/ForeignMemoryTest.kt

@@ -0,0 +1,57 @@
+/*
+ * Copyright 2018-2021 KMath contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+ */
+
+package space.kscience.kmath.memory
+
+import space.kscience.kmath.memory.foreign.allocateAsForeign
+import kotlin.test.Test
+import kotlin.test.assertEquals
+
+internal class ForeignMemoryTest {
+    private fun getMemory(int: Int) = Memory.allocateAsForeign(int)
+
+    @Test
+    fun size() {
+        val mem = getMemory(66666)
+        assertEquals(66666, mem.size)
+    }
+
+    @Test
+    fun view() {
+        val mem = getMemory(4242)
+        val sub = mem.view(10, 10)
+        sub.write { writeInt(0, 1000000) }
+        assertEquals(10, sub.size)
+        assertEquals(1000000, mem.read { readInt(10) })
+        assertEquals(1000000, sub.read { readInt(0) })
+    }
+
+    @Test
+    fun copy() {
+        val mem = getMemory(8)
+        mem.write { writeDouble(0, 12.0) }
+        val copy = mem.copy()
+        assertEquals(12.0, copy.read { readDouble(0) })
+    }
+
+    @Test
+    fun reader() {
+        val mem = getMemory(8)
+        val rd = mem.reader()
+        assertEquals(0, rd.readLong(0))
+        rd.release()
+    }
+
+    @Test
+    fun writer() {
+        val mem = getMemory(4)
+        val wr = mem.writer()
+        wr.writeFloat(0, 6f)
+        val rd = mem.reader()
+        assertEquals(6f, rd.readFloat(0))
+        rd.release()
+        wr.release()
+    }
+}