futhark 0.21.4 → 0.21.5

raw patch · 28 files changed

+650/−298 lines, 28 filesPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

API changes (from Hackage documentation)

Files

• docs/c-api.rst +26/−3
• docs/conf.py +1/−1
• docs/performance.rst +20/−8
• docs/server-protocol.rst +3/−3
• docs/usage.rst +1/−1
• futhark.cabal +1/−1
• rts/c/cuda.h +3/−4
• rts/c/errors.h +3/−0
• src/Futhark/Analysis/PrimExp.hs +28/−6
• src/Futhark/CodeGen/Backends/CCUDA.hs +1/−1
• src/Futhark/CodeGen/Backends/CCUDA/Boilerplate.hs +3/−1
• src/Futhark/CodeGen/Backends/COpenCL.hs +1/−1
• src/Futhark/CodeGen/Backends/COpenCL/Boilerplate.hs +3/−1
• src/Futhark/CodeGen/Backends/GenericC.hs +31/−10
• src/Futhark/CodeGen/Backends/MulticoreC.hs +2/−2
• src/Futhark/CodeGen/ImpGen/GPU/Base.hs +336/−108
• src/Futhark/CodeGen/ImpGen/GPU/SegMap.hs +3/−2
• src/Futhark/CodeGen/ImpGen/Multicore/Base.hs +16/−18
• src/Futhark/CodeGen/RTS/C.hs +6/−0
• src/Futhark/IR/GPU/Op.hs +1/−1
• src/Futhark/IR/Parse.hs +2/−1
• src/Futhark/IR/SOACS/Simplify.hs +17/−17
• src/Futhark/IR/SegOp.hs +20/−1
• src/Futhark/Optimise/BlkRegTiling.hs +90/−87
• src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs +5/−6
• src/Futhark/Optimise/TileLoops.hs +3/−5
• src/Futhark/Optimise/TileLoops/Shared.hs +22/−6
• src/Language/Futhark/Parser/Parser.y +2/−3

@@ -13,9 +13,9 @@ most other operations, such as calling Futhark functions.  Most functions that can fail return an integer: 0 on success and a-non-zero value on error.  Others return a ``NULL`` pointer.  Use-:c:func:`futhark_context_get_error` to get a (possibly) more precise-error message.+non-zero value on error, as documented below.  Others return a+``NULL`` pointer.  Use :c:func:`futhark_context_get_error` to get a+(possibly) more precise error message.  .. c:macro:: FUTHARK_BACKEND_foo @@ -23,6 +23,29 @@    generate the code; e.g. ``c``, ``opencl``, or ``cuda``.  This can    be used for conditional compilation of code that only works with    specific backends.++Error codes+-----------++Most errors are result in a not otherwise specified nonzero return+code, but a few classes of errors have distinct error codes.++.. c:macro:: FUTHARK_SUCCESS++   Defined as ``0``.  Returned in case of success.++.. c:macro:: FUTHARK_PROGRAM_ERROR++   Defined as ``2``.  Returned when the program fails due to+   out-of-bounds, an invalid size coercion, invalid entry point+   arguments, or similar misuse.++.. c:macro:: FUTHARK_OUT_OF_MEMORY++   Defined as ``3``.  Returned when the program fails to allocate+   memory.  This is (somewhat) reliable only for GPU memory - due to+   overcommit and other VM tricks, you should not expect running out+   of main memory to be reported gracefully.  Configuration -------------

@@ -118,7 +118,7 @@  lexers['futhark'] = FutharkLexer() -highlight_language = 'futhark'+highlight_language = 'text'  # A list of ignored prefixes for module index sorting. #modindex_common_prefix = []

@@ -45,14 +45,18 @@ The ``scan`` and ``reduce`` SOACs are rather inefficient when their operators are on arrays.  If possible, use tuples instead (see :ref:`performance-small-arrays`).  The one exception is when the-operator is a ``map2`` or equivalent.  Example::+operator is a ``map2`` or equivalent.  Example: +.. code-block:: futhark+   reduce (map2 (+)) (replicate n 0) xss  Such "vectorised" operators are typically handled quite efficiently. Although to be on the safe side, you can rewrite the above by-interchanging the ``reduce`` and ``map``::+interchanging the ``reduce`` and ``map``: +.. code-block:: futhark+   map (reduce (+) 0) (transpose xss)  Avoid reductions over tiny arrays, e.g. ``reduce (+) 0 [x,y,z]``.  In@@ -85,8 +89,10 @@  Futhark allows nested parallelism, understood as a parallel construct used inside some other parallel construct.  The simplest example is-nested SOACs.  Example::+nested SOACs.  Example: +.. code-block:: futhark+   map (\xs -> reduce (+) 0 xs) xss  Nested parallelism is allowed and encouraged, but its compilation to@@ -126,8 +132,10 @@ The main restriction is that the GPU backends can only handle *regular* nested parallelism, meaning that the sizes of inner parallel dimensions are invariant to the outer parallel dimensions.  For-example, this expression contains *irregular* nested parallelism::+example, this expression contains *irregular* nested parallelism: +.. code-block:: futhark+   map (\i -> reduce (+) 0 (iota i)) is  This is because the size of the nested parallel construct is ``i``,@@ -182,8 +190,8 @@ Sum Types ~~~~~~~~~ -A starting point, a sum type is turned into a tuple containing all the-payload components in order, prefixed with an `i8` tag to identify the+A sum type value is represented as a tuple containing all the payload+components in order, prefixed with an `i8` tag to identify the constructor.  For example,  .. code-block:: futhark@@ -345,8 +353,10 @@  *Horizontal fusion* occurs when two SOACs take as input the same array, but are not themselves in a producer-consumer relationship.-Example::+Example: +.. code-block:: futhark+    (map f xs, map g xs)  Such cases are fused into a single operation that traverses ``xs``@@ -418,7 +428,9 @@ performance.  In many cases this is currently unavoidable, but sometimes the program can be rewritten such that instead of calling the same function in multiple places, it is called in a single place,-in a loop.  E.g. we might rewrite ``f x (f y (f z v))`` as::+in a loop.  E.g. we might rewrite ``f x (f y (f z v))`` as:++.. code-block:: futhark    loop acc = v for a in [z,y,x] do     f a acc

@@ -16,7 +16,7 @@ server executables.  A server executable is started like any other executable, and supports-most of the same command line options.+most of the same command line options (:ref:`executable-options`).  Basics ------@@ -33,8 +33,8 @@ command fails, the server will print ``%%% FAILURE`` followed by the error message, and then ``%%% OK`` when it is ready for more input. Some output may also precede ``%%% FAILURE``, e.g. logging statements-that occured before failure was detected.  Fatal errors (that lead to-server shutdown) may be printed to stderr.+that occured before failure was detected.  Fatal errors that lead to+server shutdown may be printed to stderr.  Variables ---------

@@ -333,7 +333,7 @@  Return types follow the rules, with one addition: -* If the return type is an *m*-element tuple, then the the function+* If the return type is an *m*-element tuple, then the function   returns *m* values, mapped according to the rules above (but not   including this one - nested tuples are not mapped directly).  This   rule does not apply when the entry point has been given a return

