futhark 0.25.30 → 0.25.31

raw patch · 84 files changed

+2337/−1233 lines, 84 filesPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

API changes (from Hackage documentation)

Files

• CHANGELOG.md +49/−0
• docs/language-reference.rst +18/−0
• futhark.cabal +5/−2
• prelude/math.fut +20/−20
• prelude/soacs.fut +1/−1
• rts/c/atomics.h +0/−519
• rts/c/atomics16.h +177/−0
• rts/c/atomics32.h +235/−0
• rts/c/atomics64.h +285/−0
• rts/c/atomics8.h +145/−0
• rts/c/gpu.h +4/−0
• rts/cuda/prelude.cu +8/−0
• rts/python/opencl.py +3/−1
• src/Futhark/AD/Fwd.hs +6/−6
• src/Futhark/AD/Rev.hs +6/−6
• src/Futhark/AD/Rev/Hist.hs +10/−9
• src/Futhark/AD/Rev/Reduce.hs +1/−1
• src/Futhark/AD/Rev/Scan.hs +2/−2
• src/Futhark/Analysis/AccessPattern.hs +4/−3
• src/Futhark/Analysis/HORep/MapNest.hs +4/−3
• src/Futhark/Analysis/HORep/SOAC.hs +29/−58
• src/Futhark/Analysis/SymbolTable.hs +33/−28
• src/Futhark/CLI/Profile.hs +8/−5
• src/Futhark/CLI/Test.hs +3/−3
• src/Futhark/CodeGen/Backends/GenericC.hs +1/−1
• src/Futhark/CodeGen/Backends/SimpleRep.hs +9/−2
• src/Futhark/CodeGen/ImpGen.hs +1/−1
• src/Futhark/CodeGen/ImpGen/GPU.hs +22/−1
• src/Futhark/CodeGen/ImpGen/GPU/Base.hs +15/−11
• src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs +74/−27
• src/Futhark/CodeGen/RTS/C.hs +9/−1
• src/Futhark/Doc/Generator.hs +3/−1
• src/Futhark/Fmt/Printer.hs +7/−2
• src/Futhark/IR/Mem.hs +4/−4
• src/Futhark/IR/Mem/Simplify.hs +33/−10
• src/Futhark/IR/Parse.hs +9/−5
• src/Futhark/IR/Pretty.hs +21/−15
• src/Futhark/IR/Prop/Aliases.hs +2/−2
• src/Futhark/IR/Prop/Names.hs +3/−0
• src/Futhark/IR/Prop/Reshape.hs +266/−4
• src/Futhark/IR/Prop/TypeOf.hs +5/−9
• src/Futhark/IR/SOACS/Simplify.hs +38/−32
• src/Futhark/IR/Syntax.hs +46/−13
• src/Futhark/IR/Syntax/Core.hs +14/−3
• src/Futhark/IR/Traversals.hs +13/−13
• src/Futhark/IR/TypeCheck.hs +27/−22
• src/Futhark/Internalise/Exps.hs +46/−27
• src/Futhark/Internalise/Monomorphise.hs +2/−1
• src/Futhark/Internalise/ReplaceRecords.hs +6/−1
• src/Futhark/Optimise/ArrayLayout/Transform.hs +1/−1
• src/Futhark/Optimise/ArrayShortCircuiting/MemRefAggreg.hs +1/−1
• src/Futhark/Optimise/ArrayShortCircuiting/TopdownAnalysis.hs +7/−5
• src/Futhark/Optimise/BlkRegTiling.hs +6/−10
• src/Futhark/Optimise/EntryPointMem.hs +1/−1
• src/Futhark/Optimise/Fusion.hs +2/−2
• src/Futhark/Optimise/Fusion/GraphRep.hs +1/−1
• src/Futhark/Optimise/Fusion/RulesWithAccs.hs +4/−3
• src/Futhark/Optimise/Fusion/TryFusion.hs +4/−4
• src/Futhark/Optimise/GenRedOpt.hs +1/−1
• src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs +3/−3
• src/Futhark/Optimise/Simplify/Engine.hs +1/−1
• src/Futhark/Optimise/Simplify/Rules/BasicOp.hs +75/−41
• src/Futhark/Optimise/Simplify/Rules/Index.hs +22/−14
• src/Futhark/Optimise/Simplify/Rules/Simple.hs +39/−51
• src/Futhark/Optimise/TileLoops.hs +2/−2
• src/Futhark/Optimise/TileLoops/Shared.hs +19/−6
• src/Futhark/Pass/ExplicitAllocations.hs +3/−3
• src/Futhark/Pass/ExplicitAllocations/GPU.hs +1/−1
• src/Futhark/Pass/ExtractKernels/DistributeNests.hs +8/−7
• src/Futhark/Pass/ExtractKernels/ISRWIM.hs +11/−18
• src/Futhark/Pass/LiftAllocations.hs +51/−37
• src/Futhark/Tools.hs +3/−3
• src/Futhark/Transform/Substitute.hs +4/−2
• src/Language/Futhark/Interpreter.hs +16/−13
• src/Language/Futhark/Interpreter/Values.hs +10/−0
• src/Language/Futhark/Parser/Parser.y +3/−1
• src/Language/Futhark/Pretty.hs +37/−30
• src/Language/Futhark/Semantic.hs +6/−3
• src/Language/Futhark/TypeChecker.hs +8/−64
• src/Language/Futhark/TypeChecker/Modules.hs +32/−11
• src/Language/Futhark/TypeChecker/Terms.hs +17/−8
• src/Language/Futhark/TypeChecker/Types.hs +9/−5
• src/Language/Futhark/TypeChecker/Unify.hs +1/−1
• unittests/Futhark/IR/Prop/ReshapeTests.hs +196/−4

@@ -5,6 +5,55 @@ 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.31]++### Added++* GPU backends: more efficient atomic operations on 8-bit and 16-bit quantities.+  This helps histograms on these types, as well as AD on programs that use+  `f16`.++* Improved handling of long chains of `flatten`/`unflatten`/`transpose`+  operations.++* New attributes: `#[blank]` and `#[scratch]`.++* A module type `with`-refinement may now have an existentially quantified size+  on its right-hand side.++* Value specs in module types can now use section binding notation for symbolic+  names, and in fact this is the preferred form that is also used by `futhark+  fmt`. (#2266)++* `futhark profile` now also prints proportion of total runtime for each cost centre.++* Futhark no longer warns about entry points with opaque types.++* Types such as `foo.bar` are now turned into `foo_bar` in the C API, rather+  than an ugly hash.++### Fixed++* Interpreter: some tricky aspects of size-lifted types (#2258).++* Incorrect unused-name warning for named parameters in module types.++* Size-lifted abstract types with hidden sizes could result in different sizes+  being incorrectly treated as the same size.++* It was possible to make size-lifted types appear unlifted by using parametric+  types (#2268).++* The same type would be mentioned twice in some type errors.++* The type checker neglected to detect some cases of invalid references from+  return types to names bound in parameter patterns. (#2271)++* Incorrect handling of projections used in size expressions.++* Subtle interactions of modules and sizes in the interpreter and compiler+  (#2273).+ ## [0.25.30]  ### Added

@@ -1645,6 +1645,7 @@  .. productionlist::    spec:   "val" `name` `type_param`* ":" `type`+       : | "val" "(" `symbol` ")" ":" `type`        : | "val" `symbol` `type_param`* ":" `type`        : | ("type" | "type^" | "type~") `name` `type_param`* "=" `type`        : | ("type" | "type^" | "type~") `name` `type_param`*@@ -1742,6 +1743,16 @@  The following expression attributes are supported. +``blank``+.........++Indicates that the value computed by the expression does not matter, and that+the expression can be replaced with an arbitrary other expression of the same+type. This is useful for constructing arrays that will eventually be filled with+``scatter`` or similar operations. Note that this can subvert type-based+invariants safety if the blank value is used, but it cannot subvert memory+safety.+ ``trace`` ......... @@ -1801,6 +1812,13 @@ Do not inline the attributed function application.  If used within a parallel construct (e.g. ``map``), this will likely prevent the GPU backends from generating working code.++``scratch``+...........++Like ``blank``, but the resulting values (if arrays) will comprise initialised+memory. Reading from such arrays is potentially dangerous, as the elements are+completely undefined until they are updated with a ``scatter`` or similar.  ``sequential`` ..............

@@ -1,6 +1,6 @@ cabal-version: 2.4 name:           futhark-version:        0.25.30+version:        0.25.31 synopsis:       An optimising compiler for a functional, array-oriented language.  description:    Futhark is a small programming language designed to be compiled to@@ -36,7 +36,10 @@ extra-source-files: -- Cabal's recompilation tracking doesn't work when we use wildcards -- here, so for now we spell out every single file.-    rts/c/atomics.h+    rts/c/atomics8.h+    rts/c/atomics16.h+    rts/c/atomics32.h+    rts/c/atomics64.h     rts/c/context.h     rts/c/context_prototypes.h     rts/c/backends/c.h

