futhark 0.17.3 → 0.18.1

raw patch · 111 files changed

+6264/−1208 lines, 111 filesdep ~versionsPVP ok

version bump matches the API change (PVP)

Dependency ranges changed: versions

API changes (from Hackage documentation)

Files

• docs/conf.py +1/−0
• docs/index.rst +1/−0
• docs/man/futhark-multicore.rst +109/−0
• docs/usage.rst +2/−2
• futhark.cabal +15/−1
• prelude/array.fut +16/−16
• prelude/math.fut +3/−8
• prelude/soacs.fut +7/−7
• rts/c/chaselev.h +180/−0
• rts/c/multicore_defs.h +108/−0
• rts/c/multicore_util.h +103/−0
• rts/c/scheduler.h +278/−0
• rts/c/scheduler_common.h +244/−0
• rts/c/scheduler_tune.h +127/−0
• rts/c/subtask_queue.h +247/−0
• rts/c/timing.h +24/−0
• rts/c/util.h +21/−0
• rts/python/opencl.py +2/−2
• src/Futhark/Actions.hs +53/−0
• src/Futhark/Analysis/HORep/SOAC.hs +6/−6
• src/Futhark/Analysis/Metrics.hs +4/−0
• src/Futhark/Analysis/PrimExp/Convert.hs +13/−3
• src/Futhark/Analysis/SymbolTable.hs +9/−9
• src/Futhark/CLI/Dev.hs +47/−4
• src/Futhark/CLI/Multicore.hs +17/−0
• src/Futhark/CLI/Test.hs +1/−1
• src/Futhark/CodeGen/Backends/CCUDA/Boilerplate.hs +2/−2
• src/Futhark/CodeGen/Backends/COpenCL/Boilerplate.hs +4/−3
• src/Futhark/CodeGen/Backends/GenericC.hs +35/−3
• src/Futhark/CodeGen/Backends/GenericPython.hs +1/−0
• src/Futhark/CodeGen/Backends/MulticoreC.hs +703/−0
• src/Futhark/CodeGen/Backends/PyOpenCL/Boilerplate.hs +1/−0
• src/Futhark/CodeGen/Backends/SimpleRep.hs +1/−0
• src/Futhark/CodeGen/ImpCode.hs +4/−2
• src/Futhark/CodeGen/ImpCode/Kernels.hs +11/−11
• src/Futhark/CodeGen/ImpCode/Multicore.hs +126/−0
• src/Futhark/CodeGen/ImpGen.hs +49/−57
• src/Futhark/CodeGen/ImpGen/Kernels.hs +2/−2
• src/Futhark/CodeGen/ImpGen/Kernels/Base.hs +135/−126
• src/Futhark/CodeGen/ImpGen/Kernels/SegHist.hs +156/−141
• src/Futhark/CodeGen/ImpGen/Kernels/SegMap.hs +8/−7
• src/Futhark/CodeGen/ImpGen/Kernels/SegRed.hs +84/−76
• src/Futhark/CodeGen/ImpGen/Kernels/SegScan.hs +60/−55
• src/Futhark/CodeGen/ImpGen/Kernels/ToOpenCL.hs +6/−2
• src/Futhark/CodeGen/ImpGen/Multicore.hs +101/−0
• src/Futhark/CodeGen/ImpGen/Multicore/Base.hs +458/−0
• src/Futhark/CodeGen/ImpGen/Multicore/SegHist.hs +359/−0
• src/Futhark/CodeGen/ImpGen/Multicore/SegMap.hs +62/−0
• src/Futhark/CodeGen/ImpGen/Multicore/SegRed.hs +267/−0
• src/Futhark/CodeGen/ImpGen/Multicore/SegScan.hs +298/−0
• src/Futhark/Construct.hs +6/−6
• src/Futhark/IR/Kernels/Kernel.hs +10/−10
• src/Futhark/IR/Kernels/Sizes.hs +4/−4
• src/Futhark/IR/MC.hs +82/−0
• src/Futhark/IR/MC/Op.hs +175/−0
• src/Futhark/IR/MCMem.hs +84/−0
• src/Futhark/IR/Mem.hs +21/−21
• src/Futhark/IR/Pretty.hs +1/−0
• src/Futhark/IR/Prop/TypeOf.hs +2/−2
• src/Futhark/IR/Prop/Types.hs +2/−2
• src/Futhark/IR/SOACS/SOAC.hs +13/−13
• src/Futhark/IR/SOACS/Simplify.hs +6/−6
• src/Futhark/IR/SegOp.hs +19/−19
• src/Futhark/IR/Syntax/Core.hs +10/−3
• src/Futhark/Internalise.hs +97/−97
• src/Futhark/Internalise/AccurateSizes.hs +2/−2
• src/Futhark/Internalise/Bindings.hs +2/−2
• src/Futhark/Internalise/Defunctionalise.hs +2/−2
• src/Futhark/Internalise/Lambdas.hs +9/−9
• src/Futhark/Internalise/Monomorphise.hs +8/−8
• src/Futhark/Internalise/TypesValues.hs +1/−1
• src/Futhark/Optimise/CSE.hs +14/−0
• src/Futhark/Optimise/DoubleBuffer.hs +94/−48
• src/Futhark/Optimise/Fusion.hs +5/−5
• src/Futhark/Optimise/Fusion/LoopKernel.hs +2/−2
• src/Futhark/Optimise/InPlaceLowering.hs +19/−3
• src/Futhark/Optimise/Simplify/ClosedForm.hs +3/−3
• src/Futhark/Optimise/Simplify/Rules.hs +30/−31
• src/Futhark/Optimise/Sink.hs +110/−60
• src/Futhark/Optimise/TileLoops.hs +34/−39
• src/Futhark/Optimise/Unstream.hs +118/−38
• src/Futhark/Pass/ExpandAllocations.hs +22/−35
• src/Futhark/Pass/ExplicitAllocations.hs +21/−21
• src/Futhark/Pass/ExplicitAllocations/Kernels.hs +10/−13
• src/Futhark/Pass/ExplicitAllocations/MC.hs +37/−0
• src/Futhark/Pass/ExplicitAllocations/SegOp.hs +10/−6
• src/Futhark/Pass/ExplicitAllocations/Seq.hs +0/−1
• src/Futhark/Pass/ExtractKernels.hs +4/−4
• src/Futhark/Pass/ExtractKernels/BlockedKernel.hs +3/−3
• src/Futhark/Pass/ExtractKernels/DistributeNests.hs +5/−5
• src/Futhark/Pass/ExtractKernels/Distribution.hs +1/−1
• src/Futhark/Pass/ExtractKernels/ISRWIM.hs +1/−1
• src/Futhark/Pass/ExtractKernels/Intragroup.hs +7/−7
• src/Futhark/Pass/ExtractKernels/StreamKernel.hs +16/−21
• src/Futhark/Pass/ExtractKernels/ToKernels.hs +1/−0
• src/Futhark/Pass/ExtractMulticore.hs +406/−0
• src/Futhark/Pass/KernelBabysitting.hs +9/−10
• src/Futhark/Pass/Simplify.hs +10/−0
• src/Futhark/Passes.hs +34/−3
• src/Futhark/Transform/FirstOrderTransform.hs +8/−8
• src/Futhark/Transform/Rename.hs +8/−0
• src/Futhark/TypeCheck.hs +10/−9
• src/Language/Futhark/Interpreter.hs +34/−34
• src/Language/Futhark/Parser/Parser.y +1/−1
• src/Language/Futhark/Pretty.hs +1/−1
• src/Language/Futhark/Prop.hs +12/−12
• src/Language/Futhark/Syntax.hs +1/−1
• src/Language/Futhark/TypeChecker.hs +1/−1
• src/Language/Futhark/TypeChecker/Monad.hs +2/−2
• src/Language/Futhark/TypeChecker/Terms.hs +31/−26
• src/futhark.hs +2/−0

