Lazy and concurrent command buffer evaluation (#68)

Author: MarconZet

cyfra-compiler/src/main/scala/io/computenode/cyfra/spirv/SpirvTypes.scala

@@ -54,6 +54,7 @@ private[cyfra] object SpirvTypes:
     case LGBooleanTag       => 4
     case v if v <:< LVecTag =>
       vecSize(v) * typeStride(v.typeArgs.head)
+    case _ => 4
 
   def typeStride(tag: Tag[?]): Int = typeStride(tag.tag)
 

cyfra-core/src/main/scala/io/computenode/cyfra/core/Allocation.scala

@@ -10,6 +10,8 @@ import izumi.reflect.Tag
 import java.nio.ByteBuffer
 
 trait Allocation:
+  def submitLayout[L <: Layout: LayoutBinding](layout: L): Unit
+
   extension (buffer: GBinding[?])
     def read(bb: ByteBuffer, offset: Int = 0): Unit
 

cyfra-core/src/main/scala/io/computenode/cyfra/core/GBufferRegion.scala

@@ -1,16 +1,23 @@
 package io.computenode.cyfra.core
 
 import io.computenode.cyfra.core.Allocation
+import io.computenode.cyfra.core.GBufferRegion.MapRegion
 import io.computenode.cyfra.core.GProgram.BufferLengthSpec
 import io.computenode.cyfra.core.layout.{Layout, LayoutBinding}
 import io.computenode.cyfra.dsl.Value
 import io.computenode.cyfra.dsl.Value.FromExpr
 import io.computenode.cyfra.dsl.binding.GBuffer
 import izumi.reflect.Tag
 
+import scala.util.chaining.given
 import java.nio.ByteBuffer
 
-sealed trait GBufferRegion[ReqAlloc <: Layout: LayoutBinding, ResAlloc <: Layout: LayoutBinding]
+sealed trait GBufferRegion[ReqAlloc <: Layout: LayoutBinding, ResAlloc <: Layout: LayoutBinding]:
+  def reqAllocBinding: LayoutBinding[ReqAlloc] = summon[LayoutBinding[ReqAlloc]]
+  def resAllocBinding: LayoutBinding[ResAlloc] = summon[LayoutBinding[ResAlloc]]
+
+  def map[NewAlloc <: Layout: LayoutBinding](f: Allocation ?=> ResAlloc => NewAlloc): GBufferRegion[ReqAlloc, NewAlloc] =
+    MapRegion(this, (alloc: Allocation) => (resAlloc: ResAlloc) => f(using alloc)(resAlloc))
 
 object GBufferRegion:
 
@@ -24,20 +31,17 @@ object GBufferRegion:
   ) extends GBufferRegion[ReqAlloc, ResAlloc]
 
   extension [ReqAlloc <: Layout: LayoutBinding, ResAlloc <: Layout: LayoutBinding](region: GBufferRegion[ReqAlloc, ResAlloc])
-    def map[NewAlloc <: Layout: LayoutBinding](f: Allocation ?=> ResAlloc => NewAlloc): GBufferRegion[ReqAlloc, NewAlloc] =
-      MapRegion(region, (alloc: Allocation) => (resAlloc: ResAlloc) => f(using alloc)(resAlloc))
-
     def runUnsafe(init: Allocation ?=> ReqAlloc, onDone: Allocation ?=> ResAlloc => Unit)(using cyfraRuntime: CyfraRuntime): Unit =
       cyfraRuntime.withAllocation: allocation =>
 
         // noinspection ScalaRedundantCast
-        val steps: Seq[Allocation => Layout => Layout] = Seq.unfold(region: GBufferRegion[?, ?]):
-          case _: AllocRegion[?] => None
+        val steps: Seq[(Allocation => Layout => Layout, LayoutBinding[Layout])] = Seq.unfold(region: GBufferRegion[?, ?]):
+          case AllocRegion()     => None
           case MapRegion(req, f) =>
-            Some((f.asInstanceOf[Allocation => Layout => Layout], req))
+            Some(((f.asInstanceOf[Allocation => Layout => Layout], req.resAllocBinding.asInstanceOf[LayoutBinding[Layout]]), req))
 