@@ -29,21 +29,21 @@ module type numeric = {   include from_prim -  val + : t -> t -> t-  val - : t -> t -> t-  val * : t -> t -> t-  val / : t -> t -> t-  val % : t -> t -> t-  val ** : t -> t -> t+  val (+) : t -> t -> t+  val (-) : t -> t -> t+  val (*) : t -> t -> t+  val (/) : t -> t -> t+  val (%) : t -> t -> t+  val (**) : t -> t -> t    val to_i64 : t -> i64 -  val == : t -> t -> bool-  val < : t -> t -> bool-  val > : t -> t -> bool-  val <= : t -> t -> bool-  val >= : t -> t -> bool-  val != : t -> t -> bool+  val (==) : t -> t -> bool+  val (<) : t -> t -> bool+  val (>) : t -> t -> bool+  val (<=) : t -> t -> bool+  val (>=) : t -> t -> bool+  val (!=) : t -> t -> bool    -- | Arithmetic negation (use `!` for bitwise negation).   val neg : t -> t@@ -83,32 +83,32 @@    -- | Like `/`@term, but rounds towards zero.  This only matters when   -- one of the operands is negative.  May be more efficient.-  val // : t -> t -> t+  val (//) : t -> t -> t    -- | Like `%`@term, but rounds towards zero.  This only matters when   -- one of the operands is negative.  May be more efficient.-  val %% : t -> t -> t+  val (%%) : t -> t -> t    -- | Bitwise and.-  val & : t -> t -> t+  val (&) : t -> t -> t    -- | Bitwise or.-  val | : t -> t -> t+  val (|) : t -> t -> t    -- | Bitwise xor.-  val ^ : t -> t -> t+  val (^) : t -> t -> t    -- | Bitwise negation.   val not : t -> t    -- | Left shift; inserting zeroes.-  val << : t -> t -> t+  val (<<) : t -> t -> t    -- | Arithmetic right shift, using sign extension for the leftmost bits.-  val >> : t -> t -> t+  val (>>) : t -> t -> t    -- | Logical right shift, inserting zeroes for the leftmost bits.-  val >>> : t -> t -> t+  val (>>>) : t -> t -> t    val num_bits : i32   val get_bit : i32 -> t -> i32

@@ -255,6 +255,6 @@   let flags = map (\x -> if p x then 1 else 0) as   let offsets = scan (+) 0 flags   let m = if n == 0 then 0 else offsets[n - 1]-  in scatter (map (\x -> x) as[:m])+  in scatter (#[scratch] map (\x -> x) as[:m])              (map2 (\f o -> if f == 1 then o - 1 else -1) flags offsets)              as

@@ -1,519 +0,0 @@-// 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_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_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_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_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_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_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_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_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_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_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_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)-  return atomicExch((int32_t*)p, x);-#else-  return atomic_xor(p, x);-#endif-}--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-  return atomic_xor(p, x);-#endif-}--SCALAR_FUN_ATTR int32_t atomic_cmpxchg_i32_global(volatile __global int32_t *p,-                                                         int32_t cmp, int32_t val) {-#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)-  return atomicCAS((int32_t*)p, cmp, val);-#else-  return atomic_cmpxchg(p, cmp, val);-#endif-}--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);-#else-  return atomic_cmpxchg(p, cmp, val);-#endif-}--SCALAR_FUN_ATTR int32_t atomic_add_i32_global(volatile __global int32_t *p, int32_t x) {-#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)-  return atomicAdd((int32_t*)p, x);-#else-  return atomic_add(p, x);-#endif-}--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-  return atomic_add(p, x);-#endif-}--SCALAR_FUN_ATTR float atomic_fadd_f32_global(volatile __global float *p, float x) {-#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)-  return atomicAdd((float*)p, x);-  // On OpenCL, use technique from-  // https://pipinspace.github.io/blog/atomic-float-addition-in-opencl.html-#elif defined(cl_nv_pragma_unroll)-  // use hardware-supported atomic addition on Nvidia GPUs with inline-  // PTX assembly-  float ret;-  asm volatile("atom.global.add.f32 %0,[%1],%2;":"=f"(ret):"l"(p),"f"(x):"memory");-  return ret;-#elif defined(__opencl_c_ext_fp32_global_atomic_add)-  // use hardware-supported atomic addition on some Intel GPUs-  return atomic_fetch_add_explicit((volatile __global atomic_float*)p,-                                   x,-                                   memory_order_relaxed);-#elif __has_builtin(__builtin_amdgcn_global_atomic_fadd_f32)-  // use hardware-supported atomic addition on some AMD GPUs-  return __builtin_amdgcn_global_atomic_fadd_f32(p, x);-#else-  // fallback emulation:-  // https://forums.developer.nvidia.com/t/atomicadd-float-float-atomicmul-float-float/14639/5-  float old = x;-  float ret;-  while ((old=atomic_xchg(p, ret=atomic_xchg(p, 0.0f)+old))!=0.0f);-  return ret;-#endif-}--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-  union { int32_t i; float f; } old;-  union { int32_t i; float f; } assumed;-  old.f = *p;-  do {-    assumed.f = old.f;-    old.f = old.f + x;-    old.i = atomic_cmpxchg_i32_shared((volatile __local int32_t*)p, assumed.i, old.i);-  } while (assumed.i != old.i);-  return old.f;-#endif-}--SCALAR_FUN_ATTR int32_t atomic_smax_i32_global(volatile __global int32_t *p, int32_t x) {-#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)-  return atomicMax((int32_t*)p, x);-#else-  return atomic_max(p, x);-#endif-}--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-  return atomic_max(p, x);-#endif-}--SCALAR_FUN_ATTR int32_t atomic_smin_i32_global(volatile __global int32_t *p, int32_t x) {-#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)-  return atomicMin((int32_t*)p, x);-#else-  return atomic_min(p, x);-#endif-}--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-  return atomic_min(p, x);-#endif-}--SCALAR_FUN_ATTR uint32_t atomic_umax_i32_global(volatile __global uint32_t *p, uint32_t x) {-#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)-  return atomicMax((uint32_t*)p, x);-#else-  return atomic_max(p, x);-#endif-}--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-  return atomic_max(p, x);-#endif-}--SCALAR_FUN_ATTR uint32_t atomic_umin_i32_global(volatile __global uint32_t *p, uint32_t x) {-#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)-  return atomicMin((uint32_t*)p, x);-#else-  return atomic_min(p, x);-#endif-}--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-  return atomic_min(p, x);-#endif-}--SCALAR_FUN_ATTR int32_t atomic_and_i32_global(volatile __global int32_t *p, int32_t x) {-#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)-  return atomicAnd((int32_t*)p, x);-#else-  return atomic_and(p, x);-#endif-}--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-  return atomic_and(p, x);-#endif-}--SCALAR_FUN_ATTR int32_t atomic_or_i32_global(volatile __global int32_t *p, int32_t x) {-#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)-  return atomicOr((int32_t*)p, x);-#else-  return atomic_or(p, x);-#endif-}--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-  return atomic_or(p, x);-#endif-}--SCALAR_FUN_ATTR int32_t atomic_xor_i32_global(volatile __global int32_t *p, int32_t x) {-#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)-  return atomicXor((int32_t*)p, x);-#else-  return atomic_xor(p, x);-#endif-}--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-  return atomic_xor(p, x);-#endif-}--// Start of 64 bit atomics--#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_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_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_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_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_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_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_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_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_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_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_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) {-#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)-  return atomicExch((unsigned long long*)p, x);-#else-  return atom_xor(p, x);-#endif-}--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((unsigned long long*)p, x);-#else-  return atom_xor(p, x);-#endif-}--SCALAR_FUN_ATTR int64_t atomic_cmpxchg_i64_global(volatile __global int64_t *p,-                                                         int64_t cmp, int64_t val) {-#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)-  return atomicCAS((unsigned long long*)p, cmp, val);-#else-  return atom_cmpxchg(p, cmp, val);-#endif-}--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((unsigned long long*)p, cmp, val);-#else-  return atom_cmpxchg(p, cmp, val);-#endif-}--SCALAR_FUN_ATTR int64_t atomic_add_i64_global(volatile __global int64_t *p, int64_t x) {-#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)-  return atomicAdd((unsigned long long*)p, x);-#else-  return atom_add(p, x);-#endif-}--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((unsigned long long*)p, x);-#else-  return atom_add(p, x);-#endif-}--#ifdef FUTHARK_F64_ENABLED--SCALAR_FUN_ATTR double atomic_fadd_f64_global(volatile __global double *p, double x) {-#if defined(FUTHARK_CUDA) && __CUDA_ARCH__ >= 600 || defined(FUTHARK_HIP)-  return atomicAdd((double*)p, x);-  // On OpenCL, use technique from-  // https://pipinspace.github.io/blog/atomic-float-addition-in-opencl.html-#elif defined(cl_nv_pragma_unroll)-  // use hardware-supported atomic addition on Nvidia GPUs with inline-  // PTX assembly-  double ret;-  asm volatile("atom.global.add.f64 %0,[%1],%2;":"=d"(ret):"l"(p),"d"(x):"memory");-  return ret;-#elif __has_builtin(__builtin_amdgcn_global_atomic_fadd_f64)-  // use hardware-supported atomic addition on some AMD GPUs-  return __builtin_amdgcn_global_atomic_fadd_f64(p, x);-#else-  // fallback emulation:-  // https://forums.developer.nvidia.com/t/atomicadd-float-float-atomicmul-float-float/14639/5-  union {int64_t i; double f;} old;-  union {int64_t i; double f;} ret;-  old.f = x;-  while (1) {-    ret.i = atom_xchg((volatile __global int64_t*)p, (int64_t)0);-    ret.f += old.f;-    old.i = atom_xchg((volatile __global int64_t*)p, ret.i);-    if (old.i == 0) {-      break;-    }-  }-  return ret.f;-#endif-}--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-  union { int64_t i; double f; } old;-  union { int64_t i; double f; } assumed;-  old.f = *p;-  do {-    assumed.f = old.f;-    old.f = old.f + x;-    old.i = atomic_cmpxchg_i64_shared((volatile __local int64_t*)p, assumed.i, old.i);-  } while (assumed.i != old.i);-  return old.f;-#endif-}--#endif--SCALAR_FUN_ATTR int64_t atomic_smax_i64_global(volatile __global int64_t *p, int64_t x) {-#if defined(FUTHARK_CUDA)-  return atomicMax((long long*)p, x);-#elif defined(FUTHARK_HIP)-  // Currentely missing in HIP; probably a temporary oversight.-  int64_t old = *p, assumed;-  do {-    assumed = old;-    old = smax64(old, x);-    old = atomic_cmpxchg_i64_global((volatile __global int64_t*)p, assumed, old);-  } while (assumed != old);-  return old;-#else-  return atom_max(p, x);-#endif-}--SCALAR_FUN_ATTR int64_t atomic_smax_i64_shared(volatile __local int64_t *p, int64_t x) {-#if defined(FUTHARK_CUDA)-  return atomicMax((long long*)p, x);-#elif defined(FUTHARK_HIP)-  // Currentely missing in HIP; probably a temporary oversight.-  int64_t old = *p, assumed;-  do {-    assumed = old;-    old = smax64(old, x);-    old = atomic_cmpxchg_i64_shared((volatile __local int64_t*)p, assumed, old);-  } while (assumed != old);-  return old;-#else-  return atom_max(p, x);-#endif-}--SCALAR_FUN_ATTR int64_t atomic_smin_i64_global(volatile __global int64_t *p, int64_t x) {-#if defined(FUTHARK_CUDA)-  return atomicMin((long long*)p, x);-#elif defined(FUTHARK_HIP)-  // Currentely missing in HIP; probably a temporary oversight.-  int64_t old = *p, assumed;-  do {-    assumed = old;-    old = smin64(old, x);-    old = atomic_cmpxchg_i64_global((volatile __global int64_t*)p, assumed, old);-  } while (assumed != old);-  return old;-#else-  return atom_min(p, x);-#endif-}--SCALAR_FUN_ATTR int64_t atomic_smin_i64_shared(volatile __local int64_t *p, int64_t x) {-#if defined(FUTHARK_CUDA)-  return atomicMin((long long*)p, x);-#elif defined(FUTHARK_HIP)-  // Currentely missing in HIP; probably a temporary oversight.-  int64_t old = *p, assumed;-  do {-    assumed = old;-    old = smin64(old, x);-    old = atomic_cmpxchg_i64_shared((volatile __local int64_t*)p, assumed, old);-  } while (assumed != old);-  return old;-#else-  return atom_min(p, x);-#endif-}--SCALAR_FUN_ATTR uint64_t atomic_umax_i64_global(volatile __global uint64_t *p, uint64_t x) {-#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)-  return atomicMax((unsigned long long*)p, x);-#else-  return atom_max(p, x);-#endif-}--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((unsigned long long*)p, x);-#else-  return atom_max(p, x);-#endif-}--SCALAR_FUN_ATTR uint64_t atomic_umin_i64_global(volatile __global uint64_t *p, uint64_t x) {-#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)-  return atomicMin((unsigned long long*)p, x);-#else-  return atom_min(p, x);-#endif-}--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((unsigned long long*)p, x);-#else-  return atom_min(p, x);-#endif-}--SCALAR_FUN_ATTR int64_t atomic_and_i64_global(volatile __global int64_t *p, int64_t x) {-#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)-  return atomicAnd((unsigned long long*)p, x);-#else-  return atom_and(p, x);-#endif-}--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((unsigned long long*)p, x);-#else-  return atom_and(p, x);-#endif-}--SCALAR_FUN_ATTR int64_t atomic_or_i64_global(volatile __global int64_t *p, int64_t x) {-#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)-  return atomicOr((unsigned long long*)p, x);-#else-  return atom_or(p, x);-#endif-}--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((unsigned long long*)p, x);-#else-  return atom_or(p, x);-#endif-}--SCALAR_FUN_ATTR int64_t atomic_xor_i64_global(volatile __global int64_t *p, int64_t x) {-#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)-  return atomicXor((unsigned long long*)p, x);-#else-  return atom_xor(p, x);-#endif-}--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((unsigned long long*)p, x);-#else-  return atom_xor(p, x);-#endif-}--#endif // defined(FUTHARK_CUDA) || defined(FUTHARK_HIP) || defined(cl_khr_int64_base_atomics) && defined(cl_khr_int64_extended_atomics)--// End of atomics.h

@@ -0,0 +1,177 @@+// Start of atomics16.h++SCALAR_FUN_ATTR int16_t atomic_cmpxchg_i16_global(volatile __global int16_t *p,+                                                  int16_t cmp, int16_t val);+SCALAR_FUN_ATTR int16_t atomic_cmpxchg_i16_shared(volatile __local int16_t *p,+                                                  int16_t cmp, int16_t val);+SCALAR_FUN_ATTR int16_t atomic_add_i16_global(volatile __global int16_t *p, int16_t x);+SCALAR_FUN_ATTR int16_t atomic_add_i16_shared(volatile __local int16_t *p, int16_t x);+SCALAR_FUN_ATTR f16 atomic_fadd_f16_global(volatile __global uint16_t *p, f16 x);+SCALAR_FUN_ATTR f16 atomic_fadd_f16_shared(volatile __local uint16_t *p, f16 x);+SCALAR_FUN_ATTR int16_t atomic_smax_i16_global(volatile __global int16_t *p, int16_t x);+SCALAR_FUN_ATTR int16_t atomic_smax_i16_shared(volatile __local int16_t *p, int16_t x);+SCALAR_FUN_ATTR int16_t atomic_smin_i16_global(volatile __global int16_t *p, int16_t x);+SCALAR_FUN_ATTR int16_t atomic_smin_i16_shared(volatile __local int16_t *p, int16_t x);+SCALAR_FUN_ATTR uint16_t atomic_umax_i16_global(volatile __global uint16_t *p, uint16_t x);+SCALAR_FUN_ATTR uint16_t atomic_umax_i16_shared(volatile __local uint16_t *p, uint16_t x);+SCALAR_FUN_ATTR uint16_t atomic_umin_i16_global(volatile __global uint16_t *p, uint16_t x);+SCALAR_FUN_ATTR uint16_t atomic_umin_i16_shared(volatile __local uint16_t *p, uint16_t x);+SCALAR_FUN_ATTR int16_t atomic_and_i16_global(volatile __global int16_t *p, int16_t x);+SCALAR_FUN_ATTR int16_t atomic_and_i16_shared(volatile __local int16_t *p, int16_t x);+SCALAR_FUN_ATTR int16_t atomic_or_i16_global(volatile __global int16_t *p, int16_t x);+SCALAR_FUN_ATTR int16_t atomic_or_i16_shared(volatile __local int16_t *p, int16_t x);+SCALAR_FUN_ATTR int16_t atomic_xor_i16_global(volatile __global int16_t *p, int16_t x);+SCALAR_FUN_ATTR int16_t atomic_xor_i16_shared(volatile __local int16_t *p, int16_t x);++SCALAR_FUN_ATTR int16_t atomic_cmpxchg_i16_global(volatile __global int16_t *p,+                                                  int16_t cmp, int16_t val) {+  int offset = ((uintptr_t)p >> 1 & 1);+  volatile __global int32_t *p32 = (volatile __global int32_t*)((uintptr_t)p & ~0x3);++  int shift = offset * 16;+  int32_t mask = 0xffff << shift;+  int32_t shifted_val = val << shift;+  int32_t shifted_cmp = cmp << shift;++  uint32_t old = shifted_cmp;+  uint32_t upd = shifted_val;+  uint32_t got;++  while ((got=atomic_cmpxchg_i32_global(p32, old, upd)) != old) {+    old = got;+    upd = (old & ~mask) | shifted_val;+  }++  return old >> shift;+}++SCALAR_FUN_ATTR int16_t atomic_cmpxchg_i16_shared(volatile __local int16_t *p,+                                                  int16_t cmp, int16_t val) {+  int offset = ((uintptr_t)p >> 1 & 1);+  volatile __local int32_t *p32 = (volatile __local int32_t*)((uintptr_t)p & ~0x3);++  int shift = offset * 16;+  int32_t mask = 0xffff << shift;+  int32_t shifted_val = val << shift;+  int32_t shifted_cmp = cmp << shift;++  uint32_t old = shifted_cmp;+  uint32_t upd = shifted_val;+  uint32_t got;++  while ((got=atomic_cmpxchg_i32_shared(p32, old, upd)) != old) {+    old = got;+    upd = (old & ~mask) | shifted_val;+  }++  return old >> shift;+}++// Convenience macro for arithmetic.+#define DEFINE_16BIT_ATOMIC(name, T, op)                                \+  SCALAR_FUN_ATTR T                                                     \+  atomic_##name##_i16_global(volatile __global T *p, T val) {           \+    int offset = ((uintptr_t)p >> 1 & 1);                               \+    volatile __global int32_t *p32 = (volatile __global int32_t*)((uintptr_t)p & ~0x3); \+    int shift = offset * 16;                                            \+    int32_t mask = 0xffff << shift;                                     \+    int32_t old = 0;                                                    \+    int32_t upd = mask & (op(old >> shift, val) << shift);              \+    int32_t saw;                                                        \+    while ((saw=atomic_cmpxchg_i32_global(p32, old, upd)) != old) {     \+      old = saw;                                                        \+      upd = (old & ~mask) | ((op(old >> shift, val)) << shift);         \+    }                                                                   \+    return old >> shift;                                                \+  }                                                                     \+  SCALAR_FUN_ATTR T                                                     \+  atomic_##name##_i16_shared(volatile __local T *p, T val) {            \+    int offset = ((uintptr_t)p >> 1 & 1);                               \+    volatile __local int32_t *p32 = (volatile __local int32_t*)((uintptr_t)p & ~0x3); \+    int shift = offset * 16;                                            \+    int32_t mask = 0xffff << shift;                                     \+    int32_t old = 0;                                                    \+    int32_t upd = mask & ((op(old >> shift, val)) << shift);            \+    int32_t saw;                                                        \+    while ((saw=atomic_cmpxchg_i32_shared(p32, old, upd)) != old) {     \+      old = saw;                                                        \+      upd = (old & ~mask) | ((op(old >> shift, val)) << shift);         \+    }                                                                   \+    return old >> shift;                                                \+  }++DEFINE_16BIT_ATOMIC(add, int16_t, add16);+DEFINE_16BIT_ATOMIC(smax, int16_t, smax16);+DEFINE_16BIT_ATOMIC(smin, int16_t, smin16);+DEFINE_16BIT_ATOMIC(umax, uint16_t, umax16);+DEFINE_16BIT_ATOMIC(umin, uint16_t, umin16);++SCALAR_FUN_ATTR int16_t atomic_and_i16_global(volatile __global int16_t *p, int16_t val) {+  volatile __global int32_t *p32 = (volatile __global int32_t*)((uintptr_t)p & ~0x3);+  int shift = ((uintptr_t)p >> 1 & 1) * 16;+  int32_t mask = 0xffff << shift;+  return atomic_and_i32_global(p32, ~mask | (val<<shift)) >> shift;+}++SCALAR_FUN_ATTR int16_t atomic_and_i16_shared(volatile __local int16_t *p, int16_t val) {+  volatile __local int32_t *p32 = (volatile __local int32_t*)((uintptr_t)p & ~0x3);+  int shift = ((uintptr_t)p >> 1 & 1) * 16;+  int32_t mask = 0xffff << shift;+  return atomic_and_i32_shared(p32, ~mask | (val<<shift)) >> shift;+}++SCALAR_FUN_ATTR int16_t atomic_or_i16_global(volatile __global int16_t *p, int16_t val) {+  volatile __global int32_t *p32 = (volatile __global int32_t*)((uintptr_t)p & ~0x3);+  int shift = ((uintptr_t)p >> 1 & 1) * 16;+  return atomic_or_i32_global(p32, (uint16_t)val<<shift) >> shift;+}++SCALAR_FUN_ATTR int16_t atomic_or_i16_shared(volatile __local int16_t *p, int16_t val) {+  volatile __local int32_t *p32 = (volatile __local int32_t*)((uintptr_t)p & ~0x3);+  int shift = ((uintptr_t)p >> 1 & 1) * 16;+  return atomic_or_i32_shared(p32, (uint16_t)val<<shift) >> shift;+}++SCALAR_FUN_ATTR int16_t atomic_xor_i16_global(volatile __global int16_t *p, int16_t val) {+  volatile __global int32_t *p32 = (volatile __global int32_t*)((uintptr_t)p & ~0x3);+  int shift = ((uintptr_t)p >> 1 & 1) * 16;+  return atomic_xor_i32_global(p32, (uint16_t)val<<shift) >> shift;+}++SCALAR_FUN_ATTR int16_t atomic_xor_i16_shared(volatile __local int16_t *p, int16_t val) {+  volatile __local int32_t *p32 = (volatile __local int32_t*)((uintptr_t)p & ~0x3);+  int shift = ((uintptr_t)p >> 1 & 1) * 16;+  return atomic_xor_i32_shared(p32, (uint16_t)val<<shift) >> shift;+}++SCALAR_FUN_ATTR f16 atomic_fadd_f16_global(volatile __global uint16_t *p, f16 val) {+  int offset = ((uintptr_t)p >> 1 & 1);+  volatile __global int32_t *p32 = (volatile __global int32_t*)((uintptr_t)p & ~0x3);+  int shift = offset * 16;+  int32_t mask = 0xffff << shift;+  int32_t old = 0;+  int32_t upd = mask & ((int32_t)futrts_to_bits16(val) << shift);+  int32_t saw;+  while ((saw=atomic_cmpxchg_i32_global(p32, old, upd)) != old) {+    old = saw;+    upd = (old & ~mask) | (int32_t)futrts_to_bits16(futrts_from_bits16((uint32_t)old >> shift) + val) << shift;+  }+  return futrts_from_bits16((uint32_t)old >> shift);+}++SCALAR_FUN_ATTR f16 atomic_fadd_f16_shared(volatile __local uint16_t *p, f16 val) {+  int offset = ((uintptr_t)p >> 1 & 1);+  volatile __local int32_t *p32 = (volatile __local int32_t*)((uintptr_t)p & ~0x3);+  int shift = offset * 16;+  int32_t mask = 0xffff << shift;+  int32_t old = 0;+  int32_t upd = mask & ((int32_t)futrts_to_bits16(val) << shift);+  int32_t saw;+  while ((saw=atomic_cmpxchg_i32_shared(p32, old, upd)) != old) {+    old = saw;+    upd = (old & ~mask) | (int32_t)futrts_to_bits16(futrts_from_bits16((uint32_t)old >> shift) + val) << shift;+  }+  return futrts_from_bits16((uint32_t)old >> shift);+}++// End of atomics16.h

@@ -0,0 +1,235 @@+// 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_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_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_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_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_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_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_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_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_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_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_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)+  return atomicExch((int32_t*)p, x);+#else+  return atomic_xor(p, x);+#endif+}++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+  return atomic_xor(p, x);+#endif+}++SCALAR_FUN_ATTR int32_t atomic_cmpxchg_i32_global(volatile __global int32_t *p,+                                                  int32_t cmp, int32_t val) {+#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)+  return atomicCAS((int32_t*)p, cmp, val);+#else+  return atomic_cmpxchg(p, cmp, val);+#endif+}++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);+#else+  return atomic_cmpxchg(p, cmp, val);+#endif+}++SCALAR_FUN_ATTR int32_t atomic_add_i32_global(volatile __global int32_t *p, int32_t x) {+#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)+  return atomicAdd((int32_t*)p, x);+#else+  return atomic_add(p, x);+#endif+}++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+  return atomic_add(p, x);+#endif+}++SCALAR_FUN_ATTR float atomic_fadd_f32_global(volatile __global float *p, float x) {+#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)+  return atomicAdd((float*)p, x);+  // On OpenCL, use technique from+  // https://pipinspace.github.io/blog/atomic-float-addition-in-opencl.html+#elif defined(cl_nv_pragma_unroll)+  // use hardware-supported atomic addition on Nvidia GPUs with inline+  // PTX assembly+  float ret;+  asm volatile("atom.global.add.f32 %0,[%1],%2;":"=f"(ret):"l"(p),"f"(x):"memory");+  return ret;+#elif defined(__opencl_c_ext_fp32_global_atomic_add)+  // use hardware-supported atomic addition on some Intel GPUs+  return atomic_fetch_add_explicit((volatile __global atomic_float*)p,+                                   x,+                                   memory_order_relaxed);+#elif __has_builtin(__builtin_amdgcn_global_atomic_fadd_f32)+  // use hardware-supported atomic addition on some AMD GPUs+  return __builtin_amdgcn_global_atomic_fadd_f32(p, x);+#else+  // fallback emulation:+  // https://forums.developer.nvidia.com/t/atomicadd-float-float-atomicmul-float-float/14639/5+  float old = x;+  float ret;+  while ((old=atomic_xchg(p, ret=atomic_xchg(p, 0.0f)+old))!=0.0f);+  return ret;+#endif+}++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+  union { int32_t i; float f; } old;+  union { int32_t i; float f; } assumed;+  old.f = *p;+  do {+    assumed.f = old.f;+    old.f = old.f + x;+    old.i = atomic_cmpxchg_i32_shared((volatile __local int32_t*)p, assumed.i, old.i);+  } while (assumed.i != old.i);+  return old.f;+#endif+}++SCALAR_FUN_ATTR int32_t atomic_smax_i32_global(volatile __global int32_t *p, int32_t x) {+#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)+  return atomicMax((int32_t*)p, x);+#else+  return atomic_max(p, x);+#endif+}++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+  return atomic_max(p, x);+#endif+}++SCALAR_FUN_ATTR int32_t atomic_smin_i32_global(volatile __global int32_t *p, int32_t x) {+#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)+  return atomicMin((int32_t*)p, x);+#else+  return atomic_min(p, x);+#endif+}++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+  return atomic_min(p, x);+#endif+}++SCALAR_FUN_ATTR uint32_t atomic_umax_i32_global(volatile __global uint32_t *p, uint32_t x) {+#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)+  return atomicMax((uint32_t*)p, x);+#else+  return atomic_max(p, x);+#endif+}++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+  return atomic_max(p, x);+#endif+}++SCALAR_FUN_ATTR uint32_t atomic_umin_i32_global(volatile __global uint32_t *p, uint32_t x) {+#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)+  return atomicMin((uint32_t*)p, x);+#else+  return atomic_min(p, x);+#endif+}++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+  return atomic_min(p, x);+#endif+}++SCALAR_FUN_ATTR int32_t atomic_and_i32_global(volatile __global int32_t *p, int32_t x) {+#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)+  return atomicAnd((int32_t*)p, x);+#else+  return atomic_and(p, x);+#endif+}++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+  return atomic_and(p, x);+#endif+}++SCALAR_FUN_ATTR int32_t atomic_or_i32_global(volatile __global int32_t *p, int32_t x) {+#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)+  return atomicOr((int32_t*)p, x);+#else+  return atomic_or(p, x);+#endif+}++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+  return atomic_or(p, x);+#endif+}++SCALAR_FUN_ATTR int32_t atomic_xor_i32_global(volatile __global int32_t *p, int32_t x) {+#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)+  return atomicXor((int32_t*)p, x);+#else+  return atomic_xor(p, x);+#endif+}++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+  return atomic_xor(p, x);+#endif+}++// End of atomics.h

@@ -0,0 +1,285 @@+// Start of atomics64.h++#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_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_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_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_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_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_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_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_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_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_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_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) {+#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)+  return atomicExch((unsigned long long*)p, x);+#else+  return atom_xor(p, x);+#endif+}++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((unsigned long long*)p, x);+#else+  return atom_xor(p, x);+#endif+}++SCALAR_FUN_ATTR int64_t atomic_cmpxchg_i64_global(volatile __global int64_t *p,+                                                         int64_t cmp, int64_t val) {+#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)+  return atomicCAS((unsigned long long*)p, cmp, val);+#else+  return atom_cmpxchg(p, cmp, val);+#endif+}++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((unsigned long long*)p, cmp, val);+#else+  return atom_cmpxchg(p, cmp, val);+#endif+}++SCALAR_FUN_ATTR int64_t atomic_add_i64_global(volatile __global int64_t *p, int64_t x) {+#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)+  return atomicAdd((unsigned long long*)p, x);+#else+  return atom_add(p, x);+#endif+}++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((unsigned long long*)p, x);+#else+  return atom_add(p, x);+#endif+}++#ifdef FUTHARK_F64_ENABLED++SCALAR_FUN_ATTR double atomic_fadd_f64_global(volatile __global double *p, double x) {+#if defined(FUTHARK_CUDA) && __CUDA_ARCH__ >= 600 || defined(FUTHARK_HIP)+  return atomicAdd((double*)p, x);+  // On OpenCL, use technique from+  // https://pipinspace.github.io/blog/atomic-float-addition-in-opencl.html+#elif defined(cl_nv_pragma_unroll)+  // use hardware-supported atomic addition on Nvidia GPUs with inline+  // PTX assembly+  double ret;+  asm volatile("atom.global.add.f64 %0,[%1],%2;":"=d"(ret):"l"(p),"d"(x):"memory");+  return ret;+#elif __has_builtin(__builtin_amdgcn_global_atomic_fadd_f64)+  // use hardware-supported atomic addition on some AMD GPUs+  return __builtin_amdgcn_global_atomic_fadd_f64(p, x);+#else+  // fallback emulation:+  // https://forums.developer.nvidia.com/t/atomicadd-float-float-atomicmul-float-float/14639/5+  union {int64_t i; double f;} old;+  union {int64_t i; double f;} ret;+  old.f = x;+  while (1) {+    ret.i = atom_xchg((volatile __global int64_t*)p, (int64_t)0);+    ret.f += old.f;+    old.i = atom_xchg((volatile __global int64_t*)p, ret.i);+    if (old.i == 0) {+      break;+    }+  }+  return ret.f;+#endif+}++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+  union { int64_t i; double f; } old;+  union { int64_t i; double f; } assumed;+  old.f = *p;+  do {+    assumed.f = old.f;+    old.f = old.f + x;+    old.i = atomic_cmpxchg_i64_shared((volatile __local int64_t*)p, assumed.i, old.i);+  } while (assumed.i != old.i);+  return old.f;+#endif+}++#endif++SCALAR_FUN_ATTR int64_t atomic_smax_i64_global(volatile __global int64_t *p, int64_t x) {+#if defined(FUTHARK_CUDA)+  return atomicMax((long long*)p, x);+#elif defined(FUTHARK_HIP)+  // Currentely missing in HIP; probably a temporary oversight.+  int64_t old = *p, assumed;+  do {+    assumed = old;+    old = smax64(old, x);+    old = atomic_cmpxchg_i64_global((volatile __global int64_t*)p, assumed, old);+  } while (assumed != old);+  return old;+#else+  return atom_max(p, x);+#endif+}++SCALAR_FUN_ATTR int64_t atomic_smax_i64_shared(volatile __local int64_t *p, int64_t x) {+#if defined(FUTHARK_CUDA)+  return atomicMax((long long*)p, x);+#elif defined(FUTHARK_HIP)+  // Currentely missing in HIP; probably a temporary oversight.+  int64_t old = *p, assumed;+  do {+    assumed = old;+    old = smax64(old, x);+    old = atomic_cmpxchg_i64_shared((volatile __local int64_t*)p, assumed, old);+  } while (assumed != old);+  return old;+#else+  return atom_max(p, x);+#endif+}++SCALAR_FUN_ATTR int64_t atomic_smin_i64_global(volatile __global int64_t *p, int64_t x) {+#if defined(FUTHARK_CUDA)+  return atomicMin((long long*)p, x);+#elif defined(FUTHARK_HIP)+  // Currentely missing in HIP; probably a temporary oversight.+  int64_t old = *p, assumed;+  do {+    assumed = old;+    old = smin64(old, x);+    old = atomic_cmpxchg_i64_global((volatile __global int64_t*)p, assumed, old);+  } while (assumed != old);+  return old;+#else+  return atom_min(p, x);+#endif+}++SCALAR_FUN_ATTR int64_t atomic_smin_i64_shared(volatile __local int64_t *p, int64_t x) {+#if defined(FUTHARK_CUDA)+  return atomicMin((long long*)p, x);+#elif defined(FUTHARK_HIP)+  // Currentely missing in HIP; probably a temporary oversight.+  int64_t old = *p, assumed;+  do {+    assumed = old;+    old = smin64(old, x);+    old = atomic_cmpxchg_i64_shared((volatile __local int64_t*)p, assumed, old);+  } while (assumed != old);+  return old;+#else+  return atom_min(p, x);+#endif+}++SCALAR_FUN_ATTR uint64_t atomic_umax_i64_global(volatile __global uint64_t *p, uint64_t x) {+#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)+  return atomicMax((unsigned long long*)p, x);+#else+  return atom_max(p, x);+#endif+}++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((unsigned long long*)p, x);+#else+  return atom_max(p, x);+#endif+}++SCALAR_FUN_ATTR uint64_t atomic_umin_i64_global(volatile __global uint64_t *p, uint64_t x) {+#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)+  return atomicMin((unsigned long long*)p, x);+#else+  return atom_min(p, x);+#endif+}++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((unsigned long long*)p, x);+#else+  return atom_min(p, x);+#endif+}++SCALAR_FUN_ATTR int64_t atomic_and_i64_global(volatile __global int64_t *p, int64_t x) {+#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)+  return atomicAnd((unsigned long long*)p, x);+#else+  return atom_and(p, x);+#endif+}++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((unsigned long long*)p, x);+#else+  return atom_and(p, x);+#endif+}++SCALAR_FUN_ATTR int64_t atomic_or_i64_global(volatile __global int64_t *p, int64_t x) {+#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)+  return atomicOr((unsigned long long*)p, x);+#else+  return atom_or(p, x);+#endif+}++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((unsigned long long*)p, x);+#else+  return atom_or(p, x);+#endif+}++SCALAR_FUN_ATTR int64_t atomic_xor_i64_global(volatile __global int64_t *p, int64_t x) {+#if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)+  return atomicXor((unsigned long long*)p, x);+#else+  return atom_xor(p, x);+#endif+}++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((unsigned long long*)p, x);+#else+  return atom_xor(p, x);+#endif+}++#endif // defined(FUTHARK_CUDA) || defined(FUTHARK_HIP) || defined(cl_khr_int64_base_atomics) && defined(cl_khr_int64_extended_atomics)++// End of atomics64.h

@@ -0,0 +1,145 @@+// Start of atomics8.h++SCALAR_FUN_ATTR int8_t atomic_cmpxchg_i8_global(volatile __global int8_t *p,+                                                int8_t cmp, int8_t val);+SCALAR_FUN_ATTR int8_t atomic_cmpxchg_i8_shared(volatile __local int8_t *p,+                                                int8_t cmp, int8_t val);+SCALAR_FUN_ATTR int8_t atomic_add_i8_global(volatile __global int8_t *p, int8_t x);+SCALAR_FUN_ATTR int8_t atomic_add_i8_shared(volatile __local int8_t *p, int8_t x);+SCALAR_FUN_ATTR int8_t atomic_smax_i8_global(volatile __global int8_t *p, int8_t x);+SCALAR_FUN_ATTR int8_t atomic_smax_i8_shared(volatile __local int8_t *p, int8_t x);+SCALAR_FUN_ATTR int8_t atomic_smin_i8_global(volatile __global int8_t *p, int8_t x);+SCALAR_FUN_ATTR int8_t atomic_smin_i8_shared(volatile __local int8_t *p, int8_t x);+SCALAR_FUN_ATTR uint8_t atomic_umax_i8_global(volatile __global uint8_t *p, uint8_t x);+SCALAR_FUN_ATTR uint8_t atomic_umax_i8_shared(volatile __local uint8_t *p, uint8_t x);+SCALAR_FUN_ATTR uint8_t atomic_umin_i8_global(volatile __global uint8_t *p, uint8_t x);+SCALAR_FUN_ATTR uint8_t atomic_umin_i8_shared(volatile __local uint8_t *p, uint8_t x);+SCALAR_FUN_ATTR int8_t atomic_and_i8_global(volatile __global int8_t *p, int8_t x);+SCALAR_FUN_ATTR int8_t atomic_and_i8_shared(volatile __local int8_t *p, int8_t x);+SCALAR_FUN_ATTR int8_t atomic_or_i8_global(volatile __global int8_t *p, int8_t x);+SCALAR_FUN_ATTR int8_t atomic_or_i8_shared(volatile __local int8_t *p, int8_t x);+SCALAR_FUN_ATTR int8_t atomic_xor_i8_global(volatile __global int8_t *p, int8_t x);+SCALAR_FUN_ATTR int8_t atomic_xor_i8_shared(volatile __local int8_t *p, int8_t x);++SCALAR_FUN_ATTR int8_t atomic_cmpxchg_i8_global(volatile __global int8_t *p,+                                                int8_t cmp, int8_t val) {+  int offset = ((uintptr_t)p & 3);+  volatile __global int32_t *p32 = (volatile __global int32_t*)((uintptr_t)p & ~0x3);++  int shift = offset * 8;+  int32_t mask = 0xff << shift;+  int32_t shifted_val = val << shift;+  int32_t shifted_cmp = cmp << shift;++  uint32_t old = shifted_cmp;+  uint32_t upd = shifted_val;+  uint32_t got;++  while ((got=atomic_cmpxchg_i32_global(p32, old, upd)) != old) {+    old = got;+    upd = (old & ~mask) | shifted_val;+  }++  return old >> shift;+}++SCALAR_FUN_ATTR int8_t atomic_cmpxchg_i8_shared(volatile __local int8_t *p,+                                                int8_t cmp, int8_t val) {+  int offset = ((uintptr_t)p >> 1 & 3);+  volatile __local int32_t *p32 = (volatile __local int32_t*)((uintptr_t)p & ~0x3);++  int shift = offset * 8;+  int32_t mask = 0xff << shift;+  int32_t shifted_val = val << shift;+  int32_t shifted_cmp = cmp << shift;++  uint32_t old = shifted_cmp;+  uint32_t upd = shifted_val;+  uint32_t got;++  while ((got=atomic_cmpxchg_i32_shared(p32, old, upd)) != old) {+    old = got;+    upd = (old & ~mask) | shifted_val;+  }++  return old >> shift;+}++// Convenience macro for arithmetic.+#define DEFINE_8BIT_ATOMIC(name, T, op)                                 \+  SCALAR_FUN_ATTR T                                                     \+  atomic_##name##_i8_global(volatile __global T *p, T val) {            \+    int offset = ((uintptr_t)p & 3);                                    \+    volatile __global int32_t *p32 = (volatile __global int32_t*)((uintptr_t)p & ~0x3); \+    int shift = offset * 8;                                             \+    int32_t mask = 0xff << shift;                                       \+    int32_t old = 0;                                                    \+    int32_t upd = mask & (op(old >> shift, val) << shift);              \+    int32_t saw;                                                        \+    while ((saw=atomic_cmpxchg_i32_global(p32, old, upd)) != old) {     \+      old = saw;                                                        \+      upd = (old & ~mask) | ((op(old >> shift, val)) << shift);         \+    }                                                                   \+    return old >> shift;                                                \+  }                                                                     \+  SCALAR_FUN_ATTR T                                                     \+  atomic_##name##_i8_shared(volatile __local T *p, T val) {             \+    int offset = ((uintptr_t)p & 3);                                    \+    volatile __local int32_t *p32 = (volatile __local int32_t*)((uintptr_t)p & ~0x3); \+    int shift = offset * 8;                                             \+    int32_t mask = 0xff << shift;                                       \+    int32_t old = 0;                                                    \+    int32_t upd = mask & ((op(old >> shift, val)) << shift);            \+    int32_t saw;                                                        \+    while ((saw=atomic_cmpxchg_i32_shared(p32, old, upd)) != old) {     \+      old = saw;                                                        \+      upd = (old & ~mask) | ((op(old >> shift, val)) << shift);         \+    }                                                                   \+    return old >> shift;                                                \+  }++DEFINE_8BIT_ATOMIC(add, int8_t, add8);+DEFINE_8BIT_ATOMIC(smax, int8_t, smax8);+DEFINE_8BIT_ATOMIC(smin, int8_t, smin8);+DEFINE_8BIT_ATOMIC(umax, uint8_t, umax8);+DEFINE_8BIT_ATOMIC(umin, uint8_t, umin8);++SCALAR_FUN_ATTR int8_t atomic_and_i8_global(volatile __global int8_t *p, int8_t val) {+  volatile __global int32_t *p32 = (volatile __global int32_t*)((uintptr_t)p & ~0x3);+  int shift = ((uintptr_t)p & 3) * 8;+  int32_t mask = 0xff << shift;+  return atomic_and_i32_global(p32, ~mask | (val<<shift)) >> shift;+}++SCALAR_FUN_ATTR int8_t atomic_and_i8_shared(volatile __local int8_t *p, int8_t val) {+  volatile __local int32_t *p32 = (volatile __local int32_t*)((uintptr_t)p & ~0x3);+  int shift = ((uintptr_t)p & 3) * 8;+  int32_t mask = 0xff << shift;+  return atomic_and_i32_shared(p32, ~mask | (val<<shift)) >> shift;+}++SCALAR_FUN_ATTR int8_t atomic_or_i8_global(volatile __global int8_t *p, int8_t val) {+  volatile __global int32_t *p32 = (volatile __global int32_t*)((uintptr_t)p & ~0x3);+  int shift = ((uintptr_t)p & 3) * 8;+  return atomic_or_i32_global(p32, (uint8_t)val<<shift) >> shift;+}++SCALAR_FUN_ATTR int8_t atomic_or_i8_shared(volatile __local int8_t *p, int8_t val) {+  volatile __local int32_t *p32 = (volatile __local int32_t*)((uintptr_t)p & ~0x3);+  int shift = ((uintptr_t)p & 3) * 8;+  return atomic_or_i32_shared(p32, (uint8_t)val<<shift) >> shift;+}++SCALAR_FUN_ATTR int8_t atomic_xor_i8_global(volatile __global int8_t *p, int8_t val) {+  volatile __global int32_t *p32 = (volatile __global int32_t*)((uintptr_t)p & ~0x3);+  int shift = ((uintptr_t)p & 3) * 8;+  return atomic_xor_i32_global(p32, (uint8_t)val<<shift) >> shift;+}++SCALAR_FUN_ATTR int8_t atomic_xor_i8_shared(volatile __local int8_t *p, int8_t val) {+  volatile __local int32_t *p32 = (volatile __local int32_t*)((uintptr_t)p & ~0x3);+  int shift = ((uintptr_t)p & 3) * 8;+  return atomic_xor_i32_shared(p32, (uint8_t)val<<shift) >> shift;+}++// End of atomics8.h

@@ -242,6 +242,10 @@     min_size = sizeof(int);   } +  // Round up the allocation to be at least divisible by 4, because that is+  // assumed by the code generator.+  min_size = (min_size+3) & ~3;+   gpu_mem* memptr;   if (free_list_find(&ctx->gpu_free_list, min_size, tag, size_out, (fl_mem*)&memptr) == 0) {     // Successfully found a free block.  Is it big enough?

@@ -25,6 +25,14 @@ typedef unsigned long uint64_t; #endif +#if defined(__CUDACC_RTC__)+typedef uint64_t uintptr_t;+#endif++#if defined(__HIPCC_RTC__)+typedef uint64_t uintptr_t;+#endif+ #define __global #define __local #define __private

@@ -368,8 +368,10 @@   def opencl_alloc(self, min_size, tag):-    min_size = 1 if min_size == 0 else min_size+    min_size = 4 if min_size == 0 else min_size     assert min_size > 0+    # Round up to a multiple of four.+    min_size = ((min_size + 3) // 4) * 4     return self.pool.allocate(min_size)  

@@ -226,9 +226,9 @@       arr_tan <- tangent arr       arrs_tans <- mapM tangent arrs       addStm $ Let pat_tan aux $ BasicOp $ Concat d (arr_tan :| arrs_tans) w-    Manifest ds arr -> do+    Manifest arr ds -> do       arr_tan <- tangent arr-      addStm $ Let pat_tan aux $ BasicOp $ Manifest ds arr_tan+      addStm $ Let pat_tan aux $ BasicOp $ Manifest arr_tan ds     Iota n _ _ it -> do       addStm $ Let pat_tan aux $ BasicOp $ Replicate (Shape [n]) (intConst it 0)     Replicate n x -> do@@ -236,12 +236,12 @@       addStm $ Let pat_tan aux $ BasicOp $ Replicate n x_tan     Scratch t shape ->       addStm $ Let pat_tan aux $ BasicOp $ Scratch t shape-    Reshape k reshape arr -> do+    Reshape arr reshape -> do       arr_tan <- tangent arr-      addStm $ Let pat_tan aux $ BasicOp $ Reshape k reshape arr_tan-    Rearrange perm arr -> do+      addStm $ Let pat_tan aux $ BasicOp $ Reshape arr_tan reshape+    Rearrange arr perm -> do       arr_tan <- tangent arr-      addStm $ Let pat_tan aux $ BasicOp $ Rearrange perm arr_tan+      addStm $ Let pat_tan aux $ BasicOp $ Rearrange arr_tan perm     _ -> error $ "basicFwd: Unsupported op " ++ prettyString op  fwdLambda :: Lambda SOACS -> ADM (Lambda SOACS)

@@ -117,20 +117,20 @@       (_pat_v, pat_adj) <- commonBasicOp pat aux e m       returnSweepCode $ updateSubExpAdj se pat_adj     ---    Reshape k _ arr -> do+    Reshape arr newshape -> do       (_pat_v, pat_adj) <- commonBasicOp pat aux e m       returnSweepCode $ do         arr_shape <- arrayShape <$> lookupType arr         void $           updateAdj arr <=< letExp "adj_reshape" . BasicOp $-            Reshape k arr_shape pat_adj+            Reshape pat_adj (reshapeAll (newShape newshape) arr_shape)     ---    Rearrange perm arr -> do+    Rearrange arr perm -> do       (_pat_v, pat_adj) <- commonBasicOp pat aux e m       returnSweepCode $         void $           updateAdj arr <=< letExp "adj_rearrange" . BasicOp $-            Rearrange (rearrangeInverse perm) pat_adj+            Rearrange pat_adj (rearrangeInverse perm)     --     Replicate (Shape []) (Var se) -> do       (_pat_v, pat_adj) <- commonBasicOp pat aux e m@@ -145,7 +145,7 @@         n <- letSubExp "rep_size" =<< foldBinOp (Mul Int64 OverflowUndef) (intConst Int64 1) ns         pat_adj_flat <-           letExp (baseString pat_adj <> "_flat") . BasicOp $-            Reshape ReshapeArbitrary (Shape $ n : arrayDims x_t) pat_adj+            Reshape pat_adj (reshapeAll (Shape ns) (Shape $ n : arrayDims x_t))         reduce <- reduceSOAC [Reduce Commutative lam [ne]]         updateSubExpAdj x           =<< letExp "rep_contrib" (Op $ Screma n [pat_adj_flat] reduce)@@ -170,7 +170,7 @@          zipWithM_ updateAdj (arr : arrs) slices     ---    Manifest _ se -> do+    Manifest se _ -> do       (_pat_v, pat_adj) <- commonBasicOp pat aux e m       returnSweepCode $ void $ updateAdj se pat_adj     --

@@ -304,8 +304,10 @@   scatter_inps <- do     -- traverse (letExp "flat" . BasicOp . Reshape [DimNew q]) $ inds ++ [vs_bar_p]     -- ToDo: Cosmin asks: is the below the correct translation of the line above?-    traverse (letExp "flat" . BasicOp . Reshape ReshapeArbitrary (Shape [q])) $-      inds ++ [vs_bar_p]+    forM (inds ++ [vs_bar_p]) $ \v -> do+      v_t <- lookupType v+      letExp "flat" . BasicOp . Reshape v $+        reshapeAll (arrayShape v_t) (Shape [q])    f'' <- mkIdentityLambda $ replicate nr_dims (Prim int64) ++ [Prim t]   vs_bar' <-@@ -553,8 +555,8 @@     rank <- arrayRank <$> lookupType vss     let dims = [1, 0] ++ drop 2 [0 .. rank - 1] -    dstT <- letExp "dstT" $ BasicOp $ Rearrange dims dst-    vssT <- letExp "vssT" $ BasicOp $ Rearrange dims vss+    dstT <- letExp "dstT" $ BasicOp $ Rearrange dst dims+    vssT <- letExp "vssT" $ BasicOp $ Rearrange vss dims     t_dstT <- lookupType dstT     t_vssT <- lookupType vssT     t_nes <- lookupType nes@@ -578,11 +580,10 @@               [HistOp (Shape [w]) rf [dst_col_cpy] [Var $ paramName ne] op]               f     histT <--      letExp "histT" $-        Op $-          Screma (arraySize 0 t_dstT) [dstT, vssT, nes] $-            mapSOAC map_lam-    auxing aux . letBindNames [x] . BasicOp $ Rearrange dims histT+      letExp "histT" . Op $+        Screma (arraySize 0 t_dstT) [dstT, vssT, nes] $+          mapSOAC map_lam+    auxing aux . letBindNames [x] . BasicOp $ Rearrange histT dims   foldr (vjpStm ops) m stms  --

@@ -208,7 +208,7 @@     rank <- arrayRank <$> lookupType as     let rear = [1, 0] ++ drop 2 [0 .. rank - 1] -    tran_as <- letExp "tran_as" $ BasicOp $ Rearrange rear as+    tran_as <- letExp "tran_as" $ BasicOp $ Rearrange as rear     ts <- lookupType tran_as     t_ne <- lookupType ne 

@@ -446,7 +446,7 @@      transp_as <-       forM as $ \a ->-        letExp (baseString a ++ "_transp") $ BasicOp $ Rearrange rear a+        letExp (baseString a ++ "_transp") $ BasicOp $ Rearrange a rear      ts <- traverse lookupType transp_as     let n = arraysSize 0 ts@@ -465,7 +465,7 @@         Screma n (transp_as ++ subExpVars ne) (mapSOAC map_lam)      forM (zip ys transp_ys) $ \(y, x) ->-      auxing aux $ letBindNames [y] $ BasicOp $ Rearrange rear x+      auxing aux $ letBindNames [y] $ BasicOp $ Rearrange x rear    foldr (vjpStm ops) m stmts 

@@ -478,12 +478,13 @@           error $ "unhandled: Update (This should NEVER happen) onto " ++ prettyString name         -- Technically, do we need this case?         Concat _ _ length_subexp -> varInfoFromSubExp length_subexp-        Manifest _dim name -> varInfoFromNames ctx $ oneName name+        Manifest name _dim -> varInfoFromNames ctx $ oneName name         Iota end start stride _ -> concatVariableInfos mempty [end, start, stride]         Replicate (Shape shape) value' -> concatVariableInfos mempty (value' : shape)         Scratch _ sers -> concatVariableInfos mempty sers-        Reshape _ (Shape shape_subexp) name -> concatVariableInfos (oneName name) shape_subexp-        Rearrange _ name -> varInfoFromNames ctx $ oneName name+        Reshape name newshape ->+          concatVariableInfos (oneName name) (shapeDims (newShape newshape))+        Rearrange name _ -> varInfoFromNames ctx $ oneName name         UpdateAcc _ name lsubexprs rsubexprs ->           concatVariableInfos (oneName name) (lsubexprs ++ rsubexprs)         FlatIndex name _ -> varInfoFromNames ctx $ oneName name

