Add SetConstantCmd and DispatchIndirectCmd for compute

Author: hyazinthh

src/Aardvark.Rendering.GL/Instructions/OpenGL.fs

@@ -456,6 +456,7 @@ module OpenGl =
         let BindProgram  = getProcAddress "glUseProgram"
 
         let DispatchCompute = getProcAddress "glDispatchCompute"
+        let DispatchComputeIndirect = getProcAddress "glDispatchComputeIndirect"
         let GetInteger = getProcAddress "glGetIntegerv"
         let GetFloat = getProcAddress "glGetFloatv"
         let GetDouble = getProcAddress "glGetDoublev"

src/Aardvark.Rendering.GL/Management/AssemblerExtensions.fs

@@ -55,6 +55,7 @@ type ICommandStream =
     abstract member UseProgram : m : nativeptr<int> -> unit
     abstract member DispatchCompute : gx : int * gy : int * gz : int -> unit
     abstract member DispatchCompute : groups : nativeptr<V3i> -> unit
+    abstract member DispatchComputeIndirect : indirect: nativeint -> unit
 
     abstract member MemoryBarrier : MemoryBarrierFlags -> unit
 
@@ -499,6 +500,10 @@ module GLAssemblerExtensions =
             s.PushIntArg (NativePtr.toNativeInt groups + 0n)
             s.Call OpenGl.Pointers.DispatchCompute
 
+        member this.DispatchComputeIndirect(indirect : nativeint) =
+            s.BeginCall(1)
+            s.PushArg indirect
+            s.Call OpenGl.Pointers.DispatchComputeIndirect
 
         member this.Enable(v : int) =
              s.BeginCall(1)
@@ -1012,6 +1017,7 @@ module GLAssemblerExtensions =
             member this.Disable(v: nativeptr<int>) = this.Disable(v)
             member this.DispatchCompute(gx: int, gy: int, gz: int) = this.DispatchCompute(gx, gy, gz)
             member this.DispatchCompute(groups: nativeptr<V3i>) = this.DispatchCompute(groups)
+            member this.DispatchComputeIndirect(indirect: nativeint) = this.DispatchComputeIndirect(indirect)
             member this.DrawArrays(stats: nativeptr<V2i>, isActive: nativeptr<int>, beginMode: nativeptr<GLBeginMode>, calls: nativeptr<DrawCallInfoList>) = this.DrawArrays(stats, isActive, beginMode, calls)
             member this.DrawArraysIndirect(stats: nativeptr<V2i>, isActive: nativeptr<int>, beginMode: nativeptr<GLBeginMode>, indirect: nativeptr<IndirectDrawArgs>) = this.DrawArraysIndirect(stats, isActive, beginMode, indirect)
             member this.DrawElements(stats: nativeptr<V2i>, isActive: nativeptr<int>, beginMode: nativeptr<GLBeginMode>, indexType: int, calls: nativeptr<DrawCallInfoList>) = this.DrawElements(stats, isActive, beginMode, indexType, calls)
@@ -1135,6 +1141,7 @@ module GLAssemblerExtensions =
             member x.Disable(v: nativeptr<int>) = inner.Disable(v); x.Append("Disable", v)
             member x.DispatchCompute(gx: int, gy: int, gz: int) = inner.DispatchCompute(gx, gy, gz); x.Append("DispatchCompute", gx, gy, gz)
             member x.DispatchCompute(groups: nativeptr<V3i>) = inner.DispatchCompute(groups); x.Append("DispatchCompute", groups)