@@ -1,6 +1,6 @@ cabal-version: 2.4 name:           futhark-version:        0.21.4+version:        0.21.5 synopsis:       An optimising compiler for a functional, array-oriented language.  description:    Futhark is a small programming language designed to be compiled to

@@ -4,8 +4,9 @@ #define CUDA_SUCCEED_NONFATAL(x) cuda_api_succeed_nonfatal(x, #x, __FILE__, __LINE__) #define NVRTC_SUCCEED_FATAL(x) nvrtc_api_succeed_fatal(x, #x, __FILE__, __LINE__) #define NVRTC_SUCCEED_NONFATAL(x) nvrtc_api_succeed_nonfatal(x, #x, __FILE__, __LINE__)--#define SUCCEED_OR_RETURN(serror) {               \+// Take care not to override an existing error.+#define CUDA_SUCCEED_OR_RETURN(e) {               \+    char *serror = CUDA_SUCCEED_NONFATAL(e);      \     if (serror) {                                 \       if (!ctx->error) {                          \         ctx->error = serror;                      \@@ -15,8 +16,6 @@       }                                           \     }                                             \   }--#define CUDA_SUCCEED_OR_RETURN(e) SUCCEED_OR_RETURN(CUDA_SUCCEED_NONFATAL(e))  // CUDA_SUCCEED_OR_RETURN returns the value of the variable 'bad' in // scope.  By default, it will be this one.  Create a local variable

@@ -0,0 +1,3 @@+#define FUTHARK_SUCCESS 0+#define FUTHARK_PROGRAM_ERROR 2+#define FUTHARK_OUT_OF_MEMORY 3

@@ -28,7 +28,7 @@     -- * Construction     module Futhark.IR.Primitive,     NumExp (..),-    IntExp,+    IntExp (..),     FloatExp (..),     sExt,     zExt,@@ -53,6 +53,7 @@     sExt64,     zExt32,     zExt64,+    sExtAs,     fMin64,     fMax64,   )@@ -231,31 +232,37 @@  -- | The class of integer types that can be used for constructing -- 'TPrimExp's.-class NumExp t => IntExp t+class NumExp t => IntExp t where+  -- | The type of an expression, known to be an integer type.+  expIntType :: TPrimExp t v -> IntType  instance NumExp Int8 where   fromInteger' = isInt8 . ValueExp . IntValue . Int8Value . fromInteger   fromBoolExp = isInt8 . ConvOpExp (BToI Int8) . untyped -instance IntExp Int8+instance IntExp Int8 where+  expIntType = const Int8  instance NumExp Int16 where   fromInteger' = isInt16 . ValueExp . IntValue . Int16Value . fromInteger   fromBoolExp = isInt16 . ConvOpExp (BToI Int16) . untyped -instance IntExp Int16+instance IntExp Int16 where+  expIntType = const Int16  instance NumExp Int32 where   fromInteger' = isInt32 . ValueExp . IntValue . Int32Value . fromInteger   fromBoolExp = isInt32 . ConvOpExp (BToI Int32) . untyped -instance IntExp Int32+instance IntExp Int32 where+  expIntType = const Int32  instance NumExp Int64 where   fromInteger' = isInt64 . ValueExp . IntValue . Int64Value . fromInteger   fromBoolExp = isInt64 . ConvOpExp (BToI Int64) . untyped -instance IntExp Int64+instance IntExp Int64 where+  expIntType = const Int64  -- | The class of floating-point types that can be used for -- constructing 'TPrimExp's.@@ -263,6 +270,9 @@   -- | Construct a typed expression from a rational.   fromRational' :: Rational -> TPrimExp t v +  -- | The type of an expression, known to be a floating-point type.+  expFloatType :: TPrimExp t v -> FloatType+ instance NumExp Half where   fromInteger' = isF16 . ValueExp . FloatValue . Float16Value . fromInteger   fromBoolExp = isF16 . ConvOpExp (SIToFP Int16 Float16) . ConvOpExp (BToI Int16) . untyped@@ -277,12 +287,15 @@  instance FloatExp Half where   fromRational' = TPrimExp . ValueExp . FloatValue . Float16Value . fromRational+  expFloatType = const Float16  instance FloatExp Float where   fromRational' = TPrimExp . ValueExp . FloatValue . Float32Value . fromRational+  expFloatType = const Float32  instance FloatExp Double where   fromRational' = TPrimExp . ValueExp . FloatValue . Float64Value . fromRational+  expFloatType = const Float64  instance (NumExp t, Pretty v) => Num (TPrimExp t v) where   TPrimExp x + TPrimExp y@@ -617,6 +630,15 @@ -- | 64-bit float maximum. fMax64 :: TPrimExp Double v -> TPrimExp Double v -> TPrimExp Double v fMax64 x y = TPrimExp $ BinOpExp (FMax Float64) (untyped x) (untyped y)++-- | Convert result of some integer expression to have the same type+-- as another, using sign extension.+sExtAs ::+  (IntExp to, IntExp from) =>+  TPrimExp from v ->+  TPrimExp to v ->+  TPrimExp to v+sExtAs from to = TPrimExp $ sExt (expIntType to) (untyped from)  -- Prettyprinting instances 

@@ -177,7 +177,7 @@  allocateCUDABuffer :: GC.Allocate OpenCL () allocateCUDABuffer mem size tag "device" =-  GC.stm [C.cstm|CUDA_SUCCEED_OR_RETURN(cuda_alloc(&ctx->cuda, (size_t)$exp:size, $exp:tag, &$exp:mem));|]+  GC.stm [C.cstm|ctx->error = CUDA_SUCCEED_NONFATAL(cuda_alloc(&ctx->cuda, (size_t)$exp:size, $exp:tag, &$exp:mem));|] allocateCUDABuffer _ _ _ space =   error $ "Cannot allocate in '" ++ space ++ "' memory space." 

@@ -420,6 +420,8 @@       [C.cedecl|void $id:s(struct $id:ctx* ctx) {                                  $stms:free_fields                                  free_constants(ctx);+                                 cuMemFree(ctx->global_failure);+                                 cuMemFree(ctx->global_failure_args);                                  cuda_cleanup(&ctx->cuda);                                  free_lock(&ctx->lock);                                  ctx->cfg->in_use = 0;@@ -458,7 +460,7 @@                       $stm:(failureSwitch failures) -                     return 1;+                     return FUTHARK_PROGRAM_ERROR;                    }                  }                  CUDA_SUCCEED_OR_RETURN(cuCtxPopCurrent(&ctx->cuda.cu_ctx));

@@ -213,7 +213,7 @@  allocateOpenCLBuffer :: GC.Allocate OpenCL () allocateOpenCLBuffer mem size tag "device" =-  GC.stm [C.cstm|OPENCL_SUCCEED_OR_RETURN(opencl_alloc(&ctx->opencl, (size_t)$exp:size, $exp:tag, &$exp:mem));|]+  GC.stm [C.cstm|ctx->error = OPENCL_SUCCEED_NONFATAL(opencl_alloc(&ctx->opencl, (size_t)$exp:size, $exp:tag, &$exp:mem));|] allocateOpenCLBuffer _ _ _ space =   error $ "Cannot allocate in '" ++ space ++ "' space." 