@@ -215,14 +215,15 @@ -- | Reshape a map nest. It is assumed that any validity tests have -- already been done. Will automatically reshape the inputs -- appropriately.-reshape :: (MonadFreshNames m) => Certs -> Shape -> MapNest -> m MapNest-reshape cs shape (MapNest _ map_lam _ inps) =+reshape :: (MonadFreshNames m) => StmAux () -> Shape -> MapNest -> m MapNest+reshape aux shape (MapNest _ map_lam _ inps) =   descend [] $ stripDims 1 shape   where     w = shapeSize 0 shape     transform p inp =       let shape' = shape <> arrayShape p-          tr = SOAC.Reshape cs ReshapeArbitrary shape'+          inp_shape = arrayShape (SOAC.inputType inp)+          tr = SOAC.Reshape aux $ reshapeAll inp_shape shape'        in SOAC.addTransform tr inp     inps' = zipWith transform (map paramType $ lambdaParams map_lam) inps 

@@ -100,36 +100,26 @@ -- create a list, use 'ArrayTransforms' instead. data ArrayTransform   = -- | A permutation of an otherwise valid input.-    Rearrange Certs [Int]+    Rearrange (StmAux ()) [Int]   | -- | A reshaping of an otherwise valid input.-    Reshape Certs ReshapeKind Shape-  | -- | A reshaping of the outer dimension.-    ReshapeOuter Certs ReshapeKind Shape-  | -- | A reshaping of everything but the outer dimension.-    ReshapeInner Certs ReshapeKind Shape+    Reshape (StmAux ()) (NewShape SubExp)   | -- | Replicate the rows of the array a number of times.-    Replicate Certs Shape+    Replicate (StmAux ()) Shape   | -- | An array indexing operation.-    Index Certs (Slice SubExp)+    Index (StmAux ()) (Slice SubExp)   deriving (Show, Eq, Ord)  instance FreeIn ArrayTransform where   freeIn' (Rearrange cs _) = freeIn' cs-  freeIn' (Reshape cs _ shape) = freeIn' cs <> freeIn' shape-  freeIn' (ReshapeOuter cs _ shape) = freeIn' cs <> freeIn' shape-  freeIn' (ReshapeInner cs _ shape) = freeIn' cs <> freeIn' shape+  freeIn' (Reshape cs shape) = freeIn' cs <> freeIn' shape   freeIn' (Replicate cs shape) = freeIn' cs <> freeIn' shape   freeIn' (Index cs slice) = freeIn' cs <> freeIn' slice  instance Substitute ArrayTransform where   substituteNames substs (Rearrange cs xs) =     Rearrange (substituteNames substs cs) xs-  substituteNames substs (Reshape cs k ses) =-    Reshape (substituteNames substs cs) k (substituteNames substs ses)-  substituteNames substs (ReshapeOuter cs k ses) =-    ReshapeOuter (substituteNames substs cs) k (substituteNames substs ses)-  substituteNames substs (ReshapeInner cs k ses) =-    ReshapeInner (substituteNames substs cs) k (substituteNames substs ses)+  substituteNames substs (Reshape cs newshape) =+    Reshape (substituteNames substs cs) (substituteNames substs newshape)   substituteNames substs (Replicate cs se) =     Replicate (substituteNames substs cs) (substituteNames substs se)   substituteNames substs (Index cs slice) =@@ -242,11 +232,11 @@ -- an input transformation of an array variable.  If so, return the -- variable and the transformation.  Only 'Rearrange' and 'Reshape' -- are possible to express this way.-transformFromExp :: Certs -> Exp rep -> Maybe (VName, ArrayTransform)-transformFromExp cs (BasicOp (Futhark.Rearrange perm v)) =+transformFromExp :: StmAux () -> Exp rep -> Maybe (VName, ArrayTransform)+transformFromExp cs (BasicOp (Futhark.Rearrange v perm)) =   Just (v, Rearrange cs perm)-transformFromExp cs (BasicOp (Futhark.Reshape k shape v)) =-  Just (v, Reshape cs k shape)+transformFromExp cs (BasicOp (Futhark.Reshape v shape)) =+  Just (v, Reshape cs shape) transformFromExp cs (BasicOp (Futhark.Replicate shape (Var v))) =   Just (v, Replicate cs shape) transformFromExp cs (BasicOp (Futhark.Index v slice)) =@@ -254,20 +244,14 @@ transformFromExp _ _ = Nothing  -- | Turn an array transform on an array back into an expression.-transformToExp :: (Monad m, HasScope rep m) => ArrayTransform -> VName -> m (Certs, Exp rep)+transformToExp :: (Monad m, HasScope rep m) => ArrayTransform -> VName -> m (StmAux (), Exp rep) transformToExp (Replicate cs n) ia =   pure (cs, BasicOp $ Futhark.Replicate n (Var ia)) transformToExp (Rearrange cs perm) ia = do   r <- arrayRank <$> lookupType ia-  pure (cs, BasicOp $ Futhark.Rearrange (perm ++ [length perm .. r - 1]) ia)-transformToExp (Reshape cs k shape) ia = do-  pure (cs, BasicOp $ Futhark.Reshape k shape ia)-transformToExp (ReshapeOuter cs k shape) ia = do-  shape' <- reshapeOuter shape 1 . arrayShape <$> lookupType ia-  pure (cs, BasicOp $ Futhark.Reshape k shape' ia)-transformToExp (ReshapeInner cs k shape) ia = do-  shape' <- reshapeInner shape 1 . arrayShape <$> lookupType ia-  pure (cs, BasicOp $ Futhark.Reshape k shape' ia)+  pure (cs, BasicOp $ Futhark.Rearrange ia (perm ++ [length perm .. r - 1]))+transformToExp (Reshape cs shape) ia = do+  pure (cs, BasicOp $ Futhark.Reshape ia shape) transformToExp (Index cs slice) ia = do   pure (cs, BasicOp $ Futhark.Index ia slice) @@ -310,8 +294,10 @@ isVarishInput :: Input -> Maybe VName isVarishInput (Input ts v t)   | nullTransforms ts = Just v-  | Reshape cs ReshapeCoerce (Shape [_]) :< ts' <- viewf ts,-    cs == mempty =+  | Reshape aux newshape :< ts' <- viewf ts,+    ReshapeCoerce <- reshapeKind newshape,+    1 <- shapeRank $ newShape newshape,+    stmAuxCerts aux == mempty =       isVarishInput $ Input ts' v t isVarishInput _ = Nothing @@ -327,15 +313,13 @@  applyTransform :: (MonadBuilder m) => ArrayTransform -> VName -> m VName applyTransform tr ia = do-  (cs, e) <- transformToExp tr ia-  certifying cs $ letExp s e+  (aux, e) <- transformToExp tr ia+  auxing aux $ letExp s e   where     s = case tr of       Replicate {} -> "replicate"       Rearrange {} -> "rearrange"       Reshape {} -> "reshape"-      ReshapeOuter {} -> "reshape_outer"-      ReshapeInner {} -> "reshape_inner"       Index {} -> "index"  applyTransforms :: (MonadBuilder m) => ArrayTransforms -> VName -> m VName@@ -367,14 +351,8 @@       arrayOfShape t shape     transformType t (Rearrange _ perm) =       rearrangeType perm t-    transformType t (Reshape _ _ shape) =-      t `setArrayShape` shape-    transformType t (ReshapeOuter _ _ shape) =-      let Shape oldshape = arrayShape t-       in t `setArrayShape` Shape (shapeDims shape ++ drop 1 oldshape)-    transformType t (ReshapeInner _ _ shape) =-      let Shape oldshape = arrayShape t-       in t `setArrayShape` Shape (take 1 oldshape ++ shapeDims shape)+    transformType t (Reshape _ shape) =+      t `setArrayShape` newShape shape     transformType t (Index _ slice) =       t `setArrayShape` sliceShape slice @@ -394,8 +372,11 @@   where     transformRows' inp (Rearrange cs perm) =       addTransform (Rearrange cs (0 : map (+ 1) perm)) inp-    transformRows' inp (Reshape cs k shape) =-      addTransform (ReshapeInner cs k shape) inp+    transformRows' inp (Reshape cs shape) =+      addTransform (Reshape cs newshape) inp+      where+        newshape = reshapeAll inp_shape $ Shape [shapeSize 0 inp_shape] <> newShape shape+        inp_shape = arrayShape $ inputType inp     transformRows' inp (Replicate cs n)       | inputRank inp == 1 =           Rearrange mempty [1, 0]@@ -714,18 +695,8 @@ ppArrayTransform :: PP.Doc a -> ArrayTransform -> PP.Doc a ppArrayTransform e (Rearrange cs perm) =   "rearrange" <> pretty cs <> PP.apply [PP.apply (map pretty perm), e]-ppArrayTransform e (Reshape cs ReshapeArbitrary shape) =+ppArrayTransform e (Reshape cs shape) =   "reshape" <> pretty cs <> PP.apply [pretty shape, e]-ppArrayTransform e (ReshapeOuter cs ReshapeArbitrary shape) =-  "reshape_outer" <> pretty cs <> PP.apply [pretty shape, e]-ppArrayTransform e (ReshapeInner cs ReshapeArbitrary shape) =-  "reshape_inner" <> pretty cs <> PP.apply [pretty shape, e]-ppArrayTransform e (Reshape cs ReshapeCoerce shape) =-  "coerce" <> pretty cs <> PP.apply [pretty shape, e]-ppArrayTransform e (ReshapeOuter cs ReshapeCoerce shape) =-  "coerce_outer" <> pretty cs <> PP.apply [pretty shape, e]-ppArrayTransform e (ReshapeInner cs ReshapeCoerce shape) =-  "coerce_inner" <> pretty cs <> PP.apply [pretty shape, e] ppArrayTransform e (Replicate cs ne) =   "replicate" <> pretty cs <> PP.apply [pretty ne, e] ppArrayTransform e (Index cs slice) =

@@ -287,20 +287,6 @@ subExpAvailable (Var name) = available name subExpAvailable Constant {} = const True -index ::-  (ASTRep rep) =>-  VName ->-  [SubExp] ->-  SymbolTable rep ->-  Maybe Indexed-index name is table = do-  is' <- mapM asPrimExp is-  index' name is' table-  where-    asPrimExp i = do-      Prim t <- lookupSubExpType i table-      pure $ TPrimExp $ primExpFromSubExp t i- index' ::   VName ->   [TPrimExp Int64 VName] ->@@ -311,14 +297,33 @@   case entryType entry of     LetBound entry'       | Just k <--          elemIndex name . patNames . stmPat $-            letBoundStm entry' ->+          elemIndex name . patNames . stmPat $ letBoundStm entry' ->           letBoundIndex entry' k is     FreeVar entry' ->       freeVarIndex entry' name is     LParam entry' -> lparamIndex entry' is     _ -> Nothing +-- | @index arr is vtable@ fully indexes the array @arr@ at position @is@ using+-- information in @vtable@, and produces the symbolic result of the indexing if+-- it can be expressed. This is essentially a form of pull-array indexing.+index ::+  VName ->+  [SubExp] ->+  SymbolTable rep ->+  Maybe Indexed+index name = index' name . map pe64++-- | Like 'index'', but always succeeds, simply returning an 'IndexedArray' of+-- the input if nothing else is possible.+indexNext ::+  VName ->+  [TPrimExp Int64 VName] ->+  SymbolTable rep ->+  Indexed+indexNext name is vtable =+  fromMaybe (IndexedArray mempty name is) $ index' name is vtable+ class IndexOp op where   indexOp ::     (ASTRep rep, IndexOp (Op rep)) =>@@ -335,17 +340,16 @@   (IndexOp (Op rep), ASTRep rep) =>   SymbolTable rep ->   Exp rep ->+  -- | Index of result being indexed in case the expression produces more than+  -- one.   Int ->   IndexArray indexExp vtable (Op op) k is =   indexOp vtable k op is indexExp _ (BasicOp (Iota _ x s to_it)) _ [i] =-  Just $-    Indexed mempty $-      ( sExt to_it (untyped i)-          `mul` primExpFromSubExp (IntType to_it) s-      )-        `add` primExpFromSubExp (IntType to_it) x+  Just . Indexed mempty $+    (sExt to_it (untyped i) `mul` primExpFromSubExp (IntType to_it) s)+      `add` primExpFromSubExp (IntType to_it) x   where     mul = BinOpExp (Mul to_it OverflowWrap)     add = BinOpExp (Add to_it OverflowWrap)@@ -355,20 +359,21 @@       Just $ Indexed mempty $ primExpFromSubExp t v indexExp table (BasicOp (Replicate s (Var v))) _ is = do   guard $ v `available` table-  guard $ s /= mempty-  index' v (drop (shapeRank s) is) table-indexExp table (BasicOp (Reshape _ newshape v)) _ is+  Just $ indexNext v (drop (shapeRank s) is) table+indexExp table (BasicOp (Reshape v newshape)) _ is   | Just oldshape <- arrayDims <$> lookupType v table =       -- TODO: handle coercions more efficiently.       let is' =             reshapeIndex               (map pe64 oldshape)-              (map pe64 $ shapeDims newshape)+              (map pe64 $ shapeDims $ newShape newshape)               is-       in index' v is' table+       in Just $ indexNext v is' table+indexExp table (BasicOp (Rearrange v perm)) _ is =+  Just $ indexNext v (rearrangeShape (rearrangeInverse perm) is) table indexExp table (BasicOp (Index v slice)) _ is = do   guard $ v `available` table-  index' v (adjust (unSlice slice) is) table+  Just $ indexNext v (adjust (unSlice slice) is) table   where     adjust (DimFix j : js') is' =       pe64 j : adjust js' is'

@@ -54,18 +54,19 @@     numpad = 15     mkRows rows =       let longest = foldl max numpad $ map (T.length . fst) rows+          total = sum $ map (evSum . snd) rows           header = headerRow longest           splitter = T.map (const '-') header           bottom =             T.unwords               [ showText (sum (map (evCount . snd) rows)),                 "events with a total runtime of",-                T.pack $ printf "%.2fμs" $ sum $ map (evSum . snd) rows+                T.pack $ printf "%.2fμs" total               ]        in T.unlines $             header               : splitter-              : map (mkRow longest) rows+              : map (mkRow longest total) rows                 <> [splitter, bottom]     headerRow longest =       T.unwords@@ -74,16 +75,18 @@           padLeft numpad "sum",           padLeft numpad "avg",           padLeft numpad "min",-          padLeft numpad "max"+          padLeft numpad "max",+          padLeft numpad "fraction"         ]-    mkRow longest (name, ev) =+    mkRow longest total (name, ev) =       T.unwords         [ padRight longest name,           padLeft numpad (showText (evCount ev)),           padLeft numpad $ T.pack $ printf "%.2fμs" (evSum ev),           padLeft numpad $ T.pack $ printf "%.2fμs" $ evSum ev / fromInteger (evCount ev),           padLeft numpad $ T.pack $ printf "%.2fμs" (evMin ev),-          padLeft numpad $ T.pack $ printf "%.2fμs" (evMax ev)+          padLeft numpad $ T.pack $ printf "%.2fμs" (evMax ev),+          padLeft numpad $ T.pack $ printf "%.4f" (evSum ev / total)         ]  timeline :: [ProfilingEvent] -> T.Text

@@ -194,7 +194,7 @@     maybePipeline SeqMemPipeline = "(seq-mem) "     maybePipeline GpuMemPipeline = "(gpu-mem) "     maybePipeline MCMemPipeline = "(mc-mem) "-    maybePipeline NoPipeline = ""+    maybePipeline NoPipeline = " "      ok (AstMetrics metrics) (name, expected_occurences) =       case M.lookup name metrics of@@ -203,7 +203,7 @@               throwError $                 name                   <> maybePipeline pipeline-                  <> " should have occurred "+                  <> "should have occurred "                   <> showText expected_occurences                   <> " times, but did not occur at all in optimised program."         Just actual_occurences@@ -211,7 +211,7 @@               throwError $                 name                   <> maybePipeline pipeline-                  <> " should have occurred "+                  <> "should have occurred "                   <> showText expected_occurences                   <> " times, but occurred "                   <> showText actual_occurences

@@ -516,7 +516,7 @@        headerDecl InitDecl [C.cedecl|struct futhark_context;|]       headerDecl InitDecl [C.cedecl|struct futhark_context* futhark_context_new(struct futhark_context_config* cfg);|]-      headerDecl InitDecl [C.cedecl|void futhark_context_free(struct futhark_context* cfg);|]+      headerDecl InitDecl [C.cedecl|void futhark_context_free(struct futhark_context* ctx);|]       headerDecl MiscDecl [C.cedecl|int futhark_context_sync(struct futhark_context* ctx);|]        generateTuningParams params

@@ -146,15 +146,22 @@   either (Left . errorBundlePretty) Right . parse (p <* eof) "type name"   where     p :: Parsec Void T.Text T.Text-    p = choice [pArr, pTup, pAtom]+    p = choice [pArr, pTup, pQual]     pArr = do       dims <- some "[]"       (("arr" <> showText (length dims) <> "d_") <>) <$> p     pTup = between "(" ")" $ do       ts <- p `sepBy` pComma       pure $ "tup" <> showText (length ts) <> "_" <> T.intercalate "_" ts-    pAtom = T.pack <$> some (satisfy (`notElem` ("[]{}()," :: String)))+    pAtom = T.pack <$> some (satisfy (`notElem` ("[]{}(),." :: String)))     pComma = void $ "," <* space+    -- Rewrite 'x.y' to 'x_y'.+    pQual = do+      x <- pAtom+      choice+        [ "." >> ((x <> "_") <>) <$> pAtom,+          pure x+        ]  -- | The name of exposed opaque types. opaqueName :: Name -> T.Text

@@ -941,7 +941,7 @@         BinOpExp (Add it OverflowUndef) e' $           BinOpExp (Mul it OverflowUndef) i' s'     copyDWIMFix (patElemName pe) [i] (Var (tvVar x)) []-defCompileBasicOp (Pat [pe]) (Manifest _ src) =+defCompileBasicOp (Pat [pe]) (Manifest src _) =   copyDWIM (patElemName pe) [] (Var src) [] defCompileBasicOp (Pat [pe]) (Concat i (x :| ys) _) = do   offs_glb <- dPrimV "tmp_offs" 0

@@ -65,7 +65,28 @@         (Or Int32, Imp.AtomicOr Int32),         (Xor Int32, Imp.AtomicXor Int32)       ]-    opencl = opencl32 ++ opencl64+    opencl16 =+      [ (Add Int16 OverflowUndef, Imp.AtomicAdd Int16),+        (FAdd Float16, Imp.AtomicFAdd Float16),+        (SMax Int16, Imp.AtomicSMax Int16),+        (SMin Int16, Imp.AtomicSMin Int16),+        (UMax Int16, Imp.AtomicUMax Int16),+        (UMin Int16, Imp.AtomicUMin Int16),+        (And Int16, Imp.AtomicAnd Int16),+        (Or Int16, Imp.AtomicOr Int16),+        (Xor Int16, Imp.AtomicXor Int16)+      ]+    opencl8 =+      [ (Add Int8 OverflowUndef, Imp.AtomicAdd Int8),+        (SMax Int8, Imp.AtomicSMax Int8),+        (SMin Int8, Imp.AtomicSMin Int8),+        (UMax Int8, Imp.AtomicUMax Int8),+        (UMin Int8, Imp.AtomicUMin Int8),+        (And Int8, Imp.AtomicAnd Int8),+        (Or Int8, Imp.AtomicOr Int8),+        (Xor Int8, Imp.AtomicXor Int8)+      ]+    opencl = opencl8 <> opencl16 <> opencl32 <> opencl64     cuda = opencl  compileProg ::

@@ -744,13 +744,13 @@   AtomicUpdate GPUMem KernelEnv atomicUpdateLocking atomicBinOp lam   | Just ops_and_ts <- lamIsBinOp lam,-    all (\(_, t, _, _) -> primBitSize t `elem` [32, 64]) ops_and_ts =+    all (\(_, t, _, _) -> primBitSize t `elem` [8, 16, 32, 64]) ops_and_ts =       primOrCas ops_and_ts $ \space arrs bucket ->-        -- If the operator is a vectorised binary operator on 32/64-bit-        -- values, we can use a particularly efficient-        -- implementation. If the operator has an atomic implementation-        -- we use that, otherwise it is still a binary operator which-        -- can be implemented by atomic compare-and-swap if 32/64 bits.+        -- If the operator is a vectorised binary operator on single values, we+        -- can use a particularly efficient implementation. If the operator has+        -- an atomic implementation we use that, otherwise it is still a binary+        -- operator which can be implemented by atomic compare-and-swap if 32/64+        -- bits.         forM_ (zip arrs ops_and_ts) $ \(a, (op, t, x, y)) -> do           -- Common variables.           old <- dPrimS "old" t@@ -771,15 +771,19 @@       | all isPrim ops = AtomicPrim       | otherwise = AtomicCAS -    isPrim (op, _, _, _) = isJust $ atomicBinOp op+    -- Only operators of at least 32-bit integers are actually truly atomic with+    -- our current GPU backends - the rest are emulated with CAS-loops in their+    -- implementation.+    isPrim (op, _, _, _) =+      isJust (atomicBinOp op)+        && primByteSize (binOpType op) >= (4 :: Int) --- If the operator functions purely on single 32/64-bit values, we can--- use an implementation based on CAS, no matter what the operator--- does.+-- If the operator functions purely on single single values, we can use an+-- implementation based on CAS, no matter what the operator does. atomicUpdateLocking _ op   | [Prim t] <- lambdaReturnType op,     [xp, _] <- lambdaParams op,-    primBitSize t `elem` [32, 64] = AtomicCAS $ \space [arr] bucket -> do+    primBitSize t `elem` [8, 16, 32, 64] = AtomicCAS $ \space [arr] bucket -> do       old <- dPrimS "old" t       atomicUpdateCAS space t arr old bucket (paramName xp) $         compileBody' [xp] (lambdaBody op)

