futhark 0.25.15 → 0.25.16

raw patch · 60 files changed

+2635/−655 lines, 60 filesPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

API changes (from Hackage documentation)

Files

• CHANGELOG.md +29/−0
• docs/man/futhark-test.rst +8/−1
• docs/usage.rst +2/−2
• futhark.cabal +11/−2
• rts/c/backends/cuda.h +7/−1
• rts/c/backends/hip.h +7/−0
• rts/c/backends/opencl.h +6/−0
• rts/c/context.h +46/−0
• rts/c/tuning.h +10/−0
• src/Futhark/AD/Fwd.hs +2/−2
• src/Futhark/AD/Rev.hs +1/−1
• src/Futhark/AD/Rev/Monad.hs +3/−3
• src/Futhark/Actions.hs +11/−0
• src/Futhark/Analysis/AccessPattern.hs +744/−0
• src/Futhark/Analysis/PrimExp/Table.hs +169/−0
• src/Futhark/CLI/Dev.hs +59/−5
• src/Futhark/CLI/Eval.hs +1/−0
• src/Futhark/CLI/Literate.hs +5/−0
• src/Futhark/CLI/Script.hs +1/−0
• src/Futhark/CLI/Test.hs +35/−3
• src/Futhark/CodeGen/Backends/GenericC.hs +13/−80
• src/Futhark/CodeGen/ImpGen.hs +6/−2
• src/Futhark/CodeGen/ImpGen/GPU/Base.hs +9/−5
• src/Futhark/CodeGen/ImpGen/GPU/Block.hs +2/−2
• src/Futhark/CodeGen/ImpGen/Multicore.hs +9/−5
• src/Futhark/IR/Parse.hs +4/−3
• src/Futhark/IR/Pretty.hs +6/−2
• src/Futhark/IR/Prop/Aliases.hs +1/−1
• src/Futhark/IR/Prop/TypeOf.hs +1/−1
• src/Futhark/IR/SegOp.hs +16/−0
• src/Futhark/IR/Syntax.hs +5/−3
• src/Futhark/IR/Traversals.hs +3/−3
• src/Futhark/IR/TypeCheck.hs +1/−1
• src/Futhark/Internalise/Exps.hs +1/−1
• src/Futhark/Internalise/FullNormalise.hs +1/−0
• src/Futhark/Internalise/Monomorphise.hs +103/−89
• src/Futhark/Internalise/TypesValues.hs +2/−1
• src/Futhark/MonadFreshNames.hs +1/−1
• src/Futhark/Optimise/ArrayLayout.hs +42/−0
• src/Futhark/Optimise/ArrayLayout/Layout.hs +265/−0
• src/Futhark/Optimise/ArrayLayout/Transform.hs +270/−0
• src/Futhark/Optimise/BlkRegTiling.hs +1/−1
• src/Futhark/Optimise/GenRedOpt.hs +2/−2
• src/Futhark/Optimise/HistAccs.hs +2/−2
• src/Futhark/Optimise/Simplify/Rules.hs +1/−1
• src/Futhark/Optimise/Simplify/Rules/BasicOp.hs +1/−1
• src/Futhark/Pass/ExtractKernels/Interchange.hs +3/−3
• src/Futhark/Pass/KernelBabysitting.hs +0/−407
• src/Futhark/Passes.hs +8/−5
• src/Futhark/Pipeline.hs +2/−2
• src/Futhark/Util.hs +5/−0
• src/Futhark/Version.hs +7/−4
• src/Language/Futhark/Interpreter/Values.hs +3/−3
• src/Language/Futhark/Prop.hs +1/−1
• unittests/Futhark/Analysis/PrimExp/TableTests.hs +257/−0
• unittests/Futhark/Internalise/TypesValuesTests.hs +2/−2
• unittests/Futhark/Optimise/ArrayLayout/AnalyseTests.hs +241/−0
• unittests/Futhark/Optimise/ArrayLayout/LayoutTests.hs +163/−0
• unittests/Futhark/Optimise/ArrayLayoutTests.hs +15/−0
• unittests/futhark_tests.hs +3/−1

@@ -5,6 +5,35 @@ 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.16]++### Added++* ``futhark test``: `--no-terminal` now prints status messages even when+  no failures occur.++* ``futhark test`` no longer runs ``structure`` tests by default. Pass+  ``-s`` to run them.++* Rewritten array layout optimisation pass by Bjarke Pedersen and+  Oscar Nelin. Minor speedup for some programs, but is more+  importantly a principled foundation for further improvements.++* Better error message when exceeding shared memory limits.++* Better dead code removal for the GPU representation (minor impact on+  some programs).++### Fixed++* Bugs related to deduplication of array payloads in sum types.+  Unfortunately, fixed by just not deduplicating in those cases.++* Frontend bug related to turning size expressions into variables+  (#2136).++* Another exotic monomorphisation bug.+ ## [0.25.15]  ### Added

@@ -164,13 +164,20 @@ -t   Type-check the programs, but do not run them. +-s++  Run ``structure`` tests. These are not run by default. When this+  option is passed, no other testing is done.+ --futhark=program    The program used to perform operations (eg. compilation).  Defaults   to the binary running ``futhark test`` itself.  --no-terminal-  Print each result on a line by itself, without line buffering.++  Change the output format to be suitable for noninteractive+  terminals. Prints a status message roughly every minute.  --no-tuning 

@@ -33,7 +33,7 @@  This makes use of the ``futhark c`` compiler, but any other will work as well.  The compiler will automatically invoke ``cc`` to produce an-executable binary called ``prog``.  If we had used ``futhark py``+executable binary called ``prog``.  If we had used ``futhark python`` instead of ``futhark c``, the ``prog`` file would instead have contained Python code, along with a `shebang`_ for easy execution.  In general, when compiling file ``foo.fut``, the result will be written@@ -535,7 +535,7 @@ Generating Python ^^^^^^^^^^^^^^^^^ -The ``futhark py`` and ``futhark pyopencl`` compilers both support+The ``futhark python`` and ``futhark pyopencl`` compilers both support generating reusable Python code, although the latter of these generates code of sufficient performance to be worthwhile.  The following mentions options and parameters only available for

@@ -1,6 +1,6 @@ cabal-version: 2.4 name:           futhark-version:        0.25.15+version:        0.25.16 synopsis:       An optimising compiler for a functional, array-oriented language.  description:    Futhark is a small programming language designed to be compiled to@@ -122,6 +122,7 @@       Futhark.AD.Rev.Scan       Futhark.AD.Rev.Scatter       Futhark.AD.Rev.SOAC+      Futhark.Analysis.AccessPattern       Futhark.Analysis.AlgSimplify       Futhark.Analysis.Alias       Futhark.Analysis.CallGraph@@ -137,6 +138,7 @@       Futhark.Analysis.PrimExp.Convert       Futhark.Analysis.PrimExp.Parse       Futhark.Analysis.PrimExp.Simplify+      Futhark.Analysis.PrimExp.Table       Futhark.Analysis.SymbolTable       Futhark.Analysis.UsageTable       Futhark.Bench@@ -323,6 +325,9 @@       Futhark.Optimise.ArrayShortCircuiting.MemRefAggreg       Futhark.Optimise.ArrayShortCircuiting.TopdownAnalysis       Futhark.Optimise.MergeGPUBodies+      Futhark.Optimise.ArrayLayout+      Futhark.Optimise.ArrayLayout.Transform+      Futhark.Optimise.ArrayLayout.Layout       Futhark.Optimise.ReduceDeviceSyncs       Futhark.Optimise.ReduceDeviceSyncs.MigrationTable       Futhark.Optimise.ReduceDeviceSyncs.MigrationTable.Graph@@ -360,7 +365,6 @@       Futhark.Pass.ExtractKernels.ToGPU       Futhark.Pass.ExtractMulticore       Futhark.Pass.FirstOrderTransform-      Futhark.Pass.KernelBabysitting       Futhark.Pass.LiftAllocations       Futhark.Pass.LowerAllocations       Futhark.Pass.Simplify@@ -501,6 +505,7 @@   other-modules:       Futhark.AD.DerivativesTests       Futhark.Analysis.AlgSimplifyTests+      Futhark.Analysis.PrimExp.TableTests       Futhark.BenchTests       Futhark.IR.GPUTests       Futhark.IR.MCTests@@ -515,6 +520,9 @@       Futhark.IR.SyntaxTests       Futhark.Internalise.TypesValuesTests       Futhark.Optimise.MemoryBlockMerging.GreedyColoringTests+      Futhark.Optimise.ArrayLayout.AnalyseTests+      Futhark.Optimise.ArrayLayout.LayoutTests+      Futhark.Optimise.ArrayLayoutTests       Futhark.Pkg.SolveTests       Futhark.ProfileTests       Language.Futhark.CoreTests@@ -525,6 +533,7 @@       Paths_futhark   build-depends:       QuickCheck >=2.8+      , mtl >=2.2.1     , base     , containers     , free

