TemporaryLoop.hs

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

module Tests.Control.TemporaryLoop where

-- base
import Data.Maybe
  ( fromJust )

-- vector-sized
import qualified Data.Vector.Sized as Vector

-- fir
import FIR

------------------------------------------------
-- program

-- A local copy of an array made inside a loop body (in order to index it
-- with a run-time index) must not be reused after the loop:
-- the loop body might not run at all.

table :: Array 4 Float
table = MkArray ( fromJust $ Vector.fromList [ 1, 2, 3, 4 ] )

type Defs =
  '[ "ubo"  ':-> Uniform '[ DescriptorSet 0, Binding 0 ]
                   ( Struct '[ "n" ':-> Word32 ] )
   , "out"  ':-> StorageBuffer '[ DescriptorSet 0, Binding 1 ]
                   ( Struct '[ "res" ':-> Float ] )
   , "main" ':-> EntryPoint '[ LocalSize 1 1 1 ] Compute
   ]

program :: Module Defs
program = Module do

  entryPoint @"main" @Compute do

    let
      arr :: Code ( Array 4 Float )
      arr = Lit table
    n <- use @( Name "ubo" :.: Name "n" )

    _ <- def @"i"   @RW @Word32 0
    _ <- def @"acc" @RW @Float  0
    while ( ( < n ) <<$>> get @"i" ) do
      i <- get @"i"
      modify @"acc" ( + view @( AnIndex ( Code Word32 ) ) i arr )
      put @"i" ( i + 1 )

    acc <- get @"acc"
    assign @( Name "out" :.: Name "res" ) ( acc + view @( AnIndex ( Code Word32 ) ) n arr )