spectral-norm Haskell GHC #4 program

--
-- The Computer Language Benchmarks Game
-- https://salsa.debian.org/benchmarksgame-team/benchmarksgame/
--
-- Modified by Gabriel Gonzalez
-- Modified by Ryan Trinkle: 1) change from divInt# to uncheckedIShiftRA#
--                           2) changed -optc-O to -optc-O3
--                           3) added -optc-ffast-math
-- Translation from Clean by Don Stewart
-- Parallelized by Louis Wasserman
--
-- Should be compiled with:
-- 	-threaded -funbox-strict-fields -O2 -fvia-c -optc-O3 
-- 	-fexcess-precision -optc-ffast-math 
-- Should be run with:
-- 	+RTS -N<number of cores>

import System.Environment
import Foreign.Marshal.Array
import Foreign
import Text.Printf
import Control.Concurrent
import Control.Monad
import GHC.Base
import GHC.Conc

type Reals = Ptr Double

main = do
    n <- getArgs >>= readIO . head
    allocaArray n $ \ u -> allocaArray n $ \ v -> do
      forM_ [0..n-1] $ \i -> pokeElemOff u i 1 >> pokeElemOff v i 0

      powerMethod 10 n u v
      printf "%.9f\n" =<< eigenvalue n u v 0 0 0

------------------------------------------------------------------------

eigenvalue :: Int -> Reals -> Reals -> Int -> Double -> Double -> IO Double
eigenvalue !n !u !v !i !vBv !vv
    | i < n     = do	ui <- peekElemOff u i
			vi <- peekElemOff v i
			eigenvalue n u v (i+1) (vBv + ui * vi) (vv + vi * vi)
    | otherwise = return $! sqrt $! vBv / vv

------------------------------------------------------------------------

-- Essentially borrowed from the Java implementation.
data CyclicBarrier = Cyclic !Int !(MVar (Int, [MVar ()]))

await :: CyclicBarrier -> IO ()
await (Cyclic k waitsVar) = do
	(x, waits) <- takeMVar waitsVar
	if x <= 1 then do
		mapM_ (`putMVar` ()) waits
		putMVar waitsVar (k, [])
	  else do
	  	var <- newEmptyMVar
	  	putMVar waitsVar (x-1,var:waits)
	  	takeMVar var

newCyclicBarrier :: Int -> IO CyclicBarrier
newCyclicBarrier k = liftM (Cyclic k) (newMVar (k, []))

powerMethod :: Int -> Int -> Reals -> Reals -> IO ()
powerMethod z n u v = allocaArray n $ \ !t -> do
	let chunk = (n + numCapabilities - 1) `quotInt` numCapabilities
	!barrier <- newCyclicBarrier $! (n + chunk - 1) `quotInt` chunk
	let timesAtAv !s !d l r = do
		timesAv n s t l r
		await barrier
		timesAtv n t d l r
		await barrier
	let thread !l !r = GHC.Base.foldr (>>) (return ()) $ replicate z $ do
		timesAtAv u v l r
		timesAtAv v u l r
	let go l = case l + chunk of
		r | r < n	-> forkIO (thread l r) >> go r
		  | otherwise	-> thread l n
	go 0

timesAv :: Int -> Reals -> Reals -> Int -> Int -> IO ()
timesAv !n !u !au !l !r = go l where
    go :: Int -> IO ()
    go !i = when (i < r) $ do
	let avsum !j !acc
		| j < n = do
			!uj <- peekElemOff u j
			avsum (j+1) (acc + ((aij i j) * uj))
		| otherwise = pokeElemOff au i acc >> go (i+1)
	avsum 0 0

timesAtv :: Int -> Reals -> Reals -> Int -> Int -> IO ()
timesAtv !n !u !a !l !r = go l
  where
    go :: Int -> IO ()
    go !i = when (i < r) $ do
	let atvsum !j !acc 
		| j < n	= do	!uj <- peekElemOff u j
				atvsum (j+1) (acc + ((aij j i) * uj))
		| otherwise = pokeElemOff a i acc >> go (i+1)
	atvsum 0 0

--
-- manually unbox the inner loop:
-- aij i j = 1 / fromIntegral ((i+j) * (i+j+1) `div` 2 + i + 1)
--
aij (I# i) (I# j) = D# (
    case i +# j of
        n -> 1.0## /## int2Double# 
        	(((n *# (n+#1#)) `uncheckedIShiftRA#` 1#) +# (i +# 1#)))
    

NOTES:
64-bit Ubuntu quad core
The Glorious Glasgow Haskell
Compilation System,
version 9.4.8
LLVM version 17.0.2


 Thu, 07 Mar 2024 23:53:55 GMT

MAKE:
mv spectralnorm.ghc-4.ghc spectralnorm.ghc-4.hs
~/.ghcup/bin/ghc --make -optlo-passes='module(default<O2>)' -O2 -fno-llvm-tbaa -XBangPatterns -threaded -rtsopts -XMagicHash spectralnorm.ghc-4.hs -o spectralnorm.ghc-4.ghc_run
Loaded package environment from /home/dunham/.ghc/x86_64-linux-9.4.8/environments/default
[1 of 2] Compiling Main             ( spectralnorm.ghc-4.hs, spectralnorm.ghc-4.o )

spectralnorm.ghc-4.hs:41:21: warning: [-Wtabs]
    Tab character found here, and in 68 further locations.
    Suggested fix: Please use spaces instead.
   |
41 |     | i < n     = do    ui <- peekElemOff u i
   |                     ^^^^
[2 of 2] Linking spectralnorm.ghc-4.ghc_run
rm spectralnorm.ghc-4.hs

11.13s to complete and log all make actions

COMMAND LINE:
 ./spectralnorm.ghc-4.ghc_run +RTS -N4 -RTS 5500

PROGRAM OUTPUT:
1.274224153