@@ -1056,8 +1056,14 @@                              void* args[num_args],                              size_t args_sizes[num_args]) {   (void) args_sizes;-  int64_t time_start = 0, time_end = 0; +  if (shared_mem_bytes > ctx->max_shared_memory) {+    set_error(ctx, msgprintf("Kernel %s with %d bytes of memory exceeds device limit of %d\n",+                             name, shared_mem_bytes, (int)ctx->max_shared_memory));+    return 1;+  }++  int64_t time_start = 0, time_end = 0;   if (ctx->debugging) {     time_start = get_wall_time();   }

@@ -913,6 +913,13 @@                              void* args[num_args],                              size_t args_sizes[num_args]) {   (void) args_sizes;++  if (shared_mem_bytes > ctx->max_shared_memory) {+    set_error(ctx, msgprintf("Kernel %s with %d bytes of memory exceeds device limit of %d\n",+                             name, shared_mem_bytes, (int)ctx->max_shared_memory));+    return 1;+  }+   int64_t time_start = 0, time_end = 0;   if (ctx->debugging) {     time_start = get_wall_time();

@@ -1314,6 +1314,12 @@                              int num_args,                              void* args[num_args],                              size_t args_sizes[num_args]) {+  if (shared_mem_bytes > ctx->max_shared_memory) {+    set_error(ctx, msgprintf("Kernel %s with %d bytes of memory exceeds device limit of %d\n",+                             name, shared_mem_bytes, (int)ctx->max_shared_memory));+    return 1;+  }+   int64_t time_start = 0, time_end = 0;    cl_event* event = opencl_event_new(ctx);

@@ -79,6 +79,52 @@   add_event_to_list(&ctx->event_list, name, description, data, f); } +char *futhark_context_get_error(struct futhark_context *ctx) {+  char *error = ctx->error;+  ctx->error = NULL;+  return error;+}++void futhark_context_config_set_debugging(struct futhark_context_config *cfg, int flag) {+    cfg->profiling = cfg->logging = cfg->debugging = flag;+}++void futhark_context_config_set_profiling(struct futhark_context_config *cfg, int flag) {+    cfg->profiling = flag;+}++void futhark_context_config_set_logging(struct futhark_context_config *cfg, int flag) {+    cfg->logging = flag;+}++void futhark_context_config_set_cache_file(struct futhark_context_config *cfg, const char *f) {+  cfg->cache_fname = strdup(f);+}++int futhark_get_tuning_param_count(void) {+  return num_tuning_params;+}++const char *futhark_get_tuning_param_name(int i) {+  return tuning_param_names[i];+}++const char *futhark_get_tuning_param_class(int i) {+    return tuning_param_classes[i];+}++void futhark_context_set_logging_file(struct futhark_context *ctx, FILE *f){+  ctx->log = f;+}++void futhark_context_pause_profiling(struct futhark_context *ctx) {+  ctx->profiling_paused = 1;+}++void futhark_context_unpause_profiling(struct futhark_context *ctx) {+  ctx->profiling_paused = 0;+}+ struct futhark_context_config* futhark_context_config_new(void) {   struct futhark_context_config* cfg = malloc(sizeof(struct futhark_context_config));   if (cfg == NULL) {

@@ -1,5 +1,12 @@ // Start of tuning.h. ++int is_blank_line_or_comment(const char *s) {+  size_t i = strspn(s, " \t");+  return s[i] == '\0' || // Line is blank.+         strncmp(s + i, "--", 2) == 0; // Line is comment.+}+ static char* load_tuning_file(const char *fname,                               void *cfg,                               int (*set_tuning_param)(void*, const char*, size_t)) {@@ -16,6 +23,9 @@   int lineno = 0;   while (fgets(line, max_line_len, f) != NULL) {     lineno++;+    if (is_blank_line_or_comment(line)) {+      continue;+    }     char *eql = strstr(line, "=");     if (eql) {       *eql = 0;

@@ -362,11 +362,11 @@   error "fwdSOAC: nested VJP not allowed."  fwdStm :: Stm SOACS -> ADM ()-fwdStm (Let pat aux (BasicOp (UpdateAcc acc i x))) = do+fwdStm (Let pat aux (BasicOp (UpdateAcc safety acc i x))) = do   pat' <- bundleNewPat pat   x' <- bundleTangents x   acc_tan <- tangent acc-  addStm $ Let pat' aux $ BasicOp $ UpdateAcc acc_tan i x'+  addStm $ Let pat' aux $ BasicOp $ UpdateAcc safety acc_tan i x' fwdStm stm@(Let pat aux (BasicOp e)) = do   -- XXX: this has to be too naive.   unless (any isAcc $ patTypes pat) $

@@ -212,7 +212,7 @@           updateAdj arr             =<< letExp "update_src_adj" (BasicOp $ Update safety pat_adj slice zeroes)     -- See Note [Adjoints of accumulators]-    UpdateAcc _ is vs -> do+    UpdateAcc _ _ is vs -> do       addStm $ Let pat aux $ BasicOp e       m       pat_adjs <- mapM lookupAdjVal (patNames pat)

@@ -394,7 +394,7 @@           ~[v_adj'] <-             tabNest (length dims) [d, v_adj] $ \is [d', v_adj'] ->               letTupExp "acc" . BasicOp $-                UpdateAcc v_adj' (map Var is) [Var d']+                UpdateAcc Safe v_adj' (map Var is) [Var d']           insAdj v v_adj'         _ -> do           v_adj' <- letExp (baseString v <> "_adj") =<< addExp v_adj d@@ -417,7 +417,7 @@               fixSlice (fmap pe64 slice) $                 map le64 is           letTupExp (baseString v_adj') . BasicOp $-            UpdateAcc v_adj' slice' [Var d']+            UpdateAcc Safe v_adj' slice' [Var d']       pure v_adj'     _ -> do       v_adjslice <-@@ -458,7 +458,7 @@                 ~[v_adj'] <-                   tabNest (length dims) [se_v, v_adj] $ \is [se_v', v_adj'] ->                     letTupExp "acc" . BasicOp $-                      UpdateAcc v_adj' (i : map Var is) [Var se_v']+                      UpdateAcc Safe v_adj' (i : map Var is) [Var se_v']                 pure v_adj'           _ -> do             let stms s = do

@@ -7,6 +7,7 @@     printFusionGraph,     printInterferenceGPU,     printMemAliasGPU,+    printMemoryAccessAnalysis,     callGraphAction,     impCodeGenAction,     kernelImpCodeGenAction,@@ -33,6 +34,7 @@ import Data.Maybe (fromMaybe) import Data.Text qualified as T import Data.Text.IO qualified as T+import Futhark.Analysis.AccessPattern import Futhark.Analysis.Alias import Futhark.Analysis.CallGraph (buildCallGraph) import Futhark.Analysis.Interference qualified as Interference@@ -131,6 +133,15 @@     { actionName = "print mem alias gpu",       actionDescription = "Print memory alias information on gpu.",       actionProcedure = liftIO . print . MemAlias.analyzeGPUMem+    }++-- | Print result of array access analysis on the IR+printMemoryAccessAnalysis :: (Analyse rep) => Action rep+printMemoryAccessAnalysis =+  Action+    { actionName = "array-access-analysis",+      actionDescription = "Prettyprint the array access analysis to standard output.",+      actionProcedure = liftIO . putStrLn . prettyString . analyseDimAccesses     }  -- | Print call graph to stdout.

@@ -0,0 +1,744 @@+{-# LANGUAGE LambdaCase #-}++module Futhark.Analysis.AccessPattern+  ( analyseDimAccesses,+    analyseFunction,+    vnameFromSegOp,+    analysisPropagateByTransitivity,+    isInvariant,+    Analyse,+    IndexTable,+    ArrayName,+    DimAccess (..),+    IndexExprName,+    BodyType (..),+    SegOpName (SegmentedMap, SegmentedRed, SegmentedScan, SegmentedHist),+    Context (..),+    analyseIndex,+    VariableInfo (..),+    VarType (..),+    isCounter,+    Dependency (..),+  )+where++import Data.Bifunctor+import Data.Foldable+import Data.List qualified as L+import Data.Map.Strict qualified as M+import Data.Maybe+import Futhark.IR.Aliases+import Futhark.IR.GPU+import Futhark.IR.GPUMem+import Futhark.IR.MC+import Futhark.IR.MCMem+import Futhark.IR.SOACS+import Futhark.IR.Seq+import Futhark.IR.SeqMem+import Futhark.Util.Pretty++-- | Name of a SegOp, used to identify the SegOp that an array access is+-- contained in.+data SegOpName+  = SegmentedMap {vnameFromSegOp :: VName}+  | SegmentedRed {vnameFromSegOp :: VName}+  | SegmentedScan {vnameFromSegOp :: VName}+  | SegmentedHist {vnameFromSegOp :: VName}+  deriving (Eq, Ord, Show)++-- | Name of an array indexing expression. Taken from the pattern of+-- the expression.+type IndexExprName = VName++data BodyType+  = SegOpName SegOpName+  | LoopBodyName VName+  | CondBodyName VName+  deriving (Show, Ord, Eq)++-- | Stores the name of an array, the nest of loops, kernels,+-- conditionals in which it is constructed, and the existing layout of+-- the array. The latter is currently largely unused and not+-- trustworthy, but might be useful in the future.+type ArrayName = (VName, [BodyType], [Int])++-- | Tuple of patternName and nested `level` it index occurred at, as well as+-- what the actual iteration type is.+data Dependency = Dependency+  { lvl :: Int,+    varType :: VarType+  }+  deriving (Eq, Show)++-- | Collect all features of access to a specific dimension of an array.+data DimAccess rep = DimAccess+  { -- | Set of VNames of iteration variables (gtids, loop counters, etc.)+    -- that some access is variant to.+    -- An empty set indicates that the access is invariant.+    dependencies :: M.Map VName Dependency,+    -- | Used to store the name of the original expression from which `dependencies`+    -- was computed. `Nothing` if it is a constant.+    originalVar :: Maybe VName+  }+  deriving (Eq, Show)++instance Semigroup (DimAccess rep) where+  adeps <> bdeps =+    DimAccess+      (dependencies adeps <> dependencies bdeps)+      ( case originalVar adeps of+          Nothing -> originalVar bdeps+          _ -> originalVar adeps+      )++instance Monoid (DimAccess rep) where+  mempty = DimAccess mempty Nothing++isInvariant :: DimAccess rep -> Bool+isInvariant = null . dependencies++-- | For each array access in a program, this data structure stores the+-- dependencies of each dimension in the access, the array name, and the+-- name of the SegOp that the access is contained in.+-- Each DimAccess element corresponds to an access to a given dimension+-- in the given array, in the same order of the dimensions.+type IndexTable rep =+  M.Map SegOpName (M.Map ArrayName (M.Map IndexExprName [DimAccess rep]))++unionIndexTables :: IndexTable rep -> IndexTable rep -> IndexTable rep+unionIndexTables = M.unionWith (M.unionWith M.union)++-- | Make segops on arrays transitive, ie. if+-- > let A = segmap (..) xs -- A indexes into xs+-- > let B = segmap (..) A  -- B indexes into A+-- Then B also derives all A's array-accesses, like xs.+-- Runs in n²+analysisPropagateByTransitivity :: IndexTable rep -> IndexTable rep+analysisPropagateByTransitivity idx_table =+  M.map foldlArrayNameMap idx_table+  where+    aggregateResults arr_name =+      maybe+        mempty+        foldlArrayNameMap+        (M.mapKeys vnameFromSegOp idx_table M.!? arr_name)++    foldlArrayNameMap aMap =+      foldl (M.unionWith M.union) aMap $+        map (aggregateResults . \(a, _, _) -> a) $+          M.keys aMap++--+-- Helper types and functions to perform the analysis.+--++-- | Used during the analysis to keep track of the dependencies of patterns+-- encountered so far.+data Context rep = Context+  { -- | A mapping from patterns occuring in Let expressions to their dependencies+    --  and iteration types.+    assignments :: M.Map VName (VariableInfo rep),+    -- | Maps from sliced arrays to their respective access patterns.+    slices :: M.Map IndexExprName (ArrayName, [VName], [DimAccess rep]),+    -- | A list of the segMaps encountered during the analysis in the order they+    -- were encountered.+    parents :: [BodyType],+    -- | Current level of recursion, also just `length parents`+    currentLevel :: Int+  }+  deriving (Show, Eq)++instance Monoid (Context rep) where+  mempty =+    Context+      { assignments = mempty,+        slices = mempty,+        parents = [],+        currentLevel = 0+      }++instance Semigroup (Context rep) where+  Context ass0 slices0 lastBody0 lvl0 <> Context ass1 slices1 lastBody1 lvl1 =+    Context+      (ass0 <> ass1)+      (slices0 <> slices1)+      (lastBody0 <> lastBody1)+      (max lvl0 lvl1)++-- | Extend a context with another context.+-- We never have to consider the case where VNames clash in the context, since+-- they are unique.+extend :: Context rep -> Context rep -> Context rep+extend = (<>)++allSegMap :: Context rep -> [SegOpName]+allSegMap (Context _ _ parents _) = mapMaybe f parents+  where+    f (SegOpName o) = Just o+    f _ = Nothing++-- | Context Value (VariableInfo) is the type used in the context to categorize+-- assignments. For example, a pattern might depend on a function parameter, a+-- gtid, or some other pattern.+data VariableInfo rep = VariableInfo+  { deps :: Names,+    level :: Int,+    parents_nest :: [BodyType],+    variableType :: VarType+  }+  deriving (Show, Eq)++data VarType+  = ConstType+  | Variable+  | ThreadID+  | LoopVar+  deriving (Show, Eq)++isCounter :: VarType -> Bool+isCounter LoopVar = True+isCounter ThreadID = True+isCounter _ = False++varInfoFromNames :: Context rep -> Names -> VariableInfo rep+varInfoFromNames ctx names = do+  VariableInfo names (currentLevel ctx) (parents ctx) Variable++-- | Wrapper around the constructur of Context.+oneContext :: VName -> VariableInfo rep -> Context rep+oneContext name var_info =+  Context+    { assignments = M.singleton name var_info,+      slices = mempty,+      parents = [],+      currentLevel = 0+    }++-- | Create a singular varInfo with no dependencies.+varInfoZeroDeps :: Context rep -> VariableInfo rep+varInfoZeroDeps ctx =+  VariableInfo mempty (currentLevel ctx) (parents ctx) Variable++-- | Create a singular context from a segspace+contextFromNames :: Context rep -> VariableInfo rep -> [VName] -> Context rep+contextFromNames ctx var_info = foldl' extend ctx . map (`oneContext` var_info)++-- | A representation where we can analyse access patterns.+class Analyse rep where+  -- | Analyse the op for this representation.+  analyseOp :: Op rep -> Context rep -> [VName] -> (Context rep, IndexTable rep)++-- | Analyse each `entry` and accumulate the results.+analyseDimAccesses :: (Analyse rep) => Prog rep -> IndexTable rep+analyseDimAccesses = foldMap' analyseFunction . progFuns++-- | Analyse each statement in a function body.+analyseFunction :: (Analyse rep) => FunDef rep -> IndexTable rep+analyseFunction func =+  let stms = stmsToList . bodyStms $ funDefBody func+      -- Create a context containing the function parameters+      ctx = contextFromNames mempty (varInfoZeroDeps ctx) $ map paramName $ funDefParams func+   in snd $ analyseStmsPrimitive ctx stms++-- | Analyse each statement in a list of statements.+analyseStmsPrimitive :: (Analyse rep) => Context rep -> [Stm rep] -> (Context rep, IndexTable rep)+analyseStmsPrimitive ctx =+  -- Fold over statements in body+  foldl'+    (\(c, r) stm -> second (unionIndexTables r) $ analyseStm c stm)+    (ctx, mempty)++-- | Same as analyseStmsPrimitive, but change the resulting context into+-- a varInfo, mapped to pattern.+analyseStms :: (Analyse rep) => Context rep -> (VName -> BodyType) -> [VName] -> [Stm rep] -> (Context rep, IndexTable rep)+analyseStms ctx body_constructor pats body = do+  -- 0. Recurse into body with ctx+  let (ctx'', indexTable) = analyseStmsPrimitive recContext body++  -- 0.1 Get all new slices+  let slices_new = M.difference (slices ctx'') (slices ctx)+  -- 0.2 Make "IndexExpressions" of the slices+  let slices_indices =+        foldl unionIndexTables indexTable+          $ mapMaybe+            ( uncurry $ \_idx_expression (array_name, patterns, dim_indices) ->+                Just . snd $+                  -- Should we use recContex instead of ctx''?+                  analyseIndex' ctx'' patterns array_name dim_indices+            )+          $ M.toList slices_new++  -- 1. We do not want the returned context directly.+  --    however, we do want pat to map to the names what was hit in body.+  --    therefore we need to subtract the old context from the returned one,+  --    and discard all the keys within it.++  -- assignments :: M.Map VName (VariableInfo rep),+  let in_scope_dependencies_from_body =+        rmOutOfScopeDeps ctx'' $+          M.difference (assignments ctx'') (assignments recContext)++  -- 2. We are ONLY interested in the rhs of assignments (ie. the+  --    dependencies of pat :) )+  let ctx' = foldl extend ctx $ concatVariableInfo in_scope_dependencies_from_body -- . map snd $ M.toList varInfos+  -- 3. Now we have the correct context and result+  (ctx' {parents = parents ctx, currentLevel = currentLevel ctx, slices = slices ctx}, slices_indices)+  where+    -- Extracts and merges `Names` in `VariableInfo`s, and makes a new VariableInfo. This+    -- MAY throw away needed information, but it was my best guess at a solution+    -- at the time of writing.+    concatVariableInfo dependencies =+      map (\pat -> oneContext pat (varInfoFromNames ctx dependencies)) pats++    -- Context used for "recursion" into analyseStmsPrimitive+    recContext =+      ctx+        { parents = parents ctx <> concatMap (\pat -> [body_constructor pat]) pats,+          currentLevel = currentLevel ctx + 1+        }++    -- Recursively looks up dependencies, until they're in scope or empty set.+    rmOutOfScopeDeps :: Context rep -> M.Map VName (VariableInfo rep) -> Names+    rmOutOfScopeDeps ctx' new_assignments =+      let throwaway_assignments = assignments ctx'+          local_assignments = assignments ctx+          f result a var_info =+            -- if the VName of the assignment exists in the context, we are good+            if a `M.member` local_assignments+              then result <> oneName a+              else -- Otherwise, recurse on its dependencies;+              -- 0. Add dependencies in ctx to result++                let (deps_in_ctx, deps_not_in_ctx) =+                      L.partition (`M.member` local_assignments) $+                        namesToList (deps var_info)+                    deps_not_in_ctx' =+                      M.fromList $+                        mapMaybe+                          (\d -> (d,) <$> M.lookup d throwaway_assignments)+                          deps_not_in_ctx+                 in result+                      <> namesFromList deps_in_ctx+                      <> rmOutOfScopeDeps ctx' deps_not_in_ctx'+       in M.foldlWithKey f mempty new_assignments++-- | Analyse a rep statement and return the updated context and array index+-- descriptors.+analyseStm :: (Analyse rep) => Context rep -> Stm rep -> (Context rep, IndexTable rep)+analyseStm ctx (Let pats _ e) = do+  -- Get the name of the first element in a pattern+  let pattern_names = map patElemName $ patElems pats++  -- Construct the result and Context from the subexpression. If the+  -- subexpression is a body, we recurse into it.+  case e of+    BasicOp (Index name (Slice dim_subexp)) ->+      analyseIndex ctx pattern_names name dim_subexp+    BasicOp (Update _ name (Slice dim_subexp) _subexp) ->+      analyseIndex ctx pattern_names name dim_subexp+    BasicOp op ->+      analyseBasicOp ctx op pattern_names+    Match conds cases default_body _ ->+      analyseMatch ctx' pattern_names default_body $ map caseBody cases+      where+        ctx' =+          contextFromNames ctx (varInfoZeroDeps ctx) $+            concatMap (namesToList . freeIn) conds+    Loop bindings loop body ->+      analyseLoop ctx bindings loop body pattern_names+    Apply _name diets _ _ ->+      analyseApply ctx pattern_names diets+    WithAcc _ _ ->+      (ctx, mempty) -- ignored+    Op op ->+      analyseOp op ctx pattern_names++-- If left, this is just a regular index. If right, a slice happened.+getIndexDependencies :: Context rep -> [DimIndex SubExp] -> Either [DimAccess rep] [DimAccess rep]+getIndexDependencies ctx dims =+  fst $+    foldr+      ( \idx (a, i) ->+          ( either (matchDimIndex idx) (either Right Right . matchDimIndex idx) a,+            i - 1+          )+      )+      (Left [], length dims - 1)+      dims+  where+    matchDimIndex (DimFix subExpression) accumulator =+      Left $ consolidate ctx subExpression : accumulator+    -- If we encounter a DimSlice, add it to a map of `DimSlice`s and check+    -- result later.+    matchDimIndex (DimSlice offset num_elems stride) accumulator =+      -- We assume that a slice is iterated sequentially, so we have to+      -- create a fake dependency for the slice.+      let dimAccess' = DimAccess (M.singleton (VName "slice" 0) $ Dependency (currentLevel ctx) LoopVar) (Just $ VName "slice" 0)+          cons = consolidate ctx+          dimAccess = dimAccess' <> cons offset <> cons num_elems <> cons stride+       in Right $ dimAccess : accumulator++-- | Gets the dependencies of each dimension and either returns a result, or+-- adds a slice to the context.+analyseIndex :: Context rep -> [VName] -> VName -> [DimIndex SubExp] -> (Context rep, IndexTable rep)+analyseIndex ctx pats arr_name dim_indices =+  -- Get the dependendencies of each dimension+  let dependencies = getIndexDependencies ctx dim_indices+      -- Extend the current context with current pattern(s) and its deps+      ctx' = analyseIndexContextFromIndices ctx dim_indices pats++      -- The bodytype(s) are used in the result construction+      array_name' =+        -- For now, we assume the array is in row-major-order, hence the+        -- identity permutation. In the future, we might want to infer its+        -- layout, for example, if the array is the result of a transposition.+        let layout = [0 .. length dim_indices - 1]+         in -- 2. If the arrayname was not in assignments, it was not an immediately+            --    allocated array.+            fromMaybe (arr_name, [], layout)+              -- 1. Maybe find the array name, and the "stack" of body types that the+              -- array was allocated in.+              . L.find (\(n, _, _) -> n == arr_name)+              -- 0. Get the "stack" of bodytypes for each assignment+              $ map (\(n, vi) -> (n, parents_nest vi, layout)) (M.toList $ assignments ctx')+   in either (index ctx' array_name') (slice ctx' array_name') dependencies+  where+    slice :: Context rep -> ArrayName -> [DimAccess rep] -> (Context rep, IndexTable rep)+    slice context array_name dims =+      (context {slices = M.insert (head pats) (array_name, pats, dims) $ slices context}, mempty)++    index :: Context rep -> ArrayName -> [DimAccess rep] -> (Context rep, IndexTable rep)+    index context array_name@(name, _, _) dim_access =+      -- If the arrayname is a `DimSlice` we want to fixup the access+      case M.lookup name $ slices context of+        Nothing -> analyseIndex' context pats array_name dim_access+        Just (arr_name', pats', slice_access) ->+          analyseIndex'+            context+            pats'+            arr_name'+            (init slice_access ++ [head dim_access <> last slice_access] ++ drop 1 dim_access)++analyseIndexContextFromIndices :: Context rep -> [DimIndex SubExp] -> [VName] -> Context rep+analyseIndexContextFromIndices ctx dim_accesses pats =+  let subexprs =+        mapMaybe+          ( \case+              DimFix (Var v) -> Just v+              DimFix (Constant _) -> Nothing+              DimSlice _offs _n _stride -> Nothing+          )+          dim_accesses++      -- Add each non-constant DimIndex as a dependency to the index expression+      var_info = varInfoFromNames ctx $ namesFromList subexprs+   in -- Extend context with the dependencies index expression+      foldl' extend ctx $ map (`oneContext` var_info) pats++analyseIndex' ::+  Context rep ->+  [VName] ->+  ArrayName ->+  [DimAccess rep] ->+  (Context rep, IndexTable rep)+analyseIndex' ctx _ _ [] = (ctx, mempty)+analyseIndex' ctx _ _ [_] = (ctx, mempty)+analyseIndex' ctx pats arr_name dim_accesses =+  -- Get the name of all segmaps in the current "callstack"+  let segmaps = allSegMap ctx+      idx_expr_name = pats --                                                IndexExprName+      -- For each pattern, create a mapping to the dimensional indices+      map_ixd_expr = map (`M.singleton` dim_accesses) idx_expr_name --       IndexExprName |-> [DimAccess]+      -- For each pattern -> [DimAccess] mapping, create a mapping from the array+      -- name that was indexed.+      map_array = map (M.singleton arr_name) map_ixd_expr --   ArrayName |-> IndexExprName |-> [DimAccess]+      -- ∀ (arr_name -> IdxExp -> [DimAccess]) mappings, create a mapping from all+      -- segmaps in current callstack (segThread & segGroups alike).+      results = concatMap (\ma -> map (`M.singleton` ma) segmaps) map_array++      res = foldl' unionIndexTables mempty results+   in (ctx, res)++analyseBasicOp :: Context rep -> BasicOp -> [VName] -> (Context rep, IndexTable rep)+analyseBasicOp ctx expression pats =+  -- Construct a VariableInfo from the subexpressions+  let ctx_val = case expression of+        SubExp se -> varInfoFromSubExpr se+        Opaque _ se -> varInfoFromSubExpr se+        ArrayLit ses _t -> concatVariableInfos mempty ses+        UnOp _ se -> varInfoFromSubExpr se+        BinOp _ lsubexp rsubexp -> concatVariableInfos mempty [lsubexp, rsubexp]+        CmpOp _ lsubexp rsubexp -> concatVariableInfos mempty [lsubexp, rsubexp]+        ConvOp _ se -> varInfoFromSubExpr se+        Assert se _ _ -> varInfoFromSubExpr se+        Index name _ ->+          error $ "unhandled: Index (This should NEVER happen) into " ++ prettyString name+        Update _ name _slice _subexp ->+          error $ "unhandled: Update (This should NEVER happen) onto " ++ prettyString name+        -- Technically, do we need this case?+        Concat _ _ length_subexp -> varInfoFromSubExpr length_subexp+        Manifest _dim name -> varInfoFromNames ctx $ oneName name+        Iota end start stride _ -> concatVariableInfos mempty [end, start, stride]+        Replicate (Shape shape) value' -> concatVariableInfos mempty (value' : shape)+        Scratch _ sers -> concatVariableInfos mempty sers+        Reshape _ (Shape shape_subexp) name -> concatVariableInfos (oneName name) shape_subexp+        Rearrange _ name -> varInfoFromNames ctx $ oneName name+        UpdateAcc _ name lsubexprs rsubexprs ->+          concatVariableInfos (oneName name) (lsubexprs ++ rsubexprs)+        FlatIndex name _ -> varInfoFromNames ctx $ oneName name+        FlatUpdate name _ source -> varInfoFromNames ctx $ namesFromList [name, source]+      ctx' = foldl' extend ctx $ map (`oneContext` ctx_val) pats+   in (ctx', mempty)+  where+    concatVariableInfos ne nn =+      varInfoFromNames ctx (ne <> mconcat (map (analyseSubExpr pats ctx) nn))++    varInfoFromSubExpr (Constant _) = (varInfoFromNames ctx mempty) {variableType = ConstType}+    varInfoFromSubExpr (Var v) =+      case M.lookup v (assignments ctx) of+        Just _ -> (varInfoFromNames ctx $ oneName v) {variableType = Variable}+        Nothing ->+          error $+            "Failed to lookup variable \""+              ++ prettyString v+              ++ "\npat: "+              ++ prettyString pats+              ++ "\n\nContext\n"+              ++ show ctx++analyseMatch :: (Analyse rep) => Context rep -> [VName] -> Body rep -> [Body rep] -> (Context rep, IndexTable rep)+analyseMatch ctx pats body parents =+  let ctx'' = ctx {currentLevel = currentLevel ctx - 1}+   in foldl+        ( \(ctx', res) b ->+            -- This Little Maneuver's Gonna Cost Us 51 Years+            bimap constLevel (unionIndexTables res)+              . analyseStms ctx' CondBodyName pats+              . stmsToList+              $ bodyStms b+        )+        (ctx'', mempty)+        (body : parents)+  where+    constLevel context = context {currentLevel = currentLevel ctx - 1}++analyseLoop :: (Analyse rep) => Context rep -> [(FParam rep, SubExp)] -> LoopForm -> Body rep -> [VName] -> (Context rep, IndexTable rep)+analyseLoop ctx bindings loop body pats = do+  let next_level = currentLevel ctx+  let ctx'' = ctx {currentLevel = next_level}+  let ctx' =+        contextFromNames ctx'' ((varInfoZeroDeps ctx) {variableType = LoopVar}) $+          case loop of+            WhileLoop iv -> iv : map (paramName . fst) bindings+            ForLoop iv _ _ -> iv : map (paramName . fst) bindings++  -- Extend context with the loop expression+  analyseStms ctx' LoopBodyName pats $ stmsToList $ bodyStms body++analyseApply :: Context rep -> [VName] -> [(SubExp, Diet)] -> (Context rep, IndexTable rep)+analyseApply ctx pats diets =+  ( foldl' extend ctx $ map (\pat -> oneContext pat $ varInfoFromNames ctx $ mconcat $ map (freeIn . fst) diets) pats,+    mempty+  )++segOpType :: SegOp lvl rep -> VName -> SegOpName+segOpType (SegMap {}) = SegmentedMap+segOpType (SegRed {}) = SegmentedRed+segOpType (SegScan {}) = SegmentedScan+segOpType (SegHist {}) = SegmentedHist++analyseSegOp :: (Analyse rep) => SegOp lvl rep -> Context rep -> [VName] -> (Context rep, IndexTable rep)+analyseSegOp op ctx pats =+  let next_level = currentLevel ctx + length (unSegSpace $ segSpace op) - 1+      ctx' = ctx {currentLevel = next_level}+      segspace_context =+        foldl' extend ctx'+          . map (\(n, i) -> oneContext n $ VariableInfo mempty (currentLevel ctx + i) (parents ctx') ThreadID)+          . (\segspace_params -> zip segspace_params [0 ..])+          -- contextFromNames ctx' Parallel+          . map fst+          . unSegSpace+          $ segSpace op+   in -- Analyse statements in the SegOp body+      analyseStms segspace_context (SegOpName . segOpType op) pats . stmsToList . kernelBodyStms $ segBody op++analyseSizeOp :: SizeOp -> Context rep -> [VName] -> (Context rep, IndexTable rep)+analyseSizeOp op ctx pats =+  let ctx' = case op of+        CmpSizeLe _name _class subexp -> subexprsToContext [subexp]+        CalcNumBlocks lsubexp _name rsubexp -> subexprsToContext [lsubexp, rsubexp]+        _ -> ctx+      -- Add sizeOp to context+      ctx'' =+        foldl' extend ctx' $+          map+            (\pat -> oneContext pat $ (varInfoZeroDeps ctx) {parents_nest = parents ctx'})+            pats+   in (ctx'', mempty)+  where+    subexprsToContext =+      contextFromNames ctx (varInfoZeroDeps ctx)+        . concatMap (namesToList . analyseSubExpr pats ctx)++-- | Analyse statements in a rep body.+analyseGPUBody :: (Analyse rep) => Body rep -> Context rep -> (Context rep, IndexTable rep)+analyseGPUBody body ctx =+  analyseStmsPrimitive ctx $ stmsToList $ bodyStms body++analyseOtherOp :: Context rep -> [VName] -> (Context rep, IndexTable rep)+analyseOtherOp ctx _ = (ctx, mempty)++-- | Returns an intmap of names, to be used as dependencies in construction of+-- VariableInfos.+-- Throws an error if SubExp contains a name not in context. This behaviour+-- might be thrown out in the future, as it is mostly just a very verbose way to+-- ensure that we capture all necessary variables in the context at the moment+-- of development.+analyseSubExpr :: [VName] -> Context rep -> SubExp -> Names+analyseSubExpr _ _ (Constant _) = mempty+analyseSubExpr pp ctx (Var v) =+  case M.lookup v (assignments ctx) of+    (Just _) -> oneName v+    Nothing ->+      error $+        "Failed to lookup variable \""+          ++ prettyString v+          ++ "\npat: "+          ++ prettyString pp+          ++ "\n\nContext\n"+          ++ show ctx++-- | Reduce a DimFix into its set of dependencies+consolidate :: Context rep -> SubExp -> DimAccess rep+consolidate _ (Constant _) = mempty+consolidate ctx (Var v) = DimAccess (reduceDependencies ctx v) (Just v)++-- | Recursively lookup vnames until vars with no deps are reached.+reduceDependencies :: Context rep -> VName -> M.Map VName Dependency+reduceDependencies ctx v =+  case M.lookup v (assignments ctx) of+    Nothing -> error $ "Unable to find " ++ prettyString v+    Just (VariableInfo deps lvl _parents t) ->+      -- We detect whether it is a threadID or loop counter by checking+      -- whether or not it has any dependencies+      case t of+        ThreadID -> M.fromList [(v, Dependency lvl t)]+        LoopVar -> M.fromList [(v, Dependency lvl t)]+        Variable -> mconcat $ map (reduceDependencies ctx) $ namesToList deps+        ConstType -> mempty++-- Misc functions++-- Instances for AST types that we actually support+instance Analyse GPU where+  analyseOp gpu_op+    | (SegOp op) <- gpu_op = analyseSegOp op+    | (SizeOp op) <- gpu_op = analyseSizeOp op+    | (GPUBody _ body) <- gpu_op = pure . analyseGPUBody body+    | (Futhark.IR.GPU.OtherOp _) <- gpu_op = analyseOtherOp++instance Analyse MC where+  analyseOp mc_op+    | ParOp Nothing seq_segop <- mc_op = analyseSegOp seq_segop+    | ParOp (Just segop) seq_segop <- mc_op = \ctx name -> do+        let (ctx', res') = analyseSegOp segop ctx name+        let (ctx'', res'') = analyseSegOp seq_segop ctx name+        (ctx' <> ctx'', unionIndexTables res' res'')+    | Futhark.IR.MC.OtherOp _ <- mc_op = analyseOtherOp++-- Unfortunately we need these instances, even though they may never appear.+instance Analyse GPUMem where+  analyseOp _ = error $ notImplementedYet "GPUMem"++instance Analyse MCMem where+  analyseOp _ = error "Unexpected?"++instance Analyse Seq where+  analyseOp _ = error $ notImplementedYet "Seq"++instance Analyse SeqMem where+  analyseOp _ = error $ notImplementedYet "SeqMem"++instance Analyse SOACS where+  analyseOp _ = error $ notImplementedYet "SOACS"++notImplementedYet :: String -> String+notImplementedYet s = "Access pattern analysis for the " ++ s ++ " backend is not implemented."++instance Pretty (IndexTable rep) where+  pretty = stack . map f . M.toList :: IndexTable rep -> Doc ann+    where+      f (segop, arrNameToIdxExprMap) = pretty segop <+> colon <+> g arrNameToIdxExprMap++      g maps = lbrace </> indent 4 (mapprintArray $ M.toList maps) </> rbrace++      mapprintArray :: [(ArrayName, M.Map IndexExprName [DimAccess rep])] -> Doc ann+      mapprintArray [] = ""+      mapprintArray [m] = printArrayMap m+      mapprintArray (m : mm) = printArrayMap m </> mapprintArray mm++      printArrayMap :: (ArrayName, M.Map IndexExprName [DimAccess rep]) -> Doc ann+      printArrayMap ((name, _, layout), maps) =+        "(arr)"+          <+> pretty name+          <+> colon+          <+> pretty layout+          <+> lbrace+          </> indent 4 (mapprintIdxExpr (M.toList maps))+          </> rbrace++      mapprintIdxExpr :: [(IndexExprName, [DimAccess rep])] -> Doc ann+      mapprintIdxExpr [] = ""+      mapprintIdxExpr [m] = printIdxExpMap m+      mapprintIdxExpr (m : mm) = printIdxExpMap m </> mapprintIdxExpr mm++      printIdxExpMap (name, mems) = "(idx)" <+> pretty name <+> ":" </> indent 4 (printDimAccess mems)++      printDimAccess :: [DimAccess rep] -> Doc ann+      printDimAccess dim_accesses = stack $ zipWith (curry printDim) [0 ..] dim_accesses++      printDim :: (Int, DimAccess rep) -> Doc ann+      printDim (i, m) = pretty i <+> ":" <+> indent 0 (pretty m)++instance Pretty (DimAccess rep) where+  pretty dim_access =+    -- Instead of using `brackets $` we manually enclose with `[`s, to add+    -- spacing between the enclosed elements+    if case originalVar dim_access of+      Nothing -> True+      Just n ->+        length (dependencies dim_access) == 1 && n == head (map fst $ M.toList $ dependencies dim_access)+        -- Only print the original name if it is different from the first (and single) dependency+      then+        "dependencies"+          <+> equals+          <+> align (prettyDeps $ dependencies dim_access)+      else+        "dependencies"+          <+> equals+          <+> pretty (originalVar dim_access)+          <+> "->"+          <+> align (prettyDeps $ dependencies dim_access)+    where+      prettyDeps = braces . commasep . map printPair . M.toList+      printPair (name, Dependency lvl vtype) = pretty name <+> pretty lvl <+> pretty vtype++instance Pretty SegOpName where+  pretty (SegmentedMap name) = "(segmap)" <+> pretty name+  pretty (SegmentedRed name) = "(segred)" <+> pretty name+  pretty (SegmentedScan name) = "(segscan)" <+> pretty name+  pretty (SegmentedHist name) = "(seghist)" <+> pretty name++instance Pretty BodyType where+  pretty (SegOpName (SegmentedMap name)) = pretty name <+> colon <+> "segmap"+  pretty (SegOpName (SegmentedRed name)) = pretty name <+> colon <+> "segred"+  pretty (SegOpName (SegmentedScan name)) = pretty name <+> colon <+> "segscan"+  pretty (SegOpName (SegmentedHist name)) = pretty name <+> colon <+> "seghist"+  pretty (LoopBodyName name) = pretty name <+> colon <+> "loop"+  pretty (CondBodyName name) = pretty name <+> colon <+> "cond"++instance Pretty VarType where+  pretty ConstType = "const"+  pretty Variable = "var"+  pretty ThreadID = "tid"+  pretty LoopVar = "iter"

@@ -0,0 +1,169 @@+-- | Compute a mapping from variables to their corresponding (fully+-- expanded) PrimExps.+module Futhark.Analysis.PrimExp.Table+  ( primExpTable,+    PrimExpTable,++    -- * Extensibility+    PrimExpAnalysis (..),++    -- * Testing+    stmToPrimExps,+  )+where++import Control.Monad.State.Strict+import Data.Foldable+import Data.Map.Strict qualified as M+import Futhark.Analysis.PrimExp+import Futhark.Analysis.PrimExp.Convert+import Futhark.IR.Aliases+import Futhark.IR.GPU+import Futhark.IR.GPUMem+import Futhark.IR.MC+import Futhark.IR.MCMem++-- | Maps variables to maybe PrimExps. Will map to nothing if it+-- cannot be resolved to a PrimExp. For all uses of this analysis atm.+-- a variable can be considered inscrutable if it cannot be resolved+-- to a primexp.+type PrimExpTable = M.Map VName (Maybe (PrimExp VName))++-- | A class for extracting PrimExps from what is inside an op.+class PrimExpAnalysis rep where+  opPrimExp :: Scope rep -> Op rep -> State PrimExpTable ()++primExpTable :: (PrimExpAnalysis rep, RepTypes rep) => Prog rep -> PrimExpTable+primExpTable prog = initialState <> foldMap' (uncurry funToPrimExp) scopesAndFuns+  where+    scopesAndFuns = do+      let fun_defs = progFuns prog+      let scopes = map getScope fun_defs+      zip scopes fun_defs++    getScope funDef = scopeOf (progConsts prog) <> scopeOfFParams (funDefParams funDef)++    -- We need to have the dummy "slice" in the analysis for our "slice hack".+    initialState =+      M.singleton (VName "slice" 0) $ Just $ LeafExp (VName "slice" 0) $ IntType Int64++funToPrimExp ::+  (PrimExpAnalysis rep, RepTypes rep) =>+  Scope rep ->+  FunDef rep ->+  PrimExpTable+funToPrimExp scope fundef = execState (bodyToPrimExps scope (funDefBody fundef)) mempty++-- | Adds the statements of a body to the PrimExpTable+bodyToPrimExps ::+  (PrimExpAnalysis rep, RepTypes rep) =>+  Scope rep ->+  Body rep ->+  State PrimExpTable ()+bodyToPrimExps scope body = mapM_ (stmToPrimExps scope') (bodyStms body)+  where+    scope' = scope <> scopeOf (bodyStms body)++-- | Adds the statements of a kernel body to the PrimExpTable+kernelToBodyPrimExps ::+  (PrimExpAnalysis rep, RepTypes rep) =>+  Scope rep ->+  KernelBody rep ->+  State PrimExpTable ()+kernelToBodyPrimExps scope kbody = mapM_ (stmToPrimExps scope') (kernelBodyStms kbody)+  where+    scope' = scope <> scopeOf (kernelBodyStms kbody)++-- | Adds a statement to the PrimExpTable. If it can't be resolved as a `PrimExp`,+-- it will be added as `Nothing`.+stmToPrimExps ::+  forall rep.+  (PrimExpAnalysis rep, RepTypes rep) =>+  Scope rep ->+  Stm rep ->+  State PrimExpTable ()+stmToPrimExps scope stm = do+  table <- get+  case stm of+    (Let (Pat pat_elems) _ e)+      | Just primExp <- primExpFromExp (toPrimExp scope table) e ->+          -- The statement can be resolved as a `PrimExp`.+          -- For each pattern element, insert the PrimExp in the table+          forM_ pat_elems $ \pe ->+            modify $ M.insert (patElemName pe) (Just primExp)+      | otherwise -> do+          -- The statement can't be resolved as a `PrimExp`.+          walk $ stmExp stm -- Traverse the rest of the AST Get the+          -- updated PrimExpTable after traversing the AST+          table' <- get++          -- Add pattern elements that can't be resolved as `PrimExp`+          -- to the `PrimExpTable` as `Nothing`+          forM_ pat_elems $ \pe ->+            case M.lookup (patElemName pe) table' of+              Nothing -> modify $ M.insert (patElemName pe) Nothing+              Just _ -> pure ()+  where+    walk e = do+      -- Handle most cases using the walker+      walkExpM walker e+      -- Additionally, handle loop parameters+      case e of+        Loop _ (ForLoop i t _) _ ->+          modify $ M.insert i $ Just $ LeafExp i $ IntType t+        _ -> pure ()++    walker =+      (identityWalker @rep)+        { walkOnBody = \body_scope -> bodyToPrimExps (scope <> body_scope),+          walkOnOp = opPrimExp scope,+          walkOnFParam = paramToPrimExp -- Loop parameters+        }++    -- Adds a loop parameter to the PrimExpTable+    paramToPrimExp :: FParam rep -> State PrimExpTable ()+    paramToPrimExp param = do+      let name = paramName param+      -- Construct a `PrimExp` from the type of the parameter+      -- and add it to the `PrimExpTable`+      case typeOf $ paramDec param of+        -- TODO: Handle other types?+        Prim pt ->+          modify $ M.insert name (Just $ LeafExp name pt)+        _ -> pure ()++-- | Checks if a name is in the PrimExpTable and construct a `PrimExp`+-- if it is not+toPrimExp :: (RepTypes rep) => Scope rep -> PrimExpTable -> VName -> Maybe (PrimExp VName)+toPrimExp scope table name = case M.lookup name table of+  Just maybePrimExp+    | Just primExp <- maybePrimExp -> Just primExp -- Already in the table+  _ -> case fmap typeOf . M.lookup name $ scope of+    (Just (Prim pt)) -> Just $ LeafExp name pt+    _ -> Nothing++-- | Adds the parameters of a SegOp as well as the statements in its+-- body to the PrimExpTable+segOpToPrimExps :: (PrimExpAnalysis rep, RepTypes rep) => Scope rep -> SegOp lvl rep -> State PrimExpTable ()+segOpToPrimExps scope op = do+  forM_ (map fst $ unSegSpace $ segSpace op) $ \name ->+    modify $ M.insert name $ Just $ LeafExp name int64+  kernelToBodyPrimExps scope (segBody op)++instance PrimExpAnalysis GPU where+  opPrimExp scope gpu_op+    | (SegOp op) <- gpu_op = segOpToPrimExps scope op+    | (SizeOp _) <- gpu_op = pure ()+    | (GPUBody _ body) <- gpu_op = bodyToPrimExps scope body+    | (Futhark.IR.GPUMem.OtherOp _) <- gpu_op = pure ()++instance PrimExpAnalysis MC where+  opPrimExp scope mc_op+    | (ParOp maybe_par_segop seq_segop) <- mc_op = do+        -- Add the statements in the parallel part of the ParOp to the PrimExpTable+        case maybe_par_segop of+          Nothing -> pure ()+          Just _ -> forM_ maybe_par_segop $ segOpToPrimExps scope+        -- Add the statements in the sequential part of the ParOp to the PrimExpTable+        segOpToPrimExps scope seq_segop+    | (Futhark.IR.MCMem.OtherOp _) <- mc_op = pure ()

@@ -10,6 +10,7 @@ import Data.Text qualified as T import Data.Text.IO qualified as T import Futhark.Actions+import Futhark.Analysis.AccessPattern (Analyse) import Futhark.Analysis.Alias qualified as Alias import Futhark.Analysis.Metrics (OpMetrics) import Futhark.Compiler.CLI hiding (compilerMain)@@ -31,6 +32,7 @@ import Futhark.Internalise.LiftLambdas as LiftLambdas import Futhark.Internalise.Monomorphise as Monomorphise import Futhark.Internalise.ReplaceRecords as ReplaceRecords+import Futhark.Optimise.ArrayLayout import Futhark.Optimise.ArrayShortCircuiting qualified as ArrayShortCircuiting import Futhark.Optimise.CSE import Futhark.Optimise.DoubleBuffer@@ -51,7 +53,6 @@ import Futhark.Pass.ExtractKernels import Futhark.Pass.ExtractMulticore import Futhark.Pass.FirstOrderTransform-import Futhark.Pass.KernelBabysitting import Futhark.Pass.LiftAllocations as LiftAllocations import Futhark.Pass.LowerAllocations as LowerAllocations import Futhark.Pass.Simplify@@ -159,7 +160,8 @@   | PolyAction       ( forall (rep :: Data.Kind.Type).         ( AliasableRep rep,-          (OpMetrics (Op rep))+          (OpMetrics (Op rep)),+          Analyse rep         ) =>         Action rep       )@@ -233,6 +235,13 @@   externalErrorS $     "Pass '" <> name <> "' expects SeqMem representation, but got " <> representation rep +mcProg :: String -> UntypedPassState -> FutharkM (Prog MC.MC)+mcProg _ (MC prog) =+  pure prog+mcProg name rep =+  externalErrorS $+    "Pass " ++ name ++ " expects MC representation, but got " ++ representation rep+ mcMemProg :: String -> UntypedPassState -> FutharkM (Prog MCMem.MCMem) mcMemProg _ (MCMem prog) =   pure prog@@ -267,6 +276,13 @@ kernelsPassOption =   typedPassOption kernelsProg GPU +mcPassOption ::+  Pass MC.MC MC.MC ->+  String ->+  FutharkOption+mcPassOption =+  typedPassOption mcProg MC+ seqMemPassOption ::   Pass SeqMem.SeqMem SeqMem.SeqMem ->   String ->@@ -352,6 +368,21 @@     long = [passLongOption pass]     pass = performCSE True :: Pass SOACS.SOACS SOACS.SOACS +sinkOption :: String -> FutharkOption+sinkOption short =+  passOption (passDescription pass) (UntypedPass perform) short long+  where+    perform (GPU prog) config =+      GPU <$> runPipeline (onePass sinkGPU) config prog+    perform (MC prog) config =+      MC <$> runPipeline (onePass sinkMC) config prog+    perform s _ =+      externalErrorS $+        "Pass '" ++ passDescription pass ++ "' cannot operate on " ++ representation s++    long = [passLongOption pass]+    pass = sinkGPU+ pipelineOption ::   (UntypedPassState -> Maybe (Prog fromrep)) ->   String ->@@ -383,6 +414,23 @@   FutharkOption soacsPipelineOption = pipelineOption getSOACSProg "SOACS" SOACS +unstreamOption :: String -> FutharkOption+unstreamOption short =+  passOption (passDescription pass) (UntypedPass perform) short long+  where+    perform (GPU prog) config =+      GPU+        <$> runPipeline (onePass unstreamGPU) config prog+    perform (MC prog) config =+      MC+        <$> runPipeline (onePass unstreamMC) config prog+    perform s _ =+      externalErrorS $+        "Pass '" ++ passDescription pass ++ "' cannot operate on " ++ representation s++    long = [passLongOption pass]+    pass = unstreamGPU+ commandLineOptions :: [FutharkOption] commandLineOptions =   [ Option@@ -512,6 +560,11 @@       )       "Print memory alias information.",     Option+      "z"+      ["memory-access-pattern"]+      (NoArg $ Right $ \opts -> opts {futharkAction = PolyAction printMemoryAccessAnalysis})+      "Print the result of analysing memory access patterns. Currently only for --gpu --mc.",+    Option       []       ["call-graph"]       (NoArg $ Right $ \opts -> opts {futharkAction = SOACSAction callGraphAction})@@ -591,11 +644,12 @@     soacsPassOption removeDeadFunctions [],     soacsPassOption applyAD [],     soacsPassOption applyADInnermost [],-    kernelsPassOption babysitKernels [],+    kernelsPassOption optimiseArrayLayoutGPU [],+    mcPassOption optimiseArrayLayoutMC [],     kernelsPassOption tileLoops [],     kernelsPassOption histAccsGPU [],-    kernelsPassOption unstreamGPU [],-    kernelsPassOption sinkGPU [],+    unstreamOption [],+    sinkOption [],     kernelsPassOption reduceDeviceSyncs [],     typedPassOption soacsProg GPU extractKernels [],     typedPassOption soacsProg MC extractMulticore [],

@@ -26,6 +26,7 @@ import System.IO import Prelude +-- | Run @futhark eval@. main :: String -> [String] -> IO () main = mainWithOptions interpreterConfig options "options... <exprs...>" run   where

@@ -673,6 +673,7 @@     scriptStopOnError :: Bool   } +-- | The configuration before any user-provided options are processed. initialOptions :: Options initialOptions =   Options@@ -1072,6 +1073,7 @@   cleanupImgDir env $ mconcat files   pure (foldl' min Success failures, T.intercalate "\n" outputs) +-- | Common command line options that transform 'Options'. scriptCommandLineOptions :: [FunOptDescr Options] scriptCommandLineOptions =   [ Option@@ -1139,6 +1141,9 @@            "Stop and do not produce output file if any directive fails."        ] +-- | Start up (and eventually shut down) a Futhark server+-- corresponding to the provided program. If the program has a @.fut@+-- extension, it will be compiled automatically. prepareServer :: FilePath -> Options -> (ScriptServer -> IO a) -> IO a prepareServer prog opts f = do   futhark <- maybe getExecutablePath pure $ scriptFuthark opts

@@ -1,3 +1,4 @@+-- | @futhark script@ module Futhark.CLI.Script (main) where  import Control.Monad.Except

@@ -18,6 +18,7 @@ import Data.Text qualified as T import Data.Text.Encoding qualified as T import Data.Text.IO qualified as T+import Data.Time.Clock.System (SystemTime (..), getSystemTime) import Futhark.Analysis.Metrics.Type import Futhark.Server import Futhark.Test@@ -113,6 +114,8 @@     Compiled   | -- | Test interpreted code.     Interpreted+  | -- | Perform structure tests.+    Structure   deriving (Eq, Show)  data TestCase = TestCase@@ -249,7 +252,7 @@           ]   case testAction testcase of     CompileTimeFailure expected_error ->-      unless (mode == Internalise) . context checkctx $ do+      unless (mode `elem` [Structure, Internalise]) . context checkctx $ do         (code, _, err) <-           liftIO $ readProcessWithExitCode futhark ["check", program] ""         case code of@@ -285,10 +288,12 @@           -- so force just one data set at a time here.           ensureReferenceOutput (Just 1) (FutharkExe futhark) "c" program ios +      when (mode == Structure) $+        mapM_ (testMetrics progs program) structures+       when (mode `elem` [Compile, Compiled]) $         context ("Compiling with --backend=" <> T.pack backend) $ do           compileTestProgram extra_compiler_options (FutharkExe futhark) backend program warnings-          mapM_ (testMetrics progs program) structures           unless (mode == Compile) $ do             (tuning_opts, _) <-               liftIO $ determineTuning (configTuning progs) program@@ -509,6 +514,26 @@     running = labelstr <> (T.unwords . reverse . map (T.pack . testCaseProgram) . testStatusRun) ts     labelstr = "Now testing: " +reportLine :: MVar SystemTime -> TestStatus -> IO ()+reportLine time_mvar ts =+  modifyMVar_ time_mvar $ \time -> do+    time_now <- getSystemTime+    if systemSeconds time_now - systemSeconds time >= period+      then do+        T.putStrLn $+          "("+            <> showText (testStatusFail ts)+            <> " failed, "+            <> showText (testStatusPass ts)+            <> " passed, "+            <> showText num_remain+            <> " to go)."+        pure time_now+      else pure time+  where+    num_remain = length $ testStatusRemain ts+    period = 60+ moveCursorToTableTop :: IO () moveCursorToTableTop = cursorUpLine tableLines @@ -531,11 +556,13 @@   let (excluded, included) = partition (excludedTest config) all_tests   _ <- forkIO $ mapM_ (putMVar testmvar . excludeCases config) included +  time_mvar <- newMVar $ MkSystemTime 0 0+   let fancy = not (configLineOutput config) && fancyTerminal        report         | fancy = reportTable-        | otherwise = const (pure ())+        | otherwise = reportLine time_mvar       clear         | fancy = clearFromCursorToScreenEnd         | otherwise = pure ()@@ -709,6 +736,11 @@       ["compile"]       (NoArg $ Right $ \config -> config {configTestMode = Compile})       "Only compile, do not run.",+    Option+      "s"+      ["structure"]+      (NoArg $ Right $ \config -> config {configTestMode = Structure})+      "Perform structure tests.",     Option       "I"       ["internalise"]

