futhark 0.25.1 → 0.25.2

raw patch · 52 files changed

+1787/−1776 lines, 52 filesdep ~lspdep ~lsp-typesPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

Dependency ranges changed: lsp, lsp-types

API changes (from Hackage documentation)

Files

• docs/man/futhark-repl.rst +2/−1
• docs/usage.rst +2/−2
• futhark.cabal +4/−3
• rts/c/atomics.h +57/−0
• rts/c/backends/cuda.h +35/−38
• src/Futhark/CLI/C.hs +2/−2
• src/Futhark/CLI/CUDA.hs +2/−2
• src/Futhark/CLI/Dev.hs +5/−5
• src/Futhark/CLI/Misc.hs +2/−2
• src/Futhark/CLI/Multicore.hs +2/−2
• src/Futhark/CLI/MulticoreISPC.hs +2/−2
• src/Futhark/CLI/MulticoreWASM.hs +2/−2
• src/Futhark/CLI/OpenCL.hs +2/−2
• src/Futhark/CLI/PyOpenCL.hs +1/−1
• src/Futhark/CLI/Python.hs +1/−1
• src/Futhark/CLI/REPL.hs +9/−8
• src/Futhark/CLI/WASM.hs +2/−2
• src/Futhark/CodeGen/Backends/CCUDA.hs +98/−73
• src/Futhark/CodeGen/Backends/CCUDA/Boilerplate.hs +3/−3
• src/Futhark/CodeGen/Backends/COpenCL.hs +80/−85
• src/Futhark/CodeGen/Backends/SimpleRep.hs +3/−3
• src/Futhark/CodeGen/ImpCode.hs +4/−1
• src/Futhark/CodeGen/ImpGen.hs +79/−61
• src/Futhark/CodeGen/ImpGen/GPU.hs +141/−35
• src/Futhark/CodeGen/ImpGen/GPU/Base.hs +2/−30
• src/Futhark/CodeGen/ImpGen/GPU/Copy.hs +81/−0
• src/Futhark/CodeGen/ImpGen/GPU/Group.hs +4/−3
• src/Futhark/IR/Mem.hs +52/−38
• src/Futhark/IR/Mem/IxFun.hs +129/−510
• src/Futhark/IR/Mem/LMAD.hs +301/−100
• src/Futhark/IR/Parse.hs +5/−17
• src/Futhark/Internalise/ReplaceRecords.hs +1/−0
• src/Futhark/Internalise/TypesValues.hs +1/−0
• src/Futhark/LSP/Handlers.hs +2/−2
• src/Futhark/Optimise/ArrayShortCircuiting/ArrayCoalescing.hs +13/−15
• src/Futhark/Optimise/ArrayShortCircuiting/MemRefAggreg.hs +8/−8
• src/Futhark/Optimise/ArrayShortCircuiting/TopdownAnalysis.hs +9/−8
• src/Futhark/Optimise/BlkRegTiling.hs +4/−4
• src/Futhark/Optimise/Simplify/Rules/Simple.hs +10/−3
• src/Futhark/Optimise/TileLoops/Shared.hs +13/−17
• src/Futhark/Pass/ExpandAllocations.hs +43/−25
• src/Futhark/Pass/ExplicitAllocations.hs +79/−46
• src/Futhark/Passes.hs +18/−18
• src/Futhark/Util/Pretty.hs +1/−1
• src/Language/Futhark/Parser.hs +9/−25
• src/Language/Futhark/Parser/Lexer.x +112/−52
• src/Language/Futhark/Parser/Lexer/Tokens.hs +53/−81
• src/Language/Futhark/Parser/Lexer/Wrapper.hs +9/−75
• src/Language/Futhark/Parser/Monad.hs +30/−90
• src/Language/Futhark/Parser/Parser.y +17/−16
• unittests/Futhark/IR/Mem/IxFunTests.hs +232/−248
• unittests/Futhark/IR/Mem/IxFunWrapper.hs +9/−8

@@ -17,7 +17,8 @@ Start an interactive Futhark session.  This will let you interactively enter expressions and declarations which are then immediately interpreted.  If the entered line can be either a declaration or an-expression, it is assumed to be a declaration.+expression, it is assumed to be a declaration.  The input msut fit on+a single line.  Futhark source files can be loaded using the ``:load`` command.  This will erase any interactively entered definitions.  Use the ``:help``

@@ -482,9 +482,9 @@ When using the OpenCL backend, extra API functions are provided for directly accessing or providing the OpenCL objects used by Futhark. Take care when using these functions.  In particular, a Futhark-context can now be provided with the command queue to use::+context can now be configured with the command queue to use:: -  struct futhark_context *futhark_context_new_with_command_queue(struct futhark_context_config *cfg, cl_command_queue queue);+  void futhark_context_config_set_command_queue(struct futhark_context_config *cfg, cl_command_queue queue);  As a ``cl_command_queue`` specifies an OpenCL device, this is also how manual platform and device selection is possible.  A function is also

@@ -1,6 +1,6 @@ cabal-version: 2.4 name:           futhark-version:        0.25.1+version:        0.25.2 synopsis:       An optimising compiler for a functional, array-oriented language.  description:    Futhark is a small programming language designed to be compiled to@@ -248,6 +248,7 @@       Futhark.CodeGen.ImpGen.CUDA       Futhark.CodeGen.ImpGen.GPU       Futhark.CodeGen.ImpGen.GPU.Base+      Futhark.CodeGen.ImpGen.GPU.Copy       Futhark.CodeGen.ImpGen.GPU.Group       Futhark.CodeGen.ImpGen.GPU.SegHist       Futhark.CodeGen.ImpGen.GPU.SegMap@@ -494,8 +495,8 @@     , haskeline     , language-c-quote >= 0.12     , lens-    , lsp >= 2.0.0.0-    , lsp-types >= 2.0.0.0+    , lsp >= 2.1.0.0+    , lsp-types >= 2.0.1.0     , mainland-pretty >=0.7.1     , cmark-gfm >=0.2.1     , megaparsec >=9.0.0

@@ -1,5 +1,30 @@ // Start of atomics.h +inline int32_t atomic_xchg_i32_global(volatile __global int32_t *p, int32_t x);+inline int32_t atomic_xchg_i32_local(volatile __local int32_t *p, int32_t x);+inline int32_t atomic_cmpxchg_i32_global(volatile __global int32_t *p,+                                         int32_t cmp, int32_t val);+inline int32_t atomic_cmpxchg_i32_local(volatile __local int32_t *p,+                                        int32_t cmp, int32_t val);+inline int32_t atomic_add_i32_global(volatile __global int32_t *p, int32_t x);+inline int32_t atomic_add_i32_local(volatile __local int32_t *p, int32_t x);+inline float atomic_fadd_f32_global(volatile __global float *p, float x);+inline float atomic_fadd_f32_local(volatile __local float *p, float x);+inline int32_t atomic_smax_i32_global(volatile __global int32_t *p, int32_t x);+inline int32_t atomic_smax_i32_local(volatile __local int32_t *p, int32_t x);+inline int32_t atomic_smin_i32_global(volatile __global int32_t *p, int32_t x);+inline int32_t atomic_smin_i32_local(volatile __local int32_t *p, int32_t x);+inline uint32_t atomic_umax_i32_global(volatile __global uint32_t *p, uint32_t x);+inline uint32_t atomic_umax_i32_local(volatile __local uint32_t *p, uint32_t x);+inline uint32_t atomic_umin_i32_global(volatile __global uint32_t *p, uint32_t x);+inline uint32_t atomic_umin_i32_local(volatile __local uint32_t *p, uint32_t x);+inline int32_t atomic_and_i32_global(volatile __global int32_t *p, int32_t x);+inline int32_t atomic_and_i32_local(volatile __local int32_t *p, int32_t x);+inline int32_t atomic_or_i32_global(volatile __global int32_t *p, int32_t x);+inline int32_t atomic_or_i32_local(volatile __local int32_t *p, int32_t x);+inline int32_t atomic_xor_i32_global(volatile __global int32_t *p, int32_t x);+inline int32_t atomic_xor_i32_local(volatile __local int32_t *p, int32_t x);+ inline int32_t atomic_xchg_i32_global(volatile __global int32_t *p, int32_t x) { #ifdef FUTHARK_CUDA   return atomicExch((int32_t*)p, x);@@ -196,6 +221,36 @@  // Start of 64 bit atomics +#if defined(FUTHARK_CUDA) || defined(cl_khr_int64_base_atomics) && defined(cl_khr_int64_extended_atomics)++inline int64_t atomic_xchg_i64_global(volatile __global int64_t *p, int64_t x);+inline int64_t atomic_xchg_i64_local(volatile __local int64_t *p, int64_t x);+inline int64_t atomic_cmpxchg_i64_global(volatile __global int64_t *p,+                                         int64_t cmp, int64_t val);+inline int64_t atomic_cmpxchg_i64_local(volatile __local int64_t *p,+                                        int64_t cmp, int64_t val);+inline int64_t atomic_add_i64_global(volatile __global int64_t *p, int64_t x);+inline int64_t atomic_add_i64_local(volatile __local int64_t *p, int64_t x);+inline int64_t atomic_smax_i64_global(volatile __global int64_t *p, int64_t x);+inline int64_t atomic_smax_i64_local(volatile __local int64_t *p, int64_t x);+inline int64_t atomic_smin_i64_global(volatile __global int64_t *p, int64_t x);+inline int64_t atomic_smin_i64_local(volatile __local int64_t *p, int64_t x);+inline uint64_t atomic_umax_i64_global(volatile __global uint64_t *p, uint64_t x);+inline uint64_t atomic_umax_i64_local(volatile __local uint64_t *p, uint64_t x);+inline uint64_t atomic_umin_i64_global(volatile __global uint64_t *p, uint64_t x);+uint64_t atomic_umin_i64_local(volatile __local uint64_t *p, uint64_t x);+inline int64_t atomic_and_i64_global(volatile __global int64_t *p, int64_t x);+inline int64_t atomic_and_i64_local(volatile __local int64_t *p, int64_t x);+inline int64_t atomic_or_i64_global(volatile __global int64_t *p, int64_t x);+inline int64_t atomic_or_i64_local(volatile __local int64_t *p, int64_t x);+inline int64_t atomic_xor_i64_global(volatile __global int64_t *p, int64_t x);+inline int64_t atomic_xor_i64_local(volatile __local int64_t *p, int64_t x);++#ifdef FUTHARK_F64_ENABLED+inline double atomic_fadd_f64_global(volatile __global double *p, double x);+inline double atomic_fadd_f64_local(volatile __local double *p, double x);+#endif+ inline int64_t atomic_xchg_i64_global(volatile __global int64_t *p, int64_t x) { #ifdef FUTHARK_CUDA   return atomicExch((uint64_t*)p, x);@@ -393,5 +448,7 @@   return atom_xor(p, x); #endif }++#endif // defined(FUTHARK_CUDA) || defined(cl_khr_int64_base_atomics) && defined(cl_khr_int64_extended_atomics)  // End of atomics.h

@@ -251,10 +251,10 @@   lock_t error_lock;   FILE *log;   struct constants *constants;-  struct free_list cu_free_list;+  struct free_list free_list;   int64_t peak_mem_usage_default;   int64_t cur_mem_usage_default;-  // Uniform above+  // Uniform fields above.    CUdeviceptr global_failure;   CUdeviceptr global_failure_args;@@ -270,8 +270,9 @@   CUdevice dev;   CUcontext cu_ctx;   CUmodule module;+  CUstream stream; -  struct free_list free_list;+  struct free_list cu_free_list;    size_t max_block_size;   size_t max_grid_size;@@ -307,8 +308,8 @@   struct futhark_context_config *cfg = ctx->cfg;   char name[256];   int count, chosen = -1, best_cc = -1;-  int cc_major_best, cc_minor_best;-  int cc_major, cc_minor;+  int cc_major_best = 0, cc_minor_best = 0;+  int cc_major = 0, cc_minor = 0;   CUdevice dev;    CUDA_SUCCEED_FATAL(cuDeviceGetCount(&count));@@ -385,11 +386,11 @@ }  static const char *cuda_nvrtc_get_arch(CUdevice dev) {-  struct {+  static struct {     int major;     int minor;     const char *arch_str;-  } static const x[] = {+  } const x[] = {     { 3, 0, "compute_30" },     { 3, 2, "compute_32" },     { 3, 5, "compute_35" },@@ -730,38 +731,15 @@   return NULL; } -static char* cuda_setup(struct futhark_context *ctx, const char *src_fragments[],-                        const char *extra_opts[], const char* cache_fname) {-  CUDA_SUCCEED_FATAL(cuInit(0));--  if (cuda_device_setup(ctx) != 0) {-    futhark_panic(-1, "No suitable CUDA device found.\n");-  }-  CUDA_SUCCEED_FATAL(cuCtxCreate(&ctx->cu_ctx, 0, ctx->dev));--  free_list_init(&ctx->cu_free_list);--  ctx->max_shared_memory = device_query(ctx->dev, MAX_SHARED_MEMORY_PER_BLOCK);-  ctx->max_block_size = device_query(ctx->dev, MAX_THREADS_PER_BLOCK);-  ctx->max_grid_size = device_query(ctx->dev, MAX_GRID_DIM_X);-  ctx->max_tile_size = sqrt(ctx->max_block_size);-  ctx->max_threshold = 0;-  ctx->max_bespoke = 0;-  ctx->lockstep_width = device_query(ctx->dev, WARP_SIZE);--  cuda_size_setup(ctx);-  return cuda_module_setup(ctx, src_fragments, extra_opts, cache_fname);-}- // Count up the runtime all the profiling_records that occured during execution. // Also clears the buffer of profiling_records.-static cudaError_t cuda_tally_profiling_records(struct futhark_context *ctx) {-  cudaError_t err;+static CUresult cuda_tally_profiling_records(struct futhark_context *ctx) {+  CUresult err;   for (int i = 0; i < ctx->profiling_records_used; i++) {     struct profiling_record record = ctx->profiling_records[i];      float ms;-    if ((err = cudaEventElapsedTime(&ms, record.events[0], record.events[1])) != cudaSuccess) {+    if ((err = cuEventElapsedTime(&ms, record.events[0], record.events[1])) != CUDA_SUCCESS) {       return err;     } @@ -769,10 +747,10 @@     *record.runs += 1;     *record.runtime += ms*1000; -    if ((err = cudaEventDestroy(record.events[0])) != cudaSuccess) {+    if ((err = cuEventDestroy(record.events[0])) != CUDA_SUCCESS) {       return err;     }-    if ((err = cudaEventDestroy(record.events[1])) != cudaSuccess) {+    if ((err = cuEventDestroy(record.events[1])) != CUDA_SUCCESS) {       return err;     } @@ -781,7 +759,7 @@    ctx->profiling_records_used = 0; -  return cudaSuccess;+  return CUDA_SUCCESS; }  // Returns pointer to two events.@@ -918,10 +896,28 @@   ctx->peak_mem_usage_device = 0;   ctx->cur_mem_usage_device = 0; -  ctx->error = cuda_setup(ctx, cuda_program, ctx->cfg->nvrtc_opts, ctx->cfg->cache_fname);+  CUDA_SUCCEED_FATAL(cuInit(0));+  if (cuda_device_setup(ctx) != 0) {+    futhark_panic(-1, "No suitable CUDA device found.\n");+  }+  CUDA_SUCCEED_FATAL(cuCtxCreate(&ctx->cu_ctx, 0, ctx->dev)); +  free_list_init(&ctx->cu_free_list);++  ctx->max_shared_memory = device_query(ctx->dev, MAX_SHARED_MEMORY_PER_BLOCK);+  ctx->max_block_size = device_query(ctx->dev, MAX_THREADS_PER_BLOCK);+  ctx->max_grid_size = device_query(ctx->dev, MAX_GRID_DIM_X);+  ctx->max_tile_size = sqrt(ctx->max_block_size);+  ctx->max_threshold = 0;+  ctx->max_bespoke = 0;+  ctx->lockstep_width = device_query(ctx->dev, WARP_SIZE);+  CUDA_SUCCEED_FATAL(cuStreamCreate(&ctx->stream, CU_STREAM_DEFAULT));+  cuda_size_setup(ctx);+  ctx->error = cuda_module_setup(ctx, cuda_program,+                                 ctx->cfg->nvrtc_opts, ctx->cfg->cache_fname);+   if (ctx->error != NULL) {-    futhark_panic(1, "%s\n", ctx->error);+    futhark_panic(1, "During CUDA initialisation:\n%s\n", ctx->error);   }    int32_t no_error = -1;@@ -938,6 +934,7 @@   CUDA_SUCCEED_FATAL(cuda_free_all(ctx));   (void)cuda_tally_profiling_records(ctx);   free(ctx->profiling_records);+  CUDA_SUCCEED_FATAL(cuStreamDestroy(ctx->stream));   CUDA_SUCCEED_FATAL(cuModuleUnload(ctx->module));   CUDA_SUCCEED_FATAL(cuCtxDestroy(ctx->cu_ctx)); }

@@ -3,7 +3,7 @@  import Futhark.Actions (compileCAction) import Futhark.Compiler.CLI-import Futhark.Passes (sequentialCpuPipeline)+import Futhark.Passes (seqmemPipeline)  -- | Run @futhark c@ main :: String -> [String] -> IO ()@@ -12,6 +12,6 @@   []   "Compile sequential C"   "Generate sequential C code from optimised Futhark program."-  sequentialCpuPipeline+  seqmemPipeline   $ \fcfg () mode outpath prog ->     actionProcedure (compileCAction fcfg mode outpath) prog

@@ -3,7 +3,7 @@  import Futhark.Actions (compileCUDAAction) import Futhark.Compiler.CLI-import Futhark.Passes (gpuPipeline)+import Futhark.Passes (gpumemPipeline)  -- | Run @futhark cuda@. main :: String -> [String] -> IO ()@@ -12,6 +12,6 @@   []   "Compile CUDA"   "Generate CUDA/C code from optimised Futhark program."-  gpuPipeline+  gpumemPipeline   $ \fcfg () mode outpath prog ->     actionProcedure (compileCUDAAction fcfg mode outpath) prog

@@ -639,7 +639,7 @@       "GPU"       GPU       "Run the default optimised kernels pipeline"-      kernelsPipeline+      gpuPipeline       []       ["gpu"],     pipelineOption@@ -647,7 +647,7 @@       "GPUMem"       GPUMem       "Run the full GPU compilation pipeline"-      gpuPipeline+      gpumemPipeline       []       ["gpu-mem"],     pipelineOption@@ -655,7 +655,7 @@       "Seq"       Seq       "Run the sequential CPU compilation pipeline"-      sequentialPipeline+      seqPipeline       []       ["seq"],     pipelineOption@@ -663,7 +663,7 @@       "SeqMem"       SeqMem       "Run the sequential CPU+memory compilation pipeline"-      sequentialCpuPipeline+      seqmemPipeline       []       ["seq-mem"],     pipelineOption@@ -679,7 +679,7 @@       "MCMem"       MCMem       "Run the multicore+memory compilation pipeline"-      multicorePipeline+      mcmemPipeline       []       ["mc-mem"]   ]

@@ -22,7 +22,7 @@ import Futhark.Util (hashText, interactWithFileSafely) import Futhark.Util.Options import Futhark.Util.Pretty (prettyTextOneLine)-import Language.Futhark.Parser.Lexer (scanTokens)+import Language.Futhark.Parser.Lexer (scanTokensText) import Language.Futhark.Prop (isBuiltin) import Language.Futhark.Semantic (includeToString) import System.Environment (getExecutablePath)@@ -115,7 +115,7 @@ mainTokens = mainWithOptions () [] "program" $ \args () ->   case args of     [file] -> Just $ do-      res <- interactWithFileSafely (scanTokens (startPos file) <$> BS.readFile file)+      res <- interactWithFileSafely (scanTokensText (startPos file) <$> T.readFile file)       case res of         Nothing -> do           hPutStrLn stderr $ file <> ": file not found."

@@ -3,7 +3,7 @@  import Futhark.Actions (compileMulticoreAction) import Futhark.Compiler.CLI-import Futhark.Passes (multicorePipeline)+import Futhark.Passes (mcmemPipeline)  -- | Run @futhark multicore@. main :: String -> [String] -> IO ()@@ -12,6 +12,6 @@   []   "Compile to multicore C"   "Generate multicore C code from optimised Futhark program."-  multicorePipeline+  mcmemPipeline   $ \fcfg () mode outpath prog ->     actionProcedure (compileMulticoreAction fcfg mode outpath) prog

@@ -3,7 +3,7 @@  import Futhark.Actions (compileMulticoreToISPCAction) import Futhark.Compiler.CLI-import Futhark.Passes (multicorePipeline)+import Futhark.Passes (mcmemPipeline)  -- | Run @futhark multicore@. main :: String -> [String] -> IO ()@@ -12,6 +12,6 @@   []   "Compile to multicore ISPC"   "Generate multicore ISPC code from optimised Futhark program."-  multicorePipeline+  mcmemPipeline   $ \fcfg () mode outpath prog ->     actionProcedure (compileMulticoreToISPCAction fcfg mode outpath) prog

@@ -3,7 +3,7 @@  import Futhark.Actions (compileMulticoreToWASMAction) import Futhark.Compiler.CLI-import Futhark.Passes (multicorePipeline)+import Futhark.Passes (mcmemPipeline)  -- | Run @futhark c@ main :: String -> [String] -> IO ()@@ -12,6 +12,6 @@   []   "Compile to multicore WASM"   "Generate multicore WASM with the multicore C backend code from optimised Futhark program."-  multicorePipeline+  mcmemPipeline   $ \fcfg () mode outpath prog ->     actionProcedure (compileMulticoreToWASMAction fcfg mode outpath) prog

@@ -3,7 +3,7 @@  import Futhark.Actions (compileOpenCLAction) import Futhark.Compiler.CLI-import Futhark.Passes (gpuPipeline)+import Futhark.Passes (gpumemPipeline)  -- | Run @futhark opencl@ main :: String -> [String] -> IO ()@@ -12,6 +12,6 @@   []   "Compile OpenCL"   "Generate OpenCL/C code from optimised Futhark program."-  gpuPipeline+  gpumemPipeline   $ \fcfg () mode outpath prog ->     actionProcedure (compileOpenCLAction fcfg mode outpath) prog

@@ -12,6 +12,6 @@   []   "Compile PyOpenCL"   "Generate Python + OpenCL code from optimised Futhark program."-  gpuPipeline+  gpumemPipeline   $ \fcfg () mode outpath prog ->     actionProcedure (compilePyOpenCLAction fcfg mode outpath) prog

@@ -12,6 +12,6 @@   []   "Compile sequential Python"   "Generate sequential Python code from optimised Futhark program."-  sequentialCpuPipeline+  seqmemPipeline   $ \fcfg () mode outpath prog ->     actionProcedure (compilePythonAction fcfg mode outpath) prog

@@ -20,7 +20,7 @@ import Futhark.MonadFreshNames import Futhark.Util (fancyTerminal) import Futhark.Util.Options-import Futhark.Util.Pretty (AnsiStyle, Color (..), Doc, align, annotate, bgColorDull, bold, brackets, color, docText, docTextForHandle, hardline, oneLine, pretty, putDoc, putDocLn, unAnnotate, (<+>))+import Futhark.Util.Pretty (AnsiStyle, Color (..), Doc, align, annotate, bgColorDull, bold, brackets, color, docText, docTextForHandle, hardline, italicized, oneLine, pretty, putDoc, putDocLn, unAnnotate, (<+>)) import Futhark.Version import Language.Futhark import Language.Futhark.Interpreter qualified as I@@ -220,9 +220,7 @@               <> mconcat (intersperse ", " (map fst matches))     _ -> do       -- Read a declaration or expression.-      maybe_dec_or_e <- parseDecOrExpIncrM (inputLine "  ") prompt line--      case maybe_dec_or_e of+      case parseDecOrExp prompt line of         Left (SyntaxError _ err) -> liftIO $ T.putStrLn err         Right (Left d) -> onDec d         Right (Right e) -> onExp e@@ -396,10 +394,13 @@  typeCommand :: Command typeCommand = genTypeCommand parseExp T.checkExp $ \(ps, e) ->-  oneLine (pretty e)-    <> foldMap ((" " <>) . pretty) ps-    <> " : "-    <> oneLine (pretty (typeOf e))+  oneLine (pretty (typeOf e))+    <> if not (null ps)+      then+        annotate italicized $+          "\n\nPolymorphic in"+            <+> mconcat (intersperse " " $ map pretty ps) <> "."+      else mempty  mtypeCommand :: Command mtypeCommand = genTypeCommand parseModExp T.checkModExp $ pretty . fst

@@ -3,7 +3,7 @@  import Futhark.Actions (compileCtoWASMAction) import Futhark.Compiler.CLI-import Futhark.Passes (sequentialCpuPipeline)+import Futhark.Passes (seqmemPipeline)  -- | Run @futhark c@ main :: String -> [String] -> IO ()@@ -12,6 +12,6 @@   []   "Compile to WASM"   "Generate WASM with the sequential C backend code from optimised Futhark program."-  sequentialCpuPipeline+  seqmemPipeline   $ \fcfg () mode outpath prog ->     actionProcedure (compileCtoWASMAction fcfg mode outpath) prog

