<commit_before>
module Main where

import Common.Body
import Common.World
import Random.Position
import Solver

import Data.Array.Accelerate                                        as A
import Data.Array.Accelerate.Array.Data                             as A
import Data.Array.Accelerate.System.Random.SFMT

import Data.Array.Accelerate.Interpreter                            as I
import Data.Array.Accelerate.CUDA                                   as CUDA ( run1 )
import Data.Array.Accelerate.LLVM.Multi                             as Multi
import Data.Array.Accelerate.LLVM.Native                            as CPU
import Data.Array.Accelerate.LLVM.PTX                               as PTX

import Foreign.CUDA.Driver                                          as CUDA
import Foreign.CUDA.Driver.Profiler                                 as Prof

import Criterion.Main

import Control.Monad
import System.Environment
import Text.Printf
import Prelude                                                      as P

import GHC.Conc

maybeEnv :: Read a => String -> a -> IO a
maybeEnv var def = do
  ms <- lookupEnv var
  case ms of
    Just s | [(v,[])] <- reads s -> return v
    _                            -> return def

main :: IO ()
main = do
  n       <- maybeEnv "N" 1000
  pin     <- maybeEnv "PINNED" False

  printf "N-body simulation:\n"
  printf "  bodies:              %d\n" n
  printf "  number of threads:   %d\n" =<< getNumCapabilities
  printf "  using pinned memory: %s\n" (show pin)
  printf "\n"

  when pin $ do
    CUDA.initialise []
    dev <- CUDA.device 0
    ctx <- CUDA.create dev []
    registerForeignPtrAllocator (CUDA.mallocHostForeignPtr [])

  let mMax    = 40
      rMax    = 500
      epsilon = 50
      extent  = 1000
      speed   = 1

  positions <- randomArray (cloud (extent,extent) rMax) (Z :. n)
  masses    <- randomArray (uniformR (1,mMax)) (Z :. n)

  let bodies  = I.run
              $ A.map (setStartVelOfBody speed)
              $ A.zipWith setMassOfBody (A.use masses)
              $ A.map unitBody (A.use positions)

      dt      = fromList Z [0.1]
      advance = advanceBodies (calcAccels epsilon) (use dt)

  defaultMain
    [ bench "cuda"       $ whnf (CUDA.run1 advance) bodies
    , bench "llvm-cpu"   $ whnf (CPU.run1 advance) bodies
    , bench "llvm-ptx"   $ whnf (PTX.run1 advance) bodies
    , bench "llvm-multi" $ whnf (Multi.run1 advance) bodies
    ]