@@ -669,7 +669,7 @@       ind' <- GC.compileExp $ untyped $ unCount ind       val' <- GC.compileExp val       cast <- atomicCast s ty-      GC.stm [C.cstm|$id:old = $id:op'(&(($ty:cast *)$id:arr)[$exp:ind'], ($ty:ty) $exp:val');|]+      GC.stm [C.cstm|$id:old = $id:op'(&(($ty:cast *)$id:arr)[$exp:ind'], $exp:val');|]       where         op' = op ++ "_" ++ prettyString t ++ "_" ++ atomicSpace s @@ -688,7 +688,7 @@       GC.stm [C.cstm|$id:old = $id:op(&(($ty:cast *)$id:arr)[$exp:ind'], $exp:val');|]       where         op = "atomic_chg_" ++ prettyString t ++ "_" ++ atomicSpace s-    -- First the 64-bit operations.+    -- 64-bit operations     atomicOps s (AtomicAdd Int64 old arr ind val) =       doAtomic s Int64 old arr ind val "atomic_add" [C.cty|typename int64_t|]     atomicOps s (AtomicFAdd Float64 old arr ind val) =@@ -698,9 +698,9 @@     atomicOps s (AtomicSMin Int64 old arr ind val) =       doAtomic s Int64 old arr ind val "atomic_smin" [C.cty|typename int64_t|]     atomicOps s (AtomicUMax Int64 old arr ind val) =-      doAtomic s Int64 old arr ind val "atomic_umax" [C.cty|unsigned int64_t|]+      doAtomic s Int64 old arr ind val "atomic_umax" [C.cty|typename uint64_t|]     atomicOps s (AtomicUMin Int64 old arr ind val) =-      doAtomic s Int64 old arr ind val "atomic_umin" [C.cty|unsigned int64_t|]+      doAtomic s Int64 old arr ind val "atomic_umin" [C.cty|typename uint64_t|]     atomicOps s (AtomicAnd Int64 old arr ind val) =       doAtomic s Int64 old arr ind val "atomic_and" [C.cty|typename int64_t|]     atomicOps s (AtomicOr Int64 old arr ind val) =@@ -711,29 +711,74 @@       doAtomicCmpXchg s (IntType Int64) old arr ind cmp val [C.cty|typename int64_t|]     atomicOps s (AtomicXchg (IntType Int64) old arr ind val) =       doAtomicXchg s (IntType Int64) old arr ind val [C.cty|typename int64_t|]-    ---    atomicOps s (AtomicAdd t old arr ind val) =-      doAtomic s t old arr ind val "atomic_add" [C.cty|int|]-    atomicOps s (AtomicFAdd t old arr ind val) =-      doAtomic s t old arr ind val "atomic_fadd" [C.cty|float|]-    atomicOps s (AtomicSMax t old arr ind val) =-      doAtomic s t old arr ind val "atomic_smax" [C.cty|int|]-    atomicOps s (AtomicSMin t old arr ind val) =-      doAtomic s t old arr ind val "atomic_smin" [C.cty|int|]-    atomicOps s (AtomicUMax t old arr ind val) =-      doAtomic s t old arr ind val "atomic_umax" [C.cty|unsigned int|]-    atomicOps s (AtomicUMin t old arr ind val) =-      doAtomic s t old arr ind val "atomic_umin" [C.cty|unsigned int|]-    atomicOps s (AtomicAnd t old arr ind val) =-      doAtomic s t old arr ind val "atomic_and" [C.cty|int|]-    atomicOps s (AtomicOr t old arr ind val) =-      doAtomic s t old arr ind val "atomic_or" [C.cty|int|]-    atomicOps s (AtomicXor t old arr ind val) =-      doAtomic s t old arr ind val "atomic_xor" [C.cty|int|]-    atomicOps s (AtomicCmpXchg t old arr ind cmp val) =-      doAtomicCmpXchg s t old arr ind cmp val [C.cty|int|]-    atomicOps s (AtomicXchg t old arr ind val) =-      doAtomicXchg s t old arr ind val [C.cty|int|]+    -- 32 bit operations+    atomicOps s (AtomicAdd Int32 old arr ind val) =+      doAtomic s Int32 old arr ind val "atomic_add" [C.cty|int|]+    atomicOps s (AtomicFAdd Float32 old arr ind val) =+      doAtomic s Float32 old arr ind val "atomic_fadd" [C.cty|float|]+    atomicOps s (AtomicSMax Int32 old arr ind val) =+      doAtomic s Int32 old arr ind val "atomic_smax" [C.cty|int|]+    atomicOps s (AtomicSMin Int32 old arr ind val) =+      doAtomic s Int32 old arr ind val "atomic_smin" [C.cty|int|]+    atomicOps s (AtomicUMax Int32 old arr ind val) =+      doAtomic s Int32 old arr ind val "atomic_umax" [C.cty|unsigned int|]+    atomicOps s (AtomicUMin Int32 old arr ind val) =+      doAtomic s Int32 old arr ind val "atomic_umin" [C.cty|unsigned int|]+    atomicOps s (AtomicAnd Int32 old arr ind val) =+      doAtomic s Int32 old arr ind val "atomic_and" [C.cty|int|]+    atomicOps s (AtomicOr Int32 old arr ind val) =+      doAtomic s Int32 old arr ind val "atomic_or" [C.cty|int|]+    atomicOps s (AtomicXor Int32 old arr ind val) =+      doAtomic s Int32 old arr ind val "atomic_xor" [C.cty|int|]+    atomicOps s (AtomicCmpXchg (IntType Int32) old arr ind cmp val) =+      doAtomicCmpXchg s Int32 old arr ind cmp val [C.cty|int|]+    atomicOps s (AtomicXchg (IntType Int32) old arr ind val) =+      doAtomicXchg s Int32 old arr ind val [C.cty|int|]+    -- 16 bit operations+    atomicOps s (AtomicAdd Int16 old arr ind val) =+      doAtomic s Int16 old arr ind val "atomic_add" [C.cty|typename int16_t|]+    atomicOps s (AtomicFAdd Float16 old arr ind val) =+      doAtomic s Float16 old arr ind val "atomic_fadd" [C.cty|typename uint16_t|]+    atomicOps s (AtomicSMax Int16 old arr ind val) =+      doAtomic s Int16 old arr ind val "atomic_smax" [C.cty|typename int16_t|]+    atomicOps s (AtomicSMin Int16 old arr ind val) =+      doAtomic s Int16 old arr ind val "atomic_smin" [C.cty|typename int16_t|]+    atomicOps s (AtomicUMax Int16 old arr ind val) =+      doAtomic s Int16 old arr ind val "atomic_umax" [C.cty|typename uint16_t|]+    atomicOps s (AtomicUMin Int16 old arr ind val) =+      doAtomic s Int16 old arr ind val "atomic_umin" [C.cty|typename uint16_t|]+    atomicOps s (AtomicAnd Int16 old arr ind val) =+      doAtomic s Int16 old arr ind val "atomic_and" [C.cty|typename int16_t|]+    atomicOps s (AtomicOr Int16 old arr ind val) =+      doAtomic s Int16 old arr ind val "atomic_or" [C.cty|typename int16_t|]+    atomicOps s (AtomicXor Int16 old arr ind val) =+      doAtomic s Int16 old arr ind val "atomic_xor" [C.cty|typename int16_t|]+    atomicOps s (AtomicCmpXchg (IntType Int16) old arr ind cmp val) =+      doAtomicCmpXchg s Int16 old arr ind cmp val [C.cty|typename int16_t|]+    atomicOps s (AtomicXchg (IntType Int16) old arr ind val) =+      doAtomicXchg s Int16 old arr ind val [C.cty|typename int16_t|]+    -- 8 bit operations+    atomicOps s (AtomicAdd Int8 old arr ind val) =+      doAtomic s Int8 old arr ind val "atomic_add" [C.cty|typename int8_t|]+    atomicOps s (AtomicSMax Int8 old arr ind val) =+      doAtomic s Int8 old arr ind val "atomic_smax" [C.cty|typename int8_t|]+    atomicOps s (AtomicSMin Int8 old arr ind val) =+      doAtomic s Int8 old arr ind val "atomic_smin" [C.cty|typename int8_t|]+    atomicOps s (AtomicUMax Int8 old arr ind val) =+      doAtomic s Int8 old arr ind val "atomic_umax" [C.cty|typename uint8_t|]+    atomicOps s (AtomicUMin Int8 old arr ind val) =+      doAtomic s Int8 old arr ind val "atomic_umin" [C.cty|typename uint8_t|]+    atomicOps s (AtomicAnd Int8 old arr ind val) =+      doAtomic s Int8 old arr ind val "atomic_and" [C.cty|typename int8_t|]+    atomicOps s (AtomicOr Int8 old arr ind val) =+      doAtomic s Int8 old arr ind val "atomic_or" [C.cty|typename int8_t|]+    atomicOps s (AtomicXor Int8 old arr ind val) =+      doAtomic s Int8 old arr ind val "atomic_xor" [C.cty|typename int8_t|]+    atomicOps s (AtomicCmpXchg (IntType Int8) old arr ind cmp val) =+      doAtomicCmpXchg s Int8 old arr ind cmp val [C.cty|typename int8_t|]+    atomicOps s (AtomicXchg (IntType Int8) old arr ind val) =+      doAtomicXchg s Int8 old arr ind val [C.cty|typename int8_t|]+    -- General     atomicOps s (AtomicWrite t arr ind val) = do       ind' <- GC.compileExp $ untyped $ unCount ind       val' <- toStorage t <$> GC.compileExp val@@ -745,6 +790,8 @@         case s of           Space "shared" -> [C.cstm|mem_fence_local();|]           _ -> [C.cstm|mem_fence_global();|]+    atomicOps _ op =+      error $ "atomicOp: unsupported " <> show op      cannotAllocate :: GC.Allocate KernelOp KernelState     cannotAllocate _ =

@@ -40,7 +40,15 @@  -- | @rts/c/atomics.h@ atomicsH :: T.Text-atomicsH = $(embedStringFile "rts/c/atomics.h")+atomicsH =+  -- The order matters, as e.g. atomics16.h is implemented in terms of 32-bit+  -- atomics.+  mconcat+    [ $(embedStringFile "rts/c/atomics64.h"),+      $(embedStringFile "rts/c/atomics32.h"),+      $(embedStringFile "rts/c/atomics16.h"),+      $(embedStringFile "rts/c/atomics8.h")+    ] {-# NOINLINE atomicsH #-}  -- | @rts/c/uniform.h@

@@ -372,7 +372,9 @@   pure $     H.span ! A.id (fromString ("synopsis:" <> vname_id)) $       H.a ! A.href (fromString ("#" ++ vname_id)) $-        renderName (baseName vname)+        if symbolName (baseName vname)+          then parens (renderName (baseName vname))+          else renderName (baseName vname)  synopsisValBind :: ValBind -> Maybe (DocM Html) synopsisValBind vb = Just $ do

@@ -1,3 +1,4 @@+-- | The actual implementation of @futhark fmt@. module Futhark.Fmt.Printer   ( fmtToText,     fmtToDoc,@@ -510,8 +511,12 @@     addComments loc $ fmt doc <> "val" <+> sub <+> ":" </> stdIndent (fmt te)     where       sub-        | null ps = fmtName bindingStyle name-        | otherwise = fmtName bindingStyle name <+> align (sep space $ map fmt ps)+        | null ps = name'+        | otherwise = name' <+> align (sep space $ map fmt ps)+      name' =+        if symbolName name+          then parens $ fmtName bindingStyle name+          else fmtName bindingStyle name   fmt (ModSpec name mte doc loc) =     addComments loc $ fmt doc <> "module" <+> fmtName bindingStyle name <> ":" <+> fmt mte   fmt (IncludeSpec mte loc) = addComments loc $ "include" <+> fmt mte

@@ -1046,12 +1046,12 @@   Just . pure <$> subExpReturns se expReturns (BasicOp (Opaque _ (Var v))) =   Just . pure <$> varReturns v-expReturns (BasicOp (Reshape k newshape v)) = do+expReturns (BasicOp (Reshape v newshape)) = do   (et, _, mem, lmad) <- arrayVarReturns v-  case reshaper k lmad $ map pe64 $ shapeDims newshape of+  case reshaper (reshapeKind newshape) lmad $ map pe64 $ shapeDims $ newShape newshape of     Just lmad' ->       pure . Just $-        [ MemArray et (fmap Free newshape) NoUniqueness . Just $+        [ MemArray et (Free <$> newShape newshape) NoUniqueness . Just $             ReturnsInBlock mem (existentialiseLMAD [] lmad')         ]     Nothing -> pure Nothing@@ -1060,7 +1060,7 @@       LMAD.reshape lmad     reshaper ReshapeCoerce lmad =       Just . LMAD.coerce lmad-expReturns (BasicOp (Rearrange perm v)) = do+expReturns (BasicOp (Rearrange v perm)) = do   (et, Shape dims, mem, lmad) <- arrayVarReturns v   let lmad' = LMAD.permute lmad perm       dims' = rearrangeShape perm dims

@@ -101,16 +101,6 @@       Engine.isAllocation = isAlloc mempty mempty     } --- | Standard collection of simplification rules for representations--- with memory.-memRuleBook :: (SimplifyMemory rep inner) => RuleBook (Wise rep)-memRuleBook =-  standardRules-    <> ruleBook-      [ RuleOp decertifySafeAlloc-      ]-      []- -- If an allocation is statically known to be safe, then we can remove -- the certificates on it.  This can help hoist things that would -- otherwise be stuck inside loops or branches.@@ -121,3 +111,36 @@     safeOp op =       Simplify $ attributing attrs $ letBind pat $ Op op decertifySafeAlloc _ _ _ _ = Skip++--+-- copy(reshape(manifest(v0),s)) can be rewritten to just reshape(manifest(v0),s).+--+-- This is a pattern that can be produced by ExplicitAllocations when the+-- reshape would otherwise produce a layout that is not representable as an+-- LMAD. We have to be careful that the manifest writes to the same memory that+-- the original copy put it in.+copyManifest :: (SimplifyMemory rep inner) => TopDownRuleBasicOp (Wise rep)+copyManifest vtable pat aux (Replicate (Shape []) (Var v2))+  | Just (Reshape v1 s, v2_cs) <- ST.lookupBasicOp v2 vtable,+    Just (Manifest v0 perm, v1_cs) <- ST.lookupBasicOp v1 vtable,+    Pat [PatElem _ (_, MemArray _ _ _ (ArrayIn mem _))] <- pat =+      Simplify $ do+        ~(MemArray pt shape u (ArrayIn _ v1_lmad)) <- lookupMemInfo v1+        v0' <- newVName (baseString v1 <> "_manifest")+        let manifest_pat =+              Pat [PatElem v0' $ MemArray pt shape u $ ArrayIn mem v1_lmad]+            stm = mkWiseStm manifest_pat mempty $ BasicOp $ Manifest v0 perm+        certifying (v1_cs <> v2_cs) $ addStm stm+        auxing aux $ letBind pat $ BasicOp $ Reshape v0' s+copyManifest _ _ _ _ = Skip++-- | Standard collection of simplification rules for representations+-- with memory.+memRuleBook :: (SimplifyMemory rep inner) => RuleBook (Wise rep)+memRuleBook =+  standardRules+    <> ruleBook+      [ RuleOp decertifySafeAlloc,+        RuleBasicOp copyManifest+      ]+      []

@@ -260,6 +260,12 @@               <*> pure t           ) +pDimSplice :: Parser (DimSplice SubExp)+pDimSplice = DimSplice <$> pInt <* lexeme "::" <*> pInt <* lexeme "=>" <*> pShape++pNewShape :: Parser (NewShape SubExp)+pNewShape = parens $ NewShape <$> (pDimSplice `sepBy` pComma) <* pSemi <*> pShape+ pBasicOp :: Parser BasicOp pBasicOp =   choice@@ -280,17 +286,15 @@       keyword "replicate"         *> parens (Replicate <$> pShape <* pComma <*> pSubExp),       keyword "reshape"-        *> parens (Reshape ReshapeArbitrary <$> pShape <* pComma <*> pVName),-      keyword "coerce"-        *> parens (Reshape ReshapeCoerce <$> pShape <* pComma <*> pVName),+        *> parens (Reshape <$> pVName <* pComma <*> pNewShape),       keyword "scratch"         *> parens (Scratch <$> pPrimType <*> many (pComma *> pSubExp)),       keyword "rearrange"         *> parens-          (Rearrange <$> parens (pInt `sepBy` pComma) <* pComma <*> pVName),+          (Rearrange <$> pVName <* pComma <*> parens (pInt `sepBy` pComma)),       keyword "manifest"         *> parens-          (Manifest <$> parens (pInt `sepBy` pComma) <* pComma <*> pVName),+          (Manifest <$> pVName <* pComma <*> parens (pInt `sepBy` pComma)),       keyword "concat" *> do         d <- "@" *> L.decimal         parens $ do