@@ -11,6 +11,7 @@ where  import Control.Monad+import Data.List (unzip4) import Data.Maybe (catMaybes) import Data.Text qualified as T import Futhark.CodeGen.Backends.CCUDA.Boilerplate@@ -152,7 +153,7 @@                 cuMemcpyHtoDAsync($exp:mem + $exp:idx * sizeof($ty:t),                                   &$id:val',                                   sizeof($ty:t),-                                  0));+                                  ctx->stream));               $items:aft              }|] writeCUDAScalar mem idx t "device" _ val = do@@ -226,9 +227,9 @@     memcpyFun _ (Space "device") (Space "device") =       ([C.cexp|cuMemcpy($exp:dst, $exp:src, $exp:nbytes)|], copyDevToDev)     memcpyFun GC.CopyNoBarrier DefaultSpace (Space "device") =-      ([C.cexp|cuMemcpyDtoHAsync($exp:dst, $exp:src, $exp:nbytes, 0)|], copyDevToHost)+      ([C.cexp|cuMemcpyDtoHAsync($exp:dst, $exp:src, $exp:nbytes, ctx->stream)|], copyDevToHost)     memcpyFun GC.CopyNoBarrier (Space "device") DefaultSpace =-      ([C.cexp|cuMemcpyHtoDAsync($exp:dst, $exp:src, $exp:nbytes, 0)|], copyHostToDev)+      ([C.cexp|cuMemcpyHtoDAsync($exp:dst, $exp:src, $exp:nbytes, ctx->stream)|], copyHostToDev)     memcpyFun _ _ _ =       error $         "Cannot copy to '"@@ -274,21 +275,81 @@   let e = kernelConstToExp $ SizeMaxConst size_class   GC.stm [C.cstm|$id:v = $exp:e;|] callKernel (LaunchKernel safety kernel_name args num_blocks block_size) = do-  args_arr <- newVName "kernel_args"-  time_start <- newVName "time_start"-  time_end <- newVName "time_end"-  (args', shared_vars) <- mapAndUnzipM mkArgs args+  (arg_params, arg_params_inits, call_args, shared_vars) <-+    unzip4 <$> zipWithM mkArgs [(0 :: Int) ..] args   let (shared_sizes, shared_offsets) = unzip $ catMaybes shared_vars       shared_offsets_sc = mkOffsets shared_sizes       shared_args = zip shared_offsets shared_offsets_sc-      shared_tot = last shared_offsets_sc+      shared_bytes = last shared_offsets_sc   forM_ shared_args $ \(arg, offset) ->     GC.decl [C.cdecl|unsigned int $id:arg = $exp:offset;|]    (grid_x, grid_y, grid_z) <- mkDims <$> mapM GC.compileExp num_blocks   (block_x, block_y, block_z) <- mkDims <$> mapM compileGroupDim block_size-  let perm_args-        | length num_blocks == 3 = [[C.cinit|&perm[0]|], [C.cinit|&perm[1]|], [C.cinit|&perm[2]|]]++  let need_perm = length num_blocks == 3+  kernel_fname <- genKernelFunction kernel_name safety need_perm arg_params arg_params_inits++  GC.stm+    [C.cstm|{+           err = $id:kernel_fname(ctx,+                                  $exp:grid_x,$exp:grid_y,$exp:grid_z,+                                  $exp:block_x, $exp:block_y, $exp:block_z,+                                  $exp:shared_bytes,+                                  $args:call_args);+           if (err != FUTHARK_SUCCESS) { goto cleanup; }+           }|]++  when (safety >= SafetyFull) $+    GC.stm [C.cstm|ctx->failure_is_an_option = 1;|]+  where+    mkDims [] = ([C.cexp|0|], [C.cexp|0|], [C.cexp|0|])+    mkDims [x] = (x, [C.cexp|1|], [C.cexp|1|])+    mkDims [x, y] = (x, y, [C.cexp|1|])+    mkDims (x : y : z : _) = (x, y, z)+    addExp x y = [C.cexp|$exp:x + $exp:y|]+    alignExp e = [C.cexp|$exp:e + ((8 - ($exp:e % 8)) % 8)|]+    mkOffsets = scanl (\a b -> a `addExp` alignExp b) [C.cexp|0|]+    mkArgs i (ValueKArg e t) = do+      e' <- GC.compileExp e+      pure+        ( [C.cparam|$ty:(primStorageType t) $id:("arg" <> show i)|],+          [C.cinit|&$id:("arg" <> show i)|],+          toStorage t e',+          Nothing+        )+    mkArgs i (MemKArg v) = do+      v' <- GC.rawMem v+      pure+        ( [C.cparam|typename CUdeviceptr $id:("arg" <> show i)|],+          [C.cinit|&$id:("arg" <> show i)|],+          v',+          Nothing+        )+    mkArgs i (SharedMemoryKArg (Count c)) = do+      num_bytes <- GC.compileExp c+      size <- newVName "shared_size"+      offset <- newVName "shared_offset"+      GC.decl [C.cdecl|unsigned int $id:size = $exp:num_bytes;|]+      pure+        ( [C.cparam|unsigned int $id:("arg" <> show i)|],+          [C.cinit|&$id:("arg" <> show i)|],+          [C.cexp|$id:offset|],+          Just (size, offset)+        )++genKernelFunction ::+  KernelName ->+  KernelSafety ->+  Bool ->+  [C.Param] ->+  [C.Initializer] ->+  GC.CompilerM op s Name+genKernelFunction kernel_name safety need_perm arg_params arg_params_inits = do+  let kernel_fname = "gpu_kernel_" <> kernel_name+      (bef, aft) = profilingEnclosure kernel_name+      perm_args+        | need_perm = [[C.cinit|&perm[0]|], [C.cinit|&perm[1]|], [C.cinit|&perm[2]|]]         | otherwise = []       failure_args =         take@@ -297,92 +358,56 @@             [C.cinit|&ctx->failure_is_an_option|],             [C.cinit|&ctx->global_failure_args|]           ]-      args'' = perm_args ++ failure_args ++ [[C.cinit|&$id:a|] | a <- args']-      sizes_nonzero =-        expsNotZero-          [ grid_x,-            grid_y,-            grid_z,-            block_x,-            block_y,-            block_z-          ]-      (bef, aft) = profilingEnclosure kernel_name -  GC.stm-    [C.cstm|-    if ($exp:sizes_nonzero) {+  GC.libDecl+    [C.cedecl|static int $id:kernel_fname(+                struct futhark_context* ctx, unsigned int grid_x, unsigned int grid_y, unsigned int grid_z,+                unsigned int block_x, unsigned int block_y, unsigned int block_z,+                unsigned int shared_bytes, $params:arg_params) {+    if (grid_x * grid_y * grid_z * block_x * block_y * block_z != 0) {       int perm[3] = { 0, 1, 2 }; -      if ($exp:grid_y >= (1<<16)) {+      if (grid_y >= (1<<16)) {         perm[1] = perm[0];         perm[0] = 1;       } -      if ($exp:grid_z >= (1<<16)) {+      if (grid_z >= (1<<16)) {         perm[2] = perm[0];         perm[0] = 2;       }        size_t grid[3];-      grid[perm[0]] = $exp:grid_x;-      grid[perm[1]] = $exp:grid_y;-      grid[perm[2]] = $exp:grid_z;+      grid[perm[0]] = grid_x;+      grid[perm[1]] = grid_y;+      grid[perm[2]] = grid_z; -      void *$id:args_arr[] = { $inits:args'' };-      typename int64_t $id:time_start = 0, $id:time_end = 0;+      void *all_args[] = { $inits:(perm_args ++ failure_args ++ arg_params_inits) };+      typename int64_t time_start = 0, time_end = 0;       if (ctx->debugging) {-        fprintf(ctx->log, "Launching %s with grid size [%ld, %ld, %ld] and block size [%ld, %ld, %ld]; shared memory: %d bytes.\n",+        fprintf(ctx->log, "Launching %s with grid size [%d, %d, %d] and block size [%d, %d, %d]; shared memory: %d bytes.\n",                 $string:(prettyString kernel_name),-                (long int)$exp:grid_x, (long int)$exp:grid_y, (long int)$exp:grid_z,-                (long int)$exp:block_x, (long int)$exp:block_y, (long int)$exp:block_z,-                (int)$exp:shared_tot);-        $id:time_start = get_wall_time();+                grid_x, grid_y, grid_z,+                block_x, block_y, block_z,+                shared_bytes);+        time_start = get_wall_time();       }       $items:bef       CUDA_SUCCEED_OR_RETURN(         cuLaunchKernel(ctx->program->$id:kernel_name,                        grid[0], grid[1], grid[2],-                       $exp:block_x, $exp:block_y, $exp:block_z,-                       $exp:shared_tot, NULL,-                       $id:args_arr, NULL));+                       block_x, block_y, block_z,+                       shared_bytes, ctx->stream,+                       all_args, NULL));       $items:aft       if (ctx->debugging) {         CUDA_SUCCEED_FATAL(cuCtxSynchronize());-        $id:time_end = get_wall_time();+        time_end = get_wall_time();         fprintf(ctx->log, "Kernel %s runtime: %ldus\n",-                $string:(prettyString kernel_name), $id:time_end - $id:time_start);+                $string:(prettyString kernel_name), time_end - time_start);       }-    }|]+    }+    return FUTHARK_SUCCESS;+  }|] -  when (safety >= SafetyFull) $-    GC.stm [C.cstm|ctx->failure_is_an_option = 1;|]-  where-    mkDims [] = ([C.cexp|0|], [C.cexp|0|], [C.cexp|0|])-    mkDims [x] = (x, [C.cexp|1|], [C.cexp|1|])-    mkDims [x, y] = (x, y, [C.cexp|1|])-    mkDims (x : y : z : _) = (x, y, z)-    addExp x y = [C.cexp|$exp:x + $exp:y|]-    alignExp e = [C.cexp|$exp:e + ((8 - ($exp:e % 8)) % 8)|]-    mkOffsets = scanl (\a b -> a `addExp` alignExp b) [C.cexp|0|]-    expNotZero e = [C.cexp|$exp:e != 0|]-    expAnd a b = [C.cexp|$exp:a && $exp:b|]-    expsNotZero = foldl expAnd [C.cexp|1|] . map expNotZero-    mkArgs (ValueKArg e t@(FloatType Float16)) = do-      arg <- newVName "kernel_arg"-      e' <- GC.compileExp e-      GC.item [C.citem|$ty:(primStorageType t) $id:arg = $exp:(toStorage t e');|]-      pure (arg, Nothing)-    mkArgs (ValueKArg e t) =-      (,Nothing) <$> GC.compileExpToName "kernel_arg" t e-    mkArgs (MemKArg v) = do-      v' <- GC.rawMem v-      arg <- newVName "kernel_arg"-      GC.decl [C.cdecl|typename CUdeviceptr $id:arg = $exp:v';|]-      pure (arg, Nothing)-    mkArgs (SharedMemoryKArg (Count c)) = do-      num_bytes <- GC.compileExp c-      size <- newVName "shared_size"-      offset <- newVName "shared_offset"-      GC.decl [C.cdecl|unsigned int $id:size = $exp:num_bytes;|]-      pure (offset, Just (size, offset))+  pure kernel_fname

@@ -37,17 +37,17 @@ profilingEnclosure :: Name -> ([C.BlockItem], [C.BlockItem]) profilingEnclosure name =   ( [C.citems|-      typename cudaEvent_t *pevents = NULL;+      typename CUevent *pevents = NULL;       if (ctx->profiling && !ctx->profiling_paused) {         pevents = cuda_get_events(ctx,                                   &ctx->program->$id:(kernelRuns name),                                   &ctx->program->$id:(kernelRuntime name));-        CUDA_SUCCEED_FATAL(cudaEventRecord(pevents[0], 0));+        CUDA_SUCCEED_FATAL(cuEventRecord(pevents[0], ctx->stream));       }       |],     [C.citems|       if (pevents != NULL) {-        CUDA_SUCCEED_FATAL(cudaEventRecord(pevents[1], 0));+        CUDA_SUCCEED_FATAL(cuEventRecord(pevents[1], ctx->stream));       }       |]   )

@@ -11,7 +11,6 @@ where  import Control.Monad hiding (mapM)-import Data.List (intercalate) import Data.Text qualified as T import Futhark.CodeGen.Backends.COpenCL.Boilerplate import Futhark.CodeGen.Backends.GenericC qualified as GC@@ -320,111 +319,107 @@                                               &ctx->failure_is_an_option));     |] -  zipWithM_ setKernelArg [numFailureParams safety ..] args+  (arg_params, arg_set, call_args) <-+    unzip3 <$> zipWithM onArg [(0 :: Int) ..] args+   num_workgroups' <- mapM GC.compileExp num_workgroups   workgroup_size' <- mapM compileGroupDim workgroup_size   local_bytes <- foldM localBytes [C.cexp|0|] args -  launchKernel name num_workgroups' workgroup_size' local_bytes+  kernel_fname <- genKernelFunction name safety arg_params arg_set +  let grid_x : grid_y : grid_z : _ = num_workgroups' ++ repeat [C.cexp|1|]+      group_x : group_y : group_z : _ = workgroup_size' ++ repeat [C.cexp|1|]++  GC.stm+    [C.cstm|{+           err = $id:kernel_fname(ctx,+                                  $exp:grid_x,$exp:grid_y,$exp:grid_z,+                                  $exp:group_x, $exp:group_y, $exp:group_z,+                                  $exp:local_bytes,+                                  $args:call_args);+           if (err != FUTHARK_SUCCESS) { goto cleanup; }+           }|]+   when (safety >= SafetyFull) $     GC.stm [C.cstm|ctx->failure_is_an_option = 1;|]   where-    setKernelArg i (ValueKArg e pt) = do-      v <- case pt of-        -- We always transfer f16 values to the kernel as 16 bits, but-        -- the actual host type may be typedef'd to a 32-bit float.-        -- This requires some care.-        FloatType Float16 -> do-          v <- newVName "kernel_arg"-          e' <- toStorage pt <$> GC.compileExp e-          GC.decl [C.cdecl|$ty:(primStorageType pt) $id:v = $e';|]-          pure v-        _ -> GC.compileExpToName "kernel_arg" pt e-      GC.stm-        [C.cstm|-            OPENCL_SUCCEED_OR_RETURN(clSetKernelArg(ctx->program->$id:name, $int:i, sizeof($id:v), &$id:v));-          |]-    setKernelArg i (MemKArg v) = do-      v' <- GC.rawMem v-      GC.stm-        [C.cstm|-            OPENCL_SUCCEED_OR_RETURN(clSetKernelArg(ctx->program->$id:name, $int:i, sizeof($exp:v'), &$exp:v'));-          |]-    setKernelArg i (SharedMemoryKArg num_bytes) = do-      num_bytes' <- GC.compileExp $ unCount num_bytes-      GC.stm-        [C.cstm|-            OPENCL_SUCCEED_OR_RETURN(clSetKernelArg(ctx->program->$id:name, $int:i, (size_t)$exp:num_bytes', NULL));-            |]-     localBytes cur (SharedMemoryKArg num_bytes) = do       num_bytes' <- GC.compileExp $ unCount num_bytes       pure [C.cexp|$exp:cur + $exp:num_bytes'|]     localBytes cur _ = pure cur -launchKernel ::-  C.ToExp a =>-  KernelName ->-  [a] ->-  [a] ->-  a ->-  GC.CompilerM op s ()-launchKernel kernel_name num_workgroups workgroup_dims local_bytes = do-  global_work_size <- newVName "global_work_size"-  time_start <- newVName "time_start"-  time_end <- newVName "time_end"-  time_diff <- newVName "time_diff"-  local_work_size <- newVName "local_work_size"--  let (debug_str, debug_args) = debugPrint global_work_size local_work_size+    onArg i (ValueKArg e t) = do+      let arg = "arg" <> show i+      e' <- GC.compileExp e+      pure+        ( [C.cparam|$ty:(primStorageType t) $id:arg|],+          ([C.cexp|sizeof($id:arg)|], [C.cexp|&$id:arg|]),+          toStorage t e'+        )+    onArg i (MemKArg v) = do+      let arg = "arg" <> show i+      v' <- GC.rawMem v+      pure+        ( [C.cparam|typename cl_mem $id:arg|],+          ([C.cexp|sizeof($id:arg)|], [C.cexp|&$id:arg|]),+          v'+        )+    onArg i (SharedMemoryKArg (Count c)) = do+      let arg = "arg" <> show i+      num_bytes <- GC.compileExp c+      pure+        ( [C.cparam|unsigned int $id:arg|],+          ([C.cexp|$id:arg|], [C.cexp|NULL|]),+          num_bytes+        ) -  GC.stm-    [C.cstm|-    if ($exp:total_elements != 0) {-      const size_t $id:global_work_size[$int:kernel_rank] = {$inits:kernel_dims'};-      const size_t $id:local_work_size[$int:kernel_rank] = {$inits:workgroup_dims'};-      typename int64_t $id:time_start = 0, $id:time_end = 0;+genKernelFunction ::+  KernelName ->+  KernelSafety ->+  [C.Param] ->+  [(C.Exp, C.Exp)] ->+  GC.CompilerM op s Name+genKernelFunction kernel_name safety arg_params arg_set = do+  let kernel_fname = "gpu_kernel_" <> kernel_name+  GC.libDecl+    [C.cedecl|static int $id:kernel_fname(+                struct futhark_context* ctx,+                unsigned int grid_x, unsigned int grid_y, unsigned int grid_z,+                unsigned int block_x, unsigned int block_y, unsigned int block_z,+                unsigned int local_bytes, $params:arg_params) {+    (void)local_bytes;+    if (grid_x * grid_y * grid_z * block_x * block_y * block_z != 0) {+      const size_t global_work_size[3] = {grid_x*block_x, grid_y*block_y, grid_z*block_z};+      const size_t local_work_size[3] = {block_x, block_y, block_z};+      typename int64_t time_start = 0, time_end = 0;+      $stms:set_args       if (ctx->debugging) {-        fprintf(ctx->log, $string:debug_str, $args:debug_args);-        $id:time_start = get_wall_time();+        fprintf(ctx->log, "Launching %s with grid size [%d, %d, %d] and group size [%d, %d, %d]; local memory: %d bytes.\n",+                $string:(prettyString kernel_name),+                grid_x, grid_y, grid_z,+                block_x, block_y, block_z,+                local_bytes);+        time_start = get_wall_time();       }       typename cl_event *pevent = $exp:(profilingEvent kernel_name);       OPENCL_SUCCEED_OR_RETURN(-        clEnqueueNDRangeKernel(ctx->queue, ctx->program->$id:kernel_name, $int:kernel_rank, NULL,-                               $id:global_work_size, $id:local_work_size,+        clEnqueueNDRangeKernel(ctx->queue, ctx->program->$id:kernel_name, 3, NULL,+                               global_work_size, local_work_size,                                0, NULL, pevent));       if (ctx->debugging) {         OPENCL_SUCCEED_FATAL(clFinish(ctx->queue));-        $id:time_end = get_wall_time();-        long int $id:time_diff = $id:time_end - $id:time_start;+        time_end = get_wall_time();+        long int time_diff = time_end - time_start;         fprintf(ctx->log, "kernel %s runtime: %ldus\n",-                $string:(prettyString kernel_name), $id:time_diff);+                $string:(prettyString kernel_name), time_diff);       }-    }|]-  where-    kernel_rank = length kernel_dims-    kernel_dims = zipWith multExp (map toSize num_workgroups) (map toSize workgroup_dims)-    kernel_dims' = map toInit kernel_dims-    workgroup_dims' = map (toInit . toSize) workgroup_dims-    total_elements = foldl multExp [C.cexp|1|] kernel_dims--    toInit e = [C.cinit|$exp:e|]-    multExp x y = [C.cexp|$exp:x * $exp:y|]-    toSize e = [C.cexp|(size_t)$exp:e|]+    }+    return FUTHARK_SUCCESS;+  }|] -    debugPrint :: VName -> VName -> (String, [C.Exp])-    debugPrint global_work_size local_work_size =-      ( "Launching %s with global work size "-          ++ dims-          ++ " and local work size "-          ++ dims-          ++ "; local memory: %d bytes.\n",-        [C.cexp|$string:(prettyString kernel_name)|]-          : map (kernelDim global_work_size) [0 .. kernel_rank - 1]-          ++ map (kernelDim local_work_size) [0 .. kernel_rank - 1]-          ++ [[C.cexp|(int)$exp:local_bytes|]]-      )-      where-        dims = "[" ++ intercalate ", " (replicate kernel_rank "%zu") ++ "]"-        kernelDim arr i = [C.cexp|$id:arr[$int:i]|]+  pure kernel_fname+  where+    set_args = zipWith setKernelArg [numFailureParams safety ..] arg_set+    setKernelArg i (size, e) =+      [C.cstm|OPENCL_SUCCEED_OR_RETURN(clSetKernelArg(ctx->program->$id:kernel_name, $int:i, $exp:size, $exp:e));|]

@@ -209,11 +209,11 @@ instance C.ToExp FloatValue where   toExp (Float16Value x) _     | isInfinite x =-        if x > 0 then [C.cexp|INFINITY|] else [C.cexp|-INFINITY|]+        if x > 0 then [C.cexp|(typename f16)INFINITY|] else [C.cexp|(typename f16)-INFINITY|]     | isNaN x =-        [C.cexp|NAN|]+        [C.cexp|(typename f16)NAN|]     | otherwise =-        [C.cexp|$float:(fromRational (toRational x))|]+        [C.cexp|(typename f16)$float:(fromRational (toRational x))|]   toExp (Float32Value x) _     | isInfinite x =         if x > 0 then [C.cexp|INFINITY|] else [C.cexp|-INFINITY|]

@@ -602,8 +602,11 @@         If {} -> pretty fbranch         _ ->           "{" </> indent 2 (pretty fbranch) </> "}"+  pretty (Call [] fname args) =+    "call" <+> pretty fname <> parens (commasep $ map pretty args)   pretty (Call dests fname args) =-    commasep (map pretty dests)+    "call"+      <+> commasep (map pretty dests)       <+> "<-"       <+> pretty fname <> parens (commasep $ map pretty args)   pretty (Comment s code) =

