<commit_before>{-# LANGUAGE GADTs #-}
-- |
-- Module      : Data.Array.Accelerate.CUDA.Compile
-- Copyright   : [2008..2010] <NAME>, <NAME>, <NAME>
-- License     : BSD3
--
-- Maintainer  : <NAME> <<EMAIL>>
-- Stability   : experimental
-- Portability : non-partable (GHC extensions)
--

module Data.Array.Accelerate.CUDA.Compile (accToKey, compile)
  where

import Prelude hiding (id, (.), mod)
import Control.Category

import Data.Maybe
import Control.Arrow
import Control.Applicative
import Control.Monad.IO.Class
import Control.Exception.Extensible
import Language.C
import Text.PrettyPrint
import qualified Data.Map                               as M

import System.Directory
import System.FilePath
import System.Posix.Process
import System.IO

import Data.Array.Accelerate.AST
import Data.Array.Accelerate.CUDA.State
import Data.Array.Accelerate.CUDA.CodeGen

import Foreign.CUDA.Analysis.Device


-- This is going to be annoying to maintain...
--
accToInt :: OpenAcc aenv a -> Int
accToInt (Let  _ _)          = 0
accToInt (Let2 _ _)          = 1
accToInt (Avar _)            = 2
accToInt (Use  _)            = 3
accToInt (Unit _)            = 4
accToInt (Reshape _ _)       = 5
accToInt (Replicate _ _ _)   = 6
accToInt (Index _ _ _)       = 7
accToInt (Map _ _)           = 8
accToInt (ZipWith _ _ _)     = 9
accToInt (Fold _ _ _)        = 10
accToInt (FoldSeg _ _ _ _)   = 11
accToInt (Scan _ _ _)        = 12
--accToInt (Scanr _ _ _)       = 13
accToInt (Permute _ _ _ _)   = 14
accToInt (Backpermute _ _ _) = 15


-- Maybe the sub-expression (for example, the function argument to `fold`) is
-- sufficient, and/or accToInt is unnecessary.
--
accToKey :: OpenAcc aenv a -> Key
accToKey = accToInt &&& show


-- |
-- Generate and compile code for an array expression
--
compile :: OpenAcc aenv a -> CIO ()
compile acc = do
  nvcc   <- fromMaybe (error "nvcc: command not found") <$> liftIO (findExecutable "nvcc")
  dir    <- liftIO outputDir
  cufile <- outputName acc (dir </> "dragon.cu")        -- here be dragons!
  flags  <- compileFlags nvcc
  pid    <- liftIO . withFilePath dir $ do
              writeCode cufile $ codeGenAcc acc
              forkProcess      $ executeFile nvcc False (takeFileName cufile : flags) Nothing

  modM kernelEntry $ M.insert (accToKey acc) (KernelEntry cufile (Left pid))


-- Write the generated code to file, exporting C symbols
--
-- TLM: The nasty __global__ attribute resolution issue is only present for
--      cuda < 3.0, might want to add #undefs?
--
writeCode :: FilePath -> CTranslUnit -> IO ()
writeCode f code
  = withFile f WriteMode $ \hdl -> do
      hPutStrLn hdl "extern \"C\" {"
      printDoc PageMode hdl (pretty code)
      hPutStrLn hdl "}"

-- stolen from $fptools/ghc/compiler/utils/Pretty.lhs
--
-- This code has a BSD-style license<commit_msg>traverse and generate code for the entire expression tree<commit_after>{-# LANGUAGE GADTs #-}
-- |
-- Module      : Data.Array.Accelerate.CUDA.Compile
-- Copyright   : [2008..2010] <NAME>, <NAME>, <NAME>
-- License     : BSD3
--
-- Maintainer  : <NAME> <<EMAIL>>
-- Stability   : experimental
-- Portability : non-partable (GHC extensions)
--

module Data.Array.Accelerate.CUDA.Compile (accToKey, compile)
  where

import Prelude hiding (id, (.), mod)
import Control.Category

import Data.Maybe
import Control.Arrow
import Control.Monad
import Control.Applicative
import Control.Monad.IO.Class
import Control.Exception.Extensible
import Language.C
import Text.PrettyPrint
import qualified Data.Map                               as M

import System.Directory
import System.FilePath
import System.Posix.Process
import System.IO

import Data.Array.Accelerate.AST
import Data.Array.Accelerate.CUDA.State
import Data.Array.Accelerate.CUDA.CodeGen

import Foreign.CUDA.Analysis.Device


-- This is going to be annoying to maintain...
--
accToInt :: OpenAcc aenv a -> Int
accToInt (Let  _ _)          = 0
accToInt (Let2 _ _)          = 1
accToInt (Avar _)            = 2
accToInt (Use  _)            = 3
accToInt (Unit _)            = 4
accToInt (Reshape _ _)       = 5
accToInt (Replicate _ _ _)   = 6
accToInt (Index _ _ _)       = 7
accToInt (Map _ _)           = 8
accToInt (ZipWith _ _ _)     = 9
accToInt (Fold _ _ _)        = 10
accToInt (FoldSeg _ _ _ _)   = 11
accToInt (Scan _ _ _)        = 12
--accToInt (Scanr _ _ _)       = 13
accToInt (Permute _ _ _ _)   = 14
accToInt (Backpermute _ _ _) = 15


-- Maybe the sub-expression (for example, the function argument to `fold`) is
-- sufficient, and/or accToInt is unnecessary.
--
accToKey :: OpenAcc aenv a -> Key
accToKey = accToInt &&& show


-- |
-- Traverse the expression tree, generating and compiling code
--
compile :: OpenAcc aenv a -> CIO ()
compile (Let  xs ys)   = compile xs >> compile ys
compile (Let2 xs ys)   = compile xs >> compile ys
compile (Avar _)       = return ()      -- TLM: ??
compile (Use  _)       = return ()
compile (Unit _)       = return ()      -- TLM: ??
compile (Reshape _ xs) = compile xs     -- TLM: ??
compile (Index _ xs _) = compile xs     -- TLM: ??

compile acc@(Replicate _ _ xs)   = compile xs >> compile' acc
compile acc@(Map _ xs)           = compile xs >> compile' acc
compile acc@(ZipWith _ xs ys)    = compile xs >> compile ys >> compile' acc
compile acc@(Fold _ _ xs)        = compile xs >> compile' acc
compile acc@(FoldSeg _ _ xs ys)  = compile xs >> compile ys >> compile' acc
compile acc@(Scan _ _ xs)        = compile xs >> compile' acc
compile acc@(Permute _ xs _ ys)  = compile xs >> compile ys >> compile' acc
compile acc@(Backpermute _ _ xs) = compile xs >> compile' acc


-- Generate and compile code for an array expression
--
compile' :: OpenAcc aenv a -> CIO ()
compile' acc = do
  compiled <- M.member (accToKey acc) <$> getM kernelEntry
  when (not compiled) $ do
    nvcc   <- fromMaybe (error "nvcc: command not found") <$> liftIO (findExecutable "nvcc")
    dir    <- liftIO outputDir
    cufile <- outputName acc (dir </> "dragon.cu")        -- here be dragons!
    flags  <- compileFlags nvcc
    pid    <- liftIO . withFilePath dir $ do
                writeCode cufile $ codeGenAcc acc
                forkProcess      $ executeFile nvcc False (takeFileName cufile : flags) Nothing

    modM kernelEntry $ M.insert (accToKey acc) (KernelEntry cufile (Left pid))


-- Write the generated code to file, exporting C symbols
--
writeCode :: FilePath -> CTranslUnit -> IO ()
writeCode f code
  = withFile f WriteMode $ \hdl -> do
      hPutStrLn hdl "extern \"C\" {"
      printDoc PageMode hdl (pretty code)
      hPutStrLn hdl "}"

-- stolen from $fptools/ghc/compiler/utils/Pretty.lhs
--
-- This code has a BSD-style license
{'fraction_non_alphanumeric': 0.11042760690172543, 'fraction_numerical': 0.01005251312828207, 'mean_word_length': 2.5974096060442524, 'pattern_counts': {'":': 0, '<': 32, '<?xml version=': 0, '>': 59, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4537749', 'n_tokens_mistral': 2279, 'n_tokens_neox': 2150, 'n_words': 674}