JBAI-4017 [ndarray, core, utils] Typealias to marker arrays as StateM…

…arker, changed PlatformQueue to ConcurrentQueue, commented code cleaned up

inference/inference-core/src/commonMain/kotlin/io/kinference.core/operators/layer/attention/Attention.kt

@@ -65,7 +65,7 @@ sealed class Attention(name: String, info: OperatorInfo, attributes: Map<String,
             val vMarker = v.array.marker
 
             val resultBlocks: Array<FloatArray>
-            val resultMarker: Array<() -> Unit>
+            val resultMarker: Array<StateMarker>
 
             if (past == null || past.linearSize == 0) {
                 resultBlocks = kBlocks.plus(vBlocks)
@@ -84,7 +84,7 @@ sealed class Attention(name: String, info: OperatorInfo, attributes: Map<String,
 
                 val rowsSize = batchSize * numHeads
                 val futureRes = arrayOfNulls<FloatArray>(2 * batchSize * numHeads * presentDims[3] * blocksInRow)
-                val futureResMarker = arrayOfNulls<() -> Unit>(2 * batchSize * numHeads * presentDims[3] * blocksInRow)
+                val futureResMarker = arrayOfNulls<StateMarker>(2 * batchSize * numHeads * presentDims[3] * blocksInRow)
 
                 var resBlockIdx = 0
                 var pastBlocIdx = 0
@@ -109,7 +109,7 @@ sealed class Attention(name: String, info: OperatorInfo, attributes: Map<String,
                     }
                 }
                 resultBlocks = futureRes as Array<FloatArray>
-                resultMarker = futureResMarker as Array<() -> Unit>
+                resultMarker = futureResMarker as Array<StateMarker>
             }
 
             return FloatNDArray(FloatTiledArray(resultBlocks, resultMarker), Strides(presentDims))

ndarray/ndarray-api/src/commonMain/kotlin/io/kinference/ndarray/arrays/ArrayDispatcherUtils.kt

@@ -1,5 +1,7 @@
 package io.kinference.ndarray.arrays
 
+typealias StateMarker = () -> Unit
+
 enum class ArrayTypes(val index: Int) {
     ByteArray(0),
     UByteArray(1),

ndarray/ndarray-core/src/commonMain/kotlin/io/kinference/ndarray/arrays/memory/ArrayDispatcher.kt

@@ -1,13 +1,13 @@
 package io.kinference.ndarray.arrays.memory
 
 import io.kinference.ndarray.arrays.ArrayTypes
-import io.kinference.utils.PlatformQueue
+import io.kinference.utils.ConcurrentQueue
 
 internal object ArrayDispatcher {
     private const val INIT_SIZE_VALUE: Int = 2
     private val typeSize: Int = ArrayTypes.entries.size
 
-    private val unusedArrays: PlatformQueue<ArrayStorage> = PlatformQueue()
+    private val unusedArrays: ConcurrentQueue<ArrayStorage> = ConcurrentQueue()
 
     fun getStorage(): ArrayStorage {
         return unusedArrays.removeFirstOrNull() ?: ArrayStorage(typeSize, INIT_SIZE_VALUE)

ndarray/ndarray-core/src/commonMain/kotlin/io/kinference/ndarray/arrays/tiled/PrimitiveTiledArray.kt

@@ -17,7 +17,7 @@ import kotlin.math.min
 
 @GenerateNameFromPrimitives
 @MakePublic
-internal class PrimitiveTiledArray(val blocks: Array<PrimitiveArray>, val marker: Array<()->Unit> = emptyMarker) {
+internal class PrimitiveTiledArray(val blocks: Array<PrimitiveArray>, val marker: Array<StateMarker> = emptyMarker) {
     val size: Int
     val blockSize: Int = if (blocks.isEmpty()) 0 else blocks.first().size
     val blocksNum: Int = blocks.size
@@ -28,7 +28,7 @@ internal class PrimitiveTiledArray(val blocks: Array<PrimitiveArray>, val marker
 
     companion object {
         val type: ArrayTypes = ArrayTypes.valueOf(PrimitiveArray::class.simpleName!!)
-        private val emptyMarker: Array<()->Unit> = arrayOf()
+        private val emptyMarker: Array<StateMarker> = arrayOf()
 
         suspend operator fun invoke(strides: Strides): PrimitiveTiledArray {
             val blockSize = blockSizeByStrides(strides)

ndarray/ndarray-core/src/commonMain/kotlin/io/kinference/ndarray/extensions/gelu/BiasGeluPrimitive.kt

@@ -25,23 +25,10 @@ internal suspend fun computeGeluPrimitive(input: PrimitiveNDArray, bias: Primiti
 
     val blockSize = input.array.blockSize
 
-//    // This approach when arrays acquired before parallelizeByBlocks() is faster
-//    val coroutineContext = coroutineContext[ModelContext.Key]!!
-//    val modelName = coroutineContext.modelName
-//    val inferenceCycle = coroutineContext.cycleId
-//
-//    val coroutineCount = countCoroutinesByData(blockSize, inputBlocks.size, 2048)
-//    val containerTemporaryBlockArrays = ArrayDispatcher.getArrayContainers(PrimitiveTiledArray.type, blockSize, coroutineCount)
-//    val containerTemporaryBlockAbsArrays = ArrayDispatcher.getArrayContainers(PrimitiveTiledArray.type, blockSize, coroutineCount)
-//    val temporaryBlockArrays = Array(containerTemporaryBlockArrays.size) { i -> (containerTemporaryBlockArrays[i] as PrimitiveArrayContainer).array }
-//    val temporaryBlockAbsArrays = Array(containerTemporaryBlockAbsArrays.size) { i -> (containerTemporaryBlockAbsArrays[i] as PrimitiveArrayContainer).array }
-
     // Constant 2048 was precomputed on M1 Max processor
     // With this constant two launches work faster than single thread without launches
     // TODO: (cupertank) Remove constants
     parallelizeByBlocks(blockSize, inputBlocks.size, 2048) { blockStart, blockEnd, coroutineIndex ->
-//        val temporaryBlock = temporaryBlockArrays[coroutineIndex]
-//        val temporaryBlockAbs = temporaryBlockAbsArrays[coroutineIndex]
         val temporaryBlock = PrimitiveArray(blockSize)
         val temporaryBlockAbs = PrimitiveArray(blockSize)
 

ndarray/ndarray-core/src/commonMain/kotlin/io/kinference/ndarray/extensions/gelu/FastGeluPrimitive.kt

@@ -29,15 +29,10 @@ internal suspend fun fastGeluPrimitive(input: PrimitiveNDArray, bias: PrimitiveN
 
     val blockSize = input.array.blockSize
 
-//    val coroutineCount = countCoroutinesByData(blockSize, inputBlocks.size, 2048)
-//    val containerArray = ArrayDispatcher.getArrayContainers(PrimitiveTiledArray.type, blockSize, coroutineCount)
-//    val temporaryBlockExpArrays = Array(containerArray.size) { i -> (containerArray[i] as PrimitiveArrayContainer).array }
-
     // Constant 2048 was precomputed on M1 Max processor
     // With this constant two launches work faster than single thread without launches
     // TODO: (cupertank) Remove constants
-    parallelizeByBlocks(blockSize, inputBlocks.size, 2048) { blockStart, blockEnd, coroutineIndex ->
-//        val temporaryBlockExp = temporaryBlockExpArrays[coroutineIndex]
+    parallelizeByBlocks(blockSize, inputBlocks.size, 2048) { blockStart, blockEnd, _ ->
         val temporaryBlockExp = PrimitiveArray(blockSize)
         for (blockIdx in blockStart until blockEnd) {
             val outputBlock = outputBlocks[blockIdx]

utils/utils-common/src/commonMain/kotlin/io/kinference/utils/PlatformQueue.kt → utils/utils-common/src/commonMain/kotlin/io/kinference/utils/ConcurrentQueue.kt

@@ -1,6 +1,6 @@
 package io.kinference.utils
 
-expect class PlatformQueue<T>() {
+expect class ConcurrentQueue<T>() {
     fun removeFirstOrNull(): T?
     fun addLast(element: T)
 }

utils/utils-common/src/jsMain/kotlin/io/kinference/utils/PlatformQueue.kt → utils/utils-common/src/jsMain/kotlin/io/kinference/utils/ConcurrentQueue.kt

@@ -1,6 +1,6 @@
 package io.kinference.utils
 
-actual class PlatformQueue<T> actual constructor() {
+actual class ConcurrentQueue<T> actual constructor() {
     private val queue: ArrayDeque<T> = ArrayDeque()
 
     actual fun removeFirstOrNull(): T? {

utils/utils-common/src/jvmMain/kotlin/io/kinference/utils/PlatformQueue.kt → utils/utils-common/src/jvmMain/kotlin/io/kinference/utils/ConcurrentQueue.kt

@@ -2,7 +2,7 @@ package io.kinference.utils
 
 import java.util.concurrent.ConcurrentLinkedQueue
 
-actual class PlatformQueue<T> actual constructor() {
+actual class ConcurrentQueue<T> actual constructor() {
     private val queue: ConcurrentLinkedQueue<T> = ConcurrentLinkedQueue()
 
     actual fun removeFirstOrNull(): T? {