futhark 0.21.15 → 0.22.1

raw patch · 51 files changed

+2571/−3689 lines, 51 filesPVP ok

version bump matches the API change (PVP)

API changes (from Hackage documentation)

Files

• docs/error-index.rst +4/−5
• futhark.cabal +2/−1
• prelude/soacs.fut +0/−48
• src/Futhark/AD/Fwd.hs +2/−8
• src/Futhark/Analysis/HORep/SOAC.hs +17/−21
• src/Futhark/CodeGen/Backends/GenericC.hs +2/−10
• src/Futhark/CodeGen/Backends/GenericC/Code.hs +51/−25
• src/Futhark/CodeGen/ImpGen/GPU/Base.hs +1382/−2154
• src/Futhark/CodeGen/ImpGen/GPU/Group.hs +735/−0
• src/Futhark/CodeGen/ImpGen/GPU/SegMap.hs +1/−0
• src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs +52/−9
• src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs +9/−13
• src/Futhark/CodeGen/ImpGen/Multicore/Base.hs +0/−2
• src/Futhark/Construct.hs +0/−35
• src/Futhark/Doc/Generator.hs +4/−4
• src/Futhark/IR/GPU/Op.hs +1/−50
• src/Futhark/IR/GPU/Simplify.hs +0/−11
• src/Futhark/IR/Mem.hs +1/−1
• src/Futhark/IR/Mem/Simplify.hs +15/−12
• src/Futhark/IR/Parse.hs +1/−61
• src/Futhark/IR/SOACS.hs +1/−1
• src/Futhark/IR/SOACS/SOAC.hs +20/−86
• src/Futhark/IR/SOACS/Simplify.hs +6/−18
• src/Futhark/IR/SegOp.hs +0/−82
• src/Futhark/Internalise/Exps.hs +0/−84
• src/Futhark/Internalise/Lambdas.hs +0/−45
• src/Futhark/Optimise/Fusion.hs +1/−1
• src/Futhark/Optimise/Fusion/Composing.hs +0/−12
• src/Futhark/Optimise/Fusion/GraphRep.hs +2/−2
• src/Futhark/Optimise/Fusion/TryFusion.hs +47/−81
• src/Futhark/Optimise/ReduceDeviceSyncs.hs +36/−20
• src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs +24/−18
• src/Futhark/Pass/ExplicitAllocations.hs +37/−75
• src/Futhark/Pass/ExplicitAllocations/GPU.hs +2/−21
• src/Futhark/Pass/ExplicitAllocations/MC.hs +1/−2
• src/Futhark/Pass/ExplicitAllocations/SegOp.hs +1/−2
• src/Futhark/Pass/ExtractKernels.hs +10/−90
• src/Futhark/Pass/ExtractKernels/DistributeNests.hs +1/−1
• src/Futhark/Pass/ExtractKernels/Interchange.hs +1/−1
• src/Futhark/Pass/ExtractKernels/Intragroup.hs +1/−1
• src/Futhark/Pass/ExtractKernels/StreamKernel.hs +0/−230
• src/Futhark/Pass/ExtractMulticore.hs +1/−98
• src/Futhark/Pass/KernelBabysitting.hs +12/−126
• src/Futhark/Test/Spec.hs +4/−3
• src/Futhark/Transform/FirstOrderTransform.hs +1/−1
• src/Language/Futhark/Parser/Lexer.x +1/−1
• src/Language/Futhark/Parser/Lexer/Tokens.hs +1/−0
• src/Language/Futhark/Parser/Monad.hs +3/−0
• src/Language/Futhark/Parser/Parser.y +76/−73
• src/Language/Futhark/Prop.hs +0/−42
• unittests/Futhark/IR/Mem/IxFunTests.hs +2/−2

@@ -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 

@@ -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

@@ -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. --

@@ -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

@@ -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, [])

@@ -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

@@ -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

@@ -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."

@@ -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')

@@ -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)

@@ -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]

@@ -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

@@ -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 {} =

@@ -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))

@@ -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)

@@ -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

@@ -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)) =

@@ -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 

@@ -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 =

@@ -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

@@ -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

@@ -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 <length> <inputs> <lambda> <outputs>@     --     -- 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 ::

@@ -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)

@@ -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) =

@@ -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 ->

@@ -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

@@ -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

@@ -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)

@@ -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

@@ -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. 

@@ -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.

@@ -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 ::

@@ -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) ->

@@ -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

@@ -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

@@ -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 =>

@@ -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 []

@@ -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

@@ -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 ->

@@ -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) <-

@@ -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

@@ -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 <-

@@ -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. 

@@ -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

@@ -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

@@ -67,7 +67,7 @@                                       map (T.drop 3 . T.stripStart) .                                            T.split (== '\n') . ("--"<>) .                                            T.drop 4 }-  "--".*                            ;+  "--".*                   { tokenS COMMENT }   "="                      { tokenC EQU }   "("                      { tokenC LPAR }   ")"                      { tokenC RPAR }

@@ -54,6 +54,7 @@ -- with a source position. data Token   = ID Name+  | COMMENT T.Text   | INDEXING Name   | QUALINDEXING [Name] Name   | QUALPAREN [Name] Name

@@ -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

@@ -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