-        val initAlloc = init(using allocation)
+        val initAlloc = init(using allocation).tap(allocation.submitLayout)
         val bodyAlloc = steps.foldLeft[Layout](initAlloc): (acc, step) =>
-          step(allocation)(acc)
+          step._1(allocation)(acc).tap(allocation.submitLayout(_)(using step._2))
 
         onDone(using allocation)(bodyAlloc.asInstanceOf[ResAlloc])

cyfra-examples/src/main/scala/io/computenode/cyfra/samples/TestingStuff.scala

@@ -12,9 +12,16 @@ import io.computenode.cyfra.runtime.VkCyfraRuntime
 import org.lwjgl.BufferUtils
 import org.lwjgl.system.MemoryUtil
 
+import java.nio.ByteBuffer
 import java.util.concurrent.atomic.AtomicInteger
 import scala.collection.parallel.CollectionConverters.given
 
+def printBuffer(bb: ByteBuffer): Unit =
+  val l = bb.asIntBuffer()
+  val a = new Array[Int](l.remaining())
+  l.get(a)
+  println(a.mkString(" "))
+
 object TestingStuff:
 
   given GContext = GContext()
@@ -115,6 +122,7 @@ object TestingStuff:
     )
     runtime.close()
 
+    printBuffer(rbb)
     val actual = (0 until 2 * 1024).map(i => result.get(i * 1) != 0)
     val expected = (0 until 1024).flatMap(x => Seq.fill(emitFilterParams.emitN)(x)).map(_ == emitFilterParams.filterValue)
     expected
@@ -191,7 +199,7 @@ object TestingStuff:
   def testAddProgram10Times =
     given runtime: VkCyfraRuntime = VkCyfraRuntime()
     val bufferSize = 1280
-    val params = AddProgramParams(bufferSize, addA = 0, addB = 1)
+    val params = AddProgramParams(bufferSize, addA = 5, addB = 10)
     val region = GBufferRegion
       .allocate[AddProgramExecLayout]
       .map: region =>
@@ -226,6 +234,8 @@ object TestingStuff:
       },
     )
     runtime.close()