@@ -1,9 +1,8 @@ {-# OPTIONS_GHC -fno-warn-orphans #-} --- | Futhark prettyprinter.  This module defines 'Pretty' instances--- for the AST defined in "Futhark.IR.Syntax",--- but also a number of convenience functions if you don't want to use--- the interface from 'Pretty'.+-- | Futhark prettyprinter. This module defines 'Pretty' instances for the AST+-- defined in "Futhark.IR.Syntax", but also a number of convenience functions if+-- you don't want to use the interface from 'Pretty'. module Futhark.IR.Pretty   ( prettyTuple,     prettyTupleLines,@@ -43,7 +42,7 @@   pretty Commutative = "commutative"   pretty Noncommutative = "noncommutative" -instance Pretty Shape where+instance (Pretty d) => Pretty (ShapeBase d) where   pretty = mconcat . map (brackets . pretty) . shapeDims  instance Pretty Rank where@@ -53,9 +52,6 @@   pretty (Free e) = pretty e   pretty (Ext x) = "?" <> pretty (show x) -instance Pretty ExtShape where-  pretty = mconcat . map (brackets . pretty) . shapeDims- instance Pretty Space where   pretty DefaultSpace = mempty   pretty (Space s) = "@" <> pretty s@@ -152,6 +148,10 @@ instance Pretty Attrs where   pretty = hsep . attrAnnots +instance (Pretty dec) => Pretty (StmAux dec) where+  pretty (StmAux cs attrs dec) =+    braces $ mconcat $ punctuate semi [pretty cs, pretty attrs, pretty dec]+ instance (Pretty t) => Pretty (Pat t) where   pretty (Pat xs) = braces $ commastack $ map pretty xs @@ -187,6 +187,13 @@ instance (Pretty a) => Pretty (FlatSlice a) where   pretty (FlatSlice offset xs) = brackets (pretty offset <> ";" <+> commasep (map pretty xs)) +instance (Pretty d) => Pretty (DimSplice d) where+  pretty (DimSplice i k shape) = pretty i <> "::" <> pretty k <> "=>" <> pretty shape++instance (Pretty d) => Pretty (NewShape d) where+  pretty (NewShape ds shape) =+    parens $ align $ commastack (map pretty ds) <> semi </> pretty shape+ instance Pretty BasicOp where   pretty (SubExp se) = pretty se   pretty (Opaque OpaqueNil e) = "opaque" <> apply [pretty e]@@ -228,15 +235,14 @@     "replicate" <> apply [pretty ne, align (pretty ve)]   pretty (Scratch t shape) =     "scratch" <> apply (pretty t : map pretty shape)-  pretty (Reshape ReshapeArbitrary shape e) =-    "reshape" <> apply [pretty shape, pretty e]-  pretty (Reshape ReshapeCoerce shape e) =-    "coerce" <> apply [pretty shape, pretty e]-  pretty (Rearrange perm e) =-    "rearrange" <> apply [apply (map pretty perm), pretty e]+  pretty (Reshape reshape e) =+    "reshape" <> parens (align $ commastack [pretty reshape, pretty e])+  pretty (Rearrange v perm) =+    "rearrange" <> apply [pretty v, apply (map pretty perm)]   pretty (Concat i (x :| xs) w) =     "concat" <> "@" <> pretty i <> apply (pretty w : pretty x : map pretty xs)-  pretty (Manifest perm e) = "manifest" <> apply [apply (map pretty perm), pretty e]+  pretty (Manifest v perm) =+    "manifest" <> apply [pretty v, apply (map pretty perm)]   pretty (Assert e msg (loc, _)) =     "assert" <> apply [pretty e, pretty msg, pretty $ show $ locStr loc]   pretty (UpdateAcc safety acc is v) =

@@ -71,8 +71,8 @@ basicOpAliases Iota {} = [mempty] basicOpAliases Replicate {} = [mempty] basicOpAliases Scratch {} = [mempty]-basicOpAliases (Reshape _ _ e) = [vnameAliases e]-basicOpAliases (Rearrange _ e) = [vnameAliases e]+basicOpAliases (Reshape v _) = [vnameAliases v]+basicOpAliases (Rearrange v _) = [vnameAliases v] basicOpAliases Concat {} = [mempty] basicOpAliases Manifest {} = [mempty] basicOpAliases Assert {} = [mempty]

@@ -345,6 +345,9 @@ instance (FreeIn d) => FreeIn (ShapeBase d) where   freeIn' = freeIn' . shapeDims +instance (FreeIn d) => FreeIn (NewShape d) where+  freeIn' = foldMap freeIn'+ instance (FreeIn d) => FreeIn (Ext d) where   freeIn' (Free x) = freeIn' x   freeIn' (Ext _) = mempty

@@ -3,30 +3,56 @@ module Futhark.IR.Prop.Reshape   ( -- * Construction     shapeCoerce,+    reshapeAll,+    reshapeCoerce,      -- * Execution     reshapeOuter,     reshapeInner,+    newshapeInner,+    applySplice,      -- * Simplification+    flipReshapeRearrange,+    flipRearrangeReshape,+    simplifyNewShape,      -- * Shape calculations     reshapeIndex,     flattenIndex,     unflattenIndex,     sliceSizes,++    -- * Analysis+    ReshapeKind (..),+    reshapeKind,+    newShape,   ) where +import Control.Monad (guard, mplus) import Data.Foldable+import Data.Maybe+import Futhark.IR.Prop.Rearrange (isMapTranspose, rearrangeInverse, rearrangeShape) import Futhark.IR.Syntax+import Futhark.Util (focusNth, mapAccumLM, takeLast) import Futhark.Util.IntegralExp import Prelude hiding (product, quot, sum) +-- | Construct a 'NewShape' that completely reshapes the initial shape.+reshapeAll :: (ArrayShape old) => old -> ShapeBase new -> NewShape new+reshapeAll old new = NewShape [DimSplice 0 (shapeRank old) new] new++-- | Construct a 'NewShape' that coerces the shape.+reshapeCoerce :: ShapeBase new -> NewShape new+reshapeCoerce shape = NewShape (zipWith dim (shapeDims shape) [0 ..]) shape+  where+    dim x i = DimSplice i 1 $ Shape [x]+ -- | Construct a 'Reshape' that is a 'ReshapeCoerce'. shapeCoerce :: [SubExp] -> VName -> Exp rep shapeCoerce newdims arr =-  BasicOp $ Reshape ReshapeCoerce (Shape newdims) arr+  BasicOp $ Reshape arr (reshapeCoerce (Shape newdims))  -- | @reshapeOuter newshape n oldshape@ returns a 'Reshape' expression -- that replaces the outer @n@ dimensions of @oldshape@ with @newshape@.@@ -34,13 +60,24 @@ reshapeOuter newshape n oldshape =   newshape <> Shape (drop n (shapeDims oldshape)) --- | @reshapeInner newshape n oldshape@ returns a 'Reshape' expression--- that replaces the inner @m-n@ dimensions (where @m@ is the rank of--- @oldshape@) of @src@ with @newshape@.+-- | @reshapeInner newshape n oldshape@ produces a shape that replaces the inner+-- @m-n@ dimensions (where @m@ is the rank of @oldshape@) of @src@ with+-- @newshape@. reshapeInner :: Shape -> Int -> Shape -> Shape reshapeInner newshape n oldshape =   Shape (take n (shapeDims oldshape)) <> newshape +-- | @newshapeInner outershape newshape@ bumps all the dimensions in @newshape@+-- by the rank of @outershape@, essentially making them operate on the inner+-- dimensions of a larger array, and also updates the shape of @newshape@ to+-- have @outershape@ outermost.+newshapeInner :: Shape -> NewShape SubExp -> NewShape SubExp+newshapeInner outershape (NewShape ss oldshape) =+  NewShape (map f ss) (outershape <> oldshape)+  where+    r = shapeRank outershape+    f (DimSplice i k shape) = DimSplice (r + i) k shape+ -- | @reshapeIndex to_dims from_dims is@ transforms the index list -- @is@ (which is into an array of shape @from_dims@) into an index -- list @is'@, which is into an array of shape @to_dims@.  @is@ must@@ -101,3 +138,228 @@   product (n : ns) : sliceSizes ns  {- HLINT ignore sliceSizes -}++-- | Interchange a reshape and rearrange. Essentially, rewrite composition+--+-- @+-- let v1 = reshape(v0, v1_shape)+-- let v2 = rearrange(v1, perm)+-- @+--+-- into+--+-- @+-- let v1' = rearrange(v0, perm')+-- let v2' = reshape(v1', v1_shape')+-- @+--+-- The function is given the shape of @v0@, @v1@, and the @perm@, and returns+-- @perm'@. This is a meaningful operation when @v2@ is itself reshaped, as the+-- reshape-reshape can be fused. This can significantly simplify long chains of+-- reshapes and rearranges.+flipReshapeRearrange ::+  (Eq d) =>+  [d] ->+  [d] ->+  [Int] ->+  Maybe [Int]+flipReshapeRearrange v0_shape v1_shape perm = do+  (num_map_dims, num_a_dims, num_b_dims) <- isMapTranspose perm+  guard $ num_a_dims == 1+  guard $ num_b_dims == 1+  let map_dims = take num_map_dims v0_shape+      num_b_dims_expanded = length v0_shape - num_map_dims - num_a_dims+      num_a_dims_expanded = length v0_shape - num_map_dims - num_b_dims+      caseA = do+        guard $ take num_a_dims v0_shape == take num_b_dims v1_shape+        let perm' =+              [0 .. num_map_dims - 1]+                ++ map (+ num_map_dims) ([1 .. num_b_dims_expanded] ++ [0])+        Just perm'+      caseB = do+        guard $ takeLast num_b_dims v0_shape == takeLast num_b_dims v1_shape+        let perm' =+              [0 .. num_map_dims - 1]+                ++ map+                  (+ num_map_dims)+                  (num_a_dims_expanded : [0 .. num_a_dims_expanded - 1])+        Just perm'++  guard $ map_dims == take num_map_dims v1_shape++  caseA `mplus` caseB++-- | Interchange a reshape and rearrange. Essentially, rewrite composition+--+-- @+-- let v1 = rearrange(v0, perm)+-- let v2 = reshape(v1, v1_shape)+-- @+--+-- into+--+-- @+-- let v1' = reshape(v0, v1_shape')+-- let v2' = rearrange(v1', perm')+-- @+--+-- The function is given @perm@ and @v1_shape@, and returns @perm'@ and+-- @v1_shape'@. This is a meaningful operation when @v2@ is itself rearranged+-- (or @v0@ the result of a reshape), as this enables fusion.+flipRearrangeReshape :: [Int] -> NewShape d -> Maybe (NewShape d, [Int])+flipRearrangeReshape orig_perm (NewShape ss shape) = do+  (perm', ss') <- mapAccumLM f orig_perm ss+  let shape' = Shape $ rearrangeShape (rearrangeInverse perm') (shapeDims shape)+  Just (NewShape ss' shape', perm')+  where+    f perm (DimSplice i 1 s) = do+      (perm_bef, j, perm_aft) <- focusNth i perm+      let adj l = if l > j then l + length s - 1 else l+      Just+        ( map adj perm_bef ++ [j .. j + length s - 1] ++ map adj perm_aft,+          DimSplice j 1 s+        )+    f _ _ = Nothing++-- | Which kind of reshape is this?+data ReshapeKind+  = -- | New shape is dynamically same as original.+    ReshapeCoerce+  | -- | Any kind of reshaping.+    ReshapeArbitrary+  deriving (Eq, Ord, Show)++-- | Determine whether this might be a coercion.+reshapeKind :: NewShape SubExp -> ReshapeKind+reshapeKind (NewShape ss _)+  | all unit ss = ReshapeCoerce+  | otherwise = ReshapeArbitrary+  where+    unit (DimSplice _ 1 (Shape [_])) = True+    unit _ = False++-- | Apply the splice to a shape.+applySplice :: ShapeBase d -> DimSplice d -> ShapeBase d+applySplice shape_bef (DimSplice i k shape) =+  takeDims i shape_bef <> shape <> stripDims (i + k) shape_bef++-- | @dimSpan i n s@ gets @n@ dimensions starting from @i@ from @s@.+dimSpan :: Int -> Int -> ShapeBase d -> ShapeBase d+dimSpan i n = takeDims n . dropDims i++next ::+  (Eq d) =>+  ShapeBase d ->+  DimSplice d ->+  DimSplice d ->+  [DimSplice d] ->+  Maybe [DimSplice d]+next shape x y ss =+  (x :) <$> move (applySplice shape x, y) ss++move ::+  (Eq d) =>+  (ShapeBase d, DimSplice d) ->+  [DimSplice d] ->+  Maybe [DimSplice d]+--+-- A coercion that does not do anything.+move (shape_bef, DimSplice i1 n1 shape) ss+  | dimSpan i1 n1 shape_bef == shape =+      Just ss+--+-- See if we can find some redundancy.+move (shape, DimSplice i1 n1 s1) ss+  -- Check for redundant prefix.+  | match <-+      takeWhile (uncurry (==)) $+        zip (shapeDims (dimSpan i1 n1 shape)) (shapeDims s1),+    not $ null match,+    length match /= n1 =+      let k = length match+       in Just $ DimSplice (i1 + k) (n1 - k) (dropDims k s1) : ss+  -- Check for redundant suffix.+  | match <-+      takeWhile (uncurry (==)) $+        zip+          (reverse (shapeDims (dimSpan i1 n1 shape)))+          (reverse (shapeDims s1)),+    not $ null match,+    length match /= n1 =+      let k = length match+       in Just $ DimSplice i1 (n1 - k) (takeDims (length s1 - k) s1) : ss+--+-- Base case.+move _ [] = Nothing+--+-- A coercion can be fused with anything.+move (_, DimSplice i1 1 (Shape [_])) (DimSplice i2 n2 s2 : ss)+  | i1 == i2 =+      Just $ DimSplice i2 n2 s2 : ss+--+-- A flatten with an inverse unflatten turns into nothing.+move (shape_bef, DimSplice i1 n1 _s1) (DimSplice i2 _n2 s2 : ss)+  | i1 == i2,+    dimSpan i1 n1 shape_bef == s2 =+      Just ss+--+-- An unflatten where one of the dimensions is then further unflattened.+move (_, DimSplice i1 n1 s1) (DimSplice i2 n2 s2 : ss)+  | i2 >= i1,+    i2 < i1 + length s1,+    n1 == 1,+    n2 == 1 =+      Just $ DimSplice i1 1 (s1_bef <> s2 <> s1_aft) : ss+  where+    s1_bef = takeDims (i2 - i1) s1+    s1_aft = dropDims (i2 - i1 + 1) s1++--+-- Flatten followed by a flattening of overlapping dimensions.+move (_, DimSplice i1 n1 s1) (DimSplice i2 n2 s2 : ss)+  | length s1 == 1,+    length s2 == 1,+    i1 == i2 + 1,+    n2 > 1 =+      Just $ DimSplice i2 (n2 + n1 - length s1) s2 : ss+--+-- Flatten into an unflatten.+move (_, DimSplice i1 n1 (Shape [_])) (DimSplice i2 1 s2 : ss)+  | i1 == i2 =+      Just $ DimSplice i1 n1 s2 : ss+--+-- These cases are for updating dimensions as we move across intervening+-- operations.+move (shape, DimSplice i1 n1 s1) (DimSplice i2 n2 s2 : ss)+  | i1 > i2 + n2 =+      next shape (DimSplice i2 n2 s2) (DimSplice (i1 - n2 + length s2) n1 s1) ss+  | i2 > i1 + n1 =+      next shape (DimSplice (i2 - n1 + length s1) n2 s2) (DimSplice i1 n2 s1) ss+  | otherwise = Nothing++-- This is a quadratic-time function that looks for a DimSplice that can be+-- combined with a move DimSlice (and then does so). Since these lists are+-- usually small, this should not be a problem. It is called to convergence by+-- 'improve'.+improveOne :: (Eq d) => ShapeBase d -> [DimSplice d] -> Maybe [DimSplice d]+improveOne _ [] = Nothing+improveOne shape (s : ss) =+  move (shape, s) ss `mplus` ((s :) <$> improveOne (applySplice shape s) ss)++-- | Try to simplify the given 'NewShape'. Returns 'Nothing' if no improvement+-- is possible.+simplifyNewShape :: (Eq d) => ShapeBase d -> NewShape d -> Maybe (NewShape d)+simplifyNewShape shape_bef (NewShape ss shape) =+  NewShape <$> (improve <$> improveOne shape_bef ss) <*> pure shape+  where+    improve ss' = maybe ss' improve $ improveOne shape_bef ss'++{-# NOINLINE flipReshapeRearrange #-}++{-# NOINLINE flipRearrangeReshape #-}++{-# NOINLINE reshapeKind #-}++{-# NOINLINE simplifyNewShape #-}++{-# NOINLINE newshapeInner #-}

@@ -101,23 +101,19 @@   pure . flip arrayOfShape shape <$> subExpType e basicOpType (Scratch t shape) =   pure [arrayOf (Prim t) (Shape shape) NoUniqueness]-basicOpType (Reshape _ (Shape []) e) =-  result <$> lookupType e-  where-    result t = [Prim $ elemType t]-basicOpType (Reshape _ shape e) =+basicOpType (Reshape e shape) =   result <$> lookupType e   where-    result t = [t `setArrayShape` shape]-basicOpType (Rearrange perm e) =-  result <$> lookupType e+    result t = [t `setArrayShape` newShape shape]+basicOpType (Rearrange v perm) =+  result <$> lookupType v   where     result t = [rearrangeType perm t] basicOpType (Concat i (x :| _) ressize) =   result <$> lookupType x   where     result xt = [setDimSize i xt ressize]-basicOpType (Manifest _ v) =+basicOpType (Manifest v _) =   pure <$> lookupType v basicOpType Assert {} =   pure [Prim Unit]

@@ -35,7 +35,7 @@ import Futhark.Analysis.SymbolTable qualified as ST import Futhark.Analysis.UsageTable qualified as UT import Futhark.IR.Prop.Aliases-import Futhark.IR.SOACS+import Futhark.IR.SOACS hiding (reshapeInner) import Futhark.MonadFreshNames import Futhark.Optimise.Simplify qualified as Simplify import Futhark.Optimise.Simplify.Engine qualified as Engine@@ -295,12 +295,10 @@                   { lambdaBody = (lambdaBody fun) {bodyResult = subExpsRes ses'},                     lambdaReturnType = rettype'                   }-          mapM_ (uncurry letBind) invariant-          auxing aux $-            letBindNames (map patElemName pat') $-              Op $-                soacOp $-                  Screma w arrs (mapSOAC fun')+          auxing aux $ do+            mapM_ (uncurry letBind) invariant+            letBindNames (map patElemName pat') . Op $+              soacOp (Screma w arrs (mapSOAC fun')) liftIdentityMapping _ _ _ _ = Skip  liftIdentityStreaming :: BottomUpRuleOp (Wise SOACS)@@ -459,9 +457,10 @@                       { lambdaBody = (lambdaBody fun) {bodyResult = ses'},                         lambdaReturnType = ts'                       }-              auxing aux $ letBind (Pat pes') $ Op $ Screma w arrs $ mapSOAC fun'-              forM_ copies $ \(from, to) ->-                letBind (Pat [to]) $ BasicOp $ Replicate mempty $ Var $ patElemName from+              auxing aux $ do+                letBind (Pat pes') $ Op $ Screma w arrs $ mapSOAC fun'+                forM_ copies $ \(from, to) ->+                  letBind (Pat [to]) $ BasicOp $ Replicate mempty $ Var $ patElemName from   where     checkForDuplicates (ses_ts_pes', copies) (se, t, pe)       | Just (_, _, pe') <- find (\(x, _, _) -> resSubExp x == resSubExp se) ses_ts_pes' =@@ -471,27 +470,33 @@       | otherwise = (ses_ts_pes' ++ [(se, t, pe)], copies) removeDuplicateMapOutput _ _ _ _ = Skip +reshapeInner :: SubExp -> NewShape SubExp -> NewShape SubExp+reshapeInner w new_shape =+  reshapeCoerce outer <> newshapeInner outer new_shape+  where+    outer = Shape [w]+ -- Mapping some operations becomes an extension of that operation. mapOpToOp :: BottomUpRuleOp (Wise SOACS) mapOpToOp (_, used) pat aux1 e-  | Just (map_pe, cs, w, BasicOp (Reshape k newshape reshape_arr), [p], [arr]) <-+  | Just (map_pe, cs, w, BasicOp (Reshape reshape_arr newshape), [p], [arr]) <-       isMapWithOp pat e,     paramName p == reshape_arr,     not $ UT.isConsumed (patElemName map_pe) used = Simplify $ do       certifying (stmAuxCerts aux1 <> cs) . letBind pat . BasicOp $-        Reshape k (Shape [w] <> newshape) arr+        Reshape arr (reshapeInner w newshape)   | Just (_, cs, _, BasicOp (Concat d (arr :| arrs) dw), ps, outer_arr : outer_arrs) <-       isMapWithOp pat e,     (arr : arrs) == map paramName ps =       Simplify . certifying (stmAuxCerts aux1 <> cs) . letBind pat . BasicOp $         Concat (d + 1) (outer_arr :| outer_arrs) dw   | Just-      (map_pe, cs, _, BasicOp (Rearrange perm rearrange_arr), [p], [arr]) <-+      (map_pe, cs, _, BasicOp (Rearrange rearrange_arr perm), [p], [arr]) <-       isMapWithOp pat e,     paramName p == rearrange_arr,     not $ UT.isConsumed (patElemName map_pe) used =       Simplify . certifying (stmAuxCerts aux1 <> cs) . letBind pat . BasicOp $-        Rearrange (0 : map (1 +) perm) arr+        Rearrange arr (0 : map (1 +) perm) mapOpToOp _ _ _ _ = Skip  isMapWithOp ::@@ -625,9 +630,10 @@         y_ws <- mapM sizeOf ys         guard $ all (x_w ==) y_ws         pure (x_w, x : ys, cs)-      Just (BasicOp (Reshape ReshapeCoerce _ arr), cs) -> do-        (a, b, cs') <- isConcat arr-        pure (a, b, cs <> cs')+      Just (BasicOp (Reshape arr newshape), cs)+        | ReshapeCoerce <- reshapeKind newshape -> do+            (a, b, cs') <- isConcat arr+            pure (a, b, cs <> cs')       _ -> Nothing fuseConcatScatter _ _ _ _ = Skip @@ -718,7 +724,7 @@ data ArrayOp   = ArrayIndexing Certs VName (Slice SubExp)   | ArrayRearrange Certs VName [Int]-  | ArrayReshape Certs VName ReshapeKind Shape+  | ArrayReshape Certs VName (NewShape SubExp)   | ArrayCopy Certs VName   | -- | Never constructed.     ArrayVar Certs VName@@ -727,24 +733,24 @@ arrayOpArr :: ArrayOp -> VName arrayOpArr (ArrayIndexing _ arr _) = arr arrayOpArr (ArrayRearrange _ arr _) = arr-arrayOpArr (ArrayReshape _ arr _ _) = arr+arrayOpArr (ArrayReshape _ arr _) = arr arrayOpArr (ArrayCopy _ arr) = arr arrayOpArr (ArrayVar _ arr) = arr  arrayOpCerts :: ArrayOp -> Certs arrayOpCerts (ArrayIndexing cs _ _) = cs arrayOpCerts (ArrayRearrange cs _ _) = cs-arrayOpCerts (ArrayReshape cs _ _ _) = cs+arrayOpCerts (ArrayReshape cs _ _) = cs arrayOpCerts (ArrayCopy cs _) = cs arrayOpCerts (ArrayVar cs _) = cs  isArrayOp :: Certs -> Exp rep -> Maybe ArrayOp isArrayOp cs (BasicOp (Index arr slice)) =   Just $ ArrayIndexing cs arr slice-isArrayOp cs (BasicOp (Rearrange perm arr)) =+isArrayOp cs (BasicOp (Rearrange arr perm)) =   Just $ ArrayRearrange cs arr perm-isArrayOp cs (BasicOp (Reshape k new_shape arr)) =-  Just $ ArrayReshape cs arr k new_shape+isArrayOp cs (BasicOp (Reshape arr new_shape)) =+  Just $ ArrayReshape cs arr new_shape isArrayOp cs (BasicOp (Replicate (Shape []) (Var arr))) =   Just $ ArrayCopy cs arr isArrayOp _ _ =@@ -752,8 +758,8 @@  fromArrayOp :: ArrayOp -> (Certs, Exp rep) fromArrayOp (ArrayIndexing cs arr slice) = (cs, BasicOp $ Index arr slice)-fromArrayOp (ArrayRearrange cs arr perm) = (cs, BasicOp $ Rearrange perm arr)-fromArrayOp (ArrayReshape cs arr k new_shape) = (cs, BasicOp $ Reshape k new_shape arr)+fromArrayOp (ArrayRearrange cs arr perm) = (cs, BasicOp $ Rearrange arr perm)+fromArrayOp (ArrayReshape cs arr new_shape) = (cs, BasicOp $ Reshape arr new_shape) fromArrayOp (ArrayCopy cs arr) = (cs, BasicOp $ Replicate mempty $ Var arr) fromArrayOp (ArrayVar cs arr) = (cs, BasicOp $ SubExp $ Var arr) @@ -953,7 +959,7 @@       arr `elem` map_param_names         && all (`ST.elem` vtable) (namesToList $ freeIn cs)         && not (null perm)-    arrayIsMapParam (_, ArrayReshape cs arr _ new_shape) =+    arrayIsMapParam (_, ArrayReshape cs arr new_shape) =       arr `elem` map_param_names         && all (`ST.elem` vtable) (namesToList $ freeIn cs <> freeIn new_shape)     arrayIsMapParam (_, ArrayCopy cs arr) =@@ -973,9 +979,9 @@                 ArrayIndexing _ _ (Slice slice) ->                   BasicOp $ Index arr $ Slice $ whole_dim : slice                 ArrayRearrange _ _ perm ->-                  BasicOp $ Rearrange (0 : map (+ 1) perm) arr-                ArrayReshape _ _ k new_shape ->-                  BasicOp $ Reshape k (Shape [w] <> new_shape) arr+                  BasicOp $ Rearrange arr (0 : map (+ 1) perm)+                ArrayReshape _ _ new_shape ->+                  BasicOp $ Reshape arr $ reshapeInner w new_shape                 ArrayCopy {} ->                   BasicOp $ Replicate mempty $ Var arr                 ArrayVar {} ->@@ -1046,12 +1052,12 @@      invariantToMap = all (`ST.elem` vtable) . namesToList . freeIn -    onStm (transformed, map_infos, stms) (Let (Pat [pe]) aux (BasicOp (Reshape k new_shape arr)))+    onStm (transformed, map_infos, stms) (Let (Pat [pe]) aux (BasicOp (Reshape arr new_shape)))       | ([(res, _, screma_pe)], map_pesres') <- partition matches map_infos,         Just t <- typeOf <$> M.lookup arr scope,-        invariantToMap t =+        invariantToMap (t, new_shape) =           let cs = stmAuxCerts aux <> resCerts res-              transform = (arr, cs, BasicOp . Reshape k (Shape [w] <> new_shape))+              transform = (arr, cs, BasicOp . flip Reshape (reshapeInner w new_shape))            in ((t, screma_pe, transform) : transformed, map_pesres', stms)       where         matches (r, _, _) = resSubExp r == Var (patElemName pe)

@@ -99,7 +99,9 @@ module Futhark.IR.Syntax   ( module Language.Futhark.Core,     prettyString,+    prettyStringOneLine,     prettyText,+    prettyTextOneLine,     Pretty,     module Futhark.IR.Rep,     module Futhark.IR.Syntax.Core,@@ -130,7 +132,8 @@     CmpOp (..),     ConvOp (..),     OpaqueOp (..),-    ReshapeKind (..),+    DimSplice (..),+    NewShape (..),     WithAccInput,     Exp (..),     Case (..),@@ -173,7 +176,7 @@ import Data.Traversable (fmapDefault, foldMapDefault) import Futhark.IR.Rep import Futhark.IR.Syntax.Core-import Futhark.Util.Pretty (Pretty, prettyString, prettyText)+import Futhark.Util.Pretty (Pretty, prettyString, prettyStringOneLine, prettyText, prettyTextOneLine) import Language.Futhark.Core import Prelude hiding (id, (.)) @@ -205,6 +208,9 @@   }   deriving (Ord, Show, Eq) +instance (Monoid dec) => Monoid (StmAux dec) where+  mempty = StmAux mempty mempty mempty+ instance (Semigroup dec) => Semigroup (StmAux dec) where   StmAux cs1 attrs1 dec1 <> StmAux cs2 attrs2 dec2 =     StmAux (cs1 <> cs2) (attrs1 <> attrs2) (dec1 <> dec2)@@ -305,14 +311,41 @@     OpaqueTrace T.Text   deriving (Eq, Ord, Show) --- | Which kind of reshape is this?-data ReshapeKind-  = -- | New shape is dynamically same as original.-    ReshapeCoerce-  | -- | Any kind of reshaping.-    ReshapeArbitrary-  deriving (Eq, Ord, Show)+-- | Split or join a range of dimensions. A reshaping operation consists of a+-- sequence of these. The purpose is to maintain information about the original+-- operations (flatten/unflatten), which can then be used for algebraic+-- optimisations.+data DimSplice d+  = -- | @DimSplice i k s@ modifies dimensions @i@ to @i+k-1@ to instead have+    -- shape @s@.+    --+    -- If @k@ is 1 and the rank of @s@ is greater than 1, then this is+    -- equivalent to unflattening a dimension.+    --+    -- If @k@ is greater than 1 and the rank of @s@ is 1, then this is+    -- equivalent to flattening adjacent dimensions.+    --+    -- If @k@ is 1 and the rank of @s@ is 1, then it is a coercion - a change+    -- that only affects the type, but does not have any semantic effect.+    --+    -- Other cases can do arbitrary changes, but are harder for the compiler to+    -- analyse.+    DimSplice Int Int (ShapeBase d)+  deriving (Eq, Ord, Show, Functor, Foldable, Traversable) +-- | A reshaping operation consists of a sequence of splices, as well as an+-- annotation indicating the final shape.+data NewShape d = NewShape+  { -- | The changes to perform.+    dimSplices :: [DimSplice d],+    -- | The resulting shape.+    newShape :: ShapeBase d+  }+  deriving (Eq, Ord, Show, Functor, Foldable, Traversable)++instance Semigroup (NewShape d) where+  NewShape ss1 _ <> NewShape ss2 shape = NewShape (ss1 <> ss2) shape+ -- | A primitive operation that returns something of known size and -- does not itself contain any bindings. data BasicOp@@ -365,7 +398,7 @@     Concat Int (NonEmpty VName) SubExp   | -- | Manifest an array with dimensions represented in the given     -- order.  The result will not alias anything.-    Manifest [Int] VName+    Manifest VName [Int]   | -- Array construction.      -- | @iota(n, x, s) = [x,x+s,..,x+(n-1)*s]@.@@ -378,13 +411,13 @@     Replicate Shape SubExp   | -- | Create array of given type and shape, with undefined elements.     Scratch PrimType [SubExp]-  | -- | 1st arg is the new shape, 2nd arg is the input array.-    Reshape ReshapeKind Shape VName+  | -- | 1st arg is the input array, 2nd arg is new shape.+    Reshape VName (NewShape SubExp)   | -- | Permute the dimensions of the input array.  The list     -- of integers is a list of dimensions (0-indexed), which     -- must be a permutation of @[0,n-1]@, where @n@ is the     -- number of dimensions in the input array.-    Rearrange [Int] VName+    Rearrange VName [Int]   | -- | Update an accumulator at the given index with the given     -- value. Consumes the accumulator and produces a new one. If     -- 'Safe', perform a run-time bounds check and ignore the write if

@@ -15,6 +15,8 @@     ShapeBase (..),     Shape,     stripDims,+    dropDims,+    takeDims,     Ext (..),     ExtSize,     ExtShape,@@ -123,10 +125,19 @@ instance Monoid (ShapeBase d) where   mempty = Shape mempty --- | @stripDims n shape@ strips the outer @n@ dimensions from--- @shape@.+-- | Alias for 'dropDims' stripDims :: Int -> ShapeBase d -> ShapeBase d-stripDims n (Shape dims) = Shape $ drop n dims+stripDims = dropDims++-- | @dropDims n shape@ strips the outer @n@ dimensions from+-- @shape@.+dropDims :: Int -> ShapeBase d -> ShapeBase d+dropDims n (Shape dims) = Shape $ drop n dims++-- | @takeDims n shape@ takes the outer @n@ dimensions from @shape@, up to the+-- number of dimensions in @shape@.+takeDims :: Int -> ShapeBase d -> ShapeBase d+takeDims n (Shape dims) = Shape $ take n dims  -- | The size of an array as a list of subexpressions.  If a variable, -- that variable must be in scope where this array is used.

@@ -149,21 +149,21 @@   BasicOp <$> (Replicate <$> mapOnShape tv shape <*> mapOnSubExp tv vexp) mapExpM tv (BasicOp (Scratch t shape)) =   BasicOp <$> (Scratch t <$> mapM (mapOnSubExp tv) shape)-mapExpM tv (BasicOp (Reshape kind shape arrexp)) =+mapExpM tv (BasicOp (Reshape arrexp newshape)) =   BasicOp-    <$> ( Reshape kind-            <$> mapM (mapOnSubExp tv) shape-            <*> mapOnVName tv arrexp+    <$> ( Reshape+            <$> mapOnVName tv arrexp+            <*> mapM (mapOnSubExp tv) newshape         )-mapExpM tv (BasicOp (Rearrange perm e)) =-  BasicOp <$> (Rearrange perm <$> mapOnVName tv e)+mapExpM tv (BasicOp (Rearrange e perm)) =+  BasicOp <$> (Rearrange <$> mapOnVName tv e <*> pure perm) mapExpM tv (BasicOp (Concat i (x :| ys) size)) = do   x' <- mapOnVName tv x   ys' <- mapM (mapOnVName tv) ys   size' <- mapOnSubExp tv size   pure $ BasicOp $ Concat i (x' :| ys') size'-mapExpM tv (BasicOp (Manifest perm e)) =-  BasicOp <$> (Manifest perm <$> mapOnVName tv e)+mapExpM tv (BasicOp (Manifest v perm)) =+  BasicOp <$> (Manifest <$> mapOnVName tv v <*> pure perm) mapExpM tv (BasicOp (Assert e msg loc)) =   BasicOp <$> (Assert <$> mapOnSubExp tv e <*> traverse (mapOnSubExp tv) msg <*> pure loc) mapExpM tv (BasicOp (Opaque op e)) =@@ -319,14 +319,14 @@   walkOnShape tv shape >> walkOnSubExp tv vexp walkExpM tv (BasicOp (Scratch _ shape)) =   mapM_ (walkOnSubExp tv) shape-walkExpM tv (BasicOp (Reshape _ shape arrexp)) =+walkExpM tv (BasicOp (Reshape arrexp shape)) =   mapM_ (walkOnSubExp tv) shape >> walkOnVName tv arrexp-walkExpM tv (BasicOp (Rearrange _ e)) =-  walkOnVName tv e+walkExpM tv (BasicOp (Rearrange v _)) =+  walkOnVName tv v walkExpM tv (BasicOp (Concat _ (x :| ys) size)) =   walkOnVName tv x >> mapM_ (walkOnVName tv) ys >> walkOnSubExp tv size-walkExpM tv (BasicOp (Manifest _ e)) =-  walkOnVName tv e+walkExpM tv (BasicOp (Manifest v _)) =+  walkOnVName tv v walkExpM tv (BasicOp (Assert e msg _)) =   walkOnSubExp tv e >> traverse_ (walkOnSubExp tv) msg walkExpM tv (BasicOp (Opaque _ e)) =

@@ -73,7 +73,7 @@   | UnknownVariableError VName   | UnknownFunctionError Name   | ParameterMismatch (Maybe Name) [Type] [Type]-  | SlicingError Int Int+  | SlicingError Shape Int   | BadAnnotation String Type Type   | ReturnAliased Name VName   | UniqueReturnAliased Name@@ -143,7 +143,7 @@       ngot = length got       fname' = maybe "anonymous function" (("function " ++) . nameToString) fname   show (SlicingError dims got) =-    show got ++ " indices given, but type of indexee has " ++ show dims ++ " dimension(s)."+    show got ++ " indices given, but type of indexee has shape " ++ prettyString dims   show (BadAnnotation desc expected got) =     "Annotation of \""       ++ desc@@ -834,9 +834,12 @@ checkSlice :: (Checkable rep) => Type -> Slice SubExp -> TypeM rep () checkSlice vt (Slice idxes) = do   when (arrayRank vt /= length idxes) . bad $-    SlicingError (arrayRank vt) (length idxes)+    SlicingError (arrayShape vt) (length idxes)   mapM_ (traverse $ require [Prim int64]) idxes +checkShape :: (Checkable rep) => Shape -> TypeM rep ()+checkShape = mapM_ (require [Prim int64])+ checkBasicOp :: (Checkable rep) => BasicOp -> TypeM rep () checkBasicOp (SubExp es) =   void $ checkSubExp es@@ -883,15 +886,11 @@ checkBasicOp (FlatIndex ident slice) = do   vt <- lookupType ident   observe ident-  when (arrayRank vt /= 1) $-    bad $-      SlicingError (arrayRank vt) 1+  when (arrayRank vt /= 1) $ bad $ SlicingError (arrayShape vt) 1   checkFlatSlice slice checkBasicOp (FlatUpdate src slice v) = do   (src_shape, src_pt) <- checkArrIdent src-  when (shapeRank src_shape /= 1) $-    bad $-      SlicingError (shapeRank src_shape) 1+  when (shapeRank src_shape /= 1) $ bad $ SlicingError src_shape 1    v_aliases <- lookupAliases v   when (src `nameIn` v_aliases) $@@ -909,17 +908,23 @@   mapM_ (require [Prim int64]) dims   void $ checkSubExp valexp checkBasicOp (Scratch _ shape) =-  mapM_ checkSubExp shape-checkBasicOp (Reshape k newshape arrexp) = do-  rank <- shapeRank . fst <$> checkArrIdent arrexp-  mapM_ (require [Prim int64]) $ shapeDims newshape-  case k of-    ReshapeCoerce ->-      when (shapeRank newshape /= rank) . bad $-        TypeError "Coercion changes rank of array."-    ReshapeArbitrary ->-      pure ()-checkBasicOp (Rearrange perm arr) = do+  checkShape $ Shape shape+checkBasicOp (Reshape arrexp newshape) = do+  (arr_shape, _) <- checkArrIdent arrexp+  checkShape $ newShape newshape+  spliced_shape <- foldM checkSplice arr_shape $ dimSplices newshape+  when (spliced_shape /= newShape newshape) . bad . TypeError $+    "Splice produces shape "+      <> prettyText spliced_shape+      <> " but annotation is shape "+      <> prettyText (newShape newshape)+  where+    checkSplice arr_shape sp@(DimSplice i k shape) = do+      checkShape shape+      when (i < 0 || i + k > shapeRank arr_shape) . bad . TypeError $+        "Splice " <> prettyText sp <> " cannot be applied to shape " <> prettyText arr_shape+      pure $ applySplice arr_shape sp+checkBasicOp (Rearrange arr perm) = do   arrt <- lookupType arr   let rank = arrayRank arrt   when (length perm /= rank || sort perm /= [0 .. rank - 1]) $@@ -933,8 +938,8 @@     bad $       TypeError "Types of arguments to concat do not match."   require [Prim int64] ressize-checkBasicOp (Manifest perm arr) =-  checkBasicOp $ Rearrange perm arr -- Basically same thing!+checkBasicOp (Manifest arr perm) =+  checkBasicOp $ Rearrange arr perm -- Basically same thing! checkBasicOp (Assert e (ErrorMsg parts) _) = do   require [Prim Bool] e   mapM_ checkPart parts

@@ -518,12 +518,11 @@           forM_ (zip mergepat' mergeinit) $ \(p, se) ->             unless (se == I.Var (I.paramName p)) $               letBindNames [I.paramName p] $-                BasicOp $-                  case se of-                    I.Var v-                      | not $ primType $ paramType p ->-                          Reshape I.ReshapeCoerce (I.arrayShape $ paramType p) v-                    _ -> SubExp se+                case se of+                  I.Var v+                    | not $ primType $ paramType p ->+                        shapeCoerce (I.arrayDims $ paramType p) v+                  _ -> BasicOp $ SubExp se            -- As the condition expression is inserted twice, we have to           -- avoid shadowing (#1935).@@ -551,12 +550,11 @@             forM_ (zip mergepat' ses) $ \(p, se) ->               unless (se == I.Var (I.paramName p)) $                 letBindNames [I.paramName p] $-                  BasicOp $-                    case se of-                      I.Var v-                        | not $ primType $ paramType p ->-                            Reshape I.ReshapeCoerce (I.arrayShape $ paramType p) v-                      _ -> SubExp se+                  case se of+                    I.Var v+                      | not $ primType $ paramType p ->+                          shapeCoerce (I.arrayDims $ paramType p) v+                    _ -> BasicOp $ SubExp se             subExpsRes <$> internaliseExp "loop_cond" cond           loop_end_cond <- bodyBind loop_end_cond_body @@ -667,7 +665,7 @@                 (I.Shape $ map (intConst Int64 . toInteger) new_shape)                 1                 $ I.arrayShape flat_arr_t-        letSubExp desc $ I.BasicOp $ I.Reshape I.ReshapeArbitrary new_shape' flat_arr+        letSubExp desc $ I.BasicOp $ I.Reshape flat_arr (reshapeAll (I.arrayShape flat_arr_t) new_shape')   | otherwise = do       es' <- mapM (internaliseExp "arr_elem") es       let arr_t_ext = foldMap toList $ internaliseType $ E.toStruct arr_t@@ -794,6 +792,25 @@       traceRes (nameToText tag) e'     "opaque" ->       mapM (letSubExp desc . BasicOp . Opaque OpaqueNil) e'+    "scratch" -> do+      ts <- mapM subExpType e'+      forM (zip ts e') $ \(t, se) ->+        case t of+          I.Array pt shape _ ->+            letSubExp desc $ I.BasicOp $ I.Scratch pt $ I.shapeDims shape+          I.Prim pt ->+            pure $ constant $ blankPrimValue pt+          _ -> pure se+    "blank" -> do+      ts <- mapM subExpType e'+      forM (zip ts e') $ \(t, se) ->+        case t of+          I.Array pt shape _ ->+            letSubExp desc . I.BasicOp . I.Replicate shape . constant $+              blankPrimValue pt+          I.Prim pt ->+            pure $ constant $ blankPrimValue pt+          _ -> pure se     _ ->       pure e'   where@@ -1582,9 +1599,11 @@                     x' <- letExp "x" $ I.BasicOp $ I.SubExp x                     y' <- letExp "x" $ I.BasicOp $ I.SubExp y                     x_flat <--                      letExp "x_flat" $ I.BasicOp $ I.Reshape I.ReshapeArbitrary (I.Shape [x_num_elems]) x'+                      letExp "x_flat" . I.BasicOp $+                        I.Reshape x' (reshapeAll (I.arrayShape x_t) (I.Shape [x_num_elems]))                     y_flat <--                      letExp "y_flat" $ I.BasicOp $ I.Reshape I.ReshapeArbitrary (I.Shape [x_num_elems]) y'+                      letExp "y_flat" . I.BasicOp $+                        I.Reshape y' (reshapeAll (I.arrayShape x_t) (I.Shape [x_num_elems]))                      -- Compare the elements.                     cmp_lam <- cmpOpLambda $ I.CmpEq (elemType x_t)@@ -1760,17 +1779,17 @@         forM arrs $ \arr' -> do           arr_t <- lookupType arr'           letSubExp desc . I.BasicOp $-            I.Reshape-              I.ReshapeArbitrary-              (reshapeOuter (I.Shape [n', m']) 1 $ I.arrayShape arr_t)-              arr'+            I.Reshape arr' $+              reshapeAll (I.arrayShape arr_t) $+                reshapeOuter (I.Shape [n', m']) 1 $+                  I.arrayShape arr_t     handleRest [arr] "manifest" = Just $ \desc -> do       arrs <- internaliseExpToVars "flatten_arr" arr       forM arrs $ \arr' -> do         r <- I.arrayRank <$> lookupType arr'         if r == 0           then pure $ I.Var arr'-          else letSubExp desc $ I.BasicOp $ I.Manifest [0 .. r - 1] arr'+          else letSubExp desc $ I.BasicOp $ I.Manifest arr' [0 .. r - 1]     handleRest [arr] "flatten" = Just $ \desc -> do       arrs <- internaliseExpToVars "flatten_arr" arr       forM arrs $ \arr' -> do@@ -1779,10 +1798,10 @@             m = arraySize 1 arr_t         k <- letSubExp "flat_dim" $ I.BasicOp $ I.BinOp (Mul Int64 I.OverflowUndef) n m         letSubExp desc . I.BasicOp $-          I.Reshape-            I.ReshapeArbitrary-            (reshapeOuter (I.Shape [k]) 2 $ I.arrayShape arr_t)-            arr'+          I.Reshape arr' $+            reshapeAll (I.arrayShape arr_t) $+              reshapeOuter (I.Shape [k]) 2 $+                I.arrayShape arr_t     handleRest [x, y] "concat" = Just $ \desc -> do       xs <- internaliseExpToVars "concat_x" x       ys <- internaliseExpToVars "concat_y" y@@ -1801,7 +1820,7 @@     handleRest [e] "transpose" = Just $ \desc ->       internaliseOperation desc e $ \v -> do         r <- I.arrayRank <$> lookupType v-        pure $ I.Rearrange ([1, 0] ++ [2 .. r - 1]) v+        pure $ I.Rearrange v ([1, 0] ++ [2 .. r - 1])     handleRest [x, y] "zip" = Just $ \desc ->       mapM (letSubExp "zip_copy" . BasicOp . Replicate mempty . I.Var)         =<< ( (++)@@ -1884,8 +1903,8 @@             "length of index and value array does not match"             loc         certifying c $-          letExp (baseString sv ++ "_write_sv") . I.BasicOp $-            I.Reshape I.ReshapeCoerce (reshapeOuter (I.Shape [si_w]) 1 sv_shape) sv+          letExp (baseString sv ++ "_write_sv") $+            shapeCoerce (I.shapeDims (reshapeOuter (I.Shape [si_w]) 1 sv_shape)) sv        indexType <- fmap rowType <$> mapM lookupType si'       indexName <- mapM (\_ -> newVName "write_index") indexType

@@ -322,7 +322,8 @@         Just prev ->           pure $ MonoKnown prev         Nothing -> do-          put (i + 1, M.insert d i m)+          -- Ensure that each instance of anySize is treated distinctly.+          put (i + 1, if d == anySize then m else M.insert d i m)           pure $ MonoKnown i  -- Mapping from function name and instance list to a new function name in case

@@ -182,7 +182,12 @@   pure $ Lambda params' body' retdecl ret loc transformExp e = astMap m e   where-    m = identityMapper {mapOnExp = transformExp}+    m =+      identityMapper+        { mapOnExp = transformExp,+          mapOnStructType = transformStructType,+          mapOnParamType = transformParamType+        }  onValBind :: ValBind -> RecordM ValBind onValBind vb = do

@@ -236,7 +236,7 @@      manifest perm array =       letExp (baseString array ++ "_coalesced") $-        BasicOp (Manifest perm array)+        BasicOp (Manifest array perm)  lookupPermutation :: LayoutTable -> VName -> IndexExprName -> VName -> Maybe Permutation lookupPermutation perm_table seg_name idx_name arr_name =

@@ -127,7 +127,7 @@   let ws = mapMaybe (getDirAliasedIxfn td_env coal_tab . patElemName) ys       rs = mapMaybe (getDirAliasedIxfn td_env coal_tab) (a : bs)    in Just (ws, ws ++ rs)-getUseSumFromStm td_env coal_tab (Let (Pat ys) _ (BasicOp (Manifest _perm x))) =+getUseSumFromStm td_env coal_tab (Let (Pat ys) _ (BasicOp (Manifest x _perm))) =   let ws = mapMaybe (getDirAliasedIxfn td_env coal_tab . patElemName) ys       rs = mapMaybe (getDirAliasedIxfn td_env coal_tab) [x]    in Just (ws, ws ++ rs)

@@ -71,10 +71,12 @@ getDirAliasFromExp :: Exp (Aliases rep) -> Maybe (VName, DirAlias) getDirAliasFromExp (BasicOp (SubExp (Var x))) = Just (x, Just) getDirAliasFromExp (BasicOp (Opaque _ (Var x))) = Just (x, Just)-getDirAliasFromExp (BasicOp (Reshape ReshapeCoerce shp x)) =-  Just (x, Just . (`LMAD.coerce` shapeDims (fmap pe64 shp)))-getDirAliasFromExp (BasicOp (Reshape ReshapeArbitrary shp x)) =-  Just (x, (`LMAD.reshape` shapeDims (fmap pe64 shp)))+getDirAliasFromExp (BasicOp (Reshape x shp)) =+  case reshapeKind shp of+    ReshapeCoerce ->+      Just (x, Just . (`LMAD.coerce` fmap pe64 (shapeDims $ newShape shp)))+    ReshapeArbitrary ->+      Just (x, (`LMAD.reshape` fmap pe64 (shapeDims $ newShape shp))) getDirAliasFromExp (BasicOp (Rearrange _ _)) =   Nothing getDirAliasFromExp (BasicOp (Index x slc)) =@@ -109,7 +111,7 @@ getInvAliasFromExp (BasicOp (SubExp (Var _))) = Just id getInvAliasFromExp (BasicOp (Opaque _ (Var _))) = Just id getInvAliasFromExp (BasicOp Update {}) = Just id-getInvAliasFromExp (BasicOp (Rearrange perm _)) =+getInvAliasFromExp (BasicOp (Rearrange _ perm)) =   Just (`LMAD.permute` rearrangeInverse perm) getInvAliasFromExp _ = Nothing 

@@ -48,13 +48,9 @@  isInnerCoal :: Env -> VName -> Stm GPU -> Bool isInnerCoal (_, ixfn_env) slc_X (Let (Pat [pe]) _ (BasicOp (Index x _)))-  | slc_X == patElemName pe,-    Nothing <- M.lookup x ixfn_env =-      True -- if not in the table, we assume not-transposed!-isInnerCoal (_, ixfn_env) slc_X (Let (Pat [pe]) _ (BasicOp (Index x _)))-  | slc_X == patElemName pe,-    Just lmad <- M.lookup x ixfn_env =-      innerHasStride1 lmad+  | slc_X == patElemName pe =+      -- if not in the table, we assume not-transposed!+      maybe True innerHasStride1 $ M.lookup x ixfn_env   where     innerHasStride1 lmad =       let lmad_dims = LMAD.dims lmad@@ -626,7 +622,7 @@                   ones = map (const $ intConst Int64 1) rem_outer_dims                   new_shape = Shape $ concat [ones, block_dims, ones, rest_dims]               letExp "res_reshaped" . BasicOp $-                Reshape ReshapeArbitrary new_shape epilogue_res+                Reshape epilogue_res (reshapeAll (arrayShape epilogue_t) new_shape)         pure [RegTileReturns mempty regtile_ret_dims epilogue_res'] mmBlkRegTilingNrm _ _ = pure Nothing @@ -1278,7 +1274,7 @@                     ones = map (const se1) rem_outer_dims                     new_shape = Shape $ concat [ones, block_dims, ones, rest_dims]                 letExp "res_reshaped" . BasicOp $-                  Reshape ReshapeArbitrary new_shape res+                  Reshape res (reshapeAll (arrayShape res_tp') new_shape)            pure $ map (RegTileReturns mempty regtile_ret_dims) epilogue_res'         -- END (ret_seggroup, stms_seggroup) <- runBuilder $ do@@ -1312,7 +1308,7 @@               arr_tp <- lookupType arr_nm               let perm = [i + 1 .. arrayRank arr_tp - 1] ++ [0 .. i]               let arr_tr_str = baseString arr_nm ++ "_transp"-              arr_tr_nm <- letExp arr_tr_str $ BasicOp $ Manifest perm arr_nm+              arr_tr_nm <- letExp arr_tr_str $ BasicOp $ Manifest arr_nm perm               let e_ind' = BasicOp $ Index arr_tr_nm slc               let stm' = Let patt yy e_ind'               pure (tab_inn, M.insert p_nm (ptp, stm') tab_out)

@@ -42,7 +42,7 @@   where     table = consts_table <> mkTable (bodyStms (funDefBody fd))     mkSubst (Var v0)-      | Just (MemArray _ _ _ (ArrayIn mem0 lmad0), BasicOp (Manifest _ v1)) <-+      | Just (MemArray _ _ _ (ArrayIn mem0 lmad0), BasicOp (Manifest v1 _)) <-           varInfo v0 table,         Just (MemArray _ _ _ (ArrayIn mem1 lmad1), _) <-           varInfo v1 table,

@@ -92,8 +92,8 @@       letBindNames [output] . BasicOp . SubExp . Var =<< H.applyTransforms ots v   ResNode _ -> pure mempty   TransNode output tr ia -> do-    (cs, e) <- H.transformToExp tr ia-    runBuilder_ $ certifying cs $ letBindNames [output] e+    (aux, e) <- H.transformToExp tr ia+    runBuilder_ $ auxing aux $ letBindNames [output] e   FreeNode _ -> pure mempty   DoNode stm lst -> do     lst' <- mapM (finalizeNode . fst) lst

@@ -297,7 +297,7 @@     pure $ MatchNode s []   e     | [output] <- patNames pat,-      Just (ia, tr) <- H.transformFromExp (stmAuxCerts aux) e ->+      Just (ia, tr) <- H.transformFromExp aux e ->         pure $ TransNode output tr ia   _ -> pure n nodeToSoacNode n = pure n

@@ -170,15 +170,16 @@     getRepRshpArr :: ((H.Input, NodeT), LParam SOACS) -> Maybe (RshpInp, Certs)     getRepRshpArr ((H.Input outtrsf arr_nm arr_tp, _nt), farg)       | rshp_trsfm H.:< other_trsfms <- H.viewf outtrsf,-        (H.Reshape c ReshapeArbitrary shp_flat) <- rshp_trsfm,+        H.Reshape aux shp_flat <- rshp_trsfm,+        ReshapeArbitrary <- reshapeKind shp_flat,         other_trsfms == mempty,         eltp <- paramDec farg,-        Just shp_flat' <- checkShp eltp shp_flat,+        Just shp_flat' <- checkShp eltp $ newShape shp_flat,         Array _ptp shp_unflat _ <- arr_tp,         Just shp_unflat' <- checkShp eltp shp_unflat,         shapeRank shp_flat' == 1,         shapeRank shp_flat' < shapeRank shp_unflat' =-          Just (((arr_nm, farg), (shp_flat', shp_unflat', eltp)), c)+          Just (((arr_nm, farg), (shp_flat', shp_unflat', eltp)), stmAuxCerts aux)     getRepRshpArr _ = Nothing     --     checkShp (Prim _) shp_arr = Just shp_arr

@@ -558,7 +558,7 @@   TryFusion (SOAC, SOAC.ArrayTransforms) iswim _ (SOAC.Screma w arrs form) ots   | Just [Futhark.Scan scan_fun nes] <- Futhark.isScanSOAC form,-    Just (map_pat, map_cs, map_w, map_fun) <- rwimPossible scan_fun,+    Just (map_pat, map_aux, map_w, map_fun) <- rwimPossible scan_fun,     Just nes_names <- mapM subExpVar nes = do       let nes_idents = zipWith Ident nes_names $ lambdaReturnType scan_fun           map_nes = map SOAC.identInput nes_idents@@ -594,7 +594,7 @@        pure         ( SOAC.Screma map_w map_arrs' (mapSOAC map_fun'),-          ots SOAC.|> SOAC.Rearrange map_cs perm+          ots SOAC.|> SOAC.Rearrange map_aux perm         ) iswim _ _ _ =   fail "ISWIM does not apply."@@ -768,7 +768,7 @@ pullReshape :: SOAC -> SOAC.ArrayTransforms -> TryFusion (SOAC, SOAC.ArrayTransforms) pullReshape soac ots = do   Just mapnest <- MapNest.fromSOAC soac-  SOAC.Reshape cs _kind newshape SOAC.:< ots' <- pure $ SOAC.viewf ots+  SOAC.Reshape cs newshape SOAC.:< ots' <- pure $ SOAC.viewf ots   -- This handles only the easy case where the underlying lambda is   -- scalar. The more complicated cases could also be handled, but   -- requires more tricky checks.@@ -776,7 +776,7 @@     all       ((== MapNest.depth mapnest) . arrayRank)       (MapNest.typeOf mapnest)-  mapnest' <- MapNest.reshape cs newshape mapnest+  mapnest' <- MapNest.reshape cs (newShape newshape) mapnest   soac' <- MapNest.toSOAC mapnest'   pure (soac', ots') 

@@ -261,7 +261,7 @@         ii /= length dims - 1,         perm <- [0 .. ii - 1] ++ [ii + 1 .. length dims - 1] ++ [ii] = do           (arr_tr, stms_tr) <- runBuilderT' $ do-            arr' <- letExp (baseString arr ++ "_trsp") $ BasicOp $ Rearrange perm arr -- Manifest [1,0] arr+            arr' <- letExp (baseString arr ++ "_trsp") $ BasicOp $ Rearrange arr perm             letExp (baseString arr' ++ "_opaque") $ BasicOp $ Opaque OpaqueNil $ Var arr'           let tab' = M.insert arr (perm, arr_tr, stms_tr) tab               slc' = Slice $ map (dims !!) perm

@@ -463,10 +463,10 @@       -- Can be replaced with 'graphHostOnly e' to disable migration.       -- A fix can be verified by enabling tests/migration/reuse4_scratch.fut       graphInefficientReturn s e-    BasicOp (Reshape _ s arr) -> do-      graphInefficientReturn (shapeDims s) e+    BasicOp (Reshape arr s) -> do+      graphInefficientReturn (shapeDims $ newShape s) e       one bs `reuses` arr-    BasicOp (Rearrange _ arr) -> do+    BasicOp (Rearrange arr _) -> do       graphInefficientReturn [] e       one bs `reuses` arr     -- Expressions with a cost linear to the size of their result arrays are

@@ -1039,7 +1039,7 @@       _ -> pure v  instance (Simplifiable d) => Simplifiable (ShapeBase d) where-  simplify = fmap Shape . simplify . shapeDims+  simplify = traverse simplify  instance Simplifiable ExtSize where   simplify (Free se) = Free <$> simplify se

@@ -86,11 +86,11 @@           letExp "concat_rearrange" $             BasicOp $               Concat 0 (x' :| xs') new_d-      letBind pat $ BasicOp $ Rearrange perm concat_rearrange+      letBind pat $ BasicOp $ Rearrange concat_rearrange perm   where     transposedBy perm1 v =       case ST.lookupExp v vtable of-        Just (BasicOp (Rearrange perm2 v'), vcs)+        Just (BasicOp (Rearrange v' perm2), vcs)           | perm1 == perm2 -> Just (v', vcs)         _ -> Nothing @@ -191,9 +191,11 @@     isFullSlice (arrayShape dest_t) is = Simplify . auxing aux $       case se of         Var v | not $ null $ sliceDims is -> do+          v_t <- lookupType v           v_reshaped <-             letSubExp (baseString v ++ "_reshaped") . BasicOp $-              Reshape ReshapeArbitrary (arrayShape dest_t) v+              Reshape v $+                reshapeAll (arrayShape v_t) (arrayShape dest_t)           letBind pat $ BasicOp $ Replicate mempty v_reshaped         _ -> letBind pat $ BasicOp $ ArrayLit [se] $ rowType dest_t ruleBasicOp vtable pat (StmAux cs1 attrs _) (Update safety1 dest1 is1 (Var v1))@@ -254,29 +256,24 @@          in letBind pat $ BasicOp $ Replicate (Shape [n]) se ruleBasicOp vtable pat aux (Index idd slice)   | Just inds <- sliceIndices slice,-    Just (BasicOp (Reshape k newshape idd2), idd_cs) <- ST.lookupExp idd vtable,-    length newshape == length inds =-      Simplify $-        case k of-          ReshapeCoerce ->-            certifying idd_cs . auxing aux . letBind pat . BasicOp $-              Index idd2 slice-          ReshapeArbitrary -> do-            -- Linearise indices and map to old index space.-            oldshape <- arrayDims <$> lookupType idd2-            let new_inds =-                  reshapeIndex-                    (map pe64 oldshape)-                    (map pe64 $ shapeDims newshape)-                    (map pe64 inds)-            new_inds' <--              mapM (toSubExp "new_index") new_inds-            certifying idd_cs . auxing aux $-              letBind pat $-                BasicOp $-                  Index idd2 $-                    Slice $-                      map DimFix new_inds'+    Just (BasicOp (Reshape idd2 newshape), idd_cs) <- ST.lookupExp idd vtable,+    shapeRank (newShape newshape) == length inds = Simplify $+      case reshapeKind newshape of+        ReshapeCoerce ->+          certifying idd_cs . auxing aux . letBind pat . BasicOp $+            Index idd2 slice+        ReshapeArbitrary -> do+          -- Linearise indices and map to old index space.+          oldshape <- arrayDims <$> lookupType idd2+          let new_inds =+                reshapeIndex+                  (map pe64 oldshape)+                  (map pe64 $ shapeDims $ newShape newshape)+                  (map pe64 inds)+          new_inds' <-+            mapM (toSubExp "new_index") new_inds+          certifying idd_cs . auxing aux . letBind pat . BasicOp $+            Index idd2 (Slice $ map DimFix new_inds')  -- Copying an iota is pointless; just make it an iota instead. ruleBasicOp vtable pat aux (Replicate (Shape []) (Var v))@@ -286,18 +283,16 @@           BasicOp $             Iota n x s it -- Handle identity permutation.-ruleBasicOp _ pat _ (Rearrange perm v)+ruleBasicOp _ pat _ (Rearrange v perm)   | sort perm == perm =       Simplify $ letBind pat $ BasicOp $ SubExp $ Var v-ruleBasicOp vtable pat aux (Rearrange perm v)-  | Just (BasicOp (Rearrange perm2 e), v_cs) <- ST.lookupExp v vtable =+ruleBasicOp vtable pat aux (Rearrange v perm)+  | Just (BasicOp (Rearrange e perm2), v_cs) <- ST.lookupExp v vtable =       -- Rearranging a rearranging: compose the permutations.-      Simplify . certifying v_cs . auxing aux $-        letBind pat $-          BasicOp $-            Rearrange (perm `rearrangeCompose` perm2) e+      Simplify . certifying v_cs . auxing aux . letBind pat . BasicOp $+        Rearrange e (perm `rearrangeCompose` perm2) -- Rearranging a replicate where the outer dimension is left untouched.-ruleBasicOp vtable pat aux (Rearrange perm v1)+ruleBasicOp vtable pat aux (Rearrange v1 perm)   | Just (BasicOp (Replicate dims (Var v2)), v1_cs) <- ST.lookupExp v1 vtable,     num_dims <- shapeRank dims,     (rep_perm, rest_perm) <- splitAt num_dims perm,@@ -307,9 +302,8 @@         certifying v1_cs $           auxing aux $ do             v <--              letSubExp "rearrange_replicate" $-                BasicOp $-                  Rearrange (map (subtract num_dims) rest_perm) v2+              letSubExp "rearrange_replicate" . BasicOp $+                Rearrange v2 (map (subtract num_dims) rest_perm)             letBind pat $ BasicOp $ Replicate dims v  -- Simplify away 0<=i when 'i' is from a loop of form 'for i < n'.@@ -357,15 +351,55 @@       Simplify . auxing aux $ letBind pat $ BasicOp $ SubExp $ Var acc -- Manifest of a a copy (or another Manifest) can be simplified to -- manifesting the original array, if it is still available.-ruleBasicOp vtable pat aux (Manifest perm v1)+ruleBasicOp vtable pat aux (Manifest v1 perm)   | Just (Replicate (Shape []) (Var v2), cs) <- ST.lookupBasicOp v1 vtable,     ST.available v2 vtable =       Simplify . auxing aux . certifying cs . letBind pat . BasicOp $-        Manifest perm v2-  | Just (Manifest _ v2, cs) <- ST.lookupBasicOp v1 vtable,+        Manifest v2 perm+  | Just (Manifest v2 _, cs) <- ST.lookupBasicOp v1 vtable,     ST.available v2 vtable =       Simplify . auxing aux . certifying cs . letBind pat . BasicOp $-        Manifest perm v2+        Manifest v2 perm+ruleBasicOp vtable pat aux (Reshape v2 v3_shape)+  | ReshapeArbitrary <- reshapeKind v3_shape,+    Just (Rearrange v1 perm, v2_cs) <- ST.lookupBasicOp v2 vtable,+    Just (Reshape v0 v1_shape, v1_cs) <- ST.lookupBasicOp v1 vtable,+    ReshapeArbitrary <- reshapeKind v1_shape,+    Just v0_shape <- arrayShape <$> ST.lookupType v0 vtable =+      case ( flipReshapeRearrange (shapeDims v0_shape) (shapeDims (newShape v1_shape)) perm,+             flipRearrangeReshape perm v3_shape+           ) of+        (Just perm', _) -> Simplify $ do+          v1' <- letExp (baseString v1) $ BasicOp $ Rearrange v0 perm'+          v1_shape' <- arrayShape <$> lookupType v1'+          auxing aux . certifying (v1_cs <> v2_cs) . letBind pat $+            BasicOp (Reshape v1' (reshapeAll v1_shape' (newShape v3_shape)))+        (_, Just (v3_shape', perm')) -> Simplify $ do+          v2' <-+            auxing aux . certifying (v1_cs <> v2_cs) . letExp (baseString v2) $+              BasicOp (Reshape v1 v3_shape')+          letBind pat $ BasicOp (Rearrange v2' perm')+        _ ->+          Skip+-- Reshaping or transposing a copy is almost always better done by copying the+-- result instead, because that improves the likelihood that the copy will be+-- eliminated.+ruleBasicOp vtable pat aux (Reshape v2 newshape)+  | Just (Replicate (Shape []) (Var v1), cs) <- ST.lookupBasicOp v2 vtable,+    ST.available v1 vtable =+      Simplify $ do+        v1' <-+          certifying cs . auxing aux . letExp (baseString v1) . BasicOp $+            Reshape v1 newshape+        letBind pat $ BasicOp $ Replicate (Shape []) (Var v1')+ruleBasicOp vtable pat aux (Rearrange v2 perm)+  | Just (Replicate (Shape []) (Var v1), cs) <- ST.lookupBasicOp v2 vtable,+    ST.available v1 vtable =+      Simplify $ do+        v1' <-+          certifying cs . auxing aux . letExp (baseString v1) . BasicOp $+            Rearrange v1 perm+        letBind pat $ BasicOp $ Replicate (Shape []) (Var v1') ruleBasicOp _ _ _ _ =   Skip 

@@ -56,6 +56,18 @@   Maybe (m IndexResult) simplifyIndexing vtable seType idd (Slice inds) consuming consumed =   case defOf idd of+    -- FIXME: This is a special case to avoid simplifying away a slice of a+    -- rearrange. This is because register tiling cannot otherwise properly+    -- detect what is going on.+    Just (Rearrange src perm, cs)+      | rearrangeReach perm <= length (takeWhile isIndex inds) ->+          let inds' = rearrangeShape (rearrangeInverse perm) inds+           in Just $ pure $ IndexResult cs src $ Slice inds'+      | any isIndex inds ->+          Nothing+      where+        isIndex DimFix {} = True+        isIndex _ = False     _       | Just t <- seType (Var idd),         Slice inds == fullSlice t [] ->@@ -76,6 +88,7 @@               <$> mapM (toSubExp "index_primexp") inds''       | Just (ST.IndexedArray cs arr inds'') <-           ST.index' idd (fixSlice (pe64 <$> Slice inds) (map fst matches)) vtable,+        length inds == length inds'',         all (worthInlining . untyped) inds'',         arr `ST.available` vtable,         all (`ST.elem` vtable) (unCerts cs),@@ -92,7 +105,7 @@                     =<< sequence inds'''               arr_sliced_tr <-                 letSubExp (baseString arr_sliced <> "_tr") $-                  BasicOp (Rearrange perm arr_sliced)+                  BasicOp (Rearrange arr_sliced perm)               pure $ SubExpResult mempty arr_sliced_tr       where         matches = zip fakeIndices $ zip [0 :: Int ..] $ sliceDims $ Slice inds@@ -164,13 +177,6 @@       where         index DimFix {} = Nothing         index (DimSlice _ n s) = Just (n, DimSlice (constant (0 :: Int64)) n s)-    Just (Rearrange perm src, cs)-      | rearrangeReach perm <= length (takeWhile isIndex inds) ->-          let inds' = rearrangeShape (rearrangeInverse perm) inds-           in Just $ pure $ IndexResult cs src $ Slice inds'-      where-        isIndex DimFix {} = True-        isIndex _ = False     Just (Replicate (Shape []) (Var src), cs)       | Just dims <- arrayDims <$> seType (Var src),         length inds == length dims,@@ -184,17 +190,19 @@         not consuming,         ST.available src vtable ->           Just $ pure $ IndexResult cs src $ Slice inds-    Just (Reshape ReshapeCoerce newshape src, cs)-      | Just olddims <- arrayDims <$> seType (Var src),-        changed_dims <- zipWith (/=) (shapeDims newshape) olddims,+    Just (Reshape src newshape, cs)+      | ReshapeCoerce <- reshapeKind newshape,+        Just olddims <- arrayDims <$> seType (Var src),+        changed_dims <- zipWith (/=) (shapeDims (newShape newshape)) olddims,         not $ or $ drop (length inds) changed_dims ->           Just $ pure $ IndexResult cs src $ Slice inds       | Just olddims <- arrayDims <$> seType (Var src),         length newshape == length inds,-        length olddims == length (shapeDims newshape) ->+        length olddims == length (shapeDims (newShape newshape)) ->           Just $ pure $ IndexResult cs src $ Slice inds-    Just (Reshape _ (Shape [_]) v2, cs)-      | Just [_] <- arrayDims <$> seType (Var v2) ->+    Just (Reshape v2 newshape, cs)+      | Shape [_] <- newShape newshape,+        Just [_] <- arrayDims <$> seType (Var v2) ->           Just $ pure $ IndexResult cs v2 $ Slice inds     Just (Concat d (x :| xs) _, cs)       | -- HACK: simplifying the indexing of an N-array concatenation

@@ -268,54 +268,45 @@ simplifyAssert _ _ _ =   Nothing --- No-op reshape.-simplifyIdentityReshape :: SimpleRule rep-simplifyIdentityReshape _ seType (Reshape _ newshape v)+simplifyReshape :: SimpleRule rep+simplifyReshape defOf seType (Reshape v newshape)+  -- No-op reshape   | Just t <- seType $ Var v,-    newshape == arrayShape t =+    newShape newshape == arrayShape t =       resIsSubExp $ Var v-simplifyIdentityReshape _ _ _ = Nothing--simplifyReshapeReshape :: SimpleRule rep-simplifyReshapeReshape defOf _ (Reshape k1 newshape v)-  | Just (BasicOp (Reshape k2 _ v2), v_cs) <- defOf v =-      Just (Reshape (max k1 k2) newshape v2, v_cs)-simplifyReshapeReshape _ _ _ = Nothing--simplifyReshapeScratch :: SimpleRule rep-simplifyReshapeScratch defOf _ (Reshape _ newshape v)+  -- Simplifying a reshape+  | Just shape <- arrayShape <$> seType (Var v),+    Just newshape' <- simplifyNewShape shape newshape =+      Just (Reshape v newshape', mempty)+  -- Reshape-of-reshape+  | Just (BasicOp (Reshape v2 oldnewshape), v_cs) <- defOf v =+      Just (Reshape v2 (oldnewshape <> newshape), v_cs)+  -- Reshape-of-scratch   | Just (BasicOp (Scratch bt _), v_cs) <- defOf v =-      Just (Scratch bt $ shapeDims newshape, v_cs)-simplifyReshapeScratch _ _ _ = Nothing--simplifyReshapeReplicate :: SimpleRule rep-simplifyReshapeReplicate defOf seType (Reshape _ newshape v)+      Just (Scratch bt $ shapeDims $ newShape newshape, v_cs)+  -- Reshape-of-replicate   | Just (BasicOp (Replicate _ se), v_cs) <- defOf v,     Just oldshape <- arrayShape <$> seType se,-    shapeDims oldshape `isSuffixOf` shapeDims newshape =+    newshape' <- newShape newshape,+    shapeDims oldshape `isSuffixOf` shapeDims newshape' =       let new =-            take (length newshape - shapeRank oldshape) $-              shapeDims newshape+            take (shapeRank newshape' - shapeRank oldshape) $+              shapeDims newshape'        in Just (Replicate (Shape new) se, v_cs)-simplifyReshapeReplicate _ _ _ = Nothing--simplifyReshapeIota :: SimpleRule rep-simplifyReshapeIota defOf _ (Reshape _ newshape v)+  -- Reshape-of-iota   | Just (BasicOp (Iota _ offset stride it), v_cs) <- defOf v,-    [n] <- shapeDims newshape =+    [n] <- shapeDims (newShape newshape) =       Just (Iota n offset stride it, v_cs)-simplifyReshapeIota _ _ _ = Nothing--simplifyReshapeConcat :: SimpleRule rep-simplifyReshapeConcat defOf seType (Reshape ReshapeCoerce newshape v) = do-  (BasicOp (Concat d arrs _), v_cs) <- defOf v-  (bef, w', aft) <- focusNth d $ shapeDims newshape-  (arr_bef, _, arr_aft) <--    focusNth d <=< fmap arrayDims $ seType $ Var $ NE.head arrs-  guard $ arr_bef == bef-  guard $ arr_aft == aft-  Just (Concat d arrs w', v_cs)-simplifyReshapeConcat _ _ _ = Nothing+  -- Reshape-of-concat+  | ReshapeCoerce <- reshapeKind newshape = do+      (BasicOp (Concat d arrs _), v_cs) <- defOf v+      (bef, w', aft) <- focusNth d $ shapeDims $ newShape newshape+      (arr_bef, _, arr_aft) <-+        focusNth d <=< fmap arrayDims $ seType $ Var $ NE.head arrs+      guard $ arr_bef == bef+      guard $ arr_aft == aft+      Just (Concat d arrs w', v_cs)+simplifyReshape _ _ _ = Nothing  reshapeSlice :: [DimIndex d] -> [d] -> [DimIndex d] reshapeSlice (DimFix i : slice') scs =@@ -327,9 +318,10 @@ -- If we are size-coercing a slice, then we might as well just use a -- different slice instead. simplifyReshapeIndex :: SimpleRule rep-simplifyReshapeIndex defOf _ (Reshape ReshapeCoerce newshape v)-  | Just (BasicOp (Index v' slice), v_cs) <- defOf v,-    slice' <- Slice $ reshapeSlice (unSlice slice) $ shapeDims newshape,+simplifyReshapeIndex defOf _ (Reshape v newshape)+  | ReshapeCoerce <- reshapeKind newshape,+    Just (BasicOp (Index v' slice), v_cs) <- defOf v,+    slice' <- Slice $ reshapeSlice (unSlice slice) $ shapeDims $ newShape newshape,     slice' /= slice =       Just (Index v' slice', v_cs) simplifyReshapeIndex _ _ _ = Nothing@@ -338,7 +330,8 @@ -- instead use the original array and update the slice dimensions. simplifyUpdateReshape :: SimpleRule rep simplifyUpdateReshape defOf seType (Update safety dest slice (Var v))-  | Just (BasicOp (Reshape ReshapeCoerce _ v'), v_cs) <- defOf v,+  | Just (BasicOp (Reshape v' newshape), v_cs) <- defOf v,+    ReshapeCoerce <- reshapeKind newshape,     Just ds <- arrayDims <$> seType (Var v'),     slice' <- Slice $ reshapeSlice (unSlice slice) ds,     slice' /= slice =@@ -357,9 +350,9 @@       case defOf v of         Just (BasicOp Scratch {}, cs) ->           Just cs-        Just (BasicOp (Rearrange _ v'), cs) ->+        Just (BasicOp (Rearrange v' _), cs) ->           (cs <>) <$> isActuallyScratch v'-        Just (BasicOp (Reshape _ _ v'), cs) ->+        Just (BasicOp (Reshape v' _), cs) ->           (cs <>) <$> isActuallyScratch v'         _ -> Nothing repScratchToScratch _ _ _ =@@ -373,12 +366,7 @@     simplifyConvOp,     simplifyAssert,     repScratchToScratch,-    simplifyIdentityReshape,-    simplifyReshapeReshape,-    simplifyReshapeScratch,-    simplifyReshapeReplicate,-    simplifyReshapeIota,-    simplifyReshapeConcat,+    simplifyReshape,     simplifyReshapeIndex,     simplifyUpdateReshape   ]

@@ -31,7 +31,7 @@     onStms scope stms =       modifyNameSource $         runState $-          runReaderT (optimiseStms (M.empty, M.empty) stms) scope+          runReaderT (optimiseStms (M.empty, initialIxFnEnv scope) stms) scope  optimiseBody :: Env -> Body GPU -> TileM (Body GPU) optimiseBody env (Body () stms res) =@@ -741,7 +741,7 @@       else do         let new_shape = Shape $ unit_dims ++ arrayDims arr_t         letExp (baseString arr) . BasicOp $-          Reshape ReshapeArbitrary new_shape arr+          Reshape arr (reshapeAll (arrayShape arr_t) new_shape)   let tile_dims = zip (map snd dims_on_top) unit_dims ++ dims   pure $ TileReturns mempty tile_dims arr' 

@@ -13,6 +13,7 @@     varianceInStms,     isTileableRedomap,     changeEnv,+    initialIxFnEnv,     TileKind (..),   ) where@@ -275,6 +276,16 @@   ixfn_env' <- changeIxFnEnv ixfn_env y e   pure (with_env', ixfn_env') +-- | Construct an initial 'IxFnEnv' where it is assumed that every array+-- parameter in the scope has a row-major index function.+initialIxFnEnv :: Scope GPU -> IxFnEnv+initialIxFnEnv = M.mapMaybe f+  where+    f info =+      case typeOf info of+        Array _ shape _ -> Just $ LMAD.iota 0 $ map pe64 $ shapeDims shape+        _ -> Nothing+ changeWithEnv :: WithEnv -> Exp GPU -> TileM WithEnv changeWithEnv with_env (WithAcc accum_decs inner_lam) = do   let bindings = map mapfun accum_decs@@ -302,11 +313,13 @@         _ -> env  changeIxFnEnv :: IxFnEnv -> VName -> Exp GPU -> TileM IxFnEnv-changeIxFnEnv env y (BasicOp (Reshape ReshapeArbitrary shp_chg x)) =-  composeIxfuns env y x (`LMAD.reshape` fmap ExpMem.pe64 (shapeDims shp_chg))-changeIxFnEnv env y (BasicOp (Reshape ReshapeCoerce shp_chg x)) =-  composeIxfuns env y x (Just . (`LMAD.coerce` fmap ExpMem.pe64 (shapeDims shp_chg)))-changeIxFnEnv env y (BasicOp (Manifest perm x)) = do+changeIxFnEnv env y (BasicOp (Reshape x shp_chg)) =+  case reshapeKind shp_chg of+    ReshapeCoerce ->+      composeIxfuns env y x (Just . (`LMAD.coerce` fmap ExpMem.pe64 (shapeDims $ newShape shp_chg)))+    ReshapeArbitrary ->+      composeIxfuns env y x (`LMAD.reshape` fmap ExpMem.pe64 (shapeDims $ newShape shp_chg))+changeIxFnEnv env y (BasicOp (Manifest x perm)) = do   tp <- lookupType x   case tp of     Array _ptp shp _u -> do@@ -314,7 +327,7 @@       let ixfn = LMAD.permute (LMAD.iota 0 shp') perm       pure $ M.insert y ixfn env     _ -> error "In TileLoops/Shared.hs, changeIxFnEnv: manifest applied to a non-array!"-changeIxFnEnv env y (BasicOp (Rearrange perm x)) =+changeIxFnEnv env y (BasicOp (Rearrange x perm)) =   composeIxfuns env y x (Just . (`LMAD.permute` perm)) changeIxFnEnv env y (BasicOp (Index x slc)) =   composeIxfuns env y x (Just . (`LMAD.slice` Slice (map (fmap ExpMem.pe64) $ unSlice slc)))

@@ -177,11 +177,11 @@   (Allocable fromrep torep inner) =>   Exp torep ->   AllocM fromrep torep (Exp torep)-repairExpression (BasicOp (Reshape k shape v)) = do+repairExpression (BasicOp (Reshape v shape)) = do   v_mem <- fst <$> lookupArraySummary v   space <- lookupMemSpace v_mem   v' <- snd <$> ensureDirectArray (Just space) v-  pure $ BasicOp $ Reshape k shape v'+  pure $ BasicOp $ Reshape v' shape repairExpression e =   error $ "repairExpression:\n" <> prettyString e @@ -535,7 +535,7 @@             MemArray pt shape u . ArrayIn mem $               LMAD.permute (LMAD.iota 0 $ map pe64 $ arrayDims t) perm           pat = Pat [PatElem v' info]-      addStm $ Let pat (defAux ()) $ BasicOp $ Manifest perm v+      addStm $ Let pat (defAux ()) $ BasicOp $ Manifest v perm       pure (mem, v')     _ ->       error $ "allocPermArray: " ++ prettyString t

@@ -97,7 +97,7 @@   fmap (Inner . GPUBody ts) . buildBody_ . allocInStms stms $ pure res  kernelExpHints :: Exp GPUMem -> AllocM GPU GPUMem [ExpHint]-kernelExpHints (BasicOp (Manifest perm v)) = do+kernelExpHints (BasicOp (Manifest v perm)) = do   dims <- arrayDims <$> lookupType v   let perm_inv = rearrangeInverse perm       dims' = rearrangeShape perm dims

@@ -576,14 +576,14 @@         -- Move the to-be-replicated dimensions outermost.         arr_tr <-           letExp (baseString arr <> "_tr") . BasicOp $-            Rearrange ([nest_r .. arr_r - 1] ++ [0 .. nest_r - 1]) arr+            Rearrange arr ([nest_r .. arr_r - 1] ++ [0 .. nest_r - 1])         -- Replicate the now-outermost dimensions appropriately.         arr_tr_rep <-           letExp (baseString arr <> "_tr_rep") . BasicOp $             Replicate shape (Var arr_tr)         -- Move the replicated dimensions back where they belong.         letBind outerpat . BasicOp $-          Rearrange ([res_r - nest_r .. res_r - 1] ++ [0 .. res_r - nest_r - 1]) arr_tr_rep+          Rearrange arr_tr_rep ([res_r - nest_r .. res_r - 1] ++ [0 .. res_r - nest_r - 1]) maybeDistributeStm stm@(Let _ aux (BasicOp (Replicate shape v))) acc = do   distributeSingleStm acc stm >>= \case     Just (kernels, _, nest, acc')@@ -602,7 +602,7 @@   | not $ primType $ typeOf pe =       distributeSingleUnaryStm acc stm stm_arr $ \_ outerpat arr ->         pure $ oneStm $ Let outerpat aux $ BasicOp $ Replicate mempty $ Var arr-maybeDistributeStm stm@(Let _ aux (BasicOp (Rearrange perm stm_arr))) acc =+maybeDistributeStm stm@(Let _ aux (BasicOp (Rearrange stm_arr perm))) acc =   distributeSingleUnaryStm acc stm stm_arr $ \nest outerpat arr -> do     let r = length (snd nest) + 1         perm' = [0 .. r - 1] ++ map (+ r) perm@@ -613,12 +613,13 @@     pure $       stmsFromList         [ Let (Pat [PatElem arr' arr_t]) aux $ BasicOp $ Replicate mempty $ Var arr,-          Let outerpat aux $ BasicOp $ Rearrange perm' arr'+          Let outerpat aux $ BasicOp $ Rearrange arr' perm'         ]-maybeDistributeStm stm@(Let _ aux (BasicOp (Reshape k reshape stm_arr))) acc =+maybeDistributeStm stm@(Let _ aux (BasicOp (Reshape stm_arr reshape))) acc =   distributeSingleUnaryStm acc stm stm_arr $ \nest outerpat arr -> do-    let reshape' = Shape (kernelNestWidths nest) <> reshape-    pure $ oneStm $ Let outerpat aux $ BasicOp $ Reshape k reshape' arr+    let outer = Shape (kernelNestWidths nest)+        reshape' = reshapeCoerce outer <> newshapeInner outer reshape+    pure $ oneStm $ Let outerpat aux $ BasicOp $ Reshape arr reshape' maybeDistributeStm stm@(Let pat aux (BasicOp (Update _ arr slice (Var v)))) acc   | not $ null $ sliceDims slice =       distributeSingleStm acc stm >>= \case

@@ -24,7 +24,7 @@   [(SubExp, VName)] ->   Maybe (m ()) iswim res_pat w scan_fun scan_input-  | Just (map_pat, map_cs, map_w, map_fun) <- rwimPossible scan_fun = Just $ do+  | Just (map_pat, map_aux, map_w, map_fun) <- rwimPossible scan_fun = Just $ do       let (accs, arrs) = unzip scan_input       arrs' <- transposedArrays arrs       accs' <- mapM (letExp "acc" . BasicOp . SubExp) accs@@ -65,18 +65,14 @@           mapM (newIdent' (<> "_transposed") . transposeIdentType) $             patIdents res_pat -      addStm $-        Let res_pat' (StmAux map_cs mempty ()) $-          Op $-            Screma map_w map_arrs' (mapSOAC map_fun')+      addStm . Let res_pat' map_aux . Op $+        Screma map_w map_arrs' (mapSOAC map_fun')        forM_ (zip (patIdents res_pat) (patIdents res_pat')) $ \(to, from) -> do         let perm = [1, 0] ++ [2 .. arrayRank (identType from) - 1]         addStm $-          Let (basicPat [to]) (defAux ()) $-            BasicOp $-              Rearrange perm $-                identName from+          Let (basicPat [to]) (defAux ()) . BasicOp $+            Rearrange (identName from) perm   | otherwise = Nothing  -- | Interchange Reduce With Inner Map. Tries to turn a @reduce(map)@ into a@@ -90,7 +86,7 @@   [(SubExp, VName)] ->   Maybe (m ()) irwim res_pat w comm red_fun red_input-  | Just (map_pat, map_cs, map_w, map_fun) <- rwimPossible red_fun = Just $ do+  | Just (map_pat, map_aux, map_w, map_fun) <- rwimPossible red_fun = Just $ do       let (accs, arrs) = unzip red_input       arrs' <- transposedArrays arrs       -- FIXME?  Can we reasonably assume that the accumulator is a@@ -136,18 +132,15 @@        let map_fun' = Lambda map_params map_rettype map_body -      addStm $-        Let res_pat (StmAux map_cs mempty ()) $-          Op $-            Screma map_w arrs' $-              mapSOAC map_fun'+      addStm . Let res_pat map_aux . Op . Screma map_w arrs' $+        mapSOAC map_fun'   | otherwise = Nothing  -- | Does this reduce operator contain an inner map, and if so, what -- does that map look like? rwimPossible ::   Lambda SOACS ->-  Maybe (Pat Type, Certs, SubExp, Lambda SOACS)+  Maybe (Pat Type, StmAux (), SubExp, Lambda SOACS) rwimPossible fun   | Body _ stms res <- lambdaBody fun,     [stm] <- stmsToList stms, -- Body has a single binding@@ -156,7 +149,7 @@     Op (Screma map_w map_arrs form) <- stmExp stm,     Just map_fun <- isMapSOAC form,     map paramName (lambdaParams fun) == map_arrs =-      Just (map_pat, stmCerts stm, map_w, map_fun)+      Just (map_pat, stmAux stm, map_w, map_fun)   | otherwise =       Nothing @@ -164,7 +157,7 @@ transposedArrays arrs = forM arrs $ \arr -> do   t <- lookupType arr   let perm = [1, 0] ++ [2 .. arrayRank t - 1]-  letExp (baseString arr) $ BasicOp $ Rearrange perm arr+  letExp (baseString arr) $ BasicOp $ Rearrange arr perm  removeParamOuterDim :: LParam SOACS -> LParam SOACS removeParamOuterDim param =

@@ -13,6 +13,7 @@ where  import Control.Monad.Reader+import Data.Map qualified as M import Data.Sequence (Seq (..)) import Futhark.Analysis.Alias (aliasAnalysis) import Futhark.IR.Aliases@@ -32,7 +33,7 @@     }   where     onFun f = f {funDefBody = onBody (funDefBody f)}-    onBody body = runReader (liftAllocationsInBody body) (Env onOp)+    onBody body = runReader (liftAllocationsInBody body) (Env onOp mempty)  liftAllocationsSeqMem :: Pass SeqMem SeqMem liftAllocationsSeqMem =@@ -49,8 +50,10 @@   Pass "lift allocations mc" "lift allocations mc" $     pure . liftInProg liftAllocationsInMCOp -newtype Env inner = Env-  {onInner :: inner -> LiftM inner inner}+data Env inner = Env+  { onInner :: inner -> LiftM inner inner,+    envAliases :: AliasesAndConsumed+  }  type LiftM inner a = Reader (Env inner) a @@ -59,7 +62,7 @@   Body rep ->   LiftM (inner rep) (Body rep) liftAllocationsInBody body = do-  stms <- liftAllocationsInStms (bodyStms body) mempty mempty mempty+  stms <- liftAllocationsInStms (bodyStms body)   pure $ body {bodyStms = stms}  liftInsideStm ::@@ -80,54 +83,65 @@ liftInsideStm stm = pure stm  liftAllocationsInStms ::+  forall rep inner.   (Mem rep inner, Aliased rep) =>-  -- | The input stms   Stms rep ->-  -- | The lifted allocations and associated statements-  Stms rep ->-  -- | The other statements processed so far-  Stms rep ->-  -- | (Names we need to lift, consumed names)-  (Names, Names) ->   LiftM (inner rep) (Stms rep)-liftAllocationsInStms Empty lifted acc _ = pure $ lifted <> acc-liftAllocationsInStms (stms :|> stm) lifted acc (to_lift, consumed) = do-  stm' <- liftInsideStm stm-  case stmExp stm' of-    BasicOp Assert {} -> liftStm stm'-    Op Alloc {} -> liftStm stm'-    _ -> do-      let pat_names = namesFromList $ patNames $ stmPat stm'-      if (pat_names `namesIntersect` to_lift)-        || namesIntersect consumed (freeIn stm)-        then liftStm stm'-        else dontLiftStm stm'-  where-    liftStm stm' =-      liftAllocationsInStms stms (stm' :<| lifted) acc (to_lift', consumed')-      where-        to_lift' =-          freeIn stm'-            <> (to_lift `namesSubtract` namesFromList (patNames (stmPat stm')))-        consumed' = consumed <> consumedInStm stm'-    dontLiftStm stm' =-      liftAllocationsInStms stms lifted (stm' :<| acc) (to_lift, consumed)+liftAllocationsInStms stms_orig = do+  outer_aliases <- asks envAliases+  let aliases = foldl trackAliases outer_aliases stms_orig+      go ::+        -- The input stms+        Stms rep ->+        -- The lifted allocations and associated statements+        Stms rep ->+        -- The other statements processed so far+        Stms rep ->+        -- (Names we need to lift, consumed names)+        (Names, Names) ->+        LiftM (inner rep) (Stms rep)+      go Empty lifted acc _ = pure $ lifted <> acc+      go (stms :|> stm) lifted acc (to_lift, consumed) = do+        stm' <- liftInsideStm stm+        case stmExp stm' of+          BasicOp Assert {} -> liftStm stm'+          Op Alloc {} -> liftStm stm'+          _ -> do+            let pat_names = namesFromList $ patNames $ stmPat stm'+                free_in_stm = freeIn stm+                expand v = maybe [v] namesToList $ M.lookup v $ fst aliases+            if (pat_names `namesIntersect` to_lift)+              || any (`nameIn` free_in_stm) (foldMap expand $ namesToList consumed)+              then liftStm stm'+              else dontLiftStm stm'+        where+          liftStm stm' =+            go stms (stm' :<| lifted) acc (to_lift', consumed')+            where+              to_lift' =+                freeIn stm'+                  <> (to_lift `namesSubtract` namesFromList (patNames (stmPat stm')))+              consumed' = consumed <> consumedInStm stm'+          dontLiftStm stm' =+            go stms lifted (stm' :<| acc) (to_lift, consumed)+  local (\env -> env {envAliases = aliases}) $+    go stms_orig mempty mempty mempty  liftAllocationsInSegOp ::   (Mem rep inner, Aliased rep) =>   SegOp lvl rep ->   LiftM (inner rep) (SegOp lvl rep) liftAllocationsInSegOp (SegMap lvl sp tps body) = do-  stms <- liftAllocationsInStms (kernelBodyStms body) mempty mempty mempty+  stms <- liftAllocationsInStms (kernelBodyStms body)   pure $ SegMap lvl sp tps $ body {kernelBodyStms = stms} liftAllocationsInSegOp (SegRed lvl sp tps body binops) = do-  stms <- liftAllocationsInStms (kernelBodyStms body) mempty mempty mempty+  stms <- liftAllocationsInStms (kernelBodyStms body)   pure $ SegRed lvl sp tps (body {kernelBodyStms = stms}) binops liftAllocationsInSegOp (SegScan lvl sp tps body binops) = do-  stms <- liftAllocationsInStms (kernelBodyStms body) mempty mempty mempty+  stms <- liftAllocationsInStms (kernelBodyStms body)   pure $ SegScan lvl sp tps (body {kernelBodyStms = stms}) binops liftAllocationsInSegOp (SegHist lvl sp tps body histops) = do-  stms <- liftAllocationsInStms (kernelBodyStms body) mempty mempty mempty+  stms <- liftAllocationsInStms (kernelBodyStms body)   pure $ SegHist lvl sp tps (body {kernelBodyStms = stms}) histops  liftAllocationsInHostOp ::