@@ -450,6 +450,8 @@                                  free_constants(ctx);                                  free_lock(&ctx->lock);                                  $stms:(map releaseKernel (M.toList kernels))+                                 OPENCL_SUCCEED_FATAL(clReleaseMemObject(ctx->global_failure));+                                 OPENCL_SUCCEED_FATAL(clReleaseMemObject(ctx->global_failure_args));                                  teardown_opencl(&ctx->opencl);                                  ctx->cfg->in_use = 0;                                  free(ctx);@@ -492,7 +494,7 @@                     $stm:(failureSwitch failures) -                   return 1;+                   return FUTHARK_PROGRAM_ERROR;                  }                  return 0;                }|]

@@ -96,7 +96,7 @@ import Futhark.CodeGen.Backends.GenericC.Server (serverDefs) import Futhark.CodeGen.Backends.SimpleRep import Futhark.CodeGen.ImpCode-import Futhark.CodeGen.RTS.C (halfH, lockH, timingH, utilH)+import Futhark.CodeGen.RTS.C (errorsH, halfH, lockH, timingH, utilH) import Futhark.IR.Prop (isBuiltInFunction) import qualified Futhark.Manifest as Manifest import Futhark.MonadFreshNames@@ -258,7 +258,7 @@   items     [C.citems|ctx->error = msgprintf($string:formatstr', $args:formatargs, $string:stacktrace);               $items:free_all_mem-              err = 1;+              err = FUTHARK_PROGRAM_ERROR;               goto cleanup;|]  defCall :: CallCompiler op s@@ -673,12 +673,15 @@   }   int ret = $id:(fatMemUnRef space)(ctx, block, desc); -  ctx->$id:usagename += size;+  if (ret != FUTHARK_SUCCESS) {+    return ret;+  }+   if (ctx->detail_memory) {     fprintf(ctx->log, "Allocating %lld bytes for %s in %s (then allocated: %lld bytes)",             (long long) size,             desc, $string:spacedesc,-            (long long) ctx->$id:usagename);+            (long long) ctx->$id:usagename + size);   }   if (ctx->$id:usagename > ctx->$id:peakname) {     ctx->$id:peakname = ctx->$id:usagename;@@ -690,11 +693,28 @@   }    $items:alloc-  block->references = (int*) malloc(sizeof(int));-  *(block->references) = 1;-  block->size = size;-  block->desc = desc;-  return ret;++  if (ctx->error == NULL) {+    block->references = (int*) malloc(sizeof(int));+    *(block->references) = 1;+    block->size = size;+    block->desc = desc;+    ctx->$id:usagename += size;+    return FUTHARK_SUCCESS;+  } else {+    // We are naively assuming that any memory allocation error is due to OOM.+    // We preserve the original error so that a savvy user can perhaps find+    // glory despite our naiveté.++    char *old_error = ctx->error;+    ctx->error = msgprintf("Failed to allocate memory in %s.\nAttempted allocation: %12lld bytes\nCurrently allocated:  %12lld bytes\n%s",+                           $string:spacedesc,+                           (long long) size,+                           (long long) ctx->$id:usagename,+                           old_error);+    free(old_error);+    return FUTHARK_OUT_OF_MEMORY;+  }   }|]    -- Memory setting - unreference the destination and increase the@@ -1066,7 +1086,7 @@             shapearr = "shape_" ++ show i             dims = [[C.cexp|$id:shapearr[$int:j]|] | j <- [0 .. rank - 1]]             num_elems = cproduct dims-        item [C.citem|typename int64_t $id:shapearr[$int:rank];|]+        item [C.citem|typename int64_t $id:shapearr[$int:rank] = {0};|]         stms $ loadValueHeader sign pt rank [C.cexp|$id:shapearr|] [C.cexp|src|]         item [C.citem|const void* $id:dataptr = src;|]         stm [C.cstm|obj->$id:field = NULL;|]@@ -1543,6 +1563,7 @@ // Miscellaneous $miscdecls #define FUTHARK_BACKEND_$backend+$errorsH  #ifdef __cplusplus }

@@ -597,7 +597,7 @@         [C.citems|int $id:ftask_err = scheduler_prepare_task(&ctx->scheduler, &$id:ftask_name);                   if ($id:ftask_err != 0) {                     $items:free_all_mem;-                    err = 1;+                    err = $id:ftask_err;                     goto cleanup;                   }|] @@ -665,7 +665,7 @@       [C.citems|int $id:ftask_err = scheduler_execute_task(&ctx->scheduler,                                                            &$id:ftask_name);                if ($id:ftask_err != 0) {-                 err = 1;+                 err = $id:ftask_err;                  goto cleanup;                }|] 

@@ -23,10 +23,11 @@     compileThreadResult,     compileGroupResult,     virtualiseGroups,-    groupLoop,     kernelLoop,     groupCoverSpace,-    precomputeSegOpIDs,+    Precomputed,+    precomputeConstants,+    precomputedConstants,     atomicUpdateLocking,     AtomicBinOp,     Locking (..),@@ -36,12 +37,14 @@ where  import Control.Monad.Except-import Data.List (foldl', zip4)+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@@ -92,30 +95,64 @@     kernelWaveSize :: Imp.TExp Int32,     kernelThreadActive :: Imp.TExp Bool,     -- | A mapping from dimensions of nested SegOps to already-    -- computed local thread IDs.-    kernelLocalIdMap :: M.Map [SubExp] [Imp.TExp Int32]+    -- 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)   } -segOpSizes :: Stms GPUMem -> S.Set [SubExp]+-- | The sizes of nested iteration spaces in the kernel.+type SegOpSizes = S.Set [SubExp]++-- | Find the sizes of nested parallelism in a 'SegOp' body.+segOpSizes :: Stms GPUMem -> SegOpSizes segOpSizes = onStms   where     onStms = foldMap (onExp . stmExp)     onExp (Op (Inner (SegOp op))) =-      S.singleton $ map snd $ unSegSpace $ segSpace 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 (If _ tbranch fbranch _) =       onStms (bodyStms tbranch) <> onStms (bodyStms fbranch)     onExp (DoLoop _ _ body) =       onStms (bodyStms body)     onExp _ = mempty -precomputeSegOpIDs :: Stms GPUMem -> InKernelGen a -> InKernelGen a-precomputeSegOpIDs stms m = do+-- | 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 (segOpSizes stms))+  new_ids <- M.fromList <$> mapM (mkMap ltid) (S.toList (pcSegOpSizes pre))   let f env =         env           { kernelConstants =-              (kernelConstants env) {kernelLocalIdMap = new_ids}+              (kernelConstants env)+                { kernelLocalIdMap = new_ids,+                  kernelChunkItersMap = pcChunkItersMap pre+                }           }   localEnv f m   where@@ -221,24 +258,24 @@     if n == num_threads       then f tid       else do-        -- Compute how many elements this thread is responsible for.-        -- Formula: (n - tid) / num_threads (rounded up).-        let elems_for_this = (n - tid) `divUp` num_threads--        sFor "i" elems_for_this $ \i -> f $ i * num_threads + tid+        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 ::-  Imp.TExp Int64 ->-  (Imp.TExp Int64 -> InKernelGen ()) ->+  IntExp t =>+  Imp.TExp t ->+  (Imp.TExp t -> InKernelGen ()) ->   InKernelGen () groupLoop n f = do   constants <- kernelConstants <$> askEnv   kernelLoop-    (sExt64 $ kernelLocalThreadId constants)-    (kernelGroupSize constants)+    (kernelLocalThreadId constants `sExtAs` n)+    (kernelGroupSize constants `sExtAs` n)     n     f @@ -246,12 +283,65 @@ -- all threads in the group participate.  The passed-in function is -- invoked with a (symbolic) point in the index space. groupCoverSpace ::-  [Imp.TExp Int64] ->-  ([Imp.TExp Int64] -> InKernelGen ()) ->+  IntExp t =>+  [Imp.TExp t] ->+  ([Imp.TExp t] -> InKernelGen ()) ->   InKernelGen () groupCoverSpace ds f =   groupLoop (product ds) $ f . unflattenIndex ds +localThreadIDs :: [SubExp] -> InKernelGen [Imp.TExp Int64]+localThreadIDs dims = do+  ltid <- sExt64 . kernelLocalThreadId . kernelConstants <$> askEnv+  let dims' = map toInt64Exp 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.@@ -263,9 +353,11 @@ compileGroupExp _ (BasicOp (UpdateAcc acc is vs)) =   updateAcc acc is vs compileGroupExp (Pat [dest]) (BasicOp (Replicate ds se)) = do-  let ds' = map toInt64Exp $ shapeDims ds-  groupCoverSpace ds' $ \is ->-    copyDWIMFix (patElemName dest) is se (drop (shapeRank ds) is)+  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 (Iota n e s it)) = do   n' <- toExp n@@ -305,21 +397,9 @@ sanityCheckLevel SegGroup {} =   error "compileGroupOp: unexpected group-level SegOp." -localThreadIDs :: [SubExp] -> InKernelGen [Imp.TExp Int64]-localThreadIDs dims = do-  ltid <- sExt64 . kernelLocalThreadId . kernelConstants <$> askEnv-  let dims' = map toInt64Exp dims-  maybe (unflattenIndex dims' ltid) (map sExt64)-    . M.lookup dims-    . kernelLocalIdMap-    . kernelConstants-    <$> askEnv--compileGroupSpace :: SegLevel -> SegSpace -> InKernelGen ()-compileGroupSpace lvl space = do+compileFlatId :: SegLevel -> SegSpace -> InKernelGen ()+compileFlatId lvl space = do   sanityCheckLevel lvl-  let (ltids, dims) = unzip $ unSegSpace space-  zipWithM_ dPrimV_ ltids =<< localThreadIDs dims   ltid <- kernelLocalThreadId . kernelConstants <$> askEnv   dPrimV_ (segFlat space) ltid @@ -359,17 +439,6 @@            return (Just l', f l' (Space "local") local_subhistos) -whenActive :: SegLevel -> SegSpace -> InKernelGen () -> InKernelGen ()-whenActive lvl space m-  | SegNoVirtFull <- segVirt lvl = m-  | otherwise = do-    group_size <- kernelGroupSize . kernelConstants <$> askEnv-    -- XXX: the following check is too naive - we should also handle-    -- the multi-dimensional case.-    if [group_size] == map (toInt64Exp . snd) (unSegSpace space)-      then m-      else sWhen (isActive $ unSegSpace space) m- -- Which fence do we need to protect shared access to this memory space? fenceForSpace :: Space -> Imp.Fence fenceForSpace (Space "local") = Imp.FenceLocal@@ -385,26 +454,144 @@         . 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 (DimNew . 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 '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-  void $ compileGroupSpace lvl space+  compileFlatId lvl space -  whenActive lvl space . localOps threadOperations $+  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-  compileGroupSpace lvl space+  compileFlatId lvl space+   let (ltids, dims) = unzip $ unSegSpace space       dims' = map toInt64Exp dims -  whenActive lvl space . localOps threadOperations $+  groupCoverSegSpace (segVirt lvl) space $     compileStms mempty (kernelBodyStms body) $       forM_ (zip (patNames pat) $ kernelBodyResult body) $ \(dest, res) ->         copyDWIMFix@@ -424,41 +611,35 @@   -- array of scan elements, so we invent some new flattened arrays   -- here.   dims_flat <- dPrimV "dims_flat" $ product dims'-  let flattened pe = do-        MemLoc mem _ ixfun <--          entryArrayLoc <$> lookupArray (patElemName pe)-        let pe_t = typeOf pe-            arr_dims = Var (tvVar dims_flat) : drop (length dims') (arrayDims pe_t)-        sArray-          (baseString (patElemName pe) ++ "_flat")-          (elemType pe_t)-          (Shape arr_dims)-          mem-          $ IxFun.reshape ixfun $ map (DimNew . pe64) arr_dims--      num_scan_results = sum $ map (length . segBinOpNeutral) scans--  arrs_flat <- mapM flattened $ take num_scan_results $ patElems pat+  let scan = head scans+      num_scan_results = length $ segBinOpNeutral scan+  arrs_flat <-+    mapM (flattenArray (length dims') dims_flat) $+      take num_scan_results $ patNames pat -  forM_ scans $ \scan -> do-    let scan_op = segBinOpLambda scan-    groupScan (Just crossesSegment) (product dims') (product dims') scan_op arrs_flat+  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-  compileGroupSpace lvl space--  let (ltids, dims) = unzip $ unSegSpace space-      (red_pes, map_pes) =-        splitAt (segBinOpResults ops) $ patElems pat--      dims' = map toInt64Exp dims+  compileFlatId lvl space +  let dims' = map toInt64Exp dims       mkTempArr t =         sAllocArray "red_arr" (elemType t) (Shape dims <> arrayShape t) $ Space "local"    tmp_arrs <- mapM mkTempArr $ concatMap (lambdaReturnType . segBinOpLambda) ops-  let tmps_for_ops = chunks (map (length . segBinOpNeutral) ops) tmp_arrs--  whenActive lvl space . localOps threadOperations $+  groupCoverSegSpace (segVirt lvl) space $     compileStms mempty (kernelBodyStms body) $ do       let (red_res, map_res) =             splitAt (segBinOpResults ops) $ kernelBodyResult body@@ -468,18 +649,72 @@    sOp $ Imp.ErrorSync Imp.FenceLocal -  case dims' of-    -- Nonsegmented case (or rather, a single segment) - this we can-    -- handle directly with a group-level reduction.-    [dim'] -> do-      forM_ (zip ops tmps_for_ops) $ \(op, tmps) ->-        groupReduce (sExt32 dim') (segBinOpLambda op) tmps+  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 tmp_arrs) $ \(pe, arr) ->+      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]-    _ -> do+    --+    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.@@ -488,21 +723,12 @@       -- involve copying anything; merely playing with the index       -- function.       dims_flat <- dPrimV "dims_flat" $ product dims'-      let flatten arr = do-            ArrayEntry arr_loc pt <- lookupArray arr-            let flat_shape =-                  Shape $-                    Var (tvVar dims_flat) :-                    drop (length ltids) (memLocShape arr_loc)-            sArray "red_arr_flat" pt flat_shape (memLocName arr_loc) $-              IxFun.iota $ map pe64 $ shapeDims flat_shape-       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 flatten tmps+        tmps_flat <- mapM (flattenArray (length dims') dims_flat) tmps         groupScan           (Just crossesSegment)           (product dims')@@ -512,7 +738,7 @@        sOp $ Imp.ErrorSync Imp.FenceLocal -      forM_ (zip red_pes tmp_arrs) $ \(pe, arr) ->+      forM_ (zip red_pes $ concat tmps_for_ops) $ \(pe, arr) ->         copyDWIM           (patElemName pe)           []@@ -521,8 +747,8 @@        sOp $ Imp.Barrier Imp.FenceLocal compileGroupOp pat (Inner (SegOp (SegHist lvl space ops _ kbody))) = do-  compileGroupSpace lvl space-  let ltids = map fst $ unSegSpace space+  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@@ -536,7 +762,7 @@   -- Ensure that all locks have been initialised.   sOp $ Imp.Barrier Imp.FenceLocal -  whenActive lvl space . localOps threadOperations $+  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@@ -908,6 +1134,7 @@           (Imp.le32 wave_size)           true           mempty+          mempty    let set_constants = do         dPrim_ global_tid int32@@ -1333,6 +1560,7 @@         (sExt32 (group_size * num_groups))         0         (Imp.le64 thread_gtid .<. kernel_size)+        mempty         mempty,       set_constants     )@@ -1452,11 +1680,11 @@         (sliceMemLoc destloc destslice')         (sliceMemLoc srcloc srcslice')     _ -> do-      groupCoverSpace dims $ \is ->+      groupCoverSpace (map sExt32 dims) $ \is ->         copyElementWise           pt-          (sliceMemLoc destloc (Slice $ map DimFix is))-          (sliceMemLoc srcloc (Slice $ map DimFix is))+          (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

@@ -48,10 +48,11 @@             compileStms mempty (kernelBodyStms kbody) $               zipWithM_ (compileThreadResult space) (patElems pat) $                 kernelBodyResult kbody-    SegGroup {} ->+    SegGroup {} -> do+      pc <- precomputeConstants group_size' $ kernelBodyStms kbody       sKernelGroup "segmap_intragroup" num_groups' group_size' (segFlat space) $ do         let virt_num_groups = sExt32 $ product dims'-        precomputeSegOpIDs (kernelBodyStms kbody) $+        precomputedConstants pc $           virtualiseGroups (segVirt lvl) virt_num_groups $ \group_id -> do             dIndexSpace (zip is dims') $ sExt64 group_id 

@@ -312,29 +312,27 @@   -- Critical section   let try_acquire_lock = do         old <-- (0 :: Imp.TExp Int32)-        sOp $-          Imp.Atomic $-            Imp.AtomicCmpXchg-              int32-              (tvVar old)-              locks'-              (sExt32 <$> locks_offset)-              (tvVar continue)-              (untyped (lockingToLock locking))+        sOp . Imp.Atomic $+          Imp.AtomicCmpXchg+            int32+            (tvVar old)+            locks'+            (sExt32 <$> locks_offset)+            (tvVar continue)+            (untyped (lockingToLock locking))       lock_acquired = tvExp continue       -- Even the releasing is done with an atomic rather than a       -- simple write, for memory coherency reasons.       release_lock = do         old <-- lockingToLock locking-        sOp $-          Imp.Atomic $-            Imp.AtomicCmpXchg-              int32-              (tvVar old)-              locks'-              (sExt32 <$> locks_offset)-              (tvVar continue)-              (untyped (lockingToUnlock locking))+        sOp . Imp.Atomic $+          Imp.AtomicCmpXchg+            int32+            (tvVar old)+            locks'+            (sExt32 <$> locks_offset)+            (tvVar continue)+            (untyped (lockingToUnlock locking))    -- Preparing parameters. It is assumed that the caller has already   -- filled the arr_params. We copy the current value to the

@@ -17,6 +17,7 @@     tuningH,     utilH,     valuesH,+    errorsH,   ) where @@ -100,3 +101,8 @@ valuesH :: T.Text valuesH = $(embedStringFile "rts/c/values.h") {-# NOINLINE valuesH #-}++-- | @rts/c/errors.h@+errorsH :: T.Text+errorsH = $(embedStringFile "rts/c/errors.h")+{-# NOINLINE errorsH #-}

@@ -74,7 +74,7 @@         SegGroup {} -> "group"       virt = case segVirt lvl of         SegNoVirt -> mempty-        SegNoVirtFull -> PP.semi <+> text "full"+        SegNoVirtFull dims -> PP.semi <+> text "full" <+> ppr (segSeqDims dims)         SegVirt -> PP.semi <+> text "virtualise"  instance Engine.Simplifiable SegLevel where

@@ -826,7 +826,8 @@       <*> choice         [ pSemi             *> choice-              [ keyword "full" $> SegOp.SegNoVirtFull,+              [ keyword "full" $> SegOp.SegNoVirtFull+                  <*> (SegOp.SegSeqDims <$> brackets (pInt `sepBy` pComma)),                 keyword "virtualise" $> SegOp.SegVirt               ],           pure SegOp.SegNoVirt

@@ -746,20 +746,20 @@ replaceArrayOps ::   forall rep.   (Buildable rep, BuilderOps rep, HasSOAC rep) =>-  M.Map ArrayOp ArrayOp ->+  M.Map (AST.Pat rep, ArrayOp) ArrayOp ->   AST.Body rep ->   AST.Body rep replaceArrayOps substs (Body _ stms res) =   mkBody (fmap onStm stms) res   where     onStm (Let pat aux e) =-      let (cs', e') = onExp (stmAuxCerts aux) e+      let (cs', e') = onExp pat (stmAuxCerts aux) e        in certify cs' $ mkLet' (patIdents pat) aux e'-    onExp cs e+    onExp pat cs e       | Just op <- isArrayOp cs e,-        Just op' <- M.lookup op substs =+        Just op' <- M.lookup (pat, op) substs =         fromArrayOp op'-    onExp cs e = (cs, mapExp mapper e)+    onExp _ cs e = (cs, mapExp mapper e)     mapper =       identityMapper         { mapOnBody = const $ return . replaceArrayOps substs,@@ -792,11 +792,11 @@   forall rep.   (Buildable rep, BuilderOps rep, HasSOAC rep) =>   TopDownRuleOp rep-simplifyMapIota vtable pat aux op+simplifyMapIota vtable screma_pat aux op   | Just (Screma w arrs (ScremaForm scan reduce map_lam) :: SOAC rep) <- asSOAC op,     Just (p, _) <- find isIota (zip (lambdaParams map_lam) arrs),     indexings <--      mapMaybe (indexesWith (paramName p) . snd) . S.toList $+      mapMaybe (indexesWith (paramName p)) . S.toList $         arrayOps $ lambdaBody map_lam,     not $ null indexings = Simplify $ do     -- For each indexing with iota, add the corresponding array to@@ -810,7 +810,7 @@               lambdaBody = replaceArrayOps substs $ lambdaBody map_lam             } -    auxing aux . letBind pat . Op . soacOp $+    auxing aux . letBind screma_pat . Op . soacOp $       Screma w (arrs <> more_arrs) (ScremaForm scan reduce map_lam')   where     isIota (_, arr) = case ST.lookupBasicOp arr vtable of@@ -825,12 +825,12 @@       | otherwise = (j :) <$> fixWith i slice     fixWith _ _ = Nothing -    indexesWith v idx@(ArrayIndexing cs arr (Slice js))+    indexesWith v (pat, idx@(ArrayIndexing cs arr (Slice js)))       | arr `ST.elem` vtable,         all (`ST.elem` vtable) $ unCerts cs,         Just js' <- fixWith v js,         all (`ST.elem` vtable) $ namesToList $ freeIn js' =-        Just (js', idx)+        Just (pat, js', idx)     indexesWith _ _ = Nothing      properArr [] arr = pure arr@@ -838,7 +838,7 @@       arr_t <- lookupType arr       letExp (baseString arr) $ BasicOp $ Index arr $ fullSlice arr_t $ map DimFix js -    mapOverArr w (js, ArrayIndexing cs arr slice) = do+    mapOverArr w (pat, js, ArrayIndexing cs arr slice) = do       arr' <- properArr js arr       arr_t <- lookupType arr'       arr'' <-@@ -852,7 +852,7 @@         Just           ( arr'',             arr_elem_param,-            ( ArrayIndexing cs arr slice,+            ( (pat, ArrayIndexing cs arr slice),               ArrayIndexing cs (paramName arr_elem_param) (Slice (drop (length js + 1) (unSlice slice)))             )           )@@ -864,8 +864,8 @@ -- corresponding to that transformation performed on the rows of the -- full array. moveTransformToInput :: TopDownRuleOp (Wise SOACS)-moveTransformToInput vtable pat aux soac@(Screma w arrs (ScremaForm scan reduce map_lam))-  | ops <- map snd $ filter arrayIsMapParam $ S.toList $ arrayOps $ lambdaBody map_lam,+moveTransformToInput vtable screma_pat aux soac@(Screma w arrs (ScremaForm scan reduce map_lam))+  | ops <- filter arrayIsMapParam $ S.toList $ arrayOps $ lambdaBody map_lam,     not $ null ops = Simplify $ do     (more_arrs, more_params, replacements) <-       unzip3 . catMaybes <$> mapM mapOverArr ops@@ -879,7 +879,7 @@             }      auxing aux $-      letBind pat $ Op $ Screma w (arrs <> more_arrs) (ScremaForm scan reduce map_lam')+      letBind screma_pat $ Op $ Screma w (arrs <> more_arrs) (ScremaForm scan reduce map_lam')   where     -- It is not safe to move the transform if the root array is being     -- consumed by the Screma.  This is a bit too conservative - it's@@ -913,7 +913,7 @@     arrayIsMapParam (_, ArrayVar {}) =       False -    mapOverArr op+    mapOverArr (pat, op)       | Just (_, arr) <- find ((== arrayOpArr op) . fst) (zip map_param_names arrs),         not $ arr `nameIn` consumed = do         arr_t <- lookupType arr@@ -937,7 +937,7 @@           Just             ( arr_transformed,               Param mempty arr_transformed_row (rowType arr_transformed_t),-              (op, ArrayVar mempty arr_transformed_row)+              ((pat, op), ArrayVar mempty arr_transformed_row)             )     mapOverArr _ = return Nothing moveTransformToInput _ _ _ _ =

@@ -13,6 +13,7 @@ module Futhark.IR.SegOp   ( SegOp (..),     SegVirt (..),+    SegSeqDims (..),     segLevel,     segBody,     segSpace,@@ -494,6 +495,24 @@       onDim (dim, blk_tile, reg_tile) =         ppr dim <+> "/" <+> parens (ppr blk_tile <+> "*" <+> ppr reg_tile) +-- | These dimensions (indexed from 0, outermost) of the corresponding+-- 'SegSpace' should not be parallelised, but instead iterated+-- sequentially.  For example, with a 'SegSeqDims' of @[0]@ and a+-- 'SegSpace' with dimensions @[n][m]@, there will be an outer loop+-- with @n@ iterations, while the @m@ dimension will be parallelised.+--+-- Semantically, this has no effect, but it may allow reductions in+-- memory usage or other low-level optimisations.  Operationally, the+-- guarantee is that for a SegSeqDims of e.g. @[i,j,k]@, threads+-- running at any given moment will always have the same indexes along+-- the dimensions specified by @[i,j,k]@.+--+-- At the moment, this is only supported for 'SegNoVirtFull'+-- intra-group parallelism in GPU code, as we have not yet found it+-- useful anywhere else.+newtype SegSeqDims = SegSeqDims {segSeqDims :: [Int]}+  deriving (Eq, Ord, Show)+ -- | Do we need group-virtualisation when generating code for the -- segmented operation?  In most cases, we do, but for some simple -- kernels, we compute the full number of groups in advance, and then@@ -507,7 +526,7 @@   | -- | Not only do we not need virtualisation, but we _guarantee_     -- that all physical threads participate in the work.  This can     -- save some checks in code generation.-    SegNoVirtFull+    SegNoVirtFull SegSeqDims   deriving (Eq, Ord, Show)  -- | Index space of a 'SegOp'.

@@ -126,7 +126,7 @@               map snd rem_outer_dims_rev       grid_size <- letSubExp "grid_size" =<< toExp grid_pexp       group_size <- letSubExp "group_size" =<< toExp (pe64 ty * pe64 tx)-      let segthd_lvl = SegThread (Count grid_size) (Count group_size) SegNoVirtFull+      let segthd_lvl = SegThread (Count grid_size) (Count group_size) (SegNoVirtFull (SegSeqDims []))        gid_x <- newVName "gid_x"       gid_y <- newVName "gid_y"@@ -151,90 +151,90 @@         a_loc_init <- scratch "A_loc" map_t1 [a_loc_sz]         b_loc_init <- scratch "B_loc" map_t2 [b_loc_sz] -        let kkLoopBody kk0 (thd_res_merge, a_loc_init', b_loc_init') epilogue = do+        let kkLoopBody tkind kk0 (thd_res_merge, a_loc_init', b_loc_init') epilogue = do               kk <- letExp "kk" =<< toExp (le64 kk0 * pe64 tk)-              a_loc <- forLoop ry [a_loc_init'] $ \i0 [a_loc_merge] -> do-                loop_a_loc <- forLoop tk_div_tx [a_loc_merge] $ \k0 [a_loc_merge'] -> do-                  scatter_a_loc <- segScatter2D "A_glb2loc" a_loc_sz a_loc_merge' segthd_lvl (ty, tx) $-                    \(thd_y, thd_x) -> do-                      k <- letExp "k" =<< toExp (le64 thd_x + le64 k0 * pe64 tx)-                      i <- letExp "i" =<< toExp (le64 thd_y + le64 i0 * pe64 ty)+              a_loc <- segScatter2D+                "A_glb2loc"+                a_loc_sz+                a_loc_init'+                segthd_lvl+                [ry, tk_div_tx]+                (ty, tx)+                $ \[i0, k0] (thd_y, thd_x) -> do+                  k <- letExp "k" =<< toExp (le64 thd_x + le64 k0 * pe64 tx)+                  i <- letExp "i" =<< toExp (le64 thd_y + le64 i0 * pe64 ty) -                      letBindNames [gtid_y] =<< toExp (le64 iii + le64 i)-                      a_col_idx <- letExp "A_col_idx" =<< toExp (le64 kk + le64 k)+                  letBindNames [gtid_y] =<< toExp (le64 iii + le64 i)+                  a_col_idx <- letExp "A_col_idx" =<< toExp (le64 kk + le64 k) -                      a_elem <--                        letSubExp "A_elem"-                          =<< eIf-                            ( toExp $-                                le64 gtid_y .<. pe64 height_A-                                  .&&. if epilogue-                                    then le64 a_col_idx .<. pe64 common_dim-                                    else true-                            )-                            ( do-                                addStm load_A-                                res <- index "A_elem" inp_A [a_col_idx]-                                resultBodyM [Var res]-                            )-                            (eBody [eBlank $ Prim map_t1])-                      a_loc_ind <--                        letSubExp "a_loc_ind"-                          =<< eIf-                            (toExp $ le64 k .<. pe64 tk)-                            ( toExp (le64 k + le64 i * pe64 tk)-                                >>= letTupExp' "loc_fi"-                                >>= resultBodyM-                            )-                            (eBody [pure $ BasicOp $ SubExp $ intConst Int64 (-1)])-                      return (a_elem, a_loc_ind)-                  resultBodyM $ map Var scatter_a_loc-                resultBodyM [Var loop_a_loc]+                  a_elem <-+                    letSubExp "A_elem"+                      =<< eIf+                        ( toExp $+                            le64 gtid_y .<. pe64 height_A+                              .&&. if epilogue+                                then le64 a_col_idx .<. pe64 common_dim+                                else true+                        )+                        ( do+                            addStm load_A+                            res <- index "A_elem" inp_A [a_col_idx]+                            resultBodyM [Var res]+                        )+                        (eBody [eBlank $ Prim map_t1])+                  a_loc_ind <-+                    letSubExp "a_loc_ind"+                      =<< eIf+                        (toExp $ le64 k .<. pe64 tk)+                        ( toExp (le64 k + le64 i * pe64 tk)+                            >>= letTupExp' "loc_fi"+                            >>= resultBodyM+                        )+                        (eBody [pure $ BasicOp $ SubExp $ intConst Int64 (-1)])+                  return (a_elem, a_loc_ind)                -- copy B from global to shared memory-              b_loc <- forLoop tk_div_ty [b_loc_init'] $ \k0 [b_loc_merge] -> do-                loop_b_loc <- forLoop rx [b_loc_merge] $ \j0 [b_loc_merge'] -> do-                  scatter_b_loc <- segScatter2D-                    "B_glb2loc"-                    b_loc_sz-                    b_loc_merge'-                    segthd_lvl-                    (ty, tx)-                    $ \(thd_y, thd_x) -> do-                      k <- letExp "k" =<< toExp (le64 thd_y + le64 k0 * pe64 ty)-                      j <- letExp "j" =<< toExp (le64 thd_x + le64 j0 * pe64 tx)+              b_loc <- segScatter2D+                "B_glb2loc"+                b_loc_sz+                b_loc_init'+                segthd_lvl+                [tk_div_ty, rx]+                (ty, tx)+                $ \[k0, j0] (thd_y, thd_x) ->+                  do+                    k <- letExp "k" =<< toExp (le64 thd_y + le64 k0 * pe64 ty)+                    j <- letExp "j" =<< toExp (le64 thd_x + le64 j0 * pe64 tx) -                      letBindNames [gtid_x] =<< toExp (le64 jjj + le64 j)-                      b_row_idx <- letExp "B_row_idx" =<< toExp (le64 kk + le64 k)+                    letBindNames [gtid_x] =<< toExp (le64 jjj + le64 j)+                    b_row_idx <- letExp "B_row_idx" =<< toExp (le64 kk + le64 k) -                      b_elem <--                        letSubExp "B_elem"-                          =<< eIf-                            ( toExp $-                                le64 gtid_x .<. pe64 width_B-                                  .&&. if epilogue-                                    then le64 b_row_idx .<. pe64 common_dim-                                    else true-                            )-                            ( do-                                addStm load_B-                                res <- index "B_elem" inp_B [b_row_idx]-                                resultBodyM [Var res]-                            )-                            (eBody [eBlank $ Prim map_t2])+                    b_elem <-+                      letSubExp "B_elem"+                        =<< eIf+                          ( toExp $+                              le64 gtid_x .<. pe64 width_B+                                .&&. if epilogue+                                  then le64 b_row_idx .<. pe64 common_dim+                                  else true+                          )+                          ( do+                              addStm load_B+                              res <- index "B_elem" inp_B [b_row_idx]+                              resultBodyM [Var res]+                          )+                          (eBody [eBlank $ Prim map_t2]) -                      b_loc_ind <--                        letSubExp "b_loc_ind"-                          =<< eIf-                            (toExp $ le64 k .<. pe64 tk)-                            ( toExp (le64 j + le64 k * pe64 tx_rx)-                                >>= letTupExp' "loc_fi"-                                >>= resultBodyM-                            )-                            (eBody [pure $ BasicOp $ SubExp $ intConst Int64 (-1)])-                      return (b_elem, b_loc_ind)-                  resultBodyM $ map Var scatter_b_loc-                resultBodyM [Var loop_b_loc]+                    b_loc_ind <-+                      letSubExp "b_loc_ind"+                        =<< eIf+                          (toExp $ le64 k .<. pe64 tk)+                          ( toExp (le64 j + le64 k * pe64 tx_rx)+                              >>= letTupExp' "loc_fi"+                              >>= resultBodyM+                          )+                          (eBody [pure $ BasicOp $ SubExp $ intConst Int64 (-1)])+                    return (b_elem, b_loc_ind)                -- inner loop updating this thread's accumulator (loop k in mmm_kernels).               thd_acc <- forLoop tk [thd_res_merge] $ \k [acc_merge] ->@@ -291,15 +291,18 @@                             css_init <- index "css_init" acc_merge [ltid_y, ltid_x]                              css <- forLoop ry [css_init] $ \i [css_merge] -> do-                              css <- forLoop rx [css_merge] $ \j [css_merge'] ->+                              css <- forLoop rx [css_merge] $ \j [css_merge'] -> do+                                let cond =+                                      toExp $ case tkind of+                                        TileFull -> true+                                        TilePartial ->+                                          le64 iii + le64 i + pe64 ry * le64 ltid_y+                                            .<. pe64 height_A+                                              .&&. le64 jjj + le64 j + pe64 rx * le64 ltid_x+                                            .<. pe64 width_B                                 resultBodyM =<< letTupExp' "foo"                                   =<< eIf-                                    ( toExp $-                                        le64 iii + le64 i + pe64 ry * le64 ltid_y-                                          .<. pe64 height_A-                                            .&&. le64 jjj + le64 j + pe64 rx * le64 ltid_x-                                          .<. pe64 width_B-                                    )+                                    cond                                     ( do                                         a <- index "a" as [i]                                         b <- index "b" bs [j]@@ -339,14 +342,14 @@           forLoop' (Var full_tiles) [cssss, a_loc_init, b_loc_init] $             \kk0 [thd_res_merge, a_loc_merge, b_loc_merge] -> do               process_full_tiles <--                kkLoopBody kk0 (thd_res_merge, a_loc_merge, b_loc_merge) False+                kkLoopBody TileFull kk0 (thd_res_merge, a_loc_merge, b_loc_merge) False                resultBodyM $ map Var process_full_tiles          let prologue_res : a_loc_reuse : b_loc_reuse : _ = prologue_res_list          -- build epilogue.-        epilogue_res_list <- kkLoopBody full_tiles (prologue_res, a_loc_reuse, b_loc_reuse) True+        epilogue_res_list <- kkLoopBody TilePartial full_tiles (prologue_res, a_loc_reuse, b_loc_reuse) True          let redomap_res : _ = epilogue_res_list @@ -779,7 +782,7 @@       let grid_pexp = product $ gridxyz_pexp : map (pe64 . snd) rem_outer_dims_rev       grid_size <- letSubExp "grid_size_tile3d" =<< toExp grid_pexp       group_size <- letSubExp "group_size_tile3d" =<< toExp (pe64 ty * pe64 tx)-      let segthd_lvl = SegThread (Count grid_size) (Count group_size) SegNoVirtFull+      let segthd_lvl = SegThread (Count grid_size) (Count group_size) (SegNoVirtFull (SegSeqDims []))        count_shmem <- letSubExp "count_shmem" =<< ceilDiv rz group_size 

@@ -248,11 +248,9 @@     mkMerge summary       | Just (update, mergename, mergedec) <- relatedUpdate summary = do         source <- newVName "modified_source"+        precopy <- newVName $ baseString (updateValue update) <> "_precopy"         let source_t = snd $ updateType update-            elmident =-              Ident-                (updateValue update)-                (source_t `setArrayDims` sliceDims (updateIndices update))+            elm_t = source_t `setArrayDims` sliceDims (updateIndices update)         tell           ( [ mkLet [Ident source source_t] . BasicOp $                 Update@@ -261,10 +259,11 @@                   (fullSlice source_t $ unSlice $ updateIndices update)                   $ snd $ mergeParam summary             ],-            [ mkLet [elmident] . BasicOp $+            [ mkLet [Ident precopy elm_t] . BasicOp $                 Index                   (updateName update)-                  (fullSlice source_t $ unSlice $ updateIndices update)+                  (fullSlice source_t $ unSlice $ updateIndices update),+              mkLet [Ident (updateValue update) elm_t] $ BasicOp $ Copy precopy             ]           )         return $

@@ -723,8 +723,6 @@       -- Create a SegMap that takes care of the postlude for every thread.       postludeGeneric tiling privstms pat accs' poststms poststms_res res_ts -data TileKind = TilePartial | TileFull- mkReadPreludeValues :: [VName] -> [VName] -> ReadPrelude mkReadPreludeValues prestms_live_arrs prestms_live slice =   fmap mconcat $@@ -1012,7 +1010,7 @@   fmap (inputsToTiles inputs)     . segMap2D       "full_tile"-      (SegThread num_groups group_size SegNoVirtFull)+      (SegThread num_groups group_size (SegNoVirtFull (SegSeqDims [])))       ResultNoSimplify       (tile_size, tile_size)     $ \(ltid_x, ltid_y) -> do@@ -1092,7 +1090,7 @@    segMap2D     "acc"-    (SegThread num_groups group_size SegNoVirtFull)+    (SegThread num_groups group_size (SegNoVirtFull (SegSeqDims [])))     ResultPrivate     (tile_size, tile_size)     $ \(ltid_x, ltid_y) -> do@@ -1227,7 +1225,7 @@         (num_groups_y : map snd dims_on_top)    gid_flat <- newVName "gid_flat"-  let lvl = SegGroup (Count num_groups) (Count group_size) SegNoVirtFull+  let lvl = SegGroup (Count num_groups) (Count group_size) (SegNoVirtFull (SegSeqDims []))       space =         SegSpace gid_flat $           dims_on_top ++ [(gid_x, num_groups_x), (gid_y, num_groups_y)]

@@ -7,6 +7,7 @@     VarianceTable,     varianceInStms,     isTileableRedomap,+    TileKind (..),   ) where @@ -99,24 +100,36 @@   SubExp -> -- arr_size   VName ->   SegLevel -> -- lvl+  [SubExp] -> -- dims of sequential loop on top   (SubExp, SubExp) -> -- (dim_y, dim_x)-  ((VName, VName) -> Builder GPU (SubExp, SubExp)) -> -- f-  Builder GPU [VName]-segScatter2D desc arr_size updt_arr lvl (dim_x, dim_y) f = do+  ([VName] -> (VName, VName) -> Builder GPU (SubExp, SubExp)) -> -- f+  Builder GPU VName+segScatter2D desc arr_size updt_arr lvl seq_dims (dim_x, dim_y) f = do   ltid_x <- newVName "ltid_x"   ltid_y <- newVName "ltid_y"   ltid_flat <- newVName "ltid_flat"-  let segspace = SegSpace ltid_flat [(ltid_x, dim_x), (ltid_y, dim_y)] +  seq_is <- replicateM (length seq_dims) (newVName "ltid_seq")+  let seq_space = zip seq_is seq_dims++  let segspace = SegSpace ltid_flat $ seq_space ++ [(ltid_x, dim_x), (ltid_y, dim_y)]+      lvl' =+        SegThread+          (segNumGroups lvl)+          (segGroupSize lvl)+          (SegNoVirtFull (SegSeqDims [0 .. length seq_dims -1]))+   ((t_v, res_v, res_i), stms) <- runBuilder $ do-    (res_v, res_i) <- f (ltid_x, ltid_y)+    (res_v, res_i) <-+      localScope (scopeOfSegSpace segspace) $+        f seq_is (ltid_x, ltid_y)     t_v <- subExpType res_v     return (t_v, res_v, res_i)    let ret = WriteReturns mempty (Shape [arr_size]) updt_arr [(Slice [DimFix res_i], res_v)]   let body = KernelBody () stms [ret] -  letTupExp desc <=< renameExp $ Op $ SegOp $ SegMap lvl segspace [t_v] body+  letExp desc <=< renameExp $ Op $ SegOp $ SegMap lvl' segspace [t_v] body  -- | The variance table keeps a mapping from a variable name -- (something produced by a 'Stm') to the kernel thread indices@@ -181,3 +194,6 @@  varianceInStms :: VarianceTable -> Stms GPU -> VarianceTable varianceInStms = foldl' varianceInStm++-- | Are we working with full or partial tiles?+data TileKind = TilePartial | TileFull

@@ -445,8 +445,6 @@            { TEUnique $2 (srcspan $1 $>) }          | '[' DimExp ']' TypeExpTerm %prec indexprec            { TEArray $4 $2 (srcspan $1 $>) }-         | '['  ']' TypeExpTerm %prec indexprec-           { TEArray $3 DimExpAny (srcspan $1 $>) }          | TypeExpApply %prec sumprec { $1 }           -- Errors@@ -498,7 +496,6 @@  TypeArg :: { TypeArgExp Name }          : '[' DimExp ']' { TypeArgExpDim $2 (srcspan $1 $>) }-         | '[' ']'         { TypeArgExpDim DimExpAny (srcspan $1 $>) }          | TypeExpAtom     { TypeArgExpType $1 }  FieldType :: { (Name, UncheckedTypeExp) }@@ -518,6 +515,8 @@         | intlit           { let L loc (INTLIT n) = $1             in DimExpConst (fromIntegral n) loc }+        |+          { DimExpAny }  FunParam :: { PatBase NoInfo Name } FunParam : InnerPat { $1 }