+
+    printBuffer(rbbList(0))
     val expected = inData.map(_ + 11 * (params.addA + params.addB))
     outBuffers.foreach { buf =>
       (0 until bufferSize).foreach { i =>

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

@@ -29,7 +29,7 @@ import io.computenode.cyfra.vulkan.util.Util.{check, pushStack}
 import izumi.reflect.Tag
 import org.lwjgl.vulkan.VK10.*
 import org.lwjgl.vulkan.VK13.{VK_ACCESS_2_SHADER_READ_BIT, VK_ACCESS_2_SHADER_WRITE_BIT, VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT, vkCmdPipelineBarrier2}
-import org.lwjgl.vulkan.{VkCommandBuffer, VkCommandBufferBeginInfo, VkDependencyInfo, VkMemoryBarrier2, VkSubmitInfo}
+import org.lwjgl.vulkan.{VK13, VkCommandBuffer, VkCommandBufferBeginInfo, VkDependencyInfo, VkMemoryBarrier2, VkSubmitInfo}
 
 import scala.collection.mutable
 
@@ -51,7 +51,7 @@ class ExecutionHandler(runtime: VkCyfraRuntime, threadContext: VulkanThreadConte
           .zip(layout)
           .map:
             case (set, bindings) =>
-              set.update(bindings.map(x => VkAllocation.getUnderlying(x.binding)))
+              set.update(bindings.map(x => VkAllocation.getUnderlying(x.binding).buffer))
               set
 
     val dispatches: Seq[Dispatch] = shaderCalls
@@ -67,19 +67,15 @@ class ExecutionHandler(runtime: VkCyfraRuntime, threadContext: VulkanThreadConte
         else (steps.appended(step), dirty ++ bindings)
 
     val commandBuffer = recordCommandBuffer(executeSteps)
-    pushStack: stack =>
-      val pCommandBuffer = stack.callocPointer(1).put(0, commandBuffer)
-      val submitInfo = VkSubmitInfo
-        .calloc(stack)
-        .sType$Default()
-        .pCommandBuffers(pCommandBuffer)
-
-      val fence = new Fence()
-      timed("Vulkan render command"):
-        check(vkQueueSubmit(commandPool.queue.get, submitInfo, fence.get), "Failed to submit command buffer to queue")
-        fence.block().destroy()
-    commandPool.freeCommandBuffer(commandBuffer)
-    descriptorSets.flatten.foreach(dsManager.free)
+    val cleanup = () =>
+      descriptorSets.flatten.foreach(dsManager.free)
+      commandPool.freeCommandBuffer(commandBuffer)
+
+    val externalBindings = getAllBindings(executeSteps).map(VkAllocation.getUnderlying)
+    val deps = externalBindings.flatMap(_.execution.fold(Seq(_), _.toSeq))
+    val pe = new PendingExecution(commandBuffer, deps, cleanup)
+    summon[VkAllocation].addExecution(pe)
+    externalBindings.foreach(_.execution = Left(pe)) // TODO we assume all accesses are read-write
     result
 
   private def interpret[Params, EL <: Layout: LayoutBinding, RL <: Layout: LayoutBinding](
@@ -202,7 +198,6 @@ class ExecutionHandler(runtime: VkCyfraRuntime, threadContext: VulkanThreadConte
       .flags(0)
 
     check(vkBeginCommandBuffer(commandBuffer, commandBufferBeginInfo), "Failed to begin recording command buffer")
-
     steps.foreach:
       case PipelineBarrier =>
         val memoryBarrier = VkMemoryBarrier2 // TODO don't synchronise everything
@@ -228,11 +223,18 @@ class ExecutionHandler(runtime: VkCyfraRuntime, threadContext: VulkanThreadConte
 
         dispatch match
           case Direct(x, y, z)          => vkCmdDispatch(commandBuffer, x, y, z)
-          case Indirect(buffer, offset) => vkCmdDispatchIndirect(commandBuffer, VkAllocation.getUnderlying(buffer).get, offset)
+          case Indirect(buffer, offset) => vkCmdDispatchIndirect(commandBuffer, VkAllocation.getUnderlying(buffer).buffer.get, offset)
 
     check(vkEndCommandBuffer(commandBuffer), "Failed to finish recording command buffer")
     commandBuffer
 
+  private def getAllBindings(steps: Seq[ExecutionStep]): Seq[GBinding[?]] =
+    steps
+      .flatMap:
+        case Dispatch(_, layout, _, _) => layout.flatten.map(_.binding)
+        case PipelineBarrier           => Seq.empty
+      .distinct
+
 object ExecutionHandler:
   case class ShaderCall(pipeline: ComputePipeline, layout: ShaderLayout, dispatch: DispatchType)
 

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

@@ -0,0 +1,118 @@
+package io.computenode.cyfra.runtime
+
+import io.computenode.cyfra.vulkan.command.{CommandPool, Fence, Semaphore}
+import io.computenode.cyfra.vulkan.core.{Device, Queue}
+import io.computenode.cyfra.vulkan.util.Util.{check, pushStack}
+import io.computenode.cyfra.vulkan.util.VulkanObject
+import org.lwjgl.vulkan.VK10.VK_TRUE
+import org.lwjgl.vulkan.VK13.{VK_PIPELINE_STAGE_2_COPY_BIT, vkQueueSubmit2}
+import org.lwjgl.vulkan.{VK13, VkCommandBuffer, VkCommandBufferSubmitInfo, VkSemaphoreSubmitInfo, VkSubmitInfo2}
+
+import scala.collection.mutable
+
+/** A command buffer that is pending execution, along with its dependencies and cleanup actions.
+  *
+  * You can call `close()` only when `isFinished || isPending` is true
+  *
+  * You can call `destroy()` only when all dependants are `isClosed`
+  */
+class PendingExecution(protected val handle: VkCommandBuffer, val dependencies: Seq[PendingExecution], cleanup: () => Unit)(using Device):
+  private val semaphore: Semaphore = Semaphore()
+  private var fence: Option[Fence] = None
+
+  def isPending: Boolean = fence.isEmpty
+  def isRunning: Boolean = fence.exists(f => f.isAlive && !f.isSignaled)
+  def isFinished: Boolean = fence.exists(f => !f.isAlive || f.isSignaled)
+
+  def block(): Unit = fence.foreach(_.block())
+
+  private var closed = false
+  def isClosed: Boolean = closed
+  private def close(): Unit =
+    assert(isFinished || isPending, "Cannot close a PendingExecution that is not finished or pending")
+    if closed then return
+    cleanup()
+    closed = true
+
+  private var destroyed = false
+  def destroy(): Unit =
+    if destroyed then return
+    close()
+    semaphore.destroy()
+    fence.foreach(x => if x.isAlive then x.destroy())
+    destroyed = true
+
+  /** Gathers all command buffers and their semaphores for submission to the queue, in the correct order.
+    *
+    * When you call this method, you are expected to submit the command buffers to the queue, and signal the provided fence when done.
+    * @param f
+    *   The fence to signal when the command buffers are done executing.
+    * @return
+    *   A sequence of tuples, each containing a command buffer, semaphore to signal, and a set of semaphores to wait on.
+    */
+  private def gatherForSubmission(f: Fence): Seq[((VkCommandBuffer, Semaphore), Set[Semaphore])] =
+    if !isPending then return Seq.empty
+    val mySubmission = ((handle, semaphore), dependencies.map(_.semaphore).toSet)
+    fence = Some(f)
+    dependencies.flatMap(_.gatherForSubmission(f)).appended(mySubmission)
+
+object PendingExecution:
+  def executeAll(executions: Seq[PendingExecution], queue: Queue)(using Device): Fence = pushStack: stack =>
+    assert(executions.forall(_.isPending), "All executions must be pending")
+    assert(executions.nonEmpty, "At least one execution must be provided")
+
+    val fence = Fence()
+
+    val exec: Seq[(Set[Semaphore], Set[(VkCommandBuffer, Semaphore)])] =
+      val gathered = executions.flatMap(_.gatherForSubmission(fence))
+      val ordering = gathered.zipWithIndex.map(x => (x._1._1._1, x._2)).toMap
+      gathered.toSet.groupMap(_._2)(_._1).toSeq.sortBy(x => x._2.map(_._1).map(ordering).min)
+
+    val submitInfos = VkSubmitInfo2.calloc(exec.size, stack)
+    exec.foreach: (semaphores, executions) =>
+      val pCommandBuffersSI = VkCommandBufferSubmitInfo.calloc(executions.size, stack)
+      val signalSemaphoreSI = VkSemaphoreSubmitInfo.calloc(executions.size, stack)
+      executions.foreach: (cb, s) =>
+        pCommandBuffersSI
+          .get()
+          .sType$Default()
+          .commandBuffer(cb)
+          .deviceMask(0)
+        signalSemaphoreSI
+          .get()
+          .sType$Default()
+          .semaphore(s.get)
+          .stageMask(VK13.VK_PIPELINE_STAGE_2_COPY_BIT | VK13.VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT)
+
+      pCommandBuffersSI.flip()
+      signalSemaphoreSI.flip()
+
+      val waitSemaphoreSI = VkSemaphoreSubmitInfo.calloc(semaphores.size, stack)
+      semaphores.foreach: s =>
+        waitSemaphoreSI
+          .get()
+          .sType$Default()
+          .semaphore(s.get)
+          .stageMask(VK13.VK_PIPELINE_STAGE_2_COPY_BIT | VK13.VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT)
+
+      waitSemaphoreSI.flip()
+
+      submitInfos
+        .get()
+        .sType$Default()
+        .flags(0)
+        .pCommandBufferInfos(pCommandBuffersSI)
+        .pSignalSemaphoreInfos(signalSemaphoreSI)
+        .pWaitSemaphoreInfos(waitSemaphoreSI)
+
+    submitInfos.flip()
+
+    check(vkQueueSubmit2(queue.get, submitInfos, fence.get), "Failed to submit command buffer to queue")
+    fence
+
+  def cleanupAll(executions: Seq[PendingExecution]): Unit =
+    def cleanupRec(ex: PendingExecution): Unit =
+      if !ex.isClosed then return
+      ex.close()
+      ex.dependencies.foreach(cleanupRec)
+    executions.foreach(cleanupRec)