@@ -535,6 +535,19 @@        mapM_ earlyDecl $ concat memfuns       type_funs <- generateAPITypes arr_space types++      headerDecl InitDecl [C.cedecl|void futhark_context_config_set_debugging(struct futhark_context_config* cfg, int flag);|]+      headerDecl InitDecl [C.cedecl|void futhark_context_config_set_profiling(struct futhark_context_config* cfg, int flag);|]+      headerDecl InitDecl [C.cedecl|void futhark_context_config_set_logging(struct futhark_context_config* cfg, int flag);|]+      headerDecl MiscDecl [C.cedecl|void futhark_context_config_set_cache_file(struct futhark_context_config* cfg, const char *f);|]+      headerDecl InitDecl [C.cedecl|int futhark_get_tuning_param_count(void);|]+      headerDecl InitDecl [C.cedecl|const char* futhark_get_tuning_param_name(int);|]+      headerDecl InitDecl [C.cedecl|const char* futhark_get_tuning_param_class(int);|]+      headerDecl MiscDecl [C.cedecl|char* futhark_context_get_error(struct futhark_context* ctx);|]+      headerDecl MiscDecl [C.cedecl|void futhark_context_set_logging_file(struct futhark_context* ctx, typename FILE* f);|]+      headerDecl MiscDecl [C.cedecl|void futhark_context_pause_profiling(struct futhark_context* ctx);|]+      headerDecl MiscDecl [C.cedecl|void futhark_context_unpause_profiling(struct futhark_context* ctx);|]+       generateCommonLibFuns memreport        pure@@ -583,58 +596,8 @@ generateCommonLibFuns :: [C.BlockItem] -> CompilerM op s () generateCommonLibFuns memreport = do   ctx <- contextType-  cfg <- configType   ops <- asks envOperations -  publicDef_ "context_config_set_debugging" InitDecl $ \s ->-    ( [C.cedecl|void $id:s($ty:cfg* cfg, int flag);|],-      [C.cedecl|void $id:s($ty:cfg* cfg, int flag) {-                         cfg->profiling = cfg->logging = cfg->debugging = flag;-                       }|]-    )--  publicDef_ "context_config_set_profiling" InitDecl $ \s ->-    ( [C.cedecl|void $id:s($ty:cfg* cfg, int flag);|],-      [C.cedecl|void $id:s($ty:cfg* cfg, int flag) {-                         cfg->profiling = flag;-                       }|]-    )--  publicDef_ "context_config_set_logging" InitDecl $ \s ->-    ( [C.cedecl|void $id:s($ty:cfg* cfg, int flag);|],-      [C.cedecl|void $id:s($ty:cfg* cfg, int flag) {-                         cfg->logging = flag;-                       }|]-    )--  publicDef_ "context_config_set_cache_file" MiscDecl $ \s ->-    ( [C.cedecl|void $id:s($ty:cfg* cfg, const char *f);|],-      [C.cedecl|void $id:s($ty:cfg* cfg, const char *f) {-                 cfg->cache_fname = strdup(f);-               }|]-    )--  publicDef_ "get_tuning_param_count" InitDecl $ \s ->-    ( [C.cedecl|int $id:s(void);|],-      [C.cedecl|int $id:s(void) {-                return num_tuning_params;-              }|]-    )--  publicDef_ "get_tuning_param_name" InitDecl $ \s ->-    ( [C.cedecl|const char* $id:s(int);|],-      [C.cedecl|const char* $id:s(int i) {-                return tuning_param_names[i];-              }|]-    )--  publicDef_ "get_tuning_param_class" InitDecl $ \s ->-    ( [C.cedecl|const char* $id:s(int);|],-      [C.cedecl|const char* $id:s(int i) {-                return tuning_param_classes[i];-              }|]-    )-   sync <- publicName "context_sync"   let comma = [C.citem|str_builder_char(&builder, ',');|]   publicDef_ "context_report" MiscDecl $ \s ->@@ -657,36 +620,6 @@                    str_builder_str(&builder, "]}");                    return builder.str;                  }-               }|]-    )--  publicDef_ "context_get_error" MiscDecl $ \s ->-    ( [C.cedecl|char* $id:s($ty:ctx* ctx);|],-      [C.cedecl|char* $id:s($ty:ctx* ctx) {-                         char* error = ctx->error;-                         ctx->error = NULL;-                         return error;-                       }|]-    )--  publicDef_ "context_set_logging_file" MiscDecl $ \s ->-    ( [C.cedecl|void $id:s($ty:ctx* ctx, typename FILE* f);|],-      [C.cedecl|void $id:s($ty:ctx* ctx, typename FILE* f) {-                  ctx->log = f;-                }|]-    )--  publicDef_ "context_pause_profiling" MiscDecl $ \s ->-    ( [C.cedecl|void $id:s($ty:ctx* ctx);|],-      [C.cedecl|void $id:s($ty:ctx* ctx) {-                 ctx->profiling_paused = 1;-               }|]-    )--  publicDef_ "context_unpause_profiling" MiscDecl $ \s ->-    ( [C.cedecl|void $id:s($ty:ctx* ctx);|],-      [C.cedecl|void $id:s($ty:ctx* ctx) {-                 ctx->profiling_paused = 0;                }|]     ) 