@@ -146,8 +146,8 @@   -- Finally, the array parameters are set to the arrays (but reshaped   -- to make the types work out; this will be simplified rapidly).   forM_ (zip arr_params arrs) $ \(p, arr) ->-    letBindNames [paramName p] . BasicOp $-      Reshape ReshapeCoerce (arrayShape $ paramType p) arr+    letBindNames [paramName p] $+      shapeCoerce (arrayDims $ paramType p) arr    -- Then we just inline the lambda body.   mapM_ addStm $ bodyStms $ lambdaBody lam@@ -159,7 +159,7 @@     certifying cs $ case (arrayDims $ patElemType pe, se) of       (dims, Var v)         | not $ null dims ->-            letBindNames [patElemName pe] $ BasicOp $ Reshape ReshapeCoerce (Shape dims) v+            letBindNames [patElemName pe] $ shapeCoerce dims v       _ -> letBindNames [patElemName pe] $ BasicOp $ SubExp se  -- | Split the parameters of a stream reduction lambda into the chunk

@@ -142,8 +142,10 @@   substituteNames _ = id  instance (Substitute d) => Substitute (ShapeBase d) where-  substituteNames substs (Shape es) =-    Shape $ map (substituteNames substs) es+  substituteNames substs = fmap (substituteNames substs)++instance (Substitute d) => Substitute (NewShape d) where+  substituteNames substs = fmap (substituteNames substs)  instance (Substitute d) => Substitute (Ext d) where   substituteNames substs (Free x) = Free $ substituteNames substs x

