futhark 0.25.9 → 0.25.10

raw patch · 62 files changed

+1843/−2459 lines, 62 filesPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

API changes (from Hackage documentation)

Files

• CHANGELOG.md +33/−0
• docs/c-api.rst +1/−3
• docs/language-reference.rst +14/−11
• futhark.cabal +1/−4
• rts/c/backends/cuda.h +33/−2
• rts/c/backends/hip.h +23/−1
• rts/c/backends/opencl.h +42/−0
• rts/c/context.h +1/−1
• rts/c/gpu_prototypes.h +3/−0
• rts/python/opencl.py +15/−0
• src/Futhark/Actions.hs +1/−1
• src/Futhark/CLI/Defs.hs +1/−1
• src/Futhark/CLI/Dev.hs +0/−19
• src/Futhark/CLI/LSP.hs +1/−1
• src/Futhark/CLI/Literate.hs +17/−15
• src/Futhark/CLI/Misc.hs +3/−1
• src/Futhark/CodeGen/Backends/CCUDA.hs +6/−4
• src/Futhark/CodeGen/Backends/COpenCL.hs +5/−3
• src/Futhark/CodeGen/Backends/GPU.hs +33/−9
• src/Futhark/CodeGen/Backends/GenericC/Code.hs +0/−10
• src/Futhark/CodeGen/Backends/GenericPython.hs +0/−19
• src/Futhark/CodeGen/Backends/HIP.hs +5/−3
• src/Futhark/CodeGen/Backends/PyOpenCL.hs +9/−7
• src/Futhark/CodeGen/Backends/PyOpenCL/Boilerplate.hs +25/−3
• src/Futhark/CodeGen/ImpCode/GPU.hs +11/−10
• src/Futhark/CodeGen/ImpCode/OpenCL.hs +5/−1
• src/Futhark/CodeGen/ImpGen.hs +6/−11
• src/Futhark/CodeGen/ImpGen/GPU.hs +8/−8
• src/Futhark/CodeGen/ImpGen/GPU/Base.hs +101/−58
• src/Futhark/CodeGen/ImpGen/GPU/Group.hs +0/−1
• src/Futhark/CodeGen/ImpGen/GPU/SegHist.hs +8/−12
• src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs +1032/−842
• src/Futhark/CodeGen/ImpGen/GPU/SegScan/SinglePass.hs +77/−98
• src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs +27/−29
• src/Futhark/CodeGen/ImpGen/Multicore/Base.hs +0/−1
• src/Futhark/CodeGen/ImpGen/Multicore/SegHist.hs +0/−1
• src/Futhark/Compiler/Program.hs +3/−2
• src/Futhark/Doc/Generator.hs +34/−34
• src/Futhark/IR/GPU/Sizes.hs +8/−0
• src/Futhark/IR/SegOp.hs +7/−7
• src/Futhark/Internalise/Defunctorise.hs +3/−3
• src/Futhark/LSP/Handlers.hs +6/−11
• src/Futhark/Optimise/InPlaceLowering.hs +0/−416
• src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs +0/−376
• src/Futhark/Optimise/InPlaceLowering/SubstituteIndices.hs +0/−191
• src/Futhark/Passes.hs +2/−6
• src/Futhark/Util/IntegralExp.hs +3/−0
• src/Language/Futhark/Interpreter.hs +25/−18
• src/Language/Futhark/Parser/Monad.hs +1/−1
• src/Language/Futhark/Parser/Parser.y +22/−22
• src/Language/Futhark/Pretty.hs +10/−10
• src/Language/Futhark/Primitive.hs +2/−2
• src/Language/Futhark/Prop.hs +29/−29
• src/Language/Futhark/Query.hs +29/−29
• src/Language/Futhark/Semantic.hs +8/−8
• src/Language/Futhark/Syntax.hs +33/−33
• src/Language/Futhark/TypeChecker.hs +40/−40
• src/Language/Futhark/TypeChecker/Modules.hs +13/−13
• src/Language/Futhark/TypeChecker/Monad.hs +2/−2
• src/Language/Futhark/TypeChecker/Terms.hs +2/−2
• src/Language/Futhark/TypeChecker/Terms/Pat.hs +1/−1
• src/Language/Futhark/TypeChecker/Types.hs +13/−13

@@ -5,6 +5,39 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/) and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.html). +## [0.25.10]++### Added++* Faster non-commutative reductions in the GPU backends. Work by+  Anders Holst and Christian Påbøl Jacobsen.++### Fixed++* Interpreter crash for certain complicated size expressions involving+  internal bindings (#2053).++* Incorrect type checking of `let` binding with explicit size+  quantification, where size appears in type of body (#2048).++* GPU code generation for non-commutative non-segmented reductions+  with array operands (#2051).++* Histogram with non-vectorised reduction operators (#2056). (But it+  is probably not a good idea to write such programs.)++* Futhark's LSP server should work better with Eglot.++* Incorrect copy removal inside histograms could cause compiler error+  (#2058).++* CUDA backend now correctly queries for available shared memory,+  which affects performance (hopefully positively).++* `futhark literate` now switches to the directory containing the+  `.fut` file before executing its contents. This fixes accessing+  files through relative paths.+ ## [0.25.9]  ### Added

@@ -141,9 +141,7 @@    cache was hit succesfully, but you can enable logging to see what    happens. -   The lifespan of ``fname`` must exceed the lifespan of the-   configuration object.  Pass ``NULL`` to disable caching (this is-   the default).+   Pass ``NULL`` to disable caching (this is the default).  Context -------

@@ -576,6 +576,9 @@   ``t``.  To pass a single array-typed parameter, enclose it in   parens. +* The bodies of ``let``, ``if``, and ``loop`` extend as far to the+  right as possible.+ * The following table describes the precedence and associativity of   infix operators in both expressions and type expressions.  All   operators in the same row have the same precedence.  The rows are@@ -836,10 +839,10 @@      Company any two values of numeric type for equality. -  ```symbol```+  ```qualname``` -    Use ``symbol``, which may be any non-operator function name, as an-    infix operator.+    Use ``qualname``, which may be any non-operator function name, as+    an infix operator.  ``x && y`` ..........@@ -1499,14 +1502,14 @@    mod_param: "(" `name` ":" `mod_type_exp` ")"    mod_type_bind: "module" "type" `name` "=" `mod_type_exp` -Futhark supports an ML-style higher-order module system.  *Modules*-can contain types, functions, and other modules and module types.-*Module types* are used to classify the contents of modules, and-*parametric modules* are used to abstract over modules (essentially-module-level functions).  In Standard ML, modules, module types and-parametric modules are called structs, signatures, and functors,-respectively.  Module names exist in the same name space as values,-but module types are their own name space.+Futhark supports an ML-style higher-order module system. *Modules* can+contain types, functions, and other modules and module types. *Module+types* are used to classify the contents of modules, and *parametric+modules* are used to abstract over modules (essentially module-level+functions). In Standard ML, modules, module types and parametric+modules are called *structs*, *signatures*, and *functors*,+respectively. Module names exist in the same name space as values, but+module types are their own name space.  Module bindings ~~~~~~~~~~~~~~~

@@ -1,6 +1,6 @@ cabal-version: 2.4 name:           futhark-version:        0.25.9+version:        0.25.10 synopsis:       An optimising compiler for a functional, array-oriented language.  description:    Futhark is a small programming language designed to be compiled to@@ -313,9 +313,6 @@       Futhark.Optimise.Fusion.TryFusion       Futhark.Optimise.GenRedOpt       Futhark.Optimise.HistAccs-      Futhark.Optimise.InPlaceLowering-      Futhark.Optimise.InPlaceLowering.LowerIntoStm-      Futhark.Optimise.InPlaceLowering.SubstituteIndices       Futhark.Optimise.InliningDeadFun       Futhark.Optimise.MemoryBlockMerging       Futhark.Optimise.MemoryBlockMerging.GreedyColoring

@@ -289,6 +289,8 @@   size_t max_threshold;   size_t max_local_memory;   size_t max_bespoke;+  size_t max_registers;+  size_t max_cache;    size_t lockstep_width; @@ -429,6 +431,10 @@   int arch_set = 0, num_extra_opts;   struct futhark_context_config *cfg = ctx->cfg; +  char** macro_names;+  int64_t* macro_vals;+  int num_macros = gpu_macros(ctx, &macro_names, &macro_vals);+   // nvrtc cannot handle multiple -arch options.  Hence, if one of the   // extra_opts is -arch, we have to be careful not to do our usual   // automatic generation.@@ -441,7 +447,7 @@     }   } -  size_t i = 0, n_opts_alloc = 20 + num_extra_opts + cfg->num_tuning_params;+  size_t i = 0, n_opts_alloc = 20 + num_macros + num_extra_opts + cfg->num_tuning_params;   char **opts = (char**) malloc(n_opts_alloc * sizeof(char *));   if (!arch_set) {     opts[i++] = strdup("-arch");@@ -457,6 +463,17 @@   opts[i++] = msgprintf("-D%s=%d",                         "max_group_size",                         (int)ctx->max_group_size);+  opts[i++] = msgprintf("-D%s=%d",+                        "max_local_memory",+                        (int)ctx->max_local_memory);+  opts[i++] = msgprintf("-D%s=%d",+                        "max_registers",+                        (int)ctx->max_registers);++  for (int j = 0; j < num_macros; j++) {+    opts[i++] = msgprintf("-D%s=%zu", macro_names[j], macro_vals[j]);+  }+   for (int j = 0; j < cfg->num_tuning_params; j++) {     opts[i++] = msgprintf("-D%s=%zu", cfg->tuning_param_vars[j],                           cfg->tuning_params[j]);@@ -485,6 +502,9 @@   opts[i++] = msgprintf("-DTR_TILE_DIM=%d", TR_TILE_DIM);   opts[i++] = msgprintf("-DTR_ELEMS_PER_THREAD=%d", TR_ELEMS_PER_THREAD); +  free(macro_names);+  free(macro_vals);+   *n_opts = i;   *opts_out = opts; }@@ -798,12 +818,19 @@    free_list_init(&ctx->gpu_free_list); -  ctx->max_local_memory = device_query(ctx->dev, MAX_SHARED_MEMORY_PER_BLOCK);+  // MAX_SHARED_MEMORY_PER_BLOCK gives bogus numbers (48KiB); probably+  // for backwards compatibility.  Add _OPTIN and you seem to get the+  // right number.+  ctx->max_local_memory =+    device_query(ctx->dev, MAX_SHARED_MEMORY_PER_BLOCK_OPTIN) -+    device_query(ctx->dev, RESERVED_SHARED_MEMORY_PER_BLOCK);   ctx->max_group_size = device_query(ctx->dev, MAX_THREADS_PER_BLOCK);   ctx->max_grid_size = device_query(ctx->dev, MAX_GRID_DIM_X);   ctx->max_tile_size = sqrt(ctx->max_group_size);   ctx->max_threshold = 0;   ctx->max_bespoke = 0;+  ctx->max_registers = device_query(ctx->dev, MAX_REGISTERS_PER_BLOCK);+  ctx->max_cache = device_query(ctx->dev, L2_CACHE_SIZE);   ctx->lockstep_width = device_query(ctx->dev, WARP_SIZE);   CUDA_SUCCEED_FATAL(cuStreamCreate(&ctx->stream, CU_STREAM_DEFAULT));   cuda_size_setup(ctx);@@ -853,6 +880,10 @@     fprintf(ctx->log, "Creating kernel %s.\n", name);   }   CUDA_SUCCEED_FATAL(cuModuleGetFunction(kernel, ctx->module, name));+  // Unless the below is set, the kernel is limited to 48KiB of memory.+  CUDA_SUCCEED_FATAL(cuFuncSetAttribute(*kernel,+                                        cudaFuncAttributeMaxDynamicSharedMemorySize,+                                        ctx->max_local_memory)); }  static void gpu_free_kernel(struct futhark_context *ctx,

@@ -263,6 +263,8 @@   size_t max_threshold;   size_t max_local_memory;   size_t max_bespoke;+  size_t max_registers;+  size_t max_cache;    size_t lockstep_width; @@ -452,6 +454,10 @@   int arch_set = 0, num_extra_opts;   struct futhark_context_config *cfg = ctx->cfg; +  char** macro_names;+  int64_t* macro_vals;+  int num_macros = gpu_macros(ctx, &macro_names, &macro_vals);+   for (num_extra_opts = 0; extra_opts[num_extra_opts] != NULL; num_extra_opts++) {     if (strstr(extra_opts[num_extra_opts], "--gpu-architecture")         == extra_opts[num_extra_opts]) {@@ -459,7 +465,7 @@     }   } -  size_t i = 0, n_opts_alloc = 20 + num_extra_opts + cfg->num_tuning_params;+  size_t i = 0, n_opts_alloc = 20 + num_macros + num_extra_opts + cfg->num_tuning_params;   char **opts = (char**) malloc(n_opts_alloc * sizeof(char *));   if (!arch_set) {     hipDeviceProp_t props;@@ -473,6 +479,17 @@   opts[i++] = msgprintf("-D%s=%d",                         "max_group_size",                         (int)ctx->max_group_size);+  opts[i++] = msgprintf("-D%s=%d",+                        "max_local_memory",+                        (int)ctx->max_local_memory);+  opts[i++] = msgprintf("-D%s=%d",+                        "max_registers",+                        (int)ctx->max_registers);++  for (int j = 0; j < num_macros; j++) {+    opts[i++] = msgprintf("-D%s=%zu", macro_names[j], macro_vals[j]);+  }+   for (int j = 0; j < cfg->num_tuning_params; j++) {     opts[i++] = msgprintf("-D%s=%zu", cfg->tuning_param_vars[j],                           cfg->tuning_params[j]);@@ -488,6 +505,9 @@   opts[i++] = msgprintf("-DTR_TILE_DIM=%d", TR_TILE_DIM);   opts[i++] = msgprintf("-DTR_ELEMS_PER_THREAD=%d", TR_ELEMS_PER_THREAD); +  free(macro_names);+  free(macro_vals);+   *n_opts = i;   *opts_out = opts; }@@ -659,6 +679,8 @@   ctx->max_tile_size = sqrt(ctx->max_group_size);   ctx->max_threshold = 0;   ctx->max_bespoke = 0;+  ctx->max_registers = device_query(ctx->dev, hipDeviceAttributeMaxRegistersPerBlock);+  ctx->max_cache = device_query(ctx->dev, hipDeviceAttributeL2CacheSize);   // FIXME: in principle we should query hipDeviceAttributeWarpSize   // from the device, which will provide 64 on AMD GPUs.   // Unfortunately, we currently do nasty implicit intra-warp

@@ -528,6 +528,8 @@   size_t max_tile_size;   size_t max_threshold;   size_t max_local_memory;+  size_t max_registers;+  size_t max_cache;    size_t lockstep_width; @@ -579,10 +581,18 @@     compile_opts_size += strlen(ctx->cfg->tuning_param_names[i]) + 20;   } +  char** macro_names;+  int64_t* macro_vals;+  int num_macros = gpu_macros(ctx, &macro_names, &macro_vals);+   for (int i = 0; extra_build_opts[i] != NULL; i++) {     compile_opts_size += strlen(extra_build_opts[i] + 1);   } +  for (int i = 0; i < num_macros; i++) {+    compile_opts_size += strlen(macro_names[i]) + 1 + 20;+  }+   char *compile_opts = (char*) malloc(compile_opts_size);    int w = snprintf(compile_opts, compile_opts_size,@@ -594,6 +604,16 @@                 "max_group_size",                 (int)ctx->max_group_size); +  w += snprintf(compile_opts+w, compile_opts_size-w,+                "-D%s=%d ",+                "max_local_memory",+                (int)ctx->max_local_memory);++  w += snprintf(compile_opts+w, compile_opts_size-w,+                "-D%s=%d ",+                "max_registers",+                (int)ctx->max_registers);+   for (int i = 0; i < ctx->cfg->num_tuning_params; i++) {     w += snprintf(compile_opts+w, compile_opts_size-w,                   "-D%s=%d ",@@ -606,6 +626,11 @@                   "%s ", extra_build_opts[i]);   } +  for (int i = 0; i < num_macros; i++) {+    w += snprintf(compile_opts+w, compile_opts_size-w,+                  "-D%s=%zu ", macro_names[i], macro_vals[i]);+  }+   w += snprintf(compile_opts+w, compile_opts_size-w,                 "-DTR_BLOCK_DIM=%d -DTR_TILE_DIM=%d -DTR_ELEMS_PER_THREAD=%d ",                 TR_BLOCK_DIM, TR_TILE_DIM, TR_ELEMS_PER_THREAD);@@ -616,6 +641,9 @@     w += snprintf(compile_opts+w, compile_opts_size-w, "-DEMULATE_F16 ");   } +  free(macro_names);+  free(macro_vals);+   return compile_opts; } @@ -784,6 +812,20 @@     }     ctx->cfg->default_tile_size = max_tile_size;   }+++  cl_ulong cache_size;+  OPENCL_SUCCEED_FATAL(clGetDeviceInfo(device_option.device, CL_DEVICE_GLOBAL_MEM_CACHE_SIZE,+                                       sizeof(cache_size), &cache_size, NULL));++  if (cache_size == 0) {+    // Some code assumes nonzero cache.+    cache_size = 1024*1024;+  }++  ctx->max_cache = cache_size;++  ctx->max_registers = 1<<16; // I cannot find a way to query for this.    ctx->max_group_size = max_group_size;   ctx->max_tile_size = max_tile_size; // No limit.

@@ -130,8 +130,8 @@   ctx->profiling_paused = 0;   ctx->error = NULL;   ctx->log = stderr;+  set_tuning_params(ctx);   if (backend_context_setup(ctx) == 0) {-    set_tuning_params(ctx);     setup_program(ctx);     init_constants(ctx);     (void)futhark_context_clear_caches(ctx);

@@ -5,6 +5,9 @@ #define TR_TILE_DIM (TR_BLOCK_DIM*2) #define TR_ELEMS_PER_THREAD 8 +// Must be defined by the user.+static int gpu_macros(struct futhark_context *ctx, char*** names, int64_t** values);+ struct builtin_kernels* init_builtin_kernels(struct futhark_context* ctx); void free_builtin_kernels(struct futhark_context* ctx, struct builtin_kernels* kernels); static int gpu_free_all(struct futhark_context *ctx);

@@ -102,6 +102,16 @@     return sizes  +def to_c_str_rep(x):+    if type(x) is bool or type(x) is np.bool_:+        if x:+            return "true"+        else:+            return "false"+    else:+        return str(x)++ def initialise_opencl_object(     self,     program_src="",@@ -119,6 +129,7 @@     required_types=[],     all_sizes={},     user_sizes={},+    constants=[], ):     if command_queue is None:         self.ctx = get_prefered_context(@@ -283,6 +294,10 @@                 v,             )             for (s, v) in self.sizes.items()+        ]++        build_options += [+            "-D{}={}".format(s, to_c_str_rep(f())) for (s, f) in constants         ]          if self.platform.name == "Oclgrind":

@@ -166,7 +166,7 @@   Action     { actionName = "Compile imperative kernels",       actionDescription = "Translate program into imperative IL with kernels and write it on standard output.",-      actionProcedure = liftIO . putStrLn . prettyString . snd <=< ImpGenGPU.compileProgOpenCL+      actionProcedure = liftIO . putStrLn . prettyString . snd <=< ImpGenGPU.compileProgHIP     }  -- | Convert the program to CPU multicore ImpCode and print it to stdout.

@@ -33,7 +33,7 @@     defsInDec (LocalDec d _) = defsInDec d     defsInDec (OpenDec me _) = defsInModExp me     defsInDec (ModDec mb) = defsInModExp $ modExp mb-    defsInDec SigDec {} = mempty+    defsInDec ModTypeDec {} = mempty     defsInDec ImportDec {} = mempty      defsInModExp ModVar {} = mempty

@@ -35,7 +35,6 @@ import Futhark.Optimise.DoubleBuffer import Futhark.Optimise.Fusion import Futhark.Optimise.HistAccs-import Futhark.Optimise.InPlaceLowering import Futhark.Optimise.InliningDeadFun import Futhark.Optimise.MemoryBlockMerging qualified as MemoryBlockMerging import Futhark.Optimise.ReduceDeviceSyncs (reduceDeviceSyncs)@@ -330,23 +329,6 @@     long = [passLongOption pass]     pass = Seq.explicitAllocations -iplOption :: String -> FutharkOption-iplOption short =-  passOption (passDescription pass) (UntypedPass perform) short long-  where-    perform (GPU prog) config =-      GPU-        <$> runPipeline (onePass inPlaceLoweringGPU) config prog-    perform (Seq prog) config =-      Seq-        <$> runPipeline (onePass inPlaceLoweringSeq) config prog-    perform s _ =-      externalErrorS $-        "Pass '" ++ passDescription pass ++ "' cannot operate on " ++ representation s--    long = [passLongOption pass]-    pass = inPlaceLoweringSeq- cseOption :: String -> FutharkOption cseOption short =   passOption (passDescription pass) (UntypedPass perform) short long@@ -616,7 +598,6 @@     kernelsPassOption reduceDeviceSyncs [],     typedPassOption soacsProg GPU extractKernels [],     typedPassOption soacsProg MC extractMulticore [],-    iplOption [],     allocateOption "a",     kernelsMemPassOption doubleBufferGPU [],     mcMemPassOption doubleBufferMC [],

@@ -39,7 +39,7 @@ syncOptions :: TextDocumentSyncOptions syncOptions =   TextDocumentSyncOptions-    { _openClose = Just True,+    { _openClose = Just False,       _change = Just TextDocumentSyncKind_Incremental,       _willSave = Just False,       _willSaveWaitUntil = Just False,

@@ -45,7 +45,9 @@   ( copyFile,     createDirectoryIfMissing,     doesFileExist,+    getCurrentDirectory,     removePathForcibly,+    setCurrentDirectory,   ) import System.Environment (getExecutablePath) import System.Exit@@ -675,18 +677,15 @@  data Env = Env   { envImgDir :: FilePath,-    -- | Image dir relative to program.-    envRelImgDir :: FilePath,     envOpts :: Options,     envServer :: ScriptServer,     envHash :: T.Text   } -newFileWorker :: Env -> (Maybe FilePath, FilePath) -> (FilePath -> ScriptM ()) -> ScriptM (FilePath, FilePath)-newFileWorker env (fname_desired, template) m = do+newFile :: Env -> (Maybe FilePath, FilePath) -> (FilePath -> ScriptM ()) -> ScriptM FilePath+newFile env (fname_desired, template) m = do   let fname_base = fromMaybe (T.unpack (envHash env) <> "-" <> template) fname_desired       fname = envImgDir env </> fname_base-      fname_rel = envRelImgDir env </> fname_base   exists <- liftIO $ doesFileExist fname   liftIO $ createDirectoryIfMissing True $ envImgDir env   when (exists && scriptVerbose (envOpts env) > 0) $@@ -698,14 +697,11 @@         "Generating new file: " <> T.pack fname     m fname   modify $ \s -> s {stateFiles = S.insert fname $ stateFiles s}-  pure (fname, fname_rel)--newFile :: Env -> (Maybe FilePath, FilePath) -> (FilePath -> ScriptM ()) -> ScriptM FilePath-newFile env f m = snd <$> newFileWorker env f m+  pure fname  newFileContents :: Env -> (Maybe FilePath, FilePath) -> (FilePath -> ScriptM ()) -> ScriptM T.Text newFileContents env f m =-  liftIO . T.readFile . fst =<< newFileWorker env f m+  liftIO . T.readFile =<< newFile env f m  processDirective :: Env -> Directive -> ScriptM T.Text processDirective env (DirectiveBrief d) =@@ -1165,8 +1161,8 @@           system futhark ["hash", prog] mempty        let mdfile = fromMaybe (prog `replaceExtension` "md") $ scriptOutput opts-          imgdir_rel = dropExtension (takeFileName mdfile) <> "-img"-          imgdir = takeDirectory mdfile </> imgdir_rel+          prog_dir = takeDirectory prog+          imgdir = dropExtension (takeFileName mdfile) <> "-img"           run_options = scriptExtraOptions opts           onLine "call" l = T.putStrLn l           onLine _ _ = pure ()@@ -1178,16 +1174,22 @@                     else const . const $ pure ()               } +      orig_dir <- getCurrentDirectory+       withScriptServer cfg $ \server -> do         let env =               Env                 { envServer = server,                   envOpts = opts,                   envHash = proghash,-                  envImgDir = imgdir,-                  envRelImgDir = imgdir_rel+                  envImgDir = imgdir                 }++        when (scriptVerbose opts > 0) $ do+          T.hPutStrLn stderr $ "Executing from " <> T.pack prog_dir+        setCurrentDirectory prog_dir+         (failure, md) <- processScript env script-        T.writeFile mdfile md+        T.writeFile (orig_dir </> mdfile) md         when (failure == Failure) exitFailure     _ -> Nothing

@@ -107,7 +107,9 @@       [ "You're welcome!",         "Tell all your friends about Futhark!",         "Likewise!",-        "And thank you in return for trying the language!"+        "And thank you in return for trying the language!",+        "It's our pleasure!",+        "Have fun with Futhark!"       ]  -- | @futhark tokens@

@@ -29,13 +29,15 @@  mkBoilerplate ::   T.Text ->+  [(Name, KernelConstExp)] ->   M.Map Name KernelSafety ->   [PrimType] ->   [FailureMsg] ->   GC.CompilerM OpenCL () ()-mkBoilerplate cuda_program kernels types failures = do+mkBoilerplate cuda_program macros kernels types failures = do   generateGPUBoilerplate     cuda_program+    macros     backendsCudaH     (M.keys kernels)     types@@ -67,7 +69,7 @@                           fprintf(stderr, "%s: %s\n", optarg, strerror(errno));                           exit(1);                         }-                        exit(1);}|]+                        exit(0);}|]            },          Option            { optionLongName = "load-cuda",@@ -113,7 +115,7 @@ compileProg :: (MonadFreshNames m) => T.Text -> Prog GPUMem -> m (ImpGen.Warnings, GC.CParts) compileProg version prog = do   ( ws,-    Program cuda_code cuda_prelude kernels types params failures prog'+    Program cuda_code cuda_prelude macros kernels types params failures prog'     ) <-     ImpGen.compileProg prog   (ws,)@@ -122,7 +124,7 @@       version       params       operations-      (mkBoilerplate (cuda_prelude <> cuda_code) kernels types failures)+      (mkBoilerplate (cuda_prelude <> cuda_code) macros kernels types failures)       cuda_includes       (Space "device", [Space "device", DefaultSpace])       cliOptions

@@ -77,13 +77,15 @@  mkBoilerplate ::   T.Text ->+  [(Name, KernelConstExp)] ->   M.Map Name KernelSafety ->   [PrimType] ->   [FailureMsg] ->   GC.CompilerM OpenCL () ()-mkBoilerplate opencl_program kernels types failures = do+mkBoilerplate opencl_program macros kernels types failures = do   generateGPUBoilerplate     opencl_program+    macros     backendsOpenclH     (M.keys kernels)     types@@ -187,7 +189,7 @@ compileProg :: (MonadFreshNames m) => T.Text -> Prog GPUMem -> m (ImpGen.Warnings, GC.CParts) compileProg version prog = do   ( ws,-    Program opencl_code opencl_prelude kernels types params failures prog'+    Program opencl_code opencl_prelude macros kernels types params failures prog'     ) <-     ImpGen.compileProg prog   (ws,)@@ -196,7 +198,7 @@       version       params       operations-      (mkBoilerplate (opencl_prelude <> opencl_code) kernels types failures)+      (mkBoilerplate (opencl_prelude <> opencl_code) macros kernels types failures)       opencl_includes       (Space "device", [Space "device", DefaultSpace])       cliOptions

@@ -12,12 +12,13 @@ where  import Control.Monad+import Control.Monad.Identity import Data.Bifunctor (bimap) import Data.Map qualified as M import Data.Text qualified as T import Futhark.CodeGen.Backends.GenericC qualified as GC import Futhark.CodeGen.Backends.GenericC.Options-import Futhark.CodeGen.Backends.GenericC.Pretty (idText)+import Futhark.CodeGen.Backends.GenericC.Pretty (expText, idText) import Futhark.CodeGen.Backends.SimpleRep (primStorageType, toStorage) import Futhark.CodeGen.ImpCode.OpenCL import Futhark.CodeGen.RTS.C (gpuH, gpuPrototypesH)@@ -66,9 +67,12 @@           ([C.cinit|&ctx->global_failure_args|], [C.cinit|sizeof(ctx->global_failure_args)|])         ] +getParamByKey :: Name -> C.Exp+getParamByKey key = [C.cexp|*ctx->tuning_params.$id:key|]+ kernelConstToExp :: KernelConst -> C.Exp-kernelConstToExp (SizeConst key) =-  [C.cexp|*ctx->tuning_params.$id:key|]+kernelConstToExp (SizeConst key _) =+  getParamByKey key kernelConstToExp (SizeMaxConst size_class) =   [C.cexp|ctx->$id:field|]   where@@ -133,13 +137,11 @@         )  callKernel :: GC.OpCompiler OpenCL ()-callKernel (GetSize v key) = do-  let e = kernelConstToExp $ SizeConst key-  GC.stm [C.cstm|$id:v = $exp:e;|]+callKernel (GetSize v key) =+  GC.stm [C.cstm|$id:v = $exp:(getParamByKey key);|] callKernel (CmpSizeLe v key x) = do-  let e = kernelConstToExp $ SizeConst key   x' <- GC.compileExp x-  GC.stm [C.cstm|$id:v = $exp:e <= $exp:x';|]+  GC.stm [C.cstm|$id:v = $exp:(getParamByKey key) <= $exp:x';|]   -- Output size information if logging is enabled.  The autotuner   -- depends on the format of this output, so use caution if changing   -- it.@@ -381,16 +383,20 @@                   return strdup("Unknown error.  This is a compiler bug.");                 }|] +compileConstExp :: KernelConstExp -> C.Exp+compileConstExp e = runIdentity $ GC.compilePrimExp (pure . kernelConstToExp) e+ -- | Called after most code has been generated to generate the bulk of -- the boilerplate. generateGPUBoilerplate ::   T.Text ->+  [(Name, KernelConstExp)] ->   T.Text ->   [Name] ->   [PrimType] ->   [FailureMsg] ->   GC.CompilerM OpenCL () ()-generateGPUBoilerplate gpu_program backendH kernels types failures = do+generateGPUBoilerplate gpu_program macros backendH kernels types failures = do   createKernels kernels   let gpu_program_fragments =         -- Some C compilers limit the size of literal strings, so@@ -402,6 +408,13 @@         | FloatType Float64 `elem` types = [C.cexp|1|]         | otherwise = [C.cexp|0|]       max_failure_args = foldl max 0 $ map (errorMsgNumArgs . failureError) failures++      setMacro i (name, e) =+        [C.cstm|{names[$int:i] = $string:(nameToString name);+                 values[$int:i] = $esc:e';}|]+        where+          e' = T.unpack $ expText $ compileConstExp e+   mapM_     GC.earlyDecl     [C.cunit|static const int max_failure_args = $int:max_failure_args;@@ -410,6 +423,17 @@              $esc:(T.unpack gpuPrototypesH)              $esc:(T.unpack backendH)              $esc:(T.unpack gpuH)+             static int gpu_macros(struct futhark_context *ctx, char*** names_out, typename int64_t** values_out) {+               int num_macros = $int:(length macros);+               char** names = malloc(num_macros * sizeof(char*));+               typename int64_t* values = malloc(num_macros * sizeof(int64_t));++               $stms:(zipWith setMacro [(0::Int)..] macros)++               *names_out = names;+               *values_out = values;+               return num_macros;+             }             |]   GC.earlyDecl $ failureMsgFunction failures 

