diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,70 +6,77 @@ project adheres to the [Haskell Package Versioning Policy (PVP)](https://pvp.haskell.org) -## [1.2.0.1] - 2019-04-29 -### Added - * support for GHC-8.6 - * support for LLVM-7 - * support for LLVM-8 +## [1.3.0.0] - 2018-08-27 +### Changed + * Code generation improvements for stencil operations + +### Fixed + * Segmented folds crash or give inconsistent results ([accelerate#423]) + * Synchronisation problems on SM7+ [#436] + ### Contributors Special thanks to those who contributed patches as part of this release: - * Trevor L. McDonell (@tmcdonell) - * Viktor Kronvall (@considerate) + * Trevor L. McDonell (@tmcdonell) + * Josh Meredith (@JoshMeredith) + * Ivo Gabe de Wolff (@ivogabe) + * Lars van den Haak (@sakehl) + * Joshua Meredith (@JoshMeredith) + ## [1.2.0.0] - 2018-04-03 ### Changed - * `run` variants which do not take an explicit execution context now execute on - the first available device in an exclusive fashion. Multi-GPU systems can - specify the default set of GPUs to use with environment variable - `ACCELERATE_LLVM_PTX_DEVICES` as a list of device ordinals. + * `run` variants which do not take an explicit execution context now execute on + the first available device in an exclusive fashion. Multi-GPU systems can + specify the default set of GPUs to use with environment variable + `ACCELERATE_LLVM_PTX_DEVICES` as a list of device ordinals. ### Added - * support for half-precision floats - * support for struct-of-array-of-struct representations - * support 64-bit atomic-add instruction in forward permutations ([#363]) - * support for LLVM-6.0 - * support for GHC-8.4 + * support for half-precision floats + * support for struct-of-array-of-struct representations + * support 64-bit atomic-add instruction in forward permutations ([#363]) + * support for LLVM-6.0 + * support for GHC-8.4 ### Contributors Special thanks to those who contributed patches as part of this release: - * Trevor L. McDonell (@tmcdonell) - * Moritz Kiefer (@cocreature) + * Trevor L. McDonell (@tmcdonell) + * Moritz Kiefer (@cocreature) ## [1.1.0.1] - 2018-01-08 ### Fixed - * add support for building with CUDA-9.x + * add support for building with CUDA-9.x ## [1.1.0.0] - 2017-09-21 ### Added - * support for GHC-8.2 - * caching of compilation results ([accelerate-llvm#17]) - * support for ahead-of-time compilation (`runQ` and `runQAsync`) + * support for GHC-8.2 + * caching of compilation results ([accelerate-llvm#17]) + * support for ahead-of-time compilation (`runQ` and `runQAsync`) ### Changed - * generalise `run1*` to polyvariadic `runN*` + * generalise `run1*` to polyvariadic `runN*` ### Fixed - * Fixed synchronisation bug in multidimensional reduction + * Fixed synchronisation bug in multidimensional reduction ## [1.0.0.1] - 2017-05-25 ### Fixed - * device kernel image is invalid ([#386]) + * device kernel image is invalid ([#386]) ## [1.0.0.0] - 2017-03-31 - * initial release + * initial release -[1.2.0.1]: https://github.com/AccelerateHS/accelerate-llvm/compare/v1.2.0.0...v1.2.0.1 -[1.2.0.0]: https://github.com/AccelerateHS/accelerate-llvm/compare/1.1.0.1-ptx...v1.2.0.0 +[1.3.0.0]: https://github.com/AccelerateHS/accelerate-llvm/compare/1.2.0.0...v1.3.0.0 +[1.2.0.0]: https://github.com/AccelerateHS/accelerate-llvm/compare/1.1.0.1-ptx...1.2.0.0 [1.1.0.1]: https://github.com/AccelerateHS/accelerate-llvm/compare/1.1.0.0...1.1.0.1-ptx [1.1.0.0]: https://github.com/AccelerateHS/accelerate-llvm/compare/1.0.0.0...1.1.0.0 [1.0.0.1]: https://github.com/AccelerateHS/accelerate-llvm/compare/1.0.0.0...1.0.0.1 @@ -78,5 +85,7 @@ [#386]: https://github.com/AccelerateHS/accelerate/issues/386 [#363]: https://github.com/AccelerateHS/accelerate/issues/363 +[#436]: https://github.com/AccelerateHS/accelerate/issues/436 [accelerate-llvm#17]: https://github.com/AccelerateHS/accelerate-llvm/issues/17 +[accelerate#423]: https://github.com/AccelerateHS/accelerate/issues/423 diff --git a/LICENSE b/LICENSE --- a/LICENSE +++ b/LICENSE @@ -1,4 +1,4 @@ -Copyright (c) [2014..2017] The Accelerate Team. All rights reserved. +Copyright (c) [2014..2020] The Accelerate Team. All rights reserved. Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -1,14 +1,20 @@ -An LLVM backend for the Accelerate Array Language -================================================= +
+henlo, my name is Theia -[![Travis](https://img.shields.io/travis/AccelerateHS/accelerate-llvm/master.svg?label=linux)](https://travis-ci.org/AccelerateHS/accelerate-llvm) -[![AppVeyor](https://img.shields.io/appveyor/ci/tmcdonell/accelerate-llvm/master.svg?label=windows)](https://ci.appveyor.com/project/tmcdonell/accelerate-llvm) +# LLVM backends for the Accelerate array language + +[![GitHub CI](https://github.com/tmcdonell/accelerate-llvm/workflows/CI/badge.svg)](https://github.com/tmcdonell/accelerate-llvm/actions) +[![Gitter](https://img.shields.io/gitter/room/nwjs/nw.js.svg)](https://gitter.im/AccelerateHS/Lobby) +
[![Stackage LTS](https://stackage.org/package/accelerate-llvm/badge/lts)](https://stackage.org/lts/package/accelerate-llvm) [![Stackage Nightly](https://stackage.org/package/accelerate-llvm/badge/nightly)](https://stackage.org/nightly/package/accelerate-llvm) [![Hackage](https://img.shields.io/hackage/v/accelerate-llvm.svg)](https://hackage.haskell.org/package/accelerate-llvm) +
[![Docker Automated build](https://img.shields.io/docker/automated/tmcdonell/accelerate-llvm.svg)](https://hub.docker.com/r/tmcdonell/accelerate-llvm/) [![Docker status](https://images.microbadger.com/badges/image/tmcdonell/accelerate-llvm.svg)](https://microbadger.com/images/tmcdonell/accelerate-llvm) +
+ This package compiles Accelerate code to LLVM IR, and executes that code on multicore CPUs as well as NVIDIA GPUs. This avoids the need to go through `nvcc` or `clang`. For details on Accelerate, refer to the [main repository][GitHub]. @@ -67,7 +73,7 @@ Example using [Homebrew](http://brew.sh) on macOS: ```sh -$ brew install llvm-hs/llvm/llvm-8 +$ brew install llvm-hs/llvm/llvm-9 ``` ## Debian/Ubuntu @@ -78,17 +84,17 @@ then: ```sh -$ apt-get install llvm-8-dev +$ apt-get install llvm-9-dev ``` ## Building from source If your OS does not have an appropriate LLVM distribution available, you can also build from source. Detailed build instructions are available on the [LLVM.org website](http://releases.llvm.org/6.0.0/docs/CMake.html). Note that you will require at least [CMake 3.4.3](http://www.cmake.org/cmake/resources/software.html) and a recent C++ compiler; at least Clang 3.1, GCC 4.8, or Visual Studio 2015 (update 3). - 1. Download and unpack the [LLVM-8.0 source code](http://releases.llvm.org/8.0.0/llvm-8.0.0.src.tar.xz). We'll refer to + 1. Download and unpack the [LLVM-9 source code](https://github.com/llvm/llvm-project/releases/download/llvmorg-9.0.1/llvm-9.0.1.src.tar.xz). We'll refer to the path that the source tree was unpacked to as `LLVM_SRC`. Only the main LLVM source tree is required, but you can optionally add other components - such as the Clang compiler or Polly loop optimiser. See the [LLVM releases](http://releases.llvm.org/download.html#8.0.0) + such as the Clang compiler or Polly loop optimiser. See the [LLVM releases](http://releases.llvm.org/download.html#9.0.1) page for the complete list. 2. Create a temporary build directory and `cd` into it, for example: @@ -116,7 +122,7 @@ to [System Integrity Protection](https://en.wikipedia.org/wiki/System_Integrity_Protection): ```sh cd $INSTALL_PREFIX/lib - ln -s libLLVM.dylib libLLVM-8.0.dylib + ln -s libLLVM.dylib libLLVM-9.dylib install_name_tool -id $PWD/libLTO.dylib libLTO.dylib install_name_tool -id $PWD/libLLVM.dylib libLLVM.dylib install_name_tool -change '@rpath/libLLVM.dylib' $PWD/libLLVM.dylib libLTO.dylib @@ -131,13 +137,13 @@ For example, installation using [`stack`](http://docs.haskellstack.org/en/stable/README.html) just requires you to point it to the appropriate configuration file: ```sh -$ ln -s stack-8.6.yaml stack.yaml +$ ln -s stack-8.8.yaml stack.yaml $ stack setup $ stack install ``` Note that the version of [`llvm-hs`](https://hackage.haskell.org/package/llvm-hs) -used must match the installed version of LLVM, which is currently 8.0. +used must match the installed version of LLVM, which is currently 9.0. ## libNVVM @@ -153,20 +159,23 @@ is also based on LLVM, and typically lags LLVM by several releases, you must install `accelerate-llvm` with a "compatible" version of LLVM, which will depend on the version of the CUDA toolkit you have installed. The following table shows -some combinations: +combinations which have been tested: -| | LLVM-3.3 | LLVM-3.4 | LLVM-3.5 | LLVM-3.8 | LLVM-3.9 | LLVM-4.0 | LLVM-5.0 | LLVM-6.0 | LLVM-7.0 | LLVM-8.0 | -|:-------------:|:--------:|:--------:|:--------:|:--------:|:--------:|:--------:|:--------:|:--------:|:--------:|:--------:| -| **CUDA-7.0** | ⭕ | ❌ | | | | | | | | | -| **CUDA-7.5** | | ⭕ | ⭕ | ❌ | | | | | | | -| **CUDA-8.0** | | | ⭕ | ⭕ | ❌ | ❌ | | | | | -| **CUDA-9.0** | | | | | | ❌ | ❌ | | | | -| **CUDA-9.1** | | | | | | | | | | | -| **CUDA-9.2** | | | | | | | | | | | -| **CUDA-10.0** | | | | | | | | | | | -| **CUDA-10.1** | | | | | | | | | | | +| | LLVM-3.3 | LLVM-3.4 | LLVM-3.5 | LLVM-3.8 | LLVM-3.9 | LLVM-4.0 | LLVM-5.0 | LLVM-6.0 | LLVM-7 | LLVM-8 | LLVM-9 | +| ------------- | :------: | :------: | :------: | :------: | :------: | :------: | :------: | :------: | :----: | :----: | :----: | +| **CUDA-7.0** | ⭕ | ❌ | | | | | | | | | | +| **CUDA-7.5** | | ⭕ | ⭕ | ❌ | | | | | | | | +| **CUDA-8.0** | | | ⭕ | ⭕ | ❌ | ❌ | | | | | | +| **CUDA-9.0** | | | | | | ❌ | ❌ | | | | | +| **CUDA-9.1** | | | | | | | | | | | | +| **CUDA-9.2** | | | | | | | | | | | | +| **CUDA-10.0** | | | | | | | | | | | | +| **CUDA-10.1** | | | | | | | | | | | | Where ⭕ = Works, and ❌ = Does not work. + +The above table is incomplete! If you try a particular combination and find that +it does or does not work, please let us know! Note that the above restrictions on CUDA and LLVM version exist _only_ if you want to use the NVVM component. Otherwise, you should be free to use any diff --git a/accelerate-llvm-ptx.cabal b/accelerate-llvm-ptx.cabal --- a/accelerate-llvm-ptx.cabal +++ b/accelerate-llvm-ptx.cabal @@ -1,7 +1,7 @@ name: accelerate-llvm-ptx -version: 1.2.0.1 +version: 1.3.0.0 cabal-version: >= 1.10 -tested-with: GHC >= 8.0 +tested-with: GHC >= 8.6 build-type: Simple synopsis: Accelerate backend for NVIDIA GPUs @@ -26,7 +26,7 @@ . Example using Homebrew on macOS: . - > brew install llvm-hs/llvm/llvm-8 + > brew install llvm-hs/llvm/llvm-9 . /Debian & Ubuntu/ . @@ -35,41 +35,25 @@ instructions for adding the correct package database for your OS version, and then: . - > apt-get install llvm-8-dev + > apt-get install llvm-9-dev . /Building from source/ . If your OS does not have an appropriate LLVM distribution available, you can also build from source. Detailed build instructions are available on - . Make sure to + . Make sure to include the cmake build options @-DLLVM_BUILD_LLVM_DYLIB=ON -DLLVM_LINK_LLVM_DYLIB=ON@ so that the @libLLVM@ shared library will be built. Also ensure that the @LLVM_TARGETS_TO_BUILD@ option includes the @NVPTX@ target (if not specified all targets are built). . - [/Installing accelerate-llvm/] - . - To use @accelerate-llvm@ it is important that the @llvm-hs@ package is - installed against the @libLLVM@ shared library, rather than statically - linked, so that we can use LLVM from GHCi and Template Haskell. This is the - default configuration, but you can also enforce this explicitly by adding - the following to your @stack.yaml@ file: - . - > flags: - > llvm-hs: - > shared-llvm: true - . - Or by specifying the @shared-llvm@ flag to cabal: - . - > cabal install llvm-hs -fshared-llvm - . license: BSD3 license-file: LICENSE author: Trevor L. McDonell -maintainer: Trevor L. McDonell +maintainer: Trevor L. McDonell bug-reports: https://github.com/AccelerateHS/accelerate/issues -category: Compilers/Interpreters, Concurrency, Data, Parallelism +category: Accelerate, Compilers/Interpreters, Concurrency, Data, Parallelism extra-source-files: CHANGELOG.md @@ -121,8 +105,9 @@ Data.Array.Accelerate.LLVM.PTX.CodeGen.Loop Data.Array.Accelerate.LLVM.PTX.CodeGen.Map Data.Array.Accelerate.LLVM.PTX.CodeGen.Permute - -- Data.Array.Accelerate.LLVM.PTX.CodeGen.Queue Data.Array.Accelerate.LLVM.PTX.CodeGen.Scan + Data.Array.Accelerate.LLVM.PTX.CodeGen.Stencil + Data.Array.Accelerate.LLVM.PTX.CodeGen.Transform Data.Array.Accelerate.LLVM.PTX.Compile Data.Array.Accelerate.LLVM.PTX.Compile.Cache @@ -137,32 +122,43 @@ Data.Array.Accelerate.LLVM.PTX.Embed System.Process.Extra + GHC.Heap.NormalForm Paths_accelerate_llvm_ptx build-depends: - base >= 4.7 && < 4.13 - , accelerate == 1.2.* - , accelerate-llvm == 1.2.* + base >= 4.10 && < 5 + , accelerate == 1.3.* + , accelerate-llvm == 1.3.* , bytestring >= 0.10.4 - , containers >= 0.5 && <0.7 - , cuda >= 0.9 + , containers >= 0.5 && < 0.7 + , cuda >= 0.10 , deepseq >= 1.3 , directory >= 1.0 , dlist >= 0.6 , file-embed >= 0.0.8 , filepath >= 1.0 , hashable >= 1.2 - , llvm-hs >= 4.1 && < 8.1 - , llvm-hs-pure >= 4.1 && < 8.1 + , llvm-hs >= 4.1 && < 9.1 + , llvm-hs-pure >= 4.1 && < 9.1 , mtl >= 2.2.1 - , nvvm >= 0.7.5 + , nvvm >= 0.9 , pretty >= 1.1 , process >= 1.4.3 , template-haskell - , time >= 1.4 , unordered-containers >= 0.2 + if impl(ghc >= 8.6) + build-depends: + ghc-heap + else + build-depends: + ghc + , ghci + + cpp-options: + -DUNBOXED_TUPLES=1 + hs-source-dirs: src @@ -193,7 +189,7 @@ main-is: Main.hs build-depends: - base >= 4.7 + base >= 4.10 , accelerate , accelerate-llvm-ptx @@ -207,6 +203,7 @@ -rtsopts -with-rtsopts=-A128M -with-rtsopts=-n4M + -- -with-rtsopts=-N source-repository head @@ -215,7 +212,7 @@ source-repository this type: git - tag: v1.2.0.1 + tag: v1.3.0.0 location: https://github.com/AccelerateHS/accelerate-llvm.git -- vim: nospell diff --git a/src/Data/Array/Accelerate/LLVM/PTX.hs b/src/Data/Array/Accelerate/LLVM/PTX.hs --- a/src/Data/Array/Accelerate/LLVM/PTX.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX.hs @@ -1,17 +1,20 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} {-# LANGUAGE BangPatterns #-} {-# LANGUAGE CPP #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GADTs #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE TypeSynonymInstances #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX --- Copyright : [2014..2017] Trevor L. McDonell --- [2014..2014] Vinod Grover (NVIDIA Corporation) +-- Copyright : [2014..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- @@ -51,36 +54,36 @@ ) where --- accelerate -import Data.Array.Accelerate.AST ( PreOpenAfun(..) ) -import Data.Array.Accelerate.Array.Sugar ( Arrays ) -import Data.Array.Accelerate.Async -import Data.Array.Accelerate.Debug as Debug +import Data.Array.Accelerate.AST ( PreOpenAfun(..), arraysR, liftALeftHandSide ) +import Data.Array.Accelerate.AST.LeftHandSide ( lhsToTupR ) +import Data.Array.Accelerate.Array.Data +import Data.Array.Accelerate.Async ( Async, asyncBound, wait, poll, cancel ) +import Data.Array.Accelerate.Debug import Data.Array.Accelerate.Error +import Data.Array.Accelerate.Representation.Array ( liftArraysR ) import Data.Array.Accelerate.Smart ( Acc ) +import Data.Array.Accelerate.Sugar.Array ( Arrays, toArr, fromArr, ArraysR ) import Data.Array.Accelerate.Trafo +import Data.Array.Accelerate.Trafo.Delayed +import Data.Array.Accelerate.Trafo.Sharing ( Afunction(..), AfunctionRepr(..), afunctionRepr ) +import qualified Data.Array.Accelerate.Sugar.Array as Sugar -import Data.Array.Accelerate.LLVM.Execute.Async ( AsyncR(..) ) -import Data.Array.Accelerate.LLVM.Execute.Environment ( AvalR(..) ) +import Data.Array.Accelerate.LLVM.PTX.Array.Data import Data.Array.Accelerate.LLVM.PTX.Compile -import Data.Array.Accelerate.LLVM.PTX.Embed ( embedOpenAcc ) +import Data.Array.Accelerate.LLVM.PTX.Context +import Data.Array.Accelerate.LLVM.PTX.Embed import Data.Array.Accelerate.LLVM.PTX.Execute -import Data.Array.Accelerate.LLVM.PTX.Execute.Environment ( Aval ) +import Data.Array.Accelerate.LLVM.PTX.Execute.Async ( Par, evalPar, getArrays ) +import Data.Array.Accelerate.LLVM.PTX.Execute.Environment import Data.Array.Accelerate.LLVM.PTX.Link import Data.Array.Accelerate.LLVM.PTX.State import Data.Array.Accelerate.LLVM.PTX.Target -import Data.Array.Accelerate.LLVM.State -import qualified Data.Array.Accelerate.LLVM.PTX.Array.Data as AD -import qualified Data.Array.Accelerate.LLVM.PTX.Context as CT -import qualified Data.Array.Accelerate.LLVM.PTX.Execute.Async as E import Foreign.CUDA.Driver as CUDA ( CUDAException, mallocHostForeignPtr ) --- standard library import Control.Exception import Control.Monad.Trans import Data.Maybe -import Data.Typeable import System.IO.Unsafe import Text.Printf import qualified Language.Haskell.TH as TH @@ -104,18 +107,13 @@ -- /NOTE:/ it is recommended to use 'runN' or 'runQ' whenever possible. -- run :: Arrays a => Acc a -> a -run a - = unsafePerformIO - $ wait =<< runAsync a +run a = unsafePerformIO (runIO a) -- | As 'run', but execute using the specified target rather than using the -- default, automatically selected device. -- runWith :: Arrays a => PTX -> Acc a -> a -runWith target a - = unsafePerformIO - $ wait =<< runAsyncWith target a - +runWith target a = unsafePerformIO (runWithIO target a) -- | As 'run', but run the computation asynchronously and return immediately -- without waiting for the result. The status of the computation can be queried @@ -125,33 +123,30 @@ -- a specific device, use 'runAsyncWith'. -- runAsync :: Arrays a => Acc a -> IO (Async a) -runAsync a = asyncBound execute - where - !acc = convertAccWith config a - execute = do - dumpGraph acc - withPool defaultTargetPool $ \target -> - evalPTX target $ do - build <- phase "compile" (compileAcc acc) >>= dumpStats - exec <- phase "link" (linkAcc build) - res <- phase "execute" (executeAcc exec >>= AD.copyToHostLazy) - return res +runAsync a = asyncBound (runIO a) -- | As 'runWith', but execute asynchronously. Be sure not to destroy the context, -- or attempt to attach it to a different host thread, before all outstanding -- operations have completed. -- runAsyncWith :: Arrays a => PTX -> Acc a -> IO (Async a) -runAsyncWith target a = asyncBound execute +runAsyncWith target a = asyncBound (runWithIO target a) + + +runIO :: Arrays a => Acc a -> IO a +runIO a = withPool defaultTargetPool (\target -> runWithIO target a) + +runWithIO :: forall a. Arrays a => PTX -> Acc a -> IO a +runWithIO target a = execute where - !acc = convertAccWith config a - execute = do + !acc = convertAcc a + execute = do dumpGraph acc evalPTX target $ do build <- phase "compile" (compileAcc acc) >>= dumpStats exec <- phase "link" (linkAcc build) - res <- phase "execute" (executeAcc exec >>= AD.copyToHostLazy) - return res + res <- phase "execute" (evalPar (executeAcc exec >>= copyToHostLazy (Sugar.arraysR @a))) + return $ toArr res -- | This is 'runN', specialised to an array program of one argument. @@ -213,10 +208,10 @@ -- See also 'runQ', which compiles the Accelerate program at _Haskell_ compile -- time, thus eliminating the runtime overhead altogether. -- -runN :: Afunction f => f -> AfunctionR f +runN :: forall f. Afunction f => f -> AfunctionR f runN f = exec where - !acc = convertAfunWith config f + !acc = convertAfun f !exec = unsafeWithPool defaultTargetPool $ \target -> fromJust (lookup (ptxContext target) afun) @@ -229,19 +224,19 @@ -- depending on which GPU gets scheduled. -- !afun = flip map (unmanaged defaultTargetPool) - $ \target -> (ptxContext target, runNWith' target acc) + $ \target -> (ptxContext target, runNWith' @f target acc) -- | As 'runN', but execute using the specified target device. -- -runNWith :: Afunction f => PTX -> f -> AfunctionR f +runNWith :: forall f. Afunction f => PTX -> f -> AfunctionR f runNWith target f = exec where - !acc = convertAfunWith config f - !exec = runNWith' target acc + !acc = convertAfun f + !exec = runNWith' @f target acc -runNWith' :: PTX -> DelayedAfun f -> f -runNWith' target acc = exec +runNWith' :: forall f. Afunction f => PTX -> DelayedAfun (ArraysFunctionR f) -> AfunctionR f +runNWith' target acc = go (afunctionRepr @f) afun (return Empty) where !afun = unsafePerformIO $ do dumpGraph acc @@ -249,18 +244,22 @@ build <- phase "compile" (compileAfun acc) >>= dumpStats link <- phase "link" (linkAfun build) return link - !exec = go afun (return Aempty) - go :: ExecOpenAfun PTX aenv t -> LLVM PTX (Aval aenv) -> t - go (Alam l) k = \ !arrs -> - let k' = do aenv <- k - AsyncR _ a <- E.async (AD.useRemoteAsync arrs) - return (aenv `Apush` a) - in go l k' - go (Abody b) k = unsafePerformIO . phase "execute" . evalPTX target $ do + go :: forall aenv t r trepr. + AfunctionRepr t r trepr + -> ExecOpenAfun PTX aenv trepr + -> Par PTX (Val aenv) + -> r + go (AfunctionReprLam repr) (Alam lhs l) k = \ !arrs -> + let k' = do aenv <- k + a <- useRemoteAsync (lhsToTupR lhs) $ fromArr arrs + return (aenv `push` (lhs, a)) + in go repr l k' + go AfunctionReprBody (Abody b) k = unsafePerformIO . phase "execute" . evalPTX target . evalPar $ do aenv <- k - r <- E.async (executeOpenAcc b aenv) - AD.copyToHostLazy =<< E.get r + fut <- executeOpenAcc b aenv + toArr <$> copyToHostLazy (Sugar.arraysR @r) fut + go _ _ _ = error "But that's not right, oh, no, what's the story?" -- | As 'run1', but the computation is executed asynchronously. @@ -276,10 +275,10 @@ -- | As 'runN', but execute asynchronously. -- -runNAsync :: (Afunction f, RunAsync r, AfunctionR f ~ RunAsyncR r) => f -> r +runNAsync :: (Afunction f, RunAsync r, ArraysFunctionR f ~ RunAsyncR r) => f -> r runNAsync f = exec where - !acc = convertAfunWith config f + !acc = convertAfun f !exec = unsafeWithPool defaultTargetPool $ \target -> fromJust (lookup (ptxContext target) afun) @@ -289,42 +288,44 @@ -- | As 'runNWith', but execute asynchronously. -- -runNAsyncWith :: (Afunction f, RunAsync r, AfunctionR f ~ RunAsyncR r) => PTX -> f -> r +runNAsyncWith :: (Afunction f, RunAsync r, ArraysFunctionR f ~ RunAsyncR r) => PTX -> f -> r runNAsyncWith target f = exec where - !acc = convertAfunWith config f + !acc = convertAfun f !exec = runNAsyncWith' target acc runNAsyncWith' :: RunAsync f => PTX -> DelayedAfun (RunAsyncR f) -> f -runNAsyncWith' target acc = runAsync' target afun (return Aempty) +runNAsyncWith' target acc = exec where !afun = unsafePerformIO $ do dumpGraph acc evalPTX target $ do build <- phase "compile" (compileAfun acc) >>= dumpStats - exec <- phase "link" (linkAfun build) - return exec + link <- phase "link" (linkAfun build) + return link + !exec = runAsync' target afun (return Empty) class RunAsync f where type RunAsyncR f - runAsync' :: PTX -> ExecOpenAfun PTX aenv (RunAsyncR f) -> LLVM PTX (Aval aenv) -> f + runAsync' :: PTX -> ExecOpenAfun PTX aenv (RunAsyncR f) -> Par PTX (Val aenv) -> f -instance RunAsync b => RunAsync (a -> b) where - type RunAsyncR (a -> b) = a -> RunAsyncR b +instance (Arrays a, RunAsync b) => RunAsync (a -> b) where + type RunAsyncR (a -> b) = ArraysR a -> RunAsyncR b runAsync' _ Abody{} _ _ = error "runAsync: function oversaturated" - runAsync' target (Alam l) k !arrs = - let k' = do aenv <- k - AsyncR _ a <- E.async (AD.useRemoteAsync arrs) - return (aenv `Apush` a) + runAsync' target (Alam lhs l) k !arrs = + let k' = do aenv <- k + a <- useRemoteAsync (Sugar.arraysR @a) $ fromArr arrs + return (aenv `push` (lhs, a)) in runAsync' target l k' -instance RunAsync (IO (Async b)) where - type RunAsyncR (IO (Async b)) = b +instance Arrays b => RunAsync (IO (Async b)) where + type RunAsyncR (IO (Async b)) = ArraysR b runAsync' _ Alam{} _ = error "runAsync: function not fully applied" - runAsync' target (Abody b) k = asyncBound . phase "execute" . evalPTX target $ do + runAsync' target (Abody b) k = asyncBound . phase "execute" . evalPTX target . evalPar $ do aenv <- k - r <- E.async (executeOpenAcc b aenv) - AD.copyToHostLazy =<< E.get r + ans <- executeOpenAcc b aenv + arrs <- getArrays (arraysR b) ans + return $ toArr arrs -- | Stream a lazily read list of input arrays through the given program, @@ -418,7 +419,7 @@ -- @since 1.1.0.0 -- runQAsync :: Afunction f => f -> TH.ExpQ -runQAsync = runQ' [| async |] +runQAsync = runQ' [| asyncBound |] -- | Ahead-of-time analogue of 'runNAsyncWith'. See 'runQWith' for more information. -- @@ -431,14 +432,14 @@ runQAsyncWith :: Afunction f => f -> TH.ExpQ runQAsyncWith f = do target <- TH.newName "target" - TH.lamE [TH.varP target] (runQWith' [| async |] (TH.varE target) f) + TH.lamE [TH.varP target] (runQWith' [| asyncBound |] (TH.varE target) f) runQ' :: Afunction f => TH.ExpQ -> f -> TH.ExpQ -runQ' using = runQ'_main using (\go -> [| withPool defaultTargetPool $ \target -> evalPTX target $go |]) +runQ' using = runQ'_ using (\go -> [| withPool defaultTargetPool (\target -> evalPTX target (evalPar $go)) |]) runQWith' :: Afunction f => TH.ExpQ -> TH.ExpQ -> f -> TH.ExpQ -runQWith' using target = runQ'_main using (TH.appE [| evalPTX $target |]) +runQWith' using target = runQ'_ using (\go -> [| evalPTX $target (evalPar $go) |]) -- Generate a template haskell expression for the given function to be embedded -- into the current program. The supplied continuation specifies how to execute @@ -457,40 +458,40 @@ -- With runN this will automatically be recompiled for each new architecture -- (at runtime). -- -runQ'_main :: Afunction f => TH.ExpQ -> (TH.ExpQ -> TH.ExpQ) -> f -> TH.ExpQ -runQ'_main using k f = do - afun <- let acc = convertAfunWith config f +runQ'_ :: Afunction f => TH.ExpQ -> (TH.ExpQ -> TH.ExpQ) -> f -> TH.ExpQ +runQ'_ using k f = do + afun <- let acc = convertAfun f in TH.runIO $ do dumpGraph acc evalPTX defaultTarget $ phase "compile" (compileAfun acc) >>= dumpStats let - go :: Typeable aenv => CompiledOpenAfun PTX aenv t -> [TH.PatQ] -> [TH.ExpQ] -> [TH.StmtQ] -> TH.ExpQ - go (Alam lam) xs as stmts = do + go :: CompiledOpenAfun PTX aenv t -> [TH.PatQ] -> [TH.ExpQ] -> [TH.StmtQ] -> TH.ExpQ + go (Alam lhs l) xs as stmts = do x <- TH.newName "x" -- lambda bound variable a <- TH.newName "a" -- local array name - s <- TH.bindS (TH.conP 'AsyncR [TH.wildP, TH.varP a]) [| E.async (AD.useRemoteAsync $(TH.varE x)) |] - go lam (TH.bangP (TH.varP x) : xs) (TH.varE a : as) (return s : stmts) + s <- TH.bindS (TH.varP a) [| useRemoteAsync $(TH.unTypeQ $ liftArraysR (lhsToTupR lhs)) (fromArr $(TH.varE x)) |] + go l (TH.bangP (TH.varP x) : xs) ([| ($(TH.unTypeQ $ liftALeftHandSide lhs), $(TH.varE a)) |] : as) (return s : stmts) - go (Abody body) xs as stmts = - let aenv = foldr (\a gamma -> [| $gamma `Apush` $a |] ) [| Aempty |] as - eval = TH.noBindS [| AD.copyToHostLazy =<< E.get =<< E.async (executeOpenAcc $(TH.unTypeQ (embedOpenAcc defaultTarget body)) $aenv) |] - in - TH.lamE (reverse xs) (TH.appE using [| phase "execute" $(k (TH.doE (reverse (eval : stmts)))) |]) + go (Abody b) xs as stmts = do + r <- TH.newName "r" -- result + s <- TH.newName "s" + let + aenv = foldr (\a gamma -> [| $gamma `push` $a |] ) [| Empty |] as + body = embedOpenAcc defaultTarget b + -- + TH.lamE (reverse xs) + [| $using (phase "execute" $(k ( + TH.doE ( reverse stmts ++ + [ TH.bindS (TH.varP r) [| executeOpenAcc $(TH.unTypeQ body) $aenv |] + , TH.bindS (TH.varP s) [| copyToHostLazy $(TH.unTypeQ (liftArraysR (arraysR b))) $(TH.varE r) |] + , TH.noBindS [| return $ toArr $(TH.varE s) |] + ])))) + |] -- go afun [] [] [] --- How the Accelerate program should be evaluated. --- --- TODO: make sharing/fusion runtime configurable via debug flags or otherwise. --- -config :: Phase -config = phases - { convertOffsetOfSegment = True - } - - -- Controlling host-side allocation -- -------------------------------- @@ -513,19 +514,19 @@ -- Note that since the amount of available pageable memory will be reduced, -- overall system performance can suffer. -- -registerPinnedAllocatorWith :: PTX -> IO () +registerPinnedAllocatorWith :: HasCallStack => PTX -> IO () registerPinnedAllocatorWith target = - AD.registerForeignPtrAllocator $ \bytes -> - CT.withContext (ptxContext target) (CUDA.mallocHostForeignPtr [] bytes) + registerForeignPtrAllocator $ \bytes -> + withContext (ptxContext target) (CUDA.mallocHostForeignPtr [] bytes) `catch` - \e -> $internalError "registerPinnedAlocator" (show (e :: CUDAException)) + \e -> internalError (show (e :: CUDAException)) -- Debugging -- ========= dumpStats :: MonadIO m => a -> m a -dumpStats x = dumpSimplStats >> return x +dumpStats x = liftIO dumpSimplStats >> return x phase :: MonadIO m => String -> m a -> m a phase n go = timed dump_phases (\wall cpu -> printf "phase %s: %s" n (elapsed wall cpu)) go diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Analysis/Device.hs b/src/Data/Array/Accelerate/LLVM/PTX/Analysis/Device.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/Analysis/Device.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/Analysis/Device.hs @@ -1,11 +1,10 @@ {-# LANGUAGE ScopedTypeVariables #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Analysis.Device --- Copyright : [2008..2017] Manuel M T Chakravarty, Gabriele Keller --- [2009..2017] Trevor L. McDonell +-- Copyright : [2008..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Analysis/Launch.hs b/src/Data/Array/Accelerate/LLVM/PTX/Analysis/Launch.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/Analysis/Launch.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/Analysis/Launch.hs @@ -2,11 +2,10 @@ {-# LANGUAGE TemplateHaskell #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Analysis.Launch --- Copyright : [2008..2017] Manuel M T Chakravarty, Gabriele Keller --- [2009..2017] Trevor L. McDonell +-- Copyright : [2008..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Array/Data.hs b/src/Data/Array/Accelerate/LLVM/PTX/Array/Data.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/Array/Data.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/Array/Data.hs @@ -1,14 +1,19 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE UnboxedTuples #-} {-# OPTIONS_GHC -fno-warn-orphans #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Array.Data --- Copyright : [2014..2017] Trevor L. McDonell --- [2014..2014] Vinod Grover (NVIDIA Corporation) +-- Copyright : [2014..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- @@ -20,11 +25,14 @@ ) where --- accelerate -import Data.Array.Accelerate.Array.Sugar -import Data.Array.Accelerate.Array.Unique ( UniqueArray(..) ) -import Data.Array.Accelerate.Lifetime ( Lifetime(..) ) -import qualified Data.Array.Accelerate.Array.Representation as R +import Data.Array.Accelerate.Array.Data +import Data.Array.Accelerate.Array.Unique +import Data.Array.Accelerate.Error +import Data.Array.Accelerate.Lifetime +import Data.Array.Accelerate.Representation.Array +import Data.Array.Accelerate.Representation.Shape +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Type import Data.Array.Accelerate.LLVM.Array.Data import Data.Array.Accelerate.LLVM.State @@ -32,55 +40,36 @@ import Data.Array.Accelerate.LLVM.PTX.State import Data.Array.Accelerate.LLVM.PTX.Target import Data.Array.Accelerate.LLVM.PTX.Execute.Async -import qualified Data.Array.Accelerate.LLVM.PTX.Array.Prim as Prim +import qualified Data.Array.Accelerate.LLVM.PTX.Array.Prim as Prim --- standard library import Control.Applicative -import Control.Monad.State ( liftIO, gets ) -import Data.Typeable -import Foreign.Ptr -import Foreign.Storable +import Control.Monad.Reader +import Control.Monad.State ( gets ) import System.IO.Unsafe import Prelude +import GHC.Heap.NormalForm --- Instance of remote array memory management for the PTX target + +-- | Remote memory management for the PTX target. Data can be copied +-- asynchronously using multiple execution engines whenever possible. -- instance Remote PTX where - - {-# INLINEABLE allocateRemote #-} - allocateRemote !sh = do - let !n = size sh - arr <- liftIO $ allocateArray sh - runArray arr (\m ad -> Prim.mallocArray (n*m) ad >> return ad) - - {-# INLINEABLE useRemoteR #-} - useRemoteR !n !mst !ad = do - case mst of - Nothing -> Prim.useArray n ad - Just st -> Prim.useArrayAsync st n ad - - {-# INLINEABLE copyToRemoteR #-} - copyToRemoteR !from !n !mst !ad = do - case mst of - Nothing -> Prim.pokeArrayR from n ad - Just st -> Prim.pokeArrayAsyncR st from n ad - - {-# INLINEABLE copyToHostR #-} - copyToHostR !from !n !mst !ad = do - case mst of - Nothing -> Prim.peekArrayR from n ad - Just st -> Prim.peekArrayAsyncR st from n ad - - {-# INLINEABLE copyToPeerR #-} - copyToPeerR !from !n !dst !mst !ad = do - case mst of - Nothing -> Prim.copyArrayPeerR (ptxContext dst) (ptxMemoryTable dst) from n ad - Just st -> Prim.copyArrayPeerAsyncR (ptxContext dst) (ptxMemoryTable dst) st from n ad + {-# INLINEABLE allocateRemote #-} + {-# INLINEABLE indexRemoteAsync #-} + {-# INLINEABLE useRemoteR #-} + {-# INLINEABLE copyToHostR #-} + {-# INLINEABLE copyToRemoteR #-} + allocateRemote repr@(ArrayR shr tp) !sh = do + let !n = size shr sh + arr <- liftIO $ allocateArray repr sh -- shadow array on the host + liftPar $ runArray tp arr (\m t ad -> Prim.mallocArray t (n*m) ad >> return ad) - {-# INLINEABLE indexRemote #-} - indexRemote arr i = - runIndexArray Prim.indexArray arr i + indexRemoteAsync = runIndexArrayAsync Prim.indexArrayAsync + useRemoteR = Prim.useArrayAsync + copyToHostR = Prim.peekArrayAsync + copyToRemoteR = Prim.pokeArrayAsync + copyToPeerR = internalError "not supported yet" -- | Copy an array from the remote device to the host. Although the Accelerate @@ -97,130 +86,119 @@ -- an array element or using 'deepseq' to force to normal form is required -- to actually transfer the data. -- +{-# INLINEABLE copyToHostLazy #-} copyToHostLazy - :: Arrays arrs - => arrs - -> LLVM PTX arrs -copyToHostLazy arrs = do - ptx <- gets llvmTarget - liftIO $ runArrays arrs $ \(Array sh adata) -> + :: HasCallStack + => ArraysR arrs + -> FutureArraysR PTX arrs + -> Par PTX arrs +copyToHostLazy TupRunit () = return () +copyToHostLazy (TupRpair r1 r2) (f1, f2) = do + a1 <- copyToHostLazy r1 f1 + a2 <- copyToHostLazy r2 f2 + return (a1, a2) +copyToHostLazy (TupRsingle (ArrayR shr tp)) future = do + ptx <- gets llvmTarget + liftIO $ do + Array sh adata <- wait future + + -- Note: [Lazy device-host transfers] + -- + -- This needs must be non-strict at the leaves of the datatype (that + -- is, the UniqueArray pointers). This means we can traverse the + -- ArrayData constructors (in particular, the spine defined by Unit + -- and Pair) until we reach the array we care about, without forcing + -- the other fields. + -- + -- https://github.com/AccelerateHS/accelerate/issues/437 + -- + -- Furthermore, we only want to transfer the data if the host pointer + -- is currently unevaluated. This situation can occur for example if + -- the argument to 'use' or 'unit' is returned as part of the result + -- of a 'run'. Peek at GHC's underlying closure representation and + -- check whether the pointer is a thunk, and only initiate the + -- transfer if so. + -- let - peekR :: (ArrayElt e, ArrayPtrs e ~ Ptr a, Storable a, Typeable a, Typeable e) - => ArrayData e - -> UniqueArray a - -> Int - -> IO (UniqueArray a) - peekR ad (UniqueArray uid (Lifetime ref weak fp)) n = do - fp' <- unsafeInterleaveIO $ - evalPTX ptx $ do - s <- fork - copyToHostR 0 n (Just s) ad - e <- checkpoint s - block e - join s - return fp - return $ UniqueArray uid (Lifetime ref weak fp') + peekR :: SingleType e + -> ArrayData e + -> Int + -> IO (ArrayData e) + peekR t ad m + | SingleArrayDict <- singleArrayDict t + , UniqueArray uid (Lifetime lft weak fp) <- ad + = unsafeInterleaveIO $ do + yes <- isNormalForm fp + fp' <- if yes + then return fp + else unsafeInterleaveIO . evalPTX ptx . evalPar $ do + !_ <- block =<< Prim.peekArrayAsync t m ad + return fp + -- + return $ UniqueArray uid (Lifetime lft weak fp') - runR :: ArrayEltR e -> ArrayData e -> Int -> IO (ArrayData e) - runR ArrayEltRunit AD_Unit _ = return AD_Unit - runR (ArrayEltRpair aeR2 aeR1) (AD_Pair ad2 ad1) n = AD_Pair <$> runR aeR2 ad2 n <*> runR aeR1 ad1 n - runR (ArrayEltRvec2 aeR) (AD_V2 ad) n = AD_V2 <$> runR aeR ad (n*2) - runR (ArrayEltRvec3 aeR) (AD_V3 ad) n = AD_V3 <$> runR aeR ad (n*3) - runR (ArrayEltRvec4 aeR) (AD_V4 ad) n = AD_V4 <$> runR aeR ad (n*4) - runR (ArrayEltRvec8 aeR) (AD_V8 ad) n = AD_V8 <$> runR aeR ad (n*8) - runR (ArrayEltRvec16 aeR) (AD_V16 ad) n = AD_V16 <$> runR aeR ad (n*16) - -- - runR ArrayEltRint ad@(AD_Int ua) n = AD_Int <$> peekR ad ua n - runR ArrayEltRint8 ad@(AD_Int8 ua) n = AD_Int8 <$> peekR ad ua n - runR ArrayEltRint16 ad@(AD_Int16 ua) n = AD_Int16 <$> peekR ad ua n - runR ArrayEltRint32 ad@(AD_Int32 ua) n = AD_Int32 <$> peekR ad ua n - runR ArrayEltRint64 ad@(AD_Int64 ua) n = AD_Int64 <$> peekR ad ua n - runR ArrayEltRword ad@(AD_Word ua) n = AD_Word <$> peekR ad ua n - runR ArrayEltRword8 ad@(AD_Word8 ua) n = AD_Word8 <$> peekR ad ua n - runR ArrayEltRword16 ad@(AD_Word16 ua) n = AD_Word16 <$> peekR ad ua n - runR ArrayEltRword32 ad@(AD_Word32 ua) n = AD_Word32 <$> peekR ad ua n - runR ArrayEltRword64 ad@(AD_Word64 ua) n = AD_Word64 <$> peekR ad ua n - runR ArrayEltRcshort ad@(AD_CShort ua) n = AD_CShort <$> peekR ad ua n - runR ArrayEltRcushort ad@(AD_CUShort ua) n = AD_CUShort <$> peekR ad ua n - runR ArrayEltRcint ad@(AD_CInt ua) n = AD_CInt <$> peekR ad ua n - runR ArrayEltRcuint ad@(AD_CUInt ua) n = AD_CUInt <$> peekR ad ua n - runR ArrayEltRclong ad@(AD_CLong ua) n = AD_CLong <$> peekR ad ua n - runR ArrayEltRculong ad@(AD_CULong ua) n = AD_CULong <$> peekR ad ua n - runR ArrayEltRcllong ad@(AD_CLLong ua) n = AD_CLLong <$> peekR ad ua n - runR ArrayEltRcullong ad@(AD_CULLong ua) n = AD_CULLong <$> peekR ad ua n - runR ArrayEltRhalf ad@(AD_Half ua) n = AD_Half <$> peekR ad ua n - runR ArrayEltRfloat ad@(AD_Float ua) n = AD_Float <$> peekR ad ua n - runR ArrayEltRdouble ad@(AD_Double ua) n = AD_Double <$> peekR ad ua n - runR ArrayEltRcfloat ad@(AD_CFloat ua) n = AD_CFloat <$> peekR ad ua n - runR ArrayEltRcdouble ad@(AD_CDouble ua) n = AD_CDouble <$> peekR ad ua n - runR ArrayEltRbool ad@(AD_Bool ua) n = AD_Bool <$> peekR ad ua n - runR ArrayEltRchar ad@(AD_Char ua) n = AD_Char <$> peekR ad ua n - runR ArrayEltRcchar ad@(AD_CChar ua) n = AD_CChar <$> peekR ad ua n - runR ArrayEltRcschar ad@(AD_CSChar ua) n = AD_CSChar <$> peekR ad ua n - runR ArrayEltRcuchar ad@(AD_CUChar ua) n = AD_CUChar <$> peekR ad ua n - in - Array sh <$> runR arrayElt adata (R.size sh) + n = size shr sh + runR :: TypeR e -> ArrayData e -> IO (ArrayData e) + runR TupRunit !() = return () + runR (TupRpair !t1 !t2) (!ad1, !ad2) = (,) <$> runR t1 ad1 <*> runR t2 ad2 + runR (TupRsingle !t) !ad = + case t of + SingleScalarType s -> peekR s ad n + VectorScalarType (VectorType w s) + | SingleArrayDict <- singleArrayDict s -> peekR s ad (n * w) + Array sh <$> runR tp adata + -- | Clone an array into a newly allocated array on the device. -- cloneArrayAsync - :: (Shape sh, Elt e) - => Stream + :: ArrayR (Array sh e) -> Array sh e - -> LLVM PTX (Array sh e) -cloneArrayAsync stream arr@(Array _ src) = do - out@(Array _ dst) <- allocateRemote sh - copyR arrayElt src dst (size sh) - return out + -> Par PTX (Future (Array sh e)) +cloneArrayAsync repr@(ArrayR shr tp) arr@(Array _ src) = do + Array _ dst <- allocateRemote repr sh + Array sh `liftF` copyR tp src dst where - sh = shape arr + sh = shape arr + n = size shr sh - copyR :: ArrayEltR e -> ArrayData e -> ArrayData e -> Int -> LLVM PTX () - copyR ArrayEltRunit _ _ _ = return () - copyR (ArrayEltRpair aeR1 aeR2) ad1 ad2 n = copyR aeR1 (fstArrayData ad1) (fstArrayData ad2) n >> - copyR aeR2 (sndArrayData ad1) (sndArrayData ad2) n - -- - copyR (ArrayEltRvec2 aeR) (AD_V2 ad1) (AD_V2 ad2) n = copyR aeR ad1 ad2 (n*2) - copyR (ArrayEltRvec3 aeR) (AD_V3 ad1) (AD_V3 ad2) n = copyR aeR ad1 ad2 (n*3) - copyR (ArrayEltRvec4 aeR) (AD_V4 ad1) (AD_V4 ad2) n = copyR aeR ad1 ad2 (n*4) - copyR (ArrayEltRvec8 aeR) (AD_V8 ad1) (AD_V8 ad2) n = copyR aeR ad1 ad2 (n*8) - copyR (ArrayEltRvec16 aeR) (AD_V16 ad1) (AD_V16 ad2) n = copyR aeR ad1 ad2 (n*16) - -- - copyR ArrayEltRint ad1 ad2 n = copyPrim ad1 ad2 n - copyR ArrayEltRint8 ad1 ad2 n = copyPrim ad1 ad2 n - copyR ArrayEltRint16 ad1 ad2 n = copyPrim ad1 ad2 n - copyR ArrayEltRint32 ad1 ad2 n = copyPrim ad1 ad2 n - copyR ArrayEltRint64 ad1 ad2 n = copyPrim ad1 ad2 n - copyR ArrayEltRword ad1 ad2 n = copyPrim ad1 ad2 n - copyR ArrayEltRword8 ad1 ad2 n = copyPrim ad1 ad2 n - copyR ArrayEltRword16 ad1 ad2 n = copyPrim ad1 ad2 n - copyR ArrayEltRword32 ad1 ad2 n = copyPrim ad1 ad2 n - copyR ArrayEltRword64 ad1 ad2 n = copyPrim ad1 ad2 n - copyR ArrayEltRhalf ad1 ad2 n = copyPrim ad1 ad2 n - copyR ArrayEltRfloat ad1 ad2 n = copyPrim ad1 ad2 n - copyR ArrayEltRdouble ad1 ad2 n = copyPrim ad1 ad2 n - copyR ArrayEltRbool ad1 ad2 n = copyPrim ad1 ad2 n - copyR ArrayEltRchar ad1 ad2 n = copyPrim ad1 ad2 n - copyR ArrayEltRcshort ad1 ad2 n = copyPrim ad1 ad2 n - copyR ArrayEltRcushort ad1 ad2 n = copyPrim ad1 ad2 n - copyR ArrayEltRcint ad1 ad2 n = copyPrim ad1 ad2 n - copyR ArrayEltRcuint ad1 ad2 n = copyPrim ad1 ad2 n - copyR ArrayEltRclong ad1 ad2 n = copyPrim ad1 ad2 n - copyR ArrayEltRculong ad1 ad2 n = copyPrim ad1 ad2 n - copyR ArrayEltRcllong ad1 ad2 n = copyPrim ad1 ad2 n - copyR ArrayEltRcullong ad1 ad2 n = copyPrim ad1 ad2 n - copyR ArrayEltRcfloat ad1 ad2 n = copyPrim ad1 ad2 n - copyR ArrayEltRcdouble ad1 ad2 n = copyPrim ad1 ad2 n - copyR ArrayEltRcchar ad1 ad2 n = copyPrim ad1 ad2 n - copyR ArrayEltRcschar ad1 ad2 n = copyPrim ad1 ad2 n - copyR ArrayEltRcuchar ad1 ad2 n = copyPrim ad1 ad2 n + copyR :: TypeR s -> ArrayData s -> ArrayData s -> Par PTX (Future (ArrayData s)) + copyR TupRunit !_ !_ = newFull () + copyR (TupRpair !t1 !t2) !(ad1, ad2) !(ad1', ad2') = liftF2 (,) (copyR t1 ad1 ad1') (copyR t2 ad2 ad2') + copyR (TupRsingle !t) !ad !ad' = + case t of + SingleScalarType s -> copyPrim s ad ad' n + VectorScalarType (VectorType w s) + | SingleArrayDict <- singleArrayDict s -> copyPrim s ad ad' (n * w) copyPrim - :: (ArrayElt e, ArrayPtrs e ~ Ptr a, Typeable e, Storable a, Typeable a) - => ArrayData e - -> ArrayData e + :: SingleType s + -> ArrayData s + -> ArrayData s -> Int - -> LLVM PTX () - copyPrim !a1 !a2 !m = Prim.copyArrayAsync stream m a1 a2 + -> Par PTX (Future (ArrayData s)) + copyPrim !s !a1 !a2 !m = Prim.copyArrayAsync s m a1 a2 + + liftF :: Async arch + => (a -> b) + -> Par arch (FutureR arch a) + -> Par arch (FutureR arch b) + liftF f x = do + r <- new + x' <- x + put r . f =<< get x' -- don't create a new execution stream for this + return r + + liftF2 :: Async arch + => (a -> b -> c) + -> Par arch (FutureR arch a) + -> Par arch (FutureR arch b) + -> Par arch (FutureR arch c) + liftF2 f x y = do + r <- new + x' <- spawn x + y' <- spawn y + fork $ put r =<< liftM2 f (get x') (get y') + return r diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Array/Prim.hs b/src/Data/Array/Accelerate/LLVM/PTX/Array/Prim.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/Array/Prim.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/Array/Prim.hs @@ -2,17 +2,17 @@ {-# LANGUAGE CPP #-} {-# LANGUAGE DataKinds #-} {-# LANGUAGE GADTs #-} +{-# LANGUAGE MagicHash #-} {-# LANGUAGE RecordWildCards #-} {-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} {-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UnboxedTuples #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Array.Prim --- Copyright : [2014..2017] Trevor L. McDonell --- [2014..2014] Vinod Grover (NVIDIA Corporation) +-- Copyright : [2014..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- @@ -20,62 +20,59 @@ module Data.Array.Accelerate.LLVM.PTX.Array.Prim ( mallocArray, - memsetArray, memsetArrayAsync, - useArray, useArrayAsync, - indexArray, - peekArray, peekArrayR, peekArrayAsync, peekArrayAsyncR, - pokeArray, pokeArrayR, pokeArrayAsync, pokeArrayAsyncR, - copyArray, copyArrayR, copyArrayAsync, copyArrayAsyncR, - copyArrayPeer, copyArrayPeerR, copyArrayPeerAsync, copyArrayPeerAsyncR, + useArrayAsync, + indexArrayAsync, + peekArrayAsync, + pokeArrayAsync, + copyArrayAsync, + -- copyArrayPeerAsync, + memsetArrayAsync, withDevicePtr, ) where --- accelerate import Data.Array.Accelerate.Array.Data +import Data.Array.Accelerate.Array.Unique import Data.Array.Accelerate.Error -import Data.Array.Accelerate.Lifetime +import Data.Array.Accelerate.Lifetime hiding ( withLifetime ) +import Data.Array.Accelerate.Representation.Elt +import Data.Array.Accelerate.Representation.Type import Data.Array.Accelerate.Type import Data.Array.Accelerate.LLVM.State -import Data.Array.Accelerate.LLVM.PTX.Context import Data.Array.Accelerate.LLVM.PTX.Target +import Data.Array.Accelerate.LLVM.PTX.Execute.Async import Data.Array.Accelerate.LLVM.PTX.Execute.Event import Data.Array.Accelerate.LLVM.PTX.Execute.Stream -import Data.Array.Accelerate.LLVM.PTX.Array.Table import Data.Array.Accelerate.LLVM.PTX.Array.Remote as Remote import qualified Data.Array.Accelerate.LLVM.PTX.Debug as Debug --- CUDA import qualified Foreign.CUDA.Driver as CUDA import qualified Foreign.CUDA.Driver.Stream as CUDA --- standard library -import Control.Exception import Control.Monad -import Control.Monad.State -import Data.Typeable -import Foreign.Ptr -import Foreign.Storable -import GHC.TypeLits +import Control.Monad.Reader +import Data.IORef +import GHC.Base import Text.Printf -import Prelude hiding ( lookup ) +import Prelude -- | Allocate a device-side array associated with the given host array. If the -- allocation fails due to a memory error, we attempt some last-ditch memory --- cleanup before trying again. +-- cleanup before trying again. If it still fails; error. -- {-# INLINEABLE mallocArray #-} mallocArray - :: forall e a. (ArrayElt e, ArrayPtrs e ~ Ptr a, Storable a, Typeable a, Typeable e) - => Int + :: HasCallStack + => SingleType e + -> Int -> ArrayData e -> LLVM PTX () -mallocArray !n !ad = do - message ("mallocArray: " ++ showBytes (n * sizeOf (undefined::a))) - void $ malloc ad n False +mallocArray !t !n !ad = do + message ("mallocArray: " ++ showBytes (n * bytesElt (TupRsingle (SingleScalarType t)))) + void $ Remote.malloc t ad n False -- | A combination of 'mallocArray' and 'pokeArray', that allocates remotes @@ -83,259 +80,150 @@ -- allocator that the host-side array is frozen, and thus it is safe to evict -- the remote memory and re-upload the data at any time. -- -{-# INLINEABLE useArray #-} -useArray - :: forall e a. (ArrayElt e, ArrayPtrs e ~ Ptr a, Storable a, Typeable a, Typeable e) - => Int - -> ArrayData e - -> LLVM PTX () -useArray !n !ad = - blocking $ \st -> useArrayAsync st n ad - {-# INLINEABLE useArrayAsync #-} useArrayAsync - :: forall e a. (ArrayElt e, ArrayPtrs e ~ Ptr a, Storable a, Typeable a, Typeable e) - => Stream + :: HasCallStack + => SingleType e -> Int -> ArrayData e - -> LLVM PTX () -useArrayAsync !st !n !ad = do - alloc <- malloc ad n True - when alloc $ pokeArrayAsync st n ad + -> Par PTX (Future (ArrayData e)) +useArrayAsync !t !n !ad = do + message ("useArrayAsync: " ++ showBytes (n * bytesElt (TupRsingle (SingleScalarType t)))) + alloc <- liftPar $ Remote.malloc t ad n True + if alloc + then pokeArrayAsync t n ad + else newFull ad -- | Copy data from the host to an existing array on the device -- -{-# INLINEABLE pokeArray #-} -pokeArray - :: forall e a. (ArrayElt e, ArrayPtrs e ~ Ptr a, Typeable e, Storable a, Typeable a) - => Int - -> ArrayData e - -> LLVM PTX () -pokeArray !n !ad = - blocking $ \st -> pokeArrayAsync st n ad - {-# INLINEABLE pokeArrayAsync #-} pokeArrayAsync - :: forall e a. (ArrayElt e, ArrayPtrs e ~ Ptr a, Typeable e, Storable a, Typeable a) - => Stream - -> Int - -> ArrayData e - -> LLVM PTX () -pokeArrayAsync !stream !n !ad = do - let !src = CUDA.HostPtr (ptrsOfArrayData ad) - !bytes = n * sizeOf (undefined :: a) - !st = unsafeGetValue stream - -- - withDevicePtr ad $ \dst -> - nonblocking stream $ - transfer "pokeArray" bytes (Just st) $ CUDA.pokeArrayAsync n src dst (Just st) - liftIO (touchLifetime stream) - liftIO (Debug.didCopyBytesToRemote (fromIntegral bytes)) - - -{-# INLINEABLE pokeArrayR #-} -pokeArrayR - :: forall e a. (ArrayElt e, ArrayPtrs e ~ Ptr a, Typeable e, Typeable a, Storable a) - => Int - -> Int - -> ArrayData e - -> LLVM PTX () -pokeArrayR !from !n !ad = - blocking $ \st -> pokeArrayAsyncR st from n ad - -{-# INLINEABLE pokeArrayAsyncR #-} -pokeArrayAsyncR - :: forall e a. (ArrayElt e, ArrayPtrs e ~ Ptr a, Typeable e, Typeable a, Storable a) - => Stream - -> Int + :: HasCallStack + => SingleType e -> Int -> ArrayData e - -> LLVM PTX () -pokeArrayAsyncR !stream !from !n !ad = do - let !bytes = n * sizeOf (undefined :: a) - !offset = from * sizeOf (undefined :: a) - !src = CUDA.HostPtr (ptrsOfArrayData ad) - !st = unsafeGetValue stream - -- - withDevicePtr ad $ \dst -> - nonblocking stream $ - transfer "pokeArray" bytes (Just st) $ - CUDA.pokeArrayAsync n (src `CUDA.plusHostPtr` offset) (dst `CUDA.plusDevPtr` offset) (Just st) - liftIO (touchLifetime stream) - liftIO (Debug.didCopyBytesToRemote (fromIntegral bytes)) + -> Par PTX (Future (ArrayData e)) +pokeArrayAsync !t !n !ad + | SingleArrayDict <- singleArrayDict t + , SingleDict <- singleDict t + = do + let !src = CUDA.HostPtr (unsafeUniqueArrayPtr ad) + !bytes = n * bytesElt (TupRsingle (SingleScalarType t)) + -- + stream <- asks ptxStream + result <- liftPar $ + withLifetime stream $ \st -> + withDevicePtr t ad $ \dst -> + nonblocking stream $ do + transfer "pokeArray" bytes (Just st) $ do + CUDA.pokeArrayAsync n src dst (Just st) + Debug.didCopyBytesToRemote (fromIntegral bytes) + return ad + -- + return result --- | Read elements from an array at the given row-major index +-- | Read an element from an array at the given row-major index. -- -{-# INLINEABLE indexArray #-} -indexArray - :: forall e a. (ArrayElt e, ArrayPtrs e ~ Ptr a, Typeable e, Typeable a, Storable a) - => ArrayData e - -> Int - -> LLVM PTX a -indexArray !ad !i = - blocking $ \stream -> - withDevicePtr ad $ \src -> liftIO $ - bracket (CUDA.mallocHostArray [] 1) CUDA.freeHost $ \dst -> do - let !st = unsafeGetValue stream - !bytes = sizeOf (undefined::a) +-- This copies the data via a temporary array on the host, so that packed AoS +-- elements can be copied in a single transfer. +-- +{-# INLINEABLE indexArrayAsync #-} +indexArrayAsync + :: HasCallStack + => Int -- actual number of values per element (i.e. this is >1 for SIMD types) + -> SingleType e + -> ArrayData e + -> Int -- element index + -> Par PTX (Future (ArrayData e)) +indexArrayAsync !n !t !ad_src !i + | SingleArrayDict <- singleArrayDict t + , SingleDict <- singleDict t + = do + ad_dst <- liftIO $ newArrayData (TupRsingle $ SingleScalarType t) n + let !bytes = n * bytesElt (TupRsingle (SingleScalarType t)) + !dst = CUDA.HostPtr (unsafeUniqueArrayPtr ad_dst) -- - message $ "indexArray: " ++ showBytes bytes - Debug.didCopyBytesFromRemote (fromIntegral bytes) - CUDA.peekArrayAsync 1 (src `CUDA.advanceDevPtr` i) dst (Just st) - CUDA.block st - touchLifetime stream - r <- peek (CUDA.useHostPtr dst) - return (Nothing, r) + stream <- asks ptxStream + result <- liftPar $ + withLifetime stream $ \st -> + withDevicePtr t ad_src $ \src -> + nonblocking stream $ do + transfer "indexArray" bytes (Just st) $ do + CUDA.peekArrayAsync n (src `CUDA.advanceDevPtr` (i*n)) dst (Just st) + Debug.didCopyBytesFromRemote (fromIntegral bytes) + return ad_dst + -- + return result -- | Copy data from the device into the associated host-side Accelerate array -- -{-# INLINEABLE peekArray #-} -peekArray - :: forall e a. (ArrayElt e, ArrayPtrs e ~ Ptr a, Typeable e, Typeable a, Storable a) - => Int - -> ArrayData e - -> LLVM PTX () -peekArray !n !ad = - blocking $ \st -> peekArrayAsync st n ad - {-# INLINEABLE peekArrayAsync #-} peekArrayAsync - :: forall e a. (ArrayElt e, ArrayPtrs e ~ Ptr a, Typeable e, Typeable a, Storable a) - => Stream - -> Int - -> ArrayData e - -> LLVM PTX () -peekArrayAsync !stream !n !ad = do - let !bytes = n * sizeOf (undefined :: a) - !dst = CUDA.HostPtr (ptrsOfArrayData ad) - !st = unsafeGetValue stream - -- - withDevicePtr ad $ \src -> - nonblocking stream $ - transfer "peekArray" bytes (Just st) $ CUDA.peekArrayAsync n src dst (Just st) - liftIO (touchLifetime stream) - liftIO (Debug.didCopyBytesFromRemote (fromIntegral bytes)) - -{-# INLINEABLE peekArrayR #-} -peekArrayR - :: forall e a. (ArrayElt e, ArrayPtrs e ~ Ptr a, Typeable a, Typeable e, Storable a) - => Int - -> Int - -> ArrayData e - -> LLVM PTX () -peekArrayR !from !n !ad = - blocking $ \st -> peekArrayAsyncR st from n ad - -{-# INLINEABLE peekArrayAsyncR #-} -peekArrayAsyncR - :: forall e a. (ArrayElt e, ArrayPtrs e ~ Ptr a, Typeable e, Typeable a, Storable a) - => Stream - -> Int + :: HasCallStack + => SingleType e -> Int -> ArrayData e - -> LLVM PTX () -peekArrayAsyncR !stream !from !n !ad = do - let !bytes = n * sizeOf (undefined :: a) - !offset = from * sizeOf (undefined :: a) - !dst = CUDA.HostPtr (ptrsOfArrayData ad) - !st = unsafeGetValue stream - -- - withDevicePtr ad $ \src -> - nonblocking stream $ - transfer "peekArray" bytes (Just st) $ - CUDA.peekArrayAsync n (src `CUDA.plusDevPtr` offset) (dst `CUDA.plusHostPtr` offset) (Just st) - liftIO (touchLifetime stream) - liftIO (Debug.didCopyBytesFromRemote (fromIntegral bytes)) + -> Par PTX (Future (ArrayData e)) +peekArrayAsync !t !n !ad + | SingleArrayDict <- singleArrayDict t + , SingleDict <- singleDict t + = do + let !bytes = n * bytesElt (TupRsingle (SingleScalarType t)) + !dst = CUDA.HostPtr (unsafeUniqueArrayPtr ad) + -- + stream <- asks ptxStream + result <- liftPar $ + withLifetime stream $ \st -> + withDevicePtr t ad $ \src -> + nonblocking stream $ do + transfer "peekArray" bytes (Just st) $ do + CUDA.peekArrayAsync n src dst (Just st) + Debug.didCopyBytesFromRemote (fromIntegral bytes) + return ad + -- + return result -- | Copy data between arrays in the same context -- -{-# INLINEABLE copyArray #-} -copyArray - :: forall e a. (ArrayElt e, ArrayPtrs e ~ Ptr a, Typeable e, Storable a, Typeable a) - => Int - -> ArrayData e - -> ArrayData e - -> LLVM PTX () -copyArray !n !src !dst = - blocking $ \st -> copyArrayAsync st n src dst - {-# INLINEABLE copyArrayAsync #-} copyArrayAsync - :: forall e a. (ArrayElt e, ArrayPtrs e ~ Ptr a, Typeable e, Storable a, Typeable a) - => Stream - -> Int - -> ArrayData e - -> ArrayData e - -> LLVM PTX () -copyArrayAsync !stream !n !ad_src !ad_dst = do - let !bytes = n * sizeOf (undefined :: a) - !st = unsafeGetValue stream - -- - withDevicePtr ad_src $ \src -> do - e <- withDevicePtr ad_dst $ \dst -> do - (e,()) <- nonblocking stream - $ transfer "copyArray" bytes (Just st) $ CUDA.copyArrayAsync n src dst (Just st) - return (e,e) - return (e,()) - liftIO (touchLifetime stream) - -{-# INLINEABLE copyArrayR #-} -copyArrayR - :: forall e a. (ArrayElt e, ArrayPtrs e ~ Ptr a, Typeable e, Storable a, Typeable a) - => Int - -> Int - -> ArrayData e - -> ArrayData e - -> LLVM PTX () -copyArrayR !from !n !src !dst = - blocking $ \st -> copyArrayAsyncR st from n src dst - -{-# INLINEABLE copyArrayAsyncR #-} -copyArrayAsyncR - :: forall e a. (ArrayElt e, ArrayPtrs e ~ Ptr a, Typeable e, Storable a, Typeable a) - => Stream - -> Int + :: HasCallStack + => SingleType e -> Int -> ArrayData e -> ArrayData e - -> LLVM PTX () -copyArrayAsyncR !stream !from !n !ad_src !ad_dst = do - let !bytes = n * sizeOf (undefined :: a) - !offset = from * sizeOf (undefined :: a) - !st = unsafeGetValue stream - -- - withDevicePtr ad_src $ \src -> do - e <- withDevicePtr ad_dst $ \dst -> do - (e,()) <- nonblocking stream - $ transfer "copyArray" bytes (Just st) - $ CUDA.copyArrayAsync n (src `CUDA.plusDevPtr` offset) (dst `CUDA.plusDevPtr` offset) (Just st) - return (e,e) - return (e,()) - liftIO (touchLifetime stream) + -> Par PTX (Future (ArrayData e)) +copyArrayAsync !t !n !ad_src !ad_dst + | SingleArrayDict <- singleArrayDict t + , SingleDict <- singleDict t + = do + let !bytes = n * bytesElt (TupRsingle (SingleScalarType t)) + -- + stream <- asks ptxStream + result <- liftPar $ + withLifetime stream $ \st -> + withDevicePtr t ad_src $ \src -> + withDevicePtr t ad_dst $ \dst -> do + (e,r) <- nonblocking stream $ do + transfer "copyArray" bytes (Just st) $ CUDA.copyArrayAsync n src dst (Just st) + return ad_dst + return (e, (e,r)) + -- + return result +{-- -- | Copy data from one device context into a _new_ array on the second context. -- It is an error if the destination array already exists. -- -{-# INLINEABLE copyArrayPeer #-} -copyArrayPeer - :: forall e a. (ArrayElt e, ArrayPtrs e ~ Ptr a, Storable a, Typeable a) - => Context -- destination context - -> MemoryTable -- destination memory table - -> Int - -> ArrayData e - -> LLVM PTX () -copyArrayPeer !ctx2 !mt2 !n !ad = - blocking $ \st -> copyArrayPeerAsync ctx2 mt2 st n ad - {-# INLINEABLE copyArrayPeerAsync #-} copyArrayPeerAsync - :: forall e a. (ArrayElt e, ArrayPtrs e ~ Ptr a, Storable a, Typeable a) - => Context -- destination context + :: SingleType e + -> Context -- destination context -> MemoryTable -- destination memory table -> Stream -> Int @@ -343,8 +231,8 @@ -> LLVM PTX () copyArrayPeerAsync = error "copyArrayPeerAsync" {-- -copyArrayPeerAsync !ctx2 !mt2 !st !n !ad = do - let !bytes = n * sizeOf (undefined :: a) +copyArrayPeerAsync !t !ctx2 !mt2 !st !n !ad = do + let !bytes = n * sizeOfSingleType t src <- devicePtr mt1 ad dst <- mallocArray ctx2 mt2 n ad transfer "copyArrayPeer" bytes (Just st) $ @@ -354,72 +242,76 @@ -- | Copy part of an array from one device context to another. Both source and -- destination arrays must exist. -- -{-# INLINEABLE copyArrayPeerR #-} -copyArrayPeerR - :: forall e a. (ArrayElt e, ArrayPtrs e ~ Ptr a, Storable a, Typeable a) - => Context -- destination context - -> MemoryTable -- destination memory table - -> Int - -> Int - -> ArrayData e - -> LLVM PTX () -copyArrayPeerR !ctx2 !mt2 !from !n !ad = - blocking $ \st -> copyArrayPeerAsyncR ctx2 mt2 st from n ad - {-# INLINEABLE copyArrayPeerAsyncR #-} -copyArrayPeerAsyncR - :: forall e a. (ArrayElt e, ArrayPtrs e ~ Ptr a, Storable a, Typeable a) - => Context -- destination context +copyArrayPeerAsync + :: SingleType e + -> Context -- destination context -> MemoryTable -- destination memory table -> Stream -> Int -> Int -> ArrayData e -> LLVM PTX () -copyArrayPeerAsyncR = error "copyArrayPeerAsyncR" +copyArrayPeerAsync = error "copyArrayPeerAsyncR" {-- -copyArrayPeerAsyncR !ctx2 !mt2 !st !from !n !ad = do - let !bytes = n * sizeOf (undefined :: a) - !offset = from * sizeOf (undefined :: a) +copyArrayPeerAsyncR !t !ctx2 !mt2 !st !from !n !ad = do + let !bytes = n * sizeOfSingleType t + !offset = from * sizeOfSingleType t src <- devicePtr mt1 ad :: IO (CUDA.DevicePtr a) dst <- devicePtr mt2 ad :: IO (CUDA.DevicePtr a) transfer "copyArrayPeer" bytes (Just st) $ CUDA.copyArrayPeerAsync n (src `CUDA.plusDevPtr` offset) (deviceContext ctx1) (dst `CUDA.plusDevPtr` offset) (deviceContext ctx2) (Just st) --} - +--} -- | Set elements of the array to the specified value. Only 8-, 16-, and 32-bit -- values are supported. -- -{-# INLINEABLE memsetArray #-} -memsetArray - :: forall e a. (ArrayElt e, ArrayPtrs e ~ Ptr a, Typeable e, Typeable a, Storable a, BitSize a <= 32) - => Int - -> a - -> ArrayData e - -> LLVM PTX () -memsetArray !n !v !ad = - blocking $ \st -> memsetArrayAsync st n v ad - {-# INLINEABLE memsetArrayAsync #-} memsetArrayAsync - :: forall e a. (ArrayElt e, ArrayPtrs e ~ Ptr a, Typeable e, Typeable a, Storable a, BitSize a <= 32) - => Stream + :: HasCallStack + => SingleType e -> Int - -> a + -> ScalarArrayDataR e -> ArrayData e - -> LLVM PTX () -memsetArrayAsync !stream !n !v !ad = do - let !bytes = n * sizeOf (undefined :: a) - !st = unsafeGetValue stream - -- - withDevicePtr ad $ \ptr -> - nonblocking stream $ - transfer "memset" bytes (Just st) $ CUDA.memsetAsync ptr n v (Just st) - liftIO (touchLifetime stream) + -> Par PTX (Future (ArrayData e)) +memsetArrayAsync !t !n !v !ad + | SingleArrayDict <- singleArrayDict t + , SingleDict <- singleDict t + = do + let !bytes = n * bytesElt (TupRsingle (SingleScalarType t)) + -- + stream <- asks ptxStream + result <- liftPar $ + withLifetime stream $ \st -> + withDevicePtr t ad $ \ptr -> + nonblocking stream $ do + transfer "memset" bytes (Just st) $ CUDA.memsetAsync ptr n v (Just st) + return ad + -- + return result +-- Auxiliary +-- --------- + +-- | Lookup the device memory associated with a given host array and do +-- something with it. +-- +{-# INLINEABLE withDevicePtr #-} +withDevicePtr + :: HasCallStack + => SingleType e + -> ArrayData e + -> (CUDA.DevicePtr (ScalarArrayDataR e) -> LLVM PTX (Maybe Event, r)) + -> LLVM PTX r +withDevicePtr !t !ad !f = do + mr <- withRemote t ad f + case mr of + Nothing -> internalError "array does not exist on the device" + Just r -> return r + {-- -- | Lookup the device memory associated with a given host array -- @@ -438,45 +330,36 @@ --} --} --- Auxiliary --- --------- - --- | Lookup the device memory associated with a given host array and do --- something with it. --- -{-# INLINEABLE withDevicePtr #-} -withDevicePtr - :: (ArrayElt e, ArrayPtrs e ~ Ptr a, Typeable e, Typeable a, Storable a) - => ArrayData e - -> (CUDA.DevicePtr a -> LLVM PTX (Maybe Event, r)) - -> LLVM PTX r -withDevicePtr !ad !f = do - mr <- withRemote ad f - case mr of - Nothing -> $internalError "withDevicePtr" "array does not exist on the device" - Just r -> return r - --- | Execute the given operation in a new stream, and wait for the operation to --- complete before returning. --- -{-# INLINE blocking #-} -blocking :: (Stream -> LLVM PTX a) -> LLVM PTX a -blocking !f = - streaming f $ \e r -> do - liftIO $ block e - return r - -- | Execute a (presumable asynchronous) operation and return the result -- together with an event recorded immediately afterwards in the given stream. -- {-# INLINE nonblocking #-} -nonblocking :: Stream -> LLVM PTX a -> LLVM PTX (Maybe Event, a) -nonblocking !stream !f = do - r <- f - e <- waypoint stream - return (Just e, r) +nonblocking :: Stream -> LLVM PTX a -> LLVM PTX (Maybe Event, Future a) +nonblocking !stream !action = do + result <- action + event <- waypoint stream + ready <- liftIO (query event) + if ready + then do + future <- Future <$> liftIO (newIORef (Full result)) + return (Nothing, future) + else do + future <- Future <$> liftIO (newIORef (Pending event Nothing result)) + return (Just event, future) +{-# INLINE withLifetime #-} +withLifetime :: MonadIO m => Lifetime a -> (a -> m b) -> m b +withLifetime (Lifetime ref _ a) f = do + r <- f a + liftIO (touchIORef ref) + return r + +{-# INLINE touchIORef #-} +touchIORef :: IORef a -> IO () +touchIORef r = IO $ \s -> case touch# r s of s' -> (# s', () #) + + -- Debug -- ----- @@ -494,13 +377,13 @@ {-# INLINE transfer #-} transfer :: MonadIO m => String -> Int -> Maybe CUDA.Stream -> IO () -> m () -transfer name bytes stream action - = let showRate x t = Debug.showFFloatSIBase (Just 3) 1024 (fromIntegral x / t) "B/s" - msg wall cpu gpu = printf "gc: %s: %s bytes @ %s, %s" - name - (showBytes bytes) - (showRate bytes wall) - (Debug.elapsed wall cpu gpu) - in - liftIO (Debug.timed Debug.dump_gc msg stream action) +transfer name bytes stream action = + let showRate x t = Debug.showFFloatSIBase (Just 3) 1024 (fromIntegral x / t) "B/s" + msg wall cpu gpu = printf "gc: %s: %s @ %s, %s" + name + (showBytes bytes) + (showRate bytes wall) + (Debug.elapsed wall cpu gpu) + in + liftIO (Debug.timed Debug.dump_gc msg stream action) diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Array/Remote.hs b/src/Data/Array/Accelerate/LLVM/PTX/Array/Remote.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/Array/Remote.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/Array/Remote.hs @@ -1,15 +1,17 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MagicHash #-} {-# LANGUAGE RecordWildCards #-} {-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeFamilies #-} {-# OPTIONS_GHC -fno-warn-orphans #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Array.Remote --- Copyright : [2014..2017] Trevor L. McDonell +-- Copyright : [2014..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- @@ -25,8 +27,12 @@ import {-# SOURCE #-} Data.Array.Accelerate.LLVM.PTX.Execute.Event import {-# SOURCE #-} Data.Array.Accelerate.LLVM.PTX.Execute.Stream -import Data.Array.Accelerate.Lifetime import Data.Array.Accelerate.Array.Data +import Data.Array.Accelerate.Array.Unique +import Data.Array.Accelerate.Lifetime +import Data.Array.Accelerate.Representation.Elt +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Type import qualified Data.Array.Accelerate.Array.Remote as Remote import qualified Data.Array.Accelerate.LLVM.PTX.Debug as Debug @@ -37,12 +43,12 @@ import Control.Exception import Control.Monad.State -import Data.Typeable -import Foreign.Ptr -import Foreign.Storable import Text.Printf +import GHC.Base +import GHC.Int + -- Events signal once a computation has completed -- instance Remote.Task (Maybe Event) where @@ -57,31 +63,35 @@ | otherwise = liftIO $ do ep <- try (CUDA.mallocArray n) case ep of - Right p -> do liftIO (Debug.didAllocateBytesRemote (fromIntegral n)) + Right p -> do liftIO (Debug.didAllocateBytesRemote (i64 n)) return (Just p) Left (ExitCode OutOfMemory) -> do return Nothing Left e -> do message ("malloc failed with error: " ++ show e) throwIO e - peekRemote n src ad = - let bytes = n * sizeOfPtr src - dst = CUDA.HostPtr (ptrsOfArrayData ad) - in - blocking $ \stream -> - withLifetime stream $ \st -> do - Debug.didCopyBytesFromRemote (fromIntegral bytes) - transfer "peekRemote" bytes (Just st) $ CUDA.peekArrayAsync n src dst (Just st) + peekRemote t n src ad + | SingleArrayDict <- singleArrayDict t + , SingleDict <- singleDict t + = let bytes = n * bytesElt (TupRsingle (SingleScalarType t)) + dst = CUDA.HostPtr (unsafeUniqueArrayPtr ad) + in + blocking $ \stream -> + withLifetime stream $ \st -> do + Debug.didCopyBytesFromRemote (i64 bytes) + transfer "peekRemote" bytes (Just st) $ CUDA.peekArrayAsync n src dst (Just st) - pokeRemote n dst ad = - let bytes = n * sizeOfPtr dst - src = CUDA.HostPtr (ptrsOfArrayData ad) - in - blocking $ \stream -> - withLifetime stream $ \st -> do - Debug.didCopyBytesToRemote (fromIntegral bytes) - transfer "pokeRemote" bytes (Just st) $ CUDA.pokeArrayAsync n src dst (Just st) + pokeRemote t n dst ad + | SingleArrayDict <- singleArrayDict t + , SingleDict <- singleDict t + = let bytes = n * bytesElt (TupRsingle (SingleScalarType t)) + src = CUDA.HostPtr (unsafeUniqueArrayPtr ad) + in + blocking $ \stream -> + withLifetime stream $ \st -> do + Debug.didCopyBytesToRemote (i64 bytes) + transfer "pokeRemote" bytes (Just st) $ CUDA.pokeArrayAsync n src dst (Just st) - castRemotePtr _ = CUDA.castDevPtr + castRemotePtr = CUDA.castDevPtr availableRemoteMem = liftIO $ fst `fmap` CUDA.getMemInfo totalRemoteMem = liftIO $ snd `fmap` CUDA.getMemInfo remoteAllocationSize = return 4096 @@ -94,27 +104,27 @@ -- {-# INLINEABLE malloc #-} malloc - :: (ArrayElt e, ArrayPtrs e ~ Ptr a, Typeable e, Typeable a, Storable a) - => ArrayData e + :: SingleType e + -> ArrayData e -> Int -> Bool -> LLVM PTX Bool -malloc !ad !n !frozen = do +malloc !tp !ad !n !frozen = do PTX{..} <- gets llvmTarget - Remote.malloc ptxMemoryTable ad frozen n + Remote.malloc ptxMemoryTable tp ad frozen n -- | Lookup up the remote array pointer for the given host-side array -- {-# INLINEABLE withRemote #-} withRemote - :: (ArrayElt e, ArrayPtrs e ~ Ptr a, Typeable e, Typeable a, Storable a) - => ArrayData e - -> (CUDA.DevicePtr a -> LLVM PTX (Maybe Event, r)) + :: SingleType e + -> ArrayData e + -> (CUDA.DevicePtr (ScalarArrayDataR e) -> LLVM PTX (Maybe Event, r)) -> LLVM PTX (Maybe r) -withRemote !ad !f = do +withRemote !tp !ad !f = do PTX{..} <- gets llvmTarget - Remote.withRemote ptxMemoryTable ad f + Remote.withRemote ptxMemoryTable tp ad f -- Auxiliary @@ -130,16 +140,21 @@ liftIO $ block e return r -{-# INLINE sizeOfPtr #-} -sizeOfPtr :: forall a. Storable a => CUDA.DevicePtr a -> Int -sizeOfPtr _ = sizeOf (undefined :: a) +{-# INLINE i64 #-} +i64 :: Int -> Int64 +i64 (I# i#) = I64# i# +{-# INLINE double #-} +double :: Int -> Double +double (I# i#) = D# (int2Double# i#) + + -- Debugging -- --------- {-# INLINE showBytes #-} showBytes :: Int -> String -showBytes x = Debug.showFFloatSIBase (Just 0) 1024 (fromIntegral x :: Double) "B" +showBytes x = Debug.showFFloatSIBase (Just 0) 1024 (double x) "B" {-# INLINE trace #-} trace :: String -> IO a -> IO a @@ -152,7 +167,7 @@ {-# INLINE transfer #-} transfer :: String -> Int -> Maybe CUDA.Stream -> IO () -> IO () transfer name bytes stream action - = let showRate x t = Debug.showFFloatSIBase (Just 3) 1024 (fromIntegral x / t) "B/s" + = let showRate x t = Debug.showFFloatSIBase (Just 3) 1024 (double x / t) "B/s" msg wall cpu gpu = printf "gc: %s: %s bytes @ %s, %s" name (showBytes bytes) diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Array/Table.hs b/src/Data/Array/Accelerate/LLVM/PTX/Array/Table.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/Array/Table.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/Array/Table.hs @@ -1,10 +1,10 @@ {-# LANGUAGE BangPatterns #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Array.Table --- Copyright : [2014..2017] Trevor L. McDonell +-- Copyright : [2014..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- @@ -56,5 +56,4 @@ {-# INLINE message #-} message :: String -> IO () message s = s `trace` return () - diff --git a/src/Data/Array/Accelerate/LLVM/PTX/CodeGen.hs b/src/Data/Array/Accelerate/LLVM/PTX/CodeGen.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/CodeGen.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/CodeGen.hs @@ -1,11 +1,10 @@ {-# OPTIONS_GHC -fno-warn-orphans #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.CodeGen --- Copyright : [2014..2017] Trevor L. McDonell --- [2014..2014] Vinod Grover (NVIDIA Corporation) +-- Copyright : [2014..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- @@ -27,21 +26,20 @@ import Data.Array.Accelerate.LLVM.PTX.CodeGen.Map import Data.Array.Accelerate.LLVM.PTX.CodeGen.Permute import Data.Array.Accelerate.LLVM.PTX.CodeGen.Scan +import Data.Array.Accelerate.LLVM.PTX.CodeGen.Stencil +import Data.Array.Accelerate.LLVM.PTX.CodeGen.Transform import Data.Array.Accelerate.LLVM.PTX.Target instance Skeleton PTX where - map ptx _ = mkMap ptx - generate ptx _ = mkGenerate ptx - fold ptx _ = mkFold ptx - fold1 ptx _ = mkFold1 ptx - foldSeg ptx _ = mkFoldSeg ptx - fold1Seg ptx _ = mkFold1Seg ptx - scanl ptx _ = mkScanl ptx - scanl1 ptx _ = mkScanl1 ptx - scanl' ptx _ = mkScanl' ptx - scanr ptx _ = mkScanr ptx - scanr1 ptx _ = mkScanr1 ptx - scanr' ptx _ = mkScanr' ptx - permute ptx _ = mkPermute ptx + map _ = mkMap + generate _ = mkGenerate + transform _ = mkTransform + fold _ = mkFold + foldSeg _ = mkFoldSeg + scan _ = mkScan + scan' _ = mkScan' + permute _ = mkPermute + stencil1 _ = mkStencil1 + stencil2 _ = mkStencil2 diff --git a/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Base.hs b/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Base.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Base.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Base.hs @@ -1,15 +1,17 @@ +{-# LANGUAGE CPP #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE ViewPatterns #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.CodeGen.Base --- Copyright : [2014..2017] Trevor L. McDonell +-- Copyright : [2014..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- @@ -22,15 +24,16 @@ -- Thread identifiers blockDim, gridDim, threadIdx, blockIdx, warpSize, gridSize, globalThreadIdx, - gangParam, -- Other intrinsics laneId, warpId, laneMask_eq, laneMask_lt, laneMask_le, laneMask_gt, laneMask_ge, atomicAdd_f, + nanosleep, -- Barriers and synchronisation - __syncthreads, + __syncthreads, __syncthreads_count, __syncthreads_and, __syncthreads_or, + __syncwarp, __syncwarp_mask, __threadfence_block, __threadfence_grid, -- Shared memory @@ -44,48 +47,54 @@ ) where --- llvm +import Data.Array.Accelerate.Error +import Data.Array.Accelerate.LLVM.CodeGen.Arithmetic as A +import Data.Array.Accelerate.LLVM.CodeGen.Base +import Data.Array.Accelerate.LLVM.CodeGen.Constant +import Data.Array.Accelerate.LLVM.CodeGen.IR +import Data.Array.Accelerate.LLVM.CodeGen.Module +import Data.Array.Accelerate.LLVM.CodeGen.Monad +import Data.Array.Accelerate.LLVM.CodeGen.Ptr +import Data.Array.Accelerate.LLVM.CodeGen.Sugar +import Data.Array.Accelerate.LLVM.PTX.Analysis.Launch +import Data.Array.Accelerate.LLVM.PTX.Target +import Data.Array.Accelerate.Representation.Array +import Data.Array.Accelerate.Representation.Elt +import Data.Array.Accelerate.Representation.Shape +import Data.Array.Accelerate.Representation.Type + +import Foreign.CUDA.Analysis ( Compute(..), computeCapability ) +import qualified Foreign.CUDA.Analysis as CUDA + import LLVM.AST.Type.AddrSpace import LLVM.AST.Type.Constant -import LLVM.AST.Type.Global +import LLVM.AST.Type.Downcast +import LLVM.AST.Type.Function +import LLVM.AST.Type.InlineAssembly import LLVM.AST.Type.Instruction import LLVM.AST.Type.Instruction.Volatile import LLVM.AST.Type.Metadata import LLVM.AST.Type.Name import LLVM.AST.Type.Operand import LLVM.AST.Type.Representation -import qualified LLVM.AST.Global as LLVM import qualified LLVM.AST.Constant as LLVM hiding ( type' ) +import qualified LLVM.AST.Global as LLVM import qualified LLVM.AST.Linkage as LLVM import qualified LLVM.AST.Name as LLVM import qualified LLVM.AST.Type as LLVM --- accelerate -import Data.Array.Accelerate.Analysis.Type -import Data.Array.Accelerate.Array.Sugar ( Elt, Vector, eltType ) -import Data.Array.Accelerate.Error - -import Data.Array.Accelerate.LLVM.CodeGen.Arithmetic as A -import Data.Array.Accelerate.LLVM.CodeGen.Base -import Data.Array.Accelerate.LLVM.CodeGen.Constant -import Data.Array.Accelerate.LLVM.CodeGen.Downcast -import Data.Array.Accelerate.LLVM.CodeGen.IR -import Data.Array.Accelerate.LLVM.CodeGen.Module -import Data.Array.Accelerate.LLVM.CodeGen.Monad -import Data.Array.Accelerate.LLVM.CodeGen.Ptr -import Data.Array.Accelerate.LLVM.CodeGen.Sugar -import Data.Array.Accelerate.LLVM.CodeGen.Type - -import Data.Array.Accelerate.LLVM.PTX.Analysis.Launch -import Data.Array.Accelerate.LLVM.PTX.Context -import Data.Array.Accelerate.LLVM.PTX.Target - --- standard library import Control.Applicative import Control.Monad ( void ) +import Control.Monad.State ( gets ) +import Prelude as P + +#if MIN_VERSION_llvm_hs(10,0,0) +import qualified LLVM.AST.Type.Instruction.RMW as RMW +import LLVM.AST.Type.Instruction.Atomic +#elif !MIN_VERSION_llvm_hs(9,0,0) import Data.String import Text.Printf -import Prelude as P +#endif -- Thread identifiers @@ -95,21 +104,21 @@ -- -- -- -specialPTXReg :: Label -> CodeGen (IR Int32) +specialPTXReg :: Label -> CodeGen PTX (Operands Int32) specialPTXReg f = - call (Body type' f) [NoUnwind, ReadNone] + call (Body type' (Just Tail) f) [NoUnwind, ReadNone] -blockDim, gridDim, threadIdx, blockIdx, warpSize :: CodeGen (IR Int32) +blockDim, gridDim, threadIdx, blockIdx, warpSize :: CodeGen PTX (Operands Int32) blockDim = specialPTXReg "llvm.nvvm.read.ptx.sreg.ntid.x" gridDim = specialPTXReg "llvm.nvvm.read.ptx.sreg.nctaid.x" threadIdx = specialPTXReg "llvm.nvvm.read.ptx.sreg.tid.x" blockIdx = specialPTXReg "llvm.nvvm.read.ptx.sreg.ctaid.x" warpSize = specialPTXReg "llvm.nvvm.read.ptx.sreg.warpsize" -laneId :: CodeGen (IR Int32) +laneId :: CodeGen PTX (Operands Int32) laneId = specialPTXReg "llvm.nvvm.read.ptx.sreg.laneid" -laneMask_eq, laneMask_lt, laneMask_le, laneMask_gt, laneMask_ge :: CodeGen (IR Int32) +laneMask_eq, laneMask_lt, laneMask_le, laneMask_gt, laneMask_ge :: CodeGen PTX (Operands Int32) laneMask_eq = specialPTXReg "llvm.nvvm.read.ptx.sreg.lanemask.eq" laneMask_lt = specialPTXReg "llvm.nvvm.read.ptx.sreg.lanemask.lt" laneMask_le = specialPTXReg "llvm.nvvm.read.ptx.sreg.lanemask.le" @@ -124,16 +133,13 @@ -- -- http://docs.nvidia.com/cuda/parallel-thread-execution/index.html#special-registers-warpid -- --- We might consider passing in the (constant) warp size from device properties, --- so that the division can be optimised to a shift. --- -warpId :: CodeGen (IR Int32) +warpId :: CodeGen PTX (Operands Int32) warpId = do + dev <- liftCodeGen $ gets ptxDeviceProperties tid <- threadIdx - ws <- warpSize - A.quot integralType tid ws + A.quot integralType tid (A.liftInt32 (P.fromIntegral (CUDA.warpSize dev))) -_warpId :: CodeGen (IR Int32) +_warpId :: CodeGen PTX (Operands Int32) _warpId = specialPTXReg "llvm.ptx.read.warpid" @@ -141,7 +147,7 @@ -- -- > gridDim.x * blockDim.x -- -gridSize :: CodeGen (IR Int32) +gridSize :: CodeGen PTX (Operands Int32) gridSize = do ncta <- gridDim nt <- blockDim @@ -152,7 +158,7 @@ -- -- > blockDim.x * blockIdx.x + threadIdx.x -- -globalThreadIdx :: CodeGen (IR Int32) +globalThreadIdx :: CodeGen PTX (Operands Int32) globalThreadIdx = do ntid <- blockDim ctaid <- blockIdx @@ -163,37 +169,88 @@ return v +{-- -- | Generate function parameters that will specify the first and last (linear) -- index of the array this kernel should evaluate. -- -gangParam :: (IR Int, IR Int, [LLVM.Parameter]) +gangParam :: (Operands Int, Operands Int, [LLVM.Parameter]) gangParam = - let t = scalarType - start = "ix.start" - end = "ix.end" + let start = "ix.start" + end = "ix.end" in - (local t start, local t end, [ scalarParameter t start, scalarParameter t end ] ) + (local start, local end, parameter start ++ parameter end ) +--} -- Barriers and synchronisation -- ---------------------------- --- | Call a builtin CUDA synchronisation intrinsic +-- | Call a built-in CUDA synchronisation intrinsic -- -barrier :: Label -> CodeGen () -barrier f = void $ call (Body VoidType f) [NoUnwind, NoDuplicate, Convergent] +barrier :: Label -> CodeGen PTX () +barrier f = void $ call (Body VoidType (Just Tail) f) [NoUnwind, NoDuplicate, Convergent] +barrier_op :: Label -> Operands Int32 -> CodeGen PTX (Operands Int32) +barrier_op f x = call (Lam primType (op integralType x) (Body type' (Just Tail) f)) [NoUnwind, NoDuplicate, Convergent] --- | Wait until all threads in the thread block have reached this point and all --- global and shared memory accesses made by these threads prior to + +-- | Wait until all threads in the thread block have reached this point, and all +-- global and shared memory accesses made by these threads prior to the -- __syncthreads() are visible to all threads in the block. -- -- -- -__syncthreads :: CodeGen () +__syncthreads :: CodeGen PTX () __syncthreads = barrier "llvm.nvvm.barrier0" +-- | Identical to __syncthreads() with the additional feature that it returns +-- the number of threads in the block for which the predicate evaluates to +-- non-zero. +-- +__syncthreads_count :: Operands Int32 -> CodeGen PTX (Operands Int32) +__syncthreads_count = barrier_op "llvm.nvvm.barrier0.popc" +-- | Identical to __syncthreads() with the additional feature that it returns +-- non-zero iff the predicate evaluates to non-zero for all threads in the +-- block. +-- +__syncthreads_and :: Operands Int32 -> CodeGen PTX (Operands Int32) +__syncthreads_and = barrier_op "llvm.nvvm.barrier0.and" + +-- | Identical to __syncthreads() with the additional feature that it returns +-- non-zero iff the predicate evaluates to non-zero for any thread in the block. +-- +__syncthreads_or :: Operands Int32 -> CodeGen PTX (Operands Int32) +__syncthreads_or = barrier_op "llvm.nvvm.barrier0.or" + + +-- | Wait until all warp lanes have reached this point. +-- +__syncwarp :: HasCallStack => CodeGen PTX () +__syncwarp = __syncwarp_mask (liftWord32 0xffffffff) + +-- | Wait until all warp lanes named in the mask have executed a __syncwarp() +-- with the same mask. All non-exited threads named in the mask must execute +-- a corresponding __syncwarp with the same mask, or the result is undefined. +-- +-- This guarantees memory ordering among threads participating in the barrier. +-- +-- Requires LLVM-6.0 or higher. +-- Only required for devices of SM7 and later. +-- +__syncwarp_mask :: HasCallStack => Operands Word32 -> CodeGen PTX () +__syncwarp_mask mask = do + dev <- liftCodeGen $ gets ptxDeviceProperties + if computeCapability dev < Compute 7 0 + then return () + else +#if !MIN_VERSION_llvm_hs(6,0,0) + internalError "LLVM-6.0 or above is required for Volta devices and later" +#else + void $ call (Lam primType (op primType mask) (Body VoidType (Just Tail) "llvm.nvvm.bar.warp.sync")) [NoUnwind, NoDuplicate, Convergent] +#endif + + -- | Ensure that all writes to shared and global memory before the call to -- __threadfence_block() are observed by all threads in the *block* of the -- calling thread as occurring before all writes to shared and global memory @@ -201,14 +258,14 @@ -- -- -- -__threadfence_block :: CodeGen () +__threadfence_block :: CodeGen PTX () __threadfence_block = barrier "llvm.nvvm.membar.cta" -- | As __threadfence_block(), but the synchronisation is for *all* thread blocks. -- In CUDA this is known simply as __threadfence(). -- -__threadfence_grid :: CodeGen () +__threadfence_grid :: CodeGen PTX () __threadfence_grid = barrier "llvm.nvvm.membar.gl" @@ -219,35 +276,53 @@ -- additional support for atomic add on floating point types, which can be -- accessed through the following intrinsics. -- --- Double precision is only supported on Compute 6.0 devices and later. LLVM-4.0 --- currently lacks support for this intrinsic, however it may be possible to use --- inline assembly. +-- Double precision is supported on Compute 6.0 devices and later. Half +-- precision is supported on Compute 7.0 devices and later. -- +-- LLVM-4.0 currently lacks support for this intrinsic, however it is +-- accessible via inline assembly. +-- +-- LLVM-9 integrated floating-point atomic operations into the AtomicRMW +-- instruction, but this functionality is missing from llvm-hs-9. We access +-- it via inline assembly.. +-- -- -- -atomicAdd_f :: FloatingType a -> Operand (Ptr a) -> Operand a -> CodeGen () +atomicAdd_f :: HasCallStack => FloatingType a -> Operand (Ptr a) -> Operand a -> CodeGen PTX () atomicAdd_f t addr val = +#if MIN_VERSION_llvm_hs(10,0,0) + void . instr' $ AtomicRMW (FloatingNumType t) NonVolatile RMW.FAdd addr val (CrossThread, AcquireRelease) +#else let - width :: Int - width = + _width :: Int + _width = case t of - TypeHalf{} -> 16 - TypeFloat{} -> 32 - TypeDouble{} -> 64 - TypeCFloat{} -> 32 - TypeCDouble{} -> 64 + TypeHalf -> 16 + TypeFloat -> 32 + TypeDouble -> 64 - addrspace :: Word32 - (t_addr, t_val, addrspace) = + (t_addr, t_val, _addrspace) = case typeOf addr of PrimType ta@(PtrPrimType (ScalarPrimType tv) (AddrSpace as)) -> (ta, tv, as) - _ -> $internalError "atomicAdd" "unexpected operand type" + _ -> internalError "unexpected operand type" t_ret = PrimType (ScalarPrimType t_val) - fun = fromString $ printf "llvm.nvvm.atomic.load.add.f%d.p%df%d" width addrspace width +#if MIN_VERSION_llvm_hs(9,0,0) || !MIN_VERSION_llvm_hs(6,0,0) + asm = + case t of + -- assuming .address_size 64 + TypeHalf -> InlineAssembly "atom.add.noftz.f16 $0, [$1], $2;" "=c,l,c" True False ATTDialect + TypeFloat -> InlineAssembly "atom.global.add.f32 $0, [$1], $2;" "=f,l,f" True False ATTDialect + TypeDouble -> InlineAssembly "atom.global.add.f64 $0, [$1], $2;" "=d,l,d" True False ATTDialect in - void $ call (Lam t_addr addr (Lam (ScalarPrimType t_val) val (Body t_ret fun))) [NoUnwind] + void $ instr (Call (Lam t_addr addr (Lam (ScalarPrimType t_val) val (Body t_ret (Just Tail) (Left asm)))) [Right NoUnwind]) +#else + fun = fromString $ printf "llvm.nvvm.atomic.load.add.f%d.p%df%d" _width (_addrspace :: Word32) _width + in + void $ call (Lam t_addr addr (Lam (ScalarPrimType t_val) val (Body t_ret (Just Tail) fun))) [NoUnwind] +#endif +#endif -- Shared memory @@ -264,21 +339,22 @@ -- space, with enough storage to contain the given number of elements. -- staticSharedMem - :: forall e. Elt e - => Word64 - -> CodeGen (IRArray (Vector e)) -staticSharedMem n = do - ad <- go (eltType (undefined::e)) - return $ IRArray { irArrayShape = IR (OP_Pair OP_Unit (OP_Int (integral integralType (P.fromIntegral n)))) - , irArrayData = IR ad + :: TypeR e + -> Word64 + -> CodeGen PTX (IRArray (Vector e)) +staticSharedMem tp n = do + ad <- go tp + return $ IRArray { irArrayRepr = ArrayR dim1 tp + , irArrayShape = OP_Pair OP_Unit $ OP_Int $ integral integralType $ P.fromIntegral n + , irArrayData = ad , irArrayAddrSpace = sharedMemAddrSpace , irArrayVolatility = sharedMemVolatility } where - go :: TupleType s -> CodeGen (Operands s) - go TypeRunit = return OP_Unit - go (TypeRpair t1 t2) = OP_Pair <$> go t1 <*> go t2 - go tt@(TypeRscalar t) = do + go :: TypeR s -> CodeGen PTX (Operands s) + go TupRunit = return OP_Unit + go (TupRpair t1 t2) = OP_Pair <$> go t1 <*> go t2 + go tt@(TupRsingle t) = do -- Declare a new global reference for the statically allocated array -- located in the __shared__ memory space. nm <- freshName @@ -288,7 +364,7 @@ , LLVM.type' = LLVM.ArrayType n (downcast t) , LLVM.linkage = LLVM.External , LLVM.name = downcast nm - , LLVM.alignment = 4 `P.max` P.fromIntegral (sizeOf tt) + , LLVM.alignment = 4 `P.max` P.fromIntegral (bytesElt tt) } -- Return a pointer to the first element of the __shared__ memory array. @@ -297,7 +373,7 @@ p <- instr' $ GetElementPtr sm [num numType 0, num numType 0 :: Operand Int32] q <- instr' $ PtrCast (PtrPrimType (ScalarPrimType t) sharedMemAddrSpace) p - return $ ir' t (unPtr q) + return $ ir t (unPtr q) -- External declaration in shared memory address space. This must be declared in @@ -306,7 +382,7 @@ -- -- > @__shared__ = external addrspace(3) global [0 x i8] -- -initialiseDynamicSharedMemory :: CodeGen (Operand (Ptr Word8)) +initialiseDynamicSharedMemory :: CodeGen PTX (Operand (Ptr Word8)) initialiseDynamicSharedMemory = do declare $ LLVM.globalVariableDefaults { LLVM.addrSpace = sharedMemAddrSpace @@ -322,33 +398,54 @@ -- with enough space to contain the given number of elements. -- dynamicSharedMem - :: forall e int. (Elt e, IsIntegral int) - => IR int -- number of array elements - -> IR int -- #bytes of shared memory the have already been allocated - -> CodeGen (IRArray (Vector e)) -dynamicSharedMem n@(op integralType -> m) (op integralType -> offset) = do - smem <- initialiseDynamicSharedMemory + :: forall e int. + TypeR e + -> IntegralType int + -> Operands int -- number of array elements + -> Operands int -- #bytes of shared memory the have already been allocated + -> CodeGen PTX (IRArray (Vector e)) +dynamicSharedMem tp int n@(op int -> m) (op int -> offset) + | IntegralDict <- integralDict int = do + smem <- initialiseDynamicSharedMemory + let + numTp = IntegralNumType int + + go :: TypeR s -> Operand int -> CodeGen PTX (Operand int, Operands s) + go TupRunit i = return (i, OP_Unit) + go (TupRpair t2 t1) i0 = do + (i1, p1) <- go t1 i0 + (i2, p2) <- go t2 i1 + return $ (i2, OP_Pair p2 p1) + go (TupRsingle t) i = do + p <- instr' $ GetElementPtr smem [num numTp 0, i] -- TLM: note initial zero index!! + q <- instr' $ PtrCast (PtrPrimType (ScalarPrimType t) sharedMemAddrSpace) p + a <- instr' $ Mul numTp m (integral int (P.fromIntegral (bytesElt (TupRsingle t)))) + b <- instr' $ Add numTp i a + return (b, ir t (unPtr q)) + -- + (_, ad) <- go tp offset + sz <- A.fromIntegral int (numType :: NumType Int) n + return $ IRArray { irArrayRepr = ArrayR dim1 tp + , irArrayShape = OP_Pair OP_Unit sz + , irArrayData = ad + , irArrayAddrSpace = sharedMemAddrSpace + , irArrayVolatility = sharedMemVolatility + } + + +-- Other functions +-- --------------- + +-- Sleep the thread for (approximately) the given number of nanoseconds. +-- Requires compute capability >= 7.0 +-- +nanosleep :: Operands Int32 -> CodeGen PTX () +nanosleep ns = let - go :: TupleType s -> Operand int -> CodeGen (Operand int, Operands s) - go TypeRunit i = return (i, OP_Unit) - go (TypeRpair t2 t1) i0 = do - (i1, p1) <- go t1 i0 - (i2, p2) <- go t2 i1 - return $ (i2, OP_Pair p2 p1) - go (TypeRscalar t) i = do - p <- instr' $ GetElementPtr smem [num numType 0, i] -- TLM: note initial zero index!! - q <- instr' $ PtrCast (PtrPrimType (ScalarPrimType t) sharedMemAddrSpace) p - a <- instr' $ Mul numType m (integral integralType (P.fromIntegral (sizeOf (TypeRscalar t)))) - b <- instr' $ Add numType i a - return (b, ir' t (unPtr q)) - -- - (_, ad) <- go (eltType (undefined::e)) offset - IR sz <- A.fromIntegral integralType (numType :: NumType Int) n - return $ IRArray { irArrayShape = IR $ OP_Pair OP_Unit sz - , irArrayData = IR ad - , irArrayAddrSpace = sharedMemAddrSpace - , irArrayVolatility = sharedMemVolatility - } + attrs = [NoUnwind, Convergent] + asm = InlineAssembly "nanosleep.u32 $0;" "r" True False ATTDialect + in + void $ instr (Call (Lam primType (op integralType ns) (Body VoidType (Just Tail) (Left asm))) (map Right attrs)) -- Global kernel definitions @@ -365,13 +462,13 @@ -- | Create a single kernel program with the default launch configuration. -- makeOpenAcc - :: PTX - -> Label + :: Label -> [LLVM.Parameter] - -> CodeGen () - -> CodeGen (IROpenAcc PTX aenv a) -makeOpenAcc (deviceProperties . ptxContext -> dev) = - makeOpenAccWith (simpleLaunchConfig dev) + -> CodeGen PTX () + -> CodeGen PTX (IROpenAcc PTX aenv a) +makeOpenAcc name param kernel = do + dev <- liftCodeGen $ gets ptxDeviceProperties + makeOpenAccWith (simpleLaunchConfig dev) name param kernel -- | Create a single kernel program with the given launch analysis information. -- @@ -379,8 +476,8 @@ :: LaunchConfig -> Label -> [LLVM.Parameter] - -> CodeGen () - -> CodeGen (IROpenAcc PTX aenv a) + -> CodeGen PTX () + -> CodeGen PTX (IROpenAcc PTX aenv a) makeOpenAccWith config name param kernel = do body <- makeKernel config name param kernel return $ IROpenAcc [body] @@ -388,7 +485,12 @@ -- | Create a complete kernel function by running the code generation process -- specified in the final parameter. -- -makeKernel :: LaunchConfig -> Label -> [LLVM.Parameter] -> CodeGen () -> CodeGen (Kernel PTX aenv a) +makeKernel + :: LaunchConfig + -> Label + -> [LLVM.Parameter] + -> CodeGen PTX () + -> CodeGen PTX (Kernel PTX aenv a) makeKernel config name@(Label l) param kernel = do _ <- kernel code <- createBlocks diff --git a/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Fold.hs b/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Fold.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Fold.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Fold.hs @@ -4,14 +4,15 @@ {-# LANGUAGE RecordWildCards #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeOperators #-} {-# LANGUAGE ViewPatterns #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.CodeGen.Fold --- Copyright : [2016..2017] Trevor L. McDonell +-- Copyright : [2016..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- @@ -19,16 +20,15 @@ module Data.Array.Accelerate.LLVM.PTX.CodeGen.Fold where --- accelerate -import Data.Array.Accelerate.Analysis.Match -import Data.Array.Accelerate.Analysis.Type -import Data.Array.Accelerate.Array.Sugar ( Array, Scalar, Vector, Shape, Z, (:.), Elt(..) ) +import Data.Array.Accelerate.Representation.Array +import Data.Array.Accelerate.Representation.Elt +import Data.Array.Accelerate.Representation.Shape hiding ( size ) +import Data.Array.Accelerate.Representation.Type --- accelerate-llvm-* -import Data.Array.Accelerate.LLVM.Analysis.Match import Data.Array.Accelerate.LLVM.CodeGen.Arithmetic as A import Data.Array.Accelerate.LLVM.CodeGen.Array import Data.Array.Accelerate.LLVM.CodeGen.Base +import Data.Array.Accelerate.LLVM.CodeGen.Constant import Data.Array.Accelerate.LLVM.CodeGen.Environment import Data.Array.Accelerate.LLVM.CodeGen.Exp import Data.Array.Accelerate.LLVM.CodeGen.IR @@ -39,16 +39,14 @@ import Data.Array.Accelerate.LLVM.PTX.Analysis.Launch import Data.Array.Accelerate.LLVM.PTX.CodeGen.Base import Data.Array.Accelerate.LLVM.PTX.CodeGen.Generate -import Data.Array.Accelerate.LLVM.PTX.Context import Data.Array.Accelerate.LLVM.PTX.Target import LLVM.AST.Type.Representation --- cuda import qualified Foreign.CUDA.Analysis as CUDA -import Control.Applicative ( (<$>), (<*>) ) import Control.Monad ( (>=>) ) +import Control.Monad.State ( gets ) import Data.String ( fromString ) import Data.Bits as P import Prelude as P @@ -62,43 +60,20 @@ -- a neutral element {(+), 0} -- mkFold - :: forall aenv sh e. (Shape sh, Elt e) - => PTX - -> Gamma aenv - -> IRFun2 PTX aenv (e -> e -> e) - -> IRExp PTX aenv e - -> IRDelayed PTX aenv (Array (sh :. Int) e) - -> CodeGen (IROpenAcc PTX aenv (Array sh e)) -mkFold ptx@(deviceProperties . ptxContext -> dev) aenv f z acc - | Just Refl <- matchShapeType (undefined::sh) (undefined::Z) - = (+++) <$> mkFoldAll dev aenv f (Just z) acc - <*> mkFoldFill ptx aenv z - - | otherwise - = (+++) <$> mkFoldDim dev aenv f (Just z) acc - <*> mkFoldFill ptx aenv z - - --- Reduce a non-empty array along the innermost dimension. The reduction --- function must be associative to allow for an efficient parallel --- implementation. --- --- TODO: Specialise for commutative operations (such as (+)) and those with --- a neutral element {(+), 0} --- -mkFold1 - :: forall aenv sh e. (Shape sh, Elt e) - => PTX - -> Gamma aenv - -> IRFun2 PTX aenv (e -> e -> e) - -> IRDelayed PTX aenv (Array (sh :. Int) e) - -> CodeGen (IROpenAcc PTX aenv (Array sh e)) -mkFold1 (deviceProperties . ptxContext -> dev) aenv f acc - | Just Refl <- matchShapeType (undefined::sh) (undefined::Z) - = mkFoldAll dev aenv f Nothing acc - - | otherwise - = mkFoldDim dev aenv f Nothing acc + :: forall aenv sh e. + Gamma aenv + -> ArrayR (Array sh e) + -> IRFun2 PTX aenv (e -> e -> e) + -> Maybe (IRExp PTX aenv e) + -> MIRDelayed PTX aenv (Array (sh, Int) e) + -> CodeGen PTX (IROpenAcc PTX aenv (Array sh e)) +mkFold aenv repr f z acc = case z of + Just z' -> (+++) <$> codeFold <*> mkFoldFill aenv repr z' + Nothing -> codeFold + where + codeFold = case repr of + ArrayR ShapeRz tp -> mkFoldAll aenv tp f z acc + _ -> mkFoldDim aenv repr f z acc -- Reduce an array to a single element. @@ -120,17 +95,18 @@ -- element per thread) to a single element, which is stored to the output array. -- mkFoldAll - :: forall aenv e. Elt e - => DeviceProperties -- ^ properties of the target GPU - -> Gamma aenv -- ^ array environment - -> IRFun2 PTX aenv (e -> e -> e) -- ^ combination function - -> Maybe (IRExp PTX aenv e) -- ^ seed element, if this is an exclusive reduction - -> IRDelayed PTX aenv (Vector e) -- ^ input data - -> CodeGen (IROpenAcc PTX aenv (Scalar e)) -mkFoldAll dev aenv combine mseed acc = - foldr1 (+++) <$> sequence [ mkFoldAllS dev aenv combine mseed acc - , mkFoldAllM1 dev aenv combine acc - , mkFoldAllM2 dev aenv combine mseed + :: forall aenv e. + Gamma aenv -- ^ array environment + -> TypeR e + -> IRFun2 PTX aenv (e -> e -> e) -- ^ combination function + -> MIRExp PTX aenv e -- ^ (optional) initial element for exclusive reductions + -> MIRDelayed PTX aenv (Vector e) -- ^ input data + -> CodeGen PTX (IROpenAcc PTX aenv (Scalar e)) +mkFoldAll aenv tp combine mseed macc = do + dev <- liftCodeGen $ gets ptxDeviceProperties + foldr1 (+++) <$> sequence [ mkFoldAllS dev aenv tp combine mseed macc + , mkFoldAllM1 dev aenv tp combine macc + , mkFoldAllM2 dev aenv tp combine mseed ] @@ -138,34 +114,38 @@ -- processed by a single thread block. -- mkFoldAllS - :: forall aenv e. Elt e - => DeviceProperties -- ^ properties of the target GPU - -> Gamma aenv -- ^ array environment - -> IRFun2 PTX aenv (e -> e -> e) -- ^ combination function - -> Maybe (IRExp PTX aenv e) - -> IRDelayed PTX aenv (Vector e) -- ^ input data - -> CodeGen (IROpenAcc PTX aenv (Scalar e)) -mkFoldAllS dev aenv combine mseed IRDelayed{..} = + :: forall aenv e. + DeviceProperties -- ^ properties of the target GPU + -> Gamma aenv -- ^ array environment + -> TypeR e + -> IRFun2 PTX aenv (e -> e -> e) -- ^ combination function + -> MIRExp PTX aenv e -- ^ (optional) initial element for exclusive reductions + -> MIRDelayed PTX aenv (Vector e) -- ^ input data + -> CodeGen PTX (IROpenAcc PTX aenv (Scalar e)) +mkFoldAllS dev aenv tp combine mseed marr = let - (start, end, paramGang) = gangParam - (arrOut, paramOut) = mutableArray ("out" :: Name (Scalar e)) - paramEnv = envParam aenv + (arrOut, paramOut) = mutableArray (ArrayR dim0 tp) "out" + (arrIn, paramIn) = delayedArray "in" marr + paramEnv = envParam aenv -- - config = launchConfig dev (CUDA.incWarp dev) smem multipleOf multipleOfQ - smem n = warps * (1 + per_warp) * bytes + config = launchConfig dev (CUDA.incWarp dev) smem multipleOf multipleOfQ + smem n = warps * (1 + per_warp) * bytes where ws = CUDA.warpSize dev warps = n `P.quot` ws per_warp = ws + ws `P.quot` 2 - bytes = sizeOf (eltType (undefined :: e)) + bytes = bytesElt tp in - makeOpenAccWith config "foldAllS" (paramGang ++ paramOut ++ paramEnv) $ do + makeOpenAccWith config "foldAllS" (paramOut ++ paramIn ++ paramEnv) $ do tid <- threadIdx bd <- blockDim + sh <- delayedExtent arrIn + end <- shapeSize dim1 sh + -- We can assume that there is only a single thread block - start' <- i32 start + start' <- return (liftInt32 0) end' <- i32 end i0 <- A.add numType start' tid sz <- A.sub numType end' start' @@ -173,13 +153,13 @@ -- Thread reads initial element and then participates in block-wide -- reduction. - x0 <- app1 delayedLinearIndex =<< int i0 - r0 <- if A.eq singleType sz bd - then reduceBlockSMem dev combine Nothing x0 - else reduceBlockSMem dev combine (Just sz) x0 + x0 <- app1 (delayedLinearIndex arrIn) =<< int i0 + r0 <- if (tp, A.eq singleType sz bd) + then reduceBlockSMem dev tp combine Nothing x0 + else reduceBlockSMem dev tp combine (Just sz) x0 - when (A.eq singleType tid (lift 0)) $ - writeArray arrOut tid =<< + when (A.eq singleType tid (liftInt32 0)) $ + writeArray TypeInt32 arrOut tid =<< case mseed of Nothing -> return r0 Just z -> flip (app2 combine) r0 =<< z -- Note: initial element on the left @@ -192,27 +172,29 @@ -- blocks. -- mkFoldAllM1 - :: forall aenv e. Elt e - => DeviceProperties -- ^ properties of the target GPU - -> Gamma aenv -- ^ array environment - -> IRFun2 PTX aenv (e -> e -> e) -- ^ combination function - -> IRDelayed PTX aenv (Vector e) -- ^ input data - -> CodeGen (IROpenAcc PTX aenv (Scalar e)) -mkFoldAllM1 dev aenv combine IRDelayed{..} = + :: forall aenv e. + DeviceProperties -- ^ properties of the target GPU + -> Gamma aenv -- ^ array environment + -> TypeR e + -> IRFun2 PTX aenv (e -> e -> e) -- ^ combination function + -> MIRDelayed PTX aenv (Vector e) -- ^ input data + -> CodeGen PTX (IROpenAcc PTX aenv (Scalar e)) +mkFoldAllM1 dev aenv tp combine marr = let - (start, end, paramGang) = gangParam - (arrTmp, paramTmp) = mutableArray ("tmp" :: Name (Vector e)) - paramEnv = envParam aenv + (arrTmp, paramTmp) = mutableArray (ArrayR dim1 tp) "tmp" + (arrIn, paramIn) = delayedArray "in" marr + paramEnv = envParam aenv + start = liftInt 0 -- - config = launchConfig dev (CUDA.incWarp dev) smem const [|| const ||] - smem n = warps * (1 + per_warp) * bytes + config = launchConfig dev (CUDA.incWarp dev) smem const [|| const ||] + smem n = warps * (1 + per_warp) * bytes where ws = CUDA.warpSize dev warps = n `P.quot` ws per_warp = ws + ws `P.quot` 2 - bytes = sizeOf (eltType (undefined :: e)) + bytes = bytesElt tp in - makeOpenAccWith config "foldAllM1" (paramGang ++ paramTmp ++ paramEnv) $ do + makeOpenAccWith config "foldAllM1" (paramTmp ++ paramIn ++ paramEnv) $ do -- Each thread block cooperatively reduces a stripe of the input and stores -- that value into a temporary array at a corresponding index. Since the @@ -221,7 +203,8 @@ -- tid <- threadIdx bd <- int =<< blockDim - sz <- indexHead <$> delayedExtent + sz <- indexHead <$> delayedExtent arrIn + end <- shapeSize dim1 (irArrayShape arrTmp) imapFromTo start end $ \seg -> do @@ -234,9 +217,9 @@ to <- A.min singleType sz step -- Threads cooperatively reduce this stripe - reduceFromTo dev from to combine - (app1 delayedLinearIndex) - (when (A.eq singleType tid (lift 0)) . writeArray arrTmp seg) + reduceFromTo dev tp from to combine + (app1 (delayedLinearIndex arrIn)) + (when (A.eq singleType tid (liftInt32 0)) . writeArray TypeInt arrTmp seg) return_ @@ -245,28 +228,29 @@ -- reduction algorithm. -- mkFoldAllM2 - :: forall aenv e. Elt e - => DeviceProperties - -> Gamma aenv - -> IRFun2 PTX aenv (e -> e -> e) - -> Maybe (IRExp PTX aenv e) - -> CodeGen (IROpenAcc PTX aenv (Scalar e)) -mkFoldAllM2 dev aenv combine mseed = + :: forall aenv e. + DeviceProperties + -> Gamma aenv + -> TypeR e + -> IRFun2 PTX aenv (e -> e -> e) + -> MIRExp PTX aenv e + -> CodeGen PTX (IROpenAcc PTX aenv (Scalar e)) +mkFoldAllM2 dev aenv tp combine mseed = let - (start, end, paramGang) = gangParam - (arrTmp, paramTmp) = mutableArray ("tmp" :: Name (Vector e)) - (arrOut, paramOut) = mutableArray ("out" :: Name (Vector e)) - paramEnv = envParam aenv + (arrTmp, paramTmp) = mutableArray (ArrayR dim1 tp) "tmp" + (arrOut, paramOut) = mutableArray (ArrayR dim1 tp) "out" + paramEnv = envParam aenv + start = liftInt 0 -- - config = launchConfig dev (CUDA.incWarp dev) smem const [|| const ||] - smem n = warps * (1 + per_warp) * bytes + config = launchConfig dev (CUDA.incWarp dev) smem const [|| const ||] + smem n = warps * (1 + per_warp) * bytes where ws = CUDA.warpSize dev warps = n `P.quot` ws per_warp = ws + ws `P.quot` 2 - bytes = sizeOf (eltType (undefined :: e)) + bytes = bytesElt tp in - makeOpenAccWith config "foldAllM2" (paramGang ++ paramTmp ++ paramOut ++ paramEnv) $ do + makeOpenAccWith config "foldAllM2" (paramTmp ++ paramOut ++ paramEnv) $ do -- Threads cooperatively reduce a stripe of the input (temporary) array -- output from the first phase, storing the results into another temporary. @@ -277,6 +261,7 @@ gd <- gridDim bd <- int =<< blockDim sz <- return $ indexHead (irArrayShape arrTmp) + end <- shapeSize dim1 (irArrayShape arrOut) imapFromTo start end $ \seg -> do @@ -289,12 +274,12 @@ to <- A.min singleType sz step -- Threads cooperatively reduce this stripe - reduceFromTo dev from to combine (readArray arrTmp) $ \r -> - when (A.eq singleType tid (lift 0)) $ - writeArray arrOut seg =<< + reduceFromTo dev tp from to combine (readArray TypeInt arrTmp) $ \r -> + when (A.eq singleType tid (liftInt32 0)) $ + writeArray TypeInt arrOut seg =<< case mseed of Nothing -> return r - Just z -> if A.eq singleType gd (lift 1) + Just z -> if (tp, A.eq singleType gd (liftInt32 1)) then flip (app2 combine) r =<< z -- Note: initial element on the left else return r @@ -310,37 +295,42 @@ -- is known a priori). -- mkFoldDim - :: forall aenv sh e. (Shape sh, Elt e) - => DeviceProperties -- ^ properties of the target GPU - -> Gamma aenv -- ^ array environment - -> IRFun2 PTX aenv (e -> e -> e) -- ^ combination function - -> Maybe (IRExp PTX aenv e) -- ^ seed element, if this is an exclusive reduction - -> IRDelayed PTX aenv (Array (sh :. Int) e) -- ^ input data - -> CodeGen (IROpenAcc PTX aenv (Array sh e)) -mkFoldDim dev aenv combine mseed IRDelayed{..} = + :: forall aenv sh e. + Gamma aenv -- ^ array environment + -> ArrayR (Array sh e) + -> IRFun2 PTX aenv (e -> e -> e) -- ^ combination function + -> MIRExp PTX aenv e -- ^ (optional) seed element, if this is an exclusive reduction + -> MIRDelayed PTX aenv (Array (sh, Int) e) -- ^ input data + -> CodeGen PTX (IROpenAcc PTX aenv (Array sh e)) +mkFoldDim aenv repr@(ArrayR shr tp) combine mseed marr = do + dev <- liftCodeGen $ gets ptxDeviceProperties + -- let - (start, end, paramGang) = gangParam - (arrOut, paramOut) = mutableArray ("out" :: Name (Array sh e)) - paramEnv = envParam aenv + (arrOut, paramOut) = mutableArray repr "out" + (arrIn, paramIn) = delayedArray "in" marr + paramEnv = envParam aenv -- - config = launchConfig dev (CUDA.incWarp dev) smem const [|| const ||] - smem n = warps * (1 + per_warp) * bytes + config = launchConfig dev (CUDA.incWarp dev) smem const [|| const ||] + smem n = warps * (1 + per_warp) * bytes where ws = CUDA.warpSize dev warps = n `P.quot` ws per_warp = ws + ws `P.quot` 2 - bytes = sizeOf (eltType (undefined :: e)) - in - makeOpenAccWith config "fold" (paramGang ++ paramOut ++ paramEnv) $ do + bytes = bytesElt tp + -- + makeOpenAccWith config "fold" (paramOut ++ paramIn ++ paramEnv) $ do -- If the innermost dimension is smaller than the number of threads in the -- block, those threads will never contribute to the output. tid <- threadIdx - sz <- indexHead <$> delayedExtent + sz <- indexHead <$> delayedExtent arrIn sz' <- i32 sz when (A.lt singleType tid sz') $ do + start <- return (liftInt 0) + end <- shapeSize shr (irArrayShape arrOut) + -- Thread blocks iterate over the outer dimensions, each thread block -- cooperatively reducing along each outermost index to a single value. -- @@ -356,16 +346,16 @@ to <- A.add numType from sz -- last linear index this block will reduce (exclusive) i0 <- A.add numType from =<< int tid - x0 <- app1 delayedLinearIndex i0 + x0 <- app1 (delayedLinearIndex arrIn) i0 bd <- blockDim - r0 <- if A.gte singleType sz' bd - then reduceBlockSMem dev combine Nothing x0 - else reduceBlockSMem dev combine (Just sz') x0 + r0 <- if (tp, A.gte singleType sz' bd) + then reduceBlockSMem dev tp combine Nothing x0 + else reduceBlockSMem dev tp combine (Just sz') x0 -- Step 2: keep walking over the input bd' <- int bd next <- A.add numType from bd' - r <- iterFromStepTo next bd' to r0 $ \offset r -> do + r <- iterFromStepTo tp next bd' to r0 $ \offset r -> do -- Wait for all threads to catch up before starting the next stripe __syncthreads @@ -373,12 +363,12 @@ -- Threads cooperatively reduce this stripe of the input i <- A.add numType offset =<< int tid v' <- A.sub numType to offset - r' <- if A.gte singleType v' bd' + r' <- if (tp, A.gte singleType v' bd') -- All threads of the block are valid, so we can avoid -- bounds checks. then do - x <- app1 delayedLinearIndex i - y <- reduceBlockSMem dev combine Nothing x + x <- app1 (delayedLinearIndex arrIn) i + y <- reduceBlockSMem dev tp combine Nothing x return y -- Otherwise, we require bounds checks when reading the input @@ -387,22 +377,29 @@ -- reduction, we must still have all threads enter the -- reduction procedure to avoid synchronisation divergence. else do - x <- if A.lt singleType i to - then app1 delayedLinearIndex i - else return r + x <- if (tp, A.lt singleType i to) + then app1 (delayedLinearIndex arrIn) i + else let + go :: TypeR a -> Operands a + go TupRunit = OP_Unit + go (TupRpair a b) = OP_Pair (go a) (go b) + go (TupRsingle t) = ir t (undef t) + in + return $ go tp + v <- i32 v' - y <- reduceBlockSMem dev combine (Just v) x + y <- reduceBlockSMem dev tp combine (Just v) x return y - if A.eq singleType tid (lift 0) + if (tp, A.eq singleType tid (liftInt32 0)) then app2 combine r r' else return r' -- Step 3: Thread 0 writes the aggregate reduction of this dimension to -- memory. If this is an exclusive fold, combine with the initial element. -- - when (A.eq singleType tid (lift 0)) $ - writeArray arrOut seg =<< + when (A.eq singleType tid (liftInt32 0)) $ + writeArray TypeInt arrOut seg =<< case mseed of Nothing -> return r Just z -> flip (app2 combine) r =<< z -- Note: initial element on the left @@ -414,13 +411,12 @@ -- dimensions with the initial element. -- mkFoldFill - :: (Shape sh, Elt e) - => PTX - -> Gamma aenv - -> IRExp PTX aenv e - -> CodeGen (IROpenAcc PTX aenv (Array sh e)) -mkFoldFill ptx aenv seed = - mkGenerate ptx aenv (IRFun1 (const seed)) + :: Gamma aenv + -> ArrayR (Array sh e) + -> IRExp PTX aenv e + -> CodeGen PTX (IROpenAcc PTX aenv (Array sh e)) +mkFoldFill aenv repr seed = + mkGenerate aenv repr (IRFun1 (const seed)) -- Efficient threadblock-wide reduction using the specified operator. The @@ -432,61 +428,62 @@ -- Example: https://github.com/NVlabs/cub/blob/1.5.2/cub/block/specializations/block_reduce_warp_reductions.cuh -- reduceBlockSMem - :: forall aenv e. Elt e - => DeviceProperties -- ^ properties of the target device - -> IRFun2 PTX aenv (e -> e -> e) -- ^ combination function - -> Maybe (IR Int32) -- ^ number of valid elements (may be less than block size) - -> IR e -- ^ calling thread's input element - -> CodeGen (IR e) -- ^ thread-block-wide reduction using the specified operator (lane 0 only) -reduceBlockSMem dev combine size = warpReduce >=> warpAggregate + :: forall aenv e. + DeviceProperties -- ^ properties of the target device + -> TypeR e + -> IRFun2 PTX aenv (e -> e -> e) -- ^ combination function + -> Maybe (Operands Int32) -- ^ number of valid elements (may be less than block size) + -> Operands e -- ^ calling thread's input element + -> CodeGen PTX (Operands e) -- ^ thread-block-wide reduction using the specified operator (lane 0 only) +reduceBlockSMem dev tp combine size = warpReduce >=> warpAggregate where - int32 :: Integral a => a -> IR Int32 - int32 = lift . P.fromIntegral + int32 :: Integral a => a -> Operands Int32 + int32 = liftInt32 . P.fromIntegral -- Temporary storage required for each warp - bytes = sizeOf (eltType (undefined::e)) + bytes = bytesElt tp warp_smem_elems = CUDA.warpSize dev + (CUDA.warpSize dev `P.quot` 2) -- Step 1: Reduction in every warp -- - warpReduce :: IR e -> CodeGen (IR e) + warpReduce :: Operands e -> CodeGen PTX (Operands e) warpReduce input = do -- Allocate (1.5 * warpSize) elements of shared memory for each warp wid <- warpId skip <- A.mul numType wid (int32 (warp_smem_elems * bytes)) - smem <- dynamicSharedMem (int32 warp_smem_elems) skip + smem <- dynamicSharedMem tp TypeInt32 (int32 warp_smem_elems) skip -- Are we doing bounds checking for this warp? -- case size of -- The entire thread block is valid, so skip bounds checks. Nothing -> - reduceWarpSMem dev combine smem Nothing input + reduceWarpSMem dev tp combine smem Nothing input -- Otherwise check how many elements are valid for this warp. If it is -- full then we can still skip bounds checks for it. Just n -> do offset <- A.mul numType wid (int32 (CUDA.warpSize dev)) valid <- A.sub numType n offset - if A.gte singleType valid (int32 (CUDA.warpSize dev)) - then reduceWarpSMem dev combine smem Nothing input - else reduceWarpSMem dev combine smem (Just valid) input + if (tp, A.gte singleType valid (int32 (CUDA.warpSize dev))) + then reduceWarpSMem dev tp combine smem Nothing input + else reduceWarpSMem dev tp combine smem (Just valid) input -- Step 2: Aggregate per-warp reductions -- - warpAggregate :: IR e -> CodeGen (IR e) + warpAggregate :: Operands e -> CodeGen PTX (Operands e) warpAggregate input = do -- Allocate #warps elements of shared memory bd <- blockDim warps <- A.quot integralType bd (int32 (CUDA.warpSize dev)) skip <- A.mul numType warps (int32 (warp_smem_elems * bytes)) - smem <- dynamicSharedMem warps skip + smem <- dynamicSharedMem tp TypeInt32 warps skip -- Share the per-lane aggregates wid <- warpId lane <- laneId - when (A.eq singleType lane (lift 0)) $ do - writeArray smem wid input + when (A.eq singleType lane (liftInt32 0)) $ do + writeArray TypeInt32 smem wid input -- Wait for each warp to finish its local reduction __syncthreads @@ -495,7 +492,7 @@ -- done in CUB), but we could also do this cooperatively (better for -- larger thread blocks?) tid <- threadIdx - if A.eq singleType tid (lift 0) + if (tp, A.eq singleType tid (liftInt32 0)) then do steps <- case size of Nothing -> return warps @@ -503,8 +500,8 @@ a <- A.add numType n (int32 (CUDA.warpSize dev - 1)) b <- A.quot integralType a (int32 (CUDA.warpSize dev)) return b - iterFromStepTo (lift 1) (lift 1) steps input $ \step x -> - app2 combine x =<< readArray smem step + iterFromStepTo tp (liftInt32 1) (liftInt32 1) steps input $ \step x -> + app2 combine x =<< readArray TypeInt32 smem step else return input @@ -519,14 +516,15 @@ -- Example: https://github.com/NVlabs/cub/blob/1.5.2/cub/warp/specializations/warp_reduce_smem.cuh#L128 -- reduceWarpSMem - :: forall aenv e. Elt e - => DeviceProperties -- ^ properties of the target device - -> IRFun2 PTX aenv (e -> e -> e) -- ^ combination function - -> IRArray (Vector e) -- ^ temporary storage array in shared memory (1.5 warp size elements) - -> Maybe (IR Int32) -- ^ number of items that will be reduced by this warp, otherwise all lanes are valid - -> IR e -- ^ calling thread's input element - -> CodeGen (IR e) -- ^ warp-wide reduction using the specified operator (lane 0 only) -reduceWarpSMem dev combine smem size = reduce 0 + :: forall aenv e. + DeviceProperties -- ^ properties of the target device + -> TypeR e + -> IRFun2 PTX aenv (e -> e -> e) -- ^ combination function + -> IRArray (Vector e) -- ^ temporary storage array in shared memory (1.5 warp size elements) + -> Maybe (Operands Int32) -- ^ number of items that will be reduced by this warp, otherwise all lanes are valid + -> Operands e -- ^ calling thread's input element + -> CodeGen PTX (Operands e) -- ^ warp-wide reduction using the specified operator (lane 0 only) +reduceWarpSMem dev tp combine smem size = reduce 0 where log2 :: Double -> Double log2 = P.logBase 2 @@ -538,24 +536,30 @@ -- optimised away. valid i = case size of - Nothing -> return (lift True) + Nothing -> return (liftBool True) Just n -> A.lt singleType i n -- Unfold the reduction as a recursive code generation function. - reduce :: Int -> IR e -> CodeGen (IR e) + reduce :: Int -> Operands e -> CodeGen PTX (Operands e) reduce step x - | step >= steps = return x - | offset <- 1 `P.shiftL` step = do + | step >= steps = return x + | otherwise = do + let offset = liftInt32 (1 `P.shiftL` step) + -- share input through buffer lane <- laneId - writeArray smem lane x + writeArray TypeInt32 smem lane x + __syncwarp + -- update input if in range - i <- A.add numType lane (lift offset) - x' <- if valid i - then app2 combine x =<< readArray smem i + i <- A.add numType lane offset + x' <- if (tp, valid i) + then app2 combine x =<< readArray TypeInt32 smem i else return x + __syncwarp + reduce (step+1) x' @@ -565,8 +569,8 @@ -- -- reduceWarpShfl -- :: IRFun2 PTX aenv (e -> e -> e) -- ^ combination function --- -> IR e -- ^ this thread's input value --- -> CodeGen (IR e) -- ^ final result +-- -> Operands e -- ^ this thread's input value +-- -> CodeGen (Operands e) -- ^ final result -- reduceWarpShfl combine input = -- error "TODO: PTX.reduceWarpShfl" @@ -575,15 +579,15 @@ -- --------------- reduceFromTo - :: Elt a - => DeviceProperties - -> IR Int -- ^ starting index - -> IR Int -- ^ final index (exclusive) + :: DeviceProperties + -> TypeR a + -> Operands Int -- ^ starting index + -> Operands Int -- ^ final index (exclusive) -> (IRFun2 PTX aenv (a -> a -> a)) -- ^ combination function - -> (IR Int -> CodeGen (IR a)) -- ^ function to retrieve element at index - -> (IR a -> CodeGen ()) -- ^ what to do with the value - -> CodeGen () -reduceFromTo dev from to combine get set = do + -> (Operands Int -> CodeGen PTX (Operands a)) -- ^ function to retrieve element at index + -> (Operands a -> CodeGen PTX ()) -- ^ what to do with the value + -> CodeGen PTX () +reduceFromTo dev tp from to combine get set = do tid <- int =<< threadIdx bd <- int =<< blockDim @@ -591,44 +595,43 @@ valid <- A.sub numType to from i <- A.add numType from tid - _ <- if A.gte singleType valid bd + _ <- if (TupRunit, A.gte singleType valid bd) then do -- All threads in the block will participate in the reduction, so -- we can avoid bounds checks x <- get i - r <- reduceBlockSMem dev combine Nothing x + r <- reduceBlockSMem dev tp combine Nothing x set r - return (IR OP_Unit :: IR ()) -- unsightly, but free + return (lift TupRunit ()) else do -- Only in-bounds threads can read their input and participate in -- the reduction when (A.lt singleType i to) $ do x <- get i v <- i32 valid - r <- reduceBlockSMem dev combine (Just v) x + r <- reduceBlockSMem dev tp combine (Just v) x set r - return (IR OP_Unit :: IR ()) + return (lift TupRunit ()) return () - -- Utilities -- --------- -i32 :: IR Int -> CodeGen (IR Int32) +i32 :: Operands Int -> CodeGen PTX (Operands Int32) i32 = A.fromIntegral integralType numType -int :: IR Int32 -> CodeGen (IR Int) +int :: Operands Int32 -> CodeGen PTX (Operands Int) int = A.fromIntegral integralType numType imapFromTo - :: IR Int - -> IR Int - -> (IR Int -> CodeGen ()) - -> CodeGen () + :: Operands Int + -> Operands Int + -> (Operands Int -> CodeGen PTX ()) + -> CodeGen PTX () imapFromTo start end body = do bid <- int =<< blockIdx gd <- int =<< gridDim diff --git a/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/FoldSeg.hs b/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/FoldSeg.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/FoldSeg.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/FoldSeg.hs @@ -3,14 +3,15 @@ {-# LANGUAGE RebindableSyntax #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeOperators #-} {-# LANGUAGE ViewPatterns #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.CodeGen.FoldSeg --- Copyright : [2016..2017] Trevor L. McDonell +-- Copyright : [2016..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- @@ -18,12 +19,10 @@ module Data.Array.Accelerate.LLVM.PTX.CodeGen.FoldSeg where --- accelerate -import Data.Array.Accelerate.Analysis.Type -import Data.Array.Accelerate.Array.Sugar ( Array, Segments, Shape(rank), (:.), Elt(..) ) - --- accelerate-llvm-* -import LLVM.AST.Type.Representation +import Data.Array.Accelerate.Representation.Array +import Data.Array.Accelerate.Representation.Elt +import Data.Array.Accelerate.Representation.Shape +import Data.Array.Accelerate.Representation.Type import Data.Array.Accelerate.LLVM.CodeGen.Arithmetic as A import Data.Array.Accelerate.LLVM.CodeGen.Array @@ -39,15 +38,14 @@ import Data.Array.Accelerate.LLVM.PTX.Analysis.Launch import Data.Array.Accelerate.LLVM.PTX.CodeGen.Base import Data.Array.Accelerate.LLVM.PTX.CodeGen.Fold ( reduceBlockSMem, reduceWarpSMem, imapFromTo ) --- import Data.Array.Accelerate.LLVM.PTX.CodeGen.Queue -import Data.Array.Accelerate.LLVM.PTX.Context import Data.Array.Accelerate.LLVM.PTX.Target --- cuda +import LLVM.AST.Type.Representation + import qualified Foreign.CUDA.Analysis as CUDA -import Control.Applicative ( (<$>), (<*>) ) import Control.Monad ( void ) +import Control.Monad.State ( gets ) import Data.String ( fromString ) import Prelude as P @@ -56,33 +54,18 @@ -- reduction per segment of the source array. -- mkFoldSeg - :: forall aenv sh i e. (Shape sh, IsIntegral i, Elt i, Elt e) - => PTX - -> Gamma aenv - -> IRFun2 PTX aenv (e -> e -> e) - -> IRExp PTX aenv e - -> IRDelayed PTX aenv (Array (sh :. Int) e) - -> IRDelayed PTX aenv (Segments i) - -> CodeGen (IROpenAcc PTX aenv (Array (sh :. Int) e)) -mkFoldSeg (deviceProperties . ptxContext -> dev) aenv combine seed arr seg = - (+++) <$> mkFoldSegP_block dev aenv combine (Just seed) arr seg - <*> mkFoldSegP_warp dev aenv combine (Just seed) arr seg - - --- Segmented reduction along the innermost dimension of an array, where /all/ --- segments are non-empty. --- -mkFold1Seg - :: forall aenv sh i e. (Shape sh, IsIntegral i, Elt i, Elt e) - => PTX - -> Gamma aenv - -> IRFun2 PTX aenv (e -> e -> e) - -> IRDelayed PTX aenv (Array (sh :. Int) e) - -> IRDelayed PTX aenv (Segments i) - -> CodeGen (IROpenAcc PTX aenv (Array (sh :. Int) e)) -mkFold1Seg (deviceProperties . ptxContext -> dev) aenv combine arr seg = - (+++) <$> mkFoldSegP_block dev aenv combine Nothing arr seg - <*> mkFoldSegP_warp dev aenv combine Nothing arr seg + :: forall aenv sh i e. + Gamma aenv + -> ArrayR (Array (sh, Int) e) + -> IntegralType i + -> IRFun2 PTX aenv (e -> e -> e) + -> Maybe (IRExp PTX aenv e) + -> MIRDelayed PTX aenv (Array (sh, Int) e) + -> MIRDelayed PTX aenv (Segments i) + -> CodeGen PTX (IROpenAcc PTX aenv (Array (sh, Int) e)) +mkFoldSeg aenv repr intTp combine seed arr seg = + (+++) <$> mkFoldSegP_block aenv repr intTp combine seed arr seg + <*> mkFoldSegP_warp aenv repr intTp combine seed arr seg -- This implementation assumes that the segments array represents the offset @@ -93,29 +76,33 @@ -- worry about inter-block synchronisation. -- mkFoldSegP_block - :: forall aenv sh i e. (Shape sh, IsIntegral i, Elt i, Elt e) - => DeviceProperties - -> Gamma aenv - -> IRFun2 PTX aenv (e -> e -> e) - -> Maybe (IRExp PTX aenv e) - -> IRDelayed PTX aenv (Array (sh :. Int) e) - -> IRDelayed PTX aenv (Segments i) - -> CodeGen (IROpenAcc PTX aenv (Array (sh :. Int) e)) -mkFoldSegP_block dev aenv combine mseed arr seg = + :: forall aenv sh i e. + Gamma aenv + -> ArrayR (Array (sh, Int) e) + -> IntegralType i + -> IRFun2 PTX aenv (e -> e -> e) + -> MIRExp PTX aenv e + -> MIRDelayed PTX aenv (Array (sh, Int) e) + -> MIRDelayed PTX aenv (Segments i) + -> CodeGen PTX (IROpenAcc PTX aenv (Array (sh, Int) e)) +mkFoldSegP_block aenv repr@(ArrayR shr tp) intTp combine mseed marr mseg = do + dev <- liftCodeGen $ gets ptxDeviceProperties + -- let - (start, end, paramGang) = gangParam - (arrOut, paramOut) = mutableArray ("out" :: Name (Array (sh :. Int) e)) - paramEnv = envParam aenv + (arrOut, paramOut) = mutableArray repr "out" + (arrIn, paramIn) = delayedArray "in" marr + (arrSeg, paramSeg) = delayedArray "seg" mseg + paramEnv = envParam aenv -- - config = launchConfig dev (CUDA.decWarp dev) dsmem const [|| const ||] - dsmem n = warps * (1 + per_warp) * bytes + config = launchConfig dev (CUDA.decWarp dev) dsmem const [|| const ||] + dsmem n = warps * (1 + per_warp) * bytes where ws = CUDA.warpSize dev warps = n `P.quot` ws per_warp = ws + ws `P.quot` 2 - bytes = sizeOf (eltType (undefined :: e)) - in - makeOpenAccWith config "foldSeg_block" (paramGang ++ paramOut ++ paramEnv) $ do + bytes = bytesElt tp + -- + makeOpenAccWith config "foldSeg_block" (paramOut ++ paramIn ++ paramSeg ++ paramEnv) $ do -- We use a dynamically scheduled work queue in order to evenly distribute -- the uneven workload, due to the variable length of each segment, over the @@ -130,7 +117,7 @@ -- a single coalesced read, since they will be sequential in the -- segment-offset array. -- - smem <- staticSharedMem 2 + smem <- staticSharedMem (TupRsingle scalarTypeInt) 2 -- Compute the number of segments and size of the innermost dimension. These -- are required if we are reducing a rank-2 or higher array, to properly @@ -139,52 +126,56 @@ -- 'scanl (+) 0' of the segment length array, so its size has increased by -- one. -- - sz <- indexHead <$> delayedExtent arr - ss <- do n <- indexHead <$> delayedExtent seg - A.sub numType n (lift 1) + sz <- indexHead <$> delayedExtent arrIn + ss <- do n <- indexHead <$> delayedExtent arrSeg + A.sub numType n (liftInt 1) -- Each thread block cooperatively reduces a segment. -- s0 <- dequeue queue (lift 1) -- for s0 (\s -> A.lt singleType s end) (\_ -> dequeue queue (lift 1)) $ \s -> do + + start <- return (liftInt 0) + end <- shapeSize shr (irArrayShape arrOut) + imapFromTo start end $ \s -> do -- The first two threads of the block determine the indices of the -- segments array that we will reduce between and distribute those values -- to the other threads in the block. tid <- threadIdx - when (A.lt singleType tid (lift 2)) $ do - i <- case rank (undefined::sh) of - 0 -> return s + when (A.lt singleType tid (liftInt32 2)) $ do + i <- case shr of + ShapeRsnoc ShapeRz -> return s _ -> A.rem integralType s ss j <- A.add numType i =<< int tid - v <- app1 (delayedLinearIndex seg) j - writeArray smem tid =<< int v + v <- app1 (delayedLinearIndex arrSeg) j + writeArray TypeInt32 smem tid =<< A.fromIntegral intTp numType v -- Once all threads have caught up, begin work on the new segment. __syncthreads - u <- readArray smem (lift 0 :: IR Int32) - v <- readArray smem (lift 1 :: IR Int32) + u <- readArray TypeInt32 smem (liftInt32 0) + v <- readArray TypeInt32 smem (liftInt32 1) -- Determine the index range of the input array we will reduce over. -- Necessary for multidimensional segmented reduction. - (inf,sup) <- A.unpair <$> case rank (undefined::sh) of - 0 -> return (A.pair u v) + (inf,sup) <- A.unpair <$> case shr of + ShapeRsnoc ShapeRz -> return (A.pair u v) _ -> do q <- A.quot integralType s ss a <- A.mul numType q sz A.pair <$> A.add numType u a <*> A.add numType v a void $ - if A.eq singleType inf sup + if (TupRunit, A.eq singleType inf sup) -- This segment is empty. If this is an exclusive reduction the -- first thread writes out the initial element for this segment. then do case mseed of - Nothing -> return (IR OP_Unit :: IR ()) + Nothing -> return (lift TupRunit ()) Just z -> do - when (A.eq singleType tid (lift 0)) $ writeArray arrOut s =<< z - return (IR OP_Unit) + when (A.eq singleType tid (liftInt32 0)) $ writeArray TypeInt arrOut s =<< z + return (lift TupRunit ()) -- This is a non-empty segment. else do @@ -200,38 +191,38 @@ -- a bug in NVPTX / ptxas. -- i0 <- A.add numType inf =<< int tid - x0 <- if A.lt singleType i0 sup - then app1 (delayedLinearIndex arr) i0 + x0 <- if (tp, A.lt singleType i0 sup) + then app1 (delayedLinearIndex arrIn) i0 else let - go :: TupleType a -> Operands a - go TypeRunit = OP_Unit - go (TypeRpair a b) = OP_Pair (go a) (go b) - go (TypeRscalar t) = ir' t (undef t) + go :: TypeR a -> Operands a + go TupRunit = OP_Unit + go (TupRpair a b) = OP_Pair (go a) (go b) + go (TupRsingle t) = ir t (undef t) in - return . IR $ go (eltType (undefined::e)) + return $ go tp bd <- int =<< blockDim v0 <- A.sub numType sup inf v0' <- i32 v0 - r0 <- if A.gte singleType v0 bd - then reduceBlockSMem dev combine Nothing x0 - else reduceBlockSMem dev combine (Just v0') x0 + r0 <- if (tp, A.gte singleType v0 bd) + then reduceBlockSMem dev tp combine Nothing x0 + else reduceBlockSMem dev tp combine (Just v0') x0 -- Step 2: keep walking over the input nxt <- A.add numType inf bd - r <- iterFromStepTo nxt bd sup r0 $ \offset r -> do + r <- iterFromStepTo tp nxt bd sup r0 $ \offset r -> do -- Wait for threads to catch up before starting the next stripe __syncthreads i' <- A.add numType offset =<< int tid v' <- A.sub numType sup offset - r' <- if A.gte singleType v' bd + r' <- if (tp, A.gte singleType v' bd) -- All threads in the block are in bounds, so we -- can avoid bounds checks. then do - x <- app1 (delayedLinearIndex arr) i' - y <- reduceBlockSMem dev combine Nothing x + x <- app1 (delayedLinearIndex arrIn) i' + y <- reduceBlockSMem dev tp combine Nothing x return y -- Not all threads are valid. Note that we still @@ -239,29 +230,36 @@ -- to avoid thread divergence on synchronisation -- points, similar to the above NOTE. else do - x <- if A.lt singleType i' sup - then app1 (delayedLinearIndex arr) i' - else return r + x <- if (tp, A.lt singleType i' sup) + then app1 (delayedLinearIndex arrIn) i' + else let + go :: TypeR a -> Operands a + go TupRunit = OP_Unit + go (TupRpair a b) = OP_Pair (go a) (go b) + go (TupRsingle t) = ir t (undef t) + in + return $ go tp + z <- i32 v' - y <- reduceBlockSMem dev combine (Just z) x + y <- reduceBlockSMem dev tp combine (Just z) x return y -- first thread incorporates the result from the previous -- iteration - if A.eq singleType tid (lift 0) + if (tp, A.eq singleType tid (liftInt32 0)) then app2 combine r r' else return r' -- Step 3: Thread zero writes the aggregate reduction for this -- segment to memory. If this is an exclusive fold combine with the -- initial element as well. - when (A.eq singleType tid (lift 0)) $ - writeArray arrOut s =<< + when (A.eq singleType tid (liftInt32 0)) $ + writeArray TypeInt arrOut s =<< case mseed of Nothing -> return r Just z -> flip (app2 combine) r =<< z -- Note: initial element on the left - return (IR OP_Unit) + return (lift TupRunit ()) return_ @@ -274,37 +272,40 @@ -- about inter- or intra-block synchronisation. -- mkFoldSegP_warp - :: forall aenv sh i e. (Shape sh, IsIntegral i, Elt i, Elt e) - => DeviceProperties - -> Gamma aenv - -> IRFun2 PTX aenv (e -> e -> e) - -> Maybe (IRExp PTX aenv e) - -> IRDelayed PTX aenv (Array (sh :. Int) e) - -> IRDelayed PTX aenv (Segments i) - -> CodeGen (IROpenAcc PTX aenv (Array (sh :. Int) e)) -mkFoldSegP_warp dev aenv combine mseed arr seg = + :: forall aenv sh i e. + Gamma aenv + -> ArrayR (Array (sh, Int) e) + -> IntegralType i + -> IRFun2 PTX aenv (e -> e -> e) + -> MIRExp PTX aenv e + -> MIRDelayed PTX aenv (Array (sh, Int) e) + -> MIRDelayed PTX aenv (Segments i) + -> CodeGen PTX (IROpenAcc PTX aenv (Array (sh, Int) e)) +mkFoldSegP_warp aenv repr@(ArrayR shr tp) intTp combine mseed marr mseg = do + dev <- liftCodeGen $ gets ptxDeviceProperties + -- let - (start, end, paramGang) = gangParam - (arrOut, paramOut) = mutableArray ("out" :: Name (Array (sh :. Int) e)) - paramEnv = envParam aenv + (arrOut, paramOut) = mutableArray repr "out" + (arrIn, paramIn) = delayedArray "in" marr + (arrSeg, paramSeg) = delayedArray "seg" mseg + paramEnv = envParam aenv -- - config = launchConfig dev (CUDA.decWarp dev) dsmem grid gridQ - dsmem n = warps * (2 + per_warp_elems) * bytes + config = launchConfig dev (CUDA.decWarp dev) dsmem grid gridQ + dsmem n = warps * per_warp_bytes where - warps = n `P.quot` ws + warps = (n + ws - 1) `P.quot` ws -- - grid n m = multipleOf n (m `P.quot` ws) - gridQ = [|| \n m -> $$multipleOfQ n (m `P.quot` ws) ||] + grid n m = multipleOf n (m `P.quot` ws) + gridQ = [|| \n m -> $$multipleOfQ n (m `P.quot` ws) ||] -- - per_warp_bytes = per_warp_elems * bytes - per_warp_elems = ws + (ws `P.quot` 2) - ws = CUDA.warpSize dev - bytes = sizeOf (eltType (undefined :: e)) + per_warp_bytes = (per_warp_elems * bytesElt tp) `P.max` (2 * bytesElt tp) + per_warp_elems = ws + (ws `P.quot` 2) + ws = CUDA.warpSize dev - int32 :: Integral a => a -> IR Int32 - int32 = lift . P.fromIntegral - in - makeOpenAccWith config "foldSeg_warp" (paramGang ++ paramOut ++ paramEnv) $ do + int32 :: Integral a => a -> Operands Int32 + int32 = liftInt32 . P.fromIntegral + -- + makeOpenAccWith config "foldSeg_warp" (paramOut ++ paramIn ++ paramSeg ++ paramEnv) $ do -- Each warp works independently. -- Determine the ID of this warp within the thread block. @@ -329,18 +330,24 @@ -- coalesced read, as these elements will be adjacent in the segment-offset -- array. -- + -- Note that this is aliased with the memory used to communicate reduction + -- values within the warp. + -- lim <- do - a <- A.mul numType wid (int32 (2 * bytes)) - b <- dynamicSharedMem (lift 2) a + a <- A.mul numType wid (int32 per_warp_bytes) + b <- dynamicSharedMem (TupRsingle scalarTypeInt) TypeInt32 (liftInt32 2) a return b - -- Allocate (1.5 * warpSize) elements of share memory for each warp + -- Allocate (1.5 * warpSize) elements of shared memory for each warp to + -- communicate reduction values. + -- + -- Note that this is aliased with the memory used to communicate the start + -- and end indices of this segment. + -- smem <- do - a <- A.mul numType wpb (int32 (2 * bytes)) - b <- A.mul numType wid (int32 per_warp_bytes) - c <- A.add numType a b - d <- dynamicSharedMem (int32 per_warp_elems) c - return d + a <- A.mul numType wid (int32 per_warp_bytes) + b <- dynamicSharedMem tp TypeInt32 (int32 per_warp_elems) a + return b -- Compute the number of segments and size of the innermost dimension. These -- are required if we are reducing a rank-2 or higher array, to properly @@ -348,55 +355,59 @@ -- reduces. Note that this is a segment-offset array computed by 'scanl (+) 0' -- of the segment length array, so its size has increased by one. -- - sz <- indexHead <$> delayedExtent arr - ss <- do a <- indexHead <$> delayedExtent seg - b <- A.sub numType a (lift 1) + sz <- indexHead <$> delayedExtent arrIn + ss <- do a <- indexHead <$> delayedExtent arrSeg + b <- A.sub numType a (liftInt 1) return b -- Each thread reduces a segment independently - s0 <- A.add numType start =<< int gwid + s0 <- int gwid gd <- int =<< gridDim wpb' <- int wpb step <- A.mul numType wpb' gd + end <- shapeSize shr (irArrayShape arrOut) imapFromStepTo s0 step end $ \s -> do + __syncwarp + -- The first two threads of the warp determine the indices of the segments -- array that we will reduce between and distribute those values to the -- other threads in the warp lane <- laneId - when (A.lt singleType lane (lift 2)) $ do - a <- case rank (undefined::sh) of - 0 -> return s + when (A.lt singleType lane (liftInt32 2)) $ do + a <- case shr of + ShapeRsnoc ShapeRz -> return s _ -> A.rem integralType s ss b <- A.add numType a =<< int lane - c <- app1 (delayedLinearIndex seg) b - writeArray lim lane =<< int c + c <- app1 (delayedLinearIndex arrSeg) b + writeArray TypeInt32 lim lane =<< A.fromIntegral intTp numType c + __syncwarp + -- Determine the index range of the input array we will reduce over. -- Necessary for multidimensional segmented reduction. (inf,sup) <- do - u <- readArray lim (lift 0 :: IR Int32) - v <- readArray lim (lift 1 :: IR Int32) - A.unpair <$> case rank (undefined::sh) of - 0 -> return (A.pair u v) + u <- readArray TypeInt32 lim (liftInt32 0) + v <- readArray TypeInt32 lim (liftInt32 1) + A.unpair <$> case shr of + ShapeRsnoc ShapeRz -> return (A.pair u v) _ -> do q <- A.quot integralType s ss a <- A.mul numType q sz A.pair <$> A.add numType u a <*> A.add numType v a - -- TLM: I don't think this should be necessary... - __syncthreads + __syncwarp void $ - if A.eq singleType inf sup + if (TupRunit, A.eq singleType inf sup) -- This segment is empty. If this is an exclusive reduction the first -- lane writes out the initial element for this segment. then do case mseed of - Nothing -> return (IR OP_Unit :: IR ()) + Nothing -> return (lift TupRunit ()) Just z -> do - when (A.eq singleType lane (lift 0)) $ writeArray arrOut s =<< z - return (IR OP_Unit) + when (A.eq singleType lane (liftInt32 0)) $ writeArray TypeInt arrOut s =<< z + return (lift TupRunit ()) -- This is a non-empty segment. else do @@ -405,69 +416,76 @@ -- See comment above why we initialise the loop in this way -- i0 <- A.add numType inf =<< int lane - x0 <- if A.lt singleType i0 sup - then app1 (delayedLinearIndex arr) i0 + x0 <- if (tp, A.lt singleType i0 sup) + then app1 (delayedLinearIndex arrIn) i0 else let - go :: TupleType a -> Operands a - go TypeRunit = OP_Unit - go (TypeRpair a b) = OP_Pair (go a) (go b) - go (TypeRscalar t) = ir' t (undef t) + go :: TypeR a -> Operands a + go TupRunit = OP_Unit + go (TupRpair a b) = OP_Pair (go a) (go b) + go (TupRsingle t) = ir t (undef t) in - return . IR $ go (eltType (undefined::e)) + return $ go tp v0 <- A.sub numType sup inf v0' <- i32 v0 - r0 <- if A.gte singleType v0 (lift ws) - then reduceWarpSMem dev combine smem Nothing x0 - else reduceWarpSMem dev combine smem (Just v0') x0 + r0 <- if (tp, A.gte singleType v0 (liftInt ws)) + then reduceWarpSMem dev tp combine smem Nothing x0 + else reduceWarpSMem dev tp combine smem (Just v0') x0 -- Step 2: Keep walking over the rest of the segment - nx <- A.add numType inf (lift ws) - r <- iterFromStepTo nx (lift ws) sup r0 $ \offset r -> do + nx <- A.add numType inf (liftInt ws) + r <- iterFromStepTo tp nx (liftInt ws) sup r0 $ \offset r -> do - -- TLM: Similarly, I think this is unnecessary... - __syncthreads + -- __syncwarp + __syncthreads -- TLM: why is this necessary? i' <- A.add numType offset =<< int lane v' <- A.sub numType sup offset - r' <- if A.gte singleType v' (lift ws) + r' <- if (tp, A.gte singleType v' (liftInt ws)) then do -- All lanes are in bounds, so avoid bounds checks - x <- app1 (delayedLinearIndex arr) i' - y <- reduceWarpSMem dev combine smem Nothing x + x <- app1 (delayedLinearIndex arrIn) i' + y <- reduceWarpSMem dev tp combine smem Nothing x return y else do - x <- if A.lt singleType i' sup - then app1 (delayedLinearIndex arr) i' - else return r + x <- if (tp, A.lt singleType i' sup) + then app1 (delayedLinearIndex arrIn) i' + else let + go :: TypeR a -> Operands a + go TupRunit = OP_Unit + go (TupRpair a b) = OP_Pair (go a) (go b) + go (TupRsingle t) = ir t (undef t) + in + return $ go tp + z <- i32 v' - y <- reduceWarpSMem dev combine smem (Just z) x + y <- reduceWarpSMem dev tp combine smem (Just z) x return y -- The first lane incorporates the result from the previous -- iteration - if A.eq singleType lane (lift 0) + if (tp, A.eq singleType lane (liftInt32 0)) then app2 combine r r' else return r' -- Step 3: Lane zero writes the aggregate reduction for this -- segment to memory. If this is an exclusive reduction, also -- combine with the initial element - when (A.eq singleType lane (lift 0)) $ - writeArray arrOut s =<< + when (A.eq singleType lane (liftInt32 0)) $ + writeArray TypeInt arrOut s =<< case mseed of Nothing -> return r Just z -> flip (app2 combine) r =<< z -- Note: initial element on the left - return (IR OP_Unit) + return (lift TupRunit ()) return_ -i32 :: IsIntegral i => IR i -> CodeGen (IR Int32) +i32 :: IsIntegral i => Operands i -> CodeGen PTX (Operands Int32) i32 = A.fromIntegral integralType numType -int :: IsIntegral i => IR i -> CodeGen (IR Int) +int :: IsIntegral i => Operands i -> CodeGen PTX (Operands Int) int = A.fromIntegral integralType numType diff --git a/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Generate.hs b/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Generate.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Generate.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Generate.hs @@ -1,11 +1,12 @@ +{-# LANGUAGE GADTs #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.CodeGen.Generate --- Copyright : [2014..2017] Trevor L. McDonell +-- Copyright : [2014..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- @@ -16,8 +17,10 @@ import Prelude hiding ( fromIntegral ) -- accelerate -import Data.Array.Accelerate.Array.Sugar ( Array, Shape, Elt ) +import Data.Array.Accelerate.Representation.Array +import Data.Array.Accelerate.Type +import Data.Array.Accelerate.LLVM.CodeGen.Arithmetic import Data.Array.Accelerate.LLVM.CodeGen.Array import Data.Array.Accelerate.LLVM.CodeGen.Base import Data.Array.Accelerate.LLVM.CodeGen.Environment @@ -34,24 +37,24 @@ -- processes multiple adjacent elements. -- mkGenerate - :: forall aenv sh e. (Shape sh, Elt e) - => PTX - -> Gamma aenv - -> IRFun1 PTX aenv (sh -> e) - -> CodeGen (IROpenAcc PTX aenv (Array sh e)) -mkGenerate ptx aenv apply = + :: Gamma aenv + -> ArrayR (Array sh e) + -> IRFun1 PTX aenv (sh -> e) + -> CodeGen PTX (IROpenAcc PTX aenv (Array sh e)) +mkGenerate aenv repr@(ArrayR shr _) apply = let - (start, end, paramGang) = gangParam - (arrOut, paramOut) = mutableArray ("out" :: Name (Array sh e)) - paramEnv = envParam aenv + (arrOut, paramOut) = mutableArray repr "out" + paramEnv = envParam aenv in - makeOpenAcc ptx "generate" (paramGang ++ paramOut ++ paramEnv) $ do + makeOpenAcc "generate" (paramOut ++ paramEnv) $ do + start <- return (liftInt 0) + end <- shapeSize shr (irArrayShape arrOut) + imapFromTo start end $ \i -> do - ix <- indexOfInt (irArrayShape arrOut) i -- convert to multidimensional index + ix <- indexOfInt shr (irArrayShape arrOut) i -- convert to multidimensional index r <- app1 apply ix -- apply generator function - writeArray arrOut i r -- store result + writeArray TypeInt arrOut i r -- store result return_ - diff --git a/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Intrinsic.hs b/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Intrinsic.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Intrinsic.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Intrinsic.hs @@ -2,10 +2,10 @@ {-# OPTIONS_GHC -fno-warn-orphans #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.CodeGen.Intrinsic --- Copyright : [2017] Trevor L. McDonell +-- Copyright : [2017..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- @@ -26,7 +26,7 @@ instance Intrinsic PTX where - intrinsicForTarget _ = libdeviceIndex + intrinsicForTarget = libdeviceIndex -- The list of functions implemented by libdevice. These are all more-or-less -- named consistently based on the standard mathematical functions they diff --git a/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Loop.hs b/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Loop.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Loop.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Loop.hs @@ -1,9 +1,9 @@ -- | -- Module : Data.Array.Accelerate.LLVM.PTX.CodeGen.Loop --- Copyright : [2015..2017] Trevor L. McDonell +-- Copyright : [2015..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- @@ -20,6 +20,7 @@ import qualified Data.Array.Accelerate.LLVM.CodeGen.Loop as Loop import Data.Array.Accelerate.LLVM.PTX.CodeGen.Base +import Data.Array.Accelerate.LLVM.PTX.Target -- | A standard loop where the CUDA threads cooperatively step over an index @@ -37,7 +38,7 @@ -- boundary. This might not always be the case, so provide a version that -- explicitly aligns reads to the warp boundary. -- -imapFromTo :: IR Int -> IR Int -> (IR Int -> CodeGen ()) -> CodeGen () +imapFromTo :: Operands Int -> Operands Int -> (Operands Int -> CodeGen PTX ()) -> CodeGen PTX () imapFromTo start end body = do step <- A.fromIntegral integralType numType =<< gridSize tid <- A.fromIntegral integralType numType =<< globalThreadIdx diff --git a/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Map.hs b/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Map.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Map.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Map.hs @@ -1,13 +1,13 @@ {-# LANGUAGE GADTs #-} {-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE RecordWildCards #-} {-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.CodeGen.Map --- Copyright : [2014..2017] Trevor L. McDonell +-- Copyright : [2014..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- @@ -15,14 +15,16 @@ module Data.Array.Accelerate.LLVM.PTX.CodeGen.Map where -import Prelude hiding ( fromIntegral ) - -- accelerate -import Data.Array.Accelerate.Array.Sugar ( Array, Elt ) +import Data.Array.Accelerate.Representation.Array +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Type +import Data.Array.Accelerate.LLVM.CodeGen.Arithmetic import Data.Array.Accelerate.LLVM.CodeGen.Array import Data.Array.Accelerate.LLVM.CodeGen.Base import Data.Array.Accelerate.LLVM.CodeGen.Environment +import Data.Array.Accelerate.LLVM.CodeGen.Exp import Data.Array.Accelerate.LLVM.CodeGen.Monad import Data.Array.Accelerate.LLVM.CodeGen.Sugar @@ -34,24 +36,26 @@ -- Apply a unary function to each element of an array. Each thread processes -- multiple elements, striding the array by the grid size. -- -mkMap :: forall aenv sh a b. Elt b - => PTX - -> Gamma aenv - -> IRFun1 PTX aenv (a -> b) - -> IRDelayed PTX aenv (Array sh a) - -> CodeGen (IROpenAcc PTX aenv (Array sh b)) -mkMap ptx aenv apply IRDelayed{..} = +mkMap :: Gamma aenv + -> ArrayR (Array sh a) + -> TypeR b + -> IRFun1 PTX aenv (a -> b) + -> CodeGen PTX (IROpenAcc PTX aenv (Array sh b)) +mkMap aenv repr@(ArrayR shr _) tp' apply = let - (start, end, paramGang) = gangParam - (arrOut, paramOut) = mutableArray ("out" :: Name (Array sh b)) - paramEnv = envParam aenv + (arrOut, paramOut) = mutableArray (ArrayR shr tp') "out" + (arrIn, paramIn) = mutableArray repr "in" + paramEnv = envParam aenv in - makeOpenAcc ptx "map" (paramGang ++ paramOut ++ paramEnv) $ do + makeOpenAcc "map" (paramOut ++ paramIn ++ paramEnv) $ do + start <- return (liftInt 0) + end <- shapeSize shr (irArrayShape arrIn) + imapFromTo start end $ \i -> do - xs <- app1 delayedLinearIndex i + xs <- readArray TypeInt arrIn i ys <- app1 apply xs - writeArray arrOut i ys + writeArray TypeInt arrOut i ys return_ diff --git a/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Permute.hs b/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Permute.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Permute.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Permute.hs @@ -3,14 +3,14 @@ {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE RecordWildCards #-} {-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeApplications #-} {-# LANGUAGE ViewPatterns #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.CodeGen.Permute --- Copyright : [2016..2017] Trevor L. McDonell +-- Copyright : [2016..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- @@ -21,11 +21,12 @@ ) where --- accelerate -import Data.Array.Accelerate.Analysis.Type -import Data.Array.Accelerate.Array.Sugar ( Array, Vector, Shape, Elt, eltType ) +import Data.Array.Accelerate.AST import Data.Array.Accelerate.Error -import qualified Data.Array.Accelerate.Array.Sugar as S +import Data.Array.Accelerate.Representation.Array +import Data.Array.Accelerate.Representation.Elt +import Data.Array.Accelerate.Representation.Shape +import Data.Array.Accelerate.Representation.Type import Data.Array.Accelerate.LLVM.CodeGen.Arithmetic as A import Data.Array.Accelerate.LLVM.CodeGen.Array @@ -41,7 +42,6 @@ import Data.Array.Accelerate.LLVM.PTX.CodeGen.Base import Data.Array.Accelerate.LLVM.PTX.CodeGen.Loop -import Data.Array.Accelerate.LLVM.PTX.Context import Data.Array.Accelerate.LLVM.PTX.Target import LLVM.AST.Type.AddrSpace @@ -54,8 +54,8 @@ import Foreign.CUDA.Analysis -import Data.Typeable import Control.Monad ( void ) +import Control.Monad.State ( gets ) import Prelude @@ -78,21 +78,18 @@ -- a queue or some such. -- mkPermute - :: forall aenv sh sh' e. (Shape sh, Shape sh', Elt e) - => PTX - -> Gamma aenv + :: HasCallStack + => Gamma aenv + -> ArrayR (Array sh e) + -> ShapeR sh' -> IRPermuteFun PTX aenv (e -> e -> e) - -> IRFun1 PTX aenv (sh -> sh') - -> IRDelayed PTX aenv (Array sh e) - -> CodeGen (IROpenAcc PTX aenv (Array sh' e)) -mkPermute ptx aenv IRPermuteFun{..} project arr = - let - bytes = sizeOf (eltType (undefined :: e)) - sizeOk = bytes == 4 || bytes == 8 - in + -> IRFun1 PTX aenv (sh -> PrimMaybe sh') + -> MIRDelayed PTX aenv (Array sh e) + -> CodeGen PTX (IROpenAcc PTX aenv (Array sh' e)) +mkPermute aenv repr shr' IRPermuteFun{..} project arr = case atomicRMW of - Just (rmw, f) | sizeOk -> mkPermute_rmw ptx aenv rmw f project arr - _ -> mkPermute_mutex ptx aenv combine project arr + Just (rmw, f) -> mkPermute_rmw aenv repr shr' rmw f project arr + _ -> mkPermute_mutex aenv repr shr' combine project arr -- Parallel forward permutation function which uses atomic instructions to @@ -102,68 +99,74 @@ -- the element type and compute capability of the target hardware we may need to -- emulate the operation using atomic compare-and-swap. -- --- Int32 Int64 Float32 Float64 --- +---------------------------------------- --- (+) | 2.0 2.0 2.0 6.0 --- (-) | 2.0 2.0 x x +-- Int32 Int64 Float16 Float32 Float64 +-- +------------------------------------------------- +-- (+) | 2.0 2.0 7.0 2.0 6.0 +-- (-) | 2.0 2.0 x x x -- (.&.) | 2.0 3.2 -- (.|.) | 2.0 3.2 -- xor | 2.0 3.2 --- min | 2.0 3.2 x x --- max | 2.0 3.2 x x +-- min | 2.0 3.2 x x x +-- max | 2.0 3.2 x x x -- CAS | 2.0 2.0 -- -- Note that NVPTX requires at least compute 2.0, so we can always implement the -- lockfree update operations in terms of compare-and-swap. -- mkPermute_rmw - :: forall aenv sh sh' e. (Shape sh, Shape sh', Elt e) - => PTX - -> Gamma aenv + :: HasCallStack + => Gamma aenv + -> ArrayR (Array sh e) + -> ShapeR sh' -> RMWOperation - -> IRFun1 PTX aenv (e -> e) - -> IRFun1 PTX aenv (sh -> sh') - -> IRDelayed PTX aenv (Array sh e) - -> CodeGen (IROpenAcc PTX aenv (Array sh' e)) -mkPermute_rmw ptx@(deviceProperties . ptxContext -> dev) aenv rmw update project IRDelayed{..} = + -> IRFun1 PTX aenv (e -> e) + -> IRFun1 PTX aenv (sh -> PrimMaybe sh') + -> MIRDelayed PTX aenv (Array sh e) + -> CodeGen PTX (IROpenAcc PTX aenv (Array sh' e)) +mkPermute_rmw aenv (ArrayR shr tp) shr' rmw update project marr = do + dev <- liftCodeGen $ gets ptxDeviceProperties + -- let - (start, end, paramGang) = gangParam - (arrOut, paramOut) = mutableArray ("out" :: Name (Array sh' e)) - paramEnv = envParam aenv + outR = ArrayR shr' tp + (arrOut, paramOut) = mutableArray outR "out" + (arrIn, paramIn) = delayedArray "in" marr + paramEnv = envParam aenv + start = liftInt 0 -- - bytes = sizeOf (eltType (undefined :: e)) - compute = computeCapability dev - compute32 = Compute 3 2 - compute60 = Compute 6 0 - in - makeOpenAcc ptx "permute_rmw" (paramGang ++ paramOut ++ paramEnv) $ do + bytes = bytesElt tp + compute = computeCapability dev + compute32 = Compute 3 2 + compute60 = Compute 6 0 + compute70 = Compute 7 0 + -- + makeOpenAcc "permute_rmw" (paramOut ++ paramIn ++ paramEnv) $ do - sh <- delayedExtent + shIn <- delayedExtent arrIn + end <- shapeSize shr shIn imapFromTo start end $ \i -> do - ix <- indexOfInt sh i + ix <- indexOfInt shr shIn i ix' <- app1 project ix - unless (ignore ix') $ do - j <- intOfIndex (irArrayShape arrOut) ix' - x <- app1 delayedLinearIndex i + when (isJust ix') $ do + j <- intOfIndex shr' (irArrayShape arrOut) =<< fromJust ix' + x <- app1 (delayedLinearIndex arrIn) i r <- app1 update x case rmw of Exchange - -> writeArray arrOut j r + -> writeArray TypeInt arrOut j r -- - _ | TypeRscalar (SingleScalarType s) <- eltType (undefined::e) - , Just adata <- gcast (irArrayData arrOut) - , Just r' <- gcast r + _ | TupRsingle (SingleScalarType s) <- tp + , adata <- irArrayData arrOut -> do addr <- instr' $ GetElementPtr (asPtr defaultAddrSpace (op s adata)) [op integralType j] -- let - rmw_integral :: IntegralType t -> Operand (Ptr t) -> Operand t -> CodeGen () + rmw_integral :: IntegralType t -> Operand (Ptr t) -> Operand t -> CodeGen PTX () rmw_integral t ptr val - | primOk = void . instr' $ AtomicRMW t NonVolatile rmw ptr val (CrossThread, AcquireRelease) + | primOk = void . instr' $ AtomicRMW (IntegralNumType t) NonVolatile rmw ptr val (CrossThread, AcquireRelease) | otherwise = case rmw of RMW.And -> atomicCAS_rmw s' (A.band t (ir t val)) ptr @@ -171,7 +174,7 @@ RMW.Xor -> atomicCAS_rmw s' (A.xor t (ir t val)) ptr RMW.Min -> atomicCAS_cmp s' A.lt ptr val RMW.Max -> atomicCAS_cmp s' A.gt ptr val - _ -> $internalError "mkPermute_rmw.integral" "unexpected transition" + _ -> internalError "unexpected transition" where s' = NumSingleType (IntegralNumType t) primOk = compute >= compute32 @@ -181,7 +184,7 @@ RMW.Sub -> True _ -> False - rmw_floating :: FloatingType t -> Operand (Ptr t) -> Operand t -> CodeGen () + rmw_floating :: FloatingType t -> Operand (Ptr t) -> Operand t -> CodeGen PTX () rmw_floating t ptr val = case rmw of RMW.Min -> atomicCAS_cmp s' A.lt ptr val @@ -190,30 +193,20 @@ RMW.Add | primAdd -> atomicAdd_f t ptr val | otherwise -> atomicCAS_rmw s' (A.add n (ir t val)) ptr - _ -> $internalError "mkPermute_rmw.floating" "unexpected transition" + _ -> internalError "unexpected transition" where n = FloatingNumType t s' = NumSingleType n - primAdd = bytes == 4 - -- Available directly in LLVM-6 and later; - -- earlier versions could use inline assembly -#if MIN_VERSION_llvm_hs_pure(6,0,0) - || compute >= compute60 -#endif - - rmw_nonnum :: NonNumType t -> Operand (Ptr t) -> Operand t -> CodeGen () - rmw_nonnum TypeChar{} ptr val = do - ptr32 <- instr' $ PtrCast (primType :: PrimType (Ptr Word32)) ptr - val32 <- instr' $ BitCast (scalarType :: ScalarType Word32) val - void $ instr' $ AtomicRMW (integralType :: IntegralType Word32) NonVolatile rmw ptr32 val32 (CrossThread, AcquireRelease) - rmw_nonnum _ _ _ = -- C character types are 8-bit, and thus not supported - $internalError "mkPermute_rmw.nonnum" "unexpected transition" + primAdd = + case t of + TypeHalf -> compute >= compute70 + TypeFloat -> True + TypeDouble -> compute >= compute60 case s of - NumSingleType (IntegralNumType t) -> rmw_integral t addr (op t r') - NumSingleType (FloatingNumType t) -> rmw_floating t addr (op t r') - NonNumSingleType t -> rmw_nonnum t addr (op t r') + NumSingleType (IntegralNumType t) -> rmw_integral t addr (op t r) + NumSingleType (FloatingNumType t) -> rmw_floating t addr (op t r) -- - _ -> $internalError "mkPermute_rmw" "unexpected transition" + _ -> internalError "unexpected transition" return_ @@ -222,42 +215,134 @@ -- a mutex before updating the value at that location. -- mkPermute_mutex - :: forall aenv sh sh' e. (Shape sh, Shape sh', Elt e) - => PTX - -> Gamma aenv - -> IRFun2 PTX aenv (e -> e -> e) - -> IRFun1 PTX aenv (sh -> sh') - -> IRDelayed PTX aenv (Array sh e) - -> CodeGen (IROpenAcc PTX aenv (Array sh' e)) -mkPermute_mutex ptx aenv combine project IRDelayed{..} = + :: Gamma aenv + -> ArrayR (Array sh e) + -> ShapeR sh' + -> IRFun2 PTX aenv (e -> e -> e) + -> IRFun1 PTX aenv (sh -> PrimMaybe sh') + -> MIRDelayed PTX aenv (Array sh e) + -> CodeGen PTX (IROpenAcc PTX aenv (Array sh' e)) +mkPermute_mutex aenv (ArrayR shr tp) shr' combine project marr = let - (start, end, paramGang) = gangParam - (arrOut, paramOut) = mutableArray ("out" :: Name (Array sh' e)) - (arrLock, paramLock) = mutableArray ("lock" :: Name (Vector Word32)) - paramEnv = envParam aenv + outR = ArrayR shr' tp + lockR = ArrayR (ShapeRsnoc ShapeRz) (TupRsingle scalarTypeWord32) + (arrOut, paramOut) = mutableArray outR "out" + (arrLock, paramLock) = mutableArray lockR "lock" + (arrIn, paramIn) = delayedArray "in" marr + paramEnv = envParam aenv + start = liftInt 0 in - makeOpenAcc ptx "permute_mutex" (paramGang ++ paramOut ++ paramLock ++ paramEnv) $ do + makeOpenAcc "permute_mutex" (paramOut ++ paramLock ++ paramIn ++ paramEnv) $ do - sh <- delayedExtent + shIn <- delayedExtent arrIn + end <- shapeSize shr shIn imapFromTo start end $ \i -> do - ix <- indexOfInt sh i + ix <- indexOfInt shr shIn i ix' <- app1 project ix -- project element onto the destination array and (atomically) update - unless (ignore ix') $ do - j <- intOfIndex (irArrayShape arrOut) ix' - x <- app1 delayedLinearIndex i + when (isJust ix') $ do + j <- intOfIndex shr' (irArrayShape arrOut) =<< fromJust ix' + x <- app1 (delayedLinearIndex arrIn) i atomically arrLock j $ do - y <- readArray arrOut j + y <- readArray TypeInt arrOut j r <- app2 combine x y - writeArray arrOut j r + writeArray TypeInt arrOut j r return_ +-- Atomically execute the critical section only when the lock at the given +-- array indexed is obtained. +-- +atomically + :: IRArray (Vector Word32) + -> Operands Int + -> CodeGen PTX a + -> CodeGen PTX a +atomically barriers i action = do + dev <- liftCodeGen $ gets ptxDeviceProperties + if computeCapability dev >= Compute 7 0 + then atomically_thread barriers i action + else atomically_warp barriers i action + + +-- Atomically execute the critical section only when the lock at the given +-- array index is obtained. The thread spins waiting for the lock to be +-- released with exponential backoff on failure in case the lock is +-- contended. +-- +-- > uint32_t ns = 8; +-- > while ( atomic_exchange(&lock[i], 1) == 1 ) { +-- > __nanosleep(ns); +-- > if ( ns < 256 ) { +-- > ns *= 2; +-- > } +-- > } +-- +-- Requires independent thread scheduling features of SM7+. +-- +atomically_thread + :: IRArray (Vector Word32) + -> Operands Int + -> CodeGen PTX a + -> CodeGen PTX a +atomically_thread barriers i action = do + let + lock = integral integralType 1 + unlock = integral integralType 0 + unlock' = ir TypeWord32 unlock + i32 = TupRsingle scalarTypeInt32 + -- + entry <- newBlock "spinlock.entry" + sleep <- newBlock "spinlock.backoff" + moar <- newBlock "spinlock.backoff-moar" + start <- newBlock "spinlock.critical-start" + end <- newBlock "spinlock.critical-end" + exit <- newBlock "spinlock.exit" + ns <- fresh i32 + + addr <- instr' $ GetElementPtr (asPtr defaultAddrSpace (op integralType (irArrayData barriers))) [op integralType i] + top <- br entry + + -- Loop until this thread has completed its critical section. If the slot + -- was unlocked we just acquired the lock and the thread can perform its + -- critical section, otherwise sleep the thread and try again later. + setBlock entry + old <- instr $ AtomicRMW numType NonVolatile Exchange addr lock (CrossThread, Acquire) + ok <- A.eq singleType old unlock' + _ <- cbr ok start sleep + + -- We did not acquire the lock. Sleep the thread for a small amount of + -- time and (possibly) increase the sleep duration for the next round + setBlock sleep + _ <- nanosleep ns + p <- A.lt singleType ns (ir TypeInt32 (integral integralType 256)) + _ <- cbr p moar entry + + setBlock moar + ns' <- A.mul numType ns (ir TypeInt32 (integral integralType 2)) + _ <- phi' i32 entry ns [(ir TypeInt32 (integral (integralType) 8), top), (ns, sleep), (ns', moar)] + _ <- br entry + + -- If we just acquired the lock, execute the critical section, then + -- release the lock and continue with your day. + setBlock start + r <- action + _ <- br end + + setBlock end + _ <- instr $ AtomicRMW numType NonVolatile Exchange addr unlock (CrossThread, AcquireRelease) + _ <- __threadfence_grid -- TODO: why is this required? + _ <- br exit + + setBlock exit + return r + + -- Atomically execute the critical section only when the lock at the given array -- index is obtained. The thread spins waiting for the lock to be released and -- there is no backoff strategy in case the lock is contended. @@ -302,68 +387,50 @@ -- > } -- > } while ( done == 0 ); -- -atomically +atomically_warp :: IRArray (Vector Word32) - -> IR Int - -> CodeGen a - -> CodeGen a -atomically barriers i action = do + -> Operands Int + -> CodeGen PTX a + -> CodeGen PTX a +atomically_warp barriers i action = do let lock = integral integralType 1 unlock = integral integralType 0 - unlock' = lift 0 + unlock' = ir TypeWord32 unlock -- - spin <- newBlock "spinlock.entry" - crit <- newBlock "spinlock.critical-start" - skip <- newBlock "spinlock.critical-end" - exit <- newBlock "spinlock.exit" + entry <- newBlock "spinlock.entry" + start <- newBlock "spinlock.critical-start" + end <- newBlock "spinlock.critical-end" + exit <- newBlock "spinlock.exit" addr <- instr' $ GetElementPtr (asPtr defaultAddrSpace (op integralType (irArrayData barriers))) [op integralType i] - _ <- br spin + _ <- br entry -- Loop until this thread has completed its critical section. If the slot was -- unlocked then we just acquired the lock and the thread can perform the -- critical section, otherwise skip to the bottom of the critical section. - setBlock spin - old <- instr $ AtomicRMW integralType NonVolatile Exchange addr lock (CrossThread, Acquire) + setBlock entry + old <- instr $ AtomicRMW numType NonVolatile Exchange addr lock (CrossThread, Acquire) ok <- A.eq singleType old unlock' - no <- cbr ok crit skip + no <- cbr ok start end -- If we just acquired the lock, execute the critical section - setBlock crit + setBlock start r <- action - _ <- instr $ AtomicRMW integralType NonVolatile Exchange addr unlock (CrossThread, Release) - yes <- br skip + _ <- instr $ AtomicRMW numType NonVolatile Exchange addr unlock (CrossThread, AcquireRelease) + yes <- br end -- At the base of the critical section, threads participate in a memory fence -- to ensure the lock state is committed to memory. Depending on which -- incoming edge the thread arrived at this block from determines whether they -- have completed their critical section. - setBlock skip - done <- phi [(lift True, yes), (lift False, no)] + setBlock end + res <- freshName + done <- phi1 end res [(boolean True, yes), (boolean False, no)] __syncthreads - _ <- cbr done exit spin + _ <- cbr (OP_Bool done) exit entry setBlock exit return r - - --- Helper functions --- ---------------- - --- Test whether the given index is the magic value 'ignore'. This operates --- strictly rather than performing short-circuit (&&). --- -ignore :: forall ix. Shape ix => IR ix -> CodeGen (IR Bool) -ignore (IR ix) = go (S.eltType (undefined::ix)) (S.fromElt (S.ignore::ix)) ix - where - go :: TupleType t -> t -> Operands t -> CodeGen (IR Bool) - go TypeRunit () OP_Unit = return (lift True) - go (TypeRpair tsh tsz) (ish, isz) (OP_Pair sh sz) = do x <- go tsh ish sh - y <- go tsz isz sz - land' x y - go (TypeRscalar s) ig sz = case s of - SingleScalarType t -> A.eq t (ir t (single t ig)) (ir t (op' t sz)) - VectorScalarType{} -> $internalError "ignore" "unexpected shape type" diff --git a/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Scan.hs b/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Scan.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Scan.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Scan.hs @@ -5,1360 +5,1316 @@ {-# LANGUAGE RecordWildCards #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE ViewPatterns #-} --- | --- Module : Data.Array.Accelerate.LLVM.PTX.CodeGen.Scan --- Copyright : [2016..2017] Trevor L. McDonell --- License : BSD3 --- --- Maintainer : Trevor L. McDonell --- Stability : experimental --- Portability : non-portable (GHC extensions) --- - -module Data.Array.Accelerate.LLVM.PTX.CodeGen.Scan ( - - mkScanl, mkScanl1, mkScanl', - mkScanr, mkScanr1, mkScanr', - -) where - --- accelerate -import Data.Array.Accelerate.Analysis.Type -import Data.Array.Accelerate.Array.Sugar - -import Data.Array.Accelerate.LLVM.Analysis.Match -import Data.Array.Accelerate.LLVM.CodeGen.Arithmetic as A -import Data.Array.Accelerate.LLVM.CodeGen.Array -import Data.Array.Accelerate.LLVM.CodeGen.Base -import Data.Array.Accelerate.LLVM.CodeGen.Constant -import Data.Array.Accelerate.LLVM.CodeGen.Environment -import Data.Array.Accelerate.LLVM.CodeGen.Exp -import Data.Array.Accelerate.LLVM.CodeGen.IR -import Data.Array.Accelerate.LLVM.CodeGen.Loop -import Data.Array.Accelerate.LLVM.CodeGen.Monad -import Data.Array.Accelerate.LLVM.CodeGen.Sugar -import Data.Array.Accelerate.LLVM.PTX.Analysis.Launch -import Data.Array.Accelerate.LLVM.PTX.CodeGen.Base -import Data.Array.Accelerate.LLVM.PTX.CodeGen.Generate -import Data.Array.Accelerate.LLVM.PTX.Context -import Data.Array.Accelerate.LLVM.PTX.Target - -import LLVM.AST.Type.Representation - -import qualified Foreign.CUDA.Analysis as CUDA - -import Control.Applicative -import Control.Monad ( (>=>), void ) -import Data.String ( fromString ) -import Data.Coerce as Safe -import Data.Bits as P -import Prelude as P hiding ( last ) - - -data Direction = L | R - --- 'Data.List.scanl' style left-to-right exclusive scan, but with the --- restriction that the combination function must be associative to enable --- efficient parallel implementation. --- --- > scanl (+) 10 (use $ fromList (Z :. 10) [0..]) --- > --- > ==> Array (Z :. 11) [10,10,11,13,16,20,25,31,38,46,55] --- -mkScanl - :: forall aenv sh e. (Shape sh, Elt e) - => PTX - -> Gamma aenv - -> IRFun2 PTX aenv (e -> e -> e) - -> IRExp PTX aenv e - -> IRDelayed PTX aenv (Array (sh:.Int) e) - -> CodeGen (IROpenAcc PTX aenv (Array (sh:.Int) e)) -mkScanl ptx@(deviceProperties . ptxContext -> dev) aenv combine seed arr - | Just Refl <- matchShapeType (undefined::sh) (undefined::Z) - = foldr1 (+++) <$> sequence [ mkScanAllP1 L dev aenv combine (Just seed) arr - , mkScanAllP2 L dev aenv combine - , mkScanAllP3 L dev aenv combine (Just seed) - , mkScanFill ptx aenv seed - ] - -- - | otherwise - = (+++) <$> mkScanDim L dev aenv combine (Just seed) arr - <*> mkScanFill ptx aenv seed - - --- 'Data.List.scanl1' style left-to-right inclusive scan, but with the --- restriction that the combination function must be associative to enable --- efficient parallel implementation. The array must not be empty. --- --- > scanl1 (+) (use $ fromList (Z :. 10) [0..]) --- > --- > ==> Array (Z :. 10) [0,1,3,6,10,15,21,28,36,45] --- -mkScanl1 - :: forall aenv sh e. (Shape sh, Elt e) - => PTX - -> Gamma aenv - -> IRFun2 PTX aenv (e -> e -> e) - -> IRDelayed PTX aenv (Array (sh:.Int) e) - -> CodeGen (IROpenAcc PTX aenv (Array (sh:.Int) e)) -mkScanl1 (deviceProperties . ptxContext -> dev) aenv combine arr - | Just Refl <- matchShapeType (undefined::sh) (undefined::Z) - = foldr1 (+++) <$> sequence [ mkScanAllP1 L dev aenv combine Nothing arr - , mkScanAllP2 L dev aenv combine - , mkScanAllP3 L dev aenv combine Nothing - ] - -- - | otherwise - = mkScanDim L dev aenv combine Nothing arr - - --- Variant of 'scanl' where the final result is returned in a separate array. --- --- > scanr' (+) 10 (use $ fromList (Z :. 10) [0..]) --- > --- > ==> ( Array (Z :. 10) [10,10,11,13,16,20,25,31,38,46] --- , Array Z [55] --- ) --- -mkScanl' - :: forall aenv sh e. (Shape sh, Elt e) - => PTX - -> Gamma aenv - -> IRFun2 PTX aenv (e -> e -> e) - -> IRExp PTX aenv e - -> IRDelayed PTX aenv (Array (sh:.Int) e) - -> CodeGen (IROpenAcc PTX aenv (Array (sh:.Int) e, Array sh e)) -mkScanl' ptx@(deviceProperties . ptxContext -> dev) aenv combine seed arr - | Just Refl <- matchShapeType (undefined::sh) (undefined::Z) - = foldr1 (+++) <$> sequence [ mkScan'AllP1 L dev aenv combine seed arr - , mkScan'AllP2 L dev aenv combine - , mkScan'AllP3 L dev aenv combine - , mkScan'Fill ptx aenv seed - ] - -- - | otherwise - = (+++) <$> mkScan'Dim L dev aenv combine seed arr - <*> mkScan'Fill ptx aenv seed - - --- 'Data.List.scanr' style right-to-left exclusive scan, but with the --- restriction that the combination function must be associative to enable --- efficient parallel implementation. --- --- > scanr (+) 10 (use $ fromList (Z :. 10) [0..]) --- > --- > ==> Array (Z :. 11) [55,55,54,52,49,45,40,34,27,19,10] --- -mkScanr - :: forall aenv sh e. (Shape sh, Elt e) - => PTX - -> Gamma aenv - -> IRFun2 PTX aenv (e -> e -> e) - -> IRExp PTX aenv e - -> IRDelayed PTX aenv (Array (sh:.Int) e) - -> CodeGen (IROpenAcc PTX aenv (Array (sh:.Int) e)) -mkScanr ptx@(deviceProperties . ptxContext -> dev) aenv combine seed arr - | Just Refl <- matchShapeType (undefined::sh) (undefined::Z) - = foldr1 (+++) <$> sequence [ mkScanAllP1 R dev aenv combine (Just seed) arr - , mkScanAllP2 R dev aenv combine - , mkScanAllP3 R dev aenv combine (Just seed) - , mkScanFill ptx aenv seed - ] - -- - | otherwise - = (+++) <$> mkScanDim R dev aenv combine (Just seed) arr - <*> mkScanFill ptx aenv seed - - --- 'Data.List.scanr1' style right-to-left inclusive scan, but with the --- restriction that the combination function must be associative to enable --- efficient parallel implementation. The array must not be empty. --- --- > scanr (+) 10 (use $ fromList (Z :. 10) [0..]) --- > --- > ==> Array (Z :. 10) [45,45,44,42,39,35,30,24,17,9] --- -mkScanr1 - :: forall aenv sh e. (Shape sh, Elt e) - => PTX - -> Gamma aenv - -> IRFun2 PTX aenv (e -> e -> e) - -> IRDelayed PTX aenv (Array (sh:.Int) e) - -> CodeGen (IROpenAcc PTX aenv (Array (sh:.Int) e)) -mkScanr1 (deviceProperties . ptxContext -> dev) aenv combine arr - | Just Refl <- matchShapeType (undefined::sh) (undefined::Z) - = foldr1 (+++) <$> sequence [ mkScanAllP1 R dev aenv combine Nothing arr - , mkScanAllP2 R dev aenv combine - , mkScanAllP3 R dev aenv combine Nothing - ] - -- - | otherwise - = mkScanDim R dev aenv combine Nothing arr - - --- Variant of 'scanr' where the final result is returned in a separate array. --- --- > scanr' (+) 10 (use $ fromList (Z :. 10) [0..]) --- > --- > ==> ( Array (Z :. 10) [55,54,52,49,45,40,34,27,19,10] --- , Array Z [55] --- ) --- -mkScanr' - :: forall aenv sh e. (Shape sh, Elt e) - => PTX - -> Gamma aenv - -> IRFun2 PTX aenv (e -> e -> e) - -> IRExp PTX aenv e - -> IRDelayed PTX aenv (Array (sh:.Int) e) - -> CodeGen (IROpenAcc PTX aenv (Array (sh:.Int) e, Array sh e)) -mkScanr' ptx@(deviceProperties . ptxContext -> dev) aenv combine seed arr - | Just Refl <- matchShapeType (undefined::sh) (undefined::Z) - = foldr1 (+++) <$> sequence [ mkScan'AllP1 R dev aenv combine seed arr - , mkScan'AllP2 R dev aenv combine - , mkScan'AllP3 R dev aenv combine - , mkScan'Fill ptx aenv seed - ] - -- - | otherwise - = (+++) <$> mkScan'Dim R dev aenv combine seed arr - <*> mkScan'Fill ptx aenv seed - - --- Device wide scans --- ----------------- --- --- This is a classic two-pass algorithm which proceeds in two phases and --- requires ~4n data movement to global memory. In future we would like to --- replace this with a single pass algorithm. --- - --- Parallel scan, step 1. --- --- Threads scan a stripe of the input into a temporary array, incorporating the --- initial element and any fused functions on the way. The final reduction --- result of this chunk is written to a separate array. --- -mkScanAllP1 - :: forall aenv e. Elt e - => Direction - -> DeviceProperties -- ^ properties of the target GPU - -> Gamma aenv -- ^ array environment - -> IRFun2 PTX aenv (e -> e -> e) -- ^ combination function - -> Maybe (IRExp PTX aenv e) -- ^ seed element, if this is an exclusive scan - -> IRDelayed PTX aenv (Vector e) -- ^ input data - -> CodeGen (IROpenAcc PTX aenv (Vector e)) -mkScanAllP1 dir dev aenv combine mseed IRDelayed{..} = - let - (start, end, paramGang) = gangParam - (arrOut, paramOut) = mutableArray ("out" :: Name (Vector e)) - (arrTmp, paramTmp) = mutableArray ("tmp" :: Name (Vector e)) - paramEnv = envParam aenv - -- - config = launchConfig dev (CUDA.incWarp dev) smem const [|| const ||] - smem n = warps * (1 + per_warp) * bytes - where - ws = CUDA.warpSize dev - warps = n `P.quot` ws - per_warp = ws + ws `P.quot` 2 - bytes = sizeOf (eltType (undefined :: e)) - in - makeOpenAccWith config "scanP1" (paramGang ++ paramTmp ++ paramOut ++ paramEnv) $ do - - -- Size of the input array - sz <- indexHead <$> delayedExtent - - -- A thread block scans a non-empty stripe of the input, storing the final - -- block-wide aggregate into a separate array - -- - -- For exclusive scans, thread 0 of segment 0 must incorporate the initial - -- element into the input and output. Threads shuffle their indices - -- appropriately. - -- - bid <- blockIdx - gd <- gridDim - gd' <- int gd - s0 <- A.add numType start =<< int bid - - -- iterating over thread-block-wide segments - imapFromStepTo s0 gd' end $ \chunk -> do - - bd <- blockDim - bd' <- int bd - inf <- A.mul numType chunk bd' - - -- index i* is the index that this thread will read data from. Recall that - -- the supremum index is exclusive - tid <- threadIdx - tid' <- int tid - i0 <- case dir of - L -> A.add numType inf tid' - R -> do x <- A.sub numType sz inf - y <- A.sub numType x tid' - z <- A.sub numType y (lift 1) - return z - - -- index j* is the index that we write to. Recall that for exclusive scans - -- the output array is one larger than the input; the initial element will - -- be written into this spot by thread 0 of the first thread block. - j0 <- case mseed of - Nothing -> return i0 - Just _ -> case dir of - L -> A.add numType i0 (lift 1) - R -> return i0 - - -- If this thread has input, read data and participate in thread-block scan - let valid i = case dir of - L -> A.lt singleType i sz - R -> A.gte singleType i (lift 0) - - when (valid i0) $ do - x0 <- app1 delayedLinearIndex i0 - x1 <- case mseed of - Nothing -> return x0 - Just seed -> - if A.eq singleType tid (lift 0) `A.land` A.eq singleType chunk (lift 0) - then do - z <- seed - case dir of - L -> writeArray arrOut (lift 0 :: IR Int32) z >> app2 combine z x0 - R -> writeArray arrOut sz z >> app2 combine x0 z - else - return x0 - - n <- A.sub numType sz inf - n' <- i32 n - x2 <- if A.gte singleType n bd' - then scanBlockSMem dir dev combine Nothing x1 - else scanBlockSMem dir dev combine (Just n') x1 - - -- Write this thread's scan result to memory - writeArray arrOut j0 x2 - - -- The last thread also writes its result---the aggregate for this - -- thread block---to the temporary partial sums array. This is only - -- necessary for full blocks in a multi-block scan; the final - -- partially-full tile does not have a successor block. - last <- A.sub numType bd (lift 1) - when (A.gt singleType gd (lift 1) `land` A.eq singleType tid last) $ - case dir of - L -> writeArray arrTmp chunk x2 - R -> do u <- A.sub numType end chunk - v <- A.sub numType u (lift 1) - writeArray arrTmp v x2 - - return_ - - --- Parallel scan, step 2 --- --- A single thread block performs a scan of the per-block aggregates computed in --- step 1. This gives the per-block prefix which must be added to each element --- in step 3. --- -mkScanAllP2 - :: forall aenv e. Elt e - => Direction - -> DeviceProperties -- ^ properties of the target GPU - -> Gamma aenv -- ^ array environment - -> IRFun2 PTX aenv (e -> e -> e) -- ^ combination function - -> CodeGen (IROpenAcc PTX aenv (Vector e)) -mkScanAllP2 dir dev aenv combine = - let - (start, end, paramGang) = gangParam - (arrTmp, paramTmp) = mutableArray ("tmp" :: Name (Vector e)) - paramEnv = envParam aenv - -- - config = launchConfig dev (CUDA.incWarp dev) smem grid gridQ - grid _ _ = 1 - gridQ = [|| \_ _ -> 1 ||] - smem n = warps * (1 + per_warp) * bytes - where - ws = CUDA.warpSize dev - warps = n `P.quot` ws - per_warp = ws + ws `P.quot` 2 - bytes = sizeOf (eltType (undefined :: e)) - in - makeOpenAccWith config "scanP2" (paramGang ++ paramTmp ++ paramEnv) $ do - - -- The first and last threads of the block need to communicate the - -- block-wide aggregate as a carry-in value across iterations. - -- - -- TODO: We could optimise this a bit if we can get access to the shared - -- memory area used by 'scanBlockSMem', and from there directly read the - -- value computed by the last thread. - carry <- staticSharedMem 1 - - bd <- blockDim - bd' <- int bd - imapFromStepTo start bd' end $ \offset -> do - - -- Index of the partial sums array that this thread will process. - tid <- threadIdx - tid' <- int tid - i0 <- case dir of - L -> A.add numType offset tid' - R -> do x <- A.sub numType end offset - y <- A.sub numType x tid' - z <- A.sub numType y (lift 1) - return z - - let valid i = case dir of - L -> A.lt singleType i end - R -> A.gte singleType i start - - when (valid i0) $ do - - __syncthreads - - x0 <- readArray arrTmp i0 - x1 <- if A.gt singleType offset (lift 0) `land` A.eq singleType tid (lift 0) - then do - c <- readArray carry (lift 0 :: IR Int32) - case dir of - L -> app2 combine c x0 - R -> app2 combine x0 c - else do - return x0 - - n <- A.sub numType end offset - n' <- i32 n - x2 <- if A.gte singleType n bd' - then scanBlockSMem dir dev combine Nothing x1 - else scanBlockSMem dir dev combine (Just n') x1 - - -- Update the temporary array with this thread's result - writeArray arrTmp i0 x2 - - -- The last thread writes the carry-out value. If the last thread is not - -- active, then this must be the last stripe anyway. - last <- A.sub numType bd (lift 1) - when (A.eq singleType tid last) $ - writeArray carry (lift 0 :: IR Int32) x2 - - return_ - - --- Parallel scan, step 3. --- --- Threads combine every element of the partial block results with the carry-in --- value computed in step 2. --- -mkScanAllP3 - :: forall aenv e. Elt e - => Direction - -> DeviceProperties -- ^ properties of the target GPU - -> Gamma aenv -- ^ array environment - -> IRFun2 PTX aenv (e -> e -> e) -- ^ combination function - -> Maybe (IRExp PTX aenv e) -- ^ seed element, if this is an exclusive scan - -> CodeGen (IROpenAcc PTX aenv (Vector e)) -mkScanAllP3 dir dev aenv combine mseed = - let - (start, end, paramGang) = gangParam - (arrOut, paramOut) = mutableArray ("out" :: Name (Vector e)) - (arrTmp, paramTmp) = mutableArray ("tmp" :: Name (Vector e)) - paramEnv = envParam aenv - -- - stride = local scalarType ("ix.stride" :: Name Int) - paramStride = scalarParameter scalarType ("ix.stride" :: Name Int) - -- - config = launchConfig dev (CUDA.incWarp dev) (const 0) const [|| const ||] - in - makeOpenAccWith config "scanP3" (paramGang ++ paramTmp ++ paramOut ++ paramStride : paramEnv) $ do - - sz <- return $ indexHead (irArrayShape arrOut) - tid <- int =<< threadIdx - - -- Threads that will never contribute can just exit immediately. The size of - -- each chunk is set by the block dimension of the step 1 kernel, which may - -- be different from the block size of this kernel. - when (A.lt singleType tid stride) $ do - - -- Iterate over the segments computed in phase 1. Note that we have one - -- fewer chunk to process because the first has no carry-in. - bid <- int =<< blockIdx - gd <- int =<< gridDim - c0 <- A.add numType start bid - imapFromStepTo c0 gd end $ \chunk -> do - - -- Determine the start and end indicies of this chunk to which we will - -- carry-in the value. Returned for left-to-right traversal. - (inf,sup) <- case dir of - L -> do - a <- A.add numType chunk (lift 1) - b <- A.mul numType stride a - case mseed of - Just{} -> do - c <- A.add numType b (lift 1) - d <- A.add numType c stride - e <- A.min singleType d sz - return (c,e) - Nothing -> do - c <- A.add numType b stride - d <- A.min singleType c sz - return (b,d) - R -> do - a <- A.sub numType end chunk - b <- A.mul numType stride a - c <- A.sub numType sz b - case mseed of - Just{} -> do - d <- A.sub numType c (lift 1) - e <- A.sub numType d stride - f <- A.max singleType e (lift 0) - return (f,d) - Nothing -> do - d <- A.sub numType c stride - e <- A.max singleType d (lift 0) - return (e,c) - - -- Read the carry-in value - carry <- case dir of - L -> readArray arrTmp chunk - R -> do - a <- A.add numType chunk (lift 1) - b <- readArray arrTmp a - return b - - -- Apply the carry-in value to each element in the chunk - bd <- int =<< blockDim - i0 <- A.add numType inf tid - imapFromStepTo i0 bd sup $ \i -> do - v <- readArray arrOut i - u <- case dir of - L -> app2 combine carry v - R -> app2 combine v carry - writeArray arrOut i u - - return_ - - --- Parallel scan', step 1. --- --- Similar to mkScanAllP1. Threads scan a stripe of the input into a temporary --- array, incorporating the initial element and any fused functions on the way. --- The final reduction result of this chunk is written to a separate array. --- -mkScan'AllP1 - :: forall aenv e. Elt e - => Direction - -> DeviceProperties - -> Gamma aenv - -> IRFun2 PTX aenv (e -> e -> e) - -> IRExp PTX aenv e - -> IRDelayed PTX aenv (Vector e) - -> CodeGen (IROpenAcc PTX aenv (Vector e, Scalar e)) -mkScan'AllP1 dir dev aenv combine seed IRDelayed{..} = - let - (start, end, paramGang) = gangParam - (arrOut, paramOut) = mutableArray ("out" :: Name (Vector e)) - (arrTmp, paramTmp) = mutableArray ("tmp" :: Name (Vector e)) - paramEnv = envParam aenv - -- - config = launchConfig dev (CUDA.incWarp dev) smem const [|| const ||] - smem n = warps * (1 + per_warp) * bytes - where - ws = CUDA.warpSize dev - warps = n `P.quot` ws - per_warp = ws + ws `P.quot` 2 - bytes = sizeOf (eltType (undefined :: e)) - in - makeOpenAccWith config "scanP1" (paramGang ++ paramTmp ++ paramOut ++ paramEnv) $ do - - -- Size of the input array - sz <- indexHead <$> delayedExtent - - -- A thread block scans a non-empty stripe of the input, storing the partial - -- result and the final block-wide aggregate - bid <- int =<< blockIdx - gd <- int =<< gridDim - s0 <- A.add numType start bid - - -- iterate over thread-block wide segments - imapFromStepTo s0 gd end $ \seg -> do - - bd <- int =<< blockDim - inf <- A.mul numType seg bd - - -- i* is the index that this thread will read data from - tid <- int =<< threadIdx - i0 <- case dir of - L -> A.add numType inf tid - R -> do x <- A.sub numType sz inf - y <- A.sub numType x tid - z <- A.sub numType y (lift 1) - return z - - -- j* is the index this thread will write to. This is just shifted by one - -- to make room for the initial element - j0 <- case dir of - L -> A.add numType i0 (lift 1) - R -> A.sub numType i0 (lift 1) - - -- If this thread has input it participates in the scan - let valid i = case dir of - L -> A.lt singleType i sz - R -> A.gte singleType i (lift 0) - - when (valid i0) $ do - x0 <- app1 delayedLinearIndex i0 - - -- Thread 0 of the first segment must also evaluate and store the - -- initial element - ti <- threadIdx - x1 <- if A.eq singleType ti (lift 0) `A.land` A.eq singleType seg (lift 0) - then do - z <- seed - writeArray arrOut i0 z - case dir of - L -> app2 combine z x0 - R -> app2 combine x0 z - else - return x0 - - -- Block-wide scan - n <- A.sub numType sz inf - n' <- i32 n - x2 <- if A.gte singleType n bd - then scanBlockSMem dir dev combine Nothing x1 - else scanBlockSMem dir dev combine (Just n') x1 - - -- Write this thread's scan result to memory. Recall that we had to make - -- space for the initial element, so the very last thread does not store - -- its result here. - case dir of - L -> when (A.lt singleType j0 sz) $ writeArray arrOut j0 x2 - R -> when (A.gte singleType j0 (lift 0)) $ writeArray arrOut j0 x2 - - -- Last active thread writes its result to the partial sums array. These - -- will be used to compute the carry-in value in step 2. - m <- do x <- A.min singleType n bd - y <- A.sub numType x (lift 1) - return y - when (A.eq singleType tid m) $ - case dir of - L -> writeArray arrTmp seg x2 - R -> do x <- A.sub numType end seg - y <- A.sub numType x (lift 1) - writeArray arrTmp y x2 - - return_ - - --- Parallel scan', step 2 --- --- A single thread block performs an inclusive scan of the partial sums array to --- compute the per-block carry-in values, as well as the final reduction result. --- -mkScan'AllP2 - :: forall aenv e. Elt e - => Direction - -> DeviceProperties - -> Gamma aenv - -> IRFun2 PTX aenv (e -> e -> e) - -> CodeGen (IROpenAcc PTX aenv (Vector e, Scalar e)) -mkScan'AllP2 dir dev aenv combine = - let - (start, end, paramGang) = gangParam - (arrTmp, paramTmp) = mutableArray ("tmp" :: Name (Vector e)) - (arrSum, paramSum) = mutableArray ("sum" :: Name (Scalar e)) - paramEnv = envParam aenv - -- - config = launchConfig dev (CUDA.incWarp dev) smem grid gridQ - grid _ _ = 1 - gridQ = [|| \_ _ -> 1 ||] - smem n = warps * (1 + per_warp) * bytes - where - ws = CUDA.warpSize dev - warps = n `P.quot` ws - per_warp = ws + ws `P.quot` 2 - bytes = sizeOf (eltType (undefined :: e)) - in - makeOpenAccWith config "scanP2" (paramGang ++ paramTmp ++ paramSum ++ paramEnv) $ do - - -- The first and last threads of the block need to communicate the - -- block-wide aggregate as a carry-in value across iterations. - carry <- staticSharedMem 1 - - -- A single thread block iterates over the per-block partial results from - -- step 1 - tid <- threadIdx - tid' <- int tid - bd <- int =<< blockDim - imapFromStepTo start bd end $ \offset -> do - - i0 <- case dir of - L -> A.add numType offset tid' - R -> do x <- A.sub numType end offset - y <- A.sub numType x tid' - z <- A.sub numType y (lift 1) - return z - - let valid i = case dir of - L -> A.lt singleType i end - R -> A.gte singleType i start - - when (valid i0) $ do - - -- wait for the carry-in value to be updated - __syncthreads - - x0 <- readArray arrTmp i0 - x1 <- if A.gt singleType offset (lift 0) `A.land` A.eq singleType tid (lift 0) - then do - c <- readArray carry (lift 0 :: IR Int32) - case dir of - L -> app2 combine c x0 - R -> app2 combine x0 c - else - return x0 - - n <- A.sub numType end offset - n' <- i32 n - x2 <- if A.gte singleType n bd - then scanBlockSMem dir dev combine Nothing x1 - else scanBlockSMem dir dev combine (Just n') x1 - - -- Update the partial results array - writeArray arrTmp i0 x2 - - -- The last active thread saves its result as the carry-out value. - m <- do x <- A.min singleType bd n - y <- A.sub numType x (lift 1) - z <- i32 y - return z - when (A.eq singleType tid m) $ - writeArray carry (lift 0 :: IR Int32) x2 - - -- First thread stores the final carry-out values at the final reduction - -- result for the entire array - __syncthreads - - when (A.eq singleType tid (lift 0)) $ - writeArray arrSum (lift 0 :: IR Int32) =<< readArray carry (lift 0 :: IR Int32) - - return_ - - --- Parallel scan', step 3. --- --- Threads combine every element of the partial block results with the carry-in --- value computed in step 2. --- -mkScan'AllP3 - :: forall aenv e. Elt e - => Direction - -> DeviceProperties -- ^ properties of the target GPU - -> Gamma aenv -- ^ array environment - -> IRFun2 PTX aenv (e -> e -> e) -- ^ combination function - -> CodeGen (IROpenAcc PTX aenv (Vector e, Scalar e)) -mkScan'AllP3 dir dev aenv combine = - let - (start, end, paramGang) = gangParam - (arrOut, paramOut) = mutableArray ("out" :: Name (Vector e)) - (arrTmp, paramTmp) = mutableArray ("tmp" :: Name (Vector e)) - paramEnv = envParam aenv - -- - stride = local scalarType ("ix.stride" :: Name Int) - paramStride = scalarParameter scalarType ("ix.stride" :: Name Int) - -- - config = launchConfig dev (CUDA.incWarp dev) (const 0) const [|| const ||] - in - makeOpenAccWith config "scanP3" (paramGang ++ paramTmp ++ paramOut ++ paramStride : paramEnv) $ do - - sz <- return $ indexHead (irArrayShape arrOut) - tid <- int =<< threadIdx - - when (A.lt singleType tid stride) $ do - - bid <- int =<< blockIdx - gd <- int =<< gridDim - c0 <- A.add numType start bid - imapFromStepTo c0 gd end $ \chunk -> do - - (inf,sup) <- case dir of - L -> do - a <- A.add numType chunk (lift 1) - b <- A.mul numType stride a - c <- A.add numType b (lift 1) - d <- A.add numType c stride - e <- A.min singleType d sz - return (c,e) - R -> do - a <- A.sub numType end chunk - b <- A.mul numType stride a - c <- A.sub numType sz b - d <- A.sub numType c (lift 1) - e <- A.sub numType d stride - f <- A.max singleType e (lift 0) - return (f,d) - - carry <- case dir of - L -> readArray arrTmp chunk - R -> do - a <- A.add numType chunk (lift 1) - b <- readArray arrTmp a - return b - - -- Apply the carry-in value to each element in the chunk - bd <- int =<< blockDim - i0 <- A.add numType inf tid - imapFromStepTo i0 bd sup $ \i -> do - v <- readArray arrOut i - u <- case dir of - L -> app2 combine carry v - R -> app2 combine v carry - writeArray arrOut i u - - return_ - - --- Multidimensional scans --- ---------------------- - --- Multidimensional scan along the innermost dimension --- --- A thread block individually computes along each innermost dimension. This is --- a single-pass operation. --- --- * We can assume that the array is non-empty; exclusive scans with empty --- innermost dimension will be instead filled with the seed element via --- 'mkScanFill'. --- --- * Small but non-empty innermost dimension arrays (size << thread --- block size) will have many threads which do no work. --- -mkScanDim - :: forall aenv sh e. (Shape sh, Elt e) - => Direction - -> DeviceProperties -- ^ properties of the target GPU - -> Gamma aenv -- ^ array environment - -> IRFun2 PTX aenv (e -> e -> e) -- ^ combination function - -> Maybe (IRExp PTX aenv e) -- ^ seed element, if this is an exclusive scan - -> IRDelayed PTX aenv (Array (sh:.Int) e) -- ^ input data - -> CodeGen (IROpenAcc PTX aenv (Array (sh:.Int) e)) -mkScanDim dir dev aenv combine mseed IRDelayed{..} = - let - (start, end, paramGang) = gangParam - (arrOut, paramOut) = mutableArray ("out" :: Name (Array (sh:.Int) e)) - paramEnv = envParam aenv - -- - config = launchConfig dev (CUDA.incWarp dev) smem const [|| const ||] - smem n = warps * (1 + per_warp) * bytes - where - ws = CUDA.warpSize dev - warps = n `P.quot` ws - per_warp = ws + ws `P.quot` 2 - bytes = sizeOf (eltType (undefined :: e)) - in - makeOpenAccWith config "scan" (paramGang ++ paramOut ++ paramEnv) $ do - - -- The first and last threads of the block need to communicate the - -- block-wide aggregate as a carry-in value across iterations. - -- - -- TODO: we could optimise this a bit if we can get access to the shared - -- memory area used by 'scanBlockSMem', and from there directly read the - -- value computed by the last thread. - carry <- staticSharedMem 1 - - -- Size of the input array - sz <- indexHead <$> delayedExtent - - -- Thread blocks iterate over the outer dimensions. Threads in a block - -- cooperatively scan along one dimension, but thread blocks do not - -- communicate with each other. - -- - bid <- int =<< blockIdx - gd <- int =<< gridDim - s0 <- A.add numType start bid - imapFromStepTo s0 gd end $ \seg -> do - - -- Index this thread reads from - tid <- threadIdx - tid' <- int tid - i0 <- case dir of - L -> do x <- A.mul numType seg sz - y <- A.add numType x tid' - return y - - R -> do x <- A.add numType seg (lift 1) - y <- A.mul numType x sz - z <- A.sub numType y tid' - w <- A.sub numType z (lift 1) - return w - - -- Index this thread writes to - j0 <- case mseed of - Nothing -> return i0 - Just{} -> do szp1 <- return $ indexHead (irArrayShape arrOut) - case dir of - L -> do x <- A.mul numType seg szp1 - y <- A.add numType x tid' - return y - - R -> do x <- A.add numType seg (lift 1) - y <- A.mul numType x szp1 - z <- A.sub numType y tid' - w <- A.sub numType z (lift 1) - return w - - -- Stride indices by block dimension - bd <- blockDim - bd' <- int bd - let next ix = case dir of - L -> A.add numType ix bd' - R -> A.sub numType ix bd' - - -- Initialise this scan segment - -- - -- If this is an exclusive scan then the first thread just evaluates the - -- seed element and stores this value into the carry-in slot. All threads - -- shift their write-to index (j) by one, to make space for this element. - -- - -- If this is an inclusive scan then do a block-wide scan. The last thread - -- in the block writes the carry-in value. - -- - r <- - case mseed of - Just seed -> do - when (A.eq singleType tid (lift 0)) $ do - z <- seed - writeArray arrOut j0 z - writeArray carry (lift 0 :: IR Int32) z - j1 <- case dir of - L -> A.add numType j0 (lift 1) - R -> A.sub numType j0 (lift 1) - return $ A.trip sz i0 j1 - - Nothing -> do - when (A.lt singleType tid' sz) $ do - n' <- i32 sz - x0 <- app1 delayedLinearIndex i0 - r0 <- if A.gte singleType sz bd' - then scanBlockSMem dir dev combine Nothing x0 - else scanBlockSMem dir dev combine (Just n') x0 - writeArray arrOut j0 r0 - - ll <- A.sub numType bd (lift 1) - when (A.eq singleType tid ll) $ - writeArray carry (lift 0 :: IR Int32) r0 - - n1 <- A.sub numType sz bd' - i1 <- next i0 - j1 <- next j0 - return $ A.trip n1 i1 j1 - - -- Iterate over the remaining elements in this segment - void $ while - (\(A.fst3 -> n) -> A.gt singleType n (lift 0)) - (\(A.untrip -> (n,i,j)) -> do - - -- Wait for the carry-in value from the previous iteration to be updated - __syncthreads - - -- Compute and store the next element of the scan - -- - -- NOTE: As with 'foldSeg' we require all threads to participate in - -- every iteration of the loop otherwise they will die prematurely. - -- Out-of-bounds threads return 'undef' at this point, which is really - -- unfortunate ): - -- - x <- if A.lt singleType tid' n - then app1 delayedLinearIndex i - else let - go :: TupleType a -> Operands a - go TypeRunit = OP_Unit - go (TypeRpair a b) = OP_Pair (go a) (go b) - go (TypeRscalar t) = ir' t (undef t) - in - return . IR $ go (eltType (undefined::e)) - - -- Thread zero incorporates the carry-in element - y <- if A.eq singleType tid (lift 0) - then do - c <- readArray carry (lift 0 :: IR Int32) - case dir of - L -> app2 combine c x - R -> app2 combine x c - else - return x - - -- Perform the scan and write the result to memory - m <- i32 n - z <- if A.gte singleType n bd' - then scanBlockSMem dir dev combine Nothing y - else scanBlockSMem dir dev combine (Just m) y - - when (A.lt singleType tid' n) $ do - writeArray arrOut j z - - -- The last thread of the block writes its result as the carry-out - -- value. If this thread is not active then we are on the last - -- iteration of the loop and it will not be needed. - w <- A.sub numType bd (lift 1) - when (A.eq singleType tid w) $ - writeArray carry (lift 0 :: IR Int32) z - - -- Update indices for the next iteration - n' <- A.sub numType n bd' - i' <- next i - j' <- next j - return $ A.trip n' i' j') - r - - return_ - - --- Multidimensional scan' along the innermost dimension --- --- A thread block individually computes along each innermost dimension. This is --- a single-pass operation. --- --- * We can assume that the array is non-empty; exclusive scans with empty --- innermost dimension will be instead filled with the seed element via --- 'mkScan'Fill'. --- --- * Small but non-empty innermost dimension arrays (size << thread --- block size) will have many threads which do no work. --- -mkScan'Dim - :: forall aenv sh e. (Shape sh, Elt e) - => Direction - -> DeviceProperties -- ^ properties of the target GPU - -> Gamma aenv -- ^ array environment - -> IRFun2 PTX aenv (e -> e -> e) -- ^ combination function - -> IRExp PTX aenv e -- ^ seed element - -> IRDelayed PTX aenv (Array (sh:.Int) e) -- ^ input data - -> CodeGen (IROpenAcc PTX aenv (Array (sh:.Int) e, Array sh e)) -mkScan'Dim dir dev aenv combine seed IRDelayed{..} = - let - (start, end, paramGang) = gangParam - (arrOut, paramOut) = mutableArray ("out" :: Name (Array (sh:.Int) e)) - (arrSum, paramSum) = mutableArray ("sum" :: Name (Array sh e)) - paramEnv = envParam aenv - -- - config = launchConfig dev (CUDA.incWarp dev) smem const [|| const ||] - smem n = warps * (1 + per_warp) * bytes - where - ws = CUDA.warpSize dev - warps = n `P.quot` ws - per_warp = ws + ws `P.quot` 2 - bytes = sizeOf (eltType (undefined :: e)) - in - makeOpenAccWith config "scan" (paramGang ++ paramOut ++ paramSum ++ paramEnv) $ do - - -- The first and last threads of the block need to communicate the - -- block-wide aggregate as a carry-in value across iterations. - -- - -- TODO: we could optimise this a bit if we can get access to the shared - -- memory area used by 'scanBlockSMem', and from there directly read the - -- value computed by the last thread. - carry <- staticSharedMem 1 - - -- Size of the input array - sz <- indexHead <$> delayedExtent - - -- If the innermost dimension is smaller than the number of threads in the - -- block, those threads will never contribute to the output. - tid <- threadIdx - tid' <- int tid - when (A.lte singleType tid' sz) $ do - - -- Thread blocks iterate over the outer dimensions, each thread block - -- cooperatively scanning along each outermost index. - bid <- int =<< blockIdx - gd <- int =<< gridDim - s0 <- A.add numType start bid - imapFromStepTo s0 gd end $ \seg -> do - - -- Not necessary to wait for threads to catch up before starting this segment - -- __syncthreads - - -- Linear index bounds for this segment - inf <- A.mul numType seg sz - sup <- A.add numType inf sz - - -- Index that this thread will read from. Recall that the supremum index - -- is exclusive. - i0 <- case dir of - L -> A.add numType inf tid' - R -> do x <- A.sub numType sup tid' - y <- A.sub numType x (lift 1) - return y - - -- The index that this thread will write to. This is just shifted along - -- by one to make room for the initial element. - j0 <- case dir of - L -> A.add numType i0 (lift 1) - R -> A.sub numType i0 (lift 1) - - -- Evaluate the initial element. Store it into the carry-in slot as well - -- as to the array as the first element. This is always valid because if - -- the input array is empty then we will be evaluating via mkScan'Fill. - when (A.eq singleType tid (lift 0)) $ do - z <- seed - writeArray arrOut i0 z - writeArray carry (lift 0 :: IR Int32) z - - bd <- blockDim - bd' <- int bd - let next ix = case dir of - L -> A.add numType ix bd' - R -> A.sub numType ix bd' - - -- Now, threads iterate over the elements along the innermost dimension. - -- At each iteration the first thread incorporates the carry-in value - -- from the previous step. - -- - -- The index tracks how many elements remain for the thread block, since - -- indices i* and j* are local to each thread - n0 <- A.sub numType sup inf - void $ while - (\(A.fst3 -> n) -> A.gt singleType n (lift 0)) - (\(A.untrip -> (n,i,j)) -> do - - -- Wait for threads to catch up to ensure the carry-in value from - -- the last iteration has been updated - __syncthreads - - -- If all threads in the block will participate this round we can - -- avoid (almost) all bounds checks. - _ <- if A.gte singleType n bd' - -- All threads participate. No bounds checks required but - -- the last thread needs to update the carry-in value. - then do - x <- app1 delayedLinearIndex i - y <- if A.eq singleType tid (lift 0) - then do - c <- readArray carry (lift 0 :: IR Int32) - case dir of - L -> app2 combine c x - R -> app2 combine x c - else - return x - z <- scanBlockSMem dir dev combine Nothing y - - -- Write results to the output array. Note that if we - -- align directly on the boundary of the array this is not - -- valid for the last thread. - case dir of - L -> when (A.lt singleType j sup) $ writeArray arrOut j z - R -> when (A.gte singleType j inf) $ writeArray arrOut j z - - -- Last thread of the block also saves its result as the - -- carry-in value - bd1 <- A.sub numType bd (lift 1) - when (A.eq singleType tid bd1) $ - writeArray carry (lift 0 :: IR Int32) z - - return (IR OP_Unit :: IR ()) - - -- Only threads that are in bounds can participate. This is - -- the last iteration of the loop. The last active thread - -- still needs to store its value into the carry-in slot. - else do - when (A.lt singleType tid' n) $ do - x <- app1 delayedLinearIndex i - y <- if A.eq singleType tid (lift 0) - then do - c <- readArray carry (lift 0 :: IR Int32) - case dir of - L -> app2 combine c x - R -> app2 combine x c - else - return x - l <- i32 n - z <- scanBlockSMem dir dev combine (Just l) y - - m <- A.sub numType n (lift 1) - _ <- if A.lt singleType tid' m - then writeArray arrOut j z >> return (IR OP_Unit :: IR ()) - else writeArray carry (lift 0 :: IR Int32) z >> return (IR OP_Unit :: IR ()) - - return () - return (IR OP_Unit :: IR ()) - - A.trip <$> A.sub numType n bd' <*> next i <*> next j) - (A.trip n0 i0 j0) - - -- Wait for the carry-in value to be updated - __syncthreads - - -- Store the carry-in value to the separate final results array - when (A.eq singleType tid (lift 0)) $ - writeArray arrSum seg =<< readArray carry (lift 0 :: IR Int32) - - return_ - - - --- Parallel scan, auxiliary --- --- If this is an exclusive scan of an empty array, we just fill the result with --- the seed element. --- -mkScanFill - :: (Shape sh, Elt e) - => PTX - -> Gamma aenv - -> IRExp PTX aenv e - -> CodeGen (IROpenAcc PTX aenv (Array sh e)) -mkScanFill ptx aenv seed = - mkGenerate ptx aenv (IRFun1 (const seed)) - -mkScan'Fill - :: forall aenv sh e. (Shape sh, Elt e) - => PTX - -> Gamma aenv - -> IRExp PTX aenv e - -> CodeGen (IROpenAcc PTX aenv (Array (sh:.Int) e, Array sh e)) -mkScan'Fill ptx aenv seed = - Safe.coerce <$> (mkGenerate ptx aenv (IRFun1 (const seed)) :: CodeGen (IROpenAcc PTX aenv (Array sh e))) - - --- Block wide scan --- --------------- - --- Efficient block-wide (inclusive) scan using the specified operator. --- --- Each block requires (#warps * (1 + 1.5*warp size)) elements of dynamically --- allocated shared memory. --- --- Example: https://github.com/NVlabs/cub/blob/1.5.4/cub/block/specializations/block_scan_warp_scans.cuh --- -scanBlockSMem - :: forall aenv e. Elt e - => Direction - -> DeviceProperties -- ^ properties of the target device - -> IRFun2 PTX aenv (e -> e -> e) -- ^ combination function - -> Maybe (IR Int32) -- ^ number of valid elements (may be less than block size) - -> IR e -- ^ calling thread's input element - -> CodeGen (IR e) -scanBlockSMem dir dev combine nelem = warpScan >=> warpPrefix - where - int32 :: Integral a => a -> IR Int32 - int32 = lift . P.fromIntegral - - -- Temporary storage required for each warp - warp_smem_elems = CUDA.warpSize dev + (CUDA.warpSize dev `P.quot` 2) - warp_smem_bytes = warp_smem_elems * sizeOf (eltType (undefined::e)) - - -- Step 1: Scan in every warp - warpScan :: IR e -> CodeGen (IR e) - warpScan input = do - -- Allocate (1.5 * warpSize) elements of shared memory for each warp - -- (individually addressable by each warp) - wid <- warpId - skip <- A.mul numType wid (int32 warp_smem_bytes) - smem <- dynamicSharedMem (int32 warp_smem_elems) skip - scanWarpSMem dir dev combine smem input - - -- Step 2: Collect the aggregate results of each warp to compute the prefix - -- values for each warp and combine with the partial result to compute each - -- thread's final value. - warpPrefix :: IR e -> CodeGen (IR e) - warpPrefix input = do - -- Allocate #warps elements of shared memory - bd <- blockDim - warps <- A.quot integralType bd (int32 (CUDA.warpSize dev)) - skip <- A.mul numType warps (int32 warp_smem_bytes) - smem <- dynamicSharedMem warps skip - - -- Share warp aggregates - wid <- warpId - lane <- laneId - when (A.eq singleType lane (int32 (CUDA.warpSize dev - 1))) $ do - writeArray smem wid input - - -- Wait for each warp to finish its local scan and share the aggregate - __syncthreads - - -- Compute the prefix value for this warp and add to the partial result. - -- This step is not required for the first warp, which has no carry-in. - if A.eq singleType wid (lift 0) - then return input - else do - -- Every thread sequentially scans the warp aggregates to compute - -- their prefix value. We do this sequentially, but could also have - -- warp 0 do it cooperatively if we limit thread block sizes to - -- (warp size ^ 2). - steps <- case nelem of - Nothing -> return wid - Just n -> A.min singleType wid =<< A.quot integralType n (int32 (CUDA.warpSize dev)) - - p0 <- readArray smem (lift 0 :: IR Int32) - prefix <- iterFromStepTo (lift 1) (lift 1) steps p0 $ \step x -> do - y <- readArray smem step - case dir of - L -> app2 combine x y - R -> app2 combine y x - - case dir of - L -> app2 combine prefix input - R -> app2 combine input prefix - - --- Warp-wide scan --- -------------- - --- Efficient warp-wide (inclusive) scan using the specified operator. --- --- Each warp requires 48 (1.5 x warp size) elements of shared memory. The --- routine assumes that it is allocated individually per-warp (i.e. can be --- indexed in the range [0, warp size)). --- --- Example: https://github.com/NVlabs/cub/blob/1.5.4/cub/warp/specializations/warp_scan_smem.cuh --- -scanWarpSMem - :: forall aenv e. Elt e - => Direction - -> DeviceProperties -- ^ properties of the target device - -> IRFun2 PTX aenv (e -> e -> e) -- ^ combination function - -> IRArray (Vector e) -- ^ temporary storage array in shared memory (1.5 x warp size elements) - -> IR e -- ^ calling thread's input element - -> CodeGen (IR e) -scanWarpSMem dir dev combine smem = scan 0 - where - log2 :: Double -> Double - log2 = P.logBase 2 - - -- Number of steps required to scan warp - steps = P.floor (log2 (P.fromIntegral (CUDA.warpSize dev))) - halfWarp = P.fromIntegral (CUDA.warpSize dev `P.quot` 2) - - -- Unfold the scan as a recursive code generation function - scan :: Int -> IR e -> CodeGen (IR e) - scan step x - | step >= steps = return x - | offset <- 1 `P.shiftL` step = do - -- share partial result through shared memory buffer - lane <- laneId - i <- A.add numType lane (lift halfWarp) - writeArray smem i x - - -- update partial result if in range - x' <- if A.gte singleType lane (lift offset) - then do - i' <- A.sub numType i (lift offset) -- lane + HALF_WARP - offset - x' <- readArray smem i' - case dir of - L -> app2 combine x' x - R -> app2 combine x x' - - else - return x - - scan (step+1) x' - - --- Utilities --- --------- - -i32 :: IR Int -> CodeGen (IR Int32) -i32 = A.fromIntegral integralType numType - -int :: IR Int32 -> CodeGen (IR Int) +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE ViewPatterns #-} +-- | +-- Module : Data.Array.Accelerate.LLVM.PTX.CodeGen.Scan +-- Copyright : [2016..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.LLVM.PTX.CodeGen.Scan ( + + mkScan, mkScan', + +) where + +import Data.Array.Accelerate.AST ( Direction(..) ) +import Data.Array.Accelerate.Representation.Array +import Data.Array.Accelerate.Representation.Elt +import Data.Array.Accelerate.Representation.Shape +import Data.Array.Accelerate.Representation.Type + +import Data.Array.Accelerate.LLVM.CodeGen.Arithmetic as A +import Data.Array.Accelerate.LLVM.CodeGen.Array +import Data.Array.Accelerate.LLVM.CodeGen.Base +import Data.Array.Accelerate.LLVM.CodeGen.Constant +import Data.Array.Accelerate.LLVM.CodeGen.Environment +import Data.Array.Accelerate.LLVM.CodeGen.Exp +import Data.Array.Accelerate.LLVM.CodeGen.IR +import Data.Array.Accelerate.LLVM.CodeGen.Loop +import Data.Array.Accelerate.LLVM.CodeGen.Monad +import Data.Array.Accelerate.LLVM.CodeGen.Sugar +import Data.Array.Accelerate.LLVM.PTX.Analysis.Launch +import Data.Array.Accelerate.LLVM.PTX.CodeGen.Base +import Data.Array.Accelerate.LLVM.PTX.CodeGen.Generate +import Data.Array.Accelerate.LLVM.PTX.Target + +import LLVM.AST.Type.Representation + +import qualified Foreign.CUDA.Analysis as CUDA + +import Control.Applicative +import Control.Monad ( (>=>), void ) +import Control.Monad.State ( gets ) +import Data.String ( fromString ) +import Data.Coerce as Safe +import Data.Bits as P +import Prelude as P hiding ( last ) + + + +-- 'Data.List.scanl' or 'Data.List.scanl1' style exclusive scan, but with the +-- restriction that the combination function must be associative to enable +-- efficient parallel implementation. +-- +-- > scanl (+) 10 (use $ fromList (Z :. 10) [0..]) +-- > +-- > ==> Array (Z :. 11) [10,10,11,13,16,20,25,31,38,46,55] +-- +mkScan + :: forall aenv sh e. + Gamma aenv + -> ArrayR (Array (sh, Int) e) + -> Direction + -> IRFun2 PTX aenv (e -> e -> e) + -> Maybe (IRExp PTX aenv e) + -> MIRDelayed PTX aenv (Array (sh, Int) e) + -> CodeGen PTX (IROpenAcc PTX aenv (Array (sh, Int) e)) +mkScan aenv repr dir combine seed arr + = foldr1 (+++) <$> sequence (codeScan ++ codeFill) + + where + codeScan = case repr of + ArrayR (ShapeRsnoc ShapeRz) tp -> [ mkScanAllP1 dir aenv tp combine seed arr + , mkScanAllP2 dir aenv tp combine + , mkScanAllP3 dir aenv tp combine seed + ] + _ -> [ mkScanDim dir aenv repr combine seed arr + ] + codeFill = case seed of + Just s -> [ mkScanFill aenv repr s ] + Nothing -> [] + +-- Variant of 'scanl' where the final result is returned in a separate array. +-- +-- > scanr' (+) 10 (use $ fromList (Z :. 10) [0..]) +-- > +-- > ==> ( Array (Z :. 10) [10,10,11,13,16,20,25,31,38,46] +-- , Array Z [55] +-- ) +-- +mkScan' + :: forall aenv sh e. + Gamma aenv + -> ArrayR (Array (sh, Int) e) + -> Direction + -> IRFun2 PTX aenv (e -> e -> e) + -> IRExp PTX aenv e + -> MIRDelayed PTX aenv (Array (sh, Int) e) + -> CodeGen PTX (IROpenAcc PTX aenv (Array (sh, Int) e, Array sh e)) +mkScan' aenv repr dir combine seed arr + | ArrayR (ShapeRsnoc ShapeRz) tp <- repr + = foldr1 (+++) <$> sequence [ mkScan'AllP1 dir aenv tp combine seed arr + , mkScan'AllP2 dir aenv tp combine + , mkScan'AllP3 dir aenv tp combine + , mkScan'Fill aenv repr seed + ] + -- + | otherwise + = (+++) <$> mkScan'Dim dir aenv repr combine seed arr + <*> mkScan'Fill aenv repr seed + + +-- Device wide scans +-- ----------------- +-- +-- This is a classic two-pass algorithm which proceeds in two phases and +-- requires ~4n data movement to global memory. In future we would like to +-- replace this with a single pass algorithm. +-- + +-- Parallel scan, step 1. +-- +-- Threads scan a stripe of the input into a temporary array, incorporating the +-- initial element and any fused functions on the way. The final reduction +-- result of this chunk is written to a separate array. +-- +mkScanAllP1 + :: forall aenv e. + Direction + -> Gamma aenv -- ^ array environment + -> TypeR e + -> IRFun2 PTX aenv (e -> e -> e) -- ^ combination function + -> MIRExp PTX aenv e -- ^ seed element, if this is an exclusive scan + -> MIRDelayed PTX aenv (Vector e) -- ^ input data + -> CodeGen PTX (IROpenAcc PTX aenv (Vector e)) +mkScanAllP1 dir aenv tp combine mseed marr = do + dev <- liftCodeGen $ gets ptxDeviceProperties + -- + let + (arrOut, paramOut) = mutableArray (ArrayR dim1 tp) "out" + (arrTmp, paramTmp) = mutableArray (ArrayR dim1 tp) "tmp" + (arrIn, paramIn) = delayedArray "in" marr + end = indexHead (irArrayShape arrTmp) + paramEnv = envParam aenv + -- + config = launchConfig dev (CUDA.incWarp dev) smem const [|| const ||] + smem n = warps * (1 + per_warp) * bytes + where + ws = CUDA.warpSize dev + warps = n `P.quot` ws + per_warp = ws + ws `P.quot` 2 + bytes = bytesElt tp + -- + makeOpenAccWith config "scanP1" (paramTmp ++ paramOut ++ paramIn ++ paramEnv) $ do + + -- Size of the input array + sz <- indexHead <$> delayedExtent arrIn + + -- A thread block scans a non-empty stripe of the input, storing the final + -- block-wide aggregate into a separate array + -- + -- For exclusive scans, thread 0 of segment 0 must incorporate the initial + -- element into the input and output. Threads shuffle their indices + -- appropriately. + -- + bid <- blockIdx + gd <- gridDim + gd' <- int gd + s0 <- int bid + + -- iterating over thread-block-wide segments + imapFromStepTo s0 gd' end $ \chunk -> do + + bd <- blockDim + bd' <- int bd + inf <- A.mul numType chunk bd' + + -- index i* is the index that this thread will read data from. Recall that + -- the supremum index is exclusive + tid <- threadIdx + tid' <- int tid + i0 <- case dir of + LeftToRight -> A.add numType inf tid' + RightToLeft -> do x <- A.sub numType sz inf + y <- A.sub numType x tid' + z <- A.sub numType y (liftInt 1) + return z + + -- index j* is the index that we write to. Recall that for exclusive scans + -- the output array is one larger than the input; the initial element will + -- be written into this spot by thread 0 of the first thread block. + j0 <- case mseed of + Nothing -> return i0 + Just _ -> case dir of + LeftToRight -> A.add numType i0 (liftInt 1) + RightToLeft -> return i0 + + -- If this thread has input, read data and participate in thread-block scan + let valid i = case dir of + LeftToRight -> A.lt singleType i sz + RightToLeft -> A.gte singleType i (liftInt 0) + + when (valid i0) $ do + x0 <- app1 (delayedLinearIndex arrIn) i0 + x1 <- case mseed of + Nothing -> return x0 + Just seed -> + if (tp, A.eq singleType tid (liftInt32 0) `A.land'` A.eq singleType chunk (liftInt 0)) + then do + z <- seed + case dir of + LeftToRight -> writeArray TypeInt32 arrOut (liftInt32 0) z >> app2 combine z x0 + RightToLeft -> writeArray TypeInt arrOut sz z >> app2 combine x0 z + else + return x0 + + n <- A.sub numType sz inf + n' <- i32 n + x2 <- if (tp, A.gte singleType n bd') + then scanBlockSMem dir dev tp combine Nothing x1 + else scanBlockSMem dir dev tp combine (Just n') x1 + + -- Write this thread's scan result to memory + writeArray TypeInt arrOut j0 x2 + + -- The last thread also writes its result---the aggregate for this + -- thread block---to the temporary partial sums array. This is only + -- necessary for full blocks in a multi-block scan; the final + -- partially-full tile does not have a successor block. + last <- A.sub numType bd (liftInt32 1) + when (A.gt singleType gd (liftInt32 1) `land'` A.eq singleType tid last) $ + case dir of + LeftToRight -> writeArray TypeInt arrTmp chunk x2 + RightToLeft -> do u <- A.sub numType end chunk + v <- A.sub numType u (liftInt 1) + writeArray TypeInt arrTmp v x2 + + return_ + + +-- Parallel scan, step 2 +-- +-- A single thread block performs a scan of the per-block aggregates computed in +-- step 1. This gives the per-block prefix which must be added to each element +-- in step 3. +-- +mkScanAllP2 + :: forall aenv e. + Direction + -> Gamma aenv -- ^ array environment + -> TypeR e + -> IRFun2 PTX aenv (e -> e -> e) -- ^ combination function + -> CodeGen PTX (IROpenAcc PTX aenv (Vector e)) +mkScanAllP2 dir aenv tp combine = do + dev <- liftCodeGen $ gets ptxDeviceProperties + -- + let + (arrTmp, paramTmp) = mutableArray (ArrayR dim1 tp) "tmp" + paramEnv = envParam aenv + start = liftInt 0 + end = indexHead (irArrayShape arrTmp) + -- + config = launchConfig dev (CUDA.incWarp dev) smem grid gridQ + grid _ _ = 1 + gridQ = [|| \_ _ -> 1 ||] + smem n = warps * (1 + per_warp) * bytes + where + ws = CUDA.warpSize dev + warps = n `P.quot` ws + per_warp = ws + ws `P.quot` 2 + bytes = bytesElt tp + -- + makeOpenAccWith config "scanP2" (paramTmp ++ paramEnv) $ do + + -- The first and last threads of the block need to communicate the + -- block-wide aggregate as a carry-in value across iterations. + -- + -- TODO: We could optimise this a bit if we can get access to the shared + -- memory area used by 'scanBlockSMem', and from there directly read the + -- value computed by the last thread. + carry <- staticSharedMem tp 1 + + bd <- blockDim + bd' <- int bd + + imapFromStepTo start bd' end $ \offset -> do + + -- Index of the partial sums array that this thread will process. + tid <- threadIdx + tid' <- int tid + i0 <- case dir of + LeftToRight -> A.add numType offset tid' + RightToLeft -> do x <- A.sub numType end offset + y <- A.sub numType x tid' + z <- A.sub numType y (liftInt 1) + return z + + let valid i = case dir of + LeftToRight -> A.lt singleType i end + RightToLeft -> A.gte singleType i start + + when (valid i0) $ do + + -- wait for the carry-in value to be updated + __syncthreads + + x0 <- readArray TypeInt arrTmp i0 + x1 <- if (tp, A.gt singleType offset (liftInt 0) `land'` A.eq singleType tid (liftInt32 0)) + then do + c <- readArray TypeInt32 carry (liftInt32 0) + case dir of + LeftToRight -> app2 combine c x0 + RightToLeft -> app2 combine x0 c + else do + return x0 + + n <- A.sub numType end offset + n' <- i32 n + x2 <- if (tp, A.gte singleType n bd') + then scanBlockSMem dir dev tp combine Nothing x1 + else scanBlockSMem dir dev tp combine (Just n') x1 + + -- Update the temporary array with this thread's result + writeArray TypeInt arrTmp i0 x2 + + -- The last thread writes the carry-out value. If the last thread is not + -- active, then this must be the last stripe anyway. + last <- A.sub numType bd (liftInt32 1) + when (A.eq singleType tid last) $ + writeArray TypeInt32 carry (liftInt32 0) x2 + + return_ + + +-- Parallel scan, step 3. +-- +-- Threads combine every element of the partial block results with the carry-in +-- value computed in step 2. +-- +mkScanAllP3 + :: forall aenv e. + Direction + -> Gamma aenv -- ^ array environment + -> TypeR e + -> IRFun2 PTX aenv (e -> e -> e) -- ^ combination function + -> MIRExp PTX aenv e -- ^ seed element, if this is an exclusive scan + -> CodeGen PTX (IROpenAcc PTX aenv (Vector e)) +mkScanAllP3 dir aenv tp combine mseed = do + dev <- liftCodeGen $ gets ptxDeviceProperties + -- + let + (arrOut, paramOut) = mutableArray (ArrayR dim1 tp) "out" + (arrTmp, paramTmp) = mutableArray (ArrayR dim1 tp) "tmp" + paramEnv = envParam aenv + -- + stride = local (TupRsingle scalarTypeInt) "ix.stride" + paramStride = parameter (TupRsingle scalarTypeInt) "ix.stride" + -- + config = launchConfig dev (CUDA.incWarp dev) (const 0) const [|| const ||] + -- + makeOpenAccWith config "scanP3" (paramTmp ++ paramOut ++ paramStride ++ paramEnv) $ do + + sz <- return $ indexHead (irArrayShape arrOut) + tid <- int =<< threadIdx + + -- Threads that will never contribute can just exit immediately. The size of + -- each chunk is set by the block dimension of the step 1 kernel, which may + -- be different from the block size of this kernel. + when (A.lt singleType tid stride) $ do + + -- Iterate over the segments computed in phase 1. Note that we have one + -- fewer chunk to process because the first has no carry-in. + bid <- int =<< blockIdx + gd <- int =<< gridDim + end <- A.sub numType (indexHead (irArrayShape arrTmp)) (liftInt 1) + + imapFromStepTo bid gd end $ \chunk -> do + + -- Determine the start and end indicies of this chunk to which we will + -- carry-in the value. Returned for left-to-right traversal. + (inf,sup) <- case dir of + LeftToRight -> do + a <- A.add numType chunk (liftInt 1) + b <- A.mul numType stride a + case mseed of + Just{} -> do + c <- A.add numType b (liftInt 1) + d <- A.add numType c stride + e <- A.min singleType d sz + return (c,e) + Nothing -> do + c <- A.add numType b stride + d <- A.min singleType c sz + return (b,d) + RightToLeft -> do + a <- A.sub numType end chunk + b <- A.mul numType stride a + c <- A.sub numType sz b + case mseed of + Just{} -> do + d <- A.sub numType c (liftInt 1) + e <- A.sub numType d stride + f <- A.max singleType e (liftInt 0) + return (f,d) + Nothing -> do + d <- A.sub numType c stride + e <- A.max singleType d (liftInt 0) + return (e,c) + + -- Read the carry-in value + carry <- case dir of + LeftToRight -> readArray TypeInt arrTmp chunk + RightToLeft -> do + a <- A.add numType chunk (liftInt 1) + b <- readArray TypeInt arrTmp a + return b + + -- Apply the carry-in value to each element in the chunk + bd <- int =<< blockDim + i0 <- A.add numType inf tid + imapFromStepTo i0 bd sup $ \i -> do + v <- readArray TypeInt arrOut i + u <- case dir of + LeftToRight -> app2 combine carry v + RightToLeft -> app2 combine v carry + writeArray TypeInt arrOut i u + + return_ + + +-- Parallel scan', step 1. +-- +-- Similar to mkScanAllP1. Threads scan a stripe of the input into a temporary +-- array, incorporating the initial element and any fused functions on the way. +-- The final reduction result of this chunk is written to a separate array. +-- +mkScan'AllP1 + :: forall aenv e. + Direction + -> Gamma aenv + -> TypeR e + -> IRFun2 PTX aenv (e -> e -> e) + -> IRExp PTX aenv e + -> MIRDelayed PTX aenv (Vector e) + -> CodeGen PTX (IROpenAcc PTX aenv (Vector e, Scalar e)) +mkScan'AllP1 dir aenv tp combine seed marr = do + dev <- liftCodeGen $ gets ptxDeviceProperties + -- + let + (arrOut, paramOut) = mutableArray (ArrayR dim1 tp) "out" + (arrTmp, paramTmp) = mutableArray (ArrayR dim1 tp) "tmp" + (arrIn, paramIn) = delayedArray "in" marr + end = indexHead (irArrayShape arrTmp) + paramEnv = envParam aenv + -- + config = launchConfig dev (CUDA.incWarp dev) smem const [|| const ||] + smem n = warps * (1 + per_warp) * bytes + where + ws = CUDA.warpSize dev + warps = n `P.quot` ws + per_warp = ws + ws `P.quot` 2 + bytes = bytesElt tp + -- + makeOpenAccWith config "scanP1" (paramTmp ++ paramOut ++ paramIn ++ paramEnv) $ do + + -- Size of the input array + sz <- indexHead <$> delayedExtent arrIn + + -- A thread block scans a non-empty stripe of the input, storing the partial + -- result and the final block-wide aggregate + bid <- int =<< blockIdx + gd <- int =<< gridDim + + -- iterate over thread-block wide segments + imapFromStepTo bid gd end $ \seg -> do + + bd <- int =<< blockDim + inf <- A.mul numType seg bd + + -- i* is the index that this thread will read data from + tid <- int =<< threadIdx + i0 <- case dir of + LeftToRight -> A.add numType inf tid + RightToLeft -> do x <- A.sub numType sz inf + y <- A.sub numType x tid + z <- A.sub numType y (liftInt 1) + return z + + -- j* is the index this thread will write to. This is just shifted by one + -- to make room for the initial element + j0 <- case dir of + LeftToRight -> A.add numType i0 (liftInt 1) + RightToLeft -> A.sub numType i0 (liftInt 1) + + -- If this thread has input it participates in the scan + let valid i = case dir of + LeftToRight -> A.lt singleType i sz + RightToLeft -> A.gte singleType i (liftInt 0) + + when (valid i0) $ do + x0 <- app1 (delayedLinearIndex arrIn) i0 + + -- Thread 0 of the first segment must also evaluate and store the + -- initial element + ti <- threadIdx + x1 <- if (tp, A.eq singleType ti (liftInt32 0) `A.land'` A.eq singleType seg (liftInt 0)) + then do + z <- seed + writeArray TypeInt arrOut i0 z + case dir of + LeftToRight -> app2 combine z x0 + RightToLeft -> app2 combine x0 z + else + return x0 + + -- Block-wide scan + n <- A.sub numType sz inf + n' <- i32 n + x2 <- if (tp, A.gte singleType n bd) + then scanBlockSMem dir dev tp combine Nothing x1 + else scanBlockSMem dir dev tp combine (Just n') x1 + + -- Write this thread's scan result to memory. Recall that we had to make + -- space for the initial element, so the very last thread does not store + -- its result here. + case dir of + LeftToRight -> when (A.lt singleType j0 sz) $ writeArray TypeInt arrOut j0 x2 + RightToLeft -> when (A.gte singleType j0 (liftInt 0)) $ writeArray TypeInt arrOut j0 x2 + + -- Last active thread writes its result to the partial sums array. These + -- will be used to compute the carry-in value in step 2. + m <- do x <- A.min singleType n bd + y <- A.sub numType x (liftInt 1) + return y + when (A.eq singleType tid m) $ + case dir of + LeftToRight -> writeArray TypeInt arrTmp seg x2 + RightToLeft -> do x <- A.sub numType end seg + y <- A.sub numType x (liftInt 1) + writeArray TypeInt arrTmp y x2 + + return_ + + +-- Parallel scan', step 2 +-- +-- A single thread block performs an inclusive scan of the partial sums array to +-- compute the per-block carry-in values, as well as the final reduction result. +-- +mkScan'AllP2 + :: forall aenv e. + Direction + -> Gamma aenv + -> TypeR e + -> IRFun2 PTX aenv (e -> e -> e) + -> CodeGen PTX (IROpenAcc PTX aenv (Vector e, Scalar e)) +mkScan'AllP2 dir aenv tp combine = do + dev <- liftCodeGen $ gets ptxDeviceProperties + -- + let + (arrTmp, paramTmp) = mutableArray (ArrayR dim1 tp) "tmp" + (arrSum, paramSum) = mutableArray (ArrayR dim0 tp) "sum" + paramEnv = envParam aenv + start = liftInt 0 + end = indexHead (irArrayShape arrTmp) + -- + config = launchConfig dev (CUDA.incWarp dev) smem grid gridQ + grid _ _ = 1 + gridQ = [|| \_ _ -> 1 ||] + smem n = warps * (1 + per_warp) * bytes + where + ws = CUDA.warpSize dev + warps = n `P.quot` ws + per_warp = ws + ws `P.quot` 2 + bytes = bytesElt tp + -- + makeOpenAccWith config "scanP2" (paramTmp ++ paramSum ++ paramEnv) $ do + + -- The first and last threads of the block need to communicate the + -- block-wide aggregate as a carry-in value across iterations. + carry <- staticSharedMem tp 1 + + -- A single thread block iterates over the per-block partial results from + -- step 1 + tid <- threadIdx + tid' <- int tid + bd <- int =<< blockDim + + imapFromStepTo start bd end $ \offset -> do + + i0 <- case dir of + LeftToRight -> A.add numType offset tid' + RightToLeft -> do x <- A.sub numType end offset + y <- A.sub numType x tid' + z <- A.sub numType y (liftInt 1) + return z + + let valid i = case dir of + LeftToRight -> A.lt singleType i end + RightToLeft -> A.gte singleType i start + + -- wait for the carry-in value to be updated + __syncthreads + + x0 <- if (tp, valid i0) + then readArray TypeInt arrTmp i0 + else + let go :: TypeR a -> Operands a + go TupRunit = OP_Unit + go (TupRpair a b) = OP_Pair (go a) (go b) + go (TupRsingle t) = ir t (undef t) + in + return $ go tp + + x1 <- if (tp, A.gt singleType offset (liftInt 0) `A.land'` A.eq singleType tid (liftInt32 0)) + then do + c <- readArray TypeInt32 carry (liftInt32 0) + case dir of + LeftToRight -> app2 combine c x0 + RightToLeft -> app2 combine x0 c + else + return x0 + + n <- A.sub numType end offset + n' <- i32 n + x2 <- if (tp, A.gte singleType n bd) + then scanBlockSMem dir dev tp combine Nothing x1 + else scanBlockSMem dir dev tp combine (Just n') x1 + + -- Update the partial results array + when (valid i0) $ + writeArray TypeInt arrTmp i0 x2 + + -- The last active thread saves its result as the carry-out value. + m <- do x <- A.min singleType bd n + y <- A.sub numType x (liftInt 1) + z <- i32 y + return z + when (A.eq singleType tid m) $ + writeArray TypeInt32 carry (liftInt32 0) x2 + + -- First thread stores the final carry-out values at the final reduction + -- result for the entire array + __syncthreads + + when (A.eq singleType tid (liftInt32 0)) $ + writeArray TypeInt32 arrSum (liftInt32 0) =<< readArray TypeInt32 carry (liftInt32 0) + + return_ + + +-- Parallel scan', step 3. +-- +-- Threads combine every element of the partial block results with the carry-in +-- value computed in step 2. +-- +mkScan'AllP3 + :: forall aenv e. + Direction + -> Gamma aenv -- ^ array environment + -> TypeR e + -> IRFun2 PTX aenv (e -> e -> e) -- ^ combination function + -> CodeGen PTX (IROpenAcc PTX aenv (Vector e, Scalar e)) +mkScan'AllP3 dir aenv tp combine = do + dev <- liftCodeGen $ gets ptxDeviceProperties + -- + let + (arrOut, paramOut) = mutableArray (ArrayR dim1 tp) "out" + (arrTmp, paramTmp) = mutableArray (ArrayR dim1 tp) "tmp" + paramEnv = envParam aenv + -- + stride = local (TupRsingle scalarTypeInt) "ix.stride" + paramStride = parameter (TupRsingle scalarTypeInt) "ix.stride" + -- + config = launchConfig dev (CUDA.incWarp dev) (const 0) const [|| const ||] + -- + makeOpenAccWith config "scanP3" (paramTmp ++ paramOut ++ paramStride ++ paramEnv) $ do + + sz <- return $ indexHead (irArrayShape arrOut) + tid <- int =<< threadIdx + + when (A.lt singleType tid stride) $ do + + bid <- int =<< blockIdx + gd <- int =<< gridDim + end <- A.sub numType (indexHead (irArrayShape arrTmp)) (liftInt 1) + + imapFromStepTo bid gd end $ \chunk -> do + + (inf,sup) <- case dir of + LeftToRight -> do + a <- A.add numType chunk (liftInt 1) + b <- A.mul numType stride a + c <- A.add numType b (liftInt 1) + d <- A.add numType c stride + e <- A.min singleType d sz + return (c,e) + RightToLeft -> do + a <- A.sub numType end chunk + b <- A.mul numType stride a + c <- A.sub numType sz b + d <- A.sub numType c (liftInt 1) + e <- A.sub numType d stride + f <- A.max singleType e (liftInt 0) + return (f,d) + + carry <- case dir of + LeftToRight -> readArray TypeInt arrTmp chunk + RightToLeft -> do + a <- A.add numType chunk (liftInt 1) + b <- readArray TypeInt arrTmp a + return b + + -- Apply the carry-in value to each element in the chunk + bd <- int =<< blockDim + i0 <- A.add numType inf tid + imapFromStepTo i0 bd sup $ \i -> do + v <- readArray TypeInt arrOut i + u <- case dir of + LeftToRight -> app2 combine carry v + RightToLeft -> app2 combine v carry + writeArray TypeInt arrOut i u + + return_ + + +-- Multidimensional scans +-- ---------------------- + +-- Multidimensional scan along the innermost dimension +-- +-- A thread block individually computes along each innermost dimension. This is +-- a single-pass operation. +-- +-- * We can assume that the array is non-empty; exclusive scans with empty +-- innermost dimension will be instead filled with the seed element via +-- 'mkScanFill'. +-- +-- * Small but non-empty innermost dimension arrays (size << thread +-- block size) will have many threads which do no work. +-- +mkScanDim + :: forall aenv sh e. + Direction + -> Gamma aenv -- ^ array environment + -> ArrayR (Array (sh, Int) e) + -> IRFun2 PTX aenv (e -> e -> e) -- ^ combination function + -> MIRExp PTX aenv e -- ^ seed element, if this is an exclusive scan + -> MIRDelayed PTX aenv (Array (sh, Int) e) -- ^ input data + -> CodeGen PTX (IROpenAcc PTX aenv (Array (sh, Int) e)) +mkScanDim dir aenv repr@(ArrayR (ShapeRsnoc shr) tp) combine mseed marr = do + dev <- liftCodeGen $ gets ptxDeviceProperties + -- + let + (arrOut, paramOut) = mutableArray repr "out" + (arrIn, paramIn) = delayedArray "in" marr + paramEnv = envParam aenv + -- + config = launchConfig dev (CUDA.incWarp dev) smem const [|| const ||] + smem n = warps * (1 + per_warp) * bytes + where + ws = CUDA.warpSize dev + warps = n `P.quot` ws + per_warp = ws + ws `P.quot` 2 + bytes = bytesElt tp + -- + makeOpenAccWith config "scan" (paramOut ++ paramIn ++ paramEnv) $ do + + -- The first and last threads of the block need to communicate the + -- block-wide aggregate as a carry-in value across iterations. + -- + -- TODO: we could optimise this a bit if we can get access to the shared + -- memory area used by 'scanBlockSMem', and from there directly read the + -- value computed by the last thread. + carry <- staticSharedMem tp 1 + + -- Size of the input array + sz <- indexHead <$> delayedExtent arrIn + + -- Thread blocks iterate over the outer dimensions. Threads in a block + -- cooperatively scan along one dimension, but thread blocks do not + -- communicate with each other. + -- + bid <- int =<< blockIdx + gd <- int =<< gridDim + end <- shapeSize shr (indexTail (irArrayShape arrOut)) + + imapFromStepTo bid gd end $ \seg -> do + + -- Index this thread reads from + tid <- threadIdx + tid' <- int tid + i0 <- case dir of + LeftToRight -> do x <- A.mul numType seg sz + y <- A.add numType x tid' + return y + + RightToLeft -> do x <- A.add numType seg (liftInt 1) + y <- A.mul numType x sz + z <- A.sub numType y tid' + w <- A.sub numType z (liftInt 1) + return w + + -- Index this thread writes to + j0 <- case mseed of + Nothing -> return i0 + Just{} -> do szp1 <- return $ indexHead (irArrayShape arrOut) + case dir of + LeftToRight -> do x <- A.mul numType seg szp1 + y <- A.add numType x tid' + return y + + RightToLeft -> do x <- A.add numType seg (liftInt 1) + y <- A.mul numType x szp1 + z <- A.sub numType y tid' + w <- A.sub numType z (liftInt 1) + return w + + -- Stride indices by block dimension + bd <- blockDim + bd' <- int bd + let next ix = case dir of + LeftToRight -> A.add numType ix bd' + RightToLeft -> A.sub numType ix bd' + + -- Initialise this scan segment + -- + -- If this is an exclusive scan then the first thread just evaluates the + -- seed element and stores this value into the carry-in slot. All threads + -- shift their write-to index (j) by one, to make space for this element. + -- + -- If this is an inclusive scan then do a block-wide scan. The last thread + -- in the block writes the carry-in value. + -- + r <- + case mseed of + Just seed -> do + when (A.eq singleType tid (liftInt32 0)) $ do + z <- seed + writeArray TypeInt arrOut j0 z + writeArray TypeInt32 carry (liftInt32 0) z + j1 <- case dir of + LeftToRight -> A.add numType j0 (liftInt 1) + RightToLeft -> A.sub numType j0 (liftInt 1) + return $ A.trip sz i0 j1 + + Nothing -> do + when (A.lt singleType tid' sz) $ do + n' <- i32 sz + x0 <- app1 (delayedLinearIndex arrIn) i0 + r0 <- if (tp, A.gte singleType sz bd') + then scanBlockSMem dir dev tp combine Nothing x0 + else scanBlockSMem dir dev tp combine (Just n') x0 + writeArray TypeInt arrOut j0 r0 + + ll <- A.sub numType bd (liftInt32 1) + when (A.eq singleType tid ll) $ + writeArray TypeInt32 carry (liftInt32 0) r0 + + n1 <- A.sub numType sz bd' + i1 <- next i0 + j1 <- next j0 + return $ A.trip n1 i1 j1 + + -- Iterate over the remaining elements in this segment + void $ while + (TupRunit `TupRpair` TupRsingle scalarTypeInt `TupRpair` TupRsingle scalarTypeInt `TupRpair` TupRsingle scalarTypeInt) + (\(A.fst3 -> n) -> A.gt singleType n (liftInt 0)) + (\(A.untrip -> (n,i,j)) -> do + + -- Wait for the carry-in value from the previous iteration to be updated + __syncthreads + + -- Compute and store the next element of the scan + -- + -- NOTE: As with 'foldSeg' we require all threads to participate in + -- every iteration of the loop otherwise they will die prematurely. + -- Out-of-bounds threads return 'undef' at this point, which is really + -- unfortunate ): + -- + x <- if (tp, A.lt singleType tid' n) + then app1 (delayedLinearIndex arrIn) i + else let + go :: TypeR a -> Operands a + go TupRunit = OP_Unit + go (TupRpair a b) = OP_Pair (go a) (go b) + go (TupRsingle t) = ir t (undef t) + in + return $ go tp + + -- Thread zero incorporates the carry-in element + y <- if (tp, A.eq singleType tid (liftInt32 0)) + then do + c <- readArray TypeInt32 carry (liftInt32 0) + case dir of + LeftToRight -> app2 combine c x + RightToLeft -> app2 combine x c + else + return x + + -- Perform the scan and write the result to memory + m <- i32 n + z <- if (tp, A.gte singleType n bd') + then scanBlockSMem dir dev tp combine Nothing y + else scanBlockSMem dir dev tp combine (Just m) y + + when (A.lt singleType tid' n) $ do + writeArray TypeInt arrOut j z + + -- The last thread of the block writes its result as the carry-out + -- value. If this thread is not active then we are on the last + -- iteration of the loop and it will not be needed. + w <- A.sub numType bd (liftInt32 1) + when (A.eq singleType tid w) $ + writeArray TypeInt32 carry (liftInt32 0) z + + -- Update indices for the next iteration + n' <- A.sub numType n bd' + i' <- next i + j' <- next j + return $ A.trip n' i' j') + r + + return_ + + +-- Multidimensional scan' along the innermost dimension +-- +-- A thread block individually computes along each innermost dimension. This is +-- a single-pass operation. +-- +-- * We can assume that the array is non-empty; exclusive scans with empty +-- innermost dimension will be instead filled with the seed element via +-- 'mkScan'Fill'. +-- +-- * Small but non-empty innermost dimension arrays (size << thread +-- block size) will have many threads which do no work. +-- +mkScan'Dim + :: forall aenv sh e. + Direction + -> Gamma aenv -- ^ array environment + -> ArrayR (Array (sh, Int) e) + -> IRFun2 PTX aenv (e -> e -> e) -- ^ combination function + -> IRExp PTX aenv e -- ^ seed element + -> MIRDelayed PTX aenv (Array (sh, Int) e) -- ^ input data + -> CodeGen PTX (IROpenAcc PTX aenv (Array (sh, Int) e, Array sh e)) +mkScan'Dim dir aenv repr@(ArrayR (ShapeRsnoc shr) tp) combine seed marr = do + dev <- liftCodeGen $ gets ptxDeviceProperties + -- + let + (arrSum, paramSum) = mutableArray (reduceRank repr) "sum" + (arrOut, paramOut) = mutableArray repr "out" + (arrIn, paramIn) = delayedArray "in" marr + paramEnv = envParam aenv + -- + config = launchConfig dev (CUDA.incWarp dev) smem const [|| const ||] + smem n = warps * (1 + per_warp) * bytes + where + ws = CUDA.warpSize dev + warps = n `P.quot` ws + per_warp = ws + ws `P.quot` 2 + bytes = bytesElt tp + -- + makeOpenAccWith config "scan" (paramOut ++ paramSum ++ paramIn ++ paramEnv) $ do + + -- The first and last threads of the block need to communicate the + -- block-wide aggregate as a carry-in value across iterations. + -- + -- TODO: we could optimise this a bit if we can get access to the shared + -- memory area used by 'scanBlockSMem', and from there directly read the + -- value computed by the last thread. + carry <- staticSharedMem tp 1 + + -- Size of the input array + sz <- indexHead <$> delayedExtent arrIn + + -- If the innermost dimension is smaller than the number of threads in the + -- block, those threads will never contribute to the output. + tid <- threadIdx + tid' <- int tid + when (A.lte singleType tid' sz) $ do + + -- Thread blocks iterate over the outer dimensions, each thread block + -- cooperatively scanning along each outermost index. + bid <- int =<< blockIdx + gd <- int =<< gridDim + end <- shapeSize shr (irArrayShape arrSum) + + imapFromStepTo bid gd end $ \seg -> do + + -- Not necessary to wait for threads to catch up before starting this segment + -- __syncthreads + + -- Linear index bounds for this segment + inf <- A.mul numType seg sz + sup <- A.add numType inf sz + + -- Index that this thread will read from. Recall that the supremum index + -- is exclusive. + i0 <- case dir of + LeftToRight -> A.add numType inf tid' + RightToLeft -> do x <- A.sub numType sup tid' + y <- A.sub numType x (liftInt 1) + return y + + -- The index that this thread will write to. This is just shifted along + -- by one to make room for the initial element. + j0 <- case dir of + LeftToRight -> A.add numType i0 (liftInt 1) + RightToLeft -> A.sub numType i0 (liftInt 1) + + -- Evaluate the initial element. Store it into the carry-in slot as well + -- as to the array as the first element. This is always valid because if + -- the input array is empty then we will be evaluating via mkScan'Fill. + when (A.eq singleType tid (liftInt32 0)) $ do + z <- seed + writeArray TypeInt arrOut i0 z + writeArray TypeInt32 carry (liftInt32 0) z + + bd <- blockDim + bd' <- int bd + let next ix = case dir of + LeftToRight -> A.add numType ix bd' + RightToLeft -> A.sub numType ix bd' + + -- Now, threads iterate over the elements along the innermost dimension. + -- At each iteration the first thread incorporates the carry-in value + -- from the previous step. + -- + -- The index tracks how many elements remain for the thread block, since + -- indices i* and j* are local to each thread + n0 <- A.sub numType sup inf + void $ while + (TupRunit `TupRpair` TupRsingle scalarTypeInt `TupRpair` TupRsingle scalarTypeInt `TupRpair` TupRsingle scalarTypeInt) + (\(A.fst3 -> n) -> A.gt singleType n (liftInt 0)) + (\(A.untrip -> (n,i,j)) -> do + + -- Wait for threads to catch up to ensure the carry-in value from + -- the last iteration has been updated + __syncthreads + + -- If all threads in the block will participate this round we can + -- avoid (almost) all bounds checks. + _ <- if (TupRunit, A.gte singleType n bd') + -- All threads participate. No bounds checks required but + -- the last thread needs to update the carry-in value. + then do + x <- app1 (delayedLinearIndex arrIn) i + y <- if (tp, A.eq singleType tid (liftInt32 0)) + then do + c <- readArray TypeInt32 carry (liftInt32 0) + case dir of + LeftToRight -> app2 combine c x + RightToLeft -> app2 combine x c + else + return x + z <- scanBlockSMem dir dev tp combine Nothing y + + -- Write results to the output array. Note that if we + -- align directly on the boundary of the array this is not + -- valid for the last thread. + case dir of + LeftToRight -> when (A.lt singleType j sup) $ writeArray TypeInt arrOut j z + RightToLeft -> when (A.gte singleType j inf) $ writeArray TypeInt arrOut j z + + -- Last thread of the block also saves its result as the + -- carry-in value + bd1 <- A.sub numType bd (liftInt32 1) + when (A.eq singleType tid bd1) $ + writeArray TypeInt32 carry (liftInt32 0) z + + return (lift TupRunit ()) + + -- Only threads that are in bounds can participate. This is + -- the last iteration of the loop. The last active thread + -- still needs to store its value into the carry-in slot. + -- + -- Note that all threads must call the block-wide scan. + -- SEE: [Synchronisation problems with SM_70 and greater] + else do + x <- if (tp, A.lt singleType tid' n) + then do + x <- app1 (delayedLinearIndex arrIn) i + y <- if (tp, A.eq singleType tid (liftInt32 0)) + then do + c <- readArray TypeInt32 carry (liftInt32 0) + case dir of + LeftToRight -> app2 combine c x + RightToLeft -> app2 combine x c + else + return x + return y + else + let + go :: TypeR a -> Operands a + go TupRunit = OP_Unit + go (TupRpair a b) = OP_Pair (go a) (go b) + go (TupRsingle t) = ir t (undef t) + in + return $ go tp + + l <- i32 n + y <- scanBlockSMem dir dev tp combine (Just l) x + + m <- A.sub numType l (liftInt32 1) + when (A.lt singleType tid m) $ writeArray TypeInt arrOut j y + when (A.eq singleType tid m) $ writeArray TypeInt32 carry (liftInt32 0) y + + return (lift TupRunit ()) + + A.trip <$> A.sub numType n bd' <*> next i <*> next j) + (A.trip n0 i0 j0) + + -- Wait for the carry-in value to be updated + __syncthreads + + -- Store the carry-in value to the separate final results array + when (A.eq singleType tid (liftInt32 0)) $ + writeArray TypeInt arrSum seg =<< readArray TypeInt32 carry (liftInt32 0) + + return_ + + + +-- Parallel scan, auxiliary +-- +-- If this is an exclusive scan of an empty array, we just fill the result with +-- the seed element. +-- +mkScanFill + :: Gamma aenv + -> ArrayR (Array sh e) + -> IRExp PTX aenv e + -> CodeGen PTX (IROpenAcc PTX aenv (Array sh e)) +mkScanFill aenv repr seed = + mkGenerate aenv repr (IRFun1 (const seed)) + +mkScan'Fill + :: Gamma aenv + -> ArrayR (Array (sh, Int) e) + -> IRExp PTX aenv e + -> CodeGen PTX (IROpenAcc PTX aenv (Array (sh, Int) e, Array sh e)) +mkScan'Fill aenv repr seed = + Safe.coerce <$> mkGenerate aenv (reduceRank repr) (IRFun1 (const seed)) + + +-- Block wide scan +-- --------------- + +-- Efficient block-wide (inclusive) scan using the specified operator. +-- +-- Each block requires (#warps * (1 + 1.5*warp size)) elements of dynamically +-- allocated shared memory. +-- +-- Example: https://github.com/NVlabs/cub/blob/1.5.4/cub/block/specializations/block_scan_warp_scans.cuh +-- +-- NOTE: [Synchronisation problems with SM_70 and greater] +-- +-- This operation uses thread synchronisation. When calling this operation, it +-- is important that all active (that is, non-exited) threads of the thread +-- block participate. It seems that sm_70+ (devices with independent thread +-- scheduling) are stricter about the requirement that all non-existed threads +-- participate in every barrier. +-- +-- See: https://github.com/AccelerateHS/accelerate/issues/436 +-- +scanBlockSMem + :: forall aenv e. + Direction + -> DeviceProperties -- ^ properties of the target device + -> TypeR e + -> IRFun2 PTX aenv (e -> e -> e) -- ^ combination function + -> Maybe (Operands Int32) -- ^ number of valid elements (may be less than block size) + -> Operands e -- ^ calling thread's input element + -> CodeGen PTX (Operands e) +scanBlockSMem dir dev tp combine nelem = warpScan >=> warpPrefix + where + int32 :: Integral a => a -> Operands Int32 + int32 = liftInt32 . P.fromIntegral + + -- Temporary storage required for each warp + warp_smem_elems = CUDA.warpSize dev + (CUDA.warpSize dev `P.quot` 2) + warp_smem_bytes = warp_smem_elems * bytesElt tp + + -- Step 1: Scan in every warp + warpScan :: Operands e -> CodeGen PTX (Operands e) + warpScan input = do + -- Allocate (1.5 * warpSize) elements of shared memory for each warp + -- (individually addressable by each warp) + wid <- warpId + skip <- A.mul numType wid (int32 warp_smem_bytes) + smem <- dynamicSharedMem tp TypeInt32 (int32 warp_smem_elems) skip + scanWarpSMem dir dev tp combine smem input + + -- Step 2: Collect the aggregate results of each warp to compute the prefix + -- values for each warp and combine with the partial result to compute each + -- thread's final value. + warpPrefix :: Operands e -> CodeGen PTX (Operands e) + warpPrefix input = do + -- Allocate #warps elements of shared memory + bd <- blockDim + warps <- A.quot integralType bd (int32 (CUDA.warpSize dev)) + skip <- A.mul numType warps (int32 warp_smem_bytes) + smem <- dynamicSharedMem tp TypeInt32 warps skip + + -- Share warp aggregates + wid <- warpId + lane <- laneId + when (A.eq singleType lane (int32 (CUDA.warpSize dev - 1))) $ do + writeArray TypeInt32 smem wid input + + -- Wait for each warp to finish its local scan and share the aggregate + __syncthreads + + -- Compute the prefix value for this warp and add to the partial result. + -- This step is not required for the first warp, which has no carry-in. + if (tp, A.eq singleType wid (liftInt32 0)) + then return input + else do + -- Every thread sequentially scans the warp aggregates to compute + -- their prefix value. We do this sequentially, but could also have + -- warp 0 do it cooperatively if we limit thread block sizes to + -- (warp size ^ 2). + steps <- case nelem of + Nothing -> return wid + Just n -> A.min singleType wid =<< A.quot integralType n (int32 (CUDA.warpSize dev)) + + p0 <- readArray TypeInt32 smem (liftInt32 0) + prefix <- iterFromStepTo tp (liftInt32 1) (liftInt32 1) steps p0 $ \step x -> do + y <- readArray TypeInt32 smem step + case dir of + LeftToRight -> app2 combine x y + RightToLeft -> app2 combine y x + + case dir of + LeftToRight -> app2 combine prefix input + RightToLeft -> app2 combine input prefix + + +-- Warp-wide scan +-- -------------- + +-- Efficient warp-wide (inclusive) scan using the specified operator. +-- +-- Each warp requires 48 (1.5 x warp size) elements of shared memory. The +-- routine assumes that it is allocated individually per-warp (i.e. can be +-- indexed in the range [0, warp size)). +-- +-- Example: https://github.com/NVlabs/cub/blob/1.5.4/cub/warp/specializations/warp_scan_smem.cuh +-- +scanWarpSMem + :: forall aenv e. + Direction + -> DeviceProperties -- ^ properties of the target device + -> TypeR e + -> IRFun2 PTX aenv (e -> e -> e) -- ^ combination function + -> IRArray (Vector e) -- ^ temporary storage array in shared memory (1.5 x warp size elements) + -> Operands e -- ^ calling thread's input element + -> CodeGen PTX (Operands e) +scanWarpSMem dir dev tp combine smem = scan 0 + where + log2 :: Double -> Double + log2 = P.logBase 2 + + -- Number of steps required to scan warp + steps = P.floor (log2 (P.fromIntegral (CUDA.warpSize dev))) + halfWarp = P.fromIntegral (CUDA.warpSize dev `P.quot` 2) + + -- Unfold the scan as a recursive code generation function + scan :: Int -> Operands e -> CodeGen PTX (Operands e) + scan step x + | step >= steps = return x + | otherwise = do + let offset = liftInt32 (1 `P.shiftL` step) + + -- share partial result through shared memory buffer + lane <- laneId + i <- A.add numType lane (liftInt32 halfWarp) + writeArray TypeInt32 smem i x + + __syncwarp + + -- update partial result if in range + x' <- if (tp, A.gte singleType lane offset) + then do + i' <- A.sub numType i offset -- lane + HALF_WARP - offset + x' <- readArray TypeInt32 smem i' + case dir of + LeftToRight -> app2 combine x' x + RightToLeft -> app2 combine x x' + + else + return x + + __syncwarp + + scan (step+1) x' + + +-- Utilities +-- --------- + +i32 :: Operands Int -> CodeGen PTX (Operands Int32) +i32 = A.fromIntegral integralType numType + +int :: Operands Int32 -> CodeGen PTX (Operands Int) int = A.fromIntegral integralType numType diff --git a/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Stencil.hs b/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Stencil.hs new file mode 100644 --- /dev/null +++ b/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Stencil.hs @@ -0,0 +1,187 @@ +{-# LANGUAGE GADTs #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeApplications #-} +-- | +-- Module : Data.Array.Accelerate.LLVM.PTX.CodeGen.Stencil +-- Copyright : [2018..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.LLVM.PTX.CodeGen.Stencil ( + + mkStencil1, + mkStencil2, + +) where + +import Data.Array.Accelerate.Representation.Array +import Data.Array.Accelerate.Representation.Shape +import Data.Array.Accelerate.Representation.Stencil +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Type + +import Data.Array.Accelerate.LLVM.CodeGen.Arithmetic +import Data.Array.Accelerate.LLVM.CodeGen.Array +import Data.Array.Accelerate.LLVM.CodeGen.Base +import Data.Array.Accelerate.LLVM.CodeGen.Environment +import Data.Array.Accelerate.LLVM.CodeGen.Exp +import Data.Array.Accelerate.LLVM.CodeGen.IR +import Data.Array.Accelerate.LLVM.CodeGen.Monad +import Data.Array.Accelerate.LLVM.CodeGen.Stencil +import Data.Array.Accelerate.LLVM.CodeGen.Sugar + +import Data.Array.Accelerate.LLVM.PTX.CodeGen.Base +import Data.Array.Accelerate.LLVM.PTX.CodeGen.Loop +import Data.Array.Accelerate.LLVM.PTX.Target ( PTX ) + +import qualified LLVM.AST.Global as LLVM + +import Control.Monad + + +-- The stencil function is similar to a map, but has access to surrounding +-- elements as specified by the stencil pattern. +-- +-- This generates two functions: +-- +-- * stencil_inside: does not apply boundary conditions, assumes all element +-- accesses are valid +-- +-- * stencil_border: applies boundary condition check to each array access +-- +mkStencil1 + :: Gamma aenv + -> StencilR sh a stencil + -> TypeR b + -> IRFun1 PTX aenv (stencil -> b) + -> IRBoundary PTX aenv (Array sh a) + -> MIRDelayed PTX aenv (Array sh a) + -> CodeGen PTX (IROpenAcc PTX aenv (Array sh b)) +mkStencil1 aenv stencil tp fun bnd marr = + let repr = ArrayR shr tp + (shr, halo) = stencilHalo stencil + (arrIn, paramIn) = delayedArray "in" marr + in + (+++) <$> mkInside aenv repr halo (IRFun1 $ app1 fun <=< stencilAccess stencil Nothing arrIn) paramIn + <*> mkBorder aenv repr (IRFun1 $ app1 fun <=< stencilAccess stencil (Just bnd) arrIn) paramIn + + +mkStencil2 + :: Gamma aenv + -> StencilR sh a stencil1 + -> StencilR sh b stencil2 + -> TypeR c + -> IRFun2 PTX aenv (stencil1 -> stencil2 -> c) + -> IRBoundary PTX aenv (Array sh a) + -> MIRDelayed PTX aenv (Array sh a) + -> IRBoundary PTX aenv (Array sh b) + -> MIRDelayed PTX aenv (Array sh b) + -> CodeGen PTX (IROpenAcc PTX aenv (Array sh c)) +mkStencil2 aenv stencil1 stencil2 tp f bnd1 marr1 bnd2 marr2 = + let + repr = ArrayR shr tp + (arrIn1, paramIn1) = delayedArray "in1" marr1 + (arrIn2, paramIn2) = delayedArray "in2" marr2 + + inside = IRFun1 $ \ix -> do + s1 <- stencilAccess stencil1 Nothing arrIn1 ix + s2 <- stencilAccess stencil2 Nothing arrIn2 ix + app2 f s1 s2 + -- + border = IRFun1 $ \ix -> do + s1 <- stencilAccess stencil1 (Just bnd1) arrIn1 ix + s2 <- stencilAccess stencil2 (Just bnd2) arrIn2 ix + app2 f s1 s2 + + (shr, halo1) = stencilHalo stencil1 + (_, halo2) = stencilHalo stencil2 + halo = union shr halo1 halo2 + in + (+++) <$> mkInside aenv repr halo inside (paramIn1 ++ paramIn2) + <*> mkBorder aenv repr border (paramIn1 ++ paramIn2) + + +mkInside + :: Gamma aenv + -> ArrayR (Array sh e) + -> sh + -> IRFun1 PTX aenv (sh -> e) + -> [LLVM.Parameter] + -> CodeGen PTX (IROpenAcc PTX aenv (Array sh e)) +mkInside aenv repr@(ArrayR shr _) halo apply paramIn = + let + (arrOut, paramOut) = mutableArray repr "out" + paramInside = parameter (shapeType shr) "shInside" + shInside = local (shapeType shr) "shInside" + shOut = irArrayShape arrOut + paramEnv = envParam aenv + -- + in + makeOpenAcc "stencil_inside" (paramInside ++ paramOut ++ paramIn ++ paramEnv) $ do + + start <- return (liftInt 0) + end <- shapeSize shr shInside + + -- iterate over the inside region as a linear index space + -- + imapFromTo start end $ \i -> do + + ixIn <- indexOfInt shr shInside i -- convert to multidimensional index of inside region + ixOut <- offset shr ixIn (lift (shapeType shr) halo) -- shift to multidimensional index of outside region + r <- app1 apply ixOut -- apply generator function + j <- intOfIndex shr shOut ixOut + writeArray TypeInt arrOut j r + + return_ + + +mkBorder + :: Gamma aenv + -> ArrayR (Array sh e) + -> IRFun1 PTX aenv (sh -> e) + -> [LLVM.Parameter] + -> CodeGen PTX (IROpenAcc PTX aenv (Array sh e)) +mkBorder aenv repr@(ArrayR shr _) apply paramIn = + let + (arrOut, paramOut) = mutableArray repr "out" + paramFrom = parameter (shapeType shr) "shFrom" + shFrom = local (shapeType shr) "shFrom" + paramInside = parameter (shapeType shr) "shInside" + shInside = local (shapeType shr) "shInside" + shOut = irArrayShape arrOut + paramEnv = envParam aenv + -- + in + makeOpenAcc "stencil_border" (paramFrom ++ paramInside ++ paramOut ++ paramIn ++ paramEnv) $ do + + start <- return (liftInt 0) + end <- shapeSize shr shInside + + imapFromTo start end $ \i -> do + + ixIn <- indexOfInt shr shInside i -- convert to multidimensional index of inside region + ixOut <- offset shr ixIn shFrom -- shift to multidimensional index of outside region + r <- app1 apply ixOut -- apply generator function + j <- intOfIndex shr shOut ixOut + writeArray TypeInt arrOut j r + + return_ + + +offset :: ShapeR sh -> Operands sh -> Operands sh -> CodeGen PTX (Operands sh) +offset shr sh1 sh2 = go shr sh1 sh2 + where + go :: ShapeR t -> Operands t -> Operands t -> CodeGen PTX (Operands t) + go ShapeRz OP_Unit OP_Unit + = return OP_Unit + + go (ShapeRsnoc t) (OP_Pair sa1 sb1) (OP_Pair sa2 sb2) + = do x <- add (numType :: NumType Int) sb1 sb2 + OP_Pair <$> go t sa1 sa2 <*> return x + diff --git a/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Transform.hs b/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Transform.hs new file mode 100644 --- /dev/null +++ b/src/Data/Array/Accelerate/LLVM/PTX/CodeGen/Transform.hs @@ -0,0 +1,65 @@ +{-# LANGUAGE GADTs #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +-- | +-- Module : Data.Array.Accelerate.LLVM.PTX.CodeGen.Transform +-- Copyright : [2014..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.LLVM.PTX.CodeGen.Transform + where + +-- accelerate +import Data.Array.Accelerate.Representation.Array +import Data.Array.Accelerate.Type + +import Data.Array.Accelerate.LLVM.CodeGen.Arithmetic +import Data.Array.Accelerate.LLVM.CodeGen.Array +import Data.Array.Accelerate.LLVM.CodeGen.Base +import Data.Array.Accelerate.LLVM.CodeGen.Environment +import Data.Array.Accelerate.LLVM.CodeGen.Exp +import Data.Array.Accelerate.LLVM.CodeGen.Monad +import Data.Array.Accelerate.LLVM.CodeGen.Sugar + +import Data.Array.Accelerate.LLVM.PTX.CodeGen.Base +import Data.Array.Accelerate.LLVM.PTX.CodeGen.Loop +import Data.Array.Accelerate.LLVM.PTX.Target ( PTX ) + + +-- Apply a unary function to each element of an array. Each thread processes +-- multiple elements, striding the array by the grid size. +-- +mkTransform + :: Gamma aenv + -> ArrayR (Array sh a) + -> ArrayR (Array sh' b) + -> IRFun1 PTX aenv (sh' -> sh) + -> IRFun1 PTX aenv (a -> b) + -> CodeGen PTX (IROpenAcc PTX aenv (Array sh' b)) +mkTransform aenv repr@(ArrayR shr _) repr'@(ArrayR shr' _) p f = + let + (arrOut, paramOut) = mutableArray repr' "out" + (arrIn, paramIn) = mutableArray repr "in" + paramEnv = envParam aenv + in + makeOpenAcc "transform" (paramOut ++ paramIn ++ paramEnv) $ do + + let start = liftInt 0 + end <- shapeSize shr' (irArrayShape arrOut) + + imapFromTo start end $ \i' -> do + ix' <- indexOfInt shr' (irArrayShape arrOut) i' + ix <- app1 p ix' + i <- intOfIndex shr (irArrayShape arrIn) ix + a <- readArray TypeInt arrIn i + b <- app1 f a + writeArray TypeInt arrOut i' b + + return_ + diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Compile.hs b/src/Data/Array/Accelerate/LLVM/PTX/Compile.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/Compile.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/Compile.hs @@ -6,11 +6,10 @@ {-# OPTIONS_GHC -fno-warn-orphans #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Compile --- Copyright : [2014..2017] Trevor L. McDonell --- [2014..2014] Vinod Grover (NVIDIA Corporation) +-- Copyright : [2014..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- @@ -22,27 +21,16 @@ ) where --- llvm-hs -import qualified LLVM.AST as AST -import qualified LLVM.AST.Name as LLVM -import qualified LLVM.Context as LLVM -import qualified LLVM.Module as LLVM -import qualified LLVM.PassManager as LLVM -import qualified LLVM.Target as LLVM -import qualified LLVM.Internal.Module as LLVM.Internal -import qualified LLVM.Internal.FFI.LLVMCTypes as LLVM.Internal.FFI -import qualified LLVM.Analysis as LLVM - --- accelerate -import Data.Array.Accelerate.Error ( internalError ) -import Data.Array.Accelerate.Trafo ( DelayedOpenAcc ) +import Data.Array.Accelerate.AST ( PreOpenAcc ) +import Data.Array.Accelerate.Error +import Data.Array.Accelerate.Trafo.Delayed import Data.Array.Accelerate.LLVM.CodeGen import Data.Array.Accelerate.LLVM.CodeGen.Environment ( Gamma ) import Data.Array.Accelerate.LLVM.CodeGen.Module ( Module(..) ) import Data.Array.Accelerate.LLVM.Compile +import Data.Array.Accelerate.LLVM.Extra import Data.Array.Accelerate.LLVM.State -import Data.Array.Accelerate.LLVM.Util import Data.Array.Accelerate.LLVM.PTX.Analysis.Launch import Data.Array.Accelerate.LLVM.PTX.CodeGen @@ -52,12 +40,20 @@ import Data.Array.Accelerate.LLVM.PTX.Target import qualified Data.Array.Accelerate.LLVM.PTX.Debug as Debug --- cuda import Foreign.CUDA.Path import qualified Foreign.CUDA.Analysis as CUDA import qualified Foreign.NVVM as NVVM --- standard library +import qualified LLVM.AST as AST +import qualified LLVM.AST.Name as LLVM +import qualified LLVM.Context as LLVM +import qualified LLVM.Module as LLVM +import qualified LLVM.PassManager as LLVM +import qualified LLVM.Target as LLVM +import qualified LLVM.Internal.Module as LLVM.Internal +import qualified LLVM.Internal.FFI.LLVMCTypes as LLVM.Internal.FFI +import qualified LLVM.Analysis as LLVM + import Control.DeepSeq import Control.Exception import Control.Monad.Except @@ -81,7 +77,6 @@ import qualified Data.ByteString as B import qualified Data.ByteString.Char8 as B8 import qualified Data.ByteString.Internal as B -import qualified Data.ByteString.Short.Char8 as S8 import Prelude as P @@ -98,16 +93,15 @@ -- This generates the target code together with a list of each kernel function -- defined in the module paired with its occupancy information. -- -compile :: DelayedOpenAcc aenv a -> Gamma aenv -> LLVM PTX (ObjectR PTX) -compile acc aenv = do - target <- gets llvmTarget - (uid, cacheFile) <- cacheOfOpenAcc acc +compile :: HasCallStack => PreOpenAcc DelayedOpenAcc aenv a -> Gamma aenv -> LLVM PTX (ObjectR PTX) +compile pacc aenv = do -- Generate code for this Acc operation -- - let Module ast md = llvmOfOpenAcc target uid acc aenv - dev = ptxDeviceProperties target - config = [ (f,x) | (LLVM.Name f, KM_PTX x) <- Map.toList md ] + dev <- gets ptxDeviceProperties + (uid, cacheFile) <- cacheOfPreOpenAcc pacc + Module ast md <- llvmOfPreOpenAcc uid pacc aenv + let config = [ (f,x) | (LLVM.Name f, KM_PTX x) <- Map.toList md ] -- Lower the generated LLVM into a CUBIN object code. -- @@ -139,9 +133,9 @@ compilePTX :: CUDA.DeviceProperties -> LLVM.Context -> AST.Module -> IO ByteString compilePTX dev ctx ast = do #ifdef ACCELERATE_USE_NVVM - ptx <- withLibdeviceNVVM dev ctx ast (compileModuleNVVM dev (AST.moduleName ast)) + ptx <- withLibdeviceNVVM dev ctx ast (_compileModuleNVVM dev (AST.moduleName ast)) #else - ptx <- withLibdeviceNVPTX dev ctx ast (compileModuleNVPTX dev) + ptx <- withLibdeviceNVPTX dev ctx ast (_compileModuleNVPTX dev) #endif Debug.when Debug.dump_asm $ Debug.traceIO Debug.verbose (B8.unpack ptx) return ptx @@ -150,7 +144,7 @@ -- | Compile the given PTX assembly to a CUBIN file (SASS object code). The -- compiled code will be stored at the given FilePath. -- -compileCUBIN :: CUDA.DeviceProperties -> FilePath -> ByteString -> IO ByteString +compileCUBIN :: HasCallStack => CUDA.DeviceProperties -> FilePath -> ByteString -> IO ByteString compileCUBIN dev sass ptx = do _verbose <- if Debug.debuggingIsEnabled then Debug.getFlag Debug.verbose else return False _debug <- if Debug.debuggingIsEnabled then Debug.getFlag Debug.debug else return False @@ -188,7 +182,7 @@ -- wait on the process ex <- waitForProcess ph case ex of - ExitFailure r -> $internalError "compile" (printf "ptxas %s (exit %d)\n%s" (unwords flags) r info) + ExitFailure r -> internalError (printf "ptxas %s (exit %d)\n%s" (unwords flags) r info) ExitSuccess -> return () when _verbose $ @@ -202,8 +196,8 @@ -- Compile and optimise the module to PTX using the (closed source) NVVM -- library. This _may_ produce faster object code than the LLVM NVPTX compiler. -- -compileModuleNVVM :: CUDA.DeviceProperties -> ShortByteString -> [(String, ByteString)] -> LLVM.Module -> IO ByteString -compileModuleNVVM dev name libdevice mdl = do +_compileModuleNVVM :: HasCallStack => CUDA.DeviceProperties -> ShortByteString -> [(ShortByteString, ByteString)] -> LLVM.Module -> IO ByteString +_compileModuleNVVM dev name libdevice mdl = do _debug <- if Debug.debuggingIsEnabled then Debug.getFlag Debug.debug else return False -- let arch = CUDA.computeCapability dev @@ -223,7 +217,7 @@ header = case bitSize (undefined::Int) of 32 -> "target triple = \"nvptx-nvidia-cuda\"\ntarget datalayout = \"e-p:32:32:32-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v16:16:16-v32:32:32-v64:64:64-v128:128:128-n16:32:64\"" 64 -> "target triple = \"nvptx64-nvidia-cuda\"\ntarget datalayout = \"e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v16:16:16-v32:32:32-v64:64:64-v128:128:128-n16:32:64\"" - _ -> $internalError "compileModuleNVVM" "I don't know what architecture I am" + _ -> internalError "I don't know what architecture I am" Debug.when Debug.dump_cc $ do Debug.when Debug.verbose $ do @@ -233,11 +227,7 @@ -- Lower the generated module to bitcode, then compile and link together with -- the shim header and libdevice library (if necessary) bc <- LLVM.moduleBitcode mdl -#if MIN_VERSION_nvvm(0,9,0) - ptx <- NVVM.compileModules (("",header) : (name,bc) : fmap (\(n,b) -> (S8.pack n, b)) libdevice) flags -#else - ptx <- NVVM.compileModules (("",header) : (S8.unpack name,bc) : libdevice) flags -#endif + ptx <- NVVM.compileModules (("",header) : (name,bc) : libdevice) flags unless (B.null (NVVM.compileLog ptx)) $ do Debug.traceIO Debug.dump_cc $ "llvm: " ++ B8.unpack (NVVM.compileLog ptx) @@ -248,8 +238,8 @@ -- Compiling with the NVPTX backend uses LLVM-3.3 and above -- -compileModuleNVPTX :: CUDA.DeviceProperties -> LLVM.Module -> IO ByteString -compileModuleNVPTX dev mdl = +_compileModuleNVPTX :: CUDA.DeviceProperties -> LLVM.Module -> IO ByteString +_compileModuleNVPTX dev mdl = withPTXTargetMachine dev $ \nvptx -> do when Debug.internalChecksAreEnabled $ LLVM.verify mdl diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Compile/Cache.hs b/src/Data/Array/Accelerate/LLVM/PTX/Compile/Cache.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/Compile/Cache.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/Compile/Cache.hs @@ -2,10 +2,10 @@ {-# OPTIONS_GHC -fno-warn-orphans #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Compile.Cache --- Copyright : [2017] Trevor L. McDonell +-- Copyright : [2017..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- @@ -23,7 +23,7 @@ import Data.Version import Foreign.CUDA.Analysis import System.FilePath -import qualified Data.ByteString.Char8 as B8 +import Text.Printf import qualified Data.ByteString.Short.Char8 as S8 import Paths_accelerate_llvm_ptx @@ -31,12 +31,10 @@ instance Persistent PTX where targetCacheTemplate = do - dev <- gets ptxDeviceProperties - let Compute m n = computeCapability dev - -- + Compute m n <- gets (computeCapability . ptxDeviceProperties) return $ "accelerate-llvm-ptx-" ++ showVersion version "llvm-hs-" ++ VERSION_llvm_hs S8.unpack ptxTargetTriple - B8.unpack (ptxISAVersion m n) + printf "sm%d%d" m n "morp.sass" diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Compile/Libdevice.hs b/src/Data/Array/Accelerate/LLVM/PTX/Compile/Libdevice.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/Compile/Libdevice.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/Compile/Libdevice.hs @@ -4,11 +4,10 @@ {-# LANGUAGE ViewPatterns #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Compile.Libdevice --- Copyright : [2014..2017] Trevor L. McDonell --- [2014..2014] Vinod Grover (NVIDIA Corporation) +-- Copyright : [2014..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- @@ -119,7 +118,7 @@ :: DeviceProperties -> Context -> Module - -> ([(String, ByteString)] -> LLVM.Module -> IO a) + -> ([(ShortByteString, ByteString)] -> LLVM.Module -> IO a) -> IO a withLibdeviceNVVM dev ctx ast next = LLVM.withModuleFromAST ctx ast $ \mdl -> do diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Compile/Libdevice/Load.hs b/src/Data/Array/Accelerate/LLVM/PTX/Compile/Libdevice/Load.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/Compile/Libdevice/Load.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/Compile/Libdevice/Load.hs @@ -5,11 +5,10 @@ {-# LANGUAGE TupleSections #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Compile.Libdevice.Load --- Copyright : [2014..2017] Trevor L. McDonell --- [2014..2014] Vinod Grover (NVIDIA Corporation) +-- Copyright : [2014..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- @@ -20,24 +19,24 @@ ) where --- llvm-hs import LLVM.Context import LLVM.Module as LLVM import LLVM.AST as AST ( Module(..) ) --- accelerate import Data.Array.Accelerate.Error import Data.Array.Accelerate.LLVM.PTX.Compile.Libdevice.TH import Data.Array.Accelerate.LLVM.PTX.Execute.Event ( ) -- GHC#1012 import Data.Array.Accelerate.LLVM.PTX.Execute.Stream ( ) -- GHC#1012 --- cuda import Foreign.CUDA.Analysis import qualified Foreign.CUDA.Driver as CUDA --- standard library -import Data.ByteString ( ByteString ) import System.IO.Unsafe +import Data.ByteString ( ByteString ) +import Data.ByteString.Short.Char8 ( ShortByteString ) +import qualified Data.ByteString.Short.Char8 as S8 +import qualified Language.Haskell.TH as TH +import qualified Language.Haskell.TH.Syntax as TH -- NVVM Reflect @@ -49,7 +48,7 @@ instance NVVMReflect AST.Module where nvvmReflect = nvvmReflectModule -instance NVVMReflect (String, ByteString) where +instance NVVMReflect (ShortByteString, ByteString) where nvvmReflect = $$( nvvmReflectBitcode nvvmReflectModule ) @@ -64,80 +63,64 @@ class Libdevice a where libdevice :: Compute -> a -instance Libdevice AST.Module where - libdevice _ - | CUDA.libraryVersion >= 9000 - = libdevice_50_mdl - -- - libdevice (Compute n m) = - case (n,m) of - (2,_) -> libdevice_20_mdl -- 2.0, 2.1 - (3,x) | x < 5 -> libdevice_30_mdl -- 3.0, 3.2 - | otherwise -> libdevice_35_mdl -- 3.5, 3.7 - (5,_) -> libdevice_50_mdl -- 5.x - (6,_) -> libdevice_50_mdl -- 6.x - _ -> $internalError "libdevice" "no binary for this architecture" - -instance Libdevice (String, ByteString) where - libdevice _ - | CUDA.libraryVersion >= 9000 - = libdevice_50_bc - -- - libdevice (Compute n m) = - case (n,m) of - (2,_) -> libdevice_20_bc -- 2.0, 2.1 - (3,x) | x < 5 -> libdevice_30_bc -- 3.0, 3.2 - | otherwise -> libdevice_35_bc -- 3.5, 3.7 - (5,_) -> libdevice_50_bc -- 5.x - (6,_) -> libdevice_50_bc -- 6.x - _ -> $internalError "libdevice" "no binary for this architecture" - - -- Load the libdevice bitcode files as an LLVM AST module. The top-level --- unsafePerformIO ensures that the data is only read from disk once per program --- execution. +-- unsafePerformIO ensures that the data is only read from disk once per +-- program execution. -- --- TLM: As of CUDA-9.0, libdevice is no longer split into multiple files --- depending on the target compute architecture. The function 'libdeviceBitcode' --- knows this and ignores the architecture parameter, and in the above instances --- we only refer to the 5.0 module below. Although the TH splices will be run --- 4 times (and read in the same file 4 times) hopefully GHC is smart enough to --- remove the unused bindings as dead code... +-- As of CUDA-9.0, libdevice is no longer split into multiple files +-- depending on the target compute architecture. -- -{-# NOINLINE libdevice_20_mdl #-} -{-# NOINLINE libdevice_30_mdl #-} -{-# NOINLINE libdevice_35_mdl #-} -{-# NOINLINE libdevice_50_mdl #-} -libdevice_20_mdl, libdevice_30_mdl, libdevice_35_mdl, libdevice_50_mdl :: AST.Module -libdevice_20_mdl = unsafePerformIO $ libdeviceModule (Compute 2 0) -libdevice_30_mdl = unsafePerformIO $ libdeviceModule (Compute 3 0) -libdevice_35_mdl = unsafePerformIO $ libdeviceModule (Compute 3 5) -libdevice_50_mdl = unsafePerformIO $ libdeviceModule (Compute 5 0) +$( let + libdeviceModule :: TH.ExpQ + libdeviceModule = [| \(name, bc) -> + unsafePerformIO $ + withContext $ \ctx -> + withModuleFromBitcode ctx (S8.unpack name, bc) moduleAST |] + in + if CUDA.libraryVersion < 9000 + then + [d| {-# NOINLINE libdevice_20_mdl #-} + {-# NOINLINE libdevice_30_mdl #-} + {-# NOINLINE libdevice_35_mdl #-} + {-# NOINLINE libdevice_50_mdl #-} + libdevice_20_mdl, libdevice_30_mdl, libdevice_35_mdl, libdevice_50_mdl :: AST.Module + libdevice_20_mdl = $libdeviceModule libdevice_20_bc + libdevice_30_mdl = $libdeviceModule libdevice_30_bc + libdevice_35_mdl = $libdeviceModule libdevice_35_bc + libdevice_50_mdl = $libdeviceModule libdevice_50_bc --- Load the libdevice bitcode files as raw binary data. --- -libdevice_20_bc, libdevice_30_bc, libdevice_35_bc, libdevice_50_bc :: (String,ByteString) -libdevice_20_bc = $$( libdeviceBitcode (Compute 2 0) ) -libdevice_30_bc = $$( libdeviceBitcode (Compute 3 0) ) -libdevice_35_bc = $$( libdeviceBitcode (Compute 3 5) ) -libdevice_50_bc = $$( libdeviceBitcode (Compute 5 0) ) + libdevice_20_bc, libdevice_30_bc, libdevice_35_bc, libdevice_50_bc :: (ShortByteString,ByteString) + libdevice_20_bc = $( TH.unTypeQ $ libdeviceBitcode (Compute 2 0) ) + libdevice_30_bc = $( TH.unTypeQ $ libdeviceBitcode (Compute 3 0) ) + libdevice_35_bc = $( TH.unTypeQ $ libdeviceBitcode (Compute 3 5) ) + libdevice_50_bc = $( TH.unTypeQ $ libdeviceBitcode (Compute 5 0) ) + instance Libdevice AST.Module where + libdevice compute = + case compute of + Compute 2 _ -> libdevice_20_mdl -- 2.0, 2.1 + Compute 3 x + | x < 5 -> libdevice_30_mdl -- 3.0, 3.2 + | otherwise -> libdevice_35_mdl -- 3.5, 3.7 + Compute 5 _ -> libdevice_50_mdl -- 5.x + _ -> internalError + $ unlines [ "This device (compute capability " ++ show compute ++ ") is not supported by this version of the CUDA toolkit (" ++ show CUDA.libraryVersion ++ ")" + , "Please upgrade to the latest version of the CUDA toolkit and reinstall the 'cuda' package." + ] + |] + else + [d| {-# NOINLINE libdevice_mdl #-} + libdevice_mdl :: AST.Module + libdevice_mdl = $libdeviceModule libdevice_bc --- Load the libdevice bitcode file for the given compute architecture, and raise --- it to a Haskell AST that can be kept for future use. The name of the bitcode --- files follows: --- --- libdevice.compute_XX.YY.bc --- --- Where XX represents the compute capability, and YY represents a version(?) We --- search the libdevice PATH for all files of the appropriate compute capability --- and load the most recent. --- -libdeviceModule :: Compute -> IO AST.Module -libdeviceModule arch = do - let bc :: (String, ByteString) - bc = libdevice arch - -- - withContext $ \ctx -> - withModuleFromBitcode ctx bc moduleAST + libdevice_bc :: (ShortByteString,ByteString) + libdevice_bc = $( TH.unTypeQ $ libdeviceBitcode undefined ) + + instance Libdevice AST.Module where + libdevice _ = libdevice_mdl + + instance Libdevice (ShortByteString,ByteString) where + libdevice _ = libdevice_bc + |] + ) diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Compile/Libdevice/TH.hs b/src/Data/Array/Accelerate/LLVM/PTX/Compile/Libdevice/TH.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/Compile/Libdevice/TH.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/Compile/Libdevice/TH.hs @@ -1,11 +1,13 @@ +{-# LANGUAGE CPP #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeApplications #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Compile.Libdevice.TH --- Copyright : [2017] Trevor L. McDonell +-- Copyright : [2017..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- @@ -23,26 +25,35 @@ import qualified LLVM.Context as LLVM import qualified LLVM.Module as LLVM +import LLVM.AST.Type.Downcast import LLVM.AST.Type.Representation import Data.Array.Accelerate.Error import Data.Array.Accelerate.LLVM.CodeGen.Base -import Data.Array.Accelerate.LLVM.CodeGen.Downcast import Data.Array.Accelerate.LLVM.PTX.Target -import Foreign.CUDA.Path import Foreign.CUDA.Analysis import qualified Foreign.CUDA.Driver as CUDA +#if MIN_VERSION_nvvm(0,10,0) +import Foreign.NVVM.Path +#else +import Foreign.CUDA.Path +#endif import Data.ByteString ( ByteString ) +import Data.ByteString.Short ( ShortByteString ) import Data.FileEmbed import Data.List import Data.Maybe -import Language.Haskell.TH.Syntax hiding ( Name ) +import Language.Haskell.TH.Syntax ( Q, TExp ) import System.Directory import System.FilePath import Text.Printf import qualified Data.ByteString.Short as BS +import qualified Data.ByteString.Short.Char8 as S8 +import qualified Data.ByteString.Short.Extra as BS +import qualified Language.Haskell.TH as TH +import qualified Language.Haskell.TH.Syntax as TH -- This is a hacky module that can be linked against in order to provide the @@ -71,8 +82,8 @@ AST.Module { AST.moduleName = "nvvm-reflect" , AST.moduleSourceFileName = BS.empty - , AST.moduleDataLayout = targetDataLayout (undefined::PTX) - , AST.moduleTargetTriple = targetTriple (undefined::PTX) + , AST.moduleDataLayout = targetDataLayout @PTX + , AST.moduleTargetTriple = targetTriple @PTX , AST.moduleDefinitions = [AST.GlobalDefinition $ AST.G.functionDefaults { AST.G.name = AST.Name "__nvvm_reflect" , AST.G.returnType = downcast (integralType :: IntegralType Int32) @@ -85,14 +96,14 @@ -- Lower the given NVVM Reflect module into bitcode. -- -nvvmReflectBitcode :: AST.Module -> Q (TExp (String, ByteString)) +nvvmReflectBitcode :: AST.Module -> Q (TExp (ShortByteString, ByteString)) nvvmReflectBitcode mdl = do let name = "__nvvm_reflect" -- - bs <- runIO $ LLVM.withContext $ \ctx -> do - LLVM.withModuleFromAST ctx mdl LLVM.moduleLLVMAssembly - be <- bsToExp bs - return . TExp $ TupE [ LitE (StringL name), be ] + bs <- TH.runIO $ LLVM.withContext $ \ctx -> + LLVM.withModuleFromAST ctx mdl LLVM.moduleLLVMAssembly + TH.unsafeTExpCoerce $ TH.tupE [ TH.unTypeQ (BS.liftSBS name) + , bsToExp bs ] -- Load the libdevice bitcode file for the given compute architecture. The name @@ -104,85 +115,27 @@ -- search the libdevice PATH for all files of the appropriate compute capability -- and load the "most recent" (by sort order). -- -libdeviceBitcode :: Compute -> Q (TExp (String, ByteString)) -libdeviceBitcode (Compute m n) = do +libdeviceBitcode :: HasCallStack => Compute -> Q (TExp (ShortByteString, ByteString)) +libdeviceBitcode compute = do let basename - | CUDA.libraryVersion < 9000 = printf "libdevice.compute_%d%d" m n + | CUDA.libraryVersion < 9000 + , Compute m n <- compute = printf "libdevice.compute_%d%d" m n | otherwise = "libdevice" -- - err = $internalError "libdevice" (printf "not found: %s.YY.bc" basename) + err = internalError (printf "not found: %s.YY.bc" basename) best f = basename `isPrefixOf` f && takeExtension f == ".bc" +#if MIN_VERSION_nvvm(0,10,0) + base = nvvmDeviceLibraryPath +#else base = cudaInstallPath "nvvm" "libdevice" +#endif + -- - files <- runIO $ getDirectoryContents base + files <- TH.runIO $ getDirectoryContents base -- let name = fromMaybe err . listToMaybe . sortBy (flip compare) $ filter best files path = base name -- - bc <- embedFile path - return . TExp $ TupE [ LitE (StringL name), bc ] - - --- Determine the location of the libdevice bitcode libraries. We search for the --- location of the 'nvcc' executable in the PATH. From that, we assume the --- location of the libdevice bitcode files. --- --- libdevicePath :: IO FilePath --- libdevicepath = do --- nvcc <- fromMaybe (error "could not find 'nvcc' in PATH") `fmap` findExecutable "nvcc" --- -- --- let ccvn = reverse (splitPath nvcc) --- dir = "libdevice" : "nvvm" : drop 2 ccvn --- -- --- return (joinPath (reverse dir)) - - --- With these instances it is possible to also write TH function to raise the --- libNVVM modules to an AST. However, generating those large ASTs results in --- awful compile times. --- --- $( deriveLift ''AST.AddrSpace ) --- $( deriveLift ''AST.AlignType ) --- $( deriveLift ''AST.AlignmentInfo ) --- $( deriveLift ''AST.BasicBlock ) --- $( deriveLift ''AST.CallingConvention ) --- $( deriveLift ''AST.Constant ) --- $( deriveLift ''AST.DataLayout ) --- $( deriveLift ''AST.Definition ) --- $( deriveLift ''AST.Dialect ) --- $( deriveLift ''AST.Endianness ) --- $( deriveLift ''AST.FastMathFlags ) --- $( deriveLift ''AST.FloatingPointFormat ) --- $( deriveLift ''AST.FloatingPointPredicate ) --- $( deriveLift ''AST.FunctionAttribute ) --- $( deriveLift ''AST.Global ) --- $( deriveLift ''AST.GroupID ) --- $( deriveLift ''AST.InlineAssembly ) --- $( deriveLift ''AST.Instruction ) --- $( deriveLift ''AST.IntegerPredicate ) --- $( deriveLift ''AST.LandingPadClause ) --- $( deriveLift ''AST.Linkage ) --- $( deriveLift ''AST.Mangling ) --- $( deriveLift ''AST.MemoryOrdering ) --- $( deriveLift ''AST.Metadata ) --- $( deriveLift ''AST.MetadataNode ) --- $( deriveLift ''AST.MetadataNodeID ) --- $( deriveLift ''AST.Model ) --- $( deriveLift ''AST.Module ) --- $( deriveLift ''AST.Name ) --- $( deriveLift ''AST.Named ) --- $( deriveLift ''AST.Operand ) --- $( deriveLift ''AST.Parameter ) --- $( deriveLift ''AST.ParameterAttribute ) --- $( deriveLift ''AST.RMWOperation ) --- $( deriveLift ''AST.SelectionKind ) --- $( deriveLift ''AST.SomeFloat ) --- $( deriveLift ''AST.StorageClass ) --- $( deriveLift ''AST.SynchronizationScope ) --- $( deriveLift ''AST.TailCallKind ) --- $( deriveLift ''AST.Terminator ) --- $( deriveLift ''AST.Type ) --- $( deriveLift ''AST.UnnamedAddr ) --- $( deriveLift ''AST.Visibility ) --- $( deriveLift ''NonEmpty ) + TH.unsafeTExpCoerce $ TH.tupE [ TH.unTypeQ (BS.liftSBS (S8.pack name)) + , embedFile path ] diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Context.hs b/src/Data/Array/Accelerate/LLVM/PTX/Context.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/Context.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/Context.hs @@ -2,11 +2,10 @@ {-# LANGUAGE RecordWildCards #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Context --- Copyright : [2014..2017] Trevor L. McDonell --- [2014..2014] Vinod Grover (NVIDIA Corporation) +-- Copyright : [2014..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- @@ -22,11 +21,10 @@ import Data.Array.Accelerate.LLVM.PTX.Analysis.Device import qualified Data.Array.Accelerate.LLVM.PTX.Debug as Debug -import qualified Foreign.CUDA.Analysis as CUDA -import qualified Foreign.CUDA.Driver as CUDA import qualified Foreign.CUDA.Driver.Device as CUDA import qualified Foreign.CUDA.Driver.Context as CUDA +import Control.Concurrent import Control.Exception import Control.Monad import Data.Hashable @@ -102,9 +100,10 @@ -- {-# INLINE withContext #-} withContext :: Context -> IO a -> IO a -withContext Context{..} action = - withLifetime deviceContext $ \ctx -> - bracket_ (push ctx) pop action +withContext Context{..} action + = runInBoundThread + $ withLifetime deviceContext $ \ctx -> + bracket_ (push ctx) pop action {-# INLINE push #-} push :: CUDA.Context -> IO () diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Debug.hs b/src/Data/Array/Accelerate/LLVM/PTX/Debug.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/Debug.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/Debug.hs @@ -1,10 +1,10 @@ +{-# LANGUAGE ForeignFunctionInterface #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Debug --- Copyright : [2014..2017] Trevor L. McDonell --- [2014..2014] Vinod Grover (NVIDIA Corporation) +-- Copyright : [2014..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- @@ -23,7 +23,6 @@ import Control.Monad.Trans import Control.Concurrent -import Data.Time.Clock import System.CPUTime import Text.Printf @@ -57,13 +56,13 @@ then do gpuBegin <- liftIO $ Event.create [] gpuEnd <- liftIO $ Event.create [] - wallBegin <- liftIO $ getCurrentTime + wallBegin <- liftIO $ getMonotonicTime cpuBegin <- liftIO $ getCPUTime _ <- liftIO $ Event.record gpuBegin stream result <- action _ <- liftIO $ Event.record gpuEnd stream cpuEnd <- liftIO $ getCPUTime - wallEnd <- liftIO $ getCurrentTime + wallEnd <- liftIO $ getMonotonicTime -- Wait for the GPU to finish executing then display the timing execution -- message. Do this in a separate thread so that the remaining kernels can @@ -74,7 +73,7 @@ diff <- Event.elapsedTime gpuBegin gpuEnd let gpuTime = float2Double $ diff * 1E-3 -- milliseconds cpuTime = fromIntegral (cpuEnd - cpuBegin) * 1E-12 -- picoseconds - wallTime = realToFrac (diffUTCTime wallEnd wallBegin) + wallTime = wallEnd - wallBegin -- seconds Event.destroy gpuBegin Event.destroy gpuEnd @@ -94,4 +93,7 @@ (showFFloatSIBase (Just 3) 1000 wallTime "s") (showFFloatSIBase (Just 3) 1000 cpuTime "s") (showFFloatSIBase (Just 3) 1000 gpuTime "s") + +-- accelerate/cbits/clock.c +foreign import ccall unsafe "clock_gettime_monotonic_seconds" getMonotonicTime :: IO Double diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Embed.hs b/src/Data/Array/Accelerate/LLVM/PTX/Embed.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/Embed.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/Embed.hs @@ -1,12 +1,13 @@ +{-# LANGUAGE CPP #-} {-# LANGUAGE QuasiQuotes #-} {-# LANGUAGE TemplateHaskell #-} {-# OPTIONS_GHC -fno-warn-orphans #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Embed --- Copyright : [2017] Trevor L. McDonell +-- Copyright : [2017..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- @@ -17,8 +18,7 @@ ) where -import Data.ByteString.Short.Char8 as S8 -import Data.ByteString.Short.Internal as BS +import Data.ByteString.Short.Extra as BS import Data.Array.Accelerate.Lifetime @@ -30,7 +30,6 @@ import Data.Array.Accelerate.LLVM.PTX.Target import Data.Array.Accelerate.LLVM.PTX.Context --- import qualified Foreign.CUDA.Analysis as CUDA import qualified Foreign.CUDA.Driver as CUDA import Foreign.Ptr @@ -70,17 +69,10 @@ linkQ :: TH.Name -> (Kernel, Q (TExp (Int -> Int))) -> Q (TExp Kernel) linkQ jit (Kernel name _ dsmem cta _, grid) = [|| unsafePerformIO $ do - f <- CUDA.getFun (CUDA.jitModule $$(TH.unsafeTExpCoerce (TH.varE jit))) $$(TH.unsafeTExpCoerce (TH.lift (S8.unpack name))) + f <- CUDA.getFun (CUDA.jitModule $$(TH.unsafeTExpCoerce (TH.varE jit))) $$(liftSBS name) return $ Kernel $$(liftSBS name) f dsmem cta $$grid ||] listE :: [Q (TExp a)] -> Q (TExp [a]) listE xs = TH.unsafeTExpCoerce (TH.listE (map TH.unTypeQ xs)) - - liftSBS :: ShortByteString -> Q (TExp ShortByteString) - liftSBS bs = - let bytes = BS.unpack bs - len = BS.length bs - in - [|| unsafePerformIO $ BS.createFromPtr $$( TH.unsafeTExpCoerce [| Ptr $(TH.litE (TH.StringPrimL bytes)) |]) len ||] diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Execute.hs b/src/Data/Array/Accelerate/LLVM/PTX/Execute.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/Execute.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/Execute.hs @@ -5,35 +5,36 @@ {-# LANGUAGE RecordWildCards #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeOperators #-} +{-# LANGUAGE ViewPatterns #-} {-# OPTIONS_GHC -fno-warn-orphans #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Execute --- Copyright : [2014..2017] Trevor L. McDonell --- [2014..2014] Vinod Grover (NVIDIA Corporation) +-- Copyright : [2014..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- module Data.Array.Accelerate.LLVM.PTX.Execute ( - executeAcc, executeAfun, + executeAcc, executeOpenAcc, ) where --- accelerate import Data.Array.Accelerate.Analysis.Match -import Data.Array.Accelerate.Array.Sugar import Data.Array.Accelerate.Error import Data.Array.Accelerate.Lifetime +import Data.Array.Accelerate.Representation.Array +import Data.Array.Accelerate.Representation.Shape +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Type -import Data.Array.Accelerate.LLVM.Analysis.Match import Data.Array.Accelerate.LLVM.Execute -import Data.Array.Accelerate.LLVM.State import Data.Array.Accelerate.LLVM.PTX.Analysis.Launch ( multipleOf ) import Data.Array.Accelerate.LLVM.PTX.Array.Data @@ -41,28 +42,28 @@ import Data.Array.Accelerate.LLVM.PTX.Execute.Async import Data.Array.Accelerate.LLVM.PTX.Execute.Environment import Data.Array.Accelerate.LLVM.PTX.Execute.Marshal +import Data.Array.Accelerate.LLVM.PTX.Execute.Stream ( Stream ) import Data.Array.Accelerate.LLVM.PTX.Link import Data.Array.Accelerate.LLVM.PTX.Target import qualified Data.Array.Accelerate.LLVM.PTX.Debug as Debug - -import Data.Range ( Range(..) ) -import Control.Parallel.Meta ( runExecutable ) +import qualified Data.Array.Accelerate.LLVM.PTX.Execute.Event as Event --- cuda import qualified Foreign.CUDA.Driver as CUDA --- library -import Control.Monad ( when ) -import Control.Monad.State ( gets, liftIO ) +import Control.Monad ( when, forM_ ) +import Control.Monad.Reader ( asks, local ) +import Control.Monad.State ( liftIO ) import Data.ByteString.Short.Char8 ( ShortByteString, unpack ) +import qualified Data.DList as DL import Data.List ( find ) import Data.Maybe ( fromMaybe ) -import Data.Word ( Word32 ) import Text.Printf ( printf ) import Prelude hiding ( exp, map, sum, scanl, scanr ) -import qualified Prelude as P +{-# SPECIALISE INLINE executeAcc :: ExecAcc PTX a -> Par PTX (FutureArraysR PTX a) #-} +{-# SPECIALISE INLINE executeOpenAcc :: ExecOpenAcc PTX aenv a -> Val aenv -> Par PTX (FutureArraysR PTX a) #-} + -- Array expression evaluation -- --------------------------- @@ -80,22 +81,30 @@ -- code. -- instance Execute PTX where - map = simpleOp - generate = simpleOp - transform = simpleOp - backpermute = simpleOp - fold = foldOp - fold1 = fold1Op - foldSeg = foldSegOp - fold1Seg = foldSegOp - scanl = scanOp - scanl1 = scan1Op - scanl' = scan'Op - scanr = scanOp - scanr1 = scan1Op - scanr' = scan'Op + {-# INLINE map #-} + {-# INLINE generate #-} + {-# INLINE transform #-} + {-# INLINE backpermute #-} + {-# INLINE fold #-} + {-# INLINE foldSeg #-} + {-# INLINE scan #-} + {-# INLINE scan' #-} + {-# INLINE permute #-} + {-# INLINE stencil1 #-} + {-# INLINE stencil2 #-} + {-# INLINE aforeign #-} + map = mapOp + generate = generateOp + transform = transformOp + backpermute = backpermuteOp + fold True = foldOp + fold False = fold1Op + foldSeg i _ = foldSegOp i + scan _ True = scanOp + scan _ False = scan1Op + scan' _ = scan'Op permute = permuteOp - stencil1 = simpleOp + stencil1 = stencil1Op stencil2 = stencil2Op aforeign = aforeignOp @@ -105,40 +114,97 @@ -- Simple kernels just need to know the shape of the output array -- +{-# INLINE simpleOp #-} simpleOp - :: (Shape sh, Elt e) - => ExecutableR PTX + :: HasCallStack + => ShortByteString + -> ArrayR (Array sh e) + -> ExecutableR PTX -> Gamma aenv - -> Aval aenv - -> Stream + -> Val aenv -> sh - -> LLVM PTX (Array sh e) -simpleOp exe gamma aenv stream sh = withExecutable exe $ \ptxExecutable -> do - let kernel = case functionTable ptxExecutable of - k:_ -> k - _ -> $internalError "simpleOp" "no kernels found" - -- - out <- allocateRemote sh - ptx <- gets llvmTarget - liftIO $ executeOp ptx kernel gamma aenv stream (IE 0 (size sh)) out - return out + -> Par PTX (Future (Array sh e)) +simpleOp name repr exe gamma aenv sh = + withExecutable exe $ \ptxExecutable -> do + future <- new + result <- allocateRemote repr sh + -- + let paramR = TupRsingle $ ParamRarray repr + executeOp (ptxExecutable !# name) gamma aenv (arrayRshape repr) sh paramR result + put future result + return future -simpleNamed - :: (Shape sh, Elt e) - => ShortByteString +-- Mapping over an array can ignore the dimensionality of the array and +-- treat it as its underlying linear representation. +-- +{-# INLINE mapOp #-} +mapOp + :: HasCallStack + => Maybe (a :~: b) + -> ArrayR (Array sh a) + -> TypeR b -> ExecutableR PTX -> Gamma aenv - -> Aval aenv - -> Stream + -> Val aenv + -> Array sh a + -> Par PTX (Future (Array sh b)) +mapOp inplace repr tp exe gamma aenv input@(shape -> sh) = + withExecutable exe $ \ptxExecutable -> do + let reprOut = ArrayR (arrayRshape repr) tp + future <- new + result <- case inplace of + Just Refl -> return input + Nothing -> allocateRemote reprOut sh + -- + let paramsR = TupRsingle (ParamRarray reprOut) `TupRpair` TupRsingle (ParamRarray repr) + executeOp (ptxExecutable !# "map") gamma aenv (arrayRshape repr) sh paramsR (result, input) + put future result + return future + +{-# INLINE generateOp #-} +generateOp + :: HasCallStack + => ArrayR (Array sh e) + -> ExecutableR PTX + -> Gamma aenv + -> Val aenv -> sh - -> LLVM PTX (Array sh e) -simpleNamed fun exe gamma aenv stream sh = withExecutable exe $ \ptxExecutable -> do - out <- allocateRemote sh - ptx <- gets llvmTarget - liftIO $ executeOp ptx (ptxExecutable !# fun) gamma aenv stream (IE 0 (size sh)) out - return out + -> Par PTX (Future (Array sh e)) +generateOp = simpleOp "generate" +{-# INLINE transformOp #-} +transformOp + :: HasCallStack + => ArrayR (Array sh a) + -> ArrayR (Array sh' b) + -> ExecutableR PTX + -> Gamma aenv + -> Val aenv + -> sh' + -> Array sh a + -> Par PTX (Future (Array sh' b)) +transformOp repr repr' exe gamma aenv sh' input = + withExecutable exe $ \ptxExecutable -> do + future <- new + result <- allocateRemote repr' sh' + let paramsR = TupRsingle (ParamRarray repr') `TupRpair` TupRsingle (ParamRarray repr) + executeOp (ptxExecutable !# "transform") gamma aenv (arrayRshape repr') sh' paramsR (result, input) + put future result + return future +{-# INLINE backpermuteOp #-} +backpermuteOp + :: HasCallStack + => ArrayR (Array sh e) + -> ShapeR sh' + -> ExecutableR PTX + -> Gamma aenv + -> Val aenv + -> sh' + -> Array sh e + -> Par PTX (Future (Array sh' e)) +backpermuteOp (ArrayR shr tp) shr' = transformOp (ArrayR shr tp) (ArrayR shr' tp) + -- There are two flavours of fold operation: -- -- 1. If we are collapsing to a single value, then multiple thread blocks are @@ -153,378 +219,561 @@ -- block. Currently we always use the first, but require benchmarking to -- determine when to select each. -- +{-# INLINE fold1Op #-} fold1Op - :: (Shape sh, Elt e) - => ExecutableR PTX + :: HasCallStack + => ArrayR (Array sh e) + -> ExecutableR PTX -> Gamma aenv - -> Aval aenv - -> Stream - -> (sh :. Int) - -> LLVM PTX (Array sh e) -fold1Op exe gamma aenv stream sh@(sx :. sz) - = $boundsCheck "fold1" "empty array" (sz > 0) - $ case size sh of - 0 -> allocateRemote sx -- empty, but possibly with one or more non-zero dimensions - _ -> foldCore exe gamma aenv stream sh + -> Val aenv + -> Delayed (Array (sh, Int) e) + -> Par PTX (Future (Array sh e)) +fold1Op repr exe gamma aenv arr@(delayedShape -> sh@(sx, sz)) + = boundsCheck "empty array" (sz > 0) + $ case size (ShapeRsnoc $ arrayRshape repr) sh of + 0 -> newFull =<< allocateRemote repr sx -- empty, but possibly with one or more non-zero dimensions + _ -> foldCore repr exe gamma aenv arr +{-# INLINE foldOp #-} foldOp - :: (Shape sh, Elt e) - => ExecutableR PTX + :: HasCallStack + => ArrayR (Array sh e) + -> ExecutableR PTX -> Gamma aenv - -> Aval aenv - -> Stream - -> (sh :. Int) - -> LLVM PTX (Array sh e) -foldOp exe gamma aenv stream sh@(sx :. _) - = case size sh of - 0 -> simpleNamed "generate" exe gamma aenv stream sx - _ -> foldCore exe gamma aenv stream sh + -> Val aenv + -> Delayed (Array (sh, Int) e) + -> Par PTX (Future (Array sh e)) +foldOp repr exe gamma aenv arr@(delayedShape -> sh@(sx, _)) + = case size (ShapeRsnoc $ arrayRshape repr) sh of + 0 -> generateOp repr exe gamma aenv sx + _ -> foldCore repr exe gamma aenv arr +{-# INLINE foldCore #-} foldCore - :: forall aenv sh e. (Shape sh, Elt e) - => ExecutableR PTX + :: HasCallStack + => ArrayR (Array sh e) + -> ExecutableR PTX -> Gamma aenv - -> Aval aenv - -> Stream - -> (sh :. Int) - -> LLVM PTX (Array sh e) -foldCore exe gamma aenv stream sh - | Just Refl <- matchShapeType (undefined::sh) (undefined::Z) - = foldAllOp exe gamma aenv stream sh + -> Val aenv + -> Delayed (Array (sh, Int) e) + -> Par PTX (Future (Array sh e)) +foldCore repr exe gamma aenv arr + | ArrayR ShapeRz tp <- repr + = foldAllOp tp exe gamma aenv arr -- | otherwise - = foldDimOp exe gamma aenv stream sh - + = foldDimOp repr exe gamma aenv arr +{-# INLINE foldAllOp #-} foldAllOp - :: forall aenv e. Elt e - => ExecutableR PTX + :: forall aenv e. HasCallStack + => TypeR e + -> ExecutableR PTX -> Gamma aenv - -> Aval aenv - -> Stream - -> DIM1 - -> LLVM PTX (Scalar e) -foldAllOp exe gamma aenv stream (Z :. n) = withExecutable exe $ \ptxExecutable -> do - ptx <- gets llvmTarget - let - ks = ptxExecutable !# "foldAllS" - km1 = ptxExecutable !# "foldAllM1" - km2 = ptxExecutable !# "foldAllM2" - -- - if kernelThreadBlocks ks n == 1 - then do - -- The array is small enough that we can compute it in a single step - out <- allocateRemote Z - liftIO $ executeOp ptx ks gamma aenv stream (IE 0 n) out - return out + -> Val aenv + -> Delayed (Vector e) + -> Par PTX (Future (Scalar e)) +foldAllOp tp exe gamma aenv input = + withExecutable exe $ \ptxExecutable -> do + future <- new + let + ks = ptxExecutable !# "foldAllS" + km1 = ptxExecutable !# "foldAllM1" + km2 = ptxExecutable !# "foldAllM2" + sh@((), n) = delayedShape input + paramsRinput = TupRsingle $ ParamRmaybe $ ParamRarray $ ArrayR dim1 tp + paramsRdim0 = TupRsingle $ ParamRarray $ ArrayR dim0 tp + paramsRdim1 = TupRsingle $ ParamRarray $ ArrayR dim1 tp + -- + if kernelThreadBlocks ks n == 1 + then do + -- The array is small enough that we can compute it in a single step + result <- allocateRemote (ArrayR dim0 tp) () + let paramsR = paramsRdim0 `TupRpair` paramsRinput + executeOp ks gamma aenv dim1 sh paramsR (result, manifest input) + put future result - else do - -- Multi-kernel reduction to a single element. The first kernel integrates - -- any delayed elements, and the second is called recursively until - -- reaching a single element. - let - rec :: Vector e -> LLVM PTX (Scalar e) - rec tmp@(Array ((),m) adata) - | m <= 1 = return $ Array () adata - | otherwise = do - let s = m `multipleOf` kernelThreadBlockSize km2 - out <- allocateRemote (Z :. s) - liftIO $ executeOp ptx km2 gamma aenv stream (IE 0 s) (tmp, out) - rec out - -- - let s = n `multipleOf` kernelThreadBlockSize km1 - tmp <- allocateRemote (Z :. s) - liftIO $ executeOp ptx km1 gamma aenv stream (IE 0 s) tmp - rec tmp + else do + -- Multi-kernel reduction to a single element. The first kernel integrates + -- any delayed elements, and the second is called recursively until + -- reaching a single element. + let + rec :: Vector e -> Par PTX () + rec tmp@(Array ((),m) adata) + | m <= 1 = put future (Array () adata) + | otherwise = do + let sh' = ((), m `multipleOf` kernelThreadBlockSize km2) + out <- allocateRemote (ArrayR dim1 tp) sh' + let paramsR2 = paramsRdim1 `TupRpair` paramsRdim1 + executeOp km2 gamma aenv dim1 sh' paramsR2 (tmp, out) + rec out + -- + let sh' = ((), n `multipleOf` kernelThreadBlockSize km1) + tmp <- allocateRemote (ArrayR dim1 tp) sh' + let paramsR1 = paramsRdim1 `TupRpair` paramsRinput + executeOp km1 gamma aenv dim1 sh' paramsR1 (tmp, manifest input) + rec tmp + -- + return future +{-# INLINE foldDimOp #-} foldDimOp - :: forall aenv sh e. (Shape sh, Elt e) - => ExecutableR PTX + :: HasCallStack + => ArrayR (Array sh e) + -> ExecutableR PTX -> Gamma aenv - -> Aval aenv - -> Stream - -> (sh :. Int) - -> LLVM PTX (Array sh e) -foldDimOp exe gamma aenv stream (sh :. sz) = withExecutable exe $ \ptxExecutable -> do - let - kernel = if sz > 0 - then ptxExecutable !# "fold" - else ptxExecutable !# "generate" - -- - out <- allocateRemote sh - ptx <- gets llvmTarget - liftIO $ executeOp ptx kernel gamma aenv stream (IE 0 (size sh)) out - return out + -> Val aenv + -> Delayed (Array (sh, Int) e) + -> Par PTX (Future (Array sh e)) +foldDimOp repr@(ArrayR shr tp) exe gamma aenv input@(delayedShape -> (sh, sz)) + | sz == 0 = generateOp repr exe gamma aenv sh + | otherwise = + withExecutable exe $ \ptxExecutable -> do + future <- new + result <- allocateRemote repr sh + -- + let paramsR = TupRsingle (ParamRarray repr) `TupRpair` TupRsingle (ParamRmaybe $ ParamRarray $ ArrayR (ShapeRsnoc shr) tp) + executeOp (ptxExecutable !# "fold") gamma aenv shr sh paramsR (result, manifest input) + put future result + return future +{-# INLINE foldSegOp #-} foldSegOp - :: (Shape sh, Elt e) - => ExecutableR PTX + :: HasCallStack + => IntegralType i + -> ArrayR (Array (sh, Int) e) + -> ExecutableR PTX -> Gamma aenv - -> Aval aenv - -> Stream - -> (sh :. Int) - -> (Z :. Int) - -> LLVM PTX (Array (sh :. Int) e) -foldSegOp exe gamma aenv stream (sh :. sz) (Z :. ss) = withExecutable exe $ \ptxExecutable -> do - let - n = ss - 1 -- segments array has been 'scanl (+) 0'`ed - m = size sh * n - foldseg = if (sz`quot`ss) < (2 * kernelThreadBlockSize foldseg_cta) - then foldseg_warp - else foldseg_cta - -- - foldseg_cta = ptxExecutable !# "foldSeg_block" - foldseg_warp = ptxExecutable !# "foldSeg_warp" - -- qinit = ptxExecutable !# "qinit" - -- - out <- allocateRemote (sh :. n) - ptx <- gets llvmTarget - liftIO $ executeOp ptx foldseg gamma aenv stream (IE 0 m) out - return out + -> Val aenv + -> Delayed (Array (sh, Int) e) + -> Delayed (Segments i) + -> Par PTX (Future (Array (sh, Int) e)) +foldSegOp intTp repr exe gamma aenv input@(delayedShape -> (sh, sz)) segments@(delayedShape -> ((), ss)) = + withExecutable exe $ \ptxExecutable -> do + let + ArrayR (ShapeRsnoc shr') _ = repr + reprSeg = ArrayR dim1 $ TupRsingle $ SingleScalarType $ NumSingleType $ IntegralNumType intTp + n = ss - 1 -- segments array has been 'scanl (+) 0'`ed + m = size shr' sh * n + foldseg = if (sz`quot`ss) < (2 * kernelThreadBlockSize foldseg_cta) + then foldseg_warp + else foldseg_cta + -- + foldseg_cta = ptxExecutable !# "foldSeg_block" + foldseg_warp = ptxExecutable !# "foldSeg_warp" + -- qinit = ptxExecutable !# "qinit" + -- + future <- new + result <- allocateRemote repr (sh, n) + let paramsR = TupRsingle (ParamRarray repr) `TupRpair` TupRsingle (ParamRmaybe $ ParamRarray repr) `TupRpair` TupRsingle (ParamRmaybe $ ParamRarray reprSeg) + executeOp foldseg gamma aenv dim1 ((), m) paramsR ((result, manifest input), manifest segments) + put future result + return future +{-# INLINE scanOp #-} scanOp - :: (Shape sh, Elt e) - => ExecutableR PTX + :: HasCallStack + => ArrayR (Array (sh, Int) e) + -> ExecutableR PTX -> Gamma aenv - -> Aval aenv - -> Stream - -> sh :. Int - -> LLVM PTX (Array (sh:.Int) e) -scanOp exe gamma aenv stream (sz :. n) = + -> Val aenv + -> Delayed (Array (sh, Int) e) + -> Par PTX (Future (Array (sh, Int) e)) +scanOp repr exe gamma aenv input@(delayedShape -> (sz, n)) = case n of - 0 -> simpleNamed "generate" exe gamma aenv stream (sz :. 1) - _ -> scanCore exe gamma aenv stream sz n (n+1) + 0 -> generateOp repr exe gamma aenv (sz, 1) + _ -> scanCore repr exe gamma aenv (n+1) input +{-# INLINE scan1Op #-} scan1Op - :: (Shape sh, Elt e) - => ExecutableR PTX + :: HasCallStack + => ArrayR (Array (sh, Int) e) + -> ExecutableR PTX -> Gamma aenv - -> Aval aenv - -> Stream - -> sh :. Int - -> LLVM PTX (Array (sh:.Int) e) -scan1Op exe gamma aenv stream (sz :. n) - = $boundsCheck "scan1" "empty array" (n > 0) - $ scanCore exe gamma aenv stream sz n n + -> Val aenv + -> Delayed (Array (sh, Int) e) + -> Par PTX (Future (Array (sh, Int) e)) +scan1Op repr exe gamma aenv input@(delayedShape -> (_, n)) + = boundsCheck "empty array" (n > 0) + $ scanCore repr exe gamma aenv n input +{-# INLINE scanCore #-} scanCore - :: forall aenv sh e. (Shape sh, Elt e) - => ExecutableR PTX + :: HasCallStack + => ArrayR (Array (sh, Int) e) + -> ExecutableR PTX -> Gamma aenv - -> Aval aenv - -> Stream - -> sh - -> Int -- input size - -> Int -- output size - -> LLVM PTX (Array (sh:.Int) e) -scanCore exe gamma aenv stream sz n m - | Just Refl <- matchShapeType (undefined::sh) (undefined::Z) - = scanAllOp exe gamma aenv stream n m + -> Val aenv + -> Int -- output size of innermost dimension + -> Delayed (Array (sh, Int) e) + -> Par PTX (Future (Array (sh, Int) e)) +scanCore repr exe gamma aenv m input + | ArrayR (ShapeRsnoc ShapeRz) tp <- repr + = scanAllOp tp exe gamma aenv m input -- | otherwise - = scanDimOp exe gamma aenv stream sz m - + = scanDimOp repr exe gamma aenv m input +{-# INLINE scanAllOp #-} scanAllOp - :: forall aenv e. Elt e - => ExecutableR PTX + :: HasCallStack + => TypeR e + -> ExecutableR PTX -> Gamma aenv - -> Aval aenv - -> Stream - -> Int -- input size + -> Val aenv -> Int -- output size - -> LLVM PTX (Vector e) -scanAllOp exe gamma aenv stream n m = withExecutable exe $ \ptxExecutable -> do - let - k1 = ptxExecutable !# "scanP1" - k2 = ptxExecutable !# "scanP2" - k3 = ptxExecutable !# "scanP3" - -- - c = kernelThreadBlockSize k1 - s = n `multipleOf` c - -- - ptx <- gets llvmTarget - out <- allocateRemote (Z :. m) - - -- Step 1: Independent thread-block-wide scans of the input. Small arrays - -- which can be computed by a single thread block will require no - -- additional work. - tmp <- allocateRemote (Z :. s) :: LLVM PTX (Vector e) - liftIO $ executeOp ptx k1 gamma aenv stream (IE 0 s) (tmp, out) + -> Delayed (Vector e) + -> Par PTX (Future (Vector e)) +scanAllOp tp exe gamma aenv m input@(delayedShape -> ((), n)) = + withExecutable exe $ \ptxExecutable -> do + let + k1 = ptxExecutable !# "scanP1" + k2 = ptxExecutable !# "scanP2" + k3 = ptxExecutable !# "scanP3" + -- + c = kernelThreadBlockSize k1 + s = n `multipleOf` c + -- + repr = ArrayR dim1 tp + paramR = TupRsingle $ ParamRarray repr + paramsR1 = paramR `TupRpair` paramR `TupRpair` TupRsingle (ParamRmaybe $ ParamRarray repr) + paramsR3 = paramR `TupRpair` paramR `TupRpair` TupRsingle ParamRint + -- + future <- new + result <- allocateRemote repr ((), m) - -- Step 2: Multi-block reductions need to compute the per-block prefix, - -- then apply those values to the partial results. - when (s > 1) $ do - liftIO $ executeOp ptx k2 gamma aenv stream (IE 0 s) tmp - liftIO $ executeOp ptx k3 gamma aenv stream (IE 0 (s-1)) (tmp, out, c) + -- Step 1: Independent thread-block-wide scans of the input. Small arrays + -- which can be computed by a single thread block will require no + -- additional work. + tmp <- allocateRemote repr ((), s) + executeOp k1 gamma aenv dim1 ((), s) paramsR1 ((tmp, result), manifest input) - return out + -- Step 2: Multi-block reductions need to compute the per-block prefix, + -- then apply those values to the partial results. + when (s > 1) $ do + executeOp k2 gamma aenv dim1 ((), s) paramR tmp + executeOp k3 gamma aenv dim1 ((), s-1) paramsR3 ((tmp, result), c) + put future result + return future +{-# INLINE scanDimOp #-} scanDimOp - :: forall aenv sh e. (Shape sh, Elt e) - => ExecutableR PTX + :: HasCallStack + => ArrayR (Array (sh, Int) e) + -> ExecutableR PTX -> Gamma aenv - -> Aval aenv - -> Stream - -> sh + -> Val aenv -> Int - -> LLVM PTX (Array (sh:.Int) e) -scanDimOp exe gamma aenv stream sz m = withExecutable exe $ \ptxExecutable -> do - ptx <- gets llvmTarget - out <- allocateRemote (sz :. m) - liftIO $ executeOp ptx (ptxExecutable !# "scan") gamma aenv stream (IE 0 (size sz)) out - return out + -> Delayed (Array (sh, Int) e) + -> Par PTX (Future (Array (sh, Int) e)) +scanDimOp repr exe gamma aenv m input@(delayedShape -> (sz, _)) = + withExecutable exe $ \ptxExecutable -> do + let ArrayR (ShapeRsnoc shr') _ = repr + future <- new + result <- allocateRemote repr (sz, m) + let paramsR = TupRsingle (ParamRarray repr) `TupRpair` TupRsingle (ParamRmaybe $ ParamRarray repr) + executeOp (ptxExecutable !# "scan") gamma aenv dim1 ((), size shr' sz) paramsR (result, manifest input) + put future result + return future +{-# INLINE scan'Op #-} scan'Op - :: forall aenv sh e. (Shape sh, Elt e) - => ExecutableR PTX + :: HasCallStack + => ArrayR (Array (sh, Int) e) + -> ExecutableR PTX -> Gamma aenv - -> Aval aenv - -> Stream - -> sh :. Int - -> LLVM PTX (Array (sh:.Int) e, Array sh e) -scan'Op exe gamma aenv stream sh@(sz :. n) = + -> Val aenv + -> Delayed (Array (sh, Int) e) + -> Par PTX (Future (Array (sh, Int) e, Array sh e)) +scan'Op repr exe gamma aenv input@(delayedShape -> (sz, n)) = case n of - 0 -> do out <- allocateRemote (sz :. 0) - sum <- simpleNamed "generate" exe gamma aenv stream sz - return (out, sum) - _ -> scan'Core exe gamma aenv stream sh + 0 -> do + future <- new + result <- allocateRemote repr (sz, 0) + sums <- generateOp (reduceRank repr) exe gamma aenv sz + fork $ do sums' <- get sums + put future (result, sums') + return future + -- + _ -> scan'Core repr exe gamma aenv input +{-# INLINE scan'Core #-} scan'Core - :: forall aenv sh e. (Shape sh, Elt e) - => ExecutableR PTX + :: HasCallStack + => ArrayR (Array (sh, Int) e) + -> ExecutableR PTX -> Gamma aenv - -> Aval aenv - -> Stream - -> sh :. Int - -> LLVM PTX (Array (sh:.Int) e, Array sh e) -scan'Core exe gamma aenv stream sh - | Just Refl <- matchShapeType (undefined::sh) (undefined::Z) - = scan'AllOp exe gamma aenv stream sh + -> Val aenv + -> Delayed (Array (sh, Int) e) + -> Par PTX (Future (Array (sh, Int) e, Array sh e)) +scan'Core repr exe gamma aenv input + | ArrayR (ShapeRsnoc ShapeRz) tp <- repr + = scan'AllOp tp exe gamma aenv input -- | otherwise - = scan'DimOp exe gamma aenv stream sh + = scan'DimOp repr exe gamma aenv input +{-# INLINE scan'AllOp #-} scan'AllOp - :: forall aenv e. Elt e - => ExecutableR PTX + :: HasCallStack + => TypeR e + -> ExecutableR PTX -> Gamma aenv - -> Aval aenv - -> Stream - -> DIM1 - -> LLVM PTX (Vector e, Scalar e) -scan'AllOp exe gamma aenv stream (Z :. n) = withExecutable exe $ \ptxExecutable -> do - let - k1 = ptxExecutable !# "scanP1" - k2 = ptxExecutable !# "scanP2" - k3 = ptxExecutable !# "scanP3" - -- - c = kernelThreadBlockSize k1 - s = n `multipleOf` c - -- - ptx <- gets llvmTarget - out <- allocateRemote (Z :. n) - tmp <- allocateRemote (Z :. s) :: LLVM PTX (Vector e) + -> Val aenv + -> Delayed (Vector e) + -> Par PTX (Future (Vector e, Scalar e)) +scan'AllOp tp exe gamma aenv input@(delayedShape -> ((), n)) = + withExecutable exe $ \ptxExecutable -> do + let + repr = ArrayR dim1 tp + paramRdim0 = TupRsingle $ ParamRarray $ ArrayR dim0 tp + paramRdim1 = TupRsingle $ ParamRarray repr + k1 = ptxExecutable !# "scanP1" + k2 = ptxExecutable !# "scanP2" + k3 = ptxExecutable !# "scanP3" + -- + c = kernelThreadBlockSize k1 + s = n `multipleOf` c + -- + future <- new + result <- allocateRemote repr ((), n) + tmp <- allocateRemote repr ((), s) - -- Step 1: independent thread-block-wide scans. Each block stores its partial - -- sum to a temporary array. - liftIO $ executeOp ptx k1 gamma aenv stream (IE 0 s) (tmp, out) + -- Step 1: independent thread-block-wide scans. Each block stores its partial + -- sum to a temporary array. + let paramsR1 = paramRdim1 `TupRpair` paramRdim1 `TupRpair` TupRsingle (ParamRmaybe $ ParamRarray repr) + executeOp k1 gamma aenv dim1 ((), s) paramsR1 ((tmp, result), manifest input) - -- If this was a small array that was processed by a single thread block then - -- we are done, otherwise compute the per-block prefix and apply those values - -- to the partial results. - if s == 1 - then case tmp of - Array _ ad -> return (out, Array () ad) - else do - sum <- allocateRemote Z - liftIO $ executeOp ptx k2 gamma aenv stream (IE 0 s) (tmp, sum) - liftIO $ executeOp ptx k3 gamma aenv stream (IE 0 (s-1)) (tmp, out, c) - return (out, sum) + -- If this was a small array that was processed by a single thread block then + -- we are done, otherwise compute the per-block prefix and apply those values + -- to the partial results. + if s == 1 + then + case tmp of + Array _ ad -> put future (result, Array () ad) + else do + sums <- allocateRemote (ArrayR dim0 tp) () + let paramsR2 = paramRdim1 `TupRpair` paramRdim0 + let paramsR3 = paramRdim1 `TupRpair` paramRdim1 `TupRpair` TupRsingle ParamRint + executeOp k2 gamma aenv dim1 ((), s) paramsR2 (tmp, sums) + executeOp k3 gamma aenv dim1 ((), s-1) paramsR3 ((tmp, result), c) + put future (result, sums) + -- + return future +{-# INLINE scan'DimOp #-} scan'DimOp - :: forall aenv sh e. (Shape sh, Elt e) - => ExecutableR PTX + :: HasCallStack + => ArrayR (Array (sh, Int) e) + -> ExecutableR PTX -> Gamma aenv - -> Aval aenv - -> Stream - -> sh :. Int - -> LLVM PTX (Array (sh:.Int) e, Array sh e) -scan'DimOp exe gamma aenv stream sh@(sz :. _) = withExecutable exe $ \ptxExecutable -> do - ptx <- gets llvmTarget - out <- allocateRemote sh - sum <- allocateRemote sz - liftIO $ executeOp ptx (ptxExecutable !# "scan") gamma aenv stream (IE 0 (size sz)) (out,sum) - return (out,sum) + -> Val aenv + -> Delayed (Array (sh, Int) e) + -> Par PTX (Future (Array (sh, Int) e, Array sh e)) +scan'DimOp repr@(ArrayR (ShapeRsnoc shr') _) exe gamma aenv input@(delayedShape -> sh@(sz, _)) = + withExecutable exe $ \ptxExecutable -> do + future <- new + result <- allocateRemote repr sh + sums <- allocateRemote (reduceRank repr) sz + let paramsR = TupRsingle (ParamRarray repr) `TupRpair` TupRsingle (ParamRarray $ reduceRank repr) `TupRpair` TupRsingle (ParamRmaybe $ ParamRarray repr) + executeOp (ptxExecutable !# "scan") gamma aenv dim1 ((), size shr' sz) paramsR ((result, sums), manifest input) + put future (result, sums) + return future +{-# INLINE permuteOp #-} permuteOp - :: (Shape sh, Shape sh', Elt e) - => ExecutableR PTX + :: HasCallStack + => Bool + -> ArrayR (Array sh e) + -> ShapeR sh' + -> ExecutableR PTX -> Gamma aenv - -> Aval aenv - -> Stream - -> Bool - -> sh + -> Val aenv -> Array sh' e - -> LLVM PTX (Array sh' e) -permuteOp exe gamma aenv stream inplace shIn dfs = withExecutable exe $ \ptxExecutable -> do - let - n = size shIn - m = size (shape dfs) - kernel = case functionTable ptxExecutable of - k:_ -> k - _ -> $internalError "permute" "no kernels found" - -- - ptx <- gets llvmTarget - out <- if inplace - then return dfs - else cloneArrayAsync stream dfs - -- - case kernelName kernel of - "permute_rmw" -> liftIO $ executeOp ptx kernel gamma aenv stream (IE 0 n) out - "permute_mutex" -> do - barrier@(Array _ ad) <- allocateRemote (Z :. m) :: LLVM PTX (Vector Word32) - memsetArrayAsync stream m 0 ad - liftIO $ executeOp ptx kernel gamma aenv stream (IE 0 n) (out, barrier) - _ -> $internalError "permute" "unexpected kernel image" - -- - return out + -> Delayed (Array sh e) + -> Par PTX (Future (Array sh' e)) +permuteOp inplace repr@(ArrayR shr tp) shr' exe gamma aenv defaults@(shape -> shOut) input@(delayedShape -> shIn) = + withExecutable exe $ \ptxExecutable -> do + let + n = size shr shIn + m = size shr' shOut + repr' = ArrayR shr' tp + reprLock = ArrayR dim1 $ TupRsingle $ scalarTypeWord32 + paramR = TupRsingle $ ParamRmaybe $ ParamRarray repr + paramR' = TupRsingle $ ParamRarray repr' + kernel = case functionTable ptxExecutable of + k:_ -> k + _ -> internalError "no kernels found" + -- + future <- new + result <- if inplace + then Debug.trace Debug.dump_exec "exec: permute/inplace" $ return defaults + else Debug.trace Debug.dump_exec "exec: permute/clone" $ get =<< cloneArrayAsync repr' defaults + -- + case kernelName kernel of + -- execute directly using atomic operations + "permute_rmw" -> + let paramsR = paramR' `TupRpair` paramR + in executeOp kernel gamma aenv dim1 ((), n) paramsR (result, manifest input) + -- a temporary array is required for spin-locks around the critical section + "permute_mutex" -> do + barrier <- new :: Par PTX (Future (Vector Word32)) + Array _ ad <- allocateRemote reprLock ((), m) + fork $ do fill <- memsetArrayAsync (NumSingleType $ IntegralNumType TypeWord32) m 0 ad + put barrier . Array ((), m) =<< get fill + -- + let paramsR = paramR' `TupRpair` TupRsingle (ParamRfuture $ ParamRarray reprLock) `TupRpair` paramR + executeOp kernel gamma aenv dim1 ((), n) paramsR ((result, barrier), manifest input) + _ -> internalError "unexpected kernel image" + -- + put future result + return future + + +{-# INLINE stencil1Op #-} +stencil1Op + :: HasCallStack + => TypeR a + -> ArrayR (Array sh b) + -> sh + -> ExecutableR PTX + -> Gamma aenv + -> Val aenv + -> Delayed (Array sh a) + -> Par PTX (Future (Array sh b)) +stencil1Op tp repr@(ArrayR shr _) halo exe gamma aenv input@(delayedShape -> sh) = + stencilCore repr exe gamma aenv halo sh paramsR (manifest input) + where paramsR = TupRsingle $ ParamRmaybe $ ParamRarray $ ArrayR shr tp + -- Using the defaulting instances for stencil operations (for now). -- +{-# INLINE stencil2Op #-} stencil2Op - :: (Shape sh, Elt e) - => ExecutableR PTX + :: HasCallStack + => TypeR a + -> TypeR b + -> ArrayR (Array sh c) + -> sh + -> ExecutableR PTX -> Gamma aenv - -> Aval aenv - -> Stream + -> Val aenv + -> Delayed (Array sh a) + -> Delayed (Array sh b) + -> Par PTX (Future (Array sh c)) +stencil2Op tpA tpB repr@(ArrayR shr _) halo exe gamma aenv input1@(delayedShape -> sh1) input2@(delayedShape -> sh2) = + stencilCore repr exe gamma aenv halo (intersect (arrayRshape repr) sh1 sh2) paramsR (manifest input1, manifest input2) + where paramsR = TupRsingle (ParamRmaybe $ ParamRarray $ ArrayR shr tpA) `TupRpair` TupRsingle (ParamRmaybe $ ParamRarray $ ArrayR shr tpB) + +{-# INLINE stencilCore #-} +stencilCore + :: forall aenv sh e params. HasCallStack + => ArrayR (Array sh e) + -> ExecutableR PTX + -> Gamma aenv + -> Val aenv + -> sh -- border dimensions (i.e. index of first interior element) + -> sh -- output array size + -> ParamsR PTX params + -> params + -> Par PTX (Future (Array sh e)) +stencilCore repr@(ArrayR shr _) exe gamma aenv halo shOut paramsR params = + withExecutable exe $ \ptxExecutable -> do + let + inside = ptxExecutable !# "stencil_inside" + border = ptxExecutable !# "stencil_border" + + shIn :: sh + shIn = trav (\x y -> x - 2*y) shOut halo + + trav :: (Int -> Int -> Int) -> sh -> sh -> sh + trav f a b = go (arrayRshape repr) a b + where + go :: ShapeR t -> t -> t -> t + go ShapeRz () () = () + go (ShapeRsnoc shr') (xa,xb) (ya,yb) = (go shr' xa ya, f xb yb) + -- + future <- new + result <- allocateRemote repr shOut + parent <- asks ptxStream + + -- interior (no bounds checking) + let paramsRinside = TupRsingle (ParamRshape shr) `TupRpair` TupRsingle (ParamRarray repr) `TupRpair` paramsR + executeOp inside gamma aenv shr shIn paramsRinside ((shIn, result), params) + + -- halo regions (bounds checking) + -- executed in separate streams so that they might overlap the main stencil + -- and each other, as individually they will not saturate the device + forM_ (stencilBorders (arrayRshape repr) shOut halo) $ \(u, v) -> + fork $ do + -- launch in a separate stream + let sh = trav (-) v u + let paramsRborder = TupRsingle (ParamRshape shr) `TupRpair` TupRsingle (ParamRshape shr) + `TupRpair` TupRsingle (ParamRarray repr) + `TupRpair` paramsR + executeOp border gamma aenv shr sh paramsRborder (((u, sh), result), params) + + -- synchronisation with main stream + child <- asks ptxStream + event <- liftPar (Event.waypoint child) + ready <- liftIO (Event.query event) + if ready then return () + else liftIO (Event.after event parent) + + put future result + return future + +-- Compute the stencil border regions, where we may need to evaluate the +-- boundary conditions. +-- +{-# INLINE stencilBorders #-} +stencilBorders + :: forall sh. HasCallStack + => ShapeR sh -> sh -> sh - -> LLVM PTX (Array sh e) -stencil2Op exe gamma aenv stream sh1 sh2 = - simpleOp exe gamma aenv stream (sh1 `intersect` sh2) + -> [(sh, sh)] +stencilBorders shr sh halo = [ face i | i <- [0 .. (2 * rank shr - 1)] ] + where + face :: Int -> (sh, sh) + face n = go n shr sh halo + go :: Int -> ShapeR t -> t -> t -> (t, t) + go _ ShapeRz () () = ((), ()) + go n (ShapeRsnoc shr') (sha, sza) (shb, szb) + = let + (sha', shb') = go (n-2) shr' sha shb + (sza', szb') + | n < 0 = (0, sza) + | n == 0 = (0, szb) + | n == 1 = (sza-szb, sza) + | otherwise = (szb, sza-szb) + in + ((sha', sza'), (shb', szb')) + -- Foreign functions -- +{-# INLINE aforeignOp #-} aforeignOp - :: (Arrays as, Arrays bs) + :: HasCallStack => String - -> (Stream -> as -> LLVM PTX bs) - -> Stream + -> ArraysR as + -> ArraysR bs + -> (as -> Par PTX (Future bs)) -> as - -> LLVM PTX bs -aforeignOp name asm stream arr = - Debug.monitorProcTime query msg (Just (unsafeGetValue stream)) $ - asm stream arr + -> Par PTX (Future bs) +aforeignOp name _ _ asm arr = do + stream <- asks ptxStream + Debug.monitorProcTime query msg (Just (unsafeGetValue stream)) (asm arr) where query = if Debug.monitoringIsEnabled then return True @@ -539,48 +788,60 @@ -- Skeleton execution -- ------------------ --- TODO: Calculate this from the device properties, say [a multiple of] the --- maximum number of in-flight threads that the device supports. --- -defaultPPT :: Int -defaultPPT = 32768 - -- | Retrieve the named kernel -- -(!#) :: FunctionTable -> ShortByteString -> Kernel +(!#) :: HasCallStack => FunctionTable -> ShortByteString -> Kernel (!#) exe name - = fromMaybe ($internalError "lookupFunction" ("function not found: " ++ unpack name)) + = fromMaybe (internalError ("function not found: " ++ unpack name)) $ lookupKernel name exe lookupKernel :: ShortByteString -> FunctionTable -> Maybe Kernel lookupKernel name ptxExecutable = find (\k -> kernelName k == name) (functionTable ptxExecutable) +delayedShape :: Delayed (Array sh e) -> sh +delayedShape (Delayed sh) = sh +delayedShape (Manifest a) = shape a +manifest :: Delayed (Array sh e) -> Maybe (Array sh e) +manifest (Manifest a) = Just a +manifest Delayed{} = Nothing + +-- | Execute some operation with the supplied executable functions +-- +withExecutable :: HasCallStack => ExecutableR PTX -> (FunctionTable -> Par PTX b) -> Par PTX b +withExecutable PTXR{..} f = + local (\(s,_) -> (s,Just ptxExecutable)) $ do + r <- f (unsafeGetValue ptxExecutable) + liftIO $ touchLifetime ptxExecutable + return r + + -- Execute the function implementing this kernel. -- executeOp - :: Marshalable args - => PTX - -> Kernel + :: HasCallStack + => Kernel -> Gamma aenv - -> Aval aenv - -> Stream - -> Range - -> args - -> IO () -executeOp ptx@PTX{..} kernel@Kernel{..} gamma aenv stream r args = - runExecutable fillP kernelName defaultPPT r $ \start end _ -> do - argv <- marshal ptx stream (start, end, args, (gamma,aenv)) - launch kernel stream (end-start) argv + -> Val aenv + -> ShapeR sh + -> sh + -> ParamsR PTX params + -> params + -> Par PTX () +executeOp kernel gamma aenv shr sh paramsR params = + let n = size shr sh + in when (n > 0) $ do + stream <- asks ptxStream + argv <- marshalParams' @PTX (paramsR `TupRpair` TupRsingle (ParamRenv gamma)) (params, aenv) + liftIO $ launch kernel stream n $ DL.toList argv -- Execute a device function with the given thread configuration and function -- parameters. -- -launch :: Kernel -> Stream -> Int -> [CUDA.FunParam] -> IO () +launch :: HasCallStack => Kernel -> Stream -> Int -> [CUDA.FunParam] -> IO () launch Kernel{..} stream n args = - when (n > 0) $ withLifetime stream $ \st -> Debug.monitorProcTime query msg (Just st) $ CUDA.launchKernel kernelFun grid cta smem (Just st) args diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Execute/Async.hs b/src/Data/Array/Accelerate/LLVM/PTX/Execute/Async.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/Execute/Async.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/Execute/Async.hs @@ -1,63 +1,163 @@ -{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeSynonymInstances #-} {-# OPTIONS_GHC -fno-warn-orphans #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Execute.Async --- Copyright : [2014..2017] Trevor L. McDonell --- [2014..2014] Vinod Grover (NVIDIA Corporation) +-- Copyright : [2014..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- module Data.Array.Accelerate.LLVM.PTX.Execute.Async ( - Async, Stream, Event, module Data.Array.Accelerate.LLVM.Execute.Async, + module Data.Array.Accelerate.LLVM.PTX.Execute.Async, ) where -- accelerate -import Data.Array.Accelerate.LLVM.Execute.Async hiding ( Async ) -import qualified Data.Array.Accelerate.LLVM.Execute.Async as A +import Data.Array.Accelerate.Error +import Data.Array.Accelerate.Lifetime +import Data.Array.Accelerate.LLVM.Execute.Async +import Data.Array.Accelerate.LLVM.State + import Data.Array.Accelerate.LLVM.PTX.Target -import Data.Array.Accelerate.LLVM.PTX.Execute.Event ( Event ) -import Data.Array.Accelerate.LLVM.PTX.Execute.Stream ( Stream ) -import qualified Data.Array.Accelerate.LLVM.PTX.Execute.Event as Event -import qualified Data.Array.Accelerate.LLVM.PTX.Execute.Stream as Stream +import Data.Array.Accelerate.LLVM.PTX.Execute.Event ( Event ) +import Data.Array.Accelerate.LLVM.PTX.Execute.Stream ( Stream ) +import Data.Array.Accelerate.LLVM.PTX.Link.Object ( FunctionTable ) +import qualified Data.Array.Accelerate.LLVM.PTX.Execute.Event as Event +import qualified Data.Array.Accelerate.LLVM.PTX.Execute.Stream as Stream -- standard library import Control.Monad.State +import Control.Monad.Reader +import Data.IORef --- Asynchronous arrays in the CUDA backend are tagged with an Event that will be --- filled once the kernel implementing that array has completed. +-- | Evaluate a parallel computation -- -type Async a = AsyncR PTX a +{-# INLINE evalPar #-} +evalPar :: Par PTX a -> LLVM PTX a +evalPar p = do + s <- Stream.create + r <- runReaderT (runPar p) (s, Nothing) + return r -instance A.Async PTX where - type StreamR PTX = Stream - type EventR PTX = Event - {-# INLINEABLE fork #-} - fork = - Stream.create +type ParState = (Stream, Maybe (Lifetime FunctionTable)) - {-# INLINEABLE join #-} - join stream = - liftIO $! Stream.destroy stream +ptxStream :: ParState -> Stream +ptxStream = fst - {-# INLINEABLE checkpoint #-} - checkpoint stream = - Event.waypoint stream +ptxKernel :: ParState -> Maybe (Lifetime FunctionTable) +ptxKernel = snd - {-# INLINEABLE after #-} - after stream event = - liftIO $! Event.after event stream +-- Implementation +-- -------------- + +data Future a = Future {-# UNPACK #-} !(IORef (IVar a)) + +data IVar a + = Full !a + | Pending {-# UNPACK #-} !Event !(Maybe (Lifetime FunctionTable)) !a + | Empty + + +instance Async PTX where + type FutureR PTX = Future + + newtype Par PTX a = Par { runPar :: ReaderT ParState (LLVM PTX) a } + deriving ( Functor, Applicative, Monad, MonadIO, MonadReader ParState, MonadState PTX ) + + {-# INLINEABLE new #-} + {-# INLINEABLE newFull #-} + new = Future <$> liftIO (newIORef Empty) + newFull v = Future <$> liftIO (newIORef (Full v)) + + {-# INLINEABLE spawn #-} + spawn m = do + s' <- liftPar Stream.create + r <- local (const (s', Nothing)) m + liftIO (Stream.destroy s') + return r + + {-# INLINEABLE fork #-} + fork m = do + s' <- liftPar (Stream.create) + () <- local (const (s', Nothing)) m + liftIO (Stream.destroy s') + + -- When we call 'put' the actual work may not have been evaluated yet; get + -- a new event in the current execution stream and once that is filled we can + -- transition the IVar to Full. + -- + {-# INLINEABLE put #-} + put (Future ref) v = do + stream <- asks ptxStream + kernel <- asks ptxKernel + event <- liftPar (Event.waypoint stream) + ready <- liftIO (Event.query event) + liftIO . modifyIORef' ref $ \case + Empty -> if ready then Full v + else Pending event kernel v + _ -> internalError "multiple put" + + -- Get the value of Future. Since the actual cross-stream synchronisation + -- happens on the device, we should never have to block/reschedule the main + -- thread waiting on a value; if we get an empty IVar at this point, something + -- has gone wrong. + -- + {-# INLINEABLE get #-} + get (Future ref) = do + stream <- asks ptxStream + liftIO $ do + ivar <- readIORef ref + case ivar of + Full v -> return v + Pending event k v -> do + ready <- Event.query event + if ready + then do + writeIORef ref (Full v) + case k of + Just f -> touchLifetime f + Nothing -> return () + else + Event.after event stream + return v + Empty -> internalError "blocked on an IVar" + {-# INLINEABLE block #-} - block event = - liftIO $! Event.block event + block = liftIO . wait + + {-# INLINE liftPar #-} + liftPar = Par . lift + + +-- | Block the calling _host_ thread until the value offered by the future is +-- available. +-- +{-# INLINEABLE wait #-} +wait :: Future a -> IO a +wait (Future ref) = do + ivar <- readIORef ref + case ivar of + Full v -> return v + Pending event k v -> do + Event.block event + writeIORef ref (Full v) + case k of + Just f -> touchLifetime f + Nothing -> return () + return v + Empty -> internalError "blocked on an IVar" diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Execute/Environment.hs b/src/Data/Array/Accelerate/LLVM/PTX/Execute/Environment.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/Execute/Environment.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/Execute/Environment.hs @@ -1,22 +1,22 @@ -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Execute.Environment --- Copyright : [2014..2017] Trevor L. McDonell --- [2014..2014] Vinod Grover (NVIDIA Corporation) +-- Copyright : [2014..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- module Data.Array.Accelerate.LLVM.PTX.Execute.Environment ( - Aval, aprj + module Data.Array.Accelerate.LLVM.Execute.Environment, + module Data.Array.Accelerate.LLVM.PTX.Execute.Environment, ) where import Data.Array.Accelerate.LLVM.PTX.Target import Data.Array.Accelerate.LLVM.Execute.Environment -type Aval = AvalR PTX +type Val = ValR PTX diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Execute/Event.hs b/src/Data/Array/Accelerate/LLVM/PTX/Execute/Event.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/Execute/Event.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/Execute/Event.hs @@ -1,11 +1,10 @@ {-# LANGUAGE NamedFieldPuns #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Execute.Event --- Copyright : [2014..2017] Trevor L. McDonell --- [2014..2014] Vinod Grover (NVIDIA Corporation) +-- Copyright : [2014..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Execute/Event.hs-boot b/src/Data/Array/Accelerate/LLVM/PTX/Execute/Event.hs-boot --- a/src/Data/Array/Accelerate/LLVM/PTX/Execute/Event.hs-boot +++ b/src/Data/Array/Accelerate/LLVM/PTX/Execute/Event.hs-boot @@ -1,9 +1,9 @@ -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Execute.Event-boot --- Copyright : [2016..2017] Trevor L. McDonell +-- Copyright : [2016..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Execute/Marshal.hs b/src/Data/Array/Accelerate/LLVM/PTX/Execute/Marshal.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/Execute/Marshal.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/Execute/Marshal.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE CPP #-} {-# LANGUAGE BangPatterns #-} {-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE FlexibleContexts #-} @@ -6,150 +5,59 @@ {-# LANGUAGE GADTs #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeFamilies #-} {-# OPTIONS_GHC -fno-warn-orphans #-} -#if __GLASGOW_HASKELL__ <= 708 -{-# LANGUAGE OverlappingInstances #-} -{-# OPTIONS_GHC -fno-warn-unrecognised-pragmas #-} -#endif -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Execute.Marshal --- Copyright : [2014..2017] Trevor L. McDonell --- [2014..2014] Vinod Grover (NVIDIA Corporation) +-- Copyright : [2014..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- module Data.Array.Accelerate.LLVM.PTX.Execute.Marshal ( - Marshalable, M.marshal + module Data.Array.Accelerate.LLVM.Execute.Marshal ) where --- accelerate -import Data.Array.Accelerate.LLVM.CodeGen.Environment ( Gamma, Idx'(..) ) -import qualified Data.Array.Accelerate.LLVM.Execute.Marshal as M +import Data.Array.Accelerate.LLVM.State +import Data.Array.Accelerate.LLVM.Execute.Marshal -import Data.Array.Accelerate.LLVM.PTX.State import Data.Array.Accelerate.LLVM.PTX.Target -import Data.Array.Accelerate.LLVM.PTX.Array.Data -import Data.Array.Accelerate.LLVM.PTX.Execute.Async ( Async, AsyncR(..) ) -import Data.Array.Accelerate.LLVM.PTX.Execute.Event ( after ) -import Data.Array.Accelerate.LLVM.PTX.Execute.Environment +import Data.Array.Accelerate.LLVM.PTX.Execute.Async import qualified Data.Array.Accelerate.LLVM.PTX.Array.Prim as Prim --- cuda +import Data.Array.Accelerate.Type +import Data.Array.Accelerate.Array.Data + import qualified Foreign.CUDA.Driver as CUDA --- libraries -import Control.Monad -import Data.Int -import Data.DList ( DList ) -import Data.Typeable -import Foreign.Ptr -import Foreign.Storable ( Storable ) import qualified Data.DList as DL -import qualified Data.IntMap as IM --- Instances for the PTX backend --- -type Marshalable args = M.Marshalable PTX args -type instance M.ArgR PTX = CUDA.FunParam - - --- Instances for handling concrete types in this backend, namely shapes and --- array data. --- -instance M.Marshalable PTX Int where - marshal' _ _ x = return $ DL.singleton (CUDA.VArg x) - -instance M.Marshalable PTX Int32 where - marshal' _ _ x = return $ DL.singleton (CUDA.VArg x) - -instance {-# OVERLAPS #-} M.Marshalable PTX (Gamma aenv, Aval aenv) where - marshal' ptx stream (gamma, aenv) - = fmap DL.concat - $ mapM (\(_, Idx' idx) -> M.marshal' ptx stream =<< sync (aprj idx aenv)) (IM.elems gamma) - where - -- HAXORZ~ D: - -- - -- The 'Async' class functions need to run in the LLVM monad, but the - -- marshalling functions must run in IO because they will be executed in - -- the lower-level scheduling code. - -- - -- We hack around this impedance mismatch by calling the 'after' - -- implementation directly. - -- - sync :: Async a -> IO a - sync (AsyncR event arr) = after event stream >> return arr - -instance ArrayElt e => M.Marshalable PTX (ArrayData e) where - marshal' ptx _ adata = do - let marshalP :: forall e' a. (ArrayElt e', ArrayPtrs e' ~ Ptr a, Typeable e', Typeable a, Storable a) - => ArrayData e' - -> IO (DList CUDA.FunParam) - marshalP ad = - fmap (DL.singleton . CUDA.VArg) - (unsafeGetDevicePtr ptx ad :: IO (CUDA.DevicePtr a)) - - marshalR :: ArrayEltR e' -> ArrayData e' -> IO (DList CUDA.FunParam) - marshalR ArrayEltRunit _ = return DL.empty - marshalR (ArrayEltRpair aeR1 aeR2) ad = - return DL.append `ap` marshalR aeR1 (fstArrayData ad) - `ap` marshalR aeR2 (sndArrayData ad) - -- - marshalR (ArrayEltRvec2 aeR) (AD_V2 ad) = marshalR aeR ad - marshalR (ArrayEltRvec3 aeR) (AD_V3 ad) = marshalR aeR ad - marshalR (ArrayEltRvec4 aeR) (AD_V4 ad) = marshalR aeR ad - marshalR (ArrayEltRvec8 aeR) (AD_V8 ad) = marshalR aeR ad - marshalR (ArrayEltRvec16 aeR) (AD_V16 ad) = marshalR aeR ad - -- - marshalR ArrayEltRint ad = marshalP ad - marshalR ArrayEltRint8 ad = marshalP ad - marshalR ArrayEltRint16 ad = marshalP ad - marshalR ArrayEltRint32 ad = marshalP ad - marshalR ArrayEltRint64 ad = marshalP ad - marshalR ArrayEltRword ad = marshalP ad - marshalR ArrayEltRword8 ad = marshalP ad - marshalR ArrayEltRword16 ad = marshalP ad - marshalR ArrayEltRword32 ad = marshalP ad - marshalR ArrayEltRword64 ad = marshalP ad - marshalR ArrayEltRhalf ad = marshalP ad - marshalR ArrayEltRfloat ad = marshalP ad - marshalR ArrayEltRdouble ad = marshalP ad - marshalR ArrayEltRchar ad = marshalP ad - marshalR ArrayEltRcshort ad = marshalP ad - marshalR ArrayEltRcushort ad = marshalP ad - marshalR ArrayEltRcint ad = marshalP ad - marshalR ArrayEltRcuint ad = marshalP ad - marshalR ArrayEltRclong ad = marshalP ad - marshalR ArrayEltRculong ad = marshalP ad - marshalR ArrayEltRcllong ad = marshalP ad - marshalR ArrayEltRcullong ad = marshalP ad - marshalR ArrayEltRcchar ad = marshalP ad - marshalR ArrayEltRcschar ad = marshalP ad - marshalR ArrayEltRcuchar ad = marshalP ad - marshalR ArrayEltRcfloat ad = marshalP ad - marshalR ArrayEltRcdouble ad = marshalP ad - marshalR ArrayEltRbool ad = marshalP ad - - marshalR arrayElt adata +instance Marshal PTX where + type ArgR PTX = CUDA.FunParam + marshalInt = CUDA.VArg + marshalScalarData' t + | SingleArrayDict <- singleArrayDict t + = liftPar . fmap (DL.singleton . CUDA.VArg) . unsafeGetDevicePtr t -- TODO FIXME !!! -- --- We will probably need to change marshal to be a bracketed function. We may --- also want to reconsider whether to continue to restrict it to IO. +-- We will probably need to change marshal to be a bracketed function, so that +-- the garbage collector does not try to evict the array in the middle of +-- a computation. -- unsafeGetDevicePtr - :: (ArrayElt e, ArrayPtrs e ~ Ptr a, Typeable e, Typeable a, Storable a) - => PTX + :: SingleType e -> ArrayData e - -> IO (CUDA.DevicePtr a) -unsafeGetDevicePtr !ptx !ad = - evalPTX ptx $ Prim.withDevicePtr ad (\p -> return (Nothing,p)) + -> LLVM PTX (CUDA.DevicePtr (ScalarArrayDataR e)) +unsafeGetDevicePtr !t !ad = + Prim.withDevicePtr t ad (\p -> return (Nothing, p)) diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Execute/Stream.hs b/src/Data/Array/Accelerate/LLVM/PTX/Execute/Stream.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/Execute/Stream.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/Execute/Stream.hs @@ -3,11 +3,10 @@ {-# LANGUAGE RecordWildCards #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Execute.Stream --- Copyright : [2014..2017] Trevor L. McDonell --- [2014..2014] Vinod Grover (NVIDIA Corporation) +-- Copyright : [2014..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- @@ -63,7 +62,6 @@ -> (Event -> a -> LLVM PTX b) -> LLVM PTX b streaming !action !after = do - PTX{..} <- gets llvmTarget stream <- create first <- action stream end <- Event.waypoint stream diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Execute/Stream.hs-boot b/src/Data/Array/Accelerate/LLVM/PTX/Execute/Stream.hs-boot --- a/src/Data/Array/Accelerate/LLVM/PTX/Execute/Stream.hs-boot +++ b/src/Data/Array/Accelerate/LLVM/PTX/Execute/Stream.hs-boot @@ -1,9 +1,9 @@ -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Execute.Stream-boot --- Copyright : [2016..2017] Trevor L. McDonell +-- Copyright : [2016..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Execute/Stream/Reservoir.hs b/src/Data/Array/Accelerate/LLVM/PTX/Execute/Stream/Reservoir.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/Execute/Stream/Reservoir.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/Execute/Stream/Reservoir.hs @@ -1,10 +1,10 @@ {-# LANGUAGE BangPatterns #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Execute.Stream.Reservoir --- Copyright : [2016..2017] Trevor L. McDonell +-- Copyright : [2016..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Foreign.hs b/src/Data/Array/Accelerate/LLVM/PTX/Foreign.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/Foreign.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/Foreign.hs @@ -2,13 +2,14 @@ {-# LANGUAGE GADTs #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeApplications #-} {-# OPTIONS_GHC -fno-warn-orphans #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Foreign --- Copyright : [2016..2017] Trevor L. McDonell +-- Copyright : [2016..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- @@ -27,10 +28,12 @@ withDevicePtr, module Data.Array.Accelerate.LLVM.PTX.Array.Data, module Data.Array.Accelerate.LLVM.PTX.Execute.Async, + module Data.Array.Accelerate.LLVM.PTX.Execute.Event, + module Data.Array.Accelerate.LLVM.PTX.Execute.Stream, ) where -import qualified Data.Array.Accelerate.Array.Sugar as S +import qualified Data.Array.Accelerate.Sugar.Foreign as S import Data.Array.Accelerate.LLVM.State import Data.Array.Accelerate.LLVM.CodeGen.Sugar @@ -41,19 +44,23 @@ import Data.Array.Accelerate.LLVM.PTX.Context import Data.Array.Accelerate.LLVM.PTX.Execute.Async import Data.Array.Accelerate.LLVM.PTX.Target +import Data.Array.Accelerate.LLVM.PTX.Execute.Stream ( Stream ) +import Data.Array.Accelerate.LLVM.PTX.Execute.Event ( Event, waypoint, query ) import Control.Monad.State import Data.Typeable instance Foreign PTX where - foreignAcc _ (ff :: asm (a -> b)) - | Just (ForeignAcc _ asm :: ForeignAcc (a -> b)) <- cast ff = Just asm - | otherwise = Nothing + foreignAcc (ff :: asm (a -> b)) + | Just Refl <- eqT @asm @ForeignAcc + , ForeignAcc _ asm <- ff = Just asm + | otherwise = Nothing - foreignExp _ (ff :: asm (x -> y)) - | Just (ForeignExp _ asm :: ForeignExp (x -> y)) <- cast ff = Just asm - | otherwise = Nothing + foreignExp (ff :: asm (x -> y)) + | Just Refl <- eqT @asm @ForeignExp + , ForeignExp _ asm <- ff = Just asm + | otherwise = Nothing instance S.Foreign ForeignAcc where strForeign (ForeignAcc s _) = s @@ -66,7 +73,7 @@ -- data ForeignAcc f where ForeignAcc :: String - -> (Stream -> a -> LLVM PTX b) + -> (a -> Par PTX (Future b)) -> ForeignAcc (a -> b) -- Foreign expressions in the PTX backend. diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Link.hs b/src/Data/Array/Accelerate/LLVM/PTX/Link.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/Link.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/Link.hs @@ -4,10 +4,10 @@ {-# OPTIONS_GHC -fno-warn-orphans #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Link --- Copyright : [2017] Trevor L. McDonell +-- Copyright : [2017..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- @@ -16,7 +16,6 @@ module Data.Array.Accelerate.LLVM.Link, ExecutableR(..), FunctionTable(..), Kernel(..), ObjectCode, - withExecutable, linkFunctionQ, ) where @@ -71,7 +70,7 @@ -- Finalise the module by unloading it from the CUDA context addFinalizer oc $ do - Debug.traceIO Debug.dump_gc ("gc: unload module: " ++ show nm) + Debug.traceIO Debug.dump_ld ("ld: unload module: " ++ show nm) withContext (ptxContext target) (CUDA.unload mdl) return (nm, oc) @@ -98,11 +97,7 @@ -> LaunchConfig -> IO (Kernel, Q (TExp (Int -> Int))) linkFunctionQ mdl name configure = do -#if MIN_VERSION_nvvm(0,9,0) f <- CUDA.getFun mdl name -#else - f <- CUDA.getFun mdl (unpack name) -#endif regs <- CUDA.requires f CUDA.NumRegs ssmem <- CUDA.requires f CUDA.SharedSizeBytes cmem <- CUDA.requires f CUDA.ConstSizeBytes @@ -122,17 +117,8 @@ (CUDA.activeWarps occ) (CUDA.activeThreadBlocks occ) - Debug.traceIO Debug.dump_cc (printf "cc: %s\n ... %s" msg1 msg2) + Debug.traceIO Debug.dump_cc (printf "cc: %s\n %s" msg1 msg2) return (Kernel name f dsmem cta grid, gridQ) - - --- | Execute some operation with the supplied executable functions --- -withExecutable :: ExecutableR PTX -> (FunctionTable -> LLVM PTX b) -> LLVM PTX b -withExecutable PTXR{..} f = do - r <- f (unsafeGetValue ptxExecutable) - liftIO $ touchLifetime ptxExecutable - return r {-- diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Link/Cache.hs b/src/Data/Array/Accelerate/LLVM/PTX/Link/Cache.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/Link/Cache.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/Link/Cache.hs @@ -1,9 +1,9 @@ -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Link.Cache --- Copyright : [2017] Trevor L. McDonell +-- Copyright : [2017..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Link/Object.hs b/src/Data/Array/Accelerate/LLVM/PTX/Link/Object.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/Link/Object.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/Link/Object.hs @@ -1,9 +1,9 @@ -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Link.Object --- Copyright : [2017] Trevor L. McDonell +-- Copyright : [2017..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Pool.hs b/src/Data/Array/Accelerate/LLVM/PTX/Pool.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/Pool.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/Pool.hs @@ -1,10 +1,10 @@ {-# LANGUAGE ScopedTypeVariables #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Pool --- Copyright : [2017] Trevor L. McDonell +-- Copyright : [2017..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- diff --git a/src/Data/Array/Accelerate/LLVM/PTX/State.hs b/src/Data/Array/Accelerate/LLVM/PTX/State.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/State.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/State.hs @@ -1,13 +1,11 @@ {-# LANGUAGE CPP #-} {-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.State --- Copyright : [2014..2017] Trevor L. McDonell --- [2014..2014] Vinod Grover (NVIDIA Corporation) +-- Copyright : [2014..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- @@ -24,7 +22,6 @@ ) where --- accelerate import Data.Array.Accelerate.Error import Data.Array.Accelerate.LLVM.State @@ -36,15 +33,10 @@ import qualified Data.Array.Accelerate.LLVM.PTX.Link.Cache as LC import qualified Data.Array.Accelerate.LLVM.PTX.Pool as Pool -import Data.Range ( Range(..) ) -import Control.Parallel.Meta ( Executable(..) ) - --- standard library -import Control.Concurrent ( runInBoundThread ) import Control.Exception ( try, catch ) import Data.Maybe ( fromMaybe, catMaybes ) import System.Environment ( lookupEnv ) -import System.IO.Unsafe ( unsafePerformIO ) +import System.IO.Unsafe ( unsafePerformIO, unsafeInterleaveIO ) import Text.Printf ( printf ) import Text.Read ( readMaybe ) import Foreign.CUDA.Driver.Error @@ -56,9 +48,9 @@ -- evalPTX :: PTX -> LLVM PTX a -> IO a evalPTX ptx acc = - runInBoundThread (CT.withContext (ptxContext ptx) (evalLLVM ptx acc)) + CT.withContext (ptxContext ptx) (evalLLVM ptx acc) `catch` - \e -> $internalError "unhandled" (show (e :: CUDAException)) + \e -> internalError (show (e :: CUDAException)) -- | Create a new PTX execution target for the given device @@ -98,15 +90,7 @@ mt <- MT.new ctx lc <- LC.new st <- ST.new ctx - return $! PTX ctx mt lc st simpleIO - - -{-# INLINE simpleIO #-} -simpleIO :: Executable -simpleIO = Executable $ \_name _ppt range action -> - case range of - Empty -> return () - IE u v -> action u v 0 + return $! PTX ctx mt lc st -- Shared execution contexts @@ -166,7 +150,7 @@ -- Spin up the GPU at the given ordinal. -- boot :: Int -> IO (Maybe PTX) - boot i = do + boot i = unsafeInterleaveIO $ do dev <- CUDA.device i prp <- CUDA.props dev r <- try $ createTargetForDevice dev prp [CUDA.SchedAuto] diff --git a/src/Data/Array/Accelerate/LLVM/PTX/Target.hs b/src/Data/Array/Accelerate/LLVM/PTX/Target.hs --- a/src/Data/Array/Accelerate/LLVM/PTX/Target.hs +++ b/src/Data/Array/Accelerate/LLVM/PTX/Target.hs @@ -1,15 +1,14 @@ {-# LANGUAGE CPP #-} {-# LANGUAGE EmptyDataDecls #-} {-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeFamilies #-} -- | -- Module : Data.Array.Accelerate.LLVM.PTX.Target --- Copyright : [2014..2017] Trevor L. McDonell --- [2014..2014] Vinod Grover (NVIDIA Corporation) +-- Copyright : [2014..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- @@ -21,7 +20,7 @@ ) where --- llvm-general +-- llvm-hs import LLVM.AST.AddrSpace import LLVM.AST.DataLayout import LLVM.Target hiding ( Target ) @@ -33,22 +32,22 @@ -- accelerate import Data.Array.Accelerate.Error +import Data.Array.Accelerate.LLVM.Extra import Data.Array.Accelerate.LLVM.Target -import Data.Array.Accelerate.LLVM.Util -import Control.Parallel.Meta ( Executable ) import Data.Array.Accelerate.LLVM.PTX.Array.Table ( MemoryTable ) import Data.Array.Accelerate.LLVM.PTX.Context ( Context, deviceProperties ) import Data.Array.Accelerate.LLVM.PTX.Execute.Stream.Reservoir ( Reservoir ) import Data.Array.Accelerate.LLVM.PTX.Link.Cache ( KernelTable ) -- CUDA -import qualified Foreign.CUDA.Driver as CUDA +import Foreign.CUDA.Analysis.Device -- standard library import Data.ByteString ( ByteString ) import Data.ByteString.Short ( ShortByteString ) import Data.String +import Debug.Trace import System.IO.Unsafe import Text.Printf import qualified Data.Map as Map @@ -69,22 +68,21 @@ , ptxMemoryTable :: {-# UNPACK #-} !MemoryTable , ptxKernelTable :: {-# UNPACK #-} !KernelTable , ptxStreamReservoir :: {-# UNPACK #-} !Reservoir - , fillP :: {-# UNPACK #-} !Executable } instance Target PTX where - targetTriple _ = Just ptxTargetTriple + targetTriple = Just ptxTargetTriple #if ACCELERATE_USE_NVVM - targetDataLayout _ = Nothing -- see note: [NVVM and target data layout] + targetDataLayout = Nothing -- see note: [NVVM and target data layout] #else - targetDataLayout _ = Just ptxDataLayout + targetDataLayout = Just ptxDataLayout #endif -- | Extract the properties of the device the current PTX execution state is -- executing on. -- -ptxDeviceProperties :: PTX -> CUDA.DeviceProperties +ptxDeviceProperties :: PTX -> DeviceProperties ptxDeviceProperties = deviceProperties . ptxContext @@ -121,53 +119,80 @@ -- | String that describes the target host. -- -ptxTargetTriple :: ShortByteString +ptxTargetTriple :: HasCallStack => ShortByteString ptxTargetTriple = case bitSize (undefined::Int) of 32 -> "nvptx-nvidia-cuda" 64 -> "nvptx64-nvidia-cuda" - _ -> $internalError "ptxTargetTriple" "I don't know what architecture I am" + _ -> internalError "I don't know what architecture I am" -- | Bracket creation and destruction of the NVVM TargetMachine. -- withPTXTargetMachine - :: CUDA.DeviceProperties + :: HasCallStack + => DeviceProperties -> (TargetMachine -> IO a) -> IO a withPTXTargetMachine dev go = - let CUDA.Compute m n = CUDA.computeCapability dev - isa = CPUFeature (ptxISAVersion m n) - sm = fromString (printf "sm_%d%d" m n) + let (sm, isa) = ptxTargetVersion (computeCapability dev) in withTargetOptions $ \options -> do withTargetMachine - ptxTarget - ptxTargetTriple - sm - (Map.singleton isa True) -- CPU features - options -- target options - R.Default -- relocation model - CM.Default -- code model - CGO.Default -- optimisation level - go + ptxTarget + ptxTargetTriple + sm -- CPU + (Map.singleton (CPUFeature isa) True) -- CPU features + options -- target options + R.Default -- relocation model + CM.Default -- code model + CGO.Default -- optimisation level + go --- Some libdevice functions require at least ptx40, even though devices at --- that compute capability also accept older ISA versions. +-- Compile using the earliest version of the SM target PTX ISA supported by +-- the given compute device and this version of LLVM. -- --- https://github.com/llvm-mirror/llvm/blob/master/lib/Target/NVPTX/NVPTX.td#L72 +-- Note that we require at least ptx40 for some libnvvm device functions. -- -ptxISAVersion :: Int -> Int -> ByteString -ptxISAVersion 2 _ = "ptx40" -ptxISAVersion 3 7 = "ptx41" -ptxISAVersion 3 _ = "ptx40" -ptxISAVersion 5 0 = "ptx40" -ptxISAVersion 5 2 = "ptx41" -ptxISAVersion 5 3 = "ptx42" -ptxISAVersion 6 _ = "ptx50" -ptxISAVersion 7 _ = "ptx60" -ptxISAVersion _ _ = "ptx40" - +-- See table NVPTX supported processors: +-- +-- https://github.com/llvm-mirror/llvm/blob/master/lib/Target/NVPTX/NVPTX.td +-- +-- PTX ISA verison history: +-- +-- https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#release-notes +-- +ptxTargetVersion :: Compute -> (ByteString, ByteString) +ptxTargetVersion compute@(Compute m n) +#if MIN_VERSION_llvm_hs(8,0,0) + | m >= 7 && n >= 5 = ("sm_75", "ptx63") +#endif +#if MIN_VERSION_llvm_hs(7,0,0) + | m >= 7 && n >= 2 = ("sm_72", "ptx61") +#endif +#if MIN_VERSION_llvm_hs(6,0,0) + | m >= 7 = ("sm_70", "ptx60") +#endif + | m > 6 = ("sm_62", "ptx50") -- fallthrough + -- + | m == 6 && n == 2 = ("sm_62", "ptx50") + | m == 6 && n == 1 = ("sm_61", "ptx50") + | m == 6 = ("sm_60", "ptx50") + | m == 5 && n == 3 = ("sm_53", "ptx42") + | m == 5 && n == 2 = ("sm_52", "ptx41") + | m == 5 = ("sm_50", "ptx40") + | m == 3 && n == 7 = ("sm_37", "ptx41") + | m == 3 && n == 5 = ("sm_35", "ptx40") + | m == 3 && n == 2 = ("sm_32", "ptx40") + | m == 3 = ("sm_30", "ptx40") + | m == 2 && n == 1 = ("sm_21", "ptx40") + | m == 2 = ("sm_20", "ptx40") + -- + | otherwise + = trace warning (fromString (printf "sm_%d%d" m n), "ptx40") + where + warning = unlines [ "*** Warning: Unhandled CUDA device compute capability: " ++ show compute + , "*** Please submit a bug report at https://github.com/AccelerateHS/accelerate/issues" ] -- | The NVPTX target for this host. -- diff --git a/src/GHC/Heap/NormalForm.hs b/src/GHC/Heap/NormalForm.hs new file mode 100644 --- /dev/null +++ b/src/GHC/Heap/NormalForm.hs @@ -0,0 +1,93 @@ +-- | +-- Module : GHC.Heap.NormalForm +-- Copyright : [2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- +-- https://github.com/input-output-hk/cardano-prelude/blob/96e8dcb29dc3c29eee99c0d020152fad6071af6d/src/Cardano/Prelude/GHC/Heap/NormalForm.hs +-- +-- This code has been adapted from the module "GHC.AssertNF" of the package +-- +-- () authored by +-- Joachim Breitner. +-- +-- To avoid space leaks and unwanted evaluation behaviour, the programmer +-- might want his data to be fully evaluated at certain positions in the +-- code. This can be enforced, for example, by ample use of +-- "Control.DeepSeq", but this comes at a cost. +-- +-- Experienced users hence use 'Control.DeepSeq.deepseq' only to find out +-- about the existence of space leaks and optimize their code to not create +-- the thunks in the first place, until the code no longer shows better +-- performance with 'deepseq'. +-- + +module GHC.Heap.NormalForm ( + + isHeadNormalForm, + isNormalForm, + +) where + +import GHC.Exts.Heap + +-- Everything is in normal form, unless it is a thunk explicitly marked as +-- such. Indirection are also considered to be in HNF. +-- +isHeadNormalForm :: Closure -> IO Bool +isHeadNormalForm c = do + case c of + ThunkClosure{} -> return False + APClosure{} -> return False + SelectorClosure{} -> return False + BCOClosure{} -> return False + _ -> return True + +-- | The function 'isNormalForm' checks whether its argument is fully evaluated +-- and deeply evaluated. +-- +-- NOTE 1: If you want to override the behaviour of 'isNormalForm' for specific +-- types (in particular, for specific types that may be /nested/ somewhere +-- inside the @a@), consider using +-- 'Cardano.Prelude.GHC.Heap.NormalForm.Classy.noUnexpectedThunks' instead. +-- +-- NOTE 2: The normal form check can be quite brittle, especially with @-O0@. +-- For example, writing something like +-- +-- > let !(Value x) = ... in .... +-- +-- might translate to +-- +-- > let !.. = ... in ... (case ... of Value x -> x) +-- +-- which would trivially be @False@. In general, 'isNormalForm' should probably +-- only be used with @-O1@, but even then the answer may still depend on +-- internal decisions made by ghc during compilation. +-- +isNormalForm :: a -> IO Bool +isNormalForm x = isNormalFormBoxed (asBox x) + +isNormalFormBoxed :: Box -> IO Bool +isNormalFormBoxed b = do + c <- getBoxedClosureData b + nf <- isHeadNormalForm c + if nf + then do + c' <- getBoxedClosureData b + allM isNormalFormBoxed (allClosures c') + else do + return False + +-- From Control.Monad.Loops in monad-loops +-- +allM :: Monad m => (a -> m Bool) -> [a] -> m Bool +allM _ [] = return True +allM p (x : xs) = do + q <- p x + if q + then allM p xs + else return False + diff --git a/src/System/Process/Extra.hs b/src/System/Process/Extra.hs --- a/src/System/Process/Extra.hs +++ b/src/System/Process/Extra.hs @@ -1,10 +1,10 @@ {-# LANGUAGE RecordWildCards #-} -- | -- Module : System.Process.Extra --- Copyright : [2017] Trevor L. McDonell +-- Copyright : [2017..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) -- diff --git a/test/nofib/Main.hs b/test/nofib/Main.hs --- a/test/nofib/Main.hs +++ b/test/nofib/Main.hs @@ -1,9 +1,9 @@ -- | -- Module : nofib-llvm-ptx --- Copyright : [2017] Trevor L. McDonell +-- Copyright : [2017..2020] The Accelerate Team -- License : BSD3 -- --- Maintainer : Trevor L. McDonell +-- Maintainer : Trevor L. McDonell -- Stability : experimental -- Portability : non-portable (GHC extensions) --