@@ -161,12 +161,7 @@   pure x  extEnv :: EvalM Env-extEnv = valEnv . M.map f <$> getExts-  where-    f v =-      ( Nothing,-        v-      )+extEnv = valEnv . M.map (Nothing,) <$> getExts  valueStructType :: ValueType -> StructType valueStructType = first $ flip sizeFromInteger mempty . toInteger@@ -181,7 +176,7 @@   pretty (SizeClosure _ e) = pretty e  instance Pretty (F.Shape SizeClosure) where-  pretty = mconcat . map (braces . pretty) . shapeDims+  pretty = mconcat . map (brackets . pretty) . shapeDims  -- | A type where the sizes are unevaluated expressions. type EvalType = TypeBase SizeClosure NoUniqueness@@ -337,8 +332,9 @@ lookupInEnv onEnv qv env = f env $ qualQuals qv   where     f m (q : qs) =-      case M.lookup q $ envTerm m of-        Just (TermModule (Module mod)) -> f mod qs+      case (M.lookup (qualLeaf qv) $ onEnv m, M.lookup q $ envTerm m) of+        (Just x, _) -> Just x+        (Nothing, Just (TermModule (Module mod))) -> f mod qs         _ -> Nothing     f m [] = M.lookup (qualLeaf qv) $ onEnv m @@ -646,7 +642,7 @@           "Index ["             <> T.intercalate ", " (map prettyText is)             <> "] out of bounds for array of shape "-            <> prettyText (valueShape arr)+            <> prettyText (arrayValueShape arr)             <> "."   maybe oob pure $ indexArray is arr @@ -705,7 +701,7 @@       let free = fvVars $ freeInExp e        in not $ any (`S.member` bound) free || any (`S.member` outer_bound) free --- | Evaluate all sizes, and it better work. This implies it must be a+-- | Evaluate all sizes, and it better work. This implies it must not be a -- size-dependent function type, or one that has existentials. evalTypeFully :: EvalType -> EvalM ValueType evalTypeFully t = do@@ -755,7 +751,7 @@         env'' <- linkMissingSizes missing_sizes p v <$> matchPat env' p v         etaExpand (v : vs) env'' rt     etaExpand vs env' _ = do-      f <- localExts $ eval env' body+      f <- eval env' body       foldM (apply noLoc mempty) f $ reverse vs evalFunction env missing_sizes (p : ps) body rettype =   pure . ValueFun $ \v -> do@@ -1207,7 +1203,14 @@       let res_env = case res_mod of             Module x -> x             _ -> mempty-      pure (f_env <> e_env <> res_env, res_mod)+      -- The following environment handles the case where rsubst refers to names+      -- that are not actually defined in the module itself, but merely+      -- inherited from an outer environment (see #2273).+      let env_substs = (`Env` mempty) $ M.fromList $ do+            (to, from) <- M.toList rsubst+            x <- maybeToList $ M.lookup from $ envTerm env+            pure (to, x)+      pure (f_env <> e_env <> res_env <> env_substs, res_mod)     _ -> error "Expected ModuleFun."  evalDec :: Env -> Dec -> EvalM Env

@@ -12,6 +12,7 @@     -- * Values     Value (..),     valueShape,+    arrayValueShape,     prettyValue,     valueText,     valueAccum,@@ -194,6 +195,15 @@ valueShape (ValueRecord fs) = ShapeRecord $ M.map valueShape fs valueShape (ValueSum shape _ _) = shape valueShape _ = ShapeLeaf++-- | Retrieve the part of the value shape that corresponds to outer array+-- dimensions. This is used for reporting shapes in those cases where the full+-- shape is not important, namely in indexing errors.+arrayValueShape :: Value m -> ValueShape+arrayValueShape = outer . valueShape+  where+    outer (ShapeDim d s) = ShapeDim d $ outer s+    outer _ = ShapeLeaf  -- TODO: Perhaps there is some clever way to reuse the code between -- valueAccum and valueAccumLM

@@ -239,7 +239,7 @@         | ModTypeExp '->' ModTypeExp  { ModTypeArrow Nothing $1 $3 (srcspan $1 $>) }  TypeRef :: { TypeRefBase NoInfo Name }-         : QualName TypeParams '=' TypeExpTerm+         : QualName TypeParams '=' TypeExp            { TypeRef (fst $1) $2 $4 (srcspan (snd $1) $>) }  ModTypeBind :: { ModTypeBindBase NoInfo Name }@@ -301,6 +301,8 @@           in ValSpec name $3 $5 NoInfo Nothing (srcspan $1 $>) }       | val BindingBinOp TypeParams ':' TypeExp         { ValSpec $2 $3 $5 NoInfo Nothing (srcspan $1 $>) }+      | val '(' BindingBinOp ')' TypeParams ':' TypeExp+        { ValSpec $3 $5 $7 NoInfo Nothing (srcspan $1 $>) }       | TypeAbbr         { TypeAbbrSpec $1 } 

