Refactor vulkan buffer gmem executor

cyfra-e2e-test/src/test/scala/io/computenode/cyfra/vulkan/SequenceExecutorTest.scala

@@ -1,12 +1,14 @@
 package io.computenode.cyfra.vulkan
 
-
 import io.computenode.cyfra.vulkan.compute.{Binding, ComputePipeline, InputBufferSize, LayoutInfo, LayoutSet, Shader}
 import io.computenode.cyfra.vulkan.executor.BufferAction.{LoadFrom, LoadTo}
 import io.computenode.cyfra.vulkan.executor.SequenceExecutor
 import io.computenode.cyfra.vulkan.executor.SequenceExecutor.{ComputationSequence, Compute, Dependency, LayoutLocation}
+import io.computenode.cyfra.vulkan.memory.Buffer
 import munit.FunSuite
 import org.lwjgl.BufferUtils
+import org.lwjgl.vulkan.VK10.* 
+import org.lwjgl.util.vma.Vma.* 
 
 class SequenceExecutorTest extends FunSuite:
   private val vulkanContext = VulkanContext(true)
@@ -24,10 +26,31 @@ class SequenceExecutorTest extends FunSuite:
     )
     val sequenceExecutor = new SequenceExecutor(sequence, vulkanContext)
     val input = 0 until 1024
-    val buffer = BufferUtils.createByteBuffer(input.length * 4)
-    input.foreach(buffer.putInt)
-    buffer.flip()
-    val res = sequenceExecutor.execute(Seq(buffer), input.length)
-    val output = input.map(_ => res.head.getInt)
+
+    val inputBuffer = new Buffer(
+      input.length * 4, // 4 bytes per int
+      VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
+      VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
+      VMA_MEMORY_USAGE_CPU_ONLY,
+      vulkanContext.allocator
+    )
+
+    val mappedBuffer = inputBuffer.map()
+    input.foreach(mappedBuffer.putInt)
+    inputBuffer.unmap()
+
+    val res = sequenceExecutor.execute(Seq(inputBuffer), input.length)
+
+    val outputMappedBuffer = res.head.map()
+    val output = (0 until input.length).map(_ => outputMappedBuffer.getInt)
+    res.head.unmap()
 
     assertEquals(input.map(_ + 20000).toList, output.toList)
+
+    // Clean up
+    inputBuffer.destroy()
+    res.foreach(_.destroy())
+    sequenceExecutor.destroy()
+    copy1.destroy()
+    copy2.destroy()
+    shader.destroy()

cyfra-foton/src/main/scala/io/computenode/cyfra/foton/animation/AnimatedFunctionRenderer.scala

@@ -32,9 +32,10 @@ class AnimatedFunctionRenderer(params: AnimatedFunctionRenderer.Parameters) exte
 
   protected override def renderFrame(scene: AnimatedFunction, time: Float32, fn: RenderFn): Array[fRGBA] =
     val mem = Array.fill(params.width * params.height)((0.5f, 0.5f, 0.5f, 0.5f))
-    UniformContext.withUniform(AnimationIteration(time)):
+    val uniformStruct = AnimationIteration(time) 
+    UniformContext.withUniform(uniformStruct): 
       val fmem = Vec4FloatMem(mem)
