diff --git a/docs/error-index.rst b/docs/error-index.rst --- a/docs/error-index.rst +++ b/docs/error-index.rst @@ -475,11 +475,10 @@ "Parameter *x* used as size would go out of scope." ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -This error tends to happen because higher-order functions are used in -a way that causes a size requirement to become impossible to -constrait. Real programs that run into this issue are quite complex, -but to illustrate the problem, consider the following contrived -function: +This error tends to happen when higher-order functions are used in a +way that causes a size requirement to become impossible to express. +Real programs that encounter this issue tend to be complicated, but to +illustrate the problem, consider the following contrived function: .. code-block:: futhark diff --git a/futhark.cabal b/futhark.cabal --- a/futhark.cabal +++ b/futhark.cabal @@ -1,6 +1,6 @@ cabal-version: 2.4 name: futhark -version: 0.21.15 +version: 0.22.1 synopsis: An optimising compiler for a functional, array-oriented language. description: Futhark is a small programming language designed to be compiled to @@ -190,6 +190,7 @@ Futhark.CodeGen.ImpGen.CUDA Futhark.CodeGen.ImpGen.GPU Futhark.CodeGen.ImpGen.GPU.Base + Futhark.CodeGen.ImpGen.GPU.Group Futhark.CodeGen.ImpGen.GPU.SegHist Futhark.CodeGen.ImpGen.GPU.SegMap Futhark.CodeGen.ImpGen.GPU.SegRed diff --git a/prelude/soacs.fut b/prelude/soacs.fut --- a/prelude/soacs.fut +++ b/prelude/soacs.fut @@ -195,54 +195,6 @@ let (as', is) = intrinsics.partition (3, p', as) in (as'[0:is[0]], as'[is[0]:is[0]+is[1]], as'[is[0]+is[1]:n]) --- | `reduce_stream op f as` splits `as` into chunks, applies `f` to each --- of these in parallel, and uses `op` (which must be associative) to --- combine the per-chunk results into a final result. The `i64` --- passed to `f` is the size of the chunk. This SOAC is useful when --- `f` can be given a particularly work-efficient sequential --- implementation. Operationally, we can imagine that `as` is divided --- among as many threads as necessary to saturate the machine, with --- each thread operating sequentially. --- --- A chunk may be empty, and `f 0 []` must produce the neutral element for --- `op`. --- --- **Work:** *O(n ✕ W(op) + W(f))* --- --- **Span:** *O(log(n) ✕ W(op))* -def reduce_stream [n] 'a 'b (op: b -> b -> b) (f: (k: i64) -> [k]a -> b) (as: [n]a): b = - intrinsics.reduce_stream (op, f, as) - --- | As `reduce_stream`@term, but the chunks do not necessarily --- correspond to subsequences of the original array (they may be --- interleaved). --- --- **Work:** *O(n ✕ W(op) + W(f))* --- --- **Span:** *O(log(n) ✕ W(op))* -def reduce_stream_per [n] 'a 'b (op: b -> b -> b) (f: (k: i64) -> [k]a -> b) (as: [n]a): b = - intrinsics.reduce_stream_per (op, f, as) - --- | Similar to `reduce_stream`@term, except that each chunk must produce --- an array *of the same size*. The per-chunk results are --- concatenated. --- --- **Work:** *O(n ✕ W(f))* --- --- **Span:** *O(S(f))* -def map_stream [n] 'a 'b (f: (k: i64) -> [k]a -> [k]b) (as: [n]a): *[n]b = - intrinsics.map_stream (f, as) - --- | Similar to `map_stream`@term, but the chunks do not necessarily --- correspond to subsequences of the original array (they may be --- interleaved). --- --- **Work:** *O(n ✕ W(f))* --- --- **Span:** *O(S(f))* -def map_stream_per [n] 'a 'b (f: (k: i64) -> [k]a -> [k]b) (as: [n]a): *[n]b = - intrinsics.map_stream_per (f, as) - -- | Return `true` if the given function returns `true` for all -- elements in the array. -- diff --git a/src/Futhark/AD/Fwd.hs b/src/Futhark/AD/Fwd.hs --- a/src/Futhark/AD/Fwd.hs +++ b/src/Futhark/AD/Fwd.hs @@ -311,19 +311,13 @@ redLambda = op', redNeutral = redNeutral red `interleave` map Var neutral_tans } -fwdSOAC pat aux (Stream size xs form nes lam) = do +fwdSOAC pat aux (Stream size xs nes lam) = do pat' <- bundleNew pat lam' <- fwdStreamLambda lam xs' <- bundleTan xs nes_tan <- mapM (fmap Var . zeroFromSubExp) nes let nes' = interleave nes nes_tan - case form of - Sequential -> - addStm $ Let pat' aux $ Op $ Stream size xs' Sequential nes' lam' - Parallel o comm lam0 -> do - lam0' <- fwdLambda lam0 - let form' = Parallel o comm lam0' - addStm $ Let pat' aux $ Op $ Stream size xs' form' nes' lam' + addStm $ Let pat' aux $ Op $ Stream size xs' nes' lam' fwdSOAC pat aux (Hist w arrs ops bucket_fun) = do pat' <- bundleNew pat ops' <- mapM fwdHist ops diff --git a/src/Futhark/Analysis/HORep/SOAC.hs b/src/Futhark/Analysis/HORep/SOAC.hs --- a/src/Futhark/Analysis/HORep/SOAC.hs +++ b/src/Futhark/Analysis/HORep/SOAC.hs @@ -85,8 +85,6 @@ import Futhark.IR.SOACS.SOAC ( HistOp (..), ScremaForm (..), - StreamForm (..), - StreamOrd (..), scremaType, ) import qualified Futhark.IR.SOACS.SOAC as Futhark @@ -383,7 +381,7 @@ -- | A definite representation of a SOAC expression. data SOAC rep - = Stream SubExp (StreamForm rep) (Lambda rep) [SubExp] [Input] + = Stream SubExp (Lambda rep) [SubExp] [Input] | Scatter SubExp (Lambda rep) [Input] [(Shape, Int, VName)] | Screma SubExp (ScremaForm rep) [Input] | Hist SubExp [HistOp rep] (Lambda rep) [Input] @@ -412,20 +410,20 @@ instance PrettyRep rep => PP.Pretty (SOAC rep) where ppr (Screma w form arrs) = Futhark.ppScrema w arrs form ppr (Hist len ops bucket_fun imgs) = Futhark.ppHist len imgs ops bucket_fun - ppr (Stream w form lam nes arrs) = Futhark.ppStream w arrs form nes lam + ppr (Stream w lam nes arrs) = Futhark.ppStream w arrs nes lam ppr (Scatter w lam arrs dests) = Futhark.ppScatter w arrs lam dests -- | Returns the inputs used in a SOAC. inputs :: SOAC rep -> [Input] -inputs (Stream _ _ _ _ arrs) = arrs +inputs (Stream _ _ _ arrs) = arrs inputs (Scatter _len _lam ivs _as) = ivs inputs (Screma _ _ arrs) = arrs inputs (Hist _ _ _ inps) = inps -- | Set the inputs to a SOAC. setInputs :: [Input] -> SOAC rep -> SOAC rep -setInputs arrs (Stream w form lam nes _) = - Stream (newWidth arrs w) form lam nes arrs +setInputs arrs (Stream w lam nes _) = + Stream (newWidth arrs w) lam nes arrs setInputs arrs (Scatter w lam _ivs as) = Scatter (newWidth arrs w) lam arrs as setInputs arrs (Screma w form _) = @@ -439,15 +437,15 @@ -- | The lambda used in a given SOAC. lambda :: SOAC rep -> Lambda rep -lambda (Stream _ _ lam _ _) = lam +lambda (Stream _ lam _ _) = lam lambda (Scatter _len lam _ivs _as) = lam lambda (Screma _ (ScremaForm _ _ lam) _) = lam lambda (Hist _ _ lam _) = lam -- | Set the lambda used in the SOAC. setLambda :: Lambda rep -> SOAC rep -> SOAC rep -setLambda lam (Stream w form _ nes arrs) = - Stream w form lam nes arrs +setLambda lam (Stream w _ nes arrs) = + Stream w lam nes arrs setLambda lam (Scatter len _lam ivs as) = Scatter len lam ivs as setLambda lam (Screma w (ScremaForm scan red _) arrs) = @@ -457,7 +455,7 @@ -- | The return type of a SOAC. typeOf :: SOAC rep -> [Type] -typeOf (Stream w _ lam nes _) = +typeOf (Stream w lam nes _) = let accrtps = take (length nes) $ lambdaReturnType lam arrtps = [ arrayOf (stripArray 1 t) (Shape [w]) NoUniqueness @@ -479,7 +477,7 @@ -- | The "width" of a SOAC is the expected outer size of its array -- inputs _after_ input-transforms have been carried out. width :: SOAC rep -> SubExp -width (Stream w _ _ _ _) = w +width (Stream w _ _ _) = w width (Scatter len _lam _ivs _as) = len width (Screma w _ _) = w width (Hist w _ _ _) = w @@ -493,8 +491,8 @@ -- | Convert a SOAC to a Futhark-level SOAC. toSOAC :: MonadBuilder m => SOAC (Rep m) -> m (Futhark.SOAC (Rep m)) -toSOAC (Stream w form lam nes inps) = - Futhark.Stream w <$> inputsToSubExps inps <*> pure form <*> pure nes <*> pure lam +toSOAC (Stream w lam nes inps) = + Futhark.Stream w <$> inputsToSubExps inps <*> pure nes <*> pure lam toSOAC (Scatter w lam ivs dests) = Futhark.Scatter w <$> inputsToSubExps ivs <*> pure lam <*> pure dests toSOAC (Screma w form arrs) = @@ -516,8 +514,8 @@ (Op rep ~ Futhark.SOAC rep, HasScope rep m) => Exp rep -> m (Either NotSOAC (SOAC rep)) -fromExp (Op (Futhark.Stream w as form nes lam)) = - Right . Stream w form lam nes <$> traverse varInput as +fromExp (Op (Futhark.Stream w as nes lam)) = + Right . Stream w lam nes <$> traverse varInput as fromExp (Op (Futhark.Scatter w ivs lam as)) = Right <$> (Scatter w lam <$> traverse varInput ivs <*> pure as) fromExp (Op (Futhark.Screma w arrs form)) = @@ -566,9 +564,8 @@ strmbdy = mkBody (oneStm insstm) $ map (subExpRes . Var . identName) strm_resids strmpar = chunk_param : strm_inpids strmlam = Lambda strmpar strmbdy loutps - empty_lam = Lambda [] (mkBody mempty []) [] -- map(f,a) creates a stream with NO accumulators - pure (Stream w (Parallel Disorder Commutative empty_lam) strmlam [] inps, []) + pure (Stream w strmlam [] inps, []) | Just (scans, _) <- Futhark.isScanomapSOAC form, Futhark.Scan scan_lam nes <- Futhark.singleScan scans -> do -- scanomap(scan_lam,nes,map_lam,a) => is translated in strem's body to: @@ -628,7 +625,7 @@ addStms addlelstm pure $ addlelres ++ map (subExpRes . Var) (strm_resids ++ map_resids) pure - ( Stream w Sequential strmlam nes inps, + ( Stream w strmlam nes inps, map paramIdent inpacc_ids ) | Just (reds, _) <- Futhark.isRedomapSOAC form, @@ -666,8 +663,7 @@ addaccres ++ map (subExpRes . Var . identName) strm_resids strmpar = chunk_param : inpacc_ids ++ strm_inpids strmlam = Lambda strmpar strmbdy (accrtps ++ loutps') - lam0 <- renameLambda lamin - pure (Stream w (Parallel InOrder comm lam0) strmlam nes inps, []) + pure (Stream w strmlam nes inps, []) -- Otherwise it cannot become a stream. _ -> pure (soac, []) diff --git a/src/Futhark/CodeGen/Backends/GenericC.hs b/src/Futhark/CodeGen/Backends/GenericC.hs --- a/src/Futhark/CodeGen/Backends/GenericC.hs +++ b/src/Futhark/CodeGen/Backends/GenericC.hs @@ -38,7 +38,6 @@ import Futhark.CodeGen.Backends.GenericC.Types import Futhark.CodeGen.ImpCode import Futhark.CodeGen.RTS.C (cacheH, contextH, contextPrototypesH, errorsH, halfH, lockH, timingH, utilH) -import Futhark.IR.Prop (isBuiltInFunction) import qualified Futhark.Manifest as Manifest import Futhark.MonadFreshNames import Futhark.Util.Pretty (prettyText) @@ -49,15 +48,8 @@ defCall :: CallCompiler op s defCall dests fname args = do let out_args = [[C.cexp|&$id:d|] | d <- dests] - args' - | isBuiltInFunction fname = args - | otherwise = [C.cexp|ctx|] : out_args ++ args - case dests of - [dest] - | isBuiltInFunction fname -> - stm [C.cstm|$id:dest = $id:(funName fname)($args:args');|] - _ -> - item [C.citem|if ($id:(funName fname)($args:args') != 0) { err = 1; goto cleanup; }|] + args' = [C.cexp|ctx|] : out_args ++ args + item [C.citem|if ($id:(funName fname)($args:args') != 0) { err = 1; goto cleanup; }|] defError :: ErrorCompiler op s defError msg stacktrace = do diff --git a/src/Futhark/CodeGen/Backends/GenericC/Code.hs b/src/Futhark/CodeGen/Backends/GenericC/Code.hs --- a/src/Futhark/CodeGen/Backends/GenericC/Code.hs +++ b/src/Futhark/CodeGen/Backends/GenericC/Code.hs @@ -21,6 +21,7 @@ import Data.Maybe import Futhark.CodeGen.Backends.GenericC.Monad import Futhark.CodeGen.ImpCode +import Futhark.IR.Prop (isBuiltInFunction) import Futhark.MonadFreshNames import qualified Language.C.Quote.OpenCL as C import qualified Language.C.Syntax as C @@ -134,6 +135,37 @@ assignmentOperator Mul {} = Just $ \d e -> [C.cexp|$id:d *= $exp:e|] assignmentOperator _ = Nothing +compileRead :: + VName -> + Count u (TPrimExp t VName) -> + PrimType -> + Space -> + Volatility -> + CompilerM op s C.Exp +compileRead _ _ Unit _ _ = + pure [C.cexp|$exp:(UnitValue)|] +compileRead src (Count iexp) restype DefaultSpace vol = do + src' <- rawMem src + fmap (fromStorage restype) $ + derefPointer src' + <$> compileExp (untyped iexp) + <*> pure [C.cty|$tyquals:(volQuals vol) $ty:(primStorageType restype)*|] +compileRead src (Count iexp) restype (Space space) vol = + fmap (fromStorage restype) . join $ + asks (opsReadScalar . envOperations) + <*> rawMem src + <*> compileExp (untyped iexp) + <*> pure (primStorageType restype) + <*> pure space + <*> pure vol +compileRead src (Count iexp) _ ScalarSpace {} _ = do + iexp' <- compileExp $ untyped iexp + pure [C.cexp|$id:src[$exp:iexp']|] + +compileArg :: Arg -> CompilerM op s C.Exp +compileArg (MemArg m) = pure [C.cexp|$exp:m|] +compileArg (ExpArg e) = compileExp e + compileCode :: Code op -> CompilerM op s () compileCode (Op op) = join $ asks (opsCompiler . envOperations) <*> pure op @@ -175,6 +207,19 @@ e' <- compileExp e item [C.citem|$tyquals:(volQuals vol) $ty:ct $id:name = $exp:e';|] go code + go (DeclareScalar name vol t : Read dest src i restype space read_vol : code) + | name == dest = do + let ct = primTypeToCType t + e <- compileRead src i restype space read_vol + item [C.citem|$tyquals:(volQuals vol) $ty:ct $id:name = $exp:e;|] + go code + go (DeclareScalar name vol t : Call [dest] fname args : code) + | name == dest, + isBuiltInFunction fname = do + let ct = primTypeToCType t + args' <- mapM compileArg args + item [C.citem|$tyquals:(volQuals vol) $ty:ct $id:name = $id:(funName fname)($args:args');|] + go code go (x : xs) = compileCode x >> go xs go [] = pure () compileCode (Assert e msg (loc, locs)) = do @@ -276,29 +321,9 @@ <*> pure space <*> pure vol <*> (toStorage elemtype <$> compileExp elemexp) -compileCode (Read x _ _ Unit __ _) = - stm [C.cstm|$id:x = $exp:(UnitValue);|] -compileCode (Read x src (Count iexp) restype DefaultSpace vol) = do - src' <- rawMem src - e <- - fmap (fromStorage restype) $ - derefPointer src' - <$> compileExp (untyped iexp) - <*> pure [C.cty|$tyquals:(volQuals vol) $ty:(primStorageType restype)*|] - stm [C.cstm|$id:x = $exp:e;|] -compileCode (Read x src (Count iexp) restype (Space space) vol) = do - e <- - fmap (fromStorage restype) . join $ - asks (opsReadScalar . envOperations) - <*> rawMem src - <*> compileExp (untyped iexp) - <*> pure (primStorageType restype) - <*> pure space - <*> pure vol +compileCode (Read x src i restype space vol) = do + e <- compileRead src i restype space vol stm [C.cstm|$id:x = $exp:e;|] -compileCode (Read x src (Count iexp) _ ScalarSpace {} _) = do - iexp' <- compileExp $ untyped iexp - stm [C.cstm|$id:x = $id:src[$exp:iexp'];|] compileCode (DeclareMem name space) = declMem name space compileCode (DeclareScalar name vol t) = do @@ -345,12 +370,13 @@ stm [C.cstm|$id:dest = $exp:src';|] compileCode (SetMem dest src space) = setMem dest src space +compileCode (Call [dest] fname args) + | isBuiltInFunction fname = do + args' <- mapM compileArg args + stm [C.cstm|$id:dest = $id:(funName fname)($args:args');|] compileCode (Call dests fname args) = join $ asks (opsCall . envOperations) <*> pure dests <*> pure fname <*> mapM compileArg args - where - compileArg (MemArg m) = pure [C.cexp|$exp:m|] - compileArg (ExpArg e) = compileExp e diff --git a/src/Futhark/CodeGen/ImpGen/GPU/Base.hs b/src/Futhark/CodeGen/ImpGen/GPU/Base.hs --- a/src/Futhark/CodeGen/ImpGen/GPU/Base.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU/Base.hs @@ -4,2157 +4,1385 @@ module Futhark.CodeGen.ImpGen.GPU.Base ( KernelConstants (..), - keyWithEntryPoint, - CallKernelGen, - InKernelGen, - Locks (..), - HostEnv (..), - Target (..), - KernelEnv (..), - computeThreadChunkSize, - groupReduce, - groupScan, - isActive, - sKernelThread, - sKernelGroup, - KernelAttrs (..), - defKernelAttrs, - sReplicate, - sIota, - sRotateKernel, - sCopy, - compileThreadResult, - compileGroupResult, - virtualiseGroups, - kernelLoop, - groupCoverSpace, - Precomputed, - precomputeConstants, - precomputedConstants, - atomicUpdateLocking, - AtomicBinOp, - Locking (..), - AtomicUpdate (..), - DoAtomicUpdate, - ) -where - -import Control.Monad.Except -import Data.Bifunctor -import Data.List (foldl', partition, zip4) -import qualified Data.Map.Strict as M -import Data.Maybe -import qualified Data.Set as S -import qualified Futhark.CodeGen.ImpCode.GPU as Imp -import Futhark.CodeGen.ImpGen -import Futhark.Construct (fullSliceNum) -import Futhark.Error -import Futhark.IR.GPUMem -import qualified Futhark.IR.Mem.IxFun as IxFun -import Futhark.MonadFreshNames -import Futhark.Transform.Rename -import Futhark.Util (chunks, dropLast, mapAccumLM, nubOrd, splitFromEnd, takeLast) -import Futhark.Util.IntegralExp (divUp, quot, rem) -import Prelude hiding (quot, rem) - --- | Which target are we ultimately generating code for? While most --- of the kernels code is the same, there are some cases where we --- generate special code based on the ultimate low-level API we are --- targeting. -data Target = CUDA | OpenCL - --- | Information about the locks available for accumulators. -data Locks = Locks - { locksArray :: VName, - locksCount :: Int - } - -data HostEnv = HostEnv - { hostAtomics :: AtomicBinOp, - hostTarget :: Target, - hostLocks :: M.Map VName Locks - } - -data KernelEnv = KernelEnv - { kernelAtomics :: AtomicBinOp, - kernelConstants :: KernelConstants, - kernelLocks :: M.Map VName Locks - } - -type CallKernelGen = ImpM GPUMem HostEnv Imp.HostOp - -type InKernelGen = ImpM GPUMem KernelEnv Imp.KernelOp - -data KernelConstants = KernelConstants - { kernelGlobalThreadId :: Imp.TExp Int32, - kernelLocalThreadId :: Imp.TExp Int32, - kernelGroupId :: Imp.TExp Int32, - kernelGlobalThreadIdVar :: VName, - kernelLocalThreadIdVar :: VName, - kernelGroupIdVar :: VName, - kernelNumGroupsCount :: Count NumGroups SubExp, - kernelGroupSizeCount :: Count GroupSize SubExp, - kernelNumGroups :: Imp.TExp Int64, - kernelGroupSize :: Imp.TExp Int64, - kernelNumThreads :: Imp.TExp Int32, - kernelWaveSize :: Imp.TExp Int32, - -- | A mapping from dimensions of nested SegOps to already - -- computed local thread IDs. Only valid in non-virtualised case. - kernelLocalIdMap :: M.Map [SubExp] [Imp.TExp Int32], - -- | Mapping from dimensions of nested SegOps to how many - -- iterations the virtualisation loop needs. - kernelChunkItersMap :: M.Map [SubExp] (Imp.TExp Int32) - } - --- | The sizes of nested iteration spaces in the kernel. -type SegOpSizes = S.Set [SubExp] - --- | Find the sizes of nested parallelism in a t'SegOp' body. -segOpSizes :: Stms GPUMem -> SegOpSizes -segOpSizes = onStms - where - onStms = foldMap (onExp . stmExp) - onExp (Op (Inner (SegOp op))) = - case segVirt $ segLevel op of - SegNoVirtFull seq_dims -> - S.singleton $ map snd $ snd $ partitionSeqDims seq_dims $ segSpace op - _ -> S.singleton $ map snd $ unSegSpace $ segSpace op - onExp (BasicOp (Replicate shape _)) = - S.singleton $ shapeDims shape - onExp (Match _ cases defbody _) = - foldMap (onStms . bodyStms . caseBody) cases <> onStms (bodyStms defbody) - onExp (DoLoop _ _ body) = - onStms (bodyStms body) - onExp _ = mempty - --- | Various useful precomputed information. -data Precomputed = Precomputed - { pcSegOpSizes :: SegOpSizes, - pcChunkItersMap :: M.Map [SubExp] (Imp.TExp Int32) - } - --- | Precompute various constants and useful information. -precomputeConstants :: Count GroupSize (Imp.TExp Int64) -> Stms GPUMem -> CallKernelGen Precomputed -precomputeConstants group_size stms = do - let sizes = segOpSizes stms - iters_map <- M.fromList <$> mapM mkMap (S.toList sizes) - pure $ Precomputed sizes iters_map - where - mkMap dims = do - let n = product $ map Imp.pe64 dims - num_chunks <- dPrimVE "num_chunks" $ sExt32 $ n `divUp` unCount group_size - pure (dims, num_chunks) - --- | Make use of various precomputed constants. -precomputedConstants :: Precomputed -> InKernelGen a -> InKernelGen a -precomputedConstants pre m = do - ltid <- kernelLocalThreadId . kernelConstants <$> askEnv - new_ids <- M.fromList <$> mapM (mkMap ltid) (S.toList (pcSegOpSizes pre)) - let f env = - env - { kernelConstants = - (kernelConstants env) - { kernelLocalIdMap = new_ids, - kernelChunkItersMap = pcChunkItersMap pre - } - } - localEnv f m - where - mkMap ltid dims = do - let dims' = map pe64 dims - ids' <- dIndexSpace' "ltid_pre" dims' (sExt64 ltid) - pure (dims, map sExt32 ids') - -keyWithEntryPoint :: Maybe Name -> Name -> Name -keyWithEntryPoint fname key = - nameFromString $ maybe "" ((++ ".") . nameToString) fname ++ nameToString key - -allocLocal :: AllocCompiler GPUMem r Imp.KernelOp -allocLocal mem size = - sOp $ Imp.LocalAlloc mem size - -kernelAlloc :: - Pat LetDecMem -> - SubExp -> - Space -> - InKernelGen () -kernelAlloc (Pat [_]) _ ScalarSpace {} = - -- Handled by the declaration of the memory block, which is then - -- translated to an actual scalar variable during C code generation. - pure () -kernelAlloc (Pat [mem]) size (Space "local") = - allocLocal (patElemName mem) $ Imp.bytes $ pe64 size -kernelAlloc (Pat [mem]) _ _ = - compilerLimitationS $ "Cannot allocate memory block " ++ pretty mem ++ " in kernel." -kernelAlloc dest _ _ = - error $ "Invalid target for in-kernel allocation: " ++ show dest - -splitSpace :: - Pat LetDecMem -> - SplitOrdering -> - SubExp -> - SubExp -> - SubExp -> - ImpM rep r op () -splitSpace (Pat [size]) o w i elems_per_thread = do - num_elements <- Imp.elements . TPrimExp <$> toExp w - let i' = pe64 i - elems_per_thread' <- Imp.elements . TPrimExp <$> toExp elems_per_thread - computeThreadChunkSize o i' elems_per_thread' num_elements (mkTV (patElemName size) int64) -splitSpace pat _ _ _ _ = - error $ "Invalid target for splitSpace: " ++ pretty pat - -updateAcc :: VName -> [SubExp] -> [SubExp] -> InKernelGen () -updateAcc acc is vs = sComment "UpdateAcc" $ do - -- See the ImpGen implementation of UpdateAcc for general notes. - let is' = map pe64 is - (c, space, arrs, dims, op) <- lookupAcc acc is' - sWhen (inBounds (Slice (map DimFix is')) dims) $ - case op of - Nothing -> - forM_ (zip arrs vs) $ \(arr, v) -> copyDWIMFix arr is' v [] - Just lam -> do - dLParams $ lambdaParams lam - let (_x_params, y_params) = - splitAt (length vs) $ map paramName $ lambdaParams lam - forM_ (zip y_params vs) $ \(yp, v) -> copyDWIM yp [] v [] - atomics <- kernelAtomics <$> askEnv - case atomicUpdateLocking atomics lam of - AtomicPrim f -> f space arrs is' - AtomicCAS f -> f space arrs is' - AtomicLocking f -> do - c_locks <- M.lookup c . kernelLocks <$> askEnv - case c_locks of - Just (Locks locks num_locks) -> do - let locking = - Locking locks 0 1 0 $ - pure . (`rem` fromIntegral num_locks) . flattenIndex dims - f locking space arrs is' - Nothing -> - error $ "Missing locks for " ++ pretty acc - -compileThreadExp :: ExpCompiler GPUMem KernelEnv Imp.KernelOp -compileThreadExp (Pat [pe]) (BasicOp (Opaque _ se)) = - -- Cannot print in GPU code. - copyDWIM (patElemName pe) [] se [] -compileThreadExp (Pat [dest]) (BasicOp (ArrayLit es _)) = - forM_ (zip [0 ..] es) $ \(i, e) -> - copyDWIMFix (patElemName dest) [fromIntegral (i :: Int64)] e [] -compileThreadExp _ (BasicOp (UpdateAcc acc is vs)) = - updateAcc acc is vs -compileThreadExp dest e = - defCompileExp dest e - --- | Assign iterations of a for-loop to all threads in the kernel. --- The passed-in function is invoked with the (symbolic) iteration. --- The body must contain thread-level code. For multidimensional --- loops, use 'groupCoverSpace'. -kernelLoop :: - IntExp t => - Imp.TExp t -> - Imp.TExp t -> - Imp.TExp t -> - (Imp.TExp t -> InKernelGen ()) -> - InKernelGen () -kernelLoop tid num_threads n f = - localOps threadOperations $ - if n == num_threads - then f tid - else do - num_chunks <- dPrimVE "num_chunks" $ n `divUp` num_threads - sFor "chunk_i" num_chunks $ \chunk_i -> do - i <- dPrimVE "i" $ chunk_i * num_threads + tid - sWhen (i .<. n) $ f i - --- | Assign iterations of a for-loop to threads in the workgroup. The --- passed-in function is invoked with the (symbolic) iteration. For --- multidimensional loops, use 'groupCoverSpace'. -groupLoop :: - IntExp t => - Imp.TExp t -> - (Imp.TExp t -> InKernelGen ()) -> - InKernelGen () -groupLoop n f = do - constants <- kernelConstants <$> askEnv - kernelLoop - (kernelLocalThreadId constants `sExtAs` n) - (kernelGroupSize constants `sExtAs` n) - n - f - --- | Iterate collectively though a multidimensional space, such that --- all threads in the group participate. The passed-in function is --- invoked with a (symbolic) point in the index space. -groupCoverSpace :: - IntExp t => - [Imp.TExp t] -> - ([Imp.TExp t] -> InKernelGen ()) -> - InKernelGen () -groupCoverSpace ds f = do - constants <- kernelConstants <$> askEnv - let group_size = kernelGroupSize constants - case splitFromEnd 1 ds of - -- Optimise the case where the inner dimension of the space is - -- equal to the group size. - (ds', [last_d]) - | last_d == (group_size `sExtAs` last_d) -> do - let ltid = kernelLocalThreadId constants `sExtAs` last_d - sLoopSpace ds' $ \ds_is -> - f $ ds_is ++ [ltid] - _ -> - groupLoop (product ds) $ f . unflattenIndex ds - -localThreadIDs :: [SubExp] -> InKernelGen [Imp.TExp Int64] -localThreadIDs dims = do - ltid <- sExt64 . kernelLocalThreadId . kernelConstants <$> askEnv - let dims' = map pe64 dims - maybe (dIndexSpace' "ltid" dims' ltid) (pure . map sExt64) - . M.lookup dims - . kernelLocalIdMap - . kernelConstants - =<< askEnv - -partitionSeqDims :: SegSeqDims -> SegSpace -> ([(VName, SubExp)], [(VName, SubExp)]) -partitionSeqDims (SegSeqDims seq_is) space = - bimap (map fst) (map fst) $ - partition ((`elem` seq_is) . snd) (zip (unSegSpace space) [0 ..]) - -groupCoverSegSpace :: SegVirt -> SegSpace -> InKernelGen () -> InKernelGen () -groupCoverSegSpace virt space m = do - let (ltids, dims) = unzip $ unSegSpace space - dims' = map pe64 dims - - constants <- kernelConstants <$> askEnv - let group_size = kernelGroupSize constants - -- Maybe we can statically detect that this is actually a - -- SegNoVirtFull and generate ever-so-slightly simpler code. - let virt' = if dims' == [group_size] then SegNoVirtFull (SegSeqDims []) else virt - case virt' of - SegVirt -> do - iters <- M.lookup dims . kernelChunkItersMap . kernelConstants <$> askEnv - case iters of - Nothing -> do - iterations <- dPrimVE "iterations" $ product $ map sExt32 dims' - groupLoop iterations $ \i -> do - dIndexSpace (zip ltids dims') $ sExt64 i - m - Just num_chunks -> do - let ltid = kernelLocalThreadId constants - sFor "chunk_i" num_chunks $ \chunk_i -> do - i <- dPrimVE "i" $ chunk_i * sExt32 group_size + ltid - dIndexSpace (zip ltids dims') $ sExt64 i - sWhen (inBounds (Slice (map (DimFix . le64) ltids)) dims') m - SegNoVirt -> localOps threadOperations $ do - zipWithM_ dPrimV_ ltids =<< localThreadIDs dims - sWhen (isActive $ zip ltids dims) m - SegNoVirtFull seq_dims -> do - let ((ltids_seq, dims_seq), (ltids_par, dims_par)) = - bimap unzip unzip $ partitionSeqDims seq_dims space - sLoopNest (Shape dims_seq) $ \is_seq -> do - zipWithM_ dPrimV_ ltids_seq is_seq - localOps threadOperations $ do - zipWithM_ dPrimV_ ltids_par =<< localThreadIDs dims_par - m - -compileGroupExp :: ExpCompiler GPUMem KernelEnv Imp.KernelOp -compileGroupExp (Pat [pe]) (BasicOp (Opaque _ se)) = - -- Cannot print in GPU code. - copyDWIM (patElemName pe) [] se [] --- The static arrays stuff does not work inside kernels. -compileGroupExp (Pat [dest]) (BasicOp (ArrayLit es _)) = - forM_ (zip [0 ..] es) $ \(i, e) -> - copyDWIMFix (patElemName dest) [fromIntegral (i :: Int64)] e [] -compileGroupExp _ (BasicOp (UpdateAcc acc is vs)) = - updateAcc acc is vs -compileGroupExp (Pat [dest]) (BasicOp (Replicate ds se)) = do - flat <- newVName "rep_flat" - is <- replicateM (shapeRank ds) (newVName "rep_i") - let is' = map le64 is - groupCoverSegSpace SegVirt (SegSpace flat $ zip is $ shapeDims ds) $ - copyDWIMFix (patElemName dest) is' se [] - sOp $ Imp.Barrier Imp.FenceLocal -compileGroupExp (Pat [dest]) (BasicOp (Rotate rs arr)) = do - ds <- map pe64 . arrayDims <$> lookupType arr - groupCoverSpace ds $ \is -> do - is' <- sequence $ zipWith3 rotate ds rs is - copyDWIMFix (patElemName dest) is (Var arr) is' - sOp $ Imp.Barrier Imp.FenceLocal - where - rotate d r i = dPrimVE "rot_i" $ rotateIndex d (pe64 r) i -compileGroupExp (Pat [dest]) (BasicOp (Iota n e s it)) = do - n' <- toExp n - e' <- toExp e - s' <- toExp s - groupLoop (TPrimExp n') $ \i' -> do - x <- - dPrimV "x" $ - TPrimExp $ - BinOpExp (Add it OverflowUndef) e' $ - BinOpExp (Mul it OverflowUndef) (untyped i') s' - copyDWIMFix (patElemName dest) [i'] (Var (tvVar x)) [] - sOp $ Imp.Barrier Imp.FenceLocal - --- When generating code for a scalar in-place update, we must make --- sure that only one thread performs the write. When writing an --- array, the group-level copy code will take care of doing the right --- thing. -compileGroupExp (Pat [pe]) (BasicOp (Update safety _ slice se)) - | null $ sliceDims slice = do - sOp $ Imp.Barrier Imp.FenceLocal - ltid <- kernelLocalThreadId . kernelConstants <$> askEnv - sWhen (ltid .==. 0) $ - case safety of - Unsafe -> write - Safe -> sWhen (inBounds slice' dims) write - sOp $ Imp.Barrier Imp.FenceLocal - where - slice' = fmap pe64 slice - dims = map pe64 $ arrayDims $ patElemType pe - write = copyDWIM (patElemName pe) (unSlice slice') se [] -compileGroupExp dest e = - defCompileExp dest e - -sanityCheckLevel :: SegLevel -> InKernelGen () -sanityCheckLevel SegThread {} = pure () -sanityCheckLevel SegGroup {} = - error "compileGroupOp: unexpected group-level SegOp." - -compileFlatId :: SegLevel -> SegSpace -> InKernelGen () -compileFlatId lvl space = do - sanityCheckLevel lvl - ltid <- kernelLocalThreadId . kernelConstants <$> askEnv - dPrimV_ (segFlat space) ltid - --- Construct the necessary lock arrays for an intra-group histogram. -prepareIntraGroupSegHist :: - Count GroupSize SubExp -> - [HistOp GPUMem] -> - InKernelGen [[Imp.TExp Int64] -> InKernelGen ()] -prepareIntraGroupSegHist group_size = - fmap snd . mapAccumLM onOp Nothing - where - onOp l op = do - constants <- kernelConstants <$> askEnv - atomicBinOp <- kernelAtomics <$> askEnv - - let local_subhistos = histDest op - - case (l, atomicUpdateLocking atomicBinOp $ histOp op) of - (_, AtomicPrim f) -> pure (l, f (Space "local") local_subhistos) - (_, AtomicCAS f) -> pure (l, f (Space "local") local_subhistos) - (Just l', AtomicLocking f) -> pure (l, f l' (Space "local") local_subhistos) - (Nothing, AtomicLocking f) -> do - locks <- newVName "locks" - - let num_locks = pe64 $ unCount group_size - dims = map pe64 $ shapeDims (histOpShape op <> histShape op) - l' = Locking locks 0 1 0 (pure . (`rem` num_locks) . flattenIndex dims) - locks_t = Array int32 (Shape [unCount group_size]) NoUniqueness - - locks_mem <- sAlloc "locks_mem" (typeSize locks_t) $ Space "local" - dArray locks int32 (arrayShape locks_t) locks_mem $ - IxFun.iota . map pe64 . arrayDims $ - locks_t - - sComment "All locks start out unlocked" $ - groupCoverSpace [kernelGroupSize constants] $ \is -> - copyDWIMFix locks is (intConst Int32 0) [] - - pure (Just l', f l' (Space "local") local_subhistos) - --- Which fence do we need to protect shared access to this memory space? -fenceForSpace :: Space -> Imp.Fence -fenceForSpace (Space "local") = Imp.FenceLocal -fenceForSpace _ = Imp.FenceGlobal - --- If we are touching these arrays, which kind of fence do we need? -fenceForArrays :: [VName] -> InKernelGen Imp.Fence -fenceForArrays = fmap (foldl' max Imp.FenceLocal) . mapM need - where - need arr = - fmap (fenceForSpace . entryMemSpace) - . lookupMemory - . memLocName - . entryArrayLoc - =<< lookupArray arr - -groupChunkLoop :: - Imp.TExp Int32 -> - (Imp.TExp Int32 -> TV Int64 -> InKernelGen ()) -> - InKernelGen () -groupChunkLoop w m = do - constants <- kernelConstants <$> askEnv - let max_chunk_size = sExt32 $ kernelGroupSize constants - num_chunks <- dPrimVE "num_chunks" $ w `divUp` max_chunk_size - sFor "chunk_i" num_chunks $ \chunk_i -> do - chunk_start <- - dPrimVE "chunk_start" $ chunk_i * max_chunk_size - chunk_end <- - dPrimVE "chunk_end" $ sMin32 w (chunk_start + max_chunk_size) - chunk_size <- - dPrimV "chunk_size" $ sExt64 $ chunk_end - chunk_start - m chunk_start chunk_size - -sliceArray :: Imp.TExp Int64 -> TV Int64 -> VName -> ImpM rep r op VName -sliceArray start size arr = do - MemLoc mem _ ixfun <- entryArrayLoc <$> lookupArray arr - arr_t <- lookupType arr - let slice = - fullSliceNum - (map Imp.pe64 (arrayDims arr_t)) - [DimSlice start (tvExp size) 1] - sArray - (baseString arr ++ "_chunk") - (elemType arr_t) - (arrayShape arr_t `setOuterDim` Var (tvVar size)) - mem - $ IxFun.slice ixfun slice - --- | @flattenArray k flat arr@ flattens the outer @k@ dimensions of --- @arr@ to @flat@. (Make sure @flat@ is the sum of those dimensions --- or you'll have a bad time.) -flattenArray :: Int -> TV Int64 -> VName -> ImpM rep r op VName -flattenArray k flat arr = do - ArrayEntry arr_loc pt <- lookupArray arr - let flat_shape = Shape $ Var (tvVar flat) : drop k (memLocShape arr_loc) - sArray (baseString arr ++ "_flat") pt flat_shape (memLocName arr_loc) $ - IxFun.reshape (memLocIxFun arr_loc) $ - map pe64 $ - shapeDims flat_shape - --- | @applyLambda lam dests args@ emits code that: --- --- 1. Binds each parameter of @lam@ to the corresponding element of --- @args@, interpreted as a (name,slice) pair (as in 'copyDWIM'). --- Use an empty list for a scalar. --- --- 2. Executes the body of @lam@. --- --- 3. Binds the t'SubExp's that are the 'Result' of @lam@ to the --- provided @dest@s, again interpreted as the destination for a --- 'copyDWIM'. -applyLambda :: - Mem rep inner => - Lambda rep -> - [(VName, [DimIndex (Imp.TExp Int64)])] -> - [(SubExp, [DimIndex (Imp.TExp Int64)])] -> - ImpM rep r op () -applyLambda lam dests args = do - dLParams $ lambdaParams lam - forM_ (zip (lambdaParams lam) args) $ \(p, (arg, arg_slice)) -> - copyDWIM (paramName p) [] arg arg_slice - compileStms mempty (bodyStms $ lambdaBody lam) $ do - let res = map resSubExp $ bodyResult $ lambdaBody lam - forM_ (zip dests res) $ \((dest, dest_slice), se) -> - copyDWIM dest dest_slice se [] - --- | As applyLambda, but first rename the names in the lambda. This --- makes it safe to apply it in multiple places. (It might be safe --- anyway, but you have to be more careful - use this if you are in --- doubt.) -applyRenamedLambda :: - Mem rep inner => - Lambda rep -> - [(VName, [DimIndex (Imp.TExp Int64)])] -> - [(SubExp, [DimIndex (Imp.TExp Int64)])] -> - ImpM rep r op () -applyRenamedLambda lam dests args = do - lam_renamed <- renameLambda lam - applyLambda lam_renamed dests args - -virtualisedGroupScan :: - Maybe (Imp.TExp Int32 -> Imp.TExp Int32 -> Imp.TExp Bool) -> - Imp.TExp Int32 -> - Lambda GPUMem -> - [VName] -> - InKernelGen () -virtualisedGroupScan seg_flag w lam arrs = do - groupChunkLoop w $ \chunk_start chunk_size -> do - constants <- kernelConstants <$> askEnv - let ltid = kernelLocalThreadId constants - crosses_segment = - case seg_flag of - Nothing -> false - Just flag_true -> - flag_true (sExt32 (chunk_start - 1)) (sExt32 chunk_start) - sComment "possibly incorporate carry" $ - sWhen (chunk_start .>. 0 .&&. ltid .==. 0 .&&. bNot crosses_segment) $ do - carry_idx <- dPrimVE "carry_idx" $ sExt64 chunk_start - 1 - applyRenamedLambda - lam - (zip arrs $ repeat [DimFix $ sExt64 chunk_start]) - ( zip (map Var arrs) (repeat [DimFix carry_idx]) - ++ zip (map Var arrs) (repeat [DimFix $ sExt64 chunk_start]) - ) - - arrs_chunks <- mapM (sliceArray (sExt64 chunk_start) chunk_size) arrs - - sOp $ Imp.ErrorSync Imp.FenceLocal - - groupScan - seg_flag - (sExt64 w) - (tvExp chunk_size) - lam - arrs_chunks - -compileGroupOp :: OpCompiler GPUMem KernelEnv Imp.KernelOp -compileGroupOp pat (Alloc size space) = - kernelAlloc pat size space -compileGroupOp pat (Inner (SizeOp (SplitSpace o w i elems_per_thread))) = - splitSpace pat o w i elems_per_thread -compileGroupOp pat (Inner (SegOp (SegMap lvl space _ body))) = do - compileFlatId lvl space - - groupCoverSegSpace (segVirt lvl) space $ - compileStms mempty (kernelBodyStms body) $ - zipWithM_ (compileThreadResult space) (patElems pat) $ - kernelBodyResult body - sOp $ Imp.ErrorSync Imp.FenceLocal -compileGroupOp pat (Inner (SegOp (SegScan lvl space scans _ body))) = do - compileFlatId lvl space - - let (ltids, dims) = unzip $ unSegSpace space - dims' = map pe64 dims - - groupCoverSegSpace (segVirt lvl) space $ - compileStms mempty (kernelBodyStms body) $ - forM_ (zip (patNames pat) $ kernelBodyResult body) $ \(dest, res) -> - copyDWIMFix - dest - (map Imp.le64 ltids) - (kernelResultSubExp res) - [] - - fence <- fenceForArrays $ patNames pat - sOp $ Imp.ErrorSync fence - - let segment_size = last dims' - crossesSegment from to = - (sExt64 to - sExt64 from) .>. (sExt64 to `rem` segment_size) - - -- groupScan needs to treat the scan output as a one-dimensional - -- array of scan elements, so we invent some new flattened arrays - -- here. - dims_flat <- dPrimV "dims_flat" $ product dims' - let scan = head scans - num_scan_results = length $ segBinOpNeutral scan - arrs_flat <- - mapM (flattenArray (length dims') dims_flat) $ - take num_scan_results $ - patNames pat - - case segVirt lvl of - SegVirt -> - virtualisedGroupScan - (Just crossesSegment) - (sExt32 $ tvExp dims_flat) - (segBinOpLambda scan) - arrs_flat - _ -> - groupScan - (Just crossesSegment) - (product dims') - (product dims') - (segBinOpLambda scan) - arrs_flat -compileGroupOp pat (Inner (SegOp (SegRed lvl space ops _ body))) = do - compileFlatId lvl space - - let dims' = map pe64 dims - mkTempArr t = - sAllocArray "red_arr" (elemType t) (Shape dims <> arrayShape t) $ Space "local" - - tmp_arrs <- mapM mkTempArr $ concatMap (lambdaReturnType . segBinOpLambda) ops - groupCoverSegSpace (segVirt lvl) space $ - compileStms mempty (kernelBodyStms body) $ do - let (red_res, map_res) = - splitAt (segBinOpResults ops) $ kernelBodyResult body - forM_ (zip tmp_arrs red_res) $ \(dest, res) -> - copyDWIMFix dest (map Imp.le64 ltids) (kernelResultSubExp res) [] - zipWithM_ (compileThreadResult space) map_pes map_res - - sOp $ Imp.ErrorSync Imp.FenceLocal - - let tmps_for_ops = chunks (map (length . segBinOpNeutral) ops) tmp_arrs - case segVirt lvl of - SegVirt -> virtCase dims' tmps_for_ops - _ -> nonvirtCase dims' tmps_for_ops - where - (ltids, dims) = unzip $ unSegSpace space - (red_pes, map_pes) = splitAt (segBinOpResults ops) $ patElems pat - - virtCase [dim'] tmps_for_ops = do - ltid <- kernelLocalThreadId . kernelConstants <$> askEnv - groupChunkLoop (sExt32 dim') $ \chunk_start chunk_size -> do - sComment "possibly incorporate carry" $ - sWhen (chunk_start .>. 0 .&&. ltid .==. 0) $ - forM_ (zip ops tmps_for_ops) $ \(op, tmps) -> - applyRenamedLambda - (segBinOpLambda op) - (zip tmps $ repeat [DimFix $ sExt64 chunk_start]) - ( zip (map (Var . patElemName) red_pes) (repeat []) - ++ zip (map Var tmps) (repeat [DimFix $ sExt64 chunk_start]) - ) - - sOp $ Imp.ErrorSync Imp.FenceLocal - - forM_ (zip ops tmps_for_ops) $ \(op, tmps) -> do - tmps_chunks <- mapM (sliceArray (sExt64 chunk_start) chunk_size) tmps - groupReduce (sExt32 (tvExp chunk_size)) (segBinOpLambda op) tmps_chunks - - sOp $ Imp.ErrorSync Imp.FenceLocal - - forM_ (zip red_pes $ concat tmps_for_ops) $ \(pe, arr) -> - copyDWIMFix (patElemName pe) [] (Var arr) [sExt64 chunk_start] - virtCase dims' tmps_for_ops = do - dims_flat <- dPrimV "dims_flat" $ product dims' - let segment_size = last dims' - crossesSegment from to = - (sExt64 to - sExt64 from) .>. (sExt64 to `rem` sExt64 segment_size) - - forM_ (zip ops tmps_for_ops) $ \(op, tmps) -> do - tmps_flat <- mapM (flattenArray (length dims') dims_flat) tmps - virtualisedGroupScan - (Just crossesSegment) - (sExt32 $ tvExp dims_flat) - (segBinOpLambda op) - tmps_flat - - sOp $ Imp.ErrorSync Imp.FenceLocal - - forM_ (zip red_pes $ concat tmps_for_ops) $ \(pe, arr) -> - copyDWIM - (patElemName pe) - [] - (Var arr) - (map (unitSlice 0) (init dims') ++ [DimFix $ last dims' - 1]) - - sOp $ Imp.Barrier Imp.FenceLocal - - nonvirtCase [dim'] tmps_for_ops = do - -- Nonsegmented case (or rather, a single segment) - this we can - -- handle directly with a group-level reduction. - forM_ (zip ops tmps_for_ops) $ \(op, tmps) -> - groupReduce (sExt32 dim') (segBinOpLambda op) tmps - sOp $ Imp.ErrorSync Imp.FenceLocal - forM_ (zip red_pes $ concat tmps_for_ops) $ \(pe, arr) -> - copyDWIMFix (patElemName pe) [] (Var arr) [0] - -- - nonvirtCase dims' tmps_for_ops = do - -- Segmented intra-group reductions are turned into (regular) - -- segmented scans. It is possible that this can be done - -- better, but at least this approach is simple. - - -- groupScan operates on flattened arrays. This does not - -- involve copying anything; merely playing with the index - -- function. - dims_flat <- dPrimV "dims_flat" $ product dims' - let segment_size = last dims' - crossesSegment from to = - (sExt64 to - sExt64 from) .>. (sExt64 to `rem` sExt64 segment_size) - - forM_ (zip ops tmps_for_ops) $ \(op, tmps) -> do - tmps_flat <- mapM (flattenArray (length dims') dims_flat) tmps - groupScan - (Just crossesSegment) - (product dims') - (product dims') - (segBinOpLambda op) - tmps_flat - - sOp $ Imp.ErrorSync Imp.FenceLocal - - forM_ (zip red_pes $ concat tmps_for_ops) $ \(pe, arr) -> - copyDWIM - (patElemName pe) - [] - (Var arr) - (map (unitSlice 0) (init dims') ++ [DimFix $ last dims' - 1]) - - sOp $ Imp.Barrier Imp.FenceLocal -compileGroupOp pat (Inner (SegOp (SegHist lvl space ops _ kbody))) = do - compileFlatId lvl space - let (ltids, _dims) = unzip $ unSegSpace space - - -- We don't need the red_pes, because it is guaranteed by our type - -- rules that they occupy the same memory as the destinations for - -- the ops. - let num_red_res = length ops + sum (map (length . histNeutral) ops) - (_red_pes, map_pes) = - splitAt num_red_res $ patElems pat - - ops' <- prepareIntraGroupSegHist (segGroupSize lvl) ops - - -- Ensure that all locks have been initialised. - sOp $ Imp.Barrier Imp.FenceLocal - - groupCoverSegSpace (segVirt lvl) space $ - compileStms mempty (kernelBodyStms kbody) $ do - let (red_res, map_res) = splitAt num_red_res $ kernelBodyResult kbody - (red_is, red_vs) = splitAt (length ops) $ map kernelResultSubExp red_res - zipWithM_ (compileThreadResult space) map_pes map_res - - let vs_per_op = chunks (map (length . histDest) ops) red_vs - - forM_ (zip4 red_is vs_per_op ops' ops) $ - \(bin, op_vs, do_op, HistOp dest_shape _ _ _ shape lam) -> do - let bin' = pe64 bin - dest_shape' = map pe64 $ shapeDims dest_shape - bin_in_bounds = inBounds (Slice (map DimFix [bin'])) dest_shape' - bin_is = map Imp.le64 (init ltids) ++ [bin'] - vs_params = takeLast (length op_vs) $ lambdaParams lam - - sComment "perform atomic updates" $ - sWhen bin_in_bounds $ do - dLParams $ lambdaParams lam - sLoopNest shape $ \is -> do - forM_ (zip vs_params op_vs) $ \(p, v) -> - copyDWIMFix (paramName p) [] v is - do_op (bin_is ++ is) - - sOp $ Imp.ErrorSync Imp.FenceLocal -compileGroupOp pat _ = - compilerBugS $ "compileGroupOp: cannot compile rhs of binding " ++ pretty pat - -compileThreadOp :: OpCompiler GPUMem KernelEnv Imp.KernelOp -compileThreadOp pat (Alloc size space) = - kernelAlloc pat size space -compileThreadOp pat (Inner (SizeOp (SplitSpace o w i elems_per_thread))) = - splitSpace pat o w i elems_per_thread -compileThreadOp pat _ = - compilerBugS $ "compileThreadOp: cannot compile rhs of binding " ++ pretty pat - --- | Locking strategy used for an atomic update. -data Locking = Locking - { -- | Array containing the lock. - lockingArray :: VName, - -- | Value for us to consider the lock free. - lockingIsUnlocked :: Imp.TExp Int32, - -- | What to write when we lock it. - lockingToLock :: Imp.TExp Int32, - -- | What to write when we unlock it. - lockingToUnlock :: Imp.TExp Int32, - -- | A transformation from the logical lock index to the - -- physical position in the array. This can also be used - -- to make the lock array smaller. - lockingMapping :: [Imp.TExp Int64] -> [Imp.TExp Int64] - } - --- | A function for generating code for an atomic update. Assumes --- that the bucket is in-bounds. -type DoAtomicUpdate rep r = - Space -> [VName] -> [Imp.TExp Int64] -> ImpM rep r Imp.KernelOp () - --- | The mechanism that will be used for performing the atomic update. --- Approximates how efficient it will be. Ordered from most to least --- efficient. -data AtomicUpdate rep r - = -- | Supported directly by primitive. - AtomicPrim (DoAtomicUpdate rep r) - | -- | Can be done by efficient swaps. - AtomicCAS (DoAtomicUpdate rep r) - | -- | Requires explicit locking. - AtomicLocking (Locking -> DoAtomicUpdate rep r) - --- | Is there an atomic t'BinOp' corresponding to this t'BinOp'? -type AtomicBinOp = - BinOp -> - Maybe (VName -> VName -> Count Imp.Elements (Imp.TExp Int64) -> Imp.Exp -> Imp.AtomicOp) - --- | Do an atomic update corresponding to a binary operator lambda. -atomicUpdateLocking :: - AtomicBinOp -> - Lambda GPUMem -> - AtomicUpdate GPUMem KernelEnv -atomicUpdateLocking atomicBinOp lam - | Just ops_and_ts <- lamIsBinOp lam, - all (\(_, t, _, _) -> primBitSize t `elem` [32, 64]) ops_and_ts = - primOrCas ops_and_ts $ \space arrs bucket -> - -- If the operator is a vectorised binary operator on 32/64-bit - -- values, we can use a particularly efficient - -- implementation. If the operator has an atomic implementation - -- we use that, otherwise it is still a binary operator which - -- can be implemented by atomic compare-and-swap if 32/64 bits. - forM_ (zip arrs ops_and_ts) $ \(a, (op, t, x, y)) -> do - -- Common variables. - old <- dPrim "old" t - - (arr', _a_space, bucket_offset) <- fullyIndexArray a bucket - - case opHasAtomicSupport space (tvVar old) arr' bucket_offset op of - Just f -> sOp $ f $ Imp.var y t - Nothing -> - atomicUpdateCAS space t a (tvVar old) bucket x $ - x <~~ Imp.BinOpExp op (Imp.var x t) (Imp.var y t) - where - opHasAtomicSupport space old arr' bucket' bop = do - let atomic f = Imp.Atomic space . f old arr' bucket' - atomic <$> atomicBinOp bop - - primOrCas ops - | all isPrim ops = AtomicPrim - | otherwise = AtomicCAS - - isPrim (op, _, _, _) = isJust $ atomicBinOp op - --- If the operator functions purely on single 32/64-bit values, we can --- use an implementation based on CAS, no matter what the operator --- does. -atomicUpdateLocking _ op - | [Prim t] <- lambdaReturnType op, - [xp, _] <- lambdaParams op, - primBitSize t `elem` [32, 64] = AtomicCAS $ \space [arr] bucket -> do - old <- dPrim "old" t - atomicUpdateCAS space t arr (tvVar old) bucket (paramName xp) $ - compileBody' [xp] $ - lambdaBody op -atomicUpdateLocking _ op = AtomicLocking $ \locking space arrs bucket -> do - old <- dPrim "old" int32 - continue <- dPrimVol "continue" Bool true - - -- Correctly index into locks. - (locks', _locks_space, locks_offset) <- - fullyIndexArray (lockingArray locking) $ lockingMapping locking bucket - - -- Critical section - let try_acquire_lock = - sOp $ - Imp.Atomic space $ - Imp.AtomicCmpXchg - int32 - (tvVar old) - locks' - locks_offset - (untyped $ lockingIsUnlocked locking) - (untyped $ lockingToLock locking) - lock_acquired = tvExp old .==. lockingIsUnlocked locking - -- Even the releasing is done with an atomic rather than a - -- simple write, for memory coherency reasons. - release_lock = - sOp $ - Imp.Atomic space $ - Imp.AtomicCmpXchg - int32 - (tvVar old) - locks' - locks_offset - (untyped $ lockingToLock locking) - (untyped $ lockingToUnlock locking) - break_loop = continue <-- false - - -- Preparing parameters. It is assumed that the caller has already - -- filled the arr_params. We copy the current value to the - -- accumulator parameters. - -- - -- Note the use of 'everythingVolatile' when reading and writing the - -- buckets. This was necessary to ensure correct execution on a - -- newer NVIDIA GPU (RTX 2080). The 'volatile' modifiers likely - -- make the writes pass through the (SM-local) L1 cache, which is - -- necessary here, because we are really doing device-wide - -- synchronisation without atomics (naughty!). - let (acc_params, _arr_params) = splitAt (length arrs) $ lambdaParams op - bind_acc_params = - everythingVolatile $ - sComment "bind lhs" $ - forM_ (zip acc_params arrs) $ \(acc_p, arr) -> - copyDWIMFix (paramName acc_p) [] (Var arr) bucket - - let op_body = - sComment "execute operation" $ - compileBody' acc_params $ - lambdaBody op - - do_hist = - everythingVolatile $ - sComment "update global result" $ - zipWithM_ (writeArray bucket) arrs $ - map (Var . paramName) acc_params - - fence = sOp $ Imp.MemFence $ fenceForSpace space - - -- While-loop: Try to insert your value - sWhile (tvExp continue) $ do - try_acquire_lock - sWhen lock_acquired $ do - dLParams acc_params - bind_acc_params - op_body - do_hist - fence - release_lock - break_loop - fence - where - writeArray bucket arr val = copyDWIMFix arr bucket val [] - -atomicUpdateCAS :: - Space -> - PrimType -> - VName -> - VName -> - [Imp.TExp Int64] -> - VName -> - InKernelGen () -> - InKernelGen () -atomicUpdateCAS space t arr old bucket x do_op = do - -- Code generation target: - -- - -- old = d_his[idx]; - -- do { - -- assumed = old; - -- x = do_op(assumed, y); - -- old = atomicCAS(&d_his[idx], assumed, tmp); - -- } while(assumed != old); - assumed <- tvVar <$> dPrim "assumed" t - run_loop <- dPrimV "run_loop" true - - -- XXX: CUDA may generate really bad code if this is not a volatile - -- read. Unclear why. The later reads are volatile, so maybe - -- that's it. - everythingVolatile $ copyDWIMFix old [] (Var arr) bucket - - (arr', _a_space, bucket_offset) <- fullyIndexArray arr bucket - - -- While-loop: Try to insert your value - let (toBits, fromBits) = - case t of - FloatType Float16 -> - ( \v -> Imp.FunExp "to_bits16" [v] int16, - \v -> Imp.FunExp "from_bits16" [v] t - ) - FloatType Float32 -> - ( \v -> Imp.FunExp "to_bits32" [v] int32, - \v -> Imp.FunExp "from_bits32" [v] t - ) - FloatType Float64 -> - ( \v -> Imp.FunExp "to_bits64" [v] int64, - \v -> Imp.FunExp "from_bits64" [v] t - ) - _ -> (id, id) - - int - | primBitSize t == 16 = int16 - | primBitSize t == 32 = int32 - | otherwise = int64 - - sWhile (tvExp run_loop) $ do - assumed <~~ Imp.var old t - x <~~ Imp.var assumed t - do_op - old_bits_v <- newVName "old_bits" - dPrim_ old_bits_v int - let old_bits = Imp.var old_bits_v int - sOp . Imp.Atomic space $ - Imp.AtomicCmpXchg - int - old_bits_v - arr' - bucket_offset - (toBits (Imp.var assumed t)) - (toBits (Imp.var x t)) - old <~~ fromBits old_bits - let won = CmpOpExp (CmpEq int) (toBits (Imp.var assumed t)) old_bits - sWhen (isBool won) (run_loop <-- false) - -computeKernelUses :: - FreeIn a => - a -> - [VName] -> - CallKernelGen [Imp.KernelUse] -computeKernelUses kernel_body bound_in_kernel = do - let actually_free = freeIn kernel_body `namesSubtract` namesFromList bound_in_kernel - -- Compute the variables that we need to pass to the kernel. - nubOrd <$> readsFromSet actually_free - -readsFromSet :: Names -> CallKernelGen [Imp.KernelUse] -readsFromSet = fmap catMaybes . mapM f . namesToList - where - f var = do - t <- lookupType var - vtable <- getVTable - case t of - Array {} -> pure Nothing - Acc {} -> pure Nothing - Mem (Space "local") -> pure Nothing - Mem {} -> pure $ Just $ Imp.MemoryUse var - Prim bt -> - isConstExp vtable (Imp.var var bt) >>= \case - Just ce -> pure $ Just $ Imp.ConstUse var ce - Nothing -> pure $ Just $ Imp.ScalarUse var bt - -isConstExp :: - VTable GPUMem -> - Imp.Exp -> - ImpM rep r op (Maybe Imp.KernelConstExp) -isConstExp vtable size = do - fname <- askFunction - let onLeaf name _ = lookupConstExp name - lookupConstExp name = - constExp =<< hasExp =<< M.lookup name vtable - constExp (Op (Inner (SizeOp (GetSize key _)))) = - Just $ LeafExp (Imp.SizeConst $ keyWithEntryPoint fname key) int32 - constExp e = primExpFromExp lookupConstExp e - pure $ replaceInPrimExpM onLeaf size - where - hasExp (ArrayVar e _) = e - hasExp (AccVar e _) = e - hasExp (ScalarVar e _) = e - hasExp (MemVar e _) = e - -computeThreadChunkSize :: - SplitOrdering -> - Imp.TExp Int64 -> - Imp.Count Imp.Elements (Imp.TExp Int64) -> - Imp.Count Imp.Elements (Imp.TExp Int64) -> - TV Int64 -> - ImpM rep r op () -computeThreadChunkSize (SplitStrided stride) thread_index elements_per_thread num_elements chunk_var = - chunk_var - <-- sMin64 - (Imp.unCount elements_per_thread) - ((Imp.unCount num_elements - thread_index) `divUp` pe64 stride) -computeThreadChunkSize SplitContiguous thread_index elements_per_thread num_elements chunk_var = do - starting_point <- - dPrimV "starting_point" $ - thread_index * Imp.unCount elements_per_thread - remaining_elements <- - dPrimV "remaining_elements" $ - Imp.unCount num_elements - tvExp starting_point - - let no_remaining_elements = tvExp remaining_elements .<=. 0 - beyond_bounds = Imp.unCount num_elements .<=. tvExp starting_point - - sIf - (no_remaining_elements .||. beyond_bounds) - (chunk_var <-- 0) - ( sIf - is_last_thread - (chunk_var <-- Imp.unCount last_thread_elements) - (chunk_var <-- Imp.unCount elements_per_thread) - ) - where - last_thread_elements = - num_elements - Imp.elements thread_index * elements_per_thread - is_last_thread = - Imp.unCount num_elements - .<. (thread_index + 1) - * Imp.unCount elements_per_thread - -kernelInitialisationSimple :: - Count NumGroups SubExp -> - Count GroupSize SubExp -> - CallKernelGen (KernelConstants, InKernelGen ()) -kernelInitialisationSimple num_groups group_size = do - global_tid <- newVName "global_tid" - local_tid <- newVName "local_tid" - group_id <- newVName "group_tid" - wave_size <- newVName "wave_size" - inner_group_size <- newVName "group_size" - let num_groups' = Imp.pe64 (unCount num_groups) - group_size' = Imp.pe64 (unCount group_size) - constants = - KernelConstants - { kernelGlobalThreadId = Imp.le32 global_tid, - kernelLocalThreadId = Imp.le32 local_tid, - kernelGroupId = Imp.le32 group_id, - kernelGlobalThreadIdVar = global_tid, - kernelLocalThreadIdVar = local_tid, - kernelNumGroupsCount = num_groups, - kernelGroupSizeCount = group_size, - kernelGroupIdVar = group_id, - kernelNumGroups = num_groups', - kernelGroupSize = group_size', - kernelNumThreads = sExt32 (group_size' * num_groups'), - kernelWaveSize = Imp.le32 wave_size, - kernelLocalIdMap = mempty, - kernelChunkItersMap = mempty - } - - let set_constants = do - dPrim_ local_tid int32 - dPrim_ inner_group_size int64 - dPrim_ wave_size int32 - dPrim_ group_id int32 - - sOp (Imp.GetLocalId local_tid 0) - sOp (Imp.GetLocalSize inner_group_size 0) - sOp (Imp.GetLockstepWidth wave_size) - sOp (Imp.GetGroupId group_id 0) - dPrimV_ global_tid $ le32 group_id * le32 inner_group_size + le32 local_tid - - pure (constants, set_constants) - -isActive :: [(VName, SubExp)] -> Imp.TExp Bool -isActive limit = case actives of - [] -> true - x : xs -> foldl (.&&.) x xs - where - (is, ws) = unzip limit - actives = zipWith active is $ map pe64 ws - active i = (Imp.le64 i .<.) - --- | Change every memory block to be in the global address space, --- except those who are in the local memory space. This only affects --- generated code - we still need to make sure that the memory is --- actually present on the device (and declared as variables in the --- kernel). -makeAllMemoryGlobal :: CallKernelGen a -> CallKernelGen a -makeAllMemoryGlobal = - localDefaultSpace (Imp.Space "global") . localVTable (M.map globalMemory) - where - globalMemory (MemVar _ entry) - | entryMemSpace entry /= Space "local" = - MemVar Nothing entry {entryMemSpace = Imp.Space "global"} - globalMemory entry = - entry - -groupReduce :: - Imp.TExp Int32 -> - Lambda GPUMem -> - [VName] -> - InKernelGen () -groupReduce w lam arrs = do - offset <- dPrim "offset" int32 - groupReduceWithOffset offset w lam arrs - -groupReduceWithOffset :: - TV Int32 -> - Imp.TExp Int32 -> - Lambda GPUMem -> - [VName] -> - InKernelGen () -groupReduceWithOffset offset w lam arrs = do - constants <- kernelConstants <$> askEnv - - let local_tid = kernelLocalThreadId constants - global_tid = kernelGlobalThreadId constants - - barrier - | all primType $ lambdaReturnType lam = sOp $ Imp.Barrier Imp.FenceLocal - | otherwise = sOp $ Imp.Barrier Imp.FenceGlobal - - readReduceArgument param arr - | Prim _ <- paramType param = do - let i = local_tid + tvExp offset - copyDWIMFix (paramName param) [] (Var arr) [sExt64 i] - | otherwise = do - let i = global_tid + tvExp offset - copyDWIMFix (paramName param) [] (Var arr) [sExt64 i] - - writeReduceOpResult param arr - | Prim _ <- paramType param = - copyDWIMFix arr [sExt64 local_tid] (Var $ paramName param) [] - | otherwise = - pure () - - let (reduce_acc_params, reduce_arr_params) = splitAt (length arrs) $ lambdaParams lam - - skip_waves <- dPrimV "skip_waves" (1 :: Imp.TExp Int32) - dLParams $ lambdaParams lam - - offset <-- (0 :: Imp.TExp Int32) - - comment "participating threads read initial accumulator" $ - sWhen (local_tid .<. w) $ - zipWithM_ readReduceArgument reduce_acc_params arrs - - let do_reduce = do - comment "read array element" $ - zipWithM_ readReduceArgument reduce_arr_params arrs - comment "apply reduction operation" $ - compileBody' reduce_acc_params $ - lambdaBody lam - comment "write result of operation" $ - zipWithM_ writeReduceOpResult reduce_acc_params arrs - in_wave_reduce = everythingVolatile do_reduce - - wave_size = kernelWaveSize constants - group_size = kernelGroupSize constants - wave_id = local_tid `quot` wave_size - in_wave_id = local_tid - wave_id * wave_size - num_waves = (sExt32 group_size + wave_size - 1) `quot` wave_size - arg_in_bounds = local_tid + tvExp offset .<. w - - doing_in_wave_reductions = - tvExp offset .<. wave_size - apply_in_in_wave_iteration = - (in_wave_id .&. (2 * tvExp offset - 1)) .==. 0 - in_wave_reductions = do - offset <-- (1 :: Imp.TExp Int32) - sWhile doing_in_wave_reductions $ do - sWhen - (arg_in_bounds .&&. apply_in_in_wave_iteration) - in_wave_reduce - offset <-- tvExp offset * 2 - - doing_cross_wave_reductions = - tvExp skip_waves .<. num_waves - is_first_thread_in_wave = - in_wave_id .==. 0 - wave_not_skipped = - (wave_id .&. (2 * tvExp skip_waves - 1)) .==. 0 - apply_in_cross_wave_iteration = - arg_in_bounds .&&. is_first_thread_in_wave .&&. wave_not_skipped - cross_wave_reductions = - sWhile doing_cross_wave_reductions $ do - barrier - offset <-- tvExp skip_waves * wave_size - sWhen - apply_in_cross_wave_iteration - do_reduce - skip_waves <-- tvExp skip_waves * 2 - - in_wave_reductions - cross_wave_reductions - -groupScan :: - Maybe (Imp.TExp Int32 -> Imp.TExp Int32 -> Imp.TExp Bool) -> - Imp.TExp Int64 -> - Imp.TExp Int64 -> - Lambda GPUMem -> - [VName] -> - InKernelGen () -groupScan seg_flag arrs_full_size w lam arrs = do - constants <- kernelConstants <$> askEnv - renamed_lam <- renameLambda lam - - let ltid32 = kernelLocalThreadId constants - ltid = sExt64 ltid32 - (x_params, y_params) = splitAt (length arrs) $ lambdaParams lam - - dLParams (lambdaParams lam ++ lambdaParams renamed_lam) - - ltid_in_bounds <- dPrimVE "ltid_in_bounds" $ ltid .<. w - - fence <- fenceForArrays arrs - - -- The scan works by splitting the group into blocks, which are - -- scanned separately. Typically, these blocks are smaller than - -- the lockstep width, which enables barrier-free execution inside - -- them. - -- - -- We hardcode the block size here. The only requirement is that - -- it should not be less than the square root of the group size. - -- With 32, we will work on groups of size 1024 or smaller, which - -- fits every device Troels has seen. Still, it would be nicer if - -- it were a runtime parameter. Some day. - let block_size = 32 - simd_width = kernelWaveSize constants - block_id = ltid32 `quot` block_size - in_block_id = ltid32 - block_id * block_size - doInBlockScan seg_flag' active = - inBlockScan - constants - seg_flag' - arrs_full_size - simd_width - block_size - active - arrs - barrier - array_scan = not $ all primType $ lambdaReturnType lam - barrier - | array_scan = - sOp $ Imp.Barrier Imp.FenceGlobal - | otherwise = - sOp $ Imp.Barrier fence - - group_offset = sExt64 (kernelGroupId constants) * kernelGroupSize constants - - writeBlockResult p arr - | primType $ paramType p = - copyDWIM arr [DimFix $ sExt64 block_id] (Var $ paramName p) [] - | otherwise = - copyDWIM arr [DimFix $ group_offset + sExt64 block_id] (Var $ paramName p) [] - - readPrevBlockResult p arr - | primType $ paramType p = - copyDWIM (paramName p) [] (Var arr) [DimFix $ sExt64 block_id - 1] - | otherwise = - copyDWIM (paramName p) [] (Var arr) [DimFix $ group_offset + sExt64 block_id - 1] - - doInBlockScan seg_flag ltid_in_bounds lam - barrier - - let is_first_block = block_id .==. 0 - when array_scan $ do - sComment "save correct values for first block" $ - sWhen is_first_block $ - forM_ (zip x_params arrs) $ \(x, arr) -> - unless (primType $ paramType x) $ - copyDWIM arr [DimFix $ arrs_full_size + group_offset + sExt64 block_size + ltid] (Var $ paramName x) [] - - barrier - - let last_in_block = in_block_id .==. block_size - 1 - sComment "last thread of block 'i' writes its result to offset 'i'" $ - sWhen (last_in_block .&&. ltid_in_bounds) $ - everythingVolatile $ - zipWithM_ writeBlockResult x_params arrs - - barrier - - let first_block_seg_flag = do - flag_true <- seg_flag - Just $ \from to -> - flag_true (from * block_size + block_size - 1) (to * block_size + block_size - 1) - comment - "scan the first block, after which offset 'i' contains carry-in for block 'i+1'" - $ doInBlockScan first_block_seg_flag (is_first_block .&&. ltid_in_bounds) renamed_lam - - barrier - - when array_scan $ do - sComment "move correct values for first block back a block" $ - sWhen is_first_block $ - forM_ (zip x_params arrs) $ \(x, arr) -> - unless (primType $ paramType x) $ - copyDWIM - arr - [DimFix $ arrs_full_size + group_offset + ltid] - (Var arr) - [DimFix $ arrs_full_size + group_offset + sExt64 block_size + ltid] - - barrier - - no_carry_in <- dPrimVE "no_carry_in" $ is_first_block .||. bNot ltid_in_bounds - - let read_carry_in = sUnless no_carry_in $ do - forM_ (zip x_params y_params) $ \(x, y) -> - copyDWIM (paramName y) [] (Var (paramName x)) [] - zipWithM_ readPrevBlockResult x_params arrs - - op_to_x - | Nothing <- seg_flag = - sUnless no_carry_in $ compileBody' x_params $ lambdaBody lam - | Just flag_true <- seg_flag = do - inactive <- - dPrimVE "inactive" $ flag_true (block_id * block_size - 1) ltid32 - sUnless no_carry_in . sWhen inactive . forM_ (zip x_params y_params) $ \(x, y) -> - copyDWIM (paramName x) [] (Var (paramName y)) [] - -- The convoluted control flow is to ensure all threads - -- hit this barrier (if applicable). - when array_scan barrier - sUnless no_carry_in $ sUnless inactive $ compileBody' x_params $ lambdaBody lam - - write_final_result = - forM_ (zip x_params arrs) $ \(p, arr) -> - when (primType $ paramType p) $ - copyDWIM arr [DimFix ltid] (Var $ paramName p) [] - - sComment "carry-in for every block except the first" $ do - sComment "read operands" read_carry_in - sComment "perform operation" op_to_x - sComment "write final result" $ sUnless no_carry_in write_final_result - - barrier - - sComment "restore correct values for first block" $ - sWhen (is_first_block .&&. ltid_in_bounds) $ - forM_ (zip3 x_params y_params arrs) $ \(x, y, arr) -> - if primType (paramType y) - then copyDWIM arr [DimFix ltid] (Var $ paramName y) [] - else copyDWIM (paramName x) [] (Var arr) [DimFix $ arrs_full_size + group_offset + ltid] - - barrier - -inBlockScan :: - KernelConstants -> - Maybe (Imp.TExp Int32 -> Imp.TExp Int32 -> Imp.TExp Bool) -> - Imp.TExp Int64 -> - Imp.TExp Int32 -> - Imp.TExp Int32 -> - Imp.TExp Bool -> - [VName] -> - InKernelGen () -> - Lambda GPUMem -> - InKernelGen () -inBlockScan constants seg_flag arrs_full_size lockstep_width block_size active arrs barrier scan_lam = everythingVolatile $ do - skip_threads <- dPrim "skip_threads" int32 - let actual_params = lambdaParams scan_lam - (x_params, y_params) = - splitAt (length actual_params `div` 2) actual_params - y_to_x = - forM_ (zip x_params y_params) $ \(x, y) -> - when (primType (paramType x)) $ - copyDWIM (paramName x) [] (Var (paramName y)) [] - - -- Set initial y values - sComment "read input for in-block scan" $ - sWhen active $ do - zipWithM_ readInitial y_params arrs - -- Since the final result is expected to be in x_params, we may - -- need to copy it there for the first thread in the block. - sWhen (in_block_id .==. 0) y_to_x - - when array_scan barrier - - let op_to_x in_block_thread_active - | Nothing <- seg_flag = - sWhen in_block_thread_active $ - compileBody' x_params $ - lambdaBody scan_lam - | Just flag_true <- seg_flag = do - inactive <- - dPrimVE "inactive" $ flag_true (ltid32 - tvExp skip_threads) ltid32 - sWhen (in_block_thread_active .&&. inactive) $ - forM_ (zip x_params y_params) $ \(x, y) -> - copyDWIM (paramName x) [] (Var (paramName y)) [] - -- The convoluted control flow is to ensure all threads - -- hit this barrier (if applicable). - when array_scan barrier - sWhen in_block_thread_active . sUnless inactive $ - compileBody' x_params $ - lambdaBody scan_lam - - maybeBarrier = - sWhen - (lockstep_width .<=. tvExp skip_threads) - barrier - - sComment "in-block scan (hopefully no barriers needed)" $ do - skip_threads <-- 1 - sWhile (tvExp skip_threads .<. block_size) $ do - thread_active <- - dPrimVE "thread_active" $ tvExp skip_threads .<=. in_block_id .&&. active - - sWhen thread_active . sComment "read operands" $ - zipWithM_ (readParam (sExt64 $ tvExp skip_threads)) x_params arrs - sComment "perform operation" $ op_to_x thread_active - - maybeBarrier - - sWhen thread_active . sComment "write result" $ - sequence_ $ - zipWith3 writeResult x_params y_params arrs - - maybeBarrier - - skip_threads <-- tvExp skip_threads * 2 - where - block_id = ltid32 `quot` block_size - in_block_id = ltid32 - block_id * block_size - ltid32 = kernelLocalThreadId constants - ltid = sExt64 ltid32 - gtid = sExt64 $ kernelGlobalThreadId constants - array_scan = not $ all primType $ lambdaReturnType scan_lam - - readInitial p arr - | primType $ paramType p = - copyDWIM (paramName p) [] (Var arr) [DimFix ltid] - | otherwise = - copyDWIM (paramName p) [] (Var arr) [DimFix gtid] - - readParam behind p arr - | primType $ paramType p = - copyDWIM (paramName p) [] (Var arr) [DimFix $ ltid - behind] - | otherwise = - copyDWIM (paramName p) [] (Var arr) [DimFix $ gtid - behind + arrs_full_size] - - writeResult x y arr - | primType $ paramType x = do - copyDWIM arr [DimFix ltid] (Var $ paramName x) [] - copyDWIM (paramName y) [] (Var $ paramName x) [] - | otherwise = - copyDWIM (paramName y) [] (Var $ paramName x) [] - -simpleKernelGroups :: - Imp.TExp Int64 -> - Imp.TExp Int64 -> - CallKernelGen (Imp.TExp Int32, Count NumGroups SubExp, Count GroupSize SubExp) -simpleKernelGroups max_num_groups kernel_size = do - group_size <- dPrim "group_size" int64 - fname <- askFunction - let group_size_key = keyWithEntryPoint fname $ nameFromString $ pretty $ tvVar group_size - sOp $ Imp.GetSize (tvVar group_size) group_size_key Imp.SizeGroup - virt_num_groups <- dPrimVE "virt_num_groups" $ kernel_size `divUp` tvExp group_size - num_groups <- dPrimV "num_groups" $ virt_num_groups `sMin64` max_num_groups - pure (sExt32 virt_num_groups, Count $ tvSize num_groups, Count $ tvSize group_size) - -simpleKernelConstants :: - Imp.TExp Int64 -> - String -> - CallKernelGen - ( (Imp.TExp Int64 -> InKernelGen ()) -> InKernelGen (), - KernelConstants - ) -simpleKernelConstants kernel_size desc = do - -- For performance reasons, codegen assumes that the thread count is - -- never more than will fit in an i32. This means we need to cap - -- the number of groups here. The cap is set much higher than any - -- GPU will possibly need. Feel free to come back and laugh at me - -- in the future. - let max_num_groups = 1024 * 1024 - thread_gtid <- newVName $ desc ++ "_gtid" - thread_ltid <- newVName $ desc ++ "_ltid" - group_id <- newVName $ desc ++ "_gid" - inner_group_size <- newVName "group_size" - (virt_num_groups, num_groups, group_size) <- - simpleKernelGroups max_num_groups kernel_size - let group_size' = Imp.pe64 $ unCount group_size - num_groups' = Imp.pe64 $ unCount num_groups - - constants = - KernelConstants - { kernelGlobalThreadId = Imp.le32 thread_gtid, - kernelLocalThreadId = Imp.le32 thread_ltid, - kernelGroupId = Imp.le32 group_id, - kernelGlobalThreadIdVar = thread_gtid, - kernelLocalThreadIdVar = thread_ltid, - kernelGroupIdVar = group_id, - kernelNumGroupsCount = num_groups, - kernelGroupSizeCount = group_size, - kernelNumGroups = num_groups', - kernelGroupSize = group_size', - kernelNumThreads = sExt32 (group_size' * num_groups'), - kernelWaveSize = 0, - kernelLocalIdMap = mempty, - kernelChunkItersMap = mempty - } - - wrapKernel m = do - dPrim_ thread_ltid int32 - dPrim_ inner_group_size int64 - dPrim_ group_id int32 - sOp (Imp.GetLocalId thread_ltid 0) - sOp (Imp.GetLocalSize inner_group_size 0) - sOp (Imp.GetGroupId group_id 0) - dPrimV_ thread_gtid $ le32 group_id * le32 inner_group_size + le32 thread_ltid - virtualiseGroups SegVirt virt_num_groups $ \virt_group_id -> do - global_tid <- - dPrimVE "global_tid" $ - sExt64 virt_group_id * sExt64 (le32 inner_group_size) - + sExt64 (kernelLocalThreadId constants) - m global_tid - - pure (wrapKernel, constants) - --- | For many kernels, we may not have enough physical groups to cover --- the logical iteration space. Some groups thus have to perform --- double duty; we put an outer loop to accomplish this. The --- advantage over just launching a bazillion threads is that the cost --- of memory expansion should be proportional to the number of --- *physical* threads (hardware parallelism), not the amount of --- application parallelism. -virtualiseGroups :: - SegVirt -> - Imp.TExp Int32 -> - (Imp.TExp Int32 -> InKernelGen ()) -> - InKernelGen () -virtualiseGroups SegVirt required_groups m = do - constants <- kernelConstants <$> askEnv - phys_group_id <- dPrim "phys_group_id" int32 - sOp $ Imp.GetGroupId (tvVar phys_group_id) 0 - iterations <- - dPrimVE "iterations" $ - (required_groups - tvExp phys_group_id) `divUp` sExt32 (kernelNumGroups constants) - - sFor "i" iterations $ \i -> do - m . tvExp - =<< dPrimV - "virt_group_id" - (tvExp phys_group_id + i * sExt32 (kernelNumGroups constants)) - -- Make sure the virtual group is actually done before we let - -- another virtual group have its way with it. - sOp $ Imp.Barrier Imp.FenceGlobal -virtualiseGroups _ _ m = do - gid <- kernelGroupIdVar . kernelConstants <$> askEnv - m $ Imp.le32 gid - --- | Various extra configuration of the kernel being generated. -data KernelAttrs = KernelAttrs - { -- | Can this kernel execute correctly even if previous kernels failed? - kAttrFailureTolerant :: Bool, - -- | Does whatever launch this kernel check for local memory capacity itself? - kAttrCheckLocalMemory :: Bool, - -- | Number of groups. - kAttrNumGroups :: Count NumGroups SubExp, - -- | Group size. - kAttrGroupSize :: Count GroupSize SubExp - } - --- | The default kernel attributes. -defKernelAttrs :: - Count NumGroups SubExp -> - Count GroupSize SubExp -> - KernelAttrs -defKernelAttrs num_groups group_size = - KernelAttrs - { kAttrFailureTolerant = False, - kAttrCheckLocalMemory = True, - kAttrNumGroups = num_groups, - kAttrGroupSize = group_size - } - -sKernel :: - Operations GPUMem KernelEnv Imp.KernelOp -> - (KernelConstants -> Imp.TExp Int32) -> - String -> - VName -> - KernelAttrs -> - InKernelGen () -> - CallKernelGen () -sKernel ops flatf name v attrs f = do - (constants, set_constants) <- - kernelInitialisationSimple (kAttrNumGroups attrs) (kAttrGroupSize attrs) - name' <- nameForFun $ name ++ "_" ++ show (baseTag v) - sKernelOp attrs constants ops name' $ do - set_constants - dPrimV_ v $ flatf constants - f - -sKernelThread :: - String -> - VName -> - KernelAttrs -> - InKernelGen () -> - CallKernelGen () -sKernelThread = sKernel threadOperations kernelGlobalThreadId - -sKernelGroup :: - String -> - VName -> - KernelAttrs -> - InKernelGen () -> - CallKernelGen () -sKernelGroup = sKernel groupOperations kernelGroupId - -sKernelOp :: - KernelAttrs -> - KernelConstants -> - Operations GPUMem KernelEnv Imp.KernelOp -> - Name -> - InKernelGen () -> - CallKernelGen () -sKernelOp attrs constants ops name m = do - HostEnv atomics _ locks <- askEnv - body <- makeAllMemoryGlobal $ subImpM_ (KernelEnv atomics constants locks) ops m - uses <- computeKernelUses body mempty - emit . Imp.Op . Imp.CallKernel $ - Imp.Kernel - { Imp.kernelBody = body, - Imp.kernelUses = uses, - Imp.kernelNumGroups = [untyped $ kernelNumGroups constants], - Imp.kernelGroupSize = [untyped $ kernelGroupSize constants], - Imp.kernelName = name, - Imp.kernelFailureTolerant = kAttrFailureTolerant attrs, - Imp.kernelCheckLocalMemory = kAttrCheckLocalMemory attrs - } - -sKernelFailureTolerant :: - Bool -> - Operations GPUMem KernelEnv Imp.KernelOp -> - KernelConstants -> - Name -> - InKernelGen () -> - CallKernelGen () -sKernelFailureTolerant tol ops constants name m = do - sKernelOp attrs constants ops name m - where - attrs = - ( defKernelAttrs - (kernelNumGroupsCount constants) - (kernelGroupSizeCount constants) - ) - { kAttrFailureTolerant = tol - } - -copyInGroup :: CopyCompiler GPUMem KernelEnv Imp.KernelOp -copyInGroup pt destloc srcloc = do - dest_space <- entryMemSpace <$> lookupMemory (memLocName destloc) - src_space <- entryMemSpace <$> lookupMemory (memLocName srcloc) - - let src_ixfun = memLocIxFun srcloc - dims = IxFun.shape src_ixfun - rank = length dims - - case (dest_space, src_space) of - (ScalarSpace destds _, ScalarSpace srcds _) -> do - let fullDim d = DimSlice 0 d 1 - destslice' = - Slice $ - replicate (rank - length destds) (DimFix 0) - ++ takeLast (length destds) (map fullDim dims) - srcslice' = - Slice $ - replicate (rank - length srcds) (DimFix 0) - ++ takeLast (length srcds) (map fullDim dims) - copyElementWise - pt - (sliceMemLoc destloc destslice') - (sliceMemLoc srcloc srcslice') - _ -> do - groupCoverSpace (map sExt32 dims) $ \is -> - copyElementWise - pt - (sliceMemLoc destloc (Slice $ map (DimFix . sExt64) is)) - (sliceMemLoc srcloc (Slice $ map (DimFix . sExt64) is)) - sOp $ Imp.Barrier Imp.FenceLocal - -threadOperations, groupOperations :: Operations GPUMem KernelEnv Imp.KernelOp -threadOperations = - (defaultOperations compileThreadOp) - { opsCopyCompiler = copyElementWise, - opsExpCompiler = compileThreadExp, - opsStmsCompiler = \_ -> defCompileStms mempty, - opsAllocCompilers = - M.fromList [(Space "local", allocLocal)] - } -groupOperations = - (defaultOperations compileGroupOp) - { opsCopyCompiler = copyInGroup, - opsExpCompiler = compileGroupExp, - opsStmsCompiler = \_ -> defCompileStms mempty, - opsAllocCompilers = - M.fromList [(Space "local", allocLocal)] - } - --- | Perform a Replicate with a kernel. -sReplicateKernel :: VName -> SubExp -> CallKernelGen () -sReplicateKernel arr se = do - t <- subExpType se - ds <- dropLast (arrayRank t) . arrayDims <$> lookupType arr - - let dims = map pe64 $ ds ++ arrayDims t - n <- dPrimVE "replicate_n" $ product $ map sExt64 dims - (virtualise, constants) <- simpleKernelConstants n "replicate" - - fname <- askFunction - let name = - keyWithEntryPoint fname $ - nameFromString $ - "replicate_" ++ show (baseTag $ kernelGlobalThreadIdVar constants) - - sKernelFailureTolerant True threadOperations constants name $ - virtualise $ \gtid -> do - is' <- dIndexSpace' "rep_i" dims gtid - sWhen (gtid .<. n) $ - copyDWIMFix arr is' se $ - drop (length ds) is' - -replicateName :: PrimType -> String -replicateName bt = "replicate_" ++ pretty bt - -replicateForType :: PrimType -> CallKernelGen Name -replicateForType bt = do - let fname = nameFromString $ "builtin#" <> replicateName bt - - exists <- hasFunction fname - unless exists $ do - mem <- newVName "mem" - num_elems <- newVName "num_elems" - val <- newVName "val" - - let params = - [ Imp.MemParam mem (Space "device"), - Imp.ScalarParam num_elems int64, - Imp.ScalarParam val bt - ] - shape = Shape [Var num_elems] - function fname [] params $ do - arr <- - sArray "arr" bt shape mem $ IxFun.iota $ map pe64 $ shapeDims shape - sReplicateKernel arr $ Var val - - pure fname - -replicateIsFill :: VName -> SubExp -> CallKernelGen (Maybe (CallKernelGen ())) -replicateIsFill arr v = do - ArrayEntry (MemLoc arr_mem arr_shape arr_ixfun) _ <- lookupArray arr - v_t <- subExpType v - case v_t of - Prim v_t' - | IxFun.isLinear arr_ixfun -> pure $ - Just $ do - fname <- replicateForType v_t' - emit $ - Imp.Call - [] - fname - [ Imp.MemArg arr_mem, - Imp.ExpArg $ untyped $ product $ map pe64 arr_shape, - Imp.ExpArg $ toExp' v_t' v - ] - _ -> pure Nothing - --- | Perform a Replicate with a kernel. -sReplicate :: VName -> SubExp -> CallKernelGen () -sReplicate arr se = do - -- If the replicate is of a particularly common and simple form - -- (morally a memset()/fill), then we use a common function. - is_fill <- replicateIsFill arr se - - case is_fill of - Just m -> m - Nothing -> sReplicateKernel arr se - --- | Perform an Iota with a kernel. -sIotaKernel :: - VName -> - Imp.TExp Int64 -> - Imp.Exp -> - Imp.Exp -> - IntType -> - CallKernelGen () -sIotaKernel arr n x s et = do - destloc <- entryArrayLoc <$> lookupArray arr - (virtualise, constants) <- simpleKernelConstants n "iota" - - fname <- askFunction - let name = - keyWithEntryPoint fname $ - nameFromString $ - "iota_" - ++ pretty et - ++ "_" - ++ show (baseTag $ kernelGlobalThreadIdVar constants) - - sKernelFailureTolerant True threadOperations constants name $ - virtualise $ \gtid -> - sWhen (gtid .<. n) $ do - (destmem, destspace, destidx) <- fullyIndexArray' destloc [gtid] - - emit $ - Imp.Write destmem destidx (IntType et) destspace Imp.Nonvolatile $ - BinOpExp - (Add et OverflowWrap) - (BinOpExp (Mul et OverflowWrap) (Imp.sExt et $ untyped gtid) s) - x - -iotaName :: IntType -> String -iotaName bt = "iota_" ++ pretty bt - -iotaForType :: IntType -> CallKernelGen Name -iotaForType bt = do - let fname = nameFromString $ "builtin#" <> iotaName bt - - exists <- hasFunction fname - unless exists $ do - mem <- newVName "mem" - n <- newVName "n" - x <- newVName "x" - s <- newVName "s" - - let params = - [ Imp.MemParam mem (Space "device"), - Imp.ScalarParam n int32, - Imp.ScalarParam x $ IntType bt, - Imp.ScalarParam s $ IntType bt - ] - shape = Shape [Var n] - n' = Imp.le64 n - x' = Imp.var x $ IntType bt - s' = Imp.var s $ IntType bt - - function fname [] params $ do - arr <- - sArray "arr" (IntType bt) shape mem $ - IxFun.iota $ - map pe64 $ - shapeDims shape - sIotaKernel arr (sExt64 n') x' s' bt - - pure fname - --- | Perform an Iota with a kernel. -sIota :: - VName -> - Imp.TExp Int64 -> - Imp.Exp -> - Imp.Exp -> - IntType -> - CallKernelGen () -sIota arr n x s et = do - ArrayEntry (MemLoc arr_mem _ arr_ixfun) _ <- lookupArray arr - if IxFun.isLinear arr_ixfun - then do - fname <- iotaForType et - emit $ - Imp.Call - [] - fname - [Imp.MemArg arr_mem, Imp.ExpArg $ untyped n, Imp.ExpArg x, Imp.ExpArg s] - else sIotaKernel arr n x s et - -sCopy :: CopyCompiler GPUMem HostEnv Imp.HostOp -sCopy pt destloc@(MemLoc destmem _ _) srcloc@(MemLoc srcmem srcdims _) = do - -- Note that the shape of the destination and the source are - -- necessarily the same. - let shape = map pe64 srcdims - kernel_size = product shape - - (virtualise, constants) <- simpleKernelConstants kernel_size "copy" - - fname <- askFunction - let name = - keyWithEntryPoint fname $ - nameFromString $ - "copy_" ++ show (baseTag $ kernelGlobalThreadIdVar constants) - - sKernelFailureTolerant True threadOperations constants name $ - virtualise $ \gtid -> do - is <- dIndexSpace' "copy_i" shape gtid - - (_, destspace, destidx) <- fullyIndexArray' destloc is - (_, srcspace, srcidx) <- fullyIndexArray' srcloc is - - sWhen (gtid .<. kernel_size) $ do - tmp <- tvVar <$> dPrim "tmp" pt - emit $ Imp.Read tmp srcmem srcidx pt srcspace Imp.Nonvolatile - emit $ Imp.Write destmem destidx pt destspace Imp.Nonvolatile $ Imp.var tmp pt - --- | Perform a Rotate with a kernel. -sRotateKernel :: VName -> [Imp.TExp Int64] -> VName -> CallKernelGen () -sRotateKernel dest rs src = do - t <- lookupType src - let ds = map pe64 $ arrayDims t - n <- dPrimVE "rotate_n" $ product ds - (virtualise, constants) <- simpleKernelConstants n "rotate" - - fname <- askFunction - let name = - keyWithEntryPoint fname $ - nameFromString $ - "rotate_" ++ show (baseTag $ kernelGlobalThreadIdVar constants) - - sKernelFailureTolerant True threadOperations constants name $ - virtualise $ \gtid -> sWhen (gtid .<. n) $ do - is' <- dIndexSpace' "rep_i" ds gtid - is'' <- sequence $ zipWith3 rotate ds rs is' - copyDWIMFix dest is' (Var src) is'' - where - rotate d r i = dPrimVE "rot_i" $ rotateIndex d r i - -compileGroupResult :: - SegSpace -> - PatElem LetDecMem -> - KernelResult -> - InKernelGen () -compileGroupResult _ pe (TileReturns _ [(w, per_group_elems)] what) = do - n <- pe64 . arraySize 0 <$> lookupType what - - constants <- kernelConstants <$> askEnv - let ltid = sExt64 $ kernelLocalThreadId constants - offset = - pe64 per_group_elems - * sExt64 (kernelGroupId constants) - - -- Avoid loop for the common case where each thread is statically - -- known to write at most one element. - localOps threadOperations $ - if pe64 per_group_elems == kernelGroupSize constants - then - sWhen (ltid + offset .<. pe64 w) $ - copyDWIMFix (patElemName pe) [ltid + offset] (Var what) [ltid] - else sFor "i" (n `divUp` kernelGroupSize constants) $ \i -> do - j <- dPrimVE "j" $ kernelGroupSize constants * i + ltid - sWhen (j + offset .<. pe64 w) $ - copyDWIMFix (patElemName pe) [j + offset] (Var what) [j] -compileGroupResult space pe (TileReturns _ dims what) = do - let gids = map fst $ unSegSpace space - out_tile_sizes = map (pe64 . snd) dims - group_is = zipWith (*) (map Imp.le64 gids) out_tile_sizes - local_is <- localThreadIDs $ map snd dims - is_for_thread <- - mapM (dPrimV "thread_out_index") $ - zipWith (+) group_is local_is - - localOps threadOperations $ - sWhen (isActive $ zip (map tvVar is_for_thread) $ map fst dims) $ - copyDWIMFix (patElemName pe) (map tvExp is_for_thread) (Var what) local_is -compileGroupResult space pe (RegTileReturns _ dims_n_tiles what) = do - constants <- kernelConstants <$> askEnv - - let gids = map fst $ unSegSpace space - (dims, group_tiles, reg_tiles) = unzip3 dims_n_tiles - group_tiles' = map pe64 group_tiles - reg_tiles' = map pe64 reg_tiles - - -- Which group tile is this group responsible for? - let group_tile_is = map Imp.le64 gids - - -- Within the group tile, which register tile is this thread - -- responsible for? - reg_tile_is <- - dIndexSpace' "reg_tile_i" group_tiles' $ sExt64 $ kernelLocalThreadId constants - - -- Compute output array slice for the register tile belonging to - -- this thread. - let regTileSliceDim (group_tile, group_tile_i) (reg_tile, reg_tile_i) = do - tile_dim_start <- - dPrimVE "tile_dim_start" $ - reg_tile * (group_tile * group_tile_i + reg_tile_i) - pure $ DimSlice tile_dim_start reg_tile 1 - reg_tile_slices <- - Slice - <$> zipWithM - regTileSliceDim - (zip group_tiles' group_tile_is) - (zip reg_tiles' reg_tile_is) - - localOps threadOperations $ - sLoopNest (Shape reg_tiles) $ \is_in_reg_tile -> do - let dest_is = fixSlice reg_tile_slices is_in_reg_tile - src_is = reg_tile_is ++ is_in_reg_tile - sWhen (foldl1 (.&&.) $ zipWith (.<.) dest_is $ map pe64 dims) $ - copyDWIMFix (patElemName pe) dest_is (Var what) src_is -compileGroupResult space pe (Returns _ _ what) = do - constants <- kernelConstants <$> askEnv - in_local_memory <- arrayInLocalMemory what - let gids = map (Imp.le64 . fst) $ unSegSpace space - - if not in_local_memory - then - localOps threadOperations $ - sWhen (kernelLocalThreadId constants .==. 0) $ - copyDWIMFix (patElemName pe) gids what [] - else -- If the result of the group is an array in local memory, we - -- store it by collective copying among all the threads of the - -- group. TODO: also do this if the array is in global memory - -- (but this is a bit more tricky, synchronisation-wise). - copyDWIMFix (patElemName pe) gids what [] -compileGroupResult _ _ WriteReturns {} = - compilerLimitationS "compileGroupResult: WriteReturns not handled yet." -compileGroupResult _ _ ConcatReturns {} = - compilerLimitationS "compileGroupResult: ConcatReturns not handled yet." - -compileThreadResult :: - SegSpace -> - PatElem LetDecMem -> - KernelResult -> - InKernelGen () -compileThreadResult _ _ RegTileReturns {} = - compilerLimitationS "compileThreadResult: RegTileReturns not yet handled." -compileThreadResult space pe (Returns _ _ what) = do - let is = map (Imp.le64 . fst) $ unSegSpace space - copyDWIMFix (patElemName pe) is what [] -compileThreadResult _ pe (ConcatReturns _ SplitContiguous _ per_thread_elems what) = do - constants <- kernelConstants <$> askEnv - let offset = - pe64 per_thread_elems - * sExt64 (kernelGlobalThreadId constants) - n <- pe64 . arraySize 0 <$> lookupType what - copyDWIM (patElemName pe) [DimSlice offset n 1] (Var what) [] -compileThreadResult _ pe (ConcatReturns _ (SplitStrided stride) _ _ what) = do - offset <- sExt64 . kernelGlobalThreadId . kernelConstants <$> askEnv - n <- pe64 . arraySize 0 <$> lookupType what - copyDWIM (patElemName pe) [DimSlice offset n $ pe64 stride] (Var what) [] -compileThreadResult _ pe (WriteReturns _ (Shape rws) _arr dests) = do - let rws' = map pe64 rws - forM_ dests $ \(slice, e) -> do - let slice' = fmap pe64 slice - write = inBounds slice' rws' - sWhen write $ copyDWIM (patElemName pe) (unSlice slice') e [] -compileThreadResult _ _ TileReturns {} = - compilerBugS "compileThreadResult: TileReturns unhandled." - -arrayInLocalMemory :: SubExp -> InKernelGen Bool -arrayInLocalMemory (Var name) = do - res <- lookupVar name - case res of - ArrayVar _ entry -> - (Space "local" ==) . entryMemSpace - <$> lookupMemory (memLocName (entryArrayLoc entry)) - _ -> pure False -arrayInLocalMemory Constant {} = pure False + threadOperations, + keyWithEntryPoint, + CallKernelGen, + InKernelGen, + Locks (..), + HostEnv (..), + Target (..), + KernelEnv (..), + groupReduce, + groupScan, + groupLoop, + isActive, + sKernel, + sKernelThread, + KernelAttrs (..), + defKernelAttrs, + allocLocal, + kernelAlloc, + compileThreadResult, + virtualiseGroups, + kernelLoop, + groupCoverSpace, + fenceForArrays, + updateAcc, + + -- * Host-level bulk operations + sReplicate, + sIota, + sRotateKernel, + sCopy, + + -- * Atomics + AtomicBinOp, + atomicUpdateLocking, + Locking (..), + AtomicUpdate (..), + DoAtomicUpdate, + ) +where + +import Control.Monad.Except +import Data.List (foldl') +import qualified Data.Map.Strict as M +import Data.Maybe +import qualified Futhark.CodeGen.ImpCode.GPU as Imp +import Futhark.CodeGen.ImpGen +import Futhark.Error +import Futhark.IR.GPUMem +import qualified Futhark.IR.Mem.IxFun as IxFun +import Futhark.MonadFreshNames +import Futhark.Transform.Rename +import Futhark.Util (dropLast, nubOrd, splitFromEnd) +import Futhark.Util.IntegralExp (divUp, quot, rem) +import Prelude hiding (quot, rem) + +-- | Which target are we ultimately generating code for? While most +-- of the kernels code is the same, there are some cases where we +-- generate special code based on the ultimate low-level API we are +-- targeting. +data Target = CUDA | OpenCL + +-- | Information about the locks available for accumulators. +data Locks = Locks + { locksArray :: VName, + locksCount :: Int + } + +data HostEnv = HostEnv + { hostAtomics :: AtomicBinOp, + hostTarget :: Target, + hostLocks :: M.Map VName Locks + } + +data KernelEnv = KernelEnv + { kernelAtomics :: AtomicBinOp, + kernelConstants :: KernelConstants, + kernelLocks :: M.Map VName Locks + } + +type CallKernelGen = ImpM GPUMem HostEnv Imp.HostOp + +type InKernelGen = ImpM GPUMem KernelEnv Imp.KernelOp + +data KernelConstants = KernelConstants + { kernelGlobalThreadId :: Imp.TExp Int32, + kernelLocalThreadId :: Imp.TExp Int32, + kernelGroupId :: Imp.TExp Int32, + kernelGlobalThreadIdVar :: VName, + kernelLocalThreadIdVar :: VName, + kernelGroupIdVar :: VName, + kernelNumGroupsCount :: Count NumGroups SubExp, + kernelGroupSizeCount :: Count GroupSize SubExp, + kernelNumGroups :: Imp.TExp Int64, + kernelGroupSize :: Imp.TExp Int64, + kernelNumThreads :: Imp.TExp Int32, + kernelWaveSize :: Imp.TExp Int32, + -- | A mapping from dimensions of nested SegOps to already + -- computed local thread IDs. Only valid in non-virtualised case. + kernelLocalIdMap :: M.Map [SubExp] [Imp.TExp Int32], + -- | Mapping from dimensions of nested SegOps to how many + -- iterations the virtualisation loop needs. + kernelChunkItersMap :: M.Map [SubExp] (Imp.TExp Int32) + } + +keyWithEntryPoint :: Maybe Name -> Name -> Name +keyWithEntryPoint fname key = + nameFromString $ maybe "" ((++ ".") . nameToString) fname ++ nameToString key + +allocLocal :: AllocCompiler GPUMem r Imp.KernelOp +allocLocal mem size = + sOp $ Imp.LocalAlloc mem size + +kernelAlloc :: + Pat LetDecMem -> + SubExp -> + Space -> + InKernelGen () +kernelAlloc (Pat [_]) _ ScalarSpace {} = + -- Handled by the declaration of the memory block, which is then + -- translated to an actual scalar variable during C code generation. + pure () +kernelAlloc (Pat [mem]) size (Space "local") = + allocLocal (patElemName mem) $ Imp.bytes $ pe64 size +kernelAlloc (Pat [mem]) _ _ = + compilerLimitationS $ "Cannot allocate memory block " ++ pretty mem ++ " in kernel." +kernelAlloc dest _ _ = + error $ "Invalid target for in-kernel allocation: " ++ show dest + +updateAcc :: VName -> [SubExp] -> [SubExp] -> InKernelGen () +updateAcc acc is vs = sComment "UpdateAcc" $ do + -- See the ImpGen implementation of UpdateAcc for general notes. + let is' = map pe64 is + (c, space, arrs, dims, op) <- lookupAcc acc is' + sWhen (inBounds (Slice (map DimFix is')) dims) $ + case op of + Nothing -> + forM_ (zip arrs vs) $ \(arr, v) -> copyDWIMFix arr is' v [] + Just lam -> do + dLParams $ lambdaParams lam + let (_x_params, y_params) = + splitAt (length vs) $ map paramName $ lambdaParams lam + forM_ (zip y_params vs) $ \(yp, v) -> copyDWIM yp [] v [] + atomics <- kernelAtomics <$> askEnv + case atomicUpdateLocking atomics lam of + AtomicPrim f -> f space arrs is' + AtomicCAS f -> f space arrs is' + AtomicLocking f -> do + c_locks <- M.lookup c . kernelLocks <$> askEnv + case c_locks of + Just (Locks locks num_locks) -> do + let locking = + Locking locks 0 1 0 $ + pure . (`rem` fromIntegral num_locks) . flattenIndex dims + f locking space arrs is' + Nothing -> + error $ "Missing locks for " ++ pretty acc + +compileThreadExp :: ExpCompiler GPUMem KernelEnv Imp.KernelOp +compileThreadExp (Pat [pe]) (BasicOp (Opaque _ se)) = + -- Cannot print in GPU code. + copyDWIM (patElemName pe) [] se [] +compileThreadExp (Pat [dest]) (BasicOp (ArrayLit es _)) = + forM_ (zip [0 ..] es) $ \(i, e) -> + copyDWIMFix (patElemName dest) [fromIntegral (i :: Int64)] e [] +compileThreadExp _ (BasicOp (UpdateAcc acc is vs)) = + updateAcc acc is vs +compileThreadExp dest e = + defCompileExp dest e + +-- | Assign iterations of a for-loop to all threads in the kernel. +-- The passed-in function is invoked with the (symbolic) iteration. +-- The body must contain thread-level code. For multidimensional +-- loops, use 'groupCoverSpace'. +kernelLoop :: + IntExp t => + Imp.TExp t -> + Imp.TExp t -> + Imp.TExp t -> + (Imp.TExp t -> InKernelGen ()) -> + InKernelGen () +kernelLoop tid num_threads n f = + localOps threadOperations $ + if n == num_threads + then f tid + else do + num_chunks <- dPrimVE "num_chunks" $ n `divUp` num_threads + sFor "chunk_i" num_chunks $ \chunk_i -> do + i <- dPrimVE "i" $ chunk_i * num_threads + tid + sWhen (i .<. n) $ f i + +-- | Assign iterations of a for-loop to threads in the workgroup. The +-- passed-in function is invoked with the (symbolic) iteration. For +-- multidimensional loops, use 'groupCoverSpace'. +groupLoop :: + IntExp t => + Imp.TExp t -> + (Imp.TExp t -> InKernelGen ()) -> + InKernelGen () +groupLoop n f = do + constants <- kernelConstants <$> askEnv + kernelLoop + (kernelLocalThreadId constants `sExtAs` n) + (kernelGroupSize constants `sExtAs` n) + n + f + +-- | Iterate collectively though a multidimensional space, such that +-- all threads in the group participate. The passed-in function is +-- invoked with a (symbolic) point in the index space. +groupCoverSpace :: + IntExp t => + [Imp.TExp t] -> + ([Imp.TExp t] -> InKernelGen ()) -> + InKernelGen () +groupCoverSpace ds f = do + constants <- kernelConstants <$> askEnv + let group_size = kernelGroupSize constants + case splitFromEnd 1 ds of + -- Optimise the case where the inner dimension of the space is + -- equal to the group size. + (ds', [last_d]) + | last_d == (group_size `sExtAs` last_d) -> do + let ltid = kernelLocalThreadId constants `sExtAs` last_d + sLoopSpace ds' $ \ds_is -> + f $ ds_is ++ [ltid] + _ -> + groupLoop (product ds) $ f . unflattenIndex ds + +-- Which fence do we need to protect shared access to this memory space? +fenceForSpace :: Space -> Imp.Fence +fenceForSpace (Space "local") = Imp.FenceLocal +fenceForSpace _ = Imp.FenceGlobal + +-- | If we are touching these arrays, which kind of fence do we need? +fenceForArrays :: [VName] -> InKernelGen Imp.Fence +fenceForArrays = fmap (foldl' max Imp.FenceLocal) . mapM need + where + need arr = + fmap (fenceForSpace . entryMemSpace) + . lookupMemory + . memLocName + . entryArrayLoc + =<< lookupArray arr + +inBlockScan :: + KernelConstants -> + Maybe (Imp.TExp Int32 -> Imp.TExp Int32 -> Imp.TExp Bool) -> + Imp.TExp Int64 -> + Imp.TExp Int32 -> + Imp.TExp Int32 -> + Imp.TExp Bool -> + [VName] -> + InKernelGen () -> + Lambda GPUMem -> + InKernelGen () +inBlockScan constants seg_flag arrs_full_size lockstep_width block_size active arrs barrier scan_lam = everythingVolatile $ do + skip_threads <- dPrim "skip_threads" int32 + let actual_params = lambdaParams scan_lam + (x_params, y_params) = + splitAt (length actual_params `div` 2) actual_params + y_to_x = + forM_ (zip x_params y_params) $ \(x, y) -> + when (primType (paramType x)) $ + copyDWIM (paramName x) [] (Var (paramName y)) [] + + -- Set initial y values + sComment "read input for in-block scan" $ + sWhen active $ do + zipWithM_ readInitial y_params arrs + -- Since the final result is expected to be in x_params, we may + -- need to copy it there for the first thread in the block. + sWhen (in_block_id .==. 0) y_to_x + + when array_scan barrier + + let op_to_x in_block_thread_active + | Nothing <- seg_flag = + sWhen in_block_thread_active $ + compileBody' x_params $ + lambdaBody scan_lam + | Just flag_true <- seg_flag = do + inactive <- + dPrimVE "inactive" $ flag_true (ltid32 - tvExp skip_threads) ltid32 + sWhen (in_block_thread_active .&&. inactive) $ + forM_ (zip x_params y_params) $ \(x, y) -> + copyDWIM (paramName x) [] (Var (paramName y)) [] + -- The convoluted control flow is to ensure all threads + -- hit this barrier (if applicable). + when array_scan barrier + sWhen in_block_thread_active . sUnless inactive $ + compileBody' x_params $ + lambdaBody scan_lam + + maybeBarrier = + sWhen + (lockstep_width .<=. tvExp skip_threads) + barrier + + sComment "in-block scan (hopefully no barriers needed)" $ do + skip_threads <-- 1 + sWhile (tvExp skip_threads .<. block_size) $ do + thread_active <- + dPrimVE "thread_active" $ tvExp skip_threads .<=. in_block_id .&&. active + + sWhen thread_active . sComment "read operands" $ + zipWithM_ (readParam (sExt64 $ tvExp skip_threads)) x_params arrs + sComment "perform operation" $ op_to_x thread_active + + maybeBarrier + + sWhen thread_active . sComment "write result" $ + sequence_ $ + zipWith3 writeResult x_params y_params arrs + + maybeBarrier + + skip_threads <-- tvExp skip_threads * 2 + where + block_id = ltid32 `quot` block_size + in_block_id = ltid32 - block_id * block_size + ltid32 = kernelLocalThreadId constants + ltid = sExt64 ltid32 + gtid = sExt64 $ kernelGlobalThreadId constants + array_scan = not $ all primType $ lambdaReturnType scan_lam + + readInitial p arr + | primType $ paramType p = + copyDWIM (paramName p) [] (Var arr) [DimFix ltid] + | otherwise = + copyDWIM (paramName p) [] (Var arr) [DimFix gtid] + + readParam behind p arr + | primType $ paramType p = + copyDWIM (paramName p) [] (Var arr) [DimFix $ ltid - behind] + | otherwise = + copyDWIM (paramName p) [] (Var arr) [DimFix $ gtid - behind + arrs_full_size] + + writeResult x y arr + | primType $ paramType x = do + copyDWIM arr [DimFix ltid] (Var $ paramName x) [] + copyDWIM (paramName y) [] (Var $ paramName x) [] + | otherwise = + copyDWIM (paramName y) [] (Var $ paramName x) [] + +groupScan :: + Maybe (Imp.TExp Int32 -> Imp.TExp Int32 -> Imp.TExp Bool) -> + Imp.TExp Int64 -> + Imp.TExp Int64 -> + Lambda GPUMem -> + [VName] -> + InKernelGen () +groupScan seg_flag arrs_full_size w lam arrs = do + constants <- kernelConstants <$> askEnv + renamed_lam <- renameLambda lam + + let ltid32 = kernelLocalThreadId constants + ltid = sExt64 ltid32 + (x_params, y_params) = splitAt (length arrs) $ lambdaParams lam + + dLParams (lambdaParams lam ++ lambdaParams renamed_lam) + + ltid_in_bounds <- dPrimVE "ltid_in_bounds" $ ltid .<. w + + fence <- fenceForArrays arrs + + -- The scan works by splitting the group into blocks, which are + -- scanned separately. Typically, these blocks are smaller than + -- the lockstep width, which enables barrier-free execution inside + -- them. + -- + -- We hardcode the block size here. The only requirement is that + -- it should not be less than the square root of the group size. + -- With 32, we will work on groups of size 1024 or smaller, which + -- fits every device Troels has seen. Still, it would be nicer if + -- it were a runtime parameter. Some day. + let block_size = 32 + simd_width = kernelWaveSize constants + block_id = ltid32 `quot` block_size + in_block_id = ltid32 - block_id * block_size + doInBlockScan seg_flag' active = + inBlockScan + constants + seg_flag' + arrs_full_size + simd_width + block_size + active + arrs + barrier + array_scan = not $ all primType $ lambdaReturnType lam + barrier + | array_scan = + sOp $ Imp.Barrier Imp.FenceGlobal + | otherwise = + sOp $ Imp.Barrier fence + + group_offset = sExt64 (kernelGroupId constants) * kernelGroupSize constants + + writeBlockResult p arr + | primType $ paramType p = + copyDWIM arr [DimFix $ sExt64 block_id] (Var $ paramName p) [] + | otherwise = + copyDWIM arr [DimFix $ group_offset + sExt64 block_id] (Var $ paramName p) [] + + readPrevBlockResult p arr + | primType $ paramType p = + copyDWIM (paramName p) [] (Var arr) [DimFix $ sExt64 block_id - 1] + | otherwise = + copyDWIM (paramName p) [] (Var arr) [DimFix $ group_offset + sExt64 block_id - 1] + + doInBlockScan seg_flag ltid_in_bounds lam + barrier + + let is_first_block = block_id .==. 0 + when array_scan $ do + sComment "save correct values for first block" $ + sWhen is_first_block $ + forM_ (zip x_params arrs) $ \(x, arr) -> + unless (primType $ paramType x) $ + copyDWIM arr [DimFix $ arrs_full_size + group_offset + sExt64 block_size + ltid] (Var $ paramName x) [] + + barrier + + let last_in_block = in_block_id .==. block_size - 1 + sComment "last thread of block 'i' writes its result to offset 'i'" $ + sWhen (last_in_block .&&. ltid_in_bounds) $ + everythingVolatile $ + zipWithM_ writeBlockResult x_params arrs + + barrier + + let first_block_seg_flag = do + flag_true <- seg_flag + Just $ \from to -> + flag_true (from * block_size + block_size - 1) (to * block_size + block_size - 1) + comment + "scan the first block, after which offset 'i' contains carry-in for block 'i+1'" + $ doInBlockScan first_block_seg_flag (is_first_block .&&. ltid_in_bounds) renamed_lam + + barrier + + when array_scan $ do + sComment "move correct values for first block back a block" $ + sWhen is_first_block $ + forM_ (zip x_params arrs) $ \(x, arr) -> + unless (primType $ paramType x) $ + copyDWIM + arr + [DimFix $ arrs_full_size + group_offset + ltid] + (Var arr) + [DimFix $ arrs_full_size + group_offset + sExt64 block_size + ltid] + + barrier + + no_carry_in <- dPrimVE "no_carry_in" $ is_first_block .||. bNot ltid_in_bounds + + let read_carry_in = sUnless no_carry_in $ do + forM_ (zip x_params y_params) $ \(x, y) -> + copyDWIM (paramName y) [] (Var (paramName x)) [] + zipWithM_ readPrevBlockResult x_params arrs + + op_to_x + | Nothing <- seg_flag = + sUnless no_carry_in $ compileBody' x_params $ lambdaBody lam + | Just flag_true <- seg_flag = do + inactive <- + dPrimVE "inactive" $ flag_true (block_id * block_size - 1) ltid32 + sUnless no_carry_in . sWhen inactive . forM_ (zip x_params y_params) $ \(x, y) -> + copyDWIM (paramName x) [] (Var (paramName y)) [] + -- The convoluted control flow is to ensure all threads + -- hit this barrier (if applicable). + when array_scan barrier + sUnless no_carry_in $ sUnless inactive $ compileBody' x_params $ lambdaBody lam + + write_final_result = + forM_ (zip x_params arrs) $ \(p, arr) -> + when (primType $ paramType p) $ + copyDWIM arr [DimFix ltid] (Var $ paramName p) [] + + sComment "carry-in for every block except the first" $ do + sComment "read operands" read_carry_in + sComment "perform operation" op_to_x + sComment "write final result" $ sUnless no_carry_in write_final_result + + barrier + + sComment "restore correct values for first block" $ + sWhen (is_first_block .&&. ltid_in_bounds) $ + forM_ (zip3 x_params y_params arrs) $ \(x, y, arr) -> + if primType (paramType y) + then copyDWIM arr [DimFix ltid] (Var $ paramName y) [] + else copyDWIM (paramName x) [] (Var arr) [DimFix $ arrs_full_size + group_offset + ltid] + + barrier + +groupReduce :: + Imp.TExp Int32 -> + Lambda GPUMem -> + [VName] -> + InKernelGen () +groupReduce w lam arrs = do + offset <- dPrim "offset" int32 + groupReduceWithOffset offset w lam arrs + +groupReduceWithOffset :: + TV Int32 -> + Imp.TExp Int32 -> + Lambda GPUMem -> + [VName] -> + InKernelGen () +groupReduceWithOffset offset w lam arrs = do + constants <- kernelConstants <$> askEnv + + let local_tid = kernelLocalThreadId constants + global_tid = kernelGlobalThreadId constants + + barrier + | all primType $ lambdaReturnType lam = sOp $ Imp.Barrier Imp.FenceLocal + | otherwise = sOp $ Imp.Barrier Imp.FenceGlobal + + readReduceArgument param arr + | Prim _ <- paramType param = do + let i = local_tid + tvExp offset + copyDWIMFix (paramName param) [] (Var arr) [sExt64 i] + | otherwise = do + let i = global_tid + tvExp offset + copyDWIMFix (paramName param) [] (Var arr) [sExt64 i] + + writeReduceOpResult param arr + | Prim _ <- paramType param = + copyDWIMFix arr [sExt64 local_tid] (Var $ paramName param) [] + | otherwise = + pure () + + let (reduce_acc_params, reduce_arr_params) = splitAt (length arrs) $ lambdaParams lam + + skip_waves <- dPrimV "skip_waves" (1 :: Imp.TExp Int32) + dLParams $ lambdaParams lam + + offset <-- (0 :: Imp.TExp Int32) + + comment "participating threads read initial accumulator" $ + sWhen (local_tid .<. w) $ + zipWithM_ readReduceArgument reduce_acc_params arrs + + let do_reduce = do + comment "read array element" $ + zipWithM_ readReduceArgument reduce_arr_params arrs + comment "apply reduction operation" $ + compileBody' reduce_acc_params $ + lambdaBody lam + comment "write result of operation" $ + zipWithM_ writeReduceOpResult reduce_acc_params arrs + in_wave_reduce = everythingVolatile do_reduce + + wave_size = kernelWaveSize constants + group_size = kernelGroupSize constants + wave_id = local_tid `quot` wave_size + in_wave_id = local_tid - wave_id * wave_size + num_waves = (sExt32 group_size + wave_size - 1) `quot` wave_size + arg_in_bounds = local_tid + tvExp offset .<. w + + doing_in_wave_reductions = + tvExp offset .<. wave_size + apply_in_in_wave_iteration = + (in_wave_id .&. (2 * tvExp offset - 1)) .==. 0 + in_wave_reductions = do + offset <-- (1 :: Imp.TExp Int32) + sWhile doing_in_wave_reductions $ do + sWhen + (arg_in_bounds .&&. apply_in_in_wave_iteration) + in_wave_reduce + offset <-- tvExp offset * 2 + + doing_cross_wave_reductions = + tvExp skip_waves .<. num_waves + is_first_thread_in_wave = + in_wave_id .==. 0 + wave_not_skipped = + (wave_id .&. (2 * tvExp skip_waves - 1)) .==. 0 + apply_in_cross_wave_iteration = + arg_in_bounds .&&. is_first_thread_in_wave .&&. wave_not_skipped + cross_wave_reductions = + sWhile doing_cross_wave_reductions $ do + barrier + offset <-- tvExp skip_waves * wave_size + sWhen + apply_in_cross_wave_iteration + do_reduce + skip_waves <-- tvExp skip_waves * 2 + + in_wave_reductions + cross_wave_reductions + +compileThreadOp :: OpCompiler GPUMem KernelEnv Imp.KernelOp +compileThreadOp pat (Alloc size space) = + kernelAlloc pat size space +compileThreadOp pat _ = + compilerBugS $ "compileThreadOp: cannot compile rhs of binding " ++ pretty pat + +-- | Locking strategy used for an atomic update. +data Locking = Locking + { -- | Array containing the lock. + lockingArray :: VName, + -- | Value for us to consider the lock free. + lockingIsUnlocked :: Imp.TExp Int32, + -- | What to write when we lock it. + lockingToLock :: Imp.TExp Int32, + -- | What to write when we unlock it. + lockingToUnlock :: Imp.TExp Int32, + -- | A transformation from the logical lock index to the + -- physical position in the array. This can also be used + -- to make the lock array smaller. + lockingMapping :: [Imp.TExp Int64] -> [Imp.TExp Int64] + } + +-- | A function for generating code for an atomic update. Assumes +-- that the bucket is in-bounds. +type DoAtomicUpdate rep r = + Space -> [VName] -> [Imp.TExp Int64] -> ImpM rep r Imp.KernelOp () + +-- | The mechanism that will be used for performing the atomic update. +-- Approximates how efficient it will be. Ordered from most to least +-- efficient. +data AtomicUpdate rep r + = -- | Supported directly by primitive. + AtomicPrim (DoAtomicUpdate rep r) + | -- | Can be done by efficient swaps. + AtomicCAS (DoAtomicUpdate rep r) + | -- | Requires explicit locking. + AtomicLocking (Locking -> DoAtomicUpdate rep r) + +-- | Is there an atomic t'BinOp' corresponding to this t'BinOp'? +type AtomicBinOp = + BinOp -> + Maybe (VName -> VName -> Count Imp.Elements (Imp.TExp Int64) -> Imp.Exp -> Imp.AtomicOp) + +-- | Do an atomic update corresponding to a binary operator lambda. +atomicUpdateLocking :: + AtomicBinOp -> + Lambda GPUMem -> + AtomicUpdate GPUMem KernelEnv +atomicUpdateLocking atomicBinOp lam + | Just ops_and_ts <- lamIsBinOp lam, + all (\(_, t, _, _) -> primBitSize t `elem` [32, 64]) ops_and_ts = + primOrCas ops_and_ts $ \space arrs bucket -> + -- If the operator is a vectorised binary operator on 32/64-bit + -- values, we can use a particularly efficient + -- implementation. If the operator has an atomic implementation + -- we use that, otherwise it is still a binary operator which + -- can be implemented by atomic compare-and-swap if 32/64 bits. + forM_ (zip arrs ops_and_ts) $ \(a, (op, t, x, y)) -> do + -- Common variables. + old <- dPrim "old" t + + (arr', _a_space, bucket_offset) <- fullyIndexArray a bucket + + case opHasAtomicSupport space (tvVar old) arr' bucket_offset op of + Just f -> sOp $ f $ Imp.var y t + Nothing -> + atomicUpdateCAS space t a (tvVar old) bucket x $ + x <~~ Imp.BinOpExp op (Imp.var x t) (Imp.var y t) + where + opHasAtomicSupport space old arr' bucket' bop = do + let atomic f = Imp.Atomic space . f old arr' bucket' + atomic <$> atomicBinOp bop + + primOrCas ops + | all isPrim ops = AtomicPrim + | otherwise = AtomicCAS + + isPrim (op, _, _, _) = isJust $ atomicBinOp op + +-- If the operator functions purely on single 32/64-bit values, we can +-- use an implementation based on CAS, no matter what the operator +-- does. +atomicUpdateLocking _ op + | [Prim t] <- lambdaReturnType op, + [xp, _] <- lambdaParams op, + primBitSize t `elem` [32, 64] = AtomicCAS $ \space [arr] bucket -> do + old <- dPrim "old" t + atomicUpdateCAS space t arr (tvVar old) bucket (paramName xp) $ + compileBody' [xp] $ + lambdaBody op +atomicUpdateLocking _ op = AtomicLocking $ \locking space arrs bucket -> do + old <- dPrim "old" int32 + continue <- dPrimVol "continue" Bool true + + -- Correctly index into locks. + (locks', _locks_space, locks_offset) <- + fullyIndexArray (lockingArray locking) $ lockingMapping locking bucket + + -- Critical section + let try_acquire_lock = + sOp $ + Imp.Atomic space $ + Imp.AtomicCmpXchg + int32 + (tvVar old) + locks' + locks_offset + (untyped $ lockingIsUnlocked locking) + (untyped $ lockingToLock locking) + lock_acquired = tvExp old .==. lockingIsUnlocked locking + -- Even the releasing is done with an atomic rather than a + -- simple write, for memory coherency reasons. + release_lock = + sOp $ + Imp.Atomic space $ + Imp.AtomicCmpXchg + int32 + (tvVar old) + locks' + locks_offset + (untyped $ lockingToLock locking) + (untyped $ lockingToUnlock locking) + break_loop = continue <-- false + + -- Preparing parameters. It is assumed that the caller has already + -- filled the arr_params. We copy the current value to the + -- accumulator parameters. + -- + -- Note the use of 'everythingVolatile' when reading and writing the + -- buckets. This was necessary to ensure correct execution on a + -- newer NVIDIA GPU (RTX 2080). The 'volatile' modifiers likely + -- make the writes pass through the (SM-local) L1 cache, which is + -- necessary here, because we are really doing device-wide + -- synchronisation without atomics (naughty!). + let (acc_params, _arr_params) = splitAt (length arrs) $ lambdaParams op + bind_acc_params = + everythingVolatile $ + sComment "bind lhs" $ + forM_ (zip acc_params arrs) $ \(acc_p, arr) -> + copyDWIMFix (paramName acc_p) [] (Var arr) bucket + + let op_body = + sComment "execute operation" $ + compileBody' acc_params $ + lambdaBody op + + do_hist = + everythingVolatile $ + sComment "update global result" $ + zipWithM_ (writeArray bucket) arrs $ + map (Var . paramName) acc_params + + fence = sOp $ Imp.MemFence $ fenceForSpace space + + -- While-loop: Try to insert your value + sWhile (tvExp continue) $ do + try_acquire_lock + sWhen lock_acquired $ do + dLParams acc_params + bind_acc_params + op_body + do_hist + fence + release_lock + break_loop + fence + where + writeArray bucket arr val = copyDWIMFix arr bucket val [] + +atomicUpdateCAS :: + Space -> + PrimType -> + VName -> + VName -> + [Imp.TExp Int64] -> + VName -> + InKernelGen () -> + InKernelGen () +atomicUpdateCAS space t arr old bucket x do_op = do + -- Code generation target: + -- + -- old = d_his[idx]; + -- do { + -- assumed = old; + -- x = do_op(assumed, y); + -- old = atomicCAS(&d_his[idx], assumed, tmp); + -- } while(assumed != old); + assumed <- tvVar <$> dPrim "assumed" t + run_loop <- dPrimV "run_loop" true + + -- XXX: CUDA may generate really bad code if this is not a volatile + -- read. Unclear why. The later reads are volatile, so maybe + -- that's it. + everythingVolatile $ copyDWIMFix old [] (Var arr) bucket + + (arr', _a_space, bucket_offset) <- fullyIndexArray arr bucket + + -- While-loop: Try to insert your value + let (toBits, fromBits) = + case t of + FloatType Float16 -> + ( \v -> Imp.FunExp "to_bits16" [v] int16, + \v -> Imp.FunExp "from_bits16" [v] t + ) + FloatType Float32 -> + ( \v -> Imp.FunExp "to_bits32" [v] int32, + \v -> Imp.FunExp "from_bits32" [v] t + ) + FloatType Float64 -> + ( \v -> Imp.FunExp "to_bits64" [v] int64, + \v -> Imp.FunExp "from_bits64" [v] t + ) + _ -> (id, id) + + int + | primBitSize t == 16 = int16 + | primBitSize t == 32 = int32 + | otherwise = int64 + + sWhile (tvExp run_loop) $ do + assumed <~~ Imp.var old t + x <~~ Imp.var assumed t + do_op + old_bits_v <- newVName "old_bits" + dPrim_ old_bits_v int + let old_bits = Imp.var old_bits_v int + sOp . Imp.Atomic space $ + Imp.AtomicCmpXchg + int + old_bits_v + arr' + bucket_offset + (toBits (Imp.var assumed t)) + (toBits (Imp.var x t)) + old <~~ fromBits old_bits + let won = CmpOpExp (CmpEq int) (toBits (Imp.var assumed t)) old_bits + sWhen (isBool won) (run_loop <-- false) + +computeKernelUses :: + FreeIn a => + a -> + [VName] -> + CallKernelGen [Imp.KernelUse] +computeKernelUses kernel_body bound_in_kernel = do + let actually_free = freeIn kernel_body `namesSubtract` namesFromList bound_in_kernel + -- Compute the variables that we need to pass to the kernel. + nubOrd <$> readsFromSet actually_free + +readsFromSet :: Names -> CallKernelGen [Imp.KernelUse] +readsFromSet = fmap catMaybes . mapM f . namesToList + where + f var = do + t <- lookupType var + vtable <- getVTable + case t of + Array {} -> pure Nothing + Acc {} -> pure Nothing + Mem (Space "local") -> pure Nothing + Mem {} -> pure $ Just $ Imp.MemoryUse var + Prim bt -> + isConstExp vtable (Imp.var var bt) >>= \case + Just ce -> pure $ Just $ Imp.ConstUse var ce + Nothing -> pure $ Just $ Imp.ScalarUse var bt + +isConstExp :: + VTable GPUMem -> + Imp.Exp -> + ImpM rep r op (Maybe Imp.KernelConstExp) +isConstExp vtable size = do + fname <- askFunction + let onLeaf name _ = lookupConstExp name + lookupConstExp name = + constExp =<< hasExp =<< M.lookup name vtable + constExp (Op (Inner (SizeOp (GetSize key _)))) = + Just $ LeafExp (Imp.SizeConst $ keyWithEntryPoint fname key) int32 + constExp e = primExpFromExp lookupConstExp e + pure $ replaceInPrimExpM onLeaf size + where + hasExp (ArrayVar e _) = e + hasExp (AccVar e _) = e + hasExp (ScalarVar e _) = e + hasExp (MemVar e _) = e + +kernelInitialisationSimple :: + Count NumGroups SubExp -> + Count GroupSize SubExp -> + CallKernelGen (KernelConstants, InKernelGen ()) +kernelInitialisationSimple num_groups group_size = do + global_tid <- newVName "global_tid" + local_tid <- newVName "local_tid" + group_id <- newVName "group_tid" + wave_size <- newVName "wave_size" + inner_group_size <- newVName "group_size" + let num_groups' = Imp.pe64 (unCount num_groups) + group_size' = Imp.pe64 (unCount group_size) + constants = + KernelConstants + { kernelGlobalThreadId = Imp.le32 global_tid, + kernelLocalThreadId = Imp.le32 local_tid, + kernelGroupId = Imp.le32 group_id, + kernelGlobalThreadIdVar = global_tid, + kernelLocalThreadIdVar = local_tid, + kernelNumGroupsCount = num_groups, + kernelGroupSizeCount = group_size, + kernelGroupIdVar = group_id, + kernelNumGroups = num_groups', + kernelGroupSize = group_size', + kernelNumThreads = sExt32 (group_size' * num_groups'), + kernelWaveSize = Imp.le32 wave_size, + kernelLocalIdMap = mempty, + kernelChunkItersMap = mempty + } + + let set_constants = do + dPrim_ local_tid int32 + dPrim_ inner_group_size int64 + dPrim_ wave_size int32 + dPrim_ group_id int32 + + sOp (Imp.GetLocalId local_tid 0) + sOp (Imp.GetLocalSize inner_group_size 0) + sOp (Imp.GetLockstepWidth wave_size) + sOp (Imp.GetGroupId group_id 0) + dPrimV_ global_tid $ le32 group_id * le32 inner_group_size + le32 local_tid + + pure (constants, set_constants) + +isActive :: [(VName, SubExp)] -> Imp.TExp Bool +isActive limit = case actives of + [] -> true + x : xs -> foldl (.&&.) x xs + where + (is, ws) = unzip limit + actives = zipWith active is $ map pe64 ws + active i = (Imp.le64 i .<.) + +-- | Change every memory block to be in the global address space, +-- except those who are in the local memory space. This only affects +-- generated code - we still need to make sure that the memory is +-- actually present on the device (and declared as variables in the +-- kernel). +makeAllMemoryGlobal :: CallKernelGen a -> CallKernelGen a +makeAllMemoryGlobal = + localDefaultSpace (Imp.Space "global") . localVTable (M.map globalMemory) + where + globalMemory (MemVar _ entry) + | entryMemSpace entry /= Space "local" = + MemVar Nothing entry {entryMemSpace = Imp.Space "global"} + globalMemory entry = + entry + +simpleKernelGroups :: + Imp.TExp Int64 -> + Imp.TExp Int64 -> + CallKernelGen (Imp.TExp Int32, Count NumGroups SubExp, Count GroupSize SubExp) +simpleKernelGroups max_num_groups kernel_size = do + group_size <- dPrim "group_size" int64 + fname <- askFunction + let group_size_key = keyWithEntryPoint fname $ nameFromString $ pretty $ tvVar group_size + sOp $ Imp.GetSize (tvVar group_size) group_size_key Imp.SizeGroup + virt_num_groups <- dPrimVE "virt_num_groups" $ kernel_size `divUp` tvExp group_size + num_groups <- dPrimV "num_groups" $ virt_num_groups `sMin64` max_num_groups + pure (sExt32 virt_num_groups, Count $ tvSize num_groups, Count $ tvSize group_size) + +simpleKernelConstants :: + Imp.TExp Int64 -> + String -> + CallKernelGen + ( (Imp.TExp Int64 -> InKernelGen ()) -> InKernelGen (), + KernelConstants + ) +simpleKernelConstants kernel_size desc = do + -- For performance reasons, codegen assumes that the thread count is + -- never more than will fit in an i32. This means we need to cap + -- the number of groups here. The cap is set much higher than any + -- GPU will possibly need. Feel free to come back and laugh at me + -- in the future. + let max_num_groups = 1024 * 1024 + thread_gtid <- newVName $ desc ++ "_gtid" + thread_ltid <- newVName $ desc ++ "_ltid" + group_id <- newVName $ desc ++ "_gid" + inner_group_size <- newVName "group_size" + (virt_num_groups, num_groups, group_size) <- + simpleKernelGroups max_num_groups kernel_size + let group_size' = Imp.pe64 $ unCount group_size + num_groups' = Imp.pe64 $ unCount num_groups + + constants = + KernelConstants + { kernelGlobalThreadId = Imp.le32 thread_gtid, + kernelLocalThreadId = Imp.le32 thread_ltid, + kernelGroupId = Imp.le32 group_id, + kernelGlobalThreadIdVar = thread_gtid, + kernelLocalThreadIdVar = thread_ltid, + kernelGroupIdVar = group_id, + kernelNumGroupsCount = num_groups, + kernelGroupSizeCount = group_size, + kernelNumGroups = num_groups', + kernelGroupSize = group_size', + kernelNumThreads = sExt32 (group_size' * num_groups'), + kernelWaveSize = 0, + kernelLocalIdMap = mempty, + kernelChunkItersMap = mempty + } + + wrapKernel m = do + dPrim_ thread_ltid int32 + dPrim_ inner_group_size int64 + dPrim_ group_id int32 + sOp (Imp.GetLocalId thread_ltid 0) + sOp (Imp.GetLocalSize inner_group_size 0) + sOp (Imp.GetGroupId group_id 0) + dPrimV_ thread_gtid $ le32 group_id * le32 inner_group_size + le32 thread_ltid + virtualiseGroups SegVirt virt_num_groups $ \virt_group_id -> do + global_tid <- + dPrimVE "global_tid" $ + sExt64 virt_group_id * sExt64 (le32 inner_group_size) + + sExt64 (kernelLocalThreadId constants) + m global_tid + + pure (wrapKernel, constants) + +-- | For many kernels, we may not have enough physical groups to cover +-- the logical iteration space. Some groups thus have to perform +-- double duty; we put an outer loop to accomplish this. The +-- advantage over just launching a bazillion threads is that the cost +-- of memory expansion should be proportional to the number of +-- *physical* threads (hardware parallelism), not the amount of +-- application parallelism. +virtualiseGroups :: + SegVirt -> + Imp.TExp Int32 -> + (Imp.TExp Int32 -> InKernelGen ()) -> + InKernelGen () +virtualiseGroups SegVirt required_groups m = do + constants <- kernelConstants <$> askEnv + phys_group_id <- dPrim "phys_group_id" int32 + sOp $ Imp.GetGroupId (tvVar phys_group_id) 0 + iterations <- + dPrimVE "iterations" $ + (required_groups - tvExp phys_group_id) `divUp` sExt32 (kernelNumGroups constants) + + sFor "i" iterations $ \i -> do + m . tvExp + =<< dPrimV + "virt_group_id" + (tvExp phys_group_id + i * sExt32 (kernelNumGroups constants)) + -- Make sure the virtual group is actually done before we let + -- another virtual group have its way with it. + sOp $ Imp.Barrier Imp.FenceGlobal +virtualiseGroups _ _ m = do + gid <- kernelGroupIdVar . kernelConstants <$> askEnv + m $ Imp.le32 gid + +-- | Various extra configuration of the kernel being generated. +data KernelAttrs = KernelAttrs + { -- | Can this kernel execute correctly even if previous kernels failed? + kAttrFailureTolerant :: Bool, + -- | Does whatever launch this kernel check for local memory capacity itself? + kAttrCheckLocalMemory :: Bool, + -- | Number of groups. + kAttrNumGroups :: Count NumGroups SubExp, + -- | Group size. + kAttrGroupSize :: Count GroupSize SubExp + } + +-- | The default kernel attributes. +defKernelAttrs :: + Count NumGroups SubExp -> + Count GroupSize SubExp -> + KernelAttrs +defKernelAttrs num_groups group_size = + KernelAttrs + { kAttrFailureTolerant = False, + kAttrCheckLocalMemory = True, + kAttrNumGroups = num_groups, + kAttrGroupSize = group_size + } + +sKernel :: + Operations GPUMem KernelEnv Imp.KernelOp -> + (KernelConstants -> Imp.TExp Int32) -> + String -> + VName -> + KernelAttrs -> + InKernelGen () -> + CallKernelGen () +sKernel ops flatf name v attrs f = do + (constants, set_constants) <- + kernelInitialisationSimple (kAttrNumGroups attrs) (kAttrGroupSize attrs) + name' <- nameForFun $ name ++ "_" ++ show (baseTag v) + sKernelOp attrs constants ops name' $ do + set_constants + dPrimV_ v $ flatf constants + f + +sKernelThread :: + String -> + VName -> + KernelAttrs -> + InKernelGen () -> + CallKernelGen () +sKernelThread = sKernel threadOperations kernelGlobalThreadId + +sKernelOp :: + KernelAttrs -> + KernelConstants -> + Operations GPUMem KernelEnv Imp.KernelOp -> + Name -> + InKernelGen () -> + CallKernelGen () +sKernelOp attrs constants ops name m = do + HostEnv atomics _ locks <- askEnv + body <- makeAllMemoryGlobal $ subImpM_ (KernelEnv atomics constants locks) ops m + uses <- computeKernelUses body mempty + emit . Imp.Op . Imp.CallKernel $ + Imp.Kernel + { Imp.kernelBody = body, + Imp.kernelUses = uses, + Imp.kernelNumGroups = [untyped $ kernelNumGroups constants], + Imp.kernelGroupSize = [untyped $ kernelGroupSize constants], + Imp.kernelName = name, + Imp.kernelFailureTolerant = kAttrFailureTolerant attrs, + Imp.kernelCheckLocalMemory = kAttrCheckLocalMemory attrs + } + +sKernelFailureTolerant :: + Bool -> + Operations GPUMem KernelEnv Imp.KernelOp -> + KernelConstants -> + Name -> + InKernelGen () -> + CallKernelGen () +sKernelFailureTolerant tol ops constants name m = do + sKernelOp attrs constants ops name m + where + attrs = + ( defKernelAttrs + (kernelNumGroupsCount constants) + (kernelGroupSizeCount constants) + ) + { kAttrFailureTolerant = tol + } + +threadOperations :: Operations GPUMem KernelEnv Imp.KernelOp +threadOperations = + (defaultOperations compileThreadOp) + { opsCopyCompiler = copyElementWise, + opsExpCompiler = compileThreadExp, + opsStmsCompiler = \_ -> defCompileStms mempty, + opsAllocCompilers = + M.fromList [(Space "local", allocLocal)] + } + +-- | Perform a Replicate with a kernel. +sReplicateKernel :: VName -> SubExp -> CallKernelGen () +sReplicateKernel arr se = do + t <- subExpType se + ds <- dropLast (arrayRank t) . arrayDims <$> lookupType arr + + let dims = map pe64 $ ds ++ arrayDims t + n <- dPrimVE "replicate_n" $ product $ map sExt64 dims + (virtualise, constants) <- simpleKernelConstants n "replicate" + + fname <- askFunction + let name = + keyWithEntryPoint fname $ + nameFromString $ + "replicate_" ++ show (baseTag $ kernelGlobalThreadIdVar constants) + + sKernelFailureTolerant True threadOperations constants name $ + virtualise $ \gtid -> do + is' <- dIndexSpace' "rep_i" dims gtid + sWhen (gtid .<. n) $ + copyDWIMFix arr is' se $ + drop (length ds) is' + +replicateName :: PrimType -> String +replicateName bt = "replicate_" ++ pretty bt + +replicateForType :: PrimType -> CallKernelGen Name +replicateForType bt = do + let fname = nameFromString $ "builtin#" <> replicateName bt + + exists <- hasFunction fname + unless exists $ do + mem <- newVName "mem" + num_elems <- newVName "num_elems" + val <- newVName "val" + + let params = + [ Imp.MemParam mem (Space "device"), + Imp.ScalarParam num_elems int64, + Imp.ScalarParam val bt + ] + shape = Shape [Var num_elems] + function fname [] params $ do + arr <- + sArray "arr" bt shape mem $ IxFun.iota $ map pe64 $ shapeDims shape + sReplicateKernel arr $ Var val + + pure fname + +replicateIsFill :: VName -> SubExp -> CallKernelGen (Maybe (CallKernelGen ())) +replicateIsFill arr v = do + ArrayEntry (MemLoc arr_mem arr_shape arr_ixfun) _ <- lookupArray arr + v_t <- subExpType v + case v_t of + Prim v_t' + | IxFun.isLinear arr_ixfun -> pure $ + Just $ do + fname <- replicateForType v_t' + emit $ + Imp.Call + [] + fname + [ Imp.MemArg arr_mem, + Imp.ExpArg $ untyped $ product $ map pe64 arr_shape, + Imp.ExpArg $ toExp' v_t' v + ] + _ -> pure Nothing + +-- | Perform a Replicate with a kernel. +sReplicate :: VName -> SubExp -> CallKernelGen () +sReplicate arr se = do + -- If the replicate is of a particularly common and simple form + -- (morally a memset()/fill), then we use a common function. + is_fill <- replicateIsFill arr se + + case is_fill of + Just m -> m + Nothing -> sReplicateKernel arr se + +-- | Perform an Iota with a kernel. +sIotaKernel :: + VName -> + Imp.TExp Int64 -> + Imp.Exp -> + Imp.Exp -> + IntType -> + CallKernelGen () +sIotaKernel arr n x s et = do + destloc <- entryArrayLoc <$> lookupArray arr + (virtualise, constants) <- simpleKernelConstants n "iota" + + fname <- askFunction + let name = + keyWithEntryPoint fname $ + nameFromString $ + "iota_" + ++ pretty et + ++ "_" + ++ show (baseTag $ kernelGlobalThreadIdVar constants) + + sKernelFailureTolerant True threadOperations constants name $ + virtualise $ \gtid -> + sWhen (gtid .<. n) $ do + (destmem, destspace, destidx) <- fullyIndexArray' destloc [gtid] + + emit $ + Imp.Write destmem destidx (IntType et) destspace Imp.Nonvolatile $ + BinOpExp + (Add et OverflowWrap) + (BinOpExp (Mul et OverflowWrap) (Imp.sExt et $ untyped gtid) s) + x + +iotaName :: IntType -> String +iotaName bt = "iota_" ++ pretty bt + +iotaForType :: IntType -> CallKernelGen Name +iotaForType bt = do + let fname = nameFromString $ "builtin#" <> iotaName bt + + exists <- hasFunction fname + unless exists $ do + mem <- newVName "mem" + n <- newVName "n" + x <- newVName "x" + s <- newVName "s" + + let params = + [ Imp.MemParam mem (Space "device"), + Imp.ScalarParam n int32, + Imp.ScalarParam x $ IntType bt, + Imp.ScalarParam s $ IntType bt + ] + shape = Shape [Var n] + n' = Imp.le64 n + x' = Imp.var x $ IntType bt + s' = Imp.var s $ IntType bt + + function fname [] params $ do + arr <- + sArray "arr" (IntType bt) shape mem $ + IxFun.iota $ + map pe64 $ + shapeDims shape + sIotaKernel arr (sExt64 n') x' s' bt + + pure fname + +-- | Perform an Iota with a kernel. +sIota :: + VName -> + Imp.TExp Int64 -> + Imp.Exp -> + Imp.Exp -> + IntType -> + CallKernelGen () +sIota arr n x s et = do + ArrayEntry (MemLoc arr_mem _ arr_ixfun) _ <- lookupArray arr + if IxFun.isLinear arr_ixfun + then do + fname <- iotaForType et + emit $ + Imp.Call + [] + fname + [Imp.MemArg arr_mem, Imp.ExpArg $ untyped n, Imp.ExpArg x, Imp.ExpArg s] + else sIotaKernel arr n x s et + +sCopy :: CopyCompiler GPUMem HostEnv Imp.HostOp +sCopy pt destloc@(MemLoc destmem _ _) srcloc@(MemLoc srcmem srcdims _) = do + -- Note that the shape of the destination and the source are + -- necessarily the same. + let shape = map pe64 srcdims + kernel_size = product shape + + (virtualise, constants) <- simpleKernelConstants kernel_size "copy" + + fname <- askFunction + let name = + keyWithEntryPoint fname $ + nameFromString $ + "copy_" ++ show (baseTag $ kernelGlobalThreadIdVar constants) + + sKernelFailureTolerant True threadOperations constants name $ + virtualise $ \gtid -> do + is <- dIndexSpace' "copy_i" shape gtid + + (_, destspace, destidx) <- fullyIndexArray' destloc is + (_, srcspace, srcidx) <- fullyIndexArray' srcloc is + + sWhen (gtid .<. kernel_size) $ do + tmp <- tvVar <$> dPrim "tmp" pt + emit $ Imp.Read tmp srcmem srcidx pt srcspace Imp.Nonvolatile + emit $ Imp.Write destmem destidx pt destspace Imp.Nonvolatile $ Imp.var tmp pt + +-- | Perform a Rotate with a kernel. +sRotateKernel :: VName -> [Imp.TExp Int64] -> VName -> CallKernelGen () +sRotateKernel dest rs src = do + t <- lookupType src + let ds = map pe64 $ arrayDims t + n <- dPrimVE "rotate_n" $ product ds + (virtualise, constants) <- simpleKernelConstants n "rotate" + + fname <- askFunction + let name = + keyWithEntryPoint fname $ + nameFromString $ + "rotate_" ++ show (baseTag $ kernelGlobalThreadIdVar constants) + + sKernelFailureTolerant True threadOperations constants name $ + virtualise $ \gtid -> sWhen (gtid .<. n) $ do + is' <- dIndexSpace' "rep_i" ds gtid + is'' <- sequence $ zipWith3 rotate ds rs is' + copyDWIMFix dest is' (Var src) is'' + where + rotate d r i = dPrimVE "rot_i" $ rotateIndex d r i + +compileThreadResult :: + SegSpace -> + PatElem LetDecMem -> + KernelResult -> + InKernelGen () +compileThreadResult _ _ RegTileReturns {} = + compilerLimitationS "compileThreadResult: RegTileReturns not yet handled." +compileThreadResult space pe (Returns _ _ what) = do + let is = map (Imp.le64 . fst) $ unSegSpace space + copyDWIMFix (patElemName pe) is what [] +compileThreadResult _ pe (WriteReturns _ (Shape rws) _arr dests) = do + let rws' = map pe64 rws + forM_ dests $ \(slice, e) -> do + let slice' = fmap pe64 slice + write = inBounds slice' rws' + sWhen write $ copyDWIM (patElemName pe) (unSlice slice') e [] +compileThreadResult _ _ TileReturns {} = + compilerBugS "compileThreadResult: TileReturns unhandled." diff --git a/src/Futhark/CodeGen/ImpGen/GPU/Group.hs b/src/Futhark/CodeGen/ImpGen/GPU/Group.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpGen/GPU/Group.hs @@ -0,0 +1,735 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} + +-- | Generation of kernels with group-level bodies. +module Futhark.CodeGen.ImpGen.GPU.Group + ( sKernelGroup, + compileGroupResult, + groupOperations, + + -- * Precomputation + Precomputed, + precomputeConstants, + precomputedConstants, + atomicUpdateLocking, + ) +where + +import Control.Monad.Except +import Data.Bifunctor +import Data.List (partition, zip4) +import qualified Data.Map.Strict as M +import Data.Maybe +import qualified Data.Set as S +import qualified Futhark.CodeGen.ImpCode.GPU as Imp +import Futhark.CodeGen.ImpGen +import Futhark.CodeGen.ImpGen.GPU.Base +import Futhark.Construct (fullSliceNum) +import Futhark.Error +import Futhark.IR.GPUMem +import qualified Futhark.IR.Mem.IxFun as IxFun +import Futhark.MonadFreshNames +import Futhark.Transform.Rename +import Futhark.Util (chunks, mapAccumLM, takeLast) +import Futhark.Util.IntegralExp (divUp, rem) +import Prelude hiding (quot, rem) + +-- | @flattenArray k flat arr@ flattens the outer @k@ dimensions of +-- @arr@ to @flat@. (Make sure @flat@ is the sum of those dimensions +-- or you'll have a bad time.) +flattenArray :: Int -> TV Int64 -> VName -> ImpM rep r op VName +flattenArray k flat arr = do + ArrayEntry arr_loc pt <- lookupArray arr + let flat_shape = Shape $ Var (tvVar flat) : drop k (memLocShape arr_loc) + sArray (baseString arr ++ "_flat") pt flat_shape (memLocName arr_loc) $ + IxFun.reshape (memLocIxFun arr_loc) $ + map pe64 $ + shapeDims flat_shape + +sliceArray :: Imp.TExp Int64 -> TV Int64 -> VName -> ImpM rep r op VName +sliceArray start size arr = do + MemLoc mem _ ixfun <- entryArrayLoc <$> lookupArray arr + arr_t <- lookupType arr + let slice = + fullSliceNum + (map Imp.pe64 (arrayDims arr_t)) + [DimSlice start (tvExp size) 1] + sArray + (baseString arr ++ "_chunk") + (elemType arr_t) + (arrayShape arr_t `setOuterDim` Var (tvVar size)) + mem + $ IxFun.slice ixfun slice + +-- | @applyLambda lam dests args@ emits code that: +-- +-- 1. Binds each parameter of @lam@ to the corresponding element of +-- @args@, interpreted as a (name,slice) pair (as in 'copyDWIM'). +-- Use an empty list for a scalar. +-- +-- 2. Executes the body of @lam@. +-- +-- 3. Binds the t'SubExp's that are the 'Result' of @lam@ to the +-- provided @dest@s, again interpreted as the destination for a +-- 'copyDWIM'. +applyLambda :: + Mem rep inner => + Lambda rep -> + [(VName, [DimIndex (Imp.TExp Int64)])] -> + [(SubExp, [DimIndex (Imp.TExp Int64)])] -> + ImpM rep r op () +applyLambda lam dests args = do + dLParams $ lambdaParams lam + forM_ (zip (lambdaParams lam) args) $ \(p, (arg, arg_slice)) -> + copyDWIM (paramName p) [] arg arg_slice + compileStms mempty (bodyStms $ lambdaBody lam) $ do + let res = map resSubExp $ bodyResult $ lambdaBody lam + forM_ (zip dests res) $ \((dest, dest_slice), se) -> + copyDWIM dest dest_slice se [] + +-- | As applyLambda, but first rename the names in the lambda. This +-- makes it safe to apply it in multiple places. (It might be safe +-- anyway, but you have to be more careful - use this if you are in +-- doubt.) +applyRenamedLambda :: + Mem rep inner => + Lambda rep -> + [(VName, [DimIndex (Imp.TExp Int64)])] -> + [(SubExp, [DimIndex (Imp.TExp Int64)])] -> + ImpM rep r op () +applyRenamedLambda lam dests args = do + lam_renamed <- renameLambda lam + applyLambda lam_renamed dests args + +groupChunkLoop :: + Imp.TExp Int32 -> + (Imp.TExp Int32 -> TV Int64 -> InKernelGen ()) -> + InKernelGen () +groupChunkLoop w m = do + constants <- kernelConstants <$> askEnv + let max_chunk_size = sExt32 $ kernelGroupSize constants + num_chunks <- dPrimVE "num_chunks" $ w `divUp` max_chunk_size + sFor "chunk_i" num_chunks $ \chunk_i -> do + chunk_start <- + dPrimVE "chunk_start" $ chunk_i * max_chunk_size + chunk_end <- + dPrimVE "chunk_end" $ sMin32 w (chunk_start + max_chunk_size) + chunk_size <- + dPrimV "chunk_size" $ sExt64 $ chunk_end - chunk_start + m chunk_start chunk_size + +virtualisedGroupScan :: + Maybe (Imp.TExp Int32 -> Imp.TExp Int32 -> Imp.TExp Bool) -> + Imp.TExp Int32 -> + Lambda GPUMem -> + [VName] -> + InKernelGen () +virtualisedGroupScan seg_flag w lam arrs = do + groupChunkLoop w $ \chunk_start chunk_size -> do + constants <- kernelConstants <$> askEnv + let ltid = kernelLocalThreadId constants + crosses_segment = + case seg_flag of + Nothing -> false + Just flag_true -> + flag_true (sExt32 (chunk_start - 1)) (sExt32 chunk_start) + sComment "possibly incorporate carry" $ + sWhen (chunk_start .>. 0 .&&. ltid .==. 0 .&&. bNot crosses_segment) $ do + carry_idx <- dPrimVE "carry_idx" $ sExt64 chunk_start - 1 + applyRenamedLambda + lam + (zip arrs $ repeat [DimFix $ sExt64 chunk_start]) + ( zip (map Var arrs) (repeat [DimFix carry_idx]) + ++ zip (map Var arrs) (repeat [DimFix $ sExt64 chunk_start]) + ) + + arrs_chunks <- mapM (sliceArray (sExt64 chunk_start) chunk_size) arrs + + sOp $ Imp.ErrorSync Imp.FenceLocal + + groupScan seg_flag (sExt64 w) (tvExp chunk_size) lam arrs_chunks + +copyInGroup :: CopyCompiler GPUMem KernelEnv Imp.KernelOp +copyInGroup pt destloc srcloc = do + dest_space <- entryMemSpace <$> lookupMemory (memLocName destloc) + src_space <- entryMemSpace <$> lookupMemory (memLocName srcloc) + + let src_ixfun = memLocIxFun srcloc + dims = IxFun.shape src_ixfun + rank = length dims + + case (dest_space, src_space) of + (ScalarSpace destds _, ScalarSpace srcds _) -> do + let fullDim d = DimSlice 0 d 1 + destslice' = + Slice $ + replicate (rank - length destds) (DimFix 0) + ++ takeLast (length destds) (map fullDim dims) + srcslice' = + Slice $ + replicate (rank - length srcds) (DimFix 0) + ++ takeLast (length srcds) (map fullDim dims) + copyElementWise + pt + (sliceMemLoc destloc destslice') + (sliceMemLoc srcloc srcslice') + _ -> do + groupCoverSpace (map sExt32 dims) $ \is -> + copyElementWise + pt + (sliceMemLoc destloc (Slice $ map (DimFix . sExt64) is)) + (sliceMemLoc srcloc (Slice $ map (DimFix . sExt64) is)) + sOp $ Imp.Barrier Imp.FenceLocal + +localThreadIDs :: [SubExp] -> InKernelGen [Imp.TExp Int64] +localThreadIDs dims = do + ltid <- sExt64 . kernelLocalThreadId . kernelConstants <$> askEnv + let dims' = map pe64 dims + maybe (dIndexSpace' "ltid" dims' ltid) (pure . map sExt64) + . M.lookup dims + . kernelLocalIdMap + . kernelConstants + =<< askEnv + +partitionSeqDims :: SegSeqDims -> SegSpace -> ([(VName, SubExp)], [(VName, SubExp)]) +partitionSeqDims (SegSeqDims seq_is) space = + bimap (map fst) (map fst) $ + partition ((`elem` seq_is) . snd) (zip (unSegSpace space) [0 ..]) + +sanityCheckLevel :: SegLevel -> InKernelGen () +sanityCheckLevel SegThread {} = pure () +sanityCheckLevel SegGroup {} = + error "compileGroupOp: unexpected group-level SegOp." + +compileFlatId :: SegLevel -> SegSpace -> InKernelGen () +compileFlatId lvl space = do + sanityCheckLevel lvl + ltid <- kernelLocalThreadId . kernelConstants <$> askEnv + dPrimV_ (segFlat space) ltid + +-- Construct the necessary lock arrays for an intra-group histogram. +prepareIntraGroupSegHist :: + Count GroupSize SubExp -> + [HistOp GPUMem] -> + InKernelGen [[Imp.TExp Int64] -> InKernelGen ()] +prepareIntraGroupSegHist group_size = + fmap snd . mapAccumLM onOp Nothing + where + onOp l op = do + constants <- kernelConstants <$> askEnv + atomicBinOp <- kernelAtomics <$> askEnv + + let local_subhistos = histDest op + + case (l, atomicUpdateLocking atomicBinOp $ histOp op) of + (_, AtomicPrim f) -> pure (l, f (Space "local") local_subhistos) + (_, AtomicCAS f) -> pure (l, f (Space "local") local_subhistos) + (Just l', AtomicLocking f) -> pure (l, f l' (Space "local") local_subhistos) + (Nothing, AtomicLocking f) -> do + locks <- newVName "locks" + + let num_locks = pe64 $ unCount group_size + dims = map pe64 $ shapeDims (histOpShape op <> histShape op) + l' = Locking locks 0 1 0 (pure . (`rem` num_locks) . flattenIndex dims) + locks_t = Array int32 (Shape [unCount group_size]) NoUniqueness + + locks_mem <- sAlloc "locks_mem" (typeSize locks_t) $ Space "local" + dArray locks int32 (arrayShape locks_t) locks_mem $ + IxFun.iota . map pe64 . arrayDims $ + locks_t + + sComment "All locks start out unlocked" $ + groupCoverSpace [kernelGroupSize constants] $ \is -> + copyDWIMFix locks is (intConst Int32 0) [] + + pure (Just l', f l' (Space "local") local_subhistos) + +groupCoverSegSpace :: SegVirt -> SegSpace -> InKernelGen () -> InKernelGen () +groupCoverSegSpace virt space m = do + let (ltids, dims) = unzip $ unSegSpace space + dims' = map pe64 dims + + constants <- kernelConstants <$> askEnv + let group_size = kernelGroupSize constants + -- Maybe we can statically detect that this is actually a + -- SegNoVirtFull and generate ever-so-slightly simpler code. + let virt' = if dims' == [group_size] then SegNoVirtFull (SegSeqDims []) else virt + case virt' of + SegVirt -> do + iters <- M.lookup dims . kernelChunkItersMap . kernelConstants <$> askEnv + case iters of + Nothing -> do + iterations <- dPrimVE "iterations" $ product $ map sExt32 dims' + groupLoop iterations $ \i -> do + dIndexSpace (zip ltids dims') $ sExt64 i + m + Just num_chunks -> do + let ltid = kernelLocalThreadId constants + sFor "chunk_i" num_chunks $ \chunk_i -> do + i <- dPrimVE "i" $ chunk_i * sExt32 group_size + ltid + dIndexSpace (zip ltids dims') $ sExt64 i + sWhen (inBounds (Slice (map (DimFix . le64) ltids)) dims') m + SegNoVirt -> localOps threadOperations $ do + zipWithM_ dPrimV_ ltids =<< localThreadIDs dims + sWhen (isActive $ zip ltids dims) m + SegNoVirtFull seq_dims -> do + let ((ltids_seq, dims_seq), (ltids_par, dims_par)) = + bimap unzip unzip $ partitionSeqDims seq_dims space + sLoopNest (Shape dims_seq) $ \is_seq -> do + zipWithM_ dPrimV_ ltids_seq is_seq + localOps threadOperations $ do + zipWithM_ dPrimV_ ltids_par =<< localThreadIDs dims_par + m + +compileGroupExp :: ExpCompiler GPUMem KernelEnv Imp.KernelOp +compileGroupExp (Pat [pe]) (BasicOp (Opaque _ se)) = + -- Cannot print in GPU code. + copyDWIM (patElemName pe) [] se [] +-- The static arrays stuff does not work inside kernels. +compileGroupExp (Pat [dest]) (BasicOp (ArrayLit es _)) = + forM_ (zip [0 ..] es) $ \(i, e) -> + copyDWIMFix (patElemName dest) [fromIntegral (i :: Int64)] e [] +compileGroupExp _ (BasicOp (UpdateAcc acc is vs)) = + updateAcc acc is vs +compileGroupExp (Pat [dest]) (BasicOp (Replicate ds se)) = do + flat <- newVName "rep_flat" + is <- replicateM (shapeRank ds) (newVName "rep_i") + let is' = map le64 is + groupCoverSegSpace SegVirt (SegSpace flat $ zip is $ shapeDims ds) $ + copyDWIMFix (patElemName dest) is' se [] + sOp $ Imp.Barrier Imp.FenceLocal +compileGroupExp (Pat [dest]) (BasicOp (Rotate rs arr)) = do + ds <- map pe64 . arrayDims <$> lookupType arr + groupCoverSpace ds $ \is -> do + is' <- sequence $ zipWith3 rotate ds rs is + copyDWIMFix (patElemName dest) is (Var arr) is' + sOp $ Imp.Barrier Imp.FenceLocal + where + rotate d r i = dPrimVE "rot_i" $ rotateIndex d (pe64 r) i +compileGroupExp (Pat [dest]) (BasicOp (Iota n e s it)) = do + n' <- toExp n + e' <- toExp e + s' <- toExp s + groupLoop (TPrimExp n') $ \i' -> do + x <- + dPrimV "x" $ + TPrimExp $ + BinOpExp (Add it OverflowUndef) e' $ + BinOpExp (Mul it OverflowUndef) (untyped i') s' + copyDWIMFix (patElemName dest) [i'] (Var (tvVar x)) [] + sOp $ Imp.Barrier Imp.FenceLocal + +-- When generating code for a scalar in-place update, we must make +-- sure that only one thread performs the write. When writing an +-- array, the group-level copy code will take care of doing the right +-- thing. +compileGroupExp (Pat [pe]) (BasicOp (Update safety _ slice se)) + | null $ sliceDims slice = do + sOp $ Imp.Barrier Imp.FenceLocal + ltid <- kernelLocalThreadId . kernelConstants <$> askEnv + sWhen (ltid .==. 0) $ + case safety of + Unsafe -> write + Safe -> sWhen (inBounds slice' dims) write + sOp $ Imp.Barrier Imp.FenceLocal + where + slice' = fmap pe64 slice + dims = map pe64 $ arrayDims $ patElemType pe + write = copyDWIM (patElemName pe) (unSlice slice') se [] +compileGroupExp dest e = + defCompileExp dest e + +compileGroupOp :: OpCompiler GPUMem KernelEnv Imp.KernelOp +compileGroupOp pat (Alloc size space) = + kernelAlloc pat size space +compileGroupOp pat (Inner (SegOp (SegMap lvl space _ body))) = do + compileFlatId lvl space + + groupCoverSegSpace (segVirt lvl) space $ + compileStms mempty (kernelBodyStms body) $ + zipWithM_ (compileThreadResult space) (patElems pat) $ + kernelBodyResult body + sOp $ Imp.ErrorSync Imp.FenceLocal +compileGroupOp pat (Inner (SegOp (SegScan lvl space scans _ body))) = do + compileFlatId lvl space + + let (ltids, dims) = unzip $ unSegSpace space + dims' = map pe64 dims + + groupCoverSegSpace (segVirt lvl) space $ + compileStms mempty (kernelBodyStms body) $ + forM_ (zip (patNames pat) $ kernelBodyResult body) $ \(dest, res) -> + copyDWIMFix + dest + (map Imp.le64 ltids) + (kernelResultSubExp res) + [] + + fence <- fenceForArrays $ patNames pat + sOp $ Imp.ErrorSync fence + + let segment_size = last dims' + crossesSegment from to = + (sExt64 to - sExt64 from) .>. (sExt64 to `rem` segment_size) + + -- groupScan needs to treat the scan output as a one-dimensional + -- array of scan elements, so we invent some new flattened arrays + -- here. + dims_flat <- dPrimV "dims_flat" $ product dims' + let scan = head scans + num_scan_results = length $ segBinOpNeutral scan + arrs_flat <- + mapM (flattenArray (length dims') dims_flat) $ + take num_scan_results $ + patNames pat + + case segVirt lvl of + SegVirt -> + virtualisedGroupScan + (Just crossesSegment) + (sExt32 $ tvExp dims_flat) + (segBinOpLambda scan) + arrs_flat + _ -> + groupScan + (Just crossesSegment) + (product dims') + (product dims') + (segBinOpLambda scan) + arrs_flat +compileGroupOp pat (Inner (SegOp (SegRed lvl space ops _ body))) = do + compileFlatId lvl space + + let dims' = map pe64 dims + mkTempArr t = + sAllocArray "red_arr" (elemType t) (Shape dims <> arrayShape t) $ Space "local" + + tmp_arrs <- mapM mkTempArr $ concatMap (lambdaReturnType . segBinOpLambda) ops + groupCoverSegSpace (segVirt lvl) space $ + compileStms mempty (kernelBodyStms body) $ do + let (red_res, map_res) = + splitAt (segBinOpResults ops) $ kernelBodyResult body + forM_ (zip tmp_arrs red_res) $ \(dest, res) -> + copyDWIMFix dest (map Imp.le64 ltids) (kernelResultSubExp res) [] + zipWithM_ (compileThreadResult space) map_pes map_res + + sOp $ Imp.ErrorSync Imp.FenceLocal + + let tmps_for_ops = chunks (map (length . segBinOpNeutral) ops) tmp_arrs + case segVirt lvl of + SegVirt -> virtCase dims' tmps_for_ops + _ -> nonvirtCase dims' tmps_for_ops + where + (ltids, dims) = unzip $ unSegSpace space + (red_pes, map_pes) = splitAt (segBinOpResults ops) $ patElems pat + + virtCase [dim'] tmps_for_ops = do + ltid <- kernelLocalThreadId . kernelConstants <$> askEnv + groupChunkLoop (sExt32 dim') $ \chunk_start chunk_size -> do + sComment "possibly incorporate carry" $ + sWhen (chunk_start .>. 0 .&&. ltid .==. 0) $ + forM_ (zip ops tmps_for_ops) $ \(op, tmps) -> + applyRenamedLambda + (segBinOpLambda op) + (zip tmps $ repeat [DimFix $ sExt64 chunk_start]) + ( zip (map (Var . patElemName) red_pes) (repeat []) + ++ zip (map Var tmps) (repeat [DimFix $ sExt64 chunk_start]) + ) + + sOp $ Imp.ErrorSync Imp.FenceLocal + + forM_ (zip ops tmps_for_ops) $ \(op, tmps) -> do + tmps_chunks <- mapM (sliceArray (sExt64 chunk_start) chunk_size) tmps + groupReduce (sExt32 (tvExp chunk_size)) (segBinOpLambda op) tmps_chunks + + sOp $ Imp.ErrorSync Imp.FenceLocal + + forM_ (zip red_pes $ concat tmps_for_ops) $ \(pe, arr) -> + copyDWIMFix (patElemName pe) [] (Var arr) [sExt64 chunk_start] + virtCase dims' tmps_for_ops = do + dims_flat <- dPrimV "dims_flat" $ product dims' + let segment_size = last dims' + crossesSegment from to = + (sExt64 to - sExt64 from) .>. (sExt64 to `rem` sExt64 segment_size) + + forM_ (zip ops tmps_for_ops) $ \(op, tmps) -> do + tmps_flat <- mapM (flattenArray (length dims') dims_flat) tmps + virtualisedGroupScan + (Just crossesSegment) + (sExt32 $ tvExp dims_flat) + (segBinOpLambda op) + tmps_flat + + sOp $ Imp.ErrorSync Imp.FenceLocal + + forM_ (zip red_pes $ concat tmps_for_ops) $ \(pe, arr) -> + copyDWIM + (patElemName pe) + [] + (Var arr) + (map (unitSlice 0) (init dims') ++ [DimFix $ last dims' - 1]) + + sOp $ Imp.Barrier Imp.FenceLocal + + nonvirtCase [dim'] tmps_for_ops = do + -- Nonsegmented case (or rather, a single segment) - this we can + -- handle directly with a group-level reduction. + forM_ (zip ops tmps_for_ops) $ \(op, tmps) -> + groupReduce (sExt32 dim') (segBinOpLambda op) tmps + sOp $ Imp.ErrorSync Imp.FenceLocal + forM_ (zip red_pes $ concat tmps_for_ops) $ \(pe, arr) -> + copyDWIMFix (patElemName pe) [] (Var arr) [0] + -- + nonvirtCase dims' tmps_for_ops = do + -- Segmented intra-group reductions are turned into (regular) + -- segmented scans. It is possible that this can be done + -- better, but at least this approach is simple. + + -- groupScan operates on flattened arrays. This does not + -- involve copying anything; merely playing with the index + -- function. + dims_flat <- dPrimV "dims_flat" $ product dims' + let segment_size = last dims' + crossesSegment from to = + (sExt64 to - sExt64 from) .>. (sExt64 to `rem` sExt64 segment_size) + + forM_ (zip ops tmps_for_ops) $ \(op, tmps) -> do + tmps_flat <- mapM (flattenArray (length dims') dims_flat) tmps + groupScan + (Just crossesSegment) + (product dims') + (product dims') + (segBinOpLambda op) + tmps_flat + + sOp $ Imp.ErrorSync Imp.FenceLocal + + forM_ (zip red_pes $ concat tmps_for_ops) $ \(pe, arr) -> + copyDWIM + (patElemName pe) + [] + (Var arr) + (map (unitSlice 0) (init dims') ++ [DimFix $ last dims' - 1]) + + sOp $ Imp.Barrier Imp.FenceLocal +compileGroupOp pat (Inner (SegOp (SegHist lvl space ops _ kbody))) = do + compileFlatId lvl space + let (ltids, _dims) = unzip $ unSegSpace space + + -- We don't need the red_pes, because it is guaranteed by our type + -- rules that they occupy the same memory as the destinations for + -- the ops. + let num_red_res = length ops + sum (map (length . histNeutral) ops) + (_red_pes, map_pes) = + splitAt num_red_res $ patElems pat + + ops' <- prepareIntraGroupSegHist (segGroupSize lvl) ops + + -- Ensure that all locks have been initialised. + sOp $ Imp.Barrier Imp.FenceLocal + + groupCoverSegSpace (segVirt lvl) space $ + compileStms mempty (kernelBodyStms kbody) $ do + let (red_res, map_res) = splitAt num_red_res $ kernelBodyResult kbody + (red_is, red_vs) = splitAt (length ops) $ map kernelResultSubExp red_res + zipWithM_ (compileThreadResult space) map_pes map_res + + let vs_per_op = chunks (map (length . histDest) ops) red_vs + + forM_ (zip4 red_is vs_per_op ops' ops) $ + \(bin, op_vs, do_op, HistOp dest_shape _ _ _ shape lam) -> do + let bin' = pe64 bin + dest_shape' = map pe64 $ shapeDims dest_shape + bin_in_bounds = inBounds (Slice (map DimFix [bin'])) dest_shape' + bin_is = map Imp.le64 (init ltids) ++ [bin'] + vs_params = takeLast (length op_vs) $ lambdaParams lam + + sComment "perform atomic updates" $ + sWhen bin_in_bounds $ do + dLParams $ lambdaParams lam + sLoopNest shape $ \is -> do + forM_ (zip vs_params op_vs) $ \(p, v) -> + copyDWIMFix (paramName p) [] v is + do_op (bin_is ++ is) + + sOp $ Imp.ErrorSync Imp.FenceLocal +compileGroupOp pat _ = + compilerBugS $ "compileGroupOp: cannot compile rhs of binding " ++ pretty pat + +groupOperations :: Operations GPUMem KernelEnv Imp.KernelOp +groupOperations = + (defaultOperations compileGroupOp) + { opsCopyCompiler = copyInGroup, + opsExpCompiler = compileGroupExp, + opsStmsCompiler = \_ -> defCompileStms mempty, + opsAllocCompilers = + M.fromList [(Space "local", allocLocal)] + } + +arrayInLocalMemory :: SubExp -> InKernelGen Bool +arrayInLocalMemory (Var name) = do + res <- lookupVar name + case res of + ArrayVar _ entry -> + (Space "local" ==) . entryMemSpace + <$> lookupMemory (memLocName (entryArrayLoc entry)) + _ -> pure False +arrayInLocalMemory Constant {} = pure False + +sKernelGroup :: + String -> + VName -> + KernelAttrs -> + InKernelGen () -> + CallKernelGen () +sKernelGroup = sKernel groupOperations kernelGroupId + +compileGroupResult :: + SegSpace -> + PatElem LetDecMem -> + KernelResult -> + InKernelGen () +compileGroupResult _ pe (TileReturns _ [(w, per_group_elems)] what) = do + n <- pe64 . arraySize 0 <$> lookupType what + + constants <- kernelConstants <$> askEnv + let ltid = sExt64 $ kernelLocalThreadId constants + offset = + pe64 per_group_elems + * sExt64 (kernelGroupId constants) + + -- Avoid loop for the common case where each thread is statically + -- known to write at most one element. + localOps threadOperations $ + if pe64 per_group_elems == kernelGroupSize constants + then + sWhen (ltid + offset .<. pe64 w) $ + copyDWIMFix (patElemName pe) [ltid + offset] (Var what) [ltid] + else sFor "i" (n `divUp` kernelGroupSize constants) $ \i -> do + j <- dPrimVE "j" $ kernelGroupSize constants * i + ltid + sWhen (j + offset .<. pe64 w) $ + copyDWIMFix (patElemName pe) [j + offset] (Var what) [j] +compileGroupResult space pe (TileReturns _ dims what) = do + let gids = map fst $ unSegSpace space + out_tile_sizes = map (pe64 . snd) dims + group_is = zipWith (*) (map Imp.le64 gids) out_tile_sizes + local_is <- localThreadIDs $ map snd dims + is_for_thread <- + mapM (dPrimV "thread_out_index") $ + zipWith (+) group_is local_is + + localOps threadOperations $ + sWhen (isActive $ zip (map tvVar is_for_thread) $ map fst dims) $ + copyDWIMFix (patElemName pe) (map tvExp is_for_thread) (Var what) local_is +compileGroupResult space pe (RegTileReturns _ dims_n_tiles what) = do + constants <- kernelConstants <$> askEnv + + let gids = map fst $ unSegSpace space + (dims, group_tiles, reg_tiles) = unzip3 dims_n_tiles + group_tiles' = map pe64 group_tiles + reg_tiles' = map pe64 reg_tiles + + -- Which group tile is this group responsible for? + let group_tile_is = map Imp.le64 gids + + -- Within the group tile, which register tile is this thread + -- responsible for? + reg_tile_is <- + dIndexSpace' "reg_tile_i" group_tiles' $ sExt64 $ kernelLocalThreadId constants + + -- Compute output array slice for the register tile belonging to + -- this thread. + let regTileSliceDim (group_tile, group_tile_i) (reg_tile, reg_tile_i) = do + tile_dim_start <- + dPrimVE "tile_dim_start" $ + reg_tile * (group_tile * group_tile_i + reg_tile_i) + pure $ DimSlice tile_dim_start reg_tile 1 + reg_tile_slices <- + Slice + <$> zipWithM + regTileSliceDim + (zip group_tiles' group_tile_is) + (zip reg_tiles' reg_tile_is) + + localOps threadOperations $ + sLoopNest (Shape reg_tiles) $ \is_in_reg_tile -> do + let dest_is = fixSlice reg_tile_slices is_in_reg_tile + src_is = reg_tile_is ++ is_in_reg_tile + sWhen (foldl1 (.&&.) $ zipWith (.<.) dest_is $ map pe64 dims) $ + copyDWIMFix (patElemName pe) dest_is (Var what) src_is +compileGroupResult space pe (Returns _ _ what) = do + constants <- kernelConstants <$> askEnv + in_local_memory <- arrayInLocalMemory what + let gids = map (Imp.le64 . fst) $ unSegSpace space + + if not in_local_memory + then + localOps threadOperations $ + sWhen (kernelLocalThreadId constants .==. 0) $ + copyDWIMFix (patElemName pe) gids what [] + else -- If the result of the group is an array in local memory, we + -- store it by collective copying among all the threads of the + -- group. TODO: also do this if the array is in global memory + -- (but this is a bit more tricky, synchronisation-wise). + copyDWIMFix (patElemName pe) gids what [] +compileGroupResult _ _ WriteReturns {} = + compilerLimitationS "compileGroupResult: WriteReturns not handled yet." + +-- | The sizes of nested iteration spaces in the kernel. +type SegOpSizes = S.Set [SubExp] + +-- | Various useful precomputed information for group-level SegOps. +data Precomputed = Precomputed + { pcSegOpSizes :: SegOpSizes, + pcChunkItersMap :: M.Map [SubExp] (Imp.TExp Int32) + } + +-- | Find the sizes of nested parallelism in a t'SegOp' body. +segOpSizes :: Stms GPUMem -> SegOpSizes +segOpSizes = onStms + where + onStms = foldMap (onExp . stmExp) + onExp (Op (Inner (SegOp op))) = + case segVirt $ segLevel op of + SegNoVirtFull seq_dims -> + S.singleton $ map snd $ snd $ partitionSeqDims seq_dims $ segSpace op + _ -> S.singleton $ map snd $ unSegSpace $ segSpace op + onExp (BasicOp (Replicate shape _)) = + S.singleton $ shapeDims shape + onExp (Match _ cases defbody _) = + foldMap (onStms . bodyStms . caseBody) cases <> onStms (bodyStms defbody) + onExp (DoLoop _ _ body) = + onStms (bodyStms body) + onExp _ = mempty + +-- | Precompute various constants and useful information. +precomputeConstants :: Count GroupSize (Imp.TExp Int64) -> Stms GPUMem -> CallKernelGen Precomputed +precomputeConstants group_size stms = do + let sizes = segOpSizes stms + iters_map <- M.fromList <$> mapM mkMap (S.toList sizes) + pure $ Precomputed sizes iters_map + where + mkMap dims = do + let n = product $ map Imp.pe64 dims + num_chunks <- dPrimVE "num_chunks" $ sExt32 $ n `divUp` unCount group_size + pure (dims, num_chunks) + +-- | Make use of various precomputed constants. +precomputedConstants :: Precomputed -> InKernelGen a -> InKernelGen a +precomputedConstants pre m = do + ltid <- kernelLocalThreadId . kernelConstants <$> askEnv + new_ids <- M.fromList <$> mapM (mkMap ltid) (S.toList (pcSegOpSizes pre)) + let f env = + env + { kernelConstants = + (kernelConstants env) + { kernelLocalIdMap = new_ids, + kernelChunkItersMap = pcChunkItersMap pre + } + } + localEnv f m + where + mkMap ltid dims = do + let dims' = map pe64 dims + ids' <- dIndexSpace' "ltid_pre" dims' (sExt64 ltid) + pure (dims, map sExt32 ids') diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegMap.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegMap.hs --- a/src/Futhark/CodeGen/ImpGen/GPU/SegMap.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU/SegMap.hs @@ -11,6 +11,7 @@ import qualified Futhark.CodeGen.ImpCode.GPU as Imp import Futhark.CodeGen.ImpGen import Futhark.CodeGen.ImpGen.GPU.Base +import Futhark.CodeGen.ImpGen.GPU.Group import Futhark.IR.GPUMem import Futhark.Util.IntegralExp (divUp) import Prelude hiding (quot, rem) diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs --- a/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs @@ -224,7 +224,7 @@ num_elements global_tid elems_per_thread - (tvVar num_threads) + (tvExp num_threads) slugs body @@ -470,7 +470,7 @@ num_elements global_tid elems_per_thread - (tvVar threads_per_segment) + (tvExp threads_per_segment) slugs body @@ -576,13 +576,53 @@ | otherwise = pure (param_arr, [sExt64 local_tid, sExt64 group_id]) +computeThreadChunkSize :: + Commutativity -> + Imp.TExp Int64 -> + Imp.TExp Int64 -> + Imp.Count Imp.Elements (Imp.TExp Int64) -> + Imp.Count Imp.Elements (Imp.TExp Int64) -> + TV Int64 -> + ImpM rep r op () +computeThreadChunkSize Commutative threads_per_segment thread_index elements_per_thread num_elements chunk_var = + chunk_var + <-- sMin64 + (Imp.unCount elements_per_thread) + ((Imp.unCount num_elements - thread_index) `divUp` threads_per_segment) +computeThreadChunkSize Noncommutative _ thread_index elements_per_thread num_elements chunk_var = do + starting_point <- + dPrimV "starting_point" $ + thread_index * Imp.unCount elements_per_thread + remaining_elements <- + dPrimV "remaining_elements" $ + Imp.unCount num_elements - tvExp starting_point + + let no_remaining_elements = tvExp remaining_elements .<=. 0 + beyond_bounds = Imp.unCount num_elements .<=. tvExp starting_point + + sIf + (no_remaining_elements .||. beyond_bounds) + (chunk_var <-- 0) + ( sIf + is_last_thread + (chunk_var <-- Imp.unCount last_thread_elements) + (chunk_var <-- Imp.unCount elements_per_thread) + ) + where + last_thread_elements = + num_elements - Imp.elements thread_index * elements_per_thread + is_last_thread = + Imp.unCount num_elements + .<. (thread_index + 1) + * Imp.unCount elements_per_thread + reductionStageZero :: KernelConstants -> [(VName, Imp.TExp Int64)] -> Imp.Count Imp.Elements (Imp.TExp Int64) -> Imp.TExp Int64 -> Imp.Count Imp.Elements (Imp.TExp Int64) -> - VName -> + Imp.TExp Int64 -> [SegBinOpSlug] -> DoSegBody -> InKernelGen ([Lambda GPUMem], InKernelGen ()) @@ -593,10 +633,13 @@ -- Figure out how many elements this thread should process. chunk_size <- dPrim "chunk_size" int64 - let ordering = case slugsComm slugs of - Commutative -> SplitStrided $ Var threads_per_segment - Noncommutative -> SplitContiguous - computeThreadChunkSize ordering (sExt64 global_tid) elems_per_thread num_elements chunk_size + computeThreadChunkSize + (slugsComm slugs) + threads_per_segment + (sExt64 global_tid) + elems_per_thread + num_elements + chunk_size dScope Nothing $ scopeOfLParams $ concatMap slugParams slugs @@ -645,7 +688,7 @@ gtid <-- case comm of Commutative -> - global_tid + Imp.le64 threads_per_segment * i + global_tid + threads_per_segment * i Noncommutative -> let index_in_segment = global_tid `quot` kernelGroupSize constants in sExt64 local_tid @@ -696,7 +739,7 @@ Imp.Count Imp.Elements (Imp.TExp Int64) -> Imp.TExp Int64 -> Imp.Count Imp.Elements (Imp.TExp Int64) -> - VName -> + Imp.TExp Int64 -> [SegBinOpSlug] -> DoSegBody -> InKernelGen [Lambda GPUMem] diff --git a/src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs b/src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs --- a/src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs @@ -24,7 +24,6 @@ import qualified Futhark.CodeGen.ImpCode.OpenCL as ImpOpenCL import Futhark.CodeGen.RTS.C (atomicsH, halfH) import Futhark.Error (compilerLimitationS) -import Futhark.IR.Prop (isBuiltInFunction) import Futhark.MonadFreshNames import Futhark.Util (zEncodeString) import Futhark.Util.Pretty (prettyOneLine, prettyText) @@ -787,20 +786,17 @@ then [C.citems|return 1;|] else [C.citems|return;|] - callInKernel dests fname args - | isBuiltInFunction fname = - GC.opsCall GC.defaultOperations dests fname args - | otherwise = do - let out_args = [[C.cexp|&$id:d|] | d <- dests] - args' = - [C.cexp|global_failure|] - : [C.cexp|global_failure_args|] - : out_args - ++ args + callInKernel dests fname args = do + let out_args = [[C.cexp|&$id:d|] | d <- dests] + args' = + [C.cexp|global_failure|] + : [C.cexp|global_failure_args|] + : out_args + ++ args - what_next <- whatNext + what_next <- whatNext - GC.item [C.citem|if ($id:(funName fname)($args:args') != 0) { $items:what_next; }|] + GC.item [C.citem|if ($id:(funName fname)($args:args') != 0) { $items:what_next; }|] errorInKernel msg@(ErrorMsg parts) backtrace = do n <- length . kernelFailures <$> GC.getUserState diff --git a/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs b/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs --- a/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs +++ b/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs @@ -136,8 +136,6 @@ compileThreadResult space pe (Returns _ _ what) = do let is = map (Imp.le64 . fst) $ unSegSpace space copyDWIMFix (patElemName pe) is what [] -compileThreadResult _ _ ConcatReturns {} = - compilerBugS "compileThreadResult: ConcatReturn unhandled." compileThreadResult _ _ WriteReturns {} = compilerBugS "compileThreadResult: WriteReturns unhandled." compileThreadResult _ _ TileReturns {} = diff --git a/src/Futhark/Construct.hs b/src/Futhark/Construct.hs --- a/src/Futhark/Construct.hs +++ b/src/Futhark/Construct.hs @@ -84,8 +84,6 @@ eCopy, eBody, eLambda, - eRoundToMultipleOf, - eSliceArray, eBlank, eAll, eAny, @@ -386,39 +384,6 @@ bodyBind $ lambdaBody lam where bindParam param arg = letBindNames [paramName param] =<< arg - --- | @eRoundToMultipleOf t x d@ produces an expression that rounds the --- integer expression @x@ upwards to be a multiple of @d@, with @t@ --- being the integer type of the expressions. -eRoundToMultipleOf :: - MonadBuilder m => - IntType -> - m (Exp (Rep m)) -> - m (Exp (Rep m)) -> - m (Exp (Rep m)) -eRoundToMultipleOf t x d = - ePlus x (eMod (eMinus d (eMod x d)) d) - where - eMod = eBinOp (SMod t Unsafe) - eMinus = eBinOp (Sub t OverflowWrap) - ePlus = eBinOp (Add t OverflowWrap) - --- | Construct an 'Index' expressions that slices an array with unit stride. -eSliceArray :: - MonadBuilder m => - Int -> - VName -> - m (Exp (Rep m)) -> - m (Exp (Rep m)) -> - m (Exp (Rep m)) -eSliceArray d arr i n = do - arr_t <- lookupType arr - let skips = map (slice (constant (0 :: Int64))) $ take d $ arrayDims arr_t - i' <- letSubExp "slice_i" =<< i - n' <- letSubExp "slice_n" =<< n - pure $ BasicOp $ Index arr $ fullSlice arr_t $ skips ++ [slice i' n'] - where - slice j m = DimSlice j m (constant (1 :: Int64)) -- | @eInBoundsForDim w i@ produces @0 <= i < w@. eDimInBounds :: MonadBuilder m => m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m)) diff --git a/src/Futhark/Doc/Generator.hs b/src/Futhark/Doc/Generator.hs --- a/src/Futhark/Doc/Generator.hs +++ b/src/Futhark/Doc/Generator.hs @@ -80,7 +80,7 @@ type DocM = ReaderT Context (WriterT Documented (Writer Warnings)) -data IndexWhat = IndexValue | IxFun | IndexModule | IndexModuleType | IndexType +data IndexWhat = IndexValue | IndexFunction | IndexModule | IndexModuleType | IndexType -- | We keep a mapping of the names we have actually documented, so we -- can generate an index. @@ -302,7 +302,7 @@ baseString name what' = case what of IndexValue -> "value" - IxFun -> "function" + IndexFunction -> "function" IndexType -> "type" IndexModuleType -> "module type" IndexModule -> "module" @@ -849,7 +849,7 @@ orderZero t = IndexValue | otherwise = - IxFun + IndexFunction describeSpecs :: [Spec] -> DocM Html describeSpecs specs = @@ -864,7 +864,7 @@ where what = case t of - TEArrow {} -> IxFun + TEArrow {} -> IndexFunction _ -> IndexValue describeSpec (TypeAbbrSpec vb) = describeGeneric (typeAlias vb) IndexType (typeDoc vb) (`typeBindHtml` vb) diff --git a/src/Futhark/IR/GPU/Op.hs b/src/Futhark/IR/GPU/Op.hs --- a/src/Futhark/IR/GPU/Op.hs +++ b/src/Futhark/IR/GPU/Op.hs @@ -114,29 +114,7 @@ -- | A simple size-level query or computation. data SizeOp - = -- | @SplitSpace o w i elems_per_thread@. - -- - -- Computes how to divide array elements to - -- threads in a kernel. Returns the number of - -- elements in the chunk that the current thread - -- should take. - -- - -- @w@ is the length of the outer dimension in - -- the array. @i@ is the current thread - -- index. Each thread takes at most - -- @elems_per_thread@ elements. - -- - -- If the order @o@ is 'SplitContiguous', thread with index @i@ - -- should receive elements - -- @i*elems_per_tread, i*elems_per_thread + 1, - -- ..., i*elems_per_thread + (elems_per_thread-1)@. - -- - -- If the order @o@ is @'SplitStrided' stride@, - -- the thread will receive elements @i, - -- i+stride, i+2*stride, ..., - -- i+(elems_per_thread-1)*stride@. - SplitSpace SplitOrdering SubExp SubExp SubExp - | -- | Produce some runtime-configurable size. + = -- | Produce some runtime-configurable size. GetSize Name SizeClass | -- | The maximum size of some class. GetSizeMax SizeClass @@ -150,12 +128,6 @@ deriving (Eq, Ord, Show) instance Substitute SizeOp where - substituteNames subst (SplitSpace o w i elems_per_thread) = - SplitSpace - (substituteNames subst o) - (substituteNames subst w) - (substituteNames subst i) - (substituteNames subst elems_per_thread) substituteNames substs (CmpSizeLe name sclass x) = CmpSizeLe name sclass (substituteNames substs x) substituteNames substs (CalcNumGroups w max_num_groups group_size) = @@ -166,12 +138,6 @@ substituteNames _ op = op instance Rename SizeOp where - rename (SplitSpace o w i elems_per_thread) = - SplitSpace - <$> rename o - <*> rename w - <*> rename i - <*> rename elems_per_thread rename (CmpSizeLe name sclass x) = CmpSizeLe name sclass <$> rename x rename (CalcNumGroups w max_num_groups group_size) = @@ -183,7 +149,6 @@ cheapOp _ = True instance TypedOp SizeOp where - opType SplitSpace {} = pure [Prim int64] opType (GetSize _ _) = pure [Prim int64] opType (GetSizeMax _) = pure [Prim int64] opType CmpSizeLe {} = pure [Prim Bool] @@ -194,19 +159,11 @@ consumedInOp _ = mempty instance FreeIn SizeOp where - freeIn' (SplitSpace o w i elems_per_thread) = - freeIn' o <> freeIn' [w, i, elems_per_thread] freeIn' (CmpSizeLe _ _ x) = freeIn' x freeIn' (CalcNumGroups w _ group_size) = freeIn' w <> freeIn' group_size freeIn' _ = mempty instance PP.Pretty SizeOp where - ppr (SplitSpace SplitContiguous w i elems_per_thread) = - text "split_space" - <> parens (commasep [ppr w, ppr i, ppr elems_per_thread]) - ppr (SplitSpace (SplitStrided stride) w i elems_per_thread) = - text "split_space_strided" - <> parens (commasep [ppr stride, ppr w, ppr i, ppr elems_per_thread]) ppr (GetSize name size_class) = text "get_size" <> parens (commasep [ppr name, ppr size_class]) ppr (GetSizeMax size_class) = @@ -219,18 +176,12 @@ text "calc_num_groups" <> parens (commasep [ppr w, ppr max_num_groups, ppr group_size]) instance OpMetrics SizeOp where - opMetrics SplitSpace {} = seen "SplitSpace" opMetrics GetSize {} = seen "GetSize" opMetrics GetSizeMax {} = seen "GetSizeMax" opMetrics CmpSizeLe {} = seen "CmpSizeLe" opMetrics CalcNumGroups {} = seen "CalcNumGroups" typeCheckSizeOp :: TC.Checkable rep => SizeOp -> TC.TypeM rep () -typeCheckSizeOp (SplitSpace o w i elems_per_thread) = do - case o of - SplitContiguous -> pure () - SplitStrided stride -> TC.require [Prim int64] stride - mapM_ (TC.require [Prim int64]) [w, i, elems_per_thread] typeCheckSizeOp GetSize {} = pure () typeCheckSizeOp GetSizeMax {} = pure () typeCheckSizeOp (CmpSizeLe _ _ x) = TC.require [Prim int64] x diff --git a/src/Futhark/IR/GPU/Simplify.hs b/src/Futhark/IR/GPU/Simplify.hs --- a/src/Futhark/IR/GPU/Simplify.hs +++ b/src/Futhark/IR/GPU/Simplify.hs @@ -56,17 +56,6 @@ simplifyKernelOp _ (SegOp op) = do (op', hoisted) <- simplifySegOp op pure (SegOp op', hoisted) -simplifyKernelOp _ (SizeOp (SplitSpace o w i elems_per_thread)) = - (,) - <$> ( SizeOp - <$> ( SplitSpace - <$> Engine.simplify o - <*> Engine.simplify w - <*> Engine.simplify i - <*> Engine.simplify elems_per_thread - ) - ) - <*> pure mempty simplifyKernelOp _ (SizeOp (GetSize key size_class)) = pure (SizeOp $ GetSize key size_class, mempty) simplifyKernelOp _ (SizeOp (GetSizeMax size_class)) = diff --git a/src/Futhark/IR/Mem.hs b/src/Futhark/IR/Mem.hs --- a/src/Futhark/IR/Mem.hs +++ b/src/Futhark/IR/Mem.hs @@ -1124,7 +1124,7 @@ & MemArray et (Shape (flatSliceDims slice)) NoUniqueness & pure -class TypedOp op => OpReturns op where +class IsOp op => OpReturns op where opReturns :: (Mem rep inner, Monad m, HasScope rep m) => op -> m [ExpReturns] opReturns op = extReturns <$> opType op diff --git a/src/Futhark/IR/Mem/Simplify.hs b/src/Futhark/IR/Mem/Simplify.hs --- a/src/Futhark/IR/Mem/Simplify.hs +++ b/src/Futhark/IR/Mem/Simplify.hs @@ -19,6 +19,7 @@ import Futhark.Construct import Futhark.IR.Mem import qualified Futhark.IR.Mem.IxFun as IxFun +import Futhark.IR.Prop.Aliases (AliasedOp) import qualified Futhark.Optimise.Simplify as Simplify import qualified Futhark.Optimise.Simplify.Engine as Engine import Futhark.Optimise.Simplify.Rep @@ -28,6 +29,20 @@ import Futhark.Pass.ExplicitAllocations (simplifiable) import Futhark.Util +-- | Some constraints that must hold for the simplification rules to work. +type SimplifyMemory rep inner = + ( Simplify.SimplifiableRep rep, + LetDec rep ~ LetDecMem, + ExpDec rep ~ (), + BodyDec rep ~ (), + CanBeWise (Op rep), + BuilderOps (Wise rep), + OpReturns (OpWithWisdom inner), + ST.IndexOp (OpWithWisdom inner), + AliasedOp (OpWithWisdom inner), + Mem rep inner + ) + simpleGeneric :: (SimplifyMemory rep inner) => (OpWithWisdom inner -> UT.UsageTable) -> @@ -92,18 +107,6 @@ Engine.blockHoistSeq = isResultAlloc, Engine.isAllocation = isAlloc mempty mempty } - --- | Some constraints that must hold for the simplification rules to work. -type SimplifyMemory rep inner = - ( Simplify.SimplifiableRep rep, - LetDec rep ~ LetDecMem, - ExpDec rep ~ (), - BodyDec rep ~ (), - CanBeWise (Op rep), - BuilderOps (Wise rep), - OpReturns (OpWithWisdom inner), - Mem rep inner - ) callKernelRules :: SimplifyMemory rep inner => RuleBook (Wise rep) callKernelRules = diff --git a/src/Futhark/IR/Parse.hs b/src/Futhark/IR/Parse.hs --- a/src/Futhark/IR/Parse.hs +++ b/src/Futhark/IR/Parse.hs @@ -736,28 +736,7 @@ <*> braces (pSubExp `sepBy` pComma) <* pComma <*> pLambda pr - pStream = - choice - [ keyword "streamParComm" *> pStreamPar SOAC.InOrder Commutative, - keyword "streamPar" *> pStreamPar SOAC.InOrder Noncommutative, - keyword "streamParPerComm" *> pStreamPar SOAC.Disorder Commutative, - keyword "streamParPer" *> pStreamPar SOAC.Disorder Noncommutative, - keyword "streamSeq" *> pStreamSeq - ] - pStreamPar order comm = - parens $ - SOAC.Stream - <$> pSubExp - <* pComma - <*> braces (pVName `sepBy` pComma) - <* pComma - <*> pParForm order comm - <* pComma - <*> braces (pSubExp `sepBy` pComma) - <* pComma - <*> pLambda pr - pParForm order comm = - SOAC.Parallel order comm <$> pLambda pr + pStream = keyword "streamSeq" *> pStreamSeq pStreamSeq = parens $ SOAC.Stream @@ -765,7 +744,6 @@ <* pComma <*> braces (pVName `sepBy` pComma) <* pComma - <*> pure SOAC.Sequential <*> braces (pSubExp `sepBy` pComma) <* pComma <*> pLambda pr @@ -832,26 +810,6 @@ <*> pName <* pComma <*> pSubExp - ), - keyword "split_space" - *> parens - ( GPU.SplitSpace GPU.SplitContiguous - <$> pSubExp - <* pComma - <*> pSubExp - <* pComma - <*> pSubExp - ), - keyword "split_space_strided" - *> parens - ( GPU.SplitSpace - <$> (GPU.SplitStrided <$> pSubExp) - <* pComma - <*> pSubExp - <* pComma - <*> pSubExp - <* pComma - <*> pSubExp ) ] @@ -888,24 +846,6 @@ <*> pVName, try "blkreg_tile" *> parens (SegOp.RegTileReturns cs <$> (pRegTile `sepBy` pComma)) - <*> pVName, - keyword "concat" - *> parens - ( SegOp.ConcatReturns cs SegOp.SplitContiguous - <$> pSubExp - <* pComma - <*> pSubExp - ) - <*> pVName, - keyword "concat_strided" - *> parens - ( SegOp.ConcatReturns cs - <$> (SegOp.SplitStrided <$> pSubExp) - <* pComma - <*> pSubExp - <* pComma - <*> pSubExp - ) <*> pVName ] where diff --git a/src/Futhark/IR/SOACS.hs b/src/Futhark/IR/SOACS.hs --- a/src/Futhark/IR/SOACS.hs +++ b/src/Futhark/IR/SOACS.hs @@ -57,7 +57,7 @@ lamUsesAD = bodyUsesAD . lambdaBody expUsesAD (Op JVP {}) = True expUsesAD (Op VJP {}) = True - expUsesAD (Op (Stream _ _ _ _ lam)) = lamUsesAD lam + expUsesAD (Op (Stream _ _ _ lam)) = lamUsesAD lam expUsesAD (Op (Screma _ _ (ScremaForm scans reds lam))) = lamUsesAD lam || any (lamUsesAD . scanLambda) scans diff --git a/src/Futhark/IR/SOACS/SOAC.hs b/src/Futhark/IR/SOACS/SOAC.hs --- a/src/Futhark/IR/SOACS/SOAC.hs +++ b/src/Futhark/IR/SOACS/SOAC.hs @@ -9,8 +9,6 @@ -- the main form of parallelism in the early stages of the compiler. module Futhark.IR.SOACS.SOAC ( SOAC (..), - StreamOrd (..), - StreamForm (..), ScremaForm (..), HistOp (..), Scan (..), @@ -81,7 +79,7 @@ -- | A second-order array combinator (SOAC). data SOAC rep - = Stream SubExp [VName] (StreamForm rep) [SubExp] (Lambda rep) + = Stream SubExp [VName] [SubExp] (Lambda rep) | -- | @Scatter @ -- -- Scatter maps values from a set of input arrays to indices and values of a @@ -145,19 +143,6 @@ } deriving (Eq, Ord, Show) --- | Is the stream chunk required to correspond to a contiguous --- subsequence of the original input ('InOrder') or not? 'Disorder' --- streams can be more efficient, but not all algorithms work with --- this. -data StreamOrd = InOrder | Disorder - deriving (Eq, Ord, Show) - --- | What kind of stream is this? -data StreamForm rep - = Parallel StreamOrd Commutativity (Lambda rep) - | Sequential - deriving (Eq, Ord, Show) - -- | The essential parts of a 'Screma' factored out (everything -- except the input arrays). data ScremaForm rep @@ -416,18 +401,12 @@ <$> mapOnSOACLambda tv lam <*> mapM (mapOnSOACSubExp tv) args <*> mapM (mapOnSOACSubExp tv) vec -mapSOACM tv (Stream size arrs form accs lam) = +mapSOACM tv (Stream size arrs accs lam) = Stream <$> mapOnSOACSubExp tv size <*> mapM (mapOnSOACVName tv) arrs - <*> mapOnStreamForm form <*> mapM (mapOnSOACSubExp tv) accs <*> mapOnSOACLambda tv lam - where - mapOnStreamForm (Parallel o comm lam0) = - Parallel o comm <$> mapOnSOACLambda tv lam0 - mapOnStreamForm Sequential = - pure Sequential mapSOACM tv (Scatter w ivs lam as) = Scatter <$> mapOnSOACSubExp tv w @@ -494,10 +473,6 @@ freeIn' (ScremaForm scans reds lam) = freeIn' scans <> freeIn' reds <> freeIn' lam -instance ASTRep rep => FreeIn (StreamForm rep) where - freeIn' Sequential = mempty - freeIn' (Parallel _ _ lam) = freeIn' lam - instance ASTRep rep => FreeIn (HistOp rep) where freeIn' (HistOp w rf dests nes lam) = freeIn' w <> freeIn' rf <> freeIn' dests <> freeIn' nes <> freeIn' lam @@ -537,7 +512,7 @@ soacType (VJP lam _ _) = lambdaReturnType lam ++ map paramType (lambdaParams lam) -soacType (Stream outersize _ _ accs lam) = +soacType (Stream outersize _ accs lam) = map (substNamesInType substs) rtp where nms = map paramName $ take (1 + length accs) params @@ -570,14 +545,8 @@ where consumedArray v = fromMaybe v $ lookup v params_to_arrs params_to_arrs = zip (map paramName $ lambdaParams map_lam) arrs - consumedInOp (Stream _ arrs form accs lam) = - namesFromList $ - subExpVars $ - case form of - Sequential -> - map consumedArray $ namesToList $ consumedByLambda lam - Parallel {} -> - map consumedArray $ namesToList $ consumedByLambda lam + consumedInOp (Stream _ arrs accs lam) = + namesFromList $ subExpVars $ map consumedArray $ namesToList $ consumedByLambda lam where consumedArray v = fromMaybe (Var v) $ lookup v paramsToInput -- Drop the chunk parameter, which cannot alias anything. @@ -608,13 +577,8 @@ JVP (Alias.analyseLambda aliases lam) args vec addOpAliases aliases (VJP lam args vec) = VJP (Alias.analyseLambda aliases lam) args vec - addOpAliases aliases (Stream size arr form accs lam) = - Stream size arr (analyseStreamForm form) accs $ - Alias.analyseLambda aliases lam - where - analyseStreamForm (Parallel o comm lam0) = - Parallel o comm (Alias.analyseLambda aliases lam0) - analyseStreamForm Sequential = Sequential + addOpAliases aliases (Stream size arr accs lam) = + Stream size arr accs $ Alias.analyseLambda aliases lam addOpAliases aliases (Scatter len arrs lam dests) = Scatter len arrs (Alias.analyseLambda aliases lam) dests addOpAliases aliases (Hist w arrs ops bucket_fun) = @@ -720,7 +684,7 @@ PP.indent 2 (ppr $ map TC.argType args') "does not match type of seed vector" PP.indent 2 (ppr vec_ts) -typeCheckSOAC (Stream size arrexps form accexps lam) = do +typeCheckSOAC (Stream size arrexps accexps lam) = do TC.require [Prim int64] size accargs <- mapM TC.checkArg accexps arrargs <- mapM lookupType arrexps @@ -734,21 +698,6 @@ unless (map TC.argType accargs == lamrtp) $ TC.bad $ TC.TypeError "Stream with inconsistent accumulator type in lambda." - -- check reduce's lambda, if any - _ <- case form of - Parallel _ _ lam0 -> do - let acct = map TC.argType accargs - outerRetType = lambdaReturnType lam0 - TC.checkLambda lam0 $ map TC.noArgAliases $ accargs ++ accargs - unless (acct == outerRetType) $ - TC.bad $ - TC.TypeError $ - "Initial value is of type " - ++ prettyTuple acct - ++ ", but stream's reduce lambda returns type " - ++ prettyTuple outerRetType - ++ "." - Sequential -> pure () -- just get the dflow of lambda on the fakearg, which does not alias -- arr, so we can later check that aliases of arr are not used inside lam. let fake_lamarrs' = map asArg lamarrs' @@ -901,7 +850,7 @@ inside "VJP" $ lambdaMetrics lam opMetrics (JVP lam _ _) = inside "JVP" $ lambdaMetrics lam - opMetrics (Stream _ _ _ _ lam) = + opMetrics (Stream _ _ _ lam) = inside "Stream" $ lambdaMetrics lam opMetrics (Scatter _len _ lam _) = inside "Scatter" $ lambdaMetrics lam @@ -930,8 +879,8 @@ PP.braces (commasep $ map ppr args) <> comma PP.braces (commasep $ map ppr vec) ) - ppr (Stream size arrs form acc lam) = - ppStream size arrs form acc lam + ppr (Stream size arrs acc lam) = + ppStream size arrs acc lam ppr (Scatter w arrs lam dests) = ppScatter w arrs lam dests ppr (Hist w arrs ops bucket_fun) = @@ -978,30 +927,15 @@ -- | Prettyprint the given Stream. ppStream :: - (PrettyRep rep, Pretty inp) => SubExp -> [inp] -> StreamForm rep -> [SubExp] -> Lambda rep -> Doc -ppStream size arrs form acc lam = - case form of - Parallel o comm lam0 -> - let ord_str = if o == Disorder then "Per" else "" - comm_str = case comm of - Commutative -> "Comm" - Noncommutative -> "" - in text ("streamPar" ++ ord_str ++ comm_str) - <> parens - ( ppr size <> comma - ppTuple' arrs <> comma - ppr lam0 <> comma - ppTuple' acc <> comma - ppr lam - ) - Sequential -> - text "streamSeq" - <> parens - ( ppr size <> comma - ppTuple' arrs <> comma - ppTuple' acc <> comma - ppr lam - ) + (PrettyRep rep, Pretty inp) => SubExp -> [inp] -> [SubExp] -> Lambda rep -> Doc +ppStream size arrs acc lam = + text "streamSeq" + <> parens + ( ppr size <> comma + ppTuple' arrs <> comma + ppTuple' acc <> comma + ppr lam + ) -- | Prettyprint the given Scatter. ppScatter :: diff --git a/src/Futhark/IR/SOACS/Simplify.hs b/src/Futhark/IR/SOACS/Simplify.hs --- a/src/Futhark/IR/SOACS/Simplify.hs +++ b/src/Futhark/IR/SOACS/Simplify.hs @@ -96,22 +96,12 @@ arr' <- mapM Engine.simplify arr vec' <- mapM Engine.simplify vec pure (JVP lam' arr' vec', hoisted) -simplifySOAC (Stream outerdim arr form nes lam) = do +simplifySOAC (Stream outerdim arr nes lam) = do outerdim' <- Engine.simplify outerdim - (form', form_hoisted) <- simplifyStreamForm form nes' <- mapM Engine.simplify nes arr' <- mapM Engine.simplify arr (lam', lam_hoisted) <- Engine.enterLoop $ Engine.simplifyLambda lam - pure - ( Stream outerdim' arr' form' nes' lam', - form_hoisted <> lam_hoisted - ) - where - simplifyStreamForm (Parallel o comm lam0) = do - (lam0', hoisted) <- Engine.simplifyLambda lam0 - pure (Parallel o comm lam0', hoisted) - simplifyStreamForm Sequential = - pure (Sequential, mempty) + pure (Stream outerdim' arr' nes' lam', lam_hoisted) simplifySOAC (Scatter w ivs lam as) = do w' <- Engine.simplify w (lam', hoisted) <- Engine.enterLoop $ Engine.simplifyLambda lam @@ -317,7 +307,7 @@ liftIdentityMapping _ _ _ _ = Skip liftIdentityStreaming :: BottomUpRuleOp (Wise SOACS) -liftIdentityStreaming _ (Pat pes) aux (Stream w arrs form nes lam) +liftIdentityStreaming _ (Pat pes) aux (Stream w arrs nes lam) | (variant_map, invariant_map) <- partitionEithers $ map isInvariantRes $ zip3 map_ts map_pes map_res, not $ null invariant_map = Simplify $ do @@ -331,10 +321,8 @@ lambdaReturnType = fold_ts ++ variant_map_ts } - auxing aux $ - letBind (Pat $ fold_pes ++ variant_map_pes) $ - Op $ - Stream w arrs form nes lam' + auxing aux . letBind (Pat $ fold_pes ++ variant_map_pes) . Op $ + Stream w arrs nes lam' where num_folds = length nes (fold_pes, map_pes) = splitAt num_folds pes @@ -680,7 +668,7 @@ zipWithM_ bindResult red_pes red_res zipWithM_ bindArrayResult map_pes map_res simplifyKnownIterationSOAC _ pat _ op - | Just (Stream (Constant k) arrs _ nes fold_lam) <- asSOAC op, + | Just (Stream (Constant k) arrs nes fold_lam) <- asSOAC op, oneIsh k = Simplify $ do let (chunk_param, acc_params, slice_params) = partitionChunkedFoldParameters (length nes) (lambdaParams fold_lam) diff --git a/src/Futhark/IR/SegOp.hs b/src/Futhark/IR/SegOp.hs --- a/src/Futhark/IR/SegOp.hs +++ b/src/Futhark/IR/SegOp.hs @@ -36,7 +36,6 @@ KernelResult (..), kernelResultCerts, kernelResultSubExp, - SplitOrdering (..), -- ** Generic traversal SegOpMapper (..), @@ -105,28 +104,6 @@ import qualified Futhark.Util.Pretty as PP import Prelude hiding (id, (.)) --- | How an array is split into chunks. -data SplitOrdering - = SplitContiguous - | SplitStrided SubExp - deriving (Eq, Ord, Show) - -instance FreeIn SplitOrdering where - freeIn' SplitContiguous = mempty - freeIn' (SplitStrided stride) = freeIn' stride - -instance Substitute SplitOrdering where - substituteNames _ SplitContiguous = - SplitContiguous - substituteNames subst (SplitStrided stride) = - SplitStrided $ substituteNames subst stride - -instance Rename SplitOrdering where - rename SplitContiguous = - pure SplitContiguous - rename (SplitStrided stride) = - SplitStrided <$> rename stride - -- | An operator for 'SegHist'. data HistOp rep = HistOp { histShape :: Shape, @@ -227,13 +204,6 @@ Shape -- Size of array. Must match number of dims. VName -- Which array [(Slice SubExp, SubExp)] - | -- Arbitrary number of index/value pairs. - ConcatReturns - Certs - SplitOrdering -- Permuted? - SubExp -- The final size. - SubExp -- Per-thread/group (max) chunk size. - VName -- Chunk by this worker. | TileReturns Certs [(SubExp, SubExp)] -- Total/tile for each dimension @@ -255,7 +225,6 @@ kernelResultCerts :: KernelResult -> Certs kernelResultCerts (Returns _ cs _) = cs kernelResultCerts (WriteReturns cs _ _ _) = cs -kernelResultCerts (ConcatReturns cs _ _ _ _) = cs kernelResultCerts (TileReturns cs _ _) = cs kernelResultCerts (RegTileReturns cs _ _) = cs @@ -263,15 +232,12 @@ kernelResultSubExp :: KernelResult -> SubExp kernelResultSubExp (Returns _ _ se) = se kernelResultSubExp (WriteReturns _ _ arr _) = Var arr -kernelResultSubExp (ConcatReturns _ _ _ _ v) = Var v kernelResultSubExp (TileReturns _ _ v) = Var v kernelResultSubExp (RegTileReturns _ _ v) = Var v instance FreeIn KernelResult where freeIn' (Returns _ cs what) = freeIn' cs <> freeIn' what freeIn' (WriteReturns cs rws arr res) = freeIn' cs <> freeIn' rws <> freeIn' arr <> freeIn' res - freeIn' (ConcatReturns cs o w per_thread_elems v) = - freeIn' cs <> freeIn' o <> freeIn' w <> freeIn' per_thread_elems <> freeIn' v freeIn' (TileReturns cs dims v) = freeIn' cs <> freeIn' dims <> freeIn' v freeIn' (RegTileReturns cs dims_n_tiles v) = @@ -299,13 +265,6 @@ (substituteNames subst rws) (substituteNames subst arr) (substituteNames subst res) - substituteNames subst (ConcatReturns cs o w per_thread_elems v) = - ConcatReturns - (substituteNames subst cs) - (substituteNames subst o) - (substituteNames subst w) - (substituteNames subst per_thread_elems) - (substituteNames subst v) substituteNames subst (TileReturns cs dims v) = TileReturns (substituteNames subst cs) @@ -412,18 +371,6 @@ ++ pretty shape ++ ", but destination array has type " ++ pretty arr_t - checkKernelResult (ConcatReturns cs o w per_thread_elems v) t = do - TC.checkCerts cs - case o of - SplitContiguous -> pure () - SplitStrided stride -> TC.require [Prim int64] stride - TC.require [Prim int64] w - TC.require [Prim int64] per_thread_elems - vt <- lookupType v - unless (vt == t `arrayOfRow` arraySize 0 vt) $ - TC.bad $ - TC.TypeError $ - "Invalid type for ConcatReturns " ++ pretty v checkKernelResult (TileReturns cs dims v) t = do TC.checkCerts cs forM_ dims $ \(dim, tile) -> do @@ -484,20 +431,6 @@ ] where ppRes (slice, e) = ppr slice <+> text "=" <+> ppr e - ppr (ConcatReturns cs SplitContiguous w per_thread_elems v) = - PP.spread $ - certAnnots cs - ++ [ "concat" - <> parens (commasep [ppr w, ppr per_thread_elems]) - <+> ppr v - ] - ppr (ConcatReturns cs (SplitStrided stride) w per_thread_elems v) = - PP.spread $ - certAnnots cs - ++ [ "concat_strided" - <> parens (commasep [ppr stride, ppr w, ppr per_thread_elems]) - <+> ppr v - ] ppr (TileReturns cs dims v) = PP.spread $ certAnnots cs ++ ["tile" <> parens (commasep $ map onDim dims) <+> ppr v] where @@ -613,8 +546,6 @@ t `arrayOfShape` shape segResultShape space t Returns {} = foldr (flip arrayOfRow) t $ segSpaceDims space -segResultShape _ t (ConcatReturns _ _ w _ _) = - t `arrayOfRow` w segResultShape _ t (TileReturns _ dims _) = t `arrayOfShape` Shape (map fst dims) segResultShape _ t (RegTileReturns _ dims_n_tiles _) = @@ -1123,12 +1054,6 @@ --- Simplification -instance Engine.Simplifiable SplitOrdering where - simplify SplitContiguous = - pure SplitContiguous - simplify (SplitStrided stride) = - SplitStrided <$> Engine.simplify stride - instance Engine.Simplifiable SegSpace where simplify (SegSpace phys dims) = SegSpace phys <$> mapM (traverse Engine.simplify) dims @@ -1142,13 +1067,6 @@ <*> Engine.simplify ws <*> Engine.simplify a <*> Engine.simplify res - simplify (ConcatReturns cs o w pte what) = - ConcatReturns - <$> Engine.simplify cs - <*> Engine.simplify o - <*> Engine.simplify w - <*> Engine.simplify pte - <*> Engine.simplify what simplify (TileReturns cs dims what) = TileReturns <$> Engine.simplify cs <*> Engine.simplify dims <*> Engine.simplify what simplify (RegTileReturns cs dims_n_tiles what) = diff --git a/src/Futhark/Internalise/Exps.hs b/src/Futhark/Internalise/Exps.hs --- a/src/Futhark/Internalise/Exps.hs +++ b/src/Futhark/Internalise/Exps.hs @@ -1165,82 +1165,6 @@ letValExp' desc . I.Op $ I.Hist w_img (buckets' ++ img') [HistOp shape_hist rf' hist' ne_shp op'] lam' -internaliseStreamMap :: - String -> - StreamOrd -> - E.Exp -> - E.Exp -> - InternaliseM [SubExp] -internaliseStreamMap desc o lam arr = do - arrs <- internaliseExpToVars "stream_input" arr - lam' <- internaliseStreamMapLambda internaliseLambda lam $ map I.Var arrs - w <- arraysSize 0 <$> mapM lookupType arrs - let form = I.Parallel o Commutative (I.Lambda [] (mkBody mempty []) []) - letValExp' desc $ I.Op $ I.Stream w arrs form [] lam' - -internaliseStreamRed :: - String -> - StreamOrd -> - Commutativity -> - E.Exp -> - E.Exp -> - E.Exp -> - InternaliseM [SubExp] -internaliseStreamRed desc o comm lam0 lam arr = do - arrs <- internaliseExpToVars "stream_input" arr - rowts <- mapM (fmap I.rowType . lookupType) arrs - (lam_params, lam_body) <- - internaliseStreamLambda internaliseLambda lam rowts - let (chunk_param, _, lam_val_params) = - partitionChunkedFoldParameters 0 lam_params - - -- Synthesize neutral elements by applying the fold function - -- to an empty chunk. - letBindNames [I.paramName chunk_param] $ - I.BasicOp $ - I.SubExp $ - constant (0 :: Int64) - forM_ lam_val_params $ \p -> - letBindNames [I.paramName p] $ - I.BasicOp . I.Scratch (I.elemType $ I.paramType p) $ - I.arrayDims $ - I.paramType p - nes <- bodyBind =<< renameBody lam_body - - nes_ts <- mapM I.subExpResType nes - outsz <- arraysSize 0 <$> mapM lookupType arrs - let acc_arr_tps = [I.arrayOf t (I.Shape [outsz]) NoUniqueness | t <- nes_ts] - lam0' <- internaliseFoldLambda internaliseLambda lam0 nes_ts acc_arr_tps - - let lam0_acc_params = take (length nes) $ I.lambdaParams lam0' - lam_acc_params <- forM lam0_acc_params $ \p -> do - name <- newVName $ baseString $ I.paramName p - pure p {I.paramName = name} - - -- Make sure the chunk size parameter comes first. - let lam_params' = chunk_param : lam_acc_params ++ lam_val_params - - lam' <- mkLambda lam_params' $ do - lam_res <- bodyBind lam_body - lam_res' <- - ensureArgShapes - "shape of chunk function result does not match shape of initial value" - (srclocOf lam) - [] - (map I.typeOf $ I.lambdaParams lam0') - (map resSubExp lam_res) - ensureResultShape - "shape of result does not match shape of initial value" - (srclocOf lam0) - nes_ts - =<< ( eLambda lam0' . map eSubExp $ - map (I.Var . paramName) lam_acc_params ++ lam_res' - ) - - let form = I.Parallel o comm lam0' - w <- arraysSize 0 <$> mapM lookupType arrs - letValExp' desc $ I.Op $ I.Stream w arrs form (map resSubExp nes) lam' - internaliseStreamAcc :: String -> E.Exp -> @@ -1694,14 +1618,6 @@ where reduce w scan_lam nes arrs = I.Screma w arrs <$> I.scanSOAC [Scan scan_lam nes] - handleSOACs [TupLit [op, f, arr] _] "reduce_stream" = Just $ \desc -> - internaliseStreamRed desc InOrder Noncommutative op f arr - handleSOACs [TupLit [op, f, arr] _] "reduce_stream_per" = Just $ \desc -> - internaliseStreamRed desc Disorder Commutative op f arr - handleSOACs [TupLit [f, arr] _] "map_stream" = Just $ \desc -> - internaliseStreamMap desc InOrder f arr - handleSOACs [TupLit [f, arr] _] "map_stream_per" = Just $ \desc -> - internaliseStreamMap desc Disorder f arr handleSOACs [TupLit [rf, dest, op, ne, buckets, img] _] "hist_1d" = Just $ \desc -> internaliseHist 1 desc rf dest op ne buckets img loc handleSOACs [TupLit [rf, dest, op, ne, buckets, img] _] "hist_2d" = Just $ \desc -> diff --git a/src/Futhark/Internalise/Lambdas.hs b/src/Futhark/Internalise/Lambdas.hs --- a/src/Futhark/Internalise/Lambdas.hs +++ b/src/Futhark/Internalise/Lambdas.hs @@ -2,9 +2,7 @@ module Futhark.Internalise.Lambdas ( InternaliseLambda, - internaliseStreamMapLambda, internaliseFoldLambda, - internaliseStreamLambda, internalisePartitionLambda, ) where @@ -18,31 +16,6 @@ type InternaliseLambda = E.Exp -> [I.Type] -> InternaliseM ([I.LParam SOACS], I.Body SOACS, [I.Type]) -internaliseStreamMapLambda :: - InternaliseLambda -> - E.Exp -> - [I.SubExp] -> - InternaliseM (I.Lambda SOACS) -internaliseStreamMapLambda internaliseLambda lam args = do - chunk_size <- newVName "chunk_size" - let chunk_param = I.Param mempty chunk_size (I.Prim int64) - outer = (`setOuterSize` I.Var chunk_size) - localScope (scopeOfLParams [chunk_param]) $ do - argtypes <- mapM I.subExpType args - (lam_params, orig_body, rettype) <- - internaliseLambda lam $ I.Prim int64 : map outer argtypes - let orig_chunk_param : params = lam_params - body <- runBodyBuilder $ do - letBindNames [paramName orig_chunk_param] $ I.BasicOp $ I.SubExp $ I.Var chunk_size - pure orig_body - mkLambda (chunk_param : params) $ do - letBindNames [paramName orig_chunk_param] $ I.BasicOp $ I.SubExp $ I.Var chunk_size - ensureResultShape - (ErrorMsg [ErrorString "not all iterations produce same shape"]) - (srclocOf lam) - (map outer rettype) - =<< bodyBind body - internaliseFoldLambda :: InternaliseLambda -> E.Exp -> @@ -64,24 +37,6 @@ (srclocOf lam) rettype' =<< bodyBind body - -internaliseStreamLambda :: - InternaliseLambda -> - E.Exp -> - [I.Type] -> - InternaliseM ([LParam SOACS], Body SOACS) -internaliseStreamLambda internaliseLambda lam rowts = do - chunk_size <- newVName "chunk_size" - let chunk_param = I.Param mempty chunk_size $ I.Prim int64 - chunktypes = map (`arrayOfRow` I.Var chunk_size) rowts - localScope (scopeOfLParams [chunk_param]) $ do - (lam_params, orig_body, _) <- - internaliseLambda lam $ I.Prim int64 : chunktypes - let orig_chunk_param : params = lam_params - body <- runBodyBuilder $ do - letBindNames [paramName orig_chunk_param] $ I.BasicOp $ I.SubExp $ I.Var chunk_size - pure orig_body - pure (chunk_param : params, body) -- Given @k@ lambdas, this will return a lambda that returns an -- (k+2)-element tuple of integers. The first element is the diff --git a/src/Futhark/Optimise/Fusion.hs b/src/Futhark/Optimise/Fusion.hs --- a/src/Futhark/Optimise/Fusion.hs +++ b/src/Futhark/Optimise/Fusion.hs @@ -414,7 +414,7 @@ SoacNode ots pat soac aux -> do let lam = H.lambda soac lam' <- localScope (scopeOf lam) $ case soac of - H.Stream _ Sequential {} _ _ _ -> + H.Stream {} -> dontFuseScans $ doFusionLambda lam _ -> doFuseScans $ doFusionLambda lam diff --git a/src/Futhark/Optimise/Fusion/Composing.hs b/src/Futhark/Optimise/Fusion/Composing.hs --- a/src/Futhark/Optimise/Fusion/Composing.hs +++ b/src/Futhark/Optimise/Fusion/Composing.hs @@ -13,7 +13,6 @@ module Futhark.Optimise.Fusion.Composing ( fuseMaps, fuseRedomap, - mergeReduceOps, ) where @@ -279,14 +278,3 @@ ++ extra_map_ts } in (res_lam', new_inp) - -mergeReduceOps :: Lambda rep -> Lambda rep -> Lambda rep -mergeReduceOps (Lambda par1 bdy1 rtp1) (Lambda par2 bdy2 rtp2) = - let body' = - Body - (bodyDec bdy1) - (bodyStms bdy1 <> bodyStms bdy2) - (bodyResult bdy1 ++ bodyResult bdy2) - (len1, len2) = (length rtp1, length rtp2) - par' = take len1 par1 ++ take len2 par2 ++ drop len1 par1 ++ drop len2 par2 - in Lambda par' body' (rtp1 ++ rtp2) diff --git a/src/Futhark/Optimise/Fusion/GraphRep.hs b/src/Futhark/Optimise/Fusion/GraphRep.hs --- a/src/Futhark/Optimise/Fusion/GraphRep.hs +++ b/src/Futhark/Optimise/Fusion/GraphRep.hs @@ -386,8 +386,8 @@ Futhark.Screma w is form -> inputs is <> freeClassifications (w, form) Futhark.Hist w is ops lam -> inputs is <> freeClassifications (w, ops, lam) Futhark.Scatter w is lam iws -> inputs is <> freeClassifications (w, lam, iws) - Futhark.Stream w is form nes lam -> - inputs is <> freeClassifications (w, form, nes, lam) + Futhark.Stream w is nes lam -> + inputs is <> freeClassifications (w, nes, lam) Futhark.JVP {} -> freeClassifications soac Futhark.VJP {} -> freeClassifications soac where diff --git a/src/Futhark/Optimise/Fusion/TryFusion.hs b/src/Futhark/Optimise/Fusion/TryFusion.hs --- a/src/Futhark/Optimise/Fusion/TryFusion.hs +++ b/src/Futhark/Optimise/Fusion/TryFusion.hs @@ -389,12 +389,8 @@ ---------------------------- -- Stream-Stream Fusions: -- ---------------------------- - (SOAC.Stream _ Sequential _ _ _, SOAC.Stream _ Sequential _ _ _) -> do - -- fuse two SEQUENTIAL streams - (res_nms, res_stream) <- fuseStreamHelper (fsOutNames ker) unfus_set outVars outPairs soac_c soac_p - success res_nms res_stream (SOAC.Stream {}, SOAC.Stream {}) -> do - -- fuse two PARALLEL streams + -- fuse two SEQUENTIAL streams (res_nms, res_stream) <- fuseStreamHelper (fsOutNames ker) unfus_set outVars outPairs soac_c soac_p success res_nms res_stream ------------------------------------------------------------------- @@ -406,23 +402,17 @@ --- we could run in an infinite recursion, i.e., repeatedly --- --- fusing map o scan into an infinity of Stream levels! --- ------------------------------------------------------------------- - (SOAC.Stream _ form2 _ _ _, _) -> do + (SOAC.Stream {}, _) -> do -- If this rule is matched then soac_p is NOT a stream. -- To fuse a stream kernel, we transform soac_p to a stream, which -- borrows the sequential/parallel property of the soac_c Stream, -- and recursively perform stream-stream fusion. (soac_p', newacc_ids) <- SOAC.soacToStream soac_p - soac_p'' <- case form2 of - Sequential {} -> maybe (fail "not a stream") pure (toSeqStream soac_p') - _ -> pure soac_p' - if soac_p' == soac_p - then fail "SOAC could not be turned into stream." - else - fuseSOACwithKer - (namesFromList (map identName newacc_ids) <> unfus_set) - (map identName newacc_ids ++ outVars) - soac_p'' - ker + fuseSOACwithKer + (namesFromList (map identName newacc_ids) <> unfus_set) + (map identName newacc_ids ++ outVars) + soac_p' + ker (_, SOAC.Screma _ form _) | Just _ <- Futhark.isScanomapSOAC form -> do -- A Scan soac can be currently only fused as a (sequential) stream, -- hence it is first translated to a (sequential) Stream and then @@ -436,33 +426,27 @@ soac_p' ker else fail "SOAC could not be turned into stream." - (_, SOAC.Stream _ form_p _ _ _) -> do + (_, SOAC.Stream {}) -> do -- If it reached this case then soac_c is NOT a Stream kernel, -- hence transform the kernel's soac to a stream and attempt -- stream-stream fusion recursivelly. -- The newly created stream corresponding to soac_c borrows the -- sequential/parallel property of the soac_p stream. (soac_c', newacc_ids) <- SOAC.soacToStream soac_c - when (soac_c' == soac_c) $ fail "SOAC could not be turned into stream." - soac_c'' <- case form_p of - Sequential -> maybe (fail "not a stream") pure (toSeqStream soac_c') - _ -> pure soac_c' - - fuseSOACwithKer - (namesFromList (map identName newacc_ids) <> unfus_set) - outVars - soac_p - $ ker {fsSOAC = soac_c'', fsOutNames = map identName newacc_ids ++ fsOutNames ker} + if soac_c' /= soac_c + then + fuseSOACwithKer + (namesFromList (map identName newacc_ids) <> unfus_set) + outVars + soac_p + $ ker {fsSOAC = soac_c', fsOutNames = map identName newacc_ids ++ fsOutNames ker} + else fail "SOAC could not be turned into stream." --------------------------------- --- DEFAULT, CANNOT FUSE CASE --- --------------------------------- _ -> fail "Cannot fuse" -getStreamOrder :: StreamForm rep -> StreamOrd -getStreamOrder (Parallel o _ _) = o -getStreamOrder Sequential = InOrder - fuseStreamHelper :: [VName] -> Names -> @@ -476,56 +460,38 @@ unfus_set outVars outPairs - (SOAC.Stream w2 form2 lam2 nes2 inp2_arr) - (SOAC.Stream _ form1 lam1 nes1 inp1_arr) = - if getStreamOrder form2 /= getStreamOrder form1 - then fail "fusion conditions not met!" - else do - -- very similar to redomap o redomap composition, but need - -- to remove first the `chunk' parameters of streams' - -- lambdas and put them in the resulting stream lambda. - let chunk1 = head $ lambdaParams lam1 - chunk2 = head $ lambdaParams lam2 - hmnms = M.fromList [(paramName chunk2, paramName chunk1)] - lam20 = substituteNames hmnms lam2 - lam1' = lam1 {lambdaParams = tail $ lambdaParams lam1} - lam2' = lam20 {lambdaParams = tail $ lambdaParams lam20} - (res_lam', new_inp) = - fuseRedomap - unfus_set - outVars - lam1' - [] - nes1 - inp1_arr - outPairs - lam2' - [] - nes2 - inp2_arr - res_lam'' = res_lam' {lambdaParams = chunk1 : lambdaParams res_lam'} - unfus_accs = take (length nes1) outVars - unfus_arrs = filter (`notElem` unfus_accs) $ filter (`nameIn` unfus_set) outVars - let res_form = mergeForms form2 form1 - pure - ( unfus_accs ++ out_kernms ++ unfus_arrs, - SOAC.Stream w2 res_form res_lam'' (nes1 ++ nes2) new_inp - ) + (SOAC.Stream w2 lam2 nes2 inp2_arr) + (SOAC.Stream _ lam1 nes1 inp1_arr) = do + -- very similar to redomap o redomap composition, but need + -- to remove first the `chunk' parameters of streams' + -- lambdas and put them in the resulting stream lambda. + let chunk1 = head $ lambdaParams lam1 + chunk2 = head $ lambdaParams lam2 + hmnms = M.fromList [(paramName chunk2, paramName chunk1)] + lam20 = substituteNames hmnms lam2 + lam1' = lam1 {lambdaParams = tail $ lambdaParams lam1} + lam2' = lam20 {lambdaParams = tail $ lambdaParams lam20} + (res_lam', new_inp) = + fuseRedomap + unfus_set + outVars + lam1' + [] + nes1 + inp1_arr + outPairs + lam2' + [] + nes2 + inp2_arr + res_lam'' = res_lam' {lambdaParams = chunk1 : lambdaParams res_lam'} + unfus_accs = take (length nes1) outVars + unfus_arrs = filter (`notElem` unfus_accs) $ filter (`nameIn` unfus_set) outVars + pure + ( unfus_accs ++ out_kernms ++ unfus_arrs, + SOAC.Stream w2 res_lam'' (nes1 ++ nes2) new_inp + ) fuseStreamHelper _ _ _ _ _ _ = fail "Cannot Fuse Streams!" - -mergeForms :: StreamForm SOACS -> StreamForm SOACS -> StreamForm SOACS -mergeForms Sequential Sequential = Sequential -mergeForms (Parallel _ comm2 lam2r) (Parallel o1 comm1 lam1r) = - Parallel o1 (comm1 <> comm2) (mergeReduceOps lam1r lam2r) -mergeForms _ _ = error "Fusing sequential to parallel stream disallowed!" - --- | If a Stream is passed as argument then it converts it to a --- Sequential Stream. -toSeqStream :: SOAC -> Maybe SOAC -toSeqStream s@(SOAC.Stream _ Sequential _ _ _) = Just s -toSeqStream (SOAC.Stream w Parallel {} l acc inps) = - Just $ SOAC.Stream w Sequential l acc inps -toSeqStream _ = Nothing -- Here follows optimizations and transforms to expose fusability. diff --git a/src/Futhark/Optimise/ReduceDeviceSyncs.hs b/src/Futhark/Optimise/ReduceDeviceSyncs.hs --- a/src/Futhark/Optimise/ReduceDeviceSyncs.hs +++ b/src/Futhark/Optimise/ReduceDeviceSyncs.hs @@ -15,7 +15,7 @@ import Data.Bifunctor (second) import Data.Foldable import qualified Data.IntMap.Strict as IM -import Data.List (transpose) +import Data.List (transpose, zip4) import qualified Data.Map.Strict as M import Data.Sequence ((<|), (><), (|>)) import qualified Data.Text as T @@ -170,8 +170,9 @@ else do -- Otherwise, ensure all results are migrated. (all_stms', arrs) <- - fmap unzip $ forM (zip all_stms reses) $ \(stms, res) -> - storeScalar stms (resSubExp res) (patElemType pe) + fmap unzip $ + forM (zip all_stms reses) $ \(stms, res) -> + storeScalar stms (resSubExp res) (patElemType pe) pe' <- arrayizePatElem pe let bt' = staticShapes1 (patElemType pe') @@ -202,40 +203,55 @@ -- Update statement bound variables and parameters if their values -- have been migrated to device. - let lmerge (res, stms) (pe, (Param _ pn pt, pval), MoveToDevice) = do - -- Rewrite the bound variable. + let lmerge (res, stms, rebinds) (pe, param, StayOnHost) = + pure ((pe, param) : res, stms, rebinds) + lmerge (res, stms, rebinds) (pe, (Param _ pn pt, pval), _) = do + -- Migrate the bound variable. pe' <- arrayizePatElem pe - -- Move the initial value to device if not already there. + -- Move the initial value to device if not already there to + -- ensure that the parameter argument and loop return value + -- converge. (stms', arr) <- storeScalar stms pval (fromDecl pt) - -- Rewrite the parameter. + -- Migrate the parameter. pn' <- newName pn let pt' = toDecl (patElemType pe') Nonunique let pval' = Var arr let param' = (Param mempty pn' pt', pval') - -- Record the migration. - Ident pn (fromDecl pt) `movedTo` pn' + -- Record the migration and rebind the parameter inside the + -- loop body if necessary. + rebinds' <- (pe {patElemName = pn}) `migratedTo` (pn', rebinds) - pure ((pe', param') : res, stms') - lmerge _ (_, _, UsedOnHost) = - -- Initial loop parameter value and loop result should have - -- been made available on host instead. - compilerBugS "optimizeStm: unhandled host usage of loop param" - lmerge (res, stms) (pe, param, StayOnHost) = - pure ((pe, param) : res, stms) + pure ((pe', param') : res, stms', rebinds') mt <- ask - let pes = patElems (stmPat stm) let mss = map (\(Param _ n _, _) -> statusOf n mt) params - (zipped', out') <- foldM lmerge ([], out) (zip3 pes params mss) + (zipped', out', rebinds) <- + foldM lmerge ([], out, mempty) (zip3 pes params mss) let (pes', params') = unzip (reverse zipped') + -- Rewrite body. + let body1 = body {bodyStms = rebinds >< bodyStms body} + body2 <- optimizeBody body1 + let zipped = + zip4 + mss + (bodyResult body2) + (map resSubExp $ bodyResult body) + (map patElemType pes) + let rstore (bstms, res) (StayOnHost, r, _, _) = + pure (bstms, r : res) + rstore (bstms, res) (_, SubExpRes certs _, se, t) = do + (bstms', dev) <- storeScalar bstms se t + pure (bstms', SubExpRes certs (Var dev) : res) + (bstms, res) <- foldM rstore (bodyStms body2, []) zipped + let body3 = body2 {bodyStms = bstms, bodyResult = reverse res} + -- Rewrite statement. - body' <- optimizeBody body - let e' = DoLoop params' lform body' + let e' = DoLoop params' lform body3 let stm' = Let (Pat pes') (stmAux stm) e' -- Read migrated scalars that are used on host. diff --git a/src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs b/src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs --- a/src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs +++ b/src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs @@ -443,9 +443,10 @@ BasicOp (Index arr s) -> do graphInefficientReturn (sliceDims s) e one bs `reuses` arr - BasicOp (Update _ arr _ _) -> do - graphInefficientReturn [] e - one bs `reuses` arr + BasicOp (Update _ arr slice _) + | isFixed slice -> do + graphInefficientReturn [] e + one bs `reuses` arr BasicOp (FlatIndex arr s) -> do -- Migrating a FlatIndex leads to a memory allocation error. -- @@ -494,6 +495,8 @@ -- Whether the rows are primitive constants or arrays, without any scalar -- variable operands such ArrayLit cannot directly prevent a scalar read. graphHostOnly e + BasicOp Update {} -> + graphHostOnly e BasicOp Concat {} -> -- Is unlikely to prevent a scalar read as the only SubExp operand in -- practice is a computation of host-only size variables. @@ -692,7 +695,7 @@ compilerBugS "Loop statement bound no variable; should have been eliminated." graphLoop (b : bs) params lform body = do -- Graph loop params and body while capturing statistics. - g0 <- getGraph + g <- getGraph stats <- captureBodyStats (subgraphId `graphIdFor` graphTheLoop) -- Record aliases for copyable memory backing returned arrays. @@ -701,18 +704,23 @@ let results = map resSubExp (bodyResult body) may_copy_results <- reusesBranches (b : bs) [args, results] - -- Connect loop condition to a sink if the loop cannot be migrated. - -- The migration status of the condition is what determines whether the - -- loop may be migrated as a whole or not. See 'shouldMoveStm'. + -- Connect the loop condition to a sink if the loop cannot be migrated, + -- ensuring that it will be available to the host. The migration status + -- of the condition is what determines whether the loop may be migrated + -- as a whole or not. See 'shouldMoveStm'. let may_migrate = not (bodyHostOnly stats) && may_copy_results unless may_migrate $ case lform of ForLoop _ _ (Var n) _ -> connectToSink (nameToId n) - WhileLoop n -> connectToSink (nameToId n) + WhileLoop n + | (_, p, _, res) <- loopValueFor n -> do + connectToSink p + case res of + Var v -> connectToSink (nameToId v) + _ -> pure () _ -> pure () -- Connect graphed return values to their loop parameters. - g1 <- getGraph - mapM_ (mergeLoopParam g1) loopValues + mapM_ mergeLoopParam loopValues -- Route the sources within the loop body in isolation. -- The loop graph must not be altered after this point. @@ -724,8 +732,8 @@ -- If a device read is delayed from one iteration to the next the -- corresponding variables bound by the statement must be treated as -- sources. - g2 <- getGraph - let (dbs, rbc) = foldl' (deviceBindings g2) (IS.empty, MG.none) srcs + g' <- getGraph + let (dbs, rbc) = foldl' (deviceBindings g') (IS.empty, MG.none) srcs modifySources $ second (IS.toList dbs <>) -- Connect operands to sinks if they can reach a sink within the loop. @@ -733,7 +741,7 @@ -- reach and exhaust their normal entry edges into the loop. -- This means a read can be delayed through a loop but not into it if -- that would increase the number of reads done by any given iteration. - let ops = IS.filter (`MG.member` g0) (bodyOperands stats) + let ops = IS.filter (`MG.member` g) (bodyOperands stats) foldM_ connectOperand rbc (IS.elems ops) -- It might be beneficial to move the whole loop to device, to avoid @@ -807,14 +815,12 @@ ops <- onlyGraphedScalarSubExp arg addEdges (MG.oneEdge p) ops - mergeLoopParam :: Graph -> LoopValue -> Grapher () - mergeLoopParam g (_, p, _, res) + mergeLoopParam :: LoopValue -> Grapher () + mergeLoopParam (_, p, _, res) | Var n <- res, ret <- nameToId n, ret /= p = - if MG.isSinkConnected p g - then connectToSink ret - else addEdges (MG.oneEdge p) (IS.singleton ret) + addEdges (MG.oneEdge p) (IS.singleton ret) | otherwise = pure () deviceBindings :: diff --git a/src/Futhark/Pass/ExplicitAllocations.hs b/src/Futhark/Pass/ExplicitAllocations.hs --- a/src/Futhark/Pass/ExplicitAllocations.hs +++ b/src/Futhark/Pass/ExplicitAllocations.hs @@ -7,6 +7,7 @@ {-# LANGUAGE TupleSections #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UndecidableInstances #-} +{-# OPTIONS_GHC -fno-warn-redundant-constraints #-} -- | A generic transformation for adding memory allocations to a -- Futhark program. Specialised by specific representations in @@ -26,8 +27,6 @@ arraySizeInBytesExp, mkLetNamesB', mkLetNamesB'', - dimAllocationSize, - ChunkMap, -- * Module re-exports @@ -53,9 +52,11 @@ import qualified Data.Map.Strict as M import Data.Maybe import qualified Data.Set as S +import Futhark.Analysis.SymbolTable (IndexOp) import qualified Futhark.Analysis.UsageTable as UT import Futhark.IR.Mem import qualified Futhark.IR.Mem.IxFun as IxFun +import Futhark.IR.Prop.Aliases (AliasedOp) import Futhark.MonadFreshNames import Futhark.Optimise.Simplify.Engine (SimpleOps (..)) import qualified Futhark.Optimise.Simplify.Engine as Engine @@ -64,16 +65,6 @@ import Futhark.Tools import Futhark.Util (maybeNth, splitAt3) --- | The subexpression giving the number of elements we should --- allocate space for. See 'ChunkMap' comment. -dimAllocationSize :: ChunkMap -> SubExp -> SubExp -dimAllocationSize chunkmap (Var v) = - -- It is important to recurse here, as the substitution may itself - -- be a chunk size. - maybe (Var v) (dimAllocationSize chunkmap) $ M.lookup v chunkmap -dimAllocationSize _ size = - size - type Allocable fromrep torep inner = ( PrettyRep fromrep, PrettyRep torep, @@ -90,15 +81,8 @@ BuilderOps torep ) --- | A mapping from chunk names to their maximum size. XXX FIXME --- HACK: This is part of a hack to add loop-invariant allocations to --- reduce kernels, because memory expansion does not use range --- analysis yet (it should). -type ChunkMap = M.Map VName SubExp - data AllocEnv fromrep torep = AllocEnv - { chunkMap :: ChunkMap, - -- | Aggressively try to reuse memory in do-loops - + { -- | Aggressively try to reuse memory in do-loops - -- should be True inside kernels, False outside. aggressiveReuse :: Bool, -- | When allocating memory, put it in this memory space. @@ -132,9 +116,8 @@ mkLetNamesM names e = do def_space <- askDefaultSpace - chunkmap <- asks chunkMap hints <- expHints e - pat <- patWithAllocations def_space chunkmap names e hints + pat <- patWithAllocations def_space names e hints pure $ Let pat (defAux ()) e mkBodyM stms res = pure $ Body () stms res @@ -161,8 +144,7 @@ where env = AllocEnv - { chunkMap = mempty, - aggressiveReuse = False, + { aggressiveReuse = False, allocSpace = DefaultSpace, envConsts = mempty, allocInOp = handleOp, @@ -176,26 +158,25 @@ arraySizeInBytesExp t = untyped $ foldl' (*) (elemSize t) $ map pe64 (arrayDims t) -arraySizeInBytesExpM :: MonadBuilder m => ChunkMap -> Type -> m (PrimExp VName) -arraySizeInBytesExpM chunkmap t = do - let dim_prod_i64 = product $ map (pe64 . dimAllocationSize chunkmap) (arrayDims t) +arraySizeInBytesExpM :: MonadBuilder m => Type -> m (PrimExp VName) +arraySizeInBytesExpM t = do + let dim_prod_i64 = product $ map pe64 (arrayDims t) elm_size_i64 = elemSize t pure $ BinOpExp (SMax Int64) (ValueExp $ IntValue $ Int64Value 0) $ untyped $ dim_prod_i64 * elm_size_i64 -arraySizeInBytes :: MonadBuilder m => ChunkMap -> Type -> m SubExp -arraySizeInBytes chunkmap = letSubExp "bytes" <=< toExp <=< arraySizeInBytesExpM chunkmap +arraySizeInBytes :: MonadBuilder m => Type -> m SubExp +arraySizeInBytes = letSubExp "bytes" <=< toExp <=< arraySizeInBytesExpM allocForArray' :: (MonadBuilder m, Op (Rep m) ~ MemOp inner) => - ChunkMap -> Type -> Space -> m VName -allocForArray' chunkmap t space = do - size <- arraySizeInBytes chunkmap t +allocForArray' t space = do + size <- arraySizeInBytes t letExp "mem" $ Op $ Alloc size space -- | Allocate memory for a value of the given type. @@ -205,8 +186,7 @@ Space -> AllocM fromrep torep VName allocForArray t space = do - chunkmap <- asks chunkMap - allocForArray' chunkmap t space + allocForArray' t space allocsForStm :: (Allocable fromrep torep inner) => @@ -215,26 +195,24 @@ AllocM fromrep torep (Stm torep) allocsForStm idents e = do def_space <- askDefaultSpace - chunkmap <- asks chunkMap hints <- expHints e rts <- expReturns e - pes <- allocsForPat def_space chunkmap idents rts hints + pes <- allocsForPat def_space idents rts hints dec <- mkExpDecM (Pat pes) e pure $ Let (Pat pes) (defAux dec) e patWithAllocations :: (MonadBuilder m, Mem (Rep m) inner) => Space -> - ChunkMap -> [VName] -> Exp (Rep m) -> [ExpHint] -> m (Pat LetDecMem) -patWithAllocations def_space chunkmap names e hints = do +patWithAllocations def_space names e hints = do ts' <- instantiateShapes' names <$> expExtType e let idents = zipWith Ident names ts' rts <- expReturns e - Pat <$> allocsForPat def_space chunkmap idents rts hints + Pat <$> allocsForPat def_space idents rts hints mkMissingIdents :: MonadFreshNames m => [Ident] -> [ExpReturns] -> m [Ident] mkMissingIdents idents rts = @@ -247,19 +225,18 @@ allocsForPat :: (MonadBuilder m, Op (Rep m) ~ MemOp inner) => Space -> - ChunkMap -> [Ident] -> [ExpReturns] -> [ExpHint] -> m [PatElem LetDecMem] -allocsForPat def_space chunkmap some_idents rts hints = do +allocsForPat def_space some_idents rts hints = do idents <- mkMissingIdents some_idents rts forM (zip3 idents rts hints) $ \(ident, rt, hint) -> do let ident_shape = arrayShape $ identType ident case rt of MemPrim _ -> do - summary <- summaryForBindage def_space chunkmap (identType ident) hint + summary <- summaryForBindage def_space (identType ident) hint pure $ PatElem (identName ident) summary MemMem space -> pure $ PatElem (identName ident) $ MemMem space @@ -268,7 +245,7 @@ pure . PatElem (identName ident) . MemArray bt ident_shape u $ ArrayIn mem ixfn MemArray _ extshape _ Nothing | Just _ <- knownShape extshape -> do - summary <- summaryForBindage def_space chunkmap (identType ident) hint + summary <- summaryForBindage def_space (identType ident) hint pure $ PatElem (identName ident) summary MemArray bt _ u (Just (ReturnsNewBlock _ i extixfn)) -> do let ixfn = instantiateExtIxFun idents extixfn @@ -302,20 +279,19 @@ summaryForBindage :: (MonadBuilder m, Op (Rep m) ~ MemOp inner) => Space -> - ChunkMap -> Type -> ExpHint -> m (MemBound NoUniqueness) -summaryForBindage _ _ (Prim bt) _ = +summaryForBindage _ (Prim bt) _ = pure $ MemPrim bt -summaryForBindage _ _ (Mem space) _ = +summaryForBindage _ (Mem space) _ = pure $ MemMem space -summaryForBindage _ _ (Acc acc ispace ts u) _ = +summaryForBindage _ (Acc acc ispace ts u) _ = pure $ MemAcc acc ispace ts u -summaryForBindage def_space chunkmap t@(Array pt shape u) NoHint = do - m <- allocForArray' chunkmap t def_space +summaryForBindage def_space t@(Array pt shape u) NoHint = do + m <- allocForArray' t def_space pure $ MemArray pt shape u $ ArrayIn m $ IxFun.iota $ map pe64 $ arrayDims t -summaryForBindage _ _ t@(Array pt _ _) (Hint ixfun space) = do +summaryForBindage _ t@(Array pt _ _) (Hint ixfun space) = do bytes <- letSubExp "bytes" <=< toExp . untyped $ product @@ -424,11 +400,10 @@ -- _must_ be in ScalarSpace and have the right index function. (res_mem, res_ixfun) <- lift $ lookupArraySummary res res_mem_space <- lift $ lookupMemSpace res_mem - chunkmap <- asks chunkMap (res_mem', res') <- if (res_mem_space, res_ixfun) == (v_mem_space, v_ixfun) then pure (res_mem, res) - else lift $ arrayWithIxFun chunkmap v_mem_space v_ixfun (fromDecl param_t) res + else lift $ arrayWithIxFun v_mem_space v_ixfun (fromDecl param_t) res tell ([Var res_mem'], []) pure $ Var res' scalarRes _ _ _ se = pure se @@ -500,15 +475,14 @@ arrayWithIxFun :: (MonadBuilder m, Op (Rep m) ~ MemOp inner, LetDec (Rep m) ~ LetDecMem) => - ChunkMap -> Space -> IxFun -> Type -> VName -> m (VName, VName) -arrayWithIxFun chunkmap space ixfun v_t v = do +arrayWithIxFun space ixfun v_t v = do let Array pt shape u = v_t - mem <- allocForArray' chunkmap v_t space + mem <- allocForArray' v_t space v_copy <- newVName $ baseString v <> "_scalcopy" letBind (Pat [PatElem v_copy $ MemArray pt shape u $ ArrayIn mem ixfun]) $ BasicOp $ Copy v pure (mem, v_copy) @@ -705,13 +679,8 @@ allocInStms' (stm : stms) = do allocstms <- collectStms_ $ auxing (stmAux stm) $ allocInStm stm addStms allocstms - let stms_substs = foldMap sizeSubst allocstms - stms_consts = foldMap stmConsts allocstms - f env = - env - { chunkMap = stms_substs <> chunkMap env, - envConsts = stms_consts <> envConsts env - } + let stms_consts = foldMap stmConsts allocstms + f env = env {envConsts = stms_consts <> envConsts env} local f $ allocInStms' stms allocInStm :: @@ -1035,24 +1004,15 @@ pure (p {paramDec = MemAcc acc ispace ts u}, a) class SizeSubst op where - opSizeSubst :: Pat dec -> op -> ChunkMap opIsConst :: op -> Bool opIsConst = const False -instance SizeSubst () where - opSizeSubst _ _ = mempty +instance SizeSubst () instance SizeSubst op => SizeSubst (MemOp op) where - opSizeSubst pat (Inner op) = opSizeSubst pat op - opSizeSubst _ _ = mempty - opIsConst (Inner op) = opIsConst op opIsConst _ = False -sizeSubst :: SizeSubst (Op rep) => Stm rep -> ChunkMap -sizeSubst (Let pat _ (Op op)) = opSizeSubst pat op -sizeSubst _ = mempty - stmConsts :: SizeSubst (Op rep) => Stm rep -> S.Set VName stmConsts (Let pat _ (Op op)) | opIsConst op = S.fromList $ patNames pat @@ -1069,7 +1029,7 @@ Exp (Rep m) -> m (Stm (Rep m)) mkLetNamesB' dec names e = do - pat <- patWithAllocations DefaultSpace mempty names e nohints + pat <- patWithAllocations DefaultSpace names e nohints pure $ Let pat (defAux dec) e where nohints = map (const NoHint) names @@ -1088,7 +1048,7 @@ Exp (Engine.Wise rep) -> m (Stm (Engine.Wise rep)) mkLetNamesB'' names e = do - pat <- patWithAllocations DefaultSpace mempty names e nohints + pat <- patWithAllocations DefaultSpace names e nohints let pat' = Engine.addWisdomToPat pat e dec = Engine.mkWiseExpDec pat' () e pure $ Let pat' (defAux dec) e @@ -1097,10 +1057,12 @@ simplifiable :: ( Engine.SimplifiableRep rep, + LetDec rep ~ LetDecMem, ExpDec rep ~ (), BodyDec rep ~ (), - LetDec rep ~ LetDecMem, OpReturns (Engine.OpWithWisdom inner), + AliasedOp (Engine.OpWithWisdom inner), + IndexOp (Engine.OpWithWisdom inner), Mem rep inner ) => (Engine.OpWithWisdom inner -> UT.UsageTable) -> diff --git a/src/Futhark/Pass/ExplicitAllocations/GPU.hs b/src/Futhark/Pass/ExplicitAllocations/GPU.hs --- a/src/Futhark/Pass/ExplicitAllocations/GPU.hs +++ b/src/Futhark/Pass/ExplicitAllocations/GPU.hs @@ -10,7 +10,6 @@ ) where -import qualified Data.Map as M import qualified Data.Set as S import Futhark.IR.GPU import Futhark.IR.GPUMem @@ -19,10 +18,6 @@ import Futhark.Pass.ExplicitAllocations.SegOp instance SizeSubst (HostOp rep op) where - opSizeSubst (Pat [size]) (SizeOp (SplitSpace _ _ _ elems_per_thread)) = - M.singleton (patElemName size) elems_per_thread - opSizeSubst _ _ = mempty - opIsConst (SizeOp GetSize {}) = True opIsConst (SizeOp GetSizeMax {}) = True opIsConst _ = False @@ -102,11 +97,8 @@ Type -> KernelResult -> AllocM GPU GPUMem ExpHint -mapResultHint lvl space = hint +mapResultHint _lvl space = hint where - num_threads = - pe64 (unCount $ segNumGroups lvl) * pe64 (unCount $ segGroupSize lvl) - -- Heuristic: do not rearrange for returned arrays that are -- sufficiently small. coalesceReturnOfShape _ [] = False @@ -115,19 +107,8 @@ hint t Returns {} | coalesceReturnOfShape (primByteSize (elemType t)) $ arrayDims t = do - chunkmap <- asks chunkMap let space_dims = segSpaceDims space - t_dims = map (dimAllocationSize chunkmap) $ arrayDims t - pure $ Hint (innermost space_dims t_dims) DefaultSpace - hint t (ConcatReturns _ SplitStrided {} w _ _) = do - chunkmap <- asks chunkMap - let t_dims = map (dimAllocationSize chunkmap) $ arrayDims t - pure $ Hint (innermost [w] t_dims) DefaultSpace - hint Prim {} (ConcatReturns _ SplitContiguous w elems_per_thread _) = do - let ixfun_base = IxFun.iota [sExt64 num_threads, pe64 elems_per_thread] - ixfun_tr = IxFun.permute ixfun_base [1, 0] - ixfun = IxFun.reshape ixfun_tr [pe64 w] - pure $ Hint ixfun DefaultSpace + pure $ Hint (innermost space_dims (arrayDims t)) DefaultSpace hint _ _ = pure NoHint innermost :: [SubExp] -> [SubExp] -> IxFun diff --git a/src/Futhark/Pass/ExplicitAllocations/MC.hs b/src/Futhark/Pass/ExplicitAllocations/MC.hs --- a/src/Futhark/Pass/ExplicitAllocations/MC.hs +++ b/src/Futhark/Pass/ExplicitAllocations/MC.hs @@ -11,8 +11,7 @@ import Futhark.Pass.ExplicitAllocations import Futhark.Pass.ExplicitAllocations.SegOp -instance SizeSubst (MCOp rep op) where - opSizeSubst _ _ = mempty +instance SizeSubst (MCOp rep op) handleSegOp :: SegOp () MC -> AllocM MC MCMem (SegOp () MCMem) handleSegOp op = do diff --git a/src/Futhark/Pass/ExplicitAllocations/SegOp.hs b/src/Futhark/Pass/ExplicitAllocations/SegOp.hs --- a/src/Futhark/Pass/ExplicitAllocations/SegOp.hs +++ b/src/Futhark/Pass/ExplicitAllocations/SegOp.hs @@ -13,8 +13,7 @@ import qualified Futhark.IR.Mem.IxFun as IxFun import Futhark.Pass.ExplicitAllocations -instance SizeSubst (SegOp lvl rep) where - opSizeSubst _ _ = mempty +instance SizeSubst (SegOp lvl rep) allocInKernelBody :: Allocable fromrep torep inner => diff --git a/src/Futhark/Pass/ExtractKernels.hs b/src/Futhark/Pass/ExtractKernels.hs --- a/src/Futhark/Pass/ExtractKernels.hs +++ b/src/Futhark/Pass/ExtractKernels.hs @@ -280,7 +280,7 @@ bodyStms body -- XXX - our notion of balancing is probably still too naive. - unbalancedStm bound (Op (Stream w _ _ _ _)) = + unbalancedStm bound (Op (Stream w _ _ _)) = w `subExpBound` bound unbalancedStm bound (Op (Screma w _ _)) = w `subExpBound` bound @@ -449,96 +449,13 @@ inner_stms <- innerParallelBody ((outer_suff_key, False) : path) (suff_stms <>) <$> kernelAlternatives pat inner_stms [(outer_suff, outer_stms)] - --- Streams can be handled in two different ways - either we --- sequentialise the body or we keep it parallel and distribute. -transformStm path (Let pat aux@(StmAux cs _ _) (Op (Stream w arrs Parallel {} [] map_fun))) - | not ("sequential_inner" `inAttrs` stmAuxAttrs aux) = do - -- No reduction part. Remove the stream and leave the body - -- parallel. It will be distributed. - types <- asksScope scopeForSOACs - transformStms path . stmsToList . snd - =<< runBuilderT (certifying cs $ sequentialStreamWholeArray pat w [] map_fun arrs) types -transformStm path (Let pat aux@(StmAux cs _ _) (Op (Stream w arrs (Parallel o comm red_fun) nes fold_fun))) - | "sequential_inner" `inAttrs` stmAuxAttrs aux = - paralleliseOuter path - | otherwise = do - ((outer_suff, outer_suff_key), suff_stms) <- - sufficientParallelism "suff_outer_stream" [w] path Nothing - - outer_stms <- outerParallelBody ((outer_suff_key, True) : path) - inner_stms <- innerParallelBody ((outer_suff_key, False) : path) - - (suff_stms <>) - <$> kernelAlternatives pat inner_stms [(outer_suff, outer_stms)] - where - paralleliseOuter path' - | not $ all primType $ lambdaReturnType red_fun = do - -- Split into a chunked map and a reduction, with the latter - -- further transformed. - let fold_fun' = soacsLambdaToGPU fold_fun - - let (red_pat_elems, concat_pat_elems) = - splitAt (length nes) $ patElems pat - red_pat = Pat red_pat_elems - - ((num_threads, red_results), stms) <- - streamMap - segThreadCapped - (map (baseString . patElemName) red_pat_elems) - concat_pat_elems - w - Noncommutative - fold_fun' - nes - arrs - - reduce_soac <- reduceSOAC [Reduce comm' red_fun nes] - - (stms <>) - <$> inScopeOf - stms - ( transformStm path' $ - Let red_pat aux {stmAuxAttrs = mempty} $ - Op (Screma num_threads red_results reduce_soac) - ) - | otherwise = do - let red_fun_sequential = soacsLambdaToGPU red_fun - fold_fun_sequential = soacsLambdaToGPU fold_fun - fmap (certify cs) - <$> streamRed - segThreadCapped - pat - w - comm' - red_fun_sequential - fold_fun_sequential - nes - arrs - - outerParallelBody path' = - renameBody - =<< (mkBody <$> paralleliseOuter path' <*> pure (varsRes (patNames pat))) - - paralleliseInner path' = do - types <- asksScope scopeForSOACs - transformStms path' . fmap (certify cs) . stmsToList . snd - =<< runBuilderT (sequentialStreamWholeArray pat w nes fold_fun arrs) types - - innerParallelBody path' = - renameBody - =<< (mkBody <$> paralleliseInner path' <*> pure (varsRes (patNames pat))) - - comm' - | commutativeLambda red_fun, o /= InOrder = Commutative - | otherwise = comm transformStm path (Let pat (StmAux cs _ _) (Op (Screma w arrs form))) = do -- This screma is too complicated for us to immediately do -- anything, so split it up and try again. scope <- asksScope scopeForSOACs transformStms path . map (certify cs) . stmsToList . snd =<< runBuilderT (dissectScrema pat w form arrs) scope -transformStm path (Let pat _ (Op (Stream w arrs Sequential nes fold_fun))) = do +transformStm path (Let pat _ (Op (Stream w arrs nes fold_fun))) = do -- Remove the stream and leave the body parallel. It will be -- distributed. types <- asksScope scopeForSOACs @@ -618,7 +535,7 @@ (map (bodyInterest . caseBody) cases) | Op (Screma w _ (ScremaForm _ _ lam')) <- stmExp stm = zeroIfTooSmall w + bodyInterest (lambdaBody lam') - | Op (Stream _ _ Sequential _ lam') <- stmExp stm = + | Op (Stream _ _ _ lam') <- stmExp stm = bodyInterest $ lambdaBody lam' | otherwise = 0 @@ -755,9 +672,11 @@ types <- askScope + let only_intra = onlyExploitIntra (stmAuxAttrs aux) + may_intra = worthIntraGroup lam && mayExploitIntra attrs + intra <- - if onlyExploitIntra (stmAuxAttrs aux) - || (worthIntraGroup lam && mayExploitIntra attrs) + if only_intra || may_intra then flip runReaderT types $ intraGroupParallelise loopnest lam else pure Nothing @@ -767,7 +686,8 @@ kernelAlternatives pat seq_body [] -- Nothing - | Just m <- mkSeqAlts -> do + | not only_intra, + Just m <- mkSeqAlts -> do (outer_suff, outer_suff_key, outer_suff_stms, seq_body) <- m par_body <- renameBody @@ -781,7 +701,7 @@ kernelAlternatives pat par_body [] -- Just intra'@(_, _, log, intra_prelude, intra_stms) - | onlyExploitIntra attrs -> do + | only_intra -> do addLog log group_par_body <- renameBody $ mkBody intra_stms res (intra_prelude <>) <$> kernelAlternatives pat group_par_body [] diff --git a/src/Futhark/Pass/ExtractKernels/DistributeNests.hs b/src/Futhark/Pass/ExtractKernels/DistributeNests.hs --- a/src/Futhark/Pass/ExtractKernels/DistributeNests.hs +++ b/src/Futhark/Pass/ExtractKernels/DistributeNests.hs @@ -339,7 +339,7 @@ distributeMapBodyStms orig_acc = distribute <=< onStms orig_acc . stmsToList where onStms acc [] = pure acc - onStms acc (Let pat (StmAux cs _ _) (Op (Stream w arrs Sequential accs lam)) : stms) = do + onStms acc (Let pat (StmAux cs _ _) (Op (Stream w arrs accs lam)) : stms) = do types <- asksScope scopeForSOACs stream_stms <- snd <$> runBuilderT (sequentialStreamWholeArray pat w accs lam arrs) types diff --git a/src/Futhark/Pass/ExtractKernels/Interchange.hs b/src/Futhark/Pass/ExtractKernels/Interchange.hs --- a/src/Futhark/Pass/ExtractKernels/Interchange.hs +++ b/src/Futhark/Pass/ExtractKernels/Interchange.hs @@ -46,7 +46,7 @@ Let pat (defAux ()) $ DoLoop merge form body interchangeLoop :: - (MonadBuilder m, LocalScope SOACS m) => + (MonadBuilder m, Rep m ~ SOACS) => (VName -> Maybe VName) -> SeqLoop -> LoopNesting -> diff --git a/src/Futhark/Pass/ExtractKernels/Intragroup.hs b/src/Futhark/Pass/ExtractKernels/Intragroup.hs --- a/src/Futhark/Pass/ExtractKernels/Intragroup.hs +++ b/src/Futhark/Pass/ExtractKernels/Intragroup.hs @@ -275,7 +275,7 @@ certifying (stmAuxCerts aux) $ addStms =<< segHist lvl' pat w [] [] ops' bucket_fun' arrs parallelMin [w] - Op (Stream w arrs Sequential accs lam) + Op (Stream w arrs accs lam) | chunk_size_param : _ <- lambdaParams lam -> do types <- asksScope castScope ((), stream_stms) <- diff --git a/src/Futhark/Pass/ExtractKernels/StreamKernel.hs b/src/Futhark/Pass/ExtractKernels/StreamKernel.hs --- a/src/Futhark/Pass/ExtractKernels/StreamKernel.hs +++ b/src/Futhark/Pass/ExtractKernels/StreamKernel.hs @@ -4,13 +4,10 @@ module Futhark.Pass.ExtractKernels.StreamKernel ( segThreadCapped, - streamRed, - streamMap, ) where import Control.Monad -import Control.Monad.Writer import Data.List () import Futhark.Analysis.PrimExp import Futhark.IR @@ -60,233 +57,6 @@ BasicOp $ BinOp (Mul Int64 OverflowUndef) num_groups group_size pure (num_groups, num_threads) - -blockedKernelSize :: - (MonadBuilder m, Rep m ~ GPU) => - String -> - SubExp -> - m KernelSize -blockedKernelSize desc w = do - group_size <- getSize (desc ++ "_group_size") SizeGroup - - (_, num_threads) <- numberOfGroups desc w group_size - - per_thread_elements <- - letSubExp "per_thread_elements" - =<< eBinOp (SDivUp Int64 Unsafe) (eSubExp w) (eSubExp num_threads) - - pure $ KernelSize per_thread_elements num_threads - -splitArrays :: - (MonadBuilder m, Rep m ~ GPU) => - VName -> - [VName] -> - SplitOrdering -> - SubExp -> - SubExp -> - SubExp -> - [VName] -> - m () -splitArrays chunk_size split_bound ordering w i elems_per_i arrs = do - letBindNames [chunk_size] $ Op $ SizeOp $ SplitSpace ordering w i elems_per_i - case ordering of - SplitContiguous -> do - offset <- letSubExp "slice_offset" $ BasicOp $ BinOp (Mul Int64 OverflowUndef) i elems_per_i - zipWithM_ (contiguousSlice offset) split_bound arrs - SplitStrided stride -> zipWithM_ (stridedSlice stride) split_bound arrs - where - contiguousSlice offset slice_name arr = do - arr_t <- lookupType arr - let slice = fullSlice arr_t [DimSlice offset (Var chunk_size) (constant (1 :: Int64))] - letBindNames [slice_name] $ BasicOp $ Index arr slice - - stridedSlice stride slice_name arr = do - arr_t <- lookupType arr - let slice = fullSlice arr_t [DimSlice i (Var chunk_size) stride] - letBindNames [slice_name] $ BasicOp $ Index arr slice - -partitionChunkedKernelFoldParameters :: - Int -> - [Param dec] -> - (VName, Param dec, [Param dec], [Param dec]) -partitionChunkedKernelFoldParameters num_accs (i_param : chunk_param : params) = - let (acc_params, arr_params) = splitAt num_accs params - in (paramName i_param, chunk_param, acc_params, arr_params) -partitionChunkedKernelFoldParameters _ _ = - error "partitionChunkedKernelFoldParameters: lambda takes too few parameters" - -blockedPerThread :: - (MonadBuilder m, Rep m ~ GPU) => - VName -> - SubExp -> - KernelSize -> - StreamOrd -> - Lambda (Rep m) -> - Int -> - [VName] -> - m ([PatElem Type], [PatElem Type]) -blockedPerThread thread_gtid w kernel_size ordering lam num_nonconcat arrs = do - let (_, chunk_size, [], arr_params) = - partitionChunkedKernelFoldParameters 0 $ lambdaParams lam - - ordering' = - case ordering of - InOrder -> SplitContiguous - Disorder -> SplitStrided $ kernelNumThreads kernel_size - red_ts = take num_nonconcat $ lambdaReturnType lam - map_ts = map rowType $ drop num_nonconcat $ lambdaReturnType lam - - per_thread <- asIntS Int64 $ kernelElementsPerThread kernel_size - splitArrays - (paramName chunk_size) - (map paramName arr_params) - ordering' - w - (Var thread_gtid) - per_thread - arrs - - chunk_red_pes <- forM red_ts $ \red_t -> do - pe_name <- newVName "chunk_fold_red" - pure $ PatElem pe_name red_t - chunk_map_pes <- forM map_ts $ \map_t -> do - pe_name <- newVName "chunk_fold_map" - pure $ PatElem pe_name $ map_t `arrayOfRow` Var (paramName chunk_size) - - let (chunk_red_ses, chunk_map_ses) = - splitAt num_nonconcat $ bodyResult $ lambdaBody lam - - addStms $ - bodyStms (lambdaBody lam) - <> stmsFromList - [ certify cs $ Let (Pat [pe]) (defAux ()) $ BasicOp $ SubExp se - | (pe, SubExpRes cs se) <- zip chunk_red_pes chunk_red_ses - ] - <> stmsFromList - [ certify cs $ Let (Pat [pe]) (defAux ()) $ BasicOp $ SubExp se - | (pe, SubExpRes cs se) <- zip chunk_map_pes chunk_map_ses - ] - - pure (chunk_red_pes, chunk_map_pes) - --- | Given a chunked fold lambda that takes its initial accumulator --- value as parameters, bind those parameters to the neutral element --- instead. -kerneliseLambda :: - MonadFreshNames m => - [SubExp] -> - Lambda GPU -> - m (Lambda GPU) -kerneliseLambda nes lam = do - thread_index_param <- newParam "thread_index" $ Prim int64 - let (fold_chunk_param, fold_acc_params, fold_inp_params) = - partitionChunkedFoldParameters (length nes) $ lambdaParams lam - - mkAccInit p (Var v) - | not $ primType $ paramType p = - mkLet [paramIdent p] $ BasicOp $ Copy v - mkAccInit p x = mkLet [paramIdent p] $ BasicOp $ SubExp x - acc_init_stms = stmsFromList $ zipWith mkAccInit fold_acc_params nes - pure - lam - { lambdaBody = insertStms acc_init_stms $ lambdaBody lam, - lambdaParams = thread_index_param : fold_chunk_param : fold_inp_params - } - -prepareStream :: - (MonadBuilder m, Rep m ~ GPU) => - KernelSize -> - [(VName, SubExp)] -> - SubExp -> - Commutativity -> - Lambda GPU -> - [SubExp] -> - [VName] -> - m (SubExp, SegSpace, [Type], KernelBody GPU) -prepareStream size ispace w comm fold_lam nes arrs = do - let (KernelSize elems_per_thread num_threads) = size - let (ordering, split_ordering) = - case comm of - Commutative -> (Disorder, SplitStrided num_threads) - Noncommutative -> (InOrder, SplitContiguous) - - fold_lam' <- kerneliseLambda nes fold_lam - - gtid <- newVName "gtid" - space <- mkSegSpace $ ispace ++ [(gtid, num_threads)] - kbody <- fmap (uncurry (flip (KernelBody ()))) $ - runBuilder $ - localScope (scopeOfSegSpace space) $ do - (chunk_red_pes, chunk_map_pes) <- - blockedPerThread gtid w size ordering fold_lam' (length nes) arrs - let concatReturns pe = - ConcatReturns mempty split_ordering w elems_per_thread $ patElemName pe - pure - ( map (Returns ResultMaySimplify mempty . Var . patElemName) chunk_red_pes - ++ map concatReturns chunk_map_pes - ) - - let (redout_ts, mapout_ts) = splitAt (length nes) $ lambdaReturnType fold_lam - ts = redout_ts ++ map rowType mapout_ts - - pure (num_threads, space, ts, kbody) - -streamRed :: - (MonadFreshNames m, HasScope GPU m) => - MkSegLevel GPU m -> - Pat Type -> - SubExp -> - Commutativity -> - Lambda GPU -> - Lambda GPU -> - [SubExp] -> - [VName] -> - m (Stms GPU) -streamRed mk_lvl pat w comm red_lam fold_lam nes arrs = runBuilderT'_ $ do - -- The strategy here is to rephrase the stream reduction as a - -- non-segmented SegRed that does explicit chunking within its body. - -- First, figure out how many threads to use for this. - size <- blockedKernelSize "stream_red" w - - let (redout_pes, mapout_pes) = splitAt (length nes) $ patElems pat - (redout_pat, ispace, read_dummy) <- dummyDim $ Pat redout_pes - let pat' = Pat $ patElems redout_pat ++ mapout_pes - - (_, kspace, ts, kbody) <- prepareStream size ispace w comm fold_lam nes arrs - - lvl <- mk_lvl [w] "stream_red" $ NoRecommendation SegNoVirt - letBind pat' . Op . SegOp $ - SegRed lvl kspace [SegBinOp comm red_lam nes mempty] ts kbody - - read_dummy - --- Similar to streamRed, but without the last reduction. -streamMap :: - (MonadFreshNames m, HasScope GPU m) => - MkSegLevel GPU m -> - [String] -> - [PatElem Type] -> - SubExp -> - Commutativity -> - Lambda GPU -> - [SubExp] -> - [VName] -> - m ((SubExp, [VName]), Stms GPU) -streamMap mk_lvl out_desc mapout_pes w comm fold_lam nes arrs = runBuilderT' $ do - size <- blockedKernelSize "stream_map" w - - (threads, kspace, ts, kbody) <- prepareStream size [] w comm fold_lam nes arrs - - let redout_ts = take (length nes) ts - - redout_pes <- forM (zip out_desc redout_ts) $ \(desc, t) -> - PatElem <$> newVName desc <*> pure (t `arrayOfRow` threads) - - let pat = Pat $ redout_pes ++ mapout_pes - lvl <- mk_lvl [w] "stream_map" $ NoRecommendation SegNoVirt - letBind pat $ Op $ SegOp $ SegMap lvl kspace ts kbody - - pure (threads, map patElemName redout_pes) -- | Like 'segThread', but cap the thread count to the input size. -- This is more efficient for small kernels, e.g. summing a small diff --git a/src/Futhark/Pass/ExtractMulticore.hs b/src/Futhark/Pass/ExtractMulticore.hs --- a/src/Futhark/Pass/ExtractMulticore.hs +++ b/src/Futhark/Pass/ExtractMulticore.hs @@ -25,12 +25,10 @@ Stm, ) import qualified Futhark.IR.SOACS as SOACS -import qualified Futhark.IR.SOACS.Simplify as SOACS import Futhark.Pass import Futhark.Pass.ExtractKernels.DistributeNests import Futhark.Pass.ExtractKernels.ToGPU (injectSOACS) import Futhark.Tools -import qualified Futhark.Transform.FirstOrderTransform as FOT import Futhark.Transform.Rename (Rename, renameSomething) import Futhark.Util (takeLast) import Futhark.Util.Log @@ -163,56 +161,6 @@ body' <- localScope (scopeOfFParams params) $ transformBody body pure $ FunDef entry attrs name rettype params body' --- Sets the chunk size to one. -unstreamLambda :: Attrs -> [SubExp] -> Lambda SOACS -> ExtractM (Lambda SOACS) -unstreamLambda attrs nes lam = do - let (chunk_param, acc_params, slice_params) = - partitionChunkedFoldParameters (length nes) (lambdaParams lam) - - inp_params <- forM slice_params $ \(Param _ p t) -> - newParam (baseString p) (rowType t) - - body <- runBodyBuilder $ - localScope (scopeOfLParams inp_params) $ do - letBindNames [paramName chunk_param] $ - BasicOp $ - SubExp $ - intConst Int64 1 - - forM_ (zip acc_params nes) $ \(p, ne) -> - letBindNames [paramName p] $ BasicOp $ SubExp ne - - forM_ (zip slice_params inp_params) $ \(slice, v) -> - letBindNames [paramName slice] $ - BasicOp $ - ArrayLit [Var $ paramName v] (paramType v) - - (red_res, map_res) <- splitAt (length nes) <$> bodyBind (lambdaBody lam) - - map_res' <- forM map_res $ \(SubExpRes cs se) -> do - v <- letExp "map_res" $ BasicOp $ SubExp se - v_t <- lookupType v - certifying cs . letSubExp "chunk" . BasicOp $ - Index v $ - fullSlice v_t [DimFix $ intConst Int64 0] - - pure $ mkBody mempty $ red_res <> subExpsRes map_res' - - let (red_ts, map_ts) = splitAt (length nes) $ lambdaReturnType lam - map_lam = - Lambda - { lambdaReturnType = red_ts ++ map rowType map_ts, - lambdaParams = inp_params, - lambdaBody = body - } - - soacs_scope <- castScope <$> askScope - map_lam' <- runReaderT (SOACS.simplifyLambda map_lam) soacs_scope - - if "sequential_inner" `inAttrs` attrs - then FOT.transformLambda map_lam' - else pure map_lam' - -- Code generation for each parallel basic block is parameterised over -- how we handle parallelism in the body (whether it's sequentialised -- by keeping it as SOACs, or turned into SegOps). @@ -270,25 +218,6 @@ SegHist () space hists' (lambdaReturnType map_lam) kbody pure (hists_stms, op') -transformParStream :: - NeedsRename -> - (Body SOACS -> ExtractM (Body MC)) -> - SubExp -> - Commutativity -> - Lambda SOACS -> - [SubExp] -> - Lambda SOACS -> - [VName] -> - ExtractM (Stms MC, SegOp () MC) -transformParStream rename onBody w comm red_lam red_nes map_lam arrs = do - (gtid, space) <- mkSegSpace w - kbody <- mapLambdaToKernelBody onBody gtid map_lam arrs - (red_stms, red) <- reduceToSegBinOp $ Reduce comm red_lam red_nes - op <- - renameIfNeeded rename $ - SegRed () space [red] (lambdaReturnType map_lam) kbody - pure (red_stms, op) - transformSOAC :: Pat Type -> Attrs -> SOAC SOACS -> ExtractM (Stms MC) transformSOAC _ _ JVP {} = error "transformSOAC: unhandled JVP" @@ -368,33 +297,7 @@ pure $ mconcat seq_hist_stms <> oneStm (Let pat (defAux ()) $ Op $ ParOp Nothing seq_op) -transformSOAC pat attrs (Stream w arrs (Parallel _ comm red_lam) red_nes fold_lam) - | not $ null red_nes = do - map_lam <- unstreamLambda attrs red_nes fold_lam - (seq_red_stms, seq_op) <- - transformParStream - DoNotRename - sequentialiseBody - w - comm - red_lam - red_nes - map_lam - arrs - - if lambdaContainsParallelism map_lam - then do - (par_red_stms, par_op) <- - transformParStream DoRename transformBody w comm red_lam red_nes map_lam arrs - pure $ - seq_red_stms - <> par_red_stms - <> oneStm (Let pat (defAux ()) $ Op $ ParOp (Just par_op) seq_op) - else - pure $ - seq_red_stms - <> oneStm (Let pat (defAux ()) $ Op $ ParOp Nothing seq_op) -transformSOAC pat _ (Stream w arrs _ nes lam) = do +transformSOAC pat _ (Stream w arrs nes lam) = do -- Just remove the stream and transform the resulting stms. soacs_scope <- castScope <$> askScope stream_stms <- diff --git a/src/Futhark/Pass/KernelBabysitting.hs b/src/Futhark/Pass/KernelBabysitting.hs --- a/src/Futhark/Pass/KernelBabysitting.hs +++ b/src/Futhark/Pass/KernelBabysitting.hs @@ -8,7 +8,6 @@ import Control.Monad.State.Strict import Data.Foldable import Data.List (elemIndex, isPrefixOf, sort) -import Data.List.NonEmpty (NonEmpty (..)) import qualified Data.Map.Strict as M import Data.Maybe import Futhark.IR @@ -91,7 +90,7 @@ let mapper = identitySegOpMapper { mapOnSegOpBody = - transformKernelBody expmap (segLevel op) (segSpace op) + transformKernelBody expmap (segSpace op) } op' <- mapSegOpM mapper op let stm' = Let pat aux $ Op $ SegOp op' @@ -109,30 +108,22 @@ transformKernelBody :: ExpMap -> - SegLevel -> SegSpace -> KernelBody GPU -> BabysitM (KernelBody GPU) -transformKernelBody expmap lvl space kbody = do +transformKernelBody expmap space kbody = do -- Go spelunking for accesses to arrays that are defined outside the -- kernel body and where the indices are kernel thread indices. scope <- askScope let thread_gids = map fst $ unSegSpace space thread_local = namesFromList $ segFlat space : thread_gids free_ker_vars = freeIn kbody `namesSubtract` getKerVariantIds space - num_threads <- - letSubExp "num_threads" $ - BasicOp $ - BinOp - (Mul Int64 OverflowUndef) - (unCount $ segNumGroups lvl) - (unCount $ segGroupSize lvl) evalStateT ( traverseKernelBodyArrayIndexes free_ker_vars thread_local (scope <> scopeOfSegSpace space) - (ensureCoalescedAccess expmap (unSegSpace space) num_threads) + (ensureCoalescedAccess expmap (unSegSpace space)) kbody ) mempty @@ -143,7 +134,6 @@ Names -> (VName -> Bool) -> -- thread local? (VName -> SubExp -> Bool) -> -- variant to a certain gid (given as first param)? - (SubExp -> Maybe SubExp) -> -- split substitution? Scope GPU -> -- type environment VName -> Slice SubExp -> @@ -162,29 +152,27 @@ <$> mapM ( onStm ( varianceInStms mempty kstms, - mkSizeSubsts kstms, outer_scope ) ) (stmsToList kstms) <*> pure kres where - onLambda (variance, szsubst, scope) lam = + onLambda (variance, scope) lam = (\body' -> lam {lambdaBody = body'}) - <$> onBody (variance, szsubst, scope') (lambdaBody lam) + <$> onBody (variance, scope') (lambdaBody lam) where scope' = scope <> scopeOfLParams (lambdaParams lam) - onBody (variance, szsubst, scope) (Body bdec stms bres) = do - stms' <- stmsFromList <$> mapM (onStm (variance', szsubst', scope')) (stmsToList stms) + onBody (variance, scope) (Body bdec stms bres) = do + stms' <- stmsFromList <$> mapM (onStm (variance', scope')) (stmsToList stms) pure $ Body bdec stms' bres where variance' = varianceInStms variance stms - szsubst' = mkSizeSubsts stms <> szsubst scope' = scope <> scopeOf stms - onStm (variance, szsubst, _) (Let pat dec (BasicOp (Index arr is))) = - Let pat dec . oldOrNew <$> f free_ker_vars isThreadLocal isGidVariant sizeSubst outer_scope arr is + onStm (variance, _) (Let pat dec (BasicOp (Index arr is))) = + Let pat dec . oldOrNew <$> f free_ker_vars isThreadLocal isGidVariant outer_scope arr is where oldOrNew Nothing = BasicOp $ Index arr is @@ -198,14 +186,8 @@ isThreadLocal v = thread_variant `namesIntersect` M.findWithDefault (oneName v) v variance - - sizeSubst (Constant v) = Just $ Constant v - sizeSubst (Var v) - | v `M.member` outer_scope = Just $ Var v - | Just v' <- M.lookup v szsubst = sizeSubst v' - | otherwise = Nothing - onStm (variance, szsubst, scope) (Let pat dec e) = - Let pat dec <$> mapExpM (mapper (variance, szsubst, scope)) e + onStm (variance, scope) (Let pat dec e) = + Let pat dec <$> mapExpM (mapper (variance, scope)) e onOp ctx (OtherOp soac) = OtherOp <$> mapSOACM identitySOACMapper {mapOnSOACLambda = onLambda ctx} soac @@ -217,28 +199,19 @@ mapOnOp = onOp ctx } - mkSizeSubsts = foldMap mkStmSizeSubst - where - mkStmSizeSubst (Let (Pat [pe]) _ (Op (SizeOp (SplitSpace _ _ _ elems_per_i)))) = - M.singleton (patElemName pe) elems_per_i - mkStmSizeSubst _ = mempty - type Replacements = M.Map (VName, Slice SubExp) VName ensureCoalescedAccess :: MonadBuilder m => ExpMap -> [(VName, SubExp)] -> - SubExp -> ArrayIndexTransform (StateT Replacements m) ensureCoalescedAccess expmap thread_space - num_threads free_ker_vars isThreadLocal isGidVariant - sizeSubst outer_scope arr slice = do @@ -302,28 +275,6 @@ let perm = coalescingPermutation (length is) $ arrayRank t replace =<< lift (rearrangeInput (nonlinearInMemory arr expmap) perm arr) - -- We are taking a slice of the array with a unit stride. We - -- assume that the slice will be traversed sequentially. - -- - -- We will really want to treat the sliced dimension like two - -- dimensions so we can transpose them. This may require - -- padding. - | (is, rem_slice) <- splitSlice slice, - and $ zipWith (==) is $ map Var thread_gids, - DimSlice offset len (Constant stride) : _ <- unSlice rem_slice, - isThreadLocalSubExp offset, - Just {} <- sizeSubst len, - oneIsh stride -> do - let num_chunks = - if null is - then untyped $ pe32 num_threads - else - untyped $ - product $ - map pe64 $ - drop (length is) thread_gdims - replace =<< lift (rearrangeSlice (length is) (arraySize (length is) t) num_chunks arr) - -- Everything is fine... assuming that the array is in row-major -- order! Make sure that is the case. | Just {} <- nonlinearInMemory arr expmap -> @@ -336,15 +287,12 @@ _ -> replace =<< lift (rowMajorArray arr) _ -> pure Nothing where - (thread_gids, thread_gdims) = unzip thread_space + (thread_gids, _thread_gdims) = unzip thread_space replace arr' = do modify $ M.insert (arr, slice) arr' pure $ Just (arr', slice) - isThreadLocalSubExp (Var v) = isThreadLocal v - isThreadLocalSubExp Constant {} = False - -- Heuristic for avoiding rearranging too small arrays. tooSmallSlice :: Int32 -> Slice SubExp -> Bool tooSmallSlice bs = fst . foldl comb (True, bs) . sliceDims @@ -449,68 +397,6 @@ rowMajorArray arr = do rank <- arrayRank <$> lookupType arr letExp (baseString arr ++ "_rowmajor") $ BasicOp $ Manifest [0 .. rank - 1] arr - -rearrangeSlice :: - MonadBuilder m => - Int -> - SubExp -> - PrimExp VName -> - VName -> - m VName -rearrangeSlice d w num_chunks arr = do - num_chunks' <- toSubExp "num_chunks" num_chunks - - (w_padded, padding) <- paddedScanReduceInput w num_chunks' - - per_chunk <- - letSubExp "per_chunk" $ - BasicOp $ - BinOp (SQuot Int64 Unsafe) w_padded num_chunks' - arr_t <- lookupType arr - arr_padded <- padArray w_padded padding arr_t - rearrange num_chunks' w_padded per_chunk (baseString arr) arr_padded arr_t - where - padArray w_padded padding arr_t = do - let arr_shape = arrayShape arr_t - padding_shape = setDim d arr_shape padding - arr_padding <- - letExp (baseString arr <> "_padding") . BasicOp $ - Scratch (elemType arr_t) (shapeDims padding_shape) - letExp (baseString arr <> "_padded") . BasicOp $ - Concat d (arr :| [arr_padding]) w_padded - - rearrange num_chunks' w_padded per_chunk arr_name arr_padded arr_t = do - let arr_dims = arrayDims arr_t - pre_dims = take d arr_dims - post_dims = drop (d + 1) arr_dims - extradim_shape = Shape $ pre_dims ++ [num_chunks', per_chunk] ++ post_dims - tr_perm = [0 .. d - 1] ++ map (+ d) ([1] ++ [2 .. shapeRank extradim_shape - 1 - d] ++ [0]) - arr_extradim <- - letExp (arr_name <> "_extradim") . BasicOp $ - Reshape ReshapeArbitrary extradim_shape arr_padded - arr_extradim_tr <- - letExp (arr_name <> "_extradim_tr") . BasicOp $ - Manifest tr_perm arr_extradim - arr_inv_tr <- - letExp (arr_name <> "_inv_tr") . BasicOp $ - Reshape - ReshapeArbitrary - (Shape $ pre_dims ++ w_padded : post_dims) - arr_extradim_tr - letExp (arr_name <> "_inv_tr_init") - =<< eSliceArray d arr_inv_tr (eSubExp $ constant (0 :: Int64)) (eSubExp w) - -paddedScanReduceInput :: - MonadBuilder m => - SubExp -> - SubExp -> - m (SubExp, SubExp) -paddedScanReduceInput w stride = do - w_padded <- - letSubExp "padded_size" - =<< eRoundToMultipleOf Int64 (eSubExp w) (eSubExp stride) - padding <- letSubExp "padding" $ BasicOp $ BinOp (Sub Int64 OverflowUndef) w_padded w - pure (w_padded, padding) --- Computing variance. diff --git a/src/Futhark/Test/Spec.hs b/src/Futhark/Test/Spec.hs --- a/src/Futhark/Test/Spec.hs +++ b/src/Futhark/Test/Spec.hs @@ -45,6 +45,7 @@ import Futhark.Util.Pretty (prettyOneLine) import System.Exit import System.FilePath +import System.IO import System.IO.Error import Text.Megaparsec hiding (many, some) import Text.Megaparsec.Char @@ -407,7 +408,7 @@ spec_or_err <- testSpecFromProgram prog case spec_or_err of Left err -> do - putStrLn err + hPutStrLn stderr err exitFailure Right spec -> pure spec @@ -443,7 +444,7 @@ specs_or_err <- testSpecsFromPaths dirs case specs_or_err of Left err -> do - putStrLn err + hPutStrLn stderr err exitFailure Right specs -> pure specs @@ -462,6 +463,6 @@ spec_or_err <- testSpecFromFile dirs case spec_or_err of Left err -> do - putStrLn err + hPutStrLn stderr err exitFailure Right spec -> pure spec diff --git a/src/Futhark/Transform/FirstOrderTransform.hs b/src/Futhark/Transform/FirstOrderTransform.hs --- a/src/Futhark/Transform/FirstOrderTransform.hs +++ b/src/Futhark/Transform/FirstOrderTransform.hs @@ -230,7 +230,7 @@ (++ patNames pat) <$> replicateM (length scanacc_params) (newVName "discard") letBindNames names $ DoLoop merge loopform loop_body -transformSOAC pat (Stream w arrs _ nes lam) = do +transformSOAC pat (Stream w arrs nes lam) = do -- Create a loop that repeatedly applies the lambda body to a -- chunksize of 1. Hopefully this will lead to this outer loop -- being the only one, as all the innermost one can be simplified diff --git a/src/Language/Futhark/Parser/Lexer.x b/src/Language/Futhark/Parser/Lexer.x --- a/src/Language/Futhark/Parser/Lexer.x +++ b/src/Language/Futhark/Parser/Lexer.x @@ -67,7 +67,7 @@ map (T.drop 3 . T.stripStart) . T.split (== '\n') . ("--"<>) . T.drop 4 } - "--".* ; + "--".* { tokenS COMMENT } "=" { tokenC EQU } "(" { tokenC LPAR } ")" { tokenC RPAR } diff --git a/src/Language/Futhark/Parser/Lexer/Tokens.hs b/src/Language/Futhark/Parser/Lexer/Tokens.hs --- a/src/Language/Futhark/Parser/Lexer/Tokens.hs +++ b/src/Language/Futhark/Parser/Lexer/Tokens.hs @@ -54,6 +54,7 @@ -- with a source position. data Token = ID Name + | COMMENT T.Text | INDEXING Name | QUALINDEXING [Name] Name | QUALPAREN [Name] Name diff --git a/src/Language/Futhark/Parser/Monad.hs b/src/Language/Futhark/Parser/Monad.hs --- a/src/Language/Futhark/Parser/Monad.hs +++ b/src/Language/Futhark/Parser/Monad.hs @@ -246,6 +246,9 @@ xs -> do putTokens (xs, pos') lexer cont + (L _ (COMMENT _) : xs) -> do + putTokens (xs, pos) + lexer cont (x : xs) -> do putTokens (xs, pos) cont x diff --git a/src/Language/Futhark/Parser/Parser.y b/src/Language/Futhark/Parser/Parser.y --- a/src/Language/Futhark/Parser/Parser.y +++ b/src/Language/Futhark/Parser/Parser.y @@ -39,7 +39,7 @@ import Language.Futhark.Syntax hiding (ID) import Language.Futhark.Prop import Language.Futhark.Pretty -import Language.Futhark.Parser.Lexer +import Language.Futhark.Parser.Lexer (Token(..)) import Futhark.Util.Pretty import Futhark.Util.Loc import Language.Futhark.Parser.Monad @@ -564,66 +564,22 @@ | Exp2 %prec ':' { $1 } Exp2 :: { UncheckedExp } - : if Exp then Exp else Exp %prec ifprec - { AppExp (If $2 $4 $6 (srcspan $1 $>)) NoInfo } - - | loop Pat LoopForm do Exp %prec ifprec - {% fmap (\t -> AppExp (DoLoop [] $2 t $3 $5 (srcspan $1 $>)) NoInfo) (patternExp $2) } - - | loop Pat '=' Exp LoopForm do Exp %prec ifprec - { AppExp (DoLoop [] $2 $4 $5 $7 (srcspan $1 $>)) NoInfo } - + : IfExp { $1 } + | LoopExp { $1 } | LetExp %prec letprec { $1 } - - | MatchExp { $1 } + | MatchExp { $1 } | assert Atom Atom { Assert $2 $3 NoInfo (srcspan $1 $>) } | '#[' AttrInfo ']' Exp %prec bottom { Attr $2 $4 (srcspan $1 $>) } - | Exp2 '+...' Exp2 { binOp $1 $2 $3 } - | Exp2 '-...' Exp2 { binOp $1 $2 $3 } - | Exp2 '-' Exp2 { binOp $1 (L $2 (SYMBOL Minus [] (nameFromString "-"))) $3 } - | Exp2 '*...' Exp2 { binOp $1 $2 $3 } - | Exp2 '*' Exp2 { binOp $1 (L $2 (SYMBOL Times [] (nameFromString "*"))) $3 } - | Exp2 '/...' Exp2 { binOp $1 $2 $3 } - | Exp2 '%...' Exp2 { binOp $1 $2 $3 } - | Exp2 '//...' Exp2 { binOp $1 $2 $3 } - | Exp2 '%%...' Exp2 { binOp $1 $2 $3 } - | Exp2 '**...' Exp2 { binOp $1 $2 $3 } - | Exp2 '>>...' Exp2 { binOp $1 $2 $3 } - | Exp2 '<<...' Exp2 { binOp $1 $2 $3 } - | Exp2 '&...' Exp2 { binOp $1 $2 $3 } - | Exp2 '|...' Exp2 { binOp $1 $2 $3 } - | Exp2 '|' Exp2 { binOp $1 (L $2 (SYMBOL Bor [] (nameFromString "|"))) $3 } - | Exp2 '&&...' Exp2 { binOp $1 $2 $3 } - | Exp2 '||...' Exp2 { binOp $1 $2 $3 } - | Exp2 '^...' Exp2 { binOp $1 $2 $3 } - | Exp2 '^' Exp2 { binOp $1 (L $2 (SYMBOL Xor [] (nameFromString "^"))) $3 } - | Exp2 '==...' Exp2 { binOp $1 $2 $3 } - | Exp2 '!=...' Exp2 { binOp $1 $2 $3 } - | Exp2 '<...' Exp2 { binOp $1 $2 $3 } - | Exp2 '<=...' Exp2 { binOp $1 $2 $3 } - | Exp2 '>...' Exp2 { binOp $1 $2 $3 } - | Exp2 '>=...' Exp2 { binOp $1 $2 $3 } - | Exp2 '|>...' Exp2 { binOp $1 $2 $3 } - | Exp2 '<|...' Exp2 { binOp $1 $2 $3 } - - | Exp2 '<' Exp2 { binOp $1 (L $2 (SYMBOL Less [] (nameFromString "<"))) $3 } - | Exp2 '`' QualName '`' Exp2 { AppExp (BinOp (second srclocOf $3) NoInfo ($1, NoInfo) ($5, NoInfo) (srcspan $1 $>)) NoInfo } - - | Exp2 '...' Exp2 { AppExp (Range $1 Nothing (ToInclusive $3) (srcspan $1 $>)) NoInfo } - | Exp2 '..<' Exp2 { AppExp (Range $1 Nothing (UpToExclusive $3) (srcspan $1 $>)) NoInfo } - | Exp2 '..>' Exp2 { AppExp (Range $1 Nothing (DownToExclusive $3) (srcspan $1 $>)) NoInfo } - | Exp2 '..' Exp2 '...' Exp2 { AppExp (Range $1 (Just $3) (ToInclusive $5) (srcspan $1 $>)) NoInfo } - | Exp2 '..' Exp2 '..<' Exp2 { AppExp (Range $1 (Just $3) (UpToExclusive $5) (srcspan $1 $>)) NoInfo } - | Exp2 '..' Exp2 '..>' Exp2 { AppExp (Range $1 (Just $3) (DownToExclusive $5) (srcspan $1 $>)) NoInfo } + | BinOpExp { $1 } + | RangeExp { $1 } | Exp2 '..' Atom {% twoDotsRange $2 } | Atom '..' Exp2 {% twoDotsRange $2 } | '-' Exp2 %prec juxtprec { Negate $2 (srcspan $1 $>) } | '!' Exp2 %prec juxtprec { Not $2 (srcspan $1 $>) } - | Exp2 with '[' DimIndices ']' '=' Exp2 { Update $1 $4 $7 (srcspan $1 $>) } @@ -633,10 +589,7 @@ | '\\' FunParams1 maybeAscription(TypeExpTerm) '->' Exp %prec letprec { Lambda (fst $2 : snd $2) $5 $3 NoInfo (srcspan $1 $>) } - | Apply_ { $1 } - -Apply_ :: { UncheckedExp } - : ApplyList {% applyExp $1 } + | ApplyList {% applyExp $1 } ApplyList :: { [UncheckedExp] } : ApplyList Atom %prec juxtprec @@ -676,25 +629,7 @@ { let L loc (QUALPAREN qs name) = $1 in QualParens (QualName qs name, srclocOf loc) $2 (srcspan $1 $>) } - -- Operator sections. - | '(' '-' ')' - { OpSection (qualName (nameFromString "-")) NoInfo (srcspan $1 $>) } - | '(' Exp2 '-' ')' - { OpSectionLeft (qualName (nameFromString "-")) - NoInfo $2 (NoInfo, NoInfo) (NoInfo, NoInfo) (srcspan $1 $>) } - | '(' BinOp Exp2 ')' - { OpSectionRight (fst $2) NoInfo $3 (NoInfo, NoInfo) NoInfo (srcspan $1 $>) } - | '(' Exp2 BinOp ')' - { OpSectionLeft (fst $3) NoInfo $2 (NoInfo, NoInfo) (NoInfo, NoInfo) (srcspan $1 $>) } - | '(' BinOp ')' - { OpSection (fst $2) NoInfo (srcspan $1 $>) } - - | '(' FieldAccess FieldAccesses ')' - { ProjectSection (map fst ($2:$3)) NoInfo (srcspan $1 $>) } - - | '(' '.' '[' DimIndices ']' ')' - { IndexSection $4 NoInfo (srcspan $1 $>) } - + | SectionExp { $1 } NumLit :: { (PrimValue, Loc) } : i8lit { let L loc (I8LIT num) = $1 in (SignedValue $ Int8Value num, loc) } @@ -771,6 +706,74 @@ | def {% parseErrorAt $1 (Just "Unexpected \"def\" - missing \"in\"?") } | type {% parseErrorAt $1 (Just "Unexpected \"type\" - missing \"in\"?") } | module {% parseErrorAt $1 (Just "Unexpected \"module\" - missing \"in\"?") } + +BinOpExp :: { UncheckedExp } + : Exp2 '+...' Exp2 { binOp $1 $2 $3 } + | Exp2 '-...' Exp2 { binOp $1 $2 $3 } + | Exp2 '-' Exp2 { binOp $1 (L $2 (SYMBOL Minus [] (nameFromString "-"))) $3 } + | Exp2 '*...' Exp2 { binOp $1 $2 $3 } + | Exp2 '*' Exp2 { binOp $1 (L $2 (SYMBOL Times [] (nameFromString "*"))) $3 } + | Exp2 '/...' Exp2 { binOp $1 $2 $3 } + | Exp2 '%...' Exp2 { binOp $1 $2 $3 } + | Exp2 '//...' Exp2 { binOp $1 $2 $3 } + | Exp2 '%%...' Exp2 { binOp $1 $2 $3 } + | Exp2 '**...' Exp2 { binOp $1 $2 $3 } + | Exp2 '>>...' Exp2 { binOp $1 $2 $3 } + | Exp2 '<<...' Exp2 { binOp $1 $2 $3 } + | Exp2 '&...' Exp2 { binOp $1 $2 $3 } + | Exp2 '|...' Exp2 { binOp $1 $2 $3 } + | Exp2 '|' Exp2 { binOp $1 (L $2 (SYMBOL Bor [] (nameFromString "|"))) $3 } + | Exp2 '&&...' Exp2 { binOp $1 $2 $3 } + | Exp2 '||...' Exp2 { binOp $1 $2 $3 } + | Exp2 '^...' Exp2 { binOp $1 $2 $3 } + | Exp2 '^' Exp2 { binOp $1 (L $2 (SYMBOL Xor [] (nameFromString "^"))) $3 } + | Exp2 '==...' Exp2 { binOp $1 $2 $3 } + | Exp2 '!=...' Exp2 { binOp $1 $2 $3 } + | Exp2 '<...' Exp2 { binOp $1 $2 $3 } + | Exp2 '<=...' Exp2 { binOp $1 $2 $3 } + | Exp2 '>...' Exp2 { binOp $1 $2 $3 } + | Exp2 '>=...' Exp2 { binOp $1 $2 $3 } + | Exp2 '|>...' Exp2 { binOp $1 $2 $3 } + | Exp2 '<|...' Exp2 { binOp $1 $2 $3 } + | Exp2 '<' Exp2 { binOp $1 (L $2 (SYMBOL Less [] (nameFromString "<"))) $3 } + | Exp2 '`' QualName '`' Exp2 { AppExp (BinOp (second srclocOf $3) NoInfo ($1, NoInfo) ($5, NoInfo) (srcspan $1 $>)) NoInfo } + +SectionExp :: { UncheckedExp } + : '(' '-' ')' + { OpSection (qualName (nameFromString "-")) NoInfo (srcspan $1 $>) } + | '(' Exp2 '-' ')' + { OpSectionLeft (qualName (nameFromString "-")) + NoInfo $2 (NoInfo, NoInfo) (NoInfo, NoInfo) (srcspan $1 $>) } + | '(' BinOp Exp2 ')' + { OpSectionRight (fst $2) NoInfo $3 (NoInfo, NoInfo) NoInfo (srcspan $1 $>) } + | '(' Exp2 BinOp ')' + { OpSectionLeft (fst $3) NoInfo $2 (NoInfo, NoInfo) (NoInfo, NoInfo) (srcspan $1 $>) } + | '(' BinOp ')' + { OpSection (fst $2) NoInfo (srcspan $1 $>) } + + | '(' FieldAccess FieldAccesses ')' + { ProjectSection (map fst ($2:$3)) NoInfo (srcspan $1 $>) } + + | '(' '.' '[' DimIndices ']' ')' + { IndexSection $4 NoInfo (srcspan $1 $>) } + +RangeExp :: { UncheckedExp } + : Exp2 '...' Exp2 { AppExp (Range $1 Nothing (ToInclusive $3) (srcspan $1 $>)) NoInfo } + | Exp2 '..<' Exp2 { AppExp (Range $1 Nothing (UpToExclusive $3) (srcspan $1 $>)) NoInfo } + | Exp2 '..>' Exp2 { AppExp (Range $1 Nothing (DownToExclusive $3) (srcspan $1 $>)) NoInfo } + | Exp2 '..' Exp2 '...' Exp2 { AppExp (Range $1 (Just $3) (ToInclusive $5) (srcspan $1 $>)) NoInfo } + | Exp2 '..' Exp2 '..<' Exp2 { AppExp (Range $1 (Just $3) (UpToExclusive $5) (srcspan $1 $>)) NoInfo } + | Exp2 '..' Exp2 '..>' Exp2 { AppExp (Range $1 (Just $3) (DownToExclusive $5) (srcspan $1 $>)) NoInfo } + +IfExp :: { UncheckedExp } + : if Exp then Exp else Exp %prec ifprec + { AppExp (If $2 $4 $6 (srcspan $1 $>)) NoInfo } + +LoopExp :: { UncheckedExp } + : loop Pat LoopForm do Exp %prec ifprec + {% fmap (\t -> AppExp (DoLoop [] $2 t $3 $5 (srcspan $1 $>)) NoInfo) (patternExp $2) } + | loop Pat '=' Exp LoopForm do Exp %prec ifprec + { AppExp (DoLoop [] $2 $4 $5 $7 (srcspan $1 $>)) NoInfo } MatchExp :: { UncheckedExp } : match Exp Cases diff --git a/src/Language/Futhark/Prop.hs b/src/Language/Futhark/Prop.hs --- a/src/Language/Futhark/Prop.hs +++ b/src/Language/Futhark/Prop.hs @@ -919,46 +919,6 @@ ] ) ), - ( "map_stream", - IntrinsicPolyFun - [tp_a, tp_b, sp_n] - [ Scalar (Prim $ Signed Int64) `karr` (arr_ka `arr` arr_kb), - arr_a $ shape [n] - ] - $ RetType [] - $ uarr_b - $ shape [n] - ), - ( "map_stream_per", - IntrinsicPolyFun - [tp_a, tp_b, sp_n] - [ Scalar (Prim $ Signed Int64) `karr` (arr_ka `arr` arr_kb), - arr_a $ shape [n] - ] - $ RetType [] - $ uarr_b - $ shape [n] - ), - ( "reduce_stream", - IntrinsicPolyFun - [tp_a, tp_b, sp_n] - [ Scalar t_b `arr` (Scalar t_b `arr` Scalar t_b), - Scalar (Prim $ Signed Int64) `karr` (arr_ka `arr` Scalar t_b), - arr_a $ shape [n] - ] - $ RetType [] - $ Scalar t_b - ), - ( "reduce_stream_per", - IntrinsicPolyFun - [tp_a, tp_b, sp_n] - [ Scalar t_b `arr` (Scalar t_b `arr` Scalar t_b), - Scalar (Prim $ Signed Int64) `karr` (arr_ka `arr` Scalar t_b), - arr_a $ shape [n] - ] - $ RetType [] - $ Scalar t_b - ), ( "acc_write", IntrinsicPolyFun [sp_k, tp_a] @@ -1143,8 +1103,6 @@ arr_ka = Array () Nonunique (Shape [NamedSize $ qualName k]) t_a uarr_ka = Array () Unique (Shape [NamedSize $ qualName k]) t_a - arr_kb = Array () Nonunique (Shape [NamedSize $ qualName k]) t_b - karr x y = Scalar $ Arrow mempty (Named k) x (RetType [] y) accType t = TypeVar () Unique (qualName (fst intrinsicAcc)) [TypeArgType t mempty] diff --git a/unittests/Futhark/IR/Mem/IxFunTests.hs b/unittests/Futhark/IR/Mem/IxFunTests.hs --- a/unittests/Futhark/IR/Mem/IxFunTests.hs +++ b/unittests/Futhark/IR/Mem/IxFunTests.hs @@ -5,7 +5,7 @@ ) where -import qualified Data.List as DL +import qualified Data.List as L import qualified Futhark.IR.Mem.IxFun as IxFunLMAD import qualified Futhark.IR.Mem.IxFun.Alg as IxFunAlg import Futhark.IR.Mem.IxFunWrapper @@ -30,7 +30,7 @@ allPoints :: [Int] -> [[Int]] allPoints dims = let total = product dims - strides = drop 1 $ DL.reverse $ scanl (*) 1 $ DL.reverse dims + strides = drop 1 $ L.reverse $ scanl (*) 1 $ L.reverse dims in map (unflatInd strides) [0 .. total - 1] where unflatInd :: [Int] -> Int -> [Int]