diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 diff --git a/docs/language-reference.rst b/docs/language-reference.rst --- a/docs/language-reference.rst +++ b/docs/language-reference.rst @@ -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`` .............. diff --git a/futhark.cabal b/futhark.cabal --- a/futhark.cabal +++ b/futhark.cabal @@ -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 diff --git a/prelude/math.fut b/prelude/math.fut --- a/prelude/math.fut +++ b/prelude/math.fut @@ -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 diff --git a/prelude/soacs.fut b/prelude/soacs.fut --- a/prelude/soacs.fut +++ b/prelude/soacs.fut @@ -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 diff --git a/rts/c/atomics.h b/rts/c/atomics.h deleted file mode 100644 --- a/rts/c/atomics.h +++ /dev/null @@ -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 diff --git a/rts/c/atomics16.h b/rts/c/atomics16.h new file mode 100644 --- /dev/null +++ b/rts/c/atomics16.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; +} + +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; +} + +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; +} + +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; +} + +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; +} + +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; +} + +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 diff --git a/rts/c/atomics32.h b/rts/c/atomics32.h new file mode 100644 --- /dev/null +++ b/rts/c/atomics32.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 diff --git a/rts/c/atomics64.h b/rts/c/atomics64.h new file mode 100644 --- /dev/null +++ b/rts/c/atomics64.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 diff --git a/rts/c/atomics8.h b/rts/c/atomics8.h new file mode 100644 --- /dev/null +++ b/rts/c/atomics8.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; +} + +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; +} + +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; +} + +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; +} + +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; +} + +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; +} + +// End of atomics8.h diff --git a/rts/c/gpu.h b/rts/c/gpu.h --- a/rts/c/gpu.h +++ b/rts/c/gpu.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? diff --git a/rts/cuda/prelude.cu b/rts/cuda/prelude.cu --- a/rts/cuda/prelude.cu +++ b/rts/cuda/prelude.cu @@ -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 diff --git a/rts/python/opencl.py b/rts/python/opencl.py --- a/rts/python/opencl.py +++ b/rts/python/opencl.py @@ -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) diff --git a/src/Futhark/AD/Fwd.hs b/src/Futhark/AD/Fwd.hs --- a/src/Futhark/AD/Fwd.hs +++ b/src/Futhark/AD/Fwd.hs @@ -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) diff --git a/src/Futhark/AD/Rev.hs b/src/Futhark/AD/Rev.hs --- a/src/Futhark/AD/Rev.hs +++ b/src/Futhark/AD/Rev.hs @@ -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 -- diff --git a/src/Futhark/AD/Rev/Hist.hs b/src/Futhark/AD/Rev/Hist.hs --- a/src/Futhark/AD/Rev/Hist.hs +++ b/src/Futhark/AD/Rev/Hist.hs @@ -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 -- diff --git a/src/Futhark/AD/Rev/Reduce.hs b/src/Futhark/AD/Rev/Reduce.hs --- a/src/Futhark/AD/Rev/Reduce.hs +++ b/src/Futhark/AD/Rev/Reduce.hs @@ -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 diff --git a/src/Futhark/AD/Rev/Scan.hs b/src/Futhark/AD/Rev/Scan.hs --- a/src/Futhark/AD/Rev/Scan.hs +++ b/src/Futhark/AD/Rev/Scan.hs @@ -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 diff --git a/src/Futhark/Analysis/AccessPattern.hs b/src/Futhark/Analysis/AccessPattern.hs --- a/src/Futhark/Analysis/AccessPattern.hs +++ b/src/Futhark/Analysis/AccessPattern.hs @@ -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 diff --git a/src/Futhark/Analysis/HORep/MapNest.hs b/src/Futhark/Analysis/HORep/MapNest.hs --- a/src/Futhark/Analysis/HORep/MapNest.hs +++ b/src/Futhark/Analysis/HORep/MapNest.hs @@ -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 diff --git a/src/Futhark/Analysis/HORep/SOAC.hs b/src/Futhark/Analysis/HORep/SOAC.hs --- a/src/Futhark/Analysis/HORep/SOAC.hs +++ b/src/Futhark/Analysis/HORep/SOAC.hs @@ -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) = diff --git a/src/Futhark/Analysis/SymbolTable.hs b/src/Futhark/Analysis/SymbolTable.hs --- a/src/Futhark/Analysis/SymbolTable.hs +++ b/src/Futhark/Analysis/SymbolTable.hs @@ -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' diff --git a/src/Futhark/CLI/Profile.hs b/src/Futhark/CLI/Profile.hs --- a/src/Futhark/CLI/Profile.hs +++ b/src/Futhark/CLI/Profile.hs @@ -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 diff --git a/src/Futhark/CLI/Test.hs b/src/Futhark/CLI/Test.hs --- a/src/Futhark/CLI/Test.hs +++ b/src/Futhark/CLI/Test.hs @@ -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 diff --git a/src/Futhark/CodeGen/Backends/GenericC.hs b/src/Futhark/CodeGen/Backends/GenericC.hs --- a/src/Futhark/CodeGen/Backends/GenericC.hs +++ b/src/Futhark/CodeGen/Backends/GenericC.hs @@ -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 diff --git a/src/Futhark/CodeGen/Backends/SimpleRep.hs b/src/Futhark/CodeGen/Backends/SimpleRep.hs --- a/src/Futhark/CodeGen/Backends/SimpleRep.hs +++ b/src/Futhark/CodeGen/Backends/SimpleRep.hs @@ -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 diff --git a/src/Futhark/CodeGen/ImpGen.hs b/src/Futhark/CodeGen/ImpGen.hs --- a/src/Futhark/CodeGen/ImpGen.hs +++ b/src/Futhark/CodeGen/ImpGen.hs @@ -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 diff --git a/src/Futhark/CodeGen/ImpGen/GPU.hs b/src/Futhark/CodeGen/ImpGen/GPU.hs --- a/src/Futhark/CodeGen/ImpGen/GPU.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU.hs @@ -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 :: diff --git a/src/Futhark/CodeGen/ImpGen/GPU/Base.hs b/src/Futhark/CodeGen/ImpGen/GPU/Base.hs --- a/src/Futhark/CodeGen/ImpGen/GPU/Base.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU/Base.hs @@ -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) diff --git a/src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs b/src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs --- a/src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs @@ -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 _ = diff --git a/src/Futhark/CodeGen/RTS/C.hs b/src/Futhark/CodeGen/RTS/C.hs --- a/src/Futhark/CodeGen/RTS/C.hs +++ b/src/Futhark/CodeGen/RTS/C.hs @@ -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@ diff --git a/src/Futhark/Doc/Generator.hs b/src/Futhark/Doc/Generator.hs --- a/src/Futhark/Doc/Generator.hs +++ b/src/Futhark/Doc/Generator.hs @@ -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 diff --git a/src/Futhark/Fmt/Printer.hs b/src/Futhark/Fmt/Printer.hs --- a/src/Futhark/Fmt/Printer.hs +++ b/src/Futhark/Fmt/Printer.hs @@ -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 diff --git a/src/Futhark/IR/Mem.hs b/src/Futhark/IR/Mem.hs --- a/src/Futhark/IR/Mem.hs +++ b/src/Futhark/IR/Mem.hs @@ -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 diff --git a/src/Futhark/IR/Mem/Simplify.hs b/src/Futhark/IR/Mem/Simplify.hs --- a/src/Futhark/IR/Mem/Simplify.hs +++ b/src/Futhark/IR/Mem/Simplify.hs @@ -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 + ] + [] diff --git a/src/Futhark/IR/Parse.hs b/src/Futhark/IR/Parse.hs --- a/src/Futhark/IR/Parse.hs +++ b/src/Futhark/IR/Parse.hs @@ -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 diff --git a/src/Futhark/IR/Pretty.hs b/src/Futhark/IR/Pretty.hs --- a/src/Futhark/IR/Pretty.hs +++ b/src/Futhark/IR/Pretty.hs @@ -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) = diff --git a/src/Futhark/IR/Prop/Aliases.hs b/src/Futhark/IR/Prop/Aliases.hs --- a/src/Futhark/IR/Prop/Aliases.hs +++ b/src/Futhark/IR/Prop/Aliases.hs @@ -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] diff --git a/src/Futhark/IR/Prop/Names.hs b/src/Futhark/IR/Prop/Names.hs --- a/src/Futhark/IR/Prop/Names.hs +++ b/src/Futhark/IR/Prop/Names.hs @@ -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 diff --git a/src/Futhark/IR/Prop/Reshape.hs b/src/Futhark/IR/Prop/Reshape.hs --- a/src/Futhark/IR/Prop/Reshape.hs +++ b/src/Futhark/IR/Prop/Reshape.hs @@ -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 #-} diff --git a/src/Futhark/IR/Prop/TypeOf.hs b/src/Futhark/IR/Prop/TypeOf.hs --- a/src/Futhark/IR/Prop/TypeOf.hs +++ b/src/Futhark/IR/Prop/TypeOf.hs @@ -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] diff --git a/src/Futhark/IR/SOACS/Simplify.hs b/src/Futhark/IR/SOACS/Simplify.hs --- a/src/Futhark/IR/SOACS/Simplify.hs +++ b/src/Futhark/IR/SOACS/Simplify.hs @@ -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) diff --git a/src/Futhark/IR/Syntax.hs b/src/Futhark/IR/Syntax.hs --- a/src/Futhark/IR/Syntax.hs +++ b/src/Futhark/IR/Syntax.hs @@ -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 diff --git a/src/Futhark/IR/Syntax/Core.hs b/src/Futhark/IR/Syntax/Core.hs --- a/src/Futhark/IR/Syntax/Core.hs +++ b/src/Futhark/IR/Syntax/Core.hs @@ -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. diff --git a/src/Futhark/IR/Traversals.hs b/src/Futhark/IR/Traversals.hs --- a/src/Futhark/IR/Traversals.hs +++ b/src/Futhark/IR/Traversals.hs @@ -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)) = diff --git a/src/Futhark/IR/TypeCheck.hs b/src/Futhark/IR/TypeCheck.hs --- a/src/Futhark/IR/TypeCheck.hs +++ b/src/Futhark/IR/TypeCheck.hs @@ -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 diff --git a/src/Futhark/Internalise/Exps.hs b/src/Futhark/Internalise/Exps.hs --- a/src/Futhark/Internalise/Exps.hs +++ b/src/Futhark/Internalise/Exps.hs @@ -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 diff --git a/src/Futhark/Internalise/Monomorphise.hs b/src/Futhark/Internalise/Monomorphise.hs --- a/src/Futhark/Internalise/Monomorphise.hs +++ b/src/Futhark/Internalise/Monomorphise.hs @@ -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 diff --git a/src/Futhark/Internalise/ReplaceRecords.hs b/src/Futhark/Internalise/ReplaceRecords.hs --- a/src/Futhark/Internalise/ReplaceRecords.hs +++ b/src/Futhark/Internalise/ReplaceRecords.hs @@ -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 diff --git a/src/Futhark/Optimise/ArrayLayout/Transform.hs b/src/Futhark/Optimise/ArrayLayout/Transform.hs --- a/src/Futhark/Optimise/ArrayLayout/Transform.hs +++ b/src/Futhark/Optimise/ArrayLayout/Transform.hs @@ -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 = diff --git a/src/Futhark/Optimise/ArrayShortCircuiting/MemRefAggreg.hs b/src/Futhark/Optimise/ArrayShortCircuiting/MemRefAggreg.hs --- a/src/Futhark/Optimise/ArrayShortCircuiting/MemRefAggreg.hs +++ b/src/Futhark/Optimise/ArrayShortCircuiting/MemRefAggreg.hs @@ -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) diff --git a/src/Futhark/Optimise/ArrayShortCircuiting/TopdownAnalysis.hs b/src/Futhark/Optimise/ArrayShortCircuiting/TopdownAnalysis.hs --- a/src/Futhark/Optimise/ArrayShortCircuiting/TopdownAnalysis.hs +++ b/src/Futhark/Optimise/ArrayShortCircuiting/TopdownAnalysis.hs @@ -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 diff --git a/src/Futhark/Optimise/BlkRegTiling.hs b/src/Futhark/Optimise/BlkRegTiling.hs --- a/src/Futhark/Optimise/BlkRegTiling.hs +++ b/src/Futhark/Optimise/BlkRegTiling.hs @@ -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) diff --git a/src/Futhark/Optimise/EntryPointMem.hs b/src/Futhark/Optimise/EntryPointMem.hs --- a/src/Futhark/Optimise/EntryPointMem.hs +++ b/src/Futhark/Optimise/EntryPointMem.hs @@ -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, diff --git a/src/Futhark/Optimise/Fusion.hs b/src/Futhark/Optimise/Fusion.hs --- a/src/Futhark/Optimise/Fusion.hs +++ b/src/Futhark/Optimise/Fusion.hs @@ -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 diff --git a/src/Futhark/Optimise/Fusion/GraphRep.hs b/src/Futhark/Optimise/Fusion/GraphRep.hs --- a/src/Futhark/Optimise/Fusion/GraphRep.hs +++ b/src/Futhark/Optimise/Fusion/GraphRep.hs @@ -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 diff --git a/src/Futhark/Optimise/Fusion/RulesWithAccs.hs b/src/Futhark/Optimise/Fusion/RulesWithAccs.hs --- a/src/Futhark/Optimise/Fusion/RulesWithAccs.hs +++ b/src/Futhark/Optimise/Fusion/RulesWithAccs.hs @@ -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 diff --git a/src/Futhark/Optimise/Fusion/TryFusion.hs b/src/Futhark/Optimise/Fusion/TryFusion.hs --- a/src/Futhark/Optimise/Fusion/TryFusion.hs +++ b/src/Futhark/Optimise/Fusion/TryFusion.hs @@ -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') diff --git a/src/Futhark/Optimise/GenRedOpt.hs b/src/Futhark/Optimise/GenRedOpt.hs --- a/src/Futhark/Optimise/GenRedOpt.hs +++ b/src/Futhark/Optimise/GenRedOpt.hs @@ -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 diff --git a/src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs b/src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs --- a/src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs +++ b/src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs @@ -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 diff --git a/src/Futhark/Optimise/Simplify/Engine.hs b/src/Futhark/Optimise/Simplify/Engine.hs --- a/src/Futhark/Optimise/Simplify/Engine.hs +++ b/src/Futhark/Optimise/Simplify/Engine.hs @@ -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 diff --git a/src/Futhark/Optimise/Simplify/Rules/BasicOp.hs b/src/Futhark/Optimise/Simplify/Rules/BasicOp.hs --- a/src/Futhark/Optimise/Simplify/Rules/BasicOp.hs +++ b/src/Futhark/Optimise/Simplify/Rules/BasicOp.hs @@ -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 diff --git a/src/Futhark/Optimise/Simplify/Rules/Index.hs b/src/Futhark/Optimise/Simplify/Rules/Index.hs --- a/src/Futhark/Optimise/Simplify/Rules/Index.hs +++ b/src/Futhark/Optimise/Simplify/Rules/Index.hs @@ -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 diff --git a/src/Futhark/Optimise/Simplify/Rules/Simple.hs b/src/Futhark/Optimise/Simplify/Rules/Simple.hs --- a/src/Futhark/Optimise/Simplify/Rules/Simple.hs +++ b/src/Futhark/Optimise/Simplify/Rules/Simple.hs @@ -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 ] diff --git a/src/Futhark/Optimise/TileLoops.hs b/src/Futhark/Optimise/TileLoops.hs --- a/src/Futhark/Optimise/TileLoops.hs +++ b/src/Futhark/Optimise/TileLoops.hs @@ -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' diff --git a/src/Futhark/Optimise/TileLoops/Shared.hs b/src/Futhark/Optimise/TileLoops/Shared.hs --- a/src/Futhark/Optimise/TileLoops/Shared.hs +++ b/src/Futhark/Optimise/TileLoops/Shared.hs @@ -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))) diff --git a/src/Futhark/Pass/ExplicitAllocations.hs b/src/Futhark/Pass/ExplicitAllocations.hs --- a/src/Futhark/Pass/ExplicitAllocations.hs +++ b/src/Futhark/Pass/ExplicitAllocations.hs @@ -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 diff --git a/src/Futhark/Pass/ExplicitAllocations/GPU.hs b/src/Futhark/Pass/ExplicitAllocations/GPU.hs --- a/src/Futhark/Pass/ExplicitAllocations/GPU.hs +++ b/src/Futhark/Pass/ExplicitAllocations/GPU.hs @@ -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 diff --git a/src/Futhark/Pass/ExtractKernels/DistributeNests.hs b/src/Futhark/Pass/ExtractKernels/DistributeNests.hs --- a/src/Futhark/Pass/ExtractKernels/DistributeNests.hs +++ b/src/Futhark/Pass/ExtractKernels/DistributeNests.hs @@ -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 diff --git a/src/Futhark/Pass/ExtractKernels/ISRWIM.hs b/src/Futhark/Pass/ExtractKernels/ISRWIM.hs --- a/src/Futhark/Pass/ExtractKernels/ISRWIM.hs +++ b/src/Futhark/Pass/ExtractKernels/ISRWIM.hs @@ -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 = diff --git a/src/Futhark/Pass/LiftAllocations.hs b/src/Futhark/Pass/LiftAllocations.hs --- a/src/Futhark/Pass/LiftAllocations.hs +++ b/src/Futhark/Pass/LiftAllocations.hs @@ -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 :: diff --git a/src/Futhark/Tools.hs b/src/Futhark/Tools.hs --- a/src/Futhark/Tools.hs +++ b/src/Futhark/Tools.hs @@ -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 diff --git a/src/Futhark/Transform/Substitute.hs b/src/Futhark/Transform/Substitute.hs --- a/src/Futhark/Transform/Substitute.hs +++ b/src/Futhark/Transform/Substitute.hs @@ -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 diff --git a/src/Language/Futhark/Interpreter.hs b/src/Language/Futhark/Interpreter.hs --- a/src/Language/Futhark/Interpreter.hs +++ b/src/Language/Futhark/Interpreter.hs @@ -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 diff --git a/src/Language/Futhark/Interpreter/Values.hs b/src/Language/Futhark/Interpreter/Values.hs --- a/src/Language/Futhark/Interpreter/Values.hs +++ b/src/Language/Futhark/Interpreter/Values.hs @@ -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 diff --git a/src/Language/Futhark/Parser/Parser.y b/src/Language/Futhark/Parser/Parser.y --- a/src/Language/Futhark/Parser/Parser.y +++ b/src/Language/Futhark/Parser/Parser.y @@ -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 } diff --git a/src/Language/Futhark/Pretty.hs b/src/Language/Futhark/Pretty.hs --- a/src/Language/Futhark/Pretty.hs +++ b/src/Language/Futhark/Pretty.hs @@ -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 -> diff --git a/src/Language/Futhark/Semantic.hs b/src/Language/Futhark/Semantic.hs --- a/src/Language/Futhark/Semantic.hs +++ b/src/Language/Futhark/Semantic.hs @@ -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 diff --git a/src/Language/Futhark/TypeChecker.hs b/src/Language/Futhark/TypeChecker.hs --- a/src/Language/Futhark/TypeChecker.hs +++ b/src/Language/Futhark/TypeChecker.hs @@ -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 diff --git a/src/Language/Futhark/TypeChecker/Modules.hs b/src/Language/Futhark/TypeChecker/Modules.hs --- a/src/Language/Futhark/TypeChecker/Modules.hs +++ b/src/Language/Futhark/TypeChecker/Modules.hs @@ -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 diff --git a/src/Language/Futhark/TypeChecker/Terms.hs b/src/Language/Futhark/TypeChecker/Terms.hs --- a/src/Language/Futhark/TypeChecker/Terms.hs +++ b/src/Language/Futhark/TypeChecker/Terms.hs @@ -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 "") (srclocOf e) [] [] $ toRes Nonunique t - fixOverloadedTypes $ typeVars t + (tparams, _, RetType _ t') <- + letGeneralise (nameFromString "") (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' diff --git a/src/Language/Futhark/TypeChecker/Types.hs b/src/Language/Futhark/TypeChecker/Types.hs --- a/src/Language/Futhark/TypeChecker/Types.hs +++ b/src/Language/Futhark/TypeChecker/Types.hs @@ -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 $ diff --git a/src/Language/Futhark/TypeChecker/Unify.hs b/src/Language/Futhark/TypeChecker/Unify.hs --- a/src/Language/Futhark/TypeChecker/Unify.hs +++ b/src/Language/Futhark/TypeChecker/Unify.hs @@ -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) diff --git a/unittests/Futhark/IR/Prop/ReshapeTests.hs b/unittests/Futhark/IR/Prop/ReshapeTests.hs --- a/unittests/Futhark/IR/Prop/ReshapeTests.hs +++ b/unittests/Futhark/IR/Prop/ReshapeTests.hs @@ -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] + ]