@@ -262,6 +262,7 @@ man_pages = [     ('man/futhark', 'futhark', 'a parallel functional array language', [], 1),     ('man/futhark-c', 'futhark-c', 'compile Futhark to sequential C', [], 1),+    ('man/futhark-multicore', 'futhark-multicore', 'compile Futhark to multithreaded C', [], 1),     ('man/futhark-opencl', 'futhark-opencl', 'compile Futhark to OpenCL', [], 1),     ('man/futhark-cuda', 'futhark-cuda', 'compile Futhark to CUDA', [], 1),     ('man/futhark-python', 'futhark-python', 'compile Futhark to sequential Python', [], 1),

@@ -45,6 +45,7 @@    man/futhark-autotune.rst    man/futhark-bench.rst    man/futhark-c.rst+   man/futhark-multicore.rst    man/futhark-cuda.rst    man/futhark-dataset.rst    man/futhark-doc.rst

@@ -0,0 +1,109 @@+.. role:: ref(emphasis)++.. _futhark-multicore(1):++=================+futhark-multicore+=================++SYNOPSIS+========++futhark multicore [options...] <program.fut>++DESCRIPTION+===========++``futhark multicore`` translates a Futhark program to multithreaded C+code, and either compiles that C code with gcc(1) to an executable+binary program, or produces a ``.h`` and ``.c`` file that can be+linked with other code..  The standard Futhark optimisation pipeline+is used, and GCC is invoked with ``-O3 -lm -std=c11 -pthread``.++The resulting program will read the arguments to the entry point+(``main`` by default) from standard input and print its return value+on standard output.  The arguments are read and printed in Futhark+syntax.++OPTIONS+=======++-h+  Print help text to standard output and exit.++--library+  Generate a library instead of an executable.  Appends ``.c``/``.h``+  to the name indicated by the ``-o`` option to determine output+  file names.++-o outfile+  Where to write the result.  If the source program is named+  ``foo.fut``, this defaults to ``foo``.++--safe+  Ignore ``unsafe`` in program and perform safety checks unconditionally.++-v verbose+  Enable debugging output.  If compilation fails due to a compiler+  error, the result of the last successful compiler step will be+  printed to standard error.++-V+  Print version information on standard output and exit.++-W+  Do not print any warnings.++--Werror+  Treat warnings as errors.++EXECUTABLE OPTIONS+==================++The following options are accepted by executables generated by+``futhark multicore``.++-h, --help++  Print help text to standard output and exit.++-b, --binary-output++  Print the program result in the binary output format.  The default+  is human-readable text, which is very slow.++-D, --debugging++  Perform possibly expensive internal correctness checks and verbose+  logging.  Implies ``-L``.++-e, --entry-point=FUN++  The entry point to run.  Defaults to ``main``.++-L, --log++  Print various low-overhead logging information to stderr while+  running.++-r, --runs=NUM++  Perform NUM runs of the program.  With ``-t``, the runtime for each+  individual run will be printed.  Additionally, a single leading+  warmup run will be performed (not counted).  Only the final run will+  have its result written to stdout.++-t, --write-runtime-to=FILE++  Print the time taken to execute the program to the indicated file, an+  integral number of microseconds.++BUGS+====++Currently works only on Unix-like systems.++SEE ALSO+========++:ref:`futhark-c(1)`, :ref:`futhark-test(1)`

@@ -107,8 +107,8 @@     (:ref:`binary-data-format`).  For large outputs, this is     significantly faster and takes up less space. -Parallel Options-~~~~~~~~~~~~~~~~+GPU Options+~~~~~~~~~~~  The following options are supported by executables generated with the GPU backends (``opencl``, ``pyopencl``, and ``cuda``).

@@ -1,7 +1,7 @@ cabal-version: 2.4  name:           futhark-version:        0.17.3+version:        0.18.1 synopsis:       An optimising compiler for a functional, array-oriented language.  description:    Futhark is a small programming language designed to be compiled to@@ -79,6 +79,7 @@       Futhark.CLI.Dev       Futhark.CLI.Doc       Futhark.CLI.Misc+      Futhark.CLI.Multicore       Futhark.CLI.OpenCL       Futhark.CLI.Pkg       Futhark.CLI.PyOpenCL@@ -97,6 +98,7 @@       Futhark.CodeGen.Backends.GenericPython.AST       Futhark.CodeGen.Backends.GenericPython.Definitions       Futhark.CodeGen.Backends.GenericPython.Options+      Futhark.CodeGen.Backends.MulticoreC       Futhark.CodeGen.Backends.PyOpenCL       Futhark.CodeGen.Backends.PyOpenCL.Boilerplate       Futhark.CodeGen.Backends.SequentialC@@ -104,6 +106,7 @@       Futhark.CodeGen.Backends.SimpleRep       Futhark.CodeGen.ImpCode       Futhark.CodeGen.ImpCode.Kernels+      Futhark.CodeGen.ImpCode.Multicore       Futhark.CodeGen.ImpCode.OpenCL       Futhark.CodeGen.ImpCode.Sequential       Futhark.CodeGen.ImpGen@@ -116,6 +119,12 @@       Futhark.CodeGen.ImpGen.Kernels.SegScan       Futhark.CodeGen.ImpGen.Kernels.ToOpenCL       Futhark.CodeGen.ImpGen.Kernels.Transpose+      Futhark.CodeGen.ImpGen.Multicore+      Futhark.CodeGen.ImpGen.Multicore.Base+      Futhark.CodeGen.ImpGen.Multicore.SegHist+      Futhark.CodeGen.ImpGen.Multicore.SegMap+      Futhark.CodeGen.ImpGen.Multicore.SegRed+      Futhark.CodeGen.ImpGen.Multicore.SegScan       Futhark.CodeGen.ImpGen.OpenCL       Futhark.CodeGen.ImpGen.Sequential       Futhark.CodeGen.ImpGen.Transpose@@ -136,6 +145,9 @@       Futhark.IR.Kernels.Simplify       Futhark.IR.Kernels.Sizes       Futhark.IR.KernelsMem+      Futhark.IR.MC+      Futhark.IR.MC.Op+      Futhark.IR.MCMem       Futhark.IR.Mem       Futhark.IR.Mem.IxFun       Futhark.IR.Mem.Simplify@@ -195,6 +207,7 @@       Futhark.Pass.ExplicitAllocations.Kernels       Futhark.Pass.ExplicitAllocations.SegOp       Futhark.Pass.ExplicitAllocations.Seq+      Futhark.Pass.ExplicitAllocations.MC       Futhark.Pass.ExtractKernels       Futhark.Pass.ExtractKernels.BlockedKernel       Futhark.Pass.ExtractKernels.DistributeNests@@ -204,6 +217,7 @@       Futhark.Pass.ExtractKernels.Intragroup       Futhark.Pass.ExtractKernels.StreamKernel       Futhark.Pass.ExtractKernels.ToKernels+      Futhark.Pass.ExtractMulticore       Futhark.Pass.FirstOrderTransform       Futhark.Pass.KernelBabysitting       Futhark.Pass.Simplify

@@ -24,13 +24,13 @@ let init [n] 't (x: [n]t) = x[0:n-1]  -- | Take some number of elements from the head of the array.-let take [n] 't (i: i32) (x: [n]t): [i]t = x[0:i]+let take [n] 't (i: i64) (x: [n]t): [i]t = x[0:i]  -- | Remove some number of elements from the head of the array.-let drop [n] 't (i: i32) (x: [n]t) = x[i:]+let drop [n] 't (i: i64) (x: [n]t) = x[i:]  -- | Split an array at a given position.-let split [n] 't (i: i32) (xs: [n]t): ([i]t, []t) =+let split [n] 't (i: i64) (xs: [n]t): ([i]t, []t) =   (xs[:i] :> [i]t, xs[i:])  -- | Return the elements of the array in reverse order.@@ -46,28 +46,28 @@ -- | Concatenation where the result has a predetermined size.  If the -- provided size is wrong, the function will fail with a run-time -- error.-let concat_to [n] [m] 't (k: i32) (xs: [n]t) (ys: [m]t): *[k]t = xs ++ ys :> [k]t+let concat_to [n] [m] 't (k: i64) (xs: [n]t) (ys: [m]t): *[k]t = xs ++ ys :> [k]t  -- | Rotate an array some number of elements to the left.  A negative -- rotation amount is also supported. -- -- For example, if `b==rotate r a`, then `b[x+r] = a[x]`.-let rotate [n] 't (r: i32) (xs: [n]t): [n]t = intrinsics.rotate (r, xs) :> [n]t+let rotate [n] 't (r: i64) (xs: [n]t): [n]t = intrinsics.rotate (r, xs) :> [n]t  -- | Construct an array of consecutive integers of the given length, -- starting at 0.-let iota (n: i32): *[n]i32 =+let iota (n: i64): *[n]i64 =   0..1..<n  -- | Construct an array comprising valid indexes into some other -- array, starting at 0.-let indices [n] 't (_: [n]t) : *[n]i32 =+let indices [n] 't (_: [n]t) : *[n]i64 =   iota n  -- | Construct an array of the given length containing the given -- value.-let replicate 't (n: i32) (x: t): *[n]t =-  map (\_ -> x) (iota n)+let replicate 't (n: i64) (x: t): *[n]t =+  map (const x) (iota n)  -- | Copy a value.  The result will not alias anything. let copy 't (a: t): *t =@@ -79,7 +79,7 @@  -- | Like `flatten`@term, but where the final size is known.  Fails at -- runtime if the provided size is wrong.-let flatten_to [n][m] 't (l: i32) (xs: [n][m]t): [l]t =+let flatten_to [n][m] 't (l: i64) (xs: [n][m]t): [l]t =   flatten xs :> [l]t  -- | Combines the outer three dimensions of an array.@@ -91,15 +91,15 @@   flatten (flatten_3d xs)  -- | Splits the outer dimension of an array in two.-let unflatten [p] 't (n: i32) (m: i32) (xs: [p]t): [n][m]t =+let unflatten [p] 't (n: i64) (m: i64) (xs: [p]t): [n][m]t =   intrinsics.unflatten (n, m, xs) :> [n][m]t  -- | Splits the outer dimension of an array in three.-let unflatten_3d [p] 't (n: i32) (m: i32) (l: i32) (xs: [p]t): [n][m][l]t =+let unflatten_3d [p] 't (n: i64) (m: i64) (l: i64) (xs: [p]t): [n][m][l]t =   unflatten n m (unflatten (n*m) l xs)  -- | Splits the outer dimension of an array in four.-let unflatten_4d [p] 't (n: i32) (m: i32) (l: i32) (k: i32) (xs: [p]t): [n][m][l][k]t =+let unflatten_4d [p] 't (n: i64) (m: i64) (l: i64) (k: i64) (xs: [p]t): [n][m][l][k]t =   unflatten n m (unflatten_3d (n*m) l k xs)  let transpose [n] [m] 't (a: [n][m]t): [m][n]t =@@ -122,13 +122,13 @@   foldl (flip f) acc (reverse bs)  -- | Create a value for each point in a one-dimensional index space.-let tabulate 'a (n: i32) (f: i32 -> a): *[n]a =+let tabulate 'a (n: i64) (f: i64 -> a): *[n]a =   map1 f (iota n)  -- | Create a value for each point in a two-dimensional index space.-let tabulate_2d 'a (n: i32) (m: i32) (f: i32 -> i32 -> a): *[n][m]a =+let tabulate_2d 'a (n: i64) (m: i64) (f: i64 -> i64 -> a): *[n][m]a =   map1 (f >-> tabulate m) (iota n)  -- | Create a value for each point in a three-dimensional index space.-let tabulate_3d 'a (n: i32) (m: i32) (o: i32) (f: i32 -> i32 -> i32 -> a): *[n][m][o]a =+let tabulate_3d 'a (n: i64) (m: i64) (o: i64) (f: i64 -> i64 -> i64 -> a): *[n][m][o]a =   map1 (f >-> tabulate_2d m o) (iota n)

@@ -2,8 +2,6 @@  import "soacs" -local let const 'a 'b (x: a) (_: b): a = x- -- | Describes types of values that can be created from the primitive -- numeric types (and bool). module type from_prim = {@@ -122,8 +120,7 @@ module type real = {   include numeric -  val from_fraction: i32 -> i32 -> t-  val to_i32: t -> i32+  val from_fraction: i64 -> i64 -> t   val to_i64: t -> i64   val to_f64: t -> f64 @@ -852,8 +849,7 @@    let bool (x: bool) = if x then 1f64 else 0f64 -  let from_fraction (x: i32) (y: i32) = i32 x / i32 y-  let to_i32 (x: f64) = intrinsics.fptosi_f64_i32 x+  let from_fraction (x: i64) (y: i64) = i64 x / i64 y   let to_i64 (x: f64) = intrinsics.fptosi_f64_i64 x   let to_f64 (x: f64) = x @@ -960,8 +956,7 @@    let bool (x: bool) = if x then 1f32 else 0f32 -  let from_fraction (x: i32) (y: i32) = i32 x / i32 y-  let to_i32 (x: f32) = intrinsics.fptosi_f32_i32 x+  let from_fraction (x: i64) (y: i64) = i64 x / i64 y   let to_i64 (x: f32) = intrinsics.fptosi_f32_i64 x   let to_f64 (x: f32) = intrinsics.fpconv_f32_f64 x 

@@ -118,7 +118,7 @@ -- -- In practice, the *O(n)* behaviour only occurs if *m* is also very -- large.-let reduce_by_index 'a [m] [n] (dest : *[m]a) (f : a -> a -> a) (ne : a) (is : [n]i32) (as : [n]a) : *[m]a =+let reduce_by_index 'a [m] [n] (dest : *[m]a) (f : a -> a -> a) (ne : a) (is : [n]i64) (as : [n]a) : *[m]a =   intrinsics.hist (1, dest, f, ne, is, as) :> *[m]a  -- | Inclusive prefix scan.  Has the same caveats with respect to@@ -163,7 +163,7 @@  -- | `reduce_stream op f as` splits `as` into chunks, applies `f` to each -- of these in parallel, and uses `op` (which must be associative) to--- combine the per-chunk results into a final result.  The `i32`+-- combine the per-chunk results into a final result.  The `i64` -- passed to `f` is the size of the chunk.  This SOAC is useful when -- `f` can be given a particularly work-efficient sequential -- implementation.  Operationally, we can imagine that `as` is divided@@ -176,7 +176,7 @@ -- **Work:** *O(n)* -- -- **Span:** *O(log(n))*-let reduce_stream [n] 'a 'b (op: b -> b -> b) (f: (k: i32) -> [k]a -> b) (as: [n]a): b =+let reduce_stream [n] 'a 'b (op: b -> b -> b) (f: (k: i64) -> [k]a -> b) (as: [n]a): b =   intrinsics.reduce_stream (op, f, as)  -- | As `reduce_stream`@term, but the chunks do not necessarily@@ -186,7 +186,7 @@ -- **Work:** *O(n)* -- -- **Span:** *O(log(n))*-let reduce_stream_per [n] 'a 'b (op: b -> b -> b) (f: (k: i32) -> [k]a -> b) (as: [n]a): b =+let reduce_stream_per [n] 'a 'b (op: b -> b -> b) (f: (k: i64) -> [k]a -> b) (as: [n]a): b =   intrinsics.reduce_stream_per (op, f, as)  -- | Similar to `reduce_stream`@term, except that each chunk must produce@@ -196,7 +196,7 @@ -- **Work:** *O(n)* -- -- **Span:** *O(1)*-let map_stream [n] 'a 'b (f: (k: i32) -> [k]a -> [k]b) (as: [n]a): *[n]b =+let map_stream [n] 'a 'b (f: (k: i64) -> [k]a -> [k]b) (as: [n]a): *[n]b =   intrinsics.map_stream (f, as) :> *[n]b  -- | Similar to `map_stream`@term, but the chunks do not necessarily@@ -206,7 +206,7 @@ -- **Work:** *O(n)* -- -- **Span:** *O(1)*-let map_stream_per [n] 'a 'b (f: (k: i32) -> [k]a -> [k]b) (as: [n]a): *[n]b =+let map_stream_per [n] 'a 'b (f: (k: i64) -> [k]a -> [k]b) (as: [n]a): *[n]b =   intrinsics.map_stream_per (f, as) :> *[n]b  -- | Return `true` if the given function returns `true` for all@@ -252,5 +252,5 @@ -- **Work:** *O(n)* -- -- **Span:** *O(1)*-let scatter 't [m] [n] (dest: *[m]t) (is: [n]i32) (vs: [n]t): *[m]t =+let scatter 't [m] [n] (dest: *[m]t) (is: [n]i64) (vs: [n]t): *[m]t =   intrinsics.scatter (dest, is, vs) :> *[m]t

@@ -0,0 +1,180 @@+// Start of chaselev.h++#ifndef _CHASELEV_H_+#define _CHASELEV_H_++/* Implementation of Chase-lev's concurrent lock-free deque+   from ``Dynamic Circular Work-Stealing Deque`` (2005)+   This implementation was ported from+   https://github.com/deepsea-inria/heartbeat++   !!!+   This implementation leaks memory,+   if the circular array is grown+   as we don't maintain a list of the old buffers.+   However, we can't safely free it either as a stealing thread might+   be reading from it.+   !!!+ */++#if defined(MCCHASELEV)+#include <stdlib.h>+#include <assert.h>+#include <string.h>+++static struct subtask* const STEAL_RES_EMPTY = (struct subtask*) 0;+static struct subtask* const STEAL_RES_ABORT = (struct subtask*) 1;++static const int strong = 0;+static const int backoff_nb_cycles = 1l << 10;+++static inline struct subtask* cb_get(struct subtask **buf, int64_t capacity, int64_t i)  {+  return (struct subtask*)__atomic_load_n(&buf[i % capacity], __ATOMIC_RELAXED);+}++static inline void cb_put (struct subtask **buf, int64_t capacity, int64_t i, struct subtask* x) {+  __atomic_store_n(&buf[i % capacity], x, __ATOMIC_RELAXED);+}++struct deque_buffer* grow(struct subtask **old_array,+                          int64_t old_capacity,+                          int64_t new_capacity,+                          int64_t b,+                          int64_t t)+{+  struct deque_buffer* new_deque_buffer = malloc(sizeof(struct deque_buffer));+  new_deque_buffer->size = new_capacity;+  new_deque_buffer->array = calloc(new_capacity,  sizeof(struct subtask*));++  for (int64_t i = t; i < b; i++) {+    cb_put(new_deque_buffer->array, new_capacity, i, cb_get(old_array, old_capacity, i));+  }+  return new_deque_buffer;+}++static inline int deque_init(struct deque *q, int64_t capacity) {+  assert(q != NULL);+  memset(q, 0, sizeof(struct deque));++  q->buffer = malloc(sizeof(struct deque_buffer));+  q->buffer->array = calloc(capacity, sizeof(struct subtask*));+  q->buffer->size = capacity;++  q->dead = 0;++  if (q->buffer->array == NULL) {+    return -1;+  }++  if (q->buffer == NULL) {+    return -1;+  }+  return 0;+}++static inline void deque_destroy(struct deque* q)+{+  q->dead = 1;+  free(q->buffer->array);+  free(q->buffer);+}++static inline int cas_top (struct deque *q, int64_t old_val, int64_t new_val) {+  int64_t ov = old_val;+  if(__atomic_compare_exchange_n(&q->top, &ov, new_val, strong,+                                 __ATOMIC_SEQ_CST, __ATOMIC_RELAXED)) {+    return 1;+  }+  spin_for(backoff_nb_cycles);+  return 0;+}+++void push_back(struct deque *q, struct subtask*subtask)+{+  assert(subtask != NULL);+  assert(q != NULL);++  int64_t b = __atomic_load_n(&q->bottom, __ATOMIC_RELAXED); // load atomically+  int64_t t = __atomic_load_n(&q->top, __ATOMIC_ACQUIRE);    // load atomically+  struct deque_buffer *buffer = __atomic_load_n(&q->buffer, __ATOMIC_RELAXED);+  if (b-t >= (buffer->size - 1)) {+    // grow_queue+    struct subtask **old_array = buffer->array;+    int64_t old_capacity = __atomic_load_n(&buffer->size, __ATOMIC_RELAXED);+    int64_t new_capacity = old_capacity * 2;+    struct deque_buffer *new_buffer = grow(old_array, old_capacity, new_capacity, b, t);+    __atomic_store_n(&q->buffer, new_buffer, __ATOMIC_RELEASE);+    buffer = __atomic_load_n(&q->buffer, __ATOMIC_RELAXED);+    memset(old_array, 0, sizeof(struct subtask*) * old_capacity);+    // free(old_array);  Not safe!!+  }++  cb_put(buffer->array, buffer->size, b, subtask);+  __atomic_thread_fence(__ATOMIC_RELEASE);+  __atomic_store_n(&q->bottom, b+1, __ATOMIC_RELAXED);+  return;+}+++struct subtask * pop_back(struct deque *q)+{+  int64_t b = __atomic_load_n(&q->bottom, __ATOMIC_RELAXED) - 1; // load atomically+  struct deque_buffer *buffer = __atomic_load_n(&q->buffer, __ATOMIC_RELAXED);+  __atomic_store_n(&q->bottom, b, __ATOMIC_RELAXED);+	__atomic_thread_fence(__ATOMIC_SEQ_CST);+  int64_t t = __atomic_load_n(&q->top, __ATOMIC_RELAXED);+  if (b < t) {+    __atomic_store_n(&q->bottom, t, __ATOMIC_RELAXED);+    return NULL;+  }+  struct subtask* item = cb_get(buffer->array, buffer->size, b);+  if (b > t) {+    return item;+  }++  // else there's only one item left+  // Did we win the race?+  if (!cas_top(q, t, t + 1)) {+    item = NULL;+  }+  __atomic_store_n(&q->bottom, t+1, __ATOMIC_RELAXED);+  return item;+}++struct subtask* steal(struct deque *q)+{+  assert(q != NULL);++  int64_t t = __atomic_load_n(&q->top, __ATOMIC_ACQUIRE);    // load atomically+  __atomic_thread_fence(__ATOMIC_SEQ_CST);+  int64_t b = __atomic_load_n(&q->bottom, __ATOMIC_ACQUIRE); // load atomically+  if (t >= b) {+    return STEAL_RES_EMPTY;+  }++  struct deque_buffer *buffer = __atomic_load_n(&q->buffer, __ATOMIC_CONSUME);+  struct subtask* item = cb_get(buffer->array, buffer->size, t);+  if (!cas_top(q, t, t + 1)) {+    return STEAL_RES_ABORT;+  }++  return item;+}+++static inline size_t nb_subtasks(struct deque *q)+{+  return (size_t)__atomic_load_n(&q->bottom, __ATOMIC_RELAXED) - __atomic_load_n(&q->top, __ATOMIC_RELAXED);+}++static inline int empty(struct deque *q)+{+  return nb_subtasks(q) < 1;+}++#endif+#endif+// end of chaselev.h

@@ -0,0 +1,108 @@+// start of multicore_defs.h++#ifndef MULTICORE_DEFS+#define MULTICORE_DEFS++#include <signal.h>++/* #define MCPROFILE */++// Which queue implementation to use+#define MCJOBQUEUE+// NOTE! MCCHASELEV has been removed from multicore branch+// Switch to multicore_deque branch to use chase-lev deque+/* #define MCCHASELEV */+++#if defined(_WIN32)+#include <windows.h>+#elif defined(__APPLE__)+#include <sys/sysctl.h>+// For getting cpu usage of threads+#include <mach/mach.h>+#include <sys/resource.h>+#elif defined(__linux__)+#include <sys/sysinfo.h>+#include <sys/resource.h>+#include <signal.h>+#endif+++// Forward declarations+// Scheduler definitions+struct scheduler;+struct scheduler_info;+struct scheduler_subtask;+struct scheduler_task;+++struct subtask_queue {+  int capacity;             // Size of the buffer.+  int first;                // Index of the start of the ring buffer.+  int num_used;             // Number of used elements in the buffer.+  struct subtask **buffer;++  pthread_mutex_t mutex;    // Mutex used for synchronisation.+  pthread_cond_t cond;      // Condition variable used for synchronisation.+  int dead;++#if defined(MCPROFILE)+  /* Profiling fields */+  uint64_t time_enqueue;+  uint64_t time_dequeue;+  uint64_t n_dequeues;+  uint64_t n_enqueues;+#endif+};++++// Function definitions+typedef int (*segop_fn)(void* args, int64_t iterations, int tid, struct scheduler_info info);+typedef int (*parloop_fn)(void* args, int64_t start, int64_t end, int subtask_id, int tid);+++/* A subtask that can be executed by a worker */+struct subtask {+  /* The parloop function */+  parloop_fn fn;+  /* Execution parameters */+  void* args;+  int64_t start, end;+  int id;++  /* Dynamic scheduling parameters */+  int chunkable;+  int64_t chunk_size;++  /* Shared variables across subtasks */+  volatile int *counter; // Counter for ongoing subtasks+  // Shared task timers and iterators+  int64_t *task_time;+  int64_t *task_iter;++  /* For debugging */+  const char *name;+};+++struct worker {+  pthread_t thread;+  struct scheduler *scheduler;  /* Reference to the scheduler struct the worker belongs to*/+  struct subtask_queue q;+  int dead;+  int tid;                      /* Just a thread id */++  /* "thread local" time fields used for online algorithm */+  uint64_t timer;+  uint64_t total;+  int nested; /* How nested the current computation is */++  // Profiling fields+  int output_usage;            /* Whether to dump thread usage */+  uint64_t time_spent_working; /* Time spent in parloop functions */+};++#endif++// end of multicore_defs.h

@@ -0,0 +1,103 @@+// start of multicore_util.h++/* Multicore Utility functions */++#ifndef _MULTICORE_UTIL_H_+#define _MULTICORE_UTIL_H_++/* A wrapper for getting rusage on Linux and MacOS */+/* TODO maybe figure out this for windows */+static inline int getrusage_thread(struct rusage *rusage)+{+  int err = -1;+#if  defined(__APPLE__)+    thread_basic_info_data_t info = { 0 };+    mach_msg_type_number_t info_count = THREAD_BASIC_INFO_COUNT;+    kern_return_t kern_err;++    kern_err = thread_info(mach_thread_self(),+                           THREAD_BASIC_INFO,+                           (thread_info_t)&info,+                           &info_count);+    if (kern_err == KERN_SUCCESS) {+        memset(rusage, 0, sizeof(struct rusage));+        rusage->ru_utime.tv_sec = info.user_time.seconds;+        rusage->ru_utime.tv_usec = info.user_time.microseconds;+        rusage->ru_stime.tv_sec = info.system_time.seconds;+        rusage->ru_stime.tv_usec = info.system_time.microseconds;+        err = 0;+    } else {+        errno = EINVAL;+    }+#elif defined(__linux__)+    err = getrusage(RUSAGE_THREAD, rusage);+#endif+    return err;+}++/* returns the number of logical cores */+static int num_processors()+{+#if  defined(_WIN32)+/* https://docs.microsoft.com/en-us/windows/win32/api/sysinfoapi/ns-sysinfoapi-system_info */+    SYSTEM_INFO sysinfo;+    GetSystemInfo(&sysinfo);+    int ncores = sysinfo.dwNumberOfProcessors;+    fprintf(stderr, "Found %d cores on your Windows machine\n Is that correct?\n", ncores);+    return ncores;+#elif defined(__APPLE__)+    int ncores;+    size_t ncores_size = sizeof(ncores);+    CHECK_ERRNO(sysctlbyname("hw.logicalcpu", &ncores, &ncores_size, NULL, 0),+                "sysctlbyname (hw.logicalcpu)");+    return ncores;+#elif defined(__linux__)+  return get_nprocs();+#else+  fprintf(stderr, "operating system not recognised\n");+  return -1;+#endif+}++static inline void output_thread_usage(struct worker *worker)+{+  struct rusage usage;+  CHECK_ERRNO(getrusage_thread(&usage), "getrusage_thread");+  struct timeval user_cpu_time = usage.ru_utime;+  struct timeval sys_cpu_time = usage.ru_stime;+  fprintf(stderr, "tid: %2d - work time %10llu us - user time: %10llu us - sys: %10llu us\n",+          worker->tid,+          (long long unsigned)worker->time_spent_working / 1000,+          (long long unsigned)(user_cpu_time.tv_sec * 1000000 + user_cpu_time.tv_usec),+          (long long unsigned)(sys_cpu_time.tv_sec * 1000000 + sys_cpu_time.tv_usec));+}+++static unsigned int g_seed;++// Used to seed the generator.+static inline void fast_srand(unsigned int seed) {+    g_seed = seed;+}++// Compute a pseudorandom integer.+// Output value in range [0, 32767]+static inline unsigned int fast_rand(void) {+    g_seed = (214013*g_seed+2531011);+    return (g_seed>>16)&0x7FFF;+}+++int64_t min_int64(int64_t x, int64_t y)+{+  return x < y ? x : y;+}++int64_t max_int64(int64_t x, int64_t y)+{+  return x > y ? x : y;+}+++#endif+// end of multicore_util.h

@@ -0,0 +1,278 @@+// start of scheduler.h+#ifndef _SCHEDULER_H_+#define _SCHEDULER_H_+++static int dummy_counter = 0;+static int64_t dummy_timer = 0;+static int64_t dummy_iter = 0;++static int dummy_fn(void *args, int64_t start, int64_t end, int subtask_id, int tid) {+  (void)args;+  (void)start;+  (void)end;+  (void)subtask_id;+  (void)tid;+  return 0;+}++// Wake up threads, who are blocking by pushing a dummy task+// onto their queue+static inline void wake_up_threads(struct scheduler *scheduler, int start_tid, int end_tid) {++#if defined(MCDEBUG)+  assert(start_tid >= 1);+  assert(end_tid <= scheduler->num_threads);+#endif+  for (int i = start_tid; i < end_tid; i++) {+    struct subtask *subtask = create_subtask(dummy_fn, NULL, "dummy_fn",+                                            &dummy_counter,+                                            &dummy_timer, &dummy_iter,+                                            0, 0,+                                            0, 0,+                                            0);+    CHECK_ERR(subtask_queue_enqueue(&scheduler->workers[i], subtask), "subtask_queue_enqueue");+  }+}++static inline int is_finished(struct worker *worker) {+  return worker->dead && subtask_queue_is_empty(&worker->q);+}++// Try to steal from a random queue+static inline int steal_from_random_worker(struct worker* worker)+{+  int my_id = worker->tid;+  struct scheduler* scheduler = worker->scheduler;+  int k = random_other_worker(scheduler, my_id);+  struct worker *worker_k = &scheduler->workers[k];+  struct subtask* subtask =  NULL;+  int retval = subtask_queue_steal(worker_k, &subtask);+  if (retval == 0) {+    subtask_queue_enqueue(worker, subtask);+    return 1;+  }+  return 0;+}+++static inline void *scheduler_worker(void* args)+{+  struct worker *worker = (struct worker*) args;+  worker_local = worker;+  struct subtask * subtask = NULL;+  while(!is_finished(worker)) {+    if (!subtask_queue_is_empty(&worker->q)) {+      int retval = subtask_queue_dequeue(worker, &subtask, 0);+      if (retval == 0) {+        assert(subtask != NULL);+        CHECK_ERR(run_subtask(worker, subtask), "run_subtask");+      } // else someone stole our work++    } else if (active_work) { /* steal */+      while (!is_finished(worker) && active_work) {+        if (steal_from_random_worker(worker)) {+          break;+        }+      }+    } else { /* go back to sleep and wait for work */+      int retval = subtask_queue_dequeue(worker, &subtask, 1);+      if (retval == 0) {+        assert(subtask != NULL);+        CHECK_ERR(run_subtask(worker, subtask), "run_subtask");+      }+    }+  }++  assert(subtask_queue_is_empty(&worker->q));+  __atomic_fetch_sub(&num_workers, 1, __ATOMIC_RELAXED);+#if defined(MCPROFILE)+  if (worker->output_usage)+    output_thread_usage(worker);+#endif+  return NULL;+}+++static inline int scheduler_execute_parloop(struct scheduler *scheduler,+                                            struct scheduler_parloop *task,+                                            int64_t *timer)+{++  struct worker * worker = worker_local;++  struct scheduler_info info = task->info;+  int64_t iter_pr_subtask = info.iter_pr_subtask;+  int64_t remainder = info.remainder;+  int nsubtasks = info.nsubtasks;+  volatile int join_counter = nsubtasks;++  // Shared timer used to sum up all+  // sequential work from each subtask+  int64_t task_timer = 0;+  int64_t task_iter = 0;++  enum scheduling sched = info.sched;+  /* If each subtasks should be processed in chunks */+  int chunkable = sched == STATIC ? 0 : 1;+  int64_t chunk_size = 1; // The initial chunk size when no info is avaliable+++  if (info.wake_up_threads || sched == DYNAMIC)+    __atomic_add_fetch(&active_work, nsubtasks, __ATOMIC_RELAXED);++  int64_t start = 0;+  int64_t end = iter_pr_subtask + (int64_t)(remainder != 0);+  for (int subtask_id = 0; subtask_id < nsubtasks; subtask_id++) {+    struct subtask *subtask = create_subtask(task->fn, task->args, task->name,+                                              &join_counter,+                                              &task_timer, &task_iter,+                                              start, end,+                                              chunkable, chunk_size,+                                              subtask_id);+    assert(subtask != NULL);+    if (worker->nested){+      CHECK_ERR(subtask_queue_enqueue(&scheduler->workers[worker->tid], subtask),+                "subtask_queue_enqueue");+    } else {+      CHECK_ERR(subtask_queue_enqueue(&scheduler->workers[subtask_id], subtask),+                "subtask_queue_enqueue");+    }+    // Update range params+    start = end;+    end += iter_pr_subtask + ((subtask_id + 1) < remainder);+  }++  if (info.wake_up_threads) {+    wake_up_threads(scheduler, nsubtasks, scheduler->num_threads);+  }++  // Join (wait for subtasks to finish)+  while(join_counter != 0) {+    if (!subtask_queue_is_empty(&worker->q)) {+      struct subtask *subtask = NULL;+      int err = subtask_queue_dequeue(worker, &subtask, 0);+      if (err == 0 ) {+        CHECK_ERR(run_subtask(worker, subtask), "run_subtask");+      }+    } else {+      if (steal_from_random_worker(worker)) {+        struct subtask *subtask = NULL;+        int err = subtask_queue_dequeue(worker, &subtask, 0);+        if (err == 0) {+          CHECK_ERR(run_subtask(worker, subtask), "run_subtask");+        }+      }+    }+  }+++  if (info.wake_up_threads || sched == DYNAMIC) {+    __atomic_sub_fetch(&active_work, nsubtasks, __ATOMIC_RELAXED);+  }++  // Write back timing results of all sequential work+  (*timer) += task_timer;+  return scheduler_error;+}+++static inline int scheduler_execute_task(struct scheduler *scheduler,+                                         struct scheduler_parloop *task)+{++  struct worker *worker = worker_local;++  int err = 0;+  if (task->iterations == 0) {+    return err;+  }++  // How much sequential work was performed by the task+  int64_t task_timer = 0;++  /* Execute task sequential or parallel based on decision made earlier */+  if (task->info.nsubtasks == 1) {+    int64_t start = get_wall_time_ns();+    err = task->fn(task->args, 0, task->iterations, 0, worker->tid);+    int64_t end = get_wall_time_ns();+    task_timer = end - start;+    worker->time_spent_working += task_timer;+    // Report time measurements+    // TODO the update of both of these should really be a single atomic!!+    __atomic_fetch_add(task->info.task_time, task_timer, __ATOMIC_RELAXED);+    __atomic_fetch_add(task->info.task_iter, task->iterations, __ATOMIC_RELAXED);+  } else {+    // Add "before" time if we already are inside a task+    int64_t time_before = 0;+    if (worker->nested > 0) {+      time_before = total_now(worker->total, worker->timer);+    }++    err = scheduler_execute_parloop(scheduler, task, &task_timer);++    // Report time measurements+    // TODO the update of both of these should really be a single atomic!!+    __atomic_fetch_add(task->info.task_time, task_timer, __ATOMIC_RELAXED);+    __atomic_fetch_add(task->info.task_iter, task->iterations, __ATOMIC_RELAXED);++    // Update timers to account for new timings+    worker->total = time_before + task_timer;+    worker->timer = get_wall_time_ns();+  }+++  return err;+}++/* Decide on how schedule the incoming task i.e. how many subtasks and+   to run sequential or (potentially nested) parallel code body */+static inline int scheduler_prepare_task(struct scheduler* scheduler,+                                         struct scheduler_segop *task)+{+  assert(task != NULL);++  struct worker *worker = worker_local;+  struct scheduler_info info;+  info.task_time = task->task_time;+  info.task_iter = task->task_iter;++  int nsubtasks;+  // Decide if task should be scheduled sequentially+  if (is_small(task, scheduler->num_threads, &nsubtasks)) {+    info.iter_pr_subtask = task->iterations;+    info.remainder = 0;+    info.nsubtasks = nsubtasks;+    return task->top_level_fn(task->args, task->iterations, worker->tid, info);+  } else {+    info.iter_pr_subtask = task->iterations / nsubtasks;+    info.remainder = task->iterations % nsubtasks;+    info.sched = task->sched;+    switch (task->sched) {+    case STATIC:+      info.nsubtasks = info.iter_pr_subtask == 0 ? info.remainder : ((task->iterations - info.remainder) / info.iter_pr_subtask);+      break;+    case DYNAMIC:+      // As any thread can take any subtasks, we are being safe with using+      // an upper bound on the number of tasks such that the task allocate enough memory+      info.nsubtasks = info.iter_pr_subtask == 0 ? info.remainder : nsubtasks;+      break;+    default:+      assert(!"Got unknown scheduling");+    }+  }++  info.wake_up_threads = 0;+  // We only use the nested parallel segop function if we can't exchaust all cores+  // using the outer most level+  if (task->nested_fn != NULL && info.nsubtasks < scheduler->num_threads && info.nsubtasks == task->iterations) {+    if (worker->nested == 0)+      info.wake_up_threads = 1;+    return task->nested_fn(task->args, task->iterations, worker->tid, info);+  }++  return task->top_level_fn(task->args, task->iterations, worker->tid, info);+}++#endif+// End of scheduler.h

@@ -0,0 +1,244 @@+// start of scheduler_common.h++#ifndef _SCHEDULER_COMMON_H_+#define _SCHEDULER_COMMON_H_++#include <float.h>++/* Scheduler definitions */+enum scheduling {+  DYNAMIC,+  STATIC+};++/* How a given task should be executed */+/* Filled out by the scheduler+   and passed to the segop function+*/+struct scheduler_info {+  int64_t iter_pr_subtask;+  int64_t remainder;+  int nsubtasks;+  enum scheduling sched;+  int wake_up_threads;++  int64_t *task_time;+  int64_t *task_iter;+};++struct scheduler {+  struct worker *workers;+  int num_threads;+};++/* A parallel parloop task  */+struct scheduler_parloop {+  const char* name;+  parloop_fn fn;+  void* args;+  int64_t iterations;+  struct scheduler_info info;+};+++/* A task for the scheduler to execute */+struct scheduler_segop {+  void *args;+  segop_fn top_level_fn;+  segop_fn nested_fn;+  int64_t iterations;+  enum scheduling sched;++  // Pointers to timer and iter associated with the task+  int64_t *task_time;+  int64_t *task_iter;++  // For debugging+  const char* name;+};++// If there is work to steal => active_work > 0+static volatile int active_work = 0;+// Number of alive workers+static volatile sig_atomic_t num_workers;++// Thread local variable worker struct+// Note that, accesses to tls variables are expensive+// Minimize direct references to this variable+__thread struct worker* worker_local = NULL;++/* Only one error can be returned at the time now+   Maybe we can provide a stack like structure for pushing errors onto+   if we wish to backpropagte multiple errors */+static volatile sig_atomic_t scheduler_error = 0;++// kappa time unit in nanoseconds+static double kappa = 5.1f * 1000;++int64_t total_now(int64_t total, int64_t time) {+  return total + (get_wall_time_ns() - time);+}++int random_other_worker(struct scheduler *scheduler, int my_id) {+  (void)scheduler;+  int my_num_workers = __atomic_load_n(&num_workers, __ATOMIC_RELAXED);+  assert(my_num_workers != 1);+  int i = fast_rand() % (my_num_workers - 1);+  if (i >= my_id) {+    i++;+  }+#ifdef MCDEBUG+  assert(i >= 0);+  assert(i < my_num_workers);+  assert(i != my_id);+#endif++  return i;+}+++static inline int64_t compute_chunk_size(struct subtask* subtask)+{+  double C = (double)*subtask->task_time / (double)*subtask->task_iter;+  if (C == 0.0F) C += DBL_EPSILON;+  return max_int64((int64_t)(kappa / C), 1);+}++/* Takes a chunk from subtask and enqueues the remaining iterations onto the worker's queue */+/* A no-op if the subtask is not chunkable */+static inline struct subtask* chunk_subtask(struct worker* worker, struct subtask *subtask)+{+  if (subtask->chunkable) {+    // Do we have information from previous runs avaliable+    if (*subtask->task_iter > 0) {+      subtask->chunk_size = compute_chunk_size(subtask);+      assert(subtask->chunk_size > 0);+    }+    int64_t remaining_iter = subtask->end - subtask->start;+    assert(remaining_iter > 0);+    if (remaining_iter > subtask->chunk_size) {+      struct subtask *new_subtask = malloc(sizeof(struct subtask));+      *new_subtask = *subtask;+      // increment the subtask join counter to account for new subtask+      __atomic_fetch_add(subtask->counter, 1, __ATOMIC_RELAXED);+      // Update range parameters+      subtask->end = subtask->start + subtask->chunk_size;+      new_subtask->start = subtask->end;+      subtask_queue_enqueue(worker, new_subtask);+    }+  }+  return subtask;+}++static inline int run_subtask(struct worker* worker, struct subtask* subtask)+{+  assert(subtask != NULL);+  assert(worker != NULL);++  subtask = chunk_subtask(worker, subtask);+  worker->total = 0;+  worker->timer = get_wall_time_ns();+#if defined(MCPROFILE)+  int64_t start = worker->timer;+#endif+  worker->nested++;+  int err = subtask->fn(subtask->args, subtask->start, subtask->end,+                        subtask->chunkable ? worker->tid : subtask->id,+                        worker->tid);+  worker->nested--;+  // Some error occured during some other subtask+  // so we just clean-up and return+  if (scheduler_error != 0) {+    // Even a failed task counts as finished.+    __atomic_fetch_sub(subtask->counter, 1, __ATOMIC_RELAXED);+    free(subtask);+    return 0;+  }+  if (err != 0) {+    __atomic_store_n(&scheduler_error, err, __ATOMIC_RELAXED);+  }+  // Total sequential time spent+  int64_t time_elapsed = total_now(worker->total, worker->timer);+#if defined(MCPROFILE)+  worker->time_spent_working += get_wall_time_ns() - start;+#endif+  int64_t iter = subtask->end - subtask->start;+  // report measurements+  // These updates should really be done using a single atomic CAS operation+  __atomic_fetch_add(subtask->task_time, time_elapsed, __ATOMIC_RELAXED);+  __atomic_fetch_add(subtask->task_iter, iter, __ATOMIC_RELAXED);+  // We need a fence here, since if the counter is decremented before either+  // of the two above are updated bad things can happen, e.g. if they are stack-allocated+  __atomic_thread_fence(__ATOMIC_SEQ_CST);+  __atomic_fetch_sub(subtask->counter, 1, __ATOMIC_RELAXED);+  free(subtask);+  return 0;+}+++static inline int is_small(struct scheduler_segop *task, int nthreads, int *nsubtasks)+{+  int64_t time = *task->task_time;+  int64_t iter = *task->task_iter;++  if (task->sched == DYNAMIC || iter == 0) {+    *nsubtasks = nthreads;+    return 0;+  }++  // Estimate the constant C+  double C = (double)time / (double)iter;+  double cur_task_iter = (double) task->iterations;++  // Returns true if the task is small i.e.+  // if the number of iterations times C is smaller+  // than the overhead of subtask creation+  if (C == 0.0F || C * cur_task_iter < kappa) {+    *nsubtasks = 1;+    return 1;+  }++  // Else compute how many subtasks this tasks should create+  int64_t min_iter_pr_subtask = max_int64((int64_t)(kappa / C), 1);+  *nsubtasks = (int)min_int64(max_int64(task->iterations / min_iter_pr_subtask, 1), nthreads);++  return 0;+}++// TODO make this prettier+static inline struct subtask* create_subtask(parloop_fn fn,+                                             void* args,+                                             const char* name,+                                             volatile int* counter,+                                             int64_t *timer,+                                             int64_t *iter,+                                             int64_t start, int64_t end,+                                             int chunkable,+                                             int64_t chunk_size,+                                             int id)+{+  struct subtask* subtask = malloc(sizeof(struct subtask));+  if (subtask == NULL) {+    assert(!"malloc failed in create_subtask");+    return NULL;+  }+  subtask->fn         = fn;+  subtask->args       = args;++  subtask->counter    = counter;+  subtask->task_time  = timer;+  subtask->task_iter  = iter;++  subtask->start      = start;+  subtask->end        = end;+  subtask->id         = id;+  subtask->chunkable  = chunkable;+  subtask->chunk_size = chunk_size;++  subtask->name       = name;+  return subtask;+}+++#endif+// end of scheduler_common.h

@@ -0,0 +1,127 @@+/* The self-tuning program to estimate $\kappa$ */++struct futhark_mc_segred_stage_1_struct {+  struct futhark_context *ctx;+  int32_t *free_tuning_res;+  int32_t *array;+};++/* Reduction function over an integer array */+int futhark_mc_tuning_segred_stage_1(void *args, int64_t start, int64_t end,+                                     int flat_tid, int tid) {+  (void)flat_tid;+  (void)tid;++  int err = 0;+  struct futhark_mc_segred_stage_1_struct *futhark_mc_segred_stage_1_struct = (struct futhark_mc_segred_stage_1_struct *) args;+  struct futhark_context *ctx = futhark_mc_segred_stage_1_struct->ctx;+  int32_t *array = futhark_mc_segred_stage_1_struct->array;+  int32_t *tuning_res = futhark_mc_segred_stage_1_struct->free_tuning_res;++  int32_t sum = 0;+  for (int i = start; i < end; i++) {+    int32_t y = array[i];+    sum = add32(sum, y);+  }+  *tuning_res = sum;+  return err;+}++/* The main entry point for the tuning process */+/* Sets the global variable ``kappa`` */+int futhark_segred_tuning_program(struct futhark_context *ctx)+{+  int err = 0;++  int64_t iterations = 100000000;+  int64_t tuning_time = 0;+  int64_t tuning_iter = 0;++  int32_t *array = malloc(sizeof(int32_t) * iterations);+  for (int64_t i = 0; i < iterations; i++) {+    array[i] = fast_rand();+  }++  int64_t start_tuning = get_wall_time_ns();+  /* **************************** */+  /* Run sequential reduce first' */+  /* **************************** */+  int64_t tuning_sequentiual_start = get_wall_time_ns();+  struct futhark_mc_segred_stage_1_struct futhark_mc_segred_stage_1_struct;+  int32_t tuning_res;+  futhark_mc_segred_stage_1_struct.ctx = ctx;+  futhark_mc_segred_stage_1_struct.free_tuning_res = &tuning_res;+  futhark_mc_segred_stage_1_struct.array = array;++  err = futhark_mc_tuning_segred_stage_1(&futhark_mc_segred_stage_1_struct, 0, iterations, 0, 0);+  int64_t tuning_sequentiual_end = get_wall_time_ns();+  int64_t sequential_elapsed = tuning_sequentiual_end - tuning_sequentiual_start;++  double C = (double)sequential_elapsed / (double)iterations;+  fprintf(stderr, " Time for sequential run is %lld - Found C %f\n", (long long)sequential_elapsed, C);++  /* ********************** */+  /* Now run tuning process */+  /* ********************** */+  // Setup a scheduler with a single worker+  int num_threads = ctx->scheduler.num_threads;+  ctx->scheduler.num_threads = 1;+  ctx->scheduler.workers = malloc(sizeof(struct worker));+  worker_local = &ctx->scheduler.workers[0];+  worker_local->tid = 0;+  CHECK_ERR(subtask_queue_init(&ctx->scheduler.workers[0].q, 1024), "failed to init queue for worker %d\n", 0);++  // Start tuning for kappa+  double kappa_tune = 1000; // Initial kappa is 1 us+  double ratio;+  int64_t time_elapsed;+  while(1) {+    int64_t min_iter_pr_subtask = (int64_t) (kappa_tune / C) == 0 ? 1 : (kappa_tune / C);+    int nsubtasks = iterations / min_iter_pr_subtask;+    struct scheduler_info info;+    info.iter_pr_subtask = min_iter_pr_subtask;++    info.nsubtasks = iterations / min_iter_pr_subtask;+    info.remainder = iterations % min_iter_pr_subtask;+    info.task_time = &tuning_time;+    info.task_iter = &tuning_iter;+    info.sched = STATIC;++    struct scheduler_parloop futhark_segred_tuning_scheduler_parloop;+    futhark_segred_tuning_scheduler_parloop.name = "futhark_mc_tuning_segred_stage_1";+    futhark_segred_tuning_scheduler_parloop.fn = futhark_mc_tuning_segred_stage_1;+    futhark_segred_tuning_scheduler_parloop.args = &futhark_mc_segred_stage_1_struct;+    futhark_segred_tuning_scheduler_parloop.iterations = iterations;+    futhark_segred_tuning_scheduler_parloop.info = info;++    int64_t tuning_chunked_start = get_wall_time_ns();+    int futhark_segred_tuning_program_err =+      scheduler_execute_task(&ctx->scheduler,+                             &futhark_segred_tuning_scheduler_parloop);+    assert(futhark_segred_tuning_program_err == 0);+    int64_t tuning_chunked_end = get_wall_time_ns();+    time_elapsed =  tuning_chunked_end - tuning_chunked_start;++    ratio = (double)time_elapsed / (double)sequential_elapsed;+    if (ratio < 1.055) {+      break;+    }+    kappa_tune += 100; // Increase by 100 ns at the time+    fprintf(stderr, "nsubtask %d - kappa %f - ratio %f\n", nsubtasks, kappa_tune, ratio);+  }++  int64_t end_tuning = get_wall_time_ns();+  fprintf(stderr, "tuning took %lld ns and found kappa %f - time %lld - ratio %f\n",+          (long long)end_tuning - start_tuning,+          kappa_tune,+          (long long)time_elapsed,+          ratio);+  kappa = kappa_tune;++  // Clean-up+  CHECK_ERR(subtask_queue_destroy(&ctx->scheduler.workers[0].q), "failed to destroy queue");+  free(array);+  free(ctx->scheduler.workers);+  ctx->scheduler.num_threads = num_threads;+  return err;+}

@@ -0,0 +1,247 @@+// start of subtask_queue.h++#ifndef SUBTASK_QUEUE_H+#define SUBTASK_QUEUE_H++/* Doubles the size of the queue */+static inline int subtask_queue_grow_queue(struct subtask_queue *subtask_queue) {++  int new_capacity = 2 * subtask_queue->capacity;+#ifdef MCDEBUG+  fprintf(stderr, "Growing queue to %d\n", subtask_queue->capacity * 2);+#endif++  struct subtask **new_buffer = calloc(new_capacity, sizeof(struct subtask*));+  for (int i = 0; i < subtask_queue->num_used; i++) {+    new_buffer[i] = subtask_queue->buffer[(subtask_queue->first + i) % subtask_queue->capacity];+  }++  free(subtask_queue->buffer);+  subtask_queue->buffer = new_buffer;+  subtask_queue->capacity = new_capacity;+  subtask_queue->first = 0;++  return 0;+}++// Initialise a job queue with the given capacity.  The queue starts out+// empty.  Returns non-zero on error.+static inline int subtask_queue_init(struct subtask_queue *subtask_queue, int capacity)+{+  assert(subtask_queue != NULL);+  memset(subtask_queue, 0, sizeof(struct subtask_queue));++  subtask_queue->capacity = capacity;+  subtask_queue->buffer = calloc(capacity, sizeof(struct subtask*));+  if (subtask_queue->buffer == NULL) {+    return -1;+  }++  CHECK_ERRNO(pthread_mutex_init(&subtask_queue->mutex, NULL), "pthread_mutex_init");+  CHECK_ERRNO(pthread_cond_init(&subtask_queue->cond, NULL), "pthread_cond_init");++  return 0;+}++// Destroy the job queue.  Blocks until the queue is empty before it+// is destroyed.+static inline int subtask_queue_destroy(struct subtask_queue *subtask_queue)+{+  assert(subtask_queue != NULL);++  CHECK_ERR(pthread_mutex_lock(&subtask_queue->mutex), "pthread_mutex_lock");++  while (subtask_queue->num_used != 0) {+    CHECK_ERR(pthread_cond_wait(&subtask_queue->cond, &subtask_queue->mutex), "pthread_cond_wait");+  }++  // Queue is now empty.  Let's kill it!+  subtask_queue->dead = 1;+  free(subtask_queue->buffer);+  CHECK_ERR(pthread_cond_broadcast(&subtask_queue->cond), "pthread_cond_broadcast");+  CHECK_ERR(pthread_mutex_unlock(&subtask_queue->mutex), "pthread_mutex_unlock");++  return 0;+}++static inline void dump_queue(struct worker *worker)+{+  struct subtask_queue *subtask_queue = &worker->q;+  CHECK_ERR(pthread_mutex_lock(&subtask_queue->mutex), "pthread_mutex_lock");+  for (int i = 0; i < subtask_queue->num_used; i++) {+    struct subtask * subtask = subtask_queue->buffer[(subtask_queue->first + i) % subtask_queue->capacity];+    printf("queue tid %d with %d task %s\n", worker->tid, i, subtask->name);+  }+  CHECK_ERR(pthread_cond_broadcast(&subtask_queue->cond), "pthread_cond_broadcast");+  CHECK_ERR(pthread_mutex_unlock(&subtask_queue->mutex), "pthread_mutex_unlock");+}++// Push an element onto the end of the job queue.  Blocks if the+// subtask_queue is full (its size is equal to its capacity).  Returns+// non-zero on error.  It is an error to push a job onto a queue that+// has been destroyed.+static inline int subtask_queue_enqueue(struct worker *worker, struct subtask *subtask )+{+  assert(worker != NULL);+  struct subtask_queue *subtask_queue = &worker->q;++#ifdef MCPROFILE+  uint64_t start = get_wall_time();+#endif++  CHECK_ERR(pthread_mutex_lock(&subtask_queue->mutex), "pthread_mutex_lock");+  // Wait until there is room in the subtask_queue.+  while (subtask_queue->num_used == subtask_queue->capacity && !subtask_queue->dead) {+    if (subtask_queue->num_used == subtask_queue->capacity) {+      CHECK_ERR(subtask_queue_grow_queue(subtask_queue), "subtask_queue_grow_queue");+      continue;+    }+    CHECK_ERR(pthread_cond_wait(&subtask_queue->cond, &subtask_queue->mutex), "pthread_cond_wait");+  }++  if (subtask_queue->dead) {+    CHECK_ERR(pthread_mutex_unlock(&subtask_queue->mutex), "pthread_mutex_unlock");+    return -1;+  }++  // If we made it past the loop, there is room in the subtask_queue.+  subtask_queue->buffer[(subtask_queue->first + subtask_queue->num_used) % subtask_queue->capacity] = subtask;+  subtask_queue->num_used++;++#ifdef MCPROFILE+  uint64_t end = get_wall_time();+  subtask_queue->time_enqueue += (end - start);+  subtask_queue->n_enqueues++;+#endif+  // Broadcast a reader (if any) that there is now an element.+  CHECK_ERR(pthread_cond_broadcast(&subtask_queue->cond), "pthread_cond_broadcast");+  CHECK_ERR(pthread_mutex_unlock(&subtask_queue->mutex), "pthread_mutex_unlock");++  return 0;+}+++/* Like subtask_queue_dequeue, but with two differences:+   1) the subtask is stolen from the __front__ of the queue+   2) returns immediately if there is no subtasks queued,+      as we dont' want to block on another workers queue and+*/+static inline int subtask_queue_steal(struct worker *worker,+                                      struct subtask **subtask)+{+  struct subtask_queue *subtask_queue = &worker->q;+  assert(subtask_queue != NULL);++#ifdef MCPROFILE+  uint64_t start = get_wall_time();+#endif+  CHECK_ERR(pthread_mutex_lock(&subtask_queue->mutex), "pthread_mutex_lock");++  if (subtask_queue->num_used == 0) {+    CHECK_ERR(pthread_cond_broadcast(&subtask_queue->cond), "pthread_cond_broadcast");+    CHECK_ERR(pthread_mutex_unlock(&subtask_queue->mutex), "pthread_mutex_unlock");+    return 1;+  }++  if (subtask_queue->dead) {+    CHECK_ERR(pthread_mutex_unlock(&subtask_queue->mutex), "pthread_mutex_unlock");+    return -1;+  }++  // Tasks gets stolen from the "front"+  struct subtask *cur_back = subtask_queue->buffer[subtask_queue->first];+  struct subtask *new_subtask = NULL;+  int remaining_iter = cur_back->end - cur_back->start;+  // If subtask is chunkable, we steal half of the iterations+  if (cur_back->chunkable && remaining_iter > 1) {+      int64_t half = remaining_iter / 2;+      new_subtask = malloc(sizeof(struct subtask));+      *new_subtask = *cur_back;+      new_subtask->start = cur_back->end - half;+      cur_back->end = new_subtask->start;+      __atomic_fetch_add(cur_back->counter, 1, __ATOMIC_RELAXED);+  } else {+    new_subtask = cur_back;+    subtask_queue->num_used--;+    subtask_queue->first = (subtask_queue->first + 1) % subtask_queue->capacity;+  }+  *subtask = new_subtask;++  if (*subtask == NULL) {+    CHECK_ERR(pthread_mutex_unlock(&subtask_queue->mutex), "pthred_mutex_unlock");+    return 1;+  }++#ifdef MCPROFILE+  uint64_t end = get_wall_time();+  subtask_queue->time_dequeue += (end - start);+  subtask_queue->n_dequeues++;+#endif++  // Broadcast a writer (if any) that there is now room for more.+  CHECK_ERR(pthread_cond_broadcast(&subtask_queue->cond), "pthread_cond_broadcast");+  CHECK_ERR(pthread_mutex_unlock(&subtask_queue->mutex), "pthread_mutex_unlock");++  return 0;+}+++// Pop an element from the back of the job queue.+// Optional argument can be provided to block or not+static inline int subtask_queue_dequeue(struct worker *worker,+                                        struct subtask **subtask, int blocking)+{+  assert(worker != NULL);+  struct subtask_queue *subtask_queue = &worker->q;++#ifdef MCPROFILE+  uint64_t start = get_wall_time();+#endif++  CHECK_ERR(pthread_mutex_lock(&subtask_queue->mutex), "pthread_mutex_lock");+  if (subtask_queue->num_used == 0 && !blocking) {+    CHECK_ERR(pthread_mutex_unlock(&subtask_queue->mutex), "pthread_mutex_unlock");+    return 1;+  }+  // Try to steal some work while the subtask_queue is empty+  while (subtask_queue->num_used == 0 && !subtask_queue->dead) {+    pthread_cond_wait(&subtask_queue->cond, &subtask_queue->mutex);+  }++  if (subtask_queue->dead) {+    CHECK_ERR(pthread_mutex_unlock(&subtask_queue->mutex), "pthread_mutex_unlock");+    return -1;+  }++  // dequeue pops from the back+  *subtask = subtask_queue->buffer[(subtask_queue->first + subtask_queue->num_used - 1) % subtask_queue->capacity];+  subtask_queue->num_used--;++  if (*subtask == NULL) {+    assert(!"got NULL ptr");+    CHECK_ERR(pthread_mutex_unlock(&subtask_queue->mutex), "pthred_mutex_unlock");+    return -1;+  }++#ifdef MCPROFILE+  uint64_t end = get_wall_time();+  subtask_queue->time_dequeue += (end - start);+  subtask_queue->n_dequeues++;+#endif++  // Broadcast a writer (if any) that there is now room for more.+  CHECK_ERR(pthread_cond_broadcast(&subtask_queue->cond), "pthread_cond_broadcast");+  CHECK_ERR(pthread_mutex_unlock(&subtask_queue->mutex), "pthread_mutex_unlock");++  return 0;+}++static inline int subtask_queue_is_empty(struct subtask_queue *subtask_queue)+{+  return subtask_queue->num_used == 0;+}+++#endif++// End of subtask_queue.h

@@ -26,6 +26,30 @@   return time.tv_sec * 1000000 + time.tv_usec; } +static int64_t get_wall_time_ns(void) {+  struct timespec time;+  assert(clock_gettime(CLOCK_REALTIME, &time) == 0);+  return time.tv_sec * 1000000000 + time.tv_nsec;+}+++static inline uint64_t rdtsc() {+  unsigned int hi, lo;+  __asm__ __volatile__("rdtsc" : "=a"(lo), "=d"(hi));+  return  ((uint64_t) lo) | (((uint64_t) hi) << 32);+}++static inline void rdtsc_wait(uint64_t n) {+  const uint64_t start = rdtsc();+  while (rdtsc() < (start + n)) {+    __asm__("PAUSE");+  }+}+static inline void spin_for(uint64_t nb_cycles) {+  rdtsc_wait(nb_cycles);+}++ #endif  // End of timing.h.

@@ -2,6 +2,9 @@ // // Various helper functions that are useful in all generated C code. +#include <errno.h>+#include <string.h>+ static const char *fut_progname = "(embedded Futhark)";  static void futhark_panic(int eval, const char *fmt, ...) {@@ -24,6 +27,24 @@   vsnprintf(buffer, needed, s, vl);   return buffer; }+++static inline void check_err(int errval, int sets_errno, const char *fun, int line,+                            const char *msg, ...)+{+  if (errval) {+    char str[256];+    char errnum[10];+    sprintf(errnum, "%d", errval);+    sprintf(str, "ERROR: %s in %s() at line %d with error code %s\n", msg, fun, line,+            sets_errno ? strerror(errno) : errnum);+    fprintf(stderr, "%s", str);+    exit(errval);+  }+}++#define CHECK_ERR(err, msg...) check_err(err, 0, __func__, __LINE__, msg)+#define CHECK_ERRNO(err, msg...) check_err(err, 1, __func__, __LINE__, msg)  // Read a file into a NUL-terminated string; returns NULL on error. static void* slurp_file(const char *filename, size_t *size) {

@@ -122,7 +122,7 @@      self.global_failure = self.pool.allocate(np.int32().itemsize)     cl.enqueue_fill_buffer(self.queue, self.global_failure, np.int32(-1), 0, np.int32().itemsize)-    self.global_failure_args = self.pool.allocate(np.int32().itemsize *+    self.global_failure_args = self.pool.allocate(np.int64().itemsize *                                                   (self.global_failure_args_max+1))     self.failure_is_an_option = np.int32(0) @@ -225,7 +225,7 @@         cl.enqueue_fill_buffer(self.queue, self.global_failure, np.int32(-1), 0, np.int32().itemsize)          # Read failure args.-        failure_args = np.empty(self.global_failure_args_max+1, dtype=np.int32)+        failure_args = np.empty(self.global_failure_args_max+1, dtype=np.int64)         cl.enqueue_copy(self.queue, failure_args, self.global_failure_args, is_blocking=True)          raise Exception(self.failure_msgs[failure[0]].format(*failure_args))

@@ -6,10 +6,12 @@   ( printAction,     impCodeGenAction,     kernelImpCodeGenAction,+    multicoreImpCodeGenAction,     metricsAction,     compileCAction,     compileOpenCLAction,     compileCUDAAction,+    compileMulticoreAction,     sexpAction,   ) where@@ -21,12 +23,15 @@ import Futhark.Analysis.Metrics import qualified Futhark.CodeGen.Backends.CCUDA as CCUDA import qualified Futhark.CodeGen.Backends.COpenCL as COpenCL+import qualified Futhark.CodeGen.Backends.MulticoreC as MulticoreC import qualified Futhark.CodeGen.Backends.SequentialC as SequentialC import qualified Futhark.CodeGen.ImpGen.Kernels as ImpGenKernels+import qualified Futhark.CodeGen.ImpGen.Multicore as ImpGenMulticore import qualified Futhark.CodeGen.ImpGen.Sequential as ImpGenSequential import Futhark.Compiler.CLI import Futhark.IR import Futhark.IR.KernelsMem (KernelsMem)+import Futhark.IR.MCMem (MCMem) import Futhark.IR.Prop.Aliases import Futhark.IR.SeqMem (SeqMem) import Futhark.Util (runProgramWithExitCode)@@ -71,6 +76,14 @@       actionProcedure = liftIO . putStrLn . pretty . snd <=< ImpGenKernels.compileProgOpenCL     } +multicoreImpCodeGenAction :: Action MCMem+multicoreImpCodeGenAction =+  Action+    { actionName = "Compile to imperative multicore",+      actionDescription = "Translate program into imperative multicore IL and write it on standard output.",+      actionProcedure = liftIO . putStrLn . pretty . snd <=< ImpGenMulticore.compileProg+    }+ -- | Print metrics about AST node counts to stdout. sexpAction :: ASTLore lore => Action lore sexpAction =@@ -207,6 +220,46 @@                 [cpath, "-O", "-std=c99", "-lm", "-o", outpath]                   ++ extra_options           ret <- liftIO $ runProgramWithExitCode "gcc" args mempty+          case ret of+            Left err ->+              externalErrorS $ "Failed to run gcc: " ++ show err+            Right (ExitFailure code, _, gccerr) ->+              externalErrorS $+                "gcc failed with code "+                  ++ show code+                  ++ ":\n"+                  ++ gccerr+            Right (ExitSuccess, _, _) ->+              return ()++-- | The @futhark multicore@ action.+compileMulticoreAction :: FutharkConfig -> CompilerMode -> FilePath -> Action MCMem+compileMulticoreAction fcfg mode outpath =+  Action+    { actionName = "Compile to multicore",+      actionDescription = "Compile to multicore",+      actionProcedure = helper+    }+  where+    helper prog = do+      cprog <- handleWarnings fcfg $ MulticoreC.compileProg prog+      let cpath = outpath `addExtension` "c"+          hpath = outpath `addExtension` "h"++      case mode of+        ToLibrary -> do+          let (header, impl) = MulticoreC.asLibrary cprog+          liftIO $ writeFile hpath header+          liftIO $ writeFile cpath impl+        ToExecutable -> do+          liftIO $ writeFile cpath $ MulticoreC.asExecutable cprog+          -- let debug_flags = ["-g", "-fno-omit-frame-pointer", "-fsanitize=address", "-fsanitize=integer", "-fsanitize=undefined", "-fno-sanitize-recover=null"]+          ret <-+            liftIO $+              runProgramWithExitCode+                "gcc"+                [cpath, "-O3", "-pthread", "-std=c11", "-lm", "-o", outpath]+                mempty           case ret of             Left err ->               externalErrorS $ "Failed to run gcc: " ++ show err

@@ -526,7 +526,7 @@   SOAC lore ->   m (SOAC lore, [Ident]) soacToStream soac = do-  chunk_param <- newParam "chunk" $ Prim int32+  chunk_param <- newParam "chunk" $ Prim int64   let chvar = Futhark.Var $ paramName chunk_param       (lam, inps) = (lambda soac, inputs soac)       w = width soac@@ -579,7 +579,7 @@         lastel_tmp_ids <- mapM (newIdent "lstel_tmp") accrtps         empty_arr <- newIdent "empty_arr" $ Prim Bool         inpacc_ids <- mapM (newParam "inpacc") accrtps-        outszm1id <- newIdent "szm1" $ Prim int32+        outszm1id <- newIdent "szm1" $ Prim int64         -- 1. let (scan0_ids,map_resids)  = scanomap(scan_lam,nes,map_lam,a_ch)         let insbnd =               mkLet [] (scan0_ids ++ map_resids) $@@ -591,17 +591,17 @@               mkLet [] [outszm1id] $                 BasicOp $                   BinOp-                    (Sub Int32 OverflowUndef)+                    (Sub Int64 OverflowUndef)                     (Futhark.Var $ paramName chunk_param)-                    (constant (1 :: Int32))+                    (constant (1 :: Int64))             -- 3. let lasteel_ids = ...             empty_arr_bnd =               mkLet [] [empty_arr] $                 BasicOp $                   CmpOp-                    (CmpSlt Int32)+                    (CmpSlt Int64)                     (Futhark.Var $ identName outszm1id)-                    (constant (0 :: Int32))+                    (constant (0 :: Int64))             leltmpbnds =               zipWith                 ( \lid arrid ->

@@ -48,6 +48,10 @@ class OpMetrics op where   opMetrics :: op -> MetricsM () +instance OpMetrics a => OpMetrics (Maybe a) where+  opMetrics Nothing = return ()+  opMetrics (Just x) = opMetrics x+ instance OpMetrics () where   opMetrics () = return () 

@@ -7,6 +7,8 @@     primExpFromSubExp,     pe32,     le32,+    pe64,+    le64,     primExpFromSubExpM,     replaceInPrimExp,     replaceInPrimExpM,@@ -92,6 +94,14 @@ le32 :: a -> TPrimExp Int32 a le32 = isInt32 . flip LeafExp int32 +-- | Shorthand for constructing a 'TPrimExp' of type 'Int64'.+pe64 :: SubExp -> TPrimExp Int64 VName+pe64 = isInt64 . primExpFromSubExp int64++-- | Shorthand for constructing a 'TPrimExp' of type 'Int64', from a leaf.+le64 :: a -> TPrimExp Int64 a+le64 = isInt64 . flip LeafExp int64+ -- | Applying a monadic transformation to the leaves in a 'PrimExp'. replaceInPrimExpM ::   Monad m =>@@ -133,9 +143,9 @@   fromMaybe (LeafExp v t) $ M.lookup v tab  -- | Convert a 'SubExp' slice to a 'PrimExp' slice.-primExpSlice :: Slice SubExp -> Slice (TPrimExp Int32 VName)-primExpSlice = map $ fmap $ isInt32 . primExpFromSubExp int32+primExpSlice :: Slice SubExp -> Slice (TPrimExp Int64 VName)+primExpSlice = map $ fmap pe64  -- | Convert a 'PrimExp' slice to a 'SubExp' slice.-subExpSlice :: MonadBinder m => Slice (TPrimExp Int32 VName) -> m (Slice SubExp)+subExpSlice :: MonadBinder m => Slice (TPrimExp Int64 VName) -> m (Slice SubExp) subExpSlice = mapM $ traverse $ toSubExp "slice"

@@ -111,7 +111,7 @@     Indexed Certificates (PrimExp VName)   | -- | The indexing corresponds to another (perhaps more     -- advantageous) array.-    IndexedArray Certificates VName [TPrimExp Int32 VName]+    IndexedArray Certificates VName [TPrimExp Int64 VName]  indexedAddCerts :: Certificates -> Indexed -> Indexed indexedAddCerts cs1 (Indexed cs2 v) = Indexed (cs1 <> cs2) v@@ -122,7 +122,7 @@   freeIn' (IndexedArray cs arr v) = freeIn' cs <> freeIn' arr <> freeIn' v  -- | Indexing a delayed array if possible.-type IndexArray = [TPrimExp Int32 VName] -> Maybe Indexed+type IndexArray = [TPrimExp Int64 VName] -> Maybe Indexed  data Entry lore = Entry   { -- | True if consumed.@@ -265,7 +265,7 @@  index' ::   VName ->-  [TPrimExp Int32 VName] ->+  [TPrimExp Int64 VName] ->   SymbolTable lore ->   Maybe Indexed index' name is vtable = do@@ -288,7 +288,7 @@     SymbolTable lore ->     Int ->     op ->-    [TPrimExp Int32 VName] ->+    [TPrimExp Int64 VName] ->     Maybe Indexed   indexOp _ _ _ _ = Nothing @@ -322,18 +322,18 @@   | Just oldshape <- arrayDims <$> lookupType v table =     let is' =           reshapeIndex-            (map pe32 oldshape)-            (map pe32 $ newDims newshape)+            (map pe64 oldshape)+            (map pe64 $ newDims newshape)             is      in index' v is' table indexExp table (BasicOp (Index v slice)) _ is =   index' v (adjust slice is) table   where     adjust (DimFix j : js') is' =-      pe32 j : adjust js' is'+      pe64 j : adjust js' is'     adjust (DimSlice j _ s : js') (i : is') =-      let i_t_s = i * pe32 s-          j_p_i_t_s = pe32 j + i_t_s+      let i_t_s = i * pe64 s+          j_p_i_t_s = pe64 j + i_t_s        in j_p_i_t_s : adjust js' is'     adjust _ _ = [] indexExp _ _ _ _ = Nothing

@@ -18,6 +18,8 @@ import Futhark.IR (ASTLore, Op, Prog, pretty) import qualified Futhark.IR.Kernels as Kernels import qualified Futhark.IR.KernelsMem as KernelsMem+import qualified Futhark.IR.MC as MC+import qualified Futhark.IR.MCMem as MCMem import Futhark.IR.Prop.Aliases (CanBeAliased) import qualified Futhark.IR.SOACS as SOACS import qualified Futhark.IR.Seq as Seq@@ -38,6 +40,7 @@ import qualified Futhark.Pass.ExplicitAllocations.Kernels as Kernels import qualified Futhark.Pass.ExplicitAllocations.Seq as Seq import Futhark.Pass.ExtractKernels+import Futhark.Pass.ExtractMulticore import Futhark.Pass.FirstOrderTransform import Futhark.Pass.KernelBabysitting import Futhark.Pass.Simplify@@ -91,8 +94,10 @@ data UntypedPassState   = SOACS (Prog SOACS.SOACS)   | Kernels (Prog Kernels.Kernels)+  | MC (Prog MC.MC)   | Seq (Prog Seq.Seq)   | KernelsMem (Prog KernelsMem.KernelsMem)+  | MCMem (Prog MCMem.MCMem)   | SeqMem (Prog SeqMem.SeqMem)  getSOACSProg :: UntypedPassState -> Maybe (Prog SOACS.SOACS)@@ -107,15 +112,19 @@ instance Representation UntypedPassState where   representation (SOACS _) = "SOACS"   representation (Kernels _) = "Kernels"+  representation (MC _) = "MC"   representation (Seq _) = "Seq"   representation (KernelsMem _) = "KernelsMem"+  representation (MCMem _) = "MCMem"   representation (SeqMem _) = "SeqMEm"  instance PP.Pretty UntypedPassState where   ppr (SOACS prog) = PP.ppr prog   ppr (Kernels prog) = PP.ppr prog+  ppr (MC prog) = PP.ppr prog   ppr (Seq prog) = PP.ppr prog   ppr (SeqMem prog) = PP.ppr prog+  ppr (MCMem prog) = PP.ppr prog   ppr (KernelsMem prog) = PP.ppr prog  newtype UntypedPass@@ -129,6 +138,7 @@   = SOACSAction (Action SOACS.SOACS)   | KernelsAction (Action Kernels.Kernels)   | KernelsMemAction (FilePath -> Action KernelsMem.KernelsMem)+  | MCMemAction (FilePath -> Action MCMem.MCMem)   | SeqMemAction (FilePath -> Action SeqMem.SeqMem)   | PolyAction       ( forall lore.@@ -144,12 +154,14 @@ untypedActionName (KernelsAction a) = actionName a untypedActionName (SeqMemAction a) = actionName $ a "" untypedActionName (KernelsMemAction a) = actionName $ a ""+untypedActionName (MCMemAction a) = actionName $ a "" untypedActionName (PolyAction a) = actionName (a :: Action SOACS.SOACS)  instance Representation UntypedAction where   representation (SOACSAction _) = "SOACS"   representation (KernelsAction _) = "Kernels"   representation (KernelsMemAction _) = "KernelsMem"+  representation (MCMemAction _) = "MCMem"   representation (SeqMemAction _) = "SeqMem"   representation PolyAction {} = "<any>" @@ -247,12 +259,16 @@       SOACS <$> runPipeline (onePass simplifySOACS) config prog     perform (Kernels prog) config =       Kernels <$> runPipeline (onePass simplifyKernels) config prog+    perform (MC prog) config =+      MC <$> runPipeline (onePass simplifyMC) config prog     perform (Seq prog) config =       Seq <$> runPipeline (onePass simplifySeq) config prog     perform (SeqMem prog) config =       SeqMem <$> runPipeline (onePass simplifySeqMem) config prog     perform (KernelsMem prog) config =       KernelsMem <$> runPipeline (onePass simplifyKernelsMem) config prog+    perform (MCMem prog) config =+      MCMem <$> runPipeline (onePass simplifyMCMem) config prog      long = [passLongOption pass]     pass = simplifySOACS@@ -299,12 +315,16 @@       SOACS <$> runPipeline (onePass $ performCSE True) config prog     perform (Kernels prog) config =       Kernels <$> runPipeline (onePass $ performCSE True) config prog+    perform (MC prog) config =+      MC <$> runPipeline (onePass $ performCSE True) config prog     perform (Seq prog) config =       Seq <$> runPipeline (onePass $ performCSE True) config prog     perform (SeqMem prog) config =       SeqMem <$> runPipeline (onePass $ performCSE False) config prog     perform (KernelsMem prog) config =       KernelsMem <$> runPipeline (onePass $ performCSE False) config prog+    perform (MCMem prog) config =+      MCMem <$> runPipeline (onePass $ performCSE False) config prog      long = [passLongOption pass]     pass = performCSE True :: Pass SOACS.SOACS SOACS.SOACS@@ -384,6 +404,14 @@       "Translate program into the imperative IL with kernels and write it on standard output.",     Option       []+      ["compile-imperative-multicore"]+      ( NoArg $+          Right $ \opts ->+            opts {futharkAction = MCMemAction $ const multicoreImpCodeGenAction}+      )+      "Translate program into the imperative IL with kernels and write it on standard output.",+    Option+      []       ["compile-opencl"]       ( NoArg $           Right $ \opts ->@@ -443,12 +471,13 @@     soacsPassOption inlineFunctions [],     kernelsPassOption babysitKernels [],     kernelsPassOption tileLoops [],-    kernelsPassOption unstream [],-    kernelsPassOption sink [],+    kernelsPassOption unstreamKernels [],+    kernelsPassOption sinkKernels [],     typedPassOption soacsProg Kernels extractKernels [],+    typedPassOption soacsProg MC extractMulticore [],     iplOption [],     allocateOption "a",-    kernelsMemPassOption doubleBuffer [],+    kernelsMemPassOption doubleBufferKernels [],     kernelsMemPassOption expandAllocations [],     cseOption [],     simplifyOption "e",@@ -480,7 +509,15 @@       "Run the sequential CPU compilation pipeline"       sequentialCpuPipeline       []-      ["cpu"]+      ["cpu"],+    pipelineOption+      getSOACSProg+      "MCMem"+      MCMem+      "Run the multicore compilation pipeline"+      multicorePipeline+      []+      ["multicore"]   ]  incVerbosity :: Maybe FilePath -> FutharkConfig -> FutharkConfig@@ -593,15 +630,21 @@       actionProcedure (action base) prog     (KernelsMem prog, KernelsMemAction action) ->       actionProcedure (action base) prog+    (MCMem prog, MCMemAction action) ->+      actionProcedure (action base) prog     (SOACS soacs_prog, PolyAction acs) ->       actionProcedure acs soacs_prog     (Kernels kernels_prog, PolyAction acs) ->       actionProcedure acs kernels_prog+    (MC mc_prog, PolyAction acs) ->+      actionProcedure acs mc_prog     (Seq seq_prog, PolyAction acs) ->       actionProcedure acs seq_prog     (KernelsMem mem_prog, PolyAction acs) ->       actionProcedure acs mem_prog     (SeqMem mem_prog, PolyAction acs) ->+      actionProcedure acs mem_prog+    (MCMem mem_prog, PolyAction acs) ->       actionProcedure acs mem_prog     (_, action) ->       externalErrorS $

@@ -0,0 +1,17 @@+{-# LANGUAGE FlexibleContexts #-}++module Futhark.CLI.Multicore (main) where++import Futhark.Actions (compileMulticoreAction)+import Futhark.Compiler.CLI+import Futhark.Passes (multicorePipeline)++main :: String -> [String] -> IO ()+main = compilerMain+  ()+  []+  "Compile to multicore C"+  "Generate multicore C code from optimised Futhark program."+  multicorePipeline+  $ \fcfg () mode outpath prog ->+    actionProcedure (compileMulticoreAction fcfg mode outpath) prog

@@ -579,7 +579,7 @@           { configBackend = "c",             configFuthark = Nothing,             configRunner = "",-            configExtraOptions = [],+            configExtraOptions = ["-b"],             configExtraCompilerOptions = [],             configTuning = Just "tuning"           },

@@ -392,7 +392,7 @@                  CUDA_SUCCEED(cuMemAlloc(&ctx->global_failure, sizeof(no_error)));                  CUDA_SUCCEED(cuMemcpyHtoD(ctx->global_failure, &no_error, sizeof(no_error)));                  // The +1 is to avoid zero-byte allocations.-                 CUDA_SUCCEED(cuMemAlloc(&ctx->global_failure_args, sizeof(int32_t)*($int:max_failure_args+1)));+                 CUDA_SUCCEED(cuMemAlloc(&ctx->global_failure_args, sizeof(int64_t)*($int:max_failure_args+1)));                   $stms:init_kernel_fields @@ -442,7 +442,7 @@                                     &no_failure,                                     sizeof(int32_t))); -                     typename int32_t args[$int:max_failure_args+1];+                     typename int64_t args[$int:max_failure_args+1];                      CUDA_SUCCEED(                        cuMemcpyDtoH(&args,                                     ctx->global_failure_args,

@@ -41,7 +41,8 @@             escapeChar c = [c]          in concatMap escapeChar       onPart (ErrorString s) = printfEscape s-      onPart ErrorInt32 {} = "%d"+      onPart ErrorInt32 {} = "%lld"+      onPart ErrorInt64 {} = "%lld"       onFailure i (FailureMsg emsg@(ErrorMsg parts) backtrace) =         let msg = concatMap onPart parts ++ "\n" ++ printfEscape backtrace             msgargs = [[C.cexp|args[$int:j]|] | j <- [0 .. errorMsgNumArgs emsg -1]]@@ -375,7 +376,7 @@                      ctx->global_failure_args =                        clCreateBuffer(ctx->opencl.ctx,                                       CL_MEM_READ_WRITE,-                                      sizeof(cl_int)*($int:max_failure_args+1), NULL, &error);+                                      sizeof(int64_t)*($int:max_failure_args+1), NULL, &error);                      OPENCL_SUCCEED_OR_RETURN(error);                       // Load all the kernels.@@ -472,7 +473,7 @@                                          0, sizeof(cl_int), &no_failure,                                          0, NULL, NULL)); -                   typename cl_int args[$int:max_failure_args+1];+                   typename int64_t args[$int:max_failure_args+1];                    OPENCL_SUCCEED_OR_RETURN(                      clEnqueueReadBuffer(ctx->opencl.queue,                                          ctx->global_failure_args,

@@ -33,12 +33,13 @@      -- * Monadic compiler interface     CompilerM,-    CompilerState (compUserState),+    CompilerState (compUserState, compNameSrc),     getUserState,     modifyUserState,     contextContents,     contextFinalInits,     runCompilerM,+    inNewFunction,     cachingMemory,     blockScope,     compileFun,@@ -64,9 +65,16 @@     publicName,     contextType,     contextField,+    memToCType,+    cacheMem,+    fatMemory,+    rawMemCType,+    cproduct,+    fatMemType,      -- * Building Blocks     primTypeToCType,+    intTypeToCType,     copyMemoryDefaultSpace,   ) where@@ -209,6 +217,7 @@   free_all_mem <- collect $ mapM_ (uncurry unRefMem) =<< gets compDeclaredMem   let onPart (ErrorString s) = return ("%s", [C.cexp|$string:s|])       onPart (ErrorInt32 x) = ("%d",) <$> compileExp x+      onPart (ErrorInt64 x) = ("%lld",) <$> compileExp x   (formatstrs, formatargs) <- unzip <$> mapM onPart parts   let formatstr = "Error: " ++ concat formatstrs ++ "\n\nBacktrace:\n%s"   items@@ -403,6 +412,22 @@       const w {accItems = mempty}     ) +-- | Used when we, inside an existing 'CompilerM' action, want to+-- generate code for a new function.  Use this so that the compiler+-- understands that previously declared memory doesn't need to be+-- freed inside this action.+inNewFunction :: Bool -> CompilerM op s a -> CompilerM op s a+inNewFunction keep_cached m = do+  old_mem <- gets compDeclaredMem+  modify $ \s -> s {compDeclaredMem = mempty}+  x <- local noCached m+  modify $ \s -> s {compDeclaredMem = old_mem}+  return x+  where+    noCached env+      | keep_cached = env+      | otherwise = env {envCachedMem = mempty}+ item :: C.BlockItem -> CompilerM op s () item x = tell $ mempty {accItems = DL.singleton x} @@ -675,7 +700,9 @@   let setdef =         [C.cedecl|static int $id:(fatMemSet space) ($ty:ctx_ty *ctx, $ty:mty *lhs, $ty:mty *rhs, const char *lhs_desc) {   int ret = $id:(fatMemUnRef space)(ctx, lhs, lhs_desc);-  (*(rhs->references))++;+  if (rhs->references != NULL) {+    (*(rhs->references))++;+  }   *lhs = *rhs;   return ret; }@@ -1617,6 +1644,11 @@   let headerdefs =         [C.cunit| $esc:("// Headers\n")+/* We need to define _GNU_SOURCE before+   _any_ headers files are imported to get+   the usage statistics of a thread (i.e. have RUSAGE_THREAD) on GNU/Linux+   https://manpages.courier-mta.org/htmlman2/getrusage.2.html */+$esc:("#define _GNU_SOURCE") $esc:("#include <stdint.h>") $esc:("#include <stddef.h>") $esc:("#include <stdbool.h>")@@ -2097,7 +2129,7 @@       iexp' <- compileExp $ untyped iexp       return [C.cexp|$id:src[$exp:iexp']|]     compileLeaf (SizeOf t) =-      return [C.cexp|(typename int64_t)sizeof($ty:t')|]+      return [C.cexp|(typename int32_t)sizeof($ty:t')|]       where         t' = primTypeToCType t 

@@ -1132,6 +1132,7 @@   e' <- compileExp e   let onPart (Imp.ErrorString s) = return ("%s", String s)       onPart (Imp.ErrorInt32 x) = ("%d",) <$> compileExp x+      onPart (Imp.ErrorInt64 x) = ("%d",) <$> compileExp x   (formatstrs, formatargs) <- unzip <$> mapM onPart parts   stm $     Assert

@@ -0,0 +1,703 @@+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE QuasiQuotes #-}+{-# LANGUAGE TemplateHaskell #-}++-- | C code generator.  This module can convert a correct ImpCode+-- program to an equivalent C program.+module Futhark.CodeGen.Backends.MulticoreC+  ( compileProg,+    GC.CParts (..),+    GC.asLibrary,+    GC.asExecutable,+  )+where++import Control.Monad+import Data.FileEmbed+import qualified Data.Map as M+import Data.Maybe+import qualified Futhark.CodeGen.Backends.GenericC as GC+import Futhark.CodeGen.Backends.GenericC.Options+import Futhark.CodeGen.Backends.SimpleRep+import Futhark.CodeGen.ImpCode.Multicore+import qualified Futhark.CodeGen.ImpGen.Multicore as ImpGen+import Futhark.IR.MCMem (MCMem, Prog)+import Futhark.MonadFreshNames+import qualified Language.C.Quote.OpenCL as C+import qualified Language.C.Syntax as C++compileProg ::+  MonadFreshNames m =>+  Prog MCMem ->+  m (ImpGen.Warnings, GC.CParts)+compileProg =+  traverse+    ( GC.compileProg+        "multicore"+        operations+        generateContext+        ""+        [DefaultSpace]+        cliOptions+    )+    <=< ImpGen.compileProg+  where+    generateContext = do+      let multicore_defs_h = $(embedStringFile "rts/c/multicore_defs.h")+          multicore_util_h = $(embedStringFile "rts/c/multicore_util.h")+          subtask_queue_h = $(embedStringFile "rts/c/subtask_queue.h")+          scheduler_common_h = $(embedStringFile "rts/c/scheduler_common.h")+          scheduler_h = $(embedStringFile "rts/c/scheduler.h")+          scheduler_tune_h = $(embedStringFile "rts/c/scheduler_tune.h")++      mapM_+        GC.earlyDecl+        [C.cunit|+                              $esc:multicore_defs_h+                              $esc:multicore_util_h+                              $esc:subtask_queue_h+                              $esc:scheduler_common_h+                              $esc:scheduler_h+                             |]++      mapM_ GC.earlyDecl [C.cunit|int futhark_segred_tuning_program(struct futhark_context *ctx);|]+      mapM_ GC.libDecl [C.cunit|$esc:scheduler_tune_h|]++      cfg <- GC.publicDef "context_config" GC.InitDecl $ \s ->+        ( [C.cedecl|struct $id:s;|],+          [C.cedecl|struct $id:s { int debugging; int profiling; };|]+        )++      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 = (struct $id:cfg*) malloc(sizeof(struct $id:cfg));+                                 if (cfg == NULL) {+                                   return NULL;+                                 }+                                 cfg->debugging = 0;+                                 cfg->profiling = 0;+                                 return cfg;+                               }|]+        )++      GC.publicDef_ "context_config_free" GC.InitDecl $ \s ->+        ( [C.cedecl|void $id:s(struct $id:cfg* cfg);|],+          [C.cedecl|void $id:s(struct $id:cfg* cfg) {+                                 free(cfg);+                               }|]+        )++      GC.publicDef_ "context_config_set_debugging" GC.InitDecl $ \s ->+        ( [C.cedecl|void $id:s(struct $id:cfg* cfg, int flag);|],+          [C.cedecl|void $id:s(struct $id:cfg* cfg, int detail) {+                          cfg->debugging = detail;+                        }|]+        )++      GC.publicDef_ "context_config_set_profiling" GC.InitDecl $ \s ->+        ( [C.cedecl|void $id:s(struct $id:cfg* cfg, int flag);|],+          [C.cedecl|void $id:s(struct $id:cfg* cfg, int flag) {+                          cfg->profiling = flag;+                        }|]+        )++      GC.publicDef_ "context_config_set_logging" GC.InitDecl $ \s ->+        ( [C.cedecl|void $id:s(struct $id:cfg* cfg, int flag);|],+          [C.cedecl|void $id:s(struct $id:cfg* cfg, int detail) {+                                 /* Does nothing for this backend. */+                                 (void)cfg; (void)detail;+                               }|]+        )++      (fields, init_fields) <- GC.contextContents++      ctx <- GC.publicDef "context" GC.InitDecl $ \s ->+        ( [C.cedecl|struct $id:s;|],+          [C.cedecl|struct $id:s {+                          struct scheduler scheduler;+                          int detail_memory;+                          int debugging;+                          int profiling;+                          int profiling_paused;+                          typename lock_t lock;+                          char *error;+                          int total_runs;+                          long int total_runtime;+                          $sdecls:fields++                          // Tuning parameters+                          typename int64_t tuning_timing;+                          typename int64_t tuning_iter;+                        };|]+        )++      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 = (struct $id:ctx*) malloc(sizeof(struct $id:ctx));+                 if (ctx == NULL) {+                   return NULL;+                 }++                 // Initialize rand()+                 fast_srand(time(0));+                 ctx->detail_memory = cfg->debugging;+                 ctx->debugging = cfg->debugging;+                 ctx->profiling = cfg->profiling;+                 ctx->profiling_paused = 0;+                 ctx->error = NULL;+                 create_lock(&ctx->lock);+                 ctx->scheduler.num_threads = num_processors();+                 if (ctx->scheduler.num_threads < 1) return NULL;++                 $stms:init_fields++                 // futhark_segred_tuning_program(ctx);++                 ctx->scheduler.workers = calloc(ctx->scheduler.num_threads, sizeof(struct worker));+                 if (ctx->scheduler.workers == NULL) return NULL;+                 num_workers = ctx->scheduler.num_threads;++                 worker_local = &ctx->scheduler.workers[0];+                 worker_local->tid = 0;+                 worker_local->scheduler = &ctx->scheduler;+                 CHECK_ERR(subtask_queue_init(&worker_local->q, 1024), "failed to init queue for worker %d\n", 0);+++                 for (int i = 1; i < ctx->scheduler.num_threads; i++) {+                   struct worker *cur_worker = &ctx->scheduler.workers[i];+                   memset(cur_worker, 0, sizeof(struct worker));+                   cur_worker->tid = i;+                   cur_worker->output_usage = 0;+                   cur_worker->scheduler = &ctx->scheduler;+                   CHECK_ERR(subtask_queue_init(&cur_worker->q, 1024), "failed to init queue for worker %d\n", i);++                   CHECK_ERR(pthread_create(&cur_worker->thread, NULL, &scheduler_worker,+                                            cur_worker),+                             "Failed to create worker %d\n", i);+                 }++                 init_constants(ctx);+++                 return ctx;+              }|]+        )++      GC.publicDef_ "context_free" GC.InitDecl $ \s ->+        ( [C.cedecl|void $id:s(struct $id:ctx* ctx);|],+          [C.cedecl|void $id:s(struct $id:ctx* ctx) {+                 free_constants(ctx);++                 // output_thread_usage(worker_local);+                 for (int i = 1; i < ctx->scheduler.num_threads; i++)+                 {+                   struct worker *cur_worker = &ctx->scheduler.workers[i];+                   cur_worker->dead = 1;+                   subtask_queue_destroy(&cur_worker->q);+                   CHECK_ERR(pthread_join(ctx->scheduler.workers[i].thread, NULL), "pthread_join");+                 }+++                 free(ctx->scheduler.workers);+                 free_lock(&ctx->lock);+                 free(ctx);+               }|]+        )++      GC.publicDef_ "context_sync" GC.InitDecl $ \s ->+        ( [C.cedecl|int $id:s(struct $id:ctx* ctx);|],+          [C.cedecl|int $id:s(struct $id:ctx* ctx) {+                                 (void)ctx;+                                 return 0;+                               }|]+        )++cliOptions :: [Option]+cliOptions =+  [ Option+      { optionLongName = "profile",+        optionShortName = Just 'P',+        optionArgument = NoArgument,+        optionAction = [C.cstm|futhark_context_config_set_profiling(cfg, 1);|],+        optionDescription = "Gather profiling information."+      }+  ]++operations :: GC.Operations Multicore ()+operations =+  GC.defaultOperations+    { GC.opsCompiler = compileOp,+      GC.opsCopy = copyMulticoreMemory,+      GC.opsCritical =+        -- The thread entering an API function is always considered+        -- the "first worker" - note that this might differ from the+        -- thread that created the context!  This likely only matters+        -- for entry points, since they are the only API functions+        -- that contain parallel operations.+        ( [C.citems|worker_local = &ctx->scheduler.workers[0];|],+          []+        )+    }++copyMulticoreMemory :: GC.Copy Multicore ()+copyMulticoreMemory destmem destidx DefaultSpace srcmem srcidx DefaultSpace nbytes =+  GC.copyMemoryDefaultSpace destmem destidx srcmem srcidx nbytes+copyMulticoreMemory _ _ destspace _ _ srcspace _ =+  error $ "Cannot copy to " ++ show destspace ++ " from " ++ show srcspace++closureFreeStructField :: VName -> Name+closureFreeStructField v =+  nameFromString "free_" <> nameFromString (pretty v)++closureRetvalStructField :: VName -> Name+closureRetvalStructField v =+  nameFromString "retval_" <> nameFromString (pretty v)++data ValueType = Prim | MemBlock | RawMem++compileFreeStructFields :: [VName] -> [(C.Type, ValueType)] -> [C.FieldGroup]+compileFreeStructFields = zipWith field+  where+    field name (ty, Prim) =+      [C.csdecl|$ty:ty $id:(closureFreeStructField name);|]+    field name (_, _) =+      [C.csdecl|$ty:defaultMemBlockType $id:(closureFreeStructField name);|]++compileRetvalStructFields :: [VName] -> [(C.Type, ValueType)] -> [C.FieldGroup]+compileRetvalStructFields = zipWith field+  where+    field name (ty, Prim) =+      [C.csdecl|$ty:ty *$id:(closureRetvalStructField name);|]+    field name (_, _) =+      [C.csdecl|$ty:defaultMemBlockType $id:(closureRetvalStructField name);|]++compileSetStructValues ::+  C.ToIdent a =>+  a ->+  [VName] ->+  [(C.Type, ValueType)] ->+  [C.Stm]+compileSetStructValues struct = zipWith field+  where+    field name (_, Prim) =+      [C.cstm|$id:struct.$id:(closureFreeStructField name)=$id:name;|]+    field name (_, MemBlock) =+      [C.cstm|$id:struct.$id:(closureFreeStructField name)=$id:name.mem;|]+    field name (_, RawMem) =+      [C.cstm|$id:struct.$id:(closureFreeStructField name)=$id:name;|]++compileSetRetvalStructValues ::+  C.ToIdent a =>+  a ->+  [VName] ->+  [(C.Type, ValueType)] ->+  [C.Stm]+compileSetRetvalStructValues struct = zipWith field+  where+    field name (_, Prim) =+      [C.cstm|$id:struct.$id:(closureRetvalStructField name)=&$id:name;|]+    field name (_, MemBlock) =+      [C.cstm|$id:struct.$id:(closureRetvalStructField name)=$id:name.mem;|]+    field name (_, RawMem) =+      [C.cstm|$id:struct.$id:(closureRetvalStructField name)=$id:name;|]++compileGetRetvalStructVals :: C.ToIdent a => a -> [VName] -> [(C.Type, ValueType)] -> [C.InitGroup]+compileGetRetvalStructVals struct = zipWith field+  where+    field name (ty, Prim) =+      [C.cdecl|$ty:ty $id:name = *$id:struct->$id:(closureRetvalStructField name);|]+    field name (ty, _) =+      [C.cdecl|$ty:ty $id:name =+                 {.desc = $string:(pretty name),+                 .mem = $id:struct->$id:(closureRetvalStructField name),+                 .size = 0, .references = NULL};|]++compileGetStructVals ::+  C.ToIdent a =>+  a ->+  [VName] ->+  [(C.Type, ValueType)] ->+  [C.InitGroup]+compileGetStructVals struct = zipWith field+  where+    field name (ty, Prim) =+      [C.cdecl|$ty:ty $id:name = $id:struct->$id:(closureFreeStructField name);|]+    field name (ty, _) =+      [C.cdecl|$ty:ty $id:name =+                 {.desc = $string:(pretty name),+                  .mem = $id:struct->$id:(closureFreeStructField name),+                  .size = 0, .references = NULL};|]++compileWriteBackResVals :: C.ToIdent a => a -> [VName] -> [(C.Type, ValueType)] -> [C.Stm]+compileWriteBackResVals struct = zipWith field+  where+    field name (_, Prim) =+      [C.cstm|*$id:struct->$id:(closureRetvalStructField name) = $id:name;|]+    field name (_, _) =+      [C.cstm|$id:struct->$id:(closureRetvalStructField name) = $id:name.mem;|]++paramToCType :: Param -> GC.CompilerM op s (C.Type, ValueType)+paramToCType (ScalarParam _ pt) = do+  let t = GC.primTypeToCType pt+  return (t, Prim)+paramToCType (MemParam name space') = mcMemToCType name space'++mcMemToCType :: VName -> Space -> GC.CompilerM op s (C.Type, ValueType)+mcMemToCType v space = do+  refcount <- GC.fatMemory space+  cached <- isJust <$> GC.cacheMem v+  return+    ( GC.fatMemType space,+      if refcount && not cached+        then MemBlock+        else RawMem+    )++functionRuntime :: Name -> C.Id+functionRuntime = (`C.toIdent` mempty) . (<> "_total_runtime")++functionRuns :: Name -> C.Id+functionRuns = (`C.toIdent` mempty) . (<> "_runs")++functionIter :: Name -> C.Id+functionIter = (`C.toIdent` mempty) . (<> "_iter")++multiCoreReport :: [(Name, Bool)] -> [C.BlockItem]+multiCoreReport names = report_kernels+  where+    report_kernels = concatMap reportKernel names+    max_name_len_pad = 40+    format_string name True =+      let name_s = nameToString name+          padding = replicate (max_name_len_pad - length name_s) ' '+       in unwords ["tid %2d -", name_s ++ padding, "ran %10d times; avg: %10ldus; total: %10ldus; time pr. iter %9.6f; iters %9ld; avg %ld\n"]+    format_string name False =+      let name_s = nameToString name+          padding = replicate (max_name_len_pad - length name_s) ' '+       in unwords ["        ", name_s ++ padding, "ran %10d times; avg: %10ldus; total: %10ldus; time pr. iter %9.6f; iters %9ld; avg %ld\n"]+    reportKernel (name, is_array) =+      let runs = functionRuns name+          total_runtime = functionRuntime name+          iters = functionIter name+       in if is_array+            then+              [ [C.citem|+                     for (int i = 0; i < ctx->scheduler.num_threads; i++) {+                       fprintf(stderr,+                         $string:(format_string name is_array),+                         i,+                         ctx->$id:runs[i],+                         (long int) ctx->$id:total_runtime[i] / (ctx->$id:runs[i] != 0 ? ctx->$id:runs[i] : 1),+                         (long int) ctx->$id:total_runtime[i],+                         (double) ctx->$id:total_runtime[i] /  (ctx->$id:iters[i] == 0 ? 1 : (double)ctx->$id:iters[i]),+                         (long int) (ctx->$id:iters[i]),+                         (long int) (ctx->$id:iters[i]) / (ctx->$id:runs[i] != 0 ? ctx->$id:runs[i] : 1)+                         );+                     }+                   |]+              ]+            else+              [ [C.citem|+                    fprintf(stderr,+                       $string:(format_string name is_array),+                       ctx->$id:runs,+                       (long int) ctx->$id:total_runtime / (ctx->$id:runs != 0 ? ctx->$id:runs : 1),+                       (long int) ctx->$id:total_runtime,+                       (double) ctx->$id:total_runtime /  (ctx->$id:iters == 0 ? 1 : (double)ctx->$id:iters),+                       (long int) (ctx->$id:iters),+                       (long int) (ctx->$id:iters) / (ctx->$id:runs != 0 ? ctx->$id:runs : 1));+                   |],+                [C.citem|ctx->total_runtime += ctx->$id:total_runtime;|],+                [C.citem|ctx->total_runs += ctx->$id:runs;|]+              ]++addBenchmarkFields :: Name -> Maybe VName -> GC.CompilerM op s ()+addBenchmarkFields name (Just _) = do+  GC.contextField (functionRuntime name) [C.cty|typename int64_t*|] $ Just [C.cexp|calloc(sizeof(typename int64_t), ctx->scheduler.num_threads)|]+  GC.contextField (functionRuns name) [C.cty|int*|] $ Just [C.cexp|calloc(sizeof(int), ctx->scheduler.num_threads)|]+  GC.contextField (functionIter name) [C.cty|typename int64_t*|] $ Just [C.cexp|calloc(sizeof(sizeof(typename int64_t)), ctx->scheduler.num_threads)|]+addBenchmarkFields name Nothing = do+  GC.contextField (functionRuntime name) [C.cty|typename int64_t|] $ Just [C.cexp|0|]+  GC.contextField (functionRuns name) [C.cty|int|] $ Just [C.cexp|0|]+  GC.contextField (functionIter name) [C.cty|typename int64_t|] $ Just [C.cexp|0|]++benchmarkCode :: Name -> Maybe VName -> [C.BlockItem] -> GC.CompilerM op s [C.BlockItem]+benchmarkCode name tid code = do+  addBenchmarkFields name tid+  return+    [C.citems|+     typename uint64_t $id:start;+     if (ctx->profiling && !ctx->profiling_paused) {+       $id:start = get_wall_time();+     }+     $items:code+     if (ctx->profiling && !ctx->profiling_paused) {+       typename uint64_t $id:end = get_wall_time();+       typename uint64_t elapsed = $id:end - $id:start;+       $items:(updateFields tid)+     }+     |]+  where+    start = name <> "_start"+    end = name <> "_end"+    updateFields Nothing =+      [C.citems|__atomic_fetch_add(&ctx->$id:(functionRuns name), 1, __ATOMIC_RELAXED);+                                            __atomic_fetch_add(&ctx->$id:(functionRuntime name), elapsed, __ATOMIC_RELAXED);+                                            __atomic_fetch_add(&ctx->$id:(functionIter name), iterations, __ATOMIC_RELAXED);|]+    updateFields (Just _tid') =+      [C.citems|ctx->$id:(functionRuns name)[tid]++;+                                            ctx->$id:(functionRuntime name)[tid] += elapsed;+                                            ctx->$id:(functionIter name)[tid] += iterations;|]++functionTiming :: Name -> C.Id+functionTiming = (`C.toIdent` mempty) . (<> "_total_time")++functionIterations :: Name -> C.Id+functionIterations = (`C.toIdent` mempty) . (<> "_total_iter")++addTimingFields :: Name -> GC.CompilerM op s ()+addTimingFields name = do+  GC.contextField (functionTiming name) [C.cty|typename int64_t|] $ Just [C.cexp|0|]+  GC.contextField (functionIterations name) [C.cty|typename int64_t|] $ Just [C.cexp|0|]++multicoreName :: String -> GC.CompilerM op s Name+multicoreName s = do+  s' <- newVName ("futhark_mc_" ++ s)+  return $ nameFromString $ baseString s' ++ "_" ++ show (baseTag s')++multicoreDef :: String -> (Name -> GC.CompilerM op s C.Definition) -> GC.CompilerM op s Name+multicoreDef s f = do+  s' <- multicoreName s+  GC.libDecl =<< f s'+  return s'++generateFunction ::+  C.ToIdent a =>+  M.Map VName Space ->+  String ->+  Code ->+  a ->+  [(VName, (C.Type, ValueType))] ->+  [(VName, (C.Type, ValueType))] ->+  VName ->+  VName ->+  GC.CompilerM Multicore s Name+generateFunction lexical basename code fstruct free retval tid ntasks = do+  let (fargs, fctypes) = unzip free+  let (retval_args, retval_ctypes) = unzip retval+  multicoreDef basename $ \s -> do+    fbody <- benchmarkCode s (Just tid) <=< GC.inNewFunction False $+      GC.cachingMemory lexical $+        \decl_cached free_cached -> GC.blockScope $ do+          mapM_ GC.item [C.citems|$decls:(compileGetStructVals fstruct fargs fctypes)|]+          mapM_ GC.item [C.citems|$decls:(compileGetRetvalStructVals fstruct retval_args retval_ctypes)|]+          mapM_ GC.item decl_cached+          code' <- GC.blockScope $ GC.compileCode code+          mapM_ GC.item [C.citems|$items:code'|]+          mapM_ GC.stm free_cached+    return+      [C.cedecl|int $id:s(void *args, typename int64_t iterations, int tid, struct scheduler_info info) {+                           int err = 0;+                           int $id:tid = tid;+                           int $id:ntasks = info.nsubtasks;+                           struct $id:fstruct *$id:fstruct = (struct $id:fstruct*) args;+                           struct futhark_context *ctx = $id:fstruct->ctx;+                           $items:fbody+                           $stms:(compileWriteBackResVals fstruct retval_args retval_ctypes)+                           cleanup: {}+                           return err;+                      }|]++prepareTaskStruct ::+  String ->+  [VName] ->+  [(C.Type, ValueType)] ->+  [VName] ->+  [(C.Type, ValueType)] ->+  GC.CompilerM Multicore s Name+prepareTaskStruct name free_args free_ctypes retval_args retval_ctypes = do+  fstruct <- multicoreDef name $ \s ->+    return+      [C.cedecl|struct $id:s {+                       struct futhark_context *ctx;+                       $sdecls:(compileFreeStructFields free_args free_ctypes)+                       $sdecls:(compileRetvalStructFields retval_args retval_ctypes)+                     };|]+  GC.decl [C.cdecl|struct $id:fstruct $id:fstruct;|]+  GC.stm [C.cstm|$id:fstruct.ctx = ctx;|]+  GC.stms [C.cstms|$stms:(compileSetStructValues fstruct free_args free_ctypes)|]+  GC.stms [C.cstms|$stms:(compileSetRetvalStructValues fstruct retval_args retval_ctypes)|]+  return fstruct++-- Generate a segop function for top_level and potentially nested SegOp code+compileOp :: GC.OpCompiler Multicore ()+compileOp (Segop name params seq_task par_task retvals (SchedulerInfo nsubtask e sched)) = do+  let (ParallelTask seq_code tid) = seq_task+  free_ctypes <- mapM paramToCType params+  retval_ctypes <- mapM paramToCType retvals+  let free_args = map paramName params+      retval_args = map paramName retvals+      free = zip free_args free_ctypes+      retval = zip retval_args retval_ctypes++  e' <- GC.compileExp e++  let lexical = lexicalMemoryUsage $ Function False [] params seq_code [] []++  fstruct <-+    prepareTaskStruct "task" free_args free_ctypes retval_args retval_ctypes++  fpar_task <- generateFunction lexical (name ++ "_task") seq_code fstruct free retval tid nsubtask+  addTimingFields fpar_task++  let ftask_name = fstruct <> "_task"+  GC.decl [C.cdecl|struct scheduler_segop $id:ftask_name;|]+  GC.stm [C.cstm|$id:ftask_name.args = &$id:fstruct;|]+  GC.stm [C.cstm|$id:ftask_name.top_level_fn = $id:fpar_task;|]+  GC.stm [C.cstm|$id:ftask_name.name = $string:(nameToString fpar_task);|]+  GC.stm [C.cstm|$id:ftask_name.iterations = $exp:e';|]+  -- Create the timing fields for the task+  GC.stm [C.cstm|$id:ftask_name.task_time = &ctx->$id:(functionTiming fpar_task);|]+  GC.stm [C.cstm|$id:ftask_name.task_iter = &ctx->$id:(functionIterations fpar_task);|]++  case sched of+    Dynamic -> GC.stm [C.cstm|$id:ftask_name.sched = DYNAMIC;|]+    Static -> GC.stm [C.cstm|$id:ftask_name.sched = STATIC;|]++  -- Generate the nested segop function if available+  fnpar_task <- case par_task of+    Just (ParallelTask nested_code nested_tid) -> do+      let lexical_nested = lexicalMemoryUsage $ Function False [] params nested_code [] []+      fnpar_task <- generateFunction lexical_nested (name ++ "_nested_task") nested_code fstruct free retval nested_tid nsubtask+      GC.stm [C.cstm|$id:ftask_name.nested_fn = $id:fnpar_task;|]+      return $ zip [fnpar_task] [True]+    Nothing -> do+      GC.stm [C.cstm|$id:ftask_name.nested_fn=NULL;|]+      return mempty++  let ftask_err = fpar_task <> "_err"+  let code =+        [C.citems|int $id:ftask_err = scheduler_prepare_task(&ctx->scheduler, &$id:ftask_name);+                  if ($id:ftask_err != 0) {+                    err = 1; goto cleanup;+                  }|]++  mapM_ GC.item code++  -- Add profile fields for -P option+  mapM_ GC.profileReport $ multiCoreReport $ (fpar_task, True) : fnpar_task+compileOp (ParLoop s' i prebody body postbody free tid) = do+  free_ctypes <- mapM paramToCType free+  let free_args = map paramName free++  let lexical =+        lexicalMemoryUsage $+          Function False [] free (prebody <> body) [] []++  fstruct <-+    prepareTaskStruct (s' ++ "_parloop_struct") free_args free_ctypes mempty mempty++  ftask <- multicoreDef (s' ++ "_parloop") $ \s -> do+    fbody <- benchmarkCode s (Just tid)+      <=< GC.inNewFunction False+      $ GC.cachingMemory lexical $+        \decl_cached free_cached -> GC.blockScope $ do+          mapM_+            GC.item+            [C.citems|$decls:(compileGetStructVals fstruct free_args free_ctypes)|]++          mapM_ GC.item decl_cached++          GC.decl [C.cdecl|typename int64_t iterations = end - start;|]+          GC.decl [C.cdecl|typename int64_t $id:i = start;|]+          GC.compileCode prebody+          body' <- GC.blockScope $ GC.compileCode body+          GC.stm+            [C.cstm|for (; $id:i < end; $id:i++) {+                       $items:body'+                     }|]+          GC.compileCode postbody+          GC.stm [C.cstm|cleanup: {}|]+          mapM_ GC.stm free_cached++    return+      [C.cedecl|int $id:s(void *args, typename int64_t start, typename int64_t end, int $id:tid, int tid) {+                       int err = 0;+                       struct $id:fstruct *$id:fstruct = (struct $id:fstruct*) args;+                       struct futhark_context *ctx = $id:fstruct->ctx;+                       $items:fbody+                       return err;+                     }|]++  let ftask_name = ftask <> "_task"+  GC.decl [C.cdecl|struct scheduler_parloop $id:ftask_name;|]+  GC.stm [C.cstm|$id:ftask_name.name = $string:(nameToString ftask);|]+  GC.stm [C.cstm|$id:ftask_name.fn = $id:ftask;|]+  GC.stm [C.cstm|$id:ftask_name.args = &$id:fstruct;|]+  GC.stm [C.cstm|$id:ftask_name.iterations = iterations;|]+  GC.stm [C.cstm|$id:ftask_name.info = info;|]++  let ftask_err = ftask <> "_err"+      ftask_total = ftask <> "_total"+  code' <-+    benchmarkCode+      ftask_total+      Nothing+      [C.citems|int $id:ftask_err = scheduler_execute_task(&ctx->scheduler, &$id:ftask_name);+               if ($id:ftask_err != 0) {+                 err = 1;+                 goto cleanup;+               }|]++  mapM_ GC.item code'+  mapM_ GC.profileReport $ multiCoreReport $ zip [ftask, ftask_total] [True, False]+compileOp (Atomic aop) =+  atomicOps aop++doAtomic ::+  (C.ToIdent a1, C.ToIdent a2) =>+  a1 ->+  a2 ->+  Count u (TExp Int32) ->+  Exp ->+  String ->+  C.Type ->+  GC.CompilerM op s ()+doAtomic old arr ind val op ty = do+  ind' <- GC.compileExp $ untyped $ unCount ind+  val' <- GC.compileExp val+  GC.stm [C.cstm|$id:old = $id:op(&(($ty:ty*)$id:arr.mem)[$exp:ind'], ($ty:ty) $exp:val', __ATOMIC_RELAXED);|]++atomicOps :: AtomicOp -> GC.CompilerM op s ()+atomicOps (AtomicCmpXchg t old arr ind res val) = do+  ind' <- GC.compileExp $ untyped $ unCount ind+  new_val' <- GC.compileExp val+  let cast = [C.cty|$ty:(GC.primTypeToCType t)*|]+  GC.stm+    [C.cstm|$id:res = $id:op(&(($ty:cast)$id:arr.mem)[$exp:ind'],+                ($ty:cast)&$id:old,+                 $exp:new_val',+                 0, __ATOMIC_SEQ_CST, __ATOMIC_RELAXED);|]+  where+    op :: String+    op = "__atomic_compare_exchange_n"+atomicOps (AtomicXchg t old arr ind val) = do+  ind' <- GC.compileExp $ untyped $ unCount ind+  val' <- GC.compileExp val+  let cast = [C.cty|$ty:(GC.primTypeToCType t)*|]+  GC.stm [C.cstm|$id:old = $id:op(&(($ty:cast)$id:arr.mem)[$exp:ind'], $exp:val', __ATOMIC_SEQ_CST);|]+  where+    op :: String+    op = "__atomic_exchange_n"+atomicOps (AtomicAdd t old arr ind val) =+  doAtomic old arr ind val "__atomic_fetch_add" [C.cty|$ty:(GC.intTypeToCType t)|]+atomicOps (AtomicSub t old arr ind val) =+  doAtomic old arr ind val "__atomic_fetch_sub" [C.cty|$ty:(GC.intTypeToCType t)|]+atomicOps (AtomicAnd t old arr ind val) =+  doAtomic old arr ind val "__atomic_fetch_and" [C.cty|$ty:(GC.intTypeToCType t)|]+atomicOps (AtomicOr t old arr ind val) =+  doAtomic old arr ind val "__atomic_fetch_or" [C.cty|$ty:(GC.intTypeToCType t)|]+atomicOps (AtomicXor t old arr ind val) =+  doAtomic old arr ind val "__atomic_fetch_xor" [C.cty|$ty:(GC.intTypeToCType t)|]

@@ -82,6 +82,7 @@      onPart (ErrorString s) = formatEscape s     onPart ErrorInt32 {} = "{}"+    onPart ErrorInt64 {} = "{}"  sizeClassesToPython :: M.Map Name SizeClass -> PyExp sizeClassesToPython = Dict . map f . M.toList

@@ -6,6 +6,7 @@   ( tupleField,     funName,     defaultMemBlockType,+    intTypeToCType,     primTypeToCType,     signedPrimTypeToCType, 

@@ -41,6 +41,7 @@     ErrorMsgPart (..),     errorMsgArgTypes,     ArrayContents (..),+    declaredIn,     lexicalMemoryUsage,     calledFuncs, @@ -302,7 +303,8 @@     go f (Comment _ x) = f x     go _ _ = mempty -    declared (DeclareMem mem space) = M.singleton mem space+    declared (DeclareMem mem space) =+      M.singleton mem space     declared x = go declared x      set (SetMem x y _) = namesFromList [x, y]@@ -364,7 +366,7 @@  -- | Convert a count of elements into a count of bytes, given the -- per-element size.-withElemType :: Count Elements (TExp Int32) -> PrimType -> Count Bytes (TExp Int64)+withElemType :: Count Elements (TExp Int64) -> PrimType -> Count Bytes (TExp Int64) withElemType (Count e) t =   bytes $ sExt64 e * isInt64 (LeafExp (SizeOf t) (IntType Int64)) 

@@ -165,17 +165,17 @@ -- This old value is stored in the first 'VName'.  The second 'VName' -- is the memory block to update.  The 'Exp' is the new value. data AtomicOp-  = AtomicAdd IntType VName VName (Count Elements (Imp.TExp Int32)) Exp-  | AtomicFAdd FloatType VName VName (Count Elements (Imp.TExp Int32)) Exp-  | AtomicSMax IntType VName VName (Count Elements (Imp.TExp Int32)) Exp-  | AtomicSMin IntType VName VName (Count Elements (Imp.TExp Int32)) Exp-  | AtomicUMax IntType VName VName (Count Elements (Imp.TExp Int32)) Exp-  | AtomicUMin IntType VName VName (Count Elements (Imp.TExp Int32)) Exp-  | AtomicAnd IntType VName VName (Count Elements (Imp.TExp Int32)) Exp-  | AtomicOr IntType VName VName (Count Elements (Imp.TExp Int32)) Exp-  | AtomicXor IntType VName VName (Count Elements (Imp.TExp Int32)) Exp-  | AtomicCmpXchg PrimType VName VName (Count Elements (Imp.TExp Int32)) Exp Exp-  | AtomicXchg PrimType VName VName (Count Elements (Imp.TExp Int32)) Exp+  = AtomicAdd IntType VName VName (Count Elements (Imp.TExp Int64)) Exp+  | AtomicFAdd FloatType VName VName (Count Elements (Imp.TExp Int64)) Exp+  | AtomicSMax IntType VName VName (Count Elements (Imp.TExp Int64)) Exp+  | AtomicSMin IntType VName VName (Count Elements (Imp.TExp Int64)) Exp+  | AtomicUMax IntType VName VName (Count Elements (Imp.TExp Int64)) Exp+  | AtomicUMin IntType VName VName (Count Elements (Imp.TExp Int64)) Exp+  | AtomicAnd IntType VName VName (Count Elements (Imp.TExp Int64)) Exp+  | AtomicOr IntType VName VName (Count Elements (Imp.TExp Int64)) Exp+  | AtomicXor IntType VName VName (Count Elements (Imp.TExp Int64)) Exp+  | AtomicCmpXchg PrimType VName VName (Count Elements (Imp.TExp Int64)) Exp Exp+  | AtomicXchg PrimType VName VName (Count Elements (Imp.TExp Int64)) Exp   deriving (Show)  instance FreeIn AtomicOp where

@@ -0,0 +1,126 @@+-- | Multicore imperative code.+module Futhark.CodeGen.ImpCode.Multicore+  ( Program,+    Function,+    FunctionT (Function),+    Code,+    Multicore (..),+    Scheduling (..),+    SchedulerInfo (..),+    AtomicOp (..),+    ParallelTask (..),+    module Futhark.CodeGen.ImpCode,+  )+where++import Futhark.CodeGen.ImpCode hiding (Code, Function)+import qualified Futhark.CodeGen.ImpCode as Imp+import Futhark.Util.Pretty++-- | An imperative program.+type Program = Imp.Functions Multicore++-- | An imperative function.+type Function = Imp.Function Multicore++-- | A piece of imperative code, with multicore operations inside.+type Code = Imp.Code Multicore++-- | A multicore operation.+data Multicore+  = Segop String [Param] ParallelTask (Maybe ParallelTask) [Param] SchedulerInfo+  | ParLoop String VName Code Code Code [Param] VName+  | Atomic AtomicOp++-- | Atomic operations return the value stored before the update.+-- This old value is stored in the first 'VName'.  The second 'VName'+-- is the memory block to update.  The 'Exp' is the new value.+data AtomicOp+  = AtomicAdd IntType VName VName (Count Elements (Imp.TExp Int32)) Exp+  | AtomicSub IntType VName VName (Count Elements (Imp.TExp Int32)) Exp+  | AtomicAnd IntType VName VName (Count Elements (Imp.TExp Int32)) Exp+  | AtomicOr IntType VName VName (Count Elements (Imp.TExp Int32)) Exp+  | AtomicXor IntType VName VName (Count Elements (Imp.TExp Int32)) Exp+  | AtomicXchg PrimType VName VName (Count Elements (Imp.TExp Int32)) Exp+  | AtomicCmpXchg PrimType VName VName (Count Elements (Imp.TExp Int32)) VName Exp+  deriving (Show)++instance FreeIn AtomicOp where+  freeIn' (AtomicAdd _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x+  freeIn' (AtomicSub _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x+  freeIn' (AtomicAnd _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x+  freeIn' (AtomicOr _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x+  freeIn' (AtomicXor _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x+  freeIn' (AtomicCmpXchg _ _ arr i retval x) = freeIn' arr <> freeIn' i <> freeIn' x <> freeIn' retval+  freeIn' (AtomicXchg _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x++data SchedulerInfo = SchedulerInfo+  { nsubtasks :: VName, -- The variable that describes how many subtasks the scheduler created+    iterations :: Imp.Exp, -- The number of total iterations for a task+    scheduling :: Scheduling -- The type scheduling for the task+  }++data ParallelTask = ParallelTask+  { task_code :: Code,+    flatTid :: VName -- The variable for the thread id execution the code+  }++-- | Whether the Scheduler should schedule the tasks as Dynamic+-- or it is restainted to Static+data Scheduling+  = Dynamic+  | Static++instance Pretty Scheduling where+  ppr Dynamic = text "Dynamic"+  ppr Static = text "Static"++-- TODO fix all of this!+instance Pretty SchedulerInfo where+  ppr (SchedulerInfo nsubtask i sched) =+    text "SchedulingInfo"+      <+> text "number of subtasks"+      <+> ppr nsubtask+      <+> text "scheduling"+      <+> ppr sched+      <+> text "iter"+      <+> ppr i++instance Pretty ParallelTask where+  ppr (ParallelTask code _) =+    ppr code++instance Pretty Multicore where+  ppr (Segop s free _par_code seq_code retval scheduler) =+    text "parfor"+      <+> ppr scheduler+      <+> ppr free+      <+> text s+      <+> text "seq_code"+      <+> nestedBlock "{" "}" (ppr seq_code)+      <+> text "retvals"+      <+> ppr retval+  ppr (ParLoop s i prebody body postbody params info) =+    text "parloop" <+> ppr s <+> ppr i+      <+> ppr prebody+      <+> ppr params+      <+> ppr info+      <+> langle+      <+> nestedBlock "{" "}" (ppr body)+      <+> ppr postbody+  ppr (Atomic _) = text "AtomicOp"++instance FreeIn SchedulerInfo where+  freeIn' (SchedulerInfo nsubtask iter _) =+    freeIn' iter <> freeIn' nsubtask++instance FreeIn ParallelTask where+  freeIn' (ParallelTask code _) =+    freeIn' code++instance FreeIn Multicore where+  freeIn' (Segop _ _ par_code seq_code _ info) =+    freeIn' par_code <> freeIn' seq_code <> freeIn' info+  freeIn' (ParLoop _ _ prebody body postbody _ _) =+    freeIn' prebody <> fvBind (Imp.declaredIn prebody) (freeIn' $ body <> postbody)+  freeIn' (Atomic aop) = freeIn' aop

@@ -156,9 +156,9 @@ type CopyCompiler lore r op =   PrimType ->   MemLocation ->-  Slice (Imp.TExp Int32) ->+  Slice (Imp.TExp Int64) ->   MemLocation ->-  Slice (Imp.TExp Int32) ->+  Slice (Imp.TExp Int64) ->   ImpM lore r op ()  -- | An alternate way of compiling an allocation.@@ -191,7 +191,7 @@ data MemLocation = MemLocation   { memLocationName :: VName,     memLocationShape :: [Imp.DimSize],-    memLocationIxFun :: IxFun.IxFun (Imp.TExp Int32)+    memLocationIxFun :: IxFun.IxFun (Imp.TExp Int64)   }   deriving (Eq, Show) @@ -621,7 +621,7 @@         Nothing -> do           out <- imp $ newVName "out_arrsize"           tell-            ( [Imp.ScalarParam out int32],+            ( [Imp.ScalarParam out int64],               M.singleton x $ ScalarDestination out             )           put (memseen, M.insert x out arrseen)@@ -773,7 +773,7 @@     ForLoop i _ bound loopvars -> do       let setLoopParam (p, a)             | Prim _ <- paramType p =-              copyDWIM (paramName p) [] (Var a) [DimFix $ Imp.vi32 i]+              copyDWIM (paramName p) [] (Var a) [DimFix $ Imp.vi64 i]             | otherwise =               return () @@ -828,22 +828,22 @@     uncurry warn loc "Safety check required at run-time." defCompileBasicOp (Pattern _ [pe]) (Index src slice)   | Just idxs <- sliceIndices slice =-    copyDWIM (patElemName pe) [] (Var src) $ map (DimFix . toInt32Exp) idxs+    copyDWIM (patElemName pe) [] (Var src) $ map (DimFix . toInt64Exp) idxs defCompileBasicOp _ Index {} =   return () defCompileBasicOp (Pattern _ [pe]) (Update _ slice se) =-  sUpdate (patElemName pe) (map (fmap toInt32Exp) slice) se+  sUpdate (patElemName pe) (map (fmap toInt64Exp) slice) se defCompileBasicOp (Pattern _ [pe]) (Replicate (Shape ds) se) = do   ds' <- mapM toExp ds   is <- replicateM (length ds) (newVName "i")-  copy_elem <- collect $ copyDWIM (patElemName pe) (map (DimFix . Imp.vi32) is) se []+  copy_elem <- collect $ copyDWIM (patElemName pe) (map (DimFix . Imp.vi64) is) se []   emit $ foldl (.) id (zipWith Imp.For is ds') copy_elem defCompileBasicOp _ Scratch {} =   return () defCompileBasicOp (Pattern [] [pe]) (Iota n e s it) = do   e' <- toExp e   s' <- toExp s-  sFor "i" (toInt32Exp n) $ \i -> do+  sFor "i" (toInt64Exp n) $ \i -> do     let i' = sExt it $ untyped i     x <-       dPrimV "x" $@@ -856,16 +856,16 @@ defCompileBasicOp (Pattern _ [pe]) (Manifest _ src) =   copyDWIM (patElemName pe) [] (Var src) [] defCompileBasicOp (Pattern _ [pe]) (Concat i x ys _) = do-  offs_glb <- dPrimV "tmp_offs" (0 :: Imp.TExp Int32)+  offs_glb <- dPrimV "tmp_offs" 0    forM_ (x : ys) $ \y -> do     y_dims <- arrayDims <$> lookupType y     let rows = case drop i y_dims of           [] -> error $ "defCompileBasicOp Concat: empty array shape for " ++ pretty y-          r : _ -> toInt32Exp r+          r : _ -> toInt64Exp r         skip_dims = take i y_dims         sliceAllDim d = DimSlice 0 d 1-        skip_slices = map (sliceAllDim . toInt32Exp) skip_dims+        skip_slices = map (sliceAllDim . toInt64Exp) skip_dims         destslice = skip_slices ++ [DimSlice (tvExp offs_glb) rows 1]     copyDWIM (patElemName pe) destslice (Var y) []     offs_glb <-- tvExp offs_glb + rows@@ -877,7 +877,7 @@     static_array <- newVNameForFun "static_array"     emit $ Imp.DeclareArray static_array dest_space t $ Imp.ArrayValues vs     let static_src =-          MemLocation static_array [intConst Int32 $ fromIntegral $ length es] $+          MemLocation static_array [intConst Int64 $ fromIntegral $ length es] $             IxFun.iota [fromIntegral $ length es]         entry = MemVar Nothing $ MemEntry dest_space     addVar static_array entry@@ -1216,7 +1216,7 @@  fullyIndexArray ::   VName ->-  [Imp.TExp Int32] ->+  [Imp.TExp Int64] ->   ImpM lore r op (VName, Imp.Space, Count Elements (Imp.TExp Int64)) fullyIndexArray name indices = do   arr <- lookupArray name@@ -1224,7 +1224,7 @@  fullyIndexArray' ::   MemLocation ->-  [Imp.TExp Int32] ->+  [Imp.TExp Int64] ->   ImpM lore r op (VName, Imp.Space, Count Elements (Imp.TExp Int64)) fullyIndexArray' (MemLocation mem _ ixfun) indices = do   space <- entryMemSpace <$> lookupMemory mem@@ -1233,13 +1233,10 @@           let (zero_is, is) = splitFromEnd (length ds) indices            in map (const 0) zero_is ++ is         _ -> indices--      ixfun64 = fmap sExt64 ixfun-      indices64 = fmap sExt64 indices'   return     ( mem,       space,-      elements $ IxFun.index ixfun64 indices64+      elements $ IxFun.index ixfun indices'     )  -- More complicated read/write operations that use index functions.@@ -1253,15 +1250,15 @@ isMapTransposeCopy ::   PrimType ->   MemLocation ->-  Slice (Imp.TExp Int32) ->+  Slice (Imp.TExp Int64) ->   MemLocation ->-  Slice (Imp.TExp Int32) ->+  Slice (Imp.TExp Int64) ->   Maybe-    ( Imp.TExp Int32,-      Imp.TExp Int32,-      Imp.TExp Int32,-      Imp.TExp Int32,-      Imp.TExp Int32+    ( Imp.TExp Int64,+      Imp.TExp Int64,+      Imp.TExp Int64,+      Imp.TExp Int64,+      Imp.TExp Int64     ) isMapTransposeCopy   bt@@ -1334,16 +1331,16 @@         $ transposeArgs           pt           destmem-          (bytes $ sExt64 destoffset)+          (bytes destoffset)           srcmem-          (bytes $ sExt64 srcoffset)-          (sExt64 num_arrays)-          (sExt64 size_x)-          (sExt64 size_y)+          (bytes srcoffset)+          num_arrays+          size_x+          size_y   | Just destoffset <--      IxFun.linearWithOffset (IxFun.slice dest_ixfun64 destslice64) pt_size,+      IxFun.linearWithOffset (IxFun.slice dest_ixfun destslice) pt_size,     Just srcoffset <--      IxFun.linearWithOffset (IxFun.slice src_ixfun64 srcslice64) pt_size = do+      IxFun.linearWithOffset (IxFun.slice src_ixfun srcslice) pt_size = do     srcspace <- entryMemSpace <$> lookupMemory srcmem     destspace <- entryMemSpace <$> lookupMemory destmem     if isScalarSpace srcspace || isScalarSpace destspace@@ -1367,11 +1364,6 @@     MemLocation destmem _ dest_ixfun = dest     MemLocation srcmem _ src_ixfun = src -    dest_ixfun64 = fmap sExt64 dest_ixfun-    destslice64 = map (fmap sExt64) destslice-    src_ixfun64 = fmap sExt64 src_ixfun-    srcslice64 = map (fmap sExt64) srcslice-     isScalarSpace ScalarSpace {} = True     isScalarSpace _ = False @@ -1379,7 +1371,7 @@ copyElementWise bt dest destslice src srcslice = do   let bounds = sliceDims srcslice   is <- replicateM (length bounds) (newVName "i")-  let ivars = map Imp.vi32 is+  let ivars = map Imp.vi64 is   (destmem, destspace, destidx) <-     fullyIndexArray' dest $ fixSlice destslice ivars   (srcmem, srcspace, srcidx) <-@@ -1395,9 +1387,9 @@ copyArrayDWIM ::   PrimType ->   MemLocation ->-  [DimIndex (Imp.TExp Int32)] ->+  [DimIndex (Imp.TExp Int64)] ->   MemLocation ->-  [DimIndex (Imp.TExp Int32)] ->+  [DimIndex (Imp.TExp Int64)] ->   ImpM lore r op (Imp.Code op) copyArrayDWIM   bt@@ -1419,9 +1411,9 @@           Imp.index srcmem srcoffset bt srcspace vol     | otherwise = do       let destslice' =-            fullSliceNum (map toInt32Exp destshape) destslice+            fullSliceNum (map toInt64Exp destshape) destslice           srcslice' =-            fullSliceNum (map toInt32Exp srcshape) srcslice+            fullSliceNum (map toInt64Exp srcshape) srcslice           destrank = length $ sliceDims destslice'           srcrank = length $ sliceDims srcslice'       if destrank /= srcrank@@ -1445,9 +1437,9 @@ -- instead of a variable name. copyDWIMDest ::   ValueDestination ->-  [DimIndex (Imp.TExp Int32)] ->+  [DimIndex (Imp.TExp Int64)] ->   SubExp ->-  [DimIndex (Imp.TExp Int32)] ->+  [DimIndex (Imp.TExp Int64)] ->   ImpM lore r op () copyDWIMDest _ _ (Constant v) (_ : _) =   error $@@ -1539,9 +1531,9 @@ -- Thing.  Both destination and source must be in scope. copyDWIM ::   VName ->-  [DimIndex (Imp.TExp Int32)] ->+  [DimIndex (Imp.TExp Int64)] ->   SubExp ->-  [DimIndex (Imp.TExp Int32)] ->+  [DimIndex (Imp.TExp Int64)] ->   ImpM lore r op () copyDWIM dest dest_slice src src_slice = do   dest_entry <- lookupVar dest@@ -1558,9 +1550,9 @@ -- | As 'copyDWIM', but implicitly 'DimFix'es the indexes. copyDWIMFix ::   VName ->-  [Imp.TExp Int32] ->+  [Imp.TExp Int64] ->   SubExp ->-  [Imp.TExp Int32] ->+  [Imp.TExp Int64] ->   ImpM lore r op () copyDWIMFix dest dest_is src src_is =   copyDWIM dest (map DimFix dest_is) src (map DimFix src_is)@@ -1589,7 +1581,7 @@ typeSize t =   Imp.bytes $     isInt64 (Imp.LeafExp (Imp.SizeOf $ elemType t) int64)-      * product (map (sExt64 . toInt32Exp) (arrayDims t))+      * product (map toInt64Exp (arrayDims t))  --- Building blocks for constructing code. @@ -1664,14 +1656,14 @@ sArrayInMem name pt shape mem =   sArray name pt shape $     ArrayIn mem $-      IxFun.iota $ map (isInt32 . primExpFromSubExp int32) $ shapeDims shape+      IxFun.iota $ map (isInt64 . primExpFromSubExp int64) $ shapeDims shape  -- | Like 'sAllocArray', but permute the in-memory representation of the indices as specified. sAllocArrayPerm :: String -> PrimType -> ShapeBase SubExp -> Space -> [Int] -> ImpM lore r 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 (isInt32 . primExpFromSubExp int32) permuted_dims+  let iota_ixfun = IxFun.iota $ map (isInt64 . primExpFromSubExp int64) permuted_dims   sArray name pt shape $     ArrayIn mem $ IxFun.permute iota_ixfun $ rearrangeInverse perm @@ -1686,30 +1678,30 @@   let num_elems = case vs of         Imp.ArrayValues vs' -> length vs'         Imp.ArrayZeros n -> fromIntegral n-      shape = Shape [intConst Int32 $ toInteger num_elems]+      shape = Shape [intConst Int64 $ toInteger num_elems]   mem <- newVNameForFun $ name ++ "_mem"   emit $ Imp.DeclareArray mem space pt vs   addVar mem $ MemVar Nothing $ MemEntry space   sArray name pt shape $ ArrayIn mem $ IxFun.iota [fromIntegral num_elems] -sWrite :: VName -> [Imp.TExp Int32] -> Imp.Exp -> ImpM lore r op ()+sWrite :: VName -> [Imp.TExp Int64] -> Imp.Exp -> ImpM lore r op () sWrite arr is v = do   (mem, space, offset) <- fullyIndexArray arr is   vol <- asks envVolatility   emit $ Imp.Write mem offset (primExpType v) space vol v -sUpdate :: VName -> Slice (Imp.TExp Int32) -> SubExp -> ImpM lore r op ()+sUpdate :: VName -> Slice (Imp.TExp Int64) -> SubExp -> ImpM lore r op () sUpdate arr slice v = copyDWIM arr slice v []  sLoopNest ::   Shape ->-  ([Imp.TExp Int32] -> ImpM lore r op ()) ->+  ([Imp.TExp Int64] -> ImpM lore r op ()) ->   ImpM lore r op () sLoopNest = sLoopNest' [] . shapeDims   where     sLoopNest' is [] f = f $ reverse is     sLoopNest' is (d : ds) f =-      sFor "nest_i" (toInt32Exp d) $ \i -> sLoopNest' (i : is) ds f+      sFor "nest_i" (toInt64Exp d) $ \i -> sLoopNest' (i : is) ds f  -- | Untyped assignment. (<~~) :: VName -> Imp.Exp -> ImpM lore r op ()

@@ -188,7 +188,7 @@   x' <- toExp x   s' <- toExp s -  sIota (patElemName pe) (toInt32Exp n) x' s' et+  sIota (patElemName pe) (toInt64Exp n) x' s' et expCompiler (Pattern _ [pe]) (BasicOp (Replicate _ se)) =   sReplicate (patElemName pe) se -- Allocation in the "local" space is just a placeholder.@@ -243,7 +243,7 @@         IxFun.linearWithOffset (IxFun.slice destIxFun destslice) bt_size,       Just srcoffset <-         IxFun.linearWithOffset (IxFun.slice srcIxFun srcslice) bt_size = do-      let num_elems = Imp.elements $ product $ map toInt32Exp srcshape+      let num_elems = Imp.elements $ product $ map toInt64Exp srcshape       srcspace <- entryMemSpace <$> lookupMemory srcmem       destspace <- entryMemSpace <$> lookupMemory destmem       emit $

@@ -68,8 +68,8 @@     kernelGlobalThreadIdVar :: VName,     kernelLocalThreadIdVar :: VName,     kernelGroupIdVar :: VName,-    kernelNumGroups :: Imp.TExp Int32,-    kernelGroupSize :: Imp.TExp Int32,+    kernelNumGroups :: Imp.TExp Int64,+    kernelGroupSize :: Imp.TExp Int64,     kernelNumThreads :: Imp.TExp Int32,     kernelWaveSize :: Imp.TExp Int32,     kernelThreadActive :: Imp.TExp Bool,@@ -102,7 +102,7 @@   localEnv f m   where     mkMap ltid dims = do-      let dims' = map toInt32Exp dims+      let dims' = map (sExt32 . toInt64Exp) dims       ids' <- mapM (dPrimVE "ltid_pre") $ unflattenIndex dims' ltid       return (dims, ids') @@ -140,16 +140,16 @@   ImpM lore r op () splitSpace (Pattern [] [size]) o w i elems_per_thread = do   num_elements <- Imp.elements . TPrimExp <$> toExp w-  let i' = toInt32Exp i+  let i' = toInt64Exp i   elems_per_thread' <- Imp.elements . TPrimExp <$> toExp elems_per_thread-  computeThreadChunkSize o i' elems_per_thread' num_elements (mkTV (patElemName size) int32)+  computeThreadChunkSize o i' elems_per_thread' num_elements (mkTV (patElemName size) int64) splitSpace pat _ _ _ _ =   error $ "Invalid target for splitSpace: " ++ pretty pat  compileThreadExp :: ExpCompiler KernelsMem KernelEnv Imp.KernelOp compileThreadExp (Pattern _ [dest]) (BasicOp (ArrayLit es _)) =   forM_ (zip [0 ..] es) $ \(i, e) ->-    copyDWIMFix (patElemName dest) [fromIntegral (i :: Int32)] e []+    copyDWIMFix (patElemName dest) [fromIntegral (i :: Int64)] e [] compileThreadExp dest e =   defCompileExp dest e @@ -179,13 +179,13 @@ -- passed-in function is invoked with the (symbolic) iteration.  For -- multidimensional loops, use 'groupCoverSpace'. groupLoop ::-  Imp.TExp Int32 ->-  (Imp.TExp Int32 -> InKernelGen ()) ->+  Imp.TExp Int64 ->+  (Imp.TExp Int64 -> InKernelGen ()) ->   InKernelGen () groupLoop n f = do   constants <- kernelConstants <$> askEnv   kernelLoop-    (kernelLocalThreadId constants)+    (sExt64 $ kernelLocalThreadId constants)     (kernelGroupSize constants)     n     f@@ -194,8 +194,8 @@ -- all threads in the group participate.  The passed-in function is -- invoked with a (symbolic) point in the index space. groupCoverSpace ::-  [Imp.TExp Int32] ->-  ([Imp.TExp Int32] -> InKernelGen ()) ->+  [Imp.TExp Int64] ->+  ([Imp.TExp Int64] -> InKernelGen ()) ->   InKernelGen () groupCoverSpace ds f =   groupLoop (product ds) $ f . unflattenIndex ds@@ -204,9 +204,9 @@ -- The static arrays stuff does not work inside kernels. compileGroupExp (Pattern _ [dest]) (BasicOp (ArrayLit es _)) =   forM_ (zip [0 ..] es) $ \(i, e) ->-    copyDWIMFix (patElemName dest) [fromIntegral (i :: Int32)] e []+    copyDWIMFix (patElemName dest) [fromIntegral (i :: Int64)] e [] compileGroupExp (Pattern _ [dest]) (BasicOp (Replicate ds se)) = do-  let ds' = map toInt32Exp $ shapeDims ds+  let ds' = map toInt64Exp $ shapeDims ds   groupCoverSpace ds' $ \is ->     copyDWIMFix (patElemName dest) is se (drop (shapeRank ds) is)   sOp $ Imp.Barrier Imp.FenceLocal@@ -232,7 +232,7 @@     sOp $ Imp.Barrier Imp.FenceLocal     ltid <- kernelLocalThreadId . kernelConstants <$> askEnv     sWhen (ltid .==. 0) $-      copyDWIM (patElemName pe) (map (fmap toInt32Exp) slice) se []+      copyDWIM (patElemName pe) (map (fmap toInt64Exp) slice) se []     sOp $ Imp.Barrier Imp.FenceLocal compileGroupExp dest e =   defCompileExp dest e@@ -242,11 +242,11 @@ sanityCheckLevel SegGroup {} =   error "compileGroupOp: unexpected group-level SegOp." -localThreadIDs :: [SubExp] -> InKernelGen [Imp.TExp Int32]+localThreadIDs :: [SubExp] -> InKernelGen [Imp.TExp Int64] localThreadIDs dims = do-  ltid <- kernelLocalThreadId . kernelConstants <$> askEnv-  let dims' = map toInt32Exp dims-  fromMaybe (unflattenIndex dims' ltid)+  ltid <- sExt64 . kernelLocalThreadId . kernelConstants <$> askEnv+  let dims' = map toInt64Exp dims+  maybe (unflattenIndex dims' ltid) (map sExt64)     . M.lookup dims     . kernelLocalIdMap     . kernelConstants@@ -264,7 +264,7 @@ prepareIntraGroupSegHist ::   Count GroupSize SubExp ->   [HistOp KernelsMem] ->-  InKernelGen [[Imp.TExp Int32] -> InKernelGen ()]+  InKernelGen [[Imp.TExp Int64] -> InKernelGen ()] prepareIntraGroupSegHist group_size =   fmap snd . mapAccumLM onOp Nothing   where@@ -281,8 +281,8 @@         (Nothing, AtomicLocking f) -> do           locks <- newVName "locks" -          let num_locks = toInt32Exp $ unCount group_size-              dims = map toInt32Exp $ shapeDims (histShape op) ++ [histWidth op]+          let num_locks = toInt64Exp $ unCount group_size+              dims = map toInt64Exp $ shapeDims (histShape op) ++ [histWidth op]               l' = Locking locks 0 1 0 (pure . (`rem` num_locks) . flattenIndex dims)               locks_t = Array int32 (Shape [unCount group_size]) NoUniqueness @@ -290,7 +290,7 @@           dArray locks int32 (arrayShape locks_t) $             ArrayIn locks_mem $               IxFun.iota $-                map pe32 $ arrayDims locks_t+                map pe64 $ arrayDims locks_t            sComment "All locks start out unlocked" $             groupCoverSpace [kernelGroupSize constants] $ \is ->@@ -321,21 +321,22 @@ compileGroupOp pat (Inner (SegOp (SegScan lvl space scans _ body))) = do   compileGroupSpace lvl space   let (ltids, dims) = unzip $ unSegSpace space-      dims' = map toInt32Exp dims+      dims' = map toInt64Exp dims    whenActive lvl space $     compileStms mempty (kernelBodyStms body) $       forM_ (zip (patternNames pat) $ kernelBodyResult body) $ \(dest, res) ->         copyDWIMFix           dest-          (map Imp.vi32 ltids)+          (map Imp.vi64 ltids)           (kernelResultSubExp res)           []    sOp $ Imp.ErrorSync Imp.FenceLocal    let segment_size = last dims'-      crossesSegment from to = (to - from) .>. (to `rem` segment_size)+      crossesSegment from to =+        (sExt64 to - sExt64 from) .>. (sExt64 to `rem` segment_size)    -- groupScan needs to treat the scan output as a one-dimensional   -- array of scan elements, so we invent some new flattened arrays@@ -351,7 +352,7 @@           (baseString (patElemName pe) ++ "_flat")           (elemType pe_t)           (Shape arr_dims)-          $ ArrayIn mem $ IxFun.iota $ map pe32 arr_dims+          $ ArrayIn mem $ IxFun.iota $ map pe64 arr_dims        num_scan_results = sum $ map (length . segBinOpNeutral) scans @@ -367,7 +368,7 @@       (red_pes, map_pes) =         splitAt (segBinOpResults ops) $ patternElements pat -      dims' = map toInt32Exp dims+      dims' = map toInt64Exp dims        mkTempArr t =         sAllocArray "red_arr" (elemType t) (Shape dims <> arrayShape t) $ Space "local"@@ -380,7 +381,7 @@       let (red_res, map_res) =             splitAt (segBinOpResults ops) $ kernelBodyResult body       forM_ (zip tmp_arrs red_res) $ \(dest, res) ->-        copyDWIMFix dest (map Imp.vi32 ltids) (kernelResultSubExp res) []+        copyDWIMFix dest (map Imp.vi64 ltids) (kernelResultSubExp res) []       zipWithM_ (compileThreadResult space) map_pes map_res    sOp $ Imp.ErrorSync Imp.FenceLocal@@ -390,7 +391,7 @@     -- handle directly with a group-level reduction.     [dim'] -> do       forM_ (zip ops tmps_for_ops) $ \(op, tmps) ->-        groupReduce dim' (segBinOpLambda op) tmps+        groupReduce (sExt32 dim') (segBinOpLambda op) tmps        sOp $ Imp.ErrorSync Imp.FenceLocal @@ -413,10 +414,11 @@                     drop (length ltids) (memLocationShape arr_loc)             sArray "red_arr_flat" pt flat_shape $               ArrayIn (memLocationName arr_loc) $-                IxFun.iota $ map pe32 $ shapeDims flat_shape+                IxFun.iota $ map pe64 $ shapeDims flat_shape        let segment_size = last dims'-          crossesSegment from to = (to - from) .>. (to `rem` segment_size)+          crossesSegment from to =+            (sExt64 to - sExt64 from) .>. (sExt64 to `rem` sExt64 segment_size)        forM_ (zip ops tmps_for_ops) $ \(op, tmps) -> do         tmps_flat <- mapM flatten tmps@@ -463,10 +465,10 @@        forM_ (zip4 red_is vs_per_op ops' ops) $         \(bin, op_vs, do_op, HistOp dest_w _ _ _ shape lam) -> do-          let bin' = toInt32Exp bin-              dest_w' = toInt32Exp dest_w+          let bin' = toInt64Exp bin+              dest_w' = toInt64Exp dest_w               bin_in_bounds = 0 .<=. bin' .&&. bin' .<. dest_w'-              bin_is = map Imp.vi32 (init ltids) ++ [bin']+              bin_is = map Imp.vi64 (init ltids) ++ [bin']               vs_params = takeLast (length op_vs) $ lambdaParams lam            sComment "perform atomic updates" $@@ -502,13 +504,13 @@     -- | A transformation from the logical lock index to the     -- physical position in the array.  This can also be used     -- to make the lock array smaller.-    lockingMapping :: [Imp.TExp Int32] -> [Imp.TExp Int32]+    lockingMapping :: [Imp.TExp Int64] -> [Imp.TExp Int64]   }  -- | A function for generating code for an atomic update.  Assumes -- that the bucket is in-bounds. type DoAtomicUpdate lore r =-  Space -> [VName] -> [Imp.TExp Int32] -> ImpM lore r Imp.KernelOp ()+  Space -> [VName] -> [Imp.TExp Int64] -> ImpM lore r Imp.KernelOp ()  -- | The mechanism that will be used for performing the atomic update. -- Approximates how efficient it will be.  Ordered from most to least@@ -524,7 +526,7 @@ -- | Is there an atomic t'BinOp' corresponding to this t'BinOp'? type AtomicBinOp =   BinOp ->-  Maybe (VName -> VName -> Count Imp.Elements (Imp.TExp Int32) -> Imp.Exp -> Imp.AtomicOp)+  Maybe (VName -> VName -> Count Imp.Elements (Imp.TExp Int64) -> Imp.Exp -> Imp.AtomicOp)  -- | Do an atomic update corresponding to a binary operator lambda. atomicUpdateLocking ::@@ -546,7 +548,7 @@          (arr', _a_space, bucket_offset) <- fullyIndexArray a bucket -        case opHasAtomicSupport space (tvVar old) arr' (sExt32 <$> bucket_offset) op of+        case opHasAtomicSupport space (tvVar old) arr' bucket_offset op of           Just f -> sOp $ f $ Imp.var y t           Nothing ->             atomicUpdateCAS space t a (tvVar old) bucket x $@@ -588,7 +590,7 @@               int32               (tvVar old)               locks'-              (sExt32 <$> locks_offset)+              locks_offset               (untyped $ lockingIsUnlocked locking)               (untyped $ lockingToLock locking)       lock_acquired = tvExp old .==. lockingIsUnlocked locking@@ -601,7 +603,7 @@               int32               (tvVar old)               locks'-              (sExt32 <$> locks_offset)+              locks_offset               (untyped $ lockingToLock locking)               (untyped $ lockingToUnlock locking)       break_loop = continue <-- false@@ -656,7 +658,7 @@   PrimType ->   VName ->   VName ->-  [Imp.TExp Int32] ->+  [Imp.TExp Int64] ->   VName ->   InKernelGen () ->   InKernelGen ()@@ -698,7 +700,7 @@           int32           (tvVar old_bits)           arr'-          (sExt32 <$> bucket_offset)+          bucket_offset           (toBits (Imp.var assumed t))           (toBits (Imp.var x t))     old <~~ fromBits (untyped $ tvExp old_bits)@@ -773,16 +775,16 @@  computeThreadChunkSize ::   SplitOrdering ->-  Imp.TExp Int32 ->-  Imp.Count Imp.Elements (Imp.TExp Int32) ->-  Imp.Count Imp.Elements (Imp.TExp Int32) ->-  TV Int32 ->+  Imp.TExp Int64 ->+  Imp.Count Imp.Elements (Imp.TExp Int64) ->+  Imp.Count Imp.Elements (Imp.TExp Int64) ->+  TV Int64 ->   ImpM lore r op () computeThreadChunkSize (SplitStrided stride) thread_index elements_per_thread num_elements chunk_var =   chunk_var-    <-- sMin32+    <-- sMin64       (Imp.unCount elements_per_thread)-      ((Imp.unCount num_elements - thread_index) `divUp` toInt32Exp stride)+      ((Imp.unCount num_elements - thread_index) `divUp` toInt64Exp stride) computeThreadChunkSize SplitContiguous thread_index elements_per_thread num_elements chunk_var = do   starting_point <-     dPrimV "starting_point" $@@ -796,7 +798,7 @@    sIf     (no_remaining_elements .||. beyond_bounds)-    (chunk_var <-- (0 :: Imp.TExp Int32))+    (chunk_var <-- 0)     ( sIf         is_last_thread         (chunk_var <-- Imp.unCount last_thread_elements)@@ -810,8 +812,8 @@         .<. (thread_index + 1) * Imp.unCount elements_per_thread  kernelInitialisationSimple ::-  Count NumGroups (Imp.TExp Int32) ->-  Count GroupSize (Imp.TExp Int32) ->+  Count NumGroups (Imp.TExp Int64) ->+  Count GroupSize (Imp.TExp Int64) ->   CallKernelGen (KernelConstants, InKernelGen ()) kernelInitialisationSimple (Count num_groups) (Count group_size) = do   global_tid <- newVName "global_tid"@@ -829,7 +831,7 @@           group_id           num_groups           group_size-          (group_size * num_groups)+          (sExt32 (group_size * num_groups))           (Imp.vi32 wave_size)           true           mempty@@ -837,7 +839,7 @@   let set_constants = do         dPrim_ global_tid int32         dPrim_ local_tid int32-        dPrim_ inner_group_size int32+        dPrim_ inner_group_size int64         dPrim_ wave_size int32         dPrim_ group_id int32 @@ -855,8 +857,8 @@   x : xs -> foldl (.&&.) x xs   where     (is, ws) = unzip limit-    actives = zipWith active is $ map toInt32Exp ws-    active i = (Imp.vi32 i .<.)+    actives = zipWith active is $ map toInt64Exp ws+    active i = (Imp.vi64 i .<.)  -- | Change every memory block to be in the global address space, -- except those who are in the local memory space.  This only affects@@ -901,20 +903,20 @@       readReduceArgument param arr         | Prim _ <- paramType param = do           let i = local_tid + tvExp offset-          copyDWIMFix (paramName param) [] (Var arr) [i]+          copyDWIMFix (paramName param) [] (Var arr) [sExt64 i]         | otherwise = do           let i = global_tid + tvExp offset-          copyDWIMFix (paramName param) [] (Var arr) [i]+          copyDWIMFix (paramName param) [] (Var arr) [sExt64 i]        writeReduceOpResult param arr         | Prim _ <- paramType param =-          copyDWIMFix arr [local_tid] (Var $ paramName param) []+          copyDWIMFix arr [sExt64 local_tid] (Var $ paramName param) []         | otherwise =           return ()    let (reduce_acc_params, reduce_arr_params) = splitAt (length arrs) $ lambdaParams lam -  skip_waves <- dPrim "skip_waves" int32+  skip_waves <- dPrimV "skip_waves" (1 :: Imp.TExp Int32)   dLParams $ lambdaParams lam    offset <-- (0 :: Imp.TExp Int32)@@ -936,7 +938,7 @@       group_size = kernelGroupSize constants       wave_id = local_tid `quot` wave_size       in_wave_id = local_tid - wave_id * wave_size-      num_waves = (group_size + wave_size - 1) `quot` wave_size+      num_waves = (sExt32 group_size + wave_size - 1) `quot` wave_size       arg_in_bounds = local_tid + tvExp offset .<. w        doing_in_wave_reductions =@@ -959,8 +961,7 @@         (wave_id .&. (2 * tvExp skip_waves - 1)) .==. 0       apply_in_cross_wave_iteration =         arg_in_bounds .&&. is_first_thread_in_wave .&&. wave_not_skipped-      cross_wave_reductions = do-        skip_waves <-- (1 :: Imp.TExp Int32)+      cross_wave_reductions =         sWhile doing_cross_wave_reductions $ do           barrier           offset <-- tvExp skip_waves * wave_size@@ -974,8 +975,8 @@  groupScan ::   Maybe (Imp.TExp Int32 -> Imp.TExp Int32 -> Imp.TExp Bool) ->-  Imp.TExp Int32 ->-  Imp.TExp Int32 ->+  Imp.TExp Int64 ->+  Imp.TExp Int64 ->   Lambda KernelsMem ->   [VName] ->   InKernelGen ()@@ -983,11 +984,14 @@   constants <- kernelConstants <$> askEnv   renamed_lam <- renameLambda lam -  let ltid = kernelLocalThreadId constants+  let ltid32 = kernelLocalThreadId constants+      ltid = sExt64 ltid32       (x_params, y_params) = splitAt (length arrs) $ lambdaParams lam    dLParams (lambdaParams lam ++ lambdaParams renamed_lam) +  ltid_in_bounds <- dPrimVE "ltid_in_bounds" $ ltid .<. w+   -- The scan works by splitting the group into blocks, which are   -- scanned separately.  Typically, these blocks are smaller than   -- the lockstep width, which enables barrier-free execution inside@@ -1000,8 +1004,8 @@   -- it were a runtime parameter.  Some day.   let block_size = 32       simd_width = kernelWaveSize constants-      block_id = ltid `quot` block_size-      in_block_id = ltid - block_id * block_size+      block_id = ltid32 `quot` block_size+      in_block_id = ltid32 - block_id * block_size       doInBlockScan seg_flag' active =         inBlockScan           constants@@ -1012,7 +1016,6 @@           active           arrs           barrier-      ltid_in_bounds = ltid .<. w       array_scan = not $ all primType $ lambdaReturnType lam       barrier         | array_scan =@@ -1020,19 +1023,19 @@         | otherwise =           sOp $ Imp.Barrier Imp.FenceLocal -      group_offset = kernelGroupId constants * kernelGroupSize constants+      group_offset = sExt64 (kernelGroupId constants) * kernelGroupSize constants        writeBlockResult p arr         | primType $ paramType p =-          copyDWIM arr [DimFix block_id] (Var $ paramName p) []+          copyDWIM arr [DimFix $ sExt64 block_id] (Var $ paramName p) []         | otherwise =-          copyDWIM arr [DimFix $ group_offset + block_id] (Var $ paramName p) []+          copyDWIM arr [DimFix $ group_offset + sExt64 block_id] (Var $ paramName p) []        readPrevBlockResult p arr         | primType $ paramType p =-          copyDWIM (paramName p) [] (Var arr) [DimFix $ block_id - 1]+          copyDWIM (paramName p) [] (Var arr) [DimFix $ sExt64 block_id - 1]         | otherwise =-          copyDWIM (paramName p) [] (Var arr) [DimFix $ group_offset + block_id - 1]+          copyDWIM (paramName p) [] (Var arr) [DimFix $ group_offset + sExt64 block_id - 1]    doInBlockScan seg_flag ltid_in_bounds lam   barrier@@ -1043,7 +1046,7 @@       sWhen is_first_block $         forM_ (zip x_params arrs) $ \(x, arr) ->           unless (primType $ paramType x) $-            copyDWIM arr [DimFix $ arrs_full_size + group_offset + block_size + ltid] (Var $ paramName x) []+            copyDWIM arr [DimFix $ arrs_full_size + group_offset + sExt64 block_size + ltid] (Var $ paramName x) []      barrier @@ -1074,7 +1077,7 @@               arr               [DimFix $ arrs_full_size + group_offset + ltid]               (Var arr)-              [DimFix $ arrs_full_size + group_offset + block_size + ltid]+              [DimFix $ arrs_full_size + group_offset + sExt64 block_size + ltid]      barrier @@ -1092,7 +1095,7 @@           compileBody' x_params $ lambdaBody lam         | Just flag_true <- seg_flag = do           inactive <--            dPrimVE "inactive" $ flag_true (block_id * block_size -1) ltid+            dPrimVE "inactive" $ flag_true (block_id * block_size -1) ltid32           sWhen inactive y_to_x           when array_scan barrier           sUnless inactive $ compileBody' x_params $ lambdaBody lam@@ -1122,7 +1125,7 @@ inBlockScan ::   KernelConstants ->   Maybe (Imp.TExp Int32 -> Imp.TExp Int32 -> Imp.TExp Bool) ->-  Imp.TExp Int32 ->+  Imp.TExp Int64 ->   Imp.TExp Int32 ->   Imp.TExp Int32 ->   Imp.TExp Bool ->@@ -1158,7 +1161,7 @@         | Just flag_true <- seg_flag = do           inactive <-             dPrimVE "inactive" $-              flag_true (ltid - tvExp skip_threads) ltid+              flag_true (ltid32 - tvExp skip_threads) ltid32           sWhen inactive y_to_x           when array_scan barrier           sUnless inactive $ compileBody' x_params $ lambdaBody scan_lam@@ -1169,11 +1172,11 @@           barrier    sComment "in-block scan (hopefully no barriers needed)" $ do-    skip_threads <-- (1 :: Imp.TExp Int32)+    skip_threads <-- 1     sWhile (tvExp skip_threads .<. block_size) $ do       sWhen (in_block_thread_active .&&. active) $ do         sComment "read operands" $-          zipWithM_ (readParam (tvExp skip_threads)) x_params arrs+          zipWithM_ (readParam (sExt64 $ tvExp skip_threads)) x_params arrs         sComment "perform operation" op_to_x        maybeBarrier@@ -1186,10 +1189,11 @@        skip_threads <-- tvExp skip_threads * 2   where-    block_id = ltid `quot` block_size-    in_block_id = ltid - block_id * block_size-    ltid = kernelLocalThreadId constants-    gtid = kernelGlobalThreadId constants+    block_id = ltid32 `quot` block_size+    in_block_id = ltid32 - block_id * block_size+    ltid32 = kernelLocalThreadId constants+    ltid = sExt64 ltid32+    gtid = sExt64 $ kernelGlobalThreadId constants     array_scan = not $ all primType $ lambdaReturnType scan_lam      readInitial p arr@@ -1211,13 +1215,13 @@       | otherwise =         copyDWIM (paramName y) [] (Var $ paramName x) [] -computeMapKernelGroups :: Imp.TExp Int64 -> CallKernelGen (Imp.TExp Int64, Imp.TExp Int32)+computeMapKernelGroups :: Imp.TExp Int64 -> CallKernelGen (Imp.TExp Int64, Imp.TExp Int64) computeMapKernelGroups kernel_size = do-  group_size <- dPrim "group_size" int32+  group_size <- dPrim "group_size" int64   fname <- askFunction   let group_size_key = keyWithEntryPoint fname $ nameFromString $ pretty $ tvVar group_size   sOp $ Imp.GetSize (tvVar group_size) group_size_key Imp.SizeGroup-  num_groups <- dPrimV "num_groups" $ kernel_size `divUp` sExt64 (tvExp group_size)+  num_groups <- dPrimV "num_groups" $ kernel_size `divUp` tvExp group_size   return (tvExp num_groups, tvExp group_size)  simpleKernelConstants ::@@ -1245,9 +1249,9 @@         thread_gtid         thread_ltid         group_id-        (sExt32 num_groups)+        num_groups         group_size-        (group_size * sExt32 num_groups)+        (sExt32 (group_size * num_groups))         0         (Imp.vi64 thread_gtid .<. kernel_size)         mempty,@@ -1272,13 +1276,13 @@   sOp $ Imp.GetGroupId (tvVar phys_group_id) 0   let iterations =         (required_groups - tvExp phys_group_id)-          `divUp` kernelNumGroups constants+          `divUp` sExt32 (kernelNumGroups constants)    sFor "i" iterations $ \i -> do     m . tvExp       =<< dPrimV         "virt_group_id"-        (tvExp phys_group_id + i * kernelNumGroups constants)+        (tvExp phys_group_id + i * sExt32 (kernelNumGroups constants))     -- Make sure the virtual group is actually done before we let     -- another virtual group have its way with it.     sOp $ Imp.Barrier Imp.FenceGlobal@@ -1288,8 +1292,8 @@  sKernelThread ::   String ->-  Count NumGroups (Imp.TExp Int32) ->-  Count GroupSize (Imp.TExp Int32) ->+  Count NumGroups (Imp.TExp Int64) ->+  Count GroupSize (Imp.TExp Int64) ->   VName ->   InKernelGen () ->   CallKernelGen ()@@ -1297,8 +1301,8 @@  sKernelGroup ::   String ->-  Count NumGroups (Imp.TExp Int32) ->-  Count GroupSize (Imp.TExp Int32) ->+  Count NumGroups (Imp.TExp Int64) ->+  Count GroupSize (Imp.TExp Int64) ->   VName ->   InKernelGen () ->   CallKernelGen ()@@ -1331,8 +1335,8 @@   Operations KernelsMem KernelEnv Imp.KernelOp ->   (KernelConstants -> Imp.TExp Int32) ->   String ->-  Count NumGroups (Imp.TExp Int32) ->-  Count GroupSize (Imp.TExp Int32) ->+  Count NumGroups (Imp.TExp Int64) ->+  Count GroupSize (Imp.TExp Int64) ->   VName ->   InKernelGen () ->   CallKernelGen ()@@ -1392,7 +1396,7 @@   t <- subExpType se   ds <- dropLast (arrayRank t) . arrayDims <$> lookupType arr -  let dims = map toInt32Exp $ ds ++ arrayDims t+  let dims = map toInt64Exp $ ds ++ arrayDims t   (constants, set_constants) <-     simpleKernelConstants (product $ map sExt64 dims) "replicate" @@ -1401,7 +1405,7 @@         keyWithEntryPoint fname $           nameFromString $             "replicate_" ++ show (baseTag $ kernelGlobalThreadIdVar constants)-      is' = unflattenIndex dims $ kernelGlobalThreadId constants+      is' = unflattenIndex dims $ sExt64 $ kernelGlobalThreadId constants    sKernelFailureTolerant True threadOperations constants name $ do     set_constants@@ -1432,7 +1436,7 @@         sArray "arr" bt shape $           ArrayIn mem $             IxFun.iota $-              map pe32 $ shapeDims shape+              map pe64 $ shapeDims shape       sReplicateKernel arr $ Var val    return fname@@ -1451,7 +1455,7 @@               []               fname               [ Imp.MemArg arr_mem,-                Imp.ExpArg $ untyped $ product $ map toInt32Exp arr_shape,+                Imp.ExpArg $ untyped $ product $ map toInt64Exp arr_shape,                 Imp.ExpArg $ toExp' v_t' v               ]     _ -> return Nothing@@ -1488,7 +1492,7 @@    sKernelFailureTolerant True threadOperations constants name $ do     set_constants-    let gtid = kernelGlobalThreadId constants+    let gtid = sExt64 $ kernelGlobalThreadId constants     sWhen (kernelThreadActive constants) $ do       (destmem, destspace, destidx) <- fullyIndexArray' destloc [gtid] @@ -1520,7 +1524,7 @@             Imp.ScalarParam s $ IntType bt           ]         shape = Shape [Var n]-        n' = Imp.vi32 n+        n' = Imp.vi64 n         x' = Imp.var x $ IntType bt         s' = Imp.var s $ IntType bt @@ -1529,7 +1533,7 @@         sArray "arr" (IntType bt) shape $           ArrayIn mem $             IxFun.iota $-              map pe32 $ shapeDims shape+              map pe64 $ shapeDims shape       sIotaKernel arr (sExt64 n') x' s' bt    return fname@@ -1537,7 +1541,7 @@ -- | Perform an Iota with a kernel. sIota ::   VName ->-  Imp.TExp Int32 ->+  Imp.TExp Int64 ->   Imp.Exp ->   Imp.Exp ->   IntType ->@@ -1552,7 +1556,7 @@           []           fname           [Imp.MemArg arr_mem, Imp.ExpArg $ untyped n, Imp.ExpArg x, Imp.ExpArg s]-    else sIotaKernel arr (sExt64 n) x s et+    else sIotaKernel arr n x s et  sCopy :: CopyCompiler KernelsMem HostEnv Imp.HostOp sCopy@@ -1565,7 +1569,7 @@       -- Note that the shape of the destination and the source are       -- necessarily the same.       let shape = sliceDims srcslice-          kernel_size = product $ map sExt64 shape+          kernel_size = product shape        (constants, set_constants) <- simpleKernelConstants kernel_size "copy" @@ -1578,7 +1582,7 @@       sKernelFailureTolerant True threadOperations constants name $ do         set_constants -        let gtid = kernelGlobalThreadId constants+        let gtid = sExt64 $ kernelGlobalThreadId constants             dest_is = unflattenIndex shape gtid             src_is = dest_is @@ -1587,7 +1591,7 @@         (_, srcspace, srcidx) <-           fullyIndexArray' srcloc $ fixSlice srcslice src_is -        sWhen (gtid .<. sExt32 kernel_size) $+        sWhen (gtid .<. kernel_size) $           emit $             Imp.Write destmem destidx bt destspace Imp.Nonvolatile $               Imp.index srcmem srcidx bt srcspace Imp.Nonvolatile@@ -1598,26 +1602,29 @@   KernelResult ->   InKernelGen () compileGroupResult _ pe (TileReturns [(w, per_group_elems)] what) = do-  n <- toInt32Exp . arraySize 0 <$> lookupType what+  n <- toInt64Exp . arraySize 0 <$> lookupType what    constants <- kernelConstants <$> askEnv-  let ltid = kernelLocalThreadId constants-      offset = toInt32Exp per_group_elems * kernelGroupId constants+  let ltid = sExt64 $ kernelLocalThreadId constants+      offset =+        toInt64Exp per_group_elems+          * sExt64 (kernelGroupId constants)    -- Avoid loop for the common case where each thread is statically   -- known to write at most one element.   localOps threadOperations $-    if toInt32Exp per_group_elems == kernelGroupSize constants+    if toInt64Exp per_group_elems == kernelGroupSize constants       then-        sWhen (offset + ltid .<. toInt32Exp w) $+        sWhen (ltid + offset .<. toInt64Exp w) $           copyDWIMFix (patElemName pe) [ltid + offset] (Var what) [ltid]       else sFor "i" (n `divUp` kernelGroupSize constants) $ \i -> do         j <- dPrimVE "j" $ kernelGroupSize constants * i + ltid-        sWhen (j .<. n) $ copyDWIMFix (patElemName pe) [j + offset] (Var what) [j]+        sWhen (j + offset .<. toInt64Exp w) $+          copyDWIMFix (patElemName pe) [j + offset] (Var what) [j] compileGroupResult space pe (TileReturns dims what) = do   let gids = map fst $ unSegSpace space-      out_tile_sizes = map (toInt32Exp . snd) dims-      group_is = zipWith (*) (map Imp.vi32 gids) out_tile_sizes+      out_tile_sizes = map (toInt64Exp . snd) dims+      group_is = zipWith (*) (map Imp.vi64 gids) out_tile_sizes   local_is <- localThreadIDs $ map snd dims   is_for_thread <-     mapM (dPrimV "thread_out_index") $@@ -1629,7 +1636,7 @@ compileGroupResult space pe (Returns _ what) = do   constants <- kernelConstants <$> askEnv   in_local_memory <- arrayInLocalMemory what-  let gids = map (Imp.vi32 . fst) $ unSegSpace space+  let gids = map (Imp.vi64 . fst) $ unSegSpace space    if not in_local_memory     then@@ -1652,22 +1659,24 @@   KernelResult ->   InKernelGen () compileThreadResult space pe (Returns _ what) = do-  let is = map (Imp.vi32 . fst) $ unSegSpace space+  let is = map (Imp.vi64 . fst) $ unSegSpace space   copyDWIMFix (patElemName pe) is what [] compileThreadResult _ pe (ConcatReturns SplitContiguous _ per_thread_elems what) = do   constants <- kernelConstants <$> askEnv-  let offset = toInt32Exp per_thread_elems * kernelGlobalThreadId constants-  n <- toInt32Exp . arraySize 0 <$> lookupType what+  let offset =+        toInt64Exp per_thread_elems+          * sExt64 (kernelGlobalThreadId constants)+  n <- toInt64Exp . arraySize 0 <$> lookupType what   copyDWIM (patElemName pe) [DimSlice offset n 1] (Var what) [] compileThreadResult _ pe (ConcatReturns (SplitStrided stride) _ _ what) = do-  offset <- kernelGlobalThreadId . kernelConstants <$> askEnv-  n <- toInt32Exp . arraySize 0 <$> lookupType what-  copyDWIM (patElemName pe) [DimSlice offset n $ toInt32Exp stride] (Var what) []+  offset <- sExt64 . kernelGlobalThreadId . kernelConstants <$> askEnv+  n <- toInt64Exp . arraySize 0 <$> lookupType what+  copyDWIM (patElemName pe) [DimSlice offset n $ toInt64Exp stride] (Var what) [] compileThreadResult _ pe (WriteReturns rws _arr dests) = do   constants <- kernelConstants <$> askEnv-  let rws' = map toInt32Exp rws+  let rws' = map toInt64Exp rws   forM_ dests $ \(slice, e) -> do-    let slice' = map (fmap toInt32Exp) slice+    let slice' = map (fmap toInt64Exp) slice         condInBounds (DimFix i) rw =           0 .<=. i .&&. i .<. rw         condInBounds (DimSlice i n s) rw =

@@ -62,23 +62,22 @@  data SegHistSlug = SegHistSlug   { slugOp :: HistOp KernelsMem,-    slugNumSubhistos :: TV Int32,+    slugNumSubhistos :: TV Int64,     slugSubhistos :: [SubhistosInfo],     slugAtomicUpdate :: AtomicUpdate KernelsMem KernelEnv   }  histoSpaceUsage ::   HistOp KernelsMem ->-  Imp.Count Imp.Bytes (Imp.TExp Int32)+  Imp.Count Imp.Bytes (Imp.TExp Int64) histoSpaceUsage op =-  fmap sExt32 $-    sum $-      map-        ( typeSize-            . (`arrayOfRow` histWidth op)-            . (`arrayOfShape` histShape op)-        )-        $ lambdaReturnType $ histOp op+  sum $+    map+      ( typeSize+          . (`arrayOfRow` histWidth op)+          . (`arrayOfShape` histShape op)+      )+      $ lambdaReturnType $ histOp op  -- | Figure out how much memory is needed per histogram, both -- segmented and unsegmented,, and compute some other auxiliary@@ -87,8 +86,8 @@   SegSpace ->   HistOp KernelsMem ->   CallKernelGen-    ( Imp.Count Imp.Bytes (Imp.TExp Int32),-      Imp.Count Imp.Bytes (Imp.TExp Int32),+    ( Imp.Count Imp.Bytes (Imp.TExp Int64),+      Imp.Count Imp.Bytes (Imp.TExp Int64),       SegHistSlug     ) computeHistoUsage space op = do@@ -111,7 +110,7 @@         subhistos_membind =           ArrayIn subhistos_mem $             IxFun.iota $-              map pe32 $ shapeDims subhistos_shape+              map pe64 $ shapeDims subhistos_shape     subhistos <-       sArray         (baseString dest ++ "_subhistos")@@ -128,8 +127,8 @@              multiHistoCase = do               let num_elems =-                    foldl' (*) (tvExp num_subhistos) $-                      map toInt32Exp $ arrayDims dest_t+                    foldl' (*) (sExt64 $ tvExp num_subhistos) $+                      map toInt64Exp $ arrayDims dest_t                let subhistos_mem_size =                     Imp.bytes $@@ -139,15 +138,15 @@               sReplicate subhistos ne               subhistos_t <- lookupType subhistos               let slice =-                    fullSliceNum (map toInt32Exp $ arrayDims subhistos_t) $-                      map (unitSlice 0 . toInt32Exp . snd) segment_dims+                    fullSliceNum (map toInt64Exp $ arrayDims subhistos_t) $+                      map (unitSlice 0 . toInt64Exp . snd) segment_dims                         ++ [DimFix 0]               sUpdate subhistos slice $ Var dest          sIf (tvExp num_subhistos .==. 1) unitHistoCase multiHistoCase    let h = histoSpaceUsage op-      segmented_h = h * product (map (Imp.bytes . toInt32Exp) $ init $ segSpaceDims space)+      segmented_h = h * product (map (Imp.bytes . toInt64Exp) $ init $ segSpaceDims space)    atomics <- hostAtomics <$> askEnv @@ -164,7 +163,7 @@   SegHistSlug ->   CallKernelGen     ( Maybe Locking,-      [Imp.TExp Int32] -> InKernelGen ()+      [Imp.TExp Int64] -> InKernelGen ()     ) prepareAtomicUpdateGlobal l dests slug =   -- We need a separate lock array if the operators are not all of a@@ -183,7 +182,7 @@       -- algorithm to ensure good distribution of locks.       let num_locks = 100151           dims =-            map toInt32Exp $+            map toInt64Exp $               shapeDims (histShape (slugOp slug))                 ++ [ tvSize (slugNumSubhistos slug),                      histWidth (slugOp slug)@@ -208,11 +207,11 @@ prepareIntermediateArraysGlobal ::   Passage ->   Imp.TExp Int32 ->-  Imp.TExp Int32 ->+  Imp.TExp Int64 ->   [SegHistSlug] ->   CallKernelGen     ( Imp.TExp Int32,-      [[Imp.TExp Int32] -> InKernelGen ()]+      [[Imp.TExp Int64] -> InKernelGen ()]     ) prepareIntermediateArraysGlobal passage hist_T hist_N slugs = do   -- The paper formulae assume there is only one histogram, but in our@@ -223,11 +222,11 @@   -- paper.    -- The sum of all Hs.-  hist_H <- dPrimVE "hist_H" $ sum $ map (toInt32Exp . histWidth . slugOp) slugs+  hist_H <- dPrimVE "hist_H" $ sum $ map (toInt64Exp . histWidth . slugOp) slugs    hist_RF <-     dPrimVE "hist_RF" $-      sum (map (r64 . toInt32Exp . histRaceFactor . slugOp) slugs)+      sum (map (r64 . toInt64Exp . histRaceFactor . slugOp) slugs)         / genericLength slugs    hist_el_size <- dPrimVE "hist_el_size" $ sum $ map slugElAvgSize slugs@@ -238,7 +237,7 @@    hist_M_min <-     dPrimVE "hist_M_min" $-      sMax32 1 $ t64 $ r64 hist_T / hist_C_max+      sMax32 1 $ sExt32 $ t64 $ r64 hist_T / hist_C_max    -- Querying L2 cache size is not reliable.  Instead we provide a   -- tunable knob with a hopefully sane default.@@ -268,8 +267,9 @@     $ hist_S       <-- case passage of         MayBeMultiPass ->-          (hist_M_min * hist_H * hist_el_size)-            `divUp` t64 (hist_F_L2 * r64 (tvExp hist_L2) * hist_RACE_exp)+          sExt32 $+            (sExt64 hist_M_min * hist_H * sExt64 hist_el_size)+              `divUp` t64 (hist_F_L2 * r64 (tvExp hist_L2) * hist_RACE_exp)         MustBeSinglePass ->           1 @@ -289,7 +289,7 @@     hist_k_RF = 0.75 -- Chosen experimentally     hist_F_L2 = 0.4 -- Chosen experimentally     r64 = isF64 . ConvOpExp (SIToFP Int32 Float64) . untyped-    t64 = isInt32 . ConvOpExp (FPToSI Float64 Int32) . untyped+    t64 = isInt64 . ConvOpExp (FPToSI Float64 Int64) . untyped      -- "Average element size" as computed by a formula that also takes     -- locking into account.@@ -319,9 +319,9 @@      onOp hist_L2 hist_M_min hist_S hist_RACE_exp l slug = do       let SegHistSlug op num_subhistos subhisto_info do_op = slug-          hist_H = toInt32Exp $ histWidth op+          hist_H = toInt64Exp $ histWidth op -      hist_H_chk <- dPrimVE "hist_H_chk" $ hist_H `divUp` hist_S+      hist_H_chk <- dPrimVE "hist_H_chk" $ hist_H `divUp` sExt64 hist_S        emit $ Imp.DebugPrint "Chunk size (H_chk)" $ Just $ untyped hist_H_chk @@ -345,14 +345,14 @@       hist_M <- dPrimVE "hist_M" $         case slugAtomicUpdate slug of           AtomicPrim {} -> 1-          _ -> sMax32 hist_M_min $ t64 $ r64 hist_T / hist_C+          _ -> sMax32 hist_M_min $ sExt32 $ t64 $ r64 hist_T / hist_C        emit $ Imp.DebugPrint "Elements/thread in L2 cache (k_max)" $ Just $ untyped hist_k_max       emit $ Imp.DebugPrint "Multiplication degree (M)" $ Just $ untyped hist_M       emit $ Imp.DebugPrint "Cooperation level (C)" $ Just $ untyped hist_C        -- num_subhistos is the variable we use to communicate back.-      num_subhistos <-- hist_M+      num_subhistos <-- sExt64 hist_M        -- Initialise sub-histograms.       --@@ -384,22 +384,22 @@  histKernelGlobalPass ::   [PatElem KernelsMem] ->-  Count NumGroups (Imp.TExp Int32) ->-  Count GroupSize (Imp.TExp Int32) ->+  Count NumGroups (Imp.TExp Int64) ->+  Count GroupSize (Imp.TExp Int64) ->   SegSpace ->   [SegHistSlug] ->   KernelBody KernelsMem ->-  [[Imp.TExp Int32] -> InKernelGen ()] ->+  [[Imp.TExp Int64] -> InKernelGen ()] ->   Imp.TExp Int32 ->   Imp.TExp Int32 ->   CallKernelGen () histKernelGlobalPass map_pes num_groups group_size space slugs kbody histograms hist_S chk_i = do   let (space_is, space_sizes) = unzip $ unSegSpace space-      space_sizes_64 = map (sExt64 . toInt32Exp) space_sizes+      space_sizes_64 = map (sExt64 . toInt64Exp) space_sizes       total_w_64 = product space_sizes_64    hist_H_chks <- forM (map (histWidth . slugOp) slugs) $ \w ->-    dPrimVE "hist_H_chk" $ toInt32Exp w `divUp` hist_S+    dPrimVE "hist_H_chk" $ toInt64Exp w `divUp` sExt64 hist_S    sKernelThread "seghist_global" num_groups group_size (segFlat space) $ do     constants <- kernelConstants <$> askEnv@@ -408,7 +408,9 @@     subhisto_inds <- forM slugs $ \slug ->       dPrimVE "subhisto_ind" $         kernelGlobalThreadId constants-          `quot` (kernelNumThreads constants `divUp` tvExp (slugNumSubhistos slug))+          `quot` ( kernelNumThreads constants+                     `divUp` sExt32 (tvExp (slugNumSubhistos slug))+                 )      -- Loop over flat offsets into the input and output.  The     -- calculation is done with 64-bit integers to avoid overflow,@@ -434,7 +436,7 @@             forM_ (zip map_pes map_res) $ \(pe, res) ->               copyDWIMFix                 (patElemName pe)-                (map (Imp.vi32 . fst) $ unSegSpace space)+                (map (Imp.vi64 . fst) $ unSegSpace space)                 (kernelResultSubExp res)                 [] @@ -450,9 +452,9 @@                  subhisto_ind,                  hist_H_chk                  ) -> do-                  let chk_beg = chk_i * hist_H_chk-                      bucket' = toInt32Exp $ kernelResultSubExp bucket-                      dest_w' = toInt32Exp dest_w+                  let chk_beg = sExt64 chk_i * hist_H_chk+                      bucket' = toInt64Exp $ kernelResultSubExp bucket+                      dest_w' = toInt64Exp dest_w                       bucket_in_bounds =                         chk_beg .<=. bucket'                           .&&. bucket' .<. (chk_beg + hist_H_chk)@@ -461,8 +463,8 @@                    sWhen bucket_in_bounds $ do                     let bucket_is =-                          map Imp.vi32 (init space_is)-                            ++ [subhisto_ind, bucket']+                          map Imp.vi64 (init space_is)+                            ++ [sExt64 subhisto_ind, bucket']                     dLParams $ lambdaParams lam                     sLoopNest shape $ \is -> do                       forM_ (zip vs_params vs') $ \(p, res) ->@@ -478,10 +480,10 @@   KernelBody KernelsMem ->   CallKernelGen () histKernelGlobal map_pes num_groups group_size space slugs kbody = do-  let num_groups' = fmap toInt32Exp num_groups-      group_size' = fmap toInt32Exp group_size+  let num_groups' = fmap toInt64Exp num_groups+      group_size' = fmap toInt64Exp group_size   let (_space_is, space_sizes) = unzip $ unSegSpace space-      num_threads = unCount num_groups' * unCount group_size'+      num_threads = sExt32 $ unCount num_groups' * unCount group_size'    emit $ Imp.DebugPrint "## Using global memory" Nothing @@ -489,7 +491,7 @@     prepareIntermediateArraysGlobal       (bodyPassage kbody)       num_threads-      (toInt32Exp $ last space_sizes)+      (toInt64Exp $ last space_sizes)       slugs    sFor "chk_i" hist_S $ \chk_i ->@@ -509,25 +511,25 @@       SubExp ->       InKernelGen         ( [VName],-          [Imp.TExp Int32] -> InKernelGen ()+          [Imp.TExp Int64] -> InKernelGen ()         )     )   ]  prepareIntermediateArraysLocal ::   TV Int32 ->-  Count NumGroups (Imp.TExp Int32) ->+  Count NumGroups (Imp.TExp Int64) ->   SegSpace ->   [SegHistSlug] ->   CallKernelGen InitLocalHistograms prepareIntermediateArraysLocal num_subhistos_per_group groups_per_segment space slugs = do   num_segments <-     dPrimVE "num_segments" $-      product $ map (toInt32Exp . snd) $ init $ unSegSpace space+      product $ map (toInt64Exp . snd) $ init $ unSegSpace space   mapM (onOp num_segments) slugs   where     onOp num_segments (SegHistSlug op num_subhistos subhisto_info do_op) = do-      num_subhistos <-- unCount groups_per_segment * num_segments+      num_subhistos <-- sExt64 (unCount groups_per_segment) * num_segments        emit $         Imp.DebugPrint "Number of subhistograms in global memory" $@@ -544,7 +546,7 @@                     shapeDims (histShape op)                       ++ [hist_H_chk] -            let dims = map toInt32Exp $ shapeDims lock_shape+            let dims = map toInt64Exp $ shapeDims lock_shape              locks <- sAllocArray "locks" int32 lock_shape $ Space "local" @@ -581,10 +583,10 @@  histKernelLocalPass ::   TV Int32 ->-  Count NumGroups (Imp.TExp Int32) ->+  Count NumGroups (Imp.TExp Int64) ->   [PatElem KernelsMem] ->-  Count NumGroups (Imp.TExp Int32) ->-  Count GroupSize (Imp.TExp Int32) ->+  Count NumGroups (Imp.TExp Int64) ->+  Count GroupSize (Imp.TExp Int64) ->   SegSpace ->   [SegHistSlug] ->   KernelBody KernelsMem ->@@ -609,33 +611,46 @@         segment_dims = init space_sizes         (i_in_segment, segment_size) = last $ unSegSpace space         num_subhistos_per_group = tvExp num_subhistos_per_group_var-        segment_size' = toInt32Exp segment_size+        segment_size' = toInt64Exp segment_size      num_segments <-       dPrimVE "num_segments" $-        product $ map toInt32Exp segment_dims+        product $ map toInt64Exp segment_dims      hist_H_chks <- forM (map (histWidth . slugOp) slugs) $ \w ->-      dPrimV "hist_H_chk" $ toInt32Exp w `divUp` hist_S+      dPrimV "hist_H_chk" $ toInt64Exp w `divUp` sExt64 hist_S +    histo_sizes <- forM (zip slugs hist_H_chks) $ \(slug, hist_H_chk) -> do+      let histo_dims =+            tvExp hist_H_chk :+            map toInt64Exp (shapeDims (histShape (slugOp slug)))+      histo_size <-+        dPrimVE "histo_size" $ product histo_dims+      let group_hists_size =+            sExt64 num_subhistos_per_group * histo_size+      init_per_thread <-+        dPrimVE "init_per_thread" $ sExt32 $ group_hists_size `divUp` unCount group_size+      return (histo_dims, histo_size, init_per_thread)+     sKernelThread "seghist_local" num_groups group_size (segFlat space) $-      virtualiseGroups SegVirt (unCount groups_per_segment * num_segments) $ \group_id -> do+      virtualiseGroups SegVirt (sExt32 $ unCount groups_per_segment * num_segments) $ \group_id -> do         constants <- kernelConstants <$> askEnv -        flat_segment_id <- dPrimVE "flat_segment_id" $ group_id `quot` unCount groups_per_segment-        gid_in_segment <- dPrimVE "gid_in_segment" $ group_id `rem` unCount groups_per_segment+        flat_segment_id <- dPrimVE "flat_segment_id" $ group_id `quot` sExt32 (unCount groups_per_segment)+        gid_in_segment <- dPrimVE "gid_in_segment" $ group_id `rem` sExt32 (unCount groups_per_segment)         -- This pgtid is kind of a "virtualised physical" gtid - not the         -- same thing as the gtid used for the SegHist itself.         pgtid_in_segment <-           dPrimVE "pgtid_in_segment" $-            gid_in_segment * kernelGroupSize constants + kernelLocalThreadId constants+            gid_in_segment * sExt32 (kernelGroupSize constants)+              + kernelLocalThreadId constants         threads_per_segment <-           dPrimVE "threads_per_segment" $-            unCount groups_per_segment * kernelGroupSize constants+            sExt32 $ unCount groups_per_segment * kernelGroupSize constants          -- Set segment indices.         zipWithM_ dPrimV_ segment_is $-          unflattenIndex (map toInt32Exp segment_dims) flat_segment_id+          unflattenIndex (map toInt64Exp segment_dims) $ sExt64 flat_segment_id          histograms <- forM (zip init_histograms hist_H_chks) $           \((glob_subhistos, init_local_subhistos), hist_H_chk) -> do@@ -649,38 +664,32 @@           dPrimVE "thread_local_subhisto_i" $             kernelLocalThreadId constants `rem` num_subhistos_per_group -        let onSlugs f = forM_ (zip slugs histograms) $ \(slug, (dests, hist_H_chk, _)) -> do-              let histo_dims =-                    tvExp hist_H_chk :-                    map toInt32Exp (shapeDims (histShape (slugOp slug)))-              histo_size <- dPrimVE "histo_size" $ product histo_dims-              f slug dests (tvExp hist_H_chk) histo_dims histo_size+        let onSlugs f =+              forM_ (zip3 slugs histograms histo_sizes) $+                \(slug, (dests, hist_H_chk, _), (histo_dims, histo_size, init_per_thread)) ->+                  f slug dests (tvExp hist_H_chk) histo_dims histo_size init_per_thread          let onAllHistograms f =-              onSlugs $ \slug dests hist_H_chk histo_dims histo_size -> do-                let group_hists_size = num_subhistos_per_group * histo_size-                init_per_thread <--                  dPrimVE "init_per_thread" $-                    group_hists_size-                      `divUp` kernelGroupSize constants+              onSlugs $ \slug dests hist_H_chk histo_dims histo_size init_per_thread -> do+                let group_hists_size = num_subhistos_per_group * sExt32 histo_size                  forM_ (zip dests (histNeutral $ slugOp slug)) $                   \((dest_global, dest_local), ne) ->                     sFor "local_i" init_per_thread $ \i -> do                       j <-                         dPrimVE "j" $-                          i * kernelGroupSize constants+                          i * sExt32 (kernelGroupSize constants)                             + kernelLocalThreadId constants                       j_offset <-                         dPrimVE "j_offset" $-                          num_subhistos_per_group * histo_size * gid_in_segment + j+                          num_subhistos_per_group * sExt32 histo_size * gid_in_segment + j -                      local_subhisto_i <- dPrimVE "local_subhisto_i" $ j `quot` histo_size-                      let local_bucket_is = unflattenIndex histo_dims $ j `rem` histo_size+                      local_subhisto_i <- dPrimVE "local_subhisto_i" $ j `quot` sExt32 histo_size+                      let local_bucket_is = unflattenIndex histo_dims $ sExt64 $ j `rem` sExt32 histo_size                           global_bucket_is =-                            head local_bucket_is + chk_i * hist_H_chk :+                            head local_bucket_is + sExt64 chk_i * hist_H_chk :                             tail local_bucket_is-                      global_subhisto_i <- dPrimVE "global_subhisto_i" $ j_offset `quot` histo_size+                      global_subhisto_i <- dPrimVE "global_subhisto_i" $ j_offset `quot` sExt32 histo_size                        sWhen (j .<. group_hists_size) $                         f@@ -696,8 +705,8 @@         sComment "initialize histograms in local memory" $           onAllHistograms $ \dest_local dest_global op ne local_subhisto_i global_subhisto_i local_bucket_is global_bucket_is ->             sComment "First subhistogram is initialised from global memory; others with neutral element." $ do-              let global_is = map Imp.vi32 segment_is ++ [0] ++ global_bucket_is-                  local_is = local_subhisto_i : local_bucket_is+              let global_is = map Imp.vi64 segment_is ++ [0] ++ global_bucket_is+                  local_is = sExt64 local_subhisto_i : local_bucket_is               sIf                 (global_subhisto_i .==. 0)                 (copyDWIMFix dest_local local_is (Var dest_global) global_is)@@ -707,7 +716,7 @@          sOp $ Imp.Barrier Imp.FenceLocal -        kernelLoop pgtid_in_segment threads_per_segment segment_size' $ \ie -> do+        kernelLoop pgtid_in_segment threads_per_segment (sExt32 segment_size') $ \ie -> do           dPrimV_ i_in_segment ie            -- We execute the bucket function once and update each histogram@@ -726,7 +735,7 @@                 forM_ (zip map_pes map_res) $ \(pe, se) ->                   copyDWIMFix                     (patElemName pe)-                    (map Imp.vi32 space_is)+                    (map Imp.vi64 space_is)                     se                     [] @@ -736,14 +745,14 @@                  bucket,                  vs'                  ) -> do-                  let chk_beg = chk_i * tvExp hist_H_chk-                      bucket' = toInt32Exp bucket-                      dest_w' = toInt32Exp dest_w+                  let chk_beg = sExt64 chk_i * tvExp hist_H_chk+                      bucket' = toInt64Exp bucket+                      dest_w' = toInt64Exp dest_w                       bucket_in_bounds =                         bucket' .<. dest_w'                           .&&. chk_beg .<=. bucket'                           .&&. bucket' .<. (chk_beg + tvExp hist_H_chk)-                      bucket_is = [thread_local_subhisto_i, bucket' - chk_beg]+                      bucket_is = [sExt64 thread_local_subhisto_i, bucket' - chk_beg]                       vs_params = takeLast (length vs') $ lambdaParams lam                    sComment "perform atomic updates" $@@ -757,30 +766,28 @@         sOp $ Imp.ErrorSync Imp.FenceGlobal          sComment "Compact the multiple local memory subhistograms to result in global memory" $-          onSlugs $ \slug dests hist_H_chk histo_dims histo_size -> do-            bins_per_thread <--              dPrimVE "init_per_thread" $-                histo_size `divUp` kernelGroupSize constants-+          onSlugs $ \slug dests hist_H_chk histo_dims _histo_size bins_per_thread -> do             trunc_H <-               dPrimV "trunc_H" $-                sMin32 hist_H_chk $-                  toInt32Exp (histWidth (slugOp slug)) - chk_i * head histo_dims+                sMin64 hist_H_chk $+                  toInt64Exp (histWidth (slugOp slug))+                    - sExt64 chk_i * head histo_dims             let trunc_histo_dims =                   tvExp trunc_H :-                  map toInt32Exp (shapeDims (histShape (slugOp slug)))-            trunc_histo_size <- dPrimVE "histo_size" $ product trunc_histo_dims+                  map toInt64Exp (shapeDims (histShape (slugOp slug)))+            trunc_histo_size <- dPrimVE "histo_size" $ sExt32 $ product trunc_histo_dims              sFor "local_i" bins_per_thread $ \i -> do               j <-                 dPrimVE "j" $-                  i * kernelGroupSize constants + kernelLocalThreadId constants+                  i * sExt32 (kernelGroupSize constants)+                    + kernelLocalThreadId constants               sWhen (j .<. trunc_histo_size) $ do                 -- We are responsible for compacting the flat bin 'j', which                 -- we immediately unflatten.-                let local_bucket_is = unflattenIndex histo_dims j+                let local_bucket_is = unflattenIndex histo_dims $ sExt64 j                     global_bucket_is =-                      head local_bucket_is + chk_i * hist_H_chk :+                      head local_bucket_is + sExt64 chk_i * hist_H_chk :                       tail local_bucket_is                 dLParams $ lambdaParams $ histOp $ slugOp slug                 let (global_dests, local_dests) = unzip dests@@ -803,20 +810,20 @@                         (paramName yp)                         []                         (Var subhisto)-                        (subhisto_id + 1 : local_bucket_is)+                        (sExt64 subhisto_id + 1 : local_bucket_is)                     compileBody' xparams $ lambdaBody $ histOp $ slugOp slug                  sComment "Put final bucket value in global memory." $ do                   let global_is =-                        map Imp.vi32 segment_is-                          ++ [group_id `rem` unCount groups_per_segment]+                        map Imp.vi64 segment_is+                          ++ [sExt64 group_id `rem` unCount groups_per_segment]                           ++ global_bucket_is                   forM_ (zip xparams global_dests) $ \(xp, global_dest) ->                     copyDWIMFix global_dest global_is (Var $ paramName xp) []  histKernelLocal ::   TV Int32 ->-  Count NumGroups (Imp.TExp Int32) ->+  Count NumGroups (Imp.TExp Int64) ->   [PatElem KernelsMem] ->   Count NumGroups SubExp ->   Count GroupSize SubExp ->@@ -826,8 +833,8 @@   KernelBody KernelsMem ->   CallKernelGen () histKernelLocal num_subhistos_per_group_var groups_per_segment map_pes num_groups group_size space hist_S slugs kbody = do-  let num_groups' = fmap toInt32Exp num_groups-      group_size' = fmap toInt32Exp group_size+  let num_groups' = fmap toInt64Exp num_groups+      group_size' = fmap toInt64Exp group_size       num_subhistos_per_group = tvExp num_subhistos_per_group_var    emit $@@ -864,9 +871,9 @@   [PatElem KernelsMem] ->   Imp.TExp Int32 ->   SegSpace ->-  Imp.TExp Int32 ->-  Imp.TExp Int32 ->-  Imp.TExp Int32 ->+  Imp.TExp Int64 ->+  Imp.TExp Int64 ->+  Imp.TExp Int64 ->   Imp.TExp Int32 ->   [SegHistSlug] ->   KernelBody KernelsMem ->@@ -885,20 +892,20 @@   num_groups <-     fmap (Imp.Count . tvSize) $       dPrimV "num_groups" $-        hist_T `divUp` toInt32Exp (unCount group_size)-  let num_groups' = toInt32Exp <$> num_groups-      group_size' = toInt32Exp <$> group_size+        sExt64 hist_T `divUp` toInt64Exp (unCount group_size)+  let num_groups' = toInt64Exp <$> num_groups+      group_size' = toInt64Exp <$> group_size -  let r64 = isF64 . ConvOpExp (SIToFP Int32 Float64) . untyped-      t64 = isInt32 . ConvOpExp (FPToSI Float64 Int32) . untyped+  let r64 = isF64 . ConvOpExp (SIToFP Int64 Float64) . untyped+      t64 = isInt64 . ConvOpExp (FPToSI Float64 Int64) . untyped    -- M approximation.   hist_m' <-     dPrimVE "hist_m_prime" $       r64-        ( sMin32-            (tvExp hist_L `quot` hist_el_size)-            (hist_N `divUp` unCount num_groups')+        ( sMin64+            (sExt64 (tvExp hist_L `quot` hist_el_size))+            (hist_N `divUp` sExt64 (unCount num_groups'))         )         / r64 hist_H @@ -907,15 +914,15 @@   -- M in the paper, but not adjusted for asymptotic efficiency.   hist_M0 <-     dPrimVE "hist_M0" $-      sMax32 1 $ sMin32 (t64 hist_m') hist_B+      sMax64 1 $ sMin64 (t64 hist_m') hist_B    -- Minimal sequential chunking factor.   let q_small = 2    -- The number of segments/histograms produced..-  hist_Nout <- dPrimVE "hist_Nout" $ product $ map toInt32Exp segment_dims+  hist_Nout <- dPrimVE "hist_Nout" $ product $ map toInt64Exp segment_dims -  hist_Nin <- dPrimVE "hist_Nin" $ toInt32Exp $ last space_sizes+  hist_Nin <- dPrimVE "hist_Nin" $ toInt64Exp $ last space_sizes    -- Maximum M for work efficiency.   work_asymp_M_max <-@@ -928,9 +935,9 @@                 `divUp` sExt64 hist_Nout          -- Number of groups, rounded up.-        let r = hist_T_hist_min `divUp` hist_B+        let r = hist_T_hist_min `divUp` sExt32 hist_B -        dPrimVE "work_asymp_M_max" $ hist_Nin `quot` (r * hist_H)+        dPrimVE "work_asymp_M_max" $ hist_Nin `quot` (sExt64 r * hist_H)       else         dPrimVE "work_asymp_M_max" $           (hist_Nout * hist_N)@@ -939,7 +946,7 @@                    )    -- Number of subhistograms per result histogram.-  hist_M <- dPrimV "hist_M" $ sMin32 hist_M0 work_asymp_M_max+  hist_M <- dPrimV "hist_M" $ sExt32 $ sMin64 hist_M0 work_asymp_M_max    -- hist_M may be zero (which we'll check for below), but we need it   -- for some divisions first, so crudely make a nonzero form.@@ -949,7 +956,7 @@   -- working on the same (sub)histogram.   hist_C <-     dPrimVE "hist_C" $-      hist_B `divUp` hist_M_nonzero+      hist_B `divUp` sExt64 hist_M_nonzero    emit $ Imp.DebugPrint "local hist_M0" $ Just $ untyped hist_M0   emit $ Imp.DebugPrint "local work asymp M max" $ Just $ untyped work_asymp_M_max@@ -958,18 +965,30 @@   emit $ Imp.DebugPrint "local M" $ Just $ untyped $ tvExp hist_M   emit $     Imp.DebugPrint "local memory needed" $-      Just $ untyped $ hist_H * hist_el_size * tvExp hist_M+      Just $ untyped $ hist_H * hist_el_size * sExt64 (tvExp hist_M)    -- local_mem_needed is what we need to keep a single bucket in local   -- memory - this is an absolute minimum.  We can fit anything else   -- by doing multiple passes, although more than a few is   -- (heuristically) not efficient.-  local_mem_needed <- dPrimVE "local_mem_needed" $ hist_el_size * tvExp hist_M-  hist_S <- dPrimVE "hist_S" $ (hist_H * local_mem_needed) `divUp` tvExp hist_L+  local_mem_needed <-+    dPrimVE "local_mem_needed" $+      hist_el_size * sExt64 (tvExp hist_M)+  hist_S <-+    dPrimVE "hist_S" $+      sExt32 $+        (hist_H * local_mem_needed) `divUp` tvExp hist_L   let max_S = case bodyPassage kbody of         MustBeSinglePass -> 1         MayBeMultiPass -> fromIntegral $ maxinum $ map slugMaxLocalMemPasses slugs +  groups_per_segment <-+    if segmented+      then+        fmap Count $+          dPrimVE "groups_per_segment" $ unCount num_groups' `divUp` hist_Nout+      else pure num_groups'+   -- We only use local memory if the number of updates per histogram   -- at least matches the histogram size, as otherwise it is not   -- asymptotically efficient.  This mostly matters for the segmented@@ -981,10 +1000,6 @@           .&&. hist_C .<=. hist_B           .&&. tvExp hist_M .>. 0 -      groups_per_segment-        | segmented = num_groups' `divUp` Imp.Count hist_Nout-        | otherwise = num_groups'-       run = do         emit $ Imp.DebugPrint "## Using local memory" Nothing         emit $ Imp.DebugPrint "Histogram size (H)" $ Just $ untyped hist_H@@ -1020,9 +1035,9 @@   -- 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).-  let num_groups' = fmap toInt32Exp num_groups-      group_size' = fmap toInt32Exp group_size-      dims = map toInt32Exp $ segSpaceDims space+  let num_groups' = fmap toInt64Exp num_groups+      group_size' = fmap toInt64Exp group_size+      dims = map toInt64Exp $ segSpaceDims space        num_red_res = length ops + sum (map (length . histNeutral) ops)       (all_red_pes, map_pes) = splitAt num_red_res pes@@ -1038,7 +1053,7 @@     let hist_B = unCount group_size'      -- Size of a histogram.-    hist_H <- dPrimVE "hist_H" $ sum $ map (toInt32Exp . histWidth) ops+    hist_H <- dPrimVE "hist_H" $ sum $ map (toInt64Exp . histWidth) ops      -- Size of a single histogram element.  Actually the weighted     -- average of histogram elements in cases where we have more than@@ -1060,7 +1075,7 @@         sum (map (toInt32Exp . histRaceFactor . slugOp) slugs)           `quot` genericLength slugs -    let hist_T = unCount num_groups' * unCount group_size'+    let hist_T = sExt32 $ unCount num_groups' * unCount group_size'     emit $ Imp.DebugPrint "\n# SegHist" Nothing     emit $ Imp.DebugPrint "Number of threads (T)" $ Just $ untyped hist_T     emit $ Imp.DebugPrint "Desired group size (B)" $ Just $ untyped hist_B@@ -1068,7 +1083,7 @@     emit $ Imp.DebugPrint "Input elements per histogram (N)" $ Just $ untyped hist_N     emit $       Imp.DebugPrint "Number of segments" $-        Just $ untyped $ product $ map (toInt32Exp . snd) segment_dims+        Just $ untyped $ product $ map (toInt64Exp . snd) segment_dims     emit $ Imp.DebugPrint "Histogram element size (el_size)" $ Just $ untyped hist_el_size     emit $ Imp.DebugPrint "Race factor (RF)" $ Just $ untyped hist_RF     emit $ Imp.DebugPrint "Memory per set of subhistograms per segment" $ Just $ untyped h@@ -1126,7 +1141,7 @@           red_cont $             flip map subhistos $ \subhisto ->               ( Var subhisto,-                map Imp.vi32 $+                map Imp.vi64 $                   map fst segment_dims ++ [subhistogram_id, bucket_id] ++ vector_ids               )   where

@@ -24,14 +24,15 @@   CallKernelGen () compileSegMap pat lvl space kbody = do   let (is, dims) = unzip $ unSegSpace space-      dims' = map toInt32Exp dims-      num_groups' = toInt32Exp <$> segNumGroups lvl-      group_size' = toInt32Exp <$> segGroupSize lvl+      dims' = map toInt64Exp dims+      num_groups' = toInt64Exp <$> segNumGroups lvl+      group_size' = toInt64Exp <$> segGroupSize lvl    case lvl of     SegThread {} -> do       emit $ Imp.DebugPrint "\n# SegMap" Nothing-      let virt_num_groups = product dims' `divUp` unCount group_size'+      let virt_num_groups =+            sExt32 $ product dims' `divUp` unCount group_size'       sKernelThread "segmap" num_groups' group_size' (segFlat space) $         virtualiseGroups (segVirt lvl) virt_num_groups $ \group_id -> do           local_tid <- kernelLocalThreadId . kernelConstants <$> askEnv@@ -40,7 +41,7 @@                   + sExt64 local_tid            zipWithM_ dPrimV_ is $-            map sExt32 $ unflattenIndex (map sExt64 dims') global_tid+            map sExt64 $ unflattenIndex (map sExt64 dims') global_tid            sWhen (isActive $ unSegSpace space) $             compileStms mempty (kernelBodyStms kbody) $@@ -48,10 +49,10 @@                 kernelBodyResult kbody     SegGroup {} ->       sKernelGroup "segmap_intragroup" num_groups' group_size' (segFlat space) $ do-        let virt_num_groups = product dims'+        let virt_num_groups = sExt32 $ product dims'         precomputeSegOpIDs (kernelBodyStms kbody) $           virtualiseGroups (segVirt lvl) virt_num_groups $ \group_id -> do-            zipWithM_ dPrimV_ is $ unflattenIndex dims' group_id+            zipWithM_ dPrimV_ is $ unflattenIndex dims' $ sExt64 group_id              compileStms mempty (kernelBodyStms kbody) $               zipWithM_ (compileGroupResult space) (patternElements pat) $

@@ -72,7 +72,7 @@ -- for saving the results of the body.  The results should be -- represented as a pairing of a t'SubExp' along with a list of -- indexes into that 'SubExp' for reading the result.-type DoSegBody = ([(SubExp, [Imp.TExp Int32])] -> InKernelGen ()) -> InKernelGen ()+type DoSegBody = ([(SubExp, [Imp.TExp Int64])] -> InKernelGen ()) -> InKernelGen ()  -- | Compile 'SegRed' instance to host-level code with calls to -- various kernels.@@ -106,11 +106,11 @@   | genericLength reds > maxNumOps =     compilerLimitationS $       "compileSegRed': at most " ++ show maxNumOps ++ " reduction operators are supported."-  | [(_, Constant (IntValue (Int32Value 1))), _] <- unSegSpace space =+  | [(_, Constant (IntValue (Int64Value 1))), _] <- unSegSpace space =     nonsegmentedReduction pat num_groups group_size space reds body   | otherwise = do-    let group_size' = toInt32Exp $ unCount group_size-        segment_size = toInt32Exp $ last $ segSpaceDims space+    let group_size' = toInt64Exp $ unCount group_size+        segment_size = toInt64Exp $ last $ segSpaceDims space         use_small_segments = segment_size * 2 .<. group_size'     sIf       use_small_segments@@ -139,7 +139,7 @@       MemArray pt shape _ (ArrayIn mem _) -> do         let shape' = Shape [num_threads] <> shape         sArray "red_arr" pt shape' $-          ArrayIn mem $ IxFun.iota $ map pe32 $ shapeDims shape'+          ArrayIn mem $ IxFun.iota $ map pe64 $ shapeDims shape'       _ -> do         let pt = elemType $ paramType p             shape = Shape [group_size]@@ -176,9 +176,9 @@   CallKernelGen () nonsegmentedReduction segred_pat num_groups group_size space reds body = do   let (gtids, dims) = unzip $ unSegSpace space-      dims' = map toInt32Exp dims-      num_groups' = fmap toInt32Exp num_groups-      group_size' = fmap toInt32Exp group_size+      dims' = map toInt64Exp dims+      num_groups' = fmap toInt64Exp num_groups+      group_size' = fmap toInt64Exp group_size       global_tid = Imp.vi32 $ segFlat space       w = last dims' @@ -201,10 +201,12 @@      -- Since this is the nonsegmented case, all outer segment IDs must     -- necessarily be 0.-    forM_ gtids $ \v -> dPrimV_ v (0 :: Imp.TExp Int32)+    forM_ gtids $ \v -> dPrimV_ v (0 :: Imp.TExp Int64)      let num_elements = Imp.elements w-    let elems_per_thread = num_elements `divUp` Imp.elements (kernelNumThreads constants)+        elems_per_thread =+          num_elements+            `divUp` Imp.elements (sExt64 (kernelNumThreads constants))      slugs <-       mapM@@ -253,7 +255,7 @@             0             [0]             0-            (kernelNumGroups constants)+            (sExt64 $ kernelNumGroups constants)             slug             red_x_params             red_y_params@@ -276,19 +278,19 @@   CallKernelGen () smallSegmentsReduction (Pattern _ segred_pes) num_groups group_size space reds body = do   let (gtids, dims) = unzip $ unSegSpace space-      dims' = map toInt32Exp dims+      dims' = map toInt64Exp dims       segment_size = last dims'    -- Careful to avoid division by zero now.   segment_size_nonzero <--    dPrimVE "segment_size_nonzero" $ sMax32 1 segment_size+    dPrimVE "segment_size_nonzero" $ sMax64 1 segment_size -  let num_groups' = fmap toInt32Exp num_groups-      group_size' = fmap toInt32Exp group_size+  let num_groups' = fmap toInt64Exp num_groups+      group_size' = fmap toInt64Exp group_size   num_threads <- dPrimV "num_threads" $ unCount num_groups' * unCount group_size'   let num_segments = product $ init dims'       segments_per_group = unCount group_size' `quot` segment_size_nonzero-      required_groups = num_segments `divUp` segments_per_group+      required_groups = sExt32 $ num_segments `divUp` segments_per_group    emit $ Imp.DebugPrint "\n# SegRed-small" Nothing   emit $ Imp.DebugPrint "num_segments" $ Just $ untyped num_segments@@ -307,8 +309,10 @@       -- 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.-      let ltid = kernelLocalThreadId constants-          segment_index = (ltid `quot` segment_size_nonzero) + (group_id' * segments_per_group)+      let ltid = sExt64 $ kernelLocalThreadId constants+          segment_index =+            (ltid `quot` segment_size_nonzero)+              + (sExt64 group_id' * sExt64 segments_per_group)           index_within_segment = ltid `rem` segment_size        zipWithM_ dPrimV_ (init gtids) $ unflattenIndex (init dims') segment_index@@ -336,13 +340,14 @@           out_of_bounds        sOp $ Imp.ErrorSync Imp.FenceLocal -- Also implicitly barrier.-      let crossesSegment from to = (to - from) .>. (to `rem` segment_size)+      let crossesSegment from to =+            (sExt64 to - sExt64 from) .>. (sExt64 to `rem` segment_size)       sWhen (segment_size .>. 0) $         sComment "perform segmented scan to imitate reduction" $           forM_ (zip reds reds_arrs) $ \(SegBinOp _ red_op _ _, red_arrs) ->             groupScan               (Just crossesSegment)-              (tvExp num_threads)+              (sExt64 $ tvExp num_threads)               (segment_size * segments_per_group)               red_op               red_arrs@@ -351,13 +356,15 @@        sComment "save final values of segments" $         sWhen-          ( group_id' * segments_per_group + ltid .<. num_segments+          ( sExt64 group_id' * segments_per_group + sExt64 ltid .<. num_segments               .&&. ltid .<. segments_per_group           )           $ 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+            let flat_segment_index =+                  sExt64 group_id' * segments_per_group + sExt64 ltid+                gtids' =+                  unflattenIndex (init dims') flat_segment_index             copyDWIMFix               (patElemName pe)               gtids'@@ -378,11 +385,11 @@   CallKernelGen () largeSegmentsReduction segred_pat num_groups group_size space reds body = do   let (gtids, dims) = unzip $ unSegSpace space-      dims' = map toInt32Exp dims+      dims' = map toInt64Exp dims       num_segments = product $ init dims'       segment_size = last dims'-      num_groups' = fmap toInt32Exp num_groups-      group_size' = fmap toInt32Exp group_size+      num_groups' = fmap toInt64Exp num_groups+      group_size' = fmap toInt64Exp group_size    (groups_per_segment, elems_per_thread) <-     groupsPerSegmentAndElementsPerThread@@ -436,26 +443,26 @@     -- 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 SegVirt (tvExp virt_num_groups) $ \group_id -> do+    virtualiseGroups SegVirt (sExt32 (tvExp virt_num_groups)) $ \group_id -> do       let segment_gtids = init gtids           w = last dims           local_tid = kernelLocalThreadId constants        flat_segment_id <-         dPrimVE "flat_segment_id" $-          group_id `quot` groups_per_segment+          group_id `quot` sExt32 groups_per_segment        global_tid <-         dPrimVE "global_tid" $-          (group_id * unCount group_size' + local_tid)-            `rem` (unCount group_size' * groups_per_segment)+          (sExt64 group_id * sExt64 (unCount group_size') + sExt64 local_tid)+            `rem` (sExt64 (unCount group_size') * groups_per_segment) -      let first_group_for_segment = flat_segment_id * groups_per_segment+      let first_group_for_segment = sExt64 flat_segment_id * groups_per_segment        zipWithM_ dPrimV_ segment_gtids $-        unflattenIndex (init dims') flat_segment_id-      dPrim_ (last gtids) int32-      let num_elements = Imp.elements $ toInt32Exp w+        unflattenIndex (init dims') $ sExt64 flat_segment_id+      dPrim_ (last gtids) int64+      let num_elements = Imp.elements $ toInt64Exp w        slugs <-         mapM (segBinOpSlug local_tid group_id) $@@ -465,7 +472,7 @@           constants           (zip gtids dims')           num_elements-          global_tid+          (sExt32 global_tid)           elems_per_thread           (tvVar threads_per_segment)           slugs@@ -501,8 +508,8 @@                     pes                     group_id                     flat_segment_id-                    (map Imp.vi32 segment_gtids)-                    first_group_for_segment+                    (map Imp.vi64 segment_gtids)+                    (sExt64 first_group_for_segment)                     groups_per_segment                     slug                     red_x_params@@ -521,25 +528,25 @@               forM_ (zip slugs segred_pes) $ \(slug, pes) ->                 sWhen (local_tid .==. 0) $                   forM_ (zip pes (slugAccs slug)) $ \(v, (acc, acc_is)) ->-                    copyDWIMFix (patElemName v) (map Imp.vi32 segment_gtids) (Var acc) acc_is+                    copyDWIMFix (patElemName v) (map Imp.vi64 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.  We have at least one group -- per segment. groupsPerSegmentAndElementsPerThread ::-  Imp.TExp Int32 ->-  Imp.TExp Int32 ->-  Count NumGroups (Imp.TExp Int32) ->-  Count GroupSize (Imp.TExp Int32) ->+  Imp.TExp Int64 ->+  Imp.TExp Int64 ->+  Count NumGroups (Imp.TExp Int64) ->+  Count GroupSize (Imp.TExp Int64) ->   CallKernelGen-    ( Imp.TExp Int32,-      Imp.Count Imp.Elements (Imp.TExp Int32)+    ( Imp.TExp Int64,+      Imp.Count Imp.Elements (Imp.TExp Int64)     ) groupsPerSegmentAndElementsPerThread segment_size num_segments num_groups_hint group_size = do   groups_per_segment <-     dPrimVE "groups_per_segment" $-      unCount num_groups_hint `divUp` sMax32 1 num_segments+      unCount num_groups_hint `divUp` sMax64 1 num_segments   elements_per_thread <-     dPrimVE "elements_per_thread" $       segment_size `divUp` (unCount group_size * groups_per_segment)@@ -552,7 +559,7 @@     -- (either local or global memory).     slugArrs :: [VName],     -- | Places to store accumulator in stage 1 reduction.-    slugAccs :: [(VName, [Imp.TExp Int32])]+    slugAccs :: [(VName, [Imp.TExp Int64])]   }  slugBody :: SegBinOpSlug -> Body KernelsMem@@ -585,29 +592,29 @@         acc <- dPrim (baseString (paramName p) <> "_acc") t         return (tvVar acc, [])       | otherwise =-        return (param_arr, [local_tid, group_id])+        return (param_arr, [sExt64 local_tid, sExt64 group_id])  reductionStageZero ::   KernelConstants ->-  [(VName, Imp.TExp Int32)] ->-  Imp.Count Imp.Elements (Imp.TExp Int32) ->+  [(VName, Imp.TExp Int64)] ->+  Imp.Count Imp.Elements (Imp.TExp Int64) ->   Imp.TExp Int32 ->-  Imp.Count Imp.Elements (Imp.TExp Int32) ->+  Imp.Count Imp.Elements (Imp.TExp Int64) ->   VName ->   [SegBinOpSlug] ->   DoSegBody ->   InKernelGen ([Lambda KernelsMem], InKernelGen ()) reductionStageZero constants ispace num_elements global_tid elems_per_thread threads_per_segment slugs body = do   let (gtids, _dims) = unzip ispace-      gtid = mkTV (last gtids) int32-      local_tid = kernelLocalThreadId constants+      gtid = mkTV (last gtids) int64+      local_tid = sExt64 $ kernelLocalThreadId constants    -- Figure out how many elements this thread should process.-  chunk_size <- dPrim "chunk_size" int32+  chunk_size <- dPrim "chunk_size" int64   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+  computeThreadChunkSize ordering (sExt64 global_tid) elems_per_thread num_elements chunk_size    dScope Nothing $ scopeOfLParams $ concatMap slugParams slugs @@ -631,7 +638,7 @@                   copyDWIMFix arr [local_tid] (Var $ paramName p) []              sOp $ Imp.ErrorSync Imp.FenceLocal -- Also implicitly barrier.-            groupReduce (kernelGroupSize constants) slug_op_renamed (slugArrs slug)+            groupReduce (sExt32 (kernelGroupSize constants)) slug_op_renamed (slugArrs slug)              sOp $ Imp.Barrier Imp.FenceLocal @@ -656,13 +663,13 @@     gtid       <-- case comm of         Commutative ->-          global_tid-            + Imp.vi32 threads_per_segment * i+          sExt64 global_tid+            + Imp.vi64 threads_per_segment * i         Noncommutative ->-          let index_in_segment = global_tid `quot` kernelGroupSize constants-           in local_tid-                + (index_in_segment * Imp.unCount elems_per_thread + i)-                * kernelGroupSize constants+          let index_in_segment = global_tid `quot` sExt32 (kernelGroupSize constants)+           in sExt64 local_tid+                + (sExt64 index_in_segment * Imp.unCount elems_per_thread + i)+                * sExt64 (kernelGroupSize constants)      check_bounds $       sComment "apply map function" $@@ -704,10 +711,10 @@  reductionStageOne ::   KernelConstants ->-  [(VName, Imp.TExp Int32)] ->-  Imp.Count Imp.Elements (Imp.TExp Int32) ->+  [(VName, Imp.TExp Int64)] ->+  Imp.Count Imp.Elements (Imp.TExp Int64) ->   Imp.TExp Int32 ->-  Imp.Count Imp.Elements (Imp.TExp Int32) ->+  Imp.Count Imp.Elements (Imp.TExp Int64) ->   VName ->   [SegBinOpSlug] ->   DoSegBody ->@@ -730,9 +737,9 @@   [PatElem KernelsMem] ->   Imp.TExp Int32 ->   Imp.TExp Int32 ->-  [Imp.TExp Int32] ->-  Imp.TExp Int32 ->-  Imp.TExp Int32 ->+  [Imp.TExp Int64] ->+  Imp.TExp Int64 ->+  Imp.TExp Int64 ->   SegBinOpSlug ->   [LParam KernelsMem] ->   [LParam KernelsMem] ->@@ -770,13 +777,14 @@     (counter_mem, _, counter_offset) <-       fullyIndexArray         counter-        [ counter_i * num_counters-            + flat_segment_id `rem` num_counters+        [ sExt64 $+            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 (slugAccs slug)) $ \(v, (acc, acc_is)) ->-          copyDWIMFix v [0, group_id] (Var acc) acc_is+          copyDWIMFix v [0, sExt64 group_id] (Var acc) acc_is         sOp $ Imp.MemFence Imp.FenceGlobal         -- Increment the counter, thus stating that our result is         -- available.@@ -786,7 +794,7 @@               Int32               (tvVar old_counter)               counter_mem-              (sExt32 <$> counter_offset)+              counter_offset               $ untyped (1 :: Imp.TExp Int32)         -- Now check if we were the last group to write our result.  If         -- so, it is our responsibility to produce the final result.@@ -806,7 +814,7 @@       sWhen (local_tid .==. 0) $         sOp $           Imp.Atomic DefaultSpace $-            Imp.AtomicAdd Int32 (tvVar old_counter) counter_mem (sExt32 <$> counter_offset) $+            Imp.AtomicAdd Int32 (tvVar old_counter) counter_mem counter_offset $               untyped $ negate groups_per_segment        sLoopNest (slugShape slug) $ \vec_is -> do@@ -818,7 +826,7 @@         comment "read in the per-group-results" $ do           read_per_thread <-             dPrimVE "read_per_thread" $-              groups_per_segment `divUp` group_size+              groups_per_segment `divUp` sExt64 group_size            forM_ (zip red_x_params nes) $ \(p, ne) ->             copyDWIMFix (paramName p) [] ne []@@ -826,7 +834,7 @@           sFor "i" read_per_thread $ \i -> do             group_res_id <-               dPrimVE "group_res_id" $-                local_tid * read_per_thread + i+                sExt64 local_tid * read_per_thread + i             index_of_group_res <-               dPrimVE "index_of_group_res" $                 first_group_for_segment + group_res_id@@ -846,12 +854,12 @@          forM_ (zip red_x_params red_arrs) $ \(p, arr) ->           when (primType $ paramType p) $-            copyDWIMFix arr [local_tid] (Var $ paramName p) []+            copyDWIMFix arr [sExt64 local_tid] (Var $ paramName p) []          sOp $ Imp.Barrier Imp.FenceLocal          sComment "reduce the per-group results" $ do-          groupReduce group_size red_op_renamed red_arrs+          groupReduce (sExt32 group_size) red_op_renamed red_arrs            sComment "and back to memory with the final result" $             sWhen (local_tid .==. 0) $

@@ -44,7 +44,7 @@               arr <-                 lift $                   sArray "scan_arr" pt shape' $-                    ArrayIn mem $ IxFun.iota $ map pe32 $ shapeDims shape'+                    ArrayIn mem $ IxFun.iota $ map pe64 $ shapeDims shape'               return (arr, [])             _ -> do               let pt = elemType $ paramType p@@ -69,13 +69,13 @@           mem <- lift $ sDeclareMem "scan_arr_mem" $ Space "local"           return ([size], mem) -type CrossesSegment = Maybe (Imp.TExp Int32 -> Imp.TExp Int32 -> Imp.TExp Bool)+type CrossesSegment = Maybe (Imp.TExp Int64 -> Imp.TExp Int64 -> Imp.TExp Bool) -localArrayIndex :: KernelConstants -> Type -> Imp.TExp Int32+localArrayIndex :: KernelConstants -> Type -> Imp.TExp Int64 localArrayIndex constants t =   if primType t-    then kernelLocalThreadId constants-    else kernelGlobalThreadId constants+    then sExt64 (kernelLocalThreadId constants)+    else sExt64 (kernelGlobalThreadId constants)  barrierFor :: Lambda KernelsMem -> (Bool, Imp.Fence, InKernelGen ()) barrierFor scan_op = (array_scan, fence, sOp $ Imp.Barrier fence)@@ -100,7 +100,7 @@     forM_ (zip pes scan_res) $ \(pe, res) ->       copyDWIMFix         (patElemName pe)-        (map Imp.vi32 gtids)+        (map Imp.vi64 gtids)         (kernelResultSubExp res)         []   | otherwise =@@ -108,7 +108,7 @@       copyDWIMFix (paramName p) [] (kernelResultSubExp res) []  readToScanValues ::-  [Imp.TExp Int32] ->+  [Imp.TExp Int64] ->   [PatElem KernelsMem] ->   SegBinOp KernelsMem ->   InKernelGen ()@@ -120,9 +120,9 @@     return ()  readCarries ::-  Imp.TExp Int32 ->-  [Imp.TExp Int32] ->-  [Imp.TExp Int32] ->+  Imp.TExp Int64 ->+  [Imp.TExp Int64] ->+  [Imp.TExp Int64] ->   [PatElem KernelsMem] ->   SegBinOp KernelsMem ->   InKernelGen ()@@ -152,16 +152,16 @@   SegSpace ->   [SegBinOp KernelsMem] ->   KernelBody KernelsMem ->-  CallKernelGen (TV Int32, Imp.TExp Int32, CrossesSegment)+  CallKernelGen (TV Int32, Imp.TExp Int64, CrossesSegment) scanStage1 (Pattern _ all_pes) num_groups group_size space scans kbody = do-  let num_groups' = fmap toInt32Exp num_groups-      group_size' = fmap toInt32Exp group_size-  num_threads <- dPrimV "num_threads" $ unCount num_groups' * unCount group_size'+  let num_groups' = fmap toInt64Exp num_groups+      group_size' = fmap toInt64Exp group_size+  num_threads <- dPrimV "num_threads" $ sExt32 $ unCount num_groups' * unCount group_size'    let (gtids, dims) = unzip $ unSegSpace space-      dims' = map toInt32Exp dims+      dims' = map toInt64Exp dims   let num_elements = product dims'-      elems_per_thread = num_elements `divUp` tvExp num_threads+      elems_per_thread = num_elements `divUp` sExt64 (tvExp num_threads)       elems_per_group = unCount group_size' * elems_per_thread    let crossesSegment =@@ -184,18 +184,18 @@     sFor "j" elems_per_thread $ \j -> do       chunk_offset <-         dPrimV "chunk_offset" $-          kernelGroupSize constants * j-            + kernelGroupId constants * elems_per_group+          sExt64 (kernelGroupSize constants) * j+            + sExt64 (kernelGroupId constants) * elems_per_group       flat_idx <-         dPrimV "flat_idx" $-          tvExp chunk_offset + kernelLocalThreadId constants+          tvExp chunk_offset + sExt64 (kernelLocalThreadId constants)       -- Construct segment indices.       zipWithM_ dPrimV_ gtids $ unflattenIndex dims' $ tvExp flat_idx        let per_scan_pes = segBinOpChunks scans all_pes            in_bounds =-            foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi32 gtids) dims'+            foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi64 gtids) dims'            when_in_bounds = compileStms mempty (kernelBodyStms kbody) $ do             let (all_scan_res, map_res) =@@ -211,7 +211,7 @@               forM_ (zip (takeLast (length map_res) all_pes) map_res) $ \(pe, se) ->                 copyDWIMFix                   (patElemName pe)-                  (map Imp.vi32 gtids)+                  (map Imp.vi64 gtids)                   (kernelResultSubExp se)                   [] @@ -232,7 +232,7 @@                 sIf                   in_bounds                   ( do-                      readToScanValues (map Imp.vi32 gtids ++ vec_is) pes scan+                      readToScanValues (map Imp.vi64 gtids ++ vec_is) pes scan                       readCarries (tvExp chunk_offset) dims' vec_is pes scan                   )                   ( forM_ (zip (yParams scan) (segBinOpNeutral scan)) $ \(p, ne) ->@@ -242,13 +242,14 @@               sComment "combine with carry and write to local memory" $                 compileStms mempty (bodyStms $ lambdaBody scan_op) $                   forM_ (zip3 rets local_arrs (bodyResult $ lambdaBody scan_op)) $-                    \(t, arr, se) -> copyDWIMFix arr [localArrayIndex constants t] se []+                    \(t, arr, se) ->+                      copyDWIMFix arr [localArrayIndex constants t] se []                let crossesSegment' = do                     f <- crossesSegment                     Just $ \from to ->-                      let from' = from + tvExp chunk_offset-                          to' = to + tvExp chunk_offset+                      let from' = sExt64 from + tvExp chunk_offset+                          to' = sExt64 to + tvExp chunk_offset                        in f from' to'                sOp $ Imp.ErrorSync fence@@ -257,8 +258,8 @@               scan_op_renamed <- renameLambda scan_op               groupScan                 crossesSegment'-                (tvExp num_threads)-                (kernelGroupSize constants)+                (sExt64 $ tvExp num_threads)+                (sExt64 $ kernelGroupSize constants)                 scan_op_renamed                 local_arrs @@ -267,7 +268,7 @@                   forM_ (zip3 rets pes local_arrs) $ \(t, pe, arr) ->                     copyDWIMFix                       (patElemName pe)-                      (map Imp.vi32 gtids ++ vec_is)+                      (map Imp.vi64 gtids ++ vec_is)                       (Var arr)                       [localArrayIndex constants t] @@ -280,8 +281,10 @@                         []                         (Var arr)                         [ if primType $ paramType p-                            then kernelGroupSize constants - 1-                            else (kernelGroupId constants + 1) * kernelGroupSize constants - 1+                            then sExt64 (kernelGroupSize constants) - 1+                            else+                              (sExt64 (kernelGroupId constants) + 1)+                                * sExt64 (kernelGroupSize constants) - 1                         ]                   load_neutral =                     forM_ (zip nes scan_x_params) $ \(ne, p) ->@@ -294,10 +297,10 @@                     Just f ->                       f                         ( tvExp chunk_offset-                            + kernelGroupSize constants -1+                            + sExt64 (kernelGroupSize constants) -1                         )                         ( tvExp chunk_offset-                            + kernelGroupSize constants+                            + sExt64 (kernelGroupSize constants)                         )                 should_load_carry <-                   dPrimVE "should_load_carry" $@@ -313,7 +316,7 @@ scanStage2 ::   Pattern KernelsMem ->   TV Int32 ->-  Imp.TExp Int32 ->+  Imp.TExp Int64 ->   Count NumGroups SubExp ->   CrossesSegment ->   SegSpace ->@@ -321,16 +324,18 @@   CallKernelGen () scanStage2 (Pattern _ all_pes) stage1_num_threads elems_per_group num_groups crossesSegment space scans = do   let (gtids, dims) = unzip $ unSegSpace space-      dims' = map toInt32Exp dims+      dims' = map toInt64Exp dims    -- Our group size is the number of groups for the stage 1 kernel.   let group_size = Count $ unCount num_groups-      group_size' = fmap toInt32Exp group_size+      group_size' = fmap toInt64Exp group_size    let crossesSegment' = do         f <- crossesSegment         Just $ \from to ->-          f ((from + 1) * elems_per_group - 1) ((to + 1) * elems_per_group - 1)+          f+            ((sExt64 from + 1) * elems_per_group - 1)+            ((sExt64 to + 1) * elems_per_group - 1)    sKernelThread "scan_stage2" 1 group_size' (segFlat space) $ do     constants <- kernelConstants <$> askEnv@@ -340,17 +345,17 @@      flat_idx <-       dPrimV "flat_idx" $-        (kernelLocalThreadId constants + 1) * elems_per_group - 1+        (sExt64 (kernelLocalThreadId constants) + 1) * elems_per_group - 1     -- Construct segment indices.     zipWithM_ dPrimV_ gtids $ unflattenIndex dims' $ tvExp flat_idx      forM_ (zip4 scans per_scan_local_arrs per_scan_rets per_scan_pes) $       \(SegBinOp _ scan_op nes vec_shape, local_arrs, rets, pes) ->         sLoopNest vec_shape $ \vec_is -> do-          let glob_is = map Imp.vi32 gtids ++ vec_is+          let glob_is = map Imp.vi64 gtids ++ vec_is                in_bounds =-                foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi32 gtids) dims'+                foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi64 gtids) dims'                when_in_bounds = forM_ (zip3 rets local_arrs pes) $ \(t, arr, pe) ->                 copyDWIMFix@@ -371,8 +376,8 @@            groupScan             crossesSegment'-            (tvExp stage1_num_threads)-            (kernelGroupSize constants)+            (sExt64 $ tvExp stage1_num_threads)+            (sExt64 $ kernelGroupSize constants)             scan_op             local_arrs @@ -389,19 +394,19 @@   Pattern KernelsMem ->   Count NumGroups SubExp ->   Count GroupSize SubExp ->-  Imp.TExp Int32 ->+  Imp.TExp Int64 ->   CrossesSegment ->   SegSpace ->   [SegBinOp KernelsMem] ->   CallKernelGen () scanStage3 (Pattern _ all_pes) num_groups group_size elems_per_group crossesSegment space scans = do-  let num_groups' = fmap toInt32Exp num_groups-      group_size' = fmap toInt32Exp group_size+  let num_groups' = fmap toInt64Exp num_groups+      group_size' = fmap toInt64Exp group_size       (gtids, dims) = unzip $ unSegSpace space-      dims' = map toInt32Exp dims+      dims' = map toInt64Exp dims   required_groups <-     dPrimVE "required_groups" $-      product dims' `divUp` unCount group_size'+      sExt32 $ product dims' `divUp` sExt64 (unCount group_size')    sKernelThread "scan_stage3" num_groups' group_size' (segFlat space) $     virtualiseGroups SegVirt required_groups $ \virt_group_id -> do@@ -410,8 +415,8 @@       -- Compute our logical index.       flat_idx <-         dPrimVE "flat_idx" $-          virt_group_id * unCount group_size'-            + kernelLocalThreadId constants+          sExt64 virt_group_id * sExt64 (unCount group_size')+            + sExt64 (kernelLocalThreadId constants)       zipWithM_ dPrimV_ gtids $ unflattenIndex dims' flat_idx        -- Figure out which group this element was originally in.@@ -428,7 +433,7 @@       -- then the carry was updated in stage 2), and we are not crossing       -- a segment boundary.       let in_bounds =-            foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi32 gtids) dims'+            foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi64 gtids) dims'           crosses_segment =             fromMaybe false $               crossesSegment@@ -459,14 +464,14 @@                     (paramName p)                     []                     (Var $ patElemName pe)-                    (map Imp.vi32 gtids ++ vec_is)+                    (map Imp.vi64 gtids ++ vec_is)                  compileBody' scan_x_params $ lambdaBody scan_op                  forM_ (zip scan_x_params pes) $ \(p, pe) ->                   copyDWIMFix                     (patElemName pe)-                    (map Imp.vi32 gtids ++ vec_is)+                    (map Imp.vi64 gtids ++ vec_is)                     (Var $ paramName p)                     [] @@ -485,14 +490,14 @@   -- Since stage 2 involves a group size equal to the number of groups   -- used for stage 1, we have to cap this number to the maximum group   -- size.-  stage1_max_num_groups <- dPrim "stage1_max_num_groups" int32+  stage1_max_num_groups <- dPrim "stage1_max_num_groups" int64   sOp $ Imp.GetSizeMax (tvVar stage1_max_num_groups) SizeGroup    stage1_num_groups <-     fmap (Imp.Count . tvSize) $       dPrimV "stage1_num_groups" $-        sMin32 (tvExp stage1_max_num_groups) $-          toInt32Exp $ Imp.unCount $ segNumGroups lvl+        sMin64 (tvExp stage1_max_num_groups) $+          toInt64Exp $ Imp.unCount $ segNumGroups lvl    (stage1_num_threads, elems_per_group, crossesSegment) <-     scanStage1 pat stage1_num_groups (segGroupSize lvl) space scans kbody@@ -502,4 +507,4 @@   scanStage2 pat stage1_num_threads elems_per_group stage1_num_groups crossesSegment space scans   scanStage3 pat (segNumGroups lvl) (segGroupSize lvl) elems_per_group crossesSegment space scans   where-    n = product $ map toInt32Exp $ segSpaceDims space+    n = product $ map toInt64Exp $ segSpaceDims space

@@ -180,7 +180,7 @@    let params =         [ [C.cparam|__global int *global_failure|],-          [C.cparam|__global int *global_failure_args|]+          [C.cparam|__global typename int64_t *global_failure_args|]         ]       (func, cstate) =         genGPUCode FunMode (functionBody device_func) failures $@@ -312,7 +312,7 @@       failure_params =         [ [C.cparam|__global int *global_failure|],           [C.cparam|int failure_is_an_option|],-          [C.cparam|__global int *global_failure_args|]+          [C.cparam|__global typename int64_t *global_failure_args|]         ]        params =@@ -780,6 +780,10 @@       let setArgs _ [] = return []           setArgs i (ErrorString {} : parts') = setArgs i parts'           setArgs i (ErrorInt32 x : parts') = do+            x' <- GC.compileExp x+            stms <- setArgs (i + 1) parts'+            return $ [C.cstm|global_failure_args[$int:i] = (typename int64_t)$exp:x';|] : stms+          setArgs i (ErrorInt64 x : parts') = do             x' <- GC.compileExp x             stms <- setArgs (i + 1) parts'             return $ [C.cstm|global_failure_args[$int:i] = $exp:x';|] : stms

@@ -0,0 +1,101 @@+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE TypeFamilies #-}++module Futhark.CodeGen.ImpGen.Multicore+  ( Futhark.CodeGen.ImpGen.Multicore.compileProg,+    Warnings,+  )+where++import qualified Futhark.CodeGen.ImpCode.Multicore as Imp+import Futhark.CodeGen.ImpGen+import Futhark.CodeGen.ImpGen.Multicore.Base+import Futhark.CodeGen.ImpGen.Multicore.SegHist+import Futhark.CodeGen.ImpGen.Multicore.SegMap+import Futhark.CodeGen.ImpGen.Multicore.SegRed+import Futhark.CodeGen.ImpGen.Multicore.SegScan+import Futhark.IR.MCMem+import Futhark.MonadFreshNames+import Prelude hiding (quot, rem)++-- GCC supported primitve atomic Operations+-- TODO: Add support for 1, 2, and 16 bytes too+gccAtomics :: AtomicBinOp+gccAtomics = flip lookup cpu+  where+    cpu =+      [ (Add Int32 OverflowUndef, Imp.AtomicAdd Int32),+        (Sub Int32 OverflowUndef, Imp.AtomicSub Int32),+        (And Int32, Imp.AtomicAnd Int32),+        (Xor Int32, Imp.AtomicXor Int32),+        (Or Int32, Imp.AtomicOr Int32),+        (Add Int64 OverflowUndef, Imp.AtomicAdd Int64),+        (Sub Int64 OverflowUndef, Imp.AtomicSub Int64),+        (And Int64, Imp.AtomicAnd Int64),+        (Xor Int64, Imp.AtomicXor Int64),+        (Or Int64, Imp.AtomicOr Int64)+      ]++compileProg ::+  MonadFreshNames m =>+  Prog MCMem ->+  m (Warnings, Imp.Definitions Imp.Multicore)+compileProg = Futhark.CodeGen.ImpGen.compileProg (HostEnv gccAtomics) ops Imp.DefaultSpace+  where+    ops = defaultOperations opCompiler+    opCompiler dest (Alloc e space) = compileAlloc dest e space+    opCompiler dest (Inner op) = compileMCOp dest op++compileMCOp ::+  Pattern MCMem ->+  MCOp MCMem () ->+  ImpM MCMem HostEnv Imp.Multicore ()+compileMCOp _ (OtherOp ()) = pure ()+compileMCOp pat (ParOp par_op op) = do+  let space = getSpace op+  dPrimV_ (segFlat space) (0 :: Imp.TExp Int32)+  iterations <- getIterationDomain op space+  nsubtasks <- dPrim "num_tasks" $ IntType Int32+  seq_code <- compileSegOp pat op nsubtasks+  retvals <- getReturnParams pat op++  let scheduling_info = Imp.SchedulerInfo (tvVar nsubtasks) (untyped iterations)++  par_code <- case par_op of+    Just nested_op -> do+      let space' = getSpace nested_op+      dPrimV_ (segFlat space') (0 :: Imp.TExp Int32)+      compileSegOp pat nested_op nsubtasks+    Nothing -> return mempty++  let par_task = case par_op of+        Just nested_op -> Just $ Imp.ParallelTask par_code $ segFlat $ getSpace nested_op+        Nothing -> Nothing++  let non_free =+        ( [segFlat space, tvVar nsubtasks]+            ++ map Imp.paramName retvals+        )+          ++ case par_op of+            Just nested_op ->+              [segFlat $ getSpace nested_op]+            Nothing -> []++  s <- segOpString op+  free_params <- freeParams (par_code <> seq_code) non_free+  let seq_task = Imp.ParallelTask seq_code (segFlat space)+  emit $ Imp.Op $ Imp.Segop s free_params seq_task par_task retvals $ scheduling_info (decideScheduling' op seq_code)++compileSegOp ::+  Pattern MCMem ->+  SegOp () MCMem ->+  TV Int32 ->+  ImpM MCMem HostEnv Imp.Multicore Imp.Code+compileSegOp pat (SegHist _ space histops _ kbody) ntasks =+  compileSegHist pat space histops kbody ntasks+compileSegOp pat (SegScan _ space scans _ kbody) ntasks =+  compileSegScan pat space scans kbody ntasks+compileSegOp pat (SegRed _ space reds _ kbody) ntasks =+  compileSegRed pat space reds kbody ntasks+compileSegOp pat (SegMap _ space _ kbody) _ =+  compileSegMap pat space kbody

@@ -0,0 +1,458 @@+module Futhark.CodeGen.ImpGen.Multicore.Base+  ( toParam,+    compileKBody,+    extractAllocations,+    compileThreadResult,+    HostEnv (..),+    AtomicBinOp,+    MulticoreGen,+    decideScheduling,+    decideScheduling',+    groupResultArrays,+    renameSegBinOp,+    resultArrays,+    freeParams,+    renameHistOpLambda,+    atomicUpdateLocking,+    AtomicUpdate (..),+    Locking (..),+    getSpace,+    getIterationDomain,+    getReturnParams,+    segOpString,+  )+where++import Control.Monad+import Data.Bifunctor+import Data.List (elemIndex, find)+import Data.Maybe+import qualified Futhark.CodeGen.ImpCode.Multicore as Imp+import Futhark.CodeGen.ImpGen+import Futhark.Error+import Futhark.IR.MCMem+import Futhark.Transform.Rename+import Futhark.Util (maybeNth)+import Prelude hiding (quot, rem)++-- | Is there an atomic t'BinOp' corresponding to this t'BinOp'?+type AtomicBinOp =+  BinOp ->+  Maybe (VName -> VName -> Imp.Count Imp.Elements (Imp.TExp Int32) -> Imp.Exp -> Imp.AtomicOp)++newtype HostEnv = HostEnv+  {hostAtomics :: AtomicBinOp}++type MulticoreGen = ImpM MCMem HostEnv Imp.Multicore++segOpString :: SegOp () MCMem -> MulticoreGen String+segOpString SegMap {} = return "segmap"+segOpString SegRed {} = return "segred"+segOpString SegScan {} = return "segscan"+segOpString SegHist {} = return "seghist"++toParam :: VName -> TypeBase shape u -> MulticoreGen Imp.Param+toParam name (Prim pt) = return $ Imp.ScalarParam name pt+toParam name (Mem space) = return $ Imp.MemParam name space+toParam name Array {} = do+  name_entry <- lookupVar name+  case name_entry of+    ArrayVar _ (ArrayEntry (MemLocation mem _ _) _) ->+      return $ Imp.MemParam mem DefaultSpace+    _ -> error $ "[toParam] Could not handle array for " ++ show name++getSpace :: SegOp () MCMem -> SegSpace+getSpace (SegHist _ space _ _ _) = space+getSpace (SegRed _ space _ _ _) = space+getSpace (SegScan _ space _ _ _) = space+getSpace (SegMap _ space _ _) = space++getIterationDomain :: SegOp () MCMem -> SegSpace -> MulticoreGen (Imp.TExp Int64)+getIterationDomain SegMap {} space = do+  let ns = map snd $ unSegSpace space+      ns_64 = map (sExt64 . toInt32Exp) ns+  return $ product ns_64+getIterationDomain _ space = do+  let ns = map snd $ unSegSpace space+      ns_64 = map (sExt64 . toInt32Exp) ns+  case unSegSpace space of+    [_] -> return $ product ns_64+    -- A segmented SegOp is over the segments+    -- so we drop the last dimension, which is+    -- executed sequentially+    _ -> return $ product $ init ns_64++-- When the SegRed's return value is a scalar+-- we perform a call by value-result in the segop function+getReturnParams :: Pattern MCMem -> SegOp () MCMem -> MulticoreGen [Imp.Param]+getReturnParams pat SegRed {} = do+  let retvals = map patElemName $ patternElements pat+  retvals_ts <- mapM lookupType retvals+  zipWithM toParam retvals retvals_ts+getReturnParams _ _ = return mempty++renameSegBinOp :: [SegBinOp MCMem] -> MulticoreGen [SegBinOp MCMem]+renameSegBinOp segbinops =+  forM segbinops $ \(SegBinOp comm lam ne shape) -> do+    lam' <- renameLambda lam+    return $ SegBinOp comm lam' ne shape++compileKBody ::+  KernelBody MCMem ->+  ([(SubExp, [Imp.Exp])] -> ImpM MCMem () Imp.Multicore ()) ->+  ImpM MCMem () Imp.Multicore ()+compileKBody kbody red_cont =+  compileStms (freeIn $ kernelBodyResult kbody) (kernelBodyStms kbody) $ do+    let red_res = kernelBodyResult kbody+    red_cont $ zip (map kernelResultSubExp red_res) $ repeat []++compileThreadResult ::+  SegSpace ->+  PatElem MCMem ->+  KernelResult ->+  MulticoreGen ()+compileThreadResult space pe (Returns _ what) = do+  let is = map (Imp.vi64 . fst) $ unSegSpace space+  copyDWIMFix (patElemName pe) is what []+compileThreadResult _ _ ConcatReturns {} =+  compilerBugS "compileThreadResult: ConcatReturn unhandled."+compileThreadResult _ _ WriteReturns {} =+  compilerBugS "compileThreadResult: WriteReturns unhandled."+compileThreadResult _ _ TileReturns {} =+  compilerBugS "compileThreadResult: TileReturns unhandled."++freeVariables :: Imp.Code -> [VName] -> [VName]+freeVariables code names =+  namesToList $ freeIn code `namesSubtract` namesFromList names++freeParams :: Imp.Code -> [VName] -> MulticoreGen [Imp.Param]+freeParams code names = do+  let freeVars = freeVariables code names+  ts <- mapM lookupType freeVars+  zipWithM toParam freeVars ts++-- | Arrays for storing group results.+resultArrays :: String -> [SegBinOp MCMem] -> MulticoreGen [[VName]]+resultArrays s segops =+  forM segops $ \(SegBinOp _ lam _ shape) ->+    forM (lambdaReturnType lam) $ \t -> do+      let pt = elemType t+          full_shape = shape <> arrayShape t+      sAllocArray s pt full_shape DefaultSpace++-- | Arrays for storing group results shared between threads+groupResultArrays ::+  String ->+  SubExp ->+  [SegBinOp MCMem] ->+  MulticoreGen [[VName]]+groupResultArrays s num_threads reds =+  forM reds $ \(SegBinOp _ lam _ shape) ->+    forM (lambdaReturnType lam) $ \t -> do+      let pt = elemType t+          full_shape = Shape [num_threads] <> shape <> arrayShape t+      sAllocArray s pt full_shape DefaultSpace++isLoadBalanced :: Imp.Code -> Bool+isLoadBalanced (a Imp.:>>: b) = isLoadBalanced a && isLoadBalanced b+isLoadBalanced (Imp.For _ _ a) = isLoadBalanced a+isLoadBalanced (Imp.If _ a b) = isLoadBalanced a && isLoadBalanced b+isLoadBalanced (Imp.Comment _ a) = isLoadBalanced a+isLoadBalanced Imp.While {} = False+isLoadBalanced (Imp.Op (Imp.ParLoop _ _ _ code _ _ _)) = isLoadBalanced code+isLoadBalanced _ = True++segBinOpComm' :: [SegBinOp lore] -> Commutativity+segBinOpComm' = mconcat . map segBinOpComm++decideScheduling' :: SegOp () lore -> Imp.Code -> Imp.Scheduling+decideScheduling' SegHist {} _ = Imp.Static+decideScheduling' SegScan {} _ = Imp.Static+decideScheduling' (SegRed _ _ reds _ _) code =+  case segBinOpComm' reds of+    Commutative -> decideScheduling code+    Noncommutative -> Imp.Static+decideScheduling' SegMap {} code = decideScheduling code++decideScheduling :: Imp.Code -> Imp.Scheduling+decideScheduling code =+  if isLoadBalanced code+    then Imp.Static+    else Imp.Dynamic++-- | Try to extract invariant allocations.  If we assume that the+-- given 'Code' is the body of a 'SegOp', then it is always safe to+-- move the immediate allocations to the prebody.+extractAllocations :: Imp.Code -> (Imp.Code, Imp.Code)+extractAllocations segop_code = f segop_code+  where+    declared = Imp.declaredIn segop_code+    f (Imp.DeclareMem name space) =+      -- Hoisting declarations out is always safe.+      (Imp.DeclareMem name space, mempty)+    f (Imp.Allocate name size space)+      | not $ freeIn size `namesIntersect` declared =+        (Imp.Allocate name size space, mempty)+    f (x Imp.:>>: y) = f x <> f y+    f (Imp.While cond body) =+      (mempty, Imp.While cond body)+    f (Imp.For i bound body) =+      (mempty, Imp.For i bound body)+    f (Imp.Comment s code) =+      second (Imp.Comment s) (f code)+    f Imp.Free {} =+      mempty+    f (Imp.If cond tcode fcode) =+      let (ta, tcode') = f tcode+          (fa, fcode') = f fcode+       in (ta <> fa, Imp.If cond tcode' fcode')+    f (Imp.Op (Imp.ParLoop s i prebody body postbody free info)) =+      let (body_allocs, body') = extractAllocations body+          (free_allocs, here_allocs) = f body_allocs+          free' =+            filter+              ( not+                  . (`nameIn` Imp.declaredIn body_allocs)+                  . Imp.paramName+              )+              free+       in ( free_allocs,+            here_allocs+              <> Imp.Op (Imp.ParLoop s i prebody body' postbody free' info)+          )+    f code =+      (mempty, code)++-------------------------------+------- SegHist helpers -------+-------------------------------+renameHistOpLambda :: [HistOp MCMem] -> MulticoreGen [HistOp MCMem]+renameHistOpLambda hist_ops =+  forM hist_ops $ \(HistOp w rf dest neutral shape lam) -> do+    lam' <- renameLambda lam+    return $ HistOp w rf dest neutral shape lam'++-- | Locking strategy used for an atomic update.+data Locking = Locking+  { -- | Array containing the lock.+    lockingArray :: VName,+    -- | Value for us to consider the lock free.+    lockingIsUnlocked :: Imp.TExp Int32,+    -- | What to write when we lock it.+    lockingToLock :: Imp.TExp Int32,+    -- | What to write when we unlock it.+    lockingToUnlock :: Imp.TExp Int32,+    -- | A transformation from the logical lock index to the+    -- physical position in the array.  This can also be used+    -- to make the lock array smaller.+    lockingMapping :: [Imp.TExp Int64] -> [Imp.TExp Int64]+  }++-- | A function for generating code for an atomic update.  Assumes+-- that the bucket is in-bounds.+type DoAtomicUpdate lore r =+  [VName] -> [Imp.TExp Int64] -> MulticoreGen ()++-- | The mechanism that will be used for performing the atomic update.+-- Approximates how efficient it will be.  Ordered from most to least+-- efficient.+data AtomicUpdate lore r+  = AtomicPrim (DoAtomicUpdate lore r)+  | -- | Can be done by efficient swaps.+    AtomicCAS (DoAtomicUpdate lore r)+  | -- | Requires explicit locking.+    AtomicLocking (Locking -> DoAtomicUpdate lore r)++atomicUpdateLocking ::+  AtomicBinOp ->+  Lambda MCMem ->+  AtomicUpdate MCMem ()+atomicUpdateLocking atomicBinOp lam+  | Just ops_and_ts <- splitOp lam,+    all (\(_, t, _, _) -> supportedPrims $ primBitSize t) ops_and_ts =+    primOrCas ops_and_ts $ \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+      -- is still a binary operator which can be implemented by atomic+      -- compare-and-swap if 32 bits.+      forM_ (zip arrs ops_and_ts) $ \(a, (op, t, x, y)) -> do+        -- Common variables.+        old <- dPrim "old" t++        (arr', _a_space, bucket_offset) <- fullyIndexArray a bucket++        case opHasAtomicSupport (tvVar old) arr' (sExt32 <$> bucket_offset) op of+          Just f -> sOp $ f $ Imp.var y t+          Nothing ->+            atomicUpdateCAS t a (tvVar 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'+      atomic <$> atomicBinOp bop++    primOrCas ops+      | all isPrim ops = AtomicPrim+      | otherwise = AtomicCAS++    isPrim (op, _, _, _) = isJust $ atomicBinOp op+atomicUpdateLocking _ op+  | [Prim t] <- lambdaReturnType op,+    [xp, _] <- lambdaParams op,+    supportedPrims (primBitSize t) = AtomicCAS $ \[arr] bucket -> do+    old <- dPrim "old" t+    atomicUpdateCAS t arr (tvVar old) bucket (paramName xp) $+      compileBody' [xp] $ lambdaBody op+atomicUpdateLocking _ op = AtomicLocking $ \locking arrs bucket -> do+  old <- dPrim "old" int32+  continue <- dPrimVol "continue" int32 (0 :: Imp.TExp Int32)++  -- Correctly index into locks.+  (locks', _locks_space, locks_offset) <-+    fullyIndexArray (lockingArray locking) $ lockingMapping locking bucket++  -- Critical section+  let try_acquire_lock = do+        old <-- (0 :: Imp.TExp Int32)+        sOp $+          Imp.Atomic $+            Imp.AtomicCmpXchg+              int32+              (tvVar old)+              locks'+              (sExt32 <$> locks_offset)+              (tvVar continue)+              (untyped (lockingToLock locking))+      lock_acquired = tvExp continue+      -- Even the releasing is done with an atomic rather than a+      -- simple write, for memory coherency reasons.+      release_lock = do+        old <-- lockingToLock locking+        sOp $+          Imp.Atomic $+            Imp.AtomicCmpXchg+              int32+              (tvVar old)+              locks'+              (sExt32 <$> locks_offset)+              (tvVar continue)+              (untyped (lockingToUnlock locking))++  -- Preparing parameters. It is assumed that the caller has already+  -- filled the arr_params. We copy the current value to the+  -- accumulator parameters.+  let (acc_params, _arr_params) = splitAt (length arrs) $ lambdaParams op+      bind_acc_params =+        everythingVolatile $+          sComment "bind lhs" $+            forM_ (zip acc_params arrs) $ \(acc_p, arr) ->+              copyDWIMFix (paramName acc_p) [] (Var arr) bucket++  let op_body =+        sComment "execute operation" $+          compileBody' acc_params $ lambdaBody op++      do_hist =+        everythingVolatile $+          sComment "update global result" $+            zipWithM_ (writeArray bucket) arrs $ map (Var . paramName) acc_params++  -- While-loop: Try to insert your value+  sWhile (tvExp continue .==. 0) $ do+    try_acquire_lock+    sUnless (lock_acquired .==. 0) $ do+      dLParams acc_params+      bind_acc_params+      op_body+      do_hist+      release_lock+  where+    writeArray bucket arr val = copyDWIMFix arr bucket val []++atomicUpdateCAS ::+  PrimType ->+  VName ->+  VName ->+  [Imp.TExp Int64] ->+  VName ->+  MulticoreGen () ->+  MulticoreGen ()+atomicUpdateCAS t arr old bucket x do_op = do+  -- Code generation target:+  --+  -- old = d_his[idx];+  -- do {+  --   assumed = old;+  --   x = do_op(assumed, y);+  --   old = atomicCAS(&d_his[idx], assumed, tmp);+  -- } while(assumed != old);+  run_loop <- dPrimV "run_loop" (0 :: Imp.TExp Int32)+  everythingVolatile $ copyDWIMFix old [] (Var arr) bucket+  (arr', _a_space, bucket_offset) <- fullyIndexArray arr bucket++  bytes <- toIntegral $ primBitSize t+  (to, from) <- getBitConvertFunc $ primBitSize t+  -- While-loop: Try to insert your value+  let (toBits, _fromBits) =+        case t of+          FloatType _ ->+            ( \v -> Imp.FunExp to [v] bytes,+              \v -> Imp.FunExp from [v] t+            )+          _ -> (id, id)++  sWhile (tvExp run_loop .==. 0) $ do+    x <~~ Imp.var old t+    do_op -- Writes result into x+    sOp $+      Imp.Atomic $+        Imp.AtomicCmpXchg+          bytes+          old+          arr'+          (sExt32 <$> bucket_offset)+          (tvVar run_loop)+          (toBits (Imp.var x t))++-- | Horizontally fission a lambda that models a binary operator.+splitOp :: ASTLore lore => Lambda lore -> Maybe [(BinOp, PrimType, VName, VName)]+splitOp lam = mapM splitStm $ bodyResult $ lambdaBody lam+  where+    n = length $ lambdaReturnType lam+    splitStm (Var res) = do+      Let (Pattern [] [pe]) _ (BasicOp (BinOp op (Var x) (Var y))) <-+        find (([res] ==) . patternNames . stmPattern) $+          stmsToList $ bodyStms $ lambdaBody lam+      i <- Var res `elemIndex` bodyResult (lambdaBody lam)+      xp <- maybeNth i $ lambdaParams lam+      yp <- maybeNth (n + i) $ lambdaParams lam+      guard $ paramName xp == x+      guard $ paramName yp == y+      Prim t <- Just $ patElemType pe+      return (op, t, paramName xp, paramName yp)+    splitStm _ = Nothing++-- TODO for supporting 8 and 16 bits (and 128)+-- we need a functions for converting to and from bits+getBitConvertFunc :: Int -> MulticoreGen (String, String)+-- getBitConvertFunc 8 = return $ ("to_bits8, from_bits8")+-- getBitConvertFunc 16 = return $ ("to_bits8, from_bits8")+getBitConvertFunc 32 = return ("to_bits32", "from_bits32")+getBitConvertFunc 64 = return ("to_bits64", "from_bits64")+getBitConvertFunc b = error $ "number of bytes is not supported " ++ pretty b++supportedPrims :: Int -> Bool+supportedPrims 8 = True+supportedPrims 16 = True+supportedPrims 32 = True+supportedPrims 64 = True+supportedPrims _ = False++-- Supported bytes lengths by GCC (and clang) compiler+toIntegral :: Int -> MulticoreGen PrimType+toIntegral 8 = return int8+toIntegral 16 = return int16+toIntegral 32 = return int32+toIntegral 64 = return int64+toIntegral b = error $ "number of bytes is not supported for CAS - " ++ pretty b

@@ -0,0 +1,359 @@+module Futhark.CodeGen.ImpGen.Multicore.SegHist+  ( compileSegHist,+  )+where++import Control.Monad+import Data.List (zip4, zip5)+import qualified Futhark.CodeGen.ImpCode.Multicore as Imp+import Futhark.CodeGen.ImpGen+import Futhark.CodeGen.ImpGen.Multicore.Base+import Futhark.CodeGen.ImpGen.Multicore.SegRed (compileSegRed')+import Futhark.IR.MCMem+import Futhark.MonadFreshNames+import Futhark.Util (chunks, splitFromEnd, takeLast)+import Futhark.Util.IntegralExp (rem)+import Prelude hiding (quot, rem)++compileSegHist ::+  Pattern MCMem ->+  SegSpace ->+  [HistOp MCMem] ->+  KernelBody MCMem ->+  TV Int32 ->+  MulticoreGen Imp.Code+compileSegHist pat space histops kbody nsubtasks+  | [_] <- unSegSpace space =+    nonsegmentedHist pat space histops kbody nsubtasks+  | otherwise =+    segmentedHist pat space histops kbody++-- | Split some list into chunks equal to the number of values+-- returned by each 'SegBinOp'+segHistOpChunks :: [HistOp lore] -> [a] -> [[a]]+segHistOpChunks = chunks . map (length . histNeutral)++nonsegmentedHist ::+  Pattern MCMem ->+  SegSpace ->+  [HistOp MCMem] ->+  KernelBody MCMem ->+  TV Int32 ->+  MulticoreGen Imp.Code+nonsegmentedHist pat space histops kbody num_histos = do+  let ns = map snd $ unSegSpace space+      ns_64 = map toInt64Exp ns+      num_histos' = tvExp num_histos+      hist_width = toInt64Exp $ histWidth $ head histops+      use_subhistogram = sExt64 num_histos' * hist_width .<=. product ns_64++  histops' <- renameHistOpLambda histops++  collect $ do+    flat_idx <- dPrim "iter" int64+    sIf+      use_subhistogram+      (subHistogram pat flat_idx space histops num_histos kbody)+      (atomicHistogram pat flat_idx space histops' kbody)++-- |+-- Atomic Histogram approach+-- The implementation has three sub-strategies depending on the+-- type of the operator+-- 1. If values are integral scalars, a direct-supported atomic update is used.+-- 2. If values are on one memory location, e.g. a float, then a+-- CAS operation is used to perform the update, where the float is+-- casted to an integral scalar.+-- 1. and 2. currently only works for 32-bit and 64-bit types,+-- but GCC has support for 8-, 16- and 128- bit types as well.+-- 3. Otherwise a locking based approach is used+onOpAtomic :: HistOp MCMem -> MulticoreGen ([VName] -> [Imp.TExp Int64] -> MulticoreGen ())+onOpAtomic op = do+  atomics <- hostAtomics <$> askEnv+  let lambda = histOp op+      do_op = atomicUpdateLocking atomics lambda+  case do_op of+    AtomicPrim f -> return f+    AtomicCAS f -> return f+    AtomicLocking f -> do+      -- Allocate a static array of locks+      -- as in the GPU backend+      let num_locks = 100151 -- This number is taken from the GPU backend+          dims =+            map toInt64Exp $+              shapeDims (histShape op) ++ [histWidth op]+      locks <-+        sStaticArray "hist_locks" DefaultSpace int32 $+          Imp.ArrayZeros num_locks+      let l' = Locking locks 0 1 0 (pure . (`rem` fromIntegral num_locks) . flattenIndex dims)+      return $ f l'++atomicHistogram ::+  Pattern MCMem ->+  TV Int64 ->+  SegSpace ->+  [HistOp MCMem] ->+  KernelBody MCMem ->+  MulticoreGen ()+atomicHistogram pat flat_idx space histops kbody = do+  let (is, ns) = unzip $ unSegSpace space+      ns_64 = map toInt64Exp ns+  let num_red_res = length histops + sum (map (length . histNeutral) histops)+      (all_red_pes, map_pes) = splitAt num_red_res $ patternValueElements pat++  atomicOps <- mapM onOpAtomic histops++  body <- collect $ do+    zipWithM_ dPrimV_ is $ unflattenIndex ns_64 $ tvExp flat_idx+    compileStms mempty (kernelBodyStms kbody) $ do+      let (red_res, map_res) = splitFromEnd (length map_pes) $ kernelBodyResult kbody+          perOp = chunks $ map (length . histDest) histops+          (buckets, vs) = splitAt (length histops) red_res++      let pes_per_op = chunks (map (length . histDest) histops) all_red_pes+      forM_ (zip5 histops (perOp vs) buckets atomicOps pes_per_op) $+        \(HistOp dest_w _ _ _ shape lam, vs', bucket, do_op, dest_res) -> do+          let (_is_params, vs_params) = splitAt (length vs') $ lambdaParams lam+              dest_w' = toInt64Exp dest_w+              bucket' = toInt64Exp $ kernelResultSubExp bucket+              bucket_in_bounds = bucket' .<. dest_w' .&&. 0 .<=. bucket'++          sComment "save map-out results" $+            forM_ (zip map_pes map_res) $ \(pe, res) ->+              copyDWIMFix (patElemName pe) (map Imp.vi64 is) (kernelResultSubExp res) []++          sComment "perform updates" $+            sWhen bucket_in_bounds $ do+              let bucket_is = map Imp.vi64 (init is) ++ [bucket']+              dLParams $ lambdaParams lam+              sLoopNest shape $ \is' -> do+                forM_ (zip vs_params vs') $ \(p, res) ->+                  copyDWIMFix (paramName p) [] (kernelResultSubExp res) is'+                do_op (map patElemName dest_res) (bucket_is ++ is')++  free_params <- freeParams body (segFlat space : [tvVar flat_idx])+  emit $ Imp.Op $ Imp.ParLoop "atomic_seg_hist" (tvVar flat_idx) mempty body mempty free_params $ segFlat space++updateHisto :: HistOp MCMem -> [VName] -> [Imp.TExp Int64] -> MulticoreGen ()+updateHisto op arrs bucket = do+  let acc_params = take (length arrs) $ lambdaParams $ histOp op+      bind_acc_params =+        forM_ (zip acc_params arrs) $ \(acc_p, arr) ->+          copyDWIMFix (paramName acc_p) [] (Var arr) bucket+      op_body = compileBody' [] $ lambdaBody $ histOp op+      writeArray arr val = copyDWIMFix arr bucket val []+      do_hist = zipWithM_ writeArray arrs $ bodyResult $ lambdaBody $ histOp op++  sComment "Start of body" $ do+    dLParams acc_params+    bind_acc_params+    op_body+    do_hist++-- Generates num_histos sub-histograms of the size+-- of the destination histogram+-- Then for each chunk of the input each subhistogram+-- is computed and finally combined through a segmented reduction+-- across the histogram indicies.+-- This is expected to be fast if len(histDest) is small+subHistogram ::+  Pattern MCMem ->+  TV Int64 ->+  SegSpace ->+  [HistOp MCMem] ->+  TV Int32 ->+  KernelBody MCMem ->+  MulticoreGen ()+subHistogram pat flat_idx space histops num_histos kbody = do+  emit $ Imp.DebugPrint "subHistogram segHist" Nothing++  let (is, ns) = unzip $ unSegSpace space+      ns_64 = map toInt64Exp ns++  let pes = patternElements pat+      num_red_res = length histops + sum (map (length . histNeutral) histops)+      map_pes = drop num_red_res pes+      per_red_pes = segHistOpChunks histops $ patternValueElements pat++  -- Allocate array of subhistograms in the calling thread.  Each+  -- tasks will work in its own private allocations (to avoid false+  -- sharing), but this is where they will ultimately copy their+  -- results.+  global_subhistograms <- forM histops $ \histop ->+    forM (histType histop) $ \t -> do+      let shape = Shape [tvSize num_histos] <> arrayShape t+      sAllocArray "subhistogram" (elemType t) shape DefaultSpace++  let tid' = Imp.vi64 $ segFlat space+      flat_idx' = tvExp flat_idx++  (local_subhistograms, prebody) <- collect' $ do+    zipWithM_ dPrimV_ is $ unflattenIndex ns_64 $ sExt64 flat_idx'++    forM (zip per_red_pes histops) $ \(pes', histop) -> do+      op_local_subhistograms <- forM (histType histop) $ \t ->+        sAllocArray "subhistogram" (elemType t) (arrayShape t) DefaultSpace++      forM_ (zip3 pes' op_local_subhistograms (histNeutral histop)) $ \(pe, hist, ne) ->+        -- First thread initializes histogram with dest vals. Others+        -- initialize with neutral element+        sIf+          (tid' .==. 0)+          (copyDWIMFix hist [] (Var $ patElemName pe) [])+          ( sFor "i" (toInt64Exp $ histWidth histop) $ \i ->+              sLoopNest (histShape histop) $ \vec_is ->+                copyDWIMFix hist (i : vec_is) ne []+          )++      return op_local_subhistograms++  -- Generate loop body of parallel function+  body <- collect $ do+    zipWithM_ dPrimV_ is $ unflattenIndex ns_64 $ sExt64 flat_idx'+    compileStms mempty (kernelBodyStms kbody) $ do+      let (red_res, map_res) = splitFromEnd (length map_pes) $ kernelBodyResult kbody+          (buckets, vs) = splitAt (length histops) red_res+          perOp = chunks $ map (length . histDest) histops++      sComment "save map-out results" $+        forM_ (zip map_pes map_res) $ \(pe, res) ->+          copyDWIMFix+            (patElemName pe)+            (map Imp.vi64 is)+            (kernelResultSubExp res)+            []++      forM_ (zip4 histops local_subhistograms buckets (perOp vs)) $+        \( histop@(HistOp dest_w _ _ _ shape lam),+           histop_subhistograms,+           bucket,+           vs'+           ) -> do+            let bucket' = toInt64Exp $ kernelResultSubExp bucket+                dest_w' = toInt64Exp dest_w+                bucket_in_bounds = bucket' .<. dest_w' .&&. 0 .<=. bucket'+                vs_params = takeLast (length vs') $ lambdaParams lam+                bucket_is = [bucket']++            sComment "perform updates" $+              sWhen bucket_in_bounds $ do+                dLParams $ lambdaParams lam+                sLoopNest shape $ \is' -> do+                  forM_ (zip vs_params vs') $ \(p, res) ->+                    copyDWIMFix (paramName p) [] (kernelResultSubExp res) is'+                  updateHisto histop histop_subhistograms (bucket_is ++ is')++  -- Copy the task-local subhistograms to the global subhistograms,+  -- where they will be combined.+  postbody <- collect $+    forM_ (zip (concat global_subhistograms) (concat local_subhistograms)) $+      \(global, local) -> copyDWIMFix global [tid'] (Var local) []++  free_params <- freeParams (prebody <> body <> postbody) (segFlat space : [tvVar flat_idx])+  let (body_allocs, body') = extractAllocations body+  emit $ Imp.Op $ Imp.ParLoop "seghist_stage_1" (tvVar flat_idx) (body_allocs <> prebody) body' postbody free_params $ segFlat space++  -- Perform a segmented reduction over the subhistograms+  forM_ (zip3 per_red_pes global_subhistograms histops) $ \(red_pes, hists, op) -> do+    bucket_id <- newVName "bucket_id"+    subhistogram_id <- newVName "subhistogram_id"++    let num_buckets = histWidth op+        segred_space =+          SegSpace (segFlat space) $+            segment_dims+              ++ [(bucket_id, num_buckets)]+              ++ [(subhistogram_id, tvSize num_histos)]++        segred_op = SegBinOp Noncommutative (histOp op) (histNeutral op) (histShape op)++    nsubtasks_red <- dPrim "num_tasks" $ IntType Int32+    red_code <- compileSegRed' (Pattern [] red_pes) segred_space [segred_op] nsubtasks_red $ \red_cont ->+      red_cont $+        flip map hists $ \subhisto ->+          ( Var subhisto,+            map Imp.vi64 $+              map fst segment_dims ++ [subhistogram_id, bucket_id]+          )++    let ns_red = map (toInt64Exp . snd) $ unSegSpace segred_space+        iterations = product $ init ns_red -- The segmented reduction is sequential over the inner most dimension+        scheduler_info = Imp.SchedulerInfo (tvVar nsubtasks_red) (untyped iterations) Imp.Static+        red_task = Imp.ParallelTask red_code $ segFlat space+    free_params_red <- freeParams red_code [segFlat space, tvVar nsubtasks_red]+    emit $ Imp.Op $ Imp.Segop "seghist_red" free_params_red red_task Nothing mempty scheduler_info+  where+    segment_dims = init $ unSegSpace space++-- This implementation for a Segmented Hist only+-- parallelize over the segments,+-- where each segment is updated sequentially.+segmentedHist ::+  Pattern MCMem ->+  SegSpace ->+  [HistOp MCMem] ->+  KernelBody MCMem ->+  MulticoreGen Imp.Code+segmentedHist pat space histops kbody = do+  emit $ Imp.DebugPrint "Segmented segHist" Nothing+  -- Iteration variable over the segments+  segments_i <- dPrim "segment_iter" $ IntType Int64+  collect $ do+    par_body <- compileSegHistBody (tvExp segments_i) pat space histops kbody+    free_params <- freeParams par_body [segFlat space, tvVar segments_i]+    let (body_allocs, body') = extractAllocations par_body+    emit $ Imp.Op $ Imp.ParLoop "segmented_hist" (tvVar segments_i) body_allocs body' mempty free_params $ segFlat space++compileSegHistBody ::+  Imp.TExp Int64 ->+  Pattern MCMem ->+  SegSpace ->+  [HistOp MCMem] ->+  KernelBody MCMem ->+  MulticoreGen Imp.Code+compileSegHistBody idx pat space histops kbody = do+  let (is, ns) = unzip $ unSegSpace space+      ns_64 = map toInt64Exp ns++  let num_red_res = length histops + sum (map (length . histNeutral) histops)+      map_pes = drop num_red_res $ patternValueElements pat+      per_red_pes = segHistOpChunks histops $ patternValueElements pat++  collect $ do+    let inner_bound = last ns_64+    sFor "i" inner_bound $ \i -> do+      zipWithM_ dPrimV_ (init is) $ unflattenIndex (init ns_64) idx+      dPrimV_ (last is) i++      compileStms mempty (kernelBodyStms kbody) $ do+        let (red_res, map_res) =+              splitFromEnd (length map_pes) $+                map kernelResultSubExp $ kernelBodyResult kbody+            (buckets, vs) = splitAt (length histops) red_res+            perOp = chunks $ map (length . histDest) histops++        forM_ (zip4 per_red_pes histops (perOp vs) buckets) $+          \(red_pes, HistOp dest_w _ _ _ shape lam, vs', bucket) -> do+            let (is_params, vs_params) = splitAt (length vs') $ lambdaParams lam+                bucket' = toInt64Exp bucket+                dest_w' = toInt64Exp dest_w+                bucket_in_bounds = bucket' .<. dest_w' .&&. 0 .<=. bucket'++            sComment "save map-out results" $+              forM_ (zip map_pes map_res) $ \(pe, res) ->+                copyDWIMFix (patElemName pe) (map Imp.vi64 is) res []++            sComment "perform updates" $+              sWhen bucket_in_bounds $ do+                dLParams $ lambdaParams lam+                sLoopNest shape $ \vec_is -> do+                  -- Index+                  let buck = toInt64Exp bucket+                  forM_ (zip red_pes is_params) $ \(pe, p) ->+                    copyDWIMFix (paramName p) [] (Var $ patElemName pe) (map Imp.vi64 (init is) ++ [buck] ++ vec_is)+                  -- Value at index+                  forM_ (zip vs_params vs') $ \(p, v) ->+                    copyDWIMFix (paramName p) [] v vec_is+                  compileStms mempty (bodyStms $ lambdaBody lam) $+                    forM_ (zip red_pes $ bodyResult $ lambdaBody lam) $+                      \(pe, se) -> copyDWIMFix (patElemName pe) (map Imp.vi64 (init is) ++ [buck] ++ vec_is) se []