diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a7a48dfd1..36ad5f105 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -42,18 +42,12 @@ jobs: - '9.6' - '9.4' - '9.2' - - '9.0' - - '8.10' mode: - 'debug' - 'release' exclude: - os: 'windows-latest' mode: 'debug' - - os: 'macOS-latest' # GHC < 9.2 is broken on ARM64 macs - ghc: '8.10' - - os: 'macOS-latest' - ghc: '9.0' steps: - uses: actions/checkout@v3 diff --git a/.gitignore b/.gitignore index 922dbbc64..87a76e6a1 100644 --- a/.gitignore +++ b/.gitignore @@ -16,3 +16,5 @@ /docs/_build *.hi *.o +/gurobi.log +/.vscode diff --git a/README.md b/README.md index 4bef6973b..4c85efbdf 100644 --- a/README.md +++ b/README.md @@ -30,6 +30,11 @@ Chapter 6 of Simon Marlow's book [Parallel and Concurrent Programming in Haskell [Trevor's PhD thesis][Trevor-thesis] details the design and implementation of frontend optimisations and CUDA backend. +**New pipeline** + +This branch contains a new, work-in-progress, version of the Accelerate compiler. The Native backend is near completion and the PTX backend is in active development. + +To use this new version, use the latest commit from this branch (new-pipeline) of accelerate and [accelerate-llvm](https://github.com/ivogabe/accelerate-llvm/tree/new-pipeline). Either add them as a git dependency to your cabal or stack configuration, or clone them manually and add them as a local dependency. Install `clang` and `cbc` and add them to your PATH. Accelerate uses clang and cbc via the command line (instead of linking with them). The new version of Accelerate should then be working. Note that some functionality is not implemented yet (including the FFI and the trace functions). **Table of Contents** diff --git a/Setup.hs b/Setup.hs index d27a2a5fd..bd53110f3 100755 --- a/Setup.hs +++ b/Setup.hs @@ -1,11 +1,8 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE MultiWayIf #-} +{-# LANGUAGE CPP #-} {-# OPTIONS -Wall #-} module Main where import Distribution.Extra.Doctest -import Distribution.PackageDescription hiding (Flag) import Distribution.Simple import Distribution.Simple.LocalBuildInfo import Distribution.Simple.Setup @@ -13,6 +10,12 @@ import Distribution.Simple.Utils import Distribution.System import Distribution.Verbosity +#if MIN_VERSION_Cabal(3,4,0) +import Distribution.PackageDescription +#else +import Distribution.PackageDescription hiding (Flag) +#endif + #if MIN_VERSION_Cabal(3,14,0) -- Note [Cabal 3.14] -- @@ -120,7 +123,6 @@ tracyMode config_flags = fromMaybe False $ lookupFlagAssignment (mkFlagName "tracy") (configConfigurationsFlags config_flags) - rawSystemExit' :: Verbosity -> FilePath -> [String] -> IO () #if MIN_VERSION_Cabal(3,14,0) diff --git a/accelerate.cabal b/accelerate.cabal index 25da7494b..900d40a2a 100644 --- a/accelerate.cabal +++ b/accelerate.cabal @@ -133,6 +133,7 @@ stability: Experimental extra-source-files: cbits/*.c + cbits/*.cpp cbits/*.h -- XKCP cbits/xkcp/*.c @@ -141,15 +142,15 @@ extra-source-files: cbits/xkcp/*.inc -- TRACY -- These are referenced directly using the FFI - cbits/tracy/public/*.cpp - cbits/tracy/public/tracy/*.h - cbits/tracy/public/tracy/*.hpp - cbits/tracy/public/common/*.h - cbits/tracy/public/common/*.hpp - cbits/tracy/public/common/*.cpp - cbits/tracy/public/client/*.h - cbits/tracy/public/client/*.hpp - cbits/tracy/public/client/*.cpp + -- cbits/tracy/public/*.cpp + -- cbits/tracy/public/tracy/*.h + -- cbits/tracy/public/tracy/*.hpp + -- cbits/tracy/public/common/*.h + -- cbits/tracy/public/common/*.hpp + -- cbits/tracy/public/common/*.cpp + -- cbits/tracy/public/client/*.h + -- cbits/tracy/public/client/*.hpp + -- cbits/tracy/public/client/*.cpp -- These are used to build Tracy's client tools in Setup.hs cbits/tracy/cmake/*.cmake cbits/tracy/cmake/*.patch @@ -174,6 +175,13 @@ extra-source-files: cbits/tracy/server/*.cpp cbits/tracy/server/*.h cbits/tracy/server/*.hpp + -- These are used to disable/enable Tracy + cbits/flag-no-tracy/flag_tracy.h + cbits/flag-tracy/flag_tracy.h + -- ParlayLib + cbits/parlaylib/include/parlay/*.h + cbits/parlaylib/include/parlay/internal/*.h + cbits/parlaylib/include/parlay/internal/concurrency/*.h extra-doc-files: images/*.png @@ -297,8 +305,15 @@ flag tracy flag bounds-checks manual: True - default: True - description: Enable bounds checking in the interpreter + -- For now disabled by default. We may enable it by default when we know the + -- performance overhead of assertions. + default: False + description: Enable bounds checking + +flag memory-counter + manual: True + default: False + description: Enable counters for benchmarking memory usage -- This flag is currently completely unused, so let's not expose it to users -- flag unsafe-checks @@ -322,8 +337,12 @@ library build-depends: base >= 4.12 && < 4.23 , ansi-terminal >= 0.6.2 + , array >= 0.4 , base-orphans >= 0.3 + , bifunctors , bytestring >= 0.10.2 + -- , comfort-array >= 0.5.5 + , composition , containers >= 0.3 , deepseq >= 1.3 , directory >= 1.0 @@ -337,23 +356,26 @@ library , hashtables >= 1.2.3 , hedgehog >= 0.5 , microlens >= 0.4 + , microlens-th >= 0.4 + , microlens-mtl >= 0.2 + , MIP >= 0.2 , mtl >= 2.0 , prettyprinter >= 1.7 , prettyprinter-ansi-terminal >= 1.1.2 , primitive >= 0.6.4 + , scientific >= 0.3.6.2 , tasty >= 0.11 , template-haskell < 2.25 , terminal-size >= 0.3 , text >= 1.2.4 , transformers >= 0.3 - , unique , unordered-containers >= 0.2 , vector >= 0.10 - exposed-modules: -- The core language and reference implementation Data.Array.Accelerate Data.Array.Accelerate.Interpreter + Data.Array.Accelerate.Interpreter.Simple -- Prelude-like Data.Array.Accelerate.Control.Monad @@ -370,25 +392,44 @@ library -- For backend development (hidden) Data.Array.Accelerate.AST + Data.Array.Accelerate.AST.CountEnv Data.Array.Accelerate.AST.Environment + Data.Array.Accelerate.AST.Execute + Data.Array.Accelerate.AST.Exp + Data.Array.Accelerate.AST.Graph Data.Array.Accelerate.AST.Idx + Data.Array.Accelerate.AST.IdxSet + Data.Array.Accelerate.AST.Kernel Data.Array.Accelerate.AST.LeftHandSide + Data.Array.Accelerate.AST.Operation + Data.Array.Accelerate.AST.Partitioned + Data.Array.Accelerate.AST.Schedule + Data.Array.Accelerate.AST.Schedule.Sequential + Data.Array.Accelerate.AST.Schedule.Uniform Data.Array.Accelerate.AST.Var + Data.Array.Accelerate.Analysis.Exp Data.Array.Accelerate.Analysis.Hash + Data.Array.Accelerate.Analysis.Hash.Exp + Data.Array.Accelerate.Analysis.Hash.Operation + Data.Array.Accelerate.Analysis.Hash.Schedule.Uniform Data.Array.Accelerate.Analysis.Match - Data.Array.Accelerate.Array.Data - Data.Array.Accelerate.Array.Remote - Data.Array.Accelerate.Array.Remote.Class - Data.Array.Accelerate.Array.Remote.LRU - Data.Array.Accelerate.Array.Remote.Table - Data.Array.Accelerate.Array.Unique + Data.Array.Accelerate.Array.Buffer + -- Data.Array.Accelerate.Array.Remote + -- Data.Array.Accelerate.Array.Remote.Class + -- Data.Array.Accelerate.Array.Remote.LRU + -- Data.Array.Accelerate.Array.Remote.Table Data.Array.Accelerate.Async + Data.Array.Accelerate.Backend Data.Array.Accelerate.Error Data.Array.Accelerate.Debug.Internal Data.Array.Accelerate.Lifetime Data.Array.Accelerate.Pretty + Data.Array.Accelerate.Pretty.Schedule + Data.Array.Accelerate.Pretty.Schedule.Sequential + Data.Array.Accelerate.Pretty.Schedule.Uniform Data.Array.Accelerate.Representation.Array Data.Array.Accelerate.Representation.Elt + Data.Array.Accelerate.Representation.Ground Data.Array.Accelerate.Representation.Shape Data.Array.Accelerate.Representation.Slice Data.Array.Accelerate.Representation.Stencil @@ -403,18 +444,49 @@ library Data.Array.Accelerate.Sugar.Vec Data.Array.Accelerate.Trafo Data.Array.Accelerate.Trafo.Config - Data.Array.Accelerate.Trafo.Delayed - Data.Array.Accelerate.Trafo.Fusion + Data.Array.Accelerate.Trafo.Lowering + Data.Array.Accelerate.Trafo.Exp.Shrink + Data.Array.Accelerate.Trafo.Exp.Simplify + Data.Array.Accelerate.Trafo.Exp.Substitution Data.Array.Accelerate.Trafo.LetSplit + Data.Array.Accelerate.Trafo.LiveVars + Data.Array.Accelerate.Trafo.NewNewFusion + Data.Array.Accelerate.Trafo.Operation.Bounds + Data.Array.Accelerate.Trafo.Operation.Simplify + Data.Array.Accelerate.Trafo.Operation.Substitution + Data.Array.Accelerate.Trafo.Operation.LiveVars + Data.Array.Accelerate.Trafo.Partitioning.ILP + Data.Array.Accelerate.Trafo.Partitioning.ILP.Clustering + Data.Array.Accelerate.Trafo.Partitioning.ILP.ConstraintLanguage + Data.Array.Accelerate.Trafo.Partitioning.ILP.Encoding + Data.Array.Accelerate.Trafo.Partitioning.ILP.Lower + Data.Array.Accelerate.Trafo.Partitioning.ILP.Graph + Data.Array.Accelerate.Trafo.Partitioning.ILP.Labels + Data.Array.Accelerate.Trafo.Partitioning.ILP.LinearConstraint + -- Data.Array.Accelerate.Trafo.Partitioning.ILP.HiGHS + Data.Array.Accelerate.Trafo.Partitioning.ILP.MIP + Data.Array.Accelerate.Trafo.Partitioning.ILP.NameGeneration + Data.Array.Accelerate.Trafo.Partitioning.ILP.Solver + Data.Array.Accelerate.Trafo.Partitioning.ILP.Var + Data.Array.Accelerate.Trafo.Partitioning.Stream + Data.Array.Accelerate.Trafo.Schedule.Partial + Data.Array.Accelerate.Trafo.Schedule.Uniform + Data.Array.Accelerate.Trafo.Schedule.Uniform.Future + Data.Array.Accelerate.Trafo.Schedule.Uniform.LiveVars + Data.Array.Accelerate.Trafo.Schedule.Uniform.Simplify + Data.Array.Accelerate.Trafo.Schedule.Uniform.Substitution Data.Array.Accelerate.Trafo.Sharing - Data.Array.Accelerate.Trafo.Simplify + Data.Array.Accelerate.Trafo.SkipEnvironment Data.Array.Accelerate.Trafo.Substitution Data.Array.Accelerate.Trafo.Var + Data.Array.Accelerate.Trafo.WeakenedEnvironment Data.Array.Accelerate.Type -- For testing Data.Array.Accelerate.Test.NoFib Data.Array.Accelerate.Test.Similar + Data.Array.Accelerate.Pretty.Operation + Data.Array.Accelerate.Pretty.Exp -- Other Data.BitSet @@ -459,9 +531,13 @@ library Data.Array.Accelerate.Pretty.Graphviz Data.Array.Accelerate.Pretty.Graphviz.Monad Data.Array.Accelerate.Pretty.Graphviz.Type + Data.Array.Accelerate.Pretty.Partitioned Data.Array.Accelerate.Pretty.Print - Data.Array.Accelerate.Trafo.Algebra + Data.Array.Accelerate.Pretty.Type + Data.Array.Accelerate.Trafo.Exp.Algebra Data.Array.Accelerate.Trafo.Environment + Data.Array.Accelerate.Trafo.Operation.Bounds.Algebra + Data.Array.Accelerate.Trafo.Operation.Bounds.Environment Data.Array.Accelerate.Trafo.Shrink Data.Atomic @@ -541,7 +617,45 @@ library Data.Array.Accelerate.Data.Semigroup default-language: - Haskell2010 + GHC2021 + + -- Note OverloadedStrings is enabled by default to ensure that every module + -- is compatible with it, as many people configure GHCi to use it. + default-extensions: + DataKinds + DerivingStrategies + ExplicitNamespaces + GADTs + InstanceSigs + LambdaCase + OverloadedStrings + RoleAnnotations + + other-extensions: + AllowAmbiguousTypes + BlockArguments + CPP + DefaultSignatures + DeriveAnyClass + FunctionalDependencies + ImpredicativeTypes + MagicHash + MultiWayIf + NoImplicitPrelude + OverloadedLists + ParallelListComp + PatternSynonyms + QuantifiedConstraints + RebindableSyntax + RecordWildCards + TemplateHaskell + TemplateHaskellQuotes + TypeFamilies + TypeFamilyDependencies + UnboxedTuples + UndecidableInstances + UnliftedFFITypes + ViewPatterns hs-source-dirs: src @@ -553,7 +667,7 @@ library cxx-options: -O3 -Wall - -std=c++11 + -std=c++20 ghc-options: -O2 @@ -570,7 +684,7 @@ library ghc-options: -optcxx=-O3 -optcxx=-Wall - -optcxx=-std=c++11 + -optcxx=-std=c++20 if flag(debug) cc-options: @@ -579,6 +693,16 @@ library cpp-options: -DACCELERATE_DEBUG + -- Export the tracy flag via a header + -- so dependent libraries can use it + -- without having to redeclare the flag. + if flag(tracy) + include-dirs: cbits/flag-tracy + install-includes: flag_tracy.h + else + include-dirs: cbits/flag-no-tracy + install-includes: flag_tracy.h + if flag(tracy) if !flag(debug) build-depends: accelerate-tracy-flag-depends-on-debug-flag < 0 @@ -616,6 +740,13 @@ library cpp-options: -DACCELERATE_BOUNDS_CHECKS + if flag(memory-counter) + cpp-options: + -DACCELERATE_MEMORY_COUNTER + + cc-options: + -DACCELERATE_MEMORY_COUNTER + -- if flag(unsafe-checks) -- cpp-options: -- -DACCELERATE_UNSAFE_CHECKS @@ -654,34 +785,34 @@ library if impl(ghc >= 8.10) ghc-options: -optcxx=-mtune=native - -- Don't add the extensions list here. Instead, place individual LANGUAGE - -- pragmas in the files that require a specific extension. This means the - -- project loads in GHCi, and avoids extension clashes. - -- - -- Extensions: - test-suite doctest type: exitcode-stdio-1.0 default-language: Haskell2010 hs-source-dirs: test/doctest main-is: Main.hs - other-modules: Build_doctests - autogen-modules: Build_doctests + -- other-modules: Build_doctests + -- autogen-modules: Build_doctests build-depends: base >= 4.10 , accelerate - , doctest >= 0.11 + -- , doctest >= 0.11 + , criterion + , binary + , normaldistribution + , array + , random + , directory ghc-options: -Wall - -threaded + -- -threaded -rtsopts - x-doctest-options: - -freduction-depth=100 - -fspec-constr-count=50 + -- x-doctest-options: + -- -freduction-depth=100 + -- -fspec-constr-count=50 if os(windows) ghc-options: diff --git a/cbits/alloc.cpp b/cbits/alloc.cpp new file mode 100644 index 000000000..9fea50ebf --- /dev/null +++ b/cbits/alloc.cpp @@ -0,0 +1,31 @@ +#ifdef _WIN32 +#include +#include +#else +#include "./parlaylib/include/parlay/alloc.h" +#endif + +// Windows uses `_aligned_malloc` and `_aligned_free` from MinGW rather than ParlayLib. +// ParlayLib should compile with MinGW, see: https://github.com/cmuparlay/parlaylib +// But we have not gotten it to work. +// TODO: Maybe try getting Parlaylib to work on windows. + +extern "C" { + + void* accelerate_raw_alloc(uint64_t size, uint64_t align) { +#ifdef _WIN32 + return _aligned_malloc(size, align); +#else + return parlay::p_malloc(size, align); +#endif + } + + void accelerate_raw_free(void *ptr) { +#ifdef _WIN32 + _aligned_free(ptr); +#else + parlay::p_free(ptr); +#endif + } + +} diff --git a/cbits/flag-no-tracy/flag_tracy.h b/cbits/flag-no-tracy/flag_tracy.h new file mode 100644 index 000000000..e69de29bb diff --git a/cbits/flag-tracy/flag_tracy.h b/cbits/flag-tracy/flag_tracy.h new file mode 100644 index 000000000..1c213b732 --- /dev/null +++ b/cbits/flag-tracy/flag_tracy.h @@ -0,0 +1 @@ +#define ACCELERATE_TRACY diff --git a/cbits/memory.c b/cbits/memory.c new file mode 100644 index 000000000..a3bec603f --- /dev/null +++ b/cbits/memory.c @@ -0,0 +1,136 @@ +#include +#include +#include +#include +#include +#ifdef ACCELERATE_MEMORY_COUNTER +#include +#endif // ACCELERATE_MEMORY_COUNTER +#ifdef ACCELERATE_TRACY +#include "tracy/public/tracy/TracyC.h" +#endif // ACCELERATE_TRACY + +#include "align.h" +#include "flags.h" + +void* accelerate_raw_alloc(uint64_t size, uint64_t align); +void accelerate_raw_free(void *ptr); + +struct ObjectHeader { + _Atomic uint64_t reference_count; + uint64_t byte_size; // Size of the array. Does not include the header, and is not necessarily a multiple of the alignment. +} CACHE_ALIGNED; + +#ifdef ACCELERATE_MEMORY_COUNTER + +// For memory counting: +pthread_mutex_t mutex = PTHREAD_MUTEX_INITIALIZER; +uint64_t total_allocated = 0; +uint64_t current_allocated = 0; +uint64_t init_allocated = 0; +uint64_t max_allocated = 0; + +void accelerate_memory_counter_reset() { + pthread_mutex_lock(&mutex); + total_allocated = 0; + init_allocated = current_allocated; + max_allocated = current_allocated; + pthread_mutex_unlock(&mutex); +} + +uint64_t accelerate_memory_counter_total() { + pthread_mutex_lock(&mutex); + uint64_t ret = total_allocated; + pthread_mutex_unlock(&mutex); + return ret; +} + +uint64_t accelerate_memory_counter_init() { + pthread_mutex_lock(&mutex); + uint64_t ret = init_allocated; + pthread_mutex_unlock(&mutex); + return ret; +} + +uint64_t accelerate_memory_counter_max() { + pthread_mutex_lock(&mutex); + uint64_t ret = max_allocated; + pthread_mutex_unlock(&mutex); + return ret; +} + +#endif // ACCELERATE_MEMORY_COUNTER + + +void* accelerate_buffer_alloc(uint64_t byte_size) { + uint64_t size = byte_size + sizeof(struct ObjectHeader); + // Align size to the next multiple of CACHE_LINE_SIZE + // Note: it is apparently a requirement on macOS that the size is a multiple + // of the alignment, but it is not required on Linux. + size = (size + CACHE_LINE_SIZE - 1) / CACHE_LINE_SIZE * CACHE_LINE_SIZE; + void *ptr = accelerate_raw_alloc(size, CACHE_LINE_SIZE); + + #ifdef ACCELERATE_TRACY + ___tracy_emit_memory_alloc(ptr, size, 0); + #endif // ACCELERATE_TRACY + + #ifdef ACCELERATE_MEMORY_COUNTER + // For benchmarking memory usage: + pthread_mutex_lock(&mutex); + total_allocated += byte_size; + current_allocated += byte_size; + if (current_allocated > max_allocated) + max_allocated = current_allocated; + pthread_mutex_unlock(&mutex); + #endif // ACCELERATE_MEMORY_COUNTER + + struct ObjectHeader* header = (struct ObjectHeader*) ptr; + header->reference_count = 1; + header->byte_size = byte_size; + + void* interior = ptr + sizeof(struct ObjectHeader); + return interior; +} + +uint64_t accelerate_buffer_byte_size(void* interior) { + struct ObjectHeader* header = (struct ObjectHeader*) (interior - sizeof(struct ObjectHeader)); + return header->byte_size; +} + +void accelerate_buffer_retain_by(void* interior, uint64_t amount) { + if (amount == 0) return; + struct ObjectHeader* header = (struct ObjectHeader*) (interior - sizeof(struct ObjectHeader)); + atomic_fetch_add_explicit(&header->reference_count, amount, memory_order_acquire); +} + +void accelerate_buffer_retain(void* interior) { + accelerate_buffer_retain_by(interior, 1); +} + +void accelerate_buffer_release_by(void* interior, uint64_t amount) { + if (amount == 0) return; + struct ObjectHeader* header = (struct ObjectHeader*) (interior - sizeof(struct ObjectHeader)); + // Release is always needed, acquire only when the reference count drops to zero. + uint64_t old = atomic_fetch_add_explicit(&header->reference_count, -amount, memory_order_acq_rel); + if (old == amount) { + + #ifdef ACCELERATE_TRACY + ___tracy_emit_memory_free(header, 0); + #endif // ACCELERATE_TRACY + + #ifdef ACCELERATE_MEMORY_COUNTER + pthread_mutex_lock(&mutex); + current_allocated -= header->byte_size; + pthread_mutex_unlock(&mutex); + #endif // ACCELERATE_MEMORY_COUNTER + + accelerate_raw_free(header); + } + if (old < amount) { + printf("Reference count underflow\n"); + } +} + +void accelerate_buffer_release(void* interior) { + accelerate_buffer_release_by(interior, 1); +} diff --git a/cbits/parlaylib/include/parlay/alloc.h b/cbits/parlaylib/include/parlay/alloc.h new file mode 100644 index 000000000..7d919814c --- /dev/null +++ b/cbits/parlaylib/include/parlay/alloc.h @@ -0,0 +1,291 @@ +#ifndef PARLAY_ALLOC_H +#define PARLAY_ALLOC_H + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "type_traits.h" // IWYU pragma: keep +#include "utilities.h" + +#include "internal/block_allocator.h" +#include "internal/memory_size.h" +#include "internal/pool_allocator.h" + +// IWYU pragma: no_forward_declare is_trivially_relocatable + +namespace parlay { + +// ---------------------------------------------------------------------------- +// internal pool allocator +// +// allocates pools of blocks with powers of two sizes. Used internally to provide +// allocation for the container allocator, parlay::allocator; and for the free +// functions, parlay::p_malloc and parlay::p_free. +// ---------------------------------------------------------------------------- + +namespace internal { + +// The size of the largest pool used by the memory allocator +inline const size_t default_allocator_max_pool_size = getMemorySize() / 64; + +// these are bucket sizes used by the default allocator. They are powers +// of two starting at 16 and going until SystemMemory / 64. +// +// Note that them being powers of two is very important for correctness +// of the high-level allocators defined here! +// - being powers of two means that if a block is suitably aligned for an object +// of type T, the next block is guaranteed to also be suitably aligned for T +// - p_malloc optimizes its header by only storing log(size) of the allocation, +// which works because that is enough information to know which pool it came from +// +inline std::vector default_allocator_sizes() { + size_t log_min_size = 4; + size_t log_max_size = log2_up(default_allocator_max_pool_size); + + std::vector sizes; + for (size_t i = log_min_size; i <= log_max_size; i++) + sizes.push_back(size_t{1} << i); + return sizes; +} + +extern inline pool_allocator& get_default_allocator() { + static pool_allocator default_allocator(default_allocator_sizes()); + return default_allocator; +} + +// pair of total currently used space, and total unused space the allocator has in reserve +inline std::pair memory_usage() { + return get_default_allocator().stats(); +} + +inline void memory_clear() { + return get_default_allocator().clear(); +} + +} // namespace internal + +// ---------------------------------------------------------------------------- +// Free allocation functions +// +// Allocates headered blocks so the user doesn't need to remember the size +// ---------------------------------------------------------------------------- + +namespace internal { + +// Header used by p_malloc. Stores the requested size of the buffer and +// the offset of the underlying buffer that we need to free later. +struct p_malloc_header { + uint64_t log_size : 8; // We can store the log of the size because the default + uint64_t offset : 48; // allocator uses only powers of two pool sizes +}; +static_assert(sizeof(p_malloc_header) <= 8); + +// Minimum size of the padding for p_malloc. For large allocations, uses a larger +// padding to guarantee at least 64-byte alignment. Needs to be at least sizeof(size_t). +inline size_t alloc_padding_size(size_t n) { // in bytes + return (n >= 1024) ? 64 : (n & 15) ? 8 : (n & 63) ? 16 : 64; +} + +} // namespace internal + +// GCC doesn't like placement new into the offset buffer +#if defined(__GNUC__) && !defined(__clang__) +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wplacement-new" +#endif + +// Allocate size bytes of uninitialized storage. Optionally ask for an +// aligned buffer of memory with the given alignment. +// +// By default, the alignment is at least alignof(std::max_align_t), and may +// be larger for larger size allocations. +inline void* p_malloc(size_t size, size_t align = alignof(std::max_align_t)) { + assert(align > 0 && (align & (align - 1)) == 0 && "Requested alignment must be a power of two"); + + // allocates and tags with a header that contains the size and offset + size_t pad_size = std::max(internal::alloc_padding_size(size), align); + void* buffer = internal::get_default_allocator().allocate(size + pad_size); + void* offset_buffer = static_cast(static_cast(buffer) + sizeof(internal::p_malloc_header)); + size_t space = size + pad_size - sizeof(internal::p_malloc_header); + void* new_buffer = std::align(pad_size, size, offset_buffer, space); + assert(new_buffer != nullptr && new_buffer == offset_buffer); + size_t offset = (static_cast(new_buffer) - static_cast(buffer)); + assert(offset == (size + pad_size) - space); + new (static_cast(new_buffer) - sizeof(internal::p_malloc_header)) + internal::p_malloc_header{log2_up(size + pad_size), offset}; + return static_cast(new_buffer); +} + +// Free a block of memory obtained by p_malloc +inline void p_free(void* ptr) { + // reads the header to determine the offset and size, then frees + auto h = *from_bytes(static_cast(ptr) - sizeof(internal::p_malloc_header)); + if (h.log_size > 48u || h.offset > 1ull << 48) { + std::cerr << "corrupted header in p_free\n"; + std::abort(); + } + auto buffer = static_cast(static_cast(ptr) - h.offset); + internal::get_default_allocator().deallocate(buffer, size_t{1} << size_t(h.log_size)); +} + +#if defined(__GNUC__) && !defined(__clang__) +#pragma GCC diagnostic pop +#endif + +// ---------------------------------------------------------------------------- +// Container allocator +// ---------------------------------------------------------------------------- + +// A general purpose container allocator for arrays of type T. +// +// Matches the c++ Allocator specification (minimally) +// https://en.cppreference.com/w/cpp/named_req/Allocator +// Can therefore be used for STL containers, e.g.: +// std::vector> +// +template +struct allocator { + using value_type = T; + T* allocate(size_t n) { + // Use headerless blocks straight from the pool if the alignment is suitable, + // otherwise fall back to manually-aligned headered blocks from p_malloc + if constexpr (alignof(T) > internal::pool_allocator::max_alignment) { + return static_cast(p_malloc(n * sizeof(T), alignof(T))); + } + else { + // Should be suitably aligned for alignments up to pool_allocator::max_alignment + void* buffer = internal::get_default_allocator().allocate(n * sizeof(T)); + assert(reinterpret_cast(buffer) % alignof(T) == 0); + return static_cast(buffer); + } + } + void deallocate(T* ptr, [[maybe_unused]] size_t n) { + assert(reinterpret_cast(ptr) % alignof(T) == 0); + if constexpr (alignof(T) > internal::pool_allocator::max_alignment) { + p_free(static_cast(ptr)); + } + else { + internal::get_default_allocator().deallocate(static_cast(ptr), n * sizeof(T)); + } + } + + constexpr allocator() { internal::get_default_allocator(); } + template /* implicit */ constexpr allocator(const allocator&) noexcept { } +}; + +// Allocator should be trivially copyable since it is stateless and has no user-provided copy +// constructor. This should guarantee that it is also trivially relocatable. +static_assert(std::is_trivially_copyable_v>); +static_assert(is_trivially_relocatable_v>); + + +template +bool operator==(const allocator&, const allocator&) { return true; } +template +bool operator!=(const allocator&, const allocator&) { return false; } + + +// ---------------------------------------------------------------------------- +// Static allocator for single items of a given type, e.g. +// using long_allocator = type_allocator; +// long* foo = long_allocator::alloc(); +// new (foo) long{23}; +// long_allocator::free(foo); +// +// Uses block allocator internally, and is headerless +// ---------------------------------------------------------------------------- + +namespace internal { +template +extern inline block_allocator& get_block_allocator() { + static block_allocator a(Size, Align); + return a; +} +} + +// A static allocator for allocating storage for single objects of a fixed +// type. Much more efficient than using p_malloc or parlay::allocator for +// individual objects. +// +// Can be used to allocate raw uninitialized storage via alloc()/free(ptr), +// or to perform combined allocation and construction with create(args...) +// followed by destroy(ptr) to destroy and deallocate. +// +// alloc() -> T* : returns uninitialized storage for an object of type T +// free(T*) : deallocates storage obtained by alloc +// +// create(args...) -> T* : allocates storage for and constructs a T using args... +// destroy(T*) : destroys and deallocates a T obtained from create(...) +// +// All members are static, so it is not required to create an instance of +// type_allocator to use it. +// +template +class type_allocator { + + static internal::block_allocator& get_allocator() { + return internal::get_block_allocator(); + } + +public: + + // Allocate uninitialized storage appropriate for storing an object of type T + static T* alloc() { + void* buffer = get_allocator().alloc(); + assert(reinterpret_cast(buffer) % alignof(T) == 0); + return static_cast(buffer); + } + + // Free storage obtained by alloc() + static void free(T* ptr) { + assert(ptr != nullptr); + assert(reinterpret_cast(ptr) % alignof(T) == 0); + get_allocator().free(static_cast(ptr)); + } + + // Allocate storage for and then construct an object of type T using args... + template + static T* create(Args... args) { + static_assert(std::is_constructible_v); + return new (alloc()) T(std::forward(args)...); + } + + // Destroy an object obtained by create(...) and deallocate its storage + static void destroy(T* ptr) { + assert(ptr != nullptr); + ptr->~T(); + free(ptr); + } + + // for backward compatibility ----------------------------------------------- + static constexpr inline size_t default_alloc_size = 0; + static constexpr inline bool initialized = true; + + template + static T* allocate(Args... args) { return create(std::forward(args)...); } + static void retire(T* ptr) { destroy(ptr); } + static void init(size_t, size_t) {} + static void init() {} + static void reserve([[maybe_unused]] size_t n = default_alloc_size) { } + static void finish() { get_allocator().clear(); } + static size_t block_size () { return get_allocator().get_block_size(); } + static size_t num_allocated_blocks() { return get_allocator().num_allocated_blocks(); } + static size_t num_used_blocks() { return get_allocator().num_used_blocks(); } + static size_t num_used_bytes() { return num_used_blocks() * block_size(); } + static void print_stats() { get_allocator().print_stats(); } +}; + + +} // namespace parlay + +#endif // PARLAY_ALLOC_H diff --git a/cbits/parlaylib/include/parlay/internal/block_allocator.h b/cbits/parlaylib/include/parlay/internal/block_allocator.h new file mode 100644 index 000000000..eedbb57c1 --- /dev/null +++ b/cbits/parlaylib/include/parlay/internal/block_allocator.h @@ -0,0 +1,233 @@ +// A concurrent allocator for blocks of a fixed size +// +// Keeps a local pool per thread, and grabs list_length elements from a +// global pool if empty, and returns list_length elements to the global +// pool when local pool=2*list_length. +// +// Keeps track of number of allocated elements. Much more efficient +// than a general purpose allocator. +// +// Not generally intended for users. Users should use "type_allocator" +// which is a convenient wrapper around block_allocator that helps +// to manage memory for a specific type + +#ifndef PARLAY_INTERNAL_BLOCK_ALLOCATOR_H_ +#define PARLAY_INTERNAL_BLOCK_ALLOCATOR_H_ + +#include +#include + +#include +#include +#include +#include +#include + +#include "../utilities.h" +#include "../thread_specific.h" + +#include "concurrency/hazptr_stack.h" + +#include "memory_size.h" + +// IWYU pragma: no_include +// IWYU pragma: no_include + +namespace parlay { +namespace internal { + +struct block_allocator { + private: + + static inline constexpr size_t default_list_bytes = (1 << 18) - 64; // in bytes + static inline constexpr size_t min_alignment = 128; // for cache line padding + + struct block { + block* next; + }; + + struct alignas(128) local_list { + size_t sz; + block* head; + block* mid; + local_list() : sz(0), head(nullptr), mid(nullptr) {} + }; + + hazptr_stack allocated_buffers; + hazptr_stack global_stack; + ThreadSpecific my_local_list; + + size_t block_size; + std::align_val_t block_align; + size_t list_length; + size_t max_blocks; + std::atomic blocks_allocated; + + block* get_block(std::byte* buffer, size_t i) const { + // Since block is an aggregate type, it has implicit lifetime, so the following code + // is defined behaviour in C++20 even if we haven't yet called a constructor of block. + // It is still UB prior to C++20, but this is hard to avoid without losing performance. + return from_bytes(buffer + i * block_size); + } + + public: + block_allocator(const block_allocator&) = delete; + block_allocator(block_allocator&&) = delete; + block_allocator& operator=(const block_allocator&) = delete; + block_allocator& operator=(block_allocator&&) = delete; + + size_t get_block_size() const { return block_size; } + size_t num_allocated_blocks() const { return blocks_allocated.load(); } + + // Allocate a new list of list_length elements + + auto initialize_list(std::byte* buffer) const -> block* { + for (size_t i=0; i < list_length - 1; i++) { + new (buffer + i * block_size) block{get_block(buffer, i+1)}; + } + new (buffer + (list_length - 1) * block_size) block{nullptr}; + return get_block(buffer, 0); + } + + size_t num_used_blocks() { + size_t free_blocks = global_stack.size()*list_length; + my_local_list.for_each([&](auto&& list) { + free_blocks += list.sz; + }); + return blocks_allocated.load() - free_blocks; + } + + auto allocate_blocks(size_t num_blocks) -> std::byte* { + auto buffer = static_cast(::operator new(num_blocks * block_size, block_align)); + assert(buffer != nullptr); + + blocks_allocated.fetch_add(num_blocks); + assert(blocks_allocated.load() <= max_blocks); + + allocated_buffers.push(buffer); // keep track so can free later + return buffer; + } + + // Either grab a list from the global pool, or if there is none + // then allocate a new list + auto get_list() -> block* { + std::optional rem = global_stack.pop(); + if (rem) return *rem; + std::byte* buffer = allocate_blocks(list_length); + return initialize_list(buffer); + } + + // Allocate n elements across however many lists are needed (rounded up) + [[deprecated]] void reserve(size_t) { } + + void print_stats() { + size_t used = num_used_blocks(); + size_t allocated = num_allocated_blocks(); + size_t sz = get_block_size(); + std::cout << "Used: " << used << ", allocated: " << allocated + << ", block size: " << sz + << ", bytes: " << sz*allocated << "\n"; + } + + explicit block_allocator(size_t block_size_, + size_t block_align_ = alignof(std::max_align_t), + [[maybe_unused]] size_t reserved_blocks = 0, + size_t list_length_ = 0, + size_t max_blocks_ = 0) : + my_local_list(), // Each block needs to be at least + block_size(std::max(block_size_, sizeof(block))), // <------------- // large enough to hold the struct + block_align(std::align_val_t{std::max(block_align_, min_alignment)}), // representing a free block. + list_length(list_length_ == 0 ? (default_list_bytes + block_size + 1) / block_size : list_length_), + max_blocks(max_blocks_ == 0 ? (3 * getMemorySize() / block_size) / 4 : max_blocks_), + blocks_allocated(0) { + + } + + // Clears all memory ever allocated by this allocator. All allocated blocks + // must be returned before calling this function. + // + // This operation is not safe to perform concurrently with any other operation + // + // Returns false if there exists blocks that haven't been returned, in which case + // the operation fails and nothing is cleared. Returns true if no lingering + // blocks were not freed, in which case the operation is successful + bool clear() { + if (num_used_blocks() > 0) { + return false; + } + else { + // clear lists + my_local_list.for_each([](auto&& list) { + list.sz = 0; + }); + + // throw away all allocated memory + std::optional x; + while ((x = allocated_buffers.pop())) ::operator delete(*x, block_align); + global_stack.clear(); + blocks_allocated.store(0); + return true; + } + } + + ~block_allocator() { + [[maybe_unused]] auto cleared = clear(); +#if !defined(NDEBUG) && !defined(PARLAY_ALLOC_ALLOW_LEAK) + if (!cleared) { + std::cerr << "There are un-freed blocks obtained from block_allocator. If this is intentional you may" + "suppress this messsage with -DPARLAY_ALLOC_ALLOW_LEAK\n"; + } +#endif + } + + void free(void* ptr) { + + if (my_local_list->sz == list_length+1) { + my_local_list->mid = my_local_list->head; + } else if (my_local_list->sz == 2*list_length) { + global_stack.push(my_local_list->mid->next); + my_local_list->mid->next = nullptr; + my_local_list->sz = list_length; + } + + auto new_node = new (ptr) block{my_local_list->head}; + my_local_list->head = new_node; + my_local_list->sz++; + } + + inline void* alloc() { + if (my_local_list->sz == 0) { + auto new_list = get_list(); + + // !! Critical problem !! If this task got stolen during get_list(), the + // running thread may have changed, so we can't assume we are looking at + // the same local list, so we have to double-check the (possibly different) + // local list of the (possibly changed) thread + + if (my_local_list->sz == 0) { + my_local_list->head = new_list; + my_local_list->sz = list_length; + } + else { + // Looks like the task got stolen and the new thread already had a + // non-empty local list, so we can push the new one into the global + // pool for someone else to use in the future + global_stack.push(new_list); + } + } + + block* p = my_local_list->head; + my_local_list->head = my_local_list->head->next; + my_local_list->sz--; + + // Note: block is trivial, so it is legal to not call its destructor + // before returning it and allowing its storage to be reused + return static_cast(p); + } + +}; + +} // namespace internal +} // namespace parlay + +#endif // PARLAY_INTERNAL_BLOCK_ALLOCATOR_H_ diff --git a/cbits/parlaylib/include/parlay/internal/concurrency/acquire_retire.h b/cbits/parlaylib/include/parlay/internal/concurrency/acquire_retire.h new file mode 100644 index 000000000..1b643840e --- /dev/null +++ b/cbits/parlaylib/include/parlay/internal/concurrency/acquire_retire.h @@ -0,0 +1,171 @@ + +#ifndef PARLAY_INTERNAL_CONCURRENCY_ACQUIRE_RETIRE_H +#define PARLAY_INTERNAL_CONCURRENCY_ACQUIRE_RETIRE_H + +#include +#include + +#include // IWYU pragma: keep +#include +#include +#include // IWYU pragma: keep +#include +#include // IWYU pragma: keep + +#include "../../portability.h" +#include "../../thread_specific.h" + +// IWYU pragma: no_forward_declare parlay::padded + +namespace parlay { +namespace internal { + +// intrusive_acquire_retire works on types T that expose a T* via +// the ADL findable free functions intrusive_get_next(ptr) & intrusive_set_next(ptr) +template, size_t delay = 1> +class intrusive_acquire_retire { + + struct RetiredList { + + constexpr RetiredList() noexcept = default; + + ~RetiredList() { + assert(size == 0 && "RetiredLists must be emptied via cleanup() before destruction."); + } + + void push(T* p) noexcept { + intrusive_set_next(p, std::exchange(head, p)); + size++; + } + + [[nodiscard]] std::size_t get_size() const noexcept { + return size; + } + + template + void cleanup(F&& is_protected, Deleter& destroy) { + + while (head && !is_protected(head)) { + T* old = std::exchange(head, intrusive_get_next(head)); + destroy(old); + size--; + } + + if (head) { + T* prev = head; + T* current = intrusive_get_next(head); + while (current) { + if (!is_protected(current)) { + auto next = intrusive_get_next(current); + intrusive_set_next(prev, next); + destroy(current); + current = next; + size--; + } + else { + prev = std::exchange(current, intrusive_get_next(current)); + } + } + } + } + + private: + T* head = nullptr; + std::size_t size = 0; + }; + + public: + explicit intrusive_acquire_retire(Deleter deleter_ = {}) : + data(), deleter(std::move(deleter_)) {} + + template + U acquire(const std::atomic& p) { + static_assert(std::is_convertible_v, "acquire must read from a type that is convertible to T*"); + std::atomic& slot = data->announcement; + U result; + do { + result = p.load(std::memory_order_seq_cst); + PARLAY_PREFETCH(result, 0, 0); + slot.store(static_cast(result), std::memory_order_seq_cst); + } while (p.load(std::memory_order_seq_cst) != result); + return result; + } + + void release() { + data->announcement.store(nullptr); + } + + template + void retire(U p) { + static_assert(std::is_convertible_v, "retire must take a type that is convertible to T*"); + data->retired.push(static_cast(p)); + work_toward_ejects(); + } + + // Perform any remaining deferred destruction. Need to be very careful + // about additional objects being queued for deferred destruction by + // an object that was just destructed. + ~intrusive_acquire_retire() { + // Not a data race since no thread should be accessing the intrusive_acquire_retire + // instance while its destructor is being run. That would be a race from the user. + data.for_each([](auto& td) { + td.in_progress = true; + }); + + // Perform cleanup in a loop in case deleting a retired value causes + // something else to get retired (e.g., cleaning up a linked list) + bool retired = true; + while (retired) { + retired = false; + data.for_each([&](auto&& local_data) { + if (local_data.retired.get_size() > 0) { + retired = true; + local_data.retired.cleanup([](auto&&) { return false; }, deleter); + } + }); + } + } + + private: + // Apply the function f to every currently announced value + template + void scan_slots(F &&f) { + std::atomic_thread_fence(std::memory_order_seq_cst); + data.for_each([&](auto&& local_data) { + auto x = local_data.announcement.load(std::memory_order_seq_cst); + if (x != nullptr) f(x); + }); + } + + void work_toward_ejects(size_t work = 1) { + data->amortized_work += work; + auto threshold = std::max(30, delay * num_thread_ids()); // Always attempt at least 30 ejects + while (!(data->in_progress) && data->amortized_work >= threshold) { + data->amortized_work = 0; + if (data->retired.get_size() == 0) break; // nothing to collect + data->in_progress = true; + + std::unordered_set announced; + scan_slots([&](T* reserved) { announced.insert(reserved); }); + data->retired.cleanup([&](T* p) { return announced.count(p) != 0; }, deleter); + + data->in_progress = false; + } + } + + struct ThreadData { + std::atomic announcement{nullptr}; + bool in_progress{false}; // cppcheck-suppress unusedStructMember + size_t amortized_work{0}; // cppcheck-suppress unusedStructMember + RetiredList retired{}; + }; + + ThreadSpecific data; + PARLAY_NO_UNIQUE_ADDR Deleter deleter; +}; + + +} // namespace internal +} // namespace parlay + +#endif // PARLAY_INTERNAL_CONCURRENCY_ACQUIRE_RETIRE_H diff --git a/cbits/parlaylib/include/parlay/internal/concurrency/hazptr_stack.h b/cbits/parlaylib/include/parlay/internal/concurrency/hazptr_stack.h new file mode 100644 index 000000000..379974851 --- /dev/null +++ b/cbits/parlaylib/include/parlay/internal/concurrency/hazptr_stack.h @@ -0,0 +1,132 @@ + +#ifndef PARLAY_INTERNAL_CONCURRENCY_HAZPTR_STACK_H +#define PARLAY_INTERNAL_CONCURRENCY_HAZPTR_STACK_H + +#include + +#include +#include +#include // IWYU pragma: keep + +#include "acquire_retire.h" + +namespace parlay { +namespace internal { + +// A minimal, low-level, linearizable, lock-free concurrent +// stack that uses hazard-pointers for garbage collection. +// +// The stack keeps track of its current length +// +// Template parameters: +// T: The type of the element stored in the stack +template +class hazptr_stack { + + struct Node { + T t; + Node* next; + size_t length; + + // Allow the intrusive garbage collector to use the node's own next pointer + // instead of having to allocate its own memory to store the retired nodes + friend Node* intrusive_get_next(Node* node) { + return node->next; + } + + friend void intrusive_set_next(Node* node, Node* next_) { + node->next = next_; + } + + explicit Node(T t_) noexcept(std::is_nothrow_move_constructible_v) + : t(std::move(t_)), next(nullptr), length(1) { } + }; + + public: + hazptr_stack() : head(nullptr) { + // Touch the hazard pointer instance to force its initialization + // before this object, since it must outlive this object. Without + // this line, the hazard pointer instance might be destroyed before + // the stack, in which case subsequent uses of the stack may try + // to use the hazard pointer instance after it has been destroyed! + hazptr_instance(); + } + + // Insert the object t at the top of the stack + // + // This operation is safe to perform concurrently with other operations + void push(T t) { + auto p = new Node(std::move(t)); + do { + auto h = hazptr_instance().acquire(head); + p->next = h; + p->length = h ? h->length + 1 : 1; + } while (!head.compare_exchange_weak(p->next, p)); + hazptr_instance().release(); + } + + // Remove and return the object at the top of the stack. If the stack + // is empty, returns an empty optional. + // + // This operation is safe to perform concurrently with other operations + std::optional pop() { + Node* p; + do { + p = hazptr_instance().acquire(head); + if (p == nullptr) return {}; + } while (!head.compare_exchange_weak(p, p->next)); + auto val = std::make_optional(std::move(p->t)); + hazptr_instance().retire(p); + hazptr_instance().release(); + return val; + } + + // Return the current size of the stack + [[nodiscard]] size_t size() const noexcept { + auto p = hazptr_instance().acquire(head); + auto result = p ? p->length : 0; + hazptr_instance().release(); + return result; + } + + // Returns true if the stack is currently empty + [[nodiscard]] bool empty() const noexcept { + return head.load() == nullptr; + } + + ~hazptr_stack() { + // More efficient than clear since we can assume there are + // no concurrent accesses and hence we don't need to use + // pop, which acquires hazard pointers and uses a CAS loop + Node* p = head.load(); + Node* prev = nullptr; + while (p) { + prev = p; + p = p->next; + delete prev; + } + } + + // Empties the stack + // + // Can be safely invoked concurrently with other operations, though it + // is possible that this operation can be stalled indefinitely if other + // threads continue to push onto the stack forever, then this thread + // will just keep popping and never manage to empty it. + void clear() { + while (pop()) { } + } + + private: + std::atomic head; + + static intrusive_acquire_retire& hazptr_instance() { + static intrusive_acquire_retire instance; + return instance; + } +}; + +} // namespace internal +} // namespace parlay + +#endif // PARLAY_INTERNAL_CONCURRENCY_HAZPTR_STACK_H diff --git a/cbits/parlaylib/include/parlay/internal/debug_uninitialized.h b/cbits/parlaylib/include/parlay/internal/debug_uninitialized.h new file mode 100644 index 000000000..e061f544d --- /dev/null +++ b/cbits/parlaylib/include/parlay/internal/debug_uninitialized.h @@ -0,0 +1,143 @@ +#ifndef PARLAY_INTERNAL_DEBUG_UNINITIALIZED_H_ +#define PARLAY_INTERNAL_DEBUG_UNINITIALIZED_H_ + +#include + +namespace parlay { +namespace internal { + +// A simple type to help look for uninitialized memory bugs. +// UninitializedTracker is essentially an integer type, but +// also tracks whether it is currently in an initialized +// or uninitialized state. +// +// Attempting to assign to an uninitialized state will +// trigger an assertion failure. Attempting to copy or +// move an uninitialized object will also cause failure. +// +// This code will technically invoke undefined behaviour +// at some point, since we set the initialized flag to +// false in the destructor for the sole purpose of +// checking whether it is indeed false the next time +// someone comes along and wants to do an uninitialized +// in place construction at that memory location. +// +// Some compilers or debugging tools might therefore not +// work correctly with this class. Defining the initialized +// flag as volatile seems to prevent compilers from optimizing +// it out, and therefore leaves its value in memory after +// the object is destroyed, but this can not be guaranteed. +// +// For correctness, this type should only ever be used +// if its memory is allocated inside a parlay::sequence +// or parlay::uninitialized_sequence since they know how +// to set the initialized/uninitialized flag +struct UninitializedTracker { + + UninitializedTracker() : x(0), initialized(true) {} + + // cppcheck-suppress noExplicitConstructor + /* implicit */ UninitializedTracker(int _x) : x(_x), initialized(true) {} + + UninitializedTracker(const UninitializedTracker &other) { + assert(other.initialized && "Attempting to copy an uninitialized object!"); + x = other.x; + initialized = true; + } + + UninitializedTracker &operator=(const UninitializedTracker &other) { + assert(initialized && "Attempting to assign to an uninitialized object!"); + assert(other.initialized && "Copy assigning an uninitialized object!"); + x = other.x; + return *this; + } + + ~UninitializedTracker() { + assert(initialized && "Destructor called on uninitialized object!"); + initialized = false; + } + + bool operator==(const UninitializedTracker& other) const { + assert(initialized && "Trying to compare an uninitialized object!"); + assert(other.initialized && "Trying to compare against an uninitialized object!"); + return x == other.x; + } + + bool operator<(const UninitializedTracker& other) const { + assert(initialized && "Trying to compare an uninitialized object!"); + assert(other.initialized && "Trying to compare against an uninitialized object!"); + return x < other.x; + } + + void swap(UninitializedTracker& other) { + assert(initialized && "Trying to swap uninitialized object!"); + assert(other.initialized && "Trying to swap with an uninitialized object!"); + std::swap(x, other.x); + } + + friend void swap(parlay::internal::UninitializedTracker& a, parlay::internal::UninitializedTracker& b) { + a.swap(b); + } + + int x; + volatile bool initialized; // Volatile required to prevent optimizing out the store in the destructor +}; + +} // namespace internal +} // namespace parlay + + +// Macros for testing initialized/uninitializedness + +#ifdef PARLAY_DEBUG_UNINITIALIZED + +// Checks that the given UninitializedTracker object is uninitialized +// +// Usage: +// PARLAY_ASSERT_UNINITIALIZED(x) +// Result: +// Terminates the program if the expression (x) refers to an object of +// type UninitializedTracker such that (x).initialized is true, i.e. +// the object is in an initialized state +// +// Note: This macro should only ever be used on memory that was allocated +// by a parlay::sequence or a parlay::uninitialized_sequence, since they +// are required to appropriately flag all newly allocated memory as being +// uninitialized. False positives will occur if this macro is invoked on +// memory that is not managed by one of these containers. +#define PARLAY_ASSERT_UNINITIALIZED(x) \ +do { \ + using PARLAY_AU_T = std::remove_cv_t>; \ + if constexpr (std::is_same_v) { \ + assert(!((x).initialized) && "Memory required to be uninitialized is initialized!"); \ + } \ +} while (0) + +// Checks that the given UninitializedTracker object is initialized +// +// Usage: +// PARLAY_ASSERT_INITIALIZED(x) +// Result: +// Terminates the program if the expression (x) refers to an object of +// type UninitializedTracker such that (x).initialized is false, i.e. +// the object is in an uninitialized state +// +// Note: This macro should only ever be used on memory that was allocated +// by a parlay::sequence or a parlay::uninitialized_sequence, since they +// are required to appropriately flag all newly allocated memory as being +// uninitialized. False positives will occur if this macro is invoked on +// memory that is not managed by one of these containers. +#define PARLAY_ASSERT_INITIALIZED(x) \ +do { \ + using PARLAY_AI_T = std::remove_cv_t>; \ + if constexpr (std::is_same_v) { \ + assert((x).initialized && "Memory required to be initialized is uninitialized!"); \ + } \ +} while (0) + +#else +#define PARLAY_ASSERT_UNINITIALIZED(x) +#define PARLAY_ASSERT_INITIALIZED(x) +#endif // PARLAY_DEBUG_UNINITIALIZED + +#endif // PARLAY_INTERNAL_DEBUG_UNINITIALIZED_H_ diff --git a/cbits/parlaylib/include/parlay/internal/memory_size.h b/cbits/parlaylib/include/parlay/internal/memory_size.h new file mode 100644 index 000000000..5243e253a --- /dev/null +++ b/cbits/parlaylib/include/parlay/internal/memory_size.h @@ -0,0 +1,104 @@ +/* + * Author: David Robert Nadeau + * Site: http://NadeauSoftware.com/ + * License: Creative Commons Attribution 3.0 Unported License + * http://creativecommons.org/licenses/by/3.0/deed.en_US + */ + +#pragma once + +#if defined(_WIN32) + +#ifndef NOMINMAX +#define PARLAY_DEFINED_NOMINMAX +#define NOMINMAX +#endif + +#include + +#ifdef PARLAY_DEFINED_NOMINMAX +#undef NOMINMAX +#endif + +#elif defined(__unix__) || defined(__unix) || defined(unix) || \ + (defined(__APPLE__) && defined(__MACH__)) +#include +#include +#include +#if defined(BSD) +#include +#endif + +#else +#error "Unable to define getMemorySize( ) for an unknown OS." +#endif + +/** + * Returns the size of physical memory (RAM) in bytes. + */ +static size_t getMemorySize() { +#if defined(_WIN32) && (defined(__CYGWIN__) || defined(__CYGWIN32__)) + /* Cygwin under Windows. ------------------------------------ */ + /* New 64-bit MEMORYSTATUSEX isn't available. Use old 32.bit */ + MEMORYSTATUS status; + status.dwLength = sizeof(status); + GlobalMemoryStatus(&status); + return (size_t)status.dwTotalPhys; + +#elif defined(_WIN32) + /* Windows. ------------------------------------------------- */ + /* Use new 64-bit MEMORYSTATUSEX, not old 32-bit MEMORYSTATUS */ + MEMORYSTATUSEX status; + status.dwLength = sizeof(status); + GlobalMemoryStatusEx(&status); + return (size_t)status.ullTotalPhys; + +#elif defined(__unix__) || defined(__unix) || defined(unix) || \ + (defined(__APPLE__) && defined(__MACH__)) +/* UNIX variants. ------------------------------------------- */ +/* Prefer sysctl() over sysconf() except sysctl() HW_REALMEM and HW_PHYSMEM */ + +#if defined(CTL_HW) && (defined(HW_MEMSIZE) || defined(HW_PHYSMEM64)) + int mib[2]; + mib[0] = CTL_HW; +#if defined(HW_MEMSIZE) + mib[1] = HW_MEMSIZE; /* OSX. --------------------- */ +#elif defined(HW_PHYSMEM64) + mib[1] = HW_PHYSMEM64; /* NetBSD, OpenBSD. --------- */ +#endif + int64_t size = 0; /* 64-bit */ + size_t len = sizeof(size); + if (sysctl(mib, 2, &size, &len, NULL, 0) == 0) return (size_t)size; + return 0L; /* Failed? */ + +#elif defined(_SC_AIX_REALMEM) + /* AIX. ----------------------------------------------------- */ + return (size_t)sysconf(_SC_AIX_REALMEM) * (size_t)1024L; + +#elif defined(_SC_PHYS_PAGES) && defined(_SC_PAGESIZE) + /* FreeBSD, Linux, OpenBSD, and Solaris. -------------------- */ + return (size_t)sysconf(_SC_PHYS_PAGES) * (size_t)sysconf(_SC_PAGESIZE); + +#elif defined(_SC_PHYS_PAGES) && defined(_SC_PAGE_SIZE) + /* Legacy. -------------------------------------------------- */ + return (size_t)sysconf(_SC_PHYS_PAGES) * (size_t)sysconf(_SC_PAGE_SIZE); + +#elif defined(CTL_HW) && (defined(HW_PHYSMEM) || defined(HW_REALMEM)) + /* DragonFly BSD, FreeBSD, NetBSD, OpenBSD, and OSX. -------- */ + int mib[2]; + mib[0] = CTL_HW; +#if defined(HW_REALMEM) + mib[1] = HW_REALMEM; /* FreeBSD. ----------------- */ +#elif defined(HW_PYSMEM) + mib[1] = HW_PHYSMEM; /* Others. ------------------ */ +#endif + unsigned int size = 0; /* 32-bit */ + size_t len = sizeof(size); + if (sysctl(mib, 2, &size, &len, NULL, 0) == 0) return (size_t)size; + return 0L; /* Failed? */ +#endif /* sysctl and sysconf variants */ + +#else + return 0L; /* Unknown OS. */ +#endif +} diff --git a/cbits/parlaylib/include/parlay/internal/pool_allocator.h b/cbits/parlaylib/include/parlay/internal/pool_allocator.h new file mode 100644 index 000000000..2b8c2ac19 --- /dev/null +++ b/cbits/parlaylib/include/parlay/internal/pool_allocator.h @@ -0,0 +1,200 @@ + +#ifndef PARLAY_INTERNAL_POOL_ALLOCATOR_H +#define PARLAY_INTERNAL_POOL_ALLOCATOR_H + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "../portability.h" +#include "../utilities.h" + +#include "block_allocator.h" + +#include "concurrency/hazptr_stack.h" + +// IWYU pragma: no_include + +namespace parlay { +namespace internal { + +// **************************************** +// pool_allocator +// **************************************** + +// Allocates headerless blocks from pools of different sizes. +// A vector of pool sizes is given to the constructor. +// Sizes must be at least 8, and must increase. +// For pools of small blocks (below large_threshold) each thread keeps a +// thread local list of elements from each pool using the block_allocator. +// For large blocks there is only one pool shared by all threads. For +// blocks larger than the maximum pool size, allocation and deallocation +// is performed directly by operator new. +struct pool_allocator { + + // Maximum alignment guaranteed by the allocator + static inline constexpr size_t max_alignment = 128; + + private: + static inline constexpr size_t large_threshold = (1 << 18); + + size_t num_buckets; + size_t num_small; + size_t max_small; + size_t max_size; + std::atomic large_allocated{0}; + std::atomic large_used{0}; + + std::unique_ptr sizes; + std::unique_ptr[]> large_buckets; + unique_array small_allocators; + + void* allocate_large(size_t n) { + + size_t bucket = num_small; + size_t alloc_size; + large_used += n; + + if (n <= max_size) { + while (n > sizes[bucket]) bucket++; + std::optional r = large_buckets[bucket-num_small].pop(); + if (r) return *r; + alloc_size = sizes[bucket]; + } else alloc_size = n; + + // Alloc size must be a multiple of the alignment + // Round up to the next multiple. + if (alloc_size % max_alignment != 0) { + alloc_size += (max_alignment - (alloc_size % max_alignment)); + } + + void* a = ::operator new(alloc_size, std::align_val_t{max_alignment}); + + large_allocated += n; + return a; + } + + void deallocate_large(void* ptr, size_t n) { + large_used -= n; + if (n > max_size) { + ::operator delete(ptr, std::align_val_t{max_alignment}); + large_allocated -= n; + } else { + size_t bucket = num_small; + while (n > sizes[bucket]) bucket++; + large_buckets[bucket-num_small].push(ptr); + } + } + + public: + pool_allocator(const pool_allocator&) = delete; + pool_allocator(pool_allocator&&) = delete; + pool_allocator& operator=(const pool_allocator&) = delete; + pool_allocator& operator=(pool_allocator&&) = delete; + + ~pool_allocator() { + clear(); + } + + explicit pool_allocator(const std::vector& sizes_) : + num_buckets(sizes_.size()), + sizes(std::make_unique(num_buckets)) { + + std::copy(std::begin(sizes_), std::end(sizes_), sizes.get()); + max_size = sizes[num_buckets-1]; + num_small = 0; + while (num_small < num_buckets && sizes[num_small] < large_threshold) + num_small++; + max_small = (num_small > 0) ? sizes[num_small - 1] : 0; + + large_buckets = std::make_unique[]>(num_buckets-num_small); + small_allocators = make_unique_array(num_small, [&](size_t i) { + return internal::block_allocator(sizes[i], max_alignment); + }); + +#ifndef NDEBUG + size_t prev_bucket_size = 0; + for (size_t i = 0; i < num_small; i++) { + size_t bucket_size = sizes[i]; + assert(bucket_size >= 8); + assert(bucket_size > prev_bucket_size); + prev_bucket_size = bucket_size; + } +#endif + } + + void* allocate(size_t n) { + if (n > max_small) return allocate_large(n); + size_t bucket = 0; + while (n > sizes[bucket]) bucket++; + return small_allocators[bucket].alloc(); + } + + void deallocate(void* ptr, size_t n) { + if (n > max_small) deallocate_large(ptr, n); + else { + size_t bucket = 0; + while (n > sizes[bucket]) bucket++; + small_allocators[bucket].free(ptr); + } + } + + // allocate, touch, and free to make sure space for small blocks is paged in + [[deprecated]] void reserve(size_t) { } + + void print_stats() { + size_t total_a = 0; + size_t total_u = 0; + for (size_t i = 0; i < num_small; i++) { + size_t bucket_size = sizes[i]; + size_t allocated = small_allocators[i].num_allocated_blocks(); + size_t used = small_allocators[i].num_used_blocks(); + total_a += allocated * bucket_size; + total_u += used * bucket_size; + std::cout << "size = " << bucket_size << ", allocated = " << allocated + << ", used = " << used << std::endl; + } + std::cout << "Large allocated = " << large_allocated << std::endl; + std::cout << "Total bytes allocated = " << total_a + large_allocated << std::endl; + std::cout << "Total bytes used = " << total_u << std::endl; + } + + // pair of total currently used space, and total unused space the allocator has in reserve + std::pair stats() { + size_t total_a = large_allocated; + size_t total_u = large_used; + for (size_t i = 0; i < num_small; i++) { + size_t bucket_size = sizes[i]; + size_t allocated = small_allocators[i].num_allocated_blocks(); + size_t used = small_allocators[i].num_used_blocks(); + total_a += allocated * bucket_size; + total_u += used * bucket_size; + } + return {total_u, total_a-total_u}; + } + + void clear() { + for (size_t i = num_small; i < num_buckets; i++) { + std::optional r = large_buckets[i-num_small].pop(); + while (r) { + large_allocated -= sizes[i]; + ::operator delete(*r, std::align_val_t{max_alignment}); + r = large_buckets[i-num_small].pop(); + } + } + } +}; + +} // namespace internal +} // namespace parlay + +#endif // PARLAY_INTERNAL_POOL_ALLOCATOR_H diff --git a/cbits/parlaylib/include/parlay/internal/thread_id_pool.h b/cbits/parlaylib/include/parlay/internal/thread_id_pool.h new file mode 100644 index 000000000..5d8619674 --- /dev/null +++ b/cbits/parlaylib/include/parlay/internal/thread_id_pool.h @@ -0,0 +1,160 @@ + +#ifndef PARLAY_INTERNAL_THREAD_ID_POOL_H_ +#define PARLAY_INTERNAL_THREAD_ID_POOL_H_ + +#include +#include + +#include +#include +#include +#include + +namespace parlay { +namespace internal { + +using thread_id_type = unsigned int; + +// A ThreadIdPool hands out and maintains available unique dense IDs for active threads. +// Each thread that requests an ID will get one in the range from [0...get_num_thread_ids()). +// When the pool runs out of available IDs, it will allocate new ones, increasing the result +// of get_num_thread_ids(). Threads that die will return their ID to the pool for re-use by +// a subsequently spawned thread. +// +// There is a global singleton instance of ThreadIdPool given by ThreadIdPool::instance(), +// however this function is private and should not be called by the outside world. The public +// API through which the world can access thread IDs is limited to the free functions: +// +// - get_thread_id() -> size_t: Returns the thread ID of the current thread. Will assign +// one if this thread doesn't have one yet. +// - get_num_thread_ids() -> size_t: Returns the number of unique thread IDs that have +// been handed out. +// +class ThreadIdPool : public std::enable_shared_from_this { + + // Prevent public construction since this class is meant as a global singleton + struct private_constructor { + explicit private_constructor() = default; + }; + + public: + + // Returns a unique thread ID for the current thread in the range [0...get_num_thread_ids()) + friend thread_id_type get_thread_id(); + + // Returns the number of assigned thread IDs in the range [0...get_num_thread_ids()) + friend thread_id_type get_num_thread_ids(); + + + ~ThreadIdPool() noexcept { + [[maybe_unused]] size_t num_destroyed = 0; + for (auto current = available_ids.load(std::memory_order_relaxed); current; num_destroyed++) { + auto old = std::exchange(current, current->next); + delete old; + } + assert(num_destroyed == num_thread_ids.load(std::memory_order_relaxed)); + } + + // The constructor must be public since we make_shared it, but we protect it with a private parameter type + explicit ThreadIdPool(private_constructor) noexcept : num_thread_ids(0), available_ids(nullptr) { } + + ThreadIdPool(const ThreadIdPool&) = delete; + ThreadIdPool& operator=(const ThreadIdPool&) = delete; + + private: + + // A ThreadId corresponds to a unique ID number in the range [0...num_thread_ids). When it is + // not in use (the thread that owned it dies), it is returned to the global pool which maintains + // a linked list of available ones. + class ThreadId { + friend class ThreadIdPool; + + explicit ThreadId(const thread_id_type id_) noexcept : id(id_), next(nullptr) { } + + public: + const thread_id_type id; + private: + ThreadId* next; + }; + + // A ThreadIdOwner indicates that a thread is currently in possession of the given ThreadID. + // Each thread has a static thread_local ThreadIdOwner containing the ID that it owns. + // On construction, it acquires an available ThreadID, and on destruction, it releases + // it back to the pool. The ThreadIdOwner stores a shared_ptr to the pool to guarantee + // that the pool does not get destroyed before a detached thread returns its ID. + class ThreadIdOwner { + friend class ThreadIdPool; + + explicit ThreadIdOwner(ThreadIdPool& pool_) + : pool(pool_.shared_from_this()), node(pool->acquire()), id(node->id) { } + + ~ThreadIdOwner() { pool->relinquish(node); } + + private: + const std::shared_ptr pool; + ThreadId* const node; + + public: + const thread_id_type id; + }; + + // Grab a free ID from the available list, or if there are none available, allocate a new one. + ThreadId* acquire() { + if (available_ids.load(std::memory_order_relaxed)) { + // We only take the lock if there are available IDs in the pool. In the common case + // where there are no relinquished IDs available for re-use we don't need the lock. + static std::mutex m_; + std::lock_guard g_{m_}; + + ThreadId* current = available_ids.load(std::memory_order_relaxed); + while (current && !available_ids.compare_exchange_weak(current, current->next, + std::memory_order_acquire, std::memory_order_relaxed)) {} + if (current) { return current; } + } + return new ThreadId(num_thread_ids.fetch_add(1)); + } + + // Given the ID back to the global pool for reuse + void relinquish(ThreadId* p) { + p->next = available_ids.load(std::memory_order_relaxed); + while (!available_ids.compare_exchange_weak(p->next, p, + std::memory_order_release, std::memory_order_relaxed)) {} + } + + static inline const ThreadIdOwner& get_local_thread_id() { + static const thread_local ThreadIdPool::ThreadIdOwner my_id(instance()); + return my_id; + } + + static inline ThreadIdPool& instance() { + // We hold the global thread id pool inside a shared_ptr because it is possible + // for threads to be spawned *before* the ID pool has been initialized, which + // means that they may outlive this static variable. Each ThreadId holds onto + // a copy of the shared_ptr to ensure that the pool stays alive long enough + // for the IDs to relinquish themselves back to the pool. + // + // I think it is still possible to cause a segfault by spawning a new thread + // *after* the static destructors have run... so please do not spawn threads + // inside your static destructors :) + static const std::shared_ptr pool = std::make_shared(private_constructor{}); + return *pool; + } + + std::atomic num_thread_ids; + std::atomic available_ids; +}; + +inline thread_id_type get_thread_id() { + return ThreadIdPool::get_local_thread_id().id; +} + +inline thread_id_type get_num_thread_ids() { + return ThreadIdPool::instance().num_thread_ids.load(); +} + + +} // namespace internal +} // namespace parlay + + +#endif // PARLAY_INTERNAL_THREAD_ID_POOL_H_ diff --git a/cbits/parlaylib/include/parlay/portability.h b/cbits/parlaylib/include/parlay/portability.h new file mode 100644 index 000000000..88a371889 --- /dev/null +++ b/cbits/parlaylib/include/parlay/portability.h @@ -0,0 +1,112 @@ + +#ifndef PARLAY_PORTABILITY_H_ +#define PARLAY_PORTABILITY_H_ + +#if defined(_WIN32) +#ifndef NOMINMAX +#define PARLAY_DEFINED_NOMINMAX +#define NOMINMAX +#endif +#include +#ifdef PARLAY_DEFINED_NOMINMAX +#undef NOMINMAX +#endif +#endif + +#include + +#include + +namespace parlay { + +// PARLAY_INLINE: Ask the compiler politely to inline the given function. +#if defined(__GNUC__) +#define PARLAY_INLINE inline __attribute__((__always_inline__)) +#elif defined(_MSC_VER) +#define PARLAY_INLINE __forceinline +#else +#define PARLAY_INLINE inline +#endif + +// PARLAY_NOINLINE: Ask the compiler to *not* inline the given function +#if defined(__GNUC__) +#define PARLAY_NOINLINE __attribute__((__noinline__)) +#elif defined(_MSC_VER) +#define PARLAY_NOINLINE __declspec(noinline) +#else +#define PARLAY_NOINLINE +#endif + +// PARLAY_COLD: Ask the compiler to place the given function far away from other code +#if defined(__GNUC__) +#define PARLAY_COLD __attribute__((__cold__)) +#elif defined(_MSC_VER) +#define PARLAY_COLD +#else +#define PARLAY_COLD +#endif + + +// PARLAY_PACKED: Ask the compiler to pack a struct into less memory by not padding +#if defined(__GNUC__) +#define PARLAY_PACKED __attribute__((packed)) +#else +#define PARLAY_PACKED +#endif + +// PARLAY_NO_UNIQUE_ADDR: Allow a member object to occupy no space +#if defined(__has_cpp_attribute) +#if __has_cpp_attribute(no_unique_address) +#define PARLAY_NO_UNIQUE_ADDR [[no_unique_address]] +#else +#define PARLAY_NO_UNIQUE_ADDR +#endif +#else +#define PARLAY_NO_UNIQUE_ADDR +#endif + +// PARLAY_PREFETCH: Prefetch data into cache +#if defined(__GNUC__) +#define PARLAY_PREFETCH(addr, rw, locality) __builtin_prefetch ((addr), (rw), (locality)) +#elif defined(_MSC_VER) +#define PARLAY_PREFETCH(addr, rw, locality) \ + PreFetchCacheLine(((locality) ? PF_TEMPORAL_LEVEL_1 : PF_NON_TEMPORAL_LEVEL_ALL), (addr)) +#else +#define PARLAY_PREFETCH(addr, rw, locality) +#endif + + +#if defined(__cplusplus) && __cplusplus >= 202002L +#define PARLAY_LIKELY [[likely]] +#define PARLAY_UNLIKELY [[unlikely]] +#else +#define PARLAY_LIKELY +#define PARLAY_UNLIKELY +#endif + +// Check for exceptions. The standard suggests __cpp_exceptions. Clang/GCC defined __EXCEPTIONS. +// MSVC disables them with _HAS_EXCEPTIONS=0. Might not cover obscure compilers/STLs. +// +// Exceptions can be explicitly disabled in Parlay with PARLAY_NO_EXCEPTIONS. +#if !defined(PARLAY_NO_EXCEPTIONS) && \ + ((defined(__cpp_exceptions) && __cpp_exceptions != 0) || \ + (defined(__EXCEPTIONS)) || \ + (defined(_HAS_EXCEPTIONS) && _HAS_EXCEPTIONS == 1) || \ + (defined(_MSC_VER) && !defined(_HAS_EXCEPTIONS))) +#define PARLAY_EXCEPTIONS_ENABLED +#endif + +template +[[noreturn]] PARLAY_NOINLINE PARLAY_COLD void throw_exception_or_terminate(Args&&... args) { +#if defined(PARLAY_EXCEPTIONS_ENABLED) + throw Exception{std::forward(args)...}; +#else + std::cerr << Exception{std::forward(args)...}.what() << "\n"; + std::terminate(); +#endif +} + + +} // namespace parlay + +#endif // PARLAY_PORTABILITY_H_ diff --git a/cbits/parlaylib/include/parlay/range.h b/cbits/parlaylib/include/parlay/range.h new file mode 100644 index 000000000..387bdeed3 --- /dev/null +++ b/cbits/parlaylib/include/parlay/range.h @@ -0,0 +1,598 @@ +// Type traits and concepts for range properties +// +// This file defines C++17 "concepts-lite" style traits +// for detecting properties of ranges, and for determining +// whether given types are a particular type of range. +// +// The concept traits do not exhaustively check that all +// requirements of the concept are satisfied, just most of +// the ones that are do-able enough in C++17 without +// going overboard. They should be more than thorough +// enough to use for SFINAE, and pretty good for sanity +// checking that an implementation of a container or range +// adapter has implemented the required features. +// +// Range Traits: +// range_iterator_type : the type of a range's iterator, i.e. the type of std::begin(r) +// range_sentinel_type : the type of a range's sentinel, i.e., the type of std::end(r) +// range_value_type : the value type of the contents of a range +// range_reference_type : the type obtained by dereferencing a range's iterator +// +// Iterator Concepts: +// is_iterator : true if the given type is an iterator +// is_sentinel_for : true if the second type is a sentinel for the first (iterator) type +// is_contiguous_iterator : true if the given type is a contiguous iterator +// is_random_access_iterator : true if the given type is a random access iterator +// is_bidirectional_iterator : true if the given type is a bidirectional iterator +// is_forward_iterator : true if the given type is a forward iterator +// is_input_iterator : true if the given type is an input iterator +// is_output_iterator : true if the given type is an output iterator +// +// Range Concepts: +// is_range : true if the given type is a range +// is_common_range : true if the given type is a common range (iterator and sentinel are the same type) +// is_contiguous_range : true if the given type is a contiguous range +// is_random_access_range : true if the given type is a random access range +// is_bidirectional_range : true if the given type is a bidirectional range +// is_forward_range : true if the given type is a forward range +// is_input_ranger : true if the given type is an input range +// is_output_range : true if the given type is an output range +// is_block_iterable_range : true if the given type is a block-iterable range +// +// Range Operations: +// size : return the size of a range +// + +#ifndef PARLAY_RANGE_H_ +#define PARLAY_RANGE_H_ + +#include + +#include // IWYU pragma: keep +#include +#include + +#include "type_traits.h" + +// Old macros that we used to use to conditionally +// insert type constraints if concepts were enabled. +// Here for backwards compatibility. +#define PARLAY_RANGE_TYPE typename +#define PARLAY_RANGE auto + +namespace parlay { + +/* ----------------------------- Iterator traits -------------------------------- + Type traits that deduce useful types about an iterator. + + iterator_category : the category tag type of the iterator + iterator_value_type : the value type of the iterator + iterator_reference_type : the type obtained by dereferencing an iterator + iterator_difference_type : a type that can represent the difference between two iterators +*/ + +// Deduce the iterator category of the given iterator type +template +using iterator_category = type_identity::iterator_category>; + +template +using iterator_category_t = typename iterator_category::type; + +// Deduce the value type of the given iterator type +template +using iterator_value_type = type_identity::value_type>; + +// Deduce the value type of the given iterator type +template +using iterator_value_type_t = typename iterator_value_type::type; + +// Deduce the reference type of the given iterator type +template +using iterator_reference_type = type_identity::reference>; + +// Deduce the reference type of the given iterator type +template +using iterator_reference_type_t = typename iterator_reference_type::type; + +// Deduce the difference type of the given iterator type +template +using iterator_difference_type = type_identity::difference_type>; + +// Deduce the difference type of the given iterator type +template +using iterator_difference_type_t = typename iterator_difference_type::type; + +/* ----------------------------- Range traits -------------------------------- + Type traits that deduce useful types about a range. + + range_iterator_type : the type of a range's iterator, i.e. the type of std::begin(r) + range_sentinel_type : the type of a range's sentinel, i.e., the type of std::end(r) + range_value_type : the value type of the contents of a range + range_reference_type : the type obtained by dereferencing a range's iterator + range_difference_type : a type that can represent the difference between two iterators +*/ + +// Deduce the iterator type of the range type T +template +using range_iterator_type = type_identity()))>; + +// Deduce the iterator type of the range type T +template +using range_iterator_type_t = typename range_iterator_type::type; + +// Deduce the sentinel (end iterator) type of the range type T +template +using range_sentinel_type = type_identity()))>; + +// Deduce the sentinel (end iterator) type of the range type T +template +using range_sentinel_type_t = typename range_sentinel_type::type; + +// Deduce the underlying value type of a range. This should correspond to a +// type that can be used to safely copy a value obtained by one of its iterators +template +using range_value_type = type_identity>>; + +// Deduce the underlying value type of a range. This should correspond to a +// type that can be used to safely copy a value obtained by one of its iterators +template +using range_value_type_t = typename range_value_type::type; + +// Deduce the reference type of a range. This is the type obtained by dereferencing one of its iterators. +template +using range_reference_type = type_identity>>; + +// Deduce the reference type of a range. This is the type obtained by dereferencing one of its iterators. +template +using range_reference_type_t = typename range_reference_type::type; + +// Deduce the difference type of a range. This is a type that can be used to +// represent a difference between two iterators +template +using range_difference_type = type_identity>>; + +// Deduce the difference type of a range. This is a type that can be used to +//// represent a difference between two iterators +template +using range_difference_type_t = typename range_difference_type::type; + +// --------------------- Iterator concept traits ----------------------- + +// Provides the member value true if It_ is an iterator type +template> +struct is_iterator : public std::false_type {}; + +template +struct is_iterator, + decltype( ++std::declval() ), + decltype( *std::declval() ), + std::enable_if_t< std::is_swappable_v >, + std::enable_if_t< std::is_same_v()), It_&> >, + std::enable_if_t< std::is_copy_constructible_v >, + std::enable_if_t< std::is_copy_assignable_v >, + std::enable_if_t< std::is_destructible_v > +>> : public std::true_type {}; + +// True if It_ is an iterator type +template +inline constexpr bool is_iterator_v = is_iterator::value; + +// Provides the member value true if Sentinel_ is a sentinel type +// for the iterator type It_ +template> +struct is_sentinel_for : public std::false_type {}; + +template +struct is_sentinel_for >, + std::enable_if_t< is_equality_comparable_v >, + std::enable_if_t< std::is_default_constructible_v >, + std::enable_if_t< std::is_copy_constructible_v >, + std::enable_if_t< std::is_copy_assignable_v > +>> : public std::true_type {}; + +// True if Sentinel_ is a sentinel type for the iterator It_ +template +inline constexpr bool is_sentinel_for_v = is_sentinel_for::value; + +// Provides the member value true if Sentinel_ is a sized sentinel type +// for the iterator type It_ +template> +struct is_sized_sentinel_for : public std::false_type {}; + +template +struct is_sized_sentinel_for >, + std::enable_if_t< std::is_same_v() - std::declval()), iterator_difference_type_t> >, + std::enable_if_t< std::is_same_v() - std::declval()), iterator_difference_type_t> > +>> : std::true_type {}; + +// True if Sentinel_ is a sized sentinel type for the iterator It_ +template +inline constexpr bool is_sized_sentinel_for_v = is_sized_sentinel_for::value; + +// Defines a member constant value true if the iterator It_ is at least an input iterator +template> +struct is_input_iterator : public std::false_type {}; + +template +struct is_input_iterator()++ ), + std::enable_if_t< is_iterator_v >, + std::enable_if_t< is_equality_comparable_v >, + std::enable_if_t< std::is_base_of_v> >, + std::enable_if_t< std::is_same_v()), iterator_reference_type_t> > +>> : public std::true_type {}; + +// true if the iterator type It_ is at least an input iterator +template +inline constexpr bool is_input_iterator_v = is_input_iterator::value; + +// Defines a member constant value true if the iterator It_ is an output iterator +template> +struct is_output_iterator : public std::false_type {}; + +template +struct is_output_iterator >, + std::enable_if_t< std::is_base_of_v> > +>> : public std::true_type {}; + +// true if the iterator type It_ is an output iterator +template +inline constexpr bool is_output_iterator_v = is_output_iterator::value; + +// Defines a member constant value true if the iterator It_ is at least a forward iterator +template> +struct is_forward_iterator : public std::false_type {}; + +template +struct is_forward_iterator >, + std::enable_if_t< std::is_base_of_v> >, + std::enable_if_t< std::is_default_constructible_v >, + std::enable_if_t< std::is_convertible_v()++), const std::remove_reference_t&> >, + std::enable_if_t< std::is_same_v()++)), iterator_reference_type_t> > +>> : public std::true_type {}; + +// true if the iterator type It_ is at least a forward iterator +template +inline constexpr bool is_forward_iterator_v = is_forward_iterator::value; + +// Defines a member constant value true if the iterator It_ is at least a bidirectional iterator +template> +struct is_bidirectional_iterator : public std::false_type {}; + +template +struct is_bidirectional_iterator() ), + decltype( std::declval()-- ), + std::enable_if_t< is_forward_iterator_v >, + std::enable_if_t< std::is_base_of_v> >, + std::enable_if_t< std::is_convertible_v()--), const std::remove_reference_t&> >, + std::enable_if_t< std::is_same_v()--)), iterator_reference_type_t> > +>> : public std::true_type {}; + +// true if the iterator type It is at least a bidirectional iterator +template +inline constexpr bool is_bidirectional_iterator_v = is_bidirectional_iterator::value; + +// Defines a member constant value true if the iterator It_ is at least a random-access iterator +template> +struct is_random_access_iterator : public std::false_type {}; + +template +struct is_random_access_iterator >, + std::enable_if_t< std::is_base_of_v> >, + std::enable_if_t< std::is_same_v() += std::declval>()), It_&> >, + std::enable_if_t< std::is_same_v() + std::declval>()), It_> >, + std::enable_if_t< std::is_same_v() -= std::declval>()), It_&> >, + std::enable_if_t< std::is_same_v() - std::declval>()), It_> >, + std::enable_if_t< std::is_same_v() - std::declval()), iterator_difference_type_t> >, + std::enable_if_t< std::is_same_v()[std::declval>()]), iterator_reference_type_t> >, + std::enable_if_t< std::is_convertible_v() < std::declval()), bool> >, + std::enable_if_t< std::is_convertible_v() > std::declval()), bool> >, + std::enable_if_t< std::is_convertible_v() >= std::declval()), bool> >, + std::enable_if_t< std::is_convertible_v() <= std::declval()), bool> > +>> : public std::true_type {}; + +// true if the iterator type It_ is at least a random-access iterator +template +inline constexpr bool is_random_access_iterator_v = is_random_access_iterator::value; + +// Defines a member constant value true if the iterator type It_ is a contiguous iterator +template +using is_contiguous_iterator = std::is_pointer; + +// true if the iterator type It_ corresponds to a contiguous iterator type +template +inline constexpr bool is_contiguous_iterator_v = is_contiguous_iterator::value; + + +// --------------------- Range concept traits ----------------------- + +// Provides the member value true if Range_ is a range +template> +struct is_range : public std::false_type {}; + +template +struct is_range()) ), + decltype( std::end(std::declval()) ), + std::enable_if_t< is_sentinel_for_v, range_sentinel_type_t> > +>> : public std::true_type {}; + +// True if the given type Range_ is a range +template +inline constexpr bool is_range_v = is_range::value; + +// Provides the member value true if Range_ is a common range, +// i.e., a range whose iterator and sentinel types are the same +template> +struct is_common_range : public std::false_type {}; + +template +struct is_common_range >, + std::enable_if_t< std::is_same_v, range_sentinel_type_t> > +>> : public std::true_type {}; + +// True if the given type is a common range, i.e., a range +// whose iterator and sentinel types are the same +template +inline constexpr bool is_common_range_v = is_common_range::value; + +// Defines a member constant value true if the range type Range_ is an input range +template> +struct is_input_range : public std::false_type {}; + +template +struct is_input_range >, + std::enable_if_t< is_input_iterator_v> > +>> : public std::true_type {}; + +// true if the range type Range_ is an input range +template +inline constexpr bool is_input_range_v = is_input_range::value; + +// Defines a member constant value true if the range type Range_ is an output range +template> +struct is_output_range : public std::false_type {}; + +template +struct is_output_range >, + std::enable_if_t< is_output_iterator_v> > +>> : public std::true_type {}; + +// true if the range type Range_ is an output range +template +inline constexpr bool is_output_range_v = is_output_range::value; + +// Defines a member constant value true if the range type Range_ is a forward range +template> +struct is_forward_range : public std::false_type {}; + +template +struct is_forward_range >, + std::enable_if_t< is_forward_iterator_v> > +>> : public std::true_type {}; + +// true if the range type Range_ is a forward range +template +inline constexpr bool is_forward_range_v = is_forward_range::value; + +// Defines a member constant value true if the range type Range_ is a bidirectional range +template> +struct is_bidirectional_range : public std::false_type {}; + +template +struct is_bidirectional_range >, + std::enable_if_t< is_bidirectional_iterator_v> > +>> : public std::true_type {}; + +// true if the range type Range_ is a bidirectional range +template +inline constexpr bool is_bidirectional_range_v = is_bidirectional_range::value; + +// Defines a member constant value true if the range type Range_ is a random-access range +template> +struct is_random_access_range : public std::false_type {}; + +template +struct is_random_access_range >, + std::enable_if_t< is_random_access_iterator_v> > +>> : public std::true_type {}; + +// true if the range type Range_ is a random-access range +template +inline constexpr bool is_random_access_range_v = is_random_access_range::value; + +// Defines a member constant value true if the range Range_ is a contiguous range +template> +struct is_contiguous_range : public std::false_type {}; + +template +struct is_contiguous_range >, + std::enable_if_t< is_contiguous_iterator_v> > +>> : public std::true_type {}; + +// true if the range type Range_ corresponds to a contiguous range +template +inline constexpr bool is_contiguous_range_v = is_contiguous_range::value; + +namespace internal { + +// A type satisfies the block-iterable interface if it supports: +// +// size() const -> size_t +// get_num_blocks() const -> size_t +// get_begin_block(size_t) -> iterator +// get_end_block(size_t) -> iterator +// get_begin_block(size_t) const -> const_iterator +// get_end_block(size_t) const -> const_iterator +// +// where iterator and const_iterator is the common iterator type for the range when +// it is non-const or const respectively. iterator and const_iterator may be the same +// if that is semantically appropriate. +// +// Note that the const overloads are optional if the view can not be const. +// +template> +struct has_block_iterable_interface : public std::false_type {}; + +template +struct has_block_iterable_interface().get_begin_block((size_t)0) ), + decltype( std::declval().get_end_block((size_t)0) ), + decltype( std::declval().size() ), + decltype( std::declval().get_num_blocks() ), + std::enable_if_t< std::is_same_v().get_begin_block((size_t)0)), range_iterator_type_t> >, + std::enable_if_t< std::is_same_v().get_end_block((size_t)0)), range_iterator_type_t> >, + std::enable_if_t< std::is_convertible_v().size()), size_t> >, + std::enable_if_t< std::is_convertible_v().get_num_blocks()), size_t> > +>> : public std::true_type {}; + +template +inline constexpr bool has_block_iterable_interface_v = has_block_iterable_interface::value; + +} // namespace internal + +// Provides a member constant value true if the range Range is block iterable. +// A range is block-iterable if it is a common range, and is either random access +// or it supports the block-iterable interface, meaning it provides the member +// functions size(), get_num_blocks(), get_begin_block(size_t), get_end_block(size_t) +template> +struct is_block_iterable_range : public std::false_type {}; + +template +struct is_block_iterable_range >, + std::enable_if_t< is_random_access_range_v || internal::has_block_iterable_interface_v > +>> : public std::true_type {}; + +// true if the range Range_ is block iterable +template +inline constexpr bool is_block_iterable_range_v = is_block_iterable_range::value; + +/* --------------------- Range size ----------------------- + size(r) -> size_t : returns the size of a range +*/ + +// Defines the member value true if the given type is an array of known bound +template +struct is_bounded_array: std::false_type {}; + +template +struct is_bounded_array : std::true_type {}; + +// true if the given type is an array of known bound +template +inline constexpr bool is_bounded_array_v = is_bounded_array::value; + +namespace internal { + +// Defines the member value true if the given type has a member function size() +// which takes no arguments +template> +struct has_size_member : public std::false_type {}; + +template +struct has_size_member().size() ) +>> : std::true_type {}; + +// true if the given type has a member function size() which takes no arguments +template +inline constexpr bool has_size_member_v = has_size_member::value; + +namespace nonmember_size_impl { + +// Deleting the size function in this scope forces overload +// resolution to look for ADL matches for the size function + +template +void size(T&) = delete; + +template +void size(const T&) = delete; + +// Defines the member value true if the given type has a non- +// member function size(r) that can be found via ADL +template> +struct has_nonmember_size : public std::false_type {}; + +template +struct has_nonmember_size()) ) +>> : std::true_type {}; + +// true if the given type has a non-member function size(r) that can be found via ADL +template +inline constexpr bool has_nonmember_size_v = has_nonmember_size::value; + +// Invokes the non-member size function for the given range +template +auto invoke_size(Range_&& r) { + static_assert(has_nonmember_size_v); + return size(std::forward(r)); +} + +} // namespace nonmember_size_impl + +// The implementation of parlay::size is hidden inside a namespace +// and pulled in with a using directive to hide it from ADL +namespace size_impl { + +// Return the size (number of elements) of the range r. This is: +// +// * std::extent_v if Range_ is a bounded array type +// * FORWARD(r).size() if Range_ is a type with a .size() member +// * size(FORWARD(r)) if a non-member size function is found by ADL +// * std::end(r) - std::begin(r) if begin(r) and end(r) model a sized-sentinel pair +// +// In the future, this could be replaced by C++20 std::ranges::size +template +auto size(Range_&& r) { + static_assert(is_range_v); + static_assert(is_bounded_array_v || + has_size_member_v || + nonmember_size_impl::has_nonmember_size_v || + is_sized_sentinel_for_v, range_sentinel_type_t>); + + if constexpr (is_bounded_array_v) { + return std::extent_v; + } + else if constexpr (has_size_member_v) { + return std::forward(r).size(); + } + else if constexpr (nonmember_size_impl::has_nonmember_size_v) { + return nonmember_size_impl::invoke_size(std::forward(r)); + } + else { + return std::make_unsigned_t>(std::end(r) - std::begin(r)); + } +} + +} // size_impl + +} // namespace internal + +using internal::size_impl::size; + +// A functional that takes a range r and returns its size as given by parlay::size(FORWARD(r)) +struct size_of { + template + decltype(auto) operator()(Range_&& r) const { return parlay::size(std::forward(r)); } +}; + +} // namespace parlay + +#endif // PARLAY_RANGE_H_ diff --git a/cbits/parlaylib/include/parlay/thread_specific.h b/cbits/parlaylib/include/parlay/thread_specific.h new file mode 100644 index 000000000..0215b1289 --- /dev/null +++ b/cbits/parlaylib/include/parlay/thread_specific.h @@ -0,0 +1,464 @@ + +#ifndef PARLAY_THREAD_SPECIFIC_H_ +#define PARLAY_THREAD_SPECIFIC_H_ + +#include +#include + +#include +#include // IWYU pragma: keep +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "internal/thread_id_pool.h" + +#include "portability.h" +#include "range.h" +#include "type_traits.h" + + +namespace parlay { + +using internal::thread_id_type; + +// Returns a unique thread ID for the current running thread +// in the range of [0...num_thread_ids()). Thread IDs are +// guaranteed to be unique for all *live* threads, but they +// are re-used after a thread dies and another is spawned. +inline thread_id_type my_thread_id() { + return internal::get_thread_id(); +} + +// Return the number of thread IDs that have been assigned to +// threads. All thread IDs are in the range [0...num_thread_ids()). +// +// Important note: Thread IDs are assigned lazily when a thread +// first requests one. Therefore num_thread_ids() is *not* +// guaranteed to be as large as the number of live threads if +// those threads have never called my_thread_id(). +inline thread_id_type num_thread_ids() { + return internal::get_num_thread_ids(); +} + +namespace internal { + +class ThreadListChunkData { + + public: + + // This is just std::bit_ceil(std::thread::hardware_concurrency()) but we don't assume C++20 + const static inline std::size_t thread_list_chunk_size = []() { + std::size_t size = 4; + while (size < std::thread::hardware_concurrency()) + size *= 2; + return size; + }(); + + // Used by ThreadSpecific which stores a chunked sequence of items that is at least as large + // as the number of active threads. Given a thread ID, items are split into chunks of size: + // + // P, P, 2P, 4P, 8P, ... + // + // where P is the lowest power of two that is at least as large as the number of hardware threads. + static std::size_t compute_chunk_id(thread_id_type id) { + std::size_t k = thread_list_chunk_size; + std::size_t chunk = 0; + while (k <= id) { + chunk++; + k *= 2; + } + return chunk; + } + + static std::size_t compute_chunk_position(thread_id_type id, std::size_t chunk_id) { + if (chunk_id == 0) + return id; + else { + auto high_bit = thread_list_chunk_size << (chunk_id - 1); + assert(id & high_bit); + return id - high_bit; + } + } + + explicit ThreadListChunkData(thread_id_type thread_id_) noexcept : thread_id(thread_id_), + chunk_id(compute_chunk_id(thread_id)), chunk_position(compute_chunk_position(thread_id, chunk_id)) { } + + const thread_id_type thread_id; + const std::size_t chunk_id; + const std::size_t chunk_position; +}; + +extern inline const ThreadListChunkData& get_chunk_data() { + static thread_local const ThreadListChunkData data{get_thread_id()}; + return data; +} + +template +struct Uninitialized { + union { + alignas(64) std::monostate empty; + T value; + }; + + Uninitialized() noexcept { } + + T& operator*() { return value; } + + T* get() { return std::addressof(value); } + + ~Uninitialized() { value.~T(); } +}; + +} // namespace internal + +// A ThreadSpecific stores a list of objects of type T such that there +// is a unique object for each active thread. The list automatically grows +// when additional threads are spawned and attempt to access it. Threads +// may also traverse the entire list if they need to. +// +// By default, list elements are all value initialized, roughly meaning +// that class types are default constructed, and builtin types are zero +// initialized. For custom initialization, you can pass a constructor +// function which returns the desired value. The constructor function +// can take zero or one arguments. If it takes one argument, it will be +// passed the thread ID that it is constructing for. Note that the +// elements are not guaranteed to be constructed by the thread that +// they belong to, and they may be constructed in advance of any thread +// actually taking ownership of that ID. +// +// A few things to note: +// +// - Thread IDs are always unique for the set of currently live threads, +// but not unique over the course of the entire program. A thread that +// dies will give up its ID to be claimed by a new thread later. +// +// - The list elements are *not* destroyed when the thread that "owns" +// them is destroyed. A new thread that reclaims a previously-used ID +// will find the item at that position in the same state that it was +// left by the previous thread. Elements are only destroyed when the +// entire ThreadSpecific is destroyed. +// +// Therefore, threads are responsible for manually cleaning up the +// contents of a ThreadSpecific and/or resetting it to a default value +// for the next thread that might claim the spot if they need to. +// +template +class ThreadSpecific { + + // 25 chunks guarantees enough slots for any machine + // with up to 2^48 bytes of addressable virtual memory, + // assuming that threads are 8MB large. + static constexpr std::size_t n_chunks = 25; + + public: + + using reference = T&; + using value_type = T; + using size_type = std::size_t; + using difference_type = std::ptrdiff_t; + using pointer = T*; + + ThreadSpecific() : constructor([](std::size_t) { return T{}; }) { + initialize(); + } + + template && !std::is_invocable_v, int> = 0> + explicit ThreadSpecific(F&& constructor_) + : constructor([f = std::forward(constructor_)](std::size_t) { return f(); }) { + initialize(); + } + + template, int> = 0> + explicit ThreadSpecific(F&& constructor_) : constructor(std::forward(constructor_)) { + initialize(); + } + + ThreadSpecific(const ThreadSpecific&) = delete; + ThreadSpecific& operator=(const ThreadSpecific&) = delete; + ThreadSpecific(ThreadSpecific&&) = delete; + + ~ThreadSpecific() { + for (internal::Uninitialized* chunk : chunks) { + delete[] chunk; + } + } + + T& operator*() { return get(); } + T* operator->() { return std::addressof(get()); } + + T& get() { + auto chunk_data = internal::get_chunk_data(); + return get_by_index(chunk_data.chunk_id, chunk_data.chunk_position); + } + + const T& operator*() const { return get(); } + T const* operator->() const { return std::addressof(get()); } + + const T& get() const { + auto chunk_data = internal::get_chunk_data(); + return get_by_index(chunk_data.chunk_id, chunk_data.chunk_position); + } + + template + void for_each(F&& f) { + static_assert(std::is_invocable_v); + + auto num_threads = num_thread_ids(); + thread_id_type tid = 0; + internal::Uninitialized* chunk = chunks[0].load(std::memory_order_relaxed); + + for (std::size_t chunk_id = 0; tid < num_threads; chunk = chunks[++chunk_id].load(std::memory_order_acquire)) { + auto chunk_size = get_chunk_size(chunk_id); + if (!chunk) PARLAY_UNLIKELY { + ensure_chunk_exists(chunk_id); + chunk = chunks[chunk_id].load(std::memory_order_relaxed); + } + for (std::size_t i = 0; tid < num_threads && i < chunk_size; i++, tid++) { + f(*chunk[i]); + } + } + } + + // Allow looping over all thread's data + template + class iterator_t { + friend class ThreadSpecific; + + using parent_type = maybe_const_t>; + + iterator_t(std::size_t chunk_id_, std::size_t position_, parent_type* parent_) : + chunk_id(chunk_id_), position(position_), parent(parent_) { } + + public: + using iterator_category = std::random_access_iterator_tag; + using reference = std::add_lvalue_reference_t>; + using value_type = T; + using size_type = std::size_t; + using difference_type = std::ptrdiff_t; + using pointer = std::add_pointer_t>; + + iterator_t() = default; + + /* implicit */ iterator_t(const iterator_t& other) // cppcheck-suppress noExplicitConstructor // NOLINT + : chunk_id(other.chunk_id), position(other.position), parent(other.parent) { } + + reference operator*() const { return parent->get_by_index_nocheck(chunk_id, position); } + + reference operator[](std::size_t p) const { + auto tmp = *this; + tmp += p; + return *tmp; + } + + iterator_t& operator++() { + position++; + if (position == get_chunk_size(chunk_id)) { + if (++chunk_id < n_chunks && parent->chunks[chunk_id].load(std::memory_order_acquire) == nullptr) PARLAY_UNLIKELY + parent->ensure_chunk_exists(chunk_id); + position = 0; + } + return *this; + } + + iterator_t operator++(int) { auto tmp = *this; ++(*this); return tmp; } //NOLINT + + iterator_t& operator--() { + if (position == 0) { + position = get_chunk_size(--chunk_id) - 1; + if (parent->chunks[chunk_id].load(std::memory_order_acquire) == nullptr) PARLAY_UNLIKELY + parent->ensure_chunk_exists(chunk_id); + } + else { + position--; + } + return *this; + } + + iterator_t operator--(int) { auto tmp = *this; --(*this); return tmp; } //NOLINT + + iterator_t& operator+=(difference_type diff) { + if (diff < 0) return *this -= (-diff); + assert(diff >= 0); + position += diff; + if (position >= get_chunk_size(chunk_id)) { + do { + position -= get_chunk_size(chunk_id++); + } while (position >= get_chunk_size(chunk_id)); + if (parent->chunks[chunk_id].load(std::memory_order_acquire) == nullptr) PARLAY_UNLIKELY + parent->ensure_chunk_exists(chunk_id); + } + return *this; + } + + iterator_t operator+(difference_type diff) const { + auto result = *this; + result += diff; + return result; + } + + iterator_t& operator-=(difference_type diff) { + if (diff < 0) return *this += (-diff); + assert(diff >= 0); + auto pos = static_cast(position); + pos -= diff; + if (pos < 0) { + do { + pos += static_cast(get_chunk_size(--chunk_id)); + } while (pos < 0); + if (parent->chunks[chunk_id].load(std::memory_order_acquire) == nullptr) PARLAY_UNLIKELY + parent->ensure_chunk_exists(chunk_id); + } + assert(pos >= 0); + position = static_cast(pos); + return *this; + } + + iterator_t operator-(difference_type diff) const { + auto result = *this; + result -= diff; + return result; + } + + difference_type operator-(const iterator_t& other) const { + if (other > *this) return -(other - *this); + assert(other <= *this); + auto result = static_cast(position) - static_cast(other.position); + auto chunk_id_ = other.chunk_id; + while (chunk_id_ < chunk_id) { + result += static_cast(get_chunk_size(chunk_id_++)); + } + return result; + } + + bool operator==(const iterator_t& other) const { + return chunk_id == other.chunk_id && position == other.position; + } + + bool operator!=(const iterator_t& other) const { + return chunk_id != other.chunk_id || position != other.position; + } + + bool operator<(const iterator_t& other) const { + return chunk_id < other.chunk_id || (chunk_id == other.chunk_id && position < other.position); + } + + bool operator<=(const iterator_t& other) const { + return chunk_id < other.chunk_id || (chunk_id == other.chunk_id && position <= other.position); + } + + bool operator>(const iterator_t& other) const { + return chunk_id > other.chunk_id || (chunk_id == other.chunk_id && position > other.position); + } + + bool operator>=(const iterator_t& other) const { + return chunk_id > other.chunk_id || (chunk_id == other.chunk_id && position >= other.position); + } + + friend void swap(iterator_t& left, iterator_t& right) { + std::swap(left.chunk_id, right.chunk_id); + std::swap(left.position, right.position); + std::swap(left.parent, right.parent); + } + + std::size_t chunk_id{n_chunks}; + std::size_t position{0}; + parent_type* parent{nullptr}; + }; + + using iterator = iterator_t; + using const_iterator = iterator_t; + + static_assert(is_random_access_iterator_v); + static_assert(is_random_access_iterator_v); + + [[nodiscard]] iterator begin() { + return iterator{0,0,this}; + } + + [[nodiscard]] const_iterator begin() const { + return const_iterator{0,0,this}; + } + + [[nodiscard]] iterator end() { + internal::ThreadListChunkData data{num_thread_ids()}; + return iterator{data.chunk_id, data.chunk_position, this}; + } + + [[nodiscard]] const_iterator end() const { + internal::ThreadListChunkData data{num_thread_ids()}; + return const_iterator{data.chunk_id, data.chunk_position, this}; + } + + private: + + void initialize() { + internal::get_chunk_data(); // Force static initialization before any ThreadLocals are constructed + chunks[0].store(new internal::Uninitialized[internal::ThreadListChunkData::thread_list_chunk_size], std::memory_order_relaxed); + std::fill(chunks.begin() + 1, chunks.end(), nullptr); + auto chunk = chunks[0].load(std::memory_order_relaxed); + for (std::size_t i = 0; i < internal::ThreadListChunkData::thread_list_chunk_size; i++) { + new (static_cast(chunk[i].get())) T(constructor(i)); + } + } + + static std::size_t get_chunk_size(std::size_t chunk_id) { + assert(chunk_id < n_chunks); + if (chunk_id == 0) return internal::ThreadListChunkData::thread_list_chunk_size; + else return internal::ThreadListChunkData::thread_list_chunk_size << (chunk_id - 1); + } + + T& get_by_index(std::size_t chunk_id, std::size_t chunk_position) { + if (chunk_id > 0 && chunks[chunk_id].load(std::memory_order_acquire) == nullptr) PARLAY_UNLIKELY + ensure_chunk_exists(chunk_id); + return get_by_index_nocheck(chunk_id, chunk_position); + } + + const T& get_by_index(std::size_t chunk_id, std::size_t chunk_position) const { + if (chunk_id > 0 && chunks[chunk_id].load(std::memory_order_acquire) == nullptr) PARLAY_UNLIKELY + ensure_chunk_exists(chunk_id); + return get_by_index_nocheck(chunk_id, chunk_position); + } + + T& get_by_index_nocheck(std::size_t chunk_id, std::size_t chunk_position) { + assert(chunks[chunk_id].load() != nullptr); + return *(chunks[chunk_id].load(std::memory_order_relaxed)[chunk_position]); + } + + const T& get_by_index_nocheck(std::size_t chunk_id, std::size_t chunk_position) const { + assert(chunks[chunk_id].load() != nullptr); + return *(chunks[chunk_id].load(std::memory_order_relaxed)[chunk_position]); + } + + void ensure_chunk_exists(std::size_t chunk_id) const { + std::lock_guard lock(growing_mutex); + if (chunks[chunk_id].load(std::memory_order_relaxed) == nullptr) { + auto chunk_size = get_chunk_size(chunk_id); + auto chunk = new internal::Uninitialized[chunk_size]; + for (std::size_t i = 0; i < chunk_size; i++) { + new (static_cast(chunk[i].get())) T(constructor(chunk_size + i)); + } + chunks[chunk_id].store(chunk, std::memory_order_release); + } + } + + mutable std::function constructor; + mutable std::mutex growing_mutex; + mutable std::array*>, n_chunks> chunks; +}; + +static_assert(is_random_access_range_v>); +static_assert(is_random_access_range_v>); + +} // namespace parlay + + +#endif // PARLAY_THREAD_SPECIFIC_H_ diff --git a/cbits/parlaylib/include/parlay/type_traits.h b/cbits/parlaylib/include/parlay/type_traits.h new file mode 100644 index 000000000..a75586eee --- /dev/null +++ b/cbits/parlaylib/include/parlay/type_traits.h @@ -0,0 +1,444 @@ +// Useful type traits used mostly internally by Parlay +// +// Many inspired by this video: +// - https://www.youtube.com/watch?v=MWBfmmg8-Yo +// +// Notable inclusions: +// - is_trivially_relocatable +// + +#ifndef PARLAY_TYPE_TRAITS_H_ +#define PARLAY_TYPE_TRAITS_H_ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include // IWYU pragma: keep +#include + +// IWYU pragma: no_include + +namespace parlay { + +// Provides the member type T +template +struct type_identity { + using type = T; +}; + +// Equal to the type T, i.e., the identity transformation +template +using type_identity_t = typename type_identity::type; + +// Provides the member type equivalent to T with its cv-ref qualifiers removed +template +struct remove_cvref : type_identity>> { }; + +// Equal to T with its cv-ref qualifiers removed +template +using remove_cvref_t = typename remove_cvref::type; + +// Given a pointer-to-member (object or function), returns +// the type of the class in which the member lives +template +struct member_pointer_class; + +template +struct member_pointer_class : public type_identity {}; + +template +using member_pointer_class_t = typename member_pointer_class::type; + +// Provides the member type std::add_const_t if Const is +// true, otherwise provides the member type T +template +using maybe_const = std::conditional, T>; + +// Adds const to the given type if Const is true +template +using maybe_const_t = typename maybe_const::type; + +// Provides the member type std::decay_t if Decay is +// true, otherwise provides the member type T +template +using maybe_decay = std::conditional, T>; + +// Decays the given type if Decay is true +template +using maybe_decay_t = typename maybe_decay::type; + +// Provides the member value true if the given type is an instance of std::optional +template +struct is_optional : std::false_type {}; + +template +struct is_optional> : std::true_type {}; + +// true if the given type is an instance of std::optional +template +inline constexpr bool is_optional_v = is_optional::value; + +template +using is_less_than_comparable = std::conjunction< + std::is_invocable_r, T, U>, + std::is_invocable_r, U, T> + >; + +template +inline constexpr bool is_less_than_comparable_v = is_less_than_comparable::value; + +template +using is_equality_comparable = std::conjunction< + std::is_invocable_r, T, U>, + std::is_invocable_r, U, T>, + std::is_invocable_r, T, U>, + std::is_invocable_r, U, T> + >; + +template +inline constexpr bool is_equality_comparable_v = is_equality_comparable::value; + +// Defines a member value true if the given type BinaryOperator_ can be invoked on types +// T1&& and T2 to yield a result of a type that is convertible to T1. +// +// This requirement corresponds to the needs of a left fold over the operator BinaryOperator_ +// with an identity and result type of T1, where the intermediate elements being reduced over +// are potentially of type T2. +template +struct is_binary_operator_for : public std::false_type {}; + +template +struct is_binary_operator_for >, + std::enable_if_t< std::is_invocable_r_v >, + std::enable_if_t< std::is_invocable_r_v >, + std::enable_if_t< std::is_invocable_r_v >, + std::enable_if_t< std::is_invocable_r_v > +>, std::enable_if_t>> : public std::true_type{}; + +// Handle the case where BinaryOperator_ is a member function pointer +template +struct is_binary_operator_for >, + std::enable_if_t< std::is_invocable_r_v&, T1&&, T1&&> >, + std::enable_if_t< std::is_invocable_r_v&, T1&&, T2> >, + std::enable_if_t< std::is_invocable_r_v&, T2, T2> >, + std::enable_if_t< std::is_invocable_r_v&, T2, T1&&> > +>, std::enable_if_t>> : public std::true_type{}; + +// True if the given type BinaryOperator_ can be invoked on types T1&& and T2 to yield a result +// of a type that is convertible to T1. T2 defaults to T1&& if not specified. +// +// This requirement corresponds to the needs of a left fold over the operator BinaryOperator_ +// with an identity and result type of T1, where the intermediate elements being reduced over +// are potentially of type T2. +template +inline constexpr bool is_binary_operator_for_v = is_binary_operator_for::value; + +// Defines the member value true if T is a pair or a tuple of length two +template +struct is_pair : public std::false_type {}; + +template +struct is_pair(std::declval()) ), + decltype( std::get<1>(std::declval()) ), + std::enable_if_t< 2 == std::tuple_size_v> > +>> : public std::true_type {}; + +// True if T is a pair or a tuple of length two +template +inline constexpr bool is_pair_v = is_pair::value; + +/* --------------------- Priority tags. ------------------------- + Priority tags are an easy way to force template resolution to + pick the "best" option in the presence of multiple valid + choices. It works because of the facts that priority_tag + is a subtype of priority_tag, and template resolution + will always pick the most specialised option when faced with + a choice, so it will prefer priority_tag over + priority_tag +*/ + +template +struct priority_tag : priority_tag {}; + +template<> +struct priority_tag<0> {}; + + +/* ----------------- Trivial allocators. --------------------- + Allocator-aware containers and algorithms need to know whether + they can construct/destruct objects directly inside memory given + to them by an allocator, or whether the allocator has custom + behaviour. Since some optimizations require us to circumvent + custom allocator behaviour, we need to detect when an allocator + does not do this. + + Specifically, an allocator-aware algorithm must construct objects + inside memory returned by an allocator by writing + + std::allocator_traits::construct(allocator, p, args); + + if the allocator type defines a method .construct, then this results + in forwarding the construction to that method. Otherwise, this just + results in a call to + + new (p) T(std::forward(args)...) + + If we wish to circumvent calling the constructor, for example, + for a trivially relocatable type in which we would prefer to + copy directly via memcpy, we must ensure that the allocator + does not have a custom .construct method. Otherwise, we can + not optimize, and must continue to use the allocator's own + construct method. + + The same discussion is true for destruction as well. + + See https://www.youtube.com/watch?v=MWBfmmg8-Yo for more info. +*/ + +namespace internal { + +// Detect the existence of the .destroy method of the type Alloc +template +auto trivial_allocator(Alloc& a, T* p, priority_tag<2>) + -> decltype(void(a.destroy(p)), std::false_type()); + +// Detect the existence of the .construct method of the type Alloc +template +auto trivial_allocator(Alloc& a, T* p, priority_tag<1>) + -> decltype(void(a.construct(p, std::declval())), std::false_type()); + +// By default, if no .construct or .destroy methods are found, assume +// that the allocator is trivial +template +auto trivial_allocator(Alloc& a, T* p, priority_tag<0>) + -> std::true_type; + +} // namespace internal + +template +struct is_trivial_allocator + : decltype(internal::trivial_allocator(std::declval(), nullptr, priority_tag<2>())) {}; + +template +inline constexpr bool is_trivial_allocator_v = is_trivial_allocator::value; + +// Manually specialize std::allocator since it is trivial, but +// some (maybe all?) implementations still provide a .construct +// and .destroy method anyway. +template +struct is_trivial_allocator, T> : std::true_type {}; + + +/* ----------------- Trivially relocatable. --------------------- + + A type T is called trivially relocatable if, given a pointer p + to an object of type T, and a pointer q to uninitialized memory + large enough for an object of type T, then + + ::new (voidify(*q)) T(std::move(*p)); + std::destroy_at(p); + + is equivalent to + + std::memcpy(p, q, sizeof(T)); + + Any type that is trivially copyable is trivially relocatable, but many types + that are not trivially copyable are trivially relocatable. User-defined types + that are not obviously trivially relocatable can be annotated as such (with + care). This is used to optimize parlay::sequence so that it can reallocate + much faster, and is also used by sorting algorithms for faster swapping of + ranges of elements and individual elements. + + See proposal P1144 for copious details: + https://www.open-std.org/jtc1/sc22/wg21/docs/papers/2023/p1144r9.html +*/ + +// Standard language support for trivially relocatable (currently P1144) +#if defined(__cpp_lib_trivially_relocatable) + +template +using is_trivially_relocatable = std::is_trivially_relocatable; + +// Clang's builtin with some (very limited) support for trivial relocatability +#elif defined(__has_builtin) +#if __has_builtin(__is_trivially_relocatable) + +template +struct is_trivially_relocatable : std::bool_constant<__is_trivially_relocatable(T)> { }; + +#define PARLAY_MUST_SPECIALIZE_IS_TRIVIALLY_RELOCATABLE + +#else + +template +struct is_trivially_relocatable : std::is_trivially_copyable { }; + +#define PARLAY_MUST_SPECIALIZE_IS_TRIVIALLY_RELOCATABLE + +#endif +#else + +template +struct is_trivially_relocatable : std::is_trivially_copyable { }; + +#define PARLAY_MUST_SPECIALIZE_IS_TRIVIALLY_RELOCATABLE + +#endif + +template +inline constexpr bool is_trivially_relocatable_v = is_trivially_relocatable::value; + +#if defined(__cpp_impl_trivially_relocatable) +#define PARLAY_TRIVIALLY_RELOCATABLE [[trivially_relocatable]] +#else +#define PARLAY_TRIVIALLY_RELOCATABLE +#endif + + +#if defined(PARLAY_MUST_SPECIALIZE_IS_TRIVIALLY_RELOCATABLE) + +#define PARLAY_ASSUME_TRIVIALLY_RELOCATABLE(...) \ + struct is_trivially_relocatable<__VA_ARGS__> : std::true_type {}; + +#define PARLAY_ASSUME_TRIVIALLY_RELOCATABLE_IF(cond, ...) \ + struct is_trivially_relocatable<__VA_ARGS__> : std::bool_constant {}; + + +// We list all the types that ought to be trivially relocatable under any +// sensible implementation. It is possible that some of these are wrong. + +template +PARLAY_ASSUME_TRIVIALLY_RELOCATABLE(std::allocator); + +template +PARLAY_ASSUME_TRIVIALLY_RELOCATABLE(std::shared_ptr); + +template +PARLAY_ASSUME_TRIVIALLY_RELOCATABLE(std::weak_ptr); + +template +PARLAY_ASSUME_TRIVIALLY_RELOCATABLE_IF((is_trivially_relocatable_v), std::optional); + +template +PARLAY_ASSUME_TRIVIALLY_RELOCATABLE_IF((is_trivially_relocatable_v && + is_trivially_relocatable_v::pointer>), std::unique_ptr); + +template +PARLAY_ASSUME_TRIVIALLY_RELOCATABLE_IF((is_trivially_relocatable_v && + is_trivially_relocatable_v), std::pair); + +template +PARLAY_ASSUME_TRIVIALLY_RELOCATABLE_IF((is_trivially_relocatable_v), std::array); + +template +PARLAY_ASSUME_TRIVIALLY_RELOCATABLE_IF((is_trivially_relocatable_v && ...), std::tuple); + +template +PARLAY_ASSUME_TRIVIALLY_RELOCATABLE_IF((is_trivially_relocatable_v && ...), std::variant); + + +// Container implementations can differ widely by vendor, so we don't want to specialize these +// with a broad brush. Instead, only specialize them for stdlibs that we are confident about, +// or, more specifically, stdlibs that Arthur is confident about :D +// +// https://quuxplusone.github.io/blog/2019/02/20/p1144-what-types-are-relocatable/ + +// Specializations for libc++. +#if defined(_LIBCPP_VERSION) + +template +PARLAY_ASSUME_TRIVIALLY_RELOCATABLE(std::deque); + +template +PARLAY_ASSUME_TRIVIALLY_RELOCATABLE(std::forward_list); + +// std::vector and std::string are not trivially relocatable in libc++ debug mode +#if !defined(_LIBCPP_DEBUG_MODE) + +template +PARLAY_ASSUME_TRIVIALLY_RELOCATABLE(std::vector); + +template<> +PARLAY_ASSUME_TRIVIALLY_RELOCATABLE(std::string); + +#endif + +template +PARLAY_ASSUME_TRIVIALLY_RELOCATABLE(std::unordered_set); + +template +PARLAY_ASSUME_TRIVIALLY_RELOCATABLE(std::unordered_map); + +template +PARLAY_ASSUME_TRIVIALLY_RELOCATABLE(std::unordered_multiset); + +template +PARLAY_ASSUME_TRIVIALLY_RELOCATABLE(std::unordered_multimap); + +template +PARLAY_ASSUME_TRIVIALLY_RELOCATABLE(std::stack); + +template +PARLAY_ASSUME_TRIVIALLY_RELOCATABLE(std::queue); + +template +PARLAY_ASSUME_TRIVIALLY_RELOCATABLE(std::priority_queue); + + +// Specializations for GCC. +#elif defined(__GLIBCXX__) + +// GCC std::string is *not* trivially relocatable because when the string is +// in short-size mode, the data pointer points into itself :'( + +template +PARLAY_ASSUME_TRIVIALLY_RELOCATABLE(std::deque); + +template +PARLAY_ASSUME_TRIVIALLY_RELOCATABLE(std::forward_list); + +template +PARLAY_ASSUME_TRIVIALLY_RELOCATABLE(std::vector); + + +// Specializations for Microsoft STL. +#elif defined(_MSC_VER) + +template +PARLAY_ASSUME_TRIVIALLY_RELOCATABLE(std::deque); + +template +PARLAY_ASSUME_TRIVIALLY_RELOCATABLE(std::forward_list); + +template<> +PARLAY_ASSUME_TRIVIALLY_RELOCATABLE(std::string); + +template +PARLAY_ASSUME_TRIVIALLY_RELOCATABLE(std::vector); + +#endif + + +#endif + +} // namespace parlay + +#endif //PARLAY_TYPE_TRAITS_H_ diff --git a/cbits/parlaylib/include/parlay/utilities.h b/cbits/parlaylib/include/parlay/utilities.h new file mode 100644 index 000000000..319d1b440 --- /dev/null +++ b/cbits/parlaylib/include/parlay/utilities.h @@ -0,0 +1,557 @@ + +#ifndef PARLAY_UTILITIES_H_ +#define PARLAY_UTILITIES_H_ + +#include +#include +#include +#include +#include + +#include // IWYU pragma: keep +#include +#include +#include +#include +#include +#include +#include + +#include "portability.h" +#include "type_traits.h" + +#include "internal/debug_uninitialized.h" + +// IWYU pragma: no_include + +namespace parlay { + +struct empty {}; + +template +struct uninitialized { + union { + empty empty_; + T value; + }; + uninitialized() : empty_{} { } + ~uninitialized() { } +}; + +static_assert(std::is_default_constructible_v>>); +static_assert(std::is_destructible_v>>); + +// Returns a raw void* to the given object, stripping cv qualifiers. +template +PARLAY_INLINE constexpr void* voidify(T& obj) noexcept { + return const_cast(static_cast(std::addressof(obj))); +} + +// Obtains a pointer to an object of type T located at the address represented +// by p. Essentially performs std::launder(reinterpret_cast(p)). +template +[[nodiscard]] PARLAY_INLINE constexpr T* from_bytes(std::byte* p) noexcept { + // std::launder not available on older compilers +#ifdef __cpp_lib_launder + return std::launder(reinterpret_cast(p)); +#else + return reinterpret_cast(p); +#endif +} + + +typedef uint32_t flags; +const flags no_flag = 0; +const flags fl_sequential = 1; +const flags fl_debug = 2; +const flags fl_time = 4; +const flags fl_conservative = 8; +const flags fl_inplace = 16; + +template +PARLAY_INLINE void assign_uninitialized(T& a, const type_identity_t& b) { + PARLAY_ASSERT_UNINITIALIZED(a); + new (static_cast(std::addressof(a))) T(b); +} + +template +PARLAY_INLINE auto assign_uninitialized(T& a, type_identity_t&& b) { + PARLAY_ASSERT_UNINITIALIZED(a); + new (static_cast(std::addressof(a))) T(std::move(b)); +} + +template +PARLAY_INLINE void move_uninitialized(T& a, type_identity_t& b) { + PARLAY_ASSERT_UNINITIALIZED(a); + new (static_cast(std::addressof(a))) T(std::move(b)); +} + +// Relocate the object located at source into dest (required to be uninitialized memory), +// leaving source uninitialized. +// +// Effects: Equivalent to: +// +// struct guard { T *t; ~guard() { destroy_at(t); } } g{source}; +// return ::new (voidify(*dest)) T(std::move(*source)); + +// except that if T is trivially relocatable, side effects associated with the relocation +// of the value of *source might not happen. +// +#if defined(__cpp_lib_trivially_relocatable) +using std::relocate_at; +#else +template +PARLAY_INLINE T* relocate_at(T* source, T* dest) + noexcept(is_trivially_relocatable_v || std::is_nothrow_move_constructible_v) { + if constexpr (is_trivially_relocatable::value) { + std::memcpy(voidify(*dest), voidify(*source), sizeof(T)); + return dest; + } + else { + static_assert(std::is_move_constructible::value); + static_assert(std::is_destructible::value); + PARLAY_ASSERT_UNINITIALIZED(*dest); + struct guard { T *t; ~guard() { std::destroy_at(t); PARLAY_ASSERT_UNINITIALIZED(*t); } } g{source}; + return ::new (voidify(*dest)) T(std::move(*source)); + } +} +#endif + +// Relocate the object located at source into the return value, leaving source uninitialized. +// +// Effects: Equivalent to: +// +// remove_cv_t t = std::move(source); +// destroy_at(source); +// return t; + +// except that if T is trivially relocatable, side effects associated with the relocation +// of the object’s value might not happen. +// +#if defined(__cpp_lib_trivially_relocatable) +using std::relocate; +#else +template +[[nodiscard]] PARLAY_INLINE std::remove_cv_t relocate(T* source) + noexcept(std::is_nothrow_move_constructible_v) { + + static_assert(std::is_move_constructible::value); + static_assert(std::is_destructible::value); + std::remove_cv_t t = std::move(*source); + std::destroy_at(source); + return t; +} +#endif + +/* Hashing functions for various integer types */ + +// a 32-bit hash function +inline uint32_t hash32(uint32_t a) { + a = (a + 0x7ed55d16) + (a << 12); + a = (a ^ 0xc761c23c) ^ (a >> 19); + a = (a + 0x165667b1) + (a << 5); + a = (a + 0xd3a2646c) ^ (a << 9); + a = (a + 0xfd7046c5) + (a << 3); + a = (a ^ 0xb55a4f09) ^ (a >> 16); + return a; +} + +inline uint32_t hash32_2(uint32_t a) { + uint32_t z = (a + 0x6D2B79F5UL); + z = (z ^ (z >> 15)) * (z | 1UL); + z ^= z + (z ^ (z >> 7)) * (z | 61UL); + return z ^ (z >> 14); +} + +inline uint32_t hash32_3(uint32_t a) { + uint32_t z = a + 0x9e3779b9; + z ^= z >> 15; // 16 for murmur3 + z *= 0x85ebca6b; + z ^= z >> 13; + z *= 0xc2b2ae3d; // 0xc2b2ae35 for murmur3 + return z ^ (z >> 16); +} + +// from numerical recipes +inline uint64_t hash64(uint64_t u) { + uint64_t v = u * 3935559000370003845ul + 2691343689449507681ul; + v ^= v >> 21; + v ^= v << 37; + v ^= v >> 4; + v *= 4768777513237032717ul; + v ^= v << 20; + v ^= v >> 41; + v ^= v << 5; + return v; +} + +// a slightly cheaper, but possibly not as good version +// based on splitmix64 +inline uint64_t hash64_2(uint64_t x) { + x = (x ^ (x >> 30)) * UINT64_C(0xbf58476d1ce4e5b9); + x = (x ^ (x >> 27)) * UINT64_C(0x94d049bb133111eb); + x = x ^ (x >> 31); + return x; +} + +// Combine two hash values into a single hash value. Formula borrowed from Boost +inline void hash_combine(size_t& seed, size_t v) { + seed ^= v + 0x9e3779b9 + (seed << 6) + (seed >> 2); +} + +// parlay::hash provides a customization point for hashing that allows +// us to override the hash for standard types, or to provide specializations +// for standard types that do not have a std::hash specialization themselves. +template +struct hash : public std::hash { }; + +template +struct hash>> { + size_t operator()(const T& p) const { + return p * UINT64_C(0xbf58476d1ce4e5b9); + } +}; + +// Hashing for std::pairs +template +struct hash> { + size_t operator()(const std::pair& p) const { + size_t h = parlay::hash{}(p.first); + hash_combine(h, parlay::hash{}(p.second)); + return h; + } +}; + +template +struct hash> { + size_t operator()(const std::array& a) const { + size_t h = 1; + for (int i=0; i < n; i++) { + hash_combine(h, parlay::hash{}(a[i])); + } + return h; + } +}; + +/* Atomic write-add, write-min, and write-max */ + +template +inline void write_add(std::atomic* a, EV b) { + T newV, oldV; + do { + oldV = a->load(); + newV = oldV + b; + } while (!std::atomic_compare_exchange_weak(a, &oldV, newV)); +} + +template +inline bool write_min(std::atomic* a, T b, F less) { + T c; + bool r = false; + do { + c = a->load(); + } while (less(b, c) && !(r = std::atomic_compare_exchange_weak(a, &c, b))); + return r; +} + +template +inline bool write_max(std::atomic* a, T b, F less) { + T c; + bool r = false; + do { + c = a->load(); + } while (less(c, b) && !(r = std::atomic_compare_exchange_weak(a, &c, b))); + return r; +} + +// returns the log base 2 rounded up (works on ints or longs or unsigned +// versions) +template +size_t log2_up(T i) { + assert(i > 0); + size_t a = 0; + T b = i - 1; + while (b > 0) { + b = b >> 1; + a++; + } + return a; +} + + +inline size_t granularity(size_t n) { + return (n > 100) ? static_cast(std::ceil(std::sqrt(n))) : 100; +} + +// A copyable_function_wrapper allows an object to store a function (e.g., a lambda, +// or any callable object) as a member while still being default copy-assignable. +// +// As long as the function type is copy constructible (which it needs to be in +// order to even initialize the member), this wrapper will be copy assignable. +template +struct copyable_function_wrapper { + + static_assert(std::is_copy_constructible_v); + + explicit copyable_function_wrapper(F _f) : f(std::move(_f)) {} + + copyable_function_wrapper(const copyable_function_wrapper&) = default; + copyable_function_wrapper(copyable_function_wrapper&&) + noexcept(std::is_nothrow_move_constructible_v) = default; + + copyable_function_wrapper& operator=(const copyable_function_wrapper& other) { + if constexpr (std::is_copy_assignable_v) { + f = other.f; + } + else { + f.~F(); + new (std::addressof(f)) F(other.f); + } + return *this; + } + + copyable_function_wrapper& operator=(copyable_function_wrapper&& other) + noexcept(std::is_nothrow_move_assignable_v && std::is_nothrow_move_constructible_v) { + if constexpr (std::is_move_assignable_v) { + f = std::move(other.f); + } + else { + f.~F(); + new (std::addressof(f)) F(std::move(other.f)); + } + return *this; + } + + ~copyable_function_wrapper() = default; + + template + PARLAY_INLINE decltype(auto) operator()(Args&&... args) noexcept(std::is_nothrow_invocable_v) { + static_assert(std::is_invocable_v); + return std::invoke(f, std::forward(args)...); + } + + template + PARLAY_INLINE decltype(auto) operator()(Args&&... args) const noexcept(std::is_nothrow_invocable_v) { + static_assert(std::is_invocable_v); + return std::invoke(f, std::forward(args)...); + } + + // Special case for when the argument is just a size_t. This is the common case since this is what + // all the functions in delayed_sequence are. This should in theory make absolutely no difference + // compared to the templates whatsoever but for some reason it's faster??? + PARLAY_INLINE decltype(auto) operator()(size_t i) noexcept(std::is_nothrow_invocable_v) { + static_assert(std::is_invocable_v); + return std::invoke(f, i); + } + + PARLAY_INLINE decltype(auto) operator()(size_t i) const noexcept(std::is_nothrow_invocable_v) { + static_assert(std::is_invocable_v); + return std::invoke(f, i); + } + + F* get() { return std::addressof(f); } + const F* get() const { return std::addressof(f); } + + private: + F f; +}; + + +/* For inplace sorting / merging, we sometimes need to move values + around and sometimes we want to make copies. We use tag dispatch + to choose between moving and copying, so that the move algorithm + can be written agnostic to which one it uses. We also account for + moving / copying between uninitialized memory. + + Usage: + assign_dispatch(dest, val, tag_type()) + + Tag types: + move_assign_tag: The value is moved assigned into dest from val + uninitialized_move_tag: The value is move constructed from val into dest + copy_assign_tag: The value is copy assigned into dest from val + uninitialized_copy_tag: The value is copy constructed from val into dest + uninitialized_relocate_tag: The value is destructively moved from val into dest + +*/ + +struct move_assign_tag {}; +struct uninitialized_move_tag {}; +struct copy_assign_tag {}; +struct uninitialized_copy_tag {}; +struct uninitialized_relocate_tag {}; + +// Move dispatch -- move val into dest +template +void assign_dispatch(T& dest, type_identity_t& val, move_assign_tag) { + dest = std::move(val); +} + +// Copy dispatch -- copy val into dest +template +void assign_dispatch(T& dest, const type_identity_t& val, copy_assign_tag) { + dest = val; +} + +// Uninitialized move dispatch -- move construct dest with val +template +void assign_dispatch(T& dest, type_identity_t& val, uninitialized_move_tag) { + assign_uninitialized(dest, std::move(val)); +} + +// Uninitialized copy dispatch -- copy initialize dest with val +template +void assign_dispatch(T& dest, const type_identity_t& val, uninitialized_copy_tag) { + assign_uninitialized(dest, val); +} + +// Uninitialized relocate dispatch -- destructively move val into dest +template +void assign_dispatch(T& dest, T& val, uninitialized_relocate_tag) { + relocate_at(&val, &dest); +} + + +namespace internal { + +// The deleter used by unique_array / make_unique_array +template +struct unique_array_deleter { + void operator()(T* ptr) { + if (ptr != nullptr) { + if constexpr (!std::is_trivially_destructible_v) { + for (size_t i = 0; i < n; i++) { + ptr[i].~T(); + } + } + ::operator delete (static_cast(ptr), std::align_val_t{alignof(T)}); + } + } + size_t n; +}; + +} // namespace internal + +// An alias of unique_ptr as returned by make_unique_array. +// Stores a dynamically allocated array of type T[] that is automatically freed when +// the unique_ptr goes out of scope. Unlike regular unique_ptr, this type with +// make_unique_array can be used to construct an array of non-default-constructible T +template +using unique_array = std::unique_ptr>; + +template +std::size_t size(const unique_array& arr) { + return arr.get_deleter().n; +} + +template +T* begin(unique_array& arr) { + return &arr[0]; +} + +template +T* end(unique_array& arr) { + return &arr[0] + size(arr); +} + +template +T const* begin(const unique_array& arr) { + return &arr[0]; +} + +template +T const* end(const unique_array& arr) { + return &arr[0] + size(arr); +} + +// Creates a unique_ptr consisting of n objects of +// type T initialized by T(init(i)) for i from 0 to n - 1. The objects are +// automatically destroyed and freed when the unique_ptr goes out of scope. +// +// Useful in place of unique_ptr when T is not default constructible +// and hence make_unique can not be used. +template +auto make_unique_array(size_t n, F init) { + // Either T is constructible from init(i), or init(i) returns a prvalue T which means + // we get guaranteed copy elision and can construct T even if it isn't copyable/movable + static_assert(std::is_constructible_v> || + std::is_same_v>); + + if (n == 0) return unique_array(nullptr, {0}); + auto buffer = static_cast(::operator new(n * sizeof(T), std::align_val_t{alignof(T)})); + auto first = new (buffer) T(init(0)); + for (size_t i = 1; i < n; i++) { + new (buffer + i * sizeof(T)) T(init(i)); // NOLINT + } + return unique_array(first, {n}); +} + +/* + Padded types for avoiding false sharing by over-aligning to a cache-line boundary. + + parlay::padded should be a drop-in replacement for T in just about any scenario, + without having to change the code that makes use of the T object. padded can be + initialized with any initializer that would work on T, including copy and move + initialization, and any constructor of T. Assignment operations should also work. + + You can also bind a padded to const/non-const reference/rvalue-reference variables + of type T, e.g., + + padded p{1}; + int& x = p; // refers to the padded int 1 + x = 2; // p now contains 2 + + This also applies to function parameters. For class types, you should also be able + to invoke overloaded operations and member functions, e.g., + + padded> v{1,2,3}; + v.push_back(4); + v[0] = 42; + + You can even invoke function pointers!! + + padded fp = &func; + int x = fp(); + + The only thing that doesn't seem to work is directly invoking member-function pointers... + + */ + +// Padded types to avoid false sharing. Takes a type T and makes its size +// at least Size bytes by over-aligning it. Size must be a valid alignment. +template +struct padded; + +// Use user-defined conversions to pad primitive types +template +struct alignas(Size) padded>> { + padded() : x{} {} + /* implicit */ padded(T _x) : x(_x) { } // cppcheck-suppress noExplicitConstructor // NOLINT + /* implicit */ operator T&() & { return x; } // NOLINT + /* implicit */ operator const T&() const& { return x; } // NOLINT + /* implicit */ operator T&&() && { return std::move(x); } // NOLINT + /* implicit */ operator const T&&() const&& { return std::move(x); } // NOLINT + + // Allow padded function pointers to be directly invoked! + template auto operator()(Ts... args) const noexcept(std::is_nothrow_invocable_v) + -> std::enable_if_t, std::invoke_result_t> { + return x(std::forward(args)...); + } + + T x; +}; + +// Use inheritance to pad class types +template +struct alignas(Size) padded>> : public T { + static_assert(std::is_same_v>, + "padded requires T be non-const-volatile qualified. Use const padded instead of padded"); + using T::T; // inherit (non-special) constructors + padded() = default; // implement non-inherited special constructors + /* implicit */ padded(const T& t) : T(t) {} // cppcheck-suppress noExplicitConstructor // NOLINT + /* implicit */ padded(T&& t) : T(std::move(t)) {} // cppcheck-suppress noExplicitConstructor // NOLINT +}; + +} // namespace parlay + +#endif // PARLAY_UTILITIES_H_ diff --git a/cbits/parlaylib/readme.md b/cbits/parlaylib/readme.md new file mode 100644 index 000000000..f1a5c7311 --- /dev/null +++ b/cbits/parlaylib/readme.md @@ -0,0 +1,23 @@ +This folder contains an extraction of [ParlayLib](https://github.com/cmuparlay/parlaylib). We use their scalable allocator. + +ParlayLib is licensed under the MIT License + +Copyright (c) 2020-present Guy Blelloch and other contributors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/hie.yaml b/hie.yaml new file mode 100644 index 000000000..96ca33ebb --- /dev/null +++ b/hie.yaml @@ -0,0 +1,11 @@ +# Tooling support. Technically not needed, but makes things explicit and exposes missing modules in accelerate.cabal. +cradle: + stack: + - path: "src" + component: "accelerate:lib" + + - path: "test/doctest" + component: "accelerate:test:doctest" + + - path: "test/nofib" + component: "accelerate:test:nofib-interpreter" diff --git a/icebox/Sugar/Tag.hs b/icebox/Sugar/Tag.hs index 39ece0bbe..633a20deb 100644 --- a/icebox/Sugar/Tag.hs +++ b/icebox/Sugar/Tag.hs @@ -1,4 +1,5 @@ {-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE CPP #-} {-# LANGUAGE DefaultSignatures #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} @@ -68,7 +69,7 @@ instance Tagged Bool instance Tagged Ordering instance Tagged (Maybe a) instance Tagged (Either a b) - +#ifndef __GHCIDE__ runQ $ do let mkTuple :: Int -> Q Dec @@ -81,3 +82,24 @@ runQ $ do -- mapM mkTuple [2..16] +#else +instance Tagged (x0, x1) +instance Tagged (x0, x1, x2) +instance Tagged (x0, x1, x2, x3) +instance Tagged (x0, x1, x2, x3, x4) +instance Tagged (x0, x1, x2, x3, x4, x5) +instance Tagged (x0, x1, x2, x3, x4, x5, x6) +instance Tagged (x0, x1, x2, x3, x4, x5, x6, x7) +instance Tagged (x0, x1, x2, x3, x4, x5, x6, x7, x8) +instance Tagged (x0, x1, x2, x3, x4, x5, x6, x7, x8, x9) +instance Tagged (x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10) +instance Tagged (x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11) +instance Tagged (x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, + x12) +instance Tagged (x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, + x12, x13) +instance Tagged (x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, + x12, x13, x14) +instance Tagged (x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, + x12, x13, x14, x15) +#endif diff --git a/src/Data/Array/Accelerate/Array/Unique.hs b/icebox/Unique.hs similarity index 100% rename from src/Data/Array/Accelerate/Array/Unique.hs rename to icebox/Unique.hs diff --git a/src/Crypto/Hash/XKCP.hs b/src/Crypto/Hash/XKCP.hs index c0e22c87f..4b84e06c8 100644 --- a/src/Crypto/Hash/XKCP.hs +++ b/src/Crypto/Hash/XKCP.hs @@ -1,10 +1,8 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE CPP #-} -{-# LANGUAGE ForeignFunctionInterface #-} -{-# LANGUAGE MagicHash #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE UnboxedTuples #-} -{-# LANGUAGE UnliftedFFITypes #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE UnboxedTuples #-} +{-# LANGUAGE UnliftedFFITypes #-} {-# OPTIONS_GHC -fobject-code #-} {-# OPTIONS_HADDOCK hide #-} -- | @@ -19,18 +17,20 @@ module Crypto.Hash.XKCP ( - SHA3_256, - hash, hashlazy, + SHA3_256(..), + hash, hashlazy, hashIncrement, ) where import Control.Monad +import Data.List (mapAccumR) import Data.Word import Foreign.C.Types import Foreign.Ptr import Language.Haskell.TH.Syntax -import System.IO.Unsafe +import System.IO.Unsafe +import Data.Array.Byte import qualified Data.ByteString as S import qualified Data.ByteString.Lazy as L import qualified Data.ByteString.Lazy.Internal as L @@ -105,6 +105,16 @@ hashlazy lbs = unsafePerformIO $! do go cs go lbs +-- Increments a hash value. To be used carefully. Designed for generating +-- multiple variants of the same kernel: either for different phases of the +-- kernel (like initialization, parallel work and finalization), or to +-- specialize the kernel for e.g. different input sizes. +hashIncrement :: Word8 -> SHA3_256 -> SHA3_256 +hashIncrement inc (SHA3_256 input#) = SHA3_256 result# + where + !(ByteArray result#) = fromList $ snd $ mapAccumR + (\carry x -> let y = carry + x in (if y < x then 1 else 0, y)) + inc $ toList $ ByteArray input# -- Internals -- ----------------------------------------------------------------------------- @@ -151,14 +161,14 @@ keccak_HashFinal_SHA3_256 (Keccak_HashInstance hash_instance#) = -- #ifndef __GHCIDE__ -foreign import ccall unsafe "SHA3_256" c_sha3_256 :: MutableByteArray# RealWorld -> Ptr Word8 -> CSize -> IO CInt +foreign import ccall unsafe "SHA3_256" c_sha3_256 :: MutableByteArray# RealWorld -> Ptr Word8 -> CSize -> IO CInt foreign import ccall unsafe "Keccak_HashInitialize" c_keccak_hash_initialise :: MutableByteArray# RealWorld -> CUInt -> CUInt -> CUInt -> CUChar -> IO CInt -foreign import ccall unsafe "Keccak_HashUpdate" c_keccak_hash_update :: MutableByteArray# RealWorld -> Ptr Word8 -> CSize -> IO CInt -foreign import ccall unsafe "Keccak_HashFinal" c_keccak_hash_final :: MutableByteArray# RealWorld -> MutableByteArray# RealWorld -> IO CInt +foreign import ccall unsafe "Keccak_HashUpdate" c_keccak_hash_update :: MutableByteArray# RealWorld -> Ptr Word8 -> CSize -> IO CInt +foreign import ccall unsafe "Keccak_HashFinal" c_keccak_hash_final :: MutableByteArray# RealWorld -> MutableByteArray# RealWorld -> IO CInt #else -c_sha3_256 :: Ptr Word8 -> Ptr Word8 -> CSize -> IO CInt +c_sha3_256 :: MutableByteArray# RealWorld -> Ptr Word8 -> CSize -> IO CInt c_sha3_256 = undefined c_keccak_hash_initialise :: MutableByteArray# RealWorld -> CUInt -> CUInt -> CUInt -> CUChar -> IO CInt diff --git a/src/Data/.gitignore b/src/Data/.gitignore new file mode 100644 index 000000000..1f94316f2 --- /dev/null +++ b/src/Data/.gitignore @@ -0,0 +1,2 @@ +*.o +*.hi diff --git a/src/Data/Array/Accelerate.hs b/src/Data/Array/Accelerate.hs index 7936caa68..b27db6ddc 100644 --- a/src/Data/Array/Accelerate.hs +++ b/src/Data/Array/Accelerate.hs @@ -1,8 +1,5 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE ExplicitNamespaces #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE PatternSynonyms #-} -- | -- Module : Data.Array.Accelerate -- Description : The Accelerate standard prelude @@ -211,7 +208,7 @@ module Data.Array.Accelerate ( -- ** Composition -- *** Flow control - (?|), acond, awhile, + (?|), acond, awhile, awhileSpeculative, awhileAndOne, IfThenElse(..), -- *** Controlling execution @@ -227,8 +224,14 @@ module Data.Array.Accelerate ( -- *** Zipping zipWith, zipWith3, zipWith4, zipWith5, zipWith6, zipWith7, zipWith8, zipWith9, + zipWithChecked, zipWithChecked3, zipWithChecked4, zipWithChecked5, + zipWithChecked6, zipWithChecked7, zipWithChecked8, zipWithChecked9, izipWith, izipWith3, izipWith4, izipWith5, izipWith6, izipWith7, izipWith8, izipWith9, + izipWithChecked, izipWithChecked3, izipWithChecked4, izipWithChecked5, + izipWithChecked6, izipWithChecked7, izipWithChecked8, izipWithChecked9, zip, zip3, zip4, zip5, zip6, zip7, zip8, zip9, + zipChecked, zipChecked3, zipChecked4, zipChecked5, + zipChecked6, zipChecked7, zipChecked8, zipChecked9, -- *** Unzipping unzip, unzip3, unzip4, unzip5, unzip6, unzip7, unzip8, unzip9, @@ -247,7 +250,7 @@ module Data.Array.Accelerate ( -- *** Permutations -- **** Forward permutation (scatter) - permute, + permute, permute', permuteUnique, permuteUnique', scatter, -- **** Backward permutation (gather) @@ -370,13 +373,14 @@ module Data.Array.Accelerate ( fst, afst, snd, asnd, curry, uncurry, -- *** Flow control - (?), match, cond, while, iterate, + (?), match, cond, select, while, iterate, + Assert(assert, assertMessage), -- *** Scalar reduction sfoldl, -- *** Logical operations - (&&), (||), not, + (&&), (&&!), (||), (||!), not, -- *** Numeric operations subtract, even, odd, gcd, lcm, (^), (^^), @@ -435,8 +439,13 @@ module Data.Array.Accelerate ( CShort, CUShort, CInt, CUInt, CLong, CULong, CLLong, CULLong, CChar, CSChar, CUChar, + runAt, runWithAt, + run1At, run1WithAt, + runNAt, runNWithAt, + inspectCompiler, inspectCompiler' ) where +import Data.Array.Accelerate.Backend import Data.Array.Accelerate.Classes.Bounded import Data.Array.Accelerate.Classes.Enum import Data.Array.Accelerate.Classes.Eq @@ -456,6 +465,7 @@ import Data.Array.Accelerate.Language import Data.Array.Accelerate.Pattern import Data.Array.Accelerate.Pattern.TH import Data.Array.Accelerate.Prelude +import Data.Array.Accelerate.Trafo ( inspectCompiler' ) import Data.Array.Accelerate.Pretty () -- show instances import Data.Array.Accelerate.Smart import Data.Array.Accelerate.Sugar.Array ( Array, Arrays, Scalar, Vector, Matrix, Segments, fromFunction, fromFunctionM, toList, fromList ) @@ -527,6 +537,21 @@ arraySize = S.size . S.shape arrayReshape :: (Shape sh, Shape sh') => sh -> Array sh' e -> Array sh e arrayReshape = S.reshape +{- +-- Use the Interpreter's fusion rules to show an Accelerate computation +-- We should also add a way to do this with other fusion rules +instance Arrays arrs => Show (Acc arrs) where + show = withSimplStats . show . convertAcc @InterpretOp + +instance Afunction (Acc a -> f) => Show (Acc a -> f) where + show = withSimplStats . show . convertAfun @_ @InterpretOp + +instance Elt e => Show (Exp e) where + show = withSimplStats . show . convertExp + +instance Function (Exp a -> f) => Show (Exp a -> f) where + show = withSimplStats . show . convertFun +-} -- Named documentation chunks -- -------------------------- @@ -696,4 +721,3 @@ arrayReshape = S.reshape -- * : binary serialisation of arrays using -- * : efficient boxed and unboxed one-dimensional arrays -- - diff --git a/src/Data/Array/Accelerate/AST.hs b/src/Data/Array/Accelerate/AST.hs index 97eaa57da..8f583fdee 100644 --- a/src/Data/Array/Accelerate/AST.hs +++ b/src/Data/Array/Accelerate/AST.hs @@ -1,14 +1,4 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE TemplateHaskell #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.AST @@ -83,18 +73,16 @@ module Data.Array.Accelerate.AST ( -- ** Scalar expressions ELeftHandSide, ExpVar, ExpVars, expVars, undefs, - Fun, OpenFun(..), - Exp, OpenExp(..), + Fun, OpenFun, PreOpenFun(..), + Exp, OpenExp, PreOpenExp(..), + ArrayInstr(..), IsArrayInstr(..), NoArrayInstr, Boundary(..), - PrimConst(..), - PrimFun(..), - PrimBool, - PrimMaybe, + PrimFun(..), Cmp(..), + PrimBool, PrimMaybe, -- ** Extracting type information HasArraysR(..), arrayR, expType, - primConstType, primFunType, -- ** Normal-form @@ -109,7 +97,6 @@ module Data.Array.Accelerate.AST ( rnfExpVar, rnfBoundary, rnfConst, - rnfPrimConst, rnfPrimFun, -- ** Template Haskell @@ -123,31 +110,28 @@ module Data.Array.Accelerate.AST ( liftELeftHandSide, liftExpVar, liftBoundary, - liftPrimConst, liftPrimFun, liftMessage, -- ** Miscellaneous + encodeArrayVar, formatPreAccOp, - formatExpOp, ) where import Data.Array.Accelerate.AST.Idx import Data.Array.Accelerate.AST.LeftHandSide import Data.Array.Accelerate.AST.Var -import Data.Array.Accelerate.Error +import Data.Array.Accelerate.AST.Exp +import Data.Array.Accelerate.Analysis.Hash.Exp as Hash import Data.Array.Accelerate.Representation.Array import Data.Array.Accelerate.Representation.Elt import Data.Array.Accelerate.Representation.Shape import Data.Array.Accelerate.Representation.Slice import Data.Array.Accelerate.Representation.Stencil -import Data.Array.Accelerate.Representation.Tag import Data.Array.Accelerate.Representation.Type -import Data.Array.Accelerate.Representation.Vec import Data.Array.Accelerate.Sugar.Foreign import Data.Array.Accelerate.Type -import Data.Primitive.Vec import Control.DeepSeq ( NFData(..) ) import Data.Kind @@ -158,8 +142,7 @@ import Formatting import Language.Haskell.TH.Extra ( CodeQ ) import qualified Language.Haskell.TH.Extra as TH import qualified Language.Haskell.TH.Syntax as TH - -import GHC.TypeLits +import Data.Typeable ( (:~:)(..) ) -- Array expressions @@ -185,7 +168,7 @@ type Afun = OpenAfun () -- Vanilla open array computations -- -newtype OpenAcc aenv t = OpenAcc (PreOpenAcc OpenAcc aenv t) +newtype OpenAcc aenv t = OpenAcc { runOpenAcc :: PreOpenAcc OpenAcc aenv t } -- | Closed array expression aka an array program -- @@ -197,10 +180,6 @@ type ALeftHandSide = LeftHandSide ArrayR type ArrayVar = Var ArrayR type ArrayVars aenv = Vars ArrayR aenv --- Bool is not a primitive type -type PrimBool = TAG -type PrimMaybe a = (TAG, ((), a)) - -- Trace messages data Message a where Message :: (a -> String) -- embedded show @@ -250,15 +229,19 @@ data PreOpenAcc (acc :: Type -> Type -> Type) aenv a where Anil :: PreOpenAcc acc aenv () - -- Array-function application. - -- - -- The array function is not closed at the core level because we need access - -- to free variables introduced by 'run1' style evaluators. See Issue#95. - -- - Apply :: ArraysR arrs2 - -> PreOpenAfun acc aenv (arrs1 -> arrs2) - -> acc aenv arrs1 - -> PreOpenAcc acc aenv arrs2 + -- Fusion boundary. + -- Corresponds to 'compute' in the user-facing language. + Manifest :: acc aenv as + -> PreOpenAcc acc aenv as + + Aassume :: Exp aenv PrimBool + -> acc aenv as + -> PreOpenAcc acc aenv as + + Aassert :: Text + -> Exp aenv PrimBool + -> acc aenv as + -> PreOpenAcc acc aenv as -- Apply a backend-specific foreign function to an array, with a pure -- Accelerate version for use with other backends. The functions must be @@ -402,30 +385,28 @@ data PreOpenAcc (acc :: Type -> Type -> Type) aenv a where -> acc aenv (Array (sh, Int) e) -> PreOpenAcc acc aenv (Array (sh, Int) e, Array sh e) - -- Generalised forward permutation is characterised by a permutation function - -- that determines for each element of the source array where it should go in - -- the output. The permutation can be between arrays of varying shape and - -- dimensionality. + -- Generalised forward permutation. The source array contains values combined + -- with the indices they should be written to. -- - -- Other characteristics of the permutation function 'f': + -- Characteristics of the permutation mapping: -- - -- 1. 'f' is a (morally) partial function: only the elements of the domain - -- for which the function evaluates to a 'Just' value are mapped in the - -- result. Other elements are dropped. + -- 1. it is a (morally) partial mapping: only the elements of the domain + -- for which the source array contains a 'Just' value are mapped in the + -- result. -- - -- 2. 'f' is not surjective: positions in the target array need not be + -- 2. it is not surjective: positions in the target array need not be -- picked up by the permutation function, so the target array must first -- be initialised from an array of default values. -- - -- 3. 'f' is not injective: distinct elements of the domain may map to the + -- 3. it is not injective: distinct elements of the domain may map to the -- same position in the target array. In this case the combination -- function is used to combine elements, which needs to be /associative/ - -- and /commutative/. + -- and /commutative/. When the combination function is missing, + -- we assume that 'f' is injective. -- - Permute :: Fun aenv (e -> e -> e) -- combination function - -> acc aenv (Array sh' e) -- default values - -> Fun aenv (sh -> PrimMaybe sh') -- permutation function - -> acc aenv (Array sh e) -- source array + Permute :: Maybe (Fun aenv (e -> e -> e)) -- combination function + -> acc aenv (Array sh' e) -- default values + -> acc aenv (Array sh (PrimMaybe (((),sh'), e))) -- source array with destination indices -> PreOpenAcc acc aenv (Array sh' e) -- Generalised multi-dimensional backwards permutation; the permutation can @@ -459,11 +440,6 @@ data PreOpenAcc (acc :: Type -> Type -> Type) aenv a where -> acc aenv (Array sh b) -- source array #2 -> PreOpenAcc acc aenv (Array sh c) - -data Direction = LeftToRight | RightToLeft - deriving Eq - - -- | Vanilla boundary condition specification for stencil operations -- data Boundary aenv t where @@ -485,38 +461,47 @@ data Boundary aenv t where -> Boundary aenv (Array sh e) --- Embedded expressions --- -------------------- +data ArrayInstr aenv f where + -- Array shape. + Shape :: ArrayVar aenv (Array dim t) -> ArrayInstr aenv (() -> dim) --- | Vanilla open function abstraction --- -data OpenFun env aenv t where - Body :: OpenExp env aenv t -> OpenFun env aenv t - Lam :: ELeftHandSide a env env' -> OpenFun env' aenv t -> OpenFun env aenv (a -> t) + -- Project a single scalar from an array. + Index :: ArrayVar aenv (Array dim t) -> ArrayInstr aenv (dim -> t) + LinearIndex :: ArrayVar aenv (Array dim t) -> ArrayInstr aenv (Int -> t) --- | Vanilla function without free scalar variables --- -type Fun = OpenFun () +instance IsArrayInstr (ArrayInstr aenv) where + arrayInstrType (Shape (Var (ArrayR shr _) _)) = shapeType shr + arrayInstrType (Index (Var (ArrayR _ tp) _)) = tp + arrayInstrType (LinearIndex (Var (ArrayR _ tp) _)) = tp --- | Vanilla expression without free scalar variables --- -type Exp = OpenExp () + liftArrayInstr (Shape var) = [|| Shape $$(liftArrayVar var) ||] + liftArrayInstr (Index var) = [|| Index $$(liftArrayVar var) ||] + liftArrayInstr (LinearIndex var) = [|| LinearIndex $$(liftArrayVar var) ||] --- Types for scalar bindings --- -type ELeftHandSide = LeftHandSide ScalarType -type ExpVar = Var ScalarType -type ExpVars env = Vars ScalarType env + rnfArrayInstr (Shape var) = rnfArrayVar var + rnfArrayInstr (Index var) = rnfArrayVar var + rnfArrayInstr (LinearIndex var) = rnfArrayVar var + + showArrayInstrOp Shape{} = "Shape" + showArrayInstrOp Index{} = "Index" + showArrayInstrOp LinearIndex{} = "LinearIndex" + + matchArrayInstr (Shape v1) (Shape v2) | Just Refl <- matchVar v1 v2 = Just Refl + matchArrayInstr (Index v1) (Index v2) | Just Refl <- matchVar v1 v2 = Just Refl + matchArrayInstr (LinearIndex v1) (LinearIndex v2) | Just Refl <- matchVar v1 v2 = Just Refl + matchArrayInstr _ _ = Nothing + + encodeArrayInstr = \case + Index a -> intHost $(hashQ ("Index" :: String)) <> encodeArrayVar a + LinearIndex a -> intHost $(hashQ ("LinearIndex" :: String)) <> encodeArrayVar a + Shape a -> intHost $(hashQ ("Shape" :: String)) <> encodeArrayVar a -expVars :: ExpVars env t -> OpenExp env aenv t -expVars TupRunit = Nil -expVars (TupRsingle var) = Evar var -expVars (TupRpair v1 v2) = expVars v1 `Pair` expVars v2 + inlineArrayInstr = \case + Shape _ -> True + _ -> False -undefs :: TypeR t -> OpenExp env aenv t -undefs (TupRsingle tp) = Undef tp -undefs (TupRpair t1 t2) = undefs t1 `Pair` undefs t2 -undefs TupRunit = Nil +encodeArrayVar :: ArrayVar aenv a -> Hash.Builder +encodeArrayVar (Var repr v) = encodeArrayType repr <> encodeIdx v -- | Vanilla open expressions using de Bruijn indices for variables ranging -- over tuples of scalars and arrays of tuples. All code, except Cond, is @@ -525,236 +510,19 @@ undefs TupRunit = Nil -- The data type is parametrised over the representation type (not the -- surface types). -- -data OpenExp env aenv t where - - -- Local binding of a scalar expression - Let :: ELeftHandSide bnd_t env env' - -> OpenExp env aenv bnd_t - -> OpenExp env' aenv body_t - -> OpenExp env aenv body_t - - -- Variable index, ranging _only_ over scalars - Evar :: ExpVar env t - -> OpenExp env aenv t - - -- Apply a backend-specific foreign function - Foreign :: Foreign asm - => TypeR y - -> asm (x -> y) -- foreign function - -> Fun () (x -> y) -- alternate implementation (for other backends) - -> OpenExp env aenv x - -> OpenExp env aenv y - - -- Tuples - Pair :: OpenExp env aenv t1 - -> OpenExp env aenv t2 - -> OpenExp env aenv (t1, t2) - - Nil :: OpenExp env aenv () - - -- SIMD vectors - VecPack :: KnownNat n - => VecR n s tup - -> OpenExp env aenv tup - -> OpenExp env aenv (Vec n s) - - VecUnpack :: KnownNat n - => VecR n s tup - -> OpenExp env aenv (Vec n s) - -> OpenExp env aenv tup - - -- Array indices & shapes - IndexSlice :: SliceIndex slix sl co sh - -> OpenExp env aenv slix - -> OpenExp env aenv sh - -> OpenExp env aenv sl - - IndexFull :: SliceIndex slix sl co sh - -> OpenExp env aenv slix - -> OpenExp env aenv sl - -> OpenExp env aenv sh - - -- Shape and index conversion - ToIndex :: ShapeR sh - -> OpenExp env aenv sh -- shape of the array - -> OpenExp env aenv sh -- index into the array - -> OpenExp env aenv Int - - FromIndex :: ShapeR sh - -> OpenExp env aenv sh -- shape of the array - -> OpenExp env aenv Int -- index into linear representation - -> OpenExp env aenv sh - - -- Case statement - Case :: OpenExp env aenv TAG - -> [(TAG, OpenExp env aenv b)] -- list of equations - -> Maybe (OpenExp env aenv b) -- default case - -> OpenExp env aenv b - - -- Conditional expression (non-strict in 2nd and 3rd argument) - Cond :: OpenExp env aenv PrimBool - -> OpenExp env aenv t - -> OpenExp env aenv t - -> OpenExp env aenv t - - -- Value recursion - While :: OpenFun env aenv (a -> PrimBool) -- continue while true - -> OpenFun env aenv (a -> a) -- function to iterate - -> OpenExp env aenv a -- initial value - -> OpenExp env aenv a - - -- Constant values - Const :: ScalarType t - -> t - -> OpenExp env aenv t - - PrimConst :: PrimConst t - -> OpenExp env aenv t - - -- Primitive scalar operations - PrimApp :: PrimFun (a -> r) - -> OpenExp env aenv a - -> OpenExp env aenv r +type OpenExp env aenv = PreOpenExp (ArrayInstr aenv) env - -- Project a single scalar from an array. - -- The array expression can not contain any free scalar variables. - Index :: ArrayVar aenv (Array dim t) - -> OpenExp env aenv dim - -> OpenExp env aenv t - - LinearIndex :: ArrayVar aenv (Array dim t) - -> OpenExp env aenv Int - -> OpenExp env aenv t - - -- Array shape. - -- The array expression can not contain any free scalar variables. - Shape :: ArrayVar aenv (Array dim e) - -> OpenExp env aenv dim - - -- Number of elements of an array given its shape - ShapeSize :: ShapeR dim - -> OpenExp env aenv dim - -> OpenExp env aenv Int - - -- Unsafe operations (may fail or result in undefined behaviour) - -- An unspecified bit pattern - Undef :: ScalarType t - -> OpenExp env aenv t - - -- Reinterpret the bits of a value as a different type - Coerce :: BitSizeEq a b - => ScalarType a - -> ScalarType b - -> OpenExp env aenv a - -> OpenExp env aenv b - --- |Primitive constant values +-- | Vanilla open function abstraction -- -data PrimConst ty where - - -- constants from Bounded - PrimMinBound :: BoundedType a -> PrimConst a - PrimMaxBound :: BoundedType a -> PrimConst a - - -- constant from Floating - PrimPi :: FloatingType a -> PrimConst a - +type OpenFun env aenv = PreOpenFun (ArrayInstr aenv) env --- |Primitive scalar operations +-- | Vanilla function without free scalar variables -- -data PrimFun sig where - - -- operators from Num - PrimAdd :: NumType a -> PrimFun ((a, a) -> a) - PrimSub :: NumType a -> PrimFun ((a, a) -> a) - PrimMul :: NumType a -> PrimFun ((a, a) -> a) - PrimNeg :: NumType a -> PrimFun (a -> a) - PrimAbs :: NumType a -> PrimFun (a -> a) - PrimSig :: NumType a -> PrimFun (a -> a) - - -- operators from Integral - PrimQuot :: IntegralType a -> PrimFun ((a, a) -> a) - PrimRem :: IntegralType a -> PrimFun ((a, a) -> a) - PrimQuotRem :: IntegralType a -> PrimFun ((a, a) -> (a, a)) - PrimIDiv :: IntegralType a -> PrimFun ((a, a) -> a) - PrimMod :: IntegralType a -> PrimFun ((a, a) -> a) - PrimDivMod :: IntegralType a -> PrimFun ((a, a) -> (a, a)) - - -- operators from Bits & FiniteBits - PrimBAnd :: IntegralType a -> PrimFun ((a, a) -> a) - PrimBOr :: IntegralType a -> PrimFun ((a, a) -> a) - PrimBXor :: IntegralType a -> PrimFun ((a, a) -> a) - PrimBNot :: IntegralType a -> PrimFun (a -> a) - PrimBShiftL :: IntegralType a -> PrimFun ((a, Int) -> a) - PrimBShiftR :: IntegralType a -> PrimFun ((a, Int) -> a) - PrimBRotateL :: IntegralType a -> PrimFun ((a, Int) -> a) - PrimBRotateR :: IntegralType a -> PrimFun ((a, Int) -> a) - PrimPopCount :: IntegralType a -> PrimFun (a -> Int) - PrimCountLeadingZeros :: IntegralType a -> PrimFun (a -> Int) - PrimCountTrailingZeros :: IntegralType a -> PrimFun (a -> Int) - - -- operators from Fractional and Floating - PrimFDiv :: FloatingType a -> PrimFun ((a, a) -> a) - PrimRecip :: FloatingType a -> PrimFun (a -> a) - PrimSin :: FloatingType a -> PrimFun (a -> a) - PrimCos :: FloatingType a -> PrimFun (a -> a) - PrimTan :: FloatingType a -> PrimFun (a -> a) - PrimAsin :: FloatingType a -> PrimFun (a -> a) - PrimAcos :: FloatingType a -> PrimFun (a -> a) - PrimAtan :: FloatingType a -> PrimFun (a -> a) - PrimSinh :: FloatingType a -> PrimFun (a -> a) - PrimCosh :: FloatingType a -> PrimFun (a -> a) - PrimTanh :: FloatingType a -> PrimFun (a -> a) - PrimAsinh :: FloatingType a -> PrimFun (a -> a) - PrimAcosh :: FloatingType a -> PrimFun (a -> a) - PrimAtanh :: FloatingType a -> PrimFun (a -> a) - PrimExpFloating :: FloatingType a -> PrimFun (a -> a) - PrimSqrt :: FloatingType a -> PrimFun (a -> a) - PrimLog :: FloatingType a -> PrimFun (a -> a) - PrimFPow :: FloatingType a -> PrimFun ((a, a) -> a) - PrimLogBase :: FloatingType a -> PrimFun ((a, a) -> a) - - -- FIXME: add missing operations from RealFrac & RealFloat - - -- operators from RealFrac - PrimTruncate :: FloatingType a -> IntegralType b -> PrimFun (a -> b) - PrimRound :: FloatingType a -> IntegralType b -> PrimFun (a -> b) - PrimFloor :: FloatingType a -> IntegralType b -> PrimFun (a -> b) - PrimCeiling :: FloatingType a -> IntegralType b -> PrimFun (a -> b) - -- PrimProperFraction :: FloatingType a -> IntegralType b -> PrimFun (a -> (b, a)) - - -- operators from RealFloat - PrimAtan2 :: FloatingType a -> PrimFun ((a, a) -> a) - PrimIsNaN :: FloatingType a -> PrimFun (a -> PrimBool) - PrimIsInfinite :: FloatingType a -> PrimFun (a -> PrimBool) - - -- relational and equality operators - PrimLt :: SingleType a -> PrimFun ((a, a) -> PrimBool) - PrimGt :: SingleType a -> PrimFun ((a, a) -> PrimBool) - PrimLtEq :: SingleType a -> PrimFun ((a, a) -> PrimBool) - PrimGtEq :: SingleType a -> PrimFun ((a, a) -> PrimBool) - PrimEq :: SingleType a -> PrimFun ((a, a) -> PrimBool) - PrimNEq :: SingleType a -> PrimFun ((a, a) -> PrimBool) - PrimMax :: SingleType a -> PrimFun ((a, a) -> a) - PrimMin :: SingleType a -> PrimFun ((a, a) -> a) - - -- logical operators - -- - -- Note that these operators are strict in both arguments. That is, the - -- second argument of PrimLAnd is always evaluated even when the first - -- argument is false. - -- - -- We define (surface level) (&&) and (||) using if-then-else to enable - -- short-circuiting, while (&&!) and (||!) are strict versions of these - -- operators, which are defined using PrimLAnd and PrimLOr. - -- - PrimLAnd :: PrimFun ((PrimBool, PrimBool) -> PrimBool) - PrimLOr :: PrimFun ((PrimBool, PrimBool) -> PrimBool) - PrimLNot :: PrimFun (PrimBool -> PrimBool) +type Fun aenv = OpenFun () aenv - -- general conversion between types - PrimFromIntegral :: IntegralType a -> NumType b -> PrimFun (a -> b) - PrimToFloating :: NumType a -> FloatingType b -> PrimFun (a -> b) +-- | Vanilla expression without free scalar variables +-- +type Exp aenv = OpenExp () aenv -- Type utilities @@ -776,7 +544,9 @@ instance HasArraysR acc => HasArraysR (PreOpenAcc acc) where arraysR (Apair as bs) = TupRpair (arraysR as) (arraysR bs) arraysR Anil = TupRunit arraysR (Atrace _ _ bs) = arraysR bs - arraysR (Apply aR _ _) = aR + arraysR (Manifest as) = arraysR as + arraysR (Aassert _ _ as) = arraysR as + arraysR (Aassume _ as) = arraysR as arraysR (Aforeign r _ _ _) = r arraysR (Acond _ a _) = arraysR a arraysR (Awhile _ (Alam lhs _) _) = lhsToTupR lhs @@ -801,7 +571,7 @@ instance HasArraysR acc => HasArraysR (PreOpenAcc acc) where arraysR (Scan _ _ _ a) = arraysR a arraysR (Scan' _ _ _ a) = let aR@(ArrayR (ShapeRsnoc sh) tR) = arrayR a in TupRsingle aR `TupRpair` TupRsingle (ArrayR sh tR) - arraysR (Permute _ a _ _) = arraysR a + arraysR (Permute _ a _) = arraysR a arraysR (Backpermute sh _ _ a) = let ArrayR _ tR = arrayR a in arraysRarray sh tR arraysR (Stencil _ tR _ _ a) = let ArrayR sh _ = arrayR a @@ -809,145 +579,6 @@ instance HasArraysR acc => HasArraysR (PreOpenAcc acc) where arraysR (Stencil2 _ _ tR _ _ a _ _) = let ArrayR sh _ = arrayR a in arraysRarray sh tR -expType :: HasCallStack => OpenExp aenv env t -> TypeR t -expType = \case - Let _ _ body -> expType body - Evar (Var tR _) -> TupRsingle tR - Foreign tR _ _ _ -> tR - Pair e1 e2 -> TupRpair (expType e1) (expType e2) - Nil -> TupRunit - VecPack vecR _ -> TupRsingle $ VectorScalarType $ vecRvector vecR - VecUnpack vecR _ -> vecRtuple vecR - IndexSlice si _ _ -> shapeType $ sliceShapeR si - IndexFull si _ _ -> shapeType $ sliceDomainR si - ToIndex{} -> TupRsingle scalarTypeInt - FromIndex shr _ _ -> shapeType shr - Case _ ((_,e):_) _ -> expType e - Case _ [] (Just e) -> expType e - Case{} -> internalError "empty case encountered" - Cond _ e _ -> expType e - While _ (Lam lhs _) _ -> lhsToTupR lhs - While{} -> error "What's the matter, you're running in the shadows" - Const tR _ -> TupRsingle tR - PrimConst c -> TupRsingle $ SingleScalarType $ primConstType c - PrimApp f _ -> snd $ primFunType f - Index (Var repr _) _ -> arrayRtype repr - LinearIndex (Var repr _) _ -> arrayRtype repr - Shape (Var repr _) -> shapeType $ arrayRshape repr - ShapeSize{} -> TupRsingle scalarTypeInt - Undef tR -> TupRsingle tR - Coerce _ tR _ -> TupRsingle tR - -primConstType :: PrimConst a -> SingleType a -primConstType = \case - PrimMinBound t -> bounded t - PrimMaxBound t -> bounded t - PrimPi t -> floating t - where - bounded :: BoundedType a -> SingleType a - bounded (IntegralBoundedType t) = NumSingleType $ IntegralNumType t - - floating :: FloatingType t -> SingleType t - floating = NumSingleType . FloatingNumType - -primFunType :: PrimFun (a -> b) -> (TypeR a, TypeR b) -primFunType = \case - -- Num - PrimAdd t -> binary' $ num t - PrimSub t -> binary' $ num t - PrimMul t -> binary' $ num t - PrimNeg t -> unary' $ num t - PrimAbs t -> unary' $ num t - PrimSig t -> unary' $ num t - - -- Integral - PrimQuot t -> binary' $ integral t - PrimRem t -> binary' $ integral t - PrimQuotRem t -> unary' $ integral t `TupRpair` integral t - PrimIDiv t -> binary' $ integral t - PrimMod t -> binary' $ integral t - PrimDivMod t -> unary' $ integral t `TupRpair` integral t - - -- Bits & FiniteBits - PrimBAnd t -> binary' $ integral t - PrimBOr t -> binary' $ integral t - PrimBXor t -> binary' $ integral t - PrimBNot t -> unary' $ integral t - PrimBShiftL t -> (integral t `TupRpair` tint, integral t) - PrimBShiftR t -> (integral t `TupRpair` tint, integral t) - PrimBRotateL t -> (integral t `TupRpair` tint, integral t) - PrimBRotateR t -> (integral t `TupRpair` tint, integral t) - PrimPopCount t -> unary (integral t) tint - PrimCountLeadingZeros t -> unary (integral t) tint - PrimCountTrailingZeros t -> unary (integral t) tint - - -- Fractional, Floating - PrimFDiv t -> binary' $ floating t - PrimRecip t -> unary' $ floating t - PrimSin t -> unary' $ floating t - PrimCos t -> unary' $ floating t - PrimTan t -> unary' $ floating t - PrimAsin t -> unary' $ floating t - PrimAcos t -> unary' $ floating t - PrimAtan t -> unary' $ floating t - PrimSinh t -> unary' $ floating t - PrimCosh t -> unary' $ floating t - PrimTanh t -> unary' $ floating t - PrimAsinh t -> unary' $ floating t - PrimAcosh t -> unary' $ floating t - PrimAtanh t -> unary' $ floating t - PrimExpFloating t -> unary' $ floating t - PrimSqrt t -> unary' $ floating t - PrimLog t -> unary' $ floating t - PrimFPow t -> binary' $ floating t - PrimLogBase t -> binary' $ floating t - - -- RealFrac - PrimTruncate a b -> unary (floating a) (integral b) - PrimRound a b -> unary (floating a) (integral b) - PrimFloor a b -> unary (floating a) (integral b) - PrimCeiling a b -> unary (floating a) (integral b) - - -- RealFloat - PrimAtan2 t -> binary' $ floating t - PrimIsNaN t -> unary (floating t) tbool - PrimIsInfinite t -> unary (floating t) tbool - - -- Relational and equality - PrimLt t -> compare' t - PrimGt t -> compare' t - PrimLtEq t -> compare' t - PrimGtEq t -> compare' t - PrimEq t -> compare' t - PrimNEq t -> compare' t - PrimMax t -> binary' $ single t - PrimMin t -> binary' $ single t - - -- Logical - PrimLAnd -> binary' tbool - PrimLOr -> binary' tbool - PrimLNot -> unary' tbool - - -- general conversion between types - PrimFromIntegral a b -> unary (integral a) (num b) - PrimToFloating a b -> unary (num a) (floating b) - - where - unary a b = (a, b) - unary' a = unary a a - binary a b = (a `TupRpair` a, b) - binary' a = binary a a - compare' a = binary (single a) tbool - - single = TupRsingle . SingleScalarType - num = TupRsingle . SingleScalarType . NumSingleType - integral = num . IntegralNumType - floating = num . FloatingNumType - - tbool = TupRsingle scalarTypeWord8 - tint = TupRsingle scalarTypeInt - - -- Normal form data -- ================ @@ -992,7 +623,7 @@ rnfPreOpenAcc rnfA pacc = rnfB = rnfBoundary rnfM :: Message a -> () - rnfM (Message f g msg) = f `seq` rnfMaybe (\x -> x `seq` ()) g `seq` rnf msg + rnfM (Message f g msg) = f `seq` rnfMaybe (`seq` ()) g `seq` rnf msg in case pacc of Alet lhs bnd body -> rnfALeftHandSide lhs `seq` rnfA bnd `seq` rnfA body @@ -1000,7 +631,9 @@ rnfPreOpenAcc rnfA pacc = Apair as bs -> rnfA as `seq` rnfA bs Anil -> () Atrace msg as bs -> rnfM msg `seq` rnfA as `seq` rnfA bs - Apply repr afun acc -> rnfTupR rnfArrayR repr `seq` rnfAF afun `seq` rnfA acc + Manifest as -> rnfA as + Aassert _ cond as -> rnfE cond `seq` rnfA as + Aassume cond as -> rnfE cond `seq` rnfA as Aforeign repr asm afun a -> rnfTupR rnfArrayR repr `seq` rnf (strForeign asm) `seq` rnfAF afun `seq` rnfA a Acond p a1 a2 -> rnfE p `seq` rnfA a1 `seq` rnfA a2 Awhile p f a -> rnfAF p `seq` rnfAF f `seq` rnfA a @@ -1017,7 +650,7 @@ rnfPreOpenAcc rnfA pacc = FoldSeg i f z a s -> rnfIntegralType i `seq` rnfF f `seq` rnfMaybe rnfE z `seq` rnfA a `seq` rnfA s Scan d f z a -> d `seq` rnfF f `seq` rnfMaybe rnfE z `seq` rnfA a Scan' d f z a -> d `seq` rnfF f `seq` rnfE z `seq` rnfA a - Permute f d p a -> rnfF f `seq` rnfA d `seq` rnfF p `seq` rnfA a + Permute f d a -> rnfMaybe rnfF f `seq` rnfA d `seq` rnfA a Backpermute shr sh f a -> rnfShapeR shr `seq` rnfE sh `seq` rnfF f `seq` rnfA a Stencil sr tp f b a -> let @@ -1044,135 +677,6 @@ rnfBoundary _ Wrap = () rnfBoundary (ArrayR _ tR) (Constant c) = rnfConst tR c rnfBoundary _ (Function f) = rnfOpenFun f -rnfMaybe :: (a -> ()) -> Maybe a -> () -rnfMaybe _ Nothing = () -rnfMaybe f (Just x) = f x - -rnfList :: (a -> ()) -> [a] -> () -rnfList r = go - where - go [] = () - go (x:xs) = r x `seq` go xs - -rnfOpenFun :: OpenFun env aenv t -> () -rnfOpenFun (Body b) = rnfOpenExp b -rnfOpenFun (Lam lhs f) = rnfELeftHandSide lhs `seq` rnfOpenFun f - -rnfOpenExp :: forall env aenv t. OpenExp env aenv t -> () -rnfOpenExp topExp = - let - rnfF :: OpenFun env' aenv' t' -> () - rnfF = rnfOpenFun - - rnfE :: OpenExp env' aenv' t' -> () - rnfE = rnfOpenExp - in - case topExp of - Let lhs bnd body -> rnfELeftHandSide lhs `seq` rnfE bnd `seq` rnfE body - Evar v -> rnfExpVar v - Foreign tp asm f x -> rnfTypeR tp `seq` rnf (strForeign asm) `seq` rnfF f `seq` rnfE x - Const tp c -> c `seq` rnfScalarType tp -- scalars should have (nf == whnf) - Undef tp -> rnfScalarType tp - Pair a b -> rnfE a `seq` rnfE b - Nil -> () - VecPack vecr e -> rnfVecR vecr `seq` rnfE e - VecUnpack vecr e -> rnfVecR vecr `seq` rnfE e - IndexSlice slice slix sh -> rnfSliceIndex slice `seq` rnfE slix `seq` rnfE sh - IndexFull slice slix sl -> rnfSliceIndex slice `seq` rnfE slix `seq` rnfE sl - ToIndex shr sh ix -> rnfShapeR shr `seq` rnfE sh `seq` rnfE ix - FromIndex shr sh ix -> rnfShapeR shr `seq` rnfE sh `seq` rnfE ix - Case e rhs def -> rnfE e `seq` rnfList (\(t,c) -> t `seq` rnfE c) rhs `seq` rnfMaybe rnfE def - Cond p e1 e2 -> rnfE p `seq` rnfE e1 `seq` rnfE e2 - While p f x -> rnfF p `seq` rnfF f `seq` rnfE x - PrimConst c -> rnfPrimConst c - PrimApp f x -> rnfPrimFun f `seq` rnfE x - Index a ix -> rnfArrayVar a `seq` rnfE ix - LinearIndex a ix -> rnfArrayVar a `seq` rnfE ix - Shape a -> rnfArrayVar a - ShapeSize shr sh -> rnfShapeR shr `seq` rnfE sh - Coerce t1 t2 e -> rnfScalarType t1 `seq` rnfScalarType t2 `seq` rnfE e - -rnfExpVar :: ExpVar env t -> () -rnfExpVar = rnfVar rnfScalarType - -rnfELeftHandSide :: ELeftHandSide t env env' -> () -rnfELeftHandSide= rnfLeftHandSide rnfScalarType - -rnfConst :: TypeR t -> t -> () -rnfConst TupRunit () = () -rnfConst (TupRsingle t) !_ = rnfScalarType t -- scalars should have (nf == whnf) -rnfConst (TupRpair ta tb) (a,b) = rnfConst ta a `seq` rnfConst tb b - -rnfPrimConst :: PrimConst c -> () -rnfPrimConst (PrimMinBound t) = rnfBoundedType t -rnfPrimConst (PrimMaxBound t) = rnfBoundedType t -rnfPrimConst (PrimPi t) = rnfFloatingType t - -rnfPrimFun :: PrimFun f -> () -rnfPrimFun (PrimAdd t) = rnfNumType t -rnfPrimFun (PrimSub t) = rnfNumType t -rnfPrimFun (PrimMul t) = rnfNumType t -rnfPrimFun (PrimNeg t) = rnfNumType t -rnfPrimFun (PrimAbs t) = rnfNumType t -rnfPrimFun (PrimSig t) = rnfNumType t -rnfPrimFun (PrimQuot t) = rnfIntegralType t -rnfPrimFun (PrimRem t) = rnfIntegralType t -rnfPrimFun (PrimQuotRem t) = rnfIntegralType t -rnfPrimFun (PrimIDiv t) = rnfIntegralType t -rnfPrimFun (PrimMod t) = rnfIntegralType t -rnfPrimFun (PrimDivMod t) = rnfIntegralType t -rnfPrimFun (PrimBAnd t) = rnfIntegralType t -rnfPrimFun (PrimBOr t) = rnfIntegralType t -rnfPrimFun (PrimBXor t) = rnfIntegralType t -rnfPrimFun (PrimBNot t) = rnfIntegralType t -rnfPrimFun (PrimBShiftL t) = rnfIntegralType t -rnfPrimFun (PrimBShiftR t) = rnfIntegralType t -rnfPrimFun (PrimBRotateL t) = rnfIntegralType t -rnfPrimFun (PrimBRotateR t) = rnfIntegralType t -rnfPrimFun (PrimPopCount t) = rnfIntegralType t -rnfPrimFun (PrimCountLeadingZeros t) = rnfIntegralType t -rnfPrimFun (PrimCountTrailingZeros t) = rnfIntegralType t -rnfPrimFun (PrimFDiv t) = rnfFloatingType t -rnfPrimFun (PrimRecip t) = rnfFloatingType t -rnfPrimFun (PrimSin t) = rnfFloatingType t -rnfPrimFun (PrimCos t) = rnfFloatingType t -rnfPrimFun (PrimTan t) = rnfFloatingType t -rnfPrimFun (PrimAsin t) = rnfFloatingType t -rnfPrimFun (PrimAcos t) = rnfFloatingType t -rnfPrimFun (PrimAtan t) = rnfFloatingType t -rnfPrimFun (PrimSinh t) = rnfFloatingType t -rnfPrimFun (PrimCosh t) = rnfFloatingType t -rnfPrimFun (PrimTanh t) = rnfFloatingType t -rnfPrimFun (PrimAsinh t) = rnfFloatingType t -rnfPrimFun (PrimAcosh t) = rnfFloatingType t -rnfPrimFun (PrimAtanh t) = rnfFloatingType t -rnfPrimFun (PrimExpFloating t) = rnfFloatingType t -rnfPrimFun (PrimSqrt t) = rnfFloatingType t -rnfPrimFun (PrimLog t) = rnfFloatingType t -rnfPrimFun (PrimFPow t) = rnfFloatingType t -rnfPrimFun (PrimLogBase t) = rnfFloatingType t -rnfPrimFun (PrimTruncate f i) = rnfFloatingType f `seq` rnfIntegralType i -rnfPrimFun (PrimRound f i) = rnfFloatingType f `seq` rnfIntegralType i -rnfPrimFun (PrimFloor f i) = rnfFloatingType f `seq` rnfIntegralType i -rnfPrimFun (PrimCeiling f i) = rnfFloatingType f `seq` rnfIntegralType i -rnfPrimFun (PrimIsNaN t) = rnfFloatingType t -rnfPrimFun (PrimIsInfinite t) = rnfFloatingType t -rnfPrimFun (PrimAtan2 t) = rnfFloatingType t -rnfPrimFun (PrimLt t) = rnfSingleType t -rnfPrimFun (PrimGt t) = rnfSingleType t -rnfPrimFun (PrimLtEq t) = rnfSingleType t -rnfPrimFun (PrimGtEq t) = rnfSingleType t -rnfPrimFun (PrimEq t) = rnfSingleType t -rnfPrimFun (PrimNEq t) = rnfSingleType t -rnfPrimFun (PrimMax t) = rnfSingleType t -rnfPrimFun (PrimMin t) = rnfSingleType t -rnfPrimFun PrimLAnd = () -rnfPrimFun PrimLOr = () -rnfPrimFun PrimLNot = () -rnfPrimFun (PrimFromIntegral i n) = rnfIntegralType i `seq` rnfNumType n -rnfPrimFun (PrimToFloating n f) = rnfNumType n `seq` rnfFloatingType f - - -- Template Haskell -- ================ @@ -1201,6 +705,10 @@ liftPreOpenAcc liftA pacc = liftB :: ArrayR (Array sh e) -> Boundary aenv (Array sh e) -> CodeQ (Boundary aenv (Array sh e)) liftB = liftBoundary + + liftMF :: Maybe (OpenFun env aenv t) -> CodeQ (Maybe (OpenFun env aenv t)) + liftMF Nothing = [|| Nothing ||] + liftMF (Just f) = [|| Just $$(liftF f) ||] in case pacc of Alet lhs bnd body -> [|| Alet $$(liftALeftHandSide lhs) $$(liftA bnd) $$(liftA body) ||] @@ -1208,7 +716,9 @@ liftPreOpenAcc liftA pacc = Apair as bs -> [|| Apair $$(liftA as) $$(liftA bs) ||] Anil -> [|| Anil ||] Atrace msg as bs -> [|| Atrace $$(liftMessage (arraysR as) msg) $$(liftA as) $$(liftA bs) ||] - Apply repr f a -> [|| Apply $$(liftArraysR repr) $$(liftAF f) $$(liftA a) ||] + Manifest as -> [|| Manifest $$(liftA as) ||] + Aassert msg cond as -> [|| Aassert msg $$(liftE cond) $$(liftA as) ||] + Aassume cond as -> [|| Aassume $$(liftE cond) $$(liftA as) ||] Aforeign repr asm f a -> [|| Aforeign $$(liftArraysR repr) $$(liftForeign asm) $$(liftPreOpenAfun liftA f) $$(liftA a) ||] Acond p t e -> [|| Acond $$(liftE p) $$(liftA t) $$(liftA e) ||] Awhile p f a -> [|| Awhile $$(liftAF p) $$(liftAF f) $$(liftA a) ||] @@ -1225,7 +735,7 @@ liftPreOpenAcc liftA pacc = FoldSeg i f z a s -> [|| FoldSeg $$(liftIntegralType i) $$(liftF f) $$(liftMaybe liftE z) $$(liftA a) $$(liftA s) ||] Scan d f z a -> [|| Scan $$(liftDirection d) $$(liftF f) $$(liftMaybe liftE z) $$(liftA a) ||] Scan' d f z a -> [|| Scan' $$(liftDirection d) $$(liftF f) $$(liftE z) $$(liftA a) ||] - Permute f d p a -> [|| Permute $$(liftF f) $$(liftA d) $$(liftF p) $$(liftA a) ||] + Permute f d a -> [|| Permute $$(liftMF f) $$(liftA d) $$(liftA a) ||] Backpermute shr sh p a -> [|| Backpermute $$(liftShapeR shr) $$(liftE sh) $$(liftF p) $$(liftA a) ||] Stencil sr tp f b a -> let TupRsingle (ArrayR shr _) = arraysR a @@ -1261,63 +771,6 @@ liftMessage aR (Message _ fmt msg) = in [|| Message $$(fromMaybe (fmtR aR) fmt) Nothing $$(TH.unsafeCodeCoerce (TH.lift msg)) ||] -liftMaybe :: (a -> CodeQ a) -> Maybe a -> CodeQ (Maybe a) -liftMaybe _ Nothing = [|| Nothing ||] -liftMaybe f (Just x) = [|| Just $$(f x) ||] - -liftList :: (a -> CodeQ a) -> [a] -> CodeQ [a] -liftList _ [] = [|| [] ||] -liftList f (x:xs) = [|| $$(f x) : $$(liftList f xs) ||] - -liftOpenFun - :: OpenFun env aenv t - -> CodeQ (OpenFun env aenv t) -liftOpenFun (Lam lhs f) = [|| Lam $$(liftELeftHandSide lhs) $$(liftOpenFun f) ||] -liftOpenFun (Body b) = [|| Body $$(liftOpenExp b) ||] - -liftOpenExp - :: forall env aenv t. - OpenExp env aenv t - -> CodeQ (OpenExp env aenv t) -liftOpenExp pexp = - let - liftE :: OpenExp env aenv e -> CodeQ (OpenExp env aenv e) - liftE = liftOpenExp - - liftF :: OpenFun env aenv f -> CodeQ (OpenFun env aenv f) - liftF = liftOpenFun - in - case pexp of - Let lhs bnd body -> [|| Let $$(liftELeftHandSide lhs) $$(liftOpenExp bnd) $$(liftOpenExp body) ||] - Evar var -> [|| Evar $$(liftExpVar var) ||] - Foreign repr asm f x -> [|| Foreign $$(liftTypeR repr) $$(liftForeign asm) $$(liftOpenFun f) $$(liftE x) ||] - Const tp c -> [|| Const $$(liftScalarType tp) $$(liftElt (TupRsingle tp) c) ||] - Undef tp -> [|| Undef $$(liftScalarType tp) ||] - Pair a b -> [|| Pair $$(liftE a) $$(liftE b) ||] - Nil -> [|| Nil ||] - VecPack vecr e -> [|| VecPack $$(liftVecR vecr) $$(liftE e) ||] - VecUnpack vecr e -> [|| VecUnpack $$(liftVecR vecr) $$(liftE e) ||] - IndexSlice slice slix sh -> [|| IndexSlice $$(liftSliceIndex slice) $$(liftE slix) $$(liftE sh) ||] - IndexFull slice slix sl -> [|| IndexFull $$(liftSliceIndex slice) $$(liftE slix) $$(liftE sl) ||] - ToIndex shr sh ix -> [|| ToIndex $$(liftShapeR shr) $$(liftE sh) $$(liftE ix) ||] - FromIndex shr sh ix -> [|| FromIndex $$(liftShapeR shr) $$(liftE sh) $$(liftE ix) ||] - Case p rhs def -> [|| Case $$(liftE p) $$(liftList (\(t,c) -> [|| (t, $$(liftE c)) ||]) rhs) $$(liftMaybe liftE def) ||] - Cond p t e -> [|| Cond $$(liftE p) $$(liftE t) $$(liftE e) ||] - While p f x -> [|| While $$(liftF p) $$(liftF f) $$(liftE x) ||] - PrimConst t -> [|| PrimConst $$(liftPrimConst t) ||] - PrimApp f x -> [|| PrimApp $$(liftPrimFun f) $$(liftE x) ||] - Index a ix -> [|| Index $$(liftArrayVar a) $$(liftE ix) ||] - LinearIndex a ix -> [|| LinearIndex $$(liftArrayVar a) $$(liftE ix) ||] - Shape a -> [|| Shape $$(liftArrayVar a) ||] - ShapeSize shr ix -> [|| ShapeSize $$(liftShapeR shr) $$(liftE ix) ||] - Coerce t1 t2 e -> [|| Coerce $$(liftScalarType t1) $$(liftScalarType t2) $$(liftE e) ||] - -liftELeftHandSide :: ELeftHandSide t env env' -> CodeQ (ELeftHandSide t env env') -liftELeftHandSide = liftLeftHandSide liftScalarType - -liftExpVar :: ExpVar env t -> CodeQ (ExpVar env t) -liftExpVar = liftVar liftScalarType - liftBoundary :: forall aenv sh e. ArrayR (Array sh e) @@ -1329,76 +782,6 @@ liftBoundary _ Wrap = [|| Wrap ||] liftBoundary (ArrayR _ tp) (Constant v) = [|| Constant $$(liftElt tp v) ||] liftBoundary _ (Function f) = [|| Function $$(liftOpenFun f) ||] -liftPrimConst :: PrimConst c -> CodeQ (PrimConst c) -liftPrimConst (PrimMinBound t) = [|| PrimMinBound $$(liftBoundedType t) ||] -liftPrimConst (PrimMaxBound t) = [|| PrimMaxBound $$(liftBoundedType t) ||] -liftPrimConst (PrimPi t) = [|| PrimPi $$(liftFloatingType t) ||] - -liftPrimFun :: PrimFun f -> CodeQ (PrimFun f) -liftPrimFun (PrimAdd t) = [|| PrimAdd $$(liftNumType t) ||] -liftPrimFun (PrimSub t) = [|| PrimSub $$(liftNumType t) ||] -liftPrimFun (PrimMul t) = [|| PrimMul $$(liftNumType t) ||] -liftPrimFun (PrimNeg t) = [|| PrimNeg $$(liftNumType t) ||] -liftPrimFun (PrimAbs t) = [|| PrimAbs $$(liftNumType t) ||] -liftPrimFun (PrimSig t) = [|| PrimSig $$(liftNumType t) ||] -liftPrimFun (PrimQuot t) = [|| PrimQuot $$(liftIntegralType t) ||] -liftPrimFun (PrimRem t) = [|| PrimRem $$(liftIntegralType t) ||] -liftPrimFun (PrimQuotRem t) = [|| PrimQuotRem $$(liftIntegralType t) ||] -liftPrimFun (PrimIDiv t) = [|| PrimIDiv $$(liftIntegralType t) ||] -liftPrimFun (PrimMod t) = [|| PrimMod $$(liftIntegralType t) ||] -liftPrimFun (PrimDivMod t) = [|| PrimDivMod $$(liftIntegralType t) ||] -liftPrimFun (PrimBAnd t) = [|| PrimBAnd $$(liftIntegralType t) ||] -liftPrimFun (PrimBOr t) = [|| PrimBOr $$(liftIntegralType t) ||] -liftPrimFun (PrimBXor t) = [|| PrimBXor $$(liftIntegralType t) ||] -liftPrimFun (PrimBNot t) = [|| PrimBNot $$(liftIntegralType t) ||] -liftPrimFun (PrimBShiftL t) = [|| PrimBShiftL $$(liftIntegralType t) ||] -liftPrimFun (PrimBShiftR t) = [|| PrimBShiftR $$(liftIntegralType t) ||] -liftPrimFun (PrimBRotateL t) = [|| PrimBRotateL $$(liftIntegralType t) ||] -liftPrimFun (PrimBRotateR t) = [|| PrimBRotateR $$(liftIntegralType t) ||] -liftPrimFun (PrimPopCount t) = [|| PrimPopCount $$(liftIntegralType t) ||] -liftPrimFun (PrimCountLeadingZeros t) = [|| PrimCountLeadingZeros $$(liftIntegralType t) ||] -liftPrimFun (PrimCountTrailingZeros t) = [|| PrimCountTrailingZeros $$(liftIntegralType t) ||] -liftPrimFun (PrimFDiv t) = [|| PrimFDiv $$(liftFloatingType t) ||] -liftPrimFun (PrimRecip t) = [|| PrimRecip $$(liftFloatingType t) ||] -liftPrimFun (PrimSin t) = [|| PrimSin $$(liftFloatingType t) ||] -liftPrimFun (PrimCos t) = [|| PrimCos $$(liftFloatingType t) ||] -liftPrimFun (PrimTan t) = [|| PrimTan $$(liftFloatingType t) ||] -liftPrimFun (PrimAsin t) = [|| PrimAsin $$(liftFloatingType t) ||] -liftPrimFun (PrimAcos t) = [|| PrimAcos $$(liftFloatingType t) ||] -liftPrimFun (PrimAtan t) = [|| PrimAtan $$(liftFloatingType t) ||] -liftPrimFun (PrimSinh t) = [|| PrimSinh $$(liftFloatingType t) ||] -liftPrimFun (PrimCosh t) = [|| PrimCosh $$(liftFloatingType t) ||] -liftPrimFun (PrimTanh t) = [|| PrimTanh $$(liftFloatingType t) ||] -liftPrimFun (PrimAsinh t) = [|| PrimAsinh $$(liftFloatingType t) ||] -liftPrimFun (PrimAcosh t) = [|| PrimAcosh $$(liftFloatingType t) ||] -liftPrimFun (PrimAtanh t) = [|| PrimAtanh $$(liftFloatingType t) ||] -liftPrimFun (PrimExpFloating t) = [|| PrimExpFloating $$(liftFloatingType t) ||] -liftPrimFun (PrimSqrt t) = [|| PrimSqrt $$(liftFloatingType t) ||] -liftPrimFun (PrimLog t) = [|| PrimLog $$(liftFloatingType t) ||] -liftPrimFun (PrimFPow t) = [|| PrimFPow $$(liftFloatingType t) ||] -liftPrimFun (PrimLogBase t) = [|| PrimLogBase $$(liftFloatingType t) ||] -liftPrimFun (PrimTruncate ta tb) = [|| PrimTruncate $$(liftFloatingType ta) $$(liftIntegralType tb) ||] -liftPrimFun (PrimRound ta tb) = [|| PrimRound $$(liftFloatingType ta) $$(liftIntegralType tb) ||] -liftPrimFun (PrimFloor ta tb) = [|| PrimFloor $$(liftFloatingType ta) $$(liftIntegralType tb) ||] -liftPrimFun (PrimCeiling ta tb) = [|| PrimCeiling $$(liftFloatingType ta) $$(liftIntegralType tb) ||] -liftPrimFun (PrimIsNaN t) = [|| PrimIsNaN $$(liftFloatingType t) ||] -liftPrimFun (PrimIsInfinite t) = [|| PrimIsInfinite $$(liftFloatingType t) ||] -liftPrimFun (PrimAtan2 t) = [|| PrimAtan2 $$(liftFloatingType t) ||] -liftPrimFun (PrimLt t) = [|| PrimLt $$(liftSingleType t) ||] -liftPrimFun (PrimGt t) = [|| PrimGt $$(liftSingleType t) ||] -liftPrimFun (PrimLtEq t) = [|| PrimLtEq $$(liftSingleType t) ||] -liftPrimFun (PrimGtEq t) = [|| PrimGtEq $$(liftSingleType t) ||] -liftPrimFun (PrimEq t) = [|| PrimEq $$(liftSingleType t) ||] -liftPrimFun (PrimNEq t) = [|| PrimNEq $$(liftSingleType t) ||] -liftPrimFun (PrimMax t) = [|| PrimMax $$(liftSingleType t) ||] -liftPrimFun (PrimMin t) = [|| PrimMin $$(liftSingleType t) ||] -liftPrimFun PrimLAnd = [|| PrimLAnd ||] -liftPrimFun PrimLOr = [|| PrimLOr ||] -liftPrimFun PrimLNot = [|| PrimLNot ||] -liftPrimFun (PrimFromIntegral ta tb) = [|| PrimFromIntegral $$(liftIntegralType ta) $$(liftNumType tb) ||] -liftPrimFun (PrimToFloating ta tb) = [|| PrimToFloating $$(liftNumType ta) $$(liftFloatingType tb) ||] - - formatDirection :: Format r (Direction -> r) formatDirection = later $ \case LeftToRight -> singleton 'l' @@ -1410,7 +793,9 @@ formatPreAccOp = later $ \case Avar (Var _ ix) -> bformat ("Avar a" % int) (idxToInt ix) Use aR a -> bformat ("Use " % string) (showArrayShort 5 (showsElt (arrayRtype aR)) aR a) Atrace{} -> "Atrace" - Apply{} -> "Apply" + Manifest{} -> "Manifest" + Aassert{} -> "Aassert" + Aassume{} -> "Aassume" Aforeign{} -> "Aforeign" Acond{} -> "Acond" Awhile{} -> "Awhile" @@ -1432,30 +817,3 @@ formatPreAccOp = later $ \case Backpermute{} -> "Backpermute" Stencil{} -> "Stencil" Stencil2{} -> "Stencil2" - -formatExpOp :: Format r (OpenExp aenv env t -> r) -formatExpOp = later $ \case - Let{} -> "Let" - Evar (Var _ ix) -> bformat ("Var x" % int) (idxToInt ix) - Const tp c -> bformat ("Const " % string) (showElt (TupRsingle tp) c) - Undef{} -> "Undef" - Foreign{} -> "Foreign" - Pair{} -> "Pair" - Nil{} -> "Nil" - VecPack{} -> "VecPack" - VecUnpack{} -> "VecUnpack" - IndexSlice{} -> "IndexSlice" - IndexFull{} -> "IndexFull" - ToIndex{} -> "ToIndex" - FromIndex{} -> "FromIndex" - Case{} -> "Case" - Cond{} -> "Cond" - While{} -> "While" - PrimConst{} -> "PrimConst" - PrimApp{} -> "PrimApp" - Index{} -> "Index" - LinearIndex{} -> "LinearIndex" - Shape{} -> "Shape" - ShapeSize{} -> "ShapeSize" - Coerce{} -> "Coerce" - diff --git a/src/Data/Array/Accelerate/AST/CountEnv.hs b/src/Data/Array/Accelerate/AST/CountEnv.hs new file mode 100644 index 000000000..304a2e351 --- /dev/null +++ b/src/Data/Array/Accelerate/AST/CountEnv.hs @@ -0,0 +1,54 @@ +{-# OPTIONS_HADDOCK hide #-} +-- | +-- Module : Data.Array.Accelerate.AST.CountEnv +-- Copyright : [2008..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.AST.CountEnv ( + CountEnv, fromIdxSet, maxCount, findWithCount, findMostFrequent + ) where + +import Data.Array.Accelerate.AST.IdxSet (IdxSet) +import qualified Data.Array.Accelerate.AST.IdxSet as IdxSet +import Data.Array.Accelerate.AST.Environment hiding ( push ) +import Data.List (foldl') + +newtype CountEnv env = CountEnv (PartialEnv Count env) + +newtype Count a = Count Int + +fromIdxSet :: IdxSet env -> CountEnv env +fromIdxSet = CountEnv . mapPartialEnv (const $ Count 1) . IdxSet.unIdxSet + +instance Semigroup (Count a) where + Count x <> Count y = Count $ x + y + +instance Monoid (Count a) where + mempty = Count 0 + +instance Semigroup (CountEnv env) where + CountEnv a <> CountEnv b = CountEnv $ unionPartialEnv (<>) a b + +instance Monoid (CountEnv env) where + mempty = CountEnv PEnd + +maxCount :: CountEnv env -> Int +maxCount (CountEnv env) = foldl' (\x (EnvBinding _ (Count y)) -> max x y) 0 $ partialEnvToList env + +findWithCount :: CountEnv env -> Int -> IdxSet env +findWithCount (CountEnv env) count = case env of + PEnd -> IdxSet.empty + PPush env' (Count c) + | count == c -> IdxSet.push $ findWithCount (CountEnv env') count + | otherwise -> IdxSet.skip $ findWithCount (CountEnv env') count + PNone env' -> IdxSet.skip $ findWithCount (CountEnv env') count + +findMostFrequent :: CountEnv env -> (Int, IdxSet env) +findMostFrequent env = (count, findWithCount env count) + where + count = maxCount env diff --git a/src/Data/Array/Accelerate/AST/Environment.hs b/src/Data/Array/Accelerate/AST/Environment.hs index 442f4e3a2..d40919911 100644 --- a/src/Data/Array/Accelerate/AST/Environment.hs +++ b/src/Data/Array/Accelerate/AST/Environment.hs @@ -1,10 +1,6 @@ -{-# LANGUAGE EmptyCase #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE TypeFamilies #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.AST.Environment @@ -16,64 +12,431 @@ -- Portability : non-portable (GHC extensions) -- -module Data.Array.Accelerate.AST.Environment - where +module Data.Array.Accelerate.AST.Environment ( + Env(.., Push'), PartialEnv(..), Val, IdentityF(..), + push, push', prj, prj', prjVars, prjPartial, + unionPartialEnv, EnvBinding(..), partialEnvFromList, mapPartialEnv, + mapMaybePartialEnv, partialEnvValues, diffPartialEnv, diffPartialEnvWith, + intersectPartialEnv, partialEnvTail, partialEnvLast, partialEnvSkip, + partialUpdate, partialUpdateWith, + partialRemove, partialEnvToList, partialEnvSingleton, partialEnvPush, + partialEnvPushLHS, partialEnvSameKeys, partialEnvSub, partialEnvSkipLHS, + envToPartial, envFromPartialLazy, + + skipWeakenIdx, lhsSkip, + + prjUpdate', prjReplace', update', updates', mapEnv, + (:>)(Weaken), (>:>), weakenId, weakenSucc, weakenSucc', weakenEmpty, weakenReplace, + sink, (.>), sinkWithLHS, weakenWithLHS, substituteLHS, + varsGet, varsGetVal, stripWithLhs, + + Append +) where import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.Var import Data.Array.Accelerate.AST.LeftHandSide import Data.Array.Accelerate.Error - +import Data.Array.Accelerate.Representation.Type +import Data.Typeable ((:~:)(..)) +import Data.Either +import Data.List ( sortOn ) +import Data.Functor.Identity +import Data.Kind (Type) +import Data.String ( fromString ) -- Valuation for an environment -- -data Val env where - Empty :: Val () - Push :: Val env -> t -> Val (env, t) +data Env f env where + Empty :: Env f () + Push :: Env f env -> f t -> Env f (env, t) + +data PartialEnv f env where + PEnd :: PartialEnv f env + PPush :: PartialEnv f env -> f t -> PartialEnv f (env, t) + PNone :: PartialEnv f env -> PartialEnv f (env, t) + +type Val = Env Identity + +pattern Push' :: () => (env' ~ (env, t)) => Val env -> t -> Val env' +pattern Push' env t = Push env (Identity t) +{-# COMPLETE Empty, Push' #-} -- Push a set of variables into an environment -- push :: Val env -> (LeftHandSide s t env env', t) -> Val env' push env (LeftHandSideWildcard _, _ ) = env -push env (LeftHandSideSingle _ , a ) = env `Push` a +push env (LeftHandSideSingle _ , a ) = env `Push` Identity a push env (LeftHandSidePair l1 l2, (a, b)) = push env (l1, a) `push` (l2, b) +push' :: forall f s t env env'. Distributes s => Env f env -> (LeftHandSide s t env env', Distribute f t) -> Env f env' +push' env (LeftHandSideWildcard _, _ ) = env +push' env (LeftHandSidePair l1 l2, (a, b)) = push' env (l1, a) `push'` (l2, b) +push' env (LeftHandSideSingle s , a ) + | Refl <- reprIsSingle @s @t @f s = env `Push` a + -- Projection of a value from a valuation using a de Bruijn index -- prj :: Idx env t -> Val env -> t -prj ZeroIdx (Push _ v) = v -prj (SuccIdx idx) (Push val _) = prj idx val +prj ix v = runIdentity $ prj' ix v +prj' :: Idx env t -> Env f env -> f t +prj' ZeroIdx (Push _ v) = v +prj' (SuccIdx idx) (Push env _) = prj' idx env --- The type of shifting terms from one context into another +prjVars :: forall f s t env. Distributes s => Vars s env t -> Env f env -> Distribute f t +prjVars (TupRpair v1 v2) env = (prjVars v1 env, prjVars v2 env) +prjVars TupRunit _ = () +prjVars (TupRsingle var) env + | Refl <- reprIsSingle @s @t @f (varType var) = prj' (varIdx var) env + +prjPartial :: Idx env t -> PartialEnv f env -> Maybe (f t) +prjPartial ZeroIdx (PPush _ v) = Just v +prjPartial (SuccIdx idx) (PPush env _) = prjPartial idx env +prjPartial (SuccIdx idx) (PNone env) = prjPartial idx env +prjPartial _ _ = Nothing + +unionPartialEnv :: (forall t. f t -> f t -> f t) -> PartialEnv f env -> PartialEnv f env -> PartialEnv f env +unionPartialEnv _ PEnd env = env +unionPartialEnv _ env PEnd = env +unionPartialEnv g (PNone e1 ) (PNone e2 ) = PNone (unionPartialEnv g e1 e2) +unionPartialEnv g (PPush e1 a) (PNone e2 ) = PPush (unionPartialEnv g e1 e2) a +unionPartialEnv g (PNone e1 ) (PPush e2 b) = PPush (unionPartialEnv g e1 e2) b +unionPartialEnv g (PPush e1 a) (PPush e2 b) = PPush (unionPartialEnv g e1 e2) (g a b) + +intersectPartialEnv :: (forall t. f t -> g t -> h t) -> PartialEnv f env -> PartialEnv g env -> PartialEnv h env +intersectPartialEnv _ PEnd _ = PEnd +intersectPartialEnv _ _ PEnd = PEnd +intersectPartialEnv g (PNone e1 ) (PNone e2 ) = PNone (intersectPartialEnv g e1 e2) +intersectPartialEnv g (PPush e1 _) (PNone e2 ) = PNone (intersectPartialEnv g e1 e2) +intersectPartialEnv g (PNone e1 ) (PPush e2 _) = PNone (intersectPartialEnv g e1 e2) +intersectPartialEnv g (PPush e1 a) (PPush e2 b) = PPush (intersectPartialEnv g e1 e2) (g a b) + +-- Creates a partial environment containing only the identifiers which where +-- present in the first env but not in the second env. +diffPartialEnv :: PartialEnv f env -> PartialEnv g env -> PartialEnv f env +diffPartialEnv = diffPartialEnvWith (const $ const Nothing) + +diffPartialEnvWith :: (forall t. f t -> g t -> Maybe (f t)) -> PartialEnv f env -> PartialEnv g env -> PartialEnv f env +diffPartialEnvWith _ PEnd _ = PEnd +diffPartialEnvWith _ env PEnd = env +diffPartialEnvWith g (PNone e1 ) (PNone e2 ) = PNone (diffPartialEnvWith g e1 e2) +diffPartialEnvWith g (PNone e1 ) (PPush e2 _) = PNone (diffPartialEnvWith g e1 e2) +diffPartialEnvWith g (PPush e1 a) (PNone e2 ) = PPush (diffPartialEnvWith g e1 e2) a +diffPartialEnvWith g (PPush e1 a) (PPush e2 b) = case g a b of + Nothing -> PNone (diffPartialEnvWith g e1 e2) + Just x -> PPush (diffPartialEnvWith g e1 e2) x + +partialEnvTail :: PartialEnv f (env, t) -> PartialEnv f env +partialEnvTail (PPush e _) = e +partialEnvTail (PNone e ) = e +partialEnvTail PEnd = PEnd + +partialEnvLast :: PartialEnv f (env, t) -> Maybe (f t) +partialEnvLast (PPush _ a) = Just a +partialEnvLast _ = Nothing + +partialEnvSkip :: PartialEnv f env -> PartialEnv f (env, t) +partialEnvSkip PEnd = PEnd +partialEnvSkip e = PNone e + +partialEnvSkipLHS :: LeftHandSide s t env env' -> PartialEnv f env -> PartialEnv f env' +partialEnvSkipLHS (LeftHandSideSingle _) = partialEnvSkip +partialEnvSkipLHS (LeftHandSideWildcard _) = id +partialEnvSkipLHS (LeftHandSidePair l1 l2) = partialEnvSkipLHS l2 . partialEnvSkipLHS l1 + +partialEnvPush :: PartialEnv f env -> Maybe (f t) -> PartialEnv f (env, t) +partialEnvPush e Nothing = partialEnvSkip e +partialEnvPush e (Just a) = PPush e a + +partialEnvPushLHS :: LeftHandSide s t env env' -> TupR f t -> PartialEnv f env -> PartialEnv f env' +partialEnvPushLHS (LeftHandSideSingle _) (TupRsingle a) env = PPush env a +partialEnvPushLHS (LeftHandSideWildcard _) _ env = env +partialEnvPushLHS (LeftHandSidePair l1 l2) (TupRpair a1 a2) env = + partialEnvPushLHS l2 a2 $ partialEnvPushLHS l1 a1 env +partialEnvPushLHS _ _ _ = internalError "Tuple mismatch" + +partialUpdate :: f t -> Idx env t -> PartialEnv f env -> PartialEnv f env +partialUpdate = partialUpdateWith const + +partialUpdateWith :: (f t -> f t -> f t) -> f t -> Idx env t -> PartialEnv f env -> PartialEnv f env +partialUpdateWith g v ZeroIdx = \case + PEnd -> PPush PEnd v + PNone e -> PPush e v + PPush e old -> PPush e (g v old) +partialUpdateWith g v (SuccIdx idx) = \case + PEnd -> PNone (partialUpdateWith g v idx PEnd) + PNone e -> PNone (partialUpdateWith g v idx e) + PPush e a -> PPush (partialUpdateWith g v idx e) a + +partialRemove :: Idx env t -> PartialEnv f env -> PartialEnv f env +partialRemove ZeroIdx env = partialEnvSkip $ partialEnvTail env +partialRemove (SuccIdx idx) (PPush e a) = PPush (partialRemove idx e) a +partialRemove (SuccIdx idx) (PNone e ) = PNone (partialRemove idx e) +partialRemove (SuccIdx _ ) PEnd = PEnd + +partialEnvSameKeys :: PartialEnv f env -> PartialEnv g env -> Bool +partialEnvSameKeys PEnd PEnd = True +partialEnvSameKeys (PPush a _) (PPush b _) = partialEnvSameKeys a b +partialEnvSameKeys (PNone a ) (PNone b ) = partialEnvSameKeys a b +partialEnvSameKeys (PNone a ) PEnd = partialEnvSameKeys a PEnd +partialEnvSameKeys PEnd (PNone b ) = partialEnvSameKeys PEnd b +partialEnvSameKeys _ _ = False + +partialEnvSub :: PartialEnv f env -> PartialEnv g env -> Bool +partialEnvSub (PPush a _) (PPush b _) = partialEnvSub a b +partialEnvSub (PNone a) b = partialEnvSub a $ partialEnvTail b +partialEnvSub PEnd _ = True +partialEnvSub _ _ = False + +-- TODO: This is actually isomorphic to Var. Should we remove this data type? +data EnvBinding f env where + EnvBinding :: Idx env t -> f t -> EnvBinding f env + +partialEnvFromList :: forall f env. (forall t. f t -> f t -> f t) -> [EnvBinding f env] -> PartialEnv f env +-- Pair with ints to make sorting faster. +partialEnvFromList g = go SkipNone . map snd . sortOn fst . map (\b@(EnvBinding ix _) -> (idxToInt ix, b)) + where + go :: Skip env env' -> [EnvBinding f env] -> PartialEnv f env' + go _ [] = PEnd + go skip (EnvBinding idx v : bs) = go' skip (EnvBinding idx' v) bs + where + idx' = case skipIdx skip idx of + Just i -> i + Nothing -> error "List wasn't sorted properly" + + go' :: Skip env env' -> EnvBinding f env' -> [EnvBinding f env] -> PartialEnv f env' + go' skip (EnvBinding (SuccIdx idx) v) bs = PNone $ go' (SkipSucc skip) (EnvBinding idx v) bs + go' skip (EnvBinding ZeroIdx v) bs = go'' skip v bs + + go'' :: Skip env (env', t) -> f t -> [EnvBinding f env] -> PartialEnv f (env', t) + go'' _ v [] = PPush PEnd v + go'' skip v (EnvBinding idx v' : bs) = case skipIdx skip idx of + Nothing -> error "List wasn't sorted properly" + Just ZeroIdx -> go'' skip (g v v') bs + Just (SuccIdx idx') -> PPush (go' (SkipSucc skip) (EnvBinding idx' v') bs) v + +partialEnvToList :: forall f env. PartialEnv f env -> [EnvBinding f env] +partialEnvToList = go weakenId + where + go :: env' :> env -> PartialEnv f env' -> [EnvBinding f env] + go _ PEnd = [] + go k (PNone env) = go (weakenSucc k) env + go k (PPush env a) = EnvBinding (k >:> ZeroIdx) a : go (weakenSucc k) env + +mapPartialEnv :: (forall t. a t -> b t) -> PartialEnv a env -> PartialEnv b env +mapPartialEnv _ PEnd = PEnd +mapPartialEnv f (PPush env a) = PPush (mapPartialEnv f env) (f a) +mapPartialEnv f (PNone env) = PNone (mapPartialEnv f env) + +mapMaybePartialEnv :: (forall t. a t -> Maybe (b t)) -> PartialEnv a env -> PartialEnv b env +mapMaybePartialEnv _ PEnd = PEnd +mapMaybePartialEnv f (PPush env a) = case f a of + Nothing -> PNone (mapMaybePartialEnv f env) + Just b -> PPush (mapMaybePartialEnv f env) b +mapMaybePartialEnv f (PNone env) = PNone (mapMaybePartialEnv f env) + +partialEnvValues :: PartialEnv (IdentityF a) env -> [a] +partialEnvValues PEnd = [] +partialEnvValues (PNone env) = partialEnvValues env +partialEnvValues (PPush env (IdentityF a)) = a : partialEnvValues env + +partialEnvSingleton :: Idx env t -> f t -> PartialEnv f env +partialEnvSingleton ZeroIdx v = PPush PEnd v +partialEnvSingleton (SuccIdx idx) v = PNone $ partialEnvSingleton idx v + +envToPartial :: Env f env -> PartialEnv f env +envToPartial Empty = PEnd +envToPartial (Push env a) = envToPartial env `PPush` a + +-- Lazily converts a PartialEnv to an Env. +-- Laziness is here used to allow the environment to be used safely +-- even if some bindings are missing in the PartialEnv. +-- Only when a missing binding is actually used in this Env, +-- it will throw an error. +envFromPartialLazy :: HasCallStack => String -> PartialEnv f env -> Env f env +envFromPartialLazy msg = \case + PEnd -> internalError (fromString msg) + PNone env -> envFromPartialLazy msg env `Push` internalError (fromString msg) + PPush env a -> envFromPartialLazy msg env `Push` a + +-- Wrapper to put homogenous types in an Env or PartialEnv +newtype IdentityF t f = IdentityF t + +skipWeakenIdx :: Skip env env' -> env' :> env +skipWeakenIdx s = WeakenSkip s WeakenId + +lhsSkip :: forall s t env1 env2. LeftHandSide s t env1 env2 -> Skip env2 env1 +lhsSkip = (`go` SkipNone) + where + go :: LeftHandSide s t' env env' -> Skip env2 env' -> Skip env2 env + go (LeftHandSideSingle _) accum = SkipSucc accum + go (LeftHandSideWildcard _) accum = accum + go (LeftHandSidePair l1 l2) accum = go l1 $ go l2 accum + +prjUpdate' :: (f t -> (f t, a)) -> Idx env t -> Env f env -> (Env f env, a) +prjUpdate' f ZeroIdx (Push env v) = (Push env v', a) + where (v', a) = f v +prjUpdate' f (SuccIdx idx) (Push env v) = (Push env' v, a) + where (env', a) = prjUpdate' f idx env + +prjReplace' :: Idx env t -> f t -> Env f env -> (Env f env, f t) +prjReplace' ix val = prjUpdate' (val, ) ix + +update' :: (f t -> f t) -> Idx env t -> Env f env -> Env f env +update' f ZeroIdx (Push env v) = Push env (f v) +update' f (SuccIdx idx) (Push env v) = Push (update' f idx env) v + +-- Updates an environment for a list of indices. This is prefered over +-- calling update' multiple times, as that will create many 'Push' objects +-- as garbage. +-- Duplicate indices are handled only once. -- --- This is defined as a newtype, as a type synonym containing a forall --- quantifier may give issues with impredicative polymorphism, which GHC --- does not support. +updates' :: forall a f t env. (a -> f t -> f t) -> [(Idx env t, a)] -> Env f env -> Env f env +updates' _ [] env = env +updates' f indices (Push env (v :: f s)) = updates' f succs env `Push` v' + where + g :: (Idx (env', s) t, a) -> Either (f s) (Idx env' t, a) + g (ZeroIdx, a) = Left $ f a v + g (SuccIdx idx, a) = Right (idx, a) + (zeros, succs) = partitionEithers $ map g indices + v' :: f s + v' = case zeros of + v'' : _ -> v'' -- This index is in the list + [] -> v -- The index wasn't found +updates' _ _ Empty = Empty + +mapEnv :: (forall t. a t -> b t) -> Env a env -> Env b env +mapEnv _ Empty = Empty +mapEnv g (Push env f) = Push (mapEnv g env) (g f) + + +-- The type of shifting terms from one context into another -- -newtype env :> env' = Weaken { (>:>) :: forall t'. Idx env t' -> Idx env' t' } -- Weak or Weaken +-- This used to be defined as a newtype over a function (the current Weaken +-- constructor), but for higher performance we partially defunctionalized this +-- newtype into a small set of constructors. +data env :> env' where + WeakenId :: env :> env + + WeakenEmpty :: () :> env' + +#ifdef ACCELERATE_INTERNAL_CHECKS + + -- The GADT representations cannot be @UNPACK@ed. + + WeakenReplace + :: env :> env' + -> !(Idx env' t) + -> (env, t) :> env' + + WeakenKeep + :: env1 :> env2 + -> !(Keep env1 env2 env1' env2') + -> env1' :> env2' + + WeakenSkip + :: !(Skip env2 env1) + -> env2 :> env3 + -> env1 :> env3 + + WeakenSkip' + :: env1 :> env2 + -> !(Skip env3 env2) + -> env1 :> env3 + +#else + + WeakenReplace + :: env :> env' + -> {-# UNPACK #-} !(Idx env' t) + -> (env, t) :> env' + + WeakenKeep + :: env1 :> env2 + -> {-# UNPACK #-} !(Keep env1 env2 env1' env2') + -> env1' :> env2' + + WeakenSkip + :: {-# UNPACK #-} !(Skip env2 env1) + -> env2 :> env3 + -> env1 :> env3 + + WeakenSkip' + :: env1 :> env2 + -> {-# UNPACK #-} !(Skip env3 env2) + -> env1 :> env3 + +#endif + + WeakenChain + :: env1 :> env2 + -> env2 :> env3 + -> env1 :> env3 + + Weaken :: (forall (t' :: Type). Idx env t' -> Idx env' t') -> env :> env' + +(>:>) :: env :> env' -> Idx env t -> Idx env' t +WeakenId >:> idx = idx +WeakenEmpty >:> VoidIdx x = x +WeakenReplace _ idx >:> ZeroIdx = idx +WeakenReplace k _ >:> SuccIdx idx = k >:> idx +WeakenSkip s k >:> idx = k >:> unskipIdx s idx +WeakenSkip' k s >:> idx = unskipIdx s (k >:> idx) +WeakenKeep k keep >:> idx = case keepIdx keep idx of + Left idx' -> unkeepIdx keep $ k >:> idx' + Right idx' -> idx' +WeakenChain k1 k2 >:> idx = k2 >:> (k1 >:> idx) +Weaken f >:> idx = f idx weakenId :: env :> env -weakenId = Weaken id +weakenId = WeakenId weakenSucc' :: env :> env' -> env :> (env', t) -weakenSucc' (Weaken f) = Weaken (SuccIdx . f) +weakenSucc' (WeakenSkip' k s) = WeakenSkip' k (SkipSucc' s) +weakenSucc' (WeakenSkip s WeakenId) = WeakenSkip' WeakenId (SkipSucc' s) +weakenSucc' k = WeakenSkip' k (SkipSucc SkipNone) weakenSucc :: (env, t) :> env' -> env :> env' -weakenSucc (Weaken f) = Weaken (f . SuccIdx) +weakenSucc (WeakenSkip s k) = WeakenSkip (SkipSucc s) k +weakenSucc (WeakenSkip' WeakenId s) = WeakenSkip (SkipSucc s) WeakenId +weakenSucc k = WeakenSkip (SkipSucc SkipNone) k + +sink :: env :> env' -> (env, t) :> (env', t) +sink WeakenId = WeakenId +sink (WeakenKeep k keep) = WeakenKeep k $ KeepSucc keep +sink k = WeakenKeep k (KeepSucc KeepNone) weakenEmpty :: () :> env' -weakenEmpty = Weaken $ \(VoidIdx x) -> x +weakenEmpty = WeakenEmpty -sink :: forall env env' t. env :> env' -> (env, t) :> (env', t) -sink (Weaken f) = Weaken g - where - g :: Idx (env, t) t' -> Idx (env', t) t' - g ZeroIdx = ZeroIdx - g (SuccIdx ix) = SuccIdx $ f ix +weakenReplace :: forall env env' t. Idx env' t -> env :> env' -> (env, t) :> env' +weakenReplace other k = WeakenReplace k other infixr 9 .> +-- | Compose two weakenings (.>) :: env2 :> env3 -> env1 :> env2 -> env1 :> env3 -(.>) (Weaken f) (Weaken g) = Weaken (f . g) +-- Instead of always emitting a WeakenChain, we check for various simple cases +-- to return a simpler weakening. +-- Check WeakenId +(.>) WeakenId k = k +(.>) k WeakenId = k +-- Combine skips +(.>) (WeakenSkip s1 WeakenId) (WeakenSkip' k2 s2) = WeakenSkip' k2 (chainSkip s1 s2) +(.>) (WeakenSkip' WeakenId s1) (WeakenSkip' k2 s2) = WeakenSkip' k2 (chainSkip s1 s2) +(.>) (WeakenSkip s1 k1) (WeakenSkip s2 WeakenId) = WeakenSkip (chainSkip s1 s2) k1 +(.>) (WeakenSkip s1 k1) (WeakenSkip' WeakenId s2) = WeakenSkip (chainSkip s1 s2) k1 +-- Single skip +(.>) (WeakenSkip s1 WeakenId) k2 = WeakenSkip' k2 s1 +(.>) (WeakenSkip' WeakenId s1) k2 = WeakenSkip' k2 s1 +(.>) k1 (WeakenSkip s2 WeakenId) = WeakenSkip s2 k1 +(.>) k1 (WeakenSkip' WeakenId s2) = WeakenSkip s2 k1 +-- Move Skip outwards (as we might be able to combine it with a skip in a later call to (.>)). +-- Directly construct WeakenChain instead of recursing, as optimizing it further probably has no additional benefits. +(.>) k1 (WeakenSkip s2 k2) = WeakenSkip s2 (WeakenChain k2 k1) +(.>) (WeakenSkip' k1 s1) k2 = WeakenSkip' (WeakenChain k2 k1) s1 +-- Default case +(.>) k1 k2 = WeakenChain k2 k1 sinkWithLHS :: HasCallStack => LeftHandSide s t env1 env1' -> LeftHandSide s t env2 env2' -> env1 :> env2 -> env1' :> env2' sinkWithLHS (LeftHandSideWildcard _) (LeftHandSideWildcard _) k = k @@ -89,3 +452,38 @@ weakenWithLHS = go weakenId go k (LeftHandSideSingle _) = weakenSucc k go k (LeftHandSidePair l1 l2) = go (go k l2) l1 +substituteLHS :: forall s s' t env env'. LeftHandSide s t env env' -> Vars s' env t -> env' :> env +substituteLHS lhs vars = Weaken f + -- TODO: Could we convert this to a sequence of WeakenReplace constructors, + -- instead of Weaken? + where + f :: Idx env' a -> Idx env a + f ix = case go lhs vars ix of + Left ix' -> ix' + Right ix' -> ix' + + go :: LeftHandSide s u env1 env2 -> Vars s' env u -> Idx env2 a -> Either (Idx env a) (Idx env1 a) + go (LeftHandSideWildcard _) _ idx = Right idx + go (LeftHandSideSingle _) (TupRsingle (Var _ var)) idx = case idx of + ZeroIdx -> Left var + SuccIdx idx' -> Right idx' + go (LeftHandSidePair l1 l2) (TupRpair v1 v2) idx = go l2 v2 idx >>= go l1 v1 + go _ _ _ = error "LHS and tuple mismatch" + + +varsGet :: Applicative f => Vars s env t -> Env f env -> f t +varsGet TupRunit _ = pure () +varsGet (TupRsingle (Var _ idx)) env = prj' idx env +varsGet (TupRpair v1 v2) env = (,) <$> varsGet v1 env <*> varsGet v2 env + +varsGetVal :: Vars s env t -> Val env -> t +varsGetVal vars = runIdentity . varsGet vars + +stripWithLhs :: LeftHandSide f a env env' -> Env g env' -> Env g env +stripWithLhs (LeftHandSideSingle _) (Push env _) = env +stripWithLhs (LeftHandSideWildcard _) env = env +stripWithLhs (LeftHandSidePair lhs1 lhs2) env = stripWithLhs lhs1 $ stripWithLhs lhs2 env + +type family Append env1 env2 where + Append env1 () = env1 + Append env1 (env2, t) = (Append env1 env2, t) diff --git a/src/Data/Array/Accelerate/AST/Execute.hs b/src/Data/Array/Accelerate/AST/Execute.hs new file mode 100644 index 000000000..490ba99e9 --- /dev/null +++ b/src/Data/Array/Accelerate/AST/Execute.hs @@ -0,0 +1,37 @@ +{-# LANGUAGE TypeFamilies #-} + +-- | +-- Module : Data.Array.Accelerate.AST.Kernel +-- Copyright : [2012..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- +-- This module defines the interface to execute progams, to be implemented +-- in other modules for different backends. +-- + +module Data.Array.Accelerate.AST.Execute ( + Execute(..), + executeAfun, + GFunctionR(..) +) where + +import Data.Array.Accelerate.AST.Schedule +import Data.Array.Accelerate.Representation.Ground +import Data.Kind + +class Execute sched kernel where + data Linked sched kernel :: Type -> Type + -- For a backend that doesn't require a linking stage, + -- one can directly store a schedule in the Linked data type. + -- data Linked sched kernel t = MyBackendLinked (sched kernel () t) + + linkAfunSchedule :: sched kernel () t -> Linked sched kernel t + + executeAfunSchedule :: GFunctionR t -> Linked sched kernel (Scheduled sched t) -> IOFun (Scheduled sched t) + +executeAfun :: forall sched kernel t. (IsSchedule sched, Execute sched kernel) => GFunctionR t -> Linked sched kernel (Scheduled sched t) -> IOFun t +executeAfun repr sched = flattenIOFun repr $ callScheduledFun @sched repr $ executeAfunSchedule @sched @kernel repr sched diff --git a/src/Data/Array/Accelerate/AST/Exp.hs b/src/Data/Array/Accelerate/AST/Exp.hs new file mode 100644 index 000000000..b8f5d2da4 --- /dev/null +++ b/src/Data/Array/Accelerate/AST/Exp.hs @@ -0,0 +1,795 @@ +{-# LANGUAGE TemplateHaskellQuotes #-} +{-# OPTIONS_HADDOCK hide #-} +-- | +-- Module : Data.Array.Accelerate.AST.Exp +-- Copyright : [2008..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.AST.Exp ( + -- ** Scalar expressions + ELeftHandSide, ExpVar, ExpVars, expVars, undefs, + PreOpenFun(..), + PreOpenExp(..), + PrimFun(..), + Cmp(..), negateCmp, + PrimBool, + PrimMaybe, + TAG, + IsArrayInstr(..), + NoArrayInstr, + + -- ** Extracting type information + expType, + primFunType, + + -- ** Normal-form + rnfOpenFun, + rnfOpenExp, + rnfConst, + rnfPrimFun, + rnfMaybe, + rnfELeftHandSide, + rnfExpVar, + + -- ** Template Haskell + liftOpenFun, + liftOpenExp, + liftMaybe, + liftELeftHandSide, + liftExpVar, + liftPrimFun, + + -- ** Miscellaneous + mkConstant, mkBinary, unBody, + formatExpOp, Direction(..), + expIsTrivial, shapeExpVars, + +) where + +import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.AST.Var +import Data.Array.Accelerate.Error +import Data.Array.Accelerate.Representation.Elt +import Data.Array.Accelerate.Representation.Shape +import Data.Array.Accelerate.Representation.Tag +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Representation.Vec +import Data.Array.Accelerate.Sugar.Foreign +import Data.Array.Accelerate.Type +import Data.Primitive.Vec +import Data.String +import Data.Typeable ( (:~:) ) +import Data.ByteString.Builder ( Builder ) +import Data.Text ( Text ) + +import Control.DeepSeq +import Language.Haskell.TH.Extra ( CodeQ ) +import Formatting +import Prelude + +import GHC.TypeLits + + +-- Bool is not a primitive type +type PrimBool = TAG +type PrimMaybe a = (TAG, ((), a)) + +-- Embedded expressions +-- -------------------- + +-- | Vanilla open function abstraction parameterized over array instructions +-- +data PreOpenFun arr env t where + Body :: PreOpenExp arr env t -> PreOpenFun arr env t + Lam :: ELeftHandSide a env env' -> PreOpenFun arr env' t -> PreOpenFun arr env (a -> t) + +-- | Vanilla open expressions using de Bruijn indices for variables ranging +-- over tuples of scalars and arrays of tuples. All code, except Cond, is +-- evaluated eagerly. N-tuples are represented as nested pairs. +-- +-- It is parametrised over the type of array instructions, 'arr'. +-- +-- The data type is parametrised over the representation type (not the +-- surface types). +-- +data PreOpenExp arr env t where + + -- Local binding of a scalar expression + Let :: ELeftHandSide bnd_t env env' + -> PreOpenExp arr env bnd_t + -> PreOpenExp arr env' body_t + -> PreOpenExp arr env body_t + + -- Variable index, ranging _only_ over scalars + Evar :: ExpVar env t + -> PreOpenExp arr env t + + -- Apply a backend-specific foreign function + Foreign :: Foreign asm + => TypeR y + -> asm (x -> y) -- foreign function + -> PreOpenFun NoArrayInstr () (x -> y) -- alternate implementation (for other backends) + -> PreOpenExp arr env x + -> PreOpenExp arr env y + + -- Tuples + Pair :: PreOpenExp arr env t1 + -> PreOpenExp arr env t2 + -> PreOpenExp arr env (t1, t2) + + Nil :: PreOpenExp arr env () + + -- SIMD vectors + VecPack :: KnownNat n + => VecR n s tup + -> PreOpenExp arr env tup + -> PreOpenExp arr env (Vec n s) + + VecUnpack :: KnownNat n + => VecR n s tup + -> PreOpenExp arr env (Vec n s) + -> PreOpenExp arr env tup + + -- Shape and index conversion + ToIndex :: ShapeR sh + -> PreOpenExp arr env sh -- shape of the array + -> PreOpenExp arr env sh -- index into the array + -> PreOpenExp arr env Int + + FromIndex :: ShapeR sh + -> PreOpenExp arr env sh -- shape of the array + -> PreOpenExp arr env Int -- index into linear representation + -> PreOpenExp arr env sh + + -- Case statement + Case :: PreOpenExp arr env TAG + -> [(TAG, PreOpenExp arr env b)] -- list of equations + -> Maybe (PreOpenExp arr env b) -- default case + -> PreOpenExp arr env b + + -- Conditional expression (non-strict in 2nd and 3rd argument) + Cond :: PreOpenExp arr env PrimBool + -> PreOpenExp arr env t + -> PreOpenExp arr env t + -> PreOpenExp arr env t + + -- Conditional expression that chooses one value based on the condition without branching. + Select :: PreOpenExp arr env PrimBool + -> PreOpenExp arr env t + -> PreOpenExp arr env t + -> PreOpenExp arr env t + + Assert :: Text + -> PreOpenExp arr env PrimBool + -> PreOpenExp arr env t + -> PreOpenExp arr env t + + Assume :: PreOpenExp arr env PrimBool + -> PreOpenExp arr env t + -> PreOpenExp arr env t + + -- Value recursion + While :: PreOpenFun arr env (a -> PrimBool) -- continue while true + -> PreOpenFun arr env (a -> a) -- function to iterate + -> PreOpenExp arr env a -- initial value + -> PreOpenExp arr env a + + -- Constant values + Const :: ScalarType t + -> t + -> PreOpenExp arr env t + + -- Primitive scalar operations + PrimApp :: PrimFun (a -> r) + -> PreOpenExp arr env a + -> PreOpenExp arr env r + + ArrayInstr :: arr (a -> b) + -> PreOpenExp arr env a + -> PreOpenExp arr env b + + -- Number of elements of an array given its shape + ShapeSize :: ShapeR dim + -> PreOpenExp arr env dim + -> PreOpenExp arr env Int + + -- Unsafe operations (may fail or result in undefined behaviour) + -- An unspecified bit pattern + Undef :: ScalarType t + -> PreOpenExp arr env t + + -- Reinterpret the bits of a value as a different type + Coerce :: BitSizeEq a b + => ScalarType a + -> ScalarType b + -> PreOpenExp arr env a + -> PreOpenExp arr env b + +expVars :: ExpVars env t -> PreOpenExp arr env t +expVars TupRunit = Nil +expVars (TupRsingle var) = Evar var +expVars (TupRpair v1 v2) = expVars v1 `Pair` expVars v2 + +undefs :: TypeR t -> PreOpenExp arr env t +undefs (TupRsingle tp) = Undef tp +undefs (TupRpair t1 t2) = undefs t1 `Pair` undefs t2 +undefs TupRunit = Nil + +-- Types for scalar bindings +-- +type ELeftHandSide = LeftHandSide ScalarType +type ExpVar = Var ScalarType +type ExpVars env = Vars ScalarType env + +-- | The expressions are parametrised over the type of array instructions. +-- Type class 'IsArrayInstr' should be implemented for that type. +class IsArrayInstr arr where + arrayInstrType :: arr (a -> b) -> TypeR b + liftArrayInstr :: arr t -> CodeQ (arr t) + rnfArrayInstr :: arr t -> () + showArrayInstrOp :: arr t -> String + encodeArrayInstr :: arr a -> Builder + matchArrayInstr :: arr t -> arr s -> Maybe (t :~: s) + -- | Whether the array instruction may be inlined (and duplicated) + -- True for Parameter and False for Index + inlineArrayInstr :: arr t -> Bool + +-- | Allows no array instructions at all. Used for expressions +-- which cannot perform array operations. The only use case +-- is currently the fallback function of a foreign. +-- +data NoArrayInstr t where + +instance IsArrayInstr NoArrayInstr where + arrayInstrType = \case {} + liftArrayInstr = \case {} + rnfArrayInstr = \case {} + showArrayInstrOp = \case {} + encodeArrayInstr = \case {} + matchArrayInstr = \case {} + inlineArrayInstr = \case {} + +data Direction = LeftToRight | RightToLeft + deriving Eq + + +-- |Primitive scalar operations +-- +data PrimFun sig where + + -- operators from Num + PrimAdd :: NumType a -> PrimFun ((a, a) -> a) + PrimSub :: NumType a -> PrimFun ((a, a) -> a) + PrimMul :: NumType a -> PrimFun ((a, a) -> a) + PrimNeg :: NumType a -> PrimFun (a -> a) + PrimAbs :: NumType a -> PrimFun (a -> a) + PrimSig :: NumType a -> PrimFun (a -> a) + + -- operators from Integral + PrimQuot :: IntegralType a -> PrimFun ((a, a) -> a) + PrimRem :: IntegralType a -> PrimFun ((a, a) -> a) + PrimQuotRem :: IntegralType a -> PrimFun ((a, a) -> (a, a)) + PrimIDiv :: IntegralType a -> PrimFun ((a, a) -> a) + PrimMod :: IntegralType a -> PrimFun ((a, a) -> a) + PrimDivMod :: IntegralType a -> PrimFun ((a, a) -> (a, a)) + + -- operators from Bits & FiniteBits + PrimBAnd :: IntegralType a -> PrimFun ((a, a) -> a) + PrimBOr :: IntegralType a -> PrimFun ((a, a) -> a) + PrimBXor :: IntegralType a -> PrimFun ((a, a) -> a) + PrimBNot :: IntegralType a -> PrimFun (a -> a) + PrimBShiftL :: IntegralType a -> PrimFun ((a, Int) -> a) + PrimBShiftR :: IntegralType a -> PrimFun ((a, Int) -> a) + PrimBRotateL :: IntegralType a -> PrimFun ((a, Int) -> a) + PrimBRotateR :: IntegralType a -> PrimFun ((a, Int) -> a) + PrimPopCount :: IntegralType a -> PrimFun (a -> Int) + PrimCountLeadingZeros :: IntegralType a -> PrimFun (a -> Int) + PrimCountTrailingZeros :: IntegralType a -> PrimFun (a -> Int) + + -- operators from Fractional and Floating + PrimFDiv :: FloatingType a -> PrimFun ((a, a) -> a) + PrimRecip :: FloatingType a -> PrimFun (a -> a) + PrimSin :: FloatingType a -> PrimFun (a -> a) + PrimCos :: FloatingType a -> PrimFun (a -> a) + PrimTan :: FloatingType a -> PrimFun (a -> a) + PrimAsin :: FloatingType a -> PrimFun (a -> a) + PrimAcos :: FloatingType a -> PrimFun (a -> a) + PrimAtan :: FloatingType a -> PrimFun (a -> a) + PrimSinh :: FloatingType a -> PrimFun (a -> a) + PrimCosh :: FloatingType a -> PrimFun (a -> a) + PrimTanh :: FloatingType a -> PrimFun (a -> a) + PrimAsinh :: FloatingType a -> PrimFun (a -> a) + PrimAcosh :: FloatingType a -> PrimFun (a -> a) + PrimAtanh :: FloatingType a -> PrimFun (a -> a) + PrimExpFloating :: FloatingType a -> PrimFun (a -> a) + PrimSqrt :: FloatingType a -> PrimFun (a -> a) + PrimLog :: FloatingType a -> PrimFun (a -> a) + PrimFPow :: FloatingType a -> PrimFun ((a, a) -> a) + PrimLogBase :: FloatingType a -> PrimFun ((a, a) -> a) + + -- FIXME: add missing operations from RealFrac & RealFloat + + -- operators from RealFrac + PrimTruncate :: FloatingType a -> IntegralType b -> PrimFun (a -> b) + PrimRound :: FloatingType a -> IntegralType b -> PrimFun (a -> b) + PrimFloor :: FloatingType a -> IntegralType b -> PrimFun (a -> b) + PrimCeiling :: FloatingType a -> IntegralType b -> PrimFun (a -> b) + -- PrimProperFraction :: FloatingType a -> IntegralType b -> PrimFun (a -> (b, a)) + + -- operators from RealFloat + PrimAtan2 :: FloatingType a -> PrimFun ((a, a) -> a) + PrimIsNaN :: FloatingType a -> PrimFun (a -> PrimBool) + PrimIsInfinite :: FloatingType a -> PrimFun (a -> PrimBool) + + -- relational and equality operators + PrimCmp :: SingleType a -> Cmp -> PrimFun ((a, a) -> PrimBool) + PrimMax :: SingleType a -> PrimFun ((a, a) -> a) + PrimMin :: SingleType a -> PrimFun ((a, a) -> a) + + -- logical operators + -- + -- Note that these operators are strict in both arguments. That is, the + -- second argument of PrimLAnd is always evaluated even when the first + -- argument is false. + -- + -- We define (surface level) (&&) and (||) using if-then-else to enable + -- short-circuiting, while (&&!) and (||!) are strict versions of these + -- operators, which are defined using PrimLAnd and PrimLOr. + -- + PrimLAnd :: PrimFun ((PrimBool, PrimBool) -> PrimBool) + PrimLOr :: PrimFun ((PrimBool, PrimBool) -> PrimBool) + PrimLNot :: PrimFun (PrimBool -> PrimBool) + + -- general conversion between types + PrimFromIntegral :: IntegralType a -> NumType b -> PrimFun (a -> b) + PrimToFloating :: NumType a -> FloatingType b -> PrimFun (a -> b) + +-- | Comparison operator +-- Greater-than and less-than-equal are missing, as they are equal to +-- respectively less-than and greater-than-equal with their arguments swapped. +-- We kept less-than and greater-than-equal as they are each others negation +-- (similar to equal and not-equal). +-- +data Cmp = CmpLt | CmpGtEq | CmpEq | CmpNEq deriving Eq + +negateCmp :: Cmp -> Cmp +negateCmp CmpLt = CmpGtEq +negateCmp CmpGtEq = CmpLt +negateCmp CmpEq = CmpNEq +negateCmp CmpNEq = CmpEq + +expType :: (HasCallStack, IsArrayInstr arr) => PreOpenExp arr env t -> TypeR t +expType = \case + Let _ _ body -> expType body + Evar (Var tR _) -> TupRsingle tR + Foreign tR _ _ _ -> tR + Pair e1 e2 -> TupRpair (expType e1) (expType e2) + Nil -> TupRunit + VecPack vecR _ -> TupRsingle $ VectorScalarType $ vecRvector vecR + VecUnpack vecR _ -> vecRtuple vecR + ToIndex{} -> TupRsingle scalarTypeInt + FromIndex shr _ _ -> shapeType shr + Case _ ((_,e):_) _ -> expType e + Case _ [] (Just e) -> expType e + Case{} -> internalError "empty case encountered" + Cond _ e _ -> expType e + Select _ e _ -> expType e + While _ (Lam lhs _) _ -> lhsToTupR lhs + While{} -> internalError "What's the matter, you're running in the shadows" + Const tR _ -> TupRsingle tR + PrimApp f _ -> snd $ primFunType f + ArrayInstr arr _ -> arrayInstrType arr + ShapeSize{} -> TupRsingle scalarTypeInt + Undef tR -> TupRsingle tR + Coerce _ tR _ -> TupRsingle tR + Assert _ _ e2 -> expType e2 + Assume _ e2 -> expType e2 + +primFunType :: PrimFun (a -> b) -> (TypeR a, TypeR b) +primFunType = \case + -- Num + PrimAdd t -> binary' $ num t + PrimSub t -> binary' $ num t + PrimMul t -> binary' $ num t + PrimNeg t -> unary' $ num t + PrimAbs t -> unary' $ num t + PrimSig t -> unary' $ num t + + -- Integral + PrimQuot t -> binary' $ integral t + PrimRem t -> binary' $ integral t + PrimQuotRem t -> unary' $ integral t `TupRpair` integral t + PrimIDiv t -> binary' $ integral t + PrimMod t -> binary' $ integral t + PrimDivMod t -> unary' $ integral t `TupRpair` integral t + + -- Bits & FiniteBits + PrimBAnd t -> binary' $ integral t + PrimBOr t -> binary' $ integral t + PrimBXor t -> binary' $ integral t + PrimBNot t -> unary' $ integral t + PrimBShiftL t -> (integral t `TupRpair` tint, integral t) + PrimBShiftR t -> (integral t `TupRpair` tint, integral t) + PrimBRotateL t -> (integral t `TupRpair` tint, integral t) + PrimBRotateR t -> (integral t `TupRpair` tint, integral t) + PrimPopCount t -> unary (integral t) tint + PrimCountLeadingZeros t -> unary (integral t) tint + PrimCountTrailingZeros t -> unary (integral t) tint + + -- Fractional, Floating + PrimFDiv t -> binary' $ floating t + PrimRecip t -> unary' $ floating t + PrimSin t -> unary' $ floating t + PrimCos t -> unary' $ floating t + PrimTan t -> unary' $ floating t + PrimAsin t -> unary' $ floating t + PrimAcos t -> unary' $ floating t + PrimAtan t -> unary' $ floating t + PrimSinh t -> unary' $ floating t + PrimCosh t -> unary' $ floating t + PrimTanh t -> unary' $ floating t + PrimAsinh t -> unary' $ floating t + PrimAcosh t -> unary' $ floating t + PrimAtanh t -> unary' $ floating t + PrimExpFloating t -> unary' $ floating t + PrimSqrt t -> unary' $ floating t + PrimLog t -> unary' $ floating t + PrimFPow t -> binary' $ floating t + PrimLogBase t -> binary' $ floating t + + -- RealFrac + PrimTruncate a b -> unary (floating a) (integral b) + PrimRound a b -> unary (floating a) (integral b) + PrimFloor a b -> unary (floating a) (integral b) + PrimCeiling a b -> unary (floating a) (integral b) + + -- RealFloat + PrimAtan2 t -> binary' $ floating t + PrimIsNaN t -> unary (floating t) tbool + PrimIsInfinite t -> unary (floating t) tbool + + -- Relational and equality + PrimCmp t _ -> compare' t + PrimMax t -> binary' $ single t + PrimMin t -> binary' $ single t + + -- Logical + PrimLAnd -> binary' tbool + PrimLOr -> binary' tbool + PrimLNot -> unary' tbool + + -- general conversion between types + PrimFromIntegral a b -> unary (integral a) (num b) + PrimToFloating a b -> unary (num a) (floating b) + + where + unary a b = (a, b) + unary' a = unary a a + binary a b = (a `TupRpair` a, b) + binary' a = binary a a + compare' a = binary (single a) tbool + + single = TupRsingle . SingleScalarType + num = TupRsingle . SingleScalarType . NumSingleType + integral = num . IntegralNumType + floating = num . FloatingNumType + + tbool = TupRsingle scalarTypeWord8 + tint = TupRsingle scalarTypeInt + +rnfOpenFun :: forall arr env t. IsArrayInstr arr => PreOpenFun arr env t -> () +rnfOpenFun (Body b) = rnfOpenExp b +rnfOpenFun (Lam lhs f) = rnfELeftHandSide lhs `seq` rnfOpenFun f + +rnfOpenExp :: forall arr env t. IsArrayInstr arr => PreOpenExp arr env t -> () +rnfOpenExp topExp = + let + rnfF :: PreOpenFun arr env' t' -> () + rnfF = rnfOpenFun + + rnfE :: PreOpenExp arr env' t' -> () + rnfE = rnfOpenExp + in + case topExp of + Let lhs bnd body -> rnfELeftHandSide lhs `seq` rnfE bnd `seq` rnfE body + Evar v -> rnfExpVar v + Foreign tp asm f x -> rnfTypeR tp `seq` rnf (strForeign asm) `seq` rnfOpenFun f `seq` rnfE x + Const tp c -> c `seq` rnfScalarType tp -- scalars should have (nf == whnf) + Undef tp -> rnfScalarType tp + Pair a b -> rnfE a `seq` rnfE b + Nil -> () + VecPack vecr e -> rnfVecR vecr `seq` rnfE e + VecUnpack vecr e -> rnfVecR vecr `seq` rnfE e + ToIndex shr sh ix -> rnfShapeR shr `seq` rnfE sh `seq` rnfE ix + FromIndex shr sh ix -> rnfShapeR shr `seq` rnfE sh `seq` rnfE ix + Case e rhs def -> rnfE e `seq` rnfList (\(t,c) -> t `seq` rnfE c) rhs `seq` rnfMaybe rnfE def + Cond p e1 e2 -> rnfE p `seq` rnfE e1 `seq` rnfE e2 + Select p e1 e2 -> rnfE p `seq` rnfE e1 `seq` rnfE e2 + While p f x -> rnfF p `seq` rnfF f `seq` rnfE x + PrimApp f x -> rnfPrimFun f `seq` rnfE x + ArrayInstr arr e -> rnfArrayInstr arr `seq` rnfE e + ShapeSize shr sh -> rnfShapeR shr `seq` rnfE sh + Coerce t1 t2 e -> rnfScalarType t1 `seq` rnfScalarType t2 `seq` rnfE e + Assert msg e1 e2 -> rnf msg `seq` rnfE e1 `seq` rnfE e2 + Assume e1 e2 -> rnfE e1 `seq` rnfE e2 + +rnfExpVar :: ExpVar env t -> () +rnfExpVar = rnfVar rnfScalarType + +rnfELeftHandSide :: ELeftHandSide t env env' -> () +rnfELeftHandSide = rnfLeftHandSide rnfScalarType + +rnfConst :: TypeR t -> t -> () +rnfConst TupRunit () = () +rnfConst (TupRsingle t) !_ = rnfScalarType t -- scalars should have (nf == whnf) +rnfConst (TupRpair ta tb) (a,b) = rnfConst ta a `seq` rnfConst tb b + +rnfPrimFun :: PrimFun f -> () +rnfPrimFun (PrimAdd t) = rnfNumType t +rnfPrimFun (PrimSub t) = rnfNumType t +rnfPrimFun (PrimMul t) = rnfNumType t +rnfPrimFun (PrimNeg t) = rnfNumType t +rnfPrimFun (PrimAbs t) = rnfNumType t +rnfPrimFun (PrimSig t) = rnfNumType t +rnfPrimFun (PrimQuot t) = rnfIntegralType t +rnfPrimFun (PrimRem t) = rnfIntegralType t +rnfPrimFun (PrimQuotRem t) = rnfIntegralType t +rnfPrimFun (PrimIDiv t) = rnfIntegralType t +rnfPrimFun (PrimMod t) = rnfIntegralType t +rnfPrimFun (PrimDivMod t) = rnfIntegralType t +rnfPrimFun (PrimBAnd t) = rnfIntegralType t +rnfPrimFun (PrimBOr t) = rnfIntegralType t +rnfPrimFun (PrimBXor t) = rnfIntegralType t +rnfPrimFun (PrimBNot t) = rnfIntegralType t +rnfPrimFun (PrimBShiftL t) = rnfIntegralType t +rnfPrimFun (PrimBShiftR t) = rnfIntegralType t +rnfPrimFun (PrimBRotateL t) = rnfIntegralType t +rnfPrimFun (PrimBRotateR t) = rnfIntegralType t +rnfPrimFun (PrimPopCount t) = rnfIntegralType t +rnfPrimFun (PrimCountLeadingZeros t) = rnfIntegralType t +rnfPrimFun (PrimCountTrailingZeros t) = rnfIntegralType t +rnfPrimFun (PrimFDiv t) = rnfFloatingType t +rnfPrimFun (PrimRecip t) = rnfFloatingType t +rnfPrimFun (PrimSin t) = rnfFloatingType t +rnfPrimFun (PrimCos t) = rnfFloatingType t +rnfPrimFun (PrimTan t) = rnfFloatingType t +rnfPrimFun (PrimAsin t) = rnfFloatingType t +rnfPrimFun (PrimAcos t) = rnfFloatingType t +rnfPrimFun (PrimAtan t) = rnfFloatingType t +rnfPrimFun (PrimSinh t) = rnfFloatingType t +rnfPrimFun (PrimCosh t) = rnfFloatingType t +rnfPrimFun (PrimTanh t) = rnfFloatingType t +rnfPrimFun (PrimAsinh t) = rnfFloatingType t +rnfPrimFun (PrimAcosh t) = rnfFloatingType t +rnfPrimFun (PrimAtanh t) = rnfFloatingType t +rnfPrimFun (PrimExpFloating t) = rnfFloatingType t +rnfPrimFun (PrimSqrt t) = rnfFloatingType t +rnfPrimFun (PrimLog t) = rnfFloatingType t +rnfPrimFun (PrimFPow t) = rnfFloatingType t +rnfPrimFun (PrimLogBase t) = rnfFloatingType t +rnfPrimFun (PrimTruncate f i) = rnfFloatingType f `seq` rnfIntegralType i +rnfPrimFun (PrimRound f i) = rnfFloatingType f `seq` rnfIntegralType i +rnfPrimFun (PrimFloor f i) = rnfFloatingType f `seq` rnfIntegralType i +rnfPrimFun (PrimCeiling f i) = rnfFloatingType f `seq` rnfIntegralType i +rnfPrimFun (PrimIsNaN t) = rnfFloatingType t +rnfPrimFun (PrimIsInfinite t) = rnfFloatingType t +rnfPrimFun (PrimAtan2 t) = rnfFloatingType t +rnfPrimFun (PrimCmp t !_) = rnfSingleType t +rnfPrimFun (PrimMax t) = rnfSingleType t +rnfPrimFun (PrimMin t) = rnfSingleType t +rnfPrimFun PrimLAnd = () +rnfPrimFun PrimLOr = () +rnfPrimFun PrimLNot = () +rnfPrimFun (PrimFromIntegral i n) = rnfIntegralType i `seq` rnfNumType n +rnfPrimFun (PrimToFloating n f) = rnfNumType n `seq` rnfFloatingType f + +rnfMaybe :: (a -> ()) -> Maybe a -> () +rnfMaybe _ Nothing = () +rnfMaybe f (Just x) = f x + +rnfList :: (a -> ()) -> [a] -> () +rnfList r = go + where + go [] = () + go (x:xs) = r x `seq` go xs + +-- Template Haskell +-- ================ + +liftOpenFun + :: IsArrayInstr arr + => PreOpenFun arr env t + -> CodeQ (PreOpenFun arr env t) +liftOpenFun (Lam lhs f) = [|| Lam $$(liftELeftHandSide lhs) $$(liftOpenFun f) ||] +liftOpenFun (Body b) = [|| Body $$(liftOpenExp b) ||] + +liftOpenExp + :: forall arr env t. IsArrayInstr arr + => PreOpenExp arr env t + -> CodeQ (PreOpenExp arr env t) +liftOpenExp pexp = + let + liftE :: PreOpenExp arr env e -> CodeQ (PreOpenExp arr env e) + liftE = liftOpenExp + + liftF :: PreOpenFun arr env f -> CodeQ (PreOpenFun arr env f) + liftF = liftOpenFun + in + case pexp of + Let lhs bnd body -> [|| Let $$(liftELeftHandSide lhs) $$(liftOpenExp bnd) $$(liftOpenExp body) ||] + Evar var -> [|| Evar $$(liftExpVar var) ||] + Foreign repr asm f x -> [|| Foreign $$(liftTypeR repr) $$(liftForeign asm) $$(liftOpenFun f) $$(liftE x) ||] + Const tp c -> [|| Const $$(liftScalarType tp) $$(liftElt (TupRsingle tp) c) ||] + Undef tp -> [|| Undef $$(liftScalarType tp) ||] + Pair a b -> [|| Pair $$(liftE a) $$(liftE b) ||] + Nil -> [|| Nil ||] + VecPack vecr e -> [|| VecPack $$(liftVecR vecr) $$(liftE e) ||] + VecUnpack vecr e -> [|| VecUnpack $$(liftVecR vecr) $$(liftE e) ||] + ToIndex shr sh ix -> [|| ToIndex $$(liftShapeR shr) $$(liftE sh) $$(liftE ix) ||] + FromIndex shr sh ix -> [|| FromIndex $$(liftShapeR shr) $$(liftE sh) $$(liftE ix) ||] + Case p rhs def -> [|| Case $$(liftE p) $$(liftList (\(t,c) -> [|| (t, $$(liftE c)) ||]) rhs) $$(liftMaybe liftE def) ||] + Cond p t e -> [|| Cond $$(liftE p) $$(liftE t) $$(liftE e) ||] + Select p t e -> [|| Select $$(liftE p) $$(liftE t) $$(liftE e) ||] + While p f x -> [|| While $$(liftF p) $$(liftF f) $$(liftE x) ||] + PrimApp f x -> [|| PrimApp $$(liftPrimFun f) $$(liftE x) ||] + ArrayInstr arr x -> [|| ArrayInstr $$(liftArrayInstr arr) $$(liftE x) ||] + ShapeSize shr ix -> [|| ShapeSize $$(liftShapeR shr) $$(liftE ix) ||] + Coerce t1 t2 e -> [|| Coerce $$(liftScalarType t1) $$(liftScalarType t2) $$(liftE e) ||] + Assert msg e1 e2 -> [|| Assert msg $$(liftE e1) $$(liftE e2) ||] + Assume e1 e2 -> [|| Assume $$(liftE e1) $$(liftE e2) ||] + +liftPrimFun :: PrimFun f -> CodeQ (PrimFun f) +liftPrimFun (PrimAdd t) = [|| PrimAdd $$(liftNumType t) ||] +liftPrimFun (PrimSub t) = [|| PrimSub $$(liftNumType t) ||] +liftPrimFun (PrimMul t) = [|| PrimMul $$(liftNumType t) ||] +liftPrimFun (PrimNeg t) = [|| PrimNeg $$(liftNumType t) ||] +liftPrimFun (PrimAbs t) = [|| PrimAbs $$(liftNumType t) ||] +liftPrimFun (PrimSig t) = [|| PrimSig $$(liftNumType t) ||] +liftPrimFun (PrimQuot t) = [|| PrimQuot $$(liftIntegralType t) ||] +liftPrimFun (PrimRem t) = [|| PrimRem $$(liftIntegralType t) ||] +liftPrimFun (PrimQuotRem t) = [|| PrimQuotRem $$(liftIntegralType t) ||] +liftPrimFun (PrimIDiv t) = [|| PrimIDiv $$(liftIntegralType t) ||] +liftPrimFun (PrimMod t) = [|| PrimMod $$(liftIntegralType t) ||] +liftPrimFun (PrimDivMod t) = [|| PrimDivMod $$(liftIntegralType t) ||] +liftPrimFun (PrimBAnd t) = [|| PrimBAnd $$(liftIntegralType t) ||] +liftPrimFun (PrimBOr t) = [|| PrimBOr $$(liftIntegralType t) ||] +liftPrimFun (PrimBXor t) = [|| PrimBXor $$(liftIntegralType t) ||] +liftPrimFun (PrimBNot t) = [|| PrimBNot $$(liftIntegralType t) ||] +liftPrimFun (PrimBShiftL t) = [|| PrimBShiftL $$(liftIntegralType t) ||] +liftPrimFun (PrimBShiftR t) = [|| PrimBShiftR $$(liftIntegralType t) ||] +liftPrimFun (PrimBRotateL t) = [|| PrimBRotateL $$(liftIntegralType t) ||] +liftPrimFun (PrimBRotateR t) = [|| PrimBRotateR $$(liftIntegralType t) ||] +liftPrimFun (PrimPopCount t) = [|| PrimPopCount $$(liftIntegralType t) ||] +liftPrimFun (PrimCountLeadingZeros t) = [|| PrimCountLeadingZeros $$(liftIntegralType t) ||] +liftPrimFun (PrimCountTrailingZeros t) = [|| PrimCountTrailingZeros $$(liftIntegralType t) ||] +liftPrimFun (PrimFDiv t) = [|| PrimFDiv $$(liftFloatingType t) ||] +liftPrimFun (PrimRecip t) = [|| PrimRecip $$(liftFloatingType t) ||] +liftPrimFun (PrimSin t) = [|| PrimSin $$(liftFloatingType t) ||] +liftPrimFun (PrimCos t) = [|| PrimCos $$(liftFloatingType t) ||] +liftPrimFun (PrimTan t) = [|| PrimTan $$(liftFloatingType t) ||] +liftPrimFun (PrimAsin t) = [|| PrimAsin $$(liftFloatingType t) ||] +liftPrimFun (PrimAcos t) = [|| PrimAcos $$(liftFloatingType t) ||] +liftPrimFun (PrimAtan t) = [|| PrimAtan $$(liftFloatingType t) ||] +liftPrimFun (PrimSinh t) = [|| PrimSinh $$(liftFloatingType t) ||] +liftPrimFun (PrimCosh t) = [|| PrimCosh $$(liftFloatingType t) ||] +liftPrimFun (PrimTanh t) = [|| PrimTanh $$(liftFloatingType t) ||] +liftPrimFun (PrimAsinh t) = [|| PrimAsinh $$(liftFloatingType t) ||] +liftPrimFun (PrimAcosh t) = [|| PrimAcosh $$(liftFloatingType t) ||] +liftPrimFun (PrimAtanh t) = [|| PrimAtanh $$(liftFloatingType t) ||] +liftPrimFun (PrimExpFloating t) = [|| PrimExpFloating $$(liftFloatingType t) ||] +liftPrimFun (PrimSqrt t) = [|| PrimSqrt $$(liftFloatingType t) ||] +liftPrimFun (PrimLog t) = [|| PrimLog $$(liftFloatingType t) ||] +liftPrimFun (PrimFPow t) = [|| PrimFPow $$(liftFloatingType t) ||] +liftPrimFun (PrimLogBase t) = [|| PrimLogBase $$(liftFloatingType t) ||] +liftPrimFun (PrimTruncate ta tb) = [|| PrimTruncate $$(liftFloatingType ta) $$(liftIntegralType tb) ||] +liftPrimFun (PrimRound ta tb) = [|| PrimRound $$(liftFloatingType ta) $$(liftIntegralType tb) ||] +liftPrimFun (PrimFloor ta tb) = [|| PrimFloor $$(liftFloatingType ta) $$(liftIntegralType tb) ||] +liftPrimFun (PrimCeiling ta tb) = [|| PrimCeiling $$(liftFloatingType ta) $$(liftIntegralType tb) ||] +liftPrimFun (PrimIsNaN t) = [|| PrimIsNaN $$(liftFloatingType t) ||] +liftPrimFun (PrimIsInfinite t) = [|| PrimIsInfinite $$(liftFloatingType t) ||] +liftPrimFun (PrimAtan2 t) = [|| PrimAtan2 $$(liftFloatingType t) ||] +liftPrimFun (PrimCmp t c) = [|| PrimCmp $$(liftSingleType t) $$(liftCmp c) ||] +liftPrimFun (PrimMax t) = [|| PrimMax $$(liftSingleType t) ||] +liftPrimFun (PrimMin t) = [|| PrimMin $$(liftSingleType t) ||] +liftPrimFun PrimLAnd = [|| PrimLAnd ||] +liftPrimFun PrimLOr = [|| PrimLOr ||] +liftPrimFun PrimLNot = [|| PrimLNot ||] +liftPrimFun (PrimFromIntegral ta tb) = [|| PrimFromIntegral $$(liftIntegralType ta) $$(liftNumType tb) ||] +liftPrimFun (PrimToFloating ta tb) = [|| PrimToFloating $$(liftNumType ta) $$(liftFloatingType tb) ||] + +liftMaybe :: (a -> CodeQ a) -> Maybe a -> CodeQ (Maybe a) +liftMaybe _ Nothing = [|| Nothing ||] +liftMaybe f (Just x) = [|| Just $$(f x) ||] + +liftList :: (a -> CodeQ a) -> [a] -> CodeQ [a] +liftList _ [] = [|| [] ||] +liftList f (x:xs) = [|| $$(f x) : $$(liftList f xs) ||] + +liftELeftHandSide :: ELeftHandSide t env env' -> CodeQ (ELeftHandSide t env env') +liftELeftHandSide = liftLeftHandSide liftScalarType + +liftExpVar :: ExpVar env t -> CodeQ (ExpVar env t) +liftExpVar = liftVar liftScalarType + +liftCmp :: Cmp -> CodeQ Cmp +liftCmp CmpLt = [|| CmpLt ||] +liftCmp CmpGtEq = [|| CmpGtEq ||] +liftCmp CmpEq = [|| CmpEq ||] +liftCmp CmpNEq = [|| CmpNEq ||] + +mkConstant :: TypeR t -> t -> PreOpenExp arr env t +mkConstant (TupRsingle tp) c = Const tp c +mkConstant (TupRpair t1 t2) (c1, c2) = mkConstant t1 c1 `Pair` mkConstant t2 c2 +mkConstant TupRunit () = Nil + +mkBinary :: PrimFun ((a, b) -> c) -> PreOpenExp arr env a -> PreOpenExp arr env b -> PreOpenExp arr env c +mkBinary fun a b = PrimApp fun $ Pair a b + +-- Extracts the body from a PreOpenFun. The type gives evidence +-- that 't' is not a function type and the PreOpenFun thus must +-- be a Body. +-- +unBody :: TypeR t -> PreOpenFun arr env t -> PreOpenExp arr env t +unBody _ (Body e) = e +unBody tp (Lam _ _) = functionImpossible tp + +formatExpOp :: IsArrayInstr arr => Format r (PreOpenExp arr env t -> r) +formatExpOp = later $ \case + Let{} -> "Let" + Evar (Var _ ix) -> bformat ("Var x" % int) (idxToInt ix) + Const tp c -> bformat ("Const " % string) (showElt (TupRsingle tp) c) + Undef{} -> "Undef" + Foreign{} -> "Foreign" + Pair{} -> "Pair" + Nil{} -> "Nil" + VecPack{} -> "VecPack" + VecUnpack{} -> "VecUnpack" + ToIndex{} -> "ToIndex" + FromIndex{} -> "FromIndex" + Case{} -> "Case" + Cond{} -> "Cond" + Select{} -> "Select" + While{} -> "While" + PrimApp{} -> "PrimApp" + ArrayInstr ar _ -> fromString $ showArrayInstrOp ar + ShapeSize{} -> "ShapeSize" + Coerce{} -> "Coerce" + Assert{} -> "Assert" + Assume{} -> "Assume" + +expIsTrivial :: forall arr env t. (forall s. arr s -> Bool) -> PreOpenExp arr env t -> Bool +expIsTrivial arrayInstr = \case + Let _ bnd body -> trav bnd && trav body + Evar{} -> True + Const{} -> True + Undef{} -> True + Pair a b -> trav a && trav b + Nil -> True + VecPack _ e -> trav e + VecUnpack _ e -> trav e + ToIndex _ a b -> trav a && trav b + FromIndex _ a b -> trav a && trav b + Case scrutinee alts def -> trav scrutinee && all (trav . snd) alts && all trav def + Cond c t f -> trav c && trav t && trav f + PrimApp _ a -> trav a + ArrayInstr ar a -> arrayInstr ar && trav a + ShapeSize _ a -> trav a + Coerce _ _ a -> trav a + Assert _ cond e -> trav cond && trav e + Assume _ e -> trav e -- Condition is not executed, so does not have to be trivial + _ -> False + where + trav :: PreOpenExp arr env' t' -> Bool + trav = expIsTrivial arrayInstr + +shapeExpVars :: Distributes s => ShapeR sh -> Vars s env sh -> ExpVars env sh +shapeExpVars ShapeRz TupRunit = TupRunit +shapeExpVars (ShapeRsnoc shr) (sh `TupRpair` TupRsingle var) = shapeExpVars shr sh `TupRpair` TupRsingle var{varType = scalarTypeInt} +shapeExpVars ShapeRz (TupRsingle (Var tp _)) = unitImpossible tp +shapeExpVars (ShapeRsnoc _) (TupRsingle (Var tp _)) = pairImpossible tp diff --git a/src/Data/Array/Accelerate/AST/Graph.hs b/src/Data/Array/Accelerate/AST/Graph.hs new file mode 100644 index 000000000..10a099b40 --- /dev/null +++ b/src/Data/Array/Accelerate/AST/Graph.hs @@ -0,0 +1,322 @@ +{-# OPTIONS_HADDOCK hide #-} +-- | +-- Module : Data.Array.Accelerate.AST.Graph +-- Copyright : [2008..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.AST.Graph ( + Graph, InEdge(..), empty, nodeCount, pushNode, dropNode, + insertEdge, insertEdgeWith, + insertEdgesFromWith, insertEdgesToWith, + insertCartesianEdgesWith, removeEdgesOf, + out, inn, prjNode, prjEdge, updateNode, + + shortestPath, shortestPathLessThanEqual, shortestPathsFrom, shortestPathsTo +) where + +import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.Environment + +import Data.Functor.Const +import Data.Maybe +import Data.Typeable ( (:~:)(..) ) + +-- | Directed well-scoped graph. +-- 'node' is the label on a node, 'edge' the label on an edge. +-- +data Graph node edge env where + GEmpty :: Graph node edge () + GPush :: Graph node edge env -> Node node edge env t -> Graph node edge (env, t) + +-- An edge (s, t) is stored either in node s or in node t. +-- If s <= t (based on the integral value if the Idx), then it is stored in s, +-- as an outgoing edge. +-- If t < s, then it is stored in t as an incomming edge. + +data Node node edge env a = Node + -- Data of this node + (node a) + -- Incomming edges + (PartialEnv (InEdge edge a) env) + -- Outgoing edges, potentially with a self-loop + (PartialEnv (edge a) (env, a)) + +newtype InEdge edge t s = InEdge (edge s t) + +empty :: Graph node edge () +empty = GEmpty + +nodeCount :: Graph node edge env -> Int +nodeCount GEmpty = 0 +nodeCount (GPush g _) = nodeCount g + 1 + +pushNode :: Graph node edge env -> node a -> PartialEnv (InEdge edge a) env -> PartialEnv (edge a) (env, a) -> Graph node edge (env, a) +pushNode g n i o = GPush g $ Node n i o + +dropNode :: Graph node edge (env, t) -> Graph node edge env +dropNode (GPush g _) = g + +insertEdge :: Idx env s -> Idx env t -> edge s t -> Graph node edge env -> Graph node edge env +insertEdge = insertEdgeWith const + +-- | Inserts an edge in the graph. If the edge already exists, their labels will be combined +insertEdgeWith :: (edge s t -> edge s t -> edge s t) -> Idx env s -> Idx env t -> edge s t -> Graph node edge env -> Graph node edge env +insertEdgeWith f i1 i2 edge (GPush g node) = case (i1, i2) of + (SuccIdx j1, SuccIdx j2) -> insertEdgeWith f j1 j2 edge g `GPush` node + (ZeroIdx , _) + | Node n i o <- node -> GPush g $ Node n i $ partialUpdateWith f edge i2 o + (SuccIdx j1, ZeroIdx) + | Node n i o <- node -> GPush g $ Node n (partialUpdateWith f' (InEdge edge) j1 i) o + where + f' (InEdge e1) (InEdge e2) = InEdge $ f e1 e2 +insertEdgeWith _ i1 _ _ GEmpty = case i1 of {} + +insertEdgesFromWith + :: forall node edge env s'. + (forall s t. edge s t -> edge s t -> edge s t) -- Combine with existing edge, if needed + -> Idx env s' -- Starting point + -> PartialEnv (edge s') env -- End points + -> Graph node edge env + -> Graph node edge env +insertEdgesFromWith f from = insertCartesianEdgesWith + f + (\Refl edge -> edge) + (partialEnvSingleton from Refl :: PartialEnv ((:~:) s') env) + +insertEdgesToWith + :: forall node edge env t'. + (forall s t. edge s t -> edge s t -> edge s t) -- Combine with existing edge, if needed + -> PartialEnv (InEdge edge t') env -- Start points + -> Idx env t' -- End point + -> Graph node edge env + -> Graph node edge env +insertEdgesToWith f from to = insertCartesianEdgesWith + f + (\(InEdge edge) Refl -> edge) + from + (partialEnvSingleton to Refl :: PartialEnv ((:~:) t') env) + +-- Given two partial environments, adds +insertCartesianEdgesWith + :: (forall s t. edge s t -> edge s t -> edge s t) -- Combine with existing edge, if needed + -> (forall s t. from s -> to t -> edge s t) -- Construct the new edge + -> PartialEnv from env -- Starting points of edges + -> PartialEnv to env -- End points of edges + -> Graph node edge env + -> Graph node edge env +insertCartesianEdgesWith _ _ PEnd _ graph = graph +insertCartesianEdgesWith _ _ _ PEnd graph = graph +insertCartesianEdgesWith f edge (PPush starts start) (PPush ends end) (GPush graph (Node n i o)) + | i' <- unionPartialEnv (\(InEdge a) (InEdge b) -> InEdge $ f a b) i $ mapPartialEnv (\s -> InEdge $ edge s end) starts + , o' <- unionPartialEnv f o $ PNone $ mapPartialEnv (edge start) ends + = insertCartesianEdgesWith f edge starts ends graph `GPush` Node n i' o' +insertCartesianEdgesWith f edge (PPush starts start) (PNone ends) (GPush graph (Node n i o)) + | o' <- unionPartialEnv f o $ PNone $ mapPartialEnv (edge start) ends + = insertCartesianEdgesWith f edge starts ends graph `GPush` Node n i o' +insertCartesianEdgesWith f edge (PNone starts) (PPush ends end) (GPush graph (Node n i o)) + | i' <- unionPartialEnv (\(InEdge a) (InEdge b) -> InEdge $ f a b) i $ mapPartialEnv (\s -> InEdge $ edge s end) starts + = insertCartesianEdgesWith f edge starts ends graph `GPush` Node n i' o +insertCartesianEdgesWith f edge (PNone starts) (PNone ends) (GPush graph node) = + insertCartesianEdgesWith f edge starts ends graph `GPush` node + +removeEdgesOf :: Idx env s -> Graph node edge env -> Graph node edge env +removeEdgesOf ZeroIdx (GPush graph (Node n _ _)) + = GPush graph $ Node n PEnd PEnd +removeEdgesOf (SuccIdx idx) (GPush graph (Node n i o)) + = removeEdgesOf idx graph `GPush` Node n (partialRemove idx i) (partialRemove (SuccIdx idx) o) + +out :: Graph node edge env -> Idx env s -> PartialEnv (edge s) env +out (GPush g (Node _ i _)) (SuccIdx idx) = case prjPartial idx i of + Nothing -> partialEnvSkip $ out g idx + Just (InEdge edge) -> out g idx `PPush` edge +out (GPush _ (Node _ _ o)) ZeroIdx = o +out GEmpty idx = case idx of {} + +inn :: Graph node edge env -> Idx env t -> PartialEnv (InEdge edge t) env +inn (GPush g (Node _ _ o)) (SuccIdx idx) = case prjPartial (SuccIdx idx) o of + Nothing -> partialEnvSkip $ inn g idx + Just edge -> inn g idx `PPush` InEdge edge +inn (GPush _ (Node _ i _)) ZeroIdx = partialEnvSkip i +inn GEmpty idx = case idx of {} + +prjNode :: Graph node edge env -> Idx env t -> node t +prjNode (GPush g (Node n _ _)) = \case + ZeroIdx -> n + SuccIdx idx -> prjNode g idx +prjNode GEmpty = \case {} + +prjEdge :: Graph node edge env -> Idx env s -> Idx env t -> Maybe (edge s t) +prjEdge (GPush g _) (SuccIdx s) (SuccIdx t) = prjEdge g s t +prjEdge (GPush _ (Node _ _ o)) ZeroIdx t = prjPartial t o +prjEdge (GPush _ (Node _ i _)) (SuccIdx s) ZeroIdx = case prjPartial s i of + Just (InEdge edge) -> Just edge + Nothing -> Nothing +prjEdge GEmpty s _ = case s of {} + +updateNode :: Idx env t -> (node t -> node t) -> Graph node edge env -> Graph node edge env +updateNode ZeroIdx f (GPush g (Node n i o)) = GPush g $ Node (f n) i o +updateNode (SuccIdx idx) f (GPush g node) = updateNode idx f g `GPush` node + +shortestPath + :: (Show dist, Num dist, Ord dist) + => (forall a b. edge a b -> dist) + -> Graph node edge env + -> Idx env s + -> Idx env t + -> Maybe dist +shortestPath edgeDist graph from to = fmap getConst $ prjPartial to dists + where + (_, dists, _) = shortestPathGeneric From (const Nothing) edgeDist graph from + +shortestPathLessThanEqual + :: (Show dist, Num dist, Ord dist) + => (forall a b. edge a b -> dist) + -> Graph node edge env + -> Idx env s + -> Idx env t + -> dist + -> Bool +shortestPathLessThanEqual edgeDist graph from to dist = isJust res + where + (res, _, _) = shortestPathGeneric From earlyExit edgeDist graph from + earlyExit dists = case prjPartial to dists of + Just (Const d) + | d <= dist -> Just () + _ -> Nothing + +shortestPathsFrom + :: (Show dist, Num dist, Ord dist) + => (forall a b. edge a b -> dist) + -> Graph node edge env + -> Idx env s + -> PartialEnv (Const dist) env +shortestPathsFrom edgeDist graph from = dists + where + (_, dists, _) = shortestPathGeneric From (const Nothing) edgeDist graph from + +shortestPathsTo + :: (Show dist, Num dist, Ord dist) + => (forall a b. edge a b -> dist) + -> Graph node edge env + -> Idx env s + -> PartialEnv (Const dist) env +shortestPathsTo edgeDist graph to = dists + where + (_, dists, _) = shortestPathGeneric To (const Nothing) edgeDist graph to + +type HasNegativeCycle = Bool + +data StartPoint = From | To + +-- | Generic implementation of the shortest path algorithm, used for +-- shortestPath, shortestPathsFrom and shortestPathGlobalFrom. +-- We use Bellman-Ford, to support negative cycles. +shortestPathGeneric + :: forall dist node edge env res s. + (Show dist, Num dist, Ord dist) + => StartPoint + -> (PartialEnv (Const dist) env -> Maybe res) -- Early exit check + -> (forall a b. edge a b -> dist) + -> Graph node edge env + -> Idx env s + -> (Maybe res, PartialEnv (Const dist) env, HasNegativeCycle) +shortestPathGeneric startPoint earlyExit edgeDist graph startIdx = + loop (nodeCount graph) (partialEnvSingleton startIdx (Const 0)) + where + loop :: Int -> PartialEnv (Const dist) env -> (Maybe res, PartialEnv (Const dist) env, HasNegativeCycle) + -- If the tentative distances didn't stabilize after 'nodeCount graph' iterations, the graph has a negative cycle. + loop 0 dists = (Nothing, dists, True) + loop iters dists + | Just res <- earlyExit dists' = (Just res, dists', False) + | distancesEq dists dists' = (Nothing, dists', False) + | otherwise = loop (iters - 1) dists' + where + dists' = case startPoint of + From -> shortestPathFromStep edgeDist graph dists + To -> shortestPathToStep edgeDist graph dists + +-- | One iteration of Bellman-Ford for shortestPathGeneric when start = From +shortestPathFromStep + :: forall dist node edge env. + (Show dist, Num dist, Ord dist) + => (forall a b. edge a b -> dist) + -> Graph node edge env + -> PartialEnv (Const dist) env + -> PartialEnv (Const dist) env +shortestPathFromStep _ GEmpty _ = PEnd +shortestPathFromStep dist (GPush g (Node _ ins outs)) dists1 = dists4 + where + -- Update our own distance + paths :: [dist] + paths = map (\(EnvBinding _ (Const d)) -> d) + $ partialEnvToList + -- Combine all nodes with their edge to the current node + $ intersectPartialEnv (\(Const d) (InEdge edge) -> Const $ d + dist edge) (partialEnvTail dists1) ins + + own = minimum' (partialEnvLast dists1) paths + + -- Update distances based on outgoing edges out of this node + dists2 = case own of + Nothing -> dists1 + Just (Const d) -> + unionPartialEnv (\(Const a) (Const b) -> Const $ min a b) dists1 + $ mapPartialEnv + (\edge -> Const $ d + dist edge) + outs + + -- Recursion + dists3 = shortestPathFromStep dist g (partialEnvTail dists2) + + dists4 = partialEnvPush dists3 own + +-- | One iteration of Bellman-Ford for shortestPathGeneric when start = To +shortestPathToStep + :: forall dist node edge env. + (Show dist, Num dist, Ord dist) + => (forall a b. edge a b -> dist) + -> Graph node edge env + -> PartialEnv (Const dist) env + -> PartialEnv (Const dist) env +shortestPathToStep _ GEmpty _ = PEnd +shortestPathToStep dist (GPush g (Node _ ins outs)) dists1 = dists4 + where + -- Update our own distance + paths :: [dist] + paths = map (\(EnvBinding _ (Const d)) -> d) + $ partialEnvToList + -- Combine all nodes with their edge to the current node + $ intersectPartialEnv (\(Const d) edge -> Const $ d + dist edge) dists1 outs + + own = minimum' (partialEnvLast dists1) paths + + -- Update distances based on outgoing edges out of this node + dists2 = case own of + Nothing -> partialEnvTail dists1 + Just (Const d) -> + unionPartialEnv (\(Const a) (Const b) -> Const $ min a b) (partialEnvTail dists1) + $ mapPartialEnv + (\(InEdge edge) -> Const $ d + dist edge) + ins + + -- Recursion + dists3 = shortestPathToStep dist g dists2 + + dists4 = partialEnvPush dists3 own + +distancesEq :: Eq dist => PartialEnv (Const dist) env -> PartialEnv (Const dist) env -> Bool +distancesEq PEnd PEnd = True +distancesEq (PPush as a) (PPush bs b) = a == b && distancesEq as bs +distancesEq (PNone as ) (PNone bs ) = distancesEq as bs +distancesEq (PNone as ) PEnd = distancesEq as PEnd +distancesEq PEnd (PNone bs ) = distancesEq PEnd bs +distancesEq _ _ = False + +minimum' :: Ord dist => Maybe (Const dist t) -> [dist] -> Maybe (Const dist t) +minimum' m [] = m +minimum' Nothing ds = Just $ Const $ minimum ds +minimum' (Just (Const d)) ds = Just $ Const $ min d $ minimum ds diff --git a/src/Data/Array/Accelerate/AST/Idx.hs b/src/Data/Array/Accelerate/AST/Idx.hs index cd3c8c556..4225071f0 100644 --- a/src/Data/Array/Accelerate/AST/Idx.hs +++ b/src/Data/Array/Accelerate/AST/Idx.hs @@ -1,14 +1,7 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE CPP #-} -{-# LANGUAGE EmptyCase #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE KindSignatures #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE TemplateHaskellQuotes #-} +{-# LANGUAGE ViewPatterns #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.AST.Idx @@ -23,20 +16,23 @@ -- module Data.Array.Accelerate.AST.Idx ( - - Idx, pattern ZeroIdx, pattern SuccIdx, pattern VoidIdx, + Idx(ZeroIdx, SuccIdx, VoidIdx), idxToInt, - rnfIdx, liftIdx, + rnfIdx, liftIdx, matchIdx, + + PairIdx(..), - PairIdx(..) + Skip(SkipNone, SkipSucc, SkipSucc'), skipIdx, unskipIdx, chainSkip, + Keep(KeepNone, KeepSucc), keepIdx, unkeepIdx, chainKeep, ) where -import Data.Kind +import Control.DeepSeq import Language.Haskell.TH.Extra hiding ( Type ) +import Data.Type.Equality ( (:~:)(Refl) ) #ifndef ACCELERATE_INTERNAL_CHECKS -import Data.Type.Equality ( (:~:)(Refl) ) +import Data.Kind import Unsafe.Coerce ( unsafeCoerce ) #endif @@ -50,6 +46,14 @@ data Idx env t where ZeroIdx :: Idx (env, t) t SuccIdx :: Idx env t -> Idx (env, s) t +instance Eq (Idx env t) where + ZeroIdx == ZeroIdx = True + SuccIdx a == SuccIdx b = a == b + _ == _ = False + +instance Ord (Idx env t) where + compare a b = compare (idxToInt a) (idxToInt b) + idxToInt :: Idx env t -> Int idxToInt ZeroIdx = 0 idxToInt (SuccIdx idx) = 1 + idxToInt idx @@ -62,6 +66,12 @@ liftIdx :: Idx env t -> CodeQ (Idx env t) liftIdx ZeroIdx = [|| ZeroIdx ||] liftIdx (SuccIdx ix) = [|| SuccIdx $$(liftIdx ix) ||] +{-# INLINEABLE matchIdx #-} +matchIdx :: Idx env s -> Idx env t -> Maybe (s :~: t) +matchIdx ZeroIdx ZeroIdx = Just Refl +matchIdx (SuccIdx u) (SuccIdx v) = matchIdx u v +matchIdx _ _ = Nothing + #else -- | De Bruijn variable index projecting a specific type from a type @@ -76,7 +86,7 @@ liftIdx (SuccIdx ix) = [|| SuccIdx $$(liftIdx ix) ||] -- newtype Idx :: Type -> Type -> Type where UnsafeIdxConstructor :: { unsafeRunIdx :: Int } -> Idx env t - + deriving (Eq, Ord) {-# COMPLETE ZeroIdx, SuccIdx #-} pattern ZeroIdx :: forall envt t. () => forall env. (envt ~ (env, t)) => Idx envt t @@ -89,11 +99,17 @@ pattern SuccIdx idx <- (unSucc -> Just (idx, Refl)) where SuccIdx (UnsafeIdxConstructor i) = UnsafeIdxConstructor (i+1) +-- Note: env and s should actually not be universally quantified (forall), +-- so this function on its own is unsound. The integration in SuccIdx makes it +-- sound however. unSucc :: Idx envs t -> Maybe (Idx env t, envs :~: (env, s)) unSucc (UnsafeIdxConstructor i) | i < 1 = Nothing | otherwise = Just (UnsafeIdxConstructor (i-1), unsafeCoerce Refl) +-- Note: this is a separate function instead of a field of Idx, since this function is exported. +-- If we would export a field of Idx, then it could unsafely be updated via +-- idx{ idxToInt = 2 }, breaking the "safe" abstraction of this module. idxToInt :: Idx env t -> Int idxToInt = unsafeRunIdx @@ -103,11 +119,20 @@ rnfIdx !_ = () liftIdx :: Idx env t -> CodeQ (Idx env t) liftIdx (UnsafeIdxConstructor i) = [|| UnsafeIdxConstructor i ||] +{-# INLINEABLE matchIdx #-} +matchIdx :: Idx env s -> Idx env t -> Maybe (s :~: t) +matchIdx (UnsafeIdxConstructor i) (UnsafeIdxConstructor j) + | i == j = Just $ unsafeCoerce Refl + | otherwise = Nothing #endif +instance NFData (Idx env t) where + rnf = rnfIdx + -- | Despite the 'complete' pragma above, GHC can't infer that there is no -- pattern possible if the environment is empty. This can be used instead. -- +{-# COMPLETE VoidIdx #-} pattern VoidIdx :: forall env t a. (env ~ ()) => () => a -> Idx env t pattern VoidIdx a <- (\case{} -> a) @@ -117,3 +142,173 @@ data PairIdx p a where PairIdxLeft :: PairIdx (a, b) a PairIdxRight :: PairIdx (a, b) b +-- Skip: a GADT to describe that we skip over a number of bindings in an +-- environment. +#ifdef ACCELERATE_INTERNAL_CHECKS + +-- Drops some bindings of env' to result in env. +data Skip env env' where + SkipSucc :: Skip env (env', t) -> Skip env env' + SkipNone :: Skip env env + +-- Historical context: we used to have two data types, Skip and Skip'. +-- Skip was constructed via: +-- SkipSucc :: Skip env (env', t) -> Skip env env' +-- and Skip' via: +-- SkipSucc' :: Skip env env' -> Skip (env, t) env' +-- Both data types have the same functionality, but different performance, +-- similar to the difference between cons- and snoc-lists. +-- Each analysis or transformation in the compiler would need to choose between +-- these two data types. +-- To simplify this, we now only have Skip, and SkipSucc' is implemented as a +-- pattern synonym over Skip. + +skipIdx :: Skip env env' -> Idx env t -> Maybe (Idx env' t) +skipIdx SkipNone idx = Just idx +skipIdx (SkipSucc s) idx = case skipIdx s idx of + Just (SuccIdx idx') -> Just idx' + _ -> Nothing + +unskipIdx :: Skip env env' -> Idx env' t -> Idx env t +unskipIdx SkipNone idx = idx +unskipIdx (SkipSucc s) idx = unskipIdx s $ SuccIdx idx + +chainSkip :: Skip env1 env2 -> Skip env2 env3 -> Skip env1 env3 +chainSkip skipL (SkipSucc skipR) = SkipSucc $ chainSkip skipL skipR +chainSkip skipL SkipNone = skipL + +skipSucc' :: Skip env env' -> Skip (env, t) env' +skipSucc' SkipNone = SkipSucc SkipNone +skipSucc' (SkipSucc s) = SkipSucc $ skipSucc' s + +data UnSkipSucc' envt env' where + UnSkipSucc' :: Skip env env' -> UnSkipSucc' (env, t) env' + +unSkipSucc' :: Skip env env' -> Either (env :~: env') (UnSkipSucc' env env') +unSkipSucc' SkipNone = Left Refl +unSkipSucc' (SkipSucc s) = Right $ case unSkipSucc' s of + Left Refl -> UnSkipSucc' SkipNone + Right (UnSkipSucc' s') -> UnSkipSucc' $ SkipSucc s' + +pattern SkipSucc' :: forall envt env'. () => forall t env. (envt ~ (env, t)) => Skip env env' -> Skip envt env' +pattern SkipSucc' s <- (unSkipSucc' -> Right (UnSkipSucc' s)) + where + SkipSucc' s = skipSucc' s +{-# COMPLETE SkipNone, SkipSucc' #-} + +#else + +newtype Skip :: Type -> Type -> Type where + UnsafeSkipConstructor :: { skipToInt :: Int } -> Skip env env' +{-# COMPLETE SkipNone, SkipSucc #-} +{-# COMPLETE SkipNone, SkipSucc' #-} + +pattern SkipNone :: forall env env'. () => (env ~ env') => Skip env env' +pattern SkipNone <- (\x -> (skipToInt x, unsafeCoerce Refl) -> (0, Refl :: env :~: env')) + where + SkipNone = UnsafeSkipConstructor 0 + +skipIdx :: Skip env env' -> Idx env t -> Maybe (Idx env' t) +skipIdx skip idx + | i >= 0 = Just $ UnsafeIdxConstructor i + | otherwise = Nothing + where + i = idxToInt idx - skipToInt skip + +unskipIdx :: Skip env env' -> Idx env' t -> Idx env t +unskipIdx skip idx = UnsafeIdxConstructor $ skipToInt skip + idxToInt idx + +chainSkip :: Skip env1 env2 -> Skip env2 env3 -> Skip env1 env3 +chainSkip skip1 skip2 = UnsafeSkipConstructor $ skipToInt skip1 + skipToInt skip2 + +data UnSkipSucc env env' where + UnSkipSucc :: Skip env (env', t) -> UnSkipSucc env env' + +unSkipSucc :: Skip env env' -> Maybe (UnSkipSucc env env') +unSkipSucc (UnsafeSkipConstructor 0) = Nothing +unSkipSucc (UnsafeSkipConstructor i) + = unsafeCoerce $ Just $ UnSkipSucc $ UnsafeSkipConstructor $ i - 1 + +pattern SkipSucc :: forall env env'. () => forall t envt'. (envt' ~ (env', t)) => Skip env envt' -> Skip env env' +pattern SkipSucc s <- (unSkipSucc -> Just (UnSkipSucc s)) + where + SkipSucc s = UnsafeSkipConstructor $ skipToInt s + 1 + +data UnSkipSucc' envt env' where + UnSkipSucc' :: Skip env env' -> UnSkipSucc' (env, t) env' + +unSkipSucc' :: Skip env env' -> Maybe (UnSkipSucc' env env') +unSkipSucc' (UnsafeSkipConstructor 0) = Nothing +unSkipSucc' (UnsafeSkipConstructor i) + = unsafeCoerce $ Just $ UnSkipSucc' $ UnsafeSkipConstructor $ i - 1 + +pattern SkipSucc' :: forall envt env'. () => forall t env. (envt ~ (env, t)) => Skip env env' -> Skip envt env' +pattern SkipSucc' s <- (unSkipSucc' -> Just (UnSkipSucc' s)) + where + SkipSucc' s = UnsafeSkipConstructor $ skipToInt s + 1 + +#endif + +-- Keep: a GADT to describe that we preserve a number of bindings in an +-- environment. +-- The right-most n bindings of env1' and env2' are equal. +-- After dropping those bindings we get env1 and env2 respectively. +#ifdef ACCELERATE_INTERNAL_CHECKS + +data Keep env1 env2 env1' env2' where + KeepNone :: Keep env1 env2 env1 env2 + KeepSucc + :: Keep env1 env2 env1' env2' + -> Keep env1 env2 (env1', t) (env2', t) + +keepIdx :: Keep env1 env2 env1' env2' -> Idx env1' t -> Either (Idx env1 t) (Idx env2' t) +keepIdx KeepNone idx = Left idx +keepIdx (KeepSucc _) ZeroIdx = Right ZeroIdx +keepIdx (KeepSucc k) (SuccIdx idx) = SuccIdx <$> keepIdx k idx + +unkeepIdx :: Keep env1 env2 env1' env2' -> Idx env2 t -> Idx env2' t +unkeepIdx KeepNone idx = idx +unkeepIdx (KeepSucc k) idx = SuccIdx $ unkeepIdx k idx + +chainKeep :: Keep env1 env2 env1' env2' -> Keep env1' env2' env1'' env2'' -> Keep env1 env2 env1'' env2'' +chainKeep keepL KeepNone = keepL +chainKeep keepL (KeepSucc keepR) = KeepSucc $ chainKeep keepL keepR + +#else + +newtype Keep :: Type -> Type -> Type -> Type -> Type where + UnsafeKeepConstructor :: { keepToInt :: Int } -> Keep env1 env2 env1' env2' +{-# COMPLETE KeepNone, KeepSucc #-} + +pattern KeepNone :: forall env1 env2 env1' env2'. () => (env1 ~ env1', env2 ~ env2') => Keep env1 env2 env1' env2' +pattern KeepNone <- (\x -> (keepToInt x, unsafeCoerce Refl) -> (0, Refl :: (env1, env2) :~: (env1', env2'))) + where + KeepNone = UnsafeKeepConstructor 0 + +data UnKeepSucc env1 env2 env1t' env2t' where + UnKeepSucc :: Keep env1 env2 env1' env2' -> UnKeepSucc env1 env2 (env1', t) (env2', t) + +unKeepSucc :: Keep env1 env2 env1t' env2t' -> Maybe (UnKeepSucc env1 env2 env1t' env2t') +unKeepSucc (UnsafeKeepConstructor 0) = Nothing +unKeepSucc (UnsafeKeepConstructor i) + = unsafeCoerce $ Just $ UnKeepSucc $ UnsafeKeepConstructor $ i - 1 + +pattern KeepSucc :: forall env1 env2 env1t' env2t'. () => forall env1' env2' t. (env1t' ~ (env1', t), env2t' ~ (env2', t)) => Keep env1 env2 env1' env2' -> Keep env1 env2 env1t' env2t' +pattern KeepSucc k <- (unKeepSucc -> Just (UnKeepSucc k)) + where + KeepSucc k = UnsafeKeepConstructor $ keepToInt k + 1 + +keepIdx :: Keep env1 env2 env1' env2' -> Idx env1' t -> Either (Idx env1 t) (Idx env2' t) +keepIdx k idx + | i >= 0 = Left $ UnsafeIdxConstructor i + | otherwise = Right $ UnsafeIdxConstructor $ idxToInt idx + where + i = idxToInt idx - keepToInt k + +unkeepIdx :: Keep env1 env2 env1' env2' -> Idx env2 t -> Idx env2' t +unkeepIdx k idx = UnsafeIdxConstructor $ idxToInt idx + keepToInt k + +chainKeep :: Keep env1 env2 env1' env2' -> Keep env1' env2' env1'' env2'' -> Keep env1 env2 env1'' env2'' +chainKeep keep1 keep2 = UnsafeKeepConstructor $ keepToInt keep1 + keepToInt keep2 + +#endif diff --git a/src/Data/Array/Accelerate/AST/IdxSet.hs b/src/Data/Array/Accelerate/AST/IdxSet.hs new file mode 100644 index 000000000..dcd22e5de --- /dev/null +++ b/src/Data/Array/Accelerate/AST/IdxSet.hs @@ -0,0 +1,177 @@ +{-# OPTIONS_HADDOCK hide #-} +-- | +-- Module : Data.Array.Accelerate.AST.IdxSet +-- Copyright : [2008..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.AST.IdxSet ( + IdxSet(..), + member, varMember, overlaps, isSubsetOf, intersect, union, unions, (\\), (>>=), + insert, insertVar, insertVars, skip, skip', + push, push', empty, isEmpty, drop, drop', remove, partialEnvRemoveSet, partialEnvFilterSet, + fromList, fromList', fromVarList, fromVars, map, + singleton, singletonVar, first, null, + toList, rnfIdxSet +) where + +import Prelude hiding (drop, (>>=), map, null) + +import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.Var +import Data.Array.Accelerate.AST.Environment hiding ( push, push' ) +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.Representation.Type +import Data.Maybe + +newtype IdxSet env = IdxSet { unIdxSet :: PartialEnv Present env } + +data Present a = Present + +member :: Idx env t -> IdxSet env -> Bool +member idx (IdxSet set) = isJust $ prjPartial idx set + +varMember :: Var s env t -> IdxSet env -> Bool +varMember (Var _ idx) = member idx + +instance Eq (IdxSet env) where + IdxSet PEnd == b = isEmpty b + a == IdxSet PEnd = isEmpty a + IdxSet (PPush a _) == IdxSet (PPush b _) = IdxSet a == IdxSet b + IdxSet (PNone a) == IdxSet (PNone b) = IdxSet a == IdxSet b + _ == _ = False + +instance Semigroup (IdxSet env) where + (<>) = union + +instance Monoid (IdxSet env) where + mempty = empty + +overlaps :: IdxSet env -> IdxSet env -> Bool +overlaps (IdxSet (PPush _ _)) (IdxSet (PPush _ _)) = True +overlaps (IdxSet (PPush a _)) (IdxSet (PNone b )) = overlaps (IdxSet a) (IdxSet b) +overlaps (IdxSet (PNone a )) (IdxSet (PPush b _)) = overlaps (IdxSet a) (IdxSet b) +overlaps (IdxSet (PNone a )) (IdxSet (PNone b )) = overlaps (IdxSet a) (IdxSet b) +overlaps _ _ = False + +isSubsetOf :: IdxSet env -> IdxSet env -> Bool +isSubsetOf (IdxSet (PPush a _)) (IdxSet (PPush b _)) = isSubsetOf (IdxSet a) (IdxSet b) +isSubsetOf (IdxSet (PNone a )) (IdxSet (PPush b _)) = isSubsetOf (IdxSet a) (IdxSet b) +isSubsetOf (IdxSet (PPush _ _)) (IdxSet (PNone _ )) = False +isSubsetOf (IdxSet (PNone a )) (IdxSet (PNone b )) = isSubsetOf (IdxSet a) (IdxSet b) +isSubsetOf (IdxSet PEnd) _ = True +isSubsetOf a (IdxSet PEnd) = null a + +intersect :: IdxSet env -> IdxSet env -> IdxSet env +intersect (IdxSet a) (IdxSet b) = IdxSet $ intersectPartialEnv (\_ _ -> Present) a b + +union :: IdxSet env -> IdxSet env -> IdxSet env +union (IdxSet a) (IdxSet b) = IdxSet $ unionPartialEnv (\_ _ -> Present) a b + +unions :: [IdxSet env] -> IdxSet env +unions [] = empty +unions [a] = a +unions (a:as) = union a $ unions as + +(\\) :: IdxSet env -> IdxSet env -> IdxSet env +IdxSet a \\ IdxSet b = IdxSet $ diffPartialEnv a b + +(>>=) :: IdxSet env -> (forall t. Idx env t -> IdxSet env') -> IdxSet env' +set >>= f = unions $ fmap (\(Exists idx) -> f idx) $ toList set + +insert :: Idx env t -> IdxSet env -> IdxSet env +insert idx (IdxSet a) = IdxSet $ partialUpdate Present idx a + +insertVar :: Var s env t -> IdxSet env -> IdxSet env +insertVar (Var _ idx) = insert idx + +insertVars :: Vars s env t -> IdxSet env -> IdxSet env +insertVars (TupRsingle (Var _ idx)) = insert idx +insertVars (TupRpair v1 v2) = insertVars v1 . insertVars v2 +insertVars TupRunit = id + +remove :: Idx env t -> IdxSet env -> IdxSet env +remove idx (IdxSet a) = IdxSet $ partialRemove idx a + +partialEnvRemoveSet :: IdxSet env -> PartialEnv f env -> PartialEnv f env +partialEnvRemoveSet (IdxSet set) env = diffPartialEnv env set + +partialEnvFilterSet :: IdxSet env -> PartialEnv f env -> PartialEnv f env +partialEnvFilterSet (IdxSet set) env = intersectPartialEnv (\_ f -> f) set env + +skip :: IdxSet env -> IdxSet (env, t) +skip = IdxSet . PNone . unIdxSet + +skip' :: LeftHandSide s t env env' -> IdxSet env -> IdxSet env' +skip' (LeftHandSideSingle _) = skip +skip' (LeftHandSideWildcard _) = id +skip' (LeftHandSidePair l1 l2) = skip' l2 . skip' l1 + +push :: IdxSet env -> IdxSet (env, t) +push = IdxSet . flip PPush Present . unIdxSet + +push' :: LeftHandSide s t env env' -> IdxSet env -> IdxSet env' +push' (LeftHandSideSingle _) = push +push' (LeftHandSideWildcard _) = id +push' (LeftHandSidePair l1 l2) = push' l2 . push' l1 + +empty :: IdxSet env +empty = IdxSet PEnd + +isEmpty :: IdxSet env -> Bool +isEmpty (IdxSet PEnd) = True +isEmpty (IdxSet (PNone set)) = isEmpty $ IdxSet set +isEmpty _ = False + +drop :: IdxSet (env, t) -> IdxSet env +drop = IdxSet . partialEnvTail . unIdxSet + +drop' :: LeftHandSide s t env env' -> IdxSet env' -> IdxSet env +drop' (LeftHandSidePair lhs1 lhs2) = drop' lhs1 . drop' lhs2 +drop' (LeftHandSideWildcard _) = id +drop' (LeftHandSideSingle _) = drop + +toList :: IdxSet env -> [Exists (Idx env)] +toList = fmap (\(EnvBinding idx _) -> Exists idx) . partialEnvToList . unIdxSet + +fromList :: [Exists (Idx env)] -> IdxSet env +fromList = IdxSet . partialEnvFromList (\_ _ -> Present) . fmap (\(Exists idx) -> EnvBinding idx Present) + +fromList' :: [Idx env t] -> IdxSet env +fromList' = IdxSet . partialEnvFromList (\_ _ -> Present) . fmap (\idx -> EnvBinding idx Present) + +map :: (forall t. Idx env t -> Idx env' t) -> IdxSet env -> IdxSet env' +map f = fromList . fmap (\(Exists idx) -> Exists $ f idx) . toList + +fromVarList :: [Exists (Var s env)] -> IdxSet env +fromVarList = fromList . fmap (\(Exists (Var _ idx)) -> Exists idx) + +fromVars :: Vars s env t -> IdxSet env +fromVars = fromVarList . flattenTupR + +singleton :: Idx env t -> IdxSet env +singleton idx = IdxSet $ partialEnvSingleton idx Present + +singletonVar :: Var s env t -> IdxSet env +singletonVar (Var _ idx) = singleton idx + +first :: IdxSet env -> Maybe (Exists (Idx env)) +first (IdxSet (PPush _ _)) = Just $ Exists ZeroIdx +first (IdxSet (PNone env)) + | Just (Exists idx) <- first (IdxSet env) = Just $ Exists $ SuccIdx idx +first _ = Nothing + +null :: IdxSet env -> Bool +null (IdxSet (PPush _ _)) = False +null (IdxSet (PNone env)) = null (IdxSet env) +null _ = True + +rnfIdxSet :: IdxSet env -> () +rnfIdxSet = foldr (\(Exists ix) () -> rnfIdx ix) () . toList + +instance Show (IdxSet env) where + showsPrec p = showsPrec p . fmap (\(Exists idx) -> idxToInt idx) . toList diff --git a/src/Data/Array/Accelerate/AST/Kernel.hs b/src/Data/Array/Accelerate/AST/Kernel.hs new file mode 100644 index 000000000..8dc6c7493 --- /dev/null +++ b/src/Data/Array/Accelerate/AST/Kernel.hs @@ -0,0 +1,80 @@ +{-# LANGUAGE DefaultSignatures #-} +{-# LANGUAGE TypeFamilies #-} + +-- | +-- Module : Data.Array.Accelerate.AST.Kernel +-- Copyright : [2012..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- +-- This module defines the interface for kernels, to be implemented +-- in other modules. +-- + +module Data.Array.Accelerate.AST.Kernel ( + IsKernel(..), KernelArgR(..), + OpenKernelFun(..), KernelFun, + NoKernelMetadata, kernelFunKernel +) where + +import Data.Array.Accelerate.AST.Environment +import Data.Array.Accelerate.AST.Partitioned +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.Representation.Shape +import Data.Array.Accelerate.Type +import Data.Array.Accelerate.Array.Buffer +import Data.Array.Accelerate.Analysis.Hash.Exp +import Data.Kind + +class (NFData' kernel, NFData' (KernelMetadata kernel)) => IsKernel kernel where + type KernelOperation kernel :: Type -> Type + type KernelMetadata kernel :: Type -> Type + + compileKernel :: Env AccessGroundR env -> Clustered (KernelOperation kernel) args -> Args env args -> kernel env + + kernelMetadata :: KernelFun kernel f -> KernelMetadata kernel f + default kernelMetadata :: KernelMetadata kernel ~ NoKernelMetadata => KernelFun kernel f -> KernelMetadata kernel f + kernelMetadata _ = NoKernelMetadata + + -- Depending on the backend, it may either be easier + -- to return the hash/UID of an already hashed kernel, + -- or manually encode the kernel. + encodeKernel :: kernel env -> Either Hash Builder + +type KernelFun kernel = OpenKernelFun kernel () + +data OpenKernelFun kernel env t where + KernelFunLam + :: KernelArgR t r + -> OpenKernelFun kernel (env, r) f + -> OpenKernelFun kernel env (t -> f) + + KernelFunBody + :: kernel env + -> OpenKernelFun kernel env () + +data KernelArgR t r where + KernelArgRscalar :: ScalarType s -> KernelArgR (Var' s) s + KernelArgRbuffer :: Modifier m -> ScalarType s -> KernelArgR (m DIM1 s) (Buffer s) + +instance NFData' kernel => NFData' (OpenKernelFun kernel env) where + rnf' (KernelFunLam g fun) = rnfKernelArgR g `seq` rnf' fun + rnf' (KernelFunBody kernel) = rnf' kernel + +rnfKernelArgR :: KernelArgR t r -> () +rnfKernelArgR (KernelArgRscalar tp) = rnfScalarType tp +rnfKernelArgR (KernelArgRbuffer m tp) = m `seq` rnfScalarType tp + +data NoKernelMetadata f = NoKernelMetadata + +instance NFData' NoKernelMetadata where + rnf' NoKernelMetadata = () + +kernelFunKernel :: OpenKernelFun kernel env t -> Exists kernel +kernelFunKernel (KernelFunBody kernel) = Exists kernel +kernelFunKernel (KernelFunLam _ f) + | Exists k <- kernelFunKernel f + = Exists k diff --git a/src/Data/Array/Accelerate/AST/LeftHandSide.hs b/src/Data/Array/Accelerate/AST/LeftHandSide.hs index 91ec12646..dedb4f78a 100644 --- a/src/Data/Array/Accelerate/AST/LeftHandSide.hs +++ b/src/Data/Array/Accelerate/AST/LeftHandSide.hs @@ -1,11 +1,7 @@ -{-# LANGUAGE GADTs #-} -{-# LANGUAGE KindSignatures #-} {-# LANGUAGE PatternSynonyms #-} {-# LANGUAGE QuantifiedConstraints #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE StandaloneDeriving #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE TemplateHaskellQuotes #-} +{-# LANGUAGE UndecidableInstances #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.AST.LeftHandSide @@ -18,15 +14,25 @@ -- module Data.Array.Accelerate.AST.LeftHandSide + ( + Exists(..), + LeftHandSide(.., LeftHandSideUnit), leftHandSidePair, + lhsToTupR, lhsSize, leftHandSideIsVoid, + rnfLeftHandSide, liftLeftHandSide, mapLeftHandSide, flattenTupR) where import Data.Array.Accelerate.Representation.Type import Language.Haskell.TH.Extra +import Data.Typeable data Exists f where - Exists :: f a -> Exists f + Exists :: !(f a) -> Exists f + +instance (forall a . Show (f a)) => Show (Exists f) where + show :: (forall a. Show (f a)) => Exists f -> String + show (Exists e) = show e data LeftHandSide s v env env' where LeftHandSideSingle @@ -55,6 +61,11 @@ lhsToTupR (LeftHandSideSingle s) = TupRsingle s lhsToTupR (LeftHandSideWildcard r) = r lhsToTupR (LeftHandSidePair as bs) = TupRpair (lhsToTupR as) (lhsToTupR bs) +lhsSize :: LeftHandSide s v env env' -> Int +lhsSize (LeftHandSideWildcard _) = 0 +lhsSize (LeftHandSidePair l1 l2) = lhsSize l1 + lhsSize l2 +lhsSize (LeftHandSideSingle _) = 1 + rnfLeftHandSide :: (forall b. s b -> ()) -> LeftHandSide s v env env' -> () rnfLeftHandSide f (LeftHandSideWildcard r) = rnfTupR f r rnfLeftHandSide f (LeftHandSideSingle s) = f s @@ -65,3 +76,26 @@ liftLeftHandSide f (LeftHandSideSingle s) = [|| LeftHandSideSingle $$(f s) ||] liftLeftHandSide f (LeftHandSideWildcard r) = [|| LeftHandSideWildcard $$(liftTupR f r) ||] liftLeftHandSide f (LeftHandSidePair as bs) = [|| LeftHandSidePair $$(liftLeftHandSide f as) $$(liftLeftHandSide f bs) ||] +mapLeftHandSide :: (forall v. s v -> u v) -> LeftHandSide s t env env' -> LeftHandSide u t env env' +mapLeftHandSide f (LeftHandSideSingle s) = LeftHandSideSingle $ f s +mapLeftHandSide f (LeftHandSideWildcard r) = LeftHandSideWildcard $ mapTupR f r +mapLeftHandSide f (LeftHandSidePair as bs) = LeftHandSidePair (mapLeftHandSide f as) (mapLeftHandSide f bs) + +leftHandSidePair :: LeftHandSide s t1 env env' -> LeftHandSide s t2 env' env'' -> LeftHandSide s (t1, t2) env env'' +leftHandSidePair (LeftHandSideWildcard t1) (LeftHandSideWildcard t2) = LeftHandSideWildcard $ TupRpair t1 t2 +leftHandSidePair l1 l2 = LeftHandSidePair l1 l2 + +flattenTupR :: TupR s t -> [Exists s] +flattenTupR = (`go` []) + where + go :: TupR s t -> [Exists s] -> [Exists s] + go (TupRsingle s) accum = Exists s : accum + go (TupRpair t1 t2) accum = go t1 $ go t2 accum + go TupRunit accum = accum + +leftHandSideIsVoid :: LeftHandSide s t env env' -> Maybe (env :~: env') +leftHandSideIsVoid (LeftHandSideWildcard _) = Just Refl +leftHandSideIsVoid (LeftHandSidePair l1 l2) + | Just Refl <- leftHandSideIsVoid l1 + , Just Refl <- leftHandSideIsVoid l2 = Just Refl +leftHandSideIsVoid _ = Nothing diff --git a/src/Data/Array/Accelerate/AST/Operation.hs b/src/Data/Array/Accelerate/AST/Operation.hs new file mode 100644 index 000000000..85723f9bc --- /dev/null +++ b/src/Data/Array/Accelerate/AST/Operation.hs @@ -0,0 +1,762 @@ +{-# LANGUAGE QuantifiedConstraints #-} +{-# LANGUAGE TemplateHaskell #-} +{-# OPTIONS_GHC -Wno-orphans #-} +{-# OPTIONS_HADDOCK hide #-} +-- | +-- Module : Data.Array.Accelerate.AST.Operation +-- Copyright : [2008..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.AST.Operation ( + PreOpenAcc(..), PreOpenAfun(..), + OperationAcc, OperationAfun, + Uniqueness(..), Uniquenesses, shared, unique, pairUniqueness, + + GroundR(..), GroundsR, GroundVar, GroundVars, GLeftHandSide, Var(..), Vars, + HasGroundsR(..), groundToExpVar, + + PreArgs(..), Arg(..), Args, argsToList, Modifier(..), argArrayR, argVarType, + rnfArg, rnfArgs, rnfPreArgs, mapArgs, + + Var', Exp', Fun', In, Out, Mut, + + ArrayDescriptor(..), ArrayDescriptors, weakenArrayDescriptor, weakenArrayDescriptors, + + OpenExp, OpenFun, Exp, Fun, ArrayInstr(..), + expGroundVars, funGroundVars, arrayInstrsInExp, arrayInstrsInFun, + + encodeGroundR, encodeGroundsR, encodeGroundVar, encodeGroundVars, encodeArrayDescriptors, + rnfGroundR, rnfGroundsR, rnfGroundVar, rnfGroundVars, rnfUniqueness, + liftGroundR, liftGroundsR, liftGroundVar, liftGroundVars, + + bufferImpossible, groundFunctionImpossible, + + paramIn, paramsIn, paramIn', paramsIn', + + ReindexPartial, reindexArg, reindexArgs, reindexExp, reindexPreArgs, reindexVar, reindexVars, reindexArrayDescriptors, + reindexIdxSet, + weakenReindex, + arrayDescriptorsIdxSet, + argVars, argsVars, argsInputs, argsOutputs, AccessGroundR(..), + + mapAccExecutable, mapAfunExecutable, + + GroundRWithUniqueness(..), groundsRWithUniquenesses, lhsWithUniquesses, + + module Data.Array.Accelerate.AST.Exp, + + NFData'(..) +,reindexAcc,toGrounds,fromGrounds,weakenThroughReindex,fuseArgsWith,argsFromList,expToGroundVar) where + +import Data.Array.Accelerate.AST.Environment +import Data.Array.Accelerate.AST.Exp +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.Var +import Data.Array.Accelerate.Analysis.Hash.Exp +import Data.Array.Accelerate.Representation.Array +import Data.Array.Accelerate.Representation.Ground +import Data.Array.Accelerate.Representation.Shape +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Trafo.Exp.Substitution +import Data.Array.Accelerate.Trafo.Exp.Shrink +import Data.Array.Accelerate.Type +import Data.Array.Accelerate.Error +import Data.Typeable ( (:~:)(..) ) +import Data.Maybe +import Data.Text ( Text ) + +import Language.Haskell.TH.Extra ( CodeQ ) +import Data.Kind (Type) +import Control.DeepSeq (NFData (rnf)) +import Data.Functor.Identity +import Data.Array.Accelerate.AST.IdxSet (IdxSet) +import qualified Data.Array.Accelerate.AST.IdxSet as IdxSet + +-- | An intermediate representation parameterized over executable operations. +-- This data type only consists of the control flow structure and +-- bookkeeping (allocating or copying buffers). +-- +-- The data type is parameterized over the type of executable operations. In +-- the compiler pipeline, it is first instantiated with 'Execute op', which +-- means that the Exec constructor will target a single operation. Type +-- argument op is platform dependent, such that we can have different +-- operations for each backend and thus have backend specific optimizations +-- for some combinators. For instance, some backends may have a three kernel +-- scan whereas others have a single pass scan, or we may have different +-- operations based on which atomic instructions are available on the backend. +-- +-- The data type is also parameterized over the environment of the terms, +-- consisting of ground values typed by 'GroundR', and the return type, which +-- is a tuple of ground values, typed by 'GroundsR'. +-- +data PreOpenAcc (op :: Type -> Type) env a where + -- | Executes an executable operation. Such execution does not return a + -- value, the effects of the execution are only visible by the mutation of + -- buffers which were annotated with either 'Mut' or 'Out'. + -- Provides the operation arguments from the environment. + -- + Exec :: op args + -> Args env args + -> PreOpenAcc op env () + + -- | Returns the values of the given variables. + -- + Return :: GroundVars env a + -> PreOpenAcc op env a + + -- | Fusion boundary. + -- Corresponds with 'compute' in the user-facing language. Internally we use + -- a different name as 'compute' already has a different meaning (computing + -- expressions). + -- + -- This constructor may only occur before fusion. + -- + Manifest :: GroundVar env (Buffer e) + -> PreOpenAcc op env (Buffer e) + + -- | Evaluates the expression and returns its value. + -- + Compute :: Exp env t + -> PreOpenAcc op env t + + -- | Local binding of ground values. + -- As it is common in this intermediate representation to evaluate some program + -- resulting in unit, for instance when execution some operation, and then + -- and then do some other code, we write 'a; b' to denote 'let () = a in b'. + Alet :: GLeftHandSide bnd env env' + -> Uniquenesses bnd + -> PreOpenAcc op env bnd + -> PreOpenAcc op env' t + -> PreOpenAcc op env t + + -- | Allocates a new buffer of the given size. + -- + Alloc :: ShapeR sh + -> ScalarType e + -> ExpVars env sh + -> PreOpenAcc op env (Buffer e) + + -- | Buffer inlet. To pass an array constant in this data type, one may need + -- multiple 'Use' constructs because of the explicit Structure-of-Arrays. + -- Triggers (possibly) asynchronous host->device transfer if necessary. + -- + Use :: ScalarType e + -> Int -- Number of elements + -> Buffer e + -> PreOpenAcc op env (Buffer e) + + -- | Capture a scalar in a singleton buffer + -- + Unit :: ExpVar env e + -> PreOpenAcc op env (Buffer e) + + -- | If-then-else for array-level computations + -- + Acond :: ExpVar env PrimBool + -> PreOpenAcc op env a + -> PreOpenAcc op env a + -> PreOpenAcc op env a + + -- | Value-recursion for array-level computations. + -- The uniqueness guarantees are an invariant of the loop + -- and should hold before and after each iteration. + -- + Awhile :: Uniquenesses a + -> PreOpenAfun op env (a -> PrimBool) + -> PreOpenAfun op env (a -> a) + -> GroundVars env a + -> PreOpenAcc op env a + + Atrace :: Text + -> ArrayDescriptors env t + -> PreOpenAcc op env Word8 + + -- | Asserts that the given expression evaluates to true. + -- + -- The return value of this operation is intentionally not specified. + -- The value must be stored in a variable (in Alet) and that variable must be + -- used in a Fence. This ensures that the assertion runs before the + -- statements under the Fence are executed. + -- + -- Forgetting to use the return value in a Fence may cause that the assertion + -- is removed from the program (and thus not evaluated). + -- + Aassert :: Text + -> Exp env PrimBool + -> PreOpenAcc op env Word8 + + -- | Assumes that the given expression evaluates to true, + -- without checking. + -- + -- See Aassert for the description of the return value and the interaction + -- with Fence. + -- + Aassume :: Exp env PrimBool + -> PreOpenAcc op env Word8 + + -- | Compiler fence, to prevent the compiler from reordering the computation + -- of the given variables past this point. + -- + -- This is indented to be used with the return value of Aassert and Aassume, + -- to ensure an assertion runs before certain instructions. By having two + -- separate instructions (Aassert and Fence), we can more granually specify + -- which parts of the program depend on an assertion. Without this + -- granularity, assertions can block fusion in certain situations + -- + Fence :: IdxSet env + -> PreOpenAcc op env t + -> PreOpenAcc op env t + +-- | Function abstraction over parametrised array computations +-- +data PreOpenAfun op env t where + Abody :: PreOpenAcc op env t -> PreOpenAfun op env t + Alam :: GLeftHandSide a env env' -> PreOpenAfun op env' t -> PreOpenAfun op env (a -> t) + + +-- | Types for local bindings +-- +type GLeftHandSide = LeftHandSide GroundR +type GroundVar = Var GroundR +type GroundVars env = Vars GroundR env + +toGround :: ExpVar env a -> GroundVar env a +toGround (Var st idx) = Var (GroundRscalar st) idx + +toGrounds :: ExpVars env a -> GroundVars env a +toGrounds TupRunit = TupRunit +toGrounds (TupRsingle var) = TupRsingle (toGround var) +toGrounds (TupRpair x y) = TupRpair (toGrounds x) (toGrounds y) + +fromGround :: GroundVar env a -> ExpVar env a +fromGround (Var (GroundRscalar st) idx) = Var st idx +fromGround (Var (GroundRbuffer _) _) = error "Not a scalar" + +fromGrounds :: GroundVars env a -> ExpVars env a +fromGrounds TupRunit = TupRunit +fromGrounds (TupRsingle var) = TupRsingle (fromGround var) +fromGrounds (TupRpair x y) = TupRpair (fromGrounds x) (fromGrounds y) + +-- | Uniqueness annotations for buffers +-- +data Uniqueness t where + Unique :: Uniqueness (Buffer t) + Shared :: Uniqueness t + +type Uniquenesses = TupR Uniqueness + +deriving instance Show (Uniqueness t) + +instance Semigroup (Uniqueness t) where + (<>) :: Uniqueness t -> Uniqueness t -> Uniqueness t + (<>) Unique Unique = Unique + (<>) _ _ = Shared + +shared :: TupR s t -> Uniquenesses t +shared = mapTupR (const Shared) + +unique :: GroundsR t -> Uniquenesses t +unique = mapTupR f + where + f :: GroundR s -> Uniqueness s + f (GroundRbuffer _) = Unique + f _ = Shared + +pairUniqueness :: Uniquenesses (s, t) -> (Uniquenesses s, Uniquenesses t) +pairUniqueness (TupRpair u1 u2) = (u1, u2) +pairUniqueness (TupRsingle Shared) = (TupRsingle Shared, TupRsingle Shared) + +rnfUniqueness :: Uniqueness t -> () +rnfUniqueness Unique = () +rnfUniqueness Shared = () + +data ArrayDescriptor env a where + ArrayDescriptor :: ShapeR sh + -> GroundVars env sh + -> GroundVars env (Buffers e) + -> ArrayDescriptor env (Array sh e) + +type ArrayDescriptors env = TupR (ArrayDescriptor env) + +weakenArrayDescriptors :: env :> env' -> ArrayDescriptors env a -> ArrayDescriptors env' a +weakenArrayDescriptors k = mapTupR (weakenArrayDescriptor k) + +weakenArrayDescriptor :: env :> env' -> ArrayDescriptor env a -> ArrayDescriptor env' a +weakenArrayDescriptor k (ArrayDescriptor shr sh buffers) = ArrayDescriptor shr (weakenVars k sh) (weakenVars k buffers) + +arrayDescriptorIdxSet :: ArrayDescriptor env t -> IdxSet env +arrayDescriptorIdxSet (ArrayDescriptor _ sh buffers) = IdxSet.fromVars sh `IdxSet.union` IdxSet.fromVars buffers + +arrayDescriptorsIdxSet :: ArrayDescriptors env t -> IdxSet env +arrayDescriptorsIdxSet = foldMapTupR arrayDescriptorIdxSet + +encodeArrayDescriptors :: ArrayDescriptors env t -> Builder +encodeArrayDescriptors = encodeTupR encodeArrayDescriptor + +encodeArrayDescriptor :: ArrayDescriptor env t -> Builder +encodeArrayDescriptor (ArrayDescriptor shr sh buffer) = encodeShapeR shr <> encodeGroundVars sh <> encodeGroundVars buffer + +-- | The arguments to be passed to an operation of type `t`. +-- This type is represented as a cons list, separated by (->) and ending +-- in (). This function type represents the type of the operations. +-- `PreArgs a t` provides the arguments for function type `t`, where `a` +-- provides a single argument. +-- +data PreArgs a t where + ArgsNil :: PreArgs a () + (:>:) :: a s -> PreArgs a t -> PreArgs a (s -> t) +infixr 7 :>: + +deriving instance (forall s. Show (a s)) => Show (PreArgs a t) + +fuseArgsWith :: PreArgs a t -> PreArgs b t -> (forall x. a x -> b x -> c x) -> PreArgs c t +fuseArgsWith ArgsNil ArgsNil _ = ArgsNil +fuseArgsWith (a :>: as) (b :>: bs) f = f a b :>: fuseArgsWith as bs f + +argsToList :: PreArgs a t -> [Exists a] +argsToList ArgsNil = [] +argsToList (a :>: as) = Exists a : argsToList as + +argsFromList :: [Exists a] -> Exists (PreArgs a) +argsFromList [] = Exists ArgsNil +argsFromList (Exists a : as) = case argsFromList as of Exists args -> Exists (a :>: args) + +-- | A single argument to an operation. +-- +data Arg env t where + ArgVar :: ExpVars env e -> Arg env (Var' e) + + ArgExp :: Exp env e -> Arg env (Exp' e) + + ArgFun :: Fun env e -> Arg env (Fun' e) + + -- | An array is represented as scalar variables denoting the size and buffer variables. + -- We assume that all the buffer variables stored in 'ArgBuffers' have the same size, + -- as specified by the scalar variables of the size. + -- + ArgArray :: Modifier m + -> ArrayR (Array sh e) + -> GroundVars env sh + -> GroundVars env (Buffers e) + -> Arg env (m sh e) + +instance Show (Arg env a) where + show = \case + ArgVar{} -> "ArgVar" + ArgExp{} -> "ArgExp" + ArgFun{} -> "ArgFun" + ArgArray m _ _ _ -> case m of + In -> "in" + Out -> "out" + Mut -> "mut" + + +rnfPreArgs :: (forall s. a s -> ()) -> PreArgs a t -> () +rnfPreArgs f (a :>: as) = f a `seq` rnfPreArgs f as +rnfPreArgs _ ArgsNil = () + +rnfArg :: Arg env t -> () +rnfArg (ArgVar vars) = rnfTupR rnfExpVar vars +rnfArg (ArgExp e) = rnfOpenExp e +rnfArg (ArgFun f) = rnfOpenFun f +rnfArg (ArgArray m repr sh buffers) = m `seq` rnfArrayR repr `seq` rnfGroundVars sh `seq` rnfGroundVars buffers + +rnfArgs :: Args env t -> () +rnfArgs = rnfPreArgs rnfArg + +mapArgs :: (forall s. a s -> b s) -> PreArgs a t -> PreArgs b t +mapArgs f (a :>: as) = f a :>: mapArgs f as +mapArgs _ ArgsNil = ArgsNil + +-- | The arguments to be passed to an operation, in some environment. +-- +type Args env = PreArgs (Arg env) + +-- | Array arguments are annotated with an access modifier +-- +data Modifier m where + -- | The operation only reads from this array + In :: Modifier In + -- | The operation only writes to the array. The initial content of the array is thus irrelevant + Out :: Modifier Out + -- | The operation both reads and writes to the array. + Mut :: Modifier Mut +deriving instance Show (Modifier m) + +-- Empty data types, which are only used for the types of 'Arg'. +type Var' :: Type -> Type +data Var' e where + +type Exp' :: Type -> Type +data Exp' e where + +type Fun' :: Type -> Type +data Fun' t where + +type In :: Type -> Type -> Type +data In sh e where + +type Out :: Type -> Type -> Type +data Out sh e where + +type Mut :: Type -> Type -> Type +data Mut sh e where + +argVarType :: Arg env (Var' e) -> TypeR e +argVarType (ArgVar vars) = varsType vars + +argArrayR :: Arg env (m sh e) -> ArrayR (Array sh e) +argArrayR (ArgArray _ repr _ _) = repr + +-- | An intermediate representation consisting of operations. This data type +-- only consists of the control flow structure and bookkeeping (allocating or +-- copying buffers). It is parameterized over the actual type of operations. +-- +type OperationAcc op = PreOpenAcc op + +-- | Functions on an intermediate representation consisting of operations. +-- This data type only consists of the control flow structure and bookkeeping (allocating or +-- copying buffers). It is parameterized over the actual type of operations. +-- +type OperationAfun op = PreOpenAfun op + + +class HasGroundsR f where + groundsR :: f a -> GroundsR a + +instance HasGroundsR (PreOpenAcc op env) where + groundsR (Exec _ _) = TupRunit + groundsR (Return vars) = groundsR vars + groundsR (Manifest var) = groundsR var + groundsR (Compute e) = groundsR e + groundsR (Alet _ _ _ a) = groundsR a + groundsR (Alloc _ tp _) = TupRsingle $ GroundRbuffer tp + groundsR (Use tp _ _) = TupRsingle $ GroundRbuffer tp + groundsR (Unit (Var tp _)) = TupRsingle $ GroundRbuffer tp + groundsR (Acond _ a _) = groundsR a + groundsR (Awhile _ _ _ a) = groundsR a + groundsR (Aassert _ _) = TupRsingle $ GroundRscalar scalarTypeWord8 + groundsR (Aassume _) = TupRsingle $ GroundRscalar scalarTypeWord8 + groundsR (Atrace _ _) = TupRsingle $ GroundRscalar scalarTypeWord8 + groundsR (Fence _ a) = groundsR a + +instance HasGroundsR (GroundVar env) where + groundsR (Var repr _) = TupRsingle repr + +instance HasGroundsR (GroundVars env) where + groundsR = varsType + +instance HasGroundsR (OpenExp env benv) where + groundsR :: OpenExp env benv a -> GroundsR a + groundsR = typeRtoGroundsR . expType + +type OpenExp env benv = PreOpenExp (ArrayInstr benv) env +type OpenFun env benv = PreOpenFun (ArrayInstr benv) env +type Fun benv = OpenFun () benv +type Exp benv = OpenExp () benv + +data ArrayInstr benv t where + Index :: GroundVar benv (Buffer e) -> ArrayInstr benv (Int -> e) + Parameter :: ExpVar benv e -> ArrayInstr benv (() -> e) + +instance IsArrayInstr (ArrayInstr benv) where + arrayInstrType arr = TupRsingle $ case arr of + Index (Var (GroundRbuffer tp) _) -> tp + Index (Var (GroundRscalar tp) _) -> bufferImpossible tp + Parameter (Var tp _) -> tp + + liftArrayInstr :: ArrayInstr benv t -> CodeQ (ArrayInstr benv t) + liftArrayInstr (Index var) = [|| Index $$(liftGroundVar var) ||] + liftArrayInstr (Parameter var) = [|| Parameter $$(liftExpVar var) ||] + + rnfArrayInstr (Index var) = rnfGroundVar var + rnfArrayInstr (Parameter var) = rnfExpVar var + + showArrayInstrOp Index{} = "Index" + showArrayInstrOp Parameter{} = "Parameter" + + matchArrayInstr (Index v1) (Index v2) | Just Refl <- matchVar v1 v2 = Just Refl + matchArrayInstr (Parameter v1) (Parameter v2) | Just Refl <- matchVar v1 v2 = Just Refl + matchArrayInstr _ _ = Nothing + + encodeArrayInstr (Index v) = intHost $(hashQ ("Index" :: String)) <> encodeGroundVar v + encodeArrayInstr (Parameter v) = intHost $(hashQ ("Parameter" :: String)) <> encodeExpVar v + + inlineArrayInstr = \case + Parameter _ -> True + _ -> False + +encodeGroundR :: GroundR t -> Builder +encodeGroundR (GroundRscalar tp) = intHost $(hashQ ("Scalar" :: String)) <> encodeScalarType tp +encodeGroundR (GroundRbuffer tp) = intHost $(hashQ ("Buffer" :: String)) <> encodeScalarType tp + +encodeGroundsR :: GroundsR t -> Builder +encodeGroundsR = encodeTupR encodeGroundR + +encodeGroundVar :: GroundVar env t -> Builder +encodeGroundVar (Var repr ix) = encodeGroundR repr <> encodeIdx ix + +encodeGroundVars :: GroundVars env t -> Builder +encodeGroundVars = encodeTupR encodeGroundVar + +rnfGroundVar :: GroundVar env t -> () +rnfGroundVar = rnfVar rnfGroundR + +rnfGroundVars :: GroundVars env t -> () +rnfGroundVars = rnfTupR rnfGroundVar + +liftGroundR :: GroundR t -> CodeQ (GroundR t) +liftGroundR (GroundRscalar tp) = [|| GroundRscalar $$(liftScalarType tp) ||] +liftGroundR (GroundRbuffer tp) = [|| GroundRbuffer $$(liftScalarType tp) ||] + +liftGroundsR :: GroundsR t -> CodeQ (GroundsR t) +liftGroundsR = liftTupR liftGroundR + +liftGroundVar :: GroundVar env t -> CodeQ (GroundVar env t) +liftGroundVar = liftVar liftGroundR + +liftGroundVars :: GroundVars env t -> CodeQ (GroundVars env t) +liftGroundVars = liftTupR liftGroundVar + +paramIn :: ScalarType e -> GroundVar benv e -> OpenExp env benv e +paramIn t (Var _ ix) = ArrayInstr (Parameter $ Var t ix) Nil + +paramsIn :: HasCallStack => TypeR e -> GroundVars benv e -> OpenExp env benv e +paramsIn TupRunit TupRunit = Nil +paramsIn (TupRpair t1 t2) (TupRpair v1 v2) = paramsIn t1 v1 `Pair` paramsIn t2 v2 +paramsIn (TupRsingle t) (TupRsingle (Var _ ix)) = ArrayInstr (Parameter $ Var t ix) Nil +paramsIn _ _ = internalError "Tuple mismatch" + +paramIn' :: ExpVar benv e -> OpenExp env benv e +paramIn' v = ArrayInstr (Parameter v) Nil + +paramsIn' :: ExpVars benv e -> OpenExp env benv e +paramsIn' TupRunit = Nil +paramsIn' (TupRpair v1 v2) = paramsIn' v1 `Pair` paramsIn' v2 +paramsIn' (TupRsingle v) = ArrayInstr (Parameter v) Nil + +type ReindexPartial f env env' = forall a. Idx env a -> f (Idx env' a) + +-- The first argument is ReindexPartial, but without the forall and specialized to 't'. +-- This makes it usable in more situations. +-- +reindexVar :: Applicative f => (Idx env t -> f (Idx env' t)) -> Var s env t -> f (Var s env' t) +reindexVar k (Var repr ix) = Var repr <$> k ix + +reindexVars :: Applicative f => ReindexPartial f env env' -> Vars s env t -> f (Vars s env' t) +reindexVars k (TupRsingle var) = TupRsingle <$> reindexVar k var +reindexVars k (TupRpair v1 v2) = TupRpair <$> reindexVars k v1 <*> reindexVars k v2 +reindexVars _ TupRunit = pure TupRunit + +reindexArrayInstr :: Applicative f => ReindexPartial f env env' -> ArrayInstr env (s -> t) -> f (ArrayInstr env' (s -> t)) +reindexArrayInstr k (Index v) = Index <$> reindexVar k v +reindexArrayInstr k (Parameter v) = Parameter <$> reindexVar k v + +reindexExp :: (Applicative f, RebuildableExp e) => ReindexPartial f benv benv' -> e (ArrayInstr benv) env t -> f (e (ArrayInstr benv') env t) +reindexExp k = rebuildArrayInstrPartial (rebuildArrayInstrMap $ reindexArrayInstr k) + +reindexArg :: Applicative f => ReindexPartial f env env' -> Arg env t -> f (Arg env' t) +reindexArg k (ArgVar vars) = ArgVar <$> reindexVars k vars +reindexArg k (ArgExp e) = ArgExp <$> reindexExp k e +reindexArg k (ArgFun f) = ArgFun <$> reindexExp k f +reindexArg k (ArgArray m repr sh buffers) = ArgArray m repr <$> reindexVars k sh <*> reindexVars k buffers + +reindexArgs :: Applicative f => ReindexPartial f env env' -> Args env t -> f (Args env' t) +reindexArgs = reindexPreArgs reindexArg + +reindexPreArgs + :: Applicative f + => (forall f' t'. Applicative f' => ReindexPartial f' env env' -> s env t' -> f' (s env' t')) + -> ReindexPartial f env env' -> PreArgs (s env) t -> f (PreArgs (s env') t) +reindexPreArgs _ _ ArgsNil = pure ArgsNil +reindexPreArgs reindex k (a :>: as) = (:>:) <$> reindex k a <*> reindexPreArgs reindex k as + +reindexIdxSet + :: forall f env env' . Applicative f + => ReindexPartial f env env' + -> IdxSet.IdxSet env + -> f (IdxSet.IdxSet env') +reindexIdxSet k set = + foldr go (pure IdxSet.empty) (IdxSet.toList set) + where + go :: Exists (Idx env) -> f (IdxSet.IdxSet env') -> f (IdxSet.IdxSet env') + go (Exists ix) acc = IdxSet.insert <$> k ix <*> acc + +reindexAcc :: Applicative f => ReindexPartial f env env' -> PreOpenAcc op env t -> f (PreOpenAcc op env' t) +reindexAcc r (Exec opargs pa) = Exec opargs <$> reindexArgs r pa +reindexAcc r (Return tr) = Return <$> reindexVars r tr +reindexAcc r (Manifest var) = Manifest <$> reindexVar r var +reindexAcc r (Compute poe) = Compute <$> reindexExp r poe +reindexAcc r (Alet lhs tr poa poa') = reindexLHS r lhs $ \lhs' r' -> Alet lhs' tr <$> reindexAcc r poa <*> reindexAcc r' poa' +reindexAcc r (Alloc sr st tr) = Alloc sr st <$> reindexVars r tr +reindexAcc _ (Use st n bu) = pure $ Use st n bu +reindexAcc r (Unit var) = Unit <$> reindexVar r var +reindexAcc r (Acond var poa poa') = Acond <$> reindexVar r var <*> reindexAcc r poa <*> reindexAcc r poa' +reindexAcc r (Awhile tr poa poa' tr') = Awhile tr <$> reindexAfun r poa <*> reindexAfun r poa' <*> reindexVars r tr' +reindexAcc r (Aassert msg cond) = Aassert msg <$> reindexExp r cond +reindexAcc r (Aassume cond) = Aassume <$> reindexExp r cond +reindexAcc r (Atrace msg t) = Atrace msg <$> reindexArrayDescriptors r t +reindexAcc r (Fence set e) = Fence <$> reindexIdxSet r set <*> reindexAcc r e + +reindexArrayDescriptors :: (Applicative f) => ReindexPartial f env env' -> ArrayDescriptors env a -> f (ArrayDescriptors env' a) +reindexArrayDescriptors k = traverseTupR (reindexArrayDescriptor k) + +reindexArrayDescriptor :: Applicative f => ReindexPartial f env env' -> ArrayDescriptor env a -> f (ArrayDescriptor env' a) +reindexArrayDescriptor k (ArrayDescriptor shr sh buffers) = ArrayDescriptor shr <$> reindexVars k sh <*> reindexVars k buffers + +reindexAfun :: Applicative f => ReindexPartial f env env' -> PreOpenAfun op env t -> f (PreOpenAfun op env' t) +reindexAfun r (Abody poa) = Abody <$> reindexAcc r poa +reindexAfun r (Alam lhs poa) = reindexLHS r lhs $ \lhs' r' -> Alam lhs' <$> reindexAfun r' poa + +reindexLHS :: Applicative f => ReindexPartial f env env' -> LeftHandSide s t env env1 -> (forall env1'. LeftHandSide s t env' env1' -> ReindexPartial f env1 env1' -> r) -> r +reindexLHS r (LeftHandSideSingle st) k = k (LeftHandSideSingle st) $ \case + ZeroIdx -> pure ZeroIdx + SuccIdx idx -> SuccIdx <$> r idx +reindexLHS r (LeftHandSideWildcard tr) k = k (LeftHandSideWildcard tr) r +reindexLHS r (LeftHandSidePair left right) k = reindexLHS r left $ \left' r' -> + reindexLHS r' right $ \right' r'' -> + k (LeftHandSidePair left' right') r'' + +weakenReindex :: benv :> benv' -> ReindexPartial Identity benv benv' +weakenReindex k = Identity . (k >:>) + +weakenThroughReindex :: (benv :> benv') -> (ReindexPartial Identity benv benv' -> x -> Identity y) -> x -> y +weakenThroughReindex w k x = runIdentity $ k (weakenReindex w) x + +data AccessGroundR tp where + AccessGroundRscalar :: ScalarType tp -> AccessGroundR tp + AccessGroundRbuffer :: Modifier m -> ScalarType tp -> AccessGroundR (Buffer tp) + +readonlyGroundRToAccessGroundR :: Exists (GroundVar env) -> Exists (Var AccessGroundR env) +readonlyGroundRToAccessGroundR (Exists (Var (GroundRscalar tp) ix)) = Exists $ Var (AccessGroundRscalar tp) ix +readonlyGroundRToAccessGroundR (Exists (Var (GroundRbuffer tp) ix)) = Exists $ Var (AccessGroundRbuffer In tp) ix + +argVars :: Arg env t -> [Exists (Var AccessGroundR env)] +argVars (ArgVar vars) = flattenTupR $ mapTupR (\(Var tp ix) -> Var (AccessGroundRscalar tp) ix) vars +argVars (ArgExp e) = map readonlyGroundRToAccessGroundR $ expGroundVars e +argVars (ArgFun e) = map readonlyGroundRToAccessGroundR $ funGroundVars e +argVars (ArgArray m _ sh buffers) = go buffers $ map readonlyGroundRToAccessGroundR $ flattenTupR sh + where + go :: GroundVars env s -> [Exists (Var AccessGroundR env)] -> [Exists (Var AccessGroundR env)] + go TupRunit accum = accum + go (TupRsingle (Var (GroundRbuffer tp) ix)) accum = Exists (Var (AccessGroundRbuffer m tp) ix) : accum + go (TupRsingle _) _ = error "Expected buffer" + go (TupRpair v1 v2) accum = go v1 $ go v2 accum + +argsVars :: Args env t -> [Exists (Var AccessGroundR env)] +argsVars (a :>: as) = argVars a ++ argsVars as +argsVars ArgsNil = [] + +argsInputs :: Args env t -> [Exists (Idx env)] +argsInputs = mapMaybe input . argsVars + where + input :: Exists (Var AccessGroundR env) -> Maybe (Exists (Idx env)) + input (Exists (Var (AccessGroundRbuffer Out _) _)) = Nothing + input (Exists (Var _ idx)) = Just $ Exists idx + +argsOutputs :: Args env t -> [Exists (Idx env)] +argsOutputs = mapMaybe output . argsVars + where + output :: Exists (Var AccessGroundR env) -> Maybe (Exists (Idx env)) + output (Exists (Var (AccessGroundRbuffer Out _) idx)) = Just $ Exists idx + output (Exists (Var (AccessGroundRbuffer Mut _) idx)) = Just $ Exists idx + output _ = Nothing + +expGroundVars :: OpenExp env benv t -> [Exists (GroundVar benv)] +expGroundVars = map arrayInstrGroundVars . arrayInstrsInExp + +funGroundVars :: OpenFun env benv t -> [Exists (GroundVar benv)] +funGroundVars = map arrayInstrGroundVars . arrayInstrsInFun + +arrayInstrGroundVars :: Exists (ArrayInstr benv) -> Exists (GroundVar benv) +arrayInstrGroundVars (Exists (Parameter (Var tp ix))) = Exists $ Var (GroundRscalar tp) ix +arrayInstrGroundVars (Exists (Index var)) = Exists var + +mapAccExecutable :: (forall args benv'. op args -> Args benv' args -> op' args) -> PreOpenAcc op benv t -> PreOpenAcc op' benv t +mapAccExecutable f = \case + Exec op args -> Exec (f op args) args + Return vars -> Return vars + Manifest var -> Manifest var + Compute e -> Compute e + Alet lhs uniqueness bnd body -> Alet lhs uniqueness (mapAccExecutable f bnd) (mapAccExecutable f body) + Alloc shr tp sh -> Alloc shr tp sh + Use tp n buffer -> Use tp n buffer + Unit vars -> Unit vars + Acond var a1 a2 -> Acond var (mapAccExecutable f a1) (mapAccExecutable f a2) + Awhile uniqueness c g a -> Awhile uniqueness (mapAfunExecutable f c) (mapAfunExecutable f g) a + Aassert msg cond -> Aassert msg cond + Aassume cond -> Aassume cond + Atrace msg t -> Atrace msg t + Fence set e -> Fence set (mapAccExecutable f e) + +mapAfunExecutable :: (forall args benv'. op args -> Args benv' args -> op' args) -> PreOpenAfun op benv t -> PreOpenAfun op' benv t +mapAfunExecutable f (Abody a) = Abody $ mapAccExecutable f a +mapAfunExecutable f (Alam lhs a) = Alam lhs $ mapAfunExecutable f a + +groundToExpVar :: TypeR e -> GroundVars benv e -> ExpVars benv e +groundToExpVar (TupRsingle t) (TupRsingle (Var _ ix)) = TupRsingle (Var t ix) +groundToExpVar (TupRpair t1 t2) (TupRpair v1 v2) = groundToExpVar t1 v1 `TupRpair` groundToExpVar t2 v2 +groundToExpVar TupRunit TupRunit = TupRunit +groundToExpVar _ _ = internalError "Impossible pair" + +expToGroundVar :: ExpVars env e -> GroundVars env e +expToGroundVar TupRunit = TupRunit +expToGroundVar (TupRsingle (Var a ix)) = TupRsingle $ Var (GroundRscalar a) ix +expToGroundVar (TupRpair x y) = TupRpair (expToGroundVar x) (expToGroundVar y) + +class NFData' f where + rnf' :: f a -> () + +instance NFData' op => NFData (OperationAcc op env a) where + rnf (Exec op args) = rnf' op `seq` rnfArgs args + rnf (Return vars) = rnfGroundVars vars + rnf (Manifest var) = rnfGroundVar var + rnf (Compute expr) = rnfOpenExp expr + rnf (Alet lhs us bnd a) = rnfLeftHandSide rnfGroundR lhs `seq` rnfTupR rnfUniqueness us `seq` rnf bnd `seq` rnf a + rnf (Alloc shr tp sh) = rnfShapeR shr `seq` rnfScalarType tp `seq` rnfTupR rnfExpVar sh + rnf (Use tp n buffer) = n `seq` buffer `seq` rnfScalarType tp + rnf (Unit var) = rnfVar rnfScalarType var + rnf (Acond cond true false) = rnfVar rnfScalarType cond `seq` rnf true `seq` rnf false + rnf (Awhile us cond step initial) = rnfTupR rnfUniqueness us `seq` rnf cond `seq` rnf step `seq` rnfGroundVars initial + rnf (Aassert msg cond) = rnf msg `seq` rnfOpenExp cond + rnf (Aassume cond) = rnfOpenExp cond + rnf (Atrace msg t) = rnf msg `seq` rnf' t + rnf (Fence set a) = IdxSet.rnfIdxSet set `seq` rnf a + +instance NFData' (ArrayDescriptor env) where + rnf' (ArrayDescriptor shr sh buffers) = rnfShapeR shr `seq` rnfGroundVars sh `seq` rnfGroundVars buffers + +instance NFData' op => NFData (OperationAfun op env a) where + rnf (Abody a) = rnf a + rnf (Alam lhs f) = rnfLeftHandSide rnfGroundR lhs `seq` rnf f + +instance NFData' arg => NFData' (PreArgs arg) where + rnf' ArgsNil = () + rnf' (a :>: args) = rnf' a `seq` rnf' args + +instance (NFData' s) => NFData' (TupR s) where + rnf' = rnfTupR rnf' + + +data GroundRWithUniqueness t where + GroundRWithUniqueness :: GroundR t -> Uniqueness t -> GroundRWithUniqueness t + +groundsRWithUniquenesses :: GroundsR t -> Uniquenesses t -> TupR GroundRWithUniqueness t +groundsRWithUniquenesses TupRunit _ = TupRunit +groundsRWithUniquenesses (TupRpair g1 g2) (TupRpair u1 u2) = groundsRWithUniquenesses g1 u1 `TupRpair` groundsRWithUniquenesses g2 u2 +groundsRWithUniquenesses (TupRpair g1 g2) _ = groundsRWithUniquenesses g1 (TupRsingle Shared) `TupRpair` groundsRWithUniquenesses g2 (TupRsingle Shared) +groundsRWithUniquenesses (TupRsingle tp) (TupRsingle u) = TupRsingle $ GroundRWithUniqueness tp u +groundsRWithUniquenesses _ _ = internalError "Tuple mismatch" + +lhsWithUniquesses :: GLeftHandSide t env env' -> Uniquenesses t -> LeftHandSide GroundRWithUniqueness t env env' +lhsWithUniquesses (LeftHandSideWildcard g) uniquenesses = LeftHandSideWildcard $ groundsRWithUniquenesses g uniquenesses +lhsWithUniquesses (LeftHandSideSingle g) (TupRsingle u) = LeftHandSideSingle $ GroundRWithUniqueness g u +lhsWithUniquesses (LeftHandSidePair l1 l2) (TupRpair u1 u2) = LeftHandSidePair (lhsWithUniquesses l1 u1) (lhsWithUniquesses l2 u2) +lhsWithUniquesses (LeftHandSidePair l1 l2) _ = LeftHandSidePair (lhsWithUniquesses l1 $ TupRsingle Shared) (lhsWithUniquesses l2 $ TupRsingle Shared) +lhsWithUniquesses _ _ = internalError "Tuple mismatch" diff --git a/src/Data/Array/Accelerate/AST/Partitioned.hs b/src/Data/Array/Accelerate/AST/Partitioned.hs new file mode 100644 index 000000000..62aeeb86d --- /dev/null +++ b/src/Data/Array/Accelerate/AST/Partitioned.hs @@ -0,0 +1,1315 @@ +{-# LANGUAGE MultiWayIf #-} +{-# LANGUAGE TypeFamilyDependencies #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE ViewPatterns #-} +{-# OPTIONS_GHC -Wno-name-shadowing #-} +{-# OPTIONS_GHC -Wno-orphans #-} + +-- | +-- Module : Data.Array.Accelerate.AST.Partitioned +-- Copyright : [2008..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- +module Data.Array.Accelerate.AST.Partitioned ( + module Data.Array.Accelerate.AST.Operation, + module Data.Array.Accelerate.AST.Partitioned, + GroundR(..), NFData'(..), Arg(..), + AccessGroundR(..), + PreArgs(..), Modifier(..), + Node(..), Comp, GVal +) where + + +import Data.Array.Accelerate.AST.Operation hiding (OperationAcc, OperationAfun) + +import Prelude hiding ( take ) +import Data.Bifunctor +import Data.Maybe (isJust, fromMaybe) +import Data.Array.Accelerate.Trafo.Operation.Simplify (SimplifyOperation(..)) +import Data.Array.Accelerate.Representation.Array (Array, Buffer, Buffers, ArrayR (..)) +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.Representation.Shape (ShapeR (..), shapeType, rank) +import Data.Type.Equality +import Unsafe.Coerce (unsafeCoerce) +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Representation.Ground +import Data.Array.Accelerate.Type +import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.Environment +import Data.Array.Accelerate.Error + +import Data.Array.Accelerate.Trafo.Var +import Data.Array.Accelerate.Trafo.Exp.Substitution +import Data.Array.Accelerate.Trafo.Substitution +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Labels (LabelledArgs, LabelledArg (..), Node(..), ArgLabel (..), EnvLabel, Comp, GVal, EnvLabels, EnvVals, valsNodes) +import Data.List (sortOn, partition, groupBy, nubBy) +import qualified Data.Functor.Const as C +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Graph (LabelledArgOp (..), BackendClusterArg, MakesILP (..), LabelledArgsOp, BackendCluster) +import Data.Array.Accelerate.Trafo.Operation.LiveVars +import Control.Applicative ((<|>)) +import Data.Array.Accelerate.AST.Var (varsType) +import Data.Array.Accelerate.Analysis.Match + + + +type PartitionedAcc op = PreOpenAcc (Clustered op) +type PartitionedAfun op = PreOpenAfun (Clustered op) + + +data Clustered op args = Clustered (Cluster op args) (BackendCluster op args) + +instance NFData' op => NFData' (Clustered op) where + rnf' :: NFData' op => Clustered op a -> () + rnf' _c = () -- TODO + +data Cluster op args where + SingleOp :: SingleOp op args -> Node Comp -> Cluster op args + Fused :: Fusion largs rargs args + -> Cluster op largs + -> Cluster op rargs + -> Cluster op args + +data SingleOp op args where + Single :: op opArgs + -> ClusterArgs (FunToEnv args) opArgs + -> SingleOp op args + +type ClusterArgs env = PreArgs (ClusterArg env) +data ClusterArg env t where + -- Perhaps require that 't' is not 'In sh e' or 'Out sh e' + ClusterArgSingle + :: Idx env t + -> ClusterArg env t + + ClusterArgArray + -- Perhaps require that 'm' is 'In' or 'Out' + :: Modifier m -- Not Mut + -> ShapeR sh + -> TypeR e + -> ClusterArgBuffers env m sh e + -> ClusterArg env (m sh e) + +data ClusterArgBuffers env m sh t where + ClusterArgBuffersPair + :: ClusterArgBuffers env m sh l + -> ClusterArgBuffers env m sh r + -> ClusterArgBuffers env m sh (l, r) + + ClusterArgBuffersDead + :: TypeR t -- TODO: Remove? This info is already stored in ClusterArgArray + -> Idx env (Var' sh) + -> ClusterArgBuffers env Out sh t + + ClusterArgBuffersLive + :: TypeR t + -> Idx env (m sh t) + -> ClusterArgBuffers env m sh t + +type family FunToEnv f = env | env -> f where + FunToEnv () = () + FunToEnv (s -> t) = (FunToEnv t, s) + +funToEnvPrj :: PreArgs a f' -> Idx (FunToEnv f') e -> a e +funToEnvPrj (a :>: _ ) ZeroIdx = a +funToEnvPrj (_ :>: as) (SuccIdx idx) = funToEnvPrj as idx +funToEnvPrj _ _ = internalError "Empty environment" + +-- These correspond to the inference rules in the paper +-- datatype describing how to combine the arguments of two fused clusters +data Fusion largs rargs args where + EmptyF :: Fusion () () () + Vertical :: ArrayR (Array sh e) + -> Fusion l r a + -> Fusion (Out sh e -> l) (In sh e -> r) (Var' sh -> a) + Horizontal :: Fusion l r a + -> Fusion (In sh e -> l) (In sh e -> r) (In sh e -> a) + Diagonal :: Fusion l r a + -> Fusion (Out sh e -> l) (In sh e -> r) (Out sh e -> a) + IntroI1 :: Fusion l r a + -> Fusion (In sh e -> l) r (In sh e -> a) + IntroI2 :: Fusion l r a + -> Fusion l (In sh e -> r) (In sh e -> a) + IntroO1 :: Fusion l r a + -> Fusion (Out sh e -> l) r (Out sh e -> a) + IntroO2 :: Fusion l r a + -> Fusion l (Out sh e -> r) (Out sh e -> a) + -- not in the paper; not meant for array args: + -- Ivo: 'mut' arguments do go in here, right? + IntroL :: Fusion l r a -> Fusion (x -> l) r (x -> a) + IntroR :: Fusion l r a -> Fusion l (x -> r) (x -> a) +deriving instance Show (Fusion l r total) + + +combine :: Arg env (f l) -> Arg env (f r) -> Arg env (f (l,r)) +combine (ArgArray m (ArrayR shr lr) sh l) (ArgArray _ (ArrayR _ rr) _ r) = ArgArray m (ArrayR shr (TupRpair lr rr)) sh (TupRpair l r) +combine _ _ = error "SOA'd non-arrays" + +split :: Arg env (f (l,r)) -> (Arg env (f l), Arg env (f r)) +split (ArgArray In (ArrayR shr (TupRpair rl rr)) sh (TupRpair bufl bufr)) = (ArgArray In (ArrayR shr rl) sh bufl, ArgArray In (ArrayR shr rr) sh bufr) +split (ArgArray Out (ArrayR shr (TupRpair rl rr)) sh (TupRpair bufl bufr)) = (ArgArray Out (ArrayR shr rl) sh bufl, ArgArray Out (ArrayR shr rr) sh bufr) +split _ = error "non-array soa" + +splitEnvVals :: EnvVals (l,r) -> (EnvVals l, EnvVals r) +splitEnvVals (TupRpair el er, TupRpair bsl bsr, TupRpair ul ur) = ((el, bsl, ul), (er, bsr, ur)) +splitEnvVals (_, _, _) = error "pair in single" + +splitArgLabel :: ArgLabel (f (el,er)) -> (ArgLabel (f el), ArgLabel (f er)) +splitArgLabel (Arr arr sh) = bimap (`Arr` sh) (`Arr` sh) $ splitEnvVals arr +splitArgLabel (NotArr _) = error "SOA'd non-array arg" + +splitLabelledArgs :: LabelledArg env (f (l,r)) -> (LabelledArg env (f l), LabelledArg env (f r)) +splitLabelledArgs (L arg label) = uncurry bimap (bimap L L $ split arg) (splitArgLabel label) + +splitLabelledArgsOp :: LabelledArgOp op env (f (l,r)) -> (LabelledArgOp op env (f l), LabelledArgOp op env (f r)) +splitLabelledArgsOp (LOp arg label barg) = uncurry bimap (uncurry bimap (bimap LOp LOp $ split arg) (splitArgLabel label)) (barg, barg) + + + +left :: Fusion largs rargs args -> Args env args -> Args env largs +left = left' (\repr (ArgVar sh) -> ArgArray Out repr (mapTupR (\(Var t ix)->Var (GroundRscalar t) ix) sh) er) + where + er = error "accessing fused away array" +left' :: (forall sh e. ArrayR (Array sh e) -> f (Var' sh) -> f (Out sh e)) -> Fusion largs rargs args -> PreArgs f args -> PreArgs f largs +left' _ EmptyF ArgsNil = ArgsNil +left' k (Vertical repr f) (arg :>: args) = k repr arg :>: left' k f args +left' k (Horizontal f) (arg:>:args) = arg :>: left' k f args +left' k (Diagonal f) (arg:>:args) = arg :>: left' k f args +left' k (IntroI1 f) (arg:>:args) = arg :>: left' k f args +left' k (IntroI2 f) (_ :>:args) = left' k f args +left' k (IntroO1 f) (arg:>:args) = arg :>: left' k f args +left' k (IntroO2 f) (_ :>:args) = left' k f args +left' k (IntroL f) (arg:>:args) = arg :>: left' k f args +left' k (IntroR f) (_ :>:args) = left' k f args + +right :: Fusion largs rargs args -> Args env args -> Args env rargs +right = right' varToIn outToIn +right' :: (forall sh e. ArrayR (Array sh e) -> f (Var' sh) -> f (In sh e)) -> (forall sh e. f (Out sh e) -> f (In sh e)) -> Fusion largs rargs args -> PreArgs f args -> PreArgs f rargs +right' _ _ EmptyF ArgsNil = ArgsNil +right' vi oi (Vertical arr f) (arg :>: args) = vi arr arg :>: right' vi oi f args +right' vi oi (Diagonal f) (arg :>: args) = oi arg :>: right' vi oi f args +right' vi oi (Horizontal f) (arg :>: args) = arg :>: right' vi oi f args +right' vi oi (IntroI1 f) (_ :>: args) = right' vi oi f args +right' vi oi (IntroI2 f) (arg :>: args) = arg :>: right' vi oi f args +right' vi oi (IntroO1 f) (_ :>: args) = right' vi oi f args +right' vi oi (IntroO2 f) (arg :>: args) = arg :>: right' vi oi f args +right' vi oi (IntroL f) (_ :>: args) = right' vi oi f args +right' vi oi (IntroR f) (arg :>: args) = arg :>: right' vi oi f args + +varToIn :: ArrayR (Array sh e) -> Arg env (Var' sh) -> Arg env (In sh e) +varToIn (ArrayR shr ty) (ArgVar sh) + | Just Refl <- matchShapeR shr (varsToShapeR sh) = ArgArray In (ArrayR shr ty) (mapTupR (\(Var t ix)->Var (GroundRscalar t) ix) sh) er + | otherwise = error "wrong shape?" + where er = error "accessing fused away array" +outToIn :: Arg env (Out sh e) -> Arg env (In sh e) +outToIn (ArgArray Out x y z) = ArgArray In x y z +inToOut :: Arg env (In sh e) -> Arg env (Out sh e) +inToOut (ArgArray In x y z) = ArgArray Out x y z +varToOut :: ArrayR (Array sh e) -> Arg env (Var' sh) -> Arg env (Out sh e) +varToOut arr = inToOut . varToIn arr + +both :: (forall sh e. f (Out sh e) -> f (In sh e) -> f (Var' sh)) -> Fusion largs rargs args -> PreArgs f largs -> PreArgs f rargs -> PreArgs f args +both _ EmptyF ArgsNil ArgsNil = ArgsNil +both k (Vertical _ f) (l:>:ls) (r:>:rs) = k l r :>: both k f ls rs +both k (Diagonal f) (l:>:ls) (_:>:rs) = l :>: both k f ls rs +both k (Horizontal f) (l:>:ls) (_:>:rs) = l :>: both k f ls rs +both k (IntroI1 f) (l:>:ls) rs = l :>: both k f ls rs +both k (IntroI2 f) ls (r:>:rs) = r :>: both k f ls rs +both k (IntroO1 f) (l:>:ls) rs = l :>: both k f ls rs +both k (IntroO2 f) ls (r:>:rs) = r :>: both k f ls rs +both k (IntroL f) (l:>:ls) rs = l :>: both k f ls rs +both k (IntroR f) ls (r:>:rs) = r :>: both k f ls rs + + + +flattenClustered :: Clustered op args -> [Exists op] +flattenClustered (Clustered c _) = flattenCluster c + +flattenCluster :: Cluster op args -> [Exists op] +flattenCluster = (`go` []) + where + go :: Cluster op args -> [Exists op] -> [Exists op] + go (SingleOp (Single op _) _) accum = Exists op : accum + go (Fused _ a b) accum = go a $ go b accum + +varsToShapeR :: Vars ScalarType g sh -> ShapeR sh +varsToShapeR = typeRtoshapeR . varsType + +typeRtoshapeR :: TypeR sh -> ShapeR sh +typeRtoshapeR TupRunit = ShapeRz +typeRtoshapeR (TupRpair sh (TupRsingle (SingleScalarType (NumSingleType (IntegralNumType TypeInt))))) = ShapeRsnoc (typeRtoshapeR sh) +typeRtoshapeR _ = error "not a shape" + +type family InArg a = b | b -> a where + InArg (In sh e ) = Value sh e + InArg (Mut sh e) = Mut sh e + InArg (Exp' e) = Exp' e + InArg (Fun' e) = Fun' e + InArg (Var' e) = Var' e + InArg (Out sh e) = Sh sh e +type family InArgs a = b | b -> a where + InArgs () = () + InArgs (a -> x) = (InArgs x, InArg a) +type family InArgsOf a where + InArgsOf () = () + InArgsOf (a -> x) = InArg a -> InArgsOf x + +type family InArg' a where + InArg' (Value sh e) = In sh e + InArg' (Mut sh e ) = Mut sh e + InArg' (Exp' e ) = Exp' e + InArg' (Fun' e ) = Fun' e + InArg' (Var' e ) = Var' e + InArg' (Sh sh e ) = Out sh e +type family InArgs' a where + InArgs' () = () + InArgs' (x, a) = InArg' a -> InArgs' x + +inargslemma :: InArgs' (InArgs a) :~: a +inargslemma = unsafeCoerce Refl +inarglemma :: InArg (InArg' a) :~: a +inarglemma = unsafeCoerce Refl + +type family OutArgs a where + OutArgs () = () + OutArgs (Out sh e -> x) = (OutArgs x, Value sh e) + OutArgs (_ -> x) = OutArgs x +type family OutArgsOf a where + OutArgsOf () = () + OutArgsOf (Out sh e -> x) = Out sh e -> OutArgsOf x + OutArgsOf (_ -> x) = OutArgsOf x + +type family FromArg env a where + FromArg env (Exp' e) = Exp env e + FromArg env (Var' e) = ExpVars env e + FromArg env (Fun' e) = Fun env e + FromArg env (Mut sh e) = (ArrayDescriptor env (Array sh e), TypeR e) + FromArg env x = x + +type family FromArgs env a where + FromArgs env () = () + FromArgs env (a, x) = (FromArgs env a, FromArg env x) + +type FromIn env args = FromArgs env ( InArgs args) +type FromOut env args = FromArgs env (OutArgs args) + + +data Value sh e = Value e (Sh sh e) +data Sh sh e = Shape (ShapeR sh) sh + +class TupRmonoid f where + pair' :: f a -> f b -> f (a,b) + unpair' :: f (a,b) -> (f a, f b) + +fromTupR :: TupRmonoid f => f () -> TupR f a -> f a +fromTupR u TupRunit = u +fromTupR _ (TupRsingle s) = s +fromTupR u (TupRpair l r) = pair' (fromTupR u l) (fromTupR u r) + + +varsGet' :: Vars s env t -> Env f env -> TupR f t +varsGet' t env = mapTupR (\(Var _ idx) -> prj' idx env) t + +instance TupRmonoid (TupR f) where + pair' = TupRpair + unpair' (TupRpair l r) = (l, r) + unpair' _ = error "nope" + -- injL t p = TupRpair t (proofToR p) + -- injR t p = TupRpair (proofToR p) t + + + +unOpLabels' :: LabelledArgsOp op env args -> LabelledArgs env args +unOpLabels' = mapArgs unOpLabel + +unOpLabel :: LabelledArgOp op env args -> LabelledArg env args +unOpLabel (LOp arg l _) = L arg l + +data Both f g a = Both (f a) (g a) deriving (Show, Eq) +fst' :: Both f g a -> f a +fst' (Both x _) = x +snd' :: Both f g a -> g a +snd' (Both _ y) = y + +instance (TupRmonoid f, TupRmonoid g) => TupRmonoid (Both f g) where + pair' (Both a b) (Both c d) = Both (pair' a c) (pair' b d) + unpair' (Both (unpair' -> (a, c)) (unpair' -> (b, d))) = (Both a b, Both c d) + + +zipArgs :: PreArgs f a -> PreArgs g a -> PreArgs (Both f g) a +zipArgs ArgsNil ArgsNil = ArgsNil +zipArgs (f:>:fs) (g:>:gs) = Both f g :>: zipArgs fs gs + +onZipped :: (f a -> f b -> f c) -> (g a -> g b -> g c) -> (Both f g a -> Both f g b -> Both f g c) +onZipped f g (Both fa ga) (Both fb gb) = Both (f fa fb) (g ga gb) + +-- assumes that the arguments are already sorted! +fuse :: MakesILP op + => LabelledArgsOp op env l -> LabelledArgsOp op env r -> PreArgs f l -> PreArgs f r -> Clustered op l -> Clustered op r + -> (forall sh e. f (Out sh e) -> f (In sh e) -> f (Var' sh)) + -> (forall args. PreArgs f args -> Clustered op args -> result) + -> result +fuse labl labr largs rargs (Clustered cl bl) (Clustered cr br) c k = fuse' labl labr (zipArgs largs bl) (zipArgs rargs br) cl cr (onZipped c combineBackendClusterArg) + $ \args c' -> k (mapArgs fst' args) (Clustered c' $ mapArgs snd' args) + +-- assumes that the arguments are already sorted! +fuse' :: MakesILP op => LabelledArgsOp op env l -> LabelledArgsOp op env r -> PreArgs f l -> PreArgs f r -> Cluster op l -> Cluster op r + -> (forall sh e. f (Out sh e) -> f (In sh e) -> f (Var' sh)) + -> (forall args. PreArgs f args -> Cluster op args -> result) + -> result +fuse' labl labr largs rargs l r c k = + mkFused labl labr $ \f -> k (both c f largs rargs) (Fused f l r) + +mkFused :: MakesILP op => LabelledArgsOp op env l -> LabelledArgsOp op env r -> (forall args. Fusion l r args -> result) -> result +mkFused ArgsNil ArgsNil k = k EmptyF +mkFused ArgsNil (LOp r _ _ :>: rs) k = mkFused ArgsNil rs $ \f -> k (addright r f) +mkFused (LOp l _ _ :>: ls) ArgsNil k = mkFused ls ArgsNil $ \f -> k (addleft l f) +mkFused ((LOp l (NotArr _) _) :>: ls) rs k = mkFused ls rs $ \f -> k (addleft l f) +mkFused ls ((LOp r (NotArr _) _) :>: rs) k = mkFused ls rs $ \f -> k (addright r f) +mkFused ((LOp l (Arr (TupRunit, _, _) _) _) :>: ls) rs k = mkFused ls rs $ \f -> k (addleft l f) +mkFused ls ((LOp r (Arr (TupRunit, _, _) _) _) :>: rs) k = mkFused ls rs $ \f -> k (addright r f) +mkFused (l'@(LOp l _ bl) :>: ls) (r'@(LOp r _ br) :>: rs) k + | Just le <- getElabelForSort $ unOpLabel l' + , Just re <- getElabelForSort $ unOpLabel r' + = case compare le re of + LT -> mkFused ls (r':>:rs) $ \f -> k (addleft l f) + GT -> mkFused (l':>:ls) rs $ \f -> k (addright r f) + EQ -> mkFused ls rs $ \f -> if bl == br then addboth l r f k else k (addleft l (addright r f)) +mkFused ((LOp l@(ArgArray Mut _ _ _) _ _) :>: ls) rs k = mkFused ls rs $ \f -> k (addleft l f) +mkFused ls ((LOp r@(ArgArray Mut _ _ _) _ _) :>: rs) k = mkFused ls rs $ \f -> k (addright r f) +mkFused ((LOp _ (Arr (TupRpair{}, _, _) _) _) :>: _) _ _ = error "not soa'd array" +mkFused _ ((LOp _ (Arr (TupRpair{}, _, _) _) _) :>: _) _ = error "not soa'd array" +mkFused _ _ _ = error "exhaustive" + +addleft :: Arg env arg -> Fusion left right args -> Fusion (arg->left) right (arg->args) +addleft (ArgVar _ ) f = IntroL f +addleft (ArgExp _ ) f = IntroL f +addleft (ArgFun _ ) f = IntroL f +addleft (ArgArray Mut _ _ _) f = IntroL f +addleft (ArgArray In _ _ _) f = IntroI1 f +addleft (ArgArray Out _ _ _) f = IntroO1 f + +addright :: Arg env arg -> Fusion left right args -> Fusion left (arg->right) (arg->args) +addright (ArgVar _ ) f = IntroR f +addright (ArgExp _ ) f = IntroR f +addright (ArgFun _ ) f = IntroR f +addright (ArgArray Mut _ _ _) f = IntroR f +addright (ArgArray In _ _ _) f = IntroI2 f +addright (ArgArray Out _ _ _) f = IntroO2 f + +addboth :: Arg env l -> Arg env r -> Fusion left right args -> (forall arg. Fusion (l->left) (r->right) (arg -> args) -> result) -> result +addboth (ArgArray In r1 _ _) (ArgArray In r2 _ _) f k + | Just Refl <- matchArrayR r1 r2 = k $ Horizontal f + | otherwise = error "types of horizontally fused arrays are different" +addboth (ArgArray Out r1 _ _) (ArgArray In r2 _ _) f k + | Just Refl <- matchArrayR r1 r2 = k $ Diagonal f + | otherwise = error "types of diagonally fused arrays are different" +addboth (ArgArray Out _ _ _) (ArgArray Out _ _ _) _ _ = error "two producers of the same array" +addboth (ArgArray In _ _ _) (ArgArray Out _ _ _) _ _ = error "reverse vertical/diagonal" +addboth _ _ _ _ = error "fusing non-arrays" + +getElabelForSort :: LabelledArg env a -> Maybe EnvLabel +getElabelForSort (L (ArgArray m (ArrayR _ TupRsingle{}) _ _) (Arr (TupRsingle (C.Const e), _, _) _)) + | In <- m = Just e + | Out <- m = Just e +getElabelForSort _ = Nothing + +singleton + :: MakesILP op + => Node Comp + -> LabelledArgsOp op env args + -> op args + -> (forall args'. Clustered op args' -> LabelledArgsOp op env args' -> r) + -> r +singleton l largs op k + | Exists sorted <- sortAndExpandArgs largs + , opArgs <- createClusterArgs sorted largs + = k + (Clustered + (SingleOp (Single op opArgs) l) + (mapArgs getClusterArg sorted)) + sorted + +createClusterArgs + :: forall op env sorted args. + LabelledArgsOp op env sorted + -> LabelledArgsOp op env args + -> ClusterArgs (FunToEnv sorted) args +createClusterArgs sorted = go 0 + where + go + :: Int -- Number of ClusterArgSingle that we've already handled + -> LabelledArgsOp op env args' + -> ClusterArgs (FunToEnv sorted) args' + go _ ArgsNil = ArgsNil + go n (a :>: as) = a' :>: go n' as + where + a' = createClusterArg n sorted a + n' = case a' of + ClusterArgSingle _ -> n + 1 + _ -> n + +-- We ensure that environment starts with all non-in-or-out-array arguments, +-- then all array arguments. That simplifies the search for non-array arguments in 'sorted'. +createClusterArg + :: forall op env sorted arg. + Int -- Number of ClusterArgSingle that we've already handled + -> LabelledArgsOp op env sorted + -> LabelledArgOp op env arg + -> ClusterArg (FunToEnv sorted) arg +createClusterArg _ sorted (LOp (ArgArray (m :: Modifier m) (ArrayR (shr :: ShapeR sh) tp) sh _) (Arr (labels, _, _) _) _) + | inOrOut m = ClusterArgArray m shr tp $ go tp labels + where + inOrOut :: Modifier m -> Bool + inOrOut In = True + inOrOut Out = True + inOrOut _ = False + + go :: TypeR t -> EnvLabels (Buffers t) -> ClusterArgBuffers (FunToEnv sorted) m sh t + go TupRunit TupRunit + = ClusterArgBuffersLive TupRunit $ findUnit sorted + go (TupRsingle t) ((TupRsingle (C.Const label))) + = ClusterArgBuffersLive (TupRsingle t) $ findLabel (TupRsingle t) label sorted + go (TupRpair t1 t2) ((TupRpair l1 l2)) + = go t1 l1 `ClusterArgBuffersPair` go t2 l2 + go _ _ = internalError "Tuple mismatch" + + findUnit + :: LabelledArgsOp op env sorted' + -> Idx (FunToEnv sorted') (m sh ()) + findUnit = \case + LOp (ArgArray m' (ArrayR _ TupRunit) sh' _) _ _ :>: _ + | Just Refl <- matchModifier m m' + , Just Refl <- matchVars sh sh' + -> ZeroIdx + _ :>: sorted' -> SuccIdx $ findUnit sorted' + ArgsNil -> internalError "Unit array not found in sorted arguments" + + findLabel + :: TupR ScalarType t + -> EnvLabel + -> LabelledArgsOp op env sorted' + -> Idx (FunToEnv sorted') (m sh t) + findLabel tp label = \case + LOp (ArgArray m' (ArrayR _ tp') sh' _) (Arr (TupRsingle (C.Const label'), _, _) _) _ :>: _ + | label == label' + , Refl <- expectOr "Modifier didn't match" $ matchModifier m m' + , Refl <- expectOr "Shapes didn't match" $ matchVars sh sh' + , Refl <- expectOr "Array types didn't match" $ matchTypeR tp tp' + -> ZeroIdx + _ :>: sorted' -> SuccIdx $ findLabel tp label sorted' + ArgsNil -> internalError "Label not found in sorted arguments" + + expectOr _ (Just x) = x + expectOr err _ = internalError err + +-- ClusterArgSingles are not shuffled, they are kept in the original order. +-- To find the ith ClusterArgSingle, we thus search for the ith ClusterArgSingle in 'sorted'. +-- Function 'go' does this search, and skips over any In or Out arguments. +createClusterArg count sorted (LOp arg _ _) = ClusterArgSingle $ go count sorted + where + go :: Int -> LabelledArgsOp op env sorted' -> Idx (FunToEnv sorted') arg + go n (LOp (ArgArray In _ _ _) _ _ :>: args) = SuccIdx $ go n args + go n (LOp (ArgArray Out _ _ _) _ _ :>: args) = SuccIdx $ go n args + go 0 (LOp arg' _ _ :>: _) + -- We only validate the type of the argument. + -- We assume the ClusterArgSingles are not shuffled. We only do this + -- to inform GHC that this is sound, and to catch situations where things + -- are really broken. + | Just Refl <- matchArgType arg arg' = ZeroIdx + | otherwise = internalError "ClusterArgSingles are not in the right order" + go n (_ :>: args) = SuccIdx $ go (n - 1) args + go _ ArgsNil = internalError "ClusterArgSingle not found" + +-- Checks whether two types are equal, to a certain extent. +-- If this function returns Just, the types of the two arguments are equal. +-- However, the inverse is not always true. +-- Use this function if you know that the arguments are equal, but want +-- type-level proof that the types match. +matchArgType :: Arg env s -> Arg env t -> Maybe (s :~: t) +matchArgType (ArgVar v1) (ArgVar v2) + | Just Refl <- matchVars v1 v2 = Just Refl +matchArgType (ArgExp e1) (ArgExp e2) + | Just Refl <- matchTypeR (expType e1) (expType e2) = Just Refl +matchArgType (ArgFun fun1) (ArgFun fun2) + | Just Refl <- go fun1 fun2 = Just Refl + where + go :: OpenFun env1 env t1 -> OpenFun env2 env t2 -> Maybe (t1 :~: t2) + go (Lam lhs1 f1) (Lam lhs2 f2) + | Just Refl <- matchTypeR (lhsToTupR lhs1) (lhsToTupR lhs2) + , Just Refl <- go f1 f2 + = Just Refl + go (Body e1) (Body e2) = matchTypeR (expType e1) (expType e2) + go _ _ = Nothing +matchArgType (ArgArray m1 repr1 _ _) (ArgArray m2 repr2 _ _) + | Just Refl <- matchModifier m1 m2 + , Just Refl <- matchArrayR repr1 repr2 = Just Refl +matchArgType _ _ = Nothing + +-- Checks whether the two Args are the same. +-- Assumes that the Args are both arrays and that +-- the arrays both consist of a single array, or a unit array +-- This is used in createClusterArg to both find the corresponding argument, +-- and to inform the type checker that the array has the correct type. +matchArgArraySingle :: Arg env a -> Arg env b -> Maybe (a :~: b) +matchArgArraySingle (ArgArray m1 (ArrayR _ (TupRsingle t1)) sh1 (TupRsingle buffer1)) (ArgArray m2 (ArrayR _ (TupRsingle t2)) sh2 (TupRsingle buffer2)) + | Just Refl <- matchModifier m1 m2 + , Just Refl <- matchVar buffer1 buffer2 + , Just Refl <- matchVars sh1 sh2 + , Just Refl <- matchScalarType t1 t2 = Just Refl +matchArgArraySingle (ArgArray m1 (ArrayR _ TupRunit) sh1 _) (ArgArray m2 (ArrayR _ TupRunit) sh2 _) + | Just Refl <- matchModifier m1 m2 + , Just Refl <- matchVars sh1 sh2 = Just Refl +matchArgArraySingle _ _ = Nothing + +matchModifier :: Modifier m1 -> Modifier m2 -> Maybe (m1 :~: m2) +matchModifier In In = Just Refl +matchModifier Out Out = Just Refl +matchModifier Mut Mut = Just Refl +matchModifier _ _ = Nothing + +sortAndExpandArgs :: LabelledArgsOp op env args -> Exists (LabelledArgsOp op env) +sortAndExpandArgs args = argsFromList $ singles ++ unitArraysDedup ++ dedup + where + (arrays, singles) = partition isNonUnitInOut $ argsToList args + expanded = arrays >>= \(Exists arg) -> expandArg arg + sorted = sortOn fst expanded + dedup = map (snd . head) $ groupBy (\a b -> fst a == fst b) sorted + unitArraysDedup = nubBy compareUnitArrays (arrays >>= \(Exists arg) -> expandUnitArg arg) + + isNonUnitInOut :: Exists (LabelledArgOp op env) -> Bool + isNonUnitInOut (Exists (LOp (ArgArray _ (ArrayR _ TupRunit) _ _) _ _)) = False + isNonUnitInOut (Exists (LOp (ArgArray m _ _ _) _ _)) = case m of + In -> True + Out -> True + _ -> False + isNonUnitInOut _ = False + + compareUnitArrays :: Exists (LabelledArgOp op env) -> Exists (LabelledArgOp op env) -> Bool + compareUnitArrays (Exists (LOp (ArgArray m _ sh _) _ _)) (Exists (LOp (ArgArray m' _ sh' _) _ _)) + | Just _ <- matchModifier m m' + , Just _ <- matchVars sh sh' + = True + compareUnitArrays _ _ = False + +expandUnitArg :: LabelledArgOp op env t -> [Exists (LabelledArgOp op env)] +expandUnitArg (LOp (ArgArray m (ArrayR shr (TupRpair t1 t2)) sh (TupRpair b1 b2)) (Arr (TupRpair l1 l2, TupRpair bs1 bs2, TupRpair u1 u2) shval) ba) + = expandUnitArg (LOp (ArgArray m (ArrayR shr t1) sh b1) (Arr (l1, bs1, u1) shval) ba) + ++ expandUnitArg (LOp (ArgArray m (ArrayR shr t2) sh b2) (Arr (l2, bs2, u2) shval) ba) +expandUnitArg arg@(LOp _ (Arr (TupRunit, _, _) _) _) = [Exists arg] +expandUnitArg _ = [] + +expandArg :: LabelledArgOp op env t -> [(EnvLabel, Exists (LabelledArgOp op env))] +expandArg (LOp (ArgArray m (ArrayR shr (TupRpair t1 t2)) sh (TupRpair b1 b2)) (Arr (TupRpair l1 l2, TupRpair bs1 bs2, TupRpair u1 u2) shvar) ba) + = expandArg (LOp (ArgArray m (ArrayR shr t1) sh b1) (Arr (l1, bs1, u1) shvar) ba) + ++ expandArg (LOp (ArgArray m (ArrayR shr t2) sh b2) (Arr (l2, bs2, u2) shvar) ba) +expandArg arg@(LOp _ (Arr (TupRsingle (C.Const l), _, _) _) _) + = [(l, Exists arg)] +expandArg (LOp _ (Arr (TupRunit, _, _) _) _) = [] +expandArg _ = internalError "Tuple mismatch with labels" + +instance ShrinkArg (BackendClusterArg op) => SLVOperation (Clustered op) where + slvOperation (Clustered cluster b) = Just $ ShrinkOperation $ \ff' a' a -> + case slvCluster cluster ff' a' a of + ShrunkOperation' cluster' args -> + ShrunkOperation (Clustered cluster' $ shrinkArgs ff' b) args + +instance SimplifyOperation (Clustered op) + -- Default implementation, where detectCopy always returns [] + -- Any copies should already have been detected before the program was fused. + +slvClusterArg :: forall f f' t. SubArgs f f' -> ClusterArg (FunToEnv f) t -> ClusterArg (FunToEnv f') t +slvClusterArg subArgs (ClusterArgSingle idx) = case reindexSlv subArgs idx of + ReindexSlvKeep idx' -> ClusterArgSingle idx' + _ -> internalError "Expected a preserved variable, as this is not an Out argument" +slvClusterArg subArgs (ClusterArgArray m shr tp buffers) = ClusterArgArray m shr tp $ go buffers + where + go :: ClusterArgBuffers (FunToEnv f) m sh e -> ClusterArgBuffers (FunToEnv f') m sh e + go (ClusterArgBuffersPair a1 a2) = go a1 `ClusterArgBuffersPair` go a2 + go (ClusterArgBuffersDead tp idx) = case reindexSlv subArgs idx of + ReindexSlvKeep idx' -> ClusterArgBuffersDead tp idx' + go (ClusterArgBuffersLive tp idx) = case reindexSlv subArgs idx of + ReindexSlvKeep idx' -> ClusterArgBuffersLive tp idx' + ReindexSlvDead idx' -> ClusterArgBuffersDead tp idx' + +reindexSlv :: SubArgs f f' -> Idx (FunToEnv f) t -> ReindexSlv f' t +reindexSlv (SubArgsDead _) ZeroIdx = ReindexSlvDead ZeroIdx +reindexSlv (SubArgsLive arg _) ZeroIdx = case arg of + SubArgKeep -> ReindexSlvKeep ZeroIdx + SubArgOut SubTupRkeep -> ReindexSlvKeep ZeroIdx + _ -> internalError "Expected preserved value" +reindexSlv (SubArgsDead s) (SuccIdx idx) = case reindexSlv s idx of + ReindexSlvKeep idx -> ReindexSlvKeep $ SuccIdx idx + ReindexSlvDead idx -> ReindexSlvDead $ SuccIdx idx +reindexSlv (SubArgsLive _ s) (SuccIdx idx) = case reindexSlv s idx of + ReindexSlvKeep idx -> ReindexSlvKeep $ SuccIdx idx + ReindexSlvDead idx -> ReindexSlvDead $ SuccIdx idx +reindexSlv SubArgsNil idx = case idx of {} + +data ReindexSlv f t where + ReindexSlvKeep :: Idx (FunToEnv f) t -> ReindexSlv f t + ReindexSlvDead :: Idx (FunToEnv f) (Var' sh) -> ReindexSlv f (Out sh e) + +slvCluster :: Cluster op f -> SubArgs f f' -> Args env' f' -> Args env f -> ShrunkOperation' (Cluster op) env' f' +slvCluster (SingleOp (Single op args) label) sub args' _ + = ShrunkOperation' (SingleOp (Single op $ mapArgs (slvClusterArg sub) args) label) args' + +slvCluster (Fused fusion left right) sub args1' args1 = splitslvstuff fusion sub args1' args1 $ + \f' lsub largs' largs rsub rargs' rargs -> case (slvCluster left lsub largs' largs, slvCluster right rsub rargs' rargs) of + (ShrunkOperation' lop largs'', ShrunkOperation' rop rargs'') -> + ShrunkOperation' (Fused f' lop rop) (both (\x _ -> outvar x) f' largs'' rargs'') + where + splitslvstuff :: Fusion l r a + -> SubArgs a a' + -> Args env' a' + -> Args env a + -> (forall l' r'. Fusion l' r' a' -> SubArgs l l' -> Args env' l' -> Args env l -> SubArgs r r' -> Args env' r' -> Args env r -> result) + -> result + splitslvstuff EmptyF SubArgsNil ArgsNil ArgsNil k = k EmptyF SubArgsNil ArgsNil ArgsNil SubArgsNil ArgsNil ArgsNil + splitslvstuff _ (SubArgsLive (SubArgOut SubTupRskip) _) _ _ _ = error "completely removed out arg using subtupr" --splitslvstuff f (SubArgsDead subs) args' args k + splitslvstuff f (SubArgsLive (SubArgOut SubTupRkeep) subs) args' args k = splitslvstuff f (SubArgsLive SubArgKeep subs) args' args k + splitslvstuff _ (SubArgsLive (SubArgOut SubTupRpair{}) _) (_:>:_) (_:>:_) _ = error "not SOA'd array" + splitslvstuff (Diagonal f) (SubArgsDead subs) args' (ArgArray _ r sh _:>:args) k = splitslvstuff (Vertical r f) (SubArgsLive SubArgKeep subs) args' (ArgVar (groundToExpVar (shapeType $ arrayRshape r) sh) :>:args) k + splitslvstuff (IntroO1 f) (SubArgsDead subs) (arg':>:args') (arg:>:args) k = splitslvstuff f subs args' args $ \f lsubs largs' largs rsubs rargs' rargs -> k (IntroL f) (SubArgsDead lsubs) (arg':>:largs') (arg:>:largs) rsubs rargs' rargs + splitslvstuff (IntroO2 f) (SubArgsDead subs) (arg':>:args') (arg:>:args) k = splitslvstuff f subs args' args $ \f lsubs largs' largs rsubs rargs' rargs -> k (IntroR f) lsubs largs' largs (SubArgsDead rsubs) (arg':>:rargs') (arg:>:rargs) + splitslvstuff (IntroL _) (SubArgsDead _) (_:>:_) (_:>:_) _ = error "out in IntroL/R" + splitslvstuff (IntroR _) (SubArgsDead _) (_:>:_) (_:>:_) _ = error "out in IntroL/R" + splitslvstuff (Vertical r f) (SubArgsLive SubArgKeep subs) (ArgVar arg':>:args') (ArgVar arg:>:args) k = splitslvstuff f subs args' args $ \f lsubs largs' largs rsubs rargs' rargs -> k (Vertical r f) (SubArgsLive SubArgKeep lsubs) (ArgArray Out r sh' buf :>:largs') (ArgArray Out r sh buf :>:largs) (SubArgsLive SubArgKeep rsubs) (ArgArray In r sh' buf :>:rargs') (ArgArray In r sh buf :>:rargs) + where + buf = error "fused away buffer" + sh = expToGroundVar arg + sh' = expToGroundVar arg' + splitslvstuff (Diagonal f) (SubArgsLive SubArgKeep subs) (arg'@(ArgArray Out r' sh' buf'):>:args') (arg@(ArgArray Out r sh buf):>:args) k = splitslvstuff f subs args' args $ \f lsubs largs' largs rsubs rargs' rargs -> k (Diagonal f) (SubArgsLive SubArgKeep lsubs) (arg':>:largs') (arg:>:largs) (SubArgsLive SubArgKeep rsubs) (ArgArray In r' sh' buf':>:rargs') (ArgArray In r sh buf:>:rargs) + splitslvstuff (Horizontal f) (SubArgsLive SubArgKeep subs) (arg':>:args') (arg:>:args) k = splitslvstuff f subs args' args $ \f lsubs largs' largs rsubs rargs' rargs -> k (Horizontal f) (SubArgsLive SubArgKeep lsubs) ( arg':>:largs') ( arg:>:largs) (SubArgsLive SubArgKeep rsubs) ( arg':>:rargs') ( arg:>:rargs) + splitslvstuff (IntroI1 f) (SubArgsLive SubArgKeep subs) (arg':>:args') (arg:>:args) k = splitslvstuff f subs args' args $ \f lsubs largs' largs rsubs rargs' rargs -> k (IntroI1 f) (SubArgsLive SubArgKeep lsubs) ( arg':>:largs') ( arg:>:largs) rsubs rargs' rargs + splitslvstuff (IntroI2 f) (SubArgsLive SubArgKeep subs) (arg':>:args') (arg:>:args) k = splitslvstuff f subs args' args $ \f lsubs largs' largs rsubs rargs' rargs -> k (IntroI2 f) lsubs largs' largs (SubArgsLive SubArgKeep rsubs) ( arg':>:rargs') ( arg:>:rargs) + splitslvstuff (IntroO1 f) (SubArgsLive SubArgKeep subs) (arg':>:args') (arg:>:args) k = splitslvstuff f subs args' args $ \f lsubs largs' largs rsubs rargs' rargs -> k (IntroO1 f) (SubArgsLive SubArgKeep lsubs) ( arg':>:largs') ( arg:>:largs) rsubs rargs' rargs + splitslvstuff (IntroO2 f) (SubArgsLive SubArgKeep subs) (arg':>:args') (arg:>:args) k = splitslvstuff f subs args' args $ \f lsubs largs' largs rsubs rargs' rargs -> k (IntroO2 f) lsubs largs' largs (SubArgsLive SubArgKeep rsubs) ( arg':>:rargs') ( arg:>:rargs) + splitslvstuff (IntroL f) (SubArgsLive SubArgKeep subs) (arg':>:args') (arg:>:args) k = splitslvstuff f subs args' args $ \f lsubs largs' largs rsubs rargs' rargs -> k (IntroL f) (SubArgsLive SubArgKeep lsubs) ( arg':>:largs') ( arg:>:largs) rsubs rargs' rargs + splitslvstuff (IntroR f) (SubArgsLive SubArgKeep subs) (arg':>:args') (arg:>:args) k = splitslvstuff f subs args' args $ \f lsubs largs' largs rsubs rargs' rargs -> k (IntroR f) lsubs largs' largs (SubArgsLive SubArgKeep rsubs) ( arg':>:rargs') ( arg:>:rargs) + +-- Variant of ShrunkOperation where f is not an existential +data ShrunkOperation' op env f where + ShrunkOperation' :: op f -> Args env f -> ShrunkOperation' op env f + +outvar :: Arg env (Out sh e) -> Arg env (Var' sh) +outvar (ArgArray Out (ArrayR shr _) sh _) = ArgVar $ groundToExpVar (shapeType shr) sh + +prjClusterArgs :: Args env args -> ClusterArgs (FunToEnv args) s -> Args env s +prjClusterArgs args = mapArgs (prjClusterArg args) + +prjClusterArg :: forall env args s. Args env args -> ClusterArg (FunToEnv args) s -> Arg env s +prjClusterArg args (ClusterArgSingle idx) = funToEnvPrj args idx +prjClusterArg args (ClusterArgArray (m :: Modifier m) (shr :: ShapeR sh) tp buffers') = case go buffers' of + (Just sh, buffers) -> ArgArray m (ArrayR shr tp) sh buffers + _ -> internalError "Arrays only consisting of units are not supported in fusion" + where + go :: ClusterArgBuffers (FunToEnv args) m sh e -> (Maybe (GroundVars env sh), GroundVars env (Buffers e)) + go (ClusterArgBuffersPair c1 c2) + | (sh1, buffers1) <- go c1 + , (sh2, buffers2) <- go c2 + = (sh1 <|> sh2, TupRpair buffers1 buffers2) + go (ClusterArgBuffersLive _ idx) + | ArgArray _ _ sh buffers <- funToEnvPrj args idx + = (Just sh, buffers) + go (ClusterArgBuffersDead _ idx) + | ArgVar sh <- funToEnvPrj args idx + = (Just $ mapTupR f sh, TupRsingle $ internalError "Cannot access array removed by strongly-live-variable analysis") + + f :: ExpVar env t -> GroundVar env t + f (Var tp idx) = Var (GroundRscalar tp) idx + +showSorted :: LabelledArgsOp op env args -> String +showSorted ArgsNil = "" +showSorted (a :>: args) = case a of + LOp (ArgArray m _ _ _) (Arr (_,bs,_) (_,sh,_)) _ -> show m <> "{" <> show (valsNodes bs <> valsNodes sh) <> "}" <> showSorted args + _ -> showSorted args + +data FlatCluster op env where + FlatCluster + :: ShapeR sh + -> ELeftHandSide sh () idxEnv + -> GroundVars env sh + -> TupR LoopDirection sh + -> TupR LocalBufferR local + -> GLeftHandSide local env env' + -> FlatOps op env' idxEnv + -> FlatCluster op env + +data LoopDirection t where + -- Any order; there are no dependencies between iterations of the loop + LoopAny :: LoopDirection Int + -- Either ascending or descending + LoopMonotone :: LoopDirection Int + -- from 0 to n - 1 + LoopAscending :: LoopDirection Int + -- from n - 1 to 0 + LoopDescending :: LoopDirection Int + +isDescending :: LoopDirection Int -> Bool +isDescending LoopDescending = True +isDescending _ = False + +-- TODO: Migrate all uses of isDescending to chooseDirection, as +-- Direction provides more type safety than Bool? +chooseDirection :: LoopDirection Int -> Direction +chooseDirection LoopDescending = RightToLeft +chooseDirection _ = LeftToRight + +-- The representation of a buffer that is either fused away (only used within +-- a cluster) or not used at all. +data LocalBufferR b where + LocalBufferR + :: ScalarType t -> LoopDepth -> LocalBufferR (Buffer t) + +data FlatOps op env idxEnv where + FlatOpsNil :: FlatOps op env idxEnv + + FlatOpsBind + :: LoopDepth + -> ELeftHandSide t idxEnv idxEnv' + -> OpenExp idxEnv env t + -> FlatOps op env idxEnv' + -> FlatOps op env idxEnv + + FlatOpsOp + :: FlatOp op env idxEnv + -> FlatOps op env idxEnv + -> FlatOps op env idxEnv + +data FlatOp op env idxEnv where + FlatOp + :: op args + -> Args env args + -> IdxArgs idxEnv args + -> FlatOp op env idxEnv + +-- The nesting depth of something in nested loops. +-- In absence of backpermutes, this is equal to the dimensionality. +-- E.g., a zero-dimensional index or array is in nesting depth zero, +-- and a two-dimensional index or array is in nesting depth two. +type LoopDepth = Int + +instance Sink (FlatOp op) where + weaken k (FlatOp op args idxEnv) = FlatOp op (mapArgs (weaken k) args) idxEnv + +type IdxArgs idxEnv = PreArgs (IdxArg idxEnv) + +data IdxArg idxEnv t where + IdxArgIdx :: LoopDepth -> ExpVars idxEnv sh -> IdxArg idxEnv (m sh t) + + IdxArgNone :: IdxArg idxEnv t + +instance Sink IdxArg where + weaken k (IdxArgIdx depth idxs) = IdxArgIdx depth $ mapTupR (weaken k) idxs + weaken _ IdxArgNone = IdxArgNone + +data ToFlatCluster op f where + ToFlatCluster + :: TypeR a -- Fused away and dead output arrays + -- The OutR Vars should not be needed for the OutShapes, but we include them + -- anyway to use the same index transformation for the OutShapes and the FlatOps. + -> (forall env'. Args env' f -> Vars GroundR env' (Buffers a) -> [FlatOp op env' ()]) + -> ToFlatCluster op f + +data ToFlatArgs f t where + ToFlatArgs + :: TypeR a -- Fused away and dead output arrays + -> (forall env'. Args env' f -> Vars GroundR env' (Buffers a) -> Args env' t) + -> ToFlatArgs f t + +data ToFlatArg f t where + ToFlatArg + :: TypeR a -- Fused away and dead output arrays + -> (forall env'. Args env' f -> Vars GroundR env' (Buffers a) -> Arg env' t) + -> ToFlatArg f t + +data ToFlatArgBuffers f t where + ToFlatArgBuffers + :: TypeR a -- Fused away and dead output arrays + -> (forall env'. Args env' f -> Vars GroundR env' (Buffers a) -> GroundVars env' (Buffers t)) + -> ToFlatArgBuffers f t + +data ToFlatFusion fun funL funR where + ToFlatFusion + :: TypeR a -- Fused away and dead output arrays + -> (forall env'. Args env' fun -> Vars GroundR env' (Buffers a) -> Args env' funL) + -> (forall env'. Args env' fun -> Vars GroundR env' (Buffers a) -> Args env' funR) + -> ToFlatFusion fun funL funR + +toFlatClustered :: SetOpIndices op => Clustered op f -> Args env f -> FlatCluster op env +toFlatClustered (Clustered cluster _) = toFlatCluster cluster + +toFlatCluster :: SetOpIndices op => Cluster op f -> Args env f -> FlatCluster op env +toFlatCluster cluster args + | ToFlatCluster tp ops <- toFlatCluster' cluster + , DeclareVars lhs k values <- declareVars $ buffersR tp + , values' <- values weakenId + = flatOpsSetIndices tp lhs values' $ ops (mapArgs (weaken k) args) values' + +toFlatCluster' :: forall op f. Cluster op f -> ToFlatCluster op f +toFlatCluster' (Fused fusion l r) + | ToFlatCluster tp1 l' <- toFlatCluster' l + , ToFlatFusion tp2 leftArgs rightArgs <- travFusion fusion + , ToFlatCluster tp3 r' <- toFlatCluster' r + = case (tp1, tp2, tp3) of + (_, TupRunit, TupRunit) -> + ToFlatCluster tp1 + $ \args value -> + l' (leftArgs args TupRunit) value ++ r' (rightArgs args TupRunit) TupRunit + (TupRunit, _, TupRunit) -> + ToFlatCluster tp2 + $ \args value -> + l' (leftArgs args value) TupRunit ++ r' (rightArgs args value) TupRunit + (TupRunit, TupRunit, _) -> + ToFlatCluster tp3 + $ \args value -> + l' (leftArgs args TupRunit) TupRunit ++ r' (rightArgs args TupRunit) value + -- This last case is the generic case. Only having this case would be + -- sufficient, but the prior cases just prevent that we put many many + -- TupRunits in the type. + _ -> + ToFlatCluster + (tp1 `TupRpair` tp2 `TupRpair` tp3) + $ \args value -> case value of + v1 `TupRpair` v2 `TupRpair` v3 -> + l' (leftArgs args v2) v1 ++ r' (rightArgs args v2) v3 + _ -> internalError "Pair impossible" + where + travFusion :: Fusion largs rargs args -> ToFlatFusion args largs rargs + travFusion EmptyF = ToFlatFusion TupRunit (\_ _ -> ArgsNil) (\_ _ -> ArgsNil) + travFusion (IntroI1 f) = travFusion (IntroL f) + travFusion (IntroI2 f) = travFusion (IntroR f) + travFusion (IntroO1 f) = travFusion (IntroL f) + travFusion (IntroO2 f) = travFusion (IntroR f) + travFusion (IntroL f) + | ToFlatFusion tp l r <- travFusion f + = ToFlatFusion tp + (\(arg :>: args) value -> arg :>: l args value) + (\(_ :>: args) value -> r args value) + travFusion (IntroR f) + | ToFlatFusion tp l r <- travFusion f + = ToFlatFusion tp + (\(_ :>: args) value -> l args value) + (\(arg :>: args) value -> arg :>: r args value) + travFusion (Horizontal f) + | ToFlatFusion tp l r <- travFusion f + = ToFlatFusion tp + (\(arg :>: args) value -> arg :>: l args value) + (\(arg :>: args) value -> arg :>: r args value) + travFusion (Diagonal f) + | ToFlatFusion tp l r <- travFusion f + = ToFlatFusion tp + (\(arg :>: args) value -> arg :>: l args value) + (\(arg :>: args) value -> case arg of + ArgArray _ repr sh bs -> ArgArray In repr sh bs :>: r args value + ) + travFusion (Vertical repr@(ArrayR _ tp1) f) + | ToFlatFusion tp2 l r <- travFusion f + = ToFlatFusion (TupRpair tp1 tp2) + (\(arg :>: args) value -> if + | ArgVar sh <- arg + , TupRpair value1 value2 <- value -> + ArgArray Out repr + (mapTupR (\(Var tp idx) -> Var (GroundRscalar tp) idx) sh) + value1 + :>: l args value2 + | otherwise -> internalError "Pair impossible" + ) + (\(arg :>: args) value -> if + | ArgVar sh <- arg + , TupRpair value1 value2 <- value -> + ArgArray In repr + (mapTupR (\(Var tp idx) -> Var (GroundRscalar tp) idx) sh) + value1 + :>: r args value2 + | otherwise -> internalError "Pair impossible" + ) +toFlatCluster' (SingleOp (Single op opArgs) _) + | ToFlatArgs tp args' <- travArgs opArgs + = ToFlatCluster tp $ \args values -> + let args'' = args' args values + in [FlatOp op args'' $ mapArgs (const IdxArgNone) args''] + where + travArgs :: ClusterArgs (FunToEnv f) f' -> ToFlatArgs f f' + travArgs ArgsNil = ToFlatArgs TupRunit $ \_ _ -> ArgsNil + travArgs (a :>: as) = case (travArg a, travArgs as) of + (ToFlatArg TupRunit a', ToFlatArgs tp as') -> + ToFlatArgs tp $ \args values -> a' args TupRunit :>: as' args values + (ToFlatArg tp a', ToFlatArgs TupRunit as') -> + ToFlatArgs tp $ \args values -> a' args values :>: as' args TupRunit + (ToFlatArg t1 a', ToFlatArgs t2 as') -> + ToFlatArgs (TupRpair t1 t2) + $ \args values -> case values of + TupRpair v1 v2 -> a' args v1 :>: as' args v2 + TupRsingle _ -> internalError "Pair impossible" + + travArg :: ClusterArg (FunToEnv f) t -> ToFlatArg f t + travArg (ClusterArgSingle idx) = ToFlatArg TupRunit $ \args _ -> funToEnvPrj args idx + travArg (ClusterArgArray m shr tp bs) + | ToFlatArgBuffers fusedR bs' <- travBuffers bs + = ToFlatArg fusedR $ \args values -> + ArgArray + m + (ArrayR shr tp) + (findSize args bs) + (bs' args values) + + findSize :: Args env f -> ClusterArgBuffers (FunToEnv f) m sh t -> GroundVars env sh + findSize args (ClusterArgBuffersPair b1 _) = findSize args b1 + findSize args (ClusterArgBuffersLive _ idx) = case funToEnvPrj args idx of + ArgArray _ _ sh _ -> sh + findSize args (ClusterArgBuffersDead _ idx) = case funToEnvPrj args idx of + ArgVar vars -> mapTupR (\(Var tp idx) -> Var (GroundRscalar tp) idx) vars + + travBuffers :: ClusterArgBuffers (FunToEnv f) m sh t -> ToFlatArgBuffers f t + travBuffers (ClusterArgBuffersPair b1 b2) = case (travBuffers b1, travBuffers b2) of + (ToFlatArgBuffers TupRunit f1, ToFlatArgBuffers tp f2) -> + ToFlatArgBuffers tp $ \args values -> f1 args TupRunit `TupRpair` f2 args values + (ToFlatArgBuffers tp f1, ToFlatArgBuffers TupRunit f2) -> + ToFlatArgBuffers tp $ \args values -> f1 args values `TupRpair` f2 args TupRunit + (ToFlatArgBuffers t1 f1, ToFlatArgBuffers t2 f2) -> + ToFlatArgBuffers (TupRpair t1 t2) $ \args values -> case values of + TupRpair v1 v2 -> f1 args v1 `TupRpair` f2 args v2 + TupRsingle _ -> internalError "Pair impossible" + travBuffers (ClusterArgBuffersLive _ idx) = + ToFlatArgBuffers TupRunit $ \args _ -> case funToEnvPrj args idx of + ArgArray _ _ _ bs -> bs + travBuffers (ClusterArgBuffersDead tp _) = + ToFlatArgBuffers + tp + $ \_ value -> value + +class SetOpIndices op where + setOpIndices + -- Get the index variable for the given dimension / loop depth. + -- For instance, when calling this function with 0 as argument, it will return + -- the iteration variable for the outer loop. Note that the contents of that + -- loop is in LoopDepth 1, not zero. + -- If there is no loop nesting for that depth, this function returns Nothing. + -- In this case, and only in this case, setOpIndices may return Nothing. + :: (LoopDepth -> Maybe (Idx idxEnv Int)) + -> op args + -> Args env args + -- IdxArgs where only the Out positions are filled + -> IdxArgs idxEnv args + -- IdxArgs where both the In and Out positions are filled. + -- Returns Nothing if (and only if) there were not enough iteration variables, + -- i.e. when the function was called with too low dimensionality + -> Maybe (Either (IsBackpermute args) (IdxArgs idxEnv args)) + + getOpLoopDirections + :: op args + -> Args env args + -> IdxArgs idxEnv args + -> [(Idx idxEnv Int, LoopDirection Int)] + getOpLoopDirections _ _ _ = [] + +data IsBackpermute args where + IsBackpermute :: IsBackpermute (Fun' (sh' -> sh) -> In sh t -> Out sh' t -> ()) + +flatOpsSetIndices + :: forall op env env' t. + SetOpIndices op + => TypeR t + -> GLeftHandSide (Buffers t) env env' + -> GroundVars env' (Buffers t) + -> [FlatOp op env' ()] + -> FlatCluster op env +flatOpsSetIndices fusedR fusedLHS fusedVars ops = go ShapeRz + where + go :: ShapeR sh -> FlatCluster op env + go shr + | DeclareVars lhs _ value <- declareVars $ shapeType shr + , indices <- map (\(Exists idx) -> expectIntVar idx) $ flattenTupR $ value weakenId + , Just (FlatOpsSetIndices _ k bnd ops') <- flatTrySetIndices indices ops + , ops'' <- ops' weakenId + , directions <- getOpsLoopDirections ops'' + = FlatCluster + shr + lhs + (mapTupR + (weaken $ Weaken $ fromMaybe (internalError "Iteration size refers to local environment, which should only contain buffers, not sizes") . strengthenWithLHS fusedLHS) + $ findIterationSizes ops'' shr $ value k) + (mapTupR + (\(Var tp idx) -> case prjPartial idx directions of + Just dir -> dir + Nothing + | SingleScalarType (NumSingleType (IntegralNumType TypeInt)) <- tp -> LoopAny + | otherwise -> internalError "Expected Int variable" + ) + $ value k) + (localBuffersR ops'' fusedR fusedVars) + fusedLHS + $ bnd ops'' + | otherwise + = go (ShapeRsnoc shr) + + expectIntVar :: ExpVar e a -> Idx e Int + expectIntVar (Var tp idx) + | Just Refl <- matchScalarType tp scalarTypeInt + = idx + | otherwise + = internalError "All scalar types in shapeType should be Int" + + localBuffersR :: forall e idxEnv. FlatOps op env' idxEnv -> TypeR e -> GroundVars env' (Buffers e) -> TupR LocalBufferR (Buffers e) + localBuffersR _ TupRunit _ = TupRunit + localBuffersR ops'' (TupRsingle tp) (TupRsingle var) + | Refl <- reprIsSingle @ScalarType @e @Buffer tp + = TupRsingle $ LocalBufferR tp $ fromMaybe 0 $ findBufferDepth ops'' $ varIdx var + localBuffersR ops'' (TupRpair t1 t2) (TupRpair v1 v2) + = localBuffersR ops'' t1 v1 `TupRpair` localBuffersR ops'' t2 v2 + localBuffersR _ _ _ = internalError "Tuple mismatch" + +data FlatOpsSetIndices op env idxEnv where + FlatOpsSetIndices + :: PartialEnv (BufferIdx idxEnv') env + -> idxEnv :> idxEnv' + -> (FlatOps op env idxEnv' -> FlatOps op env idxEnv) + -> (forall idxEnv''. idxEnv' :> idxEnv'' -> FlatOps op env idxEnv'') + -> FlatOpsSetIndices op env idxEnv + +flatTrySetIndices + :: SetOpIndices op + => [Idx idxEnv Int] -- Indices for the different dimensions + -> [FlatOp op env ()] + -> Maybe (FlatOpsSetIndices op env idxEnv) +flatTrySetIndices _ [] = + Just $ FlatOpsSetIndices PEnd weakenId id $ const FlatOpsNil +flatTrySetIndices indices (FlatOp op args _ : ops) + | Just (FlatOpsSetIndices env k bnd ops') <- flatTrySetIndices indices ops + , Just idxArgs1 <- setOutArgsIndices (map (weaken k) indices) env args + , Just e <- setOpIndices + (fmap (weaken k) . (!?) indices) + op + args + idxArgs1 + = case e of + Left IsBackpermute + | ArgFun (Lam lhs (Body idxTransform)) + :>: _ <- args + -- Find the indices of the output + , _ :>: _ :>: IdxArgIdx depth idxs :>: _ <- idxArgs1 + -- Weaken lhs into the right environment + , Exists lhs' <- rebuildLHS lhs + -- Declare new variables for the input indices + -- (the output of the index transformation) + , DeclareVars lhs1 k1 values1 <- declareVars $ expType idxTransform + , expr <- Let lhs' (expVars idxs) (weakenE (sinkWithLHS lhs lhs' weakenEmpty) idxTransform) + -- Construct the final IdxArgs, where the input is annotated with the correct indices + , idxArgs2 <- IdxArgNone + :>: IdxArgIdx depth (values1 weakenId) + :>: IdxArgIdx depth (mapTupR (weaken k1) idxs) + :>: ArgsNil + -- Weaken all bindings in the environment + , env' <- mapPartialEnv (weaken k1) env + -- Store newly introduced index variables in environment + , env'' <- propagateArrayIndices env' args idxArgs2 + -> Just $ FlatOpsSetIndices + env'' + (k1 .> k) + (bnd . FlatOpsBind depth lhs1 expr) + $ \k0 -> FlatOpsOp + (FlatOp op args $ mapArgs (weaken k0) idxArgs2) + (ops' (k0 .> k1)) + | otherwise + -> internalError "No index available for output of backpermute" + Right idxArgs2 + | env' <- propagateArrayIndices env args idxArgs2 + -> Just $ FlatOpsSetIndices env' k bnd + $ \k0 -> FlatOpsOp (FlatOp op args $ mapArgs (weaken k0) idxArgs2) (ops' k0) + | otherwise + = Nothing + +setOutArgsIndices + :: forall env idxEnv' args. + [Idx idxEnv' Int] + -> PartialEnv (BufferIdx idxEnv') env + -> Args env args + -> Maybe (IdxArgs idxEnv' args) +setOutArgsIndices indices env = go + where + go :: Args env t -> Maybe (IdxArgs idxEnv' t) + go (ArgArray Out (ArrayR shr _) _ bs :>: as) + | Just as' <- go as = case findIndices shr bs of + Just (depth, idxs) -> Just $ IdxArgIdx depth idxs :>: as' + Nothing -> case findDefault (rank shr) shr of + Just idxs -> Just $ IdxArgIdx (rank shr) idxs :>: as' + Nothing -> Nothing + | otherwise + = Nothing + go (_ :>: as) = (IdxArgNone :>:) <$> go as + go ArgsNil = Just ArgsNil + + findIndices :: ShapeR sh -> GroundVars env b -> Maybe (LoopDepth, ExpVars idxEnv' sh) + findIndices shr (TupRsingle (Var _ idx)) + | Just (BufferIdx depth idxs) <- prjPartial idx env + = case matchTypeR (shapeType shr) (varsType idxs) of + Just Refl -> Just (depth, idxs) + Nothing -> internalError "Buffer is fused in cluster with different ranks. This should not be fused, is fusion missing a constraint?" + findIndices shr (TupRpair v1 v2) = findIndices shr v1 <|> findIndices shr v2 + findIndices _ _ = Nothing + + findDefault :: Int -> ShapeR sh -> Maybe (ExpVars idxEnv' sh) + findDefault _ ShapeRz = Just TupRunit + findDefault rank' (ShapeRsnoc shr) + | Just idx <- indices !? (rank' - 1) + , Just idxs <- findDefault (rank' - 1) shr + = Just (idxs `TupRpair` TupRsingle (Var scalarTypeInt idx)) + | otherwise = Nothing + +propagateArrayIndices + :: forall env idxEnv' args. + PartialEnv (BufferIdx idxEnv') env + -> Args env args + -> IdxArgs idxEnv' args + -> PartialEnv (BufferIdx idxEnv') env +propagateArrayIndices env (a :>: as) (i :>: is) + | IdxArgIdx depth idxs <- i + , ArgArray In _ _ bs <- a + = go env' depth idxs bs + | otherwise + = env' + where + env' = propagateArrayIndices env as is + go + :: PartialEnv (BufferIdx idxEnv') env + -> LoopDepth + -> ExpVars idxEnv' sh + -> GroundVars env t + -> PartialEnv (BufferIdx idxEnv') env + go env'' depth idxs (TupRsingle (Var _ buffer)) + = partialUpdate (BufferIdx depth idxs) buffer env'' + go env'' depth idxs (TupRpair v1 v2) + = go (go env'' depth idxs v1) depth idxs v2 + go env'' _ _ TupRunit = env'' +propagateArrayIndices env ArgsNil ArgsNil = env + +data BufferIdx idxEnv t where + -- Only used for Buffer values, but to encode that we need some more type information + -- and we don't really benefit from this guarantee. + BufferIdx :: LoopDepth -> ExpVars idxEnv sh -> BufferIdx idxEnv t + +instance Sink BufferIdx where + weaken k (BufferIdx depth vars) = BufferIdx depth $ mapTupR (weaken k) vars + +findIterationSizes + :: FlatOps op env idxEnv + -> ShapeR sh + -> ExpVars idxEnv sh + -> GroundVars env sh +findIterationSizes _ ShapeRz _ = TupRunit +findIterationSizes ops (ShapeRsnoc shr) (vars `TupRpair` TupRsingle var) + | Just sz <- findIterationSize ops (varIdx var) + , szs <- findIterationSizes ops shr vars + = szs `TupRpair` TupRsingle (Var (GroundRscalar scalarTypeInt) sz) + | otherwise + = internalError "Could not find iteration size. Is an iteration index not used?" +findIterationSizes _ _ _ = internalError "Tuple mismatch" + +findIterationSize + :: forall op env idxEnv. + FlatOps op env idxEnv + -> Idx idxEnv Int + -> Maybe (Idx env Int) +findIterationSize FlatOpsNil _ = Nothing +findIterationSize (FlatOpsBind _ lhs _ ops) idx = + findIterationSize ops $ weaken (weakenWithLHS lhs) idx +findIterationSize (FlatOpsOp (FlatOp _ args idxArgs) ops) idx + = findInArgs args idxArgs `expectSame` findIterationSize ops idx + where + findInArgs :: Args env f -> IdxArgs idxEnv f -> Maybe (Idx env Int) + findInArgs (a :>: as) (i :>: is) = findInArg a i `expectSame` findInArgs as is + findInArgs ArgsNil ArgsNil = Nothing + + findInArg :: Arg env t -> IdxArg idxEnv t -> Maybe (Idx env Int) + findInArg (ArgArray _ _ sh _) (IdxArgIdx _ idxs) = findInVars sh idxs + findInArg _ _ = Nothing + + findInVars :: GroundVars env sh -> ExpVars idxEnv sh -> Maybe (Idx env Int) + findInVars (TupRpair v1 v2) (TupRpair i1 i2) = findInVars v1 i1 `expectSame` findInVars v2 i2 + findInVars TupRunit TupRunit = Nothing + findInVars (TupRsingle v) (TupRsingle i) + | Just Refl <- matchIdx (varIdx i) idx = Just $ varIdx v + | otherwise = Nothing + findInVars _ _ = internalError "Tuple mismatch" + + expectSame :: Maybe (Idx env a) -> Maybe (Idx env a) -> Maybe (Idx env a) + expectSame Nothing m = m + expectSame m Nothing = m + expectSame (Just a) (Just b) + | a == b = Just a + | otherwise = internalError "Mismatch in iteration size: the iteration sizes of different operations in a cluster are different" + +(!?) :: [a] -> Int -> Maybe a +[] !? _ = Nothing +(x : _) !? 0 = Just x +(_ : xs) !? i + | i < 0 = Nothing + | otherwise = xs !? (i - 1) + +findBufferDepth + :: forall op env idxEnv e. + FlatOps op env idxEnv + -> Idx env (Buffer e) + -> Maybe LoopDepth +findBufferDepth FlatOpsNil _ = Nothing +findBufferDepth (FlatOpsBind _ _ _ ops) buffer = findBufferDepth ops buffer +findBufferDepth (FlatOpsOp (FlatOp _ args idxArgs) ops) buffer + = findInArgs args idxArgs `expectSame` findBufferDepth ops buffer + where + findInArgs :: Args env f -> IdxArgs idxEnv f -> Maybe LoopDepth + findInArgs (a :>: as) (i :>: is) = findInArg a i `expectSame` findInArgs as is + findInArgs ArgsNil ArgsNil = Nothing + + findInArg :: Arg env t -> IdxArg idxEnv t -> Maybe LoopDepth + findInArg (ArgArray _ _ _ buffers) (IdxArgIdx d _) + | idxElem buffers = Just d + findInArg _ _ = Nothing + + idxElem :: GroundVars env b -> Bool + idxElem TupRunit = False + idxElem (TupRsingle var) = isJust $ matchIdx buffer $ varIdx var + idxElem (TupRpair v1 v2) = idxElem v1 || idxElem v2 + + expectSame :: Maybe LoopDepth -> Maybe LoopDepth -> Maybe LoopDepth + expectSame Nothing m = m + expectSame m Nothing = m + expectSame (Just a) (Just b) + | a == b = Just a + | otherwise = internalError "Mismatch in loop depth: the loop depth of different operations on the same (fused away) buffer in a cluster are different" + +getOpsLoopDirections :: forall op env idxEnv. SetOpIndices op => FlatOps op env idxEnv -> PartialEnv LoopDirection idxEnv +getOpsLoopDirections flatOps = partialEnvFromList join $ map (uncurry EnvBinding) $ go flatOps + where + join :: LoopDirection t -> LoopDirection t -> LoopDirection t + join LoopAny d = d + join d LoopAny = d + join LoopMonotone d = d + join d LoopMonotone = d + join LoopAscending LoopAscending = LoopAscending + join LoopDescending LoopDescending = LoopDescending + join _ _ = internalError "Illegal cluster: an ascending (left-to-right) operation like scanl was fused with a descending (right-to-left) operation like scanr" + + go + :: FlatOps op env idxEnv + -> [(Idx idxEnv Int, LoopDirection Int)] + go (FlatOpsBind{}) = internalError "Did not expect a Bind here" + go (FlatOpsOp (FlatOp op args idxArgs) ops) + = getOpLoopDirections op args idxArgs ++ go ops + go FlatOpsNil = [] + +-- Returns the number of outermost dimensions (loops) +-- that have order LoopAny +flatClusterIndependentLoopDepth :: FlatCluster op env -> LoopDepth +flatClusterIndependentLoopDepth (FlatCluster _ _ _ dirs _ _ _) + = length $ takeWhile isAny $ flattenTupR dirs + where + isAny (Exists LoopAny) = True + isAny _ = False + +flatClusterLoopDepth :: FlatCluster op env -> LoopDepth +flatClusterLoopDepth (FlatCluster shr _ _ _ _ _ _) = rank shr diff --git a/src/Data/Array/Accelerate/AST/Schedule.hs b/src/Data/Array/Accelerate/AST/Schedule.hs new file mode 100644 index 000000000..bccea987c --- /dev/null +++ b/src/Data/Array/Accelerate/AST/Schedule.hs @@ -0,0 +1,147 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | +-- Module : Data.Array.Accelerate.AST.Schedule +-- Copyright : [2012..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- +-- This module defines the interface for schedules, to be implemented +-- in other modules. +-- + +module Data.Array.Accelerate.AST.Schedule ( + IsSchedule(..), + convertSchedule, + Scheduled, + reprIsBody, + IOFun, FullIOFun, FullIOFun', + flattenIOFun, unsafePerformIOFun, + generateKernelNameAndDescription +) where + +import Data.Array.Accelerate.AST.Partitioned +import Data.Array.Accelerate.AST.Kernel +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.Type +import Data.Array.Accelerate.Representation.Ground +import Data.Array.Accelerate.Representation.Type +import Data.Typeable ( (:~:)(..) ) +import Data.List (sortOn, group) +import Data.Ord +import Control.Monad +import System.IO.Unsafe + +class IsSchedule sched where + -- 'a' is a ground type (ie, can be represented using GroundR) + type ScheduleInput sched a + type ScheduleOutput sched a + + rnfSchedule :: IsKernel kernel => sched kernel env t -> () + + convertScheduleFun + :: IsKernel kernel + => PartitionedAfun (KernelOperation kernel) () t -> sched kernel () (Scheduled sched t) + + callScheduledFun :: GFunctionR t -> IOFun (Scheduled sched t) -> FullIOFun t + +convertSchedule + :: forall sched kernel t. + (IsSchedule sched, IsKernel kernel) + => PartitionedAcc (KernelOperation kernel) () t -> sched kernel () (ScheduleOutput sched t -> ()) +convertSchedule acc + | Refl <- reprIsBody @sched $ groundsR acc = convertScheduleFun (Abody acc) + +type family Scheduled sched a where + Scheduled sched (a -> b) = ScheduleInput sched a -> Scheduled sched b + Scheduled sched a = ScheduleOutput sched a -> () + +reprIsBody :: forall sched t. GroundsR t -> (Scheduled sched t, IOFun t, FullIOFun t) :~: (ScheduleOutput sched t -> (), IO t, IO t) +reprIsBody TupRunit = Refl +reprIsBody TupRpair{} = Refl +reprIsBody (TupRsingle (GroundRbuffer _)) = Refl +reprIsBody (TupRsingle (GroundRscalar (VectorScalarType _))) = Refl +reprIsBody (TupRsingle (GroundRscalar (SingleScalarType (NumSingleType t)))) = case t of + IntegralNumType TypeInt -> Refl + IntegralNumType TypeInt8 -> Refl + IntegralNumType TypeInt16 -> Refl + IntegralNumType TypeInt32 -> Refl + IntegralNumType TypeInt64 -> Refl + IntegralNumType TypeWord -> Refl + IntegralNumType TypeWord8 -> Refl + IntegralNumType TypeWord16 -> Refl + IntegralNumType TypeWord32 -> Refl + IntegralNumType TypeWord64 -> Refl + FloatingNumType TypeHalf -> Refl + FloatingNumType TypeFloat -> Refl + FloatingNumType TypeDouble -> Refl + +type family IOFun f where + IOFun (a -> b) = a -> IOFun b + IOFun t = IO t + +type FullIOFun f = IO (FullIOFun' f) + +type family FullIOFun' f where + FullIOFun' (a -> b) = a -> FullIOFun b + FullIOFun' t = t + +flattenIOFun :: GFunctionR f -> FullIOFun f -> IOFun f +flattenIOFun (GFunctionRlam _ funR) fun = \a -> flattenIOFun funR (join $ fun <*> pure a) +flattenIOFun (GFunctionRbody tp) fun + | Refl <- reprIsBody tp = fun + +unsafePerformIOFun :: GFunctionR f -> IOFun f -> f +unsafePerformIOFun (GFunctionRlam _ funR) fun = unsafePerformIOFun funR . fun +unsafePerformIOFun (GFunctionRbody tp) body + | Refl <- reprIsBody tp = unsafePerformIO body + +generateKernelNameAndDescription + -- Function that gives the priority, singular name and plural name of an + -- operation. The operations with highest priority are used in the name. + :: (forall s. op s -> (Int, String, String)) + -> Clustered op t + -- Returns a function name, a detailed description and a brief description + -> (String, String, String) +generateKernelNameAndDescription f cluster = + ( formatList False "_" "_" (take 2 sorted) ++ (if length sorted > 2 then "_etc" else "") + , if trivial then "" else "Cluster with " ++ + formatList True ", then " (if length grouped == 2 then " and then " else " and finally ") (groupAndCount ops) + , if trivial then "" else "Cluster with " ++ + formatList True ", " " and " sorted + ) + where + trivial + | [_] <- ops = True + | otherwise = False + + ops = map (\(Exists op) -> f op) $ flattenClustered cluster + grouped = groupAndCount ops + sorted = groupAndCount $ sortOn Down $ ops + + groupAndCount :: Eq a => [a] -> [(a, Int)] + groupAndCount = map g . group + where + g [] = error "impossible" + g as@(a:_) = (a, length as) + + formatList :: Bool -> String -> String -> [((Int, String, String), Int)] -> String + formatList includeNumber separator finalSeparator list = case list of + [] -> "" + [a] -> formatItem includeNumber a + [a, b] -> formatItem includeNumber a ++ finalSeparator ++ formatItem includeNumber b + (a:as) -> formatItem includeNumber a ++ separator ++ formatList includeNumber separator finalSeparator as + + formatItem :: Bool -> ((Int, String, String), Int) -> String + formatItem includeNumber ((_, singular, plural), count) + | includeNumber = show count ++ " " ++ name + | otherwise = name + where + name + | count == 1 = singular + | otherwise = plural diff --git a/src/Data/Array/Accelerate/AST/Schedule/Sequential.hs b/src/Data/Array/Accelerate/AST/Schedule/Sequential.hs new file mode 100644 index 000000000..47340948b --- /dev/null +++ b/src/Data/Array/Accelerate/AST/Schedule/Sequential.hs @@ -0,0 +1,205 @@ +{-# LANGUAGE TypeFamilies #-} +{-# OPTIONS_HADDOCK hide #-} +-- | +-- Module : Data.Array.Accelerate.AST.Schedule.Sequential +-- Copyright : [2008..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.AST.Schedule.Sequential ( + SArg(..), SArgs, + module Operation, + Cluster, + SequentialSchedule(..), + SeqSchedule(..), + SeqScheduleFun(..) +) where + +import Data.Array.Accelerate.AST.Exp +import Data.Array.Accelerate.AST.Var +import Data.Array.Accelerate.Type +import Data.Array.Accelerate.Representation.Array +import Data.Array.Accelerate.Representation.Shape +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.AST.LeftHandSide +import qualified Data.Array.Accelerate.AST.Operation as Operation +import Data.Array.Accelerate.AST.Partitioned hiding (PreOpenAcc(..), PartitionedAcc, PartitionedAfun) +import qualified Data.Array.Accelerate.AST.Partitioned as Partition +import Data.Array.Accelerate.AST.Kernel +import Data.Array.Accelerate.AST.Schedule +import Data.Array.Accelerate.AST.Schedule.Uniform ( SArg(..), SArgs ) +import Data.Array.Accelerate.Trafo.Schedule.Uniform ( compileKernel', CompiledKernel(..), rnfSArgs ) +import Data.Array.Accelerate.AST.Execute +import Control.Concurrent.MVar +import Data.Typeable ( (:~:)(..) ) +import Data.Kind (Type) +import Data.Text (Text) +import Control.DeepSeq (NFData (rnf)) + +-- Generic schedule for a uniform memory space and uniform scheduling, +-- without task parallelism. +data SequentialSchedule kernel env t where + SequentialLam + :: GLeftHandSide s env env' + -> SequentialSchedule kernel env' t + -> SequentialSchedule kernel env (s -> t) + + SequentialBody + :: SeqSchedule kernel env t + -> SequentialSchedule kernel env (MVar t -> ()) + +data SeqSchedule (kernel :: (Type -> Type)) env t where + -- | Executes a kernel. Such a kernel does not return a + -- value, the effects of the execution are only visible by the mutation of + -- buffers which were annotated with either 'Mut' or 'Out'. + -- Provides the operation arguments from the environment. + -- + Exec :: KernelMetadata kernel f + -> KernelFun kernel args + -> SArgs env args + -> SeqSchedule kernel env () + + -- | Returns the values of the given variables. + -- + Return :: GroundVars env a + -> SeqSchedule kernel env a + + -- | Evaluates the expression and returns its value. + -- + Compute :: Exp env t + -> SeqSchedule kernel env t + + -- | Local binding of ground values. + -- As it is common in this intermediate representation to evaluate some program + -- resulting in unit, for instance when execution some operation, and then + -- and then do some other code, we write 'a; b' to denote 'let () = a in b'. + Alet :: GLeftHandSide bnd env env' + -> Uniquenesses bnd + -> SeqSchedule kernel env bnd + -> SeqSchedule kernel env' t + -> SeqSchedule kernel env t + + -- | Allocates a new buffer of the given size. + -- + Alloc :: ShapeR sh + -> ScalarType e + -> ExpVars env sh + -> SeqSchedule op env (Buffer e) + + -- | Buffer inlet. To pass an array constant in this data type, one may need + -- multiple 'Use' constructs because of the explicit Structure-of-Arrays. + -- Triggers (possibly) asynchronous host->device transfer if necessary. + -- + Use :: ScalarType e + -> Int -- Number of elements + -> Buffer e + -> SeqSchedule op env (Buffer e) + + -- | Capture a scalar in a singleton buffer + -- + Unit :: ExpVar env e + -> SeqSchedule op env (Buffer e) + + -- | If-then-else for array-level computations + -- + Acond :: ExpVar env PrimBool + -> SeqSchedule op env a + -> SeqSchedule op env a + -> SeqSchedule op env a + + -- Value-recursion for array-level computations + -- The uniqueness guarantees are an invariant of the loop + -- and should hold before and after each iteration. + -- + Awhile :: Uniquenesses a + -> SeqScheduleFun op env (a -> PrimBool) + -> SeqScheduleFun op env (a -> a) + -> GroundVars env a + -> SeqSchedule op env a + + Aassert :: Text + -> Exp env PrimBool + -> SeqSchedule kernel env () + + Aassume :: Exp env PrimBool + -> SeqSchedule kernel env () + +data SeqScheduleFun kernel env t where + Slam :: GLeftHandSide s env env' + -> SeqScheduleFun kernel env' t + -> SeqScheduleFun kernel env (s -> t) + + Sbody :: SeqSchedule kernel env t + -> SeqScheduleFun kernel env t + +instance IsSchedule SequentialSchedule where + -- 'a' is a ground type (ie, can be represented using GroundR) + type ScheduleInput SequentialSchedule a = a + type ScheduleOutput SequentialSchedule a = MVar a + + rnfSchedule (SequentialLam lhs f) = rnfLeftHandSide rnfGroundR lhs `seq` rnfSchedule f + rnfSchedule (SequentialBody body) = rnfSchedule' body + + convertScheduleFun = convertScheduleFun' + + callScheduledFun (GFunctionRbody repr) f + | Refl <- reprIsBody @SequentialSchedule repr = do + destination <- newEmptyMVar + f destination + takeMVar destination + callScheduledFun (GFunctionRlam _ ret) f = do + return $ \a -> do + callScheduledFun @SequentialSchedule ret $ f a + +convertScheduleFun' + :: IsKernel kernel + => Partition.PartitionedAfun (KernelOperation kernel) env t -> SequentialSchedule kernel env (Scheduled SequentialSchedule t) +convertScheduleFun' (Operation.Alam lhs f) = SequentialLam lhs $ convertScheduleFun' f +convertScheduleFun' (Operation.Abody body) + | Refl <- reprIsBody @SequentialSchedule $ groundsR body = SequentialBody $ convertSchedule' body + +convertSchedule' :: forall kernel env t. IsKernel kernel => Partition.PartitionedAcc (KernelOperation kernel) env t -> SeqSchedule kernel env t +convertSchedule' (Operation.Exec op args) + | CompiledKernel kernel args' <- compileKernel' @env @kernel op args = Exec (kernelMetadata kernel) kernel args' +convertSchedule' (Operation.Return vars) = Return vars +convertSchedule' (Operation.Manifest var) = Return $ TupRsingle var +convertSchedule' (Operation.Compute expr) = Compute expr +convertSchedule' (Operation.Alet lhs us bnd expr) = Alet lhs us (convertSchedule' bnd) (convertSchedule' expr) +convertSchedule' (Operation.Alloc shr tp sh) = Alloc shr tp sh +convertSchedule' (Operation.Use tp n buffer) = Use tp n buffer +convertSchedule' (Operation.Unit var) = Unit var +convertSchedule' (Operation.Acond var true false) = Acond var (convertSchedule' true) (convertSchedule' false) +convertSchedule' (Operation.Awhile us condition step initial) = Awhile us (convertScheduleFun'' condition) (convertScheduleFun'' step) initial +convertSchedule' (Operation.Aassert msg cond) = Alet (LeftHandSideWildcard TupRunit) TupRunit (Aassert msg cond) $ Compute $ Const scalarTypeWord8 0 +convertSchedule' (Operation.Aassume cond) = Alet (LeftHandSideWildcard TupRunit) TupRunit (Aassume cond) $ Compute $ Const scalarTypeWord8 0 +convertSchedule' (Operation.Fence _ next) = convertSchedule' next +convertSchedule' Operation.Atrace{} = error "TODO WALL: NON-EXHAUSTIVE PATTERN MATCH" + +convertScheduleFun'' :: forall kernel env t. IsKernel kernel => Partition.PartitionedAfun (KernelOperation kernel) env t -> SeqScheduleFun kernel env t +convertScheduleFun'' (Operation.Alam lhs f) = Slam lhs $ convertScheduleFun'' f +convertScheduleFun'' (Operation.Abody body) = Sbody $ convertSchedule' body + +rnfSchedule' :: IsKernel kernel => SeqSchedule kernel env t -> () +rnfSchedule' (Exec metadata kernel args) = rnf' metadata `seq` rnf' kernel `seq` rnfSArgs args +rnfSchedule' (Return vars) = rnfGroundVars vars +rnfSchedule' (Compute expr) = rnfOpenExp expr +rnfSchedule' (Alet lhs us bnd expr) + = rnfLeftHandSide rnfGroundR lhs + `seq` rnfTupR rnfUniqueness us + `seq` rnfSchedule' bnd + `seq` rnfSchedule' expr +rnfSchedule' (Alloc shr tp sh) = rnfShapeR shr `seq` rnfScalarType tp `seq` rnfTupR rnfExpVar sh +rnfSchedule' (Use tp n buffer) = n `seq` buffer `seq` rnfScalarType tp +rnfSchedule' (Unit var) = rnfExpVar var +rnfSchedule' (Acond var true false) = rnfExpVar var `seq` rnfSchedule' true `seq` rnfSchedule' false +rnfSchedule' (Awhile us condition step initial) = rnfTupR rnfUniqueness us `seq` rnfScheduleFun condition `seq` rnfScheduleFun step `seq` rnfGroundVars initial +rnfSchedule' (Aassert msg cond) = rnf msg `seq` rnfOpenExp cond +rnfSchedule' (Aassume cond) = rnfOpenExp cond + +rnfScheduleFun :: IsKernel kernel => SeqScheduleFun kernel env t -> () +rnfScheduleFun (Slam lhs f) = rnfLeftHandSide rnfGroundR lhs `seq` rnfScheduleFun f +rnfScheduleFun (Sbody body) = rnfSchedule' body diff --git a/src/Data/Array/Accelerate/AST/Schedule/Uniform.hs b/src/Data/Array/Accelerate/AST/Schedule/Uniform.hs new file mode 100644 index 000000000..1f3cead38 --- /dev/null +++ b/src/Data/Array/Accelerate/AST/Schedule/Uniform.hs @@ -0,0 +1,499 @@ +{-# LANGUAGE TypeFamilies #-} +{-# OPTIONS_HADDOCK hide #-} +-- | +-- Module : Data.Array.Accelerate.AST.Schedule.Uniform +-- Copyright : [2008..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.AST.Schedule.Uniform ( + UniformSchedule(..), UniformScheduleFun(..), + SArg(..), SArgs, sargVars, sargOutputVars, sargBufferVars, + Input, Output, inputSingle, outputSingle, inputR, outputR, + InputOutputR(..), inputOutputInputR, inputOutputOutputR, + Binding(..), Effect(..), + BaseR(..), BasesR, BaseVar, BaseVars, BLeftHandSide, + Signal(..), SignalResolver(..), Ref(..), OutputRef(..), + module Operation, + Cluster, + await, resolve, + signalResolverImpossible, scalarSignalResolverImpossible, + rnfBaseR, + directlyAwaits1, reorder, trivialBinding, trivialSchedule, trivialEffect, + + -- ** Free variables + freeVars, funFreeVars, effectFreeVars, bindingFreeVars, +) where + +import Data.Array.Accelerate.AST.Exp +import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.IdxSet ( IdxSet ) +import qualified Data.Array.Accelerate.AST.IdxSet as IdxSet +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.AST.Var +import Data.Array.Accelerate.Type +import Data.Array.Accelerate.Representation.Array +import Data.Array.Accelerate.Representation.Shape +import Data.Array.Accelerate.Representation.Type +import qualified Data.Array.Accelerate.AST.Operation as Operation hiding (PreOpenAcc(..), PreOpenAfun(..)) +import Data.Array.Accelerate.AST.Partitioned as Partitioned hiding (PreOpenAcc(..), PreOpenAfun(..), PartitionedAcc, PartitionedAfun) +import Data.Array.Accelerate.AST.Kernel +import Data.Array.Accelerate.Trafo.Exp.Substitution +import Control.Concurrent.MVar +import Control.Monad +import Data.Either +import Data.IORef +import Data.Typeable ( (:~:)(..) ) +import Data.Text ( Text ) + +-- Generic schedule for a uniform memory space and uniform scheduling. +-- E.g., we don't have host and device memory or scheduling. +-- The schedule will exploit task parallelism. + +-- The schedule consists of bindings, effects and (parallel) control flow +-- +-- Some schedules/terms may be 'spawn-closed'. These terms guarantee that when +-- the execution of the term finishes, all spawned tasks have also finished. +-- For instance, this program is spawn-closed: +-- +-- (s0, s0') = new signal +-- spawn { +-- ... +-- resolve [s0'] +-- } +-- ... +-- await [s0] +-- +-- However, this term is not self-closed, as the spawned term may outlive the +-- outer term. +-- +-- spawn { +-- ... +-- } +-- ... +-- +-- Note that not necessarily all subterms of a spawn-closed term as spawn-closed. +-- For instance, consider: +-- +-- (s0, s0') = new signal +-- (s1, s1') = new signal +-- spawn { +-- spawn { +-- resolve [s1'] +-- } +-- resolve [s0'] +-- } +-- await [s0, s1] +-- Here, the first spawned term is not self-closed: +-- its execution may finish before the execution of the inner spawned term ends. +-- The entire term is self-closed, as it waits on the completions of both +-- spawned subterms. +data UniformSchedule kernel env where + Return :: UniformSchedule kernel env + + Alet :: BLeftHandSide t env env' + -> Binding env t + -> UniformSchedule kernel env' + -> UniformSchedule kernel env + + Effect :: Effect kernel env + -> UniformSchedule kernel env + -> UniformSchedule kernel env + + Acond :: ExpVar env PrimBool + -> UniformSchedule kernel env -- True branch + -> UniformSchedule kernel env -- False branch + -> UniformSchedule kernel env -- Operations after the if-then-else + -> UniformSchedule kernel env + + -- The step function of the loop outputs a bool to denote whether the loop should + -- proceed. If true, then the other output values should also be filled, possibly at + -- a later point in time. If it is false, then no other output values may be filled. + -- + -- An Awhile loop may be executed in parallel, meaning that we may work on multiple + -- iterations of the loop in parallel. This may be possible, typically if the state + -- of the loop consists of multiple variables. Synchronisations are handled via + -- Signal(Resolver)s in 'input' and 'output'. + -- A sequential variant is available in AwhileSeq. + Awhile :: InputOutputR input output + -- The body of the while loop. + -- Should be spawn-closed. + -> UniformScheduleFun kernel env (input -> Output PrimBool -> output -> ()) + -> BaseVars env input + -> UniformSchedule kernel env -- Operations after the while loop + -> UniformSchedule kernel env + + -- Variant of Awhile, that can be executed sequentially. This only applies to + -- parallelism between iterations of the loop. Within an iteration, parallelism + -- can be exploited as usual. + -- Iteration 'i+1' only starts when iteration 'i' is finished. Since the + -- step function is spawn-closed, this guarantees proper synchronisations. + -- Hence we don't need Signal(Resolver)s in 'input' or 'output'. + AwhileSeq + -- Should not contain InputOutputRsignal + :: InputOutputR input output + -- The body of the while loop + -- Should be spawn-closed. + -> UniformScheduleFun kernel env (input -> OutputRef PrimBool -> output -> ()) + -> BaseVars env input + -> UniformSchedule kernel env -- Operations after the while loop + -> UniformSchedule kernel env + + -- Spawns a task/thread for the first subterm, and continues with the second subterm. + -- Note that we use the name 'spawn' instead of 'fork', to imply that the first + -- subterm is not executed directly, but handed over to a different core or queued. + Spawn :: UniformSchedule kernel env + -> UniformSchedule kernel env + -> UniformSchedule kernel env + +data UniformScheduleFun kernel env t where + Slam :: BLeftHandSide s env env' + -> UniformScheduleFun kernel env' t + -> UniformScheduleFun kernel env (s -> t) + + Sbody :: UniformSchedule kernel env + -> UniformScheduleFun kernel env () + +type family Input t where + Input () = () + Input (a, b) = (Input a, Input b) + Input t = (Signal, Ref t) + +inputSingle :: forall t. GroundR t -> (Input t, Output t) :~: ((Signal, Ref t), (SignalResolver, OutputRef t)) +-- Last case of type family Input and Output. +-- We must pattern match to convince the type checker that +-- t is not () or (a, b). +inputSingle (GroundRbuffer _) = Refl +inputSingle (GroundRscalar (VectorScalarType _)) = Refl +inputSingle (GroundRscalar (SingleScalarType (NumSingleType tp))) = case tp of + IntegralNumType TypeInt -> Refl + IntegralNumType TypeInt8 -> Refl + IntegralNumType TypeInt16 -> Refl + IntegralNumType TypeInt32 -> Refl + IntegralNumType TypeInt64 -> Refl + IntegralNumType TypeWord -> Refl + IntegralNumType TypeWord8 -> Refl + IntegralNumType TypeWord16 -> Refl + IntegralNumType TypeWord32 -> Refl + IntegralNumType TypeWord64 -> Refl + FloatingNumType TypeHalf -> Refl + FloatingNumType TypeFloat -> Refl + FloatingNumType TypeDouble -> Refl + +inputR :: forall t. GroundsR t -> BasesR (Input t) +inputR TupRunit = TupRunit +inputR (TupRpair t1 t2) = TupRpair (inputR t1) (inputR t2) +inputR (TupRsingle ground) + -- Last case of type family Input. + -- We must pattern match to convince the type checker that + -- t is not () or (a, b). + | Refl <- inputSingle ground = TupRsingle BaseRsignal `TupRpair` TupRsingle (BaseRref ground) + +type family Output t where + Output () = () + Output (a, b) = (Output a, Output b) + Output t = (SignalResolver, OutputRef t) + +data SArg env t where + SArgScalar :: ExpVar env e + -> SArg env (Var' e) + + SArgBuffer :: Modifier m + -> GroundVar env (Buffer e) + -> SArg env (m DIM1 e) + +type SArgs env = PreArgs (SArg env) + +outputR :: GroundsR t -> BasesR (Output t) +outputR TupRunit = TupRunit +outputR (TupRpair t1 t2) = TupRpair (outputR t1) (outputR t2) +outputR (TupRsingle ground) + -- Last case of type family Output. + -- We must pattern match to convince the type checker that + -- t is not () or (a, b). + | Refl <- inputSingle ground = TupRsingle BaseRsignalResolver `TupRpair` TupRsingle (BaseRrefWrite ground) + +outputSingle :: GroundR t -> Output t :~: (SignalResolver, OutputRef t) +outputSingle (GroundRbuffer _) = Refl +outputSingle (GroundRscalar (VectorScalarType _)) = Refl +outputSingle (GroundRscalar (SingleScalarType tp)) = case tp of + NumSingleType (IntegralNumType t) -> case t of + TypeInt -> Refl + TypeInt8 -> Refl + TypeInt16 -> Refl + TypeInt32 -> Refl + TypeInt64 -> Refl + TypeWord -> Refl + TypeWord8 -> Refl + TypeWord16 -> Refl + TypeWord32 -> Refl + TypeWord64 -> Refl + NumSingleType (FloatingNumType t) -> case t of + TypeHalf -> Refl + TypeFloat -> Refl + TypeDouble -> Refl + +-- Relation between input and output +data InputOutputR input output where + InputOutputRsignal :: InputOutputR Signal SignalResolver + InputOutputRref :: GroundR t -> InputOutputR (Ref t) (OutputRef t) + InputOutputRpair :: InputOutputR i1 o1 + -> InputOutputR i2 o2 + -> InputOutputR (i1, i2) (o1, o2) + InputOutputRunit :: InputOutputR () () + +inputOutputInputR :: InputOutputR input output -> BasesR input +inputOutputInputR InputOutputRsignal = TupRsingle BaseRsignal +inputOutputInputR (InputOutputRref tp) = TupRsingle $ BaseRref tp +inputOutputInputR (InputOutputRpair io1 io2) = inputOutputInputR io1 `TupRpair` inputOutputInputR io2 +inputOutputInputR InputOutputRunit = TupRunit + +inputOutputOutputR :: InputOutputR input output -> BasesR output +inputOutputOutputR InputOutputRsignal = TupRsingle BaseRsignalResolver +inputOutputOutputR (InputOutputRref tp) = TupRsingle $ BaseRrefWrite tp +inputOutputOutputR (InputOutputRpair io1 io2) = inputOutputOutputR io1 `TupRpair` inputOutputOutputR io2 +inputOutputOutputR InputOutputRunit = TupRunit + +-- Bindings of instructions which have some return value. +-- They cannot perform side effects. +-- +data Binding env t where + Compute :: Exp env t + -> Binding env t + + NewSignal :: String -- Name for easier debugging + -> Binding env (Signal, SignalResolver) + + NewRef :: GroundR t + -> Binding env (Ref t, OutputRef t) + + Alloc :: ShapeR sh + -> ScalarType e + -> ExpVars env sh + -> Binding env (Buffer e) + + Use :: ScalarType e + -> Int + -> Buffer e + -> Binding env (Buffer e) + + Unit :: ExpVar env e + -> Binding env (Buffer e) + + -- undefined result if the Ref hasn't been written + RefRead :: BaseVar env (Ref t) -> Binding env t + +-- Effects do not have a return value. +data Effect kernel env where + Exec :: KernelMetadata kernel f + -> KernelFun kernel f + -> SArgs env f + -> Effect kernel env + + SignalAwait :: [Idx env Signal] -> Effect kernel env + + SignalResolve :: [Idx env SignalResolver] -> Effect kernel env + + RefWrite :: BaseVar env (OutputRef t) -> BaseVar env t -> Effect kernel env + + Aassert :: Text -> Exp env PrimBool -> Effect kernel env + + Atrace :: Text -> ArrayDescriptors env t -> Effect kernel env + +-- A base value in the schedule is a scalar, buffer, a signal (resolver) +-- or a (possibly mutable) reference +-- +data BaseR t where + BaseRground :: GroundR t -> BaseR t + BaseRsignal :: BaseR Signal + BaseRsignalResolver :: BaseR SignalResolver + BaseRref :: GroundR t -> BaseR (Ref t) + BaseRrefWrite :: GroundR t -> BaseR (OutputRef t) + +instance Distributes BaseR where + reprIsSingle (BaseRground tp) = reprIsSingle tp + reprIsSingle BaseRsignal = Refl + reprIsSingle BaseRsignalResolver = Refl + reprIsSingle (BaseRref _) = Refl + reprIsSingle (BaseRrefWrite _) = Refl + + pairImpossible (BaseRground tp) = pairImpossible tp + unitImpossible (BaseRground tp) = unitImpossible tp + +-- Tuples of base values +type BasesR = TupR BaseR + +type BaseVar = Var BaseR +type BaseVars env = Vars BaseR env + +type BLeftHandSide = LeftHandSide BaseR + +newtype Signal = Signal (MVar ()) +newtype SignalResolver = SignalResolver (MVar ()) + +-- Note: Refs should only be assigned once. They start undefined, and are at some point resolved to a value +newtype Ref t = Ref (IORef t) +newtype OutputRef t = OutputRef (IORef t) + +await :: [Idx env Signal] -> UniformSchedule kernel env -> UniformSchedule kernel env +await [] schedule = schedule +await _ Return = Return +await signals (Effect (SignalAwait signals') schedule) + = Effect (SignalAwait $ signals ++ signals') schedule +await signals schedule = Effect (SignalAwait signals) schedule + +resolve :: [Idx env SignalResolver] -> UniformSchedule kernel env -> UniformSchedule kernel env +resolve [] schedule = schedule +resolve signals (Effect (SignalResolve signals') schedule) + = Effect (SignalResolve $ signals ++ signals') schedule +resolve signals schedule = Effect (SignalResolve signals) schedule + +freeVars :: UniformSchedule kernel env -> IdxSet env +freeVars Return = IdxSet.empty +freeVars (Alet lhs bnd s) = bindingFreeVars bnd `IdxSet.union` IdxSet.drop' lhs (freeVars s) +freeVars (Effect effect s) = effectFreeVars effect `IdxSet.union` freeVars s +freeVars (Acond c t f s) + = IdxSet.insertVar c + $ IdxSet.union (freeVars t) + $ IdxSet.union (freeVars f) + $ freeVars s +freeVars (Awhile _ step ini continuation) + = IdxSet.union (funFreeVars step) + $ IdxSet.union (IdxSet.fromVarList $ flattenTupR ini) + $ freeVars continuation +freeVars (AwhileSeq _ step ini continuation) + = IdxSet.union (funFreeVars step) + $ IdxSet.union (IdxSet.fromVarList $ flattenTupR ini) + $ freeVars continuation +freeVars (Spawn s1 s2) = freeVars s1 `IdxSet.union` freeVars s2 + +funFreeVars :: UniformScheduleFun kernel env t -> IdxSet env +funFreeVars (Sbody s) = freeVars s +funFreeVars (Slam lhs f) = IdxSet.drop' lhs $ funFreeVars f + +bindingFreeVars :: Binding env t -> IdxSet env +bindingFreeVars (NewSignal _) = IdxSet.empty +bindingFreeVars (NewRef _) = IdxSet.empty +bindingFreeVars (Alloc _ _ sh) = IdxSet.fromVarList $ flattenTupR sh +bindingFreeVars (Use _ _ _) = IdxSet.empty +bindingFreeVars (Unit var) = IdxSet.singletonVar var +bindingFreeVars (RefRead var) = IdxSet.singletonVar var +bindingFreeVars (Compute e) = IdxSet.fromList $ map f $ arrayInstrsInExp e + where + f :: Exists (ArrayInstr env) -> Exists (Idx env) + f (Exists (Index (Var _ idx))) = Exists idx + f (Exists (Parameter (Var _ idx))) = Exists idx + +effectFreeVars :: Effect kernel env -> IdxSet env +effectFreeVars (Exec _ _ args) = IdxSet.fromList $ sargVars args +effectFreeVars (SignalAwait signals) = IdxSet.fromList $ map Exists $ signals +effectFreeVars (SignalResolve resolvers) = IdxSet.fromList $ map Exists resolvers +effectFreeVars (RefWrite ref value) = IdxSet.insertVar ref $ IdxSet.singletonVar value +effectFreeVars (Aassert _ cond) = bindingFreeVars $ Compute cond +effectFreeVars (Atrace _ t) = arrayDescriptorsIdxSet t + +sargVar :: SArg env t -> Exists (Idx env) +sargVar (SArgScalar v) = Exists $ varIdx v +sargVar (SArgBuffer _ v) = Exists $ varIdx v + +sargVars :: SArgs env t -> [Exists (Idx env)] +sargVars (a :>: as) = sargVar a : sargVars as +sargVars ArgsNil = [] + +sargOutputVar :: SArg env t -> Maybe (Exists (Idx env)) +sargOutputVar (SArgScalar _) = Nothing +sargOutputVar (SArgBuffer In _) = Nothing +sargOutputVar (SArgBuffer _ v) = Just $ Exists $ varIdx v + +sargOutputVars :: SArgs env t -> [Exists (Idx env)] +sargOutputVars (a :>: as) = maybe id (:) (sargOutputVar a) $ sargOutputVars as +sargOutputVars ArgsNil = [] + +sargBufferVar :: SArg env t -> Maybe (Exists (Idx env)) +sargBufferVar (SArgScalar _) = Nothing +sargBufferVar (SArgBuffer _ v) = Just $ Exists $ varIdx v + +sargBufferVars :: SArgs env t -> [Exists (Idx env)] +sargBufferVars (a :>: as) = maybe id (:) (sargBufferVar a) $ sargBufferVars as +sargBufferVars ArgsNil = [] + +signalResolverImpossible :: GroundsR SignalResolver -> a +signalResolverImpossible (TupRsingle (GroundRscalar tp)) = scalarSignalResolverImpossible tp + +scalarSignalResolverImpossible :: ScalarType SignalResolver -> a +scalarSignalResolverImpossible (SingleScalarType (NumSingleType (IntegralNumType tp))) = case tp of {} +scalarSignalResolverImpossible (SingleScalarType (NumSingleType (FloatingNumType tp))) = case tp of {} + +rnfBaseR :: BaseR t -> () +rnfBaseR (BaseRground ground) = rnfGroundR ground +rnfBaseR BaseRsignal = () +rnfBaseR BaseRsignalResolver = () +rnfBaseR (BaseRref ground) = rnfGroundR ground +rnfBaseR (BaseRrefWrite ground) = rnfGroundR ground + +-- Returns a list of signals on which the schedule waits before doing any +-- work. Assumes that the schedule starts with SignalAwait; the schedule +-- can be reordered to assure this by using 'reorder'. +directlyAwaits1 :: UniformSchedule kernel env -> ([Idx env Signal], UniformSchedule kernel env) +directlyAwaits1 (Effect (SignalAwait signals) next) = (signals, next) +directlyAwaits1 schedule = ([], schedule) + +-- Reorders a schedule. +-- Moves SignalAwait to the front of the schedule if possible. +-- Stops at a Spawn, making it safe to call this function while +-- constructing a schedule in different parts of the tree, +-- without causing quadratic traversal costs. +reorder :: forall kernel env. UniformSchedule kernel env -> UniformSchedule kernel env +reorder = uncurry await . go Just + where + go :: env' :?> env -> UniformSchedule kernel env' -> ([Idx env Signal], UniformSchedule kernel env') + -- Try to move SignalAwait to the front of the schedule + go k (Effect (SignalAwait awaits) next) = (rs ++ signals, await ls next') + where + (ls, rs) = partitionEithers $ map (try k) awaits + (signals, next') = go k next + -- Write may be postponed: a write doesn't do synchronisation, + -- that is done by a signal(resolver). + go k (Effect effect@RefWrite{} next) = (signals, Effect effect next') + where + (signals, next') = go k next + go k (Alet lhs bnd next) + | trivialBinding bnd = (signals, Alet lhs bnd next') + where + (signals, next') = go (k <=< strengthenWithLHS lhs) next + go _ sched = ([], sched) + + try :: env' :?> env -> Idx env' t -> Either (Idx env' t) (Idx env t) + try k ix = case k ix of + Just ix' -> Right ix' + Nothing -> Left ix + +-- If a binding will only take little time to execute, we say it's trivial and +-- move it (usually postpone) in the schedule. +trivialBinding :: Binding env t -> Bool +trivialBinding (NewSignal _) = True +trivialBinding (NewRef _) = True +trivialBinding (Alloc _ _ _) = True +trivialBinding (Use _ _ _) = True +trivialBinding (Unit _) = True +trivialBinding (RefRead _) = True +trivialBinding (Compute e) = expIsTrivial (const True) e + +-- If a schedule does not do blocking or slow operations, we say it's trivial +-- and don't need to spawn it as we do not gain much task parallelism from it. +trivialSchedule :: UniformSchedule kernel env -> Bool +trivialSchedule (Effect effect next) = trivialEffect effect && trivialSchedule next +trivialSchedule (Alet _ bnd next) = trivialBinding bnd && trivialSchedule next +trivialSchedule Return = True +trivialSchedule (Acond _ true false next) = trivialSchedule true && trivialSchedule false && trivialSchedule next +trivialSchedule _ = False + +trivialEffect :: Effect kernel env -> Bool +trivialEffect SignalResolve{} = True +trivialEffect RefWrite{} = True +trivialEffect (Aassert _ cond) = expIsTrivial (const True) cond +trivialEffect _ = False diff --git a/src/Data/Array/Accelerate/AST/Var.hs b/src/Data/Array/Accelerate/AST/Var.hs index e0c7ecc0d..c393c8361 100644 --- a/src/Data/Array/Accelerate/AST/Var.hs +++ b/src/Data/Array/Accelerate/AST/Var.hs @@ -1,6 +1,4 @@ -{-# LANGUAGE GADTs #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TemplateHaskellQuotes #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.AST.Var @@ -16,12 +14,14 @@ module Data.Array.Accelerate.AST.Var where import Data.Array.Accelerate.Representation.Type + ( TupR(..), mapTupR, liftTupR, rnfTupR, Distributes(..) ) import Data.Array.Accelerate.AST.Idx import Language.Haskell.TH.Extra +import Data.Typeable ( (:~:)(..) ) -data Var s env t = Var (s t) (Idx env t) +data Var s env t = Var { varType :: s t, varIdx :: Idx env t } type Vars s env = TupR (Var s env) varsType :: Vars s env t -> TupR s t @@ -29,7 +29,6 @@ varsType TupRunit = TupRunit varsType (TupRsingle (Var t _)) = TupRsingle t varsType (TupRpair a b) = TupRpair (varsType a) (varsType b) - rnfVar :: (forall b. s b -> ()) -> Var s env t -> () rnfVar f (Var t idx) = f t `seq` rnfIdx idx @@ -42,3 +41,30 @@ liftVar f (Var s idx) = [|| Var $$(f s) $$(liftIdx idx) ||] liftVars :: (forall b. s b -> CodeQ (s b)) -> Vars s env t -> CodeQ (Vars s env t) liftVars f = liftTupR (liftVar f) +{-# INLINEABLE matchVar #-} +matchVar :: Var s env t1 -> Var s env t2 -> Maybe (t1 :~: t2) +matchVar (Var _ v1) (Var _ v2) = matchIdx v1 v2 + +{-# INLINEABLE matchVars #-} +matchVars :: Vars s env t1 -> Vars s env t2 -> Maybe (t1 :~: t2) +matchVars TupRunit TupRunit = Just Refl +matchVars (TupRsingle v1) (TupRsingle v2) + | Just Refl <- matchVar v1 v2 = Just Refl +matchVars (TupRpair v w) (TupRpair x y) + | Just Refl <- matchVars v x + , Just Refl <- matchVars w y = Just Refl +matchVars _ _ = Nothing + +mapVar :: (forall u. s u -> s' u) -> Var s env t -> Var s' env t +mapVar f (Var t ix) = Var (f t) ix + +mapVars :: (forall u. s u -> s' u) -> Vars s env t -> Vars s' env t +mapVars f = mapTupR (mapVar f) + +instance Show (Var s env t) where + show (Var _ idx) = "Var " ++ show (idxToInt idx) + +instance Distributes s => Distributes (Var s env) where + reprIsSingle (Var t _) = reprIsSingle t + pairImpossible (Var t _) = pairImpossible t + unitImpossible (Var t _) = unitImpossible t diff --git a/src/Data/Array/Accelerate/Analysis/Exp.hs b/src/Data/Array/Accelerate/Analysis/Exp.hs new file mode 100644 index 000000000..a360b22c9 --- /dev/null +++ b/src/Data/Array/Accelerate/Analysis/Exp.hs @@ -0,0 +1,105 @@ +{-# OPTIONS_HADDOCK hide #-} +-- | +-- Module : Data.Array.Accelerate.Analysis.Exp +-- Copyright : [2012..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.Analysis.Exp ( + isCommutative, isLeftIdentity, isRightIdentity, findLeftIdentity, findRightIdentity +) where + +import Data.Array.Accelerate.AST.Exp +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.AST.Environment +import Data.Array.Accelerate.Analysis.Match +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Type +import Data.Array.Accelerate.Trafo.Var +import Data.Array.Accelerate.Trafo.Exp.Simplify +import Data.Array.Accelerate.Trafo.Exp.Substitution +import Data.Array.Accelerate.Error + +import Data.Maybe +import Data.List (find) +import Prelude hiding ( exp ) + +isCommutative :: IsArrayInstr arr => PreOpenFun arr () (a -> a -> b) -> Bool +isCommutative fun = isJust $ matchOpenFun fun' flipped + where + fun' = simplifyFun fun + flipped = simplifyFun $ flipFun fun' + +isLeftIdentity :: IsArrayInstr arr => PreOpenFun arr () (a -> b -> b) -> PreOpenExp arr () a -> Bool +isLeftIdentity fun arg = isJust $ isIdentity $ simplifyFun $ applyLeft fun arg + +isRightIdentity :: IsArrayInstr arr => PreOpenFun arr () (a -> b -> a) -> PreOpenExp arr () b -> Bool +isRightIdentity fun arg = isJust $ isIdentity $ simplifyFun $ applyRight fun arg + +findLeftIdentity :: IsArrayInstr arr => PreOpenFun arr () (a -> b -> b) -> Maybe a +findLeftIdentity fun = find (isLeftIdentity fun . mkConstant tp) (guesses tp) + where + tp = case fun of + Lam lhs _ -> lhsToTupR lhs + _ -> internalError "Expected binary function" + +findRightIdentity :: IsArrayInstr arr => PreOpenFun arr () (a -> b -> a) -> Maybe b +findRightIdentity fun = find (isRightIdentity fun . mkConstant tp) (guesses tp) + where + tp = case fun of + Lam _ (Lam lhs _) -> lhsToTupR lhs + _ -> internalError "Expected binary function" + +flipFun :: IsArrayInstr arr => PreOpenFun arr () (a -> b -> c) -> PreOpenFun arr () (b -> a -> c) +flipFun fun + | DeclareVars lhsB _ valueB <- declareVars tpB + , DeclareVars lhsA kA valueA <- declareVars tpA + = Lam lhsB $ Lam lhsA $ Body + $ apply fun (expVars $ valueA weakenId) (expVars $ valueB kA) + where + (tpA, tpB) = case fun of + Lam lhsA (Lam lhsB _) -> (lhsToTupR lhsA, lhsToTupR lhsB) + _ -> internalError "Expected binary function" + +apply :: IsArrayInstr arr => PreOpenFun arr () (a -> b -> c) -> PreOpenExp arr env a -> PreOpenExp arr env b -> PreOpenExp arr env c +apply fun a b + | Lam lhsA (Lam lhsB (Body body)) <- weakenE weakenEmpty fun + = Let (LeftHandSidePair lhsA lhsB) (Pair a b) body + | otherwise = internalError "Expected binary function" + +applyLeft :: IsArrayInstr arr => PreOpenFun arr () (a -> b -> c) -> PreOpenExp arr () a -> PreOpenFun arr () (b -> c) +applyLeft fun arg = applyRight (flipFun fun) arg + +applyRight :: IsArrayInstr arr => PreOpenFun arr () (a -> b -> c) -> PreOpenExp arr () b -> PreOpenFun arr () (a -> c) +applyRight (Lam lhsA (Lam lhsB (Body body))) arg + = Lam lhsA $ Body $ Let lhsB (weakenE weakenEmpty arg) body +applyRight _ _ = internalError "Expected binary function" + +data Guess = Zero | One | MinValue | MaxValue + +guesses :: TypeR t -> [t] +guesses tp = mapMaybe (\g -> guess g tp) [Zero, One, MinValue, MaxValue] + +guess :: Guess -> TypeR t -> Maybe t +guess _ TupRunit = Just () +guess g (TupRsingle t) = guessScalar g t +guess g (TupRpair t1 t2) = (,) <$> guess g t1 <*> guess g t2 + +guessScalar :: Guess -> ScalarType t -> Maybe t +guessScalar g (SingleScalarType (NumSingleType (IntegralNumType t))) + | IntegralDict <- integralDict t = case g of + Zero -> Just 0 + One -> Just 1 + MinValue -> Just minBound + MaxValue -> Just maxBound +guessScalar g (SingleScalarType (NumSingleType (FloatingNumType t))) + | FloatingDict <- floatingDict t = case g of + Zero -> Just 0 + One -> Just 1 + MinValue -> Just (-1/0) -- negative infinity + MaxValue -> Just (1/0) -- infinity +guessScalar _ (VectorScalarType _) = Nothing diff --git a/src/Data/Array/Accelerate/Analysis/Hash.hs b/src/Data/Array/Accelerate/Analysis/Hash.hs index 7e0fe88c9..dd6b61fcf 100644 --- a/src/Data/Array/Accelerate/Analysis/Hash.hs +++ b/src/Data/Array/Accelerate/Analysis/Hash.hs @@ -1,11 +1,4 @@ -{-# LANGUAGE AllowAmbiguousTypes #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE MagicHash #-} -{-# LANGUAGE PatternGuards #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TemplateHaskell #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.Analysis.Hash @@ -36,69 +29,21 @@ module Data.Array.Accelerate.Analysis.Hash ( ) where import Data.Array.Accelerate.AST -import Data.Array.Accelerate.AST.Idx -import Data.Array.Accelerate.AST.LeftHandSide import Data.Array.Accelerate.AST.Var import Data.Array.Accelerate.Analysis.Hash.TH +import Data.Array.Accelerate.Analysis.Hash.Exp import Data.Array.Accelerate.Representation.Array -import Data.Array.Accelerate.Representation.Shape -import Data.Array.Accelerate.Representation.Slice import Data.Array.Accelerate.Representation.Stencil import Data.Array.Accelerate.Representation.Type -import Data.Array.Accelerate.Type -import Data.Primitive.Vec import Crypto.Hash.XKCP import Data.ByteString.Builder -import Data.ByteString.Builder.Extra -import Data.ByteString.Short.Internal ( ShortByteString(..) ) import Data.Monoid import System.IO.Unsafe ( unsafePerformIO ) import System.Mem.StableName ( hashStableName, makeStableName ) import Prelude hiding ( exp ) import qualified Data.Hashable as Hashable - --- Hashing --- ------- - -type Hash = SHA3_256 - -data HashOptions = HashOptions - { perfect :: Bool - -- ^ Should the hash function include _all_ substructure, recursively? - -- - -- Set to true (the default) if you want a truly unique fingerprint for - -- the entire expression: - -- - -- Example: - -- - -- xs, ys :: Acc (Vector Float) - -- xs = fill (constant (Z:.10)) 1.0 - -- ys = fill (constant (Z:.20)) 1.0 - -- - -- with perfect=True: - -- - -- hash xs = 2e1f91aca4c476d13b36f22462e73c15bbdd9fcacb0d4996280f6004058e9732 - -- hash ys = 2fce5c849b6c652192b09aaeafdc8029e57b9f006c1ecd79ccf9114f349aaf9e - -- - -- However, for a code generating backend the object code used to - -- evaluate both of these expressions is likely to be identical. - -- - -- Setting perfect=False results in: - -- - -- hash xs = hash ys = f97944b0ec64ab8aa989fd60c8b50e7ec3eff759d22d2b340039d837d74dfc3c - -- - -- Note that to be useful the provided 'EncodeAcc' function must also - -- understand this option, and the consumer of the hash value must be - -- agnostic to the elided details. - } - deriving Show - -defaultHashOptions :: HashOptions -defaultHashOptions = HashOptions True - - {-# INLINEABLE hashPreOpenAcc #-} hashPreOpenAcc :: HasArraysR acc => EncodeAcc acc -> PreOpenAcc acc aenv a -> Hash hashPreOpenAcc = hashPreOpenAccWith defaultHashOptions @@ -110,20 +55,6 @@ hashPreOpenAccWith options encodeAcc . toLazyByteString . encodePreOpenAcc options encodeAcc -{-# INLINEABLE hashOpenFun #-} -hashOpenFun :: OpenFun env aenv f -> Hash -hashOpenFun - = hashlazy - . toLazyByteString - . encodeOpenFun - -{-# INLINEABLE hashOpenExp #-} -hashOpenExp :: OpenExp env aenv t -> Hash -hashOpenExp - = hashlazy - . toLazyByteString - . encodeOpenExp - -- Array computations -- ------------------ @@ -152,8 +83,8 @@ encodePreOpenAcc options encodeAcc pacc = travF = encodeOpenFun travD :: Direction -> Builder - travD LeftToRight = intHost $(hashQ "L") - travD RightToLeft = intHost $(hashQ "R") + travD LeftToRight = intHost $(hashQ ("L" :: String)) + travD RightToLeft = intHost $(hashQ ("R" :: String)) deep :: Builder -> Builder deep | perfect options = id @@ -165,34 +96,36 @@ encodePreOpenAcc options encodeAcc pacc = | otherwise = encodeTypeR $ expType e in case pacc of - Alet lhs bnd body -> intHost $(hashQ "Alet") <> encodeLeftHandSide encodeArrayType lhs <> travA bnd <> travA body - Avar (Var repr v) -> intHost $(hashQ "Avar") <> encodeArrayType repr <> deep (encodeIdx v) - Apair a1 a2 -> intHost $(hashQ "Apair") <> travA a1 <> travA a2 - Anil -> intHost $(hashQ "Anil") - Atrace (Message _ _ msg) as bs -> intHost $(hashQ "Atrace") <> intHost (Hashable.hash msg) <> travA as <> travA bs - Apply _ f a -> intHost $(hashQ "Apply") <> travAF f <> travA a - Aforeign _ _ f a -> intHost $(hashQ "Aforeign") <> travAF f <> travA a - Use repr a -> intHost $(hashQ "Use") <> encodeArrayType repr <> deep (encodeArray a) - Awhile p f a -> intHost $(hashQ "Awhile") <> travAF f <> travAF p <> travA a - Unit _ e -> intHost $(hashQ "Unit") <> travE e - Generate _ e f -> intHost $(hashQ "Generate") <> deepE e <> travF f + Alet lhs bnd body -> intHost $(hashQ ("Alet" :: String)) <> encodeLeftHandSide encodeArrayType lhs <> travA bnd <> travA body + Avar (Var repr v) -> intHost $(hashQ ("Avar" :: String)) <> encodeArrayType repr <> deep (encodeIdx v) + Apair a1 a2 -> intHost $(hashQ ("Apair" :: String)) <> travA a1 <> travA a2 + Anil -> intHost $(hashQ ("Anil" :: String)) + Atrace (Message _ _ msg) as bs -> intHost $(hashQ ("Atrace" :: String)) <> intHost (Hashable.hash msg) <> travA as <> travA bs + Manifest as -> intHost $(hashQ ("Manifest" :: String)) <> travA as + Aassert msg cond as -> intHost $(hashQ ("Aassert" :: String)) <> intHost (Hashable.hash msg) <> travE cond <> travA as + Aassume cond as -> intHost $(hashQ ("Aassume" :: String)) <> travE cond <> travA as + Aforeign _ _ f a -> intHost $(hashQ ("Aforeign" :: String)) <> travAF f <> travA a + Use repr a -> intHost $(hashQ ("Use" :: String)) <> encodeArrayType repr <> deep (encodeArray a) + Awhile p f a -> intHost $(hashQ ("Awhile" :: String)) <> travAF f <> travAF p <> travA a + Unit _ e -> intHost $(hashQ ("Unit" :: String)) <> travE e + Generate _ e f -> intHost $(hashQ ("Generate" :: String)) <> deepE e <> travF f -- We don't need to encode the type of 'e' when perfect is False, as 'e' is an expression of type Bool. -- We thus use `deep (travE e)` instead of `deepE e`. - Acond e a1 a2 -> intHost $(hashQ "Acond") <> deep (travE e) <> travA a1 <> travA a2 - Reshape _ sh a -> intHost $(hashQ "Reshape") <> deepE sh <> travA a - Backpermute _ sh f a -> intHost $(hashQ "Backpermute") <> deepE sh <> travF f <> travA a - Transform _ sh f1 f2 a -> intHost $(hashQ "Transform") <> deepE sh <> travF f1 <> travF f2 <> travA a - Replicate spec ix a -> intHost $(hashQ "Replicate") <> deepE ix <> travA a <> encodeSliceIndex spec - Slice spec a ix -> intHost $(hashQ "Slice") <> deepE ix <> travA a <> encodeSliceIndex spec - Map _ f a -> intHost $(hashQ "Map") <> travF f <> travA a - ZipWith _ f a1 a2 -> intHost $(hashQ "ZipWith") <> travF f <> travA a1 <> travA a2 - Fold f e a -> intHost $(hashQ "Fold") <> travF f <> encodeMaybe travE e <> travA a - FoldSeg _ f e a s -> intHost $(hashQ "FoldSeg") <> travF f <> encodeMaybe travE e <> travA a <> travA s - Scan d f e a -> intHost $(hashQ "Scan") <> travD d <> travF f <> encodeMaybe travE e <> travA a - Scan' d f e a -> intHost $(hashQ "Scan'") <> travD d <> travF f <> travE e <> travA a - Permute f1 a1 f2 a2 -> intHost $(hashQ "Permute") <> travF f1 <> travA a1 <> travF f2 <> travA a2 - Stencil s _ f b a -> intHost $(hashQ "Stencil") <> travF f <> encodeBoundary (stencilEltR s) b <> travA a - Stencil2 s1 s2 _ f b1 a1 b2 a2 -> intHost $(hashQ "Stencil2") <> travF f <> encodeBoundary (stencilEltR s1) b1 <> travA a1 <> encodeBoundary (stencilEltR s2) b2 <> travA a2 + Acond e a1 a2 -> intHost $(hashQ ("Acond" :: String)) <> deep (travE e) <> travA a1 <> travA a2 + Reshape _ sh a -> intHost $(hashQ ("Reshape" :: String)) <> deepE sh <> travA a + Backpermute _ sh f a -> intHost $(hashQ ("Backpermute" :: String)) <> deepE sh <> travF f <> travA a + Transform _ sh f1 f2 a -> intHost $(hashQ ("Transform" :: String)) <> deepE sh <> travF f1 <> travF f2 <> travA a + Replicate spec ix a -> intHost $(hashQ ("Replicate" :: String)) <> deepE ix <> travA a <> encodeSliceIndex spec + Slice spec a ix -> intHost $(hashQ ("Slice" :: String)) <> deepE ix <> travA a <> encodeSliceIndex spec + Map _ f a -> intHost $(hashQ ("Map" :: String)) <> travF f <> travA a + ZipWith _ f a1 a2 -> intHost $(hashQ ("ZipWith" :: String)) <> travF f <> travA a1 <> travA a2 + Fold f e a -> intHost $(hashQ ("Fold" :: String)) <> travF f <> encodeMaybe travE e <> travA a + FoldSeg _ f e a s -> intHost $(hashQ ("FoldSeg" :: String)) <> travF f <> encodeMaybe travE e <> travA a <> travA s + Scan d f e a -> intHost $(hashQ ("Scan" :: String)) <> travD d <> travF f <> encodeMaybe travE e <> travA a + Scan' d f e a -> intHost $(hashQ ("Scan'" :: String)) <> travD d <> travF f <> travE e <> travA a + Permute f a1 a2 -> intHost $(hashQ ("Permute" :: String)) <> foldMap travF f <> travA a1 <> travA a2 + Stencil s _ f b a -> intHost $(hashQ ("Stencil" :: String)) <> travF f <> encodeBoundary (stencilEltR s) b <> travA a + Stencil2 s1 s2 _ f b1 a1 b2 a2 -> intHost $(hashQ ("Stencil2" :: String)) <> travF f <> encodeBoundary (stencilEltR s1) b1 <> travA a1 <> encodeBoundary (stencilEltR s2) b2 <> travA a2 {-- {-# INLINEABLE encodePreOpenSeq #-} @@ -221,50 +154,28 @@ encodePreOpenSeq encodeAcc s = travP p = case p of StreamIn arrs -> intHost . unsafePerformIO $! hashStableName `fmap` makeStableName arrs - ToSeq spec _ acc -> intHost $(hashQ "ToSeq") <> travA acc <> stringUtf8 (show spec) - MapSeq f x -> intHost $(hashQ "MapSeq") <> travAF f <> travV x - ChunkedMapSeq f x -> intHost $(hashQ "ChunkedMapSeq") <> travAF f <> travV x - ZipWithSeq f x y -> intHost $(hashQ "ZipWithSeq") <> travAF f <> travV x <> travV y - ScanSeq f e x -> intHost $(hashQ "ScanSeq") <> travF f <> travE e <> travV x + ToSeq spec _ acc -> intHost $(hashQ ("ToSeq" :: String)) <> travA acc <> stringUtf8 (show spec) + MapSeq f x -> intHost $(hashQ ("MapSeq" :: String)) <> travAF f <> travV x + ChunkedMapSeq f x -> intHost $(hashQ ("ChunkedMapSeq" :: String)) <> travAF f <> travV x + ZipWithSeq f x y -> intHost $(hashQ ("ZipWithSeq" :: String)) <> travAF f <> travV x <> travV y + ScanSeq f e x -> intHost $(hashQ ("ScanSeq" :: String)) <> travF f <> travE e <> travV x travC :: Consumer acc aenv senv' a -> Builder travC c = case c of - FoldSeq f e x -> intHost $(hashQ "FoldSeq") <> travF f <> travE e <> travV x - FoldSeqFlatten f acc x -> intHost $(hashQ "FoldSeqFlatten") <> travAF f <> travA acc <> travV x - Stuple t -> intHost $(hashQ "Stuple") <> encodeAtuple travC t + FoldSeq f e x -> intHost $(hashQ ("FoldSeq" :: String)) <> travF f <> travE e <> travV x + FoldSeqFlatten f acc x -> intHost $(hashQ ("FoldSeqFlatten" :: String)) <> travAF f <> travA acc <> travV x + Stuple t -> intHost $(hashQ ("Stuple" :: String)) <> encodeAtuple travC t in case s of - Producer p s' -> intHost $(hashQ "Producer") <> travP p <> travS s' - Consumer c -> intHost $(hashQ "Consumer") <> travC c - Reify ix -> intHost $(hashQ "Reify") <> travV ix + Producer p s' -> intHost $(hashQ ("Producer" :: String)) <> travP p <> travS s' + Consumer c -> intHost $(hashQ ("Consumer" :: String)) <> travC c + Reify ix -> intHost $(hashQ ("Reify" :: String)) <> travV ix --} -encodeIdx :: Idx env t -> Builder -encodeIdx = intHost . idxToInt - encodeArray :: Array sh e -> Builder encodeArray ad = intHost . unsafePerformIO $! hashStableName <$> makeStableName ad -encodeTupR :: (forall b. s b -> Builder) -> TupR s a -> Builder -encodeTupR _ TupRunit = intHost $(hashQ "TupRunit") -encodeTupR f (TupRpair r1 r2) = intHost $(hashQ "TupRpair") <> encodeTupR f r1 <> encodeTupR f r2 -encodeTupR f (TupRsingle s) = intHost $(hashQ "TupRsingle") <> f s - -encodeLeftHandSide :: (forall b. s b -> Builder) -> LeftHandSide s a env env' -> Builder -encodeLeftHandSide f (LeftHandSideWildcard r) = intHost $(hashQ "LeftHandSideWildcard") <> encodeTupR f r -encodeLeftHandSide f (LeftHandSidePair r1 r2) = intHost $(hashQ "LeftHandSidePair") <> encodeLeftHandSide f r1 <> encodeLeftHandSide f r2 -encodeLeftHandSide f (LeftHandSideSingle s) = intHost $(hashQ "LeftHandSideArray") <> f s - -encodeArrayType :: ArrayR a -> Builder -encodeArrayType (ArrayR shr tp) = encodeShapeR shr <> encodeTypeR tp - -encodeArraysType :: ArraysR arrs -> Builder -encodeArraysType = encodeTupR encodeArrayType - -encodeShapeR :: ShapeR sh -> Builder -encodeShapeR = intHost . rank - encodePreOpenAfun :: forall acc aenv f. HashOptions @@ -277,230 +188,21 @@ encodePreOpenAfun options travA afun = travL lhs l = encodeLeftHandSide encodeArrayType lhs <> encodePreOpenAfun options travA l in case afun of - Abody b -> intHost $(hashQ "Abody") <> travA options b - Alam lhs l -> intHost $(hashQ "Alam") <> travL lhs l + Abody b -> intHost $(hashQ ("Abody" :: String)) <> travA options b + Alam lhs l -> intHost $(hashQ ("Alam" :: String)) <> travL lhs l encodeBoundary :: TypeR e -> Boundary aenv (Array sh e) -> Builder -encodeBoundary _ Wrap = intHost $(hashQ "Wrap") -encodeBoundary _ Clamp = intHost $(hashQ "Clamp") -encodeBoundary _ Mirror = intHost $(hashQ "Mirror") -encodeBoundary tp (Constant c) = intHost $(hashQ "Constant") <> encodeConst tp c -encodeBoundary _ (Function f) = intHost $(hashQ "Function") <> encodeOpenFun f - -encodeSliceIndex :: SliceIndex slix sl co sh -> Builder -encodeSliceIndex SliceNil = intHost $(hashQ "SliceNil") -encodeSliceIndex (SliceAll r) = intHost $(hashQ "SliceAll") <> encodeSliceIndex r -encodeSliceIndex (SliceFixed r) = intHost $(hashQ "sliceFixed") <> encodeSliceIndex r - - --- Scalar expressions --- ------------------ - -{-# INLINEABLE encodeOpenExp #-} -encodeOpenExp - :: forall env aenv exp. - OpenExp env aenv exp - -> Builder -encodeOpenExp exp = - let - travE :: forall env' aenv' e. OpenExp env' aenv' e -> Builder - travE e = encodeOpenExp e - - travF :: OpenFun env' aenv' f -> Builder - travF = encodeOpenFun - in - case exp of - Let lhs bnd body -> intHost $(hashQ "Let") <> encodeLeftHandSide encodeScalarType lhs <> travE bnd <> travE body - Evar (Var tp ix) -> intHost $(hashQ "Evar") <> encodeScalarType tp <> encodeIdx ix - Nil -> intHost $(hashQ "Nil") - Pair e1 e2 -> intHost $(hashQ "Pair") <> travE e1 <> travE e2 - VecPack _ e -> intHost $(hashQ "VecPack") <> travE e - VecUnpack _ e -> intHost $(hashQ "VecUnpack") <> travE e - Const tp c -> intHost $(hashQ "Const") <> encodeScalarConst tp c - Undef tp -> intHost $(hashQ "Undef") <> encodeScalarType tp - IndexSlice spec ix sh -> intHost $(hashQ "IndexSlice") <> travE ix <> travE sh <> encodeSliceIndex spec - IndexFull spec ix sl -> intHost $(hashQ "IndexFull") <> travE ix <> travE sl <> encodeSliceIndex spec - ToIndex _ sh i -> intHost $(hashQ "ToIndex") <> travE sh <> travE i - FromIndex _ sh i -> intHost $(hashQ "FromIndex") <> travE sh <> travE i - Case e rhs def -> intHost $(hashQ "Case") <> travE e <> mconcat [ word8 t <> travE c | (t,c) <- rhs ] <> encodeMaybe travE def - Cond c t e -> intHost $(hashQ "Cond") <> travE c <> travE t <> travE e - While p f x -> intHost $(hashQ "While") <> travF p <> travF f <> travE x - PrimApp f x -> intHost $(hashQ "PrimApp") <> encodePrimFun f <> travE x - PrimConst c -> intHost $(hashQ "PrimConst") <> encodePrimConst c - Index a ix -> intHost $(hashQ "Index") <> encodeArrayVar a <> travE ix - LinearIndex a ix -> intHost $(hashQ "LinearIndex") <> encodeArrayVar a <> travE ix - Shape a -> intHost $(hashQ "Shape") <> encodeArrayVar a - ShapeSize _ sh -> intHost $(hashQ "ShapeSize") <> travE sh - Foreign _ _ f e -> intHost $(hashQ "Foreign") <> travF f <> travE e - Coerce _ tp e -> intHost $(hashQ "Coerce") <> encodeScalarType tp <> travE e - -encodeArrayVar :: ArrayVar aenv a -> Builder -encodeArrayVar (Var repr v) = encodeArrayType repr <> encodeIdx v - -{-# INLINEABLE encodeOpenFun #-} -encodeOpenFun - :: OpenFun env aenv f - -> Builder -encodeOpenFun (Body b) = intHost $(hashQ "Body") <> encodeOpenExp b -encodeOpenFun (Lam lhs l) = intHost $(hashQ "Lam") <> encodeLeftHandSide encodeScalarType lhs <> encodeOpenFun l - +encodeBoundary _ Wrap = intHost $(hashQ ("Wrap" :: String)) +encodeBoundary _ Clamp = intHost $(hashQ ("Clamp" :: String)) +encodeBoundary _ Mirror = intHost $(hashQ ("Mirror" :: String)) +encodeBoundary tp (Constant c) = intHost $(hashQ ("Constant" :: String)) <> encodeConst tp c +encodeBoundary _ (Function f) = intHost $(hashQ ("Function" :: String)) <> encodeOpenFun f encodeConst :: TypeR t -> t -> Builder -encodeConst TupRunit () = intHost $(hashQ "nil") +encodeConst TupRunit () = intHost $(hashQ ("nil" :: String)) encodeConst (TupRsingle t) c = encodeScalarConst t c -encodeConst (TupRpair ta tb) (a,b) = intHost $(hashQ "pair") <> encodeConst ta a <> encodeConst tb b - -encodeScalarConst :: ScalarType t -> t -> Builder -encodeScalarConst (SingleScalarType t) = encodeSingleConst t -encodeScalarConst (VectorScalarType t) = encodeVectorConst t - -encodeSingleConst :: SingleType t -> t -> Builder -encodeSingleConst (NumSingleType t) = encodeNumConst t - -encodeVectorConst :: VectorType (Vec n t) -> Vec n t -> Builder -encodeVectorConst (VectorType n t) (Vec ba#) = intHost $(hashQ "Vec") <> intHost n <> encodeSingleType t <> shortByteString (SBS ba#) - -encodeNumConst :: NumType t -> t -> Builder -encodeNumConst (IntegralNumType t) = encodeIntegralConst t -encodeNumConst (FloatingNumType t) = encodeFloatingConst t - -encodeIntegralConst :: IntegralType t -> t -> Builder -encodeIntegralConst TypeInt{} x = intHost $(hashQ "Int") <> intHost x -encodeIntegralConst TypeInt8{} x = intHost $(hashQ "Int8") <> int8 x -encodeIntegralConst TypeInt16{} x = intHost $(hashQ "Int16") <> int16Host x -encodeIntegralConst TypeInt32{} x = intHost $(hashQ "Int32") <> int32Host x -encodeIntegralConst TypeInt64{} x = intHost $(hashQ "Int64") <> int64Host x -encodeIntegralConst TypeWord{} x = intHost $(hashQ "Word") <> wordHost x -encodeIntegralConst TypeWord8{} x = intHost $(hashQ "Word8") <> word8 x -encodeIntegralConst TypeWord16{} x = intHost $(hashQ "Word16") <> word16Host x -encodeIntegralConst TypeWord32{} x = intHost $(hashQ "Word32") <> word32Host x -encodeIntegralConst TypeWord64{} x = intHost $(hashQ "Word64") <> word64Host x - -encodeFloatingConst :: FloatingType t -> t -> Builder -encodeFloatingConst TypeHalf{} (Half (CUShort x)) = intHost $(hashQ "Half") <> word16Host x -encodeFloatingConst TypeFloat{} x = intHost $(hashQ "Float") <> floatHost x -encodeFloatingConst TypeDouble{} x = intHost $(hashQ "Double") <> doubleHost x - -encodePrimConst :: PrimConst c -> Builder -encodePrimConst (PrimMinBound t) = intHost $(hashQ "PrimMinBound") <> encodeBoundedType t -encodePrimConst (PrimMaxBound t) = intHost $(hashQ "PrimMaxBound") <> encodeBoundedType t -encodePrimConst (PrimPi t) = intHost $(hashQ "PrimPi") <> encodeFloatingType t - -encodePrimFun :: PrimFun f -> Builder -encodePrimFun (PrimAdd a) = intHost $(hashQ "PrimAdd") <> encodeNumType a -encodePrimFun (PrimSub a) = intHost $(hashQ "PrimSub") <> encodeNumType a -encodePrimFun (PrimMul a) = intHost $(hashQ "PrimMul") <> encodeNumType a -encodePrimFun (PrimNeg a) = intHost $(hashQ "PrimNeg") <> encodeNumType a -encodePrimFun (PrimAbs a) = intHost $(hashQ "PrimAbs") <> encodeNumType a -encodePrimFun (PrimSig a) = intHost $(hashQ "PrimSig") <> encodeNumType a -encodePrimFun (PrimQuot a) = intHost $(hashQ "PrimQuot") <> encodeIntegralType a -encodePrimFun (PrimRem a) = intHost $(hashQ "PrimRem") <> encodeIntegralType a -encodePrimFun (PrimQuotRem a) = intHost $(hashQ "PrimQuotRem") <> encodeIntegralType a -encodePrimFun (PrimIDiv a) = intHost $(hashQ "PrimIDiv") <> encodeIntegralType a -encodePrimFun (PrimMod a) = intHost $(hashQ "PrimMod") <> encodeIntegralType a -encodePrimFun (PrimDivMod a) = intHost $(hashQ "PrimDivMod") <> encodeIntegralType a -encodePrimFun (PrimBAnd a) = intHost $(hashQ "PrimBAnd") <> encodeIntegralType a -encodePrimFun (PrimBOr a) = intHost $(hashQ "PrimBOr") <> encodeIntegralType a -encodePrimFun (PrimBXor a) = intHost $(hashQ "PrimBXor") <> encodeIntegralType a -encodePrimFun (PrimBNot a) = intHost $(hashQ "PrimBNot") <> encodeIntegralType a -encodePrimFun (PrimBShiftL a) = intHost $(hashQ "PrimBShiftL") <> encodeIntegralType a -encodePrimFun (PrimBShiftR a) = intHost $(hashQ "PrimBShiftR") <> encodeIntegralType a -encodePrimFun (PrimBRotateL a) = intHost $(hashQ "PrimBRotateL") <> encodeIntegralType a -encodePrimFun (PrimBRotateR a) = intHost $(hashQ "PrimBRotateR") <> encodeIntegralType a -encodePrimFun (PrimPopCount a) = intHost $(hashQ "PrimPopCount") <> encodeIntegralType a -encodePrimFun (PrimCountLeadingZeros a) = intHost $(hashQ "PrimCountLeadingZeros") <> encodeIntegralType a -encodePrimFun (PrimCountTrailingZeros a) = intHost $(hashQ "PrimCountTrailingZeros") <> encodeIntegralType a -encodePrimFun (PrimFDiv a) = intHost $(hashQ "PrimFDiv") <> encodeFloatingType a -encodePrimFun (PrimRecip a) = intHost $(hashQ "PrimRecip") <> encodeFloatingType a -encodePrimFun (PrimSin a) = intHost $(hashQ "PrimSin") <> encodeFloatingType a -encodePrimFun (PrimCos a) = intHost $(hashQ "PrimCos") <> encodeFloatingType a -encodePrimFun (PrimTan a) = intHost $(hashQ "PrimTan") <> encodeFloatingType a -encodePrimFun (PrimAsin a) = intHost $(hashQ "PrimAsin") <> encodeFloatingType a -encodePrimFun (PrimAcos a) = intHost $(hashQ "PrimAcos") <> encodeFloatingType a -encodePrimFun (PrimAtan a) = intHost $(hashQ "PrimAtan") <> encodeFloatingType a -encodePrimFun (PrimSinh a) = intHost $(hashQ "PrimSinh") <> encodeFloatingType a -encodePrimFun (PrimCosh a) = intHost $(hashQ "PrimCosh") <> encodeFloatingType a -encodePrimFun (PrimTanh a) = intHost $(hashQ "PrimTanh") <> encodeFloatingType a -encodePrimFun (PrimAsinh a) = intHost $(hashQ "PrimAsinh") <> encodeFloatingType a -encodePrimFun (PrimAcosh a) = intHost $(hashQ "PrimAcosh") <> encodeFloatingType a -encodePrimFun (PrimAtanh a) = intHost $(hashQ "PrimAtanh") <> encodeFloatingType a -encodePrimFun (PrimExpFloating a) = intHost $(hashQ "PrimExpFloating") <> encodeFloatingType a -encodePrimFun (PrimSqrt a) = intHost $(hashQ "PrimSqrt") <> encodeFloatingType a -encodePrimFun (PrimLog a) = intHost $(hashQ "PrimLog") <> encodeFloatingType a -encodePrimFun (PrimFPow a) = intHost $(hashQ "PrimFPow") <> encodeFloatingType a -encodePrimFun (PrimLogBase a) = intHost $(hashQ "PrimLogBase") <> encodeFloatingType a -encodePrimFun (PrimAtan2 a) = intHost $(hashQ "PrimAtan2") <> encodeFloatingType a -encodePrimFun (PrimTruncate a b) = intHost $(hashQ "PrimTruncate") <> encodeFloatingType a <> encodeIntegralType b -encodePrimFun (PrimRound a b) = intHost $(hashQ "PrimRound") <> encodeFloatingType a <> encodeIntegralType b -encodePrimFun (PrimFloor a b) = intHost $(hashQ "PrimFloor") <> encodeFloatingType a <> encodeIntegralType b -encodePrimFun (PrimCeiling a b) = intHost $(hashQ "PrimCeiling") <> encodeFloatingType a <> encodeIntegralType b -encodePrimFun (PrimIsNaN a) = intHost $(hashQ "PrimIsNaN") <> encodeFloatingType a -encodePrimFun (PrimIsInfinite a) = intHost $(hashQ "PrimIsInfinite") <> encodeFloatingType a -encodePrimFun (PrimLt a) = intHost $(hashQ "PrimLt") <> encodeSingleType a -encodePrimFun (PrimGt a) = intHost $(hashQ "PrimGt") <> encodeSingleType a -encodePrimFun (PrimLtEq a) = intHost $(hashQ "PrimLtEq") <> encodeSingleType a -encodePrimFun (PrimGtEq a) = intHost $(hashQ "PrimGtEq") <> encodeSingleType a -encodePrimFun (PrimEq a) = intHost $(hashQ "PrimEq") <> encodeSingleType a -encodePrimFun (PrimNEq a) = intHost $(hashQ "PrimNEq") <> encodeSingleType a -encodePrimFun (PrimMax a) = intHost $(hashQ "PrimMax") <> encodeSingleType a -encodePrimFun (PrimMin a) = intHost $(hashQ "PrimMin") <> encodeSingleType a -encodePrimFun (PrimFromIntegral a b) = intHost $(hashQ "PrimFromIntegral") <> encodeIntegralType a <> encodeNumType b -encodePrimFun (PrimToFloating a b) = intHost $(hashQ "PrimToFloating") <> encodeNumType a <> encodeFloatingType b -encodePrimFun PrimLAnd = intHost $(hashQ "PrimLAnd") -encodePrimFun PrimLOr = intHost $(hashQ "PrimLOr") -encodePrimFun PrimLNot = intHost $(hashQ "PrimLNot") - - -encodeTypeR :: TypeR t -> Builder -encodeTypeR TupRunit = intHost $(hashQ "TupRunit") -encodeTypeR (TupRsingle t) = intHost $(hashQ "TupRsingle") <> encodeScalarType t -encodeTypeR (TupRpair a b) = intHost $(hashQ "TupRpair") <> encodeTypeR a <> intHost (depthTypeR a) - <> encodeTypeR b <> intHost (depthTypeR b) - -depthTypeR :: TypeR t -> Int -depthTypeR TupRunit = 0 -depthTypeR TupRsingle{} = 1 -depthTypeR (TupRpair a b) = depthTypeR a + depthTypeR b - -encodeScalarType :: ScalarType t -> Builder -encodeScalarType (SingleScalarType t) = intHost $(hashQ "SingleScalarType") <> encodeSingleType t -encodeScalarType (VectorScalarType t) = intHost $(hashQ "VectorScalarType") <> encodeVectorType t - -encodeSingleType :: SingleType t -> Builder -encodeSingleType (NumSingleType t) = intHost $(hashQ "NumSingleType") <> encodeNumType t - -encodeVectorType :: VectorType (Vec n t) -> Builder -encodeVectorType (VectorType n t) = intHost $(hashQ "VectorType") <> intHost n <> encodeSingleType t - -encodeBoundedType :: BoundedType t -> Builder -encodeBoundedType (IntegralBoundedType t) = intHost $(hashQ "IntegralBoundedType") <> encodeIntegralType t - -encodeNumType :: NumType t -> Builder -encodeNumType (IntegralNumType t) = intHost $(hashQ "IntegralNumType") <> encodeIntegralType t -encodeNumType (FloatingNumType t) = intHost $(hashQ "FloatingNumType") <> encodeFloatingType t - -encodeIntegralType :: IntegralType t -> Builder -encodeIntegralType TypeInt{} = intHost $(hashQ "Int") -encodeIntegralType TypeInt8{} = intHost $(hashQ "Int8") -encodeIntegralType TypeInt16{} = intHost $(hashQ "Int16") -encodeIntegralType TypeInt32{} = intHost $(hashQ "Int32") -encodeIntegralType TypeInt64{} = intHost $(hashQ "Int64") -encodeIntegralType TypeWord{} = intHost $(hashQ "Word") -encodeIntegralType TypeWord8{} = intHost $(hashQ "Word8") -encodeIntegralType TypeWord16{} = intHost $(hashQ "Word16") -encodeIntegralType TypeWord32{} = intHost $(hashQ "Word32") -encodeIntegralType TypeWord64{} = intHost $(hashQ "Word64") - -encodeFloatingType :: FloatingType t -> Builder -encodeFloatingType TypeHalf{} = intHost $(hashQ "Half") -encodeFloatingType TypeFloat{} = intHost $(hashQ "Float") -encodeFloatingType TypeDouble{} = intHost $(hashQ "Double") - -encodeMaybe :: (a -> Builder) -> Maybe a -> Builder -encodeMaybe _ Nothing = intHost $(hashQ "Nothing") -encodeMaybe f (Just x) = intHost $(hashQ "Just") <> f x - +encodeConst (TupRpair ta tb) (a,b) = intHost $(hashQ ("pair" :: String)) <> encodeConst ta a <> encodeConst tb b diff --git a/src/Data/Array/Accelerate/Analysis/Hash/Exp.hs b/src/Data/Array/Accelerate/Analysis/Hash/Exp.hs new file mode 100644 index 000000000..e79837709 --- /dev/null +++ b/src/Data/Array/Accelerate/Analysis/Hash/Exp.hs @@ -0,0 +1,339 @@ +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE TemplateHaskell #-} +{-# OPTIONS_HADDOCK hide #-} +-- | +-- Module : Data.Array.Accelerate.Analysis.Hash.Exp +-- Copyright : [2017..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.Analysis.Hash.Exp ( + Hash, + HashOptions(..), defaultHashOptions, + hashOpenFun, hashOpenExp, + encodeOpenExp, + encodeOpenFun, + + encodeExpVar, + encodeLeftHandSide, + encodeTupR, + encodeArraysType, + encodeArrayType, + encodeIdx, + encodeIdxSet, + encodeShapeR, + encodeScalarType, + encodeScalarConst, + encodeTypeR, + encodeSliceIndex, + encodeMaybe, + encodeIntegralType, + hashQ, + Builder, + intHost, + +) where + +import Data.Array.Accelerate.AST.Exp +import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.IdxSet (IdxSet) +import qualified Data.Array.Accelerate.AST.IdxSet as IdxSet +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.AST.Var +import Data.Array.Accelerate.Analysis.Hash.TH +import Data.Array.Accelerate.Representation.Array +import Data.Array.Accelerate.Representation.Shape +import Data.Array.Accelerate.Representation.Slice +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Type +import Data.Primitive.Vec + +import Crypto.Hash.XKCP +import Data.ByteString.Builder +import Data.ByteString.Builder.Extra +import Data.ByteString.Short.Internal ( ShortByteString(..) ) +import qualified Data.Hashable as Hashable +import Data.Monoid +import Prelude hiding ( exp ) + +-- Hashing +-- ------- + +type Hash = SHA3_256 + +data HashOptions = HashOptions + { perfect :: Bool + -- ^ Should the hash function include _all_ substructure, recursively? + -- + -- Set to true (the default) if you want a truly unique fingerprint for + -- the entire expression: + -- + -- Example: + -- + -- xs, ys :: Acc (Vector Float) + -- xs = fill (constant (Z:.10)) 1.0 + -- ys = fill (constant (Z:.20)) 1.0 + -- + -- with perfect=True: + -- + -- hash xs = 2e1f91aca4c476d13b36f22462e73c15bbdd9fcacb0d4996280f6004058e9732 + -- hash ys = 2fce5c849b6c652192b09aaeafdc8029e57b9f006c1ecd79ccf9114f349aaf9e + -- + -- However, for a code generating backend the object code used to + -- evaluate both of these expressions is likely to be identical. + -- + -- Setting perfect=False results in: + -- + -- hash xs = hash ys = f97944b0ec64ab8aa989fd60c8b50e7ec3eff759d22d2b340039d837d74dfc3c + -- + -- Note that to be useful the provided 'EncodeAcc' function must also + -- understand this option, and the consumer of the hash value must be + -- agnostic to the elided details. + } + deriving Show + +defaultHashOptions :: HashOptions +defaultHashOptions = HashOptions True + + +{-# INLINEABLE hashOpenFun #-} +hashOpenFun :: IsArrayInstr arr => PreOpenFun arr env f -> Hash +hashOpenFun + = hashlazy + . toLazyByteString + . encodeOpenFun + +-- {-# INLINEABLE hashOpenExp #-} throws warnings +hashOpenExp :: IsArrayInstr arr => PreOpenExp arr env t -> Hash +hashOpenExp + = hashlazy + . toLazyByteString + . encodeOpenExp + +encodeIdx :: Idx env t -> Builder +encodeIdx = intHost . idxToInt + +encodeIdxSet :: IdxSet env -> Builder +encodeIdxSet set = + intHost (length list) + <> mconcat (map (\(Exists idx) -> encodeIdx idx) list) + where + list = IdxSet.toList set + +encodeTupR :: (forall b. s b -> Builder) -> TupR s a -> Builder +encodeTupR _ TupRunit = intHost $(hashQ ("TupRunit" :: String)) +encodeTupR f (TupRpair r1 r2) = intHost $(hashQ ("TupRpair" :: String)) <> encodeTupR f r1 <> encodeTupR f r2 +encodeTupR f (TupRsingle s) = intHost $(hashQ ("TupRsingle" :: String)) <> f s + +encodeLeftHandSide :: (forall b. s b -> Builder) -> LeftHandSide s a env env' -> Builder +encodeLeftHandSide f (LeftHandSideWildcard r) = intHost $(hashQ ("LeftHandSideWildcard" :: String)) <> encodeTupR f r +encodeLeftHandSide f (LeftHandSidePair r1 r2) = intHost $(hashQ ("LeftHandSidePair" :: String)) <> encodeLeftHandSide f r1 <> encodeLeftHandSide f r2 +encodeLeftHandSide f (LeftHandSideSingle s) = intHost $(hashQ ("LeftHandSideArray" :: String)) <> f s + +encodeArrayType :: ArrayR a -> Builder +encodeArrayType (ArrayR shr tp) = encodeShapeR shr <> encodeTypeR tp + +encodeArraysType :: ArraysR arrs -> Builder +encodeArraysType = encodeTupR encodeArrayType + +encodeShapeR :: ShapeR sh -> Builder +encodeShapeR = intHost . rank + +{-# INLINEABLE encodeOpenExp #-} +encodeOpenExp + :: forall arr env exp. + IsArrayInstr arr + => PreOpenExp arr env exp + -> Builder +encodeOpenExp exp = + let + travE :: forall env' e. PreOpenExp arr env' e -> Builder + travE e = encodeOpenExp e + + travF :: PreOpenFun arr env' f -> Builder + travF = encodeOpenFun + in + case exp of + Let lhs bnd body -> intHost $(hashQ ("Let" :: String)) <> encodeLeftHandSide encodeScalarType lhs <> travE bnd <> travE body + Evar var -> intHost $(hashQ ("Evar" :: String)) <> encodeExpVar var + Nil -> intHost $(hashQ ("Nil" :: String)) + Pair e1 e2 -> intHost $(hashQ ("Pair" :: String)) <> travE e1 <> travE e2 + VecPack _ e -> intHost $(hashQ ("VecPack" :: String)) <> travE e + VecUnpack _ e -> intHost $(hashQ ("VecUnpack" :: String)) <> travE e + Const tp c -> intHost $(hashQ ("Const" :: String)) <> encodeScalarConst tp c + Undef tp -> intHost $(hashQ ("Undef" :: String)) <> encodeScalarType tp + ToIndex _ sh i -> intHost $(hashQ ("ToIndex" :: String)) <> travE sh <> travE i + FromIndex _ sh i -> intHost $(hashQ ("FromIndex" :: String)) <> travE sh <> travE i + Case e rhs def -> intHost $(hashQ ("Case" :: String)) <> travE e <> mconcat [ word8 t <> travE c | (t,c) <- rhs ] <> encodeMaybe travE def + Cond c t e -> intHost $(hashQ ("Cond" :: String)) <> travE c <> travE t <> travE e + Select c t e -> intHost $(hashQ ("Select" :: String)) <> travE c <> travE t <> travE e + While p f x -> intHost $(hashQ ("While" :: String)) <> travF p <> travF f <> travE x + PrimApp f x -> intHost $(hashQ ("PrimApp" :: String)) <> encodePrimFun f <> travE x + ArrayInstr arr e -> intHost $(hashQ ("ArrayInstr" :: String)) <> encodeArrayInstr arr <> travE e + ShapeSize _ sh -> intHost $(hashQ ("ShapeSize" :: String)) <> travE sh + Foreign _ _ f e -> intHost $(hashQ ("Foreign" :: String)) <> encodeOpenFun f <> travE e + Coerce _ tp e -> intHost $(hashQ ("Coerce" :: String)) <> encodeScalarType tp <> travE e + Assert msg e1 e2 -> intHost $(hashQ ("Assert" :: String)) <> intHost (Hashable.hash msg) <> travE e1 <> travE e2 + Assume e1 e2 -> intHost $(hashQ ("Assume" :: String)) <> travE e1 <> travE e2 + +encodeExpVar :: ExpVar env t -> Builder +encodeExpVar (Var tp ix) = encodeScalarType tp <> encodeIdx ix + +{-# INLINEABLE encodeOpenFun #-} +encodeOpenFun + :: IsArrayInstr arr + => PreOpenFun arr env f + -> Builder +encodeOpenFun (Body b) = intHost $(hashQ ("Body" :: String)) <> encodeOpenExp b +encodeOpenFun (Lam lhs l) = intHost $(hashQ ("Lam" :: String)) <> encodeLeftHandSide encodeScalarType lhs <> encodeOpenFun l + +encodeScalarConst :: ScalarType t -> t -> Builder +encodeScalarConst (SingleScalarType t) = encodeSingleConst t +encodeScalarConst (VectorScalarType t) = encodeVectorConst t + +encodeSingleConst :: SingleType t -> t -> Builder +encodeSingleConst (NumSingleType t) = encodeNumConst t + +encodeVectorConst :: VectorType (Vec n t) -> Vec n t -> Builder +encodeVectorConst (VectorType n t) (Vec ba#) = intHost $(hashQ ("Vec" :: String)) <> intHost n <> encodeSingleType t <> shortByteString (SBS ba#) + +encodeNumConst :: NumType t -> t -> Builder +encodeNumConst (IntegralNumType t) = encodeIntegralConst t +encodeNumConst (FloatingNumType t) = encodeFloatingConst t + +encodeIntegralConst :: IntegralType t -> t -> Builder +encodeIntegralConst TypeInt{} x = intHost $(hashQ ("Int" :: String)) <> intHost x +encodeIntegralConst TypeInt8{} x = intHost $(hashQ ("Int8" :: String)) <> int8 x +encodeIntegralConst TypeInt16{} x = intHost $(hashQ ("Int16" :: String)) <> int16Host x +encodeIntegralConst TypeInt32{} x = intHost $(hashQ ("Int32" :: String)) <> int32Host x +encodeIntegralConst TypeInt64{} x = intHost $(hashQ ("Int64" :: String)) <> int64Host x +encodeIntegralConst TypeWord{} x = intHost $(hashQ ("Word" :: String)) <> wordHost x +encodeIntegralConst TypeWord8{} x = intHost $(hashQ ("Word8" :: String)) <> word8 x +encodeIntegralConst TypeWord16{} x = intHost $(hashQ ("Word16" :: String)) <> word16Host x +encodeIntegralConst TypeWord32{} x = intHost $(hashQ ("Word32" :: String)) <> word32Host x +encodeIntegralConst TypeWord64{} x = intHost $(hashQ ("Word64" :: String)) <> word64Host x + +encodeFloatingConst :: FloatingType t -> t -> Builder +encodeFloatingConst TypeHalf{} (Half (CUShort x)) = intHost $(hashQ ("Half" :: String)) <> word16Host x +encodeFloatingConst TypeFloat{} x = intHost $(hashQ ("Float" :: String)) <> floatHost x +encodeFloatingConst TypeDouble{} x = intHost $(hashQ ("Double" :: String)) <> doubleHost x + +encodePrimFun :: PrimFun f -> Builder +encodePrimFun (PrimAdd a) = intHost $(hashQ ("PrimAdd" :: String)) <> encodeNumType a +encodePrimFun (PrimSub a) = intHost $(hashQ ("PrimSub" :: String)) <> encodeNumType a +encodePrimFun (PrimMul a) = intHost $(hashQ ("PrimMul" :: String)) <> encodeNumType a +encodePrimFun (PrimNeg a) = intHost $(hashQ ("PrimNeg" :: String)) <> encodeNumType a +encodePrimFun (PrimAbs a) = intHost $(hashQ ("PrimAbs" :: String)) <> encodeNumType a +encodePrimFun (PrimSig a) = intHost $(hashQ ("PrimSig" :: String)) <> encodeNumType a +encodePrimFun (PrimQuot a) = intHost $(hashQ ("PrimQuot" :: String)) <> encodeIntegralType a +encodePrimFun (PrimRem a) = intHost $(hashQ ("PrimRem" :: String)) <> encodeIntegralType a +encodePrimFun (PrimQuotRem a) = intHost $(hashQ ("PrimQuotRem" :: String)) <> encodeIntegralType a +encodePrimFun (PrimIDiv a) = intHost $(hashQ ("PrimIDiv" :: String)) <> encodeIntegralType a +encodePrimFun (PrimMod a) = intHost $(hashQ ("PrimMod" :: String)) <> encodeIntegralType a +encodePrimFun (PrimDivMod a) = intHost $(hashQ ("PrimDivMod" :: String)) <> encodeIntegralType a +encodePrimFun (PrimBAnd a) = intHost $(hashQ ("PrimBAnd" :: String)) <> encodeIntegralType a +encodePrimFun (PrimBOr a) = intHost $(hashQ ("PrimBOr" :: String)) <> encodeIntegralType a +encodePrimFun (PrimBXor a) = intHost $(hashQ ("PrimBXor" :: String)) <> encodeIntegralType a +encodePrimFun (PrimBNot a) = intHost $(hashQ ("PrimBNot" :: String)) <> encodeIntegralType a +encodePrimFun (PrimBShiftL a) = intHost $(hashQ ("PrimBShiftL" :: String)) <> encodeIntegralType a +encodePrimFun (PrimBShiftR a) = intHost $(hashQ ("PrimBShiftR" :: String)) <> encodeIntegralType a +encodePrimFun (PrimBRotateL a) = intHost $(hashQ ("PrimBRotateL" :: String)) <> encodeIntegralType a +encodePrimFun (PrimBRotateR a) = intHost $(hashQ ("PrimBRotateR" :: String)) <> encodeIntegralType a +encodePrimFun (PrimPopCount a) = intHost $(hashQ ("PrimPopCount" :: String)) <> encodeIntegralType a +encodePrimFun (PrimCountLeadingZeros a) = intHost $(hashQ ("PrimCountLeadingZeros" :: String)) <> encodeIntegralType a +encodePrimFun (PrimCountTrailingZeros a) = intHost $(hashQ ("PrimCountTrailingZeros" :: String)) <> encodeIntegralType a +encodePrimFun (PrimFDiv a) = intHost $(hashQ ("PrimFDiv" :: String)) <> encodeFloatingType a +encodePrimFun (PrimRecip a) = intHost $(hashQ ("PrimRecip" :: String)) <> encodeFloatingType a +encodePrimFun (PrimSin a) = intHost $(hashQ ("PrimSin" :: String)) <> encodeFloatingType a +encodePrimFun (PrimCos a) = intHost $(hashQ ("PrimCos" :: String)) <> encodeFloatingType a +encodePrimFun (PrimTan a) = intHost $(hashQ ("PrimTan" :: String)) <> encodeFloatingType a +encodePrimFun (PrimAsin a) = intHost $(hashQ ("PrimAsin" :: String)) <> encodeFloatingType a +encodePrimFun (PrimAcos a) = intHost $(hashQ ("PrimAcos" :: String)) <> encodeFloatingType a +encodePrimFun (PrimAtan a) = intHost $(hashQ ("PrimAtan" :: String)) <> encodeFloatingType a +encodePrimFun (PrimSinh a) = intHost $(hashQ ("PrimSinh" :: String)) <> encodeFloatingType a +encodePrimFun (PrimCosh a) = intHost $(hashQ ("PrimCosh" :: String)) <> encodeFloatingType a +encodePrimFun (PrimTanh a) = intHost $(hashQ ("PrimTanh" :: String)) <> encodeFloatingType a +encodePrimFun (PrimAsinh a) = intHost $(hashQ ("PrimAsinh" :: String)) <> encodeFloatingType a +encodePrimFun (PrimAcosh a) = intHost $(hashQ ("PrimAcosh" :: String)) <> encodeFloatingType a +encodePrimFun (PrimAtanh a) = intHost $(hashQ ("PrimAtanh" :: String)) <> encodeFloatingType a +encodePrimFun (PrimExpFloating a) = intHost $(hashQ ("PrimExpFloating" :: String)) <> encodeFloatingType a +encodePrimFun (PrimSqrt a) = intHost $(hashQ ("PrimSqrt" :: String)) <> encodeFloatingType a +encodePrimFun (PrimLog a) = intHost $(hashQ ("PrimLog" :: String)) <> encodeFloatingType a +encodePrimFun (PrimFPow a) = intHost $(hashQ ("PrimFPow" :: String)) <> encodeFloatingType a +encodePrimFun (PrimLogBase a) = intHost $(hashQ ("PrimLogBase" :: String)) <> encodeFloatingType a +encodePrimFun (PrimAtan2 a) = intHost $(hashQ ("PrimAtan2" :: String)) <> encodeFloatingType a +encodePrimFun (PrimTruncate a b) = intHost $(hashQ ("PrimTruncate" :: String)) <> encodeFloatingType a <> encodeIntegralType b +encodePrimFun (PrimRound a b) = intHost $(hashQ ("PrimRound" :: String)) <> encodeFloatingType a <> encodeIntegralType b +encodePrimFun (PrimFloor a b) = intHost $(hashQ ("PrimFloor" :: String)) <> encodeFloatingType a <> encodeIntegralType b +encodePrimFun (PrimCeiling a b) = intHost $(hashQ ("PrimCeiling" :: String)) <> encodeFloatingType a <> encodeIntegralType b +encodePrimFun (PrimIsNaN a) = intHost $(hashQ ("PrimIsNaN" :: String)) <> encodeFloatingType a +encodePrimFun (PrimIsInfinite a) = intHost $(hashQ ("PrimIsInfinite" :: String)) <> encodeFloatingType a +encodePrimFun (PrimCmp a c) = intHost $(hashQ ("PrimCmp" :: String)) <> encodeSingleType a <> encodeCmp c +encodePrimFun (PrimMax a) = intHost $(hashQ ("PrimMax" :: String)) <> encodeSingleType a +encodePrimFun (PrimMin a) = intHost $(hashQ ("PrimMin" :: String)) <> encodeSingleType a +encodePrimFun (PrimFromIntegral a b) = intHost $(hashQ ("PrimFromIntegral" :: String)) <> encodeIntegralType a <> encodeNumType b +encodePrimFun (PrimToFloating a b) = intHost $(hashQ ("PrimToFloating" :: String)) <> encodeNumType a <> encodeFloatingType b +encodePrimFun PrimLAnd = intHost $(hashQ ("PrimLAnd" :: String)) +encodePrimFun PrimLOr = intHost $(hashQ ("PrimLOr" :: String)) +encodePrimFun PrimLNot = intHost $(hashQ ("PrimLNot" :: String)) + +encodeCmp :: Cmp -> Builder +encodeCmp CmpLt = intHost $(hashQ ("CmpLt" :: String)) +encodeCmp CmpGtEq = intHost $(hashQ ("CmpGtEq" :: String)) +encodeCmp CmpEq = intHost $(hashQ ("CmpEq" :: String)) +encodeCmp CmpNEq = intHost $(hashQ ("CmpNEq" :: String)) + +encodeTypeR :: TypeR t -> Builder +encodeTypeR TupRunit = intHost $(hashQ ("TupRunit" :: String)) +encodeTypeR (TupRsingle t) = intHost $(hashQ ("TupRsingle" :: String)) <> encodeScalarType t +encodeTypeR (TupRpair a b) = intHost $(hashQ ("TupRpair" :: String)) <> encodeTypeR a <> intHost (depthTypeR a) + <> encodeTypeR b <> intHost (depthTypeR b) + +depthTypeR :: TypeR t -> Int +depthTypeR TupRunit = 0 +depthTypeR TupRsingle{} = 1 +depthTypeR (TupRpair a b) = depthTypeR a + depthTypeR b + +encodeScalarType :: ScalarType t -> Builder +encodeScalarType (SingleScalarType t) = intHost $(hashQ ("SingleScalarType" :: String)) <> encodeSingleType t +encodeScalarType (VectorScalarType t) = intHost $(hashQ ("VectorScalarType" :: String)) <> encodeVectorType t + +encodeSingleType :: SingleType t -> Builder +encodeSingleType (NumSingleType t) = intHost $(hashQ ("NumSingleType" :: String)) <> encodeNumType t + +encodeVectorType :: VectorType (Vec n t) -> Builder +encodeVectorType (VectorType n t) = intHost $(hashQ ("VectorType" :: String)) <> intHost n <> encodeSingleType t + +encodeNumType :: NumType t -> Builder +encodeNumType (IntegralNumType t) = intHost $(hashQ ("IntegralNumType" :: String)) <> encodeIntegralType t +encodeNumType (FloatingNumType t) = intHost $(hashQ ("FloatingNumType" :: String)) <> encodeFloatingType t + +encodeIntegralType :: IntegralType t -> Builder +encodeIntegralType TypeInt{} = intHost $(hashQ ("Int" :: String)) +encodeIntegralType TypeInt8{} = intHost $(hashQ ("Int8" :: String)) +encodeIntegralType TypeInt16{} = intHost $(hashQ ("Int16" :: String)) +encodeIntegralType TypeInt32{} = intHost $(hashQ ("Int32" :: String)) +encodeIntegralType TypeInt64{} = intHost $(hashQ ("Int64" :: String)) +encodeIntegralType TypeWord{} = intHost $(hashQ ("Word" :: String)) +encodeIntegralType TypeWord8{} = intHost $(hashQ ("Word8" :: String)) +encodeIntegralType TypeWord16{} = intHost $(hashQ ("Word16" :: String)) +encodeIntegralType TypeWord32{} = intHost $(hashQ ("Word32" :: String)) +encodeIntegralType TypeWord64{} = intHost $(hashQ ("Word64" :: String)) + +encodeFloatingType :: FloatingType t -> Builder +encodeFloatingType TypeHalf{} = intHost $(hashQ ("Half" :: String)) +encodeFloatingType TypeFloat{} = intHost $(hashQ ("Float" :: String)) +encodeFloatingType TypeDouble{} = intHost $(hashQ ("Double" :: String)) + +encodeSliceIndex :: SliceIndex slix sl co sh -> Builder +encodeSliceIndex SliceNil = intHost $(hashQ ("SliceNil" :: String)) +encodeSliceIndex (SliceAll r) = intHost $(hashQ ("SliceAll" :: String)) <> encodeSliceIndex r +encodeSliceIndex (SliceFixed r) = intHost $(hashQ ("sliceFixed" :: String)) <> encodeSliceIndex r + +encodeMaybe :: (a -> Builder) -> Maybe a -> Builder +encodeMaybe _ Nothing = intHost $(hashQ ("Nothing" :: String)) +encodeMaybe f (Just x) = intHost $(hashQ ("Just" :: String)) <> f x diff --git a/src/Data/Array/Accelerate/Analysis/Hash/Operation.hs b/src/Data/Array/Accelerate/Analysis/Hash/Operation.hs new file mode 100644 index 000000000..29cb53f36 --- /dev/null +++ b/src/Data/Array/Accelerate/Analysis/Hash/Operation.hs @@ -0,0 +1,93 @@ +{-# LANGUAGE TemplateHaskell #-} +{-# OPTIONS_HADDOCK hide #-} +-- | +-- Module : Data.Array.Accelerate.Analysis.Hash +-- Copyright : [2017..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.Analysis.Hash.Operation ( EncodeOperation(..), hashOperation, encodePreArgs, encodeArg ) where + +import Data.Array.Accelerate.Analysis.Hash.Exp +import Data.Array.Accelerate.AST.Partitioned +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Graph (MakesILP, encodeBackendClusterArg) + +import Crypto.Hash.XKCP +import Data.ByteString.Builder + +hashOperation :: EncodeOperation op => op t -> Args env t -> Hash +hashOperation op args + = hashlazy + $ toLazyByteString + $ encodeOperation op <> encodePreArgs encodeArg args + +class EncodeOperation op where + encodeOperation :: op t -> Builder + +instance (MakesILP op, EncodeOperation op) => EncodeOperation (Clustered op) where + encodeOperation (Clustered cluster backendCluster) = + encodePreArgs encodeBackendClusterArg backendCluster <> encodeCluster cluster + +encodePreArgs :: (forall t. arg t -> Builder) -> PreArgs arg f -> Builder +encodePreArgs f (a :>: as) = intHost $(hashQ (":>:" :: String)) <> f a <> encodePreArgs f as +encodePreArgs _f ArgsNil = intHost $(hashQ ("ArgsNil" :: String)) + +encodeArg :: Arg env t -> Builder +encodeArg (ArgArray m repr sh buffers) + = intHost $(hashQ ("Array" :: String)) <> encodeModifier m <> encodeArrayType repr <> encodeGroundVars sh <> encodeGroundVars buffers +encodeArg (ArgVar var) = intHost $(hashQ ("Var" :: String)) <> encodeTupR encodeExpVar var +encodeArg (ArgExp exp') = intHost $(hashQ ("Exp" :: String)) <> encodeOpenExp exp' +encodeArg (ArgFun fun) = intHost $(hashQ ("Fun" :: String)) <> encodeOpenFun fun + +encodeModifier :: Modifier m -> Builder +encodeModifier In = intHost $(hashQ ("In" :: String)) +encodeModifier Out = intHost $(hashQ ("Out" :: String)) +encodeModifier Mut = intHost $(hashQ ("Mut" :: String)) + +encodeCluster :: EncodeOperation op => Cluster op args -> Builder +encodeCluster (SingleOp op label) + = intHost $(hashQ ("SingleOp" :: String)) <> encodeSingleOp op <> encodeLabel label +encodeCluster (Fused fusion clusterL clusterR) + = intHost $(hashQ ("Fused" :: String)) <> encodeFusion fusion <> encodeCluster clusterL <> encodeCluster clusterR + +encodeLabel :: Node Comp -> Builder +encodeLabel (Node idx Nothing) = intHost idx <> intHost $(hashQ ("Nothing" :: String)) +encodeLabel (Node idx (Just l)) = intHost idx <> intHost $(hashQ ("Just" :: String)) <> encodeLabel l + +encodeFusion :: Fusion largs rargs args -> Builder +encodeFusion EmptyF = intHost $(hashQ ("EmptyF" :: String)) +encodeFusion (Vertical r f) = intHost $(hashQ ("Vertical" :: String)) <> encodeArrayType r <> encodeFusion f +encodeFusion (Horizontal f) = intHost $(hashQ ("Horizontal" :: String)) <> encodeFusion f +encodeFusion (Diagonal f) = intHost $(hashQ ("Diagonal" :: String)) <> encodeFusion f +encodeFusion (IntroI1 f) = intHost $(hashQ ("IntroI1" :: String)) <> encodeFusion f +encodeFusion (IntroI2 f) = intHost $(hashQ ("IntroI2" :: String)) <> encodeFusion f +encodeFusion (IntroO1 f) = intHost $(hashQ ("IntroO1" :: String)) <> encodeFusion f +encodeFusion (IntroO2 f) = intHost $(hashQ ("IntroO2" :: String)) <> encodeFusion f +encodeFusion (IntroL f) = intHost $(hashQ ("IntroL" :: String)) <> encodeFusion f +encodeFusion (IntroR f) = intHost $(hashQ ("IntroR" :: String)) <> encodeFusion f + +encodeSingleOp :: EncodeOperation op => SingleOp op args -> Builder +encodeSingleOp (Single op args) + = encodeOperation op + <> encodeClusterArgs args + +encodeClusterArgs :: ClusterArgs env f -> Builder +encodeClusterArgs = encodePreArgs encodeClusterArg + +encodeClusterArg :: ClusterArg env t -> Builder +encodeClusterArg (ClusterArgSingle idx) + = intHost $(hashQ ("ClusterArgSingle" :: String)) <> encodeIdx idx +encodeClusterArg (ClusterArgArray m _ _ buffers) + = intHost $(hashQ ("ClusterArgArray" :: String)) <> encodeModifier m <> encodeClusterArgBuffers buffers + +encodeClusterArgBuffers :: ClusterArgBuffers env m sh t -> Builder +encodeClusterArgBuffers (ClusterArgBuffersPair l r) + = intHost $(hashQ ("ClusterArgBuffersPair" :: String)) <> encodeClusterArgBuffers l <> encodeClusterArgBuffers r +encodeClusterArgBuffers (ClusterArgBuffersDead _ idx) + = intHost $(hashQ ("ClusterArgBuffersDead" :: String)) <> encodeIdx idx +encodeClusterArgBuffers (ClusterArgBuffersLive _ idx) + = intHost $(hashQ ("ClusterArgBuffersLive" :: String)) <> encodeIdx idx diff --git a/src/Data/Array/Accelerate/Analysis/Hash/Schedule/Uniform.hs b/src/Data/Array/Accelerate/Analysis/Hash/Schedule/Uniform.hs new file mode 100644 index 000000000..65f026a09 --- /dev/null +++ b/src/Data/Array/Accelerate/Analysis/Hash/Schedule/Uniform.hs @@ -0,0 +1,154 @@ +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE TemplateHaskell #-} +{-# OPTIONS_HADDOCK hide #-} +-- | +-- Module : Data.Array.Accelerate.Analysis.Hash.Exp +-- Copyright : [2017..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.Analysis.Hash.Schedule.Uniform ( + hashUniformScheduleFun +) where + +import Data.Array.Accelerate.AST.Var +import Data.Array.Accelerate.AST.Kernel +import Data.Array.Accelerate.AST.Schedule.Uniform +import Data.Array.Accelerate.Analysis.Hash.TH +import Data.Array.Accelerate.Analysis.Hash.Exp +import Data.Array.Accelerate.Analysis.Hash.Operation (encodePreArgs) + +import Crypto.Hash.XKCP +import Data.ByteString.Builder +import Data.ByteString.Short.Internal ( ShortByteString(..) ) +import qualified Data.Hashable as Hashable + +hashUniformScheduleFun :: IsKernel kernel => UniformScheduleFun kernel env f -> Hash +hashUniformScheduleFun = hashlazy . toLazyByteString . encodeUniformScheduleFun + +encodeUniformScheduleFun :: IsKernel kernel => UniformScheduleFun kernel env f -> Builder +encodeUniformScheduleFun (Slam lhs f) = intHost $(hashQ ("Slam" :: String)) <> encodeBLeftHandSide lhs <> encodeUniformScheduleFun f +encodeUniformScheduleFun (Sbody body) = intHost $(hashQ ("Sbody" :: String)) <> encodeUniformSchedule body + +encodeUniformSchedule :: IsKernel kernel => UniformSchedule kernel env -> Builder +encodeUniformSchedule Return = intHost $(hashQ ("Return" :: String)) +encodeUniformSchedule (Alet lhs bnd next) + = intHost $(hashQ ("Alet" :: String)) + <> encodeBLeftHandSide lhs + <> encodeBinding bnd + <> encodeUniformSchedule next +encodeUniformSchedule (Effect effect next) + = intHost $(hashQ ("Effect" :: String)) + <> encodeEffect effect + <> encodeUniformSchedule next +encodeUniformSchedule (Acond (Var _ idx) true false next) + = intHost $(hashQ ("Acond" :: String)) + <> encodeIdx idx + <> encodeUniformSchedule true + <> encodeUniformSchedule false + <> encodeUniformSchedule next +encodeUniformSchedule (Awhile io fn initial next) + = intHost $(hashQ ("Awhile" :: String)) + <> encodeIO io + <> encodeUniformScheduleFun fn + <> encodeTupR (\(Var _ idx) -> encodeIdx idx) initial + <> encodeUniformSchedule next +encodeUniformSchedule (AwhileSeq io fn initial next) + = intHost $(hashQ ("AwhileSeq" :: String)) + <> encodeIO io + <> encodeUniformScheduleFun fn + <> encodeTupR (\(Var _ idx) -> encodeIdx idx) initial + <> encodeUniformSchedule next +encodeUniformSchedule (Spawn a b) + = intHost $(hashQ ("Spawn" :: String)) + <> encodeUniformSchedule a + <> encodeUniformSchedule b + +encodeBLeftHandSide :: BLeftHandSide t env env' -> Builder +encodeBLeftHandSide = encodeLeftHandSide encodeBaseR + +encodeBaseR :: BaseR t -> Builder +encodeBaseR (BaseRground tp) = intHost $(hashQ ("Ground" :: String)) <> encodeGroundR tp +encodeBaseR BaseRsignal = intHost $(hashQ ("Signal" :: String)) +encodeBaseR BaseRsignalResolver = intHost $(hashQ ("SignalResolver" :: String)) +encodeBaseR (BaseRref tp) = intHost $(hashQ ("Ref" :: String)) <> encodeGroundR tp +encodeBaseR (BaseRrefWrite tp) = intHost $(hashQ ("RefWrite" :: String)) <> encodeGroundR tp + +{- TODO WALL: DEAD CODE +encodeBasesR :: BasesR t -> Builder +encodeBasesR = encodeTupR encodeBaseR +-} + +encodeBinding :: Binding env t -> Builder +encodeBinding = \case + Compute expr -> intHost $(hashQ ("Compute" :: String)) <> encodeOpenExp expr + NewSignal _ -> intHost $(hashQ ("NewSignal" :: String)) + NewRef tp -> intHost $(hashQ ("NewRef" :: String)) <> encodeGroundR tp + Alloc shr tp sh -> + intHost $(hashQ ("Alloc" :: String)) + <> encodeShapeR shr + <> encodeScalarType tp + <> encodeTupR (\(Var _ idx) -> encodeIdx idx) sh + -- Buffer is passed indirectly, via %imports_t in accelerate-llvm-native, + -- so the buffer/pointer does not need to included in the hashing. + Use tp _ _ -> intHost $(hashQ ("Use" :: String)) <> encodeScalarType tp + Unit (Var tp idx) -> intHost $(hashQ ("Unit" :: String)) <> encodeScalarType tp <> encodeIdx idx + RefRead (Var _ idx) -> intHost $(hashQ ("RefRead" :: String)) <> encodeIdx idx + +encodeEffect :: IsKernel kernel => Effect kernel env -> Builder +encodeEffect = \case + Exec _ kernel args -> encodeKernelFun kernel <> encodePreArgs encodeSArg args + SignalAwait indices -> + intHost $(hashQ ("SignalAwait" :: String)) + <> intHost (length indices) + <> mconcat (map encodeIdx indices) + SignalResolve indices -> + intHost $(hashQ ("SignalResolve" :: String)) + <> intHost (length indices) + <> mconcat (map encodeIdx indices) + RefWrite (Var _ ref) (Var _ var) -> + intHost $(hashQ ("RefWrite" :: String)) + <> encodeIdx ref + <> encodeIdx var + Aassert msg cond -> + intHost $(hashQ ("Aassert" :: String)) + <> intHost (Hashable.hash msg) + <> encodeOpenExp cond + Atrace msg t -> + intHost $(hashQ ("Atrace" :: String)) + <> intHost (Hashable.hash msg) + <> encodeArrayDescriptors t + +encodeIO :: InputOutputR input output -> Builder +encodeIO = \case + InputOutputRsignal -> intHost $(hashQ ("signal" :: String)) + InputOutputRref tp -> intHost $(hashQ ("ref" :: String)) <> encodeGroundR tp + InputOutputRpair a b -> intHost $(hashQ ("pair" :: String)) <> encodeIO a <> encodeIO b + InputOutputRunit -> intHost $(hashQ ("unit" :: String)) + +encodeSArg :: SArg env t -> Builder +encodeSArg (SArgScalar (Var tp idx)) = + intHost $(hashQ ("SArgScalar" :: String)) + <> encodeScalarType tp + <> encodeIdx idx +encodeSArg (SArgBuffer m (Var tp idx)) = + intHost $(hashQ ("SArgBuffer" :: String)) + <> m' + <> encodeGroundR tp + <> encodeIdx idx + where + m' = case m of + In -> intHost $(hashQ ("In" :: String)) + Out -> intHost $(hashQ ("Out" :: String)) + Mut -> intHost $(hashQ ("Mut" :: String)) + +encodeKernelFun :: IsKernel kernel => OpenKernelFun kernel env t -> Builder +-- Argument types are encoded via encodeSArg +encodeKernelFun (KernelFunLam _ f) = encodeKernelFun f +encodeKernelFun (KernelFunBody kernel) = case encodeKernel kernel of + Left (SHA3_256 ba#) -> shortByteString (SBS ba#) + Right builder -> builder diff --git a/src/Data/Array/Accelerate/Analysis/Match.hs b/src/Data/Array/Accelerate/Analysis/Match.hs index 98a0818a4..3f2ff396d 100644 --- a/src/Data/Array/Accelerate/Analysis/Match.hs +++ b/src/Data/Array/Accelerate/Analysis/Match.hs @@ -1,11 +1,4 @@ -{-# LANGUAGE AllowAmbiguousTypes #-} -{-# LANGUAGE CPP #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE PatternGuards #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE CPP #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.Analysis.Match @@ -29,9 +22,11 @@ module Data.Array.Accelerate.Analysis.Match ( matchPrimFun, matchPrimFun', -- auxiliary - matchIdx, matchVar, matchVars, matchArrayR, matchArraysR, matchTypeR, matchShapeR, - matchShapeType, matchIntegralType, matchFloatingType, matchNumType, matchScalarType, - matchLeftHandSide, matchALeftHandSide, matchELeftHandSide, matchSingleType, matchTupR + matchIdx, matchVar, matchVars, matchArrayR, matchArraysR, matchGroundR, + matchGroundsR, matchTypeR, matchShapeR, matchShapeType, matchIntegralType, + matchFloatingType, matchNumType, matchScalarType, matchLeftHandSide, + matchALeftHandSide, matchELeftHandSide, matchSingleType, matchTupR, + matchVecR ) where @@ -44,7 +39,10 @@ import Data.Array.Accelerate.Representation.Array import Data.Array.Accelerate.Representation.Shape import Data.Array.Accelerate.Representation.Slice import Data.Array.Accelerate.Representation.Stencil +import Data.Array.Accelerate.Representation.Tag import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Representation.Vec +import Data.Array.Accelerate.Representation.Ground import Data.Array.Accelerate.Type import Data.Primitive.Vec import qualified Data.Array.Accelerate.Sugar.Shape as Sugar @@ -55,6 +53,7 @@ import Unsafe.Coerce ( unsafeCoerce ) import System.IO.Unsafe ( unsafePerformIO ) import System.Mem.StableName import Prelude hiding ( exp ) +import Control.Monad (join) -- The type of matching array computations @@ -98,10 +97,13 @@ matchPreOpenAcc matchAcc = match match Anil Anil = Just Refl - match (Apply _ f1 a1) (Apply _ f2 a2) - | Just Refl <- matchPreOpenAfun matchAcc f1 f2 - , Just Refl <- matchAcc a1 a2 - = Just Refl + match (Manifest a1) (Manifest a2) + = matchAcc a1 a2 + + match (Aassert msg1 c1 a1) (Aassert msg2 c2 a2) + | Just Refl <- matchExp c1 c2 + , msg1 == msg2 + = matchAcc a1 a2 match (Aforeign _ ff1 f1 a1) (Aforeign _ ff2 f2 a2) | Just Refl <- matchAcc a1 a2 @@ -200,10 +202,9 @@ matchPreOpenAcc matchAcc = match , Just Refl <- matchAcc a1 a2 = Just Refl - match (Permute f1 d1 p1 a1) (Permute f2 d2 p2 a2) - | Just Refl <- matchFun f1 f2 + match (Permute f1 d1 a1) (Permute f2 d2 a2) + | Just Refl <- join $ matchFun <$> f1 <*> f2 , Just Refl <- matchAcc d1 d2 - , Just Refl <- matchFun p1 p2 , Just Refl <- matchAcc a1 a2 = Just Refl @@ -420,6 +421,16 @@ matchArrayR (ArrayR shr1 tp1) (ArrayR shr2 tp2) , Just Refl <- matchTypeR tp1 tp2 = Just Refl matchArrayR _ _ = Nothing +{-# INLINEABLE matchGroundsR #-} +matchGroundsR :: GroundsR s -> GroundsR t -> Maybe (s :~: t) +matchGroundsR = matchTupR matchGroundR + +{-# INLINEABLE matchGroundR #-} +matchGroundR :: GroundR s -> GroundR t -> Maybe (s :~: t) +matchGroundR (GroundRscalar s) (GroundRscalar t) = matchScalarType s t +matchGroundR (GroundRbuffer s) (GroundRbuffer t) + | Just Refl <- matchScalarType s t = Just Refl +matchGroundR _ _ = Nothing -- Compute the congruence of two scalar expressions. Two nodes are congruent if -- either: @@ -432,9 +443,10 @@ matchArrayR _ _ = Nothing -- {-# INLINEABLE matchOpenExp #-} matchOpenExp - :: forall env aenv s t. - OpenExp env aenv s - -> OpenExp env aenv t + :: forall arr env s t. + IsArrayInstr arr + => PreOpenExp arr env s + -> PreOpenExp arr env t -> Maybe (s :~: t) matchOpenExp (Let lhs1 x1 e1) (Let lhs2 x2 e2) @@ -475,16 +487,14 @@ matchOpenExp (Pair a1 b1) (Pair a2 b2) matchOpenExp Nil Nil = Just Refl -matchOpenExp (IndexSlice sliceIndex1 ix1 sh1) (IndexSlice sliceIndex2 ix2 sh2) - | Just Refl <- matchOpenExp ix1 ix2 - , Just Refl <- matchOpenExp sh1 sh2 - , Just Refl <- matchSliceIndex sliceIndex1 sliceIndex2 +matchOpenExp (VecPack r1 e1) (VecPack r2 e2) + | Just Refl <- matchVecR r1 r2 + , Just Refl <- matchOpenExp e1 e2 = Just Refl -matchOpenExp (IndexFull sliceIndex1 ix1 sl1) (IndexFull sliceIndex2 ix2 sl2) - | Just Refl <- matchOpenExp ix1 ix2 - , Just Refl <- matchOpenExp sl1 sl2 - , Just Refl <- matchSliceIndex sliceIndex1 sliceIndex2 +matchOpenExp (VecUnpack r1 e1) (VecUnpack r2 e2) + | Just Refl <- matchVecR r1 r2 + , Just Refl <- matchOpenExp e1 e2 = Just Refl matchOpenExp (ToIndex _ sh1 i1) (ToIndex _ sh2 i2) @@ -497,6 +507,13 @@ matchOpenExp (FromIndex _ sh1 i1) (FromIndex _ sh2 i2) , Just Refl <- matchOpenExp sh1 sh2 = Just Refl +matchOpenExp c1@(Case t1 alts1 d1) c2@(Case t2 alts2 d2) + | Just Refl <- matchOpenExp t1 t2 + , Just Refl <- matchTypeR (expType c1) (expType c2) + , matchAlts alts1 alts2 + , matchMaybeExp d1 d2 + = Just Refl + matchOpenExp (Cond p1 t1 e1) (Cond p2 t2 e2) | Just Refl <- matchOpenExp p1 p2 , Just Refl <- matchOpenExp t1 t2 @@ -509,8 +526,14 @@ matchOpenExp (While p1 f1 x1) (While p2 f2 x2) , Just Refl <- matchOpenFun f1 f2 = Just Refl -matchOpenExp (PrimConst c1) (PrimConst c2) - = matchPrimConst c1 c2 +matchOpenExp (Assert msg1 g1 e1) (Assert msg2 g2 e2) + | Just Refl <- matchOpenExp g1 g2 + , msg1 == msg2 + = matchOpenExp e1 e2 + +matchOpenExp (Assume g1 e1) (Assume g2 e2) + | Just Refl <- matchOpenExp g1 g2 + = matchOpenExp e1 e2 matchOpenExp (PrimApp f1 x1) (PrimApp f2 x2) | Just x1' <- commutes f1 x1 @@ -523,18 +546,9 @@ matchOpenExp (PrimApp f1 x1) (PrimApp f2 x2) , Just Refl <- matchPrimFun f1 f2 = Just Refl -matchOpenExp (Index a1 x1) (Index a2 x2) - | Just Refl <- matchVar a1 a2 - , Just Refl <- matchOpenExp x1 x2 - = Just Refl - -matchOpenExp (LinearIndex a1 x1) (LinearIndex a2 x2) - | Just Refl <- matchVar a1 a2 - , Just Refl <- matchOpenExp x1 x2 - = Just Refl - -matchOpenExp (Shape a1) (Shape a2) - | Just Refl <- matchVar a1 a2 +matchOpenExp (ArrayInstr arr1 e1) (ArrayInstr arr2 e2) + | Just Refl <- matchArrayInstr arr1 arr2 + , Just Refl <- matchOpenExp e1 e2 = Just Refl matchOpenExp (ShapeSize _ sh1) (ShapeSize _ sh2) @@ -544,13 +558,39 @@ matchOpenExp (ShapeSize _ sh1) (ShapeSize _ sh2) matchOpenExp _ _ = Nothing +matchAlts :: IsArrayInstr arr => [(TAG, PreOpenExp arr env t)] -> [(TAG, PreOpenExp arr env t)] -> Bool +matchAlts [] [] = True +matchAlts ((t1, e1) : alts1) ((t2, e2) : alts2) + | t1 == t2 + , Just Refl <- matchOpenExp e1 e2 + = matchAlts alts1 alts2 +matchAlts _ _ = False + +matchMaybeExp + :: IsArrayInstr arr + => Maybe (PreOpenExp arr env t) + -> Maybe (PreOpenExp arr env t) + -> Bool +matchMaybeExp Nothing Nothing = True +matchMaybeExp (Just a) (Just b) = isJust $ matchOpenExp a b +matchMaybeExp _ _ = False + +matchVecR :: VecR n1 s1 t1 -> VecR n2 s2 t2 -> Maybe ((Vec n1 s1, t1) :~: (Vec n2 s2, t2)) +matchVecR (VecRnil t1) (VecRnil t2) + | Just Refl <- matchSingleType t1 t2 + = Just Refl +matchVecR (VecRsucc r1) (VecRsucc r2) + | Just Refl <- matchVecR r1 r2 + = Just Refl +matchVecR _ _ = Nothing -- Match scalar functions -- {-# INLINEABLE matchOpenFun #-} matchOpenFun - :: OpenFun env aenv s - -> OpenFun env aenv t + :: IsArrayInstr arr + => PreOpenFun arr env s + -> PreOpenFun arr env t -> Maybe (s :~: t) matchOpenFun (Lam lhs1 s) (Lam lhs2 t) | Just Refl <- matchELeftHandSide lhs1 lhs2 @@ -582,29 +622,6 @@ evalEqNum (IntegralNumType t) | IntegralDict <- integralDict t = uncurry (==) evalEqNum (FloatingNumType t) | FloatingDict <- floatingDict t = uncurry (==) --- Environment projection indices --- -{-# INLINEABLE matchIdx #-} -matchIdx :: Idx env s -> Idx env t -> Maybe (s :~: t) -matchIdx ZeroIdx ZeroIdx = Just Refl -matchIdx (SuccIdx u) (SuccIdx v) = matchIdx u v -matchIdx _ _ = Nothing - -{-# INLINEABLE matchVar #-} -matchVar :: Var s env t1 -> Var s env t2 -> Maybe (t1 :~: t2) -matchVar (Var _ v1) (Var _ v2) = matchIdx v1 v2 - -{-# INLINEABLE matchVars #-} -matchVars :: Vars s env t1 -> Vars s env t2 -> Maybe (t1 :~: t2) -matchVars TupRunit TupRunit = Just Refl -matchVars (TupRsingle v1) (TupRsingle v2) - | Just Refl <- matchVar v1 v2 = Just Refl -matchVars (TupRpair v w) (TupRpair x y) - | Just Refl <- matchVars v x - , Just Refl <- matchVars w y = Just Refl -matchVars _ _ = Nothing - - -- Slice specifications -- matchSliceIndex :: SliceIndex slix1 sl1 co1 sh1 -> SliceIndex slix2 sl2 co2 sh2 -> Maybe (SliceIndex slix1 sl1 co1 sh1 :~: SliceIndex slix2 sl2 co2 sh2) @@ -622,14 +639,6 @@ matchSliceIndex (SliceFixed sl1) (SliceFixed sl2) matchSliceIndex _ _ = Nothing --- Primitive constants and functions --- -matchPrimConst :: PrimConst s -> PrimConst t -> Maybe (s :~: t) -matchPrimConst (PrimMinBound s) (PrimMinBound t) = matchBoundedType s t -matchPrimConst (PrimMaxBound s) (PrimMaxBound t) = matchBoundedType s t -matchPrimConst (PrimPi s) (PrimPi t) = matchFloatingType s t -matchPrimConst _ _ = Nothing - -- Covariant function matching -- @@ -684,12 +693,8 @@ matchPrimFun (PrimFloor _ s) (PrimFloor _ t) = matchIntegr matchPrimFun (PrimCeiling _ s) (PrimCeiling _ t) = matchIntegralType s t matchPrimFun (PrimIsNaN _) (PrimIsNaN _) = Just Refl matchPrimFun (PrimIsInfinite _) (PrimIsInfinite _) = Just Refl -matchPrimFun (PrimLt _) (PrimLt _) = Just Refl -matchPrimFun (PrimGt _) (PrimGt _) = Just Refl -matchPrimFun (PrimLtEq _) (PrimLtEq _) = Just Refl -matchPrimFun (PrimGtEq _) (PrimGtEq _) = Just Refl -matchPrimFun (PrimEq _) (PrimEq _) = Just Refl -matchPrimFun (PrimNEq _) (PrimNEq _) = Just Refl +matchPrimFun (PrimCmp _ c1) (PrimCmp _ c2) + | c1 == c2 = Just Refl matchPrimFun (PrimMax _) (PrimMax _) = Just Refl matchPrimFun (PrimMin _) (PrimMin _) = Just Refl matchPrimFun (PrimFromIntegral _ s) (PrimFromIntegral _ t) = matchNumType s t @@ -701,7 +706,6 @@ matchPrimFun PrimLNot PrimLNot = Just Refl matchPrimFun _ _ = Nothing - -- Contravariant function matching -- {-# INLINEABLE matchPrimFun' #-} @@ -763,29 +767,9 @@ matchPrimFun' PrimLAnd PrimLAnd = Just Refl matchPrimFun' PrimLOr PrimLOr = Just Refl matchPrimFun' PrimLNot PrimLNot = Just Refl -matchPrimFun' (PrimLt s) (PrimLt t) - | Just Refl <- matchSingleType s t - = Just Refl - -matchPrimFun' (PrimGt s) (PrimGt t) - | Just Refl <- matchSingleType s t - = Just Refl - -matchPrimFun' (PrimLtEq s) (PrimLtEq t) - | Just Refl <- matchSingleType s t - = Just Refl - -matchPrimFun' (PrimGtEq s) (PrimGtEq t) - | Just Refl <- matchSingleType s t - = Just Refl - -matchPrimFun' (PrimEq s) (PrimEq t) - | Just Refl <- matchSingleType s t - = Just Refl - -matchPrimFun' (PrimNEq s) (PrimNEq t) - | Just Refl <- matchSingleType s t - = Just Refl +matchPrimFun' (PrimCmp s c1) (PrimCmp t c2) + | c1 == c2 + , Just Refl <- matchSingleType s t = Just Refl matchPrimFun' _ _ = Nothing @@ -807,12 +791,23 @@ matchTypeR = matchTupR matchScalarType -- a known branch. -- {-# INLINEABLE matchShapeType #-} -matchShapeType :: forall s t. (Sugar.Shape s, Sugar.Shape t) => Maybe (s :~: t) +matchShapeType :: forall s t. ( +#ifdef ACCELERATE_INTERNAL_CHECKS + Typeable s, Typeable t, +#endif + Sugar.Shape s, Sugar.Shape t) => Maybe (s :~: t) matchShapeType | Just Refl <- matchShapeR (Sugar.shapeR @s) (Sugar.shapeR @t) #ifdef ACCELERATE_INTERNAL_CHECKS = gcast Refl #else + -- Ivo: I was trying to remove this unsafeCoerce. + -- I introduced an injective type family ShapeEltR which is equal to EltR for + -- shapes. In contrast to EltR, ShapeEltR can then be defined as an injective + -- type class since it is only defined over Z and (Int :. sh). + -- However, even when ShapeEltR is injective, GHC won't use that + -- ShapeEltR s ~ ShapeEltR t implies s ~ t. See the following issue: + -- https://gitlab.haskell.org/ghc/ghc/-/issues/10833 = Just (unsafeCoerce Refl) #endif | otherwise @@ -856,9 +851,11 @@ matchNumType (IntegralNumType s) (IntegralNumType t) = matchIntegralType s t matchNumType (FloatingNumType s) (FloatingNumType t) = matchFloatingType s t matchNumType _ _ = Nothing +{- TODO WALL: DEAD CODE {-# INLINEABLE matchBoundedType #-} matchBoundedType :: BoundedType s -> BoundedType t -> Maybe (s :~: t) matchBoundedType (IntegralBoundedType s) (IntegralBoundedType t) = matchIntegralType s t +-} {-# INLINEABLE matchIntegralType #-} matchIntegralType :: IntegralType s -> IntegralType t -> Maybe (s :~: t) @@ -889,26 +886,30 @@ matchFloatingType _ _ = Nothing -- a stable ordering such that matching recognises expressions modulo -- commutativity. -- +-- TODO: Should we remove this functionality? The expression simplifier +-- already converts commutative operations to a normal form, and we +-- compare expressions only after simplification. commutes - :: forall env aenv a r. - PrimFun (a -> r) - -> OpenExp env aenv a - -> Maybe (OpenExp env aenv a) + :: forall arr env a r. + IsArrayInstr arr + => PrimFun (a -> r) + -> PreOpenExp arr env a + -> Maybe (PreOpenExp arr env a) commutes f x = case f of PrimAdd{} -> Just (swizzle x) PrimMul{} -> Just (swizzle x) PrimBAnd{} -> Just (swizzle x) PrimBOr{} -> Just (swizzle x) PrimBXor{} -> Just (swizzle x) - PrimEq{} -> Just (swizzle x) - PrimNEq{} -> Just (swizzle x) + PrimCmp _ CmpEq -> Just (swizzle x) + PrimCmp _ CmpNEq -> Just (swizzle x) PrimMax{} -> Just (swizzle x) PrimMin{} -> Just (swizzle x) PrimLAnd -> Just (swizzle x) PrimLOr -> Just (swizzle x) _ -> Nothing where - swizzle :: OpenExp env aenv (a',a') -> OpenExp env aenv (a',a') + swizzle :: PreOpenExp arr env (a',a') -> PreOpenExp arr env (a',a') swizzle exp | (a `Pair` b) <- exp , hashOpenExp a > hashOpenExp b = b `Pair` a diff --git a/src/Data/Array/Accelerate/Array/Buffer.hs b/src/Data/Array/Accelerate/Array/Buffer.hs new file mode 100644 index 000000000..f02f84358 --- /dev/null +++ b/src/Data/Array/Accelerate/Array/Buffer.hs @@ -0,0 +1,432 @@ +{-# LANGUAGE CPP #-} +{-# LANGUAGE TemplateHaskell #-} +{-# OPTIONS_HADDOCK hide #-} +-- | +-- Module : Data.Array.Accelerate.Array.Buffer +-- Copyright : [2008..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- +-- This module fixes the concrete representation of Accelerate arrays. +-- We allocate memory via foreign function calls, outside of the Haskell heap. +-- This gives us flexibility to manage memory ourselves, and prevents the +-- Haskell GC to move arrays around (which would break in our foreign function +-- calls). +-- These arrays are reclaimed via reference counting. A Haskell-side reference +-- to an array has a finalizer which will decrement the reference count. +-- C or our generated LLVM code working with references to arrays will +-- manually increment and decrement the reference count. +-- + +module Data.Array.Accelerate.Array.Buffer ( + + -- * Array operations and representations + Buffers, Buffer(..), MutableBuffers, MutableBuffer(..), + runBuffers, + newBuffers, newBuffer, + indexBuffers, indexBuffers', indexBuffer, readBuffers, readBuffer, writeBuffers, writeBuffer, + touchBuffers, touchBuffer, touchMutableBuffers, touchMutableBuffer, + rnfBuffers, rnfBuffer, unsafeFreezeBuffer, unsafeFreezeBuffers, + veryUnsafeUnfreezeBuffers, bufferToList, bufferRetainAndGetRef, bufferRelease, bufferFromPtr, + memoryByteSize, + + -- * Type macros + HTYPE_INT, HTYPE_WORD, HTYPE_CLONG, HTYPE_CULONG, HTYPE_CCHAR, + + -- * Utilities for type classes + ScalarArrayDict(..), + + -- * TemplateHaskell + liftBuffers, liftBuffer, + + -- * Memory tracking + memoryCounterReset, memoryCounterReport, +) where + + +import Data.Array.Accelerate.Error +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Type + +import Data.Array.Accelerate.Debug.Internal.Flags +import Data.Array.Accelerate.Debug.Internal.Trace + +import Control.Applicative +import Data.Bits +import Data.IORef +import Data.Typeable ( (:~:)(..) ) +import Foreign.Ptr +import Foreign.ForeignPtr +import Foreign.Storable +import Foreign.Marshal.Array ( copyArray, peekArray ) +import Formatting hiding ( bytes ) +import Language.Haskell.TH.Extra hiding ( Type ) +import Language.Haskell.TH.Syntax +import System.IO.Unsafe +import Prelude hiding ( mapM ) + +import GHC.Exts hiding ( build ) + +import System.Mem + +-- | A buffer is a piece of memory representing one of the fields +-- of the SoA of an array. It does not have a multi-dimensional size, +-- e.g. the shape of the array should be stored elsewhere. +-- Replaces the former 'ScalarArrayData' type synonym. +-- +newtype Buffer e = Buffer (ForeignPtr e) + +-- | A structure of buffers represents an array, corresponding to the SoA conversion. +-- Replaces the old 'ArrayData' and 'MutableArrayData' type aliasses and the +-- 'GArrayDataR' type family. +-- +type Buffers e = Distribute Buffer e + +newtype MutableBuffer e = MutableBuffer (ForeignPtr e) +type MutableBuffers e = Distribute MutableBuffer e + +-- SEE: [HLS and GHC IDE] +-- +#ifndef __GHCIDE__ + +foreign import ccall unsafe "accelerate_buffer_alloc" memoryAlloc :: Word64 -> IO (Ptr ()) +foreign import ccall unsafe "accelerate_buffer_byte_size" memoryByteSize :: Ptr () -> IO Word64 +foreign import ccall unsafe "accelerate_buffer_retain" memoryRetain :: Ptr () -> IO () +foreign import ccall unsafe "accelerate_buffer_release" memoryRelease :: Ptr () -> IO () +foreign import ccall unsafe "&accelerate_buffer_release" memoryReleaseRef :: FunPtr (Ptr () -> IO ()) + +#else + +memoryAlloc :: Word64 -> IO (Ptr ()) +memoryAlloc = undefined + +memoryByteSize :: Ptr () -> IO Word64 +memoryByteSize = undefined + +memoryRetain :: Ptr () -> IO () +memoryRetain = undefined + +memoryRelease :: Ptr () -> IO () +memoryRelease = undefined + +memoryReleaseRef :: FunPtr (Ptr () -> IO ()) +memoryReleaseRef = undefined + +#endif + + +#if defined(ACCELERATE_MEMORY_COUNTER) && !defined(__GHCIDE__) + +foreign import ccall unsafe "accelerate_memory_counter_reset" memoryCounterReset :: IO () +foreign import ccall unsafe "accelerate_memory_counter_total" memoryCounterTotal :: IO Word64 +foreign import ccall unsafe "accelerate_memory_counter_init" memoryCounterInit :: IO Word64 +foreign import ccall unsafe "accelerate_memory_counter_max" memoryCounterMax :: IO Word64 + +#else + +memoryCounterReset :: IO () +memoryCounterReset = return () + +#endif + +#if defined(__GHCIDE__) + +memoryCounterTotal :: IO Word64 +memoryCounterTotal = return 0 + +memoryCounterInit :: IO Word64 +memoryCounterInit = return 0 + +memoryCounterMax :: IO Word64 +memoryCounterMax = return 0 + +#endif + +#if defined(ACCELERATE_MEMORY_COUNTER) || defined(__GHCIDE__) + +memoryCounterReport :: IO () +memoryCounterReport = do + totalAlloc <- memoryCounterTotal + initAlloc <- memoryCounterInit + maxAlloc <- memoryCounterMax + putStrLn $ "Total memory allocated: " ++ show totalAlloc ++ " bytes" + putStrLn $ "Initial memory allocated: " ++ show initAlloc ++ " bytes" + putStrLn $ "Maximum memory allocated: " ++ show maxAlloc ++ " bytes" + +#else + +memoryCounterReport :: IO () +memoryCounterReport = putStrLn "MEMORY COUNTER DISABLED" + +#endif + +-- +-- SEE: [linking to .c files] +-- +runQ $ do + addForeignFilePath LangC "cbits/memory.c" + addForeignFilePath LangCxx "cbits/alloc.cpp" + return [] + +data ScalarArrayDict a where + ScalarArrayDict :: ( Buffers a ~ Buffer a, Storable a ) + => ScalarArrayDict a + + +scalarArrayDict :: ScalarType a -> ScalarArrayDict a +scalarArrayDict = scalar + where + scalar :: ScalarType a -> ScalarArrayDict a + scalar (VectorScalarType t) = vector t + scalar (SingleScalarType t) + | ScalarArrayDict <- singleArrayDict t + = ScalarArrayDict + + vector :: VectorType a -> ScalarArrayDict a + vector (VectorType _ s) + | ScalarArrayDict <- singleArrayDict s + , SingleDict <- singleDict s + = ScalarArrayDict + +singleArrayDict :: SingleType a -> ScalarArrayDict a +singleArrayDict = single + where + single :: SingleType a -> ScalarArrayDict a + single (NumSingleType t) = num t + + num :: NumType a -> ScalarArrayDict a + num (IntegralNumType t) = integral t + num (FloatingNumType t) = floating t + + integral :: IntegralType a -> ScalarArrayDict a + integral TypeInt = ScalarArrayDict + integral TypeInt8 = ScalarArrayDict + integral TypeInt16 = ScalarArrayDict + integral TypeInt32 = ScalarArrayDict + integral TypeInt64 = ScalarArrayDict + integral TypeWord = ScalarArrayDict + integral TypeWord8 = ScalarArrayDict + integral TypeWord16 = ScalarArrayDict + integral TypeWord32 = ScalarArrayDict + integral TypeWord64 = ScalarArrayDict + + floating :: FloatingType a -> ScalarArrayDict a + floating TypeHalf = ScalarArrayDict + floating TypeFloat = ScalarArrayDict + floating TypeDouble = ScalarArrayDict + +-- Array operations +-- ---------------- + +newBuffers :: forall e. HasCallStack => TypeR e -> Int -> IO (MutableBuffers e) +newBuffers TupRunit !_ = return () +newBuffers (TupRpair t1 t2) !size = (,) <$> newBuffers t1 size <*> newBuffers t2 size +newBuffers (TupRsingle t) !size + | Refl <- reprIsSingle @ScalarType @e @MutableBuffer t = newBuffer t size + +newBuffer :: HasCallStack => ScalarType e -> Int -> IO (MutableBuffer e) +newBuffer tp !size + | ScalarArrayDict <- scalarArrayDict tp + = MutableBuffer <$> allocateArray size + +indexBuffers :: TypeR e -> Buffers e -> Int -> e +indexBuffers tR arr ix = unsafePerformIO $ indexBuffers' tR arr ix + +indexBuffers' :: TypeR e -> Buffers e -> Int -> IO e +indexBuffers' tR arr = readBuffers tR (veryUnsafeUnfreezeBuffers tR arr) + +indexBuffer :: ScalarType e -> Buffer e -> Int -> e +indexBuffer tR (Buffer arr) ix = unsafePerformIO $ readBuffer tR (MutableBuffer arr) ix + +readBuffers :: forall e. TypeR e -> MutableBuffers e -> Int -> IO e +readBuffers TupRunit () !_ = return () +readBuffers (TupRpair t1 t2) (a1, a2) !ix = (,) <$> readBuffers t1 a1 ix <*> readBuffers t2 a2 ix +readBuffers (TupRsingle t) !buffer !ix + | Refl <- reprIsSingle @ScalarType @e @MutableBuffer t = readBuffer t buffer ix + +readBuffer :: forall e. ScalarType e -> MutableBuffer e -> Int -> IO e +readBuffer tp !(MutableBuffer buffer) !ix + | ScalarArrayDict <- scalarArrayDict tp + = withForeignPtr buffer $ \ptr -> peekElemOff ptr ix + +writeBuffers :: forall e. TypeR e -> MutableBuffers e -> Int -> e -> IO () +writeBuffers TupRunit () !_ () = return () +writeBuffers (TupRpair t1 t2) (a1, a2) !ix (v1, v2) = writeBuffers t1 a1 ix v1 >> writeBuffers t2 a2 ix v2 +writeBuffers (TupRsingle t) arr !ix !val + | Refl <- reprIsSingle @ScalarType @e @MutableBuffer t = writeBuffer t arr ix val + +writeBuffer :: forall e. ScalarType e -> MutableBuffer e -> Int -> e -> IO () +writeBuffer tp (MutableBuffer buffer) !ix !val + | ScalarArrayDict <- scalarArrayDict tp + = withForeignPtr buffer $ \ptr -> pokeElemOff ptr ix val + +touchBuffers :: forall e. TypeR e -> Buffers e -> IO () +touchBuffers TupRunit () = return() +touchBuffers (TupRpair t1 t2) (b1, b2) = touchBuffers t1 b1 >> touchBuffers t2 b2 +touchBuffers (TupRsingle t) buffer + | Refl <- reprIsSingle @ScalarType @e @Buffer t = touchBuffer buffer + +touchMutableBuffers :: forall e. TypeR e -> MutableBuffers e -> IO () +touchMutableBuffers TupRunit () = return() +touchMutableBuffers (TupRpair t1 t2) (b1, b2) = touchMutableBuffers t1 b1 >> touchMutableBuffers t2 b2 +touchMutableBuffers (TupRsingle t) buffer + | Refl <- reprIsSingle @ScalarType @e @MutableBuffer t = touchMutableBuffer buffer + +touchBuffer :: Buffer e -> IO () +touchBuffer (Buffer buffer) = touchForeignPtr buffer + +touchMutableBuffer :: MutableBuffer e -> IO () +touchMutableBuffer (MutableBuffer buffer) = touchForeignPtr buffer + +rnfBuffers :: forall e. TypeR e -> Buffers e -> () +rnfBuffers TupRunit () = () +rnfBuffers (TupRpair t1 t2) (a1, a2) = rnfBuffers t1 a1 `seq` rnfBuffers t2 a2 +rnfBuffers (TupRsingle t) arr + | Refl <- reprIsSingle @ScalarType @e @Buffer t = rnfBuffer arr + +rnfBuffer :: Buffer e -> () +rnfBuffer !_ = () + +-- | Safe combination of creating and fast freezing of array data. +-- +runBuffers + :: TypeR e + -> IO (MutableBuffers e, e) + -> (Buffers e, e) +runBuffers tp st = unsafePerformIO $ do + (mbuffer, r) <- st + let buffer = unsafeFreezeBuffers tp mbuffer + return (buffer, r) + +unsafeFreezeBuffers :: forall e. TypeR e -> MutableBuffers e -> Buffers e +unsafeFreezeBuffers TupRunit () = () +unsafeFreezeBuffers (TupRpair t1 t2) (b1, b2) = (unsafeFreezeBuffers t1 b1, unsafeFreezeBuffers t2 b2) +unsafeFreezeBuffers (TupRsingle t) buffer + | Refl <- reprIsSingle @ScalarType @e @MutableBuffer t + , Refl <- reprIsSingle @ScalarType @e @Buffer t = unsafeFreezeBuffer buffer + +unsafeFreezeBuffer :: MutableBuffer e -> Buffer e +unsafeFreezeBuffer (MutableBuffer arr) = Buffer arr + +veryUnsafeUnfreezeBuffers :: forall e. TypeR e -> Buffers e -> MutableBuffers e +veryUnsafeUnfreezeBuffers TupRunit () = () +veryUnsafeUnfreezeBuffers (TupRpair t1 t2) (b1, b2) = (veryUnsafeUnfreezeBuffers t1 b1, veryUnsafeUnfreezeBuffers t2 b2) +veryUnsafeUnfreezeBuffers (TupRsingle t) buffer + | Refl <- reprIsSingle @ScalarType @e @MutableBuffer t + , Refl <- reprIsSingle @ScalarType @e @Buffer t = veryUnsafeUnfreezeBuffer buffer + +veryUnsafeUnfreezeBuffer :: Buffer e -> MutableBuffer e +veryUnsafeUnfreezeBuffer (Buffer arr) = MutableBuffer arr + +-- Allocate a new buffer with enough storage to hold the given number of +-- elements. +-- +-- The buffer is uninitialised. +-- +allocateArray :: forall e. (HasCallStack, Storable e) => Int -> IO (ForeignPtr e) +allocateArray !size = internalCheck "size must be >= 0" (size >= 0) $ do + let bytes = size * sizeOf (undefined :: e) + + hintGCAlloc bytes + + ptr <- memoryAlloc $ fromIntegral bytes + foreignPtr <- newForeignPtr memoryReleaseRef ptr + traceM dump_gc ("gc: allocated new host array (size=" % int % ", ptr=" % build % ")") bytes ptr + return $ castForeignPtr foreignPtr + +-- Once per 1GB of allocations, perform garbage collection. +-- Buffers are allocated outside of the Haskell world, so we have to force +-- GC to deallocate them indirectly. +hintGCAlloc :: Int -> IO () +hintGCAlloc bytes = do + let maxSize = 1024 * 1024 * 1024 + gc <- atomicModifyIORef' bytesAllocatedSinceGC $ \sz -> + let newSize = sz + bytes + in if newSize < maxSize then (newSize, False) else (newSize `mod` maxSize, True) + if gc then performGC else return () + +{-# NOINLINE bytesAllocatedSinceGC #-} +bytesAllocatedSinceGC :: IORef Int +bytesAllocatedSinceGC = unsafePerformIO $ newIORef 0 + +bufferToList :: ScalarType e -> Int -> Buffer e -> [e] +bufferToList tp n buffer = go 0 + where + go !i | i >= n = [] + | otherwise = indexBuffer tp buffer i : go (i + 1) + + +bufferRetainAndGetRef :: Buffer e -> IO (Ptr e) +bufferRetainAndGetRef (Buffer foreignPtr) = withForeignPtr foreignPtr $ \ptr -> do + memoryRetain $ castPtr ptr + return $ castPtr ptr + +-- Ptr should originate from bufferRetainAndGetRef +bufferRelease :: Ptr e -> IO () +bufferRelease = memoryRelease . castPtr + +bufferFromPtr :: Ptr e -> IO (Buffer e) +bufferFromPtr ptr = do + byteSize <- fromIntegral <$> memoryByteSize (castPtr ptr) + hintGCAlloc byteSize + fp <- newForeignPtr memoryReleaseRef $ castPtr ptr + return $ Buffer $ castForeignPtr fp + +liftBuffers :: forall e. TypeR e -> Buffers e -> CodeQ (Buffers e) +liftBuffers TupRunit () = [|| () ||] +liftBuffers (TupRpair t1 t2) (b1, b2) = [|| ($$(liftBuffers t1 b1), $$(liftBuffers t2 b2)) ||] +liftBuffers (TupRsingle s) buffer + | Refl <- reprIsSingle @ScalarType @e @Buffer s = liftBuffer s buffer + +liftBuffer :: forall e. ScalarType e -> Buffer e -> CodeQ (Buffer e) +liftBuffer tp (Buffer arr) + | ScalarArrayDict <- scalarArrayDict tp = [|| Buffer $$(liftBufferData arr) ||] + +liftBufferData :: forall a. Storable a => ForeignPtr a -> CodeQ (ForeignPtr a) +liftBufferData buffer = unsafePerformIO $ withForeignPtr buffer $ \ptr -> do + byteSize <- fromIntegral <$> memoryByteSize (castPtr ptr) + let size = byteSize `div` sizeOf (undefined::a) + bytes <- peekArray byteSize (castPtr ptr :: Ptr Word8) + + return [|| + unsafePerformIO $ do + let ptrData = Ptr $$(unsafeCodeCoerce $ litE (StringPrimL bytes)) :: Ptr Word8 + result <- allocateArray size + withForeignPtr result $ \p -> + copyArray (castPtr p) ptrData byteSize + return result + ||] + +-- Determine the underlying type of a Haskell CLong or CULong. +-- +runQ [d| type HTYPE_INT = $( + case finiteBitSize (undefined::Int) of + 32 -> [t| Int32 |] + 64 -> [t| Int64 |] + _ -> error "I don't know what architecture I am" ) |] + +runQ [d| type HTYPE_WORD = $( + case finiteBitSize (undefined::Word) of + 32 -> [t| Word32 |] + 64 -> [t| Word64 |] + _ -> error "I don't know what architecture I am" ) |] + +runQ [d| type HTYPE_CLONG = $( + case finiteBitSize (undefined::CLong) of + 32 -> [t| Int32 |] + 64 -> [t| Int64 |] + _ -> error "I don't know what architecture I am" ) |] + +runQ [d| type HTYPE_CULONG = $( + case finiteBitSize (undefined::CULong) of + 32 -> [t| Word32 |] + 64 -> [t| Word64 |] + _ -> error "I don't know what architecture I am" ) |] + +runQ [d| type HTYPE_CCHAR = $( + if isSigned (undefined::CChar) + then [t| Int8 |] + else [t| Word8 |] ) |] diff --git a/src/Data/Array/Accelerate/Array/Data.hs b/src/Data/Array/Accelerate/Array/Data.hs deleted file mode 100644 index 441b8609b..000000000 --- a/src/Data/Array/Accelerate/Array/Data.hs +++ /dev/null @@ -1,408 +0,0 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE CPP #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE MagicHash #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE UnboxedTuples #-} -{-# LANGUAGE UndecidableInstances #-} -{-# OPTIONS_HADDOCK hide #-} --- | --- Module : Data.Array.Accelerate.Array.Data --- Copyright : [2008..2020] The Accelerate Team --- License : BSD3 --- --- Maintainer : Trevor L. McDonell --- Stability : experimental --- Portability : non-portable (GHC extensions) --- --- This module fixes the concrete representation of Accelerate arrays. We --- allocate all arrays using pinned memory to enable safe direct-access by --- non-Haskell code in multi-threaded code. In particular, we can safely pass --- pointers to an array's payload to foreign code. --- - -module Data.Array.Accelerate.Array.Data ( - - -- * Array operations and representations - ArrayData, MutableArrayData, ScalarArrayData, GArrayDataR, ScalarArrayDataR, - runArrayData, - newArrayData, - indexArrayData, readArrayData, writeArrayData, - unsafeArrayDataPtr, - touchArrayData, - rnfArrayData, - - -- * Type macros - HTYPE_INT, HTYPE_WORD, HTYPE_CLONG, HTYPE_CULONG, HTYPE_CCHAR, - - -- * Allocator internals - registerForeignPtrAllocator, - - -- * Utilities for type classes - ScalarArrayDict(..), scalarArrayDict, - SingleArrayDict(..), singleArrayDict, - - -- * TemplateHaskell - liftArrayData, - -) where - -import Data.Array.Accelerate.Array.Unique -import Data.Array.Accelerate.Error -import Data.Array.Accelerate.Lifetime -import Data.Array.Accelerate.Representation.Type -import Data.Array.Accelerate.Type -import Data.Primitive.Vec - -import Data.Array.Accelerate.Debug.Internal.Flags -import Data.Array.Accelerate.Debug.Internal.Profile -import Data.Array.Accelerate.Debug.Internal.Trace - -import Control.Applicative -import Control.DeepSeq -import Control.Monad ( (<=<) ) -import Data.Bits -import Data.IORef -import Data.Primitive ( sizeOf# ) -import Foreign.ForeignPtr -import Foreign.Storable -import Formatting hiding ( bytes ) -import Language.Haskell.TH.Extra hiding ( Type ) -import System.IO.Unsafe -import Prelude hiding ( mapM ) - -import GHC.Exts hiding ( build ) -import GHC.ForeignPtr -import GHC.Types - - --- | Immutable array representation --- -type ArrayData e = MutableArrayData e - --- | Mutable array representation --- -type MutableArrayData e = GArrayDataR UniqueArray e - --- | Underlying array representation. --- --- NOTE: We use a standard (non-strict) pair to enable lazy device-host data transfers --- -type family GArrayDataR ba a where - GArrayDataR ba () = () - GArrayDataR ba (a, b) = (GArrayDataR ba a, GArrayDataR ba b) - GArrayDataR ba a = ba (ScalarArrayDataR a) - -type ScalarArrayData a = UniqueArray (ScalarArrayDataR a) - --- | Mapping from scalar type to the type as represented in memory in an --- array. --- -type family ScalarArrayDataR t where - ScalarArrayDataR Int = Int - ScalarArrayDataR Int8 = Int8 - ScalarArrayDataR Int16 = Int16 - ScalarArrayDataR Int32 = Int32 - ScalarArrayDataR Int64 = Int64 - ScalarArrayDataR Word = Word - ScalarArrayDataR Word8 = Word8 - ScalarArrayDataR Word16 = Word16 - ScalarArrayDataR Word32 = Word32 - ScalarArrayDataR Word64 = Word64 - ScalarArrayDataR Half = Half - ScalarArrayDataR Float = Float - ScalarArrayDataR Double = Double - ScalarArrayDataR (Vec n t) = ScalarArrayDataR t - - -data ScalarArrayDict a where - ScalarArrayDict :: ( ArrayData a ~ ScalarArrayData a, ScalarArrayDataR a ~ ScalarArrayDataR b ) - => {-# UNPACK #-} !Int -- vector width - -> SingleType b -- base type - -> ScalarArrayDict a - -data SingleArrayDict a where - SingleArrayDict :: ( ArrayData a ~ ScalarArrayData a, ScalarArrayDataR a ~ a ) - => SingleArrayDict a - -scalarArrayDict :: ScalarType a -> ScalarArrayDict a -scalarArrayDict = scalar - where - scalar :: ScalarType a -> ScalarArrayDict a - scalar (VectorScalarType t) = vector t - scalar (SingleScalarType t) - | SingleArrayDict <- singleArrayDict t - = ScalarArrayDict 1 t - - vector :: VectorType a -> ScalarArrayDict a - vector (VectorType w s) - | SingleArrayDict <- singleArrayDict s - = ScalarArrayDict w s - -singleArrayDict :: SingleType a -> SingleArrayDict a -singleArrayDict = single - where - single :: SingleType a -> SingleArrayDict a - single (NumSingleType t) = num t - - num :: NumType a -> SingleArrayDict a - num (IntegralNumType t) = integral t - num (FloatingNumType t) = floating t - - integral :: IntegralType a -> SingleArrayDict a - integral TypeInt = SingleArrayDict - integral TypeInt8 = SingleArrayDict - integral TypeInt16 = SingleArrayDict - integral TypeInt32 = SingleArrayDict - integral TypeInt64 = SingleArrayDict - integral TypeWord = SingleArrayDict - integral TypeWord8 = SingleArrayDict - integral TypeWord16 = SingleArrayDict - integral TypeWord32 = SingleArrayDict - integral TypeWord64 = SingleArrayDict - - floating :: FloatingType a -> SingleArrayDict a - floating TypeHalf = SingleArrayDict - floating TypeFloat = SingleArrayDict - floating TypeDouble = SingleArrayDict - - --- Array operations --- ---------------- - -newArrayData :: HasCallStack => TupR ScalarType e -> Int -> IO (MutableArrayData e) -newArrayData TupRunit !_ = return () -newArrayData (TupRpair t1 t2) !size = (,) <$> newArrayData t1 size <*> newArrayData t2 size -newArrayData (TupRsingle t) !size - | SingleScalarType s <- t - , SingleDict <- singleDict s - , SingleArrayDict <- singleArrayDict s - = allocateArray size - -- - | VectorScalarType v <- t - , VectorType w s <- v - , SingleDict <- singleDict s - , SingleArrayDict <- singleArrayDict s - = allocateArray (w * size) - -indexArrayData :: TupR ScalarType e -> ArrayData e -> Int -> e -indexArrayData tR arr ix = unsafePerformIO $ readArrayData tR arr ix - -readArrayData :: forall e. TupR ScalarType e -> MutableArrayData e -> Int -> IO e -readArrayData TupRunit () !_ = return () -readArrayData (TupRpair t1 t2) (a1, a2) !ix = (,) <$> readArrayData t1 a1 ix <*> readArrayData t2 a2 ix -readArrayData (TupRsingle t) arr !ix - | SingleScalarType s <- t - , SingleDict <- singleDict s - , SingleArrayDict <- singleArrayDict s - = unsafeReadArray arr ix - -- - | VectorScalarType v <- t - , VectorType w s <- v - , I# w# <- w - , I# ix# <- ix - , SingleDict <- singleDict s - , SingleArrayDict <- singleArrayDict s - = let - !bytes# = w# *# sizeOf# (undefined :: ScalarArrayDataR e) - !addr# = unPtr# (unsafeUniqueArrayPtr arr) `plusAddr#` (ix# *# bytes#) - in - IO $ \s0 -> - case newAlignedPinnedByteArray# bytes# 16# s0 of { (# s1, mba# #) -> - case copyAddrToByteArray# addr# mba# 0# bytes# s1 of { s2 -> - case unsafeFreezeByteArray# mba# s2 of { (# s3, ba# #) -> - (# s3, Vec ba# #) - }}} - -writeArrayData :: forall e. TupR ScalarType e -> MutableArrayData e -> Int -> e -> IO () -writeArrayData TupRunit () !_ () = return () -writeArrayData (TupRpair t1 t2) (a1, a2) !ix (v1, v2) = writeArrayData t1 a1 ix v1 >> writeArrayData t2 a2 ix v2 -writeArrayData (TupRsingle t) arr !ix !val - | SingleScalarType s <- t - , SingleDict <- singleDict s - , SingleArrayDict <- singleArrayDict s - = unsafeWriteArray arr ix val - -- - | VectorScalarType v <- t - , VectorType w s <- v - , Vec ba# <- val - , I# w# <- w - , I# ix# <- ix - , SingleDict <- singleDict s - , SingleArrayDict <- singleArrayDict s - = let - !bytes# = w# *# sizeOf# (undefined :: ScalarArrayDataR e) - !addr# = unPtr# (unsafeUniqueArrayPtr arr) `plusAddr#` (ix# *# bytes#) - in - IO $ \s0 -> case copyByteArrayToAddr# ba# 0# addr# bytes# s0 of - s1 -> (# s1, () #) - - -unsafeArrayDataPtr :: ScalarType e -> ArrayData e -> Ptr (ScalarArrayDataR e) -unsafeArrayDataPtr t arr - | ScalarArrayDict{} <- scalarArrayDict t - = unsafeUniqueArrayPtr arr - -touchArrayData :: TupR ScalarType e -> ArrayData e -> IO () -touchArrayData TupRunit () = return () -touchArrayData (TupRpair t1 t2) (a1, a2) = touchArrayData t1 a1 >> touchArrayData t2 a2 -touchArrayData (TupRsingle t) arr - | ScalarArrayDict{} <- scalarArrayDict t - = touchUniqueArray arr - -rnfArrayData :: TupR ScalarType e -> ArrayData e -> () -rnfArrayData TupRunit () = () -rnfArrayData (TupRpair t1 t2) (a1, a2) = rnfArrayData t1 a1 `seq` rnfArrayData t2 a2 `seq` () -rnfArrayData (TupRsingle t) arr = rnf (unsafeArrayDataPtr t arr) - -unPtr# :: Ptr a -> Addr# -unPtr# (Ptr addr#) = addr# - --- | Safe combination of creating and fast freezing of array data. --- -runArrayData - :: IO (MutableArrayData e, e) - -> (ArrayData e, e) -runArrayData st = unsafePerformIO $ do - (mad, r) <- st - return (mad, r) - --- Allocate a new array with enough storage to hold the given number of --- elements. --- --- The array is uninitialised and, in particular, allocated lazily. The latter --- is important because it means that for backends that have discrete memory --- spaces (e.g. GPUs), we will not increase host memory pressure simply to track --- intermediate arrays that contain meaningful data only on the device. --- -allocateArray :: forall e. (HasCallStack, Storable e) => Int -> IO (UniqueArray e) -allocateArray !size = internalCheck "size must be >= 0" (size >= 0) $ do - arr <- newUniqueArray <=< unsafeInterleaveIO $ do - let bytes = size * sizeOf (undefined :: e) - new <- readIORef __mallocForeignPtrBytes - ptr <- new bytes - traceM dump_gc ("gc: allocated new host array (size=" % int % ", ptr=" % build % ")") bytes (unsafeForeignPtrToPtr ptr) - local_memory_alloc (unsafeForeignPtrToPtr ptr) bytes - return (castForeignPtr ptr) - - -- If we can, avoid attaching a finalizer at all for performance -#ifdef ACCELERATE_TRACY - let doTracy = True -#else - let doTracy = False -#endif - haveDumpGc <- getFlag dump_gc - if doTracy || haveDumpGc - then addFinalizer (uniqueArrayData arr) $ do - traceM dump_gc ("gc: finalizer for host array " % build % ")") (unsafeUniqueArrayPtr arr) - if doTracy then local_memory_free (unsafeUniqueArrayPtr arr) else return () - else return () - - return arr - --- | Register the given function as the callback to use to allocate new array --- data on the host containing the specified number of bytes. The returned array --- must be pinned (with respect to Haskell's GC), so that it can be passed to --- foreign code. --- -registerForeignPtrAllocator - :: (Int -> IO (ForeignPtr Word8)) - -> IO () -registerForeignPtrAllocator new = do - traceM dump_gc "registering new array allocator" - atomicWriteIORef __mallocForeignPtrBytes new - -{-# NOINLINE __mallocForeignPtrBytes #-} -__mallocForeignPtrBytes :: IORef (Int -> IO (ForeignPtr Word8)) -__mallocForeignPtrBytes = unsafePerformIO $! newIORef mallocPlainForeignPtrBytesAligned - --- | Allocate the given number of bytes with 64-byte (cache line) --- alignment. This is essential for SIMD instructions. --- --- Additionally, we return a plain ForeignPtr, which unlike a regular ForeignPtr --- created with 'mallocForeignPtr' carries no finalisers. It is an error to try --- to add a finaliser to the plain ForeignPtr. For our purposes this is fine, --- since in Accelerate finalisers are handled using Lifetime --- -mallocPlainForeignPtrBytesAligned :: Int -> IO (ForeignPtr a) -mallocPlainForeignPtrBytesAligned (I# size#) = IO $ \s0 -> - case newAlignedPinnedByteArray# size# 64# s0 of - (# s1, mbarr# #) -> (# s1, ForeignPtr (byteArrayContents# (unsafeCoerce# mbarr#)) (PlainPtr mbarr#) #) - - -liftArrayData :: Int -> TypeR e -> ArrayData e -> CodeQ (ArrayData e) -liftArrayData n = tuple - where - tuple :: TypeR e -> ArrayData e -> CodeQ (ArrayData e) - tuple TupRunit () = [|| () ||] - tuple (TupRpair t1 t2) (a1, a2) = [|| ($$(tuple t1 a1), $$(tuple t2 a2)) ||] - tuple (TupRsingle s) adata = scalar s adata - - scalar :: ScalarType e -> ArrayData e -> CodeQ (ArrayData e) - scalar (SingleScalarType t) = single t - scalar (VectorScalarType t) = vector t - - vector :: forall n e. VectorType (Vec n e) -> ArrayData (Vec n e) -> CodeQ (ArrayData (Vec n e)) - vector (VectorType w t) - | SingleArrayDict <- singleArrayDict t - = liftArrayData (w * n) (TupRsingle (SingleScalarType t)) - - single :: SingleType e -> ArrayData e -> CodeQ (ArrayData e) - single (NumSingleType t) = num t - - num :: NumType e -> ArrayData e -> CodeQ (ArrayData e) - num (IntegralNumType t) = integral t - num (FloatingNumType t) = floating t - - integral :: IntegralType e -> ArrayData e -> CodeQ (ArrayData e) - integral TypeInt = liftUniqueArray n - integral TypeInt8 = liftUniqueArray n - integral TypeInt16 = liftUniqueArray n - integral TypeInt32 = liftUniqueArray n - integral TypeInt64 = liftUniqueArray n - integral TypeWord = liftUniqueArray n - integral TypeWord8 = liftUniqueArray n - integral TypeWord16 = liftUniqueArray n - integral TypeWord32 = liftUniqueArray n - integral TypeWord64 = liftUniqueArray n - - floating :: FloatingType e -> ArrayData e -> CodeQ (ArrayData e) - floating TypeHalf = liftUniqueArray n - floating TypeFloat = liftUniqueArray n - floating TypeDouble = liftUniqueArray n - --- Determine the underlying type of a Haskell CLong or CULong. --- -runQ [d| type HTYPE_INT = $( - case finiteBitSize (undefined::Int) of - 32 -> [t| Int32 |] - 64 -> [t| Int64 |] - _ -> error "I don't know what architecture I am" ) |] - -runQ [d| type HTYPE_WORD = $( - case finiteBitSize (undefined::Word) of - 32 -> [t| Word32 |] - 64 -> [t| Word64 |] - _ -> error "I don't know what architecture I am" ) |] - -runQ [d| type HTYPE_CLONG = $( - case finiteBitSize (undefined::CLong) of - 32 -> [t| Int32 |] - 64 -> [t| Int64 |] - _ -> error "I don't know what architecture I am" ) |] - -runQ [d| type HTYPE_CULONG = $( - case finiteBitSize (undefined::CULong) of - 32 -> [t| Word32 |] - 64 -> [t| Word64 |] - _ -> error "I don't know what architecture I am" ) |] - -runQ [d| type HTYPE_CCHAR = $( - if isSigned (undefined::CChar) - then [t| Int8 |] - else [t| Word8 |] ) |] - diff --git a/src/Data/Array/Accelerate/Array/Remote/Class.hs b/src/Data/Array/Accelerate/Array/Remote/Class.hs index 7a871bd1a..84c1ed320 100644 --- a/src/Data/Array/Accelerate/Array/Remote/Class.hs +++ b/src/Data/Array/Accelerate/Array/Remote/Class.hs @@ -1,5 +1,4 @@ {-# LANGUAGE AllowAmbiguousTypes #-} -{-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE TypeFamilies #-} {-# OPTIONS_HADDOCK hide #-} -- | @@ -31,7 +30,7 @@ module Data.Array.Accelerate.Array.Remote.Class ( ) where -import Data.Array.Accelerate.Array.Data +import Data.Array.Accelerate.Array.Buffer import Data.Array.Accelerate.Type import Control.Applicative @@ -54,10 +53,10 @@ class (Applicative m, Monad m, MonadCatch m, MonadMask m) => RemoteMemory m wher mallocRemote :: Int -> m (Maybe (RemotePtr m Word8)) -- | Copy the given number of elements from the host array into remote memory. - pokeRemote :: SingleType e -> Int -> RemotePtr m (ScalarArrayDataR e) -> ArrayData e -> m () + pokeRemote :: SingleType e -> Int -> RemotePtr m e -> Buffer e -> m () -- | Copy the given number of elements from remote memory to the host array. - peekRemote :: SingleType e -> Int -> RemotePtr m (ScalarArrayDataR e) -> MutableArrayData e -> m () + peekRemote :: SingleType e -> Int -> RemotePtr m e -> MutableBuffer e -> m () -- | Cast a remote pointer. castRemotePtr :: RemotePtr m a -> RemotePtr m b diff --git a/src/Data/Array/Accelerate/Array/Remote/LRU.hs b/src/Data/Array/Accelerate/Array/Remote/LRU.hs index 03d32c50b..cfb2a4040 100644 --- a/src/Data/Array/Accelerate/Array/Remote/LRU.hs +++ b/src/Data/Array/Accelerate/Array/Remote/LRU.hs @@ -1,16 +1,5 @@ {-# LANGUAGE AllowAmbiguousTypes #-} -{-# LANGUAGE BangPatterns #-} {-# LANGUAGE CPP #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE DoAndIfThenElse #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE PatternGuards #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeOperators #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.Array.Remote.LRU @@ -39,10 +28,10 @@ module Data.Array.Accelerate.Array.Remote.LRU ( ) where import Data.Array.Accelerate.Analysis.Match ( matchSingleType, (:~:)(..) ) -import Data.Array.Accelerate.Array.Data +import Data.Array.Accelerate.Array.Buffer +-- import Data.Array.Accelerate.Array.Data import Data.Array.Accelerate.Array.Remote.Class -import Data.Array.Accelerate.Array.Remote.Table ( StableArray, makeWeakArrayData, formatStableArray ) -import Data.Array.Accelerate.Array.Unique ( touchUniqueArray ) +import Data.Array.Accelerate.Array.Remote.Table ( StableBuffer, makeWeakArrayData, formatStableBuffer ) import Data.Array.Accelerate.Error ( internalError ) import Data.Array.Accelerate.Type import qualified Data.Array.Accelerate.Array.Remote.Table as Basic @@ -75,7 +64,7 @@ data MemoryTable p task = MemoryTable {-# UNPACK #-} !(Basic.MemoryTable p) {-# UNPACK #-} !(UseTable task) {-# UNPACK #-} !(Weak (UseTable task)) -type UT task = HT.CuckooHashTable StableArray (Used task) +type UT task = HT.CuckooHashTable StableBuffer (Used task) type UseTable task = MVar (UT task) data Status = Clean -- Array in remote memory matches array in host memory. @@ -88,14 +77,14 @@ data Status = Clean -- Array in remote memory matches array in host memory. type Timestamp = Integer data Used task where - Used :: ArrayData e ~ ScalarArrayData e + Used :: Buffers e ~ Buffer e => !Timestamp -> !Status -> {-# UNPACK #-} !Int -- Use count -> ![task] -- Asynchronous tasks using the array -> {-# UNPACK #-} !Int -- Number of elements -> !(SingleType e) - -> {-# UNPACK #-} !(Weak (ScalarArrayData e)) + -> {-# UNPACK #-} !(Weak (MutableBuffer e)) -> Used task -- | A Task represents a process executing asynchronously that can be polled for @@ -141,11 +130,11 @@ withRemote :: forall task m a c. (HasCallStack, Task task, RemoteMemory m, MonadIO m, Functor m) => MemoryTable (RemotePtr m) task -> SingleType a - -> ArrayData a + -> MutableBuffer a -> (RemotePtr m (ScalarArrayDataR a) -> m (task, c)) -> m (Maybe c) withRemote (MemoryTable !mt !ref _) !tp !arr run | SingleArrayDict <- singleArrayDict tp = do - key <- Basic.makeStableArray tp arr + key <- Basic.makeStableBuffer tp arr mp <- withMVar' ref $ \utbl -> do mu <- liftIO . HT.mutate utbl key $ \case Nothing -> (Nothing, Nothing) @@ -153,7 +142,7 @@ withRemote (MemoryTable !mt !ref _) !tp !arr run | SingleArrayDict <- singleArra -- case mu of Nothing -> do - message ("withRemote/array has never been malloc'd: " % formatStableArray) key + message ("withRemote/array has never been malloc'd: " % formatStableBuffer) key return Nothing -- The array was never in the table Just u -> do @@ -162,7 +151,7 @@ withRemote (MemoryTable !mt !ref _) !tp !arr run | SingleArrayDict <- singleArra Just p -> return p Nothing | isEvicted u -> copyBack utbl (incCount u) - | otherwise -> do message ("lost array " % formatStableArray) key + | otherwise -> do message ("lost array " % formatStableBuffer) key internalError "non-evicted array has been lost" return (Just ptr) -- @@ -181,7 +170,7 @@ withRemote (MemoryTable !mt !ref _) !tp !arr run | SingleArrayDict <- singleArra | Just Refl <- matchSingleType tp tp' = do message "withRemote/reuploading-evicted-array" p <- mallocWithUsage mt utbl tp arr (Used ts Clean count tasks n tp weak_arr) - pokeRemote tp n p arr + pokeRemote tp n p (unsafeFreezeBuffer arr) return p | otherwise = internalError "Type mismatch" @@ -189,12 +178,12 @@ withRemote (MemoryTable !mt !ref _) !tp !arr run | SingleArrayDict <- singleArra -- because the `permute` operation from the PTX backend requires nested -- calls to `withRemote` in order to copy the defaults array. -- - go :: (HasCallStack, ArrayData a ~ ScalarArrayData a) - => StableArray + go :: (HasCallStack, Buffer a ~ Buffers a) + => StableBuffer -> RemotePtr m (ScalarArrayDataR a) -> m c go key ptr = do - message ("withRemote/using: " % formatStableArray) key + message ("withRemote/using: " % formatStableBuffer) key (task, c) <- run ptr liftIO . withMVar ref $ \utbl -> do HT.mutateIO utbl key $ \case @@ -203,7 +192,7 @@ withRemote (MemoryTable !mt !ref _) !tp !arr run | SingleArrayDict <- singleArra u' <- updateTask u task return (Just u', ()) -- - touchUniqueArray arr + touchMutableBuffer arr return c @@ -225,13 +214,13 @@ withRemote (MemoryTable !mt !ref _) !tp !arr run | SingleArrayDict <- singleArra malloc :: forall e m task. (HasCallStack, RemoteMemory m, MonadIO m, Task task) => MemoryTable (RemotePtr m) task -> SingleType e - -> ArrayData e + -> MutableBuffer e -> Bool -- ^ True if host array is frozen. -> Int -- ^ Number of elements -> m Bool -- ^ Was the array allocated successfully? -malloc (MemoryTable mt ref weak_utbl) !tp !ad !frozen !n | SingleArrayDict <- singleArrayDict tp = do -- Required for ArrayData e ~ ScalarArrayData e +malloc (MemoryTable mt ref weak_utbl) !tp !ad !frozen !n | SingleArrayDict <- singleArrayDict tp = do -- Required for Buffer e ~ Buffers e ts <- liftIO $ getCPUTime - key <- Basic.makeStableArray tp ad + key <- Basic.makeStableBuffer tp ad -- let status = if frozen then Clean @@ -248,11 +237,11 @@ malloc (MemoryTable mt ref weak_utbl) !tp !ad !frozen !n | SingleArrayDict <- si return False mallocWithUsage - :: forall e m task. (HasCallStack, RemoteMemory m, MonadIO m, Task task, ArrayData e ~ ScalarArrayData e) + :: forall e m task. (HasCallStack, RemoteMemory m, MonadIO m, Task task, Buffers e ~ Buffer e) => Basic.MemoryTable (RemotePtr m) -> UT task -> SingleType e - -> ArrayData e + -> MutableBuffer e -> Used task -> m (RemotePtr m (ScalarArrayDataR e)) mallocWithUsage !mt !utbl !tp !ad !usage@(Used _ _ _ _ n _ _) = malloc' @@ -266,7 +255,7 @@ mallocWithUsage !mt !utbl !tp !ad !usage@(Used _ _ _ _ n _ _) = malloc' if success then malloc' else internalError "Remote memory exhausted" Just p -> liftIO $ do - key <- Basic.makeStableArray tp ad + key <- Basic.makeStableBuffer tp ad HT.insert utbl key usage return p @@ -294,7 +283,7 @@ evictLRU !utbl !mt = trace "evictLRU/evicting-eldest-array" $ do message "evictLRU/Accelerate GC interrupted by GHC GC" Just arr -> do - message ("evictLRU/evicting " % formatStableArray) sa + message ("evictLRU/evicting " % formatStableBuffer) sa copyIfNecessary status n tp arr -- liftIO $ Debug.remote_memory_evict sa (remoteBytes tp n) liftIO $ Basic.freeStable @m mt sa @@ -303,7 +292,7 @@ evictLRU !utbl !mt = trace "evictLRU/evicting-eldest-array" $ do _ -> trace "evictLRU/All arrays in use, unable to evict" $ return False where -- Find the eldest, not currently in use, array. - eldest :: (Maybe (StableArray, Used task)) -> (StableArray, Used task) -> IO (Maybe (StableArray, Used task)) + eldest :: (Maybe (StableBuffer, Used task)) -> (StableBuffer, Used task) -> IO (Maybe (StableBuffer, Used task)) eldest prev (sa, used@(Used ts status count tasks n tp weak_arr)) | count == 0 , evictable status @@ -326,7 +315,7 @@ evictLRU !utbl !mt = trace "evictLRU/evicting-eldest-array" $ do evictable Unmanaged = False evictable Evicted = False - copyIfNecessary :: Status -> Int -> SingleType e -> ArrayData e -> m () + copyIfNecessary :: Status -> Int -> SingleType e -> MutableBuffer e -> m () copyIfNecessary Clean _ _ _ = return () copyIfNecessary Unmanaged _ _ _ = return () copyIfNecessary Evicted _ _ _ = internalError "Attempting to evict already evicted array" @@ -343,12 +332,12 @@ evictLRU !utbl !mt = trace "evictLRU/evicting-eldest-array" $ do free :: forall m a task. (HasCallStack, RemoteMemory m) => MemoryTable (RemotePtr m) task -> SingleType a - -> ArrayData a + -> MutableBuffer a -> IO () free (MemoryTable !mt !ref _) !tp !arr = withMVar' ref $ \utbl -> do - key <- Basic.makeStableArray tp arr + key <- Basic.makeStableBuffer tp arr delete utbl key Basic.freeStable @m mt key @@ -362,11 +351,11 @@ insertUnmanaged :: (HasCallStack, MonadIO m, RemoteMemory m) => MemoryTable (RemotePtr m) task -> SingleType e - -> ArrayData e + -> MutableBuffer e -> RemotePtr m (ScalarArrayDataR e) -> m () insertUnmanaged (MemoryTable mt ref weak_utbl) !tp !arr !ptr | SingleArrayDict <- singleArrayDict tp = do -- Gives evidence that ArrayData e ~ ScalarArrayData e - key <- Basic.makeStableArray tp arr + key <- Basic.makeStableBuffer tp arr () <- Basic.insertUnmanaged mt tp arr ptr liftIO $ withMVar ref @@ -379,14 +368,14 @@ insertUnmanaged (MemoryTable mt ref weak_utbl) !tp !arr !ptr | SingleArrayDict < -- Removing entries -- ---------------- -finalizer :: StableArray -> Weak (UseTable task) -> IO () +finalizer :: StableBuffer -> Weak (UseTable task) -> IO () finalizer !key !weak_utbl = do mref <- deRefWeak weak_utbl case mref of Nothing -> message "finalize cache/dead table" - Just ref -> trace (bformat ("finalize cache: " % formatStableArray) key) $ withMVar' ref (`delete` key) + Just ref -> trace (bformat ("finalize cache: " % formatStableBuffer) key) $ withMVar' ref (`delete` key) -delete :: UT task -> StableArray -> IO () +delete :: UT task -> StableBuffer -> IO () delete = HT.delete diff --git a/src/Data/Array/Accelerate/Array/Remote/Nursery.hs b/src/Data/Array/Accelerate/Array/Remote/Nursery.hs index 6baee1d15..23f80d65e 100644 --- a/src/Data/Array/Accelerate/Array/Remote/Nursery.hs +++ b/src/Data/Array/Accelerate/Array/Remote/Nursery.hs @@ -1,6 +1,3 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE OverloadedStrings #-} -- | -- Module : Data.Array.Accelerate.Array.Remote.Nursery -- Copyright : [2008..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Array/Remote/Table.hs b/src/Data/Array/Accelerate/Array/Remote/Table.hs index 39f0d5b00..1207c7315 100644 --- a/src/Data/Array/Accelerate/Array/Remote/Table.hs +++ b/src/Data/Array/Accelerate/Array/Remote/Table.hs @@ -1,19 +1,8 @@ -{-# LANGUAGE AllowAmbiguousTypes #-} -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE CPP #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE GeneralizedNewtypeDeriving #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE MagicHash #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE PatternGuards #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE UnboxedTuples #-} -{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE UnboxedTuples #-} +{-# LANGUAGE ViewPatterns #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.Array.Remote.Table @@ -36,15 +25,19 @@ module Data.Array.Accelerate.Array.Remote.Table ( MemoryTable, new, lookup, malloc, free, freeStable, insertUnmanaged, reclaim, -- Internals - StableArray, makeStableArray, + StableBuffer, makeStableBuffer, makeWeakArrayData, - formatStableArray, + formatStableBuffer, ) where +-- TODO: This is currently broken due to refactoring for the new pipeline. +-- Depending on the choices we make for the new GPU backends, this file should +-- either be removed, or reimplemented. + import Control.Concurrent ( yield ) import Control.Concurrent.MVar ( MVar, newMVar, withMVar, mkWeakMVar ) -import Control.Concurrent.Unique ( Unique ) +-- import Control.Concurrent.Unique ( Unique ) import Control.Monad.IO.Class ( MonadIO, liftIO ) import Data.Functor import Data.Hashable ( hash, Hashable ) @@ -58,13 +51,14 @@ import System.Mem ( performGC import System.Mem.Weak ( Weak, deRefWeak ) import qualified Data.HashTable.IO as HT -import Data.Array.Accelerate.Array.Data +import Data.Array.Accelerate.Error ( internalError ) +import Data.Array.Accelerate.Type +-- import Data.Array.Accelerate.Array.Unique ( UniqueArray(..) ) +import Data.Array.Accelerate.Array.Buffer +-- import Data.Array.Accelerate.Array.Data import Data.Array.Accelerate.Array.Remote.Class import Data.Array.Accelerate.Array.Remote.Nursery ( Nursery(..) ) -import Data.Array.Accelerate.Array.Unique ( UniqueArray(..) ) -import Data.Array.Accelerate.Error ( internalError ) import Data.Array.Accelerate.Lifetime -import Data.Array.Accelerate.Type import qualified Data.Array.Accelerate.Array.Remote.Nursery as N import qualified Data.Array.Accelerate.Debug.Internal.Flags as Debug import qualified Data.Array.Accelerate.Debug.Internal.Profile as Debug @@ -87,7 +81,7 @@ import GHC.Stack -- semantics of weak pointers in the documentation). -- type HashTable key val = HT.CuckooHashTable key val -type MT p = MVar ( HashTable StableArray (RemoteArray p) ) +type MT p = MVar ( HashTable StableBuffer (RemoteArray p) ) data MemoryTable p = MemoryTable {-# UNPACK #-} !(MT p) {-# UNPACK #-} !(Weak (MT p)) {-# UNPACK #-} !(Nursery p) @@ -99,17 +93,17 @@ data RemoteArray p where -> {-# UNPACK #-} !(Weak ()) -- Keep track of host array liveness -> RemoteArray p --- | An untyped reference to an array, similar to a StableName. +-- | An untyped reference to a buffer, similar to a StableName. -- -newtype StableArray = StableArray Unique +newtype StableBuffer = StableBuffer () -- Unique deriving (Eq, Hashable) -instance Show StableArray where - show (StableArray u) = show (hash u) +instance Show StableBuffer where + show (StableBuffer u) = show (hash u) -formatStableArray :: Format r (StableArray -> r) -formatStableArray = later $ \case - StableArray u -> bformat int (hash u) +formatStableBuffer :: Format r (StableBuffer -> r) +formatStableBuffer = later $ \case + StableBuffer u -> bformat int (hash u) -- | Create a new memory table from host to remote arrays. -- @@ -133,18 +127,18 @@ new release = do lookup :: forall m a. (HasCallStack, RemoteMemory m) => MemoryTable (RemotePtr m) -> SingleType a - -> ArrayData a + -> MutableBuffer a -> IO (Maybe (RemotePtr m (ScalarArrayDataR a))) lookup (MemoryTable !ref _ _ _) !tp !arr | SingleArrayDict <- singleArrayDict tp = do - sa <- makeStableArray tp arr + sa <- makeStableBuffer tp arr mw <- withMVar ref (`HT.lookup` sa) case mw of - Nothing -> trace (bformat ("lookup/not found: " % formatStableArray) sa) $ return Nothing + Nothing -> trace (bformat ("lookup/not found: " % formatStableBuffer) sa) $ return Nothing Just (RemoteArray p _ w) -> do mv <- deRefWeak w case mv of - Just{} -> trace (bformat ("lookup/found: " % formatStableArray) sa) $ return (Just $ castRemotePtr @m p) + Just{} -> trace (bformat ("lookup/found: " % formatStableBuffer) sa) $ return (Just $ castRemotePtr @m p) -- Note: [Weak pointer weirdness] -- @@ -159,7 +153,7 @@ lookup (MemoryTable !ref _ _ _) !tp !arr -- above in the error message. -- Nothing -> - makeStableArray tp arr >>= \x -> internalError ("dead weak pair: " % formatStableArray) x + makeStableBuffer tp arr >>= \x -> internalError ("dead weak pair: " % formatStableBuffer) x -- | Allocate a new device array to be associated with the given host-side array. -- This may not always use the `malloc` provided by the `RemoteMemory` instance. @@ -170,7 +164,7 @@ lookup (MemoryTable !ref _ _ _) !tp !arr malloc :: forall a m. (HasCallStack, RemoteMemory m, MonadIO m) => MemoryTable (RemotePtr m) -> SingleType a - -> ArrayData a + -> MutableBuffer a -> Int -> m (Maybe (RemotePtr m (ScalarArrayDataR a))) malloc mt@(MemoryTable _ _ !nursery _) !tp !ad !n @@ -234,30 +228,30 @@ malloc mt@(MemoryTable _ _ !nursery _) !tp !ad !n free :: forall m a. (RemoteMemory m) => MemoryTable (RemotePtr m) -> SingleType a - -> ArrayData a + -> MutableBuffer a -> IO () free mt tp !arr = do - sa <- makeStableArray tp arr + sa <- makeStableBuffer tp arr freeStable @m mt sa --- | Deallocate the device array associated with the given StableArray. This +-- | Deallocate the device array associated with the given StableBuffer. This -- is useful for other memory managers built on top of the memory table. -- freeStable :: forall m. RemoteMemory m => MemoryTable (RemotePtr m) - -> StableArray + -> StableBuffer -> IO () freeStable (MemoryTable !ref _ !nrs _) !sa = withMVar ref $ \mt -> HT.mutateIO mt sa $ \mw -> do case mw of Nothing -> - message ("free/already-removed: " % formatStableArray) sa + message ("free/already-removed: " % formatStableBuffer) sa Just (RemoteArray !p !n _) -> do - message ("free/nursery: " % formatStableArray % " of " % bytes') sa n + message ("free/nursery: " % formatStableBuffer % " of " % bytes') sa n N.insert n (castRemotePtr @m p) nrs -- Debug.remote_memory_free (unsafeRemotePtrToPtr @m p) @@ -272,16 +266,16 @@ insert :: forall m a. (RemoteMemory m, MonadIO m) => MemoryTable (RemotePtr m) -> SingleType a - -> ArrayData a + -> MutableBuffer a -> RemotePtr m (ScalarArrayDataR a) -> Int -> m () -insert mt@(MemoryTable !ref _ _ _) !tp !arr !ptr !n | SingleArrayDict <- singleArrayDict tp = do - key <- makeStableArray tp arr +insert mt@(MemoryTable !ref _ _ _) !tp !arr !ptr !byteSize | SingleArrayDict <- singleArrayDict tp = do + key <- makeStableBuffer tp arr weak <- liftIO $ makeWeakArrayData tp arr () (Just $ freeStable @m mt key) - message ("insert: " % formatStableArray) key + message ("insert: " % formatStableBuffer) key -- liftIO $ Debug.remote_memory_alloc (unsafeRemotePtrToPtr @m ptr) n - liftIO $ withMVar ref $ \tbl -> HT.insert tbl key (RemoteArray (castRemotePtr @m ptr) n weak) + liftIO $ withMVar ref $ \tbl -> HT.insert tbl key (RemoteArray (castRemotePtr @m ptr) byteSize weak) -- | Record an association between a host-side array and a remote memory area @@ -294,13 +288,13 @@ insertUnmanaged :: forall m a. (MonadIO m, RemoteMemory m) => MemoryTable (RemotePtr m) -> SingleType a - -> ArrayData a + -> MutableBuffer a -> RemotePtr m (ScalarArrayDataR a) -> m () insertUnmanaged (MemoryTable !ref !weak_ref _ _) tp !arr !ptr | SingleArrayDict <- singleArrayDict tp = do - key <- makeStableArray tp arr + key <- makeStableBuffer tp arr weak <- liftIO $ makeWeakArrayData tp arr () (Just $ remoteFinalizer weak_ref key) - message ("insertUnmanaged: " % formatStableArray) key + message ("insertUnmanaged: " % formatStableBuffer) key liftIO $ withMVar ref $ \tbl -> HT.insert tbl key (RemoteArray (castRemotePtr @m ptr) 0 weak) @@ -350,28 +344,28 @@ purge (MemoryTable _ _ nursery@(Nursery nrs _) release) reclaim :: forall m. (RemoteMemory m, MonadIO m) => MemoryTable (RemotePtr m) -> m () reclaim mt = clean mt >> purge mt -remoteFinalizer :: Weak (MT p) -> StableArray -> IO () +remoteFinalizer :: Weak (MT p) -> StableBuffer -> IO () remoteFinalizer !weak_ref !key = do mr <- deRefWeak weak_ref case mr of - Nothing -> message ("finalise/dead table: " % formatStableArray) key - Just ref -> trace (bformat ("finalise: " % formatStableArray) key) $ withMVar ref (`HT.delete` key) + Nothing -> message ("finalise/dead table: " % formatStableBuffer) key + Just ref -> trace (bformat ("finalise: " % formatStableBuffer) key) $ withMVar ref (`HT.delete` key) -- Miscellaneous -- ------------- --- | Make a new 'StableArray'. +-- | Make a new 'StableBuffer'. -- -{-# INLINE makeStableArray #-} -makeStableArray +{-# INLINE makeStableBuffer #-} +makeStableBuffer :: MonadIO m => SingleType a - -> ArrayData a - -> m StableArray -makeStableArray !tp !ad + -> MutableBuffer a + -> m StableBuffer +makeStableBuffer !tp !(MutableBuffer ad) | SingleArrayDict <- singleArrayDict tp - = return $! StableArray (uniqueArrayId ad) + = return $! StableBuffer $ error "Not reimplemented yet for new allocator" -- (uniqueArrayId ad) -- Weak arrays @@ -383,16 +377,17 @@ makeStableArray !tp !ad makeWeakArrayData :: forall e c. SingleType e - -> ArrayData e + -> MutableBuffer e -> c -> Maybe (IO ()) -> IO (Weak c) -makeWeakArrayData !tp !ad !c !mf | SingleArrayDict <- singleArrayDict tp = do - let !uad = uniqueArrayData ad +makeWeakArrayData !tp !(MutableBuffer ad) !c !mf | SingleArrayDict <- singleArrayDict tp = do + {- let !uad = uniqueArrayData ad case mf of Nothing -> return () Just f -> addFinalizer uad f - mkWeak uad c + mkWeak uad c -} + error "Not reimplemented yet for new allocator" -- Debug diff --git a/src/Data/Array/Accelerate/Backend.hs b/src/Data/Array/Accelerate/Backend.hs new file mode 100644 index 000000000..320f676e5 --- /dev/null +++ b/src/Data/Array/Accelerate/Backend.hs @@ -0,0 +1,173 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE TypeFamilies #-} + +-- | +-- Module : Data.Array.Accelerate.Backend +-- Copyright : [2012..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- +-- This module defines the interface for kernels, to be implemented +-- in other modules. +-- + +module Data.Array.Accelerate.Backend ( + Backend(..), + runAt, runWithAt, + run1At, run1WithAt, + runNAt, runNWithAt, + inspectCompiler, + + -- Type classes that a backend should implement: + Lowering.LowerAcc(..), + + Operation.SLVOperation(..), + Operation.SubArgs(..), Operation.SubArg(..), + Operation.defaultSlvGenerate, Operation.defaultSlvMap, Operation.defaultSlvBackpermute, + Operation.EncodeOperation(..), Operation.hashOperation, + + Operation.SimplifyOperation(..), + Operation.CopyOperation(..), + Operation.detectMapCopies, Operation.detectBackpermuteCopies, Operation.copyOperationsForArray, + + Pretty.PrettyKernel(..), Pretty.PrettyKernelStyle(..), + + Partitioning.MakesILP(..), + IsSchedule(..), + IsKernel(..), + Pretty.PrettyOp(..), + Execute(..), + Operation.NFData'(..), + Operation.ShrinkArg(..), -- runWithObj, runNWithObj, runNBench, Benchmarking +) where + +import qualified Data.Array.Accelerate.Smart as Smart +import Data.Array.Accelerate.AST.Schedule +import Data.Array.Accelerate.AST.Kernel +import Data.Array.Accelerate.AST.Execute +import qualified Data.Array.Accelerate.Sugar.Array as Sugar +import Data.Array.Accelerate.Representation.Array +import Data.Array.Accelerate.Representation.Ground +import Data.Array.Accelerate.Trafo +import Data.Array.Accelerate.Trafo.Config +import qualified Data.Array.Accelerate.Trafo.Operation.LiveVars as Operation +import qualified Data.Array.Accelerate.Trafo.Operation.Simplify as Operation +import qualified Data.Array.Accelerate.Analysis.Hash.Operation as Operation + +import Data.Array.Accelerate.Trafo.Sharing (Afunction(..), AfunctionRepr(..), afunctionGroundR, afunctionRepr) +import qualified Data.Array.Accelerate.Trafo.Lowering as Lowering +import qualified Data.Array.Accelerate.Trafo.Partitioning.ILP.Graph as Partitioning +import qualified Data.Array.Accelerate.Pretty.Operation as Pretty +import qualified Data.Array.Accelerate.Pretty.Schedule as Pretty +import Data.Kind +import Data.Type.Equality +import System.IO.Unsafe (unsafePerformIO) +import qualified Data.Array.Accelerate.AST.Operation as Operation + + +class + ( Trafo (Schedule backend) (Kernel backend) + , Execute (Schedule backend) (Kernel backend) + ) => Backend backend where + + type Schedule backend :: (Type -> Type) -> Type -> Type -> Type + type Kernel backend :: Type -> Type + +{- TODO WALL: DEAD CODE +type Operation backend = KernelOperation (Kernel backend) +-} + +-- Backend can be chosen with an explicit type application, for instance: +-- run @Interpreter acc +runAt :: forall backend t. (Sugar.Arrays t, Backend backend) => Smart.Acc t -> t +runAt = runWithAt @backend defaultOptions + +runWithAt :: forall backend t. (Sugar.Arrays t, Backend backend) => Config -> Smart.Acc t -> t +runWithAt = runNWithAt @backend + +{- runWithObj :: forall backend t. (Sugar.Arrays t, Backend backend) => Objective -> Smart.Acc t -> t +runWithObj obj acc = Sugar.toArr $ sugarArrays repr $ unsafePerformIO $ executeAcc (lowerArraysR repr) program + where + repr = Sugar.arraysR @t + schedule = convertAccWithObj @(Schedule backend) @(Kernel backend) obj acc + program = linkAfunSchedule schedule -} + +run1At + :: forall backend s t. + (Sugar.Arrays s, Sugar.Arrays t, Backend backend) + => (Smart.Acc s -> Smart.Acc t) + -> s -> t +run1At = runNAt @backend + +run1WithAt + :: forall backend s t. + (Sugar.Arrays s, Sugar.Arrays t, Backend backend) + => Config + -> (Smart.Acc s -> Smart.Acc t) + -> s -> t +run1WithAt = runNWithAt @backend + +runNAt + :: forall backend f. + (Afunction f, Backend backend) + => f + -> AfunctionR f +runNAt = runNWithAt @backend defaultOptions + +runNWithAt + :: forall backend f. + (Afunction f, Backend backend) + => Config + -> f + -> AfunctionR f +runNWithAt config f = program `seq` sugarFunction (afunctionRepr @f) $ executeAfun (afunctionGroundR @f) program + where + schedule = convertAfunWith @(Schedule backend) @(Kernel backend) config f + program = linkAfunSchedule schedule + +{- runNWithObj + :: forall backend f. + (Afunction f, Backend backend) + => Objective + -> f + -> AfunctionR f +runNWithObj obj f = program `seq` sugarFunction (afunctionRepr @f) $ executeAfun (afunctionGroundR @f) program + where + schedule = convertAfunWithObj @(Schedule backend) @(Kernel backend) obj f + program = linkAfunSchedule schedule -} + +{- runNBench + :: forall backend f. + (Afunction f, Backend backend) + => Benchmarking + -> f + -> AfunctionR f +runNBench b f = program `seq` sugarFunction (afunctionRepr @f) $ executeAfun (afunctionGroundR @f) program + where + schedule = convertAfunBench @(Schedule backend) @(Kernel backend) b f + program = linkAfunSchedule schedule -} + +inspectCompiler + :: forall backend f. + (Afunction f, Backend backend) + => f + -> String +inspectCompiler = inspectCompiler' @(Schedule backend) @(Kernel backend) + +sugarFunction + :: forall f. + AfunctionRepr f (AfunctionR f) (ArraysFunctionR f) + -> IOFun (LoweredAfun (ArraysFunctionR f)) + -> AfunctionR f +sugarFunction AfunctionReprBody a + | Refl <- loweredAfunIsBody repr + , Refl <- reprIsBody (lowerArraysR repr) = Sugar.toArr $ sugarArrays repr $ unsafePerformIO a + where + repr = Sugar.arraysR @(AfunctionR f) +sugarFunction (AfunctionReprLam r) f = sugarFunction r . f . lowerArrays repr . Sugar.fromArr + where + repr :: forall a b. f ~ (Smart.Acc a -> b) => ArraysR (Sugar.ArraysR a) + repr = Sugar.arraysR @a diff --git a/src/Data/Array/Accelerate/Classes/Bounded.hs b/src/Data/Array/Accelerate/Classes/Bounded.hs index c567fb14c..d32c67c02 100644 --- a/src/Data/Array/Accelerate/Classes/Bounded.hs +++ b/src/Data/Array/Accelerate/Classes/Bounded.hs @@ -1,9 +1,6 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeApplications #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE TemplateHaskell #-} +{-# OPTIONS_GHC -Wno-orphans #-} -- | -- Module : Data.Array.Accelerate.Classes.Bounded -- Copyright : [2016..2020] The Accelerate Team @@ -21,7 +18,7 @@ module Data.Array.Accelerate.Classes.Bounded ( ) where -import Data.Array.Accelerate.Array.Data +import Data.Array.Accelerate.Array.Buffer import Data.Array.Accelerate.Pattern import Data.Array.Accelerate.Smart import Data.Array.Accelerate.Sugar.Elt @@ -133,7 +130,7 @@ instance P.Bounded (Exp CSChar) where instance P.Bounded (Exp CUChar) where minBound = mkBitcast (mkMinBound @Word8) maxBound = mkBitcast (mkMaxBound @Word8) - +#ifndef __GHCIDE__ $(runQ $ do let mkInstance :: Int -> Q [Dec] @@ -151,4 +148,4 @@ $(runQ $ do -- concat <$> mapM mkInstance [2..16] ) - +#endif diff --git a/src/Data/Array/Accelerate/Classes/Enum.hs b/src/Data/Array/Accelerate/Classes/Enum.hs index 84b344273..245407df2 100644 --- a/src/Data/Array/Accelerate/Classes/Enum.hs +++ b/src/Data/Array/Accelerate/Classes/Enum.hs @@ -1,8 +1,4 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MonoLocalBinds #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# OPTIONS_GHC -Wno-orphans #-} -- | -- Module : Data.Array.Accelerate.Classes.Enum -- Copyright : [2016..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Classes/Eq.hs b/src/Data/Array/Accelerate/Classes/Eq.hs index ca5ece850..e1009aca1 100644 --- a/src/Data/Array/Accelerate/Classes/Eq.hs +++ b/src/Data/Array/Accelerate/Classes/Eq.hs @@ -1,15 +1,6 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE ViewPatterns #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE TemplateHaskell #-} +{-# OPTIONS_GHC -Wno-orphans #-} -- | -- Module : Data.Array.Accelerate.Classes.Eq -- Copyright : [2016..2020] The Accelerate Team @@ -201,8 +192,8 @@ runQ $ do pat vs = conP (mkName ('T':show n)) (map varP vs) in [d| instance ($cst) => Eq $res where - $(pat xs) == $(pat ys) = $(foldr1 (\vs v -> [| $vs && $v |]) (zipWith (\x y -> [| $x == $y |]) (map varE xs) (map varE ys))) - $(pat xs) /= $(pat ys) = $(foldr1 (\vs v -> [| $vs || $v |]) (zipWith (\x y -> [| $x /= $y |]) (map varE xs) (map varE ys))) + $(pat xs) == $(pat ys) = $(foldr1 (\vs v -> [| $vs &&! $v |]) (zipWith (\x y -> [| $x == $y |]) (map varE xs) (map varE ys))) + $(pat xs) /= $(pat ys) = $(foldr1 (\vs v -> [| $vs ||! $v |]) (zipWith (\x y -> [| $x /= $y |]) (map varE xs) (map varE ys))) |] is <- mapM mkPrim integralTypes @@ -213,8 +204,8 @@ runQ $ do return $ concat (concat [is,fs,ns,cs,ts]) instance (Eq sh, Eq i) => Eq (sh :. i) where - x == y = indexHead x == indexHead y && indexTail x == indexTail y - x /= y = indexHead x /= indexHead y || indexTail x /= indexTail y + x == y = indexHead x == indexHead y &&! indexTail x == indexTail y + x /= y = indexHead x /= indexHead y ||! indexTail x /= indexTail y instance Eq Bool where x == y = mkCoerce x == (mkCoerce y :: Exp PrimBool) diff --git a/src/Data/Array/Accelerate/Classes/Floating.hs b/src/Data/Array/Accelerate/Classes/Floating.hs index b2f067af9..1c27745fd 100644 --- a/src/Data/Array/Accelerate/Classes/Floating.hs +++ b/src/Data/Array/Accelerate/Classes/Floating.hs @@ -1,9 +1,4 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeFamilies #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# OPTIONS_GHC -Wno-orphans #-} -- | -- Module : Data.Array.Accelerate.Classes.Floating -- Copyright : [2016..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Classes/Fractional.hs b/src/Data/Array/Accelerate/Classes/Fractional.hs index 52bdc61e9..f3b2c2fa0 100644 --- a/src/Data/Array/Accelerate/Classes/Fractional.hs +++ b/src/Data/Array/Accelerate/Classes/Fractional.hs @@ -1,8 +1,4 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE TypeFamilies #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# OPTIONS_GHC -Wno-orphans #-} -- | -- Module : Data.Array.Accelerate.Classes.Fractional -- Copyright : [2016..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Classes/FromIntegral.hs b/src/Data/Array/Accelerate/Classes/FromIntegral.hs index d8678ee9d..ad4782e90 100644 --- a/src/Data/Array/Accelerate/Classes/FromIntegral.hs +++ b/src/Data/Array/Accelerate/Classes/FromIntegral.hs @@ -1,8 +1,4 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE MonoLocalBinds #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TemplateHaskell #-} -- | -- Module : Data.Array.Accelerate.Classes.FromIntegral -- Copyright : [2016..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Classes/Integral.hs b/src/Data/Array/Accelerate/Classes/Integral.hs index c6cadf7cc..3c24448c3 100644 --- a/src/Data/Array/Accelerate/Classes/Integral.hs +++ b/src/Data/Array/Accelerate/Classes/Integral.hs @@ -1,9 +1,4 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MonoLocalBinds #-} -{-# LANGUAGE TypeFamilies #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# OPTIONS_GHC -Wno-orphans #-} -- | -- Module : Data.Array.Accelerate.Classes.Integral -- Copyright : [2016..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Classes/Num.hs b/src/Data/Array/Accelerate/Classes/Num.hs index 942b38673..b1487891b 100644 --- a/src/Data/Array/Accelerate/Classes/Num.hs +++ b/src/Data/Array/Accelerate/Classes/Num.hs @@ -1,8 +1,4 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE TypeFamilies #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# OPTIONS_GHC -Wno-orphans #-} -- | -- Module : Data.Array.Accelerate.Classes.Num -- Copyright : [2016..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Classes/Ord.hs b/src/Data/Array/Accelerate/Classes/Ord.hs index fa80dd9d2..fa97c3fec 100644 --- a/src/Data/Array/Accelerate/Classes/Ord.hs +++ b/src/Data/Array/Accelerate/Classes/Ord.hs @@ -1,16 +1,7 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE RebindableSyntax #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE ViewPatterns #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE RebindableSyntax #-} +{-# LANGUAGE TemplateHaskell #-} +{-# OPTIONS_GHC -Wno-orphans #-} -- | -- Module : Data.Array.Accelerate.Classes.Ord -- Copyright : [2016..2020] The Accelerate Team @@ -43,6 +34,7 @@ import Data.Array.Accelerate.Classes.Eq hiding ( (== import qualified Data.Array.Accelerate.Classes.Eq as A import Data.Char +import Data.String ( IsString(fromString) ) import Language.Haskell.TH.Extra hiding ( Exp ) import Prelude ( ($), (>>=), Ordering(..), Num(..), Maybe(..), String, show, error, unlines, return, concat, map, mapM ) import Text.Printf @@ -186,22 +178,22 @@ runQ $ do mkLt' :: [ExpQ] -> [ExpQ] -> ExpQ mkLt' [x] [y] = [| $x < $y |] - mkLt' (x:xs) (y:ys) = [| $x < $y || ( $x A.== $y && $(mkLt' xs ys) ) |] + mkLt' (x:xs) (y:ys) = [| $x < $y ||! ( $x A.== $y &&! $(mkLt' xs ys) ) |] mkLt' _ _ = error "mkLt'" mkGt' :: [ExpQ] -> [ExpQ] -> ExpQ mkGt' [x] [y] = [| $x > $y |] - mkGt' (x:xs) (y:ys) = [| $x > $y || ( $x A.== $y && $(mkGt' xs ys) ) |] + mkGt' (x:xs) (y:ys) = [| $x > $y ||! ( $x A.== $y &&! $(mkGt' xs ys) ) |] mkGt' _ _ = error "mkGt'" mkLtEq' :: [ExpQ] -> [ExpQ] -> ExpQ mkLtEq' [x] [y] = [| $x < $y |] - mkLtEq' (x:xs) (y:ys) = [| $x < $y || ( $x A.== $y && $(mkLtEq' xs ys) ) |] + mkLtEq' (x:xs) (y:ys) = [| $x < $y ||! ( $x A.== $y &&! $(mkLtEq' xs ys) ) |] mkLtEq' _ _ = error "mkLtEq'" mkGtEq' :: [ExpQ] -> [ExpQ] -> ExpQ mkGtEq' [x] [y] = [| $x > $y |] - mkGtEq' (x:xs) (y:ys) = [| $x > $y || ( $x A.== $y && $(mkGtEq' xs ys) ) |] + mkGtEq' (x:xs) (y:ys) = [| $x > $y ||! ( $x A.== $y &&! $(mkGtEq' xs ys) ) |] mkGtEq' _ _ = error "mkGtEq'" mkTup :: Int -> Q [Dec] @@ -228,14 +220,14 @@ runQ $ do return $ concat (concat [is,fs,ns,cs,ts]) instance Ord sh => Ord (sh :. Int) where - x <= y = indexHead x <= indexHead y && indexTail x <= indexTail y - x >= y = indexHead x >= indexHead y && indexTail x >= indexTail y + x <= y = indexHead x <= indexHead y &&! indexTail x <= indexTail y + x >= y = indexHead x >= indexHead y &&! indexTail x >= indexTail y x < y = indexHead x < indexHead y - && case matchTypeR (eltR @sh) (eltR @Z) of + &&! case matchTypeR (eltR @sh) (eltR @Z) of Just Refl -> constant True Nothing -> indexTail x < indexTail y x > y = indexHead x > indexHead y - && case matchTypeR (eltR @sh) (eltR @Z) of + &&! case matchTypeR (eltR @sh) (eltR @Z) of Just Refl -> constant True Nothing -> indexTail x > indexTail y diff --git a/src/Data/Array/Accelerate/Classes/Rational.hs b/src/Data/Array/Accelerate/Classes/Rational.hs index 7ec238459..6d7a2fc0d 100644 --- a/src/Data/Array/Accelerate/Classes/Rational.hs +++ b/src/Data/Array/Accelerate/Classes/Rational.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE FlexibleContexts #-} -- | -- Module : Data.Array.Accelerate.Classes.Rational -- Copyright : [2016..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Classes/Real.hs b/src/Data/Array/Accelerate/Classes/Real.hs index a6bd1b185..2e24d3229 100644 --- a/src/Data/Array/Accelerate/Classes/Real.hs +++ b/src/Data/Array/Accelerate/Classes/Real.hs @@ -1,10 +1,5 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MonoLocalBinds #-} {-# LANGUAGE UndecidableInstances #-} -{-# OPTIONS_GHC -fno-warn-missing-methods #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# OPTIONS_GHC -Wno-orphans #-} -- | -- Module : Data.Array.Accelerate.Classes.Real -- Copyright : [2016..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Classes/RealFloat.hs b/src/Data/Array/Accelerate/Classes/RealFloat.hs index 0b3366ec2..61bcb4fdd 100644 --- a/src/Data/Array/Accelerate/Classes/RealFloat.hs +++ b/src/Data/Array/Accelerate/Classes/RealFloat.hs @@ -1,14 +1,5 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE DefaultSignatures #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE NoImplicitPrelude #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE ViewPatterns #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# LANGUAGE DefaultSignatures #-} +{-# OPTIONS_GHC -Wno-orphans #-} -- | -- Module : Data.Array.Accelerate.Classes.RealFloat -- Copyright : [2016..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Classes/RealFloat.hs-boot b/src/Data/Array/Accelerate/Classes/RealFloat.hs-boot index a4f2878bd..332d3ca11 100644 --- a/src/Data/Array/Accelerate/Classes/RealFloat.hs-boot +++ b/src/Data/Array/Accelerate/Classes/RealFloat.hs-boot @@ -1,5 +1,4 @@ {-# LANGUAGE DefaultSignatures #-} -{-# LANGUAGE FlexibleContexts #-} -- | -- Module : Data.Array.Accelerate.Classes.RealFloat -- Copyright : [2019..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Classes/RealFrac.hs b/src/Data/Array/Accelerate/Classes/RealFrac.hs index 4cec0fde9..c2422a2d2 100644 --- a/src/Data/Array/Accelerate/Classes/RealFrac.hs +++ b/src/Data/Array/Accelerate/Classes/RealFrac.hs @@ -1,11 +1,5 @@ {-# LANGUAGE AllowAmbiguousTypes #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeFamilies #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# OPTIONS_GHC -Wno-orphans #-} -- | -- Module : Data.Array.Accelerate.Classes.RealFrac -- Copyright : [2016..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Classes/RealFrac.hs-boot b/src/Data/Array/Accelerate/Classes/RealFrac.hs-boot index 1bb85ec04..9426440b7 100644 --- a/src/Data/Array/Accelerate/Classes/RealFrac.hs-boot +++ b/src/Data/Array/Accelerate/Classes/RealFrac.hs-boot @@ -1,4 +1,3 @@ -{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE NoImplicitPrelude #-} -- | -- Module : Data.Array.Accelerate.Classes.RealFrac diff --git a/src/Data/Array/Accelerate/Classes/ToFloating.hs b/src/Data/Array/Accelerate/Classes/ToFloating.hs index c3f4545c6..a8b981d56 100644 --- a/src/Data/Array/Accelerate/Classes/ToFloating.hs +++ b/src/Data/Array/Accelerate/Classes/ToFloating.hs @@ -1,8 +1,4 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TemplateHaskell #-} -- | -- Module : Data.Array.Accelerate.Classes.ToFloating -- Copyright : [2016..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Data/Bits.hs b/src/Data/Array/Accelerate/Data/Bits.hs index 9696e3c3e..cd745b3c9 100644 --- a/src/Data/Array/Accelerate/Data/Bits.hs +++ b/src/Data/Array/Accelerate/Data/Bits.hs @@ -1,10 +1,3 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE ViewPatterns #-} -- | -- Module : Data.Array.Accelerate.Data.Bits -- Copyright : [2016..2020] The Accelerate Team @@ -24,7 +17,7 @@ module Data.Array.Accelerate.Data.Bits ( ) where -import Data.Array.Accelerate.Array.Data +import Data.Array.Accelerate.Array.Buffer import Data.Array.Accelerate.Language import Data.Array.Accelerate.Smart import Data.Array.Accelerate.Sugar.Elt diff --git a/src/Data/Array/Accelerate/Data/Complex.hs b/src/Data/Array/Accelerate/Data/Complex.hs index 1ac4ba771..560ed8f19 100644 --- a/src/Data/Array/Accelerate/Data/Complex.hs +++ b/src/Data/Array/Accelerate/Data/Complex.hs @@ -1,20 +1,10 @@ {-# LANGUAGE CPP #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE MagicHash #-} -{-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE PatternSynonyms #-} {-# LANGUAGE RebindableSyntax #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE TypeSynonymInstances #-} {-# LANGUAGE UndecidableInstances #-} {-# LANGUAGE ViewPatterns #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# OPTIONS_GHC -Wno-orphans #-} -- | -- Module : Data.Array.Accelerate.Data.Complex -- Copyright : [2015..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Data/Either.hs b/src/Data/Array/Accelerate/Data/Either.hs index 3c5c7401d..cd62217a2 100644 --- a/src/Data/Array/Accelerate/Data/Either.hs +++ b/src/Data/Array/Accelerate/Data/Either.hs @@ -1,19 +1,9 @@ -{-# LANGUAGE BlockArguments #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE PatternGuards #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE UndecidableInstances #-} -{-# LANGUAGE ViewPatterns #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} +{-# OPTIONS_GHC -Wno-orphans #-} -- | -- Module : Data.Array.Accelerate.Data.Either -- Copyright : [2018..2020] The Accelerate Team @@ -55,6 +45,10 @@ import Data.Array.Accelerate.Data.Semigroup import Data.Either ( Either(..) ) import Prelude ( (.), ($) ) +#ifdef ACCELERATE_INTERNAL_CHECKS +import Data.Typeable +#endif + -- | Return 'True' if the argument is a 'Left'-value -- @@ -95,7 +89,11 @@ either f g = match \case -- with a segment descriptor indicating how many elements along each dimension -- were returned. -- -lefts :: (Shape sh, Slice sh, Elt a, Elt b) +lefts :: ( +#ifdef ACCELERATE_INTERNAL_CHECKS + Typeable sh, +#endif + Shape sh, Slice sh, Elt a, Elt b) => Acc (Array (sh:.Int) (Either a b)) -> Acc (Vector a, Array sh Int) lefts es = compact (map isLeft es) (map fromLeft es) @@ -104,7 +102,11 @@ lefts es = compact (map isLeft es) (map fromLeft es) -- with a segment descriptor indicating how many elements along each dimension -- were returned. -- -rights :: (Shape sh, Slice sh, Elt a, Elt b) +rights :: ( +#ifdef ACCELERATE_INTERNAL_CHECKS + Typeable sh, +#endif + Shape sh, Slice sh, Elt a, Elt b) => Acc (Array (sh:.Int) (Either a b)) -> Acc (Vector b, Array sh Int) rights es = compact (map isRight es) (map fromRight es) diff --git a/src/Data/Array/Accelerate/Data/Fold.hs b/src/Data/Array/Accelerate/Data/Fold.hs index f38f770ac..50d4da8f7 100644 --- a/src/Data/Array/Accelerate/Data/Fold.hs +++ b/src/Data/Array/Accelerate/Data/Fold.hs @@ -1,8 +1,3 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE TypeOperators #-} -- | -- Module : Data.Array.Accelerate.Data.Fold -- Copyright : [2016..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Data/Functor.hs b/src/Data/Array/Accelerate/Data/Functor.hs index 411cb002c..7e1686e02 100644 --- a/src/Data/Array/Accelerate/Data/Functor.hs +++ b/src/Data/Array/Accelerate/Data/Functor.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE RebindableSyntax #-} -- | -- Module : Data.Array.Accelerate.Data.Functor -- Copyright : [2018..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Data/Maybe.hs b/src/Data/Array/Accelerate/Data/Maybe.hs index 60bcea2df..1f897d42b 100644 --- a/src/Data/Array/Accelerate/Data/Maybe.hs +++ b/src/Data/Array/Accelerate/Data/Maybe.hs @@ -1,19 +1,9 @@ -{-# LANGUAGE BlockArguments #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE PatternGuards #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE UndecidableInstances #-} -{-# LANGUAGE ViewPatterns #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} +{-# OPTIONS_GHC -Wno-orphans #-} -- | -- Module : Data.Array.Accelerate.Data.Maybe -- Copyright : [2018..2020] The Accelerate Team @@ -56,19 +46,16 @@ import Data.Function ( (&) ) import Data.Maybe ( Maybe(..) ) import Prelude ( ($), (.) ) +#ifdef ACCELERATE_INTERNAL_CHECKS +import Data.Typeable +#endif + -- | Returns 'True' if the argument is 'Nothing' -- isNothing :: Elt a => Exp (Maybe a) -> Exp Bool isNothing = not . isJust --- | Returns 'True' if the argument is of the form @Just _@ --- -isJust :: Elt a => Exp (Maybe a) -> Exp Bool -isJust (Exp x) = Exp $ SmartExp $ (SmartExp $ Prj PairIdxLeft x) `Pair` SmartExp Nil - -- TLM: This is a sneaky hack because we know that the tag bits for Just - -- and True are identical. - -- | The 'fromMaybe' function takes a default value and a 'Maybe' value. If the -- 'Maybe' is 'Nothing', the default value is returned; otherwise, it returns -- the value contained in the 'Maybe'. @@ -98,7 +85,11 @@ maybe d f = match \case -- | Extract from an array all of the 'Just' values, together with a segment -- descriptor indicating how many elements along each dimension were returned. -- -justs :: (Shape sh, Slice sh, Elt a) +justs :: ( +#ifdef ACCELERATE_INTERNAL_CHECKS + Typeable sh, +#endif + Shape sh, Slice sh, Elt a) => Acc (Array (sh:.Int) (Maybe a)) -> Acc (Vector a, Array sh Int) justs xs = compact (map isJust xs) (map fromJust xs) diff --git a/src/Data/Array/Accelerate/Data/Monoid.hs b/src/Data/Array/Accelerate/Data/Monoid.hs index ad9ffe130..840bf2e6c 100644 --- a/src/Data/Array/Accelerate/Data/Monoid.hs +++ b/src/Data/Array/Accelerate/Data/Monoid.hs @@ -1,16 +1,8 @@ -{-# LANGUAGE BlockArguments #-} -{-# LANGUAGE CPP #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE ViewPatterns #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE TypeFamilies #-} +{-# OPTIONS_GHC -Wno-orphans #-} #if __GLASGOW_HASKELL__ >= 806 {-# LANGUAGE UndecidableInstances #-} #endif diff --git a/src/Data/Array/Accelerate/Data/Ratio.hs b/src/Data/Array/Accelerate/Data/Ratio.hs index 190317297..fd0f2789e 100644 --- a/src/Data/Array/Accelerate/Data/Ratio.hs +++ b/src/Data/Array/Accelerate/Data/Ratio.hs @@ -1,13 +1,7 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MonoLocalBinds #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE RebindableSyntax #-} -{-# LANGUAGE StandaloneDeriving #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE UndecidableInstances #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE RebindableSyntax #-} +{-# LANGUAGE UndecidableInstances #-} +{-# OPTIONS_GHC -Wno-orphans #-} -- | -- Module : Data.Array.Accelerate.Data.Ratio -- Copyright : [2019..2020] The Accelerate Team @@ -48,6 +42,7 @@ import Data.Array.Accelerate.Classes.ToFloating import Text.Printf import Data.Ratio ( Ratio ) +import Data.String ( IsString(fromString) ) import Prelude ( ($), String, error, unlines ) import qualified Data.Ratio as P import qualified Prelude as P diff --git a/src/Data/Array/Accelerate/Data/Semigroup.hs b/src/Data/Array/Accelerate/Data/Semigroup.hs index e24f4f7c9..6cf184aa3 100644 --- a/src/Data/Array/Accelerate/Data/Semigroup.hs +++ b/src/Data/Array/Accelerate/Data/Semigroup.hs @@ -1,15 +1,8 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE RebindableSyntax #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE ViewPatterns #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE ViewPatterns #-} +{-# OPTIONS_GHC -Wno-orphans #-} #if __GLASGOW_HASKELL__ >= 806 {-# LANGUAGE UndecidableInstances #-} #endif diff --git a/src/Data/Array/Accelerate/Debug/Internal.hs b/src/Data/Array/Accelerate/Debug/Internal.hs index 598dc0091..a7961db2d 100644 --- a/src/Data/Array/Accelerate/Debug/Internal.hs +++ b/src/Data/Array/Accelerate/Debug/Internal.hs @@ -27,7 +27,9 @@ module Data.Array.Accelerate.Debug.Internal ( ) where import Data.Array.Accelerate.Debug.Internal.Flags as Debug +{- TODO WALL: DEAD CODE import Data.Array.Accelerate.Debug.Internal.Graph as Debug +-} import Data.Array.Accelerate.Debug.Internal.Profile as Debug import Data.Array.Accelerate.Debug.Internal.Stats as Debug import Data.Array.Accelerate.Debug.Internal.Timed as Debug diff --git a/src/Data/Array/Accelerate/Debug/Internal/Clock.hs b/src/Data/Array/Accelerate/Debug/Internal/Clock.hs index 0aba7b3df..14ff36e2f 100644 --- a/src/Data/Array/Accelerate/Debug/Internal/Clock.hs +++ b/src/Data/Array/Accelerate/Debug/Internal/Clock.hs @@ -1,6 +1,5 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE ForeignFunctionInterface #-} -{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE TemplateHaskell #-} {-# OPTIONS_GHC -fobject-code #-} -- | -- Module : Data.Array.Accelerate.Debug.Internal.Clock diff --git a/src/Data/Array/Accelerate/Debug/Internal/Flags.hs b/src/Data/Array/Accelerate/Debug/Internal/Flags.hs index 59cb8f0c7..67b67ae57 100644 --- a/src/Data/Array/Accelerate/Debug/Internal/Flags.hs +++ b/src/Data/Array/Accelerate/Debug/Internal/Flags.hs @@ -1,11 +1,9 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE ForeignFunctionInterface #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeOperators #-} -{-# OPTIONS_GHC -fno-warn-missing-signatures #-} -{-# OPTIONS_GHC -fno-warn-unused-imports #-} -{-# OPTIONS_GHC -fno-warn-unused-top-binds #-} -{-# OPTIONS_GHC -fobject-code #-} -- SEE: [linking to .c files] +{-# LANGUAGE CPP #-} +{-# LANGUAGE TemplateHaskell #-} +{-# OPTIONS_GHC -Wno-missing-signatures #-} +{-# OPTIONS_GHC -Wno-unused-imports #-} +{-# OPTIONS_GHC -Wno-unused-top-binds #-} +{-# OPTIONS_GHC -fobject-code #-} -- SEE: [linking to .c files] -- | -- Module : Data.Array.Accelerate.Debug.Internal.Flags -- Copyright : [2008..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Debug/Internal/Graph.hs b/src/Data/Array/Accelerate/Debug/Internal/Graph.hs index c237ffbca..9f46a9d55 100644 --- a/src/Data/Array/Accelerate/Debug/Internal/Graph.hs +++ b/src/Data/Array/Accelerate/Debug/Internal/Graph.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE CPP #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.Debug.Internal.Graph @@ -12,61 +11,61 @@ module Data.Array.Accelerate.Debug.Internal.Graph ( - dumpGraph + -- dumpGraph ) where -import Data.Array.Accelerate.Pretty.Graphviz +-- import Data.Array.Accelerate.Pretty.Graphviz -import Control.Monad.Trans ( MonadIO ) +-- import Control.Monad.Trans ( MonadIO ) -#ifdef ACCELERATE_DEBUG -import Control.Exception ( bracket ) -import Control.Monad.Trans ( liftIO ) -import System.Directory ( getTemporaryDirectory, createDirectoryIfMissing ) -import System.FilePath ( () ) -import System.IO ( Handle, openTempFile, hPutStrLn, hPrint, hClose, stderr ) -import qualified Data.Array.Accelerate.Debug.Internal.Flags as Debug +-- #ifdef ACCELERATE_DEBUG +-- import Control.Exception ( bracket ) +-- import Control.Monad.Trans ( liftIO ) +-- import System.Directory ( getTemporaryDirectory, createDirectoryIfMissing ) +-- import System.FilePath ( () ) +-- import System.IO ( Handle, openTempFile, hPutStrLn, hPrint, hClose, stderr ) +-- import qualified Data.Array.Accelerate.Debug.Internal.Flags as Debug -#if defined(mingw32_HOST_OS) -import System.Win32.Process ( ProcessId ) -#else -import System.Posix.Process ( getProcessID ) -#endif -#endif +-- #if defined(mingw32_HOST_OS) +-- import System.Win32.Process ( ProcessId ) +-- #else +-- import System.Posix.Process ( getProcessID ) +-- #endif +-- #endif --- | Write a representation of the given input (a closed array expression or --- function) to file in Graphviz dot format in the temporary directory. --- -{-# INLINEABLE dumpGraph #-} -dumpGraph :: (MonadIO m, PrettyGraph g) => g -> m () -#ifdef ACCELERATE_DEBUG -dumpGraph g = - liftIO $ do - Debug.when Debug.dump_dot $ writeGraph Full g - Debug.when Debug.dump_simpl_dot $ writeGraph Simple g -#else -dumpGraph _ = return () -#endif +-- -- | Write a representation of the given input (a closed array expression or +-- -- function) to file in Graphviz dot format in the temporary directory. +-- -- +-- {-# INLINEABLE dumpGraph #-} +-- dumpGraph :: (MonadIO m, PrettyGraph g) => g -> m () +-- #ifdef ACCELERATE_DEBUG +-- dumpGraph g = +-- liftIO $ do +-- Debug.when Debug.dump_dot $ writeGraph Full g +-- Debug.when Debug.dump_simpl_dot $ writeGraph Simple g +-- #else +-- dumpGraph _ = return () +-- #endif -#ifdef ACCELERATE_DEBUG -writeGraph :: PrettyGraph g => Detail -> g -> IO () -writeGraph simple g = do - withTemporaryFile "acc.dot" $ \path hdl -> do - hPrint hdl (ppGraph simple g) - hPutStrLn stderr ("program graph: " ++ path) +-- #ifdef ACCELERATE_DEBUG +-- writeGraph :: PrettyGraph g => Detail -> g -> IO () +-- writeGraph simple g = do +-- withTemporaryFile "acc.dot" $ \path hdl -> do +-- hPrint hdl (ppGraph simple g) +-- hPutStrLn stderr ("program graph: " ++ path) -withTemporaryFile :: String -> (FilePath -> Handle -> IO a) -> IO a -withTemporaryFile template go = do - pid <- getProcessID - tmp <- getTemporaryDirectory - let dir = tmp "accelerate-" ++ show pid - createDirectoryIfMissing True dir - bracket (openTempFile dir template) (hClose . snd) (uncurry go) +-- withTemporaryFile :: String -> (FilePath -> Handle -> IO a) -> IO a +-- withTemporaryFile template go = do +-- pid <- getProcessID +-- tmp <- getTemporaryDirectory +-- let dir = tmp "accelerate-" ++ show pid +-- createDirectoryIfMissing True dir +-- bracket (openTempFile dir template) (hClose . snd) (uncurry go) -#if defined(mingw32_HOST_OS) -getProcessID :: IO ProcessId -getProcessID = return 0xaaaa -#endif -#endif +-- #if defined(mingw32_HOST_OS) +-- getProcessID :: IO ProcessId +-- getProcessID = return 0xaaaa +-- #endif +-- #endif diff --git a/src/Data/Array/Accelerate/Debug/Internal/Profile.hs b/src/Data/Array/Accelerate/Debug/Internal/Profile.hs index 971666478..eda1781f0 100644 --- a/src/Data/Array/Accelerate/Debug/Internal/Profile.hs +++ b/src/Data/Array/Accelerate/Debug/Internal/Profile.hs @@ -1,6 +1,5 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE ForeignFunctionInterface #-} -{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE TemplateHaskell #-} {-# OPTIONS_GHC -fobject-code #-} {-# OPTIONS_HADDOCK hide #-} -- | @@ -35,7 +34,7 @@ import qualified Data.Array.Accelerate.Debug.Internal.Tracy as Tracy #endif import Data.Atomic ( Atomic ) -import qualified Data.Atomic as Atomic +import qualified Data.Atomic as Atomic import Data.Char import Foreign.C.String diff --git a/src/Data/Array/Accelerate/Debug/Internal/Stats.hs b/src/Data/Array/Accelerate/Debug/Internal/Stats.hs index 8134ba956..1dd47fa18 100644 --- a/src/Data/Array/Accelerate/Debug/Internal/Stats.hs +++ b/src/Data/Array/Accelerate/Debug/Internal/Stats.hs @@ -1,8 +1,6 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE OverloadedStrings #-} -{-# OPTIONS_GHC -fno-warn-unused-binds #-} -{-# OPTIONS_GHC -fno-warn-unused-imports #-} +{-# LANGUAGE CPP #-} +{-# OPTIONS_GHC -Wno-unused-binds #-} +{-# OPTIONS_GHC -Wno-unused-imports #-} -- | -- Module : Data.Array.Accelerate.Debug.Internal.Stats -- Copyright : [2008..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Debug/Internal/Timed.hs b/src/Data/Array/Accelerate/Debug/Internal/Timed.hs index 8e36194df..1ebe35526 100644 --- a/src/Data/Array/Accelerate/Debug/Internal/Timed.hs +++ b/src/Data/Array/Accelerate/Debug/Internal/Timed.hs @@ -1,6 +1,9 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE MagicHash #-} -{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE CPP #-} + +#ifdef ACCELERATE_DEBUG +{-# LANGUAGE MagicHash #-} +#endif + -- | -- Module : Data.Array.Accelerate.Debug.Internal.Timed -- Copyright : [2016..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Debug/Internal/Trace.hs b/src/Data/Array/Accelerate/Debug/Internal/Trace.hs index a1705c860..2bd0fdc8d 100644 --- a/src/Data/Array/Accelerate/Debug/Internal/Trace.hs +++ b/src/Data/Array/Accelerate/Debug/Internal/Trace.hs @@ -1,7 +1,4 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE CPP #-} -{-# LANGUAGE ForeignFunctionInterface #-} -{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE CPP #-} -- | -- Module : Data.Array.Accelerate.Debug.Internal.Trace -- Copyright : [2008..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Debug/Internal/Tracy.hs b/src/Data/Array/Accelerate/Debug/Internal/Tracy.hs index 19504a6a0..119511411 100644 --- a/src/Data/Array/Accelerate/Debug/Internal/Tracy.hs +++ b/src/Data/Array/Accelerate/Debug/Internal/Tracy.hs @@ -1,6 +1,9 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE ForeignFunctionInterface #-} -{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE CPP #-} + +#if defined(ACCELERATE_TRACY) && !defined(__GHCIDE__) +{-# LANGUAGE TemplateHaskell #-} +#endif + {-# OPTIONS_GHC -fobject-code #-} -- | -- Module : Data.Array.Accelerate.Debug.Internal.Tracy diff --git a/src/Data/Array/Accelerate/Debug/Trace.hs b/src/Data/Array/Accelerate/Debug/Trace.hs index 8b2f6f1d2..890936fe9 100644 --- a/src/Data/Array/Accelerate/Debug/Trace.hs +++ b/src/Data/Array/Accelerate/Debug/Trace.hs @@ -1,7 +1,4 @@ -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TemplateHaskellQuotes #-} -- | -- Module : Data.Array.Accelerate.Debug.Trace -- Copyright : [2008..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Error.hs b/src/Data/Array/Accelerate/Error.hs index 5ce86630b..2fb6038f0 100644 --- a/src/Data/Array/Accelerate/Error.hs +++ b/src/Data/Array/Accelerate/Error.hs @@ -1,6 +1,5 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE RecordWildCards #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.Error diff --git a/src/Data/Array/Accelerate/Interpreter.hs b/src/Data/Array/Accelerate/Interpreter.hs index 9d3bfdfde..2d6b760d0 100644 --- a/src/Data/Array/Accelerate/Interpreter.hs +++ b/src/Data/Array/Accelerate/Interpreter.hs @@ -1,17 +1,8 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE MagicHash #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE PatternGuards #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE RecordWildCards #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE ViewPatterns #-} -{-# OPTIONS_GHC -fno-warn-name-shadowing #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeFamilies #-} + +{-# OPTIONS_GHC -Wno-name-shadowing #-} {-# OPTIONS_HADDOCK prune #-} -- | -- Module : Data.Array.Accelerate.Interpreter @@ -29,60 +20,640 @@ -- module Data.Array.Accelerate.Interpreter ( - - Smart.Acc, Sugar.Arrays, - Afunction, AfunctionR, - - -- * Interpret an array expression - run, run1, runN, - - -- Internal (hidden) - evalPrim, evalPrimConst, evalCoerceScalar, atraceOp, - + module Data.Array.Accelerate.Interpreter, + UniformScheduleFun ) where -import Data.Array.Accelerate.AST hiding ( Boundary(..) ) -import Data.Array.Accelerate.AST.Environment -import Data.Array.Accelerate.AST.Var -import Data.Array.Accelerate.Array.Data -import Data.Array.Accelerate.Error +import Prelude hiding (take, (!!), sum, Either(..) ) +import Data.Array.Accelerate.AST.Partitioned +import Data.Array.Accelerate.AST.Kernel +import Data.Array.Accelerate.Trafo.Lowering +import qualified Data.Array.Accelerate.Debug.Internal as Debug import Data.Array.Accelerate.Representation.Array -import Data.Array.Accelerate.Representation.Elt -import Data.Array.Accelerate.Representation.Shape -import Data.Array.Accelerate.Representation.Slice -import Data.Array.Accelerate.Representation.Stencil -import Data.Array.Accelerate.Representation.Tag +import Data.Array.Accelerate.Error +import Data.Array.Accelerate.Analysis.Hash.Exp ( intHost, hashQ ) +import Data.Array.Accelerate.Analysis.Hash.Operation +import Data.Array.Accelerate.Representation.Ground import Data.Array.Accelerate.Representation.Type -import Data.Array.Accelerate.Representation.Vec -import Data.Array.Accelerate.Trafo -import Data.Array.Accelerate.Trafo.Delayed ( DelayedOpenAfun, DelayedOpenAcc ) -import Data.Array.Accelerate.Trafo.Sharing ( AfunctionR, AfunctionRepr(..), afunctionRepr ) +import Data.Array.Accelerate.Representation.Shape +import Data.Array.Accelerate.AST.Environment hiding (prjVars) import Data.Array.Accelerate.Type import Data.Primitive.Vec -import qualified Data.Array.Accelerate.AST as AST -import qualified Data.Array.Accelerate.Debug.Internal.Flags as Debug -import qualified Data.Array.Accelerate.Debug.Internal.Graph as Debug -import qualified Data.Array.Accelerate.Debug.Internal.Stats as Debug -import qualified Data.Array.Accelerate.Debug.Internal.Timed as Debug -import qualified Data.Array.Accelerate.Smart as Smart -import qualified Data.Array.Accelerate.Sugar.Array as Sugar -import qualified Data.Array.Accelerate.Sugar.Elt as Sugar -import qualified Data.Array.Accelerate.Trafo.Delayed as AST - -import Control.DeepSeq ( ($!!) ) -import Control.Exception -import Control.Monad -import Control.Monad.ST -import Data.Bits -import Data.Primitive.ByteArray import Data.Primitive.Types +import Data.Primitive.ByteArray +import Formatting hiding (int) import Data.Text.Lazy.Builder -import Formatting -import System.IO -import System.IO.Unsafe ( unsafePerformIO ) -import Unsafe.Coerce -import qualified Data.Text.IO as T -import Prelude hiding ( (!!), sum ) +import Data.Array.Accelerate.Representation.Elt +import Data.Array.Accelerate.Representation.Vec +import Data.Array.Accelerate.Trafo.Exp.Substitution +import Unsafe.Coerce (unsafeCoerce) +import Control.Monad.ST +import Data.Bits +import Data.Array.Accelerate.Backend +import Data.Array.Accelerate.Array.Buffer +import Data.Array.Accelerate.Pretty.Partitioned () +import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.LeftHandSide (LeftHandSide (LeftHandSideUnit)) +import Data.Array.Accelerate.AST.Schedule + +import Control.Concurrent (forkIO) +import Data.IORef +import Control.Concurrent.MVar + +import Data.Array.Accelerate.AST.Schedule.Uniform (UniformScheduleFun) +import Data.Array.Accelerate.Trafo.Schedule.Uniform () +import Data.Array.Accelerate.Pretty.Schedule.Uniform () +import qualified Data.Array.Accelerate.AST.Schedule.Uniform as S +import qualified Data.Map as M +import Control.Monad (when) +import Data.Array.Accelerate.Trafo.Var (DeclareVars(DeclareVars), declareVars) +import Data.Array.Accelerate.Trafo.Operation.Substitution (alet, aletUnique, weaken, LHS (LHS), mkLHS) +import Data.Map (Map) +import qualified Data.Text as T +import Data.Functor.Identity +import System.IO.Unsafe (unsafePerformIO) + +data Interpreter +-- instance Backend Interpreter where +-- type Schedule Interpreter = UniformScheduleFun +-- type Kernel Interpreter = InterpretKernel + + +(!?!) :: (Ord a1, Show a1, Show a2) => Map a1 a2 -> a1 -> a2 +map !?! key = case map M.!? key of + Just x -> x + Nothing -> error ("error: map "<> show map <> "does not contain key " <> show key) + + + +-- we can implement stencils using clamp, mirror or wrap with backpermute and zipwith(map), but for stencils using function we need a little extra. +data InterpretOp args where + IMap :: InterpretOp (Fun' (s -> t) -> In sh s -> Out sh t -> ()) + IBackpermute :: InterpretOp (Fun' (sh' -> sh) -> In sh t -> Out sh' t -> ()) + -- -- backpermute which doesn't error when the target index is out of bounds, but then evaluates the other function like a `generate` at the target index. + -- -- Generalisation of a 'backpermute with default', but not as general as a generate + -- IBackpermuteOr :: InterpretOp (Fun' (sh' -> sh) -> Fun' (sh -> t) -> In sh t -> Out sh' t -> ()) + IGenerate :: InterpretOp (Fun' (sh -> t) -> Out sh t -> ()) + IPermute :: InterpretOp (Fun' (e -> e -> e) + -> Mut sh' e + -> In sh (PrimMaybe (sh', e)) + -> ()) + IPermuteUnique::InterpretOp (Mut sh' e + -> In sh (PrimMaybe (sh', e)) + -> ()) + -- The interpreter doesn't have 'communication' with other 'threads' like the other backends do, to simulate cooperative scan/fold we use an IORef instead + IFold1 :: IORef (Map Int e) -> InterpretOp (Fun' (e -> e -> e) -> In (sh, Int) e -> Out sh e -> ()) + IScan1 :: Direction -> IORef (Map Int e) -> InterpretOp (Fun' (e -> e -> e) -> In (sh, Int) e -> Out (sh, Int) e -> ()) + -- append a number of elements to the left or right of each innermost row using the generate function + -- IAppend :: Side -> Int -> InterpretOp (Fun' ((sh, Int) -> e) -> In (sh, Int) e -> Out (sh, Int) e -> ()) + +instance LowerAcc InterpretOp where + mkMap a b c = Exec IMap (a :>: b :>: c :>: ArgsNil) + mkBackpermute a b c = Exec IBackpermute (a :>: b :>: c :>: ArgsNil) + mkGenerate a b = Exec IGenerate (a :>: b :>: ArgsNil) + mkPermute (Just a) b c = Exec IPermute (a :>: b :>: c :>: ArgsNil) + mkPermute Nothing b c = Exec IPermuteUnique (b :>: c :>: ArgsNil) + mkFold a Nothing b c = Exec (IFold1 $ unsafePerformIO $ newIORef mempty) (a :>: b :>: c :>: ArgsNil) + -- we lower a Fold with seed element into a Fold1 followed by a map which prepends the seed + mkFold a@(ArgFun f) (Just (ArgExp seed)) b@(ArgArray In (ArrayR _ tp) _ _) c@(ArgArray _ arr' sh' _) + | DeclareVars lhsTemp wTemp kTemp <- declareVars $ buffersR tp = + aletUnique lhsTemp (lowerAlloc arr' $ fromGrounds sh') $ + alet LeftHandSideUnit + (mkFold (weaken wTemp a) Nothing (weaken wTemp b) (ArgArray Out arr' (weakenVars wTemp sh') (kTemp weakenId))) $ + case mkLHS tp of + LHS lhs vars -> + mkMap + (ArgFun $ + Lam lhs $ Body $ (\f -> apply2 tp f (weakenThroughReindex wTemp reindexExp $ + weakenE weakenEmpty seed) (expVars vars)) $ weakenThroughReindex wTemp reindexExp f) + (ArgArray In arr' (weakenVars wTemp sh') (kTemp weakenId)) + (weaken wTemp c) + -- mkScan dir a Nothing b c = Exec (IScan1 dir $ unsafePerformIO $ newIORef mempty) (a :>: b :>: c :>: ArgsNil) + -- mkScan _ _ _ _ _ = error "exclusive scan not implemented" + -- we lower a Scan with seed into a scan1 followed by a map followed by an append + -- mkScan dir comb (Just (ArgExp seed)) i@(ArgArray In arr@(ArrayR shr tp) sh _) o + -- | DeclareVars lhsTemp1 wTemp kTemp1 <- declareVars $ buffersR tp + -- , DeclareVars lhsTemp2 wTemp2 kTemp2 <- declareVars $ buffersR tp + -- , wTemp1 <- wTemp2 .> wTemp = + -- aletUnique lhsTemp1 (lowerAlloc arr $ fromGrounds sh) $ + -- aletUnique lhsTemp2 (lowerAlloc arr $ fromGrounds $ weakenVars wTemp sh) $ + -- alet LeftHandSideUnit + -- (alet LeftHandSideUnit + -- (mkScan dir (weaken wTemp1 comb) Nothing (weaken wTemp1 i) (ArgArray Out arr (weakenVars wTemp1 sh) (kTemp1 wTemp2))) + -- (mkMap (ArgFun $ case mkLHS tp of + -- LHS lhs vars -> + -- Lam lhs $ Body $ (\f -> apply2 tp f (weakenThroughReindex wTemp1 reindexExp $ + -- weakenE weakenEmpty seed) (expVars vars)) $ weakenThroughReindex wTemp1 reindexExp $ (\(ArgFun f) -> f) comb) + -- (ArgArray In arr (weakenVars wTemp1 sh) (kTemp1 wTemp2)) + -- (ArgArray Out arr (weakenVars wTemp1 sh) (kTemp2 weakenId)))) $ + -- mkAppend + -- (case dir of + -- LeftToRight -> Left + -- RightToLeft -> Right) + -- 1 + -- (ArgFun $ Lam (LeftHandSideWildcard $ shapeType shr) $ Body $ weakenThroughReindex wTemp1 reindexExp seed) + -- (ArgArray In arr (weakenVars wTemp1 sh) (kTemp2 weakenId)) + -- (weaken wTemp1 o) + + -- mkBackpermuteOr + -- (case dir of + -- -- The easy direction: identity backpermute, the default case will trigger on the last element + -- RightToLeft -> case mkLHS (shapeType shr) of + -- LHS lhs vars -> + -- ArgFun $ Lam lhs $ + -- Body $ expVars vars + -- -- The difficult direction: backpermute does a +1, the default case will trigger on the first element + -- LeftToRight -> case mkLHS (shapeType shr') of + -- LHS lhs vars -> + -- ArgFun $ Lam (lhs `LeftHandSidePair` LeftHandSideSingle scalarTypeInt) $ + -- Body $ Pair (expVars $ weakenVars (weakenSucc' weakenId) vars) (PrimApp (PrimMin $ NumSingleType $ IntegralNumType TypeInt) $ Pair (Evar (Var scalarTypeInt ZeroIdx)) (Const scalarTypeInt 1)) + -- ) + +instance EncodeOperation InterpretOp where + encodeOperation IMap = intHost $(hashQ ("Map" :: String)) + encodeOperation IBackpermute = intHost $(hashQ ("Backpermute" :: String)) + encodeOperation IGenerate = intHost $(hashQ ("Generate" :: String)) + encodeOperation IPermute = intHost $(hashQ ("Permute" :: String)) + encodeOperation (IFold1 _) = intHost $(hashQ ("Fold1" :: String)) + encodeOperation (IScan1 LeftToRight _) = intHost $(hashQ ("Scanl1" :: String)) + encodeOperation (IScan1 RightToLeft _) = intHost $(hashQ ("Scanr1" :: String)) + -- encodeOperation (IAppend Left n) = intHost $(hashQ ("Appendl" :: String)) <> intHost n + -- encodeOperation (IAppend Right n) = intHost $(hashQ ("Appendr" :: String)) <> intHost n + encodeOperation IPermuteUnique{} = error "TODO WALL: NON-EXHAUSTIVE PATTERN MATCH" + +-- mkAppend :: Side -> Int -> Arg env (Fun' ((sh, Int) -> e)) -> Arg env (In (sh, Int) e) -> Arg env (Out (sh, Int) e) -> OperationAcc InterpretOp env () +-- mkAppend side i a b c = Exec (IAppend side i) (a :>: b :>: c :>: ArgsNil) + +-- mkBackpermuteOr :: Arg env (Fun' (sh' -> sh)) +-- -> Arg env (Fun' (sh -> t )) +-- -> Arg env (In sh t) +-- -> Arg env (Out sh' t) +-- -> OperationAcc InterpretOp env () +-- mkBackpermuteOr a b c d = Exec IBackpermuteOr (a :>: b :>: c :>: d :>: ArgsNil) + +instance SimplifyOperation InterpretOp where + detectCopy IMap = detectMapCopies + detectCopy IBackpermute = detectBackpermuteCopies + detectCopy _ = const [] + +instance SLVOperation InterpretOp where + slvOperation IGenerate = defaultSlvGenerate IGenerate + slvOperation IMap = defaultSlvMap IMap + slvOperation IBackpermute = defaultSlvBackpermute IBackpermute + slvOperation _ = Nothing + +instance SetOpIndices InterpretOp where + setOpIndices _ _ _ _ = error "TODO" + +data InterpretKernel env where + InterpretKernel :: Clustered InterpretOp args -> Args env args -> InterpretKernel env + +instance NFData' InterpretKernel where + rnf' (InterpretKernel cluster args) = undefined cluster `seq` rnfArgs args + +-- instance IsKernel InterpretKernel where +-- type KernelOperation InterpretKernel = InterpretOp +-- type KernelMetadata InterpretKernel = NoKernelMetadata + +-- compileKernel = const InterpretKernel + +-- encodeKernel (InterpretKernel cluster args) +-- = Prelude.Right $ encodeOperation cluster <> encodePreArgs encodeArg args + +instance PrettyKernel InterpretKernel where + -- PrettyKernelBody provides a Val but prettyOpWithArgs expects a Val', should we change them to have the + -- same type (either Val or Val'?) + prettyKernel = PrettyKernelBody False $ \env (InterpretKernel cluster args) -> prettyOpWithArgs env cluster args + +-- -- -2 is left>right, -1 is right>left, n is 'according to computation n' (e.g. Backpermute) +-- -- (Note that Labels are uniquely identified by an Int, the parent just gives extra information) +-- -- Use Output = -3 for 'cannot be fused with consumer', as that is more difficult to express (we don't know the consumer yet) +-- -- We restrict all the Inputs to >= -2. +-- data InterpreterVariables = DimensionsPerThread InOut Label +-- | IdleThreads InOut Side Label +type InterpreterVariables = () + -- deriving (Eq, Ord, Show) +data InOut = InArr | OutArr + deriving (Eq, Ord, Show) +data Side = Left | Right + deriving (Eq, Ord, Show) + + +-- instance MakesILP InterpretOp where +-- type BackendArg InterpretOp = Maybe Int +-- data BackendClusterArg InterpretOp arg where +-- -- ArrayInfo :: { dim :: Int, idleLeft :: Int, idleRight :: Int} -> BackendClusterArg InterpretOp arg -- can't do (m sh e) because vertically fused arrays get a Var' argument +-- -- NonArray :: BackendClusterArg InterpretOp arg +-- BCAI :: BackendClusterArg InterpretOp arg +-- -- each array argument gets labelled with its order, +-- -- this ensures that e.g. multiple inputs of the same array +-- -- in different orders won't fuse horizontally, and that +-- -- the correct one will be used by each consumer +-- labelLabelledArg solution l (L arg@(ArgArray In _ _ _) al) = LOp arg al . Just $ (solution !?! Graph.InDir l) --, (solution !?! BackendSpecific (DimensionsPerThread InArr l), solution !?! BackendSpecific (IdleThreads InArr Left l), solution !?! BackendSpecific (IdleThreads InArr Right l))) +-- labelLabelledArg solution l (L arg@(ArgArray Out _ _ _) al) = LOp arg al . Just $ (solution !?! Graph.OutDir l) --, (solution !?! BackendSpecific (DimensionsPerThread OutArr l), solution !?! BackendSpecific (IdleThreads OutArr Left l), solution !?! BackendSpecific (IdleThreads OutArr Right l))) +-- labelLabelledArg _ _ (L arg al) = LOp arg al Nothing + +-- getClusterArg (LOp ArgArray{} _ (Just _)) = BCAI +-- getClusterArg (LOp ArgVar{} _ (Just _)) = BCAI -- vertically fused array +-- getClusterArg (LOp ArgArray{} _ Nothing) = error "TODO: make a nice error" +-- getClusterArg (LOp _ _ Nothing) = BCAI +-- getClusterArg (LOp _ _ (Just _)) = error "TODO: make a nice error" + +-- finalize = foldMap $ \l -> timesN (manifest l) .>. c (OutDir l) + +-- mkGraph IBackpermute (_ :>: ((L _ (_, Set.toList -> lIns)) :>: _ :>: ArgsNil)) l@(Label i _) = +-- Graph.Info +-- mempty +-- ( inputConstraints l lIns +-- <> c (InDir l) .==. int i +-- -- <> c (BackendSpecific $ IdleThreads InArr Left l) .==. c (BackendSpecific $ IdleThreads OutArr Left l) +-- -- <> c (BackendSpecific $ IdleThreads InArr Right l) .==. c (BackendSpecific $ IdleThreads OutArr Right l) +-- -- TODO reconsider: this has difficult consequences, because if we're backpermuting a matrix from a vector and Dims was 1, does it make sense still? +-- -- Need to just work out some examples by hand + +-- -- seems like we might need not just #dimsperthread (shapeR), but also the size of those dims (shape) +-- -- like the backpermute functions; these are not known to the ILP but can be statically threaded through back-to-front before evaluating a cluster, using the env +-- -- or more precisely, we can thread a computation through the cluster at compile time, which can (at run time, before eval of cluster) use the env +-- -- to compute them +-- -- it's beginning to feel a lot like this is too much for the interpreter? But a nice little generalisation which allows the backends to use it too would be great. +-- -- maybe find a way to make interpret without dimsperthread first though + +-- -- <> c (BackendSpecific $ DimensionsPerThread InArr l) .==. c (BackendSpecific $ DimensionsPerThread OutArr l) +-- ) -- enforce that the backpermute follows its own rules, but the output can be anything +-- (defaultBounds l) +-- mkGraph IGenerate _ l = Graph.Info mempty mempty (defaultBounds l) -- creats some superfluous variables, oh well +-- mkGraph IMap (_ :>: L _ (_, Set.toList -> lIns) :>: _ :>: ArgsNil) l = +-- Graph.Info +-- mempty +-- ( inputConstraints l lIns +-- <> c (InDir l) .==. c (OutDir l) +-- -- <> c (BackendSpecific $ IdleThreads InArr Left l) .==. c (BackendSpecific $ IdleThreads OutArr Left l) +-- -- <> c (BackendSpecific $ IdleThreads InArr Right l) .==. c (BackendSpecific $ IdleThreads OutArr Right l) +-- -- <> c (BackendSpecific $ DimensionsPerThread InArr l) .==. c (BackendSpecific $ DimensionsPerThread OutArr l) +-- ) +-- (defaultBounds l) +-- mkGraph IPermute (_ :>: L _ (_, lTargets) :>: _ :>: L _ (_, Set.toList -> lIns) :>: ArgsNil) l@(Label i _) = +-- Graph.Info +-- ( mempty & infusibleEdges .~ Set.map (-?> l) lTargets) -- Cannot fuse with the producer of the target array +-- ( inputConstraints l lIns <> c (OutDir l) .==. int (-3-i)) -- convention meaning infusible +-- ( lower (-2) (InDir l) +-- -- <> lower 0 (BackendSpecific $ IdleThreads InArr Left l) +-- -- <> lower 0 (BackendSpecific $ IdleThreads InArr Right l) +-- -- <> equal 0 (BackendSpecific $ IdleThreads OutArr Left l) +-- -- <> equal 0 (BackendSpecific $ IdleThreads OutArr Right l) +-- -- <> equal 0 (BackendSpecific $ DimensionsPerThread OutArr l) +-- ) + +-- mkGraph (IFold1 _) (_ :>: L _ (_, Set.toList -> lIns) :>: _ :>: ArgsNil) l = +-- Graph.Info +-- mempty +-- ( inputConstraints l lIns +-- <> c (InDir l) .==. c (OutDir l) +-- -- <> c (BackendSpecific $ IdleThreads InArr Left l) .==. c (BackendSpecific $ IdleThreads OutArr Left l) +-- -- <> c (BackendSpecific $ IdleThreads InArr Right l) .==. c (BackendSpecific $ IdleThreads OutArr Right l) +-- -- <> c (BackendSpecific $ DimensionsPerThread InArr l) .+. int 1 .==. c (BackendSpecific $ DimensionsPerThread OutArr l) +-- ) +-- (defaultBounds l) + +-- mkGraph (IScan1 _ _) (_ :>: L _ (_, Set.toList -> lIns) :>: _ :>: ArgsNil) l = +-- Graph.Info +-- mempty +-- ( inputConstraints l lIns +-- <> c (InDir l) .==. c (OutDir l) +-- <> c (InDir l) .<=. int 0 +-- -- <> c (BackendSpecific $ IdleThreads InArr Left l) .==. c (BackendSpecific $ IdleThreads OutArr Left l) +-- -- <> c (BackendSpecific $ IdleThreads InArr Right l) .==. c (BackendSpecific $ IdleThreads OutArr Right l) +-- -- <> c (BackendSpecific $ DimensionsPerThread InArr l) .+. int 1 .==. c (BackendSpecific $ DimensionsPerThread OutArr l) +-- ) +-- (defaultBounds l) + +-- -- mkGraph (IAppend side n) (_ :>: L _ (_, Set.toList -> lIns) :>: _ :>: ArgsNil) l = +-- -- Graph.Info +-- -- mempty +-- -- ( inputConstraints l lIns +-- -- <> c (InDir l) .==. c (OutDir l) +-- -- -- <> c (BackendSpecific $ IdleThreads InArr Left l) .==. c (BackendSpecific $ IdleThreads OutArr Left l) .+. int (if side == Left then n else 0) +-- -- -- <> c (BackendSpecific $ IdleThreads InArr Right l) .==. c (BackendSpecific $ IdleThreads OutArr Right l) .+. int (if side == Right then n else 0) +-- -- -- <> c (BackendSpecific $ DimensionsPerThread InArr l) .==. c (BackendSpecific $ DimensionsPerThread OutArr l) +-- -- ) +-- -- (defaultBounds l) + +-- -- encodeBackendClusterArg (ArrayInfo d l r) = intHost $(hashQ ("ArrayInfo" :: String)) <> intHost d <> intHost l <> intHost r +-- -- encodeBackendClusterArg NonArray = intHost $(hashQ ("NonArray" :: String)) +-- encodeBackendClusterArg BCAI = intHost $(hashQ ("BCAI" :: String)) + +-- -- mkGraph IBackpermuteOr (_ :>: _ :>: ((L _ (_, Set.toList -> lIns)) :>: _)) l@(Label i _) = +-- -- Info +-- -- mempty +-- -- ( inputConstraints l lIns +-- -- <> c (InDir l) .==. int i +-- -- -- <> c (BackendSpecific $ IdleThreads InArr Left l) .==. c (BackendSpecific $ IdleThreads OutArr Left l) +-- -- -- <> c (BackendSpecific $ IdleThreads InArr Right l) .==. c (BackendSpecific $ IdleThreads OutArr Right l) +-- -- <> c (BackendSpecific $ DimensionsPerThread InArr l) .==. c (BackendSpecific $ DimensionsPerThread OutArr l) +-- -- ) +-- -- (defaultBounds l) + +-- -- | If l and lIn are fused, the out-order of lIn and the in-order of l should match +-- inputConstraints :: Label Comp -> [Label Comp] -> Constraint InterpretOp +-- inputConstraints l = foldMap $ \lIn -> +-- timesN (fused lIn l) .>=. c (InDir l) .-. c (OutDir lIn) +-- <> (-1) .*. timesN (fused lIn l) .<=. c (InDir l) .-. c (OutDir lIn) +-- -- <> timesN (fused lIn l) .>=. c (BackendSpecific $ DimensionsPerThread InArr l) .-. c (BackendSpecific $ DimensionsPerThread OutArr lIn) +-- -- <> (-1) .*. timesN (fused lIn l) .<=. c (BackendSpecific $ DimensionsPerThread InArr l) .-. c (BackendSpecific $ DimensionsPerThread OutArr lIn) +-- -- <> timesN (fused lIn l) .>=. c (BackendSpecific $ IdleThreads InArr Left l) .-. c (BackendSpecific $ IdleThreads OutArr Left lIn) +-- -- <> (-1) .*. timesN (fused lIn l) .<=. c (BackendSpecific $ IdleThreads InArr Left l) .-. c (BackendSpecific $ IdleThreads OutArr Left lIn) +-- -- <> timesN (fused lIn l) .>=. c (BackendSpecific $ IdleThreads InArr Right l) .-. c (BackendSpecific $ IdleThreads OutArr Right lIn) +-- -- <> (-1) .*. timesN (fused lIn l) .<=. c (BackendSpecific $ IdleThreads InArr Right l) .-. c (BackendSpecific $ IdleThreads OutArr Right lIn) + +-- defaultBounds :: Label Comp -> Bounds InterpretOp +-- defaultBounds l = lower (-2) (InDir l) <> lower (-2) (OutDir l) +-- -- <> lower 0 (BackendSpecific $ IdleThreads InArr Left l) +-- -- <> lower 0 (BackendSpecific $ IdleThreads InArr Right l) +-- -- <> lower 0 (BackendSpecific $ IdleThreads OutArr Left l) +-- -- <> lower 0 (BackendSpecific $ IdleThreads OutArr Right l) + + +-- instance NFData' (BackendClusterArg InterpretOp) where +-- -- rnf' (ArrayInfo x y z) = rnf x `seq` rnf y `seq` rnf z +-- -- rnf' NonArray = () +-- rnf' :: BackendClusterArg InterpretOp a -> () +-- rnf' BCAI = () + +-- instance ShrinkArg (BackendClusterArg InterpretOp) where +-- -- shrinkArg _ (ArrayInfo x y z) = ArrayInfo x y z +-- -- shrinkArg _ _ = error "impossible" +-- -- deadArg (ArrayInfo x y z) = ArrayInfo x y z +-- -- deadArg _ = error "impossible" +-- shrinkArg _ BCAI = BCAI +-- deadArg BCAI = BCAI + +instance NFData' InterpretOp where + rnf' (IScan1 dir lookup) = lookup `seq` dir `seq` () + rnf' (IFold1 lookup) = lookup `seq` () + -- rnf' (IAppend side n) = side `seq` n `seq` () + rnf' IMap = () + rnf' IBackpermute = () + rnf' IGenerate = () + rnf' IPermute = () + rnf' IPermuteUnique{} = error "TODO WALL: NON-EXHAUSTIVE PATTERN MATCH" + +instance PrettyOp InterpretOp where + prettyOp IMap = "map" + prettyOp IBackpermute = "backpermute" + prettyOp IGenerate = "generate" + prettyOp IPermute = "permute" + prettyOp (IFold1 _) = "fold1" + -- prettyOp IBackpermuteOr = "backpermute_with_default_generate" + prettyOp (IScan1 LeftToRight _) = "scanl1" + prettyOp (IScan1 RightToLeft _) = "scanr1" + -- prettyOp (IAppend _ _) = "append" + prettyOp IPermuteUnique{} = error "TODO WALL: NON-EXHAUSTIVE PATTERN MATCH" + +instance Execute UniformScheduleFun InterpretKernel where + data Linked UniformScheduleFun InterpretKernel t = InterpretLinked (UniformScheduleFun InterpretKernel () t) + linkAfunSchedule = InterpretLinked + executeAfunSchedule _ (InterpretLinked sched) = executeScheduleFun Empty sched + +executeScheduleFun :: Val env -> UniformScheduleFun InterpretKernel env t -> IOFun t +executeScheduleFun env (S.Sbody schedule) = executeSchedule env schedule +executeScheduleFun env (S.Slam lhs fun) = \input -> executeScheduleFun (env `push` (lhs, input)) fun + +executeSchedule :: Val env -> S.UniformSchedule InterpretKernel env -> IO () +executeSchedule !env = \case + S.Return -> return () + S.Alet lhs binding body -> do + value <- executeBinding env binding + let env' = env `push` (lhs, value) + executeSchedule env' body + S.Effect effect next -> do + executeEffect env effect + executeSchedule env next + S.Acond var true false next -> do + let value = prj (varIdx var) env + let branch = if value == 1 then true else false + executeSchedule env branch + executeSchedule env next + S.Awhile io step input next -> do + executeAwhile env io step (prjVars input env) + executeSchedule env next + -- TODO: S.AwhileSeq + S.AwhileSeq{} -> error "TODO WALL: NON-EXHAUSTIVE PATTERN MATCH" + S.Spawn a b -> do + _ <- forkIO (executeSchedule env a) + executeSchedule env b + +executeBinding :: Val env -> S.Binding env t -> IO t +executeBinding env = \case + S.Compute expr -> + return $ runIdentity $ evalExp expr (evalArrayInstrDefault env) + S.NewSignal _ -> do + mvar <- newEmptyMVar + return (S.Signal mvar, S.SignalResolver mvar) + S.NewRef _ -> do + ioref <- newIORef $ internalError "Illegal schedule: Read from ref without value. Some synchronization might be missing." + return (S.Ref ioref, S.OutputRef ioref) + S.Alloc shr tp sh -> do + let n = size shr $ prjVars sh env + MutableBuffer buffer <- newBuffer tp n + return $ Buffer buffer + S.Use _ _ buffer -> + return buffer + S.Unit (Var tp ix) -> do + mbuffer@(MutableBuffer buffer) <- newBuffer tp 1 + writeBuffer tp mbuffer 0 $ prj ix env + return $ Buffer buffer + S.RefRead ref -> do + let S.Ref ioref = prj (varIdx ref) env + readIORef ioref + +executeEffect :: forall env. Val env -> S.Effect InterpretKernel env -> IO () +executeEffect env = \case + S.Exec _ kernelFun args -> + executeKernelFun env Empty kernelFun args + S.SignalAwait signals -> mapM_ await signals + S.SignalResolve signals -> mapM_ resolve signals + S.RefWrite ref valueVar -> do + let S.OutputRef ioref = prj (varIdx ref) env + let value = prj (varIdx valueVar) env + writeIORef ioref value + S.Aassert msg cond + | runIdentity (evalExp cond $ evalArrayInstrDefault env) == 1 -> return () + | otherwise -> errorWithoutStackTrace ("\n*** Assertion failed: " ++ T.unpack msg) + S.Atrace{} -> error "TODO WALL: NON-EXHAUSTIVE PATTERN MATCH" + where + await :: Idx env S.Signal -> IO () + await idx = do + let S.Signal mvar = prj idx env + readMVar mvar + + resolve :: Idx env S.SignalResolver -> IO () + resolve idx = do + let S.SignalResolver mvar = prj idx env + putMVar mvar () + +executeKernelFun :: Val env -> Val kenv -> OpenKernelFun InterpretKernel kenv t -> S.SArgs env t -> IO () +executeKernelFun env kernelEnv (KernelFunLam KernelArgRscalar{} fun) (S.SArgScalar var :>: args) + = executeKernelFun env (Push' kernelEnv (prj (varIdx var) env)) fun args +executeKernelFun env kernelEnv (KernelFunLam KernelArgRbuffer{} fun) (S.SArgBuffer _ var :>: args) + = executeKernelFun env (Push' kernelEnv (prj (varIdx var) env)) fun args +executeKernelFun _ kernelEnv (KernelFunBody (InterpretKernel cluster args)) ArgsNil = evalClusterInterpreter cluster args kernelEnv + +executeAwhile + :: Val env + -> S.InputOutputR input output + -> UniformScheduleFun InterpretKernel env (input -> S.Output PrimBool -> output -> ()) + -> input + -> IO () +executeAwhile env io step input = do + -- Set up the output variables for this iteration (and the input for the next) + (S.Signal mvarCondition, signalResolverCondition) <- executeBinding env (S.NewSignal "interpreter signal") + (S.Ref iorefCondition, outputRefCondition) <- executeBinding env $ S.NewRef $ GroundRscalar scalarType + (output, nextInput) <- bindAwhileIO io + + -- Execute a step + executeScheduleFun env step input (signalResolverCondition, outputRefCondition) output + + -- Check condition + readMVar mvarCondition + condition <- readIORef iorefCondition + when (condition == 1) $ + executeAwhile env io step nextInput + +bindAwhileIO :: S.InputOutputR input output -> IO (output, input) +bindAwhileIO S.InputOutputRsignal = do + mvar <- newEmptyMVar + return (S.SignalResolver mvar, S.Signal mvar) +bindAwhileIO (S.InputOutputRref _) = do + ioref <- newIORef $ internalError "Illegal schedule: Read from ref without value. Some synchronization might be missing." + return (S.OutputRef ioref, S.Ref ioref) +bindAwhileIO (S.InputOutputRpair io1 io2) = do + (output1, input1) <- bindAwhileIO io1 + (output2, input2) <- bindAwhileIO io2 + return ((output1, output2), (input1, input2)) +bindAwhileIO S.InputOutputRunit = + return ((), ()) + +prjVars :: Vars s env t -> Val env -> t +prjVars TupRunit _ = () +prjVars (TupRpair v1 v2) env = (prjVars v1 env, prjVars v2 env) +prjVars (TupRsingle var) env = prj (varIdx var) env + +writeOutputInterpreter :: TypeR e -> Vars s env (Buffers e) -> Val env -> Int -> e -> IO () +writeOutputInterpreter r buf env n x = writeBuffers r (veryUnsafeUnfreezeBuffers r $ varsGetVal buf env) n x + +evalClusterInterpreter :: Clustered InterpretOp args -> Args env args -> Val env -> IO () +evalClusterInterpreter (Clustered _c _b) _args _env = error "TODO: Implement evaluator for interpreter" +-- evalClusterInterpreter c@(Cluster _ (Cluster' io _)) args env = doNTimes (iterationsize io args env) $ evalCluster c args env + +-- iterationsize (Op _) ArgsNil env = Nothing +-- iterationsize (Op _) ((ArgVar _) :>: args) env = iterationsize (Op undefined) args env +-- iterationsize (Op _) ((ArgExp _) :>: args) env = iterationsize (Op undefined) args env +-- iterationsize (Op _) ((ArgFun _) :>: args) env = iterationsize (Op undefined) args env +-- iterationsize (Op _) ((ArgArray Mut ar tr tr') :>: args) env = iterationsize (Op undefined) args env +-- iterationsize (Op _) ((ArgArray In ar tr tr') :>: args) env = _wo +-- iterationsize (Op _) ((ArgArray Out ar tr tr') :>: args) env = _wp +-- iterationsize (P.Fused fu clus clus') args env = _wd + + +-- TODO update when we add folds, reconsider when we add scans that add 1 to the size... +-- iterationsize :: ClusterIO args i o -> Args env args -> Val env -> Int +-- iterationsize io args env = case io of +-- P.Empty -> error "no size" +-- P.Output {} -> case args of ArgArray Out (ArrayR shr _) sh _ :>: _ -> arrsize shr (varsGetVal sh env) +-- P.Vertical _ _ io' -> case args of -- skip past this one +-- ArgVar _ :>: args' -> iterationsize io' args' env +-- P.Input io' -> case args of ArgArray In _ _ _ :>: args' -> iterationsize io' args' env -- -> arrsize shr (varsGetVal sh env) +-- P.MutPut io' -> case args of ArgArray Mut _ _ _ :>: args' -> iterationsize io' args' env -- arrsize shr (varsGetVal sh env) +-- P.ExpPut' io' -> case args of _ :>: args' -> iterationsize io' args' env -- skip past this one +-- P.Trivial io' -> case args of _ :>: args' -> iterationsize io' args' env + + +arrsize :: ShapeR sh -> sh -> Int +arrsize ShapeRz () = 1 +arrsize (ShapeRsnoc shr) (sh,x) = x * arrsize shr sh + + +doNTimes :: Monad m => Int -> (Int -> m ()) -> m () +doNTimes n f + | n == 0 = pure () + | otherwise = f (n-1) >> doNTimes (n-1) f + +linearIndexToSh :: ShapeR sh -> sh -> Int -> sh +linearIndexToSh ShapeRz () 0 = () +linearIndexToSh ShapeRz () _ = error "non-zero index in unit array" +linearIndexToSh (ShapeRsnoc shr) (sh, _) i = let + innerSize = arrsize shr sh + outerIndex = i `div` innerSize + innerIndex = linearIndexToSh shr sh (i `mod` innerSize) + in (innerIndex, outerIndex) + +shToLinearIndex :: ShapeR sh -> sh -> sh -> Int +shToLinearIndex ShapeRz () () = 0 +shToLinearIndex (ShapeRsnoc shr) (sh, _) (sh', y) = let + innerSize = arrsize shr sh + innerIndex = shToLinearIndex shr sh sh' + in y*innerSize+innerIndex + + +-- TODO make these think about the new info: pass if idle, and loop if nested + -- make a little combinator for that +-- evalOp _ IMap env (PushBPFA _ f infoF (PushBPFA _ (Value x (Shape shr sh)) infoX _)) = pure $ PushFA (Value (evalFun f env x) (Shape shr sh)) Empty +-- evalOp _ IBackpermute _ (PushBPFA _ _ infoF (PushBPFA _ (Value x _) infoI (PushBPFA _ sh infoO _))) = pure $ PushFA (Value x sh) Empty -- We evaluated the backpermute at the start already, now simply relabel the shape info +-- evalOp i IGenerate env (PushBPFA f' f infoF (PushBPFA _ (Shape shr sh) infoSh _)) = pure $ PushFA (Value (evalFun f env $ fromIndex shr sh (f' i)) (Shape shr sh)) Empty -- apply f', the Int->Int backpermute function, to our index before evaluating the generate function. +-- evalOp i IPermute env (PushBPFA _ comb infoC (PushBPFA _ target infoT (PushBPFA _ f infoF (PushBPFA _ (Value (e :: e) (Shape shr sh)) infoI _)))) = +-- let typer = case comb of +-- Lam lhs _ -> lhsToTupR lhs +-- _ -> error "impossible; is a function" +-- in +-- case evalFun f env (fromIndex shr sh i) of +-- (0, _) -> pure $ PushFA target Empty +-- (1, ((), sh)) -> case target of +-- ArrayDescriptor shr shvars bufvars -> do +-- let shsize = varsGetVal shvars env +-- let j = toIndex shr shsize sh +-- let buf = varsGetVal bufvars env +-- let buf' = veryUnsafeUnfreezeBuffers @e typer buf +-- x <- readBuffers @e typer buf' j +-- let x' = evalFun comb env x e +-- writeBuffers typer buf' j x' +-- return $ PushFA target Empty +-- _ -> error "PrimMaybe's tag was non-zero and non-one" +-- evalOp i (IAppend side n) env (PushBPFA f' f infoF (PushBPFA _ (Value x shX) infoX (PushBPFA _ (Shape shr sh) infoO _))) = +-- if f' i - firstOfRow (f' i) (Shape shr sh) (idleLeft infoO) <= n +-- then pure $ PushFA (Value x (Shape shr sh)) Empty +-- else pure $ PushFA (Value (evalFun f env $ fromIndex shr sh (f' i)) (Shape shr sh)) Empty + +-- -- TODO the backpermute function!!!!! +-- evalOp i (IFold1 lookup) env (PushBPFA _ f infoF (PushBPFA _ (Value x sh) infoX (PushBPFA _ outSh _ _))) +-- | i == firstOfRow i sh (idleLeft infoX) = do -- write input to output +-- lookupmap <- readIORef lookup +-- writeIORef lookup (M.insert i x lookupmap) +-- pure $ PushFA (Value x outSh) Empty +-- | otherwise = do -- combine input with output using f +-- lookupmap <- readIORef lookup +-- let y = lookupmap !?! firstOfRow i sh (idleLeft infoX) +-- let z = evalFun f env y x +-- writeIORef lookup (M.insert (firstOfRow i sh (idleLeft infoX)) z lookupmap) +-- pure $ PushFA (Value z outSh) Empty +-- evalOp i (IScan1 dir lookup) env (PushBPFA _ f infoF (PushBPFA _ (Value x sh) infoX _outputstuff)) +-- | i == firstOfRow i sh (idleLeft infoX) = do -- write input to output +-- lookupmap <- readIORef lookup +-- writeIORef lookup $ M.insert i x lookupmap +-- pure $ PushFA (Value x sh) Empty +-- | otherwise = do -- combine input with previous output using f +-- lookupmap <- readIORef lookup +-- let y = lookupmap !?! (i - 1) +-- let z = evalFun f env y x +-- writeIORef lookup $ M.insert i z lookupmap +-- pure $ PushFA (Value z sh) Empty + +firstOfRow :: Int -> Sh sh e -> Int -> Int +firstOfRow i (Shape ShapeRz ()) idleRowsLeft + | i > idleRowsLeft = i - idleRowsLeft + | otherwise = error "not in a row" +firstOfRow i (Shape (ShapeRsnoc _) (_, n)) idleRowsLeft = (i `mod` (n + idleRowsLeft)) - idleRowsLeft + + -- Program execution @@ -90,69 +661,51 @@ import Prelude hiding ( (!! -- | Run a complete embedded array program using the reference interpreter. -- -run :: (HasCallStack, Sugar.Arrays a) => Smart.Acc a -> a -run a = unsafePerformIO execute - where - !acc = convertAcc a - execute = do - Debug.dumpGraph $!! acc - Debug.dumpSimplStats - res <- phase "execute" Debug.elapsed $ evaluate $ evalOpenAcc acc Empty - return $ Sugar.toArr $ snd res - --- | This is 'runN' specialised to an array program of one argument. --- -run1 :: (HasCallStack, Sugar.Arrays a, Sugar.Arrays b) => (Smart.Acc a -> Smart.Acc b) -> a -> b -run1 = runN --- | Prepare and execute an embedded array program. --- -runN :: forall f. (HasCallStack, Afunction f) => f -> AfunctionR f -runN f = go - where - !acc = convertAfun f - !afun = unsafePerformIO $ do - Debug.dumpGraph $!! acc - Debug.dumpSimplStats - return acc - !go = eval (afunctionRepr @f) afun Empty - -- - eval :: AfunctionRepr g (AfunctionR g) (ArraysFunctionR g) - -> DelayedOpenAfun aenv (ArraysFunctionR g) - -> Val aenv - -> AfunctionR g - eval (AfunctionReprLam reprF) (Alam lhs f) aenv = \a -> eval reprF f $ aenv `push` (lhs, Sugar.fromArr a) - eval AfunctionReprBody (Abody b) aenv = unsafePerformIO $ phase "execute" Debug.elapsed (Sugar.toArr . snd <$> evaluate (evalOpenAcc b aenv)) - eval _ _aenv _ = error "Two men say they're Jesus; one of them must be wrong" - --- -- | Stream a lazily read list of input arrays through the given program, --- -- collecting results as we go +-- run :: forall a. (HasCallStack, Sugar.Arrays a) => Smart.Acc a -> a +-- run _ = unsafePerformIO execute +-- where +-- acc :: PartitionedAcc InterpretOp () (LoweredArrays (Sugar.ArraysR a)) +-- !acc = undefined -- convertAcc a +-- execute = do +-- Debug.dumpGraph $!! acc +-- Debug.dumpSimplStats +-- res <- phase "execute" Debug.elapsed $ undefined +-- return $ Sugar.toArr $ snd res + +-- -- | This is 'runN' specialised to an array program of one argument. -- -- --- streamOut :: Arrays a => Sugar.Seq [a] -> [a] --- streamOut seq = let seq' = convertSeqWith config seq --- in evalDelayedSeq defaultSeqConfig seq' +-- run1 :: (HasCallStack, Sugar.Arrays a, Sugar.Arrays b) => (Smart.Acc a -> Smart.Acc b) -> a -> b +-- run1 = runN + +-- -- | Prepare and execute an embedded array program. +-- -- +-- runN :: forall f. (HasCallStack, Afunction f) => f -> AfunctionR f +-- runN f = go +-- where +-- !acc = convertAfun f +-- !afun = unsafePerformIO $ do +-- Debug.dumpGraph $!! acc +-- Debug.dumpSimplStats +-- return acc +-- !go = eval (afunctionRepr @f) afun Empty +-- -- +-- eval :: AfunctionRepr g (AfunctionR g) (ArraysFunctionR g) +-- -> DelayedOpenAfun aenv (ArraysFunctionR g) +-- -> Val aenv +-- -> AfunctionR g +-- eval (AfunctionReprLam reprF) (Alam lhs f) aenv = \a -> eval reprF f $ aenv `push` (lhs, Sugar.fromArr a) +-- eval AfunctionReprBody (Abody b) aenv = unsafePerformIO $ phase "execute" Debug.elapsed (Sugar.toArr . snd <$> evaluate (evalOpenAcc b aenv)) +-- eval _ _aenv _ = error "Two men say they're Jesus; one of them must be wrong" -- Debugging -- --------- phase :: Builder -> Format Builder (Double -> Double -> Builder) -> IO a -> IO a -phase n fmt go = Debug.timed Debug.dump_phases (now ("phase " <> n <> ": ") % fmt) go - +phase n fmt = Debug.timed Debug.dump_phases (now ("phase " <> n <> ": ") % fmt) --- Delayed Arrays --- -------------- --- Note that in contrast to the representation used in the optimised AST, the --- delayed array representation used here is _only_ for delayed arrays --- we do --- not require an optional Manifest|Delayed data type to evaluate the program. --- -data Delayed a where - Delayed :: ArrayR (Array sh e) - -> sh - -> (sh -> e) - -> (Int -> e) - -> Delayed (Array sh e) -- Array expression evaluation @@ -163,743 +716,44 @@ type WithReprs acc = (ArraysR acc, acc) fromFunction' :: ArrayR (Array sh e) -> sh -> (sh -> e) -> WithReprs (Array sh e) fromFunction' repr sh f = (TupRsingle repr, fromFunction repr sh f) --- Evaluate an open array function --- -evalOpenAfun :: HasCallStack => DelayedOpenAfun aenv f -> Val aenv -> f -evalOpenAfun (Alam lhs f) aenv = \a -> evalOpenAfun f $ aenv `push` (lhs, a) -evalOpenAfun (Abody b) aenv = snd $ evalOpenAcc b aenv - - --- The core interpreter for optimised array programs --- -evalOpenAcc - :: forall aenv a. HasCallStack - => DelayedOpenAcc aenv a - -> Val aenv - -> WithReprs a -evalOpenAcc AST.Delayed{} _ = internalError "expected manifest array" -evalOpenAcc (AST.Manifest pacc) aenv = - let - manifest :: forall a'. HasCallStack => DelayedOpenAcc aenv a' -> WithReprs a' - manifest acc = - let (repr, a') = evalOpenAcc acc aenv - in rnfArraysR repr a' `seq` (repr, a') - - delayed :: DelayedOpenAcc aenv (Array sh e) -> Delayed (Array sh e) - delayed AST.Delayed{..} = Delayed reprD (evalE extentD) (evalF indexD) (evalF linearIndexD) - delayed a' = Delayed aR (shape a) (indexArray aR a) (linearIndexArray (arrayRtype aR) a) - where - (TupRsingle aR, a) = manifest a' - evalE :: Exp aenv t -> t - evalE exp = evalExp exp aenv - evalF :: Fun aenv f -> f - evalF fun = evalFun fun aenv - evalB :: AST.Boundary aenv t -> Boundary t - evalB bnd = evalBoundary bnd aenv - dir :: Direction -> t -> t -> t - dir LeftToRight l _ = l - dir RightToLeft _ r = r - in - case pacc of - Avar (Var repr ix) -> (TupRsingle repr, prj ix aenv) - Alet lhs acc1 acc2 -> evalOpenAcc acc2 $ aenv `push` (lhs, snd $ manifest acc1) - Apair acc1 acc2 -> let (r1, a1) = manifest acc1 - (r2, a2) = manifest acc2 - in - (TupRpair r1 r2, (a1, a2)) - Anil -> (TupRunit, ()) - Atrace msg as bs -> unsafePerformIO $ manifest bs <$ atraceOp msg (snd $ manifest as) - Apply repr afun acc -> (repr, evalOpenAfun afun aenv $ snd $ manifest acc) - Aforeign repr _ afun acc -> (repr, evalOpenAfun afun Empty $ snd $ manifest acc) - Acond p acc1 acc2 - | toBool (evalE p) -> manifest acc1 - | otherwise -> manifest acc2 - - Awhile cond body acc -> (repr, go initial) - where - (repr, initial) = manifest acc - p = evalOpenAfun cond aenv - f = evalOpenAfun body aenv - go !x - | toBool (linearIndexArray (Sugar.eltR @Word8) (p x) 0) = go (f x) - | otherwise = x - - Use repr arr -> (TupRsingle repr, arr) - Unit tp e -> unitOp tp (evalE e) - -- Collect s -> evalSeq defaultSeqConfig s aenv - - -- Producers - -- --------- - Map tp f acc -> mapOp tp (evalF f) (delayed acc) - Generate repr sh f -> generateOp repr (evalE sh) (evalF f) - Transform repr sh p f acc -> transformOp repr (evalE sh) (evalF p) (evalF f) (delayed acc) - Backpermute shr sh p acc -> backpermuteOp shr (evalE sh) (evalF p) (delayed acc) - Reshape shr sh acc -> reshapeOp shr (evalE sh) (manifest acc) - - ZipWith tp f acc1 acc2 -> zipWithOp tp (evalF f) (delayed acc1) (delayed acc2) - Replicate slice slix acc -> replicateOp slice (evalE slix) (manifest acc) - Slice slice acc slix -> sliceOp slice (manifest acc) (evalE slix) - - -- Consumers - -- --------- - Fold f (Just z) acc -> foldOp (evalF f) (evalE z) (delayed acc) - Fold f Nothing acc -> fold1Op (evalF f) (delayed acc) - FoldSeg i f (Just z) acc seg -> foldSegOp i (evalF f) (evalE z) (delayed acc) (delayed seg) - FoldSeg i f Nothing acc seg -> fold1SegOp i (evalF f) (delayed acc) (delayed seg) - Scan d f (Just z) acc -> dir d scanlOp scanrOp (evalF f) (evalE z) (delayed acc) - Scan d f Nothing acc -> dir d scanl1Op scanr1Op (evalF f) (delayed acc) - Scan' d f z acc -> dir d scanl'Op scanr'Op (evalF f) (evalE z) (delayed acc) - Permute f def p acc -> permuteOp (evalF f) (manifest def) (evalF p) (delayed acc) - Stencil s tp sten b acc -> stencilOp s tp (evalF sten) (evalB b) (delayed acc) - Stencil2 s1 s2 tp sten b1 a1 b2 a2 - -> stencil2Op s1 s2 tp (evalF sten) (evalB b1) (delayed a1) (evalB b2) (delayed a2) - - --- Array primitives --- ---------------- - -unitOp :: TypeR e -> e -> WithReprs (Scalar e) -unitOp tp e = fromFunction' (ArrayR ShapeRz tp) () (const e) - - -generateOp - :: ArrayR (Array sh e) - -> sh - -> (sh -> e) - -> WithReprs (Array sh e) -generateOp = fromFunction' - - -transformOp - :: ArrayR (Array sh' b) - -> sh' - -> (sh' -> sh) - -> (a -> b) - -> Delayed (Array sh a) - -> WithReprs (Array sh' b) -transformOp repr sh' p f (Delayed _ _ xs _) - = fromFunction' repr sh' (\ix -> f (xs $ p ix)) - - -reshapeOp - :: HasCallStack - => ShapeR sh - -> sh - -> WithReprs (Array sh' e) - -> WithReprs (Array sh e) -reshapeOp newShapeR newShape (TupRsingle (ArrayR shr tp), (Array sh adata)) - = boundsCheck "shape mismatch" (size newShapeR newShape == size shr sh) - ( TupRsingle (ArrayR newShapeR tp) - , Array newShape adata - ) - - -replicateOp - :: SliceIndex slix sl co sh - -> slix - -> WithReprs (Array sl e) - -> WithReprs (Array sh e) -replicateOp slice slix (TupRsingle repr@(ArrayR _ tp), arr) - = fromFunction' repr' sh (\ix -> (repr, arr) ! pf ix) - where - repr' = ArrayR (sliceDomainR slice) tp - (sh, pf) = extend slice slix (shape arr) - - extend :: SliceIndex slix sl co dim - -> slix - -> sl - -> (dim, dim -> sl) - extend SliceNil () () - = ((), const ()) - extend (SliceAll sliceIdx) (slx, ()) (sl, sz) - = let (dim', f') = extend sliceIdx slx sl - in ((dim', sz), \(ix, i) -> (f' ix, i)) - extend (SliceFixed sliceIdx) (slx, sz) sl - = let (dim', f') = extend sliceIdx slx sl - in ((dim', sz), \(ix, _) -> f' ix) - - -sliceOp - :: SliceIndex slix sl co sh - -> WithReprs (Array sh e) - -> slix - -> WithReprs (Array sl e) -sliceOp slice (TupRsingle repr@(ArrayR _ tp), arr) slix - = fromFunction' repr' sh' (\ix -> (repr, arr) ! pf ix) - where - repr' = ArrayR (sliceShapeR slice) tp - (sh', pf) = restrict slice slix (shape arr) - - restrict - :: HasCallStack - => SliceIndex slix sl co sh - -> slix - -> sh - -> (sl, sl -> sh) - restrict SliceNil () () - = ((), const ()) - restrict (SliceAll sliceIdx) (slx, ()) (sl, sz) - = let (sl', f') = restrict sliceIdx slx sl - in ((sl', sz), \(ix, i) -> (f' ix, i)) - restrict (SliceFixed sliceIdx) (slx, i) (sl, sz) - = let (sl', f') = restrict sliceIdx slx sl - in indexCheck i sz $ (sl', \ix -> (f' ix, i)) - - -mapOp :: TypeR b - -> (a -> b) - -> Delayed (Array sh a) - -> WithReprs (Array sh b) -mapOp tp f (Delayed (ArrayR shr _) sh xs _) - = fromFunction' (ArrayR shr tp) sh (\ix -> f (xs ix)) - - -zipWithOp - :: TypeR c - -> (a -> b -> c) - -> Delayed (Array sh a) - -> Delayed (Array sh b) - -> WithReprs (Array sh c) -zipWithOp tp f (Delayed (ArrayR shr _) shx xs _) (Delayed _ shy ys _) - = fromFunction' (ArrayR shr tp) (intersect shr shx shy) (\ix -> f (xs ix) (ys ix)) - - -foldOp - :: (e -> e -> e) - -> e - -> Delayed (Array (sh, Int) e) - -> WithReprs (Array sh e) -foldOp f z (Delayed (ArrayR (ShapeRsnoc shr) tp) (sh, n) arr _) - = fromFunction' (ArrayR shr tp) sh (\ix -> iter (ShapeRsnoc ShapeRz) ((), n) (\((), i) -> arr (ix, i)) f z) - - -fold1Op - :: HasCallStack - => (e -> e -> e) - -> Delayed (Array (sh, Int) e) - -> WithReprs (Array sh e) -fold1Op f (Delayed (ArrayR (ShapeRsnoc shr) tp) (sh, n) arr _) - = boundsCheck "empty array" (n > 0) - $ fromFunction' (ArrayR shr tp) sh (\ix -> iter1 (ShapeRsnoc ShapeRz) ((), n) (\((), i) -> arr (ix, i)) f) - - -foldSegOp - :: HasCallStack - => IntegralType i - -> (e -> e -> e) - -> e - -> Delayed (Array (sh, Int) e) - -> Delayed (Segments i) - -> WithReprs (Array (sh, Int) e) -foldSegOp itp f z (Delayed repr (sh, _) arr _) (Delayed _ ((), n) _ seg) - | IntegralDict <- integralDict itp - = boundsCheck "empty segment descriptor" (n > 0) - $ fromFunction' repr (sh, n-1) - $ \(sz, ix) -> let start = fromIntegral $ seg ix - end = fromIntegral $ seg (ix+1) - in - boundsCheck "empty segment" (end >= start) - $ iter (ShapeRsnoc ShapeRz) ((), end-start) (\((), i) -> arr (sz, start+i)) f z - - -fold1SegOp - :: HasCallStack - => IntegralType i - -> (e -> e -> e) - -> Delayed (Array (sh, Int) e) - -> Delayed (Segments i) - -> WithReprs (Array (sh, Int) e) -fold1SegOp itp f (Delayed repr (sh, _) arr _) (Delayed _ ((), n) _ seg) - | IntegralDict <- integralDict itp - = boundsCheck "empty segment descriptor" (n > 0) - $ fromFunction' repr (sh, n-1) - $ \(sz, ix) -> let start = fromIntegral $ seg ix - end = fromIntegral $ seg (ix+1) - in - boundsCheck "empty segment" (end > start) - $ iter1 (ShapeRsnoc ShapeRz) ((), end-start) (\((), i) -> arr (sz, start+i)) f - - -scanl1Op - :: forall sh e. HasCallStack - => (e -> e -> e) - -> Delayed (Array (sh, Int) e) - -> WithReprs (Array (sh, Int) e) -scanl1Op f (Delayed (ArrayR shr tp) sh ain _) - = ( TupRsingle $ ArrayR shr tp - , adata `seq` Array sh adata - ) - where - -- - (adata, _) = runArrayData @e $ do - aout <- newArrayData tp (size shr sh) - - let write (sz, 0) = writeArrayData tp aout (toIndex shr sh (sz, 0)) (ain (sz, 0)) - write (sz, i) = do - x <- readArrayData tp aout (toIndex shr sh (sz, i-1)) - let y = ain (sz, i) - writeArrayData tp aout (toIndex shr sh (sz, i)) (f x y) - - iter shr sh write (>>) (return ()) - return (aout, undefined) - - -scanlOp - :: forall sh e. - (e -> e -> e) - -> e - -> Delayed (Array (sh, Int) e) - -> WithReprs (Array (sh, Int) e) -scanlOp f z (Delayed (ArrayR shr tp) (sh, n) ain _) - = ( TupRsingle $ ArrayR shr tp - , adata `seq` Array sh' adata - ) - where - sh' = (sh, n+1) - -- - (adata, _) = runArrayData @e $ do - aout <- newArrayData tp (size shr sh') - - let write (sz, 0) = writeArrayData tp aout (toIndex shr sh' (sz, 0)) z - write (sz, i) = do - x <- readArrayData tp aout (toIndex shr sh' (sz, i-1)) - let y = ain (sz, i-1) - writeArrayData tp aout (toIndex shr sh' (sz, i)) (f x y) - - iter shr sh' write (>>) (return ()) - return (aout, undefined) - - -scanl'Op - :: forall sh e. - (e -> e -> e) - -> e - -> Delayed (Array (sh, Int) e) - -> WithReprs (Array (sh, Int) e, Array sh e) -scanl'Op f z (Delayed (ArrayR shr@(ShapeRsnoc shr') tp) (sh, n) ain _) - = ( TupRsingle (ArrayR shr tp) `TupRpair` TupRsingle (ArrayR shr' tp) - , aout `seq` asum `seq` ( Array (sh, n) aout, Array sh asum ) - ) - where - ((aout, asum), _) = runArrayData @(e, e) $ do - aout <- newArrayData tp (size shr (sh, n)) - asum <- newArrayData tp (size shr' sh) - - let write (sz, 0) - | n == 0 = writeArrayData tp asum (toIndex shr' sh sz) z - | otherwise = writeArrayData tp aout (toIndex shr (sh, n) (sz, 0)) z - write (sz, i) = do - x <- readArrayData tp aout (toIndex shr (sh, n) (sz, i-1)) - let y = ain (sz, i-1) - if i == n - then writeArrayData tp asum (toIndex shr' sh sz) (f x y) - else writeArrayData tp aout (toIndex shr (sh, n) (sz, i)) (f x y) - - iter shr (sh, n+1) write (>>) (return ()) - return ((aout, asum), undefined) - - -scanrOp - :: forall sh e. - (e -> e -> e) - -> e - -> Delayed (Array (sh, Int) e) - -> WithReprs (Array (sh, Int) e) -scanrOp f z (Delayed (ArrayR shr tp) (sz, n) ain _) - = ( TupRsingle (ArrayR shr tp) - , adata `seq` Array sh' adata - ) - where - sh' = (sz, n+1) - -- - (adata, _) = runArrayData @e $ do - aout <- newArrayData tp (size shr sh') - - let write (sz, 0) = writeArrayData tp aout (toIndex shr sh' (sz, n)) z - write (sz, i) = do - let x = ain (sz, n-i) - y <- readArrayData tp aout (toIndex shr sh' (sz, n-i+1)) - writeArrayData tp aout (toIndex shr sh' (sz, n-i)) (f x y) - - iter shr sh' write (>>) (return ()) - return (aout, undefined) - - -scanr1Op - :: forall sh e. HasCallStack - => (e -> e -> e) - -> Delayed (Array (sh, Int) e) - -> WithReprs (Array (sh, Int) e) -scanr1Op f (Delayed (ArrayR shr tp) sh@(_, n) ain _) - = ( TupRsingle $ ArrayR shr tp - , adata `seq` Array sh adata - ) - where - (adata, _) = runArrayData @e $ do - aout <- newArrayData tp (size shr sh) - - let write (sz, 0) = writeArrayData tp aout (toIndex shr sh (sz, n-1)) (ain (sz, n-1)) - write (sz, i) = do - let x = ain (sz, n-i-1) - y <- readArrayData tp aout (toIndex shr sh (sz, n-i)) - writeArrayData tp aout (toIndex shr sh (sz, n-i-1)) (f x y) - - iter shr sh write (>>) (return ()) - return (aout, undefined) - - -scanr'Op - :: forall sh e. - (e -> e -> e) - -> e - -> Delayed (Array (sh, Int) e) - -> WithReprs (Array (sh, Int) e, Array sh e) -scanr'Op f z (Delayed (ArrayR shr@(ShapeRsnoc shr') tp) (sh, n) ain _) - = ( TupRsingle (ArrayR shr tp) `TupRpair` TupRsingle (ArrayR shr' tp) - , aout `seq` asum `seq` ( Array (sh, n) aout, Array sh asum ) - ) - where - ((aout, asum), _) = runArrayData @(e, e) $ do - aout <- newArrayData tp (size shr (sh, n)) - asum <- newArrayData tp (size shr' sh) - - let write (sz, 0) - | n == 0 = writeArrayData tp asum (toIndex shr' sh sz) z - | otherwise = writeArrayData tp aout (toIndex shr (sh, n) (sz, n-1)) z - - write (sz, i) = do - let x = ain (sz, n-i) - y <- readArrayData tp aout (toIndex shr (sh, n) (sz, n-i)) - if i == n - then writeArrayData tp asum (toIndex shr' sh sz) (f x y) - else writeArrayData tp aout (toIndex shr (sh, n) (sz, n-i-1)) (f x y) - - iter shr (sh, n+1) write (>>) (return ()) - return ((aout, asum), undefined) - - -permuteOp - :: forall sh sh' e. HasCallStack - => (e -> e -> e) - -> WithReprs (Array sh' e) - -> (sh -> PrimMaybe sh') - -> Delayed (Array sh e) - -> WithReprs (Array sh' e) -permuteOp f (TupRsingle (ArrayR shr' _), def@(Array _ adef)) p (Delayed (ArrayR shr tp) sh _ ain) - = (TupRsingle $ ArrayR shr' tp, adata `seq` Array sh' adata) - where - sh' = shape def - n' = size shr' sh' - -- - (adata, _) = runArrayData @e $ do - aout <- newArrayData tp n' - - let -- initialise array with default values - init i - | i >= n' = return () - | otherwise = do - x <- readArrayData tp adef i - writeArrayData tp aout i x - init (i+1) - - -- project each element onto the destination array and update - update src - = case p src of - (0,_) -> return () - (1,((),dst)) -> do - let i = toIndex shr sh src - j = toIndex shr' sh' dst - x = ain i - -- - y <- readArrayData tp aout j - writeArrayData tp aout j (f x y) - _ -> internalError "unexpected tag" - - init 0 - iter shr sh update (>>) (return ()) - return (aout, undefined) - - -backpermuteOp - :: ShapeR sh' - -> sh' - -> (sh' -> sh) - -> Delayed (Array sh e) - -> WithReprs (Array sh' e) -backpermuteOp shr sh' p (Delayed (ArrayR _ tp) _ arr _) - = fromFunction' (ArrayR shr tp) sh' (\ix -> arr $ p ix) - - -stencilOp - :: HasCallStack - => StencilR sh a stencil - -> TypeR b - -> (stencil -> b) - -> Boundary (Array sh a) - -> Delayed (Array sh a) - -> WithReprs (Array sh b) -stencilOp stencil tp f bnd arr@(Delayed _ sh _ _) - = fromFunction' (ArrayR shr tp) sh - $ f . stencilAccess stencil (bounded shr bnd arr) - where - shr = stencilShapeR stencil - - -stencil2Op - :: HasCallStack - => StencilR sh a stencil1 - -> StencilR sh b stencil2 - -> TypeR c - -> (stencil1 -> stencil2 -> c) - -> Boundary (Array sh a) - -> Delayed (Array sh a) - -> Boundary (Array sh b) - -> Delayed (Array sh b) - -> WithReprs (Array sh c) -stencil2Op s1 s2 tp stencil bnd1 arr1@(Delayed _ sh1 _ _) bnd2 arr2@(Delayed _ sh2 _ _) - = fromFunction' (ArrayR shr tp) (intersect shr sh1 sh2) f - where - f ix = stencil (stencilAccess s1 (bounded shr bnd1 arr1) ix) - (stencilAccess s2 (bounded shr bnd2 arr2) ix) - shr = stencilShapeR s1 - -stencilAccess - :: StencilR sh e stencil - -> (sh -> e) - -> sh - -> stencil -stencilAccess stencil = goR (stencilShapeR stencil) stencil - where - -- Base cases, nothing interesting to do here since we know the lower - -- dimension is Z. - -- - goR :: ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil - goR _ (StencilRunit3 _) rf ix = - let - (z, i) = ix - rf' d = rf (z, i+d) - in - ((( () - , rf' (-1)) - , rf' 0 ) - , rf' 1 ) - - goR _ (StencilRunit5 _) rf ix = - let (z, i) = ix - rf' d = rf (z, i+d) - in - ((((( () - , rf' (-2)) - , rf' (-1)) - , rf' 0 ) - , rf' 1 ) - , rf' 2 ) - - goR _ (StencilRunit7 _) rf ix = - let (z, i) = ix - rf' d = rf (z, i+d) - in - ((((((( () - , rf' (-3)) - , rf' (-2)) - , rf' (-1)) - , rf' 0 ) - , rf' 1 ) - , rf' 2 ) - , rf' 3 ) - - goR _ (StencilRunit9 _) rf ix = - let (z, i) = ix - rf' d = rf (z, i+d) - in - ((((((((( () - , rf' (-4)) - , rf' (-3)) - , rf' (-2)) - , rf' (-1)) - , rf' 0 ) - , rf' 1 ) - , rf' 2 ) - , rf' 3 ) - , rf' 4 ) - - -- Recursive cases. Note that because the stencil pattern is defined with - -- cons ordering, whereas shapes (and indices) are defined as a snoc-list, - -- when we recurse on the stencil structure we must manipulate the - -- _left-most_ index component. - -- - goR (ShapeRsnoc shr) (StencilRtup3 s1 s2 s3) rf ix = - let (i, ix') = uncons shr ix - rf' d ds = rf (cons shr (i+d) ds) - in - ((( () - , goR shr s1 (rf' (-1)) ix') - , goR shr s2 (rf' 0) ix') - , goR shr s3 (rf' 1) ix') - - goR (ShapeRsnoc shr) (StencilRtup5 s1 s2 s3 s4 s5) rf ix = - let (i, ix') = uncons shr ix - rf' d ds = rf (cons shr (i+d) ds) - in - ((((( () - , goR shr s1 (rf' (-2)) ix') - , goR shr s2 (rf' (-1)) ix') - , goR shr s3 (rf' 0) ix') - , goR shr s4 (rf' 1) ix') - , goR shr s5 (rf' 2) ix') - - goR (ShapeRsnoc shr) (StencilRtup7 s1 s2 s3 s4 s5 s6 s7) rf ix = - let (i, ix') = uncons shr ix - rf' d ds = rf (cons shr (i+d) ds) - in - ((((((( () - , goR shr s1 (rf' (-3)) ix') - , goR shr s2 (rf' (-2)) ix') - , goR shr s3 (rf' (-1)) ix') - , goR shr s4 (rf' 0) ix') - , goR shr s5 (rf' 1) ix') - , goR shr s6 (rf' 2) ix') - , goR shr s7 (rf' 3) ix') - - goR (ShapeRsnoc shr) (StencilRtup9 s1 s2 s3 s4 s5 s6 s7 s8 s9) rf ix = - let (i, ix') = uncons shr ix - rf' d ds = rf (cons shr (i+d) ds) - in - ((((((((( () - , goR shr s1 (rf' (-4)) ix') - , goR shr s2 (rf' (-3)) ix') - , goR shr s3 (rf' (-2)) ix') - , goR shr s4 (rf' (-1)) ix') - , goR shr s5 (rf' 0) ix') - , goR shr s6 (rf' 1) ix') - , goR shr s7 (rf' 2) ix') - , goR shr s8 (rf' 3) ix') - , goR shr s9 (rf' 4) ix') - - -- Add a left-most component to an index - -- - cons :: ShapeR sh -> Int -> sh -> (sh, Int) - cons ShapeRz ix () = ((), ix) - cons (ShapeRsnoc shr) ix (sh, sz) = (cons shr ix sh, sz) - - -- Remove the left-most index of an index, and return the remainder - -- - uncons :: ShapeR sh -> (sh, Int) -> (Int, sh) - uncons ShapeRz ((), v) = (v, ()) - uncons (ShapeRsnoc shr) (v1, v2) = let (i, v1') = uncons shr v1 - in (i, (v1', v2)) - - -bounded - :: HasCallStack - => ShapeR sh - -> Boundary (Array sh e) - -> Delayed (Array sh e) - -> sh - -> e -bounded shr bnd (Delayed _ sh f _) ix = - if inside shr sh ix - then f ix - else - case bnd of - Function g -> g ix - Constant v -> v - _ -> f (bound shr sh ix) - - where - -- Whether the index (second argument) is inside the bounds of the given - -- shape (first argument). - -- - inside :: ShapeR sh -> sh -> sh -> Bool - inside ShapeRz () () = True - inside (ShapeRsnoc shr) (sh, sz) (ih, iz) = iz >= 0 && iz < sz && inside shr sh ih - - -- Return the index (second argument), updated to obey the given boundary - -- conditions when outside the bounds of the given shape (first argument) - -- - bound :: HasCallStack => ShapeR sh -> sh -> sh -> sh - bound ShapeRz () () = () - bound (ShapeRsnoc shr) (sh, sz) (ih, iz) = (bound shr sh ih, ih') - where - ih' - | iz < 0 = case bnd of - Clamp -> 0 - Mirror -> -iz - Wrap -> sz + iz - _ -> internalError "unexpected boundary condition" - | iz >= sz = case bnd of - Clamp -> sz - 1 - Mirror -> sz - (iz - sz + 2) - Wrap -> iz - sz - _ -> internalError "unexpected boundary condition" - | otherwise = iz - --- toSeqOp :: forall slix sl dim co e proxy. (Elt slix, Shape sl, Shape dim, Elt e) --- => SliceIndex (EltRepr slix) --- (EltRepr sl) --- co --- (EltRepr dim) --- -> proxy slix --- -> Array dim e --- -> [Array sl e] --- toSeqOp sliceIndex _ arr = map (sliceOp sliceIndex arr :: slix -> Array sl e) --- (enumSlices sliceIndex (shape arr)) - - --- Stencil boundary conditions --- --------------------------- +-- Scalar expression evaluation +-- ---------------------------- -data Boundary t where - Clamp :: Boundary t - Mirror :: Boundary t - Wrap :: Boundary t - Constant :: t -> Boundary (Array sh t) - Function :: (sh -> e) -> Boundary (Array sh e) - - -evalBoundary :: HasCallStack => AST.Boundary aenv t -> Val aenv -> Boundary t -evalBoundary bnd aenv = - case bnd of - AST.Clamp -> Clamp - AST.Mirror -> Mirror - AST.Wrap -> Wrap - AST.Constant v -> Constant v - AST.Function f -> Function (evalFun f aenv) - -atraceOp :: Message as -> as -> IO () -atraceOp (Message show _ msg) as = - let str = show as - in do - if null str - then T.hPutStrLn stderr msg - else hprint stderr (stext % ": " % string % "\n") msg str - hFlush stderr +newtype EvalArrayInstr m arr = EvalArrayInstr (forall s t. arr (s -> t) -> s -> m t) +evalArrayInstrDefault :: Val aenv -> EvalArrayInstr Identity (ArrayInstr aenv) +evalArrayInstrDefault aenv = EvalArrayInstr $ \instr arg -> case instr of + Index buffer -> return $ indexBuffer (groundRelt $ varType buffer) (prj (varIdx buffer) aenv) arg + Parameter var -> return $ prj (varIdx var) aenv --- Scalar expression evaluation --- ---------------------------- +evalNoArrayInstr :: EvalArrayInstr m NoArrayInstr +evalNoArrayInstr = EvalArrayInstr $ \case {} -- Evaluate a closed scalar expression -- -evalExp :: HasCallStack => Exp aenv t -> Val aenv -> t -evalExp e aenv = evalOpenExp e Empty aenv +evalExp :: (Monad m, HasCallStack) => PreOpenExp arr () t -> EvalArrayInstr m arr -> m t +evalExp e = evalOpenExp e Empty -- Evaluate a closed scalar function -- -evalFun :: HasCallStack => Fun aenv t -> Val aenv -> t -evalFun f aenv = evalOpenFun f Empty aenv +evalFun :: HasCallStack => PreOpenFun arr () t -> EvalArrayInstr Identity arr -> t +evalFun f = evalOpenFun f Empty -- Evaluate an open scalar function -- -evalOpenFun :: HasCallStack => OpenFun env aenv t -> Val env -> Val aenv -> t -evalOpenFun (Body e) env aenv = evalOpenExp e env aenv -evalOpenFun (Lam lhs f) env aenv = - \x -> evalOpenFun f (env `push` (lhs, x)) aenv +evalOpenFun :: HasCallStack => PreOpenFun arr env t -> Val env -> EvalArrayInstr Identity arr -> t +evalOpenFun (Body e) env arr = runIdentity $ evalOpenExp e env arr +evalOpenFun (Lam lhs f) env arr = + \x -> evalOpenFun f (env `push` (lhs, x)) arr +evalUnaryFun :: (Monad m, HasCallStack) => PreOpenFun arr env (s -> t) -> Val env -> EvalArrayInstr m arr -> s -> m t +evalUnaryFun (Lam lhs (Body e)) env arr x = + evalOpenExp e (env `push` (lhs, x)) arr +evalUnaryFun _ _ _ _ = internalError "Expected unary function" -- Evaluate an open scalar expression -- @@ -910,65 +764,58 @@ evalOpenFun (Lam lhs f) env aenv = -- leading to a large amount of wasteful recomputation. -- evalOpenExp - :: forall env aenv t. HasCallStack - => OpenExp env aenv t + :: forall env arr m t. (Monad m, HasCallStack) + => PreOpenExp arr env t -> Val env - -> Val aenv - -> t -evalOpenExp pexp env aenv = + -> EvalArrayInstr m arr + -> m t +evalOpenExp pexp env arr@(EvalArrayInstr runArrayInstr) = let - evalE :: OpenExp env aenv t' -> t' - evalE e = evalOpenExp e env aenv - - evalF :: OpenFun env aenv f' -> f' - evalF f = evalOpenFun f env aenv + evalE :: PreOpenExp arr env t' -> m t' + evalE e = evalOpenExp e env arr - evalA :: ArrayVar aenv a -> WithReprs a - evalA (Var repr ix) = (TupRsingle repr, prj ix aenv) + evalF :: PreOpenFun arr env (s -> t') -> s -> m t' + evalF f = evalUnaryFun f env arr in case pexp of - Let lhs exp1 exp2 -> let !v1 = evalE exp1 - env' = env `push` (lhs, v1) - in evalOpenExp exp2 env' aenv - Evar (Var _ ix) -> prj ix env - Const _ c -> c - Undef tp -> undefElt (TupRsingle tp) - PrimConst c -> evalPrimConst c - PrimApp f x -> evalPrim f (evalE x) - Nil -> () - Pair e1 e2 -> let !x1 = evalE e1 - !x2 = evalE e2 - in (x1, x2) - VecPack vecR e -> pack vecR $! evalE e - VecUnpack vecR e -> unpack vecR $! evalE e - IndexSlice slice slix sh -> restrict slice (evalE slix) - (evalE sh) + Let lhs exp1 exp2 -> do + !v1 <- evalE exp1 + let env' = env `push` (lhs, v1) + evalOpenExp exp2 env' arr + Evar (Var _ ix) -> + return $ prj ix env + Const _ c -> + return c + Undef tp -> + return $ undefElt (TupRsingle tp) + PrimApp f x -> do + !v <- evalE x + return $ evalPrim f v + Nil -> + return () + Pair e1 e2 -> do + !v1 <- evalE e1 + !v2 <- evalE e2 + return (v1, v2) + VecPack vecR e -> do + !v <- evalE e + return $ pack vecR v + VecUnpack vecR e -> do + !v <- evalE e + return $ unpack vecR v + ToIndex shr sh ix -> do + !sh' <- evalE sh + !ix' <- evalE ix + return $ toIndex shr sh' ix' + FromIndex shr sh ix -> do + !sh' <- evalE sh + !ix' <- evalE ix + return $ fromIndex shr sh' ix' + Case e rhs def -> do + !v <- evalE e + evalE $ caseof v rhs where - restrict :: SliceIndex slix sl co sh -> slix -> sh -> sl - restrict SliceNil () () = () - restrict (SliceAll sliceIdx) (slx, ()) (sl, sz) = - let sl' = restrict sliceIdx slx sl - in (sl', sz) - restrict (SliceFixed sliceIdx) (slx, _i) (sl, _sz) = - restrict sliceIdx slx sl - - IndexFull slice slix sh -> extend slice (evalE slix) - (evalE sh) - where - extend :: SliceIndex slix sl co sh -> slix -> sl -> sh - extend SliceNil () () = () - extend (SliceAll sliceIdx) (slx, ()) (sl, sz) = - let sh' = extend sliceIdx slx sl - in (sh', sz) - extend (SliceFixed sliceIdx) (slx, sz) sl = - let sh' = extend sliceIdx slx sl - in (sh', sz) - - ToIndex shr sh ix -> toIndex shr (evalE sh) (evalE ix) - FromIndex shr sh ix -> fromIndex shr (evalE sh) (evalE ix) - Case e rhs def -> evalE (caseof (evalE e) rhs) - where - caseof :: TAG -> [(TAG, OpenExp env aenv t)] -> OpenExp env aenv t + caseof :: TAG -> [(TAG, PreOpenExp arr env t)] -> PreOpenExp arr env t caseof tag = go where go ((t,c):cs) @@ -978,27 +825,47 @@ evalOpenExp pexp env aenv = | Just d <- def = d | otherwise = internalError "unmatched case" - Cond c t e - | toBool (evalE c) -> evalE t - | otherwise -> evalE e + Cond c t e -> do + !v <- evalE c + if toBool v then evalE t else evalE e - While cond body seed -> go (evalE seed) + Select c t e -> do + !v <- evalE c + -- evaluate both branches + let !t' = evalE t + let !e' = evalE e + if toBool v then t' else e' + + While cond body seed -> + evalE seed >>= go where - f = evalF body - p = evalF cond - go !x - | toBool (p x) = go (f x) - | otherwise = x - - Index acc ix -> let (TupRsingle repr, a) = evalA acc - in (repr, a) ! evalE ix - LinearIndex acc i -> let (TupRsingle repr, a) = evalA acc - ix = fromIndex (arrayRshape repr) (shape a) (evalE i) - in (repr, a) ! ix - Shape acc -> shape $ snd $ evalA acc - ShapeSize shr sh -> size shr (evalE sh) - Foreign _ _ f e -> evalOpenFun f Empty Empty $ evalE e - Coerce t1 t2 e -> evalCoerceScalar t1 t2 (evalE e) + go !x = do + c <- evalF cond x + if toBool c then + evalF body x >>= go + else + return x + + ArrayInstr instr ix -> do + ix' <- evalE ix + runArrayInstr instr ix' + ShapeSize shr sh -> do + !sh' <- evalE sh + return $ size shr sh' + Foreign _ _ f e -> do + x <- evalE e + -- We execute 'f' in the Identity monad instead of monad 'm', since + -- it does not contain array instructions. + return $ evalOpenFun f Empty evalNoArrayInstr x + Coerce t1 t2 e -> do + x <- evalE e + return $ evalCoerceScalar t1 t2 x + + Assert msg g e -> do + g' <- evalE g + (if toBool g' then evalE e else errorWithoutStackTrace ("\n*** Assertion failed: " ++ T.unpack msg)) + + Assume _ e -> evalE e -- Coercions @@ -1078,11 +945,6 @@ evalCoerceScalar VectorScalarType{} (SingleScalarType tb) a = scalar tb a -- Scalar primitives -- ----------------- -evalPrimConst :: PrimConst a -> a -evalPrimConst (PrimMinBound ty) = evalMinBound ty -evalPrimConst (PrimMaxBound ty) = evalMaxBound ty -evalPrimConst (PrimPi ty) = evalPi ty - evalPrim :: PrimFun (a -> r) -> (a -> r) evalPrim (PrimAdd ty) = evalAdd ty evalPrim (PrimSub ty) = evalSub ty @@ -1133,12 +995,10 @@ evalPrim (PrimCeiling ta tb) = evalCeiling ta tb evalPrim (PrimAtan2 ty) = evalAtan2 ty evalPrim (PrimIsNaN ty) = evalIsNaN ty evalPrim (PrimIsInfinite ty) = evalIsInfinite ty -evalPrim (PrimLt ty) = evalLt ty -evalPrim (PrimGt ty) = evalGt ty -evalPrim (PrimLtEq ty) = evalLtEq ty -evalPrim (PrimGtEq ty) = evalGtEq ty -evalPrim (PrimEq ty) = evalEq ty -evalPrim (PrimNEq ty) = evalNEq ty +evalPrim (PrimCmp ty CmpLt ) = evalLt ty +evalPrim (PrimCmp ty CmpGtEq ) = evalGtEq ty +evalPrim (PrimCmp ty CmpEq ) = evalEq ty +evalPrim (PrimCmp ty CmpNEq ) = evalNEq ty evalPrim (PrimMax ty) = evalMax ty evalPrim (PrimMin ty) = evalMin ty evalPrim PrimLAnd = evalLAnd @@ -1387,14 +1247,6 @@ evalLt :: SingleType a -> ((a, a) -> PrimBool) evalLt (NumSingleType (IntegralNumType ty)) | IntegralDict <- integralDict ty = fromBool . uncurry (<) evalLt (NumSingleType (FloatingNumType ty)) | FloatingDict <- floatingDict ty = fromBool . uncurry (<) -evalGt :: SingleType a -> ((a, a) -> PrimBool) -evalGt (NumSingleType (IntegralNumType ty)) | IntegralDict <- integralDict ty = fromBool . uncurry (>) -evalGt (NumSingleType (FloatingNumType ty)) | FloatingDict <- floatingDict ty = fromBool . uncurry (>) - -evalLtEq :: SingleType a -> ((a, a) -> PrimBool) -evalLtEq (NumSingleType (IntegralNumType ty)) | IntegralDict <- integralDict ty = fromBool . uncurry (<=) -evalLtEq (NumSingleType (FloatingNumType ty)) | FloatingDict <- floatingDict ty = fromBool . uncurry (<=) - evalGtEq :: SingleType a -> ((a, a) -> PrimBool) evalGtEq (NumSingleType (IntegralNumType ty)) | IntegralDict <- integralDict ty = fromBool . uncurry (>=) evalGtEq (NumSingleType (FloatingNumType ty)) | FloatingDict <- floatingDict ty = fromBool . uncurry (>=) @@ -1414,456 +1266,3 @@ evalMax (NumSingleType (FloatingNumType ty)) | FloatingDict <- floatingDict ty = evalMin :: SingleType a -> ((a, a) -> a) evalMin (NumSingleType (IntegralNumType ty)) | IntegralDict <- integralDict ty = uncurry min evalMin (NumSingleType (FloatingNumType ty)) | FloatingDict <- floatingDict ty = uncurry min - - -{-- --- Sequence evaluation --- --------------- - --- Position in sequence. --- -type SeqPos = Int - --- Configuration for sequence evaluation. --- -data SeqConfig = SeqConfig - { chunkSize :: Int -- Allocation limit for a sequence in - -- words. Actual runtime allocation should be the - -- maximum of this size and the size of the - -- largest element in the sequence. - } - --- Default sequence evaluation configuration for testing purposes. --- -defaultSeqConfig :: SeqConfig -defaultSeqConfig = SeqConfig { chunkSize = 2 } - -type Chunk a = Vector' a - --- The empty chunk. O(1). -emptyChunk :: Arrays a => Chunk a -emptyChunk = empty' - --- Number of arrays in chunk. O(1). --- -clen :: Arrays a => Chunk a -> Int -clen = length' - -elemsPerChunk :: SeqConfig -> Int -> Int -elemsPerChunk conf n - | n < 1 = chunkSize conf - | otherwise = - let (a,b) = chunkSize conf `quotRem` n - in a + signum b - --- Drop a number of arrays from a chunk. O(1). Note: Require keeping a --- scan of element sizes. --- -cdrop :: Arrays a => Int -> Chunk a -> Chunk a -cdrop = drop' dropOp (fst . offsetsOp) - --- Get all the shapes of a chunk of arrays. O(1). --- -chunkShapes :: Chunk (Array sh a) -> Vector sh -chunkShapes = shapes' - --- Get all the elements of a chunk of arrays. O(1). --- -chunkElems :: Chunk (Array sh a) -> Vector a -chunkElems = elements' - --- Convert a vector to a chunk of scalars. --- -vec2Chunk :: Elt e => Vector e -> Chunk (Scalar e) -vec2Chunk = vec2Vec' - --- Convert a list of arrays to a chunk. --- -fromListChunk :: Arrays a => [a] -> Vector' a -fromListChunk = fromList' concatOp - --- Convert a chunk to a list of arrays. --- -toListChunk :: Arrays a => Vector' a -> [a] -toListChunk = toList' fetchAllOp - --- fmap for Chunk. O(n). --- TODO: Use vectorised function. -mapChunk :: (Arrays a, Arrays b) - => (a -> b) - -> Chunk a -> Chunk b -mapChunk f c = fromListChunk $ map f (toListChunk c) - --- zipWith for Chunk. O(n). --- TODO: Use vectorised function. -zipWithChunk :: (Arrays a, Arrays b, Arrays c) - => (a -> b -> c) - -> Chunk a -> Chunk b -> Chunk c -zipWithChunk f c1 c2 = fromListChunk $ zipWith f (toListChunk c1) (toListChunk c2) - --- A window on a sequence. --- -data Window a = Window - { chunk :: Chunk a -- Current allocated chunk. - , wpos :: SeqPos -- Position of the window on the sequence, given - -- in number of elements. - } - --- The initial empty window. --- -window0 :: Arrays a => Window a -window0 = Window { chunk = emptyChunk, wpos = 0 } - --- Index the given window by the given index on the sequence. --- -(!#) :: Arrays a => Window a -> SeqPos -> Chunk a -w !# i - | j <- i - wpos w - , j >= 0 - = cdrop j (chunk w) - -- - | otherwise - = error $ "Window indexed before position. wpos = " ++ show (wpos w) ++ " i = " ++ show i - --- Move the give window by supplying the next chunk. --- -moveWin :: Arrays a => Window a -> Chunk a -> Window a -moveWin w c = w { chunk = c - , wpos = wpos w + clen (chunk w) - } - --- A cursor on a sequence. --- -data Cursor senv a = Cursor - { ref :: Idx senv a -- Reference to the sequence. - , cpos :: SeqPos -- Position of the cursor on the sequence, - -- given in number of elements. - } - --- Initial cursor. --- -cursor0 :: Idx senv a -> Cursor senv a -cursor0 x = Cursor { ref = x, cpos = 0 } - --- Advance cursor by a relative amount. --- -moveCursor :: Int -> Cursor senv a -> Cursor senv a -moveCursor k c = c { cpos = cpos c + k } - --- Valuation for an environment of sequence windows. --- -data Val' senv where - Empty' :: Val' () - Push' :: Val' senv -> Window t -> Val' (senv, t) - --- Projection of a window from a window valuation using a de Bruijn --- index. --- -prj' :: Idx senv t -> Val' senv -> Window t -prj' ZeroIdx (Push' _ v) = v -prj' (SuccIdx idx) (Push' val _) = prj' idx val - --- Projection of a chunk from a window valuation using a sequence --- cursor. --- -prjChunk :: Arrays a => Cursor senv a -> Val' senv -> Chunk a -prjChunk c senv = prj' (ref c) senv !# cpos c - --- An executable sequence. --- -data ExecSeq senv arrs where - ExecP :: Arrays a => Window a -> ExecP senv a -> ExecSeq (senv, a) arrs -> ExecSeq senv arrs - ExecC :: Arrays a => ExecC senv a -> ExecSeq senv a - ExecR :: Arrays a => Cursor senv a -> ExecSeq senv [a] - --- An executable producer. --- -data ExecP senv a where - ExecStreamIn :: Int - -> [a] - -> ExecP senv a - - ExecMap :: Arrays a - => (Chunk a -> Chunk b) - -> Cursor senv a - -> ExecP senv b - - ExecZipWith :: (Arrays a, Arrays b) - => (Chunk a -> Chunk b -> Chunk c) - -> Cursor senv a - -> Cursor senv b - -> ExecP senv c - - -- Stream scan skeleton. - ExecScan :: Arrays a - => (s -> Chunk a -> (Chunk r, s)) -- Chunk scanner. - -> s -- Accumulator (internal state). - -> Cursor senv a -- Input stream. - -> ExecP senv r - --- An executable consumer. --- -data ExecC senv a where - - -- Stream reduction skeleton. - ExecFold :: Arrays a - => (s -> Chunk a -> s) -- Chunk consumer function. - -> (s -> r) -- Finalizer function. - -> s -- Accumulator (internal state). - -> Cursor senv a -- Input stream. - -> ExecC senv r - - ExecStuple :: IsAtuple a - => Atuple (ExecC senv) (TupleRepr a) - -> ExecC senv a - -minCursor :: ExecSeq senv a -> SeqPos -minCursor s = travS s 0 - where - travS :: ExecSeq senv a -> Int -> SeqPos - travS s i = - case s of - ExecP _ p s' -> travP p i `min` travS s' (i+1) - ExecC c -> travC c i - ExecR _ -> maxBound - - k :: Cursor senv a -> Int -> SeqPos - k c i - | i == idxToInt (ref c) = cpos c - | otherwise = maxBound - - travP :: ExecP senv a -> Int -> SeqPos - travP p i = - case p of - ExecStreamIn _ _ -> maxBound - ExecMap _ c -> k c i - ExecZipWith _ c1 c2 -> k c1 i `min` k c2 i - ExecScan _ _ c -> k c i - - travT :: Atuple (ExecC senv) t -> Int -> SeqPos - travT NilAtup _ = maxBound - travT (SnocAtup t c) i = travT t i `min` travC c i - - travC :: ExecC senv a -> Int -> SeqPos - travC c i = - case c of - ExecFold _ _ _ cu -> k cu i - ExecStuple t -> travT t i - - -evalDelayedSeq - :: SeqConfig - -> DelayedSeq arrs - -> arrs -evalDelayedSeq cfg (DelayedSeq aenv s) | aenv' <- evalExtend aenv Empty - = evalSeq cfg s aenv' - -evalSeq :: forall aenv arrs. - SeqConfig - -> PreOpenSeq DelayedOpenAcc aenv () arrs - -> Val aenv -> arrs -evalSeq conf s aenv = evalSeq' s - where - evalSeq' :: PreOpenSeq DelayedOpenAcc aenv senv arrs -> arrs - evalSeq' (Producer _ s) = evalSeq' s - evalSeq' (Consumer _) = loop (initSeq aenv s) - evalSeq' (Reify _) = reify (initSeq aenv s) - - -- Initialize the producers and the accumulators of the consumers - -- with the given array enviroment. - initSeq :: forall senv arrs'. - Val aenv - -> PreOpenSeq DelayedOpenAcc aenv senv arrs' - -> ExecSeq senv arrs' - initSeq aenv s = - case s of - Producer p s' -> ExecP window0 (initProducer p) (initSeq aenv s') - Consumer c -> ExecC (initConsumer c) - Reify ix -> ExecR (cursor0 ix) - - -- Generate a list from the sequence. - reify :: forall arrs. ExecSeq () [arrs] - -> [arrs] - reify s = case step s Empty' of - (Just s', a) -> a ++ reify s' - (Nothing, a) -> a - - -- Iterate the given sequence until it terminates. - -- A sequence only terminates when one of the producers are exhausted. - loop :: Arrays arrs - => ExecSeq () arrs - -> arrs - loop s = - case step' s of - (Nothing, arrs) -> arrs - (Just s', _) -> loop s' - - where - step' :: ExecSeq () arrs -> (Maybe (ExecSeq () arrs), arrs) - step' s = step s Empty' - - -- One iteration of a sequence. - step :: forall senv arrs'. - ExecSeq senv arrs' - -> Val' senv - -> (Maybe (ExecSeq senv arrs'), arrs') - step s senv = - case s of - ExecP w p s' -> - let (c, mp') = produce p senv - finished = 0 == clen (w !# minCursor s') - w' = if finished then moveWin w c else w - (ms'', a) = step s' (senv `Push'` w') - in case ms'' of - Nothing -> (Nothing, a) - Just s'' | finished - , Just p' <- mp' - -> (Just (ExecP w' p' s''), a) - | not finished - -> (Just (ExecP w' p s''), a) - | otherwise - -> (Nothing, a) - ExecC c -> let (c', acc) = consume c senv - in (Just (ExecC c'), acc) - ExecR ix -> let c = prjChunk ix senv in (Just (ExecR (moveCursor (clen c) ix)), toListChunk c) - - evalA :: DelayedOpenAcc aenv a -> a - evalA acc = evalOpenAcc acc aenv - - evalAF :: DelayedOpenAfun aenv f -> f - evalAF f = evalOpenAfun f aenv - - evalE :: DelayedExp aenv t -> t - evalE exp = evalExp exp aenv - - evalF :: DelayedFun aenv f -> f - evalF fun = evalFun fun aenv - - initProducer :: forall a senv. - Producer DelayedOpenAcc aenv senv a - -> ExecP senv a - initProducer p = - case p of - StreamIn arrs -> ExecStreamIn 1 arrs - ToSeq sliceIndex slix (delayed -> Delayed sh ix _) -> - let n = R.size (R.sliceShape sliceIndex (fromElt sh)) - k = elemsPerChunk conf n - in ExecStreamIn k (toSeqOp sliceIndex slix (fromFunction sh ix)) - MapSeq f x -> ExecMap (mapChunk (evalAF f)) (cursor0 x) - ChunkedMapSeq f x -> ExecMap (evalAF f) (cursor0 x) - ZipWithSeq f x y -> ExecZipWith (zipWithChunk (evalAF f)) (cursor0 x) (cursor0 y) - ScanSeq f e x -> ExecScan scanner (evalE e) (cursor0 x) - where - scanner a c = - let v0 = chunkElems c - (v1, a') = scanl'Op (evalF f) a (delayArray v0) - in (vec2Chunk v1, fromScalar a') - - initConsumer :: forall a senv. - Consumer DelayedOpenAcc aenv senv a - -> ExecC senv a - initConsumer c = - case c of - FoldSeq f e x -> - let f' = evalF f - a0 = fromFunction (Z :. chunkSize conf) (const (evalE e)) - consumer v c = zipWith'Op f' (delayArray v) (delayArray (chunkElems c)) - finalizer = fold1Op f' . delayArray - in ExecFold consumer finalizer a0 (cursor0 x) - FoldSeqFlatten f acc x -> - let f' = evalAF f - a0 = evalA acc - consumer a c = f' a (chunkShapes c) (chunkElems c) - in ExecFold consumer id a0 (cursor0 x) - Stuple t -> - let initTup :: Atuple (Consumer DelayedOpenAcc aenv senv) t -> Atuple (ExecC senv) t - initTup NilAtup = NilAtup - initTup (SnocAtup t c) = SnocAtup (initTup t) (initConsumer c) - in ExecStuple (initTup t) - - delayed :: DelayedOpenAcc aenv (Array sh e) -> Delayed (Array sh e) - delayed AST.Manifest{} = $internalError "evalOpenAcc" "expected delayed array" - delayed AST.Delayed{..} = Delayed (evalExp extentD aenv) - (evalFun indexD aenv) - (evalFun linearIndexD aenv) - -produce :: Arrays a => ExecP senv a -> Val' senv -> (Chunk a, Maybe (ExecP senv a)) -produce p senv = - case p of - ExecStreamIn k xs -> - let (xs', xs'') = (take k xs, drop k xs) - c = fromListChunk xs' - mp = if null xs'' - then Nothing - else Just (ExecStreamIn k xs'') - in (c, mp) - ExecMap f x -> - let c = prjChunk x senv - in (f c, Just $ ExecMap f (moveCursor (clen c) x)) - ExecZipWith f x y -> - let c1 = prjChunk x senv - c2 = prjChunk y senv - k = clen c1 `min` clen c2 - in (f c1 c2, Just $ ExecZipWith f (moveCursor k x) (moveCursor k y)) - ExecScan scanner a x -> - let c = prjChunk x senv - (c', a') = scanner a c - k = clen c - in (c', Just $ ExecScan scanner a' (moveCursor k x)) - -consume :: forall senv a. ExecC senv a -> Val' senv -> (ExecC senv a, a) -consume c senv = - case c of - ExecFold f g acc x -> - let c = prjChunk x senv - acc' = f acc c - -- Even though we call g here, lazy evaluation should guarantee it is - -- only ever called once. - in (ExecFold f g acc' (moveCursor (clen c) x), g acc') - ExecStuple t -> - let consT :: Atuple (ExecC senv) t -> (Atuple (ExecC senv) t, t) - consT NilAtup = (NilAtup, ()) - consT (SnocAtup t c) | (c', acc) <- consume c senv - , (t', acc') <- consT t - = (SnocAtup t' c', (acc', acc)) - (t', acc) = consT t - in (ExecStuple t', toAtuple acc) - -evalExtend :: Extend DelayedOpenAcc aenv aenv' -> Val aenv -> Val aenv' -evalExtend BaseEnv aenv = aenv -evalExtend (PushEnv ext1 ext2) aenv | aenv' <- evalExtend ext1 aenv - = Push aenv' (evalOpenAcc ext2 aenv') - -delayArray :: Array sh e -> Delayed (Array sh e) -delayArray arr@(Array _ adata) = Delayed (shape arr) (arr!) (toElt . unsafeIndexArrayData adata) - -fromScalar :: Scalar a -> a -fromScalar = (!Z) - -concatOp :: forall e. Elt e => [Vector e] -> Vector e -concatOp = concatVectors - -fetchAllOp :: (Shape sh, Elt e) => Segments sh -> Vector e -> [Array sh e] -fetchAllOp segs elts - | (offsets, n) <- offsetsOp segs - , (n ! Z) <= size (shape elts) - = [fetch (segs ! (Z :. i)) (offsets ! (Z :. i)) | i <- [0 .. size (shape segs) - 1]] - | otherwise = error $ "illegal argument to fetchAllOp" - where - fetch sh offset = fromFunction sh (\ ix -> elts ! (Z :. ((toIndex sh ix) + offset))) - -dropOp :: Elt e => Int -> Vector e -> Vector e -dropOp i v -- TODO - -- * Implement using C-style pointer-plus. - -- ; dropOp is used often (from prjChunk), - -- so it ought to be efficient O(1). - | n <- size (shape v) - , i <= n - , i >= 0 - = fromFunction (Z :. n - i) (\ (Z :. j) -> v ! (Z :. i + j)) - | otherwise = error $ "illegal argument to drop" - -offsetsOp :: Shape sh => Segments sh -> (Vector Int, Scalar Int) -offsetsOp segs = scanl'Op (+) 0 $ delayArray (mapOp size (delayArray segs)) ---} - diff --git a/src/Data/Array/Accelerate/Interpreter/Simple.hs b/src/Data/Array/Accelerate/Interpreter/Simple.hs new file mode 100644 index 000000000..100ebc7d7 --- /dev/null +++ b/src/Data/Array/Accelerate/Interpreter/Simple.hs @@ -0,0 +1,1353 @@ +{-# LANGUAGE MagicHash #-} +{-# OPTIONS_GHC -Wno-name-shadowing #-} +{-# OPTIONS_HADDOCK prune #-} +-- | +-- Module : Data.Array.Accelerate.Interpreter.Simple +-- Description : Reference backend (interpreted) +-- Copyright : [2008..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- +-- This interpreter is meant to be a reference implementation of the +-- semantics of the embedded array language. The emphasis is on defining +-- the semantics clearly, not on performance. +-- + +module Data.Array.Accelerate.Interpreter.Simple ( + + Smart.Acc, Sugar.Arrays, + Afunction, AfunctionR, + + -- * Interpret an array expression + run, run1, runN, + + -- Internal (hidden) + evalPrim, evalCoerceScalar, atraceOp, + +) where + +import Data.Array.Accelerate.Array.Buffer +import Data.Array.Accelerate.AST hiding ( Boundary(..) ) +import Data.Array.Accelerate.AST.Environment +import Data.Array.Accelerate.AST.Var +import Data.Array.Accelerate.Error +import Data.Array.Accelerate.Representation.Array +import Data.Array.Accelerate.Representation.Elt +import Data.Array.Accelerate.Representation.Shape +import Data.Array.Accelerate.Representation.Slice +import Data.Array.Accelerate.Representation.Stencil +import Data.Array.Accelerate.Representation.Tag +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Representation.Vec +import Data.Array.Accelerate.Trafo +import Data.Array.Accelerate.Trafo.Sharing ( AfunctionR, AfunctionRepr(..), afunctionRepr ) +import qualified Data.Array.Accelerate.Trafo.Sharing as Sharing +import Data.Array.Accelerate.Type +import Data.Primitive.Vec +import qualified Data.Array.Accelerate.AST as AST +import qualified Data.Array.Accelerate.Debug.Internal.Flags as Debug +-- import qualified Data.Array.Accelerate.Debug.Internal.Graph as Debug +import qualified Data.Array.Accelerate.Debug.Internal.Stats as Debug +import qualified Data.Array.Accelerate.Debug.Internal.Timed as Debug +import qualified Data.Array.Accelerate.Smart as Smart +import qualified Data.Array.Accelerate.Sugar.Array as Sugar +import qualified Data.Array.Accelerate.Sugar.Elt as Sugar + +-- import Control.DeepSeq ( ($!!) ) +import Control.Exception +import Control.Monad +import Control.Monad.ST +import Data.Bits +import Data.Primitive.ByteArray +import Data.Primitive.Types +import qualified Data.Text as DT +import Data.Text.Lazy.Builder +import Formatting +import System.IO +import System.IO.Unsafe ( unsafePerformIO ) +import Unsafe.Coerce +import qualified Data.Text.IO as T +import Prelude hiding ( (!!), sum ) + + +-- Program execution +-- ----------------- + +-- | Run a complete embedded array program using the reference interpreter. +-- +run :: (HasCallStack, Sugar.Arrays a) => Smart.Acc a -> a +run = runN + +-- | This is 'runN' specialised to an array program of one argument. +-- +run1 :: (HasCallStack, Sugar.Arrays a, Sugar.Arrays b) => (Smart.Acc a -> Smart.Acc b) -> a -> b +run1 = runN + +-- | Prepare and execute an embedded array program. +-- +runN :: forall f. (HasCallStack, Afunction f) => f -> AfunctionR f +runN f = go + where + !acc = Sharing.convertAfun f + !afun = unsafePerformIO $ do + -- Debug.dumpGraph $!! acc + Debug.dumpSimplStats + return acc + !go = eval (afunctionRepr @f) afun Empty + -- + eval :: forall g aenv. AfunctionRepr g (AfunctionR g) (ArraysFunctionR g) + -> OpenAfun aenv (ArraysFunctionR g) + -> Val aenv + -> AfunctionR g + eval (AfunctionReprLam reprF) (Alam lhs f) aenv = \a -> eval reprF f $ aenv `push` (lhs, Sugar.fromArr a) + eval AfunctionReprBody (Abody b) aenv = unsafePerformIO $ phase "execute" Debug.elapsed (Sugar.toArr . snd <$> evaluate (evalOpenAcc b aenv)) + eval AfunctionReprLam{} (Abody b) _ = arraysRFunctionImpossible (arraysR b) + eval AfunctionReprBody Alam{} _ = arraysRFunctionImpossible (Sugar.arraysR @(AfunctionR g)) + + +-- Debugging +-- --------- + +phase :: Builder -> Format Builder (Double -> Double -> Builder) -> IO a -> IO a +phase n fmt go = Debug.timed Debug.dump_phases (now ("phase " <> n <> ": ") % fmt) go + + +-- Delayed Arrays +-- -------------- + +-- Note that in contrast to the representation used in the optimised AST, the +-- delayed array representation used here is _only_ for delayed arrays --- we do +-- not require an optional Manifest|Delayed data type to evaluate the program. +-- +data Delayed a where + Delayed :: ArrayR (Array sh e) + -> sh + -> (sh -> e) + -> (Int -> e) + -> Delayed (Array sh e) + + +-- Array expression evaluation +-- --------------------------- + +type WithReprs acc = (ArraysR acc, acc) + +fromFunction' :: ArrayR (Array sh e) -> sh -> (sh -> e) -> WithReprs (Array sh e) +fromFunction' repr sh f = (TupRsingle repr, fromFunction repr sh f) + +-- Evaluate an open array function +-- +evalOpenAfun :: HasCallStack => OpenAfun aenv f -> Val aenv -> f +evalOpenAfun (Alam lhs f) aenv = \a -> evalOpenAfun f $ aenv `push` (lhs, a) +evalOpenAfun (Abody b) aenv = snd $ evalOpenAcc b aenv + + +-- The core interpreter for optimised array programs +-- +evalOpenAcc + :: forall aenv a. HasCallStack + => OpenAcc aenv a + -> Val aenv + -> WithReprs a +evalOpenAcc (OpenAcc pacc) aenv = + let + manifest :: forall a'. HasCallStack => OpenAcc aenv a' -> WithReprs a' + manifest acc = + let (repr, a') = evalOpenAcc acc aenv + in rnfArraysR repr a' `seq` (repr, a') + + delayed :: OpenAcc aenv (Array sh e) -> Delayed (Array sh e) + delayed a' = Delayed aR (shape a) (indexArray aR a) (linearIndexArray (arrayRtype aR) a) + where + (TupRsingle aR, a) = manifest a' + + evalE :: Exp aenv t -> t + evalE exp = evalExp exp aenv + + evalF :: Fun aenv f -> f + evalF fun = evalFun fun aenv + + evalB :: AST.Boundary aenv t -> Boundary t + evalB bnd = evalBoundary bnd aenv + + dir :: Direction -> t -> t -> t + dir LeftToRight l _ = l + dir RightToLeft _ r = r + in + case pacc of + Alet lhs acc1 acc2 -> evalOpenAcc acc2 $ aenv `push` (lhs, snd $ manifest acc1) + Avar (Var repr ix) -> (TupRsingle repr, prj ix aenv) + Apair acc1 acc2 -> let (r1, a1) = manifest acc1 + (r2, a2) = manifest acc2 + in + (TupRpair r1 r2, (a1, a2)) + Anil -> (TupRunit, ()) + Manifest as -> manifest as + Aforeign repr _ afun acc -> (repr, evalOpenAfun afun Empty $ snd $ manifest acc) + Acond p acc1 acc2 + | toBool (evalE p) -> manifest acc1 + | otherwise -> manifest acc2 + + Awhile cond body acc -> (repr, go initial) + where + (repr, initial) = manifest acc + p = evalOpenAfun cond aenv + f = evalOpenAfun body aenv + go !x + | toBool (linearIndexArray (Sugar.eltR @Word8) (p x) 0) = go (f x) + | otherwise = x + + Aassert msg cond acc + | toBool (evalE cond) -> manifest acc + | otherwise -> errorWithoutStackTrace ("\n*** Assertion failed: " ++ DT.unpack msg) + + Aassume _ acc -> manifest acc + + Atrace msg as bs -> unsafePerformIO $ manifest bs <$ atraceOp msg (snd $ manifest as) + + Use repr arr -> (TupRsingle repr, arr) + Unit tp e -> unitOp tp (evalE e) + + -- Producers + -- --------- + Reshape shr sh acc -> reshapeOp shr (evalE sh) (manifest acc) + Generate repr sh f -> generateOp repr (evalE sh) (evalF f) + Transform repr sh p f acc -> transformOp repr (evalE sh) (evalF p) (evalF f) (delayed acc) + Replicate slice slix acc -> replicateOp slice (evalE slix) (manifest acc) + Slice slice acc slix -> sliceOp slice (manifest acc) (evalE slix) + Map tp f acc -> mapOp tp (evalF f) (delayed acc) + ZipWith tp f acc1 acc2 -> zipWithOp tp (evalF f) (delayed acc1) (delayed acc2) + Backpermute shr sh p acc -> backpermuteOp shr (evalE sh) (evalF p) (delayed acc) + + -- Consumers + -- --------- + Fold f (Just z) acc -> foldOp (evalF f) (evalE z) (delayed acc) + Fold f Nothing acc -> fold1Op (evalF f) (delayed acc) + FoldSeg i f (Just z) acc seg -> foldSegOp i (evalF f) (evalE z) (delayed acc) (delayed seg) + FoldSeg i f Nothing acc seg -> fold1SegOp i (evalF f) (delayed acc) (delayed seg) + Scan d f (Just z) acc -> dir d scanlOp scanrOp (evalF f) (evalE z) (delayed acc) + Scan d f Nothing acc -> dir d scanl1Op scanr1Op (evalF f) (delayed acc) + Scan' d f z acc -> dir d scanl'Op scanr'Op (evalF f) (evalE z) (delayed acc) + Permute f def acc -> permuteOp (evalF <$> f) (manifest def) (delayed acc) + Stencil s tp sten b acc -> stencilOp s tp (evalF sten) (evalB b) (delayed acc) + Stencil2 s1 s2 tp sten b1 a1 b2 a2 + -> stencil2Op s1 s2 tp (evalF sten) (evalB b1) (delayed a1) (evalB b2) (delayed a2) + + +-- Array primitives +-- ---------------- + +unitOp :: TypeR e -> e -> WithReprs (Scalar e) +unitOp tp e = fromFunction' (ArrayR ShapeRz tp) () (const e) + + +generateOp + :: ArrayR (Array sh e) + -> sh + -> (sh -> e) + -> WithReprs (Array sh e) +generateOp = fromFunction' + + +transformOp + :: ArrayR (Array sh' b) + -> sh' + -> (sh' -> sh) + -> (a -> b) + -> Delayed (Array sh a) + -> WithReprs (Array sh' b) +transformOp repr sh' p f (Delayed _ _ xs _) + = fromFunction' repr sh' (\ix -> f (xs $ p ix)) + + +reshapeOp + :: HasCallStack + => ShapeR sh + -> sh + -> WithReprs (Array sh' e) + -> WithReprs (Array sh e) +reshapeOp newShapeR newShape (TupRsingle (ArrayR shr tp), (Array sh adata)) + = boundsCheck "shape mismatch" (size newShapeR newShape == size shr sh) + ( TupRsingle (ArrayR newShapeR tp) + , Array newShape adata + ) + + +replicateOp + :: SliceIndex slix sl co sh + -> slix + -> WithReprs (Array sl e) + -> WithReprs (Array sh e) +replicateOp slice slix (TupRsingle repr@(ArrayR _ tp), arr) + = fromFunction' repr' sh (\ix -> (repr, arr) ! pf ix) + where + repr' = ArrayR (sliceDomainR slice) tp + (sh, pf) = extend slice slix (shape arr) + + extend :: SliceIndex slix sl co dim + -> slix + -> sl + -> (dim, dim -> sl) + extend SliceNil () () + = ((), const ()) + extend (SliceAll sliceIdx) (slx, ()) (sl, sz) + = let (dim', f') = extend sliceIdx slx sl + in ((dim', sz), \(ix, i) -> (f' ix, i)) + extend (SliceFixed sliceIdx) (slx, sz) sl + = let (dim', f') = extend sliceIdx slx sl + in ((dim', sz), \(ix, _) -> f' ix) + + +sliceOp + :: SliceIndex slix sl co sh + -> WithReprs (Array sh e) + -> slix + -> WithReprs (Array sl e) +sliceOp slice (TupRsingle repr@(ArrayR _ tp), arr) slix + = fromFunction' repr' sh' (\ix -> (repr, arr) ! pf ix) + where + repr' = ArrayR (sliceShapeR slice) tp + (sh', pf) = restrict slice slix (shape arr) + + restrict + :: HasCallStack + => SliceIndex slix sl co sh + -> slix + -> sh + -> (sl, sl -> sh) + restrict SliceNil () () + = ((), const ()) + restrict (SliceAll sliceIdx) (slx, ()) (sl, sz) + = let (sl', f') = restrict sliceIdx slx sl + in ((sl', sz), \(ix, i) -> (f' ix, i)) + restrict (SliceFixed sliceIdx) (slx, i) (sl, sz) + = let (sl', f') = restrict sliceIdx slx sl + in indexCheck i sz $ (sl', \ix -> (f' ix, i)) + + +mapOp :: TypeR b + -> (a -> b) + -> Delayed (Array sh a) + -> WithReprs (Array sh b) +mapOp tp f (Delayed (ArrayR shr _) sh xs _) + = fromFunction' (ArrayR shr tp) sh (\ix -> f (xs ix)) + + +zipWithOp + :: TypeR c + -> (a -> b -> c) + -> Delayed (Array sh a) + -> Delayed (Array sh b) + -> WithReprs (Array sh c) +zipWithOp tp f (Delayed (ArrayR shr _) shx xs _) (Delayed _ shy ys _) + = fromFunction' (ArrayR shr tp) (intersect shr shx shy) (\ix -> f (xs ix) (ys ix)) + + +foldOp + :: (e -> e -> e) + -> e + -> Delayed (Array (sh, Int) e) + -> WithReprs (Array sh e) +foldOp f z (Delayed (ArrayR (ShapeRsnoc shr) tp) (sh, n) arr _) + = fromFunction' (ArrayR shr tp) sh (\ix -> iter (ShapeRsnoc ShapeRz) ((), n) (\((), i) -> arr (ix, i)) f z) + + +fold1Op + :: HasCallStack + => (e -> e -> e) + -> Delayed (Array (sh, Int) e) + -> WithReprs (Array sh e) +fold1Op f (Delayed (ArrayR (ShapeRsnoc shr) tp) (sh, n) arr _) + = boundsCheck "empty array" (n > 0) + $ fromFunction' (ArrayR shr tp) sh (\ix -> iter1 (ShapeRsnoc ShapeRz) ((), n) (\((), i) -> arr (ix, i)) f) + + +foldSegOp + :: HasCallStack + => IntegralType i + -> (e -> e -> e) + -> e + -> Delayed (Array (sh, Int) e) + -> Delayed (Segments i) + -> WithReprs (Array (sh, Int) e) +foldSegOp itp f z (Delayed repr (sh, _) arr _) (Delayed _ ((), n) _ seg) + | IntegralDict <- integralDict itp + = boundsCheck "empty segment descriptor" (n > 0) + $ fromFunction' repr (sh, n-1) + $ \(sz, ix) -> let start = fromIntegral $ seg ix + end = fromIntegral $ seg (ix+1) + in + boundsCheck "empty segment" (end >= start) + $ iter (ShapeRsnoc ShapeRz) ((), end-start) (\((), i) -> arr (sz, start+i)) f z + + +fold1SegOp + :: HasCallStack + => IntegralType i + -> (e -> e -> e) + -> Delayed (Array (sh, Int) e) + -> Delayed (Segments i) + -> WithReprs (Array (sh, Int) e) +fold1SegOp itp f (Delayed repr (sh, _) arr _) (Delayed _ ((), n) _ seg) + | IntegralDict <- integralDict itp + = boundsCheck "empty segment descriptor" (n > 0) + $ fromFunction' repr (sh, n-1) + $ \(sz, ix) -> let start = fromIntegral $ seg ix + end = fromIntegral $ seg (ix+1) + in + boundsCheck "empty segment" (end > start) + $ iter1 (ShapeRsnoc ShapeRz) ((), end-start) (\((), i) -> arr (sz, start+i)) f + + +scanl1Op + :: forall sh e. HasCallStack + => (e -> e -> e) + -> Delayed (Array (sh, Int) e) + -> WithReprs (Array (sh, Int) e) +scanl1Op f (Delayed (ArrayR shr tp) sh ain _) + = ( TupRsingle $ ArrayR shr tp + , adata `seq` Array sh adata + ) + where + -- + (adata, _) = runBuffers tp $ do + aout <- newBuffers tp (size shr sh) + + let write (sz, 0) = writeBuffers tp aout (toIndex shr sh (sz, 0)) (ain (sz, 0)) + write (sz, i) = do + x <- readBuffers tp aout (toIndex shr sh (sz, i-1)) + let y = ain (sz, i) + writeBuffers tp aout (toIndex shr sh (sz, i)) (f x y) + + iter shr sh write (>>) (return ()) + return (aout, undefined) + + +scanlOp + :: forall sh e. HasCallStack + => (e -> e -> e) + -> e + -> Delayed (Array (sh, Int) e) + -> WithReprs (Array (sh, Int) e) +scanlOp f z (Delayed (ArrayR shr tp) (sh, n) ain _) + = ( TupRsingle $ ArrayR shr tp + , adata `seq` Array sh' adata + ) + where + sh' = (sh, n+1) + -- + (adata, _) = runBuffers tp $ do + aout <- newBuffers tp (size shr sh') + + let write (sz, 0) = writeBuffers tp aout (toIndex shr sh' (sz, 0)) z + write (sz, i) = do + x <- readBuffers tp aout (toIndex shr sh' (sz, i-1)) + let y = ain (sz, i-1) + writeBuffers tp aout (toIndex shr sh' (sz, i)) (f x y) + + iter shr sh' write (>>) (return ()) + return (aout, undefined) + + +scanl'Op + :: forall sh e. HasCallStack + => (e -> e -> e) + -> e + -> Delayed (Array (sh, Int) e) + -> WithReprs (Array (sh, Int) e, Array sh e) +scanl'Op f z (Delayed (ArrayR shr@(ShapeRsnoc shr') tp) (sh, n) ain _) + = ( TupRsingle (ArrayR shr tp) `TupRpair` TupRsingle (ArrayR shr' tp) + , aout `seq` asum `seq` ( Array (sh, n) aout, Array sh asum ) + ) + where + ((aout, asum), _) = runBuffers (TupRpair tp tp) $ do + aout <- newBuffers tp (size shr (sh, n)) + asum <- newBuffers tp (size shr' sh) + + let write (sz, 0) + | n == 0 = writeBuffers tp asum (toIndex shr' sh sz) z + | otherwise = writeBuffers tp aout (toIndex shr (sh, n) (sz, 0)) z + write (sz, i) = do + x <- readBuffers tp aout (toIndex shr (sh, n) (sz, i-1)) + let y = ain (sz, i-1) + if i == n + then writeBuffers tp asum (toIndex shr' sh sz) (f x y) + else writeBuffers tp aout (toIndex shr (sh, n) (sz, i)) (f x y) + + iter shr (sh, n+1) write (>>) (return ()) + return ((aout, asum), undefined) + + +scanrOp + :: forall sh e. HasCallStack + => (e -> e -> e) + -> e + -> Delayed (Array (sh, Int) e) + -> WithReprs (Array (sh, Int) e) +scanrOp f z (Delayed (ArrayR shr tp) (sz, n) ain _) + = ( TupRsingle (ArrayR shr tp) + , adata `seq` Array sh' adata + ) + where + sh' = (sz, n+1) + -- + (adata, _) = runBuffers tp $ do + aout <- newBuffers tp (size shr sh') + + let write (sz, 0) = writeBuffers tp aout (toIndex shr sh' (sz, n)) z + write (sz, i) = do + let x = ain (sz, n-i) + y <- readBuffers tp aout (toIndex shr sh' (sz, n-i+1)) + writeBuffers tp aout (toIndex shr sh' (sz, n-i)) (f x y) + + iter shr sh' write (>>) (return ()) + return (aout, undefined) + + +scanr1Op + :: forall sh e. HasCallStack + => (e -> e -> e) + -> Delayed (Array (sh, Int) e) + -> WithReprs (Array (sh, Int) e) +scanr1Op f (Delayed (ArrayR shr tp) sh@(_, n) ain _) + = ( TupRsingle $ ArrayR shr tp + , adata `seq` Array sh adata + ) + where + (adata, _) = runBuffers tp $ do + aout <- newBuffers tp (size shr sh) + + let write (sz, 0) = writeBuffers tp aout (toIndex shr sh (sz, n-1)) (ain (sz, n-1)) + write (sz, i) = do + let x = ain (sz, n-i-1) + y <- readBuffers tp aout (toIndex shr sh (sz, n-i)) + writeBuffers tp aout (toIndex shr sh (sz, n-i-1)) (f x y) + + iter shr sh write (>>) (return ()) + return (aout, undefined) + + +scanr'Op + :: forall sh e. HasCallStack + => (e -> e -> e) + -> e + -> Delayed (Array (sh, Int) e) + -> WithReprs (Array (sh, Int) e, Array sh e) +scanr'Op f z (Delayed (ArrayR shr@(ShapeRsnoc shr') tp) (sh, n) ain _) + = ( TupRsingle (ArrayR shr tp) `TupRpair` TupRsingle (ArrayR shr' tp) + , aout `seq` asum `seq` ( Array (sh, n) aout, Array sh asum ) + ) + where + ((aout, asum), _) = runBuffers (TupRpair tp tp) $ do + aout <- newBuffers tp (size shr (sh, n)) + asum <- newBuffers tp (size shr' sh) + + let write (sz, 0) + | n == 0 = writeBuffers tp asum (toIndex shr' sh sz) z + | otherwise = writeBuffers tp aout (toIndex shr (sh, n) (sz, n-1)) z + + write (sz, i) = do + let x = ain (sz, n-i) + y <- readBuffers tp aout (toIndex shr (sh, n) (sz, n-i)) + if i == n + then writeBuffers tp asum (toIndex shr' sh sz) (f x y) + else writeBuffers tp aout (toIndex shr (sh, n) (sz, n-i-1)) (f x y) + + iter shr (sh, n+1) write (>>) (return ()) + return ((aout, asum), undefined) + + +permuteOp + :: forall sh sh' e. HasCallStack + => Maybe (e -> e -> e) + -> WithReprs (Array sh' e) + -> Delayed (Array sh (PrimMaybe (((), sh'), e))) + -> WithReprs (Array sh' e) +permuteOp mf (TupRsingle (ArrayR shr' _), def) (Delayed (ArrayR shr tp') sh ain _) + = (TupRsingle $ ArrayR shr' tp, adata `seq` Array sh' adata) + where + tp = case tp' of + TupRpair _ (TupRpair _ (TupRpair _ t)) -> t + _ -> error "unreachable" + -- + sh' = shape def + n' = size shr' sh' + -- + (adata, _) = runBuffers tp $ do + aout <- newBuffers tp n' + + let -- initialise array with default values + init i + | i >= n' = return () + | otherwise = do + writeBuffers tp aout i ((tp, def) !! i) + init (i+1) + + -- project each element onto the destination array + updateScatter src + = case ain src of + (0,_) -> return () + (1,((),(((),dst),x))) -> do + let j = toIndex shr' sh' dst + writeBuffers tp aout j x + _ -> internalError "unexpected tag" + + -- project, but combine with the function f + updateWith f src + = case ain src of + (0,_) -> return () + (1,((),(((),dst),x))) -> do + let j = toIndex shr' sh' dst + y <- readBuffers tp aout j + writeBuffers tp aout j (f x y) + _ -> internalError "unexpected tag" + + init 0 + iter shr sh (maybe updateScatter updateWith mf) (>>) (return ()) + return (aout, undefined) + + +backpermuteOp + :: ShapeR sh' + -> sh' + -> (sh' -> sh) + -> Delayed (Array sh e) + -> WithReprs (Array sh' e) +backpermuteOp shr sh' p (Delayed (ArrayR _ tp) _ arr _) + = fromFunction' (ArrayR shr tp) sh' (\ix -> arr $ p ix) + + +stencilOp + :: HasCallStack + => StencilR sh a stencil + -> TypeR b + -> (stencil -> b) + -> Boundary (Array sh a) + -> Delayed (Array sh a) + -> WithReprs (Array sh b) +stencilOp stencil tp f bnd arr@(Delayed _ sh _ _) + = fromFunction' (ArrayR shr tp) sh + $ f . stencilAccess stencil (bounded shr bnd arr) + where + shr = stencilShapeR stencil + + +stencil2Op + :: HasCallStack + => StencilR sh a stencil1 + -> StencilR sh b stencil2 + -> TypeR c + -> (stencil1 -> stencil2 -> c) + -> Boundary (Array sh a) + -> Delayed (Array sh a) + -> Boundary (Array sh b) + -> Delayed (Array sh b) + -> WithReprs (Array sh c) +stencil2Op s1 s2 tp stencil bnd1 arr1@(Delayed _ sh1 _ _) bnd2 arr2@(Delayed _ sh2 _ _) + = fromFunction' (ArrayR shr tp) (intersect shr sh1 sh2) f + where + f ix = stencil (stencilAccess s1 (bounded shr bnd1 arr1) ix) + (stencilAccess s2 (bounded shr bnd2 arr2) ix) + shr = stencilShapeR s1 + +stencilAccess + :: StencilR sh e stencil + -> (sh -> e) + -> sh + -> stencil +stencilAccess stencil = goR (stencilShapeR stencil) stencil + where + -- Base cases, nothing interesting to do here since we know the lower + -- dimension is Z. + -- + goR :: ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil + goR _ (StencilRunit3 _) rf ix = + let + (z, i) = ix + rf' d = rf (z, i+d) + in + ((( () + , rf' (-1)) + , rf' 0 ) + , rf' 1 ) + + goR _ (StencilRunit5 _) rf ix = + let (z, i) = ix + rf' d = rf (z, i+d) + in + ((((( () + , rf' (-2)) + , rf' (-1)) + , rf' 0 ) + , rf' 1 ) + , rf' 2 ) + + goR _ (StencilRunit7 _) rf ix = + let (z, i) = ix + rf' d = rf (z, i+d) + in + ((((((( () + , rf' (-3)) + , rf' (-2)) + , rf' (-1)) + , rf' 0 ) + , rf' 1 ) + , rf' 2 ) + , rf' 3 ) + + goR _ (StencilRunit9 _) rf ix = + let (z, i) = ix + rf' d = rf (z, i+d) + in + ((((((((( () + , rf' (-4)) + , rf' (-3)) + , rf' (-2)) + , rf' (-1)) + , rf' 0 ) + , rf' 1 ) + , rf' 2 ) + , rf' 3 ) + , rf' 4 ) + + -- Recursive cases. Note that because the stencil pattern is defined with + -- cons ordering, whereas shapes (and indices) are defined as a snoc-list, + -- when we recurse on the stencil structure we must manipulate the + -- _left-most_ index component. + -- + goR (ShapeRsnoc shr) (StencilRtup3 s1 s2 s3) rf ix = + let (i, ix') = uncons shr ix + rf' d ds = rf (cons shr (i+d) ds) + in + ((( () + , goR shr s1 (rf' (-1)) ix') + , goR shr s2 (rf' 0) ix') + , goR shr s3 (rf' 1) ix') + + goR (ShapeRsnoc shr) (StencilRtup5 s1 s2 s3 s4 s5) rf ix = + let (i, ix') = uncons shr ix + rf' d ds = rf (cons shr (i+d) ds) + in + ((((( () + , goR shr s1 (rf' (-2)) ix') + , goR shr s2 (rf' (-1)) ix') + , goR shr s3 (rf' 0) ix') + , goR shr s4 (rf' 1) ix') + , goR shr s5 (rf' 2) ix') + + goR (ShapeRsnoc shr) (StencilRtup7 s1 s2 s3 s4 s5 s6 s7) rf ix = + let (i, ix') = uncons shr ix + rf' d ds = rf (cons shr (i+d) ds) + in + ((((((( () + , goR shr s1 (rf' (-3)) ix') + , goR shr s2 (rf' (-2)) ix') + , goR shr s3 (rf' (-1)) ix') + , goR shr s4 (rf' 0) ix') + , goR shr s5 (rf' 1) ix') + , goR shr s6 (rf' 2) ix') + , goR shr s7 (rf' 3) ix') + + goR (ShapeRsnoc shr) (StencilRtup9 s1 s2 s3 s4 s5 s6 s7 s8 s9) rf ix = + let (i, ix') = uncons shr ix + rf' d ds = rf (cons shr (i+d) ds) + in + ((((((((( () + , goR shr s1 (rf' (-4)) ix') + , goR shr s2 (rf' (-3)) ix') + , goR shr s3 (rf' (-2)) ix') + , goR shr s4 (rf' (-1)) ix') + , goR shr s5 (rf' 0) ix') + , goR shr s6 (rf' 1) ix') + , goR shr s7 (rf' 2) ix') + , goR shr s8 (rf' 3) ix') + , goR shr s9 (rf' 4) ix') + + -- Add a left-most component to an index + -- + cons :: ShapeR sh -> Int -> sh -> (sh, Int) + cons ShapeRz ix () = ((), ix) + cons (ShapeRsnoc shr) ix (sh, sz) = (cons shr ix sh, sz) + + -- Remove the left-most index of an index, and return the remainder + -- + uncons :: ShapeR sh -> (sh, Int) -> (Int, sh) + uncons ShapeRz ((), v) = (v, ()) + uncons (ShapeRsnoc shr) (v1, v2) = let (i, v1') = uncons shr v1 + in (i, (v1', v2)) + + +bounded + :: HasCallStack + => ShapeR sh + -> Boundary (Array sh e) + -> Delayed (Array sh e) + -> sh + -> e +bounded shr bnd (Delayed _ sh f _) ix = + if inside shr sh ix + then f ix + else + case bnd of + Function g -> g ix + Constant v -> v + _ -> f (bound shr sh ix) + + where + -- Whether the index (second argument) is inside the bounds of the given + -- shape (first argument). + -- + inside :: ShapeR sh -> sh -> sh -> Bool + inside ShapeRz () () = True + inside (ShapeRsnoc shr) (sh, sz) (ih, iz) = iz >= 0 && iz < sz && inside shr sh ih + + -- Return the index (second argument), updated to obey the given boundary + -- conditions when outside the bounds of the given shape (first argument) + -- + bound :: HasCallStack => ShapeR sh -> sh -> sh -> sh + bound ShapeRz () () = () + bound (ShapeRsnoc shr) (sh, sz) (ih, iz) = (bound shr sh ih, ih') + where + ih' + | iz < 0 = case bnd of + Clamp -> 0 + Mirror -> -iz + Wrap -> sz + iz + _ -> internalError "unexpected boundary condition" + | iz >= sz = case bnd of + Clamp -> sz - 1 + Mirror -> sz - (iz - sz + 2) + Wrap -> iz - sz + _ -> internalError "unexpected boundary condition" + | otherwise = iz + + +-- Stencil boundary conditions +-- --------------------------- + +data Boundary t where + Clamp :: Boundary t + Mirror :: Boundary t + Wrap :: Boundary t + Constant :: t -> Boundary (Array sh t) + Function :: (sh -> e) -> Boundary (Array sh e) + + +evalBoundary :: HasCallStack => AST.Boundary aenv t -> Val aenv -> Boundary t +evalBoundary bnd aenv = + case bnd of + AST.Clamp -> Clamp + AST.Mirror -> Mirror + AST.Wrap -> Wrap + AST.Constant v -> Constant v + AST.Function f -> Function (evalFun f aenv) + +atraceOp :: Message as -> as -> IO () +atraceOp (Message f _ msg) as = + let str = f as + in do + if null str + then T.hPutStrLn stderr msg + else hprint stderr (stext % ": " % string % "\n") msg str + hFlush stderr + + +-- Scalar expression evaluation +-- ---------------------------- + +-- Evaluate a closed scalar expression +-- +evalExp :: HasCallStack => Exp aenv t -> Val aenv -> t +evalExp e aenv = evalOpenExp e (runArrayInstr aenv) Empty + +-- Evaluate a closed scalar function +-- +evalFun :: HasCallStack => Fun aenv t -> Val aenv -> t +evalFun f aenv = evalOpenFun f (runArrayInstr aenv) Empty + +-- Evaluate an open scalar function +-- +evalOpenFun :: HasCallStack => PreOpenFun arr env t -> (forall a b. arr (a -> b) -> a -> b) -> Val env -> t +evalOpenFun (Body e) runarr env = evalOpenExp e runarr env +evalOpenFun (Lam lhs f) runarr env = + \x -> evalOpenFun f runarr (env `push` (lhs, x)) + + +-- Evaluate an open scalar expression +-- +-- NB: The implementation of 'Index' and 'Shape' demonstrate clearly why +-- array expressions must be hoisted out of scalar expressions before code +-- execution. If these operations are in the body of a function that gets +-- mapped over an array, the array argument would be evaluated many times +-- leading to a large amount of wasteful recomputation. +-- +evalOpenExp + :: forall env arr t. HasCallStack + => PreOpenExp arr env t + -> (forall a b. arr (a -> b) -> a -> b) + -> Val env + -> t +evalOpenExp pexp runarr env = + let + evalE :: PreOpenExp arr env t' -> t' + evalE e = evalOpenExp e runarr env + + evalF :: PreOpenFun arr env f' -> f' + evalF f = evalOpenFun f runarr env + in + case pexp of + Let lhs exp1 exp2 -> let !v1 = evalE exp1 + env' = env `push` (lhs, v1) + in evalOpenExp exp2 runarr env' + Evar (Var _ ix) -> prj ix env + Const _ c -> c + Undef tp -> undefElt (TupRsingle tp) + PrimApp f x -> evalPrim f (evalE x) + Nil -> () + Pair e1 e2 -> let !x1 = evalE e1 + !x2 = evalE e2 + in (x1, x2) + VecPack vecR e -> pack vecR $! evalE e + VecUnpack vecR e -> unpack vecR $! evalE e + ToIndex shr sh ix -> toIndex shr (evalE sh) (evalE ix) + FromIndex shr sh ix -> fromIndex shr (evalE sh) (evalE ix) + Case e rhs def -> evalE (caseof (evalE e) rhs) + where + caseof :: TAG -> [(TAG, PreOpenExp arr env t)] -> PreOpenExp arr env t + caseof tag = go + where + go ((t,c):cs) + | tag == t = c + | otherwise = go cs + go [] + | Just d <- def = d + | otherwise = internalError "unmatched case" + + Cond c t e + | toBool (evalE c) -> evalE t + | otherwise -> evalE e + + Select c t e -> + -- evaluate both branches + let !t' = evalE t + !e' = evalE e + in if toBool (evalE c) then t' else e' + + While cond body seed -> go (evalE seed) + where + f = evalF body + p = evalF cond + go !x + | toBool (p x) = go (f x) + | otherwise = x + + ArrayInstr instr arg -> runarr instr (evalE arg) + + ShapeSize shr sh -> size shr (evalE sh) + Foreign _ _ f e -> evalOpenFun f (\case {}) Empty $ evalE e + Coerce t1 t2 e -> evalCoerceScalar t1 t2 (evalE e) + + Assert msg c e + | toBool (evalE c) -> evalE e + | otherwise -> errorWithoutStackTrace $ "\n*** Assertion failed: " ++ DT.unpack msg + + Assume _ e -> evalE e + +runArrayInstr :: forall aenv a b. HasCallStack => Val aenv -> ArrayInstr aenv (a -> b) -> a -> b +runArrayInstr aenv instr arg = + case instr of + Shape acc -> shape $ snd $ evalA acc + Index acc -> let (TupRsingle repr, a) = evalA acc + in (repr, a) ! arg + LinearIndex acc -> let (TupRsingle repr, a) = evalA acc + ix = fromIndex (arrayRshape repr) (shape a) arg + in (repr, a) ! ix + where + evalA :: ArrayVar aenv t -> WithReprs t + evalA (Var repr ix) = (TupRsingle repr, prj ix aenv) + +-- Coercions +-- --------- + +-- Coercion between two scalar types. We require that the size of the source and +-- destination values are equal (this is not checked at this point). +-- +evalCoerceScalar :: ScalarType a -> ScalarType b -> a -> b +evalCoerceScalar SingleScalarType{} SingleScalarType{} a = unsafeCoerce a +evalCoerceScalar VectorScalarType{} VectorScalarType{} a = unsafeCoerce a -- XXX: or just unpack/repack the (Vec ba#) +evalCoerceScalar (SingleScalarType ta) VectorScalarType{} a = vector ta a + where + vector :: SingleType a -> a -> Vec n b + vector (NumSingleType t) = num t + + num :: NumType a -> a -> Vec n b + num (IntegralNumType t) = integral t + num (FloatingNumType t) = floating t + + integral :: IntegralType a -> a -> Vec n b + integral TypeInt{} = poke + integral TypeInt8{} = poke + integral TypeInt16{} = poke + integral TypeInt32{} = poke + integral TypeInt64{} = poke + integral TypeWord{} = poke + integral TypeWord8{} = poke + integral TypeWord16{} = poke + integral TypeWord32{} = poke + integral TypeWord64{} = poke + + floating :: FloatingType a -> a -> Vec n b + floating TypeHalf{} = poke + floating TypeFloat{} = poke + floating TypeDouble{} = poke + + {-# INLINE poke #-} + poke :: forall a b n. Prim a => a -> Vec n b + poke x = runST $ do + mba <- newByteArray (sizeOf (undefined::a)) + writeByteArray mba 0 x + ByteArray ba# <- unsafeFreezeByteArray mba + return $ Vec ba# + +evalCoerceScalar VectorScalarType{} (SingleScalarType tb) a = scalar tb a + where + scalar :: SingleType b -> Vec n a -> b + scalar (NumSingleType t) = num t + + num :: NumType b -> Vec n a -> b + num (IntegralNumType t) = integral t + num (FloatingNumType t) = floating t + + integral :: IntegralType b -> Vec n a -> b + integral TypeInt{} = peek + integral TypeInt8{} = peek + integral TypeInt16{} = peek + integral TypeInt32{} = peek + integral TypeInt64{} = peek + integral TypeWord{} = peek + integral TypeWord8{} = peek + integral TypeWord16{} = peek + integral TypeWord32{} = peek + integral TypeWord64{} = peek + + floating :: FloatingType b -> Vec n a -> b + floating TypeHalf{} = peek + floating TypeFloat{} = peek + floating TypeDouble{} = peek + + {-# INLINE peek #-} + peek :: Prim a => Vec n b -> a + peek (Vec ba#) = indexByteArray (ByteArray ba#) 0 + + +-- Scalar primitives +-- ----------------- + +evalPrim :: PrimFun (a -> r) -> (a -> r) +evalPrim (PrimAdd ty) = evalAdd ty +evalPrim (PrimSub ty) = evalSub ty +evalPrim (PrimMul ty) = evalMul ty +evalPrim (PrimNeg ty) = evalNeg ty +evalPrim (PrimAbs ty) = evalAbs ty +evalPrim (PrimSig ty) = evalSig ty +evalPrim (PrimQuot ty) = evalQuot ty +evalPrim (PrimRem ty) = evalRem ty +evalPrim (PrimQuotRem ty) = evalQuotRem ty +evalPrim (PrimIDiv ty) = evalIDiv ty +evalPrim (PrimMod ty) = evalMod ty +evalPrim (PrimDivMod ty) = evalDivMod ty +evalPrim (PrimBAnd ty) = evalBAnd ty +evalPrim (PrimBOr ty) = evalBOr ty +evalPrim (PrimBXor ty) = evalBXor ty +evalPrim (PrimBNot ty) = evalBNot ty +evalPrim (PrimBShiftL ty) = evalBShiftL ty +evalPrim (PrimBShiftR ty) = evalBShiftR ty +evalPrim (PrimBRotateL ty) = evalBRotateL ty +evalPrim (PrimBRotateR ty) = evalBRotateR ty +evalPrim (PrimPopCount ty) = evalPopCount ty +evalPrim (PrimCountLeadingZeros ty) = evalCountLeadingZeros ty +evalPrim (PrimCountTrailingZeros ty) = evalCountTrailingZeros ty +evalPrim (PrimFDiv ty) = evalFDiv ty +evalPrim (PrimRecip ty) = evalRecip ty +evalPrim (PrimSin ty) = evalSin ty +evalPrim (PrimCos ty) = evalCos ty +evalPrim (PrimTan ty) = evalTan ty +evalPrim (PrimAsin ty) = evalAsin ty +evalPrim (PrimAcos ty) = evalAcos ty +evalPrim (PrimAtan ty) = evalAtan ty +evalPrim (PrimSinh ty) = evalSinh ty +evalPrim (PrimCosh ty) = evalCosh ty +evalPrim (PrimTanh ty) = evalTanh ty +evalPrim (PrimAsinh ty) = evalAsinh ty +evalPrim (PrimAcosh ty) = evalAcosh ty +evalPrim (PrimAtanh ty) = evalAtanh ty +evalPrim (PrimExpFloating ty) = evalExpFloating ty +evalPrim (PrimSqrt ty) = evalSqrt ty +evalPrim (PrimLog ty) = evalLog ty +evalPrim (PrimFPow ty) = evalFPow ty +evalPrim (PrimLogBase ty) = evalLogBase ty +evalPrim (PrimTruncate ta tb) = evalTruncate ta tb +evalPrim (PrimRound ta tb) = evalRound ta tb +evalPrim (PrimFloor ta tb) = evalFloor ta tb +evalPrim (PrimCeiling ta tb) = evalCeiling ta tb +evalPrim (PrimAtan2 ty) = evalAtan2 ty +evalPrim (PrimIsNaN ty) = evalIsNaN ty +evalPrim (PrimIsInfinite ty) = evalIsInfinite ty +evalPrim (PrimCmp ty CmpLt ) = evalLt ty +evalPrim (PrimCmp ty CmpGtEq ) = evalGtEq ty +evalPrim (PrimCmp ty CmpEq ) = evalEq ty +evalPrim (PrimCmp ty CmpNEq ) = evalNEq ty +evalPrim (PrimMax ty) = evalMax ty +evalPrim (PrimMin ty) = evalMin ty +evalPrim PrimLAnd = evalLAnd +evalPrim PrimLOr = evalLOr +evalPrim PrimLNot = evalLNot +evalPrim (PrimFromIntegral ta tb) = evalFromIntegral ta tb +evalPrim (PrimToFloating ta tb) = evalToFloating ta tb + + +-- Implementation of scalar primitives +-- ----------------------------------- + +toBool :: PrimBool -> Bool +toBool 0 = False +toBool _ = True + +fromBool :: Bool -> PrimBool +fromBool False = 0 +fromBool True = 1 + +evalLAnd :: (PrimBool, PrimBool) -> PrimBool +evalLAnd (x, y) = fromBool (toBool x && toBool y) + +evalLOr :: (PrimBool, PrimBool) -> PrimBool +evalLOr (x, y) = fromBool (toBool x || toBool y) + +evalLNot :: PrimBool -> PrimBool +evalLNot = fromBool . not . toBool + +evalFromIntegral :: IntegralType a -> NumType b -> a -> b +evalFromIntegral ta (IntegralNumType tb) + | IntegralDict <- integralDict ta + , IntegralDict <- integralDict tb + = fromIntegral + +evalFromIntegral ta (FloatingNumType tb) + | IntegralDict <- integralDict ta + , FloatingDict <- floatingDict tb + = fromIntegral + +evalToFloating :: NumType a -> FloatingType b -> a -> b +evalToFloating (IntegralNumType ta) tb + | IntegralDict <- integralDict ta + , FloatingDict <- floatingDict tb + = realToFrac + +evalToFloating (FloatingNumType ta) tb + | FloatingDict <- floatingDict ta + , FloatingDict <- floatingDict tb + = realToFrac + + +-- Extract methods from reified dictionaries +-- + +-- Constant method of floating +-- + +evalSin :: FloatingType a -> (a -> a) +evalSin ty | FloatingDict <- floatingDict ty = sin + +evalCos :: FloatingType a -> (a -> a) +evalCos ty | FloatingDict <- floatingDict ty = cos + +evalTan :: FloatingType a -> (a -> a) +evalTan ty | FloatingDict <- floatingDict ty = tan + +evalAsin :: FloatingType a -> (a -> a) +evalAsin ty | FloatingDict <- floatingDict ty = asin + +evalAcos :: FloatingType a -> (a -> a) +evalAcos ty | FloatingDict <- floatingDict ty = acos + +evalAtan :: FloatingType a -> (a -> a) +evalAtan ty | FloatingDict <- floatingDict ty = atan + +evalSinh :: FloatingType a -> (a -> a) +evalSinh ty | FloatingDict <- floatingDict ty = sinh + +evalCosh :: FloatingType a -> (a -> a) +evalCosh ty | FloatingDict <- floatingDict ty = cosh + +evalTanh :: FloatingType a -> (a -> a) +evalTanh ty | FloatingDict <- floatingDict ty = tanh + +evalAsinh :: FloatingType a -> (a -> a) +evalAsinh ty | FloatingDict <- floatingDict ty = asinh + +evalAcosh :: FloatingType a -> (a -> a) +evalAcosh ty | FloatingDict <- floatingDict ty = acosh + +evalAtanh :: FloatingType a -> (a -> a) +evalAtanh ty | FloatingDict <- floatingDict ty = atanh + +evalExpFloating :: FloatingType a -> (a -> a) +evalExpFloating ty | FloatingDict <- floatingDict ty = exp + +evalSqrt :: FloatingType a -> (a -> a) +evalSqrt ty | FloatingDict <- floatingDict ty = sqrt + +evalLog :: FloatingType a -> (a -> a) +evalLog ty | FloatingDict <- floatingDict ty = log + +evalFPow :: FloatingType a -> ((a, a) -> a) +evalFPow ty | FloatingDict <- floatingDict ty = uncurry (**) + +evalLogBase :: FloatingType a -> ((a, a) -> a) +evalLogBase ty | FloatingDict <- floatingDict ty = uncurry logBase + +evalTruncate :: FloatingType a -> IntegralType b -> (a -> b) +evalTruncate ta tb + | FloatingDict <- floatingDict ta + , IntegralDict <- integralDict tb + = truncate + +evalRound :: FloatingType a -> IntegralType b -> (a -> b) +evalRound ta tb + | FloatingDict <- floatingDict ta + , IntegralDict <- integralDict tb + = round + +evalFloor :: FloatingType a -> IntegralType b -> (a -> b) +evalFloor ta tb + | FloatingDict <- floatingDict ta + , IntegralDict <- integralDict tb + = floor + +evalCeiling :: FloatingType a -> IntegralType b -> (a -> b) +evalCeiling ta tb + | FloatingDict <- floatingDict ta + , IntegralDict <- integralDict tb + = ceiling + +evalAtan2 :: FloatingType a -> ((a, a) -> a) +evalAtan2 ty | FloatingDict <- floatingDict ty = uncurry atan2 + +evalIsNaN :: FloatingType a -> (a -> PrimBool) +evalIsNaN ty | FloatingDict <- floatingDict ty = fromBool . isNaN + +evalIsInfinite :: FloatingType a -> (a -> PrimBool) +evalIsInfinite ty | FloatingDict <- floatingDict ty = fromBool . isInfinite + + +-- Methods of Num +-- + +evalAdd :: NumType a -> ((a, a) -> a) +evalAdd (IntegralNumType ty) | IntegralDict <- integralDict ty = uncurry (+) +evalAdd (FloatingNumType ty) | FloatingDict <- floatingDict ty = uncurry (+) + +evalSub :: NumType a -> ((a, a) -> a) +evalSub (IntegralNumType ty) | IntegralDict <- integralDict ty = uncurry (-) +evalSub (FloatingNumType ty) | FloatingDict <- floatingDict ty = uncurry (-) + +evalMul :: NumType a -> ((a, a) -> a) +evalMul (IntegralNumType ty) | IntegralDict <- integralDict ty = uncurry (*) +evalMul (FloatingNumType ty) | FloatingDict <- floatingDict ty = uncurry (*) + +evalNeg :: NumType a -> (a -> a) +evalNeg (IntegralNumType ty) | IntegralDict <- integralDict ty = negate +evalNeg (FloatingNumType ty) | FloatingDict <- floatingDict ty = negate + +evalAbs :: NumType a -> (a -> a) +evalAbs (IntegralNumType ty) | IntegralDict <- integralDict ty = abs +evalAbs (FloatingNumType ty) | FloatingDict <- floatingDict ty = abs + +evalSig :: NumType a -> (a -> a) +evalSig (IntegralNumType ty) | IntegralDict <- integralDict ty = signum +evalSig (FloatingNumType ty) | FloatingDict <- floatingDict ty = signum + +evalQuot :: IntegralType a -> ((a, a) -> a) +evalQuot ty | IntegralDict <- integralDict ty = uncurry quot + +evalRem :: IntegralType a -> ((a, a) -> a) +evalRem ty | IntegralDict <- integralDict ty = uncurry rem + +evalQuotRem :: IntegralType a -> ((a, a) -> (a, a)) +evalQuotRem ty | IntegralDict <- integralDict ty = uncurry quotRem + +evalIDiv :: IntegralType a -> ((a, a) -> a) +evalIDiv ty | IntegralDict <- integralDict ty = uncurry div + +evalMod :: IntegralType a -> ((a, a) -> a) +evalMod ty | IntegralDict <- integralDict ty = uncurry mod + +evalDivMod :: IntegralType a -> ((a, a) -> (a, a)) +evalDivMod ty | IntegralDict <- integralDict ty = uncurry divMod + +evalBAnd :: IntegralType a -> ((a, a) -> a) +evalBAnd ty | IntegralDict <- integralDict ty = uncurry (.&.) + +evalBOr :: IntegralType a -> ((a, a) -> a) +evalBOr ty | IntegralDict <- integralDict ty = uncurry (.|.) + +evalBXor :: IntegralType a -> ((a, a) -> a) +evalBXor ty | IntegralDict <- integralDict ty = uncurry xor + +evalBNot :: IntegralType a -> (a -> a) +evalBNot ty | IntegralDict <- integralDict ty = complement + +evalBShiftL :: IntegralType a -> ((a, Int) -> a) +evalBShiftL ty | IntegralDict <- integralDict ty = uncurry shiftL + +evalBShiftR :: IntegralType a -> ((a, Int) -> a) +evalBShiftR ty | IntegralDict <- integralDict ty = uncurry shiftR + +evalBRotateL :: IntegralType a -> ((a, Int) -> a) +evalBRotateL ty | IntegralDict <- integralDict ty = uncurry rotateL + +evalBRotateR :: IntegralType a -> ((a, Int) -> a) +evalBRotateR ty | IntegralDict <- integralDict ty = uncurry rotateR + +evalPopCount :: IntegralType a -> (a -> Int) +evalPopCount ty | IntegralDict <- integralDict ty = popCount + +evalCountLeadingZeros :: IntegralType a -> (a -> Int) +evalCountLeadingZeros ty | IntegralDict <- integralDict ty = countLeadingZeros + +evalCountTrailingZeros :: IntegralType a -> (a -> Int) +evalCountTrailingZeros ty | IntegralDict <- integralDict ty = countTrailingZeros + +evalFDiv :: FloatingType a -> ((a, a) -> a) +evalFDiv ty | FloatingDict <- floatingDict ty = uncurry (/) + +evalRecip :: FloatingType a -> (a -> a) +evalRecip ty | FloatingDict <- floatingDict ty = recip + + +evalLt :: SingleType a -> ((a, a) -> PrimBool) +evalLt (NumSingleType (IntegralNumType ty)) | IntegralDict <- integralDict ty = fromBool . uncurry (<) +evalLt (NumSingleType (FloatingNumType ty)) | FloatingDict <- floatingDict ty = fromBool . uncurry (<) + +evalGtEq :: SingleType a -> ((a, a) -> PrimBool) +evalGtEq (NumSingleType (IntegralNumType ty)) | IntegralDict <- integralDict ty = fromBool . uncurry (>=) +evalGtEq (NumSingleType (FloatingNumType ty)) | FloatingDict <- floatingDict ty = fromBool . uncurry (>=) + +evalEq :: SingleType a -> ((a, a) -> PrimBool) +evalEq (NumSingleType (IntegralNumType ty)) | IntegralDict <- integralDict ty = fromBool . uncurry (==) +evalEq (NumSingleType (FloatingNumType ty)) | FloatingDict <- floatingDict ty = fromBool . uncurry (==) + +evalNEq :: SingleType a -> ((a, a) -> PrimBool) +evalNEq (NumSingleType (IntegralNumType ty)) | IntegralDict <- integralDict ty = fromBool . uncurry (/=) +evalNEq (NumSingleType (FloatingNumType ty)) | FloatingDict <- floatingDict ty = fromBool . uncurry (/=) + +evalMax :: SingleType a -> ((a, a) -> a) +evalMax (NumSingleType (IntegralNumType ty)) | IntegralDict <- integralDict ty = uncurry max +evalMax (NumSingleType (FloatingNumType ty)) | FloatingDict <- floatingDict ty = uncurry max + +evalMin :: SingleType a -> ((a, a) -> a) +evalMin (NumSingleType (IntegralNumType ty)) | IntegralDict <- integralDict ty = uncurry min +evalMin (NumSingleType (FloatingNumType ty)) | FloatingDict <- floatingDict ty = uncurry min diff --git a/src/Data/Array/Accelerate/Language.hs b/src/Data/Array/Accelerate/Language.hs index d32cd94a5..a7e5dd977 100644 --- a/src/Data/Array/Accelerate/Language.hs +++ b/src/Data/Array/Accelerate/Language.hs @@ -1,13 +1,5 @@ -{-# LANGUAGE AllowAmbiguousTypes #-} -{-# LANGUAGE CPP #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE ViewPatterns #-} -- | -- Module : Data.Array.Accelerate.Language -- Copyright : [2008..2020] The Accelerate Team @@ -58,7 +50,7 @@ module Data.Array.Accelerate.Language ( scanl, scanl', scanl1, scanr, scanr', scanr1, -- * Permutations - permute, backpermute, + permute', permuteUnique', backpermute, -- * Stencil operations stencil, stencil2, @@ -78,8 +70,8 @@ module Data.Array.Accelerate.Language ( foreignAcc, foreignExp, - -- * Pipelining - (>->), + -- * Controlling execution + compute, -- * Index construction and destruction indexHead, indexTail, toIndex, fromIndex, @@ -87,20 +79,25 @@ module Data.Array.Accelerate.Language ( -- * Flow-control acond, awhile, - cond, while, + cond, select, while, + Assert(..), assertBounds, + + -- * Utilities for bounds and shape checks + inbounds, inboundsOf, shapeEq, -- * Array operations with a scalar result - (!), (!!), shape, size, shapeSize, + (!), (!!), shape, size, shapeSize, null, -- * Numeric functions subtract, even, odd, gcd, lcm, (^), (^^), -- * Conversions - ord, chr, boolToInt, bitcast, + ord, chr, boolToInt, bitcast, isJust, ) where import Data.Array.Accelerate.AST ( PrimFun(..) ) +import Data.Array.Accelerate.AST.Idx import Data.Array.Accelerate.Pattern import Data.Array.Accelerate.Representation.Array ( ArrayR(..) ) import Data.Array.Accelerate.Representation.Shape ( ShapeR(..) ) @@ -118,12 +115,16 @@ import Data.Array.Accelerate.Classes.Fractional import Data.Array.Accelerate.Classes.Integral import Data.Array.Accelerate.Classes.Num import Data.Array.Accelerate.Classes.Ord +import Data.Text ( Text ) +import GHC.Stack +import Data.String ( fromString ) -import Prelude ( ($), (.), Maybe(..), Char ) +import Prelude ( ($), (.), (<>), Maybe(..), Char ) #if __GLASGOW_HASKELL__ >= 904 import Data.Type.Equality #endif +import Prelude () -- $setup -- >>> :seti -XFlexibleContexts @@ -301,11 +302,11 @@ replicate = Acc $$ applyAcc (Replicate $ sliceIndex @slix) -- generate :: forall sh a. - (Shape sh, Elt a) + (HasCallStack, Shape sh, Elt a) => Exp sh -> (Exp sh -> Exp a) -> Acc (Array sh a) -generate = Acc $$ applyAcc (Generate $ arrayR @sh @a) +generate sh fun = Acc $ applyAcc (Generate $ arrayR @sh @a) (assertBounds "Shape in generate should be non-negative" isNonNegative sh) fun -- Shape manipulation -- ------------------ @@ -321,11 +322,11 @@ generate = Acc $$ applyAcc (Generate $ arrayR @sh @a) -- reshape :: forall sh sh' e. - (Shape sh, Shape sh', Elt e) + (HasCallStack, Shape sh, Shape sh', Elt e) => Exp sh -> Acc (Array sh' e) -> Acc (Array sh e) -reshape = Acc $$ applyAcc (Reshape $ shapeR @sh) +reshape sh as = Acc $ applyAcc (Reshape $ shapeR @sh) (assertBounds "Reshape should preserve the size" (\s -> shapeSize s == size as) sh) as -- Extraction of sub-arrays -- ------------------------ @@ -449,6 +450,9 @@ map = Acc $$ applyAcc (Map (eltR @a) (eltR @b)) -- 16, 18, 20, 22, 24, -- 31, 33, 35, 37, 39] -- +-- If the input arrays are expected to have the same shape, use +-- 'zipWithChecked'. +-- zipWith :: forall sh a b c. (Shape sh, Elt a, Elt b, Elt c) => (Exp a -> Exp b -> Exp c) @@ -539,11 +543,13 @@ fold f (Exp x) = Acc . applyAcc (Fold (eltR @a) (unExpBinaryFunction f) (Just x) -- efficient parallel implementation, but does not need to be commutative. -- fold1 :: forall sh a. - (Shape sh, Elt a) + (HasCallStack, Shape sh, Elt a) => (Exp a -> Exp a -> Exp a) -> Acc (Array (sh:.Int) a) -> Acc (Array sh a) -fold1 f = Acc . applyAcc (Fold (eltR @a) (unExpBinaryFunction f) Nothing) +fold1 f = + Acc . applyAcc (Fold (eltR @a) (unExpBinaryFunction f) Nothing) + . assertBounds "Input of fold1 should be non-empty" (\xs -> let Exp sh = shape xs in innerNonEmpty (shapeR @sh) sh) -- | Segmented reduction along the innermost dimension of an array. The -- segment descriptor specifies the starting index (offset) along the @@ -567,6 +573,7 @@ foldSeg' -> Acc (Segments i) -> Acc (Array (sh:.Int) a) foldSeg' f (Exp x) = Acc $$ applyAcc (FoldSeg (integralType @i) (eltR @a) (unExpBinaryFunction f) (Just x)) +-- TODO: Add assertion that segment offsets are monotone increasing (if it doesn't add too much overhead) -- | Variant of 'foldSeg'' that requires /all/ segments of the reduced -- array to be non-empty, and doesn't need a default value. The segment @@ -583,6 +590,7 @@ fold1Seg' -> Acc (Segments i) -> Acc (Array (sh:.Int) a) fold1Seg' f = Acc $$ applyAcc (FoldSeg (integralType @i) (eltR @a) (unExpBinaryFunction f) Nothing) +-- TODO: Add assertion that segment offsets are *strictly* monotone increasing (if it doesn't add too much overhead) -- Scan functions -- -------------- @@ -694,109 +702,26 @@ scanr1 :: forall sh a. -> Acc (Array (sh:.Int) a) scanr1 f (Acc a) = Acc $ SmartAcc $ Scan RightToLeft (eltR @a) (unExpBinaryFunction f) Nothing a --- Permutations --- ------------ - --- | Generalised forward permutation operation (array scatter). --- --- Forward permutation specified by a function mapping indices from the source --- array to indices in the result array. The result array is initialised with --- the given defaults and any further values that are permuted into the result --- array are added to the current value using the given combination function. --- --- The combination function must be /associative/ and /commutative/. --- Elements for which the permutation function returns 'Nothing' are --- dropped. --- --- The combination function is given the new value being permuted as its first --- argument, and the current value of the array as its second. --- --- For example, we can use 'permute' to compute the occurrence count (histogram) --- for an array of values in the range @[0,10)@: --- --- >>> :{ --- let histogram :: Acc (Vector Int) -> Acc (Vector Int) --- histogram xs = --- let zeros = fill (constant (Z:.10)) 0 --- ones = fill (shape xs) 1 --- in --- permute (+) zeros (\ix -> Just_ (I1 (xs!ix))) ones --- :} --- --- >>> let xs = fromList (Z :. 20) [0,0,1,2,1,1,2,4,8,3,4,9,8,3,2,5,5,3,1,2] :: Vector Int --- >>> run $ histogram (use xs) --- Vector (Z :. 10) [2,4,4,3,2,2,0,0,2,1] --- --- As a second example, note that the dimensionality of the source and --- destination arrays can differ. In this way, we can use 'permute' to create an --- identity matrix by overwriting elements along the diagonal: --- --- >>> :{ --- let identity :: Num a => Exp Int -> Acc (Matrix a) --- identity n = --- let zeros = fill (I2 n n) 0 --- ones = fill (I1 n) 1 --- in --- permute const zeros (\(I1 i) -> Just_ (I2 i i)) ones --- :} --- --- >>> run $ identity 5 :: Matrix Int --- Matrix (Z :. 5 :. 5) --- [ 1, 0, 0, 0, 0, --- 0, 1, 0, 0, 0, --- 0, 0, 1, 0, 0, --- 0, 0, 0, 1, 0, --- 0, 0, 0, 0, 1] --- --- [/Note:/] --- --- Regarding array fusion: --- --- 1. The 'permute' operation will always be evaluated; it can not be fused --- into a later step. --- --- 2. Since the index permutation function might not cover all positions in --- the output array (the function is not surjective), the array of default --- values must be evaluated. However, other operations may fuse into this. --- --- 3. The array of source values can fuse into the permutation operation. --- --- 4. If the array of default values is only used once, it will be updated --- in-place. This behaviour can be disabled this with @-fno-inplace@. --- --- Regarding the defaults array: --- --- If you are sure that the default values are not necessary---they are not used --- by the combination function and every element will be overwritten---a default --- array created by 'Data.Array.Accelerate.Prelude.fill'ing with the value --- 'Data.Array.Accelerate.Unsafe.undef' will give you a new uninitialised array. --- --- Regarding the combination function: --- --- The function 'const' can be used to replace elements of the defaults --- array with the new values. If the permutation function maps multiple --- values to the same location in the results array (the function is not --- injective) then this operation is non-deterministic. --- --- Since Accelerate uses an unzipped struct-of-array representation, where --- the individual components of product types (for example, pairs) are --- stored in separate arrays, storing values of product type requires --- multiple store instructions. --- --- Accelerate prior to version 1.3.0.0 performs this operation atomically, --- to ensure that the stored values are always consistent (each component --- of the product type is written by the same thread). Later versions relax --- this restriction, but this behaviour can be disabled with --- @-fno-fast-permute-const@. --- -permute - :: forall sh sh' a. (Shape sh, Shape sh', Elt a) +permute' + :: forall sh sh' a. (HasCallStack, Shape sh, Shape sh', Elt a) => (Exp a -> Exp a -> Exp a) -- ^ combination function -> Acc (Array sh' a) -- ^ array of default values - -> (Exp sh -> Exp (Maybe sh')) -- ^ index permutation function - -> Acc (Array sh a) -- ^ array of source values to be permuted + -> Acc (Array sh (Maybe (sh', a))) -- ^ array of source values to be permuted, alongside their target index -> Acc (Array sh' a) -permute = Acc $$$$ applyAcc (Permute $ arrayR @sh @a) +permute' f def src = Acc $ applyAcc + (Permute (arrayR @sh @a) $ Just $ unExpBinaryFunction f) + def + $ map (assertBounds "Write indices in permute should be in bounds of the defaults array" (\s@(Exp s') -> not (isJust s) ||! inboundsOf def (Exp $ prj1 $ prj0 $ prj0 s'))) src + +permuteUnique' + :: forall sh sh' a. (HasCallStack, Shape sh, Shape sh', Elt a) + => Acc (Array sh' a) -- ^ array of default values + -> Acc (Array sh (Maybe (sh', a))) -- ^ array of source values to be permuted, alongside their target index + -> Acc (Array sh' a) +permuteUnique' def src = Acc $ applyAcc + (Permute (arrayR @sh @a) Nothing) + def + $ map (assertBounds "Write indices in permute should be in bounds of the defaults array" (\s@(Exp s') -> not (isJust s) ||! inboundsOf def (Exp $ prj1 $ prj0 $ prj0 s'))) src -- | Generalised backward permutation operation (array gather). -- @@ -842,12 +767,15 @@ permute = Acc $$$$ applyAcc (Permute $ arrayR @sh @a) -- 9, 19, 29, 39, 49] -- backpermute - :: forall sh sh' a. (Shape sh, Shape sh', Elt a) + :: forall sh sh' a. (HasCallStack, Shape sh, Shape sh', Elt a) => Exp sh' -- ^ shape of the result array -> (Exp sh' -> Exp sh) -- ^ index permutation function -> Acc (Array sh a) -- ^ source array -> Acc (Array sh' a) -backpermute = Acc $$$ applyAcc (Backpermute $ shapeR @sh') +backpermute sz f as = + Acc $ applyAcc (Backpermute $ shapeR @sh') (assertBounds "Shape of backpermute should be non-negative" isNonNegative sz) f' as + where + f' ix = assertBounds "Indices in backpermute should be in bounds of the source array" (inboundsOf as) $ f ix -- Stencil operations -- ------------------ @@ -961,6 +889,9 @@ stencil -> Acc (Array sh a) -- ^ source array -> Acc (Array sh b) -- ^ destination array stencil f (Boundary b) (Acc a) +-- TODO: Do boundary conditions put requirements on the input size? +-- We should add an 'assert' to check whether the size is large enough for the +-- chosen boundary condition. = Acc $ SmartAcc $ Stencil (stencilR @sh @a @stencil) (eltR @b) @@ -1223,24 +1154,44 @@ foreignExp foreignExp asm f (Exp x) = mkExp $ Foreign (eltR @y) asm (unExpFunction f) x --- Composition of array computations --- --------------------------------- +-- Controlling execution +-- --------------------- --- | Pipelining of two array computations. The first argument will be fully --- evaluated before being passed to the second computation. This can be used to --- prevent the argument being fused into the function, for example. +-- | Force an array expression to be evaluated, preventing it from fusing with +-- other operations. Forcing operations to be computed to memory, rather than +-- being fused into their consuming function, can sometimes improve performance. +-- For example, computing a matrix 'transpose' could provide better memory +-- locality for the subsequent operation. Preventing fusion to split large +-- operations into several simpler steps could also help by reducing register +-- pressure. +-- +-- Preventing fusion also means that the individual operations are available to +-- be executed concurrently with other kernels. In particular, consider using +-- this if you have a series of operations that are compute bound rather than +-- memory bound. -- --- Denotationally, we have +-- Here is the synthetic example: -- --- > (acc1 >-> acc2) arrs = let tmp = acc1 arrs --- > in tmp `seq` acc2 tmp +-- > loop :: Exp Int -> Exp Int +-- > loop ticks = +-- > let clockRate = 900000 -- kHz +-- > in while (\i -> i < clockRate * ticks) (+1) 0 +-- > +-- > test :: Acc (Vector Int) +-- > test = +-- > zip3 +-- > (compute $ map loop (use $ fromList (Z:.1) [10])) +-- > (compute $ map loop (use $ fromList (Z:.1) [10])) +-- > (compute $ map loop (use $ fromList (Z:.1) [10])) +-- > -- --- For an example use of this operation see the 'Data.Array.Accelerate.compute' --- function. +-- Without the use of 'compute', the operations are fused together and the three +-- long-running loops are executed sequentially in a single kernel. Instead, the +-- individual operations can now be executed concurrently, potentially reducing +-- overall runtime. -- -infixl 1 >-> -(>->) :: forall a b c. (Arrays a, Arrays b, Arrays c) => (Acc a -> Acc b) -> (Acc b -> Acc c) -> (Acc a -> Acc c) -(>->) = Acc $$$ applyAcc $ Pipe (arraysR @a) (arraysR @b) (arraysR @c) +compute :: Arrays a => Acc a -> Acc a +compute (Acc a) = Acc $ SmartAcc $ Manifest a -- Flow control constructs @@ -1281,16 +1232,18 @@ awhile f = Acc $$ applyAcc $ Awhile (arraysR @a) (unAccFunction g) -- array. -- toIndex - :: forall sh. Shape sh + :: forall sh. (HasCallStack, Shape sh) => Exp sh -- ^ extent of the array -> Exp sh -- ^ index to remap -> Exp Int -toIndex (Exp sh) (Exp ix) = mkExp $ ToIndex (shapeR @sh) sh ix +toIndex sh@(Exp sh') ix = mkExp $ ToIndex (shapeR @sh) sh' ix' + where Exp ix' = assertBounds "Index in toIndex should be in bounds of the given shape" (inbounds sh) ix -- | Inverse of 'toIndex' -- -fromIndex :: forall sh. Shape sh => Exp sh -> Exp Int -> Exp sh -fromIndex (Exp sh) (Exp e) = mkExp $ FromIndex (shapeR @sh) sh e +fromIndex :: forall sh. (HasCallStack, Shape sh) => Exp sh -> Exp Int -> Exp sh +fromIndex sh@(Exp sh') e = mkExp $ FromIndex (shapeR @sh) sh' e' + where Exp e' = assertBounds "Integer in fromIndex should be in bounds of the given shape" (\s -> s >= 0 &&! s < shapeSize sh) e -- | Intersection of two shapes -- @@ -1333,6 +1286,15 @@ cond :: Elt t -> Exp t cond (Exp c) (Exp x) (Exp y) = mkExp $ Cond (mkCoerce' c) x y +-- | A scalar-level conditional expression +-- that chooses one value based on the condition without branching. +select :: (Elt t) + => Exp Bool -- ^ condition + -> Exp t -- ^ then-expression + -> Exp t -- ^ else-expression + -> Exp t +select (Exp c) (Exp x) (Exp y) = mkExp $ Select (mkCoerce' c) x y + -- | While construct. Continue to apply the given function, starting with the -- initial value, until the test function evaluates to 'False'. -- @@ -1346,6 +1308,104 @@ while c f (Exp e) = (mkCoerce' . unExp . c . Exp) (unExp . f . Exp) e +class Assert a where + -- | Verifies whether a predicate holds. The predicate is evaluated before + -- the result of this function is used, and the program will crash if the + -- predicate returns false. + -- + -- Note that the predicate is only guaranteed to run if the result of this + -- function is used. If the result is not used, it may (or may not) be + -- removed from the program. + -- + assert :: HasCallStack => (a -> Exp Bool) -> a -> a + + -- | Verifies whether a predicate holds. Similar to 'assert', but this + -- function also takes a message, which is printed when the assertion + -- fails. + -- + assertMessage :: HasCallStack => Text -> (a -> Exp Bool) -> a -> a + + -- | Informs the compiler that a certain property holds on the given value. + -- The compiler may use this information to optimize the program. + -- + assume :: (a -> Exp Bool) -> a -> a + +-- | Variant of 'assert' that only adds an assertion when bounds checks are +-- enabled via the flag bounds-checks. +-- +-- This function may be removed in the future and replaced by the regular +-- 'assert', if we decide to turn bounds checks on by default. +-- +assertBounds :: (HasCallStack, Assert a) => Text -> (a -> Exp Bool) -> a -> a +#ifdef ACCELERATE_BOUNDS_CHECKS +assertBounds = withFrozenCallStack assertMessage +#else +assertBounds _ _ a = a +#endif + +instance Elt t => Assert (Exp t) where + assert f (Exp e) = mkExp $ Assert (fromString (prettyCallStack callStack)) (mkCoerce' g) e + where Exp g = f $ Exp e + assertMessage msg f (Exp e) = mkExp $ Assert (msg <> "\n" <> fromString (prettyCallStack callStack)) (mkCoerce' g) e + where Exp g = f $ Exp e + assume f (Exp e) = mkExp $ Assume (mkCoerce' g) e + where Exp g = f $ Exp e + +instance Arrays t => Assert (Acc t) where + assert f (Acc a) = Acc $ SmartAcc $ Aassert (fromString (prettyCallStack callStack)) (mkCoerce' g) a + where Exp g = f $ Acc a + assertMessage msg f (Acc a) = Acc $ SmartAcc $ Aassert (msg <> "\n" <> fromString (prettyCallStack callStack)) (mkCoerce' g) a + where Exp g = f $ Acc a + assume f (Acc a) = Acc $ SmartAcc $ Aassume (mkCoerce' g) a + where Exp g = f $ Acc a + +-- Some conditions for assertions that we use for the built-in functions of Accelerate + +-- | Checks whether the innermost dimension is non-empty, +-- or if the other dimensions are empty. fold1 requires this: +-- It folds per row, so the rows must be non-empty. +-- However, if there are no rows, the size of rows is irrelevant, +-- hence we also check for that. +innerNonEmpty :: ShapeR sh -> SmartExp (sh, Int) -> Exp Bool +innerNonEmpty shr sh = isEmpty shr (prjTail sh) ||! (Exp (prj0 sh) :: Exp Int) > 0 + +isEmpty :: ShapeR sh -> SmartExp sh -> Exp Bool +isEmpty ShapeRz _ = False_ +isEmpty (ShapeRsnoc shr) sh = isEmpty shr (prjTail sh) ||! (Exp (prj0 sh) :: Exp Int) == 0 + +isNonNegative :: forall sh. Shape sh => Exp sh -> Exp Bool +isNonNegative (Exp sh) = isNonNegative' (shapeR @sh) sh + +isNonNegative' :: ShapeR sh -> SmartExp sh -> Exp Bool +isNonNegative' ShapeRz _ = True_ +isNonNegative' (ShapeRsnoc shr) sh = isNonNegative' shr (prjTail sh) &&! (Exp (prj0 sh) :: Exp Int) >= 0 + +-- | Given the size (of an array), checks whether an index is in bounds. +-- +inbounds :: forall sh. Shape sh => Exp sh -> Exp sh -> Exp Bool +inbounds (Exp shape') (Exp index) = go (shapeR @sh) shape' index + where + go :: ShapeR s -> SmartExp s -> SmartExp s -> Exp Bool + go ShapeRz _ _ = True_ + go (ShapeRsnoc shr) sh ix = + go shr (prjTail sh) (prjTail ix) &&! 0 <= (Exp (prj0 ix) :: Exp Int) &&! (Exp (prj0 ix) :: Exp Int) < Exp (prj0 sh) + +-- | Given an array, checks whether an index is in bounds of that array. +-- +inboundsOf :: (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh -> Exp Bool +inboundsOf = inbounds . shape + +-- | Checks whether two shapes are equal. +-- This function is equal to (==), but has a Shape constraint instead of an Eq +-- constraint. +-- +shapeEq :: forall sh. Shape sh => Exp sh -> Exp sh -> Exp Bool +shapeEq (Exp sh1) (Exp sh2) = shapeEq' (shapeR @sh) sh1 sh2 + +shapeEq' :: ShapeR sh -> SmartExp sh -> SmartExp sh -> Exp Bool +shapeEq' ShapeRz _ _ = True_ +shapeEq' (ShapeRsnoc shr) sh1 sh2 = + shapeEq' shr (prjTail sh1) (prjTail sh2) &&! (Exp (prj0 sh1) :: Exp Int) == (Exp (prj0 sh2) :: Exp Int) -- Array operations with a scalar result -- ------------------------------------- @@ -1366,8 +1426,9 @@ while c f (Exp e) = -- 12 -- infixl 9 ! -(!) :: forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh -> Exp e -Acc a ! Exp ix = mkExp $ Index (eltR @e) a ix +(!) :: forall sh e. (HasCallStack, Shape sh, Elt e) => Acc (Array sh e) -> Exp sh -> Exp e +a@(Acc a') ! ix = mkExp $ Index (eltR @e) a' ix' + where Exp ix' = assertBounds "Index in (!) should be in bounds of the array" (inboundsOf a) ix -- | Extract the value from an array at the specified linear index. -- Multidimensional arrays in Accelerate are stored in row-major order with @@ -1386,8 +1447,9 @@ Acc a ! Exp ix = mkExp $ Index (eltR @e) a ix -- 12 -- infixl 9 !! -(!!) :: forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp Int -> Exp e -Acc a !! Exp ix = mkExp $ LinearIndex (eltR @e) a ix +(!!) :: forall sh e. (HasCallStack, Shape sh, Elt e) => Acc (Array sh e) -> Exp Int -> Exp e +a@(Acc a') !! ix = mkExp $ LinearIndex (eltR @e) a' ix' + where Exp ix' = assertBounds "Index in (!!) should be in bounds of the array" (\i -> 0 <= i &&! i < size a) ix -- | Extract the shape (extent) of an array. -- @@ -1404,6 +1466,12 @@ size = shapeSize . shape shapeSize :: forall sh. Shape sh => Exp sh -> Exp Int shapeSize (Exp sh) = mkExp $ ShapeSize (shapeR @sh) sh +-- | Test whether an array is empty. +-- +null :: forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp Bool +null arr = isEmpty (shapeR @sh) sh + where + Exp sh = shape arr -- Numeric functions -- ----------------- @@ -1515,3 +1583,10 @@ bitcast -> Exp b bitcast = mkBitcast +-- | Returns 'True' if the argument is of the form @Just _@ +-- +isJust :: Elt a => Exp (Maybe a) -> Exp Bool +isJust (Exp x) = Exp $ SmartExp $ (SmartExp $ Prj PairIdxLeft x) `Pair` SmartExp Nil + -- TLM: This is a sneaky hack because we know that the tag bits for Just + -- and True are identical. + diff --git a/src/Data/Array/Accelerate/Lifetime.hs b/src/Data/Array/Accelerate/Lifetime.hs index 7a9dacf37..9da30b213 100644 --- a/src/Data/Array/Accelerate/Lifetime.hs +++ b/src/Data/Array/Accelerate/Lifetime.hs @@ -1,5 +1,4 @@ {-# LANGUAGE MagicHash #-} -{-# LANGUAGE TupleSections #-} {-# LANGUAGE UnboxedTuples #-} {-# OPTIONS_HADDOCK hide #-} -- | @@ -24,7 +23,7 @@ module Data.Array.Accelerate.Lifetime ( import Data.Function ( on ) import Data.IORef ( mkWeakIORef, atomicModifyIORef' ) -import Prelude +import Control.Monad.IO.Class import GHC.Base ( touch#, IO(..)) import GHC.IORef ( IORef(.. ), newIORef ) @@ -73,10 +72,10 @@ newLifetime a = do -- the function, either by returning it or by lazy IO. -- {-# INLINE withLifetime #-} -withLifetime :: Lifetime a -> (a -> IO b) -> IO b +withLifetime :: MonadIO m => Lifetime a -> (a -> m b) -> m b withLifetime (Lifetime ref _ a) f = do r <- f a - touchIORef ref + liftIO (touchIORef ref) return r -- | Ensure that the lifetime is alive at the given place in a sequence of IO diff --git a/src/Data/Array/Accelerate/Lift.hs b/src/Data/Array/Accelerate/Lift.hs index 27482a86f..434c94161 100644 --- a/src/Data/Array/Accelerate/Lift.hs +++ b/src/Data/Array/Accelerate/Lift.hs @@ -1,14 +1,6 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeFamilies #-} #if __GLASGOW_HASKELL__ >= 806 {-# LANGUAGE UndecidableInstances #-} #endif diff --git a/src/Data/Array/Accelerate/Orphans.hs b/src/Data/Array/Accelerate/Orphans.hs index 39a43d493..d63fa68c4 100644 --- a/src/Data/Array/Accelerate/Orphans.hs +++ b/src/Data/Array/Accelerate/Orphans.hs @@ -1,10 +1,5 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE DeriveGeneric #-} -{-# LANGUAGE GeneralizedNewtypeDeriving #-} -{-# LANGUAGE StandaloneDeriving #-} -{-# LANGUAGE UnboxedTuples #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# LANGUAGE CPP #-} +{-# OPTIONS_GHC -Wno-orphans #-} -- | -- Module : Data.Array.Accelerate.Orphans -- Copyright : [2008..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Pattern.hs b/src/Data/Array/Accelerate/Pattern.hs index 42d1f5ab7..08bc3d709 100644 --- a/src/Data/Array/Accelerate/Pattern.hs +++ b/src/Data/Array/Accelerate/Pattern.hs @@ -1,18 +1,8 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ParallelListComp #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE UndecidableInstances #-} -{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE ParallelListComp #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE ViewPatterns #-} -- | -- Module : Data.Array.Accelerate.Pattern -- Copyright : [2018..2020] The Accelerate Team @@ -26,6 +16,7 @@ module Data.Array.Accelerate.Pattern ( pattern Pattern, +#ifndef __GHCIDE__ pattern T2, pattern T3, pattern T4, pattern T5, pattern T6, pattern T7, pattern T8, pattern T9, pattern T10, pattern T11, pattern T12, pattern T13, pattern T14, pattern T15, pattern T16, @@ -35,7 +26,7 @@ module Data.Array.Accelerate.Pattern ( pattern I5, pattern I6, pattern I7, pattern I8, pattern I9, pattern V2, pattern V3, pattern V4, pattern V8, pattern V16, - +#endif ) where import Data.Array.Accelerate.AST.Idx @@ -111,6 +102,7 @@ instance (Elt a, Elt b) => IsPattern Exp (a :. b) (Exp a :. Exp b) where matcher (Exp t) = Exp (SmartExp $ Prj PairIdxLeft t) :. Exp (SmartExp $ Prj PairIdxRight t) +#ifndef __GHCIDE__ -- IsPattern instances for up to 16-tuples (Acc and Exp). TH takes care of -- the (unremarkable) boilerplate for us. -- @@ -229,7 +221,6 @@ runQ $ do vs <- mapM mkVecPattern [2,3,4,8,16] return $ concat (es ++ as ++ vs) - -- | Specialised pattern synonyms for tuples, which may be more convenient to -- use than 'Data.Array.Accelerate.Lift.lift' and -- 'Data.Array.Accelerate.Lift.unlift'. For example, to construct a pair: @@ -306,3 +297,4 @@ runQ $ do vs <- mapM mkV [2,3,4,8,16] return $ concat (ts ++ is ++ vs) +#endif \ No newline at end of file diff --git a/src/Data/Array/Accelerate/Pattern/Bool.hs b/src/Data/Array/Accelerate/Pattern/Bool.hs index d968aaf34..2313f59bb 100644 --- a/src/Data/Array/Accelerate/Pattern/Bool.hs +++ b/src/Data/Array/Accelerate/Pattern/Bool.hs @@ -1,9 +1,6 @@ -{-# LANGUAGE GADTs #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE ViewPatterns #-} -- | -- Module : Data.Array.Accelerate.Pattern.Bool -- Copyright : [2018..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Pattern/Either.hs b/src/Data/Array/Accelerate/Pattern/Either.hs index 67c7b3a3f..e9087b2d7 100644 --- a/src/Data/Array/Accelerate/Pattern/Either.hs +++ b/src/Data/Array/Accelerate/Pattern/Either.hs @@ -1,9 +1,6 @@ -{-# LANGUAGE GADTs #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE ViewPatterns #-} -- | -- Module : Data.Array.Accelerate.Pattern.Either -- Copyright : [2018..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Pattern/Maybe.hs b/src/Data/Array/Accelerate/Pattern/Maybe.hs index 67e341d64..debb53d33 100644 --- a/src/Data/Array/Accelerate/Pattern/Maybe.hs +++ b/src/Data/Array/Accelerate/Pattern/Maybe.hs @@ -1,9 +1,6 @@ -{-# LANGUAGE GADTs #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE ViewPatterns #-} -- | -- Module : Data.Array.Accelerate.Pattern.Maybe -- Copyright : [2018..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Pattern/Ordering.hs b/src/Data/Array/Accelerate/Pattern/Ordering.hs index 2407cf9e9..5c7952585 100644 --- a/src/Data/Array/Accelerate/Pattern/Ordering.hs +++ b/src/Data/Array/Accelerate/Pattern/Ordering.hs @@ -1,9 +1,6 @@ -{-# LANGUAGE GADTs #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE ViewPatterns #-} -- | -- Module : Data.Array.Accelerate.Pattern.Ordering -- Copyright : [2018..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Pattern/TH.hs b/src/Data/Array/Accelerate/Pattern/TH.hs index b139a14f6..ee2e39af9 100644 --- a/src/Data/Array/Accelerate/Pattern/TH.hs +++ b/src/Data/Array/Accelerate/Pattern/TH.hs @@ -1,6 +1,5 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE TemplateHaskellQuotes #-} -- | -- Module : Data.Array.Accelerate.Pattern.TH -- Copyright : [2018..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Prelude.hs b/src/Data/Array/Accelerate/Prelude.hs index 2df0ce214..a1e061f55 100644 --- a/src/Data/Array/Accelerate/Prelude.hs +++ b/src/Data/Array/Accelerate/Prelude.hs @@ -1,18 +1,7 @@ -{-# LANGUAGE AllowAmbiguousTypes #-} -{-# LANGUAGE CPP #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE PatternGuards #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE RebindableSyntax #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} -{-# OPTIONS_GHC -fno-warn-incomplete-patterns #-} -- pattern synonyms +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE RebindableSyntax #-} +{-# LANGUAGE TypeFamilies #-} -- | -- Module : Data.Array.Accelerate.Prelude -- Copyright : [2009..2020] The Accelerate Team @@ -27,6 +16,9 @@ -- module Data.Array.Accelerate.Prelude ( + -- * Permutations + permute, + permuteUnique, -- * Element-wise operations indexed, @@ -34,8 +26,14 @@ module Data.Array.Accelerate.Prelude ( -- * Zipping zipWith3, zipWith4, zipWith5, zipWith6, zipWith7, zipWith8, zipWith9, + zipWithChecked, zipWithChecked3, zipWithChecked4, zipWithChecked5, + zipWithChecked6, zipWithChecked7, zipWithChecked8, zipWithChecked9, izipWith, izipWith3, izipWith4, izipWith5, izipWith6, izipWith7, izipWith8, izipWith9, + izipWithChecked, izipWithChecked3, izipWithChecked4, izipWithChecked5, + izipWithChecked6, izipWithChecked7, izipWithChecked8, izipWithChecked9, zip, zip3, zip4, zip5, zip6, zip7, zip8, zip9, + zipChecked, zipChecked3, zipChecked4, zipChecked5, + zipChecked6, zipChecked7, zipChecked8, zipChecked9, -- * Unzipping unzip, unzip3, unzip4, unzip5, unzip6, unzip7, unzip8, unzip9, @@ -79,14 +77,14 @@ module Data.Array.Accelerate.Prelude ( init, tail, take, drop, slit, initOn, tailOn, takeOn, dropOn, slitOn, - -- * Controlling execution - compute, + -- * Composition of array computations + (>->), -- * Flow control IfThenElse(..), -- ** Array-level - (?|), + (?|), awhileSpeculative, awhileAndOne, -- ** Expression-level (?), match, @@ -108,7 +106,7 @@ module Data.Array.Accelerate.Prelude ( index0, index1, unindex1, index2, unindex2, index3, unindex3, -- * Array operations with a scalar result - the, null, length, + the, length, -- * Irregular data-parallelism expand, @@ -124,9 +122,9 @@ import Data.Array.Accelerate.Lift import Data.Array.Accelerate.Pattern import Data.Array.Accelerate.Pattern.Maybe import Data.Array.Accelerate.Smart -import Data.Array.Accelerate.Sugar.Array ( Arrays, Array, Scalar, Vector, Segments, fromList ) +import Data.Array.Accelerate.Sugar.Array ( Arrays, Array, Scalar, Vector, Segments ) import Data.Array.Accelerate.Sugar.Elt -import Data.Array.Accelerate.Sugar.Shape ( Shape, Slice, Z(..), (:.)(..), All(..), DIM1, DIM2, empty ) +import Data.Array.Accelerate.Sugar.Shape ( Shape, Slice, Z(..), (:.)(..), All(..), DIM1, DIM2 ) import Data.Array.Accelerate.Type import Data.Array.Accelerate.Classes.Eq @@ -139,10 +137,15 @@ import Data.Array.Accelerate.Data.Bits import Lens.Micro ( Lens', (&), (^.), (.~), (+~), (-~), lens, over ) import Prelude ( (.), ($), Maybe(..), const, id, flip ) +import Data.String ( fromString ) +import GHC.Stack #if __GLASGOW_HASKELL__ >= 904 import Data.Type.Equality #endif +#ifdef ACCELERATE_INTERNAL_CHECKS +import Data.Typeable +#endif -- $setup -- >>> :seti -XFlexibleContexts @@ -153,6 +156,123 @@ import Data.Type.Equality -- runExp e = indexArray (run (unit e)) Z -- :} +-- Permutations +-- ------------ + +-- | Generalised forward permutation operation (array scatter). +-- +-- Forward permutation specified by a function mapping indices from the source +-- array to indices in the result array. The result array is initialised with +-- the given defaults and any further values that are permuted into the result +-- array are added to the current value using the given combination function. +-- +-- The combination function must be /associative/ and /commutative/. +-- Elements for which the permutation function returns 'Nothing' are +-- dropped. +-- +-- The combination function is given the new value being permuted as its first +-- argument, and the current value of the array as its second. +-- +-- For example, we can use 'permute' to compute the occurrence count (histogram) +-- for an array of values in the range @[0,10)@: +-- +-- >>> :{ +-- let histogram :: Acc (Vector Int) -> Acc (Vector Int) +-- histogram xs = +-- let zeros = fill (constant (Z:.10)) 0 +-- ones = fill (shape xs) 1 +-- in +-- permute (+) zeros (\ix -> Just_ (I1 (xs!ix))) ones +-- :} +-- +-- >>> let xs = fromList (Z :. 20) [0,0,1,2,1,1,2,4,8,3,4,9,8,3,2,5,5,3,1,2] :: Vector Int +-- >>> run $ histogram (use xs) +-- Vector (Z :. 10) [2,4,4,3,2,2,0,0,2,1] +-- +-- As a second example, note that the dimensionality of the source and +-- destination arrays can differ. In this way, we can use 'permute' to create an +-- identity matrix by overwriting elements along the diagonal: +-- +-- >>> :{ +-- let identity :: Num a => Exp Int -> Acc (Matrix a) +-- identity n = +-- let zeros = fill (I2 n n) 0 +-- ones = fill (I1 n) 1 +-- in +-- permute const zeros (\(I1 i) -> Just_ (I2 i i)) ones +-- :} +-- +-- >>> run $ identity 5 :: Matrix Int +-- Matrix (Z :. 5 :. 5) +-- [ 1, 0, 0, 0, 0, +-- 0, 1, 0, 0, 0, +-- 0, 0, 1, 0, 0, +-- 0, 0, 0, 1, 0, +-- 0, 0, 0, 0, 1] +-- +-- [/Note:/] +-- +-- Regarding array fusion: +-- +-- 1. The 'permute' operation will always be evaluated; it can not be fused +-- into a later step. +-- +-- 2. Since the index permutation function might not cover all positions in +-- the output array (the function is not surjective), the array of default +-- values must be evaluated. However, other operations may fuse into this. +-- +-- 3. The array of source values can fuse into the permutation operation. +-- +-- 4. If the array of default values is only used once, it will be updated +-- in-place. This behaviour can be disabled this with @-fno-inplace@. +-- +-- Regarding the defaults array: +-- +-- If you are sure that the default values are not necessary---they are not used +-- by the combination function and every element will be overwritten---a default +-- array created by 'Data.Array.Accelerate.Prelude.fill'ing with the value +-- 'Data.Array.Accelerate.Unsafe.undef' will give you a new uninitialised array. +-- +-- Regarding the combination function: +-- +-- The function 'const' can be used to replace elements of the defaults +-- array with the new values. If the permutation function maps multiple +-- values to the same location in the results array (the function is not +-- injective) then this operation is non-deterministic. +-- +-- Since Accelerate uses an unzipped struct-of-array representation, where +-- the individual components of product types (for example, pairs) are +-- stored in separate arrays, storing values of product type requires +-- multiple store instructions. +-- +-- Accelerate prior to version 1.3.0.0 performs this operation atomically, +-- to ensure that the stored values are always consistent (each component +-- of the product type is written by the same thread). Later versions relax +-- this restriction, but this behaviour can be disabled with +-- @-fno-fast-permute-const@. +-- +permute + :: forall sh sh' a. (Shape sh, Shape sh', Elt a) + => (Exp a -> Exp a -> Exp a) -- ^ combination function + -> Acc (Array sh' a) -- ^ array of default values + -> (Exp sh -> Exp (Maybe sh')) -- ^ index permutation function + -> Acc (Array sh a) -- ^ array of source values to be permuted + -> Acc (Array sh' a) +permute c def perm src = permute' c def $ imap (\ix x -> perm ix & match \case + Nothing_ -> Nothing_ + Just_ ix' -> Just_ (T2 ix' x)) src + +permuteUnique + :: forall sh sh' a. (Shape sh, Shape sh', Elt a) + => Acc (Array sh' a) -- ^ array of default values + -> (Exp sh -> Exp (Maybe sh')) -- ^ index permutation function + -> Acc (Array sh a) -- ^ array of source values to be permuted + -> Acc (Array sh' a) +permuteUnique def perm src = permuteUnique' def $ imap (\ix x -> perm ix & match \case + Nothing_ -> Nothing_ + Just_ ix' -> Just_ (T2 ix' x)) src + + -- Element-wise operations -- ----------------------- @@ -193,6 +313,8 @@ zipWithInduction prev f as bs = prev (\(T2 a b) -> f a b) (zip as bs) -- | Zip three arrays with the given function, analogous to 'zipWith'. +-- If the input arrays are expected to have the same shape, use +-- 'zipWithChecked3'. -- zipWith3 :: (Shape sh, Elt a, Elt b, Elt c, Elt d) @@ -204,6 +326,8 @@ zipWith3 zipWith3 = zipWithInduction zipWith -- | Zip four arrays with the given function, analogous to 'zipWith'. +-- If the input arrays are expected to have the same shape, use +-- 'zipWithChecked4'. -- zipWith4 :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e) @@ -216,6 +340,8 @@ zipWith4 zipWith4 = zipWithInduction zipWith3 -- | Zip five arrays with the given function, analogous to 'zipWith'. +-- If the input arrays are expected to have the same shape, use +-- 'zipWithChecked5'. -- zipWith5 :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f) @@ -229,6 +355,8 @@ zipWith5 zipWith5 = zipWithInduction zipWith4 -- | Zip six arrays with the given function, analogous to 'zipWith'. +-- If the input arrays are expected to have the same shape, use +-- 'zipWithChecked6'. -- zipWith6 :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g) @@ -243,6 +371,8 @@ zipWith6 zipWith6 = zipWithInduction zipWith5 -- | Zip seven arrays with the given function, analogous to 'zipWith'. +-- If the input arrays are expected to have the same shape, use +-- 'zipWithChecked7'. -- zipWith7 :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h) @@ -258,6 +388,8 @@ zipWith7 zipWith7 = zipWithInduction zipWith6 -- | Zip eight arrays with the given function, analogous to 'zipWith'. +-- If the input arrays are expected to have the same shape, use +-- 'zipWithChecked8'. -- zipWith8 :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h, Elt i) @@ -274,6 +406,8 @@ zipWith8 zipWith8 = zipWithInduction zipWith7 -- | Zip nine arrays with the given function, analogous to 'zipWith'. +-- If the input arrays are expected to have the same shape, use +-- 'zipWithChecked9'. -- zipWith9 :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h, Elt i, Elt j) @@ -290,59 +424,208 @@ zipWith9 -> Acc (Array sh j) zipWith9 = zipWithInduction zipWith8 +-- | Checks whether the shapes of two arrays are equal. +-- +assertShapeEq + :: forall sh a b. (HasCallStack, Shape sh, Elt a, Elt b) + => Acc (Array sh a) -> Acc (Array sh b) -> Acc (Array sh b) +assertShapeEq a = withFrozenCallStack $ assertMessage "Shapes in checked zip should be equal" $ shapeEq (shape a) . shape --- | Used to define the izipWith functions on two or more arrays +-- | Zip two __equal-shaped__ arrays with the given function. In contrast to +-- 'zipWith', this function does check whether the input arrays have the same +-- shape. -- -izipWithInduction - :: (Shape sh, Elt a, Elt b) - => ((Exp sh -> Exp (a,b) -> rest) -> Acc (Array sh (a,b)) -> result) -- The zipWith function operating on one fewer array - -> (Exp sh -> Exp a -> Exp b -> rest) +zipWithChecked + :: (HasCallStack, Shape sh, Elt a, Elt b, Elt c) + => (Exp a -> Exp b -> Exp c) -> Acc (Array sh a) -> Acc (Array sh b) - -> result -izipWithInduction prev f as bs = prev (\ix (T2 a b) -> f ix a b) (zip as bs) + -> Acc (Array sh c) +zipWithChecked fun a b = zipWith fun a (assertShapeEq a b) +-- | Zip three __equal-shaped__ arrays with the given function. In contrast to +-- 'zipWith3', this function does check whether the input arrays have the same +-- shape. +-- +zipWithChecked3 + :: (HasCallStack, Shape sh, Elt a, Elt b, Elt c, Elt d) + => (Exp a -> Exp b -> Exp c -> Exp d) + -> Acc (Array sh a) + -> Acc (Array sh b) + -> Acc (Array sh c) + -> Acc (Array sh d) +zipWithChecked3 fun a b c = + zipWith3 fun a (assertShapeEq a b) (assertShapeEq a c) --- | Zip two arrays with a function that also takes the element index +-- | Zip four __equal-shaped__ arrays with the given function. In contrast to +-- 'zipWith4', this function does check whether the input arrays have the same +-- shape. -- -izipWith - :: (Shape sh, Elt a, Elt b, Elt c) +zipWithChecked4 + :: (HasCallStack, Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e) + => (Exp a -> Exp b -> Exp c -> Exp d -> Exp e) + -> Acc (Array sh a) + -> Acc (Array sh b) + -> Acc (Array sh c) + -> Acc (Array sh d) + -> Acc (Array sh e) +zipWithChecked4 fun a b c d = + zipWith4 fun a (assertShapeEq a b) (assertShapeEq a c) (assertShapeEq a d) + +-- | Zip five __equal-shaped__ arrays with the given function. In contrast to +-- 'zipWith5', this function does check whether the input arrays have the same +-- shape. +-- +zipWithChecked5 + :: (HasCallStack, Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f) + => (Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f) + -> Acc (Array sh a) + -> Acc (Array sh b) + -> Acc (Array sh c) + -> Acc (Array sh d) + -> Acc (Array sh e) + -> Acc (Array sh f) +zipWithChecked5 fun a b c d e = + zipWith5 fun a (assertShapeEq a b) (assertShapeEq a c) (assertShapeEq a d) + (assertShapeEq a e) + +-- | Zip six __equal-shaped__ arrays with the given function. In contrast to +-- 'zipWith6', this function does check whether the input arrays have the same +-- shape. +-- +zipWithChecked6 + :: (HasCallStack, Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g) + => (Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g) + -> Acc (Array sh a) + -> Acc (Array sh b) + -> Acc (Array sh c) + -> Acc (Array sh d) + -> Acc (Array sh e) + -> Acc (Array sh f) + -> Acc (Array sh g) +zipWithChecked6 fun a b c d e f = + zipWith6 fun a (assertShapeEq a b) (assertShapeEq a c) (assertShapeEq a d) + (assertShapeEq a e) (assertShapeEq a f) + +-- | Zip seven __equal-shaped__ arrays with the given function. In contrast to +-- 'zipWith7', this function does check whether the input arrays have the same +-- shape. +-- +zipWithChecked7 + :: (HasCallStack, Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h) + => (Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g -> Exp h) + -> Acc (Array sh a) + -> Acc (Array sh b) + -> Acc (Array sh c) + -> Acc (Array sh d) + -> Acc (Array sh e) + -> Acc (Array sh f) + -> Acc (Array sh g) + -> Acc (Array sh h) +zipWithChecked7 fun a b c d e f g = + zipWith7 fun a (assertShapeEq a b) (assertShapeEq a c) (assertShapeEq a d) + (assertShapeEq a e) (assertShapeEq a f) (assertShapeEq a g) + +-- | Zip eight __equal-shaped__ arrays with the given function. In contrast to +-- 'zipWith8', this function does check whether the input arrays have the same +-- shape. +-- +zipWithChecked8 + :: (HasCallStack, Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h, Elt i) + => (Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g -> Exp h -> Exp i) + -> Acc (Array sh a) + -> Acc (Array sh b) + -> Acc (Array sh c) + -> Acc (Array sh d) + -> Acc (Array sh e) + -> Acc (Array sh f) + -> Acc (Array sh g) + -> Acc (Array sh h) + -> Acc (Array sh i) +zipWithChecked8 fun a b c d e f g h = + zipWith8 fun a (assertShapeEq a b) (assertShapeEq a c) (assertShapeEq a d) + (assertShapeEq a e) (assertShapeEq a f) (assertShapeEq a g) + (assertShapeEq a h) + +-- | Zip nine __equal-shaped__ arrays with the given function. In contrast to +-- 'zipWith9', this function does check whether the input arrays have the same +-- shape. +-- +zipWithChecked9 + :: (HasCallStack, Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h, Elt i, Elt j) + => (Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g -> Exp h -> Exp i -> Exp j) + -> Acc (Array sh a) + -> Acc (Array sh b) + -> Acc (Array sh c) + -> Acc (Array sh d) + -> Acc (Array sh e) + -> Acc (Array sh f) + -> Acc (Array sh g) + -> Acc (Array sh h) + -> Acc (Array sh i) + -> Acc (Array sh j) +zipWithChecked9 fun a b c d e f g h i = + zipWith9 fun a (assertShapeEq a b) (assertShapeEq a c) (assertShapeEq a d) + (assertShapeEq a e) (assertShapeEq a f) (assertShapeEq a g) + (assertShapeEq a h) (assertShapeEq a i) + +-- | Variant of 'curry', without the lift type classes. Used for the +-- implementation of izipWith(Checked)(N) +-- +curry' :: (Elt x, Elt y) => (Exp x -> Exp y -> z) -> Exp (x, y) -> z +curry' f (T2 x y) = f x y + +-- | Zip two arrays with a function that also takes the element index. +-- The checked variant, 'izipWithChecked', checks that the input arrays +-- have the same shape. +-- +izipWith, izipWithChecked + :: (HasCallStack, Shape sh, Elt a, Elt b, Elt c) => (Exp sh -> Exp a -> Exp b -> Exp c) -> Acc (Array sh a) -> Acc (Array sh b) -> Acc (Array sh c) -izipWith = izipWithInduction imap +izipWith f = zipWith (curry' f) . indexed +izipWithChecked f = zipWithChecked (curry' f) . indexed -- | Zip three arrays with a function that also takes the element index, -- analogous to 'izipWith'. +-- The checked variant, 'izipWithChecked3', checks that the input arrays +-- have the same shape. -- -izipWith3 - :: (Shape sh, Elt a, Elt b, Elt c, Elt d) +izipWith3, izipWithChecked3 + :: (HasCallStack, Shape sh, Elt a, Elt b, Elt c, Elt d) => (Exp sh -> Exp a -> Exp b -> Exp c -> Exp d) -> Acc (Array sh a) -> Acc (Array sh b) -> Acc (Array sh c) -> Acc (Array sh d) -izipWith3 = izipWithInduction izipWith +izipWith3 f = zipWith3 (curry' f) . indexed +izipWithChecked3 f = zipWithChecked3 (curry' f) . indexed -- | Zip four arrays with the given function that also takes the element index, --- analogous to 'zipWith'. +-- analogous to 'izipWith'. +-- The checked variant, 'izipWithChecked4', checks that the input arrays +-- have the same shape. -- -izipWith4 - :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e) +izipWith4, izipWithChecked4 + :: (HasCallStack, Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e) => (Exp sh -> Exp a -> Exp b -> Exp c -> Exp d -> Exp e) -> Acc (Array sh a) -> Acc (Array sh b) -> Acc (Array sh c) -> Acc (Array sh d) -> Acc (Array sh e) -izipWith4 = izipWithInduction izipWith3 +izipWith4 f = zipWith4 (curry' f) . indexed +izipWithChecked4 f = zipWithChecked4 (curry' f) . indexed -- | Zip five arrays with the given function that also takes the element index, --- analogous to 'zipWith'. +-- analogous to 'izipWith'. +-- The checked variant, 'izipWithChecked5', checks that the input arrays +-- have the same shape. -- -izipWith5 - :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f) +izipWith5, izipWithChecked5 + :: (HasCallStack, Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f) => (Exp sh -> Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f) -> Acc (Array sh a) -> Acc (Array sh b) @@ -350,13 +633,16 @@ izipWith5 -> Acc (Array sh d) -> Acc (Array sh e) -> Acc (Array sh f) -izipWith5 = izipWithInduction izipWith4 +izipWith5 f = zipWith5 (curry' f) . indexed +izipWithChecked5 f = zipWithChecked5 (curry' f) . indexed -- | Zip six arrays with the given function that also takes the element index, --- analogous to 'zipWith'. +-- analogous to 'izipWith'. +-- The checked variant, 'izipWithChecked6', checks that the input arrays +-- have the same shape. -- -izipWith6 - :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g) +izipWith6, izipWithChecked6 + :: (HasCallStack, Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g) => (Exp sh -> Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g) -> Acc (Array sh a) -> Acc (Array sh b) @@ -365,13 +651,16 @@ izipWith6 -> Acc (Array sh e) -> Acc (Array sh f) -> Acc (Array sh g) -izipWith6 = izipWithInduction izipWith5 +izipWith6 f = zipWith6 (curry' f) . indexed +izipWithChecked6 f = zipWithChecked6 (curry' f) . indexed -- | Zip seven arrays with the given function that also takes the element --- index, analogous to 'zipWith'. +-- index, analogous to 'izipWith'. +-- The checked variant, 'izipWithChecked7', checks that the input arrays +-- have the same shape. -- -izipWith7 - :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h) +izipWith7, izipWithChecked7 + :: (HasCallStack, Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h) => (Exp sh -> Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g -> Exp h) -> Acc (Array sh a) -> Acc (Array sh b) @@ -381,13 +670,16 @@ izipWith7 -> Acc (Array sh f) -> Acc (Array sh g) -> Acc (Array sh h) -izipWith7 = izipWithInduction izipWith6 +izipWith7 f = zipWith7 (curry' f) . indexed +izipWithChecked7 f = zipWithChecked7 (curry' f) . indexed -- | Zip eight arrays with the given function that also takes the element --- index, analogous to 'zipWith'. +-- index, analogous to 'izipWith'. +-- The checked variant, 'izipWithChecked8', checks that the input arrays +-- have the same shape. -- -izipWith8 - :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h, Elt i) +izipWith8, izipWithChecked8 + :: (HasCallStack, Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h, Elt i) => (Exp sh -> Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g -> Exp h -> Exp i) -> Acc (Array sh a) -> Acc (Array sh b) @@ -398,13 +690,16 @@ izipWith8 -> Acc (Array sh g) -> Acc (Array sh h) -> Acc (Array sh i) -izipWith8 = izipWithInduction izipWith7 +izipWith8 f = zipWith8 (curry' f) . indexed +izipWithChecked8 f = zipWithChecked8 (curry' f) . indexed -- | Zip nine arrays with the given function that also takes the element index, --- analogous to 'zipWith'. +-- analogous to 'izipWith'. +-- The checked variant, 'izipWithChecked9', checks that the input arrays +-- have the same shape. -- -izipWith9 - :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h, Elt i, Elt j) +izipWith9, izipWithChecked9 + :: (HasCallStack, Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h, Elt i, Elt j) => (Exp sh -> Exp a -> Exp b -> Exp c -> Exp d -> Exp e -> Exp f -> Exp g -> Exp h -> Exp i -> Exp j) -> Acc (Array sh a) -> Acc (Array sh b) @@ -416,11 +711,14 @@ izipWith9 -> Acc (Array sh h) -> Acc (Array sh i) -> Acc (Array sh j) -izipWith9 = izipWithInduction izipWith8 +izipWith9 f = zipWith9 (curry' f) . indexed +izipWithChecked9 f = zipWithChecked9 (curry' f) . indexed -- | Combine the elements of two arrays pairwise. The shape of the result is the -- intersection of the two argument shapes. +-- 'zipChecked' is the checked variant of 'zip', and checks that the input +-- arrays have the same shape. -- -- >>> let m1 = fromList (Z:.5:.10) [0..] :: Matrix Int -- >>> let m2 = fromList (Z:.10:.5) [0..] :: Matrix Float @@ -432,95 +730,135 @@ izipWith9 = izipWithInduction izipWith8 -- (30,15.0), (31,16.0), (32,17.0), (33,18.0), (34,19.0), -- (40,20.0), (41,21.0), (42,22.0), (43,23.0), (44,24.0)] -- -zip :: (Shape sh, Elt a, Elt b) - => Acc (Array sh a) - -> Acc (Array sh b) - -> Acc (Array sh (a, b)) +zip + :: (Shape sh, Elt a, Elt b) + => Acc (Array sh a) + -> Acc (Array sh b) + -> Acc (Array sh (a, b)) zip = zipWith T2 --- | Take three arrays and return an array of triples, analogous to zip. --- -zip3 :: (Shape sh, Elt a, Elt b, Elt c) - => Acc (Array sh a) - -> Acc (Array sh b) - -> Acc (Array sh c) - -> Acc (Array sh (a, b, c)) +-- | Combine the elements of two __equal-sized__ arrays pairwise. In contrast +-- to 'zip', this function does check whether the input arrays have the same +-- shape. +-- +zipChecked + :: (HasCallStack, Shape sh, Elt a, Elt b) + => Acc (Array sh a) + -> Acc (Array sh b) + -> Acc (Array sh (a, b)) +zipChecked = zipWithChecked T2 + +-- | Take three arrays and return an array of triples, analogous to 'zip'. +-- 'zipChecked3' is the checked variant of 'zip3', and checks that the input +-- arrays have the same shape. +-- +zip3, zipChecked3 + :: (HasCallStack, Shape sh, Elt a, Elt b, Elt c) + => Acc (Array sh a) + -> Acc (Array sh b) + -> Acc (Array sh c) + -> Acc (Array sh (a, b, c)) zip3 = zipWith3 T3 - --- | Take four arrays and return an array of quadruples, analogous to zip. --- -zip4 :: (Shape sh, Elt a, Elt b, Elt c, Elt d) - => Acc (Array sh a) - -> Acc (Array sh b) - -> Acc (Array sh c) - -> Acc (Array sh d) - -> Acc (Array sh (a, b, c, d)) +zipChecked3 = zipWithChecked3 T3 + +-- | Take four arrays and return an array of quadruples, analogous to 'zip'. +-- 'zipChecked4' is the checked variant of 'zip4', and checks that the input +-- arrays have the same shape. +-- +zip4, zipChecked4 + :: (HasCallStack, Shape sh, Elt a, Elt b, Elt c, Elt d) + => Acc (Array sh a) + -> Acc (Array sh b) + -> Acc (Array sh c) + -> Acc (Array sh d) + -> Acc (Array sh (a, b, c, d)) zip4 = zipWith4 T4 - --- | Take five arrays and return an array of five-tuples, analogous to zip. --- -zip5 :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e) - => Acc (Array sh a) - -> Acc (Array sh b) - -> Acc (Array sh c) - -> Acc (Array sh d) - -> Acc (Array sh e) - -> Acc (Array sh (a, b, c, d, e)) +zipChecked4 = zipWithChecked4 T4 + +-- | Take five arrays and return an array of five-tuples, analogous to 'zip'. +-- 'zipChecked5' is the checked variant of 'zip5', and checks that the input +-- arrays have the same shape. +-- +zip5, zipChecked5 + :: (HasCallStack, Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e) + => Acc (Array sh a) + -> Acc (Array sh b) + -> Acc (Array sh c) + -> Acc (Array sh d) + -> Acc (Array sh e) + -> Acc (Array sh (a, b, c, d, e)) zip5 = zipWith5 T5 - --- | Take six arrays and return an array of six-tuples, analogous to zip. --- -zip6 :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f) - => Acc (Array sh a) - -> Acc (Array sh b) - -> Acc (Array sh c) - -> Acc (Array sh d) - -> Acc (Array sh e) - -> Acc (Array sh f) - -> Acc (Array sh (a, b, c, d, e, f)) +zipChecked5 = zipWithChecked5 T5 + +-- | Take six arrays and return an array of six-tuples, analogous to 'zip'. +-- 'zipChecked6' is the checked variant of 'zip6', and checks that the input +-- arrays have the same shape. +-- +zip6, zipChecked6 + :: (HasCallStack, Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f) + => Acc (Array sh a) + -> Acc (Array sh b) + -> Acc (Array sh c) + -> Acc (Array sh d) + -> Acc (Array sh e) + -> Acc (Array sh f) + -> Acc (Array sh (a, b, c, d, e, f)) zip6 = zipWith6 T6 - --- | Take seven arrays and return an array of seven-tuples, analogous to zip. --- -zip7 :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g) - => Acc (Array sh a) - -> Acc (Array sh b) - -> Acc (Array sh c) - -> Acc (Array sh d) - -> Acc (Array sh e) - -> Acc (Array sh f) - -> Acc (Array sh g) - -> Acc (Array sh (a, b, c, d, e, f, g)) +zipChecked6 = zipWithChecked6 T6 + +-- | Take seven arrays and return an array of seven-tuples, analogous to 'zip'. +-- 'zipChecked7' is the checked variant of 'zip7', and checks that the input +-- arrays have the same shape. +-- +zip7, zipChecked7 + :: (HasCallStack, Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g) + => Acc (Array sh a) + -> Acc (Array sh b) + -> Acc (Array sh c) + -> Acc (Array sh d) + -> Acc (Array sh e) + -> Acc (Array sh f) + -> Acc (Array sh g) + -> Acc (Array sh (a, b, c, d, e, f, g)) zip7 = zipWith7 T7 - --- | Take eight arrays and return an array of eight-tuples, analogous to zip. --- -zip8 :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h) - => Acc (Array sh a) - -> Acc (Array sh b) - -> Acc (Array sh c) - -> Acc (Array sh d) - -> Acc (Array sh e) - -> Acc (Array sh f) - -> Acc (Array sh g) - -> Acc (Array sh h) - -> Acc (Array sh (a, b, c, d, e, f, g, h)) +zipChecked7 = zipWithChecked7 T7 + +-- | Take eight arrays and return an array of eight-tuples, analogous to 'zip'. +-- 'zipChecked8' is the checked variant of 'zip8', and checks that the input +-- arrays have the same shape. +-- +zip8, zipChecked8 + :: (HasCallStack, Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h) + => Acc (Array sh a) + -> Acc (Array sh b) + -> Acc (Array sh c) + -> Acc (Array sh d) + -> Acc (Array sh e) + -> Acc (Array sh f) + -> Acc (Array sh g) + -> Acc (Array sh h) + -> Acc (Array sh (a, b, c, d, e, f, g, h)) zip8 = zipWith8 T8 - --- | Take nine arrays and return an array of nine-tuples, analogous to zip. --- -zip9 :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h, Elt i) - => Acc (Array sh a) - -> Acc (Array sh b) - -> Acc (Array sh c) - -> Acc (Array sh d) - -> Acc (Array sh e) - -> Acc (Array sh f) - -> Acc (Array sh g) - -> Acc (Array sh h) - -> Acc (Array sh i) - -> Acc (Array sh (a, b, c, d, e, f, g, h, i)) +zipChecked8 = zipWithChecked8 T8 + +-- | Take nine arrays and return an array of nine-tuples, analogous to 'zip'. +-- 'zipChecked9' is the checked variant of 'zip9', and checks that the input +-- arrays have the same shape. +-- +zip9, zipChecked9 + :: (HasCallStack, Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h, Elt i) + => Acc (Array sh a) + -> Acc (Array sh b) + -> Acc (Array sh c) + -> Acc (Array sh d) + -> Acc (Array sh e) + -> Acc (Array sh f) + -> Acc (Array sh g) + -> Acc (Array sh h) + -> Acc (Array sh i) + -> Acc (Array sh (a, b, c, d, e, f, g, h, i)) zip9 = zipWith9 T9 +zipChecked9 = zipWithChecked9 T9 -- | The converse of 'zip', but the shape of the two results is identical to the @@ -535,59 +873,90 @@ unzip arr = (map fst arr, map snd arr) -- | Take an array of triples and return three arrays, analogous to 'unzip'. -- -unzip3 :: (Shape sh, Elt a, Elt b, Elt c) +unzip3 :: forall sh a b c. (Shape sh, Elt a, Elt b, Elt c) => Acc (Array sh (a, b, c)) -> (Acc (Array sh a), Acc (Array sh b), Acc (Array sh c)) unzip3 xs = (map get1 xs, map get2 xs, map get3 xs) where + get1 :: Exp (a, b, c) -> Exp a get1 (T3 a _ _) = a + + get2 :: Exp (a, b, c) -> Exp b get2 (T3 _ b _) = b - get3 (T3 _ _ c) = c + get3 :: Exp (a, b, c) -> Exp c + get3 (T3 _ _ c) = c -- | Take an array of quadruples and return four arrays, analogous to 'unzip'. -- -unzip4 :: (Shape sh, Elt a, Elt b, Elt c, Elt d) +unzip4 :: forall sh a b c d. (Shape sh, Elt a, Elt b, Elt c, Elt d) => Acc (Array sh (a, b, c, d)) -> (Acc (Array sh a), Acc (Array sh b), Acc (Array sh c), Acc (Array sh d)) unzip4 xs = (map get1 xs, map get2 xs, map get3 xs, map get4 xs) where + get1 :: Exp (a, b, c, d) -> Exp a get1 (T4 a _ _ _) = a + + get2 :: Exp (a, b, c, d) -> Exp b get2 (T4 _ b _ _) = b + + get3 :: Exp (a, b, c, d) -> Exp c get3 (T4 _ _ c _) = c + + get4 :: Exp (a, b, c, d) -> Exp d get4 (T4 _ _ _ d) = d -- | Take an array of 5-tuples and return five arrays, analogous to 'unzip'. -- -unzip5 :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e) +unzip5 :: forall sh a b c d e. (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e) => Acc (Array sh (a, b, c, d, e)) -> (Acc (Array sh a), Acc (Array sh b), Acc (Array sh c), Acc (Array sh d), Acc (Array sh e)) unzip5 xs = (map get1 xs, map get2 xs, map get3 xs, map get4 xs, map get5 xs) where + get1 :: Exp (a, b, c, d, e) -> Exp a get1 (T5 a _ _ _ _) = a + + get2 :: Exp (a, b, c, d, e) -> Exp b get2 (T5 _ b _ _ _) = b + + get3 :: Exp (a, b, c, d, e) -> Exp c get3 (T5 _ _ c _ _) = c + + get4 :: Exp (a, b, c, d, e) -> Exp d get4 (T5 _ _ _ d _) = d + + get5 :: Exp (a, b, c, d, e) -> Exp e get5 (T5 _ _ _ _ e) = e -- | Take an array of 6-tuples and return six arrays, analogous to 'unzip'. -- -unzip6 :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f) +unzip6 :: forall sh a b c d e f. (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f) => Acc (Array sh (a, b, c, d, e, f)) -> ( Acc (Array sh a), Acc (Array sh b), Acc (Array sh c) , Acc (Array sh d), Acc (Array sh e), Acc (Array sh f)) unzip6 xs = (map get1 xs, map get2 xs, map get3 xs, map get4 xs, map get5 xs, map get6 xs) where + get1 :: Exp (a, b, c, d, e, f) -> Exp a get1 (T6 a _ _ _ _ _) = a + + get2 :: Exp (a, b, c, d, e, f) -> Exp b get2 (T6 _ b _ _ _ _) = b + + get3 :: Exp (a, b, c, d, e, f) -> Exp c get3 (T6 _ _ c _ _ _) = c + + get4 :: Exp (a, b, c, d, e, f) -> Exp d get4 (T6 _ _ _ d _ _) = d + + get5 :: Exp (a, b, c, d, e, f) -> Exp e get5 (T6 _ _ _ _ e _) = e + + get6 :: Exp (a, b, c, d, e, f) -> Exp f get6 (T6 _ _ _ _ _ f) = f -- | Take an array of 7-tuples and return seven arrays, analogous to 'unzip'. -- -unzip7 :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g) +unzip7 :: forall sh a b c d e f g. (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g) => Acc (Array sh (a, b, c, d, e, f, g)) -> ( Acc (Array sh a), Acc (Array sh b), Acc (Array sh c) , Acc (Array sh d), Acc (Array sh e), Acc (Array sh f) @@ -596,17 +965,30 @@ unzip7 xs = ( map get1 xs, map get2 xs, map get3 xs , map get4 xs, map get5 xs, map get6 xs , map get7 xs ) where + get1 :: Exp (a, b, c, d, e, f, g) -> Exp a get1 (T7 a _ _ _ _ _ _) = a + + get2 :: Exp (a, b, c, d, e, f, g) -> Exp b get2 (T7 _ b _ _ _ _ _) = b + + get3 :: Exp (a, b, c, d, e, f, g) -> Exp c get3 (T7 _ _ c _ _ _ _) = c + + get4 :: Exp (a, b, c, d, e, f, g) -> Exp d get4 (T7 _ _ _ d _ _ _) = d + + get5 :: Exp (a, b, c, d, e, f, g) -> Exp e get5 (T7 _ _ _ _ e _ _) = e + + get6 :: Exp (a, b, c, d, e, f, g) -> Exp f get6 (T7 _ _ _ _ _ f _) = f + + get7 :: Exp (a, b, c, d, e, f, g) -> Exp g get7 (T7 _ _ _ _ _ _ g) = g -- | Take an array of 8-tuples and return eight arrays, analogous to 'unzip'. -- -unzip8 :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h) +unzip8 :: forall sh a b c d e f g h. (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h) => Acc (Array sh (a, b, c, d, e, f, g, h)) -> ( Acc (Array sh a), Acc (Array sh b), Acc (Array sh c) , Acc (Array sh d), Acc (Array sh e), Acc (Array sh f) @@ -615,18 +997,33 @@ unzip8 xs = ( map get1 xs, map get2 xs, map get3 xs , map get4 xs, map get5 xs, map get6 xs , map get7 xs, map get8 xs ) where + get1 :: Exp (a, b, c, d, e, f, g, h) -> Exp a get1 (T8 a _ _ _ _ _ _ _) = a + + get2 :: Exp (a, b, c, d, e, f, g, h) -> Exp b get2 (T8 _ b _ _ _ _ _ _) = b + + get3 :: Exp (a, b, c, d, e, f, g, h) -> Exp c get3 (T8 _ _ c _ _ _ _ _) = c + + get4 :: Exp (a, b, c, d, e, f, g, h) -> Exp d get4 (T8 _ _ _ d _ _ _ _) = d + + get5 :: Exp (a, b, c, d, e, f, g, h) -> Exp e get5 (T8 _ _ _ _ e _ _ _) = e + + get6 :: Exp (a, b, c, d, e, f, g, h) -> Exp f get6 (T8 _ _ _ _ _ f _ _) = f + + get7 :: Exp (a, b, c, d, e, f, g, h) -> Exp g get7 (T8 _ _ _ _ _ _ g _) = g + + get8 :: Exp (a, b, c, d, e, f, g, h) -> Exp h get8 (T8 _ _ _ _ _ _ _ h) = h -- | Take an array of 9-tuples and return nine arrays, analogous to 'unzip'. -- -unzip9 :: (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h, Elt i) +unzip9 :: forall sh a b c d e f g h i. (Shape sh, Elt a, Elt b, Elt c, Elt d, Elt e, Elt f, Elt g, Elt h, Elt i) => Acc (Array sh (a, b, c, d, e, f, g, h, i)) -> ( Acc (Array sh a), Acc (Array sh b), Acc (Array sh c) , Acc (Array sh d), Acc (Array sh e), Acc (Array sh f) @@ -635,14 +1032,31 @@ unzip9 xs = ( map get1 xs, map get2 xs, map get3 xs , map get4 xs, map get5 xs, map get6 xs , map get7 xs, map get8 xs, map get9 xs ) where + get1 :: Exp (a, b, c, d, e, f, g, h, i) -> Exp a get1 (T9 a _ _ _ _ _ _ _ _) = a + + get2 :: Exp (a, b, c, d, e, f, g, h, i) -> Exp b get2 (T9 _ b _ _ _ _ _ _ _) = b + + get3 :: Exp (a, b, c, d, e, f, g, h, i) -> Exp c get3 (T9 _ _ c _ _ _ _ _ _) = c + + get4 :: Exp (a, b, c, d, e, f, g, h, i) -> Exp d get4 (T9 _ _ _ d _ _ _ _ _) = d + + get5 :: Exp (a, b, c, d, e, f, g, h, i) -> Exp e get5 (T9 _ _ _ _ e _ _ _ _) = e + + get6 :: Exp (a, b, c, d, e, f, g, h, i) -> Exp f get6 (T9 _ _ _ _ _ f _ _ _) = f + + get7 :: Exp (a, b, c, d, e, f, g, h, i) -> Exp g get7 (T9 _ _ _ _ _ _ g _ _) = g + + get8 :: Exp (a, b, c, d, e, f, g, h, i) -> Exp h get8 (T9 _ _ _ _ _ _ _ h _) = h + + get9 :: Exp (a, b, c, d, e, f, g, h, i) -> Exp i get9 (T9 _ _ _ _ _ _ _ _ i) = i @@ -662,7 +1076,11 @@ unzip9 xs = ( map get1 xs, map get2 xs, map get3 xs -- Scalar Z [1225.0] -- foldAll - :: (Shape sh, Elt a) + :: ( +#ifdef ACCELERATE_INTERNAL_CHECKS + Typeable sh, +#endif + Shape sh, Elt a) => (Exp a -> Exp a -> Exp a) -> Exp a -> Acc (Array sh a) @@ -674,7 +1092,11 @@ foldAll f e arr = fold f e (flatten arr) -- function. -- fold1All - :: (Shape sh, Elt a) + :: ( +#ifdef ACCELERATE_INTERNAL_CHECKS + Typeable sh, +#endif + Shape sh, Elt a) => (Exp a -> Exp a -> Exp a) -> Acc (Array sh a) -> Acc (Scalar a) @@ -708,13 +1130,13 @@ fold1All f arr = fold1 f (flatten arr) -- 40, 170, 0, 138] -- foldSeg - :: forall sh e i. (Shape sh, Elt e, Elt i, i ~ EltR i, IsIntegral i) + :: forall sh e i. (HasCallStack, Shape sh, Elt e, Elt i, i ~ EltR i, IsIntegral i, Ord i, Num i) => (Exp e -> Exp e -> Exp e) -> Exp e -> Acc (Array (sh:.Int) e) -> Acc (Segments i) -> Acc (Array (sh:.Int) e) -foldSeg f z arr seg = foldSeg' f z arr (scanl plus zero seg) +foldSeg f z arr seg = foldSeg' f z arr (scanl plus zero $ map (assertBounds "Segment sizes in foldSeg should be non-negative" (>= 0)) seg) where (plus, zero) = case integralType @i of @@ -735,12 +1157,12 @@ foldSeg f z arr seg = foldSeg' f z arr (scanl plus zero seg) -- descriptor species the length of each of the logical sub-arrays. -- fold1Seg - :: forall sh e i. (Shape sh, Elt e, Elt i, i ~ EltR i, IsIntegral i) + :: forall sh e i. (HasCallStack, Shape sh, Elt e, Elt i, i ~ EltR i, IsIntegral i, Ord i, Num i) => (Exp e -> Exp e -> Exp e) -> Acc (Array (sh:.Int) e) -> Acc (Segments i) -> Acc (Array (sh:.Int) e) -fold1Seg f arr seg = fold1Seg' f arr (scanl plus zero seg) +fold1Seg f arr seg = fold1Seg' f arr (scanl plus zero $ map (assertBounds "Segments in fold1Seg should be non-empty" (> 0)) seg) where plus :: Exp i -> Exp i -> Exp i zero :: Exp i @@ -807,14 +1229,14 @@ any f = or . map f and :: Shape sh => Acc (Array (sh:.Int) Bool) -> Acc (Array sh Bool) -and = fold (&&) True_ +and = fold (&&!) True_ -- | Check if any element along the innermost dimension is 'True'. -- or :: Shape sh => Acc (Array (sh:.Int) Bool) -> Acc (Array sh Bool) -or = fold (||) False_ +or = fold (||!) False_ -- | Compute the sum of elements along the innermost dimension of the array. To -- find the sum of the entire array, 'flatten' it first. @@ -977,16 +1399,14 @@ postscanr f e = map (`f` e) . scanr1 f -- 0, 40, 0, 41, 83, 126, 170, 0, 0, 45, 91, 138] -- scanlSeg - :: forall sh e i. (Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int) + :: forall sh e i. (HasCallStack, Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int) => (Exp e -> Exp e -> Exp e) -> Exp e -> Acc (Array (sh:.Int) e) -> Acc (Segments i) -> Acc (Array (sh:.Int) e) scanlSeg f z arr seg = - if null arr || null flags - then fill (sh ::. sz + length seg) z - else scanl1Seg f arr' seg' + scanl1Seg f arr' seg' where -- Segmented exclusive scan is implemented by first injecting the seed -- element at the head of each segment, and then performing a segmented @@ -998,10 +1418,16 @@ scanlSeg f z arr seg = -- sh ::. sz = shape arr seg' = map (+1) seg - arr' = permute const - (fill (sh ::. sz + length seg) z) - (\(sx ::. i) -> Just_ (sx ::. i + fromIntegral (inc ! I1 i))) - (take (length flags) arr) + arr' = permute' const + (fill (sh ::. sz + length seg) z) + (izipWith + (\(sx ::. i) inci x -> + let i' = sx ::. i + fromIntegral inci + in Just_ (T2 i' x) + ) + (replicate (lift (sh ::. All_)) inc) + arr + ) -- Each element in the segments must be shifted to the right one additional -- place for each successive segment, to make room for the seed element. @@ -1050,16 +1476,14 @@ scanlSeg f z arr seg = -- 40, 170, 0, 138] -- scanl'Seg - :: forall sh e i. (Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int) + :: forall sh e i. (HasCallStack, Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int) => (Exp e -> Exp e -> Exp e) -> Exp e -> Acc (Array (sh:.Int) e) -> Acc (Segments i) -> Acc (Array (sh:.Int) e, Array (sh:.Int) e) scanl'Seg f z arr seg = - if null arr - then T2 arr (fill (indexTail (shape arr) ::. length seg) z) - else T2 body sums + T2 body sums where -- Segmented scan' is implemented by deconstructing a segmented exclusive -- scan, to separate the final value and scan body. @@ -1090,11 +1514,10 @@ scanl'Seg f z arr seg = -- the flags align with the last element of each body section, and when -- scanned, this element will be incremented over. -- - offset = scanl1 (+) seg + offset = scanl1 (+) $ map (assertBounds "Segment sizes in foldSeg should be non-negative" (>= 0)) seg inc = scanl1 (+) - $ permute (+) (fill (I1 $ size arr + 1) 0) - (\ix -> Just_ (index1' (offset ! ix))) - (fill (shape seg) (1 :: Exp i)) + $ permute' (+) (fill (I1 $ size arr + 1) 0) + $ map (\o -> Just_ (T2 (index1' o) (1 :: Exp i))) offset len = offset ! I1 (length offset - 1) body = backpermute @@ -1136,7 +1559,7 @@ scanl'Seg f z arr seg = -- 40, 41, 83, 126, 170, 45, 91, 138] -- scanl1Seg - :: (Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int) + :: (HasCallStack, Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int) => (Exp e -> Exp e -> Exp e) -> Acc (Array (sh:.Int) e) -> Acc (Segments i) @@ -1149,7 +1572,7 @@ scanl1Seg f arr seg -- |Segmented version of 'prescanl'. -- prescanlSeg - :: (Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int) + :: (HasCallStack, Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int) => (Exp e -> Exp e -> Exp e) -> Exp e -> Acc (Array (sh:.Int) e) @@ -1162,7 +1585,7 @@ prescanlSeg f e vec seg -- |Segmented version of 'postscanl'. -- postscanlSeg - :: (Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int) + :: (HasCallStack, Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int) => (Exp e -> Exp e -> Exp e) -> Exp e -> Acc (Array (sh:.Int) e) @@ -1198,16 +1621,14 @@ postscanlSeg f e vec seg -- 42, 0, 178, 135, 91, 46, 0, 0, 144, 97, 49, 0] -- scanrSeg - :: forall sh e i. (Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int) + :: forall sh e i. (HasCallStack, Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int) => (Exp e -> Exp e -> Exp e) -> Exp e -> Acc (Array (sh:.Int) e) -> Acc (Segments i) -> Acc (Array (sh:.Int) e) scanrSeg f z arr seg = - if null arr || null flags - then fill (sh ::. sz + length seg) z - else scanr1Seg f arr' seg' + scanr1Seg f arr' seg' where sh ::. sz = shape arr @@ -1219,10 +1640,16 @@ scanrSeg f z arr seg = inc = scanl1 (+) flags seg' = map (+1) seg - arr' = permute const - (fill (sh ::. sz + length seg) z) - (\(sx ::. i) -> Just_ (sx ::. i + fromIntegral (inc !! i) - 1)) - (drop (sz - length flags) arr) + arr' = permute' const + (fill (sh ::. sz + length seg) z) + (izipWith + (\(sx ::. i) inci x -> + let i' = sx ::. i + fromIntegral inci - 1 + in Just_ (T2 i' x) + ) + (replicate (lift (sh ::. All_)) inc) + (drop (sz - length flags) arr) + ) -- | Segmented version of 'scanr''. @@ -1257,16 +1684,14 @@ scanrSeg f z arr seg = -- 42, 178, 0, 144] -- scanr'Seg - :: forall sh e i. (Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int) + :: forall sh e i. (HasCallStack, Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int) => (Exp e -> Exp e -> Exp e) -> Exp e -> Acc (Array (sh:.Int) e) -> Acc (Segments i) -> Acc (Array (sh:.Int) e, Array (sh:.Int) e) scanr'Seg f z arr seg = - if null arr - then T2 arr (fill (indexTail (shape arr) ::. length seg) z) - else T2 body sums + T2 body sums where -- Using technique described for scanl'Seg -- @@ -1313,7 +1738,7 @@ scanr'Seg f z arr seg = -- 40, 170, 129, 87, 44, 138, 93, 47] -- scanr1Seg - :: (Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int) + :: (HasCallStack, Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int) => (Exp e -> Exp e -> Exp e) -> Acc (Array (sh:.Int) e) -> Acc (Segments i) @@ -1327,7 +1752,7 @@ scanr1Seg f arr seg -- |Segmented version of 'prescanr'. -- prescanrSeg - :: (Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int) + :: (HasCallStack, Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int) => (Exp e -> Exp e -> Exp e) -> Exp e -> Acc (Array (sh:.Int) e) @@ -1340,7 +1765,7 @@ prescanrSeg f e vec seg -- |Segmented version of 'postscanr'. -- postscanrSeg - :: (Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int) + :: (HasCallStack, Shape sh, Slice sh, Elt e, Integral i, Bits i, FromIntegral i Int) => (Exp e -> Exp e -> Exp e) -> Exp e -> Acc (Array (sh:.Int) e) @@ -1362,31 +1787,31 @@ postscanrSeg f e vec seg -- data is used by exclusive segmented scan. -- mkHeadFlags - :: (Integral i, FromIntegral i Int) + :: (HasCallStack, Integral i, FromIntegral i Int) => Acc (Segments i) -> Acc (Segments i) mkHeadFlags seg = init - $ permute (+) zeros (\ix -> Just_ (index1' (offset ! ix))) ones + $ permute' (+) zeros + $ map (\o -> Just_ (T2 (index1' o) 1)) offset where - T2 offset len = scanl' (+) 0 seg + T2 offset len = scanl' (+) 0 $ map (assertBounds "Segment sizes should be non-negative" (>= 0)) seg zeros = fill (index1' $ the len + 1) 0 - ones = fill (index1 $ size offset) 1 -- | Compute tail flags vector from segment vector for right-scans. That -- is, the flag is placed at the last place in each segment. -- mkTailFlags - :: (Integral i, FromIntegral i Int) + :: (HasCallStack, Integral i, FromIntegral i Int) => Acc (Segments i) -> Acc (Segments i) mkTailFlags seg = init - $ permute (+) zeros (\ix -> Just_ (index1' (the len - 1 - offset ! ix))) ones + $ permute' (+) zeros + $ map (\o -> Just_ (T2 (index1' (the len - 1 - o)) 1)) offset where - T2 offset len = scanr' (+) 0 seg + T2 offset len = scanr' (+) 0 $ map (assertBounds "Segment sizes should be non-negative" (>= 0)) seg zeros = fill (index1' $ the len + 1) 0 - ones = fill (index1 $ size offset) 1 -- | Construct a segmented version of a function from a non-segmented -- version. The segmented apply operates on a head-flag value tuple, and @@ -1427,7 +1852,11 @@ index1' i = lift (Z :. fromIntegral i) -- | Flatten the given array of arbitrary dimension into a one-dimensional -- vector. As with 'reshape', this operation performs no work. -- -flatten :: forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Acc (Vector e) +flatten :: forall sh e. ( +#ifdef ACCELERATE_INTERNAL_CHECKS + Typeable sh, +#endif + Shape sh, Elt e) => Acc (Array sh e) -> Acc (Vector e) flatten a | Just Refl <- matchShapeType @sh @DIM1 = a @@ -1443,7 +1872,7 @@ flatten a -- >>> run $ fill (constant (Z:.10)) 0 :: Vector Float -- Vector (Z :. 10) [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0] -- -fill :: (Shape sh, Elt e) => Exp sh -> Exp e -> Acc (Array sh e) +fill :: (HasCallStack, Shape sh, Elt e) => Exp sh -> Exp e -> Acc (Array sh e) fill sh c = generate sh (const c) -- | Create an array of the given shape containing the values @x@, @x+1@, etc. @@ -1458,7 +1887,7 @@ fill sh c = generate sh (const c) -- 40, 41, 42, 43, 44, 45, 46, 47, 48, 49] -- enumFromN - :: (Shape sh, Num e, FromIntegral Int e) + :: (HasCallStack, Shape sh, Num e, FromIntegral Int e) => Exp sh -> Exp e -> Acc (Array sh e) @@ -1476,7 +1905,7 @@ enumFromN sh x = enumFromStepN sh x 1 -- 20.0, 20.5, 21.0, 21.5, 22.0, 22.5, 23.0, 23.5, 24.0, 24.5] -- enumFromStepN - :: (Shape sh, Num e, FromIntegral Int e) + :: (HasCallStack, Shape sh, Num e, FromIntegral Int e) => Exp sh -> Exp e -- ^ x: start -> Exp e -- ^ y: step @@ -1634,7 +2063,11 @@ concatOn dim xs ys = -- >>> run $ filter odd (use mat) -- (Vector (Z :. 20) [1,3,5,7,9,1,1,1,1,1,1,3,5,7,9,11,13,15,17,19],Vector (Z :. 4) [5,5,0,10]) -- -filter :: (Shape sh, Elt e) +filter :: ( +#ifdef ACCELERATE_INTERNAL_CHECKS + Typeable sh, +#endif + Shape sh, Elt e) => (Exp e -> Exp Bool) -> Acc (Array (sh:.Int) e) -> Acc (Vector e, Array sh Int) @@ -1649,7 +2082,12 @@ filter p arr = compact (map p arr) arr -- | As 'filter', but with separate arrays for the data elements and the -- flags indicating which elements of that array should be kept. -- -compact :: forall sh e. (Shape sh, Elt e) +compact :: forall sh e. ( +#ifdef ACCELERATE_INTERNAL_CHECKS + Typeable sh, + +#endif + Shape sh, Elt e) => Acc (Array (sh:.Int) Bool) -> Acc (Array (sh:.Int) e) -> Acc (Vector e, Array sh Int) @@ -1659,33 +2097,34 @@ compact keep arr | Just Refl <- matchShapeType @sh @Z = let T2 target len = scanl' (+) 0 (map boolToInt keep) - prj ix = if keep!ix - then Just_ (I1 (target!ix)) - else Nothing_ - dummy = fill (I1 (the len)) undef - result = permute const dummy prj arr + f value k t = + if k then + Just_ $ T2 (I1 t) value + else + Nothing_ + -- Allocate an array with the size of the input, instead of 'the len'. + -- We later shrink the array to 'the len'. This allows us to fuse the + -- scanl' with the permute. The backpermute to shrink the array might + -- be fused with a later consumer of the output of compact, depending + -- on the program. + dummy = fill (shape arr) undef -- fill (I1 (the len)) undef + result = permuteUnique' dummy $ zipWith3 f arr keep target in - if null arr - then T2 emptyArray (fill Z_ 0) - else - if the len == unindex1 (shape arr) - then T2 arr len - else T2 result len - + T2 + (backpermute (I1 $ the len) id result) + len compact keep arr = let sz = indexTail (shape arr) T2 target len = scanl' (+) 0 (map boolToInt keep) T2 offset valid = scanl' (+) 0 (flatten len) - prj ix = if keep!ix - then Just_ (I1 (offset !! (toIndex sz (indexTail ix)) + target!ix)) + prj ix k t x = if k + then Just_ (T2 (I1 (offset !! toIndex sz (indexTail ix) + t)) x) else Nothing_ dummy = fill (I1 (the valid)) undef - result = permute const dummy prj arr + result = permute' const dummy (izipWith3 prj keep target arr) in - if null arr - then T2 emptyArray (fill sz 0) - else T2 result len + T2 result len -- Gather operations @@ -1699,7 +2138,7 @@ compact keep arr -- Vector (Z :. 6) [9,4,1,6,2,4] -- gather - :: (Shape sh, Elt e) + :: (HasCallStack, Shape sh, Elt e) => Acc (Array sh Int) -- ^ index of source at each index to gather -> Acc (Vector e) -- ^ source values -> Acc (Array sh e) @@ -1724,19 +2163,18 @@ gather indices input = map (input !!) indices -- Vector (Z :. 6) [6.0,6.0,1.0,6.0,2.0,4.0] -- gatherIf - :: (Elt a, Elt b) + :: (HasCallStack, Elt a, Elt b) => Acc (Vector Int) -- ^ source indices to gather from -> Acc (Vector a) -- ^ mask vector -> (Exp a -> Exp Bool) -- ^ predicate function -> Acc (Vector b) -- ^ default values -> Acc (Vector b) -- ^ source values -> Acc (Vector b) -gatherIf from maskV pred defaults input = zipWith zf pf gatheredV +gatherIf from maskV pred defaults input = + zipWith3 f from maskV defaults where - zf p g = p ? (unlift g) - gatheredV = zip (gather from input) defaults - pf = map pred maskV - + -- Note that we only read from 'input' + f idx mask def = if pred mask then input ! I1 idx else def -- Scatter operations -- ------------------ @@ -1753,16 +2191,14 @@ gatherIf from maskV pred defaults input = zipWith zf pf gatheredV -- Vector (Z :. 10) [0,1,4,9,0,4,0,6,2,0] -- scatter - :: Elt e + :: (HasCallStack, Elt e) => Acc (Vector Int) -- ^ destination indices to scatter into -> Acc (Vector e) -- ^ default values -> Acc (Vector e) -- ^ source values -> Acc (Vector e) -scatter to defaults input = permute const defaults pf input' +scatter to defaults input = permuteUnique' defaults (zipWith pf to input) where - pf ix = Just_ (I1 (to ! ix)) - input' = backpermute (shape to `intersect` shape input) id input - + pf t x = Just_ (T2 (I1 t) x) -- | Conditionally overwrite elements of the destination by scattering values of -- the source array according to a given index mapping, whenever the masking @@ -1778,21 +2214,19 @@ scatter to defaults input = permute const defaults pf input' -- Vector (Z :. 10) [0,0,0,0,0,4,0,6,2,0] -- scatterIf - :: (Elt a, Elt b) + :: (HasCallStack, Elt a, Elt b) => Acc (Vector Int) -- ^ destination indices to scatter into -> Acc (Vector a) -- ^ mask vector -> (Exp a -> Exp Bool) -- ^ predicate function -> Acc (Vector b) -- ^ default values -> Acc (Vector b) -- ^ source values -> Acc (Vector b) -scatterIf to maskV pred defaults input = permute const defaults pf input' +scatterIf to maskV pred defaults input = permuteUnique' defaults (zipWith3 pf maskV to input) where - input' = backpermute (shape to `intersect` shape input) id input - pf ix = if pred (maskV ! ix) - then Just_ (I1 (to ! ix)) + pf m t x = if pred m + then Just_ (T2 (I1 t) x) else Nothing_ - -- Permutations -- ------------ @@ -2086,7 +2520,7 @@ takeOn dim n xs = sh = shape xs m = sh ^. dim in - backpermute (sh & dim .~ min m n) id xs + backpermute (sh & dim .~ min m (max 0 n)) id xs -- | Generalised version of 'drop' where the argument 'Lens'' specifies which @@ -2138,44 +2572,23 @@ slitOn slitOn dim i n = takeOn dim n . dropOn dim i --- Controlling execution --- --------------------- +-- Composition of array computations +-- --------------------------------- --- | Force an array expression to be evaluated, preventing it from fusing with --- other operations. Forcing operations to be computed to memory, rather than --- being fused into their consuming function, can sometimes improve performance. --- For example, computing a matrix 'transpose' could provide better memory --- locality for the subsequent operation. Preventing fusion to split large --- operations into several simpler steps could also help by reducing register --- pressure. +-- | Pipelining of two array computations. The first argument will be fully +-- evaluated before being passed to the second computation. This can be used to +-- prevent the argument being fused into the function, for example. -- --- Preventing fusion also means that the individual operations are available to --- be executed concurrently with other kernels. In particular, consider using --- this if you have a series of operations that are compute bound rather than --- memory bound. +-- Denotationally, we have -- --- Here is the synthetic example: --- --- > loop :: Exp Int -> Exp Int --- > loop ticks = --- > let clockRate = 900000 -- kHz --- > in while (\i -> i < clockRate * ticks) (+1) 0 --- > --- > test :: Acc (Vector Int) --- > test = --- > zip3 --- > (compute $ map loop (use $ fromList (Z:.1) [10])) --- > (compute $ map loop (use $ fromList (Z:.1) [10])) --- > (compute $ map loop (use $ fromList (Z:.1) [10])) --- > +-- > (acc1 >-> acc2) arrs = let tmp = acc1 arrs +-- > in tmp `seq` acc2 tmp -- --- Without the use of 'compute', the operations are fused together and the three --- long-running loops are executed sequentially in a single kernel. Instead, the --- individual operations can now be executed concurrently, potentially reducing --- overall runtime. --- -compute :: Arrays a => Acc a -> Acc a -compute = id >-> id +infixl 1 >-> +(>->) :: forall a b c. (Arrays a, Arrays b, Arrays c) => (Acc a -> Acc b) -> (Acc b -> Acc c) -> (Acc a -> Acc c) +acc1 >-> acc2 = acc2 . compute . acc1 +-- 'compute' used to be implemented in terms of this operator, but now that's +-- the other way around: (>->) is implemented in terms of compute. -- Flow control @@ -2219,6 +2632,56 @@ instance (Elt a, b ~ Bool) => IfThenElse (Exp b) (Exp a) where instance (Arrays a, b ~ Bool) => IfThenElse (Exp b) (Acc a) where ifThenElse = acond +-- | Variant of 'awhile' that speculatively executes one iteration ahead. +-- The result of the (last) speculatively executed iteration is ignored if the +-- condition evaluates to false, but if an error is thrown in the last +-- iteration, then the loop fails. This variant may be beneficial for GPUs, as +-- the GPU can already start the next iteration before the host has evaluated +-- the condition and scheduled the new kernel calls. The result of +-- 'awhileSpeculative' is the same as 'awhile' if the step-function doesn't +-- throw an error. +-- +-- If you want to do in-place updates with 'permute' in the loop, then you +-- should use 'awhile' or 'awhileAndOne', as 'awhileSpeculative' doesn't have +-- unique references to arrays. +-- +awhileSpeculative + :: forall a. Arrays a + => (Acc a -> Acc (Scalar Bool)) -- ^ keep evaluating while this returns 'True' + -> (Acc a -> Acc a) -- ^ function to apply + -> Acc a -- ^ initial value + -> Acc a +awhileSpeculative condition step initial + -- Perform an awhile loop with the previous and current state as the state. + -- The condition can then be performed on the old state and the step on the + -- current. At the end we return the first element of the tuple. + -- + = afst + $ awhile + (condition . afst) + (\(T2 _ current) -> T2 current $ step current) + (T2 initial $ step initial) + +-- | Variant of 'awhile' that executes one additional iteration after the +-- condition evaluates to False. This may be beneficial for GPUs, as the GPU +-- can already start the next iteration before the host has evaluated the +-- condition and scheduled the new kernel calls. The result of +-- 'awhileAndOne c f i' is equal to 'f $ awhile c f i', but the first is +-- likely more performant. +-- +awhileAndOne + :: forall a. Arrays a + => (Acc a -> Acc (Scalar Bool)) -- ^ keep evaluating while this returns 'True' + -> (Acc a -> Acc a) -- ^ function to apply + -> Acc a -- ^ initial value + -> Acc a +awhileAndOne condition step initial + -- + = asnd + $ awhile + afst + (\(T2 _ state) -> T2 (condition state) (step state)) + (T2 (unit True_) initial) -- | The 'match' operation is the core operation which enables embedded -- pattern matching. It is applied to an n-ary scalar function, and @@ -2334,7 +2797,7 @@ iterate n f z -- | Reduce along an innermost slice of an array /sequentially/, by applying a -- binary operator to a starting value and the array from left to right. -- -sfoldl :: (Shape sh, Elt a, Elt b) +sfoldl :: (HasCallStack, Shape sh, Elt a, Elt b) => (Exp a -> Exp b -> Exp a) -> Exp a -> Exp sh @@ -2442,11 +2905,6 @@ unindex3 (Z_ ::. k ::. j ::. i) = T3 k j i the :: Elt e => Acc (Scalar e) -> Exp e the = (! index0) --- | Test whether an array is empty. --- -null :: (Shape sh, Elt e) => Acc (Array sh e) -> Exp Bool -null arr = size arr == 0 - -- | Get the length of a vector. -- length :: Elt e => Acc (Vector e) -> Exp Int @@ -2516,48 +2974,48 @@ length = unindex1 . shape -- -- @since 1.3.0.0 -- -expand :: (Elt a, Elt b) +expand :: (HasCallStack, Elt a, Elt b) => (Exp a -> Exp Int) -> (Exp a -> Exp Int -> Exp b) -> Acc (Vector a) -> Acc (Vector b) expand f g xs = let - szs = map f xs - T2 offset len = scanl' (+) 0 szs - m = the len + szs = map (\a -> let sz = f a in assertBounds "Each element in expand should expand to a non-negative number of elements" (>= 0) sz) xs + T2 offsets len = scanl' (+) 0 szs + -- TODO: Make the OperationAcc simplifier better and remove 'the $ unit $'. + -- Our simplifier currently only knows that multiple array of size 'm' are + -- the same if we first store that size in a separate scalar array. We + -- should improve the simplifier by for instance doing more common + -- subexpression elimination of Compute terms. + -- + -- Until we do that, we have this tiny hack here. This way the compiler + -- detects that several arrays have the same size, which allows fusion to + -- fuse more operations. + m = the $ unit $ the len + + head_flags :: Acc (Vector Int) + head_flags = permuteUnique' (fill (I1 m) 0) + $ zipWith + (\sz offset -> if sz == 0 ||! offset >= m then Nothing_ else Just_ $ T2 (I1 offset) 1) + szs offsets + + idxs = map (subtract 1) + $ map snd + $ scanl1 (segmentedL (+)) + $ zip head_flags + $ fill (I1 m) 1 + + iotas = map snd + $ scanl1 (segmentedL const) + $ zip head_flags + $ permuteUnique' (fill (I1 m) undef) + $ zipWith3 (\ix sz offset -> if sz == 0 ||! offset >= m then Nothing_ else Just_ $ T2 (I1 offset) ix) + (enumFromN (shape xs) 0) + szs + offsets in - if length xs == 0 || m == 0 - then use $ fromList (Z:.0) [] - else - let - n = m + 1 - put ix = Just_ (I1 (offset ! ix)) - - head_flags :: Acc (Vector Int) - head_flags = permute const (fill (I1 n) 0) put (fill (shape szs) 1) - - idxs = map (subtract 1) - $ map snd - $ scanl1 (segmentedL (+)) - $ zip head_flags - $ fill (I1 m) 1 - - iotas = map snd - $ scanl1 (segmentedL const) - $ zip head_flags - $ permute const - (fill (I1 n) undef) - -- If any of the elements expand to zero new - -- elements then this would result in multiple - -- writes to the same index since the offsets are - -- also the same, which is undefined behaviour - (\ix -> if szs ! ix > 0 - then put ix - else Nothing_) - $ enumFromN (shape xs) 0 - in - zipWith g (gather iotas xs) idxs + zipWith g (gather iotas xs) idxs {-- @@ -2608,8 +3066,8 @@ generateSeq n f = toSeq (Z :. Split) (generate (index1 n) (f . unindex1)) -- Utilities -- --------- -emptyArray :: (Shape sh, Elt e) => Acc (Array sh e) -emptyArray = fill (constant empty) undef +-- emptyArray :: (Shape sh, Elt e) => Acc (Array sh e) +-- emptyArray = fill (constant empty) undef -- Lenses diff --git a/src/Data/Array/Accelerate/Pretty.hs b/src/Data/Array/Accelerate/Pretty.hs index 9aa5848e8..ea8fa3ba9 100644 --- a/src/Data/Array/Accelerate/Pretty.hs +++ b/src/Data/Array/Accelerate/Pretty.hs @@ -1,12 +1,4 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE UndecidableInstances #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# OPTIONS_GHC -Wno-orphans #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.Pretty @@ -23,27 +15,23 @@ module Data.Array.Accelerate.Pretty ( -- ** Pretty printing PrettyAcc, ExtractAcc, prettyPreOpenAcc, + prettyOpenAcc, prettyPreOpenAfun, prettyOpenExp, prettyOpenFun, + renderForTerminal, -- ** Graphviz Graph, - PrettyGraph(..), Detail(..), + PrettyGraph, Detail, graphDelayedAcc, graphDelayedAfun, ) where import Data.Array.Accelerate.AST hiding ( Acc, Exp ) -import Data.Array.Accelerate.Debug.Internal.Flags -import Data.Array.Accelerate.Error import Data.Array.Accelerate.Pretty.Graphviz import Data.Array.Accelerate.Pretty.Print hiding ( Keyword(..) ) -import Data.Array.Accelerate.Smart ( Acc, Exp ) -import Data.Array.Accelerate.Sugar.Array -import Data.Array.Accelerate.Sugar.Elt -import Data.Array.Accelerate.Trafo -import Data.Array.Accelerate.Trafo.Delayed +-- import Data.Array.Accelerate.Trafo.Delayed import Data.Maybe import Prettyprinter @@ -56,23 +44,13 @@ import qualified Data.Text.Lazy as T import qualified System.Console.ANSI as Term import qualified System.Console.Terminal.Size as Term +{- TODO WALL: DEAD CODE #if ACCELERATE_DEBUG import Control.DeepSeq +import Data.Array.Accelerate.Debug.Internal.Flags import Data.Array.Accelerate.Debug.Internal.Stats #endif - - -instance Arrays arrs => Show (Acc arrs) where - show = withSimplStats . show . convertAcc - -instance Afunction (Acc a -> f) => Show (Acc a -> f) where - show = withSimplStats . show . convertAfun - -instance Elt e => Show (Exp e) where - show = withSimplStats . show . convertExp - -instance Function (Exp a -> f) => Show (Exp a -> f) where - show = withSimplStats . show . convertFun +-} -- instance Typeable a => Show (Seq a) where -- show = withSimplStats . show . convertSeq @@ -95,13 +73,6 @@ instance PrettyEnv aenv => Show (OpenAcc aenv a) where instance PrettyEnv aenv => Show (OpenAfun aenv f) where show = renderForTerminal . prettyPreOpenAfun configPlain prettyOpenAcc (prettyEnv (pretty 'a')) -instance PrettyEnv aenv => Show (DelayedOpenAcc aenv a) where - show = let config = if shouldPrintHash then configWithHash else configPlain - in renderForTerminal . prettyDelayedOpenAcc config context0 (prettyEnv (pretty 'a')) - -instance PrettyEnv aenv => Show (DelayedOpenAfun aenv f) where - show = let config = if shouldPrintHash then configWithHash else configPlain - in renderForTerminal . prettyPreOpenAfun config prettyDelayedOpenAcc (prettyEnv (pretty 'a')) instance (PrettyEnv env, PrettyEnv aenv) => Show (OpenExp env aenv e) where show = renderForTerminal . prettyOpenExp context0 (prettyEnv (pretty 'x')) (prettyEnv (pretty 'a')) @@ -135,7 +106,7 @@ terminalLayoutOptions = unsafePerformIO $ do term <- Term.size return $ case term of - Nothing -> defaultLayoutOptions + Nothing -> LayoutOptions { layoutPageWidth = Unbounded } Just t -> LayoutOptions { layoutPageWidth = AvailablePerLine (min w 120) f } where w = Term.width t @@ -151,38 +122,26 @@ extractOpenAcc :: OpenAcc aenv a -> PreOpenAcc OpenAcc aenv a extractOpenAcc (OpenAcc pacc) = pacc -prettyDelayedOpenAcc :: HasCallStack => PrettyAcc DelayedOpenAcc -prettyDelayedOpenAcc config context aenv (Manifest pacc) - = prettyPreOpenAcc config context prettyDelayedOpenAcc extractDelayedOpenAcc aenv pacc -prettyDelayedOpenAcc _ _ aenv (Delayed _ sh f _) - = parens - $ nest shiftwidth - $ sep [ delayed "delayed" - , prettyOpenExp app Empty aenv sh - , parens $ prettyOpenFun Empty aenv f - ] - -extractDelayedOpenAcc :: HasCallStack => DelayedOpenAcc aenv a -> PreOpenAcc DelayedOpenAcc aenv a -extractDelayedOpenAcc (Manifest pacc) = pacc -extractDelayedOpenAcc Delayed{} = internalError "expected manifest array" - - --- Unfortunately, using unsafePerformIO here means that the getFlag will be --- evaluated only once when the first 'show' is performed on a Delayed value; --- afterwards, the thunk will have been evaluated, and all future pretty-print --- outputs will use the same result. --- This cannot be prevented using a NOINLINE pragma, since then the function --- itself is still a thunk that will only be evaluated once. --- --- The practical result of this is that @setFlag verbose@ will not change --- anything after a Delayed has already been printed once. -shouldPrintHash :: Bool -shouldPrintHash = unsafePerformIO $ getFlag verbose +-- prettyDelayedOpenAcc :: HasCallStack => PrettyAcc DelayedOpenAcc +-- prettyDelayedOpenAcc config context aenv (Manifest pacc) +-- = prettyPreOpenAcc config context prettyDelayedOpenAcc extractDelayedOpenAcc aenv pacc +-- prettyDelayedOpenAcc _ _ aenv (Delayed _ sh f _) +-- = parens +-- $ nest shiftwidth +-- $ sep [ delayed "delayed" +-- , prettyOpenExp app Empty aenv sh +-- , parens $ prettyOpenFun Empty aenv f +-- ] + +-- extractDelayedOpenAcc :: HasCallStack => DelayedOpenAcc aenv a -> PreOpenAcc DelayedOpenAcc aenv a +-- extractDelayedOpenAcc (Manifest pacc) = pacc +-- extractDelayedOpenAcc Delayed{} = internalError "expected manifest array" -- Debugging -- --------- +{- TODO WALL: DEAD CODE -- Attach simplifier statistics to the tail of the given string. Since the -- statistics rely on fully evaluating the expression this is difficult to do -- generally (without an additional deepseq), but easy enough for our show @@ -199,4 +158,4 @@ withSimplStats x = unsafePerformIO $ do #else withSimplStats x = x #endif - +-} diff --git a/src/Data/Array/Accelerate/Pretty/Exp.hs b/src/Data/Array/Accelerate/Pretty/Exp.hs new file mode 100644 index 000000000..ac59c0d90 --- /dev/null +++ b/src/Data/Array/Accelerate/Pretty/Exp.hs @@ -0,0 +1,517 @@ +{-# LANGUAGE RecordWildCards #-} +-- | +-- Module : Data.Array.Accelerate.Pretty.Exp +-- Copyright : [2008..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.Pretty.Exp ( + Adoc, + Keyword(..), + for_, let_, in_, case_, of_, if_, then_, else_, + prettyPreOpenFun, prettyPreOpenExp, PrettyArrayInstr, + parensIf, Operator(..), Associativity, Direction(..), + needsParens, prettyLhs, prettyLhs', prettyELhs, prettyConst, + prettyTupR, prettyText, + primOperator, isInfix, + Precedence, Fixity(..), + Val(..), PrettyEnv(..), sizeEnv, prj, prjs, + Context(..), context0, app, + shiftwidth, (?), + IdxF(..) + ) where + +import Data.Array.Accelerate.AST.Exp hiding ( Direction ) +import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.AST.Var +import Data.Array.Accelerate.Representation.Elt +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Sugar.Foreign +import Data.Array.Accelerate.Type + +import Data.Char +import Data.String +import qualified Data.Text as T +import Prettyprinter +import Prelude hiding ( exp ) +import Control.DeepSeq (NFData) + +type Adoc = Doc Keyword + +data Keyword + = Statement -- do | case of | let in + | Conditional -- if then else + | FunctionName + | Execute + | Modifier + deriving (Eq, Show) + +for_, let_, in_ :: Adoc +for_ = annotate Statement "for" +let_ = annotate Statement "let" +in_ = annotate Statement "in" + +case_, of_ :: Adoc +case_ = annotate Statement "case" +of_ = annotate Statement "of" + +if_, then_, else_ :: Adoc +if_ = annotate Statement "if" +then_ = annotate Statement "then" +else_ = annotate Statement "else" + +type PrettyArrayInstr arr = forall s. Context -> arr s -> (Context -> Adoc) -> Adoc + +prettyPreOpenFun + :: forall arr env f. + PrettyArrayInstr arr + -> Val env + -> PreOpenFun arr env f + -> Adoc +prettyPreOpenFun prettyArrayInstr env0 = next (pretty '\\') env0 + where + next :: Adoc -> Val env' -> PreOpenFun arr env' f' -> Adoc + next vs env (Body body) + -- PrimApp f x <- body + -- , op <- primOperator f + -- , isInfix op + -- , Tuple (NilTup `SnocTup` a `SnocTup` b) <- x + -- , Var (SuccIdx ZeroIdx) <- a + -- , Var ZeroIdx <- b + -- = opName op -- surrounding context will add parens + -- + = hang shiftwidth (sep [ vs <> "->" + , prettyPreOpenExp context0 prettyArrayInstr env body]) + next vs env (Lam lhs lam) = + let (env', lhs') = prettyELhs True env lhs + in next (vs <> lhs' <> space) env' lam + +prettyPreOpenExp + :: forall arr env t. + Context + -> PrettyArrayInstr arr + -> Val env + -> PreOpenExp arr env t + -> Adoc +prettyPreOpenExp ctx prettyArrayInstr env exp = + case exp of + Evar (Var _ idx) -> prj idx env + Let{} -> prettyLet ctx prettyArrayInstr env exp + PrimApp f x + | a `Pair` b <- x -> ppF2 op (ppE a) (ppE b) + | otherwise -> ppF1 op' (ppE x) + where + op = primOperator f + op' = isInfix op ? (Operator (parens (opName op)) App L 10, op) + -- + Const tp c -> prettyConst (TupRsingle tp) c + Pair{} -> prettyTuple ctx prettyArrayInstr env exp + Nil -> "()" + VecPack _ e -> ppF1 "pack" (ppE e) + VecUnpack _ e -> ppF1 "unpack" (ppE e) + Case x xs d -> prettyCase prettyArrayInstr env x xs d + Cond p t e -> flatAlt multi single + where + p' = ppE p context0 + t' = ppE t context0 + e' = ppE e context0 + -- + single = parensIf (needsParens ctx (Operator "?:" Infix N 0)) + $ sep [ p', pretty '?', t', pretty ':', e' ] + multi = parensIf (ctxPrecedence ctx > 0) + $ hang 3 + $ vsep [ if_ <+> p' + , hang shiftwidth (sep [ then_, t' ]) + , hang shiftwidth (sep [ else_, e' ]) ] + -- + Select p t e -> ppF3 "select" (ppE p) (ppE t) (ppE e) + ToIndex _ sh ix -> ppF2 "toIndex" (ppE sh) (ppE ix) + FromIndex _ sh ix -> ppF2 "fromIndex" (ppE sh) (ppE ix) + While p f x -> ppF3 "while" (ppF p) (ppF f) (ppE x) + Foreign _ ff _ e -> ppF2 "foreign" (\_ -> pretty (strForeign ff)) (ppE e) + ArrayInstr instr e -> prettyArrayInstr ctx instr (ppE e) + ShapeSize _ sh -> ppF1 "shapeSize" (ppE sh) + Coerce _ tp x -> ppF1 (Operator (withTypeRep tp "coerce") App L 10) (ppE x) + Undef tp -> withTypeRep tp "undef" + Assert msg e1 e2 -> ppF3 "assert" (\_ -> prettyText msg) (ppE e1) (ppE e2) + Assume e1 e2 -> ppF2 "assume" (ppE e1) (ppE e2) + + where + ppE :: PreOpenExp arr env e -> Context -> Adoc + ppE e c = prettyPreOpenExp c prettyArrayInstr env e + + ppF :: PreOpenFun arr env f -> Context -> Adoc + ppF f _ = parens $ prettyPreOpenFun prettyArrayInstr env f + + ppF1 :: Operator -> (Context -> Adoc) -> Adoc + ppF1 op x + = parensIf (needsParens ctx op) + $ combine [ opName op, x ctx' ] + where + ctx' = isPrefix op ? (arg op R, app) + combine = isPrefix op ? (cat, hang 2 . sep) + + ppF2 :: Operator -> (Context -> Adoc) -> (Context -> Adoc) -> Adoc + ppF2 op x y + = parensIf (needsParens ctx op) + $ if isInfix op + then sep [ x (arg op L), group (sep [opName op, y (arg op R)]) ] + else hang 2 $ sep [ opName op, x app, y app ] + + ppF3 :: Operator -> (Context -> Adoc) -> (Context -> Adoc) -> (Context -> Adoc) -> Adoc + ppF3 op x y z + = parensIf (needsParens ctx op) + $ hang 2 + $ sep [ opName op, x app, y app, z app ] + + withTypeRep :: ScalarType t -> Adoc -> Adoc + withTypeRep t op = op <+> "@" <> pretty (show t) + +prettyLet + :: forall arr env t. + Context + -> PrettyArrayInstr arr + -> Val env + -> PreOpenExp arr env t + -> Adoc +prettyLet ctx prettyArrayInstr env0 + = parensIf (ctxPrecedence ctx > 0) + . align . wrap . collect env0 + where + collect :: Val env' -> PreOpenExp arr env' e -> ([Adoc], Adoc) + collect env = + \case + Let lhs e1 e2 -> + let (env', v) = prettyELhs False env lhs + e1' = ppE env e1 + bnd | isLet e1 = nest shiftwidth (vsep [v <+> equals, e1']) + | otherwise = v <+> align (equals <+> e1') + (bnds, body) = collect env' e2 + in + (bnd:bnds, body) + -- + next -> ([], ppE env next) + + isLet :: PreOpenExp arr env' t' -> Bool + isLet Let{} = True + isLet _ = False + + ppE :: Val env' -> PreOpenExp arr env' t' -> Adoc + ppE env = prettyPreOpenExp context0 prettyArrayInstr env + + wrap :: ([Adoc], Adoc) -> Adoc + wrap ([], body) = body -- shouldn't happen! + wrap ([b], body) + = sep [ nest shiftwidth (sep [let_, b]), in_, body ] + wrap (bnds, body) + = vsep [ nest shiftwidth (vsep [let_, sepBy (flatAlt "" " ; ") bnds]) + , in_ + , body + ] + + sepBy :: Adoc -> [Adoc] -> Adoc + sepBy = encloseSep mempty mempty + +prettyTuple + :: forall arr env t. + Context + -> PrettyArrayInstr arr + -> Val env + -> PreOpenExp arr env t + -> Adoc +prettyTuple ctx prettyArrayInstr env exp = case collect exp of + Nothing -> align $ ppPair exp + Just tup -> + case tup of + [] -> "()" + _ -> align $ parensIf (ctxPrecedence ctx > 0) ("T" <> pretty (length tup) <+> align (sep tup)) + where + ppPair :: PreOpenExp arr env t' -> Adoc + ppPair (Pair e1 e2) = "(" <> ppPair e1 <> "," <+> prettyPreOpenExp context0 prettyArrayInstr env e2 <> ")" + ppPair e = prettyPreOpenExp context0 prettyArrayInstr env e + + collect :: PreOpenExp arr env t' -> Maybe [Adoc] + collect Nil = Just [] + collect (Pair e1 e2) + | Just tup <- collect e1 = Just $ tup ++ [prettyPreOpenExp app prettyArrayInstr env e2] + collect _ = Nothing + +prettyCase + :: PrettyArrayInstr arr + -> Val env + -> PreOpenExp arr env a + -> [(TAG, PreOpenExp arr env b)] + -> Maybe (PreOpenExp arr env b) + -> Adoc +prettyCase prettyArrayInstr env x xs def + = hang shiftwidth + $ vsep [ case_ <+> x' <+> of_ + , flatAlt (vcat xs') (encloseSep "{ " " }" "; " xs') + ] + where + x' = prettyPreOpenExp context0 prettyArrayInstr env x + xs' = map (\(t,e) -> pretty t <+> "->" <+> prettyPreOpenExp context0 prettyArrayInstr env e) xs + ++ case def of + Nothing -> [] + Just d -> ["_" <+> "->" <+> prettyPreOpenExp context0 prettyArrayInstr env d] + +prettyConst :: TypeR e -> e -> Adoc +prettyConst tp x = + let y = showElt tp x + in parensIf (any isSpace y) (pretty y) + +data Operator = Operator + { opName :: Adoc + , opFixity :: Fixity + , opAssociativity :: Associativity + , opPrecedence :: Precedence + } + +instance IsString Operator where + fromString s = Operator (fromString s) App L 10 + +needsParens :: Context -> Operator -> Bool +needsParens Context{..} Operator{..} + | ctxPrecedence < opPrecedence = False + | ctxPrecedence > opPrecedence = True + | ctxAssociativity /= opAssociativity = True + | otherwise = ctxPosition /= opAssociativity + +parensIf :: Bool -> Doc ann -> Doc ann +parensIf True = group . parens . align +parensIf False = id + + +prettyELhs :: Bool -> Val env -> LeftHandSide s t env env' -> (Val env', Adoc) +prettyELhs requiresParens = prettyLhs requiresParens 'x' + +prettyLhs :: forall s env env' t. Bool -> Char -> Val env -> LeftHandSide s t env env' -> (Val env', Adoc) +prettyLhs requiresParens x = prettyLhs' push Nothing requiresParens + where + push :: s t' -> Val env'' -> (Val (env'', t'), Adoc) + push _ env = (Push env v, v) + where + v = pretty x <> pretty (sizeEnv env) + +prettyLhs' + :: forall val s env env' t. + (forall t' env''. s t' -> val env'' -> (val (env'', t'), Adoc)) + -> (Maybe (Exists s -> Adoc)) + -> Bool + -> val env + -> LeftHandSide s t env env' + -> (val env', Adoc) +prettyLhs' push prettyTp requiresParens env0 lhs = case collect lhs of + Nothing -> ppPair requiresParens lhs + Just (env1, tup) -> + case tup of + [] -> (env1, "()") + _ -> (env1, parensIf requiresParens (pretty 'T' <> pretty (length tup) <+> align (sep tup))) + where + ppPair :: Bool -> LeftHandSide s arrs' env env'' -> (val env'', Adoc) + ppPair _ LeftHandSideUnit = (env0, "()") + ppPair requiresP (LeftHandSideWildcard t) + | Just prettyT <- prettyTp = (env0, parensIf requiresP $ "_: " <> align (prettyTupR (const $ prettyT . Exists) 0 t)) + | otherwise = (env0, "_") + ppPair requiresP (LeftHandSideSingle t) = (env1, v') + where + (env1, v) = push t env0 + v' + | Just prettyT <- prettyTp = parensIf requiresP $ v <> ": " <> align (prettyT $ Exists t) + | otherwise = v + ppPair _ (LeftHandSidePair a b) = (env2, tupled [doc1, doc2]) + where + (env1, doc1) = ppPair False a + (env2, doc2) = prettyLhs' push prettyTp False env1 b + + collect :: LeftHandSide s arrs' env env'' -> Maybe (val env'', [Adoc]) + collect (LeftHandSidePair l1 l2) + | Just (env1, tup ) <- collect l1 + , (env2, doc2) <- prettyLhs' push prettyTp True env1 l2 = Just (env2, tup ++ [doc2]) + collect (LeftHandSideWildcard TupRunit) = Just (env0, []) + collect _ = Nothing + +prettyTupR :: forall s t. (forall t'. Precedence -> s t' -> Adoc) -> Precedence -> TupR s t -> Adoc +prettyTupR elt p tupR + | TupRunit <- tupR = "()" + | TupRsingle e <- tupR = elt p e + | Just tuple <- collect tupR [] + = parensIf (p > 0) + $ "T" <> pretty (length tuple) <+> align (sep tuple) + | otherwise = ppPair tupR + where + -- Try to recover the 'tuple-format', i.e. a structure of pairs of the form + -- ((( (), a ), b ), c ) + collect :: TupR s t' -> [Adoc] -> Maybe [Adoc] + collect TupRunit accum = Just accum + collect (TupRpair v1 (TupRsingle v2)) accum = collect v1 (elt 10 v2 : accum) + collect _ _ = Nothing + + -- Prints a pair which is not in the tuple-format. + -- We thus do not have to traverse it again to check for this format, + -- preventing repeatedly calling 'collect' on the same structure. + -- + ppPair :: TupR s t' -> Adoc + ppPair (TupRpair v1 v2) = tupled [ppPair v1, prettyTupR elt 0 v2] + ppPair v = prettyTupR elt 0 v + +prettyText :: T.Text -> Adoc +prettyText msg + | T.length msg >= maxLength = viaShow $ T.take maxLength msg <> ellipsis + | otherwise = viaShow msg + where + ellipsis = " (...)" + maxLength = 80 - T.length ellipsis + +-- Primitive operators +-- ------------------- +-- +-- The core of the pretty printer is how to correctly handle precedence, +-- associativity, and fixity of the primitive scalar operators. +-- + +data Direction = L | N | R + deriving Eq + +data Fixity = App | Infix | Prefix + deriving Eq + +type Precedence = Int +type Associativity = Direction + +data Context = Context + { ctxAssociativity :: Associativity + , ctxPosition :: Direction + , ctxPrecedence :: Precedence + } + +context0 :: Context +context0 = Context N N 0 + +app :: Context +app = Context L N 10 + +arg :: Operator -> Direction -> Context +arg Operator{..} side = Context opAssociativity side opPrecedence + +isPrefix :: Operator -> Bool +isPrefix Operator{..} = opFixity == Prefix + +isInfix :: Operator -> Bool +isInfix Operator{..} = opFixity == Infix + +primOperator :: PrimFun a -> Operator +primOperator PrimAdd{} = Operator (pretty '+') Infix L 6 +primOperator PrimSub{} = Operator (pretty '-') Infix L 6 +primOperator PrimMul{} = Operator (pretty '*') Infix L 7 +primOperator PrimNeg{} = Operator (pretty '-') Prefix L 6 -- Haskell's only prefix operator +primOperator PrimAbs{} = Operator "abs" App L 10 +primOperator PrimSig{} = Operator "signum" App L 10 +primOperator PrimQuot{} = Operator "quot" App L 10 +primOperator PrimRem{} = Operator "rem" App L 10 +primOperator PrimQuotRem{} = Operator "quotRem" App L 10 +primOperator PrimIDiv{} = Operator "div" App L 10 +primOperator PrimMod{} = Operator "mod" App L 10 +primOperator PrimDivMod{} = Operator "divMod" App L 10 +primOperator PrimBAnd{} = Operator ".&." Infix L 7 +primOperator PrimBOr{} = Operator ".|." Infix L 5 +primOperator PrimBXor{} = Operator "xor" App L 10 +primOperator PrimBNot{} = Operator "complement" App L 10 +primOperator PrimBShiftL{} = Operator "shiftL" App L 10 +primOperator PrimBShiftR{} = Operator "shiftR" App L 10 +primOperator PrimBRotateL{} = Operator "rotateL" App L 10 +primOperator PrimBRotateR{} = Operator "rotateR" App L 10 +primOperator PrimPopCount{} = Operator "popCount" App L 10 +primOperator PrimCountLeadingZeros{} = Operator "countLeadingZeros" App L 10 +primOperator PrimCountTrailingZeros{} = Operator "countTrailingZeros" App L 10 +primOperator PrimFDiv{} = Operator (pretty '/') Infix L 7 +primOperator PrimRecip{} = Operator "recip" App L 10 +primOperator PrimSin{} = Operator "sin" App L 10 +primOperator PrimCos{} = Operator "cos" App L 10 +primOperator PrimTan{} = Operator "tan" App L 10 +primOperator PrimAsin{} = Operator "asin" App L 10 +primOperator PrimAcos{} = Operator "acos" App L 10 +primOperator PrimAtan{} = Operator "atan" App L 10 +primOperator PrimSinh{} = Operator "sinh" App L 10 +primOperator PrimCosh{} = Operator "cosh" App L 10 +primOperator PrimTanh{} = Operator "tanh" App L 10 +primOperator PrimAsinh{} = Operator "asinh" App L 10 +primOperator PrimAcosh{} = Operator "acosh" App L 10 +primOperator PrimAtanh{} = Operator "atanh" App L 10 +primOperator PrimExpFloating{} = Operator "exp" App L 10 +primOperator PrimSqrt{} = Operator "sqrt" App L 10 +primOperator PrimLog{} = Operator "log" App L 10 +primOperator PrimFPow{} = Operator "**" Infix R 8 +primOperator PrimLogBase{} = Operator "logBase" App L 10 +primOperator PrimTruncate{} = Operator "truncate" App L 10 +primOperator PrimRound{} = Operator "round" App L 10 +primOperator PrimFloor{} = Operator "floor" App L 10 +primOperator PrimCeiling{} = Operator "ceiling" App L 10 +primOperator PrimAtan2{} = Operator "atan2" App L 10 +primOperator PrimIsNaN{} = Operator "isNaN" App L 10 +primOperator PrimIsInfinite{} = Operator "isInfinite" App L 10 +primOperator (PrimCmp _ CmpLt) = Operator "<" Infix N 4 +primOperator (PrimCmp _ CmpGtEq) = Operator ">=" Infix N 4 +primOperator (PrimCmp _ CmpEq) = Operator "==" Infix N 4 +primOperator (PrimCmp _ CmpNEq) = Operator "/=" Infix N 4 +primOperator PrimMax{} = Operator "max" App L 10 +primOperator PrimMin{} = Operator "min" App L 10 +primOperator PrimLAnd = Operator "&&" Infix R 3 +primOperator PrimLOr = Operator "||" Infix R 2 +primOperator PrimLNot = Operator "not" App L 10 +primOperator PrimFromIntegral{} = Operator "fromIntegral" App L 10 +primOperator PrimToFloating{} = Operator "toFloating" App L 10 + + +-- Environments +-- ------------ + +data Val env where + Empty :: Val () + Push :: Val env -> Adoc -> Val (env, t) + +class PrettyEnv env where + prettyEnv :: Adoc -> Val env + +instance PrettyEnv () where + prettyEnv _ = Empty + +instance PrettyEnv env => PrettyEnv (env, t) where + prettyEnv v = + let env = prettyEnv v :: Val env + x = v <> pretty (sizeEnv env) + in + env `Push` x + +sizeEnv :: Val env -> Int +sizeEnv Empty = 0 +sizeEnv (Push env _) = 1 + sizeEnv env + +prj :: Idx env t -> Val env -> Adoc +prj ZeroIdx (Push _ v) = v +prj (SuccIdx ix) (Push env _) = prj ix env + +prjs :: TupR (IdxF f env) e -> Val env -> [Adoc] +prjs TupRunit _ = [] +prjs (TupRpair l r) env = prjs l env <> prjs r env +prjs (TupRsingle (IdxF ix)) env = [prj ix env] + +-- Utilities +-- --------- + +shiftwidth :: Int +shiftwidth = 2 + +infix 0 ? +(?) :: Bool -> (a, a) -> a +True ? (t,_) = t +False ? (_,f) = f + + +newtype IdxF f env a = IdxF { runIdxF :: Idx env (f a)} + deriving newtype (NFData) diff --git a/src/Data/Array/Accelerate/Pretty/Graphviz.hs b/src/Data/Array/Accelerate/Pretty/Graphviz.hs index 5cdbb098b..3731d5dc0 100644 --- a/src/Data/Array/Accelerate/Pretty/Graphviz.hs +++ b/src/Data/Array/Accelerate/Pretty/Graphviz.hs @@ -1,15 +1,3 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE PatternGuards #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE RecordWildCards #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TupleSections #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeSynonymInstances #-} -{-# LANGUAGE ViewPatterns #-} -- | -- Module : Data.Array.Accelerate.Pretty.Graphviz -- Copyright : [2015..2020] The Accelerate Team @@ -22,519 +10,532 @@ module Data.Array.Accelerate.Pretty.Graphviz ( Graph, - PrettyGraph(..), Detail(..), + PrettyGraph, Detail, graphDelayedAcc, graphDelayedAfun, ) where -import Data.Array.Accelerate.AST -import Data.Array.Accelerate.AST.Idx -import Data.Array.Accelerate.AST.LeftHandSide -import Data.Array.Accelerate.AST.Var -import Data.Array.Accelerate.Analysis.Match -import Data.Array.Accelerate.Error -import Data.Array.Accelerate.Pretty.Graphviz.Monad -import Data.Array.Accelerate.Pretty.Graphviz.Type -import Data.Array.Accelerate.Pretty.Print hiding ( Keyword(..) ) -import Data.Array.Accelerate.Representation.Array -import Data.Array.Accelerate.Representation.Stencil -import Data.Array.Accelerate.Representation.Type -import Data.Array.Accelerate.Sugar.Foreign -import Data.Array.Accelerate.Trafo.Delayed -import Data.Array.Accelerate.Trafo.Substitution - -import Control.Applicative hiding ( Const, empty ) -import Control.Arrow ( (&&&) ) -import Control.Monad.State ( modify, gets, state ) -import Data.HashSet ( HashSet ) -import Data.List ( nub, partition ) -import Data.Maybe -import Data.String -import Prettyprinter -import Prelude hiding ( exp ) -import qualified Data.HashSet as Set -import qualified Data.Sequence as Seq - - --- Configuration options --- --------------------- - -cfgIncludeShape, cfgUnique :: Bool -cfgIncludeShape = False -- draw edges for uses of shape information -cfgUnique = False -- draw a single edge per data dependency - - --- Environments --- ------------ - --- This is the standard environment typed by de Bruijn indices, where at each --- index we need to record both the pretty printed label as well its 'NodeId', --- which we use to track data dependencies. --- -data Aval env where - Aempty :: Aval () - Apush :: Aval env -> NodeId -> Label -> Aval (env, t) - --- Convert to the 'Val' used by the base pretty printing module by stripping out --- the 'NodeId' part. --- -avalToVal :: Aval aenv -> Val aenv -avalToVal Aempty = Empty -avalToVal (Apush aenv _ v) = Push (avalToVal aenv) (pretty v) - -aprj :: Idx aenv t -> Aval aenv -> (NodeId, Label) -- TLM: (Vertex, Label) ?? -aprj ZeroIdx (Apush _ n v) = (n,v) -aprj (SuccIdx ix) (Apush aenv _ _) = aprj ix aenv - - --- Graph construction --- ------------------ - -mkNode :: PNode -> Maybe Label -> Dot NodeId -mkNode (PNode ident tree deps) label = - let node = Node label ident tree - edges = Seq.fromList - $ map (\(from, to) -> Edge from (Vertex ident to)) - $ if cfgUnique then nub deps else deps - in - state $ \s -> - ( ident - , s { dotNodes = node Seq.<| dotNodes s - , dotEdges = edges Seq.>< dotEdges s - } - ) - - --- Add [T|F] ports underneath the given tree. --- -mkTF :: Tree (Maybe Port, Adoc) -> Tree (Maybe Port, Adoc) -mkTF this = - Forest [ this - , Forest [ Leaf (Just "T", "T") - , Leaf (Just "F", "F") - ] - ] - - --- Pretty Printing --- =============== --- - -class PrettyGraph g where - ppGraph :: Detail -> g -> Graph - -instance PrettyGraph (DelayedAcc a) where - ppGraph = graphDelayedAcc - -instance PrettyGraph (DelayedAfun a) where - ppGraph = graphDelayedAfun - -data Detail = Simple | Full - -simple :: Detail -> Bool -simple Simple = True -simple _ = False - --- | Generate a dependency graph for the given computation --- -{-# NOINLINE graphDelayedAcc #-} -graphDelayedAcc :: HasCallStack => Detail -> DelayedAcc a -> Graph -graphDelayedAcc detail acc = - evalDot (graphDelayedOpenAcc detail Aempty acc) - --- | Generate a dependency graph for an array function --- -{-# NOINLINE graphDelayedAfun #-} -graphDelayedAfun :: HasCallStack => Detail -> DelayedAfun f -> Graph -graphDelayedAfun detail afun = evalDot $! do - l <- prettyDelayedAfun detail Aempty afun - state $ \s -> - case Seq.viewl (dotGraph s) of - g@(Graph l' _) Seq.:< gs | l == l' -> (g, s { dotGraph = gs }) - _ -> internalError "unexpected error" - - --- Pretty-printing data-dependency graphs --- -------------------------------------- - --- Partially constructed graph nodes, consists of some body text and a list of --- vertices which we will draw edges from (and later, the port we connect into). --- -data PDoc = PDoc Adoc [Vertex] -data PNode = PNode NodeId (Tree (Maybe Port, Adoc)) [(Vertex, Maybe Port)] - -graphDelayedOpenAcc - :: HasCallStack - => Detail - -> Aval aenv - -> DelayedOpenAcc aenv a - -> Dot Graph -graphDelayedOpenAcc detail aenv acc = do - r <- prettyDelayedOpenAcc detail context0 aenv acc - i <- genNodeId - v <- mkNode r Nothing - _ <- mkNode (PNode i (Leaf (Nothing,"result")) [(Vertex v Nothing, Nothing)]) Nothing - mkGraph - --- Generate a graph for the given term. --- -prettyDelayedOpenAcc - :: forall aenv arrs. HasCallStack - => Detail -- simplified output: only print operator name - -> Context - -> Aval aenv - -> DelayedOpenAcc aenv arrs - -> Dot PNode -prettyDelayedOpenAcc _ _ _ Delayed{} = internalError "expected manifest array" -prettyDelayedOpenAcc detail ctx aenv (Manifest pacc) = - case pacc of - Avar ix -> pnode (avar ix) - Alet lhs bnd body -> do - bnd'@(PNode ident _ _) <- prettyDelayedOpenAcc detail context0 aenv bnd - (aenv1, a) <- prettyLetALeftHandSide ident aenv lhs - _ <- mkNode bnd' (Just a) - body' <- prettyDelayedOpenAcc detail context0 aenv1 body - return body' - - Acond p t e -> do - ident <- genNodeId - vt <- lift t - ve <- lift e - PDoc p' vs <- ppE p - let port = Just "P" - doc = mkTF $ Leaf (port, if simple detail then "?|" else p') - deps = (vt, Just "T") : (ve, Just "F") : map (,port) vs - return $ PNode ident doc deps - - Apply _ afun acc -> apply <$> prettyDelayedAfun detail aenv afun - <*> prettyDelayedOpenAcc detail ctx aenv acc - - Awhile p f x -> do - ident <- genNodeId - x' <- replant =<< prettyDelayedOpenAcc detail app aenv x - p' <- prettyDelayedAfun detail aenv p - f' <- prettyDelayedAfun detail aenv f - -- - case x' of - PNode _ (Leaf (Nothing,xb)) fvs -> - let loop = nest 2 (sep ["awhile", pretty p', pretty f', xb ]) - in return $ PNode ident (Leaf (Nothing,loop)) fvs - _ -> internalError "replant did not return a Leaf" - - Apair a1 a2 -> genNodeId >>= prettyDelayedApair detail aenv a1 a2 - - Anil -> "()" .$ [] - Atrace (Message _ _ msg) as bs -> "atrace" .$ [ return $ PDoc (pretty msg) [], ppA as, ppA bs ] - Use repr arr -> "use" .$ [ return $ PDoc (prettyArray repr arr) [] ] - Unit _ e -> "unit" .$ [ ppE e ] - Generate _ sh f -> "generate" .$ [ ppE sh, ppF f ] - Transform _ sh ix f xs -> "transform" .$ [ ppE sh, ppF ix, ppF f, ppA xs ] - Reshape _ sh xs -> "reshape" .$ [ ppE sh, ppA xs ] - Replicate _ty ix xs -> "replicate" .$ [ ppE ix, ppA xs ] - Slice _ty xs ix -> "slice" .$ [ ppA xs, ppE ix ] - Map _ f xs -> "map" .$ [ ppF f, ppA xs ] - ZipWith _ f xs ys -> "zipWith" .$ [ ppF f, ppA xs, ppA ys ] - Fold f (Just z) a -> "fold" .$ [ ppF f, ppE z, ppA a ] - Fold f Nothing a -> "fold1" .$ [ ppF f, ppA a ] - FoldSeg _ f (Just z) a s -> "foldSeg" .$ [ ppF f, ppE z, ppA a, ppA s ] - FoldSeg _ f Nothing a s -> "fold1Seg" .$ [ ppF f, ppA a, ppA s ] - Scan d f (Just z) a -> ppD "scan" d "" .$ [ ppF f, ppE z, ppA a ] - Scan d f Nothing a -> ppD "scan" d "1" .$ [ ppF f, ppA a ] - Scan' d f z a -> ppD "scan" d "'" .$ [ ppF f, ppE z, ppA a ] - Permute f dfts p xs -> "permute" .$ [ ppF f, ppA dfts, ppF p, ppA xs ] - Backpermute _ sh p xs -> "backpermute" .$ [ ppE sh, ppF p, ppA xs ] - Stencil s _ sten bndy xs -> "stencil" .$ [ ppF sten, ppB (stencilEltR s) bndy, ppA xs ] - Stencil2 s1 s2 _ s b1 a1 b2 a2 -> "stencil2" .$ [ ppF s, ppB (stencilEltR s1) b1, ppA a1, ppB (stencilEltR s2) b2, ppA a2 ] - Aforeign _ ff _afun xs -> "aforeign" .$ [ return (PDoc (pretty (strForeign ff)) []), {- ppAf afun, -} ppA xs ] - -- Collect{} -> error "Collect" - - where - (.$) :: Operator -> [Dot PDoc] -> Dot PNode - name .$ docs = pnode =<< fmt name docs - - fmt :: Operator -> [Dot PDoc] -> Dot PDoc - fmt name docs = do - docs' <- sequence docs - let args = [ x | PDoc x _ <- docs' ] - fvs = [ x | PDoc _ x <- docs' ] - doc = if simple detail - then manifest name - else parensIf (needsParens ctx name) - $ nest shiftwidth - $ sep ( manifest name : args ) - return $ PDoc doc (concat fvs) - - pnode :: PDoc -> Dot PNode - pnode (PDoc doc vs) = do - let port = Nothing - ident <- genNodeId - return $ PNode ident (Leaf (port, doc)) (map (,port) vs) - - -- Free variables - -- - fvF :: Fun aenv t -> [Vertex] - fvF = fvOpenFun Empty aenv - - fvE :: Exp aenv t -> [Vertex] - fvE = fvOpenExp Empty aenv - - -- Pretty-printing - -- - avar :: ArrayVar aenv t -> PDoc - avar (Var _ ix) = let (ident, v) = aprj ix aenv - in PDoc (pretty v) [Vertex ident Nothing] - - aenv' :: Val aenv - aenv' = avalToVal aenv - - ppA :: HasCallStack => DelayedOpenAcc aenv a -> Dot PDoc - ppA (Manifest (Avar ix)) = return (avar ix) - ppA acc@Manifest{} = do - -- Lift out and draw as a separate node. This can occur with the manifest - -- array arguments to permute (defaults array) and stencil[2]. - acc' <- prettyDelayedOpenAcc detail app aenv acc - v <- mkLabel - ident <- mkNode acc' (Just v) - return $ PDoc (pretty v) [Vertex ident Nothing] - ppA (Delayed _ sh f _) - | Shape a <- sh -- identical shape - , Just b <- isIdentityIndexing f -- function is `\ix -> b ! ix` - , Just Refl <- matchVar a b -- function thus is `\ix -> a ! ix` - = ppA $ Manifest $ Avar a - ppA (Delayed _ sh f _) = do - PDoc d v <- "Delayed" `fmt` [ ppE sh, ppF f ] - return $ PDoc (parens d) v - - ppB :: forall sh e. HasCallStack - => TypeR e - -> Boundary aenv (Array sh e) - -> Dot PDoc - ppB _ Clamp = return (PDoc "clamp" []) - ppB _ Mirror = return (PDoc "mirror" []) - ppB _ Wrap = return (PDoc "wrap" []) - ppB tp (Constant e) = return (PDoc (prettyConst tp e) []) - ppB _ (Function f) = ppF f - - ppF :: HasCallStack => Fun aenv t -> Dot PDoc - ppF = return . uncurry PDoc . (parens . prettyFun aenv' &&& fvF) - - ppE :: HasCallStack => Exp aenv t -> Dot PDoc - ppE = return . uncurry PDoc . (prettyExp aenv' &&& fvE) - - ppD :: String -> Direction -> String -> Operator - ppD f LeftToRight k = fromString (f <> "l" <> k) - ppD f RightToLeft k = fromString (f <> "r" <> k) - - lift :: HasCallStack => DelayedOpenAcc aenv a -> Dot Vertex - lift Delayed{} = internalError "expected manifest array" - lift (Manifest (Avar (Var _ ix))) = return $ Vertex (fst (aprj ix aenv)) Nothing - lift acc = do - acc' <- prettyDelayedOpenAcc detail context0 aenv acc - ident <- mkNode acc' Nothing - return $ Vertex ident Nothing - - apply :: Label -> PNode -> PNode - apply f (PNode ident x vs) = - let x' = case x of - Leaf (p,d) -> Leaf (p, pretty f <+> d) - Forest ts -> Forest (Leaf (Nothing,pretty f) : ts) - in - PNode ident x' vs - - --- Pretty print array functions as separate sub-graphs, and return the name of --- the sub-graph as if it can be called like a function. We will add additional --- nodes at the top of the graph to represent the bound variables. --- --- Note: [Edge placement] --- --- If a node belongs to a particular graph, so too must all its edges (and --- vertices). This means that if the subgraph references anything from the --- enclosing environment, we must lift those edges out of this subgraph, --- otherwise the referenced node will be drawn inside of the subgraph. --- -prettyDelayedAfun - :: HasCallStack - => Detail - -> Aval aenv - -> DelayedOpenAfun aenv afun - -> Dot Label -prettyDelayedAfun detail aenv afun = do - Graph _ ss <- mkSubgraph (go aenv afun) - n <- Seq.length <$> gets dotGraph - let label = "afun" <> fromString (show (n+1)) - outer = collect aenv - (lifted,ss') = - flip partition ss $ \s -> - case s of - E (Edge (Vertex ident _) _) -> Set.member ident outer - _ -> False - -- - modify $ \s -> s { dotGraph = dotGraph s Seq.|> Graph label ss' - , dotEdges = Seq.fromList [ e | E e <- lifted ] Seq.>< dotEdges s - } - return label - where - go :: Aval aenv' -> DelayedOpenAfun aenv' a' -> Dot Graph - go aenv' (Abody b) = graphDelayedOpenAcc detail aenv' b - go aenv' (Alam lhs f) = do - aenv'' <- prettyLambdaALeftHandSide aenv' lhs - go aenv'' f - - collect :: Aval aenv' -> HashSet NodeId - collect Aempty = Set.empty - collect (Apush a i _) = Set.insert i (collect a) - -prettyLetALeftHandSide - :: forall repr aenv aenv'. HasCallStack - => NodeId - -> Aval aenv - -> ALeftHandSide repr aenv aenv' - -> Dot (Aval aenv', Label) -prettyLetALeftHandSide _ aenv (LeftHandSideWildcard repr) = return (aenv, doc) - where - doc = case repr of - TupRunit -> "()" - _ -> "_" -prettyLetALeftHandSide ident aenv (LeftHandSideSingle _) = do - a <- mkLabel - return (Apush aenv ident a, a) -prettyLetALeftHandSide ident aenv (LeftHandSidePair lhs1 lhs2) = do - (aenv1, d1) <- prettyLetALeftHandSide ident aenv lhs1 - (aenv2, d2) <- prettyLetALeftHandSide ident aenv1 lhs2 - return (aenv2, "(" <> d1 <> ", " <> d2 <> ")") - -prettyLambdaALeftHandSide - :: forall repr aenv aenv'. HasCallStack - => Aval aenv - -> ALeftHandSide repr aenv aenv' - -> Dot (Aval aenv') -prettyLambdaALeftHandSide aenv (LeftHandSideWildcard _) = return aenv -prettyLambdaALeftHandSide aenv (LeftHandSideSingle _) = do - a <- mkLabel - ident <- genNodeId - _ <- mkNode (PNode ident (Leaf (Nothing, pretty a)) []) Nothing - return $ Apush aenv ident a -prettyLambdaALeftHandSide aenv (LeftHandSidePair lhs1 lhs2) = do - aenv1 <- prettyLambdaALeftHandSide aenv lhs1 - prettyLambdaALeftHandSide aenv1 lhs2 - --- Display array tuples. This is a little tricky... --- -prettyDelayedApair - :: forall aenv a1 a2. HasCallStack - => Detail - -> Aval aenv - -> DelayedOpenAcc aenv a1 - -> DelayedOpenAcc aenv a2 - -> NodeId - -> Dot PNode -prettyDelayedApair detail aenv a1 a2 ident = do - PNode id1 t1 v1 <- prettyElem a1 - PNode id2 t2 v2 <- prettyElem a2 - modify $ \s -> s { dotEdges = fmap (redirect ident [id1, id2]) (dotEdges s) } - return $ PNode ident (forest [t1, t2]) (v1 ++ v2) - where - prettyElem :: DelayedOpenAcc aenv a -> Dot PNode - prettyElem a = replant =<< prettyDelayedOpenAcc detail context0 aenv a - - -- Redirect any edges that pointed into one of the nodes now part of this - -- tuple, to instead point to the container node. - -- - redirect :: NodeId -> [NodeId] -> Edge -> Edge - redirect new subs edge@(Edge from (Vertex to port)) - | to `elem` subs = Edge from (Vertex new port) - | otherwise = edge - - -- Since we have lifted out any non-leaves into separate nodes, we can - -- simply tuple-up all of the elements. - -- - forest :: [Tree (Maybe Port, Adoc)] -> Tree (Maybe Port, Adoc) - forest leaves = Leaf (Nothing, tupled [ align d | Leaf (Nothing,d) <- leaves ]) - - --- Lift out anything that isn't a Leaf node and output it to the graph --- immediately as a new labelled node. --- -replant :: PNode -> Dot PNode -replant pnode@(PNode ident tree _) = - case tree of - Leaf (Nothing, _) -> return pnode - _ -> do - vacuous <- genNodeId - a <- mkLabel - _ <- mkNode pnode (Just a) - return $ PNode vacuous (Leaf (Nothing, pretty a)) [(Vertex ident Nothing, Nothing)] - - --- Pretty printing scalar functions and expressions --- ------------------------------------------------ --- --- This is done with the usual machinery. Note that we rely on knowing that all --- array operations will be lifted out of scalar expressions. This means that we --- don't really need to recurse into the scalar terms to uncover new graph --- nodes. --- - --- Data dependencies --- ----------------- --- --- Return the data-dependencies of the given term. This is just a tree traversal --- to extract all of the free variables. We will draw an edge from each of those --- nodes (vertices) into the current term. --- - -fvAvar :: Aval aenv -> ArrayVar aenv a -> [Vertex] -fvAvar env (Var _ ix) = [ Vertex (fst $ aprj ix env) Nothing ] - -fvOpenFun - :: forall env aenv fun. - Val env - -> Aval aenv - -> OpenFun env aenv fun - -> [Vertex] -fvOpenFun env aenv (Body b) = fvOpenExp env aenv b -fvOpenFun env aenv (Lam lhs f) = fvOpenFun env' aenv f - where - (env', _) = prettyELhs True env lhs - -fvOpenExp - :: forall env aenv exp. - Val env - -> Aval aenv - -> OpenExp env aenv exp - -> [Vertex] -fvOpenExp env aenv = fv - where - fvF :: OpenFun env aenv f -> [Vertex] - fvF = fvOpenFun env aenv - - fv :: OpenExp env aenv e -> [Vertex] - fv (Shape acc) = if cfgIncludeShape then fvAvar aenv acc else [] - fv (Index acc i) = concat [ fvAvar aenv acc, fv i ] - fv (LinearIndex acc i) = concat [ fvAvar aenv acc, fv i ] - -- - fv (Let lhs e1 e2) = concat [ fv e1, fvOpenExp env' aenv e2 ] - where - (env', _) = prettyELhs False env lhs - fv Evar{} = [] - fv Undef{} = [] - fv Const{} = [] - fv PrimConst{} = [] - fv (PrimApp _ x) = fv x - fv (Pair e1 e2) = concat [ fv e1, fv e2] - fv Nil = [] - fv (VecPack _ e) = fv e - fv (VecUnpack _ e) = fv e - fv (IndexSlice _ slix sh) = concat [ fv slix, fv sh ] - fv (IndexFull _ slix sh) = concat [ fv slix, fv sh ] - fv (ToIndex _ sh ix) = concat [ fv sh, fv ix ] - fv (FromIndex _ sh ix) = concat [ fv sh, fv ix ] - fv (ShapeSize _ sh) = fv sh - fv Foreign{} = [] - fv (Case e rhs def) = concat [ fv e, concat [ fv c | (_,c) <- rhs ], maybe [] fv def ] - fv (Cond p t e) = concat [ fv p, fv t, fv e ] - fv (While p f x) = concat [ fvF p, fvF f, fv x ] - fv (Coerce _ _ e) = fv e +import Data.Kind (Type, Constraint) + +graphDelayedAcc :: a +graphDelayedAcc = undefined + +graphDelayedAfun :: a +graphDelayedAfun = undefined + +type Graph = () +type PrettyGraph :: Type -> Constraint +type PrettyGraph a = () +type Detail = () + +-- import Data.Array.Accelerate.AST +-- import Data.Array.Accelerate.AST.Idx +-- import Data.Array.Accelerate.AST.LeftHandSide +-- import Data.Array.Accelerate.AST.Var +-- import Data.Array.Accelerate.Analysis.Match +-- import Data.Array.Accelerate.Error +-- import Data.Array.Accelerate.Pretty.Graphviz.Monad +-- import Data.Array.Accelerate.Pretty.Graphviz.Type +-- import Data.Array.Accelerate.Pretty.Print hiding ( Keyword(..) ) +-- import Data.Array.Accelerate.Representation.Array +-- import Data.Array.Accelerate.Representation.Stencil +-- import Data.Array.Accelerate.Representation.Type +-- import Data.Array.Accelerate.Sugar.Foreign +-- import Data.Array.Accelerate.Trafo.Delayed +-- import Data.Array.Accelerate.Trafo.Substitution + +-- import Control.Applicative hiding ( Const, empty ) +-- import Control.Arrow ( (&&&) ) +-- import Control.Monad.State ( modify, gets, state ) +-- import Data.HashSet ( HashSet ) +-- import Data.List ( nub, partition ) +-- import Data.Maybe +-- import Data.String +-- import Prettyprinter +-- import Prelude hiding ( exp ) +-- import qualified Data.HashSet as Set +-- import qualified Data.Sequence as Seq + + +-- -- Configuration options +-- -- --------------------- + +-- cfgIncludeShape, cfgUnique :: Bool +-- cfgIncludeShape = False -- draw edges for uses of shape information +-- cfgUnique = False -- draw a single edge per data dependency + + +-- -- Environments +-- -- ------------ + +-- -- This is the standard environment typed by de Bruijn indices, where at each +-- -- index we need to record both the pretty printed label as well its 'NodeId', +-- -- which we use to track data dependencies. +-- -- +-- data Aval env where +-- Aempty :: Aval () +-- Apush :: Aval env -> NodeId -> Label -> Aval (env, t) + +-- -- Convert to the 'Val' used by the base pretty printing module by stripping out +-- -- the 'NodeId' part. +-- -- +-- avalToVal :: Aval aenv -> Val aenv +-- avalToVal Aempty = Empty +-- avalToVal (Apush aenv _ v) = Push (avalToVal aenv) (pretty v) + +-- aprj :: Idx aenv t -> Aval aenv -> (NodeId, Label) -- TLM: (Vertex, Label) ?? +-- aprj ZeroIdx (Apush _ n v) = (n,v) +-- aprj (SuccIdx ix) (Apush aenv _ _) = aprj ix aenv + + +-- -- Graph construction +-- -- ------------------ + +-- mkNode :: PNode -> Maybe Label -> Dot NodeId +-- mkNode (PNode ident tree deps) label = +-- let node = Node label ident tree +-- edges = Seq.fromList +-- $ map (\(from, to) -> Edge from (Vertex ident to)) +-- $ if cfgUnique then nub deps else deps +-- in +-- state $ \s -> +-- ( ident +-- , s { dotNodes = node Seq.<| dotNodes s +-- , dotEdges = edges Seq.>< dotEdges s +-- } +-- ) + + +-- -- Add [T|F] ports underneath the given tree. +-- -- +-- mkTF :: Tree (Maybe Port, Adoc) -> Tree (Maybe Port, Adoc) +-- mkTF this = +-- Forest [ this +-- , Forest [ Leaf (Just "T", "T") +-- , Leaf (Just "F", "F") +-- ] +-- ] + + +-- -- Pretty Printing +-- -- =============== +-- -- + +-- class PrettyGraph g where +-- ppGraph :: Detail -> g -> Graph + +-- instance PrettyGraph (DelayedAcc a) where +-- ppGraph = graphDelayedAcc + +-- instance PrettyGraph (DelayedAfun a) where +-- ppGraph = graphDelayedAfun + +-- data Detail = Simple | Full + +-- simple :: Detail -> Bool +-- simple Simple = True +-- simple _ = False + +-- -- | Generate a dependency graph for the given computation +-- -- +-- {-# NOINLINE graphDelayedAcc #-} +-- graphDelayedAcc :: HasCallStack => Detail -> DelayedAcc a -> Graph +-- graphDelayedAcc detail acc = +-- evalDot (graphDelayedOpenAcc detail Aempty acc) + +-- -- | Generate a dependency graph for an array function +-- -- +-- {-# NOINLINE graphDelayedAfun #-} +-- graphDelayedAfun :: HasCallStack => Detail -> DelayedAfun f -> Graph +-- graphDelayedAfun detail afun = evalDot $! do +-- l <- prettyDelayedAfun detail Aempty afun +-- state $ \s -> +-- case Seq.viewl (dotGraph s) of +-- g@(Graph l' _) Seq.:< gs | l == l' -> (g, s { dotGraph = gs }) +-- _ -> internalError "unexpected error" + + +-- -- Pretty-printing data-dependency graphs +-- -- -------------------------------------- + +-- -- Partially constructed graph nodes, consists of some body text and a list of +-- -- vertices which we will draw edges from (and later, the port we connect into). +-- -- +-- data PDoc = PDoc Adoc [Vertex] +-- data PNode = PNode NodeId (Tree (Maybe Port, Adoc)) [(Vertex, Maybe Port)] + +-- graphDelayedOpenAcc +-- :: HasCallStack +-- => Detail +-- -> Aval aenv +-- -> DelayedOpenAcc aenv a +-- -> Dot Graph +-- graphDelayedOpenAcc detail aenv acc = do +-- r <- prettyDelayedOpenAcc detail context0 aenv acc +-- i <- genNodeId +-- v <- mkNode r Nothing +-- _ <- mkNode (PNode i (Leaf (Nothing,"result")) [(Vertex v Nothing, Nothing)]) Nothing +-- mkGraph + +-- -- Generate a graph for the given term. +-- -- +-- prettyDelayedOpenAcc +-- :: forall aenv arrs. HasCallStack +-- => Detail -- simplified output: only print operator name +-- -> Context +-- -> Aval aenv +-- -> DelayedOpenAcc aenv arrs +-- -> Dot PNode +-- prettyDelayedOpenAcc _ _ _ Delayed{} = internalError "expected manifest array" +-- prettyDelayedOpenAcc detail ctx aenv (Manifest pacc) = +-- case pacc of +-- Avar ix -> pnode (avar ix) +-- Alet lhs bnd body -> do +-- bnd'@(PNode ident _ _) <- prettyDelayedOpenAcc detail context0 aenv bnd +-- (aenv1, a) <- prettyLetALeftHandSide ident aenv lhs +-- _ <- mkNode bnd' (Just a) +-- body' <- prettyDelayedOpenAcc detail context0 aenv1 body +-- return body' + +-- Acond p t e -> do +-- ident <- genNodeId +-- vt <- lift t +-- ve <- lift e +-- PDoc p' vs <- ppE p +-- let port = Just "P" +-- doc = mkTF $ Leaf (port, if simple detail then "?|" else p') +-- deps = (vt, Just "T") : (ve, Just "F") : map (,port) vs +-- return $ PNode ident doc deps + +-- Apply _ afun acc -> apply <$> prettyDelayedAfun detail aenv afun +-- <*> prettyDelayedOpenAcc detail ctx aenv acc + +-- Awhile p f x -> do +-- ident <- genNodeId +-- x' <- replant =<< prettyDelayedOpenAcc detail app aenv x +-- p' <- prettyDelayedAfun detail aenv p +-- f' <- prettyDelayedAfun detail aenv f +-- -- +-- case x' of +-- PNode _ (Leaf (Nothing,xb)) fvs -> +-- let loop = nest 2 (sep ["awhile", pretty p', pretty f', xb ]) +-- in return $ PNode ident (Leaf (Nothing,loop)) fvs +-- _ -> internalError "replant did not return a Leaf" + +-- Apair a1 a2 -> genNodeId >>= prettyDelayedApair detail aenv a1 a2 + +-- Anil -> "()" .$ [] +-- Atrace (Message _ _ msg) as bs -> "atrace" .$ [ return $ PDoc (pretty msg) [], ppA as, ppA bs ] +-- Use repr arr -> "use" .$ [ return $ PDoc (prettyArray repr arr) [] ] +-- Unit _ e -> "unit" .$ [ ppE e ] +-- Generate _ sh f -> "generate" .$ [ ppE sh, ppF f ] +-- Transform _ sh ix f xs -> "transform" .$ [ ppE sh, ppF ix, ppF f, ppA xs ] +-- Reshape _ sh xs -> "reshape" .$ [ ppE sh, ppA xs ] +-- Replicate _ty ix xs -> "replicate" .$ [ ppE ix, ppA xs ] +-- Slice _ty xs ix -> "slice" .$ [ ppA xs, ppE ix ] +-- Map _ f xs -> "map" .$ [ ppF f, ppA xs ] +-- ZipWith _ f xs ys -> "zipWith" .$ [ ppF f, ppA xs, ppA ys ] +-- Fold f (Just z) a -> "fold" .$ [ ppF f, ppE z, ppA a ] +-- Fold f Nothing a -> "fold1" .$ [ ppF f, ppA a ] +-- FoldSeg _ f (Just z) a s -> "foldSeg" .$ [ ppF f, ppE z, ppA a, ppA s ] +-- FoldSeg _ f Nothing a s -> "fold1Seg" .$ [ ppF f, ppA a, ppA s ] +-- Scan d f (Just z) a -> ppD "scan" d "" .$ [ ppF f, ppE z, ppA a ] +-- Scan d f Nothing a -> ppD "scan" d "1" .$ [ ppF f, ppA a ] +-- Scan' d f z a -> ppD "scan" d "'" .$ [ ppF f, ppE z, ppA a ] +-- Permute f dfts p xs -> "permute" .$ [ ppF f, ppA dfts, ppF p, ppA xs ] +-- Backpermute _ sh p xs -> "backpermute" .$ [ ppE sh, ppF p, ppA xs ] +-- Stencil s _ sten bndy xs -> "stencil" .$ [ ppF sten, ppB (stencilEltR s) bndy, ppA xs ] +-- Stencil2 s1 s2 _ s b1 a1 b2 a2 -> "stencil2" .$ [ ppF s, ppB (stencilEltR s1) b1, ppA a1, ppB (stencilEltR s2) b2, ppA a2 ] +-- Aforeign _ ff _afun xs -> "aforeign" .$ [ return (PDoc (pretty (strForeign ff)) []), {- ppAf afun, -} ppA xs ] +-- -- Collect{} -> error "Collect" + +-- where +-- (.$) :: Operator -> [Dot PDoc] -> Dot PNode +-- name .$ docs = pnode =<< fmt name docs + +-- fmt :: Operator -> [Dot PDoc] -> Dot PDoc +-- fmt name docs = do +-- docs' <- sequence docs +-- let args = [ x | PDoc x _ <- docs' ] +-- fvs = [ x | PDoc _ x <- docs' ] +-- doc = if simple detail +-- then manifest name +-- else parensIf (needsParens ctx name) +-- $ nest shiftwidth +-- $ sep ( manifest name : args ) +-- return $ PDoc doc (concat fvs) + +-- pnode :: PDoc -> Dot PNode +-- pnode (PDoc doc vs) = do +-- let port = Nothing +-- ident <- genNodeId +-- return $ PNode ident (Leaf (port, doc)) (map (,port) vs) + +-- -- Free variables +-- -- +-- fvF :: Fun aenv t -> [Vertex] +-- fvF = fvOpenFun Empty aenv + +-- fvE :: Exp aenv t -> [Vertex] +-- fvE = fvOpenExp Empty aenv + +-- -- Pretty-printing +-- -- +-- avar :: ArrayVar aenv t -> PDoc +-- avar (Var _ ix) = let (ident, v) = aprj ix aenv +-- in PDoc (pretty v) [Vertex ident Nothing] + +-- aenv' :: Val aenv +-- aenv' = avalToVal aenv + +-- ppA :: HasCallStack => DelayedOpenAcc aenv a -> Dot PDoc +-- ppA (Manifest (Avar ix)) = return (avar ix) +-- ppA acc@Manifest{} = do +-- -- Lift out and draw as a separate node. This can occur with the manifest +-- -- array arguments to permute (defaults array) and stencil[2]. +-- acc' <- prettyDelayedOpenAcc detail app aenv acc +-- v <- mkLabel +-- ident <- mkNode acc' (Just v) +-- return $ PDoc (pretty v) [Vertex ident Nothing] +-- ppA (Delayed _ sh f _) +-- | Shape a <- sh -- identical shape +-- , Just b <- isIdentityIndexing f -- function is `\ix -> b ! ix` +-- , Just Refl <- matchVar a b -- function thus is `\ix -> a ! ix` +-- = ppA $ Manifest $ Avar a +-- ppA (Delayed _ sh f _) = do +-- PDoc d v <- "Delayed" `fmt` [ ppE sh, ppF f ] +-- return $ PDoc (parens d) v + +-- ppB :: forall sh e. HasCallStack +-- => TypeR e +-- -> Boundary aenv (Array sh e) +-- -> Dot PDoc +-- ppB _ Clamp = return (PDoc "clamp" []) +-- ppB _ Mirror = return (PDoc "mirror" []) +-- ppB _ Wrap = return (PDoc "wrap" []) +-- ppB tp (Constant e) = return (PDoc (prettyConst tp e) []) +-- ppB _ (Function f) = ppF f + +-- ppF :: HasCallStack => Fun aenv t -> Dot PDoc +-- ppF = return . uncurry PDoc . (parens . prettyFun aenv' &&& fvF) + +-- ppE :: HasCallStack => Exp aenv t -> Dot PDoc +-- ppE = return . uncurry PDoc . (prettyExp aenv' &&& fvE) + +-- ppD :: String -> Direction -> String -> Operator +-- ppD f LeftToRight k = fromString (f <> "l" <> k) +-- ppD f RightToLeft k = fromString (f <> "r" <> k) + +-- lift :: HasCallStack => DelayedOpenAcc aenv a -> Dot Vertex +-- lift Delayed{} = internalError "expected manifest array" +-- lift (Manifest (Avar (Var _ ix))) = return $ Vertex (fst (aprj ix aenv)) Nothing +-- lift acc = do +-- acc' <- prettyDelayedOpenAcc detail context0 aenv acc +-- ident <- mkNode acc' Nothing +-- return $ Vertex ident Nothing + +-- apply :: Label -> PNode -> PNode +-- apply f (PNode ident x vs) = +-- let x' = case x of +-- Leaf (p,d) -> Leaf (p, pretty f <+> d) +-- Forest ts -> Forest (Leaf (Nothing,pretty f) : ts) +-- in +-- PNode ident x' vs + + +-- -- Pretty print array functions as separate sub-graphs, and return the name of +-- -- the sub-graph as if it can be called like a function. We will add additional +-- -- nodes at the top of the graph to represent the bound variables. +-- -- +-- -- Note: [Edge placement] +-- -- +-- -- If a node belongs to a particular graph, so too must all its edges (and +-- -- vertices). This means that if the subgraph references anything from the +-- -- enclosing environment, we must lift those edges out of this subgraph, +-- -- otherwise the referenced node will be drawn inside of the subgraph. +-- -- +-- prettyDelayedAfun +-- :: HasCallStack +-- => Detail +-- -> Aval aenv +-- -> DelayedOpenAfun aenv afun +-- -> Dot Label +-- prettyDelayedAfun detail aenv afun = do +-- Graph _ ss <- mkSubgraph (go aenv afun) +-- n <- Seq.length <$> gets dotGraph +-- let label = "afun" <> fromString (show (n+1)) +-- outer = collect aenv +-- (lifted,ss') = +-- flip partition ss $ \s -> +-- case s of +-- E (Edge (Vertex ident _) _) -> Set.member ident outer +-- _ -> False +-- -- +-- modify $ \s -> s { dotGraph = dotGraph s Seq.|> Graph label ss' +-- , dotEdges = Seq.fromList [ e | E e <- lifted ] Seq.>< dotEdges s +-- } +-- return label +-- where +-- go :: Aval aenv' -> DelayedOpenAfun aenv' a' -> Dot Graph +-- go aenv' (Abody b) = graphDelayedOpenAcc detail aenv' b +-- go aenv' (Alam lhs f) = do +-- aenv'' <- prettyLambdaALeftHandSide aenv' lhs +-- go aenv'' f + +-- collect :: Aval aenv' -> HashSet NodeId +-- collect Aempty = Set.empty +-- collect (Apush a i _) = Set.insert i (collect a) + +-- prettyLetALeftHandSide +-- :: forall repr aenv aenv'. HasCallStack +-- => NodeId +-- -> Aval aenv +-- -> ALeftHandSide repr aenv aenv' +-- -> Dot (Aval aenv', Label) +-- prettyLetALeftHandSide _ aenv (LeftHandSideWildcard repr) = return (aenv, doc) +-- where +-- doc = case repr of +-- TupRunit -> "()" +-- _ -> "_" +-- prettyLetALeftHandSide ident aenv (LeftHandSideSingle _) = do +-- a <- mkLabel +-- return (Apush aenv ident a, a) +-- prettyLetALeftHandSide ident aenv (LeftHandSidePair lhs1 lhs2) = do +-- (aenv1, d1) <- prettyLetALeftHandSide ident aenv lhs1 +-- (aenv2, d2) <- prettyLetALeftHandSide ident aenv1 lhs2 +-- return (aenv2, "(" <> d1 <> ", " <> d2 <> ")") + +-- prettyLambdaALeftHandSide +-- :: forall repr aenv aenv'. HasCallStack +-- => Aval aenv +-- -> ALeftHandSide repr aenv aenv' +-- -> Dot (Aval aenv') +-- prettyLambdaALeftHandSide aenv (LeftHandSideWildcard _) = return aenv +-- prettyLambdaALeftHandSide aenv (LeftHandSideSingle _) = do +-- a <- mkLabel +-- ident <- genNodeId +-- _ <- mkNode (PNode ident (Leaf (Nothing, pretty a)) []) Nothing +-- return $ Apush aenv ident a +-- prettyLambdaALeftHandSide aenv (LeftHandSidePair lhs1 lhs2) = do +-- aenv1 <- prettyLambdaALeftHandSide aenv lhs1 +-- prettyLambdaALeftHandSide aenv1 lhs2 + +-- -- Display array tuples. This is a little tricky... +-- -- +-- prettyDelayedApair +-- :: forall aenv a1 a2. HasCallStack +-- => Detail +-- -> Aval aenv +-- -> DelayedOpenAcc aenv a1 +-- -> DelayedOpenAcc aenv a2 +-- -> NodeId +-- -> Dot PNode +-- prettyDelayedApair detail aenv a1 a2 ident = do +-- PNode id1 t1 v1 <- prettyElem a1 +-- PNode id2 t2 v2 <- prettyElem a2 +-- modify $ \s -> s { dotEdges = fmap (redirect ident [id1, id2]) (dotEdges s) } +-- return $ PNode ident (forest [t1, t2]) (v1 ++ v2) +-- where +-- prettyElem :: DelayedOpenAcc aenv a -> Dot PNode +-- prettyElem a = replant =<< prettyDelayedOpenAcc detail context0 aenv a + +-- -- Redirect any edges that pointed into one of the nodes now part of this +-- -- tuple, to instead point to the container node. +-- -- +-- redirect :: NodeId -> [NodeId] -> Edge -> Edge +-- redirect new subs edge@(Edge from (Vertex to port)) +-- | to `elem` subs = Edge from (Vertex new port) +-- | otherwise = edge + +-- -- Since we have lifted out any non-leaves into separate nodes, we can +-- -- simply tuple-up all of the elements. +-- -- +-- forest :: [Tree (Maybe Port, Adoc)] -> Tree (Maybe Port, Adoc) +-- forest leaves = Leaf (Nothing, tupled [ align d | Leaf (Nothing,d) <- leaves ]) + + +-- -- Lift out anything that isn't a Leaf node and output it to the graph +-- -- immediately as a new labelled node. +-- -- +-- replant :: PNode -> Dot PNode +-- replant pnode@(PNode ident tree _) = +-- case tree of +-- Leaf (Nothing, _) -> return pnode +-- _ -> do +-- vacuous <- genNodeId +-- a <- mkLabel +-- _ <- mkNode pnode (Just a) +-- return $ PNode vacuous (Leaf (Nothing, pretty a)) [(Vertex ident Nothing, Nothing)] + + +-- -- Pretty printing scalar functions and expressions +-- -- ------------------------------------------------ +-- -- +-- -- This is done with the usual machinery. Note that we rely on knowing that all +-- -- array operations will be lifted out of scalar expressions. This means that we +-- -- don't really need to recurse into the scalar terms to uncover new graph +-- -- nodes. +-- -- + +-- -- Data dependencies +-- -- ----------------- +-- -- +-- -- Return the data-dependencies of the given term. This is just a tree traversal +-- -- to extract all of the free variables. We will draw an edge from each of those +-- -- nodes (vertices) into the current term. +-- -- + +-- fvAvar :: Aval aenv -> ArrayVar aenv a -> [Vertex] +-- fvAvar env (Var _ ix) = [ Vertex (fst $ aprj ix env) Nothing ] + +-- fvOpenFun +-- :: forall env aenv fun. +-- Val env +-- -> Aval aenv +-- -> OpenFun env aenv fun +-- -> [Vertex] +-- fvOpenFun env aenv (Body b) = fvOpenExp env aenv b +-- fvOpenFun env aenv (Lam lhs f) = fvOpenFun env' aenv f +-- where +-- (env', _) = prettyELhs True env lhs + +-- fvOpenExp +-- :: forall env aenv exp. +-- Val env +-- -> Aval aenv +-- -> OpenExp env aenv exp +-- -> [Vertex] +-- fvOpenExp env aenv = fv +-- where +-- fvF :: OpenFun env aenv f -> [Vertex] +-- fvF = fvOpenFun env aenv + +-- fv :: OpenExp env aenv e -> [Vertex] +-- fv (Shape acc) = if cfgIncludeShape then fvAvar aenv acc else [] +-- fv (Index acc i) = concat [ fvAvar aenv acc, fv i ] +-- fv (LinearIndex acc i) = concat [ fvAvar aenv acc, fv i ] +-- -- +-- fv (Let lhs e1 e2) = concat [ fv e1, fvOpenExp env' aenv e2 ] +-- where +-- (env', _) = prettyELhs False env lhs +-- fv Evar{} = [] +-- fv Undef{} = [] +-- fv Const{} = [] +-- fv PrimConst{} = [] +-- fv (PrimApp _ x) = fv x +-- fv (Pair e1 e2) = concat [ fv e1, fv e2] +-- fv Nil = [] +-- fv (VecPack _ e) = fv e +-- fv (VecUnpack _ e) = fv e +-- fv (IndexSlice _ slix sh) = concat [ fv slix, fv sh ] +-- fv (IndexFull _ slix sh) = concat [ fv slix, fv sh ] +-- fv (ToIndex _ sh ix) = concat [ fv sh, fv ix ] +-- fv (FromIndex _ sh ix) = concat [ fv sh, fv ix ] +-- fv (ShapeSize _ sh) = fv sh +-- fv Foreign{} = [] +-- fv (Case e rhs def) = concat [ fv e, concat [ fv c | (_,c) <- rhs ], maybe [] fv def ] +-- fv (Cond p t e) = concat [ fv p, fv t, fv e ] +-- fv (While p f x) = concat [ fvF p, fvF f, fv x ] +-- fv (Coerce _ _ e) = fv e diff --git a/src/Data/Array/Accelerate/Pretty/Graphviz/Monad.hs b/src/Data/Array/Accelerate/Pretty/Graphviz/Monad.hs index 6ae231c21..da85a7fd8 100644 --- a/src/Data/Array/Accelerate/Pretty/Graphviz/Monad.hs +++ b/src/Data/Array/Accelerate/Pretty/Graphviz/Monad.hs @@ -1,5 +1,4 @@ -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE RecordWildCards #-} -- | -- Module : Data.Array.Accelerate.Pretty.Graphviz.Monad -- Copyright : [2015..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Pretty/Graphviz/Type.hs b/src/Data/Array/Accelerate/Pretty/Graphviz/Type.hs index 495ff311e..7e61b4b75 100644 --- a/src/Data/Array/Accelerate/Pretty/Graphviz/Type.hs +++ b/src/Data/Array/Accelerate/Pretty/Graphviz/Type.hs @@ -1,6 +1,3 @@ -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE StandaloneDeriving #-} -{-# LANGUAGE ViewPatterns #-} -- | -- Module : Data.Array.Accelerate.Pretty.Graphviz.Type -- Copyright : [2015..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Pretty/Operation.hs b/src/Data/Array/Accelerate/Pretty/Operation.hs new file mode 100644 index 000000000..92ed8df77 --- /dev/null +++ b/src/Data/Array/Accelerate/Pretty/Operation.hs @@ -0,0 +1,205 @@ +-- | +-- Module : Data.Array.Accelerate.Pretty.Operation +-- Copyright : [2008..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.Pretty.Operation ( + PrettyOp(..), prettyArgs, + prettyAcc, prettyOpenAcc, + prettyAfun, prettyOpenAfun, + prettyGroundR, prettyGroundRWithUniqueness, + Val, Val'(..), val, empty', + prettyVar, prettyVars, + prettyArg, prettyShapeVars, prettyModifier, prettyBuffer, + prettyFun, prettyExp, prettyExp', prettyArrayInstr, + prettyGLhsWithTypes, prettyGLhsWithUniquenessTypes, +) where + +import Data.Array.Accelerate.Pretty.Exp +import Data.Array.Accelerate.Pretty.Type +import Data.Array.Accelerate.AST.Operation +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.AST.IdxSet (IdxSet) +import qualified Data.Array.Accelerate.AST.IdxSet as IdxSet +import Data.Array.Accelerate.Array.Buffer +import Data.Array.Accelerate.Representation.Elt +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Type + +import Prettyprinter +import Data.String + +import Prelude hiding (exp) + +class PrettyOp op where + prettyOp :: op t -> Adoc + + -- Only used in OperationAcc, when printing a PartitionedAcc a cluster is printed using prettyOp. + -- The reason is that PrettyOp (Cluster' op) instance defines prettyOpWithArgs in terms of prettyOp. + -- + prettyOpWithArgs :: Val env -> op t -> Args env t -> Adoc + prettyOpWithArgs env op args = hang 2 $ group $ vsep [annotate Execute "execute", prettyOp op, prettyArgs env args] + +prettyAcc :: PrettyOp op => OperationAcc op () t -> Adoc +prettyAcc = prettyOpenAcc empty' + +prettyAfun :: PrettyOp op => OperationAfun op () f -> Adoc +prettyAfun = prettyOpenAfun empty' + +prettyOpenAfun :: PrettyOp op => Val' benv -> OperationAfun op benv f -> Adoc +prettyOpenAfun env (Abody a) = prettyOpenAcc env a +prettyOpenAfun env (Alam lhs f) = "\\" <> lhs' <+> "->" <> hardline <> indent 2 (prettyOpenAfun env' f) + where + (env', lhs') = prettyGLhsWithTypes env lhs + +prettyOpenAcc :: PrettyOp op => Val' benv -> OperationAcc op benv t -> Adoc +prettyOpenAcc env = \case + Exec op args -> prettyOpWithArgs (val env) op args + Return vars -> hang 2 $ group $ vsep [annotate Statement "return", prettyVars (val env) 10 vars] + Manifest var -> hang 2 $ group $ vsep [annotate Statement "manifest", prettyVar (val env) var] + Compute exp -> hang 2 $ group $ vsep [annotate Statement "compute", prettyExp (val env) exp] + Alet LeftHandSideUnit _ bnd body + | Exec{} <- bnd + -- Only hide the left hand side if the right hand side is a kernel execution + -> prettyOpenAcc env bnd + <> hardline + <> prettyOpenAcc env body + Alet lhs us bnd body + | (env', lhs') <- prettyGLhsWithUniquenessTypes env lhs us + -> hang 2 (group $ vsep [lhs' <+> "=", prettyOpenAcc env bnd]) + <> hardline + <> prettyOpenAcc env' body + Alloc _ tp sh -> hang 2 $ group $ vsep [annotate Statement "alloc", prettyScalarType tp <> "[" <> prettyShapeVars (val env) sh <> "]"] + Use tp n buffer -> hang 2 $ group $ vsep [annotate Statement "use" <+> prettyScalarType tp <> "[" <> pretty n <> "]", prettyBuffer tp n buffer] + Unit var -> hang 2 $ group $ vsep [annotate Statement "unit", prettyVar (val env) var] + Acond c true false + -> group $ vsep + [ hang 2 $ group $ vsep + [ if_ <+> prettyVar (val env) c <+> then_ + , prettyOpenAcc env true + ] + , hang 2 $ group $ vsep + [ else_ + , prettyOpenAcc env false + ] + ] + Awhile us condition step initial + -> "awhile" <+> prettyTupR (const prettyGroundRWithUniqueness) 10 (groundsRWithUniquenesses (mapTupR varType initial) us) + <> hardline <> hang 4 (" ( " <> prettyOpenAfun env condition) + <> hardline <> " )" + <> hardline <> hang 4 (" ( " <> prettyOpenAfun env step) + <> hardline <> " )" + <> hardline <> indent 2 (prettyVars (val env) 10 initial) + Atrace msg t -> + hang 2 (group $ vsep [annotate Statement "atrace", prettyText msg, prettyArrayDescriptor (val env) t]) + Aassert msg g -> + hang 2 (group $ vsep [annotate Statement "assert", prettyText msg, prettyExp (val env) g]) + Aassume g -> + hang 2 (group $ vsep [annotate Statement "assume", prettyExp (val env) g]) + Fence set next -> + hang 2 (group $ vsep [annotate Statement "fence", prettyIdxSet (val env) set]) + <> hardline + <> prettyOpenAcc env next + +prettyArrayDescriptor :: Val env -> ArrayDescriptors env t -> Adoc +prettyArrayDescriptor env = prettyTupR (\p (ArrayDescriptor _ sh buffer) -> prettyShapeVars env sh <+> " " <+> prettyVars env p buffer) 0 + +prettyArgs :: Val benv -> Args benv f -> Adoc +prettyArgs env args = tupled $ map (\(Exists a) -> prettyArg env a) $ argsToList args + +prettyArg :: Val benv -> Arg benv t -> Adoc +prettyArg env (ArgVar vars) = prettyVars env 10 vars +prettyArg env (ArgExp e) = prettyExp env e +prettyArg env (ArgFun f) = prettyFun env f +prettyArg env (ArgArray m _ sh buffers) = group $ vsep [prettyModifier m, "(" <> prettyShapeVars env sh <> ")", prettyVars env 0 buffers] + +prettyModifier :: Modifier m -> Adoc +prettyModifier In = annotate Modifier "in" +prettyModifier Out = annotate Modifier "out" +prettyModifier Mut = annotate Modifier "mut" + +prettyBuffer :: ScalarType tp -> Int -> Buffer tp -> Adoc +prettyBuffer _ 0 _ = "[]" +prettyBuffer tp n b = align $ group $ "[ " <> vcat (mapTail (", " <>) $ map (fromString . showElt (TupRsingle tp)) $ bufferToList tp n b) <> " ]" + +mapTail :: (a -> a) -> [a] -> [a] +mapTail f (x:xs) = x : map f xs +mapTail _ [] = [] + +-- Expressions +prettyFun :: Val env -> Fun env f -> Adoc +prettyFun env = prettyPreOpenFun (prettyArrayInstr env) Empty + +prettyExp :: Val env -> Exp env t -> Adoc +prettyExp env = prettyPreOpenExp context0 (prettyArrayInstr env) Empty + +prettyExp' :: Context -> Val env -> Exp env t -> Adoc +prettyExp' ctx env = prettyPreOpenExp ctx (prettyArrayInstr env) Empty + +prettyArrayInstr :: Val env -> PrettyArrayInstr (ArrayInstr env) +prettyArrayInstr env context (Index arr) ix + = parensIf (ctxPrecedence context < 9) + $ sep [ prj (varIdx arr) env + , group (sep ["!!", ix context']) ] + where + context' = Context L R 9 +prettyArrayInstr env _ (Parameter var) _ = prj (varIdx var) env + +-- prettyGLhs :: Val' env -> GLeftHandSide t env env' -> (Val' env', Adoc) +-- prettyGLhs = prettyLhs' push Nothing False + +prettyGLhsWithTypes :: Val' env -> GLeftHandSide t env env' -> (Val' env', Adoc) +prettyGLhsWithTypes = prettyLhs' push (Just $ \(Exists t) -> prettyGroundR t) False + +prettyGLhsWithUniquenessTypes :: Val' env -> GLeftHandSide t env env' -> Uniquenesses t -> (Val' env', Adoc) +prettyGLhsWithUniquenessTypes env lhs us + = prettyLhs' + (\(GroundRWithUniqueness g _) -> push g) + (Just $ \(Exists t) -> prettyGroundRWithUniqueness t) + False + env + (lhsWithUniquesses lhs us) + + +-- Environment +data Val' env = Val' (Val env) Int Int + +push :: GroundR t -> Val' env -> (Val' (env, t), Adoc) +push (GroundRbuffer _) (Val' env b e) = (Val' (Push env d) (b + 1) e, d) + where d = "b" <> pretty b +push (GroundRscalar _) (Val' env b e) = (Val' (Push env d) b (e + 1), d) + where d = "e" <> pretty e + +empty' :: Val' () +empty' = Val' Empty 0 0 + +val :: Val' env -> Val env +val (Val' env _ _) = env + +-- Variables +prettyVar :: Val env -> Var s env t -> Adoc +prettyVar env (Var _ ix) = prj ix env + +prettyVars :: forall env s t. Val env -> Precedence -> Vars s env t -> Adoc +prettyVars env = prettyTupR $ const $ prettyVar env + +prettyShapeVars :: Val env -> Vars s env sh -> Adoc +prettyShapeVars _ TupRunit = "Z" +prettyShapeVars env vars = encloseSep "Z :. " "" " :. " $ map (\(Exists v) -> prettyVar env v) $ flattenTupR vars + +prettyIdxSet :: Val env -> IdxSet env -> Adoc +prettyIdxSet env set = "{" <> hsep (map (\(Exists idx) -> prj idx env) $ IdxSet.toList set) <> "}" + +-- Types +prettyGroundR :: GroundR t -> Adoc +prettyGroundR (GroundRscalar tp) = prettyScalarType tp +prettyGroundR (GroundRbuffer tp) = "[" <> prettyScalarType tp <> "]" + +prettyGroundRWithUniqueness :: GroundRWithUniqueness t -> Adoc +prettyGroundRWithUniqueness (GroundRWithUniqueness tp Unique) = prettyGroundR tp <> "¹" +prettyGroundRWithUniqueness (GroundRWithUniqueness tp Shared) = prettyGroundR tp <> "" diff --git a/src/Data/Array/Accelerate/Pretty/Partitioned.hs b/src/Data/Array/Accelerate/Pretty/Partitioned.hs new file mode 100644 index 000000000..ade685a47 --- /dev/null +++ b/src/Data/Array/Accelerate/Pretty/Partitioned.hs @@ -0,0 +1,285 @@ +{-# OPTIONS_GHC -Wno-orphans #-} +-- | +-- Module : Data.Array.Accelerate.Pretty.Operation +-- Copyright : [2008..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.Pretty.Partitioned ({- instance PrettyOp (Cluster op) -}) where + +import Data.Array.Accelerate.Pretty.Exp hiding (Val(..), prj) +import Data.Array.Accelerate.Pretty.Operation +import Data.Array.Accelerate.AST.Partitioned +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.Type + +import Prettyprinter + +import Prelude hiding (exp) +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Representation.Shape +import Data.Array.Accelerate.Representation.Array +import Data.Array.Accelerate.AST.Var (varsType) + +instance (PrettyOp op, SetOpIndices op) => PrettyOp (Clustered op) where + prettyOp :: PrettyOp op => Clustered op t -> Adoc + prettyOp (Clustered c _) = prettyOp c + prettyOpWithArgs env (Clustered c _) = prettyOpWithArgs env c + +instance (PrettyOp op, SetOpIndices op) => PrettyOp (Cluster op) where + prettyOp _ = "cluster" + + prettyOpWithArgs env cluster args + {- = prettyFlatCluster env flat + where + flat = toFlatCluster cluster args-} + | isSingle cluster = body + | otherwise = annotate Execute "execute" <+> "{" <> line + <> indent 2 body <> line <> "}" + where + body = fst $ prettyCluster True cluster (mapArgs (toPrettyArg env) args) 0 + +isSingle :: Cluster op t -> Bool +isSingle SingleOp{} = True +isSingle _ = False + +data PrettyArg t where + PrettyArgArray + :: Modifier m + -> Adoc -- shape variables + -> TupR Adoc' (Buffers e) -- buffer variables + -> PrettyArg (m sh e) + + PrettyArgVarShape + :: Adoc -- Pretty printed as normal variables + -> Adoc -- Pretty printed as shape + -> PrettyArg t + + PrettyArgOther + :: Adoc + -> PrettyArg t + +type PrettyArgs = PreArgs PrettyArg + +newtype Adoc' t = Adoc' Adoc + +prettyCluster :: Bool -> PrettyOp op => Cluster op t -> PrettyArgs t -> Int -> (Adoc, Int) +prettyCluster topLevel (SingleOp singleOp _) args fresh = (prettySingleOp topLevel singleOp args, fresh) +prettyCluster _ (Fused fusion clusterL clusterR) args fresh + | (leftArgs, rightArgs, horizontals, verticals, diagonals, fresh') <- splitEnv fusion args fresh + , (leftDoc, fresh'') <- prettyCluster False clusterL leftArgs fresh' + , (rightDoc, fresh''') <- prettyCluster False clusterR rightArgs fresh'' + = ( leftDoc <> line <> + prettyFuseList (annotate Statement "fused horizontally") horizontals <> + prettyFuseList (annotate Statement "fused vertically") verticals <> + prettyFuseList (annotate Statement "fused diagonally") diagonals <> + rightDoc + , fresh''') + +prettySingleOp :: PrettyOp op => Bool -> SingleOp op t -> PrettyArgs t -> Adoc +prettySingleOp topLevel (Single op opArgs) args = + hang 2 $ group $ vsep $ + [ annotate Execute "execute" | topLevel ] + ++ + [ + prettyOp op, + tupled $ map (\(Exists a) -> prettyPrettyArg $ prettyClusterArg args a) $ argsToList opArgs + ] + +prettyClusterArg :: forall f t. PrettyArgs f -> ClusterArg (FunToEnv f) t -> PrettyArg t +prettyClusterArg pArgs = \case + ClusterArgSingle idx -> funToEnvPrj pArgs idx + ClusterArgArray m _ _ buffers -> prettyBuffers m buffers + where + prettyBuffers :: Modifier m -> ClusterArgBuffers (FunToEnv f) m sh e -> PrettyArg (m sh e) + prettyBuffers _ (ClusterArgBuffersDead _ idx) = case funToEnvPrj pArgs idx of + PrettyArgVarShape _ sh -> + PrettyArgArray Out sh (TupRsingle $ Adoc' "_") + PrettyArgOther sh -> + PrettyArgArray Out sh (TupRsingle $ Adoc' "_") + prettyBuffers _ (ClusterArgBuffersLive _ idx) = funToEnvPrj pArgs idx + prettyBuffers m (ClusterArgBuffersPair l r) = prettyBuffers m l `prettyPairArg` prettyBuffers m r + +toPrettyArg :: Val env -> Arg env t -> PrettyArg t +toPrettyArg env (ArgArray m _ sh buffers) = PrettyArgArray m (prettyShapeVars env sh) (mapTupR (Adoc' . prettyVar env) buffers) +toPrettyArg env arg@(ArgVar vars) + | Just _ <- typeShape $ varsType vars + = PrettyArgVarShape (prettyArg env arg) (prettyShapeVars env vars) +toPrettyArg env arg = PrettyArgOther $ prettyArg env arg + +prettyPrettyArg :: PrettyArg t -> Adoc +prettyPrettyArg (PrettyArgArray m sh buffers) = group $ vsep [prettyModifier m, "(" <> sh <> ")", prettyTupR (\_ (Adoc' doc) -> doc) 0 buffers] +prettyPrettyArg (PrettyArgVarShape doc _) = doc +prettyPrettyArg (PrettyArgOther doc) = doc + +prettyPairArg :: PrettyArg (f left) -> PrettyArg (f right) -> PrettyArg (f (left, right)) +prettyPairArg (PrettyArgArray m sh buffersL) (PrettyArgArray _ _ buffersR) = PrettyArgArray m sh $ TupRpair buffersL buffersR +prettyPairArg _ _ = PrettyArgOther "?" + +splitEnv :: Fusion largs rargs t -> PrettyArgs t -> Int -> (PrettyArgs largs, PrettyArgs rargs, [Adoc], [Adoc], [Adoc], Int) +splitEnv EmptyF _ fresh = (ArgsNil, ArgsNil, [], [], [], fresh) +splitEnv (Vertical _ next) (a :>: as) fresh = + let + (argsL, argsR, horizontals, verticals, diagonals, fresh') = splitEnv next as (fresh + 1) + sh = case a of + PrettyArgVarShape _ sh' -> sh' + _ -> "?" + buffer = "%" <> viaShow fresh + buffers = TupRsingle $ Adoc' buffer + in + ( PrettyArgArray Out sh buffers :>: argsL + , PrettyArgArray In sh buffers :>: argsR + , horizontals + , buffer : verticals + , diagonals + , fresh' + ) +splitEnv (Horizontal next) (a :>: as) fresh = + let + (argsL, argsR, horizontals, verticals, diagonals, fresh') = splitEnv next as fresh + buffer = case a of + PrettyArgArray _ _ (TupRsingle (Adoc' b)) -> b + _ -> prettyPrettyArg a + in + ( a :>: argsL + , a :>: argsR + , buffer : horizontals + , verticals + , diagonals + , fresh' + ) +splitEnv (Diagonal next) (a :>: as) fresh = + let + (argsL, argsR, horizontals, verticals, diagonals, fresh') = splitEnv next as fresh + a' = case a of + PrettyArgArray _ sh bs -> PrettyArgArray In sh bs + _ -> PrettyArgOther "?" + buffer = case a of + PrettyArgArray _ _ (TupRsingle (Adoc' b)) -> b + _ -> prettyPrettyArg a + in + ( a :>: argsL + , a' :>: argsR + , horizontals + , verticals + , buffer : diagonals + , fresh' + ) +splitEnv (IntroI1 next) as fresh = splitEnv (IntroL next) as fresh +splitEnv (IntroO1 next) as fresh = splitEnv (IntroL next) as fresh +splitEnv (IntroI2 next) as fresh = splitEnv (IntroR next) as fresh +splitEnv (IntroO2 next) as fresh = splitEnv (IntroR next) as fresh +splitEnv (IntroL next) (a :>: as) fresh = + let + (argsL, argsR, horizontals, verticals, diagonals, fresh') = splitEnv next as fresh + in + (a :>: argsL, argsR, horizontals, verticals, diagonals, fresh') +splitEnv (IntroR next) (a :>: as) fresh = + let + (argsL, argsR, horizontals, verticals, diagonals, fresh') = splitEnv next as fresh + in + (argsL, a :>: argsR, horizontals, verticals, diagonals, fresh') + +prettyFuseList :: Adoc -> [Adoc] -> Adoc +prettyFuseList _ [] = "" +prettyFuseList name docs = (hang 2 $ group $ vsep $ [name, tupled docs]) <> line + +{- TODO WALL: DEAD CODE +prettyFlatCluster :: PrettyOp op => Val env -> FlatCluster op env -> Adoc +prettyFlatCluster env (FlatCluster _ idxLhs sizes directions localR localLHS ops) = + annotate Execute "execute" <+> "{" <> line + <> indent 2 (forIndices <> localBuffers <> vsep ops') <> line <> "}" + where + (forIndices, idxEnv, _) = prettyIdxLHS env Pretty.Empty 0 idxLhs sizes directions + (env', _) = pushLocalBuffers env 0 localLHS + (localBuffers, _) = prettyLocalLHS 0 localLHS localR + ops' = prettyFlatOps False env' idxEnv ops + +prettyIdxLHS :: Val env0 -> Val env -> Int -> ELeftHandSide sh env env' -> GroundVars env0 sh -> TupR LoopDirection sh -> (Adoc, Val env', Int) +prettyIdxLHS _ env fresh (LeftHandSideWildcard _) _ _ = (mempty, env, fresh) +prettyIdxLHS env0 env fresh (LeftHandSideSingle _) (TupRsingle sz) (TupRsingle direction) = + (doc, env `Pretty.Push` var, fresh + 1) + where + doc = for_ <+> var <+> in_ <+> "0 .. " <> prettyVar env0 sz <> direction' <> hardline + var = "i" <> viaShow fresh + direction' = case direction of + LoopAny -> mempty + LoopMonotone -> " [monotone]" + LoopAscending -> " [ascending]" + LoopDescending -> " [descending]" +prettyIdxLHS env0 env fresh (LeftHandSidePair lhs1 lhs2) (TupRpair sz1 sz2) (TupRpair d1 d2) + | (doc1, env1, fresh1) <- prettyIdxLHS env0 env fresh lhs1 sz1 d1 + , (doc2, env2, fresh2) <- prettyIdxLHS env0 env1 fresh1 lhs2 sz2 d2 + = (doc1 <> doc2, env2, fresh2) +prettyIdxLHS _ _ _ _ _ _ = internalError "Tuple mismatch" + +pushLocalBuffers :: Val env -> Int -> GLeftHandSide t env env' -> (Val env', Int) +pushLocalBuffers env fresh (LeftHandSideWildcard _) = (env, fresh) +pushLocalBuffers env fresh (LeftHandSideSingle _) = + (env `Pretty.Push` ("%" <> viaShow fresh), fresh + 1) +pushLocalBuffers env fresh (LeftHandSidePair lhs1 lhs2) + | (env', fresh') <- pushLocalBuffers env fresh lhs1 + = pushLocalBuffers env' fresh' lhs2 + +prettyLocalLHS :: Int -> GLeftHandSide t env env' -> TupR LocalBufferR t -> (Adoc, Int) +prettyLocalLHS fresh (LeftHandSideWildcard _) _ = (mempty, fresh) +prettyLocalLHS fresh (LeftHandSideSingle _) (TupRsingle (LocalBufferR tp depth)) = + ( let_ <> " %" <> viaShow fresh <> ": [" <> prettyScalarType tp <> "] " <> comment <> hardline + , fresh + 1) + where + comment + | depth == 0 = "-- outside for loops" + | otherwise = "-- in for i" <> viaShow (depth - 1) +prettyLocalLHS fresh (LeftHandSidePair lhs1 lhs2) (TupRpair r1 r2) + | (d1, fresh1) <- prettyLocalLHS fresh lhs1 r1 + , (d2, fresh2) <- prettyLocalLHS fresh1 lhs2 r2 + = (d1 <> d2, fresh2) +prettyLocalLHS _ _ _ = internalError "Tuple mismatch" + +prettyFlatOps :: PrettyOp op => Bool -> Val env -> Val idxEnv -> FlatOps op env idxEnv -> [Adoc] +prettyFlatOps single env idxEnv = \case + FlatOpsNil -> [] + FlatOpsOp op ops -> + prettyFlatOp single env idxEnv op : prettyFlatOps single env idxEnv ops + FlatOpsBind depth lhs expr ops + | (idxEnv', lhs') <- prettyLhs False 'i' idxEnv lhs -> + ( ( if depth == 0 + then "-- Backpermute outside all for loops" + else "-- Backpermute in for loop of i" <> viaShow (depth - 1) <> hardline) + <> let_ <+> lhs' <+> "=" <+> prettyPreOpenExp context0 (prettyArrayInstr env) idxEnv expr) + : prettyFlatOps single env idxEnv' ops + +prettyFlatOp :: PrettyOp op => Bool -> Val env -> Val idxEnv -> FlatOp op env idxEnv -> Adoc +prettyFlatOp single env idxEnv (FlatOp op args idxArgs) = + hang 2 $ group $ vsep $ + [ annotate Execute "execute" | single ] + ++ + [ + prettyOp op, + tupled $ prettyArgsWithIdx env idxEnv args idxArgs + ] + +prettyArgsWithIdx :: Val env -> Val idxEnv -> Args env args -> IdxArgs idxEnv args -> [Adoc] +prettyArgsWithIdx _ _ ArgsNil _ = [] +prettyArgsWithIdx env idxEnv (a :>: as) (i :>: is) + = prettyArgWithIdx env idxEnv a i : prettyArgsWithIdx env idxEnv as is + +prettyArgWithIdx :: Val env -> Val idxEnv -> Arg env arg -> IdxArg idxEnv arg -> Adoc +prettyArgWithIdx env idxEnv arg idxArg + | ArgArray{} <- arg = case idxArg of + IdxArgIdx depth idx -> + let + idx' = map (\(Exists var) -> prettyVar idxEnv var) $ flattenTupR idx + comment + | depth == 0 = "{- outside for loops -}" + | otherwise = "{- in for i" <> viaShow (depth - 1) <> " -}" + in arg' <> " @ " <> tupled idx' <+> comment + IdxArgNone -> arg' <> " @ ?" + | otherwise = arg' + where + arg' = prettyArg env arg +-} diff --git a/src/Data/Array/Accelerate/Pretty/Print.hs b/src/Data/Array/Accelerate/Pretty/Print.hs index 965aecd99..34e309a1a 100644 --- a/src/Data/Array/Accelerate/Pretty/Print.hs +++ b/src/Data/Array/Accelerate/Pretty/Print.hs @@ -1,14 +1,4 @@ -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE PatternGuards #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE RecordWildCards #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE ViewPatterns #-} -- | -- Module : Data.Array.Accelerate.Pretty.Print -- Copyright : [2008..2020] The Accelerate Team @@ -24,8 +14,8 @@ module Data.Array.Accelerate.Pretty.Print ( PrettyAcc, ExtractAcc, prettyPreOpenAcc, prettyPreOpenAfun, - prettyOpenExp, prettyExp, - prettyOpenFun, prettyFun, + prettyExp', prettyExp, prettyOpenExp, + prettyFun, prettyOpenFun, prettyArray, prettyConst, prettyELhs, @@ -34,7 +24,7 @@ module Data.Array.Accelerate.Pretty.Print ( -- ** Configuration PrettyConfig(..), configPlain, - configWithHash, + -- configWithHash, -- ** Internals Adoc, @@ -48,7 +38,7 @@ module Data.Array.Accelerate.Pretty.Print ( shiftwidth, context0, app, - manifest, delayed, + functionName, primOperator, isInfix, prj, sizeEnv, @@ -56,21 +46,17 @@ module Data.Array.Accelerate.Pretty.Print ( ) where import Data.Array.Accelerate.AST hiding ( Direction ) -import Data.Array.Accelerate.AST.Idx import Data.Array.Accelerate.AST.LeftHandSide import Data.Array.Accelerate.AST.Var +import Data.Array.Accelerate.Pretty.Exp import Data.Array.Accelerate.Representation.Array import Data.Array.Accelerate.Representation.Elt import Data.Array.Accelerate.Representation.Stencil -import Data.Array.Accelerate.Representation.Tag import Data.Array.Accelerate.Representation.Type import Data.Array.Accelerate.Sugar.Foreign -import Data.Array.Accelerate.Type import qualified Data.Array.Accelerate.AST as AST -import qualified Data.Array.Accelerate.Analysis.Hash as Hash -import qualified Data.Array.Accelerate.Trafo.Delayed as Delayed +-- import qualified Data.Array.Accelerate.Trafo.Delayed as Delayed -import Data.Char import Data.String import Prettyprinter import Prettyprinter.Render.Terminal @@ -85,42 +71,18 @@ type PrettyAcc acc = type ExtractAcc acc = forall aenv a. acc aenv a -> PreOpenAcc acc aenv a -type Adoc = Doc Keyword - -data Keyword - = Statement -- do | case of | let in - | Conditional -- if then else - | Manifest -- collective operations (kernel functions) - | Delayed -- fused operators - deriving (Eq, Show) - -let_, in_ :: Adoc -let_ = annotate Statement "let" -in_ = annotate Statement "in" - -case_, of_ :: Adoc -case_ = annotate Statement "case" -of_ = annotate Statement "of" - -if_, then_, else_ :: Adoc -if_ = annotate Statement "if" -then_ = annotate Statement "then" -else_ = annotate Statement "else" - -manifest :: Operator -> Adoc -manifest = annotate Manifest . opName - -delayed :: Operator -> Adoc -delayed = annotate Delayed . opName +functionName :: Operator -> Adoc +functionName = annotate FunctionName . opName ansiKeyword :: Keyword -> AnsiStyle -ansiKeyword Statement = colorDull Yellow -ansiKeyword Conditional = colorDull Yellow -ansiKeyword Manifest = color Blue -ansiKeyword Delayed = color Green +ansiKeyword Statement = colorDull Yellow +ansiKeyword Conditional = colorDull Yellow +ansiKeyword FunctionName = color Blue +ansiKeyword Execute = color Blue +ansiKeyword Modifier = colorDull Green -- Configuration for the pretty-printing functions -data PrettyConfig acc +newtype PrettyConfig acc = PrettyConfig { confOperator :: forall aenv arrs. PreOpenAcc acc aenv arrs -> String @@ -129,15 +91,15 @@ data PrettyConfig acc configPlain :: PrettyConfig acc configPlain = PrettyConfig { confOperator = const fromString } -configWithHash :: PrettyConfig Delayed.DelayedOpenAcc -configWithHash = - PrettyConfig - { confOperator = \pacc name -> - let hashval = Hash.hashPreOpenAccWith - (Hash.defaultHashOptions { Hash.perfect = False }) - Delayed.encodeDelayedOpenAcc - pacc - in fromString (name ++ "_" ++ show hashval) } +-- configWithHash :: PrettyConfig Delayed.DelayedOpenAcc +-- configWithHash = +-- PrettyConfig +-- { confOperator = \pacc name -> +-- let hashval = Hash.hashPreOpenAccWith +-- (Hash.defaultHashOptions { Hash.perfect = False }) +-- Delayed.encodeDelayedOpenAcc +-- pacc +-- in fromString (name ++ "_" ++ show hashval) } -- Array computations @@ -150,7 +112,7 @@ prettyPreOpenAfun -> Val aenv -> PreOpenAfun acc aenv f -> Adoc -prettyPreOpenAfun config prettyAcc aenv0 = next (pretty '\\') aenv0 +prettyPreOpenAfun config prettyAcc = next (pretty '\\') where next :: Adoc -> Val aenv' -> PreOpenAfun acc aenv' f' -> Adoc next vs aenv (Abody body) = @@ -174,10 +136,7 @@ prettyPreOpenAcc config ctx prettyAcc extractAcc aenv pacc = Alet{} -> prettyAlet config ctx prettyAcc extractAcc aenv pacc Apair{} -> prettyAtuple config ctx prettyAcc extractAcc aenv pacc Anil -> "()" - Apply _ f a -> apply - where - op = Operator ">->" Infix L 1 - apply = sep [ ppAF f, group (sep [opName op, ppA a]) ] + Manifest as -> ppN "manifest" .$ [ ppA as ] Acond p t e -> flatAlt multi single where @@ -192,7 +151,8 @@ prettyPreOpenAcc config ctx prettyAcc extractAcc aenv pacc = , hang shiftwidth (sep [ then_, t' ]) , hang shiftwidth (sep [ else_, e' ]) ] - + Aassert msg cond a -> ppN "assert" .$ [ prettyText msg, ppE cond, ppA a] + Aassume cond a -> ppN "assume" .$ [ ppE cond, ppA a] Atrace (Message _ _ msg) as bs -> ppN "atrace" .$ [ fromString (show msg), ppA as, ppA bs ] Aforeign _ ff _ a -> ppN "aforeign" .$ [ pretty (strForeign ff), ppA a ] Awhile p f a -> ppN "awhile" .$ [ ppAF p, ppAF f, ppA a ] @@ -212,7 +172,7 @@ prettyPreOpenAcc config ctx prettyAcc extractAcc aenv pacc = Scan d f (Just z) a -> ppD "scan" d "" .$ [ ppF f, ppE z, ppA a ] Scan d f Nothing a -> ppD "scan" d "1" .$ [ ppF f, ppA a ] Scan' d f z a -> ppD "scan" d "'" .$ [ ppF f, ppE z, ppA a ] - Permute f d p s -> ppN "permute" .$ [ ppF f, ppA d, ppF p, ppA s ] + Permute f d s -> ppN "permute" .$ [ ppF f' | Just f' <- [f]] ++ [ ppA d, ppA s ] Backpermute _ sh f a -> ppN "backpermute" .$ [ ppE sh, ppF f, ppA a ] Stencil s _ f b a -> ppN "stencil" .$ [ ppF f, ppB (stencilEltR s) b, ppA a ] Stencil2 s1 s2 _ f b1 a1 b2 a2 -> ppN "stencil2" .$ [ ppF f, ppB (stencilEltR s1) b1, ppA a1, ppB (stencilEltR s2) b2, ppA a2 ] @@ -220,7 +180,7 @@ prettyPreOpenAcc config ctx prettyAcc extractAcc aenv pacc = infixr 0 .$ f .$ xs = parensIf (needsParens ctx f) - $ hang shiftwidth (sep (manifest f : xs)) + $ hang shiftwidth (sep (functionName f : xs)) ppN :: String -> Operator ppN = confOperator config pacc @@ -232,10 +192,10 @@ prettyPreOpenAcc config ctx prettyAcc extractAcc aenv pacc = ppAF = parens . prettyPreOpenAfun config prettyAcc aenv ppE :: Exp aenv t -> Adoc - ppE = prettyOpenExp app Empty aenv + ppE = prettyExp' app aenv ppF :: Fun aenv t -> Adoc - ppF = parens . prettyOpenFun Empty aenv + ppF = parens . prettyFun aenv ppB :: forall sh e. TypeR e @@ -305,7 +265,7 @@ prettyAtuple -> Val aenv -> PreOpenAcc acc aenv arrs -> Adoc -prettyAtuple config ctx prettyAcc extractAcc aenv0 acc = case collect acc of +prettyAtuple config ctx prettyAcc extractAcc aenv0 acc = case collect acc [] of Nothing -> align $ ppPair acc Just tup -> case tup of @@ -317,46 +277,15 @@ prettyAtuple config ctx prettyAcc extractAcc aenv0 acc = case collect acc of "(" <> ppPair (extractAcc a1) <> "," <+> prettyAcc config context0 aenv0 a2 <> ")" ppPair a = prettyPreOpenAcc config context0 prettyAcc extractAcc aenv0 a - collect :: PreOpenAcc acc aenv arrs' -> Maybe [Adoc] - collect Anil = Just [] - collect (Apair a1 a2) - | Just tup <- collect $ extractAcc a1 - = Just $ tup ++ [prettyAcc config app aenv0 a2] - collect _ = Nothing + collect :: PreOpenAcc acc aenv arrs' -> [Adoc] -> Maybe [Adoc] + collect Anil accum = Just accum + collect (Apair a1 a2) accum = collect (extractAcc a1) (prettyAcc config app aenv0 a2 : accum) + collect _ _ = Nothing -- TODO: Should we also print the types of the declared variables? And the types of wildcards? prettyALhs :: Bool -> Val env -> LeftHandSide s arrs env env' -> (Val env', Adoc) prettyALhs requiresParens = prettyLhs requiresParens 'a' -prettyELhs :: Bool -> Val env -> LeftHandSide s arrs env env' -> (Val env', Adoc) -prettyELhs requiresParens = prettyLhs requiresParens 'x' - -prettyLhs :: forall s env env' arrs. Bool -> Char -> Val env -> LeftHandSide s arrs env env' -> (Val env', Adoc) -prettyLhs requiresParens x env0 lhs = case collect lhs of - Nothing -> ppPair lhs - Just (env1, tup) -> - case tup of - [] -> (env1, "()") - _ -> (env1, parensIf requiresParens (pretty 'T' <> pretty (length tup) <+> align (sep tup))) - where - ppPair :: LeftHandSide s arrs' env env'' -> (Val env'', Adoc) - ppPair LeftHandSideUnit = (env0, "()") - ppPair LeftHandSideWildcard{} = (env0, "_") - ppPair LeftHandSideSingle{} = (env0 `Push` v, v) - where - v = pretty x <> pretty (sizeEnv env0) - ppPair (LeftHandSidePair a b) = (env2, tupled [doc1, doc2]) - where - (env1, doc1) = ppPair a - (env2, doc2) = prettyLhs False x env1 b - - collect :: LeftHandSide s arrs' env env'' -> Maybe (Val env'', [Adoc]) - collect (LeftHandSidePair l1 l2) - | Just (env1, tup ) <- collect l1 - , (env2, doc2) <- prettyLhs True x env1 l2 = Just (env2, tup ++ [doc2]) - collect (LeftHandSideWildcard TupRunit) = Just (env0, []) - collect _ = Nothing - prettyArray :: ArrayR (Array sh e) -> Array sh e -> Adoc prettyArray aR@(ArrayR _ eR) = parens . fromString . showArray (showsElt eR) aR @@ -365,119 +294,36 @@ prettyArray aR@(ArrayR _ eR) = parens . fromString . showArray (showsElt eR) aR -- ------------------ prettyFun :: Val aenv -> Fun aenv f -> Adoc -prettyFun = prettyOpenFun Empty +prettyFun aenv = prettyPreOpenFun (prettyArrayInstr aenv) Empty -prettyExp :: Val aenv -> Exp aenv t -> Adoc -prettyExp = prettyOpenExp context0 Empty +prettyOpenFun :: Val env -> Val aenv -> OpenFun env aenv f -> Adoc +prettyOpenFun env aenv = prettyPreOpenFun (prettyArrayInstr aenv) env -prettyOpenFun - :: forall env aenv f. - Val env - -> Val aenv - -> OpenFun env aenv f - -> Adoc -prettyOpenFun env0 aenv = next (pretty '\\') env0 +prettyExp :: Val aenv -> Exp aenv t -> Adoc +prettyExp aenv = prettyPreOpenExp context0 (prettyArrayInstr aenv) Empty + +prettyOpenExp :: Context -> Val env -> Val aenv -> OpenExp env aenv t -> Adoc +prettyOpenExp ctx env aenv = prettyPreOpenExp ctx (prettyArrayInstr aenv) env + +prettyExp' :: Context -> Val aenv -> Exp aenv t -> Adoc +prettyExp' ctx aenv = prettyPreOpenExp ctx (prettyArrayInstr aenv) Empty + +prettyArrayInstr :: Val aenv -> PrettyArrayInstr (ArrayInstr aenv) +prettyArrayInstr aenv context (Shape arr) _ + = parensIf (ctxPrecedence context < 10) + $ "shape" <+> prettyArrayVar aenv arr +prettyArrayInstr aenv context (Index arr) ix + = parensIf (ctxPrecedence context < 9) + $ sep [ prettyArrayVar aenv arr + , group (sep ["!", ix context']) ] where - next :: Adoc -> Val env' -> OpenFun env' aenv f' -> Adoc - next vs env (Body body) - -- PrimApp f x <- body - -- , op <- primOperator f - -- , isInfix op - -- , Tuple (NilTup `SnocTup` a `SnocTup` b) <- x - -- , Var (SuccIdx ZeroIdx) <- a - -- , Var ZeroIdx <- b - -- = opName op -- surrounding context will add parens - -- - = hang shiftwidth (sep [ vs <> "->" - , prettyOpenExp context0 env aenv body]) - next vs env (Lam lhs lam) = - let (env', lhs') = prettyELhs True env lhs - in next (vs <> lhs' <> space) env' lam - -prettyOpenExp - :: forall env aenv t. - Context - -> Val env - -> Val aenv - -> OpenExp env aenv t - -> Adoc -prettyOpenExp ctx env aenv exp = - case exp of - Evar (Var _ idx) -> prj idx env - Let{} -> prettyLet ctx env aenv exp - PrimApp f x - | a `Pair` b <- x -> ppF2 op (ppE a) (ppE b) - | otherwise -> ppF1 op' (ppE x) - where - op = primOperator f - op' = isInfix op ? (Operator (parens (opName op)) App L 10, op) - -- - PrimConst c -> prettyPrimConst c - Const tp c -> prettyConst (TupRsingle tp) c - Pair{} -> prettyTuple ctx env aenv exp - Nil -> "()" - VecPack _ e -> ppF1 "pack" (ppE e) - VecUnpack _ e -> ppF1 "unpack" (ppE e) - Case x xs d -> prettyCase env aenv x xs d - Cond p t e -> flatAlt multi single - where - p' = ppE p context0 - t' = ppE t context0 - e' = ppE e context0 - -- - single = parensIf (needsParens ctx (Operator "?:" Infix N 0)) - $ sep [ p', pretty '?', t', pretty ':', e' ] - multi = hang 3 - $ vsep [ if_ <+> p' - , hang shiftwidth (sep [ then_, t' ]) - , hang shiftwidth (sep [ else_, e' ]) ] - -- - IndexSlice _ slix sh -> ppF2 "indexSlice" (ppE slix) (ppE sh) - IndexFull _ slix sl -> ppF2 "indexFull" (ppE slix) (ppE sl) - ToIndex _ sh ix -> ppF2 "toIndex" (ppE sh) (ppE ix) - FromIndex _ sh ix -> ppF2 "fromIndex" (ppE sh) (ppE ix) - While p f x -> ppF3 "while" (ppF p) (ppF f) (ppE x) - Foreign _ ff _ e -> ppF2 "foreign" (\_ -> pretty (strForeign ff)) (ppE e) - Shape arr -> ppF1 "shape" (ppA arr) - ShapeSize _ sh -> ppF1 "shapeSize" (ppE sh) - Index arr ix -> ppF2 (Operator (pretty '!') Infix L 9) (ppA arr) (ppE ix) - LinearIndex arr ix -> ppF2 (Operator "!!" Infix L 9) (ppA arr) (ppE ix) - Coerce _ tp x -> ppF1 (Operator (withTypeRep tp "coerce") App L 10) (ppE x) - Undef tp -> withTypeRep tp "undef" - + context' = Context L R 9 +prettyArrayInstr aenv context (LinearIndex arr) ix + = parensIf (ctxPrecedence context < 9) + $ sep [ prettyArrayVar aenv arr + , group (sep ["!!", ix context']) ] where - ppE :: OpenExp env aenv e -> Context -> Adoc - ppE e c = prettyOpenExp c env aenv e - - ppA :: ArrayVar aenv a -> Context -> Adoc - ppA acc _ = prettyArrayVar aenv acc - - ppF :: OpenFun env aenv f -> Context -> Adoc - ppF f _ = parens $ prettyOpenFun env aenv f - - ppF1 :: Operator -> (Context -> Adoc) -> Adoc - ppF1 op x - = parensIf (needsParens ctx op) - $ combine [ opName op, x ctx' ] - where - ctx' = isPrefix op ? (arg op R, app) - combine = isPrefix op ? (cat, hang 2 . sep) - - ppF2 :: Operator -> (Context -> Adoc) -> (Context -> Adoc) -> Adoc - ppF2 op x y - = parensIf (needsParens ctx op) - $ if isInfix op - then sep [ x (arg op L), group (sep [opName op, y (arg op R)]) ] - else hang 2 $ sep [ opName op, x app, y app ] - - ppF3 :: Operator -> (Context -> Adoc) -> (Context -> Adoc) -> (Context -> Adoc) -> Adoc - ppF3 op x y z - = parensIf (needsParens ctx op) - $ hang 2 - $ sep [ opName op, x app, y app, z app ] - - withTypeRep :: ScalarType t -> Adoc -> Adoc - withTypeRep t op = op <+> "@" <> pretty (show t) + context' = Context L R 9 prettyArrayVar :: forall aenv a. @@ -485,291 +331,3 @@ prettyArrayVar -> ArrayVar aenv a -> Adoc prettyArrayVar aenv (Var _ idx) = prj idx aenv - -prettyLet - :: forall env aenv t. - Context - -> Val env - -> Val aenv - -> OpenExp env aenv t - -> Adoc -prettyLet ctx env0 aenv - = parensIf (ctxPrecedence ctx > 0) - . align . wrap . collect env0 - where - collect :: Val env' -> OpenExp env' aenv e -> ([Adoc], Adoc) - collect env = - \case - Let lhs e1 e2 -> - let (env', v) = prettyELhs False env lhs - e1' = ppE env e1 - bnd | isLet e1 = nest shiftwidth (vsep [v <+> equals, e1']) - | otherwise = v <+> align (equals <+> e1') - (bnds, body) = collect env' e2 - in - (bnd:bnds, body) - -- - next -> ([], ppE env next) - - isLet :: OpenExp env' aenv t' -> Bool - isLet Let{} = True - isLet _ = False - - ppE :: Val env' -> OpenExp env' aenv t' -> Adoc - ppE env = prettyOpenExp context0 env aenv - - wrap :: ([Adoc], Adoc) -> Adoc - wrap ([], body) = body -- shouldn't happen! - wrap ([b], body) - = sep [ nest shiftwidth (sep [let_, b]), in_, body ] - wrap (bnds, body) - = vsep [ nest shiftwidth (vsep [let_, sepBy (flatAlt "" " ; ") bnds]) - , in_ - , body - ] - - sepBy :: Adoc -> [Adoc] -> Adoc - sepBy = encloseSep mempty mempty - -prettyTuple - :: forall env aenv t. - Context - -> Val env - -> Val aenv - -> OpenExp env aenv t - -> Adoc -prettyTuple ctx env aenv exp = case collect exp of - Nothing -> align $ ppPair exp - Just tup -> - case tup of - [] -> "()" - _ -> align $ parensIf (ctxPrecedence ctx > 0) ("T" <> pretty (length tup) <+> align (sep tup)) - where - ppPair :: OpenExp env aenv t' -> Adoc - ppPair (Pair e1 e2) = "(" <> ppPair e1 <> "," <+> prettyOpenExp context0 env aenv e2 <> ")" - ppPair e = prettyOpenExp context0 env aenv e - - collect :: OpenExp env aenv t' -> Maybe [Adoc] - collect Nil = Just [] - collect (Pair e1 e2) - | Just tup <- collect e1 = Just $ tup ++ [prettyOpenExp app env aenv e2] - collect _ = Nothing - -prettyCase - :: Val env - -> Val aenv - -> OpenExp env aenv a - -> [(TAG, OpenExp env aenv b)] - -> Maybe (OpenExp env aenv b) - -> Adoc -prettyCase env aenv x xs def - = hang shiftwidth - $ vsep [ case_ <+> x' <+> of_ - , flatAlt (vcat xs') (encloseSep "{ " " }" "; " xs') - ] - where - x' = prettyOpenExp context0 env aenv x - xs' = map (\(t,e) -> pretty t <+> "->" <+> prettyOpenExp context0 env aenv e) xs - ++ case def of - Nothing -> [] - Just d -> ["_" <+> "->" <+> prettyOpenExp context0 env aenv d] - -{- - -prettyAtuple - :: forall acc aenv arrs. - PrettyAcc acc - -> ExtractAcc acc - -> Val aenv - -> PreOpenAcc acc aenv arrs - -> Adoc -prettyAtuple prettyAcc extractAcc aenv0 acc = case collect acc of - Just tup -> align $ "T" <> pretty (length tup) <+> sep tup - Nothing -> align $ ppPair acc - where - ppPair :: PreOpenAcc acc aenv arrs' -> Adoc - ppPair (Apair a1 a2) = "(" <> ppPair (extractAcc a1) <> "," <+> prettyAcc context0 aenv0 a2 <> ")" - ppPair a = prettyPreOpenAcc context0 prettyAcc extractAcc aenv0 a - - collect :: PreOpenAcc acc aenv arrs' -> Maybe [Adoc] - collect Anil = Just [] - collect (Apair a1 a2) - | Just tup <- collect $ extractAcc a1 - = Just $ tup ++ [prettyAcc app aenv0 a2] - collect _ = Nothing --} - -prettyConst :: TypeR e -> e -> Adoc -prettyConst tp x = - let y = showElt tp x - in parensIf (any isSpace y) (pretty y) - -prettyPrimConst :: PrimConst a -> Adoc -prettyPrimConst PrimMinBound{} = "minBound" -prettyPrimConst PrimMaxBound{} = "maxBound" -prettyPrimConst PrimPi{} = "pi" - - --- Primitive operators --- ------------------- --- --- The core of the pretty printer is how to correctly handle precedence, --- associativity, and fixity of the primitive scalar operators. --- - -data Direction = L | N | R - deriving Eq - -data Fixity = App | Infix | Prefix - deriving Eq - -type Precedence = Int -type Associativity = Direction - -data Context = Context - { ctxAssociativity :: Associativity - , ctxPosition :: Direction - , ctxPrecedence :: Precedence - } - -data Operator = Operator - { opName :: Adoc - , opFixity :: Fixity - , opAssociativity :: Associativity - , opPrecedence :: Precedence - } - -instance IsString Operator where - fromString s = Operator (fromString s) App L 10 - -needsParens :: Context -> Operator -> Bool -needsParens Context{..} Operator{..} - | ctxPrecedence < opPrecedence = False - | ctxPrecedence > opPrecedence = True - | ctxAssociativity /= opAssociativity = True - | otherwise = ctxPosition /= opAssociativity - -context0 :: Context -context0 = Context N N 0 - -app :: Context -app = Context L N 10 - -arg :: Operator -> Direction -> Context -arg Operator{..} side = Context opAssociativity side opPrecedence - -isPrefix :: Operator -> Bool -isPrefix Operator{..} = opFixity == Prefix - -isInfix :: Operator -> Bool -isInfix Operator{..} = opFixity == Infix - -primOperator :: PrimFun a -> Operator -primOperator PrimAdd{} = Operator (pretty '+') Infix L 6 -primOperator PrimSub{} = Operator (pretty '-') Infix L 6 -primOperator PrimMul{} = Operator (pretty '*') Infix L 7 -primOperator PrimNeg{} = Operator (pretty '-') Prefix L 6 -- Haskell's only prefix operator -primOperator PrimAbs{} = Operator "abs" App L 10 -primOperator PrimSig{} = Operator "signum" App L 10 -primOperator PrimQuot{} = Operator "quot" App L 10 -primOperator PrimRem{} = Operator "rem" App L 10 -primOperator PrimQuotRem{} = Operator "quotRem" App L 10 -primOperator PrimIDiv{} = Operator "div" App L 10 -primOperator PrimMod{} = Operator "mod" App L 10 -primOperator PrimDivMod{} = Operator "divMod" App L 10 -primOperator PrimBAnd{} = Operator ".&." Infix L 7 -primOperator PrimBOr{} = Operator ".|." Infix L 5 -primOperator PrimBXor{} = Operator "xor" App L 10 -primOperator PrimBNot{} = Operator "complement" App L 10 -primOperator PrimBShiftL{} = Operator "shiftL" App L 10 -primOperator PrimBShiftR{} = Operator "shiftR" App L 10 -primOperator PrimBRotateL{} = Operator "rotateL" App L 10 -primOperator PrimBRotateR{} = Operator "rotateR" App L 10 -primOperator PrimPopCount{} = Operator "popCount" App L 10 -primOperator PrimCountLeadingZeros{} = Operator "countLeadingZeros" App L 10 -primOperator PrimCountTrailingZeros{} = Operator "countTrailingZeros" App L 10 -primOperator PrimFDiv{} = Operator (pretty '/') Infix L 7 -primOperator PrimRecip{} = Operator "recip" App L 10 -primOperator PrimSin{} = Operator "sin" App L 10 -primOperator PrimCos{} = Operator "cos" App L 10 -primOperator PrimTan{} = Operator "tan" App L 10 -primOperator PrimAsin{} = Operator "asin" App L 10 -primOperator PrimAcos{} = Operator "acos" App L 10 -primOperator PrimAtan{} = Operator "atan" App L 10 -primOperator PrimSinh{} = Operator "sinh" App L 10 -primOperator PrimCosh{} = Operator "cosh" App L 10 -primOperator PrimTanh{} = Operator "tanh" App L 10 -primOperator PrimAsinh{} = Operator "asinh" App L 10 -primOperator PrimAcosh{} = Operator "acosh" App L 10 -primOperator PrimAtanh{} = Operator "atanh" App L 10 -primOperator PrimExpFloating{} = Operator "exp" App L 10 -primOperator PrimSqrt{} = Operator "sqrt" App L 10 -primOperator PrimLog{} = Operator "log" App L 10 -primOperator PrimFPow{} = Operator "**" Infix R 8 -primOperator PrimLogBase{} = Operator "logBase" App L 10 -primOperator PrimTruncate{} = Operator "truncate" App L 10 -primOperator PrimRound{} = Operator "round" App L 10 -primOperator PrimFloor{} = Operator "floor" App L 10 -primOperator PrimCeiling{} = Operator "ceiling" App L 10 -primOperator PrimAtan2{} = Operator "atan2" App L 10 -primOperator PrimIsNaN{} = Operator "isNaN" App L 10 -primOperator PrimIsInfinite{} = Operator "isInfinite" App L 10 -primOperator PrimLt{} = Operator "<" Infix N 4 -primOperator PrimGt{} = Operator ">" Infix N 4 -primOperator PrimLtEq{} = Operator "<=" Infix N 4 -primOperator PrimGtEq{} = Operator ">=" Infix N 4 -primOperator PrimEq{} = Operator "==" Infix N 4 -primOperator PrimNEq{} = Operator "/=" Infix N 4 -primOperator PrimMax{} = Operator "max" App L 10 -primOperator PrimMin{} = Operator "min" App L 10 -primOperator PrimLAnd = Operator "&&" Infix R 3 -primOperator PrimLOr = Operator "||" Infix R 2 -primOperator PrimLNot = Operator "not" App L 10 -primOperator PrimFromIntegral{} = Operator "fromIntegral" App L 10 -primOperator PrimToFloating{} = Operator "toFloating" App L 10 - - --- Environments --- ------------ - -data Val env where - Empty :: Val () - Push :: Val env -> Adoc -> Val (env, t) - -class PrettyEnv env where - prettyEnv :: Adoc -> Val env - -instance PrettyEnv () where - prettyEnv _ = Empty - -instance PrettyEnv env => PrettyEnv (env, t) where - prettyEnv v = - let env = prettyEnv v :: Val env - x = v <> pretty (sizeEnv env) - in - env `Push` x - -sizeEnv :: Val env -> Int -sizeEnv Empty = 0 -sizeEnv (Push env _) = 1 + sizeEnv env - -prj :: Idx env t -> Val env -> Adoc -prj ZeroIdx (Push _ v) = v -prj (SuccIdx ix) (Push env _) = prj ix env - - --- Utilities --- --------- - -shiftwidth :: Int -shiftwidth = 2 - -infix 0 ? -(?) :: Bool -> (a, a) -> a -True ? (t,_) = t -False ? (_,f) = f - -parensIf :: Bool -> Doc ann -> Doc ann -parensIf True = group . parens . align -parensIf False = id - diff --git a/src/Data/Array/Accelerate/Pretty/Schedule.hs b/src/Data/Array/Accelerate/Pretty/Schedule.hs new file mode 100644 index 000000000..4dc0ceb6e --- /dev/null +++ b/src/Data/Array/Accelerate/Pretty/Schedule.hs @@ -0,0 +1,30 @@ +-- | +-- Module : Data.Array.Accelerate.Pretty.Schedule +-- Copyright : [2008..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.Pretty.Schedule ( + PrettySchedule(..), PrettyKernel(..), PrettyKernelStyle(..), Adoc +) where + +import Data.Array.Accelerate.AST.Kernel +import Data.Array.Accelerate.Pretty.Exp + +class PrettySchedule sched where + prettySchedule :: PrettyKernel kernel => sched kernel () t -> Adoc + +class PrettyKernel kernel where + prettyKernel :: PrettyKernelStyle kernel + +-- Whether modifiers should be included in the valuation of buffer arguments. +type IncludeModifier = Bool + +data PrettyKernelStyle kernel where + PrettyKernelBody :: IncludeModifier -> (forall env. Val env -> kernel env -> Adoc) -> PrettyKernelStyle kernel + + PrettyKernelFun :: (forall t. KernelFun kernel t -> Adoc) -> PrettyKernelStyle kernel diff --git a/src/Data/Array/Accelerate/Pretty/Schedule/Sequential.hs b/src/Data/Array/Accelerate/Pretty/Schedule/Sequential.hs new file mode 100644 index 000000000..a7a6cb596 --- /dev/null +++ b/src/Data/Array/Accelerate/Pretty/Schedule/Sequential.hs @@ -0,0 +1,118 @@ +{-# OPTIONS_GHC -Wno-orphans #-} +-- | +-- Module : Data.Array.Accelerate.Pretty.Schedule.Sequential +-- Copyright : [2008..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.Pretty.Schedule.Sequential ( + +) where + +import Data.Array.Accelerate.Pretty.Exp +import Data.Array.Accelerate.Pretty.Type +import Data.Array.Accelerate.Pretty.Schedule +import Data.Array.Accelerate.AST.Kernel +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.AST.Schedule.Sequential +import Data.Array.Accelerate.Pretty.Operation +import Data.Array.Accelerate.Representation.Type + +import Prettyprinter + +import Prelude hiding (exp) + +instance PrettySchedule SequentialSchedule where + prettySchedule = prettySequentialSchedule empty' + +prettySequentialSchedule :: PrettyKernel kernel => Val' env -> SequentialSchedule kernel env t -> Adoc +prettySequentialSchedule env (SequentialLam lhs f) + | (env', lhs') <- prettyGLhsWithTypes env lhs + = "\\" <> lhs' <+> "->" <> hardline <> indent 2 (prettySequentialSchedule env' f) +prettySequentialSchedule env (SequentialBody sched) + = prettySeqSchedule env sched + +prettySeqScheduleFun :: PrettyKernel kernel => Val' env -> SeqScheduleFun kernel env t -> Adoc +prettySeqScheduleFun env (Slam lhs f) + | (env', lhs') <- prettyGLhsWithTypes env lhs + = "\\" <> lhs' <+> "->" <> hardline <> indent 2 (prettySeqScheduleFun env' f) +prettySeqScheduleFun env (Sbody sched) + = prettySeqSchedule env sched + +prettySeqSchedule :: forall kernel env t. PrettyKernel kernel => Val' env -> SeqSchedule kernel env t -> Adoc +prettySeqSchedule env = \case + Exec _ kernel args -> prettyKernelFun env kernel args + Return vars -> hang 2 $ group $ vsep [annotate Statement "return", prettyVars (val env) 10 vars] + Compute exp -> hang 2 $ group $ vsep [annotate Statement "compute", prettyExp (val env) exp] + Alet LeftHandSideUnit _ bnd body + | notReturn bnd + -- A return looks very strange if there is no explict LHS. It's uncommon, + -- but also very strange when this does happens. + -> prettySeqSchedule env bnd + <> hardline + <> prettySeqSchedule env body + Alet lhs us bnd body + | (env', lhs') <- prettyGLhsWithUniquenessTypes env lhs us + -> hang 2 (group $ vsep [lhs' <+> "=", prettySeqSchedule env bnd]) + <> hardline + <> prettySeqSchedule env' body + Alloc _ tp sh -> hang 2 $ group $ vsep [annotate Statement "alloc", prettyScalarType tp <> "[" <> prettyShapeVars (val env) sh <> "]"] + Use tp n buffer -> hang 2 $ group $ vsep [annotate Statement "use" <+> prettyScalarType tp <> "[" <> pretty n <> "]", prettyBuffer tp n buffer] + Unit var -> hang 2 $ group $ vsep [annotate Statement "unit", prettyVar (val env) var] + Acond c true false + -> group $ vsep + [ hang 2 $ group $ vsep + [ if_ <+> prettyVar (val env) c <+> then_ + , prettySeqSchedule env true + ] + , hang 2 $ group $ vsep + [ else_ + , prettySeqSchedule env false + ] + ] + Awhile us condition step initial + -> annotate Statement "awhile" <+> prettyTupR (const prettyGroundRWithUniqueness) 10 (groundsRWithUniquenesses (mapTupR varType initial) us) + <> hardline <> hang 4 (" ( " <> prettySeqScheduleFun env condition) + <> hardline <> " )" + <> hardline <> hang 4 (" ( " <> prettySeqScheduleFun env step) + <> hardline <> " )" + <> hardline <> indent 2 (prettyVars (val env) 10 initial) + Aassert msg cond + -> hang 2 $ group $ vsep [annotate Statement "assert", prettyText msg, prettyExp (val env) cond] + Aassume cond + -> hang 2 $ group $ vsep [annotate Statement "assume", prettyExp (val env) cond] + where + notReturn Return{} = False + notReturn _ = True + +prettySArgs :: Val' benv -> SArgs benv f -> Adoc +prettySArgs env args = tupled $ map (\(Exists a) -> prettySArg env a) $ argsToList args + +prettySArg :: Val' benv -> SArg benv t -> Adoc +prettySArg env (SArgScalar var) = prettyVar (val env) var +prettySArg env (SArgBuffer m var) = prettyModifier m <+> prettyVar (val env) var + +prettyKernelFun :: forall kernel env f. PrettyKernel kernel => Val' env -> KernelFun kernel f -> SArgs env f -> Adoc +prettyKernelFun env fun args = case prettyKernel of + PrettyKernelBody includeModifier prettyKernelBody -> + let + go :: Val kenv -> OpenKernelFun kernel kenv t -> SArgs env t -> Adoc + go kenv (KernelFunBody kernel) ArgsNil = prettyKernelBody kenv kernel + go kenv (KernelFunLam (KernelArgRscalar _) f) (SArgScalar a :>: as) + = go (Push kenv $ prettyVar (val env) a) f as + go kenv (KernelFunLam (KernelArgRbuffer _ _) f) (SArgBuffer mod' a :>: as) = + let + a' + | includeModifier = prettyModifier mod' <+> prettyVar (val env) a + | otherwise = prettyVar (val env) a + in + go (Push kenv a') f as + in + go Empty fun args + PrettyKernelFun prettyKernelAsFun -> + prettyKernelAsFun fun + <> hardline <> indent 2 (prettySArgs env args) diff --git a/src/Data/Array/Accelerate/Pretty/Schedule/Uniform.hs b/src/Data/Array/Accelerate/Pretty/Schedule/Uniform.hs new file mode 100644 index 000000000..a4673489b --- /dev/null +++ b/src/Data/Array/Accelerate/Pretty/Schedule/Uniform.hs @@ -0,0 +1,235 @@ +{-# OPTIONS_GHC -Wno-orphans #-} +-- | +-- Module : Data.Array.Accelerate.Pretty.Schedule.Uniform +-- Copyright : [2008..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.Pretty.Schedule.Uniform ( + prettySArg, prettySArgs, prettyKernelFun +) where + +import Data.Array.Accelerate.Pretty.Exp +import Data.Array.Accelerate.Pretty.Type +import Data.Array.Accelerate.Pretty.Schedule +import Data.Array.Accelerate.Pretty.Operation (prettyGroundR, prettyBuffer, prettyExp, prettyModifier) +import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.Kernel +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.AST.Schedule.Uniform +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Representation.Shape +import Data.Array.Accelerate.Type +import Data.String + +import Prettyprinter + +import Prelude hiding (exp) + +instance PrettySchedule UniformScheduleFun where + prettySchedule = prettyUniformScheduleFun empty' + +prettyBaseR :: BaseR t -> Adoc +prettyBaseR (BaseRground ground) = prettyGroundR ground +prettyBaseR (BaseRref ground) = "*" <> prettyGroundR ground +prettyBaseR (BaseRrefWrite ground) = "*out " <> prettyGroundR ground +prettyBaseR BaseRsignal = "Signal" +prettyBaseR BaseRsignalResolver = "Resolver" + +-- Naming: +-- * We use 'b' and 'e' as prefix for buffer and scalar variables, like we also +-- do in the pretty printer for OperationAcc. +-- * We use 's' as a prefix for signal and signal resolvers. For the latter we +-- also append a prime ('). +-- * We use 'r' as a prefix for references and writable references. For the +-- latter we append a prime ('). +-- * We use the same index for a signal and a signal resolver if they are +-- introduced in a NewSignal. Similarly, we use the same indices for a +-- reference and a writable reference if they were introduced by a +-- NewRef +data Val' env = Val' (Val env) Int Int Int Int + +val :: Val' env -> Val env +val (Val' env _ _ _ _) = env + +push :: BaseR t -> Val' env -> (Val' (env, t), Adoc) +push (BaseRground GroundRbuffer{}) (Val' env b e s r) = (Val' (Push env d) (b + 1) e s r, d) + where d = "b" <> pretty b +push (BaseRground GroundRscalar{}) (Val' env b e s r) = (Val' (Push env d) b (e + 1) s r, d) + where d = "e" <> pretty e +push BaseRsignal (Val' env b e s r) = (Val' (Push env d) b e (s + 1) r, d) + where d = "s" <> pretty s +push BaseRsignalResolver (Val' env b e s r) = (Val' (Push env d) b e (s + 1) r, d) + where d = "s" <> pretty s <> "'" +push BaseRref{} (Val' env b e s r) = (Val' (Push env d) b e s (r + 1), d) + where d = "r" <> pretty r +push BaseRrefWrite{} (Val' env b e s r) = (Val' (Push env d) b e s (r + 1), d) + where d = "r" <> pretty r <> "'" + +pushNewSignal :: Val' env -> (Val' ((env, Signal), SignalResolver), Adoc, Adoc) +pushNewSignal (Val' env b e s r) = (Val' (Push (Push env d1) d2) b e (s + 1) r, d1, d2) + where + d1 = "s" <> pretty s + d2 = "s" <> pretty s <> "'" + +pushNewRef :: Val' env -> (Val' ((env, Ref t), OutputRef t), Adoc, Adoc) +pushNewRef (Val' env b e s r) = (Val' (Push (Push env d1) d2) b e s (r + 1), d1, d2) + where + d1 = "r" <> pretty r + d2 = "r" <> pretty r <> "'" + +prettyBLhs :: Val' env -> BLeftHandSide t env env' -> (Val' env', Adoc) +prettyBLhs = prettyLhs' push Nothing False + +prettyBLhsWithTypes :: Val' env -> BLeftHandSide t env env' -> (Val' env', Adoc) +prettyBLhsWithTypes = prettyLhs' push (Just $ \(Exists t) -> prettyBaseR t) False + +prettyBLhsForBinding :: Val' env -> BLeftHandSide t env env' -> Binding env t -> (Val' env', Adoc) +prettyBLhsForBinding env (LeftHandSideSingle _ `LeftHandSidePair` LeftHandSideSingle _) binding + | NewSignal _ <- binding + , (env', a, b) <- pushNewSignal env = (env', pair a b) + | NewRef _ <- binding + , (env', a, b) <- pushNewRef env = (env', pair a b) + where pair a b = "(" <> a <> "," <+> b <> ")" +prettyBLhsForBinding env lhs binding + | bindingHasTrivialType binding = prettyBLhs env lhs + | otherwise = prettyBLhsWithTypes env lhs + +bindingHasTrivialType :: Binding env t -> Bool +bindingHasTrivialType (NewSignal _) = True +bindingHasTrivialType (NewRef _) = True +bindingHasTrivialType (Alloc _ _ _) = True +bindingHasTrivialType (Use _ _ _) = True +bindingHasTrivialType _ = False + +empty' :: Val' () +empty' = Val' Empty 0 0 0 0 + +prettyUniformScheduleFun :: PrettyKernel kernel => Val' env -> UniformScheduleFun kernel env t -> Adoc +prettyUniformScheduleFun env (Slam lhs f) + | (env', lhs') <- prettyBLhsWithTypes env lhs + = "\\" <> lhs' <+> "->" <> hardline <> indent 2 (prettyUniformScheduleFun env' f) +prettyUniformScheduleFun env (Sbody sched) + = prettyUniformSchedule env sched + +prettyUniformSchedule :: forall kernel env. PrettyKernel kernel => Val' env -> UniformSchedule kernel env -> Adoc +prettyUniformSchedule env = \case + Return + -> annotate Statement "return" + Alet lhs bnd body + | (env', lhs') <- prettyBLhsForBinding env lhs bnd + -> group (vsep [lhs' <+> "=", prettyBinding env bnd]) <> prettyNext' env' body + Effect effect next + -> prettyEffect env effect <> prettyNext next + Acond c true false next + -> group (vsep + [ hang 2 $ group $ vsep + [ if_ <+> prettyVar env c <+> then_ + , prettyUniformSchedule env true + ] + , hang 2 $ group $ vsep + [ else_ + , prettyUniformSchedule env false + ] + ]) + <> prettyNext next + Awhile _ body initial next + -> hang 2 (group (vsep [annotate Statement "awhile", prettyVars env 10 initial])) + <> hardline <> hang 4 (" (" <+> prettyUniformScheduleFun env body) + <> hardline <> " )" + <> prettyNext next + AwhileSeq _ body initial next + -> hang 2 (group (vsep [annotate Statement "awhile sequential", prettyVars env 10 initial])) + <> hardline <> hang 4 (" (" <+> prettyUniformScheduleFun env body) + <> hardline <> " )" + <> prettyNext next + Spawn (Effect (SignalAwait signals) body) next + -> annotate Statement "spawn await" <+> list (map (prettyIdx env) signals) <+> "{" + <> hardline <> indent 2 (prettyUniformSchedule env body) + <> hardline <> "}" + <> prettyNext next + Spawn body next + -> annotate Statement "spawn" <+> "{" + <> hardline <> indent 2 (prettyUniformSchedule env body) + <> hardline <> "}" + <> prettyNext next + where + prettyNext' :: Val' env' -> UniformSchedule kernel env' -> Adoc + prettyNext' _ Return = emptyDoc + prettyNext' env' sched = hardline <> prettyUniformSchedule env' sched + + prettyNext = prettyNext' env + +prettyBinding :: Val' env -> Binding env t -> Adoc +prettyBinding env = \case + Compute expr -> hang 2 $ group $ vsep [annotate Statement "compute", prettyExp (val env) expr] + NewSignal name -> annotate Statement "new signal" <+> "--" <+> fromString name + NewRef tp -> annotate Statement "ref" <+> prettyGroundR tp + Alloc _ tp sh -> hang 2 $ group $ vsep [annotate Statement "alloc", prettyScalarType tp <> "[" <> prettyShapeVars env sh <> "]"] + Use tp n buffer -> hang 2 $ group $ vsep [annotate Statement "use" <+> prettyScalarType tp <> "[" <> pretty n <> "]", prettyBuffer tp n buffer] + Unit var -> hang 2 $ group $ vsep [annotate Statement "unit", prettyVar env var] + RefRead var -> "*" <> prettyVar env var + +prettyEffect :: PrettyKernel kernel => Val' env -> Effect kernel env -> Adoc +prettyEffect env = \case + Exec _ kernel args -> prettyKernelFun env kernel args + SignalAwait signals -> hang 2 $ group $ vsep [annotate Statement "await", list $ map (prettyIdx env) signals] + SignalResolve signals -> hang 2 $ group $ vsep [annotate Statement "resolve", list $ map (prettyIdx env) signals] + RefWrite ref value -> hang 2 $ group $ vsep ["*" <> prettyVar env ref <+> "=", prettyVar env value] + Aassert msg cond -> hang 2 $ group $ vsep [annotate Statement "assert", prettyText msg, prettyExp (val env) cond] + Atrace msg t -> hang 2 $ group $ vsep [annotate Statement "atrace", prettyText msg, prettyArrayDescriptors env t] + +prettyKernelFun :: forall kernel env f. PrettyKernel kernel => Val' env -> KernelFun kernel f -> SArgs env f -> Adoc +prettyKernelFun env fun args = case prettyKernel of + PrettyKernelBody includeModifier prettyKernelBody -> + let + go :: Val kenv -> OpenKernelFun kernel kenv t -> SArgs env t -> Adoc + go kenv (KernelFunBody kernel) ArgsNil = prettyKernelBody kenv kernel + go kenv (KernelFunLam (KernelArgRscalar _) f) (SArgScalar a :>: as) + = go (Push kenv $ prettyVar env a) f as + go kenv (KernelFunLam (KernelArgRbuffer _ _) f) (SArgBuffer mod' a :>: as) = + let + a' + | includeModifier = prettyModifier mod' <+> prettyVar env a + | otherwise = prettyVar env a + in + go (Push kenv a') f as + in + go Empty fun args + PrettyKernelFun prettyKernelAsFun -> + hang 2 $ group $ vsep [annotate Execute "execute", prettyKernelAsFun fun, prettySArgs env args] + +prettySArgs :: Val' benv -> SArgs benv f -> Adoc +prettySArgs env args = tupled $ map (\(Exists a) -> prettySArg env a) $ argsToList args + +prettySArg :: Val' benv -> SArg benv t -> Adoc +prettySArg env (SArgScalar var) = prettyVar env var +prettySArg env (SArgBuffer m var) = prettyModifier m <+> prettyVar env var + +-- Variables +prettyVar :: Val' env -> Var s env t -> Adoc +prettyVar (Val' env _ _ _ _) (Var _ ix) = prj ix env + +prettyIdx :: Val' env -> Idx env t -> Adoc +prettyIdx (Val' env _ _ _ _) ix = prj ix env + +prettyVars :: forall env s t. Val' env -> Precedence -> Vars s env t -> Adoc +prettyVars env = prettyTupR $ const $ prettyVar env + +prettyShapeVars :: Val' env -> Vars s env sh -> Adoc +prettyShapeVars _ TupRunit = "Z" +prettyShapeVars env vars = encloseSep "Z :. " "" " :. " $ map (\(Exists v) -> prettyVar env v) $ flattenTupR vars + +prettyArrayDescriptors :: Val' env -> ArrayDescriptors env t -> Adoc +prettyArrayDescriptors env = prettyTupR (prettyArrayDescriptor env) 10 + +prettyArrayDescriptor :: Val' env -> Precedence -> ArrayDescriptor env t -> Adoc +prettyArrayDescriptor env p (ArrayDescriptor shr sh buffer) = parensIf (p > 0) $ prettyModifier In <+> sh' <+> prettyVars env 10 buffer + where + sh' = case shr of + ShapeRz -> "Z" + _ -> "(" <> prettyShapeVars env sh <> ")" diff --git a/src/Data/Array/Accelerate/Pretty/Type.hs b/src/Data/Array/Accelerate/Pretty/Type.hs new file mode 100644 index 000000000..ebbb0ae36 --- /dev/null +++ b/src/Data/Array/Accelerate/Pretty/Type.hs @@ -0,0 +1,49 @@ +{-# LANGUAGE CPP #-} +-- | +-- Module : Data.Array.Accelerate.Pretty.Type +-- Copyright : [2008..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.Pretty.Type where + +import Data.Array.Accelerate.Pretty.Exp +import Data.Array.Accelerate.Type +import Data.Array.Accelerate.Representation.Type + +#if MIN_VERSION_prettyprinter (1,7,1) +import Prettyprinter +#else +import Data.Text.Prettyprint.Doc +#endif + +prettyScalarType :: ScalarType t -> Adoc +prettyScalarType (SingleScalarType t) = prettySingleType t +prettyScalarType (VectorScalarType t) = prettyVectorType t + +prettyVectorType :: VectorType t -> Adoc +prettyVectorType (VectorType n t) = prettySingleType t <> "x" <> pretty n + +prettySingleType :: SingleType t -> Adoc +prettySingleType (NumSingleType (IntegralNumType tp)) = case tp of + TypeInt -> "Int" + TypeInt8 -> "Int8" + TypeInt16 -> "Int16" + TypeInt32 -> "Int32" + TypeInt64 -> "Int64" + TypeWord -> "Word" + TypeWord8 -> "Word8" + TypeWord16 -> "Word16" + TypeWord32 -> "Word32" + TypeWord64 -> "Word64" +prettySingleType (NumSingleType (FloatingNumType tp)) = case tp of + TypeHalf -> "Half" + TypeFloat -> "Float" + TypeDouble -> "Double" + +prettyType :: TypeR t -> Adoc +prettyType = prettyTupR (const prettyScalarType) 10 diff --git a/src/Data/Array/Accelerate/Representation/Array.hs b/src/Data/Array/Accelerate/Representation/Array.hs index ee02c0ad5..7bd429593 100644 --- a/src/Data/Array/Accelerate/Representation/Array.hs +++ b/src/Data/Array/Accelerate/Representation/Array.hs @@ -1,11 +1,4 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TemplateHaskellQuotes #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.Representation.Array @@ -18,9 +11,10 @@ -- module Data.Array.Accelerate.Representation.Array + ( module Data.Array.Accelerate.Representation.Array, Buffer, Buffers ) where -import Data.Array.Accelerate.Array.Data +import Data.Array.Accelerate.Array.Buffer import Data.Array.Accelerate.Error import Data.Array.Accelerate.Type import Data.Array.Accelerate.Representation.Elt @@ -42,9 +36,10 @@ import qualified Data.Vector.Unboxed as U -- data Array sh e where Array :: sh -- extent of dimensions = shape - -> ArrayData e -- array payload + -> Buffers e -- array payload -> Array sh e + -- | Segment descriptor (vector of segment lengths). -- -- To represent nested one-dimensional arrays, we use a flat array of data @@ -68,7 +63,7 @@ data ArrayR a where type ArraysR = TupR ArrayR instance Show (ArrayR a) where - show (ArrayR shR eR) = "Array DIM" ++ show (rank shR) ++ " " ++ show eR + show (ArrayR shR tp) = "Array DIM" ++ show (rank shR) ++ " " ++ show tp formatArrayR :: Format r (ArrayR a -> r) formatArrayR = later $ \case @@ -84,17 +79,21 @@ showArraysR :: ArraysR a -> ShowS showArraysR = shows arraysRarray :: ShapeR sh -> TypeR e -> ArraysR (Array sh e) -arraysRarray shR eR = TupRsingle (ArrayR shR eR) +arraysRarray shR tp = TupRsingle (ArrayR shR tp) arraysRpair :: ArrayR a -> ArrayR b -> ArraysR (((), a), b) arraysRpair a b = TupRunit `TupRpair` TupRsingle a `TupRpair` TupRsingle b +arraysRFunctionImpossible :: ArraysR (s -> t) -> a +arraysRFunctionImpossible (TupRsingle repr) = case repr of {} + -- | Creates a new, uninitialized Accelerate array. -- allocateArray :: ArrayR (Array sh e) -> sh -> IO (Array sh e) -allocateArray (ArrayR shR eR) sh = do - adata <- newArrayData eR (size shR sh) - return $! Array sh adata +allocateArray (ArrayR shR tp) sh = do + mbuffers <- newBuffers tp (size shR sh) + let buffers = unsafeFreezeBuffers tp mbuffers + return $! Array sh buffers -- | Create an array from its representation function, applied at each -- index of the array. @@ -107,63 +106,69 @@ fromFunction repr sh f = unsafePerformIO $! fromFunctionM repr sh (return . f) -- @since 1.2.0.0 -- fromFunctionM :: ArrayR (Array sh e) -> sh -> (sh -> IO e) -> IO (Array sh e) -fromFunctionM (ArrayR shR eR) sh f = do +fromFunctionM (ArrayR shR tp) sh f = do let !n = size shR sh - arr <- newArrayData eR n + mbuffers <- newBuffers tp n -- let write !i | i >= n = return () | otherwise = do v <- f (fromIndex shR sh i) - writeArrayData eR arr i v + writeBuffers tp mbuffers i v write (i+1) -- write 0 - return $! arr `seq` Array sh arr + let buffers = unsafeFreezeBuffers tp mbuffers + return $! buffers `seq` Array sh buffers -- | Convert a list into an Accelerate 'Array' in dense row-major order. -- fromList :: forall sh e. ArrayR (Array sh e) -> sh -> [e] -> Array sh e -fromList (ArrayR shR eR) sh xs = adata `seq` Array sh adata +fromList (ArrayR shR tp) sh xs = assertShapeNonNegative shR sh `seq` buffers `seq` Array sh buffers where -- Assume the array is in dense row-major order. This is safe because -- otherwise backends would not be able to directly memcpy. -- !n = size shR sh - (adata, _) = runArrayData @e $ do - arr <- newArrayData eR n + (buffers, _) = runBuffers tp $ do + mbuffers <- newBuffers tp n let go !i _ | i >= n = return () - go !i (v:vs) = writeArrayData eR arr i v >> go (i+1) vs + go !i (v:vs) = writeBuffers tp mbuffers i v >> go (i+1) vs go _ [] = error "Data.Array.Accelerate.fromList: not enough input data" -- go 0 xs - return (arr, undefined) + return (mbuffers, undefined) +assertShapeNonNegative :: ShapeR sh -> sh -> () +assertShapeNonNegative ShapeRz () = () +assertShapeNonNegative (ShapeRsnoc shr) (sh, sz) + | sz < 0 = error "Data.Array.Accelerate.fromList: negative size" + | otherwise = assertShapeNonNegative shr sh -- | Convert an accelerated 'Array' to a list in row-major order. -- toList :: ArrayR (Array sh e) -> Array sh e -> [e] -toList (ArrayR shR eR) (Array sh adata) = go 0 +toList (ArrayR shR tp) (Array sh buffers) = go 0 where -- Assume underling array is in row-major order. This is safe because -- otherwise backends would not be able to directly memcpy. -- !n = size shR sh go !i | i >= n = [] - | otherwise = indexArrayData eR adata i : go (i+1) + | otherwise = indexBuffers tp buffers i : go (i+1) concatVectors :: forall e. TypeR e -> [Vector e] -> Vector e -concatVectors tR vs = adata `seq` Array ((), len) adata +concatVectors tp vs = buffers `seq` Array ((), len) buffers where - offsets = scanl (+) 0 (map (size dim1 . shape) vs) - len = last offsets - (adata, _) = runArrayData @e $ do - arr <- newArrayData tR len - sequence_ [ writeArrayData tR arr (i + k) (indexArrayData tR ad i) + offsets = scanl (+) 0 (map (size dim1 . shape) vs) + len = last offsets + (buffers, _) = runBuffers tp $ do + mbuffers <- newBuffers tp len + sequence_ [ writeBuffers tp mbuffers (i + k) (indexBuffers tp ad i) | (Array ((), n) ad, k) <- vs `zip` offsets , i <- [0 .. n - 1] ] - return (arr, undefined) + return (mbuffers, undefined) shape :: Array sh e -> sh shape (Array sh _) = sh @@ -173,17 +178,17 @@ reshape shR sh shR' (Array sh' adata) = boundsCheck "shape mismatch" (size shR sh == size shR' sh') $ Array sh adata -(!) :: (ArrayR (Array sh e), Array sh e) -> sh -> e +(!) :: HasCallStack => (ArrayR (Array sh e), Array sh e) -> sh -> e (!) = uncurry indexArray -(!!) :: (TypeR e, Array sh e) -> Int -> e +(!!) :: HasCallStack => (TypeR e, Array sh e) -> Int -> e (!!) = uncurry linearIndexArray -indexArray :: ArrayR (Array sh e) -> Array sh e -> sh -> e -indexArray (ArrayR shR adR) (Array sh adata) ix = indexArrayData adR adata (toIndex shR sh ix) +indexArray :: HasCallStack => ArrayR (Array sh e) -> Array sh e -> sh -> e +indexArray (ArrayR shR tp) (Array sh buffers) ix = indexBuffers tp buffers (toIndex shR sh ix) linearIndexArray :: TypeR e -> Array sh e -> Int -> e -linearIndexArray adR (Array _ adata) = indexArrayData adR adata +linearIndexArray tp (Array _ buffers) = indexBuffers tp buffers showArray :: (e -> ShowS) -> ArrayR (Array sh e) -> Array sh e -> String showArray f arrR@(ArrayR shR _) arr@(Array sh _) = case shR of @@ -293,10 +298,10 @@ reduceRank :: ArrayR (Array (sh, Int) e) -> ArrayR (Array sh e) reduceRank (ArrayR (ShapeRsnoc shR) aeR) = ArrayR shR aeR rnfArray :: ArrayR a -> a -> () -rnfArray (ArrayR shR adR) (Array sh ad) = rnfShape shR sh `seq` rnfArrayData adR ad +rnfArray (ArrayR shR tp) (Array sh buffers) = rnfShape shR sh `seq` rnfBuffers tp buffers rnfArrayR :: ArrayR arr -> () -rnfArrayR (ArrayR shR tR) = rnfShapeR shR `seq` rnfTupR rnfScalarType tR +rnfArrayR (ArrayR shR tp) = rnfShapeR shR `seq` rnfTupR rnfScalarType tp rnfArraysR :: ArraysR arrs -> arrs -> () rnfArraysR TupRunit () = () @@ -304,7 +309,7 @@ rnfArraysR (TupRsingle arrR) arr = rnfArray arrR arr rnfArraysR (TupRpair aR1 aR2) (a1,a2) = rnfArraysR aR1 a1 `seq` rnfArraysR aR2 a2 liftArrayR :: ArrayR a -> CodeQ (ArrayR a) -liftArrayR (ArrayR shR tR) = [|| ArrayR $$(liftShapeR shR) $$(liftTypeR tR) ||] +liftArrayR (ArrayR shR tp) = [|| ArrayR $$(liftShapeR shR) $$(liftTypeR tp) ||] liftArraysR :: ArraysR arrs -> CodeQ (ArraysR arrs) liftArraysR TupRunit = [|| TupRunit ||] @@ -312,12 +317,9 @@ liftArraysR (TupRsingle repr) = [|| TupRsingle $$(liftArrayR repr) ||] liftArraysR (TupRpair a b) = [|| TupRpair $$(liftArraysR a) $$(liftArraysR b) ||] liftArray :: forall sh e. ArrayR (Array sh e) -> Array sh e -> CodeQ (Array sh e) -liftArray (ArrayR shR adR) (Array sh adata) = - [|| Array $$(liftElt (shapeType shR) sh) $$(liftArrayData sz adR adata) ||] `at` [t| Array $(liftTypeQ (shapeType shR)) $(liftTypeQ adR) |] +liftArray (ArrayR shR tp) (Array sh buffers) = + [|| Array $$(liftElt (shapeType shR) sh) $$(liftBuffers tp buffers) ||] `at` [t| Array $(liftTypeQ (shapeType shR)) $(liftTypeQ tp) |] where - sz :: Int - sz = size shR sh - at :: CodeQ t -> Q Type -> CodeQ t at e t = unsafeCodeCoerce $ sigE (unTypeCode e) t diff --git a/src/Data/Array/Accelerate/Representation/Elt.hs b/src/Data/Array/Accelerate/Representation/Elt.hs index d72cd7165..d805a36e9 100644 --- a/src/Data/Array/Accelerate/Representation/Elt.hs +++ b/src/Data/Array/Accelerate/Representation/Elt.hs @@ -1,7 +1,5 @@ -{-# LANGUAGE GADTs #-} -{-# LANGUAGE MagicHash #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TupleSections #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE TemplateHaskellQuotes #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.Representation.Elt @@ -41,7 +39,7 @@ undefElt = tuple vector :: VectorType t -> t vector (VectorType n t) = runST $ do - mba <- newByteArray (n * bytesElt (TupRsingle (SingleScalarType t))) + mba <- newByteArray (n * singleTypeSize t) ByteArray ba# <- unsafeFreezeByteArray mba return (Vec ba#) @@ -69,44 +67,35 @@ undefElt = tuple floating TypeFloat = 0 floating TypeDouble = 0 -bytesElt :: TypeR e -> Int -bytesElt = tuple - where - tuple :: TypeR t -> Int - tuple TupRunit = 0 - tuple (TupRpair ta tb) = tuple ta + tuple tb - tuple (TupRsingle t) = scalar t - - scalar :: ScalarType t -> Int - scalar (SingleScalarType t) = single t - scalar (VectorScalarType t) = vector t - - vector :: VectorType t -> Int - vector (VectorType n t) = n * single t - - single :: SingleType t -> Int - single (NumSingleType t) = num t - - num :: NumType t -> Int - num (IntegralNumType t) = integral t - num (FloatingNumType t) = floating t - - integral :: IntegralType t -> Int - integral TypeInt = sizeOf (undefined::Int) - integral TypeInt8 = 1 - integral TypeInt16 = 2 - integral TypeInt32 = 4 - integral TypeInt64 = 8 - integral TypeWord = sizeOf (undefined::Word) - integral TypeWord8 = 1 - integral TypeWord16 = 2 - integral TypeWord32 = 4 - integral TypeWord64 = 8 - - floating :: FloatingType t -> Int - floating TypeHalf = 2 - floating TypeFloat = 4 - floating TypeDouble = 8 +scalarTypeSize :: ScalarType e -> Int +scalarTypeSize = fst . scalarTypeSizeAlignment + +scalarTypeAlignment :: ScalarType e -> Int +scalarTypeAlignment = snd . scalarTypeSizeAlignment + +scalarTypeSizeAlignment :: ScalarType e -> (Int, Int) +scalarTypeSizeAlignment (SingleScalarType t) = (sz, sz) + where sz = singleTypeSize t +scalarTypeSizeAlignment (VectorScalarType (VectorType n t)) = (n * sz, sz) + where sz = singleTypeSize t + +-- For all SingleTypes, the alignment is equal to the size +singleTypeSize :: SingleType e -> Int +singleTypeSize (NumSingleType (IntegralNumType tp)) = case tp of + TypeInt -> sizeOf (undefined::Int) + TypeInt8 -> 1 + TypeInt16 -> 2 + TypeInt32 -> 4 + TypeInt64 -> 8 + TypeWord -> sizeOf (undefined::Word) + TypeWord8 -> 1 + TypeWord16 -> 2 + TypeWord32 -> 4 + TypeWord64 -> 8 +singleTypeSize (NumSingleType (FloatingNumType tp)) = case tp of + TypeHalf -> 2 + TypeFloat -> 4 + TypeDouble -> 8 showElt :: TypeR e -> e -> String showElt t v = showsElt t v "" diff --git a/src/Data/Array/Accelerate/Representation/Ground.hs b/src/Data/Array/Accelerate/Representation/Ground.hs new file mode 100644 index 000000000..6a2e42972 --- /dev/null +++ b/src/Data/Array/Accelerate/Representation/Ground.hs @@ -0,0 +1,104 @@ +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} +{-# OPTIONS_HADDOCK hide #-} +-- | +-- Module : Data.Array.Accelerate.Representation.Ground +-- Copyright : [2008..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.Representation.Ground + where + +import Data.Array.Accelerate.Representation.Array +import Data.Array.Accelerate.Representation.Shape +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Type +import Data.Type.Equality + +-- | Ground values are buffers or scalars. +-- +data GroundR a where + GroundRbuffer :: ScalarType e -> GroundR (Buffer e) + GroundRscalar :: ScalarType e -> GroundR e + +instance Distributes GroundR where + reprIsSingle (GroundRbuffer _) = Refl + reprIsSingle (GroundRscalar tp) = reprIsSingle tp + + pairImpossible (GroundRscalar tp) = pairImpossible tp + unitImpossible (GroundRscalar tp) = unitImpossible tp + +-- | Tuples of ground values +-- +type GroundsR = TupR GroundR + +rnfGroundR :: GroundR t -> () +rnfGroundR (GroundRscalar tp) = rnfScalarType tp +rnfGroundR (GroundRbuffer tp) = rnfScalarType tp + +rnfGroundsR :: GroundsR t -> () +rnfGroundsR = rnfTupR rnfGroundR + +-- | Conversion from arrays representation to grounds representation +lowerArraysR :: ArraysR arr -> GroundsR (LoweredArrays arr) +lowerArraysR TupRunit = TupRunit +lowerArraysR (TupRsingle repr) = lowerArrayR repr +lowerArraysR (TupRpair r1 r2) = lowerArraysR r1 `TupRpair` lowerArraysR r2 + +lowerArrayR :: ArrayR arr -> GroundsR (LoweredArrays arr) +lowerArrayR (ArrayR shr tp) = mapTupR GroundRscalar (shapeType shr) `TupRpair` buffersR tp + +buffersR :: forall e. TypeR e -> GroundsR (Buffers e) +buffersR TupRunit = TupRunit +buffersR (TupRsingle tp) + | Refl <- reprIsSingle @ScalarType @e @Buffer tp = TupRsingle (GroundRbuffer tp) +buffersR (TupRpair t1 t2) = buffersR t1 `TupRpair` buffersR t2 + +-- | Utilities for working with GroundsR +typeRtoGroundsR :: TypeR t -> GroundsR t +typeRtoGroundsR = mapTupR GroundRscalar + +bufferImpossible :: ScalarType (Buffer e) -> a +bufferImpossible (SingleScalarType (NumSingleType (IntegralNumType tp))) = case tp of {} +bufferImpossible (SingleScalarType (NumSingleType (FloatingNumType tp))) = case tp of {} + +groundFunctionImpossible :: GroundsR (s -> t) -> a +groundFunctionImpossible (TupRsingle (GroundRscalar t)) = functionImpossible (TupRsingle t) + +groundRelt :: GroundR (Buffer t) -> ScalarType t +groundRelt (GroundRbuffer tp) = tp +groundRelt (GroundRscalar tp) = bufferImpossible tp + +type family LoweredArrays a where + LoweredArrays () = () + LoweredArrays (a, b) = (LoweredArrays a, LoweredArrays b) + LoweredArrays (Array sh e) = (sh, Buffers e) + +type family LoweredAfun a where + LoweredAfun (a -> b) = LoweredArrays a -> LoweredAfun b + LoweredAfun a = LoweredArrays a + +loweredAfunIsBody :: ArraysR a -> LoweredAfun a :~: LoweredArrays a +loweredAfunIsBody (TupRsingle ArrayR{}) = Refl +loweredAfunIsBody TupRunit = Refl +loweredAfunIsBody (TupRpair _ _) = Refl + +lowerArrays :: ArraysR a -> a -> LoweredArrays a +lowerArrays TupRunit () = () +lowerArrays (TupRpair r1 r2) (a1, a2) = (lowerArrays r1 a1, lowerArrays r2 a2) +lowerArrays (TupRsingle ArrayR{}) (Array sh buffers) = (sh, buffers) + +sugarArrays :: ArraysR a -> LoweredArrays a -> a +sugarArrays TupRunit () = () +sugarArrays (TupRpair r1 r2) (d1, d2) = (sugarArrays r1 d1, sugarArrays r2 d2) +sugarArrays (TupRsingle ArrayR{}) (sh, buffers) = Array sh buffers + +data GFunctionR t where + GFunctionRlam :: GroundsR t -> GFunctionR s -> GFunctionR (t -> s) + GFunctionRbody :: GroundsR t -> GFunctionR t + diff --git a/src/Data/Array/Accelerate/Representation/Shape.hs b/src/Data/Array/Accelerate/Representation/Shape.hs index fa3651c03..b8126db9f 100644 --- a/src/Data/Array/Accelerate/Representation/Shape.hs +++ b/src/Data/Array/Accelerate/Representation/Shape.hs @@ -1,7 +1,4 @@ -{-# LANGUAGE GADTs #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TupleSections #-} +{-# LANGUAGE TemplateHaskellQuotes #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.Representation.Shape @@ -16,6 +13,7 @@ module Data.Array.Accelerate.Representation.Shape where +import Data.Array.Accelerate.AST.LeftHandSide ( Exists(..) ) import Data.Array.Accelerate.Error import Data.Array.Accelerate.Type import Data.Array.Accelerate.Representation.Type @@ -62,6 +60,12 @@ rank :: ShapeR sh -> Int rank ShapeRz = 0 rank (ShapeRsnoc shr) = rank shr + 1 +shapeRFromRank :: Int -> Exists ShapeR +shapeRFromRank 0 = Exists ShapeRz +shapeRFromRank r + | Exists shr <- shapeRFromRank (r - 1) + = Exists $ ShapeRsnoc shr + -- | Total number of elements in an array of the given shape -- size :: ShapeR sh -> sh -> Int @@ -183,6 +187,13 @@ shapeType (ShapeRsnoc shr) = `TupRpair` TupRsingle (SingleScalarType (NumSingleType (IntegralNumType TypeInt))) +typeShape :: TypeR sh -> Maybe (ShapeR sh) +typeShape TupRunit = Just ShapeRz +typeShape (tup `TupRpair` TupRsingle (SingleScalarType (NumSingleType (IntegralNumType TypeInt)))) + | Just sh <- typeShape tup + = Just $ ShapeRsnoc sh +typeShape _ = Nothing + rnfShape :: ShapeR sh -> sh -> () rnfShape ShapeRz () = () rnfShape (ShapeRsnoc shr) (sh, s) = s `seq` rnfShape shr sh diff --git a/src/Data/Array/Accelerate/Representation/Slice.hs b/src/Data/Array/Accelerate/Representation/Slice.hs index dee059a37..8b94d4669 100644 --- a/src/Data/Array/Accelerate/Representation/Slice.hs +++ b/src/Data/Array/Accelerate/Representation/Slice.hs @@ -1,9 +1,5 @@ -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TemplateHaskellQuotes #-} +{-# LANGUAGE TypeFamilies #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.Representation.Slice @@ -19,6 +15,8 @@ module Data.Array.Accelerate.Representation.Slice where import Data.Array.Accelerate.Representation.Shape +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Type import Language.Haskell.TH.Extra @@ -76,6 +74,11 @@ sliceDomain SliceNil () () = () sliceDomain (SliceAll slix) (slx, ()) (sl, sz) = (sliceDomain slix slx sl, sz) sliceDomain (SliceFixed slix) (slx, sz) sl = (sliceDomain slix slx sl, sz) +sliceEltR :: SliceIndex slix sl co dim -> TypeR slix +sliceEltR SliceNil = TupRunit +sliceEltR (SliceAll sl) = sliceEltR sl `TupRpair` TupRunit +sliceEltR (SliceFixed sl) = sliceEltR sl `TupRpair` TupRsingle scalarTypeInt + sliceShapeR :: SliceIndex slix sl co dim -> ShapeR sl sliceShapeR SliceNil = ShapeRz sliceShapeR (SliceAll sl) = ShapeRsnoc $ sliceShapeR sl diff --git a/src/Data/Array/Accelerate/Representation/Stencil.hs b/src/Data/Array/Accelerate/Representation/Stencil.hs index dd546721c..66d2ee7ed 100644 --- a/src/Data/Array/Accelerate/Representation/Stencil.hs +++ b/src/Data/Array/Accelerate/Representation/Stencil.hs @@ -1,5 +1,4 @@ -{-# LANGUAGE GADTs #-} -{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TemplateHaskellQuotes #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.Representation.Stencil diff --git a/src/Data/Array/Accelerate/Representation/Tag.hs b/src/Data/Array/Accelerate/Representation/Tag.hs index 52b1b5527..ccdb65a12 100644 --- a/src/Data/Array/Accelerate/Representation/Tag.hs +++ b/src/Data/Array/Accelerate/Representation/Tag.hs @@ -1,5 +1,4 @@ -{-# LANGUAGE GADTs #-} -{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TemplateHaskellQuotes #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.Representation.Tag diff --git a/src/Data/Array/Accelerate/Representation/Type.hs b/src/Data/Array/Accelerate/Representation/Type.hs index 9b154133f..55f9fa4ba 100644 --- a/src/Data/Array/Accelerate/Representation/Type.hs +++ b/src/Data/Array/Accelerate/Representation/Type.hs @@ -1,11 +1,6 @@ -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE QuantifiedConstraints #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE StandaloneDeriving #-} -{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE QuantifiedConstraints #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeFamilyDependencies #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.Representation.Type @@ -21,7 +16,9 @@ module Data.Array.Accelerate.Representation.Type where import Data.Array.Accelerate.Type +import Data.Array.Accelerate.Error import Data.Primitive.Vec +import Data.Type.Equality import Formatting import Language.Haskell.TH.Extra @@ -47,6 +44,14 @@ data TupR s a where deriving instance (forall a. Show (s a)) => Show (TupR s t) +instance (forall a. Semigroup (s a)) => Semigroup (TupR s t) where + (<>) :: (forall a. Semigroup (s a)) => TupR s t -> TupR s t -> TupR s t + (<>) TupRunit TupRunit = TupRunit + (<>) (TupRsingle a) (TupRsingle b) = TupRsingle (a <> b) + (<>) (TupRpair a1 a2) (TupRpair b1 b2) = TupRpair (a1 <> b1) (a2 <> b2) + (<>) _ _ = internalError "Semigroup TupR: tuple mismatch" + + formatTypeR :: Format r (TypeR a -> r) formatTypeR = later $ \case TupRunit -> "()" @@ -55,6 +60,75 @@ formatTypeR = later $ \case type TypeR = TupR ScalarType +-- class IsTypeR a where +-- typeR :: TypeR a + +-- instance (IsTypeR a, IsTypeR b) => IsTypeR (a,b) where +-- typeR = TupRpair typeR typeR + +-- instance IsTypeR () where +-- typeR = TupRunit + +-- instance IsScalar a => IsTypeR a where +-- typeR = TupRsingle scalarType + +data TupleIdx s t where + TupleIdxLeft :: TupleIdx l t -> TupleIdx (l, r) t + TupleIdxRight :: TupleIdx r t -> TupleIdx (l, r) t + TupleIdxSelf :: TupleIdx t t + +tupleLeft :: TupleIdx a (b, c) -> TupleIdx a b +tupleLeft (TupleIdxLeft idx) = TupleIdxLeft $ tupleLeft idx +tupleLeft (TupleIdxRight idx) = TupleIdxRight $ tupleLeft idx +tupleLeft TupleIdxSelf = TupleIdxLeft TupleIdxSelf + +tupleRight :: TupleIdx a (b, c) -> TupleIdx a c +tupleRight (TupleIdxLeft idx) = TupleIdxLeft $ tupleRight idx +tupleRight (TupleIdxRight idx) = TupleIdxRight $ tupleRight idx +tupleRight TupleIdxSelf = TupleIdxRight TupleIdxSelf + +-- | Distributes a type constructor over the elements of a tuple. +-- TODO: Could we make this type class injective? Then we wouldn't +-- need the type class Distributes any more. +-- Note that we must use a standard, lazy pair here, as we rely on +-- laziness in type alias Buffers to make host-device copies lazy. +-- +type family Distribute f a = b where + Distribute f () = () + Distribute f (a, b) = (Distribute f a, Distribute f b) + Distribute f a = f a + +class Distributes s where + -- Shows that a single element isn't unit or a pair + reprIsSingle :: s t -> Distribute f t :~: f t + pairImpossible :: s (t, v) -> a + unitImpossible :: s () -> a + +instance Distributes ScalarType where + reprIsSingle (VectorScalarType _) = Refl + reprIsSingle (SingleScalarType (NumSingleType tp)) = case tp of + IntegralNumType TypeInt -> Refl + IntegralNumType TypeInt8 -> Refl + IntegralNumType TypeInt16 -> Refl + IntegralNumType TypeInt32 -> Refl + IntegralNumType TypeInt64 -> Refl + IntegralNumType TypeWord -> Refl + IntegralNumType TypeWord8 -> Refl + IntegralNumType TypeWord16 -> Refl + IntegralNumType TypeWord32 -> Refl + IntegralNumType TypeWord64 -> Refl + FloatingNumType TypeHalf -> Refl + FloatingNumType TypeFloat -> Refl + FloatingNumType TypeDouble -> Refl + + pairImpossible (SingleScalarType (NumSingleType tp)) + | IntegralNumType t <- tp = case t of {} + | FloatingNumType t <- tp = case t of {} + + unitImpossible (SingleScalarType (NumSingleType tp)) + | IntegralNumType t <- tp = case t of {} + | FloatingNumType t <- tp = case t of {} + rnfTupR :: (forall b. s b -> ()) -> TupR s a -> () rnfTupR _ TupRunit = () rnfTupR f (TupRsingle s) = f s @@ -63,6 +137,22 @@ rnfTupR f (TupRpair a b) = rnfTupR f a `seq` rnfTupR f b rnfTypeR :: TypeR t -> () rnfTypeR = rnfTupR rnfScalarType +tupleSize :: TupR b t -> Int +tupleSize TupRunit = 0 +tupleSize (TupRsingle _) = 1 +tupleSize (TupRpair a b) = tupleSize a + tupleSize b + +tupleIdxToInt :: TupR b s -> TupleIdx s t -> Int +tupleIdxToInt (TupRpair a _) (TupleIdxLeft ix) = tupleIdxToInt a ix +tupleIdxToInt (TupRpair a b) (TupleIdxRight ix) = tupleSize a + tupleIdxToInt b ix +tupleIdxToInt _ _ = 0 + +prjTupleIdxR :: TupR b s -> TupleIdx s t -> TupR b t +prjTupleIdxR (TupRpair a _) (TupleIdxLeft ix) = prjTupleIdxR a ix +prjTupleIdxR (TupRpair _ b) (TupleIdxRight ix) = prjTupleIdxR b ix +prjTupleIdxR t TupleIdxSelf = t +prjTupleIdxR _ _ = internalError "Tuple mismatch" + liftTupR :: (forall b. s b -> CodeQ (s b)) -> TupR s a -> CodeQ (TupR s a) liftTupR _ TupRunit = [|| TupRunit ||] liftTupR f (TupRsingle s) = [|| TupRsingle $$(f s) ||] @@ -124,3 +214,37 @@ runQ $ in mapM mkT [2..16] +mapTupR :: (forall s. a s -> b s) -> TupR a t -> TupR b t +mapTupR f (TupRsingle a) = TupRsingle $ f a +mapTupR _ TupRunit = TupRunit +mapTupR f (TupRpair a1 a2) = mapTupR f a1 `TupRpair` mapTupR f a2 + +traverseTupR :: Applicative f => (forall s. a s -> f (b s)) -> TupR a t -> f (TupR b t) +traverseTupR f (TupRsingle a) = TupRsingle <$> f a +traverseTupR _ TupRunit = pure TupRunit +traverseTupR f (TupRpair a1 a2) = TupRpair <$> traverseTupR f a1 <*> traverseTupR f a2 + +foldMapTupR :: Monoid m => (forall s. a s -> m) -> TupR a t -> m +foldMapTupR f (TupRsingle a) = f a +foldMapTupR _ TupRunit = mempty +foldMapTupR f (TupRpair a1 a2) = foldMapTupR f a1 <> foldMapTupR f a2 + +foldMapMTupR :: (Monad m, Monoid r) => (forall s. a s -> m r) -> TupR a t -> m r +foldMapMTupR f (TupRsingle a) = f a +foldMapMTupR _ TupRunit = return mempty +foldMapMTupR f (TupRpair a1 a2) = do + r1 <- foldMapMTupR f a1 + r2 <- foldMapMTupR f a2 + return (r1 <> r2) + +functionImpossible :: TypeR (s -> t) -> a +functionImpossible (TupRsingle (SingleScalarType (NumSingleType tp))) = case tp of + IntegralNumType t -> case t of {} + FloatingNumType t -> case t of {} + +instance (forall a. Eq (s a)) => Eq (TupR s t) where + (==) :: TupR s t -> TupR s t -> Bool + (==) TupRunit TupRunit = True + (==) (TupRsingle l) (TupRsingle r) = l == r + (==) (TupRpair l1 l2) (TupRpair r1 r2) = l1 == r1 && l2 == r2 + (==) _ _ = False diff --git a/src/Data/Array/Accelerate/Representation/Vec.hs b/src/Data/Array/Accelerate/Representation/Vec.hs index 35eac3b6c..19c1c1215 100644 --- a/src/Data/Array/Accelerate/Representation/Vec.hs +++ b/src/Data/Array/Accelerate/Representation/Vec.hs @@ -1,10 +1,5 @@ -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE KindSignatures #-} -{-# LANGUAGE MagicHash #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE TemplateHaskellQuotes #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.Representation.Vec diff --git a/src/Data/Array/Accelerate/Smart.hs b/src/Data/Array/Accelerate/Smart.hs index bf0e85873..2c02cbfa6 100644 --- a/src/Data/Array/Accelerate/Smart.hs +++ b/src/Data/Array/Accelerate/Smart.hs @@ -1,17 +1,5 @@ -{-# LANGUAGE AllowAmbiguousTypes #-} -{-# LANGUAGE CPP #-} -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE TypeFamilies #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.Smart @@ -75,6 +63,7 @@ module Data.Array.Accelerate.Smart ( ($$), ($$$), ($$$$), ($$$$$), ApplyAcc(..), unAcc, unAccFunction, mkExp, unExp, unExpFunction, unExpBinaryFunction, unPair, mkPairToTuple, + prj0, prj1, prjTail, -- ** Miscellaneous formatPreAccOp, @@ -104,11 +93,11 @@ import qualified Data.Array.Accelerate.Sugar.Shape as Sugar import Data.Array.Accelerate.AST ( Direction(..), Message(..) , PrimBool, PrimMaybe - , PrimFun(..), primFunType - , PrimConst(..), primConstType ) + , PrimFun(..), Cmp(..), primFunType ) import Data.Primitive.Vec import Data.Kind +import Data.Text ( Text ) import Data.Text.Lazy.Builder import Formatting @@ -320,13 +309,8 @@ data PreSmartAcc acc exp as where -> Level -- environment size at defining occurrence -> PreSmartAcc acc exp as - Pipe :: ArraysR as - -> ArraysR bs - -> ArraysR cs - -> (SmartAcc as -> acc bs) - -> (SmartAcc bs -> acc cs) - -> acc as - -> PreSmartAcc acc exp cs + Manifest :: acc as + -> PreSmartAcc acc exp as Aforeign :: Foreign asm => ArraysR bs @@ -361,6 +345,15 @@ data PreSmartAcc acc exp as where -> acc arrs2 -> PreSmartAcc acc exp arrs2 + Aassert :: Text + -> exp PrimBool + -> acc as + -> PreSmartAcc acc exp as + + Aassume :: exp PrimBool + -> acc as + -> PreSmartAcc acc exp as + Use :: ArrayR (Array sh e) -> Array sh e -> PreSmartAcc acc exp (Array sh e) @@ -432,10 +425,9 @@ data PreSmartAcc acc exp as where -> PreSmartAcc acc exp (Array (sh, Int) e, Array sh e) Permute :: ArrayR (Array sh e) - -> (SmartExp e -> SmartExp e -> exp e) + -> Maybe (SmartExp e -> SmartExp e -> exp e) -> acc (Array sh' e) - -> (SmartExp sh -> exp (PrimMaybe sh')) - -> acc (Array sh e) + -> acc (Array sh (PrimMaybe (((),sh'), e))) -- instead `PrimMaybe (sh', e)` would be nice, but hard to link to the frontend through EltR. -> PreSmartAcc acc exp (Array sh' e) Backpermute :: ShapeR sh' @@ -539,15 +531,17 @@ data PreSmartExp acc exp t where -> exp t -> PreSmartExp acc exp t + Select :: exp PrimBool + -> exp t + -> exp t + -> PreSmartExp acc exp t + While :: TypeR t -> (SmartExp t -> exp PrimBool) -> (SmartExp t -> exp t) -> exp t -> PreSmartExp acc exp t - PrimConst :: PrimConst t - -> PreSmartExp acc exp t - PrimApp :: PrimFun (a -> r) -> exp a -> PreSmartExp acc exp r @@ -586,6 +580,15 @@ data PreSmartExp acc exp t where -> exp a -> PreSmartExp acc exp b + Assert :: Text + -> exp PrimBool + -> exp a + -> PreSmartExp acc exp a + + Assume :: exp PrimBool + -> exp a + -> PreSmartExp acc exp a + -- Smart constructors for stencils -- ------------------------------- @@ -826,7 +829,7 @@ arrayR acc = case arraysR acc of instance HasArraysR acc => HasArraysR (PreSmartAcc acc exp) where arraysR = \case Atag repr _ -> repr - Pipe _ _ repr _ _ _ -> repr + Manifest as -> arraysR as Aforeign repr _ _ _ -> repr Acond _ a _ -> arraysR a Awhile _ _ _ a -> arraysR a @@ -838,6 +841,8 @@ instance HasArraysR acc => HasArraysR (PreSmartAcc acc exp) where PairIdxRight -> t2 Aprj _ _ -> error "Ejector seat? You're joking!" Atrace _ _ a -> arraysR a + Aassert _ _ a -> arraysR a + Aassume _ a -> arraysR a Use repr _ -> TupRsingle repr Unit tp _ -> TupRsingle $ ArrayR ShapeRz $ tp Generate repr _ _ -> TupRsingle repr @@ -857,7 +862,7 @@ instance HasArraysR acc => HasArraysR (PreSmartAcc acc exp) where Scan _ _ _ _ a -> arraysR a Scan' _ _ _ _ a -> let repr@(ArrayR (ShapeRsnoc shr) tp) = arrayR a in TupRsingle repr `TupRpair` TupRsingle (ArrayR shr tp) - Permute _ _ a _ _ -> arraysR a + Permute _ _ a _ -> arraysR a Backpermute shr _ _ a -> let ArrayR _ tp = arrayR a in TupRsingle (ArrayR shr tp) Stencil s tp _ _ _ -> TupRsingle $ ArrayR (stencilShapeR s) tp @@ -889,8 +894,8 @@ instance HasTypeR exp => HasTypeR (PreSmartExp acc exp) where Case _ ((_,c):_) -> typeR c Case{} -> internalError "encountered empty case" Cond _ e _ -> typeR e + Select _ e _ -> typeR e While t _ _ _ -> t - PrimConst c -> TupRsingle $ SingleScalarType $ primConstType c PrimApp f _ -> snd $ primFunType f Index tp _ _ -> tp LinearIndex tp _ _ -> tp @@ -899,6 +904,8 @@ instance HasTypeR exp => HasTypeR (PreSmartExp acc exp) where Foreign tp _ _ _ -> tp Undef tp -> TupRsingle tp Coerce _ tp _ -> TupRsingle tp + Assert _ _ e -> typeR e + Assume _ e -> typeR e -- Smart constructors @@ -927,6 +934,14 @@ constant = Exp . go (eltR @e) . fromElt -- | 'undef' can be used anywhere a constant is expected, and indicates that the -- consumer of the value can receive an unspecified bit pattern. -- +-- 'undef' is not a pure value: each observation of it may yield a different +-- result (or not), depending on unpredictable optimisation choices in the +-- compiler, and computing with 'undef' produces 'undef'. However, computing with +-- 'undef' /is/ valid as long as each possible bit pattern would be a valid input +-- to that computation. For example, @'undef' - 'undef'@ may have any result, but +-- will not crash. This is even the case after binding 'undef' to a variable, and +-- thus @let x = 'undef' in x - x@ need not evaluate to zero. +-- -- This is useful because a store of an undefined value can be assumed to not -- have any effect; we can assume that the value is overwritten with bits that -- happen to match what was already there. However, a store /to/ an undefined @@ -977,15 +992,23 @@ indexTail (Exp x) = mkExp $ Prj PairIdxLeft x -- Smart constructor for constants -- -mkMinBound :: (Elt t, IsBounded (EltR t)) => Exp t -mkMinBound = mkExp $ PrimConst (PrimMinBound boundedType) - -mkMaxBound :: (Elt t, IsBounded (EltR t)) => Exp t -mkMaxBound = mkExp $ PrimConst (PrimMaxBound boundedType) +mkMinBound :: forall t. (Elt t, IsBounded (EltR t)) => Exp t +mkMinBound + | IntegralBoundedType tp <- boundedType @(EltR t) + , IntegralDict <- integralDict tp + = Exp $ SmartExp $ Const (SingleScalarType $ NumSingleType $ IntegralNumType tp) minBound -mkPi :: (Elt r, IsFloating (EltR r)) => Exp r -mkPi = mkExp $ PrimConst (PrimPi floatingType) +mkMaxBound :: forall t. (Elt t, IsBounded (EltR t)) => Exp t +mkMaxBound + | IntegralBoundedType tp <- boundedType @(EltR t) + , IntegralDict <- integralDict tp + = Exp $ SmartExp $ Const (SingleScalarType $ NumSingleType $ IntegralNumType tp) maxBound +mkPi :: forall r. (Elt r, IsFloating (EltR r)) => Exp r +mkPi + | tp <- floatingType @(EltR r) + , FloatingDict <- floatingDict tp + = Exp $ SmartExp $ Const (SingleScalarType $ NumSingleType $ FloatingNumType tp) pi -- Smart constructors for primitive applications -- @@ -1161,22 +1184,22 @@ mkIsInfinite = mkPrimUnaryBool $ PrimIsInfinite floatingType -- Relational and equality operators mkLt :: (Elt t, IsSingle (EltR t)) => Exp t -> Exp t -> Exp Bool -mkLt = mkPrimBinaryBool $ PrimLt singleType +mkLt = mkPrimBinaryBool $ PrimCmp singleType CmpLt mkGt :: (Elt t, IsSingle (EltR t)) => Exp t -> Exp t -> Exp Bool -mkGt = mkPrimBinaryBool $ PrimGt singleType +mkGt = flip mkLt mkLtEq :: (Elt t, IsSingle (EltR t)) => Exp t -> Exp t -> Exp Bool -mkLtEq = mkPrimBinaryBool $ PrimLtEq singleType +mkLtEq = flip mkGtEq mkGtEq :: (Elt t, IsSingle (EltR t)) => Exp t -> Exp t -> Exp Bool -mkGtEq = mkPrimBinaryBool $ PrimGtEq singleType +mkGtEq = mkPrimBinaryBool $ PrimCmp singleType CmpGtEq mkEq :: (Elt t, IsSingle (EltR t)) => Exp t -> Exp t -> Exp Bool -mkEq = mkPrimBinaryBool $ PrimEq singleType +mkEq = mkPrimBinaryBool $ PrimCmp singleType CmpEq mkNEq :: (Elt t, IsSingle (EltR t)) => Exp t -> Exp t -> Exp Bool -mkNEq = mkPrimBinaryBool $ PrimNEq singleType +mkNEq = mkPrimBinaryBool $ PrimCmp singleType CmpNEq mkMax :: (Elt t, IsSingle (EltR t)) => Exp t -> Exp t -> Exp t mkMax = mkPrimBinary $ PrimMax singleType @@ -1347,13 +1370,15 @@ formatPreAccOp :: Format r (PreSmartAcc acc exp arrs -> r) formatPreAccOp = later $ \case Atag _ i -> bformat ("Atag " % int) i Use aR a -> bformat ("Use " % string) (showArrayShort 5 (showsElt (arrayRtype aR)) aR a) - Pipe{} -> "Pipe" + Manifest{} -> "Manifest" Acond{} -> "Acond" Awhile{} -> "Awhile" Apair{} -> "Apair" Anil{} -> "Anil" Aprj{} -> "Aprj" Atrace{} -> "Atrace" + Aassert{} -> "Aassert" + Aassume{} -> "Aassume" Unit{} -> "Unit" Generate{} -> "Generate" Reshape{} -> "Reshape" @@ -1386,8 +1411,8 @@ formatPreExpOp = later $ \case FromIndex{} -> "FromIndex" Case{} -> "Case" Cond{} -> "Cond" + Select{} -> "Select" While{} -> "While" - PrimConst{} -> "PrimConst" PrimApp{} -> "PrimApp" Index{} -> "Index" LinearIndex{} -> "LinearIndex" @@ -1395,4 +1420,6 @@ formatPreExpOp = later $ \case ShapeSize{} -> "ShapeSize" Foreign{} -> "Foreign" Coerce{} -> "Coerce" + Assert{} -> "Assert" + Assume{} -> "Assume" diff --git a/src/Data/Array/Accelerate/Sugar/Array.hs b/src/Data/Array/Accelerate/Sugar/Array.hs index 87a490246..ca7bcd416 100644 --- a/src/Data/Array/Accelerate/Sugar/Array.hs +++ b/src/Data/Array/Accelerate/Sugar/Array.hs @@ -1,12 +1,8 @@ {-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE CPP #-} {-# LANGUAGE DefaultSignatures #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} {-# LANGUAGE UndecidableInstances #-} {-# OPTIONS_HADDOCK hide #-} -- | @@ -305,6 +301,8 @@ instance (Shape sh, Elt e) => Arrays (Array sh e) where fromArr (Array arr) = arr toArr = Array +#ifndef __GHCIDE__ + runQ $ do let mkTuple :: Int -> Q Dec @@ -319,3 +317,146 @@ runQ $ do -- mapM mkTuple [2..16] +#else + +instance (Arrays x0, Arrays x1) => Arrays (x0, x1) +instance (Arrays x0, Arrays x1, Arrays x2) => Arrays (x0, x1, x2) +instance (Arrays x0, Arrays x1, Arrays x2, Arrays x3) => + Arrays (x0, x1, x2, x3) +instance (Arrays x0, Arrays x1, Arrays x2, Arrays x3, Arrays x4) => + Arrays (x0, x1, x2, x3, x4) +instance (Arrays x0, + Arrays x1, + Arrays x2, + Arrays x3, + Arrays x4, + Arrays x5) => + Arrays (x0, x1, x2, x3, x4, x5) +instance (Arrays x0, + Arrays x1, + Arrays x2, + Arrays x3, + Arrays x4, + Arrays x5, + Arrays x6) => + Arrays (x0, x1, x2, x3, x4, x5, x6) +instance (Arrays x0, + Arrays x1, + Arrays x2, + Arrays x3, + Arrays x4, + Arrays x5, + Arrays x6, + Arrays x7) => + Arrays (x0, x1, x2, x3, x4, x5, x6, x7) +instance (Arrays x0, + Arrays x1, + Arrays x2, + Arrays x3, + Arrays x4, + Arrays x5, + Arrays x6, + Arrays x7, + Arrays x8) => + Arrays (x0, x1, x2, x3, x4, x5, x6, x7, x8) +instance (Arrays x0, + Arrays x1, + Arrays x2, + Arrays x3, + Arrays x4, + Arrays x5, + Arrays x6, + Arrays x7, + Arrays x8, + Arrays x9) => + Arrays (x0, x1, x2, x3, x4, x5, x6, x7, x8, x9) +instance (Arrays x0, + Arrays x1, + Arrays x2, + Arrays x3, + Arrays x4, + Arrays x5, + Arrays x6, + Arrays x7, + Arrays x8, + Arrays x9, + Arrays x10) => + Arrays (x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10) +instance (Arrays x0, + Arrays x1, + Arrays x2, + Arrays x3, + Arrays x4, + Arrays x5, + Arrays x6, + Arrays x7, + Arrays x8, + Arrays x9, + Arrays x10, + Arrays x11) => + Arrays (x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11) +instance (Arrays x0, + Arrays x1, + Arrays x2, + Arrays x3, + Arrays x4, + Arrays x5, + Arrays x6, + Arrays x7, + Arrays x8, + Arrays x9, + Arrays x10, + Arrays x11, + Arrays x12) => + Arrays (x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12) +instance (Arrays x0, + Arrays x1, + Arrays x2, + Arrays x3, + Arrays x4, + Arrays x5, + Arrays x6, + Arrays x7, + Arrays x8, + Arrays x9, + Arrays x10, + Arrays x11, + Arrays x12, + Arrays x13) => + Arrays (x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13) +instance (Arrays x0, + Arrays x1, + Arrays x2, + Arrays x3, + Arrays x4, + Arrays x5, + Arrays x6, + Arrays x7, + Arrays x8, + Arrays x9, + Arrays x10, + Arrays x11, + Arrays x12, + Arrays x13, + Arrays x14) => + Arrays (x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13, + x14) +instance (Arrays x0, + Arrays x1, + Arrays x2, + Arrays x3, + Arrays x4, + Arrays x5, + Arrays x6, + Arrays x7, + Arrays x8, + Arrays x9, + Arrays x10, + Arrays x11, + Arrays x12, + Arrays x13, + Arrays x14, + Arrays x15) => + Arrays (x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13, + x14, x15) +#endif diff --git a/src/Data/Array/Accelerate/Sugar/Elt.hs b/src/Data/Array/Accelerate/Sugar/Elt.hs index bd44c2294..7ee4a76e9 100644 --- a/src/Data/Array/Accelerate/Sugar/Elt.hs +++ b/src/Data/Array/Accelerate/Sugar/Elt.hs @@ -1,14 +1,8 @@ {-# LANGUAGE AllowAmbiguousTypes #-} -{-# LANGUAGE DataKinds #-} +{-# LANGUAGE CPP #-} {-# LANGUAGE DefaultSignatures #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TupleSections #-} -{-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} {-# LANGUAGE UndecidableInstances #-} {-# OPTIONS_HADDOCK hide #-} -- | @@ -369,6 +363,8 @@ instance Elt Char where toElt = chr . fromIntegral fromElt = fromIntegral . ord +#ifndef __GHCIDE__ + runQ $ do let -- XXX: we might want to do the digItOut trick used by FromIntegral? @@ -469,3 +465,296 @@ runQ $ do -- vs <- sequence [ mkVecElt t n | t <- integralTypes ++ floatingTypes, n <- [2,3,4,8,16] ] return (concat ss ++ concat ns ++ ts) +#else + +instance Elt Int where + type EltR Int = Int + eltR = TupRsingle scalarType + tagsR = [TagRsingle scalarType] + fromElt = id + toElt = id +instance Elt Int8 where + type EltR Int8 = Int8 + eltR = TupRsingle scalarType + tagsR = [TagRsingle scalarType] + fromElt = id + toElt = id +instance Elt Int16 where + type EltR Int16 = Int16 + eltR = TupRsingle scalarType + tagsR = [TagRsingle scalarType] + fromElt = id + toElt = id +instance Elt Int32 where + type EltR Int32 = Int32 + eltR = TupRsingle scalarType + tagsR = [TagRsingle scalarType] + fromElt = id + toElt = id +instance Elt Int64 where + type EltR Int64 = Int64 + eltR = TupRsingle scalarType + tagsR = [TagRsingle scalarType] + fromElt = id + toElt = id +instance Elt Word where + type EltR Word = Word + eltR = TupRsingle scalarType + tagsR = [TagRsingle scalarType] + fromElt = id + toElt = id +instance Elt Word8 where + type EltR Word8 = Word8 + eltR = TupRsingle scalarType + tagsR = [TagRsingle scalarType] + fromElt = id + toElt = id +instance Elt Word16 where + type EltR Word16 = Word16 + eltR = TupRsingle scalarType + tagsR = [TagRsingle scalarType] + fromElt = id + toElt = id +instance Elt Word32 where + type EltR Word32 = Word32 + eltR = TupRsingle scalarType + tagsR = [TagRsingle scalarType] + fromElt = id + toElt = id +instance Elt Word64 where + type EltR Word64 = Word64 + eltR = TupRsingle scalarType + tagsR = [TagRsingle scalarType] + fromElt = id + toElt = id +instance Elt Half where + type EltR Half = Half + eltR = TupRsingle scalarType + tagsR = [TagRsingle scalarType] + fromElt = id + toElt = id +instance Elt Float where + type EltR Float = Float + eltR = TupRsingle scalarType + tagsR = [TagRsingle scalarType] + fromElt = id + toElt = id +instance Elt Double where + type EltR Double = Double + eltR = TupRsingle scalarType + tagsR = [TagRsingle scalarType] + fromElt = id + toElt = id +instance Elt CShort where + type EltR CShort = Int16 + eltR = TupRsingle scalarType + tagsR = [TagRsingle scalarType] + fromElt (CShort x) = x + toElt = CShort +instance Elt CUShort where + type EltR CUShort = Word16 + eltR = TupRsingle scalarType + tagsR = [TagRsingle scalarType] + fromElt (CUShort x) = x + toElt = CUShort +instance Elt CInt where + type EltR CInt = Int32 + eltR = TupRsingle scalarType + tagsR = [TagRsingle scalarType] + fromElt (CInt x) = x + toElt = CInt +instance Elt CUInt where + type EltR CUInt = Word32 + eltR = TupRsingle scalarType + tagsR = [TagRsingle scalarType] + fromElt (CUInt x) = x + toElt = CUInt +instance Elt CLong where + type EltR CLong = Int64 + eltR = TupRsingle scalarType + tagsR = [TagRsingle scalarType] + fromElt (CLong x) = x + toElt = CLong +instance Elt CULong where + type EltR CULong = Word64 + eltR = TupRsingle scalarType + tagsR = [TagRsingle scalarType] + fromElt (CULong x) = x + toElt = CULong +instance Elt CLLong where + type EltR CLLong = Int64 + eltR = TupRsingle scalarType + tagsR = [TagRsingle scalarType] + fromElt (CLLong x) = x + toElt = CLLong +instance Elt CULLong where + type EltR CULLong = Word64 + eltR = TupRsingle scalarType + tagsR = [TagRsingle scalarType] + fromElt (CULLong x) = x + toElt = CULLong +instance Elt CFloat where + type EltR CFloat = Float + eltR = TupRsingle scalarType + tagsR = [TagRsingle scalarType] + fromElt (CFloat x) = x + toElt = CFloat +instance Elt CDouble where + type EltR CDouble = Double + eltR = TupRsingle scalarType + tagsR = [TagRsingle scalarType] + fromElt (CDouble x) = x + toElt = CDouble +instance Elt CChar where + type EltR CChar = Int8 + eltR = TupRsingle scalarType + tagsR = [TagRsingle scalarType] + fromElt (CChar x) = x + toElt = CChar +instance Elt CSChar where + type EltR CSChar = Int8 + eltR = TupRsingle scalarType + tagsR = [TagRsingle scalarType] + fromElt (CSChar x) = x + toElt = CSChar +instance Elt CUChar where + type EltR CUChar = Word8 + eltR = TupRsingle scalarType + tagsR = [TagRsingle scalarType] + fromElt (CUChar x) = x + toElt = CUChar +instance (Elt x0, Elt x1) => Elt (x0, x1) +instance (Elt x0, Elt x1, Elt x2) => Elt (x0, x1, x2) +instance (Elt x0, Elt x1, Elt x2, Elt x3) => Elt (x0, x1, x2, x3) +instance (Elt x0, Elt x1, Elt x2, Elt x3, Elt x4) => + Elt (x0, x1, x2, x3, x4) +instance (Elt x0, Elt x1, Elt x2, Elt x3, Elt x4, Elt x5) => + Elt (x0, x1, x2, x3, x4, x5) +instance (Elt x0, + Elt x1, + Elt x2, + Elt x3, + Elt x4, + Elt x5, + Elt x6) => + Elt (x0, x1, x2, x3, x4, x5, x6) +instance (Elt x0, + Elt x1, + Elt x2, + Elt x3, + Elt x4, + Elt x5, + Elt x6, + Elt x7) => + Elt (x0, x1, x2, x3, x4, x5, x6, x7) +instance (Elt x0, + Elt x1, + Elt x2, + Elt x3, + Elt x4, + Elt x5, + Elt x6, + Elt x7, + Elt x8) => + Elt (x0, x1, x2, x3, x4, x5, x6, x7, x8) +instance (Elt x0, + Elt x1, + Elt x2, + Elt x3, + Elt x4, + Elt x5, + Elt x6, + Elt x7, + Elt x8, + Elt x9) => + Elt (x0, x1, x2, x3, x4, x5, x6, x7, x8, x9) +instance (Elt x0, + Elt x1, + Elt x2, + Elt x3, + Elt x4, + Elt x5, + Elt x6, + Elt x7, + Elt x8, + Elt x9, + Elt x10) => + Elt (x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10) +instance (Elt x0, + Elt x1, + Elt x2, + Elt x3, + Elt x4, + Elt x5, + Elt x6, + Elt x7, + Elt x8, + Elt x9, + Elt x10, + Elt x11) => + Elt (x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11) +instance (Elt x0, + Elt x1, + Elt x2, + Elt x3, + Elt x4, + Elt x5, + Elt x6, + Elt x7, + Elt x8, + Elt x9, + Elt x10, + Elt x11, + Elt x12) => + Elt (x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12) +instance (Elt x0, + Elt x1, + Elt x2, + Elt x3, + Elt x4, + Elt x5, + Elt x6, + Elt x7, + Elt x8, + Elt x9, + Elt x10, + Elt x11, + Elt x12, + Elt x13) => + Elt (x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13) +instance (Elt x0, + Elt x1, + Elt x2, + Elt x3, + Elt x4, + Elt x5, + Elt x6, + Elt x7, + Elt x8, + Elt x9, + Elt x10, + Elt x11, + Elt x12, + Elt x13, + Elt x14) => + Elt (x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13, + x14) +instance (Elt x0, + Elt x1, + Elt x2, + Elt x3, + Elt x4, + Elt x5, + Elt x6, + Elt x7, + Elt x8, + Elt x9, + Elt x10, + Elt x11, + Elt x12, + Elt x13, + Elt x14, + Elt x15) => + Elt (x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13, + x14, x15) +#endif diff --git a/src/Data/Array/Accelerate/Sugar/Foreign.hs b/src/Data/Array/Accelerate/Sugar/Foreign.hs index df28823b4..4b07d54d7 100644 --- a/src/Data/Array/Accelerate/Sugar/Foreign.hs +++ b/src/Data/Array/Accelerate/Sugar/Foreign.hs @@ -1,5 +1,3 @@ -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE TemplateHaskell #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.Sugar.Foreign diff --git a/src/Data/Array/Accelerate/Sugar/Shape.hs b/src/Data/Array/Accelerate/Sugar/Shape.hs index 1ac8bd0c4..5e399d7e4 100644 --- a/src/Data/Array/Accelerate/Sugar/Shape.hs +++ b/src/Data/Array/Accelerate/Sugar/Shape.hs @@ -1,13 +1,6 @@ -{-# LANGUAGE AllowAmbiguousTypes #-} -{-# LANGUAGE DeriveAnyClass #-} -{-# LANGUAGE DeriveGeneric #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE TypeFamilies #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.Sugar.Shape diff --git a/src/Data/Array/Accelerate/Sugar/Vec.hs b/src/Data/Array/Accelerate/Sugar/Vec.hs index bc49eeb8b..bf32a13fb 100644 --- a/src/Data/Array/Accelerate/Sugar/Vec.hs +++ b/src/Data/Array/Accelerate/Sugar/Vec.hs @@ -1,10 +1,7 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE MagicHash #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE TypeFamilies #-} +{-# OPTIONS_GHC -Wno-orphans #-} {-# OPTIONS_HADDOCK hide #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} -- | -- Module : Data.Array.Accelerate.Sugar.Vec -- Copyright : [2008..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib.hs b/src/Data/Array/Accelerate/Test/NoFib.hs index cc01feee3..7c89cdb16 100644 --- a/src/Data/Array/Accelerate/Test/NoFib.hs +++ b/src/Data/Array/Accelerate/Test/NoFib.hs @@ -1,5 +1,4 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE CPP #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib diff --git a/src/Data/Array/Accelerate/Test/NoFib/Base.hs b/src/Data/Array/Accelerate/Test/NoFib/Base.hs index 1307c6cbb..dd5f16cb2 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Base.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Base.hs @@ -1,6 +1,3 @@ -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeOperators #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Base -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Config.hs b/src/Data/Array/Accelerate/Test/NoFib/Config.hs index 6e9ecd590..a548aacc4 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Config.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Config.hs @@ -1,7 +1,5 @@ {-# LANGUAGE AllowAmbiguousTypes #-} -{-# LANGUAGE DeriveDataTypeable #-} {-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeOperators #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Config -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Imaginary.hs b/src/Data/Array/Accelerate/Test/NoFib/Imaginary.hs index eab17e9c4..7b397c06d 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Imaginary.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Imaginary.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE RankNTypes #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Imaginary -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Imaginary/DotP.hs b/src/Data/Array/Accelerate/Test/NoFib/Imaginary/DotP.hs index b530a4d8a..989a86aa7 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Imaginary/DotP.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Imaginary/DotP.hs @@ -1,9 +1,3 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Imaginary.DotP -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Imaginary/SASUM.hs b/src/Data/Array/Accelerate/Test/NoFib/Imaginary/SASUM.hs index d9d559fb3..350e638e1 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Imaginary/SASUM.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Imaginary/SASUM.hs @@ -1,9 +1,3 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Imaginary.SASUM -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Imaginary/SAXPY.hs b/src/Data/Array/Accelerate/Test/NoFib/Imaginary/SAXPY.hs index 7537d5625..efca9c4ba 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Imaginary/SAXPY.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Imaginary/SAXPY.hs @@ -1,9 +1,3 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Imaginary.SAXPY -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Issues.hs b/src/Data/Array/Accelerate/Test/NoFib/Issues.hs index 33e41dc7f..ef723a34a 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Issues.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Issues.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE RankNTypes #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Issues -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue102.hs b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue102.hs index dedec6cf6..4c4679b00 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue102.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue102.hs @@ -1,5 +1,3 @@ -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE TypeOperators #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Issues.Issue102 -- Copyright : [2009..2020] The Accelerate Team @@ -41,7 +39,10 @@ test1 = rustride = 1 v = fill (constant (Z:.(p-1))) (constant 2) - ru' = fill (constant (Z:.(p-1))) (constant 1) + -- The original repro case had (p-1) in ru' instead of p, but that + -- results in an index-out-of-bounds error if ru' is materialised (e.g. + -- if fusion is disabled). + ru' = fill (constant (Z:.p)) (constant 1) -- generate a vector with phi(p)=p-1 elements x' = reshape (constant (Z :. lts :. (p-1) :. rts)) v diff --git a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue114.hs b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue114.hs index bad2de273..ff797a35a 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue114.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue114.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE RankNTypes #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Issues.Issue114 -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue119.hs b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue119.hs index 6bcab3d55..4d8db1c2c 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue119.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue119.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE RankNTypes #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Issues.Issue119 -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue123.hs b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue123.hs index 239c04201..d4041e279 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue123.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue123.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE RankNTypes #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Issues.Issue123 -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue137.hs b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue137.hs index 444687f44..662846b5e 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue137.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue137.hs @@ -1,6 +1,3 @@ -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeOperators #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Issues.Issue137 -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue168.hs b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue168.hs index ed096251b..323bfbd42 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue168.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue168.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE RankNTypes #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Issues.Issue168 -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue184.hs b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue184.hs index 2e8a848f3..24498512d 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue184.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue184.hs @@ -1,5 +1,4 @@ {-# LANGUAGE NoImplicitPrelude #-} -{-# LANGUAGE RankNTypes #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Issues.Issue184 -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue185.hs b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue185.hs index f14f83520..6e2056653 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue185.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue185.hs @@ -1,9 +1,4 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE OverloadedLists #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE OverloadedLists #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Issues.Issue185 -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue187.hs b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue187.hs index 9c11a5d85..88fea1b1f 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue187.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue187.hs @@ -1,5 +1,4 @@ {-# LANGUAGE NoImplicitPrelude #-} -{-# LANGUAGE RankNTypes #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Issues.Issue187 -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue228.hs b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue228.hs index 4015773c9..3ec081559 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue228.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue228.hs @@ -1,6 +1,3 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE RankNTypes #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Issues.Issue228 -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue255.hs b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue255.hs index 46aaa62fa..ca7fe7cc6 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue255.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue255.hs @@ -1,8 +1,4 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE ViewPatterns #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Issues.Issue255 -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue264.hs b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue264.hs index afdbe9294..e7ceec64b 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue264.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue264.hs @@ -1,10 +1,3 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeOperators #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Issues.Issue264 -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue286.hs b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue286.hs index 5a9feb17b..eae0a9273 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue286.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue286.hs @@ -1,5 +1,3 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE RankNTypes #-} {-# LANGUAGE TemplateHaskell #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Issues.Issue286 @@ -20,7 +18,7 @@ module Data.Array.Accelerate.Test.NoFib.Issues.Issue286 ( ) where import Data.Array.Accelerate as A hiding ( (>->), (==) ) -import Data.Array.Accelerate.Array.Data +import Data.Array.Accelerate.Array.Buffer import Data.Array.Accelerate.Array.Sugar as Sugar import Data.Array.Accelerate.Array.Unique import Data.Array.Accelerate.Error diff --git a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue287.hs b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue287.hs index 3f0216cb3..23d20ede5 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue287.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue287.hs @@ -1,7 +1,3 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE RankNTypes #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Issues.Issue287 -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue288.hs b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue288.hs index 84c18a7d5..c36019a23 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue288.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue288.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE RankNTypes #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Issues.Issue288 -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue362.hs b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue362.hs index 77365ce80..17c0ce825 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue362.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue362.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE RankNTypes #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Issues.Issue362 -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue364.hs b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue364.hs index ddc167342..7bdb3cbbc 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue364.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue364.hs @@ -1,10 +1,4 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE RebindableSyntax #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE RebindableSyntax #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Issues.Issue364 -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue407.hs b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue407.hs index 1f42aa3f9..54d1e1dd6 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue407.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue407.hs @@ -1,10 +1,5 @@ {-# LANGUAGE AllowAmbiguousTypes #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE OverloadedLists #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Issues.Issue407 -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue409.hs b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue409.hs index 3486ccb4f..1793b2fec 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue409.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue409.hs @@ -1,9 +1,4 @@ {-# LANGUAGE AllowAmbiguousTypes #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Issues.Issue409 -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue427.hs b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue427.hs index 1988d8c16..192b3ce70 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue427.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue427.hs @@ -1,6 +1,3 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ViewPatterns #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Issues.Issue427 -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue436.hs b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue436.hs index 57420f447..83ca19879 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue436.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue436.hs @@ -1,5 +1,3 @@ -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE TypeOperators #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Issues.Issue436 -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue437.hs b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue437.hs index 956b6956e..bd11c0fd6 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue437.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue437.hs @@ -1,9 +1,5 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE CPP #-} -{-# LANGUAGE ForeignFunctionInterface #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE TypeOperators #-} -{-# OPTIONS_GHC -fno-warn-unused-imports #-} +{-# LANGUAGE CPP #-} +{-# OPTIONS_GHC -Wno-unused-imports #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Issues.Issue437 -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue439.hs b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue439.hs index 3bfa18b40..2843938e4 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue439.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue439.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE RankNTypes #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Issues.Issue439 -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue517.hs b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue517.hs index 0d92a9400..66765f0cc 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue517.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue517.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE RankNTypes #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Issues.Issue439 -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue551.hs b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue551.hs index 7313aff09..e9921d910 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue551.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue551.hs @@ -1,13 +1,7 @@ -{-# LANGUAGE BlockArguments #-} -{-# LANGUAGE DeriveGeneric #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE BlockArguments #-} {-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE ViewPatterns #-} {-# OPTIONS_GHC -Wno-name-shadowing #-} module Data.Array.Accelerate.Test.NoFib.Issues.Issue551 where diff --git a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue93.hs b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue93.hs index 556b27ffb..6cf415290 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue93.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Issues/Issue93.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE RankNTypes #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Issues.Issue93 -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Misc.hs b/src/Data/Array/Accelerate/Test/NoFib/Misc.hs index 7fcd0b1e0..863b914de 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Misc.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Misc.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE RankNTypes #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Misc -- Copyright : [2025..] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Misc/Cache.hs b/src/Data/Array/Accelerate/Test/NoFib/Misc/Cache.hs index c95a98116..b3f4b922a 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Misc/Cache.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Misc/Cache.hs @@ -1,5 +1,4 @@ {-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE RankNTypes #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Misc.Cache -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Misc/Scanl1.hs b/src/Data/Array/Accelerate/Test/NoFib/Misc/Scanl1.hs index c4e739f3d..28639489d 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Misc/Scanl1.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Misc/Scanl1.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE RankNTypes #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Misc.Scanl1 -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Prelude.hs b/src/Data/Array/Accelerate/Test/NoFib/Prelude.hs index 3a3508228..97a60926f 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Prelude.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Prelude.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE RankNTypes #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Prelude -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Prelude/Backpermute.hs b/src/Data/Array/Accelerate/Test/NoFib/Prelude/Backpermute.hs index 370d52b4d..c7d3e1e97 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Prelude/Backpermute.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Prelude/Backpermute.hs @@ -1,10 +1,3 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeOperators #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Prelude.Backpermute -- Copyright : [2009..2020] The Accelerate Team @@ -152,7 +145,7 @@ takeRef -> Array (sh:.Int) e takeRef n arr = let sh :. m = S.shape arr - in fromFunction (sh :. P.min m n) (arr S.!) + in fromFunction (sh :. P.min m (P.max 0 n)) (arr S.!) dropRef :: (Shape sh, Slice sh, Elt e) diff --git a/src/Data/Array/Accelerate/Test/NoFib/Prelude/Filter.hs b/src/Data/Array/Accelerate/Test/NoFib/Prelude/Filter.hs index a3233cdeb..fbabdb510 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Prelude/Filter.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Prelude/Filter.hs @@ -1,12 +1,3 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE CPP #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE MonoLocalBinds #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeOperators #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Prelude.Filter -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Prelude/Fold.hs b/src/Data/Array/Accelerate/Test/NoFib/Prelude/Fold.hs index 5222492a7..9babb7df4 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Prelude/Fold.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Prelude/Fold.hs @@ -1,10 +1,3 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeOperators #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Prelude.Fold -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Prelude/Map.hs b/src/Data/Array/Accelerate/Test/NoFib/Prelude/Map.hs index 1d02bc194..04745227a 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Prelude/Map.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Prelude/Map.hs @@ -1,11 +1,3 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE MonoLocalBinds #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeOperators #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Prelude.Map -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Prelude/Permute.hs b/src/Data/Array/Accelerate/Test/NoFib/Prelude/Permute.hs index b9643e6f3..e2f9e4d7b 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Prelude/Permute.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Prelude/Permute.hs @@ -1,10 +1,3 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeOperators #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Prelude.Permute -- Copyright : [2009..2020] The Accelerate Team @@ -22,7 +15,7 @@ module Data.Array.Accelerate.Test.NoFib.Prelude.Permute ( ) where import Data.Array.Accelerate as A -import Data.Array.Accelerate.Array.Data +import Data.Array.Accelerate.Array.Buffer import Data.Array.Accelerate.Sugar.Array as S import Data.Array.Accelerate.Sugar.Elt as S import Data.Array.Accelerate.Sugar.Shape as S @@ -153,28 +146,29 @@ permuteRef -> (sh -> Maybe sh') -> Array sh e -> Array sh' e -permuteRef f def@(Array (R.Array _ aold)) p arr@(Array (R.Array _ anew)) = +permuteRef f def p arr = unsafePerformIO $ do - let - tp = S.eltR @e - sh = S.shape arr - sh' = S.shape def - n = S.size sh - -- - go !i - | i P.>= n = pure () - | otherwise = do - let ix = S.fromIndex sh i - case p ix of - Nothing -> pure () - Just ix' -> do - let i' = S.toIndex sh' ix' - x <- toElt P.<$> readArrayData tp anew i - x' <- toElt P.<$> readArrayData tp aold i' - writeArrayData tp aold i' (fromElt (f x x')) - -- - go (i+1) - -- - go 0 - pure def + let tp = S.eltR @e + shSrc = S.shape arr + shOut = S.shape def + + output <- newBuffers tp (S.size shOut) + + loop (S.size shOut) $ \i -> + writeBuffers tp output i (fromElt (def S.!! i)) + loop (S.size shSrc) $ \i -> do + case p (S.fromIndex shSrc i) of + Nothing -> pure () + Just ixOut -> do + let iOut = S.toIndex shOut ixOut + old <- toElt P.<$> readBuffers tp output iOut + writeBuffers tp output iOut (fromElt (f (arr S.!! i) old)) + + pure (Array (R.Array (fromElt shOut) (unsafeFreezeBuffers tp output))) + where + loop :: Int -> (Int -> IO ()) -> IO () + loop n body = + let go !i | i P.>= n = pure () + | otherwise = body i >> go (i + 1) + in go 0 diff --git a/src/Data/Array/Accelerate/Test/NoFib/Prelude/SIMD.hs b/src/Data/Array/Accelerate/Test/NoFib/Prelude/SIMD.hs index bbec663e5..733264ee0 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Prelude/SIMD.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Prelude/SIMD.hs @@ -1,9 +1,3 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Prelude.SIMD -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Prelude/Scan.hs b/src/Data/Array/Accelerate/Test/NoFib/Prelude/Scan.hs index 61a18345b..3a7b3a18e 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Prelude/Scan.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Prelude/Scan.hs @@ -1,10 +1,3 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeOperators #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Prelude.Scan -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Prelude/Stencil.hs b/src/Data/Array/Accelerate/Test/NoFib/Prelude/Stencil.hs index 690e17ba1..cf31ad1d7 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Prelude/Stencil.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Prelude/Stencil.hs @@ -1,13 +1,3 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE CPP #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE PatternGuards #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeOperators #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Prelude.Stencil -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Prelude/ZipWith.hs b/src/Data/Array/Accelerate/Test/NoFib/Prelude/ZipWith.hs index 23749281d..a50ec6ef1 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Prelude/ZipWith.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Prelude/ZipWith.hs @@ -1,11 +1,3 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE MonoLocalBinds #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeOperators #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Prelude.ZipWith -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Sharing.hs b/src/Data/Array/Accelerate/Test/NoFib/Sharing.hs index 7fada0655..ee4378b2d 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Sharing.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Sharing.hs @@ -1,8 +1,3 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE RecordWildCards #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeOperators #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Sharing -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Spectral.hs b/src/Data/Array/Accelerate/Test/NoFib/Spectral.hs index ee9246279..24bd11c1b 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Spectral.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Spectral.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE RankNTypes #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Spectral -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Spectral/BlackScholes.hs b/src/Data/Array/Accelerate/Test/NoFib/Spectral/BlackScholes.hs index e15b73231..d3a2ba095 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Spectral/BlackScholes.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Spectral/BlackScholes.hs @@ -1,11 +1,4 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE ViewPatterns #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Spectral.BlackScholes -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Spectral/RadixSort.hs b/src/Data/Array/Accelerate/Test/NoFib/Spectral/RadixSort.hs index a65c4b3ef..f2348bcc2 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Spectral/RadixSort.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Spectral/RadixSort.hs @@ -1,11 +1,3 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE MonoLocalBinds #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeOperators #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Spectral.RadixSort -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/NoFib/Spectral/SMVM.hs b/src/Data/Array/Accelerate/Test/NoFib/Spectral/SMVM.hs index 9686b503b..77a22fc09 100644 --- a/src/Data/Array/Accelerate/Test/NoFib/Spectral/SMVM.hs +++ b/src/Data/Array/Accelerate/Test/NoFib/Spectral/SMVM.hs @@ -1,10 +1,3 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeOperators #-} -- | -- Module : Data.Array.Accelerate.Test.NoFib.Spectral.SMVM -- Copyright : [2009..2020] The Accelerate Team diff --git a/src/Data/Array/Accelerate/Test/Similar.hs b/src/Data/Array/Accelerate/Test/Similar.hs index 1017408cc..b195e94fb 100644 --- a/src/Data/Array/Accelerate/Test/Similar.hs +++ b/src/Data/Array/Accelerate/Test/Similar.hs @@ -1,8 +1,4 @@ -{-# LANGUAGE BangPatterns #-} {-# LANGUAGE DefaultSignatures #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE TypeOperators #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.Test.Similar diff --git a/src/Data/Array/Accelerate/Trafo.hs b/src/Data/Array/Accelerate/Trafo.hs index ea20ce2f2..309ecece5 100644 --- a/src/Data/Array/Accelerate/Trafo.hs +++ b/src/Data/Array/Accelerate/Trafo.hs @@ -1,5 +1,5 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE CPP #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.Trafo @@ -14,12 +14,10 @@ module Data.Array.Accelerate.Trafo ( -- * HOAS -> de Bruijn conversion - -- ** Array computations - convertAcc, convertAccWith, - -- ** Array functions Afunction, ArraysFunctionR, convertAfun, convertAfunWith, + Trafo, -- ** Sequence computations -- convertSeq, convertSeqWith, @@ -28,61 +26,135 @@ module Data.Array.Accelerate.Trafo ( Function, EltFunctionR, convertExp, convertFun, + inspectCompiler', convertAfunWithObj, ) where -import Data.Array.Accelerate.Sugar.Array ( ArraysR ) -import Data.Array.Accelerate.Sugar.Elt ( EltR ) +import Data.Array.Accelerate.Sugar.Elt ( EltR ) import Data.Array.Accelerate.Smart import Data.Array.Accelerate.Trafo.Config -import Data.Array.Accelerate.Trafo.Delayed -import Data.Array.Accelerate.Trafo.Sharing ( Afunction, ArraysFunctionR, Function, EltFunctionR ) -import qualified Data.Array.Accelerate.AST as AST -import qualified Data.Array.Accelerate.Trafo.Fusion as Fusion -import qualified Data.Array.Accelerate.Trafo.LetSplit as LetSplit -import qualified Data.Array.Accelerate.Trafo.Simplify as Rewrite -import qualified Data.Array.Accelerate.Trafo.Sharing as Sharing --- import qualified Data.Array.Accelerate.Trafo.Vectorise as Vectorise - -import Control.DeepSeq -import Data.Text.Lazy.Builder +-- import Data.Array.Accelerate.Trafo.Delayed +import Data.Array.Accelerate.Trafo.Sharing ( Afunction, ArraysFunctionR, Function, EltFunctionR ) +import qualified Data.Array.Accelerate.AST as AST +import Data.Array.Accelerate.AST.Kernel +import Data.Array.Accelerate.AST.Schedule +import qualified Data.Array.Accelerate.AST.Partitioned as Partitioned +-- import qualified Data.Array.Accelerate.Trafo.Fusion as Fusion +import qualified Data.Array.Accelerate.Trafo.LetSplit as LetSplit +import qualified Data.Array.Accelerate.Trafo.Exp.Simplify as Rewrite +import qualified Data.Array.Accelerate.Trafo.Sharing as Sharing +import qualified Data.Array.Accelerate.Trafo.Operation.LiveVars as Operation +import qualified Data.Array.Accelerate.Trafo.Operation.Simplify as Operation +-- import qualified Data.Array.Accelerate.Trafo.Vectorise as Vectorise #ifdef ACCELERATE_DEBUG +import Data.Array.Accelerate.Debug.Internal import Formatting -import System.IO.Unsafe -import Data.Array.Accelerate.Debug.Internal.Flags hiding ( when ) -import Data.Array.Accelerate.Debug.Internal.Timed +import System.IO.Unsafe (unsafePerformIO) #endif +import Control.DeepSeq +import qualified Data.Array.Accelerate.Trafo.Partitioning.ILP.Graph as Partitioning +import Data.Array.Accelerate.Representation.Ground (LoweredAfun) +import Data.Array.Accelerate.Trafo.Lowering (LowerAcc, lowerAfun) +import qualified Data.Array.Accelerate.Trafo.NewNewFusion as NewNewFusion +import Prettyprinter as Pretty +import qualified Data.Array.Accelerate.Pretty as Pretty +import qualified Data.Array.Accelerate.Pretty.Operation as Pretty +import qualified Data.Array.Accelerate.Pretty.Schedule as Pretty +import Data.Array.Accelerate.Pretty (prettyOpenAcc) +import Data.Array.Accelerate.Pretty.Partitioned () +import Data.String (fromString) +import Data.Char (toUpper) +import qualified Data.Array.Accelerate.AST.Operation as Operation +import qualified Data.Array.Accelerate.Trafo.Partitioning.ILP.Graph as Graph +import Data.Array.Accelerate.Pretty.Print (configPlain, Val (Empty)) + +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Encoding (Objective(..)) +import Data.Array.Accelerate.Trafo.Partitioning.ILP (FusionType(..), defaultObjective) +import Control.Monad.Trans.Writer (runWriter, Writer, writer) +import Control.Monad ((>=>)) +import Data.Array.Accelerate.Trafo.Operation.Bounds + +inspectCompiler' + :: forall sched kernel f. + (Afunction f, Trafo sched kernel) + => f + -> String +inspectCompiler' = snd . convertAfunFullOptions @sched @kernel defaultOptions defaultObjective Pretty.renderForTerminal + +{- testWithObjective + :: forall sched kernel f. (Afunction f, LowerAcc (KernelOperation kernel), Operation.SimplifyOperation (KernelOperation kernel), Operation.SLVOperation (KernelOperation kernel), Partitioning.MakesILP (KernelOperation kernel), Partitioned.SetOpIndices (KernelOperation kernel), Pretty.PrettyOp (KernelOperation kernel), Pretty.PrettyKernel kernel, IsSchedule sched, IsKernel kernel, Pretty.PrettySchedule sched, Operation.ShrinkArg (Partitioning.BackendClusterArg (KernelOperation kernel)), Operation.NFData' (Graph.BackendClusterArg (KernelOperation kernel)), Operation.ShrinkArg (Partitioning.BackendClusterArg (KernelOperation kernel))) + => Objective + -> f + -> String +testWithObjective obj = snd . convertAfunFullOptions @sched @kernel defaultOptions (Fusion obj) Pretty.renderForTerminal + + +testBench + :: forall sched kernel f. (Afunction f, LowerAcc (KernelOperation kernel), Operation.SimplifyOperation (KernelOperation kernel), Operation.SLVOperation (KernelOperation kernel), Partitioning.MakesILP (KernelOperation kernel), Partitioned.SetOpIndices (KernelOperation kernel), Pretty.PrettyOp (KernelOperation kernel), Pretty.PrettyKernel kernel, IsSchedule sched, IsKernel kernel, Pretty.PrettySchedule sched, Operation.ShrinkArg (Partitioning.BackendClusterArg (KernelOperation kernel)), Operation.NFData' (Graph.BackendClusterArg (KernelOperation kernel)), Operation.ShrinkArg (Partitioning.BackendClusterArg (KernelOperation kernel))) + => Benchmarking + -> f + -> String +testBench bench = snd . convertAfunFullOptions @sched @kernel defaultOptions (Benchmarking bench) Pretty.renderForTerminal +-} --- HOAS -> de Bruijn conversion --- ---------------------------- - --- | Convert a closed array expression to de Bruijn form while also --- incorporating sharing observation and array fusion. --- -convertAcc :: Acc arrs -> DelayedAcc (ArraysR arrs) -convertAcc = convertAccWith defaultOptions -convertAccWith :: Config -> Acc arrs -> DelayedAcc (ArraysR arrs) -convertAccWith config - = phase "array-fusion" (Fusion.convertAccWith config) - . phase "array-split-lets" LetSplit.convertAcc - -- phase "vectorise-sequences" Vectorise.vectoriseSeqAcc `when` vectoriseSequences - . phase "sharing-recovery" (Sharing.convertAccWith config) +-- HOAS -> de Bruijn conversion +-- ---------------------------- --- | Convert a unary function over array computations, incorporating sharing +-- | Convert an Acc function over array computations, incorporating sharing -- observation and array fusion -- -convertAfun :: Afunction f => f -> DelayedAfun (ArraysFunctionR f) +-- Note that we previously had a transformations for Acc and Afunction +-- separately. Instead, we now always run the transformation via Afunction, as +-- an Acc is also a Afunction (without arguments). +-- +convertAfun + :: forall sched kernel f. + (Afunction f, Trafo sched kernel) + => f + -> sched kernel () (Scheduled sched (LoweredAfun (ArraysFunctionR f))) convertAfun = convertAfunWith defaultOptions -convertAfunWith :: Afunction f => Config -> f -> DelayedAfun (ArraysFunctionR f) -convertAfunWith config - = phase "array-fusion" (Fusion.convertAfunWith config) - . phase "array-split-lets" LetSplit.convertAfun - -- phase "vectorise-sequences" Vectorise.vectoriseSeqAfun `when` vectoriseSequences - . phase "sharing-recovery" (Sharing.convertAfunWith config) +convertAfunWith + :: forall sched kernel f. + (Afunction f, Trafo sched kernel) + => Config + -> f + -> sched kernel () (Scheduled sched (LoweredAfun (ArraysFunctionR f))) +convertAfunWith config = fst . convertAfunFullOptions config defaultObjective (const ()) + +convertAfunWithObj + :: forall sched kernel f. + (Afunction f, Trafo sched kernel) + => Objective + -> f + -> sched kernel () (Scheduled sched (LoweredAfun (ArraysFunctionR f))) +convertAfunWithObj obj = fst . convertAfunFullOptions defaultOptions (Fusion obj) (const ()) + +type Trafo sched kernel = (LowerAcc (KernelOperation kernel), Operation.SimplifyOperation (KernelOperation kernel), Operation.SLVOperation (KernelOperation kernel), OperationBounds (KernelOperation kernel), Partitioning.MakesILP (KernelOperation kernel), Partitioned.SetOpIndices (KernelOperation kernel), Pretty.PrettyOp (KernelOperation kernel), Pretty.PrettyKernel kernel, IsSchedule sched, IsKernel kernel, Pretty.PrettySchedule sched, Operation.ShrinkArg (Partitioning.BackendClusterArg (KernelOperation kernel)), Operation.NFData' (Graph.BackendClusterArg (KernelOperation kernel)), Operation.ShrinkArg (Partitioning.BackendClusterArg (KernelOperation kernel))) + +convertAfunFullOptions + :: forall sched kernel f m. + (Monoid m, Afunction f, Trafo sched kernel) + => Config -> FusionType -> (Pretty.Adoc -> m) + -> f -> (sched kernel () (Scheduled sched (LoweredAfun (ArraysFunctionR f))), m) +convertAfunFullOptions config ft pprint f = runWriter $ + ( phase'' "sharing-recovery" (Sharing.convertAfunWith config) (Pretty.prettyPreOpenAfun configPlain prettyOpenAcc Empty) + >=> phase'' "array-split-lets" LetSplit.convertAfun (Pretty.prettyPreOpenAfun configPlain prettyOpenAcc Empty) + >=> phase'' "lowering" (Operation.simplifyFun . Operation.simplifyFun . lowerAfun) Pretty.prettyAfun + >=> phase'' "operation-bounds" (Operation.simplifyFun . boundsOptimizeAfun) Pretty.prettyAfun + >=> phase'' "operation-live-vars" (Operation.simplifyFun . Operation.stronglyLiveVariablesFun) Pretty.prettyAfun + >=> phase'' "array-fusion" (Operation.simplifyFun . NewNewFusion.convertAfunWith config ft) Pretty.prettyAfun + >=> phase'' "partition-live-vars" (Operation.simplifyFun . Operation.stronglyLiveVariablesFun) Pretty.prettyAfun + -- TODO: Should we do boundsOptimizeAfun again? + >=> (\x -> let y = phase' "codegen" rnfSchedule convertScheduleFun x in writer (y, pprint $ vsep ["CODEGEN", Pretty.prettySchedule y])) + ) f + where + phase'' :: NFData b => String -> (a -> b) -> (b -> Pretty.Adoc) -> a -> Writer m b + phase'' name g pp a = let b = phase name g a in writer (b, pprint $ Pretty.vsep [fromString $ map toUpper name, pp b, mempty, mempty]) + -- | Convert a closed scalar expression, incorporating sharing observation and @@ -102,26 +174,6 @@ convertFun = phase "exp-simplify" Rewrite.simplifyFun . phase "sharing-recovery" Sharing.convertFun -{-- --- | Convert a closed sequence computation, incorporating sharing observation and --- optimisation. --- -convertSeq :: Typeable s => Seq s -> DelayedSeq s -convertSeq = convertSeqWith phases - -convertSeqWith :: Typeable s => Phase -> Seq s -> DelayedSeq s -convertSeqWith Phase{..} s - = phase "array-fusion" (Fusion.convertSeq enableAccFusion) - -- $ phase "vectorise-sequences" Vectorise.vectoriseSeq `when` vectoriseSequences - -- $ phase "rewrite-segment-offset" Rewrite.convertSegmentsSeq `when` convertOffsetOfSegment - $ phase "sharing-recovery" (Sharing.convertSeq recoverAccSharing recoverExpSharing recoverSeqSharing floatOutAccFromExp) - $ s ---} - - --- when :: (a -> a) -> Bool -> a -> a --- when f True = f --- when _ False = id -- Debugging -- --------- @@ -129,14 +181,19 @@ convertSeqWith Phase{..} s -- Execute a phase of the compiler and (possibly) print some timing/gc -- statistics. -- -phase :: NFData b => Builder -> (a -> b) -> a -> b +phase :: NFData b => String -> (a -> b) -> a -> b +phase n = phase' n rnf + +-- Execute a phase of the compiler and (possibly) print some timing/gc +-- statistics. +-- +phase' :: String -> (b -> ()) -> (a -> b) -> a -> b #ifdef ACCELERATE_DEBUG -phase n f x = unsafePerformIO $ do +phase' n rnf'' f x = unsafePerformIO $ do enabled <- getFlag dump_phases if enabled - then timed dump_phases (now ("phase " <> n <> ": ") % elapsed) (return $!! f x) + then timed dump_phases (now ("phase " <> fromString n <> ": ") % elapsed) (let y = f x in rnf'' y `seq` return y) else return (f x) #else -phase _ f = f +phase' _ _ f = f #endif - diff --git a/src/Data/Array/Accelerate/Trafo/.gitignore b/src/Data/Array/Accelerate/Trafo/.gitignore new file mode 100644 index 000000000..e74d9daab --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/.gitignore @@ -0,0 +1 @@ +/DeadBufferRemoval.hs diff --git a/src/Data/Array/Accelerate/Trafo/Config.hs b/src/Data/Array/Accelerate/Trafo/Config.hs index 40667e090..6060e1156 100644 --- a/src/Data/Array/Accelerate/Trafo/Config.hs +++ b/src/Data/Array/Accelerate/Trafo/Config.hs @@ -34,4 +34,3 @@ data Config = Config defaultOptions :: Config defaultOptions = unsafePerformIO $! Config <$> BitSet <$> peek F.__cmd_line_flags - diff --git a/src/Data/Array/Accelerate/Trafo/Delayed.hs b/src/Data/Array/Accelerate/Trafo/Delayed.hs index 045d205e0..ddc18b954 100644 --- a/src/Data/Array/Accelerate/Trafo/Delayed.hs +++ b/src/Data/Array/Accelerate/Trafo/Delayed.hs @@ -1,126 +1,121 @@ -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE RecordWildCards #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeSynonymInstances #-} -{-# OPTIONS_HADDOCK hide #-} --- | --- Module : Data.Array.Accelerate.Trafo.Delayed --- Copyright : [2012..2020] The Accelerate Team --- License : BSD3 --- --- Maintainer : Trevor L. McDonell --- Stability : experimental --- Portability : non-portable (GHC extensions) --- --- The type of delayed arrays. This representation is used to annotate the AST --- in the recursive knot to distinguish standard AST terms from operand arrays --- that should be embedded into their consumers. --- - -module Data.Array.Accelerate.Trafo.Delayed - where - -import Data.Array.Accelerate.AST -import Data.Array.Accelerate.Analysis.Hash -import Data.Array.Accelerate.Analysis.Match -import Data.Array.Accelerate.Representation.Array -import Data.Array.Accelerate.Representation.Type -import Data.Array.Accelerate.Trafo.Substitution - -import Data.Array.Accelerate.Debug.Internal.Stats as Stats - -import Control.DeepSeq -import Data.ByteString.Builder -import Data.ByteString.Builder.Extra - - -type DelayedAcc = DelayedOpenAcc () -type DelayedAfun = PreOpenAfun DelayedOpenAcc () -type DelayedOpenAfun = PreOpenAfun DelayedOpenAcc - --- type DelayedOpenSeq = PreOpenSeq DelayedOpenAcc --- data DelayedSeq t where --- DelayedSeq :: Extend DelayedOpenAcc () aenv --- -> DelayedOpenSeq aenv () t --- -> DelayedSeq t - -data DelayedOpenAcc aenv a where - Manifest :: PreOpenAcc DelayedOpenAcc aenv a - -> DelayedOpenAcc aenv a - - Delayed :: - { reprD :: ArrayR (Array sh e) - , extentD :: Exp aenv sh - , indexD :: Fun aenv (sh -> e) - , linearIndexD :: Fun aenv (Int -> e) - } -> DelayedOpenAcc aenv (Array sh e) - -instance HasArraysR DelayedOpenAcc where - arraysR (Manifest a) = arraysR a - arraysR Delayed{..} = TupRsingle reprD - -instance Rebuildable DelayedOpenAcc where - type AccClo DelayedOpenAcc = DelayedOpenAcc - rebuildPartial v = \case - Manifest pacc -> Manifest <$> rebuildPartial v pacc - Delayed{..} -> (\e i l -> Delayed reprD (unOpenAccExp e) (unOpenAccFun i) (unOpenAccFun l)) - <$> rebuildPartial v (OpenAccExp extentD) - <*> rebuildPartial v (OpenAccFun indexD) - <*> rebuildPartial v (OpenAccFun linearIndexD) - -instance Sink DelayedOpenAcc where - weaken k = Stats.substitution "weaken" . rebuildA (rebuildWeakenVar k) - -instance NFData (DelayedOpenAfun aenv t) where - rnf = rnfPreOpenAfun rnfDelayedOpenAcc - -instance NFData (DelayedOpenAcc aenv t) where - rnf = rnfDelayedOpenAcc - -encodeDelayedOpenAcc :: EncodeAcc DelayedOpenAcc -encodeDelayedOpenAcc options acc = - let - travE :: Exp aenv sh -> Builder - travE = encodeOpenExp - - travF :: Fun aenv f -> Builder - travF = encodeOpenFun - - travA :: PreOpenAcc DelayedOpenAcc aenv a -> Builder - travA = encodePreOpenAcc options encodeDelayedOpenAcc - - deepA :: PreOpenAcc DelayedOpenAcc aenv a -> Builder - deepA | perfect options = travA - | otherwise = encodeArraysType . arraysR - in - case acc of - Manifest pacc -> intHost $(hashQ ("Manifest" :: String)) <> deepA pacc - Delayed _ sh f g -> intHost $(hashQ ("Delayed" :: String)) <> travE sh <> travF f <> travF g - -matchDelayedOpenAcc :: MatchAcc DelayedOpenAcc -matchDelayedOpenAcc (Manifest pacc1) (Manifest pacc2) - = matchPreOpenAcc matchDelayedOpenAcc pacc1 pacc2 -matchDelayedOpenAcc (Delayed _ sh1 ix1 lx1) (Delayed _ sh2 ix2 lx2) - | Just Refl <- matchOpenExp sh1 sh2 - , Just Refl <- matchOpenFun ix1 ix2 - , Just Refl <- matchOpenFun lx1 lx2 - = Just Refl -matchDelayedOpenAcc _ _ - = Nothing - -rnfDelayedOpenAcc :: NFDataAcc DelayedOpenAcc -rnfDelayedOpenAcc (Manifest pacc) = - rnfPreOpenAcc rnfDelayedOpenAcc pacc -rnfDelayedOpenAcc (Delayed aR sh ix lx) = - rnfArrayR aR `seq` rnfOpenExp sh `seq` rnfOpenFun ix `seq` rnfOpenFun lx - -liftDelayedOpenAcc :: LiftAcc DelayedOpenAcc -liftDelayedOpenAcc (Manifest pacc) = - [|| Manifest $$(liftPreOpenAcc liftDelayedOpenAcc pacc) ||] -liftDelayedOpenAcc (Delayed aR sh ix lx) = - [|| Delayed $$(liftArrayR aR) $$(liftOpenExp sh) $$(liftOpenFun ix) $$(liftOpenFun lx) ||] +-- {-# LANGUAGE RecordWildCards #-} +-- {-# LANGUAGE TemplateHaskell #-} +-- {-# LANGUAGE TypeFamilies #-} +-- {-# OPTIONS_HADDOCK hide #-} +-- -- | +-- -- Module : Data.Array.Accelerate.Trafo.Delayed +-- -- Copyright : [2012..2020] The Accelerate Team +-- -- License : BSD3 +-- -- +-- -- Maintainer : Trevor L. McDonell +-- -- Stability : experimental +-- -- Portability : non-portable (GHC extensions) +-- -- +-- -- The type of delayed arrays. This representation is used to annotate the AST +-- -- in the recursive knot to distinguish standard AST terms from operand arrays +-- -- that should be embedded into their consumers. +-- -- + +-- module Data.Array.Accelerate.Trafo.Delayed +-- where + +-- import Data.Array.Accelerate.AST +-- import Data.Array.Accelerate.Analysis.Hash +-- import Data.Array.Accelerate.Analysis.Match +-- import Data.Array.Accelerate.Representation.Array +-- import Data.Array.Accelerate.Representation.Type +-- import Data.Array.Accelerate.Trafo.Substitution + +-- import Data.Array.Accelerate.Debug.Internal.Stats as Stats + +-- import Control.DeepSeq +-- import Data.ByteString.Builder +-- import Data.ByteString.Builder.Extra + + +-- type DelayedAcc = DelayedOpenAcc () +-- type DelayedAfun = PreOpenAfun DelayedOpenAcc () +-- type DelayedOpenAfun = PreOpenAfun DelayedOpenAcc + +-- -- type DelayedOpenSeq = PreOpenSeq DelayedOpenAcc +-- -- data DelayedSeq t where +-- -- DelayedSeq :: Extend DelayedOpenAcc () aenv +-- -- -> DelayedOpenSeq aenv () t +-- -- -> DelayedSeq t + +-- data DelayedOpenAcc aenv a where +-- Manifest :: PreOpenAcc DelayedOpenAcc aenv a +-- -> DelayedOpenAcc aenv a + +-- Delayed :: +-- { reprD :: ArrayR (Array sh e) +-- , extentD :: Exp aenv sh +-- , indexD :: Fun aenv (sh -> e) +-- , linearIndexD :: Fun aenv (Int -> e) +-- } -> DelayedOpenAcc aenv (Array sh e) + +-- instance HasArraysR DelayedOpenAcc where +-- arraysR (Manifest a) = arraysR a +-- arraysR Delayed{..} = TupRsingle reprD + +-- instance Rebuildable DelayedOpenAcc where +-- type AccClo DelayedOpenAcc = DelayedOpenAcc +-- rebuildPartial v = \case +-- Manifest pacc -> Manifest <$> rebuildPartial v pacc +-- Delayed{..} -> (\e i l -> Delayed reprD (unOpenAccExp e) (unOpenAccFun i) (unOpenAccFun l)) +-- <$> rebuildPartial v (OpenAccExp extentD) +-- <*> rebuildPartial v (OpenAccFun indexD) +-- <*> rebuildPartial v (OpenAccFun linearIndexD) + +-- instance Sink DelayedOpenAcc where +-- weaken k = Stats.substitution "weaken" . rebuildA (rebuildWeakenVar k) + +-- instance NFData (DelayedOpenAfun aenv t) where +-- rnf = rnfPreOpenAfun rnfDelayedOpenAcc + +-- instance NFData (DelayedOpenAcc aenv t) where +-- rnf = rnfDelayedOpenAcc + +-- encodeDelayedOpenAcc :: EncodeAcc DelayedOpenAcc +-- encodeDelayedOpenAcc options acc = +-- let +-- travE :: Exp aenv sh -> Builder +-- travE = encodeOpenExp + +-- travF :: Fun aenv f -> Builder +-- travF = encodeOpenFun + +-- travA :: PreOpenAcc DelayedOpenAcc aenv a -> Builder +-- travA = encodePreOpenAcc options encodeDelayedOpenAcc + +-- deepA :: PreOpenAcc DelayedOpenAcc aenv a -> Builder +-- deepA | perfect options = travA +-- | otherwise = encodeArraysType . arraysR +-- in +-- case acc of +-- Manifest pacc -> intHost $(hashQ ("Manifest" :: String)) <> deepA pacc +-- Delayed _ sh f g -> intHost $(hashQ ("Delayed" :: String)) <> travE sh <> travF f <> travF g + +-- matchDelayedOpenAcc :: MatchAcc DelayedOpenAcc +-- matchDelayedOpenAcc (Manifest pacc1) (Manifest pacc2) +-- = matchPreOpenAcc matchDelayedOpenAcc pacc1 pacc2 +-- matchDelayedOpenAcc (Delayed _ sh1 ix1 lx1) (Delayed _ sh2 ix2 lx2) +-- | Just Refl <- matchOpenExp sh1 sh2 +-- , Just Refl <- matchOpenFun ix1 ix2 +-- , Just Refl <- matchOpenFun lx1 lx2 +-- = Just Refl +-- matchDelayedOpenAcc _ _ +-- = Nothing + +-- rnfDelayedOpenAcc :: NFDataAcc DelayedOpenAcc +-- rnfDelayedOpenAcc (Manifest pacc) = +-- rnfPreOpenAcc rnfDelayedOpenAcc pacc +-- rnfDelayedOpenAcc (Delayed aR sh ix lx) = +-- rnfArrayR aR `seq` rnfOpenExp sh `seq` rnfOpenFun ix `seq` rnfOpenFun lx + +-- liftDelayedOpenAcc :: LiftAcc DelayedOpenAcc +-- liftDelayedOpenAcc (Manifest pacc) = +-- [|| Manifest $$(liftPreOpenAcc liftDelayedOpenAcc pacc) ||] +-- liftDelayedOpenAcc (Delayed aR sh ix lx) = +-- [|| Delayed $$(liftArrayR aR) $$(liftOpenExp sh) $$(liftOpenFun ix) $$(liftOpenFun lx) ||] diff --git a/src/Data/Array/Accelerate/Trafo/Environment.hs b/src/Data/Array/Accelerate/Trafo/Environment.hs index 05e4c9bca..00e6c0dde 100644 --- a/src/Data/Array/Accelerate/Trafo/Environment.hs +++ b/src/Data/Array/Accelerate/Trafo/Environment.hs @@ -1,8 +1,3 @@ -{-# LANGUAGE GADTs #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeOperators #-} -- | -- Module : Data.Array.Accelerate.Trafo.Environment -- Copyright : [2012..2020] The Accelerate Team @@ -27,23 +22,32 @@ import Data.Array.Accelerate.Type import Data.Array.Accelerate.Debug.Internal.Stats as Stats - -- An environment that holds let-bound scalar expressions. The second -- environment variable env' is used to project out the corresponding -- index when looking up in the environment congruent expressions. -- -data Gamma env env' aenv where - EmptyExp :: Gamma env env' aenv +data Gamma arr env env' where + EmptyExp :: Gamma arr env () - PushExp :: Gamma env env' aenv - -> WeakOpenExp env aenv t - -> Gamma env (env', t) aenv + PushExp :: Gamma arr env env' + -> WeakOpenExp arr env t + -> Gamma arr env (env', t) -data WeakOpenExp env aenv t where +data WeakOpenExp arr env t where Subst :: env :> env' - -> OpenExp env aenv t - -> OpenExp env' aenv t {- LAZY -} - -> WeakOpenExp env' aenv t + -> PreOpenExp arr env t + -> PreOpenExp arr env' t {- LAZY -} + -> WeakOpenExp arr env' t + +instance SinkExp (WeakOpenExp arr) where + weakenE k (Subst k' e1 _) = Subst + (k .> k') + e1 + (weakenE (k .> k') e1) + +instance SinkExp (Gamma arr) where + weakenE _ EmptyExp = EmptyExp + weakenE k (PushExp env e) = weakenE k env `PushExp` weakenE k e -- XXX: The simplifier calls this function every time it moves under a let -- binding. This means we have a number of calls to 'weakenE' exponential in the @@ -57,45 +61,30 @@ data WeakOpenExp env aenv t where -- -- incExp - :: Gamma env env' aenv - -> Gamma (env,s) env' aenv + :: Gamma arr env env' + -> Gamma arr (env,s) env' incExp EmptyExp = EmptyExp incExp (PushExp env w) = incExp env `PushExp` subs w where - subs :: forall env aenv s t. WeakOpenExp env aenv t -> WeakOpenExp (env,s) aenv t - subs (Subst k (e :: OpenExp env_ aenv t) _) = Subst (weakenSucc' k) e (weakenE (weakenSucc' k) e) + subs :: forall arr env s t. WeakOpenExp arr env t -> WeakOpenExp arr (env,s) t + subs (Subst k (e :: PreOpenExp arr env_ t) _) = Subst (weakenSucc' k) e (weakenE (weakenSucc' k) e) -prjExp :: HasCallStack => Idx env' t -> Gamma env env' aenv -> OpenExp env aenv t +prjExp :: HasCallStack => Idx env' t -> Gamma arr env env' -> PreOpenExp arr env t prjExp ZeroIdx (PushExp _ (Subst _ _ e)) = e prjExp (SuccIdx ix) (PushExp env _) = prjExp ix env -prjExp _ _ = internalError "inconsistent valuation" -pushExp :: Gamma env env' aenv -> OpenExp env aenv t -> Gamma env (env',t) aenv +pushExp :: Gamma arr env env' -> PreOpenExp arr env t -> Gamma arr env (env',t) pushExp env e = env `PushExp` Subst weakenId e e {-- lookupExp - :: Gamma env env' aenv - -> OpenExp env aenv t + :: Gamma arr env env' + -> PreOpenExp arr env t -> Maybe (Idx env' t) lookupExp EmptyExp _ = Nothing lookupExp (PushExp env e) x | Just Refl <- match e x = Just ZeroIdx | otherwise = SuccIdx `fmap` lookupExp env x - -weakenGamma1 - :: Gamma env env' aenv - -> Gamma env env' (aenv,t) -weakenGamma1 EmptyExp = EmptyExp -weakenGamma1 (PushExp env e) = PushExp (weakenGamma1 env) (weaken SuccIdx e) - -sinkGamma - :: Kit acc - => Extend acc aenv aenv' - -> Gamma env env' aenv - -> Gamma env env' aenv' -sinkGamma _ EmptyExp = EmptyExp -sinkGamma ext (PushExp env e) = PushExp (sinkGamma ext env) (sinkA ext e) --} -- As part of various transformations we often need to lift out array valued @@ -152,6 +141,9 @@ sinkWeaken (PushEnv e (LeftHandSideSingle _) _) = weakenSucc' $ sinkWeaken e sinkWeaken (PushEnv e (LeftHandSidePair l1 l2) _) = sinkWeaken (PushEnv (PushEnv e l1 undefined) l2 undefined) sinkWeaken BaseEnv = Stats.substitution "sink" weakenId +sinkArrayInstr :: RebuildableExp f => Extend s acc aenv aenv' -> f (ArrayInstr aenv) env t -> f (ArrayInstr aenv') env t +sinkArrayInstr env = weakenArrayInstr (sinkWeaken env) + -- Wrapper around OpenExp, with the order of type arguments env and aenv flipped newtype OpenExp' aenv env e = OpenExp' (OpenExp env aenv e) @@ -160,4 +152,3 @@ bindExps :: Extend ScalarType (OpenExp' aenv) env env' -> OpenExp env aenv e bindExps BaseEnv = id bindExps (PushEnv g lhs (OpenExp' b)) = bindExps g . Let lhs b - diff --git a/src/Data/Array/Accelerate/Trafo/Algebra.hs b/src/Data/Array/Accelerate/Trafo/Exp/Algebra.hs similarity index 80% rename from src/Data/Array/Accelerate/Trafo/Algebra.hs rename to src/Data/Array/Accelerate/Trafo/Exp/Algebra.hs index 4c0fbd397..f3cdfab6f 100644 --- a/src/Data/Array/Accelerate/Trafo/Algebra.hs +++ b/src/Data/Array/Accelerate/Trafo/Exp/Algebra.hs @@ -1,15 +1,6 @@ -{-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE PatternGuards #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TupleSections #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE ViewPatterns #-} -- | --- Module : Data.Array.Accelerate.Trafo.Algebra +-- Module : Data.Array.Accelerate.Trafo.Exp.Algebra -- Copyright : [2012..2020] The Accelerate Team -- License : BSD3 -- @@ -21,14 +12,13 @@ -- and using algebraic properties of particular operator-operand combinations. -- -module Data.Array.Accelerate.Trafo.Algebra ( +module Data.Array.Accelerate.Trafo.Exp.Algebra ( evalPrimApp, ) where -import Data.Array.Accelerate.AST -import Data.Array.Accelerate.AST.Var +import Data.Array.Accelerate.AST.Exp import Data.Array.Accelerate.Analysis.Match import Data.Array.Accelerate.Pretty.Print ( primOperator, isInfix, opName ) import Data.Array.Accelerate.Trafo.Environment @@ -51,32 +41,30 @@ import GHC.Float ( float2Doub -- or constant let bindings. Be careful not to follow self-cycles. -- propagate - :: forall env aenv exp. - Gamma env env aenv - -> OpenExp env aenv exp + :: forall arr env exp. + IsArrayInstr arr + => PreOpenExp arr env exp -> Maybe exp -propagate env = cvtE +propagate = cvtE where - cvtE :: OpenExp env aenv e -> Maybe e + cvtE :: PreOpenExp arr env e -> Maybe e cvtE exp = case exp of Const _ c -> Just c - PrimConst c -> Just (evalPrimConst c) - Evar (Var _ ix) - | e <- prjExp ix env - , Nothing <- matchOpenExp exp e -> cvtE e Nil -> Just () Pair e1 e2 -> (,) <$> cvtE e1 <*> cvtE e2 + -- Evar is already inlined if it has a constant value _ -> Nothing -- Attempt to evaluate primitive function applications -- evalPrimApp - :: forall env aenv a r. - Gamma env env aenv + :: forall arr env a r. + IsArrayInstr arr + => Gamma arr env env -> PrimFun (a -> r) - -> OpenExp env aenv a - -> (Any, OpenExp env aenv r) + -> PreOpenExp arr env a + -> (Any, PreOpenExp arr env r) evalPrimApp env f x -- First attempt to move constant values towards the left | Just r <- commutes f x env = evalPrimApp env f r @@ -135,12 +123,10 @@ evalPrimApp env f x PrimCeiling ta tb -> evalCeiling ta tb x env PrimIsNaN ty -> evalIsNaN ty x env PrimIsInfinite ty -> evalIsInfinite ty x env - PrimLt ty -> evalLt ty x env - PrimGt ty -> evalGt ty x env - PrimLtEq ty -> evalLtEq ty x env - PrimGtEq ty -> evalGtEq ty x env - PrimEq ty -> evalEq ty x env - PrimNEq ty -> evalNEq ty x env + PrimCmp ty CmpLt -> evalLt ty x env + PrimCmp ty CmpGtEq -> evalGtEq ty x env + PrimCmp ty CmpEq -> evalEq ty x env + PrimCmp ty CmpNEq -> evalNEq ty x env PrimMax ty -> evalMax ty x env PrimMin ty -> evalMin ty x env PrimLAnd -> evalLAnd x env @@ -155,27 +141,28 @@ evalPrimApp env f x -- to the left of the operator. Returning Nothing indicates no change is made. -- commutes - :: forall env aenv a r. - PrimFun (a -> r) - -> OpenExp env aenv a - -> Gamma env env aenv - -> Maybe (OpenExp env aenv a) -commutes f x env = case f of + :: forall arr env a r. + IsArrayInstr arr + => PrimFun (a -> r) + -> PreOpenExp arr env a + -> Gamma arr env env + -> Maybe (PreOpenExp arr env a) +commutes f x _env = case f of PrimAdd _ -> swizzle x PrimMul _ -> swizzle x PrimBAnd _ -> swizzle x PrimBOr _ -> swizzle x PrimBXor _ -> swizzle x - PrimEq _ -> swizzle x - PrimNEq _ -> swizzle x + PrimCmp _ CmpEq -> swizzle x + PrimCmp _ CmpNEq -> swizzle x PrimMax _ -> swizzle x PrimMin _ -> swizzle x _ -> Nothing where - swizzle :: OpenExp env aenv (b,b) -> Maybe (OpenExp env aenv (b,b)) + swizzle :: PreOpenExp arr env (b,b) -> Maybe (PreOpenExp arr env (b,b)) swizzle (Pair a b) - | Nothing <- propagate env a - , Just _ <- propagate env b + | Nothing <- propagate a + , Just _ <- propagate b = Stats.ruleFired (pprFun "commutes" f) $ Just $ Pair b a @@ -183,8 +170,8 @@ commutes f x env = case f of -- disadvantageous: for example in (x &&* y), the user might have put a -- simpler condition first that is designed to fail fast. -- --- | Nothing <- propagate env a --- , Nothing <- propagate env b +-- | Nothing <- propagate a +-- , Nothing <- propagate b -- , hashOpenExp a > hashOpenExp b -- = Just $ Tuple (NilTup `SnocTup` b `SnocTup` a) @@ -209,8 +196,8 @@ commutes f x env = case f of associates :: (Elt a, Elt r) => PrimFun (a -> r) - -> OpenExp env aenv a - -> Maybe (OpenExp env aenv r) + -> PreOpenExp arr env a + -> Maybe (PreOpenExp arr env r) associates fun exp = case fun of PrimAdd _ -> swizzle fun exp [PrimAdd ty, PrimSub ty] PrimSub _ -> swizzle fun exp [PrimAdd ty, PrimSub ty] @@ -222,7 +209,7 @@ associates fun exp = case fun of ty = undefined ops = [ PrimMul ty, PrimFDiv ty, PrimAdd ty, PrimSub ty, PrimBAnd ty, PrimBOr ty, PrimBXor ty ] - swizzle :: (Elt a, Elt r) => PrimFun (a -> r) -> OpenExp env aenv a -> [PrimFun (a -> r)] -> Maybe (OpenExp env aenv r) + swizzle :: (Elt a, Elt r) => PrimFun (a -> r) -> PreOpenExp arr env a -> [PrimFun (a -> r)] -> Maybe (PreOpenExp arr env r) swizzle f x lvl | Just Refl <- matches f ops , Just (a,bc) <- untup2 x @@ -249,17 +236,17 @@ associates fun exp = case fun of -- Helper functions -- ---------------- -type a :-> b = forall env aenv. OpenExp env aenv a -> Gamma env env aenv -> Maybe (OpenExp env aenv b) +type a :-> b = forall arr env. IsArrayInstr arr => PreOpenExp arr env a -> Gamma arr env env -> Maybe (PreOpenExp arr env b) eval1 :: SingleType b -> (a -> b) -> a :-> b -eval1 tp f x env - | Just a <- propagate env x = Stats.substitution "constant fold" . Just $ Const (SingleScalarType tp) (f a) - | otherwise = Nothing +eval1 tp f x _env + | Just a <- propagate x = Stats.substitution "constant fold" . Just $ Const (SingleScalarType tp) (f a) + | otherwise = Nothing eval2 :: SingleType c -> (a -> b -> c) -> (a,b) :-> c -eval2 tp f (untup2 -> Just (x,y)) env - | Just a <- propagate env x - , Just b <- propagate env y +eval2 tp f (untup2 -> Just (x,y)) _env + | Just a <- propagate x + , Just b <- propagate y = Stats.substitution "constant fold" $ Just $ Const (SingleScalarType tp) (f a b) eval2 _ _ _ _ @@ -274,30 +261,47 @@ toBool 0 = False toBool _ = True bool1 :: (a -> Bool) -> a :-> PrimBool -bool1 f x env - | Just a <- propagate env x +bool1 f x _env + | Just a <- propagate x = Stats.substitution "constant fold" . Just $ Const scalarTypeWord8 (fromBool (f a)) bool1 _ _ _ = Nothing bool2 :: (a -> b -> Bool) -> (a,b) :-> PrimBool -bool2 f (untup2 -> Just (x,y)) env - | Just a <- propagate env x - , Just b <- propagate env y +bool2 f (untup2 -> Just (x,y)) _env + | Just a <- propagate x + , Just b <- propagate y = Stats.substitution "constant fold" $ Just $ Const scalarTypeWord8 (fromBool (f a b)) bool2 _ _ _ = Nothing -tup2 :: (OpenExp env aenv a, OpenExp env aenv b) -> OpenExp env aenv (a, b) +bool2IfEq :: Bool -> (a -> b -> Bool) -> (a,b) :-> PrimBool +bool2IfEq ifEq f (untup2 -> Just (x,y)) _env + | Just _ <- matchOpenExp x y + = Stats.substitution "equal comparison fold" + $ Just $ Const scalarTypeWord8 (fromBool ifEq) + | Just a <- propagate x + , Just b <- propagate y + = Stats.substitution "constant fold" + $ Just $ Const scalarTypeWord8 (fromBool (f a b)) +bool2IfEq _ _ _ _ + = Nothing + +tup2 :: (PreOpenExp arr env a, PreOpenExp arr env b) -> PreOpenExp arr env (a, b) tup2 (a,b) = Pair a b -untup2 :: OpenExp env aenv (a, b) -> Maybe (OpenExp env aenv a, OpenExp env aenv b) +untup2 :: PreOpenExp arr env (a, b) -> Maybe (PreOpenExp arr env a, PreOpenExp arr env b) untup2 exp | Pair a b <- exp = Just (a, b) | otherwise = Nothing +-- If the two arguments are the same, then it will evaluate to that argument +evalIfSameArguments :: (a,a) :-> a +evalIfSameArguments (untup2 -> Just (x,y)) _ + | Just Refl <- matchOpenExp x y = Just x +evalIfSameArguments _ _ = Nothing pprFun :: Text -> PrimFun f -> Text pprFun rule f @@ -319,8 +323,8 @@ evalAdd ty@(IntegralNumType ty') | IntegralDict <- integralDict ty' = evalAdd' t evalAdd ty@(FloatingNumType ty') | FloatingDict <- floatingDict ty' = evalAdd' ty evalAdd' :: (Eq a, Num a) => NumType a -> (a,a) :-> a -evalAdd' _ (untup2 -> Just (x,y)) env - | Just a <- propagate env x +evalAdd' _ (untup2 -> Just (x,y)) _env + | Just a <- propagate x , a == 0 = Stats.ruleFired "x+0" $ Just y @@ -334,12 +338,12 @@ evalSub ty@(FloatingNumType ty') | FloatingDict <- floatingDict ty' = evalSub' t evalSub' :: forall a. (Eq a, Num a) => NumType a -> (a,a) :-> a evalSub' ty (untup2 -> Just (x,y)) env - | Just b <- propagate env y + | Just b <- propagate y , b == 0 = Stats.ruleFired "x-0" $ Just x - | Nothing <- propagate env x - , Just b <- propagate env y + | Nothing <- propagate x + , Just b <- propagate y = Stats.ruleFired "-y+x" $ Just . snd $ evalPrimApp env (PrimAdd ty) (Const tp (-b) `Pair` x) -- (Tuple $ NilTup `SnocTup` Const (fromElt (-b)) `SnocTup` x) @@ -359,9 +363,9 @@ evalMul ty@(IntegralNumType ty') | IntegralDict <- integralDict ty' = evalMul' t evalMul ty@(FloatingNumType ty') | FloatingDict <- floatingDict ty' = evalMul' ty evalMul' :: (Eq a, Num a) => NumType a -> (a,a) :-> a -evalMul' _ (untup2 -> Just (x,y)) env - | Just a <- propagate env x - , Nothing <- propagate env y +evalMul' _ (untup2 -> Just (x,y)) _env + | Just a <- propagate x + , Nothing <- propagate y = case a of 0 -> Stats.ruleFired "x*0" $ Just x 1 -> Stats.ruleFired "x*1" $ Just y @@ -404,14 +408,14 @@ evalRem _ _ _ = Nothing evalQuotRem :: forall a. IntegralType a -> (a,a) :-> (a,a) -evalQuotRem ty exp env +evalQuotRem ty exp _env | IntegralDict <- integralDict ty , Just (x, y) <- untup2 exp - , Just b <- propagate env y + , Just b <- propagate y = case b of 0 -> Nothing 1 -> Stats.ruleFired "quotRem x 1" $ Just (tup2 (x, Const tp 0)) - _ -> case propagate env x of + _ -> case propagate x of Nothing -> Nothing Just a -> Stats.substitution "constant fold" $ Just $ let (u,v) = quotRem a b @@ -439,14 +443,14 @@ evalMod _ _ _ = Nothing evalDivMod :: forall a. IntegralType a -> (a,a) :-> (a,a) -evalDivMod ty exp env +evalDivMod ty exp _env | IntegralDict <- integralDict ty , Just (x, y) <- untup2 exp - , Just b <- propagate env y + , Just b <- propagate y = case b of 0 -> Nothing 1 -> Stats.ruleFired "divMod x 1" $ Just (tup2 (x, Const tp 0)) - _ -> case propagate env x of + _ -> case propagate x of Nothing -> Nothing Just a -> Stats.substitution "constant fold" $ Just $ let (u,v) = divMod a b @@ -463,8 +467,8 @@ evalBOr :: IntegralType a -> (a,a) :-> a evalBOr ty | IntegralDict <- integralDict ty = evalBOr' ty evalBOr' :: (Eq a, Num a, Bits a) => IntegralType a -> (a,a) :-> a -evalBOr' _ (untup2 -> Just (x,y)) env - | Just 0 <- propagate env x +evalBOr' _ (untup2 -> Just (x,y)) _env + | Just 0 <- propagate x = Stats.ruleFired "x .|. 0" $ Just y evalBOr' ty arg env @@ -477,31 +481,31 @@ evalBNot :: IntegralType a -> a :-> a evalBNot ty | IntegralDict <- integralDict ty = eval1 (NumSingleType $ IntegralNumType ty) complement evalBShiftL :: IntegralType a -> (a,Int) :-> a -evalBShiftL _ (untup2 -> Just (x,i)) env - | Just 0 <- propagate env i +evalBShiftL _ (untup2 -> Just (x,i)) _env + | Just 0 <- propagate i = Stats.ruleFired "x `shiftL` 0" $ Just x evalBShiftL ty arg env | IntegralDict <- integralDict ty = eval2 (NumSingleType $ IntegralNumType ty) shiftL arg env evalBShiftR :: IntegralType a -> (a,Int) :-> a -evalBShiftR _ (untup2 -> Just (x,i)) env - | Just 0 <- propagate env i +evalBShiftR _ (untup2 -> Just (x,i)) _env + | Just 0 <- propagate i = Stats.ruleFired "x `shiftR` 0" $ Just x evalBShiftR ty arg env | IntegralDict <- integralDict ty = eval2 (NumSingleType $ IntegralNumType ty) shiftR arg env evalBRotateL :: IntegralType a -> (a,Int) :-> a -evalBRotateL _ (untup2 -> Just (x,i)) env - | Just 0 <- propagate env i +evalBRotateL _ (untup2 -> Just (x,i)) _env + | Just 0 <- propagate i = Stats.ruleFired "x `rotateL` 0" $ Just x evalBRotateL ty arg env | IntegralDict <- integralDict ty = eval2 (NumSingleType $ IntegralNumType ty) rotateL arg env evalBRotateR :: IntegralType a -> (a,Int) :-> a -evalBRotateR _ (untup2 -> Just (x,i)) env - | Just 0 <- propagate env i +evalBRotateR _ (untup2 -> Just (x,i)) _env + | Just 0 <- propagate i = Stats.ruleFired "x `rotateR` 0" $ Just x evalBRotateR ty arg env | IntegralDict <- integralDict ty = eval2 (NumSingleType $ IntegralNumType ty) rotateR arg env @@ -523,8 +527,8 @@ evalFDiv :: FloatingType a -> (a,a) :-> a evalFDiv ty | FloatingDict <- floatingDict ty = evalFDiv' ty evalFDiv' :: (Fractional a, Eq a) => FloatingType a -> (a,a) :-> a -evalFDiv' _ (untup2 -> Just (x,y)) env - | Just 1 <- propagate env y +evalFDiv' _ (untup2 -> Just (x,y)) _env + | Just 1 <- propagate y = Stats.ruleFired "x/1" $ Just x evalFDiv' ty arg env @@ -623,48 +627,58 @@ evalIsInfinite ty | FloatingDict <- floatingDict ty = bool1 isInfinite -- --------------------- evalLt :: SingleType a -> (a,a) :-> PrimBool -evalLt (NumSingleType (IntegralNumType ty)) | IntegralDict <- integralDict ty = bool2 (<) +evalLt (NumSingleType (IntegralNumType ty)) | IntegralDict <- integralDict ty = bool2IfEq False (<) evalLt (NumSingleType (FloatingNumType ty)) | FloatingDict <- floatingDict ty = bool2 (<) -evalGt :: SingleType a -> (a,a) :-> PrimBool -evalGt (NumSingleType (IntegralNumType ty)) | IntegralDict <- integralDict ty = bool2 (>) -evalGt (NumSingleType (FloatingNumType ty)) | FloatingDict <- floatingDict ty = bool2 (>) - -evalLtEq :: SingleType a -> (a,a) :-> PrimBool -evalLtEq (NumSingleType (IntegralNumType ty)) | IntegralDict <- integralDict ty = bool2 (<=) -evalLtEq (NumSingleType (FloatingNumType ty)) | FloatingDict <- floatingDict ty = bool2 (<=) - evalGtEq :: SingleType a -> (a,a) :-> PrimBool -evalGtEq (NumSingleType (IntegralNumType ty)) | IntegralDict <- integralDict ty = bool2 (>=) +evalGtEq (NumSingleType (IntegralNumType ty)) | IntegralDict <- integralDict ty = bool2IfEq True (>=) evalGtEq (NumSingleType (FloatingNumType ty)) | FloatingDict <- floatingDict ty = bool2 (>=) evalEq :: SingleType a -> (a,a) :-> PrimBool -evalEq (NumSingleType (IntegralNumType ty)) | IntegralDict <- integralDict ty = bool2 (==) +evalEq (NumSingleType (IntegralNumType ty)) | IntegralDict <- integralDict ty = bool2IfEq True (==) evalEq (NumSingleType (FloatingNumType ty)) | FloatingDict <- floatingDict ty = bool2 (==) evalNEq :: SingleType a -> (a,a) :-> PrimBool -evalNEq (NumSingleType (IntegralNumType ty)) | IntegralDict <- integralDict ty = bool2 (/=) +evalNEq (NumSingleType (IntegralNumType ty)) | IntegralDict <- integralDict ty = bool2IfEq False (/=) evalNEq (NumSingleType (FloatingNumType ty)) | FloatingDict <- floatingDict ty = bool2 (/=) evalMax :: SingleType a -> (a,a) :-> a -evalMax ty@(NumSingleType (IntegralNumType ty')) | IntegralDict <- integralDict ty' = eval2 ty max -evalMax ty@(NumSingleType (FloatingNumType ty')) | FloatingDict <- floatingDict ty' = eval2 ty max +evalMax ty args env + | Just arg <- evalIfSameArguments args env + = Stats.ruleFired "max x x" $ Just arg + + | NumSingleType (IntegralNumType ty') <- ty + , IntegralDict <- integralDict ty' + = eval2 ty max args env + + | NumSingleType (FloatingNumType ty') <- ty + , FloatingDict <- floatingDict ty' + = eval2 ty max args env evalMin :: SingleType a -> (a,a) :-> a -evalMin ty@(NumSingleType (IntegralNumType ty')) | IntegralDict <- integralDict ty' = eval2 ty min -evalMin ty@(NumSingleType (FloatingNumType ty')) | FloatingDict <- floatingDict ty' = eval2 ty min +evalMin ty args env + | Just arg <- evalIfSameArguments args env + = Stats.ruleFired "min x x" $ Just arg + + | NumSingleType (IntegralNumType ty') <- ty + , IntegralDict <- integralDict ty' + = eval2 ty min args env + + | NumSingleType (FloatingNumType ty') <- ty + , FloatingDict <- floatingDict ty' + = eval2 ty min args env -- Logical operators -- ----------------- evalLAnd :: (PrimBool,PrimBool) :-> PrimBool -evalLAnd (untup2 -> Just (x,y)) env - | Just a <- propagate env x +evalLAnd (untup2 -> Just (x,y)) _env + | Just a <- propagate x = Just $ if toBool a then Stats.ruleFired "True &&" y else Stats.ruleFired "False &&" $ Const scalarTypeWord8 0 - | Just b <- propagate env y + | Just b <- propagate y = Just $ if toBool b then Stats.ruleFired "True &&" x else Stats.ruleFired "False &&" $ Const scalarTypeWord8 0 @@ -673,13 +687,13 @@ evalLAnd _ _ = Nothing evalLOr :: (PrimBool,PrimBool) :-> PrimBool -evalLOr (untup2 -> Just (x,y)) env - | Just a <- propagate env x +evalLOr (untup2 -> Just (x,y)) _env + | Just a <- propagate x = Just $ if toBool a then Stats.ruleFired "True ||" $ Const scalarTypeWord8 1 else Stats.ruleFired "False ||" y - | Just b <- propagate env y + | Just b <- propagate y = Just $ if toBool b then Stats.ruleFired "True ||" $ Const scalarTypeWord8 1 else Stats.ruleFired "False ||" x @@ -688,8 +702,22 @@ evalLOr _ _ = Nothing evalLNot :: PrimBool :-> PrimBool -evalLNot x _ | PrimApp PrimLNot x' <- x = Stats.ruleFired "not/not" $ Just x' -evalLNot x env = bool1 (not . toBool) x env +evalLNot (PrimApp PrimLNot y) _ = Stats.ruleFired "not/not" $ Just y +evalLNot (PrimApp (PrimCmp tp c) y) _ + -- Only negate a comparison operator on integral type. + -- On floating point types this is unsound due to the NaN semantics. + | NumSingleType (IntegralNumType _) <- tp + = Stats.ruleFired "not/cmp" $ Just $ PrimApp (PrimCmp tp (negateCmp c)) y +-- Replace 'not (a && b)' by 'not a || not b', +-- so other simplifications can trigger +evalLNot (PrimApp PrimLAnd (Pair a b)) _ = Stats.ruleFired "not/and" $ Just $ PrimApp PrimLOr $ Pair + (PrimApp PrimLNot a) + (PrimApp PrimLNot b) +-- Similarly, replace 'not (a || b)' by 'not a && not b' +evalLNot (PrimApp PrimLOr (Pair a b)) _ = Stats.ruleFired "not/or" $ Just $ PrimApp PrimLAnd $ Pair + (PrimApp PrimLNot a) + (PrimApp PrimLNot b) +evalLNot x env = bool1 (not . toBool) x env evalFromIntegral :: IntegralType a -> NumType b -> a :-> b evalFromIntegral ta (IntegralNumType tb) @@ -723,22 +751,3 @@ evalToFloating (FloatingNumType ta) tb x env | FloatingDict <- floatingDict ta , FloatingDict <- floatingDict tb = eval1 (NumSingleType $ FloatingNumType tb) realToFrac x env - - --- Scalar primitives --- ----------------- - -evalPrimConst :: PrimConst a -> a -evalPrimConst (PrimMinBound ty) = evalMinBound ty -evalPrimConst (PrimMaxBound ty) = evalMaxBound ty -evalPrimConst (PrimPi ty) = evalPi ty - -evalMinBound :: BoundedType a -> a -evalMinBound (IntegralBoundedType ty) | IntegralDict <- integralDict ty = minBound - -evalMaxBound :: BoundedType a -> a -evalMaxBound (IntegralBoundedType ty) | IntegralDict <- integralDict ty = maxBound - -evalPi :: FloatingType a -> a -evalPi ty | FloatingDict <- floatingDict ty = pi - diff --git a/src/Data/Array/Accelerate/Trafo/Exp/Shrink.hs b/src/Data/Array/Accelerate/Trafo/Exp/Shrink.hs new file mode 100644 index 000000000..1716ef35b --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Exp/Shrink.hs @@ -0,0 +1,413 @@ +-- | +-- Module : Data.Array.Accelerate.Trafo.Shrink +-- Copyright : [2012..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- +-- The shrinking substitution arises as a restriction of beta-reduction to cases +-- where the bound variable is used zero (dead-code elimination) or one (linear +-- inlining) times. By simplifying terms, the shrinking reduction can expose +-- opportunities for further optimisation. +-- +-- TODO: replace with a linear shrinking algorithm; e.g. +-- +-- * Andrew Appel & Trevor Jim, "Shrinking lambda expressions in linear time". +-- +-- * Nick Benton, Andrew Kennedy, Sam Lindley and Claudio Russo, "Shrinking +-- Reductions in SML.NET" +-- + +module Data.Array.Accelerate.Trafo.Exp.Shrink ( + + -- Shrinking + shrinkExp, + shrinkFun, + + -- Occurrence counting + usesOfExp, usesOfFun, + + -- Utilities + arrayInstrsInExp, arrayInstrsInFun, + strengthenShrunkLHS, + +) where + +import Data.Array.Accelerate.AST.Environment +import Data.Array.Accelerate.AST.Exp +import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.AST.Var +import Data.Array.Accelerate.Analysis.Match +import Data.Array.Accelerate.Error +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Trafo.Substitution +import Data.Array.Accelerate.Trafo.Exp.Substitution + +import qualified Data.Array.Accelerate.Debug.Internal.Stats as Stats + +import Control.Applicative hiding ( Const ) +import Data.Maybe ( isJust ) +import Data.Monoid +import Data.Semigroup +import Prelude hiding ( exp, seq ) + +data VarsRange env = + VarsRange !(Exists (Idx env)) -- rightmost variable + {-# UNPACK #-} !Int -- count + !(Maybe RangeTuple) -- tuple + +data RangeTuple + = RTNil + | RTSingle + | RTPair !RangeTuple !RangeTuple + +lhsVarsRange :: LeftHandSide s v env env' -> Either (env :~: env') (VarsRange env') +lhsVarsRange lhs = case rightIx lhs of + Left eq -> Left eq + Right ix -> let (n, rt) = go lhs + in Right $ VarsRange ix n rt + where + rightIx :: LeftHandSide s v env env' -> Either (env :~: env') (Exists (Idx env')) + rightIx (LeftHandSideWildcard _) = Left Refl + rightIx (LeftHandSideSingle _) = Right $ Exists ZeroIdx + rightIx (LeftHandSidePair l1 l2) = case rightIx l2 of + Right ix -> Right ix + Left Refl -> rightIx l1 + + go :: LeftHandSide s v env env' -> (Int, Maybe (RangeTuple)) + go (LeftHandSideWildcard TupRunit) = (0, Just RTNil) + go (LeftHandSideWildcard _) = (0, Nothing) + go (LeftHandSideSingle _) = (1, Just RTSingle) + go (LeftHandSidePair l1 l2) = (n1 + n2, RTPair <$> t1 <*> t2) + where + (n1, t1) = go l1 + (n2, t2) = go l2 + +weakenVarsRange :: LeftHandSide s v env env' -> VarsRange env -> VarsRange env' +weakenVarsRange lhs (VarsRange ix n t) = VarsRange (go lhs ix) n t + where + go :: LeftHandSide s v env env' -> Exists (Idx env) -> Exists (Idx env') + go (LeftHandSideWildcard _) ix' = ix' + go (LeftHandSideSingle _) (Exists ix') = Exists (SuccIdx ix') + go (LeftHandSidePair l1 l2) ix' = go l2 $ go l1 ix' + +matchEVarsRange :: VarsRange env -> PreOpenExp arr env t -> Bool +matchEVarsRange (VarsRange (Exists first) _ (Just rt)) expr = isJust $ go (idxToInt first) rt expr + where + go :: Int -> RangeTuple -> PreOpenExp arr env t -> Maybe Int + go i RTNil Nil = Just i + go i RTSingle (Evar (Var _ ix)) + | checkIdx i ix = Just (i + 1) + go i (RTPair t1 t2) (Pair e1 e2) + | Just i' <- go i t2 e2 = go i' t1 e1 + go _ _ _ = Nothing + + checkIdx :: Int -> Idx env t -> Bool + checkIdx 0 ZeroIdx = True + checkIdx i (SuccIdx ix) = checkIdx (i - 1) ix + checkIdx _ _ = False +matchEVarsRange _ _ = False + +varInRange :: VarsRange env -> Var s env t -> Maybe Usages +varInRange (VarsRange (Exists rangeIx) n _) (Var _ varIx) = case go rangeIx varIx of + Nothing -> Nothing + Just j -> Just $ replicate j False ++ [True] ++ replicate (n - j - 1) False + where + -- `go ix ix'` checks whether ix <= ix' with recursion, and then checks + -- whether ix' < ix + n in go'. Returns a Just if both checks + -- are successful, containing an integer j such that ix + j = ix'. + go :: Idx env u -> Idx env t -> Maybe Int + go (SuccIdx ix) (SuccIdx ix') = go ix ix' + go ZeroIdx ix' = go' ix' 0 + go _ ZeroIdx = Nothing + + go' :: Idx env t -> Int -> Maybe Int + go' _ j | j >= n = Nothing + go' ZeroIdx j = Just j + go' (SuccIdx ix') j = go' ix' (j + 1) + +-- Describes how often the variables defined in a LHS are used together. +data Count + = Impossible !Usages + -- Cannot inline this definition. This happens when the definition + -- declares multiple variables (the right hand side returns a tuple) + -- and the variables are used seperately. + | Infinity + -- The variable is used in a loop. Inlining should only proceed if + -- the computation is cheap. + | Finite {-# UNPACK #-} !Int + +type Usages = [Bool] -- Per variable a Boolean denoting whether that variable is used. + +instance Semigroup Count where + Impossible u1 <> Impossible u2 = Impossible $ zipWith (||) u1 u2 + Impossible u <> Finite 0 = Impossible u + Finite 0 <> Impossible u = Impossible u + Impossible u <> _ = Impossible $ map (const True) u + _ <> Impossible u = Impossible $ map (const True) u + Infinity <> _ = Infinity + _ <> Infinity = Infinity + Finite a <> Finite b = Finite $ a + b + +instance Monoid Count where + mempty = Finite 0 + +loopCount :: Count -> Count +loopCount (Finite n) | n > 0 = Infinity +loopCount c = c + +shrinkLhs + :: HasCallStack + => Count + -> LeftHandSide s t env1 env2 + -> Maybe (Exists (LeftHandSide s t env1)) +shrinkLhs _ (LeftHandSideWildcard _) = Nothing -- We cannot shrink this +shrinkLhs (Finite 0) lhs = Just $ Exists $ LeftHandSideWildcard $ lhsToTupR lhs -- LHS isn't used at all, replace with a wildcard +shrinkLhs (Impossible usages) lhs = case go usages lhs of + (True , [], lhs') -> Just lhs' + (False, [], _ ) -> Nothing -- No variables were dropped. Thus lhs == lhs'. + _ -> internalError "Mismatch in length of usages array and LHS" + where + go :: HasCallStack => Usages -> LeftHandSide s t env1 env2 -> (Bool, Usages, Exists (LeftHandSide s t env1)) + go us (LeftHandSideWildcard tp) = (False, us, Exists $ LeftHandSideWildcard tp) + go (True : us) (LeftHandSideSingle tp) = (False, us, Exists $ LeftHandSideSingle tp) + go (False : us) (LeftHandSideSingle tp) = (True , us, Exists $ LeftHandSideWildcard $ TupRsingle tp) + go us (LeftHandSidePair l1 l2) + | (c2, us' , Exists l2') <- go us l2 + , (c1, us'', Exists l1') <- go us' l1 + , Exists l2'' <- rebuildLHS l2' + = let + lhs' + | LeftHandSideWildcard t1 <- l1' + , LeftHandSideWildcard t2 <- l2'' = LeftHandSideWildcard $ TupRpair t1 t2 + | otherwise = LeftHandSidePair l1' l2'' + in + (c1 || c2, us'', Exists lhs') + go _ _ = internalError "Empty array, mismatch in length of usages array and LHS" +shrinkLhs _ _ = Nothing + +-- Converts a strengthening from before a binding, to a strengthening after a binding. +-- The two LeftHandSides may have different structures. +strengthenShrunkLHS + :: Distributes s + => LeftHandSide s t env1 env2 + -> LeftHandSide s t env1' env2' + -> env1 :?> env1' + -> env2 :?> env2' +strengthenShrunkLHS (LeftHandSideWildcard _) (LeftHandSideWildcard _) k = k +strengthenShrunkLHS (LeftHandSideSingle _) (LeftHandSideSingle _) k = \ix -> case ix of + ZeroIdx -> Just ZeroIdx + SuccIdx ix' -> SuccIdx <$> k ix' +strengthenShrunkLHS (LeftHandSidePair lA hA) (LeftHandSidePair lB hB) k = strengthenShrunkLHS hA hB $ strengthenShrunkLHS lA lB k +strengthenShrunkLHS lhs (LeftHandSideWildcard _) k = \ix -> strengthenWithLHS lhs ix >>= k +strengthenShrunkLHS (LeftHandSideWildcard _) lhs k = \ix -> k ix >>= strengthenFromWeaken (weakenWithLHS lhs) +strengthenShrunkLHS (LeftHandSideSingle t) (LeftHandSidePair _ _) _ = pairImpossible t +strengthenShrunkLHS (LeftHandSidePair _ _) (LeftHandSideSingle t) _ = pairImpossible t + +strengthenFromWeaken :: env1 :> env2 -> env1 :?> env2 +strengthenFromWeaken k ix = Just $ k >:> ix + +-- Shrinking +-- ========= + +-- The shrinking substitution for scalar expressions. This is a restricted +-- instance of beta-reduction to cases where the bound variable is used zero +-- (dead-code elimination) or one (linear inlining) times. +-- +shrinkExp :: (HasCallStack, IsArrayInstr arr) => PreOpenExp arr env t -> (Bool, PreOpenExp arr env t) +shrinkExp = Stats.substitution "shrinkE" . first getAny . shrinkE + where + -- If the bound variable is used at most this many times, it will be inlined + -- into the body. In cases where it is not used at all, this is equivalent + -- to dead-code elimination. + -- + lIMIT :: Int + lIMIT = 1 + + cheap :: IsArrayInstr arr => PreOpenExp arr env t -> Bool + cheap (Evar _) = True + cheap (Pair e1 e2) = cheap e1 && cheap e2 + cheap Nil = True + cheap Const{} = True + cheap Undef{} = True + cheap (Coerce _ _ e) = cheap e + cheap (ArrayInstr arr Nil) = inlineArrayInstr arr + cheap _ = False + + shrinkE :: (HasCallStack, IsArrayInstr arr) => PreOpenExp arr env t -> (Any, PreOpenExp arr env t) + shrinkE exp = case exp of + Let (LeftHandSideSingle _) bnd@Evar{} body -> Stats.inline "Var" . yes $ shrinkE (inline body bnd) + Let lhs bnd body + | shouldInline -> case inlineVars lhs (snd body') (snd bnd') of + Just inlined -> Stats.betaReduce msg . yes $ shrinkE inlined + _ -> internalError "Unexpected failure while trying to inline some expression." + | Just (Exists lhs') <- shrinkLhs count lhs -> case strengthenE (strengthenShrunkLHS lhs lhs' Just) (snd body') of + Just body'' -> (Any True, Let lhs' (snd bnd') body'') + Nothing -> internalError "Unexpected failure in strenthenE. Variable was analysed to be unused in usesOfExp, but appeared to be used in strenthenE." + | otherwise -> Let lhs <$> bnd' <*> body' + where + shouldInline = case count of + Finite 0 -> False -- Handled by shrinkLhs + Finite n -> n <= lIMIT || cheap (snd bnd') + Infinity -> cheap (snd bnd') + Impossible _ -> False + + bnd' = shrinkE bnd + body' = shrinkE body + + -- If the lhs includes non-trivial wildcards (the last field of range is Nothing), + -- then we cannot inline the binding. We can only check which variables are not used, + -- to detect unused variables. + -- + -- If the lhs does not include non-trivial wildcards (the last field of range is a Just), + -- we can both analyse whether we can inline the binding, and check which variables are + -- not used, to detect unused variables. + -- + count = case lhsVarsRange lhs of + Left _ -> Finite 0 + Right range -> usesOfExp range (snd body') + + msg = case count of + Finite 0 -> "dead exp" + _ -> "inline exp" -- forced inlining when lIMIT > 1 + -- + Evar v -> pure (Evar v) + Const t c -> pure (Const t c) + Undef t -> pure (Undef t) + Nil -> pure Nil + Pair x y -> Pair <$> shrinkE x <*> shrinkE y + VecPack vec e -> VecPack vec <$> shrinkE e + VecUnpack vec e -> VecUnpack vec <$> shrinkE e + ToIndex shr sh ix -> ToIndex shr <$> shrinkE sh <*> shrinkE ix + FromIndex shr sh i -> FromIndex shr <$> shrinkE sh <*> shrinkE i + Case e rhs def -> Case <$> shrinkE e <*> sequenceA [ (t,) <$> shrinkE c | (t,c) <- rhs ] <*> shrinkMaybeE def + Cond p t e -> Cond <$> shrinkE p <*> shrinkE t <*> shrinkE e + Select p t e -> Select <$> shrinkE p <*> shrinkE t <*> shrinkE e + While p f x -> While <$> shrinkF p <*> shrinkF f <*> shrinkE x + PrimApp f x -> PrimApp f <$> shrinkE x + ArrayInstr arr e -> ArrayInstr arr <$> shrinkE e + ShapeSize shr sh -> ShapeSize shr <$> shrinkE sh + Foreign repr ff f e -> Foreign repr ff <$> shrinkF f <*> shrinkE e + Coerce t1 t2 e -> Coerce t1 t2 <$> shrinkE e + Assert msg e1 e2 -> Assert msg <$> shrinkE e1 <*> shrinkE e2 + Assume e1 e2 -> Assume <$> shrinkE e1 <*> shrinkE e2 + + shrinkF :: (HasCallStack, IsArrayInstr arr) => PreOpenFun arr env t -> (Any, PreOpenFun arr env t) + shrinkF = first Any . shrinkFun + + shrinkMaybeE :: (HasCallStack, IsArrayInstr arr) => Maybe (PreOpenExp arr env t) -> (Any, Maybe (PreOpenExp arr env t)) + shrinkMaybeE Nothing = pure Nothing + shrinkMaybeE (Just e) = Just <$> shrinkE e + + first :: (a -> a') -> (a,b) -> (a',b) + first f (x,y) = (f x, y) + + yes :: (Any, x) -> (Any, x) + yes (_, x) = (Any True, x) + +shrinkFun :: (HasCallStack, IsArrayInstr arr) => PreOpenFun arr env f -> (Bool, PreOpenFun arr env f) +shrinkFun (Lam lhs f) = case lhsVarsRange lhs of + Left Refl -> + let b' = case lhs of + LeftHandSideWildcard TupRunit -> b + _ -> True + in (b', Lam (LeftHandSideWildcard $ lhsToTupR lhs) f') + Right range -> + let + count = usesOfFun range f + in case shrinkLhs count lhs of + Just (Exists lhs') -> case strengthenE (strengthenShrunkLHS lhs lhs' Just) f' of + Just f'' -> (True, Lam lhs' f'') + Nothing -> internalError "Unexpected failure in strenthenE. Variable was analysed to be unused in usesOfExp, but appeared to be used in strenthenE." + Nothing -> (b, Lam lhs f') + where + (b, f') = shrinkFun f + +shrinkFun (Body b) = Body <$> shrinkExp b + +-- Occurrence Counting +-- =================== + +-- Count the number of occurrences an in-scope scalar expression bound at the +-- given variable index recursively in a term. +-- +usesOfExp :: forall arr env t. VarsRange env -> PreOpenExp arr env t -> Count +usesOfExp range = countE + where + countE :: PreOpenExp arr env e -> Count + countE exp | matchEVarsRange range exp = Finite 1 + countE exp = case exp of + Evar v -> case varInRange range v of + Just cs -> Impossible cs + Nothing -> Finite 0 + -- + Let lhs bnd body -> countE bnd <> usesOfExp (weakenVarsRange lhs range) body + Const _ _ -> Finite 0 + Undef _ -> Finite 0 + Nil -> Finite 0 + Pair e1 e2 -> countE e1 <> countE e2 + VecPack _ e -> countE e + VecUnpack _ e -> countE e + FromIndex _ sh i -> countE sh <> countE i + ToIndex _ sh e -> countE sh <> countE e + Case e rhs def -> countE e <> mconcat [ countE c | (_,c) <- rhs ] <> maybe (Finite 0) countE def + Cond p t e -> countE p <> countE t <> countE e + Select p t e -> countE p <> countE t <> countE e + While p f x -> countE x <> loopCount (usesOfFun range p) <> loopCount (usesOfFun range f) + PrimApp _ x -> countE x + ArrayInstr _ e -> countE e + ShapeSize _ sh -> countE sh + Foreign _ _ _ e -> countE e + Coerce _ _ e -> countE e + Assert _ e1 e2 -> countE e1 <> countE e2 + Assume e1 e2 -> countE e1 <> countE e2 + +usesOfFun :: VarsRange env -> PreOpenFun arr env f -> Count +usesOfFun range (Lam lhs f) = usesOfFun (weakenVarsRange lhs range) f +usesOfFun range (Body b) = usesOfExp range b + +arrayInstrsInExp :: PreOpenExp arr env t -> [Exists arr] +arrayInstrsInExp = (`travE` []) + where + travE :: PreOpenExp arr env e -> [Exists arr] -> [Exists arr] + travE !exp !acc = case exp of + Let _ bnd body -> travE bnd $ travE body acc + Evar _ -> acc + Const _ _ -> acc + Undef _ -> acc + Nil -> acc + Pair x y -> travE x $ travE y acc + VecPack _ e -> travE e acc + VecUnpack _ e -> travE e acc + ToIndex _ sh ix -> travE sh $ travE ix acc + FromIndex _ sh i -> travE sh $ travE i acc + Case e rhs def -> travE e $ travAE rhs $ travME def acc + Cond p t e -> travE p $ travE t $ travE e acc + Select p t e -> travE p $ travE t $ travE e acc + While p f x -> travF p $ travF f $ travE x acc + PrimApp _ x -> travE x acc + ArrayInstr arr e -> Exists arr : travE e acc + ShapeSize _ sh -> travE sh acc + Foreign _ _ _ e -> travE e acc + Coerce _ _ e -> travE e acc + Assert _ e1 e2 -> travE e1 $ travE e2 acc + Assume e1 e2 -> travE e1 $ travE e2 acc + + travME :: Maybe (PreOpenExp arr env e) -> [Exists arr] -> [Exists arr] + travME Nothing acc = acc + travME (Just e) acc = travE e acc + + travAE :: [(TAG, PreOpenExp arr env b)] -> [Exists arr] -> [Exists arr] + travAE [] acc = acc + travAE ((_,e) : as) acc = travE e $ travAE as acc + + travF :: PreOpenFun arr env e -> [Exists arr] -> [Exists arr] + travF (Body e) = travE e + travF (Lam _ f) = travF f + +arrayInstrsInFun :: PreOpenFun arr env t -> [Exists arr] +arrayInstrsInFun (Body e) = arrayInstrsInExp e +arrayInstrsInFun (Lam _ f) = arrayInstrsInFun f diff --git a/src/Data/Array/Accelerate/Trafo/Simplify.hs b/src/Data/Array/Accelerate/Trafo/Exp/Simplify.hs similarity index 53% rename from src/Data/Array/Accelerate/Trafo/Simplify.hs rename to src/Data/Array/Accelerate/Trafo/Exp/Simplify.hs index 6bc0f769e..c292c87c2 100644 --- a/src/Data/Array/Accelerate/Trafo/Simplify.hs +++ b/src/Data/Array/Accelerate/Trafo/Exp/Simplify.hs @@ -1,22 +1,7 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE PatternGuards #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE RecordWildCards #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TupleSections #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE TypeSynonymInstances #-} -{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE RecordWildCards #-} {-# OPTIONS_HADDOCK hide #-} -- | --- Module : Data.Array.Accelerate.Trafo.Simplify +-- Module : Data.Array.Accelerate.Trafo.Exp.Simplify -- Copyright : [2012..2020] The Accelerate Team -- License : BSD3 -- @@ -25,17 +10,14 @@ -- Portability : non-portable (GHC extensions) -- -module Data.Array.Accelerate.Trafo.Simplify ( - - simplifyAfun, - simplifyAcc, +module Data.Array.Accelerate.Trafo.Exp.Simplify ( simplifyFun, simplifyExp ) where -import Data.Array.Accelerate.AST +import Data.Array.Accelerate.AST.Exp import Data.Array.Accelerate.AST.Environment import Data.Array.Accelerate.AST.Idx import Data.Array.Accelerate.AST.LeftHandSide @@ -43,10 +25,10 @@ import Data.Array.Accelerate.AST.Var import Data.Array.Accelerate.Analysis.Hash import Data.Array.Accelerate.Analysis.Match import Data.Array.Accelerate.Error -import Data.Array.Accelerate.Representation.Array ( Array, ArrayR(..) ) -import Data.Array.Accelerate.Representation.Shape ( ShapeR(..), shapeToList ) -import Data.Array.Accelerate.Representation.Tag -import Data.Array.Accelerate.Trafo.Algebra +import Data.Array.Accelerate.Representation.Shape ( ShapeR(..), shapeToList, rank ) +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Representation.Vec +import Data.Array.Accelerate.Trafo.Exp.Algebra import Data.Array.Accelerate.Trafo.Environment import Data.Array.Accelerate.Trafo.Shrink import Data.Array.Accelerate.Trafo.Substitution @@ -57,63 +39,20 @@ import qualified Data.Array.Accelerate.Debug.Internal.Flags as Debug import qualified Data.Array.Accelerate.Debug.Internal.Trace as Debug import Control.Applicative hiding ( Const ) +import Control.Monad import Data.List ( partition ) import Data.Maybe +import Data.Text ( Text ) import Data.Monoid import Data.Text.Lazy.Builder +import Data.Primitive.Vec +import GHC.TypeNats import Formatting import Lens.Micro hiding ( ix ) import Prelude hiding ( exp, iterate ) import qualified Data.Map.Strict as Map --- Array simplification --- ==================== - -simplifyAcc :: Acc a -> Acc a -simplifyAcc = simplifyOpenAcc - -simplifyAfun :: Afun f -> Afun f -simplifyAfun = simplifyOpenAfun - - -simplifyOpenAcc :: OpenAcc aenv a -> OpenAcc aenv a -simplifyOpenAcc (OpenAcc pacc) = OpenAcc (simplifyPreOpenAcc pacc) - -simplifyOpenAfun :: PreOpenAfun OpenAcc aenv f -> PreOpenAfun OpenAcc aenv f -simplifyOpenAfun (Alam lhs f) = Alam lhs (simplifyOpenAfun f) -simplifyOpenAfun (Abody a) = Abody (simplifyOpenAcc a) - -simplifyPreOpenAcc :: PreOpenAcc OpenAcc aenv a -> PreOpenAcc OpenAcc aenv a -simplifyPreOpenAcc = \case - Alet lhs bnd body -> Alet lhs (simplifyOpenAcc bnd) (simplifyOpenAcc body) - Avar var -> Avar var - Apair a1 a2 -> Apair (simplifyOpenAcc a1) (simplifyOpenAcc a2) - Anil -> Anil - Atrace msg as bs -> Atrace msg (simplifyOpenAcc as) (simplifyOpenAcc bs) - Apply repr f a -> Apply repr (simplifyOpenAfun f) (simplifyOpenAcc a) - Aforeign repr asm f a -> Aforeign repr asm (simplifyOpenAfun f) (simplifyOpenAcc a) - Acond e a1 a2 -> Acond (simplifyExp e) (simplifyOpenAcc a1) (simplifyOpenAcc a2) - Awhile c f a -> Awhile (simplifyOpenAfun c) (simplifyOpenAfun f) (simplifyOpenAcc a) - Use repr arr -> Use repr arr - Unit eR e -> Unit eR (simplifyExp e) - Reshape shr e a -> Reshape shr (simplifyExp e) (simplifyOpenAcc a) - Generate repr e f -> Generate repr (simplifyExp e) (simplifyFun f) - Transform repr sh f g a -> Transform repr (simplifyExp sh) (simplifyFun f) (simplifyFun g) (simplifyOpenAcc a) - Replicate slix sl a -> Replicate slix (simplifyExp sl) (simplifyOpenAcc a) - Slice slix a sl -> Slice slix (simplifyOpenAcc a) (simplifyExp sl) - Map eR f a -> Map eR (simplifyFun f) (simplifyOpenAcc a) - ZipWith eR f a1 a2 -> ZipWith eR (simplifyFun f) (simplifyOpenAcc a1) (simplifyOpenAcc a2) - Fold f e a -> Fold (simplifyFun f) (fmap simplifyExp e) (simplifyOpenAcc a) - FoldSeg iR f e a s -> FoldSeg iR (simplifyFun f) (fmap simplifyExp e) (simplifyOpenAcc a) (simplifyOpenAcc s) - Scan d f e a -> Scan d (simplifyFun f) (fmap simplifyExp e) (simplifyOpenAcc a) - Scan' d f e a -> Scan' d (simplifyFun f) (simplifyExp e) (simplifyOpenAcc a) - Permute f a1 g a2 -> Permute (simplifyFun f) (simplifyOpenAcc a1) (simplifyFun g) (simplifyOpenAcc a2) - Backpermute shr sh f a -> Backpermute shr (simplifyExp sh) (simplifyFun f) (simplifyOpenAcc a) - Stencil s eR f b a -> Stencil s eR (simplifyFun f) b (simplifyOpenAcc a) - Stencil2 s1 s2 eR f b1 a1 b2 a2 -> Stencil2 s1 s2 eR (simplifyFun f) b1 (simplifyOpenAcc a1) b2 (simplifyOpenAcc a2) - - -- Scalar optimisations -- ==================== @@ -139,11 +78,11 @@ simplifyPreOpenAcc = \case -- memory because there are insufficient registers available. We sidestep this -- tricky and target-dependent issue by, for now, simply ignoring it. -- -localCSE :: (Kit acc, Elt a) - => Gamma acc env env aenv - -> OpenExp env aenv a - -> OpenExp (env,a) aenv b - -> Maybe (OpenExp env aenv b) +localCSE :: (Elt a) + => Gamma arr env env + -> PreOpenExp arr env a + -> PreOpenExp arr (env,a) b + -> Maybe (PreOpenExp arr env b) localCSE env bnd body | Just ix <- lookupExp env bnd = Stats.ruleFired "CSE" . Just $ inline body (Var ix) | otherwise = Nothing @@ -155,10 +94,10 @@ localCSE env bnd body -- -- > let x = e in .. e .. -- -globalCSE :: (Kit acc, Elt t) - => Gamma acc env env aenv - -> OpenExp env aenv t - -> Maybe (OpenExp env aenv t) +globalCSE :: (Elt t) + => Gamma arr env env + -> PreOpenExp arr env t + -> Maybe (PreOpenExp arr env t) globalCSE env exp | Just ix <- lookupExp env exp = Stats.ruleFired "CSE" . Just $ Var ix | otherwise = Nothing @@ -193,11 +132,11 @@ globalCSE env exp -- loops in its target language. -- recoverLoops - :: (Kit acc, Elt b) - => Gamma acc env env aenv - -> OpenExp env aenv a - -> OpenExp (env,a) aenv b - -> Maybe (OpenExp env aenv b) + :: (Elt b) + => Gamma arr env env + -> PreOpenExp arr env a + -> PreOpenExp arr (env,a) b + -> Maybe (PreOpenExp arr env b) recoverLoops _ bnd e3 -- To introduce scaler loops, we look for expressions of the form: -- @@ -232,15 +171,15 @@ recoverLoops _ bnd e3 = Nothing where - plus :: OpenExp env aenv Int -> OpenExp env aenv Int -> OpenExp env aenv Int + plus :: PreOpenExp arr env Int -> PreOpenExp arr env Int -> PreOpenExp arr env Int plus x y = PrimApp (PrimAdd numType) $ Tuple $ NilTup `SnocTup` x `SnocTup` y - constant :: Int -> OpenExp env aenv Int + constant :: Int -> PreOpenExp arr env Int constant i = Const ((),i) matchEnvTop :: (Elt s, Elt t) - => OpenExp (env,s) aenv f - -> OpenExp (env,t) aenv g + => PreOpenExp arr (env,s) f + -> PreOpenExp arr (env,t) g -> Maybe (s :=: t) matchEnvTop _ _ = gcast Refl --} @@ -259,86 +198,230 @@ recoverLoops _ bnd e3 -- If we do not want to do inlining, we should remove the environment here. -- simplifyOpenExp - :: forall env aenv e. - Gamma env env aenv - -> OpenExp env aenv e - -> (Bool, OpenExp env aenv e) + :: forall arr env e. + IsArrayInstr arr + => Gamma arr env env + -> PreOpenExp arr env e + -> (Bool, PreOpenExp arr env e) simplifyOpenExp env = first getAny . cvtE where - cvtE :: OpenExp env aenv t -> (Any, OpenExp env aenv t) + cvtE :: PreOpenExp arr env t -> (Any, PreOpenExp arr env t) cvtE exp = case exp of Let lhs bnd body -> (u <> v, exp') where (u, bnd') = cvtE bnd - (v, exp') = cvtLet env lhs bnd' (\env' -> cvtE' env' body) - Evar var -> pure $ Evar var + (v, exp') = cvtLet env lhs bnd' $ Subst weakenId body body + Evar var -> cvtVar var Const tp c -> pure $ Const tp c Undef tp -> pure $ Undef tp Nil -> pure Nil - Pair e1 e2 -> Pair <$> cvtE e1 <*> cvtE e2 - VecPack vec e -> VecPack vec <$> cvtE e - VecUnpack vec e -> VecUnpack vec <$> cvtE e - IndexSlice x ix sh -> IndexSlice x <$> cvtE ix <*> cvtE sh - IndexFull x ix sl -> IndexFull x <$> cvtE ix <*> cvtE sl - ToIndex shr sh ix -> toIndex shr (cvtE sh) (cvtE ix) - FromIndex shr sh ix -> fromIndex shr (cvtE sh) (cvtE ix) - Case e rhs def -> caseof (cvtE e) (sequenceA [ (t,) <$> cvtE c | (t,c) <- rhs ]) (cvtMaybeE def) - Cond p t e -> cond (cvtE p) (cvtE t) (cvtE e) - PrimConst c -> pure $ PrimConst c - PrimApp f x -> (u<>v, fx) - where - (u, x') = cvtE x - (v, fx) = evalPrimApp env f x' - Index a sh -> Index a <$> cvtE sh - LinearIndex a i -> LinearIndex a <$> cvtE i - Shape a -> shape a - ShapeSize shr sh -> shapeSize shr (cvtE sh) - Foreign tp ff f e -> Foreign tp ff <$> first Any (simplifyOpenFun EmptyExp f) <*> cvtE e + Pair e1 e2 -> hoist2 (\a b -> pure $ Pair a b) (cvtE e1) (cvtE e2) + VecPack vec e -> hoist (vecPack vec) (cvtE e) + VecUnpack vec e -> hoist (vecUnpack vec) (cvtE e) + ToIndex shr sh ix -> hoist2 (toIndex shr) (cvtE sh) (cvtE ix) + FromIndex shr sh ix -> hoist2 (fromIndex shr) (cvtE sh) (cvtE ix) + Case e rhs def -> hoist (\e' -> caseof e' (sequenceA [ (t,) <$> cvtE c | (t,c) <- rhs ]) (cvtMaybeE def)) (cvtE e) + Cond p t e -> hoist (\p' -> cond p' (cvtE t) (cvtE e)) (cvtE p) + Select p t e -> hoist (\p' -> select p' (cvtE t) (cvtE e)) (cvtE p) + PrimApp f x -> hoist (evalPrimApp env f) (cvtE x) + ArrayInstr arr e -> hoist (arrayInstr arr) (cvtE e) + ShapeSize shr sh -> hoist (shapeSize shr) (cvtE sh) + Foreign tp ff f e -> hoist (\e' -> Foreign tp ff <$> first Any (simplifyOpenFun EmptyExp f) <*> pure e') (cvtE e) While p f x -> While <$> cvtF env p <*> cvtF env f <*> cvtE x - Coerce t1 t2 e -> Coerce t1 t2 <$> cvtE e + Coerce t1 t2 e -> hoist (pure . Coerce t1 t2) (cvtE e) + Assert msg e1 e2 -> join (assert msg <$> cvtE e1 <*> cvtE e2) + Assume e1 e2 -> join (assume <$> cvtE e1 <*> cvtE e2) - cvtE' :: Gamma env' env' aenv -> OpenExp env' aenv e' -> (Any, OpenExp env' aenv e') + cvtE' :: Gamma arr env' env' -> PreOpenExp arr env' e' -> (Any, PreOpenExp arr env' e') cvtE' env' = first Any . simplifyOpenExp env' - cvtF :: Gamma env' env' aenv -> OpenFun env' aenv f -> (Any, OpenFun env' aenv f) + cvtF :: Gamma arr env' env' -> PreOpenFun arr env' f -> (Any, PreOpenFun arr env' f) cvtF env' = first Any . simplifyOpenFun env' - cvtMaybeE :: Maybe (OpenExp env aenv e') -> (Any, Maybe (OpenExp env aenv e')) + cvtMaybeE :: Maybe (PreOpenExp arr env e') -> (Any, Maybe (PreOpenExp arr env e')) cvtMaybeE Nothing = pure Nothing cvtMaybeE (Just e) = Just <$> cvtE e - cvtLet :: Gamma env' env' aenv + cvtLet :: Gamma arr env' env' -> ELeftHandSide bnd env' env'' - -> OpenExp env' aenv bnd - -> (Gamma env'' env'' aenv -> (Any, OpenExp env'' aenv t)) - -> (Any, OpenExp env' aenv t) - cvtLet env' lhs@(LeftHandSideSingle _) bnd body = Let lhs bnd <$> body (incExp $ env' `pushExp` bnd) -- Single variable on the LHS, add binding to the environment - cvtLet env' (LeftHandSideWildcard _) _ body = body env' -- Binding not used, remove let binding - cvtLet env' (LeftHandSidePair l1 l2) (Pair e1 e2) body -- Split binding to multiple bindings + -> PreOpenExp arr env' bnd -- Optimized + -> WeakOpenExp arr env'' t -- Not optimized + -> (Any, PreOpenExp arr env' t) + cvtLet env' lhs (Assert msg c bnd) body = yes $ Assert msg c $ snd $ cvtLet env' lhs bnd body + cvtLet env' lhs (Assume c bnd) body = yes $ Assume c $ snd $ cvtLet env' lhs bnd body + -- Let rotation + cvtLet env' lhs1 (Let lhs2 bnd expr) body + | Exists lhs1' <- rebuildLHS lhs1 = + yes $ snd $ cvtLet' env' lhs2 bnd $ \env'' -> + cvtLet env'' lhs1' expr $ weakenE (sinkWithLHS lhs1 lhs1' $ weakenWithLHS lhs2) body + cvtLet env' lhs bnd (Subst _ _ body) = cvtLet' env' lhs bnd $ \env'' -> cvtE' env'' body + + cvtLet' :: Gamma arr env' env' + -> ELeftHandSide bnd env' env'' + -> PreOpenExp arr env' bnd + -> (Gamma arr env'' env'' -> (Any, PreOpenExp arr env'' t)) + -> (Any, PreOpenExp arr env' t) + -- Let rotation and hoisting of assertions are already handled in cvtLet + cvtLet' env' lhs@(LeftHandSideSingle _) bnd body = Let lhs bnd <$> body (incExp $ env' `pushExp` bnd) -- Single variable on the LHS, add binding to the environment + cvtLet' env' (LeftHandSideWildcard _) _ body = body env' -- Binding not used, remove let binding + cvtLet' env' (LeftHandSidePair l1 l2) (Pair e1 e2) body -- Split binding to multiple bindings = first (const $ Any True) - $ cvtLet env' l1 e1 - $ \env'' -> cvtLet env'' l2 (weakenE (weakenWithLHS l1) e2) body - cvtLet env' lhs bnd body = Let lhs bnd <$> body (lhsExpr lhs env') -- Cannot split this binding. + $ cvtLet' env' l1 e1 + $ \env'' -> cvtLet' env'' l2 (weakenE (weakenWithLHS l1) e2) body + cvtLet' env' lhs bnd body = Let lhs bnd <$> body (lhsExpr lhs env') -- Cannot split this binding. + + cvtVar :: ExpVar env t -> (Any, PreOpenExp arr env t) + cvtVar var + | shouldInline bnd + , Nothing <- matchOpenExp bnd (Evar var) + = yes bnd + | otherwise + = pure $ Evar var + where + bnd = prjExp (varIdx var) env + + -- Note: we also do inlining in Shrink.hs, but we can easily do it here as + -- well since we already have an environment. + -- TODO: We may want to consider to remove the environment (Gamma) from + -- this function and let Shrink do the inlining. The environment was + -- annoying for let rotation for instance. + shouldInline :: PreOpenExp arr env t -> Bool + shouldInline Evar{} = True + shouldInline (ArrayInstr arr Nil) = inlineArrayInstr arr + shouldInline Const{} = True + shouldInline _ = False + + select :: PreOpenExp arr env PrimBool + -> (Any, PreOpenExp arr env t) + -> (Any, PreOpenExp arr env t) + -> (Any, PreOpenExp arr env t) + select p t@(_,t') e@(_,e') + | Const _ 1 <- p = Stats.knownBranch "True" (yes t') + | Const _ 0 <- p = Stats.knownBranch "False" (yes e') + | Undef _ <- e' = yes t' + | Undef _ <- t' = yes e' + -- Convert select over pairs to pair of selects. This may enable further + -- optimizations, and if we don't do this here we'll do it during code + -- generation anyway. + | Pair t1 t2 <- t' + , Pair e1 e2 <- e' + = if shouldInline p then + -- If the condition is simple, can directly perform this transformation + yes $ Pair (Select p t1 e1) (Select p t2 e2) + else + -- Otherwise we bind the condition to a variable + yes $ Let (LeftHandSideSingle scalarTypeWord8) p + $ Pair + (Select (Evar (Var scalarTypeWord8 ZeroIdx)) (weakenE (weakenSucc weakenId) t1) (weakenE (weakenSucc weakenId) e1)) + (Select (Evar (Var scalarTypeWord8 ZeroIdx)) (weakenE (weakenSucc weakenId) t2) (weakenE (weakenSucc weakenId) e2)) + | Just Refl <- matchOpenExp t' e' = Stats.knownBranch "redundant" (yes e') + | PrimApp PrimLNot c <- p = yes $ Select c e' t' + | otherwise = Select p <$> t <*> e -- Simplify conditional expressions, in particular by eliminating branches -- when the predicate is a known constant. -- - cond :: (Any, OpenExp env aenv PrimBool) - -> (Any, OpenExp env aenv t) - -> (Any, OpenExp env aenv t) - -> (Any, OpenExp env aenv t) - cond p@(_,p') t@(_,t') e@(_,e') - | Const _ 1 <- p' = Stats.knownBranch "True" (yes t') - | Const _ 0 <- p' = Stats.knownBranch "False" (yes e') + cond :: PreOpenExp arr env PrimBool + -> (Any, PreOpenExp arr env t) + -> (Any, PreOpenExp arr env t) + -> (Any, PreOpenExp arr env t) + cond p t@(_,t') e@(_,e') + | Const _ 1 <- p = Stats.knownBranch "True" (yes t') + | Const _ 0 <- p = Stats.knownBranch "False" (yes e') | Just Refl <- matchOpenExp t' e' = Stats.knownBranch "redundant" (yes e') - | otherwise = Cond <$> p <*> t <*> e - - caseof :: (Any, OpenExp env aenv TAG) - -> (Any, [(TAG, OpenExp env aenv b)]) - -> (Any, Maybe (OpenExp env aenv b)) - -> (Any, OpenExp env aenv b) - caseof tag@(_,tag') xs@(_,xs') md@(_,md') - | Const _ t <- tag' + | isCheap t' && isCheap e' = yes $ snd $ select p t e + | PrimApp PrimLNot c <- p = yes $ Cond c e' t' + | otherwise = Cond p <$> t <*> e + + -- Checks whether an expression is cheap, and may be evaluated eagerly + -- (i.e. lifted from within a conditional, to always be evaluated) + -- + -- If the expression performs many computations, or if the expression + -- may diverge, this function will return false. + -- + isCheap :: PreOpenExp arr env t -> Bool + isCheap = maybe False (<= maxCost) . expCost + where + maxCost = 10 + + expCost :: PreOpenExp arr env' t -> Maybe Int + expCost = \case + ArrayInstr a arg -> if inlineArrayInstr a + then Just 0 .+. expCost arg + else Nothing + Evar{} -> Just 0 + Nil -> Just 0 + Const{} -> Just 0 + Undef{} -> Just 0 + PrimApp f e -> primCost f .+. expCost e + Let _ bnd body -> expCost bnd .+. expCost body + Pair e1 e2 -> expCost e1 .+. expCost e2 + VecPack _ e -> Just 1 .+. expCost e + VecUnpack _ e -> Just 1 .+. expCost e + Coerce _ _ e -> expCost e + Select c t f -> Just 1 .+. expCost c .+. expCost t .+. expCost f + ToIndex shr sh ix -> shapeCost shr .+. expCost sh .+. expCost ix + FromIndex shr sh ix -> shapeCost shr .+. expCost sh .+. expCost ix + ShapeSize shr sh -> shapeCost shr .+. expCost sh + Foreign{} -> Nothing + Case{} -> Nothing + Cond{} -> Nothing + While{} -> Nothing + Assert{} -> Nothing + -- If we treat Assume as cheap, then it might be strictly evaluated. + -- That may imply that the assumption gets a larger scope, and we may + -- incorrectly act on the information of that assumption. + -- Example: if foo then assume foo x else y + -- If we treat Assume as cheap, this may be converted to: + -- select foo (assume foo x) y + -- The assumption will be lifted: + -- assume foo (select foo x y) + -- Bounds analysis may optimize this to: + -- assume foo x + -- Which is cleary not sound. + Assume{} -> Nothing + + shapeCost :: ShapeR sh -> Maybe Int + shapeCost = \case + ShapeRz -> Just 0 + ShapeRsnoc shr -> Just $ rank shr + + primCost :: PrimFun f -> Maybe Int + primCost = \case + PrimAdd _ -> Just 1 + PrimSub _ -> Just 1 + PrimMul _ -> Just 1 + PrimNeg _ -> Just 1 + PrimAbs _ -> Just 1 + PrimSig _ -> Just 1 + PrimBAnd _ -> Just 1 + PrimBOr _ -> Just 1 + PrimBXor _ -> Just 1 + PrimBNot _ -> Just 1 + PrimBShiftL _ -> Just 1 + PrimBShiftR _ -> Just 1 + PrimBRotateL _ -> Just 1 + PrimBRotateR _ -> Just 1 + PrimPopCount _ -> Just 1 + PrimCountLeadingZeros _ -> Just 1 + PrimCountTrailingZeros _ -> Just 1 + PrimCmp _ _ -> Just 1 + PrimMax _ -> Just 1 + PrimMin _ -> Just 1 + PrimLAnd -> Just 1 + PrimLOr -> Just 1 + PrimLNot -> Just 1 + _ -> Nothing + + (.+.) :: Maybe Int -> Maybe Int -> Maybe Int + a .+. b = (+) <$> a <*> b + + caseof :: PreOpenExp arr env TAG + -> (Any, [(TAG, PreOpenExp arr env b)]) + -> (Any, Maybe (PreOpenExp arr env b)) + -> (Any, PreOpenExp arr env b) + caseof tag xs@(_,xs') md@(_,md') + | Const _ t <- tag = Stats.caseElim "known" $ yes $ fromMaybe -- If the tag matches no alternative, then return the default (if present) -- or undef. @@ -350,16 +433,36 @@ simplifyOpenExp env = first getAny . cvtE | Just d <- md' , [(_,(_,u))] <- us , Just Refl <- matchOpenExp d u - = Stats.caseDefault "merge" $ yes (Case tag' (map snd vs) (Just u)) + = Stats.caseDefault "merge" $ yes $ Case tag (map snd vs) (Just u) | Nothing <- md' , [] <- vs , [(_,(_,u))] <- us - = Stats.caseElim "overlap" (yes u) + = Stats.caseElim "overlap" $ yes u | Nothing <- md' , [(_,(_,u))] <- us - = Stats.caseDefault "introduction" $ yes (Case tag' (map snd vs) (Just u)) + = Stats.caseDefault "introduction" $ yes $ Case tag (map snd vs) (Just u) + -- Convert case with only two alternatives to if-then-else, as we have + -- more optimizations for conditionals + | Nothing <- md' + , [(tag1, alt1), (tag2, alt2)] <- xs' + , expr <- snd $ cond (mkBinary (PrimCmp singleType CmpEq) tag $ Const scalarType tag1) (pure alt1) (pure alt2) + = if shouldInline tag then + -- Encode the information that 'tag' is either tag1 or tag2, in a way + -- that bounds analysis will understand: making them the lower and + -- upper bound of tag. + yes $ snd $ assume + (mkBinary PrimLAnd + (mkBinary (PrimCmp singleType CmpGtEq) tag (Const scalarType $ min tag1 tag2)) + (mkBinary (PrimCmp singleType CmpGtEq) (Const scalarType $ max tag1 tag2) tag) + ) + expr + else + yes expr + | Just d <- md' + , [(tag1, alt1)] <- xs' + = yes $ snd $ cond (mkBinary (PrimCmp singleType CmpEq) tag $ Const scalarType tag1) (pure alt1) (pure d) | otherwise - = Case <$> tag <*> xs <*> md + = Case tag <$> xs <*> md where tp | Just d <- md' = expType d @@ -370,49 +473,115 @@ simplifyOpenExp env = first getAny . cvtE . Map.fromListWith merge $ [ (hashOpenExp e, (1,(t, e))) | (t,e) <- xs' ] - merge :: (Int, (TAG, OpenExp env aenv b)) -> (Int, (TAG, OpenExp env aenv b)) -> (Int, (TAG, OpenExp env aenv b)) + merge :: (Int, (TAG, PreOpenExp arr env b)) -> (Int, (TAG, PreOpenExp arr env b)) -> (Int, (TAG, PreOpenExp arr env b)) merge (n,(_,a)) (m,(_,b)) = internalCheck "hashOpenExp/collision" (maybe False (const True) (matchOpenExp a b)) $ (n+m, (0xff, a)) + arrayInstr :: arr (s -> t) -> PreOpenExp arr env s -> (Any, PreOpenExp arr env t) + arrayInstr arr e = case arrayInstrType arr of + -- When an array instruction returns unit, we don't have to evaluate it. + -- This replaces the old rule which converted 'shape a' to Nil, if the array + -- had dimension zero. As we now generalized expressions over the type of + -- array instructions, we extend this rule to any array instructions returning + -- unit. Though it will likely only trigger for zero dimensional arrays. + TupRunit -> Stats.ruleFired "arrayInstr/nil" $ yes Nil + _ -> pure $ ArrayInstr arr e + -- Shape manipulations -- - shape :: ArrayVar aenv (Array sh t) -> (Any, OpenExp env aenv sh) - shape (Var (ArrayR ShapeRz _) _) - = Stats.ruleFired "shape/Z" $ yes Nil - shape a - = pure $ Shape a - - shapeSize :: ShapeR sh -> (Any, OpenExp env aenv sh) -> (Any, OpenExp env aenv Int) - shapeSize shr (_, sh) + shapeSize :: ShapeR sh -> PreOpenExp arr env sh -> (Any, PreOpenExp arr env Int) + shapeSize shr sh | Just c <- extractConstTuple sh = Stats.ruleFired "shapeSize/const" $ yes (Const scalarTypeInt (product (shapeToList shr c))) + shapeSize (ShapeRsnoc ShapeRz) (Pair _ sz) + = Stats.ruleFired "shapeSize/I1" $ yes sz shapeSize shr sh - = ShapeSize shr <$> sh + = pure $ ShapeSize shr sh toIndex :: ShapeR sh - -> (Any, OpenExp env aenv sh) - -> (Any, OpenExp env aenv sh) - -> (Any, OpenExp env aenv Int) - toIndex _ (_,sh) (_,FromIndex _ sh' ix) + -> PreOpenExp arr env sh + -> PreOpenExp arr env sh + -> (Any, PreOpenExp arr env Int) + toIndex _ sh (FromIndex _ sh' ix) | Just Refl <- matchOpenExp sh sh' = Stats.ruleFired "toIndex/fromIndex" $ yes ix - toIndex shr sh ix = ToIndex shr <$> sh <*> ix + toIndex ShapeRz _ _ = Stats.ruleFired "toIndex DIM0" $ yes $ Const scalarTypeInt 0 + toIndex (ShapeRsnoc ShapeRz) _ (Pair _ ix) + = Stats.ruleFired "toIndex DIM1" $ yes ix + toIndex (ShapeRsnoc ShapeRz) _ ix + = Stats.ruleFired "toIndex DIM1" $ yes + $ Let (LeftHandSidePair (LeftHandSideWildcard TupRunit) (LeftHandSideSingle scalarTypeInt)) + ix $ Evar $ Var scalarTypeInt ZeroIdx + toIndex shr sh ix = pure $ ToIndex shr sh ix fromIndex :: ShapeR sh - -> (Any, OpenExp env aenv sh) - -> (Any, OpenExp env aenv Int) - -> (Any, OpenExp env aenv sh) - fromIndex _ (_,sh) (_,ToIndex _ sh' ix) + -> PreOpenExp arr env sh + -> PreOpenExp arr env Int + -> (Any, PreOpenExp arr env sh) + fromIndex _ sh (ToIndex _ sh' ix) | Just Refl <- matchOpenExp sh sh' = Stats.ruleFired "fromIndex/toIndex" $ yes ix - fromIndex shr sh ix = FromIndex shr <$> sh <*> ix + fromIndex ShapeRz _ _ = Stats.ruleFired "fromIndex DIM0" $ yes Nil + fromIndex (ShapeRsnoc ShapeRz) _ ix + = Stats.ruleFired "fromIndex DIM1" $ yes $ Pair Nil ix + fromIndex shr sh ix = pure $ FromIndex shr sh ix + + vecPack :: KnownNat n => VecR n s tup -> PreOpenExp arr env tup -> (Any, PreOpenExp arr env (Vec n s)) + vecPack vecR (VecUnpack vecR' v) + | Just Refl <- matchVecR vecR vecR' = yes v + vecPack vecR e = pure $ VecPack vecR e + + vecUnpack :: KnownNat n => VecR n s tup -> PreOpenExp arr env (Vec n s) -> (Any, PreOpenExp arr env tup) + vecUnpack vecR (VecPack vecR' v) + | Just Refl <- matchVecR vecR vecR' = yes v + vecUnpack vecR e = pure $ VecUnpack vecR e + + assert :: Text -> PreOpenExp arr env PrimBool -> PreOpenExp arr env t -> (Any, PreOpenExp arr env t) + assert _ (Const _ 1) b = yes b + assert msg c@(Const _ 0) b = + let u = undefs $ expType b + in (Any (isNothing $ matchOpenExp b u), Assert msg c u) + assert msg1 (Assert msg2 c a) b = yes $ Assert msg2 c $ snd $ assert msg1 a b + assert msg (Assume c a) b = yes $ Assume c $ snd $ assert msg a b + assert msg c b = (Any False, Assert msg c b) + + assume :: PreOpenExp arr env PrimBool -> PreOpenExp arr env t -> (Any, PreOpenExp arr env t) + assume (Const _ 1) b = yes b + assume (Const _ 0) b = yes $ undefs $ expType b + assume (Assume c a) b = yes $ Assume c $ snd $ assume a b + assume c1 (Assume c2 b) = yes $ snd $ assume (PrimApp PrimLAnd (Pair c1 c2)) b + assume c b = (Any False, Assume c b) first :: (a -> a') -> (a,b) -> (a',b) first f (x,y) = (f x, y) - yes :: x -> (Any, x) - yes x = (Any True, x) - -extractConstTuple :: OpenExp env aenv t -> Maybe t +yes :: x -> (Any, x) +yes x = (Any True, x) + +hoist' + :: (PreOpenExp arr env a -> (Any, PreOpenExp arr env t)) + -> PreOpenExp arr env a + -> (Any, PreOpenExp arr env t) +hoist' f = \case + Assert msg c a -> yes $ Assert msg c $ snd $ hoist' f a + Assume c a -> yes $ Assume c $ snd $ hoist' f a + e -> f e + -- TODO: Should we also hoist let-bindings here? + -- That would make the types more difficult + +hoist + :: (PreOpenExp arr env a -> (Any, PreOpenExp arr env t)) + -> (Any, PreOpenExp arr env a) + -> (Any, PreOpenExp arr env t) +hoist f a = a >>= hoist' f + +hoist2 + :: (PreOpenExp arr env a -> PreOpenExp arr env b -> (Any, PreOpenExp arr env t)) + -> (Any, PreOpenExp arr env a) + -> (Any, PreOpenExp arr env b) + -> (Any, PreOpenExp arr env t) +hoist2 f a b = hoist (\a' -> hoist (f a') b) a + +extractConstTuple :: PreOpenExp arr env t -> Maybe t extractConstTuple Nil = Just () extractConstTuple (Pair e1 e2) = (,) <$> extractConstTuple e1 <*> extractConstTuple e2 extractConstTuple (Const _ c) = Just c @@ -421,15 +590,16 @@ extractConstTuple _ = Nothing -- Simplification for open functions -- simplifyOpenFun - :: Gamma env env aenv - -> OpenFun env aenv f - -> (Bool, OpenFun env aenv f) + :: IsArrayInstr arr + => Gamma arr env env + -> PreOpenFun arr env f + -> (Bool, PreOpenFun arr env f) simplifyOpenFun env (Body e) = Body <$> simplifyOpenExp env e simplifyOpenFun env (Lam lhs f) = Lam lhs <$> simplifyOpenFun env' f where env' = lhsExpr lhs env -lhsExpr :: ELeftHandSide t env env' -> Gamma env env aenv -> Gamma env' env' aenv +lhsExpr :: ELeftHandSide t env env' -> Gamma arr env env -> Gamma arr env' env' lhsExpr (LeftHandSideWildcard _) env = env lhsExpr (LeftHandSideSingle tp) env = incExp env `pushExp` Evar (Var tp ZeroIdx) lhsExpr (LeftHandSidePair l1 l2) env = lhsExpr l2 $ lhsExpr l1 env @@ -437,10 +607,10 @@ lhsExpr (LeftHandSidePair l1 l2) env = lhsExpr l2 $ lhsExpr l1 env -- Simplify closed expressions and functions. The process is applied -- repeatedly until no more changes are made. -- -simplifyExp :: HasCallStack => Exp aenv t -> Exp aenv t +simplifyExp :: (HasCallStack, IsArrayInstr arr) => PreOpenExp arr () t -> PreOpenExp arr () t simplifyExp = iterate summariseOpenExp matchOpenExp shrinkExp (simplifyOpenExp EmptyExp) -simplifyFun :: HasCallStack => Fun aenv f -> Fun aenv f +simplifyFun :: (HasCallStack, IsArrayInstr arr) => PreOpenFun arr () f -> PreOpenFun arr () f simplifyFun = iterate summariseOpenFun matchOpenFun shrinkFun (simplifyOpenFun EmptyExp) @@ -545,29 +715,24 @@ ops = lens _ops (\Stats{..} v -> Stats { _ops = v, ..}) {-# INLINE vars #-} {-# INLINE ops #-} -summariseOpenFun :: OpenFun env aenv f -> Stats +summariseOpenFun :: PreOpenFun arr env f -> Stats summariseOpenFun (Body e) = summariseOpenExp e & terms +~ 1 summariseOpenFun (Lam _ f) = summariseOpenFun f & terms +~ 1 & binders +~ 1 -summariseOpenExp :: OpenExp env aenv t -> Stats +summariseOpenExp :: PreOpenExp arr env t -> Stats summariseOpenExp = (terms +~ 1) . goE where zero = Stats 0 0 0 0 0 - travE :: OpenExp env aenv t -> Stats + travE :: PreOpenExp arr env t -> Stats travE = summariseOpenExp - travF :: OpenFun env aenv t -> Stats + travF :: PreOpenFun arr env t -> Stats travF = summariseOpenFun - travA :: acc aenv a -> Stats + travA :: arr t -> Stats travA _ = zero & vars +~ 1 -- assume an array index, else we should have failed elsewhere - travC :: PrimConst c -> Stats - travC (PrimMinBound t) = travBoundedType t & terms +~ 1 - travC (PrimMaxBound t) = travBoundedType t & terms +~ 1 - travC (PrimPi t) = travFloatingType t & terms +~ 1 - travIntegralType :: IntegralType t -> Stats travIntegralType _ = zero & types +~ 1 @@ -578,8 +743,10 @@ summariseOpenExp = (terms +~ 1) . goE travNumType (IntegralNumType t) = travIntegralType t & types +~ 1 travNumType (FloatingNumType t) = travFloatingType t & types +~ 1 + {- TODO WALL: DEAD CODE travBoundedType :: BoundedType t -> Stats travBoundedType (IntegralBoundedType t) = travIntegralType t & types +~ 1 + -} -- travScalarType :: ScalarType t -> Stats -- travScalarType (SingleScalarType t) = travSingleType t & types +~ 1 @@ -596,7 +763,7 @@ summariseOpenExp = (terms +~ 1) . goE -- travVectorType (Vector16Type t) = travSingleType t & types +~ 1 -- The scrutinee has already been counted - goE :: OpenExp env aenv t -> Stats + goE :: PreOpenExp arr env t -> Stats goE exp = case exp of Let _ bnd body -> travE bnd +++ travE body & binders +~ 1 @@ -608,20 +775,18 @@ summariseOpenExp = (terms +~ 1) . goE Pair e1 e2 -> travE e1 +++ travE e2 & terms +~ 1 VecPack _ e -> travE e VecUnpack _ e -> travE e - IndexSlice _ slix sh -> travE slix +++ travE sh & terms +~ 1 -- +1 for sliceIndex - IndexFull _ slix sl -> travE slix +++ travE sl & terms +~ 1 -- +1 for sliceIndex ToIndex _ sh ix -> travE sh +++ travE ix FromIndex _ sh ix -> travE sh +++ travE ix Case e rhs def -> travE e +++ mconcat [ travE c | (_,c) <- rhs ] +++ maybe zero travE def Cond p t e -> travE p +++ travE t +++ travE e + Select p t e -> travE p +++ travE t +++ travE e While p f x -> travF p +++ travF f +++ travE x - PrimConst c -> travC c - Index a ix -> travA a +++ travE ix - LinearIndex a ix -> travA a +++ travE ix - Shape a -> travA a + ArrayInstr a e -> travA a +++ travE e ShapeSize _ sh -> travE sh PrimApp f x -> travPrimFun f +++ travE x Coerce _ _ e -> travE e + Assert _ e1 e2 -> travE e1 +++ travE e2 + Assume e1 e2 -> travE e1 +++ travE e2 travPrimFun :: PrimFun f -> Stats travPrimFun = (ops +~ 1) . goF @@ -678,12 +843,7 @@ summariseOpenExp = (terms +~ 1) . goE PrimIsNaN t -> travFloatingType t PrimIsInfinite t -> travFloatingType t PrimAtan2 t -> travFloatingType t - PrimLt t -> travSingleType t - PrimGt t -> travSingleType t - PrimLtEq t -> travSingleType t - PrimGtEq t -> travSingleType t - PrimEq t -> travSingleType t - PrimNEq t -> travSingleType t + PrimCmp t _ -> travSingleType t PrimMax t -> travSingleType t PrimMin t -> travSingleType t PrimLAnd -> zero diff --git a/src/Data/Array/Accelerate/Trafo/Exp/Substitution.hs b/src/Data/Array/Accelerate/Trafo/Exp/Substitution.hs new file mode 100644 index 000000000..ce0cc070c --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Exp/Substitution.hs @@ -0,0 +1,564 @@ +{-# LANGUAGE ViewPatterns #-} +{-# OPTIONS_HADDOCK hide #-} +-- | +-- Module : Data.Array.Accelerate.Trafo.Exp.Substitution +-- Copyright : [2012..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.Trafo.Exp.Substitution ( + + -- ** Renaming & Substitution + inline, inlineVars, compose, + subTop, apply1, apply2, + + -- ** Weakening + (:>), SinkExp(..), weakenVars, + + -- ** Strengthening + (:?>), strengthenE, strengthenWithLHS, + + -- ** Rebuilding terms + RebuildableExp(..), rebuildLHS, + lhsFullVars, lhsVars, lhsIndices, lhsTake, + + RebuildArrayInstr, rebuildArrayInstrOpenExp, rebuildArrayInstrFun, + rebuildArrayInstrMap, + rebuildNoArrayInstr, mapArrayInstr, mapArrayInstrFun, + + returnExpVars, + + -- ** Checks + isIdentity, extractExpVars, + bindingIsTrivial, + +) where + +import Data.Array.Accelerate.AST.Exp +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.AST.Var +import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.IdxSet (IdxSet) +import qualified Data.Array.Accelerate.AST.IdxSet as IdxSet +import Data.Array.Accelerate.AST.Environment +import Data.Array.Accelerate.Analysis.Match +import Data.Array.Accelerate.Error +import Data.Array.Accelerate.Representation.Type +import qualified Data.Array.Accelerate.Debug.Internal.Stats as Stats + +import Data.Kind +import Data.Maybe +import Control.Applicative hiding ( Const ) +import Control.Monad +import Prelude hiding ( exp, seq ) +import Data.Functor.Identity + +infixr `compose` + + +lhsFullVars :: forall s a env1 env2. LeftHandSide s a env1 env2 -> Maybe (Vars s env2 a) +lhsFullVars = fmap snd . go weakenId + where + go :: forall env env' b. (env' :> env2) -> LeftHandSide s b env env' -> Maybe (env :> env2, Vars s env2 b) + go k (LeftHandSideWildcard TupRunit) = Just (k, TupRunit) + go k (LeftHandSideSingle s) = Just $ (weakenSucc $ k, TupRsingle $ Var s $ k >:> ZeroIdx) + go k (LeftHandSidePair l1 l2) + | Just (k', v2) <- go k l2 + , Just (k'', v1) <- go k' l1 = Just (k'', TupRpair v1 v2) + go _ _ = Nothing + +lhsVars :: forall s a env1 env2. LeftHandSide s a env1 env2 -> [Exists (Var s env2)] +lhsVars = snd . flip (go weakenId) [] + where + go :: env' :> env2 -> LeftHandSide s b env env' -> [Exists (Var s env2)] -> (env :> env2, [Exists (Var s env2)]) + go k (LeftHandSideSingle tp) accum = (weakenSucc k, Exists (Var tp $ k >:> ZeroIdx) : accum) + go k (LeftHandSideWildcard _) accum = (k, accum) + go k (LeftHandSidePair l1 l2) accum = go k' l1 accum' + where + (k', accum') = go k l2 accum + +lhsIndices :: forall s a env1 env2. LeftHandSide s a env1 env2 -> IdxSet env2 +lhsIndices = go IdxSet.empty + where + go :: IdxSet env -> LeftHandSide s b env env' -> IdxSet env' + go set (LeftHandSideSingle _) = IdxSet.push set + go set (LeftHandSideWildcard _) = set + go set (LeftHandSidePair l1 l2) = go (go set l1) l2 + +lhsTake :: LeftHandSide s t env1 env2 -> IdxSet env2 -> TupR u t -> [Exists u] +lhsTake = \lhs set values -> fst $ go [] lhs set values + where + go :: [Exists u] -> LeftHandSide s t env1 env2 -> IdxSet env2 -> TupR u t -> ([Exists u], IdxSet env1) + go accum (LeftHandSideSingle _) set (TupRsingle value) + | ZeroIdx `IdxSet.member` set = (Exists value : accum, IdxSet.drop set) + | otherwise = (accum, IdxSet.drop set) + go accum (LeftHandSidePair lhs1 lhs2) set (TupRpair values1 values2) + | (accum', set') <- go accum lhs2 set values2 + = go accum' lhs1 set' values1 + go accum (LeftHandSideWildcard _) set _ = (accum, set) + go _ _ _ _ = internalError "Tuple mismatch" + +bindingIsTrivial :: LeftHandSide s a env1 env2 -> Vars s env2 b -> Maybe (a :~: b) +bindingIsTrivial lhs vars + | Just lhsVars' <- lhsFullVars lhs + , Just Refl <- matchVars vars lhsVars' + = Just Refl +bindingIsTrivial _ _ = Nothing + +isIdentity :: PreOpenFun arr env (a -> b) -> Maybe (a :~: b) +isIdentity (Lam lhs (Body (extractExpVars -> Just vars))) = bindingIsTrivial lhs vars +isIdentity _ = Nothing + +-- | Replace the first variable with the given expression. The environment +-- shrinks. +-- +inline :: PreOpenExp arr (env, s) t + -> PreOpenExp arr env s + -> PreOpenExp arr env t +inline f g = Stats.substitution "inline" $ rebuildE (subTop g) f + +inlineVars :: forall arr env env' t1 t2. + ELeftHandSide t1 env env' + -> PreOpenExp arr env' t2 + -> PreOpenExp arr env t1 + -> Maybe (PreOpenExp arr env t2) +inlineVars lhsBound expr bound + | Just vars <- lhsFullVars lhsBound = substitute (strengthenWithLHS lhsBound) weakenId vars expr + where + substitute + :: forall env1 env2 t. + env1 :?> env2 + -> env :> env2 + -> ExpVars env1 t1 + -> PreOpenExp arr env1 t + -> Maybe (PreOpenExp arr env2 t) + substitute _ k2 vars (extractExpVars -> Just vars') + | Just Refl <- matchVars vars vars' = Just $ weakenE k2 bound + substitute k1 k2 vars topExp = case topExp of + Let lhs e1 e2 + | Exists lhs' <- rebuildLHS lhs + -> Let lhs' <$> travE e1 <*> substitute (strengthenAfter lhs lhs' k1) (weakenWithLHS lhs' .> k2) (weakenWithLHS lhs `weakenVars` vars) e2 + Evar (Var t ix) -> Evar . Var t <$> k1 ix + Foreign tp asm f e1 -> Foreign tp asm f <$> travE e1 + Pair e1 e2 -> Pair <$> travE e1 <*> travE e2 + Nil -> Just Nil + VecPack vec e1 -> VecPack vec <$> travE e1 + VecUnpack vec e1 -> VecUnpack vec <$> travE e1 + ToIndex shr e1 e2 -> ToIndex shr <$> travE e1 <*> travE e2 + FromIndex shr e1 e2 -> FromIndex shr <$> travE e1 <*> travE e2 + Case e1 rhs def -> Case <$> travE e1 <*> mapM (\(t,c) -> (t,) <$> travE c) rhs <*> travMaybeE def + Cond e1 e2 e3 -> Cond <$> travE e1 <*> travE e2 <*> travE e3 + Select e1 e2 e3 -> Select <$> travE e1 <*> travE e2 <*> travE e3 + While f1 f2 e1 -> While <$> travF f1 <*> travF f2 <*> travE e1 + Const t c -> Just $ Const t c + PrimApp p e1 -> PrimApp p <$> travE e1 + ArrayInstr arr e1 -> ArrayInstr arr <$> travE e1 + ShapeSize shr e1 -> ShapeSize shr <$> travE e1 + Undef t -> Just $ Undef t + Coerce t1 t2 e1 -> Coerce t1 t2 <$> travE e1 + Assert msg e1 e2 -> Assert msg <$> travE e1 <*> travE e2 + Assume e1 e2 -> Assume <$> travE e1 <*> travE e2 + + where + travE :: PreOpenExp arr env1 s -> Maybe (PreOpenExp arr env2 s) + travE = substitute k1 k2 vars + + travF :: PreOpenFun arr env1 s -> Maybe (PreOpenFun arr env2 s) + travF = substituteF k1 k2 vars + + travMaybeE :: Maybe (PreOpenExp arr env1 s) -> Maybe (Maybe (PreOpenExp arr env2 s)) + travMaybeE Nothing = pure Nothing + travMaybeE (Just x) = Just <$> travE x + + substituteF :: forall env1 env2 t. + env1 :?> env2 + -> env :> env2 + -> ExpVars env1 t1 + -> PreOpenFun arr env1 t + -> Maybe (PreOpenFun arr env2 t) + substituteF k1 k2 vars (Body e) = Body <$> substitute k1 k2 vars e + substituteF k1 k2 vars (Lam lhs f) + | Exists lhs' <- rebuildLHS lhs = Lam lhs' <$> substituteF (strengthenAfter lhs lhs' k1) (weakenWithLHS lhs' .> k2) (weakenWithLHS lhs `weakenVars` vars) f + +inlineVars _ _ _ = Nothing + + +-- | Composition of unary functions. +-- +compose :: HasCallStack + => PreOpenFun arr env (b -> c) + -> PreOpenFun arr env (a -> b) + -> PreOpenFun arr env (a -> c) +compose f@(Lam lhsB (Body c)) g@(Lam lhsA (Body b)) + | Stats.substitution "compose" False = undefined + | Just Refl <- isIdentity f = g -- don't rebind an identity function + | Just Refl <- isIdentity g = f + + | Exists lhsB' <- rebuildLHS lhsB + = Lam lhsA + $ Body + $ Let lhsB' b + $ weakenE (sinkWithLHS lhsB lhsB' $ weakenWithLHS lhsA) c + -- = Stats.substitution "compose" . Lam lhs2 . Body $ substitute' f g +compose _ + _ = internalError "compose: impossible evaluation" + +subTop :: PreOpenExp arr env s -> ExpVar (env, s) t -> PreOpenExp arr env t +subTop s (Var _ ZeroIdx ) = s +subTop _ (Var tp (SuccIdx ix)) = Evar $ Var tp ix + +-- Applies an argument to a unary function +-- +apply1 :: IsArrayInstr arr => TypeR t -> PreOpenFun arr () (s -> t) -> PreOpenExp arr env s -> PreOpenExp arr env t +apply1 tp fun arg = case weakenE (weakenEmpty) fun of + Body e -> functionImpossible $ expType e + Lam lhs (Body b) -> Let lhs arg b + Lam _ (Lam _ _) -> functionImpossible tp + +-- Applies arguments to a binary function +-- +apply2 :: IsArrayInstr arr => TypeR t -> PreOpenFun arr () (s -> u -> t) -> PreOpenExp arr env s -> PreOpenExp arr env u -> PreOpenExp arr env t +apply2 tp fun a1 a2 = case weakenE (weakenEmpty) fun of + Body e -> functionImpossible $ expType e + Lam _ (Body e) -> functionImpossible $ expType e + Lam l1 (Lam l2 (Body b)) -> Let (LeftHandSidePair l1 l2) (Pair a1 a2) b + Lam _ (Lam _ (Lam _ _)) -> functionImpossible tp + +-- A class for rebuilding scalar terms. +-- +class RebuildableExp f where + {-# MINIMAL rebuildPartialE, rebuildArrayInstrPartial #-} + rebuildPartialE :: (Applicative f', SyntacticExp fe) + => (forall e'. ExpVar env e' -> f' (fe arr env' e')) + -> f arr env e + -> f' (f arr env' e) + + {-# INLINEABLE rebuildE #-} + rebuildE :: SyntacticExp fe + => (forall e'. ExpVar env e' -> fe arr env' e') + -> f arr env e + -> f arr env' e + rebuildE v = runIdentity . rebuildPartialE (Identity . v) + + rebuildArrayInstrPartial :: Applicative f' + => RebuildArrayInstr f' arr arr' + -> f arr env e + -> f' (f arr' env e) + + {-# INLINEABLE rebuildArrayInstr #-} + rebuildArrayInstr :: (forall s t env'. arr (s -> t) -> PreOpenExp arr' env' s -> PreOpenExp arr' env' t) + -> f arr env e + -> f arr' env e + rebuildArrayInstr v = runIdentity . rebuildArrayInstrPartial (Identity . v) + +instance RebuildableExp PreOpenExp where + {-# INLINEABLE rebuildPartialE #-} + rebuildPartialE v x = Stats.substitution "rebuild" $ rebuildOpenExp v x + + {-# INLINEABLE rebuildArrayInstrPartial #-} + rebuildArrayInstrPartial v x = Stats.substitution "rebuildArrayInstr" $ rebuildArrayInstrOpenExp v x + +instance RebuildableExp PreOpenFun where + {-# INLINEABLE rebuildPartialE #-} + rebuildPartialE v x = Stats.substitution "rebuild" $ rebuildFun v x + + {-# INLINEABLE rebuildArrayInstrPartial #-} + rebuildArrayInstrPartial v x = Stats.substitution "rebuildArrayInstr" $ rebuildArrayInstrFun v x + +rebuildNoArrayInstr :: RebuildableExp f => f NoArrayInstr env e -> f arr env e +rebuildNoArrayInstr = rebuildArrayInstr k + where + k :: NoArrayInstr t -> a + k a = case a of {} + +-- NOTE: [Weakening] +-- +-- Weakening is something we usually take for granted: every time you learn a +-- new word, old sentences still make sense. If a conclusion is justified by a +-- hypothesis, it is still justified if you add more hypotheses. Similarly, a +-- term remains in scope if you bind more (fresh) variables. Weakening is the +-- operation of shifting things from one scope to a larger scope in which new +-- things have become meaningful, but no old things have vanished. +-- +-- When we use a named representation (or HOAS) we get weakening for free. But +-- in the de Bruijn representation weakening takes work: you have to shift all +-- variable references to make room for the new bindings. +-- + +-- TODO: Should we merge this type class with Sink? 'weaken' and 'weakenE' now have the same +-- type, because of various refactorings. +-- +class SinkExp (f :: Type -> Type -> Type) where + weakenE :: env :> env' -> f env t -> f env' t + + -- TLM: We can't use this default instance because it doesn't lead to + -- specialised code. Perhaps the INLINEABLE pragma is ignored: GHC bug? + -- + -- {-# INLINEABLE weakenE #-} + -- default weakenE :: RebuildableExp f => env :> env' -> f arr env t -> f arr env' t + -- weakenE v = Stats.substitution "weakenE" . rebuildE (IE . v) + +instance SinkExp Idx where + {-# INLINEABLE weakenE #-} + weakenE w = (>:>) w + +instance SinkExp (Var s) where + {-# INLINEABLE weakenE #-} + weakenE k (Var s ix) = Var s (k >:> ix) + +instance SinkExp (PreOpenExp arr) where + {-# INLINEABLE weakenE #-} + weakenE v = Stats.substitution "weakenE" . rebuildE (rebuildWeakenEvar v) + +instance SinkExp (PreOpenFun arr) where + {-# INLINEABLE weakenE #-} + weakenE v = Stats.substitution "weakenE" . rebuildE (rebuildWeakenEvar v) + + +-- This rewrite rule is disabled because 'weaken' is now part of a type class. +-- As such, we cannot attach a NOINLINE pragma because it has many definitions. +-- {-# RULES +-- "weakenE/weakenE" forall a (v1 :: env' :> env'') (v2 :: env :> env'). +-- weakenE v1 (weakenE v2 a) = weakenE (v1 . v2) a +-- #-} + +weakenVars :: env :> env' -> Vars s env t -> Vars s env' t +weakenVars _ TupRunit = TupRunit +weakenVars k (TupRsingle v) = TupRsingle $ weakenE k v +weakenVars k (TupRpair v w) = TupRpair (weakenVars k v) (weakenVars k w) + +rebuildWeakenEvar :: env :> env' -> ExpVar env t -> PreOpenExp arr env' t +rebuildWeakenEvar k (Var s idx) = Evar $ Var s $ k >:> idx + +-- NOTE: [Strengthening] +-- +-- Strengthening is the dual of weakening. Shifting terms from one scope to a +-- smaller scope. Of course this is not always possible. If the term contains +-- any variables not in the new environment, then it cannot be strengthened. +-- This partial behaviour is captured with 'Maybe'. +-- + +-- The type of partially shifting terms from one context into another. +type env :?> env' = forall t'. Idx env t' -> Maybe (Idx env' t') + +{-# INLINEABLE strengthenE #-} +strengthenE :: forall f arr env env' t. RebuildableExp f => env :?> env' -> f arr env t -> Maybe (f arr env' t) +strengthenE k x = Stats.substitution "strengthenE" $ rebuildPartialE @f @Maybe @IdxE (\(Var tp ix) -> fmap (IE . Var tp) $ k ix) x + +strengthenWithLHS :: LeftHandSide s t env1 env2 -> env2 :?> env1 +strengthenWithLHS (LeftHandSideWildcard _) = Just +strengthenWithLHS (LeftHandSideSingle _) = \ix -> case ix of + ZeroIdx -> Nothing + SuccIdx i -> Just i +strengthenWithLHS (LeftHandSidePair l1 l2) = strengthenWithLHS l2 >=> strengthenWithLHS l1 + +strengthenAfter :: LeftHandSide s t env1 env2 -> LeftHandSide s t env1' env2' -> env1 :?> env1' -> env2 :?> env2' +strengthenAfter (LeftHandSideWildcard _) (LeftHandSideWildcard _) k = k +strengthenAfter (LeftHandSideSingle _) (LeftHandSideSingle _) k = \ix -> case ix of + ZeroIdx -> Just ZeroIdx + SuccIdx i -> SuccIdx <$> k i +strengthenAfter (LeftHandSidePair l1 l2) (LeftHandSidePair l1' l2') k = + strengthenAfter l2 l2' $ strengthenAfter l1 l1' k +strengthenAfter _ _ _ = error "Substitution.strengthenAfter: left hand sides do not match" + + +-- Simultaneous Substitution =================================================== +-- + +-- The scalar environment +-- ------------------ + +-- SEE: [Renaming and Substitution] +-- SEE: [Weakening] +-- +class SyntacticExp f where + varIn :: ExpVar env t -> f arr env t + expOut :: f arr env t -> PreOpenExp arr env t + weakenExp :: f arr env t -> f arr (env, s) t + +newtype IdxE (arr :: Type -> Type) env t = IE { unIE :: ExpVar env t } + +instance SyntacticExp IdxE where + varIn = IE + expOut = Evar . unIE + weakenExp (IE (Var tp ix)) = IE $ Var tp $ SuccIdx ix + +instance SyntacticExp PreOpenExp where + varIn = Evar + expOut = id + weakenExp = runIdentity . rebuildOpenExp (Identity . weakenExp . IE) + +{-# INLINEABLE shiftE #-} +shiftE + :: (Applicative f, SyntacticExp fe) + => RebuildEvar f fe arr env env' + -> RebuildEvar f fe arr (env, s) (env', s) +shiftE _ (Var tp ZeroIdx) = pure $ varIn (Var tp ZeroIdx) +shiftE v (Var tp (SuccIdx ix)) = weakenExp <$> v (Var tp ix) + +{-# INLINEABLE shiftE' #-} +shiftE' + :: (Applicative f, SyntacticExp fe) + => ELeftHandSide t env1 env1' + -> ELeftHandSide t env2 env2' + -> RebuildEvar f fe arr env1 env2 + -> RebuildEvar f fe arr env1' env2' +shiftE' (LeftHandSideWildcard _) (LeftHandSideWildcard _) v = v +shiftE' (LeftHandSideSingle _) (LeftHandSideSingle _) v = shiftE v +shiftE' (LeftHandSidePair a1 b1) (LeftHandSidePair a2 b2) v = shiftE' b1 b2 $ shiftE' a1 a2 v +shiftE' _ _ _ = error "Substitution: left hand sides do not match" + +type RebuildEvar f fe (arr :: Type -> Type) env env' = + forall t'. ExpVar env t' -> f (fe arr env' t') + +{-# INLINEABLE rebuildMaybeExp #-} +rebuildMaybeExp + :: (HasCallStack, Applicative f, SyntacticExp fe) + => RebuildEvar f fe arr env env' + -> Maybe (PreOpenExp arr env t) + -> f (Maybe (PreOpenExp arr env' t)) +rebuildMaybeExp _ Nothing = pure Nothing +rebuildMaybeExp v (Just x) = Just <$> rebuildOpenExp v x + +{-# INLINEABLE rebuildOpenExp #-} +rebuildOpenExp + :: (HasCallStack, Applicative f, SyntacticExp fe) + => RebuildEvar f fe arr env env' + -> PreOpenExp arr env t + -> f (PreOpenExp arr env' t) +rebuildOpenExp v exp = + case exp of + Const t c -> pure $ Const t c + Undef t -> pure $ Undef t + Evar var -> expOut <$> v var + Let lhs a b + | Exists lhs' <- rebuildLHS lhs + -> Let lhs' <$> rebuildOpenExp v a <*> rebuildOpenExp (shiftE' lhs lhs' v) b + Pair e1 e2 -> Pair <$> rebuildOpenExp v e1 <*> rebuildOpenExp v e2 + Nil -> pure Nil + VecPack vec e -> VecPack vec <$> rebuildOpenExp v e + VecUnpack vec e -> VecUnpack vec <$> rebuildOpenExp v e + ToIndex shr sh ix -> ToIndex shr <$> rebuildOpenExp v sh <*> rebuildOpenExp v ix + FromIndex shr sh ix -> FromIndex shr <$> rebuildOpenExp v sh <*> rebuildOpenExp v ix + Case e rhs def -> Case <$> rebuildOpenExp v e <*> sequenceA [ (t,) <$> rebuildOpenExp v c | (t,c) <- rhs ] <*> rebuildMaybeExp v def + Cond p t e -> Cond <$> rebuildOpenExp v p <*> rebuildOpenExp v t <*> rebuildOpenExp v e + Select p t e -> Select <$> rebuildOpenExp v p <*> rebuildOpenExp v t <*> rebuildOpenExp v e + While p f x -> While <$> rebuildFun v p <*> rebuildFun v f <*> rebuildOpenExp v x + PrimApp f x -> PrimApp f <$> rebuildOpenExp v x + ArrayInstr arr e -> ArrayInstr arr <$> rebuildOpenExp v e + ShapeSize shr sh -> ShapeSize shr <$> rebuildOpenExp v sh + Foreign tp ff f e -> Foreign tp ff f <$> rebuildOpenExp v e + Coerce t1 t2 e -> Coerce t1 t2 <$> rebuildOpenExp v e + Assert msg e1 e2 -> Assert msg <$> rebuildOpenExp v e1 <*> rebuildOpenExp v e2 + Assume e1 e2 -> Assume <$> rebuildOpenExp v e1 <*> rebuildOpenExp v e2 + +{-# INLINEABLE rebuildFun #-} +rebuildFun + :: (HasCallStack, Applicative f, SyntacticExp fe) + => RebuildEvar f fe arr env env' + -> PreOpenFun arr env t + -> f (PreOpenFun arr env' t) +rebuildFun v fun = + case fun of + Body e -> Body <$> rebuildOpenExp v e + Lam lhs f + | Exists lhs' <- rebuildLHS lhs + -> Lam lhs' <$> rebuildFun (shiftE' lhs lhs' v) f + +{-# INLINEABLE rebuildLHS #-} +rebuildLHS :: LeftHandSide s t aenv1 aenv1' -> Exists (LeftHandSide s t aenv2) +rebuildLHS (LeftHandSideWildcard r) = Exists $ LeftHandSideWildcard r +rebuildLHS (LeftHandSideSingle s) = Exists $ LeftHandSideSingle s +rebuildLHS (LeftHandSidePair as bs) + | Exists as' <- rebuildLHS as + , Exists bs' <- rebuildLHS bs + = Exists $ LeftHandSidePair as' bs' + +type RebuildArrayInstr f arr arr' = forall s t env. arr (s -> t) -> f (PreOpenExp arr' env s -> PreOpenExp arr' env t) + +-- RebuildArrayInstr corresponds to the bind of a monad, meaning that a single array instruction can +-- be transformed to an expressions, possibly containing multiple array instructions. For many cases +-- it may be sufficient to only have the functor (map) interface. This function converts a map-like +-- function to the bind-like function. +-- +rebuildArrayInstrMap + :: Applicative f + => (forall s t. arr (s -> t) -> f (arr' (s -> t))) + -> RebuildArrayInstr f arr arr' +rebuildArrayInstrMap f arr = ArrayInstr <$> f arr + +mapArrayInstr + :: (forall s t. arr (s -> t) -> arr' (s -> t)) + -> PreOpenExp arr env e + -> PreOpenExp arr' env e +mapArrayInstr f = runIdentity . rebuildArrayInstrOpenExp (rebuildArrayInstrMap (Identity . f)) + +mapArrayInstrFun + :: (forall s t. arr (s -> t) -> arr' (s -> t)) + -> PreOpenFun arr env e + -> PreOpenFun arr' env e +mapArrayInstrFun f = runIdentity . rebuildArrayInstrFun (rebuildArrayInstrMap (Identity . f)) + +rebuildArrayInstrOpenExp + :: forall f arr arr' env t. + (HasCallStack, Applicative f) + => RebuildArrayInstr f arr arr' + -> PreOpenExp arr env t + -> f (PreOpenExp arr' env t) +rebuildArrayInstrOpenExp v = \case + Let lhs e1 e2 -> Let lhs <$> travE e1 <*> travE e2 + Evar var -> pure $ Evar var + Foreign t asm f x -> Foreign t asm f <$> travE x + Pair e1 e2 -> Pair <$> travE e1 <*> travE e2 + Nil -> pure Nil + VecPack vecr e -> VecPack vecr <$> travE e + VecUnpack vecr e -> VecUnpack vecr <$> travE e + ToIndex shr sh ix -> ToIndex shr <$> travE sh <*> travE ix + FromIndex shr sh ix -> FromIndex shr <$> travE sh <*> travE ix + Case e rhs def -> Case <$> travE e <*> sequenceA [ (t,) <$> travE c | (t,c) <- rhs ] <*> travME def + Cond e1 e2 e3 -> Cond <$> travE e1 <*> travE e2 <*> travE e3 + Select e1 e2 e3 -> Select <$> travE e1 <*> travE e2 <*> travE e3 + While c f x -> While <$> travF c <*> travF f <*> travE x + Const tp c -> pure $ Const tp c + PrimApp f x -> PrimApp f <$> travE x + ArrayInstr arr x -> v arr <*> travE x + ShapeSize shr sh -> ShapeSize shr <$> travE sh + Undef tp -> pure $ Undef tp + Coerce t1 t2 e -> Coerce t1 t2 <$> travE e + Assert msg e1 e2 -> Assert msg <$> travE e1 <*> travE e2 + Assume e1 e2 -> Assume <$> travE e1 <*> travE e2 + + where + travE :: PreOpenExp arr env' t' -> f (PreOpenExp arr' env' t') + travE = rebuildArrayInstrOpenExp v + + travF :: PreOpenFun arr env' t' -> f (PreOpenFun arr' env' t') + travF = rebuildArrayInstrFun v + + travME :: Maybe (PreOpenExp arr env' t') -> f (Maybe (PreOpenExp arr' env' t')) + travME Nothing = pure Nothing + travME (Just e) = Just <$> travE e + +rebuildArrayInstrFun + :: (HasCallStack, Applicative f) + => RebuildArrayInstr f arr arr' + -> PreOpenFun arr env t + -> f (PreOpenFun arr' env t) +rebuildArrayInstrFun v (Body e) = Body <$> rebuildArrayInstrOpenExp v e +rebuildArrayInstrFun v (Lam lhs f) = Lam lhs <$> rebuildArrayInstrFun v f + +extractExpVars :: PreOpenExp arr env a -> Maybe (ExpVars env a) +extractExpVars (Assume _ e) = extractExpVars e +extractExpVars Nil = Just TupRunit +extractExpVars (Pair e1 e2) = TupRpair <$> extractExpVars e1 <*> extractExpVars e2 +extractExpVars (Evar v) = Just $ TupRsingle v +extractExpVars _ = Nothing + +returnExpVars :: ExpVars env a -> PreOpenExp arr env a +returnExpVars TupRunit = Nil +returnExpVars (TupRpair v1 v2) = returnExpVars v1 `Pair` returnExpVars v2 +returnExpVars (TupRsingle v) = Evar v diff --git a/src/Data/Array/Accelerate/Trafo/Fusion.hs b/src/Data/Array/Accelerate/Trafo/Fusion.hs index 81487be09..dd2b2aa54 100644 --- a/src/Data/Array/Accelerate/Trafo/Fusion.hs +++ b/src/Data/Array/Accelerate/Trafo/Fusion.hs @@ -1,1770 +1,1761 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE InstanceSigs #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE PatternGuards #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE ViewPatterns #-} -{-# OPTIONS_GHC -fno-warn-name-shadowing #-} -{-# OPTIONS_HADDOCK hide #-} --- | --- Module : Data.Array.Accelerate.Trafo.Fusion --- Copyright : [2012..2020] The Accelerate Team --- License : BSD3 --- --- Maintainer : Trevor L. McDonell --- Stability : experimental --- Portability : non-portable (GHC extensions) --- --- This module implements producer/producer and consumer/producer fusion as a --- term rewriting of the Accelerate AST. --- --- The function 'quench' perform the source-to-source fusion transformation, --- while 'anneal' additionally makes the representation of embedded producers --- explicit by representing the AST as a 'DelayedAcc' of manifest and delayed --- nodes. --- - -module Data.Array.Accelerate.Trafo.Fusion ( - - convertAcc, convertAccWith, - convertAfun, convertAfunWith, - -) where - -import Data.BitSet -import Data.Array.Accelerate.AST -import Data.Array.Accelerate.AST.LeftHandSide -import Data.Array.Accelerate.AST.Environment -import Data.Array.Accelerate.AST.Var -import Data.Array.Accelerate.AST.Idx -import Data.Array.Accelerate.Error -import Data.Array.Accelerate.Trafo.Config -import Data.Array.Accelerate.Trafo.Var -import Data.Array.Accelerate.Trafo.Delayed -import Data.Array.Accelerate.Trafo.Environment -import Data.Array.Accelerate.Trafo.Shrink -import Data.Array.Accelerate.Trafo.Simplify -import Data.Array.Accelerate.Trafo.Substitution -import Data.Array.Accelerate.Representation.Array ( Array, ArrayR(..), ArraysR ) -import Data.Array.Accelerate.Representation.Shape ( ShapeR(..), shapeType ) -import Data.Array.Accelerate.Representation.Slice -import Data.Array.Accelerate.Representation.Type -import Data.Array.Accelerate.Analysis.Match -import Data.Array.Accelerate.Type - -import Data.Array.Accelerate.Debug.Internal.Flags ( array_fusion ) -import qualified Data.Array.Accelerate.Debug.Internal.Stats as Stats -#ifdef ACCELERATE_DEBUG -import System.IO.Unsafe -- for debugging -#endif - -import Data.Function -import Lens.Micro ( over, mapped, _2 ) -import Prelude hiding ( exp, until ) - - --- Delayed Array Fusion --- ==================== - --- | Apply the fusion transformation to a closed de Bruijn AST --- -convertAcc :: HasCallStack => Acc arrs -> DelayedAcc arrs -convertAcc = convertAccWith defaultOptions - -convertAccWith :: HasCallStack => Config -> Acc arrs -> DelayedAcc arrs -convertAccWith config = withSimplStats . convertOpenAcc config - --- | Apply the fusion transformation to a function of array arguments --- -convertAfun :: HasCallStack => Afun f -> DelayedAfun f -convertAfun = convertAfunWith defaultOptions - -convertAfunWith :: HasCallStack => Config -> Afun f -> DelayedAfun f -convertAfunWith config = withSimplStats . convertOpenAfun config - --- -- | Apply the fusion transformation to the array computations embedded --- -- in a sequence computation. --- --- convertSeq :: Bool -> Seq a -> DelayedSeq a --- convertSeq fuseAcc (embedSeq (embedOpenAcc fuseAcc) -> ExtendSeq aenv s) --- = withSimplStats (DelayedSeq (cvtE aenv) (convertOpenSeq fuseAcc s)) +-- {-# LANGUAGE CPP #-} +-- {-# LANGUAGE ViewPatterns #-} +-- {-# OPTIONS_GHC -Wno-name-shadowing #-} +-- {-# OPTIONS_HADDOCK hide #-} +-- -- | +-- -- Module : Data.Array.Accelerate.Trafo.Fusion +-- -- Copyright : [2012..2020] The Accelerate Team +-- -- License : BSD3 +-- -- +-- -- Maintainer : Trevor L. McDonell +-- -- Stability : experimental +-- -- Portability : non-portable (GHC extensions) +-- -- +-- -- This module implements producer/producer and consumer/producer fusion as a +-- -- term rewriting of the Accelerate AST. +-- -- +-- -- The function 'quench' perform the source-to-source fusion transformation, +-- -- while 'anneal' additionally makes the representation of embedded producers +-- -- explicit by representing the AST as a 'DelayedAcc' of manifest and delayed +-- -- nodes. +-- -- + +-- module Data.Array.Accelerate.Trafo.Fusion ( + +-- convertAcc, convertAccWith, +-- convertAfun, convertAfunWith, + +-- ) where + +-- import Data.BitSet +-- import Data.Array.Accelerate.AST +-- import Data.Array.Accelerate.AST.LeftHandSide +-- import Data.Array.Accelerate.AST.Environment +-- import Data.Array.Accelerate.AST.Var +-- import Data.Array.Accelerate.AST.Idx +-- import Data.Array.Accelerate.Error +-- import Data.Array.Accelerate.Trafo.Config +-- import Data.Array.Accelerate.Trafo.Var +-- import Data.Array.Accelerate.Trafo.Delayed +-- import Data.Array.Accelerate.Trafo.Environment +-- import Data.Array.Accelerate.Trafo.Shrink +-- import Data.Array.Accelerate.Trafo.Exp.Simplify +-- import Data.Array.Accelerate.Trafo.Substitution +-- import Data.Array.Accelerate.Representation.Array ( Array, ArrayR(..), ArraysR ) +-- import Data.Array.Accelerate.Representation.Shape ( ShapeR(..), shapeType ) +-- import Data.Array.Accelerate.Representation.Slice +-- import Data.Array.Accelerate.Representation.Type +-- import Data.Array.Accelerate.Analysis.Match +-- import Data.Array.Accelerate.Type + +-- import Data.Array.Accelerate.Debug.Internal.Flags ( array_fusion ) +-- import qualified Data.Array.Accelerate.Debug.Internal.Stats as Stats +-- #ifdef ACCELERATE_DEBUG +-- import System.IO.Unsafe -- for debugging +-- #endif + +-- import Data.Function +-- import Lens.Micro ( over, mapped, _2 ) +-- import Prelude hiding ( exp, until ) + + +-- -- Delayed Array Fusion +-- -- ==================== + +-- -- | Apply the fusion transformation to a closed de Bruijn AST +-- -- +-- convertAcc :: HasCallStack => Acc arrs -> DelayedAcc arrs +-- convertAcc = convertAccWith defaultOptions + +-- convertAccWith :: HasCallStack => Config -> Acc arrs -> DelayedAcc arrs +-- convertAccWith config = withSimplStats . convertOpenAcc config + +-- -- | Apply the fusion transformation to a function of array arguments +-- -- +-- convertAfun :: HasCallStack => Afun f -> DelayedAfun f +-- convertAfun = convertAfunWith defaultOptions + +-- convertAfunWith :: HasCallStack => Config -> Afun f -> DelayedAfun f +-- convertAfunWith config = withSimplStats . convertOpenAfun config + +-- -- -- | Apply the fusion transformation to the array computations embedded +-- -- -- in a sequence computation. +-- -- +-- -- convertSeq :: Bool -> Seq a -> DelayedSeq a +-- -- convertSeq fuseAcc (embedSeq (embedOpenAcc fuseAcc) -> ExtendSeq aenv s) +-- -- = withSimplStats (DelayedSeq (cvtE aenv) (convertOpenSeq fuseAcc s)) +-- -- where +-- -- cvtE :: Extend OpenAcc aenv aenv' -> Extend DelayedOpenAcc aenv aenv' +-- -- cvtE BaseEnv = BaseEnv +-- -- cvtE (PushEnv env a) | a' <- convertOpenAcc fuseAcc a = PushEnv (cvtE env) a' + +-- withSimplStats :: a -> a +-- #ifdef ACCELERATE_DEBUG +-- withSimplStats x = unsafePerformIO Stats.resetSimplCount `seq` x +-- #else +-- withSimplStats x = x +-- #endif + + +-- -- | Apply the fusion transformation to an AST. This consists of two phases: +-- -- +-- -- 1. A bottom-up traversal that converts nodes into the internal delayed +-- -- representation, merging adjacent producer/producer pairs. +-- -- +-- -- 2. A top-down traversal that makes the representation of fused +-- -- consumer/producer pairs explicit as a 'DelayedAcc' of manifest and +-- -- delayed nodes. +-- -- +-- -- TLM: Note that there really is no ambiguity as to which state an array will +-- -- be in following this process: an array will be either delayed or +-- -- manifest, and the two helper functions are even named as such! We should +-- -- encode this property in the type somehow... +-- -- +-- convertOpenAcc +-- :: HasCallStack +-- => Config +-- -> OpenAcc aenv arrs +-- -> DelayedOpenAcc aenv arrs +-- convertOpenAcc config = manifest config . computeAcc . embedOpenAcc config + + +-- -- Convert array computations into an embeddable delayed representation. +-- -- Reapply the embedding function from the first pass and unpack the +-- -- representation. It is safe to match on BaseEnv because the first pass +-- -- will put producers adjacent to the term consuming it. +-- -- +-- delayed +-- :: HasCallStack +-- => Config +-- -> OpenAcc aenv (Array sh e) +-- -> DelayedOpenAcc aenv (Array sh e) +-- delayed config (embedOpenAcc config -> Embed env cc) +-- | BaseEnv <- env +-- = simplify cc & \case +-- Left (Done v) -> avarsIn Manifest v +-- Right d -> d & \case +-- Yield aR sh f -> Delayed aR sh f (f `compose` fromIndex (arrayRshape aR) sh) +-- Step aR sh p f v +-- | Just Refl <- matchOpenExp sh (arrayShape v) +-- , Just Refl <- isIdentity p -> Delayed aR sh (f `compose` indexArray v) (f `compose` linearIndex v) +-- | f' <- f `compose` indexArray v `compose` p -> Delayed aR sh f' (f' `compose` fromIndex (arrayRshape aR) sh) +-- -- +-- | otherwise +-- = manifest config (computeAcc (Embed env cc)) + + +-- -- Convert array programs as manifest terms. +-- -- +-- manifest +-- :: HasCallStack +-- => Config +-- -> OpenAcc aenv a +-- -> DelayedOpenAcc aenv a +-- manifest config (OpenAcc pacc) = +-- let fusionError = internalError "unexpected fusible materials" +-- in +-- Manifest $ case pacc of +-- -- Non-fusible terms +-- -- ----------------- +-- Avar ix -> Avar ix +-- Use aR a -> Use aR a +-- Unit t e -> Unit t e +-- Alet lhs bnd body -> alet lhs (manifest config bnd) (manifest config body) +-- Acond p t e -> Acond p (manifest config t) (manifest config e) +-- Awhile p f a -> Awhile (cvtAF p) (cvtAF f) (manifest config a) +-- Apair a1 a2 -> Apair (manifest config a1) (manifest config a2) +-- Anil -> Anil +-- Atrace msg a1 a2 -> Atrace msg (manifest config a1) (manifest config a2) +-- Apply repr f a -> apply repr (cvtAF f) (manifest config a) +-- Aforeign repr ff f a -> Aforeign repr ff (cvtAF f) (manifest config a) + +-- -- Producers +-- -- --------- +-- -- +-- -- Some producers might still exist as a manifest array. Typically this +-- -- is because they are the last stage of the computation, or the result +-- -- of a let-binding to be used multiple times. The input array here +-- -- should be a evaluated array term, else something went wrong. +-- -- +-- Map t f a -> Map t f (delayed config a) +-- Generate repr sh f -> Generate repr sh f +-- Transform repr sh p f a -> Transform repr sh p f (delayed config a) +-- Backpermute shR sh p a -> Backpermute shR sh p (delayed config a) +-- Reshape slr sl a -> Reshape slr sl (manifest config a) + +-- Replicate{} -> fusionError +-- Slice{} -> fusionError +-- ZipWith{} -> fusionError + +-- -- Consumers +-- -- --------- +-- -- +-- -- Embed producers directly into the representation. For delayed terms +-- -- with local bindings, these will have been floated up above the +-- -- consumer already +-- -- +-- Fold f z a -> Fold f z (delayed config a) +-- FoldSeg i f z a s -> FoldSeg i f z (delayed config a) (delayed config s) +-- Scan d f z a -> Scan d f z (delayed config a) +-- Scan' d f z a -> Scan' d f z (delayed config a) +-- Permute f d p a -> Permute f (manifest config d) p (delayed config a) +-- Stencil s t f x a -> Stencil s t f x (delayed config a) +-- Stencil2 s1 s2 t f x a y b +-- -> Stencil2 s1 s2 t f x (delayed config a) y (delayed config b) +-- -- Collect s -> Collect (cvtS s) + +-- where +-- -- Flatten needless let-binds, which can be introduced by the +-- -- conversion to the internal embeddable representation. +-- -- +-- alet :: HasCallStack +-- => ALeftHandSide a aenv aenv' +-- -> DelayedOpenAcc aenv a +-- -> DelayedOpenAcc aenv' b +-- -> PreOpenAcc DelayedOpenAcc aenv b +-- alet lhs bnd body +-- | Just bodyVars <- extractDelayedArrayVars body +-- , Just Refl <- bindingIsTrivial lhs bodyVars +-- , Manifest x <- bnd +-- = x +-- -- +-- | otherwise +-- = Alet lhs bnd body + +-- -- Eliminate redundant application to an identity function. This +-- -- arises in the use of pipe to avoid fusion and force its argument +-- -- to be evaluated, i.e.: +-- -- +-- -- > compute :: Acc a -> Acc a +-- -- > compute = id >-> id +-- -- +-- apply :: HasCallStack +-- => ArraysR b +-- -> PreOpenAfun DelayedOpenAcc aenv (a -> b) +-- -> DelayedOpenAcc aenv a +-- -> PreOpenAcc DelayedOpenAcc aenv b +-- apply repr afun x +-- | Alam lhs (Abody body) <- afun +-- , Just bodyVars <- extractDelayedArrayVars body +-- , Just Refl <- bindingIsTrivial lhs bodyVars +-- , Manifest x' <- x +-- = Stats.ruleFired "applyD/identity" x' +-- -- +-- | otherwise +-- = Apply repr afun x + +-- cvtAF :: HasCallStack => OpenAfun aenv f -> PreOpenAfun DelayedOpenAcc aenv f +-- cvtAF (Alam lhs f) = Alam lhs (cvtAF f) +-- cvtAF (Abody b) = Abody (manifest config b) + +-- -- cvtS :: PreOpenSeq OpenAcc aenv senv s -> PreOpenSeq DelayedOpenAcc aenv senv s +-- -- cvtS = convertOpenSeq config + +-- convertOpenAfun :: HasCallStack => Config -> OpenAfun aenv f -> DelayedOpenAfun aenv f +-- convertOpenAfun c (Alam lhs f) = Alam lhs (convertOpenAfun c f) +-- convertOpenAfun c (Abody b) = Abody (convertOpenAcc c b) + +-- {-- +-- convertOpenSeq :: Config -> PreOpenSeq OpenAcc aenv senv a -> PreOpenSeq DelayedOpenAcc aenv senv a +-- convertOpenSeq config s = +-- case s of +-- Consumer c -> Consumer (cvtC c) +-- Reify ix -> Reify ix +-- Producer p s' -> Producer p' (convertOpenSeq config s') +-- where +-- p' = case p of +-- StreamIn arrs -> StreamIn arrs +-- ToSeq slix sh a -> ToSeq slix sh (delayed config a) +-- MapSeq f x -> MapSeq (cvtAF f) x +-- ChunkedMapSeq f x -> ChunkedMapSeq (cvtAF f) x +-- ZipWithSeq f x y -> ZipWithSeq (cvtAF f) x y +-- ScanSeq f e x -> ScanSeq (cvtF f) (cvtE e) x -- where --- cvtE :: Extend OpenAcc aenv aenv' -> Extend DelayedOpenAcc aenv aenv' --- cvtE BaseEnv = BaseEnv --- cvtE (PushEnv env a) | a' <- convertOpenAcc fuseAcc a = PushEnv (cvtE env) a' - -withSimplStats :: a -> a -#ifdef ACCELERATE_DEBUG -withSimplStats x = unsafePerformIO Stats.resetSimplCount `seq` x -#else -withSimplStats x = x -#endif - - --- | Apply the fusion transformation to an AST. This consists of two phases: --- --- 1. A bottom-up traversal that converts nodes into the internal delayed --- representation, merging adjacent producer/producer pairs. --- --- 2. A top-down traversal that makes the representation of fused --- consumer/producer pairs explicit as a 'DelayedAcc' of manifest and --- delayed nodes. --- --- TLM: Note that there really is no ambiguity as to which state an array will --- be in following this process: an array will be either delayed or --- manifest, and the two helper functions are even named as such! We should --- encode this property in the type somehow... --- -convertOpenAcc - :: HasCallStack - => Config - -> OpenAcc aenv arrs - -> DelayedOpenAcc aenv arrs -convertOpenAcc config = manifest config . computeAcc . embedOpenAcc config - - --- Convert array computations into an embeddable delayed representation. --- Reapply the embedding function from the first pass and unpack the --- representation. It is safe to match on BaseEnv because the first pass --- will put producers adjacent to the term consuming it. --- -delayed - :: HasCallStack - => Config - -> OpenAcc aenv (Array sh e) - -> DelayedOpenAcc aenv (Array sh e) -delayed config (embedOpenAcc config -> Embed env cc) - | BaseEnv <- env - = simplify cc & \case - Left (Done v) -> avarsIn Manifest v - Right d -> d & \case - Yield aR sh f -> Delayed aR sh f (f `compose` fromIndex (arrayRshape aR) sh) - Step aR sh p f v - | Just Refl <- matchOpenExp sh (arrayShape v) - , Just Refl <- isIdentity p -> Delayed aR sh (f `compose` indexArray v) (f `compose` linearIndex v) - | f' <- f `compose` indexArray v `compose` p -> Delayed aR sh f' (f' `compose` fromIndex (arrayRshape aR) sh) - -- - | otherwise - = manifest config (computeAcc (Embed env cc)) - - --- Convert array programs as manifest terms. --- -manifest - :: HasCallStack - => Config - -> OpenAcc aenv a - -> DelayedOpenAcc aenv a -manifest config (OpenAcc pacc) = - let fusionError = internalError "unexpected fusible materials" - in - Manifest $ case pacc of - -- Non-fusible terms - -- ----------------- - Avar ix -> Avar ix - Use aR a -> Use aR a - Unit t e -> Unit t e - Alet lhs bnd body -> alet lhs (manifest config bnd) (manifest config body) - Acond p t e -> Acond p (manifest config t) (manifest config e) - Awhile p f a -> Awhile (cvtAF p) (cvtAF f) (manifest config a) - Apair a1 a2 -> Apair (manifest config a1) (manifest config a2) - Anil -> Anil - Atrace msg a1 a2 -> Atrace msg (manifest config a1) (manifest config a2) - Apply repr f a -> apply repr (cvtAF f) (manifest config a) - Aforeign repr ff f a -> Aforeign repr ff (cvtAF f) (manifest config a) - - -- Producers - -- --------- - -- - -- Some producers might still exist as a manifest array. Typically this - -- is because they are the last stage of the computation, or the result - -- of a let-binding to be used multiple times. The input array here - -- should be a evaluated array term, else something went wrong. - -- - Map t f a -> Map t f (delayed config a) - Generate repr sh f -> Generate repr sh f - Transform repr sh p f a -> Transform repr sh p f (delayed config a) - Backpermute shR sh p a -> Backpermute shR sh p (delayed config a) - Reshape slr sl a -> Reshape slr sl (manifest config a) - - Replicate{} -> fusionError - Slice{} -> fusionError - ZipWith{} -> fusionError - - -- Consumers - -- --------- - -- - -- Embed producers directly into the representation. For delayed terms - -- with local bindings, these will have been floated up above the - -- consumer already - -- - Fold f z a -> Fold f z (delayed config a) - FoldSeg i f z a s -> FoldSeg i f z (delayed config a) (delayed config s) - Scan d f z a -> Scan d f z (delayed config a) - Scan' d f z a -> Scan' d f z (delayed config a) - Permute f d p a -> Permute f (manifest config d) p (delayed config a) - Stencil s t f x a -> Stencil s t f x (delayed config a) - Stencil2 s1 s2 t f x a y b - -> Stencil2 s1 s2 t f x (delayed config a) y (delayed config b) - -- Collect s -> Collect (cvtS s) - - where - -- Flatten needless let-binds, which can be introduced by the - -- conversion to the internal embeddable representation. - -- - alet :: HasCallStack - => ALeftHandSide a aenv aenv' - -> DelayedOpenAcc aenv a - -> DelayedOpenAcc aenv' b - -> PreOpenAcc DelayedOpenAcc aenv b - alet lhs bnd body - | Just bodyVars <- extractDelayedArrayVars body - , Just Refl <- bindingIsTrivial lhs bodyVars - , Manifest x <- bnd - = x - -- - | otherwise - = Alet lhs bnd body - - -- Eliminate redundant application to an identity function. This - -- arises in the use of pipe to avoid fusion and force its argument - -- to be evaluated, i.e.: - -- - -- > compute :: Acc a -> Acc a - -- > compute = id >-> id - -- - apply :: HasCallStack - => ArraysR b - -> PreOpenAfun DelayedOpenAcc aenv (a -> b) - -> DelayedOpenAcc aenv a - -> PreOpenAcc DelayedOpenAcc aenv b - apply repr afun x - | Alam lhs (Abody body) <- afun - , Just bodyVars <- extractDelayedArrayVars body - , Just Refl <- bindingIsTrivial lhs bodyVars - , Manifest x' <- x - = Stats.ruleFired "applyD/identity" x' - -- - | otherwise - = Apply repr afun x - - cvtAF :: HasCallStack => OpenAfun aenv f -> PreOpenAfun DelayedOpenAcc aenv f - cvtAF (Alam lhs f) = Alam lhs (cvtAF f) - cvtAF (Abody b) = Abody (manifest config b) - - -- cvtS :: PreOpenSeq OpenAcc aenv senv s -> PreOpenSeq DelayedOpenAcc aenv senv s - -- cvtS = convertOpenSeq config - -convertOpenAfun :: HasCallStack => Config -> OpenAfun aenv f -> DelayedOpenAfun aenv f -convertOpenAfun c (Alam lhs f) = Alam lhs (convertOpenAfun c f) -convertOpenAfun c (Abody b) = Abody (convertOpenAcc c b) - -{-- -convertOpenSeq :: Config -> PreOpenSeq OpenAcc aenv senv a -> PreOpenSeq DelayedOpenAcc aenv senv a -convertOpenSeq config s = - case s of - Consumer c -> Consumer (cvtC c) - Reify ix -> Reify ix - Producer p s' -> Producer p' (convertOpenSeq config s') - where - p' = case p of - StreamIn arrs -> StreamIn arrs - ToSeq slix sh a -> ToSeq slix sh (delayed config a) - MapSeq f x -> MapSeq (cvtAF f) x - ChunkedMapSeq f x -> ChunkedMapSeq (cvtAF f) x - ZipWithSeq f x y -> ZipWithSeq (cvtAF f) x y - ScanSeq f e x -> ScanSeq (cvtF f) (cvtE e) x - where - cvtC :: Consumer OpenAcc aenv senv a -> Consumer DelayedOpenAcc aenv senv a - cvtC c = - case c of - FoldSeq f e x -> FoldSeq (cvtF f) (cvtE e) x - FoldSeqFlatten f a x -> FoldSeqFlatten (cvtAF f) (manifest config a) x - Stuple t -> Stuple (cvtCT t) - - cvtCT :: Atuple (Consumer OpenAcc aenv senv) t -> Atuple (Consumer DelayedOpenAcc aenv senv) t - cvtCT NilAtup = NilAtup - cvtCT (SnocAtup t c) = SnocAtup (cvtCT t) (cvtC c) - - cvtAF :: OpenAfun aenv f -> PreOpenAfun DelayedOpenAcc aenv f - cvtAF (Alam f) = Alam (cvtAF f) - cvtAF (Abody b) = Abody (manifest config b) - - cvtE :: OpenExp env aenv t -> DelayedOpenExp env aenv t - cvtE = convertOpenExp config - - cvtF :: OpenFun env aenv f -> DelayedOpenFun env aenv f - cvtF (Lam f) = Lam (cvtF f) - cvtF (Body b) = Body (cvtE b) ---} - - -type EmbedAcc acc = forall aenv arrs. acc aenv arrs -> Embed acc aenv arrs -type ElimAcc acc = forall aenv s t. acc aenv s -> acc (aenv,s) t -> Bool - - --- | Apply the fusion transformation to the AST to combine and simplify terms. --- This converts terms into the internal delayed array representation and merges --- adjacent producer/producer terms. Using the reduced internal form limits the --- number of combinations that need to be considered. --- -embedOpenAcc :: HasCallStack => Config -> OpenAcc aenv arrs -> Embed OpenAcc aenv arrs -embedOpenAcc config (OpenAcc pacc) = - embedPreOpenAcc config matchOpenAcc (embedOpenAcc config) elimOpenAcc pacc - where - -- When does the cost of re-computation outweigh that of memory access? For - -- the moment only do the substitution on a single use of the bound array - -- into the use site, but it is likely advantageous to be far more - -- aggressive here. - -- - -- SEE: [Sharing vs. Fusion] - -- - elimOpenAcc :: ElimAcc OpenAcc - elimOpenAcc _bnd body - | count False ZeroIdx body <= lIMIT = True - | otherwise = False - where - lIMIT = 1 - - count :: UsesOfAcc OpenAcc - count no ix (OpenAcc pacc) = usesOfPreAcc no count ix pacc - - matchOpenAcc :: MatchAcc OpenAcc - matchOpenAcc (OpenAcc pacc1) (OpenAcc pacc2) = - matchPreOpenAcc matchOpenAcc pacc1 pacc2 - - -embedPreOpenAcc - :: HasCallStack - => Config - -> MatchAcc OpenAcc - -> EmbedAcc OpenAcc - -> ElimAcc OpenAcc - -> PreOpenAcc OpenAcc aenv arrs - -> Embed OpenAcc aenv arrs -embedPreOpenAcc config matchAcc embedAcc elimAcc pacc - = unembed - $ case pacc of - - -- Non-fusible terms - -- ----------------- - -- - -- Solid and semi-solid terms that we generally do not wish to fuse, such - -- as control flow (|?), array introduction (use, unit), array tupling and - -- projection, and foreign function operations. Generally we also do not - -- want to fuse past array let bindings, as this would imply work - -- duplication. SEE: [Sharing vs. Fusion] - -- - Alet lhs bnd body -> aletD embedAcc elimAcc lhs bnd body - Anil -> done $ Anil - Acond p at ae -> acondD matchAcc embedAcc (cvtE p) at ae - Apply aR f a -> done $ Apply aR (cvtAF f) (cvtA a) - Awhile p f a -> done $ Awhile (cvtAF p) (cvtAF f) (cvtA a) - Apair a1 a2 -> done $ Apair (cvtA a1) (cvtA a2) - Atrace msg a1 a2 -> done $ Atrace msg (cvtA a1) (cvtA a2) - Aforeign aR ff f a -> done $ Aforeign aR ff (cvtAF f) (cvtA a) - -- Collect s -> collectD s - - -- Array injection - Avar v -> done $ Avar v - Use aR a -> done $ Use aR a - Unit t e -> done $ Unit t (cvtE e) - - -- Producers - -- --------- - -- - -- The class of operations that given a set of zero or more input arrays, - -- produce a _single_ element for the output array by manipulating a - -- _single_ element from each input array. These can be further classified - -- as value (map, zipWith) or index space (backpermute, slice, replicate) - -- transformations. - -- - -- The critical feature is that each element of the output is produced - -- independently of all others, and so we can aggressively fuse arbitrary - -- sequences of these operations. - -- - Generate aR sh f -> generateD aR (cvtE sh) (cvtF f) - Map t f a -> mapD t (cvtF f) (embedAcc a) - ZipWith t f a b -> fuse2 (into (zipWithD t) (cvtF f)) a b - Transform aR sh p f a -> transformD aR (cvtE sh) (cvtF p) (cvtF f) (embedAcc a) - Backpermute slr sl p a -> fuse (into2 (backpermuteD slr) (cvtE sl) (cvtF p)) a - Slice slix a sl -> fuse (into (sliceD slix) (cvtE sl)) a - Replicate slix sh a -> fuse (into (replicateD slix) (cvtE sh)) a - Reshape slr sl a -> reshapeD slr (embedAcc a) (cvtE sl) - - -- Consumers - -- --------- - -- - -- Operations where each element of the output array depends on multiple - -- elements of the input array. To implement these operations efficiently in - -- parallel, we need to know how elements of the array depend on each other: - -- a parallel scan is implemented very differently from a parallel fold, for - -- example. - -- - -- In order to avoid obfuscating this crucial information required for - -- parallel implementation, fusion is separated into to phases: - -- producer/producer, implemented above, and consumer/producer, which is - -- implemented below. This will place producers adjacent to the consumer - -- node, so that the producer can be directly embedded into the consumer - -- during the code generation phase. - -- - Fold f z a -> embed aR (into2M Fold (cvtF f) (cvtE <$> z)) a - FoldSeg i f z a s -> embed2 aR (into2M (FoldSeg i) (cvtF f) (cvtE <$> z)) a s - Scan d f z a -> embed aR (into2M (Scan d) (cvtF f) (cvtE <$> z)) a - Scan' d f z a -> embed aR (into2 (Scan' d) (cvtF f) (cvtE z)) a - Permute f d p a -> embed2 aR (into2 permute (cvtF f) (cvtF p)) d a - Stencil s t f x a -> embed aR (into2 (stencil1 s t) (cvtF f) (cvtB x)) a - Stencil2 s1 s2 t f x a y b -> embed2 aR (into3 (stencil2 s1 s2 t) (cvtF f) (cvtB x) (cvtB y)) a b - - where - aR = arraysR pacc - - -- If fusion is not enabled, force terms to the manifest representation - -- - unembed :: HasCallStack => Embed OpenAcc aenv arrs -> Embed OpenAcc aenv arrs - unembed x - | array_fusion `member` options config = x - | Embed env cc <- x - , pacc <- compute cc - = case avarsOut extractOpenAcc pacc of - Just vars -> Embed env $ Done vars - _ - | DeclareVars lhs _ value <- declareVars (arraysR pacc) - -> Embed (PushEnv env lhs $ OpenAcc pacc) $ Done $ value weakenId - - cvtA :: HasCallStack => OpenAcc aenv' a -> OpenAcc aenv' a - cvtA = computeAcc . embedAcc - - cvtAF :: HasCallStack => PreOpenAfun OpenAcc aenv' f -> PreOpenAfun OpenAcc aenv' f - cvtAF (Alam lhs f) = Alam lhs (cvtAF f) - cvtAF (Abody a) = Abody (cvtA a) - - -- Helpers to shuffle the order of arguments to a constructor - -- - permute f p d a = Permute f d p a - - -- NOTE: [Stencil fusion] - -- - -- We allow stencils to delay their argument arrays with no special - -- considerations. This means that the delayed function will be evaluated - -- _at every element_ of the stencil pattern. We should do some analysis of - -- when this duplication is beneficial (keeping in mind that the stencil - -- implementations themselves may share neighbouring elements). - -- - stencil1 s t f x a = Stencil s t f x a - stencil2 s1 s2 t f x y a b = Stencil2 s1 s2 t f x a y b - - -- Conversions for closed scalar functions and expressions. This just - -- applies scalar simplifications. - -- - cvtF :: HasCallStack => Fun aenv' t -> Fun aenv' t - cvtF = simplifyFun - - cvtE :: HasCallStack => Exp aenv' t -> Exp aenv' t - cvtE = simplifyExp - - cvtB :: HasCallStack => Boundary aenv' t -> Boundary aenv' t - cvtB Clamp = Clamp - cvtB Mirror = Mirror - cvtB Wrap = Wrap - cvtB (Constant c) = Constant c - cvtB (Function f) = Function (cvtF f) - - -- Helpers to embed and fuse delayed terms - -- - into :: (HasCallStack, Sink f) - => (f env' a -> b) - -> f env a - -> Extend ArrayR OpenAcc env env' - -> b - into op a env = op (sinkA env a) - - into2 :: (HasCallStack, Sink f1, Sink f2) - => (f1 env' a -> f2 env' b -> c) - -> f1 env a - -> f2 env b - -> Extend ArrayR OpenAcc env env' - -> c - into2 op a b env = op (sinkA env a) (sinkA env b) - - into2M :: (HasCallStack, Sink f1, Sink f2) - => (f1 env' a -> Maybe (f2 env' b) -> c) - -> f1 env a - -> Maybe (f2 env b) - -> Extend ArrayR acc env env' - -> c - into2M op a b env = op (sinkA env a) (sinkA env <$> b) - - into3 :: (HasCallStack, Sink f1, Sink f2, Sink f3) - => (f1 env' a -> f2 env' b -> f3 env' c -> d) - -> f1 env a - -> f2 env b - -> f3 env c - -> Extend ArrayR OpenAcc env env' - -> d - into3 op a b c env = op (sinkA env a) (sinkA env b) (sinkA env c) - - -- Operations which can be fused into consumers. Move all of the local - -- bindings out of the way so that the fusible function operates - -- directly on the delayed representation. See also: [Representing - -- delayed arrays] - -- - fuse :: HasCallStack - => (forall r aenv'. Extend ArrayR OpenAcc aenv aenv' -> Cunctation r aenv' as -> Cunctation D aenv' bs) - -> OpenAcc aenv as - -> Embed OpenAcc aenv bs - fuse op (embedAcc -> Embed env cc) = Embed env (op env cc) - - fuse2 :: HasCallStack - => (forall r s aenv'. Extend ArrayR OpenAcc aenv aenv' -> Cunctation r aenv' as -> Cunctation s aenv' bs -> Cunctation D aenv' cs) - -> OpenAcc aenv as - -> OpenAcc aenv bs - -> Embed OpenAcc aenv cs - fuse2 op a1 a0 - | Embed env1 cc1 <- embedAcc a1 - , Embed env0 cc0 <- embedAcc (sinkA env1 a0) - , env <- env1 `append` env0 - = Embed env (op env (sinkA env0 cc1) cc0) - - -- Consumer operations which will be evaluated. - -- - -- NOTE: [Fusion and the lowest common use site] - -- - -- The AST given to us by sharing recovery will place let bindings at - -- the lowest common use site for that shared term. For example: - -- - -- fold f z (let a0 = .. - -- a1 = .. - -- in zipWith g a0 a1) - -- - -- In order to enable producer/consumer fusion for the above example, - -- it is necessary to float the let bindings above the `fold` - -- operation; SEE: [Sharing vs. Fusion] for more information. - -- - -- Furthermore, we used to maintain an invariant that all (manifest) - -- arguments were supplied as array variables, for example: - -- - -- fold1 f (let a0 = .. let a0 = .. - -- in stencil g a0) ==> a1 = stencil g a0 - -- a2 = fold1 f a1 - -- - -- However, if the argument term will be evaluated (i.e. can not be - -- fused into the producer) then it is better that we do _not_ float - -- those terms, and instead leave them under the consumer. This helps - -- to syntactically constrain the "liveness" of terms: if the argument - -- to an operation is not an array variable, we can see directly that - -- this will be the last use-site of that array. In particular, this is - -- useful for the 'permute' operation to know when it can in-place - -- update the array of default values. - -- - embed :: HasCallStack - => ArraysR bs - -> (forall aenv'. Extend ArrayR OpenAcc aenv aenv' -> OpenAcc aenv' as -> PreOpenAcc OpenAcc aenv' bs) - -> OpenAcc aenv as - -> Embed OpenAcc aenv bs - embed reprBs op (embedAcc -> Embed env cc) - | Done{} <- cc - , DeclareVars lhs _ value <- declareVars reprBs - = Embed (PushEnv BaseEnv lhs $ OpenAcc (op BaseEnv (computeAcc (Embed env cc)))) $ Done $ value weakenId - | otherwise - -- Next line is duplicated for both branches, as the type variable for the environment is instantiated differently - , DeclareVars lhs _ value <- declareVars reprBs - = Embed (PushEnv env lhs $ OpenAcc (op env (OpenAcc (compute cc)))) $ Done $ value weakenId - - embed2 :: HasCallStack - => ArraysR cs - -> (forall aenv'. Extend ArrayR OpenAcc aenv aenv' -> OpenAcc aenv' as -> OpenAcc aenv' bs -> PreOpenAcc OpenAcc aenv' cs) - -> OpenAcc aenv as - -> OpenAcc aenv bs - -> Embed OpenAcc aenv cs - embed2 reprCs op (embedAcc -> Embed env1 cc1) a0 - | Done{} <- cc1 - , a1 <- computeAcc (Embed env1 cc1) - = embed reprCs (\env0 -> op env0 (sinkA env0 a1)) a0 - -- - | Embed env0 cc0 <- embedAcc (sinkA env1 a0) - , env <- env1 `append` env0 - = case cc0 of - Done{} - | DeclareVars lhs _ value <- declareVars reprCs - -> Embed (PushEnv env1 lhs $ OpenAcc (op env1 (OpenAcc (compute cc1)) (computeAcc (Embed env0 cc0)))) - $ Done - $ value weakenId - _ - -- Next line is duplicated for both branches, as the type - -- variable for the environment is instantiated differently - | DeclareVars lhs _ value <- declareVars reprCs - -> Embed (PushEnv env lhs $ OpenAcc (op env (OpenAcc (compute (sinkA env0 cc1))) (OpenAcc (compute cc0)))) - $ Done - $ value weakenId - - -- trav1 :: (Arrays as, Arrays bs) - -- => (forall aenv'. Embed acc aenv' as -> Embed acc aenv' as) - -- -> (forall aenv'. Extend ArrayR acc aenv aenv' -> acc aenv' as -> PreOpenAcc acc aenv' bs) - -- -> acc aenv as - -- -> Embed acc aenv bs - -- trav1 f op (f . embedAcc -> Embed env cc) - -- = Embed (env `pushArrayEnv` inject (op env (inject (compute cc)))) doneZeroIdx - - -- trav2 :: (Arrays as, Arrays bs, Arrays cs) - -- => (forall aenv'. Embed acc aenv' as -> Embed acc aenv' as) - -- -> (forall aenv'. Embed acc aenv' bs -> Embed acc aenv' bs) - -- -> (forall aenv'. Extend ArrayR acc aenv aenv' -> acc aenv' as -> acc aenv' bs -> PreOpenAcc acc aenv' cs) - -- -> acc aenv as - -- -> acc aenv bs - -- -> Embed acc aenv cs - -- trav2 f1 f0 op (f1 . embedAcc -> Embed env1 cc1) (f0 . embedAcc . sinkA env1 -> Embed env0 cc0) - -- | env <- env1 `append` env0 - -- , acc1 <- inject . compute $ sinkA env0 cc1 - -- , acc0 <- inject . compute $ cc0 - -- = Embed (env `pushArrayEnv` inject (op env acc1 acc0)) doneZeroIdx - - -- force :: Arrays as => Embed acc aenv' as -> Embed acc aenv' as - -- force (Embed env cc) - -- | Done{} <- cc = Embed env cc - -- | otherwise = Embed (env `pushArrayEnv` inject (compute cc)) doneZeroIdx - - -- -- Move additional bindings for producers outside of the sequence, so that - -- -- producers may fuse with their arguments resulting in actual sequencing - -- collectD :: PreOpenSeq acc aenv () arrs - -- -> Embed acc aenv arrs - -- collectD (embedSeq embedAcc -> ExtendSeq env s') - -- = Embed (env `pushArrayEnv` inject (Collect s')) doneZeroIdx - - -{-- --- Move additional bindings for producer outside of sequence, so --- that producers may fuse with their arguments, resulting in --- actual sequencing. -embedSeq :: forall acc aenv arrs. Kit acc - => EmbedAcc acc - -> PreOpenSeq acc aenv () arrs - -> ExtendSeq acc aenv () arrs -embedSeq embedAcc s - = travS s BaseEnv - where - travS :: forall senv aenv' arrs'. - PreOpenSeq acc aenv senv arrs' - -> Extend acc aenv aenv' - -> ExtendSeq acc aenv senv arrs' - travS s env = - case s of - Producer p s - | ExtendSeq env' s' <- travS s env - , ExtendProducer env'' p' <- travP p env' - -> ExtendSeq (env' `append` env'') (Producer p' (sinkSeq env'' s')) - Consumer c - | c' <- travC c env - -> ExtendSeq env (Consumer c') - Reify ix - -> ExtendSeq env (Reify ix) - - travP :: forall arrs' aenv' senv. - Producer acc aenv senv arrs' - -> Extend acc aenv aenv' - -> ExtendProducer acc aenv' senv arrs' - travP (ToSeq slix sh a) env - | Embed env' cc <- embedAcc (sink env a) - = ExtendProducer env' (ToSeq slix sh (inject (compute cc))) - travP (StreamIn arrs) _ = ExtendProducer BaseEnv (StreamIn arrs) - travP (MapSeq f x) env = ExtendProducer BaseEnv (MapSeq (cvtAF (sink env f)) x) - travP (ChunkedMapSeq f x) env = ExtendProducer BaseEnv (ChunkedMapSeq (cvtAF (sink env f)) x) - travP (ZipWithSeq f x y) env = ExtendProducer BaseEnv (ZipWithSeq (cvtAF (sink env f)) x y) - travP (ScanSeq f e x) env = ExtendProducer BaseEnv (ScanSeq (cvtF (sink env f)) (cvtE (sink env e)) x) - - travC :: forall arrs' aenv' senv. - Consumer acc aenv senv arrs' - -> Extend acc aenv aenv' - -> Consumer acc aenv' senv arrs' - travC (FoldSeq f e x) env = FoldSeq (cvtF (sink env f)) (cvtE (sink env e)) x - travC (FoldSeqFlatten f a x) env = FoldSeqFlatten (cvtAF (sink env f)) (cvtA (sink env a)) x - travC (Stuple t) env = Stuple (cvtCT t) - where - cvtCT :: Atuple (Consumer acc aenv senv) t -> Atuple (Consumer acc aenv' senv) t - cvtCT NilAtup = NilAtup - cvtCT (SnocAtup t c) = SnocAtup (cvtCT t) (travC c env) - - cvtE :: Elt t => Exp aenv' t -> Exp aenv' t - cvtE = simplifyExp - - cvtF :: Fun aenv' t -> Fun aenv' t - cvtF = simplifyFun - - cvtA :: Arrays a => acc aenv' a -> acc aenv' a - cvtA = computeAcc . embedAcc - - cvtAF :: PreOpenAfun acc aenv' f -> PreOpenAfun acc aenv' f - cvtAF (Alam f) = Alam (cvtAF f) - cvtAF (Abody a) = Abody (cvtA a) - - --- A sequence with additional bindings -data ExtendSeq acc aenv senv arrs where - ExtendSeq :: forall acc aenv aenv' senv arrs. - Extend acc aenv aenv' - -> PreOpenSeq acc aenv' senv arrs - -> ExtendSeq acc aenv senv arrs - --- A producer with additional bindings -data ExtendProducer acc aenv senv arrs where - ExtendProducer :: forall acc aenv aenv' senv arrs. - Extend acc aenv aenv' - -> Producer acc aenv' senv arrs - -> ExtendProducer acc aenv senv arrs ---} - - --- Internal representation --- ======================= - --- NOTE: [Representing delayed arrays] --- --- During the fusion transformation we represent terms as a pair consisting of --- a collection of supplementary environment bindings and a description of how --- to construct the array. --- --- It is critical to separate these two. To create a real AST node we need both --- the environment and array term, but analysis of how to fuse terms requires --- only the array description. If the additional bindings are bundled as part of --- the representation, the existentially quantified extended environment type --- will be untouchable. This is problematic because the terms of the two arrays --- are defined with respect to this existentially quantified type, and there is --- no way to directly combine these two environments: --- --- append :: Extend env env1 -> Extend env env2 -> Extend env ??? --- --- And hence, no way to combine the terms of the delayed representation. --- --- The only way to bring terms into the same scope is to operate via the --- manifest terms. This entails a great deal of conversion between delayed and --- AST terms, but is certainly possible. --- --- However, because of the limited scope into which this existential type is --- available, we ultimately perform this process many times. In fact, complexity --- of the fusion algorithm for an AST of N terms becomes O(r^n), where r is the --- number of different rules we have for combining terms. --- -data Embed acc aenv a where - Embed :: Extend ArrayR acc aenv aenv' - -> Cunctation r aenv' a - -> Embed acc aenv a - -instance HasArraysR acc => HasArraysR (Embed acc) where - arraysR (Embed _ c) = arraysR c - - --- Delayed arrays are represented as functions -data D - --- Manifest arrays are real arrays backed by unboxed vectors -data M - --- Cunctation (n): the action or an instance of delaying; a tardy action. --- --- This describes the ways in which the fusion transformation represents --- intermediate arrays. The fusion process operates by recasting producer array --- computations in terms of a set of scalar functions used to construct an --- element at each index, and fusing successive producers by combining these --- scalar functions. --- -data Cunctation r aenv a where - - -- The base case is just a real (manifest) array term. No fusion happens here. - -- Note that the array is referenced by an index into the extended - -- environment, ensuring that the array is manifest and making the term - -- non-recursive in 'acc'. - -- - Done :: ArrayVars aenv arrs - -> Cunctation M aenv arrs - - -- We can represent an array by its shape and a function to compute an element - -- at each index. - -- - Yield :: ArrayR (Array sh e) - -> Exp aenv sh - -> Fun aenv (sh -> e) - -> Cunctation D aenv (Array sh e) - - -- A more restrictive form than 'Yield' may afford greater opportunities for - -- optimisation by a backend. This more structured form applies an index and - -- value transform to an input array. Note that the transform is applied to an - -- array stored as an environment index, so that the term is non-recursive and - -- it is always possible to embed into a collective operation. - -- - Step :: ArrayR (Array sh' b) - -> Exp aenv sh' - -> Fun aenv (sh' -> sh) - -> Fun aenv (a -> b) - -> ArrayVar aenv (Array sh a) - -> Cunctation D aenv (Array sh' b) - -instance HasArraysR (Cunctation r) where - arraysR (Done v) = varsType v - arraysR (Yield aR _ _) = TupRsingle aR - arraysR (Step aR _ _ _ _) = TupRsingle aR - -instance Sink (Cunctation r) where - weaken k = \case - Done v -> Done (weakenVars k v) - Step aR sh p f v -> Step aR (weaken k sh) (weaken k p) (weaken k f) (weaken k v) - Yield aR sh f -> Yield aR (weaken k sh) (weaken k f) - -simplify - :: HasCallStack - => Cunctation r aenv a - -> Either (Cunctation M aenv a) (Cunctation D aenv a) -simplify = \case - Done v - -> Left $ Done v - Yield aR (simplifyExp -> sh) (simplifyFun -> f) - -> Right $ Yield aR sh f - Step aR (simplifyExp -> sh) (simplifyFun -> p) (simplifyFun -> f) v - | Just Refl <- matchOpenExp sh (arrayShape v) - , Just Refl <- isIdentity p - , Just Refl <- isIdentity f - -> Left $ (Done (TupRsingle v)) - | otherwise - -> Right $ Step aR sh p f v - - --- Convert a real AST node into the internal representation --- -done :: HasCallStack => PreOpenAcc OpenAcc aenv a -> Embed OpenAcc aenv a -done pacc - | Just vars <- avarsOut extractOpenAcc pacc - = Embed BaseEnv (Done vars) - | DeclareVars lhs _ value <- declareVars (arraysR pacc) - = Embed (PushEnv BaseEnv lhs $ OpenAcc pacc) $ Done $ value weakenId - -doneZeroIdx :: ArrayR (Array sh e) -> Cunctation M (aenv, Array sh e) (Array sh e) -doneZeroIdx aR = Done (TupRsingle (Var aR ZeroIdx)) - - --- Recast a cunctation into a mapping from indices to elements --- -yield :: HasCallStack - => Cunctation r aenv (Array sh e) - -> Cunctation D aenv (Array sh e) -yield cc = - case cc of - Yield{} -> cc - Step tR sh p f v -> Yield tR sh (f `compose` indexArray v `compose` p) - Done (TupRsingle v@(Var tR _)) -> Yield tR (arrayShape v) (indexArray v) - --- Recast a cunctation into a delayed representation --- -delaying - :: HasCallStack - => Cunctation r aenv (Array sh e) - -> Cunctation D aenv (Array sh e) -delaying cc = - case cc of - Yield{} -> cc - Step{} -> cc - Done u - | TupRsingle v <- u - , Var aR _ <- v - , ArrayR shR tR <- aR - -> Step aR (arrayShape v) (identity (shapeType shR)) (identity tR) v - --- Get the shape of a delayed array --- -shape :: HasCallStack => Cunctation r aenv (Array sh e) -> Exp aenv sh -shape cc = - case delaying cc of - Step _ sh _ _ _ -> sh - Yield _ sh _ -> sh - - --- prjExtend :: Kit acc => Extend acc env env' -> Idx env' t -> PreOpenAcc acc env' t --- prjExtend (PushEnv _ v) ZeroIdx = weakenA rebuildAcc SuccIdx v --- prjExtend (PushEnv env _) (SuccIdx idx) = weakenA rebuildAcc SuccIdx $ prjExtend env idx --- prjExtend _ _ = $internalError "prjExtend" "inconsistent valuation" - -{-- --- Rearrange type arguments to fit with Sink type class. -newtype SinkSeq acc senv aenv a = SinkSeq { unSinkSeq :: PreOpenSeq acc aenv senv a } - --- sink for sequences. -sinkSeq :: Kit acc => Extend acc aenv aenv' -> PreOpenSeq acc aenv senv a -> PreOpenSeq acc aenv' senv a -sinkSeq env s = unSinkSeq $ sink env (SinkSeq s) - -instance Kit acc => Sink (SinkSeq acc senv) where - weaken :: forall aenv aenv' arrs. aenv :> aenv' -> SinkSeq acc senv aenv arrs -> SinkSeq acc senv aenv' arrs - weaken k (SinkSeq s) = SinkSeq $ - case s of - Producer p s' -> Producer (weakenP p) (weakenL s') - Consumer c -> Consumer (weakenC c) - Reify ix -> Reify ix - - where - weakenL :: forall senv' arrs'. PreOpenSeq acc aenv senv' arrs' -> PreOpenSeq acc aenv' senv' arrs' - weakenL s' = unSinkSeq (weaken k (SinkSeq s')) - - weakenP :: forall a. Producer acc aenv senv a -> Producer acc aenv' senv a - weakenP p = - case p of - StreamIn arrs -> StreamIn arrs - ToSeq slix sh a -> ToSeq slix sh (weaken k a) - MapSeq f x -> MapSeq (weaken k f) x - ChunkedMapSeq f x -> ChunkedMapSeq (weaken k f) x - ZipWithSeq f x y -> ZipWithSeq (weaken k f) x y - ScanSeq f a x -> ScanSeq (weaken k f) (weaken k a) x - - weakenC :: forall a. Consumer acc aenv senv a -> Consumer acc aenv' senv a - weakenC c = - case c of - FoldSeq f a x -> FoldSeq (weaken k f) (weaken k a) x - FoldSeqFlatten f a x -> FoldSeqFlatten (weaken k f) (weaken k a) x - Stuple t -> - let wk :: Atuple (Consumer acc aenv senv) t -> Atuple (Consumer acc aenv' senv) t - wk NilAtup = NilAtup - wk (SnocAtup t c) = wk t `SnocAtup` weakenC c - in - Stuple (wk t) ---} - --- Array fusion of a de Bruijn computation AST --- =========================================== - --- Array computations --- ------------------ - --- Evaluate a delayed computation and tie the recursive knot --- --- We do a bit of extra work to (try to) maintain that terms should be left --- at their lowest common use site. SEE: [Fusion and the lowest common use site] --- -computeAcc - :: HasCallStack - => Embed OpenAcc aenv arrs - -> OpenAcc aenv arrs -computeAcc (Embed BaseEnv cc) = OpenAcc (compute cc) -computeAcc (Embed env@(PushEnv bot lhs top) cc) = - simplify cc & \case - Left (Done v) -> bindA env (avarsIn OpenAcc v) - Right d -> d & \case - Yield repr sh f - -> bindA env (OpenAcc (Generate repr sh f)) - - Step repr sh p f v@(Var _ ix) - | Just Refl <- matchOpenExp sh (arrayShape v) - , Just Refl <- isIdentity p - -> case ix of - ZeroIdx - | LeftHandSideSingle ArrayR{} <- lhs - , Just (OpenAccFun g) <- strengthen noTop (OpenAccFun f) - -> bindA bot (OpenAcc (Map (arrayRtype repr) g top)) - _ -> bindA env (OpenAcc (Map (arrayRtype repr) f (avarIn OpenAcc v))) - - | Just Refl <- isIdentity f - -> case ix of - ZeroIdx - | LeftHandSideSingle ArrayR{} <- lhs - , Just (OpenAccFun q) <- strengthen noTop (OpenAccFun p) - , Just (OpenAccExp sz) <- strengthen noTop (OpenAccExp sh) - -> bindA bot (OpenAcc (Backpermute (arrayRshape repr) sz q top)) - _ -> bindA env (OpenAcc (Backpermute (arrayRshape repr) sh p (avarIn OpenAcc v))) - - | otherwise - -> case ix of - ZeroIdx - | LeftHandSideSingle ArrayR{} <- lhs - , Just (OpenAccFun g) <- strengthen noTop (OpenAccFun f) - , Just (OpenAccFun q) <- strengthen noTop (OpenAccFun p) - , Just (OpenAccExp sz) <- strengthen noTop (OpenAccExp sh) - -> bindA bot (OpenAcc (Transform repr sz q g top)) - _ -> bindA env (OpenAcc (Transform repr sh p f (avarIn OpenAcc v))) - - where - bindA :: HasCallStack - => Extend ArrayR OpenAcc aenv aenv' - -> OpenAcc aenv' a - -> OpenAcc aenv a - bindA BaseEnv b = b - bindA (PushEnv env lhs a) b - -- If the freshly bound value is directly, returned, we don't have to bind it in a - -- let. We can do this if the left hand side does not contain wildcards (other than - -- wildcards for unit / nil) and if the value contains the same variables. - | Just vars <- extractOpenArrayVars b - , Just Refl <- bindingIsTrivial lhs vars = bindA env a - | otherwise = bindA env (OpenAcc (Alet lhs a b)) - - noTop :: (aenv, a) :?> aenv - noTop ZeroIdx = Nothing - noTop (SuccIdx ix) = Just ix - - - --- Convert the internal representation of delayed arrays into a real AST --- node. Use the most specific version of a combinator whenever possible. --- -compute - :: HasCallStack - => Cunctation r aenv arrs - -> PreOpenAcc OpenAcc aenv arrs -compute cc = simplify cc & \case - Left (Done v) -> v & \case - TupRunit -> Anil - TupRsingle v@(Var ArrayR{} _) -> Avar v - TupRpair v1 v2 -> avarsIn OpenAcc v1 `Apair` avarsIn OpenAcc v2 - Right d -> d & \case - Yield aR sh f -> Generate aR sh f - Step (ArrayR shR tR) sh p f v - | Just Refl <- matchOpenExp sh (arrayShape v) - , Just Refl <- isIdentity p -> Map tR f (avarIn OpenAcc v) - | Just Refl <- isIdentity f -> Backpermute shR sh p (avarIn OpenAcc v) - | otherwise -> Transform (ArrayR shR tR) sh p f (avarIn OpenAcc v) - - --- Representation of a generator as a delayed array --- -generateD - :: HasCallStack - => ArrayR (Array sh e) - -> Exp aenv sh - -> Fun aenv (sh -> e) - -> Embed OpenAcc aenv (Array sh e) -generateD repr sh f - = Stats.ruleFired "generateD" - $ Embed BaseEnv (Yield repr sh f) - - --- Fuse a unary function into a delayed array. Also looks for unzips which can --- be executed in constant time; SEE [unzipD] --- -mapD :: HasCallStack - => TypeR b - -> Fun aenv (a -> b) - -> Embed OpenAcc aenv (Array sh a) - -> Embed OpenAcc aenv (Array sh b) -mapD tR f (unzipD tR f -> Just a) = a -mapD tR f (Embed env cc) - = Stats.ruleFired "mapD" - $ Embed env (go (delaying cc)) - where - go (Step (ArrayR shR _) sh ix g v) = Step (ArrayR shR tR) sh ix (sinkA env f `compose` g) v - go (Yield (ArrayR shR _) sh g) = Yield (ArrayR shR tR) sh (sinkA env f `compose` g) - - --- If we are unzipping a manifest array then force the term to be computed; --- a backend will be able to execute this in constant time. --- -unzipD - :: HasCallStack - => TypeR b - -> Fun aenv (a -> b) - -> Embed OpenAcc aenv (Array sh a) - -> Maybe (Embed OpenAcc aenv (Array sh b)) -unzipD tR f (Embed env cc@(Done v)) - | Lam lhs (Body a) <- f - , Just vars <- extractExpVars a - , ArrayR shR _ <- arrayR cc - , f' <- Lam lhs $ Body $ expVars vars - = Just $ Embed (env `pushArrayEnv` OpenAcc (Map tR f' $ avarsIn OpenAcc v)) $ doneZeroIdx $ ArrayR shR tR -unzipD _ _ _ - = Nothing - --- Fuse an index space transformation function that specifies where elements in --- the destination array read there data from in the source array. --- -backpermuteD - :: HasCallStack - => ShapeR sh' - -> Exp aenv sh' - -> Fun aenv (sh' -> sh) - -> Cunctation r aenv (Array sh e) - -> Cunctation D aenv (Array sh' e) -backpermuteD shR' sh' p = Stats.ruleFired "backpermuteD" . go . delaying - where - go (Step (ArrayR _ tR) _ q f v) = Step (ArrayR shR' tR) sh' (q `compose` p) f v - go (Yield (ArrayR _ tR) _ g) = Yield (ArrayR shR' tR) sh' (g `compose` p) - - --- Transform as a combined map and backwards permutation --- -transformD - :: HasCallStack - => ArrayR (Array sh' b) - -> Exp aenv sh' - -> Fun aenv (sh' -> sh) - -> Fun aenv (a -> b) - -> Embed OpenAcc aenv (Array sh a) - -> Embed OpenAcc aenv (Array sh' b) -transformD (ArrayR shR' tR) sh' p f - = Stats.ruleFired "transformD" - . fuse (into2 (backpermuteD shR') sh' p) - . mapD tR f - where - fuse :: HasCallStack - => (forall r aenv'. Extend ArrayR OpenAcc aenv aenv' -> Cunctation r aenv' as -> Cunctation D aenv' bs) - -> Embed OpenAcc aenv as - -> Embed OpenAcc aenv bs - fuse op (Embed env cc) = Embed env (op env cc) - - into2 :: (HasCallStack, Sink f1, Sink f2) - => (f1 env' a -> f2 env' b -> c) - -> f1 env a - -> f2 env b - -> Extend ArrayR OpenAcc env env' - -> c - into2 op a b env = op (sinkA env a) (sinkA env b) - - --- Replicate as a backwards permutation --- --- TODO: If we have a pattern such as `replicate sh (map f xs)` then in some --- cases it might be beneficial to not fuse these terms, if `f` is --- expensive and/or `sh` is large. --- -replicateD - :: HasCallStack - => SliceIndex slix sl co sh - -> Exp aenv slix - -> Cunctation r aenv (Array sl e) - -> Cunctation D aenv (Array sh e) -replicateD sliceIndex slix cc - = Stats.ruleFired "replicateD" - $ backpermuteD (sliceDomainR sliceIndex) (IndexFull sliceIndex slix (shape cc)) (extend sliceIndex slix) cc - - --- Dimensional slice as a backwards permutation --- -sliceD - :: HasCallStack - => SliceIndex slix sl co sh - -> Exp aenv slix - -> Cunctation r aenv (Array sh e) - -> Cunctation D aenv (Array sl e) -sliceD sliceIndex slix cc - = Stats.ruleFired "sliceD" - $ backpermuteD (sliceShapeR sliceIndex) (IndexSlice sliceIndex slix (shape cc)) (restrict sliceIndex slix) cc - - --- Reshape an array --- --- For delayed arrays this is implemented as an index space transformation. For --- manifest arrays this can be done with the standard Reshape operation in --- constant time without executing any array operations. This does not affect --- the fusion process since the term is already manifest. --- --- TLM: there was a runtime check to ensure the old and new shapes contained the --- same number of elements: this has been lost for the delayed cases! --- -reshapeD - :: HasCallStack - => ShapeR sl - -> Embed OpenAcc aenv (Array sh e) - -> Exp aenv sl - -> Embed OpenAcc aenv (Array sl e) -reshapeD slr (Embed env cc) (sinkA env -> sl) - | Done v <- cc - = Embed (env `pushArrayEnv` OpenAcc (Reshape slr sl (avarsIn OpenAcc v))) $ doneZeroIdx repr - - | otherwise - = Stats.ruleFired "reshapeD" - $ Embed env (backpermuteD slr sl (reindex (arrayRshape $ arrayR cc) (shape cc) slr sl) cc) - - where - ArrayR _ tR = arrayR cc - repr = ArrayR slr tR - - --- Combine two arrays element-wise with a binary function to produce a delayed --- array. --- -zipWithD - :: HasCallStack - => TypeR c - -> Fun aenv (a -> b -> c) - -> Cunctation r aenv (Array sh a) - -> Cunctation s aenv (Array sh b) - -> Cunctation D aenv (Array sh c) -zipWithD tR f cc1 cc0 - -- Two stepper functions identically accessing the same array can be kept in - -- stepping form. This might yield a simpler final term. - -- - | Step (ArrayR shR _) sh1 p1 f1 v1 <- delaying cc1 - , Step _ sh0 p0 f0 v0 <- delaying cc0 - , Just Refl <- matchVar v1 v0 - , Just Refl <- matchOpenFun p1 p0 - = Stats.ruleFired "zipWithD/step" - $ Step (ArrayR shR tR) (intersect shR sh1 sh0) p0 (combine f f1 f0) v0 - - -- Otherwise transform both delayed terms into (index -> value) mappings and - -- combine the two indexing functions that way. - -- - | Yield (ArrayR shR _) sh1 f1 <- yield cc1 - , Yield _ sh0 f0 <- yield cc0 - = Stats.ruleFired "zipWithD" - $ Yield (ArrayR shR tR) (intersect shR sh1 sh0) (combine f f1 f0) - - | otherwise - = error "work is stressing me out, I should take a break" - where - combine :: forall aenv a b c e. HasCallStack - => Fun aenv (a -> b -> c) - -> Fun aenv (e -> a) - -> Fun aenv (e -> b) - -> Fun aenv (e -> c) - combine c ixa ixb - | Lam lhs1 (Body ixa') <- ixa - , Lam lhs2 (Body ixb') <- ixb - -- The two LeftHandSides may differ in the use of wildcards. If they do not match, we must - -- combine them as done in `combineLHS`. As this will probably not occur often and requires - -- additional weakening, we do a quick check whether the left hand sides are equal. - -- - = case matchELeftHandSide lhs1 lhs2 of - Just Refl - | Lam lhsA (Lam lhsB (Body c')) <- weakenE (weakenWithLHS lhs1) c - -> Lam lhs1 $ Body $ Let lhsA ixa' $ Let lhsB (weakenE (weakenWithLHS lhsA) ixb') c' - Nothing - | CombinedLHS lhs k1 k2 <- combineLHS lhs1 lhs2 - , Lam lhsA (Lam lhsB (Body c')) <- weakenE (weakenWithLHS lhs) c - , ixa'' <- weakenE k1 ixa' - -> Lam lhs $ Body $ Let lhsA ixa'' $ Let lhsB (weakenE (weakenWithLHS lhsA .> k2) ixb') c' - _ - -> error "how's your break?" - -- - | otherwise - = error "work is stressing me out, I should get back to it" - - -data CombinedLHS s t env1' env2' env where - CombinedLHS :: LeftHandSide s t env env' - -> env1' :> env' - -> env2' :> env' - -> CombinedLHS s t env1' env2' env - -combineLHS - :: HasCallStack - => LeftHandSide s t env env1' - -> LeftHandSide s t env env2' - -> CombinedLHS s t env1' env2' env -combineLHS = go weakenId weakenId - where - go :: env1 :> env -> env2 :> env -> LeftHandSide s t env1 env1' -> LeftHandSide s t env2 env2' -> CombinedLHS s t env1' env2' env - go k1 k2 (LeftHandSideWildcard tR) (LeftHandSideWildcard _) = CombinedLHS (LeftHandSideWildcard tR) k1 k2 - go k1 k2 (LeftHandSideSingle tR) (LeftHandSideSingle _) = CombinedLHS (LeftHandSideSingle tR) (sink k1) (sink k2) - go k1 k2 (LeftHandSidePair l1 h1) (LeftHandSidePair l2 h2) - | CombinedLHS l k1' k2' <- go k1 k2 l1 l2 - , CombinedLHS h k1'' k2'' <- go k1' k2' h1 h2 = CombinedLHS (LeftHandSidePair l h) k1'' k2'' - go k1 k2 (LeftHandSideWildcard _) lhs - | Exists lhs' <- rebuildLHS lhs = CombinedLHS lhs' (weakenWithLHS lhs' .> k1) (sinkWithLHS lhs lhs' k2) - go k1 k2 lhs (LeftHandSideWildcard _) - | Exists lhs' <- rebuildLHS lhs = CombinedLHS lhs' (sinkWithLHS lhs lhs' k1) (weakenWithLHS lhs' .> k2) - go _ _ _ _ - = internalError "unexpected LHS combination" - - - --- NOTE: [Sharing vs. Fusion] --- --- The approach to array fusion is similar to that the first generation of Repa. --- It was discovered that the most immediately pressing problem with delayed --- arrays in Repa-1 was that it did not preserve sharing of collective --- operations, leading to excessive recomputation and severe repercussions on --- performance if the user did not explicitly intervene. --- --- However, as we have explicit sharing information in the term tree, so it is --- straightforward to respect sharing by not fusing let-bindings, as that --- introduces work duplication. However, sometimes we can be cleverer. --- --- let-floating: --- ------------- --- --- If the binding is of manifest data, we can instead move the let-binding to a --- different point in the program and then continue to fuse into the body. This --- is done by adding the bound term to the Extend environment. In essence this --- is covering a different occurrence of the same problem Extend was introduced --- to handle: let bindings of manifest data unnecessarily get in the way of the --- fusion process. For example: --- --- map f (zipWith g xs (map h xs)) --- --- after sharing recovery results in: --- --- map f (let a0 = xs in zipWith g a0 (map h a0)) --- --- Without allowing the binding for a0 to float outwards, `map f` will not be --- fused into the rest of the program. --- --- let-elimination: --- ---------------- --- --- Array binding points appear in the program because the array data _or_ shape --- was accessed multiple times in the source program. In general we want to fuse --- arbitrary sequences of array _data_, irrespective of how the shape component --- is used. For example, reverse is defined in the prelude as: --- --- reverse xs = let len = unindex1 (shape xs) --- pf i = len - i - 1 --- in --- backpermute (shape xs) (ilift1 pf) xs --- --- Sharing recovery introduces a let-binding for the input `xs` since it is used --- thrice in the definition, which impedes subsequent fusion. However the actual --- array data is only accessed once, with the remaining two uses querying the --- array shape. Since the delayed terms contain the shape of the array they --- represent as a scalar term, if the data component otherwise satisfies the --- rules for fusing terms, as it does in this example, we can eliminate the --- let-binding by pushing the scalar shape and value generation terms directly --- into the body. --- --- Let-elimination can also be used to _introduce_ work duplication, which may --- be beneficial if we can estimate that the cost of re-computation is less than --- the cost of completely evaluating the array and subsequently retrieving the --- data from memory. --- --- let-binding: --- ------------ --- --- Ultimately, we might not want to eliminate the binding. If so, evaluate it --- and add it to a _clean_ Extend environment for the body. If not, the Extend --- list effectively _flattens_ all bindings, so any terms required for the bound --- term get lifted out to the same scope as the body. This increases their --- lifetime and hence raises the maximum memory used. If we don't do this, we --- get terms such as: --- --- let a0 = in --- let bnd = in --- --- --- rather than the following, where the scope of a0 is clearly only availably --- when evaluating the bound term, as it should be: --- --- let bnd = --- let a0 = --- in --- in --- -aletD :: HasCallStack - => EmbedAcc OpenAcc - -> ElimAcc OpenAcc - -> ALeftHandSide arrs aenv aenv' - -> OpenAcc aenv arrs - -> OpenAcc aenv' brrs - -> Embed OpenAcc aenv brrs -aletD embedAcc elimAcc lhs (embedAcc -> Embed env1 cc1) acc0 - - -- let-floating - -- ------------ - -- - -- Immediately inline the variable referring to the bound expression into the - -- body, instead of adding to the environments and creating an indirection - -- that must be later eliminated by shrinking. - -- - | LeftHandSideSingle _ <- lhs - , Done (TupRsingle v1@(Var ArrayR{} _)) <- cc1 - , Embed env0 cc0 <- embedAcc $ rebuildA (subAtop (Avar v1) . sink1 env1) acc0 - = Stats.ruleFired "aletD/float" - $ Embed (env1 `append` env0) cc0 - - -- Ensure we only call 'embedAcc' once on the body expression - -- - | otherwise - = aletD' embedAcc elimAcc lhs (Embed env1 cc1) (embedAcc acc0) - - -aletD' :: forall aenv aenv' arrs brrs. HasCallStack - => EmbedAcc OpenAcc - -> ElimAcc OpenAcc - -> ALeftHandSide arrs aenv aenv' - -> Embed OpenAcc aenv arrs - -> Embed OpenAcc aenv' brrs - -> Embed OpenAcc aenv brrs -aletD' embedAcc elimAcc (LeftHandSideSingle ArrayR{}) (Embed env1 cc1) (Embed env0 cc0) - - -- let-binding - -- ----------- - -- - -- Check whether we can eliminate the let-binding. Note that we must inspect - -- the entire term, not just the Cunctation that would be produced by - -- embedAcc. If we don't we can be left with dead terms that don't get - -- eliminated. This problem occurred in the canny program. - -- - | acc1 <- computeAcc (Embed env1 cc1) - , False <- elimAcc acc1 acc0 - = Stats.ruleFired "aletD/bind" - $ Embed (BaseEnv `pushArrayEnv` acc1 `append` env0) cc0 - - -- let-elimination - -- --------------- - -- - -- Handle the remaining cases in a separate function. It turns out that this - -- is important so we aren't excessively sinking/delaying terms. - -- - | acc0' <- sink1 env1 acc0 - = Stats.ruleFired "aletD/eliminate" - $ case delaying cc1 of - Step{} -> eliminate env1 cc1 acc0' - Yield{} -> eliminate env1 cc1 acc0' - - where - acc0 :: OpenAcc aenv' brrs - acc0 = computeAcc (Embed env0 cc0) - - kmap :: forall aenv a b. (PreOpenAcc OpenAcc aenv a -> PreOpenAcc OpenAcc aenv b) - -> OpenAcc aenv a - -> OpenAcc aenv b - kmap f (OpenAcc pacc) = OpenAcc (f pacc) - - -- The second part of let-elimination. Splitting into two steps exposes the - -- extra type variables, and ensures we don't do extra work manipulating the - -- body when not necessary (which can lead to a complexity blowup). - -- - eliminate - :: forall r aenv aenv' sh e brrs. HasCallStack - => Extend ArrayR OpenAcc aenv aenv' - -> Cunctation r aenv' (Array sh e) - -> OpenAcc (aenv', Array sh e) brrs - -> Embed OpenAcc aenv brrs - eliminate env1 cc1 body - | Done v1 <- cc1 - , TupRsingle v1'@(Var r _) <- v1 = elim r (arrayShape v1') (indexArray v1') - | Step r sh1 p1 f1 v1 <- cc1 = elim r sh1 (f1 `compose` indexArray v1 `compose` p1) - | Yield r sh1 f1 <- cc1 = elim r sh1 f1 - where - bnd :: PreOpenAcc OpenAcc aenv' (Array sh e) - bnd = compute cc1 - - elim :: HasCallStack - => ArrayR (Array sh e) - -> Exp aenv' sh - -> Fun aenv' (sh -> e) - -> Embed OpenAcc aenv brrs - elim r sh1 f1 - | sh1' <- weaken (weakenSucc' weakenId) sh1 - , f1' <- weaken (weakenSucc' weakenId) f1 - , Embed env0' cc0' <- embedAcc $ rebuildA (subAtop bnd) $ kmap (replaceA sh1' f1' $ Var r ZeroIdx) body - = Embed (env1 `append` env0') cc0' - - -- As part of let-elimination, we need to replace uses of array variables in - -- scalar expressions with an equivalent expression that generates the - -- result directly - -- - -- TODO: when we inline bindings we ought to let bind at the first - -- occurrence and use a variable at all subsequent locations. At the - -- moment we are just hoping CSE in the simplifier phase does good - -- things, but that is limited in what it looks for. - -- - replaceE :: forall env aenv sh e t. HasCallStack - => OpenExp env aenv sh - -> OpenFun env aenv (sh -> e) - -> ArrayVar aenv (Array sh e) - -> OpenExp env aenv t - -> OpenExp env aenv t - replaceE sh' f' avar@(Var (ArrayR shR _) _) exp = - case exp of - Let lhs x y -> let k = weakenWithLHS lhs - in Let lhs (cvtE x) (replaceE (weakenE k sh') (weakenE k f') avar y) - Evar var -> Evar var - Foreign tR ff f e -> Foreign tR ff f (cvtE e) - Const tR c -> Const tR c - Undef tR -> Undef tR - Nil -> Nil - Pair e1 e2 -> Pair (cvtE e1) (cvtE e2) - VecPack vR e -> VecPack vR (cvtE e) - VecUnpack vR e -> VecUnpack vR (cvtE e) - IndexSlice x ix sh -> IndexSlice x (cvtE ix) (cvtE sh) - IndexFull x ix sl -> IndexFull x (cvtE ix) (cvtE sl) - ToIndex shR' sh ix -> ToIndex shR' (cvtE sh) (cvtE ix) - FromIndex shR' sh i -> FromIndex shR' (cvtE sh) (cvtE i) - Case e rhs def -> Case (cvtE e) (over (mapped . _2) cvtE rhs) (fmap cvtE def) - Cond p t e -> Cond (cvtE p) (cvtE t) (cvtE e) - PrimConst c -> PrimConst c - PrimApp g x -> PrimApp g (cvtE x) - ShapeSize shR' sh -> ShapeSize shR' (cvtE sh) - While p f x -> While (replaceF sh' f' avar p) (replaceF sh' f' avar f) (cvtE x) - Coerce t1 t2 e -> Coerce t1 t2 (cvtE e) - - Shape a - | Just Refl <- matchVar a avar -> Stats.substitution "replaceE/shape" sh' - | otherwise -> exp - - Index a sh - | Just Refl <- matchVar a avar - , Lam lhs (Body b) <- f' -> Stats.substitution "replaceE/!" . cvtE $ Let lhs sh b - | otherwise -> Index a (cvtE sh) - - LinearIndex a i - | Just Refl <- matchVar a avar - , Lam lhs (Body b) <- f' - -> Stats.substitution "replaceE/!!" . cvtE - $ Let lhs - (Let (LeftHandSideSingle scalarTypeInt) i - $ FromIndex shR (weakenE (weakenSucc' weakenId) sh') - $ Evar - $ Var scalarTypeInt ZeroIdx) - b - | otherwise -> LinearIndex a (cvtE i) - - where - cvtE :: OpenExp env aenv s -> OpenExp env aenv s - cvtE = replaceE sh' f' avar - - replaceF :: forall env aenv sh e t. HasCallStack - => OpenExp env aenv sh - -> OpenFun env aenv (sh -> e) - -> ArrayVar aenv (Array sh e) - -> OpenFun env aenv t - -> OpenFun env aenv t - replaceF sh' f' avar fun = - case fun of - Body e -> Body (replaceE sh' f' avar e) - Lam lhs f -> let k = weakenWithLHS lhs - in Lam lhs (replaceF (weakenE k sh') (weakenE k f') avar f) - - replaceA :: forall aenv sh e a. HasCallStack - => Exp aenv sh - -> Fun aenv (sh -> e) - -> ArrayVar aenv (Array sh e) - -> PreOpenAcc OpenAcc aenv a - -> PreOpenAcc OpenAcc aenv a - replaceA sh' f' avar pacc = - case pacc of - Avar v - | Just Refl <- matchVar v avar -> Avar avar - | otherwise -> Avar v - - Alet lhs bnd (body :: OpenAcc aenv1 a) -> - let w :: aenv :> aenv1 - w = weakenWithLHS lhs - sh'' = weaken w sh' - f'' = weaken w f' - in - Alet lhs (cvtA bnd) (kmap (replaceA sh'' f'' (weaken w avar)) body) - - Use repr arrs -> Use repr arrs - Unit tR e -> Unit tR (cvtE e) - Acond p at ae -> Acond (cvtE p) (cvtA at) (cvtA ae) - Anil -> Anil - Atrace msg a b -> Atrace msg (cvtA a) (cvtA b) - Apair a1 a2 -> Apair (cvtA a1) (cvtA a2) - Awhile p f a -> Awhile (cvtAF p) (cvtAF f) (cvtA a) - Apply repr f a -> Apply repr (cvtAF f) (cvtA a) - Aforeign repr ff f a -> Aforeign repr ff f (cvtA a) -- no sharing between f and a - Generate repr sh f -> Generate repr (cvtE sh) (cvtF f) - Map tR f a -> Map tR (cvtF f) (cvtA a) - ZipWith tR f a b -> ZipWith tR (cvtF f) (cvtA a) (cvtA b) - Backpermute shR sh p a -> Backpermute shR (cvtE sh) (cvtF p) (cvtA a) - Transform repr sh p f a -> Transform repr (cvtE sh) (cvtF p) (cvtF f) (cvtA a) - Slice slix a sl -> Slice slix (cvtA a) (cvtE sl) - Replicate slix sh a -> Replicate slix (cvtE sh) (cvtA a) - Reshape shR sl a -> Reshape shR (cvtE sl) (cvtA a) - Fold f z a -> Fold (cvtF f) (cvtE <$> z) (cvtA a) - FoldSeg i f z a s -> FoldSeg i (cvtF f) (cvtE <$> z) (cvtA a) (cvtA s) - Scan d f z a -> Scan d (cvtF f) (cvtE <$> z) (cvtA a) - Scan' d f z a -> Scan' d (cvtF f) (cvtE z) (cvtA a) - Permute f d p a -> Permute (cvtF f) (cvtA d) (cvtF p) (cvtA a) - Stencil s t f x a -> Stencil s t (cvtF f) (cvtB x) (cvtA a) - Stencil2 s1 s2 t f x a y b - -> Stencil2 s1 s2 t (cvtF f) (cvtB x) (cvtA a) (cvtB y) (cvtA b) - -- Collect seq -> Collect (cvtSeq seq) - - where - cvtA :: OpenAcc aenv s -> OpenAcc aenv s - cvtA = kmap (replaceA sh' f' avar) - - cvtE :: Exp aenv s -> Exp aenv s - cvtE = replaceE sh' f' avar - - cvtF :: Fun aenv s -> Fun aenv s - cvtF = replaceF sh' f' avar - - cvtB :: Boundary aenv s -> Boundary aenv s - cvtB Clamp = Clamp - cvtB Mirror = Mirror - cvtB Wrap = Wrap - cvtB (Constant c) = Constant c - cvtB (Function f) = Function (cvtF f) - - cvtAF :: HasCallStack => PreOpenAfun OpenAcc aenv s -> PreOpenAfun OpenAcc aenv s - cvtAF = cvt sh' f' avar - where - cvt :: forall aenv a. - Exp aenv sh -> Fun aenv (sh -> e) -> ArrayVar aenv (Array sh e) - -> PreOpenAfun OpenAcc aenv a - -> PreOpenAfun OpenAcc aenv a - cvt sh'' f'' avar' (Abody a) = Abody $ kmap (replaceA sh'' f'' avar') a - cvt sh'' f'' avar' (Alam lhs (af :: PreOpenAfun OpenAcc aenv1 b)) = - Alam lhs $ cvt (weaken w sh'') - (weaken w f'') - (weaken w avar') - af - where - w :: aenv :> aenv1 - w = weakenWithLHS lhs - --- Do not fuse bindings of multiple variables -aletD' _ _ lhs (Embed env1 cc1) (Embed env0 cc0) - = Stats.ruleFired "aletD/bind" - $ Embed (PushEnv BaseEnv lhs (computeAcc (Embed env1 cc1)) `append` env0) cc0 - -{-- - cvtSeq :: PreOpenSeq acc aenv senv s -> PreOpenSeq acc aenv senv s - cvtSeq s = - case s of - Producer p s' -> - Producer - (case p of - StreamIn arrs -> StreamIn arrs - ToSeq slix sh a -> ToSeq slix sh (cvtA a) - MapSeq f x -> MapSeq (cvtAF f) x - ChunkedMapSeq f x -> ChunkedMapSeq (cvtAF f) x - ZipWithSeq f x y -> ZipWithSeq (cvtAF f) x y - ScanSeq f e x -> ScanSeq (cvtF f) (cvtE e) x) - (cvtSeq s') - Consumer c -> - Consumer (cvtC c) - Reify ix -> Reify ix - - cvtC :: Consumer acc aenv senv s -> Consumer acc aenv senv s - cvtC c = - case c of - FoldSeq f e x -> FoldSeq (cvtF f) (cvtE e) x - FoldSeqFlatten f a x -> FoldSeqFlatten (cvtAF f) (cvtA a) x - Stuple t -> Stuple (cvtCT t) - - cvtCT :: Atuple (Consumer acc aenv senv) t -> Atuple (Consumer acc aenv senv) t - cvtCT NilAtup = NilAtup - cvtCT (SnocAtup t c) = cvtCT t `SnocAtup` cvtC c ---} - - --- Array conditionals, in particular eliminate branches when the predicate --- reduces to a known constant. --- --- Note that we take the raw unprocessed terms as input. If instead we had the --- terms for each branch in the delayed representation, this would require that --- each term has been sunk into a common environment, which implies the --- conditional has been pushed underneath the intersection of bound terms for --- both branches. This would result in redundant work processing the bindings --- for the branch not taken. --- -acondD :: HasCallStack - => MatchAcc OpenAcc - -> EmbedAcc OpenAcc - -> Exp aenv PrimBool - -> OpenAcc aenv arrs - -> OpenAcc aenv arrs - -> Embed OpenAcc aenv arrs -acondD matchAcc embedAcc p t e - | Const _ 1 <- p = Stats.knownBranch "True" $ embedAcc t - | Const _ 0 <- p = Stats.knownBranch "False" $ embedAcc e - | Just Refl <- matchAcc t e = Stats.knownBranch "redundant" $ embedAcc e - | otherwise = done $ Acond p (computeAcc (embedAcc t)) - (computeAcc (embedAcc e)) - - --- Scalar expressions --- ------------------ - -identity :: TypeR a -> OpenFun env aenv (a -> a) -identity t - | DeclareVars lhs _ value <- declareVars t - = Lam lhs $ Body $ expVars $ value weakenId - -toIndex :: ShapeR sh -> OpenExp env aenv sh -> OpenFun env aenv (sh -> Int) -toIndex shR sh - | DeclareVars lhs k value <- declareVars $ shapeType shR - = Lam lhs $ Body $ ToIndex shR (weakenE k sh) $ expVars $ value weakenId - -fromIndex :: ShapeR sh -> OpenExp env aenv sh -> OpenFun env aenv (Int -> sh) -fromIndex shR sh - = Lam (LeftHandSideSingle scalarTypeInt) - $ Body - $ FromIndex shR (weakenE (weakenSucc' weakenId) sh) - $ Evar - $ Var scalarTypeInt ZeroIdx - -intersect :: ShapeR sh -> OpenExp env aenv sh -> OpenExp env aenv sh -> OpenExp env aenv sh -intersect = mkShapeBinary f - where - f a b = PrimApp (PrimMin singleType) $ Pair a b - --- union :: ShapeR sh -> OpenExp env aenv sh -> OpenExp env aenv sh -> OpenExp env aenv sh --- union = mkShapeBinary f +-- cvtC :: Consumer OpenAcc aenv senv a -> Consumer DelayedOpenAcc aenv senv a +-- cvtC c = +-- case c of +-- FoldSeq f e x -> FoldSeq (cvtF f) (cvtE e) x +-- FoldSeqFlatten f a x -> FoldSeqFlatten (cvtAF f) (manifest config a) x +-- Stuple t -> Stuple (cvtCT t) + +-- cvtCT :: Atuple (Consumer OpenAcc aenv senv) t -> Atuple (Consumer DelayedOpenAcc aenv senv) t +-- cvtCT NilAtup = NilAtup +-- cvtCT (SnocAtup t c) = SnocAtup (cvtCT t) (cvtC c) + +-- cvtAF :: OpenAfun aenv f -> PreOpenAfun DelayedOpenAcc aenv f +-- cvtAF (Alam f) = Alam (cvtAF f) +-- cvtAF (Abody b) = Abody (manifest config b) + +-- cvtE :: OpenExp env aenv t -> DelayedOpenExp env aenv t +-- cvtE = convertOpenExp config + +-- cvtF :: OpenFun env aenv f -> DelayedOpenFun env aenv f +-- cvtF (Lam f) = Lam (cvtF f) +-- cvtF (Body b) = Body (cvtE b) +-- --} + + +-- type EmbedAcc acc = forall aenv arrs. acc aenv arrs -> Embed acc aenv arrs +-- type ElimAcc acc = forall aenv s t. acc aenv s -> acc (aenv,s) t -> Bool + + +-- -- | Apply the fusion transformation to the AST to combine and simplify terms. +-- -- This converts terms into the internal delayed array representation and merges +-- -- adjacent producer/producer terms. Using the reduced internal form limits the +-- -- number of combinations that need to be considered. +-- -- +-- embedOpenAcc :: HasCallStack => Config -> OpenAcc aenv arrs -> Embed OpenAcc aenv arrs +-- embedOpenAcc config (OpenAcc pacc) = +-- embedPreOpenAcc config matchOpenAcc (embedOpenAcc config) elimOpenAcc pacc -- where --- f a b = PrimApp (PrimMax singleType) $ Pair a b - -mkShapeBinary - :: (forall env'. OpenExp env' aenv Int -> OpenExp env' aenv Int -> OpenExp env' aenv Int) - -> ShapeR sh - -> OpenExp env aenv sh - -> OpenExp env aenv sh - -> OpenExp env aenv sh -mkShapeBinary _ ShapeRz _ _ = Nil -mkShapeBinary f (ShapeRsnoc shR) (Pair as a) (Pair bs b) = mkShapeBinary f shR as bs `Pair` f a b -mkShapeBinary f shR (Let lhs bnd a) b = Let lhs bnd $ mkShapeBinary f shR a (weakenE (weakenWithLHS lhs) b) -mkShapeBinary f shR a (Let lhs bnd b) = Let lhs bnd $ mkShapeBinary f shR (weakenE (weakenWithLHS lhs) a) b -mkShapeBinary f shR a b@Pair{} -- `a` is not Pair - | DeclareVars lhs k value <- declareVars $ shapeType shR - = Let lhs a $ mkShapeBinary f shR (expVars $ value weakenId) (weakenE k b) -mkShapeBinary f shR a b -- `b` is not a Pair - | DeclareVars lhs k value <- declareVars $ shapeType shR - = Let lhs b $ mkShapeBinary f shR (weakenE k a) (expVars $ value weakenId) - -reindex :: ShapeR sh' - -> OpenExp env aenv sh' - -> ShapeR sh - -> OpenExp env aenv sh - -> OpenFun env aenv (sh -> sh') -reindex shR' sh' shR sh - | Just Refl <- matchOpenExp sh sh' = identity (shapeType shR') - | otherwise = fromIndex shR' sh' `compose` toIndex shR sh - -extend :: SliceIndex slix sl co sh - -> Exp aenv slix - -> Fun aenv (sh -> sl) -extend sliceIndex slix - | DeclareVars lhs k value <- declareVars $ shapeType $ sliceDomainR sliceIndex - = Lam lhs $ Body $ IndexSlice sliceIndex (weakenE k slix) $ expVars $ value weakenId - -restrict :: SliceIndex slix sl co sh - -> Exp aenv slix - -> Fun aenv (sl -> sh) -restrict sliceIndex slix - | DeclareVars lhs k value <- declareVars $ shapeType $ sliceShapeR sliceIndex - = Lam lhs $ Body $ IndexFull sliceIndex (weakenE k slix) $ expVars $ value weakenId - -arrayShape :: ArrayVar aenv (Array sh e) -> Exp aenv sh -arrayShape = simplifyExp . Shape - -indexArray :: ArrayVar aenv (Array sh e) -> Fun aenv (sh -> e) -indexArray v@(Var (ArrayR shR _) _) - | DeclareVars lhs _ value <- declareVars $ shapeType shR - = Lam lhs $ Body $ Index v $ expVars $ value weakenId - -linearIndex :: ArrayVar aenv (Array sh e) -> Fun aenv (Int -> e) -linearIndex v = Lam (LeftHandSideSingle scalarTypeInt) $ Body $ LinearIndex v $ Evar $ Var scalarTypeInt ZeroIdx - - -extractOpenAcc :: ExtractAcc OpenAcc -extractOpenAcc (OpenAcc pacc) = Just pacc - -extractDelayedOpenAcc :: ExtractAcc DelayedOpenAcc -extractDelayedOpenAcc (Manifest pacc) = Just pacc -extractDelayedOpenAcc _ = Nothing - -extractOpenArrayVars - :: OpenAcc aenv a - -> Maybe (ArrayVars aenv a) -extractOpenArrayVars (OpenAcc pacc) = - avarsOut extractOpenAcc pacc - -extractDelayedArrayVars - :: DelayedOpenAcc aenv a - -> Maybe (ArrayVars aenv a) -extractDelayedArrayVars acc - | Just pacc <- extractDelayedOpenAcc acc = avarsOut extractDelayedOpenAcc pacc - | otherwise = Nothing +-- -- When does the cost of re-computation outweigh that of memory access? For +-- -- the moment only do the substitution on a single use of the bound array +-- -- into the use site, but it is likely advantageous to be far more +-- -- aggressive here. +-- -- +-- -- SEE: [Sharing vs. Fusion] +-- -- +-- elimOpenAcc :: ElimAcc OpenAcc +-- elimOpenAcc _bnd body +-- | count False ZeroIdx body <= lIMIT = True +-- | otherwise = False +-- where +-- lIMIT = 1 + +-- count :: UsesOfAcc OpenAcc +-- count no ix (OpenAcc pacc) = usesOfPreAcc no count ix pacc + +-- matchOpenAcc :: MatchAcc OpenAcc +-- matchOpenAcc (OpenAcc pacc1) (OpenAcc pacc2) = +-- matchPreOpenAcc matchOpenAcc pacc1 pacc2 + + +-- embedPreOpenAcc +-- :: HasCallStack +-- => Config +-- -> MatchAcc OpenAcc +-- -> EmbedAcc OpenAcc +-- -> ElimAcc OpenAcc +-- -> PreOpenAcc OpenAcc aenv arrs +-- -> Embed OpenAcc aenv arrs +-- embedPreOpenAcc config matchAcc embedAcc elimAcc pacc +-- = unembed +-- $ case pacc of + +-- -- Non-fusible terms +-- -- ----------------- +-- -- +-- -- Solid and semi-solid terms that we generally do not wish to fuse, such +-- -- as control flow (|?), array introduction (use, unit), array tupling and +-- -- projection, and foreign function operations. Generally we also do not +-- -- want to fuse past array let bindings, as this would imply work +-- -- duplication. SEE: [Sharing vs. Fusion] +-- -- +-- Alet lhs bnd body -> aletD embedAcc elimAcc lhs bnd body +-- Anil -> done $ Anil +-- Acond p at ae -> acondD matchAcc embedAcc (cvtE p) at ae +-- Apply aR f a -> done $ Apply aR (cvtAF f) (cvtA a) +-- Awhile p f a -> done $ Awhile (cvtAF p) (cvtAF f) (cvtA a) +-- Apair a1 a2 -> done $ Apair (cvtA a1) (cvtA a2) +-- Atrace msg a1 a2 -> done $ Atrace msg (cvtA a1) (cvtA a2) +-- Aforeign aR ff f a -> done $ Aforeign aR ff (cvtAF f) (cvtA a) +-- -- Collect s -> collectD s + +-- -- Array injection +-- Avar v -> done $ Avar v +-- Use aR a -> done $ Use aR a +-- Unit t e -> done $ Unit t (cvtE e) + +-- -- Producers +-- -- --------- +-- -- +-- -- The class of operations that given a set of zero or more input arrays, +-- -- produce a _single_ element for the output array by manipulating a +-- -- _single_ element from each input array. These can be further classified +-- -- as value (map, zipWith) or index space (backpermute, slice, replicate) +-- -- transformations. +-- -- +-- -- The critical feature is that each element of the output is produced +-- -- independently of all others, and so we can aggressively fuse arbitrary +-- -- sequences of these operations. +-- -- +-- Generate aR sh f -> generateD aR (cvtE sh) (cvtF f) +-- Map t f a -> mapD t (cvtF f) (embedAcc a) +-- ZipWith t f a b -> fuse2 (into (zipWithD t) (cvtF f)) a b +-- Transform aR sh p f a -> transformD aR (cvtE sh) (cvtF p) (cvtF f) (embedAcc a) +-- Backpermute slr sl p a -> fuse (into2 (backpermuteD slr) (cvtE sl) (cvtF p)) a +-- Slice slix a sl -> fuse (into (sliceD slix) (cvtE sl)) a +-- Replicate slix sh a -> fuse (into (replicateD slix) (cvtE sh)) a +-- Reshape slr sl a -> reshapeD slr (embedAcc a) (cvtE sl) + +-- -- Consumers +-- -- --------- +-- -- +-- -- Operations where each element of the output array depends on multiple +-- -- elements of the input array. To implement these operations efficiently in +-- -- parallel, we need to know how elements of the array depend on each other: +-- -- a parallel scan is implemented very differently from a parallel fold, for +-- -- example. +-- -- +-- -- In order to avoid obfuscating this crucial information required for +-- -- parallel implementation, fusion is separated into to phases: +-- -- producer/producer, implemented above, and consumer/producer, which is +-- -- implemented below. This will place producers adjacent to the consumer +-- -- node, so that the producer can be directly embedded into the consumer +-- -- during the code generation phase. +-- -- +-- Fold f z a -> embed aR (into2M Fold (cvtF f) (cvtE <$> z)) a +-- FoldSeg i f z a s -> embed2 aR (into2M (FoldSeg i) (cvtF f) (cvtE <$> z)) a s +-- Scan d f z a -> embed aR (into2M (Scan d) (cvtF f) (cvtE <$> z)) a +-- Scan' d f z a -> embed aR (into2 (Scan' d) (cvtF f) (cvtE z)) a +-- Permute f d p a -> embed2 aR (into2 permute (cvtF f) (cvtF p)) d a +-- Stencil s t f x a -> embed aR (into2B (stencil1 s t) (cvtF f) (cvtB x)) a +-- Stencil2 s1 s2 t f x a y b +-- -> embed2 aR (into3B (stencil2 s1 s2 t) (cvtF f) (cvtB x) (cvtB y)) a b + +-- where +-- aR = arraysR pacc + +-- -- If fusion is not enabled, force terms to the manifest representation +-- -- +-- unembed :: HasCallStack => Embed OpenAcc aenv arrs -> Embed OpenAcc aenv arrs +-- unembed x +-- | array_fusion `member` options config = x +-- | Embed env cc <- x +-- , pacc <- compute cc +-- = case avarsOut extractOpenAcc pacc of +-- Just vars -> Embed env $ Done vars +-- _ +-- | DeclareVars lhs _ value <- declareVars (arraysR pacc) +-- -> Embed (PushEnv env lhs $ OpenAcc pacc) $ Done $ value weakenId + +-- cvtA :: HasCallStack => OpenAcc aenv' a -> OpenAcc aenv' a +-- cvtA = computeAcc . embedAcc + +-- cvtAF :: HasCallStack => PreOpenAfun OpenAcc aenv' f -> PreOpenAfun OpenAcc aenv' f +-- cvtAF (Alam lhs f) = Alam lhs (cvtAF f) +-- cvtAF (Abody a) = Abody (cvtA a) + +-- -- Helpers to shuffle the order of arguments to a constructor +-- -- +-- permute f p d a = Permute f d p a + +-- -- NOTE: [Stencil fusion] +-- -- +-- -- We allow stencils to delay their argument arrays with no special +-- -- considerations. This means that the delayed function will be evaluated +-- -- _at every element_ of the stencil pattern. We should do some analysis of +-- -- when this duplication is beneficial (keeping in mind that the stencil +-- -- implementations themselves may share neighbouring elements). +-- -- +-- stencil1 s t f x a = Stencil s t f x a +-- stencil2 s1 s2 t f x y a b = Stencil2 s1 s2 t f x a y b + +-- -- Conversions for closed scalar functions and expressions. This just +-- -- applies scalar simplifications. +-- -- +-- cvtF :: HasCallStack => Fun aenv' t -> Fun aenv' t +-- cvtF = simplifyFun + +-- cvtE :: HasCallStack => Exp aenv' t -> Exp aenv' t +-- cvtE = simplifyExp + +-- cvtB :: HasCallStack => Boundary aenv' t -> Boundary aenv' t +-- cvtB Clamp = Clamp +-- cvtB Mirror = Mirror +-- cvtB Wrap = Wrap +-- cvtB (Constant c) = Constant c +-- cvtB (Function f) = Function (cvtF f) + +-- -- Helpers to embed and fuse delayed terms +-- -- +-- into :: (HasCallStack, RebuildableExp f) +-- => (f (ArrayInstr env') () a -> b) +-- -> f (ArrayInstr env) () a +-- -> Extend ArrayR OpenAcc env env' +-- -> b +-- into op a env = op (sinkArrayInstr env a) + +-- into2 :: (HasCallStack, RebuildableExp f1, RebuildableExp f2) +-- => (f1 (ArrayInstr env') () a -> f2 (ArrayInstr env') () b -> c) +-- -> f1 (ArrayInstr env) () a +-- -> f2 (ArrayInstr env) () b +-- -> Extend ArrayR OpenAcc env env' +-- -> c +-- into2 op a b env = op (sinkArrayInstr env a) (sinkArrayInstr env b) + +-- into2B :: (HasCallStack, RebuildableExp f1, Sink f2) +-- => (f1 (ArrayInstr env') () a -> f2 env' b -> c) +-- -> f1 (ArrayInstr env) () a +-- -> f2 env b +-- -> Extend ArrayR OpenAcc env env' +-- -> c +-- into2B op a b env = op (sinkArrayInstr env a) (sinkA env b) + +-- into2M :: (HasCallStack, RebuildableExp f1, RebuildableExp f2) +-- => (f1 (ArrayInstr env') () a -> Maybe (f2 (ArrayInstr env') () b) -> c) +-- -> f1 (ArrayInstr env) () a +-- -> Maybe (f2 (ArrayInstr env) () b) +-- -> Extend ArrayR acc env env' +-- -> c +-- into2M op a b env = op (sinkArrayInstr env a) (sinkArrayInstr env <$> b) + +-- into3B :: (HasCallStack, RebuildableExp f1, Sink f2, Sink f3) +-- => (f1 (ArrayInstr env') () a -> f2 env' b -> f3 env' c -> d) +-- -> f1 (ArrayInstr env) () a +-- -> f2 env b +-- -> f3 env c +-- -> Extend ArrayR OpenAcc env env' +-- -> d +-- into3B op a b c env = op (sinkArrayInstr env a) (sinkA env b) (sinkA env c) + +-- -- Operations which can be fused into consumers. Move all of the local +-- -- bindings out of the way so that the fusible function operates +-- -- directly on the delayed representation. See also: [Representing +-- -- delayed arrays] +-- -- +-- fuse :: HasCallStack +-- => (forall r aenv'. Extend ArrayR OpenAcc aenv aenv' -> Cunctation r aenv' as -> Cunctation D aenv' bs) +-- -> OpenAcc aenv as +-- -> Embed OpenAcc aenv bs +-- fuse op (embedAcc -> Embed env cc) = Embed env (op env cc) + +-- fuse2 :: HasCallStack +-- => (forall r s aenv'. Extend ArrayR OpenAcc aenv aenv' -> Cunctation r aenv' as -> Cunctation s aenv' bs -> Cunctation D aenv' cs) +-- -> OpenAcc aenv as +-- -> OpenAcc aenv bs +-- -> Embed OpenAcc aenv cs +-- fuse2 op a1 a0 +-- | Embed env1 cc1 <- embedAcc a1 +-- , Embed env0 cc0 <- embedAcc (sinkA env1 a0) +-- , env <- env1 `append` env0 +-- = Embed env (op env (sinkA env0 cc1) cc0) + +-- -- Consumer operations which will be evaluated. +-- -- +-- -- NOTE: [Fusion and the lowest common use site] +-- -- +-- -- The AST given to us by sharing recovery will place let bindings at +-- -- the lowest common use site for that shared term. For example: +-- -- +-- -- fold f z (let a0 = .. +-- -- a1 = .. +-- -- in zipWith g a0 a1) +-- -- +-- -- In order to enable producer/consumer fusion for the above example, +-- -- it is necessary to float the let bindings above the `fold` +-- -- operation; SEE: [Sharing vs. Fusion] for more information. +-- -- +-- -- Furthermore, we used to maintain an invariant that all (manifest) +-- -- arguments were supplied as array variables, for example: +-- -- +-- -- fold1 f (let a0 = .. let a0 = .. +-- -- in stencil g a0) ==> a1 = stencil g a0 +-- -- a2 = fold1 f a1 +-- -- +-- -- However, if the argument term will be evaluated (i.e. can not be +-- -- fused into the producer) then it is better that we do _not_ float +-- -- those terms, and instead leave them under the consumer. This helps +-- -- to syntactically constrain the "liveness" of terms: if the argument +-- -- to an operation is not an array variable, we can see directly that +-- -- this will be the last use-site of that array. In particular, this is +-- -- useful for the 'permute' operation to know when it can in-place +-- -- update the array of default values. +-- -- +-- embed :: HasCallStack +-- => ArraysR bs +-- -> (forall aenv'. Extend ArrayR OpenAcc aenv aenv' -> OpenAcc aenv' as -> PreOpenAcc OpenAcc aenv' bs) +-- -> OpenAcc aenv as +-- -> Embed OpenAcc aenv bs +-- embed reprBs op (embedAcc -> Embed env cc) +-- | Done{} <- cc +-- , DeclareVars lhs _ value <- declareVars reprBs +-- = Embed (PushEnv BaseEnv lhs $ OpenAcc (op BaseEnv (computeAcc (Embed env cc)))) $ Done $ value weakenId +-- | otherwise +-- -- Next line is duplicated for both branches, as the type variable for the environment is instantiated differently +-- , DeclareVars lhs _ value <- declareVars reprBs +-- = Embed (PushEnv env lhs $ OpenAcc (op env (OpenAcc (compute cc)))) $ Done $ value weakenId + +-- embed2 :: HasCallStack +-- => ArraysR cs +-- -> (forall aenv'. Extend ArrayR OpenAcc aenv aenv' -> OpenAcc aenv' as -> OpenAcc aenv' bs -> PreOpenAcc OpenAcc aenv' cs) +-- -> OpenAcc aenv as +-- -> OpenAcc aenv bs +-- -> Embed OpenAcc aenv cs +-- embed2 reprCs op (embedAcc -> Embed env1 cc1) a0 +-- | Done{} <- cc1 +-- , a1 <- computeAcc (Embed env1 cc1) +-- = embed reprCs (\env0 -> op env0 (sinkA env0 a1)) a0 +-- -- +-- | Embed env0 cc0 <- embedAcc (sinkA env1 a0) +-- , env <- env1 `append` env0 +-- = case cc0 of +-- Done{} +-- | DeclareVars lhs _ value <- declareVars reprCs +-- -> Embed (PushEnv env1 lhs $ OpenAcc (op env1 (OpenAcc (compute cc1)) (computeAcc (Embed env0 cc0)))) +-- $ Done +-- $ value weakenId +-- _ +-- -- Next line is duplicated for both branches, as the type +-- -- variable for the environment is instantiated differently +-- | DeclareVars lhs _ value <- declareVars reprCs +-- -> Embed (PushEnv env lhs $ OpenAcc (op env (OpenAcc (compute (sinkA env0 cc1))) (OpenAcc (compute cc0)))) +-- $ Done +-- $ value weakenId + +-- -- trav1 :: (Arrays as, Arrays bs) +-- -- => (forall aenv'. Embed acc aenv' as -> Embed acc aenv' as) +-- -- -> (forall aenv'. Extend ArrayR acc aenv aenv' -> acc aenv' as -> PreOpenAcc acc aenv' bs) +-- -- -> acc aenv as +-- -- -> Embed acc aenv bs +-- -- trav1 f op (f . embedAcc -> Embed env cc) +-- -- = Embed (env `pushArrayEnv` inject (op env (inject (compute cc)))) doneZeroIdx + +-- -- trav2 :: (Arrays as, Arrays bs, Arrays cs) +-- -- => (forall aenv'. Embed acc aenv' as -> Embed acc aenv' as) +-- -- -> (forall aenv'. Embed acc aenv' bs -> Embed acc aenv' bs) +-- -- -> (forall aenv'. Extend ArrayR acc aenv aenv' -> acc aenv' as -> acc aenv' bs -> PreOpenAcc acc aenv' cs) +-- -- -> acc aenv as +-- -- -> acc aenv bs +-- -- -> Embed acc aenv cs +-- -- trav2 f1 f0 op (f1 . embedAcc -> Embed env1 cc1) (f0 . embedAcc . sinkA env1 -> Embed env0 cc0) +-- -- | env <- env1 `append` env0 +-- -- , acc1 <- inject . compute $ sinkA env0 cc1 +-- -- , acc0 <- inject . compute $ cc0 +-- -- = Embed (env `pushArrayEnv` inject (op env acc1 acc0)) doneZeroIdx + +-- -- force :: Arrays as => Embed acc aenv' as -> Embed acc aenv' as +-- -- force (Embed env cc) +-- -- | Done{} <- cc = Embed env cc +-- -- | otherwise = Embed (env `pushArrayEnv` inject (compute cc)) doneZeroIdx + +-- -- -- Move additional bindings for producers outside of the sequence, so that +-- -- -- producers may fuse with their arguments resulting in actual sequencing +-- -- collectD :: PreOpenSeq acc aenv () arrs +-- -- -> Embed acc aenv arrs +-- -- collectD (embedSeq embedAcc -> ExtendSeq env s') +-- -- = Embed (env `pushArrayEnv` inject (Collect s')) doneZeroIdx + + +-- {-- +-- -- Move additional bindings for producer outside of sequence, so +-- -- that producers may fuse with their arguments, resulting in +-- -- actual sequencing. +-- embedSeq :: forall acc aenv arrs. Kit acc +-- => EmbedAcc acc +-- -> PreOpenSeq acc aenv () arrs +-- -> ExtendSeq acc aenv () arrs +-- embedSeq embedAcc s +-- = travS s BaseEnv +-- where +-- travS :: forall senv aenv' arrs'. +-- PreOpenSeq acc aenv senv arrs' +-- -> Extend acc aenv aenv' +-- -> ExtendSeq acc aenv senv arrs' +-- travS s env = +-- case s of +-- Producer p s +-- | ExtendSeq env' s' <- travS s env +-- , ExtendProducer env'' p' <- travP p env' +-- -> ExtendSeq (env' `append` env'') (Producer p' (sinkSeq env'' s')) +-- Consumer c +-- | c' <- travC c env +-- -> ExtendSeq env (Consumer c') +-- Reify ix +-- -> ExtendSeq env (Reify ix) + +-- travP :: forall arrs' aenv' senv. +-- Producer acc aenv senv arrs' +-- -> Extend acc aenv aenv' +-- -> ExtendProducer acc aenv' senv arrs' +-- travP (ToSeq slix sh a) env +-- | Embed env' cc <- embedAcc (sink env a) +-- = ExtendProducer env' (ToSeq slix sh (inject (compute cc))) +-- travP (StreamIn arrs) _ = ExtendProducer BaseEnv (StreamIn arrs) +-- travP (MapSeq f x) env = ExtendProducer BaseEnv (MapSeq (cvtAF (sink env f)) x) +-- travP (ChunkedMapSeq f x) env = ExtendProducer BaseEnv (ChunkedMapSeq (cvtAF (sink env f)) x) +-- travP (ZipWithSeq f x y) env = ExtendProducer BaseEnv (ZipWithSeq (cvtAF (sink env f)) x y) +-- travP (ScanSeq f e x) env = ExtendProducer BaseEnv (ScanSeq (cvtF (sink env f)) (cvtE (sink env e)) x) + +-- travC :: forall arrs' aenv' senv. +-- Consumer acc aenv senv arrs' +-- -> Extend acc aenv aenv' +-- -> Consumer acc aenv' senv arrs' +-- travC (FoldSeq f e x) env = FoldSeq (cvtF (sink env f)) (cvtE (sink env e)) x +-- travC (FoldSeqFlatten f a x) env = FoldSeqFlatten (cvtAF (sink env f)) (cvtA (sink env a)) x +-- travC (Stuple t) env = Stuple (cvtCT t) +-- where +-- cvtCT :: Atuple (Consumer acc aenv senv) t -> Atuple (Consumer acc aenv' senv) t +-- cvtCT NilAtup = NilAtup +-- cvtCT (SnocAtup t c) = SnocAtup (cvtCT t) (travC c env) + +-- cvtE :: Elt t => Exp aenv' t -> Exp aenv' t +-- cvtE = simplifyExp + +-- cvtF :: Fun aenv' t -> Fun aenv' t +-- cvtF = simplifyFun + +-- cvtA :: Arrays a => acc aenv' a -> acc aenv' a +-- cvtA = computeAcc . embedAcc + +-- cvtAF :: PreOpenAfun acc aenv' f -> PreOpenAfun acc aenv' f +-- cvtAF (Alam f) = Alam (cvtAF f) +-- cvtAF (Abody a) = Abody (cvtA a) + + +-- -- A sequence with additional bindings +-- data ExtendSeq acc aenv senv arrs where +-- ExtendSeq :: forall acc aenv aenv' senv arrs. +-- Extend acc aenv aenv' +-- -> PreOpenSeq acc aenv' senv arrs +-- -> ExtendSeq acc aenv senv arrs + +-- -- A producer with additional bindings +-- data ExtendProducer acc aenv senv arrs where +-- ExtendProducer :: forall acc aenv aenv' senv arrs. +-- Extend acc aenv aenv' +-- -> Producer acc aenv' senv arrs +-- -> ExtendProducer acc aenv senv arrs +-- --} + + +-- -- Internal representation +-- -- ======================= + +-- -- NOTE: [Representing delayed arrays] +-- -- +-- -- During the fusion transformation we represent terms as a pair consisting of +-- -- a collection of supplementary environment bindings and a description of how +-- -- to construct the array. +-- -- +-- -- It is critical to separate these two. To create a real AST node we need both +-- -- the environment and array term, but analysis of how to fuse terms requires +-- -- only the array description. If the additional bindings are bundled as part of +-- -- the representation, the existentially quantified extended environment type +-- -- will be untouchable. This is problematic because the terms of the two arrays +-- -- are defined with respect to this existentially quantified type, and there is +-- -- no way to directly combine these two environments: +-- -- +-- -- append :: Extend env env1 -> Extend env env2 -> Extend env ??? +-- -- +-- -- And hence, no way to combine the terms of the delayed representation. +-- -- +-- -- The only way to bring terms into the same scope is to operate via the +-- -- manifest terms. This entails a great deal of conversion between delayed and +-- -- AST terms, but is certainly possible. +-- -- +-- -- However, because of the limited scope into which this existential type is +-- -- available, we ultimately perform this process many times. In fact, complexity +-- -- of the fusion algorithm for an AST of N terms becomes O(r^n), where r is the +-- -- number of different rules we have for combining terms. +-- -- +-- data Embed acc aenv a where +-- Embed :: Extend ArrayR acc aenv aenv' +-- -> Cunctation r aenv' a +-- -> Embed acc aenv a + +-- instance HasArraysR acc => HasArraysR (Embed acc) where +-- arraysR (Embed _ c) = arraysR c + + +-- -- Delayed arrays are represented as functions +-- data D + +-- -- Manifest arrays are real arrays backed by unboxed vectors +-- data M + +-- -- Cunctation (n): the action or an instance of delaying; a tardy action. +-- -- +-- -- This describes the ways in which the fusion transformation represents +-- -- intermediate arrays. The fusion process operates by recasting producer array +-- -- computations in terms of a set of scalar functions used to construct an +-- -- element at each index, and fusing successive producers by combining these +-- -- scalar functions. +-- -- +-- data Cunctation r aenv a where + +-- -- The base case is just a real (manifest) array term. No fusion happens here. +-- -- Note that the array is referenced by an index into the extended +-- -- environment, ensuring that the array is manifest and making the term +-- -- non-recursive in 'acc'. +-- -- +-- Done :: ArrayVars aenv arrs +-- -> Cunctation M aenv arrs + +-- -- We can represent an array by its shape and a function to compute an element +-- -- at each index. +-- -- +-- Yield :: ArrayR (Array sh e) +-- -> Exp aenv sh +-- -> Fun aenv (sh -> e) +-- -> Cunctation D aenv (Array sh e) + +-- -- A more restrictive form than 'Yield' may afford greater opportunities for +-- -- optimisation by a backend. This more structured form applies an index and +-- -- value transform to an input array. Note that the transform is applied to an +-- -- array stored as an environment index, so that the term is non-recursive and +-- -- it is always possible to embed into a collective operation. +-- -- +-- Step :: ArrayR (Array sh' b) +-- -> Exp aenv sh' +-- -> Fun aenv (sh' -> sh) +-- -> Fun aenv (a -> b) +-- -> ArrayVar aenv (Array sh a) +-- -> Cunctation D aenv (Array sh' b) + +-- instance HasArraysR (Cunctation r) where +-- arraysR (Done v) = varsType v +-- arraysR (Yield aR _ _) = TupRsingle aR +-- arraysR (Step aR _ _ _ _) = TupRsingle aR + +-- instance Sink (Cunctation r) where +-- weaken k = \case +-- Done v -> Done (weakenVars k v) +-- Step repr sh p f v -> Step repr (weakenArrayInstr k sh) (weakenArrayInstr k p) (weakenArrayInstr k f) (weaken k v) +-- Yield repr sh f -> Yield repr (weakenArrayInstr k sh) (weakenArrayInstr k f) + +-- simplify +-- :: HasCallStack +-- => Cunctation r aenv a +-- -> Either (Cunctation M aenv a) (Cunctation D aenv a) +-- simplify = \case +-- Done v +-- -> Left $ Done v +-- Yield aR (simplifyExp -> sh) (simplifyFun -> f) +-- -> Right $ Yield aR sh f +-- Step aR (simplifyExp -> sh) (simplifyFun -> p) (simplifyFun -> f) v +-- | Just Refl <- matchOpenExp sh (arrayShape v) +-- , Just Refl <- isIdentity p +-- , Just Refl <- isIdentity f +-- -> Left $ (Done (TupRsingle v)) +-- | otherwise +-- -> Right $ Step aR sh p f v + + +-- -- Convert a real AST node into the internal representation +-- -- +-- done :: HasCallStack => PreOpenAcc OpenAcc aenv a -> Embed OpenAcc aenv a +-- done pacc +-- | Just vars <- avarsOut extractOpenAcc pacc +-- = Embed BaseEnv (Done vars) +-- | DeclareVars lhs _ value <- declareVars (arraysR pacc) +-- = Embed (PushEnv BaseEnv lhs $ OpenAcc pacc) $ Done $ value weakenId + +-- doneZeroIdx :: ArrayR (Array sh e) -> Cunctation M (aenv, Array sh e) (Array sh e) +-- doneZeroIdx aR = Done (TupRsingle (Var aR ZeroIdx)) + + +-- -- Recast a cunctation into a mapping from indices to elements +-- -- +-- yield :: HasCallStack +-- => Cunctation r aenv (Array sh e) +-- -> Cunctation D aenv (Array sh e) +-- yield cc = +-- case cc of +-- Yield{} -> cc +-- Step tR sh p f v -> Yield tR sh (f `compose` indexArray v `compose` p) +-- Done (TupRsingle v@(Var tR _)) -> Yield tR (arrayShape v) (indexArray v) + +-- -- Recast a cunctation into a delayed representation +-- -- +-- delaying +-- :: HasCallStack +-- => Cunctation r aenv (Array sh e) +-- -> Cunctation D aenv (Array sh e) +-- delaying cc = +-- case cc of +-- Yield{} -> cc +-- Step{} -> cc +-- Done u +-- | TupRsingle v <- u +-- , Var aR _ <- v +-- , ArrayR shR tR <- aR +-- -> Step aR (arrayShape v) (identity (shapeType shR)) (identity tR) v + +-- -- Get the shape of a delayed array +-- -- +-- shape :: HasCallStack => Cunctation r aenv (Array sh e) -> Exp aenv sh +-- shape cc = +-- case delaying cc of +-- Step _ sh _ _ _ -> sh +-- Yield _ sh _ -> sh + + +-- -- prjExtend :: Kit acc => Extend acc env env' -> Idx env' t -> PreOpenAcc acc env' t +-- -- prjExtend (PushEnv _ v) ZeroIdx = weakenA rebuildAcc SuccIdx v +-- -- prjExtend (PushEnv env _) (SuccIdx idx) = weakenA rebuildAcc SuccIdx $ prjExtend env idx +-- -- prjExtend _ _ = $internalError "prjExtend" "inconsistent valuation" + +-- {-- +-- -- Rearrange type arguments to fit with Sink type class. +-- newtype SinkSeq acc senv aenv a = SinkSeq { unSinkSeq :: PreOpenSeq acc aenv senv a } + +-- -- sink for sequences. +-- sinkSeq :: Kit acc => Extend acc aenv aenv' -> PreOpenSeq acc aenv senv a -> PreOpenSeq acc aenv' senv a +-- sinkSeq env s = unSinkSeq $ sink env (SinkSeq s) + +-- instance Kit acc => Sink (SinkSeq acc senv) where +-- weaken :: forall aenv aenv' arrs. aenv :> aenv' -> SinkSeq acc senv aenv arrs -> SinkSeq acc senv aenv' arrs +-- weaken k (SinkSeq s) = SinkSeq $ +-- case s of +-- Producer p s' -> Producer (weakenP p) (weakenL s') +-- Consumer c -> Consumer (weakenC c) +-- Reify ix -> Reify ix + +-- where +-- weakenL :: forall senv' arrs'. PreOpenSeq acc aenv senv' arrs' -> PreOpenSeq acc aenv' senv' arrs' +-- weakenL s' = unSinkSeq (weaken k (SinkSeq s')) + +-- weakenP :: forall a. Producer acc aenv senv a -> Producer acc aenv' senv a +-- weakenP p = +-- case p of +-- StreamIn arrs -> StreamIn arrs +-- ToSeq slix sh a -> ToSeq slix sh (weaken k a) +-- MapSeq f x -> MapSeq (weaken k f) x +-- ChunkedMapSeq f x -> ChunkedMapSeq (weaken k f) x +-- ZipWithSeq f x y -> ZipWithSeq (weaken k f) x y +-- ScanSeq f a x -> ScanSeq (weaken k f) (weaken k a) x + +-- weakenC :: forall a. Consumer acc aenv senv a -> Consumer acc aenv' senv a +-- weakenC c = +-- case c of +-- FoldSeq f a x -> FoldSeq (weaken k f) (weaken k a) x +-- FoldSeqFlatten f a x -> FoldSeqFlatten (weaken k f) (weaken k a) x +-- Stuple t -> +-- let wk :: Atuple (Consumer acc aenv senv) t -> Atuple (Consumer acc aenv' senv) t +-- wk NilAtup = NilAtup +-- wk (SnocAtup t c) = wk t `SnocAtup` weakenC c +-- in +-- Stuple (wk t) +-- --} + +-- -- Array fusion of a de Bruijn computation AST +-- -- =========================================== + +-- -- Array computations +-- -- ------------------ + +-- -- Evaluate a delayed computation and tie the recursive knot +-- -- +-- -- We do a bit of extra work to (try to) maintain that terms should be left +-- -- at their lowest common use site. SEE: [Fusion and the lowest common use site] +-- -- +-- computeAcc +-- :: HasCallStack +-- => Embed OpenAcc aenv arrs +-- -> OpenAcc aenv arrs +-- computeAcc (Embed BaseEnv cc) = OpenAcc (compute cc) +-- computeAcc (Embed env@(PushEnv bot lhs top) cc) = +-- simplify cc & \case +-- Left (Done v) -> bindA env (avarsIn OpenAcc v) +-- Right d -> d & \case +-- Yield repr sh f +-- -> bindA env (OpenAcc (Generate repr sh f)) + +-- Step repr sh p f v@(Var _ ix) +-- | Just Refl <- matchOpenExp sh (arrayShape v) +-- , Just Refl <- isIdentity p +-- -> case ix of +-- ZeroIdx +-- | LeftHandSideSingle ArrayR{} <- lhs +-- , Just (OpenAccFun g) <- strengthen noTop (OpenAccFun f) +-- -> bindA bot (OpenAcc (Map (arrayRtype repr) g top)) +-- _ -> bindA env (OpenAcc (Map (arrayRtype repr) f (avarIn OpenAcc v))) + +-- | Just Refl <- isIdentity f +-- -> case ix of +-- ZeroIdx +-- | LeftHandSideSingle ArrayR{} <- lhs +-- , Just (OpenAccFun q) <- strengthen noTop (OpenAccFun p) +-- , Just (OpenAccExp sz) <- strengthen noTop (OpenAccExp sh) +-- -> bindA bot (OpenAcc (Backpermute (arrayRshape repr) sz q top)) +-- _ -> bindA env (OpenAcc (Backpermute (arrayRshape repr) sh p (avarIn OpenAcc v))) + +-- | otherwise +-- -> case ix of +-- ZeroIdx +-- | LeftHandSideSingle ArrayR{} <- lhs +-- , Just (OpenAccFun g) <- strengthen noTop (OpenAccFun f) +-- , Just (OpenAccFun q) <- strengthen noTop (OpenAccFun p) +-- , Just (OpenAccExp sz) <- strengthen noTop (OpenAccExp sh) +-- -> bindA bot (OpenAcc (Transform repr sz q g top)) +-- _ -> bindA env (OpenAcc (Transform repr sh p f (avarIn OpenAcc v))) + +-- where +-- bindA :: HasCallStack +-- => Extend ArrayR OpenAcc aenv aenv' +-- -> OpenAcc aenv' a +-- -> OpenAcc aenv a +-- bindA BaseEnv b = b +-- bindA (PushEnv env lhs a) b +-- -- If the freshly bound value is directly, returned, we don't have to bind it in a +-- -- let. We can do this if the left hand side does not contain wildcards (other than +-- -- wildcards for unit / nil) and if the value contains the same variables. +-- | Just vars <- extractOpenArrayVars b +-- , Just Refl <- bindingIsTrivial lhs vars = bindA env a +-- | otherwise = bindA env (OpenAcc (Alet lhs a b)) + +-- noTop :: (aenv, a) :?> aenv +-- noTop ZeroIdx = Nothing +-- noTop (SuccIdx ix) = Just ix + + + +-- -- Convert the internal representation of delayed arrays into a real AST +-- -- node. Use the most specific version of a combinator whenever possible. +-- -- +-- compute +-- :: HasCallStack +-- => Cunctation r aenv arrs +-- -> PreOpenAcc OpenAcc aenv arrs +-- compute cc = simplify cc & \case +-- Left (Done v) -> v & \case +-- TupRunit -> Anil +-- TupRsingle v@(Var ArrayR{} _) -> Avar v +-- TupRpair v1 v2 -> avarsIn OpenAcc v1 `Apair` avarsIn OpenAcc v2 +-- Right d -> d & \case +-- Yield aR sh f -> Generate aR sh f +-- Step (ArrayR shR tR) sh p f v +-- | Just Refl <- matchOpenExp sh (arrayShape v) +-- , Just Refl <- isIdentity p -> Map tR f (avarIn OpenAcc v) +-- | Just Refl <- isIdentity f -> Backpermute shR sh p (avarIn OpenAcc v) +-- | otherwise -> Transform (ArrayR shR tR) sh p f (avarIn OpenAcc v) + + +-- -- Representation of a generator as a delayed array +-- -- +-- generateD +-- :: HasCallStack +-- => ArrayR (Array sh e) +-- -> Exp aenv sh +-- -> Fun aenv (sh -> e) +-- -> Embed OpenAcc aenv (Array sh e) +-- generateD repr sh f +-- = Stats.ruleFired "generateD" +-- $ Embed BaseEnv (Yield repr sh f) + + +-- -- Fuse a unary function into a delayed array. Also looks for unzips which can +-- -- be executed in constant time; SEE [unzipD] +-- -- +-- mapD :: HasCallStack +-- => TypeR b +-- -> Fun aenv (a -> b) +-- -> Embed OpenAcc aenv (Array sh a) +-- -> Embed OpenAcc aenv (Array sh b) +-- mapD tR f (unzipD tR f -> Just a) = a +-- mapD tR f (Embed env cc) +-- = Stats.ruleFired "mapD" +-- $ Embed env (go (delaying cc)) +-- where +-- go (Step (ArrayR shR _) sh ix g v) = Step (ArrayR shR tR) sh ix (sinkArrayInstr env f `compose` g) v +-- go (Yield (ArrayR shR _) sh g) = Yield (ArrayR shR tR) sh (sinkArrayInstr env f `compose` g) + + +-- -- If we are unzipping a manifest array then force the term to be computed; +-- -- a backend will be able to execute this in constant time. +-- -- +-- unzipD +-- :: HasCallStack +-- => TypeR b +-- -> Fun aenv (a -> b) +-- -> Embed OpenAcc aenv (Array sh a) +-- -> Maybe (Embed OpenAcc aenv (Array sh b)) +-- unzipD tR f (Embed env cc@(Done v)) +-- | Lam lhs (Body a) <- f +-- , Just vars <- extractExpVars a +-- , ArrayR shR _ <- arrayR cc +-- , f' <- Lam lhs $ Body $ expVars vars +-- = Just $ Embed (env `pushArrayEnv` OpenAcc (Map tR f' $ avarsIn OpenAcc v)) $ doneZeroIdx $ ArrayR shR tR +-- unzipD _ _ _ +-- = Nothing + +-- -- Fuse an index space transformation function that specifies where elements in +-- -- the destination array read there data from in the source array. +-- -- +-- backpermuteD +-- :: HasCallStack +-- => ShapeR sh' +-- -> Exp aenv sh' +-- -> Fun aenv (sh' -> sh) +-- -> Cunctation r aenv (Array sh e) +-- -> Cunctation D aenv (Array sh' e) +-- backpermuteD shR' sh' p = Stats.ruleFired "backpermuteD" . go . delaying +-- where +-- go (Step (ArrayR _ tR) _ q f v) = Step (ArrayR shR' tR) sh' (q `compose` p) f v +-- go (Yield (ArrayR _ tR) _ g) = Yield (ArrayR shR' tR) sh' (g `compose` p) + + +-- -- Transform as a combined map and backwards permutation +-- -- +-- transformD +-- :: HasCallStack +-- => ArrayR (Array sh' b) +-- -> Exp aenv sh' +-- -> Fun aenv (sh' -> sh) +-- -> Fun aenv (a -> b) +-- -> Embed OpenAcc aenv (Array sh a) +-- -> Embed OpenAcc aenv (Array sh' b) +-- transformD (ArrayR shR' tR) sh' p f +-- = Stats.ruleFired "transformD" +-- . fuse (into2 (backpermuteD shR') sh' p) +-- . mapD tR f +-- where +-- fuse :: HasCallStack +-- => (forall r aenv'. Extend ArrayR OpenAcc aenv aenv' -> Cunctation r aenv' as -> Cunctation D aenv' bs) +-- -> Embed OpenAcc aenv as +-- -> Embed OpenAcc aenv bs +-- fuse op (Embed env cc) = Embed env (op env cc) + +-- into2 :: (OpenExp env aenv' a -> OpenFun env aenv' b -> c) +-- -> OpenExp env aenv a +-- -> OpenFun env aenv b +-- -> Extend ArrayR OpenAcc aenv aenv' +-- -> c +-- into2 op a b env = op (sinkArrayInstr env a) (sinkArrayInstr env b) + + +-- -- Replicate as a backwards permutation +-- -- +-- -- TODO: If we have a pattern such as `replicate sh (map f xs)` then in some +-- -- cases it might be beneficial to not fuse these terms, if `f` is +-- -- expensive and/or `sh` is large. +-- -- +-- replicateD +-- :: HasCallStack +-- => SliceIndex slix sl co sh +-- -> Exp aenv slix +-- -> Cunctation r aenv (Array sl e) +-- -> Cunctation D aenv (Array sh e) +-- replicateD sliceIndex slix cc +-- = Stats.ruleFired "replicateD" +-- $ backpermuteD (sliceDomainR sliceIndex) (IndexFull sliceIndex slix (shape cc)) (extend sliceIndex slix) cc + + +-- -- Dimensional slice as a backwards permutation +-- -- +-- sliceD +-- :: HasCallStack +-- => SliceIndex slix sl co sh +-- -> Exp aenv slix +-- -> Cunctation r aenv (Array sh e) +-- -> Cunctation D aenv (Array sl e) +-- sliceD sliceIndex slix cc +-- = Stats.ruleFired "sliceD" +-- $ backpermuteD (sliceShapeR sliceIndex) (IndexSlice sliceIndex slix (shape cc)) (restrict sliceIndex slix) cc + + +-- -- Reshape an array +-- -- +-- -- For delayed arrays this is implemented as an index space transformation. For +-- -- manifest arrays this can be done with the standard Reshape operation in +-- -- constant time without executing any array operations. This does not affect +-- -- the fusion process since the term is already manifest. +-- -- +-- -- TLM: there was a runtime check to ensure the old and new shapes contained the +-- -- same number of elements: this has been lost for the delayed cases! +-- -- +-- reshapeD +-- :: HasCallStack +-- => ShapeR sl +-- -> Embed OpenAcc aenv (Array sh e) +-- -> Exp aenv sl +-- -> Embed OpenAcc aenv (Array sl e) +-- reshapeD slr (Embed env cc) (sinkArrayInstr env -> sl) +-- | Done v <- cc +-- = Embed (env `pushArrayEnv` OpenAcc (Reshape slr sl (avarsIn OpenAcc v))) $ doneZeroIdx repr + +-- | otherwise +-- = Stats.ruleFired "reshapeD" +-- $ Embed env (backpermuteD slr sl (reindex (arrayRshape $ arrayR cc) (shape cc) slr sl) cc) + +-- where +-- ArrayR _ tR = arrayR cc +-- repr = ArrayR slr tR + + +-- -- Combine two arrays element-wise with a binary function to produce a delayed +-- -- array. +-- -- +-- zipWithD +-- :: HasCallStack +-- => TypeR c +-- -> Fun aenv (a -> b -> c) +-- -> Cunctation r aenv (Array sh a) +-- -> Cunctation s aenv (Array sh b) +-- -> Cunctation D aenv (Array sh c) +-- zipWithD tR f cc1 cc0 +-- -- Two stepper functions identically accessing the same array can be kept in +-- -- stepping form. This might yield a simpler final term. +-- -- +-- | Step (ArrayR shR _) sh1 p1 f1 v1 <- delaying cc1 +-- , Step _ sh0 p0 f0 v0 <- delaying cc0 +-- , Just Refl <- matchVar v1 v0 +-- , Just Refl <- matchOpenFun p1 p0 +-- = Stats.ruleFired "zipWithD/step" +-- $ Step (ArrayR shR tR) (intersect shR sh1 sh0) p0 (combine f f1 f0) v0 + +-- -- Otherwise transform both delayed terms into (index -> value) mappings and +-- -- combine the two indexing functions that way. +-- -- +-- | Yield (ArrayR shR _) sh1 f1 <- yield cc1 +-- , Yield _ sh0 f0 <- yield cc0 +-- = Stats.ruleFired "zipWithD" +-- $ Yield (ArrayR shR tR) (intersect shR sh1 sh0) (combine f f1 f0) + +-- | otherwise +-- = error "work is stressing me out, I should take a break" +-- where +-- combine :: forall aenv a b c e. HasCallStack +-- => Fun aenv (a -> b -> c) +-- -> Fun aenv (e -> a) +-- -> Fun aenv (e -> b) +-- -> Fun aenv (e -> c) +-- combine c ixa ixb +-- | Lam lhs1 (Body ixa') <- ixa +-- , Lam lhs2 (Body ixb') <- ixb +-- -- The two LeftHandSides may differ in the use of wildcards. If they do not match, we must +-- -- combine them as done in `combineLHS`. As this will probably not occur often and requires +-- -- additional weakening, we do a quick check whether the left hand sides are equal. +-- -- +-- = case matchELeftHandSide lhs1 lhs2 of +-- Just Refl +-- | Lam lhsA (Lam lhsB (Body c')) <- weakenE (weakenWithLHS lhs1) c +-- -> Lam lhs1 $ Body $ Let lhsA ixa' $ Let lhsB (weakenE (weakenWithLHS lhsA) ixb') c' +-- Nothing +-- | CombinedLHS lhs k1 k2 <- combineLHS lhs1 lhs2 +-- , Lam lhsA (Lam lhsB (Body c')) <- weakenE (weakenWithLHS lhs) c +-- , ixa'' <- weakenE k1 ixa' +-- -> Lam lhs $ Body $ Let lhsA ixa'' $ Let lhsB (weakenE (weakenWithLHS lhsA .> k2) ixb') c' +-- _ +-- -> error "how's your break?" +-- -- +-- | otherwise +-- = error "work is stressing me out, I should get back to it" + + +-- data CombinedLHS s t env1' env2' env where +-- CombinedLHS :: LeftHandSide s t env env' +-- -> env1' :> env' +-- -> env2' :> env' +-- -> CombinedLHS s t env1' env2' env + +-- combineLHS +-- :: HasCallStack +-- => LeftHandSide s t env env1' +-- -> LeftHandSide s t env env2' +-- -> CombinedLHS s t env1' env2' env +-- combineLHS = go weakenId weakenId +-- where +-- go :: env1 :> env -> env2 :> env -> LeftHandSide s t env1 env1' -> LeftHandSide s t env2 env2' -> CombinedLHS s t env1' env2' env +-- go k1 k2 (LeftHandSideWildcard tR) (LeftHandSideWildcard _) = CombinedLHS (LeftHandSideWildcard tR) k1 k2 +-- go k1 k2 (LeftHandSideSingle tR) (LeftHandSideSingle _) = CombinedLHS (LeftHandSideSingle tR) (sink k1) (sink k2) +-- go k1 k2 (LeftHandSidePair l1 h1) (LeftHandSidePair l2 h2) +-- | CombinedLHS l k1' k2' <- go k1 k2 l1 l2 +-- , CombinedLHS h k1'' k2'' <- go k1' k2' h1 h2 = CombinedLHS (LeftHandSidePair l h) k1'' k2'' +-- go k1 k2 (LeftHandSideWildcard _) lhs +-- | Exists lhs' <- rebuildLHS lhs = CombinedLHS lhs' (weakenWithLHS lhs' .> k1) (sinkWithLHS lhs lhs' k2) +-- go k1 k2 lhs (LeftHandSideWildcard _) +-- | Exists lhs' <- rebuildLHS lhs = CombinedLHS lhs' (sinkWithLHS lhs lhs' k1) (weakenWithLHS lhs' .> k2) +-- go _ _ _ _ +-- = internalError "unexpected LHS combination" + + + +-- -- NOTE: [Sharing vs. Fusion] +-- -- +-- -- The approach to array fusion is similar to that the first generation of Repa. +-- -- It was discovered that the most immediately pressing problem with delayed +-- -- arrays in Repa-1 was that it did not preserve sharing of collective +-- -- operations, leading to excessive recomputation and severe repercussions on +-- -- performance if the user did not explicitly intervene. +-- -- +-- -- However, as we have explicit sharing information in the term tree, so it is +-- -- straightforward to respect sharing by not fusing let-bindings, as that +-- -- introduces work duplication. However, sometimes we can be cleverer. +-- -- +-- -- let-floating: +-- -- ------------- +-- -- +-- -- If the binding is of manifest data, we can instead move the let-binding to a +-- -- different point in the program and then continue to fuse into the body. This +-- -- is done by adding the bound term to the Extend environment. In essence this +-- -- is covering a different occurrence of the same problem Extend was introduced +-- -- to handle: let bindings of manifest data unnecessarily get in the way of the +-- -- fusion process. For example: +-- -- +-- -- map f (zipWith g xs (map h xs)) +-- -- +-- -- after sharing recovery results in: +-- -- +-- -- map f (let a0 = xs in zipWith g a0 (map h a0)) +-- -- +-- -- Without allowing the binding for a0 to float outwards, `map f` will not be +-- -- fused into the rest of the program. +-- -- +-- -- let-elimination: +-- -- ---------------- +-- -- +-- -- Array binding points appear in the program because the array data _or_ shape +-- -- was accessed multiple times in the source program. In general we want to fuse +-- -- arbitrary sequences of array _data_, irrespective of how the shape component +-- -- is used. For example, reverse is defined in the prelude as: +-- -- +-- -- reverse xs = let len = unindex1 (shape xs) +-- -- pf i = len - i - 1 +-- -- in +-- -- backpermute (shape xs) (ilift1 pf) xs +-- -- +-- -- Sharing recovery introduces a let-binding for the input `xs` since it is used +-- -- thrice in the definition, which impedes subsequent fusion. However the actual +-- -- array data is only accessed once, with the remaining two uses querying the +-- -- array shape. Since the delayed terms contain the shape of the array they +-- -- represent as a scalar term, if the data component otherwise satisfies the +-- -- rules for fusing terms, as it does in this example, we can eliminate the +-- -- let-binding by pushing the scalar shape and value generation terms directly +-- -- into the body. +-- -- +-- -- Let-elimination can also be used to _introduce_ work duplication, which may +-- -- be beneficial if we can estimate that the cost of re-computation is less than +-- -- the cost of completely evaluating the array and subsequently retrieving the +-- -- data from memory. +-- -- +-- -- let-binding: +-- -- ------------ +-- -- +-- -- Ultimately, we might not want to eliminate the binding. If so, evaluate it +-- -- and add it to a _clean_ Extend environment for the body. If not, the Extend +-- -- list effectively _flattens_ all bindings, so any terms required for the bound +-- -- term get lifted out to the same scope as the body. This increases their +-- -- lifetime and hence raises the maximum memory used. If we don't do this, we +-- -- get terms such as: +-- -- +-- -- let a0 = in +-- -- let bnd = in +-- -- +-- -- +-- -- rather than the following, where the scope of a0 is clearly only availably +-- -- when evaluating the bound term, as it should be: +-- -- +-- -- let bnd = +-- -- let a0 = +-- -- in +-- -- in +-- -- +-- aletD :: HasCallStack +-- => EmbedAcc OpenAcc +-- -> ElimAcc OpenAcc +-- -> ALeftHandSide arrs aenv aenv' +-- -> OpenAcc aenv arrs +-- -> OpenAcc aenv' brrs +-- -> Embed OpenAcc aenv brrs +-- aletD embedAcc elimAcc lhs (embedAcc -> Embed env1 cc1) acc0 + +-- -- let-floating +-- -- ------------ +-- -- +-- -- Immediately inline the variable referring to the bound expression into the +-- -- body, instead of adding to the environments and creating an indirection +-- -- that must be later eliminated by shrinking. +-- -- +-- | LeftHandSideSingle _ <- lhs +-- , Done (TupRsingle v1@(Var ArrayR{} _)) <- cc1 +-- , Embed env0 cc0 <- embedAcc $ rebuildA (subAtop (Avar v1) . sink1 env1) acc0 +-- = Stats.ruleFired "aletD/float" +-- $ Embed (env1 `append` env0) cc0 + +-- -- Ensure we only call 'embedAcc' once on the body expression +-- -- +-- | otherwise +-- = aletD' embedAcc elimAcc lhs (Embed env1 cc1) (embedAcc acc0) + + +-- aletD' :: forall aenv aenv' arrs brrs. HasCallStack +-- => EmbedAcc OpenAcc +-- -> ElimAcc OpenAcc +-- -> ALeftHandSide arrs aenv aenv' +-- -> Embed OpenAcc aenv arrs +-- -> Embed OpenAcc aenv' brrs +-- -> Embed OpenAcc aenv brrs +-- aletD' embedAcc elimAcc (LeftHandSideSingle ArrayR{}) (Embed env1 cc1) (Embed env0 cc0) + +-- -- let-binding +-- -- ----------- +-- -- +-- -- Check whether we can eliminate the let-binding. Note that we must inspect +-- -- the entire term, not just the Cunctation that would be produced by +-- -- embedAcc. If we don't we can be left with dead terms that don't get +-- -- eliminated. This problem occurred in the canny program. +-- -- +-- | acc1 <- computeAcc (Embed env1 cc1) +-- , False <- elimAcc acc1 acc0 +-- = Stats.ruleFired "aletD/bind" +-- $ Embed (BaseEnv `pushArrayEnv` acc1 `append` env0) cc0 + +-- -- let-elimination +-- -- --------------- +-- -- +-- -- Handle the remaining cases in a separate function. It turns out that this +-- -- is important so we aren't excessively sinking/delaying terms. +-- -- +-- | acc0' <- sink1 env1 acc0 +-- = Stats.ruleFired "aletD/eliminate" +-- $ case delaying cc1 of +-- Step{} -> eliminate env1 cc1 acc0' +-- Yield{} -> eliminate env1 cc1 acc0' + +-- where +-- acc0 :: OpenAcc aenv' brrs +-- acc0 = computeAcc (Embed env0 cc0) + +-- kmap :: forall aenv a b. (PreOpenAcc OpenAcc aenv a -> PreOpenAcc OpenAcc aenv b) +-- -> OpenAcc aenv a +-- -> OpenAcc aenv b +-- kmap f (OpenAcc pacc) = OpenAcc (f pacc) + +-- -- The second part of let-elimination. Splitting into two steps exposes the +-- -- extra type variables, and ensures we don't do extra work manipulating the +-- -- body when not necessary (which can lead to a complexity blowup). +-- -- +-- eliminate +-- :: forall r aenv aenv' sh e brrs. HasCallStack +-- => Extend ArrayR OpenAcc aenv aenv' +-- -> Cunctation r aenv' (Array sh e) +-- -> OpenAcc (aenv', Array sh e) brrs +-- -> Embed OpenAcc aenv brrs +-- eliminate env1 cc1 body +-- | Done v1 <- cc1 +-- , TupRsingle v1'@(Var r _) <- v1 = elim r (arrayShape v1') (indexArray v1') +-- | Step r sh1 p1 f1 v1 <- cc1 = elim r sh1 (f1 `compose` indexArray v1 `compose` p1) +-- | Yield r sh1 f1 <- cc1 = elim r sh1 f1 +-- where +-- bnd :: PreOpenAcc OpenAcc aenv' (Array sh e) +-- bnd = compute cc1 + +-- elim :: HasCallStack +-- => ArrayR (Array sh e) +-- -> Exp aenv' sh +-- -> Fun aenv' (sh -> e) +-- -> Embed OpenAcc aenv brrs +-- elim r sh1 f1 +-- | sh1' <- weakenArrayInstr (weakenSucc' weakenId) sh1 +-- , f1' <- weakenArrayInstr (weakenSucc' weakenId) f1 +-- , Embed env0' cc0' <- embedAcc $ rebuildA (subAtop bnd) $ kmap (replaceA sh1' f1' $ Var r ZeroIdx) body +-- = Embed (env1 `append` env0') cc0' + +-- -- As part of let-elimination, we need to replace uses of array variables in +-- -- scalar expressions with an equivalent expression that generates the +-- -- result directly +-- -- +-- -- TODO: when we inline bindings we ought to let bind at the first +-- -- occurrence and use a variable at all subsequent locations. At the +-- -- moment we are just hoping CSE in the simplifier phase does good +-- -- things, but that is limited in what it looks for. +-- -- +-- replaceE :: forall env aenv sh e t. HasCallStack +-- => OpenExp env aenv sh +-- -> OpenFun env aenv (sh -> e) +-- -> ArrayVar aenv (Array sh e) +-- -> OpenExp env aenv t +-- -> OpenExp env aenv t +-- replaceE sh' f' avar@(Var (ArrayR shR _) _) exp = +-- case exp of +-- Let lhs x y -> let k = weakenWithLHS lhs +-- in Let lhs (cvtE x) (replaceE (weakenE k sh') (weakenE k f') avar y) +-- Evar var -> Evar var +-- Foreign tR ff f e -> Foreign tR ff f (cvtE e) +-- Const tR c -> Const tR c +-- Undef tR -> Undef tR +-- Nil -> Nil +-- Pair e1 e2 -> Pair (cvtE e1) (cvtE e2) +-- VecPack vR e -> VecPack vR (cvtE e) +-- VecUnpack vR e -> VecUnpack vR (cvtE e) +-- IndexSlice x ix sh -> IndexSlice x (cvtE ix) (cvtE sh) +-- IndexFull x ix sl -> IndexFull x (cvtE ix) (cvtE sl) +-- ToIndex shR' sh ix -> ToIndex shR' (cvtE sh) (cvtE ix) +-- FromIndex shR' sh i -> FromIndex shR' (cvtE sh) (cvtE i) +-- Case e rhs def -> Case (cvtE e) (over (mapped . _2) cvtE rhs) (fmap cvtE def) +-- Cond p t e -> Cond (cvtE p) (cvtE t) (cvtE e) +-- PrimConst c -> PrimConst c +-- PrimApp g x -> PrimApp g (cvtE x) +-- ShapeSize shR' sh -> ShapeSize shR' (cvtE sh) +-- While p f x -> While (replaceF sh' f' avar p) (replaceF sh' f' avar f) (cvtE x) +-- Coerce t1 t2 e -> Coerce t1 t2 (cvtE e) + +-- ArrayInstr (Shape a) _ +-- | Just Refl <- matchVar a avar -> Stats.substitution "replaceE/shape" sh' +-- | otherwise -> exp + +-- ArrayInstr (Index a) sh +-- | Just Refl <- matchVar a avar +-- , Lam lhs (Body b) <- f' -> Stats.substitution "replaceE/!" . cvtE $ Let lhs sh b +-- | otherwise -> ArrayInstr (Index a) (cvtE sh) + +-- ArrayInstr (LinearIndex a) i +-- | Just Refl <- matchVar a avar +-- , Lam lhs (Body b) <- f' +-- -> Stats.substitution "replaceE/!!" . cvtE +-- $ Let lhs +-- (Let (LeftHandSideSingle scalarTypeInt) i +-- $ FromIndex shR (weakenE (weakenSucc' weakenId) sh') +-- $ Evar +-- $ Var scalarTypeInt ZeroIdx) +-- b +-- | otherwise -> ArrayInstr (LinearIndex a) (cvtE i) + +-- where +-- cvtE :: OpenExp env aenv s -> OpenExp env aenv s +-- cvtE = replaceE sh' f' avar + +-- replaceF :: forall env aenv sh e t. HasCallStack +-- => OpenExp env aenv sh +-- -> OpenFun env aenv (sh -> e) +-- -> ArrayVar aenv (Array sh e) +-- -> OpenFun env aenv t +-- -> OpenFun env aenv t +-- replaceF sh' f' avar fun = +-- case fun of +-- Body e -> Body (replaceE sh' f' avar e) +-- Lam lhs f -> let k = weakenWithLHS lhs +-- in Lam lhs (replaceF (weakenE k sh') (weakenE k f') avar f) + +-- replaceA :: forall aenv sh e a. HasCallStack +-- => Exp aenv sh +-- -> Fun aenv (sh -> e) +-- -> ArrayVar aenv (Array sh e) +-- -> PreOpenAcc OpenAcc aenv a +-- -> PreOpenAcc OpenAcc aenv a +-- replaceA sh' f' avar pacc = +-- case pacc of +-- Avar v +-- | Just Refl <- matchVar v avar -> Avar avar +-- | otherwise -> Avar v + +-- Alet lhs bnd (body :: OpenAcc aenv1 a) -> +-- let w :: aenv :> aenv1 +-- w = weakenWithLHS lhs +-- sh'' = weakenArrayInstr w sh' +-- f'' = weakenArrayInstr w f' +-- in +-- Alet lhs (cvtA bnd) (kmap (replaceA sh'' f'' (weaken w avar)) body) + +-- Use repr arrs -> Use repr arrs +-- Unit tR e -> Unit tR (cvtE e) +-- Acond p at ae -> Acond (cvtE p) (cvtA at) (cvtA ae) +-- Anil -> Anil +-- Atrace msg a b -> Atrace msg (cvtA a) (cvtA b) +-- Apair a1 a2 -> Apair (cvtA a1) (cvtA a2) +-- Awhile p f a -> Awhile (cvtAF p) (cvtAF f) (cvtA a) +-- Apply repr f a -> Apply repr (cvtAF f) (cvtA a) +-- Aforeign repr ff f a -> Aforeign repr ff f (cvtA a) -- no sharing between f and a +-- Generate repr sh f -> Generate repr (cvtE sh) (cvtF f) +-- Map tR f a -> Map tR (cvtF f) (cvtA a) +-- ZipWith tR f a b -> ZipWith tR (cvtF f) (cvtA a) (cvtA b) +-- Backpermute shR sh p a -> Backpermute shR (cvtE sh) (cvtF p) (cvtA a) +-- Transform repr sh p f a -> Transform repr (cvtE sh) (cvtF p) (cvtF f) (cvtA a) +-- Slice slix a sl -> Slice slix (cvtA a) (cvtE sl) +-- Replicate slix sh a -> Replicate slix (cvtE sh) (cvtA a) +-- Reshape shR sl a -> Reshape shR (cvtE sl) (cvtA a) +-- Fold f z a -> Fold (cvtF f) (cvtE <$> z) (cvtA a) +-- FoldSeg i f z a s -> FoldSeg i (cvtF f) (cvtE <$> z) (cvtA a) (cvtA s) +-- Scan d f z a -> Scan d (cvtF f) (cvtE <$> z) (cvtA a) +-- Scan' d f z a -> Scan' d (cvtF f) (cvtE z) (cvtA a) +-- Permute f d p a -> Permute (cvtF f) (cvtA d) (cvtF p) (cvtA a) +-- Stencil s t f x a -> Stencil s t (cvtF f) (cvtB x) (cvtA a) +-- Stencil2 s1 s2 t f x a y b +-- -> Stencil2 s1 s2 t (cvtF f) (cvtB x) (cvtA a) (cvtB y) (cvtA b) +-- -- Collect seq -> Collect (cvtSeq seq) + +-- where +-- cvtA :: OpenAcc aenv s -> OpenAcc aenv s +-- cvtA = kmap (replaceA sh' f' avar) + +-- cvtE :: Exp aenv s -> Exp aenv s +-- cvtE = replaceE sh' f' avar + +-- cvtF :: Fun aenv s -> Fun aenv s +-- cvtF = replaceF sh' f' avar + +-- cvtB :: Boundary aenv s -> Boundary aenv s +-- cvtB Clamp = Clamp +-- cvtB Mirror = Mirror +-- cvtB Wrap = Wrap +-- cvtB (Constant c) = Constant c +-- cvtB (Function f) = Function (cvtF f) + +-- cvtAF :: HasCallStack => PreOpenAfun OpenAcc aenv s -> PreOpenAfun OpenAcc aenv s +-- cvtAF = cvt sh' f' avar +-- where +-- cvt :: forall aenv a. +-- Exp aenv sh -> Fun aenv (sh -> e) -> ArrayVar aenv (Array sh e) +-- -> PreOpenAfun OpenAcc aenv a +-- -> PreOpenAfun OpenAcc aenv a +-- cvt sh'' f'' avar' (Abody a) = Abody $ kmap (replaceA sh'' f'' avar') a +-- cvt sh'' f'' avar' (Alam lhs (af :: PreOpenAfun OpenAcc aenv1 b)) = +-- Alam lhs $ cvt (weakenArrayInstr w sh'') +-- (weakenArrayInstr w f'') +-- (weaken w avar') +-- af +-- where +-- w :: aenv :> aenv1 +-- w = weakenWithLHS lhs + +-- -- Do not fuse bindings of multiple variables +-- aletD' _ _ lhs (Embed env1 cc1) (Embed env0 cc0) +-- = Stats.ruleFired "aletD/bind" +-- $ Embed (PushEnv BaseEnv lhs (computeAcc (Embed env1 cc1)) `append` env0) cc0 + +-- {-- +-- cvtSeq :: PreOpenSeq acc aenv senv s -> PreOpenSeq acc aenv senv s +-- cvtSeq s = +-- case s of +-- Producer p s' -> +-- Producer +-- (case p of +-- StreamIn arrs -> StreamIn arrs +-- ToSeq slix sh a -> ToSeq slix sh (cvtA a) +-- MapSeq f x -> MapSeq (cvtAF f) x +-- ChunkedMapSeq f x -> ChunkedMapSeq (cvtAF f) x +-- ZipWithSeq f x y -> ZipWithSeq (cvtAF f) x y +-- ScanSeq f e x -> ScanSeq (cvtF f) (cvtE e) x) +-- (cvtSeq s') +-- Consumer c -> +-- Consumer (cvtC c) +-- Reify ix -> Reify ix + +-- cvtC :: Consumer acc aenv senv s -> Consumer acc aenv senv s +-- cvtC c = +-- case c of +-- FoldSeq f e x -> FoldSeq (cvtF f) (cvtE e) x +-- FoldSeqFlatten f a x -> FoldSeqFlatten (cvtAF f) (cvtA a) x +-- Stuple t -> Stuple (cvtCT t) + +-- cvtCT :: Atuple (Consumer acc aenv senv) t -> Atuple (Consumer acc aenv senv) t +-- cvtCT NilAtup = NilAtup +-- cvtCT (SnocAtup t c) = cvtCT t `SnocAtup` cvtC c +-- --} + + +-- -- Array conditionals, in particular eliminate branches when the predicate +-- -- reduces to a known constant. +-- -- +-- -- Note that we take the raw unprocessed terms as input. If instead we had the +-- -- terms for each branch in the delayed representation, this would require that +-- -- each term has been sunk into a common environment, which implies the +-- -- conditional has been pushed underneath the intersection of bound terms for +-- -- both branches. This would result in redundant work processing the bindings +-- -- for the branch not taken. +-- -- +-- acondD :: HasCallStack +-- => MatchAcc OpenAcc +-- -> EmbedAcc OpenAcc +-- -> Exp aenv PrimBool +-- -> OpenAcc aenv arrs +-- -> OpenAcc aenv arrs +-- -> Embed OpenAcc aenv arrs +-- acondD matchAcc embedAcc p t e +-- | Const _ 1 <- p = Stats.knownBranch "True" $ embedAcc t +-- | Const _ 0 <- p = Stats.knownBranch "False" $ embedAcc e +-- | Just Refl <- matchAcc t e = Stats.knownBranch "redundant" $ embedAcc e +-- | otherwise = done $ Acond p (computeAcc (embedAcc t)) +-- (computeAcc (embedAcc e)) + + +-- -- Scalar expressions +-- -- ------------------ + +-- toIndex :: ShapeR sh -> OpenExp env aenv sh -> OpenFun env aenv (sh -> Int) +-- toIndex shR sh +-- | DeclareVars lhs k value <- declareVars $ shapeType shR +-- = Lam lhs $ Body $ ToIndex shR (weakenE k sh) $ expVars $ value weakenId + +-- fromIndex :: ShapeR sh -> OpenExp env aenv sh -> OpenFun env aenv (Int -> sh) +-- fromIndex shR sh +-- = Lam (LeftHandSideSingle scalarTypeInt) +-- $ Body +-- $ FromIndex shR (weakenE (weakenSucc' weakenId) sh) +-- $ Evar +-- $ Var scalarTypeInt ZeroIdx + +-- intersect :: ShapeR sh -> OpenExp env aenv sh -> OpenExp env aenv sh -> OpenExp env aenv sh +-- intersect = mkShapeBinary f +-- where +-- f a b = PrimApp (PrimMin singleType) $ Pair a b + +-- -- union :: ShapeR sh -> OpenExp env aenv sh -> OpenExp env aenv sh -> OpenExp env aenv sh +-- -- union = mkShapeBinary f +-- -- where +-- -- f a b = PrimApp (PrimMax singleType) $ Pair a b + +-- mkShapeBinary +-- :: (forall env'. OpenExp env' aenv Int -> OpenExp env' aenv Int -> OpenExp env' aenv Int) +-- -> ShapeR sh +-- -> OpenExp env aenv sh +-- -> OpenExp env aenv sh +-- -> OpenExp env aenv sh +-- mkShapeBinary _ ShapeRz _ _ = Nil +-- mkShapeBinary f (ShapeRsnoc shR) (Pair as a) (Pair bs b) = mkShapeBinary f shR as bs `Pair` f a b +-- mkShapeBinary f shR (Let lhs bnd a) b = Let lhs bnd $ mkShapeBinary f shR a (weakenE (weakenWithLHS lhs) b) +-- mkShapeBinary f shR a (Let lhs bnd b) = Let lhs bnd $ mkShapeBinary f shR (weakenE (weakenWithLHS lhs) a) b +-- mkShapeBinary f shR a b@Pair{} -- `a` is not Pair +-- | DeclareVars lhs k value <- declareVars $ shapeType shR +-- = Let lhs a $ mkShapeBinary f shR (expVars $ value weakenId) (weakenE k b) +-- mkShapeBinary f shR a b -- `b` is not a Pair +-- | DeclareVars lhs k value <- declareVars $ shapeType shR +-- = Let lhs b $ mkShapeBinary f shR (weakenE k a) (expVars $ value weakenId) + +-- reindex :: ShapeR sh' +-- -> OpenExp env aenv sh' +-- -> ShapeR sh +-- -> OpenExp env aenv sh +-- -> OpenFun env aenv (sh -> sh') +-- reindex shR' sh' shR sh +-- | Just Refl <- matchOpenExp sh sh' = identity (shapeType shR') +-- | otherwise = fromIndex shR' sh' `compose` toIndex shR sh + +-- extend :: SliceIndex slix sl co sh +-- -> Exp aenv slix +-- -> Fun aenv (sh -> sl) +-- extend sliceIndex slix +-- | DeclareVars lhs k value <- declareVars $ shapeType $ sliceDomainR sliceIndex +-- = Lam lhs $ Body $ IndexSlice sliceIndex (weakenE k slix) $ expVars $ value weakenId + +-- restrict :: SliceIndex slix sl co sh +-- -> Exp aenv slix +-- -> Fun aenv (sl -> sh) +-- restrict sliceIndex slix +-- | DeclareVars lhs k value <- declareVars $ shapeType $ sliceShapeR sliceIndex +-- = Lam lhs $ Body $ IndexFull sliceIndex (weakenE k slix) $ expVars $ value weakenId + +-- arrayShape :: ArrayVar aenv (Array sh e) -> Exp aenv sh +-- arrayShape var = simplifyExp $ ArrayInstr (Shape var) Nil + +-- indexArray :: ArrayVar aenv (Array sh e) -> Fun aenv (sh -> e) +-- indexArray v@(Var (ArrayR shR _) _) +-- | DeclareVars lhs _ value <- declareVars $ shapeType shR +-- = Lam lhs $ Body $ ArrayInstr (Index v) $ expVars $ value weakenId + +-- linearIndex :: ArrayVar aenv (Array sh e) -> Fun aenv (Int -> e) +-- linearIndex v = Lam (LeftHandSideSingle scalarTypeInt) $ Body $ ArrayInstr (LinearIndex v) $ Evar $ Var scalarTypeInt ZeroIdx + + +-- extractOpenAcc :: ExtractAcc OpenAcc +-- extractOpenAcc (OpenAcc pacc) = Just pacc + +-- extractDelayedOpenAcc :: ExtractAcc DelayedOpenAcc +-- extractDelayedOpenAcc (Manifest pacc) = Just pacc +-- extractDelayedOpenAcc _ = Nothing + +-- extractOpenArrayVars +-- :: OpenAcc aenv a +-- -> Maybe (ArrayVars aenv a) +-- extractOpenArrayVars (OpenAcc pacc) = +-- avarsOut extractOpenAcc pacc + +-- extractDelayedArrayVars +-- :: DelayedOpenAcc aenv a +-- -> Maybe (ArrayVars aenv a) +-- extractDelayedArrayVars acc +-- | Just pacc <- extractDelayedOpenAcc acc = avarsOut extractDelayedOpenAcc pacc +-- | otherwise = Nothing diff --git a/src/Data/Array/Accelerate/Trafo/LetSplit.hs b/src/Data/Array/Accelerate/Trafo/LetSplit.hs index 9d3c1250e..877fc64a4 100644 --- a/src/Data/Array/Accelerate/Trafo/LetSplit.hs +++ b/src/Data/Array/Accelerate/Trafo/LetSplit.hs @@ -1,5 +1,3 @@ -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE GADTs #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.Trafo.LetSplit @@ -14,7 +12,6 @@ module Data.Array.Accelerate.Trafo.LetSplit ( convertAfun, - convertAcc, ) where @@ -39,7 +36,9 @@ convertPreOpenAcc = \case Apair a1 a2 -> Apair (convertAcc a1) (convertAcc a2) Anil -> Anil Atrace msg as bs -> Atrace msg (convertAcc as) (convertAcc bs) - Apply repr f a -> Apply repr (convertAfun f) (convertAcc a) + Manifest as -> Manifest (convertAcc as) + Aassert msg cond as -> Aassert msg cond (convertAcc as) + Aassume cond as -> Aassume cond (convertAcc as) Aforeign repr asm f a -> Aforeign repr asm (convertAfun f) (convertAcc a) Acond e a1 a2 -> Acond e (convertAcc a1) (convertAcc a2) Awhile c f a -> Awhile (convertAfun c) (convertAfun f) (convertAcc a) @@ -56,7 +55,7 @@ convertPreOpenAcc = \case FoldSeg i f e a s -> FoldSeg i f e (convertAcc a) (convertAcc s) Scan d f e a -> Scan d f e (convertAcc a) Scan' d f e a -> Scan' d f e (convertAcc a) - Permute f a1 g a2 -> Permute f (convertAcc a1) g (convertAcc a2) + Permute f a1 a2 -> Permute f (convertAcc a1) (convertAcc a2) Backpermute shr sh f a -> Backpermute shr sh f (convertAcc a) Stencil s tp f b a -> Stencil s tp f b (convertAcc a) Stencil2 s1 s2 tp f b1 a1 b2 a2 -> Stencil2 s1 s2 tp f b1 (convertAcc a1) b2 (convertAcc a2) diff --git a/src/Data/Array/Accelerate/Trafo/LiveVars.hs b/src/Data/Array/Accelerate/Trafo/LiveVars.hs new file mode 100644 index 000000000..5b6059b4d --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/LiveVars.hs @@ -0,0 +1,607 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE MultiWayIf #-} +{-# OPTIONS_HADDOCK hide #-} +-- | +-- Module : Data.Array.Accelerate.Trafo.LiveVars +-- Copyright : [2012..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- +-- Utilties for creating a dead variable / (strongly) live +-- variable analysis on one of our IRs. +-- + +module Data.Array.Accelerate.Trafo.LiveVars + -- ANALYSIS + -- Liveness env + ( LivenessEnv, emptyLivenessEnv + , lEnvPush, lEnvPushLHS + -- Constraints + , addLiveImplies, addLiveImplications + , setLive, setIdxSetLive, setVarsLive + -- Strengthen liveness env + , lEnvStrengthen, lEnvStrengthenLHS, lEnvStrengthenLHSReturn, LHSLiveness + , ReturnImplication(..), ReturnImplications, returnImplicationsLive + , returnImplicationWeakenByLHS, returnImplicationsWeakenByLHS + , returnVars, returnIndices + -- Sub-tuples + , SubTupR(..), subTupR, subTupRpair, subTupUnit, subTupPreserves, subTupDistribute + , DeclareSubVars(..), declareSubVars, DeclareMissingVars(..), declareMissingVars + , DeclareMissingDistributedVars(..), declareMissingDistributedVars + -- TRANSFORMATION + -- Main types + , LVAnalysis(..), LVAnalysisFun(..), LVAnalysis'(..) + -- Index transformation + , ReEnv(..), reEnvIdx, reEnvIndices, reEnvIndices', reEnvVar, reEnvVars + -- Bindings + , bind, BindLiveness(..), bindSub, BindLivenessSub(..) + -- Utilities + , allDead, expectJust + , subTupExp, subTupFun + , composeSubTupR, subTup, subTupDBuf + ) where + +import Data.Maybe (fromMaybe, mapMaybe, isNothing) +import Control.DeepSeq (NFData (rnf)) +import Data.Type.Equality + +import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.IdxSet ( IdxSet(..) ) +import Data.Array.Accelerate.AST.Var +import Data.Array.Accelerate.AST.LeftHandSide +import qualified Data.Array.Accelerate.AST.IdxSet as IdxSet +import Data.Array.Accelerate.AST.Environment +import Data.Array.Accelerate.AST.Exp +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Array.Buffer +import Data.Array.Accelerate.Trafo.Var +import Data.Array.Accelerate.Trafo.Exp.Substitution +import Data.Array.Accelerate.Trafo.Substitution +import Data.Array.Accelerate.Error + +-- ANALYSIS + +-- During the analysis, a variable is either known to be live, or unknown +-- whether it is live or dead. In the last case, we store two sets: +-- * 'implies': the set of variables that are live if this variable appears to be live +-- * 'implied-by': if any implied-by variable is live, then this variable is live +-- +-- Note that type argument 't' is not used, but it is needed to give Liveness +-- the correct kind to be used in 'Env', as defined in 'LivenessEnv'. +data Liveness env t where + Live :: Liveness env t + + -- Not sure if this is live, but if it is then the given set of indices is also live, + -- i.e. the liveness of this variable implies the liveness of others. + Unknown + :: IdxSet env -- implies + -> IdxSet env -- implied by + -> Liveness env t + +dead :: Liveness env t +dead = Unknown IdxSet.empty IdxSet.empty + +data LivenessEnv env where + LEmpty :: LivenessEnv () + LPush + :: LivenessEnv env + -> Liveness env t + -> LivenessEnv (env, t) + +emptyLivenessEnv :: LivenessEnv () +emptyLivenessEnv = LEmpty + +lEnvPush :: LivenessEnv env -> LivenessEnv (env, t) +lEnvPush env = LPush env dead + +lEnvPushLHS :: LeftHandSide s t env1 env2 -> LivenessEnv env1 -> LivenessEnv env2 +lEnvPushLHS (LeftHandSideSingle _) env = lEnvPush env +lEnvPushLHS (LeftHandSideWildcard _) env = env +lEnvPushLHS (LeftHandSidePair lhs1 lhs2) env = lEnvPushLHS lhs2 $ lEnvPushLHS lhs1 env + +-- If idx1 is live, then idx2 is live +addLiveImplies :: Idx env s -> Idx env t -> LivenessEnv env -> LivenessEnv env +addLiveImplies = \idx1 idx2 env -> fromMaybe (setLive idx2 env) $ go idx1 idx2 env + where + go :: Idx env s -> Idx env t -> LivenessEnv env -> Maybe (LivenessEnv env) + go ZeroIdx ZeroIdx env = Just env -- Trivial constraint, always satisfied + go (SuccIdx idx1) (SuccIdx idx2) (LPush env l) = (`LPush` l) <$> go idx1 idx2 env + go ZeroIdx (SuccIdx idx) (LPush env l) = case l of + Live -> Nothing + Unknown implies impliedBy -> + Just $ LPush env $ Unknown (IdxSet.insert idx implies) impliedBy + go (SuccIdx idx) ZeroIdx (LPush env l) + | isLive idx env = Nothing + | otherwise = case l of + Live -> Just $ LPush env Live + Unknown implies impliedBy -> + Just $ LPush env $ Unknown implies $ IdxSet.insert idx impliedBy + +-- If any of impliedBy becomes live, then implies are live. +addLiveImplications :: IdxSet env -> IdxSet env -> LivenessEnv env -> LivenessEnv env +addLiveImplications impliedBy implies env + | anyIsLive impliedBy env + = setIdxSetLive implies env + | otherwise + = addLiveImplicationsCurrentlyUnknown impliedBy implies env + +-- If any of impliedBy becomes live, then implies are live. +-- Assumes that all indices in impliedBy are currently Unknown. +addLiveImplicationsCurrentlyUnknown :: IdxSet env -> IdxSet env -> LivenessEnv env -> LivenessEnv env +-- Base case / Early out +addLiveImplicationsCurrentlyUnknown (IdxSet PEnd) _ env = env +addLiveImplicationsCurrentlyUnknown _ (IdxSet PEnd) env = env +-- Recursive +addLiveImplicationsCurrentlyUnknown impliedBy implies (LPush env l) = LPush (addLiveImplicationsCurrentlyUnknown (IdxSet.drop impliedBy) (IdxSet.drop implies) env) l' + where + l' = case l of + Unknown implies' impliedBy' -> + Unknown + (if ZeroIdx `IdxSet.member` impliedBy + then IdxSet.drop implies `IdxSet.union` implies' + else implies') + (if ZeroIdx `IdxSet.member` implies + then IdxSet.drop impliedBy `IdxSet.union` impliedBy' + else impliedBy') + Live + | ZeroIdx `IdxSet.member` impliedBy -> error "`Unknown` may not point to an already live variable" + | otherwise -> Live + +setLive :: Idx env s -> LivenessEnv env -> LivenessEnv env +setLive = \idx env -> uncurry setIdxSetLive $ go idx env + where + go :: Idx env s -> LivenessEnv env -> (IdxSet env, LivenessEnv env) + go (SuccIdx idx) (LPush env l) = case l of + Live -> (IdxSet.skip newSet, LPush env' Live) + Unknown implies impliedBy + | idx `IdxSet.member` impliedBy -> + -- This variable may be present in the implied-by set of other variables, + -- hence we need to add it to the IdxSet and potentially mark other variables + -- live in the next iteration of setIdxSetLive. + (IdxSet.push $ implies `IdxSet.union` newSet, LPush env' Live) + | otherwise -> + (IdxSet.skip $ newSet, LPush env' $ Unknown (IdxSet.remove idx implies) impliedBy) + where + (newSet, env') = go idx env + go ZeroIdx (LPush env l) = case l of + Live -> (IdxSet.empty, LPush env Live) + Unknown implies _ -> (IdxSet.skip implies, LPush env Live) + +setIdxSetLive :: IdxSet env -> LivenessEnv env -> LivenessEnv env +setIdxSetLive = \set env -> + if IdxSet.null set then env else uncurry setIdxSetLive $ go set env + where + go :: IdxSet env -> LivenessEnv env -> (IdxSet env, LivenessEnv env) + go (IdxSet PEnd) env = (IdxSet.empty, env) + go liveSet (LPush env l) = case l of + Live -> (IdxSet.skip newSet, LPush env' Live) + Unknown implies impliedBy + -- Does this variable become live? + | ZeroIdx `IdxSet.member` liveSet || IdxSet.overlaps tailLiveSet impliedBy -> + -- This variable may be present in the implied-by set of other variables, + -- hence we need to add it to the IdxSet and potentially mark other variables + -- live in the next iteration of setIdxSetLive. + (IdxSet.push $ implies `IdxSet.union` newSet, LPush env' Live) + | otherwise -> + (IdxSet.skip newSet, LPush env' $ Unknown (implies IdxSet.\\ tailLiveSet) impliedBy) + where + tailLiveSet = IdxSet.drop liveSet + (newSet, env') = go tailLiveSet env + +setVarsLive :: Vars s env t -> LivenessEnv env -> LivenessEnv env +setVarsLive = setIdxSetLive . IdxSet.fromVars + +isLive :: Idx env t -> LivenessEnv env -> Bool +isLive (SuccIdx idx) (LPush env _) = isLive idx env +isLive ZeroIdx (LPush _ l) = case l of + Live -> True + _ -> False + +anyIsLive :: IdxSet env -> LivenessEnv env -> Bool +anyIsLive (IdxSet PEnd) _ = False +anyIsLive indices (LPush _ Live) + | ZeroIdx `IdxSet.member` indices = True +anyIsLive indices (LPush env _) = anyIsLive (IdxSet.drop indices) env + +-- Drops one entry of the liveness env. +-- Propagates the constraints imposed by the first entry on the remaining entries in the environment. +-- Returns Nothing if this variable is definitely live, or Just impliedBy if it is only live if one of the indices in impliedBy is live. +lEnvStrengthen :: LivenessEnv (env1, t) -> (Maybe (IdxSet env1), LivenessEnv env1) +lEnvStrengthen = \case + LPush env Live -> (Nothing, env) + LPush env (Unknown implies impliedBy) -> (Just impliedBy, addLiveImplicationsCurrentlyUnknown impliedBy implies env) + +-- Extension of LeftHandSide +-- Data structure to remember data between a call to 'lEnvStrengthenLHS' and 'bind' +data LHSLiveness s t env1 env2 where + LHSLivenessSingle + :: Maybe (IdxSet env1) -- Nothing if it is definitely live, or Just impliedBy if it is only live if one of the indices in impliedBy is live. + -> s t + -> LHSLiveness s t env1 (env1, t) + LHSLivenessWildcard + :: TupR s t + -> LHSLiveness s t env1 env1 + LHSLivenessPair + :: LHSLiveness s t1 env1 env2 + -> LHSLiveness s t2 env2 env3 + -> LHSLiveness s (t1, t2) env1 env3 + +lEnvStrengthenLHS :: LeftHandSide s t env1 env2 -> LivenessEnv env2 -> (LHSLiveness s t env1 env2, LivenessEnv env1) +lEnvStrengthenLHS lhs env = case lhs of + LeftHandSideWildcard tp -> (LHSLivenessWildcard tp, env) + LeftHandSideSingle tp + | (m, env') <- lEnvStrengthen env + -> (LHSLivenessSingle m tp, env') + LeftHandSidePair lhs1 lhs2 + | (lhs2', env') <- lEnvStrengthenLHS lhs2 env + , (lhs1', env'') <- lEnvStrengthenLHS lhs1 env' + -> (LHSLivenessPair lhs1' lhs2', env'') + +lEnvStrengthenLHSReturn :: LeftHandSide s t env1 env2 -> LivenessEnv env2 -> (LHSLiveness s t env1 env2, LivenessEnv env1, ReturnImplications env1 t) +lEnvStrengthenLHSReturn lhs env = case lhs of + LeftHandSideWildcard tp -> (LHSLivenessWildcard tp, env, mapTupR (const returnImplicationDead) tp) + LeftHandSideSingle tp + | (m, env') <- lEnvStrengthen env + -> (LHSLivenessSingle m tp, env', TupRsingle $ maybe ReturnLive ReturnImpliedBy m) + LeftHandSidePair lhs1 lhs2 + | (lhs2', env', r2) <- lEnvStrengthenLHSReturn lhs2 env + , (lhs1', env'', r1) <- lEnvStrengthenLHSReturn lhs1 env' + -> (LHSLivenessPair lhs1' lhs2', env'', TupRpair r1 (mapTupR (returnImplicationStrengthenByLHS lhs1') r2)) + +-- ANALYSIS for languages with a return value + +data ReturnImplication env t + -- If any of the variables in the set is live, + -- then the returned value is live. + = ReturnImpliedBy (IdxSet env) + -- This return value is always live + | ReturnLive +type ReturnImplications env = TupR (ReturnImplication env) + +returnImplicationsLive :: ReturnImplications env t +returnImplicationsLive = TupRsingle ReturnLive + +flattenReturnImplications :: ReturnImplications env t -> ReturnImplication env t +flattenReturnImplications TupRunit = returnImplicationDead +flattenReturnImplications (TupRsingle ret) = ret +flattenReturnImplications (TupRpair r1 r2) = case (flattenReturnImplications r1, flattenReturnImplications r2) of + (ReturnLive, _) -> ReturnLive + (_, ReturnLive) -> ReturnLive + (ReturnImpliedBy a, ReturnImpliedBy b) -> ReturnImpliedBy $ IdxSet.union a b + +returnImplicationDead :: ReturnImplication env t +returnImplicationDead = ReturnImpliedBy IdxSet.empty + +returnImplicationStrengthenByLHS :: LHSLiveness s t env1 env2 -> ReturnImplication env2 u -> ReturnImplication env1 u +returnImplicationStrengthenByLHS _ ReturnLive = ReturnLive +returnImplicationStrengthenByLHS lhs' (ReturnImpliedBy impliedBy1) = go lhs' impliedBy1 + where + go :: LHSLiveness s' t' env1 env2 -> IdxSet env2 -> ReturnImplication env1 t + go (LHSLivenessSingle m _) impliedBy + | ZeroIdx `IdxSet.member` impliedBy = + case m of + Nothing -> ReturnLive + Just impliedBy' -> ReturnImpliedBy $ IdxSet.drop impliedBy `IdxSet.union` impliedBy' + | otherwise = ReturnImpliedBy $ IdxSet.drop impliedBy + go (LHSLivenessWildcard _) impliedBy = ReturnImpliedBy impliedBy + go (LHSLivenessPair lhs1 lhs2) impliedBy = case go lhs2 impliedBy of + ReturnLive -> ReturnLive + ReturnImpliedBy i -> go lhs1 i + +returnImplicationWeakenByLHS :: LeftHandSide s t env1 env2 -> ReturnImplication env1 u -> ReturnImplication env2 u +returnImplicationWeakenByLHS _ ReturnLive = ReturnLive +returnImplicationWeakenByLHS l (ReturnImpliedBy impliedBy) = ReturnImpliedBy $ IdxSet.skip' l impliedBy + +returnImplicationsWeakenByLHS :: LeftHandSide s t env1 env2 -> ReturnImplications env1 u -> ReturnImplications env2 u +returnImplicationsWeakenByLHS lhs = mapTupR (returnImplicationWeakenByLHS lhs) + +returnVars :: Distributes s => ReturnImplications env t -> Vars s env t -> LivenessEnv env -> LivenessEnv env +returnVars (TupRsingle ReturnLive) vars = setVarsLive vars +returnVars (TupRsingle (ReturnImpliedBy impliedBy)) vars = addLiveImplications impliedBy $ IdxSet.fromVars vars +returnVars TupRunit TupRunit = id +returnVars (TupRpair r1 r2) (TupRpair v1 v2) = returnVars r1 v1 . returnVars r2 v2 +returnVars TupRunit (TupRsingle tp) = unitImpossible tp +returnVars TupRpair{} (TupRsingle tp) = pairImpossible tp + +returnIndices :: ReturnImplications env t -> IdxSet env -> LivenessEnv env -> LivenessEnv env +returnIndices ret indices = case flattenReturnImplications ret of + ReturnLive -> setIdxSetLive indices + ReturnImpliedBy impliedBy -> addLiveImplications impliedBy indices + +data SubTupR t t' where + SubTupRskip :: SubTupR t () + + SubTupRkeep :: SubTupR t t + + SubTupRpair :: SubTupR t1 t1' + -> SubTupR t2 t2' + -> SubTupR (t1, t2) (t1', t2') + +instance NFData (SubTupR t t') where + rnf SubTupRskip = () + rnf SubTupRkeep = () + rnf (SubTupRpair s1 s2) = rnf s1 `seq` rnf s2 + +composeSubTupR :: SubTupR b c -> SubTupR a b -> SubTupR a c +composeSubTupR bc SubTupRkeep = bc +composeSubTupR SubTupRkeep ab = ab +composeSubTupR SubTupRskip SubTupRskip = SubTupRskip +composeSubTupR SubTupRskip _ = SubTupRskip +composeSubTupR (SubTupRpair bc bc') (SubTupRpair ab ab') = subTupRpair (composeSubTupR bc ab) (composeSubTupR bc' ab') + +subTupRpair :: SubTupR t1 t1' -> SubTupR t2 t2' -> SubTupR (t1, t2) (t1', t2') +subTupRpair SubTupRkeep SubTupRkeep = SubTupRkeep +subTupRpair s t = SubTupRpair s t + +subTupR :: SubTupR t t' -> TupR s t -> TupR s t' +subTupR SubTupRskip _ = TupRunit +subTupR SubTupRkeep t = t +subTupR (SubTupRpair s1 s2) (TupRpair t1 t2) = subTupR s1 t1 `TupRpair` subTupR s2 t2 +subTupR _ _ = internalError "Tuple mismatch" + +subTup :: SubTupR t t' -> t -> t' +subTup SubTupRskip _ = () +subTup SubTupRkeep t = t +subTup (SubTupRpair l r) (t1, t2) = (subTup l t1, subTup r t2) + +subTupUnit :: SubTupR () t' -> t' :~: () +subTupUnit SubTupRskip = Refl +subTupUnit SubTupRkeep = Refl + +-- Checks whether the SubTup preserves the entire tuple. +-- This function needs a 'TupR s t' to check whether 'SubTupRskip' is used on +-- '()'. On a unit, SubTupRskip and SubTupRkeep have the same meaning. +-- +subTupPreserves :: TupR s t -> SubTupR t t' -> Maybe (t' :~: t) +subTupPreserves TupRunit s = Just $ subTupUnit s +subTupPreserves (TupRpair t1 t2) (SubTupRpair s1 s2) + | Just Refl <- subTupPreserves t1 s1 + , Just Refl <- subTupPreserves t2 s2 = Just Refl +subTupPreserves _ _ = Nothing + +-- 's' is ambiguous, so explicit type annotation is needed. +subTupDistribute :: forall s t t'. SubTupR t t' -> SubTupR (Distribute s t) (Distribute s t') +subTupDistribute SubTupRskip = SubTupRskip +subTupDistribute SubTupRkeep = SubTupRkeep +subTupDistribute (SubTupRpair s1 s2) = subTupDistribute @s s1 `SubTupRpair` subTupDistribute @s s2 + +-- TRANSFORMATION + +data ReEnv env subenv where + ReEnvEnd :: ReEnv () () + ReEnvSkip :: ReEnv env subenv -> ReEnv (env, t) subenv + ReEnvKeep :: ReEnv env subenv -> ReEnv (env, t) (subenv, t) + +reEnvIdx :: ReEnv env subenv -> env :?> subenv +reEnvIdx (ReEnvKeep _) ZeroIdx = Just ZeroIdx +reEnvIdx (ReEnvKeep r) (SuccIdx ix) = SuccIdx <$> reEnvIdx r ix +reEnvIdx (ReEnvSkip r) (SuccIdx ix) = reEnvIdx r ix +reEnvIdx _ _ = Nothing + +reEnvIndices :: ReEnv env subenv -> [Exists (Idx env)] -> [Exists (Idx subenv)] +reEnvIndices re = mapMaybe $ \(Exists idx) -> Exists <$> reEnvIdx re idx + +reEnvIndices' :: ReEnv env subenv -> [Idx env t] -> [Idx subenv t] +reEnvIndices' re = mapMaybe $ reEnvIdx re + +reEnvVar :: ReEnv env subenv -> Var s env t -> Maybe (Var s subenv t) +reEnvVar re (Var tp idx) = Var tp <$> reEnvIdx re idx + +reEnvVars :: ReEnv env subenv -> Vars s env t -> Maybe (Vars s subenv t) +reEnvVars re = traverseTupR $ reEnvVar re + +-- Fuse with lEnvStrengthenLHS (to be implemented) or let lEnvStrengthenLHS return its intermediate values in some data structure +-- Is that data structure just something like (TupR Liveness), but then with some fixes for proper environment? +bind :: LHSLiveness s t env env' -> ReEnv env subenv -> BindLiveness s t env' subenv +bind lhs re = case lhs of + LHSLivenessSingle mImplied tp + -- Was this variable already known to be live (mImplied == Nothing), + -- or is one of the implied-by variables live (mImplied == Just implied && check re implied)? + | maybe True (isImpliedLive re) mImplied -> + BindLiveness (LeftHandSideSingle tp) $ ReEnvKeep re + | otherwise -> + BindLiveness (LeftHandSideWildcard $ TupRsingle tp) $ ReEnvSkip re + LHSLivenessWildcard tp -> + BindLiveness (LeftHandSideWildcard tp) re + LHSLivenessPair lhs1 lhs2 + | BindLiveness lhs1' re1 <- bind lhs1 re + , BindLiveness lhs2' re2 <- bind lhs2 re1 + -> BindLiveness (leftHandSidePair lhs1' lhs2') re2 + +-- Given an implied-by set, checks if the variabble is live according to the ReEnv. +-- That is, it checks if one of the indices in the IdxSet is live in ReEnv. +isImpliedLive :: ReEnv env1 subenv1 -> IdxSet env1 -> Bool +isImpliedLive _ (IdxSet PEnd) = False +isImpliedLive (ReEnvKeep env) impliedBy + | ZeroIdx `IdxSet.member` impliedBy = True -- One of the members of impliedBy is live + | otherwise = isImpliedLive env $ IdxSet.drop impliedBy +isImpliedLive (ReEnvSkip env) impliedBy = isImpliedLive env $ IdxSet.drop impliedBy + +-- Captures the existentional subenv' +data BindLiveness s t env' subenv where + BindLiveness + :: LeftHandSide s t subenv subenv' + -> ReEnv env' subenv' + -> BindLiveness s t env' subenv + +bindSub :: LHSLiveness s t env env' -> ReEnv env subenv -> BindLivenessSub s t env' subenv +bindSub lhs re = case lhs of + LHSLivenessSingle mImplied tp + -- Was this variable already known to be live (mImplied == Nothing), + -- or is one of the implied-by variables live (mImplied == Just implied && check re implied)? + | maybe True (isImpliedLive re) mImplied -> + BindLivenessSub + SubTupRkeep + (LeftHandSideSingle tp) + (LeftHandSideSingle tp) + (ReEnvKeep re) + | otherwise -> + BindLivenessSub + SubTupRskip + (LeftHandSideWildcard $ TupRsingle tp) + (LeftHandSideWildcard TupRunit) + (ReEnvSkip re) + LHSLivenessWildcard tp -> + BindLivenessSub + SubTupRskip + (LeftHandSideWildcard tp) + (LeftHandSideWildcard TupRunit) + re + LHSLivenessPair lhs1 lhs2 + | BindLivenessSub subTup1 lhs1Full lhs1Sub re1 <- bindSub lhs1 re + , BindLivenessSub subTup2 lhs2Full lhs2Sub re2 <- bindSub lhs2 re1 + -> if + | LeftHandSideWildcard _ <- lhs1Sub + , LeftHandSideWildcard _ <- lhs2Sub -> + BindLivenessSub + SubTupRskip + (leftHandSidePair lhs1Full lhs2Full) + (LeftHandSideWildcard TupRunit) + re2 + + | SubTupRkeep <- subTup1 + , SubTupRkeep <- subTup2 + , lhs' <- leftHandSidePair lhs1Full lhs2Full -> + BindLivenessSub + SubTupRkeep + lhs' + lhs' + re2 + + | otherwise -> + BindLivenessSub + (SubTupRpair subTup1 subTup2) + (leftHandSidePair lhs1Full lhs2Full) + (leftHandSidePair lhs1Sub lhs2Sub) + re2 + +-- Captures the existentionals subenv' and t' +data BindLivenessSub s t env' subenv where + BindLivenessSub + :: SubTupR t t' + -> LeftHandSide s t subenv subenv' + -> LeftHandSide s t' subenv subenv' + -> ReEnv env' subenv' + -> BindLivenessSub s t env' subenv + +-- For an IR parameterized over the result type, implement a function of this +-- type: +-- +-- stronglyLiveVariables :: LivenessEnv env -> ReturnImplications env t -> SomeAcc env t -> LVAnalysis SomeAcc env t +-- +-- If the IR does have a result type, but we cannot change that type, then use: +-- +-- stronglyLiveVariables :: LivenessEnv env -> ReturnImplications env t -> SomeAcc env t -> LVAnalysisFun SomeAcc env t +-- +-- For an IR which is not parameterized over the result type, but only over the +-- environment, implement a function of this type: +-- +-- stronglyLiveVariables :: LivenessEnv env -> SomeAcc env -> LVAnalysis' SomeAcc env + +-- If this part of the returned value is returned, then this set of variables is live +data LVAnalysis ir env t where + LVAnalysis + :: LivenessEnv env + -- Depending on the binding, it may or may not be possible to restrict the + -- term to only the used parts of the tuple. + -> (forall subenv t'. ReEnv env subenv -> SubTupR t t' -> Either (ir subenv t) (ir subenv t')) + -> LVAnalysis ir env t + +data LVAnalysisFun ir env t where + LVAnalysisFun + :: LivenessEnv env + -> (forall subenv. ReEnv env subenv -> ir subenv t) + -> LVAnalysisFun ir env t + +data LVAnalysis' ir env where + LVAnalysis' + :: LivenessEnv env + -> (forall subenv. ReEnv env subenv -> ir subenv) + -> LVAnalysis' ir env + +allDead :: ReEnv env subenv -> IdxSet env -> Bool +allDead re set = all (\(Exists idx) -> isNothing $ reEnvIdx re idx) $ IdxSet.toList set + +expectJust :: HasCallStack => Maybe a -> a +expectJust (Just a) = a +expectJust Nothing = internalError "Substitution in live variable analysis failed. A variable which was assumed to be dead appeared to be live." + +-- Utilities for expressions +subTupExp :: IsArrayInstr arr => SubTupR t t' -> PreOpenExp arr env t -> PreOpenExp arr env t' +subTupExp SubTupRkeep expr = expr +subTupExp SubTupRskip _ = Nil +subTupExp s expr + | DeclareSubVars lhs _ vars <- declareSubVars (expType expr) s + = Let lhs expr $ expVars $ vars weakenId + +subTupFun :: IsArrayInstr arr => SubTupR t t' -> PreOpenFun arr env (s -> t) -> PreOpenFun arr env (s -> t') +subTupFun s (Lam lhs (Body body)) = Lam lhs $ Body $ subTupExp s body +subTupFun _ _ = internalError "Expected unary function" + +subTupDBuf :: SubTupR t t' -> TupR s (Distribute Buffer t) -> TupR s (Distribute Buffer t') +subTupDBuf SubTupRskip _ = TupRunit +subTupDBuf SubTupRkeep t = t +subTupDBuf (SubTupRpair s1 s2) (TupRpair t1 t2) = subTupDBuf s1 t1 `TupRpair` subTupDBuf s2 t2 +subTupDBuf _ _ = internalError "Tuple mismatch" + +data DeclareSubVars s t t' env where + DeclareSubVars :: LeftHandSide s t env env' + -> (env :> env') + -> (forall env''. env' :> env'' -> Vars s env'' t') + -> DeclareSubVars s t t' env + +declareSubVars :: TupR s t -> SubTupR t t' -> DeclareSubVars s t t' env +declareSubVars tp SubTupRkeep + | DeclareVars lhs k vars <- declareVars tp = DeclareSubVars lhs k vars +declareSubVars tp SubTupRskip + = DeclareSubVars (LeftHandSideWildcard tp) weakenId (const TupRunit) +declareSubVars (TupRpair t1 t2) (SubTupRpair s1 s2) + | DeclareSubVars lhs1 subst1 a1 <- declareSubVars t1 s1 + , DeclareSubVars lhs2 subst2 a2 <- declareSubVars t2 s2 + = DeclareSubVars (LeftHandSidePair lhs1 lhs2) (subst2 .> subst1) $ \k -> a1 (k .> subst2) `TupRpair` a2 k +declareSubVars _ _ = internalError "Tuple mismatch" + +data DeclareMissingVars s t t' env where + -- Captures existentials env' and t''. + DeclareMissingVars :: LeftHandSide s t'' env env' + -> (env :> env') + -> (forall env''. env' :> env'' -> Vars s env'' t) + -> DeclareMissingVars s t t' env + +declareMissingVars :: TupR s t -> SubTupR t t' -> Vars s env t' -> DeclareMissingVars s t t' env +declareMissingVars _ SubTupRkeep vars = DeclareMissingVars (LeftHandSideWildcard TupRunit) weakenId (\k -> mapTupR (weaken k) vars) +declareMissingVars tp SubTupRskip _ + | DeclareVars lhs k value <- declareVars tp + = DeclareMissingVars lhs k value +declareMissingVars (TupRpair t1 t2) (SubTupRpair s1 s2) (TupRpair v1 v2) + | DeclareMissingVars lhs1 subst1 value1 <- declareMissingVars t1 s1 v1 + , DeclareMissingVars lhs2 subst2 value2 <- declareMissingVars t2 s2 (mapTupR (weaken subst1) v2) + = DeclareMissingVars (LeftHandSidePair lhs1 lhs2) (subst2 .> subst1) $ \k -> value1 (k .> subst2) `TupRpair` value2 k +declareMissingVars _ _ _ = internalError "Tuple mismatch" + +data DeclareMissingDistributedVars f s' s t t' env where + -- Captures existentials env' and t''. + DeclareMissingDistributedVars + :: TupR s' t'' + -> LeftHandSide s (Distribute f t'') env env' + -> (env :> env') + -> (forall env''. env' :> env'' -> Vars s env'' (Distribute f t)) + -> DeclareMissingDistributedVars f s' s t t' env + +-- 'f' is ambiguous +declareMissingDistributedVars + :: forall f s' s t t' env. + TupR s' t -> TupR s (Distribute f t) -> SubTupR t t' -> Vars s env (Distribute f t') -> DeclareMissingDistributedVars f s' s t t' env +declareMissingDistributedVars _ _ SubTupRkeep vars + = DeclareMissingDistributedVars TupRunit (LeftHandSideWildcard TupRunit) weakenId (\k -> mapTupR (weaken k) vars) +declareMissingDistributedVars tp tp' SubTupRskip _ + | DeclareVars lhs k value <- declareVars tp' + = DeclareMissingDistributedVars tp lhs k value +declareMissingDistributedVars (TupRpair t1 t2) (TupRpair t1' t2') (SubTupRpair s1 s2) (TupRpair v1 v2) + | DeclareMissingDistributedVars st1 lhs1 subst1 value1 <- declareMissingDistributedVars @f t1 t1' s1 v1 + , DeclareMissingDistributedVars st2 lhs2 subst2 value2 <- declareMissingDistributedVars @f t2 t2' s2 (mapTupR (weaken subst1) v2) + = DeclareMissingDistributedVars (TupRpair st1 st2) (LeftHandSidePair lhs1 lhs2) (subst2 .> subst1) $ \k -> value1 (k .> subst2) `TupRpair` value2 k +declareMissingDistributedVars _ _ _ _ = internalError "Tuple mismatch" diff --git a/src/Data/Array/Accelerate/Trafo/Lowering.hs b/src/Data/Array/Accelerate/Trafo/Lowering.hs new file mode 100644 index 000000000..abcc1d6f4 --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Lowering.hs @@ -0,0 +1,1651 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE ViewPatterns #-} +-- | +-- Module : Data.Array.Accelerate.Trafo.Lowering +-- Copyright : [2012..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- +-- Lowers the platform-independent AST into OperationAcc +-- + +module Data.Array.Accelerate.Trafo.Lowering + where + +import qualified Data.Array.Accelerate.AST as Named +import Data.Array.Accelerate.AST.Environment hiding ( Val ) +import Data.Array.Accelerate.AST.Idx +import qualified Data.Array.Accelerate.AST.IdxSet as IdxSet +import Data.Array.Accelerate.AST.Var +import Data.Array.Accelerate.AST.Operation +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.Analysis.Match ((:~:)(..)) +import Data.Array.Accelerate.Representation.Array +import Data.Array.Accelerate.Representation.Ground +import Data.Array.Accelerate.Representation.Shape +import Data.Array.Accelerate.Representation.Slice +import Data.Array.Accelerate.Representation.Stencil +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Sugar.Foreign +import Data.Array.Accelerate.Trafo.Exp.Substitution +import Data.Array.Accelerate.Trafo.Operation.Substitution +import Data.Array.Accelerate.Trafo.Var +import Data.Array.Accelerate.Type +import Data.Array.Accelerate.Error +import Data.Kind +import Data.Maybe (isJust) + +type a $ b = a b +infixr 0 $ + +-- | Vanilla boundary condition specification for stencil operations +-- Note that this data type is not used in the actual AST, it is just +-- used to pass to 'mkStencil' or 'mkStencil2'. +-- +data Boundary env t where + -- Clamp coordinates to the extent of the array + Clamp :: Boundary env t + + -- Mirror coordinates beyond the array extent + Mirror :: Boundary env t + + -- Wrap coordinates around on each dimension + Wrap :: Boundary env t + + -- Use a constant value for outlying coordinates + Constant :: e + -> Boundary env (Array sh e) + + -- Apply the given function to outlying coordinates + Function :: Fun env (sh -> e) + -> Boundary env (Array sh e) + + +class NFData' op => LowerAcc (op :: Type -> Type) where + mkMap :: Arg env (Fun' (s -> t)) + -> Arg env (In sh s) + -> Arg env (Out sh t) + -> OperationAcc op env () + mkMap f input output = mkTransform (ArgFun $ identity $ shapeType shr) f input output + where + ArrayR shr _ = argArrayR input + + mkBackpermute :: Arg env (Fun' (sh' -> sh)) + -> Arg env (In sh t) + -> Arg env (Out sh' t) + -> OperationAcc op env () + mkBackpermute f input output = mkTransform f (ArgFun $ identity tp) input output + where + ArrayR _ tp = argArrayR input + + -- TODO why does this exist? Transform is an artifact of the old fusion, so should never occur (or should simply be a composition of map and backpermute, in any order) + mkTransform :: Arg env (Fun' (sh' -> sh)) + -> Arg env (Fun' (s -> t)) + -> Arg env (In sh s) + -> Arg env (Out sh' t) + -> OperationAcc op env () + mkTransform (ArgFun f) (ArgFun g) (ArgArray _ repr sh buffers) output = mkGenerate (ArgFun h) output + where + h = g `compose` indexFunction repr sh buffers `compose` f + + mkGenerate :: Arg env (Fun' (sh -> t)) + -> Arg env (Out sh t) + -> OperationAcc op env () + + -- Copies the contents of the (prefix of the) input to the output. + -- The size of the input array is smaller than or equal to the size output. + mkShrink :: Arg env (In sh t) + -> Arg env (Out sh t) + -> OperationAcc op env () + mkShrink input@( ArgArray _ (ArrayR shr _) sh1 _) + output@(ArgArray _ _ sh2 _) = if isJust (matchVars sh1 sh2) + then mkCopy input output + else mkBackpermute (ArgFun $ identity $ shapeType shr) input output + + -- Copies a buffer. This is used before passing a buffer to a 'Mut' argument, + -- to make sure it is unique. This may be removed during fusion to facilitate + -- in-place updates + -- TODO: Perform one map per buffer, instead of one for the whole array + -- this gives more freedom to the ILP + mkCopy :: Arg env (In sh t) + -> Arg env (Out sh t) + -> OperationAcc op env () + mkCopy input@(ArgArray _ (ArrayR _ tp) _ _) output = mkMap (ArgFun $ identity tp) input output + + mkZipWith :: Arg env (Fun' (e1 -> e2 -> e3)) + -> Arg env (In sh e1) + -> Arg env (In sh e2) + -> Arg env (Out sh e3) + -> OperationAcc op env () + mkZipWith (ArgFun f) + in1@(ArgArray _ (ArrayR _ ty1) _ _) + in2@(ArgArray _ (ArrayR _ ty2) _ _) + (ArgArray _ outR@(ArrayR shr _) sh buf) + | DeclareVars lhs w k <- declareVars $ buffersR (TupRpair ty1 ty2) = + let + typ = TupRpair ty1 ty2 + zippedArg = ArgArray Out (ArrayR shr typ) (weakenVars w sh) (k weakenId) + zippedArg' = ArgArray In (ArrayR shr typ) (weakenVars w sh) (k weakenId) in + aletUnique lhs (lowerAlloc (ArrayR shr typ) (fromGrounds sh)) $ + alet LeftHandSideUnit + (mkZip (weaken w in1) (weaken w in2) zippedArg) $ + mkMap (ArgFun . uncurry' $ weakenArrayInstr w f) zippedArg' $ ArgArray Out outR (weakenVars w sh) (weakenVars w buf) + + mkZip :: Arg env (In sh e1) + -> Arg env (In sh e2) + -> Arg env (Out sh (e1, e2)) + -> OperationAcc op env () + mkZip in1@(ArgArray _ (ArrayR shr ty1) _ _) + in2@(ArgArray _ (ArrayR _ ty2) _ _) + (ArgArray _ (ArrayR _ _) sh (TupRpair buf1 buf2)) = + alet LeftHandSideUnit + (mkShrink in1 (ArgArray Out (ArrayR shr ty1) sh buf1)) + (mkShrink in2 (ArgArray Out (ArrayR shr ty2) sh buf2)) + + + mkZip _ _ _ = error "out buffer to mkZip is not a pair" + + + mkReplicate :: SliceIndex slix sl co sh + -> Arg env (In sl e) + -> Arg env (Out sh e) + -> OperationAcc op env () + mkReplicate slix = mkBackpermute (ArgFun $ extend slix) + + mkSlice :: SliceIndex slix sl co sh + -> Arg env (Var' slix) + -> Arg env (In sh e) + -> Arg env (Out sl e) + -> OperationAcc op env () + mkSlice sliceIx slix = mkBackpermute (ArgFun $ restrict sliceIx slix) + + mkPermute :: Maybe (Arg env (Fun' (e -> e -> e))) -- If no combination function, we use const without locking + -> Arg env (Mut sh' e) + -> Arg env (In sh (PrimMaybe (sh', e))) + -> OperationAcc op env () + + mkStencil :: StencilR sh e stencil + -> Arg env (Fun' (stencil -> e')) + -> Boundary env (Array sh e) + -> Arg env (In sh e) + -> Arg env (Out sh e') + -> OperationAcc op env () + -- Default implementation for mkStencil uses 'generate'. It includes the conditional code for the boundary on all elements, making it less efficient + -- than a specialized implementation. + -- + mkStencil sr (ArgFun f) boundary input output@(ArgArray _ (ArrayR _ tp) _ _) + = mkGenerate (ArgFun $ lowerStencilFun sr tp f (Just boundary) input) output + + mkStencil2 :: StencilR sh a stencil1 + -> StencilR sh b stencil2 + -> Arg env (Fun' (stencil1 -> stencil2 -> c)) + -> Boundary env (Array sh a) + -> Arg env (In sh a) + -> Boundary env (Array sh b) + -> Arg env (In sh b) + -> Arg env (Out sh c) + -> OperationAcc op env () + -- Default implementation for mkStencil uses 'generate'. It includes the conditional code for the boundary on all elements, making it less efficient + -- than a specialized implementation. + -- XXX: Could we define this in terms of mkStencil and a zip? Maybe only if the sizes of the two input arrays are equal? + -- + mkStencil2 sr1 sr2 (ArgFun f) boundary1 input1 boundary2 input2 output@(ArgArray _ (ArrayR _ tp) _ _) + = mkGenerate (ArgFun $ lowerStencil2Fun sr1 sr2 tp f (Just boundary1) input1 (Just boundary2) input2) output + + mkFold :: Arg env (Fun' (e -> e -> e)) + -> Maybe (Arg env (Exp' e)) + -> Arg env (In (sh, Int) e) + -> Arg env (Out sh e) + -> OperationAcc op env () + mkFold f def input@(ArgArray _ _ _ _) output = mkDefaultFoldSequential f def input output + + mkFoldSeg :: IntegralType i + -> Arg env (Fun' (e -> e -> e)) + -> Maybe (Arg env (Exp' e)) + -> Arg env (In (sh, Int) e) + -> Arg env (In DIM1 i) + -> Arg env (Out (sh, Int) e) + -> OperationAcc op env () + -- Default implementation using generate. It is sequential per segment, which is inefficient for some backends. + mkFoldSeg itp f def input segments = mkGenerate (ArgFun $ mkDefaultFoldSegFunction itp f def input segments) + + mkScan :: Direction + -> Arg env (Fun' (e -> e -> e)) + -> Maybe (Arg env (Exp' e)) + -> Arg env (In (sh, Int) e) + -> Arg env (Out (sh, Int) e) + -> OperationAcc op env () + mkScan dir f (Just def) argIn@(ArgArray _ (ArrayR shr tp) _ _) argOut@(ArgArray _ _ shOut _) + | DeclareVars lhsTmp kTmp valueTmp <- declareVars $ buffersR tp -- Allocate buffers to prepend the default value + = let + shOut' = weakenVars kTmp shOut + + argG = ArgFun $ mkDefaultScanPrepend dir (weaken kTmp def) (weaken kTmp argIn) + argTmp' = ArgArray Out (ArrayR shr tp) shOut' (valueTmp weakenId) + argTmp = ArgArray In (ArrayR shr tp) shOut' (valueTmp weakenId) + in + aletUnique lhsTmp (lowerAlloc (ArrayR shr tp) $ groundToExpVar (shapeType shr) shOut) + $ alet (LeftHandSideWildcard TupRunit) (mkGenerate argG argTmp') + $ mkScan dir (weaken kTmp f) Nothing argTmp (weaken kTmp argOut) + mkScan dir f Nothing (ArgArray _ (ArrayR shr tp) sh input) argOut + -- (inc, tmp) = awhile (\(inc, _) -> inc < (n+1) / 2) (\(inc, a) -> (inc*2, generate {\i -> reduce i and i +/- inc in a})) input + -- generate {\i -> reduce i and i +/- inc in tmp} + -- + -- Note that the last iteration of the loop is decoupled, as that will output into argOut, + -- instead of a temporary array. + | DeclareVars lhsTmp kTmp valueTmp <- declareVars $ buffersR tp + , DeclareVars lhsAlloc kAlloc valueAlloc <- declareVars $ buffersR tp + = let + lhs = LeftHandSidePair (LeftHandSideSingle $ GroundRscalar scalarTypeInt) lhsTmp + argTmp = ArgArray In (ArrayR shr tp) (weakenVars (weakenSucc $ weakenSucc kTmp) sh) (valueTmp weakenId) + reduce = ArgFun $ mkDefaultScanFunction dir (weaken kTmp (Var (GroundRscalar $ scalarTypeInt) ZeroIdx)) (weaken (weakenSucc $ weakenSucc kTmp) f) argTmp + + n = case sh of + TupRpair _ n' -> paramsIn (TupRsingle scalarTypeInt) $ weakenVars (weakenSucc $ weakenSucc weakenId) n' + _ -> error "Impossible pair" + + -- Awhile condition + c = Alam (LeftHandSidePair (LeftHandSideSingle $ GroundRscalar scalarTypeInt) $ LeftHandSideWildcard $ buffersR tp) + $ Abody + $ Compute + $ mkBinary (PrimCmp singleType CmpLt) (paramIn' $ Var scalarTypeInt ZeroIdx) + $ mkBinary + (PrimBShiftR integralType) + (mkBinary (PrimAdd numType) n (mkConstant (TupRsingle scalarTypeInt) 1)) + (mkConstant (TupRsingle scalarTypeInt) 1) -- (n+1)/2 + + argAlloc = ArgArray Out (ArrayR shr tp) (weakenVars (weakenSucc $ weakenSucc $ kAlloc .> kTmp) sh) (valueAlloc weakenId) + -- Awhile step function + g = Alam lhs + $ Abody + $ aletUnique lhsAlloc (lowerAlloc (ArrayR shr tp) $ groundToExpVar (shapeType shr) $ weakenVars (weakenSucc $ weakenSucc kTmp) sh) + $ alet (LeftHandSideWildcard TupRunit) (mkGenerate (weaken kAlloc reduce) argAlloc) + $ alet (LeftHandSideSingle $ GroundRscalar scalarTypeInt) + (Compute $ mkBinary (PrimBShiftL integralType) (paramIn' $ weaken (kAlloc .> kTmp) $ Var scalarTypeInt ZeroIdx) (mkConstant (TupRsingle scalarTypeInt) 1)) -- n * 2 + $ Return $ TupRsingle (Var (GroundRscalar scalarTypeInt) ZeroIdx) `TupRpair` valueAlloc (weakenSucc weakenId) + in + alet (LeftHandSideSingle $ GroundRscalar scalarTypeInt) (Compute $ mkConstant (TupRsingle scalarTypeInt) 1) + $ alet lhs (Awhile (shared $ lhsToTupR lhs) c g (TupRsingle (Var (GroundRscalar scalarTypeInt) ZeroIdx) `TupRpair` weakenVars (weakenSucc weakenId) input)) + $ mkGenerate reduce (weaken (weakenSucc $ weakenSucc kTmp) argOut) + + mkScan' :: Direction + -> Arg env (Fun' (e -> e -> e)) + -> Arg env (Exp' e) + -> Arg env (In (sh, Int) e) + -> Arg env (Out (sh, Int) e) + -> Arg env (Out sh e) + -> OperationAcc op env () + mkScan' dir f def input@(ArgArray _ (ArrayR shr@(ShapeRsnoc shr') tp) sh _) out1 out2 + | DeclareVars lhsTmp kTmp valueTmp <- declareVars $ buffersR tp -- Allocate buffers for the output of scan + , DeclareVars lhsSh _ valueSh <- declareVars $ shapeType shr' -- Used in the function in backpermute + = let + (x, y) = case sh of + TupRpair x' y' -> (x', y') + _ -> error "Impossible pair" + + shTmp' = weakenVars (weakenSucc weakenId) x `TupRpair` TupRsingle (Var (GroundRscalar scalarTypeInt) ZeroIdx) + shTmp = weakenVars kTmp shTmp' + k = kTmp .> weakenSucc weakenId + argTmp = ArgArray In (ArrayR shr tp) shTmp (valueTmp weakenId) + argTmp' = ArgArray Out (ArrayR shr tp) shTmp (valueTmp weakenId) + + -- Fills 'out2' + argG = ArgFun $ Lam lhsSh $ Body $ Pair (expVars $ valueSh weakenId) $ case dir of + LeftToRight -> paramsIn (TupRsingle scalarTypeInt) $ weakenVars k y + RightToLeft -> mkConstant (TupRsingle scalarTypeInt) 0 + + -- Fills 'out1' in case of a RightToLeft scan + argH = ArgFun + $ Lam (LeftHandSidePair lhsSh $ LeftHandSideSingle scalarTypeInt) + $ Body + $ Pair + (expVars $ valueSh $ weakenSucc weakenId) + (mkBinary (PrimAdd numType) (Evar (Var scalarTypeInt ZeroIdx)) (mkConstant (TupRsingle scalarTypeInt) 0)) + in + alet (LeftHandSideSingle $ GroundRscalar scalarTypeInt) (Compute $ mkBinary (PrimAdd numType) (paramsIn (TupRsingle scalarTypeInt) y) $ mkConstant (TupRsingle scalarTypeInt) 1) + $ aletUnique lhsTmp (lowerAlloc (ArrayR shr tp) $ groundToExpVar (shapeType shr) shTmp') + $ alet (LeftHandSideWildcard TupRunit) (mkScan dir (weaken k f) (Just $ weaken k def) (weaken k input) argTmp') + $ alet (LeftHandSideWildcard TupRunit) (case dir of + LeftToRight -> mkShrink argTmp $ weaken k out1 + RightToLeft -> mkBackpermute argH argTmp $ weaken k out1 + ) + $ alet (LeftHandSideWildcard TupRunit) (mkBackpermute argG argTmp $ weaken k out2) + $ Return TupRunit + + mkForeign :: Foreign asm + => ArraysR bs + -> asm (as -> bs) + -> GroundVars env (LoweredArrays as) + -> Maybe (OperationAcc op env (LoweredArrays bs)) + mkForeign _ _ _ = Nothing + + -- Whether it is prefered to not use scalar operations. + -- When this is true, a 0-dimensional generate or map will + -- be converted to Compute nodes, instead of Exec nodes. + lowerPreferNoScalar :: Bool + lowerPreferNoScalar = True + +lower :: (HasCallStack, LowerAcc op) => Named.Acc a -> OperationAcc op () (LoweredArrays a) +lower = lowerOpenAcc Empty + +lowerAfun :: (HasCallStack, LowerAcc op) => Named.Afun a -> OperationAfun op () (LoweredAfun a) +lowerAfun = lowerOpenAfun Empty + +type BEnv benv = Env (ArrayDescriptor benv) + +lowerOpenAcc :: forall op benv aenv a. + (HasCallStack, LowerAcc op) + => BEnv benv aenv + -> Named.OpenAcc aenv a + -> OperationAcc op benv (LoweredArrays a) +lowerOpenAcc env = travA + where + travE :: Named.Exp aenv t -> Exp benv t + travE = lowerExp env + + travA :: Named.OpenAcc aenv arr -> OperationAcc op benv (LoweredArrays arr) + travA (Named.OpenAcc acc) = case acc of + Named.Alet lhs bnd a + | LoweredLHS env' lhs' <- lowerLHS env lhs -> alet lhs' (travA bnd) (lowerOpenAcc env' a) + Named.Avar (Var _ ix) + | ArrayDescriptor _ sh bf <- prj' ix env -> Return (TupRpair sh bf) + Named.Apair a b -> pair (travA a) (travA b) + Named.Anil -> Return TupRunit + Named.Manifest as -> makeManifest $ travA as + Named.Aassert msg cond as -> + alet + (LeftHandSideSingle $ GroundRscalar scalarTypeWord8) + (Aassert msg $ travE cond) + $ Fence (IdxSet.singleton ZeroIdx) + $ lowerOpenAcc (weakenBEnv (weakenSucc weakenId) env) as + Named.Aassume cond as -> + alet + (LeftHandSideSingle $ GroundRscalar scalarTypeWord8) + (Aassume $ travE cond) + $ Fence (IdxSet.singleton ZeroIdx) + $ lowerOpenAcc (weakenBEnv (weakenSucc weakenId) env) as + Named.Aforeign repr asm (Named.Alam lhsA _) arg + | DeclareVars lhs _ value <- declareVars $ lowerArraysR $ lhsToTupR lhsA + , Just a <- mkForeign repr asm $ value weakenId + -> alet lhs (travA arg) a + Named.Aforeign repr _ f arg -> case f of + Named.Alam lhs (Named.Abody a) + | Exists lhs' <- rebuildLHS lhs + -> travA $ Named.OpenAcc $ Named.Alet lhs' arg (weaken (sinkWithLHS lhs lhs' weakenEmpty) a) + Named.Abody a -> arraysRFunctionImpossible $ Named.arraysR a + Named.Alam _ Named.Alam{} -> arraysRFunctionImpossible repr + Named.Acond c t f -> + let + lhs = LeftHandSideSingle $ GroundRscalar scalarTypeWord8 + env' = weakenBEnv (weakenSucc weakenId) env + var = Var scalarTypeWord8 ZeroIdx + in alet lhs (Compute $ travE c) $ Acond var (lowerOpenAcc env' t) (lowerOpenAcc env' f) + Named.Awhile c f i + | repr <- Named.arraysR acc + , Refl <- loweredAfunIsBody repr + , DeclareVars lhs k value <- declareVars $ lowerArraysR repr -> + let + bufferType = GroundRbuffer $ scalarTypeWord8 + lhsScalarBool = LeftHandSidePair (LeftHandSideWildcard TupRunit) (LeftHandSideSingle bufferType) + env' = weakenBEnv k env + c' = case lowerOpenAfun @op env' c of + Alam lhs' (Abody body) -> Alam lhs' $ Abody $ alet lhsScalarBool body $ Compute $ ArrayInstr (Index $ Var bufferType ZeroIdx) $ Const scalarTypeInt 0 + Abody _ -> error "It's a long time since we last met" + f' = lowerOpenAfun env' f + in alet lhs (travA i) $ awhileRemoveScalars repr (shared $ lowerArraysR repr) c' f' (value weakenId) + Named.Use repr array -> lowerUse repr array + Named.Unit tp e + | DeclareVars lhs _ value <- declareVars tp -> + let + go :: forall benv' t. ExpVars benv' t -> OperationAcc op benv' (Buffers t) + go (TupRpair v1 v2) = go v1 `pairUnique` go v2 + go TupRunit = Return TupRunit + go (TupRsingle v@(Var repr _)) + | Refl <- reprIsSingle @ScalarType @t @Buffer repr = Unit v + in + alet (mapLeftHandSide GroundRscalar lhs) (Compute $ travE e) $ pair (Return TupRunit) $ go (value weakenId) + -- XXX: Should we check whether the sizes are equal? We could add an Assert constructor to PreOpenAcc, + -- which we can use to verify the sizes at runtime. + -- + Named.Reshape shrOut shExp a + | ArrayR shrIn tp <- Named.arrayR a + , DeclareVars lhsShIn kShIn valueShIn <- declareVars $ mapTupR GroundRscalar $ shapeType shrIn + , DeclareVars lhsIn kIn valueIn <- declareVars $ buffersR tp + , DeclareVars lhsShOut kShOut valueShOut <- declareVars $ shapeType shrOut + , DeclareVars lhsOut kOut valueOut <- declareVars $ buffersR tp + , DeclareVars lhsIx _ valueIx <- declareVars $ shapeType shrOut -> + let + shOut = mapVars GroundRscalar $ valueShOut kOut + bfOut = valueOut weakenId + argF = ArgFun + $ Lam lhsIx + $ Body + $ FromIndex shrIn (paramsIn (shapeType shrIn) $ valueShIn $ kOut .> kShOut .> kIn) + $ ToIndex shrOut (paramsIn' $ valueShOut kOut) + $ expVars $ valueIx weakenId + argIn = ArgArray In (ArrayR shrIn tp) (valueShIn (kOut .> kShOut .> kIn)) (valueIn (kOut .> kShOut)) + argOut = ArgArray Out (ArrayR shrOut tp) shOut bfOut + in + alet (LeftHandSidePair lhsShIn lhsIn) (travA a) + $ alet (mapLeftHandSide GroundRscalar lhsShOut) (Compute $ lowerExp (weakenBEnv (kIn .> kShIn) env) shExp) + $ aletUnique lhsOut (lowerAlloc (ArrayR shrOut tp) (valueShOut weakenId)) + $ alet (LeftHandSideWildcard TupRunit) (mkBackpermute argF argIn argOut) + $ Return (shOut `TupRpair` bfOut) + + -- The remaining constructors of Named.OpenAcc are compiled into operations through the mkXX functions + -- the type class LowerAcc. We must here allocate the output arrays of the appropriate size and call + -- the corresponding mkXX function. + -- + + Named.Permute c (def :: Named.OpenAcc aenv (Array sh' e)) src + | ArrayR shr' tp <- Named.arrayR def + , ArrayR shr tsht@(tag `TupRpair` (TupRunit `TupRpair` ((TupRunit `TupRpair` sht) `TupRpair` a))) <- Named.arrayR src + , DeclareVars lhsSh' kSh' valueSh' <- declareVars $ shapeType shr' + , DeclareVars lhsDef kDef valueDef <- declareVars $ buffersR tp + -- Clone defaults array, to make sure that it is unique. The clone may be removed in a later pass. + , DeclareVars lhsOut kOut valueOut <- declareVars $ buffersR tp + , DeclareVars lhsSh kSh valueSh <- declareVars $ shapeType shr + , DeclareVarsPrimMaybe lhsSrc kSrc valueSrc <- fixprimmaybepermute @_ @sh' @e $ declareVars $ buffersR tsht -> + let + srcR = tag `TupRpair` (TupRunit `TupRpair` (sht `TupRpair` a)) + lhs' = LeftHandSidePair (mapLeftHandSide GroundRscalar lhsSh') lhsDef + lhs = LeftHandSidePair (mapLeftHandSide GroundRscalar lhsSh) lhsSrc + sh' = mapVars GroundRscalar $ valueSh' (kSrc .> kSh .> kOut .> kDef) + sh = mapVars GroundRscalar $ valueSh kSrc + env' = weakenBEnv (kSrc .> kSh .> kOut .> kDef .> kSh') env + valOut = valueOut $ kSrc .> kSh + argDef = ArgArray In (ArrayR shr' tp) sh' (valueDef $ kSrc .> kSh .> kOut) + argOut = ArgArray Out (ArrayR shr' tp) sh' valOut + argMut = ArgArray Mut (ArrayR shr' tp) sh' valOut + argSrc = ArgArray In (ArrayR shr srcR) sh (valueSrc weakenId) + argC = ArgFun . lowerFun env' <$> c + in + alet lhs' (travA def) + $ aletUnique lhsOut (lowerAlloc (ArrayR shr' tp) (valueSh' kDef)) + $ alet lhs (lowerOpenAcc (weakenBEnv (kOut .> kDef .> kSh') env) src) + $ alet (LeftHandSideWildcard TupRunit) (mkCopy argDef argOut) + $ alet (LeftHandSideWildcard TupRunit) (mkPermute argC argMut argSrc) + $ Return (sh' `TupRpair` valOut) + Named.Permute {} -> error "impossible tupr" + + Named.Generate (ArrayR shr tp) sh f + -- generate sh (\_ -> undef) is a pattern commonly used for a permute + -- if the permute writes to all indices. We optimize for this pattern here. + -- In this case, we only need to allocate the result; we don't need to + -- execute a kernel. + | Lam _ (Body body) <- f + , isUndef body + , DeclareVars lhsSh _ valueSh <- declareVars $ shapeType shr + , DeclareVars lhsBf kBf valueBf <- declareVars $ buffersR tp -> + alet (mapLeftHandSide GroundRscalar lhsSh) (Compute $ travE sh) + $ aletUnique lhsBf (lowerAlloc (ArrayR shr tp) (valueSh weakenId)) + $ Return + (mapVars GroundRscalar (valueSh kBf) `TupRpair` valueBf weakenId) + Named.Generate (ArrayR ShapeRz tp) _ (Lam (LeftHandSideWildcard _) (Body expr)) + | lowerPreferNoScalar @op -> + travA $ Named.OpenAcc $ Named.Unit tp expr + Named.Generate (ArrayR shr tp) sh f + | DeclareVars lhsSh kSh valueSh <- declareVars $ shapeType shr + , DeclareVars lhsBf kBf valueBf <- declareVars $ buffersR tp -> + let + sh' = mapVars GroundRscalar $ valueSh kBf + argF = ArgFun $ lowerFun (weakenBEnv (kBf .> kSh) env) f + argOut = ArgArray Out (ArrayR shr tp) sh' (valueBf weakenId) + in + alet (mapLeftHandSide GroundRscalar lhsSh) (Compute $ travE sh) + $ aletUnique lhsBf (lowerAlloc (ArrayR shr tp) (valueSh weakenId)) + $ alet (LeftHandSideWildcard TupRunit) (mkGenerate argF argOut) + $ Return (sh' `TupRpair` valueBf weakenId) + + Named.Transform (ArrayR shrOut tpOut) shExp idxTransform valueTransform a + | ArrayR shrIn tpIn <- Named.arrayR a + , DeclareVars lhsShIn kShIn valueShIn <- declareVars $ mapTupR GroundRscalar $ shapeType shrIn + , DeclareVars lhsIn kIn valueIn <- declareVars $ buffersR tpIn + , DeclareVars lhsShOut kShOut valueShOut <- declareVars $ shapeType shrOut + , DeclareVars lhsOut kOut valueOut <- declareVars $ buffersR tpOut -> + let + shOut = mapVars GroundRscalar $ valueShOut kOut + bfOut = valueOut weakenId + env' = weakenBEnv (kOut .> kShOut .> kIn .> kShIn) env + argIdxF = ArgFun $ lowerFun env' idxTransform + argValF = ArgFun $ lowerFun env' valueTransform + argIn = ArgArray In (ArrayR shrIn tpIn) (valueShIn (kOut .> kShOut .> kIn)) (valueIn (kOut .> kShOut)) + argOut = ArgArray Out (ArrayR shrOut tpOut) shOut bfOut + in + alet (LeftHandSidePair lhsShIn lhsIn) (travA a) + $ alet (mapLeftHandSide GroundRscalar lhsShOut) (Compute $ lowerExp (weakenBEnv (kIn .> kShIn) env) shExp) + $ aletUnique lhsOut (lowerAlloc (ArrayR shrOut tpOut) (valueShOut weakenId)) + $ alet (LeftHandSideWildcard TupRunit) (mkTransform argIdxF argValF argIn argOut) + $ Return (shOut `TupRpair` bfOut) + + Named.Map tp (Lam lhs (Body expr)) a + | lowerPreferNoScalar @op + , repr@(ArrayR ShapeRz _) <- Named.arrayR a -> + travA $ Named.OpenAcc + $ Named.Alet (LeftHandSideSingle repr) a + $ Named.OpenAcc + $ Named.Unit tp + $ Let lhs (ArrayInstr (Named.Index (Var repr ZeroIdx)) Nil) + $ mapArrayInstr (weaken $ weakenSucc weakenId) expr + + Named.Map _ (Lam lhs (Body e)) a + | Just vars <- extractExpVars e + , ArrayR shr oldTp <- Named.arrayR a + , DeclareVars lhsSh _ valueSh <- declareVars $ mapTupR GroundRscalar $ shapeType shr + , DeclareVars lhsIn kIn valueIn <- declareVars $ buffersR oldTp -> + alet (LeftHandSidePair lhsSh lhsIn) (travA a) + $ Return (valueSh kIn `TupRpair` lowerUnzip (valueIn weakenId) lhs vars) + + Named.Map tp f a + | ArrayR shr oldTp <- Named.arrayR a + , DeclareVars lhsSh kSh valueSh <- declareVars $ shapeType shr + , DeclareVars lhsIn kIn valueIn <- declareVars $ buffersR oldTp + , DeclareVars lhsOut kOut valueOut <- declareVars $ buffersR tp -> + let + lhs = LeftHandSidePair (mapLeftHandSide GroundRscalar lhsSh) lhsIn + sh = mapVars GroundRscalar $ valueSh (kOut .> kIn) + argF = ArgFun $ lowerFun (weakenBEnv (kOut .> kIn .> kSh) env) f + argIn = ArgArray In (ArrayR shr oldTp) sh (valueIn kOut) + argOut = ArgArray Out (ArrayR shr tp) sh (valueOut weakenId) + in + alet lhs (travA a) + $ aletUnique lhsOut (lowerAlloc (ArrayR shr tp) (valueSh kIn)) + $ alet (LeftHandSideWildcard TupRunit) (mkMap argF argIn argOut) + $ Return (sh `TupRpair` valueOut weakenId) + + Named.Backpermute shrOut shExp f a + | ArrayR shrIn tp <- Named.arrayR a + , DeclareVars lhsShIn kShIn valueShIn <- declareVars $ mapTupR GroundRscalar $ shapeType shrIn + , DeclareVars lhsIn kIn valueIn <- declareVars $ buffersR tp + , DeclareVars lhsShOut kShOut valueShOut <- declareVars $ shapeType shrOut + , DeclareVars lhsOut kOut valueOut <- declareVars $ buffersR tp -> + let + shOut = mapVars GroundRscalar $ valueShOut kOut + bfOut = valueOut weakenId + env' = weakenBEnv (kOut .> kShOut .> kIn .> kShIn) env + argF = ArgFun $ lowerFun env' f + argIn = ArgArray In (ArrayR shrIn tp) (valueShIn (kOut .> kShOut .> kIn)) (valueIn (kOut .> kShOut)) + argOut = ArgArray Out (ArrayR shrOut tp) shOut bfOut + in + alet (LeftHandSidePair lhsShIn lhsIn) (travA a) + $ alet (mapLeftHandSide GroundRscalar lhsShOut) (Compute $ lowerExp (weakenBEnv (kIn .> kShIn) env) shExp) + $ aletUnique lhsOut (lowerAlloc (ArrayR shrOut tp) (valueShOut weakenId)) + $ alet (LeftHandSideWildcard TupRunit) (mkBackpermute argF argIn argOut) + $ Return (shOut `TupRpair` bfOut) + + -- Zip (or a combination of unzip & zip) + Named.ZipWith _ (Lam elhs1 (Lam elhs2 (Body e))) a1 a2 + | Just vars <- extractExpVars e + , ArrayR shr t1 <- Named.arrayR a1 + , ArrayR _ t2 <- Named.arrayR a2 + , DeclareVars lhsSh1 kSh1 valueSh1 <- declareVars $ shapeType shr + , DeclareVars lhsIn1 kIn1 valueIn1 <- declareVars $ buffersR t1 + , DeclareVars lhsSh2 kSh2 valueSh2 <- declareVars $ shapeType shr + , DeclareVars lhsIn2 kIn2 valueIn2 <- declareVars $ buffersR t2 + , DeclareVars lhsSh kSh valueSh <- declareVars $ shapeType shr + , DeclareVars lhsOut1 kOut1 valueOut1 <- declareVars $ buffersR t1 + , DeclareVars lhsOut2 kOut2 valueOut2 <- declareVars $ buffersR t2 -> + let + lhs1 = LeftHandSidePair (mapLeftHandSide GroundRscalar lhsSh1) lhsIn1 + lhs2 = LeftHandSidePair (mapLeftHandSide GroundRscalar lhsSh2) lhsIn2 + sh1 = mapVars GroundRscalar $ valueSh1 $ kOut2 .> kOut1 .> kSh .> kIn2 .> kSh2 .> kIn1 + sh2 = mapVars GroundRscalar $ valueSh2 $ kOut2 .> kOut1 .> kSh .> kIn2 + sh = mapVars GroundRscalar $ valueSh (kOut2 .> kOut1) + argIn1 = ArgArray In (ArrayR shr t1) sh1 (valueIn1 $ kOut2 .> kOut1 .> kSh .> kIn2 .> kSh2) + argIn2 = ArgArray In (ArrayR shr t2) sh2 (valueIn2 $ kOut2 .> kOut1 .> kSh) + argOut1 = ArgArray Out (ArrayR shr t1) sh (valueOut1 kOut2) + argOut2 = ArgArray Out (ArrayR shr t2) sh (valueOut2 weakenId) + in + alet lhs1 (travA a1) + $ alet lhs2 (lowerOpenAcc (weakenBEnv (kIn1 .> kSh1) env) a2) + $ alet (mapLeftHandSide GroundRscalar lhsSh) (Compute $ mkIntersect shr (valueSh1 $ kIn2 .> kSh2 .> kIn1) (valueSh2 kIn2)) + $ aletUnique lhsOut1 (lowerAlloc (ArrayR shr t1) (valueSh weakenId)) + $ aletUnique lhsOut2 (lowerAlloc (ArrayR shr t2) (valueSh kOut1)) + $ alet (LeftHandSideWildcard TupRunit) (mkShrink argIn1 argOut1) + $ alet (LeftHandSideWildcard TupRunit) (mkShrink argIn2 argOut2) + $ Return (sh `TupRpair` lowerUnzip (valueOut1 kOut2 `TupRpair` valueOut2 weakenId) (LeftHandSidePair elhs1 elhs2) vars) + + Named.ZipWith tp f a1 a2 + | ArrayR shr t1 <- Named.arrayR a1 + , ArrayR _ t2 <- Named.arrayR a2 + , DeclareVars lhsSh1 kSh1 valueSh1 <- declareVars $ shapeType shr + , DeclareVars lhsIn1 kIn1 valueIn1 <- declareVars $ buffersR t1 + , DeclareVars lhsSh2 kSh2 valueSh2 <- declareVars $ shapeType shr + , DeclareVars lhsIn2 kIn2 valueIn2 <- declareVars $ buffersR t2 + , DeclareVars lhsSh kSh valueSh <- declareVars $ shapeType shr + , DeclareVars lhsOut kOut valueOut <- declareVars $ buffersR tp + , DeclareVars lhsOut' kOut' valueOut' <- declareVars $ buffersR tp -> + let + lhs1 = LeftHandSidePair (mapLeftHandSide GroundRscalar lhsSh1) lhsIn1 + lhs2 = LeftHandSidePair (mapLeftHandSide GroundRscalar lhsSh2) lhsIn2 + sh1 = mapVars GroundRscalar $ valueSh1 $ kOut .> kSh .> kIn2 .> kSh2 .> kIn1 + sh1' = mapVars GroundRscalar $ valueSh1 $ kOut' .> kIn2 .> kSh2 .> kIn1 + sh2 = mapVars GroundRscalar $ valueSh2 $ kOut .> kSh .> kIn2 + sh2' = mapVars GroundRscalar $ valueSh2 $ kOut' .> kIn2 + sh = mapVars GroundRscalar $ valueSh kOut + argF = ArgFun $ lowerFun (weakenBEnv (kOut .> kSh .> kIn2 .> kSh2 .> kIn1 .> kSh1) env) f + argF' = ArgFun $ lowerFun (weakenBEnv (kOut' .> kIn2 .> kSh2 .> kIn1 .> kSh1) env) f + argIn1 = ArgArray In (ArrayR shr t1) sh1 (valueIn1 $ kOut .> kSh .> kIn2 .> kSh2) + argIn1' = ArgArray In (ArrayR shr t1) sh1' (valueIn1 $ kOut' .> kIn2 .> kSh2) + argIn2 = ArgArray In (ArrayR shr t2) sh2 (valueIn2 $ kOut .> kSh) + argIn2' = ArgArray In (ArrayR shr t2) sh2' (valueIn2 $ kOut') + argOut = ArgArray Out (ArrayR shr tp) sh (valueOut weakenId) + argOut' = ArgArray Out (ArrayR shr tp) sh1' (valueOut' weakenId) + in + alet lhs1 (travA a1) + $ alet lhs2 (lowerOpenAcc (weakenBEnv (kIn1 .> kSh1) env) a2) + $ case matchVars (valueSh1 $ kIn2 .> kSh2 .> kIn1) (valueSh2 kIn2) of + Just Refl -> + aletUnique lhsOut' (lowerAlloc (ArrayR shr tp) (valueSh2 kIn2)) + $ alet LeftHandSideUnit (mkZipWith argF' argIn1' argIn2' argOut') + $ Return (sh1' `TupRpair` valueOut' weakenId) + Nothing -> + alet (mapLeftHandSide GroundRscalar lhsSh) (Compute $ mkIntersect shr (valueSh1 $ kIn2 .> kSh2 .> kIn1) (valueSh2 kIn2)) + $ aletUnique lhsOut (lowerAlloc (ArrayR shr tp) (valueSh weakenId)) + $ alet LeftHandSideUnit (mkZipWith argF argIn1 argIn2 argOut) + $ Return (sh `TupRpair` valueOut weakenId) + Named.Slice sliceIx a slix + | ArrayR shr tp <- Named.arrayR a + , slr <- sliceShapeR sliceIx + , DeclareVars lhsSh kSh valueSh <- declareVars $ shapeType shr + , DeclareVars lhsIn kIn valueIn <- declareVars $ buffersR tp + , DeclareVars lhsSlix kSlix valueSlix <- declareVars $ sliceEltR sliceIx + , DeclareVars lhsSl kSl valueSl <- declareVars $ shapeType slr + , DeclareVars lhsOut kOut valueOut <- declareVars $ buffersR tp -> + let + lhs = LeftHandSidePair (mapLeftHandSide GroundRscalar lhsSh) lhsIn + slix' = lowerExp (weakenBEnv (kIn .> kSh) env) slix + sh = mapVars GroundRscalar $ valueSh $ kOut .> kSl .> kSlix .> kIn + sl = mapVars GroundRscalar $ valueSl kOut + argSlix = ArgVar $ valueSlix $ kOut .> kSl + argIn = ArgArray In (ArrayR shr tp) sh (valueIn $ kOut .> kSl .> kSlix) + argOut = ArgArray Out (ArrayR slr tp) sl (valueOut weakenId) + in + alet lhs (travA a) + $ alet (mapLeftHandSide GroundRscalar lhsSlix) (Compute slix') + $ alet (mapLeftHandSide GroundRscalar lhsSl) (Compute $ indexSlice sliceIx (paramsIn' $ valueSh $ kSlix .> kIn)) + $ aletUnique lhsOut (lowerAlloc (ArrayR slr tp) (valueSl weakenId)) + $ alet (LeftHandSideWildcard TupRunit) (mkSlice sliceIx argSlix argIn argOut) + $ Return (sl `TupRpair` valueOut weakenId) + + Named.Replicate sliceIx slix a + | ArrayR slr tp <- Named.arrayR a + , shr <- sliceDomainR sliceIx + , DeclareVars lhsSl kSl valueSl <- declareVars $ shapeType slr + , DeclareVars lhsIn kIn valueIn <- declareVars $ buffersR tp + , DeclareVars lhsSlix kSlix valueSlix <- declareVars $ sliceEltR sliceIx + , DeclareVars lhsSh kSh valueSh <- declareVars $ shapeType shr + , DeclareVars lhsOut kOut valueOut <- declareVars $ buffersR tp -> + let + lhs = LeftHandSidePair (mapLeftHandSide GroundRscalar lhsSl) lhsIn + slix' = lowerExp (weakenBEnv (kIn .> kSl) env) slix + sl = mapVars GroundRscalar $ valueSl $ kOut .> kSh .> kSlix .> kIn + sh = mapVars GroundRscalar $ valueSh kOut + argIn = ArgArray In (ArrayR slr tp) sl (valueIn $ kOut .> kSh .> kSlix) + argOut = ArgArray Out (ArrayR shr tp) sh (valueOut weakenId) + in + alet lhs (travA a) + $ alet (mapLeftHandSide GroundRscalar lhsSlix) (Compute slix') + $ alet (mapLeftHandSide GroundRscalar lhsSh) (Compute $ indexFull sliceIx (paramsIn' $ valueSlix weakenId) (paramsIn' $ valueSl $ kSlix .> kIn)) + $ aletUnique lhsOut (lowerAlloc (ArrayR shr tp) (valueSh weakenId)) + $ alet (LeftHandSideWildcard TupRunit) (mkReplicate sliceIx argIn argOut) + $ Return (sh `TupRpair` valueOut weakenId) + + Named.Fold f def a + | ArrayR (shapeInit -> shr) tp <- Named.arrayR a + , DeclareVars lhsSh kSh valueSh <- declareVars $ shapeType $ ShapeRsnoc shr + , DeclareVars lhsIn kIn valueIn <- declareVars $ buffersR tp + , DeclareVars lhsOut kOut valueOut <- declareVars $ buffersR tp -> + let + shIn = mapVars GroundRscalar $ valueSh (kOut .> kIn) + shOut' = case valueSh kIn of + TupRpair sh _ -> sh + TupRsingle _ -> error "Impossible pair" + shOut = case shIn of + TupRpair sh _ -> sh + TupRsingle _ -> error "Impossible pair" + k = kOut .> kIn .> kSh + argF = ArgFun $ lowerFun (weakenBEnv k env) f + argDef = fmap (ArgExp . lowerExp (weakenBEnv k env)) def + argIn = ArgArray In (ArrayR (ShapeRsnoc shr) tp) shIn (valueIn kOut) + argOut = ArgArray Out (ArrayR shr tp) shOut (valueOut weakenId) + in + alet (LeftHandSidePair (mapLeftHandSide GroundRscalar lhsSh) lhsIn) (travA a) + $ aletUnique lhsOut (lowerAlloc (ArrayR shr tp) shOut') + $ alet (LeftHandSideWildcard TupRunit) (mkFold argF argDef argIn argOut) + $ Return (shOut `TupRpair` valueOut weakenId) + + Named.FoldSeg (i :: IntegralType i) f def a segments + | ArrayR shr tp <- Named.arrayR a + , Refl <- reprIsSingle @ScalarType @i @Buffer (SingleScalarType $ NumSingleType $ IntegralNumType i) + , DeclareVars lhsSh kSh valueSh <- declareVars $ shapeType shr + , DeclareVars lhsIn kIn valueIn <- declareVars $ buffersR tp + , DeclareVars lhsOut kOut valueOut <- declareVars $ buffersR tp -> + let + itp = SingleScalarType $ NumSingleType $ IntegralNumType i + lhsSeg + = LeftHandSidePair + -- DIM1 = ((), Int) + (LeftHandSidePair (LeftHandSideWildcard TupRunit) (LeftHandSideSingle $ GroundRscalar scalarTypeInt)) + -- Buffer of integral type i + $ LeftHandSideSingle $ GroundRbuffer itp + lhsSegCount = LeftHandSideSingle $ GroundRscalar scalarTypeInt + kSeg = weakenSucc $ weakenSucc weakenId + kSegCount = weakenSucc weakenId + + shIn = mapVars GroundRscalar $ valueSh (kOut .> kSegCount .> kSeg .> kIn) + shIn' = valueSh (kSegCount .> kSeg .> kIn) + + -- Change the innermost size of the input by the number of segments + -- (which we compute in the let binding with lhsSegCount). + shOut' = case shIn' of + TupRpair s _ -> TupRpair s $ TupRsingle + $ Var scalarTypeInt ZeroIdx + TupRsingle var -> pairImpossible var + + shOut = mapTupR (\(Var t idx) -> Var (GroundRscalar t) $ weaken kOut idx) shOut' + + k = kOut .> kSegCount .> kSeg .> kIn .> kSh + argF = ArgFun $ lowerFun (weakenBEnv k env) f + argDef = fmap (ArgExp . lowerExp (weakenBEnv k env)) def + argIn = ArgArray In (ArrayR shr tp) shIn (valueIn $ kOut .> kSegCount .> kSeg) + argOut = ArgArray Out (ArrayR shr tp) shOut (valueOut weakenId) + argSeg = ArgArray In (ArrayR dim1 $ TupRsingle itp) (TupRpair TupRunit $ TupRsingle $ Var (GroundRscalar scalarTypeInt) $ weaken kOut ZeroIdx) (TupRsingle $ Var (GroundRbuffer itp) $ weaken kOut $ SuccIdx ZeroIdx) + in + alet (LeftHandSidePair (mapLeftHandSide GroundRscalar lhsSh) lhsIn) (travA a) + $ alet lhsSeg (lowerOpenAcc (weakenBEnv (kIn .> kSh) env) segments) + -- The array of segment descriptors is one longer than the number + -- of segments. Subtract its length by one to compute the number + -- of segments (which is also the output size). + $ alet lhsSegCount + (Compute $ mkBinary (PrimSub numType) + (ArrayInstr (Parameter (Var scalarTypeInt $ SuccIdx ZeroIdx)) Nil) + (mkConstant (TupRsingle scalarTypeInt) 1) + ) + $ aletUnique lhsOut (lowerAlloc (ArrayR shr tp) shOut') + $ alet (LeftHandSideWildcard TupRunit) (mkFoldSeg i argF argDef argIn argSeg argOut) + $ Return (shOut `TupRpair` valueOut weakenId) + + -- scan1 + Named.Scan dir f Nothing a + | ArrayR shr tp <- Named.arrayR a + , DeclareVars lhsSh kSh valueSh <- declareVars $ shapeType shr + , DeclareVars lhsIn kIn valueIn <- declareVars $ buffersR tp + , DeclareVars lhsOut kOut valueOut <- declareVars $ buffersR tp -> + let + sh = mapVars GroundRscalar $ valueSh (kOut .> kIn) + k = kOut .> kIn .> kSh + argF = ArgFun $ lowerFun (weakenBEnv k env) f + argIn = ArgArray In (ArrayR shr tp) sh (valueIn kOut) + argOut = ArgArray Out (ArrayR shr tp) sh (valueOut weakenId) + in + alet (LeftHandSidePair (mapLeftHandSide GroundRscalar lhsSh) lhsIn) (travA a) + $ aletUnique lhsOut (lowerAlloc (ArrayR shr tp) $ valueSh kIn) + $ alet (LeftHandSideWildcard TupRunit) (mkScan dir argF Nothing argIn argOut) + $ Return (sh `TupRpair` valueOut weakenId) + + Named.Scan dir f (Just def) a + | ArrayR shr tp <- Named.arrayR a + , DeclareVars lhsShIn kShIn valueShIn <- declareVars $ shapeType shr + , DeclareVars lhsIn kIn valueIn <- declareVars $ buffersR tp + -- The rightmost dimension is incremented by 1. This variable only contains the rightmost dimension + , DeclareVars lhsShOut kShOut valueShOut <- declareVars $ TupRsingle $ scalarTypeInt + , DeclareVars lhsOut kOut valueOut <- declareVars $ buffersR tp -> + let + shIn = mapVars GroundRscalar $ valueShIn (kOut .> kShOut .> kIn) + shIn1 = case valueShIn kIn of + TupRpair _ (TupRsingle sh) -> sh + _ -> error "Impossible pair" + shOut = case shIn of + TupRpair sh _ -> TupRpair sh $ mapVars GroundRscalar $ valueShOut kOut + _ -> error "Impossible pair" + shOut' = case valueShIn (kShOut .> kIn) of + TupRpair sh _ -> TupRpair sh $ valueShOut weakenId + _ -> error "Impossible pair" + + k = kOut .> kShOut .> kIn .> kShIn + + argF = ArgFun $ lowerFun (weakenBEnv k env) f + argDef = Just $ ArgExp $ lowerExp (weakenBEnv k env) def + argIn = ArgArray In (ArrayR shr tp) shIn (valueIn $ kOut .> kShOut) + argOut = ArgArray Out (ArrayR shr tp) shOut (valueOut weakenId) + in + alet (LeftHandSidePair (mapLeftHandSide GroundRscalar lhsShIn) lhsIn) (travA a) + $ alet (mapLeftHandSide GroundRscalar lhsShOut) (Compute $ mkBinary (PrimAdd numType) (ArrayInstr (Parameter shIn1) Nil) (mkConstant (TupRsingle scalarTypeInt) 1)) + $ aletUnique lhsOut (lowerAlloc (ArrayR shr tp) shOut') + $ alet (LeftHandSideWildcard TupRunit) (mkScan dir argF argDef argIn argOut) + $ Return (shOut `TupRpair` valueOut weakenId) + + Named.Scan' dir f def a + | ArrayR (shapeInit -> shr) tp <- Named.arrayR a + , DeclareVars lhsSh kSh valueSh <- declareVars $ shapeType $ ShapeRsnoc shr + , DeclareVars lhsIn kIn valueIn <- declareVars $ buffersR tp + -- The prefix values + , DeclareVars lhsOut1 kOut1 valueOut1 <- declareVars $ buffersR tp + -- The final reduced value + , DeclareVars lhsOut2 kOut2 valueOut2 <- declareVars $ buffersR tp -> + let + sh = mapVars GroundRscalar $ valueSh (kOut2 .> kOut1 .> kIn) + + shOut1' = valueSh kIn + shOut2' = case valueSh (kOut1 .> kIn) of + TupRpair sh' _ -> sh' + _ -> error "Impossible pair" + shOut2 = case sh of + TupRpair sh' _ -> sh' + _ -> error "Impossible pair" + + k = kOut2 .> kOut1 .> kIn .> kSh + + argF = ArgFun $ lowerFun (weakenBEnv k env) f + argDef = ArgExp $ lowerExp (weakenBEnv k env) def + argIn = ArgArray In (ArrayR (ShapeRsnoc shr) tp) sh (valueIn $ kOut2 .> kOut1) + argOut1 = ArgArray Out (ArrayR (ShapeRsnoc shr) tp) sh (valueOut1 kOut2) + argOut2 = ArgArray Out (ArrayR shr tp) shOut2 (valueOut2 weakenId) + in + alet (LeftHandSidePair (mapLeftHandSide GroundRscalar lhsSh) lhsIn) (travA a) + $ aletUnique lhsOut1 (lowerAlloc (ArrayR (ShapeRsnoc shr) tp) $ shOut1') + $ aletUnique lhsOut2 (lowerAlloc (ArrayR shr tp) $ shOut2') + $ alet (LeftHandSideWildcard TupRunit) (mkScan' dir argF argDef argIn argOut1 argOut2) + $ Return (TupRpair sh (valueOut1 kOut2) `TupRpair` TupRpair shOut2 (valueOut2 weakenId)) + + Named.Stencil sr tp f boundary a + | shr <- stencilShapeR sr + , oldTp <- stencilEltR sr + , DeclareVars lhsSh kSh valueSh <- declareVars $ shapeType shr + , DeclareVars lhsIn kIn valueIn <- declareVars $ buffersR oldTp + , DeclareVars lhsOut kOut valueOut <- declareVars $ buffersR tp -> + let + lhs = LeftHandSidePair (mapLeftHandSide GroundRscalar lhsSh) lhsIn + sh = mapVars GroundRscalar $ valueSh (kOut .> kIn) + env' = weakenBEnv (kOut .> kIn .> kSh) env + argF = ArgFun $ lowerFun env' f + argIn = ArgArray In (ArrayR shr oldTp) sh (valueIn kOut) + argOut = ArgArray Out (ArrayR shr tp) sh (valueOut weakenId) + boundary' = lowerBoundary env' boundary + in + alet lhs (travA a) + $ aletUnique lhsOut (lowerAlloc (ArrayR shr tp) (valueSh kIn)) + $ alet (LeftHandSideWildcard TupRunit) (mkStencil sr argF boundary' argIn argOut) + $ Return (sh `TupRpair` valueOut weakenId) + + Named.Stencil2 sr1 sr2 tp3 f b1 a1 b2 a2 + | shr <- stencilShapeR sr1 + , tp1 <- stencilEltR sr1 + , tp2 <- stencilEltR sr2 + , DeclareVars lhsSh1 kSh1 valueSh1 <- declareVars $ shapeType shr + , DeclareVars lhsIn1 kIn1 valueIn1 <- declareVars $ buffersR tp1 + , DeclareVars lhsSh2 kSh2 valueSh2 <- declareVars $ shapeType shr + , DeclareVars lhsIn2 kIn2 valueIn2 <- declareVars $ buffersR tp2 + , DeclareVars lhsSh kSh valueSh <- declareVars $ shapeType shr + , DeclareVars lhsOut kOut valueOut <- declareVars $ buffersR tp3 -> + let + lhs1 = LeftHandSidePair (mapLeftHandSide GroundRscalar lhsSh1) lhsIn1 + lhs2 = LeftHandSidePair (mapLeftHandSide GroundRscalar lhsSh2) lhsIn2 + sh = mapVars GroundRscalar $ valueSh kOut + sh1 = mapVars GroundRscalar $ valueSh1 (kOut .> kSh .> kIn2 .> kSh2 .> kIn1) + sh2 = mapVars GroundRscalar $ valueSh2 (kOut .> kSh .> kIn2) + env' = weakenBEnv (kOut .> kSh .> kIn2 .> kSh2 .> kIn1 .> kSh1) env + argF = ArgFun $ lowerFun env' f + b1' = lowerBoundary env' b1 + b2' = lowerBoundary env' b2 + argIn1 = ArgArray In (ArrayR shr tp1) sh1 (valueIn1 $ kOut .> kSh .> kIn2 .> kSh2) + argIn2 = ArgArray In (ArrayR shr tp2) sh2 (valueIn2 $ kOut .> kSh) + argOut = ArgArray Out (ArrayR shr tp3) sh (valueOut weakenId) + in + alet lhs1 (travA a1) + $ alet lhs2 (lowerOpenAcc (weakenBEnv (kIn1 .> kSh1) env) a2) + $ alet (mapLeftHandSide GroundRscalar lhsSh) (Compute $ mkIntersect shr (valueSh1 $ kIn2 .> kSh2 .> kIn1) (valueSh2 kIn2)) + $ aletUnique lhsOut (lowerAlloc (ArrayR shr tp3) (valueSh weakenId)) + $ alet (LeftHandSideWildcard TupRunit) (mkStencil2 sr1 sr2 argF b1' argIn1 b2' argIn2 argOut) + $ Return (sh `TupRpair` valueOut weakenId) + Named.Atrace (Named.Message _ _ text) as bs + | repr <- Named.arraysR as + , DeclareVars lhs _k value <- declareVars repr + , LoweredLHS env' lhs' <- lowerLHS env lhs -> + alet lhs' (travA as) + $ alet (LeftHandSideSingle $ GroundRscalar $ scalarTypeWord8) + (Atrace text $ mapTupR (\var -> prj' (varIdx var) env') $ value weakenId) + $ Fence (IdxSet.singleton ZeroIdx) + $ lowerOpenAcc (weakenBEnv (weakenSucc' $ weakenWithLHS lhs') env) bs + +isUndef :: Named.OpenExp aenv env a -> Bool +isUndef (Let _ _ e) = isUndef e +isUndef (Undef _) = True +isUndef Nil = True +isUndef (Pair a b) = isUndef a && isUndef b +isUndef _ = False + +fixprimmaybepermute :: forall s sh a aenv. DeclareVars s (Buffers (PrimMaybe (((),sh),a))) aenv + -> PrimMaybeDeclareVars s sh a aenv +fixprimmaybepermute (DeclareVars lhs w k) = + DeclareVarsPrimMaybe lhs w $ \w' -> case k w' of + tag `TupRpair` (TupRunit `TupRpair` ((TupRunit `TupRpair` sh) `TupRpair` a)) -> tag `TupRpair` (TupRunit `TupRpair` (sh `TupRpair` a)) + _ -> error "impossible tupr" + +data PrimMaybeDeclareVars s sh a aenv where + DeclareVarsPrimMaybe :: LeftHandSide s (Buffers (PrimMaybe (((),sh),a))) env env' + -> (env :> env') + -> (forall env''. env' :> env'' -> Vars s env'' (Buffers (PrimMaybe (sh,a)))) + -> PrimMaybeDeclareVars s sh a env + +lowerAlloc :: forall benv op sh a. ArrayR (Array sh a) -> ExpVars benv sh -> OperationAcc op benv (Buffers a) +lowerAlloc (ArrayR _ TupRunit ) _ = Return TupRunit +lowerAlloc (ArrayR shr (TupRpair t1 t2)) sh = lowerAlloc (ArrayR shr t1) sh `pairUnique` lowerAlloc (ArrayR shr t2) sh +lowerAlloc (ArrayR shr (TupRsingle tp) ) sh + | Refl <- reprIsSingle @ScalarType @a @Buffer tp = Alloc shr tp sh + +lowerBoundary :: HasCallStack + => BEnv benv aenv + -> Named.Boundary aenv (Array sh e) + -> Boundary benv (Array sh e) +lowerBoundary _ Named.Clamp = Clamp +lowerBoundary _ Named.Mirror = Mirror +lowerBoundary _ Named.Wrap = Wrap +lowerBoundary _ (Named.Constant e) = Constant e +lowerBoundary env (Named.Function f) = Function $ lowerFun env f + +lowerOpenAfun :: (HasCallStack, LowerAcc op) + => BEnv benv aenv + -> Named.OpenAfun aenv a + -> OperationAfun op benv (LoweredAfun a) +lowerOpenAfun env (Named.Abody a) = case Named.arraysR a of + -- We must pattern match on the arrays representation of the body + -- to inform the type checker that 'a' is not a function, and thus + -- that 'LoweredAfun a' is equal to 'LoweredArrays a' + TupRsingle ArrayR{} -> Abody $ addShapeAssumes env $ lowerOpenAcc (weakenBEnv (weakenSucc weakenId) env) a + TupRunit -> Abody $ addShapeAssumes env $ lowerOpenAcc (weakenBEnv (weakenSucc weakenId) env) a + TupRpair _ _ -> Abody $ addShapeAssumes env $ lowerOpenAcc (weakenBEnv (weakenSucc weakenId) env) a +lowerOpenAfun env (Named.Alam lhs f) + | LoweredLHS env' lhs' <- lowerLHS env lhs = Alam lhs' $ lowerOpenAfun env' f + +data LoweredLHS b aenv' benv where + LoweredLHS :: BEnv benv' aenv' -> GLeftHandSide b benv benv' -> LoweredLHS b aenv' benv + +lowerLHS :: BEnv benv aenv -> Named.ALeftHandSide a aenv aenv' -> LoweredLHS (LoweredArrays a) aenv' benv +lowerLHS env (LeftHandSideWildcard repr) = LoweredLHS env $ LeftHandSideWildcard $ lowerArraysR repr +lowerLHS env (LeftHandSidePair l1 l2) + | LoweredLHS env1 l1' <- lowerLHS env l1 + , LoweredLHS env2 l2' <- lowerLHS env1 l2 = LoweredLHS env2 $ LeftHandSidePair l1' l2' +lowerLHS env (LeftHandSideSingle (ArrayR shr tp)) + | DeclareVars lhsSh kSh valueSh <- declareVars $ mapTupR GroundRscalar $ shapeType shr + , DeclareVars lhsBf kBf valueBf <- declareVars $ buffersR tp + = LoweredLHS (mapEnv (weakenArrayDescriptor $ kBf .> kSh) env `Push` ArrayDescriptor shr (valueSh kBf) (valueBf weakenId)) $ LeftHandSidePair lhsSh lhsBf + +addShapeAssumes :: BEnv benv aenv -> OperationAcc op (benv, Word8) a -> OperationAcc op benv a +addShapeAssumes env acc = + alet + (LeftHandSideSingle $ GroundRscalar $ scalarTypeWord8) + (Aassume $ shapeAssumes env) + $ Fence (IdxSet.singleton ZeroIdx) acc + +shapeAssumes :: BEnv benv aenv -> Exp benv PrimBool +shapeAssumes Empty = Const scalarTypeWord8 1 +shapeAssumes (env `Push` ArrayDescriptor shr sh _) = + mkBinary PrimLAnd (shapeAssumes env) (go shr sh) + where + go :: ShapeR sh -> GroundVars benv sh -> Exp benv PrimBool + go ShapeRz _ = Const scalarTypeWord8 1 + go (ShapeRsnoc shr') (TupRpair sh' (TupRsingle sz)) = + mkBinary PrimLAnd + (go shr' sh') + (mkBinary (PrimCmp singleType CmpGtEq) (paramIn scalarTypeInt sz) (Const scalarTypeInt 0)) + go _ _ = internalError "Shape impossible" + +lowerExp :: HasCallStack + => BEnv benv aenv + -> Named.OpenExp env aenv e + -> OpenExp env benv e +lowerExp env = rebuildArrayInstr (lowerArrayInstr env) + +lowerFun :: HasCallStack + => BEnv benv aenv + -> Named.Fun aenv e + -> Fun benv e +lowerFun env = rebuildArrayInstr (lowerArrayInstr env) + +lowerArrayInstr :: BEnv benv aenv -> Named.ArrayInstr aenv (s -> t) -> OpenExp env benv s -> OpenExp env benv t +lowerArrayInstr env (Named.Shape (Var (ArrayR shr _) array)) _ + | ArrayDescriptor _ sh _ <- prj' array env = paramsIn (shapeType shr) sh +lowerArrayInstr env (Named.LinearIndex (Var (ArrayR _ tp) array)) ix + = Let lhs ix $ linearIndex tp buffers $ Var int ZeroIdx + where + ArrayDescriptor _ _ buffers = prj' array env + int = SingleScalarType $ NumSingleType $ IntegralNumType TypeInt + lhs = LeftHandSideSingle int +lowerArrayInstr env (Named.Index (Var (ArrayR shr tp) array)) ix + = Let lhs (ToIndex shr sh' ix) $ linearIndex tp buffers $ Var int ZeroIdx + where + ArrayDescriptor _ sh buffers = prj' array env + sh' = paramsIn (shapeType shr) sh + int = SingleScalarType $ NumSingleType $ IntegralNumType TypeInt + lhs = LeftHandSideSingle int + +weakenBEnv :: forall benv benv' aenv. benv :> benv' -> BEnv benv aenv -> BEnv benv' aenv +weakenBEnv k = mapEnv (weakenArrayDescriptor k) + +lowerUnzip :: GroundVars benv (Buffers s) -> ELeftHandSide s () env -> ExpVars env t -> GroundVars benv (Buffers t) +lowerUnzip _ _ TupRunit = TupRunit +lowerUnzip buffers lhs (TupRpair v1 v2) = lowerUnzip buffers lhs v1 `TupRpair` lowerUnzip buffers lhs v2 +lowerUnzip buffers lhs (TupRsingle (Var _ idx)) = case lookupVar buffers lhs idx of + Right u -> u + Left (VoidIdx void) -> void + -- Left branch is unreachable, as `Idx () y` is an empty type + where + lookupVar :: GroundVars benv (Buffers x) -> ELeftHandSide x env1 env2 -> Idx env2 y -> Either (Idx env1 y) (GroundVars benv (Buffers y)) + lookupVar _ (LeftHandSideWildcard _) ix = Left ix + lookupVar buffers' (LeftHandSideSingle _) ix = case ix of + ZeroIdx -> Right buffers' + SuccIdx ix' -> Left ix' + lookupVar (TupRpair b1 b2) (LeftHandSidePair l1 l2) ix = case lookupVar b2 l2 ix of + Right u -> Right u + Left ix' -> lookupVar b1 l1 ix' + lookupVar _ _ _ = error "Impossible pair" + +-- Utilities +-- + +index :: ArrayR (Array sh e) -> GroundVars benv sh -> GroundVars benv (Buffers e) -> OpenExp env benv sh -> OpenExp env benv e +index (ArrayR shr tp) sh buffers ix = Let lhs (ToIndex shr (paramsIn (shapeType shr) sh) ix) $ linearIndex tp buffers $ Var int ZeroIdx + where + int = SingleScalarType $ NumSingleType $ IntegralNumType TypeInt + lhs = LeftHandSideSingle int + +indexFunction :: ArrayR (Array sh e) -> GroundVars benv sh -> GroundVars benv (Buffers e) -> OpenFun env benv (sh -> e) +indexFunction repr@(ArrayR shr _) sh buffers + | DeclareVars lhs _ value <- declareVars $ shapeType shr + = Lam lhs $ Body $ index repr sh buffers (expVars $ value weakenId) + +linearIndex :: forall e benv env. HasCallStack => TypeR e -> GroundVars benv (Buffers e) -> ExpVar env Int -> OpenExp env benv e +linearIndex TupRunit TupRunit _ = Nil +linearIndex (TupRpair t1 t2) (TupRpair b1 b2) ix = linearIndex t1 b1 ix `Pair` linearIndex t2 b2 ix +linearIndex (TupRsingle t) (TupRsingle var@(Var (GroundRbuffer _) _)) ix + | Refl <- reprIsSingle @ScalarType @e @Buffer t = ArrayInstr (Index var) (Evar ix) +linearIndex _ _ _ = internalError "Tuple mismatch" + +extend :: SliceIndex slix sl co sh + -> Fun env (sh -> sl) +extend sliceIx + | DeclareVars lhs _ value <- declareVars $ shapeType $ sliceDomainR sliceIx + = Lam lhs $ Body $ indexSlice sliceIx $ expVars $ value weakenId + +restrict :: SliceIndex slix sl co sh + -> Arg env (Var' slix) + -> Fun env (sl -> sh) +restrict sliceIx (ArgVar slix) + | DeclareVars lhs _ value <- declareVars $ shapeType $ sliceShapeR sliceIx + = Lam lhs $ Body $ indexFull sliceIx (paramsIn' slix) $ expVars $ value weakenId + +-- Utility function to reduce the dimension by one. Defined as a function, +-- as doing pattern matching in the guards in lowerOpenAcc will cause +-- a false warning "pattern matches are non-exhaustive". +shapeInit :: ShapeR (sh, Int) -> ShapeR sh +shapeInit (ShapeRsnoc shr) = shr + +lowerUse :: ArrayR (Array sh e) -> Array sh e -> OperationAcc op benv (LoweredArrays (Array sh e)) +lowerUse (ArrayR shr tp) (Array sh buffers) = Compute (mkConstant (shapeType shr) sh) `pair` lowerBuffers tp (size shr sh) buffers + +lowerBuffers :: forall e op benv. TypeR e -> Int -> Buffers e -> OperationAcc op benv (Buffers e) +lowerBuffers TupRunit _ () = Return TupRunit +lowerBuffers (TupRpair t1 t2) n (b1, b2) = lowerBuffers t1 n b1 `pair` lowerBuffers t2 n b2 +lowerBuffers (TupRsingle tp) n buffer + | Refl <- reprIsSingle @ScalarType @e @Buffer tp = Use tp n buffer + +lowerBoundaryToFunction :: forall benv sh e. Boundary benv (Array sh e) -> Arg benv (In sh e) -> Fun benv (sh -> e) +lowerBoundaryToFunction boundary (ArgArray _ repr@(ArrayR shr tp) sh buffers) = case boundary of + Function f + | DeclareVars lhs _ value <- declareVars $ shapeType shr + , value' <- value weakenId + -> Lam lhs $ Body $ Cond + (inbounds shr value' sh) + (index repr sh buffers $ expVars value') + (apply1 tp f $ expVars value') + Constant e + | DeclareVars lhs _ value <- declareVars $ shapeType shr + , value' <- value weakenId + -> Lam lhs $ Body $ Cond + (inbounds shr value' sh) + (index repr sh buffers $ expVars value') + (mkConstant tp e) + Clamp -> reindex clamp + Mirror -> reindex mirror + Wrap -> reindex wrap + where + inbounds :: ShapeR sh' -> ExpVars env sh' -> GroundVars benv sh' -> OpenExp env benv PrimBool + inbounds ShapeRz _ _ = Const scalarTypeWord8 1 + inbounds (ShapeRsnoc shr') (TupRpair ixs (TupRsingle ix)) (TupRpair szs (TupRsingle sz)) = + mkBinary PrimLAnd + (mkBinary PrimLAnd + (mkBinary (PrimCmp singleType CmpGtEq) (Evar ix) $ Const scalarTypeInt 0) + (mkBinary (PrimCmp singleType CmpLt) (Evar ix) $ paramIn scalarTypeInt sz) + ) + (inbounds shr' ixs szs) + inbounds _ _ _ = internalError "Illegal tuple for shape" + + -- Note: the expressions passed as argument may be duplicated in the resulting code. + -- This doesn't explode the code size as they are always an Evar or a Parameter. + -- + clamp, mirror, wrap :: OpenExp env benv Int -> OpenExp env benv Int -> OpenExp env benv Int + -- clamp ix sz = max(0, min(ix, sz - 1)) + clamp ix sz = mkBinary (PrimMax singleType) (Const scalarTypeInt 0) $ mkBinary (PrimMin singleType) ix $ sub sz $ Const scalarTypeInt 1 + + -- if ix < 0 then -ix + mirror ix sz = Cond (mkBinary (PrimCmp singleType CmpLt) ix (Const scalarTypeInt 0)) (PrimApp (PrimNeg numType) ix) + -- else if ix >= sz then sz - ((ix - sz) + 2) + $ Cond (mkBinary (PrimCmp singleType CmpGtEq) ix sz) (sub sz (add (sub ix sz) $ Const scalarTypeInt 2)) + -- else ix + $ ix + + -- if ix < 0 then sz + ix + wrap ix sz = Cond (mkBinary (PrimCmp singleType CmpLt) ix (Const scalarTypeInt 0)) (add sz ix) + -- else if ix >= sz then ix - sz + $ Cond (mkBinary (PrimCmp singleType CmpGtEq) ix sz) (sub ix sz) + -- else ix + $ ix + + add = mkBinary (PrimAdd numType) + sub = mkBinary (PrimSub numType) + + reindex :: (forall env'. OpenExp env' benv Int -> OpenExp env' benv Int -> OpenExp env' benv Int) -> Fun benv (sh -> e) + reindex f + | DeclareVars lhs _ value <- declareVars $ shapeType shr + = Lam lhs $ Body $ index repr sh buffers $ getIndex f shr (value weakenId) sh + + getIndex :: (OpenExp env benv Int -> OpenExp env benv Int -> OpenExp env benv Int) -> ShapeR sh' -> ExpVars env sh' -> GroundVars benv sh' -> OpenExp env benv sh' + getIndex _ ShapeRz _ _ = Nil + getIndex f (ShapeRsnoc shr') (TupRpair ixs (TupRsingle ix)) (TupRpair szs (TupRsingle sz)) = getIndex f shr' ixs szs `Pair` f (Evar ix) (paramIn scalarTypeInt sz) + getIndex _ _ _ _ = internalError "Illegal tuple for shape" + +lowerStencilFun :: StencilR sh e stencil -> TypeR e' -> Fun benv (stencil -> e') -> Maybe (Boundary benv (Array sh e)) -> Arg benv (In sh e) -> Fun benv (sh -> e') +lowerStencilFun sr _ f boundary array = f `compose` stencilAccess sr boundary array + +lowerStencil2Fun :: StencilR sh a stencil1 -> StencilR sh b stencil2 -> TypeR c -> Fun benv (stencil1 -> stencil2 -> c) -> Maybe (Boundary benv (Array sh a)) -> Arg benv (In sh a) -> Maybe (Boundary benv (Array sh b)) -> Arg benv (In sh b) -> Fun benv (sh -> c) +lowerStencil2Fun sr1 sr2 tp f boundary1 array1@(ArgArray _ (ArrayR shr _) _ _) boundary2 array2 + | DeclareVars lhs _ value <- declareVars $ shapeType shr + , vars <- expVars $ value weakenId + = Lam lhs $ Body $ apply2 tp f (apply1 (stencilR sr1) (stencilAccess sr1 boundary1 array1) vars) + (apply1 (stencilR sr2) (stencilAccess sr2 boundary2 array2) vars) + +stencilAccess :: forall sh e stencil benv. StencilR sh e stencil -> Maybe (Boundary benv (Array sh e)) -> Arg benv (In sh e) -> Fun benv (sh -> stencil) +stencilAccess sr boundary array@(ArgArray _ repr@(ArrayR shr tp) sh buffers) + | DeclareVars lhs _ value <- declareVars $ shapeType shr = Lam lhs $ Body $ stencilAccess' shr sr readFunction (value weakenId) + where + readFunction :: env :> env' -> ExpVars env' sh -> OpenExp env' benv e + readFunction _ = readFunction' . expVars + + readFunction' :: OpenExp env' benv sh -> OpenExp env' benv e + readFunction' = apply1 tp $ case boundary of + Just b -> lowerBoundaryToFunction b array + Nothing -> indexFunction repr sh buffers + +stencilAccess' :: forall env benv sh e stencil. ShapeR sh -> StencilR sh e stencil -> (forall env'. env :> env' -> ExpVars env' sh -> OpenExp env' benv e) -> ExpVars env sh -> OpenExp env benv stencil +stencilAccess' shr sr rf ix = case sr of + -- Base cases, nothing interesting to do here since we know the lower + -- dimension is Z. + -- + StencilRunit3 _ -> Nil `Pair` rf' (-1) + `Pair` center + `Pair` rf' 1 + StencilRunit5 _ -> Nil `Pair` rf' (-2) + `Pair` rf' (-1) + `Pair` center + `Pair` rf' 1 + `Pair` rf' 2 + StencilRunit7 _ -> Nil `Pair` rf' (-3) + `Pair` rf' (-2) + `Pair` rf' (-1) + `Pair` center + `Pair` rf' 1 + `Pair` rf' 2 + `Pair` rf' 3 + StencilRunit9 _ -> Nil `Pair` rf' (-4) + `Pair` rf' (-3) + `Pair` rf' (-2) + `Pair` rf' (-1) + `Pair` center + `Pair` rf' 1 + `Pair` rf' 2 + `Pair` rf' 3 + `Pair` rf' 4 + + -- Recursive cases. Note that because the stencil pattern is defined with + -- a cons ordering, whereas shapes (indices) are defined as a snoc list, + -- when we recurse on the stencil structure we must manipulate the + -- _left-most_ index component + -- + StencilRtup3 s1 s2 s3 -> + Nil `Pair` go s1 (-1) + `Pair` go s2 0 + `Pair` go s3 1 + + StencilRtup5 s1 s2 s3 s4 s5 -> + Nil `Pair` go s1 (-2) + `Pair` go s2 (-1) + `Pair` go s3 0 + `Pair` go s4 1 + `Pair` go s5 2 + + StencilRtup7 s1 s2 s3 s4 s5 s6 s7 -> + Nil `Pair` go s1 (-3) + `Pair` go s2 (-2) + `Pair` go s3 (-1) + `Pair` go s4 0 + `Pair` go s5 1 + `Pair` go s6 2 + `Pair` go s7 3 + + StencilRtup9 s1 s2 s3 s4 s5 s6 s7 s8 s9 -> + Nil `Pair` go s1 (-4) + `Pair` go s2 (-3) + `Pair` go s3 (-2) + `Pair` go s4 (-1) + `Pair` go s5 0 + `Pair` go s6 1 + `Pair` go s7 2 + `Pair` go s8 3 + `Pair` go s9 4 + where + center :: sh ~ DIM1 => OpenExp env benv e + center = rf weakenId ix + + -- read function for 1 dimensional cases + rf' :: sh ~ DIM1 => Int -> OpenExp env benv e + rf' d + | TupRpair _ (TupRsingle i) <- ix = Let (LeftHandSideSingle scalarTypeInt) (mkBinary (PrimAdd numType) (Evar i) $ Const scalarTypeInt d) + $ rf (weakenSucc weakenId) $ TupRpair TupRunit $ TupRsingle $ Var scalarTypeInt ZeroIdx + rf' _ = internalError "Shape tuple mismatch" + + go :: forall sh' stencil'. sh ~ (sh', Int) => StencilR sh' e stencil' -> Int -> OpenExp env benv stencil' + go sr' d = stencilAccess' shr' sr' rf'' ix' + where + ShapeRsnoc shr' = shr + (i, ix') = uncons shr' ix + + rf'' :: (env :> env') -> ExpVars env' sh' -> OpenExp env' benv e + rf'' k ds + | d == 0 = rf k (cons shr' (weakenE k i) ds) + | otherwise = Let (LeftHandSideSingle scalarTypeInt) (mkBinary (PrimAdd numType) (Evar $ weakenE k i) $ Const scalarTypeInt d) + $ rf (weakenSucc' k) $ cons shr' (Var scalarTypeInt ZeroIdx) $ weakenVars (weakenSucc weakenId) ds + +-- Add a _left-most_ dimension to a shape +-- +cons :: ShapeR sh -> ExpVar env Int -> ExpVars env sh -> ExpVars env (sh, Int) +cons ShapeRz ix _ = TupRpair TupRunit $ TupRsingle ix +cons (ShapeRsnoc shr) ix (TupRpair sh sz) = TupRpair (cons shr ix sh) sz +cons _ _ _ = internalError "Illegal shape tuple" + + +-- Remove the _left-most_ index to a shape, and return the remainder +-- +uncons :: ShapeR sh -> ExpVars env (sh, Int) -> (ExpVar env Int, ExpVars env sh) +uncons ShapeRz (TupRpair _ (TupRsingle v2)) = (v2, TupRunit) +uncons (ShapeRsnoc shr) (TupRpair v1 v3) + = let (i, v1') = uncons shr v1 + in (i, TupRpair v1' v3) +uncons _ _ = internalError "Illegal shape tuple" + +mkIntersect :: ShapeR sh -> ExpVars benv sh -> ExpVars benv sh -> PreOpenExp (ArrayInstr benv) env sh +mkIntersect shr' x y + | Just Refl <- matchVars x y = paramsIn' x + | otherwise = mkIntersect' shr' x y + where + mkIntersect' :: ShapeR sh -> ExpVars benv sh -> ExpVars benv sh -> PreOpenExp (ArrayInstr benv) env sh + mkIntersect' ShapeRz _ _ = Nil + mkIntersect' (ShapeRsnoc shr) (TupRpair s1 x1) (TupRpair s2 x2) = mkIntersect' shr s1 s2 `Pair` mkBinary (PrimMin singleType) (paramsIn' x1) (paramsIn' x2) + mkIntersect' (ShapeRsnoc _) _ _ = error "Impossible pair" + +mkDefaultFoldSequential :: forall benv op sh e. LowerAcc op => Arg benv (Fun' (e -> e -> e)) -> Maybe (Arg benv (Exp' e)) -> Arg benv (In (sh, Int) e) -> Arg benv (Out sh e) -> OperationAcc op benv () +mkDefaultFoldSequential op def argIn = mkGenerate (mkDefaultFoldFunction op def argIn) + +mkDefaultFoldFunction :: Arg benv (Fun' (e -> e -> e)) -> Maybe (Arg benv (Exp' e)) -> Arg benv (In (sh, Int) e) -> Arg benv (Fun' (sh -> e)) +mkDefaultFoldFunction (ArgFun op) def (ArgArray _ (ArrayR (ShapeRsnoc shr) tp) (sh `TupRpair` n) buffers) + | DeclareVars lhsIdx _k1 valueIdx <- declareVars $ shapeType shr + , DeclareVars lhsVal k2 valueVal <- declareVars tp = + let + initial = case def of + Nothing -> + Pair + (mkConstant (TupRsingle scalarTypeInt) 1) + (index (ArrayR (ShapeRsnoc shr) tp) (sh `TupRpair` n) buffers (expVars (valueIdx weakenId) `Pair` mkConstant (TupRsingle scalarTypeInt) 0)) + Just (ArgExp d) -> + Pair + (mkConstant (TupRsingle scalarTypeInt) 0) + (weakenE weakenEmpty d) + + lhs = LeftHandSidePair (LeftHandSideSingle scalarTypeInt) lhsVal + -- \(idx, accum) + step = Lam lhs $ Body $ Pair + -- (idx + 1 + (mkBinary (PrimAdd numType) (Evar $ Var scalarTypeInt $ k2 >:> ZeroIdx) (mkConstant (TupRsingle scalarTypeInt) 1)) + -- , op accum (input !! idx) + $ apply2 tp op (expVars $ valueVal weakenId) + $ index (ArrayR (ShapeRsnoc shr) tp) (sh `TupRpair` n) buffers + $ expVars (valueIdx $ weakenSucc k2) `Pair` Evar (Var scalarTypeInt $ k2 >:> ZeroIdx) + + condition = + Lam (LeftHandSidePair (LeftHandSideSingle scalarTypeInt) (LeftHandSideWildcard tp)) + $ Body $ mkBinary (PrimCmp singleType CmpLt) (Evar $ Var scalarTypeInt ZeroIdx) (paramsIn (TupRsingle scalarType) n) + in + ArgFun $ Lam lhsIdx $ Body + $ Let lhs (While condition step initial) + $ expVars $ valueVal weakenId +mkDefaultFoldFunction _ _ (ArgArray _ _ (TupRsingle sh) _) = pairImpossible sh + +-- In case of a scan with a default value, prepends the initial value before the other elements +-- The default value is placed as the first value in case of a left-to-right scan, or as the +-- last value for a right-to-left scan. +mkDefaultScanPrepend :: Direction -> Arg benv (Exp' e) -> Arg benv (In (sh, Int) e) -> Fun benv ((sh, Int) -> e) +mkDefaultScanPrepend dir (ArgExp def) (ArgArray _ repr@(ArrayR (ShapeRsnoc shr) _) sh input) + | DeclareVars lhs k value <- declareVars $ shapeType shr + = let + first = case dir of + LeftToRight -> mkConstant (TupRsingle scalarTypeInt) 0 + RightToLeft -> case sh of + TupRpair _ n -> paramsIn (TupRsingle scalarTypeInt) n + _ -> error "Impossible pair" + x = case dir of + LeftToRight -> mkBinary (PrimSub numType) (Evar $ Var scalarTypeInt ZeroIdx) (mkConstant (TupRsingle scalarType) 1) + RightToLeft -> Evar $ Var scalarTypeInt ZeroIdx + in + Lam (lhs `LeftHandSidePair` LeftHandSideSingle scalarTypeInt) + $ Body + $ Cond (mkBinary (PrimCmp singleType CmpEq) (Evar $ Var scalarTypeInt ZeroIdx) first) (weakenE (weakenSucc' k) def) + $ index repr sh input + $ Pair (expVars $ value $ weakenSucc weakenId) x + +mkDefaultScanFunction :: Direction -> GroundVar benv Int -> Arg benv (Fun' (e -> e -> e)) -> Arg benv (In (sh, Int) e) -> Fun benv ((sh, Int) -> e) +mkDefaultScanFunction dir inc (ArgFun f) (ArgArray _ repr@(ArrayR (ShapeRsnoc shr) tp) sh input) + | DeclareVars lhs _ value <- declareVars $ shapeType shr + = let + op = case dir of + LeftToRight -> PrimSub + RightToLeft -> PrimAdd + x = Evar $ Var scalarTypeInt ZeroIdx + y = mkBinary (op numType) x (paramIn scalarTypeInt inc) + condition = case dir of + LeftToRight -> mkBinary (PrimCmp singleType CmpGtEq) y (mkConstant (TupRsingle scalarTypeInt) 0) + RightToLeft -> case sh of + TupRpair _ n -> mkBinary (PrimCmp singleType CmpLt) y (paramsIn (TupRsingle scalarTypeInt) n) + _ -> error "Impossible pair" + ix = expVars $ value $ weakenSucc weakenId + index' = index repr sh input . Pair ix + in + Lam (lhs `LeftHandSidePair` LeftHandSideSingle scalarTypeInt) + $ Body + $ Cond condition + ( + case dir of + LeftToRight -> apply2 tp f (index' y) (index' x) + RightToLeft -> apply2 tp f (index' x) (index' y) + ) + (index' x) + +mkDefaultFoldSegFunction + :: IntegralType i + -> Arg benv (Fun' (e -> e -> e)) + -> Maybe (Arg benv (Exp' e)) + -> Arg benv (In (sh, Int) e) + -> Arg benv (In DIM1 i) + -> Fun benv ((sh, Int) -> e) +mkDefaultFoldSegFunction itp (ArgFun f) def (ArgArray _ (ArrayR shr tp) sh input) (ArgArray _ reprSeg shSeg segments) + | DeclareVars lhsSh kSh valueSh <- declareVars $ shapeType shr + , DeclareVars lhsE kE valueE <- declareVars tp = + let + x1 = case valueSh weakenId of + TupRpair _ x' -> expVars x' + _ -> error "Impossible pair" + x2 = case valueSh (weakenSucc' kE) of + TupRpair _ x' -> expVars x' + _ -> error "Impossible pair" + start = PrimApp (PrimFromIntegral itp numType) $ index reprSeg shSeg segments $ Pair Nil x1 + end = PrimApp (PrimFromIntegral itp numType) $ index reprSeg shSeg segments $ Pair Nil $ mkBinary (PrimAdd numType) x2 (mkConstant (TupRsingle scalarTypeInt) 1) + index' = index (ArrayR shr tp) sh input + cond = Lam (lhsE `LeftHandSidePair` LeftHandSideSingle scalarTypeInt) $ Body $ mkBinary (PrimCmp singleType CmpLt) (Evar $ Var scalarTypeInt ZeroIdx) end + step = Lam (lhsE `LeftHandSidePair` LeftHandSideSingle scalarTypeInt) $ Body $ Pair + (apply2 tp f (expVars $ valueE $ weakenSucc weakenId) (index' $ Pair (expVars $ fst' $ valueSh (weakenSucc' kE)) $ Evar $ Var scalarTypeInt ZeroIdx)) + (mkBinary (PrimAdd numType) (Evar $ Var scalarTypeInt ZeroIdx) (mkConstant (TupRsingle scalarTypeInt) 1)) + initial = case def of + Just (ArgExp d) -> Pair (weakenE kSh d) start + Nothing -> Pair (index (ArrayR shr tp) sh input $ Pair (expVars $ fst' $ valueSh weakenId) start) (mkBinary (PrimAdd numType) start $ mkConstant (TupRsingle scalarTypeInt) 1) + + fst' :: ExpVars env (a, b) -> ExpVars env a + fst' (TupRpair a _) = a + fst' _ = error "Impossible pair" + in + Lam lhsSh + $ Body + $ Let (lhsE `LeftHandSidePair` LeftHandSideWildcard (TupRsingle scalarTypeInt)) (While cond step initial) + $ expVars $ valueE weakenId + +-- indexSlice and indexFull, as used by Slice and Replicate. +-- Note that these two functions used to be constructors in PreOpenExp. +-- Since they just restructure two tuples, we removed these constructors. +-- We can better track and analyse the program by just using the tuple +-- constructors (Pair and Nil). +indexSlice + :: SliceIndex slix sl co sh + -- Original IndexSlice constructor took another argument, but that was not used: + -- -> PreOpenExp arr env slix + -> PreOpenExp arr env sh + -> PreOpenExp arr env sl +indexSlice SliceNil _ = Nil +indexSlice (SliceAll sliceIdx) (Pair sl sz) = + Pair (indexSlice sliceIdx sl) sz +indexSlice (SliceFixed sliceIdx) (Pair sl _) = + indexSlice sliceIdx sl +-- If 'expr' is not a Pair, bind it in a let and then recurse +indexSlice sliceIdx expr + | tp <- shapeType $ sliceDomainR sliceIdx + , DeclareVars lhs _ value <- declareVars tp + = Let lhs expr + $ indexSlice sliceIdx $ expVars $ value weakenId + +indexFull + :: SliceIndex slix sl co t + -> PreOpenExp arr env slix + -> PreOpenExp arr env sl + -> PreOpenExp arr env t +indexFull SliceNil _ _ = Nil +indexFull (SliceAll sliceIdx) (Pair slx _) (Pair sh sz) = + Pair (indexFull sliceIdx slx sh) sz +indexFull (SliceFixed sliceIdx) (Pair slx sz) sh = + Pair (indexFull sliceIdx slx sh) sz +-- If 'sh' is not a Pair, bind it in a let and then recurse +indexFull sliceIdx@SliceAll{} slx@Pair{} sh + | tp <- shapeType $ sliceShapeR sliceIdx + , DeclareVars lhs k value <- declareVars tp + = Let lhs sh + $ indexFull sliceIdx (weakenE k slx) (expVars $ value weakenId) +-- If 'slx' is not a Pair, bind it in a let and then recurse +indexFull sliceIdx slx sh + | tp <- sliceEltR sliceIdx + , DeclareVars lhs k value <- declareVars tp + = Let lhs slx + $ indexFull sliceIdx (expVars $ value weakenId) (weakenE k sh) + +uncurry' :: Fun env (a -> b -> c) -> Fun env ((a, b) -> c) +uncurry' (Lam lhs1 (Lam lhs2 f)) = Lam (LeftHandSidePair lhs1 lhs2) f +uncurry' (Lam _ (Body _)) = error "impossible: Expression of function type" +uncurry' (Body _) = error "impossible: Expression of function type" + +-- Replaces scalar arrays with expression variables in the state of an awhile loop. + +data RemoveScalars before after where + RemoveScalarsKeep :: RemoveScalars t t + RemoveScalarsPair + :: RemoveScalars before1 after1 + -> RemoveScalars before2 after2 + -> RemoveScalars (before1, before2) (after1, after2) + RemoveScalarsScalar :: ScalarType t -> RemoveScalars (Buffer t) t + +removeScalarsForType :: ArraysR t -> Exists (RemoveScalars (LoweredArrays t)) +removeScalarsForType TupRunit = Exists RemoveScalarsKeep +removeScalarsForType (TupRsingle (ArrayR shr tp)) = case shr of + ShapeRz -> Exists $ RemoveScalarsPair RemoveScalarsKeep $ go tp + _ -> Exists RemoveScalarsKeep + where + go :: forall e. TypeR e -> RemoveScalars (Buffers e) e + go (TupRpair t1 t2) = RemoveScalarsPair (go t1) (go t2) + go TupRunit = RemoveScalarsKeep + go (TupRsingle t) + | Refl <- reprIsSingle @ScalarType @e @Buffer t = RemoveScalarsScalar t +removeScalarsForType (TupRpair l r) + | Exists RemoveScalarsKeep <- l' + , Exists RemoveScalarsKeep <- r' + = Exists RemoveScalarsKeep + | Exists l'' <- l' + , Exists r'' <- r' + = Exists $ RemoveScalarsPair l'' r'' + where + l' = removeScalarsForType l + r' = removeScalarsForType r + +data RemoveScalarsInTerm op env t where + RemoveScalarsInTerm + :: (forall r. OperationAcc op env' r -> OperationAcc op env r) + -> env :> env' + -> (forall env''. env' :> env'' -> GroundVars env'' t) + -> RemoveScalarsInTerm op env t + +awhileRemoveScalars + :: forall tp op env. + LowerAcc op + => ArraysR tp + -> Uniquenesses (LoweredArrays tp) + -> OperationAfun op env (LoweredArrays tp -> PrimBool) + -> OperationAfun op env (LoweredArrays tp -> LoweredArrays tp) + -> GroundVars env (LoweredArrays tp) + -> OperationAcc op env (LoweredArrays tp) +awhileRemoveScalars repr us cond step initial + | not $ lowerPreferNoScalar @op + = Awhile us cond step initial + | Exists RemoveScalarsKeep <- remove' + = Awhile us cond step initial + | Exists remove <- remove' + -- Convert initial value + , RemoveScalarsInTerm b1 k1 res1 <- removeScalarsInTerm' remove weakenId initial + , us' <- uniquenesses remove us + , initial' <- res1 weakenId + = removeScalarsInTermInv (varsType initial') us' remove + $ b1 + $ Awhile + us' + (weaken k1 $ removeBoundScalars remove us cond) + (weaken k1 $ removeBoundScalars remove us $ removeScalarsInFun tp us remove step) + initial' + where + remove' = removeScalarsForType repr + tp = varsType initial + + uniquenesses :: RemoveScalars b a -> Uniquenesses b -> Uniquenesses a + uniquenesses RemoveScalarsKeep u = u + uniquenesses (RemoveScalarsPair r1 r2) (TupRpair u1 u2) = uniquenesses r1 u1 `TupRpair` uniquenesses r2 u2 + uniquenesses (RemoveScalarsScalar _) _ = TupRsingle Shared + uniquenesses _ _ = internalError "Tuple mismatch" + +removeBoundScalars + :: RemoveScalars before after + -> Uniquenesses before + -> OperationAfun op env (before -> t) + -> OperationAfun op env (after -> t) +removeBoundScalars remove us (Alam lhs (Abody body)) + | DeclareVars lhsNew k0 value <- declareVars $ tp remove $ lhsToTupR lhs + , RemoveScalarsInTerm b k1 res <- removeScalarsInTermInv' remove weakenId (value weakenId) + , Exists lhs' <- rebuildLHS lhs + = Alam lhsNew + $ Abody + $ b + $ alet' lhs' us (Return $ res weakenId) + $ weaken (sinkWithLHS lhs lhs' (k1 .> k0)) body + where + tp :: RemoveScalars b a -> GroundsR b -> GroundsR a + tp RemoveScalarsKeep t = t + tp (RemoveScalarsPair r1 r2) (TupRpair t1 t2) = tp r1 t1 `TupRpair` tp r2 t2 + tp (RemoveScalarsScalar t) _ = TupRsingle $ GroundRscalar t + tp _ _ = internalError "Tuple mismatch" +removeBoundScalars _ _ _ = internalError "Expected unary function" + +removeScalarsInFun + :: GroundsR before + -> Uniquenesses before + -> RemoveScalars before after + -> OperationAfun op env (t -> before) + -> OperationAfun op env (t -> after) +removeScalarsInFun tp us r (Alam lhs (Abody body)) + = Alam lhs $ Abody $ removeScalarsInTerm tp us r body +removeScalarsInFun _ _ _ _ = internalError "Expected unary function" + +removeScalarsInTerm + :: GroundsR before + -> Uniquenesses before + -> RemoveScalars before after + -> OperationAcc op env before + -> OperationAcc op env after +removeScalarsInTerm tp us r term + | DeclareVars lhs _ value <- declareVars tp + , RemoveScalarsInTerm bind _ result <- removeScalarsInTerm' r weakenId (value weakenId) = + alet' lhs us term $ bind $ Return $ result weakenId + +removeScalarsInTerm' + :: RemoveScalars before after + -> env0 :> env + -> GroundVars env0 before + -> RemoveScalarsInTerm op env after +removeScalarsInTerm' RemoveScalarsKeep k0 vars = + RemoveScalarsInTerm id weakenId $ \k -> mapTupR (weaken (k .> k0)) vars +removeScalarsInTerm' (RemoveScalarsPair r1 r2) k0 (TupRpair v1 v2) + | RemoveScalarsInTerm b1 k1 res1 <- removeScalarsInTerm' r1 k0 v1 + , RemoveScalarsInTerm b2 k2 res2 <- removeScalarsInTerm' r2 (k1 .> k0) v2 + = RemoveScalarsInTerm + (b1 . b2) + (k2 .> k1) + $ \k -> res1 (k .> k2) `TupRpair` res2 k +removeScalarsInTerm' (RemoveScalarsScalar tp) k0 (TupRsingle buffer) = + RemoveScalarsInTerm + (alet' + (LeftHandSideSingle $ GroundRscalar tp) + (TupRsingle Shared) + (Compute $ ArrayInstr (Index $ weaken k0 buffer) $ Const scalarTypeInt 0)) + (weakenSucc weakenId) + $ \k -> TupRsingle $ Var (GroundRscalar tp) $ weaken k ZeroIdx +removeScalarsInTerm' _ _ _ = internalError "Tuple mismatch" + +removeScalarsInTermInv + :: GroundsR after + -> Uniquenesses after + -> RemoveScalars before after + -> OperationAcc op env after + -> OperationAcc op env before +removeScalarsInTermInv tp us r term + | DeclareVars lhs _ value <- declareVars tp + , RemoveScalarsInTerm bind _ result <- removeScalarsInTermInv' r weakenId (value weakenId) = + alet' lhs us term $ bind $ Return $ result weakenId + +removeScalarsInTermInv' + :: RemoveScalars before after + -> env0 :> env + -> GroundVars env0 after + -> RemoveScalarsInTerm op env before +removeScalarsInTermInv' RemoveScalarsKeep k0 vars = + RemoveScalarsInTerm id weakenId $ \k -> mapTupR (weaken (k .> k0)) vars +removeScalarsInTermInv' (RemoveScalarsPair r1 r2) k0 (TupRpair v1 v2) + | RemoveScalarsInTerm b1 k1 res1 <- removeScalarsInTermInv' r1 k0 v1 + , RemoveScalarsInTerm b2 k2 res2 <- removeScalarsInTermInv' r2 (k1 .> k0) v2 + = RemoveScalarsInTerm + (b1 . b2) + (k2 .> k1) + $ \k -> res1 (k .> k2) `TupRpair` res2 k +removeScalarsInTermInv' (RemoveScalarsScalar tp) k0 (TupRsingle var) = + RemoveScalarsInTerm + (alet' + (LeftHandSideSingle $ GroundRbuffer tp) + (TupRsingle Shared) + (Unit var')) + (weakenSucc weakenId) + $ \k -> TupRsingle $ Var (GroundRbuffer tp) $ weaken k ZeroIdx + where + var' = Var tp $ weaken k0 $ varIdx var +removeScalarsInTermInv' _ _ _ = internalError "Tuple mismatch" diff --git a/src/Data/Array/Accelerate/Trafo/NewNewFusion.hs b/src/Data/Array/Accelerate/Trafo/NewNewFusion.hs new file mode 100644 index 000000000..84af50d8d --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/NewNewFusion.hs @@ -0,0 +1,61 @@ +-- | +-- Module : Data.Array.Accelerate.Trafo.NewNewFusion +-- Copyright : [2012..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- +-- This module should implement fusion. +-- + +module Data.Array.Accelerate.Trafo.NewNewFusion ( + + convertAfun, convertAfunWith, + + Benchmarking(..) + +) where + +import Data.Array.Accelerate.AST.Operation +import Data.Array.Accelerate.AST.Partitioned +import Data.Array.Accelerate.Trafo.Config +import Data.Array.Accelerate.Error +import Data.Array.Accelerate.Trafo.Operation.Simplify +import Data.Array.Accelerate.Trafo.Partitioning.ILP +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Graph (MakesILP) +import qualified Data.Array.Accelerate.Pretty.Operation as Pretty +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Encoding (Objective (..)) + + + + + + +-- Array Fusion +-- ============ + +defaultSolver :: Solver +defaultSolver = + MIPSolver CBC + +-- | Apply the fusion transformation to a function of array arguments +-- +convertAfun :: (HasCallStack, MakesILP op, SimplifyOperation op, Pretty.PrettyOp (Cluster op)) => FusionType -> OperationAfun op () f -> PartitionedAfun op () f +convertAfun = convertAfunWith defaultOptions + +convertAfunWith :: (HasCallStack, MakesILP op, SimplifyOperation op, Pretty.PrettyOp (Cluster op)) => Config -> FusionType -> OperationAfun op () f -> PartitionedAfun op () f +convertAfunWith _ (Fusion o) = withSimplStats (ilpFusionF'' defaultSolver o) +convertAfunWith _ (Benchmarking b) = withSimplStats (benchF b FusedEdges) + + +withSimplStats :: a -> a +-- #ifdef ACCELERATE_DEBUG +-- withSimplStats x = unsafePerformIO Stats.resetSimplCount `seq` x +-- #else +withSimplStats x = x +-- #endif + +-- dontFuse :: op args -> Args env args -> Cluster' op args +-- dontFuse = unfused diff --git a/src/Data/Array/Accelerate/Trafo/Operation/Bounds.hs b/src/Data/Array/Accelerate/Trafo/Operation/Bounds.hs new file mode 100644 index 000000000..55a6c63a6 --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Operation/Bounds.hs @@ -0,0 +1,737 @@ +{-# OPTIONS_HADDOCK hide #-} +-- | +-- Module : Data.Array.Accelerate.Trafo.Operation.Bounds +-- Copyright : [2012..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.Trafo.Operation.Bounds ( + boundsOptimizeAfun, + OperationBounds(..), + -- Default implementations for OperationBounds + boundsOptimizeOpDefault, boundsOptimizeGenerate, + boundsOptimizeBackpermute, boundsOptimizeMap, + boundsOptimizeFold, boundsOptimizeFold1, + boundsOptimizeScan, boundsOptimizeScan1, boundsOptimizeScan', + -- Utilities for implementing OperationBounds + InputBounds(..), OutputBounds(..), InputBoundArgs, OutputBoundArgs, + boundsOptimizeFun, boundsOptimizeFun1, boundsOptimizeFun2, boundsOptimizeExp, +) where + +import Data.Array.Accelerate.AST.Exp +import Data.Array.Accelerate.AST.Environment +import qualified Data.Array.Accelerate.AST.Graph as Graph +import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.IdxSet (IdxSet) +import qualified Data.Array.Accelerate.AST.IdxSet as IdxSet +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.AST.Operation +import Data.Array.Accelerate.Trafo.Exp.Shrink +import Data.Array.Accelerate.Trafo.Exp.Simplify +import Data.Array.Accelerate.Trafo.Substitution +import Data.Array.Accelerate.Trafo.WeakenedEnvironment +import Data.Array.Accelerate.Representation.Array +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Type +import Data.Array.Accelerate.Error +import Data.Array.Accelerate.Analysis.Match + +import Data.Array.Accelerate.Trafo.Operation.Bounds.Algebra +import Data.Array.Accelerate.Trafo.Operation.Bounds.Environment + +import Data.Maybe (mapMaybe) +import Data.List (foldl') + +boundsOptimizeAfun + :: forall op f. + OperationBounds op + => OperationAfun op () f + -> OperationAfun op () f +boundsOptimizeAfun = snd . boundsOptimizeOpenAfun emptyEnv + +boundsOptimizeOpenAfun + :: forall benv op f. + OperationBounds op + => BoundsEnv benv () + -> OperationAfun op benv f + -> (IdxSet benv, OperationAfun op benv f) +boundsOptimizeOpenAfun env@(BoundsEnv _ _ zero _) = \case + Alam lhs f + | env1 <- env `pushBuffers` (lhs, bottomsGround zero $ lhsToTupR lhs) + , (modified, f') <- boundsOptimizeOpenAfun env1 f -> + (IdxSet.drop' lhs modified, Alam lhs f') + Abody acc + | (modified, _, _, acc') <- boundsOptimizeAcc env acc -> + (modified, Abody acc') + +boundsOptimizeAcc + :: forall benv op t. + OperationBounds op + => BoundsEnv benv () + -> OperationAcc op benv t + -- Returns the set of arrays that have been modified, + -- A new BoundsEnv (extended with new information from this term), + -- the bounds of the return value and a transformed term. + -> (IdxSet benv, BoundsEnv benv (), TermBounds (UniformEnv benv ()) t, OperationAcc op benv t) +boundsOptimizeAcc env@(BoundsEnv _ _ zero bindings) acc = case acc of + Exec op args + | modified <- IdxSet.fromList $ argsOutputs args + , input <- boundsInputs env args + -- Remove all references to 'modified' in the bounds graph and input + , (graph', input') <- boundsGraphWithArgsClearNodes (boundsGraph env, input) $ accIdxSet modified + , env' <- env{ boundsGraph = graph' } + , (output, args') <- boundsOptimizeOp op env input' args + , env'' <- envWriteArgs args' output env' -- Update env' with output + -> (modified, env'', TupRunit, Exec op args') + + Return vars -> + ( IdxSet.empty + , env + , mapTupR (boundVar zero . accIdx env . varIdx) vars + , Return vars ) + + Manifest var + | ix <- accIdx env $ varIdx var + -> (IdxSet.empty, env, TupRsingle $ boundVar zero ix, Manifest var) + + Compute expr + | (bounds, expr') <- boundsOptimizeExp env expr + -> (IdxSet.empty, env, bounds, Compute expr') + + Unit var + | bound <- boundVar zero $ accIdx env $ varIdx var + -> (IdxSet.empty, env, TupRsingle $ castTermBound bound, Unit var) + + Alet lhs uniquenesses bnd body + | (modified1, env1, bndBounds, bnd') <- boundsOptimizeAcc env bnd + , env2 <- env1 `pushBuffers` (lhs, bndBounds) + -- If this binds a single expression, an assertion or an assumption, + -- mark this binding in the environment + , env3 <- case lhs of + LeftHandSideSingle _ + | Compute expr <- bnd -> env2{ boundsBindings = boundsBindings env1 `WPushB` weaken (weakenSucc weakenId) (BindExp expr) } + | Aassert _ expr <- bnd -> env2{ boundsBindings = boundsBindings env1 `WPushB` weaken (weakenSucc weakenId) (BindAssertAssume expr) } + | Aassume expr <- bnd -> env2{ boundsBindings = boundsBindings env1 `WPushB` weaken (weakenSucc weakenId) (BindAssertAssume expr) } + _ -> env2 + , (modified2, env4, bodyBounds, body') <- boundsOptimizeAcc env3 body + , (graph, bodyBounds') <- strengthenBoundsWithBufferLHS (boundsGraph env4) lhs bodyBounds -> + ( modified1 `IdxSet.union` IdxSet.drop' lhs modified2 + , env{ boundsGraph = graph } + , bodyBounds' + , Alet lhs uniquenesses bnd' body' + ) + + Acond condVar trueOp falseOp + | cond <- case wprj (varIdx condVar) (boundsBindings env) of + BindExp expr -> expr + _ -> ArrayInstr (Parameter condVar) Nil + , (modified1, _, trueBounds, true') <- boundsOptimizeAcc (assumeTrue env cond) trueOp + , (modified2, _, falseBounds, false') <- boundsOptimizeAcc (assumeFalse env cond) falseOp + , modified <- IdxSet.union modified1 modified2 + -- We cannot easily intersect the two environments (yet), so we just remove any + -- invalid information from the original environment and return that. + -- TODO: Can we take the intersection of two environments? + , env' <- env{ boundsGraph = boundsGraphClearNodes (boundsGraph env) $ accIdxSet modified } + , bounds <- unions (makeTransitives env trueBounds) (makeTransitives env falseBounds) + -> (modified, env', bounds, Acond condVar true' false') + + Awhile uniquenesses cond step initial + | (modified1, cond') <- boundsOptimizeOpenAfun env cond + , (modified2, step') <- boundsOptimizeOpenAfun env step + , modified <- IdxSet.union modified1 modified2 + , env' <- env{ boundsGraph = boundsGraphClearNodes (boundsGraph env) $ accIdxSet modified } + -> (modified, env', bottomsGround zero $ mapTupR varType initial, Awhile uniquenesses cond' step' initial) + + Fence set body + | assertions <- + mapMaybe + (\case + Exists (BindAssertAssume term) -> Just term + _ -> Nothing + ) + $ wprjSet set (boundsBindings env) + , env' <- foldl' assumeTrue env assertions + , (modified, _, _, body') <- boundsOptimizeAcc env' body + -- We cannot propagate any information from within the body of this fence. + -- If we do that, and later act on it, then we may let other code work + -- without using the values computed here. That may make this code dead, + -- and that may cause that the assertions referenced by this fence are not + -- ran. + -- + -- TODO: If all variables in the set refer to assertions that we already + -- proved to be true, or assumes, then we can actually keep this + -- information. + -- + -- We thus return the environment from before this fence, with the + -- information of mutated buffers removed. + , env'' <- env{ boundsGraph = boundsGraphClearNodes (boundsGraph env) $ accIdxSet modified } + -> (IdxSet.empty, env'', bottomsGround zero $ groundsR acc, Fence set $ body') + + Atrace msg t -> (IdxSet.empty, env, TupRsingle $ bottom zero scalarTypeWord8, Atrace msg t) + + Aassert msg expr + | (_, expr') <- boundsOptimizeExp env expr + , expr'' <- simplifyExp expr' + -> case expr'' of + Const _ 1 -> (IdxSet.empty, env, TupRsingle $ bottom zero $ scalarTypeWord8, Compute $ Const scalarTypeWord8 1) + _ | Just idx <- lookupAssertion bindings expr'' -> + (IdxSet.empty, env, TupRsingle $ bottom zero $ scalarTypeWord8, Fence (IdxSet.singleton idx) $ Compute $ Const scalarTypeWord8 1) + _ -> (IdxSet.empty, env, TupRsingle $ bottom zero $ scalarTypeWord8, Aassert msg expr') + + Aassume expr + | (_, expr') <- boundsOptimizeExp env expr + -> (IdxSet.empty, env, TupRsingle $ bottom zero $ scalarTypeWord8, Aassume expr') + + -- Default + Alloc{} -> defaultBottom + Use{} -> defaultBottom + where + -- Default, for expressions that don't mutate buffers, + -- have no subexpressions and whose bound should be bottom. + defaultBottom :: (IdxSet benv, BoundsEnv benv (), TermBounds (UniformEnv benv ()) t, OperationAcc op benv t) + defaultBottom = (IdxSet.empty, env, bottomsGround zero $ groundsR acc, acc) + +boundsOptimizeFun + :: BoundsEnv benv env + -> OpenFun env benv f + -> OpenFun env benv f +boundsOptimizeFun env@(BoundsEnv _ _ zero _) = \case + Lam lhs f + | env' <- env `pushScalars` (lhs, bottoms zero $ lhsToTupR lhs) -> + Lam lhs $ boundsOptimizeFun env' f + Body expr -> Body $ snd $ boundsOptimizeExp env expr + +boundsOptimizeFun1 + :: forall benv s t. + BoundsEnv benv () + -> Arg benv (Fun' (s -> t)) + -> TermBounds (UniformEnv benv ()) s + -> (TermBounds (UniformEnv benv ()) t, Arg benv (Fun' (s -> t))) +boundsOptimizeFun1 env (ArgFun (Lam lhs (Body expr))) argBounds + | env' <- env `pushScalars` (lhs, argBounds) + , (retBounds, expr') <- boundsOptimizeExp env' expr = + ( snd $ strengthenBoundsWithScalarLHS @benv (boundsGraph env') lhs retBounds + , ArgFun $ Lam lhs $ Body expr' + ) +boundsOptimizeFun1 _ _ _ = internalError "Expected unary function" + +boundsOptimizeFun2 + :: forall benv r s t. + BoundsEnv benv () + -> Arg benv (Fun' (r -> s -> t)) + -> TermBounds (UniformEnv benv ()) r + -> TermBounds (UniformEnv benv ()) s + -> (TermBounds (UniformEnv benv ()) t, Arg benv (Fun' (r -> s -> t))) +boundsOptimizeFun2 env (ArgFun (Lam lhs1 (Lam lhs2 (Body expr)))) argBounds1 argBounds2 + | lhs <- LeftHandSidePair lhs1 lhs2 + , env' <- env `pushScalars` (lhs, TupRpair argBounds1 argBounds2) + , (retBounds, expr') <- boundsOptimizeExp env' expr = + ( snd $ strengthenBoundsWithScalarLHS @benv (boundsGraph env') lhs retBounds + , ArgFun $ Lam lhs1 $ Lam lhs2 $ Body expr' + ) +boundsOptimizeFun2 _ _ _ _ = internalError "Expected binary function" + +boundsOptimizeExp + :: forall benv env t. + BoundsEnv benv env + -> OpenExp env benv t + -> (TermBounds (UniformEnv benv env) t, OpenExp env benv t) +boundsOptimizeExp env@(BoundsEnv _ _ zero _) expr = detectConst env $ case expr of + Const (SingleScalarType (NumSingleType (IntegralNumType tp))) val + | IntegralDict <- integralDict tp -> + ( TupRsingle $ boundConst zero $ fromIntegral val + , expr ) + + Const tp _ -> + ( TupRsingle $ bottom zero tp + , expr ) + + Undef tp -> + ( TupRsingle $ undefBound zero tp + , expr ) + + Let lhs bnd body + | (bndBounds, bnd') <- boundsOptimizeExp env bnd + , env' <- env `pushScalars` (lhs, bndBounds) + , (bodyBounds, body') <- boundsOptimizeExp env' body -> + ( snd $ strengthenBoundsWithScalarLHS @benv (boundsGraph env') lhs bodyBounds + , Let lhs bnd' body' ) + + Evar (Var _ ix) + | ix' <- expIdx env ix -> + -- Mark this value as equal to the variable + ( TupRsingle $ boundVar zero ix', expr ) + ArrayInstr arr arg -> case arr of + Parameter (Var _ ix) + | ix' <- accIdx env ix -> + -- Mark this value as equal to the variable + ( TupRsingle $ boundVar zero ix', expr ) + + Index (Var (GroundRscalar tp) _) -> bufferImpossible tp + Index v@(Var (GroundRbuffer _) ix) -> + -- This value has the same bounds as the buffer. + ( TupRsingle $ castTermBound $ boundOfAcc env ix + , ArrayInstr (Index v) $ travE arg + ) + + Foreign tp asm f x -> (bottoms zero tp, Foreign tp asm f $ travE x) -- TODO: Optimize the fallback of Foreign + Pair a b + | (aBounds, a') <- boundsOptimizeExp env a + , (bBounds, b') <- boundsOptimizeExp env b -> + (TupRpair aBounds bBounds, Pair a' b') + Nil -> (TupRunit, Nil) + + -- Analysis doesn't (yet) track SIMD vectors + VecPack vecR e -> bottomExpr $ VecPack vecR $ travE e + VecUnpack vecR e -> bottomExpr $ VecUnpack vecR $ travE e + + ToIndex shr sh ix -> bottomExpr $ ToIndex shr (travE sh) (travE ix) + FromIndex shr sh ix + | (shBounds, sh') <- boundsOptimizeExp env sh + , (TupRsingle ixBound, ix') <- boundsOptimizeExp env ix -> + (fromIndex zero shr shBounds ixBound, FromIndex shr sh' ix') + + ShapeSize shr sh -> bottomExpr $ ShapeSize shr (travE sh) + + Case _ [] (Just def) -> boundsOptimizeExp env def + Case _ [] Nothing -> internalError "Illegal empty case" + Case tag alts def + | (bounds, alts', def') <- caseAlts alts def -> + (bounds, Case (travE tag) alts' def') + + Cond c trueExp falseExp -> case travE c of + -- Check if the condition is already known based on the bounds analysis + Const _ 1 -> boundsOptimizeExp env trueExp + Const _ 0 -> boundsOptimizeExp env falseExp + + c' + | (trueBounds, true') <- boundsOptimizeExp (assumeTrue env c') trueExp + , (falseBounds, false') <- boundsOptimizeExp (assumeFalse env c') falseExp -> + ( unions (makeTransitives env trueBounds) (makeTransitives env falseBounds) + , Cond c' true' false' ) + + Select c t f -> case travE c of + -- Check if the condition is already known based on the bounds analysis + Const _ 1 -> boundsOptimizeExp env t + Const _ 0 -> boundsOptimizeExp env f + + c' -> + let (trueBounds, t') = boundsOptimizeExp env t + (falseBounds, f') = boundsOptimizeExp env f + bs = unions (makeTransitives env trueBounds) + (makeTransitives env falseBounds) + + isBoolBounds :: TermBounds (UniformEnv benv env) Word8 -> Bool + isBoolBounds (TupRsingle bound) = boundIsBool (boundsZero env) $ makeTransitive env bound + + simplify + | Just Refl <- matchTypeR (expType t') (TupRsingle scalarTypeWord8) + , isBoolBounds trueBounds + , isBoolBounds falseBounds + = case (t', f') of + (_ , Const _ 0) -> (bs, PrimApp PrimLAnd (Pair c' t')) -- c ? t : 0 => c && t + (Const _ 0, _ ) -> (bs, PrimApp PrimLAnd (Pair (PrimApp PrimLNot c') f')) -- c ? 0 : f => (not c) && f + (Const _ 1, _ ) -> (bs, PrimApp PrimLOr (Pair c' f')) -- c ? 1 : f => c || f + (_ , Const _ 1) -> (bs, PrimApp PrimLOr (Pair (PrimApp PrimLNot c') t')) -- c ? t : 1 => (not c) || t + _ -> (bs, Select c' t' f') + | otherwise = (bs, Select c' t' f') + in simplify + + Assert msg c body -> case travE c of + Const _ 1 -> boundsOptimizeExp env body + c'@(Const _ 0) -> Assert msg c' <$> boundsOptimizeExp env (undefs $ expType body) + c' + | (_, body') <- boundsOptimizeExp (assumeTrue env c') body -> + -- Don't propagate the bounds of the body, + -- as acting on them could cause this expression to not be used any more, + -- and that could cause that the assertion is not ran at all. + bottomExpr $ Assert msg c' body' + + Assume c body -> case travE c of + Const _ 1 -> boundsOptimizeExp env body + Const _ 0 -> boundsOptimizeExp env (undefs $ expType body) + c' + | (bodyBounds, body') <- boundsOptimizeExp (assumeTrue env c') body -> + (bodyBounds, Assume c' body') + + While cond step initial + | cond' <- travF cond + , Lam lhsCond (Body bodyCond) <- cond' + , Lam lhsStep (Body bodyStep) <- step + , env' <- env `pushScalars` (lhsStep, bottoms zero $ lhsToTupR lhsStep) + , env'' <- assumeTrue' (strengthenShrunkLHS lhsCond lhsStep Just) env' bodyCond + , (stepBounds, bodyStep') <- boundsOptimizeExp env'' bodyStep + , stepBounds' <- snd $ strengthenBoundsWithScalarLHS @benv (boundsGraph env'') lhsStep stepBounds + , (initialBounds, initial') <- boundsOptimizeExp env initial -> + ( unions (makeTransitives env initialBounds) (makeTransitives env stepBounds') + , While cond' (Lam lhsStep $ Body bodyStep') initial' ) + | otherwise -> + internalError "Expected unary functions" + + PrimApp f arg + | (argBounds, arg') <- boundsOptimizeExp env arg -> + app zero f arg' argBounds (makeTransitives env argBounds) + + Coerce t1 t2 a -> case boundsOptimizeExp env a of + (TupRsingle bound, a') -> + ( TupRsingle $ cast zero t1 t2 bound + , Coerce t1 t2 a' ) + (TupRunit, _) -> unitImpossible t1 + (TupRpair _ _, _) -> pairImpossible t1 + + where + travE :: OpenExp env benv s -> OpenExp env benv s + travE = snd . boundsOptimizeExp env + + travF :: OpenFun env benv f -> OpenFun env benv f + travF = boundsOptimizeFun env + + bottomExpr :: OpenExp env benv t -> (TermBounds (UniformEnv benv env) t, OpenExp env benv t) + bottomExpr e = (bottoms zero $ expType e, e) + + caseAlts :: [(TAG, OpenExp env benv b)] -> Maybe (OpenExp env benv b) -> (TermBounds (UniformEnv benv env) b, [(TAG, OpenExp env benv b)], Maybe (OpenExp env benv b)) + caseAlts alts (Just def) + | (altsBounds, alts', _) <- caseAlts alts Nothing + , (defBounds, def') <- boundsOptimizeExp env def + = (unions altsBounds $ makeTransitives env defBounds, alts', Just def') + caseAlts alts Nothing + | (altsBoundss, alts') <- unzip $ map (\(t, e) -> let (b, e') = boundsOptimizeExp env e in (b, (t, e'))) alts + = (foldl1 unions $ map (makeTransitives env) altsBoundss, alts', Nothing) + +detectConst :: BoundsEnv benv env -> (TermBounds (UniformEnv benv env) t, OpenExp env benv t) -> (TermBounds (UniformEnv benv env) t, OpenExp env benv t) +detectConst env (bound, expr) + | TupRsingle b <- bound + , Just a <- upperConst env b + , Just a == lowerConst env b + , TupRsingle tp <- expType expr + , SingleScalarType (NumSingleType (IntegralNumType t)) <- tp + , IntegralDict <- integralDict t + = (bound, Const tp $ fromIntegral a) + | otherwise + = (bound, expr) + +boundsOptimizeArg :: BoundsEnv benv () -> Arg benv t -> Arg benv t +boundsOptimizeArg env = \case + ArgFun fun -> ArgFun $ boundsOptimizeFun env fun + ArgExp expr -> ArgExp $ snd $ boundsOptimizeExp env expr + arg -> arg -- ArgVar or ArgArray + +data InputBounds benv s where + InputIn + :: TermBounds (UniformEnv benv ()) sh + -> TermBounds (UniformEnv benv ()) t + -> InputBounds benv (In sh t) + + InputMut + :: TermBounds (UniformEnv benv ()) sh + -> TermBounds (UniformEnv benv ()) t + -> InputBounds benv (Mut sh t) + + InputOut + :: TermBounds (UniformEnv benv ()) sh + -> InputBounds benv (Out sh t) + + InputVar + :: TermBounds (UniformEnv benv ()) t + -> InputBounds benv (Var' t) + + InputExp + :: InputBounds benv (Exp' t) + + InputFun + :: InputBounds benv (Fun' f) + +data OutputBounds benv s where + OutputOut + :: TermBounds (UniformEnv benv ()) t + -> OutputBounds benv (Out sh t) + + OutputMut + :: TermBounds (UniformEnv benv ()) t + -> OutputBounds benv (Out sh t) + + OutputNone + :: OutputBounds benv s + +envWriteArgs :: forall benv args. Args benv args -> OutputBoundArgs benv args -> BoundsEnv benv () -> BoundsEnv benv () +envWriteArgs ArgsNil _ env = env +envWriteArgs (arg :>: args) (output :>: outputs) env + | ArgArray _ (ArrayR _ tp) _ buffers <- arg + , Just bounds <- case output of + OutputOut b -> Just b + OutputMut b -> Just b + OutputNone -> Nothing + = envWriteArgs args outputs $ go tp buffers bounds env + | otherwise = envWriteArgs args outputs env + where + go :: forall t. TypeR t -> GroundVars benv (Buffers t) -> TermBounds (UniformEnv benv ()) t -> BoundsEnv benv () -> BoundsEnv benv () + go (TupRsingle tp) (TupRsingle var) (TupRsingle bound) env1 + | Refl <- reprIsSingle @ScalarType @t @Buffer tp + , idx <- accIdx env $ varIdx var + , bound' <- castTermBound bound + = env1{ boundsGraph = + Graph.insertEdgesFromWith min idx (upper bound') + $ Graph.insertEdgesToWith min (lower bound') idx + $ boundsGraph env1 + } + go (TupRpair t1 t2) (TupRpair v1 v2) (TupRpair b1 b2) env1 + = go t1 v1 b1 $ go t2 v2 b2 env1 + go TupRunit TupRunit TupRunit env1 = env1 + go _ _ _ _ = internalError "Tuple mismatch" + +type InputBoundArgs benv args = PreArgs (InputBounds benv) args +type OutputBoundArgs benv args = PreArgs (OutputBounds benv) args + +boundsInputs :: BoundsEnv benv () -> Args benv args -> InputBoundArgs benv args +boundsInputs env = mapArgs (boundsInput env) + +boundsInput :: BoundsEnv benv () -> Arg benv arg -> InputBounds benv arg +boundsInput env = \case + ArgArray In (ArrayR _ tp) sh buffers -> InputIn + (mapTupR (\sz -> boundVar (boundsZero env) $ accIdx env $ varIdx sz) sh) + (unBufferTermBounds tp $ mapTupR (\buffer -> boundOfAcc env $ varIdx buffer) buffers) + ArgArray Mut (ArrayR _ tp) sh buffers -> InputMut + (mapTupR (\sz -> boundVar (boundsZero env) $ accIdx env $ varIdx sz) sh) + (unBufferTermBounds tp $ mapTupR (\buffer -> boundOfAcc env $ varIdx buffer) buffers) + ArgArray Out _ sh _ -> InputOut + (mapTupR (\sz -> boundVar (boundsZero env) $ accIdx env $ varIdx sz) sh) + ArgVar vars -> InputVar + (mapTupR (\var -> boundVar (boundsZero env) $ accIdx env $ varIdx var) vars) + ArgExp _ -> InputExp + ArgFun _ -> InputFun + +-- Variant of 'boundsGraphClearNode' that works on a graph *and Args*. +-- Assumes that the Idx refers to a Buffer. +boundsGraphWithArgsClearNode + :: forall benv args t. + (BoundsGraph (UniformEnv benv ()), InputBoundArgs benv args) + -> Idx (UniformEnv benv ()) t + -> (BoundsGraph (UniformEnv benv ()), InputBoundArgs benv args) +boundsGraphWithArgsClearNode (graph, args) idx = (boundsGraphClearNode graph idx, mapArgs f args) + where + idxLower = Graph.inn graph idx + idxUpper = Graph.out graph idx + + f :: InputBounds benv s -> InputBounds benv s + f (InputIn sh buffers) = InputIn sh $ mapTupR g buffers + f (InputMut sh buffers) = InputMut sh $ mapTupR g buffers + f bounds = bounds -- Does not contain buffers + + g :: TermBound (UniformEnv benv ()) s -> TermBound (UniformEnv benv ()) s + g bound = TermBound lower' upper' + where + lower' + | Just (Graph.InEdge (Edge d)) <- prjPartial idx $ lower bound = + partialRemove idx + $ unionPartialEnv + (\(Graph.InEdge x) (Graph.InEdge y) -> Graph.InEdge $ min x y) + (lower bound) + $ mapPartialEnv + (\(Graph.InEdge (Edge d')) -> Graph.InEdge $ Edge $ d + d') + idxLower + | otherwise = lower bound + upper' + | Just (Edge d) <- prjPartial idx $ upper bound = + partialRemove idx + $ unionPartialEnv min (upper bound) + $ mapPartialEnv (\(Edge d') -> Edge $ d + d') idxUpper + | otherwise = upper bound + +-- Variant of 'boundsGraphClearNodes' that works on a graph *and Args*. +-- Assumes that the Idx refers to a Buffer. +boundsGraphWithArgsClearNodes + :: (BoundsGraph (UniformEnv benv ()), InputBoundArgs benv args) + -> IdxSet (UniformEnv benv ()) + -> (BoundsGraph (UniformEnv benv ()), InputBoundArgs benv args) +boundsGraphWithArgsClearNodes graphArgs set = + foldl' (\ga (Exists idx) -> boundsGraphWithArgsClearNode ga idx) graphArgs + $ IdxSet.toList set + +class OperationBounds op where + boundsOptimizeOp + :: op args + -> BoundsEnv benv () + -> InputBoundArgs benv args + -> Args benv args + -> (OutputBoundArgs benv args, Args benv args) + +boundsOptimizeOpDefault + :: BoundsEnv benv () + -> InputBoundArgs benv args + -> Args benv args + -> (OutputBoundArgs benv args, Args benv args) +boundsOptimizeOpDefault env _ args = + ( mapArgs (\_ -> OutputNone) args + , mapArgs (boundsOptimizeArg env) args ) + +boundsOptimizeMap + :: args ~ (Fun' (s -> t) -> In sh s -> Out sh t -> ()) + => BoundsEnv benv () + -> InputBoundArgs benv args + -> Args benv args + -> (OutputBoundArgs benv args, Args benv args) +boundsOptimizeMap env (_ :>: InputIn _ inBounds :>: _) (f :>: input :>: output :>: _) + | (outBounds, f') <- boundsOptimizeFun1 env f inBounds = + ( OutputNone :>: OutputNone :>: OutputOut outBounds :>: ArgsNil + , f' :>: input :>: output :>: ArgsNil ) + +boundsOptimizeGenerate + :: args ~ (Fun' (sh -> t) -> Out sh t -> ()) + => BoundsEnv benv () + -> InputBoundArgs benv args + -> Args benv args + -> (OutputBoundArgs benv args, Args benv args) +boundsOptimizeGenerate env (_ :>: InputOut shBounds :>: _) (f :>: output :>: _) + | (outBounds, f') <- boundsOptimizeFun1 env f (indexBounds (boundsZero env) shBounds) = + ( OutputNone :>: OutputOut outBounds :>: ArgsNil + , f' :>: output :>: ArgsNil ) + +boundsOptimizeBackpermute + :: args ~ (Fun' (sh' -> sh) -> In sh t -> Out sh' t -> ()) + => BoundsEnv benv () + -> InputBoundArgs benv args + -> Args benv args + -> (OutputBoundArgs benv args, Args benv args) +boundsOptimizeBackpermute env (_ :>: InputIn _ inBounds :>: InputOut shBounds :>: _) (f :>: input :>: output :>: _) + | (_, f') <- boundsOptimizeFun1 env f (indexBounds (boundsZero env) shBounds) = + ( OutputNone :>: OutputNone :>: OutputOut inBounds :>: ArgsNil + , f' :>: input :>: output :>: ArgsNil ) + +-- TODO: Add default implementations for permute and permuteUnique + +boundsOptimizeFold + :: args ~ (Fun' (e -> e -> e) -> Exp' e -> In (sh, Int) e -> Out sh e -> ()) + => BoundsEnv benv () + -> InputBoundArgs benv args + -> Args benv args + -> (OutputBoundArgs benv args, Args benv args) +boundsOptimizeFold env (_ :>: _ :>: inBounds :>: _) (f :>: ArgExp seed :>: input :>: output :>: _) + | (seedBounds, seed') <- boundsOptimizeExp env seed + , f' <- boundsOptimizeArg env f + , (_, foldBounds) <- + boundsOptimizeScanlike env + (InputFun :>: inBounds :>: ArgsNil) + (f' :>: input :>: ArgsNil) + (Just (seedBounds, ArgExp seed')) + = + ( OutputNone :>: OutputNone :>: OutputNone :>: OutputOut foldBounds :>: ArgsNil + , f' :>: ArgExp seed' :>: input :>: output :>: ArgsNil + ) + +boundsOptimizeFold1 + :: args ~ (Fun' (e -> e -> e) -> In (sh, Int) e -> Out sh e -> ()) + => BoundsEnv benv () + -> InputBoundArgs benv args + -> Args benv args + -> (OutputBoundArgs benv args, Args benv args) +boundsOptimizeFold1 env (_ :>: inBounds :>: _) (f :>: input :>: output :>: _) + | f' <- boundsOptimizeArg env f + , (_, foldBounds) <- + boundsOptimizeScanlike env + (InputFun :>: inBounds :>: ArgsNil) + (f' :>: input :>: ArgsNil) + Nothing + = + ( OutputNone :>: OutputNone :>: OutputOut foldBounds :>: ArgsNil + , f' :>: input :>: output :>: ArgsNil + ) + +boundsOptimizeScan + :: args ~ (Fun' (e -> e -> e) -> Exp' e -> In (sh, Int) e -> Out (sh, Int) e -> ()) + => BoundsEnv benv () + -> InputBoundArgs benv args + -> Args benv args + -> (OutputBoundArgs benv args, Args benv args) +boundsOptimizeScan env (_ :>: _ :>: inBounds :>: _) (f :>: ArgExp seed :>: input :>: output :>: _) + | (seedBounds, seed') <- boundsOptimizeExp env seed + , f' <- boundsOptimizeArg env f + , (scanBounds, foldBounds) <- + boundsOptimizeScanlike env + (InputFun :>: inBounds :>: ArgsNil) + (f' :>: input :>: ArgsNil) + (Just (seedBounds, ArgExp seed')) + = + ( OutputNone :>: OutputNone :>: OutputNone :>: OutputOut (unions scanBounds foldBounds) :>: ArgsNil + , f' :>: ArgExp seed' :>: input :>: output :>: ArgsNil + ) + +boundsOptimizeScan' + :: args ~ (Fun' (e -> e -> e) -> Exp' e -> In (sh, Int) e -> Out (sh, Int) e -> Out sh e -> ()) + => BoundsEnv benv () + -> InputBoundArgs benv args + -> Args benv args + -> (OutputBoundArgs benv args, Args benv args) +boundsOptimizeScan' env (_ :>: _ :>: inBounds :>: _) (f :>: ArgExp seed :>: input :>: outputScan :>: outputFold :>: _) + | (seedBounds, seed') <- boundsOptimizeExp env seed + , f' <- boundsOptimizeArg env f + , (scanBounds, foldBounds) <- + boundsOptimizeScanlike env + (InputFun :>: inBounds :>: ArgsNil) + (f' :>: input :>: ArgsNil) + (Just (seedBounds, ArgExp seed')) + = + ( OutputNone :>: OutputNone :>: OutputNone :>: OutputOut scanBounds :>: OutputOut foldBounds :>: ArgsNil + , f' :>: ArgExp seed' :>: input :>: outputScan :>: outputFold :>: ArgsNil + ) + +boundsOptimizeScan1 + :: args ~ (Fun' (e -> e -> e) -> In (sh, Int) e -> Out (sh, Int) e -> ()) + => BoundsEnv benv () + -> InputBoundArgs benv args + -> Args benv args + -> (OutputBoundArgs benv args, Args benv args) +boundsOptimizeScan1 env (_ :>: inBounds :>: _) (f :>: input :>: output :>: _) + | f' <- boundsOptimizeArg env f + , (scanBounds, foldBounds) <- + boundsOptimizeScanlike env + (InputFun :>: inBounds :>: ArgsNil) + (f' :>: input :>: ArgsNil) + Nothing + = + ( OutputNone :>: OutputNone :>: OutputOut (unions scanBounds foldBounds) :>: ArgsNil + , f' :>: input :>: output :>: ArgsNil + ) + +boundsOptimizeScanlike + :: args ~ (Fun' (e -> e -> e) -> In (sh, Int) e -> ()) + => BoundsEnv benv () + -> InputBoundArgs benv args + -> Args benv args + -> Maybe (TermBounds (UniformEnv benv ()) e, Arg benv (Exp' e)) + -> (TermBounds (UniformEnv benv ()) e, TermBounds (UniformEnv benv ()) e) +boundsOptimizeScanlike env (_ :>: InputIn inShape inBounds :>: _) (f :>: input :>: _) seed + -- Detect a scan (+) over an array of zeros and ones. + | ArgArray _ (ArrayR _ (TupRsingle tp)) _ _ <- input + , SingleScalarType (NumSingleType (IntegralNumType TypeInt)) <- tp + , ArgFun (Lam LeftHandSideSingle{} (Lam LeftHandSideSingle{} (Body expr))) <- f + , PrimApp PrimAdd{} (Pair (Evar (Var _ (SuccIdx ZeroIdx))) (Evar (Var _ ZeroIdx))) <- expr + , TupRsingle inBound <- inBounds + , (l, u) <- getBoundRange (boundsZero env) TypeInt $ makeTransitive env inBound + , 0 <= l && u <= 1 + , case seed of + Nothing -> True + Just (_, ArgExp (Const _ 0)) -> True + _ -> False + , TupRpair _ (TupRsingle inSize) <- inShape + -- Yes, this is a scan with addition, whose input consists of zeros and ones. + -- Elements in the output are bounded by the input size + , scanBound <- TermBound + (lower $ boundConst (boundsZero env) 0) + (mapPartialEnv (\(Edge d) -> Edge (d - 1)) $ upper inSize) + , foldBound <- TermBound + (lower $ boundConst (boundsZero env) 0) + (upper inSize) + = + ( TupRsingle scanBound + , TupRsingle foldBound + ) +-- TODO: Detect operators that only return values from the left or right argument, +-- like min and max +boundsOptimizeScanlike env _ (_ :>: ArgArray _ (ArrayR _ tp) _ _ :>: _) _ = + ( bottoms (boundsZero env) tp + , bottoms (boundsZero env) tp + ) diff --git a/src/Data/Array/Accelerate/Trafo/Operation/Bounds/Algebra.hs b/src/Data/Array/Accelerate/Trafo/Operation/Bounds/Algebra.hs new file mode 100644 index 000000000..a49782074 --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Operation/Bounds/Algebra.hs @@ -0,0 +1,429 @@ +{-# OPTIONS_HADDOCK hide #-} +-- | +-- Module : Data.Array.Accelerate.Trafo.Operation.Bounds.Algebra +-- Copyright : [2012..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.Trafo.Operation.Bounds.Algebra where + +import Data.Array.Accelerate.AST.Environment +import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.Graph (InEdge(..)) +import Data.Array.Accelerate.AST.Operation +import Data.Array.Accelerate.Array.Buffer +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Representation.Shape hiding (union) +import Data.Array.Accelerate.Type +import Data.Array.Accelerate.Error + +import Data.Kind (Type) +import Data.Typeable ( (:~:)(..) ) + +-- x <= y becomes an edge from x to y with distance 0. +-- x < y becomes an edge from x to y with distance -1. +-- +-- x <= y + c could become an edge from x to y with distance c. +-- We however do not track them yet, as integer overflow may happen, which is +-- difficult to track. To keep the analysis sound without making it too +-- difficult, we thus ignore that (for now). +-- +data Edge s t = Edge { distance :: Integer } deriving (Eq, Ord) + +-- We use bound (singular) to refer to the combination of a lowerbound and an upperbound. +-- Bounds (plural) refers to a tuple where each value is a bound. +-- +-- t should be an IntegralType or a Buffer of an IntegralType +data TermBound (env :: Type) (t :: Type) = TermBound + -- An edge from a variable x with distance c means: + -- x <= this + c + { lower :: PartialEnv (InEdge Edge t) env + -- An edge to a variable y with distance c means: + -- this <= y + c + , upper :: PartialEnv (Edge t) env + } + +type TermBounds env = TupR (TermBound env) + +bottom :: forall t env. Idx env Int -> ScalarType t -> TermBound env t +-- We could give each integral variable a lower bound and upper bound based on +-- the range of its type. That however causes that each variable is connected +-- with each other variable in the graph, so we do not include those bounds. +-- If we ever change our minds, we can use the following code: +-- +-- bottom zero (SingleScalarType (NumSingleType (IntegralNumType tp))) +-- | IntegralDict <- integralDict tp +-- = boundRange zero (fromIntegral (minBound :: t)) (fromIntegral (maxBound :: t)) +bottom _ _ + = TermBound PEnd PEnd + +bottomGround :: Idx env Int -> GroundR t -> TermBound env t +bottomGround zero (GroundRscalar tp) = bottom zero tp +bottomGround zero (GroundRbuffer tp) = castTermBound $ bottom zero tp + +nonNegative :: forall t env. Idx env Int -> ScalarType t -> TermBound env t +nonNegative zero (SingleScalarType (NumSingleType (IntegralNumType tp))) + | IntegralDict <- integralDict tp + = boundRange zero 0 (fromIntegral (maxBound :: t)) +nonNegative _ _ + = TermBound PEnd PEnd + +undefBound :: Idx env Int -> ScalarType t -> TermBound env t +-- Not sure if the range maxBound .. minBound would work in the entire +-- analysis. 0 .. 0 should definitely be safe, but then we would convert +-- Undef to 0. Hence we use 0 .. 1, the next best option. +-- In the future we may want to include a separate 'undef' value in the +-- analysis. +undefBound zero _ = boundRange zero 0 1 + +bottoms :: Idx env Int -> TypeR t -> TermBounds env t +bottoms zero = mapTupR (bottom zero) + +bottomsGround :: Idx env Int -> GroundsR t -> TermBounds env t +bottomsGround zero = mapTupR (bottomGround zero) + +boundRange :: Idx env Int -> Integer -> Integer -> TermBound env t +boundRange zero lowerB upperB = TermBound + -- zero <= this - lowerB + -- which (for infinite-precision Integers) is equal to + -- lowerB <= this + (partialEnvSingleton zero $ InEdge $ Edge $ negate lowerB) + -- this <= zero + upperB + (partialEnvSingleton zero $ Edge upperB) + +boundConst :: Idx env Int -> Integer -> TermBound env t +boundConst zero value = boundRange zero value value + +getBoundRange :: forall t env. Idx env Int -> IntegralType t -> TermBound env t -> (Integer, Integer) +getBoundRange zero tp bound + | IntegralDict <- integralDict tp = + ( case prjPartial zero $ lower bound of + Just (InEdge (Edge distance')) -> negate distance' + Nothing -> fromIntegral (minBound :: t) + , case prjPartial zero $ upper bound of + Just (Edge distance') -> distance' + Nothing -> fromIntegral (maxBound :: t) + ) + +boundIsBool :: Idx env Int -> TermBound env Word8 -> Bool +boundIsBool zero bound = l >= 0 && u <= 1 + where (l, u) = getBoundRange zero TypeWord8 bound + +boundVar :: Idx env Int -> Idx env t -> TermBound env t +boundVar _zero ix = TermBound + (partialEnvSingleton ix $ InEdge $ Edge 0) + (partialEnvSingleton ix $ Edge 0) + +-- Casts the bounds for a type 's' to be used for a term of type 't'. +-- This does not protect for underflow or overflow; casting from Int64 to Word8 +-- is thus not always safe. It is definitely safe to cast from 't' to +-- 'Buffer t' however. +castTermBound :: TermBound env s -> TermBound env t +castTermBound (TermBound l u) = TermBound + (mapPartialEnv (\(InEdge (Edge d)) -> InEdge $ Edge d) l) + (mapPartialEnv (\(Edge d) -> Edge d) u) + +unBufferTermBounds :: forall env t. TypeR t -> TermBounds env (Buffers t) -> TermBounds env t +unBufferTermBounds (TupRsingle tp) (TupRsingle bound) + | Refl <- reprIsSingle @ScalarType @t @Buffer tp = + TupRsingle $ castTermBound bound +unBufferTermBounds (TupRpair t1 t2) (TupRpair b1 b2) = + unBufferTermBounds t1 b1 `TupRpair` unBufferTermBounds t2 b2 +unBufferTermBounds TupRunit TupRunit = TupRunit +unBufferTermBounds _ _ = internalError "Tuple mismatch" + +bufferTermBounds :: forall env t. TypeR t -> TermBounds env t -> TermBounds env (Buffers t) +bufferTermBounds (TupRsingle tp) (TupRsingle bound) + | Refl <- reprIsSingle @ScalarType @t @Buffer tp = + TupRsingle $ castTermBound bound +bufferTermBounds (TupRpair t1 t2) (TupRpair b1 b2) = bufferTermBounds t1 b1 `TupRpair` bufferTermBounds t2 b2 +bufferTermBounds TupRunit TupRunit = TupRunit +bufferTermBounds _ _ = internalError "Tuple mismatch" + +-- | Returns bounds that are valid for both arguments. +-- Some accuracy may be lost. 'makeTransitive' should be called to prevent most of that. +union :: TermBound env t -> TermBound env t -> TermBound env t +union a b = TermBound + -- Note that we use 'max' instead of 'min' here as InEdge swaps the edge around, see the comments in TermBound. + (intersectPartialEnv (\(InEdge x) (InEdge y) -> InEdge $ max x y) (lower a) (lower b)) + (intersectPartialEnv max (upper a) (upper b)) + +unions :: TermBounds env t -> TermBounds env t -> TermBounds env t +unions (TupRsingle a) (TupRsingle b) = TupRsingle $ union a b +unions (TupRpair a1 a2) (TupRpair b1 b2) = unions a1 b1 `TupRpair` unions a2 b2 +unions TupRunit TupRunit = TupRunit +unions _ _ = internalError "Tuple mismatch" + +app + :: Idx env Int -- Index of the zero node in the bounds graph + -> PrimFun (s -> t) + -> OpenExp env' benv s + -> TermBounds env s -- ^ Bounds of the argument + -- | Transitively closed bounds of the argument. Consumed lazily (only when + -- needed for a specific PrimFun), due to the time to make the bounds + -- transitively closed. + -> TermBounds env s + -> (TermBounds env t, OpenExp env' benv t) +app zero f arg (TupRsingle bound) (TupRsingle closed) = app1 zero f arg bound closed +app zero f (Pair e1 e2) (TupRpair (TupRsingle b1) (TupRsingle b2)) (TupRpair (TupRsingle c1) (TupRsingle c2)) = app2 zero f e1 e2 b1 b2 c1 c2 +app zero f arg _ _ = (bottoms zero $ snd $ primFunType f, PrimApp f arg) + +app1 + :: forall env env' benv s t. + Idx env Int + -> PrimFun (s -> t) + -> OpenExp env' benv s + -> TermBound env s + -> TermBound env s + -> (TermBounds env t, OpenExp env' benv t) +app1 zero f arg bound closed = case f of + PrimAbs tp@(IntegralNumType _) + | closed `greaterThanEqual` boundConst zero 0 -> + -- if arg >= 0, then abs arg == arg + (TupRsingle bound, arg) + | boundConst zero 0 `greaterThanEqual` closed -> + (TupRsingle $ bottom zero $ SingleScalarType $ NumSingleType tp, PrimApp (PrimNeg tp) arg) + | otherwise -> + -- Note that we cannot assume that the result of abs is non-negative. + -- For signed integers, abs minBound = minBound. For instance for Int8, + -- minBound = -128, and abs (-128) = -(-128) = -128 due to integer + -- overflow. Hence abs x >= 0 is not true for all x. + (TupRsingle $ bottom zero $ SingleScalarType $ NumSingleType tp, PrimApp (PrimAbs tp) arg) + -- TODO: PrimSig + PrimFromIntegral source (IntegralNumType target) -> + withBounds $ TupRsingle $ castIntegral zero source target closed + _ -> withBounds $ bottoms zero $ snd $ primFunType f + where + withBounds :: TermBounds env t -> (TermBounds env t, OpenExp env' benv t) + withBounds retBounds = (retBounds, PrimApp f arg) + +app2 + :: forall env env' benv a b t. + Idx env Int + -> PrimFun ((a, b) -> t) + -> OpenExp env' benv a + -> OpenExp env' benv b + -> TermBound env a + -> TermBound env b + -> TermBound env a -- Transitively closed bound of first argument, see comment in type signature of 'app' + -> TermBound env b -- Transitively closed bound of second argument + -> (TermBounds env t, OpenExp env' benv t) +app2 zero f arg1 arg2 bound1 bound2 closed1 closed2 = case f of + -- Min & max + PrimMin (NumSingleType (IntegralNumType _)) + | bound1 `greaterThanEqual` closed2 -> + (TupRsingle bound2, arg2) + | bound2 `greaterThanEqual` closed1 -> + (TupRsingle bound1, arg1) + | otherwise -> + withBounds $ TupRsingle $ TermBound + -- Note that we use 'max' instead of 'min' here as InEdge swaps the edge around, see the comments in TermBound. + (intersectPartialEnv (\(InEdge x) (InEdge y) -> InEdge $ max x y) (lower closed1) (lower closed2)) + (unionPartialEnv min (upper bound1) (upper bound2)) + PrimMax (NumSingleType (IntegralNumType _)) + | closed2 `greaterThanEqual` bound1 -> + (TupRsingle bound2, arg2) + | closed1 `greaterThanEqual` bound2 -> + (TupRsingle bound1, arg1) + | otherwise -> + withBounds $ TupRsingle $ TermBound + -- Note that we use 'min' instead of 'max' here as InEdge swaps the edge around, see the comments in TermBound. + (unionPartialEnv (\(InEdge x) (InEdge y) -> InEdge $ min x y) (lower bound1) (lower bound2)) + (intersectPartialEnv max (upper closed1) (upper closed2)) + -- Comparisons + PrimCmp _ CmpEq + | closed1 `equal` bound2 -> true zero + | closed1 `notEqual` bound2 -> false zero + PrimCmp _ CmpNEq + | closed1 `equal` bound2 -> false zero + | closed1 `notEqual` bound2 -> true zero + PrimCmp (NumSingleType (IntegralNumType _)) CmpLt + | closed1 `lessThan` bound2 -> true zero + | closed1 `greaterThanEqual` bound2 -> false zero + PrimCmp (NumSingleType (IntegralNumType _)) CmpGtEq + | closed1 `lessThan` bound2 -> false zero + | closed1 `greaterThanEqual` bound2 -> true zero + PrimCmp (NumSingleType (IntegralNumType TypeWord8)) cmp + | Const _ constant <- arg1 + , boundIsBool zero closed2 + , Just result <- cmpBool cmp constant arg2 + -> (TupRsingle $ boundRange zero 0 1, result) + | Const _ constant <- arg2 + , boundIsBool zero closed1 + -- Note that we reorder the arguments to cmpBool (constant and arg1); + -- this is sound as we only handle CmpEq and CmpNEq there + -- (and not CmpLt and CmpGtEq) + , Just result <- cmpBool cmp constant arg1 + -> (TupRsingle $ boundRange zero 0 1, result) + PrimCmp _ _ -> bool + -- Other boolean operations + PrimIsNaN _ -> bool + PrimIsInfinite _ -> bool + PrimLAnd -> bool + PrimLOr -> bool + -- Div, Mod, Quot and rem + PrimIDiv tp + | Just (divBounds, _, divExpr, _, _) <- divModOrQuotRem tp + -> (divBounds, divExpr) + PrimQuot tp + | Just (divBounds, _, divExpr, _, _) <- divModOrQuotRem tp + -> (divBounds, divExpr) + PrimMod tp + | Just (_, modBounds, _, modExpr, _) <- divModOrQuotRem tp + -> (modBounds, modExpr) + PrimRem tp + | Just (_, modBounds, _, modExpr, _) <- divModOrQuotRem tp + -> (modBounds, modExpr) + PrimDivMod tp + | Just (divBounds, modBounds, _, _, divModExpr) <- divModOrQuotRem tp + -> (TupRpair divBounds modBounds, divModExpr) + PrimQuotRem tp + | Just (divBounds, modBounds, _, _, divModExpr) <- divModOrQuotRem tp + -> (TupRpair divBounds modBounds, divModExpr) + + -- Default fallback + _ -> withBounds $ bottoms zero $ snd $ primFunType f + where + withBounds :: TermBounds env t -> (TermBounds env t, OpenExp env' benv t) + withBounds retBounds = (retBounds, PrimApp f $ Pair arg1 arg2) + + bool :: t ~ PrimBool => (TermBounds env t, OpenExp env' benv t) + bool = withBounds $ TupRsingle $ boundRange zero 0 1 + + -- Optimize comparisons of boolean expressions with constants. + -- 'arg' is thus assumed to be an expression that can only evaluate to 0 or + -- 1. Note that for the constant, we only pattern match on 0 and 1, + -- out of bounds arguments are already be handled in prior cases. + cmpBool :: Cmp -> PrimBool -> OpenExp env' benv PrimBool -> Maybe (OpenExp env' benv PrimBool) + cmpBool CmpEq 1 arg = Just arg -- True == arg --> arg + cmpBool CmpEq 0 arg = Just $ PrimApp PrimLNot arg -- False == arg --> not arg + cmpBool CmpNEq 1 arg = Just $ PrimApp PrimLNot arg -- True /= arg --> not arg + cmpBool CmpNEq 0 arg = Just arg -- False /= arg --> arg + -- Note: we don't pattern match on CmpLt and CmpGtEq, as that allows us to + -- freely reorder the arguments (the constant and the expression). + cmpBool _ _ _ = Nothing + + -- TODO: For rem, always set bounds, also if we only know that arg2 is positive? + divModOrQuotRem :: a ~ b => IntegralType a -> Maybe (TermBounds env a, TermBounds env a, OpenExp env' benv a, OpenExp env' benv a, OpenExp env' benv (a, a)) + divModOrQuotRem tp + -- If arg1 >= 0 and arg2 > 0, then divMod and quotRem behave the same way. + -- Since the hardware natively supports quot and rem, not div and mod, we + -- convert div and mod to quot and rem. + | closed1 `greaterThanEqual` boundConst zero 0 + , boundConst zero 0 `lessThan` closed2 = + if bound1 `lessThan` closed2 then + let + divExpr = case integralDict tp of + IntegralDict -> Const (SingleScalarType $ NumSingleType $ IntegralNumType tp) 0 + modExpr = arg1 + in Just + -- If the above, and arg1 < arg2, then + -- arg1 `divMod` arg2 = arg1 `quotRem` arg2 = (0, arg1) + ( TupRsingle $ boundConst zero 0 + , TupRsingle bound1 + , divExpr + , modExpr + , Pair divExpr modExpr + ) + else Just + -- Bounds of div or quot + ( TupRsingle $ nonNegative zero $ SingleScalarType $ NumSingleType $ IntegralNumType tp + -- Bounds of mod or rem + , TupRsingle $ TermBound + (partialEnvSingleton zero $ InEdge $ Edge 0) + (upper bound2) + , PrimApp (PrimQuot tp) $ Pair arg1 arg2 + , PrimApp (PrimRem tp) $ Pair arg1 arg2 + , PrimApp (PrimQuotRem tp) $ Pair arg1 arg2 + ) + | otherwise = Nothing + +true :: Idx env Int -> (TermBounds env PrimBool, OpenExp env' benv PrimBool) +true zero = (TupRsingle $ boundConst zero 1, Const scalarType 1) + +false :: Idx env Int -> (TermBounds env PrimBool, OpenExp env' benv PrimBool) +false zero = (TupRsingle $ boundConst zero 0, Const scalarType 0) + +indexBounds :: forall env sh. Idx env Int -> TermBounds env sh -> TermBounds env sh +indexBounds zero = mapTupR f + where + f :: TermBound env t -> TermBound env t + f (TermBound _ u) = TermBound + (partialEnvSingleton zero $ InEdge $ Edge 0) + (mapPartialEnv (\(Edge d) -> Edge $ d - 1) u) + +fromIndex :: forall env sh. Idx env Int -> ShapeR sh -> TermBounds env sh -> TermBound env Int -> TermBounds env sh +-- TODO: We could make this more sophisticated by using the constant bounds on the Int argument. +fromIndex zero _ sz _ = indexBounds zero sz + +-- Returns true if we can already proof that the first argument is greater than +-- or equal to the second argument +greaterThanEqual :: TermBound env a -> TermBound env a -> Bool +-- Search for a node that is a lower bound of the first argument, +-- and an upper bound of the second argument. +-- If we find such a node x, then +-- x <= arg1 + a +-- and +-- arg2 <= x + b +-- Thus arg2 <= x + b <= arg1 + a + b +-- Thus if a + b <= 0, then arg2 <= arg1 +greaterThanEqual (TermBound low _) (TermBound _ up) = + or + $ partialEnvValues + $ intersectPartialEnv + (\(InEdge (Edge a)) (Edge b) -> IdentityF $ a + b <= 0) + low + up + +lessThan :: TermBound env a -> TermBound env a -> Bool +lessThan (TermBound _ up) (TermBound low _) = + -- Similar to greaterThanEqual + or + $ partialEnvValues + $ intersectPartialEnv + (\(InEdge (Edge a)) (Edge b) -> IdentityF $ a + b < 0) + low + up + +equal :: TermBound env a -> TermBound env a -> Bool +equal a b = greaterThanEqual a b && greaterThanEqual b a + +notEqual :: TermBound env a -> TermBound env a -> Bool +notEqual a b = lessThan a b || lessThan b a + +cast :: Idx env Int -> ScalarType s -> ScalarType t -> TermBound env s -> TermBound env t +cast zero (SingleScalarType (NumSingleType (IntegralNumType source))) (SingleScalarType (NumSingleType (IntegralNumType target))) bound + = castIntegral zero source target bound +cast zero _ target _ = bottom zero target + +castIntegral :: Idx env Int -> IntegralType s -> IntegralType t -> TermBound env s -> TermBound env t +castIntegral zero source target bound + | (sourceMin, sourceMax) <- getBoundRange zero source bound + , (targetMin, targetMax) <- intRange target + , targetMin <= sourceMin -- Check if the source fits in the target + , sourceMax <= targetMax + = castTermBound bound -- It fits, we can preserve all information + | otherwise + -- It might not fit, we can't preserve bound information due to possible overflow or underflow + = bottom zero $ SingleScalarType $ NumSingleType $ IntegralNumType target + +intRange :: forall t. IntegralType t -> (Integer, Integer) +intRange tp + | IntegralDict <- integralDict tp = (fromIntegral (minBound :: t), fromIntegral (maxBound :: t)) + +-- Returns Nothing if the value is above or below the valid range of this integral type +guardOverflow :: IntegralType t -> Integer -> Maybe Integer +guardOverflow tp x + | x < a || x > b = Nothing + | otherwise = Just x + where + (a, b) = intRange tp + +saturate :: IntegralType t -> Integer -> Integer +saturate tp x = max a $ min b x + where + (a, b) = intRange tp diff --git a/src/Data/Array/Accelerate/Trafo/Operation/Bounds/Environment.hs b/src/Data/Array/Accelerate/Trafo/Operation/Bounds/Environment.hs new file mode 100644 index 000000000..d425ab6b2 --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Operation/Bounds/Environment.hs @@ -0,0 +1,349 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# OPTIONS_HADDOCK hide #-} +-- | +-- Module : Data.Array.Accelerate.Trafo.Operation.Bounds.Environment +-- Copyright : [2012..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.Trafo.Operation.Bounds.Environment where + +import Data.Array.Accelerate.AST.Environment +import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.IdxSet (IdxSet(..)) +import Data.Array.Accelerate.AST.Graph (InEdge(..)) +import qualified Data.Array.Accelerate.AST.Graph as Graph +import qualified Data.Array.Accelerate.AST.IdxSet as IdxSet +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.AST.Operation +import Data.Array.Accelerate.Analysis.Match +import Data.Array.Accelerate.Trafo.WeakenedEnvironment +import Data.Array.Accelerate.Trafo.Substitution +import Data.Array.Accelerate.Trafo.Exp.Substitution +import Data.Array.Accelerate.Trafo.Operation.Substitution () +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Type +import Data.Array.Accelerate.Error + +import qualified Data.Functor.Const as Functor + +import Data.Array.Accelerate.Trafo.Operation.Bounds.Algebra + +import Data.Maybe +import Data.List (foldl') + +type BoundsGraph = Graph.Graph Node Edge + +data BoundsEnv benv env = BoundsEnv + { boundsGraph :: BoundsGraph (UniformEnv benv env) + -- way to go from index in 'benv' to index in 'UniformEnv benv env' + , boundsSkipBenv :: Skip (UniformEnv benv env) (Append ((), Int) benv) + -- reference to the last index in 'UniformEnv benv env', representing the zero value + , boundsZero :: Idx (UniformEnv benv env) Int + , boundsBindings :: WEnv Binding benv + } + +data Binding benv t where + -- No important information + BindNone :: Binding benv t + + -- This variable is defined with a Compute node + BindExp :: Exp benv t -> Binding benv t + + -- This variable is defined with an AAssert or AAssume, + -- and can be synchronised with a Fence + BindAssertAssume :: Exp benv PrimBool -> Binding benv PrimBool + +emptyEnv :: BoundsEnv () () +emptyEnv = BoundsEnv (Graph.pushNode Graph.empty Node PEnd PEnd) SkipNone ZeroIdx WEmpty + +instance Sink Binding where + weaken _ BindNone = BindNone + weaken k (BindExp expr) = BindExp $ mapArrayInstr (weaken k) expr + weaken k (BindAssertAssume expr) = BindAssertAssume $ mapArrayInstr (weaken k) expr + +-- | Construct one global environment, consisting of: +-- * One artificial node to represent the constant value zero +-- * All identifiers in the buffer environment +-- * All identifiers in the scalar environment +type UniformEnv benv env = (Append (Append ((), Int) benv) env) + +accIdx :: BoundsEnv benv env -> Idx benv t -> Idx (UniformEnv benv env) t +accIdx (BoundsEnv _ s _ _) ix = weaken (skipWeakenIdx s) $ extendIdx @((), Int) ix + +-- This function currently assumes that the scalar environment is (), +-- as that is the case in all current uses. If this changes in the future, +-- we can generalize this function. We should then pass a 'BoundsEnv benv env' +-- as argument. +accIdxSet :: IdxSet benv -> IdxSet (UniformEnv benv ()) +accIdxSet (IdxSet penv) = IdxSet $ accPartialEnv penv + +accPartialEnv :: PartialEnv f benv -> PartialEnv f (UniformEnv benv ()) +accPartialEnv (PPush set a) = accPartialEnv set `PPush` a +accPartialEnv (PNone set) = PNone $ accPartialEnv set +accPartialEnv PEnd = PEnd + +-- First argument is not used, but by including it in the type, +-- the type becomes non-ambiguous +expIdx :: forall benv env t. BoundsEnv benv env -> Idx env t -> Idx (UniformEnv benv env) t +expIdx _ idx = extendIdx @(Append ((), Int) benv) idx + +-- This type is ambiguous. 'lbenv' must always be explicitely instantiated. +extendIdx :: forall lenv renv t. Idx renv t -> Idx (Append lenv renv) t +extendIdx ZeroIdx = ZeroIdx +extendIdx (SuccIdx idx) = SuccIdx $ extendIdx @lenv idx + +boundOfAcc :: BoundsEnv benv env -> Idx benv t -> TermBound (UniformEnv benv env) t +boundOfAcc env ix = TermBound + (mapPartialEnv (\(Graph.InEdge (Edge d)) -> Graph.InEdge $ Edge d) $ Graph.inn (boundsGraph env) ix') + (mapPartialEnv (\(Edge d) -> Edge d) $ Graph.out (boundsGraph env) ix') + where + ix' = accIdx env ix + +data Node t = Node + +pushScalar :: BoundsEnv benv env -> TermBound (UniformEnv benv env) t -> BoundsEnv benv (env, t) +pushScalar (BoundsEnv g s z bs) bound = BoundsEnv + (Graph.pushNode g Node (lower bound) (partialEnvSkip $ upper bound)) + (SkipSucc' s) + (SuccIdx z) + bs + +pushScalars :: forall benv env env' t. BoundsEnv benv env -> (ELeftHandSide t env env', TermBounds (UniformEnv benv env) t) -> BoundsEnv benv env' +pushScalars env (LeftHandSideWildcard _, _) = env +pushScalars env (LeftHandSideSingle _, TupRsingle b) = pushScalar env b +pushScalars env (LeftHandSidePair l1 l2, TupRpair p1 p2) + = pushScalars (pushScalars env (l1, p1)) (l2, mapTupR (weakenBound l1) p2) + where + weakenBound :: ELeftHandSide s env1 env2 -> TermBound (UniformEnv benv env1) a -> TermBound (UniformEnv benv env2) a + weakenBound lhs b = TermBound + (partialEnvSkipLHS' lhs $ lower b) + (partialEnvSkipLHS' lhs $ upper b) + + partialEnvSkipLHS' :: ELeftHandSide s env1 env2 -> PartialEnv u (UniformEnv benv env1) -> PartialEnv u (UniformEnv benv env2) + partialEnvSkipLHS' (LeftHandSideWildcard _) e = e + partialEnvSkipLHS' (LeftHandSideSingle _) e = partialEnvSkip e + partialEnvSkipLHS' (LeftHandSidePair h1 h2) e = partialEnvSkipLHS' h2 $ partialEnvSkipLHS' h1 e +pushScalars _ _ = internalError "Tuple mismatch" + +-- | Pushes information to the buffer environment. +-- Note that this can either be a buffer or a scalar value (in the buffer +-- environment). +pushBuffer :: BoundsEnv benv () -> TermBound (UniformEnv benv ()) t -> BoundsEnv (benv, t) () +pushBuffer (BoundsEnv g _ z bs) bound = BoundsEnv + (Graph.pushNode g Node (lower bound) (partialEnvSkip $ upper bound)) + SkipNone + (SuccIdx z) + (WPushB bs BindNone) + +pushBuffers :: BoundsEnv benv () -> (LeftHandSide r t benv benv', TermBounds (UniformEnv benv ()) t) -> BoundsEnv benv' () +pushBuffers env (LeftHandSideWildcard _, _) = env +pushBuffers env (LeftHandSideSingle _, TupRsingle b) = pushBuffer env b +pushBuffers env (LeftHandSidePair l1 l2, TupRpair p1 p2) + = pushBuffers (pushBuffers env (l1, p1)) (l2, mapTupR (weakenBound l1) p2) + where + weakenBound :: LeftHandSide r s benv1 benv2 -> TermBound (UniformEnv benv1 ()) a -> TermBound (UniformEnv benv2 ()) a + weakenBound lhs b = TermBound + (partialEnvSkipLHS' lhs $ lower b) + (partialEnvSkipLHS' lhs $ upper b) + + partialEnvSkipLHS' :: LeftHandSide r s benv1 benv2 -> PartialEnv u (UniformEnv benv1 ()) -> PartialEnv u (UniformEnv benv2 ()) + partialEnvSkipLHS' (LeftHandSideWildcard _) e = e + partialEnvSkipLHS' (LeftHandSideSingle _) e = partialEnvSkip e + partialEnvSkipLHS' (LeftHandSidePair h1 h2) e = partialEnvSkipLHS' h2 $ partialEnvSkipLHS' h1 e +pushBuffers _ _ = internalError "Tuple mismatch" + +makeTransitives :: BoundsEnv benv env -> TermBounds (UniformEnv benv env) t -> TermBounds (UniformEnv benv env) t +makeTransitives env = mapTupR (makeTransitive env) + +makeTransitive :: BoundsEnv benv env -> TermBound (UniformEnv benv env) t -> TermBound (UniformEnv benv env) t +makeTransitive env = makeTransitiveLower env . makeTransitiveUpper env + +-- | Make the lowerbound transitively closed +makeTransitiveLower :: BoundsEnv benv env -> TermBound (UniformEnv benv env) t -> TermBound (UniformEnv benv env) t +makeTransitiveLower env b = TermBound + (mapPartialEnv (\(Functor.Const d) -> InEdge $ Edge d) $ partialEnvTail dists) + (upper b) + where + BoundsEnv g _ _ _ = pushScalar env b + dists = Graph.shortestPathsTo distance g ZeroIdx + +-- | Make the upperbound transitively closed +makeTransitiveUpper :: BoundsEnv benv env -> TermBound (UniformEnv benv env) t -> TermBound (UniformEnv benv env) t +makeTransitiveUpper env b = TermBound + (lower b) + (mapPartialEnv (\(Functor.Const d) -> Edge d) $ partialEnvTail dists) + where + BoundsEnv g _ _ _ = pushScalar env b + dists = Graph.shortestPathsFrom distance g ZeroIdx + +assumeTrue + :: forall benv env. + BoundsEnv benv env + -> OpenExp env benv PrimBool + -> BoundsEnv benv env +assumeTrue env (PrimApp PrimLAnd (Pair e1 e2)) = assumeTrue (assumeTrue env e1) e2 +assumeTrue env (PrimApp PrimLNot expr) = assumeFalse env expr +assumeTrue env (PrimApp (PrimCmp tp c) (Pair e1 e2)) + | NumSingleType (IntegralNumType _) <- tp + , Just (Exists node1, dist1) <- extractArg e1 + , Just (Exists node2, dist2) <- extractArg e2 + = case c of + CmpLt -> + env{ boundsGraph = + Graph.insertEdgeWith min node1 node2 (Edge $ dist2 - dist1 - 1) $ boundsGraph env + } + + CmpGtEq -> + env{ boundsGraph = + -- Note that node2 and node1 are swapped here, as we encode greater-than as a less-than edge. + Graph.insertEdgeWith min node2 node1 (Edge $ dist1 - dist2) $ boundsGraph env + } + + CmpEq -> + env{ boundsGraph = + -- Add same edge in both directions + Graph.insertEdgeWith min node1 node2 (Edge $ dist2 - dist1) $ + Graph.insertEdgeWith min node2 node1 (Edge $ dist1 - dist2) $ + boundsGraph env + } + + CmpNEq -> env -- Cannot encode this + -- TODO: If based on the current bounds information, one operand is the + -- upper or lowerbound of the other operand, then we can convert not-equals + -- to less-than, and we can handle that in our analysis. + where + extractArg :: OpenExp env benv a -> Maybe (Exists (Idx (UniformEnv benv env)), Integer) + extractArg (Evar var) = Just (Exists $ expIdx env $ varIdx var, 0) + extractArg (ArrayInstr (Parameter var) _) = Just (Exists $ accIdx env $ varIdx var, 0) + extractArg (Const (SingleScalarType (NumSingleType (IntegralNumType t))) value) + | IntegralDict <- integralDict t = Just (Exists $ boundsZero env, fromIntegral value) + -- Note that expressions like `var + 1` are not handled. Since these + -- additions (and subtractions) may overflow, we cannot in general analyse + -- them safely. We could do a check here to see if we already know no + -- overflow will occur. Then we could track those scenarios. + -- Though I don't know how often this would trigger, and if it is thus + -- worth it to include this in the analysis. + extractArg _ = Nothing +assumeTrue env _ = env -- Cannot use the information of this expression + +assumeFalse :: BoundsEnv benv env -> OpenExp env benv PrimBool -> BoundsEnv benv env +assumeFalse env (PrimApp PrimLOr (Pair e1 e2)) = assumeFalse (assumeFalse env e1) e2 +assumeFalse env (PrimApp PrimLNot expr) = assumeTrue env expr +assumeFalse env (PrimApp (PrimCmp tp c) expr) + | NumSingleType (IntegralNumType _) <- tp + = assumeTrue env (PrimApp (PrimCmp tp (negateCmp c)) expr) +assumeFalse env _ = env -- Cannot use the information of this expression + +assumeTrue' :: env1 :?> env2 -> BoundsEnv benv env2 -> OpenExp env1 benv PrimBool -> BoundsEnv benv env2 +assumeTrue' k env (PrimApp PrimLAnd (Pair e1 e2)) = assumeTrue' k (assumeTrue' k env e1) e2 +assumeTrue' k env expr + | Just expr' <- strengthenE k expr = assumeTrue env expr' + | otherwise = env + +-- | Finds an lowerbound on the given term. +-- Only follows a direct edge between the node and zero. +lowerConst :: BoundsEnv benv env -> TermBound (UniformEnv benv env) t -> Maybe Integer +lowerConst (BoundsEnv _ _ zero _) b = (\(InEdge edge) -> negate $ distance edge) <$> prjPartial zero (lower b) + +-- | Finds an upperbound on the given term. +-- Only follows a direct edge between the node and zero. +upperConst :: BoundsEnv benv env -> TermBound (UniformEnv benv env) t -> Maybe Integer +upperConst (BoundsEnv _ _ zero _) b = distance <$> prjPartial zero (upper b) + +strengthenBound :: BoundsGraph (env, s) -> TermBound (env, s) t -> TermBound env t +strengthenBound graph bound = TermBound lower' upper' + where + lower' = case lower bound of + PPush edges (InEdge (Edge dist)) -> + unionPartialEnv (\(InEdge a) (InEdge b) -> InEdge $ min a b) edges + $ mapPartialEnv (\(InEdge (Edge d)) -> InEdge $ Edge $ d + dist) + $ partialEnvTail + $ Graph.inn graph ZeroIdx + PNone edges -> edges + PEnd -> PEnd + + upper' = case upper bound of + PPush edges (Edge dist) -> + unionPartialEnv min edges + $ mapPartialEnv (\(Edge d) -> Edge $ d + dist) + $ partialEnvTail + $ Graph.out graph ZeroIdx + PNone edges -> edges + PEnd -> PEnd + +strengthenBoundsWithScalarLHS + :: forall benv env env' r t s. + BoundsGraph (UniformEnv benv env') + -> LeftHandSide r t env env' + -> TermBounds (UniformEnv benv env') s + -> (BoundsGraph (UniformEnv benv env), TermBounds (UniformEnv benv env) s) +strengthenBoundsWithScalarLHS graph lhs bound = case lhs of + LeftHandSideWildcard _ -> (graph, bound) + LeftHandSideSingle _ -> (Graph.dropNode graph, mapTupR (strengthenBound graph) bound) + LeftHandSidePair l1 l2 + | (graph', bound') <- strengthenBoundsWithScalarLHS @benv graph l2 bound + -> strengthenBoundsWithScalarLHS @benv graph' l1 bound' + +strengthenBoundsWithBufferLHS + :: forall benv benv' r t s. + BoundsGraph (UniformEnv benv' ()) + -> LeftHandSide r t benv benv' + -> TermBounds (UniformEnv benv' ()) s + -> (BoundsGraph (UniformEnv benv ()), TermBounds (UniformEnv benv ()) s) +strengthenBoundsWithBufferLHS graph lhs bounds = case lhs of + LeftHandSideWildcard _ -> (graph, bounds) + -- Call boundsGraphClearNode before dropNode, to propagate information of the + -- first node to other nodes. See comment of boundsGraphClearNode + LeftHandSideSingle _ -> (Graph.dropNode (boundsGraphClearNode graph ZeroIdx), mapTupR (strengthenBound graph) bounds) + LeftHandSidePair l1 l2 + | (graph', bounds') <- strengthenBoundsWithBufferLHS graph l2 bounds + -> strengthenBoundsWithBufferLHS graph' l1 bounds' + +-- Removes the information of the given node in the graph. +-- Before removing it, it propagates the information on that node to other +-- nodes. For instance, if there is an edge from x to y, and an edge from y to +-- z, and y is removed, then this function adds an edge from x to z (with the +-- sum of the previous two edges). +-- +-- This function can be used when the contents of a buffer is overwritten, to +-- invalidate the existing information on that buffer. It can also be used to +-- strengthen an environment, as we then need to drop a variable (node) from +-- the environment (graph) but we still want to preserve the information it +-- implies on other variables (nodes). +boundsGraphClearNode + :: BoundsGraph env + -> Idx env t + -> BoundsGraph env +boundsGraphClearNode graph idx + = Graph.removeEdgesOf idx + $ Graph.insertCartesianEdgesWith + min + (\(InEdge (Edge a)) (Edge b) -> Edge $ a + b) + (Graph.inn graph idx) + (Graph.out graph idx) + graph + +boundsGraphClearNodes + :: BoundsGraph env + -> IdxSet env + -> BoundsGraph env +boundsGraphClearNodes graph set = + foldl' (\g (Exists idx) -> boundsGraphClearNode g idx) graph + $ IdxSet.toList set + +lookupAssertion + :: forall benv. + WEnv Binding benv + -> Exp benv PrimBool + -> Maybe (Idx benv PrimBool) +lookupAssertion wenv expr = listToMaybe $ mapMaybe f $ wenvToList wenv + where + f :: EnvBinding (Binding benv) benv -> Maybe (Idx benv PrimBool) + f (EnvBinding idx (BindAssertAssume expr')) + | Just _ <- matchOpenExp expr expr' = Just idx + f _ = Nothing diff --git a/src/Data/Array/Accelerate/Trafo/Operation/LiveVars.hs b/src/Data/Array/Accelerate/Trafo/Operation/LiveVars.hs new file mode 100644 index 000000000..93ecb1ebf --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Operation/LiveVars.hs @@ -0,0 +1,409 @@ +{-# LANGUAGE MultiWayIf #-} +{-# OPTIONS_HADDOCK hide #-} + +-- | +-- Module : Data.Array.Accelerate.Trafo.Operation.LiveVars +-- Copyright : [2012..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.Trafo.Operation.LiveVars ( + module Data.Array.Accelerate.Trafo.LiveVars, + + stronglyLiveVariables, stronglyLiveVariablesFun, + + SLVOperation(..), ShrinkOperation(..), ShrunkOperation(..), SubArgs(..), SubArg(..), + reEnvArrayInstr, + reEnvArrayDescriptors, + ShrinkArg(..), shrinkArgs, composeSubArgs, + + defaultSlvGenerate, defaultSlvMap, defaultSlvBackpermute, +) where + +import qualified Data.Array.Accelerate.AST.IdxSet as IdxSet +import Data.Array.Accelerate.AST.Var +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.AST.Operation +import Data.Array.Accelerate.AST.Environment +import Data.Array.Accelerate.Representation.Array +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Trafo.Exp.Substitution +import Data.Array.Accelerate.Trafo.LiveVars +import Data.Array.Accelerate.Error + +import Data.Type.Equality +import Data.Maybe (mapMaybe) + +stronglyLiveVariablesFun :: SLVOperation op => PreOpenAfun op () t -> PreOpenAfun op () t +stronglyLiveVariablesFun acc = acc' ReEnvEnd + where + LVAnalysisFun _ acc' = stronglyLiveVariablesFun' emptyLivenessEnv acc + +stronglyLiveVariables :: SLVOperation op => PreOpenAcc op () t -> PreOpenAcc op () t +stronglyLiveVariables acc = fromEither' $ acc' ReEnvEnd SubTupRkeep + where + LVAnalysis _ acc' = stronglyLiveVariables' emptyLivenessEnv returnImplicationsLive (mapTupR (const Shared) $ groundsR acc) acc + +stronglyLiveVariablesFun' :: SLVOperation op => LivenessEnv env -> PreOpenAfun op env t -> LVAnalysisFun (PreOpenAfun op) env t +stronglyLiveVariablesFun' liveness (Alam lhs f) + | liveness1 <- lEnvPushLHS lhs liveness + , LVAnalysisFun liveness2 f' <- stronglyLiveVariablesFun' liveness1 f + , (lhs', liveness3) <- lEnvStrengthenLHS lhs liveness2 + = LVAnalysisFun + liveness3 + $ \re -> if + | BindLiveness lhs'' re' <- bind lhs' re -> + Alam lhs'' $ f' re' +stronglyLiveVariablesFun' liveness (Abody body) + | LVAnalysis liveness1 body' <- stronglyLiveVariables' liveness returnImplicationsLive (mapTupR (const Shared) $ groundsR body) body + = LVAnalysisFun + liveness1 + $ \re -> Abody $ fromEither' $ body' re SubTupRkeep + +stronglyLiveVariablesFun'' :: SLVOperation op => LivenessEnv env -> Uniquenesses t -> PreOpenAfun op env (s -> t) -> LVAnalysisFun (PreOpenAfun op) env (s -> t) +stronglyLiveVariablesFun'' liveness us (Alam lhs (Abody body)) + | liveness1 <- lEnvPushLHS lhs liveness + , LVAnalysis liveness2 body' <- stronglyLiveVariables' liveness1 returnImplicationsLive us body + , (lhs', liveness3) <- lEnvStrengthenLHS lhs liveness2 + = LVAnalysisFun + liveness3 + $ \re -> if + | BindLiveness lhs'' re' <- bind lhs' re -> + Alam lhs'' $ Abody $ fromEither' $ body' re' SubTupRkeep +stronglyLiveVariablesFun'' _ _ _ = internalError "Expected unary function" + +fromEither' :: Either a a -> a +fromEither' (Left x) = x +fromEither' (Right x) = x + +stronglyLiveVariables' :: SLVOperation op => LivenessEnv env -> ReturnImplications env t -> Uniquenesses t -> PreOpenAcc op env t -> LVAnalysis (PreOpenAcc op) env t +stronglyLiveVariables' liveness returns us = \case + Exec op args + | Just (ShrinkOperation shrinkOp) <- slvOperation op + -- We can shrink this operation to output to part of its buffers. + , input <- IdxSet.fromList $ argsInputs args + , output <- IdxSet.fromList $ argsOutputs args + , liveness1 <- addLiveImplications output input liveness -> + LVAnalysis + liveness1 + $ \re s -> if + | Refl <- subTupUnit s + , allDead re output -> + Right $ Return TupRunit + + | Refl <- subTupUnit s + , ReEnvSubArgs subArgs args' <- reEnvSubArgs re args + , ShrunkOperation op' args'' <- shrinkOp subArgs args' args -> + Right $ Exec op' args'' + + -- We cannot shrink this operation to only output a part of its buffers. + -- Hence it's "all or nothing", if we use at least one of the output + -- buffers, then the entire operation is live. + | free <- IdxSet.fromList $ map (\(Exists (Var _ idx)) -> Exists idx) $ argsVars args + , output <- IdxSet.fromList $ argsOutputs args + , liveness1 <- addLiveImplications output free liveness -> + LVAnalysis + liveness1 + $ \re s -> if + | Refl <- subTupUnit s + , allDead re output -> + Right $ Return TupRunit -- All output vars are dead + + | Refl <- subTupUnit s + , args' <- reEnvArgs re args -> + Right $ Exec op args' -- Live + Return vars + | liveness1 <- returnVars returns vars liveness -> + LVAnalysis + liveness1 + $ \re s -> Right $ Return $ expectJust $ reEnvVars re $ subTupR s vars + Manifest var + | liveness1 <- returnVars returns (TupRsingle var) liveness -> + LVAnalysis + liveness1 + $ \re s -> case s of + SubTupRskip -> Right $ Return TupRunit + SubTupRkeep -> Right $ Manifest $ expectJust $ reEnvVar re var + Compute expr -> stronglyLiveVariablesComputeLike (const Compute) liveness returns expr + Aassert msg cond -> stronglyLiveVariablesComputeLike + (\s -> case s of + SubTupRkeep -> Aassert msg + SubTupRskip -> const $ Return TupRunit) -- Should not occur + liveness returns cond + Aassume cond -> stronglyLiveVariablesComputeLike + (\s -> case s of + SubTupRkeep -> Aassume + SubTupRskip -> const $ Return TupRunit) -- Should not occur + liveness returns cond + Alet lhs us' bnd body + | liveness1 <- lEnvPushLHS lhs liveness + , LVAnalysis liveness2 body' <- stronglyLiveVariables' liveness1 (returnImplicationsWeakenByLHS lhs returns) us body + , (lhs', liveness3, returns') <- lEnvStrengthenLHSReturn lhs liveness2 + , LVAnalysis liveness4 bnd' <- stronglyLiveVariables' liveness3 returns' us' bnd -> + LVAnalysis + liveness4 + $ \re s -> case bindSub lhs' re of + BindLivenessSub subTup' lhsFull lhsSub re' -> case (bnd' re subTup', body' re' s) of + (Left bnd'', Left body'') -> Left $ mkAlet lhsFull us' bnd'' body'' + (Left bnd'', Right body'') -> Right $ mkAlet lhsFull us' bnd'' body'' + (Right bnd'', Left body'') -> Left $ mkAlet lhsSub (subTupUniqueness subTup' us') bnd'' body'' + (Right bnd'', Right body'') -> Right $ mkAlet lhsSub (subTupUniqueness subTup' us') bnd'' body'' + Alloc shr tp sh + | free <- IdxSet.fromVars sh + , liveness1 <- returnIndices returns free liveness -> + LVAnalysis + liveness1 + $ \re s -> + case s of + SubTupRskip -> Right $ Return TupRunit + SubTupRkeep -> Right $ Alloc shr tp $ expectJust $ reEnvVars re sh + Use tp size buffer -> + LVAnalysis + liveness + $ \_ s -> + case s of + SubTupRskip -> Right $ Return TupRunit + SubTupRkeep -> Right $ Use tp size buffer + Unit var + | free <- IdxSet.singleton $ varIdx var + , liveness1 <- returnIndices returns free liveness -> + LVAnalysis + liveness1 + $ \re s -> + case s of + SubTupRskip -> Right $ Return TupRunit + SubTupRkeep -> Right $ Unit $ expectJust $ reEnvVar re var + Acond condition true false + | liveness1 <- setLive (varIdx condition) liveness + , LVAnalysis liveness2 true' <- stronglyLiveVariables' liveness1 returns us true + , LVAnalysis liveness3 false' <- stronglyLiveVariables' liveness2 returns us false -> + LVAnalysis + liveness3 + $ \re s -> + let condition' = expectJust $ reEnvVar re condition + in case (true' re s, false' re s) of + (Left true'', Left false'') -> Left $ mkAcond condition' true'' false'' + (Right true'', Right false'') -> Right $ mkAcond condition' true'' false'' + (Left true'', Right false'') + | SubTupRkeep <- s -> Left $ mkAcond condition' true'' false'' + | DeclareSubVars lhs _ vars <- declareSubVars (groundsR true) s + -> Right $ Acond condition' (Alet lhs us true'' $ Return $ vars weakenId) false'' + (Right true'', Left false'') + | SubTupRkeep <- s -> Left $ mkAcond condition' true'' false'' + | DeclareSubVars lhs _ vars <- declareSubVars (groundsR true) s + -> Right $ Acond condition' true'' (Alet lhs us false'' $ Return $ vars weakenId) + Awhile us' condition step initial + | liveness1 <- setVarsLive initial liveness + , LVAnalysisFun liveness2 condition' <- stronglyLiveVariablesFun'' liveness1 (TupRsingle Shared) condition + , LVAnalysisFun liveness3 step' <- stronglyLiveVariablesFun'' liveness2 us' step -> + LVAnalysis + liveness3 + $ \re _ -> + Left $ Awhile us' (condition' re) (step' re) $ expectJust $ reEnvVars re initial + Atrace msg t + | free <- arrayDescriptorsIdxSet t + , liveness1 <- returnIndices returns free liveness -> + LVAnalysis + liveness1 + $ \re s -> case s of + SubTupRkeep -> Right $ Atrace msg $ reEnvArrayDescriptors re t + SubTupRskip -> Right $ Return TupRunit + + Fence set next + | liveness1 <- setIdxSetLive set liveness + , LVAnalysis liveness2 next' <- stronglyLiveVariables' liveness1 returns us next -> + LVAnalysis + liveness2 + $ \re s -> + let set' = reEnvIdxSet re set + in case next' re s of + Left next'' -> Left $ Fence set' next'' + Right next'' -> Right $ Fence set' next'' + where + mkAcond :: ExpVar env' PrimBool -> PreOpenAcc op env' t' -> PreOpenAcc op env' t' -> PreOpenAcc op env' t' + mkAcond _ (Return TupRunit) (Return TupRunit) = Return TupRunit + mkAcond condition true false = Acond condition true false + + mkAlet :: GLeftHandSide bnd subenv subenv' -> Uniquenesses bnd -> PreOpenAcc op subenv bnd -> PreOpenAcc op subenv' t -> PreOpenAcc op subenv t + mkAlet (LeftHandSideWildcard TupRunit) _ (Return TupRunit) body = body + mkAlet lhs us' bnd body = Alet lhs us' bnd body + +-- Implementation of stronglyLiveVariables' for Compute, Aassert and Aassume +stronglyLiveVariablesComputeLike + :: SLVOperation op + => (forall env' t'. SubTupR t t' -> Exp env' t' -> OperationAcc op env' t') + -> LivenessEnv env + -> ReturnImplications env t + -> Exp env t + -> LVAnalysis (PreOpenAcc op) env t +stronglyLiveVariablesComputeLike constructor liveness returns expr + -- If the LHS of the binding is live, then all free variables of this + -- expression are live as well. + | free <- expGroundVars expr + , liveness1 <- returnIndices returns (IdxSet.fromVarList free) liveness = + LVAnalysis + liveness1 + $ \re s -> + let + tp = expType expr + expr' = mapArrayInstr (reEnvArrayInstr re) expr + in case s of + SubTupRskip -> Right $ Return TupRunit + SubTupRkeep -> Right $ constructor s $ expr' + _ | DeclareSubVars lhs _ vars <- declareSubVars tp s + -> Right $ constructor s $ Let lhs expr' $ returnExpVars $ vars weakenId + +reEnvIdxSet :: ReEnv env subenv -> IdxSet.IdxSet env -> IdxSet.IdxSet subenv +reEnvIdxSet re idxSet = + IdxSet.fromList $ + mapMaybe (\(Exists ix) -> fmap Exists (reEnvIdx re ix)) + (IdxSet.toList idxSet) + +class SLVOperation op where + slvOperation :: op f -> Maybe (ShrinkOperation op f) + +newtype ShrinkOperation op f = ShrinkOperation (forall f' env' env. SubArgs f f' -> Args env' f' -> Args env f -> ShrunkOperation op env') + +-- existential over f: otherwise, you couldn't change the non-array arguments. You need this e.g. for a Generate: smaller array means smaller function. +data ShrunkOperation op env where + ShrunkOperation :: op f -> Args env f -> ShrunkOperation op env + +data SubArgs f f' where + SubArgsNil :: SubArgs () () + + -- This Out argument is dead. + -- Note that implementers of 'slvOperation' may assume that at least one Out + -- or Mut argument is preserved. + SubArgsDead :: SubArgs t t' + -> SubArgs (Out sh e -> t) (Var' sh -> t') + + SubArgsLive :: SubArg s s' + -> SubArgs t t' + -> SubArgs (s -> t) (s' -> t') + +infixr 9 `SubArgsLive` + +data SubArg t t' where + SubArgKeep :: SubArg t t + + SubArgOut :: SubTupR e e' + -> SubArg (Out sh e) (Out sh e') + +class ShrinkArg arg where + shrinkArg :: SubArg t t' -> arg t -> arg t' + deadArg :: arg (Out sh e) -> arg (Var' sh) + +-- instance ShrinkArg (Arg env) where +-- shrinkArg = _ +-- deadArg = _ + +shrinkArgs :: ShrinkArg arg => SubArgs f f' -> PreArgs arg f -> PreArgs arg f' +shrinkArgs SubArgsNil ArgsNil = ArgsNil +shrinkArgs (SubArgsDead sargs) (a:>:args) = deadArg a :>: shrinkArgs sargs args +shrinkArgs (SubArgsLive sarg sargs) (a:>:args) = shrinkArg sarg a :>: shrinkArgs sargs args + +defaultSlvGenerate + :: (forall sh' t'. op (Fun' (sh' -> t') -> Out sh' t' -> ())) + -> Maybe (ShrinkOperation op (Fun' (sh -> t) -> Out sh t -> ())) +defaultSlvGenerate mkGenerate = Just $ ShrinkOperation $ \subArgs args _ -> case subArgs of + SubArgKeep `SubArgsLive` SubArgKeep `SubArgsLive` SubArgsNil + -> ShrunkOperation mkGenerate args + SubArgKeep `SubArgsLive` SubArgOut subTp `SubArgsLive` SubArgsNil + | ArgFun f :>: array :>: ArgsNil <- args + -> ShrunkOperation mkGenerate (ArgFun (subTupFun subTp f) :>: array :>: ArgsNil) + _ `SubArgsLive` SubArgsDead _ -> internalError "At least one output should be preserved" + +defaultSlvMap + :: (forall sh' s' t'. op (Fun' (s' -> t') -> In sh' s' -> Out sh' t' -> ())) + -> Maybe (ShrinkOperation op (Fun' (s -> t) -> In sh s -> Out sh t -> ())) +defaultSlvMap mkMap = Just $ ShrinkOperation $ \subArgs args _ -> case subArgs of + SubArgKeep `SubArgsLive` SubArgKeep `SubArgsLive` SubArgKeep `SubArgsLive` SubArgsNil + -> ShrunkOperation mkMap args + SubArgKeep `SubArgsLive` SubArgKeep `SubArgsLive` SubArgOut subTp `SubArgsLive` SubArgsNil + | ArgFun f :>: input :>: output :>: ArgsNil <- args + -> ShrunkOperation mkMap (ArgFun (subTupFun subTp f) :>: input :>: output :>: ArgsNil) + _ `SubArgsLive` _ `SubArgsLive` SubArgsDead _ -> internalError "At least one output should be preserved" + +defaultSlvBackpermute + :: (forall sh1' sh2' t'. op (Fun' (sh2' -> sh1') -> In sh1' t' -> Out sh2' t' -> ())) + -> Maybe (ShrinkOperation op (Fun' (sh2 -> sh1) -> In sh1 t -> Out sh2 t -> ())) +defaultSlvBackpermute mkBackpermute = Just $ ShrinkOperation $ \subArgs args _ -> case subArgs of + SubArgKeep `SubArgsLive` SubArgKeep `SubArgsLive` SubArgKeep `SubArgsLive` SubArgsNil + -> ShrunkOperation mkBackpermute args + SubArgKeep `SubArgsLive` SubArgKeep `SubArgsLive` SubArgOut s `SubArgsLive` SubArgsNil + | f :>: ArgArray In (ArrayR shr r) sh buf :>: output :>: ArgsNil <- args + -> ShrunkOperation mkBackpermute (f :>: ArgArray In (ArrayR shr (subTupR s r)) sh (subTupDBuf s buf) :>: output :>: ArgsNil) + _ `SubArgsLive` _ `SubArgsLive` SubArgsDead _ -> internalError "At least one output should be preserved" + +reEnvArrayInstr :: ReEnv env subenv -> ArrayInstr env t -> ArrayInstr subenv t +reEnvArrayInstr re (Parameter var) = Parameter $ expectJust $ reEnvVar re var +reEnvArrayInstr re (Index buffer) = Index $ expectJust $ reEnvVar re buffer + +reEnvArgs :: ReEnv env subenv -> Args env t -> Args subenv t +reEnvArgs re (a :>: as) = reEnvArg re a :>: reEnvArgs re as +reEnvArgs _ ArgsNil = ArgsNil + +reEnvArg :: ReEnv env subenv -> Arg env t -> Arg subenv t +reEnvArg re (ArgVar vars) = ArgVar $ expectJust $ reEnvVars re vars +reEnvArg re (ArgExp expr) = ArgExp $ mapArrayInstr (reEnvArrayInstr re) expr +reEnvArg re (ArgFun f) = ArgFun $ mapArrayInstrFun (reEnvArrayInstr re) f +reEnvArg re (ArgArray m repr sh buffers) = ArgArray m repr (expectJust $ reEnvVars re sh) (expectJust $ reEnvVars re buffers) + +reEnvArrayDescriptors :: ReEnv env subenv -> ArrayDescriptors env t -> TupR (ArrayDescriptor subenv) t +reEnvArrayDescriptors re = mapTupR (\(ArrayDescriptor shr sh buffer) -> ArrayDescriptor shr (expectJust $ reEnvVars re sh) (expectJust $ reEnvVars re buffer)) + +-- Captures existential f' +data ReEnvSubArgs subenv f where + ReEnvSubArgs :: SubArgs f f' + -> Args subenv f' + -> ReEnvSubArgs subenv f + +reEnvSubArgs :: ReEnv env subenv -> Args env f -> ReEnvSubArgs subenv f +reEnvSubArgs re (a :>: as) + | ReEnvSubArgs subs as' <- reEnvSubArgs re as = + case a of + ArgArray Out (ArrayR shr tp) sh buffers -> case reEnvSubBuffers re tp buffers of + ReEnvSubBuffers SubTupRskip _ + -- don't kill unit outputs + | TupRunit <- tp -> ReEnvSubArgs (SubArgsLive SubArgKeep subs) (ArgArray Out (ArrayR shr tp) (expectJust $ reEnvVars re sh) TupRunit :>: as') + | otherwise -> ReEnvSubArgs (SubArgsDead subs) (ArgVar (fromGrounds $ expectJust $ reEnvVars re sh) :>: as') + ReEnvSubBuffers SubTupRkeep buffers' -> ReEnvSubArgs (SubArgsLive SubArgKeep subs) (ArgArray Out (ArrayR shr tp) (expectJust $ reEnvVars re sh) buffers' :>: as') + ReEnvSubBuffers sub buffers' -> ReEnvSubArgs (SubArgsLive (SubArgOut sub) subs) (ArgArray Out (ArrayR shr $ subTupR sub tp) (expectJust $ reEnvVars re sh) buffers' :>: as') + _ -> ReEnvSubArgs (SubArgsLive SubArgKeep subs) (reEnvArg re a :>: as') +reEnvSubArgs _ ArgsNil = ReEnvSubArgs SubArgsNil ArgsNil + +reEnvSubBuffers :: ReEnv env subenv -> TypeR t -> GroundVars env (Buffers t) -> ReEnvSubBuffers subenv t +reEnvSubBuffers _ TupRunit TupRunit = ReEnvSubBuffers SubTupRskip TupRunit +reEnvSubBuffers re (TupRsingle _) (TupRsingle var) + | Just var' <- reEnvVar re var = ReEnvSubBuffers SubTupRkeep (TupRsingle var') + | otherwise = ReEnvSubBuffers SubTupRskip TupRunit +reEnvSubBuffers re (TupRpair t1 t2) (TupRpair v1 v2) + | ReEnvSubBuffers s1 v1' <- reEnvSubBuffers re t1 v1 + , ReEnvSubBuffers s2 v2' <- reEnvSubBuffers re t2 v2 + = case (s1, s2) of + (SubTupRskip, SubTupRskip) -> ReEnvSubBuffers SubTupRskip TupRunit + (SubTupRkeep, SubTupRkeep) -> ReEnvSubBuffers SubTupRkeep (TupRpair v1' v2') + _ -> ReEnvSubBuffers (SubTupRpair s1 s2) (TupRpair v1' v2') +reEnvSubBuffers _ _ _ = internalError "Tuple mismatch" + +data ReEnvSubBuffers subenv t where + ReEnvSubBuffers :: SubTupR t t' -> GroundVars subenv (Buffers t') -> ReEnvSubBuffers subenv t + +composeSubArgs :: SubArgs a b -> SubArgs b c -> SubArgs a c +composeSubArgs SubArgsNil SubArgsNil = SubArgsNil +composeSubArgs (SubArgsDead s1) (SubArgsLive SubArgKeep s2) = SubArgsDead $ composeSubArgs s1 s2 +composeSubArgs (SubArgsLive SubArgKeep s1) (SubArgsDead s2) = SubArgsDead $ composeSubArgs s1 s2 +composeSubArgs (SubArgsLive (SubArgOut _) s1) (SubArgsDead s2) = SubArgsDead $ composeSubArgs s1 s2 +composeSubArgs (SubArgsLive SubArgKeep s1) (SubArgsLive s s2) = SubArgsLive s $ composeSubArgs s1 s2 +composeSubArgs (SubArgsLive (SubArgOut t) s1) (SubArgsLive SubArgKeep s2) = SubArgsLive (SubArgOut t) $ composeSubArgs s1 s2 +composeSubArgs (SubArgsLive (SubArgOut t1) s1) (SubArgsLive (SubArgOut t2) s2) = SubArgsLive (SubArgOut $ composeSubTupR t2 t1) $ composeSubArgs s1 s2 + +subTupUniqueness :: SubTupR t t' -> Uniquenesses t -> Uniquenesses t' +subTupUniqueness SubTupRskip _ = TupRunit +subTupUniqueness SubTupRkeep t = t +subTupUniqueness (SubTupRpair s1 s2) (TupRpair t1 t2) = subTupUniqueness s1 t1 `TupRpair` subTupUniqueness s2 t2 +subTupUniqueness (SubTupRpair s1 s2) (TupRsingle Shared) = TupRpair (subTupUniqueness s1 (TupRsingle Shared)) (subTupUniqueness s2 (TupRsingle Shared)) --TupRsingle Shared diff --git a/src/Data/Array/Accelerate/Trafo/Operation/Simplify.hs b/src/Data/Array/Accelerate/Trafo/Operation/Simplify.hs new file mode 100644 index 000000000..d852af180 --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Operation/Simplify.hs @@ -0,0 +1,735 @@ +{-# OPTIONS_HADDOCK hide #-} +-- | +-- Module : Data.Array.Accelerate.Trafo.Operation.Simplify +-- Copyright : [2012..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.Trafo.Operation.Simplify ( + simplify, simplifyFun, SimplifyOperation(..), CopyOperation(..), isNoOp, + copyOperationsForArray, detectMapCopies, detectBackpermuteCopies +) where + +import Data.Array.Accelerate.AST.Environment +import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.IdxSet ( IdxSet ) +import qualified Data.Array.Accelerate.AST.IdxSet as IdxSet +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.AST.Operation +import Data.Array.Accelerate.Error +import Data.Array.Accelerate.Array.Buffer +import Data.Array.Accelerate.Type +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Representation.Array +import Data.Array.Accelerate.Analysis.Match +import qualified Data.Array.Accelerate.Trafo.Exp.Simplify as Exp +import Data.Array.Accelerate.Trafo.Exp.Substitution +import Data.Array.Accelerate.Trafo.Var +import Data.Array.Accelerate.Trafo.Substitution hiding ( weakenArrayInstr ) +import Data.Array.Accelerate.Trafo.WeakenedEnvironment +import Data.Array.Accelerate.Trafo.Operation.Substitution +import Data.Array.Accelerate.Trafo.LiveVars ( SubTupR(..), subTupR, subTupRpair, subTupPreserves ) +import Data.Maybe ( mapMaybe ) +import Data.List ( foldl' ) +import Data.Text ( Text ) +import Control.Monad +import Data.Functor.Identity + +class SimplifyOperation op where + detectCopy :: op f -> Args env f -> [CopyOperation env] + detectCopy _ _ = [] + +data CopyOperation env where + CopyOperation + :: Idx env (Buffer t) -- input + -> Idx env (Buffer t) -- output + -> CopyOperation env + +isNoOp :: SimplifyOperation op => op f -> Args env f -> Bool +-- An operation is a no-op if all outputs are copies of themselves. +-- For instance, 'map id xs xs' is a no-op. +-- This pattern is created by in-place updates, on the copy of the defaults +-- array of a permute. +isNoOp op args = all (\(Exists idx) -> idx `IdxSet.member` copied) $ argsOutputs args + where + copied = IdxSet.fromList + $ map (\(CopyOperation idx _) -> Exists idx) + $ filter (\(CopyOperation i o) -> i == o) + $ detectCopy op args + +copyOperationsForArray :: forall env sh sh' t. Arg env (In sh t) -> Arg env (Out sh' t) -> [CopyOperation env] +copyOperationsForArray (ArgArray _ (ArrayR _ tp) _ input) (ArgArray _ _ _ output) = go tp input output [] + where + go :: forall s. TypeR s -> GroundVars env (Buffers s) -> GroundVars env (Buffers s) -> [CopyOperation env] -> [CopyOperation env] + go (TupRpair t1 t2) (TupRpair i1 i2) (TupRpair o1 o2) = go t1 i1 o1 . go t2 i2 o2 + go (TupRsingle t) (TupRsingle (Var _ input')) (TupRsingle (Var _ output')) + | Refl <- reprIsSingle @ScalarType @s @Buffer t = (CopyOperation input' output' :) + go _ _ _ = id + +detectMapCopies :: forall genv sh t s. Args genv (Fun' (t -> s) -> In sh t -> Out sh s -> ()) -> [CopyOperation genv] +detectMapCopies (ArgFun (Lam lhs (Body body)) :>: ArgArray _ _ _ input :>: ArgArray _ _ _ output :>: ArgsNil) + = detectMapCopies' lhs body input output +detectMapCopies (ArgFun (Body f) :>: _ ) = functionImpossible $ expType f +detectMapCopies (ArgFun (Lam _ (Lam _ _)) :>: _ :>: ArgArray _ (ArrayR _ tp) _ _ :>: ArgsNil) + = functionImpossible tp + +detectMapCopies' :: forall genv env t s. ELeftHandSide t () env -> OpenExp env genv s -> GroundVars genv (Buffers t) -> GroundVars genv (Buffers s) -> [CopyOperation genv] +detectMapCopies' lhs body input output = go Just body output [] + where + go :: forall env' s'. env' :?> env -> OpenExp env' genv s' -> GroundVars genv (Buffers s') -> [CopyOperation genv] -> [CopyOperation genv] + go k (Assume _ e2) o = go k e2 o + go k (Pair e1 e2) (TupRpair o1 o2) = go k e1 o1 . go k e2 o2 + go k (Let lhs' _ expr) output' = go (strengthenWithLHS lhs' >=> k) expr output' + go k (Evar (Var tp idx)) (TupRsingle (Var _ output')) + | Just idx' <- k idx -- Check if index is bound by the function (opposed to local binding) + , Refl <- reprIsSingle @ScalarType @s' @Buffer tp + = (CopyOperation (findInput idx') output' :) + go _ _ _ = id + + findInput :: Idx env t' -> Idx genv (Buffer t') + findInput idx = case findInput' lhs input idx of + Right buffer -> buffer + Left idx' -> case idx' of {} + + findInput' :: forall u env1 env2 t'. ELeftHandSide u env1 env2 -> GroundVars genv (Buffers u) -> Idx env2 t' -> Either (Idx env1 t') (Idx genv (Buffer t')) + findInput' (LeftHandSideWildcard _) _ idx = Left idx + findInput' (LeftHandSideSingle tp) (TupRsingle (Var _ buffer)) idx = case idx of + SuccIdx idx' -> Left idx' + ZeroIdx + | Refl <- reprIsSingle @ScalarType @u @Buffer tp -> Right buffer + findInput' (LeftHandSidePair l1 l2) (TupRpair in1 in2) idx = case findInput' l2 in2 idx of + Left idx' -> findInput' l1 in1 idx' + Right buffer -> Right buffer + findInput' _ _ _ = internalError "Tuple mismatch" + +detectBackpermuteCopies :: forall env sh sh' t. Args env (Fun' (sh' -> sh) -> In sh t -> Out sh' t -> ()) -> [CopyOperation env] +detectBackpermuteCopies (ArgFun f :>: input@(ArgArray _ _ sh _) :>: output@(ArgArray _ _ sh' _) :>: ArgsNil) + | Just Refl <- matchVars sh sh' + , Just Refl <- isIdentity f = copyOperationsForArray input output +detectBackpermuteCopies _ = [] + +-- TODO for Fences: Add a way to store that one variable implies that another +-- variable is also resolved. Eg in 'let e2 = fence { e1 } assert e0 >= 0', we +-- know that after a fence on e2, e1 is also resolved. +data Info env t where + -- | This variable has a known value + InfoConst :: IdxSet env -- The set of variables one should synchronise with in a Fence + -> ScalarType t + -> t + -> Info env t -- A constant scalar + -- | This variable is alias of another variable + InfoAlias :: IdxSet env -- The set of variables one should synchronise with in a Fence + -> Idx env t + -> Info env t + -- | This is a buffer with undefined content (e.g. directly after an Alloc) + InfoUndef :: Info env (Buffer t) + -- | Information on a buffer + InfoBuffer :: Maybe (Idx env t) -- In case of a Unit, the index of the scalar variable that it contains. + -- Copy of another buffer. This is similar to an alias, but a buffer may only + -- be a copy of another buffer temporarily. A write to the original or copy + -- causes that they aren't copies any more. Hence we also keep track in + -- InfoBuffer of the buffers it was copied to. + -> Maybe (Idx env (Buffer t)) + -> [Idx env (Buffer t)] -- List of buffers where this buffer is copied to + -> Info env (Buffer t) + -- | This variable is resolved, as enforced by a Fence + InfoResolved :: Info env t + -- | No information available + InfoNone :: Info env t + +newtype InfoEnv env = InfoEnv { unInfoEnv :: WEnv Info env } + +emptySimplifyEnv :: InfoEnv () +emptySimplifyEnv = InfoEnv wempty + +instance Sink Info where + weaken k (InfoAlias set idx) = InfoAlias (IdxSet.map (weaken k) set) $ weaken k idx + weaken k (InfoConst set t c) = InfoConst (IdxSet.map (weaken k) set) t c + weaken k (InfoBuffer unitScalar copyOf copied) = InfoBuffer (fmap (weaken k) unitScalar) (fmap (weaken k) copyOf) (fmap (weaken k) copied) + weaken _ InfoUndef = InfoUndef + weaken _ InfoResolved = InfoResolved + weaken _ InfoNone = InfoNone + +infoFor :: Idx env t -> InfoEnv env -> Info env t +infoFor ix (InfoEnv env) = wprj ix env + +-- Substitutions for aliasing. +-- This substitution only assures that we use the original declaration instead +-- of the alias. It does not remove the definition of the alias, a later pass +-- should remove that (with a (strongly) live variable analysis). +-- +-- These substitutions might only be sound after synchronising with certain +-- variables. Function 'syncSubstitute' returns the set of these variables. +-- This is needed to ensure that the right assertions are evaluated before +-- using the substituted variable. This occurs for instance in this program: +-- +-- let x = { fence c; return y } +-- +-- Here we may substitute x with y, if we put a fence on c. +-- +substitute :: InfoEnv env -> env :> env +substitute env = Weaken $ \idx -> case infoFor idx env of + InfoAlias _ idx' -> idx' + InfoBuffer _ (Just idx') _ -> idx' + _ -> idx + +substituteUnlessResolved :: InfoEnv env -> Exists (Idx env) -> Maybe (Exists (Idx env)) +substituteUnlessResolved env (Exists idx) = case infoFor idx env of + InfoConst _ _ _ -> Nothing + InfoResolved -> Nothing + InfoAlias _ idx' -> Just $ Exists idx' + InfoBuffer _ (Just idx') _ -> Just $ Exists idx' + _ -> Just $ Exists idx + +-- When reading from an array, we can read from another buffer with the same +-- contents. The index of such copy is stored in InfoBuffer. When writing we +-- cannot do that. Note that after a write, the information about the copy +-- is removed in 'invalidate'. +-- +substituteOutput :: InfoEnv env -> env :> env +substituteOutput env = Weaken $ \idx -> case infoFor idx env of + InfoAlias _ idx' -> idx' + _ -> idx + +-- The variables one should synchronise with, for 'substitute' to be sound. +-- These variables should be passed to Fence. +-- +syncSubstitute :: InfoEnv env -> Idx env t -> IdxSet env +syncSubstitute env idx = case infoFor idx env of + InfoConst set _ _ -> set + InfoAlias set _ -> set + _ -> IdxSet.empty + +syncSubstitutes :: InfoEnv env -> IdxSet env -> IdxSet env +syncSubstitutes env vars = + IdxSet.unions + $ map + (\case + Exists (InfoAlias set _) -> set + _ -> IdxSet.empty + ) + $ wprjSet vars (unInfoEnv env) + +simplifyFun :: SimplifyOperation op => OperationAfun op () t -> OperationAfun op () t +simplifyFun fun = snd (simplifyFun' fun) emptySimplifyEnv + +-- Returns a set of array variables which may have been written to, and a the simplified function (given an InfoEnv) +simplifyFun' :: SimplifyOperation op => OperationAfun op env t -> (IdxSet env, InfoEnv env -> OperationAfun op env t) +simplifyFun' (Alam lhs f) = + (IdxSet.drop' lhs set, \env -> Alam lhs $ f' $ bindEnv lhs env) + where + (set, f') = simplifyFun' f +simplifyFun' (Abody body) = (set, \env -> Abody $ body' env) + where + (set, body') = simplify' (TupRsingle Shared) body + +simplifyFunWithUniqueness :: SimplifyOperation op => Uniquenesses t -> OperationAfun op env (t -> t) -> (IdxSet env, InfoEnv env -> OperationAfun op env (t -> t)) +simplifyFunWithUniqueness uniquenesses (Alam lhs (Abody body)) = + ( IdxSet.drop' lhs set + , \env -> Alam lhs $ Abody $ body' $ bindEnv lhs env + ) + where + (set, body') = simplify' uniquenesses body +simplifyFunWithUniqueness _ (Abody body) = groundFunctionImpossible $ groundsR body +simplifyFunWithUniqueness _ (Alam lhs (Alam _ _)) = groundFunctionImpossible $ lhsToTupR lhs + +simplify :: SimplifyOperation op => OperationAcc op () t -> OperationAcc op () t +simplify acc = snd (simplify' (TupRsingle Shared) acc) emptySimplifyEnv + +-- Returns the simplified program and a set of array variables which may have been written to +simplify' :: SimplifyOperation op => Uniquenesses t -> OperationAcc op env t -> (IdxSet env, InfoEnv env -> OperationAcc op env t) +simplify' uniquenesses = \case + Exec op args + | output <- outputArrays args -> + ( outputArrays args + , \env -> + let + fenceSet = syncSubstitutes env $ IdxSet.fromVarList $ argsVars args + args' = mapArgs (simplifyArg env) args + in + if isUndefCopy env output $ detectCopy op args then + Return TupRunit + else + fence fenceSet $ Exec op args') + Return vars -> + -- Note that we do not need to check for writes to variables here. + -- This construct may cause aliassing of variables, but an aliassed + -- variable cannot be unique and thus we do not need to signal the + -- original variable as mutated if it's returned. + ( variableIndices uniquenesses vars + , \env -> + fence (syncSubstitutes env $ IdxSet.fromVars vars) + $ Return $ simplifyReturnVars env uniquenesses vars + ) + Manifest var -> ( IdxSet.empty, const $ Manifest var ) + Compute expr -> + ( IdxSet.empty + , \env -> + fence (syncSubstitutes env $ IdxSet.fromVarList $ expGroundVars expr) + $ compute $ simplifyExp env expr + ) + Alet lhs us bnd body -> + let + (setBnd, bnd') = simplify' us bnd + (setBody, body') = simplify' uniquenesses body + in + ( setBnd `IdxSet.union` IdxSet.drop' lhs setBody + , \env -> + let + bnd'' = bnd' env + env' = bindingEnv setBnd IdxSet.empty lhs bnd'' (invalidate setBnd env) + in alet' lhs us bnd'' (body' env') + ) + Alloc shr tp sh -> + ( IdxSet.empty + , \env -> + fence (syncSubstitutes env $ IdxSet.fromVars sh) + $ Alloc shr tp $ mapTupR (weaken $ substitute env) sh + ) + Use tp 1 buffer -> + ( IdxSet.empty + , const + $ Alet (LeftHandSideSingle (GroundRscalar tp)) (TupRsingle Shared) (Compute $ Const tp $ indexBuffer tp buffer 0) + $ Unit (Var tp ZeroIdx) + ) + Use tp n buffer -> (IdxSet.empty, const $ Use tp n buffer) + Unit var -> + ( IdxSet.empty + , \env -> + fence (syncSubstitute env $ varIdx var) + $ Unit $ weaken (substitute env) var + ) + Acond cond true false -> + let + (setT, true') = simplify' uniquenesses true + (setF, false') = simplify' uniquenesses false + set = IdxSet.union setT setF + in + ( set + , \env -> case infoFor (varIdx cond) env of + InfoConst d _ 0 -> fence d $ false' env + InfoConst d _ _ -> fence d $ true' env + InfoAlias d ix -> fence d $ Acond (cond{varIdx = ix}) (true' env) (false' env) + _ -> Acond cond (true' env) (false' env) + ) + + Awhile us cond step initial -> + let + (setC, cond') = simplifyFun' cond + (setS, step') = simplifyFunWithUniqueness us step + set = setC `IdxSet.union` setS `IdxSet.union` variableIndices us initial + in + ( set + , \env -> + let + env' = invalidate set env + in + fence (syncSubstitutes env $ IdxSet.fromVars initial) + $ awhileSimplifyInvariant us (cond' env') (step' env') $ simplifyReturnVars env us initial + ) + + Atrace msg t -> + let + set = arrayDescriptorsIdxSet t + in + ( set + , \env -> + fence (syncSubstitutes env set) + $ Atrace msg $ simplifyArrayDescriptor env t + ) + + Aassert msg cond -> + ( IdxSet.empty + , \env -> + fence (syncSubstitutes env $ IdxSet.fromVarList $ expGroundVars cond) + $ assert msg $ simplifyExp env cond + ) + + Aassume cond -> + ( IdxSet.empty + , \env -> + fence (syncSubstitutes env $ IdxSet.fromVarList $ expGroundVars cond) + $ assume $ simplifyExp env cond) + +-- Mark all deps as evaluated in the environment +-- Remove deps from the list if they were already evaluated. + Fence deps next -> + let + (set, next') = simplify' uniquenesses next + in + ( set + , \env -> + let + deps' = IdxSet.union (syncSubstitutes env deps) $ IdxSet.fromList $ mapMaybe (substituteUnlessResolved env) $ IdxSet.toList deps + env' = InfoEnv $ wupdateSetWeakened + (\_ -> \case + InfoNone -> InfoResolved + info -> info + ) + deps' + $ unInfoEnv env + in + fence + deps' + (next' env') + ) + +simplifyArrayDescriptor :: InfoEnv env -> ArrayDescriptors env t -> ArrayDescriptors env t +simplifyArrayDescriptor env = mapTupR (\(ArrayDescriptor shr sh buffers) -> ArrayDescriptor shr (mapTupR (weaken $ substitute env) sh) (mapTupR (weaken $ substitute env) buffers)) + +-- Given an environment, the set of updated variables and a list of copies of +-- an operation, checks whether the operation copies all its outputs from +-- undefined buffers. +-- +-- This is specifically needed for permute, as it is common to use +-- `generate .. (const undef)` as defaults array. Permute introduces a map to +-- copy the defaults array and make it unique. The generate is already removed +-- in an earlier pass, and the map will be removed here. +isUndefCopy :: InfoEnv env -> IdxSet env -> [CopyOperation env] -> Bool +isUndefCopy env outputs copies + = outputs == IdxSet.fromList (map (\(CopyOperation _ o) -> Exists o) copies) + && all (\(CopyOperation i _) -> isUndef $ infoFor i env) copies + where + isUndef InfoUndef = True + isUndef _ = False + +variableIndices :: Uniquenesses t -> GroundVars env t -> IdxSet env +variableIndices (TupRsingle Unique) (TupRsingle var) = IdxSet.singleton $ varIdx var +variableIndices (TupRpair u1 u2) (TupRpair v1 v2) = variableIndices u1 v1 `IdxSet.union` variableIndices u2 v2 +variableIndices _ _ = IdxSet.empty + +simplifyReturnVars :: InfoEnv env -> Uniquenesses t -> GroundVars env t -> GroundVars env t +simplifyReturnVars env (TupRpair u1 u2) (TupRpair v1 v2) = + simplifyReturnVars env u1 v1 `TupRpair` simplifyReturnVars env u2 v2 +simplifyReturnVars env (TupRsingle Shared) v = mapTupR (weaken $ substitute env) v +simplifyReturnVars env (TupRsingle Unique) v = mapTupR (weaken $ substituteOutput env) v +simplifyReturnVars _ TupRunit _ = TupRunit +simplifyReturnVars _ _ _ = internalError "Tuple mismatch" + +findConstant :: forall env env' t. WEnv' Info env' env -> ScalarType t -> t -> Maybe (Idx env t) +findConstant env tp1 value1 = go env + where + go :: WEnv' Info env2 env1 -> Maybe (Idx env1 t) + go WEmpty = Nothing + go (WPushA _e (InfoConst d tp2 value2)) + | IdxSet.null d + , Just idx <- tryMatch tp2 value2 = Just idx + go (WPushB _e (InfoConst d tp2 value2)) + | IdxSet.null d + , Just idx <- tryMatch tp2 value2 = Just idx + go (WPushA e _) = SuccIdx <$> go e + go (WPushB e _) = SuccIdx <$> go e + go (WWeaken _ e) = go e + + tryMatch :: ScalarType s -> s -> Maybe (Idx (env3, s) t) + tryMatch tp2 value2 + | Just Refl <- matchScalarType tp1 tp2 + = case tp1 of + SingleScalarType t + | SingleDict <- singleDict t -- Gives 'Eq t' + , value1 == value2 -> Just ZeroIdx + VectorScalarType (VectorType _ t) + | SingleDict <- singleDict t + , value1 == value2 -> Just ZeroIdx + _ -> Nothing + | otherwise = Nothing + +bindingEnv :: forall op t env env'. SimplifyOperation op => IdxSet env -> IdxSet env -> GLeftHandSide t env env' -> OperationAcc op env t -> InfoEnv env -> InfoEnv env' +bindingEnv _ fenceSet lhs (Compute expr) (InfoEnv environment) = InfoEnv $ go weakenId lhs expr environment + where + go :: env :> env1 -> GLeftHandSide s env1 env2 -> Exp env s -> WEnv' Info env1 env1 -> WEnv' Info env2 env2 + go k (LeftHandSideSingle _) e env + | ArrayInstr (Parameter var) _ <- e = wpush env $ weaken (weakenSucc' k) $ InfoAlias fenceSet $ varIdx var + | Const tp c <- e = case findConstant env tp c of + Just idx -> wpush env $ InfoAlias (IdxSet.map (weaken (weakenSucc' k)) fenceSet) $ SuccIdx idx + Nothing -> wpush env $ InfoConst (IdxSet.map (weaken (weakenSucc' k)) fenceSet) tp c + | otherwise = wpush env InfoNone + + go k (LeftHandSidePair l1 l2) (Pair e1 e2) env + = go (weakenWithLHS l1 .> k) l2 e2 $ go k l1 e1 env + + go _k (LeftHandSideWildcard _) _ env = env + + go _ l _ env = goUnknown l env + + goUnknown :: GLeftHandSide s env1 env2 -> WEnv' Info env1 env1 -> WEnv' Info env2 env2 + goUnknown (LeftHandSideSingle _) env = wpush env InfoNone + goUnknown (LeftHandSideWildcard _) env = env + goUnknown (LeftHandSidePair l1 l2) env = goUnknown l2 $ goUnknown l1 env +bindingEnv _ fenceSet lhs (Return variables) (InfoEnv environment) = InfoEnv $ weaken (weakenWithLHS lhs) $ go lhs variables environment + where + go :: GLeftHandSide s env1 env2 -> GroundVars env s -> WEnv' Info env env1 -> WEnv' Info env env2 + go (LeftHandSideSingle _) (TupRsingle (Var _ ix)) env = wpush' env $ InfoAlias fenceSet ix + go (LeftHandSidePair l1 l2) (TupRpair v1 v2) env = go l2 v2 $ go l1 v1 env + go (LeftHandSideWildcard _) _ env = env + go _ _ _ = internalError "Tuple mismatch" +bindingEnv _ fenceSet (LeftHandSideSingle _) Alloc{} (InfoEnv env) + | IdxSet.null fenceSet = InfoEnv $ wpush env InfoUndef +bindingEnv _ fenceSet (LeftHandSideSingle _) (Unit (Var _ idx)) (InfoEnv env) + | IdxSet.null fenceSet = InfoEnv $ wpush env $ InfoBuffer (Just $ SuccIdx idx) Nothing [] +bindingEnv outputs fenceSet (LeftHandSideWildcard _) (Exec op args) env + | IdxSet.null fenceSet = foldl' addCopy env $ detectCopy op args + where + addCopy :: InfoEnv env -> CopyOperation env -> InfoEnv env + addCopy env' (CopyOperation input output) + | input `IdxSet.member` outputs = env' -- The operation both reads and writes to 'input'. We cannot register input and output as copies, as input will get different values + | otherwise = InfoEnv $ wupdate markCopy input $ wupdate (const $ InfoBuffer Nothing (Just input) []) output $ unInfoEnv env' + where + markCopy (InfoBuffer unitScalar copyOf list) = InfoBuffer unitScalar copyOf (output : list) + markCopy (InfoAlias _ _) = internalError "Operation contains aliased variable, which should be substituted already" + markCopy _ = (InfoBuffer Nothing Nothing [output]) +bindingEnv outputs fenceSet lhs (Fence deps next) env + = bindingEnv outputs (IdxSet.union fenceSet deps) lhs next env +bindingEnv _ _ lhs _ env = bindEnv lhs env + +-- Updates the InfoEnv, with the information that the buffers in 'indices' may +-- have been updated. This breaks the 'is-copy-of' relation between buffers. +invalidate :: forall env. IdxSet env -> InfoEnv env -> InfoEnv env +invalidate indices infoEnv@(InfoEnv env1) = + InfoEnv + $ wupdateSetWeakened dropCopyTo indicesCopiesOf + $ wupdateSetWeakened dropCopyOf indicesCopiedTo + $ wremoveSet InfoNone indices' env1 + where + indices' :: IdxSet env + indices' = IdxSet.map (weaken $ substituteOutput infoEnv) indices + + findCopies :: Exists (Idx env) -> (IdxSet env, IdxSet env) + findCopies (Exists idx) = case infoFor idx infoEnv of + InfoAlias _ _ -> internalError "Alias should be substituted already" + InfoBuffer _ (Just idx') copies -> (IdxSet.singleton idx', IdxSet.fromList' copies) + _ -> (IdxSet.empty, IdxSet.empty) + + (indicesCopiesOf', indicesCopiedTo') = unzip $ map findCopies $ IdxSet.toList indices' + indicesCopiesOf = IdxSet.unions indicesCopiesOf' + indicesCopiedTo = IdxSet.unions indicesCopiedTo' + + -- Forgets that this buffer is a copy of a buffer in indices'. + dropCopyOf :: env' :> env -> Info env' t -> Info env' t + dropCopyOf _ (InfoBuffer unitScalar _ c) + = InfoBuffer unitScalar Nothing c + dropCopyOf _ _ = internalError "Expected Buffer" + + -- Forgets that this buffer is copied to buffers in indices' + dropCopyTo :: env' :> env -> Info env' t -> Info env' t + dropCopyTo k (InfoBuffer unitScalar copyOf copiedTo') + = InfoBuffer unitScalar copyOf $ filter (\idx -> not $ k >:> idx `IdxSet.member` indices) copiedTo' + dropCopyTo _ _ = internalError "Expected Buffer" + +outputArrays :: Args env args -> IdxSet env +outputArrays = IdxSet.fromList . mapMaybe f . argsVars + where + f :: Exists (Var AccessGroundR env) -> Maybe (Exists (Idx env)) + f (Exists (Var (AccessGroundRbuffer In _) _)) = Nothing + f (Exists (Var (AccessGroundRbuffer _ _) idx)) = Just (Exists idx) -- Out or Mut + f _ = Nothing + +simplifyExp :: forall env t. InfoEnv env -> Exp env t -> Exp env t +simplifyExp env = Exp.simplifyExp . runIdentity . rebuildArrayInstrOpenExp (simplifyArrayInstr env) + +simplifyExpFun :: forall env t. InfoEnv env -> Fun env t -> Fun env t +simplifyExpFun env = Exp.simplifyFun . runIdentity . rebuildArrayInstrFun (simplifyArrayInstr env) + +simplifyArrayInstr :: InfoEnv env -> RebuildArrayInstr Identity (ArrayInstr env) (ArrayInstr env) +simplifyArrayInstr env instr@(Parameter (Var tp idx)) = case infoFor idx env of + InfoAlias _ idx' -> simplifyArrayInstr env (Parameter $ Var tp idx') + InfoConst _ _ c -> Identity $ const $ Const tp c + InfoBuffer _ _ _ -> bufferImpossible tp + _ -> Identity $ \arg -> ArrayInstr instr arg +simplifyArrayInstr env instr@(Index (Var tp idx)) = case infoFor idx env of + InfoAlias _ idx' -> simplifyArrayInstr env (Index $ Var tp idx') + InfoBuffer (Just idx') _ _ -> Identity $ const $ runIdentity (simplifyArrayInstr env $ Parameter $ Var eltTp idx') Nil -- Unit + InfoBuffer _ (Just idx') _ -> simplifyArrayInstr env (Index $ Var tp idx') -- Copy + InfoUndef -> Identity $ const $ Undef eltTp + _ -> Identity $ \arg -> ArrayInstr instr arg + where + eltTp = case tp of + GroundRscalar t -> bufferImpossible t + GroundRbuffer t -> t + +simplifyArg :: InfoEnv env -> Arg env t -> Arg env t +simplifyArg env (ArgVar var) = ArgVar $ mapTupR (weaken $ substitute env) var +simplifyArg env (ArgExp expr) = ArgExp $ simplifyExp env expr +simplifyArg env (ArgFun fun) = ArgFun $ simplifyExpFun env fun +simplifyArg env (ArgArray m repr sh buffers) + = ArgArray m repr (mapTupR (weaken $ substitute env) sh) (mapTupR (weaken $ substituteBuffer env) buffers) + where + -- Output buffers may not be substituted by buffers with the same content. + substituteBuffer + | In <- m = substitute + | Out <- m = substituteOutput + | Mut <- m = const weakenId + +bindEnv :: GLeftHandSide t env env' -> InfoEnv env -> InfoEnv env' +bindEnv lhs (InfoEnv env') = InfoEnv $ go lhs $ weaken k env' + where + k = weakenWithLHS lhs + + go :: GLeftHandSide t env1 env1' -> WEnv' Info env' env1 -> WEnv' Info env' env1' + go (LeftHandSideWildcard _) env1 = env1 + go (LeftHandSideSingle _) env1 = wpush' env1 InfoNone + go (LeftHandSidePair l1 l2) env1 = go l2 $ go l1 env1 + +unionSubTupR :: SubTupR t s -> SubTupR t s' -> Exists (SubTupR t) +unionSubTupR SubTupRskip s = Exists s +unionSubTupR s SubTupRskip = Exists s +unionSubTupR (SubTupRpair l1 r1) (SubTupRpair l2 r2) + | Exists l <- unionSubTupR l1 l2 + , Exists r <- unionSubTupR r1 r2 + = Exists $ subTupRpair l r +unionSubTupR _ _ = Exists SubTupRkeep + +-- Detects which parts of the state of an awhile loop are invariant and +-- transforms the program accordingly. +-- For instance, if the state of an awhile loop contains an array whose +-- shape doesn't change throughout the execution of the awhile loop, +-- it will remove the shape from the state of the loop and always refer +-- to the initial shape of the array. +awhileSimplifyInvariant + :: Uniquenesses a + -> PreOpenAfun op env (a -> PrimBool) + -> PreOpenAfun op env (a -> a) + -> GroundVars env a + -> PreOpenAcc op env a +awhileSimplifyInvariant us cond step initial = case awhileDropInvariantFun initial step of + Exists SubTupRkeep -> Awhile us cond step initial + Exists sub + | Just Refl <- subTupPreserves tp sub -> Awhile us cond step initial + | DeclareVars lhs k value <- declareVars $ subTupR sub tp -> + alet' lhs (subTupR sub us) + (Awhile (subTupR sub us) + (subTupFunctionArgument sub initial cond) + (subTupFunctionArgument sub initial $ subTupFunctionResult sub step) + (subTupR sub initial)) + (Return $ subTupResult sub (mapTupR (weaken k) initial) (value weakenId)) + where + tp = case cond of + Alam lhs _ -> lhsToTupR lhs + Abody body -> groundFunctionImpossible $ groundsR body + +awhileDropInvariantFun :: GroundVars env t -> OperationAfun op env (t -> t) -> Exists (SubTupR t) +awhileDropInvariantFun initial (Alam lhs (Abody body)) = + awhileDropInvariant (mapTupR (weaken $ weakenWithLHS lhs) initial) (lhsMaybeVars lhs) body +awhileDropInvariantFun _initial (Alam lhs (Alam _ _)) = groundFunctionImpossible (lhsToTupR lhs) +awhileDropInvariantFun _initial (Abody body) = groundFunctionImpossible (groundsR body) + +-- Computes a SubTupR that removes variables that are invariant in the while loop. +-- Invariant here means that the step function of the loop returns the input unchanged, +-- or returns the initial value in each step. +-- This pattern often occurs with sizes of arrays. +awhileDropInvariant :: GroundVars env t -> MaybeVars GroundR env t -> OperationAcc op env t -> Exists (SubTupR t) +awhileDropInvariant initial argument = \case + Return vars + -> matchReturn initial argument vars + Alet (LeftHandSideWildcard _) _ _ body + -> awhileDropInvariant initial argument body + Alet lhs _ _ body + -> awhileDropInvariant + (mapTupR (weaken $ weakenWithLHS lhs) initial) + (mapTupR (weaken $ weakenWithLHS lhs) argument) + body + Acond _ t f + | Exists subTupT <- awhileDropInvariant initial argument t + , Exists subTupF <- awhileDropInvariant initial argument f + -- Only remove variables if they are invariant in both branches. + -- Thus, preserve variables in the union of subTupT and subTupF. + -> unionSubTupR subTupT subTupF + _ -> Exists SubTupRkeep -- No invariant variables + where + matchReturn :: GroundVars env t' -> MaybeVars GroundR env t' -> GroundVars env t' -> Exists (SubTupR t') + matchReturn (TupRpair i1 i2) (TupRpair a1 a2) (TupRpair v1 v2) + | Exists s1 <- matchReturn i1 a1 v1 + , Exists s2 <- matchReturn i2 a2 v2 + = case (s1, s2) of + (SubTupRskip, SubTupRskip) -> Exists SubTupRskip + _ -> Exists $ subTupRpair s1 s2 + matchReturn (TupRsingle i) (TupRsingle arg) (TupRsingle var) + | Just Refl <- matchVar i var + = Exists SubTupRskip + | JustVar arg' <- arg + , Just Refl <- matchVar arg' var + = Exists SubTupRskip + | otherwise + = Exists SubTupRkeep + matchReturn TupRunit _ _ = Exists SubTupRskip + matchReturn _ _ _ = internalError "Tuple mismatch" + +subTupFunctionResult :: SubTupR t t' -> OperationAfun op env (ta -> t) -> OperationAfun op env (ta -> t') +subTupFunctionResult sub (Alam lhs (Abody body)) = Alam lhs $ Abody $ subTupAcc sub body +subTupFunctionResult _ _ = internalError "Illegal function" + +subTupAcc :: SubTupR t t' -> OperationAcc op env t -> OperationAcc op env t' +subTupAcc sub = \case + Return vars -> Return $ subTupR sub vars + Alet lhs us bnd body -> Alet lhs us bnd $ subTupAcc sub body + Acond c t f -> Acond c (subTupAcc sub t) (subTupAcc sub f) + _ -> internalError "Cannot subTup this program" + +subTupFunctionArgument :: SubTupR t t' -> GroundVars env t -> OperationAfun op env (t -> tr) -> OperationAfun op env (t' -> tr) +subTupFunctionArgument sub initial (Alam lhs body) + | SubTupSubstitution lhs' k <- subTupSubstitution sub lhs initial + = Alam lhs' $ weaken k body +subTupFunctionArgument _ _ (Abody body) = groundFunctionImpossible $ groundsR body + +subTupResult :: SubTupR t t' -> GroundVars env t -> GroundVars env t' -> GroundVars env t +subTupResult SubTupRkeep _ result = result +subTupResult SubTupRskip initial _ = initial +subTupResult (SubTupRpair s1 s2) (TupRpair i1 i2) (TupRpair r1 r2) = subTupResult s1 i1 r1 `TupRpair` subTupResult s2 i2 r2 +subTupResult _ _ _ = internalError "Tuple mismatch" + +data SubTupSubstitution env env1 t t' where + SubTupSubstitution + :: GLeftHandSide t' env1 env2 + -> env :> env2 + -> SubTupSubstitution env env1 t t' + +subTupSubstitution :: SubTupR t t' -> GLeftHandSide t env env' -> GroundVars env t -> SubTupSubstitution env' env t t' +subTupSubstitution SubTupRskip lhs vars = SubTupSubstitution (LeftHandSideWildcard TupRunit) (go lhs vars weakenId) + where + go :: GLeftHandSide s env1 env2 -> GroundVars env s -> env1 :> env -> env2 :> env + go (LeftHandSideWildcard _) _ k = k + go (LeftHandSideSingle _) (TupRsingle (Var _ ix)) k = Weaken $ \case + ZeroIdx -> ix + SuccIdx ix' -> k >:> ix' + go (LeftHandSidePair l1 l2) (TupRpair v1 v2) k = go l2 v2 $ go l1 v1 k + go _ _ _ = internalError "Tuple mismatch" +subTupSubstitution SubTupRkeep lhs _ = SubTupSubstitution lhs weakenId +subTupSubstitution (SubTupRpair s1 s2) (LeftHandSidePair l1 l2) (TupRpair v1 v2) + | SubTupSubstitution l1' k1 <- subTupSubstitution s1 l1 v1 + , Exists l2'' <- rebuildLHS l2 + , SubTupSubstitution l2' k2 <- subTupSubstitution s2 l2'' (mapTupR (weaken $ weakenWithLHS l1') v2) + = SubTupSubstitution (LeftHandSidePair l1' l2') (k2 .> sinkWithLHS l2 l2'' k1) +subTupSubstitution s (LeftHandSideWildcard t) _ + = SubTupSubstitution (LeftHandSideWildcard $ subTupR s t) weakenId +subTupSubstitution _ _ _ = internalError "Tuple mismatch" + +assert :: Text -> Exp env PrimBool -> OperationAcc op env Word8 +assert _ (Const _ 1) = Compute (Const scalarTypeWord8 1) +assert msg c = Aassert msg c + +assume :: Exp env PrimBool -> OperationAcc op env Word8 +assume (Const _ 1) = Compute (Const scalarTypeWord8 1) +assume c = Aassume c + +compute :: Exp env a -> OperationAcc op env a +compute (Assert msg c expr) = + alet + (LeftHandSideSingle $ GroundRscalar scalarTypeWord8) + (Aassert msg c) + $ Fence (IdxSet.singleton ZeroIdx) + $ compute $ mapArrayInstr (weaken $ weakenSucc weakenId) expr +compute (Assume c expr) = + alet + (LeftHandSideSingle $ GroundRscalar scalarTypeWord8) + (Aassume c) + $ Fence (IdxSet.singleton ZeroIdx) + $ compute $ mapArrayInstr (weaken $ weakenSucc weakenId) expr +compute expr + | Just vars <- extractParams expr = + Return $ mapTupR (\(Var tp ix) -> Var (GroundRscalar tp) ix) vars + | otherwise = + Compute expr diff --git a/src/Data/Array/Accelerate/Trafo/Operation/Substitution.hs b/src/Data/Array/Accelerate/Trafo/Operation/Substitution.hs new file mode 100644 index 000000000..287901e35 --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Operation/Substitution.hs @@ -0,0 +1,264 @@ +{-# OPTIONS_GHC -Wno-orphans #-} +{-# OPTIONS_HADDOCK hide #-} +-- | +-- Module : Data.Array.Accelerate.Trafo.Operation.Substitution +-- Copyright : [2012..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.Trafo.Operation.Substitution ( + Sink(..), Sink'(..), + reindexPartial, + reindexPartialAfun, + pair, pair', pairUnique, alet, aletUnique, alet', fence, + makeManifest, + weakenArrayInstr, + strengthenArrayInstr, + extractParams, + + reindexVar, reindexVars, + mkLHS, LHS(..), + ) where + +import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.Var +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.AST.Operation +import Data.Array.Accelerate.AST.Environment +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Trafo.Var +import Data.Array.Accelerate.Trafo.Substitution (Sink(..), Sink'(..)) +import Data.Array.Accelerate.Trafo.Exp.Substitution +import Data.Functor.Identity +import Data.Array.Accelerate.AST.IdxSet (IdxSet) +import qualified Data.Array.Accelerate.AST.IdxSet as IdxSet + +data SunkReindexPartial f env env' where + Sink :: SunkReindexPartial f env env' -> SunkReindexPartial f (env, s) (env', s) + ReindexF :: ReindexPartial f env env' -> SunkReindexPartial f env env' + +reindexPartial :: (Applicative f) => ReindexPartial f env env' -> PreOpenAcc op env t -> f (PreOpenAcc op env' t) +reindexPartial k = reindexA' (ReindexF k) + +reindexPartialAfun :: (Applicative f) => ReindexPartial f env env' -> PreOpenAfun op env t -> f (PreOpenAfun op env' t) +reindexPartialAfun k = reindexAfun' (ReindexF k) + +instance Sink (PreOpenAcc op) where + weaken k = runIdentity . reindexPartial (Identity . (k >:>)) + +instance Sink (PreOpenAfun op) where + weaken k = runIdentity . reindexPartialAfun (Identity . (k >:>)) + +instance Sink Arg where + weaken k = runIdentity . reindexArg (weakenReindex k) + +instance Sink ArrayInstr where + weaken k (Index v) = Index $ weaken k v + weaken k (Parameter v) = Parameter $ weaken k v + +sinkReindexWithLHS :: LeftHandSide s t env1 env1' -> LeftHandSide s t env2 env2' -> SunkReindexPartial f env1 env2 -> SunkReindexPartial f env1' env2' +sinkReindexWithLHS (LeftHandSideWildcard _) (LeftHandSideWildcard _) k = k +sinkReindexWithLHS (LeftHandSideSingle _) (LeftHandSideSingle _) k = Sink k +sinkReindexWithLHS (LeftHandSidePair a1 b1) (LeftHandSidePair a2 b2) k = sinkReindexWithLHS b1 b2 $ sinkReindexWithLHS a1 a2 k +sinkReindexWithLHS _ _ _ = error "sinkReindexWithLHS: left hand sides don't match" + +-- All functions ending in a prime work with SunkReindexPartial instead of ReindexPartial. +reindex' :: Applicative f => SunkReindexPartial f env env' -> ReindexPartial f env env' +reindex' (ReindexF f) = f +reindex' (Sink k) = \case + ZeroIdx -> pure ZeroIdx + SuccIdx ix -> SuccIdx <$> reindex' k ix + +reindexVar' :: Applicative f => SunkReindexPartial f env env' -> Var s env t -> f (Var s env' t) +reindexVar' k (Var repr ix) = Var repr <$> reindex' k ix + +reindexVars' :: Applicative f => SunkReindexPartial f env env' -> Vars s env t -> f (Vars s env' t) +reindexVars' _ TupRunit = pure $ TupRunit +reindexVars' k (TupRsingle var) = TupRsingle <$> reindexVar' k var +reindexVars' k (TupRpair v1 v2) = TupRpair <$> reindexVars' k v1 <*> reindexVars' k v2 + +reindexArrayInstr' :: Applicative f => SunkReindexPartial f env env' -> ArrayInstr env (s -> t) -> f (ArrayInstr env' (s -> t)) +reindexArrayInstr' k (Index v) = Index <$> reindexVar' k v +reindexArrayInstr' k (Parameter v) = Parameter <$> reindexVar' k v + +reindexExp' :: (Applicative f, RebuildableExp e) => SunkReindexPartial f benv benv' -> e (ArrayInstr benv) env t -> f (e (ArrayInstr benv') env t) +reindexExp' k = rebuildArrayInstrPartial (rebuildArrayInstrMap $ reindexArrayInstr' k) + +reindexArrayDescriptor' :: (Applicative f) => SunkReindexPartial f env env' -> ArrayDescriptors env t -> f (TupR (ArrayDescriptor env') t) +reindexArrayDescriptor' k = traverseTupR (\(ArrayDescriptor shr sh buffer) -> ArrayDescriptor shr <$> reindexVars' k sh <*> reindexVars' k buffer) + +reindexIdxSet' + :: forall f env env' . Applicative f + => SunkReindexPartial f env env' + -> IdxSet env + -> f (IdxSet env') +reindexIdxSet' k set = + foldr go (pure IdxSet.empty) (IdxSet.toList set) + where + go :: Exists (Idx env) -> f (IdxSet env') -> f (IdxSet env') + go (Exists ix) acc = IdxSet.insert <$> reindex' k ix <*> acc + +reindexA' :: forall op f env env' t. (Applicative f) => SunkReindexPartial f env env' -> PreOpenAcc op env t -> f (PreOpenAcc op env' t) +reindexA' k = \case + Exec op args -> Exec op <$> reindexArgs (reindex' k) args + Return vars -> Return <$> reindexVars' k vars + Manifest var -> Manifest <$> reindexVar' k var + Compute e -> Compute <$> reindexExp' k e + Alet lhs uniqueness bnd body + | Exists lhs' <- rebuildLHS lhs -> Alet lhs' uniqueness <$> travA bnd <*> reindexA' (sinkReindexWithLHS lhs lhs' k) body + Alloc shr tp sh -> Alloc shr tp <$> reindexVars' k sh + Use tp n buffer -> pure $ Use tp n buffer + Unit var -> Unit <$> reindexVar' k var + Acond c t f -> Acond <$> reindexVar' k c <*> travA t <*> travA f + Awhile uniqueness c f i -> Awhile uniqueness <$> reindexAfun' k c <*> reindexAfun' k f <*> reindexVars' k i + Atrace msg t -> Atrace msg <$> reindexArrayDescriptor' k t + Aassert msg g -> Aassert msg <$> reindexExp' k g + Aassume g -> Aassume <$> reindexExp' k g + Fence set next -> Fence <$> reindexIdxSet' k set <*> travA next + where + travA :: PreOpenAcc op env s -> f (PreOpenAcc op env' s) + travA = reindexA' k + +reindexAfun' :: (Applicative f) => SunkReindexPartial f env env' -> PreOpenAfun op env t -> f (PreOpenAfun op env' t) +reindexAfun' k (Alam lhs f) + | Exists lhs' <- rebuildLHS lhs = Alam lhs' <$> reindexAfun' (sinkReindexWithLHS lhs lhs' k) f +reindexAfun' k (Abody a) = Abody <$> reindexA' k a + +pair :: forall op env a b. PreOpenAcc op env a -> PreOpenAcc op env b -> PreOpenAcc op env (a, b) +pair = pair' False + +pairUnique :: forall op env a b. PreOpenAcc op env a -> PreOpenAcc op env b -> PreOpenAcc op env (a, b) +pairUnique = pair' True + +pair' :: forall op env a b. Bool -> PreOpenAcc op env a -> PreOpenAcc op env b -> PreOpenAcc op env (a, b) +pair' u a b = goA weakenId a + where + -- Traverse 'a' and look for a return. We can jump over let bindings + -- If we don't find a 'return', we must first bind the value in a let, + -- and then use the newly defined variables instead. + -- + goA :: env :> env' -> PreOpenAcc op env' a -> PreOpenAcc op env' (a, b) + goA k (Alet lhs uniqueness bnd x) + = Alet lhs uniqueness bnd $ goA (weakenWithLHS lhs .> k) x + goA k (Return vars) = goB vars $ weaken k b + goA k acc + | DeclareVars lhs k' value <- declareVars $ groundsR acc + = Alet lhs ((if u then unique else shared) $ groundsR acc) acc $ goB (value weakenId) $ weaken (k' .> k) b + + goB :: GroundVars env' a -> PreOpenAcc op env' b -> PreOpenAcc op env' (a, b) + goB varsA (Alet lhs uniqueness bnd x) + = Alet lhs uniqueness bnd $ goB (weakenVars (weakenWithLHS lhs) varsA) x + goB varsA (Return varsB) = Return (TupRpair varsA varsB) + goB varsA acc + | DeclareVars lhs k value <- declareVars $ groundsR b + = Alet lhs ((if u then unique else shared) $ groundsR b) acc $ Return (TupRpair (weakenVars k varsA) (value weakenId)) + +makeManifest :: PreOpenAcc op env a -> PreOpenAcc op env a +makeManifest acc = case acc of + Alet lhs uniqueness bnd x -> Alet lhs uniqueness bnd $ makeManifest x + Exec{} -> acc -- Doesn't return anything + Manifest{} -> acc -- Already manifest + Compute{} -> acc -- Doesn't return a buffer, only buffers should explicitely be marked as manifest + Aassert{} -> acc -- Same as compute + Aassume{} -> acc -- Same as compute + Alloc{} -> acc -- Can't fuse anyway + Use{} -> acc -- Can't fuse anyway + Unit{} -> acc -- Can't fuse anyway + -- Acond can't fuse, but we still mark it explicitely as manifest in case the + -- condition is evaluated at compile time. + Acond c t f -> Acond c (makeManifest t) (makeManifest f) + Awhile{} -> acc -- Can't fuse anyway + Atrace{} -> acc -- Same as compute + Return vars -> go vars + Fence set next -> Fence set $ makeManifest next + where + go :: GroundVars env t -> PreOpenAcc op env t + go TupRunit = Return TupRunit + go (TupRpair v1 v2) = go v1 `pair` go v2 + go (TupRsingle var@(Var tp _)) = case tp of + GroundRbuffer _ -> Manifest var + GroundRscalar _ -> Return $ TupRsingle var + +alet :: GLeftHandSide t env env' -> PreOpenAcc op env t -> PreOpenAcc op env' s -> PreOpenAcc op env s +alet lhs = alet' lhs $ shared $ lhsToTupR lhs + +aletUnique :: GLeftHandSide t env env' -> PreOpenAcc op env t -> PreOpenAcc op env' s -> PreOpenAcc op env s +aletUnique lhs = alet' lhs $ unique $ lhsToTupR lhs + +alet' :: GLeftHandSide t env env' -> Uniquenesses t -> PreOpenAcc op env t -> PreOpenAcc op env' s -> PreOpenAcc op env s +alet' lhs us (Fence set1 bnd) (Fence set2 body) + | set3 <- IdxSet.intersect set1 (IdxSet.drop' lhs set2) + , not $ IdxSet.null set3 + = Fence set3 $ alet' lhs us + (fence (set1 IdxSet.\\ set3) bnd) + (fence (set2 IdxSet.\\ IdxSet.skip' lhs set3) body) +alet' lhs1 us (Alet lhs2 uniqueness a1 a2) a3 + -- Perform let rotation: + -- let x = (let y = a1 in a2) in a3 + -- becomes + -- let y = a1 in let x = a2 in a3 + | Exists lhs1' <- rebuildLHS lhs1 = Alet lhs2 uniqueness a1 $ alet' lhs1' us a2 $ weaken (sinkWithLHS lhs1 lhs1' $ weakenWithLHS lhs2) a3 +-- If the binding is a fence containing a let binding, +-- then push the fence into that let binding to enable let rotation +alet' lhs1 us (Fence set (Alet lhs2 uniqueness a1 a2)) a3 = + alet' lhs1 us (Alet lhs2 uniqueness a1' a2') a3 + where + a1' = fence set a1 + a2' = fence (IdxSet.skip' lhs2 set) a2 +alet' lhs@(LeftHandSideWildcard TupRunit) _ bnd a = case bnd of + Compute _ -> a + Return _ -> a + _ -> Alet lhs TupRunit bnd a -- 'bnd' may have side effects +alet' lhs _ (Return vars) a = weaken (substituteLHS lhs vars) a +alet' lhs us (Compute e) a + | Just vars <- extractParams e = weaken (substituteLHS lhs vars) a + | LeftHandSidePair l1 l2 <- lhs + , TupRpair u1 u2 <- us + , Pair e1 e2 <- e = alet' l1 u1 (Compute e1) $ alet' l2 u2 (weaken (weakenWithLHS l1) $ Compute e2) a +alet' lhs us bnd a = Alet lhs us bnd a + +fence :: IdxSet env -> PreOpenAcc op env t -> PreOpenAcc op env t +fence set + | IdxSet.null set = id + | otherwise = \case + Fence set2 next -> fence (IdxSet.union set set2) next + -- Alet lhs us bnd body -> Alet lhs us (fence set bnd) (fence (IdxSet.skip' lhs set) body) + a -> Fence set a + +extractParams :: OpenExp env benv t -> Maybe (ExpVars benv t) +extractParams Nil = Just TupRunit +extractParams (Pair e1 e2) = TupRpair <$> extractParams e1 <*> extractParams e2 +extractParams (ArrayInstr (Parameter v) _) = Just $ TupRsingle v +extractParams _ = Nothing + +-- We cannot use 'weaken' to weaken the array environment of an 'OpenExp', +-- as OpenExp is a type synonym for 'PreOpenExp (ArrayInstr aenv) env', +-- and 'weaken' would thus affect the expression environment. Hence we +-- have a separate function for OpenExp and OpenFun. +-- +weakenArrayInstr :: RebuildableExp f => aenv :> aenv' -> f (ArrayInstr aenv) env t -> f (ArrayInstr aenv') env t +weakenArrayInstr k = rebuildArrayInstr (ArrayInstr . weaken k) + +strengthenArrayInstr :: forall f benv benv' env t. RebuildableExp f => benv :?> benv' -> f (ArrayInstr benv) env t -> Maybe (f (ArrayInstr benv') env t) +strengthenArrayInstr k = rebuildArrayInstrPartial f + where + f :: ArrayInstr benv (u -> s) -> Maybe (OpenExp env' benv' u -> OpenExp env' benv' s) + f (Parameter (Var tp ix)) = ArrayInstr . Parameter . Var tp <$> k ix + f (Index (Var tp ix)) = ArrayInstr . Index . Var tp <$> k ix + + +data LHS s v env where + LHS :: LeftHandSide s v env env' + -> Vars s env' v + -> LHS s v env + +mkLHS :: TupR s v -> LHS s v env +mkLHS TupRunit = LHS LeftHandSideUnit TupRunit +mkLHS (TupRsingle x) = LHS (LeftHandSideSingle x) (TupRsingle $ Var x ZeroIdx) +mkLHS (TupRpair x y) = case mkLHS x of + LHS lhsL varsL -> case mkLHS y of + LHS lhsR varsR -> LHS (LeftHandSidePair lhsL lhsR) (TupRpair (weakenVars (weakenWithLHS lhsR) varsL) varsR) diff --git a/src/Data/Array/Accelerate/Trafo/Partitioning/ILP.hs b/src/Data/Array/Accelerate/Trafo/Partitioning/ILP.hs new file mode 100644 index 000000000..4da9c9f1b --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Partitioning/ILP.hs @@ -0,0 +1,205 @@ +module Data.Array.Accelerate.Trafo.Partitioning.ILP where +-- No joke, this really needs to get a massive refactor... + +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Graph +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Var +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Encoding + ( interpretClusters, makeILP, splitExecs, ClusterLs, Objective (..), interpretReadDirs, interpretInplaceUpdates ) +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Clustering + ( reconstruct, reconstructF, ReadDirM, InplaceM ) +import Data.Array.Accelerate.AST.Partitioned + ( PartitionedAcc, PartitionedAfun, Cluster, groundsR ) +import Data.Array.Accelerate.AST.Operation + ( OperationAcc, OperationAfun ) +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Solver + ( ILPSolver, solve, Solution ) +import Data.Array.Accelerate.Trafo.Partitioning.ILP.MIP + ( cbc, cplex, glpsol, gurobiCl, lpSolve, scip, MIP(..) ) + +import System.IO.Unsafe (unsafePerformIO) +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Labels (Node, Comp) +import Data.Map (Map, toList, foldMapWithKey, filterWithKey) +import Data.Array.Accelerate.Trafo.Operation.Simplify +import qualified Data.Array.Accelerate.Pretty.Operation as Pretty +import Lens.Micro ((^.)) +import Data.Maybe (fromMaybe) + +data Benchmarking = GreedyUp | GreedyDown | NoFusion + deriving (Show, Eq, Bounded, Enum) + +data FusionType = Fusion Objective | Benchmarking Benchmarking + +defaultObjective :: FusionType +defaultObjective = Fusion MemoryUsage' + +-- data type that should probably be in the options +newtype Solver = MIPSolver MIPSolver +data MIPSolver = CBC | Gurobi | CPLEX | GLPSOL | LPSOLVE | SCIP + +ilpFusion'' :: (MakesILP op, SimplifyOperation op, Pretty.PrettyOp (Cluster op)) => Solver -> Objective -> OperationAcc op () a -> PartitionedAcc op () a +ilpFusion'' (MIPSolver s) = case s of + CBC -> ilpFusion (MIP cbc) + Gurobi -> ilpFusion (MIP gurobiCl) + CPLEX -> ilpFusion (MIP cplex) + GLPSOL -> ilpFusion (MIP glpsol) + LPSOLVE -> ilpFusion (MIP lpSolve) + SCIP -> ilpFusion (MIP scip) + + +ilpFusionF'' :: (MakesILP op, SimplifyOperation op, Pretty.PrettyOp (Cluster op)) => Solver -> Objective -> OperationAfun op () a -> PartitionedAfun op () a +ilpFusionF'' (MIPSolver s) = case s of + CBC -> ilpFusionF (MIP cbc) + Gurobi -> ilpFusionF (MIP gurobiCl) + CPLEX -> ilpFusionF (MIP cplex) + GLPSOL -> ilpFusionF (MIP glpsol) + LPSOLVE -> ilpFusionF (MIP lpSolve) + SCIP -> ilpFusionF (MIP scip) + + +ilpFusion :: (MakesILP op, SimplifyOperation op, ILPSolver s, Pretty.PrettyOp (Cluster op)) => s -> Objective -> OperationAcc op () a -> PartitionedAcc op () a +ilpFusion solver objective acc = ilpFusion' mkFullGraph (reconstruct (groundsR acc) False) solver objective acc + +ilpFusionF :: (MakesILP op, SimplifyOperation op, ILPSolver s, Pretty.PrettyOp (Cluster op)) => s -> Objective -> OperationAfun op () a -> PartitionedAfun op () a +ilpFusionF solver objective fun = ilpFusion' mkFullGraphF (reconstructF fun False) solver objective fun + +ilpFusion' :: (MakesILP op, SimplifyOperation op, ILPSolver s) + => (x -> FullGraph op) + -> (FusionGraph -> [ClusterLs] -> Map (Node Comp) [ClusterLs] -> Symbols op -> ReadDirM -> InplaceM -> y) + -> s + -> Objective + -> x + -> y +ilpFusion' toGraph fromGraph s obj acc = do + let fullgraph = {- traceGraph $ -} toGraph acc + let ilp = makeILP obj (fullgraph^.fusionILP) + let solution = {- traceWith ppNumInplace $ -} fromMaybe (error "Accelerate: No ILP solution found") (unsafePerformIO $ solve s ilp) + let symbols' = attachBackendLabels solution (fullgraph^.symbols) + let readDirM = interpretReadDirs solution + -- let writeDirM = interpretWriteDirs solution + let inplaceM = interpretInplaceUpdates solution + let (topClusters, subClustersM) = splitExecs (interpretClusters solution) symbols' + fromGraph (fullgraph^.fusionILP.graph) topClusters subClustersM symbols' readDirM inplaceM + +traceGraph :: FullGraph op -> FullGraph op +traceGraph g = unsafePerformIO $ do + writeFile "ilp.dot" $ toDOT (g^.fusionILP.graph) (g^.symbols) + return g + +ppNumInplace :: Solution -> String +ppNumInplace m = "Compilation performed " ++ show numInplace ++ "/" ++ show totalInplace ++ " in-place updates." + where + inplaceVars = filterWithKey (\k _ -> case k of InPlace{} -> True; _ -> False) m + totalInplace = length inplaceVars + numInplace = length $ filterWithKey (\_ v -> v == 0) inplaceVars + + +ppSolution :: Solution -> String +ppSolution solution = "solution: " ++ foldMap ppVar (toList solution) + where + ppVar :: (Var, Int) -> String + ppVar (k, v) = "\n" ++ show k ++ " == " ++ show v + -- ppVar (k, v) = case k of + -- Pi{} -> "\n" ++ show k ++ " == " ++ show v + -- Fused{} | v == 0 -> "\n" ++ show k + -- Manifest{} -> "\n" ++ show k ++ " == " ++ show v + -- InPlace{} | v == 0 -> "\n" ++ show k + -- _ -> "" + +ppList :: Show a => [a] -> String +ppList [] = "[]" +ppList [x] = "[" ++ show x ++ "]" +ppList (x:xs) = "[ " ++ show x ++ foldMap (\x' -> "\n, " ++ show x') xs ++ "\n]" + +ppScopedClusters :: (Show k, Show v) => ([v], Map k [v]) -> String +ppScopedClusters (top, sub) = "top =\n" ++ ppList top ++ foldMapWithKey (\k v -> "\n" ++ show k ++ " =\n" ++ ppList v) sub + +-- for benchmarking: make all edges infusible +-- note: does allow for horizontal fusion! +-- more rigorous is to change 'topSort' in Clustering.hs into separating each cluster completely +noFusion' :: (MakesILP op, SimplifyOperation op, ILPSolver s) + => (x -> FullGraph op) + -> (FusionGraph -> [ClusterLs] -> Map (Node Comp) [ClusterLs] -> Symbols op -> ReadDirM -> InplaceM -> y) + -> s + -> Objective + -> x + -> y +noFusion' = undefined +-- noFusion' k1 k2 s obj acc = fusedAcc +-- where +-- (fusionILP', constrM') = k1 acc +-- fusionILP'' = fusionILP' & graph.strictEdges <>~ Set.map (\(w,_,r) -> (w,r)) (fusionILP'^.graph.fusibleEdges) +-- constrM = attachBackendLabels solution constrM' +-- ilp = makeILP obj fusionILP'' +-- solution = solve' ilp +-- interpreted = interpretSolution solution +-- (labelClusters, labelClustersM) = splitExecs interpreted constrM +-- fusedAcc = k2 (fusionILP'^.graph) labelClusters labelClustersM constrM +-- solve' x = unsafePerformIO (solve s x) & \case +-- Nothing -> error "Accelerate: No ILP solution found" +-- Just y -> y + +-- for benchmarking: greedily fuse edges +-- this search is clearly inefficient, but just an easy implementation. We only benchmark its runtime. +-- note that this is perhaps still too generous. For example, anything that can fuse into 1 loop will still be fully fused! +-- it's perhaps more of an 'alternative' than a 'baseline' +greedyFusion' :: forall s op x y. (MakesILP op, SimplifyOperation op, ILPSolver s) + => (x -> FullGraph op) + -> (FusionGraph -> [ClusterLs] -> Map (Node Comp) [ClusterLs] -> Symbols op -> ReadDirM -> InplaceM -> y) + -> s + -> Benchmarking + -> Objective + -> x + -> y +greedyFusion' = undefined +-- greedyFusion' k1 k2 s b obj acc = fusedAcc +-- where +-- (info'@(FusionILP graph _ _), constrM') = k1 acc +-- nedges = (graph^.fusibleEdges) Set.\\ (graph^.infusibleEdges) & Set.size +-- go :: Int -> FusionILP op -> FusionILP op +-- go n info -- loop over all fusible edges. Try to set the current one to fused, if there's still a legal solution, keep it fused and continue. +-- | n >= nedges = info +-- | otherwise = let +-- i:->j = (graph^.fusibleEdges) Set.\\ (graph^.infusibleEdges)&Set.elemAt (case b of +-- GreedyUp -> n +-- GreedyDown -> nedges - n - 1 +-- _ -> error "nope") +-- info'' = info&constr<>~(fused i j .==. int 0) +-- in go (n+1) $ if check info'' then info'' else info +-- check :: FusionILP op -> Bool +-- check info = let +-- ilp = makeILP @op obj info +-- in isJust $ unsafePerformIO (solve s ilp) +-- info = go 0 info' +-- ilp = makeILP FusedEdges info + +-- constrM = backendConstruc solution constrM' +-- solution = solve' ilp +-- interpreted = interpretSolution solution +-- (labelClusters, labelClustersM) = splitExecs interpreted constrM +-- fusedAcc = k2 graph labelClusters labelClustersM constrM +-- solve' x = unsafePerformIO (solve s x) & \case +-- Nothing -> error "Accelerate: No ILP solution found" +-- Just y -> y + +bench :: (MakesILP op, SimplifyOperation op, Pretty.PrettyOp (Cluster op)) => Benchmarking -> Objective -> OperationAcc op () a -> PartitionedAcc op () a +bench NoFusion = no +bench b = greedy b +benchF :: (MakesILP op, SimplifyOperation op, Pretty.PrettyOp (Cluster op)) => Benchmarking -> Objective -> OperationAfun op () a -> PartitionedAfun op () a +benchF NoFusion = noF +benchF b = greedyF b +greedy :: (MakesILP op, SimplifyOperation op, Pretty.PrettyOp (Cluster op)) => Benchmarking -> Objective -> OperationAcc op () a -> PartitionedAcc op () a +greedy = greedyFusion (MIP gurobiCl) +no :: (MakesILP op, SimplifyOperation op, Pretty.PrettyOp (Cluster op)) => Objective -> OperationAcc op () a -> PartitionedAcc op () a +no = noFusion (MIP gurobiCl) +greedyF :: (MakesILP op, SimplifyOperation op, Pretty.PrettyOp (Cluster op)) => Benchmarking -> Objective -> OperationAfun op () a -> PartitionedAfun op () a +greedyF = greedyFusionF (MIP gurobiCl) +noF :: (MakesILP op, SimplifyOperation op, Pretty.PrettyOp (Cluster op)) => Objective -> OperationAfun op () a -> PartitionedAfun op () a +noF = noFusionF (MIP gurobiCl) +greedyFusion :: (MakesILP op, SimplifyOperation op, ILPSolver s, Pretty.PrettyOp (Cluster op)) => s -> Benchmarking -> Objective -> OperationAcc op () a -> PartitionedAcc op () a +greedyFusion solver b objective acc = greedyFusion' mkFullGraph (reconstruct (groundsR acc) False) solver b objective acc +greedyFusionF :: (MakesILP op, SimplifyOperation op, ILPSolver s, Pretty.PrettyOp (Cluster op)) => s -> Benchmarking -> Objective -> OperationAfun op () a -> PartitionedAfun op () a +greedyFusionF solver b objective fun = greedyFusion' mkFullGraphF (reconstructF fun False) solver b objective fun +noFusion :: (MakesILP op, SimplifyOperation op, ILPSolver s, Pretty.PrettyOp (Cluster op)) => s -> Objective -> OperationAcc op () a -> PartitionedAcc op () a +noFusion solver objective acc = noFusion' mkFullGraph (reconstruct (groundsR acc) True) solver objective acc +noFusionF :: (MakesILP op, SimplifyOperation op, ILPSolver s, Pretty.PrettyOp (Cluster op)) => s -> Objective -> OperationAfun op () a -> PartitionedAfun op () a +noFusionF solver objective fun = noFusion' mkFullGraphF (reconstructF fun True) solver objective fun diff --git a/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/Clustering.hs b/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/Clustering.hs new file mode 100644 index 000000000..cb71943b3 --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/Clustering.hs @@ -0,0 +1,398 @@ +{-# LANGUAGE ViewPatterns #-} +module Data.Array.Accelerate.Trafo.Partitioning.ILP.Clustering where + +import Data.Array.Accelerate.AST.LeftHandSide ( Exists(..), LeftHandSide (..), lhsToTupR ) +import Data.Array.Accelerate.AST.Partitioned +import Data.Array.Accelerate.AST.Var +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Type ( scalarType ) +import Data.Array.Accelerate.Trafo.Operation.Simplify +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Graph hiding (readEdges, writeEdges, strictEdges, dataflowEdges, symbols, graph) +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Labels +import Data.Array.Accelerate.Analysis.Match +import Data.Array.Accelerate.Error + +import qualified Data.Map as M +import qualified Data.Graph as G +import qualified Data.Set as S +import Data.Maybe (fromJust) +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Encoding (ClusterLs (Execs, NonExec)) +import Data.Array.Accelerate.AST.Environment (weakenWithLHS) + +import Prelude hiding ( take ) +import Lens.Micro.Extras (view) +import Data.Array.Accelerate.Representation.Array (ArrayR (ArrayR)) +import Data.Functor.Identity +import qualified Data.Tree as T +import Data.Array.Accelerate.Representation.Shape (shapeType) + +import Data.String (fromString) +import qualified Data.Array.Accelerate.Pretty.Operation as P +import qualified Data.Array.Accelerate.Pretty.Exp as P + +import Lens.Micro + +{- +Within each cluster (Labels), we do a topological sort using the edges in Graph +((a,b) means a before b in ordering). Then, we can simply cons them on top of each other. +Data.Graph (containers) has a nice topological sort. +-} + +(!??) :: (Ord a1, Show a1, HasCallStack) => M.Map a1 a2 -> a1 -> a2 +map' !?? key = case map' M.!? key of + Just x -> x + Nothing -> error ("error: map with keys " <> show (M.keys map') <> " does not contain key " <> show key) + +-- instance Show (Exists a) where +-- show (Exists x) = "exis" + +-- The caller of this function tells us what the result type should be +-- (namely, what it was before fusion), via an GroundsR. +-- Since fusion goes via an untyped ILP, during reconstruction we need to rebuild the program and temporarily +-- fulfill this contract: if something goes wrong during fusion or at the caller, bad things happen. +reconstruct :: forall op a. (MakesILP op, SimplifyOperation op) => GroundsR a -> Bool -> FusionGraph -> [ClusterLs] -> M.Map (Node Comp) [ClusterLs] -> Symbols op -> ReadDirM -> InplaceM -> PreOpenAcc (Clustered op) () a +reconstruct repr a b c d e f g = case openReconstruct a EnvNil b c d e f g of + Exists res -> expectType repr res + +reconstructF :: forall op a. (MakesILP op, SimplifyOperation op) => PreOpenAfun op () a -> Bool -> FusionGraph -> [ClusterLs] -> M.Map (Node Comp) [ClusterLs] -> Symbols op -> ReadDirM -> InplaceM -> PreOpenAfun (Clustered op) () a +reconstructF original a b c d e f g = case openReconstructF a EnvNil b c (Node 1 Nothing) d e f g of + Exists res -> expectFunTypeEqual original res + + +-- ordered list of labels +data ClusterL = ExecL [Node Comp] | NonExecL (Node Comp) + deriving Show + +foldC :: (Node Comp -> b -> b) -> b -> ClusterL -> b +foldC f x (ExecL ls) = foldr f x ls +foldC f x (NonExecL l) = f l x + +type ReadDirM = M.Map ReadEdge Int +type InplaceM = M.Map (Node GVal) (Node GVal) + +topSort :: Bool -> FusionGraph -> Nodes Comp -> ReadDirM -> [ClusterL] +topSort _ _ (S.toList -> [l]) _ = [ExecL [l]] -- If the cluster is empty. +topSort singletons (FusionGraph _ _ strictEdges dataflowEdges _ _ _) cluster readDirM = + if singletons then concatMap (map (ExecL . pure)) topsorteds else map ExecL topsorteds + where + buildGraph = + G.graphFromEdges + . map (\(a,b) -> (a,a,b)) + . M.toList + . flip (S.fold (\(x,i,y) -> M.adjust ((y,defaultDir):) (x,i))) edges + . M.fromList + . map (,[]) + . S.toList + + fedges, fpedges :: S.Set (Node Comp, Node GVal, Node Comp) + (fedges, fpedges) = S.partition (\(c1, _, c2) -> S.notMember (c1, c2) strictEdges) dataflowEdges + + -- Make a graph of all these labels and their incoming edges (for horizontal fusion)... + fpparents = S.unions $ S.map (\l -> (S.\\ cluster) $ S.map (\(w,_,_)->w) $ S.filter (\(_,_,r)->l==r) fpedges) cluster + parents = (S.\\ fpparents) $ S.unions $ S.map (\l -> (S.\\ cluster) $ S.map (\(w,_,_)->w) $ S.filter (\(_,_,r)->l==r) fedges ) cluster + parentsPlusEdges :: S.Set (Node Comp, Int, Node Comp) -- (Parent, Order, Target) + parentsPlusEdges = S.unions $ S.unions $ S.map (\l -> let relevantEdges = S.filter (\(w,_,r)->l==w && r `S.member` cluster) (fedges S.\\ fpedges) + -- TODO: why not just `ordersWithEdges = S.map (\e@(_ :->b) -> (l,readOrderOf e,b)) relevantEdges`? + orders = S.map readOrderOf relevantEdges + ordersWithEdges = S.map (\o -> S.map (\(_,_,r) -> (l,o,r)) $ S.filter (\e-> readOrderOf e == o) relevantEdges) orders + in ordersWithEdges) parents + + nodes = S.map (,defaultDir) cluster <> S.map (\(x,y,_)-> (x,y)) parentsPlusEdges + edges = S.union parentsPlusEdges $ S.map (\(a,_,b) -> (a,defaultDir,b)) dataflowEdges + (graph, getAdj, _) = buildGraph nodes + + -- .. split it into connected components and remove those parents from last step, + components = map (S.filter (\(l,_)->l `S.member` cluster) . S.fromList . map ((^._1) . getAdj) . T.flatten) $ G.components graph + -- and make a graph of each of them... + graphs = if singletons then [buildGraph $ S.map (,defaultDir) cluster] else map buildGraph components + -- .. and finally, topologically sort each of those to get the labels per cluster sorted on dependencies + topsorteds = map (\(graph', getAdj', _) -> map (view (_1 . _1) . getAdj') $ G.topSort graph') graphs + + readOrderOf :: HasCallStack => DataflowEdge -> Int + readOrderOf (_,b,r) = case readDirM M.!? (b,r) of + Just i -> i + Nothing -> error $ "can't get readorder " ++ show (b,r) + + defaultDir = 0 + + +openReconstruct :: (MakesILP op, SimplifyOperation op) + => Bool + -> Env aenv + -> FusionGraph + -> [ClusterLs] + -> M.Map (Node Comp) [ClusterLs] + -> Symbols op + -> ReadDirM + -> InplaceM + -> Exists (PreOpenAcc (Clustered op) aenv) +openReconstruct a b c d e f g h = + case openReconstruct' a b c d Nothing e f g h of + Left x -> x + Right{} -> error "TODO WALL: NON-EXHAUSTIVE PATTERN MATCH" +openReconstructF :: (MakesILP op, SimplifyOperation op) + => Bool + -> Env aenv + -> FusionGraph + -> [ClusterLs] + -> Node Comp + -> M.Map (Node Comp) [ClusterLs] + -> Symbols op + -> ReadDirM + -> InplaceM + -> Exists (PreOpenAfun (Clustered op) aenv) +openReconstructF a b c d l e f g h = + case openReconstruct' a b c d (Just l) e f g h of + Right x -> x + Left{} -> error "TODO WALL: NON-EXHAUSTIVE PATTERN MATCH" + +openReconstruct' :: forall op aenv. (MakesILP op, SimplifyOperation op) => Bool -> Env aenv -> FusionGraph -> [ClusterLs] -> Maybe (Node Comp) -> M.Map (Node Comp) [ClusterLs] -> Symbols op -> ReadDirM -> InplaceM -> Either (Exists (PreOpenAcc (Clustered op) aenv)) (Exists (PreOpenAfun (Clustered op) aenv)) +openReconstruct' singletons labelenv graph clusterslist mlab subclustersmap symbols readDirM inplaceM = + case mlab of + Just l -> Right $ makeASTF labelenv l + Nothing -> Left $ makeAST labelenv clusters + where + mkReindexPartial' :: Env env -> Env env' -> ReindexPartial Maybe env env' + mkReindexPartial' = mkReindexPartial inplaceM + + -- Make a tree of let bindings + + -- In mkFullGraph, we make sure that the bound body of a let will be in an earlier cluster. + -- Those are stored in the 'prev' argument. + -- Note also that we currently assume that the final cluster is the return argument: If all computations are relevant + -- and our analysis is sound, the return argument should always appear last. If not.. oops + makeAST :: forall env. Env env -> [ClusterL] -> Exists (PreOpenAcc (Clustered op) env) + makeAST _ [] = error "empty AST" + makeAST env [cluster] = case makeCluster env cluster of + Fold c args -> Exists $ Exec c $ unLabelOp args + InitFold o l args -> singleton l args o $ + \c args' -> + Exists $ Exec c (mapArgs (\(LOp a _ _) -> a) args') + EmptyFold -> Exists $ Return TupRunit + NotFold con -> case con of + SExe {} -> error "should be Fold/InitFold!" + SExe'{} -> error "should be Fold/InitFold!" + SUse se n be -> Exists $ Use se n be + SITE env' c t f -> case (makeAST env (subcluster t), makeAST env (subcluster f)) of + (Exists tacc, Exists facc) -> Exists $ tryBuildAcond + (fromJust $ reindexVar (mkReindexPartial' env' env) c) + tacc + facc + SWhl env' c b i u -> case (subcluster c, subcluster b) of + (findTopOfF -> c', findTopOfF -> b') -> case (makeASTF env c', makeASTF env b') of + (Exists cfun, Exists bfun) -> Exists $ tryBuildAwhile + u + cfun + bfun + (fromJust $ reindexVars (mkReindexPartial' env' env) i) + SLet {} -> error "let without scope" + SFun {} -> error "wrong type: function" + SBod {} -> error "wrong type: function" + SBlk {} -> error "wrong type: block" + SRet env' vars -> Exists $ Return (fromJust $ reindexVars (mkReindexPartial' env' env) vars) + SCmp env' expr -> Exists $ Compute (fromJust $ reindexExp (mkReindexPartial' env' env) expr) + SAtr msg env' t -> Exists $ Atrace msg (fromJust $ reindexArrayDescriptors (mkReindexPartial' env' env) t) + SAsr msg env' expr -> Exists $ Aassert msg (fromJust $ reindexExp (mkReindexPartial' env' env) expr) + SAsu env' expr -> Exists $ Aassume (fromJust $ reindexExp (mkReindexPartial' env' env) expr) + SAlc env' shr e sh -> Exists $ Alloc shr e (fromJust $ reindexVars (mkReindexPartial' env' env) sh) + SUnt env' evar -> Exists $ Unit (fromJust $ reindexVar (mkReindexPartial' env' env) evar) + SFen {} -> Exists $ Return TupRunit -- Fence with no futher terms does nothing + makeAST env (cluster:ctail) = + -- TODO: use guards to fuse these two identical cases + case makeCluster env cluster of + NotFold con + | SLet mylhs b u <- con -> + case makeAST env [NonExecL b] of + Exists bnd -> createLHS mylhs env $ \env' lhs -> + case makeAST env' ctail of + Exists scp + -> Exists $ tryBuildAlet lhs u bnd scp + | SFen env' set' <- con -> + case makeAST env ctail of + Exists next -> + Exists $ Fence (fromJust $ reindexIdxSet (mkReindexPartial' env' env) set') next + _ -> let res = makeAST env [cluster] in case cluster of + ExecL _ -> case (res, makeAST env ctail) of + (Exists exec@Exec{}, Exists scp) -> Exists $ Alet LeftHandSideUnit (shared TupRunit) exec scp + (Exists (Return TupRunit), Exists scp) -> Exists scp + _ -> error "nope" + NonExecL _ -> makeAST env ctail + + makeASTF :: forall env. Env env -> Node Comp -> Exists (PreOpenAfun (Clustered op) env) + makeASTF env l = case makeCluster env (NonExecL l) of + NotFold (SBod _l') -> case makeAST env (subcluster l) of + Exists acc -> Exists $ Abody acc + NotFold (SFun lhs l') -> createLHS lhs env $ \env' lhs' -> + case makeASTF env' l' of + Exists fun -> Exists $ Alam lhs' fun + NotFold _sym -> error $ "wrong type: acc" + _ -> error "not a notfold" + + findTopOfF :: [ClusterL] -> Node Comp + findTopOfF [] = error "empty list" + findTopOfF [NonExecL x] = x + findTopOfF (x@(NonExecL l):xs) = case symbols !?? l of + SBod _ -> findTopOfF xs + SFun _ l' -> findTopOfF $ filter (nonExecLNotEqual l') xs ++ [x] + _ -> error "should be a function" + -- findTopOfF $ filter (\(NonExecL l) -> Just l /= p) xs ++ [x] + findTopOfF _ = error "should be a function" + + nonExecLNotEqual :: Node Comp -> ClusterL -> Bool + nonExecLNotEqual l' = \case + NonExecL l'' -> l'' /= l' + ExecL{} -> error "TODO WALL: NON-EXHAUSTIVE PATTERN MATCH" + + -- do the topological sorting for each set + -- TODO: add 'backend-specific' edges to the graph for sorting, see 3.3.1 in the PLDI paper + clusters = concatMap (\case + Execs ls -> topSort singletons graph ls readDirM + NonExec l -> [NonExecL l]) clusterslist + subclusters = M.map (concatMap ( \case + Execs ls -> topSort singletons graph ls readDirM + NonExec l -> [NonExecL l])) subclustersmap + + subcluster :: HasCallStack => Node Comp -> [ClusterL] + subcluster l = subclusters !?? l + + makeCluster :: HasCallStack => Env env -> ClusterL -> FoldType op env + makeCluster env (ExecL ls) = + foldr1 (flip fuseCluster) + $ map ( \l -> case symbols !?? l of + SExe' env' args op -> + -- At first thought, this `fromJust` might error if we fuse an array away. + -- It does not: The array will still be in the environment, but after we finish + -- the `foldr1`, the input argument will dissapear. The output argument does not: + -- we clean that up in the SLV pass, if this was vertical fusion. If this is diagonal fusion, + -- it stays. + let args' = fromJust $ reindexLabelledArgsOp (mkReindexPartial inplaceM env' env) args + in + if isNoOp op (unLabelOp args') then + -- Remove operations that became a no-op by in-place updates. + -- For instance, 'map id xs ys' may become 'map id xs xs', + -- which is a no-op. + EmptyFold + else + InitFold op l args' + _ -> error "avoid this next refactor" -- c -> NotFold c + ) ls + makeCluster _ (NonExecL l) = NotFold $ symbols !?? l + + fuseCluster :: FoldType op env -> FoldType op env -> FoldType op env + fuseCluster EmptyFold f = f + fuseCluster f EmptyFold = f + fuseCluster (Fold cluster cargs) (InitFold op l largs) = + consCluster l largs op cargs cluster Fold + fuseCluster (InitFold op l largs) x = singleton l largs op $ \c cargs -> fuseCluster (Fold c cargs) x + fuseCluster Fold{} Fold{} = error "fuseCluster got non-leaf as second argument" -- Should never happen + fuseCluster NotFold{} _ = error "fuseCluster encountered NotFold" -- Should only occur in singleton clusters + fuseCluster _ NotFold{} = error "fuseCluster encountered NotFold" -- Should only occur in singleton clusters + +weakenAcc :: LeftHandSide s t env env' -> PreOpenAcc op env a -> PreOpenAcc op env' a +weakenAcc lhs = runIdentity . reindexAcc (weakenReindex $ weakenWithLHS lhs) + +-- | Internal datatype for `makeCluster`. + +data FoldType op env + = forall args. Fold (Clustered op args) (LabelledArgsOp op env args) + | forall args. InitFold (op args) (Node Comp) (LabelledArgsOp op env args) + | EmptyFold + | NotFold (Symbol op) + + +louttovar :: LabelledArgOp op env (Out sh e) -> LabelledArgOp op env (Var' sh) +louttovar (LOp a l b) = LOp (outvar a) (NotArr $ getLabelDeps l) b -- unsafe marker: maybe this NotArr ends up a problem? + +tryUpdateList :: (a -> Bool) -> (a -> a) -> [a] -> Maybe [a] +tryUpdateList _ _ [] = Nothing +tryUpdateList p f (x : xs) + | p x = Just $ f x : xs + | otherwise = tryUpdateList p f xs + + + + +consCluster :: forall env args extra op r + . MakesILP op + => Node Comp + -> LabelledArgsOp op env extra + -> op extra + -> LabelledArgsOp op env args + -> Clustered op args + -> (forall args'. Clustered op args' -> LabelledArgsOp op env args' -> r) + -> r +consCluster l lop op lcluster cluster k = singleton l lop op $ \c lop' -> + fuse + lop' + lcluster + lop' + lcluster + c + cluster + fuseVertically + $ flip k + +fuseVertically :: LabelledArgOp op env (Out sh e) -> LabelledArgOp op env (In sh e) -> LabelledArgOp op env (Var' sh) +fuseVertically + (LOp (ArgArray Out (ArrayR shr _) sh _) (getLabelDeps -> bs) b) + (LOp (ArgArray In _ _ _) (getLabelDeps -> bs') _) + = LOp (ArgVar $ groundToExpVar (shapeType shr) sh) (NotArr $ bs <> bs') b + +expectType :: HasCallStack => GroundsR t -> PreOpenAcc op env s -> PreOpenAcc op env t +expectType repr term + | Just Refl <- matchGroundsR repr $ groundsR term = term + | otherwise + = internalError $ fromString $ + "Result of fusion has incompatible type.\nExpected: " + ++ show (P.prettyTupR (const P.prettyGroundR) 0 repr) + ++ "\nActual: " + ++ show (P.prettyTupR (const P.prettyGroundR) 0 $ groundsR term) + +expectFunTypeEqual :: PreOpenAfun op1 () t -> PreOpenAfun op2 () s -> PreOpenAfun op2 () t +expectFunTypeEqual f1 f2 + | Just Refl <- matchFunType f1 f2 = f2 + | otherwise = internalError "Result of fused function has incompatible type" + +matchFunType :: PreOpenAfun op1 env t -> PreOpenAfun op2 env' s -> Maybe (t :~: s) +matchFunType (Alam lhs1 f1) (Alam lhs2 f2) + | Just Refl <- matchGroundsR (lhsToTupR lhs1) (lhsToTupR lhs2) + , Just Refl <- matchFunType f1 f2 + = Just Refl +matchFunType (Abody b1) (Abody b2) = matchGroundsR (groundsR b1) (groundsR b2) +matchFunType _ _ = Nothing + +tryBuildAcond :: HasCallStack => ExpVar env PrimBool -> PreOpenAcc op env t1 -> PreOpenAcc op env t2 -> PreOpenAcc op env t1 +tryBuildAcond cond true false + | Just Refl <- matchGroundsR (groundsR true) (groundsR false) + = Acond cond true false + | otherwise + = internalError "Cannot reconstruct Acond: branches have incompatible types" + +tryBuildAlet + :: HasCallStack + => GLeftHandSide t1 env1 env2 + -> Uniquenesses t1 + -> PreOpenAcc op env1 t2 + -> PreOpenAcc op env2 s + -> PreOpenAcc op env1 s +tryBuildAlet lhs u bnd next + | Just Refl <- matchGroundsR (lhsToTupR lhs) (groundsR bnd) + = Alet lhs u bnd next + | otherwise + = internalError "Cannot reconstruct Alet: left hand side and binding have incompatible types" + +tryBuildAwhile + :: Uniquenesses a + -> PreOpenAfun op env c + -> PreOpenAfun op env s + -> GroundVars env a + -> PreOpenAcc op env a +tryBuildAwhile u c@(Alam lhsCond (Abody cond)) s@(Alam lhsStep (Abody step)) initial + | Just Refl <- matchGroundsR tp $ lhsToTupR lhsCond + , Just Refl <- matchGroundsR (TupRsingle $ GroundRscalar $ scalarType @PrimBool) $ groundsR cond + , Just Refl <- matchGroundsR tp $ lhsToTupR lhsStep + , Just Refl <- matchGroundsR tp $ groundsR step + = Awhile u c s initial + where + tp = varsType initial +tryBuildAwhile _ _ _ _ = internalError "Cannot reconstruct Awhile: condition or step has invalid type" diff --git a/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/ConstraintLanguage.hs b/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/ConstraintLanguage.hs new file mode 100644 index 000000000..2f1eb0548 --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/ConstraintLanguage.hs @@ -0,0 +1,52 @@ +-- | A domain-specific language for the fusion ILP's constraints. +module Data.Array.Accelerate.Trafo.Partitioning.ILP.ConstraintLanguage (Constraint (..)) where + +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Labels (Comp, GVal, InplacePath, Node, ReadEdge, WriteEdge) +import Prelude hiding (pi) + +-- | A property of the fusion problem, to be lowered into linear constraints. +data Constraint + = -- | @pi i < pi j@: @i@ lands in a strictly earlier cluster than @j@. + ClusterBefore (Node Comp) (Node Comp) + | -- | @i@ and @j@ land in different clusters. + DifferentCluster (Node Comp) (Node Comp) + | -- | Iff every listed edge is fused, the buffer is not manifest. + NotManifestIfAllFused (Node GVal) [(Node Comp, Node Comp)] + | -- | Fused: @i@ and @j@ share a cluster. Unfused: @i@ comes strictly before @j@. + FusibleOrder (Node Comp) (Node Comp) + | -- | Fusing @w@ and @r@ implies that the write direction of @w@ is equal to the read direction of @r@. + FusionDirection (Node Comp) (Node GVal) (Node Comp) + | -- | Bounds @l@'s cluster index to be at most @maxCluster@. + WithinClusterCount (Node Comp) + | -- | If @p@ is in place, then the buffers must be manifest. + OnManifestIfInPlace InplacePath + | -- | In-place reuse requires read and write directions to be equal. + InPlaceDirection InplacePath + | -- | In-place use requires the read buffer to be dead by the time the write buffer is written to. + InPlaceCluster InplacePath + | -- | In-place use requires the reading and writing computations to land in the same cluster. + AcrossClusterSame InplacePath + | -- | At most one of these paths may reuse the buffer they all read from. + AtMostOneReader [InplacePath] + | -- | At most one of these paths may reuse the buffer they all write to. + AtMostOneWriter [InplacePath] + | -- | Unless one of the writers consumes this read in place, the buffer must stay live past the reading cluster. + ReadAliveThroughWriters ReadEdge [WriteEdge] + | -- | The @(buffer, consumer)@ edges out of one computation. Costs one per distinct cluster that reads from the buffer, so the solver is paid to fuse the consumers together. + HorizontalReadCost [(Node GVal, Node Comp)] + | -- | The buffer is manifest. + Manifest (Node GVal) + | -- | The buffer may not be used in place. + NoInPlace (Node GVal) + | -- | The computation reads and writes at the same fold size. + SameFoldSize (Node Comp) + | -- | The computation establishes a new fold size. + NewFoldSize (Node Comp) + | -- | If @w@ and @c@ are fused, @c@ reads at the fold size @w@ writes at. + SameFoldSizeIfFused (Node Comp) (Node Comp) + | -- | Every listed read and write happens in the same direction. + SameDirection [ReadEdge] [WriteEdge] + | -- | As 'SameDirection', but with the addition of the direction @c@. + PinnedDirection (Node Comp) [ReadEdge] [WriteEdge] + | -- | If the buffer is manifest, it is written in a negative direction. + NegativeDirIfManifest WriteEdge diff --git a/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/Encoding.hs b/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/Encoding.hs new file mode 100644 index 000000000..3b73430b2 --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/Encoding.hs @@ -0,0 +1,356 @@ +{-# LANGUAGE BlockArguments #-} +module Data.Array.Accelerate.Trafo.Partitioning.ILP.Encoding where + +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Graph hiding (graph, constraints, bounds) +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Labels (Node, parent, Nodes, Comp, GVal, DataflowEdge, StrictEdge, ReadEdge, InplacePath, WriteEdge) +import Data.Array.Accelerate.Trafo.Partitioning.ILP.LinearConstraint +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Solver hiding (finalize) + +import Data.List (groupBy, sortOn, foldl') +import Prelude hiding (sum, pi, read ) + +import qualified Data.Map as M + +-- In this file, order very often subly does matter. +-- To keep this clear, we use S.Set whenever it does not, +-- and [] only when it does. It's also often efficient +-- by removing duplicates. +import qualified Data.Set as S +import Data.Function ( on ) +import Lens.Micro ((^.), _1 ) +import Lens.Micro.Extras ( view ) +import Data.Maybe (fromJust, mapMaybe ) +import Data.Array.Accelerate.Trafo.Partitioning.ILP.ConstraintLanguage (Constraint(..)) +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Lower (LowerEnv(..), lowerAll) + +data Objective + -- Old fusion only objectives: + = NumClusters -- ^ Minimise the number of clusters. + | ArrayReads -- ^ Minimise the number of array reads. + | ArrayReadsWrites -- ^ Minimise the number of array reads and writes. + | IntermediateArrays -- ^ Minimise the number of intermediate arrays. + | FusedEdges -- ^ Minimise the number of unfused edges. + | Everything -- ^ Minimise the number of clusters and array reads/writes. + -- Old fusion objectives with in-place updates as secondary objective: + | ArrayReads' -- ^ Minimise the number of array reads, then maximise the number of in-place updates. + | ArrayReadsWrites' -- ^ Minimise the number of array reads and writes, then maximise the number of in-place updates. + | IntermediateArrays' -- ^ Minimise the number of intermediate arrays, then maximise the number of in-place updates. + | FusedEdges' -- ^ Minimise the number of unfused edges, then maximise the number of in-place updates. + -- Custom objectives: + | MemoryUsage -- ^ Minimise the number of distinct intermediate arrays. + | MemoryUsage' -- ^ Version of `MemoryUsage` that prioritizes fusion when two solutions would otherwise have the same costs. + deriving (Show, Bounded, Enum, Eq, Ord) + +-- Makes the ILP. Note that this function 'appears' to ignore the Node levels completely! +-- We could add some assertions, but if all the input is well-formed (no labels, constraints, etc +-- that reward putting non-siblings in the same cluster) this is fine: We will interpret 'cluster 3' +-- with parents `Nothing` as a different cluster than 'cluster 3' with parents `Just 5`. +makeILP :: forall op. MakesILP op => Objective -> FusionILP op -> ILP +makeILP obj (FusionILP graph constraints bounds) = + ILP minMax objFun loweredConstraints (graphBounds <> bounds) (Constants n m) + where + graphBounds = fusionBounds <> inPlaceBounds + + lowered :: (LinearConstraint, Bounds, Expression) + lowered = lowerAll (LowerEnv n) $ finalize @op graph <> fusionConstraints <> inPlaceConstraints <> constraints + + (loweredConstraints, loweredBounds, loweredCost) = lowered + + fusionConstraints = strictAcyclicConstraints + <> infusibleConstraints + <> manifestConstraints + <> fusibleAcyclicConstraints + <> fusionDirectionConstraints + <> numberOfClustersConstraints + <> horizontalReadCostConstraints + <> manifestValueConstraints + <> noInPlaceConstraints + + fusionBounds = piB <> fusedB <> manifestB <> loweredBounds + + inPlaceConstraints = if enableIU + then onManifestConstraints + <> inPlaceDirectionConstraints + <> inPlaceClusterConstraints + <> acrossClusterConstraints + <> atMostOneReaderConstraints + <> atMostOneWriterConstraints + <> readAliveThroughWritersConstraints + else mempty + + inPlaceBounds = if enableIU then pimaxB <> inplaceB else mempty + + ---------------------------------------------------------------------------- + -- Utils: + ---------------------------------------------------------------------------- + + compN :: Nodes Comp + compN = graph^.computationNodes + + buffN :: Nodes GVal + buffN = graph^.valueNodes + + n :: Int + n = S.size compN + + m :: Int + m = S.size buffN + + ---------------------------------------------------------------------------- + -- Fusion: + ---------------------------------------------------------------------------- + dataflowE :: S.Set DataflowEdge + dataflowE = graph^.dataflowEdges + + strictE :: S.Set StrictEdge + strictE = graph^.strictEdges + + fusibleE, infusibleE :: S.Set DataflowEdge + (fusibleE, infusibleE) = graph^.fusionEdges + + fusibleE', infusibleE' :: S.Set (Node Comp, Node Comp) + fusibleE' = S.map (\(i,_,j) -> (i,j)) fusibleE + infusibleE' = S.map (\(i,_,j) -> (i,j)) infusibleE + + -- objective function that maximises the number of edges we fuse, and minimises the number of array reads if you ignore horizontal fusion + numberOfUnfusedEdges = foldMap fused fusibleE' + + -- A cost function that doesn't ignore horizontal fusion. + -- Idea: Each node $x$ with $n$ outgoing edges gets $n$ extra variables. + -- Each edge (fused or not) $(x,y)$ will require that one of these variables is equal to $pi y$. + -- The number of extra variables that are equal to 0 (the thing you maximise) is exactly equal to n - the number of clusters that read from $x$. + -- Then, we also need n^2 intermediate variables just to make these disjunction of conjunctions + -- note, it's only quadratic in the number of consumers of a specific array. + -- We also check for the 'order': horizontal fusion only happens when the two fused accesses are in the same order. + horizontalReadCostConstraints :: [Constraint] + horizontalReadCostConstraints = [HorizontalReadCost consumers + | computation <- S.toList compN + , let consumers = S.toList . S.map (\(_,b,c) -> (b,c)) $ S.filter (\(c,_,_) -> c == computation) fusibleE] + + numberOfReads = loweredCost .+. numberOfUnfusedEdges + + -- objective function that maximises the number of fused away arrays, and thus minimises the number of array writes + -- using .-. instead of notB to factor the constants out of the cost function; if we use (1 - manifest l) as elsewhere Gurobi thinks the 1 is a variable name + numberOfManifestArrays = foldl' (\e b -> e .-. manifest b) (int 0) buffN + + -- objective function that minimises the total number of array reads + writes + numberOfArrayReadsWrites = numberOfReads .+. numberOfManifestArrays + + -- objective function that minimises the number of clusters only works if the constraint below it is used! + -- NOTE: this does not work remotely as well as you'd hope, because the ILP outputs clusters that get split afterwards. + -- This includes independent array operations, which might not be safe to fuse and get no real benefit from fusing, + -- but also includes independent alloc, compute, etc nodes which we don't even want to count in the first place! + -- It's possible to also give each array operation a 'exec-pi' variable, and change this to minimise the maximum of + -- these exec-pi values, in which case we are only left with the independent array operations problem. + -- To eliminate that one too, we'd need n^2 edges. + numberOfClusters = maxCluster + + -- pi_i < pi_j for all strict edges + strictAcyclicConstraints = map (uncurry ClusterBefore) $ S.toList strictE + + -- x_ij == 1 for all infusible edges + infusibleConstraints = map (uncurry DifferentCluster) $ S.toList infusibleE' + + -- forall b, iff all (w,b,r) are fused, then b is not manifest. + manifestConstraints = map (uncurry NotManifestIfAllFused) . M.toList $ foldl (flip \(i,b,j) -> M.insertWith (<>) b [(i,j)]) M.empty dataflowE + + -- x_ij <= pi_j - pi_i <= n*x_ij for all fusible edges + fusibleAcyclicConstraints = map (uncurry FusibleOrder) $ S.toList fusibleE' + + -- if (w,b,r) is fused, then d_wb == d_br + fusionDirectionConstraints = map (\(w,b,r) -> FusionDirection w b r) $ S.toList fusibleE + + -- removing this makes the ILP slightly smaller, but disables the use of this cost function + numberOfClustersConstraints = case obj of + NumClusters -> map WithinClusterCount $ S.toList compN + Everything -> map WithinClusterCount $ S.toList compN + _ -> [] + + -- b is manifest, for every program result + manifestValueConstraints = map Manifest $ S.toList $ graph^.manifestValues + + -- b stays live past every cluster + noInPlaceConstraints = map NoInPlace $ S.toList $ graph^.noInplaceValues + + -- 0 <= pi_i <= n + piB = foldMap (\i -> lowerUpper 0 (Pi i) n) compN + + -- 0 <= x_ij <= 1 + fusedB = foldMap (binary . uncurry Fused) $ S.map (\(i,_,j) -> (i,j)) dataflowE + + -- 0 <= m_i <= 1 + manifestB = foldMap (binary . IsManifest) buffN + + + ---------------------------------------------------------------------------- + -- In-place updates: + ---------------------------------------------------------------------------- + + -- NOTE: Here we need to use the number of non-inplace updates, because using 1 - x is not liked by the solver. + -- TODO: Either make the internal structure more robust by only allowing variables to appear once per expression/constraint, + -- or add a simplification step to ensure this is the case. + + readE :: S.Set ReadEdge + readE = graph^.readEdges + + inplaceP :: S.Set InplacePath + inplaceP = M.keysSet $ graph^.inplacePaths + + -- Number of in-place updates: + numberOfNonInplaceUpdates = foldMap inplace inplaceP + + -- If inplace p, then manifest b1 and manifest b2 + onManifestConstraints :: [Constraint] + onManifestConstraints = [OnManifestIfInPlace p | p <- S.toList inplaceP] + + -- If inplace p, then d_br == d_wb + inPlaceDirectionConstraints :: [Constraint] + inPlaceDirectionConstraints = [InPlaceDirection p | p <- S.toList inplaceP] + + -- If inplace p, then pimax b1 >= pi c2 + inPlaceClusterConstraints :: [Constraint] + inPlaceClusterConstraints = [InPlaceCluster p | p <- S.toList inplaceP] + + -- If inplace p, then c1 == c2 + acrossClusterConstraints :: [Constraint] + acrossClusterConstraints = [AcrossClusterSame p | p <- S.toList inplaceP] + + readerGroups = foldl (flip \p@((b,_),_) -> M.insertWith (<>) b [p]) M.empty inplaceP + writerGroups = foldl (flip \p@(_,(_,b)) -> M.insertWith (<>) b [p]) M.empty inplaceP + + -- Forall b, at most one inplace p + atMostOneReaderConstraints :: [Constraint] + atMostOneReaderConstraints = [AtMostOneReader ps | ps <- M.elems readerGroups] + + -- Forall b, at most one inplace p + atMostOneWriterConstraints :: [Constraint] + atMostOneWriterConstraints = [AtMostOneWriter ps | ps <- M.elems writerGroups] + + -- Group inplace paths by read edge: + readM = foldl (flip \(r,w) -> M.insertWith (<>) r [w]) M.empty inplaceP + + -- Iff inplace p, then pi c1 <= pimax b1 + -- Iff not inplace p, then pi c1 + 1 <= pimax b1 + readAliveThroughWritersConstraints :: [Constraint] + readAliveThroughWritersConstraints = [ReadAliveThroughWriters r (M.findWithDefault [] r readM) | r <- S.toList readE] + + -- TODO: Maybe add a constraint that c2 is the first writer to b2? + -- This would make sense because the graph doesn't acctually enforce there is only one writer per buffer. + -- For most cases there shouldn't be more than 2 writers, one of which is a let-binding, so no issues arise without this constraint. + -- However, a mutable computation would create a third writer, which would be a problem. + + -- 0 <= pimax_b + pimaxB = foldMap (\b -> lowerUpper 0 (PiMax b) (n+5)) buffN + + -- inplace b1 b2 \in {0, 1} + inplaceB = foldMap (\((b1,c1),(c2,b2)) -> binary $ InPlace b1 c1 c2 b2) inplaceP + + ---------------------------------------------------------------------------- + -- Objective function + ---------------------------------------------------------------------------- + + -- Since we want all clusters to have one 'iteration size', the final objFun should + -- take care to never reward 'fusing' disjoint clusters, and then slightly penalise it. + -- The alternative is O(n^2) edges, so this is worth the trouble! + -- + -- In the future, maybe we want this to be backend-dependent (add to MakesILP). + -- Also future: add @IVO's IPU reward here. + (enableIU, minMax, objFun) = case obj of + NumClusters -> (False, Minimise, numberOfClusters) + ArrayReads -> (False, Minimise, numberOfReads) + ArrayReadsWrites -> (False, Minimise, numberOfArrayReadsWrites) + IntermediateArrays -> (False, Minimise, numberOfManifestArrays) + FusedEdges -> (False, Minimise, numberOfUnfusedEdges) + Everything -> (False, Minimise, numberOfClusters .+. numberOfArrayReadsWrites) -- arrayreadswrites already indictly includes everything else + ArrayReads' -> (True, Minimise, (int m .*. numberOfReads) .+. numberOfNonInplaceUpdates) + ArrayReadsWrites' -> (True, Minimise, (int m .*. numberOfArrayReadsWrites) .+. numberOfNonInplaceUpdates) + IntermediateArrays' -> (True, Minimise, (int m .*. numberOfManifestArrays) .+. numberOfNonInplaceUpdates) + FusedEdges' -> (True, Minimise, (int m .*. numberOfUnfusedEdges) .+. numberOfNonInplaceUpdates) + MemoryUsage -> (True, Minimise, numberOfManifestArrays .+. numberOfNonInplaceUpdates) + MemoryUsage' -> (True, Minimise, (int (m+1) .*. numberOfManifestArrays) .+. (int m .*. numberOfNonInplaceUpdates)) -- We want to prioritise solutions that use fusion, so the weight of fusion is increased by a small factor. + + +-- | Extract the read directions from the ILP solution. +interpretReadDirs :: Solution -> M.Map ReadEdge Int +interpretReadDirs = M.fromList . mapMaybe (_1 fromReadDir) . M.toList + where + fromReadDir :: Var -> Maybe ReadEdge + fromReadDir (ReadDir b c) = Just (b, c) + fromReadDir _ = Nothing + +-- | Extract the write directions from the ILP solution. +interpretWriteDirs :: Solution -> M.Map WriteEdge Int +interpretWriteDirs = M.fromList . mapMaybe (_1 fromWriteDir) . M.toList + where + fromWriteDir :: Var -> Maybe WriteEdge + fromWriteDir (WriteDir c b) = Just (c, b) + fromWriteDir _ = Nothing + +-- | Extract the top-level clusters and the sub-scoped clusters from the ILP +-- solution. +interpretClusters :: Solution -> ([Nodes Comp], M.Map (Node Comp) [Nodes Comp]) +interpretClusters sol = do + let piVars = mapMaybe (_1 fromPi) (M.toList sol) -- Take the Pi variables. + let scopedPiVars = partition (^._1.parent) piVars -- Partition them by their parent (i.e. the scope they are in). + let clusteredPiVars = map (partition snd) scopedPiVars -- Partition them again by their cluster (i.e. the value of the variable). + let scopedClusters = map (map $ S.fromList . map fst) clusteredPiVars -- Remove the cluster numbers and convert each cluster to a set. + let topClusters = head scopedClusters -- The first scope in the list is the top-level scope. + let subScopedClusters = tail scopedClusters -- The rest of the scopes are arbitrarily nested sub-scopes. + let subScopedClustersM = M.fromList $ map (\s -> (scopeLabel s, s)) subScopedClusters + (topClusters, subScopedClustersM) + where + fromPi :: Var -> Maybe (Node Comp) + fromPi (Pi l) = Just l + fromPi _ = Nothing + + scopeLabel :: [Nodes Comp] -> Node Comp + scopeLabel = fromJust . view parent . S.findMin . head + +-- | `groupBy` except it's equivalent to SQL's `GROUP BY` clause. +partition :: Ord b => (a -> b) -> [a] -> [[a]] +partition f = groupBy ((==) `on` f) . sortOn f + +interpretInplaceUpdates :: Solution -> M.Map (Node GVal) (Node GVal) +interpretInplaceUpdates sol = M.map firstInChain inplaceM + where + -- Map from buffer to the buffer that will replace it. + inplaceM = M.fromList $ mapMaybe fromInPlace $ M.toList sol + + fromInPlace :: (Var, Int) -> Maybe (Node GVal, Node GVal) + fromInPlace (InPlace b1 _ _ b2, v) | v == 0 = Just (b2, b1) + fromInPlace _ = Nothing + + firstInChain :: Node GVal -> Node GVal + firstInChain b = maybe b firstInChain (M.lookup b inplaceM) + +-- | Cluster labels, distinguishing between execute and non-execute labels. +data ClusterLs = Execs (Nodes Comp) | NonExec (Node Comp) + deriving (Eq, Show) + +-- I think that only `let`s can still be in the same cluster as `exec`s, +-- and their bodies should all be in earlier clusters already. +-- Simply make one cluster per let, before the cluster with execs. +splitExecs :: ([Nodes Comp], M.Map (Node Comp) [Nodes Comp]) -> Symbols op -> ([ClusterLs], M.Map (Node Comp) [ClusterLs]) +splitExecs (xs, xM) symbolM = (f xs, M.map f xM) + where + f :: [Nodes Comp] -> [ClusterLs] + f = concatMap (\ls -> filter (/= Execs mempty) $ map NonExec (S.toList $ S.filter isBeforeExec ls) ++ [Execs (S.filter isExec ls)] ++ afterexecs ls) + + isExec :: Node Comp -> Bool + isExec l = case symbolM M.!? l of + Just SExe {} -> True + Just SExe'{} -> True + _ -> False + isBeforeExec l = case symbolM M.!? l of + Just SLet{} -> True + _ -> False + isAfterExec l = case symbolM M.!? l of + Just SExe{} -> False + Just SExe'{} -> False + Just SLet{} -> False + Nothing -> False + _ -> True + -- Tests say that this happens, and that it's correct anyway, but I'm unsure why. + -- The reason I doubt is because if multiple non-exec, non-lhs nodes are here, the current reconstruction code + -- (I think) ignores all but the last one. + afterexecs ls = let xs' = map NonExec (S.toList $ S.filter isAfterExec ls) in if length xs' > 1 then xs' {-error "dunno what this means"-} else xs' diff --git a/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/Graph.hs b/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/Graph.hs new file mode 100644 index 000000000..b838bb292 --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/Graph.hs @@ -0,0 +1,1305 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE FunctionalDependencies #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE TypeFamilyDependencies #-} +{-# LANGUAGE ViewPatterns #-} +{-# OPTIONS_GHC -Wno-orphans #-} +module Data.Array.Accelerate.Trafo.Partitioning.ILP.Graph where + +import Prelude hiding ( init, reads ) + +-- Accelerate imports +import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.IdxSet (IdxSet(..)) +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.AST.Operation hiding (Var) +import Data.Array.Accelerate.Analysis.Hash.Exp +import Data.Array.Accelerate.Analysis.Match +import Data.Array.Accelerate.Array.Buffer +import Data.Array.Accelerate.Error +import Data.Array.Accelerate.Representation.Shape +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Trafo.Operation.LiveVars +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Labels +import Data.Array.Accelerate.Trafo.Partitioning.ILP.ConstraintLanguage (Constraint) +import Data.Array.Accelerate.Trafo.Partitioning.ILP.LinearConstraint +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Solver +import Data.Array.Accelerate.Type + +-- Data structures +import Data.Set (Set) +import Data.Map (Map) +import Data.Text (Text) +import qualified Data.Set as S +import qualified Data.Map as M + +import Lens.Micro +import Lens.Micro.Mtl + +import Control.Monad.State.Strict (State, runState) +import Data.Foldable (Foldable (foldr'), traverse_, toList) +import Data.Kind (Type) + + +-------------------------------------------------------------------------------- +-- Fusion Graph +-------------------------------------------------------------------------------- + +-- For backwards compatitibility: +pattern (:->) :: Node Comp -> Node Comp -> DataflowEdge +pattern w :-> r <- (w,_,r) + + +-- | Program graph. +-- +-- The graph consists of read/write edges, strict ordering edges and fusible +-- edges. +-- +-- The read/write edges represent a read or write relation between a buffer and +-- a computation. In the ILP, these edges get a variable indicating in which +-- order a computation reads or writes an array in the buffer. Some of these +-- edges are duplicated, so we use a smart constructor to make sure they get the +-- same ILP variable. +-- +-- The strict ordering edges enforce a strict ordering between two computations. +-- This ordering can be due to any number of reasons, but in most cases it is to +-- prevent race conditions between two computations. +-- +-- The data-flow edges represent a flow of data between two computations over a +-- buffer. These edges are used to determine which computations can be fused. +-- Data-flow edges that are also strict edges are infusible. +-- Data-flow edges that are not strict edges are fusible. +-- +-- From the sets of data-flow and strict ordering edges we can derive: +-- 1. The set of write edges. @S.map (\(w,b,_) -> (w,b)) _dataflowEdges@ +-- 2. The set of read edges. @S.map (\(_,b,r) -> (w,b)) _dataflowEdges@ +-- 3. The set of fusible edges. @S.filter (\(w,_,r) -> S.notMember (w,r) _strictEdges) _dataflowEdges@ +-- 4. The set of infusible edges. @S.filter (\(w,_,r) -> S.member (w,r) _strictEdges) _dataflowEdges@ +-- +-- The latter two computations may be combined as such: +-- +-- @ +-- (fusible, infusible) = S.partition (\(w,_,r) -> S.notMember (w,r) _strictEdges) _dataflowEdges +-- @ +-- +data FusionGraph = FusionGraph -- TODO: Use hashmaps and hashsets in production. + { _compNodes :: Set (Node Comp) -- ^ Computation nodes. + , _valueNodes :: Set (Node GVal) -- ^ Value nodes. + , _strictEdges :: Set StrictEdge -- ^ Edges that enforce strict ordering. + , _dataflowEdges :: Set DataflowEdge -- ^ Edges that represent data-flow. + , _inplacePaths :: Map InplacePath Number -- ^ Summary paths between buffers for in-place updates + their weight. + , _manifestValues :: Set (Node GVal) -- ^ Values forced manifest. + , _noInplaceValues :: Set (Node GVal) -- ^ Values that must stay live past every cluster. + } + +instance Semigroup FusionGraph where + (<>) :: FusionGraph -> FusionGraph -> FusionGraph + (<>) (FusionGraph c1 v1 s1 d1 i1 m1 p1) (FusionGraph c2 v2 s2 d2 i2 m2 p2) + = FusionGraph (c1 <> c2) (v1 <> v2) (s1 <> s2) (d1 <> d2) (i1 <> i2) (m1 <> m2) (p1 <> p2) + +instance Monoid FusionGraph where + mempty :: FusionGraph + mempty = FusionGraph mempty mempty mempty mempty mempty mempty mempty + +-- | Class for types that contain a fusion graph. +-- +-- This is a manually written version of what microlens-th would generate when +-- using @makeClassy@. +class HasFusionGraph g where + fusionGraph :: Lens' g FusionGraph + + computationNodes :: Lens' g (Set (Node Comp)) + computationNodes = fusionGraph.computationNodes + + valueNodes :: Lens' g (Set (Node GVal)) + valueNodes = fusionGraph.valueNodes + + strictEdges :: Lens' g (Set StrictEdge) + strictEdges = fusionGraph.strictEdges + + dataflowEdges :: Lens' g (Set DataflowEdge) + dataflowEdges = fusionGraph.dataflowEdges + + inplacePaths :: Lens' g (Map InplacePath Number) + inplacePaths = fusionGraph.inplacePaths + + manifestValues :: Lens' g (Set (Node GVal)) + manifestValues = fusionGraph.manifestValues + + noInplaceValues :: Lens' g (Set (Node GVal)) + noInplaceValues = fusionGraph.noInplaceValues + +-- | Base instance of 'HasFusionGraph' for 'FusionGraph'. +-- +-- This instance cannot make use of lenses defined in 'HasFusionGraph' because +-- it is the base instance and would otherwise cause a loop. +instance HasFusionGraph FusionGraph where + fusionGraph :: Lens' FusionGraph FusionGraph + fusionGraph = id + + computationNodes :: Lens' FusionGraph (Set (Node Comp)) + computationNodes f s = f (_compNodes s) <&> \ns -> s{_compNodes = ns} + + valueNodes :: Lens' FusionGraph (Set (Node GVal)) + valueNodes f s = f (_valueNodes s) <&> \ns -> s{_valueNodes = ns} + + strictEdges :: Lens' FusionGraph (Set StrictEdge) + strictEdges f s = f (_strictEdges s) <&> \es -> s{_strictEdges = es} + + dataflowEdges :: Lens' FusionGraph (Set DataflowEdge) + dataflowEdges f s = f (_dataflowEdges s) <&> \es -> s{_dataflowEdges = es} + + inplacePaths :: Lens' FusionGraph (Map InplacePath Number) + inplacePaths f s = f (_inplacePaths s) <&> \ps -> s{_inplacePaths = ps} + + manifestValues :: Lens' FusionGraph (Set (Node GVal)) + manifestValues f s = f (_manifestValues s) <&> \ns -> s{_manifestValues = ns} + + noInplaceValues :: Lens' FusionGraph (Set (Node GVal)) + noInplaceValues f s = f (_noInplaceValues s) <&> \ns -> s{_noInplaceValues = ns} + +insertComputation :: (HasCallStack, HasFusionGraph g) => Node Comp -> g -> g +insertComputation c = computationNodes %~ S.insert c + +insertValue :: (HasCallStack, HasFusionGraph g) => Node GVal -> g -> g +insertValue v = valueNodes %~ S.insert v + +-- | Insert a strict relation between two computations. +insertStrict :: (HasCallStack, HasFusionGraph g) => StrictEdge -> g -> g +insertStrict (c1, c2) g + | c1 == c2 = internalError "reflexive edge" + | c1^.parent /= c2^.parent = internalError "different scopes" + | S.member (c2, c1) (g^.strictEdges) = internalError "cyclic edge" + | otherwise = g & strictEdges %~ S.insert (c1, c2) + +-- | Insert a fusible data-flow edge between two computations. +insertFusible :: (HasCallStack, HasFusionGraph g) => DataflowEdge -> g -> g +insertFusible (c1, b, c2) g + | c1 == c2 = internalError "reflexive edge" + | c1^.parent /= c2^.parent = internalError "different scopes" + | S.member (c2, c1) (g^.strictEdges) = internalError "cyclic edge" + | otherwise = g & dataflowEdges %~ S.insert (c1, b, c2) + +-- | Insert an infusible data-flow edge between two computations. +insertInfusible :: (HasCallStack, HasFusionGraph g) => DataflowEdge -> g -> g +insertInfusible (c1, b, c2) g + | c1 == c2 = internalError "reflexive edge" + | c1^.parent /= c2^.parent = internalError "different scopes" + | S.member (c2, c1) (g^.strictEdges) = internalError "cyclic edge" + | otherwise = g & dataflowEdges %~ S.insert (c1, b, c2) + & strictEdges %~ S.insert (c1, c2) + +-- | Gets the set of fusible edges. +fusibleEdges :: HasFusionGraph g => SimpleGetter g (Set FusibleEdge) +fusibleEdges = to (\g -> S.filter (\(w, _, r) -> S.notMember (w, r) (g^.strictEdges)) (g^.dataflowEdges)) + +-- | Gets the set of infusible edges. +infusibleEdges :: HasFusionGraph g => SimpleGetter g (Set InfusibleEdge) +infusibleEdges = to (\g -> S.filter (\(w, _, r) -> S.member (w, r) (g^.strictEdges)) (g^.dataflowEdges)) + +-- | Gets the set of fusible and infusible edges. +fusionEdges :: HasFusionGraph g => SimpleGetter g (Set FusibleEdge, Set InfusibleEdge) +fusionEdges = to (\g -> S.partition (\(w, _, r) -> S.notMember (w, r) (g^.strictEdges)) (g^.dataflowEdges)) + +-- | Gets the set of strict edges that are not data-flow edges. +orderEdges :: HasFusionGraph g => SimpleGetter g (Set StrictEdge) +orderEdges = to (\g -> let dataflowEdges' = S.map (\(w,_,r) -> (w,r)) (g^.dataflowEdges) + in S.filter (\(w, r) -> S.notMember (w, r) dataflowEdges') (g^.strictEdges)) + +readEdges :: HasFusionGraph g => SimpleGetter g (Set ReadEdge) +readEdges = to (\g -> S.map (\(_, b, r) -> (b, r)) (g^.dataflowEdges)) + +writeEdges :: HasFusionGraph g => SimpleGetter g (Set WriteEdge) +writeEdges = to (\g -> S.map (\(w, b, _) -> (w, b)) (g^.dataflowEdges)) + +-- | Gets the input edges of a computations. +inputEdgesOf :: HasFusionGraph g => Node Comp -> SimpleGetter g (Set ReadEdge) +inputEdgesOf c = to (\g -> S.filter (\(_, r) -> r == c) (g^.readEdges)) + +-- | Gets the output edges of a computations. +outputEdgesOf :: HasFusionGraph g => Node Comp -> SimpleGetter g (Set WriteEdge) +outputEdgesOf c = to (\g -> S.filter (\(w, _) -> w == c) (g^.writeEdges)) + +-- | Gets the read edges of a buffer. +readEdgesOf :: HasFusionGraph g => Node GVal -> SimpleGetter g (Set ReadEdge) +readEdgesOf b = to (\g -> S.filter (\(b',_) -> b' == b) (g^.readEdges)) + +-- | Gets the write edges of a buffer. +writeEdgesOf :: HasFusionGraph g => Node GVal -> SimpleGetter g (Set WriteEdge) +writeEdgesOf b = to (\g -> S.filter (\(_,b') -> b' == b) (g^.writeEdges)) + +-- computationNodes :: HasFusionGraph g => SimpleGetter g (Set (Node Comp)) +-- computationNodes = to (\g -> S.foldr (\(w,_,r) -> S.insert w . S.insert r) S.empty (g^.dataflowEdges) +-- <> S.foldr (\(w,r) -> S.insert w . S.insert r) S.empty (g^.strictEdges) +-- <> S.foldr (\((_,r),(w,_)) -> S.insert w . S.insert r) S.empty (M.keysSet $ g^.inplacePaths)) + +-- bufferNodes :: HasFusionGraph g => SimpleGetter g (Set (Node GVal)) +-- bufferNodes = to (\g -> S.map (\(_,b,_) -> b) (g^.dataflowEdges) +-- <> S.foldr (\((b1,_),(_,b2)) -> S.insert b1 . S.insert b2) S.empty (M.keysSet $ g^.inplacePaths)) + + + +-------------------------------------------------------------------------------- +-- The Fusion ILP. +-------------------------------------------------------------------------------- + +-- | A single block of the ILP. +-- +-- 'FusionILP' stores an fusion ILP for a single block of code. This is +-- possible because there can be no fusion between different blocks of code. +-- Separating the ILP into blocks then allows us to pass much smaller ILPs to +-- the solver, which should make the whole process faster. +-- If not, we can always merge the blocks together later. +data FusionILP (op :: Type -> Type) = FusionILP + { _graph :: FusionGraph + , _constraints :: [Constraint] + , _bounds :: Bounds + } + +instance Semigroup (FusionILP op) where + (<>) :: FusionILP op -> FusionILP op -> FusionILP op + (<>) (FusionILP g1 c1 b1) (FusionILP g2 c2 b2) = + FusionILP (g1 <> g2) (c1 <> c2) (b1 <> b2) + +instance Monoid (FusionILP op) where + mempty :: FusionILP op + mempty = FusionILP mempty mempty mempty + +-- | Class for accessing the fusion ILP field of a data structure. +-- +-- We make this because there are at least two data structures that contain a +-- fusion ILP: 'FusionGraphState' and the result of graph construction. This is +-- similar to what microlens-th would generate when using @makeFields@. +class HasFusionILP s op | s -> op where + fusionILP :: Lens' s (FusionILP op) + +graph :: Lens' (FusionILP op) FusionGraph +graph f s = f (_graph s) <&> \g -> s{_graph = g} + +constraints :: Lens' (FusionILP op) [Constraint] +constraints f s = f (_constraints s) <&> \c -> s{_constraints = c} + +bounds :: Lens' (FusionILP op) Bounds +bounds f s = f (_bounds s) <&> \b -> s{_bounds = b} + +instance HasFusionGraph (FusionILP op) where + fusionGraph :: Lens' (FusionILP op) FusionGraph + fusionGraph = graph + +-- | Safely add a strict relation between two computations. +-- +-- Add a strict edge at the 'ancestors'. +before :: HasCallStack => Node Comp -> Node Comp -> FusionILP op -> FusionILP op +before c1 c2 = maybe id insertStrict (ancestors c1 c2) + +-- | Safely add a fusible edge between two computations. +-- +-- If in the same subgraph, add a fusible edge, otherwise add an infusible edge. +fusible :: HasCallStack => Node Comp -> Node GVal -> Node Comp -> FusionILP op -> FusionILP op +fusible prod buff cons = if prod^.parent == cons^.parent + -- Same subgraph, so we can fuse things. + then insertFusible (prod, buff, cons) + -- Different subgraph. This may for instance happen if either 'prod' or + -- 'cons' is placed in an Acond. We cannot fuse over an Acond (or in general, + -- between subgraphs). Hence we mark this edge as infusible. + -- Note that 'infusible' needs to do some additional work to find nodes that + -- are in the same subgraph on which we can add the infusible edge. + else infusible prod buff cons + +-- | Safely add an infusible edge between two computations. +-- +infusible :: HasCallStack => Node Comp -> Node GVal -> Node Comp -> FusionILP op -> FusionILP op +-- We must add the edge on nodes in the same subgraph. 'prod' and 'cons' might +-- be in different subgraphs (if their 'parent' is different), which for +-- instance happens when one of them is nested in an Acond or Awhile. Using +-- 'ancestors', we find the first parent nodes that are in the same subgraph, +-- and we add the infusible edge between these nodes. +infusible prod buff cons = case ancestors prod cons of + Just (prod', cons') -> insertInfusible (prod', buff, cons') + Nothing -> id + +-- | Safely add strict ordering between multiple computations and another computation. +allBefore :: HasCallStack => Nodes Comp -> Node Comp -> FusionILP op -> FusionILP op +allBefore cs1 c2 ilp = foldr' (`before` c2) ilp cs1 + +-- | Safely add fusible edges from all producers to the consumer. +allFusible :: HasCallStack => Nodes Comp -> Node GVal -> Node Comp -> FusionILP op -> FusionILP op +allFusible prods buff cons ilp = foldr' (\prod -> fusible prod buff cons) ilp prods + +-- | Safely add infusible edges from all producers to the consumer. +allInfusible :: HasCallStack => Nodes Comp -> Node GVal -> Node Comp -> FusionILP op -> FusionILP op +allInfusible prods buff cons ilp = foldr' (\prod -> infusible prod buff cons) ilp prods + +-- | Infix synonym for 'before'. +(==|-|=>) :: HasCallStack => Node Comp -> Node Comp -> FusionILP op -> FusionILP op +(==|-|=>) = before + +-- | Infix synonym for 'fusible'. +(--|) :: HasCallStack => Node Comp -> Node GVal -> Node Comp -> FusionILP op -> FusionILP op +(--|) = fusible + +-- | Infix synonym for 'infusible'. +(==|) :: HasCallStack => Node Comp -> Node GVal -> Node Comp -> FusionILP op -> FusionILP op +(==|) = infusible + +-- | Infix synonym for 'allBefore'. +(>=|-|=>) :: HasCallStack => Nodes Comp -> Node Comp -> FusionILP op -> FusionILP op +(>=|-|=>) = allBefore + +-- | Infix synonym for 'allFusible'. +(>-|) :: HasCallStack => Nodes Comp -> Node GVal -> Node Comp -> FusionILP op -> FusionILP op +(>-|) = allFusible + +-- | Infix synonym for 'allInfusible'. +(>=|) :: HasCallStack => Nodes Comp -> Node GVal -> Node Comp -> FusionILP op -> FusionILP op +(>=|) = allInfusible + +-- | Arrow heads to complete '(--|)', '(>-|)', '(==|)' and '(>=|)'. +(|->), (|=>) :: (a -> b) -> a -> b +(|->) = ($) +(|=>) = ($) + +noInplace :: (HasCallStack, HasFusionGraph g) => Node GVal -> g -> g +noInplace node = noInplaceValues %~ S.insert node + +-------------------------------------------------------------------------------- +-- Backend specific definitions +-------------------------------------------------------------------------------- + +-- | The backend has access to a small state so it doesn't accidentally break +-- the state used by the frontend construction algorithm. +data BackendGraphState op env = BackendGraphState + { _backendFusionILP :: FusionILP op -- ^ The entire ILP. + , _backendBuffersEnv :: Env env -- ^ The buffers environment (read only). + , _backendReadersEnv :: ReadersEnv -- ^ The readers environment (read only). + , _backendWritersEnv :: WritersEnv -- ^ The writers environment (read only). + } + +instance HasFusionILP (BackendGraphState op env) op where + fusionILP :: Lens' (BackendGraphState op env) (FusionILP op) + fusionILP f s = f (_backendFusionILP s) <&> \ilp -> s{_backendFusionILP = ilp} + +instance HasBuffersEnv (BackendGraphState op env) (BackendGraphState op env') env env' where + environment :: Lens (BackendGraphState op env) (BackendGraphState op env') (Env env) (Env env') + environment f s = f (_backendBuffersEnv s) <&> \env -> s{_backendBuffersEnv = env} + +instance HasReadersEnv (BackendGraphState op env) where + readersEnv :: Lens' (BackendGraphState op env) ReadersEnv + readersEnv f s = f (_backendReadersEnv s) <&> \env -> s{_backendReadersEnv = env} + +instance HasWritersEnv (BackendGraphState op env) where + writersEnv :: Lens' (BackendGraphState op env) WritersEnv + writersEnv f s = f (_backendWritersEnv s) <&> \env -> s{_backendWritersEnv = env} + +type BackendCluster op = PreArgs (BackendClusterArg op) + +class ( ShrinkArg (BackendClusterArg op), Eq (BackendArg op) , Show (BackendArg op), Ord (BackendArg op) ) => MakesILP op where + + -- | Information that the backend attaches to arguments for use in + -- interpreting/code generation. + type BackendArg op + defaultBA :: BackendArg op + + -- | Information that the backend attaches to the cluster for use in + -- interpreting/code generation. + data BackendClusterArg op arg + combineBackendClusterArg + :: BackendClusterArg op (Out sh e) + -> BackendClusterArg op (In sh e) + -> BackendClusterArg op (Var' sh) + encodeBackendClusterArg :: BackendClusterArg op arg -> Builder + + + -- | Given an ILP solution, attach the backend-specific information to an + -- argument. + labelLabelledArg :: Solution -> Node Comp -> LabelledArg env a -> LabelledArgOp op env a + + -- | Convert a labelled argument to a cluster argument. + getClusterArg :: LabelledArgOp op env a -> BackendClusterArg op a + + -- | This function defines per-operation backend-specific fusion rules. + -- + -- When this function gets called, the majority of edges have already been + -- added to the graph. That is, we have already added read-, write-, fusible- + -- and infusible-edges such that no race conditions exist. + -- The backend is responsible for adding (or removing) edges to (or from) the + -- graph to enforce any additional constraints the implementation may have. + -- + mkGraph + :: Node Comp -- ^ The label of the operation. + -> op args -- ^ The operation. + -> LabelledArgs env args -- ^ The arguments to the operation. + -> State (BackendGraphState op env) () + + -- | This function lets the backend define additional constraints on the ILP. + finalize :: FusionGraph -> [Constraint] + +-- | Attach backend-specific information to labelled arguments. +labelLabelledArgs :: MakesILP op => Solution -> Node Comp -> LabelledArgs env args -> LabelledArgsOp op env args +labelLabelledArgs sol l (arg :>: args) = labelLabelledArg sol l arg :>: labelLabelledArgs sol l args +labelLabelledArgs _ _ ArgsNil = ArgsNil + +-------------------------------------------------------------------------------- +-- ILP Variables +-------------------------------------------------------------------------------- + + +-- | Constructor for 'Pi' variables. +pi :: Node Comp -> Expression +pi = var . Pi + +-- | No clue what this is for. +delayed :: MakesILP op => Node GVal -> Expression +delayed = notB . manifest + +-- | Constructor for 'IsManifest' variables. +manifest :: Node GVal -> Expression +manifest = var . IsManifest + +-- | Safe constructor for 'Fused' variables. +fused :: (Node Comp, Node Comp) -> Expression +fused = var . uncurry Fused + +-- | Safe constructor for 'ReadDir' variables. +readDir :: ReadEdge -> Expression +readDir = var . uncurry ReadDir + +-- | Convert a foldable structure of 'ReadEdge' to a list of 'Expression's. +readDirs :: Foldable f => f ReadEdge -> [Expression] +readDirs = map readDir . toList + +-- | Safe constructor for 'WriteDir' variables. +writeDir :: WriteEdge -> Expression +writeDir = var . uncurry WriteDir + +-- | Convert a foldable structure of 'WriteEdge' to a list of 'Expression's. +writeDirs :: Foldable f => f WriteEdge -> [Expression] +writeDirs = map writeDir . toList + +-- | Safe constructor for 'InPlace' variables. +inplace :: InplacePath -> Expression +inplace ((b1,c1),(c2,b2)) = var $ InPlace b1 c1 c2 b2 + +-- | Safe constructor for 'PiMax' variables. +pimax :: Node GVal -> Expression +pimax = var . PiMax + +maxCluster :: Expression +maxCluster = var MaxCluster + +dirToInt :: Direction -> Int +dirToInt LeftToRight = -2 +dirToInt RightToLeft = -1 + +inFoldSize :: Node Comp -> Expression +inFoldSize = var . InFoldSize + +outFoldSize :: Node Comp -> Expression +outFoldSize = var . OutFoldSize + +-------------------------------------------------------------------------------- +-- Symbol table +-------------------------------------------------------------------------------- + +data Symbol (op :: Type -> Type) where + SExe :: Env env -> LabelledArgs env args -> op args -> Symbol op + SExe' :: Env env -> LabelledArgsOp op env args -> op args -> Symbol op + SUse :: ScalarType e -> Int -> Buffer e -> Symbol op + SITE :: Env env -> ExpVar env PrimBool -> Node Comp -> Node Comp -> Symbol op + SWhl :: Env env -> Node Comp -> Node Comp -> GroundVars env bnd -> Uniquenesses bnd -> Symbol op + SLet :: BoundGLHS bnd env env' -> Node Comp -> Uniquenesses bnd -> Symbol op + SFun :: BoundGLHS bnd env env' -> Node Comp -> Symbol op + SBod :: GroundVals a -> Symbol op + SBlk :: Symbol op + SRet :: Env env -> GroundVars env a -> Symbol op + SCmp :: Env env -> Exp env a -> Symbol op + SAlc :: Env env -> ShapeR sh -> ScalarType e -> ExpVars env sh -> Symbol op + SUnt :: Env env -> ExpVar env e -> Symbol op + SAtr :: Text -> Env env -> ArrayDescriptors env t -> Symbol op + SAsr :: Text -> Env env -> Exp env PrimBool -> Symbol op + SAsu :: Env env -> Exp env PrimBool -> Symbol op + SFen :: Env env -> IdxSet env -> Symbol op + +instance Show (Symbol op) where + show :: Symbol op -> String + show (SExe {}) = "Exe" + show (SExe'{}) = "Exe'" + show (SUse {}) = "Use" + show (SITE {}) = "ITE" + show (SWhl {}) = "Whl" + show (SLet {}) = "Let" + show (SFun {}) = "Fun" + show (SBod {}) = "Bod" + show (SBlk {}) = "Blk" + show (SRet {}) = "Ret" + show (SCmp {}) = "Cmp" + show (SAlc {}) = "Alc" + show (SUnt {}) = "Unt" + show (SAtr {}) = "Atr" + show (SAsr {}) = "Asr" + show (SAsu {}) = "Asu" + show (SFen {}) = "Fen" + +-- | Mapping from labels to symbols. +type Symbols op = Map (Node Comp) (Symbol op) + +data LabelledArgOp op env a = LOp (Arg env a) (ArgLabel a) (BackendArg op) +type LabelledArgsOp op env = PreArgs (LabelledArgOp op env) + +instance Show (LabelledArgOp op env a) where + show :: LabelledArgOp op env a -> String + show (LOp _ bs _) = show bs + +unLabelOp :: LabelledArgsOp op env a -> Args env a +unLabelOp ArgsNil = ArgsNil +unLabelOp ((LOp arg _ _) :>: args) = arg :>: unLabelOp args + +reindexLabelledArgOp :: Applicative f => ReindexPartial f env env' -> LabelledArgOp op env t -> f (LabelledArgOp op env' t) +reindexLabelledArgOp k (LOp (ArgVar vars ) l o) = (\x -> LOp x l o) . ArgVar <$> reindexVars k vars +reindexLabelledArgOp k (LOp (ArgExp e ) l o) = (\x -> LOp x l o) . ArgExp <$> reindexExp k e +reindexLabelledArgOp k (LOp (ArgFun f ) l o) = (\x -> LOp x l o) . ArgFun <$> reindexExp k f +reindexLabelledArgOp k (LOp (ArgArray m repr sh buffers) l o) = (\x -> LOp x l o) <$> (ArgArray m repr <$> reindexVars k sh <*> reindexVars k buffers) + +reindexLabelledArgsOp :: Applicative f => ReindexPartial f env env' -> LabelledArgsOp op env t -> f (LabelledArgsOp op env' t) +reindexLabelledArgsOp = reindexPreArgs reindexLabelledArgOp + +attachBackendLabels :: MakesILP op => Solution -> Symbols op -> Symbols op +attachBackendLabels sol = M.mapWithKey \cases + l (SExe env largs op) -> SExe' env (labelLabelledArgs sol l largs) op + _ SExe'{} -> internalError "already converted???" + _ con -> con + + + +-------------------------------------------------------------------------------- +-- FusionGraph construction +-------------------------------------------------------------------------------- + +-- | State for the full graph construction. +-- +-- The graph is constructed inside the state monad by inserting edges into it. +-- The state also contains the symbols needed for reconstruction of the AST and +-- the current computation label. +-- +-- Computations labels and buffer labels should always be unique, so we only use +-- one counter for the computation labels and provide lenses for interpreting +-- them as buffer labels. +-- Since all labels are unique, we can use a single symbol map for all labels +-- instead of separate maps for computation and buffer labels. +-- +-- The result of the full graph construction is reserved for the return values +-- of nodes in the program, which are generally buffer labels. +-- This method makes defining the control flow easier since we do not need to +-- worry about merging the graphs in the return values as in the old approach. +-- +-- The environment is not passed as an argument to 'mkFusionGraph' since it may +-- be modified by certain computations. Specifically, when a buffer is marked as +-- mutable, a copy of the buffer is created and the original buffer is replaced +-- by the copy in the environment. +-- +-- We keep track of which computation last wrote to a buffer, i.e. the producer +-- of the buffer. Under normal circumstances a buffer has one and only one +-- producer, but when we enter an if-then-else it could be that some buffer +-- is written to by both branches. In this case the buffer is mutated by both, +-- which is safe because during execution only one branch is taken. +-- +-- The environment and return values contain sets of buffer for a similar +-- reason. An if-then-else could return different buffers of the same type +-- depending on which branch is taken. +-- +data FusionGraphState op env = FusionGraphState + { _fusionILP :: FusionILP op -- ^ The ILP information. + , _buffersEnv :: Env env -- ^ The label environment. + , _readersEnv :: ReadersEnv -- ^ Mapping from buffers to their current consumers. + , _writersEnv :: WritersEnv -- ^ Mapping from buffers to their current producers. + , _allocators :: Allocators -- ^ Mapping from buffers to their allocator. + , _symbols :: Symbols op -- ^ The symbols for the ILP. + , _currComp :: Node Comp -- ^ The current computation label. + , _currFence :: Maybe (Node Comp) -- ^ The current fence. All new nodes should be linked with an edge from currFence, so the fence is placed before that node. + , _currEnvL :: EnvLabel -- ^ The current environment label. + } + +type ReadersEnv = Map (Node GVal) (Nodes Comp) +type WritersEnv = Map (Node GVal) (Nodes Comp) +type Allocators = Map (Node GVal) (Node Comp) + +initialFusionGraphState :: FusionGraphState op () +initialFusionGraphState = FusionGraphState mempty EnvNil mempty mempty mempty mempty (Node 0 Nothing) Nothing 0 + +-- instance Show (FusionGraphState op env) where +-- show :: FusionGraphState op env -> String +-- show s = "FusionGraphState { readersEnv=" ++ show (s^.readersEnv) ++ +-- ", writersEnv=" ++ show (s^.writersEnv) ++ +-- " }" + +instance HasFusionILP (FusionGraphState op env) op where + fusionILP :: Lens' (FusionGraphState op env) (FusionILP op) + fusionILP f s = f (_fusionILP s) <&> \ilp -> s{_fusionILP = ilp} + +class HasBuffersEnv s t env env' | s -> env, t -> env' where + environment :: Lens s t (Env env) (Env env') + +instance HasBuffersEnv (FusionGraphState op env) (FusionGraphState op env') env env' where + environment :: Lens (FusionGraphState op env) (FusionGraphState op env') (Env env) (Env env') + environment f s = f (_buffersEnv s) <&> \env -> s{_buffersEnv = env} + +class HasReadersEnv s where + readersEnv :: Lens' s ReadersEnv + +instance HasReadersEnv (FusionGraphState op env) where + readersEnv :: Lens' (FusionGraphState op env) ReadersEnv + readersEnv f s = f (_readersEnv s) <&> \env -> s{_readersEnv = env} + +class HasWritersEnv s where + writersEnv :: Lens' s WritersEnv + +instance HasWritersEnv (FusionGraphState op env) where + writersEnv :: Lens' (FusionGraphState op env) WritersEnv + writersEnv f s = f (_writersEnv s) <&> \env -> s{_writersEnv = env} + +class HasAllocators s where + allocators :: Lens' s Allocators + +instance HasAllocators (FusionGraphState op env) where + allocators :: Lens' (FusionGraphState op env) Allocators + allocators f s = f (_allocators s) <&> \alloc -> s{_allocators = alloc} + +class HasSymbols s op | s -> op where + symbols :: Lens' s (Symbols op) + +instance HasSymbols (FusionGraphState op env) op where + symbols :: Lens' (FusionGraphState op env) (Symbols op) + symbols f s = f (_symbols s) <&> \sym -> s{_symbols = sym} + +currComp :: Lens' (FusionGraphState op env) (Node Comp) +currComp f s = f (_currComp s) <&> \c -> s{_currComp = c} + +currFence :: Lens' (FusionGraphState op env) (Maybe (Node Comp)) +currFence f s = f (_currFence s) <&> \c -> s{_currFence = c} + +currEnvL :: Lens' (FusionGraphState op env) EnvLabel +currEnvL f s = f (_currEnvL s) <&> \l -> s{_currEnvL = l} + +-- | Lens for creating the backend graph state. +-- +-- This lens sets new values for the readers and writers environments because +-- the backend needs to work with the environment from before the computation +-- was added. We don't need to do the same for the buffers environment, because +-- it only changes when a new variable is introduced. +-- +-- The fusion ILP is the only value that the backend may modify, so its the only +-- value that is retrieved from the backend graph state afterwards. +backendGraphState :: ReadersEnv -> WritersEnv -> Lens' (FusionGraphState op env) (BackendGraphState op env) +backendGraphState renv wenv f s = f (BackendGraphState (s^.fusionILP) (s^.environment) renv wenv) + <&> \b -> s & fusionILP .~ b^.fusionILP + +-- | Lens for getting and setting the writers of a buffer. +-- By default we throw an error if the buffer is not found in the environment. +writers :: HasWritersEnv s => Node GVal -> Lens' s (Nodes Comp) +writers b f s = f (M.findWithDefault msg b (s^.writersEnv)) <&> \cs -> s & writersEnv %~ M.insert b cs + where msg = internalError "missing writer" + +-- | Lens for getting all writers of buffers. +allWriters :: (Foldable f, HasWritersEnv s) => f (Node GVal) -> SimpleGetter s (Nodes Comp) +allWriters bs = to (\s -> foldMap (\b -> s^.writers b) bs) + +-- | Lens for getting and setting the readers of a buffer. +-- By default the set of readers is empty. +readers :: HasReadersEnv s => Node GVal -> Lens' s (Nodes Comp) +readers b f s = f (M.findWithDefault S.empty b (s^.readersEnv)) <&> \cs -> s & readersEnv %~ M.insert b cs + +-- | Lens for getting all readers of buffers. +allReaders :: (Foldable f, HasReadersEnv s) => f (Node GVal) -> SimpleGetter s (Nodes Comp) +allReaders bs = to (\s -> foldMap (\b -> s^.readers b) bs) + +-- | Lens for getting and setting symbol of a computation. +symbol :: HasSymbols s op => Node Comp -> Lens' s (Maybe (Symbol op)) +symbol c = symbols.(`M.alterF` c) + +-- | Lens for getting and setting the allocator of a buffer. 'symbol' but for +-- buffers. +allocator :: HasAllocators s => Node GVal -> Lens' s (Maybe (Node Comp)) +allocator b = allocators.(`M.alterF` b) + +-- | Lens for working under the scope of a computation. +-- +-- It first sets the parent of the current label to the supplied computation +-- label. Then it applies the function to the 'FusionGraphState' with the now +-- parented label. Finally, it sets the parent of the current label back to the +-- original parent. +scope :: Node Comp -> Lens' (FusionGraphState op env) (FusionGraphState op env) +scope c = with (currComp.parent) (Just c) + +local :: Env env' -> Lens' (FusionGraphState op env) (FusionGraphState op env') +local env' f s = (environment .~ s^.environment) <$> f (s & environment .~ env') + +-- | Fresh computation label. +freshComp :: State (FusionGraphState op env) (Node Comp) +freshComp = do + c <- zoom currComp freshL' + fence' <- use currFence + case fence' of + Nothing -> return () + Just a -> fusionILP %= a `before` c + fusionILP %= insertComputation c + return c + +-- | Fresh buffer and the corresponding computation label. +-- +-- The implementation of 'writers' makes it so by default the buffer is produced +-- by the computation that allocates it. This is possible because they have the +-- same label just, just different types. We still need to add the read edge to +-- the graph though. +freshGVal :: Node Comp -> State (FusionGraphState op env) (Node GVal) +freshGVal comp = do + v <- zoom (currComp.asBuff) freshL' + fusionILP %= insertValue v + writers v .= S.singleton comp + allocator v ?= comp + return v + + +-- | Fresh 'Val'. +freshVal :: Node Comp -> s t -> State (FusionGraphState op env) (Val s t) +freshVal comp tp = val tp <$> freshGVal comp + + +-- | Fresh 'Vals'. +freshVals :: Node Comp -> TupR s t -> State (FusionGraphState op env) (Vals s t) +freshVals comp = traverseTupR (freshVal comp) + + +-- | Read from a buffer. +readsBuffers :: HasCallStack => Node Comp -> Nodes GVal -> State (FusionGraphState op env) () +readsBuffers c = traverse_ \b -> do + ws <- use $ writers b + fusionILP %= ws >-|b|-> c + readers b %= S.insert c + +-- | Require a buffer (i.e. to index into it or pass it to a function). +-- TODO: Rename requires -> reads, reads -> input. +requiresBuffers :: HasCallStack => Node Comp -> Nodes GVal -> State (FusionGraphState op env) () +requiresBuffers c = traverse_ \b -> do + ws <- use $ writers b + fusionILP %= ws >=|b|=> c + readers b %= S.insert c + + +-- | Write to a buffer. +-- +-- For a write to be safe we need to enforce the following: +-- 1. All readers run before the computation. +-- 2. All writers run before the computation. +-- 3. We become the sole writer of the buffer. +-- 4. We clear the readers of the buffer. +writesBuffers :: HasCallStack => Node Comp -> Nodes GVal -> State (FusionGraphState op env) () +writesBuffers c = traverse_ \b -> do + rs <- use $ readers b + ws <- use $ writers b + fusionILP %= rs >=|-|=> c + fusionILP %= ws >=|-|=> c + writers b .= S.singleton c + readers b .= S.empty + +-- | Mutate a buffer. +-- +-- For a mutation to be safe we need to enforce the following: +-- 1. All readers run before this computation. +-- 2. All writers are infusible with this computation. +-- 3. We become the sole writer of the buffer. +-- 4. We clear the readers of the buffer. +mutatesBuffers :: HasCallStack => Node Comp -> Nodes GVal -> State (FusionGraphState op env) () +mutatesBuffers c = traverse_ \b -> do + rs <- use $ readers b + ws <- use $ writers b + fusionILP %= rs >=|-|=> c + fusionILP %= ws >=|b|=> c + writers b .= S.singleton c + readers b .= S.empty + +-- | Return a buffer. +-- +-- This can be interpreted as mutation with the identity function (i.e. no-op). +-- Since we don't actually change the contents of the buffer, we don't need to +-- enforce 1 and 4. +returnsBuffers :: HasCallStack => Node Comp -> Nodes GVal -> State (FusionGraphState op env) () +returnsBuffers c = traverse_ \b -> do + ws <- use $ writers b + fusionILP %= ws >=|b|=> c + writers b .= S.singleton c + +-- | Bind a buffer to a let. +-- +-- This can be interpreted as mutation with the identity function (i.e. no-op). +-- Since we don't actually change the contents of the buffer, we don't need to +-- enforce 1 and 4. We also don't enforce 2, because doing so would prevent all +-- buffers from being non-manifest. (All buffers are bound to a let and +-- infusible edges force manifestation, so all buffers would be manifest.) +bindsBuffers :: HasCallStack => Node Comp -> Nodes GVal -> State (FusionGraphState op env) () +bindsBuffers c = traverse_ \b -> do + ws <- use $ writers b + fusionILP %= ws >-|b|-> c + writers b .= S.singleton c + +-- | Registers that no in-place updates may happen on these variables. +-- +-- This is needed when the uses of a variable cannot be fully tracked. This +-- happens when a variable is returned, or used as the input of an awhile loop. +noInplaceBuffers :: Nodes GVal -> State (FusionGraphState op env) () +noInplaceBuffers = traverse_ \b -> do + fusionILP %= noInplace b + +-- -- | A computation executes another computattion. +-- executes :: HasCallStack => Node Comp -> Node Comp -> State (FusionGraphState op env) () +-- executes c1 c2 = fusionILP %= c1 ==|-|=> c2 + + + +-------------------------------------------------------------------------------- +-- Full Graph construction +-------------------------------------------------------------------------------- + +type FullGraph op = (FusionILP op, Symbols op, Allocators) + +-- The 2 instances below can be used to clean up the code in ILP.hs a bit. +instance HasFusionILP (FullGraph op) op where + fusionILP :: Lens' (FullGraph op) (FusionILP op) + fusionILP f (ilp, sym, alloc) = f ilp <&> (,sym,alloc) + +instance HasSymbols (FullGraph op) op where + symbols :: Lens' (FullGraph op) (Symbols op) + symbols f (ilp, sym, alloc) = f sym <&> (ilp,,alloc) + +instance HasAllocators (FullGraph op) where + allocators :: Lens' (FullGraph op) Allocators + allocators f (ilp, sym, alloc) = f alloc <&> (ilp,sym,) + +-- | Construct the full fusion graph for a program. +mkFullGraph :: MakesILP op => PreOpenAcc op () t -> FullGraph op +mkFullGraph acc = finalizeInplacePaths $ makeManifest (valsNodes res) (s^.fusionILP, s^.symbols, s^.allocators) + where (res, s) = runState (mkFusionGraph acc) initialFusionGraphState + +-- | Construct the full fusion graph for a function. +mkFullGraphF :: MakesILP op => PreOpenAfun op () a -> FullGraph op +mkFullGraphF acc = finalizeInplacePaths (s^.fusionILP, s^.symbols, s^.allocators) + where (_, s) = runState (mkFusionGraphF acc) initialFusionGraphState + +-- | Make the supplied value nodes manifest. +makeManifest :: HasFusionILP g op => Nodes GVal -> g -> g +makeManifest bs = fusionILP.manifestValues <>~ bs + +-------------------------------------------------------------------------------- +-- FusionGraph construction +-------------------------------------------------------------------------------- + +-- | Construct the fusion graph of a program. +mkFusionGraph :: forall op env t. MakesILP op + => PreOpenAcc op env t + -> State (FusionGraphState op env) (GroundVals t) +mkFusionGraph (Exec op args) = do + env <- use environment + renv <- use readersEnv + wenv <- use writersEnv + c <- freshComp + let labelledArgs = labelArgs args env + let inpArrs = inputArrays labelledArgs + let outArrs = outputArrays labelledArgs + let notArrs = notArrays labelledArgs + c `readsBuffers` (inpArrs `S.difference` outArrs) + c `writesBuffers` (outArrs `S.difference` inpArrs) + c `mutatesBuffers` (inpArrs `S.intersection` outArrs) + c `requiresBuffers` notArrs + zoom (backendGraphState renv wenv) (mkGraph c op labelledArgs) + symbol c ?= SExe env labelledArgs op + return TupRunit + +mkFusionGraph (Alet LeftHandSideUnit _ bnd body) + = mkFusionGraph bnd >> mkFusionGraph body + +-- In this definition I assume that whatever the right-hand side returns is +-- produced by a single computation, which is currently true because all +-- instructions attach themselves to the buffer. +mkFusionGraph (Alet lhs u bnd body) = do + c <- freshComp -- TODO: If there is an issue with reconstruction, maybe move this behind "bndRes <- mkFusionGraph bnd". The order in which labels are generate affects the order in which the clusters are interpreted. Previously let-bindings where always in a separate cluster from the bound computation, but now they are usually in the same cluster to prevent all buffers from being manifest. That said, topsort should already be taking care of this ordering issue. + env <- use environment + bndRes <- mkFusionGraph bnd + bndResW <- foldMapMTupR (\v -> use (allWriters (valNodes v))) bndRes + c `bindsBuffers` valsNodes bndRes + env' <- zoom currEnvL (weakenEnv lhs bndRes u env) + symbol c ?= SLet (bindLHS lhs env') (fromSingletonSet bndResW) u + zoom (local env') (mkFusionGraph body) + +mkFusionGraph (Return vars) = do + env <- use environment + retN <- freshComp + let (_, bs, _) = lookupVars vars env + retN `returnsBuffers` valsNodes bs + noInplaceBuffers $ valsNodes bs + symbol retN ?= SRet env vars + return bs + +mkFusionGraph (Manifest buff) = do + env <- use environment + retN <- freshComp + let (_, bs, _) = lookupVar buff env + retN `returnsBuffers` valsNodes bs + symbol retN ?= SRet env (TupRsingle buff) + return bs + +mkFusionGraph (Compute expr) = do + c <- freshComp + env <- use environment + c `requiresBuffers` getExpDeps expr env + symbol c ?= SCmp env expr + freshVals c (groundsR expr) + +mkFusionGraph (Alloc shr e sh) = do + c <- freshComp + env <- use environment + c `requiresBuffers` getVarsDeps sh env + symbol c ?= SAlc env shr e sh + TupRsingle <$> freshVal c (GroundRbuffer e) + +mkFusionGraph (Unit v) = do + c <- freshComp + env <- use environment + c `requiresBuffers` getVarDeps v env + symbol c ?= SUnt env v + TupRsingle <$> freshVal c (GroundRbuffer $ varType v) + +mkFusionGraph (Use tp n buff) = do + c <- freshComp + symbol c ?= SUse tp n buff + TupRsingle <$> freshVal c (GroundRbuffer tp) + +mkFusionGraph (Acond cond tacc facc) = do + env <- use environment + condN <- freshComp + zoom (scope condN) do + trueN <- freshComp + falseN <- freshComp + symbol condN ?= SITE env cond trueN falseN + condN `requiresBuffers` getVarDeps cond env -- If-then-else reads the condition variable, + _ <- branches -- executes "both" branches, and + (block trueN $ mkFusionGraph tacc) + (block falseN $ mkFusionGraph facc) + res <- freshVals condN (groundsR tacc) + condN `returnsBuffers` valsNodes res -- returns the unified result of both branches. + return res + +mkFusionGraph (Awhile u cond body init) = do + env <- use environment + whileN <- freshComp + res <- freshVals whileN (groundsR init) + zoom (scope whileN) do + condN <- freshComp + bodyN <- freshComp + let init_val = lookupVars init env ^. _2 + whileN `requiresBuffers` valsNodes init_val -- While reads the initial value, + noInplaceBuffers $ valsNodes init_val + _ <- branches -- executes "both" the condition and body, + (block condN $ mkFusionGraphU u cond) + (block bodyN $ mkFusionGraphU u body) + symbol whileN ?= SWhl env condN bodyN init u + return res -- to return a fresh value of the same type as the initial value. + +mkFusionGraph (Atrace msg t) = do + c <- freshComp + env <- use environment + c `requiresBuffers` getIdxSetDeps (arrayDescriptorsIdxSet t) env + symbol c ?= SAtr msg env t + TupRsingle <$> freshVal c (GroundRscalar scalarTypeWord8) + +mkFusionGraph (Aassert msg cond) = do + c <- freshComp + env <- use environment + c `requiresBuffers` getExpDeps cond env + symbol c ?= SAsr msg env cond + TupRsingle <$> freshVal c (GroundRscalar scalarTypeWord8) + +mkFusionGraph (Aassume cond) = do + c <- freshComp + env <- use environment + c `requiresBuffers` getExpDeps cond env + symbol c ?= SAsu env cond + TupRsingle <$> freshVal c (GroundRscalar scalarTypeWord8) + +mkFusionGraph (Fence deps next) = do + c <- freshComp + env <- use environment + c `requiresBuffers` getIdxSetDeps deps env + symbol c ?= SFen env deps + currFence .= Just c -- All later nodes will be placed after this fence + mkFusionGraph next + +-- | Construct the fusion graph of a single-argument function. +mkFusionGraphU :: forall op env s t. (MakesILP op) + => Uniquenesses s -> PreOpenAfun op env (s -> t) + -> State (FusionGraphState op env) (Node Comp) +mkFusionGraphU _ (Abody body) = groundFunctionImpossible $ groundsR body +mkFusionGraphU u (Alam lhs f) = do + lam <- freshComp + env <- use environment + args <- freshVals lam (lhsToTupR lhs) + env' <- zoom currEnvL (weakenEnv lhs args u env) + fun <- zoom (local env') (mkFusionGraphF f) + symbol lam ?= SFun (bindLHS lhs env') fun + return lam + + + +-- | Construct the fusion graph of a function. +mkFusionGraphF :: forall op env t. (MakesILP op) + => PreOpenAfun op env t + -> State (FusionGraphState op env) (Node Comp) +mkFusionGraphF (Alam lhs f) = mkFusionGraphU (shared $ lhsToTupR lhs) (Alam lhs f) +mkFusionGraphF (Abody acc) = do + bodyN <- freshComp + zoom (scope bodyN) do + res <- mkFusionGraph acc + symbol bodyN ?= SBod res + id %= makeManifest (valsNodes res) + bodyN `returnsBuffers` valsNodes res + return bodyN + + + +-- | A block has no effect on the graph, but the added scope they introduce is +-- required by the reconstruction algorithm. +-- The 'SBlk' symbol isn't even needed, because any time we encounter it an +-- error is thrown, but it is somewhat useful for debugging. +block :: Node Comp -> State (FusionGraphState op env) r -> State (FusionGraphState op env) r +block c f = zoom (scope c) $ symbol c ?= SBlk >> f + + + +-- | Helper for if-then-else and while, executing both branches with the same +-- readers and writers environments, then unifying said environments. +branches :: State (FusionGraphState op env) r1 + -> State (FusionGraphState op env) r2 + -> State (FusionGraphState op env) (r1, r2) +branches f1 f2 = do + -- Store the current readers and writers environments. + renv <- use readersEnv + wenv <- use writersEnv + -- Run the left branch. + res1 <- f1 + -- Retrieve the readers and writers environments after the branch, then restore them. + renv' <- readersEnv <<.= renv + wenv' <- writersEnv <<.= wenv + -- Run the right branch. + res2 <- f2 + -- Unify the readers and writers environments. + readersEnv %= M.unionWith S.union renv' + writersEnv %= M.unionWith S.union wenv' + -- Return the results of both branches. + return (res1, res2) + + + +-------------------------------------------------------------------------------- +-- In-place update path extension +-------------------------------------------------------------------------------- + +-- | Create in-place paths from a computation, one of its inputs and one of its outputs. +mkUnitInplacePaths :: HasCallStack => Number -> Node Comp -> ArgLabel (In sh e) -> ArgLabel (Out sh' e') -> Map InplacePath Number +mkUnitInplacePaths n c = mkInplacePaths n c c + + +-- | Create in-place paths from 2 computations, an input of the first and an output of the second. +mkInplacePaths :: HasCallStack => Number -> Node Comp -> Node Comp + -> ArgLabel (In sh e) -> ArgLabel (Out sh' e') -> Map InplacePath Number +mkInplacePaths n r w lIn lOut + | sameShape lIn lOut + , bsIn <- getLabelUniqueArrDeps lIn + , bsOut <- getLabelUniqueArrDeps lOut + = foldMap (\bIn -> foldMap (\bOut -> M.singleton ((bIn, r), (w, bOut)) n) bsOut) bsIn +mkInplacePaths _ _ _ _ _ = M.empty + + +-- | Filters the in-place update paths to only include those that are valid. +filterInplacePaths :: forall op. FullGraph op -> FullGraph op +filterInplacePaths g = g&fusionILP.inplacePaths %~ filterKeys sameElementType + where + -- Checks if the input and output have the same element type. + sameElementType :: InplacePath -> Bool + sameElementType ((b1, _), (_, b2)) + | Exists tp1 <- getElt b1 + , Exists tp2 <- getElt b2 + , Just Refl <- matchTypeR tp1 tp2 = True + | otherwise = False + + -- Gets the element size of a buffer. + getElt :: Node GVal -> Exists TypeR + getElt b = case (g^.allocator b) >>= (\c -> g^.symbol c) of + Just (SAlc _ _ e _) -> Exists $ TupRsingle e + Just (SUnt _ v) -> Exists $ TupRsingle $ varType v + Just (SUse e _ _) -> Exists $ TupRsingle e + Just (SFun lhs _) -> groundsRtoTypeR $ lhsToTupR $ unbindLHS lhs + Just (SWhl _ _ _ init _) -> groundsRtoTypeR $ groundsR init + Just _ -> internalError "not an array allocator" + Nothing -> internalError "no allocator found" + + -- Use exists because we don't know the exact type. + groundsRtoTypeR :: GroundsR s -> Exists TypeR + groundsRtoTypeR TupRunit = Exists TupRunit + groundsRtoTypeR (TupRsingle (GroundRscalar tp)) = Exists $ TupRsingle tp + groundsRtoTypeR (TupRsingle (GroundRbuffer tp)) = Exists $ TupRsingle tp + groundsRtoTypeR (TupRpair e1 e2) + | Exists tp1 <- groundsRtoTypeR e1 + , Exists tp2 <- groundsRtoTypeR e2 + = Exists $ TupRpair tp1 tp2 + + +-- | Finalizes the in-place update paths by combining them and filtering them. +-- +-- Using only merging of unit-size in-place paths does not yield a lot of in-place +-- update locations. +-- The better approach seems to be constructing paths from scratch. +-- It is then possible to still use the merging strategy, as it may yield some additional paths. +-- However, in testing it seems that almost no paths are missing. +finalizeInplacePaths :: MakesILP op => FullGraph op -> FullGraph op +finalizeInplacePaths = filterInplacePaths . mkInplacePathsFromClusters + + +-- | Calculate in-place update paths based on vertical clusters. +-- This approach should be able to find many more in-place update paths than +-- just combining the in-place paths of length 1, because it considers +-- computations that are not directly connected by an in-place path. +mkInplacePathsFromClusters :: forall op. MakesILP op => FullGraph op -> FullGraph op +mkInplacePathsFromClusters g = g&fusionILP.inplacePaths <>~ go initialClusters + where + initialClusters :: Set (Node Comp, Node Comp) + initialClusters = S.map (\c->(c,c)) (g^.fusionILP.computationNodes) + + go :: Set (Node Comp, Node Comp) -> Map InplacePath Number + go = foldMap $ \(r, w) -> do + let !selfPath = clusterInplacePaths r w + let !selfStepped = S.map (r,) (next w) + M.union selfPath $! go selfStepped + + -- Get the set of fusible consumers. + next :: Node Comp -> Nodes Comp + next = flip (M.findWithDefault S.empty) nextMap + + -- Map from producer to consumers. + nextMap :: Map (Node Comp) (Nodes Comp) + nextMap = foldl (flip \(c1,_,c2) -> M.insertWith (<>) c1 (S.singleton c2)) M.empty (g^.fusionILP.fusibleEdges) + + clusterInplacePaths :: Node Comp -> Node Comp -> Map InplacePath Number + clusterInplacePaths cIn cOut = case (g^.symbol cIn, g^.symbol cOut) of + (Just (SExe _ largsIn _), Just (SExe _ largsOut _)) -> + foldMapInputLabels (\lIn -> foldMapOutputLabels (mkInplacePaths 1 cIn cOut lIn) largsOut) largsIn + _ -> mempty + +-------------------------------------------------------------------------------- +-- Reconstruction +-------------------------------------------------------------------------------- + +-- | Makes a ReindexPartial, which allows us to transform indices in @env@ into indices in @env'@. +-- We cannot guarantee the index is present in env', so we use the partiality of ReindexPartial by +-- returning a Maybe. Uses unsafeCoerce to re-introduce type information implied by the EnvLabels. +mkReindexPartial :: forall env env'. Map (Node GVal) (Node GVal) -> Env env -> Env env' -> ReindexPartial Maybe env env' +mkReindexPartial m env env' idx = let node = lookupIdx idx env^._2 in case idxOf (inplaceOf node) env' of + Just idx' -> Just idx' + -- Note: we are not yet sure what happens if the variable after in-place updates is not found, + -- but the variable before in-place updates is found. It might be that this never occurs. + -- Currently we throw an error in this case. + -- Instead, we could gracefully fall back to using the original value instead of the one defined by an in-place update. + -- This is a bit unsafe, because there is no guarantee anything has been written to this value, or even that it exists at all. + -- That said, I (Timo) think this can only happen when a value is returned, thus causing both the original value and the in-place updated value to be out-of-scope. + -- The returned value is then stored in another scope, so if the return statement properly replaces the original value with the in-place updated value, then the resulting program should bind the in-place updated value instead of the original. + -- See https://github.com/ivogabe/accelerate/pull/11#discussion_r2424009960 + Nothing + | inplaceOf node /= node + , Just _ <- (idxOf node env') -> + internalError $ "mkReindexPartial: index of in-place updated buffer not found. Original buffer (without in-place updates taken into account) is available, but using that might not be sound." + | otherwise -> Nothing + where + -- Replace the buffer with the one we will actually write the data to. + inplaceOf :: GroundVals a -> GroundVals a + inplaceOf (TupRsingle (Val tp n)) = TupRsingle $ Val tp $ M.findWithDefault n n m + inplaceOf (TupRpair bs1 bs2) = TupRpair (inplaceOf bs1) (inplaceOf bs2) + inplaceOf TupRunit = TupRunit + + -- Find the corresponding EnvLabel in the new environment. + idxOf :: forall e a. GroundVals a -> Env e -> Maybe (Idx e a) + idxOf bs ((_,bs',_) :>>: rest) -- bs' is the GroundVals in the new environment + -- Here we have to convince GHC that the top element in the environment + -- really does have the same type as the one we were searching for. + -- Some literature does this stuff too: 'effect handlers in haskell, evidently' + -- and 'a monadic framework for delimited continuations' come to mind. + -- Basically: standard procedure if you're using Ints as a unique identifier + -- and want to re-introduce type information. :) + -- Type applications allow us to restrict unsafeCoerce to the return type. + | Just Refl <- matchGroundValsType bs bs' + , bs == bs' = Just ZeroIdx + -- Recurse if we did not find e' yet. + | otherwise = SuccIdx <$> idxOf bs rest + -- If we hit the end, the Elabel was not present in the environment. + -- That probably means we'll error out at a later point, but maybe there is + -- a case where we try multiple options? No need to worry about it here. + idxOf _ EnvNil = Nothing + + +-------------------------------------------------------------------------------- +-- Helpers +-------------------------------------------------------------------------------- + +-- | Groups a map by a function. +groupByMap :: (Ord k, Ord k') => (k -> v -> k') -> Map k v -> Map k' (Map k v) +groupByMap f = M.foldlWithKey (\m k v -> M.insertWith M.union (f k v) (M.singleton k v) m) M.empty + +-- | Groups a set by a function. +groupBySet :: (Ord k, Ord v) => (v -> k) -> Set v -> Map k (Set v) +groupBySet f = S.foldl (\m v -> M.insertWith S.union (f v) (S.singleton v) m) M.empty + +-- | Lens that protects a given value from being modified. +protected :: Lens' s a -> Lens' s s +protected l f s = (l .~ s^.l) <$> f s + +-- | Lens that protects all but the given value from being modified. +unprotected :: Lens' s a -> Lens' s s +unprotected l f s = (\s' -> s & l .~ s'^.l) <$> f s + +-- | Lens that temporarily uses the supplied value in place of the current +-- value, then restores the original value. +with :: Lens' s a -> a -> Lens' s s +with l a f s = (l .~ s^.l) <$> f (s & l .~ a) + +-- | Converts a singleton set into a value. +-- +-- This function is partial and will throw an error if the set is not singleton. +fromSingletonSet :: HasCallStack => Set a -> a +fromSingletonSet (S.toList -> [x]) = x +fromSingletonSet _ = internalError "set is not singleton" + +-- | Converts a triple (a, b, c) into ((a, b), c) +tripleToLeftRec :: (a, b, c) -> ((a, b), c) +tripleToLeftRec (x, y, z) = ((x, y), z) + +filterKeys :: (k -> Bool) -> Map k a -> Map k a +filterKeys p = M.filterWithKey (\k _ -> p k) + +(&&&) :: (a -> Bool) -> (a -> Bool) -> a -> Bool +(&&&) p1 p2 x = p1 x && p2 x + + + +-------------------------------------------------------------------------------- +-- Converting Graphs to DOT +-------------------------------------------------------------------------------- + +-- | Converts a graph to a DOT representation. +toDOT :: FusionGraph -> Symbols op -> String +toDOT g syms = "strict digraph {\n" ++ + concatMap (\c -> " <" ++ show c ++ "> [shape=box, label=\"" ++ show (syms M.! c) ++ tail (show c) ++ "\"];\n") (g^.computationNodes) ++ + concatMap (\b -> " <" ++ show b ++ "> [shape=ellipse, label=\"" ++ show b ++ "\"];\n") (g^.valueNodes) ++ + concatMap (\(b,c) -> " <" ++ show b ++ "> -> <" ++ show c ++ "> [];\n") (g^.readEdges) ++ + concatMap (\(c,b) -> " <" ++ show c ++ "> -> <" ++ show b ++ "> [];\n") (g^.writeEdges) ++ + concatMap (\((b1, _), (_, b2)) -> " <" ++ show b1 ++ "> -> <" ++ show b2 ++ "> [color=gray, style=dotted];\n") (M.keysSet $ g^.inplacePaths) ++ + concatMap (\(c1,_,c2) -> " <" ++ show c1 ++ "> -> <" ++ show c2 ++ "> [color=green];\n") (g^.fusibleEdges) ++ + concatMap (\(c1,_,c2) -> " <" ++ show c1 ++ "> -> <" ++ show c2 ++ "> [color=red];\n") (g^.infusibleEdges) ++ + concatMap (\(c1,c2) -> " <" ++ show c1 ++ "> -> <" ++ show c2 ++ "> [style=dashed, color=red];\n") (g^.orderEdges) ++ + "}\n" diff --git a/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/HiGHS.hs b/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/HiGHS.hs new file mode 100644 index 000000000..903218530 --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/HiGHS.hs @@ -0,0 +1,65 @@ +{-# LANGUAGE ViewPatterns #-} +module Data.Array.Accelerate.Trafo.Partitioning.ILP.HiGHS where +import Data.Array.Accelerate.Trafo.Partitioning.ILP.LinearConstraint +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Solver +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Graph (MakesILP) + +import qualified Numeric.HiGHS.LP as LP +import Numeric.LinearProgramming.Common +import GHC.Float (int2Double) +import qualified Data.Map as M +import Data.Array.Comfort.Storable +import qualified Debug.Trace + +-- DON'T USE +-- Even though HiGHS is also capable of MIP problems, +-- the highs-lp package only supports LP. +-- This is useless for us; we need ILP/MIP support. +-- Keeping this here in case we get around to forking highs-lp with ILP support. + + +data HiGHS = Highs + +instance ILPSolver HiGHS where + solvePartial Highs ilp@(ILP dir cost constraint bounds n) = pure . getSolution $ + LP.solve LP.choose bounds' constraint' (dir', cost') + where + vs = allVars ilp + bounds' = highbounds bounds [] + constraint' = highconstraints n constraint [] + dir' = case dir of + Maximise -> LP.Maximize + Minimise -> LP.Minimize + cost' = case highexpr n cost of (term,_constant) -> fromMap . M.unionWith (+) (M.fromSet (const 0) vs) $ term2Map term + +getSolution (s, Nothing) = Nothing +getSolution (s, Just (c,arr)) = Just (round <$> toMap arr) + +term2Map [] = mempty +term2Map ((Term d v):ts) = M.unionWith (+) (M.singleton v d) (term2Map ts) + +highbounds NoBounds = id +highbounds (a :<> b) = highbounds a . highbounds b +highbounds (Binary v) = (:) $ Inequality v $ Between 0 1 +highbounds (LowerUpper l v u) = (:) $ Inequality v $ Between (int2Double l) (int2Double u) +highbounds (Lower l v) = (:) $ Inequality v $ GreaterEqual (int2Double l) +highbounds (Upper v u) = (:) $ Inequality v $ LessEqual (int2Double u) + +highconstraints n TrueConstraint = id +highconstraints n (a :&& b) = highconstraints n a . highconstraints n b +highconstraints n ((a :>= b)) = (:) $ case highexpr n (a.-.b) of (terms,((-1)*)-> c) -> Inequality terms $ GreaterEqual (int2Double c) +highconstraints n ((a :== b)) = (:) $ case highexpr n (a.-.b) of (terms,((-1)*)-> c) -> Inequality terms $ Equal (int2Double c) +highconstraints n ((a :<= b)) = (:) $ case highexpr n (a.-.b) of (terms,((-1)*)-> c) -> Inequality terms $ LessEqual (int2Double c) + +highexpr n (Constant (Number m)) = ([],m n) +highexpr n ((Number m) :* v) = ([int2Double (m n) .* v],0) +highexpr n (a :+ b) = case (highexpr n a, highexpr n b) of + ((t1,c),(t2,d)) -> (t1 `merge` t2,c+d) + where + merge [] xs = xs + merge xs [] = xs + merge (Term a x:xs) (Term b y:ys) + | x == y = Term (a+b) x : merge xs ys + | x < y = Term a x : merge xs (Term b y:ys) + | y < x = Term b y : merge (Term a x:xs) ys + | otherwise = error "simple math" diff --git a/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/Labels.hs b/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/Labels.hs new file mode 100644 index 000000000..05886d4fc --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/Labels.hs @@ -0,0 +1,678 @@ +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE ViewPatterns #-} +{- | +Module : Data.Array.Accelerate.Trafo.Partitioning.ILP.LabelsNew +Description : Nodes representing nodes in the graph. + +This module provides the labels that represent nodes in the graph. A node can +either be a computation or a buffer. +-} +module Data.Array.Accelerate.Trafo.Partitioning.ILP.Labels where + +import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.IdxSet (IdxSet) +import qualified Data.Array.Accelerate.AST.IdxSet as IdxSet +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.AST.Operation +import Data.Array.Accelerate.Error +import Data.Array.Accelerate.Array.Buffer +import Data.Array.Accelerate.Representation.Type + +import Lens.Micro +import Lens.Micro.Mtl + +import Data.Set (Set) +import qualified Data.Set as S + +import Data.Hashable (Hashable, hashWithSalt) +import Prelude hiding (exp) + +import qualified Data.Functor.Const as C +import Data.Coerce +import Control.Monad.State.Strict +import Data.Kind (Type) +import Data.Maybe (fromJust) +import Data.List ( intercalate ) +import Debug.Trace +import Data.Array.Accelerate.Analysis.Match +import qualified Data.Array.Accelerate.Pretty.Operation as P +import qualified Data.Array.Accelerate.Pretty.Exp as P +import Data.String (fromString) + + +-------------------------------------------------------------------------------- +-- Nodes +-------------------------------------------------------------------------------- + +data Comp -- ^ The type of computation nodes. +data GVal -- ^ The type of ground value nodes. + + +-- | Nodes for referencing nodes. +data Node t where + Node :: Int -- ^ The computation node. + -> Parent -- ^ The parent computation. + -> Node t + + +-- | The parent of a node is either 'Nothing' (the root node) or 'Just' (a parent computation). +-- Parent nodes are set when a computation is nested within for instance an Acond or Awhile. +-- See the use of 'zoom (scope ...)' in Graph.hs +type Parent = Maybe (Node Comp) + + +-- | Lens for getting and setting the nodes unique identifier. +nodeId :: Lens' (Node t) Int +nodeId f (Node i p) = f i <&> (`Node` p) + + +-- | Lens for getting and setting the parent node. +parent :: Lens' (Node t) Parent +parent f (Node i p) = f p <&> Node i + + +-- | Lens for setting and unsafely getting the parent. +parent' :: Lens' (Node t) (Node Comp) +parent' f (Node i p) = f (fromJust p) <&> (Node i . Just) + + +-- | Lens for interpreting any node as a computation node. +asComp :: Lens' (Node t) (Node Comp) +asComp f l = coerce <$> f (coerce l) + + +-- | Lens for interpreting any node as a buffer node. +asBuff :: Lens' (Node t) (Node GVal) +asBuff f l = coerce <$> f (coerce l) + +instance Show (Node Comp) where + show :: Node Comp -> String + show c = "C" ++ intercalate "." (map show . reverse $ nodeIds c) + +instance Show (Node GVal) where + show :: Node GVal -> String + show b = "B" ++ intercalate "." (map show . reverse $ nodeIds b) + + +nodeIds :: Node t -> [Int] +nodeIds l = l^.nodeId : maybe [] nodeIds (l^.parent) + + +instance Eq (Node t) where + (==) :: Node t -> Node t -> Bool + (==) l1 l2 = (l1^.nodeId == l2^.nodeId) && checkMismatch (l1^.parent) (l2^.parent) True + + +instance Ord (Node t) where + compare :: Node t -> Node t -> Ordering + compare l1 l2 = case compare (l1^.nodeId) (l2^.nodeId) of + EQ -> checkMismatch (l1^.parent) (l2^.parent) EQ + ord -> ord + + +-- | Checks if two parents are equal and throw an error if they are not. +checkMismatch :: Parent -> Parent -> a -> a +checkMismatch (Just l1) (Just l2) | l1 == l2 = id +checkMismatch Nothing Nothing = id +checkMismatch _ _ = internalError "mismatching nodes" + + +instance Hashable (Node t) where + hashWithSalt :: Int -> Node t -> Int + hashWithSalt s l = hashWithSalt s (l^.nodeId) + + +-- | Compute the nesting level of a 'Node'. +level :: Node t -> Int +level n = case n^.parent of + Nothing -> 0 + Just p -> 1 + level p + + +-- | Compute the ancestry tree of a 'Node', staring with the ancestor underneath the root. +ancestree :: Node Comp -> [Node Comp] +ancestree n = ancestree' [n] n + + +-- | Helper function for 'ancestree'. +ancestree' :: [Node Comp] -> Node t -> [Node Comp] +ancestree' ps n = case n^.parent of + Nothing -> ps + Just p -> ancestree' (p:ps) p + + +-- | Compute the difference between the 'ancestree's of two 'Node's. +difftree :: Node Comp -> Node Comp -> ([Node Comp], [Node Comp]) +difftree n1 n2 = difftree' (ancestree n1) (ancestree n2) + + +-- | Compute the difference between two 'ancestree's. +difftree' :: [Node Comp] -> [Node Comp] -> ([Node Comp], [Node Comp]) +difftree' (x:xs) (y:ys) | x == y = difftree' xs ys +difftree' xs ys = (xs, ys) + + +-- | The first nodes in the 'ancestree's of the two 'Node's that differ, which +-- may be one or both of the original 'Node's, but always two different 'Node's. +ancestors :: Node Comp -> Node Comp -> Maybe (Node Comp, Node Comp) +ancestors n1 n2 = case ancestors' n1 n2 of + (a1, a2) | a1 /= a2 -> Just (a1, a2) + | otherwise -> Nothing + + +-- | The first nodes in the 'ancestree's of the two 'Node's that differ, which +-- may be one or both of the original 'Node's. +ancestors' :: Node Comp -> Node Comp -> (Node Comp, Node Comp) +ancestors' n1 n2 = if n1^.parent == n2^.parent then (n1, n2) + else case compare (level n1) (level n2) of + LT -> ancestors' n1 (n2^.parent') + GT -> ancestors' (n1^.parent') n2 + EQ -> ancestors' (n1^.parent') (n2^.parent') + + +-- | Create a fresh 'Node'. +freshL' :: State (Node t) (Node t) +freshL' = id <%= (nodeId +~ 1) + + +-- | Set of 'Node's. +type Nodes t = Set (Node t) + +type ReadEdge = (Node GVal, Node Comp) +type WriteEdge = (Node Comp, Node GVal) +type StrictEdge = (Node Comp, Node Comp) +type DataflowEdge = (Node Comp, Node GVal, Node Comp) +type FusibleEdge = DataflowEdge +type InfusibleEdge = DataflowEdge +type InplacePath = (ReadEdge, WriteEdge) + +-- | A value consists of its type @s t@ and and the 'Node' that represents it. +data Val s t = Val + { valType :: s t -- Type is needed to make mkReindexPartial safe + , valNode :: Node GVal + } + +-- TODO: Remove this function, and let all callers directly work with a single +-- Node instead of a Set +valNodes :: Val s t -> Nodes GVal +valNodes (Val _ n) = S.singleton n + +-- | A 'TupR' of 'Val's. +type Vals s = TupR (Val s) + + +-- | The type of ground values. +type GroundVal = Val GroundR + + +-- | A 'TupR' of 'GroundVal's. +type GroundVals = Vals GroundR + + +val :: s t -> Node GVal -> Val s t +val t n = Val t n + + +-- | Get the nodes of 'Vals'. +valsNodes :: Vals s t -> Nodes GVal +valsNodes = foldMapTupR valNodes + + +-- | Get the type of 'Vals'. +valsType :: Vals s t -> TupR s t +valsType = mapTupR valType + + +-- | Match the types of two 'GroundVals'. +matchGroundValsType :: GroundVals s -> GroundVals t -> Maybe (s :~: t) +matchGroundValsType = matchTupR matchGroundValType + + +-- | Match the types of two 'GroundVal's. +matchGroundValType :: GroundVal s -> GroundVal t -> Maybe (s :~: t) +matchGroundValType (Val t1 _) (Val t2 _) = matchGroundR t1 t2 + + +-- | Expect 'GroundVals' of the given type. +expectGroundVals :: HasCallStack => GroundsR t -> Exists GroundVals -> GroundVals t +expectGroundVals repr (Exists vals) + | Just Refl <- matchGroundsR repr (groundsR vals) = vals + | otherwise = internalError $ fromString $ + "Result of lambda has unexpected type.\nExpected: " + ++ show (P.prettyTupR (const P.prettyGroundR) 0 repr) + ++ "\nActual: " + ++ show (P.prettyTupR (const P.prettyGroundR) 0 $ groundsR vals) + + +instance HasGroundsR GroundVals where + groundsR :: GroundVals a -> GroundsR a + groundsR = valsType + + +instance HasGroundsR GroundVal where + groundsR :: GroundVal a -> GroundsR a + groundsR (Val tp _) = TupRsingle tp + + +instance Eq (Val s t) where + (==) :: Val s t -> Val s t -> Bool + (==) (Val _ n1) (Val _ n2) = n1 == n2 + +instance Show (Val s t) where + show :: Val s t -> String + show (Val _ n) = "Val " ++ show n + + +-------------------------------------------------------------------------------- +-- Labelled Environment +-------------------------------------------------------------------------------- + +-- | An 'EnvLabel' uniquely identifies an element of the environment. +newtype EnvLabel = EnvLabel Int + deriving (Eq, Ord, Show, Num) + + +-- | A 'TupR' of 'EnvLabel'. +type EnvLabels = TupR (C.Const EnvLabel) + + +-- | Create a fresh 'EnvLabel' from the current state. +freshE' :: State EnvLabel EnvLabel +freshE' = id <%= (+1) + + +-- | An environment value consists of a unique 'EnvLabel', the stored 'GroundVals' +-- and the 'Uniquenesses' of the stored values. +type EnvVal t = (EnvLabel, GroundVals t, Uniquenesses t) + + +-- | A collection of multiple 'EnvVal's. Individual elements can be accessed by +-- pattern matching on 'EnvLabels'. +type EnvVals t = (EnvLabels t, GroundVals t, Uniquenesses t) + + +-- | The environment used during graph construction. +data Env env where + -- | The empty environment. + EnvNil :: Env () + -- | The non-empty environment. + (:>>:) :: EnvVal t -- ^ See 'EnvVal'. + -> Env env -- ^ The rest of the environment. + -> Env (env, t) + + +instance Show (Env env) where + show :: Env env -> String + show EnvNil = "EnvNil" + show (envl :>>: env) = show envl ++ " :>>: " ++ show env + + +-- | Constructs a new 'Env' by prepending labels for each element in the left-hand side. +weakenEnv :: LeftHandSide s v env env' -> GroundVals v -> Uniquenesses v -> Env env -> State EnvLabel (Env env') +weakenEnv LeftHandSideWildcard{} _ _ = pure +weakenEnv LeftHandSideSingle{} bs us = \lenv -> freshE' >>= \e -> return ((e, bs, us) :>>: lenv) +weakenEnv (LeftHandSidePair l r) (TupRpair lbs rbs) (TupRpair lus rus) = weakenEnv l lbs lus >=> weakenEnv r rbs rus +weakenEnv (LeftHandSidePair _ _) _ _ = internalError "mismatching left-hand side" + + +-- | Look up 'Vars' in 'Env', returing 'EnvVals'. +lookupVars :: Vars a env b -> Env env -> EnvVals b +lookupVars TupRunit _ = (TupRunit, TupRunit, TupRunit) +lookupVars (TupRsingle var) env | (e, bs, u) <- lookupVar var env + = (TupRsingle (C.Const e), bs, u) +lookupVars (TupRpair l r) env | (el, bsl, ul) <- lookupVars l env + , (er, bsr, ur) <- lookupVars r env + = (TupRpair el er, TupRpair bsl bsr, TupRpair ul ur) + + +-- | Look up a 'Var' in 'Env', returning an 'EnvVal'. +lookupVar :: Var a env b -> Env env -> EnvVal b +lookupVar = lookupIdx . varIdx + + +-- | Look up an 'Idx' in 'Env', returning an 'EnvVal'. +lookupIdx :: Idx env t -> Env env -> EnvVal t +lookupIdx ZeroIdx (bs :>>: _) = bs +lookupIdx (SuccIdx idx) (_ :>>: env) = lookupIdx idx env + + + +-------------------------------------------------------------------------------- +-- Bound left-hand side +-------------------------------------------------------------------------------- + +-- | A 'LeftHandSide' and the 'EnvVal' bound by it. +data BoundLHS s v env env' where + BoundLHSsingle + :: EnvVal v + -> s v + -> BoundLHS s v env (env, v) + + BoundLHSwildcard + :: TupR s v + -> BoundLHS s v env env + + BoundLHSpair + :: BoundLHS s v1 env env' + -> BoundLHS s v2 env' env'' + -> BoundLHS s (v1, v2) env env'' + +instance Show (BoundLHS s v env env') where + show :: BoundLHS s v env env' -> String + show (BoundLHSsingle e _) = "BLHS(" ++ show e ++ ")" + show (BoundLHSwildcard _) = "BLHS_" + show (BoundLHSpair l r) = "BLHS(" ++ show l ++ ", " ++ show r ++ ")" + + +-- | The type of left-hand sides binding ground values. +type BoundGLHS = BoundLHS GroundR + + +-- | Bind values to the 'LeftHandSide' that produced the 'Env' +bindLHS :: LeftHandSide s v env env' -> Env env' -> BoundLHS s v env env' +bindLHS (LeftHandSideSingle sv) (l :>>: _) = BoundLHSsingle l sv +bindLHS (LeftHandSideWildcard tr) _ = BoundLHSwildcard tr +bindLHS (LeftHandSidePair l r) env = BoundLHSpair (bindLHS l (stripLHS r env)) (bindLHS r env) + + +-- | Unbind values from the 'BoundLHS'. +unbindLHS :: BoundLHS s v env env' -> LeftHandSide s v env env' +unbindLHS (BoundLHSsingle _ sv) = LeftHandSideSingle sv +unbindLHS (BoundLHSwildcard tr) = LeftHandSideWildcard tr +unbindLHS (BoundLHSpair l r) = LeftHandSidePair (unbindLHS l) (unbindLHS r) + + +-- | Strip the values bound by the 'LeftHandSide' from the 'Env'. +stripLHS :: LeftHandSide s v env env' -> Env env' -> Env env +stripLHS (LeftHandSideSingle _) (_ :>>: le') = le' +stripLHS (LeftHandSideWildcard _) le = le +stripLHS (LeftHandSidePair l r) le = stripLHS l (stripLHS r le) + + +createLHS :: BoundLHS s v env1 env2 + -> Env env + -> (forall env'. Env env' -> LeftHandSide s v env env' -> r) + -> r +createLHS (BoundLHSsingle e sv) env k = k (e :>>: env) (LeftHandSideSingle sv) +createLHS (BoundLHSwildcard tr) env k = k env (LeftHandSideWildcard tr) +createLHS (BoundLHSpair l r) env k = + createLHS l env $ \env' l' -> + createLHS r env' $ \env'' r' -> + k env'' (LeftHandSidePair l' r') + + + +-------------------------------------------------------------------------------- +-- Labelled Arguments +-------------------------------------------------------------------------------- + +-- | A label to add to function arguments. +-- +-- If the argument is an array, then we need all information about the array and its +-- shape from the environment. +-- If the argument is not an array, e.g. a scalar, function or expression, then we only +-- need 'Node's referenced by the scalar, function or expression. +data ArgLabel t where + -- | The argument is an array. + Arr :: EnvVals (Buffers e) -- ^ The array values. + -> EnvVals sh -- ^ The shape values. + -> ArgLabel (m sh e) + -- | The argument is a scalar 'Var'', 'Exp'' or 'Fun''. + NotArr :: Nodes GVal -- ^ The variables referenced by the argument. + -> ArgLabel (t e) + +deriving instance Show (ArgLabel t) + + +-- | Get the set of dependent buffers of an 'ArgLabel'. +getLabelDeps :: ArgLabel t -> Nodes GVal +getLabelDeps (Arr (_, arr, _) (_, sh, _)) = valsNodes arr <> valsNodes sh +getLabelDeps (NotArr deps) = deps + + +-- | Get the set of unique array dependencies of an 'ArgLabel'. +getLabelUniqueArrDeps :: ArgLabel t -> Nodes GVal +getLabelUniqueArrDeps (Arr (_, arr, u) _) = uniqueNodes u arr +getLabelUniqueArrDeps (NotArr _) = internalError "getLabelUniqueArrDeps: Expected Arr but got NotArr" + + +-- | Given 'Uniquenesses', get the unique nodes from 'GroundVals'. +uniqueNodes :: Uniquenesses e -> GroundVals e -> Nodes GVal +uniqueNodes TupRunit TupRunit = mempty +uniqueNodes (TupRsingle Shared) _ = mempty +uniqueNodes (TupRsingle Unique) bs = valsNodes bs +uniqueNodes (TupRpair ul ur) (TupRpair l r) = uniqueNodes ul l <> uniqueNodes ur r +uniqueNodes _ _ = internalError "uniqueNodes: Tuple mismatch " + + +-- | Get the arrays of an 'ArgLabel'. +getLabelArrays :: ArgLabel (m sh e) -> GroundVals (Buffers e) +getLabelArrays (Arr (_, arr, _) (_, _, _)) = arr +getLabelArrays (NotArr _) = internalError "getLabelArrays: Expected Arr but got NotArr" + + +-- | Get the array dependencies of an 'ArgLabel'. +getLabelArrDeps :: forall (m :: Type -> Type -> Type) sh e. ArgLabel (m sh e) -> Nodes GVal +getLabelArrDeps = valsNodes . getLabelArrays + + +-- | Get a single array dependency of an 'ArgLabel'. +getLabelArrDep :: forall (m :: Type -> Type -> Type) sh e. ArgLabel (m sh e) -> Node GVal +getLabelArrDep = foldr1 const . getLabelArrDeps + + +-- | Get the shapes of an 'ArgLabel'. +getLabelShape :: ArgLabel (m sh e) -> GroundVals sh +getLabelShape (Arr (_, _, _) (_, sh, _)) = sh +getLabelShape (NotArr _) = internalError "getLabelShape: Expected Arr but got NotArr" + + +-- | Get the shape dependencies of an 'ArgLabel'. +getLabelShDeps :: forall (m :: Type -> Type -> Type) sh e. ArgLabel (m sh e) -> Nodes GVal +getLabelShDeps = valsNodes . getLabelShape + + +-- | Check if two arguments have the same shape. +sameShape :: forall (m1 :: Type -> Type -> Type) (m2 :: Type -> Type -> Type) sh1 sh2 e1 e2. ArgLabel (m1 sh1 e1) -> ArgLabel (m2 sh2 e2) -> Bool +sameShape (getLabelShape -> sh1) (getLabelShape -> sh2) + | Just Refl <- matchGroundValsType sh1 sh2 = sh1 == sh2 + | otherwise = False + + +-- | The argument to a function paired with an 'ArgLabel' +data LabelledArg env t = L (Arg env t) (ArgLabel t) + deriving (Show) + + +-- | Arguments to a function paired with their 'ArgLabel's. +type LabelledArgs env = PreArgs (LabelledArg env) + + +-- | Node the arguments to a function using the given environment. +labelArgs :: Args env args -> Env env -> LabelledArgs env args +labelArgs (arg :>: args) env = labelArg arg env :>: labelArgs args env +labelArgs ArgsNil _ = ArgsNil + + +-- | Get the 'ArgLabels' associated with 'Arg' from 'Env'. +labelArg :: Arg env t -> Env env -> LabelledArg env t +labelArg arg env = L arg $ case arg of + (ArgVar vars) -> NotArr $ getVarsDeps vars env + (ArgExp exp) -> NotArr $ getExpDeps exp env + (ArgFun fun) -> NotArr $ getFunDeps fun env + (ArgArray _ _ sh arr) -> + Arr (lookupVars arr env) (lookupVars sh env) + + + +-- | Get the dependencies of a tuple of variables. +getVarsDeps :: Vars s env t -> Env env -> Nodes GVal +getVarsDeps vars = valsNodes . (^._2) . lookupVars vars + + +-- | Get the dependencies of a variable. +getVarDeps :: Var s env t -> Env env -> Nodes GVal +getVarDeps var = valsNodes . (^._2) . lookupVar var + +-- | Get the dependencies, given the index of a variable. +getIdxDeps :: Idx env t -> Env env -> Nodes GVal +getIdxDeps idx = valsNodes . (^._2) . lookupIdx idx + +-- | Get the dependencies of a set of indices. +getIdxSetDeps :: IdxSet env -> Env env -> Nodes GVal +getIdxSetDeps deps env = mconcat $ map (\(Exists idx) -> getIdxDeps idx env) $ IdxSet.toList deps + + +-- | Get the dependencies of an expression. +getExpDeps :: OpenExp x env y -> Env env -> Nodes GVal +getExpDeps (ArrayInstr (Index var) poe) env = getVarDeps var env <> getExpDeps poe env +getExpDeps (ArrayInstr (Parameter var) poe) env = getVarDeps var env <> getExpDeps poe env +getExpDeps (Let _ poe1 poe2) env = getExpDeps poe1 env <> getExpDeps poe2 env +getExpDeps (Evar _) _ = mempty +getExpDeps Foreign{} _ = mempty +getExpDeps (Pair poe1 poe2) env = getExpDeps poe1 env <> getExpDeps poe2 env +getExpDeps Nil _ = mempty +getExpDeps (VecPack _ poe) env = getExpDeps poe env +getExpDeps (VecUnpack _ poe) env = getExpDeps poe env +getExpDeps (ToIndex _ poe1 poe2) env = getExpDeps poe1 env <> getExpDeps poe2 env +getExpDeps (FromIndex _ poe1 poe2) env = getExpDeps poe1 env <> getExpDeps poe2 env +getExpDeps (Case poe1 poes poe2) env = getExpDeps poe1 env <> + foldMap ((`getExpDeps` env) . snd) poes <> + maybe mempty (`getExpDeps` env) poe2 +getExpDeps (Cond poe1 poe2 exp3) env = getExpDeps poe1 env <> + getExpDeps poe2 env <> + getExpDeps exp3 env +getExpDeps (Select poe1 poe2 exp3) env = getExpDeps poe1 env <> + getExpDeps poe2 env <> + getExpDeps exp3 env +getExpDeps (While pof1 pof2 poe) env = getFunDeps pof1 env <> + getFunDeps pof2 env <> + getExpDeps poe env +getExpDeps (Const _ _) _ = mempty +getExpDeps (PrimApp _ poe) env = getExpDeps poe env +getExpDeps (ShapeSize _ poe) env = getExpDeps poe env +getExpDeps (Undef _) _ = mempty +getExpDeps Coerce{} _ = mempty +getExpDeps (Assert _ g e) env = getExpDeps g env <> getExpDeps e env +getExpDeps (Assume g e) env = getExpDeps g env <> getExpDeps e env + + +-- | Get the dependencies of a function. +getFunDeps :: OpenFun x env y -> Env env -> Nodes GVal +getFunDeps (Body poe) env = getExpDeps poe env +getFunDeps (Lam _ fun) env = getFunDeps fun env + + + +-------------------------------------------------------------------------------- +-- Helpers for Labelled Arguments +-------------------------------------------------------------------------------- + +-- | Map a function over the labelled arguments. +mapLArgs :: (forall s. LabelledArg env s -> LabelledArg env s) -> LabelledArgs env t -> LabelledArgs env t +mapLArgs _ ArgsNil = ArgsNil +mapLArgs f (larg :>: largs) = f larg :>: mapLArgs f largs + +-- | Fold over the labelled arguments and combine the resulting monoidal values. +foldMapLArgs :: Monoid m => (forall s. LabelledArg env s -> m) -> LabelledArgs env t -> m +foldMapLArgs _ ArgsNil = mempty +foldMapLArgs f (larg :>: largs) = f larg <> foldMapLArgs f largs + +-- | Map a monadic function over the labelled arguments. +mapLArgsM :: Monad m => (forall s. LabelledArg env s -> m (LabelledArg env s)) -> LabelledArgs env t -> m (LabelledArgs env t) +mapLArgsM _ ArgsNil = return ArgsNil +mapLArgsM f (larg :>: largs) = do + larg' <- f larg + largs' <- mapLArgsM f largs + return (larg' :>: largs') + +-- | Flipped version of 'mapLArgsM'. +forLArgsM :: Monad m => LabelledArgs env t -> (forall s. LabelledArg env s -> m (LabelledArg env s)) -> m (LabelledArgs env t) +forLArgsM largs f = mapLArgsM f largs +{-# INLINE forLArgsM #-} + +-- | Map a monadic action over the labelled arguments and discard the result. +mapLArgsM_ :: Monad m => (forall s. LabelledArg env s -> m ()) -> LabelledArgs env t -> m () +mapLArgsM_ _ ArgsNil = return () +mapLArgsM_ f (larg :>: largs) = f larg >> mapLArgsM_ f largs + +-- | Flipped version of 'mapLArgsM_'. +forLArgsM_ :: Monad m => LabelledArgs env t -> (forall s. LabelledArg env s -> m ()) -> m () +forLArgsM_ largs f = mapLArgsM_ f largs +{-# INLINE forLArgsM_ #-} + +-- | Map a monadic function over the labelled arguments and accumulate the result. +mapAccumLArgsM :: Monad m => (forall s. a -> LabelledArg env s -> m (a, LabelledArg env s)) -> a -> LabelledArgs env t -> m (a, LabelledArgs env t) +mapAccumLArgsM _ a ArgsNil = return (a, ArgsNil) +mapAccumLArgsM f a (larg :>: largs) = do + (acc' , larg') <- f a larg + (acc'', largs') <- mapAccumLArgsM f acc' largs + return (acc'', larg' :>: largs') + +-- | Flipped version of 'mapAccumLArgsM'. +forAccumLArgsM :: Monad m => a -> LabelledArgs env t -> (forall s. a -> LabelledArg env s -> m (a, LabelledArg env s)) -> m (a, LabelledArgs env t) +forAccumLArgsM a largs f = mapAccumLArgsM f a largs +{-# INLINE forAccumLArgsM #-} + +-- | Traverse over the labelled arguments. +traverseLArgs :: Applicative f => (forall s. LabelledArg env s -> f (LabelledArg env s)) -> LabelledArgs env t -> f (LabelledArgs env t) +traverseLArgs _ ArgsNil = pure ArgsNil +traverseLArgs f (larg :>: largs) = (:>:) <$> f larg <*> traverseLArgs f largs + +-- | Flipped version of 'traverseLArgs'. +forLArgs :: Applicative f => LabelledArgs env t -> (forall s. LabelledArg env s -> f (LabelledArg env s)) -> f (LabelledArgs env t) +forLArgs largs f = traverseLArgs f largs +{-# INLINE forLArgs #-} + +-- | Traverse over the labelled arguments and discard the result. +traverseLArgs_ :: Applicative f => (forall s. LabelledArg env s -> f ()) -> LabelledArgs env t -> f () +traverseLArgs_ _ ArgsNil = pure () +traverseLArgs_ f (larg :>: largs) = f larg *> traverseLArgs_ f largs + +-- | Flipped version of 'traverseLArgs_'. +forLArgs_ :: Applicative f => LabelledArgs env t -> (forall s. LabelledArg env s -> f ()) -> f () +forLArgs_ largs f = traverseLArgs_ f largs +{-# INLINE forLArgs_ #-} + +-- | All arrays that the function reads from. +inputArrays :: LabelledArgs env t -> Nodes GVal +inputArrays = foldMapLArgs \case + L (ArgArray In _ _ _) (Arr (_,arr,_) _) -> valsNodes arr + L (ArgArray Mut _ _ _) (Arr (_,arr,_) _) -> valsNodes arr + _ -> mempty + +-- | All arrays that the function writes to. +outputArrays :: LabelledArgs env t -> Nodes GVal +outputArrays = foldMapLArgs \case + L (ArgArray Out _ _ _) (Arr (_,arr,_) _) -> valsNodes arr + L (ArgArray Mut _ _ _) (Arr (_,arr,_) _) -> valsNodes arr + _ -> mempty + +-- | All non-array arguments and array shapes. +notArrays :: LabelledArgs env t -> Nodes GVal +notArrays = foldMapLArgs \case + L _ (Arr _ (_,sh,_)) -> valsNodes sh + L _ (NotArr deps) -> deps + +-- | Fold map over all inputs. +foldMapInputLabels :: Monoid m => (forall sh e. ArgLabel (In sh e) -> m) -> LabelledArgs env t -> m +foldMapInputLabels f = foldMapLArgs \case + L (ArgArray In _ _ _) l -> f l + _ -> mempty + +-- | Fold map over all outputs. +foldMapOutputLabels :: Monoid m => (forall sh e. ArgLabel (Out sh e) -> m) -> LabelledArgs env t -> m +foldMapOutputLabels f = foldMapLArgs \case + L (ArgArray Out _ _ _) l -> f l + _ -> mempty + + + +-------------------------------------------------------------------------------- +-- Debugging +-------------------------------------------------------------------------------- + +-- | Trace a value using a function to format the output. +traceWith :: (Show a) => (a -> String) -> a -> a +traceWith f x = trace (f x) x +{-# INLINE traceWith #-} diff --git a/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/LinearConstraint.hs b/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/LinearConstraint.hs new file mode 100644 index 000000000..5a96be529 --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/LinearConstraint.hs @@ -0,0 +1,291 @@ +module Data.Array.Accelerate.Trafo.Partitioning.ILP.LinearConstraint + ( Constants (..), Number (..), Expression (..), IsNumber (..) + , LinearConstraint (..), Bounds (..), Var(..) + , (.+.), (.-.), (.*.), times, timesN, nCompsE, var + , (.>=.), (.<=.), (.==.), (.>.), (.<.), allEqual + , between, binary, lowerUpper, lower, upper + , equal, notB, impliesB, andB, allB, orB, anyB, isEqualRangeN, isEqualRange + , packB, coverB, partitionB + ) where + +import Data.Array.Accelerate.Error +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Var (Var(..)) +import Data.Foldable + +data Constants = Constants + { -- | number of computations in the ILP + nComps :: Int, + -- | number of buffers in the ILP + nBuffs :: Int + } + deriving (Show) + +-- | given `n` (for the number of nodes in the ILP), make an Int +newtype Number = Number (Constants -> Int) + +instance Show Number where + show :: Number -> String + show (Number f) = "Number {" ++ show (f $ Constants 1 1) ++ "}" + +instance Num Number where + (+) :: Number -> Number -> Number + (Number f) + (Number g) = Number (\c -> f c + g c) + + (*) :: Number -> Number -> Number + (Number f) * (Number g) = Number (\c -> f c * g c) + + negate :: Number -> Number + negate (Number f) = Number (negate . f) + + abs :: Number -> Number + abs (Number f) = Number (abs . f) + + signum :: Number -> Number + signum (Number f) = Number (signum . f) + + fromInteger :: Integer -> Number + fromInteger i = Number (\_ -> fromInteger i) + +data Expression where + Constant :: Number -> Expression + (:+) :: Expression -> Expression -> Expression + (:*) :: Number -> Var -> Expression + deriving (Show) + +instance Semigroup Expression where + (<>) :: Expression -> Expression -> Expression + (<>) a (b :+ c) = (a <> b) <> c + (<>) a b = a :+ b + +instance Monoid Expression where + mempty :: Expression + mempty = int 0 + +-- | Add two expressions. +(.+.) :: Expression -> Expression -> Expression +(.+.) = (<>) + +infixl 8 .+. + +-- | Subtract two expressions. +(.-.) :: Expression -> Expression -> Expression +e1 .-. e2 = e1 .+. ((-1) .*. e2) + +infixl 8 .-. + +-- | Multiply an expression by a constant. +(.*.) :: Number -> Expression -> Expression +i .*. (Constant j) = Constant $ i * j +i .*. (e1 :+ e2) = (:+) (i .*. e1) (i .*. e2) +i .*. (j :* v) = (:*) (i * j) v + +infixl 8 .*. + +-- | Multiply by one of the @Constants@. +times :: (Constants -> Int) -> Expression -> Expression +times f = (Number f .*.) + +-- | Multiply by @n@ (the total number of computations + some safety margin). +-- +-- This is only here because the old definitions used timesN and not all of them +-- have been replaced yet. +-- TODO: Replace all occurrences of timesN with tighter bounds. +timesN :: Expression -> Expression +timesN = times ((+ 10) . (* 2) . nComps) + +-- | Total number of computations. +nCompsE :: Expression +nCompsE = Constant $ Number nComps + +-- | Use a 'Var' in an 'Expression'. +var :: Var -> Expression +var = (Number (const 1) :*) + +class IsNumber a where + -- | Use an `Int` as `Expression` or `Number`. + int :: Int -> a + +instance IsNumber Number where + int = Number . const + +instance IsNumber Expression where + int = Constant . Number . const + +data LinearConstraint where + (:>=) :: Expression -> Expression -> LinearConstraint + (:<=) :: Expression -> Expression -> LinearConstraint + (:==) :: Expression -> Expression -> LinearConstraint + (:&&) :: LinearConstraint -> LinearConstraint -> LinearConstraint + TrueConstraint :: LinearConstraint + deriving (Show) + +instance Semigroup LinearConstraint where + (<>) :: LinearConstraint -> LinearConstraint -> LinearConstraint + (<>) TrueConstraint b = b + (<>) a TrueConstraint = a + (<>) a (b :&& c) = (a <> b) <> c + (<>) a b = a :&& b + +instance Monoid LinearConstraint where + mempty :: LinearConstraint + mempty = TrueConstraint + +-- | @x >= y@ +(.>=.) :: Expression -> Expression -> LinearConstraint +(.>=.) = (:>=) + +infixr 7 .>=. + +-- | @x <= y@ +(.<=.) :: Expression -> Expression -> LinearConstraint +(.<=.) = (:<=) + +infixr 7 .<=. + +-- | @x == y@ +(.==.) :: Expression -> Expression -> LinearConstraint +(.==.) = (:==) + +infixr 7 .==. + +-- | @x[0] == x[1] == ... == x[n-1]@ +allEqual :: [Expression] -> LinearConstraint +allEqual [] = TrueConstraint +allEqual (x : xs) = foldMap (x .==.) xs + +-- | @x < y@ +(.>.) :: Expression -> Expression -> LinearConstraint +x .>. y = x .>=. (y .+. int 1) + +infixl 7 .>. + +-- | @x < y@ +(.<.) :: Expression -> Expression -> LinearConstraint +x .<. y = (x .+. int 1) .<=. y + +infixl 7 .<. + +-- | @x <= y <= z@ +between :: Expression -> Expression -> Expression -> LinearConstraint +between x y z = x .<=. y <> y .<=. z + +data Bounds where + Binary :: Var -> Bounds + LowerUpper :: Int -> Var -> Int -> Bounds + Lower :: Int -> Var -> Bounds + Upper :: Var -> Int -> Bounds + (:<>) :: Bounds -> Bounds -> Bounds + NoBounds :: Bounds + deriving (Show) + +instance Semigroup Bounds where + (<>) :: Bounds -> Bounds -> Bounds + (<>) NoBounds b = b + (<>) a NoBounds = a + (<>) (a :<> b) c = a <> b <> c + (<>) a b = a :<> b + +instance Monoid Bounds where + mempty :: Bounds + mempty = NoBounds + +-- | 'Var' is binary. +binary :: Var -> Bounds +binary = Binary + +-- | 'Var' is bounded by lower and upper bounds. +lowerUpper :: Int -> Var -> Int -> Bounds +lowerUpper = LowerUpper + +-- | 'Var' is bounded by lower bound. +lower :: Int -> Var -> Bounds +lower = Lower + +-- | 'Var' is bounded by upper bound. +upper :: Var -> Int -> Bounds +upper = Upper + +-- | 'Var' is equal to a constant. +equal :: Int -> Var -> Bounds +equal x v = lowerUpper x v x + +-- | Not 'Expression' (i.e. 1 - 'Expression'). +notB :: Expression -> Expression +notB e = int 1 .-. e + +-- | If a is 0, then b is 0. +impliesB :: Expression -> Expression -> LinearConstraint +impliesB = (.>=.) + +-- -- | Iff a and b are 0, then r is 0. +-- andB :: Expression -> Expression -> Expression -> LinearConstraint +-- andB a b r = orB (notB a) (notB b) (notB r) + +-- | Iff a and b are 0, then r is 0. +-- +-- Source: "Formulating Integer Linear Programs: A Rogues' Gallery", B3 +andB :: Expression -> Expression -> Expression -> LinearConstraint +andB a b r = + r .<=. a .+. b + <> r .>=. a + <> r .>=. b + +-- | Iff all xs are 0, then r is 0. +allB :: (Foldable f) => f Expression -> Expression -> LinearConstraint +allB xs r + | null xs = TrueConstraint + | otherwise = + r .<=. fold xs + <> foldMap (r .>=.) xs + +-- -- | Iff a and b are 1, then r is 1. +-- -- not sure if this encoding is new, nor whether it is the simplest, but I think it works. +-- -- perhaps defining andB is easier than defining orB? +-- orB :: Expression -> Expression -> Expression -> LinearConstraint +-- orB a b r = +-- (2 .*. r .<=. a .+. b) -- r can only be 1 if both a and b are 1, so this line fixes 3/4 cases +-- <> +-- (r .+. int 1 .>=. a .+. b) -- and this line forces r to be 1 if a and b are both 1, while not restricting the other cases + +-- | Iff a and b are 1, then r is 1. +-- +-- Source: "Formulating Integer Linear Programs: A Rogues' Gallery", B2 +orB :: Expression -> Expression -> Expression -> LinearConstraint +orB a b r = + r .+. int 1 .>=. a .+. b + <> r .<=. a + <> r .<=. b + +-- | Iff all xs are 1, then r is 1. +anyB :: (Foldable f) => f Expression -> Expression -> LinearConstraint +anyB xs r + | null xs = TrueConstraint + | otherwise = + r .+. int (length xs - 1) .>=. fold xs + <> foldMap (r .<=.) xs + +isEqualRangeN :: Expression -> Expression -> Expression -> LinearConstraint +isEqualRangeN = isEqualRange timesN + +-- given a function f that multiplies by the size of the domain of a and b, r can only be 0(true) when a and b are equal +-- note that r can always be 1 +isEqualRange :: (Expression -> Expression) -> Expression -> Expression -> Expression -> LinearConstraint +isEqualRange f a b r = between (a .-. f r) b (a .+. f r) + +-- | From a set of booleans, select at most n to be 0 (true). +packB :: (Foldable f) => Int -> f Expression -> LinearConstraint +packB n xs + | n >= length xs = TrueConstraint + | n >= 0 = fold xs .>=. int (length xs - n) + | otherwise = internalError "packB: always false" + +-- | From a set of booleans, select at least n to be 0 (true). +coverB :: (Foldable f) => Int -> f Expression -> LinearConstraint +coverB n xs + | n <= 0 = TrueConstraint + | n <= length xs = fold xs .<=. int (length xs - n) + | otherwise = internalError "coverB: always false" + +-- | From a set of booleans, select exactly n to be 0 (true). +partitionB :: (Foldable f) => Int -> f Expression -> LinearConstraint +partitionB n xs = packB n xs <> coverB n xs diff --git a/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/Lower.hs b/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/Lower.hs new file mode 100644 index 000000000..b2eb95cf6 --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/Lower.hs @@ -0,0 +1,89 @@ +{-# OPTIONS_HADDOCK hide #-} +module Data.Array.Accelerate.Trafo.Partitioning.ILP.Lower (LowerEnv (..), Lower, lowerAll, lower) where + +import Control.Monad (replicateM) +import Control.Monad.State (State, evalState) +import Data.Array.Accelerate.Trafo.Partitioning.ILP.ConstraintLanguage +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Graph (fused, inFoldSize, inplace, manifest, maxCluster, outFoldSize, pi, pimax, readDir, readDirs, writeDir, writeDirs) +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Var (Var (Other)) +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Labels (Comp, GVal, Node, nodeId) +import Data.Array.Accelerate.Trafo.Partitioning.ILP.LinearConstraint (Bounds, Expression, LinearConstraint, allB, allEqual, between, binary, impliesB, int, isEqualRangeN, lowerUpper, nCompsE, notB, packB, timesN, var, (.+.), (.-.), (.<.), (.<=.), (.==.), (.>.), (.>=.)) +import Data.Array.Accelerate.Trafo.Partitioning.ILP.NameGeneration (freshName) +import Data.Foldable (fold) +import Data.Monoid (Ap (..)) +import Lens.Micro ((^.)) +import Prelude hiding (pi) + +-- | An environment for lowering, contains the number of computations +newtype LowerEnv = LowerEnv {lowerEnvNComps :: Int} + +-- | The result of lowering. Contains the generated constraints, the bounds of the variables and the cost expression. +-- State monad is used to generate fresh variable names. +type Lower = State String (LinearConstraint, Bounds, Expression) + +-- | Lower a batch of 'Constraint's under one shared name supply. +lowerAll :: LowerEnv -> [Constraint] -> (LinearConstraint, Bounds, Expression) +lowerAll env constraints = evalState (mconcat <$> mapM (lower env) constraints) "" + +-- | Lower a single 'Constraint'. +lower :: LowerEnv -> Constraint -> Lower +lower env = \case + ClusterBefore i j -> pure (pi i .<. pi j, mempty, mempty) + DifferentCluster i j -> pure (fused (i, j) .==. int 1, mempty, mempty) + NotManifestIfAllFused b pairs -> pure (allB (map fused pairs) (notB $ manifest b), mempty, mempty) + FusibleOrder i j -> pure (between (fused (i, j)) (pi j .-. pi i) (timesN $ fused (i, j)), mempty, mempty) + FusionDirection w b r -> pure (isEqualRangeN (writeDir (w, b)) (readDir (b, r)) (fused (w, r)), mempty, mempty) + WithinClusterCount l -> pure (pi l .<=. maxCluster, mempty, mempty) + OnManifestIfInPlace p@((b1, _), (_, b2)) -> pure ((inplace p `impliesB` manifest b1) <> (inplace p `impliesB` manifest b2), mempty, mempty) + InPlaceDirection p@(r, w) -> pure (isEqualRangeN (readDir r) (writeDir w) (inplace p), mempty, mempty) + InPlaceCluster p@((b1, _), (c2, _)) -> pure ((pimax b1 .-. pi c2) .<=. timesN (inplace p), mempty, mempty) + AcrossClusterSame p@((_, c1), (c2, _)) + | c1 == c2 -> pure (mempty, mempty, mempty) + | otherwise -> pure (isEqualRangeN (pi c1) (pi c2) (inplace p), mempty, mempty) + AtMostOneReader ps -> pure (packB 1 (map inplace ps), mempty, mempty) + AtMostOneWriter ps -> pure (packB 1 (map inplace ps), mempty, mempty) + ReadAliveThroughWriters r@(b1, c1) ws -> + pure (pi c1 .+. int 1 .-. foldMap (\w -> int 1 .-. inplace (r, w)) ws .<=. pimax b1, mempty, mempty) + HorizontalReadCost pairs -> lowerHorizontalReadCost env pairs + Manifest b -> pure (manifest b .==. int 0, mempty, mempty) + NoInPlace b -> pure (pimax b .>=. nCompsE, mempty, mempty) + SameFoldSize c -> pure (inFoldSize c .==. outFoldSize c, mempty, mempty) + NewFoldSize c -> pure (outFoldSize c .==. int (c ^. nodeId), mempty, mempty) + SameFoldSizeIfFused w c -> pure (isEqualRangeN (inFoldSize c) (outFoldSize w) (fused (w, c)), mempty, mempty) + SameDirection rs ws -> pure (allEqual (readDirs rs <> writeDirs ws), mempty, mempty) + PinnedDirection c rs ws -> pure (allEqual ([int (c ^. nodeId)] <> readDirs rs <> writeDirs ws), mempty, mempty) + NegativeDirIfManifest (w, b) -> pure (timesN (manifest b) .>. writeDir (w, b), mempty, mempty) + +-- | Variables for the horizontal read cost constraint. +readOrderVar, readPiVar, readPi0Var, useVar :: State String Var +readOrderVar = Other <$> freshName "ReadOrder" +readPiVar = Other <$> freshName "ReadPi" +readPi0Var = Other <$> freshName "Read0Pi" +useVar = Other <$> freshName "ReadUse" + +-- | Lower 'HorizontalReadCost' constraints. +lowerHorizontalReadCost :: LowerEnv -> [(Node GVal, Node Comp)] -> Lower +lowerHorizontalReadCost env consumers = do + let nConsumers = length consumers + n = lowerEnvNComps env + readPis <- replicateM nConsumers readPiVar + readOrders <- replicateM nConsumers readOrderVar + (subConstraint, subBounds) <- flip foldMapM consumers $ \(buff, cons) -> do + useVars <- replicateM nConsumers useVar + let c = + foldMap + ( \(uv, rp, ro) -> + isEqualRangeN (var rp) (pi cons) (var uv) + <> isEqualRangeN (var ro) (readDir (buff, cons)) (var uv) + ) + (zip3 useVars readPis readOrders) + pure (c <> foldl (.+.) (int 0) (map var useVars) .<=. int (nConsumers - 1), foldMap binary useVars) + readPi0s <- replicateM nConsumers readPi0Var + let cost = foldl (.+.) (int 0) (map var readPi0s) + constraint' = subConstraint <> fold (zipWith (\p p0 -> var p .<=. timesN (var p0)) readPis readPi0s) + bounds = subBounds <> foldMap (\v -> lowerUpper 0 v n) readPis <> foldMap binary readPi0s + pure (constraint', bounds, cost) + +-- | 'foldMap' with an effectful mapping function. +foldMapM :: (Foldable t, Applicative f, Monoid m) => (a -> f m) -> t a -> f m +foldMapM f = getAp . foldMap (Ap . f) diff --git a/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/MIP.hs b/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/MIP.hs new file mode 100644 index 000000000..e4d895f25 --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/MIP.hs @@ -0,0 +1,116 @@ +{-# LANGUAGE UndecidableInstances #-} +{-# OPTIONS_GHC -Wno-orphans #-} -- Shame on me! + +module Data.Array.Accelerate.Trafo.Partitioning.ILP.MIP ( + -- Exports default paths to 6 solvers, as well as an instance to ILPSolver for all of them + cbc, cplex, glpsol, gurobiCl, lpSolve, scip, + MIP(..) + ) where + +import Data.Array.Accelerate.Trafo.Partitioning.ILP.NameGeneration + +import Data.Array.Accelerate.Trafo.Partitioning.ILP.LinearConstraint hiding (var) +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Solver +import qualified Data.Map as M + +import Numeric.Optimization.MIP hiding (Bounds, Constraint, Var, Solution, name) +import qualified Numeric.Optimization.MIP as MIP +import qualified Numeric.Optimization.MIP.Solver.Base as MIP +import Data.Scientific ( Scientific ) +import Data.Bifunctor (bimap) + +import Numeric.Optimization.MIP.Solver + ( cbc, cplex, glpsol, gurobiCl, lpSolve, scip ) +import Data.Maybe (mapMaybe) +import Control.Monad.State ( runState ) +import Control.Monad.Reader ( Reader, asks, runReader ) +import Data.String (fromString) +import Data.Text (unpack) + +newtype MIP s = MIP s + +instance (MIP.IsSolver s IO) => ILPSolver (MIP s) where + solvePartial :: MIP s -> ILP -> IO (Maybe Solution) + solvePartial (MIP s) ilp@(ILP dir obj constr bnds n) = makeSolution names <$> MIP.solve s options problem + where + options = def { MIP.solveTimeLimit = Nothing --Just 60 + -- , MIP.solveLogger = putStrLn . ("AccILPSolver: " ++) + , MIP.solveErrorLogger = putStrLn . ("AccILPSolverError: " ++) + } --, MIP.solveCondensedSolution = False } + vs = allVars ilp + ((),(names, _)) = runState (mapM_ var' vs) ((mempty, mempty),"") + + readerProblem = Problem (Just "AccelerateILP") . mkFun dir + <$> expr n obj + <*> cons n constr + <*> pure [] + <*> pure [] + <*> (M.map (IntegerVariable,) <$> bounds bnds) + problem = runReader readerProblem names + + -- varsOf :: [MIP.Constraint Scientific] -> [MIP.Var] + -- varsOf = concatMap $ concatMap (\(Term _ vs)->vs) . (\(Expr ts)->ts) . MIP.constrExpr + + mkFun Maximise = ObjectiveFunction (Just "AccelerateObjective") OptMax + mkFun Minimise = ObjectiveFunction (Just "AccelerateObjective") OptMin + + + -- addZeroes :: MIP.Problem Scientific -> MIP.Solution Scientific -> MIP.Solution Scientific + -- addZeroes problem (MIP.Solution stat obj solmap) = + -- -- Map.union is left-biased: only values not present in the solution are added. + -- MIP.Solution stat obj $ M.union solmap (M.fromSet (const 0) (vars problem)) + +var :: Var -> Reader Names MIP.Var +var y = asks (fromString . (M.! y) . snd) + +-- MIP has a Num instance for expressions, but it's scary (because +-- you can't guarantee linearity with arbitrary multiplications). +-- We use that instance here, knowing that our own Expression can only be linear. +expr :: Constants -> Expression -> Reader Names (MIP.Expr Scientific) +expr n (Constant (Number f)) = pure $ fromIntegral (f n) +expr n (x :+ y) = (+) <$> expr n x <*> expr n y +expr n ((Number f) :* y) = (fromIntegral (f n) *) . varExpr <$> var y + +cons :: Constants -> LinearConstraint -> Reader Names [MIP.Constraint Scientific] +cons n (x :>= y) = (\a b -> [a MIP..>=. b]) <$> expr n x <*> expr n y +cons n (x :<= y) = (\a b -> [a MIP..<=. b]) <$> expr n x <*> expr n y +cons n (x :== y) = (\a b -> [a MIP..==. b]) <$> expr n x <*> expr n y + +cons n (x :&& y) = (<>) <$> cons n x <*> cons n y +cons _ TrueConstraint = pure mempty + +bounds :: Bounds -> Reader Names (M.Map MIP.Var (Extended Scientific, Extended Scientific)) +bounds (Binary v) = (`M.singleton` (Finite 0, Finite 1)) <$> var v +bounds (Lower l v ) = (`M.singleton` (Finite (fromIntegral l), PosInf )) <$> var v +bounds ( Upper v u) = (`M.singleton` (NegInf , Finite (fromIntegral u))) <$> var v +bounds (LowerUpper l v u) = (`M.singleton` (Finite (fromIntegral l), Finite (fromIntegral u))) <$> var v +bounds (x :<> y) = M.unionWith (\(l1, u1) (l2, u2) -> (max l1 l2, min u1 u2)) <$> bounds x <*> bounds y +bounds NoBounds = pure mempty + +-- -- For all variables not yet in bounds, we add infinite bounds. This is apparently required. +-- -- Potentially, it's more efficient to simply make a bounds map giving (NegInf, PosInf) to all variables (like in `allIntegers`), and then use `unionWith const`? +-- finishBounds :: M.Map MIP.Var (Extended Scientific, Extended Scientific) -> Reader Names (M.Map MIP.Var (Extended Scientific, Extended Scientific)) +-- finishBounds x = do +-- vars' <- asks $ map toVar . M.keys . fst +-- let y = M.keys x +-- return $ x <> (M.fromList . map (,(NegInf,PosInf)) . filter (not . (`elem` y)) $ vars') + +-- -- we currently have no variables that ever get less then -2 +-- extraConstraints :: Reader Names [MIP.Constraint Scientific] +-- extraConstraints = do +-- vs <- asks $ map toVar . M.keys . fst +-- return [MIP.constExpr (-5) MIP..<=. MIP.varExpr x | x <- vs] + + + +makeSolution :: Names -> MIP.Solution Scientific -> Maybe Solution +-- ------- Matching on solutions with a value: If this is Nothing, the model was infeasable or unbounded. +-- | -- Instead matching on `MIP.Solution StatusOptimal _ m` often works too, but that doesn't work for +-- v -- e.g. the identity program (which has an empty ILP). +makeSolution names (MIP.Solution _ (Just _) m) = Just . M.fromList . mapMaybe (sequence' . bimap (\v -> runReader (unvar' $ unpack $ varName v) names) round) $ M.toList m +makeSolution _ _ = Nothing + +-- tuple's traversable instance works on the second argument, we need it on the first +sequence' :: (Maybe a, b) -> Maybe (a, b) +sequence' (Nothing, _) = Nothing +sequence' (Just x, y) = Just (x, y) diff --git a/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/NameGeneration.hs b/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/NameGeneration.hs new file mode 100644 index 000000000..9e5d4f268 --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/NameGeneration.hs @@ -0,0 +1,46 @@ +module Data.Array.Accelerate.Trafo.Partitioning.ILP.NameGeneration where + +import Control.Monad.State ( State, gets, modify ) +import Data.Char ( ord ) +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Var (Var) +import qualified Data.Map as M +import Lens.Micro.Mtl (zoom) +import Lens.Micro (_2) +import Control.Monad.Reader +import Data.Bifunctor + + +-- Apparently, solvers don't appreciate variable names longer than 255 characters! +-- Instead, we generate small placeholders here and store their meaning +type Names = (M.Map String Var, M.Map Var String) +type STN = State (Names, String) +-- to avoid generating keywords, we simply prepend every name with an 'x'. This still leaves 26^254 options, more than enough! +freshName' :: STN String +freshName' = zoom _2 $ freshName "x" + +-- Iterates in the following order: +-- ["", "a", .., "z", "aa", .., "za", "ab", .., "zz", "aaa", "baa", ...] +-- The prefix disembiguates from other uses of freshName, and avoids generating keywords. +freshName :: String -> State String String +freshName prefix = do + modify increment + gets ((prefix <> "_") <>) + where + increment (char:cs) + | ord char < ord 'z' = toEnum (ord char + 1) : cs + | otherwise = 'a' : increment cs + increment "" = "a" + +var' :: Var -> STN String +var' v = do + maybeName <- gets $ (M.!? v) . snd . fst + case maybeName of + Just name -> return name + Nothing -> do + name <- freshName' + modify $ first $ bimap (M.insert name v) (M.insert v name) + return name + +unvar' :: String -> Reader Names (Maybe Var) +unvar' name = asks $ (M.!? name) . fst + diff --git a/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/Solver.hs b/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/Solver.hs new file mode 100644 index 000000000..7b7e20978 --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/Solver.hs @@ -0,0 +1,59 @@ +module Data.Array.Accelerate.Trafo.Partitioning.ILP.Solver where + +import qualified Data.Map as M +import qualified Data.Set as S +import Data.Array.Accelerate.Trafo.Partitioning.ILP.LinearConstraint + + +-- Currently the only instance is for MIP, which gives bindings to a couple of solvers. +-- Still, this way we minimise the surface that has to interact with MIP, can more easily +-- adapt if it changes, and we could easily add more bindings. +class ILPSolver ilp where + solvePartial :: ilp -> ILP -> IO (Maybe Solution) + + +solve :: ILPSolver ilp => ilp -> ILP -> IO (Maybe Solution) +solve x ilp = fmap (<> M.fromSet (const 0) (allVars ilp)) -- add zeroes to the ILP for missing variables + <$> solvePartial x (finalize ilp) + +-- adds potentially missing constraints and bounds: +-- some solvers require all variables to have a bound +-- or all variables to be in a constraint. +finalize :: ILP -> ILP +finalize ilp@(ILP dir obj constr bnds n) = + ILP dir obj (constr <> extraconstr) (bnds <> extrabnds) n + where + extraconstr = foldMap (\v -> int (-5) .<=. var v) (allVars ilp) + extrabnds = foldMap (Lower (-5)) (allVars ilp) + +data OptDir = Maximise | Minimise + deriving (Show, Eq) + +data ILP = ILP OptDir Expression LinearConstraint Bounds Constants + deriving (Show) + +type Solution = M.Map Var Int + +-- helpers for solving an ILP +allVars :: ILP -> S.Set Var +allVars (ILP _ cost constraint bounds _) = varsExpr cost <> varsConstr constraint <> varsBounds bounds + +varsExpr :: Expression -> S.Set Var +varsExpr (Constant _) = mempty +varsExpr (a :+ b) = varsExpr a <> varsExpr b +varsExpr (_ :* v) = S.singleton v + +varsConstr :: LinearConstraint -> S.Set Var +varsConstr TrueConstraint = mempty +varsConstr (a :&& b) = varsConstr a <> varsConstr b +varsConstr (a :>= b) = varsExpr a <> varsExpr b +varsConstr (a :== b) = varsExpr a <> varsExpr b +varsConstr (a :<= b) = varsExpr a <> varsExpr b + +varsBounds :: Bounds -> S.Set Var +varsBounds NoBounds = mempty +varsBounds (a :<> b) = varsBounds a <> varsBounds b +varsBounds (Binary v) = S.singleton v +varsBounds (LowerUpper _ v _) = S.singleton v +varsBounds (Lower _ v) = S.singleton v +varsBounds (Upper v _) = S.singleton v diff --git a/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/Var.hs b/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/Var.hs new file mode 100644 index 000000000..10d106673 --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Partitioning/ILP/Var.hs @@ -0,0 +1,43 @@ +-- | The decision variables of the fusion ILP. +module Data.Array.Accelerate.Trafo.Partitioning.ILP.Var (Var (..)) where + +import Data.Array.Accelerate.Trafo.Partitioning.ILP.Labels (Comp, GVal, Node) + +data Var + -- Variables used by fusion: + = Pi (Node Comp) + -- ^ Used for acyclic ordering of clusters. + -- Pi (Node x y) = z means that computation number x (possibly a subcomputation of y, see Node) is fused into cluster z (y ~ Just i -> z is a subcluster of the cluster of i) + | Fused (Node Comp) (Node Comp) + -- ^ 0 is fused (same cluster), 1 is unfused. We do *not* have one of these for all pairs, only the ones we need for constraints and/or costs! + -- Invariant: Like edges, both labels have to have the same parent: Either on top (Node _ Nothing) or as sub-computation of the same label (Node _ (Just x)). + -- In fact, this is the Var-equivalent to Edge: an infusible edge has a constraint (== 1). + | IsManifest (Node GVal) + -- ^ 0 means manifest, 1 is like a `delayed array`. + -- Binary variable; will we write the output to a manifest array, or is it fused away (i.e. all uses are in its cluster)? + | ReadDir (Node GVal) (Node Comp) + -- ^ \-3 can't fuse with anything, -2 for 'left to right', -1 for 'right to left', n for 'unknown', see computation n (currently only backpermute). + | WriteDir (Node Comp) (Node GVal) + -- ^ See 'ReadDir'. + | InFoldSize (Node Comp) -- Legacy? Probably needs per-edge equivalent + -- ^ Keeps track of the fold that's one dimension larger than this operation, and is fused in the same cluster. + -- This prevents something like @zipWith f (fold g xs) (fold g ys)@ from illegally fusing + | OutFoldSize (Node Comp) -- Legacy? Probably needs per-edge equivalent + -- ^ Keeps track of the fold that's one dimension larger than this operation, and is fused in the same cluster. + -- This prevents something like @zipWith f (fold g xs) (fold g ys)@ from illegally fusing + | Other String + -- ^ For one-shot variables that don't deserve a constructor. These are also integer variables, and the responsibility is on the user to pick a unique name! + -- It is possible to add a variation for continuous variables too, see `allIntegers` in MIP.hs. + -- We currently use this in Lower.hs for cost functions. + + -- Variables introduced for in-place updates: + | InPlace (Node GVal) (Node Comp) (Node Comp) (Node GVal) + -- ^ 0 means in-place, 1 means not in-place. The first label is an input of a cluster, the second label is an output of a cluster. + -- All 'InPlace' variables need to be unique, so we can't omit the computation labels. Taking one path through a cluster is different from taking another. + | PiMax (Node GVal) + -- ^ The cluster number of the largest reader of the buffer, since in-place updates are only allowed on the final consumer of an array/buffer. + | MaxCluster + -- ^ Upper bound on the larget 'Pi' assigned to any comptation, used by the 'NumClusters'/'Everything' objectives to approximate the cluster count. + -- | WriteDirPiMax (Node GVal) + -- -- ^ The write direction of the largest reader of the buffer. This is used to check that all reads of the buffer are in the same direction as the write. + deriving (Eq, Ord, Show) diff --git a/src/Data/Array/Accelerate/Trafo/Partitioning/Stream.hs b/src/Data/Array/Accelerate/Trafo/Partitioning/Stream.hs new file mode 100644 index 000000000..ea6fb8851 --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Partitioning/Stream.hs @@ -0,0 +1,10 @@ +module Data.Array.Accelerate.Trafo.Partitioning.Stream where + +import Data.Array.Accelerate.AST.Partitioned +import Data.Array.Accelerate.AST.Operation + + + +-- Emulate the old 'delayed arrays' approach to fusion TODO +streamfusion :: OperationAcc op () a -> PartitionedAcc op () a +streamfusion = undefined diff --git a/src/Data/Array/Accelerate/Trafo/Schedule/Partial.hs b/src/Data/Array/Accelerate/Trafo/Schedule/Partial.hs new file mode 100644 index 000000000..275b89af7 --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Schedule/Partial.hs @@ -0,0 +1,981 @@ +{-# LANGUAGE TypeFamilies #-} +{-# OPTIONS_HADDOCK hide #-} +-- | +-- Module : Data.Array.Accelerate.Trafo.Schedule.Partial +-- Copyright : [2012..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.Trafo.Schedule.Partial ( + convertAccPartial, convertAfunPartial, + PartialSchedule(..), PartialScheduleFun, + PrePartialSchedule(..), PrePartialScheduleFun(..), + Parallelism(..), + Sync(..), SyncEnv, + SyncSchedule(..), SyncScheduleFun, + Loop, + variablesToSyncEnv, + + UnitTuple, UpdateTuple(..), toPartialReturn, + + compileKernel', CompiledKernel(..), +) where + +import Data.Array.Accelerate.Analysis.Match +import Data.Array.Accelerate.AST.Environment +import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.IdxSet (IdxSet) +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.AST.Kernel +import Data.Array.Accelerate.AST.Schedule.Uniform +import Data.Array.Accelerate.Error +import Data.Array.Accelerate.Representation.Array +import Data.Array.Accelerate.Representation.Shape +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Trafo.Var +import Data.Array.Accelerate.Trafo.Substitution +import Data.Array.Accelerate.Trafo.Exp.Substitution +import Data.Array.Accelerate.Type +import qualified Data.Array.Accelerate.AST.IdxSet as IdxSet +import Data.Array.Accelerate.AST.Partitioned (PartitionedAcc, PartitionedAfun, Clustered) +import qualified Data.Array.Accelerate.AST.Partitioned as C +import Data.Text (Text) +import Data.Kind +import Data.Maybe +import Prelude hiding (id, (.), read) +import Control.Category + +import Data.Functor.Identity +-- Constructs a partial schedule. It is partial as many details still need +-- to be filled in for the actual schedule. This schedule does however +-- decide whether the binding and body of a let-binding should be executed +-- in parallel or sequentially. +-- When implemention a schedule, this module may provide a good first step. +convertAccPartial :: IsKernel kernel => PartitionedAcc (KernelOperation kernel) env t -> SyncSchedule kernel env t +convertAccPartial = analyseSyncEnv . rebuild 50 . toPartial + +convertAfunPartial :: IsKernel kernel => PartitionedAfun (KernelOperation kernel) env f -> SyncScheduleFun kernel env f +convertAfunPartial (C.Alam lhs f) = Plam lhs $ convertAfunPartial f +convertAfunPartial (C.Abody body) = Pbody $ convertAccPartial body + +-- Sync environment +type SyncEnv = PartialEnv Sync + +data Sync t where + SyncRead :: Sync (Buffer e) + SyncWrite :: Sync (Buffer e) + +instance Eq (Sync t) where + SyncRead == SyncRead = True + SyncWrite == SyncWrite = True + _ == _ = False + +instance Ord (Sync t) where + SyncRead < SyncWrite = True + _ < _ = False + + SyncWrite <= SyncRead = False + _ <= _ = True + +weakenSyncEnv :: GLeftHandSide t env env' -> SyncEnv env' -> SyncEnv env +weakenSyncEnv _ PEnd = PEnd +weakenSyncEnv (LeftHandSideWildcard _) env = env +weakenSyncEnv (LeftHandSideSingle _) (PPush env _) = env +weakenSyncEnv (LeftHandSideSingle _) (PNone env) = env +weakenSyncEnv (LeftHandSidePair l1 l2) env = weakenSyncEnv l1 $ weakenSyncEnv l2 env + +-- Intermediate representation, used during the transformation + +data Parallelism + -- This let-binding is parallel: the binding is spawned in a separate task, and the body is directly executed. + = Parallel + -- This let-binding is sequential. The binding is executed before the body. + | Sequential + deriving Eq + +data PrePartialScheduleFun (schedule :: (Type -> Type) -> Type -> Type -> Type) kernel env t where + Plam :: GLeftHandSide s env env' + -> PrePartialScheduleFun schedule kernel env' t + -> PrePartialScheduleFun schedule kernel env (s -> t) + + Pbody :: schedule kernel env t + -> PrePartialScheduleFun schedule kernel env t + +data PrePartialSchedule schedule kernel env t where + PExec + :: KernelFun kernel s + -> SArgs env s + -> PrePartialSchedule schedule kernel env () + + PLet + :: Parallelism -- Should this be executed in parallel? + -> GLeftHandSide bnd env env' + -> Uniquenesses bnd + -> schedule kernel env bnd + -> schedule kernel env' t + -> PrePartialSchedule schedule kernel env t + + PAlloc + :: ShapeR sh + -> ScalarType e + -> ExpVars env sh + -> PrePartialSchedule schedule kernel env (Buffer e) + + PUse + :: ScalarType e + -> Int -- Number of elements + -> Buffer e + -> PrePartialSchedule schedule kernel env (Buffer e) + + PUnit + :: ExpVar env e + -> PrePartialSchedule schedule kernel env (Buffer e) + + PCompute + :: Exp env e + -> PrePartialSchedule schedule kernel env e + + -- End of control flow. Returns no values, that happens in PReturnValues + PReturnEnd + :: TupR Void' t -- Tuple without TupRsingle + -> PrePartialSchedule schedule kernel env t + + PReturnValues + :: UpdateTuple env t1 t2 + -> schedule kernel env t1 + -> PrePartialSchedule schedule kernel env t2 + + PAcond + :: ExpVar env PrimBool + -> schedule kernel env t + -> schedule kernel env t + -> PrePartialSchedule schedule kernel env t + + PAwhile + :: Uniquenesses t + -> PrePartialScheduleFun schedule kernel env (t -> Loop t) + -> GroundVars env t + -> PrePartialSchedule schedule kernel env t + + PAtrace + :: Text + -> ArrayDescriptors env t + -> PrePartialSchedule schedule kernel env Word8 + + -- Signals that this while loop should do another iteration. + -- The state for the next iteration is computed by the subterm. + PContinue + :: schedule kernel env t + -> PrePartialSchedule schedule kernel env (Loop t) + + -- Stops this while loop. The result value of the loop is + -- given by the GroundVars. + PBreak + :: Uniquenesses t + -> GroundVars env t + -> PrePartialSchedule schedule kernel env (Loop t) + + PAssert + :: Text + -> Exp env PrimBool + -> PrePartialSchedule schedule kernel env PrimBool + + PFence + :: IdxSet env + -> schedule kernel env t + -> PrePartialSchedule schedule kernel env t + +data Void' a where +data Loop a where + +data UpdateTuple env t1 t2 where + UpdateKeep :: UpdateTuple env t t + UpdatePair + :: UpdateTuple env left1 left2 + -> UpdateTuple env right1 right2 + -> UpdateTuple env (left1, right1) (left2, right2) + UpdateSet + :: Uniqueness t + -> GroundVar env t + -> UpdateTuple env () t + +updatePair + :: UpdateTuple env left1 left2 + -> UpdateTuple env right1 right2 + -> UpdateTuple env (left1, right1) (left2, right2) +updatePair UpdateKeep UpdateKeep = UpdateKeep +updatePair up1 up2 = UpdatePair up1 up2 + +updateTupVarList :: forall env t1 t2. UpdateTuple env t1 t2 -> [Exists (Idx env)] +updateTupVarList = (`go` []) + where + go :: UpdateTuple env s1 s2 -> [Exists (Idx env)] -> [Exists (Idx env)] + go UpdateKeep accum = accum + go (UpdateSet _ var) accum = Exists (varIdx var) : accum + go (UpdatePair up1 up2) accum = go up1 $ go up2 accum + +updateTupSync :: forall env t1 t2. UpdateTuple env t1 t2 -> [EnvBinding Sync env] +updateTupSync = (`go` []) + where + go :: UpdateTuple env s1 s2 -> [EnvBinding Sync env] -> [EnvBinding Sync env] + go UpdateKeep accum = accum + go (UpdateSet Unique var) accum = EnvBinding (varIdx var) SyncWrite : accum + go (UpdateSet _ (Var (GroundRbuffer _) idx)) accum = EnvBinding idx SyncRead : accum + go (UpdateSet _ _) accum = accum + go (UpdatePair up1 up2) accum = go up1 $ go up2 accum + +type PartialScheduleFun = PrePartialScheduleFun PartialSchedule +newtype PartialSchedule kernel env t = PartialSchedule (PrePartialSchedule PartialSchedule kernel env t) + +type SyncScheduleFun = PrePartialScheduleFun SyncSchedule +data SyncSchedule kernel env t = + SyncSchedule { + syncEnv :: SyncEnv env, + -- Whether this term introduces no branching and no forks + syncSimple :: Bool, + syncSchedule :: (PrePartialSchedule SyncSchedule kernel env t) + } + +-- Tuple where all the fields are replaced with unit. +type family UnitTuple t where + UnitTuple (a, b) = (UnitTuple a, UnitTuple b) + UnitTuple a = () + +unitTupleForVars :: GroundVars env t -> TupR Void' (UnitTuple t) +unitTupleForVars (TupRpair v1 v2) = unitTupleForVars v1 `TupRpair` unitTupleForVars v2 +unitTupleForVars TupRunit = TupRunit +-- Pattern match on the type of the variable to proof that t isn't a pair. +unitTupleForVars (TupRsingle (Var tp _)) + | Refl <- groundNotPair tp = TupRunit + +groundNotPair :: GroundR t -> UnitTuple t :~: () +groundNotPair (GroundRbuffer _) = Refl +groundNotPair (GroundRscalar (VectorScalarType _)) = Refl +groundNotPair (GroundRscalar (SingleScalarType (NumSingleType (IntegralNumType t)))) = case t of + TypeInt -> Refl + TypeInt8 -> Refl + TypeInt16 -> Refl + TypeInt32 -> Refl + TypeInt64 -> Refl + TypeWord -> Refl + TypeWord8 -> Refl + TypeWord16 -> Refl + TypeWord32 -> Refl + TypeWord64 -> Refl +groundNotPair (GroundRscalar (SingleScalarType (NumSingleType (FloatingNumType t)))) = case t of + TypeHalf -> Refl + TypeFloat -> Refl + TypeDouble -> Refl + +data HoistReturn env env' t1 t3 where + HoistReturn + :: UpdateTuple env t2 t3 + -> UpdateTuple env' t1 t2 + -> HoistReturn env env' t1 t3 + +hoistReturn :: GLeftHandSide bnd env env' -> SyncEnv env -> UpdateTuple env' t1 t3 -> HoistReturn env env' t1 t3 +hoistReturn lhs bndSyncEnv = \case + UpdateKeep -> HoistReturn UpdateKeep UpdateKeep + UpdatePair upL upR + | HoistReturn upL1 upL2 <- hoistReturn lhs bndSyncEnv upL + , HoistReturn upR1 upR2 <- hoistReturn lhs bndSyncEnv upR -> + HoistReturn (updatePair upL1 upR1) (updatePair upL2 upR2) + UpdateSet u var + | Just idx <- strengthenWithLHS lhs $ varIdx var + , canHoist u (prjPartial idx bndSyncEnv) -> + -- We can hoist this return + HoistReturn (UpdateSet u $ Var (varType var) idx) UpdateKeep + up -> HoistReturn UpdateKeep up + where + canHoist :: Uniqueness t -> Maybe (Sync t) -> Bool + canHoist _ Nothing = True -- Variable is not used in this term, so we can hoist the return to before this term. + canHoist Shared (Just SyncRead) = True -- A shared buffer may still be read after returning it. + canHoist _ _ = False -- All other combinations are not allowed. + +-- A function that gets the available variables (i.e., the variables that don't have to be waited on) +-- and then returns the constructed schedule (PartialSchedule or PartialScheduleFun) +-- together with the set of variables the function directly waits on. +-- These two sets should have no overlap. +-- Also returns whether a term is trivial, meaning that the term is computationally inexpensive. +-- Invariant: If a schedule is trivial, then the directly-awaits set must +-- equal the set of all free variables (minus the already available variables). +type Build p (kernel :: Type -> Type) env t = IdxSet env -> Built p kernel env t + +data Built p (kernel :: Type -> Type) env t = + Built { + -- Did the program change this iteration, in a way that new information may be available + -- in the next iteration? + didChange :: !Bool, + -- The set of variables that are directly required to start this computation. + -- I.e., there is no use in starting this computation if some of these variables + -- are not available yet. This information is used to decide if a let binding + -- should be executed in parallel or sequentially. + directlyAwaits :: IdxSet env, + -- The set of buffer variables that may be written to. + writes :: IdxSet env, + -- The set of buffer variables that are released by a term at the end of its execution. + -- 'Released' here means that the term wrote to those buffers and is only ready writing + -- to them at the end of the execution. + -- This should be a subset of 'writes'. + -- This information can be used to decide if a let binding should be sequential or parallel. + -- If the binding of a let 'finally releases' a buffer variable and the body 'directly awaits' it, + -- the let should be executed sequentially. + finallyReleases :: IdxSet env, + trivial :: Bool, + term :: !(p kernel env t) + } + +data Exists' (a :: (Type -> Type -> Type) -> Type) where + Exists' :: a m -> Exists' a + +combineMod :: Modifier m -> Modifier m' -> Exists' Modifier +combineMod In In = Exists' In +combineMod Out Out = Exists' Out +combineMod _ _ = Exists' Mut +-- combineMod _ _ = error "Remove this error once we add in place updates" + +combineAccessGroundR :: AccessGroundR t -> AccessGroundR t -> AccessGroundR t +combineAccessGroundR (AccessGroundRbuffer m1 tp) (AccessGroundRbuffer m2 _) + | Exists' m <- combineMod m1 m2 = AccessGroundRbuffer m tp +combineAccessGroundR a _ = a -- Nothing has to be done when combining two scalars; they don't have access modifiers + +data CompiledKernel kenv fenv kernel where + CompiledKernel :: OpenKernelFun kernel kenv f + -> SArgs fenv f + -> CompiledKernel kenv fenv kernel + +compileKernel' + :: forall fenv kernel args. + IsKernel kernel + => Clustered (KernelOperation kernel) args + -> Args fenv args + -> CompiledKernel () fenv kernel +compileKernel' cluster args = + go + (partialEnvFromList combineAccessGroundR (map (\(Exists (Var tp ix)) -> EnvBinding ix tp) vars)) + weakenId + Empty + (\ _ x -> x) + where + vars = argsVars args + + go + :: forall fenv1 kenv. + PartialEnv AccessGroundR fenv1 + -> (fenv1 :> fenv) + -> Env AccessGroundR kenv + -> (forall kenv'. kenv :> kenv' -> PartialEnv (Idx kenv') fenv1 -> PartialEnv (Idx kenv') fenv) + -> CompiledKernel kenv fenv kernel + go PEnd _ kenv kSubst = + CompiledKernel + (KernelFunBody $ compileKernel kenv cluster args') + ArgsNil + where + kSubst' = kSubst weakenId PEnd + Identity args' = reindexArgs prjIdx args + prjIdx :: forall a. Idx fenv a -> Identity (Idx kenv a) + prjIdx idx = case prjPartial idx kSubst' of + Nothing -> internalError "Variable not found in environment. The environment was probably built incorrectly, argsVars may have missed this variable?" + Just idx' -> Identity idx' + go (PPush env a@(AccessGroundRscalar tp)) k1 kenv kSubstF + | CompiledKernel kernel sargs <- go env (weakenSucc k1) (kenv `Push` a) (\k2 kSubst -> kSubstF (weakenSucc k2) $ PPush kSubst (k2 >:> ZeroIdx)) + = CompiledKernel (KernelFunLam (KernelArgRscalar tp) kernel) (SArgScalar (Var tp $ k1 >:> ZeroIdx) :>: sargs) + go (PPush env a@(AccessGroundRbuffer mod' tp)) k1 kenv kSubstF + | CompiledKernel kernel sargs <- go env (weakenSucc k1) (kenv `Push` a) (\k2 kSubst -> kSubstF (weakenSucc k2) $ PPush kSubst (k2 >:> ZeroIdx)) + = CompiledKernel (KernelFunLam (KernelArgRbuffer mod' tp) kernel) (SArgBuffer mod' (Var (GroundRbuffer tp) $ k1 >:> ZeroIdx) :>: sargs) + go (PNone env) k1 kenv kSubstF + = go env (weakenSucc k1) kenv (\k2 kSubst -> kSubstF k2 $ PNone kSubst) + +toPartial :: IsKernel kernel => PartitionedAcc (KernelOperation kernel) env t -> PartialSchedule kernel env t +toPartial = snd . toPartial' (TupRsingle Shared) + +toPartial' :: IsKernel kernel => Uniquenesses t -> PartitionedAcc (KernelOperation kernel) env t -> (IdxSet env, PartialSchedule kernel env t) +toPartial' us = \case + C.Exec cluster args + | CompiledKernel kernel sargs <- compileKernel' cluster args -> + ( IdxSet.fromList $ sargBufferVars sargs + , PartialSchedule $ PExec kernel sargs ) + C.Return vars -> + ( IdxSet.fromVarList $ filter isBuffer $ flattenTupR vars + , returnValues (toPartialReturn us vars) $ PartialSchedule $ PReturnEnd $ unitTupleForVars vars ) + C.Manifest var -> toPartial' us (C.Return $ TupRsingle var) + C.Compute expr -> + ( IdxSet.fromList $ mapMaybe (\(Exists a) -> instrToSync a) $ arrayInstrsInExp expr + , PartialSchedule $ PCompute expr ) + C.Alet lhs usBnd bnd body + | (bndFree, bnd') <- toPartial' usBnd bnd + , (bodyFree, body') <- toPartial' us body -> + ( bndFree `IdxSet.union` IdxSet.drop' lhs bodyFree + , PartialSchedule $ PLet Parallel lhs usBnd bnd' body' ) + C.Alloc shr tp sh -> simple $ PAlloc shr tp sh + C.Use tp sz buf -> simple $ PUse tp sz buf + C.Unit var -> simple $ PUnit var + C.Acond var true false + | (trueFree, true') <- toPartial' us true + , (falseFree, false') <- toPartial' us false -> + ( trueFree `IdxSet.union` falseFree + , PartialSchedule $ PAcond var true' false' ) + C.Awhile us' (C.Alam lhsC (C.Abody cond)) (C.Alam lhsS (C.Abody step)) initial + | DeclareVars lhs _ vars <- declareVars $ lhsToTupR lhsC + , (condFree, cond') <- toPartial' (TupRsingle Shared) $ weaken (weakenToFullLHS lhs lhsC) cond + , (stepFree, step') <- toPartial' (TupRsingle Shared) $ weaken (weakenSucc' $ weakenToFullLHS lhs lhsS) step + , fn <- PartialSchedule $ PLet Sequential (LeftHandSideSingle (GroundRscalar scalarType)) (TupRsingle Shared) cond' + $ PartialSchedule $ PAcond (Var scalarType ZeroIdx) + (PartialSchedule $ PContinue step') + (PartialSchedule $ PBreak us' $ vars (weakenSucc weakenId)) + -> + ( IdxSet.drop' lhs condFree `IdxSet.union` IdxSet.drop' lhs (IdxSet.drop stepFree) `IdxSet.union` IdxSet.fromList (groundBufferVars initial) + , PartialSchedule $ PAwhile us' (Plam lhs $ Pbody fn) initial ) + C.Awhile{} -> internalError "Unary function impossible" + C.Atrace msg t -> (arrayDescriptorsIdxSet t, PartialSchedule $ PAtrace msg t) + C.Aassert msg cond -> + ( IdxSet.fromList $ mapMaybe (\(Exists a) -> instrToSync a) $ arrayInstrsInExp cond + , PartialSchedule $ PAssert msg cond ) + C.Aassume _ -> toPartial' us $ C.Compute $ Const scalarTypeWord8 0 -- Assumptions are removed at this point + C.Fence deps next + | (nextFree, next') <- toPartial' us next -> + ( deps `IdxSet.union` nextFree + , PartialSchedule $ PFence deps next' ) + where + -- For all simple cases, with no free buffer variables. + simple :: PrePartialSchedule PartialSchedule kernel env t -> (IdxSet env, PartialSchedule kernel env t) + simple schedule = (IdxSet.empty, PartialSchedule schedule) + + isBuffer :: Exists (GroundVar env) -> Bool + isBuffer (Exists (Var (GroundRbuffer _) _)) = True + isBuffer _ = False + + instrToSync :: ArrayInstr env a -> Maybe (Exists (Idx env)) + instrToSync (Index v) = Just $ Exists $ varIdx v + instrToSync (Parameter _) = Nothing -- Parameter is a scalar variable + +weakenToFullLHS + -- Full LHS + :: GLeftHandSide t env envFull + -- Sub LHS + -> GLeftHandSide t env envSub + -> envSub :> envFull +weakenToFullLHS = \lhs1 lhs2 -> go lhs1 lhs2 weakenId + where + go :: GLeftHandSide t env1 env1' -> GLeftHandSide t env2 env2' -> env2 :> env1 -> env2' :> env1' + go lhs (LeftHandSideWildcard _) k = weakenWithLHS lhs .> k + go (LeftHandSidePair lhs1 lhs1') (LeftHandSidePair lhs2 lhs2') k = go lhs1' lhs2' $ go lhs1 lhs2 k + go (LeftHandSideSingle _) (LeftHandSideSingle _) k = sink k + go (LeftHandSideWildcard _) _ _ = internalError "Expected second LHS to be contained in first LHS" + go _ _ _ = internalError "LHS mismatch" + +returnValues + :: UpdateTuple env t1 t2 + -> PartialSchedule kernel env t1 + -> PartialSchedule kernel env t2 +returnValues UpdateKeep next = next +returnValues updateTup next = PartialSchedule $ PReturnValues updateTup next + +toPartialReturn :: Uniquenesses t -> GroundVars env t -> UpdateTuple env (UnitTuple t) t +toPartialReturn _ TupRunit = UpdateKeep +toPartialReturn us (TupRpair v1 v2) = toPartialReturn us1 v1 `updatePair` toPartialReturn us2 v2 + where + (us1, us2) = pairUniqueness us +toPartialReturn (TupRsingle u) (TupRsingle var) + | Refl <- groundNotPair $ varType var + = UpdateSet u var +toPartialReturn _ _ = internalError "Tuple mismatch" + +groundBufferVars :: GroundVars env t -> [Exists (Idx env)] +groundBufferVars = (`go` []) + where + go :: GroundVars env t -> [Exists (Idx env)] -> [Exists (Idx env)] + go (TupRpair v1 v2) accum = go v1 $ go v2 accum + go (TupRsingle (Var (GroundRbuffer _) idx)) accum = Exists idx : accum + go _ accum = accum + +-- Optimizes the PartialSchedule by rebuilding it. It keeps rebuilding until +-- it reaches a fixed-point, or after it performs n iterations. +-- Assumes n >= 1 +rebuild :: Int -> PartialSchedule kernel env t -> PartialSchedule kernel env t +rebuild n t + | didChange built && n > 1 = rebuild (n - 1) $ term built + | otherwise = term built + where + built = rebuild' t IdxSet.empty + +rebuild' :: PartialSchedule kernel env t -> Build PartialSchedule kernel env t +rebuild' (PartialSchedule schedule) = case schedule of + PExec kernel args -> buildExec kernel args + PLet parallelism lhs us bnd body -> buildLet parallelism lhs us (rebuild' bnd) (rebuild' body) + PAlloc shr tp sh -> buildAlloc shr tp sh + PUse sz tp buf -> buildUse sz tp buf + PUnit var -> buildUnit var + PCompute expr -> buildCompute expr + PReturnEnd tup -> buildReturnEnd tup + PReturnValues updateTup next -> buildReturnValues updateTup (rebuild' next) + PAcond var true false -> buildAcond var (rebuild' true) (rebuild' false) + PAwhile us fn initial -> buildAwhile us (rebuildUnary fn) initial + PAtrace msg t -> buildTrace msg t + PContinue next -> buildContinue $ rebuild' next + PBreak us vars -> buildBreak us vars + PAssert msg cond -> buildAssert msg cond + PFence deps next -> buildFence deps $ rebuild' next + +rebuildUnary :: PartialScheduleFun kernel env f -> BuildUnary kernel env f +rebuildUnary (Plam lhs (Pbody body)) = BuildUnary lhs $ rebuild' body +rebuildUnary _ = internalError "Expected unary function" + +buildExec :: KernelFun kernel s -> SArgs env s -> Build PartialSchedule kernel env () +buildExec kernel args available = + Built { + didChange = False, + directlyAwaits = IdxSet.fromList (sargVars args) IdxSet.\\ available, + writes = output, + finallyReleases = output, + trivial = False, + term = PartialSchedule $ PExec kernel args + } + where + output = IdxSet.fromList (sargOutputVars args) + +buildLet + :: Parallelism + -> GLeftHandSide bnd env env' + -> Uniquenesses bnd + -> Build PartialSchedule kernel env bnd + -> Build PartialSchedule kernel env' t + -> Build PartialSchedule kernel env t +buildLet parallelismHint lhs us bnd' body' available = + Built { + didChange = didChange bnd || didChange body || parallelismHint /= parallelism, + directlyAwaits = thisAwaits, + writes = writes bnd `IdxSet.union` IdxSet.drop' lhs (writes body), + finallyReleases = thisReleases, + trivial = sequential && trivial bnd && trivial body, + term = + PartialSchedule $ PLet + parallelism + lhs us (term bnd) (term body) + } + where + bnd = bnd' available + bodyAvailable + | sequentialEarlyTest = IdxSet.push' lhs $ available `IdxSet.union` directlyAwaits bnd + | otherwise = IdxSet.skip' lhs (available IdxSet.\\ writes bnd) + body = body' bodyAvailable + + bodyAwaitsDropped = IdxSet.drop' lhs (directlyAwaits body) + bodyReleasesDropped = IdxSet.drop' lhs (finallyReleases body) + + thisAwaits + | sequential && trivial bnd = IdxSet.union (directlyAwaits bnd) bodyAwaitsDropped + | sequential = directlyAwaits bnd + | otherwise = IdxSet.intersect (directlyAwaits bnd) bodyAwaitsDropped + + thisReleases + | sequential && trivial body = IdxSet.union (finallyReleases bnd) bodyReleasesDropped + | sequential = bodyReleasesDropped + | otherwise = IdxSet.empty + + -- Checks if this let-binding should become sequential. + -- This check happens in two steps: + -- An early check runs before calling body'. Based on this check, + -- we pass an availability-set to body'. + -- A late check runs after calling body'. This uses the information from + -- body. + sequentialEarlyTest = + -- If the hint was to make this Let sequential, make it sequential + parallelismHint == Sequential + -- If the binding is trivial and doesn't wait on variables. + || trivial bnd && IdxSet.null (directlyAwaits bnd) + + thisLhsIndices = lhsIndices lhs + + -- Make it sequential if: + sequential = + -- the early test already made it sequential, + sequentialEarlyTest + -- the body directly needs all declared variables (and there is at least one such variable), + || not (IdxSet.isEmpty thisLhsIndices) && thisLhsIndices `IdxSet.isSubsetOf` directlyAwaits body + -- the binding is trivial and doesn't use other variables than the + -- body already waits on. + -- Note that the awaits set contains all free variables (minus the available variables) + -- if the program is trivial. + || trivial bnd && directlyAwaits bnd `IdxSet.isSubsetOf` bodyAwaitsDropped + -- or the binding is trivial and the body directly awaits on the result of this expression. + || trivial bnd && thisLhsIndices `IdxSet.overlaps` directlyAwaits body + + parallelism = if sequential then Sequential else Parallel + +buildAlloc :: ShapeR sh -> ScalarType e -> ExpVars env sh -> Build PartialSchedule kernel env (Buffer e) +buildAlloc shr tp vars available = + Built{ + didChange = False, + directlyAwaits = IdxSet.fromVars vars IdxSet.\\ available, + writes = IdxSet.empty, + finallyReleases = IdxSet.empty, + trivial = True, + term = PartialSchedule $ PAlloc shr tp vars + } + +buildUse :: ScalarType e -> Int -> Buffer e -> Build PartialSchedule kernel env (Buffer e) +buildUse tp sz buf _ = + Built{ + didChange = False, + directlyAwaits = IdxSet.empty, + writes = IdxSet.empty, + finallyReleases = IdxSet.empty, + trivial = True, + term = PartialSchedule $ PUse tp sz buf + } + +buildUnit :: ExpVar env e -> Build PartialSchedule kernel env (Buffer e) +buildUnit var available = + Built{ + didChange = False, + directlyAwaits = if var `IdxSet.varMember` available then IdxSet.empty else IdxSet.singletonVar var, + writes = IdxSet.empty, + finallyReleases = IdxSet.empty, + trivial = True, + term = PartialSchedule $ PUnit var + } + +buildCompute :: Exp env e -> Build PartialSchedule kernel env e +buildCompute expr available = + Built{ + didChange = False, + directlyAwaits = IdxSet.fromVarList vars IdxSet.\\ available, + writes = IdxSet.empty, + finallyReleases = IdxSet.empty, + trivial = expIsTrivial (const True) expr, + term = PartialSchedule $ PCompute expr + } + where + vars = expGroundVars expr + +buildReturnEnd :: TupR Void' t -> Build PartialSchedule kernel env t +buildReturnEnd tup _ = + Built{ + didChange = False, + directlyAwaits = IdxSet.empty, + writes = IdxSet.empty, + finallyReleases = IdxSet.empty, + trivial = True, + term = PartialSchedule $ PReturnEnd tup + } + +buildReturnValues :: forall kernel env t1 t2. UpdateTuple env t1 t2 -> Build PartialSchedule kernel env t1 -> Build PartialSchedule kernel env t2 +buildReturnValues updateTup next' available = + Built{ + didChange = didChange next, + -- A Return can be executed in parallel: if some returned variables are/become available, + -- they can directly be returned. Others can then later be returned when they become + -- available. Hence the directly-awaits set is in general empty. + -- One exception is if the return has only one variable + directlyAwaits = thisAwaits, + writes = writes next, + finallyReleases = if allAvailable then finallyReleases next else IdxSet.empty, + -- A parallel return is non-trivial. If the return has only one variable, or if + -- all variables are definitely available here, the return is sequential, + -- and we can thus treat is as trivial. + -- If the next instruction is trivial, and the return nor the next do not await on any values, + -- we can also treat it as trivial. + trivial = (nextIsEmpty && hasOne) || (allAvailable && trivial next && IdxSet.null (directlyAwaits next)), + term = PartialSchedule $ PReturnValues updateTup $ term next + } + where + next = next' available + nextIsEmpty + | PartialSchedule (PReturnEnd _) <- term next = True + | otherwise = False + + varsSet = IdxSet.fromList $ updateTupVarList updateTup + + hasOne + | One <- updateCount updateTup = True + | otherwise = False + + allAvailable = varsSet `IdxSet.isSubsetOf` available + + thisAwaits + | hasOne && nextIsEmpty = varsSet IdxSet.\\ available + | otherwise = IdxSet.empty + +buildAcond + :: ExpVar env PrimBool + -> Build PartialSchedule kernel env t + -> Build PartialSchedule kernel env t + -> Build PartialSchedule kernel env t +buildAcond var true' false' available = + Built{ + didChange = didChange true || didChange false, + directlyAwaits = if var `IdxSet.varMember` available then awaits else IdxSet.insertVar var awaits, + writes = writes true `IdxSet.union` writes false, + finallyReleases = finallyReleases true `IdxSet.intersect` finallyReleases false, + -- true and false branch may have different sets of free variables, + -- hence we cannot treat an it as a trivial term. + trivial = False, + term = PartialSchedule $ PAcond var (term true) (term false) + } + where + available' = IdxSet.insertVar var available + true = true' available' + false = false' available' + awaits = directlyAwaits true `IdxSet.intersect` directlyAwaits false + +data BuildUnary kernel env f where + BuildUnary + :: GLeftHandSide s env env' + -> Build PartialSchedule kernel env' t + -> BuildUnary kernel env (s -> t) + +buildAwhile + :: Uniquenesses t + -> BuildUnary kernel env (t -> Loop t) + -> GroundVars env t + -> Build PartialSchedule kernel env t +buildAwhile us (BuildUnary lhs fn') initial available = + Built{ + didChange = didChange fn, + directlyAwaits = IdxSet.fromVarList (lhsTake lhs (directlyAwaits fn) initial) `IdxSet.union` IdxSet.drop' lhs (directlyAwaits fn), + writes = IdxSet.drop' lhs (writes fn), + finallyReleases = IdxSet.empty, + trivial = False, + term = PartialSchedule $ PAwhile + us + (Plam lhs $ Pbody $ term fn) + initial + } + where + fn = fn' (IdxSet.skip' lhs available) + +buildTrace :: + Text -> + ArrayDescriptors env t -> + Build PartialSchedule kernel env Word8 +buildTrace msg t available = + Built{ + didChange = False, + directlyAwaits = arrayDescriptorsIdxSet t IdxSet.\\ available, + writes = IdxSet.empty, + finallyReleases = IdxSet.empty, + trivial = True, + term = PartialSchedule $ PAtrace msg t + } + +buildContinue + :: Build PartialSchedule kernel env t + -> Build PartialSchedule kernel env (Loop t) +buildContinue next' available = + Built{ + didChange = didChange next, + directlyAwaits = IdxSet.empty, + writes = writes next, + finallyReleases = finallyReleases next, + trivial = False, + term = PartialSchedule $ PContinue $ term next + } + where + next = next' available + +buildBreak + :: Uniquenesses t + -> GroundVars env t + -> Build PartialSchedule kernel env (Loop t) +buildBreak us vars _ = + Built{ + didChange = False, + directlyAwaits = IdxSet.fromVars vars, + writes = IdxSet.empty, + finallyReleases = IdxSet.empty, + trivial = False, + term = PartialSchedule $ PBreak us vars + } + +buildAssert + :: Text + -> Exp env PrimBool + -> Build PartialSchedule kernel env PrimBool +buildAssert msg cond available = + Built{ + didChange = False, + directlyAwaits = IdxSet.fromVarList vars IdxSet.\\ available, + writes = IdxSet.empty, + finallyReleases = IdxSet.empty, + trivial = expIsTrivial (const True) cond, + term = PartialSchedule $ PAssert msg cond + } + where + vars = expGroundVars cond + +buildFence + :: IdxSet env + -> Build PartialSchedule kernel env t + -> Build PartialSchedule kernel env t +buildFence deps next' available = + Built{ + didChange = didChange next, + directlyAwaits = deps' `IdxSet.union` directlyAwaits next, + writes = writes next, + finallyReleases = finallyReleases next, + trivial = trivial next, + term = PartialSchedule $ PFence deps' $ term next + } + where + next = next' (IdxSet.union available deps) + deps' = deps IdxSet.\\ available + +updateCount :: UpdateTuple env t1 t2 -> Count +updateCount UpdateKeep = Zero +updateCount (UpdateSet _ _) = One +updateCount (UpdatePair up1 up2) = updateCount up1 <> updateCount up2 + +data Count = Zero | One | Many +instance Semigroup Count where + Zero <> a = a + One <> Zero = One + _ <> _ = Many + +data ToSyncSchedule kernel env t where + ToSyncSchedule + :: UpdateTuple env t1 t2 + -> SyncSchedule kernel env t1 + -> ToSyncSchedule kernel env t2 + +arrayDescriptorsSync :: ArrayDescriptors env t -> SyncEnv env +arrayDescriptorsSync TupRunit = PEnd +arrayDescriptorsSync (TupRsingle a) = arrayDescriptorSync a +arrayDescriptorsSync (TupRpair l r) = unionPartialEnv max (arrayDescriptorsSync l) (arrayDescriptorsSync r) + +arrayDescriptorSync :: ArrayDescriptor env t -> SyncEnv env +arrayDescriptorSync (ArrayDescriptor _ _ buffers) = variablesToSyncEnv (mapTupR (const Shared) buffers) buffers + +analyseSyncEnv :: PartialSchedule kernel env t -> SyncSchedule kernel env t +analyseSyncEnv sched + | ToSyncSchedule up sched' <- analyseSyncEnv' sched + = sReturnValues up sched' + +sReturnValues + :: UpdateTuple env t1 t2 + -> SyncSchedule kernel env t1 + -> SyncSchedule kernel env t2 +sReturnValues UpdateKeep sched = sched +sReturnValues up sched = SyncSchedule + (unionPartialEnv max (syncEnv sched) $ partialEnvFromList noCombine $ updateTupSync up) + True + (PReturnValues up sched) + +analyseSyncEnv' :: PartialSchedule kernel env t -> ToSyncSchedule kernel env t +analyseSyncEnv' (PartialSchedule sched) = case sched of + PExec kernel args -> + let + bindings = mapMaybe (\(Exists a) -> argToSync a) $ argsToList args + argToSync :: SArg env a -> Maybe (EnvBinding Sync env) + argToSync (SArgBuffer In v) = Just $ EnvBinding (varIdx v) SyncRead + argToSync (SArgBuffer _ v) = Just $ EnvBinding (varIdx v) SyncWrite + argToSync _ = Nothing + in + ToSyncSchedule UpdateKeep + $ SyncSchedule (partialEnvFromList max bindings) True (PExec kernel args) + PLet par lhs us bnd body + | bnd' <- analyseSyncEnv bnd + , ToSyncSchedule up body' <- analyseSyncEnv' body + , HoistReturn up1 up2 <- hoistReturn lhs (syncEnv bnd') up -> + ToSyncSchedule up1 $ + SyncSchedule + (unionPartialEnv max (syncEnv bnd') $ weakenSyncEnv lhs $ syncEnv body') + (syncSimple bnd' && syncSimple body' && par == Sequential) + (PLet par lhs us bnd' $ sReturnValues up2 body') + PAlloc shr tp sh -> + noBuffers $ PAlloc shr tp sh + PUse tp sz buf -> + noBuffers $ PUse tp sz buf + PUnit var -> + noBuffers $ PUnit var + PCompute expr -> + let + bindings = mapMaybe (\(Exists a) -> instrToSync a) $ arrayInstrsInExp expr + in + ToSyncSchedule UpdateKeep + $ SyncSchedule + (partialEnvFromList const bindings) + True + (PCompute expr) + PReturnEnd tup -> + noBuffers $ PReturnEnd tup + PReturnValues updateTup next -> ToSyncSchedule updateTup $ analyseSyncEnv next + PAcond var true false -> + let + true' = analyseSyncEnv true + false' = analyseSyncEnv false + in + ToSyncSchedule UpdateKeep $ + SyncSchedule + (unionPartialEnv max (syncEnv true') (syncEnv false')) + False + (PAcond var true' false') + PAwhile us (Plam lhs (Pbody body)) initial -> + let + body' = analyseSyncEnv body + in + ToSyncSchedule UpdateKeep $ + SyncSchedule + (unionPartialEnv max (weakenSyncEnv lhs $ syncEnv body') $ variablesToSyncEnv us initial) + False + (PAwhile us (Plam lhs $ Pbody body') initial) + PAwhile{} -> internalError "Function impossible" + PAtrace msg t -> + ToSyncSchedule UpdateKeep $ + SyncSchedule + (arrayDescriptorsSync t) + True + (PAtrace msg t) + PContinue next -> + let + next' = analyseSyncEnv next + in + ToSyncSchedule UpdateKeep $ + SyncSchedule (syncEnv next') False $ PContinue next' + PBreak us vars -> + ToSyncSchedule UpdateKeep $ SyncSchedule (variablesToSyncEnv us vars) False $ PBreak us vars + PAssert msg cond -> + let + bindings = mapMaybe (\(Exists a) -> instrToSync a) $ arrayInstrsInExp cond + in + ToSyncSchedule UpdateKeep + $ SyncSchedule + (partialEnvFromList const bindings) + True + (PAssert msg cond) + PFence deps next -> + let + next' = analyseSyncEnv next + in + ToSyncSchedule UpdateKeep + $ SyncSchedule + (syncEnv next') + (syncSimple next') + (PFence deps next') + where + noBuffers :: PrePartialSchedule SyncSchedule kernel env t -> ToSyncSchedule kernel env t + noBuffers = ToSyncSchedule UpdateKeep . SyncSchedule PEnd True + + instrToSync :: ArrayInstr env a -> Maybe (EnvBinding Sync env) + instrToSync (Index v) = Just $ EnvBinding (varIdx v) SyncRead + instrToSync (Parameter _) = Nothing -- Parameter is a scalar variable + +variablesToSyncEnv :: Uniquenesses t -> GroundVars genv t -> SyncEnv genv +variablesToSyncEnv uniquenesses vars = partialEnvFromList noCombine $ go uniquenesses vars [] + where + go :: Uniquenesses t -> GroundVars genv t -> [EnvBinding Sync genv] -> [EnvBinding Sync genv] + go (TupRsingle Unique) (TupRsingle (Var (GroundRbuffer _) ix)) + accum = EnvBinding ix SyncWrite : accum + go (TupRsingle Shared) (TupRsingle (Var (GroundRbuffer _) ix)) + accum = EnvBinding ix SyncRead : accum + go u (TupRpair v1 v2) accum = go u1 v1 $ go u2 v2 accum + where (u1, u2) = pairUniqueness u + go _ _ accum = accum + +noCombine :: Sync s -> Sync s -> Sync s +noCombine SyncRead SyncRead = SyncRead +noCombine _ _ = internalError "A writable buffer cannot be aliassed" diff --git a/src/Data/Array/Accelerate/Trafo/Schedule/Uniform.hs b/src/Data/Array/Accelerate/Trafo/Schedule/Uniform.hs new file mode 100644 index 000000000..b67fca102 --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Schedule/Uniform.hs @@ -0,0 +1,1309 @@ +{-# LANGUAGE ImpredicativeTypes #-} +{-# LANGUAGE MultiWayIf #-} +{-# LANGUAGE TypeFamilies #-} +{-# OPTIONS_GHC -Wno-orphans #-} +{-# OPTIONS_HADDOCK hide #-} + +-- | +-- Module : Data.Array.Accelerate.Trafo.Schedule.Uniform +-- Copyright : [2012..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.Trafo.Schedule.Uniform ( + -- Exports the instance IsSchedule UniformScheduleFun + + compileKernel', CompiledKernel(..), rnfSArg, rnfSArgs, +) where + +import Data.Array.Accelerate.Analysis.Match +import Data.Array.Accelerate.AST.Environment +import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.IdxSet (IdxSet(..)) +import qualified Data.Array.Accelerate.AST.IdxSet as IdxSet +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.AST.Kernel +import Data.Array.Accelerate.AST.Schedule +import Data.Array.Accelerate.AST.Schedule.Uniform +import Data.Array.Accelerate.AST.Var +import Data.Array.Accelerate.Error +import Data.Array.Accelerate.Representation.Ground +import Data.Array.Accelerate.Representation.Shape +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Trafo.Exp.Substitution +import Data.Array.Accelerate.Trafo.Substitution +import Data.Array.Accelerate.Trafo.Var +import Data.Array.Accelerate.Trafo.Schedule.Partial +import Data.Array.Accelerate.Trafo.Schedule.Uniform.Future +import Data.Array.Accelerate.Trafo.Schedule.Uniform.Simplify +import Data.Array.Accelerate.Trafo.Schedule.Uniform.LiveVars +import Data.Array.Accelerate.Trafo.SkipEnvironment +import Data.Array.Accelerate.Type +import Data.Array.Accelerate.AST.Partitioned (PartitionedAfun) +import qualified Data.Array.Accelerate.AST.Partitioned as C +import Data.Maybe +import Prelude hiding (id, (.), read) +import Control.Category +import Control.DeepSeq (NFData(rnf)) +import Control.Concurrent +import Data.IORef +import System.IO.Unsafe (unsafePerformIO) + +instance IsSchedule UniformScheduleFun where + type ScheduleInput UniformScheduleFun a = Input a + type ScheduleOutput UniformScheduleFun a = Output a + + convertScheduleFun = transformAfun + + rnfSchedule (Slam lhs s) = rnfLeftHandSide rnfBaseR lhs `seq` rnfSchedule s + rnfSchedule (Sbody s) = rnfSchedule' s + + callScheduledFun (GFunctionRbody repr) f + | Refl <- reprIsBody @UniformScheduleFun repr = do + (output, result) <- bindOutput repr + f output + return result + where + bindOutput :: GroundsR t -> IO (Output t, t) + bindOutput TupRunit = return ((), ()) + bindOutput (TupRpair tp1 tp2) = do + (out1, result1) <- bindOutput tp1 + (out2, result2) <- bindOutput tp2 + return ((out1, out2), (result1, result2)) + bindOutput (TupRsingle tp) + | Refl <- outputSingle tp = do + ref <- newIORef $ internalError "Illegal schedule: Read from ref without value. Some synchronization might be missing." + mvar <- newEmptyMVar + -- We return the result lazily, as the value is not yet available + let value = unsafePerformIO $ do + readMVar mvar + readIORef ref + return ((SignalResolver mvar, OutputRef ref), value) + callScheduledFun (GFunctionRlam arg ret) f = do + return $ \a -> do + input <- bindInput arg a + callScheduledFun @UniformScheduleFun ret $ f input + where + bindInput :: GroundsR t -> t -> IO (Input t) + bindInput TupRunit _ = return () + bindInput (TupRpair tp1 tp2) (val1, val2) = do + in1 <- bindInput tp1 val1 + in2 <- bindInput tp2 val2 + return (in1, in2) + bindInput (TupRsingle tp) value + | Refl <- inputSingle tp = do + ref <- newIORef $ internalError "Illegal schedule: Read from ref without value. Some synchronization might be missing." + mvar <- newEmptyMVar + -- We consume the result lazily. The computation can already start + -- asynchronously, and some parts/threads of the computation may at + -- some point block on the input to be available. + _ <- forkIO $ do + -- In a separate thread, evaluate to WHNF + value `seq` return () + -- and write value to reference and resolve signal + writeIORef ref value + putMVar mvar () + return (Signal mvar, Ref ref) + +transformAfun + :: IsKernel kernel + => PartitionedAfun (KernelOperation kernel) () f + -> UniformScheduleFun kernel () (Scheduled UniformScheduleFun f) +transformAfun afun = simplify $ stronglyLiveVariablesFun $ simplify $ stronglyLiveVariablesFun $ simplify $ stronglyLiveVariablesFun $ makeFunSpawnClosed $ funConstruct (go FEnvEnd afun) weakenId BEmpty + where + go + :: IsKernel kernel + => FutureEnv fenv genv + -> PartitionedAfun (KernelOperation kernel) genv f + -> BuildScheduleFun kernel fenv (Scheduled UniformScheduleFun f) + go env (C.Alam lhs f) + | DeclareInput _ lhs' env' <- declareInput lhs env + = buildFunLam lhs' $ go env' f + go env (C.Abody body) + | Refl <- reprIsBody @UniformScheduleFun $ groundsR body + , schedule <- convertAccPartial body + , DeclareOutput skip lhs dest <- declareOutput $ groundsR body + = buildFunLam lhs $ buildFunBody $ + transformSub (fenvFSkipMany skip env) Parallel (CtxNormal $ dest SkipNone) schedule + +-- Transforms a SyncSchedule to a UniformSchedule. +-- Assumes that the FutureEnv corresponds with the syncEnv of the schedule. +transform + :: IsKernel kernel + => FutureEnv fenv genv -- Mapping from ground environment to future environment + -> Parallelism -- Whether the returns should happen in parallel or sequential. + -> Context kernel fenv t -- The destinations for the output + -> SyncSchedule kernel genv t -- The SyncSchedule that should be transformed to a UniformSchedule + -> BuildSchedule kernel fenv +transform env parallelism ctx schedule = case transform' schedule of + TransformSchedule f -> f env parallelism ctx + TransformBinding f + | TransformToBinding tp skip instr resolvers binding <- f env + , DeclareVars lhs k vars <- declareVars tp -> + case ctx of + CtxLoop{} -> internalError "Loop impossible" + CtxNormal dest -> + instr + $ buildLet (mapLeftHandSide BaseRground lhs) binding + $ buildEffect (SignalResolve $ weaken (weakenWithLHS lhs) <$> resolvers) + $ writeOutput (mapTupR (weaken k . weaken (skipWeakenIdx skip)) dest) (vars weakenId) + +-- Variant of 'transform' that also performs the subenvironment rule: +-- transform assumes that the FutureEnv corresponds with the syncEnv of the schedule, +-- but transformSub correctly handles redundant entries in the FutureEnv. +transformSub + :: IsKernel kernel + => FutureEnv fenv genv -- Mapping from ground environment to future environment + -> Parallelism + -> Context kernel fenv t + -> SyncSchedule kernel genv t -- The SyncSchedule that should be transformed to a UniformSchedule + -> BuildSchedule kernel fenv +transformSub env parallelism dest schedule = buildSpawn actions + $ transform env' parallelism dest schedule + where + (env', actions) = subFutureEnvironment env (syncEnv schedule) + +data Transform kernel genv t where + TransformSchedule + :: (forall fenv. FutureEnv fenv genv -> Parallelism -> Context kernel fenv t -> BuildSchedule kernel fenv) + -> Transform kernel genv t + + TransformBinding + :: (forall fenv. FutureEnv fenv genv -> TransformToBinding kernel fenv t) + -> Transform kernel genv t + +-- Captures existential fenv' +data TransformToBinding kernel fenv t where + TransformToBinding + :: GroundsR t + -> Skip fenv' fenv + -> (BuildSchedule kernel fenv' -> BuildSchedule kernel fenv) + -> [Idx fenv' SignalResolver] + -> Binding fenv' t + -> TransformToBinding kernel fenv t + +-- Invariants: +-- - If 'simple' then Parallelism (argument of TransformSchedule) is Sequential. +-- - If 'simple' then we may not introduce forks for this term. +-- - Parallelism is not SequentialInNextParallel +transform' + :: IsKernel kernel + => SyncSchedule kernel genv t -- The SyncSchedule that should be transformed to a UniformSchedule + -> Transform kernel genv t +transform' (SyncSchedule _ simple schedule) = case schedule of + -- Return & return values + PReturnEnd _ -> TransformSchedule $ \_ _ _ -> buildReturn + PReturnValues updateTup next -> TransformSchedule $ \fenv parallelism ctx -> + case ctx of + CtxLoop{} -> internalError "Loop impossible" + CtxNormal dest + | (dest', instr) <- returnValues fenv parallelism updateTup dest + , next' <- transformSub fenv parallelism (CtxNormal dest') next -> + instr next' + + -- Effects + PExec kernel args -> TransformSchedule $ \fenv _ _ -> + case acquireArgs fenv args of + AcquireArgs skip signals resolvers instr args' -> + buildEffect (SignalAwait signals) + $ instr + $ buildEffect (Exec (kernelMetadata kernel) kernel $ args' SkipNone) + $ buildEffect (SignalResolve $ map (weaken $ skipWeakenIdx skip) resolvers) buildReturn + PAtrace msg t -> TransformSchedule $ \fenv _ ctx -> + case acquireSome fenv $ arrayDescriptorsIdxSet t of + AcquireSome skip signals resolvers instr mapping' + | mapping <- mapping' SkipNone -> case ctx of + CtxNormal dest + | TupRsingle dest' <- dest -> + case reindexArrayDescriptors (`prjPartial` mapping) t of + Nothing -> internalError "Variable missing after acquireSome" + Just t' -> + buildEffect (SignalAwait signals) + $ instr + $ buildEffect (Atrace msg t') + $ buildEffect (SignalResolve $ map (weaken $ skipWeakenIdx skip) resolvers) + $ writeLoopCondition (weaken (skipWeakenIdx skip) dest') True + PAssert msg cond -> TransformSchedule $ \fenv _ ctx -> + case acquireSome fenv $ IdxSet.fromVarList $ expGroundVars cond of + AcquireSome skip signals resolvers instr mapping' + | mapping <- mapping' SkipNone -> case ctx of + CtxNormal dest + | TupRsingle dest' <- dest -> + buildEffect (SignalAwait signals) + $ instr + $ buildEffect (Aassert msg $ mapArrayInstr + (weaken (Weaken $ \idx -> + fromMaybe (internalError "Variable missing after acquireSome") $ prjPartial idx mapping)) + cond) + $ buildEffect (SignalResolve $ map (weaken $ skipWeakenIdx skip) resolvers) + $ writeLoopCondition (weaken (skipWeakenIdx skip) dest') True + PFence _ next -> transform' next + + -- Bindings + PAlloc shr tp sh -> TransformBinding $ \fenv -> + case acquireVars fenv sh of + AcquireVars skip instr sh' -> + TransformToBinding (TupRsingle $ GroundRbuffer tp) skip instr [] $ Alloc shr tp $ sh' SkipNone + PUse tp sz buf -> TransformBinding $ \_ -> + TransformToBinding (TupRsingle $ GroundRbuffer tp) SkipNone id [] $ Use tp sz buf + PUnit var -> TransformBinding $ \fenv -> + case fprj (varIdx var) fenv of + Just (FutureScalar tp signal (Left idx)) -> + TransformToBinding (TupRsingle $ GroundRbuffer tp) SkipNone (buildAwait [signal]) [] $ Unit $ Var tp idx + Just (FutureScalar tp signal (Right idx)) -> + TransformToBinding + (TupRsingle $ GroundRbuffer tp) + (SkipSucc SkipNone) + (buildAwait [signal] + . buildLet (LeftHandSideSingle $ BaseRground $ GroundRscalar tp) (RefRead $ Var (BaseRref $ GroundRscalar tp) idx)) + [] + (Unit $ Var tp ZeroIdx) + _ -> variableMissing + PCompute expr -> TransformBinding $ \fenv -> + case acquireSome fenv $ IdxSet.fromVarList $ expGroundVars expr of + AcquireSome skip signals resolvers instr mapping' + | mapping <- mapping' SkipNone -> + TransformToBinding + (mapTupR GroundRscalar $ expType expr) + skip + (buildEffect (SignalAwait signals) . instr) + (weaken (skipWeakenIdx skip) <$> resolvers) + (Compute $ mapArrayInstr + (weaken (Weaken $ \idx -> + fromMaybe (internalError "Variable missing after acquireSome") $ prjPartial idx mapping)) + expr) + + PLet bndParallelism (LeftHandSideWildcard tp) _ bnd body -> TransformSchedule $ \fenv parallelism ctx -> + case chainFutureEnvironmentSeqIf (bndParallelism == Sequential && simple) fenv (syncEnv bnd) (syncEnv body) of + ChainFutureEnv instr skip bndEnv bodyEnv -> + let + bnd' = transform bndEnv bndParallelism (CtxNormal $ mapTupR (const DestinationIgnore) tp) bnd + body' = transform bodyEnv parallelism (weaken (skipWeakenIdx skip) ctx) body + in instr $ buildSeqOrSpawn bndParallelism bnd' body' + + PLet Parallel lhs us bnd body -> TransformSchedule $ \fenv parallelism ctx -> + case chainFutureEnvironment SkipNone fenv (syncEnv bnd) (environmentDropLHS (syncEnv body) lhs) of + ChainFutureEnv instr1 skip1 bndEnv bodyEnv + | DeclareDestinations skip2 instr2 bndDestinations bodyEnv' <- declareDestinations Parallel bodyEnv lhs us + , bodyEnv'' <- restrictEnvForLHS lhs bodyEnv' $ syncEnv body -> + let + k = skipWeakenIdx skip2 .> skipWeakenIdx skip1 + bnd' = transform (fenvFSkipMany skip2 bndEnv) Parallel (CtxNormal $ bndDestinations SkipNone) bnd + body' = transform bodyEnv'' parallelism (weaken k ctx) body + in instr1 $ instr2 $ buildSpawn bnd' body' + + PLet Sequential lhs us bnd body -> TransformSchedule $ \fenv parallelism ctx -> + case chainFutureEnvironmentSeqIf simple fenv (syncEnv bnd) (environmentDropLHS (syncEnv body) lhs) of + ChainFutureEnv instr1 skip1 bndEnv bodyEnv -> + case transform' bnd of + TransformBinding f + | TransformToBinding _ skip2 instr2 resolvers binding <- f bndEnv + , Exists lhs' <- rebuildLHS lhs + , bodyEnv' <- fenvPushLHS lhs lhs' us $ fenvFSkipMany skip2 bodyEnv + , bodyEnv'' <- restrictEnvForLHS lhs bodyEnv' $ syncEnv body -> + instr1 $ instr2 + $ buildLet (mapLeftHandSide BaseRground lhs') binding + $ buildEffect (SignalResolve $ weaken (weakenWithLHS lhs') <$> resolvers) + $ transform bodyEnv'' parallelism (weaken (weakenWithLHS lhs' .> skipWeakenIdx skip2 .> skipWeakenIdx skip1) ctx) body + TransformSchedule f + | DeclareDestinations skip2 instr2 bndDestinations bodyEnv' <- declareDestinations Parallel bodyEnv lhs us + , bodyEnv'' <- restrictEnvForLHS lhs bodyEnv' $ syncEnv body + , MakeAvailable skip3 bodyEnv''' instr3 <- makeAvailable lhs bodyEnv'' -> + instr1 $ instr2 $ buildSpawn + (f (fenvFSkipMany skip2 bndEnv) Sequential $ CtxNormal $ bndDestinations SkipNone) + $ instr3 + $ transform bodyEnv''' parallelism (weaken (skipWeakenIdx skip3 .> skipWeakenIdx skip2 .> skipWeakenIdx skip1) ctx) body + + -- Control flow + PAcond cond true false -> TransformSchedule $ \fenv parallelism dest -> + case fprj (varIdx cond) fenv of + Just (FutureScalar tp signal ref) -> + buildAwait [signal] $ + withRef (GroundRscalar tp) ref $ \skip var -> + buildAcond (Var (varType cond) (varIdx var)) + (weaken' (skipWeakenIdx skip) $ transformSub fenv parallelism dest true) + (weaken' (skipWeakenIdx skip) $ transformSub fenv parallelism dest false) + buildReturn + _ -> variableMissing + + PAwhile us fn initial -> TransformSchedule $ \fenv _ ctx -> + transformWhile fenv us fn initial ctx + + PContinue next -> TransformSchedule $ \fenv parallelism ctx -> + case ctx of + CtxNormal _ -> internalError "Loop impossible" + CtxLoop _ destBool dest _ -> + buildSeq + (writeLoopCondition destBool True) + (transform fenv parallelism (CtxNormal dest) next) + + PBreak us vars -> TransformSchedule $ \fenv _ ctx -> + case ctx of + CtxNormal _ -> internalError "Loop impossible" + CtxLoop release destBool _ dest -> + buildSeq + (writeLoopCondition destBool False) + $ buildSpawn + release + $ snd (returnValues fenv Parallel (toPartialReturn us vars) dest) buildReturn + +transformWhile + :: IsKernel kernel => + FutureEnv fenv genv + -> Uniquenesses t + -> SyncScheduleFun kernel genv (t -> Loop t) + -> GroundVars genv t + -> Context kernel fenv t + -> BuildSchedule kernel fenv +transformWhile _ _ _ _ CtxLoop{} = internalError "Impossible nesting of loops" +transformWhile fenv us (Plam lhs (Pbody step)) groundInitial (CtxNormal finalDest) + | fenv' <- FEnvFSkip $ FEnvFSkip fenv -- Add two bindings to the environment for a new signal and resolver. + -- Split the environment into an environment for the initial state and an environment for the step function. + , ChainFutureEnv instr1 skip1 initialEnv stepEnv + <- chainFutureEnvironment SkipNone fenv' (variablesToSyncEnv us groundInitial) (environmentDropLHS (syncEnv step) lhs) + -- The index of an already resolved signal. Used to construct the initial state. + , resolved <- weaken (skipWeakenIdx skip1) $ SuccIdx ZeroIdx + -- Creates the environment for the free variables of the step function. + , LoopFutureEnv io1 initial1 localR body <- loopFutureEnv resolved stepEnv + -- Extend the environment with the arguments of the function, + -- And prepare the destination for the result of the function (state for next iteration). + , AwhileInputOutput io2 skip2 instr2 initial2 extendStepEnv' makeContinueDest + <- awhileInputOutput initialEnv resolved lhs us groundInitial + , io <- InputOutputRpair io1 io2 + -- Given the computed types, construct the left hand sides for the function. + , DeclareVars lhsInput _ varsInput <- declareVars $ inputOutputInputR io + , DeclareVars lhsOutput k2 varsOutput <- declareVars $ inputOutputOutputR io + , DeclareVars lhsLocal k3 varsLocal <- declareVars $ localBaseR localR + -- Compose the Skips for all left hand sides of the function. + , funSkip <- lhsSkip lhsLocal + `chainSkip` SkipSucc (SkipSucc $ lhsSkip lhsOutput) + `chainSkip` lhsSkip lhsInput + -- Compose that with the Skips of instr1 and instr2. + , fullSkip <- funSkip `chainSkip` skip2 `chainSkip` (SkipSucc $ SkipSucc skip1) + -- Separate the input into an input for loopFutureEnv (input1) and awhileInputOutput (input2) + , (input1, input2) <- expectPair $ varsInput (k3 .> (weakenSucc $ weakenSucc k2)) + -- Also separate the output + , (output1, output2) <- expectPair $ varsOutput k3 + -- Using this input and output, construct the things regarding loopFutureEnv for within the step function + , LoopFutureEnvInner instr3 release stepEnv' <- body (funSkip `chainSkip` skip2) (mapTupR varIdx input1) (mapTupR varIdx output1) + $ mapTupR varIdx $ varsLocal $ weakenId + -- Construct the destination and context for the step function + , ctx <- CtxLoop + release + (weaken (skipWeakenIdx (lhsSkip lhsLocal `chainSkip` lhsSkip lhsOutput)) + $ DestinationShared ZeroIdx $ SuccIdx ZeroIdx) + (makeContinueDest output2) + (mapTupR + (weaken $ skipWeakenIdx fullSkip) + finalDest) + = buildLet lhsSignal (NewSignal "resolved signal for while") + $ buildEffect (SignalResolve [ZeroIdx]) + $ instr1 + $ instr2 + $ buildAwhile + io + (buildFunLam lhsInput + $ buildFunLam (LeftHandSidePair (LeftHandSideSingle BaseRsignalResolver) (LeftHandSideSingle $ BaseRrefWrite $ GroundRscalar scalarType)) + $ buildFunLam lhsOutput + $ buildFunBody + $ declareSignals localR lhsLocal + $ buildSpawn instr3 + $ transform (extendStepEnv' input2 stepEnv') Parallel ctx step + ) + (TupRpair (mapTupR (weaken $ skipWeakenIdx skip2) initial1) (initial2 SkipNone)) + buildReturn + where + expectPair :: Distributes s => TupR s (t1, t2) -> (TupR s t1, TupR s t2) + expectPair (TupRpair t1 t2) = (t1, t2) + expectPair (TupRsingle t) = pairImpossible t +transformWhile _ _ _ _ _ = internalError "Function impossible" + +writeOutput + :: Destinations fenv t + -> GroundVars fenv t + -> BuildSchedule kernel fenv +writeOutput (TupRsingle dest) (TupRsingle (Var tp idx)) + | Just ref <- destRef dest = + buildEffect (RefWrite (Var (BaseRrefWrite tp) ref) (Var (BaseRground tp) idx)) + $ buildEffect (SignalResolve $ catMaybes [destResolver dest, destResolverWrite dest]) buildReturn + | otherwise = buildReturn -- This return value is ignored (e.g., via LeftHandSideWildcard) +writeOutput (TupRpair dest1 dest2) (TupRpair vars1 vars2) = + buildSeq (writeOutput dest1 vars1) (writeOutput dest2 vars2) +writeOutput TupRunit TupRunit = buildReturn +writeOutput _ _ = internalError "Tuple mismatch" + +writeLoopCondition + :: Destination fenv PrimBool + -> Bool + -> BuildSchedule kernel fenv +writeLoopCondition dest value = + buildLet (LeftHandSideSingle $ BaseRground tp) + (Compute $ Const scalarType $ if value then 1 else 0) + $ writeOutput (TupRsingle $ weaken (weakenSucc weakenId) dest) (TupRsingle $ Var tp ZeroIdx) + where + tp = GroundRscalar scalarType + +returnValues + :: FutureEnv fenv genv + -> Parallelism + -> UpdateTuple genv t1 t2 + -> Destinations fenv t2 + -> (Destinations fenv t1, BuildSchedule kernel fenv -> BuildSchedule kernel fenv) +returnValues fenv par tup dest = (dest', \next -> foldr (buildSeqOrSpawn par) next terms) + where + (dest', terms) = returnValuesSpawns fenv tup dest + +returnValuesSpawns + :: FutureEnv fenv genv + -> UpdateTuple genv t1 t2 + -> Destinations fenv t2 + -> (Destinations fenv t1, [BuildSchedule kernel fenv]) +returnValuesSpawns _ UpdateKeep dest = (dest, []) +returnValuesSpawns env (UpdateSet u var) (TupRsingle dest) = + (TupRunit, [returnValue env u var dest]) +returnValuesSpawns env (UpdatePair up1 up2) (TupRpair dest1 dest2) + | (dest1', s1) <- returnValuesSpawns env up1 dest1 + , (dest2', s2) <- returnValuesSpawns env up2 dest2 + = (TupRpair dest1' dest2', s1 ++ s2) +returnValuesSpawns _ _ _ = internalError "Tuple mismatch" + +returnValue + :: FutureEnv fenv genv + -> Uniqueness t + -> GroundVar genv t + -> Destination fenv t + -> BuildSchedule kernel fenv +returnValue env u var dest + | Just outRef <- destRef dest = case fprj (varIdx var) env of + Nothing -> variableMissing + -- Return scalar value + Just (FutureScalar tp signal ref) -> + buildAwait [signal] + $ withRef (GroundRscalar tp) ref + $ \skip var' -> buildEffect (RefWrite (weaken (skipWeakenIdx skip) $ Var (BaseRrefWrite $ GroundRscalar tp) outRef) var') + $ buildEffect (SignalResolve $ catMaybes [weaken (skipWeakenIdx skip) <$> destResolver dest]) + buildReturn + -- Return shared buffer + Just (FutureBuffer tp ref read Nothing) + | Unique <- u -> internalError "Error in construction of schedule: shared buffer returned as unique buffer" + | otherwise -> + buildAwait [lockSignal read] + $ withRef (GroundRbuffer tp) ref + $ \skip var' -> buildEffect (RefWrite (weaken (skipWeakenIdx skip) $ Var (BaseRrefWrite $ GroundRbuffer tp) outRef) var') + $ buildEffect (SignalResolve $ catMaybes [weaken (skipWeakenIdx skip) <$> destResolver dest]) + buildReturn + Just (FutureBuffer tp ref read (Just write)) + | Shared <- u -> internalError "Error in construction of schedule: unique buffer returned as shared buffer" + | otherwise -> + buildAwait [lockSignal read] + $ withRef (GroundRbuffer tp) ref + $ \skip var' -> buildEffect (RefWrite (weaken (skipWeakenIdx skip) $ Var (BaseRrefWrite $ GroundRbuffer tp) outRef) var') + $ buildEffect (SignalResolve $ catMaybes [weaken (skipWeakenIdx skip) <$> destResolver dest]) + $ buildAwait [weaken (skipWeakenIdx skip) <$> lockSignal write] + $ buildEffect (SignalResolve $ catMaybes [weaken (skipWeakenIdx skip) <$> destResolverWrite dest]) + buildReturn + | otherwise = buildReturn + +buildSeqOrSpawn :: Parallelism -> BuildSchedule kernel fenv -> BuildSchedule kernel fenv -> BuildSchedule kernel fenv +buildSeqOrSpawn Sequential = buildSeq +buildSeqOrSpawn _ = buildSpawn + +data Destination fenv t + = DestinationIgnore + -- In sequential mode, we only use a Ref to store the result. + -- No synchronisation for this result is needed as the next term starts after + -- this term. + | DestinationSequential (Idx fenv (OutputRef t)) + -- Destination of a shared resource, in parallel mode. + -- Signal that the resource is available. The Ref is then filled, + -- and in case of a buffer, grants read access. + | DestinationShared (Idx fenv (OutputRef t)) (Idx fenv SignalResolver) + -- Destination of a unique buffer, in parallel mode. + -- Compared to DestinationShared, it now also gets a signal to grant write + -- access. + | DestinationUnique (Idx fenv (OutputRef t)) (Idx fenv SignalResolver) (Idx fenv SignalResolver) + +destRef :: Destination fenv t -> Maybe (Idx fenv (OutputRef t)) +destRef (DestinationSequential r) = Just r +destRef (DestinationShared r _) = Just r +destRef (DestinationUnique r _ _) = Just r +destRef _ = Nothing + +destResolver :: Destination fenv t -> MaybeSignalResolver fenv +destResolver (DestinationShared _ s) = Just s +destResolver (DestinationUnique _ s _) = Just s +destResolver _ = Nothing + +destResolverWrite :: Destination fenv t -> MaybeSignalResolver fenv +destResolverWrite (DestinationUnique _ _ s) = Just s +destResolverWrite _ = Nothing + +instance Sink Destination where + weaken _ DestinationIgnore = DestinationIgnore + weaken k (DestinationSequential ref) = DestinationSequential (weaken k ref) + weaken k (DestinationShared ref s1) = DestinationShared (weaken k ref) (weaken k s1) + weaken k (DestinationUnique ref s1 s2) = DestinationUnique (weaken k ref) (weaken k s1) (weaken k s2) + +type Destinations fenv = TupR (Destination fenv) + +data Context kernel fenv t where + CtxNormal + :: Destinations fenv t + -> Context kernel fenv t + CtxLoop + :: BuildSchedule kernel fenv + -> Destination fenv PrimBool -- Condition (true = another iteration) + -> Destinations fenv t -- State for next iteration + -> Destinations fenv t -- Result of loop + -> Context kernel fenv (Loop t) + +instance Sink (Context kernel) where + weaken k (CtxNormal dest) = CtxNormal $ mapTupR (weaken k) dest + weaken k (CtxLoop release destCond destContinue destBreak) = CtxLoop + (weaken' k release) + (weaken k destCond) + (mapTupR (weaken k) destContinue) + (mapTupR (weaken k) destBreak) + +data DeclareDestinations kernel fenv genv t where + DeclareDestinations + :: Skip fenv' fenv + -> (BuildSchedule kernel fenv' -> BuildSchedule kernel fenv) + -> (forall fenv''. Skip fenv'' fenv' -> Destinations fenv'' t) + -> FutureEnv fenv' genv + -> DeclareDestinations kernel fenv genv t + +declareDestinations :: Parallelism -> FutureEnv fenv genv -> GLeftHandSide t genv genv' -> Uniquenesses t -> DeclareDestinations kernel fenv genv' t +declareDestinations Sequential env (LeftHandSideSingle tp) (TupRsingle u) = + DeclareDestinations + (SkipSucc $ SkipSucc $ SkipNone) + (buildLet (lhsRef tp) $ NewRef tp) + (\skip' -> TupRsingle $ DestinationSequential $ weaken (skipWeakenIdx skip') ZeroIdx) + (FEnvPush (FEnvFSkip $ FEnvFSkip env) future) + where + future = case tp of + GroundRscalar t -> FutureScalar t Nothing (Right $ SuccIdx ZeroIdx) + GroundRbuffer t + | Shared <- u -> FutureBuffer t (Right $ SuccIdx ZeroIdx) (Move Nothing) Nothing + | otherwise -> FutureBuffer t (Right $ SuccIdx ZeroIdx) (Move Nothing) (Just $ Move Nothing) +declareDestinations Parallel env (LeftHandSideSingle tp) (TupRsingle Shared) = + DeclareDestinations + (SkipSucc $ SkipSucc $ SkipSucc $ SkipSucc $ SkipNone) + (buildLet lhsSignal (NewSignal "declareDeclarations (Shared)") . buildLet (lhsRef tp) (NewRef tp)) + (\skip' -> TupRsingle $ DestinationShared + (weaken (skipWeakenIdx skip') ZeroIdx) + (weaken (skipWeakenIdx skip') $ SuccIdx $ SuccIdx ZeroIdx)) + (FEnvPush (FEnvFSkip $ FEnvFSkip $ FEnvFSkip $ FEnvFSkip env) future) + where + future = case tp of + GroundRscalar t -> FutureScalar t (Just $ SuccIdx $ SuccIdx $ SuccIdx ZeroIdx) (Right $ SuccIdx ZeroIdx) + GroundRbuffer t -> FutureBuffer t (Right $ SuccIdx ZeroIdx) (Move $ Just $ SuccIdx $ SuccIdx $ SuccIdx ZeroIdx) Nothing +declareDestinations Parallel env (LeftHandSideSingle tp@(GroundRbuffer t)) (TupRsingle Unique) = + DeclareDestinations + (SkipSucc $ SkipSucc $ SkipSucc $ SkipSucc $ SkipSucc $ SkipSucc $ SkipNone) + (buildLet lhsSignal (NewSignal "declareDeclarations (Unique, 1)") . buildLet lhsSignal (NewSignal "declareDeclarations (Unique, 2)") . buildLet (lhsRef tp) (NewRef tp)) + (\skip' -> TupRsingle $ DestinationUnique + (weaken (skipWeakenIdx skip') ZeroIdx) + (weaken (skipWeakenIdx skip') $ SuccIdx $ SuccIdx ZeroIdx) + (weaken (skipWeakenIdx skip') $ SuccIdx $ SuccIdx $ SuccIdx $ SuccIdx ZeroIdx)) + (FEnvPush (FEnvFSkip $ FEnvFSkip $ FEnvFSkip $ FEnvFSkip $ FEnvFSkip $ FEnvFSkip env) future) + where + future = FutureBuffer t + (Right $ SuccIdx ZeroIdx) + (Move $ Just $ SuccIdx $ SuccIdx $ SuccIdx ZeroIdx) + (Just $ Move $ Just $ SuccIdx $ SuccIdx $ SuccIdx $ SuccIdx $ SuccIdx ZeroIdx) +declareDestinations _ env (LeftHandSideWildcard tp) _ = + DeclareDestinations + SkipNone + id + (const $ mapTupR (const DestinationIgnore) tp) + env +declareDestinations par env (LeftHandSidePair lhs1 lhs2) us + | DeclareDestinations skip1 instr1 dest1 env1 <- declareDestinations par env lhs1 us1 + , DeclareDestinations skip2 instr2 dest2 env2 <- declareDestinations par env1 lhs2 us2 + = DeclareDestinations + (chainSkip skip2 skip1) + (instr1 . instr2) + (\skip -> dest1 (chainSkip skip skip2) `TupRpair` dest2 skip) + env2 + where + (us1, us2) = pairUniqueness us +declareDestinations _ _ _ _ = internalError "Tuple mismatch between LHS and uniqueness" + +data AcquireArgs kernel fenv f where + -- Captures existential fenv'. + AcquireArgs + :: Skip fenv' fenv + -> [Idx fenv Signal] + -> [Idx fenv SignalResolver] + -> (BuildSchedule kernel fenv' -> BuildSchedule kernel fenv) + -> (forall fenv''. Skip fenv'' fenv' -> SArgs fenv'' f) + -> AcquireArgs kernel fenv f + +acquireArgs :: forall kernel fenv genv f. FutureEnv fenv genv -> SArgs genv f -> AcquireArgs kernel fenv f +acquireArgs env (a :>: as) + | AcquireArgs skip2 signals2 resolvers2 instr2 as' <- acquireArgs env as + , AcquireArg skip1 signals1 resolvers1 instr1 a' <- acquireArg env skip2 a + = AcquireArgs + (chainSkip skip1 skip2) + (catMaybes signals1 ++ signals2) + (catMaybes resolvers1 ++ resolvers2) + (instr2 . instr1) + (\skip -> a' skip :>: as' (chainSkip skip skip1)) +acquireArgs _ ArgsNil = AcquireArgs SkipNone [] [] id $ const ArgsNil + +data AcquireArg kernel fenv fenv1 f where + -- Captures existential fenv'. + AcquireArg + :: Skip fenv' fenv1 + -> [MaybeSignal fenv] + -> [MaybeSignalResolver fenv] + -> (BuildSchedule kernel fenv' -> BuildSchedule kernel fenv1) + -> (forall fenv''. Skip fenv'' fenv' -> SArg fenv'' f) + -> AcquireArg kernel fenv fenv1 f + +acquireArg :: FutureEnv fenv genv -> Skip fenv1 fenv -> SArg genv t -> AcquireArg kernel fenv fenv1 t +acquireArg env skip (SArgScalar var) = case fprj (varIdx var) env of + Just (FutureScalar _ signal (Left idx)) -> + AcquireArg + SkipNone + [signal] + [] + id + (\skip' -> SArgScalar $ Var (varType var) $ + weaken (skipWeakenIdx skip') $ weaken (skipWeakenIdx skip) idx) + Just (FutureScalar tp signal (Right ref)) -> + AcquireArg + (SkipSucc SkipNone) + [signal] + [] + (buildRead (GroundRscalar tp) $ weaken (skipWeakenIdx skip) ref) + (\skip' -> SArgScalar $ Var (varType var) $ weaken (skipWeakenIdx skip') ZeroIdx) + _ -> variableMissing +acquireArg env skip (SArgBuffer modifier var) + | Just (FutureBuffer tp idxOrRef read mwrite) <- fprj (varIdx var) env + , (signals, resolvers) <- bufferSignals modifier read mwrite = case idxOrRef of + Left idx -> + AcquireArg + SkipNone + signals + resolvers + id + (\skip' -> SArgBuffer modifier $ Var (varType var) $ + weaken (skipWeakenIdx skip') $ weaken (skipWeakenIdx skip) idx) + Right ref -> + AcquireArg + (SkipSucc SkipNone) + signals + resolvers + (buildRead (GroundRbuffer tp) $ weaken (skipWeakenIdx skip) ref) + (\skip' -> SArgBuffer modifier $ Var (varType var) $ weaken (skipWeakenIdx skip') ZeroIdx) + | otherwise = variableMissing + where + bufferSignals :: Modifier m -> Lock fenv -> Maybe (Lock fenv) -> ([MaybeSignal fenv], [MaybeSignalResolver fenv]) + bufferSignals In (Lock signal resolver) _ = ([signal], [resolver]) + bufferSignals _ (Lock signal resolver) (Just (Lock wsignal wresolver)) = ([signal, wsignal], [resolver, wresolver]) + bufferSignals _ _ _ = internalError "Expected buffer with write permissions" + +data AcquireVars kernel fenv t where + -- Captures existential fenv'. + AcquireVars + :: Skip fenv' fenv + -> (BuildSchedule kernel fenv' -> BuildSchedule kernel fenv) + -> (forall fenv''. Skip fenv'' fenv' -> ExpVars fenv'' t) + -> AcquireVars kernel fenv t + +acquireVars :: forall kernel fenv genv t. FutureEnv fenv genv -> ExpVars genv t -> AcquireVars kernel fenv t +acquireVars env = go SkipNone + where + go :: Skip fenv1 fenv -> ExpVars genv s -> AcquireVars kernel fenv1 s + go skip (TupRpair vars1 vars2) + | AcquireVars skip1 instr1 vars1' <- go skip vars1 + , AcquireVars skip2 instr2 vars2' <- go (chainSkip skip1 skip) vars2 + = AcquireVars + (chainSkip skip2 skip1) + (instr1 . instr2) + (\skip' -> vars1' (chainSkip skip' skip2) `TupRpair` vars2' skip') + go _ TupRunit = AcquireVars SkipNone id (const TupRunit) + go skip (TupRsingle var) = case fprj (varIdx var) env of + Just (FutureScalar _ signal (Left idx)) -> + AcquireVars + SkipNone + (buildAwait [weaken (skipWeakenIdx skip) <$> signal]) + (\skip' -> + TupRsingle $ Var (varType var) $ + weaken (skipWeakenIdx skip') $ weaken (skipWeakenIdx skip) idx + ) + Just (FutureScalar tp signal (Right ref)) -> + AcquireVars + (SkipSucc SkipNone) + (buildAwait [weaken (skipWeakenIdx skip) <$> signal] . + buildLet (LeftHandSideSingle $ BaseRground $ GroundRscalar tp) (RefRead $ Var (BaseRref $ GroundRscalar tp) $ weaken (skipWeakenIdx skip) ref)) + (\skip' -> TupRsingle $ Var (varType var) $ weaken (skipWeakenIdx skip') ZeroIdx) + _ -> variableMissing + +data AcquireSome kernel fenv fenv1 genv where + AcquireSome + :: Skip fenv' fenv1 + -> [Idx fenv Signal] + -> [Idx fenv SignalResolver] + -> (BuildSchedule kernel fenv' -> BuildSchedule kernel fenv1) + -> (forall fenv''. Skip fenv'' fenv' -> PartialEnv (Idx fenv'') genv) + -> AcquireSome kernel fenv fenv1 genv + +-- Assumes that buffers are only require read access. +acquireSome :: forall kernel fenv genv. FutureEnv fenv genv -> IdxSet genv -> AcquireSome kernel fenv fenv genv +acquireSome = go SkipNone SkipNone + where + go :: Skip fenv fenv0 -> Skip fenv1 fenv -> FutureEnv fenv0 genv' -> IdxSet genv' -> AcquireSome kernel fenv fenv1 genv' + go _ _ _ (IdxSet PEnd) = AcquireSome SkipNone [] [] id (const PEnd) + go envSkip skip (FEnvFSkip env) indices = go (SkipSucc envSkip) skip env indices + go envSkip skip env (IdxSet (PNone indices)) + | AcquireSome skip1 signals resolvers instr penv <- go envSkip skip env' (IdxSet indices) = + AcquireSome skip1 signals resolvers instr (\skip' -> PNone $ penv skip') + where + env' = case env of + FEnvPush e _ -> e + FEnvGSkip e -> e + -- FEnvFSkip impossible, already handled in previous case of 'go'. + _ -> FEnvEnd + go envSkip skip (FEnvPush env future) (IdxSet (PPush indices _)) + | AcquireSome skip1 signals resolvers instr penv <- go envSkip skip env (IdxSet indices) = + case future of + FutureScalar _ signal (Left idx) -> + AcquireSome + skip1 + (appendMaybe (weaken (skipWeakenIdx envSkip) <$> signal) signals) + resolvers + instr + (\skip' -> PPush (penv skip') $ + weaken (skipWeakenIdx skip') $ weaken (skipWeakenIdx skip1) + $ weaken (skipWeakenIdx skip) $ weaken (skipWeakenIdx envSkip) idx) + FutureScalar tp signal (Right ref) -> + AcquireSome + (chainSkip (SkipSucc SkipNone) skip1) + (appendMaybe (weaken (skipWeakenIdx envSkip) <$> signal) signals) + resolvers + (instr . buildRead (GroundRscalar tp) + (weaken (skipWeakenIdx skip1) $ weaken (skipWeakenIdx skip) $ weaken (skipWeakenIdx envSkip) ref)) + (\skip' -> PPush (penv $ SkipSucc skip') $ weaken (skipWeakenIdx skip') ZeroIdx) + FutureBuffer _ _ _ (Just _) -> internalError "Expected read-only buffer" + FutureBuffer _ (Left idx) (Lock signal resolver) Nothing -> + AcquireSome + skip1 + (appendMaybe (weaken (skipWeakenIdx envSkip) <$> signal) signals) + (appendMaybe (weaken (skipWeakenIdx envSkip) <$> resolver) resolvers) + instr + (\skip' -> PPush (penv skip') $ + weaken (skipWeakenIdx skip') $ weaken (skipWeakenIdx skip1) + $ weaken (skipWeakenIdx skip) $ weaken (skipWeakenIdx envSkip) idx) + FutureBuffer tp (Right ref) (Lock signal resolver) Nothing -> + AcquireSome + (chainSkip (SkipSucc SkipNone) skip1) + (appendMaybe (weaken (skipWeakenIdx envSkip) <$> signal) signals) + (appendMaybe (weaken (skipWeakenIdx envSkip) <$> resolver) resolvers) + (instr . buildRead (GroundRbuffer tp) + (weaken (skipWeakenIdx skip1) $ weaken (skipWeakenIdx skip) $ weaken (skipWeakenIdx envSkip) ref)) + (\skip' -> PPush (penv $ SkipSucc skip') $ weaken (skipWeakenIdx skip') ZeroIdx) + go _ _ _ _ = variableMissing + +data DeclareOutput kernel fenv t where + DeclareOutput :: Skip fenv' fenv + -> BLeftHandSide (Output t) fenv fenv' + -> (forall fenv''. Skip fenv'' fenv' -> Destinations fenv'' t) + -> DeclareOutput kernel fenv t + +declareOutput :: GroundsR t -> DeclareOutput kernel fenv t +declareOutput (TupRpair t1 t2) + | DeclareOutput skip1 lhs1 dest1 <- declareOutput t1 + , DeclareOutput skip2 lhs2 dest2 <- declareOutput t2 + = DeclareOutput + (chainSkip skip2 skip1) + (LeftHandSidePair lhs1 lhs2) + (\skip -> dest1 (chainSkip skip skip2) `TupRpair` dest2 skip) +declareOutput (TupRsingle tp) + | Refl <- inputSingle tp + = DeclareOutput + (SkipSucc $ SkipSucc SkipNone) + (LeftHandSidePair (LeftHandSideSingle $ BaseRsignalResolver) (LeftHandSideSingle $ BaseRrefWrite tp)) + $ \skip -> TupRsingle $ DestinationShared + (weaken (skipWeakenIdx skip) ZeroIdx) + (weaken (skipWeakenIdx skip) $ SuccIdx ZeroIdx) +declareOutput TupRunit + = DeclareOutput + SkipNone + (LeftHandSideWildcard TupRunit) + $ \_ -> TupRunit + +data DeclareInput fenv genv genv' t where + DeclareInput + :: Skip fenv' fenv + -> BLeftHandSide (Input t) fenv fenv' + -> FutureEnv fenv' genv' + -> DeclareInput fenv genv genv' t + +declareInput :: GLeftHandSide t genv genv' -> FutureEnv fenv genv -> DeclareInput fenv genv genv' t +declareInput (LeftHandSidePair lhs1 lhs2) fenv + | DeclareInput skip1 lhs1' fenv1 <- declareInput lhs1 fenv + , DeclareInput skip2 lhs2' fenv2 <- declareInput lhs2 fenv1 + = DeclareInput (chainSkip skip2 skip1) (LeftHandSidePair lhs1' lhs2') fenv2 +declareInput (LeftHandSideWildcard tp) fenv = + DeclareInput + SkipNone + (LeftHandSideWildcard $ inputR tp) + fenv +declareInput (LeftHandSideSingle tp) fenv + | Refl <- inputSingle tp = + DeclareInput + (SkipSucc $ SkipSucc SkipNone) + (LeftHandSideSingle BaseRsignal `LeftHandSidePair` LeftHandSideSingle (BaseRref tp)) + $ FEnvPush (FEnvFSkip $ FEnvFSkip fenv) + $ case tp of + GroundRscalar t -> FutureScalar t (Just $ SuccIdx ZeroIdx) (Right ZeroIdx) + GroundRbuffer t -> FutureBuffer t (Right ZeroIdx) (Move $ Just $ SuccIdx ZeroIdx) Nothing + +data AwhileInputOutput kernel fenv genv genv' t where + AwhileInputOutput + :: InputOutputR input output + -- Initial state (& instructions to acquire initial state) + -> Skip fenv1 fenv + -> (BuildSchedule kernel fenv1 -> BuildSchedule kernel fenv) + -> (forall fenv'. Skip fenv' fenv1 -> BaseVars fenv' input) + -- Info for iteration function (ie the body of the loop) + -> (forall fenv'. BaseVars fenv' input -> FutureEnv fenv' genv -> FutureEnv fenv' genv') + -> (forall fenv'. BaseVars fenv' output -> Destinations fenv' t) + -> AwhileInputOutput kernel fenv genv genv' t + +awhileInputOutput + :: FutureEnv fenv genv0 + -- A signal that is already resolved. + -> Idx fenv Signal + -- Left hand side of the iteration function of the while loop. + -> GLeftHandSide t genv genv' + -- The uniquenesses of the state of the loop. + -> Uniquenesses t + -- The initial values of the loop, in the ground environment. + -> GroundVars genv0 t + -> AwhileInputOutput kernel fenv genv genv' t +awhileInputOutput _ _ (LeftHandSideWildcard TupRunit) _ _ = + AwhileInputOutput InputOutputRunit SkipNone id (const TupRunit) (flip const) (const TupRunit) +awhileInputOutput env resolved (LeftHandSidePair lhs1 lhs2) us (TupRpair vars1 vars2) + | AwhileInputOutput io1 skip1 instr1 initial1 env1 dest1 + <- awhileInputOutput env resolved lhs1 us1 vars1 + , AwhileInputOutput io2 skip2 instr2 initial2 env2 dest2 + <- awhileInputOutput (fenvFSkipMany skip1 env) (weaken (skipWeakenIdx skip1) resolved) lhs2 us2 vars2 = + AwhileInputOutput + (InputOutputRpair io1 io2) + (chainSkip skip2 skip1) + (instr1 . instr2) + (\skip -> initial1 (chainSkip skip skip2) `TupRpair` initial2 skip) + (\case + TupRpair v1 v2 -> env2 v2 . env1 v1 + _ -> internalError "Pair impossible") + (\case + TupRpair v1 v2 -> dest1 v1 `TupRpair` dest2 v2 + _ -> internalError "Pair impossible") + where + (us1, us2) = pairUniqueness us +awhileInputOutput env resolved (LeftHandSideWildcard (TupRpair t1 t2)) us vars = + awhileInputOutput env resolved (LeftHandSidePair (LeftHandSideWildcard t1) (LeftHandSideWildcard t2)) us vars +-- Unique buffer +awhileInputOutput env resolved (LeftHandSideSingle _) (TupRsingle Unique) (TupRsingle var) = case fprj (varIdx var) env of + Just (FutureBuffer tp ref (Move readSignal) (Just (Move writeSignal))) -> + let + -- Signal for read access and signal for write access + ioSignals = InputOutputRpair InputOutputRsignal InputOutputRsignal + io = InputOutputRpair ioSignals (InputOutputRref $ GroundRbuffer tp) + + -- Based on 'ref', fill in the first fields of AwhileInputOutput + firstFields = case ref of + Right idx -> + AwhileInputOutput io SkipNone id + (\skip -> TupRpair + (TupRpair + (TupRsingle $ Var BaseRsignal $ weaken (skipWeakenIdx skip) $ fromMaybe resolved readSignal) + (TupRsingle $ Var BaseRsignal $ weaken (skipWeakenIdx skip) $ fromMaybe resolved writeSignal) + ) + (TupRsingle $ Var (BaseRref $ GroundRbuffer tp) $ weaken (skipWeakenIdx skip) idx)) + Left idx -> + AwhileInputOutput io (SkipSucc $ SkipSucc SkipNone) + -- Store value in a reference, to ensure uniform representation in the body of the loop. + (buildLet (lhsRef $ GroundRbuffer tp) (NewRef $ GroundRbuffer tp) + . buildEffect (RefWrite (Var (BaseRrefWrite $ GroundRbuffer tp) ZeroIdx) $ Var (BaseRground $ GroundRbuffer tp) + $ weaken (weakenSucc $ weakenSucc weakenId) idx)) + (\skip -> TupRpair + (TupRpair + -- Weaken the existing signals by 2 + (TupRsingle $ Var BaseRsignal $ weaken (skipWeakenIdx $ SkipSucc $ SkipSucc skip) $ fromMaybe resolved readSignal) + (TupRsingle $ Var BaseRsignal $ weaken (skipWeakenIdx $ SkipSucc $ SkipSucc skip) $ fromMaybe resolved writeSignal) + ) + -- Pass the newly constructed reference + (TupRsingle $ Var (BaseRref $ GroundRbuffer tp) $ weaken (skipWeakenIdx skip) $ SuccIdx ZeroIdx)) + in + firstFields + (\case + TupRpair (TupRpair (TupRsingle signal1) (TupRsingle signal2)) (TupRsingle ref') -> + let future = FutureBuffer tp (Right $ varIdx ref') (Move $ Just $ varIdx signal1) $ Just $ Move $ Just $ varIdx signal2 + in (`FEnvPush` future) + _ -> internalError "Input tuple impossible") + (\case + TupRpair (TupRpair (TupRsingle resolver1) (TupRsingle resolver2)) (TupRsingle ref') -> + TupRsingle $ DestinationUnique (varIdx ref') (varIdx resolver1) (varIdx resolver2) + _ -> internalError "Output tuple impossible") + Just (FutureBuffer _ _ _ Nothing) -> internalError "FutureBuffer with write lock expected" + Just (FutureBuffer _ _ _ _) -> internalError "FutureBuffer with Move locks expected" + Just (FutureScalar tp _ _) -> bufferImpossible tp + Nothing -> variableMissing +-- Shared buffer +awhileInputOutput env resolved (LeftHandSideSingle (GroundRbuffer _)) _ (TupRsingle var) = case fprj (varIdx var) env of + Just (FutureBuffer tp ref (Move readSignal) Nothing) -> + let + -- Only a signal for read access + ioSignals = InputOutputRsignal + io = InputOutputRpair ioSignals (InputOutputRref $ GroundRbuffer tp) + + -- Based on 'ref', fill in the first fields of AwhileInputOutput + firstFields = case ref of + Right idx -> + AwhileInputOutput io SkipNone id + (\skip -> TupRpair + (TupRsingle $ Var BaseRsignal $ weaken (skipWeakenIdx skip) $ fromMaybe resolved readSignal) + (TupRsingle $ Var (BaseRref $ GroundRbuffer tp) $ weaken (skipWeakenIdx skip) idx)) + Left idx -> + AwhileInputOutput io (SkipSucc $ SkipSucc SkipNone) + -- Store value in a reference, to ensure uniform representation in the body of the loop. + (buildLet (lhsRef $ GroundRbuffer tp) (NewRef $ GroundRbuffer tp) + . buildEffect (RefWrite (Var (BaseRrefWrite $ GroundRbuffer tp) ZeroIdx) $ Var (BaseRground $ GroundRbuffer tp) + $ weaken (weakenSucc $ weakenSucc weakenId) idx)) + (\skip -> TupRpair + -- Weaken the existing signals by 2 + (TupRsingle $ Var BaseRsignal $ weaken (skipWeakenIdx $ SkipSucc $ SkipSucc skip) $ fromMaybe resolved readSignal) + -- Pass the newly constructed reference + (TupRsingle $ Var (BaseRref $ GroundRbuffer tp) $ weaken (skipWeakenIdx skip) $ SuccIdx ZeroIdx)) + in + firstFields + (\case + TupRpair (TupRsingle signal1) (TupRsingle ref') -> + let future = FutureBuffer tp (Right $ varIdx ref') (Move $ Just $ varIdx signal1) $ Nothing + in (`FEnvPush` future) + _ -> internalError "Input tuple impossible") + (\case + TupRpair (TupRsingle resolver1) (TupRsingle ref') -> + TupRsingle $ DestinationShared (varIdx ref') (varIdx resolver1) + _ -> internalError "Output tuple impossible") + Just (FutureBuffer _ _ _ Just{}) -> internalError "FutureBuffer without write lock expected" + Just (FutureBuffer _ _ _ _) -> internalError "FutureBuffer with Move locks expected" + Just (FutureScalar tp _ _) -> bufferImpossible tp + Nothing -> variableMissing +-- Scalar +awhileInputOutput env resolved (LeftHandSideSingle (GroundRscalar tp')) _ (TupRsingle var) = case fprj (varIdx var) env of + Just (FutureScalar tp signal ref) -> + let + -- Only a signal to announce that the reference is filled. + ioSignals = InputOutputRsignal + io = InputOutputRpair ioSignals (InputOutputRref $ GroundRscalar tp) + + -- Based on 'ref', fill in the first fields of AwhileInputOutput + firstFields = case ref of + Right idx -> + AwhileInputOutput io SkipNone id + (\skip -> TupRpair + (TupRsingle $ Var BaseRsignal $ weaken (skipWeakenIdx skip) $ fromMaybe resolved signal) + (TupRsingle $ Var (BaseRref $ GroundRscalar tp) $ weaken (skipWeakenIdx skip) idx)) + Left idx -> + AwhileInputOutput io (SkipSucc $ SkipSucc SkipNone) + -- Store value in a reference, to ensure uniform representation in the body of the loop. + (buildLet (lhsRef $ GroundRscalar tp) (NewRef $ GroundRscalar tp) + . buildEffect (RefWrite (Var (BaseRrefWrite $ GroundRscalar tp) ZeroIdx) $ Var (BaseRground $ GroundRscalar tp) + $ weaken (weakenSucc $ weakenSucc weakenId) idx)) + (\skip -> TupRpair + -- Weaken the existing signals by 2 + (TupRsingle $ Var BaseRsignal $ weaken (skipWeakenIdx $ SkipSucc $ SkipSucc skip) $ fromMaybe resolved signal) + -- Pass the newly constructed reference + (TupRsingle $ Var (BaseRref $ GroundRscalar tp) $ weaken (skipWeakenIdx skip) $ SuccIdx ZeroIdx)) + in + firstFields + (\case + TupRpair (TupRsingle signal1) (TupRsingle ref') -> + let future = FutureScalar tp (Just $ varIdx signal1) (Right $ varIdx ref') + in (`FEnvPush` future) + _ -> internalError "Input tuple impossible") + (\case + TupRpair (TupRsingle resolver1) (TupRsingle ref') -> + TupRsingle $ DestinationShared (varIdx ref') (varIdx resolver1) + _ -> internalError "Output tuple impossible") + Just (FutureBuffer _ _ _ _) -> bufferImpossible tp' + Nothing -> variableMissing +awhileInputOutput env resolved (LeftHandSideWildcard (TupRsingle tp)) us var + | AwhileInputOutput io skip instr initial _ dest <- awhileInputOutput env resolved (LeftHandSideSingle tp) us var + = AwhileInputOutput io skip instr initial (flip const) dest +awhileInputOutput _ _ _ _ _ = internalError "Tuple mismatch" + +data MakeAvailable kernel fenv genv' where + -- Captures existential fenv' + MakeAvailable + :: Skip fenv' fenv + -> FutureEnv fenv' genv' + -> (BuildSchedule kernel fenv' -> BuildSchedule kernel fenv) + -> MakeAvailable kernel fenv genv' + +makeAvailable :: GLeftHandSide t genv genv' -> FutureEnv fenv genv' -> MakeAvailable kernel fenv genv' +makeAvailable lhs = go (lhsSkip lhs) SkipNone + where + go :: Skip genv' genv -> Skip fenv2 fenv1 -> FutureEnv fenv1 genv' -> MakeAvailable kernel fenv2 genv' + go SkipNone envSkip env = MakeAvailable SkipNone (fenvFSkipMany envSkip env) id + go skip envSkip (FEnvFSkip env) = go skip (SkipSucc envSkip) env + go _ _ FEnvEnd = MakeAvailable SkipNone FEnvEnd id + go (SkipSucc' skip) envSkip (FEnvGSkip env) + | MakeAvailable s env' instr <- go skip envSkip env + = MakeAvailable s (FEnvGSkip env') instr + go (SkipSucc' skip) envSkip (FEnvPush env future) + | MakeAvailable s env' instr <- go skip envSkip env + = case future of + FutureScalar tp signal (Right ref) -> + MakeAvailable + (SkipSucc' s) + (FEnvPush (FEnvFSkip env') $ FutureScalar tp Nothing $ Left ZeroIdx) + (instr + . buildAwait [weaken (skipWeakenIdx s) . weaken (skipWeakenIdx envSkip) <$> signal] + . buildRead (GroundRscalar tp) (weaken (skipWeakenIdx s) $ weaken (skipWeakenIdx envSkip) ref)) + FutureBuffer tp (Right ref) (Lock signal resolver) Nothing -> + MakeAvailable + (SkipSucc' s) + (FEnvPush (FEnvFSkip env') $ FutureBuffer tp (Left ZeroIdx) + (Lock Nothing $ weaken (skipWeakenIdx $ SkipSucc' s) . weaken (skipWeakenIdx envSkip) <$> resolver) + Nothing) + (instr + . buildAwait [weaken (skipWeakenIdx s) . weaken (skipWeakenIdx envSkip) <$> signal] + . buildRead (GroundRbuffer tp) (weaken (skipWeakenIdx s) $ weaken (skipWeakenIdx envSkip) ref)) + FutureBuffer tp (Right ref) (Lock signal resolver) (Just (Lock wsignal wresolver)) -> + MakeAvailable + (SkipSucc' s) + (FEnvPush (FEnvFSkip env') $ FutureBuffer tp (Left ZeroIdx) + (Lock Nothing $ weaken (skipWeakenIdx $ SkipSucc' s) . weaken (skipWeakenIdx envSkip) <$> resolver) + (Just $ Lock Nothing $ weaken (skipWeakenIdx $ SkipSucc' s) . weaken (skipWeakenIdx envSkip) <$> wresolver)) + (instr + . buildAwait (fmap (fmap $ weaken (skipWeakenIdx s) . weaken (skipWeakenIdx envSkip)) [signal, wsignal]) + . buildRead (GroundRbuffer tp) (weaken (skipWeakenIdx s) $ weaken (skipWeakenIdx envSkip) ref)) + _ -> MakeAvailable s (FEnvPush env' $ weaken (skipWeakenIdx s .> skipWeakenIdx envSkip) future) instr + +environmentDropLHS :: PartialEnv f env1 -> LeftHandSide s t env env1 -> PartialEnv f env +environmentDropLHS env (LeftHandSideWildcard _) = env +environmentDropLHS env (LeftHandSideSingle _) = partialEnvTail env +environmentDropLHS env (LeftHandSidePair l1 l2) = environmentDropLHS (environmentDropLHS env l2) l1 + +appendMaybe :: Maybe a -> [a] -> [a] +appendMaybe = maybe id (:) + +variableMissing :: HasCallStack => a +variableMissing = internalError "Variable missing in FutureEnv (and possibly in SyncEnv)" + +rnfSchedule' :: IsKernel kernel => UniformSchedule kernel env -> () +rnfSchedule' Return = () +rnfSchedule' (Alet lhs bnd body) = rnfLeftHandSide rnfBaseR lhs `seq` rnfBinding bnd `seq` rnfSchedule' body +rnfSchedule' (Effect eff next) = rnfEffect eff `seq` rnfSchedule' next +rnfSchedule' (Acond c true false next) = rnfExpVar c `seq` rnfSchedule' true `seq` rnfSchedule' false `seq` rnfSchedule' next +rnfSchedule' (Awhile io body initial next) = rnfInputOutputR io `seq` rnfSchedule body `seq` rnfTupR rnfBaseVar initial `seq` rnfSchedule' next +rnfSchedule' (AwhileSeq io body ini next) = rnfInputOutputR io `seq` rnfSchedule body `seq` rnfTupR rnfBaseVar ini `seq` rnfSchedule' next +rnfSchedule' (Spawn a b) = rnfSchedule' a `seq` rnfSchedule' b + +rnfBinding :: Binding env t -> () +rnfBinding (Compute e) = rnfOpenExp e +rnfBinding (NewSignal _) = () +rnfBinding (NewRef r) = rnfGroundR r +rnfBinding (Alloc shr tp sz) = rnfShapeR shr `seq` rnfScalarType tp `seq` rnfTupR rnfExpVar sz +rnfBinding (Use tp n buffer) = buffer `seq` n `seq` rnfScalarType tp +rnfBinding (Unit var) = rnfExpVar var +rnfBinding (RefRead r) = rnfBaseVar r + +rnfEffect :: IsKernel kernel => Effect kernel env -> () +rnfEffect (Exec md kernel args) = rnf' md `seq` rnf' kernel `seq` rnfSArgs args +rnfEffect (SignalAwait signals) = rnf signals +rnfEffect (SignalResolve signals) = rnf signals +rnfEffect (RefWrite ref value) = rnfBaseVar ref `seq` rnfBaseVar value +rnfEffect (Aassert msg cond) = rnf msg `seq` rnfOpenExp cond +rnfEffect (Atrace msg desc) = rnf msg `seq` rnf' desc + +rnfBaseVar :: BaseVar env t -> () +rnfBaseVar = rnfVar rnfBaseR + +rnfSArgs :: SArgs env args -> () +rnfSArgs ArgsNil = () +rnfSArgs (arg :>: args) = rnfSArg arg `seq` rnfSArgs args + +rnfSArg :: SArg env t -> () +rnfSArg (SArgScalar var) = rnfExpVar var +rnfSArg (SArgBuffer mod' var) = mod' `seq` rnfGroundVar var + +rnfInputOutputR :: InputOutputR input output -> () +rnfInputOutputR (InputOutputRpair io1 io2) = rnfInputOutputR io1 `seq` rnfInputOutputR io2 +rnfInputOutputR _ = () -- all other options have no information + +makeFunSpawnClosed :: UniformScheduleFun kernel env t -> UniformScheduleFun kernel env t +makeFunSpawnClosed = makeFunSpawnClosed' weakenId + +makeFunSpawnClosed' :: env :> env' -> UniformScheduleFun kernel env t -> UniformScheduleFun kernel env' t +makeFunSpawnClosed' k (Slam lhs f) + | Exists lhs' <- rebuildLHS lhs + , k' <- sinkWithLHS lhs lhs' k = + Slam lhs' $ makeFunSpawnClosed' k' f +makeFunSpawnClosed' k (Sbody s) = Sbody $ makeSpawnClosed' k PEnd s + +-- Maybe we should make this a core utility, if we need this on more locations +type TypedIdxSet t = PartialEnv ((:~:) t) + +makeSpawnClosed' :: env :> env' -> TypedIdxSet Signal env' -> UniformSchedule kernel env -> UniformSchedule kernel env' +makeSpawnClosed' k awaitEnd = \case + Return -> case toList awaitEnd of + [] -> Return + list -> Effect (SignalAwait list) Return + Alet lhs bnd next + | Exists lhs' <- rebuildLHS lhs + , k' <- sinkWithLHS lhs lhs' k -> + Alet lhs' (weaken k bnd) $ makeSpawnClosed' k' (partialEnvSkipLHS lhs' awaitEnd ) next + Effect effect next -> + Effect (weaken' k effect) $ makeSpawnClosed' k awaitEnd next + Acond var true false next + | (trueCount, true') <- addSync true + , (falseCount, false') <- addSync false -> + bindSignals (trueCount `max` falseCount) $ + \skip signals resolvers -> + let + k2 = skipWeakenIdx skip .> k + resolvers' = toList resolvers + in + Acond + (weaken k2 var) + (resolve (drop trueCount resolvers') + $ true' k2 weakenId (take trueCount resolvers')) + (resolve (drop falseCount resolvers') + $ false' k2 weakenId (take falseCount resolvers')) + (makeSpawnClosed' k2 + (unionPartialEnv const signals $ skipWeakenPartialEnv skip awaitEnd) next) + -- Acond (weaken k var) (makeSpawnClosed' k PEnd true) (makeSpawnClosed' k PEnd false) (makeSpawnClosed' k awaitEnd next) + Awhile io step initial next -> + Awhile io (makeFunSpawnClosed' k step) (mapTupR (weaken k) initial) $ makeSpawnClosed' k awaitEnd next + AwhileSeq io step initial next -> + AwhileSeq io (makeFunSpawnClosed' k step) (mapTupR (weaken k) initial) $ makeSpawnClosed' k awaitEnd next + Spawn a b + | (count, a') <- addSync a -> + -- Bind signals for all spawns in 'a', and this current spawn. + bindSignals (count + 1) $ + \skip signals resolvers -> + Spawn + (a' (skipWeakenIdx skip .> k) weakenId (toList resolvers)) + (makeSpawnClosed' (skipWeakenIdx skip .> k) + (unionPartialEnv const signals $ skipWeakenPartialEnv skip awaitEnd) b) + where + bindSignals + :: Int + -> (forall env2. Skip env2 env1 -> TypedIdxSet Signal env2 -> TypedIdxSet SignalResolver env2 -> UniformSchedule kernel env2) + -> UniformSchedule kernel env1 + bindSignals 0 f = f SkipNone PEnd PEnd + bindSignals n f = bindSignals (n - 1) $ \skip signals resolvers -> + Alet lhsSignal (NewSignal "Synchronisation to make schedule spawn-closed") $ + f (SkipSucc' $ SkipSucc' skip) (PNone $ PPush signals Refl) (PPush (PNone resolvers) Refl) + + -- Returns the number of spawns and adds synchronisation points to the end of all Spawns. + addSync + :: UniformSchedule kernel env + -> ( Int -- Number of spawns, excluding spawns in the step function of Awhile loops + -- Function to do the actual transformation + -- The list of SignalResolvers is either as long as the number of spawns, + -- or one longer. + -- If the latter is the case, then the final Return of this term should resolve one of these SignalResolvers. + -- If the former is the case, then no such SignalResolver should be resolved at the end; + -- the SignalResolvers are then only needed for any spawned terms. + , forall env1 env2. env :> env1 -> env2 :> env1 -> [Idx env2 SignalResolver] -> UniformSchedule kernel env1) + addSync Return = + ( 0 + , \_ k2 -> \case + [] -> Return + [s] -> Effect (SignalResolve [weaken k2 s]) Return + _ -> internalError "makeSpawnClosed: number of Signals is incorrect." + ) + addSync (Alet lhs bnd next) = + ( count + , \k1 k2 resolvers -> if + | Exists lhs' <- rebuildLHS lhs + , k1' <- sinkWithLHS lhs lhs' k1 + , k2' <- weakenWithLHS lhs' .> k2 -> + Alet lhs' (weaken k1 bnd) $ next' k1' k2' resolvers + ) + where + (count, next') = addSync next + addSync (Effect effect next) = + ( count + , \k1 k2 resolvers -> + Effect (weaken' k1 effect) $ next' k1 k2 resolvers + ) + where + (count, next') = addSync next + addSync (Acond var true false next) = + -- Since only 'true' or 'false' is executed, not both, + -- we take the maximum of both counts + ( (trueCount `max` falseCount) + nextCount + , \k1 k2 resolvers -> + let + (branchResolvers, nextResolvers) = splitAt (trueCount `max` falseCount) resolvers + in + Acond (weaken k1 var) + -- If the true branch requires fewer signals than the false branch, + -- resolve all the extra signals here. + ( resolve (map (weaken k2) $ drop trueCount branchResolvers) + $ true' k1 k2 (take trueCount branchResolvers) + ) + -- If the false branch requires fewer signals than the true branch, + -- resolve all the extra signals here. + ( resolve (map (weaken k2) $ drop falseCount branchResolvers) + $ false' k1 k2 (take falseCount branchResolvers) + ) + ( next' k1 k2 nextResolvers ) + ) + where + (trueCount, true') = addSync true + (falseCount, false') = addSync false + (nextCount, next') = addSync next + addSync (Awhile io f initial next) = + ( count + , \k1 k2 resolvers -> + Awhile + io + (makeFunSpawnClosed' k1 f) + (mapTupR (weaken k1) initial) + $ next' k1 k2 resolvers + ) + where + (count, next') = addSync next + addSync (AwhileSeq io f initial next) = + ( count + , \k1 k2 resolvers -> + AwhileSeq + io + (makeFunSpawnClosed' k1 f) + (mapTupR (weaken k1) initial) + $ next' k1 k2 resolvers + ) + where + (count, next') = addSync next + addSync (Spawn a b) = + -- +1 for this Spawn + ( aCount + bCount + 1 + , \k1 k2 resolvers -> + Spawn + -- Give the first (aCount + 1) resolvers to term 'a', the others to 'b'. + -- Note that 'b' gets either bCount or (bCount + 1) resolvers, + -- see the comment at the start of 'addSync'. + (a' k1 k2 (take (aCount + 1) resolvers)) + (b' k1 k2 (drop (aCount + 1) resolvers)) + ) + where + (aCount, a') = addSync a + (bCount, b') = addSync b + + toList :: TypedIdxSet t env -> [Idx env t] + toList = map (\(EnvBinding idx Refl) -> idx) . partialEnvToList diff --git a/src/Data/Array/Accelerate/Trafo/Schedule/Uniform/Future.hs b/src/Data/Array/Accelerate/Trafo/Schedule/Uniform/Future.hs new file mode 100644 index 000000000..73d9b1ff3 --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Schedule/Uniform/Future.hs @@ -0,0 +1,890 @@ +{-# LANGUAGE MultiWayIf #-} +{-# LANGUAGE PatternSynonyms #-} +{-# OPTIONS_HADDOCK hide #-} +-- | +-- Module : Data.Array.Accelerate.Trafo.Schedule.Uniform.Future +-- Copyright : [2012..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.Trafo.Schedule.Uniform.Future ( + Future(..), FutureEnv(..), Lock(..), pattern Move, pattern Borrow, + withRef, buildRead, weakenEither, + fprj, fenvDrop, fenvDrop1, fenvFSkipMany, fenvPushLHS, + chainFuture, ChainFuture(..), chainFutureEnvironmentSeqIf, chainFutureEnvironment, ChainFutureEnv(..), + loopFutureEnv, LoopFutureEnv(..), LoopFutureEnvInner(..), + IsNewSignal(..), localBaseR, declareSignals, + subFutureEnvironment, restrictEnvForLHS, + MaybeSignal, MaybeSignalResolver, buildAwait, + + assertFutureEnv, +) where + +import Data.Array.Accelerate.AST.Environment +import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.AST.Schedule.Uniform +import Data.Array.Accelerate.Error +import Data.Array.Accelerate.Representation.Array +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Trafo.Substitution +import Data.Array.Accelerate.Trafo.Schedule.Partial +import Data.Array.Accelerate.Trafo.Schedule.Uniform.Simplify +import Data.Array.Accelerate.Type +import Data.Maybe +import Prelude hiding (id, (.), read) +import Control.Category + +data Future fenv t where + FutureScalar :: ScalarType t + -> MaybeSignal fenv + -> Either (Idx fenv t) (Idx fenv (Ref t)) + -> Future fenv t + + -- A buffer has signals and resolvers grouped in Locks to synchronize + -- read and write access to the buffer. + -- The Ref may be read when the read signal is resolved. + -- Informal properties / invariants: + -- - If the read signal is resolved, then we may read from + -- the array. + -- - If the signals of the read and write access are both + -- resolved, then we may destructively update the array. + -- - The read resolver may only be resolved after the read + -- signal is resolved. + -- - The write resolver may only be resolved after both + -- the read and write signals are resolved. + -- - If the write lock is move, then the read lock should also be move. + -- (See datatype Lock for the meaning of move and borrow) + FutureBuffer :: ScalarType t + -> Either (Idx fenv (Buffer t)) (Idx fenv (Ref (Buffer t))) + -> Lock fenv -- Read access + -> Maybe (Lock fenv) -- Write access, if needed + -> Future fenv (Buffer t) + +withRef + :: forall kernel fenv t. + GroundR t + -> Either (Idx fenv t) (Idx fenv (Ref t)) + -> (forall fenv'. Skip fenv' fenv -> BaseVar fenv' t -> BuildSchedule kernel fenv') + -> BuildSchedule kernel fenv +withRef tp (Left idx) f = f @fenv SkipNone $ Var (BaseRground tp) idx +withRef tp (Right ref) f = + buildLet (LeftHandSideSingle $ BaseRground tp) (RefRead (Var (BaseRref tp) ref)) + $ f @(fenv, t) (SkipSucc SkipNone) $ Var (BaseRground tp) ZeroIdx + +buildRead :: GroundR t -> Idx fenv (Ref t) -> BuildSchedule kernel (fenv, t) -> BuildSchedule kernel fenv +buildRead tp ref = buildLet + (LeftHandSideSingle $ BaseRground tp) + (RefRead $ Var (BaseRref tp) ref) + +-- We call a lock with a resolver a borrow lock, and one without a move lock. +data Lock fenv = + Lock { + lockSignal :: MaybeSignal fenv, + lockResolver :: MaybeSignalResolver fenv + } + +pattern Borrow :: MaybeSignal fenv -> Idx fenv SignalResolver -> Lock fenv +pattern Borrow signal resolver = Lock signal (Just resolver) +pattern Move :: MaybeSignal fenv -> Lock fenv +pattern Move signal = Lock signal Nothing +{-# COMPLETE Borrow, Move #-} + +instance Sink' Lock where + weaken' k (Lock s r) = Lock (weaken k <$> s) (weaken k <$> r) + +instance Sink Future where + weaken k (FutureScalar tp signal ref) = FutureScalar tp (weaken k <$> signal) (weakenEither k ref) + weaken k (FutureBuffer tp ref read write) + = FutureBuffer + tp + (weakenEither k ref) + (weaken' k read) + (weaken' k <$> write) + +weakenEither :: (Sink a, Sink b) => env :> env' -> Either (a env t) (b env s) -> Either (a env' t) (b env' s) +weakenEither k (Left a) = Left $ weaken k a +weakenEither k (Right b) = Right $ weaken k b + +data FutureEnv fenv genv where + FEnvEnd + :: FutureEnv fenv genv + FEnvFSkip + :: FutureEnv fenv genv -> FutureEnv (fenv, t) genv + FEnvGSkip + :: FutureEnv fenv genv -> FutureEnv fenv (genv, t) + FEnvPush + :: FutureEnv fenv genv + -> Future fenv t + -> FutureEnv fenv (genv, t) + +fprj :: forall fenv genv t. Idx genv t -> FutureEnv fenv genv -> Maybe (Future fenv t) +fprj = \idx env -> go idx env SkipNone + where + go :: Idx genv' t -> FutureEnv fenv' genv' -> Skip fenv fenv' -> Maybe (Future fenv t) + go idx (FEnvFSkip env) skip = go idx env (SkipSucc skip) + go ZeroIdx (FEnvPush _ future) skip = Just $ weaken (skipWeakenIdx skip) future + go (SuccIdx idx) (FEnvPush env _) skip = go idx env skip + go (SuccIdx idx) (FEnvGSkip env) skip = go idx env skip + go _ _ _ = Nothing + +fenvDrop :: Skip genv genv' -> FutureEnv fenv genv -> FutureEnv fenv genv' +fenvDrop SkipNone env = env +fenvDrop (SkipSucc skip) env = fenvDrop1 $ fenvDrop skip env + +fenvDrop1 :: FutureEnv fenv (genv', t) -> FutureEnv fenv genv' +fenvDrop1 (FEnvFSkip env) = FEnvFSkip $ fenvDrop1 env +fenvDrop1 (FEnvGSkip env) = env +fenvDrop1 (FEnvPush env _) = env +fenvDrop1 FEnvEnd = FEnvEnd + +fenvFSkipMany :: Skip fenv' fenv -> FutureEnv fenv genv -> FutureEnv fenv' genv +fenvFSkipMany (SkipSucc skip) env = fenvFSkipMany skip $ FEnvFSkip env +fenvFSkipMany SkipNone env = env + +fenvGSkip :: FutureEnv fenv genv -> FutureEnv fenv (genv, t) +fenvGSkip FEnvEnd = FEnvEnd +fenvGSkip env = FEnvGSkip env + +fenvWeakenF :: Skip fenv' fenv -> FutureEnv fenv genv -> FutureEnv fenv' genv +fenvWeakenF SkipNone env = env +fenvWeakenF (SkipSucc skip) env = fenvWeakenF skip (FEnvFSkip env) + +fRemoveBuffers :: FutureEnv fenv genv -> FutureEnv fenv genv +fRemoveBuffers FEnvEnd = FEnvEnd +fRemoveBuffers (FEnvFSkip env) = FEnvFSkip $ fRemoveBuffers env +fRemoveBuffers (FEnvGSkip env) = FEnvGSkip $ fRemoveBuffers env +fRemoveBuffers (FEnvPush env FutureBuffer{}) = FEnvGSkip $ fRemoveBuffers env +fRemoveBuffers (FEnvPush env future) = FEnvPush (fRemoveBuffers env) future + +-- Assumes that both left-hand-sides have the same shape. +fenvPushLHS :: GLeftHandSide t genv genv' -> GLeftHandSide t fenv fenv' -> Uniquenesses t -> FutureEnv fenv genv -> FutureEnv fenv' genv' +fenvPushLHS (LeftHandSideSingle (GroundRscalar tp)) (LeftHandSideSingle _) _ env = + FEnvFSkip env `FEnvPush` FutureScalar tp Nothing (Left ZeroIdx) +fenvPushLHS (LeftHandSideSingle (GroundRbuffer tp)) (LeftHandSideSingle _) (TupRsingle Unique) env = + FEnvFSkip env `FEnvPush` FutureBuffer tp (Left ZeroIdx) (Move Nothing) (Just $ Move Nothing) +fenvPushLHS (LeftHandSideSingle (GroundRbuffer tp)) (LeftHandSideSingle _) _ env = + FEnvFSkip env `FEnvPush` FutureBuffer tp (Left ZeroIdx) (Move Nothing) Nothing +fenvPushLHS (LeftHandSideWildcard _) (LeftHandSideWildcard _) _ env = env +fenvPushLHS (LeftHandSidePair lhs1 lhs2) (LeftHandSidePair lhs1' lhs2') us env + | (us1, us2) <- pairUniqueness us + = fenvPushLHS lhs2 lhs2' us2 $ fenvPushLHS lhs1 lhs1' us1 env +fenvPushLHS _ _ _ _ = internalError "LHS mismatch: expected two left hand sides of the same shape" + +-- TODO: Separate chainFuture(Env) for sequential let bindings? +-- And possibly another for sequential let bindings, where the binding guaranteed does not introduce forks? + +-- Data type for the existentially qualified type variable fenv' used in chainFuture +data ChainFutureEnv kernel fenv genv where + ChainFutureEnv + :: (BuildSchedule kernel fenv' -> BuildSchedule kernel fenv) + -> Skip fenv' fenv + -> FutureEnv fenv' genv + -> FutureEnv fenv' genv + -> ChainFutureEnv kernel fenv genv + +chainFutureEnvironmentSeqIf :: Bool -> FutureEnv fenv genv -> SyncEnv genv -> SyncEnv genv -> ChainFutureEnv kernel fenv genv +chainFutureEnvironmentSeqIf True env left right + | (left', right') <- seqChainFutureEnv env left right + = ChainFutureEnv id SkipNone left' right' +chainFutureEnvironmentSeqIf False env left right = chainFutureEnvironment SkipNone env left right + +chainFutureEnvironment :: Skip fenv' fenv -> FutureEnv fenv genv -> SyncEnv genv -> SyncEnv genv -> ChainFutureEnv kernel fenv' genv +-- If left SyncEnv is empty, no further synchronisations are necessary. We do still have to pass scalars to the left environment. +chainFutureEnvironment skip env PEnd _ = ChainFutureEnv id SkipNone (fRemoveBuffers env') env' + where env' = fenvWeakenF skip env +-- Same, but now for the right subterm. +chainFutureEnvironment skip env _ PEnd = ChainFutureEnv id SkipNone env' (fRemoveBuffers env') + where env' = fenvWeakenF skip env +chainFutureEnvironment skip (FEnvFSkip env) senvLeft senvRight = chainFutureEnvironment (SkipSucc skip) env senvLeft senvRight +-- Scalars don't require further synchronisation +chainFutureEnvironment skip (FEnvPush fenv f@FutureScalar{}) senvLeft senvRight + | ChainFutureEnv instr skip1 fenvLeft fenvRight <- chainFutureEnvironment skip fenv (partialEnvTail senvLeft) (partialEnvTail senvRight) + , f' <- weaken (skipWeakenIdx skip1) $ weaken (skipWeakenIdx skip) f + = ChainFutureEnv instr skip1 (FEnvPush fenvLeft f') (FEnvPush fenvRight f') +-- Buffer used in both subterms +chainFutureEnvironment skip (FEnvPush fenv f) (PPush senvLeft sLeft) (PPush senvRight sRight) + | ChainFuture instr1 skip1 fLeft fRight <- chainFuture (weaken (skipWeakenIdx skip) f) sLeft sRight + , ChainFutureEnv instr2 skip2 fenvLeft fenvRight <- chainFutureEnvironment (chainSkip skip1 skip) fenv senvLeft senvRight + = ChainFutureEnv + (instr1 . instr2) + (chainSkip skip2 skip1) + (FEnvPush fenvLeft $ weaken (skipWeakenIdx skip2) fLeft) + (FEnvPush fenvRight $ weaken (skipWeakenIdx skip2) fRight) +-- Only used in left subterm +chainFutureEnvironment skip (FEnvPush fenv f) (PPush senvLeft _) senvRight + | ChainFutureEnv instr skip1 fenvLeft fenvRight <- chainFutureEnvironment skip fenv senvLeft (partialEnvTail senvRight) + , f' <- weaken (skipWeakenIdx skip1) $ weaken (skipWeakenIdx skip) f + = ChainFutureEnv instr skip1 (FEnvPush fenvLeft f') (fenvGSkip fenvRight) +-- Only used in right subterm +chainFutureEnvironment skip (FEnvPush fenv f) senvLeft (PPush senvRight _) + | ChainFutureEnv instr skip1 fenvLeft fenvRight <- chainFutureEnvironment skip fenv (partialEnvTail senvLeft) senvRight + , f' <- weaken (skipWeakenIdx skip1) $ weaken (skipWeakenIdx skip) f + = ChainFutureEnv instr skip1 (fenvGSkip fenvLeft) (FEnvPush fenvRight f') +-- Index not present +chainFutureEnvironment skip (FEnvGSkip fenv) senvLeft senvRight + | ChainFutureEnv instr skip1 fenvLeft fenvRight <- chainFutureEnvironment skip fenv (partialEnvTail senvLeft) (partialEnvTail senvRight) + = ChainFutureEnv instr skip1 (fenvGSkip fenvLeft) (fenvGSkip fenvRight) +chainFutureEnvironment _ _ _ _ = internalError "Illegal case. The keys of the FutureEnv should be the union of the keys of the two SyncEnvs." + + +-- Data type for the existentially qualified type variable fenv' used in chainFuture +data ChainFuture kernel fenv t where + ChainFuture + :: (BuildSchedule kernel fenv' -> BuildSchedule kernel fenv) + -> Skip fenv' fenv + -> Future fenv' t + -> Future fenv' t + -> ChainFuture kernel fenv t + +chainFuture :: Future fenv t -> Sync t -> Sync t -> ChainFuture kernel fenv t +chainFuture (FutureScalar tp _ _) SyncRead _ = bufferImpossible tp +chainFuture (FutureScalar tp _ _) SyncWrite _ = bufferImpossible tp + +-- Read before read, without a release +-- Left Right +-- Read --> X -> X +-- \ / +-- ----- +chainFuture f@(FutureBuffer _ _ (Move _) mwrite) SyncRead SyncRead + | Just _ <- mwrite = internalError "Expected a FutureBuffer without write lock" + | Nothing <- mwrite + = ChainFuture + -- This doesn't require any additional signals + id + SkipNone + f + f + +-- Read before read +-- Left Right +-- ------- +-- / \ +-- Read --> X -> X --> +-- \ / +-- ----- +chainFuture (FutureBuffer tp ref (Borrow s r) mwrite) SyncRead SyncRead + | Just _ <- mwrite = internalError "Expected a FutureBuffer without write lock" + | Nothing <- mwrite + = ChainFuture + -- Create a pair of signal and resolver for both subterms. + -- Spawn a thread which will resolve the final read signal when the two + -- new signals have been resolved. + ( buildLet lhsSignal (NewSignal "chainFuture (read-before-write, left)") + . buildLet lhsSignal (NewSignal "chainFuture (read-before-write, right)") + . buildSpawn (buildEffect (SignalAwait [signal1, signal2]) $ buildEffect (SignalResolve [weaken k r]) buildReturn) + ) + -- Weaken all other identifiers with four, as we introduced two new signals + -- and two new signal resolvers. + ( SkipSucc $ SkipSucc $ SkipSucc $ SkipSucc SkipNone ) + ( FutureBuffer + tp + (weakenEither k ref) + (Lock (weaken k <$> s) (Just resolver1)) + Nothing + ) + ( FutureBuffer + tp + (weakenEither k ref) + (Lock (weaken k <$> s) (Just resolver2)) + Nothing + ) + where + k = weakenSucc $ weakenSucc $ weakenSucc $ weakenSucc weakenId + + signal1 = SuccIdx $ SuccIdx $ SuccIdx ZeroIdx + resolver1 = SuccIdx $ SuccIdx ZeroIdx + signal2 = SuccIdx ZeroIdx + resolver2 = ZeroIdx + +-- Write before read, without release +-- Left Right +-- Read --> X > X +-- / +-- / +-- / +-- Write --> X +-- +-- Note that the left subterm must synchronise its read and write operations itself. +chainFuture (FutureBuffer tp ref (Move readSignal) (Just (Move writeSignal))) SyncWrite SyncRead + = ChainFuture + -- Create a signal to let the read operation in the second subterm only + -- start after the write operation of the first subterm has finished. + ( buildLet lhsSignal $ NewSignal "chainFuture (write-before-read)") + ( SkipSucc $ SkipSucc SkipNone ) + -- The first subterm must resolve the new signal after finishing its write operation. + ( FutureBuffer + tp + (weakenEither k ref) + (Lock (weaken k <$> readSignal) Nothing) + (Just $ Lock (weaken k <$> writeSignal) $ Just writeResolver) + ) + ( FutureBuffer + tp + (weakenEither k ref) + (Lock (Just writeSignal2) Nothing) + Nothing + ) + where + k = weakenSucc $ weakenSucc weakenId + writeSignal2 = SuccIdx ZeroIdx + writeResolver = ZeroIdx + +-- Write before read +-- Left Right +-- ------- +-- / \ +-- Read --> X > X --> +-- / +-- / +-- / +-- Write --> X -------------> +-- Note that the left subterm must synchronise its read and write operations itself. +chainFuture (FutureBuffer tp ref (Borrow readSignal readRelease) (Just (Borrow writeSignal writeRelease))) SyncWrite SyncRead + = ChainFuture + -- Create a signal (signal1) to let the read operation in the second subterm only + -- start after the write operation of the first subterm has finished. + -- Also create signals (signal2 and signal3) to denote that the read operations + -- of respectively the left and right subterm have finished. + -- 'readRelease' will be resolved when signal2 and signal3 are both resolved. + -- 'writeRelease' will be resolved when signal1 is resolved. + ( buildLet lhsSignal (NewSignal "chainFuture (write-before-read, 1)") + . buildLet lhsSignal (NewSignal "chainFuture (write-before-read, 2)") + . buildLet lhsSignal (NewSignal "chainFuture (write-before-read, 3)") + . buildSpawn (buildEffect (SignalAwait [signal2, signal3]) $ buildEffect (SignalResolve [weaken k readRelease]) buildReturn) + . buildSpawn (buildEffect (SignalAwait [signal1]) $ buildEffect (SignalResolve [weaken k writeRelease]) buildReturn) + ) + ( SkipSucc $ SkipSucc $ SkipSucc $ SkipSucc $ SkipSucc $ SkipSucc SkipNone ) + ( FutureBuffer + tp + (weakenEither k ref) + (Lock (weaken k <$> readSignal) (Just resolver2)) + (Just $ Lock (weaken k <$> writeSignal) (Just resolver1)) + ) + ( FutureBuffer + tp + (weakenEither k ref) + (Lock (Just signal1) (Just resolver3)) + Nothing + ) + where + k = weakenSucc $ weakenSucc $ weakenSucc $ weakenSucc $ weakenSucc $ weakenSucc weakenId + + signal1 = SuccIdx $ SuccIdx $ SuccIdx $ SuccIdx $ SuccIdx ZeroIdx + resolver1 = SuccIdx $ SuccIdx $ SuccIdx $ SuccIdx ZeroIdx + signal2 = SuccIdx $ SuccIdx $ SuccIdx ZeroIdx + resolver2 = SuccIdx $ SuccIdx ZeroIdx + signal3 = SuccIdx ZeroIdx + resolver3 = ZeroIdx + +-- Write before read, with a write release +-- Left Right +-- Read --> X > X +-- / +-- / +-- / +-- Write --> X -------------> +-- Note that the left subterm must synchronise its read and write operations itself. +chainFuture (FutureBuffer tp ref (Move readSignal) (Just (Borrow writeSignal writeRelease))) SyncWrite SyncRead + = ChainFuture + -- Create a signal to let the read operation in the second subterm only + -- start after the write operation of the first subterm has finished. + -- 'writeSignal' can be resolved when this newly introduced signal + -- is resolved. + ( buildLet lhsSignal (NewSignal "chainFuture (write-before-read, with release)") + . buildSpawn (buildEffect (SignalAwait [signal1]) $ buildEffect (SignalResolve [weaken k writeRelease]) buildReturn) + ) + -- Weaken all other identifiers with two, as we introduced a new signal + -- and a new signal resolver + (SkipSucc $ SkipSucc SkipNone) + ( FutureBuffer + tp + (weakenEither k ref) + (Lock (weaken k <$> readSignal) Nothing) + (Just $ Lock (weaken k <$> writeSignal) (Just resolver1)) + ) + ( FutureBuffer + tp + (weakenEither k ref) + (Lock (Just signal1) Nothing) + Nothing + ) + where + k = weakenSucc $ weakenSucc weakenId + signal1 = SuccIdx ZeroIdx + resolver1 = ZeroIdx + +-- Invalid cases of write-before-read +chainFuture (FutureBuffer _ _ _ Nothing) SyncWrite SyncRead = internalError "Expected a FutureBuffer with write lock" +chainFuture (FutureBuffer _ _ (Borrow _ _) (Just (Move _))) SyncWrite SyncRead = internalError "Illegal FutureBuffer with borrow read access and move write access" + +-- Read before write +-- Left Right +-- ----- +-- / \ +-- Read --> X -> X --> +-- \ +-- \ +-- \ +-- Write ------------> X --> +chainFuture (FutureBuffer tp ref read mwrite) SyncRead SyncWrite + | Nothing <- mwrite = internalError "Expected a FutureBuffer with write lock" + | Just write <- mwrite + = ChainFuture + -- Create a signal to let the write operation in the second subterm only + -- start after the read operation of the first subterm has finished. + -- Also create a signal which will be resolved when the newly introduced signal + -- and the incoming write signal are both resolved. + ( buildLet lhsSignal (NewSignal "chainFuture (read-before-write, 1)") + . buildLet lhsSignal (NewSignal "chainFuture (read-before-write, 2)") + . buildSpawn (buildAwait [weaken k <$> lockSignal write, Just signal1] $ buildEffect (SignalResolve [resolver2]) buildReturn) + ) + -- Weaken all other identifiers with four, as we introduced two new signals + -- and two new signal resolvers. + (SkipSucc $ SkipSucc $ SkipSucc $ SkipSucc SkipNone) + -- The first subterm must resolve the new signal after finishing its read operation. + ( FutureBuffer + tp + (weakenEither k ref) + (Borrow (weaken k <$> lockSignal read) resolver1) + Nothing + ) + -- The second subterm must wait on the signal before it can start the write operation. + ( FutureBuffer + tp + (weakenEither k ref) + (weaken' k read) + (Just $ Lock (Just signal2) $ fmap (weaken k) $ lockResolver write) -- lockResolver write) + ) + where + k = weakenSucc $ weakenSucc $ weakenSucc $ weakenSucc weakenId + + signal1 = SuccIdx $ SuccIdx $ SuccIdx ZeroIdx + resolver1 = SuccIdx $ SuccIdx ZeroIdx + signal2 = SuccIdx ZeroIdx + resolver2 = ZeroIdx + +-- Write before write +-- Left Right +-- Read --> X > X --> +-- \ / +-- X +-- / \ +-- Write --> X > X --> +chainFuture (FutureBuffer tp ref read mwrite) SyncWrite SyncWrite + | Nothing <- mwrite = internalError "Expected a FutureBuffer with write lock" + | Just write <- mwrite + = ChainFuture + -- Create two signals (signal1 and signal2) to let the first subterm + -- inform that respectively its read or write operations have finished. + ( buildLet lhsSignal (NewSignal "chainFuture (write-before-write, 1)") + . buildLet lhsSignal (NewSignal "chainFuture (write-before-write, 1)") + ) + (SkipSucc $ SkipSucc $ SkipSucc $ SkipSucc SkipNone) + ( FutureBuffer + tp + (weakenEither k ref) + (Borrow (weaken k <$> lockSignal read) resolver1) + (Just $ Borrow (weaken k <$> lockSignal write) resolver2) + ) + ( FutureBuffer + tp + (weakenEither k ref) + (Lock (Just signal2) $ weaken k <$> lockResolver read) + (Just $ Lock (Just signal1) $ weaken k <$> lockResolver write) + ) + where + k = weakenSucc $ weakenSucc $ weakenSucc $ weakenSucc weakenId + + signal1 = SuccIdx $ SuccIdx $ SuccIdx ZeroIdx + resolver1 = SuccIdx $ SuccIdx ZeroIdx + signal2 = SuccIdx ZeroIdx + resolver2 = ZeroIdx + +-- Variant of chainFutureEnv, where there is no synchronisation needed between the two subterms. +-- This is used when the two terms are executed sequentially. +-- This case does require synchronisation with other parts of the program. +-- Assumes that the first term introduces no forks. +seqChainFutureEnv :: FutureEnv fenv genv -> SyncEnv genv -> SyncEnv genv -> (FutureEnv fenv genv, FutureEnv fenv genv) +seqChainFutureEnv FEnvEnd _ _ = (FEnvEnd, FEnvEnd) +seqChainFutureEnv env PEnd _ = (fRemoveBuffers env, env) +seqChainFutureEnv env _ PEnd = (env, fRemoveBuffers env) +seqChainFutureEnv (FEnvFSkip env) left right + | (envL, envR) <- seqChainFutureEnv env left right + = (FEnvFSkip envL, FEnvFSkip envR) +seqChainFutureEnv (FEnvGSkip env) left right + | (envL, envR) <- seqChainFutureEnv env (partialEnvTail left) (partialEnvTail right) + = (FEnvGSkip envL, FEnvGSkip envR) +seqChainFutureEnv (FEnvPush env future@FutureScalar{}) left right + | (envL, envR) <- seqChainFutureEnv env (partialEnvTail left) (partialEnvTail right) + = (FEnvPush envL future, FEnvPush envR future) +seqChainFutureEnv (FEnvPush env future) (PPush left l) (PPush right r) + | (envL, envR) <- seqChainFutureEnv env left right + , (futureL, futureR) <- seqChainFuture future l r + = (FEnvPush envL futureL, FEnvPush envR futureR) +seqChainFutureEnv (FEnvPush env future) (PPush left _) right + | (envL, envR) <- seqChainFutureEnv env left (partialEnvTail right) + = (FEnvPush envL future, FEnvGSkip envR) +seqChainFutureEnv (FEnvPush env future) left (PPush right _) + | (envL, envR) <- seqChainFutureEnv env (partialEnvTail left) right + = (FEnvGSkip envL , FEnvPush envR future) +seqChainFutureEnv (FEnvPush _ FutureBuffer{}) (PNone _) (PNone _) = internalError "FutureEnv & SyncEnv mismatch" + +seqChainFuture :: Future fenv t -> Sync t -> Sync t -> (Future fenv t, Future fenv t) +seqChainFuture (FutureBuffer tp ref (Lock signal resolver) Nothing) SyncRead SyncRead = + ( FutureBuffer tp ref (Lock signal Nothing) Nothing + , FutureBuffer tp ref (Lock Nothing resolver) Nothing + ) +seqChainFuture (FutureBuffer tp ref (Lock signal resolver) (Just (Lock wsignal wresolver))) SyncWrite SyncRead = + ( FutureBuffer tp ref (Lock signal Nothing) (Just $ Lock wsignal wresolver) + , FutureBuffer tp ref (Lock Nothing resolver) Nothing + ) +seqChainFuture (FutureBuffer tp ref (Lock signal resolver) (Just (Lock wsignal wresolver))) SyncRead SyncWrite = + ( FutureBuffer tp ref (Lock signal Nothing) Nothing + , FutureBuffer tp ref (Lock Nothing resolver) (Just $ Lock wsignal wresolver) + ) +seqChainFuture (FutureBuffer tp ref (Lock signal resolver) (Just (Lock wsignal wresolver))) SyncWrite SyncWrite = + ( FutureBuffer tp ref (Lock signal Nothing) (Just $ Lock wsignal Nothing) + , FutureBuffer tp ref (Lock Nothing resolver) (Just $ Lock Nothing wresolver) + ) +seqChainFuture _ _ _ = internalError "FutureBuffer & Sync mismatch" + +data IsNewSignal t where + IsNewSignal :: String -> IsNewSignal (Signal, SignalResolver) + +localBaseR :: TupR IsNewSignal t -> BasesR t +localBaseR (TupRsingle (IsNewSignal _)) = TPair BaseRsignal BaseRsignalResolver +localBaseR (TupRpair t1 t2) = localBaseR t1 `TupRpair` localBaseR t2 +localBaseR TupRunit = TupRunit + +declareSignals :: TupR IsNewSignal t -> BLeftHandSide t env env' -> BuildSchedule kernel env' -> BuildSchedule kernel env +declareSignals _ (LeftHandSideWildcard _) = id +declareSignals (TupRsingle (IsNewSignal name)) lhs = buildLet lhs (NewSignal name) +declareSignals (TupRpair t1 t2) (LeftHandSidePair lhs1 lhs2) = declareSignals t1 lhs1 . declareSignals t2 lhs2 +declareSignals _ _ = internalError "Tuple mismatch" + +pattern TPair :: forall s t. () => forall a b. (t ~ (a, b)) => s a -> s b -> TupR s t +pattern TPair a b = TupRsingle a `TupRpair` TupRsingle b + +data LoopFutureEnv kernel fenv1 genv where + -- Captures existentials input, output and local + LoopFutureEnv :: { + loopEnvIO :: InputOutputR input output, + loopEnvInitial :: BaseVars fenv1 input, + loopEnvLocal :: TupR IsNewSignal local, + loopEnvBody :: forall fenv2. Skip fenv2 fenv1 -> TupR (Idx fenv2) input -> TupR (Idx fenv2) output -> TupR (Idx fenv2) local -> LoopFutureEnvInner kernel fenv2 genv + } -> LoopFutureEnv kernel fenv1 genv + +data LoopFutureEnvInner kernel fenv2 genv where + LoopFutureEnvInner + -- Executed at the start of an iteration + :: BuildSchedule kernel fenv2 + -- Releases the locks at the end of the loop + -> BuildSchedule kernel fenv2 + -> FutureEnv fenv2 genv + -> LoopFutureEnvInner kernel fenv2 genv + +loopFutureEnv + -- An already resolved signal + :: forall kernel fenv1 genv. + Idx fenv1 Signal + -> FutureEnv fenv1 genv + -> LoopFutureEnv kernel fenv1 genv +loopFutureEnv resolved = go SkipNone + where + go :: Skip fenv1 fenv' -> FutureEnv fenv' genv' -> LoopFutureEnv kernel fenv1 genv' + go _ FEnvEnd = LoopFutureEnv InputOutputRunit TupRunit TupRunit $ \_ _ _ _ -> + LoopFutureEnvInner buildReturn buildReturn FEnvEnd + go envSkip (FEnvFSkip env) = go (SkipSucc envSkip) env + go envSkip (FEnvGSkip env) + | LoopFutureEnv io initial localR body <- go envSkip env = + LoopFutureEnv io initial localR $ \skip input output local -> if + | LoopFutureEnvInner instr release env' <- body skip input output local -> + LoopFutureEnvInner instr release (FEnvGSkip env') + go envSkip (FEnvPush env future) + | LoopFutureEnv io1 initial1 localR1 body1 <- go envSkip env + , LoopFuture io2 initial2 localR2 body2 <- loopFuture @kernel resolved (weaken (skipWeakenIdx envSkip) future) = if + -- Optimize for common case where the future doesn't require further synchronisation + | InputOutputRunit <- io2 + , TupRunit <- localR2 -> + LoopFutureEnv io1 initial1 localR1 $ \skip input output local -> if + | LoopFutureEnvInner instr1 release1 env' <- body1 skip input output local + , LoopFutureInner instr2 release2 future' <- body2 skip TupRunit TupRunit TupRunit -> + LoopFutureEnvInner + (buildSpawn instr1 instr2) + (buildSpawn release1 release2) + (FEnvPush env' future') + | otherwise -> + LoopFutureEnv{ + loopEnvIO = InputOutputRpair io1 io2, + loopEnvInitial = TupRpair initial1 initial2, + loopEnvLocal = TupRpair localR1 localR2, + loopEnvBody = \skip input output local -> if + | TupRpair input1 input2 <- input + , TupRpair output1 output2 <- output + , TupRpair local1 local2 <- local + , LoopFutureEnvInner instr1 release1 env' <- body1 skip input1 output1 local1 + , LoopFutureInner instr2 release2 future' <- body2 skip input2 output2 local2 -> + LoopFutureEnvInner + (buildSpawn instr1 instr2) + (buildSpawn release1 release2) + (FEnvPush env' future') + | otherwise -> internalError "Input, output or local tuple impossible" + } + + +data LoopFuture kernel fenv1 t where + -- Captures existentials input, output and local + LoopFuture :: { + loopIO :: InputOutputR input output, + loopInitial :: BaseVars fenv1 input, + loopLocal :: TupR IsNewSignal local, + loopBody :: forall fenv2. Skip fenv2 fenv1 -> TupR (Idx fenv2) input -> TupR (Idx fenv2) output -> TupR (Idx fenv2) local -> LoopFutureInner kernel fenv2 t + } -> LoopFuture kernel fenv1 t + +data LoopFutureInner kernel fenv2 t where + LoopFutureInner + -- Executed at the start of an iteration + :: BuildSchedule kernel fenv2 + -> BuildSchedule kernel fenv2 + -> Future fenv2 t + -> LoopFutureInner kernel fenv2 t + +loopFuture + -- An already resolved signal + :: forall kernel fenv1 t. + Idx fenv1 Signal + -> Future fenv1 t + -> LoopFuture kernel fenv1 t +loopFuture resolved (FutureBuffer tp ref (Move readLockSignal) (Just (Move writeLockSignal))) = + -- A moved writable buffer + -- We must add two signals (and accompanying signal resolvers) to the state + -- to synchronize read and write access. + -- Note that the next iteration gets read access when the current one + -- finished writing, and the next iteration gets write access when this + -- iteratioon has finished reading. This is similar to the write-before-write + -- case in chainFuture. + -- + LoopFuture { + loopIO = InputOutputRsignal `InputOutputRpair` InputOutputRsignal, + loopInitial = TupRsingle (Var BaseRsignal $ fromMaybe resolved $ readLockSignal) + `TupRpair` TupRsingle (Var BaseRsignal $ fromMaybe resolved $ writeLockSignal), + loopLocal = TupRunit, + loopBody = \skip input output _ -> if + | signalR `TPair` signalW <- input + , resolverR `TPair` resolverW <- output + , ref' <- weakenEither (skipWeakenIdx skip) ref -> + LoopFutureInner + buildReturn + buildReturn + $ FutureBuffer tp ref' + -- Note that the resolvers for read and write are swapped, similar + -- to the write-before-write case in chainFuture (explaination is + -- there) + (Borrow (Just signalR) resolverW) + $ Just $ Borrow (Just signalW) resolverR + | otherwise -> internalError "input or output impossible" + } +loopFuture resolved (FutureBuffer tp ref (Lock readLockSignal readLockResolver) (Just (Lock writeLockSignal writeLockResolver))) = + -- A borrowed writable buffer + -- We must add two signals (and accompanying signal resolvers) to the state + -- to synchronize read and write access. Furthermore we need to declare two + -- local signals in the body of the loop, to mimic the signals of the loop + -- state, to properly release a buffer at the end of the loop. + -- Note that the next iteration gets read access when the current one + -- finished writing, and the next iteration gets write access when this + -- iteratioon has finished reading. This is similar to the write-before-write + -- case in chainFuture. + -- + LoopFuture { + loopIO = InputOutputRsignal `InputOutputRpair` InputOutputRsignal, + loopInitial = TupRsingle (Var BaseRsignal $ fromMaybe resolved $ readLockSignal) + `TupRpair` TupRsingle (Var BaseRsignal $ fromMaybe resolved $ writeLockSignal), + loopLocal = TPair (IsNewSignal "loopFuture (borrowed writable buffer, 1)") (IsNewSignal "loopFuture (borrowed writable buffer, 2)"), + loopBody = \skip input output local -> if + | signalR `TPair` signalW <- input + , resolverR `TPair` resolverW <- output + , TPair signalL1 resolverL1 `TupRpair` TPair signalL2 resolverL2 <- local + , ref' <- weakenEither (skipWeakenIdx skip) ref -> + LoopFutureInner + -- Note that the resolvers for read and write are swapped, similar + -- to the write-before-write case in chainFuture (explaination is + -- there) + (buildEffect (SignalAwait [signalL2]) $ buildEffect (SignalResolve [resolverR]) + $ buildEffect (SignalAwait [signalL1]) $ buildEffect (SignalResolve [resolverW]) buildReturn) + (buildEffect (SignalAwait [signalL2]) + $ buildEffect (SignalResolve $ catMaybes [weaken (skipWeakenIdx skip) <$> readLockResolver]) + $ buildEffect (SignalAwait [signalL1]) + $ buildEffect (SignalResolve $ catMaybes [weaken (skipWeakenIdx skip) <$> writeLockResolver]) buildReturn) + $ FutureBuffer tp ref' + (Borrow (Just signalR) resolverL1) + $ Just $ Borrow (Just signalW) resolverL2 + | otherwise -> internalError "Input, output or local tuple impossible" + } +loopFuture resolved (FutureBuffer tp ref (Borrow readLockSignal readLockResolver) Nothing) = + -- A borrowed read-only buffer. + -- We need one signal in the state of the loop, to denote that the previous + -- iterations have finished working with the array. The accompanying signal + -- resolver is resolved when this signal is finished and the current + -- iteration is done working with the array. When the loop has finished and + -- those signals are resolved, we will also release the entire buffer by + -- resolving 'releaseRead'. + LoopFuture { + loopIO = InputOutputRsignal, + loopInitial = TupRsingle (Var BaseRsignal $ fromMaybe resolved $ readLockSignal), + loopLocal = TupRsingle (IsNewSignal "loopFuture (borrowed read-only buffer)"), + loopBody = \skip input output local -> if + | TupRsingle signalR <- input + , TupRsingle resolverR <- output + , TPair signalL1 resolverL1 <- local + , ref' <- weakenEither (skipWeakenIdx skip) ref + , future <- FutureBuffer tp ref' + (Lock (weaken (skipWeakenIdx skip) <$> readLockSignal) (Just resolverL1)) + Nothing -> + LoopFutureInner + (buildEffect (SignalAwait [signalR, signalL1]) $ buildEffect (SignalResolve [resolverR]) buildReturn) + (buildEffect (SignalAwait $ [signalR, signalL1]) -- Wait on signalL1 if left + $ buildEffect (SignalResolve [weaken (skipWeakenIdx skip) $ readLockResolver]) buildReturn) + future + | otherwise -> internalError "input, output or local impossible" + } +-- A buffer with Move read lock or a scalar, no further synchronisation needed +loopFuture _ future = + LoopFuture { + loopIO = InputOutputRunit, + loopInitial = TupRunit, + loopLocal = TupRunit, + loopBody = \skip _ _ _ -> LoopFutureInner buildReturn buildReturn $ weaken (skipWeakenIdx skip) future + } + +subFutureEnvironment :: forall fenv genv kernel. FutureEnv fenv genv -> SyncEnv genv -> (FutureEnv fenv genv, BuildSchedule kernel fenv) +subFutureEnvironment = subFutureEnvironment' SkipNone + +subFutureEnvironment' :: forall fenv fenv' genv kernel. Skip fenv' fenv -> FutureEnv fenv genv -> SyncEnv genv -> (FutureEnv fenv' genv, BuildSchedule kernel fenv') +subFutureEnvironment' _ FEnvEnd PEnd = (FEnvEnd, buildReturn) +subFutureEnvironment' skip (FEnvGSkip fenv) (PNone senv) = (FEnvGSkip fenv', actions) + where + (fenv', actions) = subFutureEnvironment' skip fenv senv +subFutureEnvironment' skip (FEnvGSkip fenv) PEnd = (FEnvGSkip fenv', actions) + where + (fenv', actions) = subFutureEnvironment' skip fenv PEnd +subFutureEnvironment' skip FEnvEnd (PNone senv) = (FEnvGSkip fenv', actions) + where + (fenv', actions) = subFutureEnvironment' skip FEnvEnd senv +subFutureEnvironment' skip (FEnvFSkip fenv) senv = subFutureEnvironment' (SkipSucc skip) fenv senv +subFutureEnvironment' skip (FEnvPush fenv f@FutureScalar{}) senv = (FEnvPush fenv' $ weaken (skipWeakenIdx skip) f, actions) + where + (fenv', actions) = subFutureEnvironment' skip fenv $ partialEnvTail senv +subFutureEnvironment' skip (FEnvPush fenv f@(FutureBuffer tp ref read write)) (PPush senv sync) = (FEnvPush fenv' $ weaken (skipWeakenIdx skip) f', actions') + where + (fenv', actions) = subFutureEnvironment' skip fenv senv + + (f', actions') = case (write, sync) of + (Nothing, SyncRead) -> (f, actions) -- No need to change + (Just _, SyncWrite) -> (f, actions) -- No need to change + (Nothing, SyncWrite) -> -- Illegal input + internalError "Got a FutureBuffer without write capabilities, but the SyncEnv asks for write permissions" + (Just (Borrow ws wr), SyncRead) -> -- Write capability not used + ( FutureBuffer tp ref read Nothing + -- Resolve the write resolver after taking both the read and write signal + , buildSpawn + (buildEffect (SignalAwait $ fmap (weaken (skipWeakenIdx skip)) $ catMaybes [lockSignal read, ws]) + $ buildEffect (SignalResolve [weaken (skipWeakenIdx skip) $ wr]) buildReturn) + actions + ) + (Just (Move _), SyncRead) -> + ( FutureBuffer tp ref read Nothing + , actions' + ) +subFutureEnvironment' skip (FEnvPush fenv (FutureBuffer _ _ read write)) (PNone senv) = (FEnvGSkip fenv', actions') + where + (fenv', actions) = subFutureEnvironment' skip fenv senv + + actions' = case (read, write) of + (Borrow rs rr, Just (Borrow ws wr)) -> + buildSpawn + (buildEffect (SignalAwait $ fmap (weaken $ skipWeakenIdx skip) $ catMaybes [rs]) + $ buildEffect (SignalResolve [weaken (skipWeakenIdx skip) rr]) + $ buildEffect (SignalAwait $ fmap (weaken $ skipWeakenIdx skip) $ catMaybes [ws]) + $ buildEffect (SignalResolve [weaken (skipWeakenIdx skip) wr]) + buildReturn) + actions + (Borrow _ _, Just (Move _)) -> internalError "Illegal FutureBuffer with borrow read access and move write access" + (Move rs, Just (Borrow ws wr)) -> + buildSpawn + (buildEffect (SignalAwait $ fmap (weaken $ skipWeakenIdx skip) $ catMaybes [rs, ws]) + $ buildEffect (SignalResolve [weaken (skipWeakenIdx skip) wr]) + buildReturn) + actions + (Move _, Just (Move _)) -> actions + (Borrow rs rr, Nothing) -> + buildSpawn + (buildEffect (SignalAwait $ fmap (weaken $ skipWeakenIdx skip) $ catMaybes [rs]) + $ buildEffect (SignalResolve [weaken (skipWeakenIdx skip) rr]) + buildReturn) + actions + (Move _, Nothing) -> actions +subFutureEnvironment' _ (FEnvGSkip _) (PPush _ _) = internalError "Variable missing in FutureEnv, but expected based on SyncEnv" +subFutureEnvironment' _ FEnvEnd (PPush _ _) = internalError "Variable missing in FutureEnv, but expected based on SyncEnv" +subFutureEnvironment' skip fenv@FEnvPush{} PEnd = subFutureEnvironment' skip fenv (PNone PEnd) + +-- When binding new variables, the LeftHandSide and Uniquenesses might not correspond to the values in SyncEnv. +-- This is a simplified variant of subFutureEnvironment, where we only visit +-- the part of the FutureEnv corresponding to the LeftHandSide, and assume that +-- those bindings only have Move locks (no Borrows). The latter is the case +-- because of the construction in declareDestinations. +restrictEnvForLHS :: GLeftHandSide t genv genv' -> FutureEnv fenv genv' -> SyncEnv genv' -> FutureEnv fenv genv' +restrictEnvForLHS lhs = go (lhsSkip lhs) + where + -- Note: Since Skip appends the environment in + go :: Skip genv' genv -> FutureEnv fenv genv' -> SyncEnv genv' -> FutureEnv fenv genv' + go SkipNone env _ = env + go skip (FEnvFSkip env) senv = FEnvFSkip $ go skip env senv + go _ FEnvEnd _ = FEnvEnd + go (SkipSucc' skip) (FEnvGSkip env) senv = FEnvGSkip $ go skip env $ partialEnvTail senv + go (SkipSucc' skip) (FEnvPush env future@FutureScalar{}) senv = go skip env (partialEnvTail senv) `FEnvPush` future + -- Unused buffer. The locks should be Move, so we don't need to release anything. + go (SkipSucc' skip) (FEnvPush env _) (PNone senv) = FEnvGSkip $ go skip env senv + go (SkipSucc' skip) (FEnvPush env future) (PPush senv sync) = go skip env senv `FEnvPush` future' + where + future' + | FutureBuffer tp ref read (Just _) <- future + , SyncRead <- sync + -- Downgrade unique buffer to shared buffer. + -- Note that 'write' will be Move, so we don't need to release anything. + = FutureBuffer tp ref read Nothing + | otherwise + = future + go skip@SkipSucc'{} env PEnd = go skip env (PNone PEnd) + +type MaybeSignal fenv = Maybe (Idx fenv Signal) +type MaybeSignalResolver fenv = Maybe (Idx fenv SignalResolver) + +buildAwait :: [MaybeSignal fenv] -> BuildSchedule kernel fenv -> BuildSchedule kernel fenv +buildAwait signals = case catMaybes signals of + [] -> id + signals' -> buildEffect (SignalAwait signals') + +-- Asserts that the FutureEnv corresponds with the SyncEnv, i.e. that it +-- provides the expected buffers, with expected (read/write) capabilities +assertFutureEnv :: SyncEnv genv -> FutureEnv fenv genv -> () +assertFutureEnv senv (FEnvFSkip env) = assertFutureEnv senv env +assertFutureEnv (PPush senv SyncWrite) (FEnvPush env (FutureBuffer _ _ _ Just{})) = assertFutureEnv senv env +assertFutureEnv (PPush _ SyncWrite) (FEnvPush _ FutureBuffer{}) = internalError "SyncWrite, read-only buffer mismatch" +assertFutureEnv (PPush senv SyncRead) (FEnvPush env (FutureBuffer _ _ _ Nothing)) = assertFutureEnv senv env +assertFutureEnv (PPush _ SyncRead) (FEnvPush _ FutureBuffer{}) = internalError "SyncRead, writeable buffer mismatch" +assertFutureEnv (PPush _ _) (FEnvPush _ FutureScalar{}) = internalError "Buffer impossible" +assertFutureEnv (PPush _ _) (FEnvGSkip _) = internalError "Buffer missing in FutureEnv" +assertFutureEnv (PPush _ _) FEnvEnd = internalError "Buffer missing in FutureEnv" +assertFutureEnv (PNone _) (FEnvPush _ FutureBuffer{}) = internalError "Redundant buffer in FutureEnv" +assertFutureEnv PEnd (FEnvPush _ FutureBuffer{}) = internalError "Redundant buffer in FutureEnv" +assertFutureEnv (PNone senv) (FEnvPush env FutureScalar{}) = assertFutureEnv senv env +assertFutureEnv (PNone senv) (FEnvGSkip env) = assertFutureEnv senv env +assertFutureEnv (PNone senv) FEnvEnd = assertFutureEnv senv FEnvEnd +assertFutureEnv PEnd FEnvEnd = () +assertFutureEnv PEnd (FEnvGSkip env) = assertFutureEnv PEnd env +assertFutureEnv PEnd (FEnvPush env FutureScalar{}) = assertFutureEnv PEnd env diff --git a/src/Data/Array/Accelerate/Trafo/Schedule/Uniform/LiveVars.hs b/src/Data/Array/Accelerate/Trafo/Schedule/Uniform/LiveVars.hs new file mode 100644 index 000000000..43806a5b4 --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Schedule/Uniform/LiveVars.hs @@ -0,0 +1,218 @@ +{-# LANGUAGE MultiWayIf #-} +{-# OPTIONS_HADDOCK hide #-} +-- | +-- Module : Data.Array.Accelerate.Trafo.Schedule.Uniform.LiveVars +-- Copyright : [2012..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.Trafo.Schedule.Uniform.LiveVars ( + stronglyLiveVariablesFun +) where + +import Data.Array.Accelerate.AST.Environment +import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.IdxSet ( IdxSet(..) ) +import qualified Data.Array.Accelerate.AST.IdxSet as IdxSet +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.AST.Schedule.Uniform +import Data.Array.Accelerate.Error +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Trafo.Exp.Substitution +import Data.Array.Accelerate.Trafo.LiveVars +import Data.Array.Accelerate.Trafo.Substitution +import Data.Array.Accelerate.Trafo.Operation.LiveVars (reEnvArrayDescriptors, reEnvArrayInstr) + +import Data.Maybe + +stronglyLiveVariablesFun :: UniformScheduleFun kernel () t -> UniformScheduleFun kernel () t +stronglyLiveVariablesFun schedule = schedule' ReEnvEnd + where + LVAnalysisFun _ schedule' = stronglyLiveVariablesFun' emptyLivenessEnv schedule + +stronglyLiveVariablesFun' :: LivenessEnv env -> UniformScheduleFun kernel env t -> LVAnalysisFun (UniformScheduleFun kernel) env t +stronglyLiveVariablesFun' liveness = \case + Sbody body + | LVAnalysis' liveness2 body' <- stronglyLiveVariables' liveness body + -> LVAnalysisFun liveness2 $ \re -> Sbody $ body' re + Slam lhs f + | liveness1 <- lEnvPushLHS lhs liveness + , liveness2 <- + setIdxSetLive + (IdxSet.fromList + $ mapMaybe (\(Exists (Var tp idx)) -> if isOutput tp then Just $ Exists idx else Nothing) + $ lhsVars lhs) + liveness1 + , LVAnalysisFun liveness3 f' <- stronglyLiveVariablesFun' liveness2 f + , (lhs', liveness4) <- lEnvStrengthenLHS lhs liveness3 + -> LVAnalysisFun + liveness4 + $ \re -> if -- A one-way "multi way if" to pattern match on a GADT + | BindLiveness lhs'' re' <- bind lhs' re + -> Slam lhs'' $ f' re' + +stronglyLiveVariables' :: LivenessEnv env -> UniformSchedule kernel env -> LVAnalysis' (UniformSchedule kernel) env +stronglyLiveVariables' liveness = \case + Return -> LVAnalysis' liveness $ const Return + Alet lhs binding body -> + let + liveness1 = analyseBinding (weakenWithLHS lhs) (lhsIndices lhs) binding $ lEnvPushLHS lhs liveness + LVAnalysis' liveness2 body' = stronglyLiveVariables' liveness1 body + (lhs', liveness3) = lEnvStrengthenLHS lhs liveness2 + in + LVAnalysis' + liveness3 + $ \re -> if -- A one-way "multi way if" to pattern match on a GADT + | BindLiveness lhs'' re' <- bind lhs' re -> case lhs'' of + LeftHandSideWildcard _ -> body' re' -- Entire binding wasn't used + _ -> Alet lhs'' (reEnvBinding re binding) (body' re') + Effect (SignalAwait signals) (Effect (SignalResolve resolvers) Return) -> + let + liveness1 = addLiveImplications (IdxSet.fromList $ map Exists resolvers) (IdxSet.fromList $ map Exists signals) liveness + in + LVAnalysis' + liveness1 + $ \re -> await (reEnvIndices' re signals) $ resolve (reEnvIndices' re resolvers) Return + Effect effect next -> + let + liveness1 = analyseEffect effect liveness + LVAnalysis' liveness2 next' = stronglyLiveVariables' liveness1 next + in + LVAnalysis' + liveness2 + $ \re -> reEnvEffect re effect $ next' re + Acond condition true false next -> + let + liveness1 = setLive (varIdx condition) liveness + LVAnalysis' liveness2 true' = stronglyLiveVariables' liveness1 true + LVAnalysis' liveness3 false' = stronglyLiveVariables' liveness2 false + LVAnalysis' liveness4 next' = stronglyLiveVariables' liveness3 next + in + LVAnalysis' + liveness4 + $ \re -> Acond + (expectJust $ reEnvVar re condition) + (true' re) + (false' re) + (next' re) + Awhile io step initial next -> + let + -- TODO: We could track which parts of the state are used + liveness1 = setVarsLive initial liveness + LVAnalysisFun liveness2 step' = stronglyLiveVariablesFun' liveness1 step + LVAnalysis' liveness3 next' = stronglyLiveVariables' liveness2 next + in + LVAnalysis' + liveness3 + $ \re -> Awhile io + (step' re) + (expectJust $ reEnvVars re initial) + (next' re) + AwhileSeq io step initial next -> + let + -- TODO: We could track which parts of the state are used + liveness1 = setVarsLive initial liveness + LVAnalysisFun liveness2 step' = stronglyLiveVariablesFun' liveness1 step + LVAnalysis' liveness3 next' = stronglyLiveVariables' liveness2 next + in + LVAnalysis' + liveness3 + $ \re -> AwhileSeq io + (step' re) + (expectJust $ reEnvVars re initial) + (next' re) + Spawn a b -> + let + LVAnalysis' liveness1 a' = stronglyLiveVariables' liveness a + LVAnalysis' liveness2 b' = stronglyLiveVariables' liveness1 b + in + LVAnalysis' + liveness2 + $ \re -> Spawn + (a' re) + (b' re) + +analyseBinding :: env :> env' -> IdxSet env' -> Binding env t -> LivenessEnv env' -> LivenessEnv env' +analyseBinding k lhs binding liveness = case binding of + Compute expr -> + let + -- If the LHS of the binding is live, then all free variables of this + -- expression are live as well. + free = map (\(Exists (Var _ idx)) -> Exists $ k >:> idx) $ expGroundVars expr + in + addLiveImplications lhs (IdxSet.fromList free) liveness + NewSignal _ + | IdxSet (_ `PPush` _ `PPush` _) <- lhs -> + -- If the signal is live, then the resolver is live as well. + addLiveImplies + (SuccIdx ZeroIdx) + (ZeroIdx) + liveness + | otherwise -> liveness + NewRef _ + | IdxSet (_ `PPush` _ `PPush` _) <- lhs -> + -- If the Ref is live, then the OutputRef is live as well. + addLiveImplies + (SuccIdx ZeroIdx) + (ZeroIdx) + liveness + | otherwise -> liveness + Alloc _ _ sh -> + -- If this buffer is live, then the shape variables are live as well. + addLiveImplications lhs (IdxSet.fromVars $ mapTupR (weaken k) sh) liveness + Use _ _ _ -> liveness + Unit (Var _ idx) -> -- If the lhs is live, then the argument of Unit is live as well. + addLiveImplications lhs (IdxSet.singleton $ k >:> idx) liveness + RefRead (Var _ idx) -> -- If the lhs is live, then the Ref is live as well. + addLiveImplications lhs (IdxSet.singleton $ k >:> idx) liveness + +reEnvBinding :: ReEnv env subenv -> Binding env t -> Binding subenv t +reEnvBinding re = \case + Compute expr -> Compute $ mapArrayInstr (reEnvArrayInstr re) expr + NewSignal name -> NewSignal name + NewRef tp -> NewRef tp + Alloc shr tp sh -> Alloc shr tp $ expectJust $ reEnvVars re sh + Use tp sh buffer -> Use tp sh buffer + Unit var -> Unit $ expectJust $ reEnvVar re var + RefRead var -> RefRead $ expectJust $ reEnvVar re var + +analyseEffect :: Effect kernel env -> LivenessEnv env -> LivenessEnv env +analyseEffect (Exec _ _ args) liveness = setIdxSetLive (IdxSet.fromList $ argsIndices args) liveness +analyseEffect (SignalAwait signals) liveness = setIdxSetLive (IdxSet.fromList $ map Exists signals) liveness +analyseEffect (SignalResolve _) liveness = liveness +analyseEffect (RefWrite ref value) liveness = addLiveImplies (varIdx ref) (varIdx value) liveness +analyseEffect (Atrace _ t) liveness = setIdxSetLive (arrayDescriptorsIdxSet t) liveness +analyseEffect (Aassert _ cond) liveness = setIdxSetLive (IdxSet.fromList free) liveness + where free = map (\(Exists (Var _ idx)) -> Exists idx) $ expGroundVars cond + +reEnvEffect :: ReEnv env subenv -> Effect kernel env -> UniformSchedule kernel subenv -> UniformSchedule kernel subenv +reEnvEffect re = \case + Exec md kernel args -> Effect $ Exec md kernel $ mapArgs (reEnvArg re) args + SignalAwait signals -> Effect $ SignalAwait $ reEnvIndices' re signals + SignalResolve resolvers -> resolve $ reEnvIndices' re resolvers + RefWrite ref value -> case (reEnvVar re ref, reEnvVar re value) of + (Just ref', Just value') -> Effect $ RefWrite ref' value' + (Nothing , _ ) -> id + (Just _ , Nothing ) -> internalError "Substitution in live variable analysis failed. A reference was live, but the value written to it was dead." + Aassert msg cond -> Effect $ Aassert msg $ mapArrayInstr (reEnvArrayInstr re) cond + Atrace msg t -> Effect $ Atrace msg $ reEnvArrayDescriptors re t + +reEnvArg :: ReEnv env subenv -> SArg env t -> SArg subenv t +reEnvArg re (SArgScalar var) = SArgScalar $ expectJust $ reEnvVar re var +reEnvArg re (SArgBuffer m var) = SArgBuffer m $ expectJust $ reEnvVar re var + +argsIndices :: SArgs env t -> [Exists (Idx env)] +argsIndices = map (\(Exists a) -> argIndex a) . argsToList + +argIndex :: SArg env t -> Exists (Idx env) +argIndex (SArgScalar var) = Exists $ varIdx var +argIndex (SArgBuffer _ var) = Exists $ varIdx var + +isOutput :: BaseR t -> Bool +isOutput BaseRrefWrite{} = True +isOutput BaseRsignalResolver = True +isOutput _ = False diff --git a/src/Data/Array/Accelerate/Trafo/Schedule/Uniform/Simplify.hs b/src/Data/Array/Accelerate/Trafo/Schedule/Uniform/Simplify.hs new file mode 100644 index 000000000..4f4c855da --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Schedule/Uniform/Simplify.hs @@ -0,0 +1,1043 @@ +{-# LANGUAGE MultiWayIf #-} +{-# LANGUAGE TypeFamilies #-} +{-# OPTIONS_HADDOCK hide #-} +-- | +-- Module : Data.Array.Accelerate.Trafo.Schedule.Uniform.Simplify +-- Copyright : [2012..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.Trafo.Schedule.Uniform.Simplify ( + -- Construction DSL + BuildSchedule, BuildScheduleFun, funConstruct, + buildAcond, buildAwhile, buildEffect, buildSpawn, + buildLet, buildReturn, buildSeq, + buildFunBody, buildFunLam, + BuildEnv(BEmpty), + + lhsSignal, lhsRef, + + -- Optimize + simplify +) where + +import Data.Array.Accelerate.AST.Environment +import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.AST.Schedule.Uniform +import Data.Array.Accelerate.Analysis.Match +import Data.Array.Accelerate.Error +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Trafo.Exp.Substitution +import Data.Array.Accelerate.Trafo.Substitution hiding ( weakenArrayInstr ) +import Data.Array.Accelerate.Trafo.Schedule.Uniform.Substitution () +import Data.Maybe +import Data.List + ( isSubsequenceOf, + sort ) + +-- TODO: Move fork combining and/or InfoEnv into this utility? +data BuildSchedule kernel env = + BuildSchedule{ + -- Sorted, without duplicates + directlyAwaits :: [Idx env Signal], + -- The SignalResolvers that this term resolves at the end of its execution. + -- 'End' here is the place where the Continuation would be placed + -- (as 'end' is otherwise ambiguous if the term has Spawns). + -- Sorted, without duplicates + finallyResolves :: [Idx env SignalResolver], + trivial :: Bool, + -- How deep Awhile(Seq)s are nested within this term + awhileHeight :: Int, + -- Constructs a schedule, but doesn't wait on the directlyAwaits signals. + -- constructFull adds that. + construct + :: forall env'. + env :> env' + -> BuildEnv env' + -> Postponed kernel env' + -> Continuation kernel env' + -> UniformSchedule kernel env' + } + +instance Sink' (BuildSchedule kernel) where + weaken' k schedule = + BuildSchedule{ + directlyAwaits = sort $ map (weaken k) $ directlyAwaits schedule, + finallyResolves = sort $ map (weaken k) $ finallyResolves schedule, + trivial = trivial schedule, + awhileHeight = awhileHeight schedule, + construct = \k' -> construct schedule (k' .> k) + } + +data BuildScheduleFun kernel env t where + BuildBody + :: BuildSchedule kernel env + -> BuildScheduleFun kernel env () + + BuildLam + :: BLeftHandSide t env env' + -> BuildScheduleFun kernel env' f + -> BuildScheduleFun kernel env (t -> f) + +funConstruct :: BuildScheduleFun kernel env t -> env :> env' -> BuildEnv env' -> UniformScheduleFun kernel env' t +funConstruct (BuildBody body) k env = Sbody $ constructFull body k env nothingPostponed ContinuationEnd +funConstruct (BuildLam lhs f) k env + | Exists lhs' <- rebuildLHS lhs + = Slam lhs' $ funConstruct f (sinkWithLHS lhs lhs' k) $ buildEnvSkip lhs' env + +funAwhileHeight :: BuildScheduleFun kernel env t -> Int +funAwhileHeight (BuildBody b) = awhileHeight b +funAwhileHeight (BuildLam _ f) = funAwhileHeight f + +funDirectlyAwaits :: BuildScheduleFun kernel env t -> [Idx env Signal] +funDirectlyAwaits (BuildBody b) = directlyAwaits b +funDirectlyAwaits (BuildLam lhs f) = mapMaybe (strengthenWithLHS lhs) $ funDirectlyAwaits f + +instance Sink (BuildScheduleFun kernel) where + weaken k (BuildBody body) = BuildBody $ weaken' k body + weaken k (BuildLam lhs f) + | Exists lhs' <- rebuildLHS lhs + = BuildLam lhs' $ weaken (sinkWithLHS lhs lhs' k) f + +data BuildSpawn kernel env = + BuildSpawn{ + -- The signals that this term resolves at the end. + -- Corresponds with finallyResolves, but the SignalResolvers are converted to Signals. + spawnFinallyResolves :: [Idx env Signal], + spawnTerm :: BuildSchedule kernel env + } + +instance Sink' (BuildSpawn kernel) where + weaken' k (BuildSpawn signals schedule) = BuildSpawn (map (weaken k) signals) (weaken' k schedule) + +data Postponed kernel env = Postponed + [BuildSpawn kernel env] -- The terms that should be spawned before doing serious work, in reverse order + [Idx env SignalResolver] -- The SignalResolvers that should be resolved before doing serious work + +nothingPostponed :: Postponed kernel env +nothingPostponed = Postponed [] [] + +instance Sink' (Postponed kernel) where + weaken' k (Postponed spawns resolvers) = Postponed + (map (weaken' k) spawns) + (map (weaken k) resolvers) + +-- Constructs a schedule, and waits on the directlyAwaits signals. +constructFull + :: BuildSchedule kernel env + -> env :> env' + -> BuildEnv env' + -> Postponed kernel env' + -> Continuation kernel env' + -> UniformSchedule kernel env' +constructFull schedule k env postponed cont + | null $ directlyAwaits schedule = construct schedule k env postponed cont + | signals' <- + -- Don't wait on already resolved signals + sortedDedup + $ sort + $ filter (\idx -> not (isResolved idx env)) + $ map (weaken k) + $ directlyAwaits schedule + , env' <- markResolved signals' env = + if null signals' then + construct schedule k env' postponed cont + else + case findDependingSpawn postponed signals' of + Nothing -> + placePostponed postponed env + $ Effect (SignalAwait signals') + $ construct schedule k env' nothingPostponed cont + Just (BuildSpawn _ prepend, postponed') -> + constructFull prepend weakenId env postponed' $ ContinuationDo k schedule weakenId cont + +placePostponed :: Postponed kernel env -> BuildEnv env -> UniformSchedule kernel env -> UniformSchedule kernel env +placePostponed (Postponed spawns resolvers) env next = + foldl + (\other (BuildSpawn _ term) -> spawnAndRotate (constructFull term weakenId env nothingPostponed ContinuationEnd) other) + next' + spawns + where + next' + | [] <- resolvers = next + | otherwise = Effect (SignalResolve resolvers) next + +spawnAndRotate :: UniformSchedule kernel env -> UniformSchedule kernel env -> UniformSchedule kernel env +spawnAndRotate (Spawn a b) c = Spawn a $ spawnAndRotate b c +spawnAndRotate a Return = a +spawnAndRotate a (Effect eff@(SignalResolve _) Return) = Effect eff $ a +spawnAndRotate a b = Spawn a b + +resolveSignalsInPostponed :: [Idx env Signal] -> [Idx env SignalResolver] -> Postponed kernel env -> Postponed kernel env +resolveSignalsInPostponed signals newResolvers (Postponed spawns resolvers) = Postponed + (map (\(BuildSpawn r s) -> BuildSpawn r s{ directlyAwaits = filter (`notElem` signals) $ directlyAwaits s }) spawns) + (newResolvers ++ resolvers) + +spawnPostponed :: forall kernel env. BuildSpawn kernel env -> Postponed kernel env -> Postponed kernel env +spawnPostponed spawn (Postponed spawns resolvers) + | Just spawns' <- tryCombine spawns = Postponed spawns' resolvers + | otherwise = Postponed (spawn : spawns) resolvers + where + -- Tries to combine 'spawn' with a BuildSpawn in spawns. + tryCombine :: [BuildSpawn kernel env] -> Maybe [BuildSpawn kernel env] + tryCombine [] = Nothing -- It wasn't possible to combine it. + tryCombine (x:xs) + -- If 'spawn' waits on a result of 'x' + | shouldCombine x spawn + = Just $ combine x spawn : xs + + -- If 'x' waits on a result of 'spawn' + | shouldCombine spawn x + = Just $ combine spawn x : xs + + | otherwise + = (x:) <$> tryCombine xs + + shouldCombine :: BuildSpawn kernel env -> BuildSpawn kernel env -> Bool + shouldCombine before after + -- If 'after' waits on a result of 'before' + = not (null $ spawnFinallyResolves before `sortedIntersection` directlyAwaits (spawnTerm after)) + -- If 'before' is trivial and does not wait on other signals than 'after' + || trivial (spawnTerm before) && directlyAwaits (spawnTerm before) `isSubsequenceOf` directlyAwaits (spawnTerm after) + + combine :: BuildSpawn kernel env -> BuildSpawn kernel env -> BuildSpawn kernel env + combine first second = BuildSpawn + (spawnFinallyResolves first) + (buildSeq (spawnTerm first) (buildSpawnReturn $ spawnTerm second)) + +-- Finds a spawned term, on which the next term (given by the directly-awaits +-- list) directly depends on. +findDependingSpawn :: forall kernel env. Postponed kernel env -> [Idx env Signal] -> Maybe (BuildSpawn kernel env, Postponed kernel env) +findDependingSpawn (Postponed spawns resolvers) nextDirectlyAwaits = case go spawns of + Just (y, ys) -> Just (y, Postponed ys resolvers) + Nothing -> Nothing + where + go :: [BuildSpawn kernel env] -> Maybe (BuildSpawn kernel env, [BuildSpawn kernel env]) + go (x:xs) + | not $ null $ spawnFinallyResolves x `sortedIntersection` nextDirectlyAwaits + = Just (x, xs) + | otherwise = case go xs of + Nothing -> Nothing + Just (y, ys) -> Just (y, x:ys) + go [] = Nothing + +data BuildEnv env where + BEmpty :: BuildEnv () + BPush :: BuildEnv env -> BuildEnvInfo env t -> BuildEnv (env, t) + +data BuildEnvInfo env t where + -- No information available. + INone + :: BuildEnvInfo env t + + -- This SignalResolver resolves the Signal at the next index in the environment. + IResolvesNext + :: BuildEnvInfo (env, Signal) SignalResolver + + -- This Signal is resolved. + IResolved + :: BuildEnvInfo env Signal + + -- This OutputRef writes to the Ref at the next index in the environment. + IWritesNext + :: BuildEnvInfo (env, Ref t) (OutputRef t) + + -- This Ref contains the value of the specified variable. + IRefContains + :: Idx env t + -> BuildEnvInfo env (Ref t) + + -- This variable has the value of the given Refs. + -- Only for Buffer and Scalar variables. + IValue + :: [Idx env (Ref t)] + -> BuildEnvInfo env t + +data BuildEnvPrj t where + BuildEnvPrj :: BuildEnvInfo env t -> BuildEnvPrj t + +buildEnvPrj :: Idx env t -> BuildEnv env -> BuildEnvPrj t +buildEnvPrj ZeroIdx (BPush _ v) = BuildEnvPrj v +buildEnvPrj (SuccIdx idx) (BPush env _) = buildEnvPrj idx env + +data BuildEnvPrj' env t where + BuildEnvPrj' :: Skip env env' -> BuildEnvInfo env' t -> BuildEnvPrj' env t + +buildEnvPrj' :: Idx env t -> BuildEnv env -> BuildEnvPrj' env t +buildEnvPrj' = go SkipNone + where + go :: Skip env env' -> Idx env' t -> BuildEnv env' -> BuildEnvPrj' env t + go skip ZeroIdx (BPush _ v) = BuildEnvPrj' (SkipSucc skip) v + go skip (SuccIdx idx) (BPush env _) = go (SkipSucc skip) idx env + +findSignal :: Idx env SignalResolver -> BuildEnv env -> Maybe (Idx env Signal) +findSignal idx env = case buildEnvPrj' idx env of + BuildEnvPrj' skip IResolvesNext -> Just $ weaken (skipWeakenIdx skip) ZeroIdx + _ -> Nothing + +findRef :: Idx env (OutputRef t) -> BuildEnv env -> Maybe (Idx env (Ref t)) +findRef idx env = case buildEnvPrj' idx env of + BuildEnvPrj' skip IWritesNext -> Just $ weaken (skipWeakenIdx skip) ZeroIdx + _ -> Nothing + +-- Given a sorted list of unique signals, marks those signals as resolved in the BuildEnv. +markResolved :: [Idx env Signal] -> BuildEnv env -> BuildEnv env +markResolved [] env = env +markResolved signals (BPush env info) + | ZeroIdx : signals' <- signals + = BPush (markResolved (map unSucc signals') env) IResolved + | otherwise + = BPush (markResolved (map unSucc signals) env) info + where + unSucc :: Idx (env, t) s -> Idx env s + unSucc ZeroIdx = internalError "markResolved: input was not sorted or contains duplicates" + unSucc (SuccIdx idx) = idx +markResolved (s:_) BEmpty = case s of {} + +isResolved :: Idx env Signal -> BuildEnv env -> Bool +isResolved signal env + | BuildEnvPrj IResolved <- buildEnvPrj signal env = True + | otherwise = False + +markRefValue :: Idx env (Ref t) -> Idx env t -> BuildEnv env -> BuildEnv env +markRefValue (SuccIdx refIdx) (SuccIdx valueIdx) (BPush env info) = BPush (markRefValue refIdx valueIdx env) info +markRefValue ZeroIdx (SuccIdx valueIdx) (BPush env _) = BPush env (IRefContains valueIdx) +markRefValue (SuccIdx refIdx) ZeroIdx (BPush env info) = BPush env $ case info of + INone + -> IValue [refIdx] + IValue refs + -> IValue (refIdx : refs) + _ -> info + +-- Finds the value of a reference, if available +findRefValue :: Idx env (Ref t) -> BuildEnv env -> Maybe (Idx env t) +findRefValue = go SkipNone + where + go :: Skip env env' -> Idx env' (Ref t) -> BuildEnv env' -> Maybe (Idx env t) + go skip ZeroIdx (BPush _ info) = case info of + IRefContains value -> Just $ weaken (skipWeakenIdx $ SkipSucc skip) value + _ -> Nothing + go skip (SuccIdx refIdx) (BPush env info) + | IValue refs <- info + , Refl : _ <- mapMaybe (matchIdx refIdx) refs = Just $ weaken (skipWeakenIdx skip) ZeroIdx + | otherwise = go (SkipSucc skip) refIdx env + +data Continuation kernel env where + ContinuationEnd + :: Continuation kernel env + + ContinuationDo + :: env1 :> env + -> BuildSchedule kernel env1 + -> env2 :> env + -> Continuation kernel env2 + -> Continuation kernel env + +instance Sink' (Continuation kernel) where + weaken' _ ContinuationEnd = ContinuationEnd + weaken' k1 (ContinuationDo k2 b k3 c) = ContinuationDo (k1 .> k2) b (k1 .> k3) c + +buildReturn :: BuildSchedule kernel env +buildReturn = BuildSchedule{ + directlyAwaits = [], + finallyResolves = [], + trivial = True, + awhileHeight = 0, + construct = \_ env postponed -> \case + ContinuationEnd -> placePostponed postponed env Return + ContinuationDo k2 build k3 cont -> constructFull build k2 env postponed $ weaken' k3 cont + } + +buildLet + :: forall kernel t env1 env2. + BLeftHandSide t env1 env2 + -> Binding env1 t + -> BuildSchedule kernel env2 + -> BuildSchedule kernel env1 +buildLet lhs binding body + | trivialBinding binding = + BuildSchedule{ + directlyAwaits = map (fromMaybe (internalError "Illegal schedule: deadlock") . strengthenWithLHS lhs) $ directlyAwaits body, + finallyResolves = mapMaybe (strengthenWithLHS lhs) $ finallyResolves body, + trivial = trivial body, + awhileHeight = awhileHeight body, + construct = constructLet False + } + | otherwise = + BuildSchedule{ + directlyAwaits = [], + finallyResolves = mapMaybe (strengthenWithLHS lhs) $ finallyResolves body, + trivial = False, + awhileHeight = awhileHeight body, + construct = constructLet True + } + where + constructLet + :: Bool + -> env1 :> env1' + -> BuildEnv env1' + -> Postponed kernel env1' + -> Continuation kernel env1' + -> UniformSchedule kernel env1' + constructLet shouldAwait k env postponed cont + -- Eliminate this let-binding, if it reads from a Ref, and we already know the value of that Ref. + | RefRead refVar <- binding + , Just valueIdx <- findRefValue (weaken k $ varIdx refVar) env + , LeftHandSideSingle _ <- lhs = + -- Note that RefRead is a trivial binding, so shouldAwait is False + construct body (weakenReplace valueIdx k) env postponed cont + | Exists lhs' <- rebuildLHS lhs + , k' <- sinkWithLHS lhs lhs' k + , binding' <- weaken k binding = + placePostponed (if shouldAwait then postponed else nothingPostponed) env + $ Alet lhs' binding' + $ constructFull (if shouldAwait then body else body{ directlyAwaits = [] }) k' + (buildEnvExtend lhs' binding' env) + (if shouldAwait then nothingPostponed else weaken' (weakenWithLHS lhs') postponed) + $ weaken' (weakenWithLHS lhs') cont + +buildLetNewSignal :: String -> [Idx env SignalResolver] -> BuildSchedule kernel ((env, Signal), SignalResolver) -> BuildSchedule kernel env +buildLetNewSignal comment resolvers body = + -- NewSignal is trivial + BuildSchedule{ + directlyAwaits = map (fromMaybe (internalError "Illegal schedule: deadlock") . strengthenWithLHS lhs) $ directlyAwaits body, + finallyResolves = mapMaybe (strengthenWithLHS lhs) $ finallyResolves body, + trivial = trivial body, + awhileHeight = awhileHeight body, + construct = \k env postponed cont -> if + | otherSignal : _ <- mapMaybe (\idx -> k >:> idx `findSignal` env) resolvers + , k' <- sink $ weakenReplace otherSignal k -> + -- Remove the index for the signal. + -- Replace all occurrences of that signal with 'otherSignal', + -- as their resolvers are resolved at the same time. + Alet lhsResolver (NewSignal comment) + $ construct body k' + (buildEnvExtend lhsResolver (NewSignal comment) env) + (weaken' (weakenSucc weakenId) postponed) + $ weaken' (weakenSucc weakenId) cont + | k' <- sink $ sink k -> + Alet lhs (NewSignal comment) + $ construct body k' + (buildEnvExtend lhs (NewSignal comment) env) + (weaken' (weakenSucc $ weakenSucc weakenId) postponed) + $ weaken' (weakenSucc $ weakenSucc weakenId) cont + } + where + lhs = LeftHandSideSingle BaseRsignal `LeftHandSidePair` LeftHandSideSingle BaseRsignalResolver + lhsResolver = LeftHandSideWildcard (TupRsingle BaseRsignal) `LeftHandSidePair` LeftHandSideSingle BaseRsignalResolver + +buildEnvExtend :: BLeftHandSide t env1 env2 -> Binding env1 t -> BuildEnv env1 -> BuildEnv env2 +buildEnvExtend (LeftHandSidePair (LeftHandSideSingle _) (LeftHandSideSingle _)) (NewSignal _) env = + env `BPush` INone `BPush` IResolvesNext +buildEnvExtend (LeftHandSidePair (LeftHandSideSingle _) (LeftHandSideSingle _)) (NewRef _) env = + env `BPush` INone `BPush` IWritesNext +buildEnvExtend (LeftHandSideSingle _) (RefRead refVar) env = env `BPush` IValue [varIdx refVar] +buildEnvExtend lhs _ env = buildEnvSkip lhs env + +buildEnvSkip :: BLeftHandSide t env1 env2 -> BuildEnv env1 -> BuildEnv env2 +buildEnvSkip lhs env = case lhs of + LeftHandSideWildcard _ -> env + LeftHandSideSingle _ + -> env `BPush` INone + LeftHandSidePair lhs1 lhs2 + -> buildEnvSkip lhs2 $ buildEnvSkip lhs1 $ env + +buildEffect + :: forall kernel env. + Effect kernel env + -> BuildSchedule kernel env + -> BuildSchedule kernel env +buildEffect (SignalResolve []) next = next +buildEffect (SignalResolve resolvers) next = + BuildSchedule{ + directlyAwaits = [], + finallyResolves = + if trivial next && null (directlyAwaits next) then + resolvers' `mergeDedup` finallyResolves next + else + finallyResolves next, + trivial = trivial next && null (directlyAwaits next), + awhileHeight = awhileHeight next, + construct = \k env postponed cont -> + let + resolvers'' = map (weaken k) resolvers' + signals = sort $ mapMaybe (\r -> findSignal r env) resolvers'' + env' = markResolved signals env + in + constructFull next k env' (resolveSignalsInPostponed signals resolvers'' postponed) cont + } + where + resolvers' = sort resolvers +buildEffect (SignalAwait signals) next = + BuildSchedule{ + directlyAwaits = sort signals `mergeDedup` directlyAwaits next, + finallyResolves = finallyResolves next, + trivial = trivial next, + awhileHeight = awhileHeight next, + construct = \k env postponed cont -> construct next k env postponed cont + } +buildEffect effect next + | canPostpone = + BuildSchedule{ + directlyAwaits = directlyAwaits next, + finallyResolves = finallyResolves next, + trivial = effectIsTrivial && trivial next, + awhileHeight = awhileHeight next, + construct = \k env postponed cont -> + Effect (weaken' k effect) + $ construct next k (updateEnv k env) postponed cont + } + | otherwise = + BuildSchedule{ + directlyAwaits = [], + finallyResolves = finallyResolves next, + trivial = effectIsTrivial && trivial next, + awhileHeight = awhileHeight next, + construct = \k env postponed cont -> + placePostponed (if effectIsTrivial then nothingPostponed else postponed) env + $ Effect (weaken' k effect) + $ constructFull next k (updateEnv k env) (if effectIsTrivial then postponed else nothingPostponed) cont + } + where + effectIsTrivial = trivialEffect effect + -- Write may be postponed: a write doesn't do synchronisation, + -- that is done by a signal(resolver). + canPostpone + | RefWrite{} <- effect = True + | otherwise = False + updateEnv :: env :> env' -> BuildEnv env' -> BuildEnv env' + updateEnv k env = case effect of + RefWrite outRefVar valueVar + | Just refIdx <- findRef (k >:> varIdx outRefVar) env + -> markRefValue refIdx (k >:> varIdx valueVar) env + _ -> env + +buildSeq :: BuildSchedule kernel env -> BuildSchedule kernel env -> BuildSchedule kernel env +buildSeq a b = + BuildSchedule { + directlyAwaits = directlyAwaits a, + finallyResolves = + if trivial b && isSubseq then + finallyResolves a `mergeDedup` finallyResolves b + else + finallyResolves b, + trivial = trivial a && trivial b && isSubseq, + awhileHeight = awhileHeight a `max` awhileHeight b, + construct = \k env postponed cont -> + construct a k env postponed $ ContinuationDo k b weakenId cont + } + where + isSubseq = directlyAwaits b `isSubsequenceOf` directlyAwaits a + +buildSpawn :: BuildSchedule kernel env -> BuildSchedule kernel env -> BuildSchedule kernel env +buildSpawn a b + | trivial a && directlyAwaits a `isSubsequenceOf` directlyAwaits b = + buildSeq a b + | trivial b && directlyAwaits b `isSubsequenceOf` directlyAwaits a = + buildSeq b a + | otherwise = + BuildSchedule{ + directlyAwaits = directlyAwaits a `sortedIntersection` directlyAwaits b, + -- Only return the resolved signals at the place where the continuation is placed. + -- We thus ignore 'a' here. + finallyResolves = finallyResolves b, + -- TODO: Make trivial if 'b' is trivial? + -- Might be a problem if the then later on (in construct) make the Spawn sequential (the true case in the if below). + -- We could perhaps add a condition to not make the Spawn sequential if 'b' is trivial and there is some continuation, + -- but that might still be fragile. + trivial = False, + awhileHeight = awhileHeight a `max` awhileHeight b, + construct = \k env postponed cont -> + let + aResolves = sort $ mapMaybe ((`findSignal` env) . weaken k) $ finallyResolves a + a' = a{directlyAwaits = directlyAwaits a `sortedMinus` directlyAwaits b} + b' = b{directlyAwaits = directlyAwaits b `sortedMinus` directlyAwaits a} + in + if not $ null (aResolves `sortedIntersection` map (weaken k) (directlyAwaits b)) then + constructFull a' k env postponed $ ContinuationDo k b' weakenId cont + else + constructFull b' k env (spawnPostponed (BuildSpawn aResolves $ weaken' k a') postponed) cont + } + +-- Builds a term of the form 'Spawn a Return'. +-- This is a separate function from buildSpawn, as 'buildSpawn a Return' will be simplified to 'a'. +-- This is sometimes not desired, like in spawnPostponed, +-- where we must have term 'a' be spawned, to set things like 'finallyResolves' without taking 'a' into account. +-- It must also guarantee that the continuation is placed in parallel with 'a', instead of after 'a'. +-- 'buildSpawn a Return' fails to provide that. +buildSpawnReturn :: BuildSchedule kernel env -> BuildSchedule kernel env +buildSpawnReturn a = + BuildSchedule{ + directlyAwaits = [], + finallyResolves = [], + trivial = True, + awhileHeight = awhileHeight a, + construct = \k env postponed cont -> + let + aResolves = sort $ mapMaybe ((`findSignal` env) . weaken k) $ finallyResolves a + in + constructFull buildReturn k env (spawnPostponed (BuildSpawn aResolves $ weaken' k a) postponed) cont + } + +buildAcond + :: ExpVar env PrimBool + -> BuildSchedule kernel env -- True branch + -> BuildSchedule kernel env -- False branch + -> BuildSchedule kernel env -- Operations after the if-then-else + -> BuildSchedule kernel env +buildAcond var true false next = + BuildSchedule{ + directlyAwaits = directlyAwaits true `sortedIntersection` directlyAwaits false, + finallyResolves = finallyResolves next, + trivial = False, + awhileHeight = awhileHeight true `max` awhileHeight false `max` awhileHeight next, + construct = \k env postponed cont -> + placePostponed postponed env + $ Acond + (weaken k var) + (constructFull true{directlyAwaits = directlyAwaits true `sortedMinus` directlyAwaits false} k env nothingPostponed ContinuationEnd) + (constructFull false{directlyAwaits = directlyAwaits false `sortedMinus` directlyAwaits true} k env nothingPostponed ContinuationEnd) + (constructFull next k env nothingPostponed cont) + } + +buildAwhile + :: InputOutputR input output + -> BuildScheduleFun kernel env (input -> Output PrimBool -> output -> ()) + -> BaseVars env input + -> BuildSchedule kernel env -- Operations after the while loop + -> BuildSchedule kernel env +buildAwhile io step initial next = + BuildSchedule{ + -- TODO: Also check if the function directly awaits parts of the input, and take the accompanying signals from 'initial' + directlyAwaits = funDirectlyAwaits step, + finallyResolves = finallyResolves next, + trivial = False, + awhileHeight = (funAwhileHeight step + 1) `max` awhileHeight next, + construct = \k env postponed cont -> + placePostponed postponed env + -- If a this awhile contains another awhile, downgrade this awhile to an awhile-seq. + -- Nested awhiles would exponentially blow up the memory required to + -- store the space of this program. + $ if funAwhileHeight step > 0 then + let + env' = env `BPush` IResolved `BPush` INone + k' = weakenSucc' $ weakenSucc' k + step' = downgradeAwhileFun (SuccIdx ZeroIdx) $ weaken k' step + in + await (weaken k <$> varsGetSignals initial) + $ Alet lhsSignal (NewSignal "Signal for awhile to awhile-seq downgrade") + $ Effect (SignalResolve [ZeroIdx]) + $ AwhileSeq + (ioRemoveSignal io) + (funConstruct step' weakenId env') + (mapTupR (weaken k') $ varsRemoveSignal initial) + (constructFull next k' env' nothingPostponed $ weaken' (weakenSucc $ weakenSucc weakenId) cont) + else + Awhile + io + (funConstruct step k env) + (mapTupR (weaken k) initial) + (constructFull next k env nothingPostponed cont) + } + +buildAwhileSeq + -- Should not contain InputOutputRsignal + :: InputOutputR input output + -> BuildScheduleFun kernel env (input -> OutputRef PrimBool -> output -> ()) + -> BaseVars env input + -> BuildSchedule kernel env -- Operations after the while loop + -> BuildSchedule kernel env +buildAwhileSeq io step initial next = + BuildSchedule{ + directlyAwaits = funDirectlyAwaits step, + finallyResolves = finallyResolves next, + trivial = False, + awhileHeight = (funAwhileHeight step + 1) `max` awhileHeight next, + construct = \k env postponed cont -> + placePostponed postponed env + $ AwhileSeq + io + (funConstruct step k env) + (mapTupR (weaken k) initial) + (constructFull next k env nothingPostponed cont) + } + +downgradeAwhileFun + :: forall kernel env input output. + Idx env Signal -- Index to a resolved signal + -> BuildScheduleFun kernel env (input -> (SignalResolver, OutputRef PrimBool) -> output -> ()) + -> BuildScheduleFun kernel env (NoSignal input -> OutputRef PrimBool -> NoSignal output -> ()) +downgradeAwhileFun signalIdx (BuildLam lhsInput (BuildLam lhsBool (BuildLam lhsOutput (BuildBody f)))) + | Exists lhsInput' <- lhsRemoveSignal lhsInput + , kInput <- declareRemovedSignal lhsInput lhsInput' signalIdx weakenId + , Exists lhsBoolOutput' <- lhsRemoveSignal (LeftHandSidePair lhsBool lhsOutput) + , DeclareRemovedResolvers termOutput kOutput <- + declareRemovedResolvers @kernel + (LeftHandSidePair lhsBool lhsOutput) + lhsBoolOutput' + kInput + = case lhsBoolOutput' of + LeftHandSidePair lhsBool' lhsOutput' -> + BuildLam lhsInput' + $ BuildLam (lhsSnd lhsBool') + $ BuildLam lhsOutput' + $ BuildBody + $ termOutput + $ weaken' kOutput f + LeftHandSideWildcard _ -> + internalError "Infinite awhile loop: condition of the loop never returns" + LeftHandSideSingle tp -> pairImpossible tp + where + lhsSnd :: BLeftHandSide ((), b) e1 e2 -> BLeftHandSide b e1 e2 + lhsSnd (LeftHandSidePair LeftHandSideWildcard{} l) = l + lhsSnd (LeftHandSidePair (LeftHandSideSingle t) _) = unitImpossible t + lhsSnd (LeftHandSideWildcard (TupRpair _ t)) = LeftHandSideWildcard t + lhsSnd (LeftHandSideWildcard (TupRsingle t)) = pairImpossible t + lhsSnd (LeftHandSideSingle t) = pairImpossible t + +buildFunLam + :: BLeftHandSide t env1 env2 + -> BuildScheduleFun kernel env2 f + -> BuildScheduleFun kernel env1 (t -> f) +buildFunLam = BuildLam + +buildFunBody :: BuildSchedule kernel env -> BuildScheduleFun kernel env () +buildFunBody = BuildBody + +lhsSignal :: BLeftHandSide (Signal, SignalResolver) fenv ((fenv, Signal), SignalResolver) +lhsSignal = LeftHandSidePair (LeftHandSideSingle BaseRsignal) (LeftHandSideSingle BaseRsignalResolver) + +lhsRef :: GroundR tp -> LeftHandSide BaseR (Ref tp, OutputRef tp) fenv ((fenv, Ref tp), OutputRef tp) +lhsRef tp = LeftHandSidePair (LeftHandSideSingle $ BaseRref tp) (LeftHandSideSingle $ BaseRrefWrite tp) + +-- Assumes that the input arrays are sorted, +-- with no duplicates. +-- Creates a sorted list containing all elements of the two input list. +-- If an element is present in both input lists, it will be added only +-- once to the output. +mergeDedup :: Ord a => [a] -> [a] -> [a] +mergeDedup as@(a:as') bs@(b:bs') + | a == b = a : mergeDedup as' bs' + | a < b = a : mergeDedup as' bs + | otherwise = b : mergeDedup as bs' +mergeDedup as [] = as +mergeDedup [] bs = bs + +sortedDedup :: Eq a => [a] -> [a] +sortedDedup = \case + [] -> [] + a : as -> go a as + where + go x (y:ys) + | x == y = go x ys + | otherwise = x : go y ys + go x [] = [x] + +-- Constructs the intersection of two lists, +-- assuming they are sorted and have no duplicates. +sortedIntersection :: Ord a => [a] -> [a] -> [a] +sortedIntersection as@(a:as') bs@(b:bs') + | a == b = a : sortedIntersection as' bs' + | a < b = sortedIntersection as' bs + | otherwise = sortedIntersection as bs' +sortedIntersection _ _ = [] + +-- Constructs the difference of two lists, +-- assuming they are sorted and have no duplicates. +-- It returns all elements present in the first list, +-- but not in the second +sortedMinus :: Ord a => [a] -> [a] -> [a] +sortedMinus as@(a:as') bs@(b:bs') + | a == b = sortedMinus as' bs' + | a < b = a : sortedMinus as' bs + | otherwise = sortedMinus as bs' +sortedMinus as [] = as +sortedMinus [] _ = [] + +-- Simplify schedule, by rebuilding it using the build functions +simplify :: UniformScheduleFun kernel () t -> UniformScheduleFun kernel () t +simplify f = funConstruct (rebuildFun f) weakenId BEmpty + +rebuildFun :: UniformScheduleFun kernel env t -> BuildScheduleFun kernel env t +rebuildFun (Slam lhs f) = buildFunLam lhs $ rebuildFun f +rebuildFun (Sbody body) = buildFunBody $ snd $ rebuild body + +rebuild :: UniformSchedule kernel env -> (SignalAnalysis env, BuildSchedule kernel env) +rebuild = \case + Return -> (SEmpty, buildReturn) + Alet lhs bnd body + | (analysis, body') <- rebuild body -> + ( analysisDrop lhs analysis + , rebuildLet lhs bnd analysis body' + ) + Effect eff next + | (analysis, next') <- rebuild next -> + ( analyseEffect eff `analysisJoin` analysis + , buildEffect eff next' + ) + Acond var true false next + | (aTrue, true') <- rebuild true + , (aFalse, false') <- rebuild false + , (aNext, next') <- rebuild next -> + ( analysisMeet aTrue aFalse `analysisJoin` aNext + , buildAcond var true' false' next' + ) + Awhile io f input next + | (analysis, next') <- rebuild next -> + ( analysis + , buildAwhile io (rebuildFun f) input next' + ) + AwhileSeq io f input next + | (analysis, next') <- rebuild next -> + ( analysis + , buildAwhileSeq io (rebuildFun f) input next' + ) + Spawn term1 term2 + | (analysis1, term1') <- rebuild term1 + , (analysis2, term2') <- rebuild term2 -> + ( analysisJoin analysis1 analysis2 + , buildSpawn term1' term2' + ) + +rebuildLet + :: BLeftHandSide t env env' + -> Binding env t + -> SignalAnalysis env' + -> BuildSchedule kernel env' + -> BuildSchedule kernel env +rebuildLet (LeftHandSidePair LeftHandSideSingle{} LeftHandSideSingle{}) (NewSignal comment) (SPush _ (SIResolvedWith resolvers)) body = buildLetNewSignal comment (map unSucc resolvers) body + where + unSucc :: Idx (env, Signal) SignalResolver -> Idx env SignalResolver + unSucc (SuccIdx idx) = idx +rebuildLet lhs bnd _ body = buildLet lhs bnd body + +-- Signal analysis +data SignalAnalysis env where + SEmpty :: SignalAnalysis env + SPush :: SignalAnalysis env -> SignalInfo env t -> SignalAnalysis (env, t) + +{- TODO WALL: DEAD CODE +spush :: SignalAnalysis env -> SignalInfo env t -> SignalAnalysis (env, t) +spush SEmpty SINone = SEmpty +spush env info = SPush env info +-} + +data SignalInfo env t where + -- This SignalResolver is resolved at the same time as the given list of SignalResolvers. + SIResolvedWith + :: [Idx env SignalResolver] + -> SignalInfo env SignalResolver + + SINone + :: SignalInfo env t + +analysisDrop :: LeftHandSide s t env env' -> SignalAnalysis env' -> SignalAnalysis env +analysisDrop _ SEmpty = SEmpty +analysisDrop LeftHandSideWildcard{} env = env +analysisDrop LeftHandSideSingle{} (SPush env _) = env +analysisDrop (LeftHandSidePair lhs1 lhs2) env = analysisDrop lhs1 $ analysisDrop lhs2 env + +-- Use this when two terms are both executed, for instance in a spawn +analysisJoin :: SignalAnalysis env -> SignalAnalysis env -> SignalAnalysis env +analysisJoin SEmpty env = env +analysisJoin env SEmpty = env +analysisJoin (SPush as a) (SPush bs b) = analysisJoin as bs `SPush` signalInfoJoin a b + +signalInfoJoin :: SignalInfo env t -> SignalInfo env t -> SignalInfo env t +signalInfoJoin SINone info = info +signalInfoJoin info SINone = info +signalInfoJoin (SIResolvedWith as) (SIResolvedWith bs) = SIResolvedWith $ as `mergeDedup` bs + +-- Use this when only one of the two terms is executed, for instance in an if-then-else +analysisMeet :: SignalAnalysis env -> SignalAnalysis env -> SignalAnalysis env +analysisMeet SEmpty _ = SEmpty +analysisMeet _ SEmpty = SEmpty +analysisMeet (SPush as a) (SPush bs b) = analysisMeet as bs `SPush` signalInfoMeet a b + +signalInfoMeet :: SignalInfo env t -> SignalInfo env t -> SignalInfo env t +signalInfoMeet SINone _ = SINone +signalInfoMeet _ SINone = SINone +signalInfoMeet (SIResolvedWith as) (SIResolvedWith bs) = SIResolvedWith $ as `sortedIntersection` bs + +analyseEffect :: Effect kernel env -> SignalAnalysis env +analyseEffect (SignalResolve resolvers) = analyseSignalResolve resolvers +analyseEffect _ = SEmpty + +analyseSignalResolve :: [Idx env SignalResolver] -> SignalAnalysis env +analyseSignalResolve = const SEmpty -- go . sort +{- TODO WALL: DEAD CODE + where + -- input is sorted from low indices to high indices + go :: [Idx env SignalResolver] -> SignalAnalysis env + go [] = SEmpty + go [_] = SEmpty + go (ZeroIdx : ids) = go ids' `SPush` SIResolvedWith ids' + where ids' = map unSucc ids + go ids@(SuccIdx _ : _) = go (map unSucc ids) `spush` SINone + + unSucc :: Idx (env, s) t -> Idx env t + unSucc (SuccIdx idx) = idx + unSucc ZeroIdx = internalError "Expected non-zero index. Is the list of indices sorted and unique?" +-} + +-- Removes Signals from the InputOutputR of an awhile loop +ioRemoveSignal :: InputOutputR input output -> InputOutputR (NoSignal input) (NoSignal output) +ioRemoveSignal InputOutputRsignal = InputOutputRunit +ioRemoveSignal (InputOutputRref tp) = InputOutputRref tp +ioRemoveSignal InputOutputRunit = InputOutputRunit +ioRemoveSignal (InputOutputRpair a b) = ioRemoveSignal a `InputOutputRpair` ioRemoveSignal b + +typeRemoveSignal :: BasesR input -> BasesR (NoSignal input) +typeRemoveSignal TupRunit = TupRunit +typeRemoveSignal (TupRsingle t) = case t of + BaseRsignal -> TupRunit + BaseRsignalResolver -> TupRunit + BaseRref _ -> TupRsingle t + BaseRrefWrite _ -> TupRsingle t + _ -> internalError "Only expected Signals, SignalResolvers, and Refs" +typeRemoveSignal (TupRpair t1 t2) = typeRemoveSignal t1 `TupRpair` typeRemoveSignal t2 + +lhsRemoveSignal :: BLeftHandSide input env1 env2 -> Exists (BLeftHandSide (NoSignal input) env3) +lhsRemoveSignal (LeftHandSideWildcard tp) = Exists $ LeftHandSideWildcard $ typeRemoveSignal tp +lhsRemoveSignal (LeftHandSideSingle tp) = case tp of + BaseRsignal -> Exists $ LeftHandSideWildcard TupRunit + BaseRsignalResolver -> Exists $ LeftHandSideWildcard TupRunit + BaseRref _ -> Exists $ LeftHandSideSingle tp + BaseRrefWrite _ -> Exists $ LeftHandSideSingle tp + _ -> internalError "Only expected Signals, SignalResolvers, and Refs" +lhsRemoveSignal (LeftHandSidePair lhs1 lhs2) + | Exists lhs1' <- lhsRemoveSignal lhs1 + , Exists lhs2' <- lhsRemoveSignal lhs2 + = Exists $ LeftHandSidePair lhs1' lhs2' + +varsRemoveSignal :: BaseVars env input -> BaseVars env (NoSignal input) +varsRemoveSignal TupRunit = TupRunit +varsRemoveSignal vars@(TupRsingle (Var t _)) = case t of + BaseRsignal -> TupRunit + BaseRsignalResolver -> TupRunit + BaseRref _ -> vars + BaseRrefWrite _ -> vars + _ -> internalError "Only expected Signals, SignalResolvers, and Refs" +varsRemoveSignal (TupRpair v1 v2) = varsRemoveSignal v1 `TupRpair` varsRemoveSignal v2 + +varsGetSignals :: BaseVars env input -> [Idx env Signal] +varsGetSignals (TupRsingle (Var BaseRsignal idx)) = [idx] +varsGetSignals (TupRpair v1 v2) = varsGetSignals v1 ++ varsGetSignals v2 +varsGetSignals _ = [] -- TupRunit or TupRsingle of a different type + +type family NoSignal input where + NoSignal Signal = () + NoSignal SignalResolver = () + NoSignal (a, b) = (NoSignal a, NoSignal b) + NoSignal input = input + +-- Substitutes the uses of a removed Signal variable to a newly introduced +-- Signal, that is directly resolved. +-- This is needed for converting an awhile to an awhile-seq. +-- Only call this function on the input of the body of an awhile loop. +-- This function expects that the left hand side only binds Refs and Signals. +-- Furthermore, it expects that 'Exists lhsNew = lhsRemoveSignal lhsOld'. +-- The newly introduced Signal (and SignalResolver) should eventually be +-- removed. The Signal is directly resolved, and our optimizer will thus +-- remove SignalAwaits waiting on that signal. This will make the variables +-- dead, and strongly live variable analysis will then remove it. +-- We go via this indirection, since BuildSchedule cannot remove variables; +-- it can only rename variables. +declareRemovedSignal + :: forall env0 env1 env1' env2' t. + BLeftHandSide t env0 env1 + -> BLeftHandSide (NoSignal t) env0 env1' + -> Idx env0 Signal -- Index of a resolved signal + -> env1' :> env2' + -> env1 :> env2' +declareRemovedSignal lhsOld lhsNew signalIdx k = + go lhsOld lhsNew k (k .> weakenWithLHS lhsNew) + where + go + :: BLeftHandSide s envA envB + -> BLeftHandSide (NoSignal s) envA' envB' + -> envB' :> env2' + -> envA :> env2' + -> envB :> env2' + -- A preserved variable. + go (LeftHandSideSingle t) (LeftHandSideSingle _) kB' kA = case t of + BaseRref _ -> weakenReplace (weaken kB' ZeroIdx) kA + _ -> internalError "Signal impossible, other types not expected" + -- A removed Signal. Use the newly introduced signal instead. + go (LeftHandSideSingle BaseRsignal) (LeftHandSideWildcard _) _ kA = + weakenReplace (weaken (k .> weakenWithLHS lhsNew) signalIdx) kA + go (LeftHandSidePair old1 old2) (LeftHandSidePair new1 new2) kB' kA = + go old2 new2 kB' $ go old1 new1 (kB' .> weakenWithLHS new2) kA + go (LeftHandSideWildcard _) (LeftHandSideWildcard _) _ kA = kA + go _ _ _ _ = internalError "LeftHandSide mismatch" + +-- Similar to declareRemovedSignal, but +-- 1. This function is about removed SignalResolvers. +-- 2. This function operates on the output and thus expects OutputRefs instead of Refs +-- 3. Instead of defining one SignalResolver and using that for each removed variable, +-- we must now define one SignalResolver per variable. +declareRemovedResolvers + :: forall kernel env0 env0' env1 env1' t. + BLeftHandSide t env0 env1 + -> BLeftHandSide (NoSignal t) env0' env1' + -> env0 :> env0' + -> DeclareRemovedResolvers kernel env1 env1' +declareRemovedResolvers lhsOld lhsNew k = + bindResolvers removed $ \skip term resolvers -> + DeclareRemovedResolvers + term + $ snd $ go resolvers lhsOld lhsNew + (skipWeakenIdx skip) + (skipWeakenIdx skip .> weakenWithLHS lhsNew .> k) + where + removed = lhsSize lhsOld - lhsSize lhsNew + + bindResolvers + :: Int + -> (forall env2'. + Skip env2' env1' + -> (BuildSchedule kernel env2' -> BuildSchedule kernel env1') + -> [Idx env2' SignalResolver] + -> DeclareRemovedResolvers kernel env1 env1') + -> DeclareRemovedResolvers kernel env1 env1' + bindResolvers 0 f = f SkipNone id [] + bindResolvers n f = bindResolvers (n - 1) $ \skip term resolvers -> + f + (SkipSucc' skip) + (term . buildLet lhsOnlyResolver (NewSignal "SignalResolver for awhile to awhile-seq downgrade")) + (ZeroIdx : map SuccIdx resolvers) + where + lhsOnlyResolver :: BLeftHandSide (Signal, SignalResolver) e (e, SignalResolver) + lhsOnlyResolver = LeftHandSideWildcard (TupRsingle BaseRsignal) `LeftHandSidePair` LeftHandSideSingle BaseRsignalResolver + + go + :: [Idx env2' SignalResolver] + -> BLeftHandSide s envA envB + -> BLeftHandSide (NoSignal s) envA' envB' + -> envB' :> env2' + -> envA :> env2' + -> ([Idx env2' SignalResolver], envB :> env2') + go resolvers (LeftHandSideSingle t) (LeftHandSideSingle _) kB' kA = case t of + BaseRrefWrite _ -> (resolvers, weakenReplace (weaken kB' ZeroIdx) kA) + _ -> internalError "SignalResolver impossible, other types not expected" + -- A removed Resolver. Use a newly introduced resolver instead. + go resolvers (LeftHandSideSingle BaseRsignalResolver) (LeftHandSideWildcard _) _ kA = case resolvers of + [] -> internalError "Not enough new SignalResolvers" + r:rs -> (rs, weakenReplace r kA) + go resolvers (LeftHandSidePair old1 old2) (LeftHandSidePair new1 new2) kB' kA + | (resolvers1, kA1) <- go resolvers old1 new1 (kB' .> weakenWithLHS new2) kA + = go resolvers1 old2 new2 kB' kA1 + go resolvers (LeftHandSideWildcard _) (LeftHandSideWildcard _) _ kA = (resolvers, kA) + go _ _ _ _ _ = internalError "LeftHandSide mismatch" + +-- Captures existential env2' +data DeclareRemovedResolvers kernel env1 env1' where + DeclareRemovedResolvers + :: (BuildSchedule kernel env2' -> BuildSchedule kernel env1') + -> env1 :> env2' + -> DeclareRemovedResolvers kernel env1 env1' diff --git a/src/Data/Array/Accelerate/Trafo/Schedule/Uniform/Substitution.hs b/src/Data/Array/Accelerate/Trafo/Schedule/Uniform/Substitution.hs new file mode 100644 index 000000000..f63eecaa5 --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/Schedule/Uniform/Substitution.hs @@ -0,0 +1,171 @@ +{-# OPTIONS_GHC -Wno-orphans #-} +{-# OPTIONS_HADDOCK hide #-} +-- | +-- Module : Data.Array.Accelerate.Trafo.Schedule.Uniform +-- Copyright : [2012..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.Trafo.Schedule.Uniform.Substitution ( + NewIdx(..), ReindexPartialN, SunkReindexPartialN(..), + reindexArgs, reindexSchedule +) where + +import Data.Array.Accelerate.AST.Environment +import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.AST.Schedule.Uniform +import Data.Array.Accelerate.Error +import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Trafo.Exp.Substitution +import Data.Array.Accelerate.Trafo.Substitution +import Data.Maybe +import Prelude hiding (id, (.), read) +import Control.Category +import Data.Functor.Identity + + +instance Sink' (UniformSchedule kernel) where + weaken' _ Return = Return + weaken' k (Alet lhs b s) + | Exists lhs' <- rebuildLHS lhs = Alet lhs' (weaken k b) (weaken' (sinkWithLHS lhs lhs' k) s) + weaken' k (Effect effect s) = Effect (weaken' k effect) (weaken' k s) + weaken' k (Acond cond true false s) = Acond (weaken k cond) (weaken' k true) (weaken' k false) (weaken' k s) + weaken' k (Awhile io f input s) = Awhile io (weaken k f) (mapTupR (weaken k) input) (weaken' k s) + weaken' k (AwhileSeq io f input s) = AwhileSeq io (weaken k f) (mapTupR (weaken k) input) (weaken' k s) + weaken' k (Spawn s1 s2) = Spawn (weaken' k s1) (weaken' k s2) + +instance Sink (UniformScheduleFun kernel) where + weaken k (Slam lhs f) + | Exists lhs' <- rebuildLHS lhs = Slam lhs' $ weaken (sinkWithLHS lhs lhs' k) f + weaken k (Sbody s) = Sbody $ weaken' k s + +instance Sink Binding where + weaken k (Compute e) = Compute $ mapArrayInstr (weaken k) e + weaken _ (NewSignal name) = NewSignal name + weaken _ (NewRef r) = NewRef r + weaken k (Alloc shr tp sz) = Alloc shr tp $ mapTupR (weaken k) sz + weaken _ (Use tp n buffer) = Use tp n buffer + weaken k (Unit var) = Unit $ weaken k var + weaken k (RefRead ref) = RefRead $ weaken k ref + +instance Sink' (Effect kernel) where + weaken' k (Exec md kernel args) = Exec md kernel $ runIdentity $ reindexSArgs' (ReindexF $ \ix -> NewIdxJust <$> weakenReindex k ix) args + weaken' k (SignalAwait vars) = SignalAwait $ map (weaken k) vars + weaken' k (SignalResolve vars) = SignalResolve $ map (weaken k) vars + weaken' k (RefWrite ref value) = RefWrite (weaken k ref) (weaken k value) + weaken' k (Aassert msg cond) = Aassert msg $ mapArrayInstr (weaken k) cond + weaken' k (Atrace msg t) = Atrace msg $ weakenArrayDescriptors k t + +instance Sink SArg where + weaken k = runIdentity . reindexSArg' (ReindexF $ \ix -> NewIdxJust <$> weakenReindex k ix) + +type ReindexPartialN f env env' = forall a. Idx env a -> f (NewIdx env' a) + +data NewIdx env t where + NewIdxNoResolver :: NewIdx env SignalResolver + NewIdxJust :: Idx env t -> NewIdx env t + +data SunkReindexPartialN f env env' where + Sink :: SunkReindexPartialN f env env' -> SunkReindexPartialN f (env, s) (env', s) + ReindexF :: ReindexPartialN f env env' -> SunkReindexPartialN f env env' + + +reindexSchedule :: (Applicative f) => ReindexPartialN f env env' -> UniformSchedule kernel env -> f (UniformSchedule kernel env') +reindexSchedule k = reindexSchedule' $ ReindexF k + +sinkReindexWithLHS :: LeftHandSide s t env1 env1' -> LeftHandSide s t env2 env2' -> SunkReindexPartialN f env1 env2 -> SunkReindexPartialN f env1' env2' +sinkReindexWithLHS (LeftHandSideWildcard _) (LeftHandSideWildcard _) k = k +sinkReindexWithLHS (LeftHandSideSingle _) (LeftHandSideSingle _) k = Sink k +sinkReindexWithLHS (LeftHandSidePair a1 b1) (LeftHandSidePair a2 b2) k = sinkReindexWithLHS b1 b2 $ sinkReindexWithLHS a1 a2 k +sinkReindexWithLHS _ _ _ = internalError "sinkReindexWithLHS: left hand sides don't match" + +reindex' :: Applicative f => SunkReindexPartialN f env env' -> ReindexPartialN f env env' +reindex' (ReindexF f) = f +reindex' (Sink k) = \case + ZeroIdx -> pure $ NewIdxJust ZeroIdx + SuccIdx ix -> + let + f NewIdxNoResolver = NewIdxNoResolver + f (NewIdxJust ix') = NewIdxJust $ SuccIdx ix' + in + f <$> reindex' k ix + +reindexSchedule' :: (Applicative f) => SunkReindexPartialN f env env' -> UniformSchedule kernel env -> f (UniformSchedule kernel env') +reindexSchedule' k = \case + Return -> pure Return + Alet lhs bnd s + | Exists lhs' <- rebuildLHS lhs -> Alet lhs' <$> reindexBinding' k bnd <*> reindexSchedule' (sinkReindexWithLHS lhs lhs' k) s + Effect effect s -> Effect <$> reindexEffect' k effect <*> reindexSchedule' k s + Acond cond t f continue -> Acond <$> reindexVarUnsafe k cond <*> reindexSchedule' k t <*> reindexSchedule' k f <*> reindexSchedule' k continue + Awhile io f initial continue -> Awhile io <$> reindexScheduleFun' k f <*> traverseTupR (reindexVarUnsafe k) initial <*> reindexSchedule' k continue + AwhileSeq io f initial continue -> AwhileSeq io <$> reindexScheduleFun' k f <*> traverseTupR (reindexVarUnsafe k) initial <*> reindexSchedule' k continue + Spawn s1 s2 -> Spawn <$> reindexSchedule' k s1 <*> reindexSchedule' k s2 + +reindexVarUnsafe :: Applicative f => SunkReindexPartialN f env env' -> Var s env t -> f (Var s env' t) +reindexVarUnsafe k (Var tp idx) = Var tp . fromNewIdxUnsafe <$> reindex' k idx + +reindexScheduleFun' :: (Applicative f) => SunkReindexPartialN f env env' -> UniformScheduleFun kernel env t -> f (UniformScheduleFun kernel env' t) +reindexScheduleFun' k = \case + Sbody s -> Sbody <$> reindexSchedule' k s + Slam lhs f + | Exists lhs' <- rebuildLHS lhs -> Slam lhs' <$> reindexScheduleFun' (sinkReindexWithLHS lhs lhs' k) f + +reindexEffect' :: forall kernel f env env'. (Applicative f) => SunkReindexPartialN f env env' -> Effect kernel env -> f (Effect kernel env') +reindexEffect' k = \case + Exec md kernel args -> Exec md kernel <$> reindexSArgs' k args + SignalAwait signals -> SignalAwait <$> traverse (fromNewIdxSignal <.> reindex' k) signals + SignalResolve resolvers -> SignalResolve . mapMaybe toMaybe <$> traverse (reindex' k) resolvers + RefWrite ref value -> RefWrite <$> reindexVar (fromNewIdxOutputRef <.> reindex' k) ref <*> reindexVar (fromNewIdxUnsafe <.> reindex' k) value + Aassert msg cond -> Aassert msg <$> reindexExp (fromNewIdxUnsafe <.> reindex' k) cond + Atrace msg t -> Atrace msg <$> reindexArrayDescriptors (fromNewIdxUnsafe <.> reindex' k) t + where + toMaybe :: NewIdx env' a -> Maybe (Idx env' a) + toMaybe (NewIdxJust idx) = Just idx + toMaybe _ = Nothing + +reindexSArgs' :: Applicative f => SunkReindexPartialN f env env' -> SArgs env t -> f (SArgs env' t) +reindexSArgs' k = \case + a :>: as -> (:>:) <$> reindexSArg' k a <*> reindexSArgs' k as + ArgsNil -> pure ArgsNil + +reindexSArg' :: Applicative f => SunkReindexPartialN f env env' -> SArg env t -> f (SArg env' t) +reindexSArg' k = \case + SArgScalar var -> SArgScalar <$> reindexVarUnsafe k var + SArgBuffer m var -> SArgBuffer m <$> reindexVarUnsafe k var + +-- For Exec we cannot have a safe function from the conversion, +-- as we cannot enforce in the type system that no SignalResolvers +-- occur in an Exec or Compute. +fromNewIdxUnsafe :: NewIdx env' a -> Idx env' a +fromNewIdxUnsafe (NewIdxJust idx) = idx +fromNewIdxUnsafe _ = error "Expected NewIdxJust" + +-- Different versions, which have different ways of getting evidence +-- that NewIdxNoResolver is impossible +fromNewIdxSignal :: NewIdx env' Signal -> Idx env' Signal +fromNewIdxSignal (NewIdxJust idx) = idx + +fromNewIdxOutputRef :: NewIdx env' (OutputRef t) -> Idx env' (OutputRef t) +fromNewIdxOutputRef (NewIdxJust idx) = idx + +fromNewIdxGround :: GroundR a -> NewIdx env' a -> Idx env' a +fromNewIdxGround _ (NewIdxJust idx) = idx +fromNewIdxGround tp NewIdxNoResolver = signalResolverImpossible (TupRsingle tp) + +reindexBinding' :: Applicative f => SunkReindexPartialN f env env' -> Binding env t -> f (Binding env' t) +reindexBinding' k = \case + Compute e -> Compute <$> reindexExp (fromNewIdxUnsafe <.> reindex' k) e + NewSignal name -> pure $ NewSignal name + NewRef tp -> pure $ NewRef tp + Alloc shr tp sh -> Alloc shr tp <$> reindexVars (fromNewIdxUnsafe <.> reindex' k) sh + Use tp n buffer -> pure $ Use tp n buffer + Unit (Var tp idx) -> Unit . Var tp <$> (fromNewIdxGround (GroundRscalar tp) <.> reindex' k) idx + RefRead ref -> RefRead <$> reindexVar (fromNewIdxUnsafe <.> reindex' k) ref + +(<.>) :: Applicative f => (b -> c) -> (a -> f b) -> a -> f c +(<.>) g h a = g <$> h a diff --git a/src/Data/Array/Accelerate/Trafo/Sharing.hs b/src/Data/Array/Accelerate/Trafo/Sharing.hs index 627ccf3f6..9c540634a 100644 --- a/src/Data/Array/Accelerate/Trafo/Sharing.hs +++ b/src/Data/Array/Accelerate/Trafo/Sharing.hs @@ -1,20 +1,7 @@ {-# LANGUAGE AllowAmbiguousTypes #-} -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} {-# LANGUAGE OverloadedLists #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE PatternGuards #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE StandaloneDeriving #-} -{-# LANGUAGE TupleSections #-} -{-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} -{-# OPTIONS_GHC -fno-warn-name-shadowing #-} +{-# OPTIONS_GHC -Wno-name-shadowing #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.Trafo.Sharing @@ -35,7 +22,7 @@ module Data.Array.Accelerate.Trafo.Sharing ( convertAcc, convertAccWith, Afunction, AfunctionR, ArraysFunctionR, AfunctionRepr(..), afunctionRepr, - convertAfun, convertAfunWith, + convertAfun, convertAfunWith, afunctionGroundR, Function, FunctionR, EltFunctionR, FunctionRepr(..), functionRepr, convertExp, convertExpWith, @@ -45,7 +32,7 @@ module Data.Array.Accelerate.Trafo.Sharing ( ) where -import Data.Array.Accelerate.AST hiding ( PreOpenAcc(..), OpenAcc(..), Acc, OpenExp(..), Exp, Boundary(..), HasArraysR(..), formatPreAccOp ) +import Data.Array.Accelerate.AST hiding ( PreOpenAcc(..), OpenAcc(..), Acc, ArrayInstr(..), PreOpenExp(..), OpenExp, Exp, Boundary(..), HasArraysR(..), formatPreAccOp ) import Data.Array.Accelerate.AST.Environment import Data.Array.Accelerate.AST.Idx import Data.Array.Accelerate.AST.LeftHandSide @@ -55,6 +42,7 @@ import Data.Array.Accelerate.Debug.Internal.Flags as Debug import Data.Array.Accelerate.Debug.Internal.Trace as Debug import Data.Array.Accelerate.Error import Data.Array.Accelerate.Representation.Array ( Array, ArraysR, ArrayR(..), formatArraysR ) +import Data.Array.Accelerate.Representation.Ground ( GFunctionR(..), LoweredAfun, lowerArraysR, loweredAfunIsBody ) import Data.Array.Accelerate.Representation.Shape hiding ( zip ) import Data.Array.Accelerate.Representation.Stencil import Data.Array.Accelerate.Representation.Tag @@ -191,14 +179,16 @@ data AfunctionRepr f ar areprr where class Afunction f where type AfunctionR f type ArraysFunctionR f - afunctionRepr :: HasCallStack => AfunctionRepr f (AfunctionR f) (ArraysFunctionR f) - convertOpenAfun :: HasCallStack => Config -> ArrayLayout aenv aenv -> f -> AST.OpenAfun aenv (ArraysFunctionR f) + afunctionRepr :: HasCallStack => AfunctionRepr f (AfunctionR f) (ArraysFunctionR f) + afunctionGroundR :: GFunctionR (LoweredAfun (ArraysFunctionR f)) + convertOpenAfun :: HasCallStack => Config -> ArrayLayout aenv aenv -> f -> AST.OpenAfun aenv (ArraysFunctionR f) instance (Arrays a, Afunction r) => Afunction (Acc a -> r) where type AfunctionR (Acc a -> r) = a -> AfunctionR r type ArraysFunctionR (Acc a -> r) = Sugar.ArraysR a -> ArraysFunctionR r afunctionRepr = AfunctionReprLam $ afunctionRepr @r + afunctionGroundR = GFunctionRlam (lowerArraysR $ Sugar.arraysR @a) $ afunctionGroundR @r convertOpenAfun config alyt f | repr <- Sugar.arraysR @a , DeclareVars lhs k value <- declareVars repr @@ -211,6 +201,9 @@ instance Arrays b => Afunction (Acc b) where type AfunctionR (Acc b) = b type ArraysFunctionR (Acc b) = Sugar.ArraysR b afunctionRepr = AfunctionReprBody + afunctionGroundR + | Refl <- loweredAfunIsBody repr = GFunctionRbody (lowerArraysR repr) + where repr = Sugar.arraysR @b convertOpenAfun config alyt (Acc body) = Abody $ convertOpenAcc config alyt body convertSmartAfun1 @@ -317,16 +310,7 @@ convertSharingAcc config alyt aenv (ScopedAcc lams (AccSharing _ preAcc)) -> let AST.OpenAcc a = avarsIn AST.OpenAcc $ prjIdx (bformat ("de Bruijn conversion tag " % int) i) formatArraysR matchArraysR repr i alyt in a - Pipe reprA reprB reprC (afun1 :: SmartAcc as -> ScopedAcc bs) (afun2 :: SmartAcc bs -> ScopedAcc cs) acc - | DeclareVars lhs k value <- declareVars reprB -> - let - noStableSharing = StableSharingAcc noStableAccName (undefined :: SharingAcc acc exp ()) - boundAcc = AST.Apply reprB (cvtAfun1 reprA afun1) (cvtA acc) - alyt' = PushLayout (incLayout k alyt) lhs (value weakenId) - bodyAcc = AST.Apply reprC - (convertSharingAfun1 config alyt' (noStableSharing : aenv') reprB afun2) - (avarsIn AST.OpenAcc $ value weakenId) - in AST.Alet lhs (AST.OpenAcc boundAcc) (AST.OpenAcc bodyAcc) + Manifest as -> AST.Manifest (cvtA as) Aforeign repr ff afun acc -> AST.Aforeign repr ff (convertSmartAfun1 config (Smart.arraysR acc) afun) (cvtA acc) @@ -338,6 +322,8 @@ convertSharingAcc config alyt aenv (ScopedAcc lams (AccSharing _ preAcc)) Aprj ix a -> let AST.OpenAcc a' = cvtAprj ix a in a' Atrace msg acc1 acc2 -> AST.Atrace msg (cvtA acc1) (cvtA acc2) + Aassert msg cond acc -> AST.Aassert msg (cvtE cond) (cvtA acc) + Aassume cond acc -> AST.Aassume (cvtE cond) (cvtA acc) Use repr array -> AST.Use repr array Unit tp e -> AST.Unit tp (cvtE e) Generate repr@(ArrayR shr _) sh f @@ -352,8 +338,8 @@ convertSharingAcc config alyt aenv (ScopedAcc lams (AccSharing _ preAcc)) FoldSeg i tp f e acc1 acc2 -> AST.FoldSeg i (cvtF2 tp tp f) (cvtE <$> e) (cvtA acc1) (cvtA acc2) Scan d tp f e acc -> AST.Scan d (cvtF2 tp tp f) (cvtE <$> e) (cvtA acc) Scan' d tp f e acc -> AST.Scan' d (cvtF2 tp tp f) (cvtE e) (cvtA acc) - Permute (ArrayR shr tp) f dftAcc perm acc - -> AST.Permute (cvtF2 tp tp f) (cvtA dftAcc) (cvtF1 (shapeType shr) perm) (cvtA acc) + Permute (ArrayR _shr tp) f dftAcc acc + -> AST.Permute (cvtF2 tp tp <$> f) (cvtA dftAcc) (cvtA acc) Backpermute shr newDim perm acc -> AST.Backpermute shr (cvtE newDim) (cvtF1 (shapeType shr) perm) (cvtA acc) Stencil stencil tp f boundary acc @@ -765,15 +751,17 @@ convertSharingExp config lyt alyt env aenv exp@(ScopedExp lams _) = cvt exp FromIndex shr sh e -> AST.FromIndex shr (cvt sh) (cvt e) Case e rhs -> cvtCase (cvt e) (over (mapped . _2) cvt rhs) Cond e1 e2 e3 -> AST.Cond (cvt e1) (cvt e2) (cvt e3) + Select e1 e2 e3 -> AST.Select (cvt e1) (cvt e2) (cvt e3) While tp p it i -> AST.While (cvtFun1 tp p) (cvtFun1 tp it) (cvt i) - PrimConst c -> AST.PrimConst c PrimApp f e -> cvtPrimFun f (cvt e) - Index _ a e -> AST.Index (cvtAvar a) (cvt e) - LinearIndex _ a i -> AST.LinearIndex (cvtAvar a) (cvt i) - Shape _ a -> AST.Shape (cvtAvar a) + Index _ a e -> AST.ArrayInstr (AST.Index $ cvtAvar a) (cvt e) + LinearIndex _ a i -> AST.ArrayInstr (AST.LinearIndex $ cvtAvar a) (cvt i) + Shape _ a -> AST.ArrayInstr (AST.Shape $ cvtAvar a) AST.Nil ShapeSize shr e -> AST.ShapeSize shr (cvt e) - Foreign repr ff f e -> AST.Foreign repr ff (convertSmartFun config (typeR e) f) (cvt e) + Foreign repr ff f e -> AST.Foreign repr ff (cvtNoArrayInstr $ convertSmartFun config (typeR e) f) (cvt e) Coerce t1 t2 e -> AST.Coerce t1 t2 (cvt e) + Assert msg g e -> AST.Assert msg (cvt g) (cvt e) + Assume g e -> AST.Assume (cvt g) (cvt e) cvtPrj :: forall a b c env1 aenv1. PairIdx (a, b) c -> AST.OpenExp env1 aenv1 (a, b) -> AST.OpenExp env1 aenv1 c cvtPrj PairIdxLeft (AST.Pair a _) = a @@ -799,6 +787,14 @@ convertSharingExp config lyt alyt env aenv exp@(ScopedExp lams _) = cvt exp in Lam lhs $ Body $ convertSharingExp config lyt' alyt env' aenv body + cvtNoArrayInstr :: AST.OpenFun env' () f -> AST.PreOpenFun NoArrayInstr env' f + cvtNoArrayInstr = rebuildArrayInstr k + where + k :: AST.ArrayInstr () (s -> u) -> AST.PreOpenExp NoArrayInstr env' s -> AST.PreOpenExp NoArrayInstr env' u + k (AST.Shape (Var _ (VoidIdx void))) = void + k (AST.Index (Var _ (VoidIdx void))) = void + k (AST.LinearIndex (Var _ (VoidIdx void))) = void + -- Push primitive function applications down through let bindings so that -- they are adjacent to their arguments. It looks a bit nicer this way. -- @@ -1489,13 +1485,9 @@ makeOccMapSharingAcc config accOccMap = traverseAcc reconstruct $ case pacc of Atag repr i -> return (Atag repr i, 0) -- height is 0! - Pipe repr1 repr2 repr3 afun1 afun2 acc - -> do - (afun1', h1) <- traverseAfun1 lvl repr1 afun1 - (afun2', h2) <- traverseAfun1 lvl repr2 afun2 - (acc', h3) <- traverseAcc lvl acc - return (Pipe repr1 repr2 repr3 afun1' afun2' acc' - , h1 `max` h2 `max` h3 + 1) + Manifest as -> do + (as', h) <- traverseAcc lvl as + return (Manifest as', h) Aforeign repr ff afun acc -> travA (Aforeign repr ff afun) acc Acond e acc1 acc2 -> do (e' , h1) <- traverseExp lvl e @@ -1520,6 +1512,14 @@ makeOccMapSharingAcc config accOccMap = traverseAcc (a', h1) <- traverseAcc lvl acc1 (b', h2) <- traverseAcc lvl acc2 return (Atrace msg a' b', h1 `max` h2 + 1) + Aassert msg cond acc -> do + (cond', h1) <- traverseExp lvl cond + (acc', h2) <- traverseAcc lvl acc + return (Aassert msg cond' acc', h1 `max` h2 + 1) + Aassume cond acc -> do + (cond', h1) <- traverseExp lvl cond + (acc', h2) <- traverseAcc lvl acc + return (Aassume cond' acc', h1 `max` h2 + 1) Use repr arr -> return (Use repr arr, 1) Unit tp e -> do (e', h) <- traverseExp lvl e @@ -1548,14 +1548,19 @@ makeOccMapSharingAcc config accOccMap = traverseAcc h1 `max` h2 `max` h3 `max` h4 + 1) Scan d tp f e acc -> travF2MEA (Scan d tp) tp tp f e acc Scan' d tp f e acc -> travF2EA (Scan' d tp) tp tp f e acc - Permute repr@(ArrayR shr tp) c acc1 p acc2 + Permute repr@(ArrayR _shr tp) (Just c) acc1 acc2 -> do (c' , h1) <- traverseFun2 lvl tp tp c - (p' , h2) <- traverseFun1 lvl (shapeType shr) p - (acc1', h3) <- traverseAcc lvl acc1 - (acc2', h4) <- traverseAcc lvl acc2 - return (Permute repr c' acc1' p' acc2', - h1 `max` h2 `max` h3 `max` h4 + 1) + (acc1', h2) <- traverseAcc lvl acc1 + (acc2', h3) <- traverseAcc lvl acc2 + return (Permute repr (Just c') acc1' acc2', + h1 `max` h2 `max` h3 + 1) + Permute repr@(ArrayR _shr _tp) Nothing acc1 acc2 + -> do + (acc1', h2) <- traverseAcc lvl acc1 + (acc2', h3) <- traverseAcc lvl acc2 + return (Permute repr Nothing acc1' acc2', + h2 `max` h3 + 1) Backpermute shr e p acc -> do (e' , h1) <- traverseExp lvl e (p' , h2) <- traverseFun1 lvl (shapeType shr) p @@ -1854,12 +1859,12 @@ makeOccMapSharingExp config accOccMap expOccMap = travE (rhs', h2) <- unzip <$> sequence [ travE1 (t,) c | (t,c) <- rhs ] return (Case e' rhs', h1 `max` maximum h2 + 1) Cond e1 e2 e3 -> travE3 Cond e1 e2 e3 + Select e1 e2 e3 -> travE3 Select e1 e2 e3 While t p iter init -> do (p' , h1) <- traverseFun1 lvl t p (iter', h2) <- traverseFun1 lvl t iter (init', h3) <- travE lvl init return (While t p' iter' init', h1 `max` h2 `max` h3 + 1) - PrimConst c -> return (PrimConst c, 1) PrimApp p e -> travE1 (PrimApp p) e Index tp a e -> travAE (Index tp) a e LinearIndex tp a i -> travAE (LinearIndex tp) a i @@ -1869,6 +1874,8 @@ makeOccMapSharingExp config accOccMap expOccMap = travE (e', h) <- travE lvl e return (Foreign tp ff f e', h+1) Coerce t1 t2 e -> travE1 (Coerce t1 t2) e + Assert msg g e -> travE2 (Assert msg) g e + Assume g e -> travE2 Assume g e where traverseAcc :: HasCallStack => Level -> SmartAcc arrs -> IO (UnscopedAcc arrs, Int) @@ -2341,14 +2348,9 @@ determineScopesSharingAcc config accOccMap = scopesAcc scopesAcc (UnscopedAcc _ (AccSharing sn pacc)) = case pacc of Atag tp i -> reconstruct (Atag tp i) noNodeCounts - Pipe repr1 repr2 repr3 afun1 afun2 acc - -> let - (afun1', accCount1) = scopesAfun1 afun1 - (afun2', accCount2) = scopesAfun1 afun2 - (acc', accCount3) = scopesAcc acc - in - reconstruct (Pipe repr1 repr2 repr3 afun1' afun2' acc') - (accCount1 +++ accCount2 +++ accCount3) + Manifest as -> let + (as', accCount) = scopesAcc as + in reconstruct (Manifest as') accCount Aforeign r ff afun acc -> let (acc', accCount) = scopesAcc acc @@ -2379,6 +2381,16 @@ determineScopesSharingAcc config accOccMap = scopesAcc reconstruct (Apair a1' a2') (accCount1 +++ accCount2) Aprj ix a -> travA (Aprj ix) a + Aassert msg cond acc -> let + (cond', accCount1) = scopesExp cond + (acc', accCount2) = scopesAcc acc + in + reconstruct (Aassert msg cond' acc') (accCount1 +++ accCount2) + Aassume cond acc -> let + (cond', accCount1) = scopesExp cond + (acc', accCount2) = scopesAcc acc + in + reconstruct (Aassume cond' acc') (accCount1 +++ accCount2) Atrace msg a1 a2 -> let (a1', accCount1) = scopesAcc a1 (a2', accCount2) = scopesAcc a2 @@ -2415,15 +2427,21 @@ determineScopesSharingAcc config accOccMap = scopesAcc (accCount1 +++ accCount2 +++ accCount3 +++ accCount4) Scan d tp f z acc -> travF2MEA (Scan d tp) f z acc Scan' d tp f z acc -> travF2EA (Scan' d tp) f z acc - Permute repr fc acc1 fp acc2 + Permute repr (Just fc) acc1 acc2 -> let (fc' , accCount1) = scopesFun2 fc (acc1', accCount2) = scopesAcc acc1 - (fp' , accCount3) = scopesFun1 fp - (acc2', accCount4) = scopesAcc acc2 + (acc2', accCount3) = scopesAcc acc2 in - reconstruct (Permute repr fc' acc1' fp' acc2') - (accCount1 +++ accCount2 +++ accCount3 +++ accCount4) + reconstruct (Permute repr (Just fc') acc1' acc2') + (accCount1 +++ accCount2 +++ accCount3) + Permute repr Nothing acc1 acc2 + -> let + (acc1', accCount2) = scopesAcc acc1 + (acc2', accCount3) = scopesAcc acc2 + in + reconstruct (Permute repr Nothing acc1' acc2') + (accCount2 +++ accCount3) Backpermute shr sh fp acc -> let (sh' , accCount1) = scopesExp sh @@ -2615,12 +2633,15 @@ determineScopesSharingAcc config accOccMap = scopesAcc :: HasCallStack => (SmartAcc a1 -> UnscopedAcc a2) -> (SmartAcc a1 -> ScopedAcc a2, NodeCounts) - scopesAfun1 f = (const (ScopedAcc ssa body'), (counts', graph)) + scopesAfun1 f = + case scopesAcc body of + (ScopedAcc [] body', (counts, graph)) -> + let (freeCounts, counts') = partition isBoundHere counts + ssa = buildInitialEnvAcc fvs [sa | AccNodeCount sa _ <- freeCounts] + in (const (ScopedAcc ssa body'), (counts', graph)) + _ -> error "TODO WALL: NON-EXHAUSTIVE PATTERN MATCH" where - body@(UnscopedAcc fvs _) = f undefined - (ScopedAcc [] body', (counts,graph)) = scopesAcc body - (freeCounts, counts') = partition isBoundHere counts - ssa = buildInitialEnvAcc fvs [sa | AccNodeCount sa _ <- freeCounts] + body@(UnscopedAcc fvs _) = f undefined isBoundHere (AccNodeCount (StableSharingAcc _ (AccSharing _ (Atag _ i))) _) = i `elem` fvs isBoundHere _ = False @@ -2692,15 +2713,16 @@ determineScopesExp -> OccMap SmartAcc -> RootExp t -> (ScopedExp t, NodeCounts) -- Root (closed) expression plus Acc node counts -determineScopesExp config accOccMap (RootExp expOccMap exp@(UnscopedExp fvs _)) - = let - (ScopedExp [] expWithScopes, (nodeCounts,graph)) = determineScopesSharingExp config accOccMap expOccMap exp - (expCounts, accCounts) = partition isExpNodeCount nodeCounts - - isExpNodeCount ExpNodeCount{} = True - isExpNodeCount _ = False - in - (ScopedExp (buildInitialEnvExp fvs [se | ExpNodeCount se _ <- expCounts]) expWithScopes, cleanCounts (accCounts,graph)) +determineScopesExp config accOccMap (RootExp expOccMap exp@(UnscopedExp fvs _)) = + case determineScopesSharingExp config accOccMap expOccMap exp of + (ScopedExp [] expWithScopes, (nodeCounts,graph)) -> + let (expCounts, accCounts) = partition isExpNodeCount nodeCounts + in (ScopedExp (buildInitialEnvExp fvs [se | ExpNodeCount se _ <- expCounts]) expWithScopes, cleanCounts (accCounts,graph)) + _ -> error "TODO WALL: NON-EXHAUSTIVE PATTERN MATCH" + where + isExpNodeCount :: NodeCount -> Bool + isExpNodeCount ExpNodeCount{} = True + isExpNodeCount _ = False determineScopesSharingExp @@ -2722,12 +2744,15 @@ determineScopesSharingExp config accOccMap expOccMap = scopesExp :: HasCallStack => (SmartExp a -> UnscopedExp b) -> (SmartExp a -> ScopedExp b, NodeCounts) - scopesFun1 f = tracePure (bformat ("LAMBDA " % list formatStableSharingExp) ssa) (bformat (list formatNodeCount) counts) (const (ScopedExp ssa body'), (counts',graph)) + scopesFun1 f = + case scopesExp body of + (ScopedExp [] body', (counts, graph)) -> + let (freeCounts, counts') = partition isBoundHere counts + ssa = buildInitialEnvExp fvs [se | ExpNodeCount se _ <- freeCounts] + in tracePure (bformat ("LAMBDA " % list formatStableSharingExp) ssa) (bformat (list formatNodeCount) counts) (const (ScopedExp ssa body'), (counts',graph)) + _ -> error "TODO WALL: NON-EXHAUSTIVE PATTERN MATCH" where body@(UnscopedExp fvs _) = f undefined - (ScopedExp [] body', (counts, graph)) = scopesExp body - (freeCounts, counts') = partition isBoundHere counts - ssa = buildInitialEnvExp fvs [se | ExpNodeCount se _ <- freeCounts] isBoundHere (ExpNodeCount (StableSharingExp _ (ExpSharing _ (Tag _ i))) _) = i `elem` fvs isBoundHere _ = False @@ -2759,11 +2784,11 @@ determineScopesSharingExp config accOccMap expOccMap = scopesExp (rhs', accCount2) = unzip [ ((t,c'), counts)| (t,c) <- rhs, let (c', counts) = scopesExp c ] in reconstruct (Case e' rhs') (foldr (+++) accCount1 accCount2) Cond e1 e2 e3 -> travE3 Cond e1 e2 e3 + Select e1 e2 e3 -> travE3 Select e1 e2 e3 While tp p it i -> let (p' , accCount1) = scopesFun1 p (it', accCount2) = scopesFun1 it (i' , accCount3) = scopesExp i in reconstruct (While tp p' it' i') (accCount1 +++ accCount2 +++ accCount3) - PrimConst c -> reconstruct (PrimConst c) noNodeCounts PrimApp p e -> travE1 (PrimApp p) e Index tp a e -> travAE (Index tp) a e LinearIndex tp a e -> travAE (LinearIndex tp) a e @@ -2771,6 +2796,9 @@ determineScopesSharingExp config accOccMap expOccMap = scopesExp ShapeSize shr e -> travE1 (ShapeSize shr) e Foreign tp ff f e -> travE1 (Foreign tp ff f) e Coerce t1 t2 e -> travE1 (Coerce t1 t2) e + Assert msg g e -> travE2 (Assert msg) g e + Assume g e -> travE2 Assume g e + where travE1 :: HasCallStack => (ScopedExp a -> PreSmartExp ScopedAcc ScopedExp t) diff --git a/src/Data/Array/Accelerate/Trafo/Shrink.hs b/src/Data/Array/Accelerate/Trafo/Shrink.hs index 11c8dde21..972f30349 100644 --- a/src/Data/Array/Accelerate/Trafo/Shrink.hs +++ b/src/Data/Array/Accelerate/Trafo/Shrink.hs @@ -1,14 +1,3 @@ -{-# LANGUAGE TupleSections #-} -{-# LANGUAGE CPP #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE PatternGuards #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE ViewPatterns #-} -- | -- Module : Data.Array.Accelerate.Trafo.Shrink -- Copyright : [2012..2020] The Accelerate Team @@ -49,300 +38,11 @@ import Data.Array.Accelerate.AST.Idx import Data.Array.Accelerate.AST.LeftHandSide import Data.Array.Accelerate.AST.Var import Data.Array.Accelerate.Analysis.Match -import Data.Array.Accelerate.Error -import Data.Array.Accelerate.Representation.Type -import Data.Array.Accelerate.Trafo.Substitution +import Data.Array.Accelerate.Trafo.Exp.Shrink -import qualified Data.Array.Accelerate.Debug.Internal.Stats as Stats - -import Control.Applicative hiding ( Const ) -import Data.Maybe ( isJust ) -import Data.Monoid -import Data.Semigroup import Prelude hiding ( exp, seq ) -data VarsRange env = - VarsRange !(Exists (Idx env)) -- rightmost variable - {-# UNPACK #-} !Int -- count - !(Maybe RangeTuple) -- tuple - -data RangeTuple - = RTNil - | RTSingle - | RTPair !RangeTuple !RangeTuple - -lhsVarsRange :: LeftHandSide s v env env' -> Either (env :~: env') (VarsRange env') -lhsVarsRange lhs = case rightIx lhs of - Left eq -> Left eq - Right ix -> let (n, rt) = go lhs - in Right $ VarsRange ix n rt - where - rightIx :: LeftHandSide s v env env' -> Either (env :~: env') (Exists (Idx env')) - rightIx (LeftHandSideWildcard _) = Left Refl - rightIx (LeftHandSideSingle _) = Right $ Exists ZeroIdx - rightIx (LeftHandSidePair l1 l2) = case rightIx l2 of - Right ix -> Right ix - Left Refl -> rightIx l1 - - go :: LeftHandSide s v env env' -> (Int, Maybe (RangeTuple)) - go (LeftHandSideWildcard TupRunit) = (0, Just RTNil) - go (LeftHandSideWildcard _) = (0, Nothing) - go (LeftHandSideSingle _) = (1, Just RTSingle) - go (LeftHandSidePair l1 l2) = (n1 + n2, RTPair <$> t1 <*> t2) - where - (n1, t1) = go l1 - (n2, t2) = go l2 - -weakenVarsRange :: LeftHandSide s v env env' -> VarsRange env -> VarsRange env' -weakenVarsRange lhs (VarsRange ix n t) = VarsRange (go lhs ix) n t - where - go :: LeftHandSide s v env env' -> Exists (Idx env) -> Exists (Idx env') - go (LeftHandSideWildcard _) ix' = ix' - go (LeftHandSideSingle _) (Exists ix') = Exists (SuccIdx ix') - go (LeftHandSidePair l1 l2) ix' = go l2 $ go l1 ix' - -matchEVarsRange :: VarsRange env -> OpenExp env aenv t -> Bool -matchEVarsRange (VarsRange (Exists first) _ (Just rt)) expr = isJust $ go (idxToInt first) rt expr - where - go :: Int -> RangeTuple -> OpenExp env aenv t -> Maybe Int - go i RTNil Nil = Just i - go i RTSingle (Evar (Var _ ix)) - | checkIdx i ix = Just (i + 1) - go i (RTPair t1 t2) (Pair e1 e2) - | Just i' <- go i t2 e2 = go i' t1 e1 - go _ _ _ = Nothing - - checkIdx :: Int -> Idx env t -> Bool - checkIdx 0 ZeroIdx = True - checkIdx i (SuccIdx ix) = checkIdx (i - 1) ix - checkIdx _ _ = False -matchEVarsRange _ _ = False - -varInRange :: VarsRange env -> Var s env t -> Maybe Usages -varInRange (VarsRange (Exists rangeIx) n _) (Var _ varIx) = case go rangeIx varIx of - Nothing -> Nothing - Just j -> Just $ replicate j False ++ [True] ++ replicate (n - j - 1) False - where - -- `go ix ix'` checks whether ix <= ix' with recursion, and then checks - -- whether ix' < ix + n in go'. Returns a Just if both checks - -- are successful, containing an integer j such that ix + j = ix'. - go :: Idx env u -> Idx env t -> Maybe Int - go (SuccIdx ix) (SuccIdx ix') = go ix ix' - go ZeroIdx ix' = go' ix' 0 - go _ ZeroIdx = Nothing - - go' :: Idx env t -> Int -> Maybe Int - go' _ j | j >= n = Nothing - go' ZeroIdx j = Just j - go' (SuccIdx ix') j = go' ix' (j + 1) - --- Describes how often the variables defined in a LHS are used together. -data Count - = Impossible !Usages - -- Cannot inline this definition. This happens when the definition - -- declares multiple variables (the right hand side returns a tuple) - -- and the variables are used seperately. - | Infinity - -- The variable is used in a loop. Inlining should only proceed if - -- the computation is cheap. - | Finite {-# UNPACK #-} !Int - -type Usages = [Bool] -- Per variable a Boolean denoting whether that variable is used. - -instance Semigroup Count where - Impossible u1 <> Impossible u2 = Impossible $ zipWith (||) u1 u2 - Impossible u <> Finite 0 = Impossible u - Finite 0 <> Impossible u = Impossible u - Impossible u <> _ = Impossible $ map (const True) u - _ <> Impossible u = Impossible $ map (const True) u - Infinity <> _ = Infinity - _ <> Infinity = Infinity - Finite a <> Finite b = Finite $ a + b - -instance Monoid Count where - mempty = Finite 0 - -loopCount :: Count -> Count -loopCount (Finite n) | n > 0 = Infinity -loopCount c = c - -shrinkLhs - :: HasCallStack - => Count - -> LeftHandSide s t env1 env2 - -> Maybe (Exists (LeftHandSide s t env1)) -shrinkLhs _ (LeftHandSideWildcard _) = Nothing -- We cannot shrink this -shrinkLhs (Finite 0) lhs = Just $ Exists $ LeftHandSideWildcard $ lhsToTupR lhs -- LHS isn't used at all, replace with a wildcard -shrinkLhs (Impossible usages) lhs = case go usages lhs of - (True , [], lhs') -> Just lhs' - (False, [], _ ) -> Nothing -- No variables were dropped. Thus lhs == lhs'. - _ -> internalError "Mismatch in length of usages array and LHS" - where - go :: HasCallStack => Usages -> LeftHandSide s t env1 env2 -> (Bool, Usages, Exists (LeftHandSide s t env1)) - go us (LeftHandSideWildcard tp) = (False, us, Exists $ LeftHandSideWildcard tp) - go (True : us) (LeftHandSideSingle tp) = (False, us, Exists $ LeftHandSideSingle tp) - go (False : us) (LeftHandSideSingle tp) = (True , us, Exists $ LeftHandSideWildcard $ TupRsingle tp) - go us (LeftHandSidePair l1 l2) - | (c2, us' , Exists l2') <- go us l2 - , (c1, us'', Exists l1') <- go us' l1 - , Exists l2'' <- rebuildLHS l2' - = let - lhs' - | LeftHandSideWildcard t1 <- l1' - , LeftHandSideWildcard t2 <- l2'' = LeftHandSideWildcard $ TupRpair t1 t2 - | otherwise = LeftHandSidePair l1' l2'' - in - (c1 || c2, us'', Exists lhs') - go _ _ = internalError "Empty array, mismatch in length of usages array and LHS" -shrinkLhs _ _ = Nothing - --- The first LHS should be 'larger' than the second, eg the second may have --- a wildcard if the first LHS does bind variables there, but not the other --- way around. --- -strengthenShrunkLHS - :: HasCallStack - => LeftHandSide s t env1 env2 - -> LeftHandSide s t env1' env2' - -> env1 :?> env1' - -> env2 :?> env2' -strengthenShrunkLHS (LeftHandSideWildcard _) (LeftHandSideWildcard _) k = k -strengthenShrunkLHS (LeftHandSideSingle _) (LeftHandSideSingle _) k = \ix -> case ix of - ZeroIdx -> Just ZeroIdx - SuccIdx ix' -> SuccIdx <$> k ix' -strengthenShrunkLHS (LeftHandSidePair lA hA) (LeftHandSidePair lB hB) k = strengthenShrunkLHS hA hB $ strengthenShrunkLHS lA lB k -strengthenShrunkLHS (LeftHandSideSingle _) (LeftHandSideWildcard _) k = \ix -> case ix of - ZeroIdx -> Nothing - SuccIdx ix' -> k ix' -strengthenShrunkLHS (LeftHandSidePair l h) (LeftHandSideWildcard t) k = strengthenShrunkLHS h (LeftHandSideWildcard t2) $ strengthenShrunkLHS l (LeftHandSideWildcard t1) k - where - (t1, t2) = case t of TupRpair t1' t2' -> (t1', t2') - _ -> internalError "Pair type was not TupRpair" -strengthenShrunkLHS (LeftHandSideWildcard _) _ _ = internalError "Second LHS defines more variables" -strengthenShrunkLHS _ _ _ = internalError "Mismatch LHS single with LHS pair" - - --- Shrinking --- ========= - --- The shrinking substitution for scalar expressions. This is a restricted --- instance of beta-reduction to cases where the bound variable is used zero --- (dead-code elimination) or one (linear inlining) times. --- -shrinkExp :: HasCallStack => OpenExp env aenv t -> (Bool, OpenExp env aenv t) -shrinkExp = Stats.substitution "shrinkE" . first getAny . shrinkE - where - -- If the bound variable is used at most this many times, it will be inlined - -- into the body. In cases where it is not used at all, this is equivalent - -- to dead-code elimination. - -- - lIMIT :: Int - lIMIT = 1 - - cheap :: OpenExp env aenv t -> Bool - cheap (Evar _) = True - cheap (Pair e1 e2) = cheap e1 && cheap e2 - cheap Nil = True - cheap Const{} = True - cheap PrimConst{} = True - cheap Undef{} = True - cheap (Coerce _ _ e) = cheap e - cheap _ = False - - shrinkE :: HasCallStack => OpenExp env aenv t -> (Any, OpenExp env aenv t) - shrinkE exp = case exp of - Let (LeftHandSideSingle _) bnd@Evar{} body -> Stats.inline "Var" . yes $ shrinkE (inline body bnd) - Let lhs bnd body - | shouldInline -> case inlineVars lhs (snd body') (snd bnd') of - Just inlined -> Stats.betaReduce msg . yes $ shrinkE inlined - _ -> internalError "Unexpected failure while trying to inline some expression." - | Just (Exists lhs') <- shrinkLhs count lhs -> case strengthenE (strengthenShrunkLHS lhs lhs' Just) (snd body') of - Just body'' -> (Any True, Let lhs' (snd bnd') body'') - Nothing -> internalError "Unexpected failure in strenthenE. Variable was analysed to be unused in usesOfExp, but appeared to be used in strenthenE." - | otherwise -> Let lhs <$> bnd' <*> body' - where - shouldInline = case count of - Finite 0 -> False -- Handled by shrinkLhs - Finite n -> n <= lIMIT || cheap (snd bnd') - Infinity -> cheap (snd bnd') - Impossible _ -> False - - bnd' = shrinkE bnd - body' = shrinkE body - - -- If the lhs includes non-trivial wildcards (the last field of range is Nothing), - -- then we cannot inline the binding. We can only check which variables are not used, - -- to detect unused variables. - -- - -- If the lhs does not include non-trivial wildcards (the last field of range is a Just), - -- we can both analyse whether we can inline the binding, and check which variables are - -- not used, to detect unused variables. - -- - count = case lhsVarsRange lhs of - Left _ -> Finite 0 - Right range -> usesOfExp range (snd body') - - msg = case count of - Finite 0 -> "dead exp" - _ -> "inline exp" -- forced inlining when lIMIT > 1 - -- - Evar v -> pure (Evar v) - Const t c -> pure (Const t c) - Undef t -> pure (Undef t) - Nil -> pure Nil - Pair x y -> Pair <$> shrinkE x <*> shrinkE y - VecPack vec e -> VecPack vec <$> shrinkE e - VecUnpack vec e -> VecUnpack vec <$> shrinkE e - IndexSlice x ix sh -> IndexSlice x <$> shrinkE ix <*> shrinkE sh - IndexFull x ix sl -> IndexFull x <$> shrinkE ix <*> shrinkE sl - ToIndex shr sh ix -> ToIndex shr <$> shrinkE sh <*> shrinkE ix - FromIndex shr sh i -> FromIndex shr <$> shrinkE sh <*> shrinkE i - Case e rhs def -> Case <$> shrinkE e <*> sequenceA [ (t,) <$> shrinkE c | (t,c) <- rhs ] <*> shrinkMaybeE def - Cond p t e -> Cond <$> shrinkE p <*> shrinkE t <*> shrinkE e - While p f x -> While <$> shrinkF p <*> shrinkF f <*> shrinkE x - PrimConst c -> pure (PrimConst c) - PrimApp f x -> PrimApp f <$> shrinkE x - Index a sh -> Index a <$> shrinkE sh - LinearIndex a i -> LinearIndex a <$> shrinkE i - Shape a -> pure (Shape a) - ShapeSize shr sh -> ShapeSize shr <$> shrinkE sh - Foreign repr ff f e -> Foreign repr ff <$> shrinkF f <*> shrinkE e - Coerce t1 t2 e -> Coerce t1 t2 <$> shrinkE e - - shrinkF :: HasCallStack => OpenFun env aenv t -> (Any, OpenFun env aenv t) - shrinkF = first Any . shrinkFun - - shrinkMaybeE :: HasCallStack => Maybe (OpenExp env aenv t) -> (Any, Maybe (OpenExp env aenv t)) - shrinkMaybeE Nothing = pure Nothing - shrinkMaybeE (Just e) = Just <$> shrinkE e - - first :: (a -> a') -> (a,b) -> (a',b) - first f (x,y) = (f x, y) - - yes :: (Any, x) -> (Any, x) - yes (_, x) = (Any True, x) - -shrinkFun :: HasCallStack => OpenFun env aenv f -> (Bool, OpenFun env aenv f) -shrinkFun (Lam lhs f) = case lhsVarsRange lhs of - Left Refl -> - let b' = case lhs of - LeftHandSideWildcard TupRunit -> b - _ -> True - in (b', Lam (LeftHandSideWildcard $ lhsToTupR lhs) f') - Right range -> - let - count = usesOfFun range f - in case shrinkLhs count lhs of - Just (Exists lhs') -> case strengthenE (strengthenShrunkLHS lhs lhs' Just) f' of - Just f'' -> (True, Lam lhs' f'') - Nothing -> internalError "Unexpected failure in strenthenE. Variable was analysed to be unused in usesOfExp, but appeared to be used in strenthenE." - Nothing -> (b, Lam lhs f') - where - (b, f') = shrinkFun f - -shrinkFun (Body b) = Body <$> shrinkExp b - -- The shrinking substitution for array computations. This is further limited to -- dead-code elimination only, primarily because linear inlining may inline -- array computations into scalar expressions, which is generally not desirable. @@ -475,46 +175,6 @@ shrinkPreAcc shrinkAcc reduceAcc = Stats.substitution "shrinkA" shrinkA -- Occurrence Counting -- =================== --- Count the number of occurrences an in-scope scalar expression bound at the --- given variable index recursively in a term. --- -usesOfExp :: forall env aenv t. VarsRange env -> OpenExp env aenv t -> Count -usesOfExp range = countE - where - countE :: OpenExp env aenv e -> Count - countE exp | matchEVarsRange range exp = Finite 1 - countE exp = case exp of - Evar v -> case varInRange range v of - Just cs -> Impossible cs - Nothing -> Finite 0 - -- - Let lhs bnd body -> countE bnd <> usesOfExp (weakenVarsRange lhs range) body - Const _ _ -> Finite 0 - Undef _ -> Finite 0 - Nil -> Finite 0 - Pair e1 e2 -> countE e1 <> countE e2 - VecPack _ e -> countE e - VecUnpack _ e -> countE e - IndexSlice _ ix sh -> countE ix <> countE sh - IndexFull _ ix sl -> countE ix <> countE sl - FromIndex _ sh i -> countE sh <> countE i - ToIndex _ sh e -> countE sh <> countE e - Case e rhs def -> countE e <> mconcat [ countE c | (_,c) <- rhs ] <> maybe (Finite 0) countE def - Cond p t e -> countE p <> countE t <> countE e - While p f x -> countE x <> loopCount (usesOfFun range p) <> loopCount (usesOfFun range f) - PrimConst _ -> Finite 0 - PrimApp _ x -> countE x - Index _ sh -> countE sh - LinearIndex _ i -> countE i - Shape _ -> Finite 0 - ShapeSize _ sh -> countE sh - Foreign _ _ _ e -> countE e - Coerce _ _ e -> countE e - -usesOfFun :: VarsRange env -> OpenFun env aenv f -> Count -usesOfFun range (Lam lhs f) = usesOfFun (weakenVarsRange lhs range) f -usesOfFun range (Body b) = usesOfExp range b - -- Count the number of occurrences of the array term bound at the given -- environment index. If the first argument is 'True' then it includes in the -- total uses of the variable for 'Shape' information, otherwise not. @@ -545,7 +205,9 @@ usesOfPreAcc withShape countAcc idx = count Apair a1 a2 -> countA a1 + countA a2 Anil -> 0 Atrace _ a1 a2 -> countA a1 + countA a2 - Apply _ f a -> countAF f idx + countA a + Manifest as -> countA as + Aassert _ c as -> countE c + countA as + Aassume c as -> countE c + countA as Aforeign _ _ _ a -> countA a Acond p t e -> countE p + countA t + countA e -- Body and condition of the while loop may be evaluated multiple times. @@ -566,39 +228,14 @@ usesOfPreAcc withShape countAcc idx = count FoldSeg _ f z a s -> countF f + countME z + countA a + countA s Scan _ f z a -> countF f + countME z + countA a Scan' _ f z a -> countF f + countE z + countA a - Permute f1 a1 f2 a2 -> countF f1 + countA a1 + countF f2 + countA a2 + Permute f1 a1 a2 -> maybe 0 countF f1 + countA a1 + countA a2 Backpermute _ sh f a -> countE sh + countF f + countA a Stencil _ _ f _ a -> countF f + countA a Stencil2 _ _ _ f _ a1 _ a2 -> countF f + countA a1 + countA a2 -- Collect s -> countS s countE :: OpenExp env aenv e -> Int - countE exp = case exp of - Let _ bnd body -> countE bnd + countE body - Evar _ -> 0 - Const _ _ -> 0 - Undef _ -> 0 - Nil -> 0 - Pair x y -> countE x + countE y - VecPack _ e -> countE e - VecUnpack _ e -> countE e - IndexSlice _ ix sh -> countE ix + countE sh - IndexFull _ ix sl -> countE ix + countE sl - ToIndex _ sh ix -> countE sh + countE ix - FromIndex _ sh i -> countE sh + countE i - Case e rhs def -> countE e + sum [ countE c | (_,c) <- rhs ] + maybe 0 countE def - Cond p t e -> countE p + countE t + countE e - While p f x -> countF p + countF f + countE x - PrimConst _ -> 0 - PrimApp _ x -> countE x - Index a sh -> countAvar a + countE sh - LinearIndex a i -> countAvar a + countE i - ShapeSize _ sh -> countE sh - Shape a - | withShape -> countAvar a - | otherwise -> 0 - Foreign _ _ _ e -> countE e - Coerce _ _ e -> countE e + countE = sum . map countArrayInstr . arrayInstrsInExp countME :: Maybe (OpenExp env aenv e) -> Int countME = maybe 0 countE @@ -615,9 +252,13 @@ usesOfPreAcc withShape countAcc idx = count countAF (Alam lhs f) v = countAF f (weakenWithLHS lhs >:> v) countAF (Abody a) v = countAcc withShape v a - countF :: OpenFun env aenv f -> Int - countF (Lam _ f) = countF f - countF (Body b) = countE b + countArrayInstr :: Exists (ArrayInstr aenv) -> Int + countArrayInstr (Exists (Shape var)) = countAvar var + countArrayInstr (Exists (Index var)) = countAvar var + countArrayInstr (Exists (LinearIndex var)) = countAvar var + + countF :: OpenFun env aenv e -> Int + countF = sum . map countArrayInstr . arrayInstrsInFun {-- countS :: PreOpenSeq acc aenv senv arrs -> Int diff --git a/src/Data/Array/Accelerate/Trafo/SkipEnvironment.hs b/src/Data/Array/Accelerate/Trafo/SkipEnvironment.hs new file mode 100644 index 000000000..2f96d9746 --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/SkipEnvironment.hs @@ -0,0 +1,133 @@ +{-# OPTIONS_HADDOCK hide #-} +-- | +-- Module : Data.Array.Accelerate.AST.SkipEnvironment +-- Copyright : [2008..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.Trafo.SkipEnvironment + ( SEnv, SEnv'(..), SEnvValue(..), sempty, spush, spush', sprj, weakenSEnv1, weakenSEnv, + strengthenSEnv, sprjUpdate, senvTop, + joinSkips, + skipWeakenPartialEnv, + skipWeakenIdxSet, + skipStrengthenPartialEnv, + skipStrengthenIdxSet, + skipTakePartialEnv, skipTakeIdxSet, + ) where + +import Data.Array.Accelerate.AST.Environment +import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.IdxSet (IdxSet(..)) + +type SEnv f env = SEnv' f env env + +data SEnv' f env' env where + SEmpty :: SEnv' f env' () + SPush :: SEnv' f env' env -> f env' t -> SEnv' f env' (env, t) + SSkip :: Skip env2 env1 -> SEnv' f env1 env -> SEnv' f env2 env + +class SEnvValue f where + -- Instead of using (:>) in weaken, we use Skip and Skip' here. This way we + -- can weaken an IdxSet more efficiently. + weakenSEnvValue :: Skip env1 env2 -> f env2 t -> f env1 t + + strengthenSEnvValue :: SEnv f env1 -> Skip env1 env2 -> f env1 t -> f env2 t + +sempty :: SEnv f () +sempty = SEmpty + +spush :: SEnv f env -> f (env, t) t -> SEnv f (env, t) +spush env v = SPush (weakenSEnv1 env) v + +spush' :: SEnv' f env' env -> f env' t -> SEnv' f env' (env, t) +spush' = SPush + +-- We assume that doing weakenSEnvValue multiple times is cheaper than doing it +-- once and combining the Skips with chainSkip'. +sprj :: SEnvValue f => Idx envB t -> SEnv' f envA envB -> f envA t +sprj idx (SSkip s env) = weakenSEnvValue s $ sprj idx env +sprj ZeroIdx (SPush _ v) = v +sprj (SuccIdx idx) (SPush env _) = sprj idx env + +sprjUpdate :: forall f env t a. SEnvValue f => (f env t -> (f env t, a)) -> Idx env t -> SEnv f env -> (SEnv f env, a) +sprjUpdate f = go SkipNone + where + -- We must remove all SSkips until we arrive at the given index. + go :: Skip env env1 -> Idx env' t -> SEnv' f env1 env' -> (SEnv' f env env', a) + go skip idx (SSkip skip' env) = go (chainSkip skip skip') idx env + go skip ZeroIdx (SPush env v) = (SPush env' v', a) + where + env' = weakenSEnv skip env + (v', a) = f $ weakenSEnvValue skip v + go skip (SuccIdx idx) (SPush env v) = (SPush env' v', a) + where + (env', a) = go skip idx env + v' = weakenSEnvValue skip v + +senvTop :: SEnvValue f => SEnv' f env' (env, t) -> (SEnv' f env' env, f env' t) +senvTop (SPush env v) = (env, v) +senvTop (SSkip skip env) = (weakenSEnv skip env', weakenSEnvValue skip v) + where + (env', v) = senvTop env + +weakenSEnv1 :: SEnv' f env1 env -> SEnv' f (env1, t) env +weakenSEnv1 (SSkip s env) = SSkip (SkipSucc' s) env +weakenSEnv1 env = SSkip (SkipSucc' SkipNone) env + +weakenSEnv :: Skip env1 env2 -> SEnv' f env2 env -> SEnv' f env1 env +weakenSEnv SkipNone env = env +weakenSEnv s1 (SSkip s2 env) = SSkip (chainSkip s1 s2) env +weakenSEnv s1 env = SSkip s1 env + +strengthenSEnv :: SEnvValue f => Skip env1 env2 -> SEnv f env1 -> SEnv f env2 +strengthenSEnv skip initialEnv = strengthenSEnv' SkipNone initialEnv skip $ go skip initialEnv -- strengthenSEnv' SkipNone env skip env + where + go :: Skip env env' -> SEnv' f env1 env -> SEnv' f env1 env' + go SkipNone env = env + go (SkipSucc' s) (SPush env _) = go s env + go s (SSkip s' env) = SSkip s' $ go s env + +strengthenSEnv' :: SEnvValue f => Skip env0 env1 -> SEnv f env0 -> Skip env1 env2 -> SEnv' f env1 env -> SEnv' f env2 env +strengthenSEnv' _ _ SkipNone env = env +strengthenSEnv' _ _ _ SEmpty = SEmpty +strengthenSEnv' initialSkip initialEnv s (SSkip s' env) = case joinSkips s' s of + Left l -> SSkip l env + Right r -> strengthenSEnv' (chainSkip initialSkip s') initialEnv r env +strengthenSEnv' initialSkip initialEnv s (SPush env v) + = SPush (strengthenSEnv' initialSkip initialEnv s env) + $ strengthenSEnvValue initialEnv (chainSkip initialSkip s) + $ weakenSEnvValue initialSkip v + +-- Returns `Right SkipNone` if env2 ~ env3 +joinSkips :: Skip env1 env2 -> Skip env1 env3 -> Either (Skip env3 env2) (Skip env2 env3) +joinSkips SkipNone r = Right r +joinSkips l SkipNone = Left l +joinSkips (SkipSucc' l) (SkipSucc' r) = joinSkips l r + +skipWeakenPartialEnv :: Skip env1 env2 -> PartialEnv f env2 -> PartialEnv f env1 +skipWeakenPartialEnv SkipNone env = env +skipWeakenPartialEnv (SkipSucc s) env = skipWeakenPartialEnv s $ partialEnvSkip env + +skipWeakenIdxSet :: Skip env1 env2 -> IdxSet env2 -> IdxSet env1 +skipWeakenIdxSet s (IdxSet env) = IdxSet $ skipWeakenPartialEnv s env + +skipStrengthenPartialEnv :: Skip env1 env2 -> PartialEnv f env1 -> PartialEnv f env2 +skipStrengthenPartialEnv SkipNone env = env +skipStrengthenPartialEnv (SkipSucc s) env = partialEnvTail $ skipStrengthenPartialEnv s env + +skipStrengthenIdxSet :: Skip env1 env2 -> IdxSet env1 -> IdxSet env2 +skipStrengthenIdxSet s (IdxSet env) = IdxSet $ skipStrengthenPartialEnv s env + +skipTakePartialEnv :: Skip env1 env2 -> PartialEnv f env1 -> PartialEnv f env1 +skipTakePartialEnv SkipNone env = env +skipTakePartialEnv (SkipSucc' _) PEnd = PEnd +skipTakePartialEnv (SkipSucc' s) (PPush env v) = PPush (skipTakePartialEnv s env) v +skipTakePartialEnv (SkipSucc' s) (PNone env) = PNone (skipTakePartialEnv s env) + +skipTakeIdxSet :: Skip env1 env2 -> IdxSet env1 -> IdxSet env1 +skipTakeIdxSet s (IdxSet env) = IdxSet $ skipTakePartialEnv s env diff --git a/src/Data/Array/Accelerate/Trafo/Substitution.hs b/src/Data/Array/Accelerate/Trafo/Substitution.hs index e1aa1176b..805c9c2b6 100644 --- a/src/Data/Array/Accelerate/Trafo/Substitution.hs +++ b/src/Data/Array/Accelerate/Trafo/Substitution.hs @@ -1,17 +1,4 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE KindSignatures #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE PatternGuards #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TupleSections #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE TypeFamilies #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.Trafo.Substitution @@ -27,10 +14,10 @@ module Data.Array.Accelerate.Trafo.Substitution ( -- ** Renaming & Substitution inline, inlineVars, compose, - subTop, subAtop, + subTop, subAtop, apply1, -- ** Weakening - (:>), Sink(..), SinkExp(..), weakenVars, + (:>), Sink(..), Sink'(..), SinkExp(..), weakenVars, weakenArrayInstr, -- ** Strengthening (:?>), strengthen, strengthenE, @@ -53,14 +40,15 @@ import Data.Array.Accelerate.AST.Idx import Data.Array.Accelerate.AST.Environment import Data.Array.Accelerate.Analysis.Match import Data.Array.Accelerate.Error -import Data.Array.Accelerate.Representation.Type import Data.Array.Accelerate.Representation.Array +import Data.Array.Accelerate.Trafo.Exp.Substitution import qualified Data.Array.Accelerate.Debug.Internal.Stats as Stats import Data.Kind import Control.Applicative hiding ( Const ) import Control.Monad import Prelude hiding ( exp, seq ) +import Data.Functor.Identity -- NOTE: [Renaming and Substitution] @@ -87,178 +75,20 @@ import Prelude hiding ( exp -- able to rebuildPartial terms. In essence, the crucial functionality is to propagate -- a class of operations on variables that is closed under shifting. -- -infixr `compose` --- infixr `substitute` - -lhsFullVars :: forall s a env1 env2. LeftHandSide s a env1 env2 -> Maybe (Vars s env2 a) -lhsFullVars = fmap snd . go weakenId - where - go :: forall env env' b. (env' :> env2) -> LeftHandSide s b env env' -> Maybe (env :> env2, Vars s env2 b) - go k (LeftHandSideWildcard TupRunit) = Just (k, TupRunit) - go k (LeftHandSideSingle s) = Just $ (weakenSucc $ k, TupRsingle $ Var s $ k >:> ZeroIdx) - go k (LeftHandSidePair l1 l2) - | Just (k', v2) <- go k l2 - , Just (k'', v1) <- go k' l1 = Just (k'', TupRpair v1 v2) - go _ _ = Nothing - -bindingIsTrivial :: LeftHandSide s a env1 env2 -> Vars s env2 b -> Maybe (a :~: b) -bindingIsTrivial lhs vars - | Just lhsVars <- lhsFullVars lhs - , Just Refl <- matchVars vars lhsVars - = Just Refl -bindingIsTrivial _ _ = Nothing - -isIdentity :: OpenFun env aenv (a -> b) -> Maybe (a :~: b) -isIdentity (Lam lhs (Body (extractExpVars -> Just vars))) = bindingIsTrivial lhs vars -isIdentity _ = Nothing -- Detects whether the function is of the form \ix -> a ! ix isIdentityIndexing :: OpenFun env aenv (a -> b) -> Maybe (ArrayVar aenv (Array a b)) isIdentityIndexing (Lam lhs (Body body)) - | Index avar ix <- body + | ArrayInstr (Index avar) ix <- body , Just vars <- extractExpVars ix , Just Refl <- bindingIsTrivial lhs vars = Just avar isIdentityIndexing _ = Nothing --- | Replace the first variable with the given expression. The environment --- shrinks. --- -inline :: OpenExp (env, s) aenv t - -> OpenExp env aenv s - -> OpenExp env aenv t -inline f g = Stats.substitution "inline" $ rebuildE (subTop g) f - -inlineVars :: forall env env' aenv t1 t2. - ELeftHandSide t1 env env' - -> OpenExp env' aenv t2 - -> OpenExp env aenv t1 - -> Maybe (OpenExp env aenv t2) -inlineVars lhsBound expr bound - | Just vars <- lhsFullVars lhsBound = substitute (strengthenWithLHS lhsBound) weakenId vars expr - where - substitute - :: forall env1 env2 t. - env1 :?> env2 - -> env :> env2 - -> ExpVars env1 t1 - -> OpenExp env1 aenv t - -> Maybe (OpenExp env2 aenv t) - substitute _ k2 vars (extractExpVars -> Just vars') - | Just Refl <- matchVars vars vars' = Just $ weakenE k2 bound - substitute k1 k2 vars topExp = case topExp of - Let lhs e1 e2 - | Exists lhs' <- rebuildLHS lhs - -> Let lhs' <$> travE e1 <*> substitute (strengthenAfter lhs lhs' k1) (weakenWithLHS lhs' .> k2) (weakenWithLHS lhs `weakenVars` vars) e2 - Evar (Var t ix) -> Evar . Var t <$> k1 ix - Foreign tp asm f e1 -> Foreign tp asm f <$> travE e1 - Pair e1 e2 -> Pair <$> travE e1 <*> travE e2 - Nil -> Just Nil - VecPack vec e1 -> VecPack vec <$> travE e1 - VecUnpack vec e1 -> VecUnpack vec <$> travE e1 - IndexSlice si e1 e2 -> IndexSlice si <$> travE e1 <*> travE e2 - IndexFull si e1 e2 -> IndexFull si <$> travE e1 <*> travE e2 - ToIndex shr e1 e2 -> ToIndex shr <$> travE e1 <*> travE e2 - FromIndex shr e1 e2 -> FromIndex shr <$> travE e1 <*> travE e2 - Case e1 rhs def -> Case <$> travE e1 <*> mapM (\(t,c) -> (t,) <$> travE c) rhs <*> travMaybeE def - Cond e1 e2 e3 -> Cond <$> travE e1 <*> travE e2 <*> travE e3 - While f1 f2 e1 -> While <$> travF f1 <*> travF f2 <*> travE e1 - Const t c -> Just $ Const t c - PrimConst c -> Just $ PrimConst c - PrimApp p e1 -> PrimApp p <$> travE e1 - Index a e1 -> Index a <$> travE e1 - LinearIndex a e1 -> LinearIndex a <$> travE e1 - Shape a -> Just $ Shape a - ShapeSize shr e1 -> ShapeSize shr <$> travE e1 - Undef t -> Just $ Undef t - Coerce t1 t2 e1 -> Coerce t1 t2 <$> travE e1 - - where - travE :: OpenExp env1 aenv s -> Maybe (OpenExp env2 aenv s) - travE = substitute k1 k2 vars - - travF :: OpenFun env1 aenv s -> Maybe (OpenFun env2 aenv s) - travF = substituteF k1 k2 vars - - travMaybeE :: Maybe (OpenExp env1 aenv s) -> Maybe (Maybe (OpenExp env2 aenv s)) - travMaybeE Nothing = pure Nothing - travMaybeE (Just x) = Just <$> travE x - - substituteF :: forall env1 env2 t. - env1 :?> env2 - -> env :> env2 - -> ExpVars env1 t1 - -> OpenFun env1 aenv t - -> Maybe (OpenFun env2 aenv t) - substituteF k1 k2 vars (Body e) = Body <$> substitute k1 k2 vars e - substituteF k1 k2 vars (Lam lhs f) - | Exists lhs' <- rebuildLHS lhs = Lam lhs' <$> substituteF (strengthenAfter lhs lhs' k1) (weakenWithLHS lhs' .> k2) (weakenWithLHS lhs `weakenVars` vars) f - -inlineVars _ _ _ = Nothing - - --- | Replace an expression that uses the top environment variable with another. --- The result of the first is let bound into the second. --- -{- substitute' :: OpenExp (env, b) aenv c - -> OpenExp (env, a) aenv b - -> OpenExp (env, a) aenv c -substitute' f g - | Stats.substitution "substitute" False = undefined - | isIdentity f = g -- don't rebind an identity function - | isIdentity g = f - | otherwise = Let g $ rebuildE split f - where - split :: Idx (env,b) c -> OpenExp ((env,a),b) aenv c - split ZeroIdx = Var ZeroIdx - split (SuccIdx ix) = Var (SuccIdx (SuccIdx ix)) - -substitute :: LeftHandSide b env envb - -> OpenExp envb c - -> LeftHandSide a env enva - -> OpenExp enva b --} - --- | Composition of unary functions. --- -compose :: HasCallStack - => OpenFun env aenv (b -> c) - -> OpenFun env aenv (a -> b) - -> OpenFun env aenv (a -> c) -compose f@(Lam lhsB (Body c)) g@(Lam lhsA (Body b)) - | Stats.substitution "compose" False = undefined - | Just Refl <- isIdentity f = g -- don't rebind an identity function - | Just Refl <- isIdentity g = f - - | Exists lhsB' <- rebuildLHS lhsB - = Lam lhsA - $ Body - $ Let lhsB' b - $ weakenE (sinkWithLHS lhsB lhsB' $ weakenWithLHS lhsA) c - -- = Stats.substitution "compose" . Lam lhs2 . Body $ substitute' f g -compose _ - _ = error "compose: impossible evaluation" - -subTop :: OpenExp env aenv s -> ExpVar (env, s) t -> OpenExp env aenv t -subTop s (Var _ ZeroIdx ) = s -subTop _ (Var tp (SuccIdx ix)) = Evar $ Var tp ix - subAtop :: PreOpenAcc acc aenv t -> ArrayVar (aenv, t) (Array sh2 e2) -> PreOpenAcc acc aenv (Array sh2 e2) subAtop t (Var _ ZeroIdx ) = t subAtop _ (Var repr (SuccIdx idx)) = Avar $ Var repr idx -data Identity a = Identity { runIdentity :: a } - -instance Functor Identity where - {-# INLINE fmap #-} - fmap f (Identity a) = Identity (f a) - -instance Applicative Identity where - {-# INLINE (<*>) #-} - {-# INLINE pure #-} - Identity f <*> Identity a = Identity (f a) - pure a = Identity a - -- A class for rebuilding terms. -- class Rebuildable f where @@ -277,22 +107,6 @@ class Rebuildable f where -> f aenv' a rebuildA av = runIdentity . rebuildPartial (Identity . av) --- A class for rebuilding scalar terms. --- -class RebuildableExp f where - {-# MINIMAL rebuildPartialE #-} - rebuildPartialE :: (Applicative f', SyntacticExp fe) - => (forall e'. ExpVar env e' -> f' (fe env' aenv e')) - -> f env aenv e - -> f' (f env' aenv e) - - {-# INLINEABLE rebuildE #-} - rebuildE :: SyntacticExp fe - => (forall e'. ExpVar env e' -> fe env' aenv e') - -> f env aenv e - -> f env' aenv e - rebuildE v = runIdentity . rebuildPartialE (Identity . v) - -- Terms that are rebuildable and also recursive closures -- type RebuildableAcc acc = (Rebuildable acc, AccClo acc ~ acc) @@ -310,12 +124,12 @@ data OpenAccFun (acc :: Type -> Type -> Type) env aenv a where instance Rebuildable (OpenAccExp acc env) where type AccClo (OpenAccExp acc env) = acc {-# INLINEABLE rebuildPartial #-} - rebuildPartial v (OpenAccExp e) = OpenAccExp <$> Stats.substitution "rebuild" (rebuildOpenExp (pure . IE) (reindexAvar v) e) + rebuildPartial v (OpenAccExp e) = OpenAccExp <$> Stats.substitution "rebuild" (rebuildArrayInstrPartial (reindexArrayInstr v) e) instance Rebuildable (OpenAccFun acc env) where type AccClo (OpenAccFun acc env) = acc {-# INLINEABLE rebuildPartial #-} - rebuildPartial v (OpenAccFun f) = OpenAccFun <$> Stats.substitution "rebuild" (rebuildFun (pure . IE) (reindexAvar v) f) + rebuildPartial v (OpenAccFun f) = OpenAccFun <$> Stats.substitution "rebuild" (rebuildArrayInstrPartial (reindexArrayInstr v) f) instance RebuildableAcc acc => Rebuildable (PreOpenAcc acc) where type AccClo (PreOpenAcc acc) = acc @@ -332,27 +146,8 @@ instance Rebuildable OpenAcc where {-# INLINEABLE rebuildPartial #-} rebuildPartial x = Stats.substitution "rebuild" $ rebuildOpenAcc x -instance RebuildableExp OpenExp where - {-# INLINEABLE rebuildPartialE #-} - rebuildPartialE v x = Stats.substitution "rebuild" $ rebuildOpenExp v (ReindexAvar pure) x - -instance RebuildableExp OpenFun where - {-# INLINEABLE rebuildPartialE #-} - rebuildPartialE v x = Stats.substitution "rebuild" $ rebuildFun v (ReindexAvar pure) x - --- NOTE: [Weakening] --- --- Weakening is something we usually take for granted: every time you learn a --- new word, old sentences still make sense. If a conclusion is justified by a --- hypothesis, it is still justified if you add more hypotheses. Similarly, a --- term remains in scope if you bind more (fresh) variables. Weakening is the --- operation of shifting things from one scope to a larger scope in which new --- things have become meaningful, but no old things have vanished. --- --- When we use a named representation (or HOAS) we get weakening for free. But --- in the de Bruijn representation weakening takes work: you have to shift all --- variable references to make room for the new bindings. --- +-- | Weakening +-- SEE: [Weakening] class Sink f where weaken :: env :> env' -> forall t. f env t -> f env' t @@ -364,28 +159,23 @@ class Sink f where -- default weaken :: Rebuildable f => env :> env' -> f env t -> f env' t -- weaken k = Stats.substitution "weaken" . rebuildA rebuildWeakenVar +class Sink' f where + weaken' :: env :> env' -> f env -> f env' + --instance Rebuildable f => Sink f where -- undecidable, incoherent -- weaken k = Stats.substitution "weaken" . rebuildA rebuildWeakenVar instance Sink Idx where {-# INLINEABLE weaken #-} - weaken = (>:>) + weaken w = (w >:>) instance Sink (Var s) where {-# INLINEABLE weaken #-} weaken k (Var s ix) = Var s (k >:> ix) -weakenVars :: env :> env' -> Vars s env t -> Vars s env' t -weakenVars _ TupRunit = TupRunit -weakenVars k (TupRsingle v) = TupRsingle $ weaken k v -weakenVars k (TupRpair v w) = TupRpair (weakenVars k v) (weakenVars k w) - rebuildWeakenVar :: env :> env' -> ArrayVar env (Array sh e) -> PreOpenAcc acc env' (Array sh e) rebuildWeakenVar k (Var s idx) = Avar $ Var s $ k >:> idx -rebuildWeakenEvar :: env :> env' -> ExpVar env t -> OpenExp env' aenv t -rebuildWeakenEvar k (Var s idx) = Evar $ Var s $ k >:> idx - instance RebuildableAcc acc => Sink (PreOpenAcc acc) where {-# INLINEABLE weaken #-} weaken k = Stats.substitution "weaken" . rebuildA (rebuildWeakenVar k) @@ -394,6 +184,15 @@ instance RebuildableAcc acc => Sink (PreOpenAfun acc) where {-# INLINEABLE weaken #-} weaken k = Stats.substitution "weaken" . rebuildA (rebuildWeakenVar k) +-- We cannot use 'weaken' to weaken the array environment of an 'OpenExp', +-- as OpenExp is a type synonym for 'PreOpenExp (ArrayInstr aenv) env', +-- and 'weaken' would thus affect the expression environment. Hence we +-- have a separate function for OpenExp and OpenFun. +-- +weakenArrayInstr :: RebuildableExp f => aenv :> aenv' -> f (ArrayInstr aenv) env t -> f (ArrayInstr aenv') env t +weakenArrayInstr k = Stats.substitution "weaken" . rebuildArrayInstr (ArrayInstr . weaken k) + +{- instance Sink (OpenExp env) where {-# INLINEABLE weaken #-} weaken k = Stats.substitution "weaken" . runIdentity . rebuildOpenExp (Identity . Evar) (ReindexAvar (Identity . weaken k)) @@ -401,6 +200,12 @@ instance Sink (OpenExp env) where instance Sink (OpenFun env) where {-# INLINEABLE weaken #-} weaken k = Stats.substitution "weaken" . runIdentity . rebuildFun (Identity . Evar) (ReindexAvar (Identity . weaken k)) +-} + +instance Sink ArrayInstr where + weaken k (Index var) = Index $ weaken k var + weaken k (LinearIndex var) = LinearIndex $ weaken k var + weaken k (Shape var) = Shape $ weaken k var instance Sink Boundary where {-# INLINEABLE weaken #-} @@ -410,7 +215,7 @@ instance Sink Boundary where Mirror -> Mirror Wrap -> Wrap Constant c -> Constant c - Function f -> Function (weaken k f) + Function f -> Function (weakenArrayInstr k f) instance Sink OpenAcc where {-# INLINEABLE weaken #-} @@ -423,171 +228,16 @@ instance Sink OpenAcc where -- weaken v1 (weaken v2 a) = weaken (v1 . v2) a -- #-} -class SinkExp f where - weakenE :: env :> env' -> f env aenv t -> f env' aenv t - - -- See comment in 'weaken' - -- - -- {-# INLINEABLE weakenE #-} - -- default weakenE :: RebuildableExp f => env :> env' -> f env aenv t -> f env' aenv t - -- weakenE v = Stats.substitution "weakenE" . rebuildE (IE . v) - -instance SinkExp OpenExp where - {-# INLINEABLE weakenE #-} - weakenE v = Stats.substitution "weakenE" . rebuildE (rebuildWeakenEvar v) -instance SinkExp OpenFun where - {-# INLINEABLE weakenE #-} - weakenE v = Stats.substitution "weakenE" . rebuildE (rebuildWeakenEvar v) - --- See above for why this is disabled. --- {-# RULES --- "weakenE/weakenE" forall a (v1 :: env' :> env'') (v2 :: env :> env'). --- weakenE v1 (weakenE v2 a) = weakenE (v1 . v2) a --- #-} - --- NOTE: [Strengthening] +-- SEE: [Strengthening] -- --- Strengthening is the dual of weakening. Shifting terms from one scope to a --- smaller scope. Of course this is not always possible. If the term contains --- any variables not in the new environment, then it cannot be strengthened. --- This partial behaviour is captured with 'Maybe'. --- - --- The type of partially shifting terms from one context into another. -type env :?> env' = forall t'. Idx env t' -> Maybe (Idx env' t') {-# INLINEABLE strengthen #-} strengthen :: forall f env env' t. Rebuildable f => env :?> env' -> f env t -> Maybe (f env' t) strengthen k x = Stats.substitution "strengthen" $ rebuildPartial @f @Maybe @IdxA (\(Var s ix) -> fmap (IA . Var s) $ k ix) x -{-# INLINEABLE strengthenE #-} -strengthenE :: forall f env env' aenv t. RebuildableExp f => env :?> env' -> f env aenv t -> Maybe (f env' aenv t) -strengthenE k x = Stats.substitution "strengthenE" $ rebuildPartialE @f @Maybe @IdxE (\(Var tp ix) -> fmap (IE . Var tp) $ k ix) x - -strengthenWithLHS :: LeftHandSide s t env1 env2 -> env2 :?> env1 -strengthenWithLHS (LeftHandSideWildcard _) = Just -strengthenWithLHS (LeftHandSideSingle _) = \ix -> case ix of - ZeroIdx -> Nothing - SuccIdx i -> Just i -strengthenWithLHS (LeftHandSidePair l1 l2) = strengthenWithLHS l2 >=> strengthenWithLHS l1 - -strengthenAfter :: LeftHandSide s t env1 env2 -> LeftHandSide s t env1' env2' -> env1 :?> env1' -> env2 :?> env2' -strengthenAfter (LeftHandSideWildcard _) (LeftHandSideWildcard _) k = k -strengthenAfter (LeftHandSideSingle _) (LeftHandSideSingle _) k = \ix -> case ix of - ZeroIdx -> Just ZeroIdx - SuccIdx i -> SuccIdx <$> k i -strengthenAfter (LeftHandSidePair l1 l2) (LeftHandSidePair l1' l2') k = - strengthenAfter l2 l2' $ strengthenAfter l1 l1' k -strengthenAfter _ _ _ = error "Substitution.strengthenAfter: left hand sides do not match" - --- Simultaneous Substitution =================================================== --- - --- The scalar environment --- ------------------ - --- SEE: [Renaming and Substitution] --- SEE: [Weakening] --- -class SyntacticExp f where - varIn :: ExpVar env t -> f env aenv t - expOut :: f env aenv t -> OpenExp env aenv t - weakenExp :: f env aenv t -> f (env, s) aenv t - -newtype IdxE env aenv t = IE { unIE :: ExpVar env t } - -instance SyntacticExp IdxE where - varIn = IE - expOut = Evar . unIE - weakenExp (IE (Var tp ix)) = IE $ Var tp $ SuccIdx ix - -instance SyntacticExp OpenExp where - varIn = Evar - expOut = id - weakenExp = runIdentity . rebuildOpenExp (Identity . weakenExp . IE) (ReindexAvar Identity) - -{-# INLINEABLE shiftE #-} -shiftE - :: (Applicative f, SyntacticExp fe) - => RebuildEvar f fe env env' aenv - -> RebuildEvar f fe (env, s) (env', s) aenv -shiftE _ (Var tp ZeroIdx) = pure $ varIn (Var tp ZeroIdx) -shiftE v (Var tp (SuccIdx ix)) = weakenExp <$> v (Var tp ix) - -{-# INLINEABLE shiftE' #-} -shiftE' - :: (Applicative f, SyntacticExp fa) - => ELeftHandSide t env1 env1' - -> ELeftHandSide t env2 env2' - -> RebuildEvar f fa env1 env2 aenv - -> RebuildEvar f fa env1' env2' aenv -shiftE' (LeftHandSideWildcard _) (LeftHandSideWildcard _) v = v -shiftE' (LeftHandSideSingle _) (LeftHandSideSingle _) v = shiftE v -shiftE' (LeftHandSidePair a1 b1) (LeftHandSidePair a2 b2) v = shiftE' b1 b2 $ shiftE' a1 a2 v -shiftE' _ _ _ = error "Substitution: left hand sides do not match" - -{-# INLINEABLE rebuildMaybeExp #-} -rebuildMaybeExp - :: (HasCallStack, Applicative f, SyntacticExp fe) - => RebuildEvar f fe env env' aenv' - -> ReindexAvar f aenv aenv' - -> Maybe (OpenExp env aenv t) - -> f (Maybe (OpenExp env' aenv' t)) -rebuildMaybeExp _ _ Nothing = pure Nothing -rebuildMaybeExp v av (Just x) = Just <$> rebuildOpenExp v av x - -{-# INLINEABLE rebuildOpenExp #-} -rebuildOpenExp - :: (HasCallStack, Applicative f, SyntacticExp fe) - => RebuildEvar f fe env env' aenv' - -> ReindexAvar f aenv aenv' - -> OpenExp env aenv t - -> f (OpenExp env' aenv' t) -rebuildOpenExp v av@(ReindexAvar reindex) exp = - case exp of - Const t c -> pure $ Const t c - PrimConst c -> pure $ PrimConst c - Undef t -> pure $ Undef t - Evar var -> expOut <$> v var - Let lhs a b - | Exists lhs' <- rebuildLHS lhs - -> Let lhs' <$> rebuildOpenExp v av a <*> rebuildOpenExp (shiftE' lhs lhs' v) av b - Pair e1 e2 -> Pair <$> rebuildOpenExp v av e1 <*> rebuildOpenExp v av e2 - Nil -> pure Nil - VecPack vec e -> VecPack vec <$> rebuildOpenExp v av e - VecUnpack vec e -> VecUnpack vec <$> rebuildOpenExp v av e - IndexSlice x ix sh -> IndexSlice x <$> rebuildOpenExp v av ix <*> rebuildOpenExp v av sh - IndexFull x ix sl -> IndexFull x <$> rebuildOpenExp v av ix <*> rebuildOpenExp v av sl - ToIndex shr sh ix -> ToIndex shr <$> rebuildOpenExp v av sh <*> rebuildOpenExp v av ix - FromIndex shr sh ix -> FromIndex shr <$> rebuildOpenExp v av sh <*> rebuildOpenExp v av ix - Case e rhs def -> Case <$> rebuildOpenExp v av e <*> sequenceA [ (t,) <$> rebuildOpenExp v av c | (t,c) <- rhs ] <*> rebuildMaybeExp v av def - Cond p t e -> Cond <$> rebuildOpenExp v av p <*> rebuildOpenExp v av t <*> rebuildOpenExp v av e - While p f x -> While <$> rebuildFun v av p <*> rebuildFun v av f <*> rebuildOpenExp v av x - PrimApp f x -> PrimApp f <$> rebuildOpenExp v av x - Index a sh -> Index <$> reindex a <*> rebuildOpenExp v av sh - LinearIndex a i -> LinearIndex <$> reindex a <*> rebuildOpenExp v av i - Shape a -> Shape <$> reindex a - ShapeSize shr sh -> ShapeSize shr <$> rebuildOpenExp v av sh - Foreign tp ff f e -> Foreign tp ff f <$> rebuildOpenExp v av e - Coerce t1 t2 e -> Coerce t1 t2 <$> rebuildOpenExp v av e - -{-# INLINEABLE rebuildFun #-} -rebuildFun - :: (HasCallStack, Applicative f, SyntacticExp fe) - => RebuildEvar f fe env env' aenv' - -> ReindexAvar f aenv aenv' - -> OpenFun env aenv t - -> f (OpenFun env' aenv' t) -rebuildFun v av fun = - case fun of - Body e -> Body <$> rebuildOpenExp v av e - Lam lhs f - | Exists lhs' <- rebuildLHS lhs - -> Lam lhs' <$> rebuildFun (shiftE' lhs lhs' v) av f - --- The array environment --- ----------------- +-- Substitutions in array environment +-- ---------------------------------- type RebuildAcc acc = forall aenv aenv' f fa a. (HasCallStack, Applicative f, SyntacticAcc fa) @@ -615,26 +265,6 @@ instance SyntacticAcc PreOpenAcc where type RebuildAvar f (fa :: (Type -> Type -> Type) -> Type -> Type -> Type) acc aenv aenv' = forall sh e. ArrayVar aenv (Array sh e) -> f (fa acc aenv' (Array sh e)) -type RebuildEvar f fe env env' aenv' = - forall t'. ExpVar env t' -> f (fe env' aenv' t') - -newtype ReindexAvar f aenv aenv' = - ReindexAvar (forall sh e. ArrayVar aenv (Array sh e) -> f (ArrayVar aenv' (Array sh e))) - -reindexAvar - :: forall f fa acc aenv aenv'. - (HasCallStack, Applicative f, SyntacticAcc fa) - => RebuildAvar f fa acc aenv aenv' - -> ReindexAvar f aenv aenv' -reindexAvar v = ReindexAvar f where - f :: forall sh e. ArrayVar aenv (Array sh e) -> f (ArrayVar aenv' (Array sh e)) - f var = g <$> v var - - g :: fa acc aenv' (Array sh e) -> ArrayVar aenv' (Array sh e) - g fa = case accOut fa of - Avar var' -> var' - _ -> internalError "An Avar which was used in an Exp was mapped to an array term other than Avar. This mapping is invalid as an Exp can only contain array variables." - {-# INLINEABLE shiftA #-} shiftA @@ -668,7 +298,8 @@ rebuildOpenAcc av (OpenAcc acc) = OpenAcc <$> rebuildPreOpenAcc rebuildOpenAcc a {-# INLINEABLE rebuildPreOpenAcc #-} rebuildPreOpenAcc - :: (HasCallStack, Applicative f, SyntacticAcc fa) + :: forall f fa acc aenv aenv' t. + (HasCallStack, Applicative f, SyntacticAcc fa) => RebuildAcc acc -> RebuildAvar f fa acc aenv aenv' -> PreOpenAcc acc aenv t @@ -681,30 +312,44 @@ rebuildPreOpenAcc k av acc = Apair as bs -> Apair <$> k av as <*> k av bs Anil -> pure Anil Atrace msg as bs -> Atrace msg <$> k av as <*> k av bs - Apply repr f a -> Apply repr <$> rebuildAfun k av f <*> k av a - Acond p t e -> Acond <$> rebuildOpenExp (pure . IE) av' p <*> k av t <*> k av e + Manifest as -> Manifest <$> k av as + Aassert msg c as -> Aassert msg <$> travE c <*> k av as + Aassume c as -> Aassume <$> travE c <*> k av as + Acond p t e -> Acond <$> travE p <*> k av t <*> k av e Awhile p f a -> Awhile <$> rebuildAfun k av p <*> rebuildAfun k av f <*> k av a - Unit tp e -> Unit tp <$> rebuildOpenExp (pure . IE) av' e - Reshape shr e a -> Reshape shr <$> rebuildOpenExp (pure . IE) av' e <*> k av a - Generate repr e f -> Generate repr <$> rebuildOpenExp (pure . IE) av' e <*> rebuildFun (pure . IE) av' f - Transform repr sh ix f a -> Transform repr <$> rebuildOpenExp (pure . IE) av' sh <*> rebuildFun (pure . IE) av' ix <*> rebuildFun (pure . IE) av' f <*> k av a - Replicate sl slix a -> Replicate sl <$> rebuildOpenExp (pure . IE) av' slix <*> k av a - Slice sl a slix -> Slice sl <$> k av a <*> rebuildOpenExp (pure . IE) av' slix - Map tp f a -> Map tp <$> rebuildFun (pure . IE) av' f <*> k av a - ZipWith tp f a1 a2 -> ZipWith tp <$> rebuildFun (pure . IE) av' f <*> k av a1 <*> k av a2 - Fold f z a -> Fold <$> rebuildFun (pure . IE) av' f <*> rebuildMaybeExp (pure . IE) av' z <*> k av a - FoldSeg itp f z a s -> FoldSeg itp <$> rebuildFun (pure . IE) av' f <*> rebuildMaybeExp (pure . IE) av' z <*> k av a <*> k av s - Scan d f z a -> Scan d <$> rebuildFun (pure . IE) av' f <*> rebuildMaybeExp (pure . IE) av' z <*> k av a - Scan' d f z a -> Scan' d <$> rebuildFun (pure . IE) av' f <*> rebuildOpenExp (pure . IE) av' z <*> k av a - Permute f1 a1 f2 a2 -> Permute <$> rebuildFun (pure . IE) av' f1 <*> k av a1 <*> rebuildFun (pure . IE) av' f2 <*> k av a2 - Backpermute shr sh f a -> Backpermute shr <$> rebuildOpenExp (pure . IE) av' sh <*> rebuildFun (pure . IE) av' f <*> k av a - Stencil sr tp f b a -> Stencil sr tp <$> rebuildFun (pure . IE) av' f <*> rebuildBoundary av' b <*> k av a + Unit tp e -> Unit tp <$> travE e + Reshape shr e a -> Reshape shr <$> travE e <*> k av a + Generate repr e f -> Generate repr <$> travE e <*> travF f + Transform repr sh ix f a -> Transform repr <$> travE sh <*> travF ix <*> travF f <*> k av a + Replicate sl slix a -> Replicate sl <$> travE slix <*> k av a + Slice sl a slix -> Slice sl <$> k av a <*> travE slix + Map tp f a -> Map tp <$> travF f <*> k av a + ZipWith tp f a1 a2 -> ZipWith tp <$> travF f <*> k av a1 <*> k av a2 + Fold f z a -> Fold <$> travF f <*> travME z <*> k av a + FoldSeg itp f z a s -> FoldSeg itp <$> travF f <*> travME z <*> k av a <*> k av s + Scan d f z a -> Scan d <$> travF f <*> travME z <*> k av a + Scan' d f z a -> Scan' d <$> travF f <*> travE z <*> k av a + Permute f1 a1 a2 -> Permute <$> travMF f1 <*> k av a1 <*> k av a2 + Backpermute shr sh f a -> Backpermute shr <$> travE sh <*> travF f <*> k av a + Stencil sr tp f b a -> Stencil sr tp <$> travF f <*> rebuildBoundary av b <*> k av a Stencil2 s1 s2 tp f b1 a1 b2 a2 - -> Stencil2 s1 s2 tp <$> rebuildFun (pure . IE) av' f <*> rebuildBoundary av' b1 <*> k av a1 <*> rebuildBoundary av' b2 <*> k av a2 + -> Stencil2 s1 s2 tp <$> travF f <*> rebuildBoundary av b1 <*> k av a1 <*> rebuildBoundary av b2 <*> k av a2 Aforeign repr ff afun as -> Aforeign repr ff afun <$> k av as -- Collect seq -> Collect <$> rebuildSeq k av seq where - av' = reindexAvar av + travE :: Exp aenv e -> f (Exp aenv' e) + travE = rebuildArrayInstrPartial (reindexArrayInstr av) + + travME :: Maybe (Exp aenv e) -> f (Maybe (Exp aenv' e)) + travME Nothing = pure Nothing + travME (Just e) = Just <$> travE e + + travF :: Fun aenv e -> f (Fun aenv' e) + travF = rebuildArrayInstrPartial (reindexArrayInstr av) + + travMF :: Maybe (Fun aenv e) -> f (Maybe (Fun aenv' e)) + travMF Nothing = pure Nothing + travMF (Just x) = Just <$> travF x {-# INLINEABLE rebuildAfun #-} rebuildAfun @@ -730,19 +375,10 @@ rebuildAlet k av lhs1 bind1 body1 | Exists lhs2 <- rebuildLHS lhs1 = Alet lhs2 <$> k av bind1 <*> k (shiftA' lhs1 lhs2 k av) body1 -{-# INLINEABLE rebuildLHS #-} -rebuildLHS :: LeftHandSide s t aenv1 aenv1' -> Exists (LeftHandSide s t aenv2) -rebuildLHS (LeftHandSideWildcard r) = Exists $ LeftHandSideWildcard r -rebuildLHS (LeftHandSideSingle s) = Exists $ LeftHandSideSingle s -rebuildLHS (LeftHandSidePair as bs) - | Exists as' <- rebuildLHS as - , Exists bs' <- rebuildLHS bs - = Exists $ LeftHandSidePair as' bs' - {-# INLINEABLE rebuildBoundary #-} rebuildBoundary - :: Applicative f - => ReindexAvar f aenv aenv' + :: (HasCallStack, Applicative f, SyntacticAcc fa) + => RebuildAvar f fa acc aenv aenv' -> Boundary aenv t -> f (Boundary aenv' t) rebuildBoundary av bndy = @@ -751,7 +387,25 @@ rebuildBoundary av bndy = Mirror -> pure Mirror Wrap -> pure Wrap Constant v -> pure (Constant v) - Function f -> Function <$> rebuildFun (pure . IE) av f + Function f -> Function <$> rebuildArrayInstrPartial (reindexArrayInstr av) f + +reindexArrayInstr + :: forall f fa acc aenv aenv'. + (HasCallStack, Applicative f, SyntacticAcc fa) + => RebuildAvar f fa acc aenv aenv' + -> RebuildArrayInstr f (ArrayInstr aenv) (ArrayInstr aenv') +reindexArrayInstr av arr = case arr of + Shape var -> ArrayInstr . Shape <$> reindex var + Index var -> ArrayInstr . Index <$> reindex var + LinearIndex var -> ArrayInstr . LinearIndex <$> reindex var + where + reindex :: forall sh e. ArrayVar aenv (Array sh e) -> f (ArrayVar aenv' (Array sh e)) + reindex var = g <$> av var + + g :: fa acc aenv' (Array sh e) -> ArrayVar aenv' (Array sh e) + g fa = case accOut fa of + Avar var' -> var' + _ -> internalError "An Avar which was used in an Exp was mapped to an array term other than Avar. This mapping is invalid as an Exp can only contain array variables." {-- {-# INLINEABLE rebuildSeq #-} @@ -798,10 +452,3 @@ rebuildC k v c = rebuildT NilAtup = pure NilAtup rebuildT (SnocAtup t s) = SnocAtup <$> (rebuildT t) <*> (rebuildC k v s) --} - -extractExpVars :: OpenExp env aenv a -> Maybe (ExpVars env a) -extractExpVars Nil = Just TupRunit -extractExpVars (Pair e1 e2) = TupRpair <$> extractExpVars e1 <*> extractExpVars e2 -extractExpVars (Evar v) = Just $ TupRsingle v -extractExpVars _ = Nothing - diff --git a/src/Data/Array/Accelerate/Trafo/Var.hs b/src/Data/Array/Accelerate/Trafo/Var.hs index aef30b672..ca52716d0 100644 --- a/src/Data/Array/Accelerate/Trafo/Var.hs +++ b/src/Data/Array/Accelerate/Trafo/Var.hs @@ -1,8 +1,3 @@ -{-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeOperators #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.Trafo.Var @@ -24,6 +19,7 @@ import Data.Array.Accelerate.AST.LeftHandSide import Data.Array.Accelerate.AST.Var import Data.Array.Accelerate.Representation.Array import Data.Array.Accelerate.Representation.Type +import Data.Array.Accelerate.Trafo.Substitution data DeclareVars s t aenv where @@ -77,3 +73,33 @@ avarsOut extract = \case -> Just (TupRpair as bs) _ -> Nothing +identity :: TypeR a -> PreOpenFun arr env (a -> a) +identity t + | DeclareVars lhs _ value <- declareVars t + = Lam lhs $ Body $ expVars $ value weakenId + +-- Used in various analyses +data MaybeVar s env t where + NoVar :: MaybeVar s env t + JustVar :: Var s env t -> MaybeVar s env t +type MaybeVars s env = TupR (MaybeVar s env) + +strengthenMaybeVar :: LeftHandSide s' t env env' -> MaybeVar s env' u -> MaybeVar s env u +strengthenMaybeVar _ NoVar = NoVar +strengthenMaybeVar (LeftHandSideWildcard _) v = v +strengthenMaybeVar (LeftHandSideSingle _) (JustVar (Var t ix)) = case ix of + SuccIdx ix' -> JustVar $ Var t ix' + ZeroIdx -> NoVar +strengthenMaybeVar (LeftHandSidePair l1 l2) v = strengthenMaybeVar l1 $ strengthenMaybeVar l2 v + +strengthenMaybeVars :: LeftHandSide s' t env env' -> MaybeVars s env' u -> MaybeVars s env u +strengthenMaybeVars lhs = mapTupR (strengthenMaybeVar lhs) + +instance Sink (MaybeVar s) where + weaken _ NoVar = NoVar + weaken k (JustVar var) = JustVar $ weaken k var + +lhsMaybeVars :: LeftHandSide s t env env' -> MaybeVars s env' t +lhsMaybeVars (LeftHandSideWildcard tp) = mapTupR (const NoVar) tp +lhsMaybeVars (LeftHandSidePair l1 l2) = mapTupR (weaken $ weakenWithLHS l2) (lhsMaybeVars l1) `TupRpair` lhsMaybeVars l2 +lhsMaybeVars (LeftHandSideSingle tp) = TupRsingle $ JustVar $ Var tp ZeroIdx diff --git a/src/Data/Array/Accelerate/Trafo/WeakenedEnvironment.hs b/src/Data/Array/Accelerate/Trafo/WeakenedEnvironment.hs new file mode 100644 index 000000000..887f042d0 --- /dev/null +++ b/src/Data/Array/Accelerate/Trafo/WeakenedEnvironment.hs @@ -0,0 +1,137 @@ +{-# OPTIONS_HADDOCK hide #-} +-- | +-- Module : Data.Array.Accelerate.AST.WeakenedEnvironment +-- Copyright : [2008..2020] The Accelerate Team +-- License : BSD3 +-- +-- Maintainer : Trevor L. McDonell +-- Stability : experimental +-- Portability : non-portable (GHC extensions) +-- + +module Data.Array.Accelerate.Trafo.WeakenedEnvironment + ( WEnv, WEnv'(..), wprj, wprj', wupdate, wupdateSetWeakened, wempty + , wpush, wpush2, wpush', wremoveSet, wupdatePrjSet, wprjSet, wenvToList + ) where + +import Data.Array.Accelerate.AST.Environment +import Data.Array.Accelerate.AST.Idx +import Data.Array.Accelerate.AST.IdxSet (IdxSet(..)) +import Data.Array.Accelerate.AST.LeftHandSide +import Data.Array.Accelerate.Trafo.Substitution + +-- Valuation for an environment +-- +type WEnv f env = WEnv' f env env + +data WEnv' f env' env where + WEmpty :: WEnv' f env' () + WPushA :: WEnv' f env' env -> f env' t -> WEnv' f env' (env, t) + WPushB :: WEnv' f env' env -> f (env', t) t -> WEnv' f (env', t) (env, t) + WWeaken :: env1 :> env2 -> WEnv' f env1 env -> WEnv' f env2 env + +wprj :: Sink f => Idx envB t -> WEnv' f envA envB -> f envA t +wprj = wprj' weakenId + +wprj' :: Sink f => env1 :> env2 -> Idx env t -> WEnv' f env1 env -> f env2 t +wprj' k idx (WWeaken k' env) = wprj' (k .> k') idx env +wprj' k ZeroIdx (WPushA _ v) = weaken k v +wprj' k ZeroIdx (WPushB _ v) = weaken k v +wprj' k (SuccIdx idx) (WPushA env _) = wprj' k idx env +wprj' k (SuccIdx idx) (WPushB env _) = wprj' (weakenSucc k) idx env + +instance Sink f => Sink (WEnv' f) where + weaken k (WWeaken k' env) = WWeaken (k .> k') env + weaken k env = WWeaken k env + +wupdate :: forall f env' env t. Sink f => (f env' t -> f env' t) -> Idx env t -> WEnv' f env' env -> WEnv' f env' env +wupdate f = go' + where + -- Specialized version of go, where env1 ~ env' + go' :: Idx env2 t -> WEnv' f env' env2 -> WEnv' f env' env2 + go' idx (WWeaken k' env) = go k' idx env + go' ZeroIdx (WPushA env v) = WPushA env $ f v + go' ZeroIdx (WPushB env v) = WPushB env $ f v + go' (SuccIdx idx) (WPushA env v) = WPushA (go' idx env) v + go' (SuccIdx idx) (WPushB env v) = WPushA (go (weakenSucc weakenId) idx env) v + + go :: (env1 :> env') -> Idx env2 t -> WEnv' f env1 env2 -> WEnv' f env' env2 + go k idx (WWeaken k' env) = go (k .> k') idx env + go k ZeroIdx (WPushA env v) = WPushA (weaken k env) $ f $ weaken k v + go k ZeroIdx (WPushB env v) = WPushA (weaken (weakenSucc k) env) $ f $ weaken k v + go k (SuccIdx idx) (WPushA env v) = WPushA (go k idx env) $ weaken k v + go k (SuccIdx idx) (WPushB env v) = WPushA (go (weakenSucc k) idx env) $ weaken k v + +wupdateSetWeakened :: forall f env' env. (forall env1 t. env1 :> env' -> f env1 t -> f env1 t) -> IdxSet env -> WEnv' f env' env -> WEnv' f env' env +wupdateSetWeakened f (IdxSet set) = go weakenId set + where + go :: (env1 :> env') -> PartialEnv g env2 -> WEnv' f env1 env2 -> WEnv' f env1 env2 + go _ PEnd env = env + go k p (WWeaken k' env) = WWeaken k' $ go (k .> k') p env + go k (PPush p _) (WPushA env v) = WPushA (go k p env) (f k v) + go k (PPush p _) (WPushB env v) = WPushB (go (weakenSucc k) p env) (f k v) + go k (PNone p) (WPushA env v) = WPushA (go k p env) v + go k (PNone p) (WPushB env v) = WPushB (go (weakenSucc k) p env) v + +wempty :: WEnv f () +wempty = WEmpty + +wpush :: WEnv f env -> f (env, t) t -> WEnv f (env, t) +wpush = WPushB + +wpush2 :: Sink f => WEnv f env -> f ((env, t), s) t -> f ((env, t), s) s -> WEnv f ((env, t), s) +wpush2 env t s = weaken (weakenSucc $ weakenSucc weakenId) env `WPushA` t `WPushA` s + +wpush' :: WEnv' f env' env -> f env' t -> WEnv' f env' (env, t) +wpush' = WPushA + +wremoveSet :: forall f env. (forall env' t. f env' t) -> IdxSet env -> WEnv f env -> WEnv f env +wremoveSet nil (IdxSet set) env = go set env + where + go :: PartialEnv g env' -> WEnv' f env1 env' -> WEnv' f env1 env' + go PEnd e = e + go p (WWeaken k e) = WWeaken k $ go p e + go (PNone p) (WPushA e f) = WPushA (go p e) f + go (PNone p) (WPushB e f) = WPushB (go p e) f + go (PPush p _) (WPushA e _) = WPushA (go p e) nil + go (PPush p _) (WPushB e _) = WPushB (go p e) nil + +wupdatePrjSet :: forall f env. Sink f => (forall env' t. f env' t -> f env' t) -> IdxSet env -> WEnv f env -> ([Exists (f env)], WEnv f env) +wupdatePrjSet update (IdxSet set) env = go weakenId set env + where + go :: env1 :> env -> PartialEnv g env' -> WEnv' f env1 env' -> ([Exists (f env)], WEnv' f env1 env') + go _ PEnd e = ([], e) + go k p (WWeaken k' e) = (l, WWeaken k' e') + where (l, e') = go (k .> k') p e + go k (PNone p) (WPushA e f) = (Exists (weaken k f) : l, WPushA e' $ update f) + where (l, e') = go k p e + go k (PNone p) (WPushB e f) = (Exists (weaken k f) : l, WPushB e' $ update f) + where (l, e') = go (weakenSucc k) p e + go k (PPush p _) (WPushA e f) = (l, WPushA e' f) + where (l, e') = go k p e + go k (PPush p _) (WPushB e f) = (l, WPushB e' f) + where (l, e') = go (weakenSucc k) p e + +wprjSet :: forall f env. Sink f => IdxSet env -> WEnv f env -> ([Exists (f env)]) +wprjSet (IdxSet set) env = go weakenId set env + where + go :: env1 :> env -> PartialEnv g env' -> WEnv' f env1 env' -> ([Exists (f env)]) + go _ PEnd _ = [] + go k p (WWeaken k' e) = go (k .> k') p e + go k (PPush p _) (WPushA e f) = Exists (weaken k f) : go k p e + go k (PPush p _) (WPushB e f) = Exists (weaken k f) : go (weakenSucc k) p e + go k (PNone p) (WPushA e _) = go k p e + go k (PNone p) (WPushB e _) = go (weakenSucc k) p e + +wenvToList :: forall f env. Sink f => WEnv f env -> [EnvBinding (f env) env] +wenvToList = go weakenId weakenId + where + go :: env1 :> env -> env2 :> env -> WEnv' f env1 env2 -> [EnvBinding (f env) env] + go _ _ WEmpty = [] + go k1 k2 (WPushA env a) = + EnvBinding (k2 >:> ZeroIdx) (weaken k1 a) + : go k1 (weakenSucc k2) env + go k1 k2 (WPushB env a) = + EnvBinding (k2 >:> ZeroIdx) (weaken k1 a) + : go (weakenSucc k1) (weakenSucc k2) env + go k1 k2 (WWeaken k env) = go (k1 .> k) k2 env diff --git a/src/Data/Array/Accelerate/Type.hs b/src/Data/Array/Accelerate/Type.hs index 94e891cc1..29eb73fa0 100644 --- a/src/Data/Array/Accelerate/Type.hs +++ b/src/Data/Array/Accelerate/Type.hs @@ -1,19 +1,6 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE MagicHash #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE RoleAnnotations #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeFamilies #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Array.Accelerate.Type @@ -89,6 +76,11 @@ import GHC.TypeLits -- Reified dictionaries -- +data ScalarDict a where + ScalarDict :: ( Eq a, Show a, Storable a ) + => ScalarDict a + + data SingleDict a where SingleDict :: ( Eq a, Ord a, Show a, Storable a, Prim a ) => SingleDict a @@ -302,6 +294,12 @@ singleDict = single floating TypeFloat = SingleDict floating TypeDouble = SingleDict +scalarDict :: ScalarType a -> ScalarDict a +scalarDict (SingleScalarType tp) + | SingleDict <- singleDict tp = ScalarDict +scalarDict (VectorScalarType (VectorType _ tp)) + | SingleDict <- singleDict tp = ScalarDict + scalarTypeInt :: ScalarType Int scalarTypeInt = SingleScalarType $ NumSingleType $ IntegralNumType TypeInt @@ -318,6 +316,9 @@ scalarTypeWord8 = SingleScalarType $ NumSingleType $ IntegralNumType TypeWord8 scalarTypeWord32 :: ScalarType Word32 scalarTypeWord32 = SingleScalarType $ NumSingleType $ IntegralNumType TypeWord32 +scalarTypeWord64 :: ScalarType Word64 +scalarTypeWord64 = SingleScalarType $ NumSingleType $ IntegralNumType TypeWord64 + rnfScalarType :: ScalarType t -> () rnfScalarType (SingleScalarType t) = rnfSingleType t rnfScalarType (VectorScalarType t) = rnfVectorType t @@ -517,4 +518,3 @@ runQ $ do vs <- mapM (uncurry mkVector) vectorTypes -- return (concat is ++ concat fs ++ concat vs) - diff --git a/src/Data/Array/Accelerate/Unsafe.hs b/src/Data/Array/Accelerate/Unsafe.hs index 289cf7a2b..822fedfbd 100644 --- a/src/Data/Array/Accelerate/Unsafe.hs +++ b/src/Data/Array/Accelerate/Unsafe.hs @@ -1,5 +1,3 @@ -{-# LANGUAGE MonoLocalBinds #-} -{-# LANGUAGE FlexibleContexts #-} -- | -- Module : Data.Array.Accelerate.Unsafe -- Copyright : [2009..2020] The Accelerate Team @@ -18,12 +16,13 @@ module Data.Array.Accelerate.Unsafe ( -- ** Unsafe operations Coerce, coerce, - undef, + undef, Assert(assume), ) where import Data.Array.Accelerate.Smart import Data.Array.Accelerate.Sugar.Elt +import Data.Array.Accelerate.Language -- | The function 'coerce' allows you to convert a value between any two types diff --git a/src/Data/Atomic.hs b/src/Data/Atomic.hs index 386dcea87..d39fd84e6 100644 --- a/src/Data/Atomic.hs +++ b/src/Data/Atomic.hs @@ -1,9 +1,6 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE ForeignFunctionInterface #-} -{-# LANGUAGE MagicHash #-} -{-# LANGUAGE NoImplicitPrelude #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE UnboxedTuples #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE NoImplicitPrelude #-} +{-# LANGUAGE TemplateHaskell #-} {-# OPTIONS_GHC -fobject-code #-} -- | -- Module : Data.Atomic diff --git a/src/Data/BitSet.hs b/src/Data/BitSet.hs index a3d18c1ae..0723d95a8 100644 --- a/src/Data/BitSet.hs +++ b/src/Data/BitSet.hs @@ -1,4 +1,4 @@ -{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} {-# LANGUAGE TypeFamilies #-} {-# OPTIONS_HADDOCK hide #-} -- | @@ -13,9 +13,14 @@ module Data.BitSet where -import Data.Bits +#if MIN_VERSION_base (4,20,0) -- foldl' is exported by Prelude from GHC 9.10 import Prelude hiding ( foldl, foldr, foldl' ) +#else +import Prelude hiding ( foldl, foldr ) +#endif + +import Data.Bits import qualified Data.List as List import GHC.Exts ( IsList, build ) diff --git a/src/Data/Primitive/Vec.hs b/src/Data/Primitive/Vec.hs index 0342f401c..0b904632f 100644 --- a/src/Data/Primitive/Vec.hs +++ b/src/Data/Primitive/Vec.hs @@ -1,15 +1,8 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE KindSignatures #-} -{-# LANGUAGE MagicHash #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE RoleAnnotations #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE UnboxedTuples #-} -{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE UnboxedTuples #-} +{-# LANGUAGE ViewPatterns #-} {-# OPTIONS_HADDOCK hide #-} -- | -- Module : Data.Primitive.Vec @@ -37,10 +30,13 @@ module Data.Primitive.Vec ( ) where import Control.Monad.ST +import Data.Kind ( Type ) import Data.Primitive.ByteArray import Data.Primitive.Types -import Language.Haskell.TH.Extra +import Language.Haskell.TH.Extra qualified as TH import Prettyprinter +import qualified Foreign.Storable as S +import Foreign.Ptr import GHC.Base ( isTrue# ) import GHC.Int @@ -81,7 +77,7 @@ import GHC.Word -- One inefficiency of this approach is that the byte array does track its size, -- which redundant for our use case (derivable from type level information). -- -data Vec (n :: Nat) a = Vec ByteArray# +data Vec (n :: Nat) (a :: Type) = Vec ByteArray# type role Vec nominal representational @@ -93,6 +89,33 @@ instance (Show a, Prim a, KnownNat n) => Show (Vec n a) where . group . encloseSep (flatAlt "< " "<") (flatAlt " >" ">") ", " . map viaShow +instance (S.Storable a, Prim a, KnownNat n) => S.Storable (Vec n a) where + sizeOf _ = S.sizeOf (undefined :: a) * fromIntegral (natVal' (proxy# :: Proxy# n)) + alignment _ = S.alignment (undefined :: a) + peek vecPtr = do + let elemPtr = castPtr vecPtr :: Ptr a + mba <- newByteArray $ sizeOf (undefined :: a) * fromIntegral (natVal' (proxy# :: Proxy# n)) + let + go i + | i == fromIntegral (natVal' (proxy# :: Proxy# n)) = return () + | otherwise = do + value <- S.peekElemOff elemPtr i + writeByteArray mba i value + go (i + 1) + go 0 + ByteArray ba# <- unsafeFreezeByteArray mba + return $! Vec ba# + poke vecPtr (Vec ba#) = do + let + elemPtr = castPtr vecPtr :: Ptr a + go i@(I# i#) + | i == fromIntegral (natVal' (proxy# :: Proxy# n)) = return () + | otherwise = do + let value = indexByteArray# ba# i# + S.pokeElemOff elemPtr i value + go (i + 1) + go 0 + listOfVec :: forall a n. (Prim a, KnownNat n) => Vec n a -> [a] listOfVec (Vec ba#) = go 0# where @@ -263,12 +286,12 @@ packVec16 a b c d e f g h i j k l m n o p = runST $ do -- to do this without copying, but I don't think the definition of ByteArray# is -- exported (or it is deeply magical). -- -liftVec :: Vec n a -> CodeQ (Vec n a) +liftVec :: Vec n a -> TH.CodeQ (Vec n a) liftVec (Vec ba#) - = unsafeCodeCoerce + = TH.unsafeCodeCoerce [| runST $ \s -> case newByteArray# $(liftInt# n#) s of { (# s1, mba# #) -> - case copyAddrToByteArray# $(litE (StringPrimL bytes)) mba# 0# $(liftInt# n#) s1 of { s2 -> + case copyAddrToByteArray# $(TH.litE (TH.StringPrimL bytes)) mba# 0# $(liftInt# n#) s1 of { s2 -> case unsafeFreezeByteArray# mba# s2 of { (# s3, ba'# #) -> (# s3, Vec ba'# #) }}} @@ -284,6 +307,6 @@ liftVec (Vec ba#) -- XXX: Typed TH does not support unlifted types -- - liftInt# :: Int# -> ExpQ - liftInt# i# = litE (IntPrimL (toInteger (I# i#))) + liftInt# :: Int# -> TH.ExpQ + liftInt# i# = TH.litE (TH.IntPrimL (toInteger (I# i#))) diff --git a/src/Language/Haskell/TH/Extra.hs b/src/Language/Haskell/TH/Extra.hs index 7688167fc..2dd5b9546 100644 --- a/src/Language/Haskell/TH/Extra.hs +++ b/src/Language/Haskell/TH/Extra.hs @@ -1,8 +1,5 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE KindSignatures #-} -{-# LANGUAGE PolyKinds #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE TemplateHaskellQuotes #-} -- | -- Module : Language.Haskell.TH.Extra -- Copyright : [2019..2020] The Accelerate Team diff --git a/stack.yaml b/stack.yaml index 18bb7bb61..e832c5b09 100644 --- a/stack.yaml +++ b/stack.yaml @@ -2,14 +2,17 @@ # For advanced use and comprehensive documentation of the format, please see: # https://docs.haskellstack.org/en/stable/yaml_configuration/ -resolver: lts-21.4 +resolver: lts-21.25 packages: - . extra-deps: -- git: https://github.com/gregorycollins/hashtables.git - commit: 0c5c8548ad7e8e90ab68bfbe6697069c969718aa +- OptDir-0.0.4 +- bytestring-encoding-0.1.2.0 +- fclabels-2.0.5.1 +- MIP-0.2.0.0 +- hashtables-1.4.2 # Override default flag values for local packages and extra-deps # flags: {} @@ -36,4 +39,3 @@ extra-deps: # compiler-check: newer-minor # # vim: nospell - diff --git a/test/doctest/Main.hs b/test/doctest/Main.hs index d2b98743e..cdc74db25 100644 --- a/test/doctest/Main.hs +++ b/test/doctest/Main.hs @@ -8,16 +8,500 @@ -- Portability : non-portable (GHC extensions) -- +-- module Main where + +-- import Build_doctests ( flags, pkgs, module_sources ) +-- import DatFoldable ( traverse_ ) +-- import Test.DocTest + +-- main :: IO () +-- main = do +-- traverse_ putStrLn args +-- doctest args +-- where +-- args = flags ++ pkgs ++ module_sources + +{-# language FlexibleInstances #-} +{-# language FlexibleContexts #-} +{-# language ScopedTypeVariables #-} +{-# language TemplateHaskell #-} +{-# language TypeApplications #-} +{-# language TypeOperators #-} +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE RebindableSyntax #-} +{-# OPTIONS_GHC -Wno-name-shadowing -Wno-unused-matches -Wno-unused-local-binds #-} module Main where -import Build_doctests ( flags, pkgs, module_sources ) -import Data.Foldable ( traverse_ ) -import Test.DocTest +import Data.Array.Accelerate as A +import qualified Prelude as P + +import Control.Applicative +import Control.Monad +import Data.Binary ( decodeFile, encodeFile ) + +import Data.Array.Accelerate.Interpreter + +import Data.Word +import Data.Random.Normal +import System.Environment +import System.IO +import System.Directory +import System.Random + +import Criterion.Main main :: IO () -main = do - traverse_ putStrLn args - doctest args +main = print . runN @Interpreter $ difficult (use $ fromList (Z:.20) $ [1..]) + +dotp :: Acc (Vector Int) -> Acc (Vector Int) -> Acc (Scalar Int) +dotp a b = fold (+) 0 $ zipWith (*) (map (+1) a) (map (`div` 2) b) + +diagonal :: (Shape sh, P.Num (Exp b), Elt b) + => Acc (Array sh b) -> Acc (Array sh b, Array sh b) +diagonal xs = + let ys = map (+2) xs + zs = map (+3) ys + in T2 ys zs + +-- Neither of the backpermutes is allowed to fuse with the map: otherwise the other backpermute cannot be computed. +-- Fusing both is possible, but only with work duplication (we still choose to never do that for now). +-- The backpermutes _are_ allowed to fuse with each other: This should however 1. not be rewarded 2. supported in codegen +difficult :: Acc (Array DIM1 Int) -> Acc (Array DIM1 (Int, Int)) +difficult acc = zip + (backpermute sh (\(I1 y) -> I1 (y `div` 2)) x) + (backpermute sh (\(I1 y) -> I1 (y + 1)) x) where - args = flags ++ pkgs ++ module_sources + x = map (+3) acc + sh = I1 10 + + +complex xs = + let as = map (+1) xs + bs = map (+2) xs + cs = gather bs as -- map (\a -> bs!a) as + ds = gather as bs -- map (\b -> as!b) bs + in (cs, ds) + +npDifficult xs = + let ys = map (+1) xs + in complex ys + +-- filter' = filter + + + + + + + + + + + + + + + +-- awhile' :: Acc (Vector Int) -> Acc (Vector Int) +-- awhile' = awhile (\x -> unit ((x ! I1 0) == 0)) P.id + +-- iffy :: Acc (Vector Int) -> Acc (Vector Int) +-- iffy acc = if size acc == 20 then twoMaps acc else reshape (Z_ ::. 1) (unit 1) +-- twoMaps :: Acc (Vector Int) -> Acc (Vector Int) +-- twoMaps = map (+1) . map (*2) + +-- foo :: Acc (Vector Int) -> Acc (Vector Int) +-- foo a = map (*2) $ if (a ! I1 0) == 2 then map (+1) a else a + + + +-- runningExample :: Acc (Array (Z :. Int) Int) -> Acc (Array (Z :. Int) Int, Array (Z :. Int) Int) +-- runningExample xs = let +-- as = map (`div` 2) xs +-- bs = map (`div` 3) xs +-- cs = map (\a -> bs ! I1 a) as +-- ds = map (\b -> as ! I1 b) bs +-- in T2 cs ds + +-- main :: P.IO () +-- main = +-- P.putStrLn $ test @UniformScheduleFun @InterpretKernel $ zipWith @DIM1 @Int (+) + +-- -- this fails spectacularly: The problem is that we think we fuse, but we don't actually, because +-- -- the Cluster AST doesn't allow zip's style of SoA trickery yet. +-- -- probably need some SoA translations, e.g. in the `take`s in Make/Reqr and in ClusterIO. +-- -- Would then also need to fix cluster construction, which detects fusion +-- -- P.print $ runN @Interpreter (zipWith (+)) (fromList (Z :. 10) [1::Int ..]) (fromList (Z :. 10) [1..]) + +-- --dotp (fromList (Z :. 10) [1..]) (fromList (Z :. 10) [1..]) +-- -- P.print $ run1 @Interpreter (permute (+) (use $ fromList (Z:.10) (P.repeat @P.Int 0)) (\(I1 x) -> Just_ (I1 (x `div` 2)))) $ fromList (Z :. 10) [1..] +-- -- P.putStrLn $ test @UniformScheduleFun @InterpretKernel $ permute (+) (use $ fromList (Z:.10) (P.repeat @P.Int 0)) (\(I1 x) -> Just_ (I1 (x `div` 2))) +-- -- doNTimes 10 P.print +-- -- import qualified Data.Array.Accelerate as A +-- -- import qualified Data.Array.Accelerate.Interpreter as A +-- -- main = do +-- -- putStrLn $ A.test @A.UniformScheduleFun @A.InterpretKernel program +-- -- print $ A.run1 @A.Interpreter program $ A.fromList (A.Z) [20 :: Int] +-- -- where +-- -- program = collatzIndex + +-- -- step :: A.Acc (A.Scalar Int, A.Scalar Int) -> A.Acc (A.Scalar Int, A.Scalar Int) +-- -- step (A.T2 x' count') = A.T2 (A.ifThenElse (x `mod` 2 A.== 0) (A.unit $ x `div` 2) (A.unit $ x * 3 + 1)) (A.unit $ count + 1) +-- -- where +-- -- x = A.the x' +-- -- count = A.the count' + +-- -- collatzIndex :: A.Acc (A.Scalar Int) -> A.Acc (A.Scalar Int) +-- -- collatzIndex input = A.asnd $ A.awhile (A.unit . (A./= 1) . A.the . A.afst) step (A.T2 input $ A.unit 0) + + + + + + + + + +-- program :: Acc (Vector Float) +-- program = let xs = use $ fromList (Z :. 4) [1,2,3,4] +-- xs'' = A.map (+ 1) (use $ fromList (Z :. 3) [0,1,2]) +-- ys = generate (I1 3) (const 2) +-- zs = scatter xs'' (fill (constant (Z :. 4)) 0) xs +-- in A.zipWith (+) ys zs + + +-- main = putStrLn $ test @UniformScheduleFun @InterpretKernel program + + + + + + + + +-- -- This implementation works on 2D points. In future, generalise this to some +-- -- arbitrary "feature vector". +-- -- +-- type Point a = (a, a) + +-- -- Clusters consist of the centroid location as well as its identifier +-- -- +-- type Id = Word32 +-- type Cluster a = (Id, (a, a)) + +-- idOfCluster :: Elt a => Exp (Cluster a) -> Exp Id +-- idOfCluster = A.fst + +-- centroidOfCluster :: Elt a => Exp (Cluster a) -> Exp (Point a) +-- centroidOfCluster = A.snd + +-- -- We'll use this as an intermediate structure; it contains the number of points +-- -- in the set as well as the sum of the x and y coordinates respectively. +-- -- +-- type PointSum a = (Word32, (a, a)) + + +-- -- Get the distance (squared) between two points. Since we only compare this +-- -- value, we elide the square root. +-- -- +-- distance :: A.Num a => Exp (Point a) -> Exp (Point a) -> Exp a +-- distance u v = +-- let (x1,y1) = unlift u +-- (x2,y2) = unlift v +-- in +-- (x1-x2) P.^ (2::Int) + (y1-y2) P.^ (2::Int) + + +-- -- For each of the given points, return the cluster Id that that that point is +-- -- closest to. +-- -- +-- findClosestCluster +-- :: forall a. (A.RealFloat a, P.RealFloat a) +-- => Acc (Vector (Cluster a)) +-- -> Acc (Vector (Point a)) +-- -> Acc (Vector Id) +-- findClosestCluster clusters points = +-- A.map (\p -> A.fst $ A.sfoldl (nearest p) z (constant Z) clusters) points +-- where +-- z = constant (-1, inf) + +-- nearest :: Exp (Point a) -> Exp (Id, a) -> Exp (Cluster a) -> Exp (Id, a) +-- nearest p st c = +-- let d = A.snd st +-- d' = distance p (centroidOfCluster c) +-- in +-- d' A.< d ? ( lift (idOfCluster c, d') , st ) + + +-- -- Given a vector of points and a vector of clusters we, we first locate the +-- -- closest cluster to each point, assign that point to their closest cluster, +-- -- and compute the centroid of the cluster. This yields the new centroid +-- -- locations. +-- -- +-- makeNewClusters +-- :: forall a. (A.RealFloat a, P.RealFloat a, A.FromIntegral Word32 a) +-- => Acc (Vector (Point a)) +-- -> Acc (Vector (Cluster a)) +-- -> Acc (Vector (Cluster a)) +-- makeNewClusters points clusters +-- = pointSumToCluster +-- . makePointSum +-- . findClosestCluster clusters +-- $ points + +-- -- TLM: This setup might be quicker, because it forces the result of +-- -- findClosestCluster to be evaluated, which overall reduces memory +-- -- traffic. However, this is hitting a sharing recovery bug in Accelerate +-- -- so can't be used right now ): +-- -- +-- -- As per tip in issue #148, we can get around the bug by making the first +-- -- argument to pipe closed. It turns out that the first version is quicker! +-- -- +-- -- = A.uncurry findClosestCluster >-> pointSumToCluster . makePointSum $ A.lift (clusters, points) +-- where +-- npts = size points +-- nclusters = size clusters + +-- -- Turn the PointSum intermediate structure into the clusters, by averaging +-- -- the cumulative (x,y) positions. +-- -- +-- pointSumToCluster :: Acc (Vector (PointSum a)) -> Acc (Vector (Cluster a)) +-- pointSumToCluster ps = +-- A.generate (A.shape ps) +-- (\ix -> lift (A.fromIntegral (unindex1 ix), average (ps A.! ix))) + +-- average :: Exp (PointSum a) -> Exp (Point a) +-- average ps = +-- let (n, xy) = unlift ps :: (Exp Word32, Exp (Point a)) +-- (x, y) = unlift xy +-- in +-- lift (x / A.fromIntegral n, y / A.fromIntegral n) -- TLM: what if there are no points in the cluster?? + +-- -- Reduce along the rows of 'pointSum' to get the cumulative (x,y) position +-- -- and number of points assigned to each centroid. +-- -- +-- makePointSum :: Acc (Vector Id) -> Acc (Vector (PointSum a)) +-- makePointSum = A.fold1 addPointSum . compute . pointSum + +-- -- The point sum is an intermediate 2D array (it gets fused away, so does +-- -- not exist in memory). The points are laid out along the innermost +-- -- dimension (rows), and down the column is each of the clusters. +-- -- +-- -- For each point, we put its (x,y) coordinates into the row corresponding +-- -- to whichever cluster it is closest to, and zeros in each of the other +-- -- rows. +-- -- +-- pointSum :: Acc (Vector Id) -> Acc (A.Array DIM2 (PointSum a)) +-- pointSum nearest = +-- A.generate (lift (Z:.nclusters:.npts)) +-- (\ix -> let Z:.i:.j = unlift ix :: Z :. Exp Int :. Exp Int +-- near = nearest A.! index1 j + +-- yes = lift (constant 1, points A.! index1 j) +-- no = constant (0, (0, 0)) +-- in +-- near A.== A.fromIntegral i ? ( yes, no )) + +-- addPointSum :: Exp (PointSum a) -> Exp (PointSum a) -> Exp (PointSum a) +-- addPointSum x y = +-- let (c1, u) = unlift x :: (Exp Word32, Exp (Point a)) +-- (c2, v) = unlift y :: (Exp Word32, Exp (Point a)) +-- (x1,y1) = unlift u :: (Exp a, Exp a) +-- (x2,y2) = unlift v :: (Exp a, Exp a) +-- in +-- lift (c1+c2, lift (x1+x2, y1+y2) :: Exp (Point a)) + +-- {-- +-- -- Alternative to computing the PointSum structure. +-- -- +-- -- This method uses a forward permutation with atomic instructions to create +-- -- the array directly (this method is closer to what one might write +-- -- sequentially). This avoids a parallel reduction, but has very high +-- -- contention. Overall performance much lower, as: +-- -- +-- -- number of clusters << number of points +-- -- +-- makePointSum :: Acc (Vector (PointSum a)) +-- makePointSum = A.permute addPointSum zeros near input +-- where +-- zeros = A.fill (constant (Z:.nclusters)) (constant (0,(0,0))) +-- input = A.zip (A.fill (A.shape points) (constant 1)) points +-- near ix = index1 (A.fromIntegral (nearest ! ix)) +-- --} + + +-- -- To complete the k-means algorithm, we loop repeatedly generating new clusters +-- -- positions, until the positions converge (or some maximum iteration limit is +-- -- reached?) +-- -- +-- kmeans :: forall a. (A.RealFloat a, P.RealFloat a, A.FromIntegral Word32 a) +-- => Acc (Vector (Point a)) -- the points to cluster +-- -> Acc (Vector (Cluster a)) -- initial cluster positions (guess) +-- -> Acc (Vector (Cluster a)) +-- kmeans points clusters +-- = A.asnd +-- $ A.awhile (A.uncurry keepGoing) +-- (\cs -> let (_, old) = unlift cs :: (Acc (Vector (Cluster a)), Acc (Vector (Cluster a))) +-- new = makeNewClusters points old +-- in +-- lift (old,new)) +-- (lift (clusters, makeNewClusters points clusters)) +-- where +-- keepGoing :: Acc (Vector (Cluster a)) -> Acc (Vector (Cluster a)) -> Acc (Scalar Bool) +-- keepGoing xs ys +-- = A.or +-- $ A.zipWith (\c1 c2 -> let (x1,y1) = unlift (centroidOfCluster c1) +-- (x2,y2) = unlift (centroidOfCluster c2) +-- in +-- abs (x1-x2) A.> 0.01 A.|| abs (y1-y2) A.> 0.01) xs ys + + +-- -- The largest non-infinite floating point number +-- -- +-- inf :: forall a. P.RealFloat a => a +-- inf = P.encodeFloat m n +-- where +-- a = undefined :: a +-- b = P.floatRadix a +-- e = P.floatDigits a +-- (_, e') = P.floatRange a +-- m = b P.^ e - 1 +-- n = e' - e + +-- main :: IO () +-- main = do +-- putStrLn $ test @UniformScheduleFun @InterpretKernel $ A.map @DIM2 (*(2::Exp Int)) +-- -- inputs <- (P.&&) <$> doesFileExist "points.bin" +-- -- <*> doesFileExist "clusters" +-- -- unless inputs $ do +-- -- print "Running the GenSamples program first to generate random data" +-- -- main' + +-- -- points' <- decodeFile "points.bin" +-- -- initial' <- read `fmap` readFile "clusters" + +-- -- let nclusters = P.length initial' +-- -- npoints = P.length points' + +-- -- solve = run1 @Interpreter (kmeans (use points)) + +-- -- initial :: Vector (Cluster Float) +-- -- initial = A.fromList (Z:.nclusters) initial' + +-- -- points :: Vector (Point Float) +-- -- points = A.fromList (Z:.npoints) points' + +-- -- -- Warm up first by printing the expected results +-- -- -- +-- -- putStrLn $ "number of points: " P.++ show npoints +-- -- putStrLn $ "final clusters:\n" P.++ +-- -- unlines (P.map show (A.toList (solve initial))) + +-- -- print $ solve initial + + +-- type Point' = (Float, Float) + +-- zeroPoint :: Point' +-- zeroPoint = (0,0) + + +-- -- Clusters +-- ----------- + +-- type Cluster' = (Word32, Point') + +-- makeCluster :: Int -> [Point'] -> Cluster' +-- makeCluster clid points +-- = ( P.fromIntegral clid +-- , (a / P.fromIntegral count, b / P.fromIntegral count)) +-- where +-- (a,b) = foldl' addPoint zeroPoint points +-- count = P.length points + +-- addPoint :: Point' -> Point' -> Point' +-- addPoint (x,y) (u,v) = (x+u,y+v) + + +-- -- Generate random points +-- -- ---------------------- + +-- minX, maxX, minY, maxY, minSD, maxSD :: Float +-- minX = -10 +-- maxX = 10 +-- minY = -10 +-- maxY = 10 +-- minSD = 1.5 +-- maxSD = 2.0 + +-- main' :: IO () +-- main' = do +-- -- n: minp: maxp: rest <- fmap (fmap read) getArgs +-- let n = 5 +-- minp = 5 +-- maxp = 5 +-- rest = [] + +-- case rest of +-- [seed] -> setStdGen (mkStdGen seed) +-- _ -> return () + +-- nps <- replicateM n (randomRIO (minp, maxp)) +-- xs <- replicateM n (randomRIO (minX, maxX)) +-- ys <- replicateM n (randomRIO (minY, maxY)) +-- sds <- replicateM n (randomRIO (minSD, maxSD)) + +-- let params = Data.List.zip5 nps xs ys sds sds + +-- -- first generate a set of points for each set of sample parameters +-- ss <- mapM (\(a,b,c,d,e) -> generate2DSamples a b c d e) params +-- let points = concat ss + +-- -- dump all the points into the file "points" +-- hsamp <- openFile "points" WriteMode +-- mapM_ (printPoint hsamp) points +-- hClose hsamp + +-- encodeFile "points.bin" points + +-- -- generate the initial clusters by assigning each point to random +-- -- cluster. +-- gen <- newStdGen +-- let +-- rand_clusters = randomRs (0,n-1) gen :: [Int] +-- arr = accumArray (flip (:)) [] (0,n-1) $ +-- P.zip rand_clusters points +-- clusters = P.map (P.uncurry makeCluster) (assocs arr) +-- writeFile "clusters" (show clusters) + +-- -- so we can tell what the answer should be: +-- writeFile "params" (show params) + + +-- printPoint :: Handle -> Point' -> IO () +-- printPoint h (x,y) = do +-- hPutStr h (show x) +-- hPutChar h ' ' +-- hPutStr h (show y) +-- hPutChar h '\n' + +-- generate2DSamples +-- :: Int -- number of samples to generate +-- -> Float -> Float -- X and Y of the mean +-- -> Float -> Float -- X and Y standard deviations +-- -> IO [Point'] +-- generate2DSamples n mx my sdx sdy = do +-- gen <- newStdGen +-- let (genx, geny) = split gen +-- xsamples = normals' (mx,sdx) genx +-- ysamples = normals' (my,sdy) geny +-- return (P.zipWith (,) (P.take n xsamples) ysamples) + + + + + diff --git a/test/nofib/Main.hs b/test/nofib/Main.hs index 03fa9e41d..9597ccbb3 100644 --- a/test/nofib/Main.hs +++ b/test/nofib/Main.hs @@ -11,7 +11,7 @@ module Main where import Data.Array.Accelerate.Test.NoFib -import Data.Array.Accelerate.Interpreter +import Data.Array.Accelerate.Interpreter.Simple import System.Environment import Test.Tasty @@ -21,4 +21,3 @@ main = do setEnv "TASTY_INTERPRETER" "True" setEnv "TASTY_HEDGEHOG_TESTS" "50" nofib runN (testGroup "runQ" []) -