@@ -4,7 +4,6 @@ module Futhark.CodeGen.Backends.GenericC.Code   ( compilePrimExp,     compileExp,-    compileExpToName,     compileCode,     compileDest,     compileArg,@@ -44,15 +43,6 @@       onPart (ErrorVal (FloatType Float64) x) = ("%f",) <$> compileExp x   (formatstrs, formatargs) <- mapAndUnzipM onPart parts   pure (mconcat formatstrs, formatargs)--compileExpToName :: String -> PrimType -> Exp -> CompilerM op s VName-compileExpToName _ _ (LeafExp v _) =-  pure v-compileExpToName desc t e = do-  desc' <- newVName desc-  e' <- compileExp e-  decl [C.cdecl|$ty:(primTypeToCType t) $id:desc' = $e';|]-  pure desc'  compileExp :: Exp -> CompilerM op s C.Exp compileExp = compilePrimExp $ \v -> pure [C.cexp|$id:v|]

@@ -14,7 +14,6 @@     compileCode,     compilePrimValue,     compilePrimType,-    compilePrimTypeExt,     compilePrimToNp,     compilePrimToExtNp,     fromStorage,@@ -1026,24 +1025,6 @@     FloatType Float64 -> "ct.c_double"     Imp.Bool -> "ct.c_bool"     Unit -> "ct.c_bool"---- | The ctypes type corresponding to a 'PrimType', taking sign into account.-compilePrimTypeExt :: PrimType -> Imp.Signedness -> String-compilePrimTypeExt t ept =-  case (t, ept) of-    (IntType Int8, Imp.Unsigned) -> "ct.c_uint8"-    (IntType Int16, Imp.Unsigned) -> "ct.c_uint16"-    (IntType Int32, Imp.Unsigned) -> "ct.c_uint32"-    (IntType Int64, Imp.Unsigned) -> "ct.c_uint64"-    (IntType Int8, _) -> "ct.c_int8"-    (IntType Int16, _) -> "ct.c_int16"-    (IntType Int32, _) -> "ct.c_int32"-    (IntType Int64, _) -> "ct.c_int64"-    (FloatType Float16, _) -> "ct.c_uint16"-    (FloatType Float32, _) -> "ct.c_float"-    (FloatType Float64, _) -> "ct.c_double"-    (Imp.Bool, _) -> "ct.c_bool"-    (Unit, _) -> "ct.c_byte"  -- | The Numpy type corresponding to a 'PrimType'. compilePrimToNp :: Imp.PrimType -> String

@@ -29,13 +29,15 @@  mkBoilerplate ::   T.Text ->+  [(Name, KernelConstExp)] ->   M.Map Name KernelSafety ->   [PrimType] ->   [FailureMsg] ->   GC.CompilerM OpenCL () ()-mkBoilerplate hip_program kernels types failures = do+mkBoilerplate hip_program macros kernels types failures = do   generateGPUBoilerplate     hip_program+    macros     backendsHipH     (M.keys kernels)     types@@ -95,7 +97,7 @@ compileProg :: (MonadFreshNames m) => T.Text -> Prog GPUMem -> m (ImpGen.Warnings, GC.CParts) compileProg version prog = do   ( ws,-    Program hip_code hip_prelude kernels types params failures prog'+    Program hip_code hip_prelude macros kernels types params failures prog'     ) <-     ImpGen.compileProg prog   (ws,)@@ -104,7 +106,7 @@       version       params       operations-      (mkBoilerplate (hip_prelude <> hip_code) kernels types failures)+      (mkBoilerplate (hip_prelude <> hip_code) macros kernels types failures)       hip_includes       (Space "device", [Space "device", DefaultSpace])       cliOptions

@@ -33,6 +33,7 @@     Imp.Program       opencl_code       opencl_prelude+      macros       kernels       types       sizes@@ -89,7 +90,7 @@             "default_threshold=default_threshold",             "sizes=sizes"           ]-          [Escape $ openClInit types assign sizes failures]+          [Escape $ openClInit macros types assign sizes failures]       options =         [ Option             { optionLongName = "platform",@@ -202,9 +203,12 @@ asLong :: PyExp -> PyExp asLong x = simpleCall "np.int64" [x] +getParamByKey :: Name -> PyExp+getParamByKey key = Index (Var "self.sizes") (IdxExp $ String $ prettyText key)+ kernelConstToExp :: Imp.KernelConst -> PyExp-kernelConstToExp (Imp.SizeConst key) =-  Index (Var "self.sizes") (IdxExp $ String $ prettyText key)+kernelConstToExp (Imp.SizeConst key _) =+  getParamByKey key kernelConstToExp (Imp.SizeMaxConst size_class) =   Var $ "self.max_" <> prettyString size_class @@ -215,13 +219,11 @@ callKernel :: OpCompiler Imp.OpenCL () callKernel (Imp.GetSize v key) = do   v' <- compileVar v-  stm $ Assign v' $ kernelConstToExp $ Imp.SizeConst key+  stm $ Assign v' $ getParamByKey key callKernel (Imp.CmpSizeLe v key x) = do   v' <- compileVar v   x' <- compileExp x-  stm $-    Assign v' $-      BinOp "<=" (kernelConstToExp (Imp.SizeConst key)) x'+  stm $ Assign v' $ BinOp "<=" (getParamByKey key) x' callKernel (Imp.GetSizeMax v size_class) = do   v' <- compileVar v   stm $ Assign v' $ kernelConstToExp $ Imp.SizeMaxConst size_class

@@ -15,6 +15,8 @@   ( ErrorMsg (..),     ErrorMsgPart (..),     FailureMsg (..),+    KernelConst (..),+    KernelConstExp,     ParamMap,     PrimType (..),     errorMsgArgTypes,@@ -28,15 +30,28 @@ errorMsgNumArgs :: ErrorMsg a -> Int errorMsgNumArgs = length . errorMsgArgTypes +getParamByKey :: Name -> PyExp+getParamByKey key = Index (Var "self.sizes") (IdxExp $ String $ prettyText key)++kernelConstToExp :: KernelConst -> PyExp+kernelConstToExp (SizeConst key _) =+  getParamByKey key+kernelConstToExp (SizeMaxConst size_class) =+  Var $ "self.max_" <> prettyString size_class++compileConstExp :: KernelConstExp -> PyExp+compileConstExp e = runIdentity $ Py.compilePrimExp (pure . kernelConstToExp) e+ -- | Python code (as a string) that calls the -- @initiatialize_opencl_object@ procedure.  Should be put in the -- class constructor.-openClInit :: [PrimType] -> String -> ParamMap -> [FailureMsg] -> T.Text-openClInit types assign sizes failures =+openClInit :: [(Name, KernelConstExp)] -> [PrimType] -> String -> ParamMap -> [FailureMsg] -> T.Text+openClInit constants types assign sizes failures =   [text| size_heuristics=$size_heuristics self.global_failure_args_max = $max_num_args self.failure_msgs=$failure_msgs+constants = $constants' program = initialise_opencl_object(self,                                    program_src=fut_opencl_src,                                    build_options=build_options,@@ -52,7 +67,8 @@                                    size_heuristics=size_heuristics,                                    required_types=$types',                                    user_sizes=sizes,-                                   all_sizes=$sizes')+                                   all_sizes=$sizes',+                                   constants=constants) $assign' |]   where@@ -62,6 +78,12 @@     sizes' = prettyText $ sizeClassesToPython sizes     max_num_args = prettyText $ foldl max 0 $ map (errorMsgNumArgs . failureError) failures     failure_msgs = prettyText $ List $ map formatFailure failures+    onConstant (name, e) =+      Tuple+        [ String (nameToText name),+          Lambda "" (compileConstExp e)+        ]+    constants' = prettyText $ List $ map onConstant constants  formatFailure :: FailureMsg -> PyExp formatFailure (FailureMsg (ErrorMsg parts) backtrace) =

@@ -34,7 +34,7 @@  -- | A run-time constant related to kernels. data KernelConst-  = SizeConst Name+  = SizeConst Name SizeClass   | SizeMaxConst SizeClass   deriving (Eq, Ord, Show) @@ -62,15 +62,14 @@     -- | A short descriptive and _unique_ name - should be     -- alphanumeric and without spaces.     kernelName :: Name,-    -- | If true, this kernel does not need to check-    -- whether we are in a failing state, as it can cope.-    -- Intuitively, it means that the kernel does not-    -- depend on any non-scalar parameters to make control-    -- flow decisions.  Replication, transpose, and copy+    -- | If true, this kernel does not need to check whether we are in+    -- a failing state, as it can cope. Intuitively, it means that the+    -- kernel does not depend on any non-scalar parameters to make+    -- control flow decisions. Replication, transpose, and copy     -- kernels are examples of this.     kernelFailureTolerant :: Bool,     -- | If true, multi-versioning branches will consider this kernel-    -- when considering the local memory requirements.  Set this to+    -- when considering the local memory requirements. Set this to     -- false for kernels that do their own checking.     kernelCheckLocalMemory :: Bool   }@@ -86,11 +85,13 @@   deriving (Eq, Ord, Show)  instance Pretty KernelConst where-  pretty (SizeConst key) = "get_size" <> parens (pretty key)-  pretty (SizeMaxConst size_class) = "get_max_size" <> parens (pretty size_class)+  pretty (SizeConst key size_class) =+    "get_size" <> parens (commasep [pretty key, pretty size_class])+  pretty (SizeMaxConst size_class) =+    "get_max_size" <> parens (pretty size_class)  instance FreeIn KernelConst where-  freeIn' (SizeConst _) = mempty+  freeIn' SizeConst {} = mempty   freeIn' (SizeMaxConst _) = mempty  instance Pretty KernelUse where

@@ -18,6 +18,7 @@     FailureMsg (..),     GroupDim,     KernelConst (..),+    KernelConstExp,     module Futhark.CodeGen.ImpCode,     module Futhark.IR.GPU.Sizes,   )@@ -26,7 +27,7 @@ import Data.Map qualified as M import Data.Text qualified as T import Futhark.CodeGen.ImpCode-import Futhark.CodeGen.ImpCode.GPU (GroupDim, KernelConst (..))+import Futhark.CodeGen.ImpCode.GPU (GroupDim, KernelConst (..), KernelConstExp) import Futhark.IR.GPU.Sizes import Futhark.Util.Pretty @@ -35,6 +36,9 @@   { openClProgram :: T.Text,     -- | Must be prepended to the program.     openClPrelude :: T.Text,+    -- | Definitions to be passed as macro definitions to the kernel+    -- compiler.+    openClMacroDefs :: [(Name, KernelConstExp)],     openClKernelNames :: M.Map KernelName KernelSafety,     -- | So we can detect whether the device is capable.     openClUsedTypes :: [PrimType],

@@ -38,7 +38,6 @@     hasFunction,     collect,     collect',-    comment,     VarEntry (..),     ArrayEntry (..), @@ -76,6 +75,7 @@     caseMatch,      -- * Constructing code.+    newVName,     dLParams,     dFParams,     addLoopVar,@@ -391,13 +391,6 @@   modify $ \s -> s {stateCode = prev_code}   pure (x, new_code) --- | Execute a code generation action, wrapping the generated code--- within a 'Imp.Comment' with the given description.-comment :: T.Text -> ImpM rep r op () -> ImpM rep r op ()-comment desc m = do-  code <- collect m-  emit $ Imp.Comment desc code- -- | Emit some generated imperative code. emit :: Imp.Code op -> ImpM rep r op () emit code = modify $ \s -> s {stateCode = stateCode s <> code}@@ -870,7 +863,7 @@   copyDWIM (patElemName pe) [] se []   case op of     OpaqueNil -> pure ()-    OpaqueTrace s -> comment ("Trace: " <> s) $ do+    OpaqueTrace s -> sComment ("Trace: " <> s) $ do       se_t <- subExpType se       case se_t of         Prim t -> tracePrim s t se@@ -937,7 +930,7 @@       dPrimV "x" . TPrimExp $         BinOpExp (Add it OverflowUndef) e' $           BinOpExp (Mul it OverflowUndef) i' s'-    copyDWIM (patElemName pe) [DimFix i] (Var (tvVar x)) []+    copyDWIMFix (patElemName pe) [i] (Var (tvVar x)) [] defCompileBasicOp (Pat [pe]) (Manifest _ src) =   copyDWIM (patElemName pe) [] (Var src) [] defCompileBasicOp (Pat [pe]) (Concat i (x :| ys) _) = do@@ -967,7 +960,7 @@       copy t dest_mem static_src   | otherwise =       forM_ (zip [0 ..] es) $ \(i, e) ->-        copyDWIM (patElemName pe) [DimFix $ fromInteger i] e []+        copyDWIMFix (patElemName pe) [fromInteger i] e []   where     isLiteral (Constant v) = Just v     isLiteral _ = Nothing@@ -1665,6 +1658,8 @@   body' <- collect body   emit $ Imp.While cond body' +-- | Execute a code generation action, wrapping the generated code+-- within a 'Imp.Comment' with the given description. sComment :: T.Text -> ImpM rep r op () -> ImpM rep r op () sComment s code = do   code' <- collect code

@@ -27,7 +27,7 @@ import Futhark.Error import Futhark.IR.GPUMem import Futhark.MonadFreshNames-import Futhark.Util.IntegralExp (divUp, rem)+import Futhark.Util.IntegralExp (divUp, nextMul) import Prelude hiding (quot, rem)  callKernelOperations :: Operations GPUMem HostEnv Imp.HostOp@@ -134,7 +134,7 @@   sKernelThread "gpuseq" tid (defKernelAttrs one one) $     compileStms (freeIn res) stms $       forM_ (zip pes res) $ \(pe, SubExpRes _ se) ->-        copyDWIM (patElemName pe) [DimFix 0] se []+        copyDWIMFix (patElemName pe) [0] se [] opCompiler pat e =   compilerBugS $     "opCompiler: Invalid pattern\n  "@@ -170,14 +170,13 @@ -- otherwise protected by their own multi-versioning branches deeper -- down.  Currently the compiler will not generate multi-versioning -- that makes this a problem, but it might in the future.-checkLocalMemoryReqs :: Imp.HostCode -> CallKernelGen (Maybe (Imp.TExp Bool))-checkLocalMemoryReqs code = do-  scope <- askScope+checkLocalMemoryReqs :: (VName -> Bool) -> Imp.HostCode -> CallKernelGen (Maybe (Imp.TExp Bool))+checkLocalMemoryReqs in_scope code = do   let alloc_sizes = map (sum . map alignedSize . localAllocSizes . Imp.kernelBody) $ getGPU code    -- If any of the sizes involve a variable that is not known at this   -- point, then we cannot check the requirements.-  if any (`M.notMember` scope) (namesToList $ freeIn alloc_sizes)+  if not $ all in_scope $ namesToList $ freeIn alloc_sizes     then pure Nothing     else do       local_memory_capacity :: TV Int32 <- dPrim "local_memory_capacity" int32@@ -200,7 +199,7 @@     -- These allocations will actually be padded to an 8-byte aligned     -- size, so we should take that into account when checking whether     -- they fit.-    alignedSize x = x + ((8 - (x `rem` 8)) `rem` 8)+    alignedSize x = nextMul x 8  withAcc ::   Pat LetDecMem ->@@ -250,9 +249,10 @@ -- always be safe (and what would we do if none of the branches would -- work?). expCompiler dest (Match cond (first_case : cases) defbranch sort@(MatchDec _ MatchEquiv)) = do+  scope <- askScope   tcode <- collect $ compileBody dest $ caseBody first_case   fcode <- collect $ expCompiler dest $ Match cond cases defbranch sort-  check <- checkLocalMemoryReqs tcode+  check <- checkLocalMemoryReqs (`M.member` scope) tcode   let matches = caseMatch cond (casePat first_case)   emit $ case check of     Nothing -> fcode

@@ -29,6 +29,9 @@     fenceForArrays,     updateAcc,     genZeroes,+    isPrimParam,+    kernelConstToExp,+    getChunkSize,      -- * Host-level bulk operations     sReplicate,@@ -52,7 +55,6 @@ import Futhark.Error import Futhark.IR.GPUMem import Futhark.IR.Mem.LMAD qualified as LMAD-import Futhark.MonadFreshNames import Futhark.Transform.Rename import Futhark.Util (dropLast, nubOrd, splitFromEnd) import Futhark.Util.IntegralExp (divUp, quot, rem)@@ -257,6 +259,44 @@         . entryArrayLoc         =<< lookupArray arr +isPrimParam :: (Typed p) => Param p -> Bool+isPrimParam = primType . paramType++kernelConstToExp :: Imp.KernelConstExp -> CallKernelGen Imp.Exp+kernelConstToExp = traverse f+  where+    f (Imp.SizeMaxConst c) = do+      v <- dPrim (prettyString c) int64+      sOp $ Imp.GetSizeMax (tvVar v) c+      pure $ tvVar v+    f (Imp.SizeConst k c) = do+      v <- dPrim (nameToString k) int64+      sOp $ Imp.GetSize (tvVar v) k c+      pure $ tvVar v++-- | Given available register and a list of parameter types, compute+-- the largest available chunk size given the parameters for which we+-- want chunking and the available resources. Used in+-- 'SegScan.SinglePass.compileSegScan', and 'SegRed.compileSegRed'+-- (with primitive non-commutative operators only).+getChunkSize :: [Type] -> Imp.KernelConstExp+getChunkSize types = do+  let max_group_size = Imp.SizeMaxConst SizeGroup+      max_group_mem = Imp.SizeMaxConst SizeLocalMemory+      max_group_reg = Imp.SizeMaxConst SizeRegisters+      k_mem = le64 max_group_mem `quot` le64 max_group_size+      k_reg = le64 max_group_reg `quot` le64 max_group_size+      types' = map elemType $ filter primType types+      sizes = map primByteSize types'++      sum_sizes = sum sizes+      sum_sizes' = sum (map (sMax64 4 . primByteSize) types') `quot` 4+      max_size = maximum sizes++      mem_constraint = max k_mem sum_sizes `quot` max_size+      reg_constraint = (k_reg - 1 - sum_sizes') `quot` (2 * sum_sizes')+  untyped $ sMax64 1 $ sMin64 mem_constraint reg_constraint+ inBlockScan ::   KernelConstants ->   Maybe (Imp.TExp Int32 -> Imp.TExp Int32 -> Imp.TExp Bool) ->@@ -275,7 +315,7 @@         splitAt (length actual_params `div` 2) actual_params       y_to_x =         forM_ (zip x_params y_params) $ \(x, y) ->-          when (primType (paramType x)) $+          when (isPrimParam x) $             copyDWIM (paramName x) [] (Var (paramName y)) []    -- Set initial y values@@ -342,23 +382,21 @@     array_scan = not $ all primType $ lambdaReturnType scan_lam      readInitial p arr-      | primType $ paramType p =-          copyDWIM (paramName p) [] (Var arr) [DimFix ltid]+      | isPrimParam p =+          copyDWIMFix (paramName p) [] (Var arr) [ltid]       | otherwise =-          copyDWIM (paramName p) [] (Var arr) [DimFix gtid]+          copyDWIMFix (paramName p) [] (Var arr) [gtid]      readParam behind p arr-      | primType $ paramType p =-          copyDWIM (paramName p) [] (Var arr) [DimFix $ ltid - behind]+      | isPrimParam p =+          copyDWIMFix (paramName p) [] (Var arr) [ltid - behind]       | otherwise =-          copyDWIM (paramName p) [] (Var arr) [DimFix $ gtid - behind + arrs_full_size]+          copyDWIMFix (paramName p) [] (Var arr) [gtid - behind + arrs_full_size] -    writeResult x y arr-      | primType $ paramType x = do-          copyDWIM arr [DimFix ltid] (Var $ paramName x) []-          copyDWIM (paramName y) [] (Var $ paramName x) []-      | otherwise =-          copyDWIM (paramName y) [] (Var $ paramName x) []+    writeResult x y arr = do+      when (isPrimParam x) $+        copyDWIMFix arr [ltid] (Var $ paramName x) []+      copyDWIM (paramName y) [] (Var $ paramName x) []  groupScan ::   Maybe (Imp.TExp Int32 -> Imp.TExp Int32 -> Imp.TExp Bool) ->@@ -421,16 +459,16 @@       group_offset = sExt64 (kernelGroupId constants) * kernelGroupSize constants        writeBlockResult p arr-        | primType $ paramType p =-            copyDWIM arr [DimFix $ sExt64 block_id] (Var $ paramName p) []+        | isPrimParam p =+            copyDWIMFix arr [sExt64 block_id] (Var $ paramName p) []         | otherwise =-            copyDWIM arr [DimFix $ group_offset + sExt64 block_id] (Var $ paramName p) []+            copyDWIMFix arr [group_offset + sExt64 block_id] (Var $ paramName p) []        readPrevBlockResult p arr-        | primType $ paramType p =-            copyDWIM (paramName p) [] (Var arr) [DimFix $ sExt64 block_id - 1]+        | isPrimParam p =+            copyDWIMFix (paramName p) [] (Var arr) [sExt64 block_id - 1]         | otherwise =-            copyDWIM (paramName p) [] (Var arr) [DimFix $ group_offset + sExt64 block_id - 1]+            copyDWIMFix (paramName p) [] (Var arr) [group_offset + sExt64 block_id - 1]    doInBlockScan seg_flag ltid_in_bounds lam   barrier@@ -440,8 +478,8 @@     sComment "save correct values for first block" $       sWhen is_first_block $         forM_ (zip x_params arrs) $ \(x, arr) ->-          unless (primType $ paramType x) $-            copyDWIM arr [DimFix $ arrs_full_size + group_offset + sExt64 block_size + ltid] (Var $ paramName x) []+          unless (isPrimParam x) $+            copyDWIMFix arr [arrs_full_size + group_offset + sExt64 block_size + ltid] (Var $ paramName x) []      barrier @@ -457,7 +495,7 @@         flag_true <- seg_flag         Just $ \from to ->           flag_true (from * block_size + block_size - 1) (to * block_size + block_size - 1)-  comment+  sComment     "scan the first block, after which offset 'i' contains carry-in for block 'i+1'"     $ doInBlockScan first_block_seg_flag (is_first_block .&&. ltid_in_bounds) renamed_lam @@ -467,12 +505,12 @@     sComment "move correct values for first block back a block" $       sWhen is_first_block $         forM_ (zip x_params arrs) $ \(x, arr) ->-          unless (primType $ paramType x) $-            copyDWIM+          unless (isPrimParam x) $+            copyDWIMFix               arr-              [DimFix $ arrs_full_size + group_offset + ltid]+              [arrs_full_size + group_offset + ltid]               (Var arr)-              [DimFix $ arrs_full_size + group_offset + sExt64 block_size + ltid]+              [arrs_full_size + group_offset + sExt64 block_size + ltid]      barrier @@ -498,8 +536,8 @@        write_final_result =         forM_ (zip x_params arrs) $ \(p, arr) ->-          when (primType $ paramType p) $-            copyDWIM arr [DimFix ltid] (Var $ paramName p) []+          when (isPrimParam p) $+            copyDWIMFix arr [ltid] (Var $ paramName p) []    sComment "carry-in for every block except the first" $     localOps threadOperations $ do@@ -512,9 +550,9 @@   sComment "restore correct values for first block" $     sWhen (is_first_block .&&. ltid_in_bounds) $       forM_ (zip3 x_params y_params arrs) $ \(x, y, arr) ->-        if primType (paramType y)-          then copyDWIM arr [DimFix ltid] (Var $ paramName y) []-          else copyDWIM (paramName x) [] (Var arr) [DimFix $ arrs_full_size + group_offset + ltid]+        if isPrimParam y+          then copyDWIMFix arr [ltid] (Var $ paramName y) []+          else copyDWIMFix (paramName x) [] (Var arr) [arrs_full_size + group_offset + ltid]    barrier @@ -550,17 +588,13 @@         let i = local_tid + tvExp offset         copyDWIMFix (paramName param) [] (Var arr) [sExt64 i] -      writeReduceOpResult param arr-        | Prim _ <- paramType param =-            copyDWIMFix arr [sExt64 local_tid] (Var $ paramName param) []-        | otherwise =-            pure ()+      writeReduceOpResult param arr =+        when (isPrimParam param) $+          copyDWIMFix arr [sExt64 local_tid] (Var $ paramName param) [] -      writeArrayOpResult param arr-        | Prim _ <- paramType param =-            pure ()-        | otherwise =-            copyDWIMFix arr [0] (Var $ paramName param) []+      writeArrayOpResult param arr =+        unless (isPrimParam param) $+          copyDWIMFix arr [sExt64 local_tid] (Var $ paramName param) []    let (reduce_acc_params, reduce_arr_params) =         splitAt (length arrs) $ lambdaParams lam@@ -570,17 +604,17 @@    offset <-- (0 :: Imp.TExp Int32) -  comment "participating threads read initial accumulator" $+  sComment "participating threads read initial accumulator" $     localOps threadOperations . sWhen (local_tid .<. w) $       zipWithM_ readReduceArgument reduce_acc_params arrs    let do_reduce = localOps threadOperations $ do-        comment "read array element" $+        sComment "read array element" $           zipWithM_ readReduceArgument reduce_arr_params arrs-        comment "apply reduction operation" $+        sComment "apply reduction operation" $           compileBody' reduce_acc_params $             lambdaBody lam-        comment "write result of operation" $+        sComment "write result of operation" $           zipWithM_ writeReduceOpResult reduce_acc_params arrs       in_wave_reduce = everythingVolatile do_reduce @@ -622,10 +656,11 @@   cross_wave_reductions   errorsync -  sComment "Copy array-typed operands to result array" $ do-    sWhen (local_tid .==. 0) $-      localOps threadOperations $-        zipWithM_ writeArrayOpResult reduce_acc_params arrs+  unless (all isPrimParam reduce_acc_params) $+    sComment "Copy array-typed operands to result array" $+      sWhen (local_tid .==. 0) $+        localOps threadOperations $+          zipWithM_ writeArrayOpResult reduce_acc_params arrs  compileThreadOp :: OpCompiler GPUMem KernelEnv Imp.KernelOp compileThreadOp pat (Alloc size space) =@@ -900,10 +935,10 @@   let onLeaf name _ = lookupConstExp name       lookupConstExp name =         constExp =<< hasExp =<< M.lookup name vtable-      constExp (Op (Inner (SizeOp (GetSize key _)))) =-        Just $ LeafExp (Imp.SizeConst $ keyWithEntryPoint fname key) int32-      constExp (Op (Inner (SizeOp (GetSizeMax size_class)))) =-        Just $ LeafExp (Imp.SizeMaxConst size_class) int32+      constExp (Op (Inner (SizeOp (GetSize key c)))) =+        Just $ LeafExp (Imp.SizeConst (keyWithEntryPoint fname key) c) int32+      constExp (Op (Inner (SizeOp (GetSizeMax c)))) =+        Just $ LeafExp (Imp.SizeMaxConst c) int32       constExp e = primExpFromExp lookupConstExp e   pure $ replaceInPrimExpM onLeaf size   where@@ -1092,7 +1127,12 @@     -- | Number of groups.     kAttrNumGroups :: Count NumGroups SubExp,     -- | Group size.-    kAttrGroupSize :: Count GroupSize SubExp+    kAttrGroupSize :: Count GroupSize SubExp,+    -- | Variables that are specially in scope inside the kernel.+    -- Operationally, these will be available at kernel compile time+    -- (which happens at run-time, with access to machine-specific+    -- information).+    kAttrConstExps :: M.Map VName Imp.KernelConstExp   }  -- | The default kernel attributes.@@ -1105,7 +1145,8 @@     { kAttrFailureTolerant = False,       kAttrCheckLocalMemory = True,       kAttrNumGroups = num_groups,-      kAttrGroupSize = group_size+      kAttrGroupSize = group_size,+      kAttrConstExps = mempty     }  getSize :: String -> SizeClass -> CallKernelGen (TV Int64)@@ -1170,12 +1211,12 @@ sKernelOp attrs constants ops name m = do   HostEnv atomics _ locks <- askEnv   body <- makeAllMemoryGlobal $ subImpM_ (KernelEnv atomics constants locks) ops m-  uses <- computeKernelUses body mempty+  uses <- computeKernelUses body $ M.keys $ kAttrConstExps attrs   group_size <- onGroupSize $ kernelGroupSize constants   emit . Imp.Op . Imp.CallKernel $     Imp.Kernel       { Imp.kernelBody = body,-        Imp.kernelUses = uses,+        Imp.kernelUses = uses <> map constToUse (M.toList (kAttrConstExps attrs)),         Imp.kernelNumGroups = [untyped $ kernelNumGroups constants],         Imp.kernelGroupSize = [group_size],         Imp.kernelName = name,@@ -1192,6 +1233,8 @@         case x of           Just (LeafExp kc _) -> Right kc           _ -> Left $ untyped e++    constToUse (v, e) = Imp.ConstUse v e  sKernelFailureTolerant ::   Bool ->

@@ -27,7 +27,6 @@ import Futhark.Error import Futhark.IR.GPUMem import Futhark.IR.Mem.LMAD qualified as LMAD-import Futhark.MonadFreshNames import Futhark.Transform.Rename import Futhark.Util (chunks, mapAccumLM, takeLast) import Futhark.Util.IntegralExp (divUp, rem)

@@ -48,8 +48,8 @@ import Futhark.Construct (fullSliceNum) import Futhark.IR.GPUMem import Futhark.IR.Mem.LMAD qualified as LMAD-import Futhark.MonadFreshNames import Futhark.Pass.ExplicitAllocations ()+import Futhark.Transform.Substitute import Futhark.Util (chunks, mapAccumLM, maxinum, splitFromEnd, takeLast) import Futhark.Util.IntegralExp (divUp, quot, rem) import Prelude hiding (quot, rem)@@ -242,17 +242,9 @@           t64 $             r64 hist_T / hist_C_max -  -- Querying L2 cache size is not reliable.  Instead we provide a-  -- tunable knob with a hopefully sane default.-  let hist_L2_def = 4 * 1024 * 1024   hist_L2 <- dPrim "L2_size" int32-  entry <- askFunction   -- Equivalent to F_L2*L2 in paper.-  sOp-    $ Imp.GetSize-      (tvVar hist_L2)-      (keyWithEntryPoint entry $ nameFromString (prettyString (tvVar hist_L2)))-    $ Imp.SizeBespoke (nameFromString "L2_for_histogram") hist_L2_def+  sOp $ Imp.GetSizeMax (tvVar hist_L2) Imp.SizeCache    let hist_L2_ln_sz = 16 * 4 -- L2 cache line size approximation   hist_RACE_exp <-@@ -1073,7 +1065,8 @@   KernelBody GPUMem ->   CallKernelGen () compileSegHist (Pat pes) lvl space ops kbody = do-  KernelAttrs _ _ num_groups group_size <- lvlKernelAttrs lvl+  KernelAttrs {kAttrNumGroups = num_groups, kAttrGroupSize = group_size} <-+    lvlKernelAttrs lvl   -- Most of this function is not the histogram part itself, but   -- rather figuring out whether to use a local or global memory   -- strategy, as well as collapsing the subhistograms produced (which@@ -1180,8 +1173,11 @@                   ++ zip bucket_ids (shapeDims (histShape op))                   ++ zip vector_ids (shapeDims $ histOpShape op)                   ++ [(subhistogram_id, Var $ tvVar num_histos)]+            -- The operator may have references to the old flat thread+            -- ID, which we must update to point at the new one.+            subst = M.singleton (segFlat space) flat_gtid -        let segred_op = SegBinOp Commutative (histOp op) (histNeutral op) mempty+        let segred_op = SegBinOp Commutative (substituteNames subst $ histOp op) (histNeutral op) mempty         compileSegRed' (Pat red_pes) grid segred_space [segred_op] $ \red_cont ->           red_cont . flip map subhistos $ \subhisto ->             ( Var subhisto,

@@ -41,845 +41,1035 @@ --   cases, we would allocate a /whole/ group per segment with the --   large strategy, but at most 50% of the threads in the group would --   have any element to read, which becomes highly inefficient.-module Futhark.CodeGen.ImpGen.GPU.SegRed-  ( compileSegRed,-    compileSegRed',-    DoSegBody,-  )-where--import Control.Monad-import Data.List (genericLength, zip7)-import Data.Maybe-import Futhark.CodeGen.ImpCode.GPU qualified as Imp-import Futhark.CodeGen.ImpGen-import Futhark.CodeGen.ImpGen.GPU.Base-import Futhark.Error-import Futhark.IR.GPUMem-import Futhark.IR.Mem.LMAD qualified as LMAD-import Futhark.Transform.Rename-import Futhark.Util (chunks)-import Futhark.Util.IntegralExp (divUp, quot, rem)-import Prelude hiding (quot, rem)---- | The maximum number of operators we support in a single SegRed.--- This limit arises out of the static allocation of counters.-maxNumOps :: Int32-maxNumOps = 10---- | Code generation for the body of the SegRed, taking a continuation--- for saving the results of the body.  The results should be--- represented as a pairing of a t'SubExp' along with a list of--- indexes into that t'SubExp' for reading the result.-type DoSegBody = ([(SubExp, [Imp.TExp Int64])] -> InKernelGen ()) -> InKernelGen ()---- | Compile 'SegRed' instance to host-level code with calls to--- various kernels.-compileSegRed ::-  Pat LetDecMem ->-  SegLevel ->-  SegSpace ->-  [SegBinOp GPUMem] ->-  KernelBody GPUMem ->-  CallKernelGen ()-compileSegRed pat lvl space reds body = do-  emit $ Imp.DebugPrint "\n# SegRed" Nothing-  KernelAttrs _ _ num_groups group_size <- lvlKernelAttrs lvl-  let grid = KernelGrid num_groups group_size-  compileSegRed' pat grid space reds $ \red_cont ->-    compileStms mempty (kernelBodyStms body) $ do-      let (red_res, map_res) = splitAt (segBinOpResults reds) $ kernelBodyResult body--      sComment "save map-out results" $ do-        let map_arrs = drop (segBinOpResults reds) $ patElems pat-        zipWithM_ (compileThreadResult space) map_arrs map_res--      red_cont $ map ((,[]) . kernelResultSubExp) red_res-  emit $ Imp.DebugPrint "" Nothing---- | Like 'compileSegRed', but where the body is a monadic action.-compileSegRed' ::-  Pat LetDecMem ->-  KernelGrid ->-  SegSpace ->-  [SegBinOp GPUMem] ->-  DoSegBody ->-  CallKernelGen ()-compileSegRed' pat grid space reds body-  | genericLength reds > maxNumOps =-      compilerLimitationS $-        "compileSegRed': at most " ++ show maxNumOps ++ " reduction operators are supported."-  | [(_, Constant (IntValue (Int64Value 1))), _] <- unSegSpace space =-      nonsegmentedReduction pat num_groups group_size space reds body-  | otherwise = do-      let group_size' = pe64 $ unCount group_size-          segment_size = pe64 $ last $ segSpaceDims space-          use_small_segments = segment_size * 2 .<. group_size'-      sIf-        use_small_segments-        (smallSegmentsReduction pat num_groups group_size space reds body)-        (largeSegmentsReduction pat num_groups group_size space reds body)-  where-    num_groups = gridNumGroups grid-    group_size = gridGroupSize grid---- | Prepare intermediate arrays for the reduction.  Prim-typed--- arguments go in local memory (so we need to do the allocation of--- those arrays inside the kernel), while array-typed arguments go in--- global memory.  Allocations for the former have already been--- performed.  This policy is baked into how the allocations are done--- in ExplicitAllocations.-intermediateArrays ::-  Count GroupSize SubExp ->-  SubExp ->-  SegBinOp GPUMem ->-  InKernelGen [VName]-intermediateArrays (Count group_size) num_threads (SegBinOp _ red_op nes _) = do-  let red_op_params = lambdaParams red_op-      (red_acc_params, _) = splitAt (length nes) red_op_params-  forM red_acc_params $ \p ->-    case paramDec p of-      MemArray pt shape _ (ArrayIn mem _) -> do-        let shape' = Shape [num_threads] <> shape-        sArray "red_arr" pt shape' mem $-          LMAD.iota 0 (map pe64 $ shapeDims shape')-      _ -> do-        let pt = elemType $ paramType p-            shape = Shape [group_size]-        sAllocArray "red_arr" pt shape $ Space "local"---- | Arrays for storing group results.------ The group-result arrays have an extra dimension because they are--- also used for keeping vectorised accumulators for first-stage--- reduction, if necessary.  If necessary, this dimension has size--- group_size, and otherwise 1.  When actually storing group results,--- the first index is set to 0.-groupResultArrays ::-  Count NumGroups SubExp ->-  Count GroupSize SubExp ->-  [SegBinOp GPUMem] ->-  CallKernelGen [[VName]]-groupResultArrays (Count virt_num_groups) (Count group_size) reds =-  forM reds $ \(SegBinOp _ lam _ shape) ->-    forM (lambdaReturnType lam) $ \t -> do-      let pt = elemType t-          extra_dim-            | primType t, shapeRank shape == 0 = intConst Int64 1-            | otherwise = group_size-          full_shape = Shape [extra_dim, virt_num_groups] <> shape <> arrayShape t-          -- Move the groupsize dimension last to ensure coalesced-          -- memory access.-          perm = [1 .. shapeRank full_shape - 1] ++ [0]-      sAllocArrayPerm "segred_tmp" pt full_shape (Space "device") perm--nonsegmentedReduction ::-  Pat LetDecMem ->-  Count NumGroups SubExp ->-  Count GroupSize SubExp ->-  SegSpace ->-  [SegBinOp GPUMem] ->-  DoSegBody ->-  CallKernelGen ()-nonsegmentedReduction segred_pat num_groups group_size space reds body = do-  let (gtids, dims) = unzip $ unSegSpace space-      dims' = map pe64 dims-      num_groups' = fmap pe64 num_groups-      group_size' = fmap pe64 group_size-      global_tid = Imp.le64 $ segFlat space-      w = last dims'--  counter <- genZeroes "counters" $ fromIntegral maxNumOps--  reds_group_res_arrs <- groupResultArrays num_groups group_size reds--  num_threads <--    dPrimV "num_threads" $-      unCount num_groups' * unCount group_size'--  sKernelThread "segred_nonseg" (segFlat space) (defKernelAttrs num_groups group_size) $ do-    constants <- kernelConstants <$> askEnv-    sync_arr <- sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local"-    reds_arrs <- mapM (intermediateArrays group_size (tvSize num_threads)) reds--    -- Since this is the nonsegmented case, all outer segment IDs must-    -- necessarily be 0.-    forM_ gtids $ \v -> dPrimV_ v (0 :: Imp.TExp Int64)--    let num_elements = Imp.elements w-        elems_per_thread =-          num_elements-            `divUp` Imp.elements (sExt64 (kernelNumThreads constants))--    slugs <--      mapM (segBinOpSlug (kernelLocalThreadId constants) (kernelGroupId constants)) $-        zip3 reds reds_arrs reds_group_res_arrs-    reds_op_renamed <--      reductionStageOne-        constants-        (zip gtids dims')-        num_elements-        global_tid-        elems_per_thread-        (tvExp num_threads)-        slugs-        body--    let segred_pes =-          chunks (map (length . segBinOpNeutral) reds) $-            patElems segred_pat-    forM_ (zip7 reds reds_arrs reds_group_res_arrs segred_pes slugs reds_op_renamed [0 ..]) $-      \(SegBinOp _ red_op nes _, red_arrs, group_res_arrs, pes, slug, red_op_renamed, i) -> do-        let (red_x_params, red_y_params) = splitAt (length nes) $ lambdaParams red_op-        reductionStageTwo-          constants-          pes-          (kernelGroupId constants)-          0-          [0]-          0-          (sExt64 $ kernelNumGroups constants)-          slug-          red_x_params-          red_y_params-          red_op_renamed-          nes-          1-          counter-          (fromInteger i)-          sync_arr-          group_res_arrs-          red_arrs--smallSegmentsReduction ::-  Pat LetDecMem ->-  Count NumGroups SubExp ->-  Count GroupSize SubExp ->-  SegSpace ->-  [SegBinOp GPUMem] ->-  DoSegBody ->-  CallKernelGen ()-smallSegmentsReduction (Pat segred_pes) num_groups group_size space reds body = do-  let (gtids, dims) = unzip $ unSegSpace space-      dims' = map pe64 dims-      segment_size = last dims'--  -- Careful to avoid division by zero now.-  segment_size_nonzero <--    dPrimVE "segment_size_nonzero" $ sMax64 1 segment_size--  let num_groups' = fmap pe64 num_groups-      group_size' = fmap pe64 group_size-  num_threads <- dPrimV "num_threads" $ unCount num_groups' * unCount group_size'-  let num_segments = product $ init dims'-      segments_per_group = unCount group_size' `quot` segment_size_nonzero-      required_groups = sExt32 $ num_segments `divUp` segments_per_group--  emit $ Imp.DebugPrint "# SegRed-small" Nothing-  emit $ Imp.DebugPrint "num_segments" $ Just $ untyped num_segments-  emit $ Imp.DebugPrint "segment_size" $ Just $ untyped segment_size-  emit $ Imp.DebugPrint "segments_per_group" $ Just $ untyped segments_per_group-  emit $ Imp.DebugPrint "required_groups" $ Just $ untyped required_groups--  sKernelThread "segred_small" (segFlat space) (defKernelAttrs num_groups group_size) $ do-    constants <- kernelConstants <$> askEnv-    reds_arrs <- mapM (intermediateArrays group_size (Var $ tvVar num_threads)) reds--    -- We probably do not have enough actual workgroups to cover the-    -- entire iteration space.  Some groups thus have to perform double-    -- duty; we put an outer loop to accomplish this.-    virtualiseGroups SegVirt required_groups $ \group_id' -> do-      -- Compute the 'n' input indices.  The outer 'n-1' correspond to-      -- the segment ID, and are computed from the group id.  The inner-      -- is computed from the local thread id, and may be out-of-bounds.-      let ltid = sExt64 $ kernelLocalThreadId constants-          segment_index =-            (ltid `quot` segment_size_nonzero)-              + (sExt64 group_id' * sExt64 segments_per_group)-          index_within_segment = ltid `rem` segment_size--      dIndexSpace (zip (init gtids) (init dims')) segment_index-      dPrimV_ (last gtids) index_within_segment--      let out_of_bounds =-            forM_ (zip reds reds_arrs) $ \(SegBinOp _ _ nes _, red_arrs) ->-              forM_ (zip red_arrs nes) $ \(arr, ne) ->-                copyDWIMFix arr [ltid] ne []--          in_bounds =-            body $ \red_res ->-              sComment "save results to be reduced" $ do-                let red_dests = map (,[ltid]) (concat reds_arrs)-                forM_ (zip red_dests red_res) $ \((d, d_is), (res, res_is)) ->-                  copyDWIMFix d d_is res res_is--      sComment "apply map function if in bounds" $-        sIf-          ( segment_size-              .>. 0-              .&&. isActive (init $ zip gtids dims)-              .&&. ltid-              .<. segment_size-              * segments_per_group-          )-          in_bounds-          out_of_bounds--      sOp $ Imp.ErrorSync Imp.FenceLocal -- Also implicitly barrier.-      let crossesSegment from to =-            (sExt64 to - sExt64 from) .>. (sExt64 to `rem` segment_size)-      sWhen (segment_size .>. 0) $-        sComment "perform segmented scan to imitate reduction" $-          forM_ (zip reds reds_arrs) $ \(SegBinOp _ red_op _ _, red_arrs) ->-            groupScan-              (Just crossesSegment)-              (sExt64 $ tvExp num_threads)-              (segment_size * segments_per_group)-              red_op-              red_arrs--      sOp $ Imp.Barrier Imp.FenceLocal--      sComment "save final values of segments"-        $ sWhen-          ( sExt64 group_id'-              * segments_per_group-              + sExt64 ltid-                .<. num_segments-                .&&. ltid-                .<. segments_per_group-          )-        $ forM_ (zip segred_pes (concat reds_arrs))-        $ \(pe, arr) -> do-          -- Figure out which segment result this thread should write...-          let flat_segment_index =-                sExt64 group_id' * segments_per_group + sExt64 ltid-              gtids' =-                unflattenIndex (init dims') flat_segment_index-          copyDWIMFix-            (patElemName pe)-            gtids'-            (Var arr)-            [(ltid + 1) * segment_size_nonzero - 1]--      -- Finally another barrier, because we will be writing to the-      -- local memory array first thing in the next iteration.-      sOp $ Imp.Barrier Imp.FenceLocal--largeSegmentsReduction ::-  Pat LetDecMem ->-  Count NumGroups SubExp ->-  Count GroupSize SubExp ->-  SegSpace ->-  [SegBinOp GPUMem] ->-  DoSegBody ->-  CallKernelGen ()-largeSegmentsReduction segred_pat num_groups group_size space reds body = do-  let (gtids, dims) = unzip $ unSegSpace space-      dims' = map pe64 dims-      num_segments = product $ init dims'-      segment_size = last dims'-      num_groups' = fmap pe64 num_groups-      group_size' = fmap pe64 group_size--  (groups_per_segment, elems_per_thread) <--    groupsPerSegmentAndElementsPerThread-      segment_size-      num_segments-      num_groups'-      group_size'-  virt_num_groups <--    dPrimV "virt_num_groups" $-      groups_per_segment * num_segments--  num_threads <--    dPrimV "num_threads" $-      unCount num_groups' * unCount group_size'--  threads_per_segment <--    dPrimV "threads_per_segment" $-      groups_per_segment * unCount group_size'--  emit $ Imp.DebugPrint "# SegRed-large" Nothing-  emit $ Imp.DebugPrint "num_segments" $ Just $ untyped num_segments-  emit $ Imp.DebugPrint "segment_size" $ Just $ untyped segment_size-  emit $ Imp.DebugPrint "virt_num_groups" $ Just $ untyped $ tvExp virt_num_groups-  emit $ Imp.DebugPrint "num_groups" $ Just $ untyped $ Imp.unCount num_groups'-  emit $ Imp.DebugPrint "group_size" $ Just $ untyped $ Imp.unCount group_size'-  emit $ Imp.DebugPrint "elems_per_thread" $ Just $ untyped $ Imp.unCount elems_per_thread-  emit $ Imp.DebugPrint "groups_per_segment" $ Just $ untyped groups_per_segment--  reds_group_res_arrs <- groupResultArrays (Count (tvSize virt_num_groups)) group_size reds--  -- In principle we should have a counter for every segment.  Since-  -- the number of segments is a dynamic quantity, we would have to-  -- allocate and zero out an array here, which is expensive.-  -- However, we exploit the fact that the number of segments being-  -- reduced at any point in time is limited by the number of-  -- workgroups. If we bound the number of workgroups, we can get away-  -- with using that many counters.  FIXME: Is this limit checked-  -- anywhere?  There are other places in the compiler that will fail-  -- if the group count exceeds the maximum group size, which is at-  -- most 1024 anyway.-  let num_counters = fromIntegral maxNumOps * 1024-  counter <- genZeroes "counters" num_counters--  sKernelThread "segred_large" (segFlat space) (defKernelAttrs num_groups group_size) $ do-    constants <- kernelConstants <$> askEnv-    reds_arrs <- mapM (intermediateArrays group_size (tvSize num_threads)) reds-    sync_arr <- sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local"--    -- We probably do not have enough actual workgroups to cover the-    -- entire iteration space.  Some groups thus have to perform double-    -- duty; we put an outer loop to accomplish this.-    virtualiseGroups SegVirt (sExt32 (tvExp virt_num_groups)) $ \group_id -> do-      let segment_gtids = init gtids-          w = last dims-          local_tid = kernelLocalThreadId constants--      flat_segment_id <--        dPrimVE "flat_segment_id" $-          group_id `quot` sExt32 groups_per_segment--      global_tid <--        dPrimVE "global_tid" $-          (sExt64 group_id * sExt64 (unCount group_size') + sExt64 local_tid)-            `rem` (sExt64 (unCount group_size') * groups_per_segment)--      let first_group_for_segment = sExt64 flat_segment_id * groups_per_segment-      dIndexSpace (zip segment_gtids (init dims')) $ sExt64 flat_segment_id-      dPrim_ (last gtids) int64-      let num_elements = Imp.elements $ pe64 w--      slugs <--        mapM (segBinOpSlug local_tid group_id) $-          zip3 reds reds_arrs reds_group_res_arrs-      reds_op_renamed <--        reductionStageOne-          constants-          (zip gtids dims')-          num_elements-          global_tid-          elems_per_thread-          (tvExp threads_per_segment)-          slugs-          body--      let segred_pes =-            chunks (map (length . segBinOpNeutral) reds) $-              patElems segred_pat--          multiple_groups_per_segment =-            forM_ (zip7 reds reds_arrs reds_group_res_arrs segred_pes slugs reds_op_renamed [0 ..]) $-              \(SegBinOp _ red_op nes _, red_arrs, group_res_arrs, pes, slug, red_op_renamed, i) -> do-                let (red_x_params, red_y_params) =-                      splitAt (length nes) $ lambdaParams red_op-                reductionStageTwo-                  constants-                  pes-                  group_id-                  flat_segment_id-                  (map Imp.le64 segment_gtids)-                  (sExt64 first_group_for_segment)-                  groups_per_segment-                  slug-                  red_x_params-                  red_y_params-                  red_op_renamed-                  nes-                  (fromIntegral num_counters)-                  counter-                  (fromInteger i)-                  sync_arr-                  group_res_arrs-                  red_arrs--          one_group_per_segment =-            comment "first thread in group saves final result to memory" $-              forM_ (zip slugs segred_pes) $ \(slug, pes) ->-                sWhen (local_tid .==. 0) $-                  forM_ (zip pes (slugAccs slug)) $ \(v, (acc, acc_is)) ->-                    copyDWIMFix (patElemName v) (map Imp.le64 segment_gtids) (Var acc) acc_is--      sIf (groups_per_segment .==. 1) one_group_per_segment multiple_groups_per_segment---- Careful to avoid division by zero here.  We have at least one group--- per segment.-groupsPerSegmentAndElementsPerThread ::-  Imp.TExp Int64 ->-  Imp.TExp Int64 ->-  Count NumGroups (Imp.TExp Int64) ->-  Count GroupSize (Imp.TExp Int64) ->-  CallKernelGen-    ( Imp.TExp Int64,-      Imp.Count Imp.Elements (Imp.TExp Int64)-    )-groupsPerSegmentAndElementsPerThread segment_size num_segments num_groups_hint group_size = do-  groups_per_segment <--    dPrimVE "groups_per_segment" $-      unCount num_groups_hint `divUp` sMax64 1 num_segments-  elements_per_thread <--    dPrimVE "elements_per_thread" $-      segment_size `divUp` (unCount group_size * groups_per_segment)-  pure (groups_per_segment, Imp.elements elements_per_thread)---- | A SegBinOp with auxiliary information.-data SegBinOpSlug = SegBinOpSlug-  { slugOp :: SegBinOp GPUMem,-    -- | The arrays used for computing the intra-group reduction-    -- (either local or global memory).-    slugArrs :: [VName],-    -- | Places to store accumulator in stage 1 reduction.-    slugAccs :: [(VName, [Imp.TExp Int64])]-  }--slugBody :: SegBinOpSlug -> Body GPUMem-slugBody = lambdaBody . segBinOpLambda . slugOp--slugParams :: SegBinOpSlug -> [LParam GPUMem]-slugParams = lambdaParams . segBinOpLambda . slugOp--slugNeutral :: SegBinOpSlug -> [SubExp]-slugNeutral = segBinOpNeutral . slugOp--slugShape :: SegBinOpSlug -> Shape-slugShape = segBinOpShape . slugOp--slugsComm :: [SegBinOpSlug] -> Commutativity-slugsComm = mconcat . map (segBinOpComm . slugOp)--accParams, nextParams :: SegBinOpSlug -> [LParam GPUMem]-accParams slug = take (length (slugNeutral slug)) $ slugParams slug-nextParams slug = drop (length (slugNeutral slug)) $ slugParams slug--segBinOpSlug :: Imp.TExp Int32 -> Imp.TExp Int32 -> (SegBinOp GPUMem, [VName], [VName]) -> InKernelGen SegBinOpSlug-segBinOpSlug local_tid group_id (op, group_res_arrs, param_arrs) =-  SegBinOpSlug op group_res_arrs-    <$> zipWithM mkAcc (lambdaParams (segBinOpLambda op)) param_arrs-  where-    mkAcc p param_arr-      | Prim t <- paramType p,-        shapeRank (segBinOpShape op) == 0 = do-          acc <- dPrim (baseString (paramName p) <> "_acc") t-          pure (tvVar acc, [])-      | otherwise =-          pure (param_arr, [sExt64 local_tid, sExt64 group_id])--computeThreadChunkSize ::-  Commutativity ->-  Imp.TExp Int64 ->-  Imp.TExp Int64 ->-  Imp.Count Imp.Elements (Imp.TExp Int64) ->-  Imp.Count Imp.Elements (Imp.TExp Int64) ->-  TV Int64 ->-  ImpM rep r op ()-computeThreadChunkSize Commutative threads_per_segment thread_index elements_per_thread num_elements chunk_var =-  chunk_var-    <-- sMin64-      (Imp.unCount elements_per_thread)-      ((Imp.unCount num_elements - thread_index) `divUp` threads_per_segment)-computeThreadChunkSize Noncommutative _ thread_index elements_per_thread num_elements chunk_var = do-  starting_point <--    dPrimV "starting_point" $-      thread_index * Imp.unCount elements_per_thread-  remaining_elements <--    dPrimV "remaining_elements" $-      Imp.unCount num_elements - tvExp starting_point--  let no_remaining_elements = tvExp remaining_elements .<=. 0-      beyond_bounds = Imp.unCount num_elements .<=. tvExp starting_point--  sIf-    (no_remaining_elements .||. beyond_bounds)-    (chunk_var <-- 0)-    ( sIf-        is_last_thread-        (chunk_var <-- Imp.unCount last_thread_elements)-        (chunk_var <-- Imp.unCount elements_per_thread)-    )-  where-    last_thread_elements =-      num_elements - Imp.elements thread_index * elements_per_thread-    is_last_thread =-      Imp.unCount num_elements-        .<. (thread_index + 1)-        * Imp.unCount elements_per_thread--reductionStageZero ::-  KernelConstants ->-  [(VName, Imp.TExp Int64)] ->-  Imp.Count Imp.Elements (Imp.TExp Int64) ->-  Imp.TExp Int64 ->-  Imp.Count Imp.Elements (Imp.TExp Int64) ->-  Imp.TExp Int64 ->-  [SegBinOpSlug] ->-  DoSegBody ->-  InKernelGen ([Lambda GPUMem], InKernelGen ())-reductionStageZero constants ispace num_elements global_tid elems_per_thread threads_per_segment slugs body = do-  let (gtids, _dims) = unzip ispace-      gtid = mkTV (last gtids) int64-      local_tid = sExt64 $ kernelLocalThreadId constants--  -- Figure out how many elements this thread should process.-  chunk_size <- dPrim "chunk_size" int64-  computeThreadChunkSize-    (slugsComm slugs)-    threads_per_segment-    (sExt64 global_tid)-    elems_per_thread-    num_elements-    chunk_size--  dScope Nothing $ scopeOfLParams $ concatMap slugParams slugs--  sComment "neutral-initialise the accumulators" $-    forM_ slugs $ \slug ->-      forM_ (zip (slugAccs slug) (slugNeutral slug)) $ \((acc, acc_is), ne) ->-        sLoopNest (slugShape slug) $ \vec_is ->-          copyDWIMFix acc (acc_is ++ vec_is) ne []--  slugs_op_renamed <- mapM (renameLambda . segBinOpLambda . slugOp) slugs--  let doTheReduction =-        forM_ (zip slugs_op_renamed slugs) $ \(slug_op_renamed, slug) ->-          sLoopNest (slugShape slug) $ \vec_is -> do-            comment "to reduce current chunk, first store our result in memory" $ do-              forM_ (zip (slugParams slug) (slugAccs slug)) $ \(p, (acc, acc_is)) ->-                copyDWIMFix (paramName p) [] (Var acc) (acc_is ++ vec_is)--              forM_ (zip (slugArrs slug) (slugParams slug)) $ \(arr, p) ->-                when (primType $ paramType p) $-                  copyDWIMFix arr [local_tid] (Var $ paramName p) []--            sOp $ Imp.ErrorSync Imp.FenceLocal -- Also implicitly barrier.-            groupReduce (sExt32 (kernelGroupSize constants)) slug_op_renamed (slugArrs slug)--            sOp $ Imp.Barrier Imp.FenceLocal--            sComment "first thread saves the result in accumulator" $-              sWhen (local_tid .==. 0) $-                forM_ (zip (slugAccs slug) (lambdaParams slug_op_renamed)) $ \((acc, acc_is), p) ->-                  copyDWIMFix acc (acc_is ++ vec_is) (Var $ paramName p) []--  -- If this is a non-commutative reduction, each thread must run the-  -- loop the same number of iterations, because we will be performing-  -- a group-wide reduction in there.-  let comm = slugsComm slugs-      (bound, check_bounds) =-        case comm of-          Commutative -> (tvExp chunk_size, id)-          Noncommutative ->-            ( Imp.unCount elems_per_thread,-              sWhen (tvExp gtid .<. Imp.unCount num_elements)-            )--  sFor "i" bound $ \i -> do-    gtid-      <-- case comm of-        Commutative ->-          global_tid + threads_per_segment * i-        Noncommutative ->-          let index_in_segment = global_tid `quot` kernelGroupSize constants-           in sExt64 local_tid-                + (index_in_segment * Imp.unCount elems_per_thread + i)-                  * kernelGroupSize constants--    check_bounds $-      sComment "apply map function" $-        body $ \all_red_res -> do-          let slugs_res = chunks (map (length . slugNeutral) slugs) all_red_res--          forM_ (zip slugs slugs_res) $ \(slug, red_res) ->-            sLoopNest (slugShape slug) $ \vec_is -> do-              sComment "load accumulator" $-                forM_ (zip (accParams slug) (slugAccs slug)) $ \(p, (acc, acc_is)) ->-                  copyDWIMFix (paramName p) [] (Var acc) (acc_is ++ vec_is)-              sComment "load new values" $-                forM_ (zip (nextParams slug) red_res) $ \(p, (res, res_is)) ->-                  copyDWIMFix (paramName p) [] res (res_is ++ vec_is)-              sComment "apply reduction operator"-                $ compileStms mempty (bodyStms $ slugBody slug)-                $ sComment "store in accumulator"-                $ forM_-                  ( zip-                      (slugAccs slug)-                      (map resSubExp $ bodyResult $ slugBody slug)-                  )-                $ \((acc, acc_is), se) ->-                  copyDWIMFix acc (acc_is ++ vec_is) se []--    case comm of-      Noncommutative -> do-        doTheReduction-        sComment "first thread keeps accumulator; others reset to neutral element" $ do-          let reset_to_neutral =-                forM_ slugs $ \slug ->-                  forM_ (zip (slugAccs slug) (slugNeutral slug)) $ \((acc, acc_is), ne) ->-                    sLoopNest (slugShape slug) $ \vec_is ->-                      copyDWIMFix acc (acc_is ++ vec_is) ne []-          sUnless (local_tid .==. 0) reset_to_neutral-      _ -> pure ()-  sOp $ Imp.ErrorSync Imp.FenceLocal-  pure (slugs_op_renamed, doTheReduction)--reductionStageOne ::-  KernelConstants ->-  [(VName, Imp.TExp Int64)] ->-  Imp.Count Imp.Elements (Imp.TExp Int64) ->-  Imp.TExp Int64 ->-  Imp.Count Imp.Elements (Imp.TExp Int64) ->-  Imp.TExp Int64 ->-  [SegBinOpSlug] ->-  DoSegBody ->-  InKernelGen [Lambda GPUMem]-reductionStageOne constants ispace num_elements global_tid elems_per_thread threads_per_segment slugs body = do-  (slugs_op_renamed, doTheReduction) <--    reductionStageZero constants ispace num_elements global_tid elems_per_thread threads_per_segment slugs body--  case slugsComm slugs of-    Noncommutative -> pure ()-    Commutative -> doTheReduction--  pure slugs_op_renamed--reductionStageTwo ::-  KernelConstants ->-  [PatElem LetDecMem] ->-  Imp.TExp Int32 ->-  Imp.TExp Int32 ->-  [Imp.TExp Int64] ->-  Imp.TExp Int64 ->-  Imp.TExp Int64 ->-  SegBinOpSlug ->-  [LParam GPUMem] ->-  [LParam GPUMem] ->-  Lambda GPUMem ->-  [SubExp] ->-  Imp.TExp Int32 ->-  VName ->-  Imp.TExp Int32 ->-  VName ->-  [VName] ->-  [VName] ->-  InKernelGen ()-reductionStageTwo-  constants-  segred_pes-  group_id-  flat_segment_id-  segment_gtids-  first_group_for_segment-  groups_per_segment-  slug-  red_x_params-  red_y_params-  red_op_renamed-  nes-  num_counters-  counter-  counter_i-  sync_arr-  group_res_arrs-  red_arrs = do-    let local_tid = kernelLocalThreadId constants-        group_size = kernelGroupSize constants-    old_counter <- dPrim "old_counter" int32-    (counter_mem, _, counter_offset) <--      fullyIndexArray-        counter-        [ sExt64 $-            counter_i * num_counters-              + flat_segment_id `rem` num_counters-        ]-    comment "first thread in group saves group result to global memory" $-      sWhen (local_tid .==. 0) $ do-        forM_ (take (length nes) $ zip group_res_arrs (slugAccs slug)) $ \(v, (acc, acc_is)) ->-          copyDWIMFix v [0, sExt64 group_id] (Var acc) acc_is-        sOp $ Imp.MemFence Imp.FenceGlobal-        -- Increment the counter, thus stating that our result is-        -- available.-        sOp-          $ Imp.Atomic DefaultSpace-          $ Imp.AtomicAdd-            Int32-            (tvVar old_counter)-            counter_mem-            counter_offset-          $ untyped (1 :: Imp.TExp Int32)-        -- Now check if we were the last group to write our result.  If-        -- so, it is our responsibility to produce the final result.-        sWrite sync_arr [0] $ untyped $ tvExp old_counter .==. groups_per_segment - 1--    sOp $ Imp.Barrier Imp.FenceGlobal--    is_last_group <- dPrim "is_last_group" Bool-    copyDWIMFix (tvVar is_last_group) [] (Var sync_arr) [0]-    sWhen (tvExp is_last_group) $ do-      -- The final group has written its result (and it was-      -- us!), so read in all the group results and perform the-      -- final stage of the reduction.  But first, we reset the-      -- counter so it is ready for next time.  This is done-      -- with an atomic to avoid warnings about write/write-      -- races in oclgrind.-      sWhen (local_tid .==. 0) $-        sOp $-          Imp.Atomic DefaultSpace $-            Imp.AtomicAdd Int32 (tvVar old_counter) counter_mem counter_offset $-              untyped $-                negate groups_per_segment--      sLoopNest (slugShape slug) $ \vec_is -> do-        unless (null $ slugShape slug) $-          sOp (Imp.Barrier Imp.FenceLocal)--        -- There is no guarantee that the number of workgroups for the-        -- segment is less than the workgroup size, so each thread may-        -- have to read multiple elements.  We do this in a sequential-        -- way that may induce non-coalesced accesses, but the total-        -- number of accesses should be tiny here.-        comment "read in the per-group-results" $ do-          read_per_thread <--            dPrimVE "read_per_thread" $-              groups_per_segment `divUp` sExt64 group_size--          forM_ (zip red_x_params nes) $ \(p, ne) ->-            copyDWIMFix (paramName p) [] ne []--          sFor "i" read_per_thread $ \i -> do-            group_res_id <--              dPrimVE "group_res_id" $-                sExt64 local_tid * read_per_thread + i-            index_of_group_res <--              dPrimVE "index_of_group_res" $-                first_group_for_segment + group_res_id--            sWhen (group_res_id .<. groups_per_segment) $ do-              forM_ (zip red_y_params group_res_arrs) $-                \(p, group_res_arr) ->-                  copyDWIMFix-                    (paramName p)-                    []-                    (Var group_res_arr)-                    ([0, index_of_group_res] ++ vec_is)--              compileStms mempty (bodyStms $ slugBody slug) $-                forM_ (zip red_x_params $ map resSubExp $ bodyResult $ slugBody slug) $ \(p, se) ->-                  copyDWIMFix (paramName p) [] se []--        forM_ (zip red_x_params red_arrs) $ \(p, arr) ->-          when (primType $ paramType p) $-            copyDWIMFix arr [sExt64 local_tid] (Var $ paramName p) []--        sOp $ Imp.ErrorSync Imp.FenceLocal--        sComment "reduce the per-group results" $ do-          groupReduce (sExt32 group_size) red_op_renamed red_arrs--          sComment "and back to memory with the final result" $-            sWhen (local_tid .==. 0) $-              forM_ (zip segred_pes $ lambdaParams red_op_renamed) $ \(pe, p) ->-                copyDWIMFix-                  (patElemName pe)-                  (segment_gtids ++ vec_is)-                  (Var $ paramName p)-                  []+--+-- An optimization specfically targeted at non-segmented and large-segments+-- segmented reductions with non-commutative is made: The stage one main loop is+-- essentially stripmined by a factor `chunk`, inserting collective copies via+-- local memory of each reduction parameter going into the intra-group (partial)+-- reductions. This saves a factor `chunk` number of intra-group reductions at+-- the cost of some overhead in collective copies.+module Futhark.CodeGen.ImpGen.GPU.SegRed+  ( compileSegRed,+    compileSegRed',+    DoSegBody,+  )+where++import Control.Monad+import Data.List (genericLength, zip4)+import Data.Map qualified as M+import Data.Maybe+import Futhark.CodeGen.ImpCode.GPU qualified as Imp+import Futhark.CodeGen.ImpGen+import Futhark.CodeGen.ImpGen.GPU.Base+import Futhark.Error+import Futhark.IR.GPUMem+import Futhark.IR.Mem.LMAD qualified as LMAD+import Futhark.Transform.Rename+import Futhark.Util (chunks, mapAccumLM)+import Futhark.Util.IntegralExp (divUp, nextMul, quot, rem)+import Prelude hiding (quot, rem)++forM2_ :: (Monad m) => [a] -> [b] -> (a -> b -> m c) -> m ()+forM2_ xs ys f = forM_ (zip xs ys) (uncurry f)++-- | The maximum number of operators we support in a single SegRed.+-- This limit arises out of the static allocation of counters.+maxNumOps :: Int+maxNumOps = 10++-- | Code generation for the body of the SegRed, taking a continuation+-- for saving the results of the body.  The results should be+-- represented as a pairing of a t'SubExp' along with a list of+-- indexes into that t'SubExp' for reading the result.+type DoSegBody = ([(SubExp, [Imp.TExp Int64])] -> InKernelGen ()) -> InKernelGen ()++-- | Datatype used to distinguish between and work with the different sets of+-- intermediate memory we need for the different ReduceKinds.+data SegRedIntermediateArrays+  = GeneralSegRedInterms+      { groupRedArrs :: [VName]+      }+  | NoncommPrimSegRedInterms+      { collCopyArrs :: [VName],+        groupRedArrs :: [VName],+        privateChunks :: [VName]+      }++-- | Compile 'SegRed' instance to host-level code with calls to+-- various kernels.+compileSegRed ::+  Pat LetDecMem ->+  SegLevel ->+  SegSpace ->+  [SegBinOp GPUMem] ->+  KernelBody GPUMem ->+  CallKernelGen ()+compileSegRed pat lvl space segbinops map_kbody = do+  emit $ Imp.DebugPrint "\n# SegRed" Nothing+  KernelAttrs {kAttrNumGroups = num_groups, kAttrGroupSize = group_size} <-+    lvlKernelAttrs lvl+  let grid = KernelGrid num_groups group_size++  compileSegRed' pat grid space segbinops $ \red_cont ->+    sComment "apply map function" $+      compileStms mempty (kernelBodyStms map_kbody) $ do+        let (red_res, map_res) = splitAt (segBinOpResults segbinops) $ kernelBodyResult map_kbody++        let mapout_arrs = drop (segBinOpResults segbinops) $ patElems pat+        unless (null mapout_arrs) $+          sComment "write map-out result(s)" $ do+            zipWithM_ (compileThreadResult space) mapout_arrs map_res++        red_cont $ map ((,[]) . kernelResultSubExp) red_res+  emit $ Imp.DebugPrint "" Nothing++paramOf :: SegBinOp GPUMem -> [Param LParamMem]+paramOf (SegBinOp _ op ne _) = take (length ne) $ lambdaParams op++isPrimSegBinOp :: SegBinOp GPUMem -> Bool+isPrimSegBinOp segbinop =+  all primType (lambdaReturnType $ segBinOpLambda segbinop)+    && shapeRank (segBinOpShape segbinop) == 0++-- | Like 'compileSegRed', but where the body is a monadic action.+compileSegRed' ::+  Pat LetDecMem ->+  KernelGrid ->+  SegSpace ->+  [SegBinOp GPUMem] ->+  DoSegBody ->+  CallKernelGen ()+compileSegRed' pat grid space segbinops map_body_cont+  | genericLength segbinops > maxNumOps =+      compilerLimitationS $+        "compileSegRed': at most " ++ show maxNumOps ++ " reduction operators are supported."+  | otherwise = do+      chunk_v <- dPrimV "chunk_size" . isInt64 =<< kernelConstToExp chunk_const+      case unSegSpace space of+        [(_, Constant (IntValue (Int64Value 1))), _] ->+          compileReduction (chunk_v, chunk_const) nonsegmentedReduction+        _ -> do+          let segment_size = pe64 $ last $ segSpaceDims space+              use_small_segments = segment_size * 2 .<. pe64 (unCount group_size) * tvExp chunk_v+          sIf+            use_small_segments+            (compileReduction (chunk_v, chunk_const) smallSegmentsReduction)+            (compileReduction (chunk_v, chunk_const) largeSegmentsReduction)+  where+    compileReduction chunk f =+      f pat num_groups group_size chunk space segbinops map_body_cont++    param_types = map paramType $ concatMap paramOf segbinops++    num_groups = gridNumGroups grid+    group_size = gridGroupSize grid++    chunk_const =+      if Noncommutative `elem` map segBinOpComm segbinops+        && all isPrimSegBinOp segbinops+        then getChunkSize param_types+        else Imp.ValueExp $ IntValue $ intValue Int64 (1 :: Int64)++-- | Prepare intermediate arrays for the reduction.  Prim-typed+-- arguments go in local memory (so we need to do the allocation of+-- those arrays inside the kernel), while array-typed arguments go in+-- global memory.  Allocations for the latter have already been+-- performed.  This policy is baked into how the allocations are done+-- in ExplicitAllocations.+--+-- For more info about the intermediate arrays used for the different reduction+-- kernels, see note [IntermArrays].+makeIntermArrays ::+  Imp.TExp Int64 ->+  SubExp ->+  SubExp ->+  [SegBinOp GPUMem] ->+  InKernelGen [SegRedIntermediateArrays]+makeIntermArrays group_id group_size chunk segbinops+  | Noncommutative <- mconcat (map segBinOpComm segbinops),+    all isPrimSegBinOp segbinops =+      noncommPrimSegRedInterms+  | otherwise =+      generalSegRedInterms group_id group_size segbinops+  where+    params = map paramOf segbinops++    noncommPrimSegRedInterms = do+      group_worksize <- tvSize <$> dPrimV "group_worksize" group_worksize_E++      -- compute total amount of lmem.+      let sum_ x y = nextMul x y + group_size_E * y+          group_reds_lmem_requirement = foldl sum_ 0 $ concat elem_sizes+          collcopy_lmem_requirement = group_worksize_E * max_elem_size+          lmem_total_size =+            Imp.bytes $+              collcopy_lmem_requirement `sMax64` group_reds_lmem_requirement++      -- offsets into the total pool of lmem for each group reduction array.+      (_, offsets) <-+        forAccumLM2D 0 elem_sizes $ \byte_offs elem_size ->+          (,byte_offs `quot` elem_size)+            <$> dPrimVE "offset" (sum_ byte_offs elem_size)++      -- total pool of local mem.+      lmem <- sAlloc "local_mem" lmem_total_size (Space "local")+      let arrInLMem ptype name len_se offset =+            sArray+              (name ++ "_" ++ prettyString ptype)+              ptype+              (Shape [len_se])+              lmem+              $ LMAD.iota offset [pe64 len_se]++      forM (zipWith zip params offsets) $ \ps_and_offsets -> do+        (coll_copy_arrs, group_red_arrs, priv_chunks) <-+          fmap unzip3 $ forM ps_and_offsets $ \(p, offset) -> do+            let ptype = elemType $ paramType p+            (,,)+              <$> arrInLMem ptype "coll_copy_arr" group_worksize 0+              <*> arrInLMem ptype "group_red_arr" group_size offset+              <*> sAllocArray+                ("chunk_" ++ prettyString ptype)+                ptype+                (Shape [chunk])+                (ScalarSpace [chunk] ptype)+        pure $ NoncommPrimSegRedInterms coll_copy_arrs group_red_arrs priv_chunks+    group_size_E = pe64 group_size+    group_worksize_E = group_size_E * pe64 chunk++    paramSize = primByteSize . elemType . paramType+    elem_sizes = map (map paramSize) params+    max_elem_size = maximum $ concat elem_sizes++    forAccumLM2D acc ls f = mapAccumLM (mapAccumLM f) acc ls++generalSegRedInterms ::+  Imp.TExp Int64 ->+  SubExp ->+  [SegBinOp GPUMem] ->+  InKernelGen [SegRedIntermediateArrays]+generalSegRedInterms group_id group_size segbinops =+  fmap (map GeneralSegRedInterms) $+    forM (map paramOf segbinops) $+      mapM $ \p ->+        case paramDec p of+          MemArray pt shape _ (ArrayIn mem _) -> do+            let shape' = Shape [group_size] <> shape+            let shape_E = map pe64 $ shapeDims shape'+            sArray ("red_arr_" ++ prettyString pt) pt shape' mem $+              LMAD.iota (group_id * product shape_E) shape_E+          _ -> do+            let pt = elemType $ paramType p+                shape = Shape [group_size]+            sAllocArray ("red_arr_" ++ prettyString pt) pt shape $ Space "local"++-- | Arrays for storing group results.+--+-- The group-result arrays have an extra dimension because they are+-- also used for keeping vectorised accumulators for first-stage+-- reduction, if necessary.  If necessary, this dimension has size+-- group_size, and otherwise 1.  When actually storing group results,+-- the first index is set to 0.+groupResultArrays ::+  SubExp ->+  SubExp ->+  [SegBinOp GPUMem] ->+  CallKernelGen [[VName]]+groupResultArrays num_virtgroups group_size segbinops =+  forM segbinops $ \(SegBinOp _ lam _ shape) ->+    forM (lambdaReturnType lam) $ \t -> do+      let pt = elemType t+          extra_dim+            | primType t, shapeRank shape == 0 = intConst Int64 1+            | otherwise = group_size+          full_shape = Shape [extra_dim, num_virtgroups] <> shape <> arrayShape t+          -- Move the groupsize dimension last to ensure coalesced+          -- memory access.+          perm = [1 .. shapeRank full_shape - 1] ++ [0]+      sAllocArrayPerm "segred_tmp" pt full_shape (Space "device") perm++type DoCompileSegRed =+  Pat LetDecMem ->+  Count NumGroups SubExp ->+  Count GroupSize SubExp ->+  (TV Int64, Imp.KernelConstExp) ->+  SegSpace ->+  [SegBinOp GPUMem] ->+  DoSegBody ->+  CallKernelGen ()++nonsegmentedReduction :: DoCompileSegRed+nonsegmentedReduction (Pat segred_pes) num_groups group_size (chunk_v, chunk_const) space segbinops map_body_cont = do+  let (gtids, dims) = unzip $ unSegSpace space+      chunk = tvExp chunk_v+      num_groups_se = unCount num_groups+      group_size_se = unCount group_size+      group_size' = pe64 group_size_se+      global_tid = Imp.le64 $ segFlat space+      n = pe64 $ last dims++  counters <- genZeroes "counters" maxNumOps++  reds_group_res_arrs <- groupResultArrays num_groups_se group_size_se segbinops++  num_threads <-+    fmap tvSize $ dPrimV "num_threads" $ pe64 num_groups_se * group_size'++  let attrs =+        (defKernelAttrs num_groups group_size)+          { kAttrConstExps = M.singleton (tvVar chunk_v) chunk_const+          }++  sKernelThread "segred_nonseg" (segFlat space) attrs $ do+    constants <- kernelConstants <$> askEnv+    let ltid = kernelLocalThreadId constants+    let group_id = kernelGroupId constants++    interms <- makeIntermArrays (sExt64 group_id) group_size_se (tvSize chunk_v) segbinops+    sync_arr <- sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local"++    -- Since this is the nonsegmented case, all outer segment IDs must+    -- necessarily be 0.+    forM_ gtids $ \v -> dPrimV_ v (0 :: Imp.TExp Int64)++    q <- dPrimVE "q" $ n `divUp` (sExt64 (kernelNumThreads constants) * chunk)++    slugs <-+      mapM (segBinOpSlug ltid group_id) $+        zip3 segbinops interms reds_group_res_arrs+    new_lambdas <-+      reductionStageOne+        gtids+        n+        global_tid+        q+        chunk+        (pe64 num_threads)+        slugs+        map_body_cont++    let segred_pess =+          chunks+            (map (length . segBinOpNeutral) segbinops)+            segred_pes++    forM_ (zip4 segred_pess slugs new_lambdas [0 ..]) $+      \(pes, slug, new_lambda, i) ->+        reductionStageTwo+          pes+          group_id+          [0]+          0+          (sExt64 $ kernelNumGroups constants)+          slug+          new_lambda+          counters+          sync_arr+          (fromInteger i)++smallSegmentsReduction :: DoCompileSegRed+smallSegmentsReduction (Pat segred_pes) num_groups group_size _ space segbinops map_body_cont = do+  let (gtids, dims) = unzip $ unSegSpace space+      dims' = map pe64 dims+      segment_size = last dims'++  -- Careful to avoid division by zero now.+  segment_size_nonzero <-+    dPrimVE "segment_size_nonzero" $ sMax64 1 segment_size++  let group_size_se = unCount group_size+      num_groups_se = unCount group_size+      num_groups' = pe64 num_groups_se+      group_size' = pe64 group_size_se+  num_threads <- fmap tvSize $ dPrimV "num_threads" $ num_groups' * group_size'+  let num_segments = product $ init dims'+      segments_per_group = group_size' `quot` segment_size_nonzero+      required_groups = sExt32 $ num_segments `divUp` segments_per_group++  emit $ Imp.DebugPrint "# SegRed-small" Nothing+  emit $ Imp.DebugPrint "num_segments" $ Just $ untyped num_segments+  emit $ Imp.DebugPrint "segment_size" $ Just $ untyped segment_size+  emit $ Imp.DebugPrint "segments_per_group" $ Just $ untyped segments_per_group+  emit $ Imp.DebugPrint "required_groups" $ Just $ untyped required_groups++  sKernelThread "segred_small" (segFlat space) (defKernelAttrs num_groups group_size) $ do+    constants <- kernelConstants <$> askEnv+    let group_id = kernelGroupSize constants+        ltid = sExt64 $ kernelLocalThreadId constants++    interms <- generalSegRedInterms group_id group_size_se segbinops+    let reds_arrs = map groupRedArrs interms++    -- We probably do not have enough actual workgroups to cover the+    -- entire iteration space.  Some groups thus have to perform double+    -- duty; we put an outer loop to accomplish this.+    virtualiseGroups SegVirt required_groups $ \virtgroup_id -> do+      -- Compute the 'n' input indices.  The outer 'n-1' correspond to+      -- the segment ID, and are computed from the group id.  The inner+      -- is computed from the local thread id, and may be out-of-bounds.+      let segment_index =+            (ltid `quot` segment_size_nonzero)+              + (sExt64 virtgroup_id * sExt64 segments_per_group)+          index_within_segment = ltid `rem` segment_size++      dIndexSpace (zip (init gtids) (init dims')) segment_index+      dPrimV_ (last gtids) index_within_segment++      let in_bounds =+            map_body_cont $ \red_res ->+              sComment "save results to be reduced" $ do+                let red_dests = map (,[ltid]) (concat reds_arrs)+                forM2_ red_dests red_res $ \(d, d_is) (res, res_is) ->+                  copyDWIMFix d d_is res res_is+          out_of_bounds =+            forM2_ segbinops reds_arrs $ \(SegBinOp _ _ nes _) red_arrs ->+              forM2_ red_arrs nes $ \arr ne ->+                copyDWIMFix arr [ltid] ne []++      sComment "apply map function if in bounds" $+        sIf+          ( segment_size+              .>. 0+              .&&. isActive (init $ zip gtids dims)+              .&&. ltid+              .<. segment_size+              * segments_per_group+          )+          in_bounds+          out_of_bounds++      sOp $ Imp.ErrorSync Imp.FenceLocal -- Also implicitly barrier.+      let crossesSegment from to =+            (sExt64 to - sExt64 from) .>. (sExt64 to `rem` segment_size)+      sWhen (segment_size .>. 0) $+        sComment "perform segmented scan to imitate reduction" $+          forM2_ segbinops reds_arrs $ \(SegBinOp _ red_op _ _) red_arrs ->+            groupScan+              (Just crossesSegment)+              (sExt64 $ pe64 num_threads)+              (segment_size * segments_per_group)+              red_op+              red_arrs++      sOp $ Imp.Barrier Imp.FenceLocal++      sComment "save final values of segments"+        $ sWhen+          ( sExt64 virtgroup_id+              * segments_per_group+              + sExt64 ltid+                .<. num_segments+                .&&. ltid+                .<. segments_per_group+          )+        $ forM2_ segred_pes (concat reds_arrs)+        $ \pe arr -> do+          -- Figure out which segment result this thread should write...+          let flat_segment_index =+                sExt64 virtgroup_id * segments_per_group + sExt64 ltid+              gtids' =+                unflattenIndex (init dims') flat_segment_index+          copyDWIMFix+            (patElemName pe)+            gtids'+            (Var arr)+            [(ltid + 1) * segment_size_nonzero - 1]++      -- Finally another barrier, because we will be writing to the+      -- local memory array first thing in the next iteration.+      sOp $ Imp.Barrier Imp.FenceLocal++largeSegmentsReduction :: DoCompileSegRed+largeSegmentsReduction (Pat segred_pes) num_groups group_size (chunk_v, chunk_const) space segbinops map_body_cont = do+  let (gtids, dims) = unzip $ unSegSpace space+      dims' = map pe64 dims+      num_segments = product $ init dims'+      segment_size = last dims'+      num_groups' = pe64 $ unCount num_groups+      group_size_se = unCount group_size+      group_size' = pe64 group_size_se+      chunk = tvExp chunk_v++  groups_per_segment <-+    dPrimVE "groups_per_segment" $+      num_groups' `divUp` sMax64 1 num_segments++  q <-+    dPrimVE "q" $+      segment_size `divUp` (group_size' * groups_per_segment * chunk)++  num_virtgroups <-+    dPrimV "num_virtgroups" $+      groups_per_segment * num_segments+  threads_per_segment <-+    dPrimVE "threads_per_segment" $+      groups_per_segment * group_size'++  emit $ Imp.DebugPrint "# SegRed-large" Nothing+  emit $ Imp.DebugPrint "num_segments" $ Just $ untyped num_segments+  emit $ Imp.DebugPrint "segment_size" $ Just $ untyped segment_size+  emit $ Imp.DebugPrint "num_virtgroups" $ Just $ untyped $ tvExp num_virtgroups+  emit $ Imp.DebugPrint "num_groups" $ Just $ untyped num_groups'+  emit $ Imp.DebugPrint "group_size" $ Just $ untyped group_size'+  emit $ Imp.DebugPrint "q" $ Just $ untyped q+  emit $ Imp.DebugPrint "groups_per_segment" $ Just $ untyped groups_per_segment++  reds_group_res_arrs <- groupResultArrays (tvSize num_virtgroups) group_size_se segbinops++  -- In principle we should have a counter for every segment.  Since+  -- the number of segments is a dynamic quantity, we would have to+  -- allocate and zero out an array here, which is expensive.+  -- However, we exploit the fact that the number of segments being+  -- reduced at any point in time is limited by the number of+  -- workgroups. If we bound the number of workgroups, we can get away+  -- with using that many counters.  FIXME: Is this limit checked+  -- anywhere?  There are other places in the compiler that will fail+  -- if the group count exceeds the maximum group size, which is at+  -- most 1024 anyway.+  let num_counters = maxNumOps * 1024+  counters <- genZeroes "counters" $ fromIntegral num_counters++  let attrs =+        (defKernelAttrs num_groups group_size)+          { kAttrConstExps = M.singleton (tvVar chunk_v) chunk_const+          }++  sKernelThread "segred_large" (segFlat space) attrs $ do+    constants <- kernelConstants <$> askEnv+    let group_id = sExt64 $ kernelGroupId constants+        ltid = kernelLocalThreadId constants++    interms <- makeIntermArrays group_id group_size_se (tvSize chunk_v) segbinops+    sync_arr <- sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local"++    -- We probably do not have enough actual workgroups to cover the+    -- entire iteration space.  Some groups thus have to perform double+    -- duty; we put an outer loop to accomplish this.+    virtualiseGroups SegVirt (sExt32 (tvExp num_virtgroups)) $ \virtgroup_id -> do+      let segment_gtids = init gtids++      flat_segment_id <-+        dPrimVE "flat_segment_id" $+          sExt64 virtgroup_id `quot` groups_per_segment++      global_tid <-+        dPrimVE "global_tid" $+          (sExt64 virtgroup_id * sExt64 group_size' + sExt64 ltid)+            `rem` threads_per_segment++      let first_group_for_segment = flat_segment_id * groups_per_segment+      dIndexSpace (zip segment_gtids (init dims')) flat_segment_id+      dPrim_ (last gtids) int64+      let n = pe64 $ last dims++      slugs <-+        mapM (segBinOpSlug ltid virtgroup_id) $+          zip3 segbinops interms reds_group_res_arrs+      new_lambdas <-+        reductionStageOne+          gtids+          n+          global_tid+          q+          chunk+          threads_per_segment+          slugs+          map_body_cont++      let segred_pess =+            chunks+              (map (length . segBinOpNeutral) segbinops)+              segred_pes++          multiple_groups_per_segment =+            forM_ (zip4 segred_pess slugs new_lambdas [0 ..]) $+              \(pes, slug, new_lambda, i) -> do+                let counter_idx =+                      fromIntegral (i * num_counters)+                        + flat_segment_id+                          `rem` fromIntegral num_counters+                reductionStageTwo+                  pes+                  virtgroup_id+                  (map Imp.le64 segment_gtids)+                  first_group_for_segment+                  groups_per_segment+                  slug+                  new_lambda+                  counters+                  sync_arr+                  counter_idx++          one_group_per_segment =+            sComment "first thread in group saves final result to memory" $+              forM2_ slugs segred_pess $ \slug pes ->+                sWhen (ltid .==. 0) $+                  forM2_ pes (slugAccs slug) $ \v (acc, acc_is) ->+                    copyDWIMFix (patElemName v) (map Imp.le64 segment_gtids) (Var acc) acc_is++      sIf (groups_per_segment .==. 1) one_group_per_segment multiple_groups_per_segment++-- | Auxiliary information for a single reduction. A slug holds the `SegBinOp`+-- operator for a single reduction, the different arrays required throughout+-- stages one and two, and a global mem destination for the final result of the+-- particular reduction.+data SegBinOpSlug = SegBinOpSlug+  { slugOp :: SegBinOp GPUMem,+    -- | Intermediate arrays needed for the given reduction.+    slugInterms :: SegRedIntermediateArrays,+    -- | Place(s) to store group accumulator(s) in stage 1 reduction.+    slugAccs :: [(VName, [Imp.TExp Int64])],+    -- | Global memory destination(s) for the final result(s) for this+    -- particular reduction.+    groupResArrs :: [VName]+  }++segBinOpSlug ::+  Imp.TExp Int32 ->+  Imp.TExp Int32 ->+  (SegBinOp GPUMem, SegRedIntermediateArrays, [VName]) ->+  InKernelGen SegBinOpSlug+segBinOpSlug ltid group_id (op, interms, group_res_arrs) = do+  accs <- zipWithM mkAcc (lambdaParams (segBinOpLambda op)) group_res_arrs+  pure $ SegBinOpSlug op interms accs group_res_arrs+  where+    mkAcc p group_res_arr+      | Prim t <- paramType p,+        shapeRank (segBinOpShape op) == 0 = do+          group_res_acc <- dPrim (baseString (paramName p) <> "_group_res_acc") t+          pure (tvVar group_res_acc, [])+      -- if this is a non-primitive reduction, the global mem result array will+      -- double as accumulator.+      | otherwise =+          pure (group_res_arr, [sExt64 ltid, sExt64 group_id])++slugLambda :: SegBinOpSlug -> Lambda GPUMem+slugLambda = segBinOpLambda . slugOp++slugBody :: SegBinOpSlug -> Body GPUMem+slugBody = lambdaBody . slugLambda++slugParams :: SegBinOpSlug -> [LParam GPUMem]+slugParams = lambdaParams . slugLambda++slugNeutral :: SegBinOpSlug -> [SubExp]+slugNeutral = segBinOpNeutral . slugOp++slugShape :: SegBinOpSlug -> Shape+slugShape = segBinOpShape . slugOp++slugsComm :: [SegBinOpSlug] -> Commutativity+slugsComm = mconcat . map (segBinOpComm . slugOp)++slugSplitParams :: SegBinOpSlug -> ([LParam GPUMem], [LParam GPUMem])+slugSplitParams slug = splitAt (length (slugNeutral slug)) $ slugParams slug++slugGroupRedArrs :: SegBinOpSlug -> [VName]+slugGroupRedArrs = groupRedArrs . slugInterms++slugPrivChunks :: SegBinOpSlug -> [VName]+slugPrivChunks = privateChunks . slugInterms++slugCollCopyArrs :: SegBinOpSlug -> [VName]+slugCollCopyArrs = collCopyArrs . slugInterms++reductionStageOne ::+  [VName] ->+  Imp.TExp Int64 ->+  Imp.TExp Int64 ->+  Imp.TExp Int64 ->+  Imp.TExp Int64 ->+  Imp.TExp Int64 ->+  [SegBinOpSlug] ->+  DoSegBody ->+  InKernelGen [Lambda GPUMem]+reductionStageOne gtids n global_tid q chunk threads_per_segment slugs body_cont = do+  constants <- kernelConstants <$> askEnv+  let glb_ind_var = mkTV (last gtids) int64+      ltid = sExt64 $ kernelLocalThreadId constants++  dScope Nothing $ scopeOfLParams $ concatMap slugParams slugs++  sComment "ne-initialise the outer (per-group) accumulator(s)" $ do+    forM_ slugs $ \slug ->+      forM2_ (slugAccs slug) (slugNeutral slug) $ \(acc, acc_is) ne ->+        sLoopNest (slugShape slug) $ \vec_is ->+          copyDWIMFix acc (acc_is ++ vec_is) ne []++  new_lambdas <- mapM (renameLambda . slugLambda) slugs+  let group_size = sExt32 $ kernelGroupSize constants+  let doGroupReduce =+        forM2_ slugs new_lambdas $ \slug new_lambda -> do+          let accs = slugAccs slug+          let params = slugParams slug+          sLoopNest (slugShape slug) $ \vec_is -> do+            let group_red_arrs = slugGroupRedArrs slug+            sComment "store accs. prims go in lmem; non-prims in params (in global mem)" $+              forM_ (zip3 group_red_arrs accs params) $+                \(arr, (acc, acc_is), p) ->+                  if isPrimParam p+                    then copyDWIMFix arr [ltid] (Var acc) (acc_is ++ vec_is)+                    else copyDWIMFix (paramName p) [] (Var acc) (acc_is ++ vec_is)++            sOp $ Imp.ErrorSync Imp.FenceLocal -- Also implicitly barrier.+            groupReduce group_size new_lambda group_red_arrs+            sOp $ Imp.Barrier Imp.FenceLocal++            sComment "thread 0 updates per-group acc(s); rest reset to ne" $ do+              sIf+                (ltid .==. 0)+                ( forM2_ accs (lambdaParams new_lambda) $+                    \(acc, acc_is) p ->+                      copyDWIMFix acc (acc_is ++ vec_is) (Var $ paramName p) []+                )+                ( forM2_ accs (slugNeutral slug) $+                    \(acc, acc_is) ne ->+                      copyDWIMFix acc (acc_is ++ vec_is) ne []+                )++  case (slugsComm slugs, all (isPrimSegBinOp . slugOp) slugs) of+    (Noncommutative, True) ->+      noncommPrimParamsStageOneBody+        slugs+        body_cont+        glb_ind_var+        global_tid+        q+        n+        chunk+        doGroupReduce+    _ ->+      generalStageOneBody+        slugs+        body_cont+        glb_ind_var+        global_tid+        q+        n+        threads_per_segment+        doGroupReduce++  pure new_lambdas++generalStageOneBody ::+  [SegBinOpSlug] ->+  DoSegBody ->+  TV Int64 ->+  Imp.TExp Int64 ->+  Imp.TExp Int64 ->+  Imp.TExp Int64 ->+  Imp.TExp Int64 ->+  InKernelGen () ->+  InKernelGen ()+generalStageOneBody slugs body_cont glb_ind_var global_tid q n threads_per_segment doGroupReduce = do+  let is_comm = slugsComm slugs == Commutative++  constants <- kernelConstants <$> askEnv+  let group_size = kernelGroupSize constants+      ltid = sExt64 $ kernelLocalThreadId constants++  -- this group's id within its designated segment, and this group's initial+  -- global offset.+  group_id_in_segment <- dPrimVE "group_id_in_segment" $ global_tid `quot` group_size+  group_base_offset <- dPrimVE "group_base_offset" $ group_id_in_segment * q * group_size++  sFor "i" q $ \i -> do+    group_offset <- dPrimVE "group_offset" $ group_base_offset + i * group_size+    glb_ind_var+      <-- if is_comm+        then global_tid + threads_per_segment * i+        else group_offset + ltid++    sWhen (tvExp glb_ind_var .<. n) $+      sComment "apply map function(s)" $+        body_cont $ \all_red_res -> do+          let maps_res = chunks (map (length . slugNeutral) slugs) all_red_res++          forM2_ slugs maps_res $ \slug map_res ->+            sLoopNest (slugShape slug) $ \vec_is -> do+              let (acc_params, next_params) = slugSplitParams slug+              sComment "load accumulator(s)" $+                forM2_ acc_params (slugAccs slug) $ \p (acc, acc_is) ->+                  copyDWIMFix (paramName p) [] (Var acc) (acc_is ++ vec_is)+              sComment "load next value(s)" $+                forM2_ next_params map_res $ \p (res, res_is) ->+                  copyDWIMFix (paramName p) [] res (res_is ++ vec_is)+              sComment "apply reduction operator(s)"+                $ compileStms mempty (bodyStms $ slugBody slug)+                $ sComment "store in accumulator(s)"+                $ forM2_+                  (slugAccs slug)+                  (map resSubExp $ bodyResult $ slugBody slug)+                $ \(acc, acc_is) se ->+                  copyDWIMFix acc (acc_is ++ vec_is) se []++    unless is_comm doGroupReduce+  sOp $ Imp.ErrorSync Imp.FenceLocal+  when is_comm doGroupReduce++noncommPrimParamsStageOneBody ::+  [SegBinOpSlug] ->+  DoSegBody ->+  TV Int64 ->+  Imp.TExp Int64 ->+  Imp.TExp Int64 ->+  Imp.TExp Int64 ->+  Imp.TExp Int64 ->+  InKernelGen () ->+  InKernelGen ()+noncommPrimParamsStageOneBody slugs body_cont glb_ind_var global_tid q n chunk doLMemGroupReduce = do+  constants <- kernelConstants <$> askEnv+  let group_size = kernelGroupSize constants+      ltid = sExt64 $ kernelLocalThreadId constants++  -- this group's id within its designated segment; the stride made per group in+  -- the outer `i < q` loop; and this group's initial global offset.+  group_id_in_segment <- dPrimVE "group_offset_in_segment" $ global_tid `quot` group_size+  group_stride <- dPrimVE "group_stride" $ group_size * chunk+  group_base_offset <- dPrimVE "group_base_offset" $ group_id_in_segment * q * group_stride++  let chunkLoop = sFor "k" chunk++  sFor "i" q $ \i -> do+    -- group offset in this iteration.+    group_offset <- dPrimVE "group_offset" $ group_base_offset + i * group_stride+    chunkLoop $ \k -> do+      loc_ind <- dPrimVE "loc_ind" $ k * group_size + ltid+      glb_ind_var <-- group_offset + loc_ind++      sIf+        (tvExp glb_ind_var .<. n)+        ( body_cont $ \all_red_res -> do+            let slugs_res = chunks (map (length . slugNeutral) slugs) all_red_res+            forM2_ slugs slugs_res $ \slug slug_res -> do+              let priv_chunks = slugPrivChunks slug+              sComment "write map result(s) to private chunk(s)" $+                forM2_ priv_chunks slug_res $ \priv_chunk (res, res_is) ->+                  copyDWIMFix priv_chunk [k] res res_is+        )+        -- if out of bounds, fill chunk(s) with neutral element(s)+        ( forM_ slugs $ \slug ->+            forM2_ (slugPrivChunks slug) (slugNeutral slug) $+              \priv_chunk ne ->+                copyDWIMFix priv_chunk [k] ne []+        )++    sOp $ Imp.ErrorSync Imp.FenceLocal++    sComment "effectualize collective copies in local memory" $ do+      forM_ slugs $ \slug -> do+        let coll_copy_arrs = slugCollCopyArrs slug+        let priv_chunks = slugPrivChunks slug++        forM2_ coll_copy_arrs priv_chunks $ \lmem_arr priv_chunk -> do+          chunkLoop $ \k -> do+            lmem_idx <- dPrimVE "lmem_idx" $ ltid + k * group_size+            copyDWIMFix lmem_arr [lmem_idx] (Var priv_chunk) [k]++          sOp $ Imp.Barrier Imp.FenceLocal++          chunkLoop $ \k -> do+            lmem_idx <- dPrimVE "lmem_idx" $ ltid * chunk + k+            copyDWIMFix priv_chunk [k] (Var lmem_arr) [lmem_idx]++          sOp $ Imp.Barrier Imp.FenceLocal++    sComment "per-thread sequential reduction of private chunk(s)" $ do+      chunkLoop $ \k ->+        forM_ slugs $ \slug -> do+          let accs = map fst $ slugAccs slug+          let (acc_ps, next_ps) = slugSplitParams slug+          let ps_accs_chunks = zip4 acc_ps next_ps accs (slugPrivChunks slug)++          sComment "load params for all reductions" $ do+            forM_ ps_accs_chunks $ \(acc_p, next_p, acc, priv_chunk) -> do+              copyDWIM (paramName acc_p) [] (Var acc) []+              copyDWIMFix (paramName next_p) [] (Var priv_chunk) [k]++          sComment "apply reduction operator(s)" $ do+            let binop_ress = map resSubExp $ bodyResult $ slugBody slug+            compileStms mempty (bodyStms $ slugBody slug) $+              forM2_ accs binop_ress $ \acc binop_res ->+                copyDWIM acc [] binop_res []+    doLMemGroupReduce+  sOp $ Imp.ErrorSync Imp.FenceLocal++reductionStageTwo ::+  [PatElem LetDecMem] ->+  Imp.TExp Int32 ->+  [Imp.TExp Int64] ->+  Imp.TExp Int64 ->+  Imp.TExp Int64 ->+  SegBinOpSlug ->+  Lambda GPUMem ->+  VName ->+  VName ->+  Imp.TExp Int64 ->+  InKernelGen ()+reductionStageTwo segred_pes group_id segment_gtids first_group_for_segment groups_per_segment slug new_lambda counters sync_arr counter_idx = do+  constants <- kernelConstants <$> askEnv++  let ltid32 = kernelLocalThreadId constants+      ltid = sExt64 ltid32+      group_size = kernelGroupSize constants++  let (acc_params, next_params) = slugSplitParams slug+  let nes = slugNeutral slug+  let red_arrs = slugGroupRedArrs slug+  let group_res_arrs = groupResArrs slug++  old_counter <- dPrim "old_counter" int32+  (counter_mem, _, counter_offset) <-+    fullyIndexArray+      counters+      [counter_idx]+  sComment "first thread in group saves group result to global memory" $+    sWhen (ltid32 .==. 0) $ do+      forM_ (take (length nes) $ zip group_res_arrs (slugAccs slug)) $ \(v, (acc, acc_is)) ->+        copyDWIMFix v [0, sExt64 group_id] (Var acc) acc_is+      sOp $ Imp.MemFence Imp.FenceGlobal+      -- Increment the counter, thus stating that our result is+      -- available.+      sOp+        $ Imp.Atomic DefaultSpace+        $ Imp.AtomicAdd+          Int32+          (tvVar old_counter)+          counter_mem+          counter_offset+        $ untyped (1 :: Imp.TExp Int32)+      -- Now check if we were the last group to write our result.  If+      -- so, it is our responsibility to produce the final result.+      sWrite sync_arr [0] $ untyped $ tvExp old_counter .==. groups_per_segment - 1++  sOp $ Imp.Barrier Imp.FenceGlobal++  is_last_group <- dPrim "is_last_group" Bool+  copyDWIMFix (tvVar is_last_group) [] (Var sync_arr) [0]+  sWhen (tvExp is_last_group) $ do+    -- The final group has written its result (and it was+    -- us!), so read in all the group results and perform the+    -- final stage of the reduction.  But first, we reset the+    -- counter so it is ready for next time.  This is done+    -- with an atomic to avoid warnings about write/write+    -- races in oclgrind.+    sWhen (ltid32 .==. 0) $+      sOp $+        Imp.Atomic DefaultSpace $+          Imp.AtomicAdd Int32 (tvVar old_counter) counter_mem counter_offset $+            untyped $+              negate groups_per_segment++    sLoopNest (slugShape slug) $ \vec_is -> do+      unless (null $ slugShape slug) $+        sOp (Imp.Barrier Imp.FenceLocal)++      -- There is no guarantee that the number of workgroups for the+      -- segment is less than the workgroup size, so each thread may+      -- have to read multiple elements.  We do this in a sequential+      -- way that may induce non-coalesced accesses, but the total+      -- number of accesses should be tiny here.+      --+      -- TODO: here we *could* insert a collective copy of the num_groups+      -- per-group results. However, it may not be beneficial, since num_groups+      -- is not necessarily larger than group_size, meaning the number of+      -- uncoalesced reads here may be insignificant. In fact, if we happen to+      -- have a num_groups < group_size, then the collective copy would add+      -- unnecessary overhead. Also, this code is only executed by a single+      -- group.+      sComment "read in the per-group-results" $ do+        read_per_thread <-+          dPrimVE "read_per_thread" $+            groups_per_segment `divUp` sExt64 group_size++        forM2_ acc_params nes $ \p ne ->+          copyDWIM (paramName p) [] ne []++        sFor "i" read_per_thread $ \i -> do+          group_res_id <-+            dPrimVE "group_res_id" $+              ltid * read_per_thread + i+          index_of_group_res <-+            dPrimVE "index_of_group_res" $+              first_group_for_segment + group_res_id++          sWhen (group_res_id .<. groups_per_segment) $ do+            forM2_ next_params group_res_arrs $+              \p group_res_arr ->+                copyDWIMFix+                  (paramName p)+                  []+                  (Var group_res_arr)+                  ([0, index_of_group_res] ++ vec_is)++            compileStms mempty (bodyStms $ slugBody slug) $+              forM2_ acc_params (map resSubExp $ bodyResult $ slugBody slug) $ \p se ->+                copyDWIMFix (paramName p) [] se []++      forM2_ acc_params red_arrs $ \p arr ->+        when (isPrimParam p) $+          copyDWIMFix arr [ltid] (Var $ paramName p) []++      sOp $ Imp.ErrorSync Imp.FenceLocal++      sComment "reduce the per-group results" $ do+        groupReduce (sExt32 group_size) new_lambda red_arrs++        sComment "and back to memory with the final result" $+          sWhen (ltid32 .==. 0) $+            forM2_ segred_pes (lambdaParams new_lambda) $ \pe p ->+              copyDWIMFix+                (patElemName pe)+                (segment_gtids ++ vec_is)+                (Var $ paramName p)+                []++-- Note [IntermArrays]+--+-- Intermediate memory for the nonsegmented and large segments non-commutative+-- reductions with all primitive parameters:+--+--   These kernels need local memory for 1) the initial collective copy, 2) the+--   (virtualized) group reductions, and (TODO: this one not implemented yet!)+--   3) the final single-group collective copy. There are no dependencies+--   between these three stages, so we can reuse the same pool of local mem for+--   all three. These intermediates all go into local mem because of the+--   assumption of primitive parameter types.+--+--   Let `elem_sizes` be a list of element type sizes for the reduction+--   operators in a given redomap fusion. Then the amount of local mem needed+--   across the three steps are:+--+--   1) The initial collective copy from global to thread-private memory+--   requires `group_size * CHUNK * max elem_sizes`, since the collective copies+--   are performed in sequence (ie. inputs to different reduction operators need+--   not be held in local mem simultaneously).+--   2) The intra-group reductions of local memory held per-thread results+--   require `group_size * sum elem_sizes` bytes, since per-thread results for+--   all fused reductions are group-reduced simultaneously.+--   3) If group_size < num_groups, then after the final single-group collective+--   copy, a thread-sequential reduction reduces the number of per-group partial+--   results from num_groups down to group_size for each reduction array, such+--   that they will each fit in the final intra-group reduction. This requires+--   `num_groups * max elem_sizes`.+--+--   In summary, the total amount of local mem needed is the maximum between:+--   1) initial collective copy: group_size * CHUNK * max elem_sizes+--   2) intra-group reductions:  group_size * sum elem_sizes+--   3) final collective copy:   num_groups * max elem_sizes+--+--   The amount of local mem will most likely be decided by 1) in most cases,+--   unless the number of fused operators is very high *or* if we have a+--   `num_groups > group_size * CHUNK`, but this is unlikely, in which case 2)+--   and 3), respectively, will dominate.+--+--   Aside from local memory, these kernels also require a CHUNK-sized array of+--   thread-private register memory per reduction operator.+--+-- For all other reductions, ie. commutative reductions, reductions with at+-- least one non-primitive operator, and small segments reductions:+--+--   These kernels use local memory only for the intra-group reductions, and+--   since they do not use chunking or CHUNK, they all require onlly `group_size+--   * max elem_sizes` bytes of local memory and no thread-private register mem.

@@ -7,6 +7,7 @@  import Control.Monad import Data.List (zip4, zip7)+import Data.Map qualified as M import Data.Maybe import Futhark.CodeGen.ImpCode.GPU qualified as Imp import Futhark.CodeGen.ImpGen@@ -15,7 +16,7 @@ import Futhark.IR.Mem.LMAD qualified as LMAD import Futhark.Transform.Rename import Futhark.Util (mapAccumLM, takeLast)-import Futhark.Util.IntegralExp (IntegralExp (mod, rem), divUp, quot)+import Futhark.Util.IntegralExp (IntegralExp (mod, rem), divUp, nextMul, quot) import Prelude hiding (mod, quot, rem)  xParams, yParams :: SegBinOp GPUMem -> [LParam GPUMem]@@ -24,9 +25,6 @@ yParams scan =   drop (length (segBinOpNeutral scan)) (lambdaParams (segBinOpLambda scan)) -alignTo :: (IntegralExp a) => a -> a -> a-alignTo x a = (x `divUp` a) * a- createLocalArrays ::   Count GroupSize SubExp ->   SubExp ->@@ -36,7 +34,7 @@   let groupSizeE = pe64 groupSize       workSize = pe64 chunk * groupSizeE       prefixArraysSize =-        foldl (\acc tySize -> alignTo acc tySize + tySize * groupSizeE) 0 $+        foldl (\acc tySize -> nextMul acc tySize + tySize * groupSizeE) 0 $           map primByteSize types       maxTransposedArraySize =         foldl1 sMax64 $ map (\ty -> workSize * primByteSize ty) types@@ -44,7 +42,7 @@       warpSize :: (Num a) => a       warpSize = 32       maxWarpExchangeSize =-        foldl (\acc tySize -> alignTo acc tySize + tySize * fromInteger warpSize) 0 $+        foldl (\acc tySize -> nextMul acc tySize + tySize * fromInteger warpSize) 0 $           map primByteSize types       maxLookbackSize = maxWarpExchangeSize + warpSize       size = Imp.bytes $ maxLookbackSize `sMax64` prefixArraysSize `sMax64` maxTransposedArraySize@@ -52,7 +50,7 @@   (_, byteOffsets) <-     mapAccumLM       ( \off tySize -> do-          off' <- dPrimVE "byte_offsets" $ alignTo off tySize + pe64 groupSize * tySize+          off' <- dPrimVE "byte_offsets" $ nextMul off tySize + pe64 groupSize * tySize           pure (off', off)       )       0@@ -61,7 +59,7 @@   (_, warpByteOffsets) <-     mapAccumLM       ( \off tySize -> do-          off' <- dPrimVE "warp_byte_offset" $ alignTo off tySize + warpSize * tySize+          off' <- dPrimVE "warp_byte_offset" $ nextMul off tySize + warpSize * tySize           pure (off', off)       )       warpSize@@ -208,12 +206,10 @@       | otherwise =           copyDWIMFix (paramName p) [] (Var arr) [gtid - behind + arrs_full_size] -    writeResult x y arr-      | primType $ paramType x = do-          copyDWIMFix arr [ltid] (Var $ paramName x) []-          copyDWIM (paramName y) [] (Var $ paramName x) []-      | otherwise =-          copyDWIM (paramName y) [] (Var $ paramName x) []+    writeResult x y arr = do+      when (isPrimParam x) $+        copyDWIMFix arr [ltid] (Var $ paramName x) []+      copyDWIM (paramName y) [] (Var $ paramName x) []  -- | Compile 'SegScan' instance to host-level code with calls to a -- single-pass kernel.@@ -232,27 +228,17 @@        n = product $ map pe64 $ segSpaceDims space -      tys = map (\(Prim pt) -> pt) $ lambdaReturnType $ segBinOpLambda scan_op-      tys_sizes = map primByteSize tys--      sumT, maxT :: Integer-      sumT = sum tys_sizes-      sumT' = sum (map (max 4 . primByteSize) tys) `div` 4-      maxT = maximum tys_sizes--      -- TODO: Make these constants dynamic by querying device-      k_reg = 64-      k_mem = 95--      mem_constraint = max k_mem sumT `div` maxT-      reg_constraint = (k_reg - 1 - sumT') `div` (2 * sumT')+      tys' = lambdaReturnType $ segBinOpLambda scan_op -      chunk :: (Num a) => a-      chunk = fromIntegral $ max 1 $ min mem_constraint reg_constraint+      tys = map elemType tys'        group_size_e = pe64 $ unCount $ kAttrGroupSize attrs       num_physgroups_e = pe64 $ unCount $ kAttrNumGroups attrs +  let chunk_const = getChunkSize tys'+  chunk_v <- dPrimV "chunk_size" . isInt64 =<< kernelConstToExp chunk_const+  let chunk = tvExp chunk_v+   num_virtgroups <-     tvSize <$> dPrimV "num_virtgroups" (n `divUp` (group_size_e * chunk))   let num_virtgroups_e = pe64 num_virtgroups@@ -266,11 +252,7 @@       not_segmented_e = fromBool $ not segmented       segment_size = last dims' -  let debug_ s v = emit $ Imp.DebugPrint s $ Just $ untyped (v :: Imp.TExp Int32)-  debug_ "Sequential elements per thread (chunk) " chunk-  debug_ "Memory constraint" $ fromIntegral mem_constraint-  debug_ "Register constraint" $ fromIntegral reg_constraint-  debug_ "sumT'" $ fromIntegral sumT'+  emit $ Imp.DebugPrint "Sequential elements per thread (chunk)" $ Just $ untyped chunk    statusFlags <- sAllocArray "status_flags" int8 (Shape [num_virtgroups]) (Space "device")   sReplicate statusFlags $ intConst Int8 statusX@@ -284,11 +266,18 @@    global_id <- genZeroes "global_dynid" 1 -  sKernelThread "segscan" (segFlat space) attrs $ do+  let attrs' = attrs {kAttrConstExps = M.singleton (tvVar chunk_v) chunk_const}++  sKernelThread "segscan" (segFlat space) attrs' $ do+    chunk32 <- dPrimVE "chunk_size_32b" $ sExt32 $ tvExp chunk_v+     constants <- kernelConstants <$> askEnv +    let ltid32 = kernelLocalThreadId constants+        ltid = sExt64 ltid32+     (sharedId, transposedArrays, prefixArrays, warpscan, exchanges) <--      createLocalArrays (kAttrGroupSize attrs) (intConst Int64 chunk) tys+      createLocalArrays (kAttrGroupSize attrs) (tvSize chunk_v) tys      -- We wrap the entire kernel body in a virtualisation loop to handle the     -- case where we do not have enough workgroups to cover the iteration space.@@ -308,7 +297,7 @@     sFor "virtloop_i" iters $ const $ do       dyn_id <- dPrim "dynamic_id" int32       sComment "First thread in block fetches this block's dynamic_id" $ do-        sWhen (kernelLocalThreadId constants .==. 0) $ do+        sWhen (ltid32 .==. 0) $ do           (globalIdMem, _, globalIdOff) <- fullyIndexArray global_id [0]           sOp $             Imp.Atomic DefaultSpace $@@ -333,34 +322,32 @@       copyDWIMFix (tvVar dyn_id) [] (Var sharedId) [0]       sOp local_barrier -      blockOff <--        dPrimV "blockOff" $-          sExt64 (tvExp dyn_id) * chunk * group_size_e -- kernelGroupSize constants-      sgmIdx <- dPrimVE "sgm_idx" $ tvExp blockOff `mod` segment_size+      block_offset <-+        dPrimVE "block_offset" $+          sExt64 (tvExp dyn_id) * chunk * group_size_e+      sgm_idx <- dPrimVE "sgm_idx" $ block_offset `mod` segment_size       boundary <-         dPrimVE "boundary" $           sExt32 $-            sMin64 (chunk * group_size_e) (segment_size - sgmIdx)+            sMin64 (chunk * group_size_e) (segment_size - sgm_idx)       segsize_compact <-         dPrimVE "segsize_compact" $           sExt32 $             sMin64 (chunk * group_size_e) segment_size-      privateArrays <-+      private_chunks <-         forM tys $ \ty ->           sAllocArray             "private"             ty-            (Shape [intConst Int64 chunk])-            (ScalarSpace [intConst Int64 chunk] ty)+            (Shape [tvSize chunk_v])+            (ScalarSpace [tvSize chunk_v] ty) +      thd_offset <- dPrimVE "thd_offset" $ block_offset + ltid+       sComment "Load and map" $         sFor "i" chunk $ \i -> do           -- The map's input index-          virt_tid <--            dPrimVE "virt_tid" $-              tvExp blockOff-                + sExt64 (kernelLocalThreadId constants)-                + i * kernelGroupSize constants+          virt_tid <- dPrimVE "virt_tid" $ thd_offset + i * group_size_e           dIndexSpace (zip gtids dims') virt_tid           -- Perform the map           let in_bounds =@@ -373,66 +360,62 @@                     copyDWIMFix (patElemName dest) (map Imp.le64 gtids) (kernelResultSubExp src) []                    -- Write to-scan results to private memory.-                  forM_ (zip privateArrays $ map kernelResultSubExp all_scan_res) $ \(dest, src) ->+                  forM_ (zip private_chunks $ map kernelResultSubExp all_scan_res) $ \(dest, src) ->                     copyDWIMFix dest [i] src []                out_of_bounds =-                forM_ (zip privateArrays scanop_nes) $ \(dest, ne) ->+                forM_ (zip private_chunks scanop_nes) $ \(dest, ne) ->                   copyDWIMFix dest [i] ne []            sIf (virt_tid .<. n) in_bounds out_of_bounds        sOp $ Imp.ErrorSync Imp.FenceLocal       sComment "Transpose scan inputs" $ do-        forM_ (zip transposedArrays privateArrays) $ \(trans, priv) -> do+        forM_ (zip transposedArrays private_chunks) $ \(trans, priv) -> do           sFor "i" chunk $ \i -> do-            sharedIdx <--              dPrimVE "sharedIdx" $-                sExt64 (kernelLocalThreadId constants)-                  + i * kernelGroupSize constants+            sharedIdx <- dPrimVE "sharedIdx" $ ltid + i * group_size_e             copyDWIMFix trans [sharedIdx] (Var priv) [i]           sOp local_barrier           sFor "i" chunk $ \i -> do-            sharedIdx <- dPrimV "sharedIdx" $ kernelLocalThreadId constants * chunk + i+            sharedIdx <- dPrimV "sharedIdx" $ ltid * chunk + i             copyDWIMFix priv [sExt64 i] (Var trans) [sExt64 $ tvExp sharedIdx]           sOp local_barrier        sComment "Per thread scan" $ do         -- We don't need to touch the first element, so only m-1         -- iterations here.-        globalIdx <--          dPrimVE "gidx" $-            (kernelLocalThreadId constants * chunk) + 1         sFor "i" (chunk - 1) $ \i -> do           let xs = map paramName $ xParams scan_op               ys = map paramName $ yParams scan_op           -- determine if start of segment           new_sgm <-             if segmented-              then dPrimVE "new_sgm" $ (globalIdx + sExt32 i - boundary) `mod` segsize_compact .==. 0+              then do+                gidx <- dPrimVE "gidx" $ (ltid32 * chunk32) + 1+                dPrimVE "new_sgm" $ (gidx + sExt32 i - boundary) `mod` segsize_compact .==. 0               else pure false           -- skip scan of first element in segment           sUnless new_sgm $ do-            forM_ (zip4 privateArrays xs ys tys) $ \(src, x, y, ty) -> do+            forM_ (zip4 private_chunks xs ys tys) $ \(src, x, y, ty) -> do               dPrim_ x ty               dPrim_ y ty               copyDWIMFix x [] (Var src) [i]               copyDWIMFix y [] (Var src) [i + 1]              compileStms mempty (bodyStms $ lambdaBody $ segBinOpLambda scan_op) $-              forM_ (zip privateArrays $ map resSubExp $ bodyResult $ lambdaBody $ segBinOpLambda scan_op) $ \(dest, res) ->+              forM_ (zip private_chunks $ map resSubExp $ bodyResult $ lambdaBody $ segBinOpLambda scan_op) $ \(dest, res) ->                 copyDWIMFix dest [i + 1] res []        sComment "Publish results in shared memory" $ do-        forM_ (zip prefixArrays privateArrays) $ \(dest, src) ->-          copyDWIMFix dest [sExt64 $ kernelLocalThreadId constants] (Var src) [chunk - 1]+        forM_ (zip prefixArrays private_chunks) $ \(dest, src) ->+          copyDWIMFix dest [ltid] (Var src) [chunk - 1]         sOp local_barrier        let crossesSegment = do             guard segmented             Just $ \from to ->-              let from' = (from + 1) * chunk - 1-                  to' = (to + 1) * chunk - 1+              let from' = (from + 1) * chunk32 - 1+                  to' = (to + 1) * chunk32 - 1                in (to' - from') .>. (to' + segsize_compact - boundary) `mod` segsize_compact        scan_op1 <- renameLambda $ segBinOpLambda scan_op@@ -441,19 +424,19 @@       sComment "Scan results (with warp scan)" $ do         groupScan           crossesSegment+          group_size_e           num_virt_threads-          (kernelGroupSize constants)           scan_op1           prefixArrays          sOp $ Imp.ErrorSync Imp.FenceLocal          let firstThread acc prefixes =-              copyDWIMFix (tvVar acc) [] (Var prefixes) [sExt64 (kernelGroupSize constants) - 1]+              copyDWIMFix (tvVar acc) [] (Var prefixes) [sExt64 group_size_e - 1]             notFirstThread acc prefixes =-              copyDWIMFix (tvVar acc) [] (Var prefixes) [sExt64 (kernelLocalThreadId constants) - 1]+              copyDWIMFix (tvVar acc) [] (Var prefixes) [ltid - 1]         sIf-          (kernelLocalThreadId constants .==. 0)+          (ltid32 .==. 0)           (zipWithM_ firstThread accs prefixArrays)           (zipWithM_ notFirstThread accs prefixArrays) @@ -461,9 +444,9 @@        prefixes <- forM (zip scanop_nes tys) $ \(ne, ty) ->         dPrimV "prefix" $ TPrimExp $ toExp' ty ne-      blockNewSgm <- dPrimVE "block_new_sgm" $ sgmIdx .==. 0+      blockNewSgm <- dPrimVE "block_new_sgm" $ sgm_idx .==. 0       sComment "Perform lookback" $ do-        sWhen (blockNewSgm .&&. kernelLocalThreadId constants .==. 0) $ do+        sWhen (blockNewSgm .&&. ltid32 .==. 0) $ do           everythingVolatile $             forM_ (zip accs incprefixArrays) $ \(acc, incprefixArray) ->               copyDWIMFix incprefixArray [tvExp dyn_id] (tvSize acc) []@@ -475,8 +458,8 @@         -- end sWhen          let warpSize = kernelWaveSize constants-        sWhen (bNot blockNewSgm .&&. kernelLocalThreadId constants .<. warpSize) $ do-          sWhen (kernelLocalThreadId constants .==. 0) $ do+        sWhen (bNot blockNewSgm .&&. ltid32 .<. warpSize) $ do+          sWhen (ltid32 .==. 0) $ do             sIf               (not_segmented_e .||. boundary .==. sExt32 (group_size_e * chunk))               ( do@@ -505,7 +488,7 @@            sIf             (tvExp status .==. statusP)-            ( sWhen (kernelLocalThreadId constants .==. 0) $+            ( sWhen (ltid32 .==. 0) $                 everythingVolatile $                   forM_ (zip prefixes incprefixArrays) $ \(prefix, incprefixArray) ->                     copyDWIMFix (tvVar prefix) [] (Var incprefixArray) [tvExp dyn_id - 1]@@ -518,11 +501,11 @@                 let loopStop = warpSize * (-1)                     sameSegment readIdx                       | segmented =-                          let startIdx = sExt64 (tvExp readIdx + 1) * kernelGroupSize constants * chunk - 1-                           in tvExp blockOff - startIdx .<=. sgmIdx+                          let startIdx = sExt64 (tvExp readIdx + 1) * group_size_e * chunk - 1+                           in block_offset - startIdx .<=. sgm_idx                       | otherwise = true                 sWhile (tvExp readOffset .>. loopStop) $ do-                  readI <- dPrimV "read_i" $ tvExp readOffset + kernelLocalThreadId constants+                  readI <- dPrimV "read_i" $ tvExp readOffset + ltid32                   aggrs <- forM (zip scanop_nes tys) $ \(ne, ty) ->                     dPrimV "aggr" $ TPrimExp $ toExp' ty ne                   flag <- dPrimV "flag" (statusX :: Imp.TExp Int8)@@ -546,12 +529,12 @@                   -- end sWhen                    forM_ (zip exchanges aggrs) $ \(exchange, aggr) ->-                    copyDWIMFix exchange [sExt64 $ kernelLocalThreadId constants] (tvSize aggr) []-                  copyDWIMFix warpscan [sExt64 $ kernelLocalThreadId constants] (tvSize flag) []+                    copyDWIMFix exchange [ltid] (tvSize aggr) []+                  copyDWIMFix warpscan [ltid] (tvSize flag) []                    -- execute warp-parallel reduction but only if the last read flag in not STATUS_P                   copyDWIMFix (tvVar flag) [] (Var warpscan) [sExt64 warpSize - 1]-                  sWhen (tvExp flag .<. (2 :: Imp.TExp Int8)) $ do+                  sWhen (tvExp flag .<. statusP) $ do                     lam' <- renameLambda scan_op1                     inBlockScanLookback                       constants@@ -573,7 +556,7 @@                     )                    -- update prefix if flag different than STATUS_X:-                  sWhen (tvExp flag .>. (statusX :: Imp.TExp Int8)) $ do+                  sWhen (tvExp flag .>. statusX) $ do                     lam <- renameLambda scan_op1                     let (xs, ys) = splitAt (length tys) $ map paramName $ lambdaParams lam                     forM_ (zip xs aggrs) $ \(x, aggr) -> dPrimV_ x (tvExp aggr)@@ -586,7 +569,7 @@            -- end sWhile           -- end sIf-          sWhen (kernelLocalThreadId constants .==. 0) $ do+          sWhen (ltid32 .==. 0) $ do             scan_op2 <- renameLambda scan_op1             let xs = map paramName $ take (length tys) $ lambdaParams scan_op2                 ys = map paramName $ drop (length tys) $ lambdaParams scan_op2@@ -629,7 +612,7 @@             dPrimV_ y' $ tvExp acc          sIf-          (kernelLocalThreadId constants * chunk .<. boundary .&&. bNot blockNewSgm)+          (ltid32 * chunk32 .<. boundary .&&. bNot blockNewSgm)           ( compileStms mempty (bodyStms $ lambdaBody scan_op4) $               forM_ (zip3 xs tys $ map resSubExp $ bodyResult $ lambdaBody scan_op4) $                 \(x, ty, res) -> x <~~ toExp' ty res@@ -639,38 +622,34 @@         -- elements left before new segment.         stop <-           dPrimVE "stopping_point" $-            segsize_compact - (kernelLocalThreadId constants * chunk - 1 + segsize_compact - boundary) `rem` segsize_compact+            segsize_compact - (ltid32 * chunk32 - 1 + segsize_compact - boundary) `rem` segsize_compact         sFor "i" chunk $ \i -> do           sWhen (sExt32 i .<. stop - 1) $ do-            forM_ (zip privateArrays ys) $ \(src, y) ->+            forM_ (zip private_chunks ys) $ \(src, y) ->               -- only include prefix for the first segment part per thread               copyDWIMFix y [] (Var src) [i]             compileStms mempty (bodyStms $ lambdaBody scan_op3) $-              forM_ (zip privateArrays $ map resSubExp $ bodyResult $ lambdaBody scan_op3) $+              forM_ (zip private_chunks $ map resSubExp $ bodyResult $ lambdaBody scan_op3) $                 \(dest, res) ->                   copyDWIMFix dest [i] res []        sComment "Transpose scan output and Write it to global memory in coalesced fashion" $ do-        forM_ (zip3 transposedArrays privateArrays $ map patElemName all_pes) $ \(locmem, priv, dest) -> do+        forM_ (zip3 transposedArrays private_chunks $ map patElemName all_pes) $ \(locmem, priv, dest) -> do           -- sOp local_barrier           sFor "i" chunk $ \i -> do             sharedIdx <-               dPrimV "sharedIdx" $-                sExt64 (kernelLocalThreadId constants * chunk) + i+                sExt64 (ltid * chunk) + i             copyDWIMFix locmem [tvExp sharedIdx] (Var priv) [i]           sOp local_barrier           sFor "i" chunk $ \i -> do-            flat_idx <--              dPrimVE "flat_idx" $-                tvExp blockOff-                  + kernelGroupSize constants * i-                  + sExt64 (kernelLocalThreadId constants)+            flat_idx <- dPrimVE "flat_idx" $ thd_offset + i * group_size_e             dIndexSpace (zip gtids dims') flat_idx             sWhen (flat_idx .<. n) $ do               copyDWIMFix                 dest                 (map Imp.le64 gtids)                 (Var locmem)-                [sExt64 $ flat_idx - tvExp blockOff]+                [sExt64 $ flat_idx - block_offset]           sOp local_barrier {-# NOINLINE compileSegScan #-}

@@ -10,7 +10,6 @@ where  import Control.Monad-import Control.Monad.Identity import Control.Monad.Reader import Control.Monad.State import Data.Bifunctor (second)@@ -58,7 +57,7 @@ translateGPU target prog =   let env = envFromProg prog       ( prog',-        ToOpenCL kernels device_funs used_types sizes failures+        ToOpenCL kernels device_funs used_types sizes failures constants         ) =           (`runState` initialOpenCL) . (`runReaderT` env) $ do             let ImpGPU.Definitions@@ -86,13 +85,15 @@             T.unlines device_defs           ]    in ImpOpenCL.Program-        opencl_code-        opencl_prelude-        kernels'-        (S.toList used_types)-        (findParamUsers env prog' (cleanSizes sizes))-        failures-        prog'+        { openClProgram = opencl_code,+          openClPrelude = opencl_prelude,+          openClMacroDefs = constants,+          openClKernelNames = kernels',+          openClUsedTypes = S.toList used_types,+          openClParams = findParamUsers env prog' (cleanSizes sizes),+          openClFailures = failures,+          hostDefinitions = prog'+        }   where     genPrelude TargetOpenCL = genOpenClPrelude     genPrelude TargetCUDA = const genCUDAPrelude@@ -172,11 +173,12 @@     clDevFuns :: M.Map Name (C.Definition, T.Text),     clUsedTypes :: S.Set PrimType,     clSizes :: M.Map Name SizeClass,-    clFailures :: [FailureMsg]+    clFailures :: [FailureMsg],+    clConstants :: [(Name, KernelConstExp)]   }  initialOpenCL :: ToOpenCL-initialOpenCL = ToOpenCL mempty mempty mempty mempty mempty+initialOpenCL = ToOpenCL mempty mempty mempty mempty mempty mempty  data Env = Env   { envFuns :: ImpGPU.Functions ImpGPU.HostOp,@@ -331,7 +333,7 @@ isConst :: GroupDim -> Maybe T.Text isConst (Left (ValueExp (IntValue x))) =   Just $ prettyText $ intToInt64 x-isConst (Right (SizeConst v)) =+isConst (Right (SizeConst v _)) =   Just $ zEncodeText $ nameToText v isConst (Right (SizeMaxConst size_class)) =   Just $ "max_" <> prettyText size_class@@ -354,7 +356,8 @@           mapM_ GC.item body       kstate = GC.compUserState cstate -      (const_defs, const_undefs) = unzip $ mapMaybe constDef $ kernelUses kernel+      (kernel_consts, (const_defs, const_undefs)) =+        second unzip $ unzip $ mapMaybe (constDef (kernelName kernel)) $ kernelUses kernel    let (local_memory_bytes, (local_memory_params, local_memory_args, local_memory_init)) =         second unzip3 $@@ -458,7 +461,8 @@     s       { clGPU = M.insert name (safety, kernel_fun) $ clGPU s,         clUsedTypes = typesInKernel kernel <> clUsedTypes s,-        clFailures = kernelFailures kstate+        clFailures = kernelFailures kstate,+        clConstants = kernel_consts <> clConstants s       }    -- The error handling stuff is automatically added later.@@ -506,17 +510,19 @@ -- Constants are #defined as macros.  Since a constant name in one -- kernel might potentially (although unlikely) also be used for -- something else in another kernel, we #undef them after the kernel.-constDef :: KernelUse -> Maybe (C.BlockItem, C.BlockItem)-constDef (ConstUse v e) =+constDef :: Name -> KernelUse -> Maybe ((Name, KernelConstExp), (C.BlockItem, C.BlockItem))+constDef kernel_name (ConstUse v e) =   Just-    ( [C.citem|$escstm:(T.unpack def)|],-      [C.citem|$escstm:(T.unpack undef)|]+    ( (nameFromText v', e),+      ( [C.citem|$escstm:(T.unpack def)|],+        [C.citem|$escstm:(T.unpack undef)|]+      )     )   where-    e' = compilePrimExp e-    def = "#define " <> idText (C.toIdent v mempty) <> " (" <> expText e' <> ")"+    v' = zEncodeText $ nameToText kernel_name <> "." <> prettyText v+    def = "#define " <> idText (C.toIdent v mempty) <> " (" <> v' <> ")"     undef = "#undef " <> idText (C.toIdent v mempty)-constDef _ = Nothing+constDef _ _ = Nothing  commonPrelude :: T.Text commonPrelude =@@ -549,14 +555,6 @@   "#define FUTHARK_HIP\n"     <> preludeCU     <> commonPrelude--compilePrimExp :: PrimExp KernelConst -> C.Exp-compilePrimExp e = runIdentity $ GC.compilePrimExp compileKernelConst e-  where-    compileKernelConst (SizeConst key) =-      pure [C.cexp|$id:(zEncodeText (prettyText key))|]-    compileKernelConst (SizeMaxConst size_class) =-      pure [C.cexp|$id:("max_" <> prettyString size_class)|]  kernelArgs :: Kernel -> [KernelArg] kernelArgs = mapMaybe useToArg . kernelUses

@@ -37,7 +37,6 @@ import Futhark.CodeGen.ImpGen import Futhark.Error import Futhark.IR.MCMem-import Futhark.MonadFreshNames import Futhark.Transform.Rename import Prelude hiding (quot, rem) 

@@ -11,7 +11,6 @@ import Futhark.CodeGen.ImpGen.Multicore.Base import Futhark.CodeGen.ImpGen.Multicore.SegRed (compileSegRed') import Futhark.IR.MCMem-import Futhark.MonadFreshNames import Futhark.Transform.Rename (renameLambda) import Futhark.Util (chunks, splitFromEnd, takeLast) import Futhark.Util.IntegralExp (rem)

@@ -155,7 +155,6 @@   -- First we try to find a file of the given name in the search path,   -- then we look at the builtin library if we have to.  For the   -- builtins, we don't use the search path.-  let filepath = includeToFilePath include   r <- contentsAndModTime filepath vfs   case (r, lookup prelude_str prelude) of     (Just (Right (s, mod_time)), _) ->@@ -167,6 +166,8 @@     (Nothing, Nothing) ->       pure $ Left $ ProgError loc $ pretty not_found   where+    filepath = includeToFilePath include+     prelude_str = "/" Posix.</> includeToString include Posix.<.> "fut"      loaded path s mod_time =@@ -178,7 +179,7 @@         }      not_found =-      "Could not find import " <> E.quote (includeToText include) <> "."+      "Could not read file " <> E.quote (T.pack filepath) <> "."  handleFile :: ReaderState -> VFS -> LoadedFile T.Text -> IO UncheckedImport handleFile state_mvar vfs (LoadedFile file_name import_name file_contents mod_time) = do

@@ -124,7 +124,7 @@          forEnv env =           mconcat (map vname' $ M.toList $ envNameMap env)-            <> mconcat (map forMty $ M.elems $ envSigTable env)+            <> mconcat (map forMty $ M.elems $ envModTypeTable env)         forMod (ModEnv env) = forEnv env         forMod ModFun {} = mempty         forMty = forMod . mtyMod@@ -373,7 +373,7 @@  synopsisDec :: S.Set VName -> FileModule -> Dec -> Maybe (DocM Html) synopsisDec visible fm dec = case dec of-  SigDec s -> synopsisModType mempty s+  ModTypeDec s -> synopsisModType mempty s   ModDec m -> synopsisMod fm m   ValDec v -> synopsisValBind v   TypeDec t -> synopsisType t@@ -386,8 +386,8 @@           pure $             fullRow $               keyword "open" <> fromString (" <" <> prettyString x <> ">")-  LocalDec (SigDec s) _-    | sigName s `S.member` visible ->+  LocalDec (ModTypeDec s) _+    | modTypeName s `S.member` visible ->         synopsisModType (keyword "local" <> " ") s   LocalDec {} -> Nothing   ImportDec {} -> Nothing@@ -404,7 +404,7 @@     keyword "import "       <> (H.a ! A.href dest) (fromString (show (includeToString file))) synopsisOpened (ModAscript _ se _ _) = Just $ do-  se' <- synopsisSigExp se+  se' <- synopsisModTypeExp se   pure $ "... : " <> se' synopsisOpened _ = Nothing @@ -427,18 +427,18 @@       mconcat (intersperse " -> " $ params' ++ [rettype'])     ) -synopsisModType :: Html -> SigBind -> Maybe (DocM Html)+synopsisModType :: Html -> ModTypeBind -> Maybe (DocM Html) synopsisModType prefix sb = Just $ do-  let name' = vnameSynopsisDef $ sigName sb+  let name' = vnameSynopsisDef $ modTypeName sb   fullRow <$> do-    se' <- synopsisSigExp $ sigExp sb+    se' <- synopsisModTypeExp $ modTypeExp sb     pure $ prefix <> keyword "module type " <> name' <> " = " <> se'  synopsisMod :: FileModule -> ModBind -> Maybe (DocM Html) synopsisMod fm (ModBind name ps sig _ _ _) =   case sig of-    Nothing -> (proceed <=< envSig) <$> M.lookup name modtable-    Just (s, _) -> Just $ proceed =<< synopsisSigExp s+    Nothing -> (proceed <=< envModType) <$> M.lookup name modtable+    Just (s, _) -> Just $ proceed =<< synopsisModTypeExp s   where     proceed sig' = do       let name' = vnameSynopsisDef name@@ -446,8 +446,8 @@       pure $ specRow (keyword "module " <> name') ": " (ps' <> sig')      FileModule _abs Env {envModTable = modtable} _ _ = fm-    envSig (ModEnv e) = renderEnv e-    envSig (ModFun (FunSig _ _ (MTy _ m))) = envSig m+    envModType (ModEnv e) = renderEnv e+    envModType (ModFun (FunModType _ _ (MTy _ m))) = envModType m  synopsisType :: TypeBind -> Maybe (DocM Html) synopsisType tb = Just $ do@@ -548,7 +548,7 @@ modParamHtml :: [ModParamBase Info VName] -> DocM Html modParamHtml [] = pure mempty modParamHtml (ModParam pname psig _ _ : mps) =-  liftM2 f (synopsisSigExp psig) (modParamHtml mps)+  liftM2 f (synopsisModTypeExp psig) (modParamHtml mps)   where     f se params =       "("@@ -558,26 +558,26 @@         <> ") -> "         <> params -synopsisSigExp :: SigExpBase Info VName -> DocM Html-synopsisSigExp e = case e of-  SigVar v _ _ -> qualNameHtml v-  SigParens e' _ -> parens <$> synopsisSigExp e'-  SigSpecs ss _ -> braces . (H.table ! A.class_ "specs") . mconcat <$> mapM synopsisSpec ss-  SigWith s (TypeRef v ps t _) _ -> do-    s' <- synopsisSigExp s+synopsisModTypeExp :: ModTypeExpBase Info VName -> DocM Html+synopsisModTypeExp e = case e of+  ModTypeVar v _ _ -> qualNameHtml v+  ModTypeParens e' _ -> parens <$> synopsisModTypeExp e'+  ModTypeSpecs ss _ -> braces . (H.table ! A.class_ "specs") . mconcat <$> mapM synopsisSpec ss+  ModTypeWith s (TypeRef v ps t _) _ -> do+    s' <- synopsisModTypeExp s     t' <- typeExpHtml t     v' <- qualNameHtml v     let ps' = mconcat $ map ((" " <>) . typeParamHtml) ps     pure $ s' <> keyword " with " <> v' <> ps' <> " = " <> t'-  SigArrow Nothing e1 e2 _ ->-    liftM2 f (synopsisSigExp e1) (synopsisSigExp e2)+  ModTypeArrow Nothing e1 e2 _ ->+    liftM2 f (synopsisModTypeExp e1) (synopsisModTypeExp e2)     where       f e1' e2' = e1' <> " -> " <> e2'-  SigArrow (Just v) e1 e2 _ ->+  ModTypeArrow (Just v) e1 e2 _ ->     do       let name = vnameHtml v-      e1' <- synopsisSigExp e1-      e2' <- noLink [v] $ synopsisSigExp e2+      e1' <- synopsisModTypeExp e1+      e2' <- noLink [v] $ synopsisModTypeExp e2       pure $ "(" <> name <> ": " <> e1' <> ") -> " <> e2'  keyword :: String -> Html@@ -625,8 +625,8 @@         (mconcat (map (" " <>) tparams') <> ": ")         rettype'   ModSpec name sig _ _ ->-    specRow (keyword "module " <> vnameSynopsisDef name) ": " <$> synopsisSigExp sig-  IncludeSpec e _ -> fullRow . (keyword "include " <>) <$> synopsisSigExp e+    specRow (keyword "module " <> vnameSynopsisDef name) ": " <$> synopsisModTypeExp sig+  IncludeSpec e _ -> fullRow . (keyword "include " <>) <$> synopsisModTypeExp e  typeExpHtml :: TypeExp Info VName -> DocM Html typeExpHtml e = case e of@@ -802,7 +802,7 @@ describeGenericMod ::   VName ->   IndexWhat ->-  SigExp ->+  ModTypeExp ->   Maybe DocComment ->   (Html -> DocM Html) ->   DocM Html@@ -812,7 +812,7 @@   decl_type <- f name'    doc' <- case se of-    SigSpecs specs _ -> (<>) <$> docHtml doc <*> describeSpecs specs+    ModTypeSpecs specs _ -> (<>) <$> docHtml doc <*> describeSpecs specs     _ -> docHtml doc    let decl_doc = H.dd ! A.class_ "desc_doc" $ doc'@@ -837,14 +837,14 @@ describeDec _ (TypeDec vb) =   Just $     describeGeneric (typeAlias vb) IndexType (typeDoc vb) (`typeBindHtml` vb)-describeDec _ (SigDec (SigBind name se doc _)) = Just $+describeDec _ (ModTypeDec (ModTypeBind name se doc _)) = Just $   describeGenericMod name IndexModuleType se doc $ \name' ->     pure $ keyword "module type " <> name' describeDec _ (ModDec mb) = Just $   describeGeneric (modName mb) IndexModule (modDoc mb) $ \name' ->     pure $ keyword "module " <> name' describeDec _ OpenDec {} = Nothing-describeDec visible (LocalDec (SigDec (SigBind name se doc _)) _)+describeDec visible (LocalDec (ModTypeDec (ModTypeBind name se doc _)) _)   | name `S.member` visible = Just $       describeGenericMod name IndexModuleType se doc $ \name' ->         pure $ keyword "local module type " <> name'@@ -883,12 +883,12 @@ describeSpec (ModSpec name se doc _) =   describeGenericMod name IndexModule se doc $ \name' ->     case se of-      SigSpecs {} -> pure $ keyword "module " <> name'+      ModTypeSpecs {} -> pure $ keyword "module " <> name'       _ -> do-        se' <- synopsisSigExp se+        se' <- synopsisModTypeExp se         pure $ keyword "module " <> name' <> ": " <> se' describeSpec (IncludeSpec sig _) = do-  sig' <- synopsisSigExp sig+  sig' <- synopsisModTypeExp sig   doc' <- docHtml Nothing   let decl_header =         (H.dt ! A.class_ "desc_header") $

@@ -38,6 +38,12 @@     SizeLocalMemory   | -- | A bespoke size with a default.     SizeBespoke Name Int64+  | -- | Amount of registers available per workgroup. Mostly+    -- meaningful for querying the maximum.+    SizeRegisters+  | -- | Amount of L2 cache memory, in bytes. Mostly meaningful for+    -- querying the maximum.+    SizeCache   deriving (Eq, Ord, Show)  instance Pretty SizeClass where@@ -54,6 +60,8 @@   pretty SizeLocalMemory = "local_memory"   pretty (SizeBespoke k def) =     "bespoke" <> parens (pretty k <> comma <+> pretty def)+  pretty SizeRegisters = "registers"+  pretty SizeCache = "cache"  -- | The default value for the size.  If 'Nothing', that means the backend gets to decide. sizeDefault :: SizeClass -> Maybe Int64

@@ -1157,8 +1157,8 @@ simplifySegOp (SegHist lvl space ops ts kbody) = do   (lvl', space', ts') <- Engine.simplify (lvl, space, ts) -  (ops', ops_hoisted) <- fmap unzip $-    forM ops $+  Engine.localVtable (flip (foldr ST.consume) $ concatMap histDest ops) $ do+    (ops', ops_hoisted) <- fmap unzip . forM ops $       \(HistOp w rf arrs nes dims lam) -> do         w' <- Engine.simplify w         rf' <- Engine.simplify rf@@ -1174,12 +1174,12 @@             op_hoisted           ) -  (kbody', body_hoisted) <- simplifyKernelBody space kbody+    (kbody', body_hoisted) <- simplifyKernelBody space kbody -  pure-    ( SegHist lvl' space' ops' ts' kbody',-      mconcat ops_hoisted <> body_hoisted-    )+    pure+      ( SegHist lvl' space' ops' ts' kbody',+        mconcat ops_hoisted <> body_hoisted+      )   where     scope = scopeOfSegSpace space     scope_vtable = ST.fromScope scope

@@ -148,7 +148,7 @@  maybeAscript ::   SrcLoc ->-  Maybe (SigExp, Info (M.Map VName VName)) ->+  Maybe (ModTypeExp, Info (M.Map VName VName)) ->   ModExp ->   ModExp maybeAscript loc (Just (mtye, substs)) me = ModAscript me mtye substs loc@@ -311,7 +311,7 @@     evalModExp       $ foldr         addParam-        (maybeAscript (srclocOf mb) (modSignature mb) $ modExp mb)+        (maybeAscript (srclocOf mb) (modType mb) $ modExp mb)       $ modParams mb   mname <- transformName $ modName mb   pure $ Scope (scopeSubsts $ modScope mod) $ M.singleton mname mod@@ -331,7 +331,7 @@       bindingNames [typeAlias tb] $ do         transformTypeBind tb         transformDecs ds'-    SigDec {} : ds' ->+    ModTypeDec {} : ds' ->       transformDecs ds'     ModDec mb : ds' ->       bindingNames [modName mb] $ do

@@ -66,18 +66,13 @@         file_path = uriToFilePath $ doc ^. uri     tryReCompile state_mvar file_path -onDocumentOpenHandler :: IORef State -> Handlers (LspM ())-onDocumentOpenHandler state_mvar =-  notificationHandler SMethod_TextDocumentDidOpen $ \msg -> do-    let TNotificationMessage _ _ (DidOpenTextDocumentParams doc) = msg-        file_path = uriToFilePath $ doc ^. uri-    logStringStderr <& ("Opened document: " ++ show (doc ^. uri))-    tryReCompile state_mvar file_path+-- Some clients (Eglot) sends open/close events whether we want them+-- or not, so we better be prepared to ignore them.+onDocumentOpenHandler :: Handlers (LspM ())+onDocumentOpenHandler = notificationHandler SMethod_TextDocumentDidOpen $ \_ -> pure ()  onDocumentCloseHandler :: Handlers (LspM ())-onDocumentCloseHandler =-  notificationHandler SMethod_TextDocumentDidClose $ \_msg ->-    logStringStderr <& "Closed document"+onDocumentCloseHandler = notificationHandler SMethod_TextDocumentDidClose $ \_msg -> pure ()  -- Sent by Eglot when first connecting - not sure when else it might -- be sent.@@ -93,7 +88,7 @@ handlers state_mvar _ =   mconcat     [ onInitializeHandler,-      onDocumentOpenHandler state_mvar,+      onDocumentOpenHandler,       onDocumentCloseHandler,       onDocumentSaveHandler state_mvar,       onDocumentChangeHandler state_mvar,

@@ -1,416 +0,0 @@-{-# LANGUAGE TypeFamilies #-}-{-# LANGUAGE UndecidableInstances #-}---- | This module implements an optimisation that moves in-place--- updates into/before loops where possible, with the end goal of--- minimising memory copies.  As an example, consider this program:------ @---   let r =---     loop (r1 = r0) = for i < n do---       let a = r1[i]---       let r1[i] = a * i---       in r1---   ...---   let x = y with [k] <- r in---   ...--- @------ We want to turn this into the following:------ @---   let x0 = y with [k] <- r0---   loop (x = x0) = for i < n do---     let a = a[k,i]---     let x[k,i] = a * i---     in x---   let r = x[k] in---   ...--- @------ The intent is that we are also going to optimise the new data--- movement (in the @x0@-binding), possibly by changing how @r0@ is--- defined.  For the above transformation to be valid, a number of--- conditions must be fulfilled:------    (1) @r@ must not be consumed after the original in-place update.------    (2) @k@ and @y@ must be available at the beginning of the loop.------    (3) @x@ must be visible whenever @r@ is visible.  (This means---    that both @x@ and @r@ must be bound in the same t'Body'.)------    (4) If @x@ is consumed at a point after the loop, @r@ must not---    be used after that point.------    (5) The size of @r1@ is invariant inside the loop.------    (6) The value @r@ must come from something that we can actually---    optimise (e.g. not a function parameter).------    (7) @y@ (or its aliases) may not be used inside the body of the---    loop.------    (8) The result of the loop may not alias the merge parameter---    @r1@.------    (9) @y@ or its aliases may not be used after the loop.------ FIXME: the implementation is not finished yet.  Specifically, not--- all of the above conditions are checked.-module Futhark.Optimise.InPlaceLowering-  ( inPlaceLoweringGPU,-    inPlaceLoweringSeq,-    inPlaceLoweringMC,-  )-where--import Control.Monad-import Control.Monad.RWS-import Data.Map.Strict qualified as M-import Data.Ord (comparing)-import Futhark.Analysis.Alias-import Futhark.Builder-import Futhark.IR.Aliases-import Futhark.IR.GPU-import Futhark.IR.MC-import Futhark.IR.Seq (Seq)-import Futhark.Optimise.InPlaceLowering.LowerIntoStm-import Futhark.Pass-import Futhark.Util (nubByOrd)---- | Apply the in-place lowering optimisation to the given program.-inPlaceLoweringGPU :: Pass GPU GPU-inPlaceLoweringGPU = inPlaceLowering onKernelOp lowerUpdateGPU---- | Apply the in-place lowering optimisation to the given program.-inPlaceLoweringSeq :: Pass Seq Seq-inPlaceLoweringSeq = inPlaceLowering pure lowerUpdate---- | Apply the in-place lowering optimisation to the given program.-inPlaceLoweringMC :: Pass MC MC-inPlaceLoweringMC = inPlaceLowering onMCOp lowerUpdate---- | Apply the in-place lowering optimisation to the given program.-inPlaceLowering ::-  (Constraints rep) =>-  OnOp rep ->-  LowerUpdate rep (ForwardingM rep) ->-  Pass rep rep-inPlaceLowering onOp lower =-  Pass "In-place lowering" "Lower in-place updates into loops" $-    fmap removeProgAliases-      . intraproceduralTransformationWithConsts optimiseConsts optimiseFunDef-      . aliasAnalysis-  where-    optimiseConsts stms =-      modifyNameSource $-        runForwardingM lower onOp $-          stmsFromList <$> optimiseStms (stmsToList stms) (pure ())--    optimiseFunDef consts fundec =-      modifyNameSource $-        runForwardingM lower onOp $-          descend (stmsToList consts) $-            bindingFParams (funDefParams fundec) $ do-              body <- optimiseBody $ funDefBody fundec-              pure $ fundec {funDefBody = body}--    descend [] m = m-    descend (stm : stms) m = bindingStm stm $ descend stms m--type Constraints rep = (Buildable rep, AliasableRep rep)--optimiseBody ::-  (Constraints rep) =>-  Body (Aliases rep) ->-  ForwardingM rep (Body (Aliases rep))-optimiseBody (Body als stms res) = do-  stms' <- deepen $ optimiseStms (stmsToList stms) $ mapM_ (seen . resSubExp) res-  pure $ Body als (stmsFromList stms') res-  where-    seen Constant {} = pure ()-    seen (Var v) = seenVar v--optimiseStms ::-  (Constraints rep) =>-  [Stm (Aliases rep)] ->-  ForwardingM rep () ->-  ForwardingM rep [Stm (Aliases rep)]-optimiseStms [] m = m >> pure []-optimiseStms (stm : stms) m = do-  (stms', bup) <- tapBottomUp $ bindingStm stm $ optimiseStms stms m-  stm' <- optimiseInStm stm-  -- XXX: unfortunate that we cannot handle duplicate update values.-  -- Would be good to improve this.  See inplacelowering6.fut.-  case nubByOrd (comparing updateValue)-    . filter ((`notNameIn` bottomUpSeen bup) . updateSource) -- (9)-    . filter ((`elem` boundHere) . updateValue)-    $ forwardThese bup of-    [] -> do-      checkIfForwardableUpdate stm'-      pure $ stm' : stms'-    updates -> do-      lower <- asks topLowerUpdate-      scope <- askScope--      -- If we forward any updates, we need to remove them from stms'.-      let updated_names =-            map updateName updates-          notUpdated =-            not . any (`elem` updated_names) . patNames . stmPat--      -- Condition (5) and (7) are assumed to be checked by-      -- lowerUpdate.-      case lower scope stm' updates of-        Just lowering -> do-          new_stms <- lowering-          new_stms' <- optimiseStms new_stms $ tell bup {forwardThese = []}-          pure $ new_stms' ++ filter notUpdated stms'-        Nothing -> do-          checkIfForwardableUpdate stm'-          pure $ stm' : stms'-  where-    boundHere = patNames $ stmPat stm--    checkIfForwardableUpdate (Let pat (StmAux cs _ _) e)-      | Pat [PatElem v dec] <- pat,-        BasicOp (Update Unsafe src slice (Var ve)) <- e =-          maybeForward ve v dec cs src slice-    checkIfForwardableUpdate stm' =-      mapM_ seenVar $ namesToList $ freeIn $ stmExp stm'--optimiseInStm :: (Constraints rep) => Stm (Aliases rep) -> ForwardingM rep (Stm (Aliases rep))-optimiseInStm (Let pat dec e) =-  Let pat dec <$> optimiseExp e--optimiseExp :: (Constraints rep) => Exp (Aliases rep) -> ForwardingM rep (Exp (Aliases rep))-optimiseExp (Loop merge form body) =-  bindingScope (scopeOfLoopForm form) . bindingFParams (map fst merge) $-    Loop merge form <$> optimiseBody body-optimiseExp (Op op) = do-  f <- asks topOnOp-  Op <$> f op-optimiseExp e = mapExpM optimise e-  where-    optimise =-      identityMapper-        { mapOnBody = const optimiseBody-        }--onSegOp ::-  (Constraints rep) =>-  SegOp lvl (Aliases rep) ->-  ForwardingM rep (SegOp lvl (Aliases rep))-onSegOp op =-  bindingScope (scopeOfSegSpace (segSpace op)) $ do-    let mapper = identitySegOpMapper {mapOnSegOpBody = onKernelBody}-        onKernelBody kbody = do-          stms <--            deepen $-              optimiseStms (stmsToList (kernelBodyStms kbody)) $-                mapM_ seenVar $-                  namesToList $-                    freeIn $-                      kernelBodyResult kbody-          pure kbody {kernelBodyStms = stmsFromList stms}-    mapSegOpM mapper op--onMCOp :: OnOp MC-onMCOp (ParOp par_op op) = ParOp <$> traverse onSegOp par_op <*> onSegOp op-onMCOp op = pure op--onKernelOp :: OnOp GPU-onKernelOp (SegOp op) = SegOp <$> onSegOp op-onKernelOp op = pure op--data Entry rep = Entry-  { entryNumber :: Int,-    entryAliases :: Names,-    entryDepth :: Int,-    entryOptimisable :: Bool,-    entryType :: NameInfo (Aliases rep)-  }--type VTable rep = M.Map VName (Entry rep)--type OnOp rep = Op (Aliases rep) -> ForwardingM rep (Op (Aliases rep))--data TopDown rep = TopDown-  { topDownCounter :: Int,-    topDownTable :: VTable rep,-    topDownDepth :: Int,-    topLowerUpdate :: LowerUpdate rep (ForwardingM rep),-    topOnOp :: OnOp rep-  }--data BottomUp rep = BottomUp-  { bottomUpSeen :: Names,-    forwardThese :: [DesiredUpdate (LetDec (Aliases rep))]-  }--instance Semigroup (BottomUp rep) where-  BottomUp seen1 forward1 <> BottomUp seen2 forward2 =-    BottomUp (seen1 <> seen2) (forward1 <> forward2)--instance Monoid (BottomUp rep) where-  mempty = BottomUp mempty mempty--newtype ForwardingM rep a = ForwardingM (RWS (TopDown rep) (BottomUp rep) VNameSource a)-  deriving-    ( Monad,-      Applicative,-      Functor,-      MonadReader (TopDown rep),-      MonadWriter (BottomUp rep),-      MonadState VNameSource-    )--instance MonadFreshNames (ForwardingM rep) where-  getNameSource = get-  putNameSource = put--instance (Constraints rep) => HasScope (Aliases rep) (ForwardingM rep) where-  askScope = M.map entryType <$> asks topDownTable--runForwardingM ::-  LowerUpdate rep (ForwardingM rep) ->-  OnOp rep ->-  ForwardingM rep a ->-  VNameSource ->-  (a, VNameSource)-runForwardingM f g (ForwardingM m) src =-  let (x, src', _) = runRWS m emptyTopDown src-   in (x, src')-  where-    emptyTopDown =-      TopDown-        { topDownCounter = 0,-          topDownTable = M.empty,-          topDownDepth = 0,-          topLowerUpdate = f,-          topOnOp = g-        }--bindingParams ::-  (dec -> NameInfo (Aliases rep)) ->-  [Param dec] ->-  ForwardingM rep a ->-  ForwardingM rep a-bindingParams f params = local $ \(TopDown n vtable d x y) ->-  let entry fparam =-        ( paramName fparam,-          Entry n mempty d False $ f $ paramDec fparam-        )-      entries = M.fromList $ map entry params-   in TopDown (n + 1) (M.union entries vtable) d x y--bindingFParams ::-  [FParam (Aliases rep)] ->-  ForwardingM rep a ->-  ForwardingM rep a-bindingFParams = bindingParams FParamName--bindingScope ::-  Scope (Aliases rep) ->-  ForwardingM rep a ->-  ForwardingM rep a-bindingScope scope = local $ \(TopDown n vtable d x y) ->-  let entries = M.map entry scope-      infoAliases (LetName (aliases, _)) = unAliases aliases-      infoAliases _ = mempty-      entry info = Entry n (infoAliases info) d False info-   in TopDown (n + 1) (entries <> vtable) d x y--bindingStm ::-  Stm (Aliases rep) ->-  ForwardingM rep a ->-  ForwardingM rep a-bindingStm (Let pat _ _) = local $ \(TopDown n vtable d x y) ->-  let entries = M.fromList $ map entry $ patElems pat-      entry patElem =-        let (aliases, _) = patElemDec patElem-         in ( patElemName patElem,-              Entry n (unAliases aliases) d True $ LetName $ patElemDec patElem-            )-   in TopDown (n + 1) (M.union entries vtable) d x y--bindingNumber :: VName -> ForwardingM rep Int-bindingNumber name = do-  res <- asks $ fmap entryNumber . M.lookup name . topDownTable-  case res of-    Just n -> pure n-    Nothing ->-      error $-        "bindingNumber: variable "-          ++ prettyString name-          ++ " not found."--deepen :: ForwardingM rep a -> ForwardingM rep a-deepen = local $ \env -> env {topDownDepth = topDownDepth env + 1}--areAvailableBefore :: Names -> VName -> ForwardingM rep Bool-areAvailableBefore names point = do-  pointN <- bindingNumber point-  nameNs <- mapM bindingNumber $ namesToList names-  pure $ all (< pointN) nameNs--isInCurrentBody :: VName -> ForwardingM rep Bool-isInCurrentBody name = do-  current <- asks topDownDepth-  res <- asks $ fmap entryDepth . M.lookup name . topDownTable-  case res of-    Just d -> pure $ d == current-    Nothing ->-      error $-        "isInCurrentBody: variable "-          ++ prettyString name-          ++ " not found."--isOptimisable :: VName -> ForwardingM rep Bool-isOptimisable name = do-  res <- asks $ fmap entryOptimisable . M.lookup name . topDownTable-  case res of-    Just b -> pure b-    Nothing ->-      error $-        "isOptimisable: variable "-          ++ prettyString name-          ++ " not found."--seenVar :: forall rep. VName -> ForwardingM rep ()-seenVar name = do-  aliases <--    asks $-      maybe mempty entryAliases-        . M.lookup name-        . topDownTable-  tell $ (mempty :: BottomUp rep) {bottomUpSeen = oneName name <> aliases}--tapBottomUp :: ForwardingM rep a -> ForwardingM rep (a, BottomUp rep)-tapBottomUp m = do-  (x, bup) <- listen m-  pure (x, bup)--maybeForward ::-  (Constraints rep) =>-  VName ->-  VName ->-  LetDec (Aliases rep) ->-  Certs ->-  VName ->-  Slice SubExp ->-  ForwardingM rep ()-maybeForward v dest_nm dest_dec cs src slice = do-  -- Checks condition (2)-  available <--    (freeIn src <> freeIn slice <> freeIn cs)-      `areAvailableBefore` v-  -- Check condition (3)-  samebody <- isInCurrentBody v-  -- Check condition (6)-  optimisable <- isOptimisable v-  not_prim <- not . primType <$> lookupType v-  when (available && samebody && optimisable && not_prim) $ do-    let fwd = DesiredUpdate dest_nm dest_dec cs src slice v-    tell mempty {forwardThese = [fwd]}

@@ -1,376 +0,0 @@-{-# LANGUAGE TypeFamilies #-}--module Futhark.Optimise.InPlaceLowering.LowerIntoStm-  ( lowerUpdateGPU,-    lowerUpdate,-    LowerUpdate,-    DesiredUpdate (..),-  )-where--import Control.Monad-import Control.Monad.Writer-import Data.Either-import Data.List (find, unzip5)-import Data.Maybe (isNothing, mapMaybe)-import Futhark.Analysis.PrimExp.Convert-import Futhark.Construct-import Futhark.IR.Aliases-import Futhark.IR.GPU-import Futhark.Optimise.InPlaceLowering.SubstituteIndices--data DesiredUpdate dec = DesiredUpdate-  { -- | Name of result.-    updateName :: VName,-    -- | Type of result.-    updateType :: dec,-    updateCerts :: Certs,-    updateSource :: VName,-    updateIndices :: Slice SubExp,-    updateValue :: VName-  }-  deriving (Show)--instance Functor DesiredUpdate where-  f `fmap` u = u {updateType = f $ updateType u}--updateHasValue :: VName -> DesiredUpdate dec -> Bool-updateHasValue name = (name ==) . updateValue--type LowerUpdate rep m =-  Scope (Aliases rep) ->-  Stm (Aliases rep) ->-  [DesiredUpdate (LetDec (Aliases rep))] ->-  Maybe (m [Stm (Aliases rep)])--lowerUpdate ::-  ( MonadFreshNames m,-    Buildable rep,-    LetDec rep ~ Type,-    AliasableRep rep-  ) =>-  LowerUpdate rep m-lowerUpdate scope (Let pat aux (Loop merge form body)) updates = do-  canDo <- lowerUpdateIntoLoop scope updates pat merge form body-  Just $ do-    (prestms, poststms, pat', merge', body') <- canDo-    pure $-      prestms-        ++ [ certify (stmAuxCerts aux) $-               mkLet pat' $-                 Loop merge' form body'-           ]-        ++ poststms-lowerUpdate-  _-  (Let pat aux (BasicOp (SubExp (Var v))))-  [DesiredUpdate bindee_nm bindee_dec cs src (Slice is) val]-    | patNames pat == [src] =-        let is' = fullSlice (typeOf bindee_dec) is-         in Just . pure $-              [ certify (stmAuxCerts aux <> cs) $-                  mkLet [Ident bindee_nm $ typeOf bindee_dec] $-                    BasicOp $-                      Update Unsafe v is' $-                        Var val-              ]-lowerUpdate _ _ _ =-  Nothing--lowerUpdateGPU :: (MonadFreshNames m) => LowerUpdate GPU m-lowerUpdateGPU-  scope-  (Let pat aux (Op (SegOp (SegMap lvl space ts kbody))))-  updates-    | all ((`elem` patNames pat) . updateValue) updates,-      not source_used_in_kbody = do-        mk <- lowerUpdatesIntoSegMap scope pat ts updates space kbody-        Just $ do-          (pat', ts', kbody', poststms) <- mk-          let cs = stmAuxCerts aux <> foldMap updateCerts updates-          pure $-            certify cs (Let pat' aux $ Op $ SegOp $ SegMap lvl space ts' kbody')-              : stmsToList poststms-    where-      -- This check is a bit more conservative than ideal.  In a perfect-      -- world, we would allow indexing a[i,j] if the update is also-      -- to exactly a[i,j], as that will not create cross-iteration-      -- dependencies.  (Although the type checker wouldn't be able to-      -- permit this anyway.)-      source_used_in_kbody =-        mconcat (map (`lookupAliases` scope) (namesToList (freeIn kbody)))-          `namesIntersect` mconcat (map ((`lookupAliases` scope) . updateSource) updates)-lowerUpdateGPU scope stm updates = lowerUpdate scope stm updates--lowerUpdatesIntoSegMap ::-  (MonadFreshNames m) =>-  Scope (Aliases GPU) ->-  Pat (LetDec (Aliases GPU)) ->-  [Type] ->-  [DesiredUpdate (LetDec (Aliases GPU))] ->-  SegSpace ->-  KernelBody (Aliases GPU) ->-  Maybe-    ( m-        ( Pat (LetDec (Aliases GPU)),-          [Type],-          KernelBody (Aliases GPU),-          Stms (Aliases GPU)-        )-    )-lowerUpdatesIntoSegMap scope pat ret_ts updates kspace kbody = do-  -- The updates are all-or-nothing.  Being more liberal would require-  -- changes to the in-place-lowering pass itself.-  mk <- mapM onRet (zip3 (patElems pat) ret_ts (kernelBodyResult kbody))-  pure $ do-    (pes, ret_ts', bodystms, krets, poststms) <- unzip5 <$> sequence mk-    pure-      ( Pat pes,-        ret_ts',-        kbody-          { kernelBodyStms = kernelBodyStms kbody <> mconcat bodystms,-            kernelBodyResult = krets-          },-        mconcat poststms-      )-  where-    (gtids, _dims) = unzip $ unSegSpace kspace--    onRet (PatElem v v_dec, _, ret)-      | Just (DesiredUpdate bindee_nm bindee_dec _cs src slice _val) <--          find ((== v) . updateValue) updates = do-          Returns _ cs se <- Just ret--          -- The slice we're writing per thread must fully cover the-          -- underlying dimensions.-          guard $-            let (dims', slice') =-                  unzip . drop (length gtids) . filter (isNothing . dimFix . snd) $-                    zip (arrayDims (typeOf bindee_dec)) (unSlice slice)-             in isFullSlice (Shape dims') (Slice slice')--          Just $ do-            (slice', bodystms) <--              flip runBuilderT scope $-                traverse (toSubExp "index") $-                  fixSlice (fmap pe64 slice) $-                    map (pe64 . Var) gtids--            let ret' = WriteReturns cs src [(fullSlice (typeOf bindee_dec) (map DimFix slice'), se)]--            v_aliased <- newName v--            pure-              ( PatElem bindee_nm bindee_dec,-                typeOf bindee_dec,-                bodystms,-                ret',-                stmsFromList-                  [ mkLet [Ident v_aliased $ typeOf v_dec] $ BasicOp $ Index bindee_nm slice,-                    mkLet [Ident v $ typeOf v_dec] $ BasicOp $ Replicate mempty $ Var v_aliased-                  ]-              )-    onRet (pe, ret_t, ret) =-      Just $ pure (pe, ret_t, mempty, ret, mempty)--lowerUpdateIntoLoop ::-  ( Buildable rep,-    BuilderOps rep,-    Aliased rep,-    LetDec rep ~ (als, Type),-    MonadFreshNames m-  ) =>-  Scope rep ->-  [DesiredUpdate (LetDec rep)] ->-  Pat (LetDec rep) ->-  [(FParam rep, SubExp)] ->-  LoopForm ->-  Body rep ->-  Maybe-    ( m-        ( [Stm rep],-          [Stm rep],-          [Ident],-          [(FParam rep, SubExp)],-          Body rep-        )-    )-lowerUpdateIntoLoop scope updates pat val form body = do-  -- Algorithm:-  ---  --   0) Map each result of the loop body to a corresponding in-place-  --      update, if one exists.-  ---  --   1) Create new merge variables corresponding to the arrays being-  --      updated; extend the pattern and the @res@ list with these,-  --      and remove the parts of the result list that have a-  --      corresponding in-place update.-  ---  --      (The creation of the new merge variable identifiers is-  --      actually done at the same time as step (0)).-  ---  --   2) Create in-place updates at the end of the loop body.-  ---  --   3) Create index expressions that read back the values written-  --      in (2).  If the merge parameter corresponding to this value-  --      is unique, also @copy@ this value.-  ---  --   4) Update the result of the loop body to properly pass the new-  --      arrays and indexed elements to the next iteration of the-  --      loop.-  ---  -- We also check that the merge parameters we work with have-  -- loop-invariant shapes.--  -- Safety condition (8).-  forM_ (zip val $ bodyAliases body) $ \((p, _), als) ->-    guard $ paramName p `notNameIn` als--  mk_in_place_map <- summariseLoop scope updates usedInBody resmap val--  Just $ do-    in_place_map <- mk_in_place_map-    (val', prestms, poststms) <- mkMerges in_place_map-    let valpat = mkResAndPat in_place_map-        idxsubsts = indexSubstitutions in_place_map-    (idxsubsts', newstms) <- substituteIndices idxsubsts $ bodyStms body-    (body_res, res_stms) <- manipulateResult in_place_map idxsubsts'-    let body' = mkBody (newstms <> res_stms) body_res-    pure (prestms, poststms, valpat, val', body')-  where-    usedInBody =-      mconcat $ map (`lookupAliases` scope) $ namesToList $ freeIn body <> freeIn form-    resmap = zip (bodyResult body) $ patIdents pat--    mkMerges ::-      (MonadFreshNames m, Buildable rep) =>-      [LoopResultSummary (als, Type)] ->-      m ([(Param DeclType, SubExp)], [Stm rep], [Stm rep])-    mkMerges summaries = do-      ((origmerge, extramerge), (prestms, poststms)) <--        runWriterT $ partitionEithers <$> mapM mkMerge summaries-      pure (origmerge ++ extramerge, prestms, poststms)--    mkMerge summary-      | Just (update, mergename, mergedec) <- relatedUpdate summary = do-          source <- newVName "modified_source"-          precopy <- newVName $ baseString (updateValue update) <> "_precopy"-          let source_t = snd $ updateType update-              elm_t = source_t `setArrayDims` sliceDims (updateIndices update)-          tell-            ( [ mkLet [Ident source source_t] . BasicOp-                  $ Update-                    Unsafe-                    (updateSource update)-                    (fullSlice source_t $ unSlice $ updateIndices update)-                  $ snd-                  $ mergeParam summary-              ],-              [ mkLet [Ident precopy elm_t] . BasicOp $-                  Index-                    (updateName update)-                    (fullSlice source_t $ unSlice $ updateIndices update),-                mkLet [Ident (updateValue update) elm_t] $ BasicOp $ Replicate mempty $ Var precopy-              ]-            )-          pure $-            Right-              ( Param mempty mergename (toDecl (typeOf mergedec) Unique),-                Var source-              )-      | otherwise = pure $ Left $ mergeParam summary--    mkResAndPat summaries =-      let (origpat, extrapat) = partitionEithers $ map mkResAndPat' summaries-       in origpat ++ extrapat--    mkResAndPat' summary-      | Just (update, _, _) <- relatedUpdate summary =-          Right (Ident (updateName update) (snd $ updateType update))-      | otherwise =-          Left (inPatAs summary)--summariseLoop ::-  ( Aliased rep,-    MonadFreshNames m-  ) =>-  Scope rep ->-  [DesiredUpdate (als, Type)] ->-  Names ->-  [(SubExpRes, Ident)] ->-  [(Param DeclType, SubExp)] ->-  Maybe (m [LoopResultSummary (als, Type)])-summariseLoop scope updates usedInBody resmap merge =-  sequence <$> zipWithM summariseLoopResult resmap merge-  where-    summariseLoopResult (se, v) (fparam, mergeinit)-      | Just update <- find (updateHasValue $ identName v) updates =-          -- Safety condition (7)-          if usedInBody `namesIntersect` lookupAliases (updateSource update) scope-            then Nothing-            else-              if hasLoopInvariantShape fparam-                then Just $ do-                  lowered_array <- newVName "lowered_array"-                  pure-                    LoopResultSummary-                      { resultSubExp = se,-                        inPatAs = v,-                        mergeParam = (fparam, mergeinit),-                        relatedUpdate =-                          Just-                            ( update,-                              lowered_array,-                              updateType update-                            )-                      }-                else Nothing-    summariseLoopResult _ _ =-      Nothing -- XXX: conservative; but this entire pass is going away.-    hasLoopInvariantShape = all loopInvariant . arrayDims . paramType--    merge_param_names = map (paramName . fst) merge--    loopInvariant (Var v) = v `notElem` merge_param_names-    loopInvariant Constant {} = True--data LoopResultSummary dec = LoopResultSummary-  { resultSubExp :: SubExpRes,-    inPatAs :: Ident,-    mergeParam :: (Param DeclType, SubExp),-    relatedUpdate :: Maybe (DesiredUpdate dec, VName, dec)-  }-  deriving (Show)--indexSubstitutions :: (Typed dec) => [LoopResultSummary dec] -> IndexSubstitutions-indexSubstitutions = mapMaybe getSubstitution-  where-    getSubstitution res = do-      (DesiredUpdate _ _ cs _ is _, nm, dec) <- relatedUpdate res-      let name = paramName $ fst $ mergeParam res-      pure (name, (cs, nm, typeOf dec, is))--manipulateResult ::-  (Buildable rep, MonadFreshNames m) =>-  [LoopResultSummary (LetDec rep)] ->-  IndexSubstitutions ->-  m (Result, Stms rep)-manipulateResult summaries substs = do-  let (orig_ses, updated_ses) = partitionEithers $ map unchangedRes summaries-  (subst_ses, res_stms) <- runWriterT $ zipWithM substRes updated_ses substs-  pure (orig_ses ++ subst_ses, stmsFromList res_stms)-  where-    unchangedRes summary =-      case relatedUpdate summary of-        Nothing -> Left $ resultSubExp summary-        Just _ -> Right $ resultSubExp summary-    substRes (SubExpRes res_cs (Var res_v)) (subst_v, (_, nm, _, _))-      | res_v == subst_v =-          pure $ SubExpRes res_cs $ Var nm-    substRes (SubExpRes res_cs res_se) (_, (cs, nm, dec, Slice is)) = do-      v' <- newIdent' (++ "_updated") $ Ident nm $ typeOf dec-      tell-        [ certify (res_cs <> cs) . mkLet [v'] . BasicOp $-            Update Unsafe nm (fullSlice (typeOf dec) is) res_se-        ]-      pure $ varRes $ identName v'

@@ -1,191 +0,0 @@-{-# LANGUAGE TypeFamilies #-}---- | This module exports facilities for transforming array accesses in--- a list of 'Stm's (intended to be the bindings in a body).  The--- idea is that you can state that some variable @x@ is in fact an--- array indexing @v[i0,i1,...]@.-module Futhark.Optimise.InPlaceLowering.SubstituteIndices-  ( substituteIndices,-    IndexSubstitution,-    IndexSubstitutions,-  )-where--import Control.Monad-import Data.Map.Strict qualified as M-import Futhark.Construct-import Futhark.IR-import Futhark.IR.Prop.Aliases-import Futhark.Transform.Substitute---- | Essentially the components of an 'Index' expression.-type IndexSubstitution = (Certs, VName, Type, Slice SubExp)---- | A mapping from variable names to the indexing operation they--- should be replaced with.-type IndexSubstitutions = [(VName, IndexSubstitution)]--typeEnvFromSubstitutions :: (LParamInfo rep ~ Type) => IndexSubstitutions -> Scope rep-typeEnvFromSubstitutions = M.fromList . map (fromSubstitution . snd)-  where-    fromSubstitution (_, name, t, _) =-      (name, LParamName t)---- | Perform the substitution.-substituteIndices ::-  ( MonadFreshNames m,-    BuilderOps rep,-    Buildable rep,-    Aliased rep-  ) =>-  IndexSubstitutions ->-  Stms rep ->-  m (IndexSubstitutions, Stms rep)-substituteIndices substs stms =-  runBuilderT (substituteIndicesInStms substs stms) types-  where-    types = typeEnvFromSubstitutions substs--substituteIndicesInStms ::-  (MonadBuilder m, Buildable (Rep m), Aliased (Rep m)) =>-  IndexSubstitutions ->-  Stms (Rep m) ->-  m IndexSubstitutions-substituteIndicesInStms = foldM substituteIndicesInStm--substituteIndicesInStm ::-  (MonadBuilder m, Buildable (Rep m), Aliased (Rep m)) =>-  IndexSubstitutions ->-  Stm (Rep m) ->-  m IndexSubstitutions-substituteIndicesInStm substs (Let pat _ (BasicOp (Rearrange perm v)))-  | Just (cs, src, src_t, is) <- lookup v substs,-    [v'] <- patNames pat = do-      let extra_dims = arrayRank src_t - length perm-          perm' = [0 .. extra_dims - 1] ++ map (+ extra_dims) perm-      src' <--        letExp (baseString v' <> "_subst") $ BasicOp $ Rearrange perm' src-      src_t' <- lookupType src'-      pure $ (v', (cs, src', src_t', is)) : substs-substituteIndicesInStm substs (Let pat rep e) = do-  e' <- substituteIndicesInExp substs e-  addStm $ Let pat rep e'-  pure substs--substituteIndicesInExp ::-  (MonadBuilder m, Buildable (Rep m), Aliased (Rep m)) =>-  IndexSubstitutions ->-  Exp (Rep m) ->-  m (Exp (Rep m))-substituteIndicesInExp substs (Op op) = do-  let used_in_op = filter ((`nameIn` freeIn op) . fst) substs-  var_substs <- fmap mconcat $-    forM used_in_op $ \(v, (cs, src, src_dec, Slice is)) -> do-      v' <--        certifying cs $-          letExp (baseString src <> "_op_idx") $-            BasicOp $-              Index src $-                fullSlice (typeOf src_dec) is-      pure $ M.singleton v v'-  pure $ Op $ substituteNames var_substs op-substituteIndicesInExp substs e = do-  substs' <- copyAnyConsumed e-  let substitute =-        identityMapper-          { mapOnSubExp = substituteIndicesInSubExp substs',-            mapOnVName = substituteIndicesInVar substs',-            mapOnBody = const $ substituteIndicesInBody substs'-          }--  mapExpM substitute e-  where-    copyAnyConsumed =-      let consumingSubst substs' v-            | Just (cs2, src2, src2dec, is2) <- lookup v substs = do-                row <--                  certifying cs2 $-                    letSubExp (baseString v ++ "_row") $-                      BasicOp $-                        Index src2 $-                          fullSlice (typeOf src2dec) (unSlice is2)-                row_copy <--                  letExp (baseString v ++ "_row_copy") . BasicOp $-                    Replicate mempty row-                pure $-                  update-                    v-                    v-                    ( mempty,-                      row_copy,-                      src2dec-                        `setType` ( typeOf src2dec-                                      `setArrayDims` sliceDims is2-                                  ),-                      Slice []-                    )-                    substs'-          consumingSubst substs' _ =-            pure substs'-       in foldM consumingSubst substs . namesToList . consumedInExp--substituteIndicesInSubExp ::-  (MonadBuilder m) =>-  IndexSubstitutions ->-  SubExp ->-  m SubExp-substituteIndicesInSubExp substs (Var v) =-  Var <$> substituteIndicesInVar substs v-substituteIndicesInSubExp _ se =-  pure se--substituteIndicesInVar ::-  (MonadBuilder m) =>-  IndexSubstitutions ->-  VName ->-  m VName-substituteIndicesInVar substs v-  | Just (cs2, src2, _, Slice []) <- lookup v substs =-      certifying cs2 $-        letExp (baseString src2) $-          BasicOp $-            SubExp $-              Var src2-  | Just (cs2, src2, src2_dec, Slice is2) <- lookup v substs =-      certifying cs2 $-        letExp (baseString src2 <> "_v_idx") $-          BasicOp $-            Index src2 $-              fullSlice (typeOf src2_dec) is2-  | otherwise =-      pure v--substituteIndicesInBody ::-  (MonadBuilder m, Buildable (Rep m), Aliased (Rep m)) =>-  IndexSubstitutions ->-  Body (Rep m) ->-  m (Body (Rep m))-substituteIndicesInBody substs (Body _ stms res) = do-  (substs', stms') <--    inScopeOf stms $-      collectStms $-        substituteIndicesInStms substs stms-  (res', res_stms) <--    inScopeOf stms' $-      collectStms $-        mapM (onSubExpRes substs') res-  mkBodyM (stms' <> res_stms) res'-  where-    onSubExpRes substs' (SubExpRes cs se) =-      SubExpRes cs <$> substituteIndicesInSubExp substs' se--update ::-  VName ->-  VName ->-  IndexSubstitution ->-  IndexSubstitutions ->-  IndexSubstitutions-update needle name subst ((othername, othersubst) : substs)-  | needle == othername = (name, subst) : substs-  | otherwise = (othername, othersubst) : update needle name subst substs-update needle _ _ [] = error $ "Cannot find substitution for " ++ prettyString needle

@@ -25,7 +25,6 @@ import Futhark.Optimise.Fusion import Futhark.Optimise.GenRedOpt import Futhark.Optimise.HistAccs-import Futhark.Optimise.InPlaceLowering import Futhark.Optimise.InliningDeadFun import Futhark.Optimise.MemoryBlockMerging qualified as MemoryBlockMerging import Futhark.Optimise.MergeGPUBodies@@ -103,7 +102,6 @@         mergeGPUBodies,         simplifyGPU, -- Cleanup merged GPUBody kernels.         sinkGPU, -- Sink reads within GPUBody kernels.-        inPlaceLoweringGPU,         babysitKernels,         -- Important to simplify after babysitting in order to fix up         -- redundant manifests.@@ -118,8 +116,7 @@   standardPipeline     >>> onePass firstOrderTransform     >>> passes-      [ simplifySeq,-        inPlaceLoweringSeq+      [ simplifySeq       ]  -- | Run 'seqPipeline', then add memory information (and@@ -183,8 +180,7 @@         unstreamMC,         performCSE True,         simplifyMC,-        sinkMC,-        inPlaceLoweringMC+        sinkMC       ]  -- | Run 'mcPipeline' and then add memory information.

@@ -39,6 +39,9 @@   divUp x y =     (x + y - 1) `Futhark.Util.IntegralExp.div` y +  nextMul :: e -> e -> e+  nextMul x y = x `divUp` y * y+ -- | This wrapper allows you to use a type that is an instance of the -- true class whenever the simile class is required. newtype Wrapped a = Wrapped {wrappedValue :: a}

@@ -90,14 +90,14 @@  type Stack = [StackFrame] -type Sizes = M.Map VName Int64+type Exts = M.Map VName Value  -- | The monad in which evaluation takes place. newtype EvalM a   = EvalM       ( ReaderT           (Stack, M.Map ImportName Env)-          (StateT Sizes (F ExtOp))+          (StateT Exts (F ExtOp))           a       )   deriving@@ -106,7 +106,7 @@       Functor,       MonadFree ExtOp,       MonadReader (Stack, M.Map ImportName Env),-      MonadState Sizes+      MonadState Exts     )  runEvalM :: M.Map ImportName Env -> EvalM a -> F ExtOp a@@ -129,11 +129,11 @@ lookupImport :: ImportName -> EvalM (Maybe Env) lookupImport f = asks $ M.lookup f . snd -putExtSize :: VName -> Int64 -> EvalM ()+putExtSize :: VName -> Value -> EvalM () putExtSize v x = modify $ M.insert v x -getSizes :: EvalM Sizes-getSizes = get+getExts :: EvalM Exts+getExts = get  -- | Disregard any existential sizes computed during this action. -- This is used so that existentials computed during one iteration of@@ -145,8 +145,13 @@   put s   pure x -extSizeEnv :: EvalM Env-extSizeEnv = i64Env <$> getSizes+extEnv :: EvalM Env+extEnv = valEnv . M.map f <$> getExts+  where+    f v =+      ( Nothing,+        v+      )  valueStructType :: ValueType -> StructType valueStructType = first $ flip sizeFromInteger mempty . toInteger@@ -296,7 +301,8 @@ -- | An expression evaluator that embeds an environment. type Eval = Exp -> EvalM Value --- | A TermValue with a 'Nothing' type annotation is an intrinsic.+-- | A TermValue with a 'Nothing' type annotation is an intrinsic or+-- an existential. data TermBinding   = TermValue (Maybe T.BoundV) Value   | -- | A polymorphic value that must be instantiated.  The@@ -636,7 +642,7 @@  evalWithExts :: Env -> EvalM Eval evalWithExts env = do-  size_env <- extSizeEnv+  size_env <- extEnv   pure $ eval $ size_env <> env  -- | Evaluate all possible sizes, except those that contain free@@ -750,16 +756,17 @@ evalArg env e ext = do   v <- eval env e   case ext of-    Just ext' -> putExtSize ext' $ asInt64 v-    Nothing -> pure ()+    Just ext' -> putExtSize ext' v+    _ -> pure ()   pure v  returned :: Env -> TypeBase Size als -> [VName] -> Value -> EvalM Value returned _ _ [] v = pure v returned env ret retext v = do-  mapM_ (uncurry putExtSize) . M.toList $-    resolveExistentials retext (expandType env $ toStruct ret) $-      valueShape v+  mapM_ (uncurry putExtSize . second (ValuePrim . SignedValue . Int64Value))+    . M.toList+    $ resolveExistentials retext (expandType env $ toStruct ret)+    $ valueShape v   pure v  evalAppExp :: Env -> AppExp -> EvalM Value@@ -1193,7 +1200,7 @@ evalDec :: Env -> Dec -> EvalM Env evalDec env (ValDec (ValBind _ v _ (Info ret) tparams ps fbody _ _ _)) = localExts $ do   binding <- evalFunctionBinding env tparams ps ret fbody-  sizes <- extSizeEnv+  sizes <- extEnv   pure $     env {envTerm = M.insert v binding $ envTerm env} <> sizes evalDec env (OpenDec me _) = do@@ -1204,7 +1211,7 @@ evalDec env (ImportDec name name' loc) =   evalDec env $ LocalDec (OpenDec (ModImport name name' loc) loc) loc evalDec env (LocalDec d _) = evalDec env d-evalDec env SigDec {} = pure env+evalDec env ModTypeDec {} = pure env evalDec env (TypeDec (TypeBind v l ps _ (Info (RetType dims t)) _ _)) = do   let abbr = T.TypeAbbr l ps . RetType dims $ expandType env t   pure env {envType = M.insert v abbr $ envType env}@@ -1941,7 +1948,7 @@     -- We need to extract any new existential sizes and add them as     -- ordinary bindings to the context, or we will not be able to     -- look up their values later.-    sizes <- extSizeEnv+    sizes <- extEnv     pure $ env <> sizes  interpretDec :: Ctx -> Dec -> F ExtOp Ctx

@@ -57,7 +57,7 @@ addDoc :: DocComment -> UncheckedDec -> UncheckedDec addDoc doc (ValDec val) = ValDec (val {valBindDoc = Just doc}) addDoc doc (TypeDec tp) = TypeDec (tp {typeDoc = Just doc})-addDoc doc (SigDec sig) = SigDec (sig {sigDoc = Just doc})+addDoc doc (ModTypeDec sig) = ModTypeDec (sig {modTypeDoc = Just doc}) addDoc doc (ModDec mod) = ModDec (mod {modDoc = Just doc}) addDoc _ dec = dec 

@@ -217,7 +217,7 @@ Dec_ :: { UncheckedDec }     : Val               { ValDec $1 }     | TypeAbbr          { TypeDec $1 }-    | SigBind           { SigDec $1 }+    | ModTypeBind       { ModTypeDec $1 }     | ModBind           { ModDec $1 }     | open ModExp       { OpenDec $2 (srclocOf $1) }     | import stringlit@@ -228,29 +228,29 @@  ; -SigExp :: { UncheckedSigExp }-        : QualName            { let (v, loc) = $1 in SigVar v NoInfo (srclocOf loc) }-        | '{' Specs '}'  { SigSpecs $2 (srcspan $1 $>) }-        | SigExp with TypeRef { SigWith $1 $3 (srcspan $1 $>) }-        | '(' SigExp ')'      { SigParens $2 (srcspan $1 $>) }-        | '(' id ':' SigExp ')' '->' SigExp-                              { let L _ (ID name) = $2-                                in SigArrow (Just name) $4 $7 (srcspan $1 $>) }-        | SigExp '->' SigExp  { SigArrow Nothing $1 $3 (srcspan $1 $>) }+ModTypeExp :: { UncheckedModTypeExp }+        : QualName                { let (v, loc) = $1 in ModTypeVar v NoInfo (srclocOf loc) }+        | '{' Specs '}'           { ModTypeSpecs $2 (srcspan $1 $>) }+        | ModTypeExp with TypeRef { ModTypeWith $1 $3 (srcspan $1 $>) }+        | '(' ModTypeExp ')'      { ModTypeParens $2 (srcspan $1 $>) }+        | '(' id ':' ModTypeExp ')' '->' ModTypeExp+                                  { let L _ (ID name) = $2+                                    in ModTypeArrow (Just name) $4 $7 (srcspan $1 $>) }+        | ModTypeExp '->' ModTypeExp  { ModTypeArrow Nothing $1 $3 (srcspan $1 $>) }  TypeRef :: { TypeRefBase NoInfo Name }          : QualName TypeParams '=' TypeExpTerm            { TypeRef (fst $1) $2 $4 (srcspan (snd $1) $>) } -SigBind :: { SigBindBase NoInfo Name }-         : module type id '=' SigExp+ModTypeBind :: { ModTypeBindBase NoInfo Name }+         : module type id '=' ModTypeExp           { let L _ (ID name) = $3-            in SigBind name $5 Nothing (srcspan $1 $>) }+            in ModTypeBind name $5 Nothing (srcspan $1 $>) }  ModExp :: { UncheckedModExp }-        : ModExp ':' SigExp+        : ModExp ':' ModTypeExp           { ModAscript $1 $3 NoInfo (srcspan $1 $>) }-        | '\\' ModParam maybeAscription(SimpleSigExp) '->' ModExp+        | '\\' ModParam maybeAscription(SimpleModTypeExp) '->' ModExp           { ModLambda $2 (fmap (,NoInfo) $3) $5 (srcspan $1 $>) }         | import stringlit           { let L _ (STRINGLIT s) = $2 in ModImport (T.unpack s) NoInfo (srcspan $1 $>) }@@ -273,18 +273,18 @@               { let (v, loc) = $1 in ModVar v (srclocOf loc) }             | '{' Decs '}' { ModDecs $2 (srcspan $1 $>) } -SimpleSigExp :: { UncheckedSigExp }-             : QualName            { let (v, loc) = $1 in SigVar v NoInfo (srclocOf loc) }-             | '(' SigExp ')'      { $2 }+SimpleModTypeExp :: { UncheckedModTypeExp }+             : QualName            { let (v, loc) = $1 in ModTypeVar v NoInfo (srclocOf loc) }+             | '(' ModTypeExp ')'      { $2 }  ModBind :: { ModBindBase NoInfo Name }-         : module id ModParams maybeAscription(SigExp) '=' ModExp+         : module id ModParams maybeAscription(ModTypeExp) '=' ModExp            { let L floc (ID fname) = $2;              in ModBind fname $3 (fmap (,NoInfo) $4) $6 Nothing (srcspan $1 $>)            }  ModParam :: { ModParamBase NoInfo Name }-          : '(' id ':' SigExp ')' { let L _ (ID name) = $2 in ModParam name $4 NoInfo (srcspan $1 $>) }+          : '(' id ':' ModTypeExp ')' { let L _ (ID name) = $2 in ModParam name $4 NoInfo (srcspan $1 $>) }  ModParams :: { [ModParamBase NoInfo Name] }            : ModParam ModParams { $1 : $2 }@@ -308,10 +308,10 @@         { let L _ (ID name) = $3           in TypeSpec $2 name $4 Nothing (srcspan $1 $>) } -      | module id ':' SigExp+      | module id ':' ModTypeExp         { let L _ (ID name) = $2           in ModSpec name $4 Nothing (srcspan $1 $>) }-      | include SigExp+      | include ModTypeExp         { IncludeSpec $2 (srcspan $1 $>) }       | Doc Spec         { addDocSpec $1 $2 }

@@ -460,7 +460,7 @@ instance (Eq vn, IsName vn, Annot f) => Pretty (DecBase f vn) where   pretty (ValDec dec) = pretty dec   pretty (TypeDec dec) = pretty dec-  pretty (SigDec sig) = pretty sig+  pretty (ModTypeDec sig) = pretty sig   pretty (ModDec sd) = pretty sd   pretty (OpenDec x _) = "open" <+> pretty x   pretty (LocalDec dec _) = "local" <+> pretty dec@@ -544,19 +544,19 @@   pretty (IncludeSpec e _) =     "include" <+> pretty e -instance (Eq vn, IsName vn, Annot f) => Pretty (SigExpBase f vn) where-  pretty (SigVar v _ _) = pretty v-  pretty (SigParens e _) = parens $ pretty e-  pretty (SigSpecs ss _) = nestedBlock "{" "}" (stack $ punctuate line $ map pretty ss)-  pretty (SigWith s (TypeRef v ps td _) _) =+instance (Eq vn, IsName vn, Annot f) => Pretty (ModTypeExpBase f vn) where+  pretty (ModTypeVar v _ _) = pretty v+  pretty (ModTypeParens e _) = parens $ pretty e+  pretty (ModTypeSpecs ss _) = nestedBlock "{" "}" (stack $ punctuate line $ map pretty ss)+  pretty (ModTypeWith s (TypeRef v ps td _) _) =     pretty s <+> "with" <+> pretty v <+> hsep (map pretty ps) <> " =" <+> pretty td-  pretty (SigArrow (Just v) e1 e2 _) =+  pretty (ModTypeArrow (Just v) e1 e2 _) =     parens (prettyName v <> colon <+> pretty e1) <+> "->" <+> pretty e2-  pretty (SigArrow Nothing e1 e2 _) =+  pretty (ModTypeArrow Nothing e1 e2 _) =     pretty e1 <+> "->" <+> pretty e2 -instance (Eq vn, IsName vn, Annot f) => Pretty (SigBindBase f vn) where-  pretty (SigBind name e _ _) =+instance (Eq vn, IsName vn, Annot f) => Pretty (ModTypeBindBase f vn) where+  pretty (ModTypeBind name e _ _) =     "module type" <+> prettyName name <+> equals <+> pretty e  instance (Eq vn, IsName vn, Annot f) => Pretty (ModParamBase f vn) where

@@ -1285,7 +1285,7 @@           FloatType Float32,           \case             [FloatValue (Float32Value x), IntValue (Int32Value y)] ->-              Just $ FloatValue $ Float32Value $ x * (2 ** fromIntegral y)+              Just $ FloatValue $ Float32Value $ ldexpf x $ fromIntegral y             _ -> Nothing         )       ),@@ -1294,7 +1294,7 @@           FloatType Float64,           \case             [FloatValue (Float64Value x), IntValue (Int32Value y)] ->-              Just $ FloatValue $ Float64Value $ x * (2 ** fromIntegral y)+              Just $ FloatValue $ Float64Value $ ldexp x $ fromIntegral y             _ -> Nothing         )       ),

@@ -89,12 +89,12 @@     UncheckedSlice,     UncheckedExp,     UncheckedModExp,-    UncheckedSigExp,+    UncheckedModTypeExp,     UncheckedTypeParam,     UncheckedPat,     UncheckedValBind,     UncheckedTypeBind,-    UncheckedSigBind,+    UncheckedModTypeBind,     UncheckedModBind,     UncheckedDec,     UncheckedSpec,@@ -110,9 +110,9 @@     Pat,     ModExp,     ModParam,-    SigExp,+    ModTypeExp,     ModBind,-    SigBind,+    ModTypeBind,     ValBind,     Dec,     Spec,@@ -1209,7 +1209,7 @@ decImports :: DecBase f vn -> [(String, Loc)] decImports (OpenDec x _) = modExpImports x decImports (ModDec md) = modExpImports $ modExp md-decImports SigDec {} = []+decImports ModTypeDec {} = [] decImports TypeDec {} = [] decImports ValDec {} = [] decImports (LocalDec d _) = decImports d@@ -1236,31 +1236,31 @@       where         onDec OpenDec {} = mempty         onDec ModDec {} = mempty-        onDec (SigDec sb) =-          M.singleton (sigName sb) (onSigExp (sigExp sb))+        onDec (ModTypeDec sb) =+          M.singleton (modTypeName sb) (onModTypeExp (modTypeExp sb))         onDec TypeDec {} = mempty         onDec ValDec {} = mempty         onDec (LocalDec d _) = onDec d         onDec ImportDec {} = mempty -        onSigExp (SigVar v _ _) = S.singleton $ qualLeaf v-        onSigExp (SigParens e _) = onSigExp e-        onSigExp (SigSpecs ss _) = foldMap onSpec ss-        onSigExp (SigWith e _ _) = onSigExp e-        onSigExp (SigArrow _ e1 e2 _) = onSigExp e1 <> onSigExp e2+        onModTypeExp (ModTypeVar v _ _) = S.singleton $ qualLeaf v+        onModTypeExp (ModTypeParens e _) = onModTypeExp e+        onModTypeExp (ModTypeSpecs ss _) = foldMap onSpec ss+        onModTypeExp (ModTypeWith e _ _) = onModTypeExp e+        onModTypeExp (ModTypeArrow _ e1 e2 _) = onModTypeExp e1 <> onModTypeExp e2          onSpec ValSpec {} = mempty         onSpec TypeSpec {} = mempty         onSpec TypeAbbrSpec {} = mempty-        onSpec (ModSpec vn e _ _) = S.singleton vn <> onSigExp e-        onSpec (IncludeSpec e _) = onSigExp e+        onSpec (ModSpec vn e _ _) = S.singleton vn <> onModTypeExp e+        onSpec (IncludeSpec e _) = onModTypeExp e      mtypes_used = foldMap onDec $ progDecs prog       where         onDec (OpenDec x _) = onModExp x         onDec (ModDec md) =-          maybe mempty (onSigExp . fst) (modSignature md) <> onModExp (modExp md)-        onDec SigDec {} = mempty+          maybe mempty (onModTypeExp . fst) (modType md) <> onModExp (modExp md)+        onDec ModTypeDec {} = mempty         onDec TypeDec {} = mempty         onDec ValDec {} = mempty         onDec LocalDec {} = mempty@@ -1271,17 +1271,17 @@         onModExp ModImport {} = mempty         onModExp (ModDecs ds _) = mconcat $ map onDec ds         onModExp (ModApply me1 me2 _ _ _) = onModExp me1 <> onModExp me2-        onModExp (ModAscript me se _ _) = onModExp me <> onSigExp se+        onModExp (ModAscript me se _ _) = onModExp me <> onModTypeExp se         onModExp (ModLambda p r me _) =-          onModParam p <> maybe mempty (onSigExp . fst) r <> onModExp me+          onModParam p <> maybe mempty (onModTypeExp . fst) r <> onModExp me -        onModParam = onSigExp . modParamType+        onModParam = onModTypeExp . modParamType -        onSigExp (SigVar v _ _) = S.singleton $ qualLeaf v-        onSigExp (SigParens e _) = onSigExp e-        onSigExp SigSpecs {} = mempty-        onSigExp (SigWith e _ _) = onSigExp e-        onSigExp (SigArrow _ e1 e2 _) = onSigExp e1 <> onSigExp e2+        onModTypeExp (ModTypeVar v _ _) = S.singleton $ qualLeaf v+        onModTypeExp (ModTypeParens e _) = onModTypeExp e+        onModTypeExp ModTypeSpecs {} = mempty+        onModTypeExp (ModTypeWith e _ _) = onModTypeExp e+        onModTypeExp (ModTypeArrow _ e1 e2 _) = onModTypeExp e1 <> onModTypeExp e2  -- | Extract a leading @((name, namespace, file), remainder)@ from a -- documentation comment string.  These are formatted as@@ -1321,7 +1321,7 @@     holesInDec (OpenDec me _) = holesInModExp me     holesInDec (LocalDec d _) = holesInDec d     holesInDec TypeDec {} = mempty-    holesInDec SigDec {} = mempty+    holesInDec ModTypeDec {} = mempty     holesInDec ImportDec {} = mempty      holesInModExp (ModDecs ds _) = foldMap holesInDec ds@@ -1455,7 +1455,7 @@ type ModBind = ModBindBase Info VName  -- | A type-checked module type binding.-type SigBind = SigBindBase Info VName+type ModTypeBind = ModTypeBindBase Info VName  -- | A type-checked module expression. type ModExp = ModExpBase Info VName@@ -1464,7 +1464,7 @@ type ModParam = ModParamBase Info VName  -- | A type-checked module type expression.-type SigExp = SigExpBase Info VName+type ModTypeExp = ModTypeExpBase Info VName  -- | A type-checked declaration. type Dec = DecBase Info VName@@ -1509,7 +1509,7 @@ type UncheckedModExp = ModExpBase NoInfo Name  -- | A module type expression with no type annotations.-type UncheckedSigExp = SigExpBase NoInfo Name+type UncheckedModTypeExp = ModTypeExpBase NoInfo Name  -- | A type parameter with no type annotations. type UncheckedTypeParam = TypeParamBase Name@@ -1524,7 +1524,7 @@ type UncheckedTypeBind = TypeBindBase NoInfo Name  -- | A module type binding with no type annotations.-type UncheckedSigBind = SigBindBase NoInfo Name+type UncheckedModTypeBind = ModTypeBindBase NoInfo Name  -- | A module binding with no type annotations. type UncheckedModBind = ModBindBase NoInfo Name

@@ -120,7 +120,7 @@ modParamDefs :: ModParam -> Defs modParamDefs (ModParam p se _ loc) =   M.singleton p (DefBound $ BoundModule $ locOf loc)-    <> sigExpDefs se+    <> modTypeExpDefs se  modExpDefs :: ModExp -> Defs modExpDefs ModVar {} =@@ -143,7 +143,7 @@   M.singleton (modName mbind) (DefBound $ BoundModule $ locOf mbind)     <> mconcat (map modParamDefs (modParams mbind))     <> modExpDefs (modExp mbind)-    <> case modSignature mbind of+    <> case modType mbind of       Nothing -> mempty       Just (_, Info substs) ->         M.map DefIndirect substs@@ -159,28 +159,28 @@       M.singleton v $ DefBound $ BoundType $ locOf loc     ModSpec v se _ loc ->       M.singleton v (DefBound $ BoundModuleType $ locOf loc)-        <> sigExpDefs se-    IncludeSpec se _ -> sigExpDefs se+        <> modTypeExpDefs se+    IncludeSpec se _ -> modTypeExpDefs se -sigExpDefs :: SigExp -> Defs-sigExpDefs se =+modTypeExpDefs :: ModTypeExp -> Defs+modTypeExpDefs se =   case se of-    SigVar _ (Info substs) _ -> M.map DefIndirect substs-    SigParens e _ -> sigExpDefs e-    SigSpecs specs _ -> mconcat $ map specDefs specs-    SigWith e _ _ -> sigExpDefs e-    SigArrow _ e1 e2 _ -> sigExpDefs e1 <> sigExpDefs e2+    ModTypeVar _ (Info substs) _ -> M.map DefIndirect substs+    ModTypeParens e _ -> modTypeExpDefs e+    ModTypeSpecs specs _ -> mconcat $ map specDefs specs+    ModTypeWith e _ _ -> modTypeExpDefs e+    ModTypeArrow _ e1 e2 _ -> modTypeExpDefs e1 <> modTypeExpDefs e2 -sigBindDefs :: SigBind -> Defs+sigBindDefs :: ModTypeBind -> Defs sigBindDefs sbind =-  M.singleton (sigName sbind) (DefBound $ BoundModuleType $ locOf sbind)-    <> sigExpDefs (sigExp sbind)+  M.singleton (modTypeName sbind) (DefBound $ BoundModuleType $ locOf sbind)+    <> modTypeExpDefs (modTypeExp sbind)  decDefs :: Dec -> Defs decDefs (ValDec vbind) = valBindDefs vbind decDefs (TypeDec vbind) = typeBindDefs vbind decDefs (ModDec mbind) = modBindDefs mbind-decDefs (SigDec mbind) = sigBindDefs mbind+decDefs (ModTypeDec mbind) = sigBindDefs mbind decDefs (OpenDec me _) = modExpDefs me decDefs (LocalDec dec _) = decDefs dec decDefs ImportDec {} = mempty@@ -312,19 +312,19 @@     ValSpec _ _ te _ _ _ -> atPosInTypeExp te pos     TypeAbbrSpec tbind -> atPosInTypeBind tbind pos     TypeSpec {} -> Nothing-    ModSpec _ se _ _ -> atPosInSigExp se pos-    IncludeSpec se _ -> atPosInSigExp se pos+    ModSpec _ se _ _ -> atPosInModTypeExp se pos+    IncludeSpec se _ -> atPosInModTypeExp se pos -atPosInSigExp :: SigExp -> Pos -> Maybe RawAtPos-atPosInSigExp se pos =+atPosInModTypeExp :: ModTypeExp -> Pos -> Maybe RawAtPos+atPosInModTypeExp se pos =   case se of-    SigVar qn _ loc -> do+    ModTypeVar qn _ loc -> do       guard $ loc `contains` pos       Just $ RawAtName qn $ locOf loc-    SigParens e _ -> atPosInSigExp e pos-    SigSpecs specs _ -> msum $ map (`atPosInSpec` pos) specs-    SigWith e _ _ -> atPosInSigExp e pos-    SigArrow _ e1 e2 _ -> atPosInSigExp e1 pos `mplus` atPosInSigExp e2 pos+    ModTypeParens e _ -> atPosInModTypeExp e pos+    ModTypeSpecs specs _ -> msum $ map (`atPosInSpec` pos) specs+    ModTypeWith e _ _ -> atPosInModTypeExp e pos+    ModTypeArrow _ e1 e2 _ -> atPosInModTypeExp e1 pos `mplus` atPosInModTypeExp e2 pos  atPosInValBind :: ValBind -> Pos -> Maybe RawAtPos atPosInValBind vbind pos =@@ -341,12 +341,12 @@     `mplus` atPosInModExp e pos     `mplus` case sig of       Nothing -> Nothing-      Just (se, _) -> atPosInSigExp se pos+      Just (se, _) -> atPosInModTypeExp se pos   where-    inParam (ModParam _ se _ _) = atPosInSigExp se pos+    inParam (ModParam _ se _ _) = atPosInModTypeExp se pos -atPosInSigBind :: SigBind -> Pos -> Maybe RawAtPos-atPosInSigBind = atPosInSigExp . sigExp+atPosInModTypeBind :: ModTypeBind -> Pos -> Maybe RawAtPos+atPosInModTypeBind = atPosInModTypeExp . modTypeExp  atPosInDec :: Dec -> Pos -> Maybe RawAtPos atPosInDec dec pos = do@@ -355,7 +355,7 @@     ValDec vbind -> atPosInValBind vbind pos     TypeDec tbind -> atPosInTypeBind tbind pos     ModDec mbind -> atPosInModBind mbind pos-    SigDec sbind -> atPosInSigBind sbind pos+    ModTypeDec sbind -> atPosInModTypeBind sbind pos     OpenDec e _ -> atPosInModExp e pos     LocalDec dec' _ -> atPosInDec dec' pos     ImportDec {} -> Nothing

@@ -12,7 +12,7 @@     Namespace (..),     Env (..),     TySet,-    FunSig (..),+    FunModType (..),     NameMap,     BoundV (..),     Mod (..),@@ -96,15 +96,15 @@ -- or a parametric module ("functor" in SML). data Mod   = ModEnv Env-  | ModFun FunSig+  | ModFun FunModType   deriving (Show)  -- | A parametric functor consists of a set of abstract types, the -- environment of its parameter, and the resulting module type.-data FunSig = FunSig-  { funSigAbs :: TySet,-    funSigMod :: Mod,-    funSigMty :: MTy+data FunModType = FunModType+  { funModTypeAbs :: TySet,+    funModTypeMod :: Mod,+    funModTypeMty :: MTy   }   deriving (Show) @@ -137,7 +137,7 @@ data Env = Env   { envVtable :: M.Map VName BoundV,     envTypeTable :: M.Map VName TypeBinding,-    envSigTable :: M.Map VName MTy,+    envModTypeTable :: M.Map VName MTy,     envModTable :: M.Map VName Mod,     envNameMap :: NameMap   }@@ -160,7 +160,7 @@  instance Pretty Mod where   pretty (ModEnv e) = pretty e-  pretty (ModFun (FunSig _ mod mty)) = pretty mod <+> "->" </> pretty mty+  pretty (ModFun (FunModType _ mod mty)) = pretty mod <+> "->" </> pretty mty  instance Pretty Env where   pretty (Env vtable ttable sigtable modtable _) =

@@ -66,9 +66,9 @@     -- * Module language     ImportName (..),     SpecBase (..),-    SigExpBase (..),+    ModTypeExpBase (..),     TypeRefBase (..),-    SigBindBase (..),+    ModTypeBindBase (..),     ModExpBase (..),     ModBindBase (..),     ModParamBase (..),@@ -1123,8 +1123,8 @@   | TypeAbbrSpec (TypeBindBase f vn)   | -- | Abstract type.     TypeSpec Liftedness vn [TypeParamBase vn] (Maybe DocComment) SrcLoc-  | ModSpec vn (SigExpBase f vn) (Maybe DocComment) SrcLoc-  | IncludeSpec (SigExpBase f vn) SrcLoc+  | ModSpec vn (ModTypeExpBase f vn) (Maybe DocComment) SrcLoc+  | IncludeSpec (ModTypeExpBase f vn) SrcLoc  deriving instance Show (SpecBase Info VName) @@ -1138,16 +1138,16 @@   locOf (IncludeSpec _ loc) = locOf loc  -- | A module type expression.-data SigExpBase f vn-  = SigVar (QualName vn) (f (M.Map VName VName)) SrcLoc-  | SigParens (SigExpBase f vn) SrcLoc-  | SigSpecs [SpecBase f vn] SrcLoc-  | SigWith (SigExpBase f vn) (TypeRefBase f vn) SrcLoc-  | SigArrow (Maybe vn) (SigExpBase f vn) (SigExpBase f vn) SrcLoc+data ModTypeExpBase f vn+  = ModTypeVar (QualName vn) (f (M.Map VName VName)) SrcLoc+  | ModTypeParens (ModTypeExpBase f vn) SrcLoc+  | ModTypeSpecs [SpecBase f vn] SrcLoc+  | ModTypeWith (ModTypeExpBase f vn) (TypeRefBase f vn) SrcLoc+  | ModTypeArrow (Maybe vn) (ModTypeExpBase f vn) (ModTypeExpBase f vn) SrcLoc -deriving instance Show (SigExpBase Info VName)+deriving instance Show (ModTypeExpBase Info VName) -deriving instance Show (SigExpBase NoInfo Name)+deriving instance Show (ModTypeExpBase NoInfo Name)  -- | A type refinement. data TypeRefBase f vn = TypeRef (QualName vn) [TypeParamBase vn] (TypeExp f vn) SrcLoc@@ -1159,27 +1159,27 @@ instance Located (TypeRefBase f vn) where   locOf (TypeRef _ _ _ loc) = locOf loc -instance Located (SigExpBase f vn) where-  locOf (SigVar _ _ loc) = locOf loc-  locOf (SigParens _ loc) = locOf loc-  locOf (SigSpecs _ loc) = locOf loc-  locOf (SigWith _ _ loc) = locOf loc-  locOf (SigArrow _ _ _ loc) = locOf loc+instance Located (ModTypeExpBase f vn) where+  locOf (ModTypeVar _ _ loc) = locOf loc+  locOf (ModTypeParens _ loc) = locOf loc+  locOf (ModTypeSpecs _ loc) = locOf loc+  locOf (ModTypeWith _ _ loc) = locOf loc+  locOf (ModTypeArrow _ _ _ loc) = locOf loc  -- | Module type binding.-data SigBindBase f vn = SigBind-  { sigName :: vn,-    sigExp :: SigExpBase f vn,-    sigDoc :: Maybe DocComment,-    sigLoc :: SrcLoc+data ModTypeBindBase f vn = ModTypeBind+  { modTypeName :: vn,+    modTypeExp :: ModTypeExpBase f vn,+    modTypeDoc :: Maybe DocComment,+    modTypeLoc :: SrcLoc   } -deriving instance Show (SigBindBase Info VName)+deriving instance Show (ModTypeBindBase Info VName) -deriving instance Show (SigBindBase NoInfo Name)+deriving instance Show (ModTypeBindBase NoInfo Name) -instance Located (SigBindBase f vn) where-  locOf = locOf . sigLoc+instance Located (ModTypeBindBase f vn) where+  locOf = locOf . modTypeLoc  -- | Canonical reference to a Futhark code file.  Does not include the -- @.fut@ extension.  This is most often a path relative to the@@ -1205,10 +1205,10 @@       (f (M.Map VName VName))       (f (M.Map VName VName))       SrcLoc-  | ModAscript (ModExpBase f vn) (SigExpBase f vn) (f (M.Map VName VName)) SrcLoc+  | ModAscript (ModExpBase f vn) (ModTypeExpBase f vn) (f (M.Map VName VName)) SrcLoc   | ModLambda       (ModParamBase f vn)-      (Maybe (SigExpBase f vn, f (M.Map VName VName)))+      (Maybe (ModTypeExpBase f vn, f (M.Map VName VName)))       (ModExpBase f vn)       SrcLoc @@ -1229,7 +1229,7 @@ data ModBindBase f vn = ModBind   { modName :: vn,     modParams :: [ModParamBase f vn],-    modSignature :: Maybe (SigExpBase f vn, f (M.Map VName VName)),+    modType :: Maybe (ModTypeExpBase f vn, f (M.Map VName VName)),     modExp :: ModExpBase f vn,     modDoc :: Maybe DocComment,     modLocation :: SrcLoc@@ -1245,7 +1245,7 @@ -- | A module parameter. data ModParamBase f vn = ModParam   { modParamName :: vn,-    modParamType :: SigExpBase f vn,+    modParamType :: ModTypeExpBase f vn,     modParamAbs :: f [VName],     modParamLocation :: SrcLoc   }@@ -1261,7 +1261,7 @@ data DecBase f vn   = ValDec (ValBindBase f vn)   | TypeDec (TypeBindBase f vn)-  | SigDec (SigBindBase f vn)+  | ModTypeDec (ModTypeBindBase f vn)   | ModDec (ModBindBase f vn)   | OpenDec (ModExpBase f vn) SrcLoc   | LocalDec (DecBase f vn) SrcLoc@@ -1274,7 +1274,7 @@ instance Located (DecBase f vn) where   locOf (ValDec d) = locOf d   locOf (TypeDec d) = locOf d-  locOf (SigDec d) = locOf d+  locOf (ModTypeDec d) = locOf d   locOf (ModDec d) = locOf d   locOf (OpenDec _ loc) = locOf loc   locOf (LocalDec _ loc) = locOf loc

@@ -178,7 +178,7 @@       check Term (valBindName vb) (srclocOf vb)     f (TypeDec (TypeBind name _ _ _ _ _ loc)) =       check Type name loc-    f (SigDec (SigBind name _ _ loc)) =+    f (ModTypeDec (ModTypeBind name _ _ loc)) =       check Signature name loc     f (ModDec (ModBind name _ _ _ _ loc)) =       check Term name loc@@ -279,7 +279,7 @@ checkSpecs (ModSpec name sig doc loc : specs) =   bindSpaced [(Term, name)] $ do     name' <- checkName Term name loc-    (_sig_abs, mty, sig') <- checkSigExp sig+    (_sig_abs, mty, sig') <- checkModTypeExp sig     let senv =           mempty             { envNameMap = M.singleton (Term, name) $ qualName name',@@ -292,7 +292,7 @@         ModSpec name' sig' doc loc : specs'       ) checkSpecs (IncludeSpec e loc : specs) = do-  (e_abs, env_abs, e_env, e') <- checkSigExpToEnv e+  (e_abs, env_abs, e_env, e') <- checkModTypeExpToEnv e    mapM_ (warnIfShadowing . fmap baseName) $ M.keys env_abs @@ -309,28 +309,28 @@     warnAbout qn =       warn loc $ "Inclusion shadows type" <+> dquotes (pretty qn) <+> "." -checkSigExp :: SigExpBase NoInfo Name -> TypeM (TySet, MTy, SigExpBase Info VName)-checkSigExp (SigParens e loc) = do-  (abs, mty, e') <- checkSigExp e-  pure (abs, mty, SigParens e' loc)-checkSigExp (SigVar name NoInfo loc) = do+checkModTypeExp :: ModTypeExpBase NoInfo Name -> TypeM (TySet, MTy, ModTypeExpBase Info VName)+checkModTypeExp (ModTypeParens e loc) = do+  (abs, mty, e') <- checkModTypeExp e+  pure (abs, mty, ModTypeParens e' loc)+checkModTypeExp (ModTypeVar name NoInfo loc) = do   (name', mty) <- lookupMTy loc name   (mty', substs) <- newNamesForMTy mty-  pure (mtyAbs mty', mty', SigVar name' (Info substs) loc)-checkSigExp (SigSpecs specs loc) = do+  pure (mtyAbs mty', mty', ModTypeVar name' (Info substs) loc)+checkModTypeExp (ModTypeSpecs specs loc) = do   checkForDuplicateSpecs specs   (abstypes, env, specs') <- checkSpecs specs-  pure (abstypes, MTy abstypes $ ModEnv env, SigSpecs specs' loc)-checkSigExp (SigWith s (TypeRef tname ps te trloc) loc) = do-  (abs, s_abs, s_env, s') <- checkSigExpToEnv s+  pure (abstypes, MTy abstypes $ ModEnv env, ModTypeSpecs specs' loc)+checkModTypeExp (ModTypeWith s (TypeRef tname ps te trloc) loc) = do+  (abs, s_abs, s_env, s') <- checkModTypeExpToEnv s   checkTypeParams ps $ \ps' -> do     (ext, te', te_t, _) <- bindingTypeParams ps' $ checkTypeDecl te     unless (null ext) $       typeError te' mempty "Anonymous dimensions are not allowed here."     (tname', s_abs', s_env') <- refineEnv loc s_abs s_env tname ps' te_t-    pure (abs, MTy s_abs' $ ModEnv s_env', SigWith s' (TypeRef tname' ps' te' trloc) loc)-checkSigExp (SigArrow maybe_pname e1 e2 loc) = do-  (e1_abs, MTy s_abs e1_mod, e1') <- checkSigExp e1+    pure (abs, MTy s_abs' $ ModEnv s_env', ModTypeWith s' (TypeRef tname' ps' te' trloc) loc)+checkModTypeExp (ModTypeArrow maybe_pname e1 e2 loc) = do+  (e1_abs, MTy s_abs e1_mod, e1') <- checkModTypeExp e1   (env_for_e2, maybe_pname') <-     case maybe_pname of       Just pname -> bindSpaced [(Term, pname)] $ do@@ -344,34 +344,34 @@           )       Nothing ->         pure (mempty, Nothing)-  (e2_abs, e2_mod, e2') <- localEnv env_for_e2 $ checkSigExp e2+  (e2_abs, e2_mod, e2') <- localEnv env_for_e2 $ checkModTypeExp e2   pure     ( e1_abs <> e2_abs,-      MTy mempty $ ModFun $ FunSig s_abs e1_mod e2_mod,-      SigArrow maybe_pname' e1' e2' loc+      MTy mempty $ ModFun $ FunModType s_abs e1_mod e2_mod,+      ModTypeArrow maybe_pname' e1' e2' loc     ) -checkSigExpToEnv ::-  SigExpBase NoInfo Name ->-  TypeM (TySet, TySet, Env, SigExpBase Info VName)-checkSigExpToEnv e = do-  (abs, MTy mod_abs mod, e') <- checkSigExp e+checkModTypeExpToEnv ::+  ModTypeExpBase NoInfo Name ->+  TypeM (TySet, TySet, Env, ModTypeExpBase Info VName)+checkModTypeExpToEnv e = do+  (abs, MTy mod_abs mod, e') <- checkModTypeExp e   case mod of     ModEnv env -> pure (abs, mod_abs, env, e')     ModFun {} -> unappliedFunctor $ srclocOf e -checkSigBind :: SigBindBase NoInfo Name -> TypeM (TySet, Env, SigBindBase Info VName)-checkSigBind (SigBind name e doc loc) = do-  (abs, env, e') <- checkSigExp e+checkModTypeBind :: ModTypeBindBase NoInfo Name -> TypeM (TySet, Env, ModTypeBindBase Info VName)+checkModTypeBind (ModTypeBind name e doc loc) = do+  (abs, env, e') <- checkModTypeExp e   bindSpaced [(Signature, name)] $ do     name' <- checkName Signature name loc     pure       ( abs,         mempty-          { envSigTable = M.singleton name' env,+          { envModTypeTable = M.singleton name' env,             envNameMap = M.singleton (Signature, name) (qualName name')           },-        SigBind name' e' doc loc+        ModTypeBind name' e' doc loc       )  checkOneModExp ::@@ -418,7 +418,7 @@       typeError loc mempty "Cannot apply non-parametric module." checkOneModExp (ModAscript me se NoInfo loc) = do   (me_abs, me_mod, me') <- checkOneModExp me-  (se_abs, se_mty, se') <- checkSigExp se+  (se_abs, se_mty, se') <- checkModTypeExp se   match_subst <- badOnLeft $ matchMTys me_mod se_mty (locOf loc)   pure (se_abs <> me_abs, se_mty, ModAscript me' se' (Info match_subst) loc) checkOneModExp (ModLambda param maybe_fsig_e body_e loc) =@@ -427,7 +427,7 @@       checkModBody (fst <$> maybe_fsig_e) body_e loc     pure       ( abs,-        MTy mempty $ ModFun $ FunSig param_abs param_mod mty,+        MTy mempty $ ModFun $ FunModType param_abs param_mod mty,         ModLambda param' maybe_fsig_e' body_e' loc       ) @@ -443,7 +443,7 @@   (ModParamBase Info VName -> TySet -> Mod -> TypeM a) ->   TypeM a withModParam (ModParam pname psig_e NoInfo loc) m = do-  (_abs, MTy p_abs p_mod, psig_e') <- checkSigExp psig_e+  (_abs, MTy p_abs p_mod, psig_e') <- checkModTypeExp psig_e   bindSpaced [(Term, pname)] $ do     pname' <- checkName Term pname loc     let in_body_env = mempty {envModTable = M.singleton pname' p_mod}@@ -460,12 +460,12 @@     withModParams ps $ \ps' -> m $ (p', pabs, pmod) : ps'  checkModBody ::-  Maybe (SigExpBase NoInfo Name) ->+  Maybe (ModTypeExpBase NoInfo Name) ->   ModExpBase NoInfo Name ->   SrcLoc ->   TypeM     ( TySet,-      Maybe (SigExp, Info (M.Map VName VName)),+      Maybe (ModTypeExp, Info (M.Map VName VName)),       ModExp,       MTy     )@@ -480,7 +480,7 @@           body_mty         )     Just fsig_e -> do-      (fsig_abs, fsig_mty, fsig_e') <- checkSigExp fsig_e+      (fsig_abs, fsig_mty, fsig_e') <- checkModTypeExp fsig_e       fsig_subst <- badOnLeft $ matchMTys body_mty fsig_mty (locOf loc)       pure         ( fsig_abs <> body_e_abs,@@ -508,13 +508,13 @@       withModParams ps $ \params_stuff -> do         let (ps', ps_abs, ps_mod) = unzip3 params_stuff         (abs, maybe_fsig_e', body_e', mty) <- checkModBody (fst <$> maybe_fsig_e) body_e loc-        let addParam (x, y) mty' = MTy mempty $ ModFun $ FunSig x y mty'+        let addParam (x, y) mty' = MTy mempty $ ModFun $ FunModType x y mty'         pure           ( abs,             p' : ps',             maybe_fsig_e',             body_e',-            FunSig p_abs p_mod $ foldr addParam mty $ zip ps_abs ps_mod+            FunModType p_abs p_mod $ foldr addParam mty $ zip ps_abs ps_mod           )   bindSpaced [(Term, name)] $ do     name' <- checkName Term name loc@@ -755,9 +755,9 @@ checkOneDec (ModDec struct) = do   (abs, modenv, struct') <- checkModBind struct   pure (abs, modenv, ModDec struct')-checkOneDec (SigDec sig) = do-  (abs, sigenv, sig') <- checkSigBind sig-  pure (abs, sigenv, SigDec sig')+checkOneDec (ModTypeDec sig) = do+  (abs, sigenv, sig') <- checkModTypeBind sig+  pure (abs, sigenv, ModTypeDec sig') checkOneDec (TypeDec tdec) = do   (tenv, tdec') <- checkTypeBind tdec   pure (mempty, tenv, TypeDec tdec')

@@ -25,8 +25,8 @@ substituteTypesInMod :: TypeSubs -> Mod -> Mod substituteTypesInMod substs (ModEnv e) =   ModEnv $ substituteTypesInEnv substs e-substituteTypesInMod substs (ModFun (FunSig abs mod mty)) =-  ModFun $ FunSig abs (substituteTypesInMod substs mod) (substituteTypesInMTy substs mty)+substituteTypesInMod substs (ModFun (FunModType abs mod mty)) =+  ModFun $ FunModType abs (substituteTypesInMod substs mod) (substituteTypesInMTy substs mty)  substituteTypesInMTy :: TypeSubs -> MTy -> MTy substituteTypesInMTy substs (MTy abs mod) = MTy abs $ substituteTypesInMod substs mod@@ -97,7 +97,7 @@            in Env                 { envVtable = vtable',                   envTypeTable = ttable',-                  envSigTable = mempty,+                  envModTypeTable = mempty,                   envModTable = mtable',                   envNameMap = M.map (fmap substitute) names                 }@@ -122,10 +122,10 @@         substituteInMod (ModEnv env) =           ModEnv $ substituteInEnv env         substituteInMod (ModFun funsig) =-          ModFun $ substituteInFunSig funsig+          ModFun $ substituteInFunModType funsig -        substituteInFunSig (FunSig abs mod mty) =-          FunSig+        substituteInFunModType (FunModType abs mod mty) =+          FunModType             (M.mapKeys (fmap substitute) abs)             (substituteInMod mod)             (substituteInMTy substs mty)@@ -169,7 +169,7 @@ modTypeAbbrs :: Mod -> M.Map VName TypeBinding modTypeAbbrs (ModEnv env) =   envTypeAbbrs env-modTypeAbbrs (ModFun (FunSig _ mod mty)) =+modTypeAbbrs (ModFun (FunModType _ mod mty)) =   modTypeAbbrs mod <> mtyTypeAbbrs mty  envTypeAbbrs :: Env -> M.Map VName TypeBinding@@ -309,8 +309,8 @@     resolveModNames (ModEnv mod_env) (ModEnv sig_env) =       resolveEnvNames mod_env sig_env     resolveModNames (ModFun mod_fun) (ModFun sig_fun) =-      resolveModNames (funSigMod mod_fun) (funSigMod sig_fun)-        <> resolveMTyNames' (funSigMty mod_fun) (funSigMty sig_fun)+      resolveModNames (funModTypeMod mod_fun) (funModTypeMod sig_fun)+        <> resolveMTyNames' (funModTypeMty mod_fun) (funModTypeMty sig_fun)     resolveModNames _ _ =       mempty @@ -453,8 +453,8 @@     matchMods       old_abs_subst_to_type       quals-      (ModFun (FunSig mod_abs mod_pmod mod_mod))-      (ModFun (FunSig sig_abs sig_pmod sig_mod))+      (ModFun (FunModType mod_abs mod_pmod mod_mod))+      (ModFun (FunModType sig_abs sig_pmod sig_mod))       loc = do         -- We need to use different substitutions when matching         -- parameter and body signatures - this is because the@@ -602,14 +602,14 @@ -- | Apply a parametric module to an argument. applyFunctor ::   Loc ->-  FunSig ->+  FunModType ->   MTy ->   TypeM     ( MTy,       M.Map VName VName,       M.Map VName VName     )-applyFunctor applyloc (FunSig p_abs p_mod body_mty) a_mty = do+applyFunctor applyloc (FunModType p_abs p_mod body_mty) a_mty = do   p_subst <- badOnLeft $ matchMTys a_mty (MTy p_abs p_mod) applyloc    -- Apply type abbreviations from a_mty to body_mty.

@@ -30,7 +30,7 @@     module Language.Futhark.Warnings,     Env (..),     TySet,-    FunSig (..),+    FunModType (..),     ImportTable,     NameMap,     BoundV (..),@@ -240,7 +240,7 @@ lookupMTy :: SrcLoc -> QualName Name -> TypeM (QualName VName, MTy) lookupMTy loc qn = do   (scope, qn'@(QualName _ name)) <- checkQualNameWithEnv Signature qn loc-  (qn',) <$> maybe explode pure (M.lookup name $ envSigTable scope)+  (qn',) <$> maybe explode pure (M.lookup name $ envModTypeTable scope)   where     explode = unknownVariable Signature qn loc 

@@ -631,7 +631,7 @@                 let f x = if x == v then Just (ExpSubst e') else Nothing                 pure (applySubst f body_t, [])           _ ->-            unscopeType loc (patNames pat') body_t+            unscopeType loc (map sizeName sizes' <> patNames pat') body_t        pure $         AppExp@@ -1695,7 +1695,7 @@     let keep_type_vars = overloadedTypeVars now_substs      cur_lvl <- curLevel-    let candidate k (lvl, _) = (k `S.notMember` keep_type_vars) && lvl >= cur_lvl+    let candidate k (lvl, _) = (k `S.notMember` keep_type_vars) && lvl >= (cur_lvl - length params)         new_substs = M.filterWithKey candidate now_substs      (tparams', RetType ret_dims restype') <-

@@ -343,7 +343,7 @@   checkTypeParams tps $ \tps' -> bindingTypeParams tps' $ do     let descend ps' (p : ps) =           checkPat [] p NoneInferred $ \p' ->-            binding (patIdents $ fmap toStruct p') $ descend (p' : ps') ps+            binding (patIdents $ fmap toStruct p') $ incLevel $ descend (p' : ps') ps         descend ps' [] = m tps' $ reverse ps'      descend [] orig_ps

@@ -547,11 +547,11 @@           }  applyType ::-  (Monoid als) =>+  (Monoid u) =>   [TypeParam] ->-  TypeBase Size als ->+  TypeBase Size u ->   [StructTypeArg] ->-  TypeBase Size als+  TypeBase Size u applyType ps t args = substTypesAny (`M.lookup` substs) t   where     substs = M.fromList $ zipWith mkSubst ps args@@ -564,10 +564,10 @@       error $ "applyType mkSubst: cannot substitute " ++ prettyString a ++ " for " ++ prettyString p  substTypesRet ::-  (Monoid as) =>-  (VName -> Maybe (Subst (RetTypeBase Size as))) ->-  TypeBase Size as ->-  RetTypeBase Size as+  (Monoid u) =>+  (VName -> Maybe (Subst (RetTypeBase Size u))) ->+  TypeBase Size u ->+  RetTypeBase Size u substTypesRet lookupSubst ot =   uncurry (flip RetType) $ runState (onType ot) []   where@@ -614,8 +614,8 @@           pure $ Scalar $ TypeVar u v targs'     onType (Scalar (Record ts)) =       Scalar . Record <$> traverse onType ts-    onType (Scalar (Arrow als v d t1 t2)) =-      Scalar <$> (Arrow als v d <$> onType t1 <*> onRetType t2)+    onType (Scalar (Arrow u v d t1 t2)) =+      Scalar <$> (Arrow u v d <$> onType t1 <*> onRetType t2)     onType (Scalar (Sum ts)) =       Scalar . Sum <$> traverse (traverse onType) ts @@ -642,10 +642,10 @@ -- | Perform substitutions, from type names to types, on a type. Works -- regardless of what shape and uniqueness information is attached to the type. substTypesAny ::-  (Monoid as) =>-  (VName -> Maybe (Subst (RetTypeBase Size as))) ->-  TypeBase Size as ->-  TypeBase Size as+  (Monoid u) =>+  (VName -> Maybe (Subst (RetTypeBase Size u))) ->+  TypeBase Size u ->+  TypeBase Size u substTypesAny lookupSubst ot =   case substTypesRet lookupSubst ot of     RetType [] ot' -> ot'