FunctionArgument.hs

{-# LANGUAGE BlockArguments      #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE PolyKinds           #-}
{-# LANGUAGE RebindableSyntax    #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeOperators       #-}

module Tests.Layout.FunctionArgument where

-- fir
import FIR

------------------------------------------------
-- passing a value loaded from a buffer to a function

type S = Struct '[ "a" ':-> Float, "b" ':-> Array 2 Float ]

type Defs =
  '[ "ubo"  ':-> Uniform '[ DescriptorSet 0, Binding 0 ]
                   ( Struct '[ "s" ':-> S ] )
   , "out"  ':-> StorageBuffer '[ DescriptorSet 0, Binding 1 ]
                   ( Struct '[ "res" ':-> Float ] )
   , "f"    ':-> Function '[ "s" ':-> Var R S ] Float
   , "main" ':-> EntryPoint '[ LocalSize 1 1 1 ] Compute
   ]

program :: Module Defs
program = Module do

  ( f :: Code S -> Code Float ) <- fundef @"f" do
    s <- get @"s"
    pure ( view @( Name "a" ) s + view @( Name "b" :.: Index 1 ) s )

  entryPoint @"main" @Compute do
    s <- use @( Name "ubo" :.: Name "s" )
    assign @( Name "out" :.: Name "res" ) ( f s )