LoopsNested.hs

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

module Tests.Control.LoopsNested where

-- fir
import FIR

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

-- Nested loops in a function, with bindings that are:
--   - modified in the inner loop ("acc", "m"),
--   - modified in the outer loop only ("j"),
--   - never modified ("k"),
--   - only modified just before breaking out of the outer loop ("b").

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

program :: Module Defs
program = Module do

  ( f :: Code Word32 -> Code Float ) <- fundef @"f" do
    _ <- def @"acc" @RW @Float  0
    _ <- def @"j"   @RW @Word32 0
    _ <- def @"m"   @RW @Word32 0
    _ <- def @"k"   @RW @Float  3
    _ <- def @"b"   @RW @Float  0
    u <- get @"u"
    while ( ( < u ) <<$>> get @"j" ) do
      j <- get @"j"
      put @"m" 0
      while ( ( < j ) <<$>> get @"m" ) do
        m <- get @"m"
        k <- get @"k"
        modify @"acc" ( + ( fromIntegral m * k ) )
        put @"m" ( m + 1 )
      acc <- get @"acc"
      if acc > 100
      then do
        put @"b" acc
        break @1
      else put @"j" ( j + 1 )
    acc <- get @"acc"
    b   <- get @"b"
    k   <- get @"k"
    pure ( acc + b + k )

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