futhark 0.10.2 → 0.11.1

raw patch · 120 files changed

+4417/−3837 lines, 120 filesdep ~aesonPVP ok

version bump matches the API change (PVP)

Dependency ranges changed: aeson

API changes (from Hackage documentation)

Files

• futhark.cabal +73/−72
• futlib/math.fut +6/−0
• futlib/prelude.fut +3/−3
• futlib/soacs.fut +11/−10
• rts/c/cuda.h +9/−6
• rts/c/free_list.h +3/−2
• rts/c/lock.h +5/−1
• rts/c/opencl.h +20/−9
• rts/c/panic.h +4/−2
• rts/c/timing.h +5/−4
• rts/c/tuning.h +3/−3
• rts/c/values.h +4/−0
• rts/csharp/opencl.cs +13/−1
• rts/csharp/scalar.cs +4/−0
• rts/python/memory.py +3/−1
• rts/python/opencl.py +1/−0
• rts/python/panic.py +4/−0
• rts/python/scalar.py +16/−1
• rts/python/tuning.py +2/−4
• rts/python/values.py +62/−19
• src/Futhark/Analysis/Alias.hs +0/−1
• src/Futhark/Analysis/DataDependencies.hs +1/−1
• src/Futhark/Analysis/HORepresentation/MapNest.hs +14/−35
• src/Futhark/Analysis/HORepresentation/SOAC.hs +4/−2
• src/Futhark/Analysis/PrimExp/Convert.hs +25/−16
• src/Futhark/Analysis/Range.hs +8/−9
• src/Futhark/Analysis/ScalExp.hs +2/−0
• src/Futhark/Analysis/SymbolTable.hs +3/−2
• src/Futhark/Analysis/Usage.hs +5/−30
• src/Futhark/Analysis/UsageTable.hs +0/−4
• src/Futhark/CLI/CUDA.hs +1/−1
• src/Futhark/CLI/OpenCL.hs +1/−1
• src/Futhark/CodeGen/Backends/CCUDA.hs +4/−6
• src/Futhark/CodeGen/Backends/CCUDA/Boilerplate.hs +4/−4
• src/Futhark/CodeGen/Backends/COpenCL.hs +4/−9
• src/Futhark/CodeGen/Backends/COpenCL/Boilerplate.hs +5/−5
• src/Futhark/CodeGen/Backends/CSOpenCL.hs +10/−13
• src/Futhark/CodeGen/Backends/GenericC.hs +55/−49
• src/Futhark/CodeGen/Backends/GenericCSharp.hs +14/−19
• src/Futhark/CodeGen/Backends/GenericPython.hs +17/−22
• src/Futhark/CodeGen/Backends/PyOpenCL.hs +5/−12
• src/Futhark/CodeGen/Backends/SequentialC.hs +2/−2
• src/Futhark/CodeGen/Backends/SimpleRepresentation.hs +25/−7
• src/Futhark/CodeGen/ImpCode.hs +23/−15
• src/Futhark/CodeGen/ImpCode/Kernels.hs +38/−36
• src/Futhark/CodeGen/ImpGen.hs +190/−199
• src/Futhark/CodeGen/ImpGen/Kernels.hs +58/−157
• src/Futhark/CodeGen/ImpGen/Kernels/Base.hs +369/−242
• src/Futhark/CodeGen/ImpGen/Kernels/SegGenRed.hs +529/−139
• src/Futhark/CodeGen/ImpGen/Kernels/SegMap.hs +26/−0
• src/Futhark/CodeGen/ImpGen/Kernels/SegRed.hs +358/−268
• src/Futhark/CodeGen/ImpGen/Kernels/SegScan.hs +258/−0
• src/Futhark/CodeGen/ImpGen/Kernels/ToOpenCL.hs +84/−70
• src/Futhark/CodeGen/ImpGen/Kernels/Transpose.hs +15/−16
• src/Futhark/CodeGen/ImpGen/Sequential.hs +1/−1
• src/Futhark/CodeGen/SetDefaultSpace.hs +4/−4
• src/Futhark/Construct.hs +1/−1
• src/Futhark/Internalise.hs +60/−27
• src/Futhark/Internalise/AccurateSizes.hs +1/−1
• src/Futhark/Internalise/Bindings.hs +11/−13
• src/Futhark/Internalise/Defunctionalise.hs +85/−73
• src/Futhark/Internalise/Lambdas.hs +18/−9
• src/Futhark/Internalise/Monomorphise.hs +13/−13
• src/Futhark/Optimise/DoubleBuffer.hs +27/−10
• src/Futhark/Optimise/Fusion.hs +44/−53
• src/Futhark/Optimise/Fusion/LoopKernel.hs +14/−17
• src/Futhark/Optimise/InPlaceLowering.hs +10/−6
• src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs +18/−11
• src/Futhark/Optimise/Simplify/Engine.hs +7/−11
• src/Futhark/Optimise/Simplify/Lore.hs +2/−4
• src/Futhark/Optimise/Simplify/Rules.hs +9/−15
• src/Futhark/Optimise/TileLoops.hs +3/−1
• src/Futhark/Optimise/TileLoops/RegTiling3D.hs +31/−33
• src/Futhark/Optimise/Unstream.hs +11/−5
• src/Futhark/Pass/ExpandAllocations.hs +66/−70
• src/Futhark/Pass/ExplicitAllocations.hs +128/−115
• src/Futhark/Pass/ExtractKernels.hs +80/−78
• src/Futhark/Pass/ExtractKernels/BlockedKernel.hs +195/−745
• src/Futhark/Pass/ExtractKernels/Distribution.hs +12/−15
• src/Futhark/Pass/ExtractKernels/ISRWIM.hs +2/−2
• src/Futhark/Pass/ExtractKernels/Interchange.hs +1/−1
• src/Futhark/Pass/ExtractKernels/Intragroup.hs +22/−33
• src/Futhark/Pass/ExtractKernels/Kernelise.hs +2/−1
• src/Futhark/Pass/ExtractKernels/Segmented.hs +0/−89
• src/Futhark/Pass/KernelBabysitting.hs +37/−26
• src/Futhark/Pass/ResolveAssertions.hs +1/−2
• src/Futhark/Representation/AST/Attributes/Names.hs +51/−67
• src/Futhark/Representation/AST/Attributes/TypeOf.hs +0/−12
• src/Futhark/Representation/AST/Attributes/Types.hs +10/−12
• src/Futhark/Representation/AST/Pretty.hs +6/−6
• src/Futhark/Representation/AST/Syntax/Core.hs +1/−1
• src/Futhark/Representation/AST/Traversals.hs +2/−18
• src/Futhark/Representation/Aliases.hs +5/−6
• src/Futhark/Representation/ExplicitMemory.hs +69/−96
• src/Futhark/Representation/ExplicitMemory/IndexFunction.hs +9/−10
• src/Futhark/Representation/ExplicitMemory/Simplify.hs +15/−23
• src/Futhark/Representation/Kernels/Kernel.hs +226/−183
• src/Futhark/Representation/Kernels/KernelExp.hs +6/−17
• src/Futhark/Representation/Kernels/Simplify.hs +116/−51
• src/Futhark/Representation/Kernels/Sizes.hs +4/−0
• src/Futhark/Representation/Primitive.hs +3/−2
• src/Futhark/Representation/SOACS/SOAC.hs +126/−91
• src/Futhark/Representation/SOACS/Simplify.hs +241/−27
• src/Futhark/Tools.hs +7/−8
• src/Futhark/Transform/FirstOrderTransform.hs +4/−4
• src/Futhark/Transform/Rename.hs +1/−2
• src/Futhark/Transform/Substitute.hs +2/−3
• src/Futhark/Util.hs +22/−0
• src/Language/Futhark/Attributes.hs +10/−30
• src/Language/Futhark/Interpreter.hs +49/−32
• src/Language/Futhark/Parser/Parser.y +7/−12
• src/Language/Futhark/Pretty.hs +17/−15
• src/Language/Futhark/Syntax.hs +6/−7
• src/Language/Futhark/Traversals.hs +6/−7
• src/Language/Futhark/TypeChecker.hs +9/−5
• src/Language/Futhark/TypeChecker/Monad.hs +9/−0
• src/Language/Futhark/TypeChecker/Terms.hs +32/−44
• src/Language/Futhark/TypeChecker/Types.hs +1/−10
• unittests/Futhark/Representation/ExplicitMemory/IndexFunction/Alg.hs +19/−19
• unittests/Futhark/Representation/ExplicitMemory/IndexFunctionTests.hs +2/−2

@@ -2,10 +2,10 @@ -- -- see: https://github.com/sol/hpack ----- hash: befc35fd82c38f76939296d95942a59561978305f93fb1599eb91ea20a6af5fe+-- hash: eb85e7d5c7284742fdf0fda53a17074c98dc9526652f211fdabddb727e2beffa  name:           futhark-version:        0.10.2+version:        0.11.1 synopsis:       An optimising compiler for a functional, array-oriented language. description:    Futhark is a small programming language designed to be compiled to                 efficient parallel code. It is a statically typed, data-parallel,@@ -59,58 +59,6 @@   location: https://github.com/diku-dk/futhark  library-  hs-source-dirs:-      src-  ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists-  build-depends:-      aeson-    , ansi-terminal >=0.6.3.1-    , array >=0.4-    , base >=4 && <5-    , binary >=0.8.3-    , blaze-html >=0.9.0.1-    , bytestring >=0.10.8-    , containers >=0.5-    , data-binary-ieee754 >=0.1-    , directory >=1.3.0.0-    , directory-tree >=0.12.1-    , dlist >=0.6.0.1-    , file-embed >=0.0.9-    , filepath >=1.4.1.1-    , free >=4.12.4-    , gitrev >=1.2.0-    , haskeline-    , http-client >=0.5.7.0-    , http-client-tls >=0.3.5.1-    , http-conduit >=2.2.4-    , language-c-quote >=0.12-    , mainland-pretty >=0.6.1-    , markdown >=0.1.16-    , megaparsec >=7.0.1-    , mtl >=2.2.1-    , neat-interpolation >=0.3-    , parallel >=3.2.1.0-    , parser-combinators >=1.0.0-    , process >=1.4.3.0-    , process-extras >=0.7.2-    , random-    , regex-tdfa >=1.2-    , srcloc >=0.4-    , template-haskell >=2.11.1-    , temporary-    , terminal-size >=0.3-    , text >=1.2.2.2-    , time >=1.6.0.1-    , transformers >=0.3-    , utf8-string >=1-    , vector >=0.12-    , vector-binary-instances >=0.2.2.0-    , versions >=3.3.1-    , zip-archive >=0.3.1.1-    , zlib >=0.6.1.2-  build-tools:-      alex-    , happy   exposed-modules:       Futhark.Actions       Futhark.Analysis.AlgSimplify@@ -179,7 +127,9 @@       Futhark.CodeGen.ImpGen.Kernels       Futhark.CodeGen.ImpGen.Kernels.Base       Futhark.CodeGen.ImpGen.Kernels.SegGenRed+      Futhark.CodeGen.ImpGen.Kernels.SegMap       Futhark.CodeGen.ImpGen.Kernels.SegRed+      Futhark.CodeGen.ImpGen.Kernels.SegScan       Futhark.CodeGen.ImpGen.Kernels.ToOpenCL       Futhark.CodeGen.ImpGen.Kernels.Transpose       Futhark.CodeGen.ImpGen.OpenCL@@ -233,7 +183,6 @@       Futhark.Pass.ExtractKernels.Intragroup       Futhark.Pass.ExtractKernels.ISRWIM       Futhark.Pass.ExtractKernels.Kernelise-      Futhark.Pass.ExtractKernels.Segmented       Futhark.Pass.FirstOrderTransform       Futhark.Pass.KernelBabysitting       Futhark.Pass.ResolveAssertions@@ -310,37 +259,74 @@       Language.Futhark.Parser.Parser       Language.Futhark.Parser.Lexer       Paths_futhark+  hs-source-dirs:+      src+  ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists+  build-depends:+      aeson >=1.0.0.0+    , ansi-terminal >=0.6.3.1+    , array >=0.4+    , base >=4 && <5+    , binary >=0.8.3+    , blaze-html >=0.9.0.1+    , bytestring >=0.10.8+    , containers >=0.5+    , data-binary-ieee754 >=0.1+    , directory >=1.3.0.0+    , directory-tree >=0.12.1+    , dlist >=0.6.0.1+    , file-embed >=0.0.9+    , filepath >=1.4.1.1+    , free >=4.12.4+    , gitrev >=1.2.0+    , haskeline+    , http-client >=0.5.7.0+    , http-client-tls >=0.3.5.1+    , http-conduit >=2.2.4+    , language-c-quote >=0.12+    , mainland-pretty >=0.6.1+    , markdown >=0.1.16+    , megaparsec >=7.0.1+    , mtl >=2.2.1+    , neat-interpolation >=0.3+    , parallel >=3.2.1.0+    , parser-combinators >=1.0.0+    , process >=1.4.3.0+    , process-extras >=0.7.2+    , random+    , regex-tdfa >=1.2+    , srcloc >=0.4+    , template-haskell >=2.11.1+    , temporary+    , terminal-size >=0.3+    , text >=1.2.2.2+    , time >=1.6.0.1+    , transformers >=0.3+    , utf8-string >=1+    , vector >=0.12+    , vector-binary-instances >=0.2.2.0+    , versions >=3.3.1+    , zip-archive >=0.3.1.1+    , zlib >=0.6.1.2+  build-tools:+      alex+    , happy   default-language: Haskell2010  executable futhark   main-is: src/futhark.hs+  other-modules:+      Paths_futhark   ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists -threaded -rtsopts "-with-rtsopts=-N -qg"   build-depends:       base     , futhark     , text-  other-modules:-      Paths_futhark   default-language: Haskell2010  test-suite unit   type: exitcode-stdio-1.0   main-is: futhark_tests.hs-  hs-source-dirs:-      unittests-  ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists-  build-depends:-      QuickCheck >=2.8-    , base-    , containers-    , futhark-    , megaparsec-    , mtl-    , parser-combinators-    , tasty-    , tasty-hunit-    , tasty-quickcheck-    , text   other-modules:       Futhark.Analysis.ScalExpTests       Futhark.Optimise.AlgSimplifyTests@@ -357,4 +343,19 @@       Language.Futhark.CoreTests       Language.Futhark.SyntaxTests       Paths_futhark+  hs-source-dirs:+      unittests+  ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists+  build-depends:+      QuickCheck >=2.8+    , base+    , containers+    , futhark+    , megaparsec+    , mtl+    , parser-combinators+    , tasty+    , tasty-hunit+    , tasty-quickcheck+    , text   default-language: Haskell2010

@@ -122,6 +122,8 @@   val acos: t -> t   val atan: t -> t   val atan2: t -> t -> t+  val gamma: t -> t+  val lgamma: t -> t    -- | Natural logarithm.   val log: t -> t@@ -828,6 +830,8 @@   let asin (x: f64) = intrinsics.asin64 x   let atan (x: f64) = intrinsics.atan64 x   let atan2 (x: f64) (y: f64) = intrinsics.atan2_64 x y+  let gamma = intrinsics.gamma64+  let lgamma = intrinsics.lgamma64    let ceil (x: f64) : f64 =     let i = to_i64 x@@ -937,6 +941,8 @@   let asin (x: f32) = intrinsics.asin32 x   let atan (x: f32) = intrinsics.atan32 x   let atan2 (x: f32) (y: f32) = intrinsics.atan2_32 x y+  let gamma = intrinsics.gamma32+  let lgamma = intrinsics.lgamma32    let ceil (x: f32) : f32 =     let i = to_i32 x

@@ -23,12 +23,12 @@ let opaque 't (x: t): t =   intrinsics.opaque x --- | Semantically just identity, but in `futharki` the argument value--- will be printed.+-- | Semantically just identity, but when run in the interpreter, the+-- argument value will be printed. let trace 't (x: t): t =   intrinsics.trace x  -- | Semantically just identity, but acts as a break point in--- `futharki`.+-- `futhark repl`. let break 't (x: t): t =   intrinsics.break x

@@ -158,19 +158,20 @@  -- | `stream_red 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.  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.+-- combine the per-chunk results into a final result.  The `i32`+-- 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, `f []` must produce the neutral element for+-- A chunk may be empty, and `f 0 []` must produce the neutral element for -- `op`. -- -- **Work:** *O(n)* -- -- **Span:** *O(log(n))*-let stream_red 'a 'b (op: b -> b -> b) (f: []a -> b) (as: []a): b =+let stream_red 'a 'b (op: b -> b -> b) (f: i32 -> []a -> b) (as: []a): b =   intrinsics.stream_red (op, f, as)  -- | As `stream_red`@term, but the chunks do not necessarily@@ -180,7 +181,7 @@ -- **Work:** *O(n)* -- -- **Span:** *O(log(n))*-let stream_red_per 'a 'b (op: b -> b -> b) (f: []a -> b) (as: []a): b =+let stream_red_per 'a 'b (op: b -> b -> b) (f: i32 -> []a -> b) (as: []a): b =   intrinsics.stream_red_per (op, f, as)  -- | Similar to `stream_red`@term, except that each chunk must produce@@ -190,7 +191,7 @@ -- **Work:** *O(n)* -- -- **Span:** *O(1)*-let stream_map 'a 'b (f: []a -> []b) (as: []a): *[]b =+let stream_map 'a 'b (f: i32 -> []a -> []b) (as: []a): *[]b =   intrinsics.stream_map (f, as)  -- | Similar to `stream_map`@term, but the chunks do not necessarily@@ -200,7 +201,7 @@ -- **Work:** *O(n)* -- -- **Span:** *O(1)*-let stream_map_per 'a 'b (f: []a -> []b) (as: []a): *[]b =+let stream_map_per 'a 'b (f: i32 -> []a -> []b) (as: []a): *[]b =   intrinsics.stream_map_per (f, as)  -- | Return `true` if the given function returns `true` for all

@@ -1,4 +1,4 @@-/* Simple CUDA runtime framework */+// Start of cuda.h.  #define CUDA_SUCCEED(x) cuda_api_succeed(x, #x, __FILE__, __LINE__) #define NVRTC_SUCCEED(x) nvrtc_api_succeed(x, #x, __FILE__, __LINE__)@@ -98,6 +98,7 @@   size_t max_grid_size;   size_t max_tile_size;   size_t max_threshold;+  size_t max_shared_memory;    size_t lockstep_width; };@@ -193,7 +194,7 @@     src_len += strlen(*p);   } -  char *src = malloc(src_len + 1);+  char *src = (char*) malloc(src_len + 1);   size_t n = 0;   for (p = src_fragments; *p; p++) {     strcpy(src + n, *p);@@ -262,7 +263,7 @@   }    size_t n_opts, i = 0, i_dyn, n_opts_alloc = 20 + num_extra_opts + ctx->cfg.num_sizes;-  const char **opts = malloc(n_opts_alloc * sizeof(const char *));+  const char **opts = (const char**) malloc(n_opts_alloc * sizeof(const char *));   if (!arch_set) {     opts[i++] = "-arch";     opts[i++] = cuda_nvrtc_get_arch(ctx->dev);@@ -302,7 +303,7 @@   if (res != NVRTC_SUCCESS) {     size_t log_size;     if (nvrtcGetProgramLogSize(prog, &log_size) == NVRTC_SUCCESS) {-      char *log = malloc(log_size);+      char *log = (char*) malloc(log_size);       if (nvrtcGetProgramLog(prog, log) == NVRTC_SUCCESS) {         fprintf(stderr,"Compilation log:\n%s\n", log);       }@@ -317,7 +318,7 @@   char *ptx;   size_t ptx_size;   NVRTC_SUCCEED(nvrtcGetPTXSize(prog, &ptx_size));-  ptx = malloc(ptx_size);+  ptx = (char*) malloc(ptx_size);   NVRTC_SUCCEED(nvrtcGetPTX(prog, ptx));    NVRTC_SUCCEED(nvrtcDestroyProgram(&prog));@@ -405,7 +406,7 @@   len = ftell(f);   assert(fseek(f, 0, SEEK_SET) == 0); -  buf = malloc(len + 1);+  buf = (char*) malloc(len + 1);   assert(fread(buf, 1, len, f) == len);   buf[len] = 0;   *obuf = buf;@@ -467,6 +468,7 @@    free_list_init(&ctx->free_list); +  ctx->max_shared_memory = device_query(ctx->dev, MAX_SHARED_MEMORY_PER_BLOCK);   ctx->max_block_size = device_query(ctx->dev, MAX_THREADS_PER_BLOCK);   ctx->max_grid_size = device_query(ctx->dev, MAX_GRID_DIM_X);   ctx->max_tile_size = sqrt(ctx->max_block_size);@@ -557,3 +559,4 @@   return CUDA_SUCCESS; } +// End of cuda.h.

@@ -1,4 +1,4 @@-/* Free list management */+// Start of free_list.h.  /* An entry in the free list.  May be invalid, to avoid having to    deallocate entries as soon as they are removed.  There is also a@@ -19,7 +19,7 @@ void free_list_init(struct free_list *l) {   l->capacity = 30; // Picked arbitrarily.   l->used = 0;-  l->entries = malloc(sizeof(struct free_list_entry) * l->capacity);+  l->entries = (struct free_list_entry*) malloc(sizeof(struct free_list_entry) * l->capacity);   for (int i = 0; i < l->capacity; i++) {     l->entries[i].valid = 0;   }@@ -108,3 +108,4 @@   return 1; } +// End of free_list.h.

@@ -1,3 +1,5 @@+// Start of lock.h.+ /* A very simple cross-platform implementation of locks.  Uses    pthreads on Unix and some Windows thing there.  Futhark's    host-level code is not multithreaded, but user code may be, so we@@ -51,7 +53,9 @@  static void free_lock(lock_t *lock) {   /* Nothing to do for pthreads. */-  lock = lock;+  (void)lock; }  #endif++// End of lock.h.

@@ -1,4 +1,4 @@-/* The simple OpenCL runtime framework used by Futhark. */+// Start of opencl.h.  #define OPENCL_SUCCEED_FATAL(e) opencl_succeed_fatal(e, #e, __FILE__, __LINE__) #define OPENCL_SUCCEED_NONFATAL(e) opencl_succeed_nonfatal(e, #e, __FILE__, __LINE__)@@ -97,6 +97,7 @@   size_t max_num_groups;   size_t max_tile_size;   size_t max_threshold;+  size_t max_local_memory;    size_t lockstep_width; };@@ -117,7 +118,7 @@  static char *strclone(const char *str) {   size_t size = strlen(str) + 1;-  char *copy = malloc(size);+  char *copy = (char*) malloc(size);   if (copy == NULL) {     return NULL;   }@@ -150,7 +151,7 @@   return s; } -static const char* opencl_error_string(unsigned int err)+static const char* opencl_error_string(cl_int err) {     switch (err) {         case CL_SUCCESS:                            return "Success!";@@ -254,7 +255,7 @@    OPENCL_SUCCEED_FATAL(clGetPlatformInfo(platform, param, 0, NULL, &req_bytes)); -  info = malloc(req_bytes);+  info = (char*) malloc(req_bytes);    OPENCL_SUCCEED_FATAL(clGetPlatformInfo(platform, param, req_bytes, info, NULL)); @@ -268,7 +269,7 @@    OPENCL_SUCCEED_FATAL(clGetDeviceInfo(device, param, 0, NULL, &req_bytes)); -  info = malloc(req_bytes);+  info = (char*) malloc(req_bytes);    OPENCL_SUCCEED_FATAL(clGetDeviceInfo(device, param, req_bytes, info, NULL)); @@ -461,7 +462,7 @@     size_t ret_val_size;     OPENCL_SUCCEED_FATAL(clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, 0, NULL, &ret_val_size)); -    build_log = malloc(ret_val_size+1);+    build_log = (char*) malloc(ret_val_size+1);     OPENCL_SUCCEED_FATAL(clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, ret_val_size, build_log, NULL));      // The spec technically does not say whether the build log is zero-terminated, so let's be careful.@@ -531,6 +532,10 @@    size_t max_tile_size = sqrt(max_group_size); +  cl_ulong max_local_memory;+  OPENCL_SUCCEED_FATAL(clGetDeviceInfo(device_option.device, CL_DEVICE_LOCAL_MEM_SIZE,+                                       sizeof(size_t), &max_local_memory, NULL));+   // Make sure this function is defined.   post_opencl_setup(ctx, &device_option); @@ -553,6 +558,7 @@   ctx->max_group_size = max_group_size;   ctx->max_tile_size = max_tile_size; // No limit.   ctx->max_threshold = ctx->max_num_groups = 0; // No limit.+  ctx->max_local_memory = max_local_memory;    // Now we go through all the sizes, clamp them to the valid range,   // or set them to the default.@@ -608,7 +614,7 @@       src_size += strlen(*src);     } -    fut_opencl_src = malloc(src_size + 1);+    fut_opencl_src = (char*) malloc(src_size + 1);      size_t n, i;     for (i = 0, n = 0; srcs && srcs[i]; i++) {@@ -644,7 +650,7 @@       compile_opts_size += strlen(extra_build_opts[i] + 1);     } -    char *compile_opts = malloc(compile_opts_size);+    char *compile_opts = (char*) malloc(compile_opts_size);      int w = snprintf(compile_opts, compile_opts_size,                      "-DLOCKSTEP_WIDTH=%d ",@@ -687,7 +693,7 @@     size_t binary_size;     OPENCL_SUCCEED_FATAL(clGetProgramInfo(prog, CL_PROGRAM_BINARY_SIZES,                                           sizeof(size_t), &binary_size, NULL));-    unsigned char *binary = malloc(binary_size);+    unsigned char *binary = (unsigned char*) malloc(binary_size);     unsigned char *binaries[1] = { binary };     OPENCL_SUCCEED_FATAL(clGetProgramInfo(prog, CL_PROGRAM_BINARIES,                                           sizeof(unsigned char*), binaries, NULL));@@ -803,6 +809,9 @@   int error = opencl_alloc_actual(ctx, min_size, mem_out);    while (error == CL_MEM_OBJECT_ALLOCATION_FAILURE) {+    if (ctx->cfg.debugging) {+      fprintf(stderr, "Out of OpenCL memory: releasing entry from the free list...\n");+    }     cl_mem mem;     if (free_list_first(&ctx->free_list, &mem) == 0) {       error = clReleaseMemObject(mem);@@ -851,3 +860,5 @@    return CL_SUCCESS; }++// End of opencl.h.

@@ -1,4 +1,4 @@-/* Crash and burn. */+// Start of panic.h.  #include <stdarg.h> @@ -21,8 +21,10 @@   va_list vl;   va_start(vl, s);   size_t needed = 1 + vsnprintf(NULL, 0, s, vl);-  char *buffer = malloc(needed);+  char *buffer = (char*) malloc(needed);   va_start(vl, s); /* Must re-init. */   vsnprintf(buffer, needed, s, vl);   return buffer; }++// End of panic.h.

@@ -1,8 +1,7 @@-/* Some simple utilities for wall-clock timing.+// Start of timing.h. -   The function get_wall_time() returns the wall time in microseconds-   (with an unspecified offset).-*/+// The function get_wall_time() returns the wall time in microseconds+// (with an unspecified offset).  #ifdef _WIN32 @@ -28,3 +27,5 @@ }  #endif++// End of timing.h.

@@ -1,10 +1,10 @@-// Start of tuning.h+// Start of tuning.h.  static char* load_tuning_file(const char *fname,                               void *cfg,                               int (*set_size)(void*, const char*, size_t)) {   const int max_line_len = 1024;-  char* line = malloc(max_line_len);+  char* line = (char*) malloc(max_line_len);    FILE *f = fopen(fname, "r"); @@ -37,4 +37,4 @@   return NULL; } -// End of tuning.h+// End of tuning.h.

@@ -1,3 +1,5 @@+// Start of values.h.+ //// Text I/O  typedef int (*writer)(FILE*, void*);@@ -831,3 +833,5 @@     return type->write_str(out, src);   } }++// End of values.h.

@@ -401,6 +401,7 @@    public int MaxNumGroups;    public int MaxTileSize;    public int MaxThreshold;+   public int MaxLocalMemory;     public int LockstepWidth; }@@ -734,7 +735,7 @@     ComputeErrorCode error;     CLPlatformHandle platform;     CLDeviceHandle device;-    int MaxGroupSize;+    int MaxGroupSize, MaxLocalMemory;      ctx.OpenCL.LockstepWidth = 0; @@ -777,6 +778,16 @@      int MaxTileSize = (int) Math.Sqrt(MaxGroupSize); +    unsafe+    {+        IntPtr throwaway1 = new IntPtr();+        OPENCL_SUCCEED(CL10.GetDeviceInfo(device,+                                          ComputeDeviceInfo.LocalMemorySize,+                                          new IntPtr(sizeof(IntPtr)),+                                          new IntPtr(&MaxLocalMemory),+                                          out throwaway1));+    }+     // Make sure this function is defined.     PostOpenCLSetup(ref ctx, ref device_option); @@ -795,6 +806,7 @@     ctx.OpenCL.MaxGroupSize = MaxGroupSize;     ctx.OpenCL.MaxTileSize = MaxTileSize; // No limit.     ctx.OpenCL.MaxThreshold = ctx.OpenCL.MaxNumGroups; // No limit.+    ctx.OpenCL.MaxLocalMemory = MaxLocalMemory;      // Now we go through all the sizes, clamp them to the valid range,     // or set them to the default.

@@ -282,6 +282,8 @@ private static double futhark_asin64(double x){return Math.Asin(x);} private static double futhark_atan64(double x){return Math.Atan(x);} private static double futhark_atan2_64(double x, double y){return Math.Atan2(x, y);}+private static double futhark_gamma64(double x){throw new NotImplementedException();}+private static double futhark_lgamma64(double x){throw new NotImplementedException();} private static bool futhark_isnan64(double x){return double.IsNaN(x);} private static bool futhark_isinf64(double x){return double.IsInfinity(x);} private static long futhark_to_bits64(double x){return BitConverter.ToInt64(BitConverter.GetBytes(x),0);}@@ -299,6 +301,8 @@ private static float futhark_asin32(float x){return (float) Math.Asin(x);} private static float futhark_atan32(float x){return (float) Math.Atan(x);} private static float futhark_atan2_32(float x, float y){return (float) Math.Atan2(x, y);}+private static float futhark_gamma32(float x){throw new NotImplementedException();}+private static float futhark_lgamma32(float x){throw new NotImplementedException();} private static bool futhark_isnan32(float x){return float.IsNaN(x);} private static bool futhark_isinf32(float x){return float.IsInfinity(x);} private static int futhark_to_bits32(float x){return BitConverter.ToInt32(BitConverter.GetBytes(x), 0);}

@@ -1,4 +1,4 @@-# Helper functions dealing with memory blocks.+# Start of memory.py.  import ctypes as ct @@ -36,3 +36,5 @@    def __repr__(self):     return "<opaque Futhark value of type {}>".format(self.desc)++# End of memory.py.

@@ -109,6 +109,7 @@     self.max_tile_size = max_tile_size     self.max_threshold = 0     self.max_num_groups = 0+    self.max_local_memory = int(self.device.local_mem_size)     self.free_list = {}      if 'default_group_size' in sizes:

@@ -1,4 +1,8 @@+# Start of panic.py.+ def panic(exitcode, fmt, *args):     sys.stderr.write('%s: ' % sys.argv[0])     sys.stderr.write(fmt % args)     sys.exit(exitcode)++# End of panic.py.

@@ -1,6 +1,7 @@-# Scalar functions.+# Start of scalar.py.  import numpy as np+import math import struct  def signed(x):@@ -295,6 +296,12 @@ def futhark_atan2_64(x, y):   return np.arctan2(x, y) +def futhark_gamma64(x):+  return np.float64(math.gamma(x))++def futhark_lgamma64(x):+  return np.float64(math.lgamma(x))+ def futhark_round64(x):   return np.round(x) @@ -348,6 +355,12 @@ def futhark_atan2_32(x, y):   return np.arctan2(x, y) +def futhark_gamma32(x):+  return np.float32(math.gamma(x))++def futhark_lgamma32(x):+  return np.float32(math.lgamma(x))+ def futhark_round32(x):   return np.round(x) @@ -364,3 +377,5 @@ def futhark_from_bits32(x):   s = struct.pack('>l', x)   return np.float32(struct.unpack('>f', s)[0])++# End of scalar.py.

@@ -1,6 +1,4 @@-### start of tuning.py-###-### Reading the .tuning file.+# Start of tuning.py  def read_tuning_file(kvs, f):     for line in f.read().splitlines():@@ -8,4 +6,4 @@         kvs[size] = int(value)     return kvs -### end of tuning.py+# End of tuning.py.

@@ -1,3 +1,5 @@+# Start of values.py.+ # Hacky parser/reader/writer for values written in Futhark syntax. # Used for reading stdin when compiling standalone programs with the # Python code generator.@@ -76,7 +78,8 @@             read.append(c)         return True     except ValueError:-        map(f.unget_char, read[::-1])+        for c in read[::-1]:+            f.unget_char(c)         raise  def optional(p, *args):@@ -112,49 +115,46 @@     s = b''     c = f.get_char()     while c != None:-        if c in string.hexdigits:+        if c in b'01234556789ABCDEFabcdef':             s += c             c = f.get_char()-        elif c == '_':+        elif c == b'_':             c = f.get_char() # skip _         else:             f.unget_char(c)             break-    return str(int(s, 16))-+    return str(int(s, 16)).encode('utf8') # ugh  def parse_int(f):     s = b''     c = f.get_char()-    if c == b'0' and f.peek_char() in [b'x', b'X']:+    if c == b'0' and f.peek_char() in b'xX':         c = f.get_char() # skip X-        s += parse_hex_int(f)+        return parse_hex_int(f)     else:         while c != None:             if c.isdigit():                 s += c                 c = f.get_char()-            elif c == '_':+            elif c == b'_':                 c = f.get_char() # skip _             else:                 f.unget_char(c)                 break-    if len(s) == 0:-        raise ValueError-    return s+        if len(s) == 0:+            raise ValueError+        return s  def parse_int_signed(f):     s = b''     c = f.get_char()      if c == b'-' and f.peek_char().isdigit():-      s = c + parse_int(f)+      return c + parse_int(f)     else:       if c != b'+':           f.unget_char(c)-      s = parse_int(f)--    return s+      return parse_int(f)  def read_str_comma(f):     skip_spaces(f)@@ -568,7 +568,7 @@             return read_scalar(reader, basetype)         return (dims, basetype) -def write_value(v, out=sys.stdout):+def write_value_text(v, out=sys.stdout):     if type(v) == np.uint8:         out.write("%uu8" % v)     elif type(v) == np.uint16:@@ -625,6 +625,49 @@     else:         raise Exception("Cannot print value of type {}: {}".format(type(v), v)) -################################################################################-### end of values.py-################################################################################+type_strs = { np.dtype('int8'): b'  i8',+              np.dtype('int16'): b' i16',+              np.dtype('int32'): b' i32',+              np.dtype('int64'): b' i64',+              np.dtype('uint8'): b'  u8',+              np.dtype('uint16'): b' u16',+              np.dtype('uint32'): b' u32',+              np.dtype('uint64'): b' u64',+              np.dtype('float32'): b' f32',+              np.dtype('float64'): b' f64',+              np.dtype('bool'): b'bool'}++def construct_binary_value(v):+    t = v.dtype+    shape = v.shape++    elems = 1+    for d in shape:+        elems *= d++    num_bytes = 1 + 1 + 1 + 4 + len(shape) * 8 + elems * t.itemsize+    bytes = bytearray(num_bytes)+    bytes[0] = np.int8(ord('b'))+    bytes[1] = 2+    bytes[2] = np.int8(len(shape))+    bytes[3:7] = type_strs[t]++    for i in range(len(shape)):+        bytes[7+i*8:7+(i+1)*8] = np.int64(shape[i]).tostring()++    bytes[7+len(shape)*8:] = np.ascontiguousarray(v).tostring()++    return bytes++def write_value_binary(v, out=sys.stdout):+    if sys.version_info >= (3,0):+        out = out.buffer+    out.write(construct_binary_value(v))++def write_value(v, out=sys.stdout, binary=False):+    if binary:+        return write_value_binary(v, out=out)+    else:+        return write_value_text(v, out=out)++# End of values.py.

@@ -50,7 +50,6 @@ analyseExp = mapExp analyse   where analyse =           Mapper { mapOnSubExp = return-                 , mapOnCertificates = return                  , mapOnVName = return                  , mapOnBody = const $ return . analyseBody                  , mapOnRetType = return

@@ -40,7 +40,7 @@           in M.unions [branchdeps, deps, tdeps, fdeps]          grow deps (Let pat _ e) =-          let free = freeIn pat <> freeInExp e+          let free = freeIn pat <> freeIn e               freeDeps = S.unions $ map (depsOfVar deps) $ S.toList free           in M.fromList [ (name, freeDeps) | name <- patternNames pat ] `M.union` deps 

@@ -13,7 +13,6 @@   ) where -import Control.Monad import Data.List import Data.Maybe import qualified Data.Map.Strict as M@@ -74,7 +73,7 @@           -> SOAC lore           -> m (Maybe (MapNest lore)) -fromSOAC' bound (SOAC.Screma w (SOAC.ScremaForm (_, []) (_, _, []) lam) inps) = do+fromSOAC' bound (SOAC.Screma w (SOAC.ScremaForm (_, []) [] lam) inps) = do   maybenest <- case (stmsToList $ bodyStms $ lambdaBody lam,                      bodyResult $ lambdaBody lam) of     ([Let pat _ e], res) | res == map Var (patternNames pat) ->@@ -91,12 +90,11 @@         unzip <$>         fixInputs w (zip (map paramName $ lambdaParams lam) inps)         (zip (params mn) inps')-      let n' = Nesting {-            nestingParamNames   = ps-            , nestingResult     = patternNames pat-            , nestingReturnType = typeOf mn-            , nestingWidth      = inner_w-            }+      let n' = Nesting { nestingParamNames = ps+                       , nestingResult     = patternNames pat+                       , nestingReturnType = typeOf mn+                       , nestingWidth      = inner_w+                       }       return $ Just $ MapNest w body' (n':ns') inps''     -- No nested MapNest it seems.     _ -> do@@ -106,11 +104,11 @@             | otherwise =               Nothing           boundUsedInBody =-            mapMaybe isBound $ S.toList $ freeInLambda lam+            mapMaybe isBound $ S.toList $ freeIn lam       newParams <- mapM (newIdent' (++"_wasfree")) boundUsedInBody       let subst = M.fromList $                   zip (map identName boundUsedInBody) (map identName newParams)-          inps' = map (substituteNames subst) inps +++          inps' = inps ++                   map (SOAC.addTransform (SOAC.Replicate mempty $ Shape [w]) . SOAC.identInput)                   boundUsedInBody           lam' =@@ -144,35 +142,16 @@ fixInputs :: MonadFreshNames m =>              SubExp -> [(VName, SOAC.Input)] -> [(VName, SOAC.Input)]           -> m [(VName, SOAC.Input)]-fixInputs w ourInps childInps =-  reverse . snd <$> foldM inspect (ourInps, []) childInps+fixInputs w ourInps = mapM inspect   where     isParam x (y, _) = x == y -    findParam :: [(VName, SOAC.Input)]-              -> VName-              -> Maybe ((VName, SOAC.Input), [(VName, SOAC.Input)])-    findParam remPs v-      | ([ourP], remPs') <- partition (isParam v) remPs = Just (ourP, remPs')-      | otherwise                                       = Nothing--    inspect :: MonadFreshNames m =>-               ([(VName, SOAC.Input)], [(VName, SOAC.Input)])-            -> (VName, SOAC.Input)-            -> m ([(VName, SOAC.Input)], [(VName, SOAC.Input)])-    inspect (remPs, newInps) (_, SOAC.Input ts v _)-      | Just ((p,pInp), remPs') <- findParam remPs v =+    inspect (_, SOAC.Input ts v _)+      | Just (p,pInp) <- find (isParam v) ourInps = do           let pInp' = SOAC.transformRows ts pInp-          in return (remPs',-                     (p, pInp') : newInps)--      | Just ((p,pInp), _) <- findParam newInps v = do-          -- The input corresponds to a variable that has already-          -- been used.           p' <- newNameFromString $ baseString p-          return (remPs, (p', pInp) : newInps)+          return (p', pInp') -    inspect (remPs, newInps) (param, SOAC.Input ts a t) = do+    inspect (param, SOAC.Input ts a t) = do       param' <- newNameFromString (baseString param ++ "_rep")-      return (remPs, (param',-                      SOAC.Input (ts SOAC.|> SOAC.Replicate mempty (Shape [w])) a t) : newInps)+      return (param', SOAC.Input (ts SOAC.|> SOAC.Replicate mempty (Shape [w])) a t)

@@ -597,7 +597,8 @@       return (Stream w (Sequential nes) strmlam inps,               map paramIdent inpacc_ids) -      | Just (comm, lamin, nes, _) <- Futhark.isRedomapSOAC form -> do+      | Just (reds, _) <- Futhark.isRedomapSOAC form,+        Futhark.Reduce comm lamin nes <- Futhark.singleReduce reds -> do       -- Redomap(+,lam,nes,a) => is translated in strem's body to:       -- 1. let (acc0_ids,strm_resids) = redomap(+,lam,nes,a_ch) in       -- 2. let acc'                   = acc + acc0_ids          in@@ -613,7 +614,8 @@       inpacc_ids <- mapM (newParam "inpacc")  accrtps       acc0_ids   <- mapM (newIdent "acc0"  )  accrtps       -- 1. let (acc0_ids,strm_resids) = redomap(+,lam,nes,a_ch) in-      let insoac = Futhark.Screma chvar (Futhark.redomapSOAC comm lamin nes foldlam) $+      let insoac = Futhark.Screma chvar+                   (Futhark.redomapSOAC [Futhark.Reduce comm lamin nes] foldlam) $                    map paramName strm_inpids           insbnd = mkLet [] (acc0_ids++strm_resids) $ Op insoac       -- 2. let acc'     = acc + acc0_ids    in

@@ -7,6 +7,7 @@   , primExpFromSubExp   , primExpFromSubExpM   , replaceInPrimExp+  , replaceInPrimExpM   , substituteInPrimExp      -- * Module reexport@@ -14,6 +15,7 @@   ) where  import qualified Control.Monad.Fail as Fail+import           Control.Monad.Identity import           Data.Loc import qualified Data.Map.Strict as M import           Data.Maybe@@ -84,23 +86,30 @@ primExpFromSubExp t (Var v)      = LeafExp v t primExpFromSubExp _ (Constant v) = ValueExp v --- | Applying a transformation to the leaves in a 'PrimExp'.-replaceInPrimExp :: (v -> PrimType -> PrimExp v) ->-                    PrimExp v -> PrimExp v-replaceInPrimExp f (LeafExp v pt) =+-- | Applying a monadic transformation to the leaves in a 'PrimExp'.+replaceInPrimExpM :: Monad m =>+                     (a -> PrimType -> m (PrimExp b)) ->+                     PrimExp a -> m (PrimExp b)+replaceInPrimExpM f (LeafExp v pt) =   f v pt-replaceInPrimExp _ (ValueExp v) =-  ValueExp v-replaceInPrimExp f (BinOpExp bop pe1 pe2) =-  constFoldPrimExp $ BinOpExp bop (replaceInPrimExp f pe1) (replaceInPrimExp f pe2)-replaceInPrimExp f (CmpOpExp cop pe1 pe2) =-  CmpOpExp cop (replaceInPrimExp f pe1) (replaceInPrimExp f pe2)-replaceInPrimExp f (UnOpExp uop pe) =-  UnOpExp uop $ replaceInPrimExp f pe-replaceInPrimExp f (ConvOpExp cop pe) =-  ConvOpExp cop $ replaceInPrimExp f pe-replaceInPrimExp f (FunExp h args t) =-  FunExp h (map (replaceInPrimExp f) args) t+replaceInPrimExpM _ (ValueExp v) =+  return $ ValueExp v+replaceInPrimExpM f (BinOpExp bop pe1 pe2) =+  constFoldPrimExp <$>+  (BinOpExp bop <$> replaceInPrimExpM f pe1 <*> replaceInPrimExpM f pe2)+replaceInPrimExpM f (CmpOpExp cop pe1 pe2) =+  CmpOpExp cop <$> replaceInPrimExpM f pe1 <*> replaceInPrimExpM f pe2+replaceInPrimExpM f (UnOpExp uop pe) =+  UnOpExp uop <$> replaceInPrimExpM f pe+replaceInPrimExpM f (ConvOpExp cop pe) =+  ConvOpExp cop <$> replaceInPrimExpM f pe+replaceInPrimExpM f (FunExp h args t) =+  FunExp h <$> mapM (replaceInPrimExpM f) args <*> pure t++replaceInPrimExp :: (a -> PrimType -> PrimExp b) ->+                    PrimExp a -> PrimExp b+replaceInPrimExp f e = runIdentity $ replaceInPrimExpM f' e+  where f' x y = return $ f x y  -- | Substituting names in a PrimExp with other PrimExps substituteInPrimExp :: Ord v => M.Map v (PrimExp v)

@@ -66,15 +66,14 @@ analyseExp = mapExpM analyse   where analyse =           Mapper { mapOnSubExp = return-                    , mapOnCertificates = return-                    , mapOnVName = return-                    , mapOnBody = const analyseBody-                    , mapOnRetType = return-                    , mapOnBranchType = return-                    , mapOnFParam = return-                    , mapOnLParam = return-                    , mapOnOp = return . addOpRanges-                    }+                 , mapOnVName = return+                 , mapOnBody = const analyseBody+                 , mapOnRetType = return+                 , mapOnBranchType = return+                 , mapOnFParam = return+                 , mapOnLParam = return+                 , mapOnOp = return . addOpRanges+                 }  analyseLambda :: (Attributes lore, CanBeRanged (Op lore)) =>                  Lambda lore

@@ -284,6 +284,8 @@                    , (Mul t, STimes)                    , (AST.SDiv t, SDiv)                    , (AST.Pow t, SPow)+                   , (AST.SMax t, \x y -> MaxMin False [x,y])+                   , (AST.SMin t, \x y -> MaxMin True [x,y])                    ]  instance FreeIn ScalExp where

@@ -566,7 +566,7 @@                 AST.FParam lore              -> SymbolTable lore              -> SymbolTable lore-insertFParam fparam = insertEntry name entry+insertFParam fparam = flip (foldr (`isAtLeast` 0)) sizes . insertEntry name entry   where name = AST.paramName fparam         entry = FParam FParamEntry { fparamRange = (Nothing, Nothing)                                    , fparamAttr = AST.paramAttr fparam@@ -574,12 +574,13 @@                                    , fparamStmDepth = 0                                    , fparamConsumed = False                                    }+        sizes = subExpVars $ arrayDims $ AST.paramType fparam  insertFParams :: Attributes lore =>                  [AST.FParam lore]               -> SymbolTable lore               -> SymbolTable lore-insertFParams fparams symtable = foldr insertFParam symtable fparams+insertFParams fparams symtable = foldl' (flip insertFParam) symtable fparams  insertLParamWithRange :: Attributes lore =>                          LParam lore -> ScalExpRange -> IndexArray -> SymbolTable lore

@@ -1,12 +1,5 @@ {-# LANGUAGE FlexibleContexts #-}-module Futhark.Analysis.Usage-       ( usageInStm-       , usageInExp-       , usageInLambda--       , UsageInOp(..)-       )-       where+module Futhark.Analysis.Usage ( usageInStm ) where  import Data.Foldable import qualified Data.Set as S@@ -15,13 +8,12 @@ import Futhark.Representation.AST.Attributes.Aliases import qualified Futhark.Analysis.UsageTable as UT -usageInStm :: (Attributes lore, Aliased lore, UsageInOp (Op lore)) =>-              Stm lore -> UT.UsageTable+usageInStm :: (Attributes lore, Aliased lore) => Stm lore -> UT.UsageTable usageInStm (Let pat lore e) =   mconcat [usageInPat,            usageInExpLore,            usageInExp e,-           UT.usages (freeInExp e)]+           UT.usages (freeIn e)]   where usageInPat =           UT.usages (mconcat (map freeIn $ patternElements pat)                      `S.difference`@@ -29,7 +21,7 @@         usageInExpLore =           UT.usages $ freeIn lore -usageInExp :: (Aliased lore, UsageInOp (Op lore)) => Exp lore -> UT.UsageTable+usageInExp :: Aliased lore => Exp lore -> UT.UsageTable usageInExp (Apply _ args _ _) =   mconcat [ mconcat $ map UT.consumedUsage $             S.toList $ subExpAliases arg@@ -44,22 +36,5 @@ usageInExp (BasicOp (Update src _ _)) =   UT.consumedUsage src usageInExp (Op op) =-  mconcat $ usageInOp op : map UT.consumedUsage (S.toList $ consumedInOp op)+  mconcat $ map UT.consumedUsage (S.toList $ consumedInOp op) usageInExp _ = UT.empty--class UsageInOp op where-  usageInOp :: op -> UT.UsageTable--instance UsageInOp () where-  usageInOp () = mempty--usageInLambda :: Aliased lore =>-                 Lambda lore -> [VName] -> UT.UsageTable-usageInLambda lam arrs =-  mconcat $-  map (UT.consumedUsage . snd) $-  filter ((`S.member` consumed_in_body) . fst) $-  zip (map paramName arr_params) arrs-  where arr_params = snd $ splitAt n $ lambdaParams lam-        consumed_in_body = consumedInBody $ lambdaBody lam-        n = length arrs

@@ -19,7 +19,6 @@   , consumedUsage   , inResultUsage   , Usages-  , leftScope   )   where @@ -131,6 +130,3 @@ -- | x - y, but for Usages. withoutU :: Usages -> Usages -> Usages withoutU (Usages x) (Usages y) = Usages $ x .&. complement y--leftScope :: UsageTable -> UsageTable-leftScope (UsageTable table) = UsageTable $ M.map (`withoutU` inResultU) table

@@ -30,7 +30,7 @@              liftIO $ writeFile cpath impl            ToExecutable -> do              liftIO $ writeFile cpath $ CCUDA.asExecutable cprog-             let args = [cpath, "-O3", "-std=c99", "-lm", "-o", outpath]+             let args = [cpath, "-O", "-std=c99", "-lm", "-o", outpath]                         ++ extra_options              ret <- liftIO $ runProgramWithExitCode "gcc" args ""              case ret of

@@ -37,7 +37,7 @@            ToExecutable -> do              liftIO $ writeFile cpath $ COpenCL.asExecutable cprog              ret <- liftIO $ runProgramWithExitCode "gcc"-                    ([cpath, "-O3", "-std=c99", "-lm", "-o", outpath] ++ extra_options) ""+                    ([cpath, "-O", "-std=c99", "-lm", "-o", outpath] ++ extra_options) ""              case ret of                Left err ->                  externalErrorS $ "Failed to run gcc: " ++ show err

@@ -30,7 +30,7 @@       let extra = generateBoilerplate cuda_code cuda_prelude                                       kernel_names sizes       in Right <$> GC.compileProg operations extra cuda_includes-                   [Space "device", Space "local", DefaultSpace] cliOptions prog'+                   [Space "device", DefaultSpace] cliOptions prog'   where     operations :: GC.Operations OpenCL ()     operations = GC.Operations@@ -106,7 +106,7 @@   val' <- newVName "write_tmp"   GC.stm [C.cstm|{$ty:t $id:val' = $exp:val;                   CUDA_SUCCEED(-                    cuMemcpyHtoD($exp:mem + $exp:idx,+                    cuMemcpyHtoD($exp:mem + $exp:idx * sizeof($ty:t),                                  &$id:val',                                  sizeof($ty:t)));                  }|]@@ -119,7 +119,7 @@   GC.decl [C.cdecl|$ty:t $id:val;|]   GC.stm [C.cstm|CUDA_SUCCEED(                    cuMemcpyDtoH(&$id:val,-                                $exp:mem + $exp:idx,+                                $exp:mem + $exp:idx * sizeof($ty:t),                                 sizeof($ty:t)));                 |]   return [C.cexp|$id:val|]@@ -129,14 +129,12 @@ allocateCUDABuffer :: GC.Allocate OpenCL () allocateCUDABuffer mem size tag "device" =   GC.stm [C.cstm|CUDA_SUCCEED(cuda_alloc(&ctx->cuda, $exp:size, $exp:tag, &$exp:mem));|]-allocateCUDABuffer _ _ _ "local" = return () allocateCUDABuffer _ _ _ space =   fail $ "Cannot allocate in '" ++ space ++ "' memory space."  deallocateCUDABuffer :: GC.Deallocate OpenCL () deallocateCUDABuffer mem tag "device" =   GC.stm [C.cstm|CUDA_SUCCEED(cuda_free(&ctx->cuda, $exp:mem, $exp:tag));|]-deallocateCUDABuffer _ _ "local" = return () deallocateCUDABuffer _ _ space =   fail $ "Cannot deallocate in '" ++ space ++ "' memory space." @@ -187,7 +185,6 @@  cudaMemoryType :: GC.MemoryType OpenCL () cudaMemoryType "device" = return [C.cty|typename CUdeviceptr|]-cudaMemoryType "local" = pure [C.cty|unsigned char|] -- dummy type cudaMemoryType space =   fail $ "CUDA backend does not support '" ++ space ++ "' memory space." @@ -206,6 +203,7 @@     cudaSizeClass SizeGroup = "block_size"     cudaSizeClass SizeNumGroups = "grid_size"     cudaSizeClass SizeTile = "tile_size"+    cudaSizeClass SizeLocalMemory = "shared_memory" callKernel (LaunchKernel name args num_blocks block_size) = do   args_arr <- newVName "kernel_args"   time_start <- newVName "time_start"

@@ -86,13 +86,13 @@   GC.publicDef_ "context_config_new" GC.InitDecl $ \s ->     ([C.cedecl|struct $id:cfg* $id:s(void);|],      [C.cedecl|struct $id:cfg* $id:s(void) {-                         struct $id:cfg *cfg = malloc(sizeof(struct $id:cfg));+                         struct $id:cfg *cfg = (struct $id:cfg*) malloc(sizeof(struct $id:cfg));                          if (cfg == NULL) {                            return NULL;                          }                           cfg->num_nvrtc_opts = 0;-                         cfg->nvrtc_opts = malloc(sizeof(const char*));+                         cfg->nvrtc_opts = (const char**) malloc(sizeof(const char*));                          cfg->nvrtc_opts[0] = NULL;                          $stms:size_value_inits                          cuda_config_init(&cfg->cu_cfg, $int:num_sizes,@@ -113,7 +113,7 @@      [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *opt) {                          cfg->nvrtc_opts[cfg->num_nvrtc_opts] = opt;                          cfg->num_nvrtc_opts++;-                         cfg->nvrtc_opts = realloc(cfg->nvrtc_opts, (cfg->num_nvrtc_opts+1) * sizeof(const char*));+                         cfg->nvrtc_opts = (const char**) realloc(cfg->nvrtc_opts, (cfg->num_nvrtc_opts+1) * sizeof(const char*));                          cfg->nvrtc_opts[cfg->num_nvrtc_opts] = NULL;                        }|]) @@ -227,7 +227,7 @@   GC.publicDef_ "context_new" GC.InitDecl $ \s ->     ([C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg);|],      [C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg) {-                          struct $id:ctx* ctx = malloc(sizeof(struct $id:ctx));+                          struct $id:ctx* ctx = (struct $id:ctx*) malloc(sizeof(struct $id:ctx));                           if (ctx == NULL) {                             return NULL;                           }

@@ -29,7 +29,7 @@     Right (Program opencl_code opencl_prelude kernel_names types sizes prog') ->       Right <$> GC.compileProg operations                 (generateBoilerplate opencl_code opencl_prelude kernel_names types sizes)-                include_opencl_h [Space "device", Space "local", DefaultSpace]+                include_opencl_h [Space "device", DefaultSpace]                 cliOptions prog'   where operations :: GC.Operations OpenCL ()         operations = GC.Operations@@ -167,7 +167,7 @@   GC.stm [C.cstm|{$item:decl                   OPENCL_SUCCEED_OR_RETURN(                     clEnqueueWriteBuffer(ctx->opencl.queue, $exp:mem, $exp:blocking,-                                         $exp:i, sizeof($ty:t),+                                         $exp:i * sizeof($ty:t), sizeof($ty:t),                                          &$id:val',                                          0, NULL, NULL));                 }|]@@ -180,7 +180,7 @@   GC.decl [C.cdecl|$ty:t $id:val;|]   GC.stm [C.cstm|OPENCL_SUCCEED_OR_RETURN(                    clEnqueueReadBuffer(ctx->opencl.queue, $exp:mem, CL_TRUE,-                                       $exp:i, sizeof($ty:t),+                                       $exp:i * sizeof($ty:t), sizeof($ty:t),                                        &$id:val,                                        0, NULL, NULL));               |]@@ -191,16 +191,12 @@ allocateOpenCLBuffer :: GC.Allocate OpenCL () allocateOpenCLBuffer mem size tag "device" =   GC.stm [C.cstm|OPENCL_SUCCEED_OR_RETURN(opencl_alloc(&ctx->opencl, $exp:size, $exp:tag, &$exp:mem));|]-allocateOpenCLBuffer _ _ _ "local" =-  return () -- Hack - these memory blocks do not actually exist. allocateOpenCLBuffer _ _ _ space =-  fail $ "Cannot allocate in '" ++ space ++ "' space"+  fail $ "Cannot allocate in '" ++ space ++ "' space."  deallocateOpenCLBuffer :: GC.Deallocate OpenCL () deallocateOpenCLBuffer mem tag "device" =   GC.stm [C.cstm|OPENCL_SUCCEED_OR_RETURN(opencl_free(&ctx->opencl, $exp:mem, $exp:tag));|]-deallocateOpenCLBuffer _ _ "local" =-  return () -- Hack - these memory blocks do not actually exist. deallocateOpenCLBuffer _ _ space =   fail $ "Cannot deallocate in '" ++ space ++ "' space" @@ -252,7 +248,6 @@  openclMemoryType :: GC.MemoryType OpenCL () openclMemoryType "device" = pure [C.cty|typename cl_mem|]-openclMemoryType "local" = pure [C.cty|unsigned char|] -- dummy type openclMemoryType space =   fail $ "OpenCL backend does not support '" ++ space ++ "' memory space." 

@@ -69,13 +69,13 @@   GC.publicDef_ "context_config_new" GC.InitDecl $ \s ->     ([C.cedecl|struct $id:cfg* $id:s(void);|],      [C.cedecl|struct $id:cfg* $id:s(void) {-                         struct $id:cfg *cfg = malloc(sizeof(struct $id:cfg));+                         struct $id:cfg *cfg = (struct $id:cfg*) malloc(sizeof(struct $id:cfg));                          if (cfg == NULL) {                            return NULL;                          }                           cfg->num_build_opts = 0;-                         cfg->build_opts = malloc(sizeof(const char*));+                         cfg->build_opts = (const char**) malloc(sizeof(const char*));                          cfg->build_opts[0] = NULL;                          $stms:size_value_inits                          opencl_config_init(&cfg->opencl, $int:num_sizes,@@ -96,7 +96,7 @@      [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *opt) {                          cfg->build_opts[cfg->num_build_opts] = opt;                          cfg->num_build_opts++;-                         cfg->build_opts = realloc(cfg->build_opts, (cfg->num_build_opts+1) * sizeof(const char*));+                         cfg->build_opts = (const char**) realloc(cfg->build_opts, (cfg->num_build_opts+1) * sizeof(const char*));                          cfg->build_opts[cfg->num_build_opts] = NULL;                        }|]) @@ -264,7 +264,7 @@   GC.publicDef_ "context_new" GC.InitDecl $ \s ->     ([C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg);|],      [C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg) {-                          struct $id:ctx* ctx = malloc(sizeof(struct $id:ctx));+                          struct $id:ctx* ctx = (struct $id:ctx*) malloc(sizeof(struct $id:ctx));                           if (ctx == NULL) {                             return NULL;                           }@@ -281,7 +281,7 @@   GC.publicDef_ "context_new_with_command_queue" GC.InitDecl $ \s ->     ([C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg, typename cl_command_queue queue);|],      [C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg, typename cl_command_queue queue) {-                          struct $id:ctx* ctx = malloc(sizeof(struct $id:ctx));+                          struct $id:ctx* ctx = (struct $id:ctx*) malloc(sizeof(struct $id:ctx));                           if (ctx == NULL) {                             return NULL;                           }

@@ -129,6 +129,7 @@                      Imp.SizeNumGroups -> "MaxNumGroups"                      Imp.SizeTile -> "MaxTileSize"                      Imp.SizeThreshold{} -> "MaxThreshold"+                     Imp.SizeLocalMemory -> "MaxLocalMemory"  callKernel (Imp.HostCode c) = CS.compileCode c @@ -254,7 +255,8 @@     , AssignTyped (PointerT VoidT) (Var ptr) (Just $ Addr $ Var scalar)     , Exp $ CS.simpleCall "CL10.EnqueueWriteBuffer"         [ Var "Ctx.OpenCL.Queue", memblockFromMem mem, Bool True-        ,CS.toIntPtr i,CS.toIntPtr $ CS.sizeOf bt',CS.toIntPtr $ Var ptr+        , CS.toIntPtr $ BinOp "*" i (CS.sizeOf bt')+        , CS.toIntPtr $ CS.sizeOf bt',CS.toIntPtr $ Var ptr     , Integer 0, Null, Null]     ] @@ -271,7 +273,8 @@     [ CS.assignScalarPointer (Var val) (Var ptr)     , Exp $ CS.simpleCall "CL10.EnqueueReadBuffer"       [ Var "Ctx.OpenCL.Queue", memblockFromMem mem , Bool True-      , CS.toIntPtr i, CS.toIntPtr $ CS.sizeOf bt', CS.toIntPtr $ Var ptr+      , CS.toIntPtr $ BinOp "*" i (CS.sizeOf bt')+      , CS.toIntPtr $ CS.sizeOf bt', CS.toIntPtr $ Var ptr       , Integer 0, Null, Null]     ]   return $ Var val@@ -388,32 +391,26 @@   fail $ "Cannot return array from " ++ sid ++ " space."  unpackArrayInput :: CS.EntryInput Imp.OpenCL ()-unpackArrayInput mem memsize "device" t _ dims e = do+unpackArrayInput mem "device" t _ dims e = do   let size = foldr (BinOp "*") (Integer 1) dims'   let t' = CS.compilePrimTypeToAST t   let nbytes = BinOp "*" (CS.sizeOf t') size   zipWithM_ (CS.unpackDim e) dims [0..]   ptr <- pretty <$> newVName "ptr" -  CS.stm $ compileMemsize memsize nbytes--  let memsize' = CS.compileDim memsize-   CS.stm $ CS.getDefaultDecl (Imp.MemParam mem (Imp.Space "device"))-  allocateOpenCLBuffer mem memsize' "device"+  allocateOpenCLBuffer mem nbytes "device"   CS.stm $ Unsafe [Fixed (Var ptr) (Addr $ Index (Field e "Item1") $ IdxExp $ Integer 0)-      [ ifNotZeroSize memsize' $+      [ ifNotZeroSize nbytes $         Exp $ CS.simpleCall "CL10.EnqueueWriteBuffer"         [ Var "Ctx.OpenCL.Queue", memblockFromMem mem, Bool True-        , CS.toIntPtr (Integer 0), CS.toIntPtr memsize', CS.toIntPtr (Var ptr)+        , CS.toIntPtr (Integer 0), CS.toIntPtr nbytes, CS.toIntPtr (Var ptr)         , Integer 0, Null, Null]       ]]    where dims' = map CS.compileDim dims-        compileMemsize (Imp.VarSize v) nbytes = Assign (Var $ CS.compileName v) nbytes-        compileMemsize _ _                    = Pass -unpackArrayInput _ _ sid _ _ _ _ =+unpackArrayInput _ sid _ _ _ _ =   fail $ "Cannot accept array from " ++ sid ++ " space."  futharkSyncContext :: CSStmt

@@ -140,13 +140,13 @@ -- | The type of a memory block in the given memory space. type MemoryType op s = SpaceId -> CompilerM op s C.Type --- | Write a scalar to the given memory block with the given index and--- in the given memory space.+-- | Write a scalar to the given memory block with the given element+-- index and in the given memory space. type WriteScalar op s =   C.Exp -> C.Exp -> C.Type -> SpaceId -> Volatility -> C.Exp -> CompilerM op s () --- | Read a scalar from the given memory block with the given index and--- in the given memory space.+-- | Read a scalar from the given memory block with the given element+-- index and in the given memory space. type ReadScalar op s =   C.Exp -> C.Exp -> C.Type -> SpaceId -> Volatility -> CompilerM op s C.Exp @@ -528,14 +528,10 @@           (long long)size, desc, $string:spacedesc, ctx->$id:usagename);   }   int ret = $id:(fatMemUnRef space)(ctx, block, desc);-  $items:alloc-  block->references = (int*) malloc(sizeof(int));-  *(block->references) = 1;-  block->size = size;-  block->desc = desc;+   ctx->$id:usagename += size;   if (ctx->detail_memory) {-    fprintf(stderr, "Allocated %lld bytes for %s in %s (now allocated: %lld bytes)",+    fprintf(stderr, "Allocating %lld bytes for %s in %s (then allocated: %lld bytes)",             (long long) size,             desc, $string:spacedesc,             (long long) ctx->$id:usagename);@@ -548,6 +544,12 @@   } else if (ctx->detail_memory) {     fprintf(stderr, ".\n");   }++  $items:alloc+  block->references = (int*) malloc(sizeof(int));+  *(block->references) = 1;+  block->size = size;+  block->desc = desc;   return ret;   }|] @@ -673,7 +675,7 @@ opaqueName s vds = "opaque_" ++ hash (zipWith xor [0..] $ map ord (s ++ concatMap p vds))   where p (ScalarValue pt signed _) =           show (pt, signed)-        p (ArrayValue _ _ space pt signed dims) =+        p (ArrayValue _ space pt signed dims) =           show (space, pt, signed, length dims)          -- FIXME: a stupid hash algorithm; may have collisions.@@ -704,7 +706,7 @@   shape_array <- publicName $ "shape_" ++ name    let shape_names = [ "dim"++show i | i <- [0..rank-1] ]-      shape_params = [ [C.cparam|int $id:k|] | k <- shape_names ]+      shape_params = [ [C.cparam|typename int64_t $id:k|] | k <- shape_names ]       arr_size = cproduct [ [C.cexp|$id:k|] | k <- shape_names ]       arr_size_array = cproduct [ [C.cexp|arr->shape[$int:i]|] | i <- [0..rank-1] ]   copy <- asks envCopy@@ -759,7 +761,7 @@   return [C.cunit|           $ty:array_type* $id:new_array($ty:ctx_ty *ctx, $ty:pt' *data, $params:shape_params) {             $ty:array_type* bad = NULL;-            $ty:array_type *arr = malloc(sizeof($ty:array_type));+            $ty:array_type *arr = ($ty:array_type*) malloc(sizeof($ty:array_type));             if (arr == NULL) {               return bad;             }@@ -770,7 +772,7 @@           $ty:array_type* $id:new_raw_array($ty:ctx_ty *ctx, $ty:memty data, int offset,                                             $params:shape_params) {             $ty:array_type* bad = NULL;-            $ty:array_type *arr = malloc(sizeof($ty:array_type));+            $ty:array_type *arr = ($ty:array_type*) malloc(sizeof($ty:array_type));             if (arr == NULL) {               return bad;             }@@ -808,7 +810,7 @@        freeComponent _ ScalarValue{} =         return ()-      freeComponent i (ArrayValue _ _ _ pt signed shape) = do+      freeComponent i (ArrayValue _ _ pt signed shape) = do         let rank = length shape         free_array <- publicName $ "free_" ++ arrayName pt signed rank         stm [C.cstm|if ((tmp = $id:free_array(ctx, obj->$id:(tupleField i))) != 0) {@@ -834,7 +836,7 @@ valueDescToCType :: ValueDesc -> CompilerM op s C.Type valueDescToCType (ScalarValue pt signed _) =   return $ signedPrimTypeToCType signed pt-valueDescToCType (ArrayValue _ _ space pt signed shape) = do+valueDescToCType (ArrayValue _ space pt signed shape) = do   let pt' = signedPrimTypeToCType signed pt       rank = length shape   exists <- gets $ lookup (pt',rank) . compArrayStructs@@ -898,15 +900,11 @@           stm [C.cstm|$id:name = $exp:src;|]           return pt' -        prepareValue src vd@(ArrayValue mem mem_size _ _ _ shape) = do+        prepareValue src vd@(ArrayValue mem _ _ _ shape) = do           ty <- valueDescToCType vd            stm [C.cstm|$exp:mem = $exp:src->mem;|]-          case mem_size of-            VarSize v -> stm [C.cstm|$id:v = $exp:src->mem.size;|]-            ConstSize _ -> return () -           let rank = length shape               maybeCopyDim (VarSize d) i =                 Just [C.cstm|$id:d = $exp:src->shape[$int:i];|]@@ -924,7 +922,7 @@            case vd of             ArrayValue{} -> do-              stm [C.cstm|assert((*$id:pname = malloc(sizeof($ty:ty))) != NULL);|]+              stm [C.cstm|assert((*$id:pname = ($ty:ty*) malloc(sizeof($ty:ty))) != NULL);|]               prepareValue [C.cexp|*$id:pname|] vd               return [C.cparam|$ty:ty **$id:pname|]             ScalarValue{} -> do@@ -936,14 +934,14 @@           ty <- opaqueToCType desc vds           vd_ts <- mapM valueDescToCType vds -          stm [C.cstm|assert((*$id:pname = malloc(sizeof($ty:ty))) != NULL);|]+          stm [C.cstm|assert((*$id:pname = ($ty:ty*) malloc(sizeof($ty:ty))) != NULL);|]             forM_ (zip3 [0..] vd_ts vds) $ \(i,ct,vd) -> do             let field = [C.cexp|(*$id:pname)->$id:(tupleField i)|]             case vd of               ScalarValue{} -> return ()-              _ -> stm [C.cstm|assert(($exp:field = malloc(sizeof($ty:ct))) != NULL);|]+              _ -> stm [C.cstm|assert(($exp:field = ($ty:ct*) malloc(sizeof($ty:ct))) != NULL);|]             prepareValue field vd            return [C.cparam|$ty:ty **$id:pname|]@@ -951,7 +949,7 @@         prepareValue dest (ScalarValue _ _ name) =           stm [C.cstm|$exp:dest = $id:name;|] -        prepareValue dest (ArrayValue mem _ _ _ _ shape) = do+        prepareValue dest (ArrayValue mem _ _ _ shape) = do           stm [C.cstm|$exp:dest->mem = $id:mem;|]            let rank = length shape@@ -1039,7 +1037,7 @@   return [C.cstm|printf("#<opaque %s>", $string:desc);|] printStm (TransparentValue (ScalarValue bt ept _)) e =   return $ printPrimStm [C.cexp|stdout|] e bt ept-printStm (TransparentValue (ArrayValue _ _ _ bt ept shape)) e = do+printStm (TransparentValue (ArrayValue _ _ bt ept shape)) e = do   values_array <- publicName $ "values_" ++ name   shape_array <- publicName $ "shape_" ++ name   let num_elems = cproduct [ [C.cexp|$id:shape_array(ctx, $exp:e)[$int:i]|] | i <- [0..rank-1] ]@@ -1076,7 +1074,7 @@   item [C.citem|$ty:(primTypeToCType t) $id:dest;|]   stm $ readPrimStm dest i t ept   return ([C.cstm|;|], [C.cstm|;|], [C.cstm|;|], [C.cexp|$id:dest|])-readInput i (TransparentValue vd@(ArrayValue _ _ _ t ept dims)) = do+readInput i (TransparentValue vd@(ArrayValue _ _ t ept dims)) = do   dest <- newVName "read_value"   shape <- newVName "read_shape"   arr <- newVName "read_arr"@@ -1124,7 +1122,7 @@             TransparentValue ScalarValue{} -> do               item [C.citem|$ty:ty $id:result;|]               return ([C.cexp|$id:result|], [C.cstm|;|])-            TransparentValue (ArrayValue _ _ _ t ept dims) -> do+            TransparentValue (ArrayValue _ _ t ept dims) -> do               let name = arrayName t ept $ length dims               free_array <- publicName $ "free_" ++ name               item [C.citem|$ty:ty *$id:result;|]@@ -1307,6 +1305,7 @@       option_parser = generateOptionParser "parse_options" $ benchmarkOptions++options    let headerdefs = [C.cunit|+$esc:("#pragma once\n") $esc:("/*\n * Headers\n*/\n") $esc:("#include <stdint.h>") $esc:("#include <stddef.h>")@@ -1463,7 +1462,8 @@   |]    return $ CParts (pretty headerdefs) (pretty utildefs) (pretty clidefs) (pretty libdefs)-  where compileProg' = do+    where+      compileProg' = do           (memstructs, memfuns, memreport) <- unzip3 <$> mapM defineMemorySpace spaces            (prototypes, definitions) <- unzip <$> mapM compileFun funs@@ -1477,29 +1477,30 @@           ctx_ty <- contextType           headerDecl MiscDecl [C.cedecl|void futhark_debugging_report($ty:ctx_ty *ctx);|]           libDecl [C.cedecl|void futhark_debugging_report($ty:ctx_ty *ctx) {-  if (ctx->detail_memory) {-    $items:memreport-  }-  if (ctx->debugging) {-    $items:debugreport-  }-}|]+                      if (ctx->detail_memory) {+                        $items:memreport+                      }+                      if (ctx->debugging) {+                        $items:debugreport+                      }+                    }|]            return (prototypes, definitions, entry_points)-        funcToDef func = C.FuncDef func loc-          where loc = case func of-                        C.OldFunc _ _ _ _ _ _ l -> l-                        C.Func _ _ _ _ _ l      -> l -        builtin = cIntOps ++ cFloat32Ops ++ cFloat64Ops ++ cFloatConvOps ++-                  cFloat32Funs ++ cFloat64Funs+      funcToDef func = C.FuncDef func loc+        where loc = case func of+                       C.OldFunc _ _ _ _ _ _ l -> l+                       C.Func _ _ _ _ _ l      -> l -        panic_h = $(embedStringFile "rts/c/panic.h")-        values_h = $(embedStringFile "rts/c/values.h")-        timing_h = $(embedStringFile "rts/c/timing.h")-        lock_h = $(embedStringFile "rts/c/lock.h")-        tuning_h = $(embedStringFile "rts/c/tuning.h")+      builtin = cIntOps ++ cFloat32Ops ++ cFloat64Ops ++ cFloatConvOps +++                cFloat32Funs ++ cFloat64Funs +      panic_h  = $(embedStringFile "rts/c/panic.h")+      values_h = $(embedStringFile "rts/c/values.h")+      timing_h = $(embedStringFile "rts/c/timing.h")+      lock_h   = $(embedStringFile "rts/c/lock.h")+      tuning_h = $(embedStringFile "rts/c/tuning.h")+ compileFun :: (Name, Function op) -> CompilerM op s (C.Definition, C.Func) compileFun (fname, Function _ outputs inputs body _ _) = do   (outparams, out_ptrs) <- unzip <$> mapM compileOutput outputs@@ -1565,7 +1566,7 @@  derefPointer :: C.Exp -> C.Exp -> C.Type -> C.Exp derefPointer ptr i res_t =-  [C.cexp|*(($ty:res_t)&($exp:ptr[$exp:i]))|]+  [C.cexp|(($ty:res_t)$exp:ptr)[$exp:i]|]  writeScalarPointerWithQuals :: PointerQuals op s -> WriteScalar op s writeScalarPointerWithQuals quals_f dest i elemtype space vol v = do@@ -1707,10 +1708,15 @@               { $items:items }              |] -compileCode (DebugPrint s _ e) = do+compileCode (DebugPrint s (Just (_, e))) = do   e' <- compileExp e   stm [C.cstm|if (ctx->debugging) {           fprintf(stderr, "%s: %d\n", $exp:s, (int)$exp:e');+       }|]++compileCode (DebugPrint s Nothing) =+  stm [C.cstm|if (ctx->debugging) {+          fprintf(stderr, "%s\n", $exp:s);        }|]  compileCode c

@@ -123,7 +123,7 @@                         CompilerM op s CSExp  -- | Unpack the array being passed to an entry point.-type EntryInput op s = VName -> Imp.MemSize -> Imp.SpaceId ->+type EntryInput op s = VName -> Imp.SpaceId ->                        PrimType -> Imp.Signedness ->                        [Imp.DimSize] ->                        CSExp ->@@ -646,12 +646,12 @@   return $ cast $ Var $ compileName name   where cast = compileTypecastExt bt ept -entryPointOutput (Imp.TransparentValue (Imp.ArrayValue mem _ Imp.DefaultSpace bt ept dims)) = do+entryPointOutput (Imp.TransparentValue (Imp.ArrayValue mem Imp.DefaultSpace bt ept dims)) = do   let src = Var $ compileName mem   let createTuple = "createTuple_" ++ compilePrimTypeExt bt ept   return $ simpleCall createTuple [src, CreateArray (Primitive $ CSInt Int64T) $ Right $ map compileDim dims] -entryPointOutput (Imp.TransparentValue (Imp.ArrayValue mem _ (Imp.Space sid) bt ept dims)) = do+entryPointOutput (Imp.TransparentValue (Imp.ArrayValue mem (Imp.Space sid) bt ept dims)) = do   unRefMem mem (Imp.Space sid)   pack_output <- asks envEntryOutput   pack_output mem sid bt ept dims@@ -666,21 +666,16 @@       cast = compileTypecast bt   stm $ Assign vname' (cast e) -entryPointInput (_, Imp.TransparentValue (Imp.ArrayValue mem memsize Imp.DefaultSpace bt _ dims), e) = do+entryPointInput (_, Imp.TransparentValue (Imp.ArrayValue mem Imp.DefaultSpace bt _ dims), e) = do   zipWithM_ (unpackDim e) dims [0..]   let arrayData = Field e "Item1"   let dest = Var $ compileName mem       unwrap_call = simpleCall "unwrapArray" [arrayData, sizeOf $ compilePrimTypeToAST bt]-  case memsize of-    Imp.VarSize sizevar ->-      stm $ Assign (Var $ compileName sizevar) $ Field e "Item2.Length"-    Imp.ConstSize _ ->-      return ()   stm $ Assign dest unwrap_call -entryPointInput (_, Imp.TransparentValue (Imp.ArrayValue mem memsize (Imp.Space sid) bt ept dims), e) = do+entryPointInput (_, Imp.TransparentValue (Imp.ArrayValue mem (Imp.Space sid) bt ept dims), e) = do   unpack_input <- asks envEntryInput-  unpack <- collect $ unpack_input mem memsize sid bt ept dims e+  unpack <- collect $ unpack_input mem sid bt ept dims e   stms unpack  extValueDescName :: Imp.ExternalValue -> String@@ -706,7 +701,7 @@  valueDescVName :: Imp.ValueDesc -> VName valueDescVName (Imp.ScalarValue _ _ vname) = vname-valueDescVName (Imp.ArrayValue vname _ _ _ _ _) = vname+valueDescVName (Imp.ArrayValue vname _ _ _ _) = vname  consoleWrite :: String -> [CSExp] -> CSExp consoleWrite str exps = simpleCall "Console.Write" $ String str:exps@@ -780,7 +775,7 @@ -- TODO: If the type identifier of 'Float32' is changed, currently the error -- messages for reading binary input will not use this new name. This is also a -- problem for the C runtime system.-readInput decl@(Imp.TransparentValue (Imp.ArrayValue _ _ _ bt ept dims)) =+readInput decl@(Imp.TransparentValue (Imp.ArrayValue _ _ bt ept dims)) =   let rank' = Var $ show $ length dims       type_enum = String $ readTypeEnum bt ept       bt' =  compilePrimTypeExt bt ept@@ -802,11 +797,11 @@ printStm :: Imp.ValueDesc -> CSExp -> CSExp -> CompilerM op s CSStmt printStm Imp.ScalarValue{} _ e =   return $ printPrimStm e-printStm (Imp.ArrayValue _ _ _ _ _ []) ind e = do+printStm (Imp.ArrayValue _ _ _ _ []) ind e = do   let e' = Index e (IdxExp (PostUnOp "++" ind))   return $ printPrimStm e' -printStm (Imp.ArrayValue mem memsize space bt ept (outer:shape)) ind e = do+printStm (Imp.ArrayValue mem space bt ept (outer:shape)) ind e = do   ptr <- newVName "shapePtr"   first <- newVName "printFirst"   let size = callMethod (CreateArray (Primitive $ CSInt Int32T) $ Right $ map compileDim $ outer:shape)@@ -816,7 +811,7 @@                              ]       emptystr = "empty(" ++ ppArrayType bt (length shape) ++ ")" -  printelem <- printStm (Imp.ArrayValue mem memsize space bt ept shape) ind e+  printelem <- printStm (Imp.ArrayValue mem space bt ept shape) ind e   return $     If (BinOp "==" size (Integer 0))       [puts emptystr]@@ -926,9 +921,9 @@         initCopy (varName, Imp.MemParam _ space) = declMem' varName space         initCopy _ = Pass -        valueDescFun (Imp.ArrayValue mem _ Imp.DefaultSpace _ _ _) =+        valueDescFun (Imp.ArrayValue mem Imp.DefaultSpace _ _ _) =             stm $ Assign (Var $ compileName mem ++ "_nbytes") (Var $ compileName mem ++ ".Length")-        valueDescFun (Imp.ArrayValue mem _ (Imp.Space _) bt _ dims) =+        valueDescFun (Imp.ArrayValue mem (Imp.Space _) bt _ dims) =             stm $ Assign (Var $ compileName mem ++ "_nbytes") $ foldr (BinOp "*" . compileDim) (sizeOf $ compilePrimTypeToAST bt) dims         valueDescFun _ = stm Pass @@ -970,7 +965,7 @@         getType' :: Imp.ValueDesc -> CSType         getType' (Imp.ScalarValue primtype signedness _) =           compilePrimTypeToASText primtype signedness-        getType' (Imp.ArrayValue _ _ _ primtype signedness _) =+        getType' (Imp.ArrayValue _ _ primtype signedness _) =           let t = compilePrimTypeToASText primtype signedness           in Composite $ SystemTupleT [Composite $ ArrayT t, Composite $ ArrayT $ Primitive $ CSInt Int64T] 

@@ -100,7 +100,7 @@                         CompilerM op s PyExp  -- | Unpack the array being passed to an entry point.-type EntryInput op s = VName -> Imp.MemSize -> Imp.SpaceId ->+type EntryInput op s = VName -> Imp.SpaceId ->                        PrimType -> Imp.Signedness ->                        [Imp.DimSize] ->                        PyExp ->@@ -272,11 +272,10 @@          , optionAction =            [ Assign (Var "entry_point") $ Var "optarg" ]          },-  -- The -b option is just a dummy for now.   Option { optionLongName = "binary-output"          , optionShortName = Just 'b'          , optionArgument = NoArgument-         , optionAction = [Pass]+         , optionAction = [Assign (Var "binary_output") $ Bool True]          },   Option { optionLongName = "tuning"          , optionShortName = Nothing@@ -353,6 +352,7 @@           Assign (Var "do_warmup_run") (Bool False) :           Assign (Var "num_runs") (Integer 1) :           Assign (Var "entry_point") (String "main") :+          Assign (Var "binary_output") (Bool False) :           generateOptionParser (standardOptions ++ options)          selectEntryPoint entry_point_names entry_points =@@ -415,10 +415,10 @@ entryPointOutput (Imp.TransparentValue (Imp.ScalarValue bt ept name)) =   return $ simpleCall tf [Var $ compileName name]   where tf = compilePrimToExtNp bt ept-entryPointOutput (Imp.TransparentValue (Imp.ArrayValue mem _ Imp.DefaultSpace bt ept dims)) = do+entryPointOutput (Imp.TransparentValue (Imp.ArrayValue mem Imp.DefaultSpace bt ept dims)) = do   let cast = Cast (Var $ compileName mem) (compilePrimTypeExt bt ept)   return $ simpleCall "createArray" [cast, Tuple $ map compileDim dims]-entryPointOutput (Imp.TransparentValue (Imp.ArrayValue mem _ (Imp.Space sid) bt ept dims)) = do+entryPointOutput (Imp.TransparentValue (Imp.ArrayValue mem (Imp.Space sid) bt ept dims)) = do   pack_output <- asks envEntryOutput   pack_output mem sid bt ept dims @@ -455,7 +455,7 @@     [Catch (Tuple [Var "TypeError", Var "AssertionError"])      [badInput i e $ prettySigned (s==Imp.TypeUnsigned) bt]] -entryPointInput (i, Imp.TransparentValue (Imp.ArrayValue mem memsize Imp.DefaultSpace t s dims), e) = do+entryPointInput (i, Imp.TransparentValue (Imp.ArrayValue mem Imp.DefaultSpace t s dims), e) = do   let type_is_wrong =         UnOp "not" $         BinOp "and"@@ -470,18 +470,11 @@   let dest = Var $ compileName mem       unwrap_call = simpleCall "unwrapArray" [e] -  case memsize of-    Imp.VarSize sizevar ->-      stm $ Assign (Var $ compileName sizevar) $-      simpleCall "np.int32" [Field e "nbytes"]-    Imp.ConstSize _ ->-      return ()-   stm $ Assign dest unwrap_call -entryPointInput (i, Imp.TransparentValue (Imp.ArrayValue mem memsize (Imp.Space sid) bt ept dims), e) = do+entryPointInput (i, Imp.TransparentValue (Imp.ArrayValue mem (Imp.Space sid) bt ept dims), e) = do   unpack_input <- asks envEntryInput-  unpack <- collect $ unpack_input mem memsize sid bt ept dims e+  unpack <- collect $ unpack_input mem sid bt ept dims e   stm $ Try unpack     [Catch (Tuple [Var "TypeError", Var "AssertionError"])      [badInput i e $ concat (replicate (length dims) "[]") ++@@ -501,7 +494,7 @@  valueDescVName :: Imp.ValueDesc -> VName valueDescVName (Imp.ScalarValue _ _ vname) = vname-valueDescVName (Imp.ArrayValue vname _ _ _ _ _) = vname+valueDescVName (Imp.ArrayValue vname _ _ _ _) = vname  -- Key into the FUTHARK_PRIMTYPES dict. readTypeEnum :: PrimType -> Imp.Signedness -> String@@ -527,7 +520,7 @@   let type_name = readTypeEnum bt ept   in Assign (Var $ extValueDescName decl) $ simpleCall "read_value" [String type_name] -readInput decl@(Imp.TransparentValue (Imp.ArrayValue _ _ _ bt ept dims)) =+readInput decl@(Imp.TransparentValue (Imp.ArrayValue _ _ bt ept dims)) =   let type_name = readTypeEnum bt ept   in Assign (Var $ extValueDescName decl) $ simpleCall "read_value"      [String $ concat (replicate (length dims) "[]") ++ type_name]@@ -542,11 +535,13 @@   where printValue' (Imp.OpaqueValue desc _) _ =           return [Exp $ simpleCall "sys.stdout.write"                   [String $ "#<opaque " ++ desc ++ ">"]]-        printValue' (Imp.TransparentValue (Imp.ArrayValue mem memsize (Space _) bt ept shape)) e =-          printValue' (Imp.TransparentValue (Imp.ArrayValue mem memsize DefaultSpace bt ept shape)) $+        printValue' (Imp.TransparentValue (Imp.ArrayValue mem (Space _) bt ept shape)) e =+          printValue' (Imp.TransparentValue (Imp.ArrayValue mem DefaultSpace bt ept shape)) $           simpleCall (pretty e ++ ".get") []         printValue' (Imp.TransparentValue _) e =-          return [Exp $ simpleCall "write_value" [e],+          return [Exp $ Call (Var "write_value")+                   [Arg e,+                    ArgKeyword "binary" (Var "binary_output")],                   Exp $ simpleCall "sys.stdout.write" [String "\n"]]  prepareEntry :: (Name, Imp.Function op) -> CompilerM op s@@ -616,7 +611,7 @@                    map desc $ Imp.functionResult func)   where desc (Imp.OpaqueValue d _) = d         desc (Imp.TransparentValue (Imp.ScalarValue pt s _)) = readTypeEnum pt s-        desc (Imp.TransparentValue (Imp.ArrayValue _ _ _ pt s dims)) =+        desc (Imp.TransparentValue (Imp.ArrayValue _ _ pt s dims)) =           concat (replicate (length dims) "[]") ++ readTypeEnum pt s  callEntryFun :: [PyStmt] -> (Name, Imp.Function op)@@ -751,7 +746,7 @@     (IntType Int64, _) -> "np.int64"     (FloatType Float32, _) -> "np.float32"     (FloatType Float64, _) -> "np.float64"-    (Imp.Bool, _) -> "np.bool"+    (Imp.Bool, _) -> "np.bool_"     (Cert, _) -> "np.byte"  compilePrimValue :: Imp.PrimValue -> PyExp

@@ -179,7 +179,7 @@               [Arg val, ArgKeyword "dtype" $ Var $ Py.compilePrimType bt]   Py.stm $ Exp $ Call (Var "cl.enqueue_copy")     [Arg $ Var "self.queue", Arg mem', Arg nparr,-     ArgKeyword "device_offset" $ asLong i,+     ArgKeyword "device_offset" $ BinOp "*" (asLong i) (Integer $ Imp.primByteSize bt),      ArgKeyword "is_blocking" $ Var "synchronous"]  writeOpenCLScalar _ _ _ space _ =@@ -196,7 +196,7 @@   Py.stm $ Assign val' nparr   Py.stm $ Exp $ Call (Var "cl.enqueue_copy")     [Arg $ Var "self.queue", Arg val', Arg mem',-     ArgKeyword "device_offset" $ asLong i,+     ArgKeyword "device_offset" $ BinOp "*" (asLong i) (Integer $ Imp.primByteSize bt),      ArgKeyword "is_blocking" $ Bool True]   return $ Index val' $ IdxExp $ Integer 0 @@ -303,7 +303,7 @@   fail $ "Cannot return array from " ++ sid ++ " space."  unpackArrayInput :: Py.EntryInput Imp.OpenCL ()-unpackArrayInput mem memsize "device" t s dims e = do+unpackArrayInput mem "device" t s dims e = do   let type_is_ok =         BinOp "and"         (BinOp "in" (Py.simpleCall "type" [e]) (List [Var "np.ndarray", Var "cl.array.Array"]))@@ -312,14 +312,7 @@    zipWithM_ (Py.unpackDim e) dims [0..] -  case memsize of-    Imp.VarSize sizevar ->-      Py.stm $ Assign (Var $ Py.compileName sizevar) $-      Py.simpleCall "np.int64" [Field e "nbytes"]-    Imp.ConstSize _ ->-      return ()--  let memsize' = Py.compileDim memsize+  let memsize' = Py.simpleCall "np.int64" [Field e "nbytes"]       pyOpenCLArrayCase =         [Assign mem_dest $ Field e "data"]   numpyArrayCase <- Py.collect $ do@@ -335,7 +328,7 @@     pyOpenCLArrayCase     numpyArrayCase   where mem_dest = Var $ Py.compileName mem-unpackArrayInput _ _ sid _ _ _ _ =+unpackArrayInput _ sid _ _ _ _ =   fail $ "Cannot accept array from " ++ sid ++ " space."  ifNotZeroSize :: PyExp -> PyStmt -> PyStmt

@@ -38,7 +38,7 @@           GC.publicDef_ "context_config_new" GC.InitDecl $ \s ->             ([C.cedecl|struct $id:cfg* $id:s();|],              [C.cedecl|struct $id:cfg* $id:s() {-                                 struct $id:cfg *cfg = malloc(sizeof(struct $id:cfg));+                                 struct $id:cfg *cfg = (struct $id:cfg*) malloc(sizeof(struct $id:cfg));                                  if (cfg == NULL) {                                    return NULL;                                  }@@ -80,7 +80,7 @@           GC.publicDef_ "context_new" GC.InitDecl $ \s ->             ([C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg);|],              [C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg) {-                                  struct $id:ctx* ctx = malloc(sizeof(struct $id:ctx));+                                  struct $id:ctx* ctx = (struct $id:ctx*) malloc(sizeof(struct $id:ctx));                                   if (ctx == NULL) {                                     return NULL;                                   }

@@ -24,7 +24,7 @@ import qualified Language.C.Quote.C as C  import Futhark.CodeGen.ImpCode-import Futhark.Util.Pretty (pretty)+import Futhark.Util.Pretty (pretty, prettyOneLine) import Futhark.Util (zEncodeString)  -- | The C type corresponding to a signed integer type.@@ -158,6 +158,12 @@         mkSlt = intCmpOp "slt" [C.cexp|x < y|]         mkSle = intCmpOp "sle" [C.cexp|x <= y|] +        -- We define some operations as macros rather than functions,+        -- because this allows us to use them as constant expressions+        -- in things like array sizes and static initialisers.+        macro name rhs =+          [C.cedecl|$esc:("#define " ++ name ++ "(x) (" ++ prettyOneLine rhs ++ ")")|]+         mkPow t =           let ct = intTypeToCType t           in [C.cedecl|static inline $ty:ct $id:(taggedI "pow" t)($ty:ct x, $ty:ct y) {@@ -172,16 +178,12 @@                          return res;               }|] -        mkSExt from_t to_t =-          [C.cedecl|static inline $ty:to_ct-                    $id:name($ty:from_ct x) { return x;} |]+        mkSExt from_t to_t = macro name [C.cexp|($ty:to_ct)(($ty:from_ct)x)|]           where name = "sext_"++pretty from_t++"_"++pretty to_t                 from_ct = intTypeToCType from_t                 to_ct = intTypeToCType to_t -        mkZExt from_t to_t =-          [C.cedecl|static inline $ty:to_ct-                    $id:name($ty:from_ct x) { return x;} |]+        mkZExt from_t to_t = macro name [C.cexp|($ty:to_ct)(($ty:from_ct)x)|]           where name = "zext_"++pretty from_t++"_"++pretty to_t                 from_ct = uintTypeToCType from_t                 to_ct = uintTypeToCType to_t@@ -315,6 +317,14 @@       return atan2(x,y);     } +    static inline float $id:(funName' "gamma32")(float x) {+      return tgamma(x);+    }++    static inline float $id:(funName' "lgamma32")(float x) {+      return lgamma(x);+    }+     static inline float $id:(funName' "round32")(float x) {       return rint(x);     }@@ -394,6 +404,14 @@      static inline double $id:(funName' "atan2_64")(double x, double y) {       return atan2(x,y);+    }++    static inline double $id:(funName' "gamma64")(double x) {+      return tgamma(x);+    }++    static inline double $id:(funName' "lgamma64")(double x) {+      return lgamma(x);     }      static inline double $id:(funName' "round64")(double x) {

@@ -31,6 +31,7 @@   , Volatility (..)   , Arg (..)   , var+  , vi32   , index   , ErrorMsg(..)   , ErrorMsgPart(..)@@ -105,7 +106,7 @@                 deriving (Eq, Show)  -- | A description of an externally meaningful value.-data ValueDesc = ArrayValue VName MemSize Space PrimType Signedness [DimSize]+data ValueDesc = ArrayValue VName Space PrimType Signedness [DimSize]                -- ^ An array with memory block, memory block size,                -- memory space, element type, signedness of element                -- type (if applicable), and shape.@@ -174,7 +175,7 @@               -- ^ Destination, offset in destination, destination               -- space, source, offset in source, offset space, number               -- of bytes.-            | Write VName (Count Bytes) PrimType Space Volatility Exp+            | Write VName (Count Elements) PrimType Space Volatility Exp             | SetScalar VName Exp             | SetMem VName VName Space               -- ^ Must be in same space.@@ -185,12 +186,13 @@               -- ^ Has the same semantics as the contained code, but               -- the comment should show up in generated code for ease               -- of inspection.-            | DebugPrint String PrimType Exp+            | DebugPrint String (Maybe (PrimType, Exp))               -- ^ Print the given value (of the given type) to the               -- screen, somehow annotated with the given string as a-              -- description.  This has no semantic meaning, but is-              -- used entirely for debugging.  Code generators are-              -- free to ignore this statement.+              -- description.  If no type/value pair, just print the+              -- string.  This has no semantic meaning, but is used+              -- entirely for debugging.  Code generators are free to+              -- ignore this statement.             | Op a             deriving (Show) @@ -208,7 +210,7 @@  data ExpLeaf = ScalarVar VName              | SizeOf PrimType-             | Index VName (Count Bytes) PrimType Space Volatility+             | Index VName (Count Elements) PrimType Space Volatility            deriving (Eq, Show)  type Exp = PrimExp ExpLeaf@@ -253,7 +255,11 @@ var :: VName -> PrimType -> Exp var = LeafExp . ScalarVar -index :: VName -> Count Bytes -> PrimType -> Space -> Volatility -> Exp+-- | Turn a 'VName' into a 'int32' 'Imp.ScalarVar'.+vi32 :: VName -> Exp+vi32 = flip var $ IntType Int32++index :: VName -> Count Elements -> PrimType -> Space -> Volatility -> Exp index arr i t s vol = LeafExp (Index arr i t s vol) t  -- Prettyprinting definitions.@@ -286,8 +292,8 @@     ppr t <+> ppr name <> ept'     where ept' = case ept of TypeUnsigned -> text " (unsigned)"                              TypeDirect   -> mempty-  ppr (ArrayValue mem memsize space et ept shape) =-    foldr f (ppr et) shape <+> text "at" <+> ppr mem <> parens (ppr memsize) <> space' <+> ept'+  ppr (ArrayValue mem space et ept shape) =+    foldr f (ppr et) shape <+> text "at" <+> ppr mem <> space' <+> ept'     where f e s = brackets $ s <> comma <> ppr e           ept' = case ept of TypeUnsigned -> text " (unsigned)"                              TypeDirect   -> mempty@@ -361,8 +367,10 @@     ppr fname <> parens (commasep $ map ppr args)   ppr (Comment s code) =     text "--" <+> text s </> ppr code-  ppr (DebugPrint desc pt e) =+  ppr (DebugPrint desc (Just (pt, e))) =     text "debug" <+> parens (commasep [text (show desc), ppr pt, ppr e])+  ppr (DebugPrint desc Nothing) =+    text "debug" <+> parens (text (show desc))  instance Pretty Arg where   ppr (MemArg m) = ppr m@@ -445,8 +453,8 @@     pure $ Call dests fname args   traverse f (Comment s code) =     Comment s <$> traverse f code-  traverse _ (DebugPrint s t e) =-    pure $ DebugPrint s t e+  traverse _ (DebugPrint s v) =+    pure $ DebugPrint s v  declaredIn :: Code a -> Names declaredIn (DeclareMem name _) = S.singleton name@@ -496,8 +504,8 @@     freeIn op   freeIn (Comment _ code) =     freeIn code-  freeIn (DebugPrint _ _ e) =-    freeIn e+  freeIn (DebugPrint _ v) =+    maybe mempty (freeIn . snd) v  instance FreeIn ExpLeaf where   freeIn (Index v e _ _ _) = freeIn v <> freeIn e

@@ -56,8 +56,6 @@ -- | A generic kernel containing arbitrary kernel code. data Kernel = Kernel               { kernelBody :: Imp.Code KernelOp-              , kernelLocalMemory :: [LocalMemoryUse]-              -- ^ The local memory used by this kernel.                , kernelUses :: [KernelUse]                 -- ^ The host variables referenced by the kernel.@@ -71,7 +69,7 @@             deriving (Show)  -- ^ In-kernel name and per-workgroup size in bytes.-type LocalMemoryUse = (VName, Either MemSize KernelConstExp)+type LocalMemoryUse = (VName, Either (Count Bytes) KernelConstExp)  data KernelUse = ScalarUse VName PrimType                | MemoryUse VName@@ -87,7 +85,7 @@         sameKernel _ _ = False  -- | Get an atomic operator corresponding to a binary operator.-atomicBinOp :: BinOp -> Maybe (VName -> VName -> Count Bytes -> Exp -> AtomicOp)+atomicBinOp :: BinOp -> Maybe (VName -> VName -> Count Elements -> Exp -> AtomicOp) atomicBinOp = flip lookup [ (Add Int32, AtomicAdd)                           , (SMax Int32, AtomicSMax)                           , (SMin Int32, AtomicSMin)@@ -140,15 +138,8 @@     text "kernel" <+> brace     (text "groups" <+> brace (ppr $ kernelNumGroups kernel) </>      text "group_size" <+> brace (ppr $ kernelGroupSize kernel) </>-     text "local_memory" <+> brace (commasep $-                                    map ppLocalMemory $-                                    kernelLocalMemory kernel) </>      text "uses" <+> brace (commasep $ map ppr $ kernelUses kernel) </>      text "body" <+> brace (ppr $ kernelBody kernel))-    where ppLocalMemory (name, Left size) =-            ppr name <+> parens (ppr size <+> text "bytes")-          ppLocalMemory (name, Right size) =-            ppr name <+> parens (ppr size <+> text "bytes (const)")  data KernelOp = GetGroupId VName Int               | GetLocalId VName Int@@ -156,24 +147,27 @@               | GetGlobalSize VName Int               | GetGlobalId VName Int               | GetLockstepWidth VName-              | Atomic AtomicOp+              | Atomic Space AtomicOp               | LocalBarrier               | GlobalBarrier-              | MemFence+              | MemFenceLocal+              | MemFenceGlobal+              | PrivateAlloc VName (Count Bytes)+              | LocalAlloc VName (Either (Count Bytes) KernelConstExp)               deriving (Show)  -- Atomic operations return the value stored before the update. -- This value is stored in the first VName.-data AtomicOp = AtomicAdd VName VName (Count Bytes) Exp-              | AtomicSMax VName VName (Count Bytes) Exp-              | AtomicSMin VName VName (Count Bytes) Exp-              | AtomicUMax VName VName (Count Bytes) Exp-              | AtomicUMin VName VName (Count Bytes) Exp-              | AtomicAnd VName VName (Count Bytes) Exp-              | AtomicOr VName VName (Count Bytes) Exp-              | AtomicXor VName VName (Count Bytes) Exp-              | AtomicCmpXchg VName VName (Count Bytes) Exp Exp-              | AtomicXchg VName VName (Count Bytes) Exp+data AtomicOp = AtomicAdd VName VName (Count Elements) Exp+              | AtomicSMax VName VName (Count Elements) Exp+              | AtomicSMin VName VName (Count Elements) Exp+              | AtomicUMax VName VName (Count Elements) Exp+              | AtomicUMin VName VName (Count Elements) Exp+              | AtomicAnd VName VName (Count Elements) Exp+              | AtomicOr VName VName (Count Elements) Exp+              | AtomicXor VName VName (Count Elements) Exp+              | AtomicCmpXchg VName VName (Count Elements) Exp Exp+              | AtomicXchg VName VName (Count Elements) Exp               deriving (Show)  instance FreeIn AtomicOp where@@ -211,41 +205,49 @@     text "local_barrier()"   ppr GlobalBarrier =     text "global_barrier()"-  ppr MemFence =-    text "mem_fence()"-  ppr (Atomic (AtomicAdd old arr ind x)) =+  ppr MemFenceLocal =+    text "mem_fence_local()"+  ppr MemFenceGlobal =+    text "mem_fence_global()"+  ppr (PrivateAlloc name size) =+    ppr name <+> equals <+> text "private_alloc" <> parens (ppr size)+  ppr (LocalAlloc name size) =+    ppr name <+> equals <+> text "local_alloc" <>+    parens (either ppr constCase size)+    where constCase e = text "(constant)" <+> ppr e+  ppr (Atomic _ (AtomicAdd old arr ind x)) =     ppr old <+> text "<-" <+> text "atomic_add" <>     parens (commasep [ppr arr <> brackets (ppr ind), ppr x])-  ppr (Atomic (AtomicSMax old arr ind x)) =+  ppr (Atomic _ (AtomicSMax old arr ind x)) =     ppr old <+> text "<-" <+> text "atomic_smax" <>     parens (commasep [ppr arr <> brackets (ppr ind), ppr x])-  ppr (Atomic (AtomicSMin old arr ind x)) =+  ppr (Atomic _ (AtomicSMin old arr ind x)) =     ppr old <+> text "<-" <+> text "atomic_smin" <>     parens (commasep [ppr arr <> brackets (ppr ind), ppr x])-  ppr (Atomic (AtomicUMax old arr ind x)) =+  ppr (Atomic _ (AtomicUMax old arr ind x)) =     ppr old <+> text "<-" <+> text "atomic_umax" <>     parens (commasep [ppr arr <> brackets (ppr ind), ppr x])-  ppr (Atomic (AtomicUMin old arr ind x)) =+  ppr (Atomic _ (AtomicUMin old arr ind x)) =     ppr old <+> text "<-" <+> text "atomic_umin" <>     parens (commasep [ppr arr <> brackets (ppr ind), ppr x])-  ppr (Atomic (AtomicAnd old arr ind x)) =+  ppr (Atomic _ (AtomicAnd old arr ind x)) =     ppr old <+> text "<-" <+> text "atomic_and" <>     parens (commasep [ppr arr <> brackets (ppr ind), ppr x])-  ppr (Atomic (AtomicOr old arr ind x)) =+  ppr (Atomic _ (AtomicOr old arr ind x)) =     ppr old <+> text "<-" <+> text "atomic_or" <>     parens (commasep [ppr arr <> brackets (ppr ind), ppr x])-  ppr (Atomic (AtomicXor old arr ind x)) =+  ppr (Atomic _ (AtomicXor old arr ind x)) =     ppr old <+> text "<-" <+> text "atomic_xor" <>     parens (commasep [ppr arr <> brackets (ppr ind), ppr x])-  ppr (Atomic (AtomicCmpXchg old arr ind x y)) =+  ppr (Atomic _ (AtomicCmpXchg old arr ind x y)) =     ppr old <+> text "<-" <+> text "atomic_cmp_xchg" <>     parens (commasep [ppr arr <> brackets (ppr ind), ppr x, ppr y])-  ppr (Atomic (AtomicXchg old arr ind x)) =+  ppr (Atomic _ (AtomicXchg old arr ind x)) =     ppr old <+> text "<-" <+> text "atomic_xchg" <>     parens (commasep [ppr arr <> brackets (ppr ind), ppr x])  instance FreeIn KernelOp where-  freeIn (Atomic op) = freeIn op+  freeIn (Atomic _ op) = freeIn op   freeIn _ = mempty  brace :: Doc -> Doc

@@ -10,6 +10,7 @@   , ExpCompiler   , CopyCompiler   , StmsCompiler+  , AllocCompiler   , Operations (..)   , defaultOperations   , ValueDestination@@ -30,6 +31,7 @@   , emitFunction   , hasFunction   , collect+  , collect'   , comment   , VarEntry (..)   , ArrayEntry (..)@@ -40,10 +42,7 @@   , lookupMemory      -- * Building Blocks-  , compileSubExp-  , compileSubExpOfType-  , compileSubExpTo-  , compilePrimExp+  , ToExp(..)   , compileAlloc   , subExpToDimSize   , everythingVolatile@@ -58,13 +57,13 @@   , strideArray   , fullyIndexArray   , fullyIndexArray'-  , varIndex   , Imp.dimSizeToExp   , dimSizeToSubExp   , copy   , copyDWIM   , copyDWIMDest   , copyElementWise+  , typeSize    -- * Constructing code.   , dLParams@@ -72,15 +71,16 @@   , dScope   , dScopes   , dArray-  , dPrim, dPrim_, dPrimV+  , dPrim, dPrim_, dPrimV_, dPrimV    , sFor, sWhile   , sComment   , sIf, sWhen, sUnless   , sOp   , sDeclareMem, sAlloc, sAlloc_-  , sArray, sAllocArray, sStaticArray+  , sArray, sAllocArray, sAllocArrayPerm, sStaticArray   , sWrite, sUpdate+  , sLoopNest   , (<--)   )   where@@ -101,7 +101,7 @@ import Futhark.CodeGen.ImpCode   (Count (..),    Bytes, Elements,-   bytes, withElemType)+   bytes, elements, withElemType) import Futhark.Representation.ExplicitMemory import Futhark.Representation.SOACS (SOACS) import qualified Futhark.Representation.ExplicitMemory.IndexFunction as IxFun@@ -114,7 +114,7 @@ type OpCompiler lore op = Pattern lore -> Op lore -> ImpM lore op ()  -- | How to compile some 'Stms'.-type StmsCompiler lore op = Names -> [Stm lore] -> ImpM lore op () -> ImpM lore op ()+type StmsCompiler lore op = Names -> Stms lore -> ImpM lore op () -> ImpM lore op ()  -- | How to compile an 'Exp'. type ExpCompiler lore op = Pattern lore -> Exp lore -> ImpM lore op ()@@ -125,10 +125,14 @@                            -> Count Elements -- ^ Number of row elements of the source.                            -> ImpM lore op () +-- | An alternate way of compiling an allocation.+type AllocCompiler lore op = VName -> Count Bytes -> ImpM lore op ()+ data Operations lore op = Operations { opsExpCompiler :: ExpCompiler lore op                                      , opsOpCompiler :: OpCompiler lore op                                      , opsStmsCompiler :: StmsCompiler lore op                                      , opsCopyCompiler :: CopyCompiler lore op+                                     , opsAllocCompilers :: M.Map Space (AllocCompiler lore op)                                      }  -- | An operations set for which the expression compiler always@@ -139,6 +143,7 @@                                    , opsOpCompiler = opc                                    , opsStmsCompiler = defCompileStms                                    , opsCopyCompiler = defaultCopy+                                   , opsAllocCompilers = mempty                                    }  -- | When an array is dared, this is where it is stored.@@ -157,10 +162,7 @@ entryArrayShape :: ArrayEntry -> [Imp.DimSize] entryArrayShape = memLocationShape . entryArrayLocation -data MemEntry = MemEntry {-      entryMemSize  :: Imp.MemSize-    , entryMemSpace :: Imp.Space-  }+newtype MemEntry = MemEntry { entryMemSpace :: Imp.Space }   deriving (Show)  newtype ScalarEntry = ScalarEntry {@@ -200,23 +202,22 @@   , envStmsCompiler :: StmsCompiler lore op   , envOpCompiler :: OpCompiler lore op   , envCopyCompiler :: CopyCompiler lore op+  , envAllocCompilers :: M.Map Space (AllocCompiler lore op)   , envDefaultSpace :: Imp.Space   , envVolatility :: Imp.Volatility-  , envFakeMemory :: [Space]-    -- ^ Do not actually generate allocations for these memory spaces.   , envFunction :: Name     -- ^ Name of the function we are compiling.   } -newEnv :: Operations lore op -> Imp.Space -> [Imp.Space] -> Name -> Env lore op-newEnv ops ds fake fname =+newEnv :: Operations lore op -> Imp.Space -> Name -> Env lore op+newEnv ops ds fname =   Env { envExpCompiler = opsExpCompiler ops       , envStmsCompiler = opsStmsCompiler ops       , envOpCompiler = opsOpCompiler ops       , envCopyCompiler = opsCopyCompiler ops+      , envAllocCompilers = mempty       , envDefaultSpace = ds       , envVolatility = Imp.Nonvolatile-      , envFakeMemory = fake       , envFunction = fname       } @@ -245,10 +246,12 @@   getNameSource = gets stateNameSource   putNameSource src = modify $ \s -> s { stateNameSource = src } +-- Cannot be an ExplicitMemory scope because the index functions have+-- the wrong leaves (VName instead of Imp.Exp). instance HasScope SOACS (ImpM lore op) where   askScope = M.map (LetInfo . entryType) <$> gets stateVTable     where entryType (MemVar _ memEntry) =-            Mem (dimSizeToSubExp $ entryMemSize memEntry) (entryMemSpace memEntry)+            Mem (entryMemSpace memEntry)           entryType (ArrayVar _ arrayEntry) =             Array             (entryArrayElemType arrayEntry)@@ -258,10 +261,9 @@             Prim $ entryScalarType scalarEntry  runImpM :: ImpM lore op a-        -> Operations lore op -> Imp.Space -> [Imp.Space] -> Name -> State lore op+        -> Operations lore op -> Imp.Space -> Name -> State lore op         -> Either InternalError (a, State lore op, Imp.Code op)-runImpM (ImpM m) comp space fake fname =-  runRWST m (newEnv comp space fake fname)+runImpM (ImpM m) ops space fname = runRWST m $ newEnv ops space fname  subImpM_ :: Operations lore' op' -> ImpM lore' op' a          -> ImpM lore op (Imp.Code op')@@ -276,6 +278,7 @@                      , envStmsCompiler = opsStmsCompiler ops                      , envCopyCompiler = opsCopyCompiler ops                      , envOpCompiler = opsOpCompiler ops+                     , envAllocCompilers = opsAllocCompilers ops                      }                  s { stateVTable = M.map scrubExps $ stateVTable s                    , stateFunctions = mempty } of@@ -321,16 +324,16 @@                                  in isJust $ lookup fname fs  compileProg :: (ExplicitMemorish lore, MonadFreshNames m) =>-               Operations lore op -> Imp.Space -> [Imp.Space]+               Operations lore op -> Imp.Space             -> Prog lore -> m (Either InternalError (Imp.Functions op))-compileProg ops space fake prog =+compileProg ops space prog =   modifyNameSource $ \src ->   case foldM compileFunDef' (newState src) (progFunctions prog) of     Left err -> (Left err, src)     Right s -> (Right $ stateFunctions s, stateNameSource s)   where compileFunDef' s fdef = do           ((), s', _) <--            runImpM (compileFunDef fdef) ops space fake (funDefName fdef) s+            runImpM (compileFunDef fdef) ops space (funDefName fdef) s           return s'  compileInParam :: ExplicitMemorish lore =>@@ -338,20 +341,18 @@ compileInParam fparam = case paramAttr fparam of   MemPrim bt ->     return $ Left $ Imp.ScalarParam name bt-  MemMem _ space ->+  MemMem space ->     return $ Left $ Imp.MemParam name space   MemArray bt shape _ (ArrayIn mem ixfun) -> do     shape' <- mapM subExpToDimSize $ shapeDims shape     return $ Right $ ArrayDecl name bt $-      MemLocation mem shape' $ fmap compilePrimExp ixfun+      MemLocation mem shape' $ fmap (toExp' int32) ixfun   where name = paramName fparam  data ArrayDecl = ArrayDecl VName PrimType MemLocation  fparamSizes :: Typed attr => Param attr -> S.Set VName-fparamSizes fparam-  | Mem (Var size) _ <- paramType fparam = S.singleton size-  | otherwise = S.fromList $ subExpVars $ arrayDims $ paramType fparam+fparamSizes = S.fromList . subExpVars . arrayDims . paramType  compileInParams :: ExplicitMemorish lore =>                    [FParam lore] -> [EntryPointType]@@ -365,21 +366,19 @@        summaries = M.fromList $ mapMaybe memSummary params         where memSummary param-                | MemMem (Constant (IntValue (Int64Value size))) space <- paramAttr param =-                    Just (paramName param, (Imp.ConstSize size, space))-                | MemMem (Var size) space <- paramAttr param =-                    Just (paramName param, (Imp.VarSize size, space))+                | MemMem space <- paramAttr param =+                    Just (paramName param, space)                 | otherwise =                     Nothing -      findMemInfo :: VName -> Maybe (Imp.MemSize, Space)+      findMemInfo :: VName -> Maybe Space       findMemInfo = flip M.lookup summaries        mkValueDesc fparam signedness =         case (findArray $ paramName fparam, paramType fparam) of           (Just (ArrayDecl _ bt (MemLocation mem shape _)), _) -> do-            (memsize, memspace) <- findMemInfo mem-            Just $ Imp.ArrayValue mem memsize memspace bt signedness shape+            memspace <- findMemInfo mem+            Just $ Imp.ArrayValue mem memspace bt signedness shape           (_, Prim bt)             | paramName fparam `S.member` sizes ->               Nothing@@ -440,18 +439,16 @@           return (Imp.ScalarValue t ept out, ScalarDestination out)         mkParam (MemArray t shape _ attr) ept = do           space <- asks envDefaultSpace-          (memout, memsize) <- case attr of-            ReturnsNewBlock _ x x_size _ixfun -> do+          memout <- case attr of+            ReturnsNewBlock _ x _ixfun -> do               memout <- imp $ newVName "out_mem"-              sizeout <- ensureMemSizeOut x_size               tell ([Imp.MemParam memout space],                     M.singleton x $ MemoryDestination memout)-              return (memout, sizeout)-            ReturnsInBlock memout _ -> do-              memsize <- imp $ entryMemSize <$> lookupMemory memout-              return (memout, memsize)+              return memout+            ReturnsInBlock memout _ ->+              return memout           resultshape <- mapM inspectExtSize $ shapeDims shape-          return (Imp.ArrayValue memout memsize space t ept resultshape,+          return (Imp.ArrayValue memout space t ept resultshape,                   ArrayDestination Nothing)          inspectExtSize (Ext x) = do@@ -468,19 +465,6 @@         inspectExtSize (Free se) =           imp $ subExpToDimSize se -        -- | Return the name of the out-parameter for the memory size-        -- 'x', creating it if it does not already exist.-        ensureMemSizeOut (Ext x) = do-          (memseen, arrseen) <- get-          case M.lookup x memseen of-            Nothing -> do sizeout <- imp $ newVName "out_memsize"-                          tell ([Imp.ScalarParam sizeout int64],-                                M.singleton x $ ScalarDestination sizeout)-                          put (M.insert x sizeout memseen, arrseen)-                          return $ Imp.VarSize sizeout-            Just sizeout -> return $ Imp.VarSize sizeout-        ensureMemSizeOut (Free v) = imp $ subExpToDimSize v- compileFunDef :: ExplicitMemorish lore =>                  FunDef lore               -> ImpM lore op ()@@ -496,7 +480,7 @@           addArrays arrayds            let Body _ stms ses = body-          compileStms (freeIn ses) (stmsToList stms) $+          compileStms (freeIn ses) stms $             forM_ (zip dests ses) $ \(d, se) -> copyDWIMDest d [] se []            return (outparams, inparams, results, args)@@ -504,50 +488,49 @@ compileBody :: (ExplicitMemorish lore) => Pattern lore -> Body lore -> ImpM lore op () compileBody pat (Body _ bnds ses) = do   Destination _ dests <- destinationFromPattern pat-  compileStms (freeIn ses) (stmsToList bnds) $+  compileStms (freeIn ses) bnds $     forM_ (zip dests ses) $ \(d, se) -> copyDWIMDest d [] se []  compileBody' :: (ExplicitMemorish lore, attr ~ LetAttr lore)              => [Param attr] -> Body lore -> ImpM lore op () compileBody' = compileBody . patternFromParams -compileLoopBody :: [VName] -> Body lore -> ImpM lore op ()-compileLoopBody mergenames (Body _ bnds ses) = do+compileLoopBody :: Typed attr => [Param attr] -> Body lore -> ImpM lore op ()+compileLoopBody mergeparams (Body _ bnds ses) = do   -- We cannot write the results to the merge parameters immediately,   -- as some of the results may actually *be* merge parameters, and   -- would thus be clobbered.  Therefore, we first copy to new   -- variables mirroring the merge parameters, and then copy this   -- buffer to the merge parameters.  This is efficient, because the   -- operations are all scalar operations.-  tmpnames <- mapM (newVName . (++"_tmp") . baseString) mergenames-  compileStms (freeIn ses) (stmsToList bnds) $ do-    copy_to_merge_params <- forM (zip3 mergenames tmpnames ses) $ \(d,tmp,se) ->-      subExpType se >>= \case-        Prim bt  -> do-          se' <- compileSubExp se-          emit $ Imp.DeclareScalar tmp bt-          emit $ Imp.SetScalar tmp se'-          return $ emit $ Imp.SetScalar d $ Imp.var tmp bt-        Mem _ space | Var v <- se -> do+  tmpnames <- mapM (newVName . (++"_tmp") . baseString . paramName) mergeparams+  compileStms (freeIn ses) bnds $ do+    copy_to_merge_params <- forM (zip3 mergeparams tmpnames ses) $ \(p,tmp,se) ->+      case typeOf p of+        Prim pt  -> do+          emit $ Imp.DeclareScalar tmp pt+          emit $ Imp.SetScalar tmp $ toExp' pt se+          return $ emit $ Imp.SetScalar (paramName p) $ Imp.var tmp pt+        Mem space | Var v <- se -> do           emit $ Imp.DeclareMem tmp space           emit $ Imp.SetMem tmp v space-          return $ emit $ Imp.SetMem d tmp space+          return $ emit $ Imp.SetMem (paramName p) tmp space         _ -> return $ return ()     sequence_ copy_to_merge_params -compileStms :: Names -> [Stm lore] -> ImpM lore op () -> ImpM lore op ()+compileStms :: Names -> Stms lore -> ImpM lore op () -> ImpM lore op () compileStms alive_after_stms all_stms m = do   cb <- asks envStmsCompiler   cb alive_after_stms all_stms m  defCompileStms :: (ExplicitMemorish lore, FreeIn op) =>-                  Names -> [Stm lore] -> ImpM lore op () -> ImpM lore op ()+                  Names -> Stms lore -> ImpM lore op () -> ImpM lore op () defCompileStms alive_after_stms all_stms m =   -- We keep track of any memory blocks produced by the statements,   -- and after the last time that memory block is used, we insert a   -- Free.  This is very conservative, but can cut down on lifetimes   -- in some cases.-  void $ compileStms' mempty all_stms+  void $ compileStms' mempty $ stmsToList all_stms   where compileStms' allocs (Let pat _ e:bs) = do           dVars (Just e) (patternElements pat) @@ -569,8 +552,8 @@          patternAllocs = S.fromList . mapMaybe isMemPatElem . patternElements         isMemPatElem pe = case patElemType pe of-                            Mem _ space -> Just (patElemName pe, space)-                            _           -> Nothing+                            Mem space -> Just (patElemName pe, space)+                            _         -> Nothing  compileExp :: Pattern lore -> Exp lore -> ImpM lore op () compileExp pat e = do@@ -581,10 +564,9 @@                  Pattern lore -> Exp lore -> ImpM lore op ()  defCompileExp pat (If cond tbranch fbranch _) = do-  cond' <- compileSubExp cond   tcode <- collect $ compileBody pat tbranch   fcode <- collect $ compileBody pat fbranch-  emit $ Imp.If cond' tcode fcode+  emit $ Imp.If (toExp' Bool cond) tcode fcode  defCompileExp pat (Apply fname args _ _) = do   dest <- destinationFromPattern pat@@ -594,7 +576,7 @@   where compileArg (se, _) = do           t <- subExpType se           case (se, t) of-            (_, Prim pt)    -> return $ Just $ Imp.ExpArg $ compileSubExpOfType pt se+            (_, Prim pt)   -> return $ Just $ Imp.ExpArg $ toExp' pt se             (Var v, Mem{}) -> return $ Just $ Imp.MemArg v             _              -> return Nothing @@ -602,25 +584,23 @@  defCompileExp pat (DoLoop ctx val form body) = do   dFParams mergepat-  forM_ merge $ \(p, se) -> do-    na <- subExpNotArray se-    when na $-      copyDWIM (paramName p) [] se []+  forM_ merge $ \(p, se) ->+    when ((==0) $ arrayRank $ paramType p) $+    copyDWIM (paramName p) [] se [] -  let doBody = compileLoopBody mergenames body+  let doBody = compileLoopBody mergepat body    case form of     ForLoop i it bound loopvars -> do-      bound' <- compileSubExp bound-       let setLoopParam (p,a)             | Prim _ <- paramType p =-                copyDWIM (paramName p) [] (Var a) [varIndex i]+                copyDWIM (paramName p) [] (Var a) [Imp.vi32 i]             | otherwise =                 return ()        dLParams $ map fst loopvars-      sFor i it bound' $ mapM_ setLoopParam loopvars >> doBody+      sFor i it (toExp' (IntType it) bound) $+        mapM_ setLoopParam loopvars >> doBody     WhileLoop cond ->       sWhile (Imp.var cond Bool) doBody @@ -630,7 +610,6 @@    where merge = ctx ++ val         mergepat = map fst merge-        mergenames = map paramName mergepat  defCompileExp pat (Op op) = do   opc <- asks envOpCompiler@@ -646,42 +625,42 @@   copyDWIM (patElemName pe) [] se []  defCompileBasicOp (Pattern _ [pe]) (UnOp op e) = do-  e' <- compileSubExp e+  e' <- toExp e   patElemName pe <-- Imp.UnOpExp op e'  defCompileBasicOp (Pattern _ [pe]) (ConvOp conv e) = do-  e' <- compileSubExp e+  e' <- toExp e   patElemName pe <-- Imp.ConvOpExp conv e'  defCompileBasicOp (Pattern _ [pe]) (BinOp bop x y) = do-  x' <- compileSubExp x-  y' <- compileSubExp y+  x' <- toExp x+  y' <- toExp y   patElemName pe <-- Imp.BinOpExp bop x' y'  defCompileBasicOp (Pattern _ [pe]) (CmpOp bop x y) = do-  x' <- compileSubExp x-  y' <- compileSubExp y+  x' <- toExp x+  y' <- toExp y   patElemName pe <-- Imp.CmpOpExp bop x' y'  defCompileBasicOp _ (Assert e msg loc) = do-  e' <- compileSubExp e-  msg' <- traverse compileSubExp msg+  e' <- toExp e+  msg' <- traverse toExp msg   emit $ Imp.Assert e' msg' loc  defCompileBasicOp (Pattern _ [pe]) (Index src slice)   | Just idxs <- sliceIndices slice =-      copyDWIM (patElemName pe) [] (Var src) $ map (compileSubExpOfType int32) idxs+      copyDWIM (patElemName pe) [] (Var src) $ map (toExp' int32) idxs  defCompileBasicOp _ Index{} =   return ()  defCompileBasicOp (Pattern _ [pe]) (Update _ slice se) =-  sUpdate (patElemName pe) (map (fmap (compileSubExpOfType int32)) slice) se+  sUpdate (patElemName pe) (map (fmap (toExp' int32)) slice) se  defCompileBasicOp (Pattern _ [pe]) (Replicate (Shape ds) se) = do-  ds' <- mapM compileSubExp ds+  ds' <- mapM toExp ds   is <- replicateM (length ds) (newVName "i")-  copy_elem <- collect $ copyDWIM (patElemName pe) (map varIndex is) se []+  copy_elem <- collect $ copyDWIM (patElemName pe) (map Imp.vi32 is) se []   emit $ foldl (.) id (zipWith (`Imp.For` Int32) is ds') copy_elem  defCompileBasicOp _ Scratch{} =@@ -690,14 +669,14 @@ defCompileBasicOp (Pattern [] [pe]) (Iota n e s et) = do   i <- newVName "i"   x <- newVName "x"-  n' <- compileSubExp n-  e' <- compileSubExp e-  s' <- compileSubExp s+  n' <- toExp n+  e' <- toExp e+  s' <- toExp s   let i' = ConvOpExp (SExt Int32 et) $ Imp.var i $ IntType Int32   dPrim_ x $ IntType et   sFor i Int32 n' $ do     x <-- e' + i' * s'-    copyDWIM (patElemName pe) [varIndex i] (Var x) []+    copyDWIM (patElemName pe) [Imp.vi32 i] (Var x) []  defCompileBasicOp (Pattern _ [pe]) (Copy src) =   copyDWIM (patElemName pe) [] (Var src) []@@ -708,14 +687,13 @@ defCompileBasicOp (Pattern _ [pe]) (Concat i x ys _) = do     MemLocation destmem destshape destixfun <-       entryArrayLocation <$> lookupArray (patElemName pe)-    xtype <- lookupType x     offs_glb <- dPrim "tmp_offs" int32     emit $ Imp.SetScalar offs_glb 0     let perm = [i] ++ [0..i-1] ++ [i+1..length destshape-1]         invperm = rearrangeInverse perm         destloc = MemLocation destmem destshape                   (IxFun.permute (IxFun.offsetIndex (IxFun.permute destixfun perm) $-                                  varIndex offs_glb)+                                  Imp.vi32 offs_glb)                    invperm)      forM_ (x:ys) $ \y -> do@@ -724,7 +702,7 @@           rows = case drop i $ entryArrayShape yentry of                   []  -> error $ "defCompileBasicOp Concat: empty array shape for " ++ pretty y                   r:_ -> innerExp $ Imp.dimSizeToExp r-      copy (elemType xtype) destloc srcloc $ arrayOuterSize yentry+      copy (elemType $ patElemType pe) destloc srcloc $ arrayOuterSize yentry       emit $ Imp.SetScalar offs_glb $ Imp.var offs_glb int32 + rows  defCompileBasicOp (Pattern [] [pe]) (ArrayLit es _)@@ -736,13 +714,12 @@       emit $ Imp.DeclareArray static_array dest_space t $ Imp.ArrayValues vs       let static_src = MemLocation static_array [Imp.ConstSize $ fromIntegral $ length es] $                        IxFun.iota [fromIntegral $ length es]-          num_bytes = Imp.ConstSize $ fromIntegral (length es) * primByteSize t-          entry = MemVar Nothing $ MemEntry num_bytes dest_space+          entry = MemVar Nothing $ MemEntry dest_space       addVar static_array entry       copy t dest_mem static_src $ fromIntegral $ length es   | otherwise =     forM_ (zip [0..] es) $ \(i,e) ->-      copyDWIM (patElemName pe) [constIndex i] e []+      copyDWIM (patElemName pe) [fromInteger i] e []    where isLiteral (Constant v) = Just v         isLiteral _ = Nothing@@ -806,6 +783,10 @@                   dPrim_ name' t                   return name' +dPrimV_ :: VName -> Imp.Exp -> ImpM lore op ()+dPrimV_ name e = do dPrim_ name $ primExpType e+                    name <-- e+ dPrimV :: String -> Imp.Exp -> ImpM lore op VName dPrimV name e = do name' <- dPrim name $ primExpType e                    name' <-- e@@ -815,14 +796,11 @@                    -> ImpM lore op (VarEntry lore) memBoundToVarEntry e (MemPrim bt) =   return $ ScalarVar e ScalarEntry { entryScalarType = bt }-memBoundToVarEntry e (MemMem size space) = do-  size' <- subExpToDimSize size-  return $ MemVar e MemEntry { entryMemSize = size'-                             , entryMemSpace = space-                             }+memBoundToVarEntry e (MemMem space) =+  return $ MemVar e $ MemEntry space memBoundToVarEntry e (MemArray bt shape _ (ArrayIn mem ixfun)) = do   shape' <- mapM subExpToDimSize $ shapeDims shape-  let location = MemLocation mem shape' $ fmap compilePrimExp ixfun+  let location = MemLocation mem shape' $ fmap (toExp' int32) ixfun   return $ ArrayVar e ArrayEntry { entryArrayLocation = location                                  , entryArrayElemType = bt                                  }@@ -879,30 +857,29 @@ subExpToDimSize Constant{} =   compilerBugS "Size subexp is not an int32 or int64 constant." -compileSubExpTo :: VName -> SubExp -> ImpM lore op ()-compileSubExpTo d se = copyDWIM d [] se []--compileSubExp :: SubExp -> ImpM lore op Imp.Exp-compileSubExp (Constant v) =-  return $ Imp.ValueExp v-compileSubExp (Var v) = do-  t <- lookupType v-  case t of-    Prim pt -> return $ Imp.var v pt-    _       -> compilerBugS $ "compileSubExp: SubExp is not a primitive type: " ++ pretty v--compileSubExpOfType :: PrimType -> SubExp -> Imp.Exp-compileSubExpOfType _ (Constant v) = Imp.ValueExp v-compileSubExpOfType t (Var v) = Imp.var v t+-- | Compile things to 'Imp.Exp'.+class ToExp a where+  -- | Compile to an 'Imp.Exp', where the type (must must still be a+  -- primitive) is deduced monadically.+  toExp :: a -> ImpM lore op Imp.Exp+  -- | Compile where we know the type in advance.+  toExp' :: PrimType -> a -> Imp.Exp -compilePrimExp :: PrimExp VName -> Imp.Exp-compilePrimExp = fmap Imp.ScalarVar+instance ToExp SubExp where+  toExp (Constant v) =+    return $ Imp.ValueExp v+  toExp (Var v) =+    lookupVar v >>= \case+    ScalarVar _ (ScalarEntry pt) ->+      return $ Imp.var v pt+    _       -> compilerBugS $ "toExp SubExp: SubExp is not a primitive type: " ++ pretty v -varIndex :: VName -> Imp.Exp-varIndex name = LeafExp (Imp.ScalarVar name) int32+  toExp' _ (Constant v) = Imp.ValueExp v+  toExp' t (Var v) = Imp.var v t -constIndex :: Int -> Imp.Exp-constIndex = fromIntegral+instance ToExp (PrimExp VName) where+  toExp = pure . fmap Imp.ScalarVar+  toExp' _ = fmap Imp.ScalarVar  addVar :: VName -> VarEntry lore -> ImpM lore op () addVar name entry =@@ -966,17 +943,17 @@               return $ ScalarDestination name  fullyIndexArray :: VName -> [Imp.Exp]-                -> ImpM lore op (VName, Imp.Space, Count Bytes)+                -> ImpM lore op (VName, Imp.Space, Count Elements) fullyIndexArray name indices = do   arr <- lookupArray name-  fullyIndexArray' (entryArrayLocation arr) indices $ entryArrayElemType arr+  fullyIndexArray' (entryArrayLocation arr) indices -fullyIndexArray' :: MemLocation -> [Imp.Exp] -> PrimType-                 -> ImpM lore op (VName, Imp.Space, Count Bytes)-fullyIndexArray' (MemLocation mem _ ixfun) indices bt = do+fullyIndexArray' :: MemLocation -> [Imp.Exp]+                 -> ImpM lore op (VName, Imp.Space, Count Elements)+fullyIndexArray' (MemLocation mem _ ixfun) indices = do   space <- entryMemSpace <$> lookupMemory mem   return (mem, space,-          bytes $ IxFun.index ixfun indices $ primByteSize bt)+          elements $ IxFun.index ixfun indices)  sliceArray :: MemLocation            -> Slice Imp.Exp@@ -999,11 +976,6 @@ strideArray (MemLocation mem shape ixfun) stride =   MemLocation mem shape $ IxFun.strideIndex ixfun stride -subExpNotArray :: SubExp -> ImpM lore op Bool-subExpNotArray se = subExpType se >>= \case-  Array {} -> return False-  _        -> return True- arrayOuterSize :: ArrayEntry -> Count Elements arrayOuterSize = arrayDimSize 0 @@ -1045,17 +1017,16 @@ copyElementWise :: CopyCompiler lore op copyElementWise bt (MemLocation destmem _ destIxFun) (MemLocation srcmem srcshape srcIxFun) n = do     is <- replicateM (IxFun.rank destIxFun) (newVName "i")-    let ivars = map varIndex is-        destidx = IxFun.index destIxFun ivars bt_size-        srcidx = IxFun.index srcIxFun ivars bt_size+    let ivars = map Imp.vi32 is+        destidx = IxFun.index destIxFun ivars+        srcidx = IxFun.index srcIxFun ivars         bounds = map innerExp $ n : drop 1 (map Imp.dimSizeToExp srcshape)     srcspace <- entryMemSpace <$> lookupMemory srcmem     destspace <- entryMemSpace <$> lookupMemory destmem     vol <- asks envVolatility     emit $ foldl (.) id (zipWith (`Imp.For` Int32) is bounds) $-      Imp.Write destmem (bytes destidx) bt destspace vol $-      Imp.index srcmem (bytes srcidx) bt srcspace vol-  where bt_size = primByteSize bt+      Imp.Write destmem (elements destidx) bt destspace vol $+      Imp.index srcmem (elements srcidx) bt srcspace vol  -- | Copy from here to there; both destination and source may be -- indexeded.@@ -1069,9 +1040,9 @@    | length srcis == length srcshape, length destis == length destshape = do   (targetmem, destspace, targetoffset) <--    fullyIndexArray' destlocation destis bt+    fullyIndexArray' destlocation destis   (srcmem, srcspace, srcoffset) <--    fullyIndexArray' srclocation srcis bt+    fullyIndexArray' srclocation srcis   vol <- asks envVolatility   return $ Imp.Write targetmem targetoffset bt destspace vol $     Imp.index srcmem srcoffset bt srcspace vol@@ -1083,7 +1054,15 @@           srclocation'  =             sliceArray srclocation $             fullSliceNum (IxFun.shape src_ixfun) $ map DimFix srcis-      if destlocation' == srclocation'+          destrank = length (memLocationShape destlocation')+          srcrank = length (memLocationShape srclocation')+      if destrank /= srcrank+        then fail $ "copyArrayDWIM: cannot copy to " +++             pretty (memLocationName destlocation') +++             " from " ++ pretty (memLocationName srclocation') +++             " because ranks do not match (" ++ pretty destrank +++             " vs " ++ pretty srcrank ++ ")"+      else if destlocation' == srclocation'         then return mempty -- Copy would be no-op.         else collect $ copy bt destlocation' srclocation' $              product $ map Imp.dimSizeToExp $@@ -1106,7 +1085,7 @@     unwords ["copyDWIMDest: constant source", pretty v, "cannot be written to memory destination."]   ArrayDestination (Just dest_loc) -> do     (dest_mem, dest_space, dest_i) <--      fullyIndexArray' dest_loc dest_is bt+      fullyIndexArray' dest_loc dest_is     vol <- asks envVolatility     emit $ Imp.Write dest_mem dest_i bt dest_space vol $ Imp.ValueExp v   ArrayDestination Nothing ->@@ -1116,7 +1095,7 @@ copyDWIMDest dest dest_is (Var src) src_is = do   src_entry <- lookupVar src   case (dest, src_entry) of-    (MemoryDestination mem, MemVar _ (MemEntry _ space)) ->+    (MemoryDestination mem, MemVar _ (MemEntry space)) ->       emit $ Imp.SetMem mem src space      (MemoryDestination{}, _) ->@@ -1142,7 +1121,7 @@     (ScalarDestination name, ArrayVar _ arr) -> do       let bt = entryArrayElemType arr       (mem, space, i) <--        fullyIndexArray' (entryArrayLocation arr) src_is bt+        fullyIndexArray' (entryArrayLocation arr) src_is       vol <- asks envVolatility       emit $ Imp.SetScalar name $ Imp.index mem i bt space vol @@ -1152,8 +1131,7 @@       emit =<< copyArrayDWIM bt dest_loc dest_is src_loc src_is      (ArrayDestination (Just dest_loc), ScalarVar _ (ScalarEntry bt)) -> do-      (dest_mem, dest_space, dest_i) <--        fullyIndexArray' dest_loc dest_is bt+      (dest_mem, dest_space, dest_i) <- fullyIndexArray' dest_loc dest_is       vol <- asks envVolatility       emit $ Imp.Write dest_mem dest_i bt dest_space vol (Imp.var src bt) @@ -1188,9 +1166,11 @@                 Pattern lore -> SubExp -> Space              -> ImpM lore op () compileAlloc (Pattern [] [mem]) e space = do-  e' <- compileSubExp e-  fake <- asks $ elem space . envFakeMemory-  unless fake $ emit $ Imp.Allocate (patElemName mem) (Imp.bytes e') space+  e' <- Imp.bytes <$> toExp e+  allocator <- asks $ M.lookup space . envAllocCompilers+  case allocator of+    Nothing -> emit $ Imp.Allocate (patElemName mem) e' space+    Just allocator' -> allocator' (patElemName mem) e' compileAlloc pat _ _ =   compilerBugS $ "compileAlloc: Invalid pattern: " ++ pretty pat @@ -1199,8 +1179,14 @@ dimSizeToSubExp (Imp.VarSize v) = Var v  dimSizeToExp :: Imp.Size -> Imp.Exp-dimSizeToExp = compilePrimExp . primExpFromSubExp int32 . dimSizeToSubExp+dimSizeToExp = toExp' int32 . primExpFromSubExp int32 . dimSizeToSubExp +-- | The number of bytes needed to represent the array in a+-- straightforward contiguous format.+typeSize :: Type -> Count Bytes+typeSize t = Imp.bytes $ Imp.LeafExp (Imp.SizeOf $ elemType t) int32 *+             product (map (toExp' int32) (arrayDims t))+ --- Building blocks for constructing code.  sFor :: VName -> IntType -> Imp.Exp -> ImpM lore op () -> ImpM lore op ()@@ -1234,33 +1220,24 @@ sOp :: op -> ImpM lore op () sOp = emit . Imp.Op -dSize :: Imp.Count u -> ImpM lore op Imp.Size-dSize size =-  case Imp.innerExp size of-    Imp.LeafExp (Imp.ScalarVar size') _ -> return $ Imp.VarSize size'-    Imp.ValueExp (IntValue (Int64Value v)) -> return $ Imp.ConstSize v-    _ -> do size_var <- dPrim "local_buf_size" int32-            size_var <-- Imp.innerExp size-            return $ Imp.VarSize size_var--sDeclareMem :: String -> Count Bytes -> Space -> ImpM lore op (VName, Imp.MemSize)-sDeclareMem name size space = do-+sDeclareMem :: String -> Space -> ImpM lore op VName+sDeclareMem name space = do   name' <- newVName name-  size' <- dSize size   emit $ Imp.DeclareMem name' space-  addVar name' $ MemVar Nothing $ MemEntry size' space-  return (name', size')+  addVar name' $ MemVar Nothing $ MemEntry space+  return name' -sAlloc_ :: VName -> Imp.MemSize -> Space -> ImpM lore op ()+sAlloc_ :: VName -> Count Bytes -> Space -> ImpM lore op () sAlloc_ name' size' space = do-  fake <- asks $ elem space . envFakeMemory-  unless fake $ emit $ Imp.Allocate name' (Imp.memSizeToExp size') space+  allocator <- asks $ M.lookup space . envAllocCompilers+  case allocator of+    Nothing -> emit $ Imp.Allocate name' size' space+    Just allocator' -> allocator' name' size'  sAlloc :: String -> Count Bytes -> Space -> ImpM lore op VName sAlloc name size space = do-  (name', size') <- sDeclareMem name size space-  sAlloc_ name' size' space+  name' <- sDeclareMem name space+  sAlloc_ name' size space   return name'  sArray :: String -> PrimType -> ShapeBase SubExp -> MemBind -> ImpM lore op VName@@ -1269,25 +1246,29 @@   dArray name' bt shape membind   return name' +-- | Like 'sAllocArray', but permute the in-memory representation of the indices as specified.+sAllocArrayPerm :: String -> PrimType -> ShapeBase SubExp -> Space -> [Int] -> ImpM lore op VName+sAllocArrayPerm name pt shape space perm = do+  let permuted_dims = rearrangeShape perm $ shapeDims shape+  mem <- sAlloc (name ++ "_mem") (typeSize (Array pt shape NoUniqueness)) space+  let iota_ixfun = IxFun.iota $ map (primExpFromSubExp int32) permuted_dims+  sArray name pt shape $+    ArrayIn mem $ IxFun.permute iota_ixfun $ rearrangeInverse perm+ -- | Uses linear/iota index function. sAllocArray :: String -> PrimType -> ShapeBase SubExp -> Space -> ImpM lore op VName-sAllocArray name pt shape space = do-  let arr_bytes = Imp.bytes $ Imp.LeafExp (Imp.SizeOf pt) int32 *-                  product (map (compileSubExpOfType int32) (shapeDims shape))-  mem <- sAlloc (name ++ "_mem") arr_bytes space-  sArray name pt shape $-    ArrayIn mem $ IxFun.iota $ map (primExpFromSubExp int32) $ shapeDims shape+sAllocArray name pt shape space =+  sAllocArrayPerm name pt shape space [0..shapeRank shape-1]  -- | Uses linear/iota index function. sStaticArray :: String -> Space -> PrimType -> Imp.ArrayContents -> ImpM lore op VName sStaticArray name space pt vs = do-  let num_elems = case vs of Imp.ArrayValues vs' -> genericLength vs'+  let num_elems = case vs of Imp.ArrayValues vs' -> length vs'                              Imp.ArrayZeros n -> fromIntegral n-      shape = Shape [constant num_elems]-      mem_size = Imp.ConstSize $ num_elems * primByteSize pt+      shape = Shape [intConst Int32 $ toInteger num_elems]   mem <- newVName $ name ++ "_mem"   emit $ Imp.DeclareArray mem space pt vs-  addVar mem $ MemVar Nothing $ MemEntry mem_size space+  addVar mem $ MemVar Nothing $ MemEntry space   sArray name pt shape $ ArrayIn mem $ IxFun.iota [fromIntegral num_elems]  sWrite :: VName -> [Imp.Exp] -> PrimExp Imp.ExpLeaf -> ImpM lore op ()@@ -1301,6 +1282,16 @@   MemLocation mem shape ixfun <- entryArrayLocation <$> lookupArray arr   let memdest = sliceArray (MemLocation mem shape ixfun) slice   copyDWIMDest (ArrayDestination $ Just memdest) [] v []++sLoopNest :: Shape+          -> ([Imp.Exp] -> ImpM lore op ())+          -> ImpM lore op ()+sLoopNest = sLoopNest' [] . shapeDims+  where sLoopNest' is [] f = f $ reverse is+        sLoopNest' is (d:ds) f = do+          i <- newVName "nest_i"+          d' <- toExp d+          sFor i Int32 d' $ sLoopNest' (Imp.var i int32:is) ds f  -- | ASsignment. (<--) :: VName -> Imp.Exp -> ImpM lore op ()

@@ -3,14 +3,13 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE ConstraintKinds #-} module Futhark.CodeGen.ImpGen.Kernels-  ( compileProg+  ( Futhark.CodeGen.ImpGen.Kernels.compileProg   )   where  import Control.Monad.Except import Control.Monad.Reader import Data.Maybe-import qualified Data.Map.Strict as M import Data.List  import Prelude hiding (quot)@@ -20,49 +19,50 @@ import Futhark.Representation.ExplicitMemory import qualified Futhark.CodeGen.ImpCode.Kernels as Imp import Futhark.CodeGen.ImpCode.Kernels (bytes)-import qualified Futhark.CodeGen.ImpGen as ImpGen+import Futhark.CodeGen.ImpGen import Futhark.CodeGen.ImpGen.Kernels.Base+import Futhark.CodeGen.ImpGen.Kernels.SegMap import Futhark.CodeGen.ImpGen.Kernels.SegRed+import Futhark.CodeGen.ImpGen.Kernels.SegScan import Futhark.CodeGen.ImpGen.Kernels.SegGenRed-import Futhark.CodeGen.ImpGen (sFor, sWhen,-                               sOp) import Futhark.CodeGen.ImpGen.Kernels.Transpose import qualified Futhark.Representation.ExplicitMemory.IndexFunction as IxFun import Futhark.CodeGen.SetDefaultSpace-import Futhark.Util.IntegralExp (quotRoundingUp, quot, IntegralExp)+import Futhark.Util.IntegralExp (quot, IntegralExp) -callKernelOperations :: ImpGen.Operations ExplicitMemory Imp.HostOp+callKernelOperations :: Operations ExplicitMemory Imp.HostOp callKernelOperations =-  ImpGen.Operations { ImpGen.opsExpCompiler = expCompiler-                    , ImpGen.opsCopyCompiler = callKernelCopy-                    , ImpGen.opsOpCompiler = opCompiler-                    , ImpGen.opsStmsCompiler = ImpGen.defCompileStms+  Operations { opsExpCompiler = expCompiler+                    , opsCopyCompiler = callKernelCopy+                    , opsOpCompiler = opCompiler+                    , opsStmsCompiler = defCompileStms+                    , opsAllocCompilers = mempty                     }  compileProg :: MonadFreshNames m => Prog ExplicitMemory -> m (Either InternalError Imp.Program) compileProg prog =   fmap (setDefaultSpace (Imp.Space "device")) <$>-  ImpGen.compileProg callKernelOperations (Imp.Space "device") [Imp.Space "local"] prog+  Futhark.CodeGen.ImpGen.compileProg callKernelOperations (Imp.Space "device") prog  opCompiler :: Pattern ExplicitMemory -> Op ExplicitMemory            -> CallKernelGen () opCompiler dest (Alloc e space) =-  ImpGen.compileAlloc dest e space+  compileAlloc dest e space opCompiler (Pattern _ [pe]) (Inner (GetSize key size_class)) = do-  fname <- asks ImpGen.envFunction+  fname <- asks envFunction   sOp $ Imp.GetSize (patElemName pe) (keyWithEntryPoint fname key) $     sizeClassWithEntryPoint fname size_class opCompiler (Pattern _ [pe]) (Inner (CmpSizeLe key size_class x)) = do-  fname <- asks ImpGen.envFunction+  fname <- asks envFunction   let size_class' = sizeClassWithEntryPoint fname size_class   sOp . Imp.CmpSizeLe (patElemName pe) (keyWithEntryPoint fname key) size_class'-    =<< ImpGen.compileSubExp x+    =<< toExp x opCompiler (Pattern _ [pe]) (Inner (GetSizeMax size_class)) =   sOp $ Imp.GetSizeMax (patElemName pe) size_class opCompiler dest (Inner (HostOp kernel)) =   kernelCompiler dest kernel opCompiler pat e =-  compilerBugS $ "ImpGen.opCompiler: Invalid pattern\n  " +++  compilerBugS $ "opCompiler: Invalid pattern\n  " ++   pretty pat ++ "\nfor expression\n  " ++ pretty e  sizeClassWithEntryPoint :: Name -> Imp.SizeClass -> Imp.SizeClass@@ -75,18 +75,7 @@                -> CallKernelGen ()  kernelCompiler pat (Kernel desc space _ kernel_body) = do-  (constants, init_constants) <- kernelInitialisation space--  kernel_body' <--    makeAllMemoryGlobal $ ImpGen.subImpM_ (inKernelOperations constants) $ do-    init_constants-    compileKernelBody pat constants kernel_body--  let bound_in_kernel =-        M.keys $-        scopeOfKernelSpace space <>-        scopeOf (kernelBodyStms kernel_body)-  (uses, local_memory) <- computeKernelUses kernel_body' bound_in_kernel+  (constants, init_constants) <- kernelInitialisationSetSpace space $ return ()    forM_ (kernelHints desc) $ \(s,v) -> do     ty <- case v of@@ -96,31 +85,40 @@                                          , " in kernel '", kernelName desc, "'"                                          , " did not have primType value." ] -    ImpGen.compileSubExp v >>= ImpGen.emit . Imp.DebugPrint s (elemType ty)+    emit $ Imp.DebugPrint s $ Just (elemType ty, toExp' (elemType ty) v) -  sOp $ Imp.CallKernel Imp.Kernel-            { Imp.kernelBody = kernel_body'-            , Imp.kernelLocalMemory = local_memory-            , Imp.kernelUses = uses-            , Imp.kernelNumGroups = [ImpGen.compileSubExpOfType int32 $ spaceNumGroups space]-            , Imp.kernelGroupSize = [ImpGen.compileSubExpOfType int32 $ spaceGroupSize space]-            , Imp.kernelName = nameFromString $ kernelName desc ++ "_" ++-                               show (baseTag $ kernelGlobalThreadIdVar constants)-            }+  let virt_groups = toExp' int32 (spaceNumVirtGroups space)+  sKernel constants (kernelName desc) $ do+    init_constants+    virtualiseGroups constants virt_groups $ \group_id -> do+      let flat_id =+            if kernelGroupIdVar constants /= group_id+            then Imp.vi32 group_id * kernelGroupSize constants + kernelLocalThreadId constants+            else kernelGlobalThreadId constants+      setSpaceIndices flat_id space+      compileKernelStms constants (kernelBodyStms kernel_body) $+        zipWithM_ (compileKernelResult constants) (patternElements pat) $+        kernelBodyResult kernel_body -kernelCompiler pat (SegRed space comm red_op nes _ body) =-  compileSegRed pat space comm red_op nes body+kernelCompiler pat (SegMap space _ body) =+  compileSegMap pat space body +kernelCompiler pat (SegRed space reds _ body) =+  compileSegRed pat space reds body++kernelCompiler pat (SegScan space red_op nes _ kbody) =+  compileSegScan pat space red_op nes kbody+ kernelCompiler pat (SegGenRed space ops _ body) =   compileSegGenRed pat space ops body -expCompiler :: ImpGen.ExpCompiler ExplicitMemory Imp.HostOp+expCompiler :: ExpCompiler ExplicitMemory Imp.HostOp  -- We generate a simple kernel for itoa and replicate. expCompiler (Pattern _ [pe]) (BasicOp (Iota n x s et)) = do-  n' <- ImpGen.compileSubExp n-  x' <- ImpGen.compileSubExp x-  s' <- ImpGen.compileSubExp s+  n' <- toExp n+  x' <- toExp x+  s' <- toExp s    sIota (patElemName pe) n' x' s' et @@ -132,19 +130,19 @@   return ()  expCompiler dest e =-  ImpGen.defCompileExp dest e+  defCompileExp dest e -callKernelCopy :: ImpGen.CopyCompiler ExplicitMemory Imp.HostOp+callKernelCopy :: CopyCompiler ExplicitMemory Imp.HostOp callKernelCopy bt-  destloc@(ImpGen.MemLocation destmem destshape destIxFun)-  srcloc@(ImpGen.MemLocation srcmem srcshape srcIxFun)+  destloc@(MemLocation destmem destshape destIxFun)+  srcloc@(MemLocation srcmem srcshape srcIxFun)   n   | Just (destoffset, srcoffset,           num_arrays, size_x, size_y,           src_elems, dest_elems) <- isMapTransposeKernel bt destloc srcloc = do        fname <- mapTransposeForType bt-      ImpGen.emit $ Imp.Call [] fname+      emit $ Imp.Call [] fname         [Imp.MemArg destmem, Imp.ExpArg destoffset,          Imp.MemArg srcmem, Imp.ExpArg srcoffset,          Imp.ExpArg num_arrays, Imp.ExpArg size_x, Imp.ExpArg size_y,@@ -159,25 +157,25 @@       IxFun.linearWithOffset destIxFun bt_size,     Just srcoffset  <-       IxFun.linearWithOffset srcIxFun bt_size = do-        let row_size = product $ map ImpGen.dimSizeToExp $ drop 1 srcshape-        srcspace <- ImpGen.entryMemSpace <$> ImpGen.lookupMemory srcmem-        destspace <- ImpGen.entryMemSpace <$> ImpGen.lookupMemory destmem-        ImpGen.emit $ Imp.Copy+        let row_size = product $ map dimSizeToExp $ drop 1 srcshape+        srcspace <- entryMemSpace <$> lookupMemory srcmem+        destspace <- entryMemSpace <$> lookupMemory destmem+        emit $ Imp.Copy           destmem (bytes destoffset) destspace           srcmem (bytes srcoffset) srcspace $           (n * row_size) `Imp.withElemType` bt    | otherwise = sCopy bt destloc srcloc n -mapTransposeForType :: PrimType -> ImpGen.ImpM ExplicitMemory Imp.HostOp Name+mapTransposeForType :: PrimType -> ImpM ExplicitMemory Imp.HostOp Name mapTransposeForType bt = do   -- XXX: The leading underscore is to avoid clashes with a   -- programmer-defined function of the same name (this is a bad   -- solution...).   let fname = nameFromString $ "_" <> mapTransposeName bt -  exists <- ImpGen.hasFunction fname-  unless exists $ ImpGen.emitFunction fname $ mapTransposeFunction bt+  exists <- hasFunction fname+  unless exists $ emitFunction fname $ mapTransposeFunction bt    return fname @@ -288,13 +286,13 @@             v32 mulx, v32 muly, v32 num_arrays,             block) bt -isMapTransposeKernel :: PrimType -> ImpGen.MemLocation -> ImpGen.MemLocation+isMapTransposeKernel :: PrimType -> MemLocation -> MemLocation                      -> Maybe (Imp.Exp, Imp.Exp,                                Imp.Exp, Imp.Exp, Imp.Exp,                                Imp.Exp, Imp.Exp) isMapTransposeKernel bt-  (ImpGen.MemLocation _ _ destIxFun)-  (ImpGen.MemLocation _ _ srcIxFun)+  (MemLocation _ _ destIxFun)+  (MemLocation _ _ srcIxFun)   | Just (dest_offset, perm_and_destshape) <- IxFun.rearrangeWithOffset destIxFun bt_size,     (perm, destshape) <- unzip perm_and_destshape,     srcshape' <- IxFun.shape srcIxFun,@@ -322,100 +320,3 @@           let (mapped, notmapped) = splitAt r1 shape               (pretrans, posttrans) = f $ splitAt r2 notmapped           in (product mapped, product pretrans, product posttrans)--compileKernelBody :: Pattern InKernel-                  -> KernelConstants-                  -> KernelBody InKernel-                  -> InKernelGen ()-compileKernelBody pat constants kbody =-  compileKernelStms constants (stmsToList $ kernelBodyStms kbody) $-  zipWithM_ (compileKernelResult constants) (patternElements pat) $-  kernelBodyResult kbody--compileKernelResult :: KernelConstants -> PatElem InKernel -> KernelResult-                    -> InKernelGen ()--compileKernelResult constants pe (ThreadsReturn OneResultPerGroup what) = do-  i <- newVName "i"--  in_local_memory <- arrayInLocalMemory what-  let me = kernelLocalThreadId constants--  if not in_local_memory then do-    who' <- ImpGen.compileSubExp $ intConst Int32 0-    sWhen (me .==. who') $-      ImpGen.copyDWIM (patElemName pe) [kernelGroupId constants] what []-    else do-      -- 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).-      ---      -- We do the reads/writes multidimensionally, but the loop is-      -- single-dimensional.-      ws <- mapM ImpGen.compileSubExp . arrayDims =<< subExpType what-      -- Compute how many elements this thread is responsible for.-      -- Formula: (w - ltid) / group_size (rounded up).-      let w = product ws-          ltid = kernelLocalThreadId constants-          group_size = kernelGroupSize constants-          to_write = (w - ltid) `quotRoundingUp` group_size-          is = unflattenIndex ws $ ImpGen.varIndex i * group_size + ltid--      sFor i Int32 to_write $-        ImpGen.copyDWIM (patElemName pe) (kernelGroupId constants : is) what is--compileKernelResult constants pe (ThreadsReturn AllThreads what) =-  ImpGen.copyDWIM (patElemName pe) [kernelGlobalThreadId constants] what []--compileKernelResult constants pe (ThreadsReturn (ThreadsPerGroup limit) what) =-  sWhen (isActive limit) $-  ImpGen.copyDWIM (patElemName pe) [kernelGroupId constants] what []--compileKernelResult constants pe (ThreadsReturn ThreadsInSpace what) = do-  let is = map (ImpGen.varIndex . fst) $ kernelDimensions constants-  sWhen (kernelThreadActive constants) $ ImpGen.copyDWIM (patElemName pe) is what []--compileKernelResult constants pe (ConcatReturns SplitContiguous _ per_thread_elems moffset what) = do-  dest_loc <- ImpGen.entryArrayLocation <$> ImpGen.lookupArray (patElemName pe)-  let dest_loc_offset = ImpGen.offsetArray dest_loc offset-      dest' = ImpGen.arrayDestination dest_loc_offset-  ImpGen.copyDWIMDest dest' [] (Var what) []-  where offset = case moffset of-                   Nothing -> ImpGen.compileSubExpOfType int32 per_thread_elems *-                              kernelGlobalThreadId constants-                   Just se -> ImpGen.compileSubExpOfType int32 se--compileKernelResult constants pe (ConcatReturns (SplitStrided stride) _ _ moffset what) = do-  dest_loc <- ImpGen.entryArrayLocation <$> ImpGen.lookupArray (patElemName pe)-  let dest_loc' = ImpGen.strideArray-                  (ImpGen.offsetArray dest_loc offset) $-                  ImpGen.compileSubExpOfType int32 stride-      dest' = ImpGen.arrayDestination dest_loc'-  ImpGen.copyDWIMDest dest' [] (Var what) []-  where offset = case moffset of-                   Nothing -> kernelGlobalThreadId constants-                   Just se -> ImpGen.compileSubExpOfType int32 se--compileKernelResult constants pe (WriteReturn rws _arr dests) = do-  rws' <- mapM ImpGen.compileSubExp rws-  forM_ dests $ \(is, e) -> do-    is' <- mapM ImpGen.compileSubExp is-    let condInBounds i rw = 0 .<=. i .&&. i .<. rw-        write = foldl (.&&.) (kernelThreadActive constants) $-                zipWith condInBounds is' rws'-    sWhen write $ ImpGen.copyDWIM (patElemName pe) (map (ImpGen.compileSubExpOfType int32) is) e []--compileKernelResult _ _ KernelInPlaceReturn{} =-  -- Already in its place... said it was a hack.-  return ()--arrayInLocalMemory :: SubExp -> InKernelGen Bool-arrayInLocalMemory (Var name) = do-  res <- ImpGen.lookupVar name-  case res of-    ImpGen.ArrayVar _ entry ->-      (Space "local"==) . ImpGen.entryMemSpace <$>-      ImpGen.lookupMemory (ImpGen.memLocationName (ImpGen.entryArrayLocation entry))-    _ -> return False-arrayInLocalMemory Constant{} = return False

@@ -4,12 +4,13 @@ module Futhark.CodeGen.ImpGen.Kernels.Base   ( KernelConstants (..)   , inKernelOperations-  , computeKernelUses   , keyWithEntryPoint   , CallKernelGen   , InKernelGen   , computeThreadChunkSize+  , simpleKernelConstants   , kernelInitialisation+  , kernelInitialisationSimple   , kernelInitialisationSetSpace   , setSpaceIndices   , makeAllMemoryGlobal@@ -22,6 +23,8 @@   , sReplicate   , sIota   , sCopy+  , compileKernelResult+  , virtualiseGroups    , atomicUpdate   , atomicUpdateLocking@@ -45,21 +48,17 @@ import Futhark.Transform.Rename import Futhark.Representation.ExplicitMemory import qualified Futhark.CodeGen.ImpCode.Kernels as Imp-import Futhark.CodeGen.ImpCode.Kernels (bytes)-import qualified Futhark.CodeGen.ImpGen as ImpGen-import Futhark.CodeGen.ImpGen ((<--),-                               sFor, sWhile, sComment, sIf, sWhen, sUnless,-                               sOp,-                               dPrim, dPrim_, dPrimV)-import Futhark.Tools (partitionChunkedKernelLambdaParameters)+import Futhark.CodeGen.ImpCode.Kernels (elements)+import Futhark.CodeGen.ImpGen import Futhark.Util.IntegralExp (quotRoundingUp, quot, rem, IntegralExp) import Futhark.Util (splitAt3, maybeNth, takeLast) -type CallKernelGen = ImpGen.ImpM ExplicitMemory Imp.HostOp-type InKernelGen = ImpGen.ImpM InKernel Imp.KernelOp+type CallKernelGen = ImpM ExplicitMemory Imp.HostOp+type InKernelGen = ImpM InKernel Imp.KernelOp  data KernelConstants = KernelConstants-                       { kernelGlobalThreadId :: Imp.Exp+                       { kernelOuterVTable :: VTable ExplicitMemory -- XXX+                       , kernelGlobalThreadId :: Imp.Exp                        , kernelLocalThreadId :: Imp.Exp                        , kernelGroupId :: Imp.Exp                        , kernelGlobalThreadIdVar :: VName@@ -76,12 +75,22 @@                        -- for unrolling.                        } -inKernelOperations :: KernelConstants -> ImpGen.Operations InKernel Imp.KernelOp-inKernelOperations constants = (ImpGen.defaultOperations $ compileInKernelOp constants)-                               { ImpGen.opsCopyCompiler = inKernelCopy-                               , ImpGen.opsExpCompiler = inKernelExpCompiler-                               , ImpGen.opsStmsCompiler = \_ -> compileKernelStms constants-                               }+inKernelOperations :: KernelConstants -> Operations InKernel Imp.KernelOp+inKernelOperations constants =+  (defaultOperations $ compileInKernelOp constants)+  { opsCopyCompiler = inKernelCopy+  , opsExpCompiler = inKernelExpCompiler+  , opsStmsCompiler = \_ -> compileKernelStms constants+  , opsAllocCompilers =+      M.fromList [ (Space "local", allocLocal)+                 , (Space "private", allocPrivate) ]+  }+  where allocLocal :: AllocCompiler InKernel Imp.KernelOp+        allocLocal mem size = do+          size' <- localMemSize (kernelOuterVTable constants) size+          sOp $ Imp.LocalAlloc mem size'+        allocPrivate mem size =+          sOp $ Imp.PrivateAlloc mem size  keyWithEntryPoint :: Name -> Name -> Name keyWithEntryPoint fname key =@@ -89,11 +98,18 @@  -- | We have no bulk copy operation (e.g. memmove) inside kernels, so -- turn any copy into a loop.-inKernelCopy :: ImpGen.CopyCompiler InKernel Imp.KernelOp-inKernelCopy = ImpGen.copyElementWise+inKernelCopy :: CopyCompiler InKernel Imp.KernelOp+inKernelCopy = copyElementWise  compileInKernelOp :: KernelConstants -> Pattern InKernel -> Op InKernel                   -> InKernelGen ()+compileInKernelOp _ (Pattern _ [mem]) (Alloc size (Space "private")) = do+  size' <- toExp size+  sOp $ Imp.PrivateAlloc (patElemName mem) $ Imp.bytes size'+compileInKernelOp constants (Pattern _ [mem]) (Alloc size (Space "local")) = do+  size' <- localMemSize (kernelOuterVTable constants) . Imp.bytes =<<+           toExp size+  sOp $ Imp.LocalAlloc (patElemName mem) size' compileInKernelOp _ (Pattern _ [mem]) Alloc{} =   compilerLimitationS $ "Cannot allocate memory block " ++ pretty mem ++ " in kernel." compileInKernelOp _ dest Alloc{} =@@ -101,7 +117,7 @@ compileInKernelOp constants pat (Inner op) =   compileKernelExp constants pat op -inKernelExpCompiler :: ImpGen.ExpCompiler InKernel Imp.KernelOp+inKernelExpCompiler :: ExpCompiler InKernel Imp.KernelOp inKernelExpCompiler _ (BasicOp (Assert _ _ (loc, locs))) =   compilerLimitationS $   unlines [ "Cannot compile assertion at " ++@@ -111,22 +127,22 @@ -- The static arrays stuff does not work inside kernels. inKernelExpCompiler (Pattern _ [dest]) (BasicOp (ArrayLit es _)) =   forM_ (zip [0..] es) $ \(i,e) ->-  ImpGen.copyDWIM (patElemName dest) [fromIntegral (i::Int32)] e []+  copyDWIM (patElemName dest) [fromIntegral (i::Int32)] e [] inKernelExpCompiler dest e =-  ImpGen.defCompileExp dest e+  defCompileExp dest e  compileKernelExp :: KernelConstants -> Pattern InKernel -> KernelExp InKernel                  -> InKernelGen ()  compileKernelExp _ pat (Barrier ses) = do   forM_ (zip (patternNames pat) ses) $ \(d, se) ->-    ImpGen.copyDWIM d [] se []+    copyDWIM d [] se []   sOp Imp.LocalBarrier  compileKernelExp _ (Pattern [] [size]) (SplitSpace o w i elems_per_thread) = do-  num_elements <- Imp.elements <$> ImpGen.compileSubExp w-  i' <- ImpGen.compileSubExp i-  elems_per_thread' <- Imp.elements <$> ImpGen.compileSubExp elems_per_thread+  num_elements <- Imp.elements <$> toExp w+  i' <- toExp i+  elems_per_thread' <- Imp.elements <$> toExp elems_per_thread   computeThreadChunkSize o i' elems_per_thread' num_elements (patElemName size)  compileKernelExp constants pat (Combine (CombineSpace scatter cspace) _ aspace body) = do@@ -166,56 +182,51 @@      -- Execute the body if we are within bounds.     sWhen (isActive cspace .&&. isActive aspace) $ allThreads constants $-      ImpGen.compileStms (freeIn $ bodyResult body) (stmsToList $ bodyStms body) $ do+      compileStms (freeIn $ bodyResult body) (bodyStms body) $ do        forM_ (zip4 scatter_ws_repl res_is res_vs scatter_pes) $         \(w, res_i, res_v, scatter_pe) -> do-          let res_i' = ImpGen.compileSubExpOfType int32 res_i-              w'     = ImpGen.compileSubExpOfType int32 w+          let res_i' = toExp' int32 res_i+              w'     = toExp' int32 w               -- We have to check that 'res_i' is in-bounds wrt. an array of size 'w'.               in_bounds = 0 .<=. res_i' .&&. res_i' .<. w'-          sWhen in_bounds $ ImpGen.copyDWIM (patElemName scatter_pe) [res_i'] res_v []+          sWhen in_bounds $ copyDWIM (patElemName scatter_pe) [res_i'] res_v []        forM_ (zip normal_pes res_normal) $ \(pe, res) ->-        ImpGen.copyDWIM (patElemName pe) local_index res []+        copyDWIM (patElemName pe) local_index res []    sOp Imp.LocalBarrier    where streamBounded (Var v)           | Just x <- lookup v $ kernelStreamed constants =               Imp.sizeToExp x-        streamBounded se = ImpGen.compileSubExpOfType int32 se+        streamBounded se = toExp' int32 se -        local_index = map (ImpGen.compileSubExpOfType int32 . Var . fst) cspace+        local_index = map (toExp' int32 . Var . fst) cspace  compileKernelExp constants (Pattern _ dests) (GroupReduce w lam input) = do-  let [my_index_param, offset_param] = take 2 $ lambdaParams lam-      lam' = lam { lambdaParams = drop 2 $ lambdaParams lam }--  dPrim_ (paramName my_index_param) int32-  dPrim_ (paramName offset_param) int32-  paramName my_index_param <-- kernelGlobalThreadId constants-  w' <- ImpGen.compileSubExp w-  groupReduceWithOffset constants (paramName offset_param) w' lam' $ map snd input+  w' <- toExp w+  groupReduce constants w' lam $ map snd input    sOp Imp.LocalBarrier    -- The final result will be stored in element 0 of the local memory array.   forM_ (zip dests input) $ \(dest, (_, arr)) ->-    ImpGen.copyDWIM (patElemName dest) [] (Var arr) [0]+    copyDWIM (patElemName dest) [] (Var arr) [0]  compileKernelExp constants _ (GroupScan w lam input) = do-  w' <- ImpGen.compileSubExp w+  w' <- toExp w   groupScan constants Nothing w' lam $ map snd input  compileKernelExp constants (Pattern _ final) (GroupStream w maxchunk lam accs _arrs) = do   let GroupStreamLambda block_size block_offset acc_params arr_params body = lam       block_offset' = Imp.var block_offset int32-  w' <- ImpGen.compileSubExp w-  max_block_size <- ImpGen.compileSubExp maxchunk+  w' <- toExp w+  max_block_size <- toExp maxchunk -  ImpGen.dLParams (acc_params++arr_params)-  zipWithM_ ImpGen.compileSubExpTo (map paramName acc_params) accs+  dLParams (acc_params++arr_params)+  forM_ (zip acc_params accs) $ \(p, acc) ->+    copyDWIM (paramName p) [] acc []   dPrim_ block_size int32    -- If the GroupStream is morally just a do-loop, generate simpler code.@@ -223,7 +234,7 @@     Just stms' | ValueExp x <- max_block_size, oneIsh x -> do       let body' = body { bodyStms = stmsFromList stms' }           body'' = allThreads constants $-                   ImpGen.compileLoopBody (map paramName acc_params) body'+                   compileLoopBody acc_params body'       block_size <-- 1        -- Check if loop is candidate for unrolling.@@ -244,7 +255,7 @@     _ -> do       dPrim_ block_offset int32       let body' = streaming constants block_size maxchunk $-                  ImpGen.compileBody' acc_params body+                  compileBody' acc_params body        block_offset <-- 0 @@ -270,23 +281,23 @@              set_block_size >> body' >> increase_offset    forM_ (zip final acc_params) $ \(pe, p) ->-    ImpGen.copyDWIM (patElemName pe) [] (Var $ paramName p) []+    copyDWIM (patElemName pe) [] (Var $ paramName p) []    where isSimpleThreadInSpace (Let _ _ Op{}) = Nothing         isSimpleThreadInSpace bnd = Just bnd  compileKernelExp _ _ (GroupGenReduce w arrs op bucket values locks) = do   -- Check if bucket is in-bounds-  bucket' <- mapM ImpGen.compileSubExp bucket-  w' <- mapM ImpGen.compileSubExp w-  num_locks <- ImpGen.compileSubExpOfType int32 . arraySize 0 <$> lookupType locks+  bucket' <- mapM toExp bucket+  w' <- mapM toExp w+  num_locks <- toExp' int32 . arraySize 0 <$> lookupType locks   let locking = Locking locks 0 1 0 $ (`rem` num_locks) . sum       values_params = takeLast (length values) $ lambdaParams op    sWhen (indexInBounds bucket' w') $ do     forM_ (zip values_params values) $ \(p, v) ->-      ImpGen.copyDWIM (paramName p) [] v []-    atomicUpdate arrs bucket' op locking+      copyDWIM (paramName p) [] v []+    atomicUpdate DefaultSpace arrs bucket' op locking   where indexInBounds inds bounds =           foldl1 (.&&.) $ zipWith checkBound inds bounds           where checkBound ind bound = 0 .<=. ind .&&. ind .<. bound@@ -297,10 +308,10 @@  streaming :: KernelConstants -> VName -> SubExp -> InKernelGen () -> InKernelGen () streaming constants chunksize bound m = do-  bound' <- ImpGen.subExpToDimSize bound+  bound' <- subExpToDimSize bound   let constants' =         constants { kernelStreamed = (chunksize, bound') : kernelStreamed constants }-  ImpGen.emit =<< ImpGen.subImpM_ (inKernelOperations constants') m+  emit =<< subImpM_ (inKernelOperations constants') m  -- | Locking strategy used for an atomic update. data Locking =@@ -321,15 +332,15 @@ -- | A function for generating code for an atomic update.  Assumes -- that the bucket is in-bounds. type AtomicUpdate lore =-  [VName] -> [Imp.Exp] -> ImpGen.ImpM lore Imp.KernelOp ()+  Space -> [VName] -> [Imp.Exp] -> ImpM lore Imp.KernelOp ()  atomicUpdate :: ExplicitMemorish lore =>-                [VName] -> [Imp.Exp] -> Lambda lore -> Locking-             -> ImpGen.ImpM lore Imp.KernelOp ()-atomicUpdate arrs bucket lam locking =+                Space -> [VName] -> [Imp.Exp] -> Lambda lore -> Locking+             -> ImpM lore Imp.KernelOp ()+atomicUpdate space arrs bucket lam locking =   case atomicUpdateLocking lam of-    Left f -> f arrs bucket-    Right f -> f locking arrs bucket+    Left f -> f space arrs bucket+    Right f -> f locking space arrs bucket  -- | 'atomicUpdate', but where it is explicitly visible whether a -- locking strategy is necessary.@@ -339,7 +350,7 @@  atomicUpdateLocking lam   | Just ops_and_ts <- splitOp lam,-    all (\(_, t, _, _) -> primBitSize t == 32) ops_and_ts = Left $ \arrs bucket ->+    all (\(_, t, _, _) -> primBitSize t == 32) ops_and_ts = Left $ \space arrs bucket ->   -- If the operator is a vectorised binary operator on 32-bit values,   -- we can use a particularly efficient implementation. If the   -- operator has an atomic implementation we use that, otherwise it@@ -350,16 +361,15 @@   -- Common variables.   old <- dPrim "old" t -  (arr', _a_space, bucket_offset) <- ImpGen.fullyIndexArray a bucket+  (arr', _a_space, bucket_offset) <- fullyIndexArray a bucket -  case opHasAtomicSupport old arr' bucket_offset op of+  case opHasAtomicSupport space old arr' bucket_offset op of     Just f -> sOp $ f $ Imp.var y t--    Nothing -> atomicUpdateCAS t a old bucket x $+    Nothing -> atomicUpdateCAS space t a old bucket x $       x <-- Imp.BinOpExp op (Imp.var x t) (Imp.var y t) -  where opHasAtomicSupport old arr' bucket' bop = do-          let atomic f = Imp.Atomic . f old arr' bucket'+  where opHasAtomicSupport space old arr' bucket' bop = do+          let atomic f = Imp.Atomic space . f old arr' bucket'           atomic <$> Imp.atomicBinOp bop  -- If the operator functions purely on single 32-bit values, we can@@ -368,28 +378,28 @@ atomicUpdateLocking op   | [Prim t] <- lambdaReturnType op,     [xp, _] <- lambdaParams op,-    primBitSize t == 32 = Left $ \[arr] bucket -> do+    primBitSize t == 32 = Left $ \space [arr] bucket -> do       old <- dPrim "old" t-      atomicUpdateCAS t arr old bucket (paramName xp) $-        ImpGen.compileBody' [xp] $ lambdaBody op+      atomicUpdateCAS space t arr old bucket (paramName xp) $+        compileBody' [xp] $ lambdaBody op -atomicUpdateLocking op = Right $ \locking arrs bucket -> do+atomicUpdateLocking op = Right $ \locking space arrs bucket -> do   old <- dPrim "old" int32   continue <- dPrimV "continue" true    -- Correctly index into locks.   (locks', _locks_space, locks_offset) <--    ImpGen.fullyIndexArray (lockingArray locking) [lockingMapping locking bucket]+    fullyIndexArray (lockingArray locking) [lockingMapping locking bucket]    -- Critical section   let try_acquire_lock =-        sOp $ Imp.Atomic $+        sOp $ Imp.Atomic space $         Imp.AtomicCmpXchg old locks' locks_offset (lockingIsUnlocked locking) (lockingToLock locking)       lock_acquired = Imp.var old int32 .==. lockingIsUnlocked locking       -- Even the releasing is done with an atomic rather than a       -- simple write, for memory coherency reasons.       release_lock =-        sOp $ Imp.Atomic $+        sOp $ Imp.Atomic space $         Imp.AtomicCmpXchg old locks' locks_offset (lockingToLock locking) (lockingToUnlock locking)       break_loop = continue <-- false @@ -405,39 +415,43 @@   -- synchronisation without atomics (naughty!).   let (acc_params, _arr_params) = splitAt (length arrs) $ lambdaParams op       bind_acc_params =-        ImpGen.everythingVolatile $-        ImpGen.sComment "bind lhs" $+        everythingVolatile $+        sComment "bind lhs" $         forM_ (zip acc_params arrs) $ \(acc_p, arr) ->-        ImpGen.copyDWIM (paramName acc_p) [] (Var arr) bucket+        copyDWIM (paramName acc_p) [] (Var arr) bucket -  let op_body = ImpGen.sComment "execute operation" $-                ImpGen.compileBody' acc_params $ lambdaBody op+  let op_body = sComment "execute operation" $+                compileBody' acc_params $ lambdaBody op        do_gen_reduce =-        ImpGen.everythingVolatile $-        ImpGen.sComment "update global result" $+        everythingVolatile $+        sComment "update global result" $         zipWithM_ (writeArray bucket) arrs $ map (Var . paramName) acc_params +      fence = case space of Space "local" -> sOp Imp.MemFenceLocal+                            _             -> sOp Imp.MemFenceGlobal++   -- While-loop: Try to insert your value   sWhile (Imp.var continue Bool) $ do     try_acquire_lock     sWhen lock_acquired $ do-      ImpGen.dLParams acc_params+      dLParams acc_params       bind_acc_params       op_body       do_gen_reduce-      sOp Imp.MemFence+      fence       release_lock       break_loop-    sOp Imp.MemFence-  where writeArray bucket arr val = ImpGen.copyDWIM arr bucket val []+    fence+  where writeArray bucket arr val = copyDWIM arr bucket val [] -atomicUpdateCAS :: PrimType+atomicUpdateCAS :: Space -> PrimType                 -> VName -> VName                 -> [Imp.Exp] -> VName-                -> ImpGen.ImpM lore Imp.KernelOp ()-                -> ImpGen.ImpM lore Imp.KernelOp ()-atomicUpdateCAS t arr old bucket x do_op = do+                -> ImpM lore Imp.KernelOp ()+                -> ImpM lore Imp.KernelOp ()+atomicUpdateCAS space t arr old bucket x do_op = do   -- Code generation target:   --   -- old = d_his[idx];@@ -448,9 +462,9 @@   -- } while(assumed != old);   assumed <- dPrim "assumed" t   run_loop <- dPrimV "run_loop" 1-  ImpGen.copyDWIM old [] (Var arr) bucket+  copyDWIM old [] (Var arr) bucket -  (arr', _a_space, bucket_offset) <- ImpGen.fullyIndexArray arr bucket+  (arr', _a_space, bucket_offset) <- fullyIndexArray arr bucket    -- While-loop: Try to insert your value   let (toBits, fromBits) =@@ -462,7 +476,7 @@     x <-- Imp.var assumed t     do_op     old_bits <- dPrim "old_bits" int32-    sOp $ Imp.Atomic $+    sOp $ Imp.Atomic space $       Imp.AtomicCmpXchg old_bits arr' bucket_offset       (toBits (Imp.var assumed t)) (toBits (Imp.var x t))     old <-- fromBits (Imp.var old_bits int32)@@ -488,62 +502,50 @@  computeKernelUses :: FreeIn a =>                      a -> [VName]-                  -> CallKernelGen ([Imp.KernelUse], [Imp.LocalMemoryUse])+                  -> CallKernelGen [Imp.KernelUse] computeKernelUses kernel_body bound_in_kernel = do   let actually_free = freeIn kernel_body `S.difference` S.fromList bound_in_kernel-   -- Compute the variables that we need to pass to the kernel.-  reads_from <- readsFromSet actually_free--  -- Are we using any local memory?-  local_memory <- computeLocalMemoryUse actually_free-  return (nub reads_from, nub local_memory)+  nub <$> readsFromSet actually_free  readsFromSet :: Names -> CallKernelGen [Imp.KernelUse] readsFromSet free =   fmap catMaybes $   forM (S.toList free) $ \var -> do     t <- lookupType var+    vtable <- getVTable     case t of       Array {} -> return Nothing-      Mem _ (Space "local") -> return Nothing+      Mem (Space "local") -> return Nothing       Mem {} -> return $ Just $ Imp.MemoryUse var       Prim bt ->-        isConstExp var >>= \case+        isConstExp vtable (Imp.var var bt) >>= \case           Just ce -> return $ Just $ Imp.ConstUse var ce           Nothing | bt == Cert -> return Nothing                   | otherwise  -> return $ Just $ Imp.ScalarUse var bt -computeLocalMemoryUse :: Names -> CallKernelGen [Imp.LocalMemoryUse]-computeLocalMemoryUse free =-  fmap catMaybes $-  forM (S.toList free) $ \var -> do-    t <- lookupType var-    case t of-      Mem memsize (Space "local") -> do-        memsize' <- localMemSize =<< ImpGen.subExpToDimSize memsize-        return $ Just (var, memsize')-      _ -> return Nothing--localMemSize :: Imp.MemSize -> CallKernelGen (Either Imp.MemSize Imp.KernelConstExp)-localMemSize (Imp.ConstSize x) =-  return $ Right $ ValueExp $ IntValue $ Int64Value x-localMemSize (Imp.VarSize v) = isConstExp v >>= \case-  Just e | isStaticExp e -> return $ Right e-  _ -> return $ Left $ Imp.VarSize v+localMemSize :: VTable ExplicitMemory -> Imp.Count Imp.Bytes+             -> ImpM lore op (Either (Imp.Count Imp.Bytes) Imp.KernelConstExp)+localMemSize vtable e = isConstExp vtable (Imp.innerExp e) >>= \case+  Just e' | isStaticExp e' -> return $ Right e'+  _ -> return $ Left e -isConstExp :: VName -> CallKernelGen (Maybe Imp.KernelConstExp)-isConstExp v = do-  vtable <- ImpGen.getVTable-  fname <- asks ImpGen.envFunction-  let lookupConstExp name = constExp =<< hasExp =<< M.lookup name vtable+isConstExp :: VTable ExplicitMemory -> Imp.Exp+           -> ImpM lore op (Maybe Imp.KernelConstExp)+isConstExp vtable size = do+  fname <- asks envFunction+  let onLeaf (Imp.ScalarVar name) _ = lookupConstExp name+      onLeaf (Imp.SizeOf pt) _ = Just $ primByteSize pt+      onLeaf Imp.Index{} _ = Nothing+      lookupConstExp name =+        constExp =<< hasExp =<< M.lookup name vtable       constExp (Op (Inner (GetSize key _))) =         Just $ LeafExp (Imp.SizeConst $ keyWithEntryPoint fname key) int32       constExp e = primExpFromExp lookupConstExp e-  return $ lookupConstExp v-  where hasExp (ImpGen.ArrayVar e _) = e-        hasExp (ImpGen.ScalarVar e _) = e-        hasExp (ImpGen.MemVar e _) = e+  return $ replaceInPrimExpM onLeaf size+  where hasExp (ArrayVar e _) = e+        hasExp (ScalarVar e _) = e+        hasExp (MemVar e _) = e  -- | Only some constant expressions qualify as *static* expressions, -- which we can use for static memory allocation.  This is a bit of a@@ -552,6 +554,8 @@ isStaticExp :: Imp.KernelConstExp -> Bool isStaticExp LeafExp{} = True isStaticExp ValueExp{} = True+isStaticExp (ConvOpExp ZExt{} x) = isStaticExp x+isStaticExp (ConvOpExp SExt{} x) = isStaticExp x isStaticExp (BinOpExp Add{} x y) = isStaticExp x && isStaticExp y isStaticExp (BinOpExp Sub{} x y) = isStaticExp x && isStaticExp y isStaticExp (BinOpExp Mul{} x y) = isStaticExp x && isStaticExp y@@ -562,9 +566,9 @@                        -> Imp.Count Imp.Elements                        -> Imp.Count Imp.Elements                        -> VName-                       -> ImpGen.ImpM lore op ()+                       -> ImpM lore op () computeThreadChunkSize (SplitStrided stride) thread_index elements_per_thread num_elements chunk_var = do-  stride' <- ImpGen.compileSubExp stride+  stride' <- toExp stride   chunk_var <--     Imp.BinOpExp (SMin Int32)     (Imp.innerExp elements_per_thread)@@ -590,36 +594,37 @@           Imp.innerExp num_elements .<.           (thread_index + 1) * Imp.innerExp elements_per_thread -kernelInitialisationSetSpace :: KernelSpace -> InKernelGen ()-                             -> ImpGen.ImpM lore op (KernelConstants, ImpGen.ImpM InKernel Imp.KernelOp ())-kernelInitialisationSetSpace space set_space = do-  group_size' <- ImpGen.compileSubExp $ spaceGroupSize space-  num_threads' <- ImpGen.compileSubExp $ spaceNumThreads space-  num_groups <- ImpGen.compileSubExp $ spaceNumGroups space--  let global_tid = spaceGlobalId space-      local_tid = spaceLocalId space-      group_id = spaceGroupId space+kernelInitialisationSimple :: Imp.Exp -> Imp.Exp+                           -> Maybe (VName, VName, VName)+                           -> CallKernelGen (KernelConstants, ImpM InKernel Imp.KernelOp ())+kernelInitialisationSimple num_groups group_size names = do+  (global_tid, local_tid, group_id) <-+    case names of Nothing ->+                    (,,)+                    <$> newVName "global_tid"+                    <*> newVName "local_tid"+                    <*> newVName "group_id"+                  Just (global_tid, local_tid, group_id) ->+                    return (global_tid, local_tid, group_id)   wave_size <- newVName "wave_size"   inner_group_size <- newVName "group_size"--  let (space_is, space_dims) = unzip $ spaceDimensions space-  space_dims' <- mapM ImpGen.compileSubExp space_dims+  vtable <- getVTable   let constants =-        KernelConstants+        KernelConstants vtable         (Imp.var global_tid int32)         (Imp.var local_tid int32)         (Imp.var group_id int32)         global_tid local_tid group_id-        group_size' num_groups num_threads'-        (Imp.var wave_size int32) (zip space_is space_dims')-        (if null (spaceDimensions space)-         then true else isActive (spaceDimensions space)) mempty+        group_size num_groups (group_size*num_groups)+        (Imp.var wave_size int32) []+        true mempty    let set_constants = do-        dPrim_ wave_size int32+        dPrim_ global_tid int32+        dPrim_ local_tid int32         dPrim_ inner_group_size int32-        ImpGen.dScope Nothing (scopeOfKernelSpace space)+        dPrim_ wave_size int32+        dPrim_ group_id int32          sOp (Imp.GetGlobalId global_tid 0)         sOp (Imp.GetLocalId local_tid 0)@@ -627,12 +632,42 @@         sOp (Imp.GetLockstepWidth wave_size)         sOp (Imp.GetGroupId group_id 0) +  return (constants, set_constants)++kernelInitialisationSetSpace :: KernelSpace -> InKernelGen ()+                             -> CallKernelGen (KernelConstants, ImpM InKernel Imp.KernelOp ())+kernelInitialisationSetSpace space set_space = do+  group_size <- toExp $ spaceGroupSize space+  num_groups <- toExp $ spaceNumGroups space++  (constants, set_constants) <-+    kernelInitialisationSimple num_groups group_size $+    Just (spaceGlobalId space, spaceLocalId space, spaceGroupId space)++  let set_constants' = do+        set_constants+        case spaceStructure space of+          FlatThreadSpace is_and_dims ->+            mapM_ ((`dPrim_` int32) . fst) is_and_dims+          NestedThreadSpace is_and_dims -> do+            let (gtids, _, ltids, _) = unzip4 is_and_dims+            mapM_ (`dPrim_` int32) $ gtids ++ ltids         set_space -  return (constants, set_constants)+  let (space_is, space_dims) = unzip $ spaceDimensions space+  space_dims' <- mapM toExp space_dims +  return (constants { kernelThreadActive =+                        if null $ spaceDimensions space+                        then true+                        else isActive $ spaceDimensions space+                    , kernelDimensions =+                        zip space_is space_dims'+                    },+          set_constants')+ kernelInitialisation :: KernelSpace-                     -> ImpGen.ImpM lore op (KernelConstants, ImpGen.ImpM InKernel Imp.KernelOp ())+                     -> CallKernelGen (KernelConstants, ImpM InKernel Imp.KernelOp ()) kernelInitialisation space =   kernelInitialisationSetSpace space $   setSpaceIndices (Imp.var (spaceGlobalId space) int32) space@@ -644,23 +679,22 @@       flatSpaceWith gtid is_and_dims     NestedThreadSpace is_and_dims -> do       let (gtids, gdims, ltids, ldims) = unzip4 is_and_dims-      gdims' <- mapM ImpGen.compileSubExp gdims-      ldims' <- mapM ImpGen.compileSubExp ldims+      gdims' <- mapM toExp gdims+      ldims' <- mapM toExp ldims       let (gtid_es, ltid_es) = unzip $ unflattenNestedIndex gdims' ldims' gtid       zipWithM_ (<--) gtids gtid_es       zipWithM_ (<--) ltids ltid_es   where flatSpaceWith base is_and_dims = do           let (is, dims) = unzip is_and_dims-          dims' <- mapM ImpGen.compileSubExp dims-          let index_expressions = unflattenIndex dims' base-          zipWithM_ (<--) is index_expressions+          dims' <- mapM toExp dims+          zipWithM_ (<--) is $ unflattenIndex dims' base  isActive :: [(VName, SubExp)] -> Imp.Exp isActive limit = case actives of                     [] -> Imp.ValueExp $ BoolValue True                     x:xs -> foldl (.&&.) x xs   where (is, ws) = unzip limit-        actives = zipWith active is $ map (ImpGen.compileSubExpOfType Bool) ws+        actives = zipWith active is $ map (toExp' Bool) ws         active i = (Imp.var i int32 .<.)  unflattenNestedIndex :: IntegralExp num => [num] -> [num] -> num -> [(num,num)]@@ -683,16 +717,16 @@ -- kernel). makeAllMemoryGlobal :: CallKernelGen a -> CallKernelGen a makeAllMemoryGlobal =-  local (\env -> env { ImpGen.envDefaultSpace = Imp.Space "global" }) .-  ImpGen.localVTable (M.map globalMemory)-  where globalMemory (ImpGen.MemVar _ entry)-          | ImpGen.entryMemSpace entry /= Space "local" =-              ImpGen.MemVar Nothing entry { ImpGen.entryMemSpace = Imp.Space "global" }+  local (\env -> env { envDefaultSpace = 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  allThreads :: KernelConstants -> InKernelGen () -> InKernelGen ()-allThreads constants = ImpGen.emit <=< ImpGen.subImpM_ (inKernelOperations constants')+allThreads constants = emit <=< subImpM_ (inKernelOperations constants')   where constants' =           constants { kernelThreadActive = Imp.ValueExp (BoolValue True) } @@ -700,34 +734,32 @@  writeParamToLocalMemory :: Typed (MemBound u) =>                            Imp.Exp -> (VName, t) -> Param (MemBound u)-                        -> ImpGen.ImpM lore op ()+                        -> ImpM lore op () writeParamToLocalMemory i (mem, _) param   | Prim t <- paramType param =-      ImpGen.emit $-      Imp.Write mem (bytes i') bt (Space "local") Imp.Volatile $+      emit $+      Imp.Write mem (elements i) bt (Space "local") Imp.Volatile $       Imp.var (paramName param) t   | otherwise =       return ()-  where i' = i * Imp.LeafExp (Imp.SizeOf bt) int32-        bt = elemType $ paramType param+  where bt = elemType $ paramType param  readParamFromLocalMemory :: Typed (MemBound u) =>-                            VName -> Imp.Exp -> Param (MemBound u) -> (VName, t)-                         -> ImpGen.ImpM lore op ()-readParamFromLocalMemory index i param (l_mem, _)+                            Imp.Exp -> Param (MemBound u) -> (VName, t)+                         -> ImpM lore op ()+readParamFromLocalMemory i param (l_mem, _)   | Prim _ <- paramType param =       paramName param <---      Imp.index l_mem (bytes i') bt (Space "local") Imp.Volatile-  | otherwise = index <-- i-  where i' = i * Imp.LeafExp (Imp.SizeOf bt) int32-        bt = elemType $ paramType param+      Imp.index l_mem (elements i) bt (Space "local") Imp.Volatile+  | otherwise = return ()+  where bt = elemType $ paramType param  groupReduce :: ExplicitMemorish lore =>                KernelConstants             -> Imp.Exp             -> Lambda lore             -> [VName]-            -> ImpGen.ImpM lore Imp.KernelOp ()+            -> ImpM lore Imp.KernelOp () groupReduce constants w lam arrs = do   offset <- dPrim "offset" int32   groupReduceWithOffset constants offset w lam arrs@@ -738,26 +770,26 @@                       -> Imp.Exp                       -> Lambda lore                       -> [VName]-                      -> ImpGen.ImpM lore Imp.KernelOp ()+                      -> ImpM lore Imp.KernelOp () groupReduceWithOffset constants offset w lam arrs = do   let (reduce_acc_params, reduce_arr_params) = splitAt (length arrs) $ lambdaParams lam    skip_waves <- dPrim "skip_waves" int32-  ImpGen.dLParams $ lambdaParams lam+  dLParams $ lambdaParams lam    offset <-- 0 -  ImpGen.comment "participating threads read initial accumulator" $+  comment "participating threads read initial accumulator" $     sWhen (local_tid .<. w) $     zipWithM_ readReduceArgument reduce_acc_params arrs -  let do_reduce = do ImpGen.comment "read array element" $+  let do_reduce = do comment "read array element" $                        zipWithM_ readReduceArgument reduce_arr_params arrs-                     ImpGen.comment "apply reduction operation" $-                       ImpGen.compileBody' reduce_acc_params $ lambdaBody lam-                     ImpGen.comment "write result of operation" $+                     comment "apply reduction operation" $+                       compileBody' reduce_acc_params $ lambdaBody lam+                     comment "write result of operation" $                        zipWithM_ writeReduceOpResult reduce_acc_params arrs-      in_wave_reduce = ImpGen.everythingVolatile do_reduce+      in_wave_reduce = everythingVolatile do_reduce        wave_size = kernelWaveSize constants       group_size = kernelGroupSize constants@@ -805,15 +837,15 @@          readReduceArgument param arr           | Prim _ <- paramType param = do-              let i = local_tid + ImpGen.varIndex offset-              ImpGen.copyDWIM (paramName param) [] (Var arr) [i]+              let i = local_tid + Imp.vi32 offset+              copyDWIM (paramName param) [] (Var arr) [i]           | otherwise = do-              let i = global_tid + ImpGen.varIndex offset-              ImpGen.copyDWIM (paramName param) [] (Var arr) [i]+              let i = global_tid + Imp.vi32 offset+              copyDWIM (paramName param) [] (Var arr) [i]          writeReduceOpResult param arr           | Prim _ <- paramType param =-              ImpGen.copyDWIM arr [local_tid] (Var $ paramName param) []+              copyDWIM arr [local_tid] (Var $ paramName param) []           | otherwise =               return () @@ -822,7 +854,7 @@           -> Imp.Exp           -> Lambda InKernel           -> [VName]-          -> ImpGen.ImpM InKernel Imp.KernelOp ()+          -> ImpM InKernel Imp.KernelOp () groupScan constants seg_flag w lam arrs = do   when (any (not . primType . paramType) $ lambdaParams lam) $     compilerLimitationS "Cannot compile parallel scans with array element type."@@ -830,16 +862,13 @@   renamed_lam <- renameLambda lam    acc_local_mem <- flip zip (repeat ()) <$>-                   mapM (fmap (ImpGen.memLocationName . ImpGen.entryArrayLocation) .-                         ImpGen.lookupArray) arrs+                   mapM (fmap (memLocationName . entryArrayLocation) .+                         lookupArray) arrs    let ltid = kernelLocalThreadId constants-      (lam_i, other_index_param, actual_params) =-        partitionChunkedKernelLambdaParameters $ lambdaParams lam-      (x_params, y_params) = splitAt (length arrs) actual_params+      (x_params, y_params) = splitAt (length arrs) $ lambdaParams lam -  ImpGen.dLParams (lambdaParams lam++lambdaParams renamed_lam)-  lam_i <-- ltid+  dLParams (lambdaParams lam++lambdaParams renamed_lam)    -- The scan works by splitting the group into blocks, which are   -- scanned separately.  Typically, these blocks are smaller than@@ -873,23 +902,22 @@         flag_true <- seg_flag         Just $ \from to ->           flag_true (from*block_size+block_size-1) (to*block_size+block_size-1)-  ImpGen.comment+  comment     "scan the first block, after which offset 'i' contains carry-in for warp 'i+1'" $     doInBlockScan first_block_seg_flag (is_first_block .&&. ltid_in_bounds) renamed_lam    sOp Imp.LocalBarrier    let read_carry_in =-        zipWithM_ (readParamFromLocalMemory-                   (paramName other_index_param) (block_id - 1))+        zipWithM_ (readParamFromLocalMemory (block_id - 1))         x_params acc_local_mem    let op_to_y         | Nothing <- seg_flag =-            ImpGen.compileBody' y_params $ lambdaBody lam+            compileBody' y_params $ lambdaBody lam         | Just flag_true <- seg_flag =             sUnless (flag_true (block_id*block_size-1) ltid) $-              ImpGen.compileBody' y_params $ lambdaBody lam+              compileBody' y_params $ lambdaBody lam       write_final_result =         zipWithM_ (writeParamToLocalMemory ltid) acc_local_mem y_params @@ -914,30 +942,27 @@             -> [(VName, t)]             -> Lambda InKernel             -> InKernelGen ()-inBlockScan seg_flag lockstep_width block_size active ltid acc_local_mem scan_lam = ImpGen.everythingVolatile $ do+inBlockScan seg_flag lockstep_width block_size active ltid acc_local_mem scan_lam = everythingVolatile $ do   skip_threads <- dPrim "skip_threads" int32   let in_block_thread_active =         Imp.var skip_threads int32 .<=. in_block_id-      (scan_lam_i, other_index_param, actual_params) =-        partitionChunkedKernelLambdaParameters $ lambdaParams scan_lam+      actual_params = lambdaParams scan_lam       (x_params, y_params) =         splitAt (length actual_params `div` 2) actual_params       read_operands =-        zipWithM_ (readParamFromLocalMemory (paramName other_index_param) $-                   ltid - Imp.var skip_threads int32)+        zipWithM_ (readParamFromLocalMemory $ ltid - Imp.var skip_threads int32)         x_params acc_local_mem    -- Set initial y values   sWhen active $-    zipWithM_ (readParamFromLocalMemory scan_lam_i ltid)-    y_params acc_local_mem+    zipWithM_ (readParamFromLocalMemory ltid) y_params acc_local_mem    let op_to_y         | Nothing <- seg_flag =-            ImpGen.compileBody' y_params $ lambdaBody scan_lam+            compileBody' y_params $ lambdaBody scan_lam         | Just flag_true <- seg_flag =             sUnless (flag_true (ltid-Imp.var skip_threads int32) ltid) $-              ImpGen.compileBody' y_params $ lambdaBody scan_lam+              compileBody' y_params $ lambdaBody scan_lam       write_operation_result =         zipWithM_ (writeParamToLocalMemory ltid) acc_local_mem y_params       maybeLocalBarrier = sWhen (lockstep_width .<=. Imp.var skip_threads int32) $@@ -962,14 +987,14 @@   where block_id = ltid `quot` block_size         in_block_id = ltid - block_id * block_size -compileKernelStms :: KernelConstants -> [Stm InKernel]+compileKernelStms :: KernelConstants -> Stms InKernel                   -> InKernelGen a                   -> InKernelGen a compileKernelStms constants ungrouped_bnds m =-  compileGroupedKernelStms' $ groupStmsByGuard constants ungrouped_bnds+  compileGroupedKernelStms' $ groupStmsByGuard constants $ stmsToList ungrouped_bnds   where compileGroupedKernelStms' [] = m         compileGroupedKernelStms' ((g, bnds):rest_bnds) = do-          ImpGen.dScopes (map ((Just . stmExp) &&& (castScope . scopeOf)) bnds)+          dScopes (map ((Just . stmExp) &&& (castScope . scopeOf)) bnds)           protect g $ mapM_ compileKernelStm bnds           compileGroupedKernelStms' rest_bnds @@ -980,7 +1005,7 @@         protect (Just g) body_m =           sWhen g $ allThreads constants body_m -        compileKernelStm (Let pat _ e) = ImpGen.compileExp pat e+        compileKernelStm (Let pat _ e) = compileExp pat e  groupStmsByGuard :: KernelConstants                      -> [Stm InKernel]@@ -1000,7 +1025,7 @@ computeMapKernelGroups :: Imp.Exp -> CallKernelGen (Imp.Exp, Imp.Exp) computeMapKernelGroups kernel_size = do   group_size <- dPrim "group_size" int32-  fname <- asks ImpGen.envFunction+  fname <- asks envFunction   let group_size_var = Imp.var group_size int32       group_size_key = keyWithEntryPoint fname $ nameFromString $ pretty group_size   sOp $ Imp.GetSize group_size group_size_key Imp.SizeGroup@@ -1008,7 +1033,7 @@   return (Imp.var group_size int32, Imp.var num_groups int32)  simpleKernelConstants :: Imp.Exp -> String-                      -> CallKernelGen (KernelConstants, ImpGen.ImpM InKernel Imp.KernelOp ())+                      -> CallKernelGen (KernelConstants, ImpM InKernel Imp.KernelOp ()) simpleKernelConstants kernel_size desc = do   thread_gtid <- newVName $ desc ++ "_gtid"   thread_ltid <- newVName $ desc ++ "_ltid"@@ -1022,7 +1047,9 @@         sOp (Imp.GetLocalId thread_ltid 0)         sOp (Imp.GetGroupId group_id 0) -  return (KernelConstants++  vtable <- getVTable+  return (KernelConstants vtable           (Imp.var thread_gtid int32) (Imp.var thread_ltid int32) (Imp.var group_id int32)           thread_gtid thread_ltid group_id           group_size num_groups (group_size*num_groups) 0@@ -1030,20 +1057,43 @@            set_constants) -sKernel :: KernelConstants -> String -> ImpGen.ImpM InKernel Imp.KernelOp a -> CallKernelGen ()+-- | 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 :: KernelConstants+                 -> Imp.Exp+                 -> (VName -> InKernelGen ())+                 -> InKernelGen ()+virtualiseGroups constants required_groups m+  | kernelNumGroups constants == required_groups =+      m $ kernelGroupIdVar constants+  | otherwise = do+  phys_group_id <- dPrim "phys_group_id" int32+  sOp $ Imp.GetGroupId phys_group_id 0+  let iterations = (required_groups - Imp.vi32 phys_group_id) `quotRoundingUp`+                   kernelNumGroups constants+  i <- newVName "i"+  sFor i Int32 iterations $+    m =<< dPrimV "virt_group_id" (Imp.vi32 phys_group_id + Imp.vi32 i * kernelNumGroups constants)++sKernel :: KernelConstants -> String -> ImpM InKernel Imp.KernelOp a -> CallKernelGen () sKernel constants name m = do-  body <- makeAllMemoryGlobal $-          ImpGen.subImpM_ (inKernelOperations constants) m-  (uses, local_memory) <- computeKernelUses body mempty-  ImpGen.emit $ Imp.Op $ Imp.CallKernel Imp.Kernel+  body <- makeAllMemoryGlobal $ subImpM_ (inKernelOperations constants) m+  uses <- computeKernelUses body mempty++  emit $ Imp.Op $ Imp.CallKernel Imp.Kernel     { Imp.kernelBody = body-    , Imp.kernelLocalMemory = local_memory     , Imp.kernelUses = uses     , Imp.kernelNumGroups = [kernelNumGroups constants]     , Imp.kernelGroupSize = [kernelGroupSize constants]     , Imp.kernelName =-        nameFromString $ name ++ "_" ++ show (baseTag $ kernelGlobalThreadIdVar constants)+        nameFromString $ name ++ "_" ++ show tag     }+  where tag = baseTag $ kernelGlobalThreadIdVar constants  -- | Perform a Replicate with a kernel. sReplicate :: VName -> Shape -> SubExp@@ -1051,7 +1101,7 @@ sReplicate arr (Shape ds) se = do   t <- subExpType se -  dims <- mapM ImpGen.compileSubExp $ ds ++ arrayDims t+  dims <- mapM toExp $ ds ++ arrayDims t   (constants, set_constants) <-     simpleKernelConstants (product dims) "replicate" @@ -1060,39 +1110,38 @@   sKernel constants "replicate" $ do     set_constants     sWhen (kernelThreadActive constants) $-      ImpGen.copyDWIM arr is' se $ drop (length ds) is'+      copyDWIM arr is' se $ drop (length ds) is'  -- | Perform an Iota with a kernel. sIota :: VName -> Imp.Exp -> Imp.Exp -> Imp.Exp -> IntType       -> CallKernelGen () sIota arr n x s et = do-  destloc <- ImpGen.entryArrayLocation <$> ImpGen.lookupArray arr+  destloc <- entryArrayLocation <$> lookupArray arr   (constants, set_constants) <- simpleKernelConstants n "iota"    sKernel constants "iota" $ do     set_constants     let gtid = kernelGlobalThreadId constants     sWhen (kernelThreadActive constants) $ do-      (destmem, destspace, destidx) <--        ImpGen.fullyIndexArray' destloc [gtid] (IntType et)+      (destmem, destspace, destidx) <- fullyIndexArray' destloc [gtid] -      ImpGen.emit $+      emit $         Imp.Write destmem destidx (IntType et) destspace Imp.Nonvolatile $         Imp.ConvOpExp (SExt Int32 et) gtid * s + x  sCopy :: PrimType-      -> ImpGen.MemLocation-      -> ImpGen.MemLocation+      -> MemLocation+      -> MemLocation       -> Imp.Count Imp.Elements       -> CallKernelGen () sCopy bt-  destloc@(ImpGen.MemLocation destmem _ _)-  srcloc@(ImpGen.MemLocation srcmem srcshape _)+  destloc@(MemLocation destmem _ _)+  srcloc@(MemLocation srcmem srcshape _)   n = do   -- Note that the shape of the destination and the source are   -- necessarily the same.   let shape = map Imp.sizeToExp srcshape-      shape_se = map (Imp.innerExp . ImpGen.dimSizeToExp) srcshape+      shape_se = map (Imp.innerExp . dimSizeToExp) srcshape       kernel_size = Imp.innerExp n * product (drop 1 shape)    (constants, set_constants) <- simpleKernelConstants kernel_size "copy"@@ -1104,9 +1153,87 @@         dest_is = unflattenIndex shape_se gtid         src_is = dest_is -    (_, destspace, destidx) <- ImpGen.fullyIndexArray' destloc dest_is bt-    (_, srcspace, srcidx) <- ImpGen.fullyIndexArray' srcloc src_is bt+    (_, destspace, destidx) <- fullyIndexArray' destloc dest_is+    (_, srcspace, srcidx) <- fullyIndexArray' srcloc src_is -    sWhen (gtid .<. kernel_size) $ ImpGen.emit $+    sWhen (gtid .<. kernel_size) $ emit $       Imp.Write destmem destidx bt destspace Imp.Nonvolatile $       Imp.index srcmem srcidx bt srcspace Imp.Nonvolatile+++compileKernelResult :: KernelConstants -> PatElem InKernel -> KernelResult+                    -> InKernelGen ()++compileKernelResult constants pe (GroupsReturn what) = do+  i <- newVName "i"++  in_local_memory <- arrayInLocalMemory what+  let me = kernelLocalThreadId constants++  if not in_local_memory then do+    who' <- toExp $ intConst Int32 0+    sWhen (me .==. who') $+      copyDWIM (patElemName pe) [kernelGroupId constants] what []+    else do+      -- 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).+      --+      -- We do the reads/writes multidimensionally, but the loop is+      -- single-dimensional.+      ws <- mapM toExp . arrayDims =<< subExpType what+      -- Compute how many elements this thread is responsible for.+      -- Formula: (w - ltid) / group_size (rounded up).+      let w = product ws+          ltid = kernelLocalThreadId constants+          group_size = kernelGroupSize constants+          to_write = (w - ltid) `quotRoundingUp` group_size+          is = unflattenIndex ws $ Imp.vi32 i * group_size + ltid++      sFor i Int32 to_write $+        copyDWIM (patElemName pe) (kernelGroupId constants : is) what is++compileKernelResult constants pe (ThreadsReturn what) = do+  let is = map (Imp.vi32 . fst) $ kernelDimensions constants+  sWhen (kernelThreadActive constants) $ copyDWIM (patElemName pe) is what []++compileKernelResult constants pe (ConcatReturns SplitContiguous _ per_thread_elems moffset what) = do+  dest_loc <- entryArrayLocation <$> lookupArray (patElemName pe)+  let dest_loc_offset = offsetArray dest_loc offset+      dest' = arrayDestination dest_loc_offset+  copyDWIMDest dest' [] (Var what) []+  where offset = case moffset of+                   Nothing -> toExp' int32 per_thread_elems *+                              kernelGlobalThreadId constants+                   Just se -> toExp' int32 se++compileKernelResult constants pe (ConcatReturns (SplitStrided stride) _ _ moffset what) = do+  dest_loc <- entryArrayLocation <$> lookupArray (patElemName pe)+  let dest_loc' = strideArray+                  (offsetArray dest_loc offset) $+                  toExp' int32 stride+      dest' = arrayDestination dest_loc'+  copyDWIMDest dest' [] (Var what) []+  where offset = case moffset of+                   Nothing -> kernelGlobalThreadId constants+                   Just se -> toExp' int32 se++compileKernelResult constants pe (WriteReturn rws _arr dests) = do+  rws' <- mapM toExp rws+  forM_ dests $ \(is, e) -> do+    is' <- mapM toExp is+    let condInBounds i rw = 0 .<=. i .&&. i .<. rw+        write = foldl (.&&.) (kernelThreadActive constants) $+                zipWith condInBounds is' rws'+    sWhen write $ copyDWIM (patElemName pe) (map (toExp' int32) is) e []++arrayInLocalMemory :: SubExp -> InKernelGen Bool+arrayInLocalMemory (Var name) = do+  res <- lookupVar name+  case res of+    ArrayVar _ entry ->+      (Space "local"==) . entryMemSpace <$>+      lookupMemory (memLocationName (entryArrayLocation entry))+    _ -> return False+arrayInLocalMemory Constant{} = return False

@@ -1,5 +1,6 @@ {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE OverloadedStrings #-} -- | Our compilation strategy for 'SegGenRed' is based around avoiding -- bin conflicts.  We do this by splitting the input into chunks, and -- for each chunk computing a single subhistogram.  Then we combine@@ -8,11 +9,40 @@ -- There are some branches around to efficiently handle the case where -- we use only a single subhistogram (because it's large), so that we -- respect the asymptotics, and do not copy the destination array.+--+-- We also use a heuristic strategy for computing subhistograms in+-- local memory when possible.  Given:+--+-- H: total size of histograms in bytes, including any lock arrays.+--+-- G: group size+--+-- T: number of bytes of local memory each thread can be given without+-- impacting occupancy (determined experimentally, e.g. 32).+--+-- LMAX: maximum amount of local memory per workgroup (hard limit).+--+-- We wish to compute:+--+-- COOP: cooperation level (number of threads per subhistogram)+--+-- LH: number of local memory subhistograms+--+-- We do this as:+--+-- COOP = ceil(H / T)+-- LH = ceil((G*T)/H)+-- if COOP <= G && H <= LMAX then+--   use local memory+-- else+--   use global memory+ module Futhark.CodeGen.ImpGen.Kernels.SegGenRed   ( compileSegGenRed )   where  import Control.Monad.Except+import Data.Either import Data.Maybe import Data.List @@ -23,25 +53,102 @@ import qualified Futhark.Representation.ExplicitMemory.IndexFunction as IxFun import Futhark.Pass.ExplicitAllocations() import qualified Futhark.CodeGen.ImpCode.Kernels as Imp-import qualified Futhark.CodeGen.ImpGen as ImpGen-import Futhark.CodeGen.ImpGen ((<--),-                               sFor, sComment, sIf, sWhen, sArray,-                               dPrim_, dPrimV)+import Futhark.CodeGen.ImpGen import Futhark.CodeGen.ImpGen.Kernels.SegRed (compileSegRed') import Futhark.CodeGen.ImpGen.Kernels.Base import Futhark.Util.IntegralExp (quotRoundingUp, quot, rem) import Futhark.Util (chunks, mapAccumLM, splitFromEnd, takeLast) import Futhark.Construct (fullSliceNum) -prepareAtomicUpdate :: Maybe Locking -> [VName] -> Lambda InKernel-                    -> CallKernelGen (Maybe Locking,-                                      [Imp.Exp] -> ImpGen.ImpM InKernel Imp.KernelOp ())-prepareAtomicUpdate l dests lam =-  -- We need a separate lock array if the opterators are not all of a+i32Toi64 :: PrimExp v -> PrimExp v+i32Toi64 = ConvOpExp (SExt Int32 Int64)++data SubhistosInfo = SubhistosInfo { subhistosArray :: VName+                                   , subhistosAlloc :: CallKernelGen ()+                                   }++data SegGenRedSlug = SegGenRedSlug+                     { slugOp :: GenReduceOp InKernel+                     , slugNumSubhistos :: VName+                     , slugSubhistos :: [SubhistosInfo]+                     }++-- | Figure out how much memory is needed per histogram, and compute+-- some other auxiliary information.+computeHistoUsage :: KernelSpace+                  -> GenReduceOp InKernel+                  -> CallKernelGen (Imp.Count Imp.Bytes, SegGenRedSlug)+computeHistoUsage space op = do+  let segment_dims = init $ spaceDimensions space+      num_segments = length segment_dims++  op_h <- fmap (sum . map typeSize) $ mapM lookupType $ genReduceDest op++  -- Create names for the intermediate array memory blocks,+  -- memory block sizes, arrays, and number of subhistograms.+  num_subhistos <- dPrim "num_subhistos" int32+  subhisto_infos <- forM (zip (genReduceDest op) (genReduceNeutral op)) $ \(dest, ne) -> do+    dest_t <- lookupType dest+    dest_mem <- entryArrayLocation <$> lookupArray dest++    subhistos_mem <-+      sDeclareMem (baseString dest ++ "_subhistos_mem") (Space "device")++    let subhistos_shape = Shape (map snd segment_dims++[Var num_subhistos]) <>+                          stripDims num_segments (arrayShape dest_t)+        subhistos_membind = ArrayIn subhistos_mem $ IxFun.iota $+                            map (primExpFromSubExp int32) $ shapeDims subhistos_shape+    subhistos <- sArray (baseString dest ++ "_subhistos")+                 (elemType dest_t) subhistos_shape subhistos_membind++    return $ SubhistosInfo subhistos $ do+      let unitHistoCase =+            emit $+            Imp.SetMem subhistos_mem (memLocationName dest_mem) $+            Space "device"++          multiHistoCase = do+            let num_elems = foldl' (*) (Imp.var num_subhistos int32) $+                            map (toExp' int32) $ arrayDims dest_t++            let subhistos_mem_size =+                  Imp.bytes $+                  Imp.innerExp (Imp.elements num_elems `Imp.withElemType` elemType dest_t)++            sAlloc_ subhistos_mem subhistos_mem_size $ Space "device"+            sReplicate subhistos (Shape (map snd segment_dims +++                                         [Var num_subhistos, genReduceWidth op]) <>+                                  genReduceShape op) ne+            subhistos_t <- lookupType subhistos+            let slice = fullSliceNum (map (toExp' int32) $ arrayDims subhistos_t) $+                        map (unitSlice 0 . toExp' int32 . snd) segment_dims +++                        [DimFix 0]+            sUpdate subhistos slice $ Var dest++      sIf (Imp.var num_subhistos int32 .==. 1) unitHistoCase multiHistoCase++  return (op_h, SegGenRedSlug op num_subhistos subhisto_infos)++localMemLockArray :: KernelSpace -> Type+localMemLockArray space = Array int32 (Shape [spaceGroupSize space]) NoUniqueness++-- | How many bytes will be spent on lock arrays if we use a local+-- memory implementation?+localMemLockUsage :: KernelSpace -> [SegGenRedSlug] -> Imp.Count Imp.Bytes+localMemLockUsage space slugs =+  if any (isRight . atomicUpdateLocking . genReduceOp . slugOp) slugs+  then typeSize $ localMemLockArray space+  else 0++prepareAtomicUpdateGlobal :: Maybe Locking -> [VName] -> SegGenRedSlug+                          -> CallKernelGen (Maybe Locking,+                                            [Imp.Exp] -> ImpM InKernel Imp.KernelOp ())+prepareAtomicUpdateGlobal l dests slug =+  -- We need a separate lock array if the operators are not all of a   -- particularly simple form that permits pure atomic operations.-  case (l, atomicUpdateLocking lam) of-    (_, Left f) -> return (l, f dests)-    (Just l', Right f) -> return (l, f l' dests)+  case (l, atomicUpdateLocking $ genReduceOp $ slugOp slug) of+    (_, Left f) -> return (l, f (Space "global") dests)+    (Just l', Right f) -> return (l, f l' (Space "global") dests)     (Nothing, Right f) -> do       -- The number of locks used here is too low, but since we are       -- currently forced to inline a huge list, I'm keeping it down@@ -51,120 +158,105 @@       -- A fun solution would also be to use a simple hashing       -- algorithm to ensure good distribution of locks.       let num_locks = 10000+          dims = map (toExp' int32) $+                 shapeDims (genReduceShape (slugOp slug)) +++                 [ Var (slugNumSubhistos slug)+                 , genReduceWidth (slugOp slug)]       locks <--        ImpGen.sStaticArray "genred_locks" (Space "device") int32 $+        sStaticArray "genred_locks" (Space "device") int32 $         Imp.ArrayZeros num_locks-      let l' = Locking locks 0 1 0 ((`rem` fromIntegral num_locks) . sum)-      return (Just l', f l' dests)+      let l' = Locking locks 0 1 0 ((`rem` fromIntegral num_locks) . flattenIndex dims)+      return (Just l', f l' (Space "global") dests) -prepareIntermediateArrays :: [SubExp] -> Imp.Exp -> [GenReduceOp InKernel]-                          -> CallKernelGen-                             [(VName,-                               [VName],-                               [Imp.Exp] -> ImpGen.ImpM InKernel Imp.KernelOp ())]-prepareIntermediateArrays segment_dims num_threads = fmap snd . mapAccumLM onOp Nothing+prepareIntermediateArraysGlobal :: Imp.Exp -> [SegGenRedSlug]+                                -> CallKernelGen+                                   [(VName,+                                     [VName],+                                     [Imp.Exp] -> ImpM InKernel Imp.KernelOp ())]+prepareIntermediateArraysGlobal num_threads = fmap snd . mapAccumLM onOp Nothing   where-    onOp l op = do+    onOp l slug@(SegGenRedSlug op num_subhistos subhisto_info) = do       -- Determining the degree of cooperation (heuristic):       -- coop_lvl   := size of histogram (Cooperation level)       -- num_histos := (threads / coop_lvl) (Number of histograms)       -- threads    := min(physical_threads, segment_size)       --       -- Careful to avoid division by zero when genReduceWidth==0.-      num_histos <- dPrimV "num_histos" $ num_threads `quotRoundingUp`-                    BinOpExp (SMax Int32) 1 (ImpGen.compileSubExpOfType int32 (genReduceWidth op))+      num_subhistos <--+        num_threads `quotRoundingUp`+        BinOpExp (SMax Int32) 1 (toExp' int32 (genReduceWidth op)) -      ImpGen.emit $ Imp.DebugPrint "num_histograms" int32 $ Imp.var num_histos int32+      emit $ Imp.DebugPrint "Number of subhistograms in global memory" $+        Just (int32, Imp.vi32 num_subhistos)        -- Initialise sub-histograms.       ---      -- If num_histos is 1, then we just reuse the original+      -- If num_subhistos is 1, then we just reuse the original       -- destination.  The idea is to avoid a copy if we are writing a       -- small number of values into a very large prior histogram.--      dests <- forM (zip (genReduceDest op) (genReduceNeutral op)) $ \(dest, ne) -> do-        dest_t <- lookupType dest-        dest_mem <- ImpGen.entryArrayLocation <$> ImpGen.lookupArray dest-        let num_elems = foldl' (*) (Imp.var num_histos int32) $-                        map (ImpGen.compileSubExpOfType int32) $-                        arrayDims dest_t-        let size = Imp.elements num_elems `Imp.withElemType` int32--        (sub_mem, size') <--          ImpGen.sDeclareMem "subhistogram_mem" size $ Space "device"+      dests <- forM (zip (genReduceDest op) subhisto_info) $ \(dest, info) -> do+        dest_mem <- entryArrayLocation <$> lookupArray dest -        let num_segments = length segment_dims-            sub_shape = Shape (segment_dims++[Var num_histos]) <>-                        stripDims num_segments (arrayShape dest_t)-            sub_membind = ArrayIn sub_mem $ IxFun.iota $-                          map (primExpFromSubExp int32) $ shapeDims sub_shape-        subhisto <- sArray "genred_dest" (elemType dest_t) sub_shape sub_membind+        sub_mem <- fmap memLocationName $+                   entryArrayLocation <$>+                   lookupArray (subhistosArray info)          let unitHistoCase =-              ImpGen.emit $-              Imp.SetMem sub_mem (ImpGen.memLocationName dest_mem) $+              emit $+              Imp.SetMem sub_mem (memLocationName dest_mem) $               Space "device" -            multiHistoCase = do-              ImpGen.sAlloc_ sub_mem size' $ Space "device"-              sReplicate subhisto (Shape $ segment_dims ++ [Var num_histos, genReduceWidth op]) ne-              subhisto_t <- lookupType subhisto-              let slice = fullSliceNum (map (ImpGen.compileSubExpOfType int32) $ arrayDims subhisto_t) $-                          map (unitSlice 0 . ImpGen.compileSubExpOfType int32) segment_dims ++-                          [DimFix 0]-              ImpGen.sUpdate subhisto slice $ Var dest+            multiHistoCase = subhistosAlloc info -        sIf (Imp.var num_histos int32 .==. 1) unitHistoCase multiHistoCase+        sIf (Imp.var num_subhistos int32 .==. 1) unitHistoCase multiHistoCase -        return subhisto+        return $ subhistosArray info -      (l', do_op) <- prepareAtomicUpdate l dests $ genReduceOp op+      (l', do_op) <- prepareAtomicUpdateGlobal l dests slug -      return (l', (num_histos, dests, do_op))+      return (l', (num_subhistos, dests, do_op)) -genRedKernel :: [PatElem ExplicitMemory]-             -> KernelSpace-             -> [GenReduceOp InKernel]-             -> Body InKernel-             -> CallKernelGen [(VName, [VName])]-genRedKernel map_pes space ops body = do+genRedKernelGlobal :: [PatElem ExplicitMemory]+                  -> KernelSpace+                  -> [SegGenRedSlug]+                  -> KernelBody InKernel+                  -> CallKernelGen ()+genRedKernelGlobal map_pes space slugs kbody = do   (base_constants, init_constants) <- kernelInitialisationSetSpace space $ return ()   let constants = base_constants { kernelThreadActive = true }       (space_is, space_sizes) = unzip $ spaceDimensions space-      i32_to_i64 = ConvOpExp (SExt Int32 Int64)-      space_sizes_64 = map (i32_to_i64 . ImpGen.compileSubExpOfType int32) space_sizes+      space_sizes_64 = map (i32Toi64 . toExp' int32) space_sizes       total_w_64 = product space_sizes_64 -  histograms <- prepareIntermediateArrays (init space_sizes) (kernelNumThreads constants) ops+  histograms <- prepareIntermediateArraysGlobal (kernelNumThreads constants) slugs    elems_per_thread_64 <- dPrimV "elems_per_thread_64" $                          total_w_64 `quotRoundingUp`                          ConvOpExp (SExt Int32 Int64) (kernelNumThreads constants) -  sKernel constants "seggenred" $ allThreads constants $ do+  sKernel constants "seggenred_global" $ allThreads constants $ do     init_constants -    i <- newVName "i"-     -- Compute subhistogram index for each thread, per histogram.     subhisto_inds <- forM histograms $ \(num_histograms, _, _) ->       dPrimV "subhisto_ind" $       kernelGlobalThreadId constants `quot`       (kernelNumThreads constants `quotRoundingUp` Imp.var num_histograms int32) -    sFor i Int64 (Imp.var elems_per_thread_64 int64) $ do+    flat_idx <- newVName "flat_idx"+    sFor flat_idx Int64 (Imp.var elems_per_thread_64 int64) $ do       -- Compute the offset into the input and output.  To this a       -- thread can add its local ID to figure out which element it is       -- responsible for.  The calculation is done with 64-bit       -- integers to avoid overflow, but the final segment indexes are       -- 32 bit.       offset <- dPrimV "offset" $-                (i32_to_i64 (kernelGroupId constants) *+                (i32Toi64 (kernelGroupId constants) *                  (Imp.var elems_per_thread_64 int64 *-                  i32_to_i64 (kernelGroupSize constants)))-                + (Imp.var i int64 * i32_to_i64 (kernelGroupSize constants))+                  i32Toi64 (kernelGroupSize constants)))+                + (Imp.var flat_idx int64 * i32Toi64 (kernelGroupSize constants)) -      j <- dPrimV "j" $ Imp.var offset int64 + i32_to_i64 (kernelLocalThreadId constants)+      j <- dPrimV "j" $ Imp.var offset int64 + i32Toi64 (kernelLocalThreadId constants)        -- Construct segment indices.       let setIndex v e = do dPrim_ v int32@@ -178,98 +270,396 @@       -- arrays.       let input_in_bounds = Imp.var j int32 .<. total_w_64 -      sWhen input_in_bounds $ ImpGen.compileStms mempty (stmsToList $ bodyStms body) $ do-        let (red_res, map_res) = splitFromEnd (length map_pes) $ bodyResult body+      sWhen input_in_bounds $ compileStms mempty (kernelBodyStms kbody) $ do+        let (red_res, map_res) = splitFromEnd (length map_pes) $ kernelBodyResult kbody          sComment "save map-out results" $-          forM_ (zip map_pes map_res) $ \(pe, se) ->-          ImpGen.copyDWIM (patElemName pe)-          (map ((`Imp.var` int32) . fst) $ kernelDimensions constants) se []+          forM_ (zip map_pes map_res) $ \(pe, res) ->+          copyDWIM (patElemName pe)+          (map ((`Imp.var` int32) . fst) $ kernelDimensions constants)+          (kernelResultSubExp res) [] -        let (buckets, vs) = splitAt (length ops) red_res-            perOp = chunks $ map (length . genReduceDest) ops+        let (buckets, vs) = splitAt (length slugs) red_res+            perOp = chunks $ map (length . genReduceDest . slugOp) slugs          sComment "perform atomic updates" $-          forM_ (zip5 ops histograms buckets (perOp vs) subhisto_inds) $+          forM_ (zip5 (map slugOp slugs) histograms buckets (perOp vs) subhisto_inds) $           \(GenReduceOp dest_w _ _ shape lam,             (_, _, do_op), bucket, vs', subhisto_ind) -> do -            let bucket' = ImpGen.compileSubExpOfType int32 bucket-                dest_w' = ImpGen.compileSubExpOfType int32 dest_w+            let bucket' = toExp' int32 $ kernelResultSubExp bucket+                dest_w' = toExp' int32 dest_w                 bucket_in_bounds = 0 .<=. bucket' .&&. bucket' .<. dest_w'                 bucket_is = map (`Imp.var` int32) (init space_is) ++                             [Imp.var subhisto_ind int32, bucket']                 vs_params = takeLast (length vs') $ lambdaParams lam              sWhen bucket_in_bounds $ do-              ImpGen.dLParams $ lambdaParams lam-              vectorLoops [] (shapeDims shape) $ \is -> do+              dLParams $ lambdaParams lam+              sLoopNest shape $ \is -> do+                forM_ (zip vs_params vs') $ \(p, res) ->+                  copyDWIM (paramName p) [] (kernelResultSubExp res) is+                do_op (bucket_is ++ is)++prepareIntermediateArraysLocal :: KernelSpace -> KernelConstants+                               -> VName -> [SegGenRedSlug]+                               -> CallKernelGen+                                  [([VName],+                                    InKernelGen ([VName],+                                                 [Imp.Exp] -> ImpM InKernel Imp.KernelOp ()))]+prepareIntermediateArraysLocal space constants num_subhistos_per_group =+  fmap snd . mapAccumLM onOp Nothing+  where+    onOp l (SegGenRedSlug op num_subhistos subhisto_info) = do++      num_subhistos <--+        toExp' int32 (spaceNumGroups space)++      emit $ Imp.DebugPrint "Number of subhistograms in global memory" $+        Just (int32, Imp.vi32 num_subhistos)++      -- Some trickery is afoot here because we need to construct a+      -- Locking structure in the CallKernelGen monad, but the actual+      -- initialisation of the locks array must happen on the device.+      -- Also, we want only one locks array, no matter how many+      -- operators need locking.+      (l', mk_op) <-+        case (l, atomicUpdateLocking $ genReduceOp op) of+          (_, Left f) -> return (l, return f)+          (Just l', Right f) -> return (l, return $ f l')+          (Nothing, Right f) -> do+            locks <- newVName "locks"+            num_locks <- toExp $ spaceGroupSize space++            let dims = map (toExp' int32) $+                       Var num_subhistos_per_group :+                       shapeDims (genReduceShape op) +++                       [genReduceWidth op]+                l' = Locking locks 0 1 0 ((`rem` num_locks) . flattenIndex dims)+                locks_t = localMemLockArray space++                mk_op = do+                  locks_mem <- sAlloc "locks_mem" (typeSize locks_t) $ Space "local"+                  dArray locks int32 (arrayShape locks_t) $+                    ArrayIn locks_mem $ IxFun.iota $+                    map (primExpFromSubExp int32) $ arrayDims locks_t++                  sComment "All locks start out unlocked" $+                    copyDWIM locks [kernelLocalThreadId constants] (intConst Int32 0) []++                  return $ f l'++            return (Just l', mk_op)++      -- Initialise local-memory sub-histograms.  These are+      -- represented as two-dimensional arrays.+      let init_local_subhistos = do+            local_subhistos <-+              forM (genReduceDest op) $ \dest -> do+                dest_t <- lookupType dest++                let sub_local_shape =+                      Shape [Var num_subhistos_per_group] <> arrayShape dest_t+                sAllocArray "subhistogram_local"+                  (elemType dest_t) sub_local_shape (Space "local")++            do_op <- mk_op++            return (local_subhistos, do_op (Space "local") local_subhistos)++      -- Initialise global-memory sub-histograms.+      glob_subhistos <- forM subhisto_info $ \info -> do+        subhistosAlloc info+        return $ subhistosArray info++      return (l', (glob_subhistos, init_local_subhistos))++genRedKernelLocal :: VName+                  -> [PatElem ExplicitMemory]+                  -> KernelSpace+                  -> [SegGenRedSlug]+                  -> KernelBody InKernel+                  -> CallKernelGen ()+genRedKernelLocal num_subhistos_per_group_var map_pes space slugs kbody = do+  (base_constants, init_constants) <- kernelInitialisationSetSpace space $ return ()+  let (space_is, space_sizes) = unzip $ spaceDimensions space+      segment_dims = init space_sizes+      num_segments = length segment_dims+      constants = base_constants { kernelThreadActive = true }+      space_sizes_64 = map (i32Toi64 . toExp' int32) space_sizes+      total_w_64 = product space_sizes_64+      num_subhistos_per_group = Imp.var num_subhistos_per_group_var int32++  emit $ Imp.DebugPrint "Number of local subhistograms per group" $ Just (int32, num_subhistos_per_group)++  init_histograms <- prepareIntermediateArraysLocal space constants num_subhistos_per_group_var slugs++  elems_per_thread_64 <- dPrimV "elems_per_thread_64" $+                         total_w_64 `quotRoundingUp`+                         ConvOpExp (SExt Int32 Int64) (kernelNumThreads constants)++  sKernel constants "seggenred_local" $ allThreads constants $ do+    init_constants++    histograms <- forM init_histograms $+                  \(glob_subhistos, init_local_subhistos) -> do+      (local_subhistos, do_op) <- init_local_subhistos+      return (zip glob_subhistos local_subhistos, do_op)++    -- Find index of local subhistograms updated by this thread.  We+    -- try to ensure, as much as possible, that threads in the same+    -- warp use different subhistograms, to avoid conflicts.+    thread_local_subhisto_i <-+      fmap (`Imp.var` int32) $ dPrimV "thread_local_subhisto_i" $+      kernelLocalThreadId constants `rem` num_subhistos_per_group++    let (red_res, map_res) = splitFromEnd (length map_pes) $+                             map kernelResultSubExp $ kernelBodyResult kbody+        (buckets, vs) = splitAt (length slugs) red_res+        perOp = chunks $ map (length . genReduceDest . slugOp) slugs++    let onSlugs f = forM_ (zip slugs histograms) $ \(slug, (dests, _)) -> do+          let histo_dims =+                map (toExp' int32) $+                segment_dims +++                genReduceWidth (slugOp slug) : shapeDims (genReduceShape (slugOp slug))+          histo_size <- fmap (`Imp.var` int32) $ dPrimV "histo_size" $+                        product histo_dims+          f slug dests histo_dims histo_size++    let onAllHistograms f =+          onSlugs $ \slug dests histo_dims histo_size -> do+            let group_hists_size = num_subhistos_per_group * histo_size+            init_per_thread <- dPrimV "init_per_thread" $+                               group_hists_size `quotRoundingUp` kernelGroupSize constants++            forM_ (zip dests (genReduceNeutral $ slugOp slug)) $ \((dest_global, dest_local), ne) -> do+              i <- newVName "local_i"+              sFor i Int32 (Imp.var init_per_thread int32) $ do+                j <- fmap (`Imp.var` int32) $ dPrimV "j" $+                     Imp.var i int32 * kernelGroupSize constants ++                     kernelLocalThreadId constants+                j_offset <- fmap (`Imp.var` int32) $ dPrimV "j_offset" $+                            num_subhistos_per_group *+                            histo_size *+                            kernelGroupId constants + j++                local_subhisto_i <- dPrimV "local_subhisto_i" $ j `quot` histo_size+                let bucket_is = unflattenIndex histo_dims $ j `rem` histo_size+                global_subhisto_i <- dPrimV "global_subhisto_i" $ j_offset `quot` histo_size++                sWhen (j .<. group_hists_size) $+                  f dest_local dest_global (slugOp slug) ne+                  (Imp.var local_subhisto_i int32) (Imp.var global_subhisto_i int32)+                  bucket_is++    sComment "initialize histograms in local memory" $+      onAllHistograms $ \dest_local dest_global op ne local_subhisto_i global_subhisto_i bucket_is ->+      sComment "First subhistogram is initialised from global memory; others with neutral element." $ do+      let global_is = take num_segments bucket_is +++                      [0] ++ drop num_segments bucket_is+          local_is = local_subhisto_i : bucket_is+      sIf (global_subhisto_i .==. 0)+        (copyDWIM dest_local local_is (Var dest_global) global_is)+        (sLoopNest (genReduceShape op) $ \is ->+            copyDWIM dest_local (local_is++is) ne [])++    sOp Imp.LocalBarrier++    flat_idx <- newVName "flat_idx"+    sFor flat_idx Int64 (Imp.var elems_per_thread_64 int64) $ do+      -- Compute the offset into the input and output.  To this a+      -- thread can add its local ID to figure out which element it is+      -- responsible for.  The calculation is done with 64-bit+      -- integers to avoid overflow, but the final segment indexes are+      -- 32 bit.+      offset <- dPrimV "offset" $+                (i32Toi64 (kernelGroupId constants) *+                 (Imp.var elems_per_thread_64 int64 *+                  i32Toi64 (kernelGroupSize constants)))+                + (Imp.var flat_idx int64 * i32Toi64 (kernelGroupSize constants))++      j <- dPrimV "j" $ Imp.var offset int64 + i32Toi64 (kernelLocalThreadId constants)++      -- Construct segment indices.+      zipWithM_ dPrimV_ space_is $+        map (ConvOpExp (SExt Int64 Int32)) . unflattenIndex space_sizes_64 $ Imp.var j int64++      -- We execute the bucket function once and update each histogram serially.+      -- We apply the bucket function if j=offset+ltid is less than+      -- num_elements.  This also involves writing to the mapout+      -- arrays.+      let input_in_bounds = Imp.var j int32 .<. total_w_64++      sWhen input_in_bounds $ compileStms mempty (kernelBodyStms kbody) $ do++        sComment "save map-out results" $+          forM_ (zip map_pes map_res) $ \(pe, se) ->+          copyDWIM (patElemName pe)+          (map ((`Imp.var` int32) . fst) $ kernelDimensions constants) se []++        forM_ (zip4 (map slugOp slugs) histograms buckets (perOp vs)) $+          \(GenReduceOp dest_w _ _ shape lam,+            (_, do_op), bucket, vs') -> do++            let bucket' = toExp' int32 bucket+                dest_w' = toExp' int32 dest_w+                bucket_in_bounds = 0 .<=. bucket' .&&. bucket' .<. dest_w'+                bucket_is = thread_local_subhisto_i :+                            map (`Imp.var` int32) (init space_is) ++ [bucket']+                vs_params = takeLast (length vs') $ lambdaParams lam++            sComment "perform atomic updates" $+              sWhen bucket_in_bounds $ do+              dLParams $ lambdaParams lam+              sLoopNest shape $ \is -> do                 forM_ (zip vs_params vs') $ \(p, v) ->-                  ImpGen.copyDWIM (paramName p) [] v is+                  copyDWIM (paramName p) [] v is                 do_op (bucket_is ++ is) -  let histogramInfo (num_histos, dests, _) = (num_histos, dests)-  return $ map histogramInfo histograms+    sOp Imp.LocalBarrier+    sOp Imp.GlobalBarrier -vectorLoops :: [Imp.Exp] -> [SubExp]-            -> ([Imp.Exp] -> ImpGen.ImpM lore op ())-            -> ImpGen.ImpM lore op ()-vectorLoops is [] f = f $ reverse is-vectorLoops is (d:ds) f = do-  i <- newVName "i"-  d' <- ImpGen.compileSubExp d-  ImpGen.sFor i Int32 d' $ vectorLoops (Imp.var i int32:is) ds f+    sComment "Compact the multiple local memory subhistograms to a single subhistogram result" $+      onSlugs $ \slug dests histo_dims histo_size -> do+      bins_per_thread <- fmap (`Imp.var` int32) $ dPrimV "init_per_thread" $+                         histo_size `quotRoundingUp` kernelGroupSize constants +      i <- newVName "local_i"+      sFor i Int32 bins_per_thread $ do+        j <- fmap (`Imp.var` int32) $ dPrimV "j" $+             Imp.var i int32 * kernelGroupSize constants ++             kernelLocalThreadId constants+        sWhen (j .<. histo_size) $ do+          -- We are responsible for compacting the flat bin 'j', which+          -- we immediately unflatten.+          let bucket_is = unflattenIndex histo_dims j+          dLParams $ lambdaParams $ genReduceOp $ slugOp slug+          let (xparams, yparams) = splitAt (length local_dests) $+                                   lambdaParams $ genReduceOp $ slugOp slug+              local_dests = map snd dests++          sComment "Read values from subhistogram 0." $+            forM_ (zip xparams local_dests) $ \(xp, subhisto) ->+            copyDWIM+            (paramName xp) []+            (Var subhisto) (0:bucket_is)++          sComment "Accumulate based on values in other subhistograms." $ do+            subhisto_id <- newVName "subhisto_id"+            sFor subhisto_id Int32 (num_subhistos_per_group - 1) $ do+              forM_ (zip yparams local_dests) $ \(yp, subhisto) ->+                copyDWIM+                (paramName yp) []+                (Var subhisto) (Imp.var subhisto_id int32 + 1 : bucket_is)+              compileBody' xparams $ lambdaBody $ genReduceOp $ slugOp slug++          sComment "Put values back in subhistogram 0." $+            forM_ (zip xparams local_dests) $ \(xp, subhisto) ->+              copyDWIM+              subhisto (0:bucket_is)+              (Var $ paramName xp) []++    sComment "Copy the first local histogram to global memory." $+      onSlugs $ \_slug dests histo_dims histo_size -> do+      write_per_thread <- dPrimV "write_per_thread" $+                          histo_size `quotRoundingUp` kernelGroupSize constants++      forM_ dests $ \(dest_global, dest_local) -> do+        i <- newVName "local_i"+        sFor i Int32 (Imp.var write_per_thread int32) $ do+          j <- fmap (`Imp.var` int32) $ dPrimV "j" $+               Imp.var i int32 * kernelGroupSize constants ++               kernelLocalThreadId constants++          sWhen (j .<. histo_size) $ do+            let bucket_is = unflattenIndex histo_dims $ j `rem` histo_size+                global_is = take num_segments bucket_is +++                            [kernelGroupId constants] +++                            drop num_segments bucket_is+                local_is = 0 : bucket_is+            copyDWIM dest_global global_is (Var dest_local) local_is++-- Most of this function is not the histogram part itself, but rather+-- figuring out whether to use a local or global memory strategy, as+-- well as collapsing the subhistograms produced (which are always in+-- global memory, but their number may vary). compileSegGenRed :: Pattern ExplicitMemory                  -> KernelSpace                  -> [GenReduceOp InKernel]-                 -> Body InKernel+                 -> KernelBody InKernel                  -> CallKernelGen ()-compileSegGenRed (Pattern _ pes) genred_space ops body = do+compileSegGenRed (Pattern _ pes) space ops kbody = do   let num_red_res = length ops + sum (map (length . genReduceNeutral) ops)       (all_red_pes, map_pes) = splitAt num_red_res pes -  infos <- genRedKernel map_pes genred_space ops body-  let pes_per_op = chunks (map (length . genReduceDest) ops) all_red_pes+  let t = 8 * 4+  g <- toExp $ spaceGroupSize space+  lmax <- dPrim "lmax" int32+  sOp $ Imp.GetSizeMax lmax Imp.SizeLocalMemory -  forM_ (zip3 infos pes_per_op ops) $ \((num_histos, subhistos), red_pes, op) -> do-    let unitHistoCase =-          -- This is OK because the memory blocks are at least as-          -- large as the ones we are supposed to use for the result.-          forM_ (zip red_pes subhistos) $ \(pe, subhisto) -> do-            pe_mem <- ImpGen.memLocationName . ImpGen.entryArrayLocation <$>-                      ImpGen.lookupArray (patElemName pe)-            subhisto_mem <- ImpGen.memLocationName . ImpGen.entryArrayLocation <$>-                            ImpGen.lookupArray subhisto-            ImpGen.emit $ Imp.SetMem pe_mem subhisto_mem $ Space "device"+  (op_hs, slugs) <- unzip <$> mapM (computeHistoUsage space) ops+  h <- fmap (`Imp.var` int32) $+       dPrimV "h" $ Imp.innerExp $ sum op_hs+  coop <- fmap (`Imp.var` int32) $+          dPrimV "coop" $ h `quotRoundingUp` t -    sIf (Imp.var num_histos int32 .==. 1) unitHistoCase $ do-      -- For the segmented reduction, we keep the segment dimensions-      -- unchanged.  To this, we add two dimensions: one over the number-      -- of buckets, and one over the number of subhistograms.  This-      -- inner dimension is the one that is collapsed in the reduction.-      let segment_dims = init $ spaceDimensions genred_space-          num_buckets = genReduceWidth op+  -- Check for emptyness to avoid division-by-zero.+  sUnless (h .==. 0) $ do+    lh <- dPrimV "lh" $ (g * t) `quotRoundingUp` h -      bucket_id <- newVName "bucket_id"-      subhistogram_id <- newVName "subhistogram_id"-      vector_ids <- mapM (const $ newVName "vector_id") $-                    shapeDims $ genReduceShape op-      gtid <- newVName $ baseString $ spaceGlobalId genred_space-      let lam = genReduceOp op-          segred_space =-            genred_space-            { spaceStructure =-                FlatThreadSpace $-                segment_dims ++-                [(bucket_id, num_buckets)] ++-                zip vector_ids (shapeDims $ genReduceShape op) ++-                [(subhistogram_id, Var num_histos)]-            , spaceGlobalId = gtid-            }+    emit $ Imp.DebugPrint "\n# SegGenRed" Nothing+    emit $ Imp.DebugPrint "Cooperation level" $ Just (int32, coop)+    emit $ Imp.DebugPrint "Memory per set of subhistograms" $ Just (int32, h)+    emit $ Imp.DebugPrint "Desired group size" $ Just (int32, g) -      compileSegRed' (Pattern [] red_pes) segred_space-        Commutative lam (genReduceNeutral op) $ \red_dests _ ->-        forM_ (zip red_dests subhistos) $ \((d, is), subhisto) ->-          ImpGen.copyDWIM d is (Var subhisto) $ map (`Imp.var` int32) $-          map fst segment_dims ++ [subhistogram_id, bucket_id] ++ vector_ids+    sIf (Imp.innerExp (localMemLockUsage space slugs) + h * Imp.vi32 lh .<=. Imp.vi32 lmax+         .&&. coop .<=. g)+      (genRedKernelLocal lh map_pes space slugs kbody)+      (genRedKernelGlobal map_pes space slugs kbody)++    let pes_per_op = chunks (map (length . genReduceDest) ops) all_red_pes++    forM_ (zip3 slugs pes_per_op ops) $ \(slug, red_pes, op) -> do+      let num_histos = slugNumSubhistos slug+          subhistos = map subhistosArray $ slugSubhistos slug++      let unitHistoCase =+            -- This is OK because the memory blocks are at least as+            -- large as the ones we are supposed to use for the result.+            forM_ (zip red_pes subhistos) $ \(pe, subhisto) -> do+              pe_mem <- memLocationName . entryArrayLocation <$>+                        lookupArray (patElemName pe)+              subhisto_mem <- memLocationName . entryArrayLocation <$>+                              lookupArray subhisto+              emit $ Imp.SetMem pe_mem subhisto_mem $ Space "device"++      sIf (Imp.var num_histos int32 .==. 1) unitHistoCase $ do+        -- For the segmented reduction, we keep the segment dimensions+        -- unchanged.  To this, we add two dimensions: one over the number+        -- of buckets, and one over the number of subhistograms.  This+        -- inner dimension is the one that is collapsed in the reduction.+        let num_buckets = genReduceWidth op++        bucket_id <- newVName "bucket_id"+        subhistogram_id <- newVName "subhistogram_id"+        vector_ids <- mapM (const $ newVName "vector_id") $+                      shapeDims $ genReduceShape op+        gtid <- newVName $ baseString $ spaceGlobalId space+        let segred_space =+              space { spaceStructure =+                        FlatThreadSpace $+                        segment_dims +++                        [(bucket_id, num_buckets)] +++                        zip vector_ids (shapeDims $ genReduceShape op) +++                        [(subhistogram_id, Var num_histos)]+                    , spaceGlobalId = gtid+                    }++        let segred_op = SegRedOp Commutative (genReduceOp op) (genReduceNeutral op) mempty+        compileSegRed' (Pattern [] red_pes) segred_space [segred_op] $ \_ red_cont ->+          red_cont $ flip map subhistos $ \subhisto ->+            (Var subhisto, map (`Imp.var` int32) $+              map fst segment_dims ++ [subhistogram_id, bucket_id] ++ vector_ids)++  where segment_dims = init $ spaceDimensions space

@@ -0,0 +1,26 @@+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE FlexibleContexts #-}+module Futhark.CodeGen.ImpGen.Kernels.SegMap+  ( compileSegMap ) where++import Control.Monad.Except++import Prelude hiding (quot, rem)++import Futhark.Representation.ExplicitMemory+import Futhark.CodeGen.ImpGen.Kernels.Base++-- | Compile 'SegMap' instance to host-level code with calls to+-- various kernels.+compileSegMap :: Pattern ExplicitMemory+              -> KernelSpace+              -> KernelBody InKernel+              -> CallKernelGen ()+compileSegMap pat space kbody = do+  (constants, init_constants) <- kernelInitialisation space++  sKernel constants "segmap" $ do+    init_constants+    compileKernelStms constants (kernelBodyStms kbody) $+      zipWithM_ (compileKernelResult constants) (patternElements pat) $+      kernelBodyResult kbody

@@ -53,139 +53,169 @@  import Prelude hiding (quot, rem) +import Futhark.Error import Futhark.MonadFreshNames import Futhark.Transform.Rename import Futhark.Representation.ExplicitMemory import qualified Futhark.CodeGen.ImpCode.Kernels as Imp-import qualified Futhark.CodeGen.ImpGen as ImpGen-import Futhark.CodeGen.ImpGen ((<--),-                               sFor, sComment, sIf, sWhen,-                               sOp,-                               dPrim, dPrimV)+import Futhark.CodeGen.ImpGen import Futhark.CodeGen.ImpGen.Kernels.Base import qualified Futhark.Representation.ExplicitMemory.IndexFunction as IxFun+import Futhark.Util (chunks) import Futhark.Util.IntegralExp (quotRoundingUp, quot, rem) --- | 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 :: KernelConstants-                 -> Imp.Exp-                 -> (Imp.Exp -> ImpGen.ImpM lore op ())-                 -> ImpGen.ImpM lore op ()-virtualiseGroups constants required_groups m = do-  let group_id = kernelGroupId constants-      iterations = (required_groups - group_id) `quotRoundingUp` kernelNumGroups constants-  i <- newVName "i"-  sFor i Int32 iterations $ m $ group_id + Imp.var i int32 * kernelNumGroups constants+-- | The maximum number of operators we support in a single SegRed.+-- This limit arises out of the static allocation of counters.+maxNumOps :: Int32+maxNumOps = 10 +type DoSegBody = (KernelConstants -> ([(SubExp, [Imp.Exp])] -> InKernelGen ()) -> InKernelGen ())+ -- | Compile 'SegRed' instance to host-level code with calls to -- various kernels. compileSegRed :: Pattern ExplicitMemory               -> KernelSpace-              -> Commutativity -> Lambda InKernel -> [SubExp]-              -> Body InKernel+              -> [SegRedOp InKernel]+              -> KernelBody InKernel               -> CallKernelGen ()-compileSegRed pat space comm red_op nes body =-  compileSegRed' pat space comm red_op nes $ \red_dests map_dests ->-  ImpGen.compileStms mempty (stmsToList $ bodyStms body) $ do-  let (red_res, map_res) = splitAt (length nes) $ bodyResult body-  sComment "save results to be reduced" $-    forM_ (zip red_dests red_res) $ \((d,is), se) -> ImpGen.copyDWIM d is se []-  sComment "save map-out results" $-    forM_ (zip map_dests map_res) $ \((d,is), se) -> ImpGen.copyDWIM d is se []+compileSegRed pat space reds body =+  compileSegRed' pat space reds $ \constants red_cont ->+  compileKernelStms constants (kernelBodyStms body) $ do+  let (red_res, map_res) = splitAt (segRedResults reds) $ kernelBodyResult body +  sComment "save map-out results" $ do+    let map_arrs = drop (segRedResults reds) $ patternElements pat+    zipWithM_ (compileKernelResult constants) map_arrs map_res++  red_cont $ zip (map kernelResultSubExp red_res) $ repeat []+ -- | Like 'compileSegRed', but where the body is a monadic action. compileSegRed' :: Pattern ExplicitMemory                -> KernelSpace-               -> Commutativity -> Lambda InKernel -> [SubExp]-               -> ([(VName, [Imp.Exp])] -> [(VName, [Imp.Exp])] -> InKernelGen ())+               -> [SegRedOp InKernel]+               -> DoSegBody                -> CallKernelGen ()-compileSegRed' pat space comm red_op nes body+compileSegRed' pat space reds body+  | genericLength reds > maxNumOps =+      compilerLimitationS $+      "compileSegRed': at most " ++ show maxNumOps ++ " reduction operators are supported."   | [(_, Constant (IntValue (Int32Value 1))), _] <- spaceDimensions space =-      nonsegmentedReduction pat space comm red_op nes body+      nonsegmentedReduction pat space reds body   | otherwise = do       segment_size <--        ImpGen.compileSubExp $ last $ map snd $ spaceDimensions space-      group_size <- ImpGen.compileSubExp $ spaceGroupSize space+        toExp $ last $ map snd $ spaceDimensions space+      group_size <- toExp $ spaceGroupSize space       let use_small_segments = segment_size * 2 .<. group_size       sIf use_small_segments-        (smallSegmentsReduction pat space red_op nes body)-        (largeSegmentsReduction pat space comm red_op nes body)--nonsegmentedReduction :: Pattern ExplicitMemory-                      -> KernelSpace-                      -> Commutativity -> Lambda InKernel -> [SubExp]-                      -> ([(VName, [Imp.Exp])] -> [(VName, [Imp.Exp])] -> InKernelGen ())-                      -> CallKernelGen ()-nonsegmentedReduction segred_pat space comm red_op nes body = do-  (base_constants, init_constants) <- kernelInitialisationSetSpace space $ return ()-  let constants = base_constants { kernelThreadActive = true }-      global_tid = kernelGlobalThreadId constants-      (_, w) = last $ spaceDimensions space+        (smallSegmentsReduction pat space reds body)+        (largeSegmentsReduction pat space reds body) +-- | Prepare intermediate arrays for the reduction.  Prim-typed+-- arguments go in local memory (so we need to do the allocation of+-- those arrays inside the kernel), while array-typed arguments go in+-- global memory.  Allocations for the former have already been+-- performed.  This policy is baked into how the allocations are done+-- in ExplicitAllocations.+intermediateArrays :: KernelSpace -> SegRedOp InKernel -> InKernelGen [VName]+intermediateArrays space (SegRedOp _ red_op nes _) = do   let red_op_params = lambdaParams red_op       (red_acc_params, _) = splitAt (length nes) red_op_params-  red_arrs <- forM red_acc_params $ \p ->+  forM red_acc_params $ \p ->     case paramAttr p of       MemArray pt shape _ (ArrayIn mem _) -> do         let shape' = Shape [spaceNumThreads space] <> shape-        ImpGen.sArray "red_arr" pt shape' $+        sArray "red_arr" pt shape' $           ArrayIn mem $ IxFun.iota $ map (primExpFromSubExp int32) $ shapeDims shape'       _ -> do         let pt = elemType $ paramType p             shape = Shape [spaceGroupSize space]-        ImpGen.sAllocArray "red_arr" pt shape $ Space "local"--  counter <--    ImpGen.sStaticArray "counter" (Space "device") int32 $-    Imp.ArrayValues $ replicate 1 $ IntValue $ Int32Value 0+        sAllocArray "red_arr" pt shape $ Space "local" -  group_res_arrs <- forM (lambdaReturnType red_op) $ \t -> do+-- | Arrays for storing group results.+--+-- The group-result arrays have an extra dimension (of size groupsize)+-- because they are also used for keeping vectorised accumulators for+-- first-stage reduction, if necessary.  When actually storing group+-- results, the first index is set to 0.+groupResultArrays :: SubExp -> SubExp -> [SegRedOp InKernel] -> CallKernelGen [[VName]]+groupResultArrays virt_num_groups group_size reds =+  forM reds $ \(SegRedOp _ lam _ shape) ->+    forM (lambdaReturnType lam) $ \t -> do     let pt = elemType t-        shape = Shape [spaceNumGroups space] <> arrayShape t-    ImpGen.sAllocArray "group_res_arr" pt shape $ Space "device"+        full_shape = Shape [group_size, virt_num_groups] <> shape <> arrayShape t+        -- Move the groupsize dimension last to ensure coalesced+        -- memory access.+        perm = [1..shapeRank full_shape-1] ++ [0]+    sAllocArrayPerm "group_res_arr" pt full_shape (Space "device") perm -  sync_arr <- ImpGen.sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local"+nonsegmentedReduction :: Pattern ExplicitMemory+                      -> KernelSpace+                      -> [SegRedOp InKernel]+                      -> DoSegBody+                      -> CallKernelGen ()+nonsegmentedReduction segred_pat space reds body = do+  (base_constants, init_constants) <- kernelInitialisationSetSpace space $ return ()+  let constants = base_constants { kernelThreadActive = true }+      global_tid = kernelGlobalThreadId constants+      (_, w) = last $ spaceDimensions space +  counter <-+    sStaticArray "counter" (Space "device") int32 $+    Imp.ArrayValues $ replicate (fromIntegral maxNumOps) $ IntValue $ Int32Value 0++  -- The group-result arrays have an extra dimension (of size+  -- groupsize) because they are also used for keeping vectorised+  -- accumulators for first-stage reduction, if necessary.  When+  -- actually storing group results, the first index is set to 0.+  reds_group_res_arrs <- groupResultArrays (spaceNumGroups space) (spaceGroupSize space) reds+   num_threads <- dPrimV "num_threads" $ kernelNumThreads constants    sKernel constants "segred_nonseg" $ allThreads constants $ do     init_constants +    sync_arr <- sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local"+    reds_arrs <- mapM (intermediateArrays space) reds+     -- Since this is the nonsegmented case, all outer segment IDs must     -- necessarily be 0.     let gtids = map fst $ spaceDimensions space     forM_ (init gtids) $ \v ->       v <-- 0 -    num_elements <- Imp.elements <$> ImpGen.compileSubExp w+    num_elements <- Imp.elements <$> toExp w     let elems_per_thread = num_elements `quotRoundingUp` Imp.elements (kernelNumThreads constants) -    (group_result_params, red_op_renamed) <--      reductionStageOne constants segred_pat num_elements+    slugs <- mapM (segRedOpSlug (kernelLocalThreadId constants) (kernelGroupId constants)) $+             zip3 reds reds_arrs reds_group_res_arrs+    reds_op_renamed <-+      reductionStageOne constants num_elements       global_tid elems_per_thread num_threads-      comm red_op nes red_arrs body+      slugs body -    reductionStageTwo constants segred_pat 0 [0] 0-      (kernelNumGroups constants) group_result_params red_acc_params red_op_renamed nes-      1 counter sync_arr group_res_arrs red_arrs+    let segred_pes = chunks (map (length . segRedNeutral) reds) $+                     patternElements segred_pat+    forM_ (zip7 reds reds_arrs reds_group_res_arrs segred_pes+           slugs reds_op_renamed [0..]) $+      \(SegRedOp _ red_op nes _,+        red_arrs, group_res_arrs, pes, slug, red_op_renamed, i) -> do+      let red_acc_params = take (length nes) $ lambdaParams red_op+      reductionStageTwo constants pes (kernelGroupId constants) 0 [0] 0+        (kernelNumGroups constants) slug red_acc_params red_op_renamed nes+        1 counter (ValueExp $ IntValue $ Int32Value i)+        sync_arr group_res_arrs red_arrs  smallSegmentsReduction :: Pattern ExplicitMemory                        -> KernelSpace-                       -> Lambda InKernel -> [SubExp]-                       -> ([(VName, [Imp.Exp])] -> [(VName, [Imp.Exp])] -> InKernelGen ())+                       -> [SegRedOp InKernel]+                       -> DoSegBody                        -> CallKernelGen ()-smallSegmentsReduction (Pattern _ segred_pes) space red_op nes body = do+smallSegmentsReduction (Pattern _ segred_pes) space reds body = do   (base_constants, init_constants) <- kernelInitialisationSetSpace space $ return ()   let constants = base_constants { kernelThreadActive = true }    let (gtids, dims) = unzip $ spaceDimensions space-  dims' <- mapM ImpGen.compileSubExp dims+  dims' <- mapM toExp dims    let segment_size = last dims'   -- Careful to avoid division by zero now.@@ -196,31 +226,22 @@       segments_per_group = kernelGroupSize constants `quot` segment_size_nonzero       required_groups = num_segments `quotRoundingUp` segments_per_group -  let red_op_params = lambdaParams red_op-      (red_acc_params, _red_next_params) = splitAt (length nes) red_op_params-  red_arrs <- forM red_acc_params $ \p ->-    case paramAttr p of-      MemArray pt shape _ (ArrayIn mem _) -> do-        let shape' = Shape [spaceNumThreads space] <> shape-        ImpGen.sArray "red_arr" pt shape' $-          ArrayIn mem $ IxFun.iota $ map (primExpFromSubExp int32) $ shapeDims shape'-      _ -> do-        let pt = elemType $ paramType p-            shape = Shape [spaceGroupSize space]-        ImpGen.sAllocArray "red_arr" pt shape $ Space "local"--  ImpGen.emit $ Imp.DebugPrint "num_segments" int32 num_segments-  ImpGen.emit $ Imp.DebugPrint "segment_size" int32 segment_size-  ImpGen.emit $ Imp.DebugPrint "segments_per_group" int32 segments_per_group-  ImpGen.emit $ Imp.DebugPrint "required_groups" int32 required_groups+  emit $ Imp.DebugPrint "\n# SegRed-small" Nothing+  emit $ Imp.DebugPrint "num_segments" $ Just (int32, num_segments)+  emit $ Imp.DebugPrint "segment_size" $ Just (int32, segment_size)+  emit $ Imp.DebugPrint "segments_per_group" $ Just (int32, segments_per_group)+  emit $ Imp.DebugPrint "required_groups" $ Just (int32, required_groups)    sKernel constants "segred_small" $ allThreads constants $ do     init_constants +    reds_arrs <- mapM (intermediateArrays space) reds+     -- We probably do not have enough actual workgroups to cover the     -- entire iteration space.  Some groups thus have to perform double     -- duty; we put an outer loop to accomplish this.-    virtualiseGroups constants required_groups $ \group_id' -> do+    virtualiseGroups constants required_groups $ \group_id_var' -> do+      let group_id' = Imp.vi32 group_id_var'       -- Compute the 'n' input indices.  The outer 'n-1' correspond to       -- the segment ID, and are computed from the group id.  The inner       -- is computed from the local thread id, and may be out-of-bounds.@@ -231,43 +252,41 @@       zipWithM_ (<--) (init gtids) $ unflattenIndex (init dims') segment_index       last gtids <-- index_within_segment -      let toLocalMemory ses =-            forM_ (zip red_arrs ses) $ \(arr, se) ->-            ImpGen.copyDWIM arr [ltid] se []+      let out_of_bounds =+            forM_ (zip reds reds_arrs) $ \(SegRedOp _ _ nes _, red_arrs) ->+            forM_ (zip red_arrs nes) $ \(arr, ne) ->+            copyDWIM arr [ltid] ne []            in_bounds =-            body (zip red_arrs $ repeat [ltid])-                 (zip (map patElemName $ drop (length nes) segred_pes) $-                  repeat $ map (`Imp.var` int32) gtids)+            body constants $ \red_res ->+            sComment "save results to be reduced" $ do+            let red_dests = zip (concat reds_arrs) $ repeat [ltid]+            forM_ (zip red_dests red_res) $ \((d,d_is), (res, res_is)) ->+              copyDWIM d d_is res res_is        sComment "apply map function if in bounds" $         sIf (segment_size .>. 0 .&&.              isActive (init $ zip gtids dims) .&&.-             ltid .<. segment_size * segments_per_group) in_bounds (toLocalMemory nes)+             ltid .<. segment_size * segments_per_group) in_bounds out_of_bounds        sOp Imp.LocalBarrier -      index_i <- newVName "index_i"-      index_j <- newVName "index_j"       let crossesSegment from to = (to-from) .>. (to `rem` segment_size)-          red_op' = red_op { lambdaParams = Param index_i (MemPrim int32) :-                                            Param index_j (MemPrim int32) :-                                            lambdaParams red_op }-       sWhen (segment_size .>. 0) $         sComment "perform segmented scan to imitate reduction" $-        groupScan constants (Just crossesSegment) (segment_size*segments_per_group) red_op' red_arrs+        forM_ (zip reds reds_arrs) $ \(SegRedOp _ red_op _ _, red_arrs) ->+        groupScan constants (Just crossesSegment) (segment_size*segments_per_group) red_op red_arrs        sOp Imp.LocalBarrier        sComment "save final values of segments" $         sWhen (group_id' * segments_per_group + ltid .<. num_segments .&&.                ltid .<. segments_per_group) $-        forM_ (zip segred_pes red_arrs) $ \(pe, arr) -> do+        forM_ (zip segred_pes (concat reds_arrs)) $ \(pe, arr) -> do         -- Figure out which segment result this thread should write...         let flat_segment_index = group_id' * segments_per_group + ltid             gtids' = unflattenIndex (init dims') flat_segment_index-        ImpGen.copyDWIM (patElemName pe) gtids'+        copyDWIM (patElemName pe) gtids'                         (Var arr) [(ltid+1) * segment_size_nonzero - 1]        -- Finally another barrier, because we will be writing to the@@ -276,58 +295,37 @@  largeSegmentsReduction :: Pattern ExplicitMemory                        -> KernelSpace-                       -> Commutativity -> Lambda InKernel -> [SubExp]-                       -> ([(VName, [Imp.Exp])] -> [(VName, [Imp.Exp])] -> InKernelGen ())+                       -> [SegRedOp InKernel]+                       -> DoSegBody                        -> CallKernelGen ()-largeSegmentsReduction segred_pat space comm red_op nes body = do+largeSegmentsReduction segred_pat space reds body = do   (base_constants, init_constants) <- kernelInitialisationSetSpace space $ return ()   let (gtids, dims) = unzip $ spaceDimensions space-  dims' <- mapM ImpGen.compileSubExp dims+  dims' <- mapM toExp dims   let segment_size = last dims'       num_segments = product $ init dims'    let (groups_per_segment, elems_per_thread) =         groupsPerSegmentAndElementsPerThread segment_size num_segments         (kernelNumGroups base_constants) (kernelGroupSize base_constants)-  num_groups <- dPrimV "num_groups" $+  virt_num_groups <- dPrimV "virt_num_groups" $     groups_per_segment * num_segments -  num_threads <- dPrimV "num_threads" $-    Imp.var num_groups int32 * kernelGroupSize base_constants-   threads_per_segment <- dPrimV "thread_per_segment" $     groups_per_segment * kernelGroupSize base_constants -  let constants = base_constants-                  { kernelThreadActive = true-                  , kernelNumGroups = Imp.var num_groups int32-                  , kernelNumThreads = Imp.var num_threads int32-                  }--  ImpGen.emit $ Imp.DebugPrint "num_segments" int32 num_segments-  ImpGen.emit $ Imp.DebugPrint "segment_size" int32 segment_size-  ImpGen.emit $ Imp.DebugPrint "num_groups" int32 (Imp.var num_groups int32)-  ImpGen.emit $ Imp.DebugPrint "group_size" int32 (kernelGroupSize constants)-  ImpGen.emit $ Imp.DebugPrint "elems_per_thread" int32 $ Imp.innerExp elems_per_thread-  ImpGen.emit $ Imp.DebugPrint "groups_per_segment" int32 groups_per_segment+  let constants = base_constants { kernelThreadActive = true } -  let red_op_params = lambdaParams red_op-      (red_acc_params, _) = splitAt (length nes) red_op_params-  red_arrs <- forM red_acc_params $ \p ->-    case paramAttr p of-      MemArray pt shape _ (ArrayIn mem _) -> do-        let shape' = Shape [Var num_threads] <> shape-        ImpGen.sArray "red_arr" pt shape' $-          ArrayIn mem $ IxFun.iota $ map (primExpFromSubExp int32) $ shapeDims shape'-      _ -> do-        let pt = elemType $ paramType p-            shape = Shape [spaceGroupSize space]-        ImpGen.sAllocArray "red_arr" pt shape $ Space "local"+  emit $ Imp.DebugPrint "\n# SegRed-large" Nothing+  emit $ Imp.DebugPrint "num_segments" $ Just (int32, num_segments)+  emit $ Imp.DebugPrint "segment_size" $ Just (int32, segment_size)+  emit $ Imp.DebugPrint "virt_num_groups" $ Just (int32, Imp.vi32 virt_num_groups)+  emit $ Imp.DebugPrint "num_groups" $ Just (int32, kernelNumGroups constants)+  emit $ Imp.DebugPrint "group_size" $ Just (int32, kernelGroupSize constants)+  emit $ Imp.DebugPrint "elems_per_thread" $ Just (int32, Imp.innerExp elems_per_thread)+  emit $ Imp.DebugPrint "groups_per_segment" $ Just (int32, groups_per_segment) -  group_res_arrs <- forM (lambdaReturnType red_op) $ \t -> do-    let pt = elemType t-        shape = Shape [Var num_groups] <> arrayShape t-    ImpGen.sAllocArray "group_res_arr" pt shape $ Space "device"+  reds_group_res_arrs <- groupResultArrays (Var virt_num_groups) (spaceGroupSize space) reds    -- In principle we should have a counter for every segment.  Since   -- the number of segments is a dynamic quantity, we would have to@@ -339,49 +337,68 @@   -- anywhere?  There are other places in the compiler that will fail   -- if the group count exceeds the maximum group size, which is at   -- most 1024 anyway.-  let num_counters = 1024+  let num_counters = fromIntegral maxNumOps * 1024   counter <--    ImpGen.sStaticArray "counter" (Space "device") int32 $+    sStaticArray "counter" (Space "device") int32 $     Imp.ArrayZeros num_counters -  sync_arr <- ImpGen.sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local"-   sKernel constants "segred_large" $ allThreads constants $ do     init_constants-    let segment_gtids = init gtids-        group_id = kernelGroupId constants-        group_size = kernelGroupSize constants-        flat_segment_id = group_id `quot` groups_per_segment-        local_tid = kernelLocalThreadId constants -        global_tid = kernelGlobalThreadId constants-                     `rem` (group_size * groups_per_segment)-        w = last dims-        first_group_for_segment = flat_segment_id * groups_per_segment+    reds_arrs <- mapM (intermediateArrays space) reds+    sync_arr <- sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local" -    zipWithM_ (<--) segment_gtids $ unflattenIndex (init dims') flat_segment_id-    num_elements <- Imp.elements <$> ImpGen.compileSubExp w+    -- We probably do not have enough actual workgroups to cover the+    -- entire iteration space.  Some groups thus have to perform double+    -- duty; we put an outer loop to accomplish this.+    virtualiseGroups constants (Imp.vi32 virt_num_groups) $ \group_id_var -> do -    (group_result_params, red_op_renamed) <--      reductionStageOne constants segred_pat num_elements-      global_tid elems_per_thread threads_per_segment-      comm red_op nes red_arrs body+      let segment_gtids = init gtids+          group_id = Imp.vi32 group_id_var+          group_size = kernelGroupSize constants+          flat_segment_id = group_id `quot` groups_per_segment+          local_tid = kernelLocalThreadId constants -    let multiple_groups_per_segment =-          reductionStageTwo constants segred_pat-          flat_segment_id (map (`Imp.var` int32) segment_gtids)-          first_group_for_segment groups_per_segment-          group_result_params red_acc_params red_op_renamed-          nes (fromIntegral num_counters) counter sync_arr group_res_arrs red_arrs+          global_tid = (group_id * group_size + local_tid)+                       `rem` (group_size * groups_per_segment)+          w = last dims+          first_group_for_segment = flat_segment_id * groups_per_segment -        one_group_per_segment =-          ImpGen.comment "first thread in group saves final result to memory" $-          sWhen (local_tid .==. 0) $-            forM_ (take (length nes) $ zip (patternNames segred_pat) group_result_params) $ \(v, p) ->-            ImpGen.copyDWIM v (map (`Imp.var` int32) segment_gtids) (Var $ paramName p) []+      zipWithM_ (<--) segment_gtids $ unflattenIndex (init dims') flat_segment_id+      num_elements <- Imp.elements <$> toExp w -    sIf (groups_per_segment .==. 1) one_group_per_segment multiple_groups_per_segment+      slugs <- mapM (segRedOpSlug local_tid group_id) $+               zip3 reds reds_arrs reds_group_res_arrs+      reds_op_renamed <-+        reductionStageOne constants num_elements+        global_tid elems_per_thread threads_per_segment+        slugs body +      let segred_pes = chunks (map (length . segRedNeutral) reds) $+                       patternElements segred_pat++          multiple_groups_per_segment =+            forM_ (zip7 reds reds_arrs reds_group_res_arrs segred_pes+                   slugs reds_op_renamed [0..]) $+            \(SegRedOp _ red_op nes _, red_arrs, group_res_arrs, pes,+              slug, red_op_renamed, i) -> do+              let red_acc_params = take (length nes) $ lambdaParams red_op+              reductionStageTwo constants pes+                group_id flat_segment_id (map (`Imp.var` int32) segment_gtids)+                first_group_for_segment groups_per_segment+                slug red_acc_params red_op_renamed nes+                (fromIntegral num_counters) counter (ValueExp $ IntValue $ Int32Value i)+                sync_arr group_res_arrs red_arrs++          one_group_per_segment =+            comment "first thread in group saves final result to memory" $+            forM_ (zip slugs segred_pes) $ \(slug, pes) ->+            sWhen (local_tid .==. 0) $+              forM_ (zip pes (slugAccs slug)) $ \(v, (acc, acc_is)) ->+              copyDWIM (patElemName v) (map (`Imp.var` int32) segment_gtids) (Var acc) acc_is++      sIf (groups_per_segment .==. 1) one_group_per_segment multiple_groups_per_segment+ -- Careful to avoid division by zero here. groupsPerSegmentAndElementsPerThread :: Imp.Exp -> Imp.Exp -> Imp.Exp -> Imp.Exp                                      -> (Imp.Exp, Imp.Count Imp.Elements)@@ -392,56 +409,107 @@         segment_size `quotRoundingUp` (group_size * groups_per_segment)   in (groups_per_segment, Imp.elements elements_per_thread) -reductionStageOne :: KernelConstants-                  -> Pattern ExplicitMemory-                  -> Imp.Count Imp.Elements-                  -> Imp.Exp-                  -> Imp.Count Imp.Elements-                  -> VName-                  -> Commutativity-                  -> LambdaT InKernel-                  -> [SubExp]-                  -> [VName]-                  -> ([(VName, [Imp.Exp])] -> [(VName, [Imp.Exp])] -> InKernelGen ())-                  -> InKernelGen ([LParam InKernel], Lambda InKernel)-reductionStageOne constants (Pattern _ segred_pes) num_elements global_tid elems_per_thread threads_per_segment comm red_op nes red_arrs body = do+-- | A SegRedOp with auxiliary information.+data SegRedOpSlug =+  SegRedOpSlug+  { slugOp :: SegRedOp InKernel+  , slugArrs :: [VName]+    -- ^ The arrays used for computing the intra-group reduction+    -- (either local or global memory).+  , slugAccs :: [(VName, [Imp.Exp])]+    -- ^ Places to store accumulator in stage 1 reduction.+  } -  let red_op_params = lambdaParams red_op-      (red_acc_params, red_next_params) = splitAt (length nes) red_op_params-      (gtids, _dims) = unzip $ kernelDimensions constants+slugBody :: SegRedOpSlug -> Body InKernel+slugBody = lambdaBody . segRedLambda . slugOp++slugParams :: SegRedOpSlug -> [LParam InKernel]+slugParams = lambdaParams . segRedLambda . slugOp++slugNeutral :: SegRedOpSlug -> [SubExp]+slugNeutral = segRedNeutral . slugOp++slugShape :: SegRedOpSlug -> Shape+slugShape = segRedShape . slugOp++slugsComm :: [SegRedOpSlug] -> Commutativity+slugsComm = mconcat . map (segRedComm . slugOp)++accParams, nextParams :: SegRedOpSlug -> [LParam InKernel]+accParams slug = take (length (slugNeutral slug)) $ slugParams slug+nextParams slug = drop (length (slugNeutral slug)) $ slugParams slug++segRedOpSlug :: Imp.Exp -> Imp.Exp -> (SegRedOp InKernel, [VName], [VName]) -> InKernelGen SegRedOpSlug+segRedOpSlug local_tid group_id (op, group_res_arrs, param_arrs) =+  SegRedOpSlug op group_res_arrs <$>+  zipWithM mkAcc (lambdaParams (segRedLambda op)) param_arrs+  where mkAcc p param_arr+          | Prim t <- paramType p,+            shapeRank (segRedShape op) == 0 = do+              acc <- dPrim (baseString (paramName p) <> "_acc") t+              return (acc, [])+          | otherwise =+              return (param_arr, [local_tid, group_id])++reductionStageZero :: KernelConstants+                   -> Imp.Count Imp.Elements+                   -> Imp.Exp+                   -> Imp.Count Imp.Elements+                   -> VName+                   -> [SegRedOpSlug]+                   -> DoSegBody+                   -> InKernelGen ([Lambda InKernel], InKernelGen ())+reductionStageZero constants num_elements global_tid elems_per_thread threads_per_segment slugs body = do++  let (gtids, _dims) = unzip $ kernelDimensions constants       gtid = last gtids       local_tid = kernelLocalThreadId constants    -- Figure out how many elements this thread should process.   chunk_size <- dPrim "chunk_size" int32-  let ordering = case comm of Commutative -> SplitStrided $ Var threads_per_segment-                              Noncommutative -> SplitContiguous+  let ordering = case slugsComm slugs of+                   Commutative -> SplitStrided $ Var threads_per_segment+                   Noncommutative -> SplitContiguous   computeThreadChunkSize ordering global_tid elems_per_thread num_elements chunk_size -  ImpGen.dScope Nothing $ scopeOfLParams $ lambdaParams red_op+  dScope Nothing $ scopeOfLParams $ concatMap slugParams slugs -  forM_ (zip red_acc_params nes) $ \(p, ne) ->-    ImpGen.copyDWIM (paramName p) [] ne []+  sComment "neutral-initialise the accumulators" $+    forM_ slugs $ \slug ->+    forM_ (zip (slugAccs slug) (slugNeutral slug)) $ \((acc, acc_is), ne) ->+    sLoopNest (slugShape slug) $ \vec_is ->+    copyDWIM acc (acc_is++vec_is) ne [] -  red_op_renamed <- renameLambda red_op+  slugs_op_renamed <- mapM (renameLambda . segRedLambda . slugOp) slugs -  let doTheReduction = do-        ImpGen.comment "to reduce current chunk, first store our result to memory" $-          forM_ (zip red_arrs red_acc_params) $ \(arr, p) ->-          when (primType $ paramType p) $-          ImpGen.copyDWIM arr [local_tid] (Var $ paramName p) []+  let doTheReduction =+        forM_ (zip slugs_op_renamed slugs) $ \(slug_op_renamed, slug) ->+        sLoopNest (slugShape slug) $ \vec_is -> do+          comment "to reduce current chunk, first store our result in memory" $ do+            forM_ (zip (slugParams slug) (slugAccs slug)) $ \(p, (acc, acc_is)) ->+              copyDWIM (paramName p) [] (Var acc) (acc_is++vec_is) -        sOp Imp.LocalBarrier+            forM_ (zip (slugArrs slug) (slugParams slug)) $ \(arr, p) ->+              when (primType $ paramType p) $+              copyDWIM arr [local_tid] (Var $ paramName p) [] -        groupReduce constants (kernelGroupSize constants) red_op_renamed red_arrs+          sOp Imp.LocalBarrier -        sOp Imp.LocalBarrier+          groupReduce constants (kernelGroupSize constants) slug_op_renamed (slugArrs slug) +          sOp Imp.LocalBarrier++          sComment "first thread saves the result in accumulator" $+            sWhen (local_tid .==. 0) $+            forM_ (zip (slugAccs slug) (lambdaParams slug_op_renamed)) $ \((acc, acc_is), p) ->+            copyDWIM acc (acc_is++vec_is) (Var $ paramName p) []+   i <- newVName "i"   -- If this is a non-commutative reduction, each thread must run the   -- loop the same number of iterations, because we will be performing   -- a group-wide reduction in there.-  let (bound, check_bounds) =+  let comm = slugsComm slugs+      (bound, check_bounds) =         case comm of           Commutative -> (Imp.var chunk_size int32, id)           Noncommutative -> (Imp.innerExp elems_per_thread,@@ -459,78 +527,100 @@              (index_in_segment * Imp.innerExp elems_per_thread + Imp.var i int32) *              kernelGroupSize constants -    let red_dests = zip (map paramName red_next_params) $ repeat []-        map_dests = zip (map patElemName $ drop (length nes) segred_pes) $-                    repeat $ map (`Imp.var` int32) gtids+    check_bounds $ sComment "apply map function" $+      body constants $ \all_red_res -> do -    check_bounds $ sComment "apply map function" $ do-      body red_dests map_dests+      let slugs_res = chunks (map (length . slugNeutral) slugs) all_red_res -      sComment "apply reduction operator" $-        ImpGen.compileBody' red_acc_params $ lambdaBody red_op+      forM_ (zip slugs slugs_res) $ \(slug, red_res) ->+        sLoopNest (slugShape slug) $ \vec_is -> do+        sComment "load accumulator" $+          forM_ (zip (accParams slug) (slugAccs slug)) $ \(p, (acc, acc_is)) ->+          copyDWIM (paramName p) [] (Var acc) (acc_is ++ vec_is)+        sComment "load new values" $+          forM_ (zip (nextParams slug) red_res) $ \(p, (res, res_is)) ->+          copyDWIM (paramName p) [] res (res_is ++ vec_is)+        sComment "apply reduction operator" $+          compileStms mempty (bodyStms $ slugBody slug) $+          sComment "store in accumulator" $+          forM_ (zip+                  (slugAccs slug)+                  (bodyResult $ slugBody slug)) $ \((acc, acc_is), se) ->+          copyDWIM acc (acc_is ++ vec_is) se []      case comm of       Noncommutative -> do         doTheReduction-        sComment "first thread takes carry-out; others neutral element" $ do-          let carry_out =-                forM_ (zip red_acc_params $ lambdaParams red_op_renamed) $ \(p_to, p_from) ->-                ImpGen.copyDWIM (paramName p_to) [] (Var $ paramName p_from) []-              reset_to_neutral =-                forM_ (zip red_acc_params nes) $ \(p, ne) ->-                ImpGen.copyDWIM (paramName p) [] ne []-          sIf (local_tid .==. 0) carry_out reset_to_neutral+        sComment "first thread keeps accumulator; others reset to neutral element" $ do+          let reset_to_neutral =+                forM_ slugs $ \slug ->+                forM_ (zip (slugAccs slug) (slugNeutral slug)) $ \((acc, acc_is), ne) ->+                sLoopNest (slugShape slug) $ \vec_is ->+                copyDWIM acc (acc_is++vec_is) ne []+          sUnless (local_tid .==. 0) reset_to_neutral       _ -> return () -  group_result_params <--    case comm of Noncommutative -> return red_acc_params-                 _              -> do doTheReduction-                                      return $ lambdaParams red_op_renamed+  return (slugs_op_renamed, doTheReduction) -  return (group_result_params, red_op_renamed)+reductionStageOne :: KernelConstants+                  -> Imp.Count Imp.Elements+                  -> Imp.Exp+                  -> Imp.Count Imp.Elements+                  -> VName+                  -> [SegRedOpSlug]+                  -> DoSegBody+                  -> InKernelGen [Lambda InKernel]+reductionStageOne constants num_elements global_tid elems_per_thread threads_per_segment slugs body = do+  (slugs_op_renamed, doTheReduction) <-+    reductionStageZero constants num_elements global_tid elems_per_thread threads_per_segment slugs body +  case slugsComm slugs of+    Noncommutative ->+      forM_ slugs $ \slug ->+      forM_ (zip (accParams slug) (slugAccs slug)) $ \(p, (acc, acc_is)) ->+      copyDWIM (paramName p) [] (Var acc) acc_is+    _ -> doTheReduction++  return slugs_op_renamed+ reductionStageTwo :: KernelConstants-                  -> Pattern ExplicitMemory+                  -> [PatElem ExplicitMemory]                   -> Imp.Exp+                  -> Imp.Exp                   -> [Imp.Exp]                   -> Imp.Exp-                  -> PrimExp Imp.ExpLeaf-                  -> [LParam InKernel]-                  -> [LParam InKernel]-                  -> Lambda InKernel-                  -> [SubExp]                   -> Imp.Exp-                  -> VName-                  -> VName-                  -> [VName]-                  -> [VName]+                  -> SegRedOpSlug+                  -> [LParam InKernel]+                  -> Lambda InKernel -> [SubExp]+                  -> Imp.Exp -> VName -> Imp.Exp -> VName -> [VName] -> [VName]                   -> InKernelGen ()-reductionStageTwo constants segred_pat-                  flat_segment_id segment_gtids first_group_for_segment groups_per_segment-                  group_result_params red_acc_params+reductionStageTwo constants segred_pes+                  group_id flat_segment_id segment_gtids first_group_for_segment groups_per_segment+                  slug red_acc_params                   red_op_renamed nes-                  num_counters counter sync_arr group_res_arrs red_arrs = do+                  num_counters counter counter_i sync_arr group_res_arrs red_arrs = do   let local_tid = kernelLocalThreadId constants-      group_id = kernelGroupId constants       group_size = kernelGroupSize constants   old_counter <- dPrim "old_counter" int32-  (counter_mem, _, counter_offset) <- ImpGen.fullyIndexArray counter [flat_segment_id `rem` num_counters]-  ImpGen.comment "first thread in group saves group result to memory" $+  (counter_mem, _, counter_offset) <- fullyIndexArray counter [counter_i * num_counters ++                                                               flat_segment_id `rem` num_counters]+  comment "first thread in group saves group result to global memory" $     sWhen (local_tid .==. 0) $ do-    forM_ (take (length nes) $ zip group_res_arrs group_result_params) $ \(v, p) ->-      ImpGen.copyDWIM v [group_id] (Var $ paramName p) []-    sOp Imp.MemFence+    forM_ (take (length nes) $ zip group_res_arrs (slugAccs slug)) $ \(v, (acc, acc_is)) ->+      copyDWIM v [0, group_id] (Var acc) acc_is+    sOp Imp.MemFenceGlobal     -- Increment the counter, thus stating that our result is     -- available.-    sOp $ Imp.Atomic $ Imp.AtomicAdd old_counter counter_mem counter_offset 1+    sOp $ Imp.Atomic DefaultSpace $ Imp.AtomicAdd old_counter counter_mem counter_offset 1     -- Now check if we were the last group to write our result.  If     -- so, it is our responsibility to produce the final result.-    ImpGen.sWrite sync_arr [0] $ Imp.var old_counter int32 .==. groups_per_segment - 1+    sWrite sync_arr [0] $ Imp.var old_counter int32 .==. groups_per_segment - 1    sOp Imp.LocalBarrier    is_last_group <- dPrim "is_last_group" Bool-  ImpGen.copyDWIM is_last_group [] (Var sync_arr) [0]+  copyDWIM is_last_group [] (Var sync_arr) [0]   sWhen (Imp.var is_last_group Bool) $ do     -- The final group has written its result (and it was     -- us!), so read in all the group results and perform the@@ -539,28 +629,28 @@     -- with an atomic to avoid warnings about write/write     -- races in oclgrind.     sWhen (local_tid .==. 0) $-      sOp $ Imp.Atomic $ Imp.AtomicAdd old_counter counter_mem counter_offset $+      sOp $ Imp.Atomic DefaultSpace $ Imp.AtomicAdd old_counter counter_mem counter_offset $       negate groups_per_segment-    ImpGen.comment "read in the per-group-results" $-      forM_ (zip4 red_acc_params red_arrs nes group_res_arrs) $-      \(p, arr, ne, group_res_arr) -> do-        let load_group_result =-              ImpGen.copyDWIM (paramName p) []-              (Var group_res_arr) [first_group_for_segment + local_tid]-            load_neutral_element =-              ImpGen.copyDWIM (paramName p) [] ne []-        ImpGen.sIf (local_tid .<. groups_per_segment)-          load_group_result load_neutral_element-        when (primType $ paramType p) $-          ImpGen.copyDWIM arr [local_tid] (Var $ paramName p) []+    sLoopNest (slugShape slug) $ \vec_is -> do+      comment "read in the per-group-results" $+        forM_ (zip4 red_acc_params red_arrs nes group_res_arrs) $+        \(p, arr, ne, group_res_arr) -> do+          let load_group_result =+                copyDWIM (paramName p) []+                (Var group_res_arr) ([0, first_group_for_segment + local_tid] ++ vec_is)+              load_neutral_element =+                copyDWIM (paramName p) [] ne []+          sIf (local_tid .<. groups_per_segment)+            load_group_result load_neutral_element+          when (primType $ paramType p) $+            copyDWIM arr [local_tid] (Var $ paramName p) [] -    sOp Imp.LocalBarrier+      sOp Imp.LocalBarrier -    sComment "reduce the per-group results" $ do-      groupReduce constants group_size red_op_renamed red_arrs+      sComment "reduce the per-group results" $ do+        groupReduce constants group_size red_op_renamed red_arrs -      sComment "and back to memory with the final result" $-        sWhen (local_tid .==. 0) $-        forM_ (take (length nes) $ zip (patternNames segred_pat) $-               lambdaParams red_op_renamed) $ \(v, p) ->-        ImpGen.copyDWIM v segment_gtids (Var $ paramName p) []+        sComment "and back to memory with the final result" $+          sWhen (local_tid .==. 0) $+          forM_ (zip segred_pes $ lambdaParams red_op_renamed) $ \(pe, p) ->+          copyDWIM (patElemName pe) (segment_gtids++vec_is) (Var $ paramName p) []

@@ -0,0 +1,258 @@+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE FlexibleContexts #-}+module Futhark.CodeGen.ImpGen.Kernels.SegScan+  ( compileSegScan )+  where++import Control.Monad.Except+import Data.Maybe+import Data.List++import Prelude hiding (quot, rem)++import Futhark.MonadFreshNames+import Futhark.Transform.Rename+import Futhark.Representation.ExplicitMemory+import qualified Futhark.CodeGen.ImpCode.Kernels as Imp+import Futhark.CodeGen.ImpGen+import Futhark.CodeGen.ImpGen.Kernels.Base+import qualified Futhark.Representation.ExplicitMemory.IndexFunction as IxFun+import Futhark.Util.IntegralExp (quotRoundingUp, quot, rem)++makeLocalArrays :: SubExp -> SubExp -> [SubExp] -> Lambda InKernel+                -> InKernelGen [VName]+makeLocalArrays group_size num_threads nes scan_op = do+  let (scan_x_params, _scan_y_params) =+        splitAt (length nes) $ lambdaParams scan_op+  forM scan_x_params $ \p ->+    case paramAttr p of+      MemArray pt shape _ (ArrayIn mem _) -> do+        let shape' = Shape [num_threads] <> shape+        sArray "scan_arr" pt shape' $+          ArrayIn mem $ IxFun.iota $ map (primExpFromSubExp int32) $ shapeDims shape'+      _ -> do+        let pt = elemType $ paramType p+            shape = Shape [group_size]+        sAllocArray "scan_arr" pt shape $ Space "local"++type CrossesSegment = Maybe (Imp.Exp -> Imp.Exp -> Imp.Exp)++-- | Produce partially scanned intervals; one per workgroup.+scanStage1 :: Pattern ExplicitMemory+           -> KernelSpace+           -> Lambda InKernel -> [SubExp]+           -> KernelBody InKernel+           -> CallKernelGen (Imp.Exp, CrossesSegment)+scanStage1 (Pattern _ pes) space scan_op nes kbody = do+  (base_constants, init_constants) <- kernelInitialisationSetSpace space $ return ()+  let (gtids, dims) = unzip $ spaceDimensions space+  dims' <- mapM toExp dims+  let constants = base_constants { kernelThreadActive = true }+      num_elements = product dims'+      elems_per_thread = num_elements `quotRoundingUp` kernelNumThreads constants+      elems_per_group = kernelGroupSize constants * elems_per_thread++  -- Squirrel away a copy of the operator with unique names that we+  -- can pass to groupScan.+  scan_op_renamed <- renameLambda scan_op++  let crossesSegment =+        case reverse dims' of+          segment_size : _ : _ -> Just $ \from to ->+            (to-from) .>. (to `rem` segment_size)+          _ -> Nothing++  sKernel constants "scan_stage1" $ allThreads constants $ do+    init_constants++    local_arrs <-+      makeLocalArrays (spaceGroupSize space) (spaceNumThreads space)+      nes scan_op++    -- The variables from scan_op will be used for the carry and such+    -- in the big chunking loop.+    dScope Nothing $ scopeOfLParams $ lambdaParams scan_op+    let (scan_x_params, scan_y_params) =+          splitAt (length nes) $ lambdaParams scan_op++    forM_ (zip scan_x_params nes) $ \(p, ne) ->+      copyDWIM (paramName p) [] ne []++    j <- newVName "j"+    sFor j Int32 elems_per_thread $ do+      chunk_offset <- dPrimV "chunk_offset" $+                      kernelGroupSize constants * Imp.var j int32 ++                      kernelGroupId constants * elems_per_group+      flat_idx <- dPrimV "flat_idx" $+                  Imp.var chunk_offset int32 + kernelLocalThreadId constants+      -- Construct segment indices.+      zipWithM_ (<--) gtids $ unflattenIndex dims' $ Imp.var flat_idx int32++      let in_bounds =+            foldl1 (.&&.) $ zipWith (.<.) (map (`Imp.var` int32) gtids) dims'+          when_in_bounds = compileStms mempty (kernelBodyStms kbody) $ do+            let (scan_res, map_res) = splitAt (length nes) $ kernelBodyResult kbody+            sComment "write to-scan values to parameters" $+              forM_ (zip scan_y_params scan_res) $ \(p, se) ->+              copyDWIM (paramName p) [] (kernelResultSubExp se) []+            sComment "write mapped values results to global memory" $+              forM_ (zip (drop (length nes) pes) map_res) $ \(pe, se) ->+              copyDWIM (patElemName pe) (map (`Imp.var` int32) gtids)+              (kernelResultSubExp se) []+          when_out_of_bounds = forM_ (zip scan_y_params nes) $ \(p, ne) ->+            copyDWIM (paramName p) [] ne []++      sComment "threads in bounds read input; others get neutral element" $+        sIf in_bounds when_in_bounds when_out_of_bounds++      sComment "combine with carry and write to local memory" $+        compileStms mempty (bodyStms $ lambdaBody scan_op) $+        forM_ (zip local_arrs $ bodyResult $ lambdaBody scan_op) $ \(arr, se) ->+          copyDWIM arr [kernelLocalThreadId constants] se []++      let crossesSegment' = do+            f <- crossesSegment+            Just $ \from to ->+              let from' = from + Imp.var chunk_offset int32+                  to' = to + Imp.var chunk_offset int32+              in f from' to'++      groupScan constants crossesSegment'+        (kernelGroupSize constants) scan_op_renamed local_arrs++      sComment "threads in bounds write partial scan result" $+        sWhen in_bounds $ forM_ (zip pes local_arrs) $ \(pe, arr) ->+        copyDWIM (patElemName pe) (map (`Imp.var` int32) gtids)+        (Var arr) [kernelLocalThreadId constants]++      sOp Imp.LocalBarrier++      let load_carry =+            forM_ (zip local_arrs scan_x_params) $ \(arr, p) ->+            copyDWIM (paramName p) [] (Var arr) [kernelGroupSize constants - 1]+          load_neutral =+            forM_ (zip nes scan_x_params) $ \(ne, p) ->+            copyDWIM (paramName p) [] ne []++      sComment "first thread reads last element as carry-in for next iteration" $+        sWhen (kernelLocalThreadId constants .==. 0) $+        case crossesSegment of Nothing -> load_carry+                               Just f -> sIf (f (Imp.var chunk_offset int32 ++                                                 kernelGroupSize constants-1)+                                                (Imp.var chunk_offset int32 ++                                                 kernelGroupSize constants))+                                         load_neutral load_carry++      sOp Imp.LocalBarrier++  return (elems_per_group, crossesSegment)++scanStage2 :: Pattern ExplicitMemory+           -> Imp.Exp -> CrossesSegment -> KernelSpace+           -> Lambda InKernel -> [SubExp]+           -> CallKernelGen ()+scanStage2 (Pattern _ pes) elems_per_group crossesSegment space scan_op nes = do+  -- A single group, with one thread for each group in stage 1.+  group_size <- toExp $ spaceNumGroups space+  (constants, init_constants) <-+    kernelInitialisationSimple 1 group_size Nothing++  let (gtids, dims) = unzip $ spaceDimensions space+  dims' <- mapM toExp dims+  let crossesSegment' = do+        f <- crossesSegment+        Just $ \from to ->+          f ((from + 1) * elems_per_group - 1) ((to + 1) * elems_per_group - 1)++  sKernel constants "scan_stage2" $ do+    init_constants++    local_arrs <- makeLocalArrays (spaceNumGroups space) (spaceNumGroups space)+                  nes scan_op++    flat_idx <- dPrimV "flat_idx" $+      (kernelLocalThreadId constants + 1) * elems_per_group - 1+    -- Construct segment indices.+    zipWithM_ dPrimV_ gtids $ unflattenIndex dims' $ Imp.var flat_idx int32++    let in_bounds =+          foldl1 (.&&.) $ zipWith (.<.) (map (`Imp.var` int32) gtids) dims'+        when_in_bounds = forM_ (zip local_arrs pes) $ \(arr, pe) ->+          copyDWIM arr [kernelLocalThreadId constants]+          (Var $ patElemName pe) $ map (`Imp.var` int32) gtids+        when_out_of_bounds = forM_ (zip local_arrs nes) $ \(arr, ne) ->+          copyDWIM arr [kernelLocalThreadId constants] ne []++    sComment "threads in bound read carries; others get neutral element" $+      sIf in_bounds when_in_bounds when_out_of_bounds++    groupScan constants crossesSegment'+      (kernelGroupSize constants) scan_op local_arrs++    sComment "threads in bounds write scanned carries" $+      sWhen in_bounds $ forM_ (zip pes local_arrs) $ \(pe, arr) ->+      copyDWIM (patElemName pe) (map (`Imp.var` int32) gtids)+      (Var arr) [kernelLocalThreadId constants]++scanStage3 :: Pattern ExplicitMemory+           -> Imp.Exp -> CrossesSegment -> KernelSpace+           -> Lambda InKernel -> [SubExp]+           -> CallKernelGen ()+scanStage3 (Pattern _ pes) elems_per_group crossesSegment space scan_op nes = do+  let (gtids, dims) = unzip $ spaceDimensions space+  dims' <- mapM toExp dims+  (constants, init_constants) <- simpleKernelConstants (product dims') "scan"+  sKernel constants "scan_stage3" $ do+    init_constants+    -- Compute our logical index.+    zipWithM_ dPrimV_ gtids $ unflattenIndex dims' $ kernelGlobalThreadId constants+    -- Figure out which group this element was originally in.+    orig_group <- dPrimV "orig_group" $+                  kernelGlobalThreadId constants `quot` elems_per_group+    -- Then the index of the carry-in of the preceding group.+    carry_in_flat_idx <- dPrimV "carry_in_flat_idx" $+                         Imp.var orig_group int32 * elems_per_group - 1+    -- Figure out the logical index of the carry-in.+    let carry_in_idx = unflattenIndex dims' $ Imp.var carry_in_flat_idx int32++    -- Apply the carry if we are not in the scan results for the first+    -- group, and are not the last element in such a group (because+    -- then the carry was updated in stage 2), and we are not crossing+    -- a segment boundary.+    let crosses_segment = fromMaybe false $+          crossesSegment <*>+            pure (Imp.var carry_in_flat_idx int32) <*>+            pure (kernelGlobalThreadId constants)+        is_a_carry = kernelGlobalThreadId constants .==.+                     (Imp.var orig_group int32 + 1) * elems_per_group - 1+        no_carry_in = Imp.var orig_group int32 .==. 0 .||. is_a_carry .||. crosses_segment++    sWhen (kernelThreadActive constants) $ sUnless no_carry_in $ do+      dScope Nothing $ scopeOfLParams $ lambdaParams scan_op+      let (scan_x_params, scan_y_params) =+            splitAt (length nes) $ lambdaParams scan_op+      forM_ (zip scan_x_params pes) $ \(p, pe) ->+        copyDWIM (paramName p) [] (Var $ patElemName pe) carry_in_idx+      forM_ (zip scan_y_params pes) $ \(p, pe) ->+        copyDWIM (paramName p) [] (Var $ patElemName pe) $ map (`Imp.var` int32) gtids+      compileBody' scan_x_params $ lambdaBody scan_op+      forM_ (zip scan_x_params pes) $ \(p, pe) ->+        copyDWIM (patElemName pe) (map (`Imp.var` int32) gtids) (Var $ paramName p) []++-- | Compile 'SegScan' instance to host-level code with calls to+-- various kernels.+compileSegScan :: Pattern ExplicitMemory+               -> KernelSpace+               -> Lambda InKernel -> [SubExp]+               -> KernelBody InKernel+               -> CallKernelGen ()+compileSegScan pat space scan_op nes kbody = do+  (elems_per_group, crossesSegment) <- scanStage1 pat space scan_op nes kbody++  emit $ Imp.DebugPrint "\n# SegScan" Nothing+  emit $ Imp.DebugPrint "elems_per_group" $ Just (int32, elems_per_group)++  scan_op' <- renameLambda scan_op+  scan_op'' <- renameLambda scan_op+  scanStage2 pat elems_per_group crossesSegment space scan_op' nes+  scanStage3 pat elems_per_group crossesSegment space scan_op'' nes

@@ -48,7 +48,7 @@       opencl_code = openClCode $ M.elems kernels       opencl_prelude = pretty $ genPrelude target requirements   return $ ImpOpenCL.Program opencl_code opencl_prelude kernel_names-    (S.toList $ kernelUsedTypes requirements) sizes $+    (S.toList $ openclUsedTypes requirements) sizes $     ImpOpenCL.Functions (M.toList extra_funs) <> prog'   where genPrelude TargetOpenCL = genOpenClPrelude         genPrelude TargetCUDA = genCUDAPrelude@@ -63,10 +63,18 @@ pointerQuals "kernel"     = return [C.ctyquals|__kernel|] pointerQuals s            = fail $ "'" ++ s ++ "' is not an OpenCL kernel address space." -type UsedFunctions = [(String,C.Func)] -- The ordering is important!+newtype KernelRequirements =+  KernelRequirements { kernelLocalMemory :: [LocalMemoryUse] } +instance Semigroup KernelRequirements where+  KernelRequirements lm1 <> KernelRequirements lm2 =+    KernelRequirements (lm1<>lm2)++instance Monoid KernelRequirements where+  mempty = KernelRequirements mempty+ newtype OpenClRequirements =-  OpenClRequirements { kernelUsedTypes :: S.Set PrimType }+  OpenClRequirements { openclUsedTypes :: S.Set PrimType }  instance Semigroup OpenClRequirements where   OpenClRequirements ts1 <> OpenClRequirements ts2 =@@ -104,16 +112,17 @@ onKernel :: KernelTarget -> Kernel -> OnKernelM OpenCL  onKernel target kernel = do-  let (kernel_body, _) =-        GenericC.runCompilerM (Functions []) inKernelOperations blankNameSource mempty $+  let (kernel_body, requirements) =+        GenericC.runCompilerM mempty inKernelOperations blankNameSource mempty $         GenericC.blockScope $ GenericC.compileCode $ kernelBody kernel        use_params = mapMaybe useAsParam $ kernelUses kernel -      (local_memory_params, local_memory_init) =-        unzip $+      (local_memory_args, local_memory_params, local_memory_init) =+        unzip3 $         flip evalState (blankNameSource :: VNameSource) $-        mapM (prepareLocalMemory target) $ kernelLocalMemory kernel+        mapM (prepareLocalMemory target) $ kernelLocalMemory $+        GenericC.compUserState requirements        -- CUDA has very strict restrictions on the number of blocks       -- permitted along the 'y' and 'z' dimensions of the grid@@ -151,27 +160,30 @@               , clRequirements = OpenClRequirements (typesInKernel kernel)               } -  return $ LaunchKernel name (kernelArgs kernel) num_groups group_size+  return $ LaunchKernel name (catMaybes local_memory_args ++ kernelArgs kernel) num_groups group_size   where name = nameToString $ kernelName kernel         num_groups = kernelNumGroups kernel         group_size = kernelGroupSize kernel -        prepareLocalMemory TargetOpenCL (mem, Left _) = do+        prepareLocalMemory TargetOpenCL (mem, Left size) = do           mem_aligned <- newVName $ baseString mem ++ "_aligned"-          return (Just [C.cparam|__local volatile typename int64_t* $id:mem_aligned|],+          return (Just $ SharedMemoryKArg size,+                  Just [C.cparam|__local volatile typename int64_t* $id:mem_aligned|],                   [C.citem|__local volatile char* restrict $id:mem = $id:mem_aligned;|])         prepareLocalMemory TargetOpenCL (mem, Right size) = do           let size' = compilePrimExp size-          return (Nothing,+          return (Nothing, Nothing,                   [C.citem|ALIGNED_LOCAL_MEMORY($id:mem, $exp:size');|])-        prepareLocalMemory TargetCUDA (mem, Left _) = do+        prepareLocalMemory TargetCUDA (mem, Left size) = do           param <- newVName $ baseString mem ++ "_offset"-          return (Just [C.cparam|uint $id:param|],+          return (Just $ SharedMemoryKArg size,+                  Just [C.cparam|uint $id:param|],                   [C.citem|volatile char *$id:mem = &shared_mem[$id:param];|])         prepareLocalMemory TargetCUDA (mem, Right size) = do+          -- We declare the shared memory array as int64_t to force alignment.           let size' = compilePrimExp size-          return (Nothing,-                  [CUDAC.citem|__shared__ volatile char $id:mem[$exp:size'];|])+          return (Nothing, Nothing,+                  [CUDAC.citem|__shared__ volatile typename int64_t $id:mem[(($exp:size' + 7) & ~7)/8];|])  useAsParam :: KernelUse -> Maybe C.Param useAsParam (ScalarUse name bt) =@@ -384,16 +396,14 @@           return [C.cexp|$id:(zEncodeString (pretty key))|]  kernelArgs :: Kernel -> [KernelArg]-kernelArgs kernel =-  mapMaybe (fmap (SharedMemoryKArg . memSizeToExp) . localMemorySize)-  (kernelLocalMemory kernel) ++-  mapMaybe useToArg (kernelUses kernel)-  where localMemorySize (_, Left size) = Just size-        localMemorySize (_, Right{}) = Nothing+kernelArgs = mapMaybe useToArg . kernelUses+  where useToArg (MemoryUse mem)  = Just $ MemKArg mem+        useToArg (ScalarUse v bt) = Just $ ValueKArg (LeafExp (ScalarVar v) bt) bt+        useToArg ConstUse{}       = Nothing  --- Generating C -inKernelOperations :: GenericC.Operations KernelOp UsedFunctions+inKernelOperations :: GenericC.Operations KernelOp KernelRequirements inKernelOperations = GenericC.Operations                      { GenericC.opsCompiler = kernelOps                      , GenericC.opsMemoryType = kernelMemoryType@@ -405,7 +415,7 @@                      , GenericC.opsStaticArray = noStaticArrays                      , GenericC.opsFatMemory = False                      }-  where kernelOps :: GenericC.OpCompiler KernelOp UsedFunctions+  where kernelOps :: GenericC.OpCompiler KernelOp KernelRequirements         kernelOps (GetGroupId v i) =           GenericC.stm [C.cstm|$id:v = get_group_id($int:i);|]         kernelOps (GetLocalId v i) =@@ -422,74 +432,83 @@           GenericC.stm [C.cstm|barrier(CLK_LOCAL_MEM_FENCE);|]         kernelOps GlobalBarrier =           GenericC.stm [C.cstm|barrier(CLK_GLOBAL_MEM_FENCE);|]-        kernelOps MemFence =+        kernelOps MemFenceLocal =+          GenericC.stm [C.cstm|mem_fence_local();|]+        kernelOps MemFenceGlobal =           GenericC.stm [C.cstm|mem_fence_global();|]-        kernelOps (Atomic aop) = atomicOps aop+        kernelOps (PrivateAlloc name size) = do+          size' <- GenericC.compileExp $ innerExp size+          name' <- newVName $ pretty name ++ "_backing"+          GenericC.item [C.citem|__private char $id:name'[$exp:size'];|]+          GenericC.stm [C.cstm|$id:name = $id:name';|]+        kernelOps (LocalAlloc name size) = do+          name' <- newVName $ pretty name ++ "_backing"+          GenericC.modifyUserState (<>KernelRequirements [(name', size)])+          GenericC.stm [C.cstm|$id:name = (__local char*) $id:name';|]+        kernelOps (Atomic space aop) = atomicOps space aop -        atomicOps (AtomicAdd old arr ind val) = do-          ind' <- GenericC.compileExp $ innerExp ind-          val' <- GenericC.compileExp val-          GenericC.stm [C.cstm|$id:old = atomic_add((volatile __global int *)&$id:arr[$exp:ind'], $exp:val');|]+        atomicCast s t = do+          let volatile = [C.ctyquals|volatile|]+          quals <- case s of DefaultSpace -> pointerQuals "global"+                             Space sid    -> pointerQuals sid+          return [C.cty|$tyquals:(volatile++quals) $ty:t|] -        atomicOps (AtomicSMax old arr ind val) = do+        doAtomic s old arr ind val op ty = do           ind' <- GenericC.compileExp $ innerExp ind           val' <- GenericC.compileExp val-          GenericC.stm [C.cstm|$id:old = atomic_max((volatile __global int *)&$id:arr[$exp:ind'], $exp:val');|]+          cast <- atomicCast s ty+          GenericC.stm [C.cstm|$id:old = $id:op(&(($ty:cast *)$id:arr)[$exp:ind'], ($ty:ty) $exp:val');|] -        atomicOps (AtomicSMin old arr ind val) = do-          ind' <- GenericC.compileExp $ innerExp ind-          val' <- GenericC.compileExp val-          GenericC.stm [C.cstm|$id:old = atomic_min((volatile __global int *)&$id:arr[$exp:ind'], $exp:val');|]+        atomicOps s (AtomicAdd old arr ind val) =+          doAtomic s old arr ind val "atomic_add" [C.cty|int|] -        atomicOps (AtomicUMax old arr ind val) = do-          ind' <- GenericC.compileExp $ innerExp ind-          val' <- GenericC.compileExp val-          GenericC.stm [C.cstm|$id:old = atomic_max((volatile __global unsigned int *)&$id:arr[$exp:ind'], (unsigned int)$exp:val');|]+        atomicOps s (AtomicSMax old arr ind val) =+          doAtomic s old arr ind val "atomic_max" [C.cty|int|] -        atomicOps (AtomicUMin old arr ind val) = do-          ind' <- GenericC.compileExp $ innerExp ind-          val' <- GenericC.compileExp val-          GenericC.stm [C.cstm|$id:old = atomic_min((volatile __global unsigned int *)&$id:arr[$exp:ind'], (unsigned int)$exp:val');|]+        atomicOps s (AtomicSMin old arr ind val) =+          doAtomic s old arr ind val "atomic_min" [C.cty|int|] -        atomicOps (AtomicAnd old arr ind val) = do-          ind' <- GenericC.compileExp $ innerExp ind-          val' <- GenericC.compileExp val-          GenericC.stm [C.cstm|$id:old = atomic_and((volatile __global unsigned int *)&$id:arr[$exp:ind'], (unsigned int)$exp:val');|]+        atomicOps s (AtomicUMax old arr ind val) =+          doAtomic s old arr ind val "atomic_max" [C.cty|unsigned int|] -        atomicOps (AtomicOr old arr ind val) = do-          ind' <- GenericC.compileExp $ innerExp ind-          val' <- GenericC.compileExp val-          GenericC.stm [C.cstm|$id:old = atomic_or((volatile __global unsigned int *)&$id:arr[$exp:ind'], (unsigned int)$exp:val');|]+        atomicOps s (AtomicUMin old arr ind val) =+          doAtomic s old arr ind val "atomic_min" [C.cty|unsigned int|] -        atomicOps (AtomicXor old arr ind val) = do-          ind' <- GenericC.compileExp $ innerExp ind-          val' <- GenericC.compileExp val-          GenericC.stm [C.cstm|$id:old = atomic_xor((volatile __global unsigned int *)&$id:arr[$exp:ind'], (unsigned int)$exp:val');|]+        atomicOps s (AtomicAnd old arr ind val) =+          doAtomic s old arr ind val "atomic_and" [C.cty|unsigned int|] -        atomicOps (AtomicCmpXchg old arr ind cmp val) = do+        atomicOps s (AtomicOr old arr ind val) =+          doAtomic s old arr ind val "atomic_or" [C.cty|unsigned int|]++        atomicOps s (AtomicXor old arr ind val) =+          doAtomic s old arr ind val "atomic_xor" [C.cty|unsigned int|]++        atomicOps s (AtomicCmpXchg old arr ind cmp val) = do           ind' <- GenericC.compileExp $ innerExp ind           cmp' <- GenericC.compileExp cmp           val' <- GenericC.compileExp val-          GenericC.stm [C.cstm|$id:old = atomic_cmpxchg((volatile __global int *)&$id:arr[$exp:ind'], $exp:cmp', $exp:val');|]+          cast <- atomicCast s [C.cty|int|]+          GenericC.stm [C.cstm|$id:old = atomic_cmpxchg(&(($ty:cast *)$id:arr)[$exp:ind'], $exp:cmp', $exp:val');|] -        atomicOps (AtomicXchg old arr ind val) = do+        atomicOps s (AtomicXchg old arr ind val) = do           ind' <- GenericC.compileExp $ innerExp ind           val' <- GenericC.compileExp val-          GenericC.stm [C.cstm|$id:old = atomic_xchg((volatile __global int *)&$id:arr[$exp:ind'], $exp:val');|]+          cast <- atomicCast s [C.cty|int|]+          GenericC.stm [C.cstm|$id:old = atomic_xchg(&(($ty:cast *)$id:arr)[$exp:ind'], $exp:val');|] -        cannotAllocate :: GenericC.Allocate KernelOp UsedFunctions+        cannotAllocate :: GenericC.Allocate KernelOp KernelRequirements         cannotAllocate _ =           fail "Cannot allocate memory in kernel" -        cannotDeallocate :: GenericC.Deallocate KernelOp UsedFunctions+        cannotDeallocate :: GenericC.Deallocate KernelOp KernelRequirements         cannotDeallocate _ _ =           fail "Cannot deallocate memory in kernel" -        copyInKernel :: GenericC.Copy KernelOp UsedFunctions+        copyInKernel :: GenericC.Copy KernelOp KernelRequirements         copyInKernel _ _ _ _ _ _ _ =           fail "Cannot bulk copy in kernel." -        noStaticArrays :: GenericC.StaticArray KernelOp UsedFunctions+        noStaticArrays :: GenericC.StaticArray KernelOp KernelRequirements         noStaticArrays _ _ _ _ =           fail "Cannot create static array in kernel." @@ -499,11 +518,6 @@  --- Checking requirements -useToArg :: KernelUse -> Maybe KernelArg-useToArg (MemoryUse mem)  = Just $ MemKArg mem-useToArg (ScalarUse v bt) = Just $ ValueKArg (LeafExp (ScalarVar v) bt) bt-useToArg ConstUse{}       = Nothing- typesInKernel :: Kernel -> S.Set PrimType typesInKernel kernel = typesInCode $ kernelBody kernel @@ -530,7 +544,7 @@   typesInExp e <> typesInCode c1 <> typesInCode c2 typesInCode (Assert e _ _) = typesInExp e typesInCode (Comment _ c) = typesInCode c-typesInCode (DebugPrint _ _ e) = typesInExp e+typesInCode (DebugPrint _ v) = maybe mempty (typesInExp . snd) v typesInCode Op{} = mempty  typesInExp :: Exp -> S.Set PrimType

@@ -80,8 +80,8 @@       , dec index_out $ v32 x_index * height + v32 y_index        , If (v32 get_global_id_0 .<. input_size)-        (Write odata (bytes $ (v32 odata_offset + v32 index_out) * tsize) t (Space "global") Nonvolatile $-         index idata (bytes $ (v32 idata_offset + v32 index_in) * tsize) t (Space "global") Nonvolatile)+        (Write odata (elements $ v32 odata_offset + v32 index_out) t (Space "global") Nonvolatile $+         index idata (elements $ v32 idata_offset + v32 index_in) t (Space "global") Nonvolatile)         mempty       ] @@ -113,10 +113,10 @@         in mconcat [ dec index_in $ (v32 y_index + i) * width + v32 x_index                    , when (v32 y_index + i .<. height .&&.                            v32 index_in .<. input_size) $-                     Write block (bytes $ ((v32 get_local_id_1 + i) * (tile_dim+1)-                                           + v32 get_local_id_0) * tsize)+                     Write block (elements $ (v32 get_local_id_1 + i) * (tile_dim+1)+                                             + v32 get_local_id_0)                      t (Space "local") Nonvolatile $-                     index idata (bytes $ (v32 idata_offset + v32 index_in) * tsize)+                     index idata (elements $ v32 idata_offset + v32 index_in)                      t (Space "global") Nonvolatile]       , Op LocalBarrier       , SetScalar x_index $ v32 get_group_id_1 * tile_dim + v32 get_local_id_0@@ -127,17 +127,16 @@         in mconcat [ dec index_out $ (v32 y_index + i) * height + v32 x_index                    , when (v32 y_index + i .<. width .&&.                            v32 index_out .<. output_size) $-                     Write odata (bytes $ (v32 odata_offset + v32 index_out) * tsize)+                     Write odata (elements $ v32 odata_offset + v32 index_out)                      t (Space "global") Nonvolatile $-                     index block (bytes $ (v32 get_local_id_0 * (tile_dim+1)-                                           +v32 get_local_id_1+i)*tsize)+                     index block (elements $ v32 get_local_id_0 * (tile_dim+1)+                                             + v32 get_local_id_1+i)                      t (Space "local") Nonvolatile                    ]       ]    where dec v e = DeclareScalar v int32 <> SetScalar v e         v32 = flip var int32-        tsize = LeafExp (SizeOf t) int32         tile_dim = 2 * block_dim          when a b = If a b mempty@@ -198,19 +197,18 @@           , dec y_index y_in_index           , dec index_in $ v32 y_index * width + v32 x_index           , when (v32 x_index .<. width .&&. v32 y_index .<. height .&&. v32 index_in .<. input_size) $-            Write block (bytes $ (v32 get_local_id_1 * (block_dim+1) + v32 get_local_id_0) * tsize)+            Write block (elements $ v32 get_local_id_1 * (block_dim+1) + v32 get_local_id_0)             t (Space "local") Nonvolatile $-            index idata (bytes $ (v32 idata_offset + v32 index_in) * tsize)+            index idata (elements $ v32 idata_offset + v32 index_in)             t (Space "global") Nonvolatile           , Op LocalBarrier           , SetScalar x_index x_out_index           , SetScalar y_index y_out_index           , dec index_out $ v32 y_index * height + v32 x_index           , when (v32 x_index .<. height .&&. v32 y_index .<. width .&&. v32 index_out .<. output_size) $-            Write odata (bytes $ (v32 odata_offset + v32 index_out) * tsize)+            Write odata (elements $ v32 odata_offset + v32 index_out)             t (Space "global") Nonvolatile $-            index block (bytes $ (v32 get_local_id_0 * (block_dim+1)-                                   +v32 get_local_id_1)*tsize)+            index block (elements $ v32 get_local_id_0 * (block_dim+1) + v32 get_local_id_1)             t (Space "local") Nonvolatile           ] @@ -218,8 +216,9 @@                    -> Kernel mapTransposeKernel desc block_dim_int args t kind =   Kernel-  { kernelBody = mapTranspose block_dim args t kind-  , kernelLocalMemory = [(block, Right block_size)]+  { kernelBody = DeclareMem block (Space "local") <>+                 Op (LocalAlloc block (Right block_size)) <>+                 mapTranspose block_dim args t kind   , kernelUses = uses   , kernelNumGroups = num_groups   , kernelGroupSize = group_size