@@ -972,7 +972,7 @@   pure () defCompileBasicOp _ Reshape {} =   pure ()-defCompileBasicOp _ (UpdateAcc acc is vs) = sComment "UpdateAcc" $ do+defCompileBasicOp _ (UpdateAcc safety acc is vs) = sComment "UpdateAcc" $ do   -- We are abusing the comment mechanism to wrap the operator in   -- braces when we end up generating code.  This is necessary because   -- we might otherwise end up declaring lambda parameters (if any)@@ -985,7 +985,11 @@   -- index parameters.   (_, _, arrs, dims, op) <- lookupAcc acc is' -  sWhen (inBounds (Slice (map DimFix is')) dims) $+  let boundsCheck =+        case safety of+          Safe -> sWhen (inBounds (Slice (map DimFix is')) dims)+          _ -> id+  boundsCheck $     case op of       Nothing ->         -- Scatter-like.

@@ -130,12 +130,16 @@ threadAlloc dest _ _ =   error $ "Invalid target for in-kernel allocation: " ++ show dest -updateAcc :: VName -> [SubExp] -> [SubExp] -> InKernelGen ()-updateAcc acc is vs = sComment "UpdateAcc" $ do+updateAcc :: Safety -> VName -> [SubExp] -> [SubExp] -> InKernelGen ()+updateAcc safety acc is vs = sComment "UpdateAcc" $ do   -- See the ImpGen implementation of UpdateAcc for general notes.   let is' = map pe64 is   (c, space, arrs, dims, op) <- lookupAcc acc is'-  sWhen (inBounds (Slice (map DimFix is')) dims) $+  let boundsCheck =+        case safety of+          Safe -> sWhen (inBounds (Slice (map DimFix is')) dims)+          _ -> id+  boundsCheck $     case op of       Nothing ->         forM_ (zip arrs vs) $ \(arr, v) -> copyDWIMFix arr is' v []@@ -176,8 +180,8 @@ compileThreadExp (Pat [dest]) (BasicOp (ArrayLit es _)) =   forM_ (zip [0 ..] es) $ \(i, e) ->     copyDWIMFix (patElemName dest) [fromIntegral (i :: Int64)] e []-compileThreadExp _ (BasicOp (UpdateAcc acc is vs)) =-  updateAcc acc is vs+compileThreadExp _ (BasicOp (UpdateAcc safety acc is vs)) =+  updateAcc safety acc is vs compileThreadExp dest e =   defCompileExp dest e 

@@ -281,9 +281,9 @@ compileBlockExp (Pat [dest]) (BasicOp (ArrayLit es _)) =   forM_ (zip [0 ..] es) $ \(i, e) ->     copyDWIMFix (patElemName dest) [fromIntegral (i :: Int64)] e []-compileBlockExp _ (BasicOp (UpdateAcc acc is vs)) = do+compileBlockExp _ (BasicOp (UpdateAcc safety acc is vs)) = do   ltid <- kernelLocalThreadId . kernelConstants <$> askEnv-  sWhen (ltid .==. 0) $ updateAcc acc is vs+  sWhen (ltid .==. 0) $ updateAcc safety acc is vs   sOp $ Imp.Barrier Imp.FenceLocal compileBlockExp (Pat [dest]) (BasicOp (Replicate ds se)) | ds /= mempty = do   flat <- newVName "rep_flat"

@@ -59,12 +59,16 @@       opsCopyCompiler = parallelCopy     } -updateAcc :: VName -> [SubExp] -> [SubExp] -> MulticoreGen ()-updateAcc acc is vs = sComment "UpdateAcc" $ do+updateAcc :: Safety -> VName -> [SubExp] -> [SubExp] -> MulticoreGen ()+updateAcc safety acc is vs = sComment "UpdateAcc" $ do   -- See the ImpGen implementation of UpdateAcc for general notes.   let is' = map pe64 is   (c, _space, arrs, dims, op) <- lookupAcc acc is'-  sWhen (inBounds (Slice (map DimFix is')) dims) $+  let boundsCheck =+        case safety of+          Safe -> sWhen (inBounds (Slice (map DimFix is')) dims)+          _ -> id+  boundsCheck $     case op of       Nothing ->         forM_ (zip arrs vs) $ \(arr, v) -> copyDWIMFix arr is' v []@@ -113,8 +117,8 @@           locksForInputs atomics inputs'  compileMCExp :: ExpCompiler MCMem HostEnv Imp.Multicore-compileMCExp _ (BasicOp (UpdateAcc acc is vs)) =-  updateAcc acc is vs+compileMCExp _ (BasicOp (UpdateAcc safety acc is vs)) =+  updateAcc safety acc is vs compileMCExp pat (WithAcc inputs lam) =   withAcc pat inputs lam compileMCExp dest e =

@@ -317,9 +317,10 @@       ArrayLit         <$> brackets (pSubExp `sepBy` pComma)         <*> (lexeme ":" *> "[]" *> pType),-      keyword "update_acc"-        *> parens-          (UpdateAcc <$> pVName <* pComma <*> pSubExps <* pComma <*> pSubExps),+      do+        safety <-+          choice [keyword "update_acc_unsafe" $> Unsafe, keyword "update_acc" $> Safe]+        parens (UpdateAcc safety <$> pVName <* pComma <*> pSubExps <* pComma <*> pSubExps),       --       pConvOp "sext" SExt pIntType pIntType,       pConvOp "zext" ZExt pIntType pIntType,

@@ -234,13 +234,17 @@   pretty (Manifest perm e) = "manifest" <> apply [apply (map pretty perm), pretty e]   pretty (Assert e msg (loc, _)) =     "assert" <> apply [pretty e, pretty msg, pretty $ show $ locStr loc]-  pretty (UpdateAcc acc is v) =-    "update_acc"+  pretty (UpdateAcc safety acc is v) =+    update_acc_str       <> apply         [ pretty acc,           ppTuple' $ map pretty is,           ppTuple' $ map pretty v         ]+    where+      update_acc_str = case safety of+        Safe -> "update_acc"+        Unsafe -> "update_acc_unsafe"  instance (Pretty a) => Pretty (ErrorMsg a) where   pretty (ErrorMsg parts) = braces $ align $ commasep $ map p parts

@@ -180,7 +180,7 @@     inputConsumed (_, arrs, _) = namesFromList arrs consumedInExp (BasicOp (Update _ src _ _)) = oneName src consumedInExp (BasicOp (FlatUpdate src _ _)) = oneName src-consumedInExp (BasicOp (UpdateAcc acc _ _)) = oneName acc+consumedInExp (BasicOp (UpdateAcc _ acc _ _)) = oneName acc consumedInExp (BasicOp _) = mempty consumedInExp (Op op) = consumedInOp op 

@@ -117,7 +117,7 @@   pure <$> lookupType v basicOpType Assert {} =   pure [Prim Unit]-basicOpType (UpdateAcc v _ _) =+basicOpType (UpdateAcc _ v _ _) =   pure <$> lookupType v  -- | The type of an expression.

@@ -60,6 +60,8 @@     intersperse,     isPrefixOf,     partition,+    unzip4,+    zip4,   ) import Data.Map.Strict qualified as M import Data.Maybe@@ -1331,6 +1333,20 @@   Rule rep -- Some SegOp results can be moved outside the SegOp, which can -- simplify further analysis.+bottomUpSegOp (_vtable, used) (Pat kpes) dec segop+  -- Remove dead results. This is a bit tricky to do with scan/red+  -- results, so we only deal with map results for now.+  | (_, kpes', kts', kres') <- unzip4 $ filter keep $ zip4 [0 ..] kpes kts kres,+    kpes' /= kpes = Simplify $ do+      kbody' <- localScope (scopeOfSegSpace space) $ mkKernelBodyM kstms kres'+      addStm $ Let (Pat kpes') dec $ Op $ segOp $ mk_segop kts' kbody'+  where+    space = segSpace segop+    (kts, KernelBody _ kstms kres, num_nonmap_results, mk_segop) =+      segOpGuts segop++    keep (i, pe, _, _) =+      i < num_nonmap_results || patElemName pe `UT.used` used bottomUpSegOp (vtable, _used) (Pat kpes) dec segop = Simplify $ do   -- Iterate through the bindings.  For each, we check whether it is   -- in kres and can be moved outside.  If so, we remove it from kres

@@ -378,9 +378,11 @@     -- must be a permutation of @[0,n-1]@, where @n@ is the     -- number of dimensions in the input array.     Rearrange [Int] VName-  | -- | Update an accumulator at the given index with the given value.-    -- Consumes the accumulator and produces a new one.-    UpdateAcc VName [SubExp] [SubExp]+  | -- | Update an accumulator at the given index with the given+    -- value. Consumes the accumulator and produces a new one. If+    -- 'Safe', perform a run-time bounds check and ignore the write if+    -- out of bounds (like @Scatter@).+    UpdateAcc Safety VName [SubExp] [SubExp]   deriving (Eq, Ord, Show)  -- | The input to a 'WithAcc' construct.  Comprises the index space of

@@ -166,9 +166,9 @@   BasicOp <$> (Assert <$> mapOnSubExp tv e <*> traverse (mapOnSubExp tv) msg <*> pure loc) mapExpM tv (BasicOp (Opaque op e)) =   BasicOp <$> (Opaque op <$> mapOnSubExp tv e)-mapExpM tv (BasicOp (UpdateAcc v is ses)) =+mapExpM tv (BasicOp (UpdateAcc safety v is ses)) =   BasicOp-    <$> ( UpdateAcc+    <$> ( UpdateAcc safety             <$> mapOnVName tv v             <*> mapM (mapOnSubExp tv) is             <*> mapM (mapOnSubExp tv) ses@@ -327,7 +327,7 @@   walkOnSubExp tv e >> traverse_ (walkOnSubExp tv) msg walkExpM tv (BasicOp (Opaque _ e)) =   walkOnSubExp tv e-walkExpM tv (BasicOp (UpdateAcc v is ses)) = do+walkExpM tv (BasicOp (UpdateAcc _ v is ses)) = do   walkOnVName tv v   mapM_ (walkOnSubExp tv) is   mapM_ (walkOnSubExp tv) ses

@@ -933,7 +933,7 @@   where     checkPart ErrorString {} = pure ()     checkPart (ErrorVal t x) = require [Prim t] x-checkBasicOp (UpdateAcc acc is ses) = do+checkBasicOp (UpdateAcc _ acc is ses) = do   (shape, ts) <- checkAccIdent acc    unless (length ses == length ts) . bad . TypeError $

@@ -1685,7 +1685,7 @@       acc' <- head <$> internaliseExpToVars "acc" acc       i' <- internaliseExp1 "acc_i" i       vs <- internaliseExp "acc_v" v-      fmap pure $ letSubExp desc $ BasicOp $ UpdateAcc acc' [i'] vs+      fmap pure $ letSubExp desc $ BasicOp $ UpdateAcc Safe acc' [i'] vs     handleAccs _ _ = Nothing      handleAD [f, x, v] fname

@@ -361,5 +361,6 @@   body' <- transformBody $ valBindBody valbind   pure $ valbind {valBindBody = body'} +-- | Fully normalise top level bindings. transformProg :: (MonadFreshNames m) => [ValBind] -> m [ValBind] transformProg = mapM transformValBind

@@ -225,7 +225,8 @@ parametrizing :: S.Set VName -> MonoM ExpReplacements parametrizing argset = do   intros <- askIntros argset-  (params, nxtBind) <- gets $ partition (not . S.disjoint intros . fvVars . freeInExp . unReplaced . fst)+  let usesIntros = not . S.disjoint intros . fvVars . freeInExp+  (params, nxtBind) <- gets $ partition (usesIntros . unReplaced . fst)   put nxtBind   pure params @@ -284,22 +285,16 @@ -- Given instantiated type of function, produce size arguments. type InferSizeArgs = StructType -> MonoM [Exp] +-- | The integer encodes an equivalence class, so we can keep+-- track of sizes that are statically identical. data MonoSize-  = -- | The integer encodes an equivalence class, so we can keep-    -- track of sizes that are statically identical.-    MonoKnown Int-  | MonoAnon-  deriving (Show)---- We treat all MonoAnon as identical.-instance Eq MonoSize where-  MonoKnown x == MonoKnown y = x == y-  MonoAnon == MonoAnon = True-  _ == _ = False+  = MonoKnown Int+  | MonoAnon Int+  deriving (Eq, Show)  instance Pretty MonoSize where   pretty (MonoKnown i) = "?" <> pretty i-  pretty MonoAnon = "?"+  pretty (MonoAnon i) = "??" <> pretty i  instance Pretty (Shape MonoSize) where   pretty (Shape ds) = mconcat (map (brackets . pretty) ds)@@ -321,10 +316,16 @@     noExtsScalar (Arrow as p d t1 (RetType _ t2)) =       Arrow as p d (noExts t1) (RetType [] (noExts t2))     noExtsScalar t = t-    onDim bound _ e+    onDim bound _ d       -- A locally bound size.-      | any (`S.member` bound) $ fvVars $ freeInExp e =-          pure MonoAnon+      | any (`S.member` bound) $ fvVars $ freeInExp d = do+          (i, m) <- get+          case M.lookup d m of+            Just prev ->+              pure $ MonoAnon prev+            Nothing -> do+              put (i + 1, M.insert d i m)+              pure $ MonoAnon i     onDim _ _ d = do       (i, m) <- get       case M.lookup d m of@@ -351,6 +352,9 @@ lookupLifted :: VName -> MonoType -> MonoM (Maybe (VName, InferSizeArgs)) lookupLifted fname t = lookup (fname, t) <$> getLifts +sizeVarName :: Exp -> String+sizeVarName e = "d<{" <> prettyString (bareExp e) <> "}>"+ -- | Creates a new expression replacement if needed, this always produces normalised sizes. -- (e.g. single variable or constant) replaceExp :: Exp -> MonoM Exp@@ -365,7 +369,7 @@         (Just vn, _) -> pure $ sizeFromName (qualName vn) (srclocOf e)         (Nothing, Just vn) -> pure $ sizeFromName (qualName vn) (srclocOf e)         (Nothing, Nothing) -> do-          vn <- newNameFromString $ "d<{" ++ prettyString (bareExp e) ++ "}>"+          vn <- newNameFromString $ sizeVarName e           modify ((e', vn) :)           pure $ sizeFromName (qualName vn) (srclocOf e)   where@@ -433,15 +437,15 @@       Record <$> traverse transformType fs     transformScalarSizes (Sum cs) =       Sum <$> (traverse . traverse) transformType cs-    transformScalarSizes (Arrow as argName d argT retT) = do-      retT' <- transformRetTypeSizes argset retT-      Arrow as argName d <$> transformType argT <*> pure retT'+    transformScalarSizes (Arrow as argName d argT retT) =+      Arrow as argName d+        <$> transformType argT+        <*> transformRetTypeSizes argset retT       where         argset =-          fvVars (freeInType argT)-            <> case argName of-              Unnamed -> mempty-              Named vn -> S.singleton vn+          case argName of+            Unnamed -> mempty+            Named vn -> S.singleton vn     transformScalarSizes (TypeVar u qn args) =       TypeVar u qn <$> mapM onArg args       where@@ -980,68 +984,71 @@   PolyBinding ->   MonoType ->   MonoM (VName, InferSizeArgs, ValBind)-monomorphiseBinding entry (PolyBinding (name, tparams, params, rettype, body, attrs, loc)) inst_t = do-  isolateNormalisation $ do-    let bind_t = funType params rettype-    (substs, t_shape_params) <--      typeSubstsM loc (noSizes bind_t) $ noNamedParams inst_t-    let shape_names = S.fromList $ map typeParamName $ shape_params ++ t_shape_params-        substs' = M.map (Subst []) substs-        substStructType =-          substTypesAny (fmap (fmap (second (const mempty))) . (`M.lookup` substs'))-        params' = map (substPat substStructType) params-    params'' <- withArgs shape_names $ mapM transformPat params'-    exp_naming <- paramGetClean+monomorphiseBinding entry (PolyBinding (name, tparams, params, rettype, body, attrs, loc)) inst_t = isolateNormalisation $ do+  let bind_t = funType params rettype+  (substs, t_shape_params) <-+    typeSubstsM loc bind_t $ noNamedParams inst_t+  let shape_names = S.fromList $ map typeParamName $ shape_params ++ t_shape_params+      substs' = M.map (Subst []) substs+      substStructType =+        substTypesAny (fmap (fmap (second (const mempty))) . (`M.lookup` substs'))+      params' = map (substPat substStructType) params+  params'' <- withArgs shape_names $ mapM transformPat params'+  exp_naming <- paramGetClean -    let args = S.fromList $ foldMap patNames params-        arg_params = map snd exp_naming+  let args = S.fromList $ foldMap patNames params+      arg_params = map snd exp_naming -    rettype' <--      withParams exp_naming $-        withArgs (args <> shape_names) $-          hardTransformRetType (applySubst (`M.lookup` substs') rettype)-    extNaming <- paramGetClean-    scope <- S.union shape_names <$> askScope'-    let (rettype'', new_params) = arrowArg scope args arg_params rettype'-        bind_t' = substTypesAny (`M.lookup` substs') bind_t-        (shape_params_explicit, shape_params_implicit) =-          partition ((`S.member` (mustBeExplicitInBinding bind_t'' `S.union` mustBeExplicitInBinding bind_t')) . typeParamName) $-            shape_params ++ t_shape_params ++ map (`TypeParamDim` mempty) (S.toList new_params)-        exp_naming' = filter ((`S.member` new_params) . snd) (extNaming <> exp_naming)+  rettype' <-+    withParams exp_naming $+      withArgs (args <> shape_names) $+        hardTransformRetType (applySubst (`M.lookup` substs') rettype)+  extNaming <- paramGetClean+  scope <- S.union shape_names <$> askScope'+  let (rettype'', new_params) = arrowArg scope args arg_params rettype'+      bind_t' = substTypesAny (`M.lookup` substs') bind_t+      mkExplicit =+        flip+          S.member+          (mustBeExplicitInBinding bind_t'' <> mustBeExplicitInBinding bind_t')+      (shape_params_explicit, shape_params_implicit) =+        partition (mkExplicit . typeParamName) $+          shape_params ++ t_shape_params ++ map (`TypeParamDim` mempty) (S.toList new_params)+      exp_naming' = filter ((`S.member` new_params) . snd) (extNaming <> exp_naming) -        bind_t'' = funType params'' rettype''-        bind_r = exp_naming <> extNaming-    body' <- updateExpTypes (`M.lookup` substs') body-    body'' <- withParams exp_naming' $ withArgs (shape_names <> args) $ transformExp body'-    scope' <- S.union (shape_names <> args) <$> askScope'-    body''' <--      expReplace exp_naming' <$> (calculateDims body'' . canCalculate scope' =<< get)+      bind_t'' = funType params'' rettype''+      bind_r = exp_naming <> extNaming+  body' <- updateExpTypes (`M.lookup` substs') body+  body'' <- withParams exp_naming' $ withArgs (shape_names <> args) $ transformExp body'+  scope' <- S.union (shape_names <> args) <$> askScope'+  body''' <-+    expReplace exp_naming' <$> (calculateDims body'' . canCalculate scope' =<< get) -    seen_before <- elem name . map (fst . fst) <$> getLifts-    name' <--      if null tparams && not entry && not seen_before-        then pure name-        else newName name+  seen_before <- elem name . map (fst . fst) <$> getLifts+  name' <-+    if null tparams && not entry && not seen_before+      then pure name+      else newName name -    pure-      ( name',-        inferSizeArgs shape_params_explicit bind_t'' bind_r,-        if entry-          then-            toValBinding-              name'-              (shape_params_explicit ++ shape_params_implicit)-              params''-              rettype''-              (entryAssert exp_naming body''')-          else-            toValBinding-              name'-              shape_params_implicit-              (map shapeParam shape_params_explicit ++ params'')-              rettype''-              body'''-      )+  pure+    ( name',+      inferSizeArgs shape_params_explicit bind_t'' bind_r,+      if entry+        then+          toValBinding+            name'+            (shape_params_explicit ++ shape_params_implicit)+            params''+            rettype''+            (entryAssert exp_naming body''')+        else+          toValBinding+            name'+            shape_params_implicit+            (map shapeParam shape_params_explicit ++ params'')+            rettype''+            body'''+    )   where     askScope' = S.filter (`notElem` retDims rettype) <$> askScope @@ -1088,7 +1095,7 @@ typeSubstsM ::   (MonadFreshNames m) =>   SrcLoc ->-  TypeBase () NoUniqueness ->+  StructType ->   MonoType ->   m (M.Map VName StructRetType, [TypeParam]) typeSubstsM loc orig_t1 orig_t2 =@@ -1100,10 +1107,13 @@     subRet t1 (RetType _ t2) =       sub t1 t2 -    sub t1@Array {} t2@Array {}-      | Just t1' <- peelArray (arrayRank t1) t1,-        Just t2' <- peelArray (arrayRank t1) t2 =-          sub t1' t2'+    sub t1@(Array _ (Shape (d1 : _)) _) t2@(Array _ (Shape (d2 : _)) _) = do+      case d2 of+        MonoAnon i -> do+          (ts, sizes) <- get+          put (ts, M.insert i d1 sizes)+        _ -> pure ()+      sub (stripArray 1 t1) (stripArray 1 t2)     sub (Scalar (TypeVar _ v _)) t =       unless (baseTag (qualLeaf v) <= maxIntrinsicTag) $         addSubst v $@@ -1136,11 +1146,15 @@         Nothing -> do           d <- lift $ lift $ newVName "d"           tell [TypeParamDim d loc]-          put (ts, M.insert i d sizes)+          put (ts, M.insert i (sizeFromName (qualName d) mempty) sizes)           pure $ sizeFromName (qualName d) mempty         Just d ->-          pure $ sizeFromName (qualName d) mempty-    onDim MonoAnon = pure anySize+          pure d+    onDim (MonoAnon i) = do+      (_, sizes) <- get+      case M.lookup i sizes of+        Nothing -> pure anySize+        Just d -> pure d  -- Perform a given substitution on the types in a pattern. substPat :: (t -> t) -> Pat t -> Pat t

@@ -299,7 +299,8 @@       where         size = sum . map length         f (ts', js, new_ts) t-          | Just (_, j) <- find ((== fmap fromDecl t) . fst) ts' =+          | all primType t,+            Just (_, j) <- find ((== fmap fromDecl t) . fst) ts' =               ( delete (fmap fromDecl t, j) ts',                 js ++ take (length t) [j ..],                 new_ts

@@ -73,7 +73,7 @@ modifyNameSource m = do   src <- getNameSource   let (x, src') = m src-  putNameSource src'+  src' `seq` putNameSource src'   pure x  -- | Produce a fresh name, using the given name as a template.

@@ -0,0 +1,42 @@+module Futhark.Optimise.ArrayLayout+  ( optimiseArrayLayoutGPU,+    optimiseArrayLayoutMC,+  )+where++import Control.Monad.State.Strict+import Futhark.Analysis.AccessPattern (Analyse, analyseDimAccesses)+import Futhark.Analysis.PrimExp.Table (primExpTable)+import Futhark.Builder+import Futhark.IR.GPU (GPU)+import Futhark.IR.MC (MC)+import Futhark.Optimise.ArrayLayout.Layout (layoutTableFromIndexTable)+import Futhark.Optimise.ArrayLayout.Transform (Transform, transformStms)+import Futhark.Pass++optimiseArrayLayout :: (Analyse rep, Transform rep, BuilderOps rep) => String -> Pass rep rep+optimiseArrayLayout s =+  Pass+    ("optimise array layout " <> s)+    "Transform array layout for locality optimisations."+    $ \prog -> do+      -- Analyse the program+      let index_table = analyseDimAccesses prog+      -- Compute primExps for all variables+      let table = primExpTable prog+      -- Compute permutations to acheive coalescence for all arrays+      let permutation_table = layoutTableFromIndexTable table index_table+      -- Insert permutations in the AST+      intraproceduralTransformation (onStms permutation_table) prog+  where+    onStms layout_table scope stms = do+      let m = transformStms layout_table mempty stms+      fmap fst $ modifyNameSource $ runState $ runBuilderT m scope++-- | The optimisation performed on the GPU representation.+optimiseArrayLayoutGPU :: Pass GPU GPU+optimiseArrayLayoutGPU = optimiseArrayLayout "gpu"++-- | The optimisation performed on the MC representation.+optimiseArrayLayoutMC :: Pass MC MC+optimiseArrayLayoutMC = optimiseArrayLayout "mc"

@@ -0,0 +1,265 @@+module Futhark.Optimise.ArrayLayout.Layout+  ( layoutTableFromIndexTable,+    Layout,+    Permutation,+    LayoutTable,++    -- * Exposed for testing+    commonPermutationEliminators,+  )+where++import Control.Monad (join)+import Data.List qualified as L+import Data.Map.Strict qualified as M+import Data.Maybe+import Futhark.Analysis.AccessPattern+import Futhark.Analysis.PrimExp.Table (PrimExpTable)+import Futhark.IR.Aliases+import Futhark.IR.GPU+import Futhark.IR.MC+import Futhark.IR.MCMem+import Futhark.Util (mininum)++type Permutation = [Int]++type LayoutTable =+  M.Map+    SegOpName+    ( M.Map+        ArrayName+        (M.Map IndexExprName Permutation)+    )++class Layout rep where+  -- | Produce a coalescing permutation that will be used to create a+  -- manifest of the array. Returns Nothing if the array is already in+  -- the optimal layout or if the array access is too complex to+  -- confidently determine the optimal layout. Map each list of+  -- 'DimAccess' in the IndexTable to a permutation in a generic way+  -- that can be handled uniquely by each backend.+  permutationFromDimAccess ::+    PrimExpTable ->+    SegOpName ->+    ArrayName ->+    IndexExprName ->+    [DimAccess rep] ->+    Maybe Permutation++isInscrutableExp :: PrimExp VName -> Bool+isInscrutableExp (LeafExp _ _) = False+isInscrutableExp (ValueExp _) = False+isInscrutableExp (BinOpExp _ a b) =+  isInscrutableExp a || isInscrutableExp b+isInscrutableExp (UnOpExp _ a) =+  isInscrutableExp a+isInscrutableExp _ = True++isInscrutable :: PrimExp VName -> Bool -> Bool+isInscrutable op@(BinOpExp {}) counter =+  if counter+    then -- Calculate stride and offset for loop-counters and thread-IDs+    case reduceStrideAndOffset op of+      -- Maximum allowable stride, might need tuning.+      Just (s, _) -> s > 8+      Nothing -> isInscrutableExp op+    else isInscrutableExp op+isInscrutable op _ = isInscrutableExp op++reduceStrideAndOffset :: PrimExp l -> Maybe (Int, Int)+reduceStrideAndOffset (LeafExp _ _) = Just (1, 0)+reduceStrideAndOffset (BinOpExp oper a b) = case (a, b) of+  (ValueExp (IntValue v), _) -> reduce v b+  (_, ValueExp (IntValue v)) -> reduce v a+  _ -> Nothing+  where+    reduce v (LeafExp _ _) =+      case oper of+        Add _ _ -> Just (1, valueIntegral v)+        Sub _ _ -> Just (1, -valueIntegral v)+        Mul _ _ -> Just (valueIntegral v, 0)+        _ -> Nothing+    reduce v op@(BinOpExp {}) =+      case reduceStrideAndOffset op of+        Nothing -> Nothing+        Just (s, o) -> case oper of+          Add _ _ -> Just (s, o + valueIntegral v)+          Sub _ _ -> Just (s, o - valueIntegral v)+          Mul _ _ -> Just (s * valueIntegral v, o * valueIntegral v)+          _ -> Nothing+    reduce _ (UnOpExp Not _) = Nothing+    reduce _ (UnOpExp (Complement _) _) = Nothing+    reduce _ (UnOpExp (Abs _) _) = Nothing+    reduce _ (UnOpExp _ sub_op) = reduceStrideAndOffset sub_op+    reduce _ (ConvOpExp _ sub_op) = reduceStrideAndOffset sub_op+    reduce _ _ = Nothing+reduceStrideAndOffset _ = Nothing++-- | Reasons common to all backends to not manifest an array.+commonPermutationEliminators :: [Int] -> [BodyType] -> Bool+commonPermutationEliminators perm nest = do+  -- Don't manifest if the permutation is the permutation is invalid+  let is_invalid_perm = not (L.sort perm `L.isPrefixOf` [0 ..])+      -- Don't manifest if the permutation is the identity permutation+      is_identity = perm `L.isPrefixOf` [0 ..]+      -- or is not a transpose.+      inefficient_transpose = isNothing $ isMapTranspose perm+      -- or if the last idx remains last+      static_last_idx = last perm == length perm - 1+      -- Don't manifest if the array is defined inside a segOp+      inside_undesired = any undesired nest++  is_invalid_perm+    || is_identity+    || inefficient_transpose+    || static_last_idx+    || inside_undesired+  where+    undesired :: BodyType -> Bool+    undesired bodyType = case bodyType of+      SegOpName _ -> True+      _ -> False++sortMC :: [(Int, DimAccess rep)] -> [(Int, DimAccess rep)]+sortMC =+  L.sortBy dimdexMCcmp+  where+    dimdexMCcmp (ia, a) (ib, b) = do+      let aggr1 =+            foldl max' Nothing $ map (f ia . snd) $ M.toList $ dependencies a+          aggr2 =+            foldl max' Nothing $ map (f ib . snd) $ M.toList $ dependencies b+      cmpIdxPat aggr1 aggr2+      where+        cmpIdxPat Nothing Nothing = EQ+        cmpIdxPat (Just _) Nothing = GT+        cmpIdxPat Nothing (Just _) = LT+        cmpIdxPat+          (Just (iterL, lvlL, original_lvl_L))+          (Just (iterR, lvlR, original_lvl_R)) =+            case (iterL, iterR) of+              (ThreadID, ThreadID) -> (lvlL, original_lvl_L) `compare` (lvlR, original_lvl_R)+              (ThreadID, _) -> LT+              (_, ThreadID) -> GT+              _ -> (lvlL, original_lvl_L) `compare` (lvlR, original_lvl_R)++        max' lhs rhs =+          case cmpIdxPat lhs rhs of+            LT -> rhs+            _ -> lhs++        f og (Dependency lvl varType) = Just (varType, lvl, og)++multicorePermutation :: PrimExpTable -> SegOpName -> ArrayName -> IndexExprName -> [DimAccess rep] -> Maybe Permutation+multicorePermutation primExpTable _segOpName (_arr_name, nest, arr_layout) _idx_name dimAccesses = do+  -- Dont accept indices where the last index is invariant+  let lastIdxIsInvariant = isInvariant $ last dimAccesses++  -- Check if any of the dependencies are too complex to reason about+  let dimAccesses' = filter (isJust . originalVar) dimAccesses+      deps = mapMaybe originalVar dimAccesses'+      counters = concatMap (map (isCounter . varType . snd) . M.toList . dependencies) dimAccesses'+      primExps = mapM (join . (`M.lookup` primExpTable)) deps+      inscrutable = maybe True (any (uncurry isInscrutable) . flip zip counters) primExps++  -- Create a candidate permutation+  let perm = map fst $ sortMC (zip arr_layout dimAccesses)++  -- Check if we want to manifest this array with the permutation+  if lastIdxIsInvariant || inscrutable || commonPermutationEliminators perm nest+    then Nothing+    else Just perm++instance Layout MC where+  permutationFromDimAccess = multicorePermutation++sortGPU :: [(Int, DimAccess rep)] -> [(Int, DimAccess rep)]+sortGPU =+  L.sortBy dimdexGPUcmp+  where+    dimdexGPUcmp (ia, a) (ib, b) = do+      let aggr1 =+            foldl max' Nothing $ map (f ia . snd) $ M.toList $ dependencies a+          aggr2 =+            foldl max' Nothing $ map (f ib . snd) $ M.toList $ dependencies b+      cmpIdxPat aggr1 aggr2+      where+        cmpIdxPat Nothing Nothing = EQ+        cmpIdxPat (Just _) Nothing = GT+        cmpIdxPat Nothing (Just _) = LT+        cmpIdxPat+          (Just (iterL, lvlL, original_lvl_L))+          (Just (iterR, lvlR, original_lvl_R)) = case (iterL, iterR) of+            (ThreadID, ThreadID) -> (lvlL, original_lvl_L) `compare` (lvlR, original_lvl_R)+            (ThreadID, _) -> GT+            (_, ThreadID) -> LT+            _ -> (lvlL, original_lvl_L) `compare` (lvlR, original_lvl_R)++        max' lhs rhs =+          case cmpIdxPat lhs rhs of+            LT -> rhs+            _ -> lhs++        f og (Dependency lvl varType) = Just (varType, lvl, og)++gpuPermutation :: PrimExpTable -> SegOpName -> ArrayName -> IndexExprName -> [DimAccess rep] -> Maybe Permutation+gpuPermutation primExpTable _segOpName (_arr_name, nest, arr_layout) _idx_name dimAccesses = do+  -- Find the outermost parallel level. XXX: this is a bit hacky. Why+  -- don't we simply know at this point the nest in which this index+  -- occurs?+  let outermost_par = mininum $ foldMap (map lvl . parDeps) dimAccesses+      invariantToPar = (< outermost_par) . lvl++  -- Do nothing if last index is invariant to segop.+  let lastIdxIsInvariant = all invariantToPar $ dependencies $ last dimAccesses++  -- Do nothing if any index is constant, because otherwise we can end+  -- up transposing a too-large array.+  let anyIsConstant = any (null . dependencies) dimAccesses++  -- Check if any of the dependencies are too complex to reason about+  let dimAccesses' = filter (isJust . originalVar) dimAccesses+      deps = mapMaybe originalVar dimAccesses'+      counters = concatMap (map (isCounter . varType . snd) . M.toList . dependencies) dimAccesses'+      primExps = mapM (join . (`M.lookup` primExpTable)) deps+      inscrutable = maybe True (any (uncurry isInscrutable) . flip zip counters) primExps++  -- Create a candidate permutation+  let perm = map fst $ sortGPU (zip arr_layout dimAccesses)++  -- Check if we want to manifest this array with the permutation+  if lastIdxIsInvariant+    || anyIsConstant+    || inscrutable+    || commonPermutationEliminators perm nest+    then Nothing+    else Just perm+  where+    parDeps = filter ((== ThreadID) . varType) . M.elems . dependencies++instance Layout GPU where+  permutationFromDimAccess = gpuPermutation++-- | like mapMaybe, but works on nested maps. Eliminates "dangling"+-- maps / rows with missing (Nothing) values.+tableMapMaybe ::+  (k0 -> k1 -> k2 -> a -> Maybe b) ->+  M.Map k0 (M.Map k1 (M.Map k2 a)) ->+  M.Map k0 (M.Map k1 (M.Map k2 b))+tableMapMaybe f =+  M.mapMaybeWithKey $ \key0 -> mapToMaybe $ mapToMaybe . f key0+  where+    maybeMap :: M.Map k a -> Maybe (M.Map k a)+    maybeMap val = if null val then Nothing else Just val++    mapToMaybe g = maybeMap . M.mapMaybeWithKey g++-- | Given an ordering function for `DimAccess`, and an IndexTable,+-- return a LayoutTable. We remove entries with no results after+-- `permutationFromDimAccess`+layoutTableFromIndexTable ::+  (Layout rep) =>+  PrimExpTable ->+  IndexTable rep ->+  LayoutTable+layoutTableFromIndexTable = tableMapMaybe . permutationFromDimAccess

@@ -0,0 +1,270 @@+-- | Do various kernel optimisations - mostly related to coalescing.+module Futhark.Optimise.ArrayLayout.Transform+  ( Transform,+    transformStms,+  )+where++import Control.Monad+import Control.Monad.State.Strict+import Data.Map.Strict qualified as M+import Futhark.Analysis.AccessPattern (IndexExprName, SegOpName (..))+import Futhark.Analysis.PrimExp.Table (PrimExpAnalysis)+import Futhark.Builder+import Futhark.Construct+import Futhark.IR.Aliases+import Futhark.IR.GPU+import Futhark.IR.MC+import Futhark.Optimise.ArrayLayout.Layout (Layout, LayoutTable, Permutation)++class (Layout rep, PrimExpAnalysis rep) => Transform rep where+  onOp ::+    (Monad m) =>+    SOACMapper rep rep m ->+    Op rep ->+    m (Op rep)+  transformOp ::+    LayoutTable ->+    ExpMap rep ->+    Stm rep ->+    Op rep ->+    TransformM rep (LayoutTable, ExpMap rep)++type TransformM rep = Builder rep++-- | A map from the name of an expression to the expression that defines it.+type ExpMap rep = M.Map VName (Stm rep)++instance Transform GPU where+  onOp soac_mapper (Futhark.IR.GPU.OtherOp soac) =+    Futhark.IR.GPU.OtherOp <$> mapSOACM soac_mapper soac+  onOp _ op = pure op+  transformOp perm_table expmap stm gpuOp+    | SegOp op <- gpuOp,+      -- TODO: handle non-segthread cases. This requires some care to+      -- avoid doing huge manifests at the block level.+      SegThread {} <- segLevel op =+        transformSegOpGPU perm_table expmap stm op+    | _ <- gpuOp = transformRestOp perm_table expmap stm++instance Transform MC where+  onOp soac_mapper (Futhark.IR.MC.OtherOp soac) =+    Futhark.IR.MC.OtherOp <$> mapSOACM soac_mapper soac+  onOp _ op = pure op+  transformOp perm_table expmap stm mcOp+    | ParOp maybe_par_segop seqSegOp <- mcOp =+        transformSegOpMC perm_table expmap stm maybe_par_segop seqSegOp+    | _ <- mcOp = transformRestOp perm_table expmap stm++transformSegOpGPU :: LayoutTable -> ExpMap GPU -> Stm GPU -> SegOp SegLevel GPU -> TransformM GPU (LayoutTable, ExpMap GPU)+transformSegOpGPU perm_table expmap stm@(Let pat aux _) op =+  -- Optimization: Only traverse the body of the SegOp if it is+  -- represented in the layout table+  case M.lookup patternName (M.mapKeys vnameFromSegOp perm_table) of+    Nothing -> do+      addStm stm+      pure (perm_table, M.fromList [(name, stm) | name <- patNames pat] <> expmap)+    Just _ -> do+      let mapper =+            identitySegOpMapper+              { mapOnSegOpBody = case segLevel op of+                  SegBlock {} -> transformSegGroupKernelBody perm_table expmap+                  _ -> transformSegThreadKernelBody perm_table patternName+              }+      op' <- mapSegOpM mapper op+      let stm' = Let pat aux $ Op $ SegOp op'+      addStm stm'+      pure (perm_table, M.fromList [(name, stm') | name <- patNames pat] <> expmap)+  where+    patternName = patElemName . head $ patElems pat++transformSegOpMC :: LayoutTable -> ExpMap MC -> Stm MC -> Maybe (SegOp () MC) -> SegOp () MC -> TransformM MC (LayoutTable, ExpMap MC)+transformSegOpMC perm_table expmap (Let pat aux _) maybe_par_segop seqSegOp+  | Nothing <- maybe_par_segop = add Nothing+  | Just par_segop <- maybe_par_segop =+      -- Optimization: Only traverse the body of the SegOp if it is+      -- represented in the layout table+      case M.lookup patternName (M.mapKeys vnameFromSegOp perm_table) of+        Nothing -> add $ Just par_segop+        Just _ -> add . Just =<< mapSegOpM mapper par_segop+  where+    add maybe_par_segop' = do+      -- Map the sequential part of the ParOp+      seqSegOp' <- mapSegOpM mapper seqSegOp+      let stm' = Let pat aux $ Op $ ParOp maybe_par_segop' seqSegOp'+      addStm stm'+      pure (perm_table, M.fromList [(name, stm') | name <- patNames pat] <> expmap)+    mapper = identitySegOpMapper {mapOnSegOpBody = transformKernelBody perm_table expmap patternName}+    patternName = patElemName . head $ patElems pat++transformRestOp :: (Transform rep, BuilderOps rep) => LayoutTable -> ExpMap rep -> Stm rep -> TransformM rep (LayoutTable, ExpMap rep)+transformRestOp perm_table expmap (Let pat aux e) = do+  e' <- mapExpM (transform perm_table expmap) e+  let stm' = Let pat aux e'+  addStm stm'+  pure (perm_table, M.fromList [(name, stm') | name <- patNames pat] <> expmap)++transform :: (Transform rep, BuilderOps rep) => LayoutTable -> ExpMap rep -> Mapper rep rep (TransformM rep)+transform perm_table expmap =+  identityMapper {mapOnBody = \scope -> localScope scope . transformBody perm_table expmap}++-- | Recursively transform the statements in a body.+transformBody :: (Transform rep, BuilderOps rep) => LayoutTable -> ExpMap rep -> Body rep -> TransformM rep (Body rep)+transformBody perm_table expmap (Body b stms res) =+  Body b <$> transformStms perm_table expmap stms <*> pure res++-- | Recursively transform the statements in the body of a SegGroup kernel.+transformSegGroupKernelBody ::+  (Transform rep, BuilderOps rep) =>+  LayoutTable ->+  ExpMap rep ->+  KernelBody rep ->+  TransformM rep (KernelBody rep)+transformSegGroupKernelBody perm_table expmap (KernelBody b stms res) =+  KernelBody b <$> transformStms perm_table expmap stms <*> pure res++-- | Transform the statements in the body of a SegThread kernel.+transformSegThreadKernelBody ::+  (Transform rep, BuilderOps rep) =>+  LayoutTable ->+  VName ->+  KernelBody rep ->+  TransformM rep (KernelBody rep)+transformSegThreadKernelBody perm_table seg_name kbody = do+  evalStateT+    ( traverseKernelBodyArrayIndexes+        seg_name+        (ensureTransformedAccess perm_table)+        kbody+    )+    mempty++transformKernelBody ::+  (Transform rep, BuilderOps rep) =>+  LayoutTable ->+  ExpMap rep ->+  VName ->+  KernelBody rep ->+  TransformM rep (KernelBody rep)+transformKernelBody perm_table expmap seg_name (KernelBody b stms res) = do+  stms' <- transformStms perm_table expmap stms+  evalStateT+    ( traverseKernelBodyArrayIndexes+        seg_name+        (ensureTransformedAccess perm_table)+        (KernelBody b stms' res)+    )+    mempty++traverseKernelBodyArrayIndexes ::+  forall m rep.+  (Monad m, Transform rep) =>+  VName -> -- seg_name+  ArrayIndexTransform m ->+  KernelBody rep ->+  m (KernelBody rep)+traverseKernelBodyArrayIndexes seg_name coalesce (KernelBody b kstms kres) =+  KernelBody b . stmsFromList+    <$> mapM onStm (stmsToList kstms)+    <*> pure kres+  where+    onLambda lam =+      (\body' -> lam {lambdaBody = body'})+        <$> onBody (lambdaBody lam)++    onBody (Body bdec stms bres) = do+      stms' <- stmsFromList <$> mapM onStm (stmsToList stms)+      pure $ Body bdec stms' bres++    onStm (Let pat dec (BasicOp (Index arr is))) =+      Let pat dec . oldOrNew <$> coalesce seg_name patternName arr is+      where+        oldOrNew Nothing =+          BasicOp $ Index arr is+        oldOrNew (Just (arr', is')) =+          BasicOp $ Index arr' is'+        patternName = patElemName . head $ patElems pat+    onStm (Let pat dec e) =+      Let pat dec <$> mapExpM mapper e++    soac_mapper =+      identitySOACMapper+        { mapOnSOACLambda = onLambda+        }++    mapper =+      (identityMapper @rep)+        { mapOnBody = const onBody,+          mapOnOp = onOp soac_mapper+        }++-- | Used to keep track of which pairs of arrays and permutations we have+-- already created manifests for, in order to avoid duplicates.+type Replacements = M.Map (VName, Permutation) VName++type ArrayIndexTransform m =+  VName -> -- seg_name (name of the SegThread expression's pattern)+  VName -> -- idx_name (name of the Index expression's pattern)+  VName -> -- arr (name of the array)+  Slice SubExp -> -- slice+  m (Maybe (VName, Slice SubExp))++ensureTransformedAccess ::+  (MonadBuilder m) =>+  LayoutTable ->+  ArrayIndexTransform (StateT Replacements m)+ensureTransformedAccess perm_table seg_name idx_name arr slice = do+  -- Check if the array has the optimal layout in memory.+  -- If it does not, replace it with a manifest to allocate+  -- it with the optimal layout+  case lookupPermutation perm_table seg_name idx_name arr of+    Nothing -> pure $ Just (arr, slice)+    Just perm -> do+      seen <- gets $ M.lookup (arr, perm)+      case seen of+        -- Already created a manifest for this array + permutation.+        -- So, just replace the name and don't make a new manifest.+        Just arr' -> pure $ Just (arr', slice)+        Nothing -> replace perm =<< lift (manifest perm arr)+  where+    replace perm arr' = do+      -- Store the fact that we have seen this array + permutation+      -- so we don't make duplicate manifests+      modify $ M.insert (arr, perm) arr'+      -- Return the new manifest+      pure $ Just (arr', slice)++    manifest perm array =+      letExp (baseString array ++ "_coalesced") $+        BasicOp (Manifest perm array)++lookupPermutation :: LayoutTable -> VName -> IndexExprName -> VName -> Maybe Permutation+lookupPermutation perm_table seg_name idx_name arr_name =+  case M.lookup seg_name (M.mapKeys vnameFromSegOp perm_table) of+    Nothing -> Nothing+    Just arrayNameMap ->+      -- Look for the current array+      case M.lookup arr_name (M.mapKeys (\(n, _, _) -> n) arrayNameMap) of+        Nothing -> Nothing+        Just idxs -> M.lookup idx_name idxs++transformStm ::+  (Transform rep, BuilderOps rep) =>+  (LayoutTable, ExpMap rep) ->+  Stm rep ->+  TransformM rep (LayoutTable, ExpMap rep)+transformStm (perm_table, expmap) (Let pat aux (Op op)) = transformOp perm_table expmap (Let pat aux (Op op)) op+transformStm (perm_table, expmap) (Let pat aux e) = do+  e' <- mapExpM (transform perm_table expmap) e+  let stm' = Let pat aux e'+  addStm stm'+  pure (perm_table, M.fromList [(name, stm') | name <- patNames pat] <> expmap)++transformStms ::+  (Transform rep, BuilderOps rep) =>+  LayoutTable ->+  ExpMap rep ->+  Stms rep ->+  TransformM rep (Stms rep)+transformStms perm_table expmap stms =+  collectStms_ $ foldM_ transformStm (perm_table, expmap) stms

@@ -436,7 +436,7 @@       foldl getAccumStm False $ reverse $ stmsToList acc_code       where         getAccumStm True _ = True-        getAccumStm False (Let (Pat [pat_el]) _aux (BasicOp (UpdateAcc _acc_nm _ind vals)))+        getAccumStm False (Let (Pat [pat_el]) _aux (BasicOp (UpdateAcc Safe _acc_nm _ind vals)))           | [v] <- vals,             patElemName pat_el == res_nm =               v == Var redomap_orig_res

@@ -105,7 +105,7 @@     -- some `code1`, followed by one accumulation, followed by some `code2`     -- UpdateAcc VName [SubExp] [SubExp]     (code1, Just accum_stmt, code2) <- matchCodeAccumCode kstms,-    Let pat_accum _aux_acc (BasicOp (UpdateAcc acc_nm acc_inds acc_vals)) <- accum_stmt,+    Let pat_accum _aux_acc (BasicOp (UpdateAcc safety acc_nm acc_inds acc_vals)) <- accum_stmt,     [pat_acc_nm] <- patNames pat_accum,     -- check that the `acc_inds` are invariant to at least one     -- parallel kernel dimensions, and return the innermost such one:@@ -153,7 +153,7 @@         let op_exp = Op (OtherOp (Screma inv_dim_len [iota] (ScremaForm [] [red] map_lam)))         res_redmap <- letTupExp "res_mapred" op_exp         letSubExp (baseString pat_acc_nm ++ "_big_update") $-          BasicOp (UpdateAcc acc_nm acc_inds $ map Var res_redmap)+          BasicOp (UpdateAcc safety acc_nm acc_inds $ map Var res_redmap)        -- 1.3. build the kernel expression and rename it!       gid_flat_1 <- newVName "gid_flat"

@@ -38,7 +38,7 @@ extractUpdate accs v stms = do   (stm, stms') <- stmsHead stms   case stm of-    Let (Pat [PatElem pe_v _]) _ (BasicOp (UpdateAcc acc is vs))+    Let (Pat [PatElem pe_v _]) _ (BasicOp (UpdateAcc _ acc is vs))       | pe_v == v -> do           acc_input <- M.lookup acc accs           Just ((acc_input, acc, is, vs), stms')@@ -82,7 +82,7 @@               map (DimFix . Var) gtids     letExp (baseString acc <> "_upd") $       BasicOp $-        UpdateAcc acc (map Var gtids) vs+        UpdateAcc Safe acc (map Var gtids) vs    acc_t <- lookupType acc   pure . Op . SegOp . SegMap lvl space [acc_t] $

@@ -210,7 +210,7 @@       stms' <- onStms $ bodyStms body       pure body {bodyStms = stms'}     onStms = traverse onStm-    onStm (Let pat@(Pat [PatElem _ dec]) aux (BasicOp (UpdateAcc acc _ _)))+    onStm (Let pat@(Pat [PatElem _ dec]) aux (BasicOp (UpdateAcc _ acc _ _)))       | Acc c _ _ _ <- typeOf dec,         c `elem` get_rid_of = do           modify (insert c)

@@ -349,7 +349,7 @@               Var v -- Remove UpdateAccs that contribute the neutral value, which is -- always a no-op.-ruleBasicOp vtable pat aux (UpdateAcc acc _ vs)+ruleBasicOp vtable pat aux (UpdateAcc _ acc _ vs)   | Pat [pe] <- pat,     Acc token _ _ _ <- patElemType pe,     Just (_, _, Just (_, ne)) <- ST.entryAccInput =<< ST.lookup token vtable,

@@ -316,14 +316,14 @@        trExp i (WithAcc acc_inputs lam) =         WithAcc acc_inputs <$> trLam i lam-      trExp i (BasicOp (UpdateAcc acc is ses)) = do+      trExp i (BasicOp (UpdateAcc safety acc is ses)) = do         acc_t <- lookupType acc         pure $ case acc_t of           Acc cert _ _ _             | cert `elem` acc_certs ->-                BasicOp $ UpdateAcc acc (i : is) ses+                BasicOp $ UpdateAcc safety acc (i : is) ses           _ ->-            BasicOp $ UpdateAcc acc is ses+            BasicOp $ UpdateAcc safety acc is ses       trExp i e = mapExpM mapper e         where           mapper =

@@ -1,407 +0,0 @@-{-# LANGUAGE TypeFamilies #-}---- | Do various kernel optimisations - mostly related to coalescing.-module Futhark.Pass.KernelBabysitting (babysitKernels) where--import Control.Monad-import Control.Monad.State.Strict-import Data.Bifunctor (first)-import Data.Foldable-import Data.List qualified as L-import Data.Map.Strict qualified as M-import Data.Maybe-import Futhark.IR.GPU-import Futhark.Pass-import Futhark.Tools-import Futhark.Util---- | The pass definition.-babysitKernels :: Pass GPU GPU-babysitKernels =-  Pass-    "babysit kernels"-    "Transpose kernel input arrays for better performance."-    $ intraproceduralTransformation onStms-  where-    onStms scope stms = do-      let m = localScope scope $ transformStms mempty stms-      fmap fst $ modifyNameSource $ runState (runBuilderT m M.empty)--type BabysitM = Builder GPU--transformStms :: ExpMap -> Stms GPU -> BabysitM (Stms GPU)-transformStms expmap stms = collectStms_ $ foldM_ transformStm expmap stms--transformBody :: ExpMap -> Body GPU -> BabysitM (Body GPU)-transformBody expmap (Body () stms res) = do-  stms' <- transformStms expmap stms-  pure $ Body () stms' res---- | Map from variable names to defining expression.  We use this to--- hackily determine whether something is transposed or otherwise--- funky in memory (and we'd prefer it not to be).  If we cannot find--- it in the map, we just assume it's all good.  HACK and FIXME, I--- suppose.  We really should do this at the memory level.-type ExpMap = M.Map VName (Stm GPU)--nonlinearInMemory :: VName -> ExpMap -> Maybe (Maybe [Int])-nonlinearInMemory name m =-  case M.lookup name m of-    Just (Let _ _ (BasicOp (Opaque _ (Var arr)))) -> nonlinearInMemory arr m-    Just (Let _ _ (BasicOp (Rearrange perm _))) -> Just $ Just $ rearrangeInverse perm-    Just (Let _ _ (BasicOp (Reshape _ _ arr))) -> nonlinearInMemory arr m-    Just (Let _ _ (BasicOp (Manifest perm _))) -> Just $ Just perm-    Just (Let pat _ (Op (SegOp (SegMap _ _ ts _)))) ->-      nonlinear-        =<< find-          ((== name) . patElemName . fst)-          (zip (patElems pat) ts)-    _ -> Nothing-  where-    nonlinear (pe, t)-      | inner_r <- arrayRank t,-        inner_r > 0 = do-          let outer_r = arrayRank (patElemType pe) - inner_r-          pure $ Just $ rearrangeInverse $ [inner_r .. inner_r + outer_r - 1] ++ [0 .. inner_r - 1]-      | otherwise = Nothing--transformStm :: ExpMap -> Stm GPU -> BabysitM ExpMap-transformStm expmap (Let pat aux (Op (SegOp op)))-  -- FIXME: We only make coalescing optimisations for SegThread-  -- SegOps, because that's what the analysis assumes.  For SegBlock-  -- we should probably look at the component SegThreads, but it-  -- apparently hasn't come up in practice yet.-  | SegThread {} <- segLevel op = do-      let mapper =-            identitySegOpMapper-              { mapOnSegOpBody =-                  transformKernelBody expmap (segSpace op)-              }-      op' <- mapSegOpM mapper op-      let stm' = Let pat aux $ Op $ SegOp op'-      addStm stm'-      pure $ M.fromList [(name, stm') | name <- patNames pat] <> expmap-transformStm expmap (Let pat aux e) = do-  e' <- mapExpM (transform expmap) e-  let stm' = Let pat aux e'-  addStm stm'-  pure $ M.fromList [(name, stm') | name <- patNames pat] <> expmap--transform :: ExpMap -> Mapper GPU GPU BabysitM-transform expmap =-  identityMapper {mapOnBody = \scope -> localScope scope . transformBody expmap}--transformKernelBody ::-  ExpMap ->-  SegSpace ->-  KernelBody GPU ->-  BabysitM (KernelBody GPU)-transformKernelBody expmap space kbody = do-  -- Go spelunking for accesses to arrays that are defined outside the-  -- kernel body and where the indices are kernel thread indices.-  scope <- askScope-  let thread_gids = map fst $ unSegSpace space-      thread_local = namesFromList $ segFlat space : thread_gids-      free_ker_vars = freeIn kbody `namesSubtract` getKerVariantIds space-  evalStateT-    ( traverseKernelBodyArrayIndexes-        free_ker_vars-        thread_local-        (scope <> scopeOfSegSpace space)-        (ensureCoalescedAccess expmap (unSegSpace space))-        kbody-    )-    mempty-  where-    getKerVariantIds = namesFromList . M.keys . scopeOfSegSpace--type ArrayIndexTransform m =-  Names ->-  (VName -> Bool) -> -- thread local?-  (VName -> SubExp -> Bool) -> -- variant to a certain gid (given as first param)?-  Scope GPU -> -- type environment-  VName ->-  Slice SubExp ->-  m (Maybe (VName, Slice SubExp))--traverseKernelBodyArrayIndexes ::-  forall f.-  (Monad f) =>-  Names ->-  Names ->-  Scope GPU ->-  ArrayIndexTransform f ->-  KernelBody GPU ->-  f (KernelBody GPU)-traverseKernelBodyArrayIndexes free_ker_vars thread_variant outer_scope f (KernelBody () kstms kres) =-  KernelBody () . stmsFromList-    <$> mapM-      ( onStm-          ( varianceInStms mempty kstms,-            outer_scope-          )-      )-      (stmsToList kstms)-    <*> pure kres-  where-    onLambda :: (VarianceTable, Scope GPU) -> Lambda GPU -> f (Lambda GPU)-    onLambda (variance, scope) lam =-      (\body' -> lam {lambdaBody = body'})-        <$> onBody (variance, scope') (lambdaBody lam)-      where-        scope' = scope <> scopeOfLParams (lambdaParams lam)--    onBody (variance, scope) (Body bdec stms bres) = do-      stms' <- stmsFromList <$> mapM (onStm (variance', scope')) (stmsToList stms)-      pure $ Body bdec stms' bres-      where-        variance' = varianceInStms variance stms-        scope' = scope <> scopeOf stms--    onStm (variance, _) (Let pat dec (BasicOp (Index arr is))) =-      Let pat dec . oldOrNew <$> f free_ker_vars isThreadLocal isGidVariant outer_scope arr is-      where-        oldOrNew Nothing =-          BasicOp $ Index arr is-        oldOrNew (Just (arr', is')) =-          BasicOp $ Index arr' is'--        isGidVariant gid (Var v) =-          gid == v || nameIn gid (M.findWithDefault (oneName v) v variance)-        isGidVariant _ _ = False--        isThreadLocal v =-          thread_variant-            `namesIntersect` M.findWithDefault (oneName v) v variance-    onStm (variance, scope) (Let pat dec e) =-      Let pat dec <$> mapExpM (mapper (variance, scope)) e--    onOp :: (VarianceTable, Scope GPU) -> Op GPU -> f (Op GPU)-    onOp ctx (OtherOp soac) =-      OtherOp <$> mapSOACM (soacMapper ctx) soac-    onOp _ op = pure op--    soacMapper ctx =-      (identitySOACMapper @GPU)-        { mapOnSOACLambda = onLambda ctx-        }--    mapper ctx =-      (identityMapper @GPU)-        { mapOnBody = const (onBody ctx),-          mapOnOp = onOp ctx-        }--type Replacements = M.Map (VName, Slice SubExp) VName--ensureCoalescedAccess ::-  (MonadBuilder m) =>-  ExpMap ->-  [(VName, SubExp)] ->-  ArrayIndexTransform (StateT Replacements m)-ensureCoalescedAccess-  expmap-  thread_space-  free_ker_vars-  isThreadLocal-  isGidVariant-  outer_scope-  arr-  slice = do-    seen <- gets $ M.lookup (arr, slice)--    case (seen, isThreadLocal arr, typeOf <$> M.lookup arr outer_scope) of-      -- Already took care of this case elsewhere.-      (Just arr', _, _) ->-        pure $ Just (arr', slice)-      (Nothing, False, Just t)-        -- We are fully indexing the array with thread IDs, but the-        -- indices are in a permuted order.-        | Just is <- sliceIndices slice,-          length is == arrayRank t,-          Just is' <- coalescedIndexes free_ker_vars isGidVariant (map Var thread_gids) is,-          Just perm <- is' `isPermutationOf` is ->-            replace =<< lift (rearrangeInput (nonlinearInMemory arr expmap) perm arr)-        -- Check whether the access is already coalesced because of a-        -- previous rearrange being applied to the current array:-        -- 1. get the permutation of the source-array rearrange-        -- 2. apply it to the slice-        -- 3. check that the innermost index is actually the gid-        --    of the innermost kernel dimension.-        -- If so, the access is already coalesced, nothing to do!-        -- (Cosmin's Heuristic.)-        | Just (Let _ _ (BasicOp (Rearrange perm _))) <- M.lookup arr expmap,-          not $ null perm,-          not $ null thread_gids,-          inner_gid <- last thread_gids,-          length slice >= length perm,-          slice' <- map (unSlice slice !!) perm,-          DimFix inner_ind <- last slice',-          not $ null thread_gids,-          isGidVariant inner_gid inner_ind ->-            pure Nothing-        -- We are not fully indexing an array, but the remaining slice-        -- is invariant to the innermost-kernel dimension. We assume-        -- the remaining slice will be sequentially streamed, hence-        -- tiling will be applied later and will solve coalescing.-        -- Hence nothing to do at this point. (Cosmin's Heuristic.)-        | (is, rem_slice) <- splitSlice slice,-          not $ null rem_slice,-          allDimAreSlice rem_slice,-          Nothing <- M.lookup arr expmap,-          pt <- elemType t,-          not $ tooSmallSlice (primByteSize pt) rem_slice,-          is /= map Var (take (length is) thread_gids) || length is == length thread_gids,-          not (null thread_gids || null is),-          last thread_gids `notNameIn` (freeIn is <> freeIn rem_slice) ->-            pure Nothing-        -- We are not fully indexing the array, and the indices are not-        -- a proper prefix of the thread indices, and some indices are-        -- thread local, so we assume (HEURISTIC!)  that the remaining-        -- dimensions will be traversed sequentially.-        | (is, rem_slice) <- splitSlice slice,-          not $ null rem_slice,-          pt <- elemType t,-          not $ tooSmallSlice (primByteSize pt) rem_slice,-          is /= map Var (take (length is) thread_gids) || length is == length thread_gids,-          any isThreadLocal (namesToList $ freeIn is) -> do-            let perm = coalescingPermutation (length is) $ arrayRank t-            replace =<< lift (rearrangeInput (nonlinearInMemory arr expmap) perm arr)--        -- Everything is fine... assuming that the array is in row-major-        -- order!  Make sure that is the case.-        | Just {} <- nonlinearInMemory arr expmap ->-            case sliceIndices slice of-              Just is-                | Just _ <- coalescedIndexes free_ker_vars isGidVariant (map Var thread_gids) is ->-                    replace =<< lift (rowMajorArray arr)-                | otherwise ->-                    pure Nothing-              _ -> replace =<< lift (rowMajorArray arr)-      _ -> pure Nothing-    where-      (thread_gids, _thread_gdims) = unzip thread_space--      replace arr' = do-        modify $ M.insert (arr, slice) arr'-        pure $ Just (arr', slice)---- Heuristic for avoiding rearranging too small arrays.-tooSmallSlice :: Int32 -> Slice SubExp -> Bool-tooSmallSlice bs = fst . foldl comb (True, bs) . sliceDims-  where-    comb (True, x) (Constant (IntValue (Int32Value d))) = (d * x < 4, d * x)-    comb (_, x) _ = (False, x)--splitSlice :: Slice SubExp -> ([SubExp], Slice SubExp)-splitSlice (Slice []) = ([], Slice [])-splitSlice (Slice (DimFix i : is)) = first (i :) $ splitSlice (Slice is)-splitSlice is = ([], is)--allDimAreSlice :: Slice SubExp -> Bool-allDimAreSlice (Slice []) = True-allDimAreSlice (Slice (DimFix _ : _)) = False-allDimAreSlice (Slice (_ : is)) = allDimAreSlice (Slice is)---- Try to move thread indexes into their proper position.-coalescedIndexes :: Names -> (VName -> SubExp -> Bool) -> [SubExp] -> [SubExp] -> Maybe [SubExp]-coalescedIndexes free_ker_vars isGidVariant tgids is-  -- Do Nothing if:-  -- 1. any of the indices is a constant or a kernel free variable-  --    (because it would transpose a bigger array then needed -- big overhead).-  -- 2. the innermost index is variant to the innermost-thread gid-  --    (because access is likely to be already coalesced)-  -- 3. the indexes are a prefix of the thread indexes, because that-  -- means multiple threads will be accessing the same element.-  | any isCt is =-      Nothing-  | any (`nameIn` free_ker_vars) (subExpVars is) =-      Nothing-  | is `L.isPrefixOf` tgids =-      Nothing-  | not (null tgids),-    not (null is),-    Var innergid <- last tgids,-    num_is > 0 && isGidVariant innergid (last is) =-      Just is-  -- 3. Otherwise try fix coalescing-  | otherwise =-      Just $ reverse $ foldl move (reverse is) $ zip [0 ..] (reverse tgids)-  where-    num_is = length is--    move is_rev (i, tgid)-      -- If tgid is in is_rev anywhere but at position i, and-      -- position i exists, we move it to position i instead.-      | Just j <- L.elemIndex tgid is_rev,-        i /= j,-        i < num_is =-          swap i j is_rev-      | otherwise =-          is_rev--    swap i j l-      | Just ix <- maybeNth i l,-        Just jx <- maybeNth j l =-          update i jx $ update j ix l-      | otherwise =-          error $ "coalescedIndexes swap: invalid indices" ++ show (i, j, l)--    update 0 x (_ : ys) = x : ys-    update i x (y : ys) = y : update (i - 1) x ys-    update _ _ [] = error "coalescedIndexes: update"--    isCt :: SubExp -> Bool-    isCt (Constant _) = True-    isCt (Var _) = False--coalescingPermutation :: Int -> Int -> [Int]-coalescingPermutation num_is rank =-  [num_is .. rank - 1] ++ [0 .. num_is - 1]--rearrangeInput ::-  (MonadBuilder m) =>-  Maybe (Maybe [Int]) ->-  [Int] ->-  VName ->-  m VName-rearrangeInput (Just (Just current_perm)) perm arr-  | current_perm == perm = pure arr -- Already has desired representation.-rearrangeInput Nothing perm arr-  | L.sort perm == perm = pure arr -- We don't know the current-  -- representation, but the indexing-  -- is linear, so let's hope the-  -- array is too.-rearrangeInput (Just Just {}) perm arr-  | L.sort perm == perm = rowMajorArray arr -- We just want a row-major array, no tricks.-rearrangeInput manifest perm arr = do-  -- We may first manifest the array to ensure that it is flat in-  -- memory.  This is sometimes unnecessary, in which case the copy-  -- will hopefully be removed by the simplifier.-  manifested <- if isJust manifest then rowMajorArray arr else pure arr-  letExp (baseString arr ++ "_coalesced") $-    BasicOp $-      Manifest perm manifested--rowMajorArray ::-  (MonadBuilder m) =>-  VName ->-  m VName-rowMajorArray arr = do-  rank <- arrayRank <$> lookupType arr-  letExp (baseString arr ++ "_rowmajor") $ BasicOp $ Manifest [0 .. rank - 1] arr----- Computing variance.--type VarianceTable = M.Map VName Names--varianceInStms :: VarianceTable -> Stms GPU -> VarianceTable-varianceInStms t = foldl varianceInStm t . stmsToList--varianceInStm :: VarianceTable -> Stm GPU -> VarianceTable-varianceInStm variance stm =-  foldl' add variance $ patNames $ stmPat stm-  where-    add variance' v = M.insert v binding_variance variance'-    look variance' v = oneName v <> M.findWithDefault mempty v variance'-    binding_variance = mconcat $ map (look variance) $ namesToList (freeIn stm)

@@ -18,6 +18,7 @@ import Futhark.IR.SOACS (SOACS, usesAD) import Futhark.IR.Seq (Seq) import Futhark.IR.SeqMem (SeqMem)+import Futhark.Optimise.ArrayLayout import Futhark.Optimise.ArrayShortCircuiting qualified as ArrayShortCircuiting import Futhark.Optimise.CSE import Futhark.Optimise.DoubleBuffer@@ -40,7 +41,6 @@ import Futhark.Pass.ExtractKernels import Futhark.Pass.ExtractMulticore import Futhark.Pass.FirstOrderTransform-import Futhark.Pass.KernelBabysitting import Futhark.Pass.LiftAllocations as LiftAllocations import Futhark.Pass.LowerAllocations as LowerAllocations import Futhark.Pass.Simplify@@ -79,7 +79,7 @@       simplifySOACS     ] --- | The pipeline used by the CUDA and OpenCL backends, but before+-- | The pipeline used by the CUDA, HIP, and OpenCL backends, but before -- adding memory information.  Includes 'standardPipeline'. gpuPipeline :: Pipeline SOACS GPU gpuPipeline =@@ -102,8 +102,8 @@         mergeGPUBodies,         simplifyGPU, -- Cleanup merged GPUBody kernels.         sinkGPU, -- Sink reads within GPUBody kernels.-        babysitKernels,-        -- Important to simplify after babysitting in order to fix up+        optimiseArrayLayoutGPU,+        -- Important to simplify after coalescing in order to fix up         -- redundant manifests.         simplifyGPU,         performCSE True@@ -180,7 +180,10 @@         unstreamMC,         performCSE True,         simplifyMC,-        sinkMC+        sinkMC,+        optimiseArrayLayoutMC,+        simplifyMC,+        performCSE True       ]  -- | Run 'mcPipeline' and then add memory information.

@@ -78,7 +78,7 @@         now <- liftIO getCurrentTime         let delta :: Double             delta = fromRational $ toRational (now `diffUTCTime` prev)-            prefix = printf "[  +%.6f] " delta+            prefix = printf "[  +%7.3f] " delta         modify $ \s -> s {futharkPrevLog = now}         when verb $ liftIO $ T.hPutStrLn stderr $ T.pack prefix <> msg @@ -160,7 +160,7 @@   where     perform cfg prog = do       when (pipelineVerbose cfg) . logMsg $-        "Running pass " <> T.pack (passName pass)+        "Running pass: " <> T.pack (passName pass)       prog' <- runPass pass prog       -- Spark validation in a separate task and speculatively execute       -- next pass.  If the next pass throws an exception, we better

@@ -11,6 +11,7 @@     nubByOrd,     mapAccumLM,     maxinum,+    mininum,     chunk,     chunks,     chunkLike,@@ -144,6 +145,10 @@ -- | Like 'maximum', but returns zero for an empty list. maxinum :: (Num a, Ord a, Foldable f) => f a -> a maxinum = foldl' max 0++-- | Like 'minimum', but returns zero for an empty list.+mininum :: (Num a, Ord a, Foldable f) => f a -> a+mininum xs = foldl' min (maxinum xs) xs  -- | @dropAt i n@ drops @n@ elements starting at element @i@. dropAt :: Int -> Int -> [a] -> [a]

@@ -29,7 +29,11 @@ -- | The version of Futhark that we are using, in human-readable form. versionString :: T.Text versionString =-  T.pack (showVersion version) <> unreleased <> gitversion $$tGitInfoCwdTry <> ghcversion+  T.pack (showVersion version)+    <> unreleased+    <> ".\n"+    <> gitversion $$tGitInfoCwdTry+    <> ghcversion   where     unreleased =       if last (versionBranch version) == 0@@ -38,11 +42,10 @@     gitversion (Left _) =       case commitIdFromFile of         Nothing -> ""-        Just commit -> "\ngit: " <> T.pack commit+        Just commit -> "git: " <> T.pack commit <> "\n"     gitversion (Right gi) =       mconcat-        [ "\n",-          "git: ",+        [ "git: ",           branch,           T.pack (take 7 $ giHash gi),           " (",

@@ -109,9 +109,9 @@  instance Show (Value m) where   show (ValuePrim v) = "ValuePrim " <> show v <> ""-  show (ValueArray shape vs) = unwords ["ValueArray", show shape, show vs]-  show (ValueRecord fs) = "ValueRecord " <> show fs-  show (ValueSum shape c vs) = unwords ["ValueSum", show shape, show c, show vs]+  show (ValueArray shape vs) = unwords ["ValueArray", "(" <> show shape <> ")", "(" <> show vs <> ")"]+  show (ValueRecord fs) = "ValueRecord " <> "(" <> show fs <> ")"+  show (ValueSum shape c vs) = unwords ["ValueSum", "(" <> show shape <> ")", show c, "(" <> show vs <> ")"]   show ValueFun {} = "ValueFun _"   show ValueAcc {} = "ValueAcc _" 

@@ -153,7 +153,7 @@ -- | Return the dimensionality of a type.  For non-arrays, this is -- zero.  For a one-dimensional array it is one, for a two-dimensional -- it is two, and so forth.-arrayRank :: TypeBase () u -> Int+arrayRank :: TypeBase d u -> Int arrayRank = shapeRank . arrayShape  -- | Return the shape of a type - for non-arrays, this is 'mempty'.

@@ -0,0 +1,257 @@+module Futhark.Analysis.PrimExp.TableTests (tests) where++import Control.Monad.State.Strict+import Data.Map.Strict qualified as M+import Futhark.Analysis.PrimExp+import Futhark.Analysis.PrimExp.Table+import Futhark.IR.GPU+import Futhark.IR.GPUTests ()+import Futhark.IR.MC+import Futhark.IR.MCTests ()+import Test.Tasty+import Test.Tasty.HUnit++tests :: TestTree+tests = testGroup "AnalyzePrim" [stmToPrimExpsTests]++stmToPrimExpsTests :: TestTree+stmToPrimExpsTests =+  testGroup+    "stmToPrimExps"+    [stmToPrimExpsTestsGPU, stmToPrimExpsTestsMC]++stmToPrimExpsTestsGPU :: TestTree+stmToPrimExpsTestsGPU =+  testGroup+    "GPU"+    $ do+      let scope =+            M.fromList+              [ ("n_5142", FParamName "i64"),+                ("m_5143", FParamName "i64"),+                ("xss_5144", FParamName "[n_5142][m_5143]i64"),+                ("segmap_group_size_5201", LetName "i64"),+                ("segmap_usable_groups_5202", LetName "i64"),+                ("defunc_0_map_res_5203", LetName "[n_5142]i64"),+                ("defunc_0_f_res_5207", LetName "i64"),+                ("i_5208", IndexName Int64),+                ("acc_5209", FParamName "i64"),+                ("b_5210", LetName "i64"),+                ("defunc_0_f_res_5211", LetName "i64")+              ]+      [ testCase "BinOp" $ do+          let stm = "let {defunc_0_f_res_5211 : i64} = add64(acc_5209, b_5210)"+          let res = execState (stmToPrimExps scope stm) mempty+          let expected =+                M.fromList+                  [ ( "defunc_0_f_res_5211",+                      Just+                        ( BinOpExp+                            (Add Int64 OverflowWrap)+                            (LeafExp "acc_5209" (IntType Int64))+                            (LeafExp "b_5210" (IntType Int64))+                        )+                    )+                  ]+          res @?= expected,+        testCase "Index" $ do+          let stm = "let {b_5210 : i64} = xss_5144[gtid_5204, i_5208]"+          let res = execState (stmToPrimExps scope stm) mempty+          let expected = M.fromList [("b_5210", Nothing)]+          res @?= expected,+        testCase "Loop" $ do+          let stm = "let {defunc_0_f_res_5207 : i64} = loop {acc_5209 : i64} = {0i64} for i_5208:i64 < m_5143 do { {defunc_0_f_res_5211} }"+          let res = execState (stmToPrimExps scope stm) mempty+          let expected =+                M.fromList+                  [ ("defunc_0_f_res_5207", Nothing),+                    ("i_5208", Just (LeafExp "i_5208" (IntType Int64))),+                    ("acc_5209", Just (LeafExp "acc_5209" (IntType Int64)))+                  ]+          res @?= expected,+        testCase "Loop body" $ do+          let stm = "let {defunc_0_f_res_5207 : i64} = loop {acc_5209 : i64} = {0i64} for i_5208:i64 < m_5143 do { let {b_5210 : i64} = xss_5144[gtid_5204, i_5208] let {defunc_0_f_res_5211 : i64} = add64(acc_5209, b_5210) in {defunc_0_f_res_5211} }"+          let res = execState (stmToPrimExps scope stm) mempty+          let expected =+                M.fromList+                  [ ("defunc_0_f_res_5207", Nothing),+                    ("i_5208", Just (LeafExp "i_5208" (IntType Int64))),+                    ("acc_5209", Just (LeafExp "acc_5209" (IntType Int64))),+                    ("b_5210", Nothing),+                    ( "defunc_0_f_res_5211",+                      Just+                        ( BinOpExp+                            (Add Int64 OverflowWrap)+                            (LeafExp "acc_5209" (IntType Int64))+                            (LeafExp "b_5210" (IntType Int64))+                        )+                    )+                  ]+          res @?= expected,+        testCase "SegMap" $+          do+            let stm =+                  "let {defunc_0_map_res_5125 : [n_5142]i64} =\+                  \  segmap(thread; ; grid=segmap_usable_groups_5124; blocksize=segmap_group_size_5123)\+                  \  (gtid_5126 < n_5142) (~phys_tid_5127) : {i64} {\+                  \  return {returns lifted_lambda_res_5129} \+                  \}"+            let res = execState (stmToPrimExps scope stm) mempty+            let expected =+                  M.fromList+                    [ ("defunc_0_map_res_5125", Nothing),+                      ("gtid_5126", Just (LeafExp "gtid_5126" (IntType Int64)))+                    ]+            res @?= expected,+        testCase "SegMap body" $+          do+            let stm :: Stm GPU+                stm =+                  "let {defunc_0_map_res_5125 : [n_5142]i64} =\+                  \  segmap(thread; ; grid=segmap_usable_groups_5124; blocksize=segmap_group_size_5123)\+                  \  (gtid_5126 < n_5142) (~phys_tid_5127) : {i64} {\+                  \    let {eta_p_5128 : i64} =\+                  \      xs_5093[gtid_5126]\+                  \    let {lifted_lambda_res_5129 : i64} =\+                  \      add64(2i64, eta_p_5128)\+                  \    return {returns lifted_lambda_res_5129}\+                  \  }"+            let res = execState (stmToPrimExps scope stm) mempty+            let expected =+                  M.fromList+                    [ ("defunc_0_map_res_5125", Nothing),+                      ("gtid_5126", Just (LeafExp "gtid_5126" (IntType Int64))),+                      ("eta_p_5128", Nothing),+                      ( "lifted_lambda_res_5129",+                        Just+                          ( BinOpExp+                              (Add Int64 OverflowWrap)+                              (ValueExp (IntValue (Int64Value 2)))+                              (LeafExp "eta_p_5128" (IntType Int64))+                          )+                      )+                    ]+            res @?= expected+        ]++stmToPrimExpsTestsMC :: TestTree+stmToPrimExpsTestsMC =+  testGroup+    "MC"+    $ do+      let scope =+            M.fromList+              [ ("n_5142", FParamName "i64"),+                ("m_5143", FParamName "i64"),+                ("xss_5144", FParamName "[n_5142][5143]i64"),+                ("segmap_group_size_5201", LetName "i64"),+                ("segmap_usable_groups_5202", LetName "i64"),+                ("defunc_0_map_res_5203", LetName "[n_5142]i64"),+                ("defunc_0_f_res_5207", LetName "i64"),+                ("i_5208", IndexName Int64),+                ("acc_5209", FParamName "i64"),+                ("b_5210", LetName "i64"),+                ("defunc_0_f_res_5211", LetName "i64")+              ]+      [ testCase "BinOp" $ do+          let stm = "let {defunc_0_f_res_5211 : i64} = add64(acc_5209, b_5210)"+          let res = execState (stmToPrimExps scope stm) mempty+          let expected =+                M.fromList+                  [ ( "defunc_0_f_res_5211",+                      Just+                        ( BinOpExp+                            (Add Int64 OverflowWrap)+                            (LeafExp "acc_5209" (IntType Int64))+                            (LeafExp "b_5210" (IntType Int64))+                        )+                    )+                  ]+          res @?= expected,+        testCase "Index" $ do+          let stm = "let {b_5210 : i64} = xss_5144[gtid_5204, i_5208]"+          let res = execState (stmToPrimExps scope stm) mempty+          let expected = M.fromList [("b_5210", Nothing)]+          res @?= expected,+        testCase "Loop" $ do+          let stm = "let {defunc_0_f_res_5207 : i64} = loop {acc_5209 : i64} = {0i64} for i_5208:i64 < m_5143 do { {defunc_0_f_res_5211} }"+          let res = execState (stmToPrimExps scope stm) mempty+          let expected =+                M.fromList+                  [ ("defunc_0_f_res_5207", Nothing),+                    ("i_5208", Just (LeafExp "i_5208" (IntType Int64))),+                    ("acc_5209", Just (LeafExp "acc_5209" (IntType Int64)))+                  ]+          res @?= expected,+        testCase "Loop body" $ do+          let stm =+                "\+                \let {defunc_0_f_res_5207 : i64} =\+                \  loop {acc_5209 : i64} = {0i64}\+                \  for i_5208:i64 < m_5143 do {\+                \    let {b_5210 : i64} =\+                \      xss_5144[gtid_5204, i_5208]\+                \    let {defunc_0_f_res_5211 : i64} =\+                \      add64(acc_5209, b_5210)\+                \    in {defunc_0_f_res_5211}\+                \  }"+          let res = execState (stmToPrimExps scope stm) mempty+          let expected =+                M.fromList+                  [ ("defunc_0_f_res_5207", Nothing),+                    ("i_5208", Just (LeafExp "i_5208" (IntType Int64))),+                    ("acc_5209", Just (LeafExp "acc_5209" (IntType Int64))),+                    ("b_5210", Nothing),+                    ( "defunc_0_f_res_5211",+                      Just+                        ( BinOpExp+                            (Add Int64 OverflowWrap)+                            (LeafExp "acc_5209" (IntType Int64))+                            (LeafExp "b_5210" (IntType Int64))+                        )+                    )+                  ]+          res @?= expected,+        testCase "SegMap" $ do+          let stm =+                "let {defunc_0_map_res_5125 : [n_5142]i64} =\+                \  segmap()\+                \  (gtid_5126 < n_5142) (~flat_tid_5112) : {i64} {\+                \    return {returns lifted_lambda_res_5129}\+                \  }"+          let res = execState (stmToPrimExps scope stm) mempty+          let expected =+                M.fromList+                  [ ("defunc_0_map_res_5125", Nothing),+                    ("gtid_5126", Just (LeafExp "gtid_5126" (IntType Int64)))+                  ]+          res @?= expected,+        testCase "SegMap body" $ do+          let stm :: Stm MC+              stm =+                "let {defunc_0_map_res_5125 : [n_5142]i64} =\+                \  segmap()\+                \  (gtid_5126 < n_5142) (~flat_tid_5112) : {i64} {\+                \    let {eta_p_5128 : i64} =\+                \      xs_5093[gtid_5126]\+                \    let {lifted_lambda_res_5129 : i64} =\+                \      add64(2i64, eta_p_5128)\+                \    return {returns lifted_lambda_res_5129}\+                \  }"+          let res = execState (stmToPrimExps scope stm) mempty+          let expected =+                M.fromList+                  [ ("defunc_0_map_res_5125", Nothing),+                    ("gtid_5126", Just (LeafExp "gtid_5126" (IntType Int64))),+                    ("eta_p_5128", Nothing),+                    ( "lifted_lambda_res_5129",+                      Just+                        ( BinOpExp+                            (Add Int64 OverflowWrap)+                            (ValueExp (IntValue (Int64Value 2)))+                            (LeafExp "eta_p_5128" (IntType Int64))+                        )+                    )+                  ]+          res @?= expected+        ]

@@ -59,9 +59,9 @@                 ("bar", [Pure "[?0]i64"])               ]           )-          @?= ( [Pure "[?0]i64"],+          @?= ( [Pure "[?0]i64", Pure "[?0]i64"],                 [ ("bar", [0]),-                  ("foo", [0])+                  ("foo", [1])                 ]               ),       testCase

@@ -0,0 +1,241 @@+module Futhark.Optimise.ArrayLayout.AnalyseTests (tests) where++import Data.Map.Strict qualified as M+import Futhark.Analysis.AccessPattern+import Futhark.IR.GPU+import Futhark.IR.GPUTests ()+import Futhark.IR.SyntaxTests ()+import Test.Tasty+import Test.Tasty.HUnit++tests :: TestTree+tests = testGroup "Analyse" [analyseStmTests]++analyseStmTests :: TestTree+analyseStmTests =+  testGroup+    "analyseStm"+    [analyseIndexTests, analyseDimAccessesTests]++analyseIndexTests :: TestTree+analyseIndexTests =+  testGroup+    "analyseIndex"+    $ do+      let arr_name = "xss_5144"+      -- ============================= TestCase0 =============================+      -- Most simple case where we want to manifest an array, hence, we record+      -- the Index in the IndexTable.+      let testCase0 = testCase "2D manifest" $ do+            let ctx =+                  mempty+                    { parents =+                        [ SegOpName (SegmentedMap "defunc_0_map_res_5204"),+                          LoopBodyName "defunc_0_f_res_5208"+                        ],+                      assignments =+                        M.fromList+                          [ ("gtid_5205", VariableInfo mempty 0 mempty ThreadID),+                            ("i_5209", VariableInfo mempty 1 mempty LoopVar)+                          ]+                    }+            let patternNames = ["b_5211"]+            let dimFixes =+                  [ DimFix (Var "gtid_5205"),+                    DimFix (Var "i_5209")+                  ]+            let indexTable =+                  M.fromList+                    [ ( SegmentedMap "defunc_0_map_res_5204",+                        M.fromList+                          [ ( (arr_name, [], [0 .. 1]),+                              M.fromList+                                [ ( "b_5211",+                                    [ DimAccess (M.fromList [("gtid_5205", Dependency 0 ThreadID)]) (Just "gtid_5205"),+                                      DimAccess (M.fromList [("i_5209", Dependency 1 LoopVar)]) (Just "i_5209")+                                    ]+                                  )+                                ]+                            )+                          ]+                      )+                    ]+            let (_, indexTable') = analyseIndex ctx patternNames arr_name dimFixes+            indexTable' @?= indexTable++      -- ============================= TestCase2 =============================+      -- We don't want to manifest an array with only one dimension, so we don't+      -- record anything in the IndexTable.+      let testCase1 = testCase "1D manifest" $ do+            let ctx =+                  mempty+                    { parents =+                        [ SegOpName (SegmentedMap "defunc_0_map_res_5204"),+                          LoopBodyName "defunc_0_f_res_5208"+                        ]+                    }+            let patternNames = ["b_5211"]+            let dimFixes = [DimFix "i_5209"]++            let (_, indexTable') = analyseIndex ctx patternNames arr_name dimFixes+            indexTable' @?= mempty++      -- ============================= TestCase1 =============================+      -- We don't want to record anything to the IndexTable when the array is+      -- not accessed inside a SegMap+      -- TODO: Create a similar one for MC with loops+      let testCase2 = testCase "Not inside SegMap" $ do+            let ctx = mempty+            let patternNames = ["b_5211"]+            let dimFixes =+                  [ DimFix "gtid_5205",+                    DimFix "i_5209"+                  ]+            let (_, indexTable') = analyseIndex ctx patternNames arr_name dimFixes+            indexTable' @?= mempty++      -- ============================= TestCase3 =============================+      -- If an array is allocated inside a loop or SegMap, we want to record that+      -- information in the ArrayName of the IndexTable.+      let testCase3 = testCase "Allocated inside SegMap" $ do+            let parents' =+                  [ SegOpName (SegmentedMap "defunc_0_map_res_5204"),+                    LoopBodyName "defunc_0_f_res_5208"+                  ]+            let ctx =+                  mempty+                    { parents = parents',+                      assignments =+                        M.fromList+                          [ ("gtid_5205", VariableInfo mempty 0 mempty ThreadID),+                            ("i_5209", VariableInfo mempty 1 mempty LoopVar),+                            (arr_name, VariableInfo mempty 0 parents' Variable)+                          ]+                    }+            let patternNames = ["b_5211"]+            let dimFixes =+                  [ DimFix "gtid_5205",+                    DimFix "i_5209"+                  ]+            let indexTable =+                  M.fromList+                    [ ( SegmentedMap "defunc_0_map_res_5204",+                        M.fromList+                          [ ( (arr_name, parents', [0 .. 1]),+                              M.fromList+                                [ ( "b_5211",+                                    [ DimAccess (M.fromList [("gtid_5205", Dependency 0 ThreadID)]) (Just "gtid_5205"),+                                      DimAccess (M.fromList [("i_5209", Dependency 1 LoopVar)]) (Just "i_5209")+                                    ]+                                  )+                                ]+                            )+                          ]+                      )+                    ]+            let (_, indexTable') = analyseIndex ctx patternNames arr_name dimFixes+            indexTable' @?= indexTable++      -- ============================= TestCase4 =============================+      -- If the vars in the index are temporaries, we want to reduce them to+      -- to the thread IDs and or loop counters they are functions of.+      let testCase4 = testCase "Reduce dependencies" $ do+            let ctx =+                  mempty+                    { parents =+                        [ SegOpName (SegmentedMap "defunc_0_map_res_5204"),+                          LoopBodyName "defunc_0_f_res_5208"+                        ],+                      assignments =+                        M.fromList+                          [ ("gtid_5205", VariableInfo mempty 0 mempty ThreadID),+                            ("i_5209", VariableInfo mempty 1 mempty LoopVar),+                            ("tmp0_5210", VariableInfo (namesFromList ["gtid_5205"]) 2 mempty Variable),+                            ("tmp1_5211", VariableInfo (namesFromList ["i_5209"]) 3 mempty Variable),+                            ("k_5212", VariableInfo mempty 1 mempty ConstType)+                          ]+                    }+            let patternNames = ["b_5211"]+            let dimFixes =+                  [ DimFix "tmp0_5210",+                    DimFix "tmp1_5211",+                    DimFix "k_5212"+                  ]+            let indexTable =+                  M.fromList+                    [ ( SegmentedMap "defunc_0_map_res_5204",+                        M.fromList+                          [ ( (arr_name, [], [0 .. 2]),+                              M.fromList+                                [ ( "b_5211",+                                    [ DimAccess (M.fromList [("gtid_5205", Dependency 0 ThreadID)]) (Just "tmp0_5210"),+                                      DimAccess (M.fromList [("i_5209", Dependency 1 LoopVar)]) (Just "tmp1_5211"),+                                      DimAccess mempty (Just "k_5212")+                                    ]+                                  )+                                ]+                            )+                          ]+                      )+                    ]+            let (_, indexTable') = analyseIndex ctx patternNames arr_name dimFixes+            indexTable' @?= indexTable++      [testCase0, testCase1, testCase2, testCase3, testCase4]++analyseDimAccessesTests :: TestTree+analyseDimAccessesTests = testGroup+  "analyseDimAccesses"+  $ do+    let testCase0 = testCase "Fold" $ do+          let indexTable =+                M.fromList+                  [ ( SegmentedMap "defunc_0_map_res_5204",+                      M.fromList+                        [ ( ("xss_5144", [], [0, 1]),+                            M.fromList+                              [ ( "b_5211",+                                  [ DimAccess (M.fromList [("gtid_5205", Dependency 0 ThreadID)]) (Just "gtid_5205"),+                                    DimAccess (M.fromList [("i_5209", Dependency 1 LoopVar)]) (Just "i_5209")+                                  ]+                                )+                              ]+                          )+                        ]+                    )+                  ]+          let indexTable' = (analyseDimAccesses @GPU) prog0+          indexTable' @?= indexTable++    [testCase0]+  where+    prog0 :: Prog GPU+    prog0 =+      "\+      \entry(\"main\",\+      \      {xss: [][]i64},\+      \      {[]i64})\+      \  entry_main (n_5142 : i64,\+      \              m_5143 : i64,\+      \              xss_5144 : [n_5142][m_5143]i64)\+      \  : {[n_5142]i64#([2], [0])} = {\+      \  let {segmap_group_size_5202 : i64} =\+      \    get_size(segmap_group_size_5190, thread_block_size)\+      \  let {segmap_usable_groups_5203 : i64} =\+      \    sdiv_up64(n_5142, segmap_group_size_5202)\+      \  let {defunc_0_map_res_5204 : [n_5142]i64} =\+      \    segmap(thread; ; grid=segmap_usable_groups_5203; blocksize=segmap_group_size_5202)\+      \    (gtid_5205 < n_5142) (~phys_tid_5206) : {i64} {\+      \      let {defunc_0_f_res_5208 : i64} =\+      \        loop {acc_5210 : i64} = {0i64}\+      \        for i_5209:i64 < m_5143 do {\+      \          let {b_5211 : i64} =\+      \            xss_5144[gtid_5205, i_5209]\+      \          let {defunc_0_f_res_5212 : i64} =\+      \            add64(acc_5210, b_5211)\+      \          in {defunc_0_f_res_5212}\+      \        }\+      \      return {returns defunc_0_f_res_5208}\+      \    }\+      \  in {defunc_0_map_res_5204}\+      \}"

@@ -0,0 +1,163 @@+module Futhark.Optimise.ArrayLayout.LayoutTests (tests) where++import Data.Map.Strict qualified as M+import Futhark.Analysis.AccessPattern+import Futhark.Analysis.PrimExp+import Futhark.FreshNames+import Futhark.IR.GPU (GPU)+import Futhark.IR.GPUTests ()+import Futhark.Optimise.ArrayLayout.Layout+import Language.Futhark.Core+import Test.Tasty+import Test.Tasty.HUnit++tests :: TestTree+tests =+  testGroup+    "Layout"+    [commonPermutationEliminatorsTests]++commonPermutationEliminatorsTests :: TestTree+commonPermutationEliminatorsTests =+  testGroup+    "commonPermutationEliminators"+    [permutationTests, nestTests, dimAccessTests, constIndexElimTests]++permutationTests :: TestTree+permutationTests =+  testGroup "Permutations" $+    do+      -- This isn't the way to test this, in reality we should provide realistic+      -- access patterns that might result in the given permutations.+      -- Luckily we only use the original access for one check atm.+      [ testCase (unwords [show perm, "->", show res]) $+          commonPermutationEliminators perm [] @?= res+        | (perm, res) <-+            [ ([0], True),+              ([1, 0], False),+              ([0, 1], True),+              ([0, 0], True),+              ([1, 1], True),+              ([1, 2, 0], False),+              ([2, 0, 1], False),+              ([0, 1, 2], True),+              ([1, 0, 2], True),+              ([2, 1, 0], True),+              ([2, 2, 0], True),+              ([2, 1, 1], True),+              ([1, 0, 1], True),+              ([0, 0, 0], True),+              ([0, 1, 2, 3, 4], True),+              ([1, 0, 2, 3, 4], True),+              ([2, 3, 0, 1, 4], True),+              ([3, 4, 2, 0, 1], True),+              ([2, 3, 4, 0, 1], False),+              ([1, 2, 3, 4, 0], False),+              ([3, 4, 0, 1, 2], False)+            ]+        ]++nestTests :: TestTree+nestTests = testGroup "Nests" $+  do+    let names = generateNames 2+    [ testCase (unwords [args, "->", show res]) $+        commonPermutationEliminators [1, 0] nest @?= res+      | (args, nest, res) <-+          [ ("[]", [], False),+            ("[CondBodyName]", [CondBodyName] <*> names, False),+            ("[SegOpName]", [SegOpName . SegmentedMap] <*> names, True),+            ("[LoopBodyName]", [LoopBodyName] <*> names, False),+            ("[SegOpName, CondBodyName]", [SegOpName . SegmentedMap, CondBodyName] <*> names, True),+            ("[CondBodyName, LoopBodyName]", [CondBodyName, LoopBodyName] <*> names, False)+          ]+      ]++dimAccessTests :: TestTree+dimAccessTests = testGroup "DimAccesses" [] -- TODO: Write tests for the part of commonPermutationEliminators that checks the complexity of the DimAccesses.++constIndexElimTests :: TestTree+constIndexElimTests =+  testGroup+    "constIndexElimTests"+    [ testCase "gpu eliminates indexes with constant in any dim" $ do+        let primExpTable =+              M.fromList+                [ ("gtid_4", Just (LeafExp "n_4" (IntType Int64))),+                  ("i_5", Just (LeafExp "n_4" (IntType Int64)))+                ]+        layoutTableFromIndexTable primExpTable accessTableGPU @?= mempty,+      testCase "gpu ignores when not last" $ do+        let primExpTable =+              M.fromList+                [ ("gtid_4", Just (LeafExp "gtid_4" (IntType Int64))),+                  ("gtid_5", Just (LeafExp "gtid_5" (IntType Int64))),+                  ("i_6", Just (LeafExp "i_6" (IntType Int64)))+                ]+        layoutTableFromIndexTable primExpTable accessTableGPUrev+          @?= M.fromList+            [ ( SegmentedMap "mapres_1",+                M.fromList+                  [ ( ("a_2", [], [0, 1, 2, 3]),+                      M.fromList [("A_3", [2, 3, 0, 1])]+                    )+                  ]+              )+            ]+    ]+  where+    accessTableGPU :: IndexTable GPU+    accessTableGPU =+      singleAccess+        [ singleParAccess 0 "gtid_4",+          DimAccess mempty Nothing,+          singleSeqAccess 1 "i_5"+        ]++    accessTableGPUrev :: IndexTable GPU+    accessTableGPUrev =+      singleAccess+        [ singleParAccess 1 "gtid_4",+          singleParAccess 2 "gtid_5",+          singleSeqAccess 0 "i_5",+          singleSeqAccess 2 "gtid_4"+        ]++singleAccess :: [DimAccess rep] -> IndexTable rep+singleAccess dims =+  M.fromList+    [ ( sgOp,+        M.fromList+          [ ( ("A_2", [], [0, 1, 2, 3]),+              M.fromList+                [ ( "a_3",+                    dims+                  )+                ]+            )+          ]+      )+    ]+  where+    sgOp = SegmentedMap {vnameFromSegOp = "mapres_1"}++singleParAccess :: Int -> VName -> DimAccess rep+singleParAccess level name =+  DimAccess+    (M.singleton name $ Dependency level ThreadID)+    (Just name)++singleSeqAccess :: Int -> VName -> DimAccess rep+singleSeqAccess level name =+  DimAccess+    (M.singleton name $ Dependency level LoopVar)+    (Just name)++generateNames :: Int -> [VName]+generateNames count = do+  let (name, source) = newName blankNameSource "i_0"+  fst $ foldl f ([name], source) [1 .. count - 1]+  where+    f (names, source) _ = do+      let (name, source') = newName source (last names)+      (names ++ [name], source')

@@ -0,0 +1,15 @@+module Futhark.Optimise.ArrayLayoutTests (tests) where++import Futhark.Analysis.PrimExp.TableTests qualified+import Futhark.Optimise.ArrayLayout.AnalyseTests qualified+import Futhark.Optimise.ArrayLayout.LayoutTests qualified+import Test.Tasty++tests :: TestTree+tests =+  testGroup+    "OptimizeArrayLayoutTests"+    [ Futhark.Optimise.ArrayLayout.AnalyseTests.tests,+      Futhark.Optimise.ArrayLayout.LayoutTests.tests,+      Futhark.Analysis.PrimExp.TableTests.tests+    ]

@@ -8,6 +8,7 @@ import Futhark.IR.PropTests qualified import Futhark.IR.Syntax.CoreTests qualified import Futhark.Internalise.TypesValuesTests qualified+import Futhark.Optimise.ArrayLayoutTests qualified import Futhark.Optimise.MemoryBlockMerging.GreedyColoringTests qualified import Futhark.Pkg.SolveTests qualified import Language.Futhark.PrimitiveTests qualified@@ -31,7 +32,8 @@       Language.Futhark.PrimitiveTests.tests,       Futhark.Optimise.MemoryBlockMerging.GreedyColoringTests.tests,       Futhark.Analysis.AlgSimplifyTests.tests,-      Language.Futhark.TypeCheckerTests.tests+      Language.Futhark.TypeCheckerTests.tests,+      Futhark.Optimise.ArrayLayoutTests.tests     ]  main :: IO ()