-      fmem.map(fn).asInstanceOf[Vec4FloatMem].toArray
+      fmem.map(uniformStruct, fn).asInstanceOf[Vec4FloatMem].toArray
 
   protected override def renderFunction(scene: AnimatedFunction): RenderFn = 
     GFunction.from2D(params.width, {

cyfra-foton/src/main/scala/io/computenode/cyfra/foton/rt/ImageRtRenderer.scala

@@ -34,10 +34,11 @@ class ImageRtRenderer(params: ImageRtRenderer.Parameters) extends RtRenderer(par
   private def render(scene: Scene, fn: GFunction[RaytracingIteration, Vec4[Float32], Vec4[Float32]]): LazyList[Array[fRGBA]] =
     val initialMem = Array.fill(params.width * params.height)((0.5f, 0.5f, 0.5f, 0.5f))
     LazyList.iterate((initialMem, 0), params.iterations + 1) { case (mem, render) =>
-      UniformContext.withUniform(RaytracingIteration(render)):
+      val uniformStruct = RaytracingIteration(render) 
+      UniformContext.withUniform(uniformStruct): 
         val fmem = Vec4FloatMem(mem)
         val result = timed(s"Rendered iteration $render")(
-          fmem.map(fn).asInstanceOf[Vec4FloatMem].toArray
+          fmem.map(uniformStruct, fn).asInstanceOf[Vec4FloatMem].toArray
         )
         (result, render + 1)
     }.drop(1).map(_._1)

cyfra-foton/src/main/scala/io/computenode/cyfra/foton/rt/animation/AnimationRtRenderer.scala

@@ -29,9 +29,10 @@ class AnimationRtRenderer(params: AnimationRtRenderer.Parameters) extends RtRend
   ): Array[fRGBA] =
     val initialMem = Array.fill(params.width * params.height)((0.5f, 0.5f, 0.5f, 0.5f))
     List.iterate((initialMem, 0), params.iterations + 1) { case (mem, render) =>
-      UniformContext.withUniform(RaytracingIteration(render, time)):
+      val uniformStruct = RaytracingIteration(render, time) 
+      UniformContext.withUniform(uniformStruct): 
         val fmem = Vec4FloatMem(mem)
-        val result = fmem.map(fn).asInstanceOf[Vec4FloatMem].toArray
+        val result = fmem.map(uniformStruct, fn).asInstanceOf[Vec4FloatMem].toArray
         (result, render + 1)
     }.map(_._1).last
 

cyfra-runtime/src/main/scala/io/computenode/cyfra/runtime/GContext.scala

@@ -11,85 +11,140 @@ import SequenceExecutor.*
 import io.computenode.cyfra.runtime.mem.GMem.totalStride
 import io.computenode.cyfra.spirv.SpirvTypes.typeStride
 import io.computenode.cyfra.spirv.compilers.DSLCompiler
-import io.computenode.cyfra.spirv.compilers.ExpressionCompiler.{UniformStructRef, WorkerIndex}
+import io.computenode.cyfra.spirv.compilers.ExpressionCompiler 
+import io.computenode.cyfra.dsl.Expression.E
 import mem.{FloatMem, GMem, Vec4FloatMem}
 import org.lwjgl.system.{Configuration, MemoryUtil}
 import izumi.reflect.Tag
-
-import java.io.FileOutputStream
+import io.computenode.cyfra.vulkan.memory.Buffer
+import org.lwjgl.vulkan.VK10.*
+import org.lwjgl.util.vma.Vma.*
 import java.nio.ByteBuffer
+import java.io.{FileOutputStream, IOException}
 import java.nio.channels.FileChannel
-import java.util.concurrent.Executors
+import scala.collection.mutable 
+import scala.collection.mutable.ListBuffer
 import scala.concurrent.{ExecutionContext, ExecutionContextExecutor}
 
 
-class GContext:
-
-  Configuration.STACK_SIZE.set(1024) // fix lwjgl stack size
+class GContext(debug: Boolean = false):
+  val vkContext = VulkanContext(debug)
+  private val pipelineCache = mutable.Map[Any, ComputePipeline]()
 
-  val vkContext = new VulkanContext(enableValidationLayers = true)
+  private def createPipeline[G <: GStruct[G] : GStructSchema, H <: Value : Tag : FromExpr, R <: Value : Tag : FromExpr](
+    function: GFunction[G, H, R]
+  ): ComputePipeline = {
+    val uniformStructSchemaImpl = summon[GStructSchema[G]]
+    val tagGImpl: Tag[G] = uniformStructSchemaImpl.structTag 
 
-  implicit val ec: ExecutionContextExecutor = ExecutionContext.fromExecutor(Executors.newFixedThreadPool(16))
-
-  def compile[
-    G <: GStruct[G] : Tag : GStructSchema,
-    H <: Value : Tag : FromExpr,
-    R <: Value : Tag : FromExpr
-  ](function: GFunction[G, H, R]): ComputePipeline = {
-    val uniformStructSchema = summon[GStructSchema[G]]
-    val uniformStruct = uniformStructSchema.fromTree(UniformStructRef)
+    val uniformStruct = uniformStructSchemaImpl.fromTree(
+      ExpressionCompiler.UniformStructRef[G](using tagGImpl).asInstanceOf[E[G]]
+    )
     val tree = function
-      .fn
+      .fn 
       .apply(
         uniformStruct,
-        WorkerIndex,
+        ExpressionCompiler.WorkerIndex, 
         GArray[H](0)
       )
-    val shaderCode = DSLCompiler.compile(tree, function.arrayInputs, function.arrayOutputs, uniformStructSchema)
+    val shaderCode = DSLCompiler.compile(tree, function.arrayInputs, function.arrayOutputs, uniformStructSchemaImpl)
     dumpSpvToFile(shaderCode, "program.spv") // TODO remove before release
-    val inOut = 0 to 1 map (Binding(_, InputBufferSize(typeStride(summon[Tag[H]]))))
-    val uniform = Option.when(uniformStructSchema.fields.nonEmpty)(Binding(2, UniformSize(totalStride(uniformStructSchema))))
-    val layoutInfo = LayoutInfo(Seq(LayoutSet(0, inOut ++ uniform)))
+
+    val inputBinding = Binding(0, InputBufferSize(typeStride(summon[Tag[H]])))
+    val outputBinding = Binding(1, InputBufferSize(typeStride(summon[Tag[R]])))
+
+    val uniformBindingOpt = Option.when(uniformStructSchemaImpl.fields.nonEmpty)(
+      Binding(2, UniformSize(GMem.totalStride(uniformStructSchemaImpl)))
+    )
+
+    val bindings = Seq(inputBinding, outputBinding) ++ uniformBindingOpt.toSeq
+    val layoutInfo = LayoutInfo(Seq(LayoutSet(0, bindings)))
+
     val shader = new Shader(shaderCode, new org.joml.Vector3i(256, 1, 1), layoutInfo, "main", vkContext.device)
     new ComputePipeline(shader, vkContext)
   }
 
   private def dumpSpvToFile(code: ByteBuffer, path: String): Unit =
-    val fc: FileChannel = new FileOutputStream("program.spv").getChannel
-    fc.write(code)
-    fc.close()
-    code.rewind()
+    try {
+      val fc: FileChannel = new FileOutputStream(path).getChannel
+      fc.write(code)
+      fc.close()
+    } catch {
+      case e: IOException => e.printStackTrace()
+    } finally {
+      code.rewind()
+    }
 
   def execute[
     G <: GStruct[G] : Tag : GStructSchema,
-    H <: Value,
-    R <: Value
-  ](mem: GMem[H], fn: GFunction[?, H, R])(using uniformContext: UniformContext[_]): GMem[R] =
-    val isUniformEmpty = uniformContext.uniform.schema.fields.isEmpty
-    val actions = Map(
-      LayoutLocation(0, 0) -> BufferAction.LoadTo,
-      LayoutLocation(0, 1) -> BufferAction.LoadFrom
-    ) ++ (
-      if isUniformEmpty then Map.empty 
-      else Map(LayoutLocation(0, 2) -> BufferAction.LoadTo)
-    )
-    val sequence = ComputationSequence(Seq(Compute(fn.pipeline, actions)), Seq.empty)
-    val executor = new SequenceExecutor(sequence, vkContext)
+    H <: Value : Tag : FromExpr, 
+    R <: Value : FromExpr : Tag 
+  ](mem: GMem[H], uniformStruct: G, fn: GFunction[G, H, R]): GMem[R] = {
+    val pipeline = pipelineCache.getOrElseUpdate(fn.fn, createPipeline(fn))
+
+    val sourceBuffersForExecutor = ListBuffer[Buffer]()
+    val bufferActions = mutable.Map[LayoutLocation, BufferAction]()
+
+    bufferActions.put(LayoutLocation(0, 0), BufferAction.LoadTo)
+    sourceBuffersForExecutor.addOne(mem.vulkanBuffer)
+
+    bufferActions.put(LayoutLocation(0, 1), BufferAction.LoadFrom) 
+
+    var uniformStagingBufferOpt: Option[Buffer] = None
+    val uniformStructSchema = summon[GStructSchema[G]]
+    if (uniformStructSchema.fields.nonEmpty) {
+      val uniformCPUByteBuffer = GMem.serializeUniform(uniformStruct)
+      val uniformStagingVkBuffer = new Buffer(
+        uniformCPUByteBuffer.remaining(), // Changed from .toLong to direct Int, or .toInt if remaining() can exceed Int
+        VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
+        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
+        VMA_MEMORY_USAGE_CPU_ONLY,
+        vkContext.allocator
+      )
+      uniformStagingVkBuffer.map { mappedUniform =>
+        mappedUniform.put(uniformCPUByteBuffer)
+      }
+
+      uniformStagingBufferOpt = Some(uniformStagingVkBuffer)
+      bufferActions.put(LayoutLocation(0, 2), BufferAction.LoadTo)
+      sourceBuffersForExecutor.addOne(uniformStagingVkBuffer)
+    }
+
+    val computeStep = Compute(pipeline, bufferActions.toMap)
+    val sequence = ComputationSequence(Seq(computeStep), dependencies = Nil) 
+    val sequenceExecutor = new SequenceExecutor(sequence, vkContext) 
+
+    val outputVulkanBuffers = sequenceExecutor.execute(sourceBuffersForExecutor.toSeq, mem.size)
 
-    val data = mem.toReadOnlyBuffer
-    val inData =
-      if isUniformEmpty then Seq(data)
-      else Seq(data, GMem.serializeUniform(uniformContext.uniform))
-    val out = executor.execute(inData, mem.size)
-    executor.destroy()
-
-    val outTags = fn.arrayOutputs
-    assert(outTags.size == 1)
-
-    outTags.head match
-      case t if t == Tag[Float32] =>
-        new FloatMem(mem.size, out.head).asInstanceOf[GMem[R]]
-      case t if t == Tag[Vec4[Float32]] =>
-        new Vec4FloatMem(mem.size, out.head).asInstanceOf[GMem[R]]
-      case _ => assert(false, "Supported output types are Float32 and Vec4[Float32]")
+    uniformStagingBufferOpt.foreach(_.destroy())
+
+    if (outputVulkanBuffers.isEmpty) {
+      throw new IllegalStateException("SequenceExecutor did not return an output buffer.")
+    }
+    val resultVulkanBuffer = outputVulkanBuffers.head
+
+    val tagR = summon[Tag[R]]
+    val resultMem = 
+      if (tagR.tag =:= Tag[Float32].tag) { 
+        new FloatMem(mem.size, resultVulkanBuffer).asInstanceOf[GMem[R]]
+      } else if (tagR.tag =:= Tag[Vec4[Float32]].tag) { 
+        new Vec4FloatMem(mem.size, resultVulkanBuffer).asInstanceOf[GMem[R]]
+      } else {
+        resultVulkanBuffer.destroy()
+        throw new UnsupportedOperationException(s"Cannot create GMem for result type ${tagR.tag}. Output buffer has been destroyed.")
+      }
+    resultMem
+  }
+
+  def execute[H <: Value : Tag : FromExpr, R <: Value : FromExpr : Tag]( 
+    mem: GMem[H],
+    fn: GFunction[GStruct.Empty, H, R]
+  ): GMem[R] =
+    execute[GStruct.Empty, H, R](mem, GStruct.Empty(), fn) 
+
+  def cleanup(): Unit = {
+    pipelineCache.values.foreach(_.destroy()) 
+    pipelineCache.clear()
+    vkContext.destroy() 
+  }
 

cyfra-runtime/src/main/scala/io/computenode/cyfra/runtime/GFunction.scala

@@ -1,38 +1,39 @@
 package io.computenode.cyfra.runtime
 
-import io.computenode.cyfra.dsl.{*, given}
+import io.computenode.cyfra.dsl.{*, given} 
 import io.computenode.cyfra.dsl.Value.Int32
-import io.computenode.cyfra.vulkan.compute.ComputePipeline
+import io.computenode.cyfra.dsl.Expression.E 
 import izumi.reflect.Tag
 
 case class GFunction[
-  G <: GStruct[G] : GStructSchema : Tag, 
-  H <: Value : Tag : FromExpr, 
+  G <: GStruct[G] : GStructSchema : Tag,
+  H <: Value : Tag : FromExpr,
   R <: Value : Tag : FromExpr
-](fn: (G, Int32, GArray[H]) => R)(implicit context: GContext){
+](
+  val fn: (G, Int32, GArray[H]) => R 
+) {
   def arrayInputs: List[Tag[_]] = List(summon[Tag[H]])
   def arrayOutputs: List[Tag[_]] = List(summon[Tag[R]])
-  val pipeline: ComputePipeline = context.compile(this)
 }
 
 object GFunction:
   def apply[
     H <: Value : Tag : FromExpr,
     R <: Value : Tag : FromExpr
-  ](fn: H => R)(using context: GContext): GFunction[GStruct.Empty, H, R] =
+  ](userSimpleFn: H => R): GFunction[GStruct.Empty, H, R] =
     new GFunction[GStruct.Empty, H, R](
-      (_, index: Int32, gArray: GArray[H]) => fn(gArray.at(index))
+      (_: GStruct.Empty, workerIdx: Int32, gArray: GArray[H]) => userSimpleFn(gArray.at(workerIdx))
     )
 
   def from2D[
     G <: GStruct[G] : GStructSchema : Tag,
     H <: Value : Tag : FromExpr,
     R <: Value : Tag : FromExpr
-  ](width: Int, fn: (G, (Int32, Int32), GArray2D[H]) => R)(using context: GContext): GFunction[G, H, R] =
-    GFunction[G, H, R](
-      (g: G, index: Int32, a: GArray[H]) =>
+  ](width: Int, userFn2D: (G, (Int32, Int32), GArray2D[H]) => R): GFunction[G, H, R] =
+    new GFunction[G, H, R](
+      (g: G, index: Int32, garray: GArray[H]) =>
         val x: Int32 = index mod width
         val y: Int32 = index / width
-        val arr = GArray2D(width, a)
-        fn(g, (x, y), arr)
+        val arr2d = GArray2D(width, garray)
+        userFn2D(g, (x, y), arr2d)
     )