futhark 0.25.10 → 0.25.11

raw patch · 75 files changed

+2284/−2109 lines, 75 filesdep ~futhark-manifestPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

Dependency ranges changed: futhark-manifest

API changes (from Hackage documentation)

Files

• CHANGELOG.md +19/−0
• docs/c-api.rst +14/−2
• docs/language-reference.rst +2/−2
• docs/man/futhark-cuda.rst +4/−5
• docs/man/futhark-opencl.rst +9/−5
• docs/usage.rst +12/−4
• futhark.cabal +2/−2
• prelude/array.fut +12/−0
• rts/c/atomics.h +48/−48
• rts/c/backends/cuda.h +37/−26
• rts/c/backends/hip.h +34/−23
• rts/c/backends/opencl.h +46/−36
• rts/c/context.h +1/−1
• rts/c/gpu.h +10/−10
• rts/c/scheduler.h +12/−13
• rts/c/server.h +2/−2
• rts/cuda/prelude.cu +4/−4
• rts/opencl/copy.cl +1/−1
• rts/opencl/prelude.cl +3/−1
• rts/opencl/transpose.cl +55/−55
• rts/python/opencl.py +19/−19
• src/Futhark/CLI/Autotune.hs +4/−4
• src/Futhark/CodeGen/Backends/CCUDA.hs +0/−5
• src/Futhark/CodeGen/Backends/COpenCL.hs +2/−7
• src/Futhark/CodeGen/Backends/GPU.hs +36/−14
• src/Futhark/CodeGen/Backends/GenericC.hs +1/−1
• src/Futhark/CodeGen/Backends/HIP.hs +0/−5
• src/Futhark/CodeGen/Backends/MulticoreC.hs +1/−0
• src/Futhark/CodeGen/Backends/PyOpenCL.hs +19/−19
• src/Futhark/CodeGen/Backends/PyOpenCL/Boilerplate.hs +4/−4
• src/Futhark/CodeGen/ImpCode/GPU.hs +17/−17
• src/Futhark/CodeGen/ImpCode/OpenCL.hs +3/−3
• src/Futhark/CodeGen/ImpGen/GPU.hs +23/−23
• src/Futhark/CodeGen/ImpGen/GPU/Base.hs +167/−168
• src/Futhark/CodeGen/ImpGen/GPU/Block.hs +743/−0
• src/Futhark/CodeGen/ImpGen/GPU/Group.hs +0/−743
• src/Futhark/CodeGen/ImpGen/GPU/SegHist.hs +97/−97
• src/Futhark/CodeGen/ImpGen/GPU/SegMap.hs +15/−15
• src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs +227/−226
• src/Futhark/CodeGen/ImpGen/GPU/SegScan/SinglePass.hs +59/−58
• src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs +48/−48
• src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs +24/−24
• src/Futhark/CodeGen/OpenCL/Heuristics.hs +16/−16
• src/Futhark/IR/GPU.hs +52/−0
• src/Futhark/IR/GPU/Op.hs +62/−62
• src/Futhark/IR/GPU/Simplify.hs +2/−2
• src/Futhark/IR/GPU/Sizes.hs +14/−13
• src/Futhark/IR/Mem.hs +7/−9
• src/Futhark/IR/Mem/IxFun.hs +24/−11
• src/Futhark/IR/Mem/LMAD.hs +5/−4
• src/Futhark/IR/Mem/Simplify.hs +2/−2
• src/Futhark/IR/Parse.hs +10/−11
• src/Futhark/IR/SOACS/SOAC.hs +14/−18
• src/Futhark/IR/SegOp.hs +3/−3
• src/Futhark/IR/TypeCheck.hs +2/−2
• src/Futhark/Internalise/Exps.hs +7/−0
• src/Futhark/Optimise/ArrayShortCircuiting/ArrayCoalescing.hs +54/−52
• src/Futhark/Optimise/BlkRegTiling.hs +16/−16
• src/Futhark/Optimise/DoubleBuffer.hs +3/−3
• src/Futhark/Optimise/Fusion/TryFusion.hs +1/−2
• src/Futhark/Optimise/MergeGPUBodies.hs +17/−17
• src/Futhark/Optimise/Simplify/Rules/Loop.hs +9/−9
• src/Futhark/Optimise/TileLoops.hs +34/−34
• src/Futhark/Optimise/TileLoops/Shared.hs +1/−1
• src/Futhark/Pass/ExpandAllocations.hs +21/−21
• src/Futhark/Pass/ExplicitAllocations.hs +16/−14
• src/Futhark/Pass/ExplicitAllocations/GPU.hs +9/−9
• src/Futhark/Pass/ExtractKernels.hs +4/−4
• src/Futhark/Pass/ExtractKernels/Intragroup.hs +11/−11
• src/Futhark/Pass/ExtractKernels/StreamKernel.hs +13/−13
• src/Futhark/Pass/ExtractKernels/ToGPU.hs +2/−2
• src/Futhark/Pass/KernelBabysitting.hs +1/−1
• src/Language/Futhark/Interpreter.hs +1/−0
• src/Language/Futhark/Prop.hs +9/−1
• unittests/Futhark/Optimise/MemoryBlockMerging/GreedyColoringTests.hs +6/−6

@@ -5,6 +5,25 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/) and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.html). +## [0.25.11]++### Added++* New prelude function: `manifest`.  For doing subtle things to memory.++* The GPU backends now handle up to 20 operators in a single fused+  reduction.++* CUDA/HIP terminology for GPU concepts (e.g. "thread block") is now+  used in all public interfaces. The OpenCL names are still supported+  for backwards compatibility.++* More fusion across array slicing.++### Fixed++* Compatibility with CUDA versions prior than 12.+ ## [0.25.10]  ### Added

@@ -461,13 +461,25 @@  The following functions are not interesting to most users. +.. c:function:: void futhark_context_config_set_default_thread_block_size(struct futhark_context_config *cfg, int size)++   Set the default number of work-items in a thread block.+ .. c:function:: void futhark_context_config_set_default_group_size(struct futhark_context_config *cfg, int size) -   Set the default number of work-items in a work-group.+   Identical to+   :c:func:`futhark_context_config_set_default_thread_block_size`;+   provided for backwards compatibility. +.. c:function:: void futhark_context_config_set_default_grid_size(struct futhark_context_config *cfg, int num)++   Set the default number of thread blocks used for kernels.+ .. c:function:: void futhark_context_config_set_default_num_groups(struct futhark_context_config *cfg, int num) -   Set the default number of work-groups used for kernels.+   Identical to+   :c:func:`futhark_context_config_set_default_grid_size`;+   provided for backwards compatibility.  .. c:function:: void futhark_context_config_set_default_tile_size(struct futhark_context_config *cfg, int num) 

@@ -1745,13 +1745,13 @@ ``incremental_flattening(no_intra)`` .................................... -When using incremental flattening, do not generate the "intra-group+When using incremental flattening, do not generate the "intra-block parallelism" version for the attributed SOACs.  ``incremental_flattening(only_intra)`` ...................................... -When using incremental flattening, *only* generate the "intra-group+When using incremental flattening, *only* generate the "intra-block parallelism" version of the attributed SOACs.  **Beware**: the resulting program will fail to run if the inner parallelism does not fit on the device.

@@ -51,9 +51,8 @@ ==================  Generated executables accept the same options as those generated by-:ref:`futhark-c(1)`.  The ``-t`` option behaves as with-:ref:`futhark-opencl(1)`.  For commonality, the options use OpenCL-nomenclature ("group" instead of "thread block").+:ref:`futhark-c(1)`. The ``-t`` option behaves as with+:ref:`futhark-opencl(1)`.  The following additional options are accepted. @@ -61,12 +60,12 @@    Print help text to standard output and exit. ---default-group-size=INT+--default-thread-block-size=INT    The default size of thread blocks that are launched.  Capped to the   hardware limit if necessary. ---default-num-groups=INT+--default-num-thread-blocks=INT    The default number of thread blocks that are launched. 

@@ -25,6 +25,10 @@ macOS).  If using ``--library``, you will need to do the same when linking the final binary. +The GPU terminology used is derived from CUDA nomenclature (e.g.+"thread block" instead of "workgroup"), but OpenCL nomenclature is+also supported for compatibility.+ OPTIONS ======= @@ -64,14 +68,14 @@   options are supported.  Be careful - some options can easily   result in invalid results. ---default-group-size=INT+--default-thread-block-size=INT, --default-group-size=INT -  The default size of OpenCL workgroups that are launched.  Capped-  to the hardware limit if necessary.+  The default size of thread blocks that are launched. Capped to the+  hardware limit if necessary. ---default-num-groups=INT+--default-num-thread-blocks, --default-num-groups=INT -  The default number of OpenCL workgroups that are launched.+  The default number of thread blocks that are launched.  --default-threshold=INT 

@@ -166,6 +166,15 @@     special string ``#k``, where ``k`` is an integer, can be used to     pick the *k*-th device, numbered from zero. +  ``--default-thread-block-size INT``++    The default size of GPU thread blocks that are launched. Capped to+    the hardware limit if necessary.++  ``--default-num-thread-blocks INT``++    The default number of GPU thread blocks that are launched.+   ``-P/--profile``      Measure the time taken by various GPU operations (such as kernels)@@ -305,9 +314,7 @@ While compiling a Futhark program to an executable is useful for testing, it is not suitable for production use.  Instead, a Futhark program should be compiled into a reusable library in some target-language, enabling integration into a larger program.  Five of the-Futhark compilers support this: ``futhark c``, ``futhark opencl``,-``futhark cuda``, ``futhark py``, and ``futhark pyopencl``.+language, enabling integration into a larger program.  General Concerns ^^^^^^^^^^^^^^^^@@ -321,7 +328,8 @@ *out*-parameters for writing the result, for target languages that do not support multiple return values from functions. -The entry point should have a name that is also a valid C identifier.+The entry point should have a name that is also a valid identifier in+the target language (usually C).  Not all Futhark types can be mapped cleanly to the target language. Arrays of tuples, for example, are a common issue.  In such cases,

@@ -1,6 +1,6 @@ cabal-version: 2.4 name:           futhark-version:        0.25.10+version:        0.25.11 synopsis:       An optimising compiler for a functional, array-oriented language.  description:    Futhark is a small programming language designed to be compiled to@@ -215,7 +215,7 @@       Futhark.CodeGen.ImpGen.CUDA       Futhark.CodeGen.ImpGen.GPU       Futhark.CodeGen.ImpGen.GPU.Base-      Futhark.CodeGen.ImpGen.GPU.Group+      Futhark.CodeGen.ImpGen.GPU.Block       Futhark.CodeGen.ImpGen.GPU.SegHist       Futhark.CodeGen.ImpGen.GPU.SegMap       Futhark.CodeGen.ImpGen.GPU.SegRed

@@ -137,6 +137,18 @@ def copy 't (a: t): *t =   ([a])[0] +-- | Copy a value. The result will not alias anything. Additionally,+-- there is a guarantee that the result will be laid out in row-major+-- order in memory. This can be used for locality optimisations in+-- cases where the compiler does not otherwise do the right thing.+--+-- **Work:** O(n).+--+-- **Span:** O(1).+#[inline]+def manifest 't (a: t): *t =+  intrinsics.manifest a+ -- | Combines the outer two dimensions of an array. -- -- **Complexity:** O(1).

@@ -1,29 +1,29 @@ // Start of atomics.h  SCALAR_FUN_ATTR int32_t atomic_xchg_i32_global(volatile __global int32_t *p, int32_t x);-SCALAR_FUN_ATTR int32_t atomic_xchg_i32_local(volatile __local int32_t *p, int32_t x);+SCALAR_FUN_ATTR int32_t atomic_xchg_i32_shared(volatile __local int32_t *p, int32_t x); SCALAR_FUN_ATTR int32_t atomic_cmpxchg_i32_global(volatile __global int32_t *p,                                                          int32_t cmp, int32_t val);-SCALAR_FUN_ATTR int32_t atomic_cmpxchg_i32_local(volatile __local int32_t *p,+SCALAR_FUN_ATTR int32_t atomic_cmpxchg_i32_shared(volatile __local int32_t *p,                                                         int32_t cmp, int32_t val); SCALAR_FUN_ATTR int32_t atomic_add_i32_global(volatile __global int32_t *p, int32_t x);-SCALAR_FUN_ATTR int32_t atomic_add_i32_local(volatile __local int32_t *p, int32_t x);+SCALAR_FUN_ATTR int32_t atomic_add_i32_shared(volatile __local int32_t *p, int32_t x); SCALAR_FUN_ATTR float atomic_fadd_f32_global(volatile __global float *p, float x);-SCALAR_FUN_ATTR float atomic_fadd_f32_local(volatile __local float *p, float x);+SCALAR_FUN_ATTR float atomic_fadd_f32_shared(volatile __local float *p, float x); SCALAR_FUN_ATTR int32_t atomic_smax_i32_global(volatile __global int32_t *p, int32_t x);-SCALAR_FUN_ATTR int32_t atomic_smax_i32_local(volatile __local int32_t *p, int32_t x);+SCALAR_FUN_ATTR int32_t atomic_smax_i32_shared(volatile __local int32_t *p, int32_t x); SCALAR_FUN_ATTR int32_t atomic_smin_i32_global(volatile __global int32_t *p, int32_t x);-SCALAR_FUN_ATTR int32_t atomic_smin_i32_local(volatile __local int32_t *p, int32_t x);+SCALAR_FUN_ATTR int32_t atomic_smin_i32_shared(volatile __local int32_t *p, int32_t x); SCALAR_FUN_ATTR uint32_t atomic_umax_i32_global(volatile __global uint32_t *p, uint32_t x);-SCALAR_FUN_ATTR uint32_t atomic_umax_i32_local(volatile __local uint32_t *p, uint32_t x);+SCALAR_FUN_ATTR uint32_t atomic_umax_i32_shared(volatile __local uint32_t *p, uint32_t x); SCALAR_FUN_ATTR uint32_t atomic_umin_i32_global(volatile __global uint32_t *p, uint32_t x);-SCALAR_FUN_ATTR uint32_t atomic_umin_i32_local(volatile __local uint32_t *p, uint32_t x);+SCALAR_FUN_ATTR uint32_t atomic_umin_i32_shared(volatile __local uint32_t *p, uint32_t x); SCALAR_FUN_ATTR int32_t atomic_and_i32_global(volatile __global int32_t *p, int32_t x);-SCALAR_FUN_ATTR int32_t atomic_and_i32_local(volatile __local int32_t *p, int32_t x);+SCALAR_FUN_ATTR int32_t atomic_and_i32_shared(volatile __local int32_t *p, int32_t x); SCALAR_FUN_ATTR int32_t atomic_or_i32_global(volatile __global int32_t *p, int32_t x);-SCALAR_FUN_ATTR int32_t atomic_or_i32_local(volatile __local int32_t *p, int32_t x);+SCALAR_FUN_ATTR int32_t atomic_or_i32_shared(volatile __local int32_t *p, int32_t x); SCALAR_FUN_ATTR int32_t atomic_xor_i32_global(volatile __global int32_t *p, int32_t x);-SCALAR_FUN_ATTR int32_t atomic_xor_i32_local(volatile __local int32_t *p, int32_t x);+SCALAR_FUN_ATTR int32_t atomic_xor_i32_shared(volatile __local int32_t *p, int32_t x);  SCALAR_FUN_ATTR int32_t atomic_xchg_i32_global(volatile __global int32_t *p, int32_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)@@ -33,7 +33,7 @@ #endif } -SCALAR_FUN_ATTR int32_t atomic_xchg_i32_local(volatile __local int32_t *p, int32_t x) {+SCALAR_FUN_ATTR int32_t atomic_xchg_i32_shared(volatile __local int32_t *p, int32_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicExch((int32_t*)p, x); #else@@ -50,7 +50,7 @@ #endif } -SCALAR_FUN_ATTR int32_t atomic_cmpxchg_i32_local(volatile __local int32_t *p,+SCALAR_FUN_ATTR int32_t atomic_cmpxchg_i32_shared(volatile __local int32_t *p,                                                         int32_t cmp, int32_t val) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicCAS((int32_t*)p, cmp, val);@@ -67,7 +67,7 @@ #endif } -SCALAR_FUN_ATTR int32_t atomic_add_i32_local(volatile __local int32_t *p, int32_t x) {+SCALAR_FUN_ATTR int32_t atomic_add_i32_shared(volatile __local int32_t *p, int32_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicAdd((int32_t*)p, x); #else@@ -91,7 +91,7 @@ #endif } -SCALAR_FUN_ATTR float atomic_fadd_f32_local(volatile __local float *p, float x) {+SCALAR_FUN_ATTR float atomic_fadd_f32_shared(volatile __local float *p, float x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicAdd((float*)p, x); #else@@ -101,7 +101,7 @@   do {     assumed.f = old.f;     old.f = old.f + x;-    old.i = atomic_cmpxchg_i32_local((volatile __local int32_t*)p, assumed.i, old.i);+    old.i = atomic_cmpxchg_i32_shared((volatile __local int32_t*)p, assumed.i, old.i);   } while (assumed.i != old.i);   return old.f; #endif@@ -115,7 +115,7 @@ #endif } -SCALAR_FUN_ATTR int32_t atomic_smax_i32_local(volatile __local int32_t *p, int32_t x) {+SCALAR_FUN_ATTR int32_t atomic_smax_i32_shared(volatile __local int32_t *p, int32_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicMax((int32_t*)p, x); #else@@ -131,7 +131,7 @@ #endif } -SCALAR_FUN_ATTR int32_t atomic_smin_i32_local(volatile __local int32_t *p, int32_t x) {+SCALAR_FUN_ATTR int32_t atomic_smin_i32_shared(volatile __local int32_t *p, int32_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicMin((int32_t*)p, x); #else@@ -147,7 +147,7 @@ #endif } -SCALAR_FUN_ATTR uint32_t atomic_umax_i32_local(volatile __local uint32_t *p, uint32_t x) {+SCALAR_FUN_ATTR uint32_t atomic_umax_i32_shared(volatile __local uint32_t *p, uint32_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicMax((uint32_t*)p, x); #else@@ -163,7 +163,7 @@ #endif } -SCALAR_FUN_ATTR uint32_t atomic_umin_i32_local(volatile __local uint32_t *p, uint32_t x) {+SCALAR_FUN_ATTR uint32_t atomic_umin_i32_shared(volatile __local uint32_t *p, uint32_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicMin((uint32_t*)p, x); #else@@ -179,7 +179,7 @@ #endif } -SCALAR_FUN_ATTR int32_t atomic_and_i32_local(volatile __local int32_t *p, int32_t x) {+SCALAR_FUN_ATTR int32_t atomic_and_i32_shared(volatile __local int32_t *p, int32_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicAnd((int32_t*)p, x); #else@@ -195,7 +195,7 @@ #endif } -SCALAR_FUN_ATTR int32_t atomic_or_i32_local(volatile __local int32_t *p, int32_t x) {+SCALAR_FUN_ATTR int32_t atomic_or_i32_shared(volatile __local int32_t *p, int32_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicOr((int32_t*)p, x); #else@@ -211,7 +211,7 @@ #endif } -SCALAR_FUN_ATTR int32_t atomic_xor_i32_local(volatile __local int32_t *p, int32_t x) {+SCALAR_FUN_ATTR int32_t atomic_xor_i32_shared(volatile __local int32_t *p, int32_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicXor((int32_t*)p, x); #else@@ -224,31 +224,31 @@ #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP) || defined(cl_khr_int64_base_atomics) && defined(cl_khr_int64_extended_atomics)  SCALAR_FUN_ATTR int64_t atomic_xchg_i64_global(volatile __global int64_t *p, int64_t x);-SCALAR_FUN_ATTR int64_t atomic_xchg_i64_local(volatile __local int64_t *p, int64_t x);+SCALAR_FUN_ATTR int64_t atomic_xchg_i64_shared(volatile __local int64_t *p, int64_t x); SCALAR_FUN_ATTR int64_t atomic_cmpxchg_i64_global(volatile __global int64_t *p,                                                          int64_t cmp, int64_t val);-SCALAR_FUN_ATTR int64_t atomic_cmpxchg_i64_local(volatile __local int64_t *p,+SCALAR_FUN_ATTR int64_t atomic_cmpxchg_i64_shared(volatile __local int64_t *p,                                                         int64_t cmp, int64_t val); SCALAR_FUN_ATTR int64_t atomic_add_i64_global(volatile __global int64_t *p, int64_t x);-SCALAR_FUN_ATTR int64_t atomic_add_i64_local(volatile __local int64_t *p, int64_t x);+SCALAR_FUN_ATTR int64_t atomic_add_i64_shared(volatile __local int64_t *p, int64_t x); SCALAR_FUN_ATTR int64_t atomic_smax_i64_global(volatile __global int64_t *p, int64_t x);-SCALAR_FUN_ATTR int64_t atomic_smax_i64_local(volatile __local int64_t *p, int64_t x);+SCALAR_FUN_ATTR int64_t atomic_smax_i64_shared(volatile __local int64_t *p, int64_t x); SCALAR_FUN_ATTR int64_t atomic_smin_i64_global(volatile __global int64_t *p, int64_t x);-SCALAR_FUN_ATTR int64_t atomic_smin_i64_local(volatile __local int64_t *p, int64_t x);+SCALAR_FUN_ATTR int64_t atomic_smin_i64_shared(volatile __local int64_t *p, int64_t x); SCALAR_FUN_ATTR uint64_t atomic_umax_i64_global(volatile __global uint64_t *p, uint64_t x);-SCALAR_FUN_ATTR uint64_t atomic_umax_i64_local(volatile __local uint64_t *p, uint64_t x);+SCALAR_FUN_ATTR uint64_t atomic_umax_i64_shared(volatile __local uint64_t *p, uint64_t x); SCALAR_FUN_ATTR uint64_t atomic_umin_i64_global(volatile __global uint64_t *p, uint64_t x);-SCALAR_FUN_ATTR uint64_t atomic_umin_i64_local(volatile __local uint64_t *p, uint64_t x);+SCALAR_FUN_ATTR uint64_t atomic_umin_i64_shared(volatile __local uint64_t *p, uint64_t x); SCALAR_FUN_ATTR int64_t atomic_and_i64_global(volatile __global int64_t *p, int64_t x);-SCALAR_FUN_ATTR int64_t atomic_and_i64_local(volatile __local int64_t *p, int64_t x);+SCALAR_FUN_ATTR int64_t atomic_and_i64_shared(volatile __local int64_t *p, int64_t x); SCALAR_FUN_ATTR int64_t atomic_or_i64_global(volatile __global int64_t *p, int64_t x);-SCALAR_FUN_ATTR int64_t atomic_or_i64_local(volatile __local int64_t *p, int64_t x);+SCALAR_FUN_ATTR int64_t atomic_or_i64_shared(volatile __local int64_t *p, int64_t x); SCALAR_FUN_ATTR int64_t atomic_xor_i64_global(volatile __global int64_t *p, int64_t x);-SCALAR_FUN_ATTR int64_t atomic_xor_i64_local(volatile __local int64_t *p, int64_t x);+SCALAR_FUN_ATTR int64_t atomic_xor_i64_shared(volatile __local int64_t *p, int64_t x);  #ifdef FUTHARK_F64_ENABLED SCALAR_FUN_ATTR double atomic_fadd_f64_global(volatile __global double *p, double x);-SCALAR_FUN_ATTR double atomic_fadd_f64_local(volatile __local double *p, double x);+SCALAR_FUN_ATTR double atomic_fadd_f64_shared(volatile __local double *p, double x); #endif  SCALAR_FUN_ATTR int64_t atomic_xchg_i64_global(volatile __global int64_t *p, int64_t x) {@@ -259,7 +259,7 @@ #endif } -SCALAR_FUN_ATTR int64_t atomic_xchg_i64_local(volatile __local int64_t *p, int64_t x) {+SCALAR_FUN_ATTR int64_t atomic_xchg_i64_shared(volatile __local int64_t *p, int64_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicExch((uint64_t*)p, x); #else@@ -276,7 +276,7 @@ #endif } -SCALAR_FUN_ATTR int64_t atomic_cmpxchg_i64_local(volatile __local int64_t *p,+SCALAR_FUN_ATTR int64_t atomic_cmpxchg_i64_shared(volatile __local int64_t *p,                                                         int64_t cmp, int64_t val) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicCAS((uint64_t*)p, cmp, val);@@ -293,7 +293,7 @@ #endif } -SCALAR_FUN_ATTR int64_t atomic_add_i64_local(volatile __local int64_t *p, int64_t x) {+SCALAR_FUN_ATTR int64_t atomic_add_i64_shared(volatile __local int64_t *p, int64_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicAdd((uint64_t*)p, x); #else@@ -319,7 +319,7 @@ #endif } -SCALAR_FUN_ATTR double atomic_fadd_f64_local(volatile __local double *p, double x) {+SCALAR_FUN_ATTR double atomic_fadd_f64_shared(volatile __local double *p, double x) { #if defined(FUTHARK_CUDA) && __CUDA_ARCH__ >= 600 || defined(FUTHARK_HIP)   return atomicAdd((double*)p, x); #else@@ -329,7 +329,7 @@   do {     assumed.f = old.f;     old.f = old.f + x;-    old.i = atomic_cmpxchg_i64_local((volatile __local int64_t*)p, assumed.i, old.i);+    old.i = atomic_cmpxchg_i64_shared((volatile __local int64_t*)p, assumed.i, old.i);   } while (assumed.i != old.i);   return old.f; #endif@@ -354,7 +354,7 @@ #endif } -SCALAR_FUN_ATTR int64_t atomic_smax_i64_local(volatile __local int64_t *p, int64_t x) {+SCALAR_FUN_ATTR int64_t atomic_smax_i64_shared(volatile __local int64_t *p, int64_t x) { #if defined(FUTHARK_CUDA)   return atomicMax((int64_t*)p, x); #elif defined(FUTHARK_HIP)@@ -363,7 +363,7 @@   do {     assumed = old;     old = smax64(old, x);-    old = atomic_cmpxchg_i64_local((volatile __local int64_t*)p, assumed, old);+    old = atomic_cmpxchg_i64_shared((volatile __local int64_t*)p, assumed, old);   } while (assumed != old);   return old; #else@@ -388,7 +388,7 @@ #endif } -SCALAR_FUN_ATTR int64_t atomic_smin_i64_local(volatile __local int64_t *p, int64_t x) {+SCALAR_FUN_ATTR int64_t atomic_smin_i64_shared(volatile __local int64_t *p, int64_t x) { #if defined(FUTHARK_CUDA)   return atomicMin((int64_t*)p, x); #elif defined(FUTHARK_HIP)@@ -397,7 +397,7 @@   do {     assumed = old;     old = smin64(old, x);-    old = atomic_cmpxchg_i64_local((volatile __local int64_t*)p, assumed, old);+    old = atomic_cmpxchg_i64_shared((volatile __local int64_t*)p, assumed, old);   } while (assumed != old);   return old; #else@@ -413,7 +413,7 @@ #endif } -SCALAR_FUN_ATTR uint64_t atomic_umax_i64_local(volatile __local uint64_t *p, uint64_t x) {+SCALAR_FUN_ATTR uint64_t atomic_umax_i64_shared(volatile __local uint64_t *p, uint64_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicMax((uint64_t*)p, x); #else@@ -429,7 +429,7 @@ #endif } -SCALAR_FUN_ATTR uint64_t atomic_umin_i64_local(volatile __local uint64_t *p, uint64_t x) {+SCALAR_FUN_ATTR uint64_t atomic_umin_i64_shared(volatile __local uint64_t *p, uint64_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicMin((uint64_t*)p, x); #else@@ -445,7 +445,7 @@ #endif } -SCALAR_FUN_ATTR int64_t atomic_and_i64_local(volatile __local int64_t *p, int64_t x) {+SCALAR_FUN_ATTR int64_t atomic_and_i64_shared(volatile __local int64_t *p, int64_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicAnd((uint64_t*)p, x); #else@@ -461,7 +461,7 @@ #endif } -SCALAR_FUN_ATTR int64_t atomic_or_i64_local(volatile __local int64_t *p, int64_t x) {+SCALAR_FUN_ATTR int64_t atomic_or_i64_shared(volatile __local int64_t *p, int64_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicOr((uint64_t*)p, x); #else@@ -477,7 +477,7 @@ #endif } -SCALAR_FUN_ATTR int64_t atomic_xor_i64_local(volatile __local int64_t *p, int64_t x) {+SCALAR_FUN_ATTR int64_t atomic_xor_i64_shared(volatile __local int64_t *p, int64_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicXor((uint64_t*)p, x); #else

@@ -186,16 +186,24 @@   cfg->load_ptx_from = strdup(path); } -void futhark_context_config_set_default_group_size(struct futhark_context_config *cfg, int size) {+void futhark_context_config_set_default_thread_block_size(struct futhark_context_config *cfg, int size) {   cfg->default_block_size = size;   cfg->default_block_size_changed = 1; } -void futhark_context_config_set_default_num_groups(struct futhark_context_config *cfg, int num) {+void futhark_context_config_set_default_grid_size(struct futhark_context_config *cfg, int num) {   cfg->default_grid_size = num;   cfg->default_grid_size_changed = 1; } +void futhark_context_config_set_default_group_size(struct futhark_context_config *cfg, int size) {+  futhark_context_config_set_default_thread_block_size(cfg, size);+}++void futhark_context_config_set_default_num_groups(struct futhark_context_config *cfg, int num) {+  futhark_context_config_set_default_grid_size(cfg, num);+}+ void futhark_context_config_set_default_tile_size(struct futhark_context_config *cfg, int size) {   cfg->default_tile_size = size;   cfg->default_tile_size_changed = 1;@@ -218,11 +226,13 @@       return 0;     }   }-  if (strcmp(param_name, "default_group_size") == 0) {+  if (strcmp(param_name, "default_thread_block_size") == 0+      || strcmp(param_name, "default_group_size") == 0) {     cfg->default_block_size = new_value;     return 0;   }-  if (strcmp(param_name, "default_num_groups") == 0) {+  if (strcmp(param_name, "default_num_groups") == 0 ||+      strcmp(param_name, "default_grid_size") == 0) {     cfg->default_grid_size = new_value;     return 0;   }@@ -283,11 +293,11 @@    struct free_list gpu_free_list; -  size_t max_group_size;+  size_t max_thread_block_size;   size_t max_grid_size;   size_t max_tile_size;   size_t max_threshold;-  size_t max_local_memory;+  size_t max_shared_memory;   size_t max_bespoke;   size_t max_registers;   size_t max_cache;@@ -461,11 +471,11 @@     opts[i++] = strdup("--disable-warnings");   }   opts[i++] = msgprintf("-D%s=%d",-                        "max_group_size",-                        (int)ctx->max_group_size);+                        "max_thread_block_size",+                        (int)ctx->max_thread_block_size);   opts[i++] = msgprintf("-D%s=%d",-                        "max_local_memory",-                        (int)ctx->max_local_memory);+                        "max_shared_memory",+                        (int)ctx->max_shared_memory);   opts[i++] = msgprintf("-D%s=%d",                         "max_registers",                         (int)ctx->max_registers);@@ -479,7 +489,7 @@                           cfg->tuning_params[j]);   }   opts[i++] = msgprintf("-DLOCKSTEP_WIDTH=%zu", ctx->lockstep_width);-  opts[i++] = msgprintf("-DMAX_THREADS_PER_BLOCK=%zu", ctx->max_group_size);+  opts[i++] = msgprintf("-DMAX_THREADS_PER_BLOCK=%zu", ctx->max_thread_block_size);    // Time for the best lines of the code in the entire compiler.   if (getenv("CUDA_HOME") != NULL) {@@ -570,13 +580,13 @@ static void cuda_size_setup(struct futhark_context *ctx) {   struct futhark_context_config *cfg = ctx->cfg;-  if (cfg->default_block_size > ctx->max_group_size) {+  if (cfg->default_block_size > ctx->max_thread_block_size) {     if (cfg->default_block_size_changed) {       fprintf(stderr,               "Note: Device limits default block size to %zu (down from %zu).\n",-              ctx->max_group_size, cfg->default_block_size);+              ctx->max_thread_block_size, cfg->default_block_size);     }-    cfg->default_block_size = ctx->max_group_size;+    cfg->default_block_size = ctx->max_thread_block_size;   }   if (cfg->default_grid_size > ctx->max_grid_size) {     if (cfg->default_grid_size_changed) {@@ -608,10 +618,10 @@     const char* size_name = cfg->tuning_param_names[i];     int64_t max_value = 0, default_value = 0; -    if (strstr(size_class, "group_size") == size_class) {-      max_value = ctx->max_group_size;+    if (strstr(size_class, "thread_block_size") == size_class) {+      max_value = ctx->max_thread_block_size;       default_value = cfg->default_block_size;-    } else if (strstr(size_class, "num_groups") == size_class) {+    } else if (strstr(size_class, "grid_size") == size_class) {       max_value = ctx->max_grid_size;       default_value = cfg->default_grid_size;       // XXX: as a quick and dirty hack, use twice as many threads for@@ -821,12 +831,13 @@   // MAX_SHARED_MEMORY_PER_BLOCK gives bogus numbers (48KiB); probably   // for backwards compatibility.  Add _OPTIN and you seem to get the   // right number.-  ctx->max_local_memory =-    device_query(ctx->dev, MAX_SHARED_MEMORY_PER_BLOCK_OPTIN) --    device_query(ctx->dev, RESERVED_SHARED_MEMORY_PER_BLOCK);-  ctx->max_group_size = device_query(ctx->dev, MAX_THREADS_PER_BLOCK);+  ctx->max_shared_memory = device_query(ctx->dev, MAX_SHARED_MEMORY_PER_BLOCK_OPTIN);+#if CUDART_VERSION >= 12000+  ctx->max_shared_memory -= device_query(ctx->dev, RESERVED_SHARED_MEMORY_PER_BLOCK);+#endif+  ctx->max_thread_block_size = device_query(ctx->dev, MAX_THREADS_PER_BLOCK);   ctx->max_grid_size = device_query(ctx->dev, MAX_GRID_DIM_X);-  ctx->max_tile_size = sqrt(ctx->max_group_size);+  ctx->max_tile_size = sqrt(ctx->max_thread_block_size);   ctx->max_threshold = 0;   ctx->max_bespoke = 0;   ctx->max_registers = device_query(ctx->dev, MAX_REGISTERS_PER_BLOCK);@@ -883,7 +894,7 @@   // Unless the below is set, the kernel is limited to 48KiB of memory.   CUDA_SUCCEED_FATAL(cuFuncSetAttribute(*kernel,                                         cudaFuncAttributeMaxDynamicSharedMemorySize,-                                        ctx->max_local_memory));+                                        ctx->max_shared_memory)); }  static void gpu_free_kernel(struct futhark_context *ctx,@@ -1015,7 +1026,7 @@                              gpu_kernel kernel, const char *name,                              const int32_t grid[3],                              const int32_t block[3],-                             unsigned int local_mem_bytes,+                             unsigned int shared_mem_bytes,                              int num_args,                              void* args[num_args],                              size_t args_sizes[num_args]) {@@ -1039,7 +1050,7 @@                         name,                         grid[0], grid[1], grid[2],                         block[0], block[1], block[2],-                        local_mem_bytes),+                        shared_mem_bytes),               event,               (event_report_fn)cuda_event_report);   }@@ -1048,7 +1059,7 @@     (cuLaunchKernel(kernel,                     grid[0], grid[1], grid[2],                     block[0], block[1], block[2],-                    local_mem_bytes, ctx->stream,+                    shared_mem_bytes, ctx->stream,                     args, NULL));    if (event != NULL) {

@@ -166,16 +166,25 @@   cfg->program = strdup(s); } -void futhark_context_config_set_default_group_size(struct futhark_context_config *cfg, int size) {+void futhark_context_config_set_default_thread_block_size(struct futhark_context_config *cfg, int size) {   cfg->default_block_size = size;   cfg->default_block_size_changed = 1; } -void futhark_context_config_set_default_num_groups(struct futhark_context_config *cfg, int num) {+void futhark_context_config_set_default_grid_size(struct futhark_context_config *cfg, int num) {   cfg->default_grid_size = num;   cfg->default_grid_size_changed = 1; } +void futhark_context_config_set_default_group_size(struct futhark_context_config *cfg, int num) {+  futhark_context_config_set_default_thread_block_size(cfg, num);+}+++void futhark_context_config_set_default_num_groups(struct futhark_context_config *cfg, int num) {+  futhark_context_config_set_default_grid_size(cfg, num);+}+ void futhark_context_config_set_default_tile_size(struct futhark_context_config *cfg, int size) {   cfg->default_tile_size = size;   cfg->default_tile_size_changed = 1;@@ -198,11 +207,13 @@       return 0;     }   }-  if (strcmp(param_name, "default_group_size") == 0) {+  if (strcmp(param_name, "default_thread_block_size") == 0 ||+      strcmp(param_name, "default_group_size") == 0) {     cfg->default_block_size = new_value;     return 0;   }-  if (strcmp(param_name, "default_num_groups") == 0) {+  if (strcmp(param_name, "default_grid_size") == 0 ||+      strcmp(param_name, "default_num_groups") == 0) {     cfg->default_grid_size = new_value;     return 0;   }@@ -257,11 +268,11 @@    struct free_list gpu_free_list; -  size_t max_group_size;+  size_t max_thread_block_size;   size_t max_grid_size;   size_t max_tile_size;   size_t max_threshold;-  size_t max_local_memory;+  size_t max_shared_memory;   size_t max_bespoke;   size_t max_registers;   size_t max_cache;@@ -342,13 +353,13 @@  static void hip_size_setup(struct futhark_context *ctx) {   struct futhark_context_config *cfg = ctx->cfg;-  if (cfg->default_block_size > ctx->max_group_size) {+  if (cfg->default_block_size > ctx->max_thread_block_size) {     if (cfg->default_block_size_changed) {       fprintf(stderr,               "Note: Device limits default block size to %zu (down from %zu).\n",-              ctx->max_group_size, cfg->default_block_size);+              ctx->max_thread_block_size, cfg->default_block_size);     }-    cfg->default_block_size = ctx->max_group_size;+    cfg->default_block_size = ctx->max_thread_block_size;   }   if (cfg->default_grid_size > ctx->max_grid_size) {     if (cfg->default_grid_size_changed) {@@ -380,10 +391,10 @@     const char* size_name = cfg->tuning_param_names[i];     int64_t max_value = 0, default_value = 0; -    if (strstr(size_class, "group_size") == size_class) {-      max_value = ctx->max_group_size;+    if (strstr(size_class, "thread_block_size") == size_class) {+      max_value = ctx->max_thread_block_size;       default_value = cfg->default_block_size;-    } else if (strstr(size_class, "num_groups") == size_class) {+    } else if (strstr(size_class, "grid_size") == size_class) {       max_value = ctx->max_grid_size;       default_value = cfg->default_grid_size;       // XXX: as a quick and dirty hack, use twice as many threads for@@ -477,11 +488,11 @@     opts[i++] = strdup("-lineinfo");   }   opts[i++] = msgprintf("-D%s=%d",-                        "max_group_size",-                        (int)ctx->max_group_size);+                        "max_thread_block_size",+                        (int)ctx->max_thread_block_size);   opts[i++] = msgprintf("-D%s=%d",-                        "max_local_memory",-                        (int)ctx->max_local_memory);+                        "max_shared_memory",+                        (int)ctx->max_shared_memory);   opts[i++] = msgprintf("-D%s=%d",                         "max_registers",                         (int)ctx->max_registers);@@ -495,7 +506,7 @@                           cfg->tuning_params[j]);   }   opts[i++] = msgprintf("-DLOCKSTEP_WIDTH=%zu", ctx->lockstep_width);-  opts[i++] = msgprintf("-DMAX_THREADS_PER_BLOCK=%zu", ctx->max_group_size);+  opts[i++] = msgprintf("-DMAX_THREADS_PER_BLOCK=%zu", ctx->max_thread_block_size);    for (int j = 0; extra_opts[j] != NULL; j++) {     opts[i++] = strdup(extra_opts[j]);@@ -673,10 +684,10 @@    free_list_init(&ctx->gpu_free_list); -  ctx->max_local_memory = device_query(ctx->dev, hipDeviceAttributeMaxSharedMemoryPerBlock);-  ctx->max_group_size = device_query(ctx->dev, hipDeviceAttributeMaxThreadsPerBlock);+  ctx->max_shared_memory = device_query(ctx->dev, hipDeviceAttributeMaxSharedMemoryPerBlock);+  ctx->max_thread_block_size = device_query(ctx->dev, hipDeviceAttributeMaxThreadsPerBlock);   ctx->max_grid_size = device_query(ctx->dev, hipDeviceAttributeMaxGridDimX);-  ctx->max_tile_size = sqrt(ctx->max_group_size);+  ctx->max_tile_size = sqrt(ctx->max_thread_block_size);   ctx->max_threshold = 0;   ctx->max_bespoke = 0;   ctx->max_registers = device_query(ctx->dev, hipDeviceAttributeMaxRegistersPerBlock);@@ -872,7 +883,7 @@                              gpu_kernel kernel, const char *name,                              const int32_t grid[3],                              const int32_t block[3],-                             unsigned int local_mem_bytes,+                             unsigned int shared_mem_bytes,                              int num_args,                              void* args[num_args],                              size_t args_sizes[num_args]) {@@ -895,7 +906,7 @@                         name,                         grid[0], grid[1], grid[2],                         block[0], block[1], block[2],-                        local_mem_bytes),+                        shared_mem_bytes),               event,               (event_report_fn)hip_event_report);   }@@ -904,7 +915,7 @@     (hipModuleLaunchKernel(kernel,                            grid[0], grid[1], grid[2],                            block[0], block[1], block[2],-                           local_mem_bytes, ctx->stream,+                           shared_mem_bytes, ctx->stream,                            args, NULL));    if (event != NULL) {

@@ -432,15 +432,23 @@   cfg->load_binary_from = strdup(path); } -void futhark_context_config_set_default_group_size(struct futhark_context_config *cfg, int size) {+void futhark_context_config_set_default_thread_block_size(struct futhark_context_config *cfg, int size) {   cfg->default_group_size = size;   cfg->default_group_size_changed = 1; } -void futhark_context_config_set_default_num_groups(struct futhark_context_config *cfg, int num) {+void futhark_context_config_set_default_group_size(struct futhark_context_config *cfg, int size) {+  futhark_context_config_set_default_thread_block_size(cfg, size);+}++void futhark_context_config_set_default_grid_size(struct futhark_context_config *cfg, int num) {   cfg->default_num_groups = num; } +void futhark_context_config_set_default_num_groups(struct futhark_context_config *cfg, int num) {+  futhark_context_config_set_default_grid_size(cfg, num);+}+ void futhark_context_config_set_default_tile_size(struct futhark_context_config *cfg, int size) {   cfg->default_tile_size = size;   cfg->default_tile_size_changed = 1;@@ -463,11 +471,13 @@       return 0;     }   }-  if (strcmp(param_name, "default_group_size") == 0) {+  if (strcmp(param_name, "default_thread_block_size") == 0 ||+      strcmp(param_name, "default_group_size") == 0) {     cfg->default_group_size = new_value;     return 0;   }-  if (strcmp(param_name, "default_num_groups") == 0) {+  if (strcmp(param_name, "default_grid_size") == 0 ||+      strcmp(param_name, "default_num_groups") == 0) {     cfg->default_num_groups = new_value;     return 0;   }@@ -523,11 +533,11 @@    struct free_list gpu_free_list; -  size_t max_group_size;+  size_t max_thread_block_size;   size_t max_num_groups;   size_t max_tile_size;   size_t max_threshold;-  size_t max_local_memory;+  size_t max_shared_memory;   size_t max_registers;   size_t max_cache; @@ -601,13 +611,13 @@    w += snprintf(compile_opts+w, compile_opts_size-w,                 "-D%s=%d ",-                "max_group_size",-                (int)ctx->max_group_size);+                "max_thread_block_size",+                (int)ctx->max_thread_block_size);    w += snprintf(compile_opts+w, compile_opts_size-w,                 "-D%s=%d ",-                "max_local_memory",-                (int)ctx->max_local_memory);+                "max_shared_memory",+                (int)ctx->max_shared_memory);    w += snprintf(compile_opts+w, compile_opts_size-w,                 "-D%s=%d ",@@ -767,18 +777,18 @@     }   } -  size_t max_group_size;+  size_t max_thread_block_size;   OPENCL_SUCCEED_FATAL(clGetDeviceInfo(device_option.device, CL_DEVICE_MAX_WORK_GROUP_SIZE,-                                       sizeof(size_t), &max_group_size, NULL));+                                       sizeof(size_t), &max_thread_block_size, NULL)); -  size_t max_tile_size = sqrt(max_group_size);+  size_t max_tile_size = sqrt(max_thread_block_size); -  cl_ulong max_local_memory;+  cl_ulong max_shared_memory;   OPENCL_SUCCEED_FATAL(clGetDeviceInfo(device_option.device, CL_DEVICE_LOCAL_MEM_SIZE,-                                       sizeof(size_t), &max_local_memory, NULL));+                                       sizeof(size_t), &max_shared_memory, NULL));    // Futhark reserves 4 bytes for bookkeeping information.-  max_local_memory -= 4;+  max_shared_memory -= 4;    // The OpenCL implementation may reserve some local memory bytes for   // various purposes.  In principle, we should use@@ -789,20 +799,20 @@   // (but which might be arbitrarily wrong).  Fortunately, we rarely   // try to really push the local memory usage.   if (strstr(device_option.platform_name, "NVIDIA CUDA") != NULL) {-    max_local_memory -= 12;+    max_shared_memory -= 12;   } else if (strstr(device_option.platform_name, "AMD") != NULL) {-    max_local_memory -= 16;+    max_shared_memory -= 16;   }    // Make sure this function is defined.   post_opencl_setup(ctx, &device_option); -  if (max_group_size < ctx->cfg->default_group_size) {+  if (max_thread_block_size < ctx->cfg->default_group_size) {     if (ctx->cfg->default_group_size_changed) {       fprintf(stderr, "Note: Device limits default group size to %zu (down from %zu).\n",-              max_group_size, ctx->cfg->default_group_size);+              max_thread_block_size, ctx->cfg->default_group_size);     }-    ctx->cfg->default_group_size = max_group_size;+    ctx->cfg->default_group_size = max_thread_block_size;   }    if (max_tile_size < ctx->cfg->default_tile_size) {@@ -827,10 +837,10 @@    ctx->max_registers = 1<<16; // I cannot find a way to query for this. -  ctx->max_group_size = max_group_size;+  ctx->max_thread_block_size = max_thread_block_size;   ctx->max_tile_size = max_tile_size; // No limit.   ctx->max_threshold = ctx->max_num_groups = 0; // No limit.-  ctx->max_local_memory = max_local_memory;+  ctx->max_shared_memory = max_shared_memory;    // Now we go through all the sizes, clamp them to the valid range,   // or set them to the default.@@ -840,11 +850,11 @@     const char* size_name = ctx->cfg->tuning_param_names[i];     int64_t max_value = 0, default_value = 0; -    if (strstr(size_class, "group_size") == size_class) {-      max_value = max_group_size;+    if (strstr(size_class, "thread_block_size") == size_class) {+      max_value = max_thread_block_size;       default_value = ctx->cfg->default_group_size;-    } else if (strstr(size_class, "num_groups") == size_class) {-      max_value = max_group_size; // Futhark assumes this constraint.+    } else if (strstr(size_class, "grid_size") == size_class) {+      max_value = max_thread_block_size; // Futhark assumes this constraint.       default_value = ctx->cfg->default_num_groups;       // XXX: as a quick and dirty hack, use twice as many threads for       // histograms by default.  We really should just be smarter@@ -853,7 +863,7 @@         default_value *= 2;       }     } else if (strstr(size_class, "tile_size") == size_class) {-      max_value = sqrt(max_group_size);+      max_value = sqrt(max_thread_block_size);       default_value = ctx->cfg->default_tile_size;     } else if (strstr(size_class, "reg_tile_size") == size_class) {       max_value = 0; // No limit.@@ -879,8 +889,8 @@    if (ctx->cfg->logging) {     fprintf(stderr, "Lockstep width: %d\n", (int)ctx->lockstep_width);-    fprintf(stderr, "Default group size: %d\n", (int)ctx->cfg->default_group_size);-    fprintf(stderr, "Default number of groups: %d\n", (int)ctx->cfg->default_num_groups);+    fprintf(stderr, "Default thread block size: %d\n", (int)ctx->cfg->default_group_size);+    fprintf(stderr, "Default number of thread blocks: %d\n", (int)ctx->cfg->default_num_groups);   }    char *compile_opts = mk_compile_opts(ctx, extra_build_opts, device_option);@@ -1265,7 +1275,7 @@                              gpu_kernel kernel, const char *name,                              const int32_t grid[3],                              const int32_t block[3],-                             unsigned int local_mem_bytes,+                             unsigned int shared_mem_bytes,                              int num_args,                              void* args[num_args],                              size_t args_sizes[num_args]) {@@ -1282,7 +1292,7 @@                         name,                         grid[0], grid[1], grid[2],                         block[0], block[1], block[2],-                        local_mem_bytes),+                        shared_mem_bytes),               event,               (event_report_fn)opencl_event_report);   }@@ -1292,12 +1302,12 @@   }    // Some implementations do not work with 0-byte local memory.-  if (local_mem_bytes == 0) {-    local_mem_bytes = 4;+  if (shared_mem_bytes == 0) {+    shared_mem_bytes = 4;   }    OPENCL_SUCCEED_OR_RETURN-    (clSetKernelArg(kernel, 0, local_mem_bytes, NULL));+    (clSetKernelArg(kernel, 0, shared_mem_bytes, NULL));   for (int i = 0; i < num_args; i++) {     OPENCL_SUCCEED_OR_RETURN       (clSetKernelArg(kernel, i+1, args_sizes[i], args[i]));@@ -1325,7 +1335,7 @@     fprintf(ctx->log, "  runtime: %ldus\n", time_diff);   }   if (ctx->logging) {-    printf("\n");+    fprintf(ctx->log, "\n");   }    return FUTHARK_SUCCESS;

@@ -123,9 +123,9 @@   ctx->peak_mem_usage_default = 0;   ctx->cur_mem_usage_default = 0;   ctx->constants = malloc(sizeof(struct constants));-  ctx->detail_memory = cfg->debugging;   ctx->debugging = cfg->debugging;   ctx->logging = cfg->logging;+  ctx->detail_memory = cfg->logging;   ctx->profiling = cfg->profiling;   ctx->profiling_paused = 0;   ctx->error = NULL;

@@ -10,7 +10,7 @@                       gpu_kernel kernel, const char *name,                       const int32_t grid[3],                       const int32_t block[3],-                      unsigned int local_mem_bytes,+                      unsigned int shared_mem_bytes,                       int num_args,                       void* args[num_args],                       size_t args_sizes[num_args]);@@ -28,9 +28,9 @@                        gpu_kernel* kernel,                        const char* name); -// Max number of groups we allow along the second or third dimension-// for transpositions.-#define MAX_TR_GROUPS 65535+// Max number of thead blocks we allow along the second or third+// dimension for transpositions.+#define MAX_TR_THREAD_BLOCKS 65535  struct builtin_kernels {   // We have a lot of ways to transpose arrays.@@ -327,12 +327,12 @@     args_sizes[8] = sizeof(int64_t);   } -  // Cap the number of groups we launch and figure out how many+  // Cap the number of thead blocks we launch and figure out how many   // repeats we need alongside each dimension.-  int32_t repeat_1 = grid[1] / MAX_TR_GROUPS;-  int32_t repeat_2 = grid[2] / MAX_TR_GROUPS;-  grid[1] = repeat_1 > 0 ? MAX_TR_GROUPS : grid[1];-  grid[2] = repeat_2 > 0 ? MAX_TR_GROUPS : grid[2];+  int32_t repeat_1 = grid[1] / MAX_TR_THREAD_BLOCKS;+  int32_t repeat_2 = grid[2] / MAX_TR_THREAD_BLOCKS;+  grid[1] = repeat_1 > 0 ? MAX_TR_THREAD_BLOCKS : grid[1];+  grid[2] = repeat_2 > 0 ? MAX_TR_THREAD_BLOCKS : grid[2];   args[9] = &repeat_1;   args[10] = &repeat_2;   args_sizes[9] = sizeof(repeat_1);@@ -412,7 +412,7 @@       args[6+i*3+2] = &zero;     }   }-  const size_t w = 256; // XXX: hardcoded workgroup size.+  const size_t w = 256; // XXX: hardcoded thread block size.    return gpu_launch_kernel(ctx, kernel, "copy_lmad_dev_to_dev",                            (const int32_t[3]) {(n+w-1)/w,1,1},

@@ -140,21 +140,20 @@ }  /* returns the number of logical cores */-static int num_processors()-{+static int num_processors(void) { #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;+  /* 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;+  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(); #elif __EMSCRIPTEN__

@@ -282,14 +282,14 @@  // Print the command-done marker, indicating that we are ready for // more input.-void ok() {+void ok(void) {   printf("%%%%%% OK\n");   fflush(stdout); }  // Print the failure marker.  Output is now an error message until the // next ok().-void failure() {+void failure(void) {   printf("%%%%%% FAILURE\n"); } 

@@ -20,7 +20,7 @@ #define __write_only #define __read_only -static inline __device__ int get_group_id(int d) {+static inline __device__ int get_tblock_id(int d) {   switch (d) {   case 0: return blockIdx.x;   case 1: return blockIdx.y;@@ -29,7 +29,7 @@   } } -static inline __device__ int get_num_groups(int d) {+static inline __device__ int get_num_tblocks(int d) {   switch(d) {   case 0: return gridDim.x;   case 1: return gridDim.y;@@ -93,9 +93,9 @@  #define NAN (0.0/0.0) #define INFINITY (1.0/0.0)-extern volatile __shared__ unsigned char local_mem[];+extern volatile __shared__ unsigned char shared_mem[]; -#define LOCAL_MEM_PARAM+#define SHARED_MEM_PARAM #define FUTHARK_KERNEL extern "C" __global__ __launch_bounds__(MAX_THREADS_PER_BLOCK) #define FUTHARK_KERNEL_SIZED(a,b,c) extern "C" __global__ __launch_bounds__(a*b*c) 

@@ -1,7 +1,7 @@ // Start of copy.cl  #define GEN_COPY_KERNEL(NAME, ELEM_TYPE) \-FUTHARK_KERNEL void lmad_copy_##NAME(LOCAL_MEM_PARAM                    \+FUTHARK_KERNEL void lmad_copy_##NAME(SHARED_MEM_PARAM                   \                                __global ELEM_TYPE *dst_mem,             \                                int64_t dst_offset,                      \                                __global ELEM_TYPE *src_mem,             \

@@ -13,6 +13,8 @@ typedef uint uint32_t; typedef ulong uint64_t; +#define get_tblock_id(d) get_group_id(d)+#define get_num_tblocks(d) get_num_groups(d)  // Clang-based OpenCL implementations need this for 'static' to work. #ifdef cl_clang_storage_class_specifiers@@ -47,7 +49,7 @@ }  // Important for this to be int64_t so it has proper alignment for any type.-#define LOCAL_MEM_PARAM __local uint64_t* local_mem,+#define SHARED_MEM_PARAM __local uint64_t* shared_mem, #define FUTHARK_KERNEL __kernel #define FUTHARK_KERNEL_SIZED(a,b,c) __attribute__((reqd_work_group_size(a, b, c))) __kernel 

@@ -2,7 +2,7 @@  #define GEN_TRANSPOSE_KERNELS(NAME, ELEM_TYPE)                          \ FUTHARK_KERNEL_SIZED(TR_BLOCK_DIM*2, TR_TILE_DIM/TR_ELEMS_PER_THREAD, 1)\-void map_transpose_##NAME(LOCAL_MEM_PARAM                               \+void map_transpose_##NAME(SHARED_MEM_PARAM                              \                           __global ELEM_TYPE *dst_mem,                  \                           int64_t dst_offset,                           \                           __global ELEM_TYPE *src_mem,                  \@@ -15,20 +15,20 @@                           int32_t repeat_1,                             \                           int32_t repeat_2) {                           \   (void)mulx; (void)muly;                                               \-  __local ELEM_TYPE* block = (__local ELEM_TYPE*)local_mem;             \-  int group_id_0 = get_group_id(0);                                     \+  __local ELEM_TYPE* block = (__local ELEM_TYPE*)shared_mem;            \+  int tblock_id_0 = get_tblock_id(0);                                   \   int global_id_0 = get_global_id(0);                                   \-  int group_id_1 = get_group_id(1);                                     \+  int tblock_id_1 = get_tblock_id(1);                                   \   int global_id_1 = get_global_id(1);                                   \   for (int i1 = 0; i1 <= repeat_1; i1++) {                              \-    int group_id_2 = get_group_id(2);                                   \+    int tblock_id_2 = get_tblock_id(2);                                 \     int global_id_2 = get_global_id(2);                                 \     for (int i2 = 0; i2 <= repeat_2; i2++) {                            \-      int32_t our_array_offset = group_id_2 * x_elems * y_elems;        \+      int32_t our_array_offset = tblock_id_2 * x_elems * y_elems;       \       int32_t odata_offset = dst_offset + our_array_offset;             \       int32_t idata_offset = src_offset + our_array_offset;             \       int32_t x_index = global_id_0;                                    \-      int32_t y_index = group_id_1 * TR_TILE_DIM + get_local_id(1);     \+      int32_t y_index = tblock_id_1 * TR_TILE_DIM + get_local_id(1);    \       if (x_index < x_elems) {                                          \         for (int32_t j = 0; j < TR_ELEMS_PER_THREAD; j++) {             \           int32_t index_i = (y_index + j * (TR_TILE_DIM/TR_ELEMS_PER_THREAD)) * x_elems + x_index; \@@ -40,8 +40,8 @@         }                                                               \       }                                                                 \       barrier_local();                                                  \-      x_index = group_id_1 * TR_TILE_DIM + get_local_id(0);             \-      y_index = group_id_0 * TR_TILE_DIM + get_local_id(1);             \+      x_index = tblock_id_1 * TR_TILE_DIM + get_local_id(0);            \+      y_index = tblock_id_0 * TR_TILE_DIM + get_local_id(1);            \       if (x_index < y_elems) {                                          \         for (int32_t j = 0; j < TR_ELEMS_PER_THREAD; j++) {             \           int32_t index_out = (y_index + j * (TR_TILE_DIM/TR_ELEMS_PER_THREAD)) * y_elems + x_index; \@@ -52,16 +52,16 @@           }                                                             \         }                                                               \       }                                                                 \-      group_id_2 += get_num_groups(2);                                  \+      tblock_id_2 += get_num_tblocks(2);                                \       global_id_2 += get_global_size(2);                                \     }                                                                   \-    group_id_1 += get_num_groups(1);                                    \+    tblock_id_1 += get_num_tblocks(1);                                  \     global_id_1 += get_global_size(1);                                  \   }                                                                     \ }                                                                       \                                                                         \ FUTHARK_KERNEL_SIZED(TR_BLOCK_DIM, TR_BLOCK_DIM, 1)                     \-void map_transpose_##NAME##_low_height(LOCAL_MEM_PARAM                  \+void map_transpose_##NAME##_low_height(SHARED_MEM_PARAM                 \                                                 __global ELEM_TYPE *dst_mem, \                                                 int64_t dst_offset,     \                                                 __global ELEM_TYPE *src_mem, \@@ -73,32 +73,32 @@                                                 int32_t muly,           \                                                 int32_t repeat_1,       \                                                 int32_t repeat_2) {     \-  __local ELEM_TYPE* block = (__local ELEM_TYPE*)local_mem; \-  int group_id_0 = get_group_id(0);                                     \+  __local ELEM_TYPE* block = (__local ELEM_TYPE*)shared_mem;            \+  int tblock_id_0 = get_tblock_id(0);                                   \   int global_id_0 = get_global_id(0);                                   \-  int group_id_1 = get_group_id(1);                                     \+  int tblock_id_1 = get_tblock_id(1);                                   \   int global_id_1 = get_global_id(1);                                   \   for (int i1 = 0; i1 <= repeat_1; i1++) {                              \-    int group_id_2 = get_group_id(2);                                   \+    int tblock_id_2 = get_tblock_id(2);                                 \     int global_id_2 = get_global_id(2);                                 \     for (int i2 = 0; i2 <= repeat_2; i2++) {                            \-      int32_t our_array_offset = group_id_2 * x_elems * y_elems;        \+      int32_t our_array_offset = tblock_id_2 * x_elems * y_elems;       \       int32_t odata_offset = dst_offset + our_array_offset;             \       int32_t idata_offset = src_offset + our_array_offset;             \       int32_t x_index =                                                 \-        group_id_0 * TR_BLOCK_DIM * mulx +                              \+        tblock_id_0 * TR_BLOCK_DIM * mulx +                             \         get_local_id(0) +                                               \         get_local_id(1)%mulx * TR_BLOCK_DIM;                            \-      int32_t y_index = group_id_1 * TR_BLOCK_DIM + get_local_id(1)/mulx; \+      int32_t y_index = tblock_id_1 * TR_BLOCK_DIM + get_local_id(1)/mulx; \       int32_t index_in = y_index * x_elems + x_index;                   \       if (x_index < x_elems && y_index < y_elems) {                     \         block[get_local_id(1) * (TR_BLOCK_DIM+1) + get_local_id(0)] =   \           src_mem[idata_offset + index_in];                             \       }                                                                 \       barrier_local();                                                  \-      x_index = group_id_1 * TR_BLOCK_DIM + get_local_id(0)/mulx;       \+      x_index = tblock_id_1 * TR_BLOCK_DIM + get_local_id(0)/mulx;      \       y_index =                                                         \-        group_id_0 * TR_BLOCK_DIM * mulx +                              \+        tblock_id_0 * TR_BLOCK_DIM * mulx +                             \         get_local_id(1) +                                               \         (get_local_id(0)%mulx) * TR_BLOCK_DIM;                          \       int32_t index_out = y_index * y_elems + x_index;                  \@@ -106,16 +106,16 @@         dst_mem[odata_offset + index_out] =                             \           block[get_local_id(0) * (TR_BLOCK_DIM+1) + get_local_id(1)];  \       }                                                                 \-      group_id_2 += get_num_groups(2);                                  \+      tblock_id_2 += get_num_tblocks(2);                                \       global_id_2 += get_global_size(2);                                \     }                                                                   \-    group_id_1 += get_num_groups(1);                                    \+    tblock_id_1 += get_num_tblocks(1);                                  \     global_id_1 += get_global_size(1);                                  \   }                                                                     \ }                                                                       \                                                                         \ FUTHARK_KERNEL_SIZED(TR_BLOCK_DIM, TR_BLOCK_DIM, 1)                     \-void map_transpose_##NAME##_low_width(LOCAL_MEM_PARAM                   \+void map_transpose_##NAME##_low_width(SHARED_MEM_PARAM                  \                                       __global ELEM_TYPE *dst_mem,      \                                       int64_t dst_offset,               \                                       __global ELEM_TYPE *src_mem,      \@@ -127,21 +127,21 @@                                       int32_t muly,                     \                                       int32_t repeat_1,                 \                                       int32_t repeat_2) {               \-  __local ELEM_TYPE* block = (__local ELEM_TYPE*)local_mem;             \-  int group_id_0 = get_group_id(0);                                     \+  __local ELEM_TYPE* block = (__local ELEM_TYPE*)shared_mem;            \+  int tblock_id_0 = get_tblock_id(0);                                   \   int global_id_0 = get_global_id(0);                                   \-  int group_id_1 = get_group_id(1);                                     \+  int tblock_id_1 = get_tblock_id(1);                                   \   int global_id_1 = get_global_id(1);                                   \   for (int i1 = 0; i1 <= repeat_1; i1++) {                              \-    int group_id_2 = get_group_id(2);                                   \+    int tblock_id_2 = get_tblock_id(2);                                 \     int global_id_2 = get_global_id(2);                                 \     for (int i2 = 0; i2 <= repeat_2; i2++) {                            \-      int32_t our_array_offset = group_id_2 * x_elems * y_elems;        \+      int32_t our_array_offset = tblock_id_2 * x_elems * y_elems;       \       int32_t odata_offset = dst_offset + our_array_offset;             \       int32_t idata_offset = src_offset + our_array_offset;             \-      int32_t x_index = group_id_0 * TR_BLOCK_DIM + get_local_id(0)/muly; \+      int32_t x_index = tblock_id_0 * TR_BLOCK_DIM + get_local_id(0)/muly; \       int32_t y_index =                                                 \-        group_id_1 * TR_BLOCK_DIM * muly +                              \+        tblock_id_1 * TR_BLOCK_DIM * muly +                             \         get_local_id(1) + (get_local_id(0)%muly) * TR_BLOCK_DIM;        \       int32_t index_in = y_index * x_elems + x_index;                   \       if (x_index < x_elems && y_index < y_elems) {                     \@@ -149,24 +149,24 @@           src_mem[idata_offset + index_in];                             \       }                                                                 \       barrier_local();                                                  \-      x_index = group_id_1 * TR_BLOCK_DIM * muly +                      \+      x_index = tblock_id_1 * TR_BLOCK_DIM * muly +                     \         get_local_id(0) + (get_local_id(1)%muly) * TR_BLOCK_DIM;        \-      y_index = group_id_0 * TR_BLOCK_DIM + get_local_id(1)/muly;       \+      y_index = tblock_id_0 * TR_BLOCK_DIM + get_local_id(1)/muly;      \       int32_t index_out = y_index * y_elems + x_index;                  \       if (x_index < y_elems && y_index < x_elems) {                     \         dst_mem[odata_offset + index_out] =                             \           block[get_local_id(0) * (TR_BLOCK_DIM+1) + get_local_id(1)];  \       }                                                                 \-      group_id_2 += get_num_groups(2);                                  \-      global_id_2 += get_num_groups(2) * get_local_size(2);             \+      tblock_id_2 += get_num_tblocks(2);                                \+      global_id_2 += get_num_tblocks(2) * get_local_size(2);            \     }                                                                   \-    group_id_1 += get_num_groups(1);                                    \-    global_id_1 += get_num_groups(1) * get_local_size(1);               \+    tblock_id_1 += get_num_tblocks(1);                                  \+    global_id_1 += get_num_tblocks(1) * get_local_size(1);              \   }                                                                     \ }                                                                       \                                                                         \ FUTHARK_KERNEL_SIZED(TR_BLOCK_DIM*TR_BLOCK_DIM, 1, 1)                   \-void map_transpose_##NAME##_small(LOCAL_MEM_PARAM                       \+void map_transpose_##NAME##_small(SHARED_MEM_PARAM                       \                                   __global ELEM_TYPE *dst_mem,          \                                   int64_t dst_offset,                   \                                   __global ELEM_TYPE *src_mem,          \@@ -179,13 +179,13 @@                                   int32_t repeat_1,                     \                                   int32_t repeat_2) {                   \   (void)mulx; (void)muly;                                               \-  __local ELEM_TYPE* block = (__local ELEM_TYPE*)local_mem;             \-  int group_id_0 = get_group_id(0);                                     \+  __local ELEM_TYPE* block = (__local ELEM_TYPE*)shared_mem;            \+  int tblock_id_0 = get_tblock_id(0);                                   \   int global_id_0 = get_global_id(0);                                   \-  int group_id_1 = get_group_id(1);                                     \+  int tblock_id_1 = get_tblock_id(1);                                   \   int global_id_1 = get_global_id(1);                                   \   for (int i1 = 0; i1 <= repeat_1; i1++) {                              \-    int group_id_2 = get_group_id(2);                                   \+    int tblock_id_2 = get_tblock_id(2);                                 \     int global_id_2 = get_global_id(2);                                 \     for (int i2 = 0; i2 <= repeat_2; i2++) {                            \       int32_t our_array_offset = global_id_0/(y_elems * x_elems) * y_elems * x_elems; \@@ -198,16 +198,16 @@       if (global_id_0 < x_elems * y_elems * num_arrays) {               \         dst_mem[odata_offset + index_out] = src_mem[idata_offset + index_in]; \       }                                                                 \-      group_id_2 += get_num_groups(2);                                  \+      tblock_id_2 += get_num_tblocks(2);                                \       global_id_2 += get_global_size(2);                                \     }                                                                   \-    group_id_1 += get_num_groups(1);                                    \+    tblock_id_1 += get_num_tblocks(1);                                  \     global_id_1 += get_global_size(1);                                  \   }                                                                     \ }                                                                       \                                                                         \ FUTHARK_KERNEL_SIZED(TR_BLOCK_DIM*2, TR_TILE_DIM/TR_ELEMS_PER_THREAD, 1)\-void map_transpose_##NAME##_large(LOCAL_MEM_PARAM                       \+void map_transpose_##NAME##_large(SHARED_MEM_PARAM                      \                                   __global ELEM_TYPE *dst_mem,          \                                   int64_t dst_offset,                   \                                   __global ELEM_TYPE *src_mem,          \@@ -220,20 +220,20 @@                                   int32_t repeat_1,                     \                                   int32_t repeat_2) {                   \   (void)mulx; (void)muly;                                               \-  __local ELEM_TYPE* block = (__local ELEM_TYPE*)local_mem;             \-  int group_id_0 = get_group_id(0);                                     \+  __local ELEM_TYPE* block = (__local ELEM_TYPE*)shared_mem;             \+  int tblock_id_0 = get_tblock_id(0);                                   \   int global_id_0 = get_global_id(0);                                   \-  int group_id_1 = get_group_id(1);                                     \+  int tblock_id_1 = get_tblock_id(1);                                   \   int global_id_1 = get_global_id(1);                                   \   for (int i1 = 0; i1 <= repeat_1; i1++) {                              \-    int group_id_2 = get_group_id(2);                                   \+    int tblock_id_2 = get_tblock_id(2);                                 \     int global_id_2 = get_global_id(2);                                 \     for (int i2 = 0; i2 <= repeat_2; i2++) {                            \-      int64_t our_array_offset = group_id_2 * x_elems * y_elems;        \+      int64_t our_array_offset = tblock_id_2 * x_elems * y_elems;       \       int64_t odata_offset = dst_offset + our_array_offset;             \       int64_t idata_offset = src_offset + our_array_offset;             \       int64_t x_index = global_id_0;                                    \-      int64_t y_index = group_id_1 * TR_TILE_DIM + get_local_id(1);     \+      int64_t y_index = tblock_id_1 * TR_TILE_DIM + get_local_id(1);    \       if (x_index < x_elems) {                                          \         for (int64_t j = 0; j < TR_ELEMS_PER_THREAD; j++) {             \           int64_t index_i = (y_index + j * (TR_TILE_DIM/TR_ELEMS_PER_THREAD)) * x_elems + x_index; \@@ -245,8 +245,8 @@         }                                                               \       }                                                                 \       barrier_local();                                                  \-      x_index = group_id_1 * TR_TILE_DIM + get_local_id(0);             \-      y_index = group_id_0 * TR_TILE_DIM + get_local_id(1);             \+      x_index = tblock_id_1 * TR_TILE_DIM + get_local_id(0);            \+      y_index = tblock_id_0 * TR_TILE_DIM + get_local_id(1);            \       if (x_index < y_elems) {                                          \         for (int64_t j = 0; j < TR_ELEMS_PER_THREAD; j++) {             \           int64_t index_out = (y_index + j * (TR_TILE_DIM/TR_ELEMS_PER_THREAD)) * y_elems + x_index; \@@ -257,10 +257,10 @@           }                                                             \         }                                                               \       }                                                                 \-      group_id_2 += get_num_groups(2);                                  \+      tblock_id_2 += get_num_tblocks(2);                                \       global_id_2 += get_global_size(2);                                \     }                                                                   \-    group_id_1 += get_num_groups(1);                                    \+    tblock_id_1 += get_num_tblocks(1);                                  \     global_id_1 += get_global_size(1);                                  \   }                                                                     \ }                                                                       \

@@ -120,7 +120,7 @@     interactive=False,     platform_pref=None,     device_pref=None,-    default_group_size=None,+    default_tblock_size=None,     default_num_groups=None,     default_tile_size=None,     default_reg_tile_size=None,@@ -146,10 +146,10 @@      check_types(self, required_types) -    max_group_size = int(self.device.max_work_group_size)-    max_tile_size = int(np.sqrt(self.device.max_work_group_size))+    max_tblock_size = int(self.device.max_work_tblock_size)+    max_tile_size = int(np.sqrt(self.device.max_work_tblock_size)) -    self.max_group_size = max_group_size+    self.max_tblock_size = max_tblock_size     self.max_tile_size = max_tile_size     self.max_threshold = 0     self.max_num_groups = 0@@ -176,9 +176,9 @@     )     self.failure_is_an_option = np.int32(0) -    if "default_group_size" in sizes:-        default_group_size = sizes["default_group_size"]-        del sizes["default_group_size"]+    if "default_tblock_size" in sizes:+        default_tblock_size = sizes["default_tblock_size"]+        del sizes["default_tblock_size"]      if "default_num_groups" in sizes:         default_num_groups = sizes["default_num_groups"]@@ -196,13 +196,13 @@         default_threshold = sizes["default_threshold"]         del sizes["default_threshold"] -    default_group_size_set = default_group_size != None+    default_tblock_size_set = default_tblock_size != None     default_tile_size_set = default_tile_size != None     default_sizes = apply_size_heuristics(         self,         size_heuristics,         {-            "group_size": default_group_size,+            "tblock_size": default_tblock_size,             "tile_size": default_tile_size,             "reg_tile_size": default_reg_tile_size,             "num_groups": default_num_groups,@@ -210,21 +210,21 @@             "threshold": default_threshold,         },     )-    default_group_size = default_sizes["group_size"]+    default_tblock_size = default_sizes["tblock_size"]     default_num_groups = default_sizes["num_groups"]     default_threshold = default_sizes["threshold"]     default_tile_size = default_sizes["tile_size"]     default_reg_tile_size = default_sizes["reg_tile_size"]     lockstep_width = default_sizes["lockstep_width"] -    if default_group_size > max_group_size:-        if default_group_size_set:+    if default_tblock_size > max_tblock_size:+        if default_tblock_size_set:             sys.stderr.write(                 "Note: Device limits group size to {} (down from {})\n".format(-                    max_tile_size, default_group_size+                    max_tile_size, default_tblock_size                 )             )-        default_group_size = max_group_size+        default_tblock_size = max_tblock_size      if default_tile_size > max_tile_size:         if default_tile_size_set:@@ -247,11 +247,11 @@      self.sizes = {}     for k, v in all_sizes.items():-        if v["class"] == "group_size":-            max_value = max_group_size-            default_value = default_group_size+        if v["class"] == "tblock_size":+            max_value = max_tblock_size+            default_value = default_tblock_size         elif v["class"] == "num_groups":-            max_value = max_group_size  # Intentional!+            max_value = max_tblock_size  # Intentional!             default_value = default_num_groups         elif v["class"] == "tile_size":             max_value = max_tile_size@@ -283,7 +283,7 @@     if len(program_src) >= 0:         build_options += ["-DLOCKSTEP_WIDTH={}".format(lockstep_width)] -        build_options += ["-D{}={}".format("max_group_size", max_group_size)]+        build_options += ["-D{}={}".format("max_tblock_size", max_tblock_size)]          build_options += [             "-D{}={}".format(

@@ -76,8 +76,8 @@  type Path = [(T.Text, Int)] -regexGroups :: Regex -> T.Text -> Maybe [T.Text]-regexGroups regex s = do+regexBlocks :: Regex -> T.Text -> Maybe [T.Text]+regexBlocks regex s = do   (_, _, _, groups) <-     matchM regex s :: Maybe (T.Text, T.Text, T.Text, [T.Text])   Just groups@@ -87,7 +87,7 @@   where     regex = makeRegex ("Compared ([^ ]+) <= (-?[0-9]+)" :: String)     isComparison l = do-      [thresh, val] <- regexGroups regex l+      [thresh, val] <- regexBlocks regex l       val' <- readMaybe $ T.unpack val       pure (thresh, val') @@ -246,7 +246,7 @@      findThreshold :: (T.Text, T.Text) -> Maybe (T.Text, [(T.Text, Bool)])     findThreshold (name, param_class) = do-      [_, grp] <- regexGroups regex param_class+      [_, grp] <- regexBlocks regex param_class       pure         ( name,           filter (not . T.null . fst)

@@ -49,11 +49,6 @@   GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_program(struct futhark_context_config *cfg, const char* s);|]   GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_dump_ptx_to(struct futhark_context_config *cfg, const char* s);|]   GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_load_ptx_from(struct futhark_context_config *cfg, const char* s);|]-  GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_default_group_size(struct futhark_context_config *cfg, int size);|]-  GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_default_num_groups(struct futhark_context_config *cfg, int size);|]-  GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_default_tile_size(struct futhark_context_config *cfg, int size);|]-  GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_default_reg_tile_size(struct futhark_context_config *cfg, int size);|]-  GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_default_threshold(struct futhark_context_config *cfg, int size);|]  cliOptions :: [Option] cliOptions =

@@ -44,8 +44,8 @@      which' = case which of       LockstepWidth -> [C.cexp|ctx->lockstep_width|]-      NumGroups -> [C.cexp|ctx->cfg->default_num_groups|]-      GroupSize -> [C.cexp|ctx->cfg->default_group_size|]+      NumBlocks -> [C.cexp|ctx->cfg->default_num_groups|]+      BlockSize -> [C.cexp|ctx->cfg->default_group_size|]       TileSize -> [C.cexp|ctx->cfg->default_tile_size|]       RegTileSize -> [C.cexp|ctx->cfg->default_reg_tile_size|]       Threshold -> [C.cexp|ctx->cfg->default_threshold|]@@ -105,11 +105,6 @@   GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_program(struct futhark_context_config *cfg, const char* s);|]   GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_dump_binary_to(struct futhark_context_config *cfg, const char* s);|]   GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_load_binary_from(struct futhark_context_config *cfg, const char* s);|]-  GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_default_group_size(struct futhark_context_config *cfg, int size);|]-  GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_default_num_groups(struct futhark_context_config *cfg, int size);|]-  GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_default_tile_size(struct futhark_context_config *cfg, int size);|]-  GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_default_reg_tile_size(struct futhark_context_config *cfg, int size);|]-  GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_default_threshold(struct futhark_context_config *cfg, int size);|]   GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_command_queue(struct futhark_context_config *cfg, typename cl_command_queue);|]   GC.headerDecl GC.MiscDecl [C.cedecl|typename cl_command_queue futhark_context_get_command_queue(struct futhark_context* ctx);|] 

@@ -78,9 +78,9 @@   where     field = "max_" <> prettyString size_class -compileGroupDim :: GroupDim -> GC.CompilerM op s C.Exp-compileGroupDim (Left e) = GC.compileExp e-compileGroupDim (Right kc) = pure $ kernelConstToExp kc+compileBlockDim :: BlockDim -> GC.CompilerM op s C.Exp+compileBlockDim (Left e) = GC.compileExp e+compileBlockDim (Right kc) = pure $ kernelConstToExp kc  genLaunchKernel ::   KernelSafety ->@@ -88,25 +88,25 @@   Count Bytes (TExp Int64) ->   [KernelArg] ->   [Exp] ->-  [GroupDim] ->+  [BlockDim] ->   GC.CompilerM op s ()-genLaunchKernel safety kernel_name local_memory args num_groups group_size = do+genLaunchKernel safety kernel_name shared_memory args num_tblocks tblock_size = do   (arg_params, arg_params_inits, call_args) <-     unzip3 <$> zipWithM mkArgs [(0 :: Int) ..] args -  (grid_x, grid_y, grid_z) <- mkDims <$> mapM GC.compileExp num_groups-  (group_x, group_y, group_z) <- mkDims <$> mapM compileGroupDim group_size+  (grid_x, grid_y, grid_z) <- mkDims <$> mapM GC.compileExp num_tblocks+  (block_x, block_y, block_z) <- mkDims <$> mapM compileBlockDim tblock_size    kernel_fname <- genKernelFunction kernel_name safety arg_params arg_params_inits -  local_memory' <- GC.compileExp $ untyped $ unCount local_memory+  shared_memory' <- GC.compileExp $ untyped $ unCount shared_memory    GC.stm     [C.cstm|{            err = $id:kernel_fname(ctx,                                   $exp:grid_x, $exp:grid_y, $exp:grid_z,-                                  $exp:group_x, $exp:group_y, $exp:group_z,-                                  $exp:local_memory',+                                  $exp:block_x, $exp:block_y, $exp:block_z,+                                  $exp:shared_memory',                                   $args:call_args);            if (err != FUTHARK_SUCCESS) { goto cleanup; }            }|]@@ -152,8 +152,8 @@ callKernel (GetSizeMax v size_class) = do   let e = kernelConstToExp $ SizeMaxConst size_class   GC.stm [C.cstm|$id:v = $exp:e;|]-callKernel (LaunchKernel safety kernel_name local_memory args num_groups group_size) =-  genLaunchKernel safety kernel_name local_memory args num_groups group_size+callKernel (LaunchKernel safety kernel_name shared_memory args num_tblocks tblock_size) =+  genLaunchKernel safety kernel_name shared_memory args num_tblocks tblock_size  copygpu2gpu :: GC.DoLMADCopy op s copygpu2gpu _ t shape dst (dstoffset, dststride) src (srcoffset, srcstride) = do@@ -323,17 +323,31 @@         optionAction = [C.cstm|futhark_context_config_set_device(cfg, optarg);|]       },     Option+      { optionLongName = "default-thread-block-size",+        optionShortName = Nothing,+        optionArgument = RequiredArgument "INT",+        optionDescription = "The default size of thread blocks that are launched.",+        optionAction = [C.cstm|futhark_context_config_set_default_thread_block_size(cfg, atoi(optarg));|]+      },+    Option+      { optionLongName = "default-grid-size",+        optionShortName = Nothing,+        optionArgument = RequiredArgument "INT",+        optionDescription = "The default number of thread blocks that are launched.",+        optionAction = [C.cstm|futhark_context_config_set_default_grid_size(cfg, atoi(optarg));|]+      },+    Option       { optionLongName = "default-group-size",         optionShortName = Nothing,         optionArgument = RequiredArgument "INT",-        optionDescription = "The default size of workgroups that are launched.",+        optionDescription = "Alias for --default-thread-block-size.",         optionAction = [C.cstm|futhark_context_config_set_default_group_size(cfg, atoi(optarg));|]       },     Option       { optionLongName = "default-num-groups",         optionShortName = Nothing,         optionArgument = RequiredArgument "INT",-        optionDescription = "The default number of workgroups that are launched.",+        optionDescription = "Alias for --default-num-thread-blocks.",         optionAction = [C.cstm|futhark_context_config_set_default_num_groups(cfg, atoi(optarg));|]       },     Option@@ -440,3 +454,11 @@   GC.generateProgramStruct    GC.onClear [C.citem|if (ctx->error == NULL) { gpu_free_all(ctx); }|]++  GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_default_thread_block_size(struct futhark_context_config *cfg, int size);|]+  GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_default_grid_size(struct futhark_context_config *cfg, int size);|]+  GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_default_group_size(struct futhark_context_config *cfg, int size);|]+  GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_default_num_groups(struct futhark_context_config *cfg, int size);|]+  GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_default_tile_size(struct futhark_context_config *cfg, int size);|]+  GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_default_reg_tile_size(struct futhark_context_config *cfg, int size);|]+  GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_default_threshold(struct futhark_context_config *cfg, int size);|]

@@ -573,7 +573,7 @@       size_default_inits = map (intinit . fromMaybe 0 . sizeDefault) param_classes       size_decls = map (\k -> [C.csdecl|typename int64_t *$id:k;|]) param_names       num_params = length params-  earlyDecl [C.cedecl|struct tuning_params { $sdecls:size_decls };|]+  earlyDecl [C.cedecl|struct tuning_params { int dummy; $sdecls:size_decls };|]   earlyDecl [C.cedecl|static const int num_tuning_params = $int:num_params;|]   earlyDecl [C.cedecl|static const char *tuning_param_names[] = { $inits:size_name_inits, NULL };|]   earlyDecl [C.cedecl|static const char *tuning_param_vars[] = { $inits:size_var_inits, NULL };|]

@@ -47,11 +47,6 @@   GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_device(struct futhark_context_config *cfg, const char* s);|]   GC.headerDecl GC.InitDecl [C.cedecl|const char* futhark_context_config_get_program(struct futhark_context_config *cfg);|]   GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_program(struct futhark_context_config *cfg, const char* s);|]-  GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_default_group_size(struct futhark_context_config *cfg, int size);|]-  GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_default_num_groups(struct futhark_context_config *cfg, int size);|]-  GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_default_tile_size(struct futhark_context_config *cfg, int size);|]-  GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_default_reg_tile_size(struct futhark_context_config *cfg, int size);|]-  GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_default_threshold(struct futhark_context_config *cfg, int size);|]  cliOptions :: [Option] cliOptions =

@@ -395,6 +395,7 @@                                  int subtask_id,                                  int tid) {                        (void)subtask_id;+                       (void)tid;                        int err = 0;                        struct $id:fstruct *$id:fstruct = (struct $id:fstruct*) args;                        struct futhark_context *ctx = $id:fstruct->ctx;

@@ -59,8 +59,8 @@           Assign (Var "build_options") $ List [],           Assign (Var "preferred_device") None,           Assign (Var "default_threshold") None,-          Assign (Var "default_group_size") None,-          Assign (Var "default_num_groups") None,+          Assign (Var "default_tblock_size") None,+          Assign (Var "default_num_tblocks") None,           Assign (Var "default_tile_size") None,           Assign (Var "default_reg_tile_size") None,           Assign (Var "fut_opencl_src") $ RawStringLiteral $ opencl_prelude <> opencl_code@@ -83,8 +83,8 @@             "interactive=False",             "platform_pref=preferred_platform",             "device_pref=preferred_device",-            "default_group_size=default_group_size",-            "default_num_groups=default_num_groups",+            "default_tblock_size=default_tblock_size",+            "default_num_tblocks=default_num_tblocks",             "default_tile_size=default_tile_size",             "default_reg_tile_size=default_reg_tile_size",             "default_threshold=default_threshold",@@ -128,14 +128,14 @@               optionShortName = Nothing,               optionArgument = RequiredArgument "int",               optionAction =-                [Assign (Var "default_group_size") $ Var "optarg"]+                [Assign (Var "default_tblock_size") $ Var "optarg"]             },           Option             { optionLongName = "default-num-groups",               optionShortName = Nothing,               optionArgument = RequiredArgument "int",               optionAction =-                [Assign (Var "default_num_groups") $ Var "optarg"]+                [Assign (Var "default_num_tblocks") $ Var "optarg"]             },           Option             { optionLongName = "default-tile-size",@@ -212,9 +212,9 @@ kernelConstToExp (Imp.SizeMaxConst size_class) =   Var $ "self.max_" <> prettyString size_class -compileGroupDim :: Imp.GroupDim -> CompilerM op s PyExp-compileGroupDim (Left e) = asLong <$> compileExp e-compileGroupDim (Right kc) = pure $ kernelConstToExp kc+compileBlockDim :: Imp.BlockDim -> CompilerM op s PyExp+compileBlockDim (Left e) = asLong <$> compileExp e+compileBlockDim (Right kc) = pure $ kernelConstToExp kc  callKernel :: OpCompiler Imp.OpenCL () callKernel (Imp.GetSize v key) = do@@ -227,14 +227,14 @@ callKernel (Imp.GetSizeMax v size_class) = do   v' <- compileVar v   stm $ Assign v' $ kernelConstToExp $ Imp.SizeMaxConst size_class-callKernel (Imp.LaunchKernel safety name local_memory args num_workgroups workgroup_size) = do-  num_workgroups' <- mapM (fmap asLong . compileExp) num_workgroups-  workgroup_size' <- mapM compileGroupDim workgroup_size-  let kernel_size = zipWith mult_exp num_workgroups' workgroup_size'+callKernel (Imp.LaunchKernel safety name shared_memory args num_threadblocks worktblock_size) = do+  num_threadblocks' <- mapM (fmap asLong . compileExp) num_threadblocks+  worktblock_size' <- mapM compileBlockDim worktblock_size+  let kernel_size = zipWith mult_exp num_threadblocks' worktblock_size'       total_elements = foldl mult_exp (Integer 1) kernel_size       cond = BinOp "!=" total_elements (Integer 0)-  local_memory' <- compileExp $ Imp.untyped $ Imp.unCount local_memory-  body <- collect $ launchKernel name safety kernel_size workgroup_size' local_memory' args+  shared_memory' <- compileExp $ Imp.untyped $ Imp.unCount shared_memory+  body <- collect $ launchKernel name safety kernel_size worktblock_size' shared_memory' args   stm $ If cond body []    when (safety >= Imp.SafetyFull) $@@ -252,9 +252,9 @@   PyExp ->   [Imp.KernelArg] ->   CompilerM op s ()-launchKernel kernel_name safety kernel_dims workgroup_dims local_memory args = do+launchKernel kernel_name safety kernel_dims threadblock_dims shared_memory args = do   let kernel_dims' = Tuple kernel_dims-      workgroup_dims' = Tuple workgroup_dims+      threadblock_dims' = Tuple threadblock_dims       kernel_name' = "self." <> zEncodeText (nameToText kernel_name) <> "_var"   args' <- mapM processKernelArg args   let failure_args =@@ -266,13 +266,13 @@           ]   stm . Exp $     simpleCall (T.unpack $ kernel_name' <> ".set_args") $-      [simpleCall "cl.LocalMemory" [simpleCall "max" [local_memory, Integer 1]]]+      [simpleCall "cl.SharedMemory" [simpleCall "max" [shared_memory, Integer 1]]]         ++ failure_args         ++ args'   stm . Exp $     simpleCall       "cl.enqueue_nd_range_kernel"-      [Var "self.queue", Var (T.unpack kernel_name'), kernel_dims', workgroup_dims']+      [Var "self.queue", Var (T.unpack kernel_name'), kernel_dims', threadblock_dims']   finishIfSynchronous   where     processKernelArg :: Imp.KernelArg -> CompilerM op s PyExp

@@ -59,8 +59,8 @@                                    interactive=interactive,                                    platform_pref=platform_pref,                                    device_pref=device_pref,-                                   default_group_size=default_group_size,-                                   default_num_groups=default_num_groups,+                                   default_tblock_size=default_tblock_size,+                                   default_num_tblocks=default_num_tblocks,                                    default_tile_size=default_tile_size,                                    default_reg_tile_size=default_reg_tile_size,                                    default_threshold=default_threshold,@@ -128,8 +128,8 @@          which' = case which of           LockstepWidth -> String "lockstep_width"-          NumGroups -> String "num_groups"-          GroupSize -> String "group_size"+          NumBlocks -> String "num_tblocks"+          BlockSize -> String "tblock_size"           TileSize -> String "tile_size"           RegTileSize -> String "reg_tile_size"           Threshold -> String "threshold"

@@ -10,7 +10,7 @@     KernelOp (..),     Fence (..),     AtomicOp (..),-    GroupDim,+    BlockDim,     Kernel (..),     KernelUse (..),     module Futhark.CodeGen.ImpCode,@@ -49,16 +49,16 @@   | GetSizeMax VName SizeClass   deriving (Show) --- | The size of one dimension of a group.-type GroupDim = Either Exp KernelConst+-- | The size of one dimension of a block.+type BlockDim = Either Exp KernelConst  -- | A generic kernel containing arbitrary kernel code. data Kernel = Kernel   { kernelBody :: Code KernelOp,     -- | The host variables referenced by the kernel.     kernelUses :: [KernelUse],-    kernelNumGroups :: [Exp],-    kernelGroupSize :: [GroupDim],+    kernelNumBlocks :: [Exp],+    kernelBlockSize :: [BlockDim],     -- | A short descriptive and _unique_ name - should be     -- alphanumeric and without spaces.     kernelName :: Name,@@ -71,7 +71,7 @@     -- | If true, multi-versioning branches will consider this kernel     -- when considering the local memory requirements. Set this to     -- false for kernels that do their own checking.-    kernelCheckLocalMemory :: Bool+    kernelCheckSharedMemory :: Bool   }   deriving (Show) @@ -134,24 +134,24 @@   freeIn' kernel =     freeIn'       ( kernelBody kernel,-        kernelNumGroups kernel,-        kernelGroupSize kernel+        kernelNumBlocks kernel,+        kernelBlockSize kernel       )  instance Pretty Kernel where   pretty kernel =     "kernel"       <+> brace-        ( "groups"-            <+> brace (pretty $ kernelNumGroups kernel)-            </> "group_size"-            <+> brace (list $ map (either pretty pretty) $ kernelGroupSize kernel)+        ( "blocks"+            <+> brace (pretty $ kernelNumBlocks kernel)+            </> "tblock_size"+            <+> brace (list $ map (either pretty pretty) $ kernelBlockSize kernel)             </> "uses"             <+> brace (commasep $ map pretty $ kernelUses kernel)             </> "failure_tolerant"             <+> brace (pretty $ kernelFailureTolerant kernel)-            </> "check_local_memory"-            <+> brace (pretty $ kernelCheckLocalMemory kernel)+            </> "check_shared_memory"+            <+> brace (pretty $ kernelCheckSharedMemory kernel)             </> "body"             <+> brace (pretty $ kernelBody kernel)         )@@ -163,7 +163,7 @@  -- | An operation that occurs within a kernel body. data KernelOp-  = GetGroupId VName Int+  = GetBlockId VName Int   | GetLocalId VName Int   | GetLocalSize VName Int   | GetLockstepWidth VName@@ -211,10 +211,10 @@   freeIn' (AtomicXchg _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x  instance Pretty KernelOp where-  pretty (GetGroupId dest i) =+  pretty (GetBlockId dest i) =     pretty dest       <+> "<-"-      <+> "get_group_id"+      <+> "get_tblock_id"       <> parens (pretty i)   pretty (GetLocalId dest i) =     pretty dest

@@ -16,7 +16,7 @@     numFailureParams,     KernelTarget (..),     FailureMsg (..),-    GroupDim,+    BlockDim,     KernelConst (..),     KernelConstExp,     module Futhark.CodeGen.ImpCode,@@ -27,7 +27,7 @@ import Data.Map qualified as M import Data.Text qualified as T import Futhark.CodeGen.ImpCode-import Futhark.CodeGen.ImpCode.GPU (GroupDim, KernelConst (..), KernelConstExp)+import Futhark.CodeGen.ImpCode.GPU (BlockDim, KernelConst (..), KernelConstExp) import Futhark.IR.GPU.Sizes import Futhark.Util.Pretty @@ -97,7 +97,7 @@  -- | Host-level OpenCL operation. data OpenCL-  = LaunchKernel KernelSafety KernelName (Count Bytes (TExp Int64)) [KernelArg] [Exp] [GroupDim]+  = LaunchKernel KernelSafety KernelName (Count Bytes (TExp Int64)) [KernelArg] [Exp] [BlockDim]   | GetSize VName Name   | CmpSizeLe VName Name Exp   | GetSizeMax VName SizeClass

@@ -108,24 +108,24 @@     =<< toExp x opCompiler (Pat [pe]) (Inner (SizeOp (GetSizeMax size_class))) =   sOp $ Imp.GetSizeMax (patElemName pe) size_class-opCompiler (Pat [pe]) (Inner (SizeOp (CalcNumGroups w64 max_num_groups_key group_size))) = do+opCompiler (Pat [pe]) (Inner (SizeOp (CalcNumBlocks w64 max_num_tblocks_key tblock_size))) = do   fname <- askFunction-  max_num_groups :: TV Int32 <- dPrim "max_num_groups" int32+  max_num_tblocks :: TV Int32 <- dPrim "max_num_tblocks" int32   sOp $-    Imp.GetSize (tvVar max_num_groups) (keyWithEntryPoint fname max_num_groups_key) $-      sizeClassWithEntryPoint fname SizeNumGroups+    Imp.GetSize (tvVar max_num_tblocks) (keyWithEntryPoint fname max_num_tblocks_key) $+      sizeClassWithEntryPoint fname SizeGrid -  -- If 'w' is small, we launch fewer groups than we normally would.-  -- We don't want any idle groups.+  -- If 'w' is small, we launch fewer blocks than we normally would.+  -- We don't want any idle blocks.   --   -- The calculations are done with 64-bit integers to avoid overflow   -- issues.-  let num_groups_maybe_zero =-        sMin64 (pe64 w64 `divUp` pe64 group_size) $-          sExt64 (tvExp max_num_groups)-  -- We also don't want zero groups.-  let num_groups = sMax64 1 num_groups_maybe_zero-  mkTV (patElemName pe) int32 <-- sExt32 num_groups+  let num_tblocks_maybe_zero =+        sMin64 (pe64 w64 `divUp` pe64 tblock_size) $+          sExt64 (tvExp max_num_tblocks)+  -- We also don't want zero blocks.+  let num_tblocks = sMax64 1 num_tblocks_maybe_zero+  mkTV (patElemName pe) int32 <-- sExt32 num_tblocks opCompiler dest (Inner (SegOp op)) =   segOpCompiler dest op opCompiler (Pat pes) (Inner (GPUBody _ (Body _ stms res))) = do@@ -170,8 +170,8 @@ -- otherwise protected by their own multi-versioning branches deeper -- down.  Currently the compiler will not generate multi-versioning -- that makes this a problem, but it might in the future.-checkLocalMemoryReqs :: (VName -> Bool) -> Imp.HostCode -> CallKernelGen (Maybe (Imp.TExp Bool))-checkLocalMemoryReqs in_scope code = do+checkSharedMemoryReqs :: (VName -> Bool) -> Imp.HostCode -> CallKernelGen (Maybe (Imp.TExp Bool))+checkSharedMemoryReqs in_scope code = do   let alloc_sizes = map (sum . map alignedSize . localAllocSizes . Imp.kernelBody) $ getGPU code    -- If any of the sizes involve a variable that is not known at this@@ -179,17 +179,17 @@   if not $ all in_scope $ namesToList $ freeIn alloc_sizes     then pure Nothing     else do-      local_memory_capacity :: TV Int32 <- dPrim "local_memory_capacity" int32-      sOp $ Imp.GetSizeMax (tvVar local_memory_capacity) SizeLocalMemory+      shared_memory_capacity :: TV Int32 <- dPrim "shared_memory_capacity" int32+      sOp $ Imp.GetSizeMax (tvVar shared_memory_capacity) SizeSharedMemory -      let local_memory_capacity_64 =-            sExt64 $ tvExp local_memory_capacity+      let shared_memory_capacity_64 =+            sExt64 $ tvExp shared_memory_capacity           fits size =-            unCount size .<=. local_memory_capacity_64+            unCount size .<=. shared_memory_capacity_64       pure $ Just $ foldl' (.&&.) true (map fits alloc_sizes)   where     getGPU = foldMap getKernel-    getKernel (Imp.CallKernel k) | Imp.kernelCheckLocalMemory k = [k]+    getKernel (Imp.CallKernel k) | Imp.kernelCheckSharedMemory k = [k]     getKernel _ = []      localAllocSizes = foldMap localAllocSize@@ -237,8 +237,8 @@       if shapeRank shape == 0         then copyDWIM (patElemName pe) [] se []         else sReplicate (patElemName pe) se--- Allocation in the "local" space is just a placeholder.-expCompiler _ (Op (Alloc _ (Space "local"))) =+-- Allocation in the "shared" space is just a placeholder.+expCompiler _ (Op (Alloc _ (Space "shared"))) =   pure () expCompiler pat (WithAcc inputs lam) =   withAcc pat inputs lam@@ -252,7 +252,7 @@   scope <- askScope   tcode <- collect $ compileBody dest $ caseBody first_case   fcode <- collect $ expCompiler dest $ Match cond cases defbranch sort-  check <- checkLocalMemoryReqs (`M.member` scope) tcode+  check <- checkSharedMemoryReqs (`M.member` scope) tcode   let matches = caseMatch cond (casePat first_case)   emit $ case check of     Nothing -> fcode

@@ -11,9 +11,9 @@     HostEnv (..),     Target (..),     KernelEnv (..),-    groupReduce,-    groupScan,-    groupLoop,+    blockReduce,+    blockScan,+    blockLoop,     isActive,     sKernel,     sKernelThread,@@ -23,9 +23,9 @@     allocLocal,     kernelAlloc,     compileThreadResult,-    virtualiseGroups,+    virtualiseBlocks,     kernelLoop,-    groupCoverSpace,+    blockCoverSpace,     fenceForArrays,     updateAcc,     genZeroes,@@ -91,14 +91,14 @@ data KernelConstants = KernelConstants   { kernelGlobalThreadId :: Imp.TExp Int32,     kernelLocalThreadId :: Imp.TExp Int32,-    kernelGroupId :: Imp.TExp Int32,+    kernelBlockId :: Imp.TExp Int32,     kernelGlobalThreadIdVar :: VName,     kernelLocalThreadIdVar :: VName,-    kernelGroupIdVar :: VName,-    kernelNumGroupsCount :: Count NumGroups SubExp,-    kernelGroupSizeCount :: Count GroupSize SubExp,-    kernelNumGroups :: Imp.TExp Int64,-    kernelGroupSize :: Imp.TExp Int64,+    kernelBlockIdVar :: VName,+    kernelNumBlocksCount :: Count NumBlocks SubExp,+    kernelBlockSizeCount :: Count BlockSize SubExp,+    kernelNumBlocks :: Imp.TExp Int64,+    kernelBlockSize :: Imp.TExp Int64,     kernelNumThreads :: Imp.TExp Int32,     kernelWaveSize :: Imp.TExp Int32,     -- | A mapping from dimensions of nested SegOps to already@@ -126,7 +126,7 @@   -- Handled by the declaration of the memory block, which is then   -- translated to an actual scalar variable during C code generation.   pure ()-kernelAlloc (Pat [mem]) size (Space "local") =+kernelAlloc (Pat [mem]) size (Space "shared") =   allocLocal (patElemName mem) $ Imp.bytes $ pe64 size kernelAlloc (Pat [mem]) _ _ =   compilerLimitationS $ "Cannot allocate memory block " ++ prettyString mem ++ " in kernel."@@ -187,7 +187,7 @@ -- | Assign iterations of a for-loop to all threads in the kernel. -- The passed-in function is invoked with the (symbolic) iteration. -- The body must contain thread-level code.  For multidimensional--- loops, use 'groupCoverSpace'.+-- loops, use 'blockCoverSpace'. kernelLoop ::   (IntExp t) =>   Imp.TExp t ->@@ -205,47 +205,47 @@           i <- dPrimVE "i" $ chunk_i * num_threads + tid           sWhen (i .<. n) $ f i --- | Assign iterations of a for-loop to threads in the workgroup.  The+-- | Assign iterations of a for-loop to threads in the threadblock.  The -- passed-in function is invoked with the (symbolic) iteration.  For--- multidimensional loops, use 'groupCoverSpace'.-groupLoop ::+-- multidimensional loops, use 'blockCoverSpace'.+blockLoop ::   (IntExp t) =>   Imp.TExp t ->   (Imp.TExp t -> InKernelGen ()) ->   InKernelGen ()-groupLoop n f = do+blockLoop n f = do   constants <- kernelConstants <$> askEnv   kernelLoop     (kernelLocalThreadId constants `sExtAs` n)-    (kernelGroupSize constants `sExtAs` n)+    (kernelBlockSize constants `sExtAs` n)     n     f  -- | Iterate collectively though a multidimensional space, such that--- all threads in the group participate.  The passed-in function is+-- all threads in the block participate.  The passed-in function is -- invoked with a (symbolic) point in the index space.-groupCoverSpace ::+blockCoverSpace ::   (IntExp t) =>   [Imp.TExp t] ->   ([Imp.TExp t] -> InKernelGen ()) ->   InKernelGen ()-groupCoverSpace ds f = do+blockCoverSpace ds f = do   constants <- kernelConstants <$> askEnv-  let group_size = kernelGroupSize constants+  let tblock_size = kernelBlockSize constants   case splitFromEnd 1 ds of     -- Optimise the case where the inner dimension of the space is-    -- equal to the group size.+    -- equal to the block size.     (ds', [last_d])-      | last_d == (group_size `sExtAs` last_d) -> do+      | last_d == (tblock_size `sExtAs` last_d) -> do           let ltid = kernelLocalThreadId constants `sExtAs` last_d           sLoopSpace ds' $ \ds_is ->             f $ ds_is ++ [ltid]     _ ->-      groupLoop (product ds) $ f . unflattenIndex ds+      blockLoop (product ds) $ f . unflattenIndex ds  -- Which fence do we need to protect shared access to this memory space? fenceForSpace :: Space -> Imp.Fence-fenceForSpace (Space "local") = Imp.FenceLocal+fenceForSpace (Space "shared") = Imp.FenceLocal fenceForSpace _ = Imp.FenceGlobal  -- | If we are touching these arrays, which kind of fence do we need?@@ -281,11 +281,11 @@ -- (with primitive non-commutative operators only). getChunkSize :: [Type] -> Imp.KernelConstExp getChunkSize types = do-  let max_group_size = Imp.SizeMaxConst SizeGroup-      max_group_mem = Imp.SizeMaxConst SizeLocalMemory-      max_group_reg = Imp.SizeMaxConst SizeRegisters-      k_mem = le64 max_group_mem `quot` le64 max_group_size-      k_reg = le64 max_group_reg `quot` le64 max_group_size+  let max_tblock_size = Imp.SizeMaxConst SizeThreadBlock+      max_block_mem = Imp.SizeMaxConst SizeSharedMemory+      max_block_reg = Imp.SizeMaxConst SizeRegisters+      k_mem = le64 max_block_mem `quot` le64 max_tblock_size+      k_reg = le64 max_block_reg `quot` le64 max_tblock_size       types' = map elemType $ filter primType types       sizes = map primByteSize types' @@ -297,7 +297,7 @@       reg_constraint = (k_reg - 1 - sum_sizes') `quot` (2 * sum_sizes')   untyped $ sMax64 1 $ sMin64 mem_constraint reg_constraint -inBlockScan ::+inChunkScan ::   KernelConstants ->   Maybe (Imp.TExp Int32 -> Imp.TExp Int32 -> Imp.TExp Bool) ->   Imp.TExp Int64 ->@@ -308,7 +308,7 @@   InKernelGen () ->   Lambda GPUMem ->   InKernelGen ()-inBlockScan constants seg_flag arrs_full_size lockstep_width block_size active arrs barrier scan_lam = everythingVolatile $ do+inChunkScan constants seg_flag arrs_full_size lockstep_width block_size active arrs barrier scan_lam = everythingVolatile $ do   skip_threads <- dPrim "skip_threads" int32   let actual_params = lambdaParams scan_lam       (x_params, y_params) =@@ -398,14 +398,14 @@         copyDWIMFix arr [ltid] (Var $ paramName x) []       copyDWIM (paramName y) [] (Var $ paramName x) [] -groupScan ::+blockScan ::   Maybe (Imp.TExp Int32 -> Imp.TExp Int32 -> Imp.TExp Bool) ->   Imp.TExp Int64 ->   Imp.TExp Int64 ->   Lambda GPUMem ->   [VName] ->   InKernelGen ()-groupScan seg_flag arrs_full_size w lam arrs = do+blockScan seg_flag arrs_full_size w lam arrs = do   constants <- kernelConstants <$> askEnv   renamed_lam <- renameLambda lam @@ -419,27 +419,26 @@    fence <- fenceForArrays arrs -  -- The scan works by splitting the group into blocks, which are-  -- scanned separately.  Typically, these blocks are smaller than-  -- the lockstep width, which enables barrier-free execution inside-  -- them.+  -- The scan works by splitting the block into chunks, which are+  -- scanned separately. Typically, these chunks are at most the+  -- lockstep width, which enables barrier-free execution inside them.   ---  -- We hardcode the block size here.  The only requirement is that-  -- it should not be less than the square root of the group size.-  -- With 32, we will work on groups of size 1024 or smaller, which-  -- fits every device Troels has seen.  Still, it would be nicer if-  -- it were a runtime parameter.  Some day.-  let block_size = 32+  -- We hardcode the chunk size here. The only requirement is that it+  -- should not be less than the square root of the block size. With+  -- 32, we will work on blocks of size 1024 or smaller, which fits+  -- every device Troels has seen. Still, it would be nicer if it were+  -- a runtime parameter. Some day.+  let chunk_size = 32       simd_width = kernelWaveSize constants-      block_id = ltid32 `quot` block_size-      in_block_id = ltid32 - block_id * block_size-      doInBlockScan seg_flag' active =-        inBlockScan+      chunk_id = ltid32 `quot` chunk_size+      in_chunk_id = ltid32 - chunk_id * chunk_size+      doInChunkScan seg_flag' active =+        inChunkScan           constants           seg_flag'           arrs_full_size           simd_width-          block_size+          chunk_size           active           arrs           barrier@@ -456,34 +455,34 @@         | otherwise =             sOp $ Imp.ErrorSync Imp.FenceLocal -      group_offset = sExt64 (kernelGroupId constants) * kernelGroupSize constants+      block_offset = sExt64 (kernelBlockId constants) * kernelBlockSize constants        writeBlockResult p arr         | isPrimParam p =-            copyDWIMFix arr [sExt64 block_id] (Var $ paramName p) []+            copyDWIMFix arr [sExt64 chunk_id] (Var $ paramName p) []         | otherwise =-            copyDWIMFix arr [group_offset + sExt64 block_id] (Var $ paramName p) []+            copyDWIMFix arr [block_offset + sExt64 chunk_id] (Var $ paramName p) []        readPrevBlockResult p arr         | isPrimParam p =-            copyDWIMFix (paramName p) [] (Var arr) [sExt64 block_id - 1]+            copyDWIMFix (paramName p) [] (Var arr) [sExt64 chunk_id - 1]         | otherwise =-            copyDWIMFix (paramName p) [] (Var arr) [group_offset + sExt64 block_id - 1]+            copyDWIMFix (paramName p) [] (Var arr) [block_offset + sExt64 chunk_id - 1] -  doInBlockScan seg_flag ltid_in_bounds lam+  doInChunkScan seg_flag ltid_in_bounds lam   barrier -  let is_first_block = block_id .==. 0+  let is_first_block = chunk_id .==. 0   when array_scan $ do     sComment "save correct values for first block" $       sWhen is_first_block $         forM_ (zip x_params arrs) $ \(x, arr) ->           unless (isPrimParam x) $-            copyDWIMFix arr [arrs_full_size + group_offset + sExt64 block_size + ltid] (Var $ paramName x) []+            copyDWIMFix arr [arrs_full_size + block_offset + sExt64 chunk_size + ltid] (Var $ paramName x) []      barrier -  let last_in_block = in_block_id .==. block_size - 1+  let last_in_block = in_chunk_id .==. chunk_size - 1   sComment "last thread of block 'i' writes its result to offset 'i'" $     sWhen (last_in_block .&&. ltid_in_bounds) $       everythingVolatile $@@ -494,10 +493,10 @@   let first_block_seg_flag = do         flag_true <- seg_flag         Just $ \from to ->-          flag_true (from * block_size + block_size - 1) (to * block_size + block_size - 1)+          flag_true (from * chunk_size + chunk_size - 1) (to * chunk_size + chunk_size - 1)   sComment     "scan the first block, after which offset 'i' contains carry-in for block 'i+1'"-    $ doInBlockScan first_block_seg_flag (is_first_block .&&. ltid_in_bounds) renamed_lam+    $ doInChunkScan first_block_seg_flag (is_first_block .&&. ltid_in_bounds) renamed_lam    errorsync @@ -508,9 +507,9 @@           unless (isPrimParam x) $             copyDWIMFix               arr-              [arrs_full_size + group_offset + ltid]+              [arrs_full_size + block_offset + ltid]               (Var arr)-              [arrs_full_size + group_offset + sExt64 block_size + ltid]+              [arrs_full_size + block_offset + sExt64 chunk_size + ltid]      barrier @@ -526,7 +525,7 @@             sUnless no_carry_in $ compileBody' x_params $ lambdaBody lam         | Just flag_true <- seg_flag = do             inactive <--              dPrimVE "inactive" $ flag_true (block_id * block_size - 1) ltid32+              dPrimVE "inactive" $ flag_true (chunk_id * chunk_size - 1) ltid32             sUnless no_carry_in . sWhen inactive . forM_ (zip x_params y_params) $ \(x, y) ->               copyDWIM (paramName x) [] (Var (paramName y)) []             -- The convoluted control flow is to ensure all threads@@ -552,26 +551,26 @@       forM_ (zip3 x_params y_params arrs) $ \(x, y, arr) ->         if isPrimParam y           then copyDWIMFix arr [ltid] (Var $ paramName y) []-          else copyDWIMFix (paramName x) [] (Var arr) [arrs_full_size + group_offset + ltid]+          else copyDWIMFix (paramName x) [] (Var arr) [arrs_full_size + block_offset + ltid]    barrier -groupReduce ::+blockReduce ::   Imp.TExp Int32 ->   Lambda GPUMem ->   [VName] ->   InKernelGen ()-groupReduce w lam arrs = do+blockReduce w lam arrs = do   offset <- dPrim "offset" int32-  groupReduceWithOffset offset w lam arrs+  blockReduceWithOffset offset w lam arrs -groupReduceWithOffset ::+blockReduceWithOffset ::   TV Int32 ->   Imp.TExp Int32 ->   Lambda GPUMem ->   [VName] ->   InKernelGen ()-groupReduceWithOffset offset w lam arrs = do+blockReduceWithOffset offset w lam arrs = do   constants <- kernelConstants <$> askEnv    let local_tid = kernelLocalThreadId constants@@ -619,10 +618,10 @@       in_wave_reduce = everythingVolatile do_reduce        wave_size = kernelWaveSize constants-      group_size = kernelGroupSize constants+      tblock_size = kernelBlockSize constants       wave_id = local_tid `quot` wave_size       in_wave_id = local_tid - wave_id * wave_size-      num_waves = (sExt32 group_size + wave_size - 1) `quot` wave_size+      num_waves = (sExt32 tblock_size + wave_size - 1) `quot` wave_size       arg_in_bounds = local_tid + tvExp offset .<. w        doing_in_wave_reductions =@@ -919,7 +918,7 @@       case t of         Array {} -> pure Nothing         Acc {} -> pure Nothing-        Mem (Space "local") -> pure Nothing+        Mem (Space "shared") -> pure Nothing         Mem {} -> pure $ Just $ Imp.MemoryUse var         Prim bt ->           isConstExp vtable (Imp.var var bt) >>= \case@@ -948,30 +947,30 @@     hasExp (MemVar e _) = e  kernelInitialisationSimple ::-  Count NumGroups SubExp ->-  Count GroupSize SubExp ->+  Count NumBlocks SubExp ->+  Count BlockSize SubExp ->   CallKernelGen (KernelConstants, InKernelGen ())-kernelInitialisationSimple num_groups group_size = do+kernelInitialisationSimple num_tblocks tblock_size = do   global_tid <- newVName "global_tid"   local_tid <- newVName "local_tid"-  group_id <- newVName "group_tid"+  tblock_id <- newVName "block_id"   wave_size <- newVName "wave_size"-  inner_group_size <- newVName "group_size"-  let num_groups' = Imp.pe64 (unCount num_groups)-      group_size' = Imp.pe64 (unCount group_size)+  inner_tblock_size <- newVName "tblock_size"+  let num_tblocks' = Imp.pe64 (unCount num_tblocks)+      tblock_size' = Imp.pe64 (unCount tblock_size)       constants =         KernelConstants           { kernelGlobalThreadId = Imp.le32 global_tid,             kernelLocalThreadId = Imp.le32 local_tid,-            kernelGroupId = Imp.le32 group_id,+            kernelBlockId = Imp.le32 tblock_id,             kernelGlobalThreadIdVar = global_tid,             kernelLocalThreadIdVar = local_tid,-            kernelNumGroupsCount = num_groups,-            kernelGroupSizeCount = group_size,-            kernelGroupIdVar = group_id,-            kernelNumGroups = num_groups',-            kernelGroupSize = group_size',-            kernelNumThreads = sExt32 (group_size' * num_groups'),+            kernelNumBlocksCount = num_tblocks,+            kernelBlockSizeCount = tblock_size,+            kernelBlockIdVar = tblock_id,+            kernelNumBlocks = num_tblocks',+            kernelBlockSize = tblock_size',+            kernelNumThreads = sExt32 (tblock_size' * num_tblocks'),             kernelWaveSize = Imp.le32 wave_size,             kernelLocalIdMap = mempty,             kernelChunkItersMap = mempty@@ -979,15 +978,15 @@    let set_constants = do         dPrim_ local_tid int32-        dPrim_ inner_group_size int64+        dPrim_ inner_tblock_size int64         dPrim_ wave_size int32-        dPrim_ group_id int32+        dPrim_ tblock_id int32          sOp (Imp.GetLocalId local_tid 0)-        sOp (Imp.GetLocalSize inner_group_size 0)+        sOp (Imp.GetLocalSize inner_tblock_size 0)         sOp (Imp.GetLockstepWidth wave_size)-        sOp (Imp.GetGroupId group_id 0)-        dPrimV_ global_tid $ le32 group_id * le32 inner_group_size + le32 local_tid+        sOp (Imp.GetBlockId tblock_id 0)+        dPrimV_ global_tid $ le32 tblock_id * le32 inner_tblock_size + le32 local_tid    pure (constants, set_constants) @@ -1010,23 +1009,23 @@   localDefaultSpace (Imp.Space "global") . localVTable (M.map globalMemory)   where     globalMemory (MemVar _ entry)-      | entryMemSpace entry /= Space "local" =+      | entryMemSpace entry /= Space "shared" =           MemVar Nothing entry {entryMemSpace = Imp.Space "global"}     globalMemory entry =       entry -simpleKernelGroups ::+simpleKernelBlocks ::   Imp.TExp Int64 ->   Imp.TExp Int64 ->-  CallKernelGen (Imp.TExp Int32, Count NumGroups SubExp, Count GroupSize SubExp)-simpleKernelGroups max_num_groups kernel_size = do-  group_size <- dPrim "group_size" int64+  CallKernelGen (Imp.TExp Int32, Count NumBlocks SubExp, Count BlockSize SubExp)+simpleKernelBlocks max_num_tblocks kernel_size = do+  tblock_size <- dPrim "tblock_size" int64   fname <- askFunction-  let group_size_key = keyWithEntryPoint fname $ nameFromString $ prettyString $ tvVar group_size-  sOp $ Imp.GetSize (tvVar group_size) group_size_key Imp.SizeGroup-  virt_num_groups <- dPrimVE "virt_num_groups" $ kernel_size `divUp` tvExp group_size-  num_groups <- dPrimV "num_groups" $ virt_num_groups `sMin64` max_num_groups-  pure (sExt32 virt_num_groups, Count $ tvSize num_groups, Count $ tvSize group_size)+  let tblock_size_key = keyWithEntryPoint fname $ nameFromString $ prettyString $ tvVar tblock_size+  sOp $ Imp.GetSize (tvVar tblock_size) tblock_size_key Imp.SizeThreadBlock+  virt_num_tblocks <- dPrimVE "virt_num_tblocks" $ kernel_size `divUp` tvExp tblock_size+  num_tblocks <- dPrimV "num_tblocks" $ virt_num_tblocks `sMin64` max_num_tblocks+  pure (sExt32 virt_num_tblocks, Count $ tvSize num_tblocks, Count $ tvSize tblock_size)  simpleKernelConstants ::   Imp.TExp Int64 ->@@ -1038,32 +1037,32 @@ simpleKernelConstants kernel_size desc = do   -- For performance reasons, codegen assumes that the thread count is   -- never more than will fit in an i32.  This means we need to cap-  -- the number of groups here.  The cap is set much higher than any+  -- the number of blocks here.  The cap is set much higher than any   -- GPU will possibly need.  Feel free to come back and laugh at me   -- in the future.-  let max_num_groups = 1024 * 1024+  let max_num_tblocks = 1024 * 1024   thread_gtid <- newVName $ desc ++ "_gtid"   thread_ltid <- newVName $ desc ++ "_ltid"-  group_id <- newVName $ desc ++ "_gid"-  inner_group_size <- newVName "group_size"-  (virt_num_groups, num_groups, group_size) <--    simpleKernelGroups max_num_groups kernel_size-  let group_size' = Imp.pe64 $ unCount group_size-      num_groups' = Imp.pe64 $ unCount num_groups+  tblock_id <- newVName $ desc ++ "_gid"+  inner_tblock_size <- newVName "tblock_size"+  (virt_num_tblocks, num_tblocks, tblock_size) <-+    simpleKernelBlocks max_num_tblocks kernel_size+  let tblock_size' = Imp.pe64 $ unCount tblock_size+      num_tblocks' = Imp.pe64 $ unCount num_tblocks        constants =         KernelConstants           { kernelGlobalThreadId = Imp.le32 thread_gtid,             kernelLocalThreadId = Imp.le32 thread_ltid,-            kernelGroupId = Imp.le32 group_id,+            kernelBlockId = Imp.le32 tblock_id,             kernelGlobalThreadIdVar = thread_gtid,             kernelLocalThreadIdVar = thread_ltid,-            kernelGroupIdVar = group_id,-            kernelNumGroupsCount = num_groups,-            kernelGroupSizeCount = group_size,-            kernelNumGroups = num_groups',-            kernelGroupSize = group_size',-            kernelNumThreads = sExt32 (group_size' * num_groups'),+            kernelBlockIdVar = tblock_id,+            kernelNumBlocksCount = num_tblocks,+            kernelBlockSizeCount = tblock_size,+            kernelNumBlocks = num_tblocks',+            kernelBlockSize = tblock_size',+            kernelNumThreads = sExt32 (tblock_size' * num_tblocks'),             kernelWaveSize = 0,             kernelLocalIdMap = mempty,             kernelChunkItersMap = mempty@@ -1071,51 +1070,51 @@        wrapKernel m = do         dPrim_ thread_ltid int32-        dPrim_ inner_group_size int64-        dPrim_ group_id int32+        dPrim_ inner_tblock_size int64+        dPrim_ tblock_id int32         sOp (Imp.GetLocalId thread_ltid 0)-        sOp (Imp.GetLocalSize inner_group_size 0)-        sOp (Imp.GetGroupId group_id 0)-        dPrimV_ thread_gtid $ le32 group_id * le32 inner_group_size + le32 thread_ltid-        virtualiseGroups SegVirt virt_num_groups $ \virt_group_id -> do+        sOp (Imp.GetLocalSize inner_tblock_size 0)+        sOp (Imp.GetBlockId tblock_id 0)+        dPrimV_ thread_gtid $ le32 tblock_id * le32 inner_tblock_size + le32 thread_ltid+        virtualiseBlocks SegVirt virt_num_tblocks $ \virt_tblock_id -> do           global_tid <-             dPrimVE "global_tid" $-              sExt64 virt_group_id * sExt64 (le32 inner_group_size)+              sExt64 virt_tblock_id * sExt64 (le32 inner_tblock_size)                 + sExt64 (kernelLocalThreadId constants)           m global_tid    pure (wrapKernel, constants) --- | For many kernels, we may not have enough physical groups to cover--- the logical iteration space.  Some groups thus have to perform+-- | For many kernels, we may not have enough physical blocks to cover+-- the logical iteration space.  Some blocks thus have to perform -- double duty; we put an outer loop to accomplish this.  The -- advantage over just launching a bazillion threads is that the cost -- of memory expansion should be proportional to the number of -- *physical* threads (hardware parallelism), not the amount of -- application parallelism.-virtualiseGroups ::+virtualiseBlocks ::   SegVirt ->   Imp.TExp Int32 ->   (Imp.TExp Int32 -> InKernelGen ()) ->   InKernelGen ()-virtualiseGroups SegVirt required_groups m = do+virtualiseBlocks SegVirt required_blocks m = do   constants <- kernelConstants <$> askEnv-  phys_group_id <- dPrim "phys_group_id" int32-  sOp $ Imp.GetGroupId (tvVar phys_group_id) 0+  phys_tblock_id <- dPrim "phys_tblock_id" int32+  sOp $ Imp.GetBlockId (tvVar phys_tblock_id) 0   iterations <-     dPrimVE "iterations" $-      (required_groups - tvExp phys_group_id) `divUp` sExt32 (kernelNumGroups constants)+      (required_blocks - tvExp phys_tblock_id) `divUp` sExt32 (kernelNumBlocks constants)    sFor "i" iterations $ \i -> do     m . tvExp       =<< dPrimV-        "virt_group_id"-        (tvExp phys_group_id + i * sExt32 (kernelNumGroups constants))-    -- Make sure the virtual group is actually done before we let-    -- another virtual group have its way with it.+        "virt_tblock_id"+        (tvExp phys_tblock_id + i * sExt32 (kernelNumBlocks constants))+    -- Make sure the virtual block is actually done before we let+    -- another virtual block have its way with it.     sOp $ Imp.ErrorSync Imp.FenceGlobal-virtualiseGroups _ _ m = do-  gid <- kernelGroupIdVar . kernelConstants <$> askEnv+virtualiseBlocks _ _ m = do+  gid <- kernelBlockIdVar . kernelConstants <$> askEnv   m $ Imp.le32 gid  -- | Various extra configuration of the kernel being generated.@@ -1123,11 +1122,11 @@   { -- | Can this kernel execute correctly even if previous kernels failed?     kAttrFailureTolerant :: Bool,     -- | Does whatever launch this kernel check for local memory capacity itself?-    kAttrCheckLocalMemory :: Bool,-    -- | Number of groups.-    kAttrNumGroups :: Count NumGroups SubExp,-    -- | Group size.-    kAttrGroupSize :: Count GroupSize SubExp,+    kAttrCheckSharedMemory :: Bool,+    -- | Number of blocks.+    kAttrNumBlocks :: Count NumBlocks SubExp,+    -- | Block size.+    kAttrBlockSize :: Count BlockSize SubExp,     -- | Variables that are specially in scope inside the kernel.     -- Operationally, these will be available at kernel compile time     -- (which happens at run-time, with access to machine-specific@@ -1137,15 +1136,15 @@  -- | The default kernel attributes. defKernelAttrs ::-  Count NumGroups SubExp ->-  Count GroupSize SubExp ->+  Count NumBlocks SubExp ->+  Count BlockSize SubExp ->   KernelAttrs-defKernelAttrs num_groups group_size =+defKernelAttrs num_tblocks tblock_size =   KernelAttrs     { kAttrFailureTolerant = False,-      kAttrCheckLocalMemory = True,-      kAttrNumGroups = num_groups,-      kAttrGroupSize = group_size,+      kAttrCheckSharedMemory = True,+      kAttrNumBlocks = num_tblocks,+      kAttrBlockSize = tblock_size,       kAttrConstExps = mempty     } @@ -1158,23 +1157,23 @@   pure v  -- | Compute kernel attributes from 'SegLevel'; including synthesising--- group-size and thread count if no grid is provided.+-- block-size and thread count if no grid is provided. lvlKernelAttrs :: SegLevel -> CallKernelGen KernelAttrs lvlKernelAttrs lvl =   case lvl of     SegThread _ Nothing -> mkGrid-    SegThread _ (Just (KernelGrid num_groups group_size)) ->-      pure $ defKernelAttrs num_groups group_size-    SegGroup _ Nothing -> mkGrid-    SegGroup _ (Just (KernelGrid num_groups group_size)) ->-      pure $ defKernelAttrs num_groups group_size-    SegThreadInGroup {} ->-      error "lvlKernelAttrs: SegThreadInGroup"+    SegThread _ (Just (KernelGrid num_tblocks tblock_size)) ->+      pure $ defKernelAttrs num_tblocks tblock_size+    SegBlock _ Nothing -> mkGrid+    SegBlock _ (Just (KernelGrid num_tblocks tblock_size)) ->+      pure $ defKernelAttrs num_tblocks tblock_size+    SegThreadInBlock {} ->+      error "lvlKernelAttrs: SegThreadInBlock"   where     mkGrid = do-      group_size <- getSize "group_size" Imp.SizeGroup-      num_groups <- getSize "num_groups" Imp.SizeNumGroups-      pure $ defKernelAttrs (Count $ tvSize num_groups) (Count $ tvSize group_size)+      tblock_size <- getSize "tblock_size" Imp.SizeThreadBlock+      num_tblocks <- getSize "num_tblocks" Imp.SizeGrid+      pure $ defKernelAttrs (Count $ tvSize num_tblocks) (Count $ tvSize tblock_size)  sKernel ::   Operations GPUMem KernelEnv Imp.KernelOp ->@@ -1186,7 +1185,7 @@   CallKernelGen () sKernel ops flatf name v attrs f = do   (constants, set_constants) <--    kernelInitialisationSimple (kAttrNumGroups attrs) (kAttrGroupSize attrs)+    kernelInitialisationSimple (kAttrNumBlocks attrs) (kAttrBlockSize attrs)   name' <- nameForFun $ name ++ "_" ++ show (baseTag v)   sKernelOp attrs constants ops name' $ do     set_constants@@ -1212,21 +1211,21 @@   HostEnv atomics _ locks <- askEnv   body <- makeAllMemoryGlobal $ subImpM_ (KernelEnv atomics constants locks) ops m   uses <- computeKernelUses body $ M.keys $ kAttrConstExps attrs-  group_size <- onGroupSize $ kernelGroupSize constants+  tblock_size <- onBlockSize $ kernelBlockSize constants   emit . Imp.Op . Imp.CallKernel $     Imp.Kernel       { Imp.kernelBody = body,         Imp.kernelUses = uses <> map constToUse (M.toList (kAttrConstExps attrs)),-        Imp.kernelNumGroups = [untyped $ kernelNumGroups constants],-        Imp.kernelGroupSize = [group_size],+        Imp.kernelNumBlocks = [untyped $ kernelNumBlocks constants],+        Imp.kernelBlockSize = [tblock_size],         Imp.kernelName = name,         Imp.kernelFailureTolerant = kAttrFailureTolerant attrs,-        Imp.kernelCheckLocalMemory = kAttrCheckLocalMemory attrs+        Imp.kernelCheckSharedMemory = kAttrCheckSharedMemory attrs       }   where     -- Figure out if this expression actually corresponds to a     -- KernelConst.-    onGroupSize e = do+    onBlockSize e = do       vtable <- getVTable       x <- isConstExp vtable $ untyped e       pure $@@ -1248,8 +1247,8 @@   where     attrs =       ( defKernelAttrs-          (kernelNumGroupsCount constants)-          (kernelGroupSizeCount constants)+          (kernelNumBlocksCount constants)+          (kernelBlockSizeCount constants)       )         { kAttrFailureTolerant = tol         }@@ -1261,7 +1260,7 @@       opsExpCompiler = compileThreadExp,       opsStmsCompiler = \_ -> defCompileStms mempty,       opsAllocCompilers =-        M.fromList [(Space "local", allocLocal)]+        M.fromList [(Space "shared", allocLocal)]     }  -- | Perform a Replicate with a kernel.

@@ -0,0 +1,743 @@+{-# LANGUAGE TypeFamilies #-}++-- | Generation of kernels with block-level bodies.+module Futhark.CodeGen.ImpGen.GPU.Block+  ( sKernelBlock,+    compileBlockResult,+    blockOperations,++    -- * Precomputation+    Precomputed,+    precomputeConstants,+    precomputedConstants,+    atomicUpdateLocking,+  )+where++import Control.Monad+import Data.Bifunctor+import Data.List (partition, zip4)+import Data.Map.Strict qualified as M+import Data.Maybe+import Data.Set qualified as S+import Futhark.CodeGen.ImpCode.GPU qualified as Imp+import Futhark.CodeGen.ImpGen+import Futhark.CodeGen.ImpGen.GPU.Base+import Futhark.Construct (fullSliceNum)+import Futhark.Error+import Futhark.IR.GPUMem+import Futhark.IR.Mem.LMAD qualified as LMAD+import Futhark.Transform.Rename+import Futhark.Util (chunks, mapAccumLM, takeLast)+import Futhark.Util.IntegralExp (divUp, rem)+import Prelude hiding (quot, rem)++-- | @flattenArray k flat arr@ flattens the outer @k@ dimensions of+-- @arr@ to @flat@.  (Make sure @flat@ is the sum of those dimensions+-- or you'll have a bad time.)+flattenArray :: Int -> TV Int64 -> VName -> ImpM rep r op VName+flattenArray k flat arr = do+  ArrayEntry arr_loc pt <- lookupArray arr+  let flat_shape = Shape $ Var (tvVar flat) : drop k (memLocShape arr_loc)+  sArray (baseString arr ++ "_flat") pt flat_shape (memLocName arr_loc) $+    fromMaybe (error "flattenArray") $+      LMAD.reshape (memLocLMAD arr_loc) (map pe64 $ shapeDims flat_shape)++sliceArray :: Imp.TExp Int64 -> TV Int64 -> VName -> ImpM rep r op VName+sliceArray start size arr = do+  MemLoc mem _ ixfun <- entryArrayLoc <$> lookupArray arr+  arr_t <- lookupType arr+  let slice =+        fullSliceNum+          (map Imp.pe64 (arrayDims arr_t))+          [DimSlice start (tvExp size) 1]+  sArray+    (baseString arr ++ "_chunk")+    (elemType arr_t)+    (arrayShape arr_t `setOuterDim` Var (tvVar size))+    mem+    $ LMAD.slice ixfun slice++-- | @applyLambda lam dests args@ emits code that:+--+-- 1. Binds each parameter of @lam@ to the corresponding element of+--    @args@, interpreted as a (name,slice) pair (as in 'copyDWIM').+--    Use an empty list for a scalar.+--+-- 2. Executes the body of @lam@.+--+-- 3. Binds the t'SubExp's that are the 'Result' of @lam@ to the+-- provided @dest@s, again interpreted as the destination for a+-- 'copyDWIM'.+applyLambda ::+  (Mem rep inner) =>+  Lambda rep ->+  [(VName, [DimIndex (Imp.TExp Int64)])] ->+  [(SubExp, [DimIndex (Imp.TExp Int64)])] ->+  ImpM rep r op ()+applyLambda lam dests args = do+  dLParams $ lambdaParams lam+  forM_ (zip (lambdaParams lam) args) $ \(p, (arg, arg_slice)) ->+    copyDWIM (paramName p) [] arg arg_slice+  compileStms mempty (bodyStms $ lambdaBody lam) $ do+    let res = map resSubExp $ bodyResult $ lambdaBody lam+    forM_ (zip dests res) $ \((dest, dest_slice), se) ->+      copyDWIM dest dest_slice se []++-- | As applyLambda, but first rename the names in the lambda.  This+-- makes it safe to apply it in multiple places.  (It might be safe+-- anyway, but you have to be more careful - use this if you are in+-- doubt.)+applyRenamedLambda ::+  (Mem rep inner) =>+  Lambda rep ->+  [(VName, [DimIndex (Imp.TExp Int64)])] ->+  [(SubExp, [DimIndex (Imp.TExp Int64)])] ->+  ImpM rep r op ()+applyRenamedLambda lam dests args = do+  lam_renamed <- renameLambda lam+  applyLambda lam_renamed dests args++blockChunkLoop ::+  Imp.TExp Int32 ->+  (Imp.TExp Int32 -> TV Int64 -> InKernelGen ()) ->+  InKernelGen ()+blockChunkLoop w m = do+  constants <- kernelConstants <$> askEnv+  let max_chunk_size = sExt32 $ kernelBlockSize constants+  num_chunks <- dPrimVE "num_chunks" $ w `divUp` max_chunk_size+  sFor "chunk_i" num_chunks $ \chunk_i -> do+    chunk_start <-+      dPrimVE "chunk_start" $ chunk_i * max_chunk_size+    chunk_end <-+      dPrimVE "chunk_end" $ sMin32 w (chunk_start + max_chunk_size)+    chunk_size <-+      dPrimV "chunk_size" $ sExt64 $ chunk_end - chunk_start+    m chunk_start chunk_size++virtualisedBlockScan ::+  Maybe (Imp.TExp Int32 -> Imp.TExp Int32 -> Imp.TExp Bool) ->+  Imp.TExp Int32 ->+  Lambda GPUMem ->+  [VName] ->+  InKernelGen ()+virtualisedBlockScan seg_flag w lam arrs = do+  blockChunkLoop w $ \chunk_start chunk_size -> do+    constants <- kernelConstants <$> askEnv+    let ltid = kernelLocalThreadId constants+        crosses_segment =+          case seg_flag of+            Nothing -> false+            Just flag_true ->+              flag_true (sExt32 (chunk_start - 1)) (sExt32 chunk_start)+    sComment "possibly incorporate carry" $+      sWhen (chunk_start .>. 0 .&&. ltid .==. 0 .&&. bNot crosses_segment) $ do+        carry_idx <- dPrimVE "carry_idx" $ sExt64 chunk_start - 1+        applyRenamedLambda+          lam+          (map (,[DimFix $ sExt64 chunk_start]) arrs)+          ( map ((,[DimFix carry_idx]) . Var) arrs+              ++ map ((,[DimFix $ sExt64 chunk_start]) . Var) arrs+          )++    arrs_chunks <- mapM (sliceArray (sExt64 chunk_start) chunk_size) arrs++    sOp $ Imp.ErrorSync Imp.FenceLocal++    blockScan seg_flag (sExt64 w) (tvExp chunk_size) lam arrs_chunks++copyInBlock :: CopyCompiler GPUMem KernelEnv Imp.KernelOp+copyInBlock pt destloc srcloc = do+  dest_space <- entryMemSpace <$> lookupMemory (memLocName destloc)+  src_space <- entryMemSpace <$> lookupMemory (memLocName srcloc)++  let src_lmad = memLocLMAD srcloc+      dims = LMAD.shape src_lmad+      rank = length dims++  case (dest_space, src_space) of+    (ScalarSpace destds _, ScalarSpace srcds _) -> do+      let fullDim d = DimSlice 0 d 1+          destslice' =+            Slice $+              replicate (rank - length destds) (DimFix 0)+                ++ takeLast (length destds) (map fullDim dims)+          srcslice' =+            Slice $+              replicate (rank - length srcds) (DimFix 0)+                ++ takeLast (length srcds) (map fullDim dims)+      lmadCopy+        pt+        (sliceMemLoc destloc destslice')+        (sliceMemLoc srcloc srcslice')+    _ -> do+      blockCoverSpace (map sExt32 dims) $ \is ->+        lmadCopy+          pt+          (sliceMemLoc destloc (Slice $ map (DimFix . sExt64) is))+          (sliceMemLoc srcloc (Slice $ map (DimFix . sExt64) is))+      sOp $ Imp.Barrier Imp.FenceLocal++localThreadIDs :: [SubExp] -> InKernelGen [Imp.TExp Int64]+localThreadIDs dims = do+  ltid <- sExt64 . kernelLocalThreadId . kernelConstants <$> askEnv+  let dims' = map pe64 dims+  maybe (dIndexSpace' "ltid" dims' ltid) (pure . map sExt64)+    . M.lookup dims+    . kernelLocalIdMap+    . kernelConstants+    =<< askEnv++partitionSeqDims :: SegSeqDims -> SegSpace -> ([(VName, SubExp)], [(VName, SubExp)])+partitionSeqDims (SegSeqDims seq_is) space =+  bimap (map fst) (map fst) $+    partition ((`elem` seq_is) . snd) (zip (unSegSpace space) [0 ..])++compileFlatId :: SegSpace -> InKernelGen ()+compileFlatId space = do+  ltid <- kernelLocalThreadId . kernelConstants <$> askEnv+  dPrimV_ (segFlat space) ltid++-- Construct the necessary lock arrays for an intra-block histogram.+prepareIntraBlockSegHist ::+  Shape ->+  Count BlockSize SubExp ->+  [HistOp GPUMem] ->+  InKernelGen [[Imp.TExp Int64] -> InKernelGen ()]+prepareIntraBlockSegHist segments tblock_size =+  fmap snd . mapAccumLM onOp Nothing+  where+    onOp l op = do+      constants <- kernelConstants <$> askEnv+      atomicBinOp <- kernelAtomics <$> askEnv++      let local_subhistos = histDest op++      case (l, atomicUpdateLocking atomicBinOp $ histOp op) of+        (_, AtomicPrim f) -> pure (l, f (Space "shared") local_subhistos)+        (_, AtomicCAS f) -> pure (l, f (Space "shared") local_subhistos)+        (Just l', AtomicLocking f) -> pure (l, f l' (Space "shared") local_subhistos)+        (Nothing, AtomicLocking f) -> do+          locks <- newVName "locks"++          let num_locks = pe64 $ unCount tblock_size+              dims = map pe64 $ shapeDims (segments <> histOpShape op <> histShape op)+              l' = Locking locks 0 1 0 (pure . (`rem` num_locks) . flattenIndex dims)+              locks_t = Array int32 (Shape [unCount tblock_size]) NoUniqueness++          locks_mem <- sAlloc "locks_mem" (typeSize locks_t) $ Space "shared"+          dArray locks int32 (arrayShape locks_t) locks_mem $+            LMAD.iota 0 . map pe64 . arrayDims $+              locks_t++          sComment "All locks start out unlocked" $+            blockCoverSpace [kernelBlockSize constants] $ \is ->+              copyDWIMFix locks is (intConst Int32 0) []++          pure (Just l', f l' (Space "shared") local_subhistos)++blockCoverSegSpace :: SegVirt -> SegSpace -> InKernelGen () -> InKernelGen ()+blockCoverSegSpace virt space m = do+  let (ltids, dims) = unzip $ unSegSpace space+      dims' = map pe64 dims++  constants <- kernelConstants <$> askEnv+  let tblock_size = kernelBlockSize constants+  -- Maybe we can statically detect that this is actually a+  -- SegNoVirtFull and generate ever-so-slightly simpler code.+  let virt' = if dims' == [tblock_size] then SegNoVirtFull (SegSeqDims []) else virt+  case virt' of+    SegVirt -> do+      iters <- M.lookup dims . kernelChunkItersMap . kernelConstants <$> askEnv+      case iters of+        Nothing -> do+          iterations <- dPrimVE "iterations" $ product $ map sExt32 dims'+          blockLoop iterations $ \i -> do+            dIndexSpace (zip ltids dims') $ sExt64 i+            m+        Just num_chunks -> localOps threadOperations $ do+          let ltid = kernelLocalThreadId constants+          sFor "chunk_i" num_chunks $ \chunk_i -> do+            i <- dPrimVE "i" $ chunk_i * sExt32 tblock_size + ltid+            dIndexSpace (zip ltids dims') $ sExt64 i+            sWhen (inBounds (Slice (map (DimFix . le64) ltids)) dims') m+    SegNoVirt -> localOps threadOperations $ do+      zipWithM_ dPrimV_ ltids =<< localThreadIDs dims+      sWhen (isActive $ zip ltids dims) m+    SegNoVirtFull seq_dims -> do+      let ((ltids_seq, dims_seq), (ltids_par, dims_par)) =+            bimap unzip unzip $ partitionSeqDims seq_dims space+      sLoopNest (Shape dims_seq) $ \is_seq -> do+        zipWithM_ dPrimV_ ltids_seq is_seq+        localOps threadOperations $ do+          zipWithM_ dPrimV_ ltids_par =<< localThreadIDs dims_par+          m++compileBlockExp :: ExpCompiler GPUMem KernelEnv Imp.KernelOp+compileBlockExp (Pat [pe]) (BasicOp (Opaque _ se)) =+  -- Cannot print in GPU code.+  copyDWIM (patElemName pe) [] se []+-- The static arrays stuff does not work inside kernels.+compileBlockExp (Pat [dest]) (BasicOp (ArrayLit es _)) =+  forM_ (zip [0 ..] es) $ \(i, e) ->+    copyDWIMFix (patElemName dest) [fromIntegral (i :: Int64)] e []+compileBlockExp _ (BasicOp (UpdateAcc acc is vs)) = do+  ltid <- kernelLocalThreadId . kernelConstants <$> askEnv+  sWhen (ltid .==. 0) $ updateAcc acc is vs+  sOp $ Imp.Barrier Imp.FenceLocal+compileBlockExp (Pat [dest]) (BasicOp (Replicate ds se)) | ds /= mempty = do+  flat <- newVName "rep_flat"+  is <- replicateM (arrayRank dest_t) (newVName "rep_i")+  let is' = map le64 is+  blockCoverSegSpace SegVirt (SegSpace flat $ zip is $ arrayDims dest_t) $+    copyDWIMFix (patElemName dest) is' se (drop (shapeRank ds) is')+  sOp $ Imp.Barrier Imp.FenceLocal+  where+    dest_t = patElemType dest+compileBlockExp (Pat [dest]) (BasicOp (Iota n e s it)) = do+  n' <- toExp n+  e' <- toExp e+  s' <- toExp s+  blockLoop (TPrimExp n') $ \i' -> do+    x <-+      dPrimV "x" $+        TPrimExp $+          BinOpExp (Add it OverflowUndef) e' $+            BinOpExp (Mul it OverflowUndef) (untyped i') s'+    copyDWIMFix (patElemName dest) [i'] (Var (tvVar x)) []+  sOp $ Imp.Barrier Imp.FenceLocal++-- When generating code for a scalar in-place update, we must make+-- sure that only one thread performs the write.  When writing an+-- array, the block-level copy code will take care of doing the right+-- thing.+compileBlockExp (Pat [pe]) (BasicOp (Update safety _ slice se))+  | null $ sliceDims slice = do+      sOp $ Imp.Barrier Imp.FenceLocal+      ltid <- kernelLocalThreadId . kernelConstants <$> askEnv+      sWhen (ltid .==. 0) $+        case safety of+          Unsafe -> write+          Safe -> sWhen (inBounds slice' dims) write+      sOp $ Imp.Barrier Imp.FenceLocal+  where+    slice' = fmap pe64 slice+    dims = map pe64 $ arrayDims $ patElemType pe+    write = copyDWIM (patElemName pe) (unSlice slice') se []+compileBlockExp dest e = do+  -- It is a messy to jump into control flow for error handling.+  -- Avoid that by always doing an error sync here.  Potential+  -- improvement: only do this if any errors are pending (this could+  -- also be handled in later codegen).+  when (doSync e) $ sOp $ Imp.ErrorSync Imp.FenceLocal+  defCompileExp dest e+  where+    doSync Loop {} = True+    doSync Match {} = True+    doSync _ = False++compileBlockOp :: OpCompiler GPUMem KernelEnv Imp.KernelOp+compileBlockOp pat (Alloc size space) =+  kernelAlloc pat size space+compileBlockOp pat (Inner (SegOp (SegMap lvl space _ body))) = do+  compileFlatId space++  blockCoverSegSpace (segVirt lvl) space $+    compileStms mempty (kernelBodyStms body) $+      zipWithM_ (compileThreadResult space) (patElems pat) $+        kernelBodyResult body+  sOp $ Imp.ErrorSync Imp.FenceLocal+compileBlockOp pat (Inner (SegOp (SegScan lvl space scans _ body))) = do+  compileFlatId space++  let (ltids, dims) = unzip $ unSegSpace space+      dims' = map pe64 dims++  blockCoverSegSpace (segVirt lvl) space $+    compileStms mempty (kernelBodyStms body) $+      forM_ (zip (patNames pat) $ kernelBodyResult body) $ \(dest, res) ->+        copyDWIMFix+          dest+          (map Imp.le64 ltids)+          (kernelResultSubExp res)+          []++  fence <- fenceForArrays $ patNames pat+  sOp $ Imp.ErrorSync fence++  let segment_size = last dims'+      crossesSegment from to =+        (sExt64 to - sExt64 from) .>. (sExt64 to `rem` segment_size)++  -- blockScan needs to treat the scan output as a one-dimensional+  -- array of scan elements, so we invent some new flattened arrays+  -- here.+  dims_flat <- dPrimV "dims_flat" $ product dims'+  let scan = head scans+      num_scan_results = length $ segBinOpNeutral scan+  arrs_flat <-+    mapM (flattenArray (length dims') dims_flat) $+      take num_scan_results $+        patNames pat++  case segVirt lvl of+    SegVirt ->+      virtualisedBlockScan+        (Just crossesSegment)+        (sExt32 $ tvExp dims_flat)+        (segBinOpLambda scan)+        arrs_flat+    _ ->+      blockScan+        (Just crossesSegment)+        (product dims')+        (product dims')+        (segBinOpLambda scan)+        arrs_flat+compileBlockOp pat (Inner (SegOp (SegRed lvl space ops _ body))) = do+  compileFlatId space++  let dims' = map pe64 dims+      mkTempArr t =+        sAllocArray "red_arr" (elemType t) (Shape dims <> arrayShape t) $ Space "shared"++  tmp_arrs <- mapM mkTempArr $ concatMap (lambdaReturnType . segBinOpLambda) ops+  blockCoverSegSpace (segVirt lvl) space $+    compileStms mempty (kernelBodyStms body) $ do+      let (red_res, map_res) =+            splitAt (segBinOpResults ops) $ kernelBodyResult body+      forM_ (zip tmp_arrs red_res) $ \(dest, res) ->+        copyDWIMFix dest (map Imp.le64 ltids) (kernelResultSubExp res) []+      zipWithM_ (compileThreadResult space) map_pes map_res++  sOp $ Imp.ErrorSync Imp.FenceLocal++  let tmps_for_ops = chunks (map (length . segBinOpNeutral) ops) tmp_arrs+  case segVirt lvl of+    SegVirt -> virtCase dims' tmps_for_ops+    _ -> nonvirtCase dims' tmps_for_ops+  where+    (ltids, dims) = unzip $ unSegSpace space+    (red_pes, map_pes) = splitAt (segBinOpResults ops) $ patElems pat++    virtCase [dim'] tmps_for_ops = do+      ltid <- kernelLocalThreadId . kernelConstants <$> askEnv+      blockChunkLoop (sExt32 dim') $ \chunk_start chunk_size -> do+        sComment "possibly incorporate carry" $+          sWhen (chunk_start .>. 0 .&&. ltid .==. 0) $+            forM_ (zip ops tmps_for_ops) $ \(op, tmps) ->+              applyRenamedLambda+                (segBinOpLambda op)+                (map (,[DimFix $ sExt64 chunk_start]) tmps)+                ( map ((,[]) . Var . patElemName) red_pes+                    ++ map ((,[DimFix $ sExt64 chunk_start]) . Var) tmps+                )++        sOp $ Imp.ErrorSync Imp.FenceLocal++        forM_ (zip ops tmps_for_ops) $ \(op, tmps) -> do+          tmps_chunks <- mapM (sliceArray (sExt64 chunk_start) chunk_size) tmps+          blockReduce (sExt32 (tvExp chunk_size)) (segBinOpLambda op) tmps_chunks++        sOp $ Imp.ErrorSync Imp.FenceLocal++        sComment "Save result of reduction." $+          forM_ (zip red_pes $ concat tmps_for_ops) $ \(pe, arr) ->+            copyDWIMFix (patElemName pe) [] (Var arr) [sExt64 chunk_start]++    --+    virtCase dims' tmps_for_ops = do+      dims_flat <- dPrimV "dims_flat" $ product dims'+      let segment_size = last dims'+          crossesSegment from to =+            (sExt64 to - sExt64 from) .>. (sExt64 to `rem` sExt64 segment_size)++      forM_ (zip ops tmps_for_ops) $ \(op, tmps) -> do+        tmps_flat <- mapM (flattenArray (length dims') dims_flat) tmps+        virtualisedBlockScan+          (Just crossesSegment)+          (sExt32 $ tvExp dims_flat)+          (segBinOpLambda op)+          tmps_flat++      sOp $ Imp.ErrorSync Imp.FenceLocal++      sComment "Save result of reduction." $+        forM_ (zip red_pes $ concat tmps_for_ops) $ \(pe, arr) ->+          copyDWIM+            (patElemName pe)+            []+            (Var arr)+            (map (unitSlice 0) (init dims') ++ [DimFix $ last dims' - 1])++      sOp $ Imp.Barrier Imp.FenceLocal++    -- Nonsegmented case (or rather, a single segment) - this we can+    -- handle directly with a block-level reduction.+    nonvirtCase [dim'] tmps_for_ops = do+      forM_ (zip ops tmps_for_ops) $ \(op, tmps) ->+        blockReduce (sExt32 dim') (segBinOpLambda op) tmps+      sOp $ Imp.ErrorSync Imp.FenceLocal+      sComment "Save result of reduction." $+        forM_ (zip red_pes $ concat tmps_for_ops) $ \(pe, arr) ->+          copyDWIMFix (patElemName pe) [] (Var arr) [0]+      sOp $ Imp.ErrorSync Imp.FenceLocal++    -- Segmented intra-block reductions are turned into (regular)+    -- segmented scans.  It is possible that this can be done+    -- better, but at least this approach is simple.+    nonvirtCase dims' tmps_for_ops = do+      -- blockScan operates on flattened arrays.  This does not+      -- involve copying anything; merely playing with the index+      -- function.+      dims_flat <- dPrimV "dims_flat" $ product dims'+      let segment_size = last dims'+          crossesSegment from to =+            (sExt64 to - sExt64 from) .>. (sExt64 to `rem` sExt64 segment_size)++      forM_ (zip ops tmps_for_ops) $ \(op, tmps) -> do+        tmps_flat <- mapM (flattenArray (length dims') dims_flat) tmps+        blockScan+          (Just crossesSegment)+          (product dims')+          (product dims')+          (segBinOpLambda op)+          tmps_flat++      sOp $ Imp.ErrorSync Imp.FenceLocal++      sComment "Save result of reduction." $+        forM_ (zip red_pes $ concat tmps_for_ops) $ \(pe, arr) ->+          copyDWIM+            (patElemName pe)+            []+            (Var arr)+            (map (unitSlice 0) (init dims') ++ [DimFix $ last dims' - 1])++      sOp $ Imp.Barrier Imp.FenceLocal+compileBlockOp pat (Inner (SegOp (SegHist lvl space ops _ kbody))) = do+  compileFlatId space+  let (ltids, dims) = unzip $ unSegSpace space++  -- We don't need the red_pes, because it is guaranteed by our type+  -- rules that they occupy the same memory as the destinations for+  -- the ops.+  let num_red_res = length ops + sum (map (length . histNeutral) ops)+      (_red_pes, map_pes) =+        splitAt num_red_res $ patElems pat++  tblock_size <- kernelBlockSizeCount . kernelConstants <$> askEnv+  ops' <- prepareIntraBlockSegHist (Shape $ init dims) tblock_size ops++  -- Ensure that all locks have been initialised.+  sOp $ Imp.Barrier Imp.FenceLocal++  blockCoverSegSpace (segVirt lvl) space $+    compileStms mempty (kernelBodyStms kbody) $ do+      let (red_res, map_res) = splitAt num_red_res $ kernelBodyResult kbody+          (red_is, red_vs) = splitAt (length ops) $ map kernelResultSubExp red_res+      zipWithM_ (compileThreadResult space) map_pes map_res++      let vs_per_op = chunks (map (length . histDest) ops) red_vs++      forM_ (zip4 red_is vs_per_op ops' ops) $+        \(bin, op_vs, do_op, HistOp dest_shape _ _ _ shape lam) -> do+          let bin' = pe64 bin+              dest_shape' = map pe64 $ shapeDims dest_shape+              bin_in_bounds = inBounds (Slice [DimFix bin']) dest_shape'+              bin_is = map Imp.le64 (init ltids) ++ [bin']+              vs_params = takeLast (length op_vs) $ lambdaParams lam++          sComment "perform atomic updates" $+            sWhen bin_in_bounds $ do+              dLParams $ lambdaParams lam+              sLoopNest shape $ \is -> do+                forM_ (zip vs_params op_vs) $ \(p, v) ->+                  copyDWIMFix (paramName p) [] v is+                do_op (bin_is ++ is)++  sOp $ Imp.ErrorSync Imp.FenceLocal+compileBlockOp pat _ =+  compilerBugS $ "compileBlockOp: cannot compile rhs of binding " ++ prettyString pat++blockOperations :: Operations GPUMem KernelEnv Imp.KernelOp+blockOperations =+  (defaultOperations compileBlockOp)+    { opsCopyCompiler = copyInBlock,+      opsExpCompiler = compileBlockExp,+      opsStmsCompiler = \_ -> defCompileStms mempty,+      opsAllocCompilers =+        M.fromList [(Space "shared", allocLocal)]+    }++arrayInSharedMemory :: SubExp -> InKernelGen Bool+arrayInSharedMemory (Var name) = do+  res <- lookupVar name+  case res of+    ArrayVar _ entry ->+      (Space "shared" ==) . entryMemSpace+        <$> lookupMemory (memLocName (entryArrayLoc entry))+    _ -> pure False+arrayInSharedMemory Constant {} = pure False++sKernelBlock ::+  String ->+  VName ->+  KernelAttrs ->+  InKernelGen () ->+  CallKernelGen ()+sKernelBlock = sKernel blockOperations kernelBlockId++compileBlockResult ::+  SegSpace ->+  PatElem LetDecMem ->+  KernelResult ->+  InKernelGen ()+compileBlockResult _ pe (TileReturns _ [(w, per_block_elems)] what) = do+  n <- pe64 . arraySize 0 <$> lookupType what++  constants <- kernelConstants <$> askEnv+  let ltid = sExt64 $ kernelLocalThreadId constants+      offset =+        pe64 per_block_elems+          * sExt64 (kernelBlockId constants)++  -- Avoid loop for the common case where each thread is statically+  -- known to write at most one element.+  localOps threadOperations $+    if pe64 per_block_elems == kernelBlockSize constants+      then+        sWhen (ltid + offset .<. pe64 w) $+          copyDWIMFix (patElemName pe) [ltid + offset] (Var what) [ltid]+      else sFor "i" (n `divUp` kernelBlockSize constants) $ \i -> do+        j <- dPrimVE "j" $ kernelBlockSize constants * i + ltid+        sWhen (j + offset .<. pe64 w) $+          copyDWIMFix (patElemName pe) [j + offset] (Var what) [j]+compileBlockResult space pe (TileReturns _ dims what) = do+  let gids = map fst $ unSegSpace space+      out_tile_sizes = map (pe64 . snd) dims+      block_is = zipWith (*) (map Imp.le64 gids) out_tile_sizes+  local_is <- localThreadIDs $ map snd dims+  is_for_thread <-+    mapM (dPrimV "thread_out_index") $+      zipWith (+) block_is local_is++  localOps threadOperations $+    sWhen (isActive $ zip (map tvVar is_for_thread) $ map fst dims) $+      copyDWIMFix (patElemName pe) (map tvExp is_for_thread) (Var what) local_is+compileBlockResult space pe (RegTileReturns _ dims_n_tiles what) = do+  constants <- kernelConstants <$> askEnv++  let gids = map fst $ unSegSpace space+      (dims, block_tiles, reg_tiles) = unzip3 dims_n_tiles+      block_tiles' = map pe64 block_tiles+      reg_tiles' = map pe64 reg_tiles++  -- Which block tile is this block responsible for?+  let block_tile_is = map Imp.le64 gids++  -- Within the block tile, which register tile is this thread+  -- responsible for?+  reg_tile_is <-+    dIndexSpace' "reg_tile_i" block_tiles' $ sExt64 $ kernelLocalThreadId constants++  -- Compute output array slice for the register tile belonging to+  -- this thread.+  let regTileSliceDim (block_tile, block_tile_i) (reg_tile, reg_tile_i) = do+        tile_dim_start <-+          dPrimVE "tile_dim_start" $+            reg_tile * (block_tile * block_tile_i + reg_tile_i)+        pure $ DimSlice tile_dim_start reg_tile 1+  reg_tile_slices <-+    Slice+      <$> zipWithM+        regTileSliceDim+        (zip block_tiles' block_tile_is)+        (zip reg_tiles' reg_tile_is)++  localOps threadOperations $+    sLoopNest (Shape reg_tiles) $ \is_in_reg_tile -> do+      let dest_is = fixSlice reg_tile_slices is_in_reg_tile+          src_is = reg_tile_is ++ is_in_reg_tile+      sWhen (foldl1 (.&&.) $ zipWith (.<.) dest_is $ map pe64 dims) $+        copyDWIMFix (patElemName pe) dest_is (Var what) src_is+compileBlockResult space pe (Returns _ _ what) = do+  constants <- kernelConstants <$> askEnv+  in_shared_memory <- arrayInSharedMemory what+  let gids = map (Imp.le64 . fst) $ unSegSpace space++  if not in_shared_memory+    then+      localOps threadOperations $+        sWhen (kernelLocalThreadId constants .==. 0) $+          copyDWIMFix (patElemName pe) gids what []+    else -- If the result of the block is an array in local memory, we+    -- store it by collective copying among all the threads of the+    -- block.  TODO: also do this if the array is in global memory+    -- (but this is a bit more tricky, synchronisation-wise).+      copyDWIMFix (patElemName pe) gids what []+compileBlockResult _ _ WriteReturns {} =+  compilerLimitationS "compileBlockResult: WriteReturns not handled yet."++-- | The sizes of nested iteration spaces in the kernel.+type SegOpSizes = S.Set [SubExp]++-- | Various useful precomputed information for block-level SegOps.+data Precomputed = Precomputed+  { pcSegOpSizes :: SegOpSizes,+    pcChunkItersMap :: M.Map [SubExp] (Imp.TExp Int32)+  }++-- | Find the sizes of nested parallelism in a t'SegOp' body.+segOpSizes :: Stms GPUMem -> SegOpSizes+segOpSizes = onStms+  where+    onStms = foldMap onStm+    onStm (Let _ _ (Op (Inner (SegOp op)))) =+      case segVirt $ segLevel op of+        SegNoVirtFull seq_dims ->+          S.singleton $ map snd $ snd $ partitionSeqDims seq_dims $ segSpace op+        _ -> S.singleton $ map snd $ unSegSpace $ segSpace op+    onStm (Let (Pat [pe]) _ (BasicOp (Replicate {}))) =+      S.singleton $ arrayDims $ patElemType pe+    onStm (Let (Pat [pe]) _ (BasicOp (Iota {}))) =+      S.singleton $ arrayDims $ patElemType pe+    onStm (Let (Pat [pe]) _ (BasicOp (Manifest {}))) =+      S.singleton $ arrayDims $ patElemType pe+    onStm (Let _ _ (Match _ cases defbody _)) =+      foldMap (onStms . bodyStms . caseBody) cases <> onStms (bodyStms defbody)+    onStm (Let _ _ (Loop _ _ body)) =+      onStms (bodyStms body)+    onStm _ = mempty++-- | Precompute various constants and useful information.+precomputeConstants :: Count BlockSize (Imp.TExp Int64) -> Stms GPUMem -> CallKernelGen Precomputed+precomputeConstants tblock_size stms = do+  let sizes = segOpSizes stms+  iters_map <- M.fromList <$> mapM mkMap (S.toList sizes)+  pure $ Precomputed sizes iters_map+  where+    mkMap dims = do+      let n = product $ map Imp.pe64 dims+      num_chunks <- dPrimVE "num_chunks" $ sExt32 $ n `divUp` unCount tblock_size+      pure (dims, num_chunks)++-- | Make use of various precomputed constants.+precomputedConstants :: Precomputed -> InKernelGen a -> InKernelGen a+precomputedConstants pre m = do+  ltid <- kernelLocalThreadId . kernelConstants <$> askEnv+  new_ids <- M.fromList <$> mapM (mkMap ltid) (S.toList (pcSegOpSizes pre))+  let f env =+        env+          { kernelConstants =+              (kernelConstants env)+                { kernelLocalIdMap = new_ids,+                  kernelChunkItersMap = pcChunkItersMap pre+                }+          }+  localEnv f m+  where+    mkMap ltid dims = do+      let dims' = map pe64 dims+      ids' <- dIndexSpace' "ltid_pre" dims' (sExt64 ltid)+      pure (dims, map sExt32 ids')

@@ -1,743 +0,0 @@-{-# LANGUAGE TypeFamilies #-}---- | Generation of kernels with group-level bodies.-module Futhark.CodeGen.ImpGen.GPU.Group-  ( sKernelGroup,-    compileGroupResult,-    groupOperations,--    -- * Precomputation-    Precomputed,-    precomputeConstants,-    precomputedConstants,-    atomicUpdateLocking,-  )-where--import Control.Monad-import Data.Bifunctor-import Data.List (partition, zip4)-import Data.Map.Strict qualified as M-import Data.Maybe-import Data.Set qualified as S-import Futhark.CodeGen.ImpCode.GPU qualified as Imp-import Futhark.CodeGen.ImpGen-import Futhark.CodeGen.ImpGen.GPU.Base-import Futhark.Construct (fullSliceNum)-import Futhark.Error-import Futhark.IR.GPUMem-import Futhark.IR.Mem.LMAD qualified as LMAD-import Futhark.Transform.Rename-import Futhark.Util (chunks, mapAccumLM, takeLast)-import Futhark.Util.IntegralExp (divUp, rem)-import Prelude hiding (quot, rem)---- | @flattenArray k flat arr@ flattens the outer @k@ dimensions of--- @arr@ to @flat@.  (Make sure @flat@ is the sum of those dimensions--- or you'll have a bad time.)-flattenArray :: Int -> TV Int64 -> VName -> ImpM rep r op VName-flattenArray k flat arr = do-  ArrayEntry arr_loc pt <- lookupArray arr-  let flat_shape = Shape $ Var (tvVar flat) : drop k (memLocShape arr_loc)-  sArray (baseString arr ++ "_flat") pt flat_shape (memLocName arr_loc) $-    fromMaybe (error "flattenArray") $-      LMAD.reshape (memLocLMAD arr_loc) (map pe64 $ shapeDims flat_shape)--sliceArray :: Imp.TExp Int64 -> TV Int64 -> VName -> ImpM rep r op VName-sliceArray start size arr = do-  MemLoc mem _ ixfun <- entryArrayLoc <$> lookupArray arr-  arr_t <- lookupType arr-  let slice =-        fullSliceNum-          (map Imp.pe64 (arrayDims arr_t))-          [DimSlice start (tvExp size) 1]-  sArray-    (baseString arr ++ "_chunk")-    (elemType arr_t)-    (arrayShape arr_t `setOuterDim` Var (tvVar size))-    mem-    $ LMAD.slice ixfun slice---- | @applyLambda lam dests args@ emits code that:------ 1. Binds each parameter of @lam@ to the corresponding element of---    @args@, interpreted as a (name,slice) pair (as in 'copyDWIM').---    Use an empty list for a scalar.------ 2. Executes the body of @lam@.------ 3. Binds the t'SubExp's that are the 'Result' of @lam@ to the--- provided @dest@s, again interpreted as the destination for a--- 'copyDWIM'.-applyLambda ::-  (Mem rep inner) =>-  Lambda rep ->-  [(VName, [DimIndex (Imp.TExp Int64)])] ->-  [(SubExp, [DimIndex (Imp.TExp Int64)])] ->-  ImpM rep r op ()-applyLambda lam dests args = do-  dLParams $ lambdaParams lam-  forM_ (zip (lambdaParams lam) args) $ \(p, (arg, arg_slice)) ->-    copyDWIM (paramName p) [] arg arg_slice-  compileStms mempty (bodyStms $ lambdaBody lam) $ do-    let res = map resSubExp $ bodyResult $ lambdaBody lam-    forM_ (zip dests res) $ \((dest, dest_slice), se) ->-      copyDWIM dest dest_slice se []---- | As applyLambda, but first rename the names in the lambda.  This--- makes it safe to apply it in multiple places.  (It might be safe--- anyway, but you have to be more careful - use this if you are in--- doubt.)-applyRenamedLambda ::-  (Mem rep inner) =>-  Lambda rep ->-  [(VName, [DimIndex (Imp.TExp Int64)])] ->-  [(SubExp, [DimIndex (Imp.TExp Int64)])] ->-  ImpM rep r op ()-applyRenamedLambda lam dests args = do-  lam_renamed <- renameLambda lam-  applyLambda lam_renamed dests args--groupChunkLoop ::-  Imp.TExp Int32 ->-  (Imp.TExp Int32 -> TV Int64 -> InKernelGen ()) ->-  InKernelGen ()-groupChunkLoop w m = do-  constants <- kernelConstants <$> askEnv-  let max_chunk_size = sExt32 $ kernelGroupSize constants-  num_chunks <- dPrimVE "num_chunks" $ w `divUp` max_chunk_size-  sFor "chunk_i" num_chunks $ \chunk_i -> do-    chunk_start <--      dPrimVE "chunk_start" $ chunk_i * max_chunk_size-    chunk_end <--      dPrimVE "chunk_end" $ sMin32 w (chunk_start + max_chunk_size)-    chunk_size <--      dPrimV "chunk_size" $ sExt64 $ chunk_end - chunk_start-    m chunk_start chunk_size--virtualisedGroupScan ::-  Maybe (Imp.TExp Int32 -> Imp.TExp Int32 -> Imp.TExp Bool) ->-  Imp.TExp Int32 ->-  Lambda GPUMem ->-  [VName] ->-  InKernelGen ()-virtualisedGroupScan seg_flag w lam arrs = do-  groupChunkLoop w $ \chunk_start chunk_size -> do-    constants <- kernelConstants <$> askEnv-    let ltid = kernelLocalThreadId constants-        crosses_segment =-          case seg_flag of-            Nothing -> false-            Just flag_true ->-              flag_true (sExt32 (chunk_start - 1)) (sExt32 chunk_start)-    sComment "possibly incorporate carry" $-      sWhen (chunk_start .>. 0 .&&. ltid .==. 0 .&&. bNot crosses_segment) $ do-        carry_idx <- dPrimVE "carry_idx" $ sExt64 chunk_start - 1-        applyRenamedLambda-          lam-          (map (,[DimFix $ sExt64 chunk_start]) arrs)-          ( map ((,[DimFix carry_idx]) . Var) arrs-              ++ map ((,[DimFix $ sExt64 chunk_start]) . Var) arrs-          )--    arrs_chunks <- mapM (sliceArray (sExt64 chunk_start) chunk_size) arrs--    sOp $ Imp.ErrorSync Imp.FenceLocal--    groupScan seg_flag (sExt64 w) (tvExp chunk_size) lam arrs_chunks--copyInGroup :: CopyCompiler GPUMem KernelEnv Imp.KernelOp-copyInGroup pt destloc srcloc = do-  dest_space <- entryMemSpace <$> lookupMemory (memLocName destloc)-  src_space <- entryMemSpace <$> lookupMemory (memLocName srcloc)--  let src_lmad = memLocLMAD srcloc-      dims = LMAD.shape src_lmad-      rank = length dims--  case (dest_space, src_space) of-    (ScalarSpace destds _, ScalarSpace srcds _) -> do-      let fullDim d = DimSlice 0 d 1-          destslice' =-            Slice $-              replicate (rank - length destds) (DimFix 0)-                ++ takeLast (length destds) (map fullDim dims)-          srcslice' =-            Slice $-              replicate (rank - length srcds) (DimFix 0)-                ++ takeLast (length srcds) (map fullDim dims)-      lmadCopy-        pt-        (sliceMemLoc destloc destslice')-        (sliceMemLoc srcloc srcslice')-    _ -> do-      groupCoverSpace (map sExt32 dims) $ \is ->-        lmadCopy-          pt-          (sliceMemLoc destloc (Slice $ map (DimFix . sExt64) is))-          (sliceMemLoc srcloc (Slice $ map (DimFix . sExt64) is))-      sOp $ Imp.Barrier Imp.FenceLocal--localThreadIDs :: [SubExp] -> InKernelGen [Imp.TExp Int64]-localThreadIDs dims = do-  ltid <- sExt64 . kernelLocalThreadId . kernelConstants <$> askEnv-  let dims' = map pe64 dims-  maybe (dIndexSpace' "ltid" dims' ltid) (pure . map sExt64)-    . M.lookup dims-    . kernelLocalIdMap-    . kernelConstants-    =<< askEnv--partitionSeqDims :: SegSeqDims -> SegSpace -> ([(VName, SubExp)], [(VName, SubExp)])-partitionSeqDims (SegSeqDims seq_is) space =-  bimap (map fst) (map fst) $-    partition ((`elem` seq_is) . snd) (zip (unSegSpace space) [0 ..])--compileFlatId :: SegSpace -> InKernelGen ()-compileFlatId space = do-  ltid <- kernelLocalThreadId . kernelConstants <$> askEnv-  dPrimV_ (segFlat space) ltid---- Construct the necessary lock arrays for an intra-group histogram.-prepareIntraGroupSegHist ::-  Shape ->-  Count GroupSize SubExp ->-  [HistOp GPUMem] ->-  InKernelGen [[Imp.TExp Int64] -> InKernelGen ()]-prepareIntraGroupSegHist segments group_size =-  fmap snd . mapAccumLM onOp Nothing-  where-    onOp l op = do-      constants <- kernelConstants <$> askEnv-      atomicBinOp <- kernelAtomics <$> askEnv--      let local_subhistos = histDest op--      case (l, atomicUpdateLocking atomicBinOp $ histOp op) of-        (_, AtomicPrim f) -> pure (l, f (Space "local") local_subhistos)-        (_, AtomicCAS f) -> pure (l, f (Space "local") local_subhistos)-        (Just l', AtomicLocking f) -> pure (l, f l' (Space "local") local_subhistos)-        (Nothing, AtomicLocking f) -> do-          locks <- newVName "locks"--          let num_locks = pe64 $ unCount group_size-              dims = map pe64 $ shapeDims (segments <> histOpShape op <> histShape op)-              l' = Locking locks 0 1 0 (pure . (`rem` num_locks) . flattenIndex dims)-              locks_t = Array int32 (Shape [unCount group_size]) NoUniqueness--          locks_mem <- sAlloc "locks_mem" (typeSize locks_t) $ Space "local"-          dArray locks int32 (arrayShape locks_t) locks_mem $-            LMAD.iota 0 . map pe64 . arrayDims $-              locks_t--          sComment "All locks start out unlocked" $-            groupCoverSpace [kernelGroupSize constants] $ \is ->-              copyDWIMFix locks is (intConst Int32 0) []--          pure (Just l', f l' (Space "local") local_subhistos)--groupCoverSegSpace :: SegVirt -> SegSpace -> InKernelGen () -> InKernelGen ()-groupCoverSegSpace virt space m = do-  let (ltids, dims) = unzip $ unSegSpace space-      dims' = map pe64 dims--  constants <- kernelConstants <$> askEnv-  let group_size = kernelGroupSize constants-  -- Maybe we can statically detect that this is actually a-  -- SegNoVirtFull and generate ever-so-slightly simpler code.-  let virt' = if dims' == [group_size] then SegNoVirtFull (SegSeqDims []) else virt-  case virt' of-    SegVirt -> do-      iters <- M.lookup dims . kernelChunkItersMap . kernelConstants <$> askEnv-      case iters of-        Nothing -> do-          iterations <- dPrimVE "iterations" $ product $ map sExt32 dims'-          groupLoop iterations $ \i -> do-            dIndexSpace (zip ltids dims') $ sExt64 i-            m-        Just num_chunks -> localOps threadOperations $ do-          let ltid = kernelLocalThreadId constants-          sFor "chunk_i" num_chunks $ \chunk_i -> do-            i <- dPrimVE "i" $ chunk_i * sExt32 group_size + ltid-            dIndexSpace (zip ltids dims') $ sExt64 i-            sWhen (inBounds (Slice (map (DimFix . le64) ltids)) dims') m-    SegNoVirt -> localOps threadOperations $ do-      zipWithM_ dPrimV_ ltids =<< localThreadIDs dims-      sWhen (isActive $ zip ltids dims) m-    SegNoVirtFull seq_dims -> do-      let ((ltids_seq, dims_seq), (ltids_par, dims_par)) =-            bimap unzip unzip $ partitionSeqDims seq_dims space-      sLoopNest (Shape dims_seq) $ \is_seq -> do-        zipWithM_ dPrimV_ ltids_seq is_seq-        localOps threadOperations $ do-          zipWithM_ dPrimV_ ltids_par =<< localThreadIDs dims_par-          m--compileGroupExp :: ExpCompiler GPUMem KernelEnv Imp.KernelOp-compileGroupExp (Pat [pe]) (BasicOp (Opaque _ se)) =-  -- Cannot print in GPU code.-  copyDWIM (patElemName pe) [] se []--- The static arrays stuff does not work inside kernels.-compileGroupExp (Pat [dest]) (BasicOp (ArrayLit es _)) =-  forM_ (zip [0 ..] es) $ \(i, e) ->-    copyDWIMFix (patElemName dest) [fromIntegral (i :: Int64)] e []-compileGroupExp _ (BasicOp (UpdateAcc acc is vs)) = do-  ltid <- kernelLocalThreadId . kernelConstants <$> askEnv-  sWhen (ltid .==. 0) $ updateAcc acc is vs-  sOp $ Imp.Barrier Imp.FenceLocal-compileGroupExp (Pat [dest]) (BasicOp (Replicate ds se)) | ds /= mempty = do-  flat <- newVName "rep_flat"-  is <- replicateM (arrayRank dest_t) (newVName "rep_i")-  let is' = map le64 is-  groupCoverSegSpace SegVirt (SegSpace flat $ zip is $ arrayDims dest_t) $-    copyDWIMFix (patElemName dest) is' se (drop (shapeRank ds) is')-  sOp $ Imp.Barrier Imp.FenceLocal-  where-    dest_t = patElemType dest-compileGroupExp (Pat [dest]) (BasicOp (Iota n e s it)) = do-  n' <- toExp n-  e' <- toExp e-  s' <- toExp s-  groupLoop (TPrimExp n') $ \i' -> do-    x <--      dPrimV "x" $-        TPrimExp $-          BinOpExp (Add it OverflowUndef) e' $-            BinOpExp (Mul it OverflowUndef) (untyped i') s'-    copyDWIMFix (patElemName dest) [i'] (Var (tvVar x)) []-  sOp $ Imp.Barrier Imp.FenceLocal---- When generating code for a scalar in-place update, we must make--- sure that only one thread performs the write.  When writing an--- array, the group-level copy code will take care of doing the right--- thing.-compileGroupExp (Pat [pe]) (BasicOp (Update safety _ slice se))-  | null $ sliceDims slice = do-      sOp $ Imp.Barrier Imp.FenceLocal-      ltid <- kernelLocalThreadId . kernelConstants <$> askEnv-      sWhen (ltid .==. 0) $-        case safety of-          Unsafe -> write-          Safe -> sWhen (inBounds slice' dims) write-      sOp $ Imp.Barrier Imp.FenceLocal-  where-    slice' = fmap pe64 slice-    dims = map pe64 $ arrayDims $ patElemType pe-    write = copyDWIM (patElemName pe) (unSlice slice') se []-compileGroupExp dest e = do-  -- It is a messy to jump into control flow for error handling.-  -- Avoid that by always doing an error sync here.  Potential-  -- improvement: only do this if any errors are pending (this could-  -- also be handled in later codegen).-  when (doSync e) $ sOp $ Imp.ErrorSync Imp.FenceLocal-  defCompileExp dest e-  where-    doSync Loop {} = True-    doSync Match {} = True-    doSync _ = False--compileGroupOp :: OpCompiler GPUMem KernelEnv Imp.KernelOp-compileGroupOp pat (Alloc size space) =-  kernelAlloc pat size space-compileGroupOp pat (Inner (SegOp (SegMap lvl space _ body))) = do-  compileFlatId space--  groupCoverSegSpace (segVirt lvl) space $-    compileStms mempty (kernelBodyStms body) $-      zipWithM_ (compileThreadResult space) (patElems pat) $-        kernelBodyResult body-  sOp $ Imp.ErrorSync Imp.FenceLocal-compileGroupOp pat (Inner (SegOp (SegScan lvl space scans _ body))) = do-  compileFlatId space--  let (ltids, dims) = unzip $ unSegSpace space-      dims' = map pe64 dims--  groupCoverSegSpace (segVirt lvl) space $-    compileStms mempty (kernelBodyStms body) $-      forM_ (zip (patNames pat) $ kernelBodyResult body) $ \(dest, res) ->-        copyDWIMFix-          dest-          (map Imp.le64 ltids)-          (kernelResultSubExp res)-          []--  fence <- fenceForArrays $ patNames pat-  sOp $ Imp.ErrorSync fence--  let segment_size = last dims'-      crossesSegment from to =-        (sExt64 to - sExt64 from) .>. (sExt64 to `rem` segment_size)--  -- groupScan needs to treat the scan output as a one-dimensional-  -- array of scan elements, so we invent some new flattened arrays-  -- here.-  dims_flat <- dPrimV "dims_flat" $ product dims'-  let scan = head scans-      num_scan_results = length $ segBinOpNeutral scan-  arrs_flat <--    mapM (flattenArray (length dims') dims_flat) $-      take num_scan_results $-        patNames pat--  case segVirt lvl of-    SegVirt ->-      virtualisedGroupScan-        (Just crossesSegment)-        (sExt32 $ tvExp dims_flat)-        (segBinOpLambda scan)-        arrs_flat-    _ ->-      groupScan-        (Just crossesSegment)-        (product dims')-        (product dims')-        (segBinOpLambda scan)-        arrs_flat-compileGroupOp pat (Inner (SegOp (SegRed lvl space ops _ body))) = do-  compileFlatId space--  let dims' = map pe64 dims-      mkTempArr t =-        sAllocArray "red_arr" (elemType t) (Shape dims <> arrayShape t) $ Space "local"--  tmp_arrs <- mapM mkTempArr $ concatMap (lambdaReturnType . segBinOpLambda) ops-  groupCoverSegSpace (segVirt lvl) space $-    compileStms mempty (kernelBodyStms body) $ do-      let (red_res, map_res) =-            splitAt (segBinOpResults ops) $ kernelBodyResult body-      forM_ (zip tmp_arrs red_res) $ \(dest, res) ->-        copyDWIMFix dest (map Imp.le64 ltids) (kernelResultSubExp res) []-      zipWithM_ (compileThreadResult space) map_pes map_res--  sOp $ Imp.ErrorSync Imp.FenceLocal--  let tmps_for_ops = chunks (map (length . segBinOpNeutral) ops) tmp_arrs-  case segVirt lvl of-    SegVirt -> virtCase dims' tmps_for_ops-    _ -> nonvirtCase dims' tmps_for_ops-  where-    (ltids, dims) = unzip $ unSegSpace space-    (red_pes, map_pes) = splitAt (segBinOpResults ops) $ patElems pat--    virtCase [dim'] tmps_for_ops = do-      ltid <- kernelLocalThreadId . kernelConstants <$> askEnv-      groupChunkLoop (sExt32 dim') $ \chunk_start chunk_size -> do-        sComment "possibly incorporate carry" $-          sWhen (chunk_start .>. 0 .&&. ltid .==. 0) $-            forM_ (zip ops tmps_for_ops) $ \(op, tmps) ->-              applyRenamedLambda-                (segBinOpLambda op)-                (map (,[DimFix $ sExt64 chunk_start]) tmps)-                ( map ((,[]) . Var . patElemName) red_pes-                    ++ map ((,[DimFix $ sExt64 chunk_start]) . Var) tmps-                )--        sOp $ Imp.ErrorSync Imp.FenceLocal--        forM_ (zip ops tmps_for_ops) $ \(op, tmps) -> do-          tmps_chunks <- mapM (sliceArray (sExt64 chunk_start) chunk_size) tmps-          groupReduce (sExt32 (tvExp chunk_size)) (segBinOpLambda op) tmps_chunks--        sOp $ Imp.ErrorSync Imp.FenceLocal--        sComment "Save result of reduction." $-          forM_ (zip red_pes $ concat tmps_for_ops) $ \(pe, arr) ->-            copyDWIMFix (patElemName pe) [] (Var arr) [sExt64 chunk_start]--    ---    virtCase dims' tmps_for_ops = do-      dims_flat <- dPrimV "dims_flat" $ product dims'-      let segment_size = last dims'-          crossesSegment from to =-            (sExt64 to - sExt64 from) .>. (sExt64 to `rem` sExt64 segment_size)--      forM_ (zip ops tmps_for_ops) $ \(op, tmps) -> do-        tmps_flat <- mapM (flattenArray (length dims') dims_flat) tmps-        virtualisedGroupScan-          (Just crossesSegment)-          (sExt32 $ tvExp dims_flat)-          (segBinOpLambda op)-          tmps_flat--      sOp $ Imp.ErrorSync Imp.FenceLocal--      sComment "Save result of reduction." $-        forM_ (zip red_pes $ concat tmps_for_ops) $ \(pe, arr) ->-          copyDWIM-            (patElemName pe)-            []-            (Var arr)-            (map (unitSlice 0) (init dims') ++ [DimFix $ last dims' - 1])--      sOp $ Imp.Barrier Imp.FenceLocal--    -- Nonsegmented case (or rather, a single segment) - this we can-    -- handle directly with a group-level reduction.-    nonvirtCase [dim'] tmps_for_ops = do-      forM_ (zip ops tmps_for_ops) $ \(op, tmps) ->-        groupReduce (sExt32 dim') (segBinOpLambda op) tmps-      sOp $ Imp.ErrorSync Imp.FenceLocal-      sComment "Save result of reduction." $-        forM_ (zip red_pes $ concat tmps_for_ops) $ \(pe, arr) ->-          copyDWIMFix (patElemName pe) [] (Var arr) [0]-      sOp $ Imp.ErrorSync Imp.FenceLocal--    -- Segmented intra-group reductions are turned into (regular)-    -- segmented scans.  It is possible that this can be done-    -- better, but at least this approach is simple.-    nonvirtCase dims' tmps_for_ops = do-      -- groupScan operates on flattened arrays.  This does not-      -- involve copying anything; merely playing with the index-      -- function.-      dims_flat <- dPrimV "dims_flat" $ product dims'-      let segment_size = last dims'-          crossesSegment from to =-            (sExt64 to - sExt64 from) .>. (sExt64 to `rem` sExt64 segment_size)--      forM_ (zip ops tmps_for_ops) $ \(op, tmps) -> do-        tmps_flat <- mapM (flattenArray (length dims') dims_flat) tmps-        groupScan-          (Just crossesSegment)-          (product dims')-          (product dims')-          (segBinOpLambda op)-          tmps_flat--      sOp $ Imp.ErrorSync Imp.FenceLocal--      sComment "Save result of reduction." $-        forM_ (zip red_pes $ concat tmps_for_ops) $ \(pe, arr) ->-          copyDWIM-            (patElemName pe)-            []-            (Var arr)-            (map (unitSlice 0) (init dims') ++ [DimFix $ last dims' - 1])--      sOp $ Imp.Barrier Imp.FenceLocal-compileGroupOp pat (Inner (SegOp (SegHist lvl space ops _ kbody))) = do-  compileFlatId space-  let (ltids, dims) = unzip $ unSegSpace space--  -- We don't need the red_pes, because it is guaranteed by our type-  -- rules that they occupy the same memory as the destinations for-  -- the ops.-  let num_red_res = length ops + sum (map (length . histNeutral) ops)-      (_red_pes, map_pes) =-        splitAt num_red_res $ patElems pat--  group_size <- kernelGroupSizeCount . kernelConstants <$> askEnv-  ops' <- prepareIntraGroupSegHist (Shape $ init dims) group_size ops--  -- Ensure that all locks have been initialised.-  sOp $ Imp.Barrier Imp.FenceLocal--  groupCoverSegSpace (segVirt lvl) space $-    compileStms mempty (kernelBodyStms kbody) $ do-      let (red_res, map_res) = splitAt num_red_res $ kernelBodyResult kbody-          (red_is, red_vs) = splitAt (length ops) $ map kernelResultSubExp red_res-      zipWithM_ (compileThreadResult space) map_pes map_res--      let vs_per_op = chunks (map (length . histDest) ops) red_vs--      forM_ (zip4 red_is vs_per_op ops' ops) $-        \(bin, op_vs, do_op, HistOp dest_shape _ _ _ shape lam) -> do-          let bin' = pe64 bin-              dest_shape' = map pe64 $ shapeDims dest_shape-              bin_in_bounds = inBounds (Slice [DimFix bin']) dest_shape'-              bin_is = map Imp.le64 (init ltids) ++ [bin']-              vs_params = takeLast (length op_vs) $ lambdaParams lam--          sComment "perform atomic updates" $-            sWhen bin_in_bounds $ do-              dLParams $ lambdaParams lam-              sLoopNest shape $ \is -> do-                forM_ (zip vs_params op_vs) $ \(p, v) ->-                  copyDWIMFix (paramName p) [] v is-                do_op (bin_is ++ is)--  sOp $ Imp.ErrorSync Imp.FenceLocal-compileGroupOp pat _ =-  compilerBugS $ "compileGroupOp: cannot compile rhs of binding " ++ prettyString pat--groupOperations :: Operations GPUMem KernelEnv Imp.KernelOp-groupOperations =-  (defaultOperations compileGroupOp)-    { opsCopyCompiler = copyInGroup,-      opsExpCompiler = compileGroupExp,-      opsStmsCompiler = \_ -> defCompileStms mempty,-      opsAllocCompilers =-        M.fromList [(Space "local", allocLocal)]-    }--arrayInLocalMemory :: SubExp -> InKernelGen Bool-arrayInLocalMemory (Var name) = do-  res <- lookupVar name-  case res of-    ArrayVar _ entry ->-      (Space "local" ==) . entryMemSpace-        <$> lookupMemory (memLocName (entryArrayLoc entry))-    _ -> pure False-arrayInLocalMemory Constant {} = pure False--sKernelGroup ::-  String ->-  VName ->-  KernelAttrs ->-  InKernelGen () ->-  CallKernelGen ()-sKernelGroup = sKernel groupOperations kernelGroupId--compileGroupResult ::-  SegSpace ->-  PatElem LetDecMem ->-  KernelResult ->-  InKernelGen ()-compileGroupResult _ pe (TileReturns _ [(w, per_group_elems)] what) = do-  n <- pe64 . arraySize 0 <$> lookupType what--  constants <- kernelConstants <$> askEnv-  let ltid = sExt64 $ kernelLocalThreadId constants-      offset =-        pe64 per_group_elems-          * sExt64 (kernelGroupId constants)--  -- Avoid loop for the common case where each thread is statically-  -- known to write at most one element.-  localOps threadOperations $-    if pe64 per_group_elems == kernelGroupSize constants-      then-        sWhen (ltid + offset .<. pe64 w) $-          copyDWIMFix (patElemName pe) [ltid + offset] (Var what) [ltid]-      else sFor "i" (n `divUp` kernelGroupSize constants) $ \i -> do-        j <- dPrimVE "j" $ kernelGroupSize constants * i + ltid-        sWhen (j + offset .<. pe64 w) $-          copyDWIMFix (patElemName pe) [j + offset] (Var what) [j]-compileGroupResult space pe (TileReturns _ dims what) = do-  let gids = map fst $ unSegSpace space-      out_tile_sizes = map (pe64 . snd) dims-      group_is = zipWith (*) (map Imp.le64 gids) out_tile_sizes-  local_is <- localThreadIDs $ map snd dims-  is_for_thread <--    mapM (dPrimV "thread_out_index") $-      zipWith (+) group_is local_is--  localOps threadOperations $-    sWhen (isActive $ zip (map tvVar is_for_thread) $ map fst dims) $-      copyDWIMFix (patElemName pe) (map tvExp is_for_thread) (Var what) local_is-compileGroupResult space pe (RegTileReturns _ dims_n_tiles what) = do-  constants <- kernelConstants <$> askEnv--  let gids = map fst $ unSegSpace space-      (dims, group_tiles, reg_tiles) = unzip3 dims_n_tiles-      group_tiles' = map pe64 group_tiles-      reg_tiles' = map pe64 reg_tiles--  -- Which group tile is this group responsible for?-  let group_tile_is = map Imp.le64 gids--  -- Within the group tile, which register tile is this thread-  -- responsible for?-  reg_tile_is <--    dIndexSpace' "reg_tile_i" group_tiles' $ sExt64 $ kernelLocalThreadId constants--  -- Compute output array slice for the register tile belonging to-  -- this thread.-  let regTileSliceDim (group_tile, group_tile_i) (reg_tile, reg_tile_i) = do-        tile_dim_start <--          dPrimVE "tile_dim_start" $-            reg_tile * (group_tile * group_tile_i + reg_tile_i)-        pure $ DimSlice tile_dim_start reg_tile 1-  reg_tile_slices <--    Slice-      <$> zipWithM-        regTileSliceDim-        (zip group_tiles' group_tile_is)-        (zip reg_tiles' reg_tile_is)--  localOps threadOperations $-    sLoopNest (Shape reg_tiles) $ \is_in_reg_tile -> do-      let dest_is = fixSlice reg_tile_slices is_in_reg_tile-          src_is = reg_tile_is ++ is_in_reg_tile-      sWhen (foldl1 (.&&.) $ zipWith (.<.) dest_is $ map pe64 dims) $-        copyDWIMFix (patElemName pe) dest_is (Var what) src_is-compileGroupResult space pe (Returns _ _ what) = do-  constants <- kernelConstants <$> askEnv-  in_local_memory <- arrayInLocalMemory what-  let gids = map (Imp.le64 . fst) $ unSegSpace space--  if not in_local_memory-    then-      localOps threadOperations $-        sWhen (kernelLocalThreadId constants .==. 0) $-          copyDWIMFix (patElemName pe) gids what []-    else -- If the result of the group is an array in local memory, we-    -- store it by collective copying among all the threads of the-    -- group.  TODO: also do this if the array is in global memory-    -- (but this is a bit more tricky, synchronisation-wise).-      copyDWIMFix (patElemName pe) gids what []-compileGroupResult _ _ WriteReturns {} =-  compilerLimitationS "compileGroupResult: WriteReturns not handled yet."---- | The sizes of nested iteration spaces in the kernel.-type SegOpSizes = S.Set [SubExp]---- | Various useful precomputed information for group-level SegOps.-data Precomputed = Precomputed-  { pcSegOpSizes :: SegOpSizes,-    pcChunkItersMap :: M.Map [SubExp] (Imp.TExp Int32)-  }---- | Find the sizes of nested parallelism in a t'SegOp' body.-segOpSizes :: Stms GPUMem -> SegOpSizes-segOpSizes = onStms-  where-    onStms = foldMap onStm-    onStm (Let _ _ (Op (Inner (SegOp op)))) =-      case segVirt $ segLevel op of-        SegNoVirtFull seq_dims ->-          S.singleton $ map snd $ snd $ partitionSeqDims seq_dims $ segSpace op-        _ -> S.singleton $ map snd $ unSegSpace $ segSpace op-    onStm (Let (Pat [pe]) _ (BasicOp (Replicate {}))) =-      S.singleton $ arrayDims $ patElemType pe-    onStm (Let (Pat [pe]) _ (BasicOp (Iota {}))) =-      S.singleton $ arrayDims $ patElemType pe-    onStm (Let (Pat [pe]) _ (BasicOp (Manifest {}))) =-      S.singleton $ arrayDims $ patElemType pe-    onStm (Let _ _ (Match _ cases defbody _)) =-      foldMap (onStms . bodyStms . caseBody) cases <> onStms (bodyStms defbody)-    onStm (Let _ _ (Loop _ _ body)) =-      onStms (bodyStms body)-    onStm _ = mempty---- | Precompute various constants and useful information.-precomputeConstants :: Count GroupSize (Imp.TExp Int64) -> Stms GPUMem -> CallKernelGen Precomputed-precomputeConstants group_size stms = do-  let sizes = segOpSizes stms-  iters_map <- M.fromList <$> mapM mkMap (S.toList sizes)-  pure $ Precomputed sizes iters_map-  where-    mkMap dims = do-      let n = product $ map Imp.pe64 dims-      num_chunks <- dPrimVE "num_chunks" $ sExt32 $ n `divUp` unCount group_size-      pure (dims, num_chunks)---- | Make use of various precomputed constants.-precomputedConstants :: Precomputed -> InKernelGen a -> InKernelGen a-precomputedConstants pre m = do-  ltid <- kernelLocalThreadId . kernelConstants <$> askEnv-  new_ids <- M.fromList <$> mapM (mkMap ltid) (S.toList (pcSegOpSizes pre))-  let f env =-        env-          { kernelConstants =-              (kernelConstants env)-                { kernelLocalIdMap = new_ids,-                  kernelChunkItersMap = pcChunkItersMap pre-                }-          }-  localEnv f m-  where-    mkMap ltid dims = do-      let dims' = map pe64 dims-      ids' <- dIndexSpace' "ltid_pre" dims' (sExt64 ltid)-      pure (dims, map sExt32 ids')

@@ -14,12 +14,12 @@ -- -- H: total size of histograms in bytes, including any lock arrays. ----- G: group size+-- G: block size -- -- T: number of bytes of local memory each thread can be given without -- impacting occupancy (determined experimentally, e.g. 32). ----- LMAX: maximum amount of local memory per workgroup (hard limit).+-- LMAX: maximum amount of local memory per threadblock (hard limit). -- -- We wish to compute: --@@ -380,8 +380,8 @@  histKernelGlobalPass ::   [PatElem LetDecMem] ->-  Count NumGroups SubExp ->-  Count GroupSize SubExp ->+  Count NumBlocks SubExp ->+  Count BlockSize SubExp ->   SegSpace ->   [SegHistSlug] ->   KernelBody GPUMem ->@@ -389,7 +389,7 @@   Imp.TExp Int32 ->   Imp.TExp Int32 ->   CallKernelGen ()-histKernelGlobalPass map_pes num_groups group_size space slugs kbody histograms hist_S chk_i = do+histKernelGlobalPass map_pes num_tblocks tblock_size space slugs kbody histograms hist_S chk_i = do   let (space_is, space_sizes) = unzip $ unSegSpace space       space_sizes_64 = map (sExt64 . pe64) space_sizes       total_w_64 = product space_sizes_64@@ -397,7 +397,7 @@   hist_H_chks <- forM (map (histSize . slugOp) slugs) $ \w ->     dPrimVE "hist_H_chk" $ w `divUp` sExt64 hist_S -  sKernelThread "seghist_global" (segFlat space) (defKernelAttrs num_groups group_size) $ do+  sKernelThread "seghist_global" (segFlat space) (defKernelAttrs num_tblocks tblock_size) $ do     constants <- kernelConstants <$> askEnv      -- Compute subhistogram index for each thread, per histogram.@@ -472,17 +472,17 @@  histKernelGlobal ::   [PatElem LetDecMem] ->-  Count NumGroups SubExp ->-  Count GroupSize SubExp ->+  Count NumBlocks SubExp ->+  Count BlockSize SubExp ->   SegSpace ->   [SegHistSlug] ->   KernelBody GPUMem ->   CallKernelGen ()-histKernelGlobal map_pes num_groups group_size space slugs kbody = do-  let num_groups' = fmap pe64 num_groups-      group_size' = fmap pe64 group_size+histKernelGlobal map_pes num_tblocks tblock_size space slugs kbody = do+  let num_tblocks' = fmap pe64 num_tblocks+      tblock_size' = fmap pe64 tblock_size   let (_space_is, space_sizes) = unzip $ unSegSpace space-      num_threads = sExt32 $ unCount num_groups' * unCount group_size'+      num_threads = sExt32 $ unCount num_tblocks' * unCount tblock_size'    emit $ Imp.DebugPrint "## Using global memory" Nothing @@ -497,8 +497,8 @@   sFor "chk_i" hist_S $ \chk_i ->     histKernelGlobalPass       map_pes-      num_groups-      group_size+      num_tblocks+      tblock_size       space       slugs       kbody@@ -518,14 +518,14 @@  prepareIntermediateArraysLocal ::   TV Int32 ->-  Count NumGroups (Imp.TExp Int64) ->+  Count NumBlocks (Imp.TExp Int64) ->   [SegHistSlug] ->   CallKernelGen InitLocalHistograms-prepareIntermediateArraysLocal num_subhistos_per_group groups_per_segment =+prepareIntermediateArraysLocal num_subhistos_per_block blocks_per_segment =   mapM onOp   where     onOp (SegHistSlug op num_subhistos subhisto_info do_op) = do-      num_subhistos <-- sExt64 (unCount groups_per_segment)+      num_subhistos <-- sExt64 (unCount blocks_per_segment)        emit $         Imp.DebugPrint "Number of subhistograms in global memory per segment" $@@ -539,14 +539,14 @@           AtomicCAS f -> pure $ const $ pure f           AtomicLocking f -> pure $ \hist_H_chk -> do             let lock_shape =-                  Shape [tvSize num_subhistos_per_group, hist_H_chk]+                  Shape [tvSize num_subhistos_per_block, hist_H_chk]              let dims = map pe64 $ shapeDims lock_shape -            locks <- sAllocArray "locks" int32 lock_shape $ Space "local"+            locks <- sAllocArray "locks" int32 lock_shape $ Space "shared"              sComment "All locks start out unlocked" $-              groupCoverSpace dims $ \is ->+              blockCoverSpace dims $ \is ->                 copyDWIMFix locks is (intConst Int32 0) []              pure $ f $ Locking locks 0 1 0 id@@ -557,17 +557,17 @@             local_subhistos <-               forM (histType op) $ \t -> do                 let sub_local_shape =-                      Shape [tvSize num_subhistos_per_group]+                      Shape [tvSize num_subhistos_per_block]                         <> setOuterDims (arrayShape t) (histRank op) (Shape [hist_H_chk])                 sAllocArray                   "subhistogram_local"                   (elemType t)                   sub_local_shape-                  (Space "local")+                  (Space "shared")              do_op' <- mk_op hist_H_chk -            pure (local_subhistos, do_op' (Space "local") local_subhistos)+            pure (local_subhistos, do_op' (Space "shared") local_subhistos)        -- Initialise global-memory sub-histograms.       glob_subhistos <- forM subhisto_info $ \info -> do@@ -578,10 +578,10 @@  histKernelLocalPass ::   TV Int32 ->-  Count NumGroups (Imp.TExp Int64) ->+  Count NumBlocks (Imp.TExp Int64) ->   [PatElem LetDecMem] ->-  Count NumGroups SubExp ->-  Count GroupSize SubExp ->+  Count NumBlocks SubExp ->+  Count BlockSize SubExp ->   SegSpace ->   [SegHistSlug] ->   KernelBody GPUMem ->@@ -590,11 +590,11 @@   Imp.TExp Int32 ->   CallKernelGen () histKernelLocalPass-  num_subhistos_per_group_var-  groups_per_segment+  num_subhistos_per_block_var+  blocks_per_segment   map_pes-  num_groups-  group_size+  num_tblocks+  tblock_size   space   slugs   kbody@@ -605,7 +605,7 @@         segment_is = init space_is         segment_dims = init space_sizes         (i_in_segment, segment_size) = last $ unSegSpace space-        num_subhistos_per_group = tvExp num_subhistos_per_group_var+        num_subhistos_per_block = tvExp num_subhistos_per_block_var         segment_size' = pe64 segment_size      num_segments <-@@ -622,29 +622,29 @@               : map pe64 (shapeDims (histOpShape (slugOp slug)))       histo_size <-         dPrimVE "histo_size" $ product histo_dims-      let group_hists_size =-            sExt64 num_subhistos_per_group * histo_size+      let block_hists_size =+            sExt64 num_subhistos_per_block * histo_size       init_per_thread <--        dPrimVE "init_per_thread" $ sExt32 $ group_hists_size `divUp` pe64 (unCount group_size)+        dPrimVE "init_per_thread" $ sExt32 $ block_hists_size `divUp` pe64 (unCount tblock_size)       pure (histo_dims, histo_size, init_per_thread) -    let attrs = (defKernelAttrs num_groups group_size) {kAttrCheckLocalMemory = False}+    let attrs = (defKernelAttrs num_tblocks tblock_size) {kAttrCheckSharedMemory = False}     sKernelThread "seghist_local" (segFlat space) attrs $-      virtualiseGroups SegVirt (sExt32 $ unCount groups_per_segment * num_segments) $ \group_id -> do+      virtualiseBlocks SegVirt (sExt32 $ unCount blocks_per_segment * num_segments) $ \tblock_id -> do         constants <- kernelConstants <$> askEnv -        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)+        flat_segment_id <- dPrimVE "flat_segment_id" $ tblock_id `quot` sExt32 (unCount blocks_per_segment)+        gid_in_segment <- dPrimVE "gid_in_segment" $ tblock_id `rem` sExt32 (unCount blocks_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 * sExt32 (kernelGroupSize constants)+            gid_in_segment * sExt32 (kernelBlockSize constants)               + kernelLocalThreadId constants         threads_per_segment <-           dPrimVE "threads_per_segment" $             sExt32 $-              unCount groups_per_segment * kernelGroupSize constants+              unCount blocks_per_segment * kernelBlockSize constants          -- Set segment indices.         zipWithM_ dPrimV_ segment_is $@@ -661,7 +661,7 @@         -- warp use different subhistograms, to avoid conflicts.         thread_local_subhisto_i <-           dPrimVE "thread_local_subhisto_i" $-            kernelLocalThreadId constants `rem` num_subhistos_per_group+            kernelLocalThreadId constants `rem` num_subhistos_per_block          let onSlugs f =               forM_ (zip3 slugs histograms histo_sizes) $@@ -670,18 +670,18 @@          let onAllHistograms f =               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+                let block_hists_size = num_subhistos_per_block * 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 * sExt32 (kernelGroupSize constants)+                          i * sExt32 (kernelBlockSize constants)                             + kernelLocalThreadId constants                       j_offset <-                         dPrimVE "j_offset" $-                          num_subhistos_per_group * sExt32 histo_size * gid_in_segment + j+                          num_subhistos_per_block * sExt32 histo_size * gid_in_segment + j                        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@@ -695,7 +695,7 @@                               ++ tail local_bucket_is                       global_subhisto_i <- dPrimVE "global_subhisto_i" $ j_offset `quot` sExt32 histo_size -                      sWhen (j .<. group_hists_size) $+                      sWhen (j .<. block_hists_size) $                         f                           dest_local                           dest_global@@ -785,7 +785,7 @@             sFor "local_i" bins_per_thread $ \i -> do               j <-                 dPrimVE "j" $-                  i * sExt32 (kernelGroupSize constants)+                  i * sExt32 (kernelBlockSize constants)                     + kernelLocalThreadId constants               sWhen (j .<. trunc_histo_size) $ do                 -- We are responsible for compacting the flat bin 'j', which@@ -815,7 +815,7 @@                       (0 : local_bucket_is)                  sComment "Accumulate based on values in other subhistograms." $-                  sFor "subhisto_id" (num_subhistos_per_group - 1) $ \subhisto_id -> do+                  sFor "subhisto_id" (num_subhistos_per_block - 1) $ \subhisto_id -> do                     forM_ (zip yparams local_dests) $ \(yp, subhisto) ->                       copyDWIMFix                         (paramName yp)@@ -827,40 +827,40 @@                 sComment "Put final bucket value in global memory." $ do                   let global_is =                         map Imp.le64 segment_is-                          ++ [sExt64 group_id `rem` unCount groups_per_segment]+                          ++ [sExt64 tblock_id `rem` unCount blocks_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 Int64) ->+  Count NumBlocks (Imp.TExp Int64) ->   [PatElem LetDecMem] ->-  Count NumGroups SubExp ->-  Count GroupSize SubExp ->+  Count NumBlocks SubExp ->+  Count BlockSize SubExp ->   SegSpace ->   Imp.TExp Int32 ->   [SegHistSlug] ->   KernelBody GPUMem ->   CallKernelGen ()-histKernelLocal num_subhistos_per_group_var groups_per_segment map_pes num_groups group_size space hist_S slugs kbody = do-  let num_subhistos_per_group = tvExp num_subhistos_per_group_var+histKernelLocal num_subhistos_per_block_var blocks_per_segment map_pes num_tblocks tblock_size space hist_S slugs kbody = do+  let num_subhistos_per_block = tvExp num_subhistos_per_block_var    emit $-    Imp.DebugPrint "Number of local subhistograms per group" $+    Imp.DebugPrint "Number of local subhistograms per block" $       Just $-        untyped num_subhistos_per_group+        untyped num_subhistos_per_block    init_histograms <--    prepareIntermediateArraysLocal num_subhistos_per_group_var groups_per_segment slugs+    prepareIntermediateArraysLocal num_subhistos_per_block_var blocks_per_segment slugs    sFor "chk_i" hist_S $ \chk_i ->     histKernelLocalPass-      num_subhistos_per_group_var-      groups_per_segment+      num_subhistos_per_block_var+      blocks_per_segment       map_pes-      num_groups-      group_size+      num_tblocks+      tblock_size       space       slugs       kbody@@ -894,30 +894,30 @@       segmented = not $ null segment_dims    hist_L <- dPrim "hist_L" int32-  sOp $ Imp.GetSizeMax (tvVar hist_L) Imp.SizeLocalMemory+  sOp $ Imp.GetSizeMax (tvVar hist_L) Imp.SizeSharedMemory -  max_group_size <- dPrim "max_group_size" int32-  sOp $ Imp.GetSizeMax (tvVar max_group_size) Imp.SizeGroup+  max_tblock_size <- dPrim "max_tblock_size" int32+  sOp $ Imp.GetSizeMax (tvVar max_tblock_size) Imp.SizeThreadBlock -  -- XXX: we need to record for later use that max_group_size is the+  -- XXX: we need to record for later use that max_tblock_size is the   -- result of GetSizeMax.  This is an ugly hack that reflects our   -- inability to track which variables are actually constants.   let withSizeMax vtable =-        case M.lookup (tvVar max_group_size) vtable of+        case M.lookup (tvVar max_tblock_size) vtable of           Just (ScalarVar _ se) ->             M.insert-              (tvVar max_group_size)-              (ScalarVar (Just (Op (Inner (SizeOp (GetSizeMax SizeGroup))))) se)+              (tvVar max_tblock_size)+              (ScalarVar (Just (Op (Inner (SizeOp (GetSizeMax SizeThreadBlock))))) se)               vtable           _ -> vtable -  let group_size = Imp.Count $ Var $ tvVar max_group_size-  num_groups <-+  let tblock_size = Imp.Count $ Var $ tvVar max_tblock_size+  num_tblocks <-     fmap (Imp.Count . tvSize) $-      dPrimV "num_groups" $-        sExt64 hist_T `divUp` pe64 (unCount group_size)-  let num_groups' = pe64 <$> num_groups-      group_size' = pe64 <$> group_size+      dPrimV "num_tblocks" $+        sExt64 hist_T `divUp` pe64 (unCount tblock_size)+  let num_tblocks' = pe64 <$> num_tblocks+      tblock_size' = pe64 <$> tblock_size    let r64 = isF64 . ConvOpExp (SIToFP Int64 Float64) . untyped       t64 = isInt64 . ConvOpExp (FPToSI Float64 Int64) . untyped@@ -928,11 +928,11 @@       r64         ( sMin64             (sExt64 (tvExp hist_L `quot` hist_el_size))-            (hist_N `divUp` sExt64 (unCount num_groups'))+            (hist_N `divUp` sExt64 (unCount num_tblocks'))         )         / r64 hist_H -  let hist_B = unCount group_size'+  let hist_B = unCount tblock_size'    -- M in the paper, but not adjusted for asymptotic efficiency.   hist_M0 <-@@ -958,14 +958,14 @@               sMin64 (sExt64 hist_Nin * sExt64 hist_Nout) (sExt64 hist_T)                 `divUp` sExt64 hist_Nout -        -- Number of groups, rounded up.+        -- Number of blocks, rounded up.         let r = hist_T_hist_min `divUp` sExt32 hist_B          dPrimVE "work_asymp_M_max" $ hist_Nin `quot` (sExt64 r * hist_H)       else         dPrimVE "work_asymp_M_max" $           (hist_Nout * hist_N)-            `quot` ( (q_small * unCount num_groups' * hist_H)+            `quot` ( (q_small * unCount num_tblocks' * hist_H)                        `quot` genericLength slugs                    ) @@ -1008,13 +1008,13 @@         MustBeSinglePass -> 1         MayBeMultiPass -> fromIntegral $ maxinum $ map slugMaxLocalMemPasses slugs -  groups_per_segment <-+  blocks_per_segment <-     if segmented       then         fmap Count $-          dPrimVE "groups_per_segment" $-            unCount num_groups' `divUp` hist_Nout-      else pure num_groups'+          dPrimVE "blocks_per_segment" $+            unCount num_tblocks' `divUp` hist_Nout+      else pure num_tblocks'    -- We only use local memory if the number of updates per histogram   -- at least matches the histogram size, as otherwise it is not@@ -1038,17 +1038,17 @@         emit $ Imp.DebugPrint "Number of chunks (S)" $ Just $ untyped hist_S         when segmented $           emit $-            Imp.DebugPrint "Groups per segment" $+            Imp.DebugPrint "Blocks per segment" $               Just $                 untyped $-                  unCount groups_per_segment+                  unCount blocks_per_segment         localVTable withSizeMax $           histKernelLocal             hist_M-            groups_per_segment+            blocks_per_segment             map_pes-            num_groups-            group_size+            num_tblocks+            tblock_size             space             hist_S             slugs@@ -1065,14 +1065,14 @@   KernelBody GPUMem ->   CallKernelGen () compileSegHist (Pat pes) lvl space ops kbody = do-  KernelAttrs {kAttrNumGroups = num_groups, kAttrGroupSize = group_size} <-+  KernelAttrs {kAttrNumBlocks = num_tblocks, kAttrBlockSize = tblock_size} <-     lvlKernelAttrs lvl   -- Most of this function is not the histogram part itself, but   -- rather figuring out whether to use a local or global memory   -- strategy, as well as collapsing the subhistograms produced (which   -- are always in global memory, but their number may vary).-  let num_groups' = fmap pe64 num_groups-      group_size' = fmap pe64 group_size+  let num_tblocks' = fmap pe64 num_tblocks+      tblock_size' = fmap pe64 tblock_size       dims = map pe64 $ segSpaceDims space        num_red_res = length ops + sum (map (length . histNeutral) ops)@@ -1085,8 +1085,8 @@    -- Check for emptyness to avoid division-by-zero.   sUnless (seg_h .==. 0) $ do-    -- Maximum group size (or actual, in this case).-    let hist_B = unCount group_size'+    -- Maximum block size (or actual, in this case).+    let hist_B = unCount tblock_size'      -- Size of a histogram.     hist_H <- dPrimVE "hist_H" $ sum $ map histSize ops@@ -1112,10 +1112,10 @@           sum (map (pe64 . histRaceFactor . slugOp) slugs)             `quot` genericLength slugs -    let hist_T = sExt32 $ unCount num_groups' * unCount group_size'+    let hist_T = sExt32 $ unCount num_tblocks' * unCount tblock_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+    emit $ Imp.DebugPrint "Desired block size (B)" $ Just $ untyped hist_B     emit $ Imp.DebugPrint "Histogram size (H)" $ Just $ untyped hist_H     emit $ Imp.DebugPrint "Input elements per histogram (N)" $ Just $ untyped hist_N     emit $@@ -1129,11 +1129,11 @@     emit $ Imp.DebugPrint "Memory per set of subhistograms per segment" $ Just $ untyped h     emit $ Imp.DebugPrint "Memory per set of subhistograms times segments" $ Just $ untyped seg_h -    (use_local_memory, run_in_local_memory) <-+    (use_shared_memory, run_in_shared_memory) <-       localMemoryCase map_pes hist_T space hist_H hist_el_size hist_N hist_RF slugs kbody -    sIf use_local_memory run_in_local_memory $-      histKernelGlobal map_pes num_groups group_size space slugs kbody+    sIf use_shared_memory run_in_shared_memory $+      histKernelGlobal map_pes num_tblocks tblock_size space slugs kbody      let pes_per_op = chunks (map (length . histDest) ops) all_red_pes @@ -1166,7 +1166,7 @@          flat_gtid <- newVName "flat_gtid" -        let grid = KernelGrid num_groups group_size+        let grid = KernelGrid num_tblocks tblock_size             segred_space =               SegSpace flat_gtid $                 segment_dims

@@ -3,14 +3,14 @@ -- | Code generation for 'SegMap' is quite straightforward.  The only -- trick is virtualisation in case the physical number of threads is -- not sufficient to cover the logical thread space.  This is handled--- by having actual workgroups run a loop to imitate multiple workgroups.+-- by having actual threadblocks run a loop to imitate multiple threadblocks. module Futhark.CodeGen.ImpGen.GPU.SegMap (compileSegMap) where  import Control.Monad import Futhark.CodeGen.ImpCode.GPU qualified as Imp import Futhark.CodeGen.ImpGen import Futhark.CodeGen.ImpGen.GPU.Base-import Futhark.CodeGen.ImpGen.GPU.Group+import Futhark.CodeGen.ImpGen.GPU.Block import Futhark.IR.GPUMem import Futhark.Util.IntegralExp (divUp) import Prelude hiding (quot, rem)@@ -27,19 +27,19 @@    let (is, dims) = unzip $ unSegSpace space       dims' = map pe64 dims-      group_size' = pe64 <$> kAttrGroupSize attrs+      tblock_size' = pe64 <$> kAttrBlockSize attrs    emit $ Imp.DebugPrint "\n# SegMap" Nothing   case lvl of     SegThread {} -> do-      virt_num_groups <- dPrimVE "virt_num_groups" $ sExt32 $ product dims' `divUp` unCount group_size'+      virt_num_tblocks <- dPrimVE "virt_num_tblocks" $ sExt32 $ product dims' `divUp` unCount tblock_size'       sKernelThread "segmap" (segFlat space) attrs $-        virtualiseGroups (segVirt lvl) virt_num_groups $ \group_id -> do+        virtualiseBlocks (segVirt lvl) virt_num_tblocks $ \tblock_id -> do           local_tid <- kernelLocalThreadId . kernelConstants <$> askEnv            global_tid <-             dPrimVE "global_tid" $-              sExt64 group_id * sExt64 (unCount group_size')+              sExt64 tblock_id * sExt64 (unCount tblock_size')                 + sExt64 local_tid            dIndexSpace (zip is dims') global_tid@@ -48,17 +48,17 @@             compileStms mempty (kernelBodyStms kbody) $               zipWithM_ (compileThreadResult space) (patElems pat) $                 kernelBodyResult kbody-    SegGroup {} -> do-      pc <- precomputeConstants group_size' $ kernelBodyStms kbody-      virt_num_groups <- dPrimVE "virt_num_groups" $ sExt32 $ product dims'-      sKernelGroup "segmap_intragroup" (segFlat space) attrs $ do+    SegBlock {} -> do+      pc <- precomputeConstants tblock_size' $ kernelBodyStms kbody+      virt_num_tblocks <- dPrimVE "virt_num_tblocks" $ sExt32 $ product dims'+      sKernelBlock "segmap_intrablock" (segFlat space) attrs $ do         precomputedConstants pc $-          virtualiseGroups (segVirt lvl) virt_num_groups $ \group_id -> do-            dIndexSpace (zip is dims') $ sExt64 group_id+          virtualiseBlocks (segVirt lvl) virt_num_tblocks $ \tblock_id -> do+            dIndexSpace (zip is dims') $ sExt64 tblock_id              compileStms mempty (kernelBodyStms kbody) $-              zipWithM_ (compileGroupResult space) (patElems pat) $+              zipWithM_ (compileBlockResult space) (patElems pat) $                 kernelBodyResult kbody-    SegThreadInGroup {} ->-      error "compileSegMap: SegThreadInGroup"+    SegThreadInBlock {} ->+      error "compileSegMap: SegThreadInBlock"   emit $ Imp.DebugPrint "" Nothing

@@ -6,12 +6,12 @@ -- deviations are: -- -- * While we still use two-phase reduction, we use only a single---   kernel, with the final workgroup to write a result (tracked via+--   kernel, with the final threadblock to write a result (tracked via --   an atomic counter) performing the final reduction as well. -- -- * Instead of depending on storage layout transformations to handle --   non-commutative reductions efficiently, we slide a---   @groupsize@-sized window over the input, and perform a parallel+--   @tblocksize@-sized window over the input, and perform a parallel --   reduction for each window.  This sacrifices the notion of --   efficient sequentialisation, but is sometimes faster and --   definitely simpler and more predictable (and uses less auxiliary@@ -21,32 +21,32 @@ -- Regular Segmented Reductions on GPU" (FHPC '17).  This involves -- having two different strategies, and dynamically deciding which one -- to use based on the number of segments and segment size. We use the--- (static) @group_size@ to decide which of the following two+-- (static) @tblock_size@ to decide which of the following two -- strategies to choose: ----- * Large: uses one or more groups to process a single segment. If---   multiple groups are used per segment, the intermediate reduction+-- * Large: uses one or more blocks to process a single segment. If+--   multiple blocks are used per segment, the intermediate reduction --   results must be recursively reduced, until there is only a single --   value per segment. -- --   Each thread /can/ read multiple elements, which will greatly --   increase performance; however, if the reduction is --   non-commutative we will have to use a less efficient traversal---   (with interim group-wide reductions) to enable coalesced memory+--   (with interim block-wide reductions) to enable coalesced memory --   accesses, just as in the non-segmented case. ----- * Small: is used to let each group process *multiple* segments---   within a group. We will only use this approach when we can---   process at least two segments within a single group. In those---   cases, we would allocate a /whole/ group per segment with the---   large strategy, but at most 50% of the threads in the group would+-- * Small: is used to let each block process *multiple* segments+--   within a block. We will only use this approach when we can+--   process at least two segments within a single block. In those+--   cases, we would allocate a /whole/ block per segment with the+--   large strategy, but at most 50% of the threads in the block would --   have any element to read, which becomes highly inefficient. -- -- An optimization specfically targeted at non-segmented and large-segments -- segmented reductions with non-commutative is made: The stage one main loop is -- essentially stripmined by a factor `chunk`, inserting collective copies via--- local memory of each reduction parameter going into the intra-group (partial)--- reductions. This saves a factor `chunk` number of intra-group reductions at+-- local memory of each reduction parameter going into the intra-block (partial)+-- reductions. This saves a factor `chunk` number of intra-block reductions at -- the cost of some overhead in collective copies. module Futhark.CodeGen.ImpGen.GPU.SegRed   ( compileSegRed,@@ -76,7 +76,7 @@ -- | The maximum number of operators we support in a single SegRed. -- This limit arises out of the static allocation of counters. maxNumOps :: Int-maxNumOps = 10+maxNumOps = 20  -- | Code generation for the body of the SegRed, taking a continuation -- for saving the results of the body.  The results should be@@ -88,11 +88,11 @@ -- intermediate memory we need for the different ReduceKinds. data SegRedIntermediateArrays   = GeneralSegRedInterms-      { groupRedArrs :: [VName]+      { blockRedArrs :: [VName]       }   | NoncommPrimSegRedInterms       { collCopyArrs :: [VName],-        groupRedArrs :: [VName],+        blockRedArrs :: [VName],         privateChunks :: [VName]       } @@ -107,9 +107,9 @@   CallKernelGen () compileSegRed pat lvl space segbinops map_kbody = do   emit $ Imp.DebugPrint "\n# SegRed" Nothing-  KernelAttrs {kAttrNumGroups = num_groups, kAttrGroupSize = group_size} <-+  KernelAttrs {kAttrNumBlocks = num_tblocks, kAttrBlockSize = tblock_size} <-     lvlKernelAttrs lvl-  let grid = KernelGrid num_groups group_size+  let grid = KernelGrid num_tblocks tblock_size    compileSegRed' pat grid space segbinops $ \red_cont ->     sComment "apply map function" $@@ -143,7 +143,8 @@ compileSegRed' pat grid space segbinops map_body_cont   | genericLength segbinops > maxNumOps =       compilerLimitationS $-        "compileSegRed': at most " ++ show maxNumOps ++ " reduction operators are supported."+        ("compileSegRed': at most " <> show maxNumOps <> " reduction operators are supported.\n")+          <> ("Pattern: " <> prettyString pat)   | otherwise = do       chunk_v <- dPrimV "chunk_size" . isInt64 =<< kernelConstToExp chunk_const       case unSegSpace space of@@ -151,19 +152,19 @@           compileReduction (chunk_v, chunk_const) nonsegmentedReduction         _ -> do           let segment_size = pe64 $ last $ segSpaceDims space-              use_small_segments = segment_size * 2 .<. pe64 (unCount group_size) * tvExp chunk_v+              use_small_segments = segment_size * 2 .<. pe64 (unCount tblock_size) * tvExp chunk_v           sIf             use_small_segments             (compileReduction (chunk_v, chunk_const) smallSegmentsReduction)             (compileReduction (chunk_v, chunk_const) largeSegmentsReduction)   where     compileReduction chunk f =-      f pat num_groups group_size chunk space segbinops map_body_cont+      f pat num_tblocks tblock_size chunk space segbinops map_body_cont      param_types = map paramType $ concatMap paramOf segbinops -    num_groups = gridNumGroups grid-    group_size = gridGroupSize grid+    num_tblocks = gridNumBlocks grid+    tblock_size = gridBlockSize grid      chunk_const =       if Noncommutative `elem` map segBinOpComm segbinops@@ -186,34 +187,34 @@   SubExp ->   [SegBinOp GPUMem] ->   InKernelGen [SegRedIntermediateArrays]-makeIntermArrays group_id group_size chunk segbinops+makeIntermArrays tblock_id tblock_size chunk segbinops   | Noncommutative <- mconcat (map segBinOpComm segbinops),     all isPrimSegBinOp segbinops =       noncommPrimSegRedInterms   | otherwise =-      generalSegRedInterms group_id group_size segbinops+      generalSegRedInterms tblock_id tblock_size segbinops   where     params = map paramOf segbinops      noncommPrimSegRedInterms = do-      group_worksize <- tvSize <$> dPrimV "group_worksize" group_worksize_E+      block_worksize <- tvSize <$> dPrimV "block_worksize" block_worksize_E        -- compute total amount of lmem.-      let sum_ x y = nextMul x y + group_size_E * y-          group_reds_lmem_requirement = foldl sum_ 0 $ concat elem_sizes-          collcopy_lmem_requirement = group_worksize_E * max_elem_size+      let sum_ x y = nextMul x y + tblock_size_E * y+          block_reds_lmem_requirement = foldl sum_ 0 $ concat elem_sizes+          collcopy_lmem_requirement = block_worksize_E * max_elem_size           lmem_total_size =             Imp.bytes $-              collcopy_lmem_requirement `sMax64` group_reds_lmem_requirement+              collcopy_lmem_requirement `sMax64` block_reds_lmem_requirement -      -- offsets into the total pool of lmem for each group reduction array.+      -- offsets into the total pool of lmem for each block reduction array.       (_, offsets) <-         forAccumLM2D 0 elem_sizes $ \byte_offs elem_size ->           (,byte_offs `quot` elem_size)             <$> dPrimVE "offset" (sum_ byte_offs elem_size)        -- total pool of local mem.-      lmem <- sAlloc "local_mem" lmem_total_size (Space "local")+      lmem <- sAlloc "local_mem" lmem_total_size (Space "shared")       let arrInLMem ptype name len_se offset =             sArray               (name ++ "_" ++ prettyString ptype)@@ -223,20 +224,20 @@               $ LMAD.iota offset [pe64 len_se]        forM (zipWith zip params offsets) $ \ps_and_offsets -> do-        (coll_copy_arrs, group_red_arrs, priv_chunks) <-+        (coll_copy_arrs, block_red_arrs, priv_chunks) <-           fmap unzip3 $ forM ps_and_offsets $ \(p, offset) -> do             let ptype = elemType $ paramType p             (,,)-              <$> arrInLMem ptype "coll_copy_arr" group_worksize 0-              <*> arrInLMem ptype "group_red_arr" group_size offset+              <$> arrInLMem ptype "coll_copy_arr" block_worksize 0+              <*> arrInLMem ptype "block_red_arr" tblock_size offset               <*> sAllocArray                 ("chunk_" ++ prettyString ptype)                 ptype                 (Shape [chunk])                 (ScalarSpace [chunk] ptype)-        pure $ NoncommPrimSegRedInterms coll_copy_arrs group_red_arrs priv_chunks-    group_size_E = pe64 group_size-    group_worksize_E = group_size_E * pe64 chunk+        pure $ NoncommPrimSegRedInterms coll_copy_arrs block_red_arrs priv_chunks+    tblock_size_E = pe64 tblock_size+    block_worksize_E = tblock_size_E * pe64 chunk      paramSize = primByteSize . elemType . paramType     elem_sizes = map (map paramSize) params@@ -249,50 +250,50 @@   SubExp ->   [SegBinOp GPUMem] ->   InKernelGen [SegRedIntermediateArrays]-generalSegRedInterms group_id group_size segbinops =+generalSegRedInterms tblock_id tblock_size segbinops =   fmap (map GeneralSegRedInterms) $     forM (map paramOf segbinops) $       mapM $ \p ->         case paramDec p of           MemArray pt shape _ (ArrayIn mem _) -> do-            let shape' = Shape [group_size] <> shape+            let shape' = Shape [tblock_size] <> shape             let shape_E = map pe64 $ shapeDims shape'             sArray ("red_arr_" ++ prettyString pt) pt shape' mem $-              LMAD.iota (group_id * product shape_E) shape_E+              LMAD.iota (tblock_id * product shape_E) shape_E           _ -> do             let pt = elemType $ paramType p-                shape = Shape [group_size]-            sAllocArray ("red_arr_" ++ prettyString pt) pt shape $ Space "local"+                shape = Shape [tblock_size]+            sAllocArray ("red_arr_" ++ prettyString pt) pt shape $ Space "shared" --- | Arrays for storing group results.+-- | Arrays for storing block results. ----- The group-result arrays have an extra dimension because they are+-- The block-result arrays have an extra dimension because they are -- also used for keeping vectorised accumulators for first-stage -- reduction, if necessary.  If necessary, this dimension has size--- group_size, and otherwise 1.  When actually storing group results,+-- tblock_size, and otherwise 1.  When actually storing block results, -- the first index is set to 0. groupResultArrays ::   SubExp ->   SubExp ->   [SegBinOp GPUMem] ->   CallKernelGen [[VName]]-groupResultArrays num_virtgroups group_size segbinops =+groupResultArrays num_virtblocks tblock_size segbinops =   forM segbinops $ \(SegBinOp _ lam _ shape) ->     forM (lambdaReturnType lam) $ \t -> do       let pt = elemType t           extra_dim             | primType t, shapeRank shape == 0 = intConst Int64 1-            | otherwise = group_size-          full_shape = Shape [extra_dim, num_virtgroups] <> shape <> arrayShape t-          -- Move the groupsize dimension last to ensure coalesced+            | otherwise = tblock_size+          full_shape = Shape [extra_dim, num_virtblocks] <> shape <> arrayShape t+          -- Move the tblocksize dimension last to ensure coalesced           -- memory access.           perm = [1 .. shapeRank full_shape - 1] ++ [0]       sAllocArrayPerm "segred_tmp" pt full_shape (Space "device") perm  type DoCompileSegRed =   Pat LetDecMem ->-  Count NumGroups SubExp ->-  Count GroupSize SubExp ->+  Count NumBlocks SubExp ->+  Count BlockSize SubExp ->   (TV Int64, Imp.KernelConstExp) ->   SegSpace ->   [SegBinOp GPUMem] ->@@ -300,34 +301,34 @@   CallKernelGen ()  nonsegmentedReduction :: DoCompileSegRed-nonsegmentedReduction (Pat segred_pes) num_groups group_size (chunk_v, chunk_const) space segbinops map_body_cont = do+nonsegmentedReduction (Pat segred_pes) num_tblocks tblock_size (chunk_v, chunk_const) space segbinops map_body_cont = do   let (gtids, dims) = unzip $ unSegSpace space       chunk = tvExp chunk_v-      num_groups_se = unCount num_groups-      group_size_se = unCount group_size-      group_size' = pe64 group_size_se+      num_tblocks_se = unCount num_tblocks+      tblock_size_se = unCount tblock_size+      tblock_size' = pe64 tblock_size_se       global_tid = Imp.le64 $ segFlat space       n = pe64 $ last dims    counters <- genZeroes "counters" maxNumOps -  reds_group_res_arrs <- groupResultArrays num_groups_se group_size_se segbinops+  reds_block_res_arrs <- groupResultArrays num_tblocks_se tblock_size_se segbinops    num_threads <--    fmap tvSize $ dPrimV "num_threads" $ pe64 num_groups_se * group_size'+    fmap tvSize $ dPrimV "num_threads" $ pe64 num_tblocks_se * tblock_size'    let attrs =-        (defKernelAttrs num_groups group_size)+        (defKernelAttrs num_tblocks tblock_size)           { kAttrConstExps = M.singleton (tvVar chunk_v) chunk_const           }    sKernelThread "segred_nonseg" (segFlat space) attrs $ do     constants <- kernelConstants <$> askEnv     let ltid = kernelLocalThreadId constants-    let group_id = kernelGroupId constants+    let tblock_id = kernelBlockId constants -    interms <- makeIntermArrays (sExt64 group_id) group_size_se (tvSize chunk_v) segbinops-    sync_arr <- sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local"+    interms <- makeIntermArrays (sExt64 tblock_id) tblock_size_se (tvSize chunk_v) segbinops+    sync_arr <- sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "shared"      -- Since this is the nonsegmented case, all outer segment IDs must     -- necessarily be 0.@@ -336,8 +337,8 @@     q <- dPrimVE "q" $ n `divUp` (sExt64 (kernelNumThreads constants) * chunk)      slugs <--      mapM (segBinOpSlug ltid group_id) $-        zip3 segbinops interms reds_group_res_arrs+      mapM (segBinOpSlug ltid tblock_id) $+        zip3 segbinops interms reds_block_res_arrs     new_lambdas <-       reductionStageOne         gtids@@ -358,10 +359,10 @@       \(pes, slug, new_lambda, i) ->         reductionStageTwo           pes-          group_id+          tblock_id           [0]           0-          (sExt64 $ kernelNumGroups constants)+          (sExt64 $ kernelNumBlocks constants)           slug           new_lambda           counters@@ -369,7 +370,7 @@           (fromInteger i)  smallSegmentsReduction :: DoCompileSegRed-smallSegmentsReduction (Pat segred_pes) num_groups group_size _ space segbinops map_body_cont = do+smallSegmentsReduction (Pat segred_pes) num_tblocks tblock_size _ space segbinops map_body_cont = do   let (gtids, dims) = unzip $ unSegSpace space       dims' = map pe64 dims       segment_size = last dims'@@ -378,39 +379,39 @@   segment_size_nonzero <-     dPrimVE "segment_size_nonzero" $ sMax64 1 segment_size -  let group_size_se = unCount group_size-      num_groups_se = unCount group_size-      num_groups' = pe64 num_groups_se-      group_size' = pe64 group_size_se-  num_threads <- fmap tvSize $ dPrimV "num_threads" $ num_groups' * group_size'+  let tblock_size_se = unCount tblock_size+      num_tblocks_se = unCount tblock_size+      num_tblocks' = pe64 num_tblocks_se+      tblock_size' = pe64 tblock_size_se+  num_threads <- fmap tvSize $ dPrimV "num_threads" $ num_tblocks' * tblock_size'   let num_segments = product $ init dims'-      segments_per_group = group_size' `quot` segment_size_nonzero-      required_groups = sExt32 $ num_segments `divUp` segments_per_group+      segments_per_block = tblock_size' `quot` segment_size_nonzero+      required_blocks = sExt32 $ num_segments `divUp` segments_per_block    emit $ Imp.DebugPrint "# SegRed-small" Nothing   emit $ Imp.DebugPrint "num_segments" $ Just $ untyped num_segments   emit $ Imp.DebugPrint "segment_size" $ Just $ untyped segment_size-  emit $ Imp.DebugPrint "segments_per_group" $ Just $ untyped segments_per_group-  emit $ Imp.DebugPrint "required_groups" $ Just $ untyped required_groups+  emit $ Imp.DebugPrint "segments_per_block" $ Just $ untyped segments_per_block+  emit $ Imp.DebugPrint "required_blocks" $ Just $ untyped required_blocks -  sKernelThread "segred_small" (segFlat space) (defKernelAttrs num_groups group_size) $ do+  sKernelThread "segred_small" (segFlat space) (defKernelAttrs num_tblocks tblock_size) $ do     constants <- kernelConstants <$> askEnv-    let group_id = kernelGroupSize constants+    let tblock_id = kernelBlockSize constants         ltid = sExt64 $ kernelLocalThreadId constants -    interms <- generalSegRedInterms group_id group_size_se segbinops-    let reds_arrs = map groupRedArrs interms+    interms <- generalSegRedInterms tblock_id tblock_size_se segbinops+    let reds_arrs = map blockRedArrs interms -    -- We probably do not have enough actual workgroups to cover the-    -- entire iteration space.  Some groups thus have to perform double+    -- We probably do not have enough actual threadblocks to cover the+    -- entire iteration space.  Some blocks thus have to perform double     -- duty; we put an outer loop to accomplish this.-    virtualiseGroups SegVirt required_groups $ \virtgroup_id -> do+    virtualiseBlocks SegVirt required_blocks $ \virttblock_id -> do       -- Compute the 'n' input indices.  The outer 'n-1' correspond to-      -- the segment ID, and are computed from the group id.  The inner+      -- the segment ID, and are computed from the block id.  The inner       -- is computed from the local thread id, and may be out-of-bounds.       let segment_index =             (ltid `quot` segment_size_nonzero)-              + (sExt64 virtgroup_id * sExt64 segments_per_group)+              + (sExt64 virttblock_id * sExt64 segments_per_block)           index_within_segment = ltid `rem` segment_size        dIndexSpace (zip (init gtids) (init dims')) segment_index@@ -434,7 +435,7 @@               .&&. isActive (init $ zip gtids dims)               .&&. ltid               .<. segment_size-              * segments_per_group+              * segments_per_block           )           in_bounds           out_of_bounds@@ -445,10 +446,10 @@       sWhen (segment_size .>. 0) $         sComment "perform segmented scan to imitate reduction" $           forM2_ segbinops reds_arrs $ \(SegBinOp _ red_op _ _) red_arrs ->-            groupScan+            blockScan               (Just crossesSegment)               (sExt64 $ pe64 num_threads)-              (segment_size * segments_per_group)+              (segment_size * segments_per_block)               red_op               red_arrs @@ -456,18 +457,18 @@        sComment "save final values of segments"         $ sWhen-          ( sExt64 virtgroup_id-              * segments_per_group+          ( sExt64 virttblock_id+              * segments_per_block               + sExt64 ltid                 .<. num_segments                 .&&. ltid-                .<. segments_per_group+                .<. segments_per_block           )         $ forM2_ segred_pes (concat reds_arrs)         $ \pe arr -> do           -- Figure out which segment result this thread should write...           let flat_segment_index =-                sExt64 virtgroup_id * segments_per_group + sExt64 ltid+                sExt64 virttblock_id * segments_per_block + sExt64 ltid               gtids' =                 unflattenIndex (init dims') flat_segment_index           copyDWIMFix@@ -481,91 +482,91 @@       sOp $ Imp.Barrier Imp.FenceLocal  largeSegmentsReduction :: DoCompileSegRed-largeSegmentsReduction (Pat segred_pes) num_groups group_size (chunk_v, chunk_const) space segbinops map_body_cont = do+largeSegmentsReduction (Pat segred_pes) num_tblocks tblock_size (chunk_v, chunk_const) space segbinops map_body_cont = do   let (gtids, dims) = unzip $ unSegSpace space       dims' = map pe64 dims       num_segments = product $ init dims'       segment_size = last dims'-      num_groups' = pe64 $ unCount num_groups-      group_size_se = unCount group_size-      group_size' = pe64 group_size_se+      num_tblocks' = pe64 $ unCount num_tblocks+      tblock_size_se = unCount tblock_size+      tblock_size' = pe64 tblock_size_se       chunk = tvExp chunk_v -  groups_per_segment <--    dPrimVE "groups_per_segment" $-      num_groups' `divUp` sMax64 1 num_segments+  blocks_per_segment <-+    dPrimVE "blocks_per_segment" $+      num_tblocks' `divUp` sMax64 1 num_segments    q <-     dPrimVE "q" $-      segment_size `divUp` (group_size' * groups_per_segment * chunk)+      segment_size `divUp` (tblock_size' * blocks_per_segment * chunk) -  num_virtgroups <--    dPrimV "num_virtgroups" $-      groups_per_segment * num_segments+  num_virtblocks <-+    dPrimV "num_virtblocks" $+      blocks_per_segment * num_segments   threads_per_segment <-     dPrimVE "threads_per_segment" $-      groups_per_segment * group_size'+      blocks_per_segment * tblock_size'    emit $ Imp.DebugPrint "# SegRed-large" Nothing   emit $ Imp.DebugPrint "num_segments" $ Just $ untyped num_segments   emit $ Imp.DebugPrint "segment_size" $ Just $ untyped segment_size-  emit $ Imp.DebugPrint "num_virtgroups" $ Just $ untyped $ tvExp num_virtgroups-  emit $ Imp.DebugPrint "num_groups" $ Just $ untyped num_groups'-  emit $ Imp.DebugPrint "group_size" $ Just $ untyped group_size'+  emit $ Imp.DebugPrint "num_virtblocks" $ Just $ untyped $ tvExp num_virtblocks+  emit $ Imp.DebugPrint "num_tblocks" $ Just $ untyped num_tblocks'+  emit $ Imp.DebugPrint "tblock_size" $ Just $ untyped tblock_size'   emit $ Imp.DebugPrint "q" $ Just $ untyped q-  emit $ Imp.DebugPrint "groups_per_segment" $ Just $ untyped groups_per_segment+  emit $ Imp.DebugPrint "blocks_per_segment" $ Just $ untyped blocks_per_segment -  reds_group_res_arrs <- groupResultArrays (tvSize num_virtgroups) group_size_se segbinops+  reds_block_res_arrs <- groupResultArrays (tvSize num_virtblocks) tblock_size_se segbinops    -- In principle we should have a counter for every segment.  Since   -- the number of segments is a dynamic quantity, we would have to   -- allocate and zero out an array here, which is expensive.   -- However, we exploit the fact that the number of segments being   -- reduced at any point in time is limited by the number of-  -- workgroups. If we bound the number of workgroups, we can get away+  -- threadblocks. If we bound the number of threadblocks, we can get away   -- with using that many counters.  FIXME: Is this limit checked   -- anywhere?  There are other places in the compiler that will fail-  -- if the group count exceeds the maximum group size, which is at+  -- if the block count exceeds the maximum block size, which is at   -- most 1024 anyway.   let num_counters = maxNumOps * 1024   counters <- genZeroes "counters" $ fromIntegral num_counters    let attrs =-        (defKernelAttrs num_groups group_size)+        (defKernelAttrs num_tblocks tblock_size)           { kAttrConstExps = M.singleton (tvVar chunk_v) chunk_const           }    sKernelThread "segred_large" (segFlat space) attrs $ do     constants <- kernelConstants <$> askEnv-    let group_id = sExt64 $ kernelGroupId constants+    let tblock_id = sExt64 $ kernelBlockId constants         ltid = kernelLocalThreadId constants -    interms <- makeIntermArrays group_id group_size_se (tvSize chunk_v) segbinops-    sync_arr <- sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local"+    interms <- makeIntermArrays tblock_id tblock_size_se (tvSize chunk_v) segbinops+    sync_arr <- sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "shared" -    -- We probably do not have enough actual workgroups to cover the-    -- entire iteration space.  Some groups thus have to perform double+    -- We probably do not have enough actual threadblocks to cover the+    -- entire iteration space.  Some blocks thus have to perform double     -- duty; we put an outer loop to accomplish this.-    virtualiseGroups SegVirt (sExt32 (tvExp num_virtgroups)) $ \virtgroup_id -> do+    virtualiseBlocks SegVirt (sExt32 (tvExp num_virtblocks)) $ \virttblock_id -> do       let segment_gtids = init gtids        flat_segment_id <-         dPrimVE "flat_segment_id" $-          sExt64 virtgroup_id `quot` groups_per_segment+          sExt64 virttblock_id `quot` blocks_per_segment        global_tid <-         dPrimVE "global_tid" $-          (sExt64 virtgroup_id * sExt64 group_size' + sExt64 ltid)+          (sExt64 virttblock_id * sExt64 tblock_size' + sExt64 ltid)             `rem` threads_per_segment -      let first_group_for_segment = flat_segment_id * groups_per_segment+      let first_block_for_segment = flat_segment_id * blocks_per_segment       dIndexSpace (zip segment_gtids (init dims')) flat_segment_id       dPrim_ (last gtids) int64       let n = pe64 $ last dims        slugs <--        mapM (segBinOpSlug ltid virtgroup_id) $-          zip3 segbinops interms reds_group_res_arrs+        mapM (segBinOpSlug ltid virttblock_id) $+          zip3 segbinops interms reds_block_res_arrs       new_lambdas <-         reductionStageOne           gtids@@ -582,7 +583,7 @@               (map (length . segBinOpNeutral) segbinops)               segred_pes -          multiple_groups_per_segment =+          multiple_blocks_per_segment =             forM_ (zip4 segred_pess slugs new_lambdas [0 ..]) $               \(pes, slug, new_lambda, i) -> do                 let counter_idx =@@ -591,24 +592,24 @@                           `rem` fromIntegral num_counters                 reductionStageTwo                   pes-                  virtgroup_id+                  virttblock_id                   (map Imp.le64 segment_gtids)-                  first_group_for_segment-                  groups_per_segment+                  first_block_for_segment+                  blocks_per_segment                   slug                   new_lambda                   counters                   sync_arr                   counter_idx -          one_group_per_segment =-            sComment "first thread in group saves final result to memory" $+          one_block_per_segment =+            sComment "first thread in block saves final result to memory" $               forM2_ slugs segred_pess $ \slug pes ->                 sWhen (ltid .==. 0) $                   forM2_ pes (slugAccs slug) $ \v (acc, acc_is) ->                     copyDWIMFix (patElemName v) (map Imp.le64 segment_gtids) (Var acc) acc_is -      sIf (groups_per_segment .==. 1) one_group_per_segment multiple_groups_per_segment+      sIf (blocks_per_segment .==. 1) one_block_per_segment multiple_blocks_per_segment  -- | Auxiliary information for a single reduction. A slug holds the `SegBinOp` -- operator for a single reduction, the different arrays required throughout@@ -618,11 +619,11 @@   { slugOp :: SegBinOp GPUMem,     -- | Intermediate arrays needed for the given reduction.     slugInterms :: SegRedIntermediateArrays,-    -- | Place(s) to store group accumulator(s) in stage 1 reduction.+    -- | Place(s) to store block accumulator(s) in stage 1 reduction.     slugAccs :: [(VName, [Imp.TExp Int64])],     -- | Global memory destination(s) for the final result(s) for this     -- particular reduction.-    groupResArrs :: [VName]+    blockResArrs :: [VName]   }  segBinOpSlug ::@@ -630,19 +631,19 @@   Imp.TExp Int32 ->   (SegBinOp GPUMem, SegRedIntermediateArrays, [VName]) ->   InKernelGen SegBinOpSlug-segBinOpSlug ltid group_id (op, interms, group_res_arrs) = do-  accs <- zipWithM mkAcc (lambdaParams (segBinOpLambda op)) group_res_arrs-  pure $ SegBinOpSlug op interms accs group_res_arrs+segBinOpSlug ltid tblock_id (op, interms, block_res_arrs) = do+  accs <- zipWithM mkAcc (lambdaParams (segBinOpLambda op)) block_res_arrs+  pure $ SegBinOpSlug op interms accs block_res_arrs   where-    mkAcc p group_res_arr+    mkAcc p block_res_arr       | Prim t <- paramType p,         shapeRank (segBinOpShape op) == 0 = do-          group_res_acc <- dPrim (baseString (paramName p) <> "_group_res_acc") t-          pure (tvVar group_res_acc, [])+          block_res_acc <- dPrim (baseString (paramName p) <> "_block_res_acc") t+          pure (tvVar block_res_acc, [])       -- if this is a non-primitive reduction, the global mem result array will       -- double as accumulator.       | otherwise =-          pure (group_res_arr, [sExt64 ltid, sExt64 group_id])+          pure (block_res_arr, [sExt64 ltid, sExt64 tblock_id])  slugLambda :: SegBinOpSlug -> Lambda GPUMem slugLambda = segBinOpLambda . slugOp@@ -665,8 +666,8 @@ slugSplitParams :: SegBinOpSlug -> ([LParam GPUMem], [LParam GPUMem]) slugSplitParams slug = splitAt (length (slugNeutral slug)) $ slugParams slug -slugGroupRedArrs :: SegBinOpSlug -> [VName]-slugGroupRedArrs = groupRedArrs . slugInterms+slugBlockRedArrs :: SegBinOpSlug -> [VName]+slugBlockRedArrs = blockRedArrs . slugInterms  slugPrivChunks :: SegBinOpSlug -> [VName] slugPrivChunks = privateChunks . slugInterms@@ -691,32 +692,32 @@    dScope Nothing $ scopeOfLParams $ concatMap slugParams slugs -  sComment "ne-initialise the outer (per-group) accumulator(s)" $ do+  sComment "ne-initialise the outer (per-block) accumulator(s)" $ do     forM_ slugs $ \slug ->       forM2_ (slugAccs slug) (slugNeutral slug) $ \(acc, acc_is) ne ->         sLoopNest (slugShape slug) $ \vec_is ->           copyDWIMFix acc (acc_is ++ vec_is) ne []    new_lambdas <- mapM (renameLambda . slugLambda) slugs-  let group_size = sExt32 $ kernelGroupSize constants-  let doGroupReduce =+  let tblock_size = sExt32 $ kernelBlockSize constants+  let doBlockReduce =         forM2_ slugs new_lambdas $ \slug new_lambda -> do           let accs = slugAccs slug           let params = slugParams slug           sLoopNest (slugShape slug) $ \vec_is -> do-            let group_red_arrs = slugGroupRedArrs slug+            let block_red_arrs = slugBlockRedArrs slug             sComment "store accs. prims go in lmem; non-prims in params (in global mem)" $-              forM_ (zip3 group_red_arrs accs params) $+              forM_ (zip3 block_red_arrs accs params) $                 \(arr, (acc, acc_is), p) ->                   if isPrimParam p                     then copyDWIMFix arr [ltid] (Var acc) (acc_is ++ vec_is)                     else copyDWIMFix (paramName p) [] (Var acc) (acc_is ++ vec_is)              sOp $ Imp.ErrorSync Imp.FenceLocal -- Also implicitly barrier.-            groupReduce group_size new_lambda group_red_arrs+            blockReduce tblock_size new_lambda block_red_arrs             sOp $ Imp.Barrier Imp.FenceLocal -            sComment "thread 0 updates per-group acc(s); rest reset to ne" $ do+            sComment "thread 0 updates per-block acc(s); rest reset to ne" $ do               sIf                 (ltid .==. 0)                 ( forM2_ accs (lambdaParams new_lambda) $@@ -738,7 +739,7 @@         q         n         chunk-        doGroupReduce+        doBlockReduce     _ ->       generalStageOneBody         slugs@@ -748,7 +749,7 @@         q         n         threads_per_segment-        doGroupReduce+        doBlockReduce    pure new_lambdas @@ -762,24 +763,24 @@   Imp.TExp Int64 ->   InKernelGen () ->   InKernelGen ()-generalStageOneBody slugs body_cont glb_ind_var global_tid q n threads_per_segment doGroupReduce = do+generalStageOneBody slugs body_cont glb_ind_var global_tid q n threads_per_segment doBlockReduce = do   let is_comm = slugsComm slugs == Commutative    constants <- kernelConstants <$> askEnv-  let group_size = kernelGroupSize constants+  let tblock_size = kernelBlockSize constants       ltid = sExt64 $ kernelLocalThreadId constants -  -- this group's id within its designated segment, and this group's initial+  -- this block's id within its designated segment, and this block's initial   -- global offset.-  group_id_in_segment <- dPrimVE "group_id_in_segment" $ global_tid `quot` group_size-  group_base_offset <- dPrimVE "group_base_offset" $ group_id_in_segment * q * group_size+  tblock_id_in_segment <- dPrimVE "tblock_id_in_segment" $ global_tid `quot` tblock_size+  block_base_offset <- dPrimVE "block_base_offset" $ tblock_id_in_segment * q * tblock_size    sFor "i" q $ \i -> do-    group_offset <- dPrimVE "group_offset" $ group_base_offset + i * group_size+    block_offset <- dPrimVE "block_offset" $ block_base_offset + i * tblock_size     glb_ind_var       <-- if is_comm         then global_tid + threads_per_segment * i-        else group_offset + ltid+        else block_offset + ltid      sWhen (tvExp glb_ind_var .<. n) $       sComment "apply map function(s)" $@@ -804,9 +805,9 @@                 $ \(acc, acc_is) se ->                   copyDWIMFix acc (acc_is ++ vec_is) se [] -    unless is_comm doGroupReduce+    unless is_comm doBlockReduce   sOp $ Imp.ErrorSync Imp.FenceLocal-  when is_comm doGroupReduce+  when is_comm doBlockReduce  noncommPrimParamsStageOneBody ::   [SegBinOpSlug] ->@@ -818,25 +819,25 @@   Imp.TExp Int64 ->   InKernelGen () ->   InKernelGen ()-noncommPrimParamsStageOneBody slugs body_cont glb_ind_var global_tid q n chunk doLMemGroupReduce = do+noncommPrimParamsStageOneBody slugs body_cont glb_ind_var global_tid q n chunk doLMemBlockReduce = do   constants <- kernelConstants <$> askEnv-  let group_size = kernelGroupSize constants+  let tblock_size = kernelBlockSize constants       ltid = sExt64 $ kernelLocalThreadId constants -  -- this group's id within its designated segment; the stride made per group in-  -- the outer `i < q` loop; and this group's initial global offset.-  group_id_in_segment <- dPrimVE "group_offset_in_segment" $ global_tid `quot` group_size-  group_stride <- dPrimVE "group_stride" $ group_size * chunk-  group_base_offset <- dPrimVE "group_base_offset" $ group_id_in_segment * q * group_stride+  -- this block's id within its designated segment; the stride made per block in+  -- the outer `i < q` loop; and this block's initial global offset.+  tblock_id_in_segment <- dPrimVE "block_offset_in_segment" $ global_tid `quot` tblock_size+  block_stride <- dPrimVE "block_stride" $ tblock_size * chunk+  block_base_offset <- dPrimVE "block_base_offset" $ tblock_id_in_segment * q * block_stride    let chunkLoop = sFor "k" chunk    sFor "i" q $ \i -> do-    -- group offset in this iteration.-    group_offset <- dPrimVE "group_offset" $ group_base_offset + i * group_stride+    -- block offset in this iteration.+    block_offset <- dPrimVE "block_offset" $ block_base_offset + i * block_stride     chunkLoop $ \k -> do-      loc_ind <- dPrimVE "loc_ind" $ k * group_size + ltid-      glb_ind_var <-- group_offset + loc_ind+      loc_ind <- dPrimVE "loc_ind" $ k * tblock_size + ltid+      glb_ind_var <-- block_offset + loc_ind        sIf         (tvExp glb_ind_var .<. n)@@ -864,7 +865,7 @@          forM2_ coll_copy_arrs priv_chunks $ \lmem_arr priv_chunk -> do           chunkLoop $ \k -> do-            lmem_idx <- dPrimVE "lmem_idx" $ ltid + k * group_size+            lmem_idx <- dPrimVE "lmem_idx" $ ltid + k * tblock_size             copyDWIMFix lmem_arr [lmem_idx] (Var priv_chunk) [k]            sOp $ Imp.Barrier Imp.FenceLocal@@ -892,7 +893,7 @@             compileStms mempty (bodyStms $ slugBody slug) $               forM2_ accs binop_ress $ \acc binop_res ->                 copyDWIM acc [] binop_res []-    doLMemGroupReduce+    doLMemBlockReduce   sOp $ Imp.ErrorSync Imp.FenceLocal  reductionStageTwo ::@@ -907,27 +908,27 @@   VName ->   Imp.TExp Int64 ->   InKernelGen ()-reductionStageTwo segred_pes group_id segment_gtids first_group_for_segment groups_per_segment slug new_lambda counters sync_arr counter_idx = do+reductionStageTwo segred_pes tblock_id segment_gtids first_block_for_segment blocks_per_segment slug new_lambda counters sync_arr counter_idx = do   constants <- kernelConstants <$> askEnv    let ltid32 = kernelLocalThreadId constants       ltid = sExt64 ltid32-      group_size = kernelGroupSize constants+      tblock_size = kernelBlockSize constants    let (acc_params, next_params) = slugSplitParams slug-  let nes = slugNeutral slug-  let red_arrs = slugGroupRedArrs slug-  let group_res_arrs = groupResArrs slug+      nes = slugNeutral slug+      red_arrs = slugBlockRedArrs slug+      block_res_arrs = blockResArrs slug    old_counter <- dPrim "old_counter" int32   (counter_mem, _, counter_offset) <-     fullyIndexArray       counters       [counter_idx]-  sComment "first thread in group saves group result to global memory" $+  sComment "first thread in block saves block result to global memory" $     sWhen (ltid32 .==. 0) $ do-      forM_ (take (length nes) $ zip group_res_arrs (slugAccs slug)) $ \(v, (acc, acc_is)) ->-        copyDWIMFix v [0, sExt64 group_id] (Var acc) acc_is+      forM_ (take (length nes) $ zip block_res_arrs (slugAccs slug)) $ \(v, (acc, acc_is)) ->+        copyDWIMFix v [0, sExt64 tblock_id] (Var acc) acc_is       sOp $ Imp.MemFence Imp.FenceGlobal       -- Increment the counter, thus stating that our result is       -- available.@@ -939,17 +940,17 @@           counter_mem           counter_offset         $ untyped (1 :: Imp.TExp Int32)-      -- Now check if we were the last group to write our result.  If+      -- Now check if we were the last block to write our result.  If       -- so, it is our responsibility to produce the final result.-      sWrite sync_arr [0] $ untyped $ tvExp old_counter .==. groups_per_segment - 1+      sWrite sync_arr [0] $ untyped $ tvExp old_counter .==. blocks_per_segment - 1    sOp $ Imp.Barrier Imp.FenceGlobal -  is_last_group <- dPrim "is_last_group" Bool-  copyDWIMFix (tvVar is_last_group) [] (Var sync_arr) [0]-  sWhen (tvExp is_last_group) $ do-    -- The final group has written its result (and it was-    -- us!), so read in all the group results and perform the+  is_last_block <- dPrim "is_last_block" Bool+  copyDWIMFix (tvVar is_last_block) [] (Var sync_arr) [0]+  sWhen (tvExp is_last_block) $ do+    -- The final block has written its result (and it was+    -- us!), so read in all the block results and perform the     -- final stage of the reduction.  But first, we reset the     -- counter so it is ready for next time.  This is done     -- with an atomic to avoid warnings about write/write@@ -959,49 +960,49 @@         Imp.Atomic DefaultSpace $           Imp.AtomicAdd Int32 (tvVar old_counter) counter_mem counter_offset $             untyped $-              negate groups_per_segment+              negate blocks_per_segment      sLoopNest (slugShape slug) $ \vec_is -> do       unless (null $ slugShape slug) $         sOp (Imp.Barrier Imp.FenceLocal) -      -- There is no guarantee that the number of workgroups for the-      -- segment is less than the workgroup size, so each thread may+      -- There is no guarantee that the number of threadblocks for the+      -- segment is less than the threadblock size, so each thread may       -- have to read multiple elements.  We do this in a sequential       -- way that may induce non-coalesced accesses, but the total       -- number of accesses should be tiny here.       ---      -- TODO: here we *could* insert a collective copy of the num_groups-      -- per-group results. However, it may not be beneficial, since num_groups-      -- is not necessarily larger than group_size, meaning the number of+      -- TODO: here we *could* insert a collective copy of the num_tblocks+      -- per-block results. However, it may not be beneficial, since num_tblocks+      -- is not necessarily larger than tblock_size, meaning the number of       -- uncoalesced reads here may be insignificant. In fact, if we happen to-      -- have a num_groups < group_size, then the collective copy would add+      -- have a num_tblocks < tblock_size, then the collective copy would add       -- unnecessary overhead. Also, this code is only executed by a single-      -- group.-      sComment "read in the per-group-results" $ do+      -- block.+      sComment "read in the per-block-results" $ do         read_per_thread <-           dPrimVE "read_per_thread" $-            groups_per_segment `divUp` sExt64 group_size+            blocks_per_segment `divUp` sExt64 tblock_size          forM2_ acc_params nes $ \p ne ->           copyDWIM (paramName p) [] ne []          sFor "i" read_per_thread $ \i -> do-          group_res_id <--            dPrimVE "group_res_id" $+          block_res_id <-+            dPrimVE "block_res_id" $               ltid * read_per_thread + i-          index_of_group_res <--            dPrimVE "index_of_group_res" $-              first_group_for_segment + group_res_id+          index_of_block_res <-+            dPrimVE "index_of_block_res" $+              first_block_for_segment + block_res_id -          sWhen (group_res_id .<. groups_per_segment) $ do-            forM2_ next_params group_res_arrs $-              \p group_res_arr ->+          sWhen (block_res_id .<. blocks_per_segment) $ do+            forM2_ next_params block_res_arrs $+              \p block_res_arr ->                 copyDWIMFix                   (paramName p)                   []-                  (Var group_res_arr)-                  ([0, index_of_group_res] ++ vec_is)+                  (Var block_res_arr)+                  ([0, index_of_block_res] ++ vec_is)              compileStms mempty (bodyStms $ slugBody slug) $               forM2_ acc_params (map resSubExp $ bodyResult $ slugBody slug) $ \p se ->@@ -1013,8 +1014,8 @@        sOp $ Imp.ErrorSync Imp.FenceLocal -      sComment "reduce the per-group results" $ do-        groupReduce (sExt32 group_size) new_lambda red_arrs+      sComment "reduce the per-block results" $ do+        blockReduce (sExt32 tblock_size) new_lambda red_arrs          sComment "and back to memory with the final result" $           sWhen (ltid32 .==. 0) $@@ -1031,8 +1032,8 @@ -- reductions with all primitive parameters: -- --   These kernels need local memory for 1) the initial collective copy, 2) the---   (virtualized) group reductions, and (TODO: this one not implemented yet!)---   3) the final single-group collective copy. There are no dependencies+--   (virtualized) block reductions, and (TODO: this one not implemented yet!)+--   3) the final single-block collective copy. There are no dependencies --   between these three stages, so we can reuse the same pool of local mem for --   all three. These intermediates all go into local mem because of the --   assumption of primitive parameter types.@@ -1042,26 +1043,26 @@ --   across the three steps are: -- --   1) The initial collective copy from global to thread-private memory---   requires `group_size * CHUNK * max elem_sizes`, since the collective copies+--   requires `tblock_size * CHUNK * max elem_sizes`, since the collective copies --   are performed in sequence (ie. inputs to different reduction operators need --   not be held in local mem simultaneously).---   2) The intra-group reductions of local memory held per-thread results---   require `group_size * sum elem_sizes` bytes, since per-thread results for---   all fused reductions are group-reduced simultaneously.---   3) If group_size < num_groups, then after the final single-group collective---   copy, a thread-sequential reduction reduces the number of per-group partial---   results from num_groups down to group_size for each reduction array, such---   that they will each fit in the final intra-group reduction. This requires---   `num_groups * max elem_sizes`.+--   2) The intra-block reductions of local memory held per-thread results+--   require `tblock_size * sum elem_sizes` bytes, since per-thread results for+--   all fused reductions are block-reduced simultaneously.+--   3) If tblock_size < num_tblocks, then after the final single-block collective+--   copy, a thread-sequential reduction reduces the number of per-block partial+--   results from num_tblocks down to tblock_size for each reduction array, such+--   that they will each fit in the final intra-block reduction. This requires+--   `num_tblocks * max elem_sizes`. -- --   In summary, the total amount of local mem needed is the maximum between:---   1) initial collective copy: group_size * CHUNK * max elem_sizes---   2) intra-group reductions:  group_size * sum elem_sizes---   3) final collective copy:   num_groups * max elem_sizes+--   1) initial collective copy: tblock_size * CHUNK * max elem_sizes+--   2) intra-block reductions:  tblock_size * sum elem_sizes+--   3) final collective copy:   num_tblocks * max elem_sizes -- --   The amount of local mem will most likely be decided by 1) in most cases, --   unless the number of fused operators is very high *or* if we have a---   `num_groups > group_size * CHUNK`, but this is unlikely, in which case 2)+--   `num_tblocks > tblock_size * CHUNK`, but this is unlikely, in which case 2) --   and 3), respectively, will dominate. -- --   Aside from local memory, these kernels also require a CHUNK-sized array of@@ -1070,6 +1071,6 @@ -- For all other reductions, ie. commutative reductions, reductions with at -- least one non-primitive operator, and small segments reductions: -----   These kernels use local memory only for the intra-group reductions, and---   since they do not use chunking or CHUNK, they all require onlly `group_size+--   These kernels use local memory only for the intra-block reductions, and+--   since they do not use chunking or CHUNK, they all require onlly `tblock_size --   * max elem_sizes` bytes of local memory and no thread-private register mem.

@@ -26,31 +26,31 @@   drop (length (segBinOpNeutral scan)) (lambdaParams (segBinOpLambda scan))  createLocalArrays ::-  Count GroupSize SubExp ->+  Count BlockSize SubExp ->   SubExp ->   [PrimType] ->   InKernelGen (VName, [VName], [VName], VName, [VName])-createLocalArrays (Count groupSize) chunk types = do-  let groupSizeE = pe64 groupSize-      workSize = pe64 chunk * groupSizeE+createLocalArrays (Count block_size) chunk types = do+  let block_sizeE = pe64 block_size+      workSize = pe64 chunk * block_sizeE       prefixArraysSize =-        foldl (\acc tySize -> nextMul acc tySize + tySize * groupSizeE) 0 $+        foldl (\acc tySize -> nextMul acc tySize + tySize * block_sizeE) 0 $           map primByteSize types       maxTransposedArraySize =         foldl1 sMax64 $ map (\ty -> workSize * primByteSize ty) types -      warpSize :: (Num a) => a-      warpSize = 32+      warp_size :: (Num a) => a+      warp_size = 32       maxWarpExchangeSize =-        foldl (\acc tySize -> nextMul acc tySize + tySize * fromInteger warpSize) 0 $+        foldl (\acc tySize -> nextMul acc tySize + tySize * fromInteger warp_size) 0 $           map primByteSize types-      maxLookbackSize = maxWarpExchangeSize + warpSize+      maxLookbackSize = maxWarpExchangeSize + warp_size       size = Imp.bytes $ maxLookbackSize `sMax64` prefixArraysSize `sMax64` maxTransposedArraySize    (_, byteOffsets) <-     mapAccumLM       ( \off tySize -> do-          off' <- dPrimVE "byte_offsets" $ nextMul off tySize + pe64 groupSize * tySize+          off' <- dPrimVE "byte_offsets" $ nextMul off tySize + pe64 block_size * tySize           pure (off', off)       )       0@@ -59,15 +59,15 @@   (_, warpByteOffsets) <-     mapAccumLM       ( \off tySize -> do-          off' <- dPrimVE "warp_byte_offset" $ nextMul off tySize + warpSize * tySize+          off' <- dPrimVE "warp_byte_offset" $ nextMul off tySize + warp_size * tySize           pure (off', off)       )-      warpSize+      warp_size       $ map primByteSize types    sComment "Allocate reusable shared memory" $ pure () -  localMem <- sAlloc "local_mem" size (Space "local")+  localMem <- sAlloc "local_mem" size (Space "shared")   transposeArrayLength <- dPrimV "trans_arr_len" workSize    sharedId <- sArrayInMem "shared_id" int32 (Shape [constant (1 :: Int32)]) localMem@@ -86,20 +86,20 @@       sArray         "local_prefix_arr"         ty-        (Shape [groupSize])+        (Shape [block_size])         localMem-        $ LMAD.iota off' [pe64 groupSize]+        $ LMAD.iota off' [pe64 block_size] -  warpscan <- sArrayInMem "warpscan" int8 (Shape [constant (warpSize :: Int64)]) localMem+  warpscan <- sArrayInMem "warpscan" int8 (Shape [constant (warp_size :: Int64)]) localMem   warpExchanges <-     forM (zip warpByteOffsets types) $ \(off, ty) -> do       let off' = off `quot` primByteSize ty       sArray         "warp_exchange"         ty-        (Shape [constant (warpSize :: Int64)])+        (Shape [constant (warp_size :: Int64)])         localMem-        $ LMAD.iota off' [warpSize]+        $ LMAD.iota off' [warp_size]    pure (sharedId, transposedArrays, prefixArrays, warpscan, warpExchanges) @@ -232,19 +232,19 @@        tys = map elemType tys' -      group_size_e = pe64 $ unCount $ kAttrGroupSize attrs-      num_physgroups_e = pe64 $ unCount $ kAttrNumGroups attrs+      tblock_size_e = pe64 $ unCount $ kAttrBlockSize attrs+      num_phys_blocks_e = pe64 $ unCount $ kAttrNumBlocks attrs    let chunk_const = getChunkSize tys'   chunk_v <- dPrimV "chunk_size" . isInt64 =<< kernelConstToExp chunk_const   let chunk = tvExp chunk_v -  num_virtgroups <--    tvSize <$> dPrimV "num_virtgroups" (n `divUp` (group_size_e * chunk))-  let num_virtgroups_e = pe64 num_virtgroups+  num_virt_blocks <-+    tvSize <$> dPrimV "num_virt_blocks" (n `divUp` (tblock_size_e * chunk))+  let num_virt_blocks_e = pe64 num_virt_blocks    num_virt_threads <--    dPrimVE "num_virt_threads" $ num_virtgroups_e * group_size_e+    dPrimVE "num_virt_threads" $ num_virt_blocks_e * tblock_size_e    let (gtids, dims) = unzip $ unSegSpace space       dims' = map pe64 dims@@ -254,15 +254,15 @@    emit $ Imp.DebugPrint "Sequential elements per thread (chunk)" $ Just $ untyped chunk -  statusFlags <- sAllocArray "status_flags" int8 (Shape [num_virtgroups]) (Space "device")+  statusFlags <- sAllocArray "status_flags" int8 (Shape [num_virt_blocks]) (Space "device")   sReplicate statusFlags $ intConst Int8 statusX    (aggregateArrays, incprefixArrays) <-     fmap unzip $       forM tys $ \ty ->         (,)-          <$> sAllocArray "aggregates" ty (Shape [num_virtgroups]) (Space "device")-          <*> sAllocArray "incprefixes" ty (Shape [num_virtgroups]) (Space "device")+          <$> sAllocArray "aggregates" ty (Shape [num_virt_blocks]) (Space "device")+          <*> sAllocArray "incprefixes" ty (Shape [num_virt_blocks]) (Space "device")    global_id <- genZeroes "global_dynid" 1 @@ -277,22 +277,23 @@         ltid = sExt64 ltid32      (sharedId, transposedArrays, prefixArrays, warpscan, exchanges) <--      createLocalArrays (kAttrGroupSize attrs) (tvSize chunk_v) tys+      createLocalArrays (kAttrBlockSize attrs) (tvSize chunk_v) tys -    -- We wrap the entire kernel body in a virtualisation loop to handle the-    -- case where we do not have enough workgroups to cover the iteration space.-    -- Dynamic group indexing has no implication on this, since each group-    -- simply fetches a new dynamic ID upon entry into the virtualisation loop.+    -- We wrap the entire kernel body in a virtualisation loop to+    -- handle the case where we do not have enough thread blocks to+    -- cover the iteration space. Dynamic block indexing has no+    -- implication on this, since each block simply fetches a new+    -- dynamic ID upon entry into the virtualisation loop.     ---    -- We could use virtualiseGroups, but this introduces a barrier which is-    -- redundant in this case, and also we don't need to base virtual group IDs+    -- We could use virtualiseBlocks, but this introduces a barrier which is+    -- redundant in this case, and also we don't need to base virtual block IDs     -- on the loop variable, but rather on the dynamic IDs.-    physgroup_id <- dPrim "physgroup_id" int32-    sOp $ Imp.GetGroupId (tvVar physgroup_id) 0+    phys_block_id <- dPrim "phys_block_id" int32+    sOp $ Imp.GetBlockId (tvVar phys_block_id) 0     iters <-       dPrimVE "virtloop_bound" $-        (num_virtgroups_e - tvExp physgroup_id)-          `divUp` num_physgroups_e+        (num_virt_blocks_e - tvExp phys_block_id)+          `divUp` num_phys_blocks_e      sFor "virtloop_i" iters $ const $ do       dyn_id <- dPrim "dynamic_id" int32@@ -311,7 +312,7 @@             copyDWIMFix sharedId [0] (tvSize dyn_id) []            sComment "First thread in last (virtual) block resets global dynamic_id" $ do-            sWhen (tvExp dyn_id .==. num_virtgroups_e - 1) $+            sWhen (tvExp dyn_id .==. num_virt_blocks_e - 1) $               copyDWIMFix global_id [0] (intConst Int32 0) []        let local_barrier = Imp.Barrier Imp.FenceLocal@@ -324,16 +325,16 @@        block_offset <-         dPrimVE "block_offset" $-          sExt64 (tvExp dyn_id) * chunk * group_size_e+          sExt64 (tvExp dyn_id) * chunk * tblock_size_e       sgm_idx <- dPrimVE "sgm_idx" $ block_offset `mod` segment_size       boundary <-         dPrimVE "boundary" $           sExt32 $-            sMin64 (chunk * group_size_e) (segment_size - sgm_idx)+            sMin64 (chunk * tblock_size_e) (segment_size - sgm_idx)       segsize_compact <-         dPrimVE "segsize_compact" $           sExt32 $-            sMin64 (chunk * group_size_e) segment_size+            sMin64 (chunk * tblock_size_e) segment_size       private_chunks <-         forM tys $ \ty ->           sAllocArray@@ -347,7 +348,7 @@       sComment "Load and map" $         sFor "i" chunk $ \i -> do           -- The map's input index-          virt_tid <- dPrimVE "virt_tid" $ thd_offset + i * group_size_e+          virt_tid <- dPrimVE "virt_tid" $ thd_offset + i * tblock_size_e           dIndexSpace (zip gtids dims') virt_tid           -- Perform the map           let in_bounds =@@ -373,7 +374,7 @@       sComment "Transpose scan inputs" $ do         forM_ (zip transposedArrays private_chunks) $ \(trans, priv) -> do           sFor "i" chunk $ \i -> do-            sharedIdx <- dPrimVE "sharedIdx" $ ltid + i * group_size_e+            sharedIdx <- dPrimVE "sharedIdx" $ ltid + i * tblock_size_e             copyDWIMFix trans [sharedIdx] (Var priv) [i]           sOp local_barrier           sFor "i" chunk $ \i -> do@@ -422,9 +423,9 @@        accs <- mapM (dPrim "acc") tys       sComment "Scan results (with warp scan)" $ do-        groupScan+        blockScan           crossesSegment-          group_size_e+          tblock_size_e           num_virt_threads           scan_op1           prefixArrays@@ -432,7 +433,7 @@         sOp $ Imp.ErrorSync Imp.FenceLocal          let firstThread acc prefixes =-              copyDWIMFix (tvVar acc) [] (Var prefixes) [sExt64 group_size_e - 1]+              copyDWIMFix (tvVar acc) [] (Var prefixes) [sExt64 tblock_size_e - 1]             notFirstThread acc prefixes =               copyDWIMFix (tvVar acc) [] (Var prefixes) [ltid - 1]         sIf@@ -457,11 +458,11 @@             copyDWIMFix (tvVar acc) [] ne []         -- end sWhen -        let warpSize = kernelWaveSize constants-        sWhen (bNot blockNewSgm .&&. ltid32 .<. warpSize) $ do+        let warp_size = kernelWaveSize constants+        sWhen (bNot blockNewSgm .&&. ltid32 .<. warp_size) $ do           sWhen (ltid32 .==. 0) $ do             sIf-              (not_segmented_e .||. boundary .==. sExt32 (group_size_e * chunk))+              (not_segmented_e .||. boundary .==. sExt32 (tblock_size_e * chunk))               ( do                   everythingVolatile $                     forM_ (zip aggregateArrays accs) $ \(aggregateArray, acc) ->@@ -498,10 +499,10 @@                   dPrimV "readOffset" $                     sExt32 $                       tvExp dyn_id - sExt64 (kernelWaveSize constants)-                let loopStop = warpSize * (-1)+                let loopStop = warp_size * (-1)                     sameSegment readIdx                       | segmented =-                          let startIdx = sExt64 (tvExp readIdx + 1) * group_size_e * chunk - 1+                          let startIdx = sExt64 (tvExp readIdx + 1) * tblock_size_e * chunk - 1                            in block_offset - startIdx .<=. sgm_idx                       | otherwise = true                 sWhile (tvExp readOffset .>. loopStop) $ do@@ -533,7 +534,7 @@                   copyDWIMFix warpscan [ltid] (tvSize flag) []                    -- execute warp-parallel reduction but only if the last read flag in not STATUS_P-                  copyDWIMFix (tvVar flag) [] (Var warpscan) [sExt64 warpSize - 1]+                  copyDWIMFix (tvVar flag) [] (Var warpscan) [sExt64 warp_size - 1]                   sWhen (tvExp flag .<. statusP) $ do                     lam' <- renameLambda scan_op1                     inBlockScanLookback@@ -544,15 +545,15 @@                       lam'                    -- all threads of the warp read the result of reduction-                  copyDWIMFix (tvVar flag) [] (Var warpscan) [sExt64 warpSize - 1]+                  copyDWIMFix (tvVar flag) [] (Var warpscan) [sExt64 warp_size - 1]                   forM_ (zip aggrs exchanges) $ \(aggr, exchange) ->-                    copyDWIMFix (tvVar aggr) [] (Var exchange) [sExt64 warpSize - 1]+                    copyDWIMFix (tvVar aggr) [] (Var exchange) [sExt64 warp_size - 1]                   -- update read offset                   sIf                     (tvExp flag .==. statusP)                     (readOffset <-- loopStop)                     ( sWhen (tvExp flag .==. statusA) $ do-                        readOffset <-- tvExp readOffset - zExt32 warpSize+                        readOffset <-- tvExp readOffset - zExt32 warp_size                     )                    -- update prefix if flag different than STATUS_X:@@ -573,7 +574,7 @@             scan_op2 <- renameLambda scan_op1             let xs = map paramName $ take (length tys) $ lambdaParams scan_op2                 ys = map paramName $ drop (length tys) $ lambdaParams scan_op2-            sWhen (boundary .==. sExt32 (group_size_e * chunk)) $ do+            sWhen (boundary .==. sExt32 (tblock_size_e * chunk)) $ do               forM_ (zip xs prefixes) $ \(x, prefix) -> dPrimV_ x $ tvExp prefix               forM_ (zip ys accs) $ \(y, acc) -> dPrimV_ y $ tvExp acc               compileStms mempty (bodyStms $ lambdaBody scan_op2) $@@ -643,7 +644,7 @@             copyDWIMFix locmem [tvExp sharedIdx] (Var priv) [i]           sOp local_barrier           sFor "i" chunk $ \i -> do-            flat_idx <- dPrimVE "flat_idx" $ thd_offset + i * group_size_e+            flat_idx <- dPrimVE "flat_idx" $ thd_offset + i * tblock_size_e             dIndexSpace (zip gtids dims') flat_idx             sWhen (flat_idx .<. n) $ do               copyDWIMFix

@@ -21,15 +21,15 @@ -- Aggressively try to reuse memory for different SegBinOps, because -- we will run them sequentially after another. makeLocalArrays ::-  Count GroupSize SubExp ->+  Count BlockSize SubExp ->   SubExp ->   [SegBinOp GPUMem] ->   InKernelGen [[VName]]-makeLocalArrays (Count group_size) num_threads scans = do+makeLocalArrays (Count tblock_size) num_threads scans = do   (arrs, mems_and_sizes) <- runStateT (mapM onScan scans) mempty   let maxSize sizes = Imp.bytes $ foldl' sMax64 1 $ map Imp.unCount sizes   forM_ mems_and_sizes $ \(sizes, mem) ->-    sAlloc_ mem (maxSize sizes) (Space "local")+    sAlloc_ mem (maxSize sizes) (Space "shared")   pure arrs   where     onScan (SegBinOp _ scan_op nes _) = do@@ -46,7 +46,7 @@               pure (arr, [])             _ -> do               let pt = elemType $ paramType p-                  shape = Shape [group_size]+                  shape = Shape [tblock_size]               (sizes, mem') <- getMem pt shape               arr <- lift $ sArrayInMem "scan_arr" pt shape mem'               pure (arr, [(sizes, mem')])@@ -64,7 +64,7 @@           put mems'           pure (size : size', mem)         (Nothing, []) -> do-          mem <- lift $ sDeclareMem "scan_arr_mem" $ Space "local"+          mem <- lift $ sDeclareMem "scan_arr_mem" $ Space "shared"           pure ([size], mem)  type CrossesSegment = Maybe (Imp.TExp Int64 -> Imp.TExp Int64 -> Imp.TExp Bool)@@ -143,25 +143,25 @@   | otherwise =       pure () --- | Produce partially scanned intervals; one per workgroup.+-- | Produce partially scanned intervals; one per threadblock. scanStage1 ::   Pat LetDecMem ->-  Count NumGroups SubExp ->-  Count GroupSize SubExp ->+  Count NumBlocks SubExp ->+  Count BlockSize SubExp ->   SegSpace ->   [SegBinOp GPUMem] ->   KernelBody GPUMem ->   CallKernelGen (TV Int32, Imp.TExp Int64, CrossesSegment)-scanStage1 (Pat all_pes) num_groups group_size space scans kbody = do-  let num_groups' = fmap pe64 num_groups-      group_size' = fmap pe64 group_size-  num_threads <- dPrimV "num_threads" $ sExt32 $ unCount num_groups' * unCount group_size'+scanStage1 (Pat all_pes) num_tblocks tblock_size space scans kbody = do+  let num_tblocks' = fmap pe64 num_tblocks+      tblock_size' = fmap pe64 tblock_size+  num_threads <- dPrimV "num_threads" $ sExt32 $ unCount num_tblocks' * unCount tblock_size'    let (gtids, dims) = unzip $ unSegSpace space       dims' = map pe64 dims   let num_elements = product dims'       elems_per_thread = num_elements `divUp` sExt64 (tvExp num_threads)-      elems_per_group = unCount group_size' * elems_per_thread+      elems_per_group = unCount tblock_size' * elems_per_thread    let crossesSegment =         case reverse dims' of@@ -169,9 +169,9 @@             (to - from) .>. (to `rem` segment_size)           _ -> Nothing -  sKernelThread "scan_stage1" (segFlat space) (defKernelAttrs num_groups group_size) $ do+  sKernelThread "scan_stage1" (segFlat space) (defKernelAttrs num_tblocks tblock_size) $ do     constants <- kernelConstants <$> askEnv-    all_local_arrs <- makeLocalArrays group_size (tvSize num_threads) scans+    all_local_arrs <- makeLocalArrays tblock_size (tvSize num_threads) scans      -- The variables from scan_op will be used for the carry and such     -- in the big chunking loop.@@ -183,8 +183,8 @@     sFor "j" elems_per_thread $ \j -> do       chunk_offset <-         dPrimV "chunk_offset" $-          sExt64 (kernelGroupSize constants) * j-            + sExt64 (kernelGroupId constants) * elems_per_group+          sExt64 (kernelBlockSize constants) * j+            + sExt64 (kernelBlockId constants) * elems_per_group       flat_idx <-         dPrimV "flat_idx" $           tvExp chunk_offset + sExt64 (kernelLocalThreadId constants)@@ -259,10 +259,10 @@                -- We need to avoid parameter name clashes.               scan_op_renamed <- renameLambda scan_op-              groupScan+              blockScan                 crossesSegment'                 (sExt64 $ tvExp num_threads)-                (sExt64 $ kernelGroupSize constants)+                (sExt64 $ kernelBlockSize constants)                 scan_op_renamed                 local_arrs @@ -284,10 +284,10 @@                         []                         (Var arr)                         [ if primType $ paramType p-                            then sExt64 (kernelGroupSize constants) - 1+                            then sExt64 (kernelBlockSize constants) - 1                             else-                              (sExt64 (kernelGroupId constants) + 1)-                                * sExt64 (kernelGroupSize constants)+                              (sExt64 (kernelBlockId constants) + 1)+                                * sExt64 (kernelBlockSize constants)                                 - 1                         ]                   load_neutral =@@ -301,11 +301,11 @@                     Just f ->                       f                         ( tvExp chunk_offset-                            + sExt64 (kernelGroupSize constants)+                            + sExt64 (kernelBlockSize constants)                             - 1                         )                         ( tvExp chunk_offset-                            + sExt64 (kernelGroupSize constants)+                            + sExt64 (kernelBlockSize constants)                         )                 should_load_carry <-                   dPrimVE "should_load_carry" $@@ -322,17 +322,17 @@   Pat LetDecMem ->   TV Int32 ->   Imp.TExp Int64 ->-  Count NumGroups SubExp ->+  Count NumBlocks SubExp ->   CrossesSegment ->   SegSpace ->   [SegBinOp GPUMem] ->   CallKernelGen ()-scanStage2 (Pat all_pes) stage1_num_threads elems_per_group num_groups crossesSegment space scans = do+scanStage2 (Pat all_pes) stage1_num_threads elems_per_group num_tblocks crossesSegment space scans = do   let (gtids, dims) = unzip $ unSegSpace space       dims' = map pe64 dims    -- Our group size is the number of groups for the stage 1 kernel.-  let group_size = Count $ unCount num_groups+  let tblock_size = Count $ unCount num_tblocks    let crossesSegment' = do         f <- crossesSegment@@ -341,9 +341,9 @@             ((sExt64 from + 1) * elems_per_group - 1)             ((sExt64 to + 1) * elems_per_group - 1) -  sKernelThread "scan_stage2" (segFlat space) (defKernelAttrs (Count (intConst Int64 1)) group_size) $ do+  sKernelThread "scan_stage2" (segFlat space) (defKernelAttrs (Count (intConst Int64 1)) tblock_size) $ do     constants <- kernelConstants <$> askEnv-    per_scan_local_arrs <- makeLocalArrays group_size (tvSize stage1_num_threads) scans+    per_scan_local_arrs <- makeLocalArrays tblock_size (tvSize stage1_num_threads) scans     let per_scan_rets = map (lambdaReturnType . segBinOpLambda) scans         per_scan_pes = segBinOpChunks scans all_pes @@ -378,10 +378,10 @@            barrier -          groupScan+          blockScan             crossesSegment'             (sExt64 $ tvExp stage1_num_threads)-            (sExt64 $ kernelGroupSize constants)+            (sExt64 $ kernelBlockSize constants)             scan_op             local_arrs @@ -396,30 +396,30 @@  scanStage3 ::   Pat LetDecMem ->-  Count NumGroups SubExp ->-  Count GroupSize SubExp ->+  Count NumBlocks SubExp ->+  Count BlockSize SubExp ->   Imp.TExp Int64 ->   CrossesSegment ->   SegSpace ->   [SegBinOp GPUMem] ->   CallKernelGen ()-scanStage3 (Pat all_pes) num_groups group_size elems_per_group crossesSegment space scans = do-  let group_size' = fmap pe64 group_size+scanStage3 (Pat all_pes) num_tblocks tblock_size elems_per_group crossesSegment space scans = do+  let tblock_size' = fmap pe64 tblock_size       (gtids, dims) = unzip $ unSegSpace space       dims' = map pe64 dims   required_groups <-     dPrimVE "required_groups" $       sExt32 $-        product dims' `divUp` sExt64 (unCount group_size')+        product dims' `divUp` sExt64 (unCount tblock_size') -  sKernelThread "scan_stage3" (segFlat space) (defKernelAttrs num_groups group_size) $-    virtualiseGroups SegVirt required_groups $ \virt_group_id -> do+  sKernelThread "scan_stage3" (segFlat space) (defKernelAttrs num_tblocks tblock_size) $+    virtualiseBlocks SegVirt required_groups $ \virt_tblock_id -> do       constants <- kernelConstants <$> askEnv        -- Compute our logical index.       flat_idx <-         dPrimVE "flat_idx" $-          sExt64 virt_group_id * sExt64 (unCount group_size')+          sExt64 virt_tblock_id * sExt64 (unCount tblock_size')             + sExt64 (kernelLocalThreadId constants)       zipWithM_ dPrimV_ gtids $ unflattenIndex dims' flat_idx @@ -494,20 +494,20 @@   -- 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" int64-  sOp $ Imp.GetSizeMax (tvVar stage1_max_num_groups) SizeGroup+  stage1_max_num_tblocks <- dPrim "stage1_max_num_tblocks" int64+  sOp $ Imp.GetSizeMax (tvVar stage1_max_num_tblocks) SizeThreadBlock -  stage1_num_groups <-+  stage1_num_tblocks <-     fmap (Imp.Count . tvSize) $-      dPrimV "stage1_num_groups" $-        sMin64 (tvExp stage1_max_num_groups) $-          pe64 . Imp.unCount . kAttrNumGroups $+      dPrimV "stage1_num_tblocks" $+        sMin64 (tvExp stage1_max_num_tblocks) $+          pe64 . Imp.unCount . kAttrNumBlocks $             attrs    (stage1_num_threads, elems_per_group, crossesSegment) <--    scanStage1 pat stage1_num_groups (kAttrGroupSize attrs) space scans kbody+    scanStage1 pat stage1_num_tblocks (kAttrBlockSize attrs) space scans kbody    emit $ Imp.DebugPrint "elems_per_group" $ Just $ untyped elems_per_group -  scanStage2 pat stage1_num_threads elems_per_group stage1_num_groups crossesSegment space scans-  scanStage3 pat (kAttrNumGroups attrs) (kAttrGroupSize attrs) elems_per_group crossesSegment space scans+  scanStage2 pat stage1_num_threads elems_per_group stage1_num_tblocks crossesSegment space scans+  scanStage3 pat (kAttrNumBlocks attrs) (kAttrBlockSize attrs) elems_per_group crossesSegment space scans

@@ -137,7 +137,7 @@  pointerQuals :: String -> [C.TypeQual] pointerQuals "global" = [C.ctyquals|__global|]-pointerQuals "local" = [C.ctyquals|__local|]+pointerQuals "shared" = [C.ctyquals|__local|] pointerQuals "private" = [C.ctyquals|__private|] pointerQuals "constant" = [C.ctyquals|__constant|] pointerQuals "write_only" = [C.ctyquals|__write_only|]@@ -149,11 +149,11 @@ pointerQuals "device" = pointerQuals "global" pointerQuals s = error $ "'" ++ s ++ "' is not an OpenCL kernel address space." --- In-kernel name and per-workgroup size in bytes.-type LocalMemoryUse = (VName, Count Bytes (TExp Int64))+-- In-kernel name and per-threadblock size in bytes.+type SharedMemoryUse = (VName, Count Bytes (TExp Int64))  data KernelState = KernelState-  { kernelLocalMemory :: [LocalMemoryUse],+  { kernelSharedMemory :: [SharedMemoryUse],     kernelFailures :: [FailureMsg],     kernelNextSync :: Int,     -- | Has a potential failure occurred sine the last@@ -330,7 +330,7 @@     toDevice :: HostOp -> KernelOp     toDevice _ = bad -isConst :: GroupDim -> Maybe T.Text+isConst :: BlockDim -> Maybe T.Text isConst (Left (ValueExp (IntValue x))) =   Just $ prettyText $ intToInt64 x isConst (Right (SizeConst v _)) =@@ -359,10 +359,10 @@       (kernel_consts, (const_defs, const_undefs)) =         second unzip $ unzip $ mapMaybe (constDef (kernelName kernel)) $ kernelUses kernel -  let (local_memory_bytes, (local_memory_params, local_memory_args, local_memory_init)) =+  let (shared_memory_bytes, (shared_memory_params, shared_memory_args, shared_memory_init)) =         second unzip3 $           evalState-            (mapAccumLM prepareLocalMemory 0 (kernelLocalMemory kstate))+            (mapAccumLM prepareSharedMemory 0 (kernelSharedMemory kstate))             blankNameSource    let (use_params, unpack_params) =@@ -417,23 +417,23 @@           [C.cparam|__global typename int64_t *global_failure_args|]         ] -      (local_memory_param, prepare_local_memory) =+      (shared_memory_param, prepare_shared_memory) =         case target of           TargetOpenCL ->-            ( [[C.cparam|__local typename uint64_t* local_mem_aligned|]],-              [C.citems|__local unsigned char* local_mem = local_mem_aligned;|]+            ( [[C.cparam|__local typename uint64_t* shared_mem_aligned|]],+              [C.citems|__local unsigned char* shared_mem = (__local unsigned char*)shared_mem_aligned;|]             )           TargetCUDA -> (mempty, mempty)           TargetHIP -> (mempty, mempty)        params =-        local_memory_param+        shared_memory_param           ++ take (numFailureParams safety) failure_params-          ++ local_memory_params+          ++ shared_memory_params           ++ use_params        attribute =-        case mapM isConst $ kernelGroupSize kernel of+        case mapM isConst $ kernelBlockSize kernel of           Just [x, y, z] ->             "FUTHARK_KERNEL_SIZED" <> prettyText (x, y, z) <> "\n"           Just [x, y] ->@@ -448,8 +448,8 @@             [C.cfun|void $id:name ($params:params) {                     $items:(mconcat unpack_params)                     $items:const_defs-                    $items:prepare_local_memory-                    $items:local_memory_init+                    $items:prepare_shared_memory+                    $items:shared_memory_init                     $items:error_init                     $items:kernel_body @@ -466,23 +466,23 @@       }    -- The error handling stuff is automatically added later.-  let args = local_memory_args ++ kernelArgs kernel+  let args = shared_memory_args ++ kernelArgs kernel -  pure $ LaunchKernel safety name local_memory_bytes args num_groups group_size+  pure $ LaunchKernel safety name shared_memory_bytes args num_tblocks tblock_size   where     name = kernelName kernel-    num_groups = kernelNumGroups kernel-    group_size = kernelGroupSize kernel+    num_tblocks = kernelNumBlocks kernel+    tblock_size = kernelBlockSize kernel     padTo8 e = e + ((8 - (e `rem` 8)) `rem` 8) -    prepareLocalMemory (Count offset) (mem, Count size) = do+    prepareSharedMemory (Count offset) (mem, Count size) = do       param <- newVName $ baseString mem ++ "_offset"       let offset' = offset + padTo8 size       pure         ( bytes offset',           ( [C.cparam|typename int64_t $id:param|],             ValueKArg (untyped offset) $ IntType Int64,-            [C.citem|volatile __local $ty:defaultMemBlockType $id:mem = &local_mem[$id:param];|]+            [C.citem|volatile __local $ty:defaultMemBlockType $id:mem = &shared_mem[$id:param];|]           )         ) @@ -614,8 +614,8 @@     fence FenceGlobal = [C.cexp|CLK_GLOBAL_MEM_FENCE | CLK_LOCAL_MEM_FENCE|]      kernelOps :: GC.OpCompiler KernelOp KernelState-    kernelOps (GetGroupId v i) =-      GC.stm [C.cstm|$id:v = get_group_id($int:i);|]+    kernelOps (GetBlockId v i) =+      GC.stm [C.cstm|$id:v = get_tblock_id($int:i);|]     kernelOps (GetLocalId v i) =       GC.stm [C.cstm|$id:v = get_local_id($int:i);|]     kernelOps (GetLocalSize v i) =@@ -632,7 +632,7 @@     kernelOps (LocalAlloc name size) = do       name' <- newVName $ prettyString name ++ "_backing"       GC.modifyUserState $ \s ->-        s {kernelLocalMemory = (name', size) : kernelLocalMemory s}+        s {kernelSharedMemory = (name', size) : kernelSharedMemory s}       GC.stm [C.cstm|$id:name = (__local unsigned char*) $id:name';|]     kernelOps (ErrorSync f) = do       label <- nextErrorLabel

@@ -1,13 +1,13 @@--- | Some OpenCL platforms have a SIMD/warp/wavefront-based execution--- model that execute groups of threads in lockstep, permitting us to--- perform cross-thread synchronisation within each such group without--- the use of barriers.  Unfortunately, there seems to be no reliable--- way to query these sizes at runtime.  Instead, we use builtin--- tables to figure out which size we should use for a specific--- platform and device.  If nothing matches here, the wave size should--- be set to one.+-- | Some GPU platforms have a SIMD/warp/wavefront-based execution+-- model that execute blocks of threads in lockstep, permitting us to+-- perform cross-thread synchronisation within each such block without+-- the use of barriers. Unfortunately, there seems to be no reliable+-- way to query these sizes at runtime. Instead, we use builtin tables+-- to figure out which size we should use for a specific platform and+-- device. If nothing matches here, the wave size should be set to+-- one. ----- We also use this to select reasonable default group sizes and group+-- We also use this to select reasonable default block sizes and block -- counts. module Futhark.CodeGen.OpenCL.Heuristics   ( SizeHeuristic (..),@@ -34,7 +34,7 @@   pretty (DeviceInfo s) = "device_info" <> parens (pretty s)  -- | A size that can be assigned a default.-data WhichSize = LockstepWidth | NumGroups | GroupSize | TileSize | RegTileSize | Threshold+data WhichSize = LockstepWidth | NumBlocks | BlockSize | TileSize | RegTileSize | Threshold  -- | A heuristic for setting the default value for something. data SizeHeuristic = SizeHeuristic@@ -50,17 +50,17 @@   [ SizeHeuristic "NVIDIA CUDA" DeviceGPU LockstepWidth 32,     SizeHeuristic "AMD Accelerated Parallel Processing" DeviceGPU LockstepWidth 32,     SizeHeuristic "" DeviceGPU LockstepWidth 1,-    -- We calculate the number of groups to aim for 1024 threads per-    -- compute unit if we also use the default group size.  This seems+    -- We calculate the number of blocks to aim for 1024 threads per+    -- compute unit if we also use the default block size.  This seems     -- to perform well in practice.-    SizeHeuristic "" DeviceGPU NumGroups $ 4 * max_compute_units,-    SizeHeuristic "" DeviceGPU GroupSize 256,+    SizeHeuristic "" DeviceGPU NumBlocks $ 4 * max_compute_units,+    SizeHeuristic "" DeviceGPU BlockSize 256,     SizeHeuristic "" DeviceGPU TileSize 16,     SizeHeuristic "" DeviceGPU RegTileSize 4,     SizeHeuristic "" DeviceGPU Threshold $ 32 * 1024,     SizeHeuristic "" DeviceCPU LockstepWidth 1,-    SizeHeuristic "" DeviceCPU NumGroups max_compute_units,-    SizeHeuristic "" DeviceCPU GroupSize 32,+    SizeHeuristic "" DeviceCPU NumBlocks max_compute_units,+    SizeHeuristic "" DeviceCPU BlockSize 32,     SizeHeuristic "" DeviceCPU TileSize 4,     SizeHeuristic "" DeviceCPU RegTileSize 1,     SizeHeuristic "" DeviceCPU Threshold max_compute_units

@@ -62,3 +62,55 @@   asSegOp (SegOp op) = Just op   asSegOp _ = Nothing   segOp = SegOp++-- Note [GPU Terminology]+--+-- For lack of a better spot to put it, this Note summarises the+-- terminology used for GPU concepts in the Futhark compiler. The+-- terminology is based on CUDA terminology, and tries to match it as+-- closely as possible. However, this was not always the case (issue+-- #2062), so you may find some code that uses e.g. OpenCL+-- terminology. In most cases there is no ambiguity, but there are a+-- few instances where the same term is used for different things.+-- Please fix any instances you find.+--+-- The terminology is as follows:+--+-- Host: Essentially the CPU; whatever is controlling the GPU.+--+-- Kernel: A GPU program that can be launched from the host.+--+-- Grid: The geometry of the thread blocks launched for a kernel. The+-- size of a grid is always in terms of the number of thread blocks+-- ("grid size"). A grid can have up to 3 dimensions, although we do+-- not make much use of it - and not at all prior to code generation.+--+-- Thread block: Just as in CUDA. "Workgroup" in OpenCL. Abbretiation:+-- tblock. Never just call this "block"; there are too many things+-- called "block". Must match the dimensionality of the grid.+--+-- Thread: Just as in CUDA.  "Workitem" in OpenCL.+--+-- Global thread identifier: A globally unique number for a thread+-- along one dimension. Abbreviation: gtid. We also use this term for+-- the identifiers bound by SegOps. In OpenCL, corresponds to+-- get_global_id(). (Except when we virtualise the thread space.)+--+-- Local thread identifier: A locally unique number (within the thread+-- block) for each thread. Abbreviation: ltid. In OpenCL, corresponds+-- to get_local_id().  In CUDA, corresponds to threadIdx.+--+-- Thread block identifier: A number unique to each thread block in a+-- single dimension.  In CUDA, corresponds to blockIdx.+--+-- Local memory: Thread-local private memory. In CUDA, this is+-- sometimes put in registers (if you are very careful in how you use+-- it). In OpenCL, this is called "private memory", and "local memory"+-- is something else entirely.+--+-- Shared memory: Just as in CUDA. Fast scratchpad memory accessible+-- to all threads within the same thread block. In OpenCL, this is+-- "local memory".+--+-- Device memory: Sometimes also called "global memory"; this is the+-- big-but-slow memory on the GPU.

@@ -61,14 +61,14 @@ -- the dimensions specified by @[i,j,k]@. -- -- At the moment, this is only supported for 'SegNoVirtFull'--- intra-group parallelism in GPU code, as we have not yet found it+-- intra-block parallelism in GPU code, as we have not yet found it -- useful anywhere else. newtype SegSeqDims = SegSeqDims {segSeqDims :: [Int]}   deriving (Eq, Ord, Show) --- | Do we need group-virtualisation when generating code for the+-- | Do we need block-virtualisation when generating code for the -- segmented operation?  In most cases, we do, but for some simple--- kernels, we compute the full number of groups in advance, and then+-- kernels, we compute the full number of blocks in advance, and then -- virtualisation is an unnecessary (but generally very small) -- overhead.  This only really matters for fairly trivial but very -- wide @map@ kernels where each thread performs constant-time work on@@ -85,23 +85,23 @@ -- | The actual, physical grid dimensions used for the GPU kernel -- running this 'SegOp'. data KernelGrid = KernelGrid-  { gridNumGroups :: Count NumGroups SubExp,-    gridGroupSize :: Count GroupSize SubExp+  { gridNumBlocks :: Count NumBlocks SubExp,+    gridBlockSize :: Count BlockSize SubExp   }   deriving (Eq, Ord, Show)  -- | At which level the *body* of a t'SegOp' executes. data SegLevel   = SegThread SegVirt (Maybe KernelGrid)-  | SegGroup SegVirt (Maybe KernelGrid)-  | SegThreadInGroup SegVirt+  | SegBlock SegVirt (Maybe KernelGrid)+  | SegThreadInBlock SegVirt   deriving (Eq, Ord, Show)  -- | The 'SegVirt' of the 'SegLevel'. segVirt :: SegLevel -> SegVirt segVirt (SegThread v _) = v-segVirt (SegGroup v _) = v-segVirt (SegThreadInGroup v) = v+segVirt (SegBlock v _) = v+segVirt (SegThreadInBlock v) = v  instance PP.Pretty SegVirt where   pretty SegNoVirt = mempty@@ -109,60 +109,60 @@   pretty SegVirt = "virtualise"  instance PP.Pretty KernelGrid where-  pretty (KernelGrid num_groups group_size) =-    "groups="-      <> pretty num_groups+  pretty (KernelGrid num_tblocks tblock_size) =+    "grid="+      <> pretty num_tblocks       <> PP.semi-        <+> "groupsize="-      <> pretty group_size+        <+> "blocksize="+      <> pretty tblock_size  instance PP.Pretty SegLevel where   pretty (SegThread virt grid) =     PP.parens ("thread" <> PP.semi <+> pretty virt <> PP.semi <+> pretty grid)-  pretty (SegGroup virt grid) =-    PP.parens ("group" <> PP.semi <+> pretty virt <> PP.semi <+> pretty grid)-  pretty (SegThreadInGroup virt) =-    PP.parens ("ingroup" <> PP.semi <+> pretty virt)+  pretty (SegBlock virt grid) =+    PP.parens ("block" <> PP.semi <+> pretty virt <> PP.semi <+> pretty grid)+  pretty (SegThreadInBlock virt) =+    PP.parens ("inblock" <> PP.semi <+> pretty virt)  instance Engine.Simplifiable KernelGrid where-  simplify (KernelGrid num_groups group_size) =+  simplify (KernelGrid num_tblocks tblock_size) =     KernelGrid-      <$> traverse Engine.simplify num_groups-      <*> traverse Engine.simplify group_size+      <$> traverse Engine.simplify num_tblocks+      <*> traverse Engine.simplify tblock_size  instance Engine.Simplifiable SegLevel where   simplify (SegThread virt grid) =     SegThread virt <$> Engine.simplify grid-  simplify (SegGroup virt grid) =-    SegGroup virt <$> Engine.simplify grid-  simplify (SegThreadInGroup virt) =-    pure $ SegThreadInGroup virt+  simplify (SegBlock virt grid) =+    SegBlock virt <$> Engine.simplify grid+  simplify (SegThreadInBlock virt) =+    pure $ SegThreadInBlock virt  instance Substitute KernelGrid where-  substituteNames substs (KernelGrid num_groups group_size) =+  substituteNames substs (KernelGrid num_tblocks tblock_size) =     KernelGrid-      (substituteNames substs num_groups)-      (substituteNames substs group_size)+      (substituteNames substs num_tblocks)+      (substituteNames substs tblock_size)  instance Substitute SegLevel where   substituteNames substs (SegThread virt grid) =     SegThread virt (substituteNames substs grid)-  substituteNames substs (SegGroup virt grid) =-    SegGroup virt (substituteNames substs grid)-  substituteNames _ (SegThreadInGroup virt) =-    SegThreadInGroup virt+  substituteNames substs (SegBlock virt grid) =+    SegBlock virt (substituteNames substs grid)+  substituteNames _ (SegThreadInBlock virt) =+    SegThreadInBlock virt  instance Rename SegLevel where   rename = substituteRename  instance FreeIn KernelGrid where-  freeIn' (KernelGrid num_groups group_size) =-    freeIn' (num_groups, group_size)+  freeIn' (KernelGrid num_tblocks tblock_size) =+    freeIn' (num_tblocks, tblock_size)  instance FreeIn SegLevel where   freeIn' (SegThread _virt grid) = freeIn' grid-  freeIn' (SegGroup _virt grid) = freeIn' grid-  freeIn' (SegThreadInGroup _virt) = mempty+  freeIn' (SegBlock _virt grid) = freeIn' grid+  freeIn' (SegThreadInBlock _virt) = mempty  -- | A simple size-level query or computation. data SizeOp@@ -172,28 +172,28 @@     GetSizeMax SizeClass   | -- | Compare size (likely a threshold) with some integer value.     CmpSizeLe Name SizeClass SubExp-  | -- | @CalcNumGroups w max_num_groups group_size@ calculates the-    -- number of GPU workgroups to use for an input of the given size.+  | -- | @CalcNumBlocks w max_num_tblocks tblock_size@ calculates the+    -- number of GPU threadblocks to use for an input of the given size.     -- The @Name@ is a size name.  Note that @w@ is an i64 to avoid     -- overflow issues.-    CalcNumGroups SubExp Name SubExp+    CalcNumBlocks SubExp Name SubExp   deriving (Eq, Ord, Show)  instance Substitute SizeOp where   substituteNames substs (CmpSizeLe name sclass x) =     CmpSizeLe name sclass (substituteNames substs x)-  substituteNames substs (CalcNumGroups w max_num_groups group_size) =-    CalcNumGroups+  substituteNames substs (CalcNumBlocks w max_num_tblocks tblock_size) =+    CalcNumBlocks       (substituteNames substs w)-      max_num_groups-      (substituteNames substs group_size)+      max_num_tblocks+      (substituteNames substs tblock_size)   substituteNames _ op = op  instance Rename SizeOp where   rename (CmpSizeLe name sclass x) =     CmpSizeLe name sclass <$> rename x-  rename (CalcNumGroups w max_num_groups group_size) =-    CalcNumGroups <$> rename w <*> pure max_num_groups <*> rename group_size+  rename (CalcNumBlocks w max_num_tblocks tblock_size) =+    CalcNumBlocks <$> rename w <*> pure max_num_tblocks <*> rename tblock_size   rename x = pure x  instance IsOp SizeOp where@@ -205,7 +205,7 @@   opType (GetSize _ _) = pure [Prim int64]   opType (GetSizeMax _) = pure [Prim int64]   opType CmpSizeLe {} = pure [Prim Bool]-  opType CalcNumGroups {} = pure [Prim int64]+  opType CalcNumBlocks {} = pure [Prim int64]  instance AliasedOp SizeOp where   opAliases _ = [mempty]@@ -213,7 +213,7 @@  instance FreeIn SizeOp where   freeIn' (CmpSizeLe _ _ x) = freeIn' x-  freeIn' (CalcNumGroups w _ group_size) = freeIn' w <> freeIn' group_size+  freeIn' (CalcNumBlocks w _ tblock_size) = freeIn' w <> freeIn' tblock_size   freeIn' _ = mempty  instance PP.Pretty SizeOp where@@ -226,22 +226,22 @@       <> parens (commasep [pretty name, pretty size_class])         <+> "<="         <+> pretty x-  pretty (CalcNumGroups w max_num_groups group_size) =-    "calc_num_groups" <> parens (commasep [pretty w, pretty max_num_groups, pretty group_size])+  pretty (CalcNumBlocks w max_num_tblocks tblock_size) =+    "calc_num_tblocks" <> parens (commasep [pretty w, pretty max_num_tblocks, pretty tblock_size])  instance OpMetrics SizeOp where   opMetrics GetSize {} = seen "GetSize"   opMetrics GetSizeMax {} = seen "GetSizeMax"   opMetrics CmpSizeLe {} = seen "CmpSizeLe"-  opMetrics CalcNumGroups {} = seen "CalcNumGroups"+  opMetrics CalcNumBlocks {} = seen "CalcNumBlocks"  typeCheckSizeOp :: (TC.Checkable rep) => SizeOp -> TC.TypeM rep () typeCheckSizeOp GetSize {} = pure () typeCheckSizeOp GetSizeMax {} = pure () typeCheckSizeOp (CmpSizeLe _ _ x) = TC.require [Prim int64] x-typeCheckSizeOp (CalcNumGroups w _ group_size) = do+typeCheckSizeOp (CalcNumBlocks w _ tblock_size) = do   TC.require [Prim int64] w-  TC.require [Prim int64] group_size+  TC.require [Prim int64] tblock_size  -- | A host-level operation; parameterised by what else it can do. data HostOp op rep@@ -364,31 +364,31 @@  checkGrid :: (TC.Checkable rep) => KernelGrid -> TC.TypeM rep () checkGrid grid = do-  TC.require [Prim int64] $ unCount $ gridNumGroups grid-  TC.require [Prim int64] $ unCount $ gridGroupSize grid+  TC.require [Prim int64] $ unCount $ gridNumBlocks grid+  TC.require [Prim int64] $ unCount $ gridBlockSize grid  checkSegLevel ::   (TC.Checkable rep) =>   Maybe SegLevel ->   SegLevel ->   TC.TypeM rep ()-checkSegLevel (Just SegGroup {}) (SegThreadInGroup _virt) =+checkSegLevel (Just SegBlock {}) (SegThreadInBlock _virt) =   pure ()-checkSegLevel _ (SegThreadInGroup _virt) =-  TC.bad $ TC.TypeError "ingroup SegOp not in group SegOp."+checkSegLevel _ (SegThreadInBlock _virt) =+  TC.bad $ TC.TypeError "inblock SegOp not in block SegOp." checkSegLevel (Just SegThread {}) _ =   TC.bad $ TC.TypeError "SegOps cannot occur when already at thread level."-checkSegLevel (Just SegThreadInGroup {}) _ =-  TC.bad $ TC.TypeError "SegOps cannot occur when already at ingroup level."+checkSegLevel (Just SegThreadInBlock {}) _ =+  TC.bad $ TC.TypeError "SegOps cannot occur when already at inblock level." checkSegLevel _ (SegThread _virt Nothing) =   pure () checkSegLevel (Just _) SegThread {} =   TC.bad $ TC.TypeError "thread-level SegOp cannot be nested" checkSegLevel Nothing (SegThread _virt grid) =   mapM_ checkGrid grid-checkSegLevel (Just _) SegGroup {} =-  TC.bad $ TC.TypeError "group-level SegOp cannot be nested"-checkSegLevel Nothing (SegGroup _virt grid) =+checkSegLevel (Just _) SegBlock {} =+  TC.bad $ TC.TypeError "block-level SegOp cannot be nested"+checkSegLevel Nothing (SegBlock _virt grid) =   mapM_ checkGrid grid  typeCheckHostOp ::

@@ -59,9 +59,9 @@ simplifyKernelOp _ (SizeOp (CmpSizeLe key size_class x)) = do   x' <- Engine.simplify x   pure (SizeOp $ CmpSizeLe key size_class x', mempty)-simplifyKernelOp _ (SizeOp (CalcNumGroups w max_num_groups group_size)) = do+simplifyKernelOp _ (SizeOp (CalcNumBlocks w max_num_tblocks tblock_size)) = do   w' <- Engine.simplify w-  pure (SizeOp $ CalcNumGroups w' max_num_groups group_size, mempty)+  pure (SizeOp $ CalcNumBlocks w' max_num_tblocks tblock_size, mempty) simplifyKernelOp _ (GPUBody ts body) = do   ts' <- Engine.simplify ts   (hoisted, body') <-

@@ -5,8 +5,8 @@     sizeDefault,     KernelPath,     Count (..),-    NumGroups,-    GroupSize,+    NumBlocks,+    BlockSize,     NumThreads,   ) where@@ -29,16 +29,17 @@ data SizeClass   = -- | A threshold with an optional default.     SizeThreshold KernelPath (Maybe Int64)-  | SizeGroup-  | SizeNumGroups+  | SizeThreadBlock+  | -- | The number of thread blocks.+    SizeGrid   | SizeTile   | SizeRegTile   | -- | Likely not useful on its own, but querying the     -- maximum can be handy.-    SizeLocalMemory+    SizeSharedMemory   | -- | A bespoke size with a default.     SizeBespoke Name Int64-  | -- | Amount of registers available per workgroup. Mostly+  | -- | Amount of registers available per threadblock. Mostly     -- meaningful for querying the maximum.     SizeRegisters   | -- | Amount of L2 cache memory, in bytes. Mostly meaningful for@@ -53,11 +54,11 @@       pStep (v, True) = pretty v       pStep (v, False) = "!" <> pretty v       def' = maybe "def" pretty def-  pretty SizeGroup = "group_size"-  pretty SizeNumGroups = "num_groups"+  pretty SizeThreadBlock = "thread_block_size"+  pretty SizeGrid = "grid_size"   pretty SizeTile = "tile_size"   pretty SizeRegTile = "reg_tile_size"-  pretty SizeLocalMemory = "local_memory"+  pretty SizeSharedMemory = "shared_memory"   pretty (SizeBespoke k def) =     "bespoke" <> parens (pretty k <> comma <+> pretty def)   pretty SizeRegisters = "registers"@@ -82,11 +83,11 @@ instance Traversable (Count u) where   traverse f (Count x) = Count <$> f x --- | Phantom type for the number of groups of some kernel.-data NumGroups+-- | Phantom type for the number of blocks of some kernel.+data NumBlocks --- | Phantom type for the group size of some kernel.-data GroupSize+-- | Phantom type for the block size of some kernel.+data BlockSize  -- | Phantom type for number of threads. data NumThreads

@@ -914,17 +914,15 @@    TC.context ("in index function " <> prettyText ixfun) $ do     traverse_ (TC.requirePrimExp int64 . untyped) ixfun-    let ixfun_rank = IxFun.rank ixfun-        ident_rank = shapeRank shape-    unless (ixfun_rank == ident_rank) $+    unless (IxFun.shape ixfun == map pe64 (shapeDims shape)) $       TC.bad $         TC.TypeError $-          "Arity of index function ("-            <> prettyText ixfun_rank-            <> ") does not match rank of array "+          "Shape of index function ("+            <> prettyText (IxFun.shape ixfun)+            <> ") does not match shape of array "             <> prettyText name             <> " ("-            <> prettyText ident_rank+            <> prettyText shape             <> ")"  bodyReturnsFromPat ::@@ -1074,7 +1072,7 @@         ( Array pt shape u,           MemArray _ _ _ (ArrayIn mem ixfun)           )-            | Just (i, mem_p) <- isMergeVar mem,+            | Just (i, mem_p) <- isLoopVar mem,               Mem space <- paramType mem_p ->                 pure $ MemArray pt shape u $ Just $ ReturnsNewBlock space i ixfun'             | otherwise ->@@ -1089,7 +1087,7 @@           pure $ MemPrim pt         (Mem space, _) ->           pure $ MemMem space-    isMergeVar v = find ((== v) . paramName . snd) $ zip [0 ..] mergevars+    isLoopVar v = find ((== v) . paramName . snd) $ zip [0 ..] mergevars     mergevars = map fst merge expReturns (Apply _ _ ret _) =   pure $ Just $ map (funReturnsToExpReturns . fst) ret

@@ -22,6 +22,7 @@     substituteInIxFun,     substituteInLMAD,     existentialize,+    existentialized,     closeEnough,     disjoint,     disjoint2,@@ -152,13 +153,13 @@ iota :: (IntegralExp num) => Shape num -> IxFun num iota = iotaOffset 0 --- | Create a single-LMAD index function that is--- existential in everything, with the provided permutation.-mkExistential :: Int -> Int -> Int -> IxFun (Ext a)-mkExistential basis_rank lmad_rank start =-  IxFun (LMAD.mkExistential lmad_rank start) basis+-- | Create a single-LMAD index function that is existential in+-- everything except shape, with the provided shape.+mkExistential :: Int -> Shape (Ext a) -> Int -> IxFun (Ext a)+mkExistential basis_rank lmad_shape start =+  IxFun (LMAD.mkExistential lmad_shape start) basis   where-    basis = take basis_rank $ map Ext [start + 1 + lmad_rank * 2 ..]+    basis = take basis_rank $ map Ext [start + 1 + length lmad_shape ..]  -- | Permute dimensions. permute ::@@ -241,18 +242,30 @@           (map onDim (LMAD.dims lmad))    in Just $ IxFun lmad' base --- | Turn all the leaves of the index function into 'Ext's.  We---  require that there's only one LMAD, that the index function is---  contiguous, and the base shape has only one dimension.+-- | Turn all the leaves of the index function into 'Ext's, except for+--  the shape, which where the leaves are simply made 'Free'. existentialize ::+  Int ->   IxFun (TPrimExp Int64 a) ->-  IxFun (TPrimExp Int64 (Ext b))-existentialize ixfun = evalState (traverse (const mkExt) ixfun) 0+  IxFun (TPrimExp Int64 (Ext a))+existentialize start (IxFun lmad base) = evalState (IxFun <$> lmad' <*> base') start   where     mkExt = do       i <- get       put $ i + 1       pure $ TPrimExp $ LeafExp (Ext i) int64+    lmad' = LMAD <$> mkExt <*> mapM onDim (dims lmad)+    base' = traverse (const mkExt) base+    onDim ld = LMADDim <$> mkExt <*> pure (fmap Free (ldShape ld))++-- | Retrieve those elements that 'existentialize' changes. That is,+-- everything except the shape (and in the same order as+-- 'existentialise' existentialises them).+existentialized :: IxFun a -> [a]+existentialized (IxFun (LMAD offset dims) base) =+  offset : concatMap onDim dims <> base+  where+    onDim (LMADDim ldstride _) = [ldstride]  -- | When comparing index functions as part of the type check in KernelsMem, -- we may run into problems caused by the simplifier. As index functions can be

@@ -276,11 +276,12 @@ iota off = iotaStrided off 1 {-# NOINLINE iota #-} --- | Create an LMAD that is existential in everything.-mkExistential :: Int -> Int -> LMAD (Ext a)-mkExistential r start = LMAD (Ext start) $ map onDim [0 .. r - 1]+-- | Create an LMAD that is existential in everything except shape.+mkExistential :: Shape (Ext a) -> Int -> LMAD (Ext a)+mkExistential shp start = LMAD (Ext start) $ zipWith onDim shp [0 .. r - 1]   where-    onDim i = LMADDim (Ext (start + 1 + i * 2)) (Ext (start + 2 + i * 2))+    r = length shp+    onDim d i = LMADDim {ldStride = Ext (start + 1 + i), ldShape = d}  -- | Permute dimensions. permute :: LMAD num -> Permutation -> LMAD num

@@ -134,8 +134,8 @@       (arr_to_mem, oldmem_to_mem) <-         fmap unzip $           forM fixable $ \(arr_pe, mem_size, oldmem, space) -> do-            size <- toSubExp "size" mem_size-            mem <- letExp "mem" $ Op $ Alloc size space+            size <- toSubExp "unext_mem_size" mem_size+            mem <- letExp "unext_mem" $ Op $ Alloc size space             pure ((patElemName arr_pe, mem), (oldmem, mem))        -- Update the branches to contain Copy expressions putting the

@@ -765,12 +765,11 @@ pSizeClass :: Parser GPU.SizeClass pSizeClass =   choice-    [ keyword "group_size" $> GPU.SizeGroup,-      keyword "num_groups" $> GPU.SizeNumGroups,-      keyword "num_groups" $> GPU.SizeNumGroups,+    [ keyword "thread_block_size" $> GPU.SizeThreadBlock,+      keyword "grid_size" $> GPU.SizeGrid,       keyword "tile_size" $> GPU.SizeTile,       keyword "reg_tile_size" $> GPU.SizeRegTile,-      keyword "local_memory" $> GPU.SizeLocalMemory,+      keyword "shared_memory" $> GPU.SizeSharedMemory,       keyword "threshold"         *> parens           ( flip GPU.SizeThreshold@@ -800,9 +799,9 @@         *> ( parens (GPU.CmpSizeLe <$> pName <* pComma <*> pSizeClass)                <*> (lexeme "<=" *> pSubExp)            ),-      keyword "calc_num_groups"+      keyword "calc_num_tblocks"         *> parens-          ( GPU.CalcNumGroups+          ( GPU.CalcNumBlocks               <$> pSubExp               <* pComma               <*> pName@@ -917,13 +916,13 @@         <*> pSegVirt         <* pSemi         <*> optional pKernelGrid,-      "group"-        $> GPU.SegGroup+      "block"+        $> GPU.SegBlock         <* pSemi         <*> pSegVirt         <* pSemi         <*> optional pKernelGrid,-      "ingroup" $> GPU.SegThreadInGroup <* pSemi <*> pSegVirt+      "inblock" $> GPU.SegThreadInBlock <* pSemi <*> pSegVirt     ]   where     pSegVirt =@@ -938,8 +937,8 @@         ]     pKernelGrid =       GPU.KernelGrid-        <$> (lexeme "groups=" $> GPU.Count <*> pSubExp <* pSemi)-        <*> (lexeme "groupsize=" $> GPU.Count <*> pSubExp)+        <$> (lexeme "grid=" $> GPU.Count <*> pSubExp <* pSemi)+        <*> (lexeme "blocksize=" $> GPU.Count <*> pSubExp)  pHostOp :: PR rep -> Parser (op rep) -> Parser (GPU.HostOp op rep) pHostOp pr pOther =

@@ -31,7 +31,6 @@     isScanSOAC,     isReduceSOAC,     isMapSOAC,-    scremaLambda,     ppScrema,     ppHist,     ppStream,@@ -143,11 +142,16 @@  -- | The essential parts of a 'Screma' factored out (everything -- except the input arrays).-data ScremaForm rep-  = ScremaForm-      [Scan rep]-      [Reduce rep]-      (Lambda rep)+data ScremaForm rep = ScremaForm+  { scremaScans :: [Scan rep],+    scremaReduces :: [Reduce rep],+    -- | The "main" lambda of the Screma. For a map, this is+    -- equivalent to 'isMapSOAC'. Note that the meaning of the return+    -- value of this lambda depends crucially on exactly which Screma+    -- this is. The parameters will correspond exactly to elements of+    -- the input arrays, however.+    scremaLambda :: Lambda rep+  }   deriving (Eq, Ord, Show)  singleBinOp :: (Buildable rep) => [Lambda rep] -> Lambda rep@@ -316,17 +320,9 @@   guard $ null reds   pure map_lam --- | Return the "main" lambda of the Screma.  For a map, this is--- equivalent to 'isMapSOAC'.  Note that the meaning of the return--- value of this lambda depends crucially on exactly which Screma this--- is.  The parameters will correspond exactly to elements of the--- input arrays, however.-scremaLambda :: ScremaForm rep -> Lambda rep-scremaLambda (ScremaForm _ _ map_lam) = map_lam- -- | @groupScatterResults <output specification> <results>@ ----- Groups the index values and result values of <results> according to the+-- Blocks the index values and result values of <results> according to the -- <output specification>. -- -- This function is used for extracting and grouping the results of a@@ -348,7 +344,7 @@  -- | @groupScatterResults' <output specification> <results>@ ----- Groups the index values and result values of <results> according to the+-- Blocks the index values and result values of <results> according to the -- output specification. This is the simpler version of @groupScatterResults@, -- which doesn't return any information about shapes or output arrays. --@@ -625,9 +621,9 @@          in map (is_flat <>) vs   opDependencies (Scatter w arrs lam outputs) =     let deps = lambdaDependencies mempty lam (depsOfArrays w arrs)-     in map flattenGroups (groupScatterResults outputs deps)+     in map flattenBlocks (groupScatterResults outputs deps)     where-      flattenGroups (_, arr, ivs) =+      flattenBlocks (_, arr, ivs) =         oneName arr <> mconcat (map (mconcat . fst) ivs) <> mconcat (map snd ivs)   opDependencies (JVP lam args vec) =     mconcat $

@@ -189,7 +189,7 @@ data KernelResult   = -- | Each "worker" in the kernel returns this.     -- Whether this is a result-per-thread or a-    -- result-per-group depends on where the 'SegOp' occurs.+    -- result-per-block depends on where the 'SegOp' occurs.     Returns ResultManifest Certs SubExp   | WriteReturns       Certs@@ -209,7 +209,7 @@           SubExp -- reg tile size for this dim.         )       ]-      VName -- Tile returned by this worker/group.+      VName -- Tile returned by this thread/block.   deriving (Eq, Show, Ord)  -- | Get the certs for this 'KernelResult'.@@ -437,7 +437,7 @@ -- as well as a *level*.  The *level* is a representation-specific bit -- of information.  For example, in GPU backends, it is used to -- indicate whether the 'SegOp' is expected to run at the thread-level--- or the group-level.+-- or the block-level. -- -- The type list is usually the type of the element returned by a -- single thread. The result of the SegOp is then an array of that