<commit_msg>n-body: add manual 2-gpu split benchmark<commit_after>{-# LANGUAGE BangPatterns     #-}
{-# LANGUAGE ParallelListComp #-}

module Main where

import Common.Body
import Common.World
import Random.Position
import Solver

import Data.Array.Accelerate                                        as A
import Data.Array.Accelerate.Array.Data                             as A
import Data.Array.Accelerate.Array.Sugar                            as S
import Data.Array.Accelerate.System.Random.SFMT

import Data.Array.Accelerate.Interpreter                            as I
import Data.Array.Accelerate.CUDA                                   as CUDA
import Data.Array.Accelerate.LLVM.Multi                             as Multi
import Data.Array.Accelerate.LLVM.Native                            as CPU
import Data.Array.Accelerate.LLVM.PTX                               as PTX
import Data.Array.Accelerate.Debug                                  ( accInit )

import Foreign.CUDA.Driver                                          as F

import Criterion.Main

import Control.Monad
import System.Environment
import System.IO
import System.IO.Unsafe
import System.CPUTime
import Text.Printf
import Prelude                                                      as P

import GHC.Conc
import GHC.Base                                                     ( quotInt, remInt )


maybeEnv :: Read a => String -> a -> IO a
maybeEnv var def = do
  ms <- lookupEnv var
  case ms of
    Just s | [(v,[])] <- reads s -> return v
    _                            -> return def

main :: IO ()
main = do
  accInit

  -- Create the contexts for the CUDA and PTX backends. Sadly, the underlying
  -- (raw) contexts can not be shared with the current API.
  --
  F.initialise []
  ngpu    <- F.count
  devs    <- mapM F.device [0 .. ngpu - 1]
  prps    <- mapM F.props devs
  cc    <- mapM     (\dev     -> CUDA.create      dev     []) devs
  pc    <- zipWithM (\dev prp -> PTX.createTarget dev prp []) devs prps

  n       <- maybeEnv "N" 10000
  pin     <- maybeEnv "PINNED" False

  printf "N-body simulation:\n"
  printf "  bodies:              %d\n" n
  printf "  number of threads:   %d\n" =<< getNumCapabilities
  printf "  number of GPUs:      %d\n" ngpu
  printf "  using pinned memory: %s\n" (show pin)
  printf "\n"

  when pin $ do
    registerForeignPtrAllocator (F.mallocHostForeignPtr [])

  let mMax    = 40
      rMax    = 500
      epsilon = 50
      extent  = 1000
      speed   = 1

  printf "generating data... "
  hFlush stdout

  t1          <- getCPUTime
  !positions  <- randomArray (cloud (extent,extent) rMax) (Z :. n)
  !masses     <- randomArray (uniformR (1,mMax)) (Z :. n)

  let !bodies       = I.run
                    $ A.map (setStartVelOfBody speed)
                    $ A.zipWith setMassOfBody (A.use masses)
                    $ A.map unitBody (A.use positions)

      dt            = fromList Z [0.1]
      advance       = advanceBodies (calcAccels epsilon) (use dt)
      split_bodies  = split ngpu bodies

      p0            = head pc
      c0            = head cc

  t2          <- bodies `seq` getCPUTime

  printf "done! (%.2fs)\n" (P.fromIntegral (t2-t1) * 1.0e-12 :: Double)
  printf "\n"

  defaultMain
    [ bench "cuda"       $ whnf (CUDA.run1With c0 advance) bodies
    , bench "llvm-ptx"   $ whnf (PTX.run1With  p0 advance) bodies
    , bench "llvm-cpu"   $ whnf (CPU.run1   advance) bodies
    , bench "llvm-multi" $ whnf (Multi.run1 advance) bodies
    , bgroup "multi"
      [ bench "cuda"     $ whnf (async [ CUDA.run1AsyncWith ctx advance | ctx <- cc]) split_bodies
      , bench "llvm-ptx" $ whnf (async [ PTX.run1AsyncWith  ptx advance | ptx <- pc]) split_bodies
      ]
    ]


async :: [a -> Async b] -> [a] -> ()
async fs xs = unsafePerformIO $! do
  let as = P.zipWith ($) fs xs
  () <- mapM_ (\a -> a `seq` return ()) as
  () <- mapM_ wait as
  return ()


split :: Elt e => Int -> Vector e -> [Vector e]
split 1      arr = [ arr ]
split pieces arr = [ range from to | from <- splitPts | to <- P.tail splitPts ]
  where
    Z :. n        = arrayShape arr
    chunk         = n `quotInt` pieces
    leftover      = n `remInt`  pieces

    splitPts      = P.map splitIx [0 .. pieces]
    splitIx i
      | i < leftover  = i * (chunk + 1)
      | otherwise     = i * chunk + leftover

    range from to = A.fromFunction
                        (Z :. to - from)
                        (\(Z :. i) -> arr S.! (Z :. i+from))

{'fraction_non_alphanumeric': 0.09868995633187773, 'fraction_numerical': 0.010189228529839884, 'mean_word_length': 1.5821119879744456, 'pattern_counts': {'":': 0, '<': 35, '<?xml version=': 0, '>': 23, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 21, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3924032', 'n_tokens_mistral': 2221, 'n_tokens_neox': 2098, 'n_words': 635}