@@ -146,6 +146,7 @@ import Futhark.Construct hiding (ToExp (..)) import Futhark.IR.Mem import Futhark.IR.Mem.IxFun qualified as IxFun+import Futhark.IR.Mem.LMAD qualified as LMAD import Futhark.IR.SOACS (SOACS) import Futhark.Util import Futhark.Util.IntegralExp@@ -207,9 +208,9 @@ sliceMemLoc (MemLoc mem shape ixfun) slice =   MemLoc mem shape $ IxFun.slice ixfun slice -flatSliceMemLoc :: MemLoc -> FlatSlice (Imp.TExp Int64) -> MemLoc+flatSliceMemLoc :: MemLoc -> FlatSlice (Imp.TExp Int64) -> Maybe MemLoc flatSliceMemLoc (MemLoc mem shape ixfun) slice =-  MemLoc mem shape $ IxFun.flatSlice ixfun slice+  MemLoc mem shape <$> IxFun.flatSlice ixfun slice  data ArrayEntry = ArrayEntry   { entryArrayLoc :: MemLoc,@@ -931,7 +932,9 @@ defCompileBasicOp (Pat [pe]) (FlatUpdate _ slice v) = do   pe_loc <- entryArrayLoc <$> lookupArray (patElemName pe)   v_loc <- entryArrayLoc <$> lookupArray v-  copy (elemType (patElemType pe)) (flatSliceMemLoc pe_loc slice') v_loc+  case flatSliceMemLoc pe_loc slice' of+    Just pe_loc' -> copy (elemType (patElemType pe)) pe_loc' v_loc+    Nothing -> error "defCompileBasicOp FlatUpdate"   where     slice' = fmap pe64 slice defCompileBasicOp (Pat [pe]) (Replicate shape se)@@ -1401,8 +1404,8 @@ copy :: CopyCompiler rep r op copy   bt-  dst@(MemLoc dst_name _ dst_ixfn@(IxFun.IxFun dst_lmads@(dst_lmad :| _) _ _))-  src@(MemLoc src_name _ src_ixfn@(IxFun.IxFun src_lmads@(src_lmad :| _) _ _)) = do+  dst@(MemLoc dst_name _ dst_ixfn@(IxFun.IxFun dst_lmad _))+  src@(MemLoc src_name _ src_ixfn@(IxFun.IxFun src_lmad _)) = do     -- If we can statically determine that the two index-functions     -- are equivalent, don't do anything     unless (dst_name == src_name && dst_ixfn `IxFun.equivalent` src_ixfn)@@ -1410,7 +1413,7 @@       -- It's also possible that we can dynamically determine that the two       -- index-functions are equivalent.       sUnless-        ( fromBool (dst_name == src_name && length dst_lmads == 1 && length src_lmads == 1)+        ( fromBool (dst_name == src_name)             .&&. IxFun.dynamicEqualsLMAD dst_lmad src_lmad         )       $ do@@ -1418,44 +1421,54 @@         cc <- asks envCopyCompiler         cc bt dst src --- | Is this copy really a mapping with transpose?+-- | Is this copy really a mapping with transpose?  Produce an+-- expression that is true if so, as well as other expressions that+-- contain information about the transpose in that case (don't trust+-- these if the boolean is false). isMapTransposeCopy ::   PrimType ->   MemLoc ->   MemLoc ->   Maybe-    ( Imp.TExp Int64,-      Imp.TExp Int64,-      Imp.TExp Int64,-      Imp.TExp Int64,-      Imp.TExp Int64+    ( Imp.TExp Bool,+      ( Imp.TExp Int64,+        Imp.TExp Int64,+        Imp.TExp Int64,+        Imp.TExp Int64,+        Imp.TExp Int64+      )     )-isMapTransposeCopy bt (MemLoc _ _ destIxFun) (MemLoc _ _ srcIxFun)-  | Just (dest_offset, perm_and_destshape) <- IxFun.rearrangeWithOffset destIxFun bt_size,-    (perm, destshape) <- unzip perm_and_destshape,-    Just src_offset <- IxFun.linearWithOffset srcIxFun bt_size,+isMapTransposeCopy pt (MemLoc _ _ destIxFun) (MemLoc _ _ srcIxFun)+  | perm <- LMAD.permutation dest_lmad,+    LMAD.permutation src_lmad == [0 .. rank - 1],     Just (r1, r2, _) <- isMapTranspose perm =-      isOk destshape swap r1 r2 dest_offset src_offset-  | Just dest_offset <- IxFun.linearWithOffset destIxFun bt_size,-    Just (src_offset, perm_and_srcshape) <- IxFun.rearrangeWithOffset srcIxFun bt_size,-    (perm, srcshape) <- unzip perm_and_srcshape,+      isOk (IxFun.shape destIxFun) swap r1 r2+  | perm <- LMAD.permutation src_lmad,+    LMAD.permutation dest_lmad == [0 .. rank - 1],     Just (r1, r2, _) <- isMapTranspose perm =-      isOk srcshape id r1 r2 dest_offset src_offset+      isOk (IxFun.shape srcIxFun) id r1 r2   | otherwise =       Nothing   where-    bt_size = primByteSize bt+    rank = IxFun.rank destIxFun     swap (x, y) = (y, x)+    dest_lmad = IxFun.ixfunLMAD destIxFun+    src_lmad = IxFun.ixfunLMAD srcIxFun+    dest_offset = LMAD.offset dest_lmad+    src_offset = LMAD.offset src_lmad -    isOk shape f r1 r2 dest_offset src_offset = do+    isOk shape f r1 r2 =       let (num_arrays, size_x, size_y) = getSizes shape f r1 r2-      pure-        ( dest_offset,-          src_offset,-          num_arrays,-          size_x,-          size_y-        )+       in Just+            ( LMAD.contiguous dest_lmad+                .&&. LMAD.contiguous src_lmad,+              ( dest_offset * primByteSize pt,+                src_offset * primByteSize pt,+                num_arrays,+                size_x,+                size_y+              )+            )      getSizes shape f r1 r2 =       let (mapped, notmapped) = splitAt r1 shape@@ -1477,35 +1490,25 @@ -- | Use 'sCopy' if possible, otherwise 'copyElementWise'. defaultCopy :: CopyCompiler rep r op defaultCopy pt dest src-  | Just (destoffset, srcoffset, num_arrays, size_x, size_y) <-+  | Just (is_transpose, (destoffset, srcoffset, num_arrays, size_x, size_y)) <-       isMapTransposeCopy pt dest src = do       fname <- mapTransposeForType pt-      emit-        $ Imp.Call-          []-          fname-        $ transposeArgs-          pt-          destmem-          (bytes destoffset)-          srcmem-          (bytes srcoffset)-          num_arrays-          size_x-          size_y-  | Just destoffset <--      IxFun.linearWithOffset dest_ixfun pt_size,-    Just srcoffset <--      IxFun.linearWithOffset src_ixfun pt_size = do-      srcspace <- entryMemSpace <$> lookupMemory srcmem-      destspace <- entryMemSpace <$> lookupMemory destmem-      if isScalarSpace srcspace || isScalarSpace destspace-        then copyElementWise pt dest src-        else sCopy destmem destoffset destspace srcmem srcoffset srcspace num_elems pt-  | otherwise =-      copyElementWise pt dest src+      sIf+        is_transpose+        ( emit . Imp.Call [] fname $+            transposeArgs+              pt+              destmem+              (bytes destoffset)+              srcmem+              (bytes srcoffset)+              num_arrays+              size_x+              size_y+        )+        nontranspose+  | otherwise = nontranspose   where-    pt_size = primByteSize pt     num_elems = Imp.elements $ product $ IxFun.shape $ memLocIxFun src      MemLoc destmem _ dest_ixfun = dest@@ -1514,6 +1517,21 @@     isScalarSpace ScalarSpace {} = True     isScalarSpace _ = False +    nontranspose = do+      srcspace <- entryMemSpace <$> lookupMemory srcmem+      destspace <- entryMemSpace <$> lookupMemory destmem+      if isScalarSpace srcspace || isScalarSpace destspace+        then copyElementWise pt dest src+        else do+          let dest_lmad = LMAD.noPermutation $ IxFun.ixfunLMAD dest_ixfun+              src_lmad = LMAD.noPermutation $ IxFun.ixfunLMAD src_ixfun+              destoffset = elements (LMAD.offset dest_lmad) `withElemType` pt+              srcoffset = elements (LMAD.offset src_lmad) `withElemType` pt+          sIf+            (LMAD.memcpyable dest_lmad src_lmad)+            (sCopy destmem destoffset destspace srcmem srcoffset srcspace num_elems pt)+            (copyElementWise pt dest src)+ copyElementWise :: CopyCompiler rep r op copyElementWise bt dest src = do   let bounds = IxFun.shape $ memLocIxFun src@@ -1889,17 +1907,17 @@  sCopy ::   VName ->-  Imp.TExp Int64 ->+  Count Bytes (Imp.TExp Int64) ->   Space ->   VName ->-  Imp.TExp Int64 ->+  Count Bytes (Imp.TExp Int64) ->   Space ->   Count Elements (Imp.TExp Int64) ->   PrimType ->   ImpM rep r op () sCopy destmem destoffset destspace srcmem srcoffset srcspace num_elems pt =   if destmem == srcmem-    then sUnless (destoffset .==. srcoffset) the_copy+    then sUnless (Imp.unCount destoffset .==. Imp.unCount srcoffset) the_copy     else the_copy   where     the_copy =@@ -1907,10 +1925,10 @@         $ Imp.Copy           pt           destmem-          (bytes destoffset)+          destoffset           destspace           srcmem-          (bytes srcoffset)+          srcoffset           srcspace         $ num_elems `withElemType` pt 

@@ -12,6 +12,8 @@ where  import Control.Monad+import Control.Monad.State+import Data.Foldable (toList) import Data.List (foldl') import Data.Map qualified as M import Data.Maybe@@ -19,6 +21,7 @@ import Futhark.CodeGen.ImpGen hiding (compileProg) import Futhark.CodeGen.ImpGen qualified import Futhark.CodeGen.ImpGen.GPU.Base+import Futhark.CodeGen.ImpGen.GPU.Copy import Futhark.CodeGen.ImpGen.GPU.SegHist import Futhark.CodeGen.ImpGen.GPU.SegMap import Futhark.CodeGen.ImpGen.GPU.SegRed@@ -27,6 +30,7 @@ import Futhark.Error import Futhark.IR.GPUMem import Futhark.IR.Mem.IxFun qualified as IxFun+import Futhark.IR.Mem.LMAD qualified as LMAD import Futhark.MonadFreshNames import Futhark.Util.IntegralExp (IntegralExp, divUp, quot, rem) import Prelude hiding (quot, rem)@@ -231,10 +235,12 @@   s' <- toExp s    sIota (patElemName pe) (pe64 n) x' s' et-expCompiler (Pat [pe]) (BasicOp (Replicate _ se))+expCompiler (Pat [pe]) (BasicOp (Replicate shape se))   | Acc {} <- patElemType pe = pure ()   | otherwise =-      sReplicate (patElemName pe) se+      if shapeRank shape == 0+        then copyDWIM (patElemName pe) [] se []+        else sReplicate (patElemName pe) se -- Allocation in the "local" space is just a placeholder. expCompiler _ (Op (Alloc _ (Space "local"))) =   pure ()@@ -257,40 +263,92 @@ expCompiler dest e =   defCompileExp dest e -callKernelCopy :: CopyCompiler GPUMem HostEnv Imp.HostOp-callKernelCopy bt destloc@(MemLoc destmem _ destIxFun) srcloc@(MemLoc srcmem srcshape srcIxFun)-  | Just (destoffset, srcoffset, num_arrays, size_x, size_y) <--      isMapTransposeCopy bt destloc srcloc = do-      fname <- mapTransposeForType bt-      emit $-        Imp.Call-          []-          fname-          [ Imp.MemArg destmem,-            Imp.ExpArg $ untyped destoffset,-            Imp.MemArg srcmem,-            Imp.ExpArg $ untyped srcoffset,-            Imp.ExpArg $ untyped num_arrays,-            Imp.ExpArg $ untyped size_x,-            Imp.ExpArg $ untyped size_y-          ]-  | bt_size <- primByteSize bt,-    Just destoffset <- IxFun.linearWithOffset destIxFun bt_size,-    Just srcoffset <- IxFun.linearWithOffset srcIxFun bt_size = do-      let num_elems = Imp.elements $ product $ map pe64 srcshape-      srcspace <- entryMemSpace <$> lookupMemory srcmem-      destspace <- entryMemSpace <$> lookupMemory destmem-      sCopy-        destmem-        (sExt64 destoffset)-        destspace-        srcmem-        (sExt64 srcoffset)-        srcspace-        num_elems-        bt-  | otherwise = sCopyKernel bt destloc srcloc+gpuCopyForType :: Rank -> PrimType -> CallKernelGen Name+gpuCopyForType r bt = do+  let fname = nameFromString $ "builtin#" <> gpuCopyName r bt +  exists <- hasFunction fname+  unless exists $ emitFunction fname $ gpuCopyFunction r bt++  pure fname++gpuCopyName :: Rank -> PrimType -> String+gpuCopyName (Rank r) bt = "gpu_copy_" <> show r <> "d_" <> prettyString bt++gpuCopyFunction :: Rank -> PrimType -> Imp.Function Imp.HostOp+gpuCopyFunction (Rank r) pt = do+  let tdesc = mconcat (replicate r "[]") <> prettyString pt+  Imp.Function Nothing [] params $+    Imp.DebugPrint ("\n# Copy " <> tdesc) Nothing+      <> copy_code+      <> Imp.DebugPrint "" Nothing+  where+    space = Space "device"+    memparam v = Imp.MemParam v space+    intparam v = Imp.ScalarParam v $ IntType Int64++    mkIxFun desc = do+      let new x = newVName $ desc <> "_" <> x+          newDim i = LMAD.LMADDim <$> new "stride" <*> new "shape" <*> pure i+      LMAD.LMAD <$> new "offset" <*> mapM newDim [0 .. r - 1]++    (params, copy_code) = do+      flip evalState blankNameSource $ do+        dest_mem <- newVName "destmem"+        dest_lmad <- mkIxFun "dest"++        src_mem <- newVName "srcmem"+        src_lmad <- mkIxFun "src"++        group_size <- newVName "group_size"+        num_groups <- newVName "num_groups"++        let kernel =+              copyKernel+                pt+                (le64 num_groups, Left $ untyped $ le64 group_size)+                (dest_mem, le64 <$> dest_lmad)+                (src_mem, le64 <$> src_lmad)++            dest_offset =+              Imp.elements (le64 (LMAD.offset dest_lmad)) `Imp.withElemType` pt++            src_offset =+              Imp.elements (le64 (LMAD.offset src_lmad)) `Imp.withElemType` pt++            num_bytes =+              Imp.elements (product (le64 <$> LMAD.shape src_lmad)) `Imp.withElemType` pt++            do_copy =+              Imp.Copy+                pt+                dest_mem+                dest_offset+                (Space "device")+                src_mem+                src_offset+                (Space "device")+                num_bytes++        pure+          ( [memparam dest_mem]+              ++ map intparam (toList dest_lmad)+              ++ [memparam src_mem]+              ++ map intparam (toList src_lmad),+            Imp.DeclareScalar group_size Imp.Nonvolatile int64+              <> Imp.DeclareScalar num_groups Imp.Nonvolatile int64+              <> Imp.Op (Imp.GetSize group_size "copy_group_size" Imp.SizeGroup)+              <> Imp.Op (Imp.GetSize num_groups "copy_num_groups" Imp.SizeNumGroups)+              <> Imp.If+                (LMAD.memcpyable (le64 <$> dest_lmad) (le64 <$> src_lmad))+                ( Imp.DebugPrint "## Simple copy" Nothing+                    <> do_copy+                )+                ( Imp.DebugPrint "## Kernel copy" Nothing+                    <> Imp.Op (Imp.CallKernel kernel)+                )+          )+ mapTransposeForType :: PrimType -> CallKernelGen Name mapTransposeForType bt = do   let fname = nameFromString $ "builtin#" <> mapTransposeName bt@@ -480,3 +538,51 @@ -- unfortunate code bloat, and it would be preferable if we could -- simply optimise the 64-bit version to make this distinction -- unnecessary.  Fortunately these kernels are quite small.++callKernelCopy :: CopyCompiler GPUMem HostEnv Imp.HostOp+callKernelCopy pt destloc@(MemLoc destmem _ dest_ixfun) srcloc@(MemLoc srcmem _ src_ixfun)+  | Just (is_transpose, (destoffset, srcoffset, num_arrays, size_x, size_y)) <-+      isMapTransposeCopy pt destloc srcloc = do+      fname <- mapTransposeForType pt+      sIf+        is_transpose+        ( emit . Imp.Call [] fname $+            [ Imp.MemArg destmem,+              Imp.ExpArg $ untyped destoffset,+              Imp.MemArg srcmem,+              Imp.ExpArg $ untyped srcoffset,+              Imp.ExpArg $ untyped num_arrays,+              Imp.ExpArg $ untyped size_x,+              Imp.ExpArg $ untyped size_y+            ]+        )+        nontranspose+  | otherwise = nontranspose+  where+    nontranspose = do+      fname <- gpuCopyForType (Rank (IxFun.rank dest_ixfun)) pt+      dest_space <- entryMemSpace <$> lookupMemory destmem+      src_space <- entryMemSpace <$> lookupMemory srcmem+      let dest_lmad = LMAD.noPermutation $ IxFun.ixfunLMAD dest_ixfun+          src_lmad = LMAD.noPermutation $ IxFun.ixfunLMAD src_ixfun+          num_elems = Imp.elements $ product $ LMAD.shape dest_lmad+      if dest_space == Space "device" && src_space == Space "device"+        then+          emit . Imp.Call [] fname $+            [Imp.MemArg destmem]+              ++ map (Imp.ExpArg . untyped) (toList dest_lmad)+              ++ [Imp.MemArg srcmem]+              ++ map (Imp.ExpArg . untyped) (toList src_lmad)+        else -- FIXME: this assumes a linear representation!+        -- Currently we never generate code where this is not the+        -- case, but we might in the future.++          sCopy+            destmem+            (Imp.elements (LMAD.offset dest_lmad) `Imp.withElemType` pt)+            dest_space+            srcmem+            (Imp.elements (LMAD.offset src_lmad) `Imp.withElemType` pt)+            src_space+            num_elems+            pt

@@ -19,7 +19,6 @@     sKernelThread,     KernelAttrs (..),     defKernelAttrs,-    sCopyKernel,     lvlKernelAttrs,     allocLocal,     kernelAlloc,@@ -1258,7 +1257,7 @@   v_t <- subExpType v   case v_t of     Prim v_t'-      | IxFun.isLinear arr_ixfun -> pure $+      | IxFun.isDirect arr_ixfun -> pure $           Just $ do             fname <- replicateForType v_t'             emit $@@ -1360,7 +1359,7 @@   CallKernelGen () sIota arr n x s et = do   ArrayEntry (MemLoc arr_mem _ arr_ixfun) _ <- lookupArray arr-  if IxFun.isLinear arr_ixfun+  if IxFun.isDirect arr_ixfun     then do       fname <- iotaForType et       emit $@@ -1369,33 +1368,6 @@           fname           [Imp.MemArg arr_mem, Imp.ExpArg $ untyped n, Imp.ExpArg x, Imp.ExpArg s]     else sIotaKernel arr n x s et--sCopyKernel :: CopyCompiler GPUMem HostEnv Imp.HostOp-sCopyKernel pt destloc@(MemLoc destmem _ _) srcloc@(MemLoc srcmem srcdims _) = do-  -- Note that the shape of the destination and the source are-  -- necessarily the same.-  let shape = map pe64 srcdims-      kernel_size = product shape--  (virtualise, constants) <- simpleKernelConstants kernel_size "copy"--  fname <- askFunction-  let name =-        keyWithEntryPoint fname $-          nameFromString $-            "copy_" ++ show (baseTag $ kernelGlobalThreadIdVar constants)--  sKernelFailureTolerant True threadOperations constants name $-    virtualise $ \gtid -> do-      is <- dIndexSpace' "copy_i" shape gtid--      (_, destspace, destidx) <- fullyIndexArray' destloc is-      (_, srcspace, srcidx) <- fullyIndexArray' srcloc is--      sWhen (gtid .<. kernel_size) $ do-        tmp <- tvVar <$> dPrim "tmp" pt-        emit $ Imp.Read tmp srcmem srcidx pt srcspace Imp.Nonvolatile-        emit $ Imp.Write destmem destidx pt destspace Imp.Nonvolatile $ Imp.var tmp pt  compileThreadResult ::   SegSpace ->

@@ -0,0 +1,81 @@+-- | General implementation of GPU copying, using LMAD representation.+-- That means the dynamic performance of this kernel depends crucially+-- on the LMAD.  In most cases we should use a more specialised kernel.+-- Written in ImpCode so we can compile it to both CUDA and OpenCL.+module Futhark.CodeGen.ImpGen.GPU.Copy (copyKernel) where++import Control.Monad+import Control.Monad.State+import Data.Foldable (toList)+import Futhark.CodeGen.ImpCode.GPU+import Futhark.IR.Mem.LMAD qualified as LMAD+import Futhark.IR.Prop.Reshape+import Futhark.MonadFreshNames+import Futhark.Util (nubOrd)+import Futhark.Util.IntegralExp (divUp)+import Prelude hiding (quot, rem)++copyKernel ::+  PrimType ->+  (TExp Int64, GroupDim) ->+  (VName, LMAD.LMAD (TExp Int64)) ->+  (VName, LMAD.LMAD (TExp Int64)) ->+  Kernel+copyKernel pt (num_groups, group_dim) (dest_mem, dest_lmad) (src_mem, src_lmad) =+  Kernel+    { kernelBody = body,+      kernelUses =+        let frees =+              nubOrd+                ( foldMap toList dest_lmad+                    <> foldMap toList src_lmad+                    <> toList num_groups+                )+         in map (`ScalarUse` IntType Int64) frees+              ++ map MemoryUse [dest_mem, src_mem],+      kernelNumGroups = [untyped num_groups],+      kernelGroupSize = [group_dim],+      kernelName = nameFromString ("copy_" <> show rank <> "d_" <> prettyString pt),+      kernelFailureTolerant = True,+      kernelCheckLocalMemory = False+    }+  where+    shape = LMAD.shape dest_lmad+    rank = length shape++    body = flip evalState (newNameSource 1000) $ do+      group_id <- newVName "group_id"+      local_id <- newVName "local_id"+      local_size <- newVName "local_size"+      global_id <- newVName "global_id"+      group_iter <- newVName "group_iter"+      let global_id_e =+            ((le64 group_id + le64 group_iter * num_groups) * le64 local_size)+              + le64 local_id++      is <- replicateM rank $ newVName "i"+      let is_e = map untyped (unflattenIndex shape (le64 global_id))+          in_bounds = foldl1 (.&&.) (zipWith (.<.) (map le64 is) shape)++      element <- newVName "element"++      let dec v = DeclareScalar v Nonvolatile $ IntType Int64+          src_o = Count (LMAD.index src_lmad (map le64 is))+          dest_o = Count (LMAD.index dest_lmad (map le64 is))+          copy_elem =+            DeclareScalar element Nonvolatile pt+              <> Read element src_mem src_o pt (Space "device") Nonvolatile+              <> Write dest_mem dest_o pt (Space "device") Nonvolatile (LeafExp element pt)+      pure $+        foldMap dec [group_id, local_id, local_size, global_id]+          <> Op (GetLocalId local_id 0)+          <> Op (GetLocalSize local_size 0)+          <> Op (GetGroupId group_id 0)+          <> For+            group_iter+            (untyped (product shape `divUp` (le64 local_size * num_groups)))+            ( SetScalar global_id (untyped global_id_e)+                <> foldMap dec is+                <> mconcat (zipWith SetScalar is is_e)+                <> If in_bounds copy_elem mempty+            )

@@ -41,9 +41,10 @@   ArrayEntry arr_loc pt <- lookupArray arr   let flat_shape = Shape $ Var (tvVar flat) : drop k (memLocShape arr_loc)   sArray (baseString arr ++ "_flat") pt flat_shape (memLocName arr_loc) $-    IxFun.reshape (memLocIxFun arr_loc) $-      map pe64 $-        shapeDims flat_shape+    fromMaybe (error "flattenArray") $+      IxFun.reshape (memLocIxFun arr_loc) $+        map pe64 $+          shapeDims flat_shape  sliceArray :: Imp.TExp Int64 -> TV Int64 -> VName -> ImpM rep r op VName sliceArray start size arr = do

@@ -584,7 +584,7 @@         MemMem {} -> pure ()         MemAcc {} -> pure ()         MemArray _ _ _ (ArrayIn _ ixfun)-          | IxFun.isLinear ixfun ->+          | IxFun.isDirect ixfun ->               pure ()           | otherwise ->               TC.bad . TC.TypeError $@@ -766,7 +766,7 @@   TC.TypeM rep () matchPatToExp pat e = do   scope <- asksScope removeScopeAliases-  rt <- runReaderT (expReturns $ removeExpAliases e) scope+  rt <- maybe illformed pure $ runReader (expReturns $ removeExpAliases e) scope    let (ctx_ids, val_ts) = unzip $ bodyReturnsFromPat $ removePatAliases pat       (ctx_map_ids, ctx_map_exts) = getExtMaps $ zip ctx_ids [0 .. 1]@@ -780,6 +780,13 @@       </> "cannot match pattern type:"       </> indent 2 (ppTupleLines' $ map pretty val_ts)   where+    illformed =+      TC.bad $+        TC.TypeError . docText $+          "Expression"+            </> indent 2 (pretty e)+            </> "cannot be assigned an index function."+     matches _ _ (MemPrim x) (MemPrim y) = x == y     matches _ _ (MemMem x_space) (MemMem y_space) =       x_space == y_space@@ -1006,49 +1013,55 @@  -- | The return information of an expression.  This can be seen as the -- "return type with memory annotations" of the expression.+--+-- This can produce Nothing, which signifies that the result is an+-- array layout that is not expressible as an index function. expReturns ::   (LocalScope rep m, Mem rep inner) =>   Exp rep ->-  m [ExpReturns]+  m (Maybe [ExpReturns]) expReturns (BasicOp (SubExp se)) =-  pure <$> subExpReturns se+  Just . pure <$> subExpReturns se expReturns (BasicOp (Opaque _ (Var v))) =-  pure <$> varReturns v+  Just . pure <$> varReturns v expReturns (BasicOp (Reshape k newshape v)) = do   (et, _, mem, ixfun) <- arrayVarReturns v-  pure-    [ MemArray et (fmap Free newshape) NoUniqueness $-        Just . ReturnsInBlock mem . existentialiseIxFun [] $-          reshaper ixfun $-            map pe64 (shapeDims newshape)-    ]+  case reshaper k ixfun $ map pe64 $ shapeDims newshape of+    Just ixfun' ->+      pure . Just $+        [ MemArray et (fmap Free newshape) NoUniqueness . Just $+            ReturnsInBlock mem (existentialiseIxFun [] ixfun')+        ]+    Nothing -> pure Nothing   where-    reshaper = case k of-      ReshapeArbitrary -> IxFun.reshape-      ReshapeCoerce -> IxFun.coerce+    reshaper ReshapeArbitrary ixfun =+      IxFun.reshape ixfun+    reshaper ReshapeCoerce ixfun =+      Just . IxFun.coerce ixfun expReturns (BasicOp (Rearrange perm v)) = do   (et, Shape dims, mem, ixfun) <- arrayVarReturns v   let ixfun' = IxFun.permute ixfun perm       dims' = rearrangeShape perm dims-  pure-    [ MemArray et (Shape $ map Free dims') NoUniqueness $-        Just $-          ReturnsInBlock mem $-            existentialiseIxFun [] ixfun'-    ]+  pure $+    Just+      [ MemArray et (Shape $ map Free dims') NoUniqueness $+          Just $+            ReturnsInBlock mem $+              existentialiseIxFun [] ixfun'+      ] expReturns (BasicOp (Index v slice)) = do-  pure . varInfoToExpReturns <$> sliceInfo v slice+  Just . pure . varInfoToExpReturns <$> sliceInfo v slice expReturns (BasicOp (Update _ v _ _)) =-  pure <$> varReturns v+  Just . pure <$> varReturns v expReturns (BasicOp (FlatIndex v slice)) = do-  pure . varInfoToExpReturns <$> flatSliceInfo v slice+  fmap (pure . varInfoToExpReturns) <$> flatSliceInfo v slice expReturns (BasicOp (FlatUpdate v _ _)) =-  pure <$> varReturns v+  Just . pure <$> varReturns v expReturns (BasicOp op) =-  extReturns . staticShapes <$> basicOpType op+  Just . extReturns . staticShapes <$> basicOpType op expReturns e@(DoLoop merge _ _) = do   t <- expExtType e-  zipWithM typeWithDec t $ map fst merge+  Just <$> zipWithM typeWithDec t (map fst merge)   where     typeWithDec t p =       case (t, paramDec p) of@@ -1073,18 +1086,20 @@     isMergeVar v = find ((== v) . paramName . snd) $ zip [0 ..] mergevars     mergevars = map fst merge expReturns (Apply _ _ ret _) =-  pure $ map (funReturnsToExpReturns . fst) ret+  pure $ Just $ map (funReturnsToExpReturns . fst) ret expReturns (Match _ _ _ (MatchDec ret _)) =-  pure $ map bodyReturnsToExpReturns ret+  pure $ Just $ map bodyReturnsToExpReturns ret expReturns (Op op) =-  opReturns op+  Just <$> opReturns op expReturns (WithAcc inputs lam) =-  (<>)-    <$> (concat <$> mapM inputReturns inputs)-    <*>-    -- XXX: this is a bit dubious because it enforces extra copies.  I-    -- think WithAcc should perhaps have a return annotation like If.-    pure (extReturns $ staticShapes $ drop num_accs $ lambdaReturnType lam)+  Just+    <$> ( (<>)+            <$> (concat <$> mapM inputReturns inputs)+            <*>+            -- XXX: this is a bit dubious because it enforces extra copies.  I+            -- think WithAcc should perhaps have a return annotation like If.+            pure (extReturns $ staticShapes $ drop num_accs $ lambdaReturnType lam)+        )   where     inputReturns (_, arrs, _) = mapM varReturns arrs     num_accs = length inputs@@ -1107,14 +1122,13 @@   (Monad m, HasScope rep m, Mem rep inner) =>   VName ->   FlatSlice SubExp ->-  m (MemInfo SubExp NoUniqueness MemBind)+  m (Maybe (MemInfo SubExp NoUniqueness MemBind)) flatSliceInfo v slice@(FlatSlice offset idxs) = do   (et, _, mem, ixfun) <- arrayVarReturns v   map (fmap pe64) idxs     & FlatSlice (pe64 offset)     & IxFun.flatSlice ixfun-    & ArrayIn mem-    & MemArray et (Shape (flatSliceDims slice)) NoUniqueness+    & fmap (MemArray et (Shape (flatSliceDims slice)) NoUniqueness . ArrayIn mem)     & pure  class IsOp op => OpReturns op where

@@ -7,7 +7,6 @@     Shape,     LMAD (..),     LMADDim (..),-    Monotonicity (..),     index,     mkExistential,     iota,@@ -19,20 +18,15 @@     flatSlice,     rebase,     shape,-    lmadShape,+    permutation,     rank,-    linearWithOffset,-    rearrangeWithOffset,     isDirect,-    isLinear,     substituteInIxFun,     substituteInLMAD,     existentialize,     closeEnough,     equivalent,-    hasOneLmad,     permuteInv,-    conservativeFlatten,     disjoint,     disjoint2,     disjoint3,@@ -43,25 +37,27 @@ import Control.Category import Control.Monad import Control.Monad.State-import Data.Function (on, (&))-import Data.List (sort, sortBy, zip4, zipWith4)-import Data.List.NonEmpty (NonEmpty (..))-import Data.List.NonEmpty qualified as NE+import Data.List (sort, zip4) import Data.Map.Strict qualified as M-import Data.Maybe (isJust) import Data.Traversable import Futhark.Analysis.PrimExp import Futhark.Analysis.PrimExp.Convert-import Futhark.IR.Mem.LMAD+import Futhark.IR.Mem.LMAD hiding+  ( flatSlice,+    index,+    iota,+    mkExistential,+    permutation,+    permute,+    reshape,+    shape,+    slice,+  )+import Futhark.IR.Mem.LMAD qualified as LMAD import Futhark.IR.Prop import Futhark.IR.Syntax-  ( DimIndex (..),-    FlatDimIndex (..),-    FlatSlice (..),+  ( FlatSlice (..),     Slice (..),-    dimFix,-    flatSliceDims,-    flatSliceStrides,     unitSlice,   ) import Futhark.IR.Syntax.Core (Ext (..))@@ -71,8 +67,6 @@ import Futhark.Util.Pretty import Prelude hiding (gcd, id, mod, (.)) -type Indices num = [num]- -- | An index function is a mapping from a multidimensional array -- index space (the domain) to a one-dimensional memory index space. -- Essentially, it explains where the element at position @[i,j,p]@ of@@ -80,23 +74,20 @@ -- constitutes its memory.  For example, we can use this to -- distinguish row-major and column-major representations. ----- An index function is represented as a sequence of 'LMAD's.+-- An index function is represented as an LMAD. data IxFun num = IxFun-  { ixfunLMADs :: NonEmpty (LMAD num),+  { ixfunLMAD :: LMAD num,     -- | the shape of the support array, i.e., the original array     --   that birthed (is the start point) of this index function.-    base :: Shape num,-    -- | ignoring permutations, is the index function contiguous?-    contiguous :: Bool+    base :: Shape num   }   deriving (Show, Eq)  instance Pretty num => Pretty (IxFun num) where-  pretty (IxFun lmads oshp cg) =+  pretty (IxFun lmad oshp) =     braces . semistack $       [ "base:" <+> brackets (commasep $ map pretty oshp),-        "contiguous:" <+> if cg then "true" else "false",-        "LMADs:" <+> brackets (commastack $ NE.toList $ NE.map pretty lmads)+        "LMAD:" <+> pretty lmad       ]  instance Substitute num => Substitute (IxFun num) where@@ -117,23 +108,8 @@ -- It is important that the traversal order here is the same as in -- mkExistential. instance Traversable IxFun where-  traverse f (IxFun lmads oshp cg) =-    IxFun <$> traverse (traverse f) lmads <*> traverse f oshp <*> pure cg--(++@) :: [a] -> NonEmpty a -> NonEmpty a-es ++@ (ne :| nes) = case es of-  e : es' -> e :| es' ++ [ne] ++ nes-  [] -> ne :| nes--(@++@) :: NonEmpty a -> NonEmpty a -> NonEmpty a-(x :| xs) @++@ (y :| ys) = x :| xs ++ [y] ++ ys--setLMADPermutation :: Permutation -> LMAD num -> LMAD num-setLMADPermutation perm lmad =-  lmad {lmadDims = zipWith (\dim p -> dim {ldPerm = p}) (lmadDims lmad) perm}--setLMADShape :: Shape num -> LMAD num -> LMAD num-setLMADShape shp lmad = lmad {lmadDims = zipWith (\dim s -> dim {ldShape = s}) (lmadDims lmad) shp}+  traverse f (IxFun lmad oshp) =+    IxFun <$> traverse f lmad <*> traverse f oshp  -- | Substitute a name with a PrimExp in an index function. substituteInIxFun ::@@ -141,44 +117,39 @@   M.Map a (TPrimExp t a) ->   IxFun (TPrimExp t a) ->   IxFun (TPrimExp t a)-substituteInIxFun tab (IxFun lmads oshp cg) =+substituteInIxFun tab (IxFun lmad oshp) =   IxFun-    (NE.map (substituteInLMAD tab) lmads)+    (substituteInLMAD tab lmad)     (map (TPrimExp . substituteInPrimExp tab' . untyped) oshp)-    cg   where     tab' = fmap untyped tab  -- | Is this is a row-major array? isDirect :: (Eq num, IntegralExp num) => IxFun num -> Bool-isDirect ixfun@(IxFun (LMAD offset dims :| []) oshp True) =+isDirect (IxFun (LMAD offset dims) oshp) =   let strides_expected = reverse $ scanl (*) 1 (reverse (tail oshp))-   in hasContiguousPerm ixfun-        && length oshp == length dims+   in length oshp == length dims         && offset == 0         && all-          (\(LMADDim s n p _, m, d, se) -> s == se && n == d && p == m)+          (\(LMADDim s n p, m, d, se) -> s == se && n == d && p == m)           (zip4 dims [0 .. length dims - 1] oshp strides_expected)-isDirect _ = False --- | Is index function "analyzable", i.e., consists of one LMAD-hasOneLmad :: IxFun num -> Bool-hasOneLmad (IxFun (_ :| []) _ _) = True-hasOneLmad _ = False- -- | Does the index function have an ascending permutation? hasContiguousPerm :: IxFun num -> Bool-hasContiguousPerm (IxFun (lmad :| []) _ _) =-  let perm = lmadPermutation lmad+hasContiguousPerm (IxFun lmad _) =+  let perm = LMAD.permutation lmad    in perm == sort perm-hasContiguousPerm _ = False  -- | The index space of the index function.  This is the same as the -- shape of arrays that the index function supports. shape :: (Eq num, IntegralExp num) => IxFun num -> Shape num-shape (IxFun (lmad :| _) _ _) =-  permuteFwd (lmadPermutation lmad) $ lmadShapeBase lmad+shape (IxFun lmad _) =+  permuteFwd (LMAD.permutation lmad) $ LMAD.shapeBase lmad +-- | The permutation of the first LMAD of the index function.+permutation :: IxFun num -> Permutation+permutation = map LMAD.ldPerm . LMAD.dims . ixfunLMAD+ -- | Compute the flat memory index for a complete set @inds@ of array indices -- and a certain element size @elem_size@. index ::@@ -186,52 +157,24 @@   IxFun num ->   Indices num ->   num-index = indexFromLMADs . ixfunLMADs-  where-    indexFromLMADs ::-      (IntegralExp num, Eq num) =>-      NonEmpty (LMAD num) ->-      Indices num ->-      num-    indexFromLMADs (lmad :| []) inds = indexLMAD lmad inds-    indexFromLMADs (lmad1 :| lmad2 : lmads) inds =-      let i_flat = indexLMAD lmad1 inds-          new_inds = unflattenIndex (permuteFwd (lmadPermutation lmad2) $ lmadShapeBase lmad2) i_flat-       in indexFromLMADs (lmad2 :| lmads) new_inds-    indexLMAD ::-      (IntegralExp num, Eq num) =>-      LMAD num ->-      Indices num ->-      num-    indexLMAD lmad@(LMAD off dims) inds =-      let prod =-            sum $-              zipWith-                flatOneDim-                (map ldStride dims)-                (permuteInv (lmadPermutation lmad) inds)-       in off + prod+index = LMAD.index . ixfunLMAD  -- | iota with offset. iotaOffset :: IntegralExp num => num -> Shape num -> IxFun num-iotaOffset o ns = IxFun (makeRotIota Inc o ns :| []) ns True+iotaOffset o ns = IxFun (LMAD.iota o ns) ns  -- | iota. iota :: IntegralExp num => Shape num -> IxFun num iota = iotaOffset 0 --- | Create a contiguous single-LMAD index function that is--- existential in everything, with the provided permutation,--- monotonicity, and contiguousness.-mkExistential :: Int -> [(Int, Monotonicity)] -> Bool -> Int -> IxFun (Ext a)-mkExistential basis_rank perm contig start =-  IxFun (NE.singleton lmad) basis contig+-- | Create a single-LMAD index function that is+-- existential in everything, with the provided permutation.+mkExistential :: Int -> [Int] -> Int -> IxFun (Ext a)+mkExistential basis_rank perm start =+  IxFun (LMAD.mkExistential perm start) basis   where     basis = take basis_rank $ map Ext [start + 1 + dims_rank * 2 ..]     dims_rank = length perm-    lmad = LMAD (Ext start) $ zipWith onDim perm [0 ..]-    onDim (p, mon) i =-      LMADDim (Ext (start + 1 + i * 2)) (Ext (start + 2 + i * 2)) p mon  -- | Permute dimensions. permute ::@@ -239,114 +182,8 @@   IxFun num ->   Permutation ->   IxFun num-permute (IxFun (lmad :| lmads) oshp cg) perm_new =-  let perm_cur = lmadPermutation lmad-      perm = map (perm_cur !!) perm_new-   in IxFun (setLMADPermutation perm lmad :| lmads) oshp cg---- | Handle the case where a slice can stay within a single LMAD.-sliceOneLMAD ::-  (Eq num, IntegralExp num) =>-  IxFun num ->-  Slice num ->-  Maybe (IxFun num)-sliceOneLMAD (IxFun (lmad@(LMAD _ ldims) :| lmads) oshp cg) (Slice is) = do-  let perm = lmadPermutation lmad-      is' = permuteInv perm is-      cg' = cg && slicePreservesContiguous lmad (Slice is')-  let lmad' = foldl sliceOne (LMAD (lmadOffset lmad) []) $ zip is' ldims-      -- need to remove the fixed dims from the permutation-      perm' =-        updatePerm perm $-          map fst $-            filter (isJust . dimFix . snd) $-              zip [0 .. length is' - 1] is'--  pure $ IxFun (setLMADPermutation perm' lmad' :| lmads) oshp cg'-  where-    updatePerm ps inds = concatMap decrease ps-      where-        decrease p =-          let f n i-                | i == p = -1-                | i > p = n-                | n /= -1 = n + 1-                | otherwise = n-              d = foldl f 0 inds-           in [p - d | d /= -1]--    -- XXX: TODO: what happens to r on a negative-stride slice; is there-    -- such a case?-    sliceOne ::-      (Eq num, IntegralExp num) =>-      LMAD num ->-      (DimIndex num, LMADDim num) ->-      LMAD num-    sliceOne (LMAD off dims) (DimFix i, LMADDim s _x _ _) =-      LMAD (off + flatOneDim s i) dims-    sliceOne (LMAD off dims) (DimSlice _ ne _, LMADDim 0 _ p _) =-      LMAD off (dims ++ [LMADDim 0 ne p Unknown])-    sliceOne (LMAD off dims) (dmind, dim@(LMADDim _ n _ _))-      | dmind == unitSlice 0 n = LMAD off (dims ++ [dim])-    sliceOne (LMAD off dims) (dmind, LMADDim s n p m)-      | dmind == DimSlice (n - 1) n (-1) =-          let off' = off + flatOneDim s (n - 1)-           in LMAD off' (dims ++ [LMADDim (s * (-1)) n p (invertMonotonicity m)])-    sliceOne (LMAD off dims) (DimSlice b ne 0, LMADDim s _ p _) =-      LMAD (off + flatOneDim s b) (dims ++ [LMADDim 0 ne p Unknown])-    sliceOne (LMAD off dims) (DimSlice bs ns ss, LMADDim s _ p m) =-      let m' = case sgn ss of-            Just 1 -> m-            Just (-1) -> invertMonotonicity m-            _ -> Unknown-       in LMAD (off + s * bs) (dims ++ [LMADDim (ss * s) ns p m'])--    slicePreservesContiguous ::-      (Eq num, IntegralExp num) =>-      LMAD num ->-      Slice num ->-      Bool-    slicePreservesContiguous (LMAD _ dims) (Slice slc) =-      -- remove from the slice the LMAD dimensions that have stride 0.-      -- If the LMAD was contiguous in mem, then these dims will not-      -- influence the contiguousness of the result.-      -- Also normalize the input slice, i.e., 0-stride and size-1-      -- slices are rewritten as DimFixed.-      let (dims', slc') =-            unzip $-              filter ((/= 0) . ldStride . fst) $-                zip dims $-                  map normIndex slc-          -- Check that:-          -- 1. a clean split point exists between Fixed and Sliced dims-          -- 2. the outermost sliced dim has +/- 1 stride.-          -- 3. the rest of inner sliced dims are full.-          (_, success) =-            foldl-              ( \(found, res) (slcdim, LMADDim _ n _ _) ->-                  case (slcdim, found) of-                    (DimFix {}, True) -> (found, False)-                    (DimFix {}, False) -> (found, res)-                    (DimSlice _ _ ds, False) ->-                      -- outermost sliced dim: +/-1 stride-                      let res' = (ds == 1 || ds == -1)-                       in (True, res && res')-                    (DimSlice _ ne ds, True) ->-                      -- inner sliced dim: needs to be full-                      let res' = (n == ne) && (ds == 1 || ds == -1)-                       in (found, res && res')-              )-              (False, True)-              $ zip slc' dims'-       in success--    normIndex ::-      (Eq num, IntegralExp num) =>-      DimIndex num ->-      DimIndex num-    normIndex (DimSlice b 1 _) = DimFix b-    normIndex (DimSlice b _ 0) = DimFix b-    normIndex d = d+permute (IxFun lmad oshp) perm_new =+  IxFun (LMAD.permute lmad perm_new) oshp  -- | Slice an index function. slice ::@@ -354,46 +191,23 @@   IxFun num ->   Slice num ->   IxFun num-slice ixfun@(IxFun (lmad@(LMAD _ _) :| lmads) oshp cg) dim_slices+slice ixfun@(IxFun lmad@(LMAD _ _) oshp) (Slice is)   -- Avoid identity slicing.-  | unSlice dim_slices == map (unitSlice 0) (shape ixfun) = ixfun-  | Just ixfun' <- sliceOneLMAD ixfun dim_slices = ixfun'+  | is == map (unitSlice 0) (shape ixfun) = ixfun   | otherwise =-      case sliceOneLMAD (iota (lmadShape lmad)) dim_slices of-        Just (IxFun (lmad' :| []) _ cg') ->-          IxFun (lmad' :| lmad : lmads) oshp (cg && cg')-        _ -> error "slice: reached impossible case"+      IxFun (LMAD.slice lmad (Slice is)) oshp  -- | Flat-slice an index function. flatSlice ::   (Eq num, IntegralExp num) =>   IxFun num ->   FlatSlice num ->-  IxFun num-flatSlice ixfun@(IxFun (LMAD offset (dim : dims) :| lmads) oshp cg) (FlatSlice new_offset is)-  | hasContiguousPerm ixfun =-      let lmad =-            LMAD-              (offset + new_offset * ldStride dim)-              (map (helper $ ldStride dim) is <> dims)-              & setLMADPermutation [0 ..]-       in IxFun (lmad :| lmads) oshp cg-  where-    helper s0 (FlatDimIndex n s) =-      let new_mon = if s0 * s == 1 then Inc else Unknown-       in LMADDim (s0 * s) n 0 new_mon-flatSlice (IxFun (lmad :| lmads) oshp cg) s@(FlatSlice new_offset _) =-  IxFun (LMAD (new_offset * base_stride) (new_dims <> tail_dims) :| lmad : lmads) oshp cg-  where-    tail_shapes = tail $ lmadShape lmad-    base_stride = product tail_shapes-    tail_strides = tail $ scanr (*) 1 tail_shapes-    tail_dims = zipWith4 LMADDim tail_strides tail_shapes [length new_shapes ..] (repeat Inc)-    new_shapes = flatSliceDims s-    new_strides = map (* base_stride) $ flatSliceStrides s-    new_dims = zipWith4 LMADDim new_strides new_shapes [0 ..] (repeat Inc)+  Maybe (IxFun num)+flatSlice (IxFun lmad oshp) s = do+  lmad' <- LMAD.flatSlice lmad s+  Just $ IxFun lmad' oshp --- | Handle the case where a reshape operation can stay inside a single LMAD.+-- | Reshape an index function. -- -- There are four conditions that all must hold for the result of a reshape -- operation to remain in the one-LMAD domain:@@ -402,82 +216,18 @@ --       the LMAD dimensions that were *not* reshape coercions. --   (2) the repetition of dimensions of the underlying LMAD must --       refer only to the coerced-dimensions of the reshape operation.---   (3) finally, the underlying memory is contiguous (and monotonous). ----- If any of these conditions do not hold, then the reshape operation will--- conservatively add a new LMAD to the list, leading to a representation that--- provides less opportunities for further analysis.-reshapeOneLMAD ::-  (Eq num, IntegralExp num) =>-  IxFun num ->-  Shape num ->-  Maybe (IxFun num)-reshapeOneLMAD ixfun@(IxFun (lmad@(LMAD off dims) :| lmads) oldbase cg) newshape = do-  let perm = lmadPermutation lmad-      dims_perm = permuteFwd perm dims-      mid_dims = take (length dims) dims_perm-      mon = ixfunMonotonicity ixfun--  guard $-    -- checking conditions (2)-    all (\(LMADDim s _ _ _) -> s /= 0) mid_dims-      &&-      -- checking condition (1)-      consecutive 0 (map ldPerm mid_dims)-      &&-      -- checking condition (3)-      hasContiguousPerm ixfun-      && cg-      && (mon == Inc || mon == Dec)--  -- make new permutation-  let rsh_len = length newshape-      diff = length newshape - length dims-      iota_shape = [0 .. length newshape - 1]-      perm' =-        map-          ( \i ->-              let ind = i - diff-               in if (i >= 0) && (i < rsh_len)-                    then i -- already checked mid_dims not affected-                    else ldPerm (dims !! ind) + diff-          )-          iota_shape-      -- split the dimensions-      (support_inds, repeat_inds) =-        foldl-          (\(sup, rpt) (shpdim, ip) -> ((ip, shpdim) : sup, rpt))-          ([], [])-          $ reverse-          $ zip newshape perm'--      (sup_inds, support) = unzip $ sortBy (compare `on` fst) support_inds-      (rpt_inds, repeats) = unzip repeat_inds-      LMAD off' dims_sup = makeRotIota mon off support-      repeats' = map (\n -> LMADDim 0 n 0 Unknown) repeats-      dims' =-        map snd $-          sortBy (compare `on` fst) $-            zip sup_inds dims_sup ++ zip rpt_inds repeats'-      lmad' = LMAD off' dims'-  pure $ IxFun (setLMADPermutation perm' lmad' :| lmads) oldbase cg-  where-    consecutive _ [] = True-    consecutive i [p] = i == p-    consecutive i ps = and $ zipWith (==) ps [i, i + 1 ..]---- | Reshape an index function.+-- If any of these conditions do not hold, then the reshape operation+-- will conservatively add a new LMAD to the list, leading to a+-- representation that provides less opportunities for further+-- analysis reshape ::   (Eq num, IntegralExp num) =>   IxFun num ->   Shape num ->-  IxFun num-reshape ixfun new_shape-  | Just ixfun' <- reshapeOneLMAD ixfun new_shape = ixfun'-reshape (IxFun (lmad0 :| lmad0s) oshp cg) new_shape =-  case iota new_shape of-    IxFun (lmad :| []) _ _ -> IxFun (lmad :| lmad0 : lmad0s) oshp cg-    _ -> error "reshape: reached impossible case"+  Maybe (IxFun num)+reshape (IxFun lmad _) new_shape =+  IxFun <$> LMAD.reshape lmad new_shape <*> pure new_shape  -- | Coerce an index function to look like it has a new shape. -- Dynamically the shape must be the same.@@ -486,18 +236,15 @@   IxFun num ->   Shape num ->   IxFun num-coerce (IxFun (lmad :| lmads) oshp cg) new_shape =-  IxFun (onLMAD lmad :| lmads) oshp cg+coerce (IxFun lmad _) new_shape =+  IxFun (onLMAD lmad) new_shape   where     onLMAD (LMAD offset dims) = LMAD offset $ zipWith onDim dims new_shape     onDim ld d = ld {ldShape = d}  -- | The number of dimensions in the domain of the input function.-rank ::-  IntegralExp num =>-  IxFun num ->-  Int-rank (IxFun (LMAD _ sss :| _) _ _) = length sss+rank :: IntegralExp num => IxFun num -> Int+rank (IxFun (LMAD _ sss) _) = length sss  -- | Essentially @rebase new_base ixfun = ixfun o new_base@ -- Core soundness condition: @base ixfun == shape new_base@@@ -509,15 +256,6 @@ -- We can often stay in that domain if the original ixfun is essentially a -- slice, e.g. `x[i, (k1,m,s1), (k2,n,s2)] = orig`. ----- XXX: TODO: handle repetitions in both lmads.------ How to handle repeated dimensions in the original?------   (a) Shave them off of the last lmad of original---   (b) Compose the result from (a) with the first---       lmad of the new base---   (c) apply a repeat operation on the result of (b).--- -- However, I strongly suspect that for in-place update what we need is actually -- the INVERSE of the rebase function, i.e., given an index function new-base -- and another one orig, compute the index function ixfun0 such that:@@ -527,177 +265,67 @@ -- -- because then I can go bottom up and compose with ixfun0 all the index -- functions corresponding to the memory block associated with ixfun.-rebaseNice ::-  (Eq num, IntegralExp num) =>-  IxFun num ->-  IxFun num ->-  Maybe (IxFun num)-rebaseNice-  new_base@(IxFun (lmad_base :| lmads_base) _ cg_base)-  ixfun@(IxFun lmads shp cg) = do-    let (lmad :| lmads') = NE.reverse lmads-        dims = lmadDims lmad-        perm = lmadPermutation lmad-        perm_base = lmadPermutation lmad_base--    guard $-      -- Core rebase condition.-      base ixfun == shape new_base-        -- Conservative safety conditions: ixfun is contiguous and has known-        -- monotonicity for all dimensions.-        && cg-        && all ((/= Unknown) . ldMon) dims-        -- XXX: We should be able to handle some basic cases where both index-        -- functions have non-trivial permutations.-        && (hasContiguousPerm ixfun || hasContiguousPerm new_base)-        -- We need the permutations to be of the same size if we want to compose-        -- them.  They don't have to be of the same size if the ixfun has a trivial-        -- permutation.  Supporting this latter case allows us to rebase when ixfun-        -- has been created by slicing with fixed dimensions.-        && (length perm == length perm_base || hasContiguousPerm ixfun)-        -- To not have to worry about ixfun having non-1 strides, we also check that-        -- it is a row-major array (modulo permutation, which is handled-        -- separately).  Accept a non-full innermost dimension.  XXX: Maybe this can-        -- be less conservative?-        && and-          ( zipWith3-              (\sn ld inner -> sn == ldShape ld || (inner && ldStride ld == 1))-              shp-              dims-              (replicate (length dims - 1) False ++ [True])-          )--    -- Compose permutations, reverse strides and adjust offset if necessary.-    let perm_base' =-          if hasContiguousPerm ixfun-            then perm_base-            else map (perm !!) perm_base-        lmad_base' = setLMADPermutation perm_base' lmad_base-        dims_base = lmadDims lmad_base'-        n_fewer_dims = length dims_base - length dims-        (dims_base', offs_contrib) =-          unzip $-            zipWith-              ( \(LMADDim s1 n1 p1 _) (LMADDim _ _ _ m2) ->-                  let (s', off')-                        | m2 == Inc = (s1, 0)-                        | otherwise = (s1 * (-1), s1 * (n1 - 1))-                   in (LMADDim s' n1 (p1 - n_fewer_dims) Inc, off')-              )-              -- If @dims@ is morally a slice, it might have fewer dimensions than-              -- @dims_base@.  Drop extraneous outer dimensions.-              (drop n_fewer_dims dims_base)-              dims-        off_base = lmadOffset lmad_base' + sum offs_contrib-        lmad_base''-          | lmadOffset lmad == 0 = LMAD off_base dims_base'-          | otherwise =-              -- If the innermost dimension of the ixfun was not full (but still-              -- had a stride of 1), add its offset relative to the new base.-              setLMADShape-                (lmadShape lmad)-                ( LMAD-                    (off_base + ldStride (last dims_base) * lmadOffset lmad)-                    dims_base'-                )-        new_base' = IxFun (lmad_base'' :| lmads_base) shp cg_base-        IxFun lmads_base' _ _ = new_base'-        lmads'' = lmads' ++@ lmads_base'-    pure $ IxFun lmads'' shp (cg && cg_base)---- | Rebase an index function on top of a new base. rebase ::   (Eq num, IntegralExp num) =>   IxFun num ->   IxFun num ->-  IxFun num-rebase new_base@(IxFun lmads_base shp_base cg_base) ixfun@(IxFun lmads shp cg)-  | Just ixfun' <- rebaseNice new_base ixfun = ixfun'-  -- In the general case just concatenate LMADs since this refers to index-  -- function composition, which is always safe.-  | otherwise =-      let (lmads_base', shp_base') =-            if base ixfun == shape new_base-              then (lmads_base, shp_base)-              else-                let IxFun lmads' shp_base'' _ = reshape new_base shp-                 in (lmads', shp_base'')-       in IxFun (lmads @++@ lmads_base') shp_base' (cg && cg_base)---- | If the memory support of the index function is contiguous and row-major--- (i.e., no transpositions, repetitions, rotates, etc.), then this should--- return the offset from which the memory-support of this index function--- starts.-linearWithOffset ::-  (Eq num, IntegralExp num) =>-  IxFun num ->-  num ->-  Maybe num-linearWithOffset ixfun@(IxFun (lmad :| []) _ cg) elem_size-  | hasContiguousPerm ixfun && cg && ixfunMonotonicity ixfun == Inc =-      Just $ lmadOffset lmad * elem_size-linearWithOffset _ _ = Nothing---- | Similar restrictions to @linearWithOffset@ except for transpositions, which--- are returned together with the offset.-rearrangeWithOffset ::-  (Eq num, IntegralExp num) =>-  IxFun num ->-  num ->-  Maybe (num, [(Int, num)])-rearrangeWithOffset (IxFun (lmad :| []) oshp cg) elem_size = do-  -- Note that @cg@ describes whether the index function is-  -- contiguous, *ignoring permutations*.  This function requires that-  -- functionality.-  let perm = lmadPermutation lmad-      perm_contig = [0 .. length perm - 1]-  offset <--    linearWithOffset-      (IxFun (setLMADPermutation perm_contig lmad :| []) oshp cg)-      elem_size-  pure (offset, zip perm (permuteFwd perm (lmadShapeBase lmad)))-rearrangeWithOffset _ _ = Nothing---- | Is this a row-major array starting at offset zero?-isLinear :: (Eq num, IntegralExp num) => IxFun num -> Bool-isLinear = (== Just 0) . flip linearWithOffset 1--flatOneDim ::-  (Eq num, IntegralExp num) =>-  num ->-  num ->-  num-flatOneDim s i-  | s == 0 = 0-  | otherwise = i * s+  Maybe (IxFun num)+rebase new_base@(IxFun lmad_base _) ixfun@(IxFun lmad shp) = do+  let dims = LMAD.dims lmad+      perm = LMAD.permutation lmad+      perm_base = LMAD.permutation lmad_base --- | Check monotonicity of an index function.-ixfunMonotonicity ::-  (Eq num, IntegralExp num) =>-  IxFun num ->-  Monotonicity-ixfunMonotonicity (IxFun (lmad :| lmads) _ _) =-  let mon0 = lmadMonotonicityRots lmad-   in if all ((== mon0) . lmadMonotonicityRots) lmads-        then mon0-        else Unknown-  where-    lmadMonotonicityRots ::-      (Eq num, IntegralExp num) =>-      LMAD num ->-      Monotonicity-    lmadMonotonicityRots (LMAD _ dims)-      | all (isMonDim Inc) dims = Inc-      | all (isMonDim Dec) dims = Dec-      | otherwise = Unknown+  guard $+    -- Core rebase condition.+    base ixfun == shape new_base+      -- XXX: We should be able to handle some basic cases where both index+      -- functions have non-trivial permutations.+      && (hasContiguousPerm ixfun || hasContiguousPerm new_base)+      -- We need the permutations to be of the same size if we want to compose+      -- them.  They don't have to be of the same size if the ixfun has a trivial+      -- permutation.  Supporting this latter case allows us to rebase when ixfun+      -- has been created by slicing with fixed dimensions.+      && (length perm == length perm_base || hasContiguousPerm ixfun)+      -- To not have to worry about ixfun having non-1 strides, we also check that+      -- it is a row-major array (modulo permutation, which is handled+      -- separately).  Accept a non-full outermost dimension.  XXX: Maybe this can+      -- be less conservative?+      && and+        ( zipWith3+            (\sn ld inner -> inner || sn == ldShape ld)+            shp+            dims+            (True : replicate (length dims - 1) False)+        ) -    isMonDim ::-      (Eq num, IntegralExp num) =>-      Monotonicity ->-      LMADDim num ->-      Bool-    isMonDim mon (LMADDim s _ _ ldmon) =-      s == 0 || mon == ldmon+  -- Compose permutations, reverse strides and adjust offset if necessary.+  let perm_base' =+        if hasContiguousPerm ixfun+          then perm_base+          else map (perm !!) perm_base+      lmad_base' = LMAD.setPermutation perm_base' lmad_base+      dims_base = LMAD.dims lmad_base'+      n_fewer_dims = length dims_base - length dims+      (dims_base', offs_contrib) =+        unzip $+          zipWith+            ( \(LMADDim s1 n1 p1) (LMADDim {}) ->+                let (s', off') = (s1, 0)+                 in (LMADDim s' n1 (p1 - n_fewer_dims), off')+            )+            -- If @dims@ is morally a slice, it might have fewer dimensions than+            -- @dims_base@.  Drop extraneous outer dimensions.+            (drop n_fewer_dims dims_base)+            dims+      off_base = LMAD.offset lmad_base' + sum offs_contrib+      lmad_base'' =+        LMAD.setShape+          (LMAD.shape lmad)+          ( LMAD+              (off_base + ldStride (last dims_base) * LMAD.offset lmad)+              dims_base'+          )+  pure $ IxFun lmad_base'' shp  -- | Turn all the leaves of the index function into 'Ext's.  We --  require that there's only one LMAD, that the index function is@@ -725,31 +353,22 @@ closeEnough :: IxFun num -> IxFun num -> Bool closeEnough ixf1 ixf2 =   (length (base ixf1) == length (base ixf2))-    && (NE.length (ixfunLMADs ixf1) == NE.length (ixfunLMADs ixf2))-    && all closeEnoughLMADs (NE.zip (ixfunLMADs ixf1) (ixfunLMADs ixf2))-    -- This treats ixf1 as the "declared type" that we are matching against.-    && (contiguous ixf1 <= contiguous ixf2)+    && closeEnoughLMADs (ixfunLMAD ixf1) (ixfunLMAD ixf2)   where-    closeEnoughLMADs :: (LMAD num, LMAD num) -> Bool-    closeEnoughLMADs (lmad1, lmad2) =-      length (lmadDims lmad1) == length (lmadDims lmad2)-        && map ldPerm (lmadDims lmad1)-          == map ldPerm (lmadDims lmad2)+    closeEnoughLMADs lmad1 lmad2 =+      length (LMAD.dims lmad1) == length (LMAD.dims lmad2)+        && map ldPerm (LMAD.dims lmad1) == map ldPerm (LMAD.dims lmad2)  -- | Returns true if two 'IxFun's are equivalent. -- -- Equivalence in this case is defined as having the same number of LMADs, with--- each pair of LMADs matching in permutation, offsets, strides and rotations.+-- each pair of LMADs matching in permutation, offsets, and strides. equivalent :: Eq num => IxFun num -> IxFun num -> Bool equivalent ixf1 ixf2 =-  NE.length (ixfunLMADs ixf1) == NE.length (ixfunLMADs ixf2)-    && all equivalentLMADs (NE.zip (ixfunLMADs ixf1) (ixfunLMADs ixf2))+  equivalentLMADs (ixfunLMAD ixf1) (ixfunLMAD ixf2)   where-    equivalentLMADs (lmad1, lmad2) =-      length (lmadDims lmad1) == length (lmadDims lmad2)-        && map ldPerm (lmadDims lmad1)-          == map ldPerm (lmadDims lmad2)-        && lmadOffset lmad1-          == lmadOffset lmad2-        && map ldStride (lmadDims lmad1)-          == map ldStride (lmadDims lmad2)+    equivalentLMADs lmad1 lmad2 =+      length (LMAD.dims lmad1) == length (LMAD.dims lmad2)+        && map ldPerm (LMAD.dims lmad1) == map ldPerm (LMAD.dims lmad2)+        && LMAD.offset lmad1 == LMAD.offset lmad2+        && map ldStride (LMAD.dims lmad1) == map ldStride (LMAD.dims lmad2)

@@ -1,40 +1,62 @@ -- | This module contains a representation of linear-memory accessor -- descriptors (LMAD); see work by Zhu, Hoeflinger and David.+--+-- This module is designed to be used as a qualified import, as the+-- exported names are quite generic. module Futhark.IR.Mem.LMAD   ( Shape,+    Indices,     LMAD (..),     LMADDim (..),-    Monotonicity (..),     Permutation,-    lmadShape,-    lmadShapeBase,+    index,+    slice,+    flatSlice,+    reshape,+    permute,+    shape,+    permutation,+    shapeBase,+    setPermutation,+    setShape,     substituteInLMAD,     permuteInv,     permuteFwd,-    conservativeFlatten,     disjoint,     disjoint2,     disjoint3,     dynamicEqualsLMAD,-    lmadPermutation,-    makeRotIota,-    invertMonotonicity,+    contiguous,+    memcpyable,+    noPermutation,+    iota,+    mkExistential,   ) where  import Control.Category import Control.Monad+import Data.Foldable (toList) import Data.Function (on, (&))-import Data.List (elemIndex, partition, sortBy, zipWith4)+import Data.List (elemIndex, partition, sort, sortBy) import Data.Map.Strict qualified as M-import Data.Maybe (fromJust, isNothing)+import Data.Maybe (fromJust, isJust, isNothing) import Data.Traversable import Futhark.Analysis.AlgSimplify qualified as AlgSimplify import Futhark.Analysis.PrimExp import Futhark.Analysis.PrimExp.Convert import Futhark.IR.Mem.Interval import Futhark.IR.Prop-import Futhark.IR.Syntax (Type)+import Futhark.IR.Syntax+  ( DimIndex (..),+    Ext (..),+    FlatDimIndex (..),+    FlatSlice (..),+    Slice (..),+    Type,+    dimFix,+    unitSlice,+  ) import Futhark.IR.Syntax.Core (VName (..)) import Futhark.Transform.Rename import Futhark.Transform.Substitute@@ -46,51 +68,26 @@ -- | The shape of an index function. type Shape num = [num] -type Permutation = [Int]+-- | Indices passed to an LMAD.  Must always match the rank of the LMAD.+type Indices num = [num] --- | The physical element ordering alongside a dimension, i.e. the--- sign of the stride.-data Monotonicity-  = -- | Increasing.-    Inc-  | -- | Decreasing.-    Dec-  | -- | Unknown.-    Unknown-  deriving (Show, Eq)+-- | A complete permutation.+type Permutation = [Int]  -- | A single dimension in an 'LMAD'. data LMADDim num = LMADDim   { ldStride :: num,     ldShape :: num,-    ldPerm :: Int,-    ldMon :: Monotonicity+    ldPerm :: Int   }-  deriving (Show, Eq)--instance Ord Monotonicity where-  (<=) _ Inc = True-  (<=) Unknown _ = True-  (<=) _ Unknown = False-  (<=) Inc Dec = False-  (<=) _ Dec = True--instance Ord num => Ord (LMADDim num) where-  (LMADDim s1 q1 p1 m1) <= (LMADDim s2 q2 p2 m2) =-    ([q1, s1] < [q2, s2])-      || ( ([q1, s1] == [q2, s2])-             && ( (p1 < p2)-                    || ( (p1 == p2)-                           && (m1 <= m2)-                       )-                )-         )+  deriving (Show, Eq, Ord) --- | LMAD's representation consists of a general offset and for each dimension a--- stride, number of elements (or shape), permutation, and--- monotonicity. Note that the permutation is not strictly necessary in that the--- permutation can be performed directly on LMAD dimensions, but then it is--- difficult to extract the permutation back from an LMAD.+-- | LMAD's representation consists of a general offset and for each+-- dimension a stride, number of elements (or shape), and+-- permutation. Note that the permutation is not strictly necessary in+-- that the permutation can be performed directly on LMAD dimensions,+-- but then it is difficult to extract the permutation back from an+-- LMAD. -- -- LMAD algebra is closed under composition w.r.t. operators such as -- permute, index and slice.  However, other operations, such as@@ -116,22 +113,18 @@ --    \{ ~ \sigma + i_1 * s_1 + \ldots + i_m * s_m ~ | ~ 0 \leq i_1 < n_1, \ldots, 0 \leq i_m < n_m ~ \} -- \] data LMAD num = LMAD-  { lmadOffset :: num,-    lmadDims :: [LMADDim num]+  { offset :: num,+    dims :: [LMADDim num]   }   deriving (Show, Eq, Ord) -instance Pretty Monotonicity where-  pretty = pretty . show- instance Pretty num => Pretty (LMAD num) where   pretty (LMAD offset dims) =     braces . semistack $       [ "offset:" <+> group (pretty offset),         "strides:" <+> p ldStride,         "shape:" <+> p ldShape,-        "permutation:" <+> p ldPerm,-        "monotonicity:" <+> p ldMon+        "permutation:" <+> p ldPerm       ]     where       p f = group $ brackets $ align $ commasep $ map (pretty . f) dims@@ -146,7 +139,7 @@   freeIn' = foldMap freeIn'  instance FreeIn num => FreeIn (LMADDim num) where-  freeIn' (LMADDim s n _ _) = freeIn' s <> freeIn' n+  freeIn' (LMADDim s n _) = freeIn' s <> freeIn' n  instance Functor LMAD where   fmap = fmapDefault@@ -158,16 +151,188 @@   traverse f (LMAD offset dims) =     LMAD <$> f offset <*> traverse f' dims     where-      f' (LMADDim s n p m) = LMADDim <$> f s <*> f n <*> pure p <*> pure m+      f' (LMADDim s n p) = LMADDim <$> f s <*> f n <*> pure p -invertMonotonicity :: Monotonicity -> Monotonicity-invertMonotonicity Inc = Dec-invertMonotonicity Dec = Inc-invertMonotonicity Unknown = Unknown+flatOneDim ::+  (Eq num, IntegralExp num) =>+  num ->+  num ->+  num+flatOneDim s i+  | s == 0 = 0+  | otherwise = i * s -lmadPermutation :: LMAD num -> Permutation-lmadPermutation = map ldPerm . lmadDims+index :: (IntegralExp num, Eq num) => LMAD num -> Indices num -> num+index lmad@(LMAD off dims) inds =+  off + sum prods+  where+    prods =+      zipWith+        flatOneDim+        (map ldStride dims)+        (permuteInv (permutation lmad) inds) +setLMADPermutation :: Permutation -> LMAD num -> LMAD num+setLMADPermutation perm lmad =+  lmad {dims = zipWith (\dim p -> dim {ldPerm = p}) (dims lmad) perm}++-- | Handle the case where a slice can stay within a single LMAD.+slice ::+  (Eq num, IntegralExp num) =>+  LMAD num ->+  Slice num ->+  LMAD num+slice lmad@(LMAD _ ldims) (Slice is) =+  let perm = permutation lmad+      is' = permuteInv perm is+      lmad' = foldl sliceOne (LMAD (offset lmad) []) $ zip is' ldims+      -- need to remove the fixed dims from the permutation+      perm' =+        updatePerm perm $+          map fst $+            filter (isJust . dimFix . snd) $+              zip [0 .. length is' - 1] is'+   in setLMADPermutation perm' lmad'+  where+    updatePerm ps inds = concatMap decrease ps+      where+        decrease p =+          let f n i+                | i == p = -1+                | i > p = n+                | n /= -1 = n + 1+                | otherwise = n+              d = foldl f 0 inds+           in [p - d | d /= -1]++    sliceOne ::+      (Eq num, IntegralExp num) =>+      LMAD num ->+      (DimIndex num, LMADDim num) ->+      LMAD num+    sliceOne (LMAD off dims) (DimFix i, LMADDim s _x _) =+      LMAD (off + flatOneDim s i) dims+    sliceOne (LMAD off dims) (DimSlice _ ne _, LMADDim 0 _ p) =+      LMAD off (dims ++ [LMADDim 0 ne p])+    sliceOne (LMAD off dims) (dmind, dim@(LMADDim _ n _))+      | dmind == unitSlice 0 n = LMAD off (dims ++ [dim])+    sliceOne (LMAD off dims) (dmind, LMADDim s n p)+      | dmind == DimSlice (n - 1) n (-1) =+          let off' = off + flatOneDim s (n - 1)+           in LMAD off' (dims ++ [LMADDim (s * (-1)) n p])+    sliceOne (LMAD off dims) (DimSlice b ne 0, LMADDim s _ p) =+      LMAD (off + flatOneDim s b) (dims ++ [LMADDim 0 ne p])+    sliceOne (LMAD off dims) (DimSlice bs ns ss, LMADDim s _ p) =+      LMAD (off + s * bs) (dims ++ [LMADDim (ss * s) ns p])++hasContiguousPerm :: LMAD num -> Bool+hasContiguousPerm lmad = perm == sort perm+  where+    perm = permutation lmad++-- | Flat-slice an LMAD.+flatSlice ::+  (IntegralExp num) =>+  LMAD num ->+  FlatSlice num ->+  Maybe (LMAD num)+flatSlice lmad@(LMAD offset (dim : dims)) (FlatSlice new_offset is)+  | hasContiguousPerm lmad =+      Just $+        LMAD+          (offset + new_offset * ldStride dim)+          (map (helper $ ldStride dim) is <> dims)+          & setLMADPermutation [0 ..]+  where+    helper s0 (FlatDimIndex n s) = LMADDim (s0 * s) n 0+flatSlice _ _ = Nothing++-- | Handle the case where a reshape operation can stay inside a+-- single LMAD.  See "Futhark.IR.Mem.IxFun.reshape" for+-- conditions.+reshape ::+  (Eq num, IntegralExp num) => LMAD num -> Shape num -> Maybe (LMAD num)+--+-- First a special case for when we are merely injecting unit+-- dimensions into a non-permuted LMAD.+reshape lmad@(LMAD off dims) newshape+  | sort (permutation lmad) == permutation lmad,+    Just dims' <- addingVacuous 0 newshape dims =+      Just $ LMAD off dims'+  where+    addingVacuous i (dnew : dnews) (dold : dolds)+      | dnew == ldShape dold =+          (dold {ldPerm = i} :) <$> addingVacuous (i + 1) dnews dolds+    addingVacuous i (1 : dnews) dolds =+      (LMADDim 0 1 i :) <$> addingVacuous (i + 1) dnews dolds+    addingVacuous _ [] [] = Just []+    addingVacuous _ _ _ = Nothing++-- Then the general case.+reshape lmad@(LMAD off dims) newshape = do+  let perm = permutation lmad+      dims_perm = permuteFwd perm dims+      mid_dims = take (length dims) dims_perm++  guard $+    -- checking conditions (2)+    all (\(LMADDim s _ _) -> s /= 0) mid_dims+      &&+      -- checking condition (1)+      consecutive 0 (map ldPerm mid_dims)+      &&+      -- checking condition (3)+      hasContiguousPerm lmad+      && all+        (\(ld, se) -> ldStride ld == se)+        (zip dims (reverse $ scanl (*) 1 (reverse (tail (shape lmad)))))++  -- make new permutation+  let rsh_len = length newshape+      diff = length newshape - length dims+      iota_shape = [0 .. length newshape - 1]+      perm' =+        map+          ( \i ->+              let ind = i - diff+               in if (i >= 0) && (i < rsh_len)+                    then i -- already checked mid_dims not affected+                    else ldPerm (dims !! ind) + diff+          )+          iota_shape+      -- split the dimensions+      (support_inds, repeat_inds) =+        foldl+          (\(sup, rpt) (shpdim, ip) -> ((ip, shpdim) : sup, rpt))+          ([], [])+          $ reverse+          $ zip newshape perm'++      (sup_inds, support) = unzip $ sortBy (compare `on` fst) support_inds+      (rpt_inds, repeats) = unzip repeat_inds+      LMAD off' dims_sup = iota off support+      repeats' = map (\n -> LMADDim 0 n 0) repeats+      dims' =+        map snd $+          sortBy (compare `on` fst) $+            zip sup_inds dims_sup ++ zip rpt_inds repeats'+      lmad' = LMAD off' dims'+  Just $ setLMADPermutation perm' lmad'+  where+    consecutive _ [] = True+    consecutive i [p] = i == p+    consecutive i ps = and $ zipWith (==) ps [i, i + 1 ..]++permutation :: LMAD num -> Permutation+permutation = map ldPerm . dims++setPermutation :: Permutation -> LMAD num -> LMAD num+setPermutation perm lmad =+  lmad {dims = zipWith (\dim p -> dim {ldPerm = p}) (dims lmad) perm}++setShape :: Shape num -> LMAD num -> LMAD num+setShape shp lmad = lmad {dims = zipWith (\dim s -> dim {ldShape = s}) (dims lmad) shp}+ -- | Substitute a name with a PrimExp in an LMAD. substituteInLMAD ::   Ord a =>@@ -175,29 +340,19 @@   LMAD (TPrimExp t a) ->   LMAD (TPrimExp t a) substituteInLMAD tab (LMAD offset dims) =-  let offset' = sub offset-      dims' =-        map-          ( \(LMADDim s n p m) ->-              LMADDim-                (sub s)-                (sub n)-                p-                m-          )-          dims-   in LMAD offset' dims'+  LMAD (sub offset) $+    map (\(LMADDim s n p) -> LMADDim (sub s) (sub n) p) dims   where     tab' = fmap untyped tab     sub = TPrimExp . substituteInPrimExp tab' . untyped  -- | Shape of an LMAD.-lmadShape :: LMAD num -> Shape num-lmadShape lmad = permuteInv (lmadPermutation lmad) $ lmadShapeBase lmad+shape :: LMAD num -> Shape num+shape lmad = permuteInv (permutation lmad) $ shapeBase lmad  -- | Shape of an LMAD, ignoring permutations.-lmadShapeBase :: LMAD num -> Shape num-lmadShapeBase = map ldShape . lmadDims+shapeBase :: LMAD num -> Shape num+shapeBase = map ldShape . dims  permuteFwd :: Permutation -> [a] -> [a] permuteFwd ps elems = map (elems !!) ps@@ -205,28 +360,37 @@ permuteInv :: Permutation -> [a] -> [a] permuteInv ps elems = map snd $ sortBy (compare `on` fst) $ zip ps elems --- | Generalised iota with user-specified offset and rotates.-makeRotIota ::+-- | Generalised iota with user-specified offset.+iota ::   IntegralExp num =>-  Monotonicity ->   -- | Offset   num ->   -- | Shape   [num] ->   LMAD num-makeRotIota mon off ns-  | mon == Inc || mon == Dec =-      let rk = length ns-          ss0 = reverse $ take rk $ scanl (*) 1 $ reverse ns-          ss =-            if mon == Inc-              then ss0-              else map (* (-1)) ss0-          ps = map fromIntegral [0 .. rk - 1]-          fi = replicate rk mon-       in LMAD off $ zipWith4 LMADDim ss ns ps fi-  | otherwise = error "makeRotIota: requires Inc or Dec"+iota off ns =+  let rk = length ns+      ss = reverse $ take rk $ scanl (*) 1 $ reverse ns+      ps = map fromIntegral [0 .. rk - 1]+   in LMAD off $ zipWith3 LMADDim ss ns ps +-- | Create an LMAD that is existential in everything, with the+-- provided permutation.+mkExistential :: [Int] -> Int -> LMAD (Ext a)+mkExistential perm start =+  lmad+  where+    lmad = LMAD (Ext start) $ zipWith onDim perm [0 ..]+    onDim p i =+      LMADDim (Ext (start + 1 + i * 2)) (Ext (start + 2 + i * 2)) p++-- | Permute dimensions.+permute :: LMAD num -> Permutation -> LMAD num+permute lmad perm_new =+  let perm_cur = permutation lmad+      perm = map (perm_cur !!) perm_new+   in setPermutation perm lmad+ -- | Computes the maximum span of an 'LMAD'. The result is the lowest and -- highest flat values representable by that 'LMAD'. flatSpan :: LMAD (TPrimExp Int64 VName) -> TPrimExp Int64 VName@@ -248,7 +412,7 @@ -- conservativeFlatten :: (IntegralExp e, Ord e, Pretty e) => LMAD e -> LMAD e conservativeFlatten :: LMAD (TPrimExp Int64 VName) -> Maybe (LMAD (TPrimExp Int64 VName)) conservativeFlatten (LMAD offset []) =-  pure $ LMAD offset [LMADDim 1 1 0 Inc]+  pure $ LMAD offset [LMADDim 1 1 0] conservativeFlatten l@(LMAD _ [_]) =   pure l conservativeFlatten l@(LMAD offset dims) = do@@ -257,7 +421,7 @@       gcd       (ldStride $ head dims)       $ map ldStride dims-  pure $ LMAD offset [LMADDim strd (shp + 1) 0 Unknown]+  pure $ LMAD offset [LMADDim strd (shp + 1) 0]   where     shp = flatSpan l @@ -440,12 +604,12 @@ lmadToIntervals :: LMAD (TPrimExp Int64 VName) -> (AlgSimplify.SofP, [Interval]) lmadToIntervals (LMAD offset []) = (AlgSimplify.simplify0 $ untyped offset, [Interval 0 1 1]) lmadToIntervals lmad@(LMAD offset dims0) =-  (offset', map helper $ permuteInv (lmadPermutation lmad) dims0)+  (offset', map helper $ permuteInv (permutation lmad) dims0)   where     offset' = AlgSimplify.simplify0 $ untyped offset      helper :: LMADDim (TPrimExp Int64 VName) -> Interval-    helper (LMADDim strd shp _ _) = do+    helper (LMADDim strd shp _) = do       Interval 0 (AlgSimplify.simplify' shp) (AlgSimplify.simplify' strd)  -- | Dynamically determine if two 'LMADDim' are equal.@@ -456,15 +620,52 @@   ldStride dim1 .==. ldStride dim2     .&&. ldShape dim1 .==. ldShape dim2     .&&. fromBool (ldPerm dim1 == ldPerm dim2)-    .&&. fromBool (ldMon dim1 == ldMon dim2)  -- | Dynamically determine if two 'LMAD' are equal. -- -- True if offset and constituent 'LMADDim' are equal. dynamicEqualsLMAD :: Eq num => LMAD (TPrimExp t num) -> LMAD (TPrimExp t num) -> TPrimExp Bool num dynamicEqualsLMAD lmad1 lmad2 =-  lmadOffset lmad1 .==. lmadOffset lmad2+  offset lmad1 .==. offset lmad2     .&&. foldr       ((.&&.) . uncurry dynamicEqualsLMADDim)       true-      (zip (lmadDims lmad1) (lmadDims lmad2))+      (zip (dims lmad1) (dims lmad2))++-- | True if these LMADs represent the same function (ignoring+-- offset).+compatible ::+  Eq num =>+  LMAD (TPrimExp Int64 num) ->+  LMAD (TPrimExp Int64 num) ->+  TPrimExp Bool num+compatible x y =+  foldl1 (.&&.) $ zipWith dynamicEqualsLMADDim (dims x) (dims y)++-- | True if this LMAD corresponds to an array without "holes".  This+-- implies it can be copied with a memcpy()-like operation.+contiguous ::+  (Pretty num, Eq num) =>+  LMAD (TPrimExp Int64 num) ->+  TPrimExp Bool num+contiguous lmad =+  foldl1 (.&&.) $ zipWith (.==.) (toList lmad) lmad'+  where+    lmad' = toList (iota (offset lmad) $ map ldShape $ dims lmad)++-- | True if these LMADs have the same contiguous representation, such+-- that one can be copied to the other with a @memcpy()@-like+-- operation.+memcpyable ::+  (Pretty num, Eq num) =>+  LMAD (TPrimExp Int64 num) ->+  LMAD (TPrimExp Int64 num) ->+  TPrimExp Bool num+memcpyable dest_lmad src_lmad =+  contiguous dest_lmad .&&. compatible dest_lmad src_lmad++-- | Remove the permutation of an LMAD by actually applying it to the+-- dimensions.+noPermutation :: LMAD t -> LMAD t+noPermutation lmad =+  lmad {dims = rearrangeShape (permutation lmad) $ dims lmad}

@@ -17,9 +17,8 @@  import Data.Char (isAlpha) import Data.Functor-import Data.List (singleton, zipWith4)+import Data.List (singleton) import Data.List.NonEmpty (NonEmpty (..))-import Data.List.NonEmpty qualified as NE import Data.Maybe import Data.Set qualified as S import Data.Text qualified as T@@ -68,9 +67,6 @@     pTag = "_" *> L.decimal <* notFollowedBy (satisfy constituent)     exprBox = ("<{" <>) . (<> "}>") <$> (chunk "<{" *> manyTill anySingle (chunk "}>")) -pBool :: Parser Bool-pBool = choice [keyword "true" $> True, keyword "false" $> False]- pInt :: Parser Int pInt = lexeme L.decimal @@ -975,24 +971,16 @@ pIxFunBase pNum =   braces $ do     base <- pLab "base" $ brackets (pNum `sepBy` pComma) <* pSemi-    ct <- pLab "contiguous" $ pBool <* pSemi-    lmads <- pLab "LMADs" $ brackets (pLMAD `sepBy1` pComma)-    pure $ IxFun.IxFun (NE.fromList lmads) base ct+    lmad <- pLab "LMAD" pLMAD+    pure $ IxFun.IxFun lmad base   where     pLab s m = keyword s *> pColon *> m-    pMon =-      choice-        [ "Inc" $> IxFun.Inc,-          "Dec" $> IxFun.Dec,-          "Unknown" $> IxFun.Unknown-        ]     pLMAD = braces $ do       offset <- pLab "offset" pNum <* pSemi       strides <- pLab "strides" $ brackets (pNum `sepBy` pComma) <* pSemi       shape <- pLab "shape" $ brackets (pNum `sepBy` pComma) <* pSemi-      perm <- pLab "permutation" $ brackets (pInt `sepBy` pComma) <* pSemi-      mon <- pLab "monotonicity" $ brackets (pMon `sepBy` pComma)-      pure $ IxFun.LMAD offset $ zipWith4 IxFun.LMADDim strides shape perm mon+      perm <- pLab "permutation" $ brackets (pInt `sepBy` pComma)+      pure $ IxFun.LMAD offset $ zipWith3 IxFun.LMADDim strides shape perm  pPrimExpLeaf :: Parser VName pPrimExpLeaf = pVName

@@ -4,6 +4,7 @@ -- around huge records of which only a tiny part is needed. module Futhark.Internalise.ReplaceRecords (transformProg) where +import Control.Monad import Control.Monad.Reader import Control.Monad.State import Data.Map.Strict qualified as M

@@ -21,6 +21,7 @@   ) where +import Control.Monad import Control.Monad.Free (Free (..)) import Control.Monad.State import Data.Bitraversable (bitraverse)

@@ -89,8 +89,8 @@ -- | Given an 'IORef' tracking the state, produce a set of handlers. -- When we want to add more features to the language server, this is -- the thing to change.-handlers :: IORef State -> Handlers (LspM ())-handlers state_mvar =+handlers :: IORef State -> ClientCapabilities -> Handlers (LspM ())+handlers state_mvar _ =   mconcat     [ onInitializeHandler,       onDocumentOpenHandler state_mvar,

@@ -17,7 +17,6 @@ import Data.Function ((&)) import Data.List qualified as L import Data.List.NonEmpty (NonEmpty (..))-import Data.List.NonEmpty qualified as NE import Data.Map.Strict qualified as M import Data.Maybe import Data.Sequence (Seq (..))@@ -421,8 +420,10 @@                 (aggSummaryMapPartial (scalarTable td_env) $ unSegSpace space0)                 (S.toList s)             Undeterminable -> pure Undeterminable-        let res = noMemOverlap td_env destination_uses source_writes-        if res+        -- Do not allow short-circuiting from a segop-local memory+        -- block (not in the topdown scope) to an outer memory block.+        if dstmem entry `M.member` scope td_env+          && noMemOverlap td_env destination_uses source_writes           then pure bu_env_f           else do             let (ac, inh) = markFailedCoal (activeCoals bu_env_f, inhibit bu_env_f) k@@ -484,8 +485,8 @@                     Just (Coalesced knd mbd@(MemBlock _ _ _ ixfn) _) -> pure $                       case freeVarSubstitutions (scope td_env) (scalarTable td_env) ixfn of                         Just fv_subst ->-                          if ixfunPermutation ixfn-                            == ixfunPermutation (ixfun $ fromJust $ getScopeMemInfo (patElemName p) $ scope td_env)+                          if IxFun.permutation ixfn+                            == IxFun.permutation (ixfun $ fromJust $ getScopeMemInfo (patElemName p) $ scope td_env)                             then                               let entry =                                     coal_entry@@ -506,9 +507,6 @@    foldM checkMergeeOverlap bu_env''' mergee_writes -ixfunPermutation :: IxFun -> [Int]-ixfunPermutation = map IxFun.ldPerm . IxFun.lmadDims . NE.head . IxFun.ixfunLMADs- -- | Given a pattern element and the corresponding kernel result, try to put the -- kernel result directly in the memory block of pattern element makeSegMapCoals ::@@ -534,8 +532,7 @@       case M.lookup pat_mem active of         Nothing ->           -- We are not in a transitive case-          case ( IxFun.hasOneLmad pat_ixf-                   && maybe False (pat_mem `nameIn`) (M.lookup return_mem inhb),+          case ( maybe False (pat_mem `nameIn`) (M.lookup return_mem inhb),                  Coalesced                    InPlaceCoal                    (MemBlock tp return_shp pat_mem $ resultSlice pat_ixf)@@ -1259,8 +1256,7 @@                         _ -> (failed, s_acc) -- fail!  ixfunToAccessSummary :: IxFun.IxFun (TPrimExp Int64 VName) -> AccessSummary-ixfunToAccessSummary (IxFun.IxFun (lmad NE.:| []) _ _) = Set $ S.singleton lmad-ixfunToAccessSummary _ = Undeterminable+ixfunToAccessSummary (IxFun.IxFun lmad _) = Set $ S.singleton lmad  -- | Check safety conditions 2 and 5 and update new substitutions: -- called on the pat-elements of loop and if-then-else expressions.@@ -1365,9 +1361,8 @@                       Nothing ->                         ind_y                  in (m_y, alias_fn ind, oneName m_x <> y_al, x_deps0, certs <> certs'')-          success0 = IxFun.hasOneLmad ind_yx           m_b_aliased_m_yx = areAnyAliased td_env m_b [m_yx] -- m_b \= m_yx-       in if success0 && not m_b_aliased_m_yx && isInScope td_env m_yx -- nameIn m_yx (alloc td_env)+       in if not m_b_aliased_m_yx && isInScope td_env m_yx -- nameIn m_yx (alloc td_env)       -- Finally update the @activeCoals@ table with a fresh       --   binding for @m_b@; if such one exists then overwrite.       -- Also, add all variables from the alias chain of @b@ to@@ -1562,7 +1557,10 @@   where     updateIndFunSlice :: IxFun -> FlatSlice SubExp -> IxFun     updateIndFunSlice ind_fun (FlatSlice offset dims) =-      IxFun.flatSlice ind_fun $ FlatSlice (pe64 offset) $ map (fmap pe64) dims+      fromMaybe (error "updateIndFunSlice") $+        IxFun.flatSlice ind_fun $+          FlatSlice (pe64 offset) $+            map (fmap pe64) dims  -- CASE b) @let x = concat(a, b^{lu})@ genCoalStmtInfo lutab _ scopetab (Let pat aux (BasicOp (Concat concat_dim (b0 :| bs) _)))

@@ -134,12 +134,12 @@         Constant _ -> Just (ws, ws)         Var x -> Just (ws, ws ++ mapMaybe (getDirAliasedIxfn td_env coal_tab) [x]) getUseSumFromStm td_env coal_tab (Let (Pat [x]) _ (BasicOp (FlatUpdate _ (FlatSlice offset slc) v)))-  | Just (m_b, m_x, x_ixfn) <- getDirAliasedIxfn td_env coal_tab (patElemName x) =-      let x_ixfn_slc = IxFun.flatSlice x_ixfn $ FlatSlice (pe64 offset) $ map (fmap pe64) slc-          r1 = (m_b, m_x, x_ixfn_slc)-       in case getDirAliasedIxfn td_env coal_tab v of-            Nothing -> Just ([r1], [r1])-            Just r2 -> Just ([r1], [r1, r2])+  | Just (m_b, m_x, x_ixfn) <- getDirAliasedIxfn td_env coal_tab (patElemName x) = do+      x_ixfn_slc <- IxFun.flatSlice x_ixfn $ FlatSlice (pe64 offset) $ map (fmap pe64) slc+      let r1 = (m_b, m_x, x_ixfn_slc)+      case getDirAliasedIxfn td_env coal_tab v of+        Nothing -> Just ([r1], [r1])+        Just r2 -> Just ([r1], [r1, r2]) -- getUseSumFromStm td_env coal_tab (Let (Pat ys) _ (BasicOp bop)) = --   let wrt = mapMaybe (getDirAliasedIxfn td_env coal_tab . patElemName) ys --    in trace ("getUseBla: " <> show bop) $ pure (wrt, wrt)@@ -251,7 +251,7 @@         <> fromMaybe mempty (M.lookup m (m_alias td_env))     mbLmad indfun       | Just subs <- freeVarSubstitutions (scope td_env) (scals bu_env) indfun,-        (IxFun.IxFun (lmad :| []) _ _) <- IxFun.substituteInIxFun subs indfun =+        (IxFun.IxFun lmad _) <- IxFun.substituteInIxFun subs indfun =           Just lmad     mbLmad _ = Nothing     addLmads wrts uses etry =@@ -458,7 +458,7 @@           new_offset = replaceIteratorWith lower_bound offset0           new_lmad =             IxFun.LMAD new_offset $-              IxFun.LMADDim new_stride spn 0 IxFun.Inc : map incPerm dims0+              IxFun.LMADDim new_stride spn 0 : map incPerm dims0       if new_var `nameIn` freeIn new_lmad         then pure Undeterminable         else pure $ Set $ S.singleton new_lmad

@@ -27,7 +27,7 @@ import Futhark.IR.Mem.IxFun qualified as IxFun import Futhark.Optimise.ArrayShortCircuiting.DataStructs -type DirAlias = IxFun -> IxFun+type DirAlias = IxFun -> Maybe IxFun -- ^ A direct aliasing transformation  type InvAlias = Maybe (IxFun -> IxFun)@@ -69,23 +69,23 @@  -- | Get alias and (direct) index function mapping from expression getDirAliasFromExp :: Exp (Aliases rep) -> Maybe (VName, DirAlias)-getDirAliasFromExp (BasicOp (SubExp (Var x))) = Just (x, id)-getDirAliasFromExp (BasicOp (Opaque _ (Var x))) = Just (x, id)+getDirAliasFromExp (BasicOp (SubExp (Var x))) = Just (x, Just)+getDirAliasFromExp (BasicOp (Opaque _ (Var x))) = Just (x, Just) getDirAliasFromExp (BasicOp (Reshape ReshapeCoerce shp x)) =-  Just (x, (`IxFun.coerce` shapeDims (fmap pe64 shp)))+  Just (x, Just . (`IxFun.coerce` shapeDims (fmap pe64 shp))) getDirAliasFromExp (BasicOp (Reshape ReshapeArbitrary shp x)) =   Just (x, (`IxFun.reshape` shapeDims (fmap pe64 shp))) getDirAliasFromExp (BasicOp (Rearrange _ _)) =   Nothing getDirAliasFromExp (BasicOp (Index x slc)) =-  Just (x, (`IxFun.slice` (Slice $ map (fmap pe64) $ unSlice slc)))-getDirAliasFromExp (BasicOp (Update _ x _ _elm)) = Just (x, id)+  Just (x, Just . (`IxFun.slice` (Slice $ map (fmap pe64) $ unSlice slc)))+getDirAliasFromExp (BasicOp (Update _ x _ _elm)) = Just (x, Just) getDirAliasFromExp (BasicOp (FlatIndex x (FlatSlice offset idxs))) =   Just     ( x,       (`IxFun.flatSlice` FlatSlice (pe64 offset) (map (fmap pe64) idxs))     )-getDirAliasFromExp (BasicOp (FlatUpdate x _ _)) = Just (x, id)+getDirAliasFromExp (BasicOp (FlatUpdate x _ _)) = Just (x, Just) getDirAliasFromExp _ = Nothing  -- | This was former @createsAliasedArrOK@ from DataStructs@@ -268,7 +268,8 @@ walkAliasTab alias_tab vtab x   | Just (x0, alias0, _) <- M.lookup x alias_tab = do       Coalesced knd (MemBlock pt shp vname ixf) substs <- walkAliasTab alias_tab vtab x0-      pure $ Coalesced knd (MemBlock pt shp vname $ alias0 ixf) substs+      ixf' <- alias0 ixf+      pure $ Coalesced knd (MemBlock pt shp vname ixf') substs walkAliasTab _ _ _ = Nothing  -- | We assume @x@ is in @vartab@ and we add the variables that @x@ aliases

@@ -20,12 +20,12 @@  import Control.Monad import Data.List qualified as L-import Data.List.NonEmpty qualified as NE import Data.Map.Strict qualified as M import Data.Maybe import Data.Sequence qualified as Seq import Futhark.IR.GPU import Futhark.IR.Mem.IxFun qualified as IxFun+import Futhark.IR.Mem.LMAD qualified as LMAD import Futhark.MonadFreshNames import Futhark.Optimise.TileLoops.Shared import Futhark.Tools@@ -143,12 +143,12 @@         | [slc_X'] <- patNames pat,           slc_X == slc_X',           Just ixf_fn <- M.lookup x ixfn_env,-          (IxFun.IxFun lmads _ _) <- ixf_fn =-            all innerHasStride1 $ NE.toList lmads+          (IxFun.IxFun lmad _) <- ixf_fn =+            innerHasStride1 lmad       isInnerCoal _ _ _ =         error "kkLoopBody.isInnerCoal: not an error, but I would like to know why!"       innerHasStride1 lmad =-        let lmad_dims = IxFun.lmadDims lmad+        let lmad_dims = LMAD.dims lmad             q = length lmad_dims             last_perm = IxFun.ldPerm $ last lmad_dims             stride = IxFun.ldStride $ last lmad_dims

@@ -68,7 +68,7 @@     op1 == op2,     Just res <- doBinOp op1 x1 x2 =       Just (BinOp op1 (Constant res) y2, cs)-simplifyBinOp look _ (BinOp Add {} e1 e2)+simplifyBinOp look _ (BinOp (Add it ovf) e1 e2)   | isCt0 e1 = resIsSubExp e2   | isCt0 e2 = resIsSubExp e1   -- x+(y-x) => y@@ -76,13 +76,20 @@     Just (BasicOp (BinOp Sub {} e2_a e2_b), cs) <- look v2,     e2_b == e1 =       Just (SubExp e2_a, cs)+  -- x+(-1*y) => x-y+  | Var v2 <- e2,+    Just (BasicOp (BinOp Mul {} (Constant (IntValue x)) e3), cs) <- look v2,+    valueIntegral x == (-1 :: Int) =+      Just (BinOp (Sub it ovf) e1 e3, cs) simplifyBinOp _ _ (BinOp FAdd {} e1 e2)   | isCt0 e1 = resIsSubExp e2   | isCt0 e2 = resIsSubExp e1 simplifyBinOp look _ (BinOp sub@(Sub t _) e1 e2)   | isCt0 e2 = resIsSubExp e1-  -- Cases for simplifying (a+b)-b and permutations.-+  | e1 == e2 = Just (SubExp (intConst t 0), mempty)+  --+  -- Below are cases for simplifying (a+b)-b and permutations.+  --   -- (e1_a+e1_b)-e1_a == e1_b   | Var v1 <- e1,     Just (BasicOp (BinOp Add {} e1_a e1_b), cs) <- look v1,

@@ -39,11 +39,7 @@ index :: MonadBuilder m => String -> VName -> [VName] -> m VName index se_desc arr outer_indices = do   arr_t <- lookupType arr-  let shape = arrayShape arr_t-      inner_dims = shapeDims $ stripDims (length outer_indices) shape-      untouched d = DimSlice (intConst Int64 0) d (intConst Int64 1)-      inner_slices = map untouched inner_dims-      slice = Slice $ map (DimFix . Var) outer_indices ++ inner_slices+  let slice = fullSlice arr_t $ map (DimFix . Var) outer_indices   letExp se_desc $ BasicOp $ Index arr slice  update :: MonadBuilder m => String -> VName -> [VName] -> SubExp -> m VName@@ -295,23 +291,23 @@        in (lam_op, ne) changeWithEnv with_env _ = pure with_env -composeIxfuns :: IxFnEnv -> VName -> VName -> (IxFun -> IxFun) -> TileM IxFnEnv+composeIxfuns :: IxFnEnv -> VName -> VName -> (IxFun -> Maybe IxFun) -> TileM IxFnEnv composeIxfuns env y x ixf_fun =-  case M.lookup x env of-    Just ixf -> pure $ M.insert y (ixf_fun ixf) env+  case ixf_fun =<< M.lookup x env of+    Just ixf -> pure $ M.insert y ixf env     Nothing -> do       tp <- lookupType x-      case tp of-        Array _ptp shp _u -> do-          let shp' = map ExpMem.pe64 (shapeDims shp)-          pure $ M.insert y (ixf_fun $ IxFun.iota shp') env-        _ -> pure env+      pure $ case tp of+        Array _ptp shp _u+          | Just ixf <- ixf_fun $ IxFun.iota $ map ExpMem.pe64 (shapeDims shp) ->+              M.insert y ixf env+        _ -> env  changeIxFnEnv :: IxFnEnv -> VName -> Exp GPU -> TileM IxFnEnv changeIxFnEnv env y (BasicOp (Reshape ReshapeArbitrary shp_chg x)) =   composeIxfuns env y x (`IxFun.reshape` fmap ExpMem.pe64 (shapeDims shp_chg)) changeIxFnEnv env y (BasicOp (Reshape ReshapeCoerce shp_chg x)) =-  composeIxfuns env y x (`IxFun.coerce` fmap ExpMem.pe64 (shapeDims shp_chg))+  composeIxfuns env y x (Just . (`IxFun.coerce` fmap ExpMem.pe64 (shapeDims shp_chg))) changeIxFnEnv env y (BasicOp (Manifest perm x)) = do   tp <- lookupType x   case tp of@@ -321,9 +317,9 @@       pure $ M.insert y ixfn env     _ -> error "In TileLoops/Shared.hs, changeIxFnEnv: manifest applied to a non-array!" changeIxFnEnv env y (BasicOp (Rearrange perm x)) =-  composeIxfuns env y x (`IxFun.permute` perm)+  composeIxfuns env y x (Just . (`IxFun.permute` perm)) changeIxFnEnv env y (BasicOp (Index x slc)) =-  composeIxfuns env y x (`IxFun.slice` (Slice $ map (fmap ExpMem.pe64) $ unSlice slc))+  composeIxfuns env y x (Just . (`IxFun.slice` Slice (map (fmap ExpMem.pe64) $ unSlice slc))) changeIxFnEnv env y (BasicOp (Opaque _ (Var x))) =-  composeIxfuns env y x id+  composeIxfuns env y x Just changeIxFnEnv env _ _ = pure env

@@ -30,7 +30,6 @@ import Futhark.Transform.Rename (renameStm) import Futhark.Transform.Substitute import Futhark.Util (mapAccumLM)-import Futhark.Util.IntegralExp import Prelude hiding (quot)  -- | The memory expansion pass definition.@@ -538,28 +537,26 @@    pure (slice_stms' <> stmsFromList alloc_stms, mconcat rebases)   where-    expand (mem, (offset, total_size, space)) = do+    expand (mem, (_offset, total_size, space)) = do       let allocpat = Pat [PatElem mem $ MemMem space]       pure         ( Let allocpat (defAux ()) $ Op $ Alloc total_size space,-          M.singleton mem $ newBase offset+          M.singleton mem newBase         )      num_threads' = pe64 num_threads     gtid = le64 $ segFlat kspace +    untouched d = DimSlice 0 d 1+     -- For the variant allocations, we add an inner dimension,     -- which is then offset by a thread-specific amount.-    newBase size_per_thread (old_shape, pt) =-      let elems_per_thread =-            pe64 size_per_thread `quot` primByteSize pt-          root_ixfun = IxFun.iota [elems_per_thread, num_threads']+    newBase (old_shape, _pt) =+      let root_ixfun = IxFun.iota $ old_shape ++ [num_threads']           offset_ixfun =             IxFun.slice root_ixfun . Slice $-              [DimSlice 0 num_threads' 1, DimFix gtid]-       in if length old_shape == 1-            then IxFun.coerce offset_ixfun old_shape-            else IxFun.reshape offset_ixfun old_shape+              map untouched old_shape ++ [DimFix gtid]+       in offset_ixfun  -- | A map from memory block names to new index function bases. type RebaseMap = M.Map VName (([TPrimExp Int64 VName], PrimType) -> IxFun)@@ -622,11 +619,16 @@ offsetMemoryInStm :: Stm GPUMem -> OffsetM (Scope GPUMem, Stm GPUMem) offsetMemoryInStm (Let pat dec e) = do   e' <- offsetMemoryInExp e-  pat' <- offsetMemoryInPat pat =<< expReturns e'+  pat' <-+    offsetMemoryInPat pat+      =<< maybe (throwError "offsetMemoryInStm: ill-typed") pure+      =<< expReturns e'   scope <- askScope   -- Try to recompute the index function.  Fall back to creating rebase   -- operations with the RebaseMap.-  rts <- runReaderT (expReturns e') scope+  rts <-+    maybe (throwError "offsetMemoryInStm: ill-typed") pure $+      runReader (expReturns e') scope   let pat'' = Pat $ zipWith pick (patElems pat') rts       stm = Let pat'' dec e'   let scope' = scopeOf stm <> scope@@ -659,32 +661,48 @@         pure . PatElem name . MemArray pt shape u . ArrayIn mem $           fmap (fmap unExt) ixfun     onPE pe _ = do-      new_dec <- offsetMemoryInMemBound $ patElemDec pe+      new_dec <- offsetMemoryInMemBound (patElemName pe) $ patElemDec pe       pure pe {patElemDec = new_dec}     unExt (Ext i) = patElemName (pes !! i)     unExt (Free v) = v  offsetMemoryInParam :: Param (MemBound u) -> OffsetM (Param (MemBound u)) offsetMemoryInParam fparam = do-  fparam' <- offsetMemoryInMemBound $ paramDec fparam+  fparam' <- offsetMemoryInMemBound (paramName fparam) $ paramDec fparam   pure fparam {paramDec = fparam'} -offsetMemoryInMemBound :: MemBound u -> OffsetM (MemBound u)-offsetMemoryInMemBound summary@(MemArray pt shape u (ArrayIn mem ixfun)) = do+offsetMemoryInMemBound :: VName -> MemBound u -> OffsetM (MemBound u)+offsetMemoryInMemBound v summary@(MemArray pt shape u (ArrayIn mem ixfun)) = do   new_base <- lookupNewBase mem (IxFun.base ixfun, pt)-  pure . fromMaybe summary $ do-    new_base' <- new_base-    pure $ MemArray pt shape u $ ArrayIn mem $ IxFun.rebase new_base' ixfun-offsetMemoryInMemBound summary = pure summary+  case new_base of+    Nothing -> pure summary+    Just new_base' -> do+      let problem =+            throwError . unlines $+              [ "offsetMemoryInMemBound",+                prettyString v,+                prettyString new_base',+                prettyString ixfun+              ]+      ixfun' <- maybe problem pure $ IxFun.rebase new_base' ixfun+      pure $ MemArray pt shape u $ ArrayIn mem ixfun'+offsetMemoryInMemBound _ summary = pure summary  offsetMemoryInBodyReturns :: BodyReturns -> OffsetM BodyReturns offsetMemoryInBodyReturns br@(MemArray pt shape u (ReturnsInBlock mem ixfun))   | Just ixfun' <- isStaticIxFun ixfun = do       new_base <- lookupNewBase mem (IxFun.base ixfun', pt)-      pure . fromMaybe br $ do-        new_base' <- new_base-        pure . MemArray pt shape u . ReturnsInBlock mem $-          IxFun.rebase (fmap (fmap Free) new_base') ixfun+      case new_base of+        Nothing -> pure br+        Just new_base' -> do+          let problem =+                throwError . unlines $+                  [ "offsetMemoryInBodyReturns",+                    prettyString new_base',+                    prettyString ixfun+                  ]+          ixfun'' <- maybe problem pure $ IxFun.rebase (fmap (fmap Free) new_base') ixfun+          pure $ MemArray pt shape u $ ReturnsInBlock mem ixfun'' offsetMemoryInBodyReturns br = pure br  offsetMemoryInLambda :: Lambda GPUMem -> OffsetM (Lambda GPUMem)

@@ -43,7 +43,6 @@ import Data.Either (partitionEithers) import Data.Foldable (toList) import Data.List (foldl', transpose, zip4)-import Data.List.NonEmpty qualified as NE import Data.Map.Strict qualified as M import Data.Maybe import Data.Set qualified as S@@ -186,6 +185,40 @@ allocForArray t space = do   allocForArray' t space +-- | Repair an expression that cannot be assigned an index function.+-- There is a simple remedy for this: normalise the input arrays and+-- try again.+repairExpression ::+  (Allocable fromrep torep inner) =>+  Exp torep ->+  AllocM fromrep torep (Exp torep)+repairExpression (BasicOp (Reshape k shape v)) = do+  v_mem <- fst <$> lookupArraySummary v+  space <- lookupMemSpace v_mem+  v' <- snd <$> ensureDirectArray (Just space) v+  pure $ BasicOp $ Reshape k shape v'+repairExpression e =+  error $ "repairExpression:\n" <> prettyString e++expReturns' ::+  (Allocable fromrep torep inner) =>+  Exp torep ->+  AllocM fromrep torep ([ExpReturns], Exp torep)+expReturns' e = do+  maybe_rts <- expReturns e+  case maybe_rts of+    Just rts -> pure (rts, e)+    Nothing -> do+      e' <- repairExpression e+      let bad =+            error . unlines $+              [ "expReturns': impossible index transformation",+                prettyString e,+                prettyString e'+              ]+      rts <- fromMaybe bad <$> expReturns e'+      pure (rts, e')+ allocsForStm ::   (Allocable fromrep torep inner) =>   [Ident] ->@@ -194,10 +227,10 @@ allocsForStm idents e = do   def_space <- askDefaultSpace   hints <- expHints e-  rts <- expReturns e+  (rts, e') <- expReturns' e   pes <- allocsForPat def_space idents rts hints-  dec <- mkExpDecM (Pat pes) e-  pure $ Let (Pat pes) (defAux dec) e+  dec <- mkExpDecM (Pat pes) e'+  pure $ Let (Pat pes) (defAux dec) e'  patWithAllocations ::   (MonadBuilder m, Mem (Rep m) inner) =>@@ -209,7 +242,7 @@ patWithAllocations def_space names e hints = do   ts' <- instantiateShapes' names <$> expExtType e   let idents = zipWith Ident names ts'-  rts <- expReturns e+  rts <- fromMaybe (error "patWithAllocations: ill-typed") <$> expReturns e   Pat <$> allocsForPat def_space idents rts hints  mkMissingIdents :: MonadFreshNames m => [Ident] -> [ExpReturns] -> m [Ident]@@ -344,11 +377,9 @@   mem_space <- lookupMemSpace mem   default_space <- askDefaultSpace   let space = fromMaybe default_space space_ok-  if numLMADs ixfun == 1-    && ixFunPerm ixfun == [0 .. IxFun.rank ixfun - 1]+  if IxFun.permutation ixfun == [0 .. IxFun.rank ixfun - 1]     && length (IxFun.base ixfun) == IxFun.rank ixfun     && maybe True (== mem_space) space_ok-    && IxFun.contiguous ixfun     then pure (mem, v)     else allocLinearArray space (baseString v) v @@ -527,6 +558,22 @@     _ ->       error $ "allocPermArray: " ++ prettyString t +ensurePermArray ::+  (Allocable fromrep torep inner) =>+  Maybe Space ->+  [Int] ->+  VName ->+  AllocM fromrep torep (VName, VName)+ensurePermArray space_ok perm v = do+  (mem, ixfun) <- lookupArraySummary v+  mem_space <- lookupMemSpace mem+  default_space <- askDefaultSpace+  if length (IxFun.base ixfun) == length (IxFun.shape ixfun)+    && IxFun.permutation ixfun == perm+    && maybe True (== mem_space) space_ok+    then pure (mem, v)+    else allocPermArray (fromMaybe default_space space_ok) perm (baseString v) v+ allocLinearArray ::   (Allocable fromrep torep inner) =>   Space ->@@ -710,25 +757,16 @@ allocInLambda params body =   mkLambda params . allocInStms (bodyStms body) $ pure $ bodyResult body -numLMADs :: IxFun -> Int-numLMADs = length . IxFun.ixfunLMADs--ixFunPerm :: IxFun -> [Int]-ixFunPerm = map IxFun.ldPerm . IxFun.lmadDims . NE.head . IxFun.ixfunLMADs--ixFunMon :: IxFun -> [IxFun.Monotonicity]-ixFunMon = map IxFun.ldMon . IxFun.lmadDims . NE.head . IxFun.ixfunLMADs- data MemReq-  = MemReq Space [Int] [IxFun.Monotonicity] Rank Bool-  | NeedsLinearisation Space+  = MemReq Space [Int] Rank+  | NeedsNormalisation Space   deriving (Eq, Show)  combMemReqs :: MemReq -> MemReq -> MemReq-combMemReqs x@NeedsLinearisation {} _ = x-combMemReqs _ y@NeedsLinearisation {} = y-combMemReqs x@(MemReq x_space _ _ _ _) y@MemReq {} =-  if x == y then x else NeedsLinearisation x_space+combMemReqs x@NeedsNormalisation {} _ = x+combMemReqs _ y@NeedsNormalisation {} = y+combMemReqs x@(MemReq x_space _ _) y@MemReq {} =+  if x == y then x else NeedsNormalisation x_space  type MemReqType = MemInfo (Ext SubExp) NoUniqueness MemReq @@ -738,13 +776,13 @@ combMemReqTypes x _ = x  contextRets :: MemReqType -> [MemInfo d u r]-contextRets (MemArray _ shape _ (MemReq space _ _ (Rank base_rank) _)) =+contextRets (MemArray _ shape _ (MemReq space _ (Rank base_rank))) =   -- Memory + offset + base_rank + (stride,size)*rank.   MemMem space     : MemPrim int64     : replicate base_rank (MemPrim int64)     ++ replicate (2 * shapeRank shape) (MemPrim int64)-contextRets (MemArray _ shape _ (NeedsLinearisation space)) =+contextRets (MemArray _ shape _ (NeedsNormalisation space)) =   -- Memory + offset + (base,stride,size)*rank.   MemMem space     : MemPrim int64@@ -771,15 +809,10 @@         (Array pt shape u, MemArray _ _ _ (ArrayIn mem ixfun)) -> do           space <- lookupMemSpace mem           pure . MemArray pt shape u $-            if numLMADs ixfun == 1-              then-                MemReq-                  space-                  (ixFunPerm ixfun)-                  (ixFunMon ixfun)-                  (Rank $ length $ IxFun.base ixfun)-                  (IxFun.contiguous ixfun)-              else NeedsLinearisation space+            MemReq+              space+              (IxFun.permutation ixfun)+              (Rank $ length $ IxFun.base ixfun)         (_, MemMem space) -> pure $ MemMem space         (_, MemPrim pt) -> pure $ MemPrim pt         (_, MemAcc acc ispace ts u) -> pure $ MemAcc acc ispace ts u@@ -800,17 +833,16 @@         res_rets_acc ++ [inspect ctx_offset req]       ) -    arrayInfo rank (NeedsLinearisation space) =-      (space, [0 .. rank - 1], repeat IxFun.Inc, rank, True)-    arrayInfo _ (MemReq space perm mon (Rank base_rank) contig) =-      (space, perm, mon, base_rank, contig)+    arrayInfo rank (NeedsNormalisation space) =+      (space, [0 .. rank - 1], rank)+    arrayInfo _ (MemReq space perm (Rank base_rank)) =+      (space, perm, base_rank)      inspect ctx_offset (MemArray pt shape u req) =       let shape' = fmap (adjustExt num_new_ctx) shape-          (space, perm, mon, base_rank, contig) = arrayInfo (shapeRank shape) req+          (space, perm, base_rank) = arrayInfo (shapeRank shape) req        in MemArray pt shape' u . ReturnsNewBlock space ctx_offset $-            convert-              <$> IxFun.mkExistential base_rank (zip perm mon) contig (ctx_offset + 1)+            convert <$> IxFun.mkExistential base_rank perm (ctx_offset + 1)     inspect _ (MemAcc acc ispace ts u) = MemAcc acc ispace ts u     inspect _ (MemPrim pt) = MemPrim pt     inspect _ (MemMem space) = MemMem space@@ -828,13 +860,14 @@   Body torep ->   AllocM fromrep torep (Body torep) addCtxToMatchBody reqs body = buildBody_ $ do-  res <- zipWithM linearIfNeeded reqs =<< bodyBind body+  res <- zipWithM normaliseIfNeeded reqs =<< bodyBind body   ctx <- concat <$> mapM resCtx res   pure $ ctx ++ res   where-    linearIfNeeded (MemArray _ _ _ (NeedsLinearisation space)) (SubExpRes cs (Var v)) =-      SubExpRes cs . Var . snd <$> ensureRowMajorArray (Just space) v-    linearIfNeeded _ res =+    normaliseIfNeeded (MemArray _ shape _ (NeedsNormalisation space)) (SubExpRes cs (Var v)) =+      SubExpRes cs . Var . snd+        <$> ensurePermArray (Just space) [0 .. shapeRank shape - 1] v+    normaliseIfNeeded _ res =       pure res      resCtx (SubExpRes _ Constant {}) =@@ -1136,7 +1169,7 @@       innerUsage inner      simplifyPat (Pat pes) e = do-      rets <- expReturns e+      rets <- fromMaybe (error "simplifyPat: ill-typed") <$> expReturns e       Pat <$> zipWithM update pes rets       where         names = map patElemName pes

@@ -1,12 +1,12 @@ -- | Optimisation pipelines. module Futhark.Passes   ( standardPipeline,-    sequentialPipeline,-    kernelsPipeline,-    sequentialCpuPipeline,+    seqPipeline,     gpuPipeline,+    seqmemPipeline,+    gpumemPipeline,     mcPipeline,-    multicorePipeline,+    mcmemPipeline,   ) where @@ -82,8 +82,8 @@  -- | The pipeline used by the CUDA and OpenCL backends, but before -- adding memory information.  Includes 'standardPipeline'.-kernelsPipeline :: Pipeline SOACS GPU-kernelsPipeline =+gpuPipeline :: Pipeline SOACS GPU+gpuPipeline =   standardPipeline     >>> onePass extractKernels     >>> passes@@ -110,8 +110,8 @@  -- | The pipeline used by the sequential backends.  Turns all -- parallelism into sequential loops.  Includes 'standardPipeline'.-sequentialPipeline :: Pipeline SOACS Seq-sequentialPipeline =+seqPipeline :: Pipeline SOACS Seq+seqPipeline =   standardPipeline     >>> onePass firstOrderTransform     >>> passes@@ -119,11 +119,11 @@         inPlaceLoweringSeq       ] --- | Run 'sequentialPipeline', then add memory information (and+-- | Run 'seqPipeline', then add memory information (and -- optimise it slightly).-sequentialCpuPipeline :: Pipeline SOACS SeqMem-sequentialCpuPipeline =-  sequentialPipeline+seqmemPipeline :: Pipeline SOACS SeqMem+seqmemPipeline =+  seqPipeline     >>> onePass Seq.explicitAllocations     >>> passes       [ performCSE False,@@ -140,11 +140,11 @@         simplifySeqMem       ] --- | Run 'kernelsPipeline', then add memory information (and optimise+-- | Run 'gpuPipeline', then add memory information (and optimise -- it a lot).-gpuPipeline :: Pipeline SOACS GPUMem-gpuPipeline =-  kernelsPipeline+gpumemPipeline :: Pipeline SOACS GPUMem+gpumemPipeline =+  gpuPipeline     >>> onePass GPU.explicitAllocations     >>> passes       [ simplifyGPUMem,@@ -185,8 +185,8 @@       ]  -- | Run 'mcPipeline' and then add memory information.-multicorePipeline :: Pipeline SOACS MCMem-multicorePipeline =+mcmemPipeline :: Pipeline SOACS MCMem+mcmemPipeline =   mcPipeline     >>> onePass MC.explicitAllocations     >>> passes

@@ -46,7 +46,7 @@ import Data.Text qualified as T import Numeric.Half import Prettyprinter-import Prettyprinter.Render.Terminal (AnsiStyle, Color (..), bgColor, bgColorDull, bold, color, colorDull)+import Prettyprinter.Render.Terminal (AnsiStyle, Color (..), bgColor, bgColorDull, bold, color, colorDull, italicized, underlined) import Prettyprinter.Render.Terminal qualified import Prettyprinter.Render.Text qualified import Prettyprinter.Symbols.Ascii

@@ -5,7 +5,7 @@     parseExp,     parseModExp,     parseType,-    parseDecOrExpIncrM,+    parseDecOrExp,     SyntaxError (..),     Comment (..),   )@@ -57,29 +57,13 @@   Either SyntaxError UncheckedTypeExp parseType = parse futharkType --- | Parse an Futhark expression incrementally from monadic actions, using the--- 'FilePath' as the source name for error messages.-parseExpIncrM ::-  Monad m =>-  m T.Text ->-  FilePath ->-  T.Text ->-  m (Either SyntaxError UncheckedExp)-parseExpIncrM fetch file program =-  getLinesFromM fetch $ fmap fst <$> parseInMonad expression file program---- | Parse either an expression or a declaration incrementally;--- favouring declarations in case of ambiguity.-parseDecOrExpIncrM ::-  Monad m =>-  m T.Text ->+-- | Parse either an expression or a declaration; favouring+-- declarations in case of ambiguity.+parseDecOrExp ::   FilePath ->   T.Text ->-  m (Either SyntaxError (Either UncheckedDec UncheckedExp))-parseDecOrExpIncrM fetch file input =-  case parseInMonad declaration file input of-    Value Left {} -> fmap Right <$> parseExpIncrM fetch file input-    Value (Right (d, _)) -> pure $ Right $ Left d-    GetLine _ -> do-      l <- fetch-      parseDecOrExpIncrM fetch file $ input <> "\n" <> l+  Either SyntaxError (Either UncheckedDec UncheckedExp)+parseDecOrExp file input =+  case parse declaration file input of+    Left {} -> Right <$> parseExp file input+    Right d -> Right $ Left d

@@ -3,7 +3,7 @@ -- | The Futhark lexer.  Takes a string, produces a list of tokens with position information. module Language.Futhark.Parser.Lexer   ( Token(..)-  , scanTokens+  , getToken   , scanTokensText   ) where @@ -12,7 +12,7 @@ import qualified Data.Text as T import qualified Data.Text.Encoding as T import qualified Data.Text.Read as T-import Data.Char (ord, toLower, digitToInt)+import Data.Char (chr, ord, toLower) import Data.Int (Int8, Int16, Int32, Int64) import Data.Word (Word8) import Data.Loc (Loc (..), L(..), Pos(..))@@ -20,7 +20,7 @@  import Language.Futhark.Core (Int8, Int16, Int32, Int64,                               Word8, Word16, Word32, Word64,-                              Name, nameFromText, nameToText)+                              Name, nameFromText, nameToText, nameFromString) import Language.Futhark.Prop (leadingOperator) import Language.Futhark.Syntax (BinOp(..)) import Language.Futhark.Parser.Lexer.Wrapper@@ -34,7 +34,6 @@ @declit = [0-9][0-9_]* @binlit = 0[bB][01][01_]* @romlit = 0[rR][IVXLCDM][IVXLCDM_]*-@intlit = @hexlit|@binlit|@declit|@romlit @reallit = (([0-9][0-9_]*("."[0-9][0-9_]*)?))([eE][\+\-]?[0-9]+)? @hexreallit = 0[xX][0-9a-fA-F][0-9a-fA-F_]*"."[0-9a-fA-F][0-9a-fA-F_]*([pP][\+\-]?[0-9_]+) @@ -54,7 +53,7 @@ tokens :-    $white+                               ;-  @doc                     { tokenM $ pure . DOC . T.intercalate "\n" .+  @doc                     { tokenS $ DOC . T.intercalate "\n" .                                       map (T.drop 3 . T.stripStart) .                                            T.split (== '\n') . ("--"<>) .                                            T.drop 4 }@@ -91,62 +90,123 @@   "$"                      { tokenC DOLLAR }   "???"                    { tokenC HOLE } -  @declit                  { tokenS $ \s -> NATLIT (nameFromText s) (readIntegral s) }-  @intlit i8               { tokenM $ pure . I8LIT . readIntegral . T.filter (/= '_') . T.takeWhile (/='i') }-  @intlit i16              { tokenM $ pure . I16LIT . readIntegral . T.filter (/= '_') . T.takeWhile (/='i') }-  @intlit i32              { tokenM $ pure . I32LIT . readIntegral . T.filter (/= '_') . T.takeWhile (/='i') }-  @intlit i64              { tokenM $ pure . I64LIT . readIntegral . T.filter (/= '_') . T.takeWhile (/='i') }-  @intlit u8               { tokenM $ pure . U8LIT . readIntegral . T.filter (/= '_') . T.takeWhile (/='u') }-  @intlit u16              { tokenM $ pure . U16LIT . readIntegral . T.filter (/= '_') . T.takeWhile (/='u') }-  @intlit u32              { tokenM $ pure . U32LIT . readIntegral . T.filter (/= '_') . T.takeWhile (/='u') }-  @intlit u64              { tokenM $ pure . U64LIT . readIntegral . T.filter (/= '_') . T.takeWhile (/='u') }-  @intlit                  { tokenM $ pure . INTLIT . readIntegral . T.filter (/= '_') }+  @declit                  { \s -> decToken (NATLIT (nameFromBS s)) s } -  [\n[^\.]] ^ @reallit f16     { tokenM $ fmap F16LIT . tryRead "f16" . suffZero . T.filter (/= '_') . T.takeWhile (/='f') }-  [\n[^\.]] ^ @reallit f32     { tokenM $ fmap F32LIT . tryRead "f32" . suffZero . T.filter (/= '_') . T.takeWhile (/='f') }-  [\n[^\.]] ^ @reallit f64     { tokenM $ fmap F64LIT . tryRead "f64" . suffZero . T.filter (/= '_') . T.takeWhile (/='f') }-  [\n[^\.]] ^ @reallit         { tokenM $ fmap FLOATLIT . tryRead "f64" . suffZero . T.filter (/= '_') }-  @hexreallit f16          { tokenM $ fmap F16LIT . readHexRealLit . T.filter (/= '_') . T.dropEnd 3 }-  @hexreallit f32          { tokenM $ fmap F32LIT . readHexRealLit . T.filter (/= '_') . T.dropEnd 3 }-  @hexreallit f64          { tokenM $ fmap F64LIT . readHexRealLit . T.filter (/= '_') . T.dropEnd 3 }-  @hexreallit              { tokenM $ fmap FLOATLIT . readHexRealLit . T.filter (/= '_') }-  "'" @charlit "'"         { tokenM $ fmap CHARLIT . tryRead "char" }-  \" @stringcharlit* \"    { tokenM $ fmap (STRINGLIT . T.pack) . tryRead "string"  }+  @declit i8               { decToken I8LIT . BS.dropEnd 2 }+  @binlit i8               { binToken I8LIT . BS.drop 2 . BS.dropEnd 2 }+  @hexlit i8               { hexToken I8LIT . BS.drop 2 . BS.dropEnd 2 }+  @romlit i8               { romToken I8LIT . BS.drop 2 . BS.dropEnd 2 } -  @identifier              { tokenS keyword }+  @declit i16              { decToken I16LIT . BS.dropEnd 3 }+  @binlit i16              { binToken I16LIT . BS.drop 2 . BS.dropEnd 3 }+  @hexlit i16              { hexToken I16LIT . BS.drop 2 . BS.dropEnd 3 }+  @romlit i16              { romToken I16LIT . BS.drop 2 . BS.dropEnd 3 }++  @declit i32              { decToken I32LIT . BS.dropEnd 3 }+  @binlit i32              { binToken I32LIT . BS.drop 2 . BS.dropEnd 3 }+  @hexlit i32              { hexToken I32LIT . BS.drop 2 . BS.dropEnd 3 }+  @romlit i32              { romToken I32LIT . BS.drop 2 . BS.dropEnd 3 }++  @declit i64              { decToken I64LIT . BS.dropEnd 3 }+  @binlit i64              { binToken I64LIT . BS.drop 2 . BS.dropEnd 3 }+  @hexlit i64              { hexToken I64LIT . BS.drop 2 . BS.dropEnd 3 }+  @romlit i64              { romToken I64LIT . BS.drop 2 . BS.dropEnd 3 }++  @declit u8               { decToken U8LIT . BS.dropEnd 2 }+  @binlit u8               { binToken U8LIT . BS.drop 2 . BS.dropEnd 2 }+  @hexlit u8               { hexToken U8LIT . BS.drop 2 . BS.dropEnd 2 }+  @romlit u8               { romToken U8LIT . BS.drop 2 . BS.dropEnd 2 }++  @declit u16              { decToken U16LIT . BS.dropEnd 3 }+  @binlit u16              { binToken U16LIT . BS.drop 2 . BS.dropEnd 3 }+  @hexlit u16              { hexToken U16LIT . BS.drop 2 . BS.dropEnd 3 }+  @romlit u16              { romToken U16LIT . BS.drop 2 . BS.dropEnd 3 }++  @declit u32              { decToken U32LIT . BS.dropEnd 3 }+  @binlit u32              { binToken U32LIT . BS.drop 2 . BS.dropEnd 3 }+  @hexlit u32              { hexToken U32LIT . BS.drop 2 . BS.dropEnd 3 }+  @romlit u32              { romToken U32LIT . BS.drop 2 . BS.dropEnd 3 }++  @declit u64              { decToken U64LIT . BS.dropEnd 3 }+  @binlit u64              { binToken U64LIT . BS.drop 2 . BS.dropEnd 3 }+  @hexlit u64              { hexToken U64LIT . BS.drop 2 . BS.dropEnd 3 }+  @romlit u64              { romToken U64LIT . BS.drop 2 . BS.dropEnd 3 }++  @declit                  { decToken INTLIT }+  @binlit                  { binToken INTLIT . BS.drop 2 }+  @hexlit                  { hexToken INTLIT . BS.drop 2 }+  @romlit                  { romToken INTLIT . BS.drop 2 }++  [\n[^\.]] ^ @reallit f16 { tokenS $ F16LIT . tryRead "f16" . suffZero . T.filter (/= '_') . T.takeWhile (/='f') }+  [\n[^\.]] ^ @reallit f32 { tokenS $ F32LIT . tryRead "f32" . suffZero . T.filter (/= '_') . T.takeWhile (/='f') }+  [\n[^\.]] ^ @reallit f64 { tokenS $ F64LIT . tryRead "f64" . suffZero . T.filter (/= '_') . T.takeWhile (/='f') }+  [\n[^\.]] ^ @reallit     { tokenS $ FLOATLIT . tryRead "f64" . suffZero . T.filter (/= '_') }+  @hexreallit f16          { tokenS $ F16LIT . readHexRealLit . T.filter (/= '_') . T.dropEnd 3 }+  @hexreallit f32          { tokenS $ F32LIT . readHexRealLit . T.filter (/= '_') . T.dropEnd 3 }+  @hexreallit f64          { tokenS $ F64LIT . readHexRealLit . T.filter (/= '_') . T.dropEnd 3 }+  @hexreallit              { tokenS $ FLOATLIT . readHexRealLit . T.filter (/= '_') }+  "'" @charlit "'"         { tokenS $ CHARLIT . tryRead "char" }+  \" @stringcharlit* \"    { tokenS $ STRINGLIT . T.pack . tryRead "string"  }++  "true"                   { tokenC TRUE }+  "false"                  { tokenC FALSE }+  "if"                     { tokenC IF }+  "then"                   { tokenC THEN }+  "else"                   { tokenC ELSE }+  "def"                    { tokenC DEF }+  "let"                    { tokenC LET }+  "loop"                   { tokenC LOOP }+  "in"                     { tokenC IN }+  "val"                    { tokenC VAL }+  "for"                    { tokenC FOR }+  "do"                     { tokenC DO }+  "with"                   { tokenC WITH }+  "local"                  { tokenC LOCAL }+  "open"                   { tokenC OPEN }+  "include"                { tokenC INCLUDE }+  "import"                 { tokenC IMPORT }+  "type"                   { tokenC TYPE }+  "entry"                  { tokenC ENTRY }+  "module"                 { tokenC MODULE }+  "while"                  { tokenC WHILE }+  "assert"                 { tokenC ASSERT }+  "match"                  { tokenC MATCH }+  "case"                   { tokenC CASE }++  @identifier              { tokenS $ ID . nameFromText }+   "#" @identifier          { tokenS $ CONSTRUCTOR . nameFromText . T.drop 1 } -  @binop                   { tokenM $ pure . symbol [] . nameFromText }-  @qualbinop               { tokenM $ \s -> do (qs,k) <- mkQualId s; pure (symbol qs k) }+  @binop                   { tokenS $ symbol [] . nameFromText }+  @qualbinop               { tokenS $ uncurry symbol . mkQualId } { -getToken :: Alex (Lexeme Token)-getToken = do-  inp@(_,_,_,n) <- alexGetInput-  sc <- alexGetStartCode-  case alexScan inp sc of-    AlexEOF -> do pos <- alexGetPos-                  pure (pos, pos, EOF)+nameFromBS :: BS.ByteString -> Name+nameFromBS = nameFromString . map (chr . fromIntegral) . BS.unpack++getToken :: AlexInput -> Either LexerError (AlexInput, (Pos, Pos, Token))+getToken state@(pos,c,s,n) =+  case alexScan state 0 of+    AlexEOF -> Right (state, (pos, pos, EOF))     AlexError (pos,_,_,_) ->-      alexError (Loc pos pos) "Invalid lexical syntax."-    AlexSkip  inp' _len -> do-      alexSetInput inp'-      getToken-    AlexToken inp'@(_,_,_,n') _ action -> let len = n'-n in do-      alexSetInput inp'-      action inp len+      Left $ LexerError (Loc pos pos) "Invalid lexical syntax."+    AlexSkip state' _len ->+      getToken state'+    AlexToken state'@(pos',_,_,n') _ action -> do+      let x = action (BS.take (n'-n) s)+      x `seq` Right (state', (pos, pos', x))  scanTokens :: Pos -> BS.ByteString -> Either LexerError ([L Token], Pos)-scanTokens pos str =-  runAlex pos str $ do-  fix $ \loop -> do-    tok <- getToken-    case tok of-      (start, end, EOF) ->-        pure ([], end)-      (start, end, t) -> do-        (rest, endpos) <- loop-        pure (L (Loc start end) t : rest, endpos)+scanTokens pos str = loop $ initialLexerState pos str+  where+   loop s = do+     (s', tok) <- getToken s+     case tok of+       (start, end, EOF) ->+         pure ([], end)+       (start, end, t) -> do+         (rest, endpos) <- loop s'+         pure (L (Loc start end) t : rest, endpos)  -- | Given a starting position, produce tokens from the given text (or -- a lexer error).  Returns the final position.

@@ -5,28 +5,27 @@ -- Also defines other useful building blocks for constructing tokens. module Language.Futhark.Parser.Lexer.Tokens   ( Token (..),-    Lexeme,     fromRoman,     symbol,     mkQualId,-    tokenPosM,-    tokenM,     tokenC,-    keyword,     tokenS,-    indexing,     suffZero,     tryRead,-    readIntegral,+    decToken,+    binToken,+    hexToken,+    romToken,+    advance,     readHexRealLit,   ) where  import Data.ByteString.Lazy qualified as BS-import Data.Char (digitToInt, ord)+import Data.Char (ord) import Data.Either-import Data.List (find, foldl')-import Data.Loc (Loc (..), Pos (..))+import Data.List (find)+import Data.Loc (Pos (..)) import Data.Text qualified as T import Data.Text.Encoding qualified as T import Data.Text.Read qualified as T@@ -41,7 +40,6 @@     Word64,     Word8,   )-import Language.Futhark.Parser.Lexer.Wrapper import Language.Futhark.Prop (leadingOperator) import Language.Futhark.Syntax (BinOp, nameFromText, nameToText) import Numeric.Half@@ -133,94 +131,68 @@   | HOLE   deriving (Show, Eq, Ord) -keyword :: T.Text -> Token-keyword s =-  case s of-    "true" -> TRUE-    "false" -> FALSE-    "if" -> IF-    "then" -> THEN-    "else" -> ELSE-    "def" -> DEF-    "let" -> LET-    "loop" -> LOOP-    "in" -> IN-    "val" -> VAL-    "for" -> FOR-    "do" -> DO-    "with" -> WITH-    "local" -> LOCAL-    "open" -> OPEN-    "include" -> INCLUDE-    "import" -> IMPORT-    "type" -> TYPE-    "entry" -> ENTRY-    "module" -> MODULE-    "while" -> WHILE-    "assert" -> ASSERT-    "match" -> MATCH-    "case" -> CASE-    _ -> ID $ nameFromText s--indexing :: (Loc, T.Text) -> Alex Name-indexing (loc, s) = case keyword s of-  ID v -> pure v-  _ -> alexError loc $ "Cannot index keyword '" <> s <> "'."--mkQualId :: T.Text -> Alex ([Name], Name)+mkQualId :: T.Text -> ([Name], Name) mkQualId s = case reverse $ T.splitOn "." s of   [] -> error "mkQualId: no components"-  k : qs -> pure (map nameFromText (reverse qs), nameFromText k)+  k : qs -> (map nameFromText (reverse qs), nameFromText k)  -- | Suffix a zero if the last character is dot. suffZero :: T.Text -> T.Text suffZero s = if T.last s == '.' then s <> "0" else s -tryRead :: Read a => String -> T.Text -> Alex a+tryRead :: Read a => String -> T.Text -> a tryRead desc s = case reads s' of-  [(x, "")] -> pure x+  [(x, "")] -> x   _ -> error $ "Invalid " ++ desc ++ " literal: `" ++ T.unpack s ++ "'."   where     s' = T.unpack s -readIntegral :: Integral a => T.Text -> a-readIntegral s-  | "0x" `T.isPrefixOf` s || "0X" `T.isPrefixOf` s = parseBase 16 (T.drop 2 s)-  | "0b" `T.isPrefixOf` s || "0B" `T.isPrefixOf` s = parseBase 2 (T.drop 2 s)-  | "0r" `T.isPrefixOf` s || "0R" `T.isPrefixOf` s = fromRoman (T.drop 2 s)-  | otherwise = parseBase 10 s-  where-    parseBase base = T.foldl (\acc c -> acc * base + fromIntegral (digitToInt c)) 0+{-# INLINE tokenC #-}+tokenC :: a -> BS.ByteString -> a+tokenC v _ = v -tokenC :: a -> (Pos, Char, BS.ByteString, Int64) -> Int64 -> Alex (Lexeme a)-tokenC v = tokenS $ const v+{-# INLINE decToken #-}+decToken :: Integral a => (a -> Token) -> BS.ByteString -> Token+decToken f = f . BS.foldl' digit 0+  where+    digit x c =+      if c >= 48 && c <= 57+        then x * 10 + fromIntegral (c - 48)+        else x -tokenS :: (T.Text -> a) -> (Pos, Char, BS.ByteString, Int64) -> Int64 -> Alex (Lexeme a)-tokenS f = tokenM $ pure . f+{-# INLINE binToken #-}+binToken :: Integral a => (a -> Token) -> BS.ByteString -> Token+binToken f = f . BS.foldl' digit 0+  where+    digit x c =+      if c >= 48 && c <= 49+        then x * 2 + fromIntegral (c - 48)+        else x -type Lexeme a = (Pos, Pos, a)+{-# INLINE hexToken #-}+hexToken :: Integral a => (a -> Token) -> BS.ByteString -> Token+hexToken f = f . BS.foldl' digit 0+  where+    digit x c+      | c >= 48 && c <= 57 =+          x * 16 + fromIntegral (c - 48)+      | c >= 65 && c <= 70 =+          x * 16 + fromIntegral (10 + c - 65)+      | c >= 97 && c <= 102 =+          x * 16 + fromIntegral (10 + c - 97)+      | otherwise =+          x -tokenM ::-  (T.Text -> Alex a) ->-  (Pos, Char, BS.ByteString, Int64) ->-  Int64 ->-  Alex (Lexeme a)-tokenM f = tokenPosM (f . snd)+{-# INLINE romToken #-}+romToken :: Integral a => (a -> Token) -> BS.ByteString -> Token+romToken f = tokenS $ f . fromRoman -tokenPosM ::-  ((Loc, T.Text) -> Alex a) ->-  (Pos, Char, BS.ByteString, Int64) ->-  Int64 ->-  Alex (Lexeme a)-tokenPosM f (pos, _, s, _) len = do-  x <- f (Loc pos pos', T.decodeUtf8 $ BS.toStrict s')-  pure (pos, pos', x)-  where-    pos' = advance pos s'-    s' = BS.take len s+{-# INLINE tokenS #-}+tokenS :: (T.Text -> a) -> BS.ByteString -> a+tokenS f = f . T.decodeUtf8 . BS.toStrict  advance :: Pos -> BS.ByteString -> Pos-advance orig_pos = foldl' advance' orig_pos . init . BS.unpack+advance orig_pos = BS.foldl' advance' orig_pos . BS.init   where     advance' (Pos f !line !col !addr) c       | c == nl = Pos f (line + 1) 1 (addr + 1)@@ -261,7 +233,7 @@     Nothing -> 0     Just (d, n) -> n + fromRoman (T.drop (T.length d) s) -readHexRealLit :: RealFloat a => T.Text -> Alex a+readHexRealLit :: RealFloat a => T.Text -> a readHexRealLit s =   let num = T.drop 2 s    in -- extract number into integer, fractional and (optional) exponent@@ -276,5 +248,5 @@                   fracLen = fromIntegral $ T.length f                   fracVal = fracPart / (16.0 ** fracLen)                   totalVal = (intPart + fracVal) * (2.0 ** exponent)-               in pure totalVal+               in totalVal             _ -> error "bad hex real literal"

@@ -1,27 +1,18 @@ {-# OPTIONS_GHC -funbox-strict-fields #-}  -- | Utility definitions used by the lexer.  None of the default Alex--- "wrappers" are precisely what we need.  The code here is based on--- the "monad-bytestring" wrapper.  The code here is completely--- Futhark-agnostic, and perhaps it can even serve as inspiration for--- other Alex lexer wrappers.+-- "wrappers" are precisely what we need.  The code here is highly+-- minimalistic.  Lexers should not be complicated! module Language.Futhark.Parser.Lexer.Wrapper-  ( runAlex,-    Alex,+  ( initialLexerState,     AlexInput,     alexInputPrevChar,-    Byte,     LexerError (..),-    alexSetInput,-    alexGetInput,     alexGetByte,-    alexGetStartCode,-    alexError,     alexGetPos,   ) where -import Control.Applicative (liftA) import Data.ByteString.Internal qualified as BS (w2c) import Data.ByteString.Lazy qualified as BS import Data.Int (Int64)@@ -61,6 +52,9 @@           n' = n + 1        in p' `seq` cs' `seq` n' `seq` Just (b, (p', c, cs', n')) +alexGetPos :: AlexInput -> Pos+alexGetPos (pos, _, _, _) = pos+ tabSize :: Int tabSize = 8 @@ -70,71 +64,11 @@ alexMove (Pos !f !l _ !a) '\n' = Pos f (l + 1) 1 (a + 1) alexMove (Pos !f !l !c !a) _ = Pos f l (c + 1) (a + 1) -data AlexState = AlexState-  { alex_pos :: !Pos, -- position at current input location-    alex_bpos :: !Int64, -- bytes consumed so far-    alex_inp :: BS.ByteString, -- the current input-    alex_chr :: !Char, -- the character before the input-    alex_scd :: !Int -- the current startcode-  }--runAlex :: Pos -> BS.ByteString -> Alex a -> Either LexerError a-runAlex start_pos input (Alex f) =-  case f-    ( AlexState-        { alex_pos = start_pos,-          alex_bpos = 0,-          alex_inp = input,-          alex_chr = '\n',-          alex_scd = 0-        }-    ) of-    Left msg -> Left msg-    Right (_, a) -> Right a--newtype Alex a = Alex {unAlex :: AlexState -> Either LexerError (AlexState, a)}+initialLexerState :: Pos -> BS.ByteString -> AlexInput+initialLexerState start_pos input =+  (start_pos, '\n', input, 0)  data LexerError = LexerError Loc T.Text  instance Show LexerError where   show (LexerError _ s) = T.unpack s--instance Functor Alex where-  fmap = liftA--instance Applicative Alex where-  pure a = Alex $ \s -> Right (s, a)-  fa <*> a = Alex $ \s -> case unAlex fa s of-    Left msg -> Left msg-    Right (s', f) -> case unAlex a s' of-      Left msg -> Left msg-      Right (s'', b) -> Right (s'', f b)--instance Monad Alex where-  m >>= k = Alex $ \s -> case unAlex m s of-    Left msg -> Left msg-    Right (s', a) -> unAlex (k a) s'--alexGetInput :: Alex AlexInput-alexGetInput =-  Alex $ \s@AlexState {alex_pos = pos, alex_bpos = bpos, alex_chr = c, alex_inp = inp} ->-    Right (s, (pos, c, inp, bpos))--alexSetInput :: AlexInput -> Alex ()-alexSetInput (pos, c, inp, bpos) =-  Alex $ \s -> case s-    { alex_pos = pos,-      alex_bpos = bpos,-      alex_chr = c,-      alex_inp = inp-    } of-    state@AlexState {} -> Right (state, ())--alexError :: Loc -> T.Text -> Alex a-alexError loc message = Alex $ const $ Left $ LexerError loc message--alexGetStartCode :: Alex Int-alexGetStartCode = Alex $ \s@AlexState {alex_scd = sc} -> Right (s, sc)--alexGetPos :: Alex Pos-alexGetPos = Alex $ \s -> Right (s, alex_pos s)

@@ -6,12 +6,9 @@ module Language.Futhark.Parser.Monad   ( ParserMonad,     ParserState,-    ReadLineMonad (..),     Comment (..),-    parseInMonad,     parse,     parseWithComments,-    getLinesFromM,     lexer,     mustBeEmpty,     arrayFromList,@@ -38,19 +35,20 @@   ) where -import Control.Applicative (liftA) import Control.Monad import Control.Monad.Except (ExceptT, MonadError (..), runExceptT) import Control.Monad.Trans.Class (lift) import Control.Monad.Trans.State import Data.Array hiding (index)+import Data.ByteString.Lazy qualified as BS import Data.List.NonEmpty qualified as NE import Data.Monoid import Data.Text qualified as T+import Data.Text.Encoding qualified as T import Futhark.Util.Loc import Futhark.Util.Pretty hiding (line, line') import Language.Futhark.Parser.Lexer-import Language.Futhark.Parser.Lexer.Wrapper (LexerError (..))+import Language.Futhark.Parser.Lexer.Wrapper (AlexInput, LexerError (..), initialLexerState) import Language.Futhark.Pretty () import Language.Futhark.Prop import Language.Futhark.Syntax@@ -105,38 +103,10 @@     parserInput :: T.Text,     -- | Note: reverse order.     parserComments :: [Comment],-    parserLexical :: ([L Token], Pos)+    parserLexerState :: AlexInput   } -type ParserMonad = ExceptT SyntaxError (StateT ParserState ReadLineMonad)--data ReadLineMonad a-  = Value a-  | GetLine (Maybe T.Text -> ReadLineMonad a)--readLineFromMonad :: ReadLineMonad (Maybe T.Text)-readLineFromMonad = GetLine Value--instance Monad ReadLineMonad where-  Value x >>= f = f x-  GetLine g >>= f = GetLine $ g >=> f--instance Functor ReadLineMonad where-  fmap = liftA--instance Applicative ReadLineMonad where-  pure = Value-  (<*>) = ap--getLinesFromM :: Monad m => m T.Text -> ReadLineMonad a -> m a-getLinesFromM _ (Value x) = pure x-getLinesFromM fetch (GetLine f) = do-  s <- fetch-  getLinesFromM fetch $ f $ Just s--getNoLines :: ReadLineMonad a -> Either SyntaxError a-getNoLines (Value x) = Right x-getNoLines (GetLine f) = getNoLines $ f Nothing+type ParserMonad = ExceptT SyntaxError (State ParserState)  arrayFromList :: [a] -> Array Int a arrayFromList l = listArray (0, length l - 1) l@@ -169,9 +139,6 @@   where     field (name, pat) = RecordFieldExplicit name <$> patternExp pat <*> pure loc -eof :: Pos -> L Token-eof pos = L (Loc pos pos) EOF- binOpName :: L Token -> (QualName Name, Loc) binOpName (L loc (SYMBOL _ qs op)) = (QualName qs op, loc) binOpName t = error $ "binOpName: unexpected " ++ show t@@ -181,12 +148,6 @@   AppExp (BinOp (QualName qs op, srclocOf loc) NoInfo (x, NoInfo) (y, NoInfo) (srcspan x y)) NoInfo binOp _ t _ = error $ "binOp: unexpected " ++ show t -getTokens :: ParserMonad ([L Token], Pos)-getTokens = lift $ gets parserLexical--putTokens :: ([L Token], Pos) -> ParserMonad ()-putTokens l = lift $ modify $ \env -> env {parserLexical = l}- putComment :: Comment -> ParserMonad () putComment c = lift $ modify $ \env ->   env {parserComments = c : parserComments env}@@ -208,42 +169,21 @@ primNegate (UnsignedValue v) = UnsignedValue $ intNegate v primNegate (BoolValue v) = BoolValue $ not v -readLine :: ParserMonad (Maybe T.Text)-readLine = do-  s <- lift $ lift readLineFromMonad-  case s of-    Just s' ->-      lift $ modify $ \env -> env {parserInput = parserInput env <> "\n" <> s'}-    Nothing -> pure ()-  pure s- lexer :: (L Token -> ParserMonad a) -> ParserMonad a lexer cont = do-  (ts, pos) <- getTokens-  case ts of-    [] -> do-      ended <- lift $ runExceptT $ cont $ eof pos-      case ended of-        Right x -> pure x-        Left parse_e -> do-          line <- readLine-          ts' <--            case line of-              Nothing -> throwError parse_e-              Just line' -> pure $ scanTokensText (advancePos pos '\n') line'-          (ts'', pos') <- either (throwError . lexerErrToParseErr) pure ts'-          case ts'' of-            [] -> cont $ eof pos-            xs -> do-              putTokens (xs, pos')-              lexer cont-    (L loc (COMMENT text) : xs) -> do-      putComment $ Comment loc text-      putTokens (xs, pos)-      lexer cont-    (x : xs) -> do-      putTokens (xs, pos)-      cont x+  ls <- lift $ gets parserLexerState+  case getToken ls of+    Left e ->+      throwError $ lexerErrToParseErr e+    Right (ls', (start, end, tok)) -> do+      let loc = Loc start end+      lift $ modify $ \s -> s {parserLexerState = ls'}+      case tok of+        COMMENT text -> do+          putComment $ Comment loc text+          lexer cont+        _ ->+          cont $ L loc tok  parseError :: (L Token, [String]) -> ParserMonad a parseError (L loc EOF, expected) =@@ -257,7 +197,7 @@ parseError (L loc _, expected) = do   input <- lift $ gets parserInput   let ~(Loc (Pos _ _ _ beg) (Pos _ _ _ end)) = locOf loc-      tok_src = T.take (end - beg + 1) $ T.drop beg input+      tok_src = T.take (end - beg) $ T.drop beg input   parseErrorAt loc . Just . T.unlines $     [ "Unexpected token: '" <> tok_src <> "'",       "Expected one of the following: " <> T.unwords (map T.pack expected)@@ -288,23 +228,23 @@ lexerErrToParseErr :: LexerError -> SyntaxError lexerErrToParseErr (LexerError loc msg) = SyntaxError loc msg -parseInMonad ::+parseWithComments ::   ParserMonad a ->   FilePath ->   T.Text ->-  ReadLineMonad (Either SyntaxError (a, [Comment]))-parseInMonad p file program =-  either-    (pure . Left . lexerErrToParseErr)-    (fmap onRes . runStateT (runExceptT p) . env)-    (scanTokensText (Pos file 1 1 0) program)+  Either SyntaxError (a, [Comment])+parseWithComments p file program =+  onRes $ runState (runExceptT p) env   where-    env = ParserState file program []+    env =+      ParserState+        file+        program+        []+        (initialLexerState start $ BS.fromStrict . T.encodeUtf8 $ program)+    start = Pos file 1 1 0     onRes (Left err, _) = Left err     onRes (Right x, s) = Right (x, reverse $ parserComments s)--parseWithComments :: ParserMonad a -> FilePath -> T.Text -> Either SyntaxError (a, [Comment])-parseWithComments p file program = join $ getNoLines $ parseInMonad p file program  parse :: ParserMonad a -> FilePath -> T.Text -> Either SyntaxError a parse p file program = fst <$> parseWithComments p file program

@@ -8,9 +8,6 @@   , futharkType   , parse   , parseWithComments-  , ReadLineMonad (..)-  , getLinesFromM-  , parseInMonad   , SyntaxError(..)   , Comment(..)   )@@ -212,9 +209,12 @@     | Doc Dec_          { addDoc $1 $2 }  Decs :: { [UncheckedDec] }-      :          { [] }-      | Dec Decs { $1 : $2 }+      : Decs_     { reverse $1 } +Decs_ :: { [UncheckedDec] }+      :           { [] }+      | Decs_ Dec { $2 : $1 }+ Dec_ :: { UncheckedDec }     : Val               { ValDec $1 }     | TypeAbbr          { TypeDec $1 }@@ -320,9 +320,12 @@         { addAttrSpec $2 $4 }  Specs :: { [SpecBase NoInfo Name] }-       : Spec Specs { $1 : $2 }-       |            { [] }+       : Specs_      { reverse $1 } +Specs_ :: { [SpecBase NoInfo Name] }+       : Specs_ Spec { $2 : $1 }+       |             { [] }+ SizeBinder :: { SizeBinder Name }             : '[' id ']'  { let L _ (ID name) = $2 in SizeBinder name (srcspan $1 $>) }             | '...[' id ']' { let L _ (ID name) = $2 in SizeBinder name (srcspan $1 $>) }@@ -600,9 +603,9 @@                         StringLit (BS.unpack (T.encodeUtf8 s)) (srclocOf loc) }      | hole           { Hole NoInfo (srclocOf $1) }      | '(' Exp ')'            { Parens $2 (srcspan $1 $>) }-     | '(' Exp ',' Exps1 ')'  { TupLit ($2 : fst $4 : snd $4) (srcspan $1 $>) }+     | '(' Exp ',' Exps1 ')'  { TupLit ($2 : $4) (srcspan $1 $>) }      | '('      ')'           { TupLit [] (srcspan $1 $>) }-     | '[' Exps1 ']'          { ArrayLit (fst $2:snd $2) NoInfo (srcspan $1 $>) }+     | '[' Exps1 ']'          { ArrayLit $2 NoInfo (srcspan $1 $>) }      | '['       ']'          { ArrayLit [] NoInfo (srcspan $1 $>) }       | id { let L loc (ID v)  = $1 in Var (QualName [] v) NoInfo (srclocOf loc) }@@ -651,14 +654,12 @@         | false  { (BoolValue False, $1) }         | NumLit { $1 } -Exps1 :: { (UncheckedExp, [UncheckedExp]) }-       : Exps1_ { case reverse (snd $1 : fst $1) of-                    []   -> (snd $1, [])-                    y:ys -> (y, ys) }+Exps1 :: { [UncheckedExp] }+       : Exps1_ { reverse $1 } -Exps1_ :: { ([UncheckedExp], UncheckedExp) }-        : Exps1_ ',' Exp { (snd $1 : fst $1, $3) }-        | Exp            { ([], $1) }+Exps1_ :: { [UncheckedExp] }+        : Exps1_ ',' Exp { $3 : $1 }+        | Exp            { [$1] }  FieldAccesses :: { [(Name, Loc)] }                : '.' FieldId FieldAccesses { $2 : $3 }

@@ -8,6 +8,7 @@ import Data.Function ((&)) import Data.List qualified as L import Data.Map qualified as M+import Data.Text qualified as T import Futhark.Analysis.PrimExp.Convert import Futhark.IR.Mem.IxFun qualified as IxFunLMAD import Futhark.IR.Mem.IxFun.Alg qualified as IxFunAlg@@ -17,7 +18,7 @@ import Futhark.IR.Syntax import Futhark.IR.Syntax.Core () import Futhark.Util.IntegralExp qualified as IE-import Futhark.Util.Pretty qualified as PR+import Futhark.Util.Pretty import Test.Tasty import Test.Tasty.HUnit import Prelude hiding (span)@@ -47,42 +48,42 @@           ([], x)           strides -compareIxFuns :: IxFunLMAD.IxFun Int -> IxFunAlg.IxFun Int -> Assertion-compareIxFuns ixfunLMAD ixfunAlg =+compareIxFuns :: Maybe (IxFunLMAD.IxFun Int) -> IxFunAlg.IxFun Int -> Assertion+compareIxFuns (Just ixfunLMAD) ixfunAlg =   let lmadShape = IxFunLMAD.shape ixfunLMAD       algShape = IxFunAlg.shape ixfunAlg       points = allPoints lmadShape       resLMAD = map (IxFunLMAD.index ixfunLMAD) points       resAlg = map (IxFunAlg.index ixfunAlg) points       errorMessage =-        "lmad ixfun:  "-          ++ PR.prettyString ixfunLMAD-          ++ "\n"-          ++ "alg ixfun:   "-          ++ PR.prettyString ixfunAlg-          ++ "\n"-          ++ "lmad shape:  "-          ++ show lmadShape-          ++ "\n"-          ++ "alg shape:   "-          ++ show algShape-          ++ "\n"-          ++ "lmad points length: "-          ++ show (length resLMAD)-          ++ "\n"-          ++ "alg points length:  "-          ++ show (length resAlg)-          ++ "\n"-          ++ "lmad points: "-          ++ show resLMAD-          ++ "\n"-          ++ "alg points:  "-          ++ show resAlg+        T.unpack . docText $+          "lmad ixfun:  " <> pretty ixfunLMAD+            </> "alg ixfun:   " <> pretty ixfunAlg+            </> "lmad shape:  " <> pretty lmadShape+            </> "alg shape:   " <> pretty algShape+            </> "lmad points length: " <> pretty (length resLMAD)+            </> "alg points length:  " <> pretty (length resAlg)+            </> "lmad points: " <> pretty resLMAD+            </> "alg points:  " <> pretty resAlg    in (lmadShape == algShape && resLMAD == resAlg) @? errorMessage+compareIxFuns Nothing ixfunAlg =+  assertFailure $+    unlines+      [ "lmad ixfun: Nothing",+        "alg ixfun:  " <> prettyString ixfunAlg+      ]  compareOps :: IxFunWrap.IxFun Int -> Assertion compareOps (ixfunLMAD, ixfunAlg) = compareIxFuns ixfunLMAD ixfunAlg +compareOpsFailure :: IxFunWrap.IxFun Int -> Assertion+compareOpsFailure (Nothing, _) = pure ()+compareOpsFailure (Just ixfunLMAD, ixfunAlg) =+  assertFailure . T.unpack . docText $+    "Not supposed to be representable as LMAD."+      </> "lmad ixfun: " <> pretty ixfunLMAD+      </> "alg ixfun:  " <> pretty ixfunAlg+ -- XXX: Clean this up. n :: Int n = 19@@ -102,10 +103,11 @@     concat       [ test_iota,         test_slice_iota,-        test_reshape_slice_iota1,+        test_slice_reshape_iota1,         test_permute_slice_iota,+        test_reshape_iota,         test_reshape_permute_iota,-        test_reshape_slice_iota2,+        test_slice_reshape_iota2,         test_reshape_slice_iota3,         test_complex1,         test_complex2,@@ -113,6 +115,8 @@         test_rebase2,         test_rebase3,         test_rebase4_5,+        test_rebase6,+        test_rebase7,         test_flatSlice_iota,         test_slice_flatSlice_iota,         test_flatSlice_flatSlice_iota,@@ -129,194 +133,161 @@  test_iota :: [TestTree] test_iota =-  singleton $-    testCase "iota" $-      compareOps $-        iota [n]+  singleton . testCase "iota" . compareOps $+    iota [n]  test_slice_iota :: [TestTree] test_slice_iota =-  singleton $-    testCase "slice . iota" $-      compareOps $-        slice (iota [n, n, n]) slice3+  singleton . testCase "slice . iota" . compareOps $+    slice (iota [n, n, n]) slice3 -test_reshape_slice_iota1 :: [TestTree]-test_reshape_slice_iota1 =-  singleton $-    testCase "reshape . slice . iota 1" $-      compareOps $-        reshape-          (slice (iota [n, n, n]) slice3)-          [n `P.div` 2, n `P.div` 3]+test_slice_reshape_iota1 :: [TestTree]+test_slice_reshape_iota1 =+  singleton . testCase "slice . reshape . iota 1" . compareOps $+    slice (reshape (iota [n, n, n]) [n `P.div` 2, n `P.div` 3, 1]) slice3  test_permute_slice_iota :: [TestTree] test_permute_slice_iota =-  singleton $-    testCase "permute . slice . iota" $-      compareOps $-        permute (slice (iota [n, n, n]) slice3) [1, 0]+  singleton . testCase "permute . slice . iota" . compareOps $+    permute (slice (iota [n, n, n]) slice3) [1, 0] +test_reshape_iota :: [TestTree]+test_reshape_iota =+  -- This tests a pattern that occurs with ScalarSpace.+  singleton . testCase "reshape . zeroslice . iota" . compareOps $+    let s = Slice [DimSlice 0 n 0, DimSlice 0 n 1]+     in reshape (slice (iota [n, n]) s) [1, n, 1, n]+ test_reshape_permute_iota :: [TestTree] test_reshape_permute_iota =   -- negative reshape test-  singleton $-    testCase "reshape . permute . iota" $-      compareOps $-        let newdims = [n * n, n]-         in reshape (permute (iota [n, n, n]) [1, 2, 0]) newdims+  singleton . testCase "reshape . permute . iota" . compareOpsFailure $+    let newdims = [n * n, n]+     in reshape (permute (iota [n, n, n]) [1, 2, 0]) newdims -test_reshape_slice_iota2 :: [TestTree]-test_reshape_slice_iota2 =-  -- negative reshape test-  singleton $-    testCase "reshape . slice . iota 2" $-      compareOps $-        let newdims = [n * n, n]-            slc =-              Slice-                [ DimFix (n `P.div` 2),-                  DimSlice (n - 1) n (-1),-                  DimSlice 0 n 1,-                  DimSlice (n - 1) n (-1)-                ]-         in reshape (slice (iota [n, n, n, n]) slc) newdims+test_slice_reshape_iota2 :: [TestTree]+test_slice_reshape_iota2 =+  singleton . testCase "slice . reshape . iota 2" . compareOps $+    let newdims = [n * n, n]+        slc =+          Slice+            [ DimFix (n `P.div` 2),+              DimSlice 0 n 1+            ]+     in slice (reshape (iota [n, n, n, n]) newdims) slc  test_reshape_slice_iota3 :: [TestTree] test_reshape_slice_iota3 =   -- negative reshape test-  singleton $-    testCase "reshape . slice . iota 3" $-      compareOps $-        let newdims = [n * n, n]-            slc =-              Slice-                [ DimFix (n `P.div` 2),-                  DimSlice 0 n 1,-                  DimSlice 0 (n `P.div` 2) 1,-                  DimSlice 0 n 1-                ]-         in reshape (slice (iota [n, n, n, n]) slc) newdims+  singleton . testCase "reshape . slice . iota 3" . compareOpsFailure $+    let newdims = [n * n, n]+        slc =+          Slice+            [ DimFix (n `P.div` 2),+              DimSlice 0 n 1,+              DimSlice 0 (n `P.div` 2) 1,+              DimSlice 0 n 1+            ]+     in reshape (slice (iota [n, n, n, n]) slc) newdims  test_complex1 :: [TestTree] test_complex1 =-  singleton $-    testCase "reshape . permute . slice . permute . slice . iota 1" $-      compareOps $-        let newdims =-              [ n,-                n,-                n,-                (n `P.div` 3) - 2-              ]-            slice33 =-              Slice-                [ DimSlice (n - 1) (n `P.div` 3) (-1),-                  DimSlice (n - 1) n (-1),-                  DimSlice (n - 1) n (-1),-                  DimSlice 0 n 1-                ]-            ixfun = permute (slice (iota [n, n, n, n, n]) slice33) [3, 1, 2, 0]-            m = n `P.div` 3-            slice1 =-              Slice-                [ DimSlice 0 n 1,-                  DimSlice (n - 1) n (-1),-                  DimSlice (n - 1) n (-1),-                  DimSlice 1 (m - 2) (-1)-                ]-            ixfun' = reshape (slice ixfun slice1) newdims-         in ixfun'+  singleton . testCase "permute . slice . permute . slice . iota 1" . compareOps $+    let slice33 =+          Slice+            [ DimSlice (n - 1) (n `P.div` 3) (-1),+              DimSlice (n - 1) n (-1),+              DimSlice (n - 1) n (-1),+              DimSlice 0 n 1+            ]+        ixfun = permute (slice (iota [n, n, n, n, n]) slice33) [3, 1, 2, 0]+        m = n `P.div` 3+        slice1 =+          Slice+            [ DimSlice 0 n 1,+              DimSlice (n - 1) n (-1),+              DimSlice (n - 1) n (-1),+              DimSlice 1 (m - 2) (-1)+            ]+        ixfun' = slice ixfun slice1+     in ixfun'  test_complex2 :: [TestTree] test_complex2 =-  singleton $-    testCase "reshape . permute . slice . permute . slice . iota 2" $-      compareOps $-        let newdims =-              [ n,-                n * n,-                (n `P.div` 3) - 2-              ]-            slc2 =-              Slice-                [ DimFix (n `P.div` 2),-                  DimSlice (n - 1) (n `P.div` 3) (-1),-                  DimSlice (n - 1) n (-1),-                  DimSlice (n - 1) n (-1),-                  DimSlice 0 n 1-                ]-            ixfun = permute (slice (iota [n, n, n, n, n]) slc2) [3, 1, 2, 0]-            m = n `P.div` 3-            slice1 =-              Slice-                [ DimSlice 0 n 1,-                  DimSlice (n - 1) n (-1),-                  DimSlice (n - 1) n (-1),-                  DimSlice 1 (m - 2) (-1)-                ]-            ixfun' = reshape (slice ixfun slice1) newdims-         in ixfun'+  singleton . testCase "permute . slice . permute . slice . iota 2" . compareOps $+    let slc2 =+          Slice+            [ DimFix (n `P.div` 2),+              DimSlice (n - 1) (n `P.div` 3) (-1),+              DimSlice (n - 1) n (-1),+              DimSlice (n - 1) n (-1),+              DimSlice 0 n 1+            ]+        ixfun = permute (slice (iota [n, n, n, n, n]) slc2) [3, 1, 2, 0]+        m = n `P.div` 3+        slice1 =+          Slice+            [ DimSlice 0 n 1,+              DimSlice (n - 1) n (-1),+              DimSlice (n - 1) n (-1),+              DimSlice 1 (m - 2) (-1)+            ]+        ixfun' = slice ixfun slice1+     in ixfun'  test_rebase1 :: [TestTree] test_rebase1 =-  singleton $-    testCase "rebase 1" $-      compareOps $-        let slice_base =-              Slice-                [ DimFix (n `P.div` 2),-                  DimSlice 2 (n - 2) 1,-                  DimSlice 3 (n - 3) 1-                ]-            ixfn_base = permute (slice (iota [n, n, n]) slice_base) [1, 0]-            ixfn_orig = permute (iota [n - 3, n - 2]) [1, 0]-            ixfn_rebase = rebase ixfn_base ixfn_orig-         in ixfn_rebase+  singleton . testCase "rebase 1" . compareOpsFailure $+    let slice_base =+          Slice+            [ DimFix (n `P.div` 2),+              DimSlice 2 (n - 2) 1,+              DimSlice 3 (n - 3) 1+            ]+        ixfn_base = permute (slice (iota [n, n, n]) slice_base) [1, 0]+        ixfn_orig = permute (iota [n - 3, n - 2]) [1, 0]+        ixfn_rebase = rebase ixfn_base ixfn_orig+     in ixfn_rebase  test_rebase2 :: [TestTree] test_rebase2 =-  singleton $-    testCase "rebase 2" $-      compareOps $-        let slice_base =-              Slice-                [ DimFix (n `P.div` 2),-                  DimSlice (n - 1) (n - 2) (-1),-                  DimSlice (n - 1) (n - 3) (-1)-                ]-            slice_orig =-              Slice-                [ DimSlice (n - 4) (n - 3) (-1),-                  DimSlice (n - 3) (n - 2) (-1)-                ]-            ixfn_base = permute (slice (iota [n, n, n]) slice_base) [1, 0]-            ixfn_orig = permute (slice (iota [n - 3, n - 2]) slice_orig) [1, 0]-            ixfn_rebase = rebase ixfn_base ixfn_orig-         in ixfn_rebase+  singleton . testCase "rebase 2" . compareOpsFailure $+    let slice_base =+          Slice+            [ DimFix (n `P.div` 2),+              DimSlice (n - 1) (n - 2) (-1),+              DimSlice (n - 1) (n - 3) (-1)+            ]+        slice_orig =+          Slice+            [ DimSlice (n - 4) (n - 3) (-1),+              DimSlice (n - 3) (n - 2) (-1)+            ]+        ixfn_base = permute (slice (iota [n, n, n]) slice_base) [1, 0]+        ixfn_orig = permute (slice (iota [n - 3, n - 2]) slice_orig) [1, 0]+        ixfn_rebase = rebase ixfn_base ixfn_orig+     in ixfn_rebase  test_rebase3 :: [TestTree] test_rebase3 =-  singleton $-    testCase "rebase full orig but not monotonic" $-      compareOps $-        let n2 = (n - 2) `P.div` 3-            n3 = (n - 3) `P.div` 2-            slice_base =-              Slice-                [ DimFix (n `P.div` 2),-                  DimSlice (n - 1) n2 (-3),-                  DimSlice (n - 1) n3 (-2)-                ]-            slice_orig =-              Slice-                [ DimSlice (n3 - 1) n3 (-1),-                  DimSlice (n2 - 1) n2 (-1)-                ]-            ixfn_base = permute (slice (iota [n, n, n]) slice_base) [1, 0]-            ixfn_orig = permute (slice (iota [n3, n2]) slice_orig) [1, 0]-            ixfn_rebase = rebase ixfn_base ixfn_orig-         in ixfn_rebase+  singleton . testCase "rebase full orig but not monotonic" . compareOpsFailure $+    let n2 = (n - 2) `P.div` 3+        n3 = (n - 3) `P.div` 2+        slice_base =+          Slice+            [ DimFix (n `P.div` 2),+              DimSlice (n - 1) n2 (-3),+              DimSlice (n - 1) n3 (-2)+            ]+        slice_orig =+          Slice+            [ DimSlice (n3 - 1) n3 (-1),+              DimSlice (n2 - 1) n2 (-1)+            ]+        ixfn_base = permute (slice (iota [n, n, n]) slice_base) [1, 0]+        ixfn_orig = permute (slice (iota [n3, n2]) slice_orig) [1, 0]+        ixfn_rebase = rebase ixfn_base ixfn_orig+     in ixfn_rebase  test_rebase4_5 :: [TestTree] test_rebase4_5 =@@ -335,54 +306,67 @@           ]       ixfn_base = permute (slice (iota [n, n, n]) slice_base) [1, 0]       ixfn_orig = permute (slice (iota [n3, n2]) slice_orig) [1, 0]-   in [ testCase "rebase mixed monotonicities" $-          compareOps $-            rebase ixfn_base ixfn_orig+   in [ testCase "rebase mixed monotonicities" . compareOpsFailure $+          rebase ixfn_base ixfn_orig       ] +-- Imitates a case from memory expansion.+test_rebase6 :: [TestTree]+test_rebase6 =+  [ testCase "rebase . slice1 . iota" . compareOps $+      rebase+        (slice (iota [n, n, n]) (Slice [DimSlice 0 n 1, DimSlice 0 n 1]))+        ( slice+            (iota [n, n])+            (Slice [DimSlice 1 (n - 1) 1, DimSlice 0 n 1])+        )+  ]++-- Imitates another case from memory expansion.+test_rebase7 :: [TestTree]+test_rebase7 =+  [ testCase "rebase . slice2 . iota" . compareOps $+      rebase+        (slice (iota [n, n, n]) (Slice [DimSlice 0 n 1, DimSlice 0 n 1]))+        ( slice+            (iota [n, n])+            (Slice [DimSlice 0 (n - 1) 1, DimSlice 0 n 1])+        )+  ]+ test_flatSlice_iota :: [TestTree] test_flatSlice_iota =-  singleton $-    testCase "flatSlice . iota" $-      compareOps $-        flatSlice (iota [n * n * n * n]) $-          FlatSlice 2 [FlatDimIndex (n * 2) 4, FlatDimIndex n 3, FlatDimIndex 1 2]+  singleton . testCase "flatSlice . iota" . compareOps $+    flatSlice (iota [n * n * n * n]) $+      FlatSlice 2 [FlatDimIndex (n * 2) 4, FlatDimIndex n 3, FlatDimIndex 1 2]  test_slice_flatSlice_iota :: [TestTree] test_slice_flatSlice_iota =-  singleton $-    testCase "slice . flatSlice . iota " $-      compareOps $-        slice (flatSlice (iota [2 + n * n * n]) flat_slice) $-          Slice [DimFix 2, DimSlice 0 n 1, DimFix 0]+  singleton . testCase "slice . flatSlice . iota " . compareOps $+    slice (flatSlice (iota [2 + n * n * n]) flat_slice) $+      Slice [DimFix 2, DimSlice 0 n 1, DimFix 0]   where     flat_slice = FlatSlice 2 [FlatDimIndex (n * n) 1, FlatDimIndex n 1, FlatDimIndex 1 1]  test_flatSlice_flatSlice_iota :: [TestTree] test_flatSlice_flatSlice_iota =-  singleton $-    testCase "flatSlice . flatSlice . iota " $-      compareOps $-        flatSlice (flatSlice (iota [10 * 10]) flat_slice_1) flat_slice_2+  singleton . testCase "flatSlice . flatSlice . iota " . compareOps $+    flatSlice (flatSlice (iota [10 * 10]) flat_slice_1) flat_slice_2   where     flat_slice_1 = FlatSlice 17 [FlatDimIndex 3 27, FlatDimIndex 3 10, FlatDimIndex 3 1]     flat_slice_2 = FlatSlice 2 [FlatDimIndex 2 (-2)]  test_flatSlice_slice_iota :: [TestTree] test_flatSlice_slice_iota =-  singleton $-    testCase "flatSlice . slice . iota " $-      compareOps $-        flatSlice (slice (iota [210, 100]) $ Slice [DimSlice 10 100 2, DimFix 10]) flat_slice_1+  singleton . testCase "flatSlice . slice . iota " . compareOps $+    flatSlice (slice (iota [210, 100]) $ Slice [DimSlice 10 100 2, DimFix 10]) flat_slice_1   where     flat_slice_1 = FlatSlice 17 [FlatDimIndex 3 27, FlatDimIndex 3 10, FlatDimIndex 3 1]  test_flatSlice_transpose_slice_iota :: [TestTree] test_flatSlice_transpose_slice_iota =-  singleton $-    testCase "flatSlice . transpose . slice . iota " $-      compareOps $-        flatSlice (permute (slice (iota [20, 20]) $ Slice [DimSlice 1 5 2, DimSlice 0 5 2]) [1, 0]) flat_slice_1+  singleton . testCase "flatSlice . transpose . slice . iota " . compareOpsFailure $+    flatSlice (permute (slice (iota [20, 20]) $ Slice [DimSlice 1 5 2, DimSlice 0 5 2]) [1, 0]) flat_slice_1   where     flat_slice_1 = FlatSlice 1 [FlatDimIndex 2 2] @@ -419,15 +403,15 @@ --             lm1 = --               IxFunLMAD.LMAD --                 256---                 [ IxFunLMAD.LMADDim 256 0 (sub64 (num_blocks_8284) 1) 0 IxFunLMAD.Inc,---                   IxFunLMAD.LMADDim 1 0 16 1 IxFunLMAD.Inc,---                   IxFunLMAD.LMADDim 16 0 16 2 IxFunLMAD.Inc+--                 [ IxFunLMAD.LMADDim 256 0 (sub64 (num_blocks_8284) 1) 0 ,+--                   IxFunLMAD.LMADDim 1 0 16 1 ,+--                   IxFunLMAD.LMADDim 16 0 16 2 --                 ] --             lm2 :: IxFunLMAD.LMAD (TPrimExp Int64 VName) --             lm2 = --               IxFunLMAD.LMAD --                 (add_nw64 (add_nw64 (add_nw64 (add_nw64 (mul_nw64 (256) (num_blocks_8284)) (256)) (mul_nw64 (gtid_8472) (mul_nw64 (256) (num_blocks_8284)))) (mul_nw64 (gtid_8473) (256))) (mul_nw64 (gtid_8474) (16)))---                 [IxFunLMAD.LMADDim 1 0 16 0 IxFunLMAD.Inc]+--                 [IxFunLMAD.LMADDim 1 0 16 0 ] --          in testCase (pretty lm1 <> " and " <> pretty lm2) $ IxFunLMAD.disjoint2 lessthans nonnegs lm1 lm2 @? "Failed" --       ] @@ -480,37 +464,37 @@               lm1 =                 IxFunLMAD.LMAD                   (add_nw64 (mul64 block_size_12121 i_12214) (mul_nw64 (add_nw64 gtid_12553 1) (sub64 (mul64 block_size_12121 n_blab) block_size_12121)))-                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) (sub_nw64 (sub_nw64 (add64 1 i_12214) gtid_12553) 1) 0 IxFunLMAD.Inc,-                    IxFunLMAD.LMADDim 1 (block_size_12121 + 1) 1 IxFunLMAD.Inc+                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) (sub_nw64 (sub_nw64 (add64 1 i_12214) gtid_12553) 1) 0,+                    IxFunLMAD.LMADDim 1 (block_size_12121 + 1) 1                   ]                lm2 =                 IxFunLMAD.LMAD                   (block_size_12121 * i_12214)-                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) gtid_12553 0 IxFunLMAD.Inc,-                    IxFunLMAD.LMADDim 1 (1 + block_size_12121) 1 IxFunLMAD.Inc+                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) gtid_12553 0,+                    IxFunLMAD.LMADDim 1 (1 + block_size_12121) 1                   ]                lm_w =                 IxFunLMAD.LMAD                   (add_nw64 (add64 (add64 1 n_blab) (mul64 block_size_12121 i_12214)) (mul_nw64 gtid_12553 (sub64 (mul64 block_size_12121 n_blab) block_size_12121)))-                  [ IxFunLMAD.LMADDim n_blab block_size_12121 0 IxFunLMAD.Inc,-                    IxFunLMAD.LMADDim 1 block_size_12121 1 IxFunLMAD.Inc+                  [ IxFunLMAD.LMADDim n_blab block_size_12121 0,+                    IxFunLMAD.LMADDim 1 block_size_12121 1                   ]                lm_blocks =                 IxFunLMAD.LMAD                   (block_size_12121 * i_12214 + n_blab + 1)-                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) (i_12214 + 1) 0 IxFunLMAD.Inc,-                    IxFunLMAD.LMADDim n_blab block_size_12121 1 IxFunLMAD.Inc,-                    IxFunLMAD.LMADDim 1 block_size_12121 2 IxFunLMAD.Inc+                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) (i_12214 + 1) 0,+                    IxFunLMAD.LMADDim n_blab block_size_12121 1,+                    IxFunLMAD.LMADDim 1 block_size_12121 2                   ]                lm_lower_per =                 IxFunLMAD.LMAD                   (block_size_12121 * i_12214)-                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) (i_12214 + 1) 0 IxFunLMAD.Inc,-                    IxFunLMAD.LMADDim 1 (block_size_12121 + 1) 1 IxFunLMAD.Inc+                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) (i_12214 + 1) 0,+                    IxFunLMAD.LMADDim 1 (block_size_12121 + 1) 1                   ]                res1 = disjointTester asserts lessthans lm1 lm_w@@ -539,36 +523,36 @@               lm1 =                 IxFunLMAD.LMAD                   (add_nw64 (add64 n_blab (sub64 (sub64 (mul64 n_blab (add64 1 (mul64 block_size_12121 (add64 1 i_12214)))) block_size_12121) 1)) (mul_nw64 (add_nw64 gtid_12553 1) (sub64 (mul64 block_size_12121 n_blab) block_size_12121)))-                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) (sub_nw64 (sub_nw64 (sub64 (sub64 (sdiv64 (sub64 n_blab 1) block_size_12121) i_12214) 1) gtid_12553) 1) 0 IxFunLMAD.Inc,-                    IxFunLMAD.LMADDim n_blab block_size_12121 1 IxFunLMAD.Inc+                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) (sub_nw64 (sub_nw64 (sub64 (sub64 (sdiv64 (sub64 n_blab 1) block_size_12121) i_12214) 1) gtid_12553) 1) 0,+                    IxFunLMAD.LMADDim n_blab block_size_12121 1                   ]                lm2 =                 IxFunLMAD.LMAD                   (add_nw64 (sub64 (sub64 (mul64 n_blab (add64 1 (mul64 block_size_12121 (add64 1 i_12214)))) block_size_12121) 1) (mul_nw64 (add_nw64 gtid_12553 1) (sub64 (mul64 block_size_12121 n_blab) block_size_12121)))-                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) (sub_nw64 (sub_nw64 (sub64 (sub64 (sdiv64 (sub64 n_blab 1) block_size_12121) i_12214) 1) gtid_12553) 1) 0 IxFunLMAD.Inc,-                    IxFunLMAD.LMADDim 1 (1 + block_size_12121) 1 IxFunLMAD.Inc+                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) (sub_nw64 (sub_nw64 (sub64 (sub64 (sdiv64 (sub64 n_blab 1) block_size_12121) i_12214) 1) gtid_12553) 1) 0,+                    IxFunLMAD.LMADDim 1 (1 + block_size_12121) 1                   ]                lm3 =                 IxFunLMAD.LMAD                   (add64 n_blab (sub64 (sub64 (mul64 n_blab (add64 1 (mul64 block_size_12121 (add64 1 i_12214)))) block_size_12121) 1))-                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) gtid_12553 0 IxFunLMAD.Inc,-                    IxFunLMAD.LMADDim n_blab block_size_12121 1 IxFunLMAD.Inc+                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) gtid_12553 0,+                    IxFunLMAD.LMADDim n_blab block_size_12121 1                   ]                lm4 =                 IxFunLMAD.LMAD                   (sub64 (sub64 (mul64 n_blab (add64 1 (mul64 block_size_12121 (add64 1 i_12214)))) block_size_12121) 1)-                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) gtid_12553 0 IxFunLMAD.Inc,-                    IxFunLMAD.LMADDim 1 (1 + block_size_12121) 1 IxFunLMAD.Inc+                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) gtid_12553 0,+                    IxFunLMAD.LMADDim 1 (1 + block_size_12121) 1                   ]                lm_w =                 IxFunLMAD.LMAD                   (add_nw64 (sub64 (mul64 n_blab (add64 2 (mul64 block_size_12121 (add64 1 i_12214)))) block_size_12121) (mul_nw64 gtid_12553 (sub64 (mul64 block_size_12121 n_blab) block_size_12121)))-                  [ IxFunLMAD.LMADDim n_blab block_size_12121 0 IxFunLMAD.Inc,-                    IxFunLMAD.LMADDim 1 block_size_12121 1 IxFunLMAD.Inc+                  [ IxFunLMAD.LMADDim n_blab block_size_12121 0,+                    IxFunLMAD.LMADDim 1 block_size_12121 1                   ]                res1 = disjointTester asserts lessthans lm1 lm_w@@ -589,29 +573,29 @@               lm1 =                 IxFunLMAD.LMAD                   (1024 * num_blocks * (1 + step) + 1024 * step)-                  [ IxFunLMAD.LMADDim (1024 * num_blocks) (num_blocks - step - 1) 0 IxFunLMAD.Inc,-                    IxFunLMAD.LMADDim 32 32 1 IxFunLMAD.Inc,-                    IxFunLMAD.LMADDim 1 32 2 IxFunLMAD.Inc+                  [ IxFunLMAD.LMADDim (1024 * num_blocks) (num_blocks - step - 1) 0,+                    IxFunLMAD.LMADDim 32 32 1,+                    IxFunLMAD.LMADDim 1 32 2                   ]                lm_w1 =                 IxFunLMAD.LMAD                   (1024 * num_blocks * step + 1024 * step)-                  [ IxFunLMAD.LMADDim 32 32 0 IxFunLMAD.Inc,-                    IxFunLMAD.LMADDim 1 32 1 IxFunLMAD.Inc+                  [ IxFunLMAD.LMADDim 32 32 0,+                    IxFunLMAD.LMADDim 1 32 1                   ]                lm_w2 =                 IxFunLMAD.LMAD                   ((1 + step) * 1024 * num_blocks + (1 + step) * 1024)-                  [ IxFunLMAD.LMADDim (1024 * num_blocks) (num_blocks - step - 1) 0 IxFunLMAD.Inc,-                    IxFunLMAD.LMADDim 1024 (num_blocks - step - 1) 1 IxFunLMAD.Inc,-                    IxFunLMAD.LMADDim 1024 1 2 IxFunLMAD.Inc,-                    IxFunLMAD.LMADDim 32 1 3 IxFunLMAD.Inc,-                    IxFunLMAD.LMADDim 128 8 4 IxFunLMAD.Inc,-                    IxFunLMAD.LMADDim 4 8 5 IxFunLMAD.Inc,-                    IxFunLMAD.LMADDim 32 4 6 IxFunLMAD.Inc,-                    IxFunLMAD.LMADDim 1 4 7 IxFunLMAD.Inc+                  [ IxFunLMAD.LMADDim (1024 * num_blocks) (num_blocks - step - 1) 0,+                    IxFunLMAD.LMADDim 1024 (num_blocks - step - 1) 1,+                    IxFunLMAD.LMADDim 1024 1 2,+                    IxFunLMAD.LMADDim 32 1 3,+                    IxFunLMAD.LMADDim 128 8 4,+                    IxFunLMAD.LMADDim 4 8 5,+                    IxFunLMAD.LMADDim 32 4 6,+                    IxFunLMAD.LMADDim 1 4 7                   ]                asserts =

@@ -12,6 +12,7 @@   ) where +import Control.Monad (join) import Futhark.IR.Mem.IxFun qualified as I import Futhark.IR.Mem.IxFun.Alg qualified as IA import Futhark.IR.Syntax (FlatSlice, Slice)@@ -21,52 +22,52 @@  type Permutation = [Int] -type IxFun num = (I.IxFun num, IA.IxFun num)+type IxFun num = (Maybe (I.IxFun num), IA.IxFun num)  iota ::   IntegralExp num =>   Shape num ->   IxFun num-iota x = (I.iota x, IA.iota x)+iota x = (Just $ I.iota x, IA.iota x)  permute ::   IntegralExp num =>   IxFun num ->   Permutation ->   IxFun num-permute (l, a) x = (I.permute l x, IA.permute a x)+permute (l, a) x = (I.permute <$> l <*> pure x, IA.permute a x)  reshape ::   (Eq num, IntegralExp num) =>   IxFun num ->   Shape num ->   IxFun num-reshape (l, a) x = (I.reshape l x, IA.reshape a x)+reshape (l, a) x = (join (I.reshape <$> l <*> pure x), IA.reshape a x)  coerce ::   (Eq num, IntegralExp num) =>   IxFun num ->   Shape num ->   IxFun num-coerce (l, a) x = (I.coerce l x, IA.coerce a x)+coerce (l, a) x = (I.coerce <$> l <*> pure x, IA.coerce a x)  slice ::   (Eq num, IntegralExp num) =>   IxFun num ->   Slice num ->   IxFun num-slice (l, a) x = (I.slice l x, IA.slice a x)+slice (l, a) x = (I.slice <$> l <*> pure x, IA.slice a x)  flatSlice ::   (Eq num, IntegralExp num) =>   IxFun num ->   FlatSlice num ->   IxFun num-flatSlice (l, a) x = (I.flatSlice l x, IA.flatSlice a x)+flatSlice (l, a) x = (join (I.flatSlice <$> l <*> pure x), IA.flatSlice a x)  rebase ::   (Eq num, IntegralExp num) =>   IxFun num ->   IxFun num ->   IxFun num-rebase (l, a) (l1, a1) = (I.rebase l l1, IA.rebase a a1)+rebase (l, a) (l1, a1) = (join (I.rebase <$> l <*> l1), IA.rebase a a1)