@@ -6,6 +6,7 @@   ( prettyString,     prettyTuple,     leadingOperator,+    symbolName,     IsName (..),     prettyNameString,     Annot (..),@@ -37,7 +38,7 @@ -- with the tag.  To avoid erroneously using the 'Pretty' instance for -- VNames, we in fact only define it inside the modules for the core -- language (as an orphan instance).-class IsName v where+class (Eq v) => IsName v where   prettyName :: v -> Doc a   toName :: v -> Name @@ -204,7 +205,7 @@ hasArrayLit (TupLit es2 _) = any hasArrayLit es2 hasArrayLit _ = False -instance (Eq vn, IsName vn, Annot f) => Pretty (DimIndexBase f vn) where+instance (IsName vn, Annot f) => Pretty (DimIndexBase f vn) where   pretty (DimFix e) = pretty e   pretty (DimSlice i j (Just s)) =     maybe mempty pretty i@@ -223,12 +224,12 @@ instance (IsName vn) => Pretty (SizeBinder vn) where   pretty (SizeBinder v _) = brackets $ prettyName v -letBody :: (Eq vn, IsName vn, Annot f) => ExpBase f vn -> Doc a+letBody :: (IsName vn, Annot f) => ExpBase f vn -> Doc a letBody body@(AppExp LetPat {} _) = pretty body letBody body@(AppExp LetFun {} _) = pretty body letBody body = "in" <+> align (pretty body) -prettyAppExp :: (Eq vn, IsName vn, Annot f) => Int -> AppExpBase f vn -> Doc a+prettyAppExp :: (IsName vn, Annot f) => Int -> AppExpBase f vn -> Doc a prettyAppExp p (BinOp (bop, _) _ (x, _) (y, _) _) = prettyBinOp p bop x y prettyAppExp _ (Match e cs _) = "match" <+> pretty e </> (stack . map pretty) (NE.toList cs) prettyAppExp _ (Loop sizeparams pat initexp form loopbody _) =@@ -308,7 +309,7 @@     prettyExp 0 f       <+> hsep (map (prettyExp 10 . snd) $ NE.toList args) -instance (Eq vn, IsName vn, Annot f) => Pretty (AppExpBase f vn) where+instance (IsName vn, Annot f) => Pretty (AppExpBase f vn) where   pretty = prettyAppExp (-1)  prettyInst :: (Annot f, Pretty t) => f t -> Doc a@@ -322,17 +323,18 @@ prettyAttr :: (Pretty a) => a -> Doc ann prettyAttr attr = "#[" <> pretty attr <> "]" -operatorName :: Name -> Bool-operatorName = (`elem` opchars) . T.head . nameToText+-- | Does this name correspond to a symbol rather than an identifier?+symbolName :: Name -> Bool+symbolName = (`elem` opchars) . T.head . nameToText   where     opchars :: String     opchars = "+-*/%=!><|&^." -prettyExp :: (Eq vn, IsName vn, Annot f) => Int -> ExpBase f vn -> Doc a+prettyExp :: (IsName vn, Annot f) => Int -> ExpBase f vn -> Doc a prettyExp _ (Var name t _)   -- The first case occurs only for programs that have been normalised   -- by the compiler.-  | operatorName (toName (qualLeaf name)) = parens $ pretty name <> prettyInst t+  | symbolName (toName (qualLeaf name)) = parens $ pretty name <> prettyInst t   | otherwise = pretty name <> prettyInst t prettyExp _ (Hole t _) = "???" <> prettyInst t prettyExp _ (Parens e _) = align $ parens $ pretty e@@ -359,7 +361,7 @@   brackets (commasep $ map pretty es) <> prettyInst t prettyExp _ (StringLit s _) =   pretty $ show $ map (chr . fromIntegral) s-prettyExp _ (Project k e _ _) = pretty e <> "." <> pretty k+prettyExp _ (Project k e _ _) = prettyExp 11 e <> "." <> pretty k prettyExp _ (Negate e _) = "-" <> pretty e prettyExp _ (Not e _) = "!" <> pretty e prettyExp _ (Update src idxs ve _) =@@ -409,7 +411,7 @@         <> parens (pretty t <> "," <+> brackets (commasep $ map prettyName ext))   | otherwise = prettyAppExp i e -instance (Eq vn, IsName vn, Annot f) => Pretty (ExpBase f vn) where+instance (IsName vn, Annot f) => Pretty (ExpBase f vn) where   pretty = prettyExp (-1)  instance (IsName vn) => Pretty (AttrAtom vn) where@@ -420,18 +422,18 @@   pretty (AttrAtom attr _) = pretty attr   pretty (AttrComp f attrs _) = pretty f <> parens (commasep $ map pretty attrs) -instance (Eq vn, IsName vn, Annot f) => Pretty (FieldBase f vn) where+instance (IsName vn, Annot f) => Pretty (FieldBase f vn) where   pretty (RecordFieldExplicit (L _ name) e _) = pretty name <> equals <> pretty e   pretty (RecordFieldImplicit (L _ name) _ _) = prettyName name -instance (Eq vn, IsName vn, Annot f) => Pretty (CaseBase f vn) where+instance (IsName vn, Annot f) => Pretty (CaseBase f vn) where   pretty (CasePat p e _) = "case" <+> pretty p <+> "->" </> indent 2 (pretty e) -instance (Eq vn, IsName vn, Annot f) => Pretty (LoopInitBase f vn) where+instance (IsName vn, Annot f) => Pretty (LoopInitBase f vn) where   pretty (LoopInitImplicit e) = maybe "_" pretty $ unAnnot e   pretty (LoopInitExplicit e) = pretty e -instance (Eq vn, IsName vn, Annot f) => Pretty (LoopFormBase f vn) where+instance (IsName vn, Annot f) => Pretty (LoopFormBase f vn) where   pretty (For i ubound) =     "for" <+> pretty i <+> "<" <+> align (pretty ubound)   pretty (ForIn x e) =@@ -444,7 +446,7 @@   pretty (PatLitFloat f) = pretty f   pretty (PatLitPrim v) = pretty v -instance (Eq vn, IsName vn, Annot f, Pretty t) => Pretty (PatBase f vn t) where+instance (IsName vn, Annot f, Pretty t) => Pretty (PatBase f vn t) where   pretty (PatAscription p t _) = pretty p <> colon <+> align (pretty t)   pretty (PatParens p _) = parens $ pretty p   pretty (Id v t _) = case unAnnot t of@@ -465,10 +467,10 @@ ppAscription Nothing = mempty ppAscription (Just t) = colon <> align (pretty t) -instance (Eq vn, IsName vn, Annot f) => Pretty (ProgBase f vn) where+instance (IsName vn, Annot f) => Pretty (ProgBase f vn) where   pretty = stack . punctuate line . map pretty . progDecs -instance (Eq vn, IsName vn, Annot f) => Pretty (DecBase f vn) where+instance (IsName vn, Annot f) => Pretty (DecBase f vn) where   pretty (ValDec dec) = pretty dec   pretty (TypeDec dec) = pretty dec   pretty (ModTypeDec sig) = pretty sig@@ -477,7 +479,7 @@   pretty (LocalDec dec _) = "local" <+> pretty dec   pretty (ImportDec x _ _) = "import" <+> pretty x -prettyModExp :: (Eq vn, IsName vn, Annot f) => Int -> ModExpBase f vn -> Doc a+prettyModExp :: (IsName vn, Annot f) => Int -> ModExpBase f vn -> Doc a prettyModExp _ (ModVar v _) =   pretty v prettyModExp _ (ModParens e _) =@@ -502,7 +504,7 @@       Nothing -> mempty       Just (sig, _) -> colon <+> pretty sig -instance (Eq vn, IsName vn, Annot f) => Pretty (ModExpBase f vn) where+instance (IsName vn, Annot f) => Pretty (ModExpBase f vn) where   pretty = prettyModExp (-1)  instance Pretty Liftedness where@@ -510,7 +512,7 @@   pretty SizeLifted = "~"   pretty Lifted = "^" -instance (Eq vn, IsName vn, Annot f) => Pretty (TypeBindBase f vn) where+instance (IsName vn, Annot f) => Pretty (TypeBindBase f vn) where   pretty (TypeBind name l params te rt _ _) =     "type"       <> pretty l@@ -518,11 +520,11 @@         <+> equals         <+> maybe (pretty te) pretty (unAnnot rt) -instance (Eq vn, IsName vn) => Pretty (TypeParamBase vn) where+instance (IsName vn) => Pretty (TypeParamBase vn) where   pretty (TypeParamDim name _) = brackets $ prettyName name   pretty (TypeParamType l name _) = "'" <> pretty l <> prettyName name -instance (Eq vn, IsName vn, Annot f) => Pretty (ValBindBase f vn) where+instance (IsName vn, Annot f) => Pretty (ValBindBase f vn) where   pretty (ValBind entry name retdecl rettype tparams args body _ attrs _) =     mconcat (map ((<> line) . prettyAttr) attrs)       <> fun@@ -544,18 +546,23 @@         Just rettype' -> [colon <+> align rettype']         Nothing -> mempty -instance (Eq vn, IsName vn, Annot f) => Pretty (SpecBase f vn) where+instance (IsName vn, Annot f) => Pretty (SpecBase f vn) where   pretty (TypeAbbrSpec tpsig) = pretty tpsig   pretty (TypeSpec l name ps _ _) =     "type" <> pretty l <+> hsep (prettyName name : map pretty ps)   pretty (ValSpec name tparams vtype _ _ _) =-    "val" <+> hsep (prettyName name : map pretty tparams) <> colon <+> pretty vtype+    "val" <+> hsep (name' : map pretty tparams) <> colon <+> pretty vtype+    where+      name' =+        if symbolName $ toName name+          then parens $ prettyName name+          else prettyName name   pretty (ModSpec name sig _ _) =     "module" <+> prettyName name <> colon <+> pretty sig   pretty (IncludeSpec e _) =     "include" <+> pretty e -instance (Eq vn, IsName vn, Annot f) => Pretty (ModTypeExpBase f vn) where+instance (IsName vn, Annot f) => Pretty (ModTypeExpBase f vn) where   pretty (ModTypeVar v _ _) = pretty v   pretty (ModTypeParens e _) = parens $ pretty e   pretty (ModTypeSpecs ss _) = nestedBlock "{" "}" (stack $ punctuate line $ map pretty ss)@@ -566,15 +573,15 @@   pretty (ModTypeArrow Nothing e1 e2 _) =     pretty e1 <+> "->" <+> pretty e2 -instance (Eq vn, IsName vn, Annot f) => Pretty (ModTypeBindBase f vn) where+instance (IsName vn, Annot f) => Pretty (ModTypeBindBase f vn) where   pretty (ModTypeBind name e _ _) =     "module type" <+> prettyName name <+> equals <+> pretty e -instance (Eq vn, IsName vn, Annot f) => Pretty (ModParamBase f vn) where+instance (IsName vn, Annot f) => Pretty (ModParamBase f vn) where   pretty (ModParam pname psig _ _) =     parens (prettyName pname <> colon <+> pretty psig) -instance (Eq vn, IsName vn, Annot f) => Pretty (ModBindBase f vn) where+instance (IsName vn, Annot f) => Pretty (ModBindBase f vn) where   pretty (ModBind name ps sig e _ _) =     "module" <+> hsep (prettyName name : map pretty ps) <> sig' <> " =" <+> pretty e     where@@ -591,7 +598,7 @@     leading = leadingOperator $ toName $ qualLeaf bop  prettyBinOp ::-  (Eq vn, IsName vn, Annot f) =>+  (IsName vn, Annot f) =>   Int ->   QualName vn ->   ExpBase f vn ->

@@ -28,7 +28,7 @@ import Language.Futhark import System.FilePath qualified as Native import System.FilePath.Posix qualified as Posix-import Prelude hiding (mod)+import Prelude hiding (abs, mod)  -- | Create an import name immediately from a file path specified by -- the user.@@ -158,7 +158,10 @@   mempty = Env mempty mempty mempty mempty mempty  instance Pretty MTy where-  pretty = pretty . mtyMod+  pretty (MTy abs mod) =+    "abstract" <> parens (hsep $ map p $ M.toList abs) </> pretty mod+    where+      p (v, l) = pretty l <> pretty v  instance Pretty Mod where   pretty (ModEnv e) = pretty e@@ -190,7 +193,7 @@         "val"           <+> prettyName name           <> mconcat (map ((" " <>) . pretty) tps)-          <> " ="+            <+> ":"             <+> pretty t       renderModType (name, _sig) =         "module type" <+> prettyName name

@@ -331,25 +331,18 @@   (abs, s_abs, s_env, s') <- checkModTypeExpToEnv s   resolveTypeParams ps $ \ps' -> do     (ext, te', te_t, _) <- bindingTypeParams ps' $ checkTypeDecl te-    unless (null ext) $-      typeError te' mempty "Anonymous dimensions are not allowed here."-    (tname', s_abs', s_env') <- refineEnv loc s_abs s_env tname ps' te_t+    (tname', s_abs', s_env') <-+      refineEnv loc s_abs s_env tname ps' $ RetType ext te_t     pure (abs, MTy s_abs' $ ModEnv s_env', ModTypeWith s' (TypeRef tname' ps' te' trloc) loc) checkModTypeExp (ModTypeArrow maybe_pname e1 e2 loc) = do   (e1_abs, MTy s_abs e1_mod, e1') <- checkModTypeExp e1-  (env_for_e2, maybe_pname') <-+  (maybe_pname', (e2_abs, e2_mod, e2')) <-     case maybe_pname of       Just pname -> bindSpaced1 Term pname loc $ \pname' ->-        pure-          ( mempty-              { envNameMap = M.singleton (Term, pname) $ qualName pname',-                envModTable = M.singleton pname' e1_mod-              },-            Just pname'-          )+        localEnv (mempty {envModTable = M.singleton pname' e1_mod}) $+          (Just pname',) <$> checkModTypeExp e2       Nothing ->-        pure (mempty, Nothing)-  (e2_abs, e2_mod, e2') <- localEnv env_for_e2 $ checkModTypeExp e2+        (Nothing,) <$> checkModTypeExp e2   pure     ( e1_abs <> e2_abs,       MTy mempty $ ModFun $ FunModType s_abs e1_mod e2_mod,@@ -644,10 +637,9 @@   SrcLoc ->   [TypeParam] ->   [Pat ParamType] ->-  Maybe (TypeExp Exp VName) ->   ResRetType ->   TypeM ()-checkEntryPoint loc tparams params maybe_tdecl rettype+checkEntryPoint loc tparams params rettype   | any isTypeParam tparams =       typeError loc mempty $         withIndexLink@@ -674,16 +666,6 @@         "Entry point size parameter "           <> pretty p           <> " only used non-constructively."-  | p : _ <- filter nastyParameter params =-      warn p $-        "Entry point parameter\n"-          </> indent 2 (pretty p)-          </> "\nwill have an opaque type, so the entry point will likely not be callable."-  | nastyReturnType maybe_tdecl rettype_t =-      warn loc $-        "Entry point return type\n"-          </> indent 2 (pretty rettype)-          </> "\nwill have an opaque type, so the result will likely not be usable."   | otherwise =       pure ()   where@@ -708,7 +690,7 @@    let entry' = Info (entryPoint params' maybe_tdecl' rettype) <$ entry   case entry' of-    Just _ -> checkEntryPoint loc tparams' params' maybe_tdecl' rettype+    Just _ -> checkEntryPoint loc tparams' params' rettype     _ -> pure ()    let vb' = ValBind entry' fname maybe_tdecl' (Info rettype) tparams' params' body' doc attrs' loc@@ -721,44 +703,6 @@         },       vb'     )--nastyType :: (Monoid als) => TypeBase dim als -> Bool-nastyType (Scalar Prim {}) = False-nastyType t@Array {} = nastyType $ stripArray 1 t-nastyType _ = True--nastyReturnType :: (Monoid als) => Maybe (TypeExp Exp VName) -> TypeBase dim als -> Bool-nastyReturnType Nothing (Scalar (Arrow _ _ _ t1 (RetType _ t2))) =-  nastyType t1 || nastyReturnType Nothing t2-nastyReturnType (Just (TEArrow _ te1 te2 _)) (Scalar (Arrow _ _ _ t1 (RetType _ t2))) =-  (not (niceTypeExp te1) && nastyType t1)-    || nastyReturnType (Just te2) t2-nastyReturnType (Just te) _-  | niceTypeExp te = False-nastyReturnType te t-  | Just ts <- isTupleRecord t =-      case te of-        Just (TETuple tes _) -> or $ zipWith nastyType' (map Just tes) ts-        _ -> any nastyType ts-  | otherwise = nastyType' te t-  where-    nastyType' (Just te') _ | niceTypeExp te' = False-    nastyType' _ t' = nastyType t'--nastyParameter :: Pat ParamType -> Bool-nastyParameter p = nastyType (patternType p) && not (ascripted p)-  where-    ascripted (PatAscription _ te _) = niceTypeExp te-    ascripted (PatParens p' _) = ascripted p'-    ascripted _ = False--niceTypeExp :: TypeExp Exp VName -> Bool-niceTypeExp (TEVar (QualName [] _) _) = True-niceTypeExp (TEApply te TypeArgExpSize {} _) = niceTypeExp te-niceTypeExp (TEArray _ te _) = niceTypeExp te-niceTypeExp (TEUnique te _) = niceTypeExp te-niceTypeExp (TEDim _ te _) = niceTypeExp te-niceTypeExp _ = False  checkOneDec :: DecBase NoInfo Name -> TypeM (TySet, Env, DecBase Info VName) checkOneDec (ModDec struct) = do

@@ -9,6 +9,7 @@ where  import Control.Monad+import Control.Monad.Identity import Data.Either import Data.Map.Strict qualified as M import Data.Maybe@@ -17,6 +18,7 @@ import Futhark.Util.Pretty import Language.Futhark import Language.Futhark.Semantic+import Language.Futhark.Traversals import Language.Futhark.TypeChecker.Monad import Language.Futhark.TypeChecker.Types import Language.Futhark.TypeChecker.Unify (doUnification)@@ -109,10 +111,21 @@         substitute v =           fromMaybe v $ M.lookup v substs -        -- For applySubst and friends.-        subst v =-          ExpSubst . flip sizeFromName mempty . qualName <$> M.lookup v substs+        substituteInExp :: Exp -> Exp+        substituteInExp = runIdentity . astMap mapper+          where+            mapper =+              ASTMapper+                { mapOnExp = pure . substituteInExp,+                  mapOnName = pure . substituteInQualName,+                  mapOnStructType = pure . substituteInType,+                  mapOnParamType = pure . substituteInType,+                  mapOnResRetType = pure . substituteInRetType+                } +        substituteInQualName (QualName qs v) =+          QualName (map substitute qs) (substitute v)+         substituteInMap f m =           let (ks, vs) = unzip $ M.toList m            in M.fromList $@@ -143,9 +156,8 @@           TypeParamType l (substitute p) loc          substituteInScalarType :: ScalarTypeBase Size u -> ScalarTypeBase Size u-        substituteInScalarType (TypeVar u (QualName qs v) targs) =-          TypeVar u (QualName (map substitute qs) $ substitute v) $-            map substituteInTypeArg targs+        substituteInScalarType (TypeVar u v targs) =+          TypeVar u (substituteInQualName v) $ map substituteInTypeArg targs         substituteInScalarType (Prim t) =           Prim t         substituteInScalarType (Record ts) =@@ -155,15 +167,18 @@         substituteInScalarType (Arrow als v d1 t1 (RetType dims t2)) =           Arrow als v d1 (substituteInType t1) $ RetType dims $ substituteInType t2 +        substituteInRetType :: RetTypeBase Size u -> RetTypeBase Size u+        substituteInRetType (RetType ext t) = RetType ext $ substituteInType t+         substituteInType :: TypeBase Size u -> TypeBase Size u         substituteInType (Scalar t) = Scalar $ substituteInScalarType t         substituteInType (Array u shape t) =           Array u (substituteInShape shape) $ substituteInScalarType t -        substituteInShape (Shape ds) = Shape $ map (applySubst subst) ds+        substituteInShape (Shape ds) = Shape $ map substituteInExp ds          substituteInTypeArg (TypeArgDim e) =-          TypeArgDim $ applySubst subst e+          TypeArgDim $ substituteInExp e         substituteInTypeArg (TypeArgType t) =           TypeArgType $ substituteInType t @@ -182,13 +197,19 @@     <> (mconcat . map modTypeAbbrs . M.elems . envModTable) env  -- | Refine the given type name in the given env.+--+-- XXX: we do not check whether this results in a meaningful module type. In+-- particular, we may refine a nonlifted type to contain a function or+-- existentially quantified sizes. However, it is still not possible to+-- construct a module that matches such malformed module types, so this is not a+-- soundness issue, merely an ergonomic issue. refineEnv ::   SrcLoc ->   TySet ->   Env ->   QualName Name ->   [TypeParam] ->-  StructType ->+  StructRetType ->   TypeM (QualName VName, TySet, Env) refineEnv loc tset env tname ps t   | Just (tname', TypeAbbr _ cur_ps (RetType _ (Scalar (TypeVar _ (QualName qs v) _)))) <-@@ -202,8 +223,8 @@               substituteTypesInEnv                 ( flip M.lookup $                     M.fromList-                      [ (qualLeaf tname', Subst cur_ps $ RetType [] t),-                        (v, Subst ps $ RetType [] t)+                      [ (qualLeaf tname', Subst cur_ps t),+                        (v, Subst ps t)                       ]                 )                 env

@@ -958,10 +958,11 @@ checkOneExp :: ExpBase NoInfo VName -> TypeM ([TypeParam], Exp) checkOneExp e = runTermTypeM checkExp $ do   e' <- checkExp e-  let t = typeOf e'-  (tparams, _, _) <--    letGeneralise (nameFromString "<exp>") (srclocOf e) [] [] $ toRes Nonunique t-  fixOverloadedTypes $ typeVars t+  (tparams, _, RetType _ t') <-+    letGeneralise (nameFromString "<exp>") (srclocOf e) [] [] $+      toRes Nonunique $+        typeOf e'+  fixOverloadedTypes $ typeVars t'   e'' <- normTypeFully e'   localChecks e''   causalityCheck e''@@ -1606,10 +1607,18 @@           loc           (filter (`elem` hidden) $ foldMap patNames params)           body_t--      let usage = mkUsage body "return type annotation"-      onFailure (CheckingReturn rettype body_t') $-        unify usage (toStruct rettype) body_t'+      case find (`elem` hidden) $ fvVars $ freeInType rettype of+        Just v ->+          typeError loc mempty $+            "The return type annotation"+              </> indent 2 (align (pretty rettype))+              </> "refers to the name"+              <+> dquotes (prettyName v)+              <+> "which is bound to an inner component of a function parameter."+        Nothing -> do+          let usage = mkUsage body "return type annotation"+          onFailure (CheckingReturn rettype body_t') $+            unify usage (toStruct rettype) body_t'     Nothing -> pure ()    pure body'

@@ -258,13 +258,16 @@           <+> pretty (length targs)           <> "."     else do-      (targs', dims, substs) <- unzip3 <$> zipWithM checkArgApply ps targs+      (targs', dims, substs, targs_ls) <-+        L.unzip4 <$> zipWithM checkArgApply ps targs       pure         ( foldl (\x y -> TEApply x y tloc) (TEVar tname tname_loc) targs',           [],           RetType (t_dims ++ mconcat dims) $             applySubst (`M.lookup` mconcat substs) t,-          l+          -- XXX: this is an overapproximation of the liftedness in case one of+          -- these type parameters is a phantom type.+          maximum $ l : targs_ls         )   where     tloc = srclocOf ote@@ -294,13 +297,14 @@      checkArgApply (TypeParamDim pv _) (TypeArgExpSize d) = do       (d', svars, subst) <- checkSizeExp d-      pure (d', svars, M.singleton pv subst)+      pure (d', svars, M.singleton pv subst, Unlifted)     checkArgApply (TypeParamType _ pv _) (TypeArgExpType te) = do-      (te', svars, RetType dims st, _) <- evalTypeExp df te+      (te', svars, RetType dims st, te_l) <- evalTypeExp df te       pure         ( TypeArgExpType te',           svars ++ dims,-          M.singleton pv $ Subst [] $ RetType [] $ toStruct st+          M.singleton pv $ Subst [] $ RetType [] $ toStruct st,+          te_l         )     checkArgApply p a =       typeError tloc mempty $

@@ -557,7 +557,7 @@     not (anyBound bound e1) || (qualLeaf v2 `elem` bound) =       linkVarToDim usage bcs (qualLeaf v2) lvl2 e1 unifySizes usage bcs _ _ e1 e2 = do-  notes <- (<>) <$> dimNotes usage e2 <*> dimNotes usage e2+  notes <- (<>) <$> dimNotes usage e1 <*> dimNotes usage e2   unifyError usage notes bcs $     "Sizes"       <+> dquotes (pretty e1)

@@ -5,12 +5,17 @@   ) where +import Data.List qualified as L import Futhark.IR.Prop.Constants import Futhark.IR.Prop.Reshape import Futhark.IR.Syntax+import Futhark.IR.SyntaxTests () import Test.Tasty import Test.Tasty.HUnit +intShape :: [Int] -> Shape+intShape = Shape . map (intConst Int32 . toInteger)+ reshapeOuterTests :: [TestTree] reshapeOuterTests =   [ testCase (unwords ["reshapeOuter", show sc, show n, show shape, "==", show sc_res]) $@@ -35,10 +40,197 @@         ]   ] -intShape :: [Int] -> Shape-intShape = Shape . map (intConst Int32 . toInteger)+dimFlatten :: Int -> Int -> d -> DimSplice d+dimFlatten i k w = DimSplice i k (Shape [w]) +dimUnflatten :: Int -> [d] -> DimSplice d+dimUnflatten i ws = DimSplice i 1 (Shape ws)++dimCoerce :: Int -> d -> DimSplice d+dimCoerce i w = DimSplice i 1 (Shape [w])++dimSplice :: Int -> Int -> [d] -> DimSplice d+dimSplice i n s = DimSplice i n $ Shape s++flipReshapeRearrangeTests :: [TestTree]+flipReshapeRearrangeTests =+  [ testCase+      ( unwords+          [ "flipReshapeRearrange",+            show v0_shape,+            show v1_shape,+            show perm+          ]+      )+      $ flipReshapeRearrange v0_shape v1_shape perm @?= res+    | (v0_shape :: [String], v1_shape, perm, res) <-+        [ ( ["A", "B", "C"],+            ["A", "BC"],+            [1, 0],+            Just [1, 2, 0]+          ),+          ( ["A", "B", "C", "D"],+            ["A", "BCD"],+            [1, 0],+            Just [1, 2, 3, 0]+          ),+          ( ["A"],+            ["B", "C"],+            [1, 0],+            Nothing+          ),+          ( ["A", "B", "C"],+            ["AB", "C"],+            [1, 0],+            Just [2, 0, 1]+          ),+          ( ["A", "B", "C", "D"],+            ["ABC", "D"],+            [1, 0],+            Just [3, 0, 1, 2]+          )+        ]+  ]++flipRearrangeReshapeTests :: [TestTree]+flipRearrangeReshapeTests =+  [ testCase+      ( unwords+          [ "flipRearrangeReshape",+            show perm,+            prettyStringOneLine newshape+          ]+      )+      $ flipRearrangeReshape perm newshape @?= res+    | (perm, newshape :: NewShape String, res) <-+        [ ( [1, 0],+            NewShape+              [dimUnflatten 1 ["B", "C"]]+              (Shape ["A", "B", "C"]),+            Just+              ( NewShape+                  [dimUnflatten 0 ["B", "C"]]+                  (Shape ["B", "C", "A"]),+                [2, 0, 1]+              )+          ),+          ( [1, 0],+            NewShape+              [dimFlatten 0 2 "AB"]+              (Shape ["AB"]),+            Nothing+          )+        ]+  ]++simplifyTests :: TestTree+simplifyTests =+  testGroup+    "simplifyNewShape"+    [ testCase "Inverse flatten and unflatten - simple case" $+        lhs+          ["A", "B"]+          [dimFlatten 0 2 "AB", dimUnflatten 0 ["A", "B"]]+          @?= Just [],+      testCase "Non-inverse flatten and unflatten - simple case" $+        lhs+          ["A", "B"]+          [dimFlatten 0 2 "AB", dimUnflatten 0 ["C", "D"]]+          @?= Just [dimSplice 0 2 ["C", "D"]],+      testCase "Inverse flatten and unflatten - separated by coercion" $+        lhs+          ["A", "B"]+          [ dimFlatten 0 2 "AB",+            dimCoerce 0 "CD",+            dimUnflatten 0 ["C", "D"]+          ]+          @?= Just [dimSplice 0 2 ["C", "D"]],+      testCase "Two unflattens - simple case" $+        lhs+          ["ABC"]+          [dimUnflatten 0 ["A", "BC"], dimUnflatten 1 ["B", "C"]]+          @?= Just [dimUnflatten 0 ["A", "B", "C"]],+      testCase "Two unflattens with unchanged prefix" $+        lhs+          ["A", "B", "C", "D", "E"]+          [ DimSplice 3 2 $ Shape ["DE"],+            DimSplice 2 2 $ Shape ["CDE"]+          ]+          @?= Just [dimFlatten 2 3 "CDE"],+      testCase "Identity coerce" $+        lhs+          ["A", "B", "C"]+          [dimCoerce 1 "B", dimCoerce 2 "C"]+          @?= Just [],+      testCase "Identity coerce (multiple dimensions)" $+        lhs+          ["A", "B", "C"]+          [DimSplice 1 2 (Shape ["B", "C"])]+          @?= Just [],+      testCase "Identity coerce (with non-identity stuff afterwards)" $+        lhs+          ["B", "CD"]+          [dimCoerce 0 "B", dimUnflatten 1 ["C", "D"]]+          @?= Just [dimUnflatten 1 ["C", "D"]],+      testCase "Get rid of a coerce before an unflatten" $+        lhs+          ["CD"]+          [dimCoerce 0 "AB", dimUnflatten 0 ["A", "B"]]+          @?= Just [dimUnflatten 0 ["A", "B"]],+      testCase "Get rid of a coerce after a flatten" $+        lhs+          ["A", "B", "C"]+          [dimFlatten 0 2 "ABC", dimCoerce 0 "K"]+          @?= Just [dimFlatten 0 2 "K"],+      testCase "Flatten and unflatten (invariant suffix)" $+        lhs+          ["A", "B", "C"]+          [dimFlatten 0 3 "ABC", dimUnflatten 0 ["D", "E", "C"]]+          @?= Just [dimSplice 0 2 ["D", "E"]],+      testCase "Flatten and unflatten (invariant prefix)" $+        lhs+          ["A", "B", "C"]+          [dimFlatten 0 3 "ABC", dimUnflatten 0 ["A", "D", "E"]]+          @?= Just [dimSplice 1 2 ["D", "E"]],+      testCase "Invariant part of splice" $+        lhs+          ["A", "B", "C", "D"]+          [DimSplice 1 3 $ Shape ["BC", "D"]]+          @?= Just [DimSplice 1 2 $ Shape ["BC"]],+      testCase "Necessary coercion" $+        lhs+          ["A", "B"]+          [dimCoerce 0 "C", dimCoerce 1 "D"]+          @?= Nothing,+      testCase "Another necessary coercion" $+        lhs+          ["A", "B", "C"]+          [dimCoerce 0 "A'", dimCoerce 1 "A'", dimCoerce 2 "A'"]+          @?= Nothing,+      testCase "Long with redundancies" $+        lhs+          ["A", "B", "C", "D"]+          [ DimSplice 1 3 $ Shape ["BC", "D"],+            dimCoerce 1 "BC",+            dimCoerce 2 "D",+            dimFlatten 1 2 "BCD",+            dimFlatten 0 2 "ABCD"+          ]+          @?= Just [dimFlatten 0 4 "ABCD"]+    ]+  where+    lhs orig_shape ss =+      let res_shape :: ShapeBase String =+            L.foldl' applySplice (Shape orig_shape) ss+       in dimSplices+            <$> simplifyNewShape (Shape orig_shape) (NewShape ss res_shape)+ tests :: TestTree tests =-  testGroup "ReshapeTests" $-    reshapeOuterTests ++ reshapeInnerTests+  testGroup "ReshapeTests" . mconcat $+    [ reshapeOuterTests,+      reshapeInnerTests,+      flipReshapeRearrangeTests,+      flipRearrangeReshapeTests,+      [simplifyTests]+    ]