diff --git a/kmath-noa/README.md b/kmath-noa/README.md index a0d1ecb9f..c9e20eed4 100644 --- a/kmath-noa/README.md +++ b/kmath-noa/README.md @@ -48,7 +48,7 @@ To load the native library you will need to add to the VM options: The library is under active development. Many more features will be available soon. -### Linear Algebra +### Tensors and Linear Algebra We implement the tensor algebra interfaces from [kmath-tensors](../kmath-tensors): @@ -70,12 +70,31 @@ NoaFloat { tensorReg.save("tensorReg.pt") } ``` + The saved tensor can be loaded in `C++` or in `python`: ```python import torch tensor_reg = list(torch.jit.load('tensorReg.pt').parameters())[0] ``` +The most efficient way pass data between the `JVM` and the native backend +is to rely on primitive arrays: +```kotlin +val array = (1..8).map { 100f * it }.toFloatArray() +val updateArray = floatArrayOf(15f, 20f) +val resArray = NoaFloat { + val tensor = copyFromArray(array, intArrayOf(2, 2, 2)) + NoaFloat { + // The call `tensor[0]` creates a native tensor instance pointing to a slice of `tensor` + // The second call `[1]` is a setter call and does not create any new instances + tensor[0][1] = updateArray + // The instance `tensor[0]` is destroyed as we move out of the scope + }!! // if the computation fails the result fill be null + tensor.copyToArray() + // the instance `tensor` is destroyed here +}!! + +``` ### Automatic Differentiation The [AutoGrad](https://pytorch.org/tutorials/beginner/blitz/autograd_tutorial.html) diff --git a/kmath-noa/src/main/java/space/kscience/kmath/noa/JNoa.java b/kmath-noa/src/main/java/space/kscience/kmath/noa/JNoa.java index 8255725e7..166eeba6c 100644 --- a/kmath-noa/src/main/java/space/kscience/kmath/noa/JNoa.java +++ b/kmath-noa/src/main/java/space/kscience/kmath/noa/JNoa.java @@ -329,4 +329,29 @@ class JNoa { public static native void saveTensor(long tensorHandle, String path); public static native void saveJitModule(long jitModuleHandle, String path); + + public static native void assignBlobDouble(long tensorHandle, double[] data); + + public static native void assignBlobFloat(long tensorHandle, float[] data); + + public static native void assignBlobLong(long tensorHandle, long[] data); + + public static native void assignBlobInt(long tensorHandle, int[] data); + + public static native void setBlobDouble(long tensorHandle, int i, double[] data); + + public static native void setBlobFloat(long tensorHandle, int i, float[] data); + + public static native void setBlobLong(long tensorHandle, int i, long[] data); + + public static native void setBlobInt(long tensorHandle, int i, int[] data); + + public static native void getBlobDouble(long tensorHandle, double[] data); + + public static native void getBlobFloat(long tensorHandle, float[] data); + + public static native void getBlobLong(long tensorHandle, long[] data); + + public static native void getBlobInt(long tensorHandle, int[] data); + } diff --git a/kmath-noa/src/main/kotlin/space/kscience/kmath/noa/algebras.kt b/kmath-noa/src/main/kotlin/space/kscience/kmath/noa/algebras.kt index 01d87715e..ef553cdea 100644 --- a/kmath-noa/src/main/kotlin/space/kscience/kmath/noa/algebras.kt +++ b/kmath-noa/src/main/kotlin/space/kscience/kmath/noa/algebras.kt @@ -22,6 +22,7 @@ protected constructor(protected val scope: NoaScope) : protected abstract fun wrap(tensorHandle: TensorHandle): TensorType + @PerformancePitfall public fun Tensor.cast(): TensorType = tensor /** @@ -38,8 +39,7 @@ protected constructor(protected val scope: NoaScope) : public abstract fun randDiscrete(low: Long, high: Long, shape: IntArray, device: Device = Device.CPU): TensorType - @PerformancePitfall - public abstract fun Tensor.copyToArray(): PrimitiveArray + public abstract fun TensorType.copyToArray(): PrimitiveArray public abstract fun copyFromArray(array: PrimitiveArray, shape: IntArray, device: Device = Device.CPU): TensorType @@ -164,6 +164,10 @@ protected constructor(protected val scope: NoaScope) : public infix fun TensorType.swap(other: TensorType): Unit = JNoa.swapTensors(tensorHandle, other.tensorHandle) + public abstract fun TensorType.assignFromArray(array: PrimitiveArray): Unit + + public abstract operator fun TensorType.set(i: Int, array: PrimitiveArray): Unit + } public sealed class NoaPartialDivisionAlgebra> @@ -345,9 +349,11 @@ protected constructor(scope: NoaScope) : override fun wrap(tensorHandle: TensorHandle): NoaDoubleTensor = NoaDoubleTensor(scope = scope, tensorHandle = tensorHandle) - @PerformancePitfall - override fun Tensor.copyToArray(): DoubleArray = - tensor.elements().map { it.second }.toList().toDoubleArray() + override fun NoaDoubleTensor.copyToArray(): DoubleArray { + val array = DoubleArray(numElements) + JNoa.getBlobDouble(tensorHandle, array) + return array + } override fun copyFromArray(array: DoubleArray, shape: IntArray, device: Device): NoaDoubleTensor = wrap(JNoa.fromBlobDouble(array, shape, device.toInt())) @@ -405,6 +411,13 @@ protected constructor(scope: NoaScope) : override fun loadTensor(path: String, device: Device): NoaDoubleTensor = wrap(JNoa.loadTensorDouble(path, device.toInt())) + + override fun NoaDoubleTensor.assignFromArray(array: DoubleArray): Unit = + JNoa.assignBlobDouble(tensorHandle, array) + + override fun NoaDoubleTensor.set(i: Int, array: DoubleArray): Unit = + JNoa.setBlobDouble(tensorHandle, i, array) + } public sealed class NoaFloatAlgebra @@ -426,9 +439,11 @@ protected constructor(scope: NoaScope) : override fun wrap(tensorHandle: TensorHandle): NoaFloatTensor = NoaFloatTensor(scope = scope, tensorHandle = tensorHandle) - @PerformancePitfall - override fun Tensor.copyToArray(): FloatArray = - tensor.elements().map { it.second }.toList().toFloatArray() + override fun NoaFloatTensor.copyToArray(): FloatArray { + val res = FloatArray(numElements) + JNoa.getBlobFloat(tensorHandle, res) + return res + } override fun copyFromArray(array: FloatArray, shape: IntArray, device: Device): NoaFloatTensor = wrap(JNoa.fromBlobFloat(array, shape, device.toInt())) @@ -486,6 +501,13 @@ protected constructor(scope: NoaScope) : override fun loadTensor(path: String, device: Device): NoaFloatTensor = wrap(JNoa.loadTensorFloat(path, device.toInt())) + + override fun NoaFloatTensor.assignFromArray(array: FloatArray): Unit = + JNoa.assignBlobFloat(tensorHandle, array) + + override fun NoaFloatTensor.set(i: Int, array: FloatArray): Unit = + JNoa.setBlobFloat(tensorHandle, i, array) + } public sealed class NoaLongAlgebra @@ -507,9 +529,11 @@ protected constructor(scope: NoaScope) : override fun wrap(tensorHandle: TensorHandle): NoaLongTensor = NoaLongTensor(scope = scope, tensorHandle = tensorHandle) - @PerformancePitfall - override fun Tensor.copyToArray(): LongArray = - tensor.elements().map { it.second }.toList().toLongArray() + override fun NoaLongTensor.copyToArray(): LongArray { + val array = LongArray(numElements) + JNoa.getBlobLong(tensorHandle, array) + return array + } override fun copyFromArray(array: LongArray, shape: IntArray, device: Device): NoaLongTensor = wrap(JNoa.fromBlobLong(array, shape, device.toInt())) @@ -553,6 +577,12 @@ protected constructor(scope: NoaScope) : override fun loadTensor(path: String, device: Device): NoaLongTensor = wrap(JNoa.loadTensorLong(path, device.toInt())) + override fun NoaLongTensor.assignFromArray(array: LongArray): Unit = + JNoa.assignBlobLong(tensorHandle, array) + + override fun NoaLongTensor.set(i: Int, array: LongArray): Unit = + JNoa.setBlobLong(tensorHandle, i, array) + } public sealed class NoaIntAlgebra @@ -574,9 +604,11 @@ protected constructor(scope: NoaScope) : override fun wrap(tensorHandle: TensorHandle): NoaIntTensor = NoaIntTensor(scope = scope, tensorHandle = tensorHandle) - @PerformancePitfall - override fun Tensor.copyToArray(): IntArray = - tensor.elements().map { it.second }.toList().toIntArray() + override fun NoaIntTensor.copyToArray(): IntArray { + val array = IntArray(numElements) + JNoa.getBlobInt(tensorHandle, array) + return array + } override fun copyFromArray(array: IntArray, shape: IntArray, device: Device): NoaIntTensor = wrap(JNoa.fromBlobInt(array, shape, device.toInt())) @@ -620,4 +652,10 @@ protected constructor(scope: NoaScope) : override fun loadTensor(path: String, device: Device): NoaIntTensor = wrap(JNoa.loadTensorInt(path, device.toInt())) + override fun NoaIntTensor.assignFromArray(array: IntArray): Unit = + JNoa.assignBlobInt(tensorHandle, array) + + override fun NoaIntTensor.set(i: Int, array: IntArray): Unit = + JNoa.setBlobInt(tensorHandle, i, array) + } diff --git a/kmath-noa/src/main/resources/space_kscience_kmath_noa_JNoa.h b/kmath-noa/src/main/resources/space_kscience_kmath_noa_JNoa.h index 72828f6b9..34db64490 100644 --- a/kmath-noa/src/main/resources/space_kscience_kmath_noa_JNoa.h +++ b/kmath-noa/src/main/resources/space_kscience_kmath_noa_JNoa.h @@ -1247,6 +1247,102 @@ JNIEXPORT void JNICALL Java_space_kscience_kmath_noa_JNoa_saveTensor JNIEXPORT void JNICALL Java_space_kscience_kmath_noa_JNoa_saveJitModule (JNIEnv *, jclass, jlong, jstring); +/* + * Class: space_kscience_kmath_noa_JNoa + * Method: assignBlobDouble + * Signature: (J[D)V + */ +JNIEXPORT void JNICALL Java_space_kscience_kmath_noa_JNoa_assignBlobDouble + (JNIEnv *, jclass, jlong, jdoubleArray); + +/* + * Class: space_kscience_kmath_noa_JNoa + * Method: assignBlobFloat + * Signature: (J[F)V + */ +JNIEXPORT void JNICALL Java_space_kscience_kmath_noa_JNoa_assignBlobFloat + (JNIEnv *, jclass, jlong, jfloatArray); + +/* + * Class: space_kscience_kmath_noa_JNoa + * Method: assignBlobLong + * Signature: (J[J)V + */ +JNIEXPORT void JNICALL Java_space_kscience_kmath_noa_JNoa_assignBlobLong + (JNIEnv *, jclass, jlong, jlongArray); + +/* + * Class: space_kscience_kmath_noa_JNoa + * Method: assignBlobInt + * Signature: (J[I)V + */ +JNIEXPORT void JNICALL Java_space_kscience_kmath_noa_JNoa_assignBlobInt + (JNIEnv *, jclass, jlong, jintArray); + +/* + * Class: space_kscience_kmath_noa_JNoa + * Method: setBlobDouble + * Signature: (JI[D)V + */ +JNIEXPORT void JNICALL Java_space_kscience_kmath_noa_JNoa_setBlobDouble + (JNIEnv *, jclass, jlong, jint, jdoubleArray); + +/* + * Class: space_kscience_kmath_noa_JNoa + * Method: setBlobFloat + * Signature: (JI[F)V + */ +JNIEXPORT void JNICALL Java_space_kscience_kmath_noa_JNoa_setBlobFloat + (JNIEnv *, jclass, jlong, jint, jfloatArray); + +/* + * Class: space_kscience_kmath_noa_JNoa + * Method: setBlobLong + * Signature: (JI[J)V + */ +JNIEXPORT void JNICALL Java_space_kscience_kmath_noa_JNoa_setBlobLong + (JNIEnv *, jclass, jlong, jint, jlongArray); + +/* + * Class: space_kscience_kmath_noa_JNoa + * Method: setBlobInt + * Signature: (JI[I)V + */ +JNIEXPORT void JNICALL Java_space_kscience_kmath_noa_JNoa_setBlobInt + (JNIEnv *, jclass, jlong, jint, jintArray); + +/* + * Class: space_kscience_kmath_noa_JNoa + * Method: getBlobDouble + * Signature: (J[D)V + */ +JNIEXPORT void JNICALL Java_space_kscience_kmath_noa_JNoa_getBlobDouble + (JNIEnv *, jclass, jlong, jdoubleArray); + +/* + * Class: space_kscience_kmath_noa_JNoa + * Method: getBlobFloat + * Signature: (J[F)V + */ +JNIEXPORT void JNICALL Java_space_kscience_kmath_noa_JNoa_getBlobFloat + (JNIEnv *, jclass, jlong, jfloatArray); + +/* + * Class: space_kscience_kmath_noa_JNoa + * Method: getBlobLong + * Signature: (J[J)V + */ +JNIEXPORT void JNICALL Java_space_kscience_kmath_noa_JNoa_getBlobLong + (JNIEnv *, jclass, jlong, jlongArray); + +/* + * Class: space_kscience_kmath_noa_JNoa + * Method: getBlobInt + * Signature: (J[I)V + */ +JNIEXPORT void JNICALL Java_space_kscience_kmath_noa_JNoa_getBlobInt + (JNIEnv *, jclass, jlong, jintArray); + #ifdef __cplusplus } #endif diff --git a/kmath-noa/src/test/kotlin/space/kscience/kmath/noa/TestTensor.kt b/kmath-noa/src/test/kotlin/space/kscience/kmath/noa/TestTensor.kt index a5a0387a0..4ecd203a0 100644 --- a/kmath-noa/src/test/kotlin/space/kscience/kmath/noa/TestTensor.kt +++ b/kmath-noa/src/test/kotlin/space/kscience/kmath/noa/TestTensor.kt @@ -38,13 +38,27 @@ internal fun NoaInt.testingViewWithNoCopy(device: Device = Device.CPU) { assertEquals(tensor[intArrayOf(0)], 10) } -internal fun NoaFloat.testingTensorSerialisation(tensorPath: String, device: Device = Device.CPU){ +internal fun NoaFloat.testingSerialisation(tensorPath: String, device: Device = Device.CPU) { val tensor = copyFromArray(floatArrayOf(45.5f, 98.6f), intArrayOf(2), device) tensor.save(tensorPath) val loadedTensor = loadTensor(tensorPath, device) assertTrue(tensor.copyToArray() contentEquals loadedTensor.copyToArray()) } +internal fun NoaFloat.testingArrayTransfer(device: Device = Device.CPU) { + val array = (1..8).map { 100f * it }.toFloatArray() + val updateArray = floatArrayOf(15f, 20f) + val resArray = NoaFloat { + val tensor = copyFromArray(array, intArrayOf(2, 2, 2), device) + NoaFloat { + tensor[0][1] = updateArray + }!! + assertEquals(tensor[intArrayOf(0, 1, 0)], 15f) + tensor.copyToArray() + }!! + assertEquals(resArray[3], 20f) +} + class TestTensor { private val resources = File("").resolve("src/test/resources") @@ -89,9 +103,16 @@ class TestTensor { }!! @Test - fun tensorSerialisation() = NoaFloat { + fun testSerialisation() = NoaFloat { withCuda { device -> - testingTensorSerialisation(tensorPath, device) + testingSerialisation(tensorPath, device) + } + }!! + + @Test + fun testArrayTransfer() = NoaFloat { + withCuda { device -> + testingArrayTransfer(device) } }!!