+            member x.DispatchComputeIndirect(indirect: nativeint) = inner.DispatchComputeIndirect(indirect); x.Append("DispatchComputeIndirect")
             member x.DrawArrays(stats: nativeptr<V2i>, isActive: nativeptr<int>, beginMode: nativeptr<GLBeginMode>, calls: nativeptr<DrawCallInfoList>) = inner.DrawArrays(stats, isActive, beginMode, calls); x.Append("DrawArrays", stats, isActive, beginMode, calls)
             member x.DrawArraysIndirect(stats: nativeptr<V2i>, isActive: nativeptr<int>, beginMode: nativeptr<GLBeginMode>, indirect: nativeptr<IndirectDrawArgs>) = inner.DrawArraysIndirect(stats, isActive, beginMode, indirect); x.Append("DrawArraysIndirect", stats, isActive, beginMode, indirect)
             member x.DrawElements(stats: nativeptr<V2i>, isActive: nativeptr<int>, beginMode: nativeptr<GLBeginMode>, indexType: int, calls: nativeptr<DrawCallInfoList>) = inner.DrawElements(stats, isActive, beginMode, indexType, calls); x.Append("DrawElements", stats, isActive, beginMode, indexType, calls)

src/Aardvark.Rendering.GL/Runtime/Compute.fs

@@ -375,11 +375,22 @@ module internal ComputeTaskInternals =
                                 s.BindImageTexture(slot, TextureAccess.ReadWrite, binding.Pointer)
                         )
 
+                    | ComputeCommand.SetConstantCmd _ ->
+                        raise <| NotSupportedException("Constants are not supported.")
+
                     | ComputeCommand.DispatchCmd groups ->
                         do! CompilerState.assemble (fun _ s ->
                             s.DispatchCompute(groups.X, groups.Y, groups.Z)
                         )
 
+                    | ComputeCommand.DispatchIndirectCmd (indirectBuffer, offset) ->
+                        let indirectBuffer = unbox<GL.Buffer> indirectBuffer
+
+                        do! CompilerState.assemble (fun _ s ->
+                            s.BindBuffer(BufferTarget.DispatchIndirectBuffer, indirectBuffer.Handle)
+                            s.DispatchComputeIndirect(nativeint offset)
+                        )
+
                     | ComputeCommand.ExecuteCmd other ->
                         do! CompilerState.execute other
 

src/Aardvark.Rendering.GL/Runtime/Runtime.fs

@@ -193,6 +193,9 @@ type Runtime(debug : IDebugConfig) =
         member x.CreateInputBinding(shader, inputs) =
             x.CreateInputBinding(shader, inputs)
 
+        member x.GetComputeConstant<'T>(_, _) : IComputeConstant<'T> =
+            raise <| NotSupportedException("Constants are not supported.")
+
         member x.CompileCompute (commands) =
             x.CompileCompute commands
 

src/Aardvark.Rendering.Vulkan/Runtime/ComputeTask.fs

@@ -13,12 +13,42 @@ module internal ComputeTaskInternals =
         {
             Shader         : IComputeShader
             DescriptorSets : INativeResourceLocation<DescriptorSetBinding>
-            PushConstants  : IConstantResourceLocation<PushConstants> voption
         }
 
         interface IComputeInputBinding with
             member x.Shader = x.Shader
 
+    type PushConstant(shader: IComputeShader, name: Symbol, inputType: Type) =
+        let layout =
+            shader.PipelineLayout.PushConstants
+            |> ValueOption.defaultWith (fun _ -> failf "Compute shader does not use any push constants.")
+
+        let field =
+            layout.Buffer.ubFields
+            |> List.tryFindV (fun f -> f.ufName = string name)
+            |> ValueOption.defaultWith (fun _ -> failf $"Compute shader does not use push constant '{name}'.")
+
+        let writer =
+            match UniformWriters.tryGetWriter 0 field.ufType inputType with
+            | Result.Ok writer -> writer
+            | Result.Error msg -> failf $"Cannot get writer for compute constant '{name}': {msg}"
+
+        let size = GLSLType.sizeof field.ufType
+
+        member _.Size = size
+
+        member _.Write(stream: VKVM.CommandStream, value: obj, buffer: nativeint) =
+            writer.WriteUnsafeValue(value, buffer)
+            stream.PushConstants(shader.PipelineLayout.Handle, layout.StageFlags, uint32 field.ufOffset, uint32 size, buffer) |> ignore
+
+        interface IComputeConstant with
+            member _.Shader = shader
+            member _.Name = name
+
+    type PushConstant<'T>(shader: IComputeShader, name: Symbol) =
+        inherit PushConstant(shader, name, typeof<'T>)
+        interface IComputeConstant<'T>
+
     type ResourceManager with
 
         member x.CreateComputeInputBinding(shader : IComputeShader, inputs : IUniformProvider) =
@@ -28,15 +58,8 @@ module internal ComputeTaskInternals =
                 let sets = x.CreateDescriptorSets(shader.PipelineLayout, provider)
                 x.CreateDescriptorSetBinding(VkPipelineBindPoint.Compute, shader.PipelineLayout, sets)
 
-            let pushConstants =
-                shader.PipelineLayout.PushConstants |> ValueOption.map (fun pc ->
-                    x.CreatePushConstants(pc, provider)
-                )
-
             { Shader         = shader
-              DescriptorSets = descriptorSets
-              PushConstants  = pushConstants }
-
+              DescriptorSets = descriptorSets }
 
     [<RequireQualifiedAccess>]
     type private HostCommand =
@@ -59,11 +82,16 @@ module internal ComputeTaskInternals =
 
     type private CompilerState =
         {
-            Commands     : CompiledCommand list
-            UsedImages   : HashSet<Image>
-            ImageLayouts : HashMap<Image, VkImageLayout>
+            Commands        : CompiledCommand list
+            UsedImages      : HashSet<Image>
+            ImageLayouts    : HashMap<Image, VkImageLayout>
+            ConstantBuffers : nativeptr<uint8> list
         }
 
+        member this.Free() =
+            for cmd in this.Commands do cmd.Dispose()
+            for buf in this.ConstantBuffers do NativePtr.free buf
+
     type private ICompiledTask =
         abstract member State : CompilerState
 
@@ -93,9 +121,10 @@ module internal ComputeTaskInternals =
         module private CompilerState =
 
             let empty =
-                { Commands     = []
-                  UsedImages   = HashSet.empty
-                  ImageLayouts = HashMap.empty }
+                { Commands        = []
+                  UsedImages      = HashSet.empty
+                  ImageLayouts    = HashMap.empty
+                  ConstantBuffers = [] }
 
             let stream =
                 State.custom (fun s ->
@@ -127,7 +156,7 @@ module internal ComputeTaskInternals =
                 State.modify (fun s -> { s with UsedImages = s.UsedImages |> HashSet.add image })
 
             let usedImages =
-                State.get |> State.map (fun s -> s.UsedImages)
+                State.get |> State.map _.UsedImages
 
             let inline layout (image : Image) =
                 State.get |> State.map (fun s ->
@@ -150,6 +179,12 @@ module internal ComputeTaskInternals =
                     return oldLayout
                 }
 
+            let inline constantBuffer (constant: PushConstant) =
+                State.custom (fun s ->
+                    let buffer = NativePtr.alloc constant.Size
+                    { s with ConstantBuffers = buffer :: s.ConstantBuffers }, buffer
+                )
+
         [<AutoOpen>]
         module private CommandStreamExtensions =
 
@@ -262,14 +297,22 @@ module internal ComputeTaskInternals =
                         let input = unbox<ComputeInputBinding> input
                         let! stream = CompilerState.stream
                         stream.IndirectBindDescriptorSets(input.DescriptorSets.Pointer) |> ignore
-                        match input.PushConstants with
-                        | ValueSome pc -> stream.PushConstants(input.Shader.PipelineLayout.Handle, pc.Handle) |> ignore
-                        | _ -> ()
+
+                    | ComputeCommand.SetConstantCmd (constant, value) ->
+                        let constant = unbox<PushConstant> constant
+                        let! stream = CompilerState.stream
+                        let! buffer = CompilerState.constantBuffer constant
+                        constant.Write(stream, value, buffer.Address)
 
                     | ComputeCommand.DispatchCmd groups ->
                         let! stream = CompilerState.stream
                         stream.Dispatch(uint32 groups.X, uint32 groups.Y, uint32 groups.Z) |> ignore
 
+                    | ComputeCommand.DispatchIndirectCmd (indirectBuffer, offset) ->
+                        let! stream = CompilerState.stream
+                        let indirectBuffer = indirectBuffer |> unbox<Buffer>
+                        stream.DispatchIndirect(indirectBuffer.Handle, offset) |> ignore
+
                     | ComputeCommand.ExecuteCmd other ->
                         let compiled = unbox<ICompiledTask> other
                         do! restoreLayouts compiled.State.UsedImages
@@ -415,7 +458,6 @@ module internal ComputeTaskInternals =
                             failf "unknown input binding type %A" (input.GetType())
 
                     resources.Add input.DescriptorSets
-                    input.PushConstants |> ValueOption.iter resources.Add
                     inputs.[index] <- input
 
                     ComputeCommand.SetInputCmd input
@@ -469,7 +511,6 @@ module internal ComputeTaskInternals =
                 // This way nothing will be released if the input just moved in the command list
                 for input in removedInputs do
                     resources.Remove input.DescriptorSets
-                    input.PushConstants |> ValueOption.iter resources.Remove
 
                 // Update all hooked compute programs
                 let mutable changed = deltas.Count > 0
@@ -479,7 +520,7 @@ module internal ComputeTaskInternals =
 
                 // Compile updated command list
                 if changed then
-                    for c in compiled.Commands do c.Dispose()
+                    compiled.Free()
                     compiled <- ComputeCommand.compile queueFlags commands
                     true
                 else
@@ -488,7 +529,6 @@ module internal ComputeTaskInternals =
             member x.Dispose() =
                 for KeyValue(_, input) in inputs do
                     resources.Remove input.DescriptorSets
-                    input.PushConstants |> ValueOption.iter resources.Remove
                 inputs.Clear()
 
                 for KeyValue(_, task) in nested do
@@ -497,7 +537,7 @@ module internal ComputeTaskInternals =
 
                 hooked.Clear()
                 commands <- IndexList.empty
-                for c in compiled.Commands do c.Dispose()
+                compiled.Free()
                 compiled <- CompilerState.empty
 
             interface IDisposable with

src/Aardvark.Rendering.Vulkan/Runtime/Runtime.fs

@@ -451,6 +451,9 @@ type Runtime(device : Device) as this =
     member x.CreateInputBinding(shader : IComputeShader, inputs : IUniformProvider) : IComputeInputBinding =
         manager.CreateComputeInputBinding(shader, inputs)
 
+    member x.GetComputeConstant<'T>(shader : IComputeShader, name : Symbol) : IComputeConstant<'T> =
+        PushConstant<'T>(shader, name)
+
     member x.CompileCompute (commands : alist<ComputeCommand>) =
         new ComputeTask(manager, commands) :> IComputeTask
 
@@ -561,6 +564,9 @@ type Runtime(device : Device) as this =
         member x.CreateInputBinding(shader : IComputeShader, inputs : IUniformProvider) =
             x.CreateInputBinding(shader, inputs)
 
+        member x.GetComputeConstant<'T>(shader : IComputeShader, name : Symbol) =
+            x.GetComputeConstant<'T>(shader, name)
+
         member x.CompileCompute (commands : alist<ComputeCommand>) =
             x.CompileCompute commands
 

src/Aardvark.Rendering/Runtime/Compute/Compute.fs

@@ -26,13 +26,21 @@ and IComputeShader =
 and IComputeInputBinding =
     abstract member Shader : IComputeShader
 
+and IComputeConstant =
+    abstract member Shader : IComputeShader
+    abstract member Name : Symbol
+
+and IComputeConstant<'T> =
+    inherit IComputeConstant
+
 and IComputeRuntime =
     inherit IBufferRuntime
     inherit ITextureRuntime
     abstract member ContextLock : IDisposable
     abstract member MaxLocalSize : V3i
     abstract member CreateComputeShader : shader: FShade.ComputeShader -> IComputeShader
     abstract member CreateInputBinding : shader: IComputeShader * inputs: IUniformProvider -> IComputeInputBinding
+    abstract member GetComputeConstant<'T> : shader: IComputeShader * name: Symbol -> IComputeConstant<'T>
     abstract member CompileCompute : commands: alist<ComputeCommand> -> IComputeTask
 
 and [<RequireQualifiedAccess>]
@@ -44,7 +52,9 @@ and [<RequireQualifiedAccess>]
     ComputeCommand =
     | BindCmd                 of shader: IComputeShader
     | SetInputCmd             of input: IComputeInputBinding
+    | SetConstantCmd          of constant: IComputeConstant * data: obj
     | DispatchCmd             of groups: V3i
+    | DispatchIndirectCmd     of indirectBuffer: IBackendBuffer * offset: uint64
     | ExecuteCmd              of task: IComputeTask
     | CopyBufferCmd           of src: IBufferRange * dst: IBufferRange
     | DownloadBufferCmd       of src: IBufferRange * dst: HostMemory
@@ -66,6 +76,9 @@ and [<RequireQualifiedAccess>]
     static member SetInput(input : IComputeInputBinding) =
         ComputeCommand.SetInputCmd input
 
+    static member SetConstant<'T>(constant : IComputeConstant<'T>, value : 'T) =
+        ComputeCommand.SetConstantCmd(constant, value)
+
     static member Dispatch(groups : V3i) =
         ComputeCommand.DispatchCmd groups
 
@@ -75,6 +88,9 @@ and [<RequireQualifiedAccess>]
     static member Dispatch(groups : int) =
         ComputeCommand.DispatchCmd (V3i(groups, 1, 1))
 
+    static member DispatchIndirect(indirectBuffer : IBackendBuffer, [<Optional; DefaultParameterValue(0UL)>] offset : uint64) =
+        ComputeCommand.DispatchIndirectCmd(indirectBuffer, offset)
+
     static member Execute(task : IComputeTask) =
         ComputeCommand.ExecuteCmd task
 

src/Aardvark.Rendering/Runtime/Compute/ComputeExtensions.fs

@@ -63,6 +63,10 @@ type IComputeRuntimeExtensions private() =
     static member Run(runtime : IComputeRuntime, commands : list<ComputeCommand>) =
         runtime.Run(commands, RenderToken.Empty)
 
+    [<Extension>]
+    static member GetComputeConstant<'T>(runtime : IComputeRuntime, shader : IComputeShader, name : string) =
+        runtime.GetComputeConstant<'T>(shader, Sym.ofString name)
+
 
 [<AbstractClass; Sealed; Extension>]
 type IComputeShaderExtensions private() =
@@ -71,6 +75,14 @@ type IComputeShaderExtensions private() =
     static member CreateInputBinding(shader : IComputeShader, inputs : IUniformProvider) =
         shader.Runtime.CreateInputBinding(shader, inputs)
 
+    [<Extension>]
+    static member GetConstant<'T>(shader : IComputeShader, name : Symbol) =
+        shader.Runtime.GetComputeConstant<'T>(shader, name)
+
+    [<Extension>]
+    static member GetConstant<'T>(shader : IComputeShader, name : string) =
+        shader.Runtime.GetComputeConstant<'T>(shader, name)
+
     [<Extension>]
     static member Invoke(shader : IComputeShader, groupCount : V3i, input : IComputeInputBinding, renderToken : RenderToken) =
         shader.Runtime.Run([

src/Tests/Aardvark.Rendering.Tests/Aardvark.Rendering.Tests.fsproj

@@ -77,6 +77,8 @@
     <Compile Include="Tests\Compute\Sorting.fs" />
     <Compile Include="Tests\Compute\Jpeg.fs" />
     <Compile Include="Tests\Compute\MutableInputBinding.fs" />
+    <Compile Include="Tests\Compute\PushConstants.fs" />
+    <Compile Include="Tests\Compute\DispatchIndirect.fs" />
     <Compile Include="Tests\Compute\ComputeTests.fs" />
     <Compile Include="Tests\Other\Camera.fs" />
     <Compile Include="Tests\Other\IndexedGeometryTests.fs" />

src/Tests/Aardvark.Rendering.Tests/Tests/Compute/ComputeTests.fs

@@ -13,6 +13,8 @@ module ``Compute Tests`` =
         ComputeSorting.tests
         ComputeJpeg.tests
         MutableInputBinding.tests
+        PushConstants.tests
+        DispatchIndirect.tests
     